kernels: add support for layer selectors

[danieldk]

Until now, one could only register layers by device (+ capability) and the kernelization mode. It can be useful to select kernels in a more fine-grained manner. For instance, the user might want to select the best kernel based on input shapes, available device memory, and other properties. This change makes it possible to register a selector for a kernel, where the selector is a user-provided function that determines the layer repository to use based on the model, device, and module being kernelized.

For example:

def select_silu_and_mul(module, *, device_type, mode):
    if device_type.type != "cuda":
        return None
    repo = LayerRepository(
        repo_id="kernels-community/activation",
        layer_name="SiluAndMul",
        version=1,
    )
    return repo, Mode.FALLBACK

with use_kernel_mapping({"SiluAndMul": select_silu_and_mul}):
    model = kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE, device="cuda")

Fixes #823.

[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: 87.9% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	15	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	130	14	89%	53, 59-60, 63-64, 102, 165-170, 199, 219
src/kernels/archs.py	56	1	98%	94
src/kernels/backends.py	213	62	71%	42, 46, 50-53, 70, 92, 110, 119, 123, 127-129, 150, 159, 163, 167-169, 190, 201, 203, 210-213, 226, 230, 234-254, 262, 285-305
src/kernels/compat.py	9	1	89%	5
src/kernels/deps.py	70	1	99%	56
src/kernels/hf_hub.py	63	2	97%	21, 23
src/kernels/importer.py	44	5	89%	80, 84, 87, 101-102
src/kernels/install.py	21	7	67%	76-100
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	85	6	93%	90, 115, 191, 311, 338, 368
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/kernelize.py	106	8	92%	340, 375, 383-384, 390, 394, 410-412
src/kernels/layer/layer.py	217	13	94%	182, 229, 256, 390, 470-471, 495, 503, 545, 549, 562, 615, 645
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	147	30	80%	36-43, 74, 90, 94, 103-104, 110, 113-116, 123-124, 130, 133-136, 143-144, 150, 153-156, 174, 257
src/kernels/load.py	71	2	97%	338, 378
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	156	2	99%	220, 226
src/kernels/status.py	50	2	96%	25, 79
src/kernels/validate.py	88	5	94%	9, 100, 167, 190-191
src/kernels/variants.py	278	19	93%	65, 96, 117, 147, 256-257, 299-302, 304, 388-394, 400-406, 437-443, 455-461
src/kernels/verify.py	127	6	95%	46, 202-204, 318-319
TOTAL	2275	275	88%

Updated by the Test kernels workflow on commit df601b432bc23eecfb69f9e176ff53ece78f9e0f.