kernels: support conditions on use_kernel_forward_from_hub

[danieldk]

Add and optional condition argument to use_kernel_forward_from_hub.
This condition accepts an instance of the layer and returns whether
kernelization of the layer should proceed. For instance,

@use_kernel_forward_from_hub(
    "SwiGLUMLP",
    condition=lambda module: module.config.hidden_act == "silu",
)
class MyMLP(nn.Module):
    ...

would only kernelize the layer using a mapping registered for SwiGLUMLP
its hidden_act configuration is silu.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

Coverage report — kernels/

Measured on: Python 3.10 / Torch 2.13.0.
Other CI configurations are not included in this number.
Hardware-gated code paths (ROCm/XPU/NPU/Darwin/Windows) are excluded or unreachable on the Linux+CUDA runner.

Total coverage: 86.3% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	14	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	78	9	88%	47, 53-54, 57-58, 97, 119, 130, 136
src/kernels/backends.py	212	62	71%	40, 44, 48-51, 68, 90, 108, 117, 121, 125-127, 148, 157, 161, 165-167, 188, 199, 201, 208-211, 224, 228, 232-252, 260, 283-303
src/kernels/compat.py	8	1	88%	5
src/kernels/deps.py	67	1	99%	57
src/kernels/hf_hub.py	62	3	95%	18, 20, 114
src/kernels/importer.py	50	3	94%	110, 114, 117
src/kernels/install.py	21	7	67%	78-102
src/kernels/layer/__init__.py	6	0	100%
src/kernels/layer/_interval_tree.py	103	4	96%	23, 52, 147, 150
src/kernels/layer/device.py	48	14	71%	42, 47-49, 91, 96-98, 101, 149, 152, 155-157
src/kernels/layer/func.py	86	6	93%	90, 121, 193, 313, 338, 366
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/kernelize.py	80	8	90%	258, 293, 301-302, 308, 312, 328-330
src/kernels/layer/layer.py	216	14	94%	177, 222, 247, 379, 459-460, 481, 489, 500, 529, 533, 546, 599, 629
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	144	42	71%	27, 33, 36-43, 63-64, 70, 73-76, 90, 94, 103-104, 110, 113-116, 123-124, 130, 133-136, 143-144, 150, 153-156, 163-164, 170, 173-176, 257
src/kernels/load.py	61	2	97%	297, 335
src/kernels/locking.py	89	64	28%	35-83, 91-98, 102-125, 137, 152-159, 165-175, 179-186
src/kernels/python_deps.py	58	6	90%	59-60, 64-65, 101, 104
src/kernels/resolver.py	135	2	99%	174, 180
src/kernels/status.py	49	2	96%	23, 81
src/kernels/variants.py	278	19	93%	64, 95, 116, 146, 255-256, 298-301, 303, 387-394, 400-406, 437-443, 455-461
src/kernels/verify.py	88	1	99%	32
TOTAL	1978	271	86%

Updated by the Test kernels workflow on commit 8b006b67e5c601e839e3528ffcc5e91208a739b9.

[vasqu]

Careful approval but I sanity checked on huggingface/transformers#48335 (llama) and it worked as expected, e.g. when passing a new config with activation == gelu, we no longer kernelize that module

Ofc, docs are missing but since it's a draft I expect it tbh 🤗

[danieldk]

Thanks for testing! I'll add the docs.

[sayakpaul]

There should also be a test to exercise the path:

if cond and not cond(module):
   ...

and check if the logging info was as expected.

[sayakpaul]

Just single comment regarding testing. Otherwise, looks good.

[drbh]

lgtm to me too - should be good to merge after the comments above