kernel-fun 0.2.0.dev1__py3-none-any.whl

kernel_fun/kda/_provenance.py

@@ -0,0 +1,18 @@
+"""Generated by tools/vendor.py — what this family was cut from.
+
+Hand-edited after vendoring (knobs frozen, dead branches removed, call cache
+unified), so these hashes identify the SOURCE, not the files here. drift.py
+compares them against the ladder's current HEAD.
+"""
+
+SOURCE_COMMIT = "6fc6309"
+SOURCE_FILES = {
+    "kernels/kda/ideas/002-cute-bwd/kernel_dhu.py": "eddf8ded92731b6376ec760eec0b2631e35188aebf6f7b32b7b1f2999ff72e2e",
+    "kernels/kda/ideas/002-cute-bwd/kernel_fwd.py": "0387d74fc85e4bac541439e54b306b112d6e70b3f52e4c71e06aa126a1bca963",
+    "kernels/kda/ideas/002-cute-bwd/kernel_intra.py": "a18915e9ed4606edadbafacb2d146b99dc7d816ea749492985f4f6b165638878",
+    "kernels/kda/ideas/002-cute-bwd/kernel_scan.py": "3aed945213f0741fda017aae9ca904009d7ae606b9a690652bdc0c57cf00817b",
+    "kernels/kda/ideas/002-cute-bwd/kernel_wy2.py": "eee08ab2008ab7bc85aa861ba17f638e42a471abe901a68e97f62f6a8ff5392c",
+    "kernels/kda/ideas/003-intra-mma/kernel_intra_cute.py": "3e82e01b381a98711d9035077279c8b888d562fa431cf0fa180eb02deab89322",
+    "kernels/kda/ideas/004-fwd-block/kernel_fwd_intra_triton.py": "5f69a5e36746f6e96fa6c173d579b71a187c11251e82e7f686fdf0d4f6a42e00",
+    "kernels/kda/ideas/005-wy-transposed/kernel_wy_t.py": "0d1ba89945c4fcdaef9a72159e4eca98184f7d291e170f3f6f2730ea178bc7f7",
+}

kernel_fun/kda/autograd.py

@@ -0,0 +1,114 @@
+"""The autograd Function behind `chunk_kda`.
+
+Residuals are fla's own set — q, k, v, g2, beta, Aqk, Akk, h0 (+ the l2norm rstds) — so
+peak activation memory matches fla's, which is a requirement for a drop-in and not an
+accident. Two deliberate differences, both fla's own behaviour rather than ours:
+
+  - with the fused gate, g2 is NOT saved. The backward recomputes it from the raw (and
+    half-size) g, exactly as fla's recompute path does. Saving it would cost 1 GiB per
+    layer at prod8192 for a tensor a single kernel launch reproduces.
+  - dq/dk come back already in q's dtype when the intra kernel can emit them that way
+    (HV == H), so the casts below are no-ops rather than two extra launches.
+
+Everything here runs under `torch.compiler.disable` at the entry point, so nothing in this
+file is traced. See ops.py for why.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .._common import support
+from . import chain
+
+
+def _make_fn():
+    """Built lazily: fla's decorators import triton, and importing this package must stay
+    cheap on a machine that will never call it."""
+    from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
+
+    class ChunkKDAFunction(torch.autograd.Function):
+        @staticmethod
+        @input_guard
+        @autocast_custom_fwd
+        def forward(ctx, q, k, v, g, beta, h0, scale, chunk_size, use_qk_l2norm,
+                    A_log, dt_bias, lower_bound, h0_was_none, warm_sig):
+            q_rstd = k_rstd = None
+            if use_qk_l2norm:
+                from fla.modules.l2norm import l2norm_fwd
+
+                q, q_rstd = l2norm_fwd(q)
+                k, k_rstd = l2norm_fwd(k)
+
+            g2 = chain.gate_cumsum(g, A_log, dt_bias, chunk_size, lower_bound)
+            o, ht, Aqk, Akk = chain.forward(q, k, v, g2, beta, h0, scale, chunk_size)
+
+            ctx.save_for_backward(
+                q, q_rstd, k, k_rstd, v,
+                None if A_log is not None else g2,
+                beta, Aqk, Akk, h0,
+                g if A_log is not None else None, A_log, dt_bias,
+            )
+            ctx.scale = scale
+            ctx.chunk_size = chunk_size
+            ctx.lower_bound = lower_bound
+            ctx.h0_was_none = h0_was_none
+            ctx.warm_sig = warm_sig
+            return o, ht
+
+        @staticmethod
+        @input_guard
+        @autocast_custom_bwd
+        def backward(ctx, do, dht):
+            (q, q_rstd, k, k_rstd, v, g2, beta, Aqk, Akk, h0,
+             g_org, A_log, dt_bias) = ctx.saved_tensors
+            if g2 is None:
+                g2 = chain.gate_cumsum(
+                    g_org, A_log, dt_bias, ctx.chunk_size, ctx.lower_bound
+                )
+            dq, dk, dv, db, dg, dh0 = chain.backward(
+                q=q, k=k, v=v, g2=g2, beta=beta, Aqk=Aqk, Akk=Akk, h0=h0,
+                do=do.contiguous(),
+                # dht is None when the caller never asked for a final state and nothing
+                # differentiated it; the dhu stage needs a tensor, and a zero one is the
+                # honest gradient. Shared and read-only, like the zero initial state.
+                dht=dht.contiguous() if dht is not None else chain.zero_state(
+                    *h0.shape, h0.device
+                ),
+                scale=ctx.scale, chunk_size=ctx.chunk_size,
+            )
+            if q_rstd is not None:
+                from fla.modules.l2norm import l2norm_bwd
+
+                dq = l2norm_bwd(q, q_rstd, dq)
+                dk = l2norm_bwd(k, k_rstd, dk)
+
+            dA = dbias = None
+            if A_log is not None:
+                dg, dA, dbias = chain.gate_backward(
+                    g_org, A_log, dt_bias, dg, ctx.lower_bound
+                )
+            g_ref = g_org if g_org is not None else g2
+            support.mark_warm(ctx.warm_sig, "bwd")  # capture may use this shape's backward now
+            return (
+                dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype),
+                dg.to(g_ref.dtype), db.to(beta.dtype),
+                None if ctx.h0_was_none else dh0,
+                None, None, None, dA, dbias, None, None, None,
+            )
+
+    return ChunkKDAFunction
+
+
+_FN = None
+
+
+def apply(q, k, v, g, beta, h0, scale, chunk_size, use_qk_l2norm,
+          A_log, dt_bias, lower_bound, h0_was_none, warm_sig):
+    global _FN
+    if _FN is None:
+        _FN = _make_fn()
+    return _FN.apply(
+        q, k, v, g, beta, h0, scale, chunk_size, use_qk_l2norm,
+        A_log, dt_bias, lower_bound, h0_was_none, warm_sig,
+    )

kernel_fun/kda/chain.py

@@ -0,0 +1,150 @@
+"""The KDA chain: which kernel runs at each stage, and in what order.
+
+This is fla 0.5.2's own decomposition of `chunk_kda`, launch for launch, with five stages
+replaced. Keeping fla's structure is deliberate — it is what makes a stage-by-stage
+comparison meaningful, and it is why the parity tests can hold to fla's own tolerances.
+
+Forward:
+    gate+cumsum   fla   kda_gate_chunk_cumsum (fused activation, or the plain cumsum)
+    intra+solve   fla   token_parallel + inter_solve_fused, preceded by OUR zero-fill of
+                        Aqk's upper triangles (fla leaves them uninitialized and masks at
+                        load; our scan contracts the full tile). The zero-fill replaces a
+                        masked_fill that read the whole 268MB tile to write an eighth of it.
+    w/u prep      fla   recompute_w_u_fwd
+    scan + o      OURS  fused state scan and readout, state resident in registers
+
+Backward (fla's recompute path):
+    1 recompute   fla   recompute_w_u_fwd
+    2 rescan      OURS  B1: forward re-scan
+    3 dAv         fla   chunk_kda_bwd_dAv
+    4 dhu         OURS  B2a: reverse dh/dv scan, dh^T resident in registers
+    5 wy_dqkg     OURS  the transposed pair (ladder 005): M=K=128 MMAs over a loop that
+                        streams at the DRAM roofline, a K-strip epilogue at 128 regs so
+                        two CTAs share an SM, and the v-side/dA chain in a side kernel.
+                        Below its CTA floor it is B2b, the full-K restructure of fla's
+                        fused kernel (same math, one K slab instead of NK)
+    6 intra       OURS  tcgen05 off-diagonals + SIMT diagonals, with the dg reverse-cumsum
+                        and the bf16 dq/dk cast folded into its epilogue
+    7 dg_cumsum   —     identity: stage 6 already emitted it
+
+Every one of ours falls back to fla's kernel off its supported shape or below the CTA
+floor, so this table is always safe to call; `is_supported` exists so a caller can know
+whether that happened instead of reading a silent 1.00x.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .._common.support import log_once
+
+
+def gate_cumsum(g, A_log, dt_bias, chunk_size, lower_bound=None):
+    """The chunk-local cumsum, with the gate activation fused in when there is one.
+
+    With use_gate_in_kernel the op receives w_g(x) raw and owes
+    -exp(A_log)*softplus(g + dt_bias) before the cumsum. Doing that in eager torch is four
+    passes over an fp32 [B,T,HV,K] tensor — at prod8192 that is ~2ms and ~2GiB of saved
+    activations PER LAYER, against a measured 1.9ms for the entire fused stage. fla fuses
+    it into the cumsum this chain launches anyway, so the fused form is free.
+    """
+    from fla.ops.utils.constant import RCP_LN2
+
+    if A_log is None:
+        from fla.ops.utils import chunk_local_cumsum
+
+        return chunk_local_cumsum(g=g, scale=RCP_LN2, chunk_size=chunk_size)
+
+    from fla.ops.kda.gate import kda_gate_chunk_cumsum
+
+    return kda_gate_chunk_cumsum(
+        g=g, A_log=A_log, dt_bias=dt_bias, scale=RCP_LN2, chunk_size=chunk_size,
+        lower_bound=lower_bound,
+    )
+
+
+def forward(q, k, v, g2, beta, h0, scale, chunk_size):
+    """Returns (o, ht, Aqk, Akk) — o, the final state, and the backward's two residuals."""
+    from fla.ops.kda.wy_fast import recompute_w_u_fwd
+
+    from ._kernels import fwd_intra_triton, fwd_state
+
+    Aqk, Akk = fwd_intra_triton.chunk_kda_fwd_intra_zerofill(
+        q, k, g2, beta, float(scale), chunk_size
+    )
+    w, u, qg, kg = recompute_w_u_fwd(k=k, v=v, beta=beta, A=Akk, q=q, gk=g2)
+    # aqk_prezeroed: the zero-fill above already cleared the upper triangles, so the scan
+    # skips the masked_fill it would otherwise do.
+    o, ht = fwd_state.kda_cute_fwd_call(
+        qg, kg, w, u, Aqk, g2, h0, float(scale), aqk_prezeroed=True
+    )
+    return o, ht, Aqk, Akk
+
+
+def backward(q, k, v, g2, beta, Aqk, Akk, h0, do, dht, scale, chunk_size):
+    """Returns (dq, dk, dv, dbeta, dg, dh0). dg is w.r.t. the pre-cumsum decay."""
+    from fla.ops.kda.chunk_bwd import chunk_kda_bwd_dAv
+    from fla.ops.kda.wy_fast import recompute_w_u_fwd
+
+    from ._kernels import bwd_dhu, bwd_intra, bwd_scan, bwd_wy_t
+
+    H, HV = q.shape[2], v.shape[2]
+
+    w, u, qg, kg = recompute_w_u_fwd(q=q, k=k, v=v, beta=beta, A=Akk, gk=g2)
+    # dq0 is B1's raw do@h^T when the dq fusion is on. It is off: the fusion measured
+    # negative, so wy_dqkg computes dq itself and B1 returns None here.
+    h, v_new, _dq0 = bwd_scan.kda_rescan_b1(kg, w, u, g2, h0, do, chunk_size)
+    dAqk, dv = chunk_kda_bwd_dAv(
+        q=q, k=k, v=v_new, do=do, A=Aqk, scale=scale, chunk_size=chunk_size,
+    )
+    dh, dh0, dv = bwd_dhu.kda_dhu_b2a(qg, kg, w, g2, h0, dht, do, dv, scale, chunk_size)
+    dq, dk, dv, db, dg, dAkk = bwd_wy_t.chunk_kda_bwd_wy_dqkg_t(
+        q=q, k=k, v=v, v_new=v_new, g=g2, beta=beta, A=Akk, h=h,
+        do=do, dh=dh, dv=dv, scale=scale, chunk_size=chunk_size,
+    )
+    # fold_dg: intra emits dg already chunk-reverse-cumsum'd, so fla's dg_cumsum stage
+    # disappears. emit_bf16: dq/dk come out in q's dtype, making the wrapper's casts no-ops
+    # — but only at HV == H, since the GVA reduction below must sum in fp32.
+    dq, dk, db, dg = bwd_intra.chunk_kda_bwd_intra_cutedsl(
+        q=q, k=k, g=g2, beta=beta, dAqk=dAqk, dAkk=dAkk,
+        dq=dq, dk=dk, db=db, dg=dg, chunk_size=chunk_size,
+        fold_dg=True, emit_bf16=HV == H,
+    )
+    # The GVA reduction sits where fla puts it: after intra, before the (folded) cumsum.
+    if HV > H:
+        G = HV // H
+        dq = dq.view(*dq.shape[:2], H, G, dq.shape[-1]).sum(dim=3)
+        dk = dk.view(*dk.shape[:2], H, G, dk.shape[-1]).sum(dim=3)
+    return dq, dk, dv, db, dg, dh0
+
+
+def gate_backward(g_org, A_log, dt_bias, dg, lower_bound=None):
+    """Close the gate: dg w.r.t. the raw input, plus dA_log and ddt_bias.
+
+    Applied to the already-reverse-cumsum'd dg, exactly where fla applies it.
+    """
+    from fla.ops.kda.gate import kda_gate_bwd
+
+    return kda_gate_bwd(
+        g=g_org, A_log=A_log, dt_bias=dt_bias, dyg=dg, lower_bound=lower_bound
+    )
+
+
+_ZEROS: dict = {}
+
+
+def zero_state(B, HV, K, V, device) -> torch.Tensor:
+    """A shared read-only zero initial state.
+
+    The scan kernels read the initial state unconditionally — there is no null-h0 branch —
+    and production almost always passes initial_state=None. Nothing writes this buffer
+    (dh0 is a separate output), so one per (shape, device) is shared across every layer
+    rather than allocated per call: 134 MiB at prod8192, once.
+    """
+    key = (B, HV, K, V, device)
+    z = _ZEROS.get(key)
+    if z is None:
+        z = torch.zeros(B, HV, K, V, device=device, dtype=torch.float32)
+        _ZEROS[key] = z
+        log_once(f"kernel-fun kda: allocated a shared zero state {tuple(z.shape)}")
+    return z

kernel_fun/kda/ops.py

@@ -0,0 +1,342 @@
+"""`chunk_kda` — a drop-in for `fla.ops.kda.chunk_kda`.
+
+Same signature, same return contract, same validation. On a supported call it runs this
+package's kernels; on anything else it forwards the call to fla verbatim, which is why the
+fallback is bit-identical rather than merely close.
+
+The gate is a WHITELIST. Every argument is either known-supported, handled here, or forces
+the fallback — a new fla flag we have never seen degrades to fla instead of being silently
+dropped. `is_supported` returns the reason, because the worst failure mode for a kernel
+port is not being slow, it is running fla while everyone believes otherwise.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from .._common import support
+from . import autograd, chain
+
+__all__ = ["chunk_kda", "is_supported", "warmup"]
+
+log = logging.getLogger(__name__)
+
+# Arguments this package understands. Anything else present and non-default -> fla.
+_HANDLED = frozenset({
+    "q", "k", "v", "g", "beta", "scale", "initial_state", "output_final_state",
+    "use_qk_l2norm_in_kernel", "use_gate_in_kernel", "use_beta_sigmoid_in_kernel",
+    "allow_neg_eigval", "lower_bound", "A_log", "dt_bias", "chunk_size",
+})
+# Present in fla's signature, and each one forces the fallback (see is_supported).
+_UNSUPPORTED = frozenset({
+    "cu_seqlens", "cu_seqlens_cpu", "cp_context", "safe_gate", "state_v_first",
+    "disable_recompute", "return_intermediate_states", "transpose_state_layout",
+})
+
+
+def _warm_sig(q: torch.Tensor, v: torch.Tensor, initial_state, kwargs: dict) -> tuple:
+    """Everything that decides which kernels a call launches — shapes, dtypes, device and
+    the kernel-selecting flags. Not the stream: see `support.capture_unsupported_reason`."""
+    return (
+        "kda", tuple(q.shape), tuple(v.shape), q.dtype, v.dtype, q.device.index,
+        initial_state is not None,
+        bool(kwargs.get("use_qk_l2norm_in_kernel", False)),
+        bool(kwargs.get("use_gate_in_kernel", False)),
+        bool(kwargs.get("use_beta_sigmoid_in_kernel", False)),
+        bool(kwargs.get("allow_neg_eigval", False)),
+        kwargs.get("lower_bound") is not None,
+    )
+
+
+def is_supported(
+    q: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    chunk_size: int = 64,
+    initial_state: torch.Tensor | None = None,
+    **kwargs,
+) -> tuple[bool, str | None]:
+    """Can this call use our kernels? Returns (ok, reason-if-not).
+
+    The reason strings are meant to end up in a training log verbatim.
+    """
+    reason = support.common_unsupported_reason(q, "kda")
+    if reason is not None:
+        return False, reason
+
+    for name in _UNSUPPORTED:
+        val = kwargs.get(name)
+        if val not in (None, False):
+            return False, f"{name}={val!r} is not implemented here"
+    for name in kwargs:
+        if name not in _HANDLED and name not in _UNSUPPORTED:
+            return False, f"unrecognized argument {name!r} (fla may have grown a flag)"
+
+    B, T, H, K = q.shape
+    HV, V = v.shape[2], v.shape[-1]
+    if chunk_size != 64:
+        return False, f"chunk_size={chunk_size} (only 64 is implemented)"
+    if T % 64 != 0:
+        return False, f"T={T} is not a multiple of the chunk size"
+    if K not in (64, 128):
+        return False, f"K={K} (only 64 and 128 are implemented)"
+    if V % 64 != 0:
+        return False, f"V={V} is not a multiple of 64"
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        return False, f"dtype {q.dtype} (only bf16 and fp16)"
+    if not (q.dtype == v.dtype):
+        return False, f"mixed dtypes q={q.dtype} v={v.dtype}"
+    if initial_state is not None and initial_state.dtype != torch.float32:
+        return False, f"initial_state must be fp32, got {initial_state.dtype}"
+    # Below this the CuTe scans underfill the GPU and fla is genuinely faster. Not a
+    # correctness gate — the kernels would run — so it is worth stating in the log. Where
+    # the crossover actually falls is a property of the workload, which is why
+    # KERNEL_FUN_KDA_MIN_CTAS can move THIS gate for a shape somebody has timed. It moves
+    # nothing else: every check above is a capability, and every stage below keeps its own.
+    ctas = B * HV * (V // 64)
+    floor = support.min_ctas("kda")
+    if ctas < floor:
+        why = "fla is faster here" if floor == support.MIN_CTAS else "configured floor"
+        return False, f"grid too small ({ctas} CTAs < {floor}); {why}"
+    if floor != support.MIN_CTAS:
+        caveat = ""
+        if floor < support.MIN_CTAS:
+            caveat = (
+                f"; per-stage floors do not follow it, so the b1 scan and dhu backwards "
+                f"are fla's below {support.MIN_CTAS} CTAs"
+            )
+        support.log_once(
+            f"kernel-fun kda: dispatch floor is {floor} CTAs, not the default "
+            f"{support.MIN_CTAS} (KERNEL_FUN_KDA_MIN_CTAS){caveat}"
+        )
+    # Last, so it only speaks for a call that would otherwise be ours.
+    needs_grad = torch.is_grad_enabled() and (q.requires_grad or v.requires_grad)
+    reason = support.capture_unsupported_reason(
+        _warm_sig(q, v, initial_state, kwargs), needs_grad
+    )
+    if reason is not None:
+        return False, reason
+    return True, None
+
+
+@torch.compiler.disable
+def chunk_kda(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    use_gate_in_kernel: bool = False,
+    use_beta_sigmoid_in_kernel: bool = False,
+    allow_neg_eigval: bool = False,
+    safe_gate: bool = False,
+    lower_bound: float | None = None,
+    disable_recompute: bool = False,
+    return_intermediate_states: bool = False,
+    state_v_first: bool = False,
+    cu_seqlens: torch.LongTensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    cp_context=None,
+    **kwargs,
+):
+    """Drop-in for ``fla.ops.kda.chunk_kda``. See that function for the full argument docs.
+
+    ``@torch.compiler.disable``: the host path drives compiled CuTe objects through ctypes
+    pointer writes and a per-layout call cache, and reads
+    ``torch.cuda.current_stream().cuda_stream`` — none of which Dynamo can trace. This
+    MOVES a graph break rather than adding one: fla's own ``chunk_kda`` is wrapped in
+    ``@dispatch('kda')``, which applies ``torch.compiler.disable`` too. Keep every tensor
+    operation inside this function for that reason — a ``.float()`` in the caller would
+    split a compiled block in two and cost more than these kernels save.
+    """
+    from fla.ops.kda import chunk_kda as fla_chunk_kda
+
+    fla_kwargs = dict(
+        q=q, k=k, v=v, g=g, beta=beta, scale=scale, initial_state=initial_state,
+        output_final_state=output_final_state,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        use_gate_in_kernel=use_gate_in_kernel,
+        use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
+        allow_neg_eigval=allow_neg_eigval, safe_gate=safe_gate, lower_bound=lower_bound,
+        disable_recompute=disable_recompute,
+        return_intermediate_states=return_intermediate_states,
+        state_v_first=state_v_first, cu_seqlens=cu_seqlens,
+        cu_seqlens_cpu=cu_seqlens_cpu, cp_context=cp_context, **kwargs,
+    )
+
+    chunk_size = kwargs.get("chunk_size", 64)
+    ok, reason = is_supported(
+        q, v, chunk_size=chunk_size, initial_state=initial_state,
+        # The kernel-selecting flags are part of the capture gate's shape signature, so
+        # they must reach is_supported exactly as chunk_kda marks them warm below.
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        use_gate_in_kernel=use_gate_in_kernel,
+        use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
+        allow_neg_eigval=allow_neg_eigval, lower_bound=lower_bound,
+        safe_gate=safe_gate, state_v_first=state_v_first,
+        disable_recompute=disable_recompute,
+        return_intermediate_states=return_intermediate_states,
+        cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, cp_context=cp_context,
+        **{k_: v_ for k_, v_ in kwargs.items() if k_ != "chunk_size"},
+    )
+    if not ok:
+        support.log_once(
+            f"kernel-fun kda: falling back to fla — {reason}", logging.WARNING
+        )
+        return fla_chunk_kda(**fla_kwargs)
+
+    # Same validation fla does, so a malformed call fails identically on both paths.
+    B, T, H, K = q.shape
+    HV = v.shape[2]
+    assert q.shape == k.shape, f"q and k must match, got {q.shape} vs {k.shape}"
+    assert HV % H == 0, f"HV={HV} must be divisible by H={H}"
+    assert g.shape == (B, T, HV, K), f"g must be {[B, T, HV, K]}, got {list(g.shape)}"
+    assert beta.shape == (B, T, HV), f"beta must be {[B, T, HV]}, got {list(beta.shape)}"
+
+    A_log, dt_bias = kwargs.get("A_log"), kwargs.get("dt_bias")
+    if use_gate_in_kernel:
+        assert A_log is not None, "A_log is required when use_gate_in_kernel=True"
+    else:
+        A_log = dt_bias = None
+    if use_beta_sigmoid_in_kernel:
+        from fla.ops.common.gate import fused_beta_sigmoid
+
+        beta = fused_beta_sigmoid(beta, scale=2.0 if allow_neg_eigval else 1.0)
+    # The intra backward types beta as q's dtype; production computes it in fp32
+    # (w_b(x).float().sigmoid()*2). Casting here rather than widening the kernel — beta is
+    # bounded in (0, 2) and dbeta carries the loosest tolerance in the op.
+    if beta.dtype != q.dtype:
+        beta = beta.to(q.dtype)
+    if scale is None:
+        scale = K ** -0.5
+
+    h0, h0_was_none = initial_state, initial_state is None
+    if h0_was_none:
+        h0 = chain.zero_state(B, HV, K, v.shape[-1], q.device)
+
+    support.log_versions_once()
+    support.log_once(
+        f"kernel-fun kda: engaged (B={B} T={T} H={H} HV={HV} K={K} V={v.shape[-1]} "
+        f"chunk={chunk_size} l2norm={use_qk_l2norm_in_kernel} gate={use_gate_in_kernel})"
+    )
+    sig = _warm_sig(q, v, initial_state, dict(
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        use_gate_in_kernel=use_gate_in_kernel,
+        use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
+        allow_neg_eigval=allow_neg_eigval, lower_bound=lower_bound,
+    ))
+    o, ht = autograd.apply(
+        q, k, v, g, beta, h0, scale, chunk_size, use_qk_l2norm_in_kernel,
+        A_log, dt_bias, lower_bound, h0_was_none, sig,
+    )
+    support.mark_warm(sig, "fwd")  # the backward marks itself, from the Function
+    return o.type_as(q), (ht if output_final_state else None)
+
+
+def warmup(
+    *,
+    K: int = 128,
+    V: int = 256,
+    HV: int = 16,
+    dtype: torch.dtype = torch.bfloat16,
+    device: str | torch.device = "cuda",
+    use_qk_l2norm: bool = True,
+    use_gate_in_kernel: bool = True,
+) -> float:
+    """Compile every kernel this op needs, so training step 1 does not.
+
+    Four `cute.compile` calls plus fla's Triton autotuning is tens of seconds. Paid inside
+    step 1 it looks exactly like a regression — that is what the last port's reported
+    "-3% tokens/sec" turned out to be.
+
+    Compile keys carry no B, T or HV, so a tiny shape compiles the production kernels — but
+    it has to clear BOTH floors, the chain's (B*HV*(V//64) >= 256) and the intra kernel's
+    own (B*(T/64)*HV >= 1024). Miss the second and three of the four CuTe kernels compile
+    while the fourth quietly warms up its Triton fallback instead. Also runs the fla
+    compatibility probe, so a version mismatch raises here rather than at step 40,000.
+
+    KERNEL_FUN_KDA_MIN_CTAS earns a second pass. Raised, the first shape has to reach the
+    new floor or this call falls back and compiles nothing. LOWERED, it puts production on
+    a chain this shape never takes — the b1 scan and dhu keep their own 256 floors, so
+    below it they are fla's, and fla autotunes them at step 1 unless something warms them
+    here. Hence one pass per grid the process will actually dispatch at.
+
+    Returns the elapsed seconds — log it.
+    """
+    import time
+
+    from .._common.compat import check_fla
+    from ._kernels.bwd_intra import _MIN_CTAS as INTRA_MIN_CTAS
+
+    def one_pass(B: int) -> None:
+        kw = dict(device=device, dtype=dtype)
+        q = torch.randn(B, T, HV, K, **kw, requires_grad=True)
+        k = torch.randn(B, T, HV, K, **kw, requires_grad=True)
+        v = torch.randn(B, T, HV, V, **kw, requires_grad=True)
+        beta = torch.rand(B, T, HV, **kw, requires_grad=True)
+        A_log = dt_bias = None
+        if use_gate_in_kernel:
+            g = torch.randn(B, T, HV, K, **kw, requires_grad=True)
+            A_log = torch.rand(HV, device=device, dtype=torch.float32).add(1).log()
+            dt_bias = torch.zeros(HV * K, device=device, dtype=torch.float32)
+        else:
+            g = torch.nn.functional.logsigmoid(
+                torch.randn(B, T, HV, K, device=device, dtype=torch.float32)
+            ).requires_grad_(True)
+        o, ht = chunk_kda(
+            q, k, v, g, beta, initial_state=None, output_final_state=True,
+            use_qk_l2norm_in_kernel=use_qk_l2norm, use_gate_in_kernel=use_gate_in_kernel,
+            A_log=A_log, dt_bias=dt_bias,
+        )
+        (o.float().square().sum() + ht.float().square().sum()).backward()
+        torch.cuda.synchronize()
+
+    check_fla()
+    T = 1024
+    per_b = HV * max(V // 64, 1)
+    dispatch = support.min_ctas("kda")
+    # The compile pass: the four CuTe kernels, so it clears the default floor (two of them
+    # have their own 256) AND the intra kernel's, and a raised dispatch floor on top.
+    B = max(
+        -(-max(support.MIN_CTAS, dispatch) // per_b),
+        -(-INTRA_MIN_CTAS // (HV * (T // 64))),
+        1,
+    )
+    grids = [B]
+    # The autotune pass, only for a lowered floor: the fla stages production is about to
+    # run there. Nothing NEW of ours compiles at this grid — the compile keys are the same
+    # — so what it costs is one small fwd+bwd and what it buys is fla's autotune.
+    if dispatch < support.MIN_CTAS:
+        B_low = max(-(-dispatch // per_b), 1)
+        if B_low < B:
+            grids.append(B_low)
+
+    t0 = time.perf_counter()
+    torch.manual_seed(0)
+    for b in grids:
+        one_pass(b)
+
+    # A warmup that silently compiled nothing is worse than none: it hides the cost it was
+    # supposed to move, and the run pays it at step 1 anyway.
+    from ._kernels import bwd_dhu, bwd_intra, bwd_scan, fwd_state
+
+    empty = [
+        m.__name__.rsplit(".", 1)[-1]
+        for m in (fwd_state, bwd_scan, bwd_dhu, bwd_intra)
+        if not getattr(m, "_COMPILE_CACHE", {})
+    ]
+    if empty:
+        raise RuntimeError(
+            f"kernel-fun kda warmup compiled nothing for {empty} — the shape "
+            f"(B={B} T={T} HV={HV} K={K} V={V}) did not reach those kernels, so the cost "
+            f"this call exists to move is still waiting in step 1. Most likely a CTA floor: "
+            f"the chain needs B*HV*(V//64) >= {support.MIN_CTAS} and the intra kernel needs "
+            f"B*(T/64)*HV >= {INTRA_MIN_CTAS}."
+        )
+    return time.perf_counter() - t0