[RFC] Load compute kernels from the Hugging Face Hub

[adarshxs]

[RFC] Load compute kernels from the Hugging Face Hub

Overview

This RFC proposes adopting huggingface/kernels as an optional kernel source for miles. Kernels are published as Hub repos with one prebuilt variant per (torch, CUDA, C++ ABI, arch, OS); the client resolves a variant at load time and imports it from the HF cache. No compiler is needed on the target machine.

miles's docker/Dockerfile currently installs flash_attn, flash_attn_3, apex, fast_hadamard_transform, causal_conv1d, mamba_ssm and transformer_engine from prebuilt wheels produced out of band, plus FlashQLA from source. Image builds might be very long and any kernel change requires a rebuild.

Scope is the FSDP backend (--train-backend fsdp), opt-in by default. kernels is already present in every miles image: SGLang v0.5.16 declares kernels>=0.14.1,<0.15 as a dependency, and ships a Hub FlashAttention-3 in sgl-project/sglang#20796.

Phase 0: attention, no code changes. Works already

--attn-implementation is a free-form str with no choices, passed directly to from_pretrained for both the policy and the reference model. transformers 5.12.1 resolves Hub repo IDs there via is_kernel() in integrations/hub_kernels.py.

python train.py --train-backend fsdp --attn-implementation kernels-community/flash-attn2@v2

• kernels-community/flash-attn2 exposes bwd and varlen_bwd, so it covers the training backward pass.
• The packed-sequence path is unchanged: cu_seqlens is derived by transformers in prepare_fa_kwargs_from_position_ids, and load_and_register_attn_kernel binds transformers' own flash_attention_forward whenever the kernel exposes flash_attn_varlen_func.

Verified on one H100, in a venv resolved from sglang==0.5.16 (torch 2.11.0+cu130, transformers 5.12.1, kernels 0.14.1) with radixark/miles @ HEAD installed from source. Qwen3-4B driven through miles' own construction sequence (apply_class_patches to apply_packing(..., "config") to from_pretrained to apply_post_load_fixups): load, forward and backward succeed, gradients finite, state_dict intact, run-to-run deterministic. Against an sdpa fp32 block-diagonal reference, per-document error does not grow with document index and hub-vs-fp32 max_abs matches the sdpa-bf16-vs-fp32 floor exactly.

Phase 1: layer kernels for FSDP

New fields on the FSDP argument dataclass:

kernel_backend: str = "native"        # {"native", "hub"}
kernel_mapping_path: str = ""         # dotted path resolved by miles.utils.misc.load_function
kernel_strict: bool = False           # raise instead of falling back to the native layer

Mappings come from a callable so users can substitute their own without patching miles. Default in miles_plugins/kernels/presets.py:

def default_mapping(args):
    from kernels import LayerRepository

    if getattr(args, "true_on_policy_mode", False) or args.deterministic_mode:
        return {}
    return {
        "MegaBlocksMoeMLP": {
            "cuda": LayerRepository(
                repo_id="kernels-community/megablocks",
                layer_name="MegaBlocksMoeMLP",
                version=1,
            ),
        },
    }

Loader in fsdp_utils/kernels/hub.py. kernels is imported lazily so a node without a matching variant cannot break import miles:

def kernelize_model(model, args, *, device: str = "cuda"):
    if args.kernel_backend != "hub":
        return model

    from kernels import Mode, kernelize, use_kernel_mapping

    mapping = load_mapping(args)
    if not mapping:
        return model

    # `device=` is required: under `_get_init_weight_context_manager()` parameters may be
    # on the meta device on non-zero ranks, and device inference would yield "meta".
    with use_kernel_mapping(mapping):
        return kernelize(model, mode=Mode.TRAINING, device=device,
                         use_fallback=not args.kernel_strict)

Mode.TRAINING unconditionally, including for the torch.no_grad() log-prob pass: old_log_probs and the training forward must use the same kernel, or exp(new_logp - old_logp) is not 1 at step 0.

Call site, in FSDPTrainRayActor.init():

 apply_packing(model, self.hf_config, "post_load")
+model = kernelize_model(model, self.args, device="cuda")
 model.train()
 full_state = model.state_dict()

 model = apply_fsdp2(model, mesh=get_parallel_state().get_mesh("fsdp"), ...)

Before apply_fsdp2. kernelize() rebinds only forward on instances, so full_state, the _no_split_modules walk, and the DTensor gather in update_weight_utils.py are unaffected. _create_ref_model takes the same call.

Phase 2: prebuilt cache for offline clusters

[tool.kernels.dependencies]
"kernels-community/flash-attn2" = 2
"kernels-community/megablocks" = 1

RUN kernels lock . && kernels download .
COPY --from=framework /root/.cache/huggingface /root/.cache/huggingface

export HF_HUB_OFFLINE=1
python train.py --train-backend fsdp --kernel-backend hub \
    --attn-implementation kernels-community/flash-attn2@v2

kernels.lock pins a commit SHA plus a per-variant SHA-256, so compute nodes need no egress and every rank loads a byte-identical kernel.

Numerics

miles reports bit-wise identical training and inference log probs under --true-on-policy-mode, which requires the training-side kernel to match SGLang's build exactly. Until that equivalence is established per kernel, we will keep the two mutually exclusive:

if args.kernel_backend == "hub" and (
    getattr(args, "true_on_policy_mode", False) or args.deterministic_mode
):
    raise ValueError(
        "--kernel-backend hub is incompatible with --true-on-policy-mode / --deterministic-mode"
    )

Follow-ups

Megatron backend via miles_plugins/models/hf_attention.py; rollout-side alignment through SGLang's existing Hub FA3 path; and publishing miles' fused-MoE Triton kernels (fsdp_utils/kernels/) to the Hub as they are pure Triton, so they build as a noarch variant with no torch-version matrix.

References

kernels docs · layers API · locking · transformers hub_kernels.py · sglang#20796

[alexnails]

Initial pass: cc@Zhichenzzz

• MegaBlocksMoeMLP is only used for GPT OSS? https://github.com/search?q=repo%3Ahuggingface%2Ftransformers+MegaBlocksMoeMLP&type=code
  • was this just example or do you have a model list in mind for first support?
• _create_ref_model will need some post actions it seems
• claude: when context_parallel_size > 1, ring_flash_attn.substitute_hf_flash_attn() / update_ring_flash_attn_params() are called and the kernels substitution would be skipped. Functionality that needs to be skipped

[adarshxs]

MegaBlocksMoeMLP is only used for GPT OSS? https://github.com/search?q=repo%3Ahuggingface%2Ftransformers+MegaBlocksMoeMLP&type=code
was this just example or do you have a model list in mind for first support?

yes it was an example, but probably out of scope here(needs more research for layer level). i think we are better off integrating modules from HF kernels as a first off.

from kernels import get_kernel
ops = get_kernel("kernels-community/flash-attn2", version=2)

function level dispatchers miles already has(--linear-attention-backend {fla,flashqla}, --dsa-attention-backend {megatron,tilelang}, --miles-dsa-topk-backend {torch,flashinfer}) can greatly benefit from this.

[Zhichenzzz]

Thanks @adarshxs and @alexnails for the support! I will look into it after this week (so busy with some stuff lol

[adarshxs]

no problem :P