Add cudnn.fla: a cuDNN drop-in for flash-linear-attention (GDN + KDA)

python/cudnn/__init__.py

@@ -332,6 +332,13 @@ def __getattr__(name: str) -> Any:
         globals()["jax"] = _jax
         return _jax
 
+    if name == "fla":
+        # `import cudnn; cudnn.fla.accelerate_fla()` works like `import cudnn.fla`.
+        # Deferred so `import cudnn` never eagerly imports torch / the FLA shim.
+        _fla = importlib.import_module(".fla", __name__)
+        globals()["fla"] = _fla
+        return _fla
+
     if name in _LAZY_OPTIONAL_IMPORTS:
         return _load_optional_symbol(name)
 

python/cudnn/fla/__init__.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""cuDNN drop-in acceleration for flash-linear-attention (FLA).
+
+``accelerate_fla()`` monkeypatches the FLA ops cuDNN can serve so an existing
+``import fla`` training/inference script gets cuDNN's Blackwell kernels with no
+code change, and transparently falls back to FLA where cuDNN has no kernel — so
+results never change. Today: ``gated_delta_rule`` (Gated DeltaNet) and ``kda``
+(Kimi Delta Attention).
+
+    import cudnn.fla
+    cudnn.fla.accelerate_fla()   # call before importing FLA layers/models
+
+Correctness is gated by ``test/python/linear_attention/test_fla_compat.py``, which
+requires cuDNN to match FLA within its own bf16 noise on the output and every
+gradient; a config that does not match must fall back rather than run.
+"""
+
+from __future__ import annotations
+
+import sys
+
+from .gated_delta_rule import make_chunk_gated_delta_rule, last_path
+from .kda import make_chunk_kda
+
+__all__ = ["accelerate_fla", "is_accelerated", "last_path"]
+
+# The FLA ops cuDNN accelerates: (import path, attribute, shim factory).
+_ACCELERATED = [
+    ("fla.ops.gated_delta_rule", "chunk_gated_delta_rule", make_chunk_gated_delta_rule),
+    ("fla.ops.kda", "chunk_kda", make_chunk_kda),
+]
+
+_ORIGINALS: dict = {}
+
+
+def is_accelerated() -> bool:
+    return bool(_ORIGINALS)
+
+
+def _rebind_everywhere(fn_name: str, original, replacement) -> None:
+    """Rebind ``fn_name`` from ``original`` to ``replacement`` in every module that
+    captured it by reference (e.g. FLA layers that did ``from ... import fn``)."""
+    for mod in list(sys.modules.values()):
+        if mod is None:
+            continue
+        try:
+            if getattr(mod, fn_name, None) is original:
+                setattr(mod, fn_name, replacement)
+        except Exception:
+            # Some modules raise on getattr of arbitrary names; skip them.
+            continue
+
+
+def accelerate_fla(verbose: bool = True) -> None:
+    """Patch the FLA ops cuDNN accelerates. Idempotent; call before FLA models load."""
+    if is_accelerated():
+        return
+    import importlib
+
+    patched_names = []
+    for mod_path, attr, maker in _ACCELERATED:
+        try:
+            mod = importlib.import_module(mod_path)
+        except ImportError:
+            continue  # this op not present in the installed FLA
+        original = getattr(mod, attr, None)
+        if original is None:
+            continue
+        _ORIGINALS[attr] = (mod_path, original)
+        _rebind_everywhere(attr, original, maker(original))
+        patched_names.append(attr)
+
+    if not patched_names:
+        raise ImportError("accelerate_fla() requires flash-linear-attention installed")
+    if verbose:
+        print(f"[cudnn.fla] accelerated FLA {', '.join(patched_names)} with cuDNN (SM100); " "unsupported configs fall back to FLA.")
+
+
+def restore_fla() -> None:
+    """Undo :func:`accelerate_fla`."""
+    import importlib
+
+    for attr, (mod_path, original) in list(_ORIGINALS.items()):
+        mod = importlib.import_module(mod_path)
+        current = getattr(mod, attr, None)
+        if current is not None and current is not original:
+            _rebind_everywhere(attr, current, original)
+        setattr(mod, attr, original)  # restore the owning module's attribute explicitly
+    _ORIGINALS.clear()

python/cudnn/fla/gated_delta_rule.py

@@ -0,0 +1,202 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""cuDNN-accelerated drop-in for ``fla.ops.gated_delta_rule.chunk_gated_delta_rule``.
+
+Maps the flash-linear-attention public signature onto cuDNN's native
+``gated_delta_net`` (Blackwell/SM100) and falls back to the wrapped FLA function
+for anything cuDNN does not serve, so results never change and never regress.
+
+FLA layout is ``[B, T, H, ...]`` batch-first; ``g``/``beta`` are indexed by the
+*value* heads ``HV`` and FLA's grouped-value attention has ``HV >= H``, so native
+``HO = max(H, HV) = HV`` and the head mapping is a plain reshape + float cast.
+
+FLA's ``GatedDeltaNet`` layer drives the kernel with fused knobs; the native op
+now fuses each transform in-kernel (fwd+bwd), so the adapter forwards the raw
+inputs and the fusion flags rather than reproducing the math in torch:
+
+* ``use_gate_in_kernel`` -> ``safe_gate`` with ``a_log``/``dt_bias``; the kernel
+  applies ``g = -exp(a_log) * softplus(g + dt_bias)`` (native ``safe_gate`` matches
+  FLA's log decay exactly). FLA may omit ``dt_bias``; native requires it, so a
+  zero bias is synthesized.
+* ``use_beta_sigmoid_in_kernel`` -> forwarded; the kernel applies ``sigmoid(beta)``
+  (raw beta stays io-dtype). ``allow_neg_eigval`` would scale by 2 and is not served.
+* ``use_qk_l2norm_in_kernel`` -> forwarded; the kernel L2-normalizes q/k.
+"""
+
+from __future__ import annotations
+
+import torch
+
+import cudnn
+from cudnn.linear_attention.ops import gated_delta_net
+
+# Native declines a graph it cannot serve with one of these; treat as a fallback.
+_DECLINE = (cudnn.cudnnGraphNotSupportedError, NotImplementedError)
+
+# Diagnostic: which path the last shimmed call took ("native" | "fallback:<reason>").
+_LAST = {"path": None}
+
+
+def last_path() -> str | None:
+    """The route the most recent shimmed call took. For tests/telemetry only."""
+    return _LAST["path"]
+
+
+class _Decline(Exception):
+    """Raised internally when a call cannot be adapted to the native op."""
+
+
+def _to_native(
+    q,
+    k,
+    v,
+    g,
+    beta,
+    scale,
+    initial_state,
+    output_final_state,
+    cu_seqlens,
+    use_qk_l2norm_in_kernel,
+    use_beta_sigmoid_in_kernel,
+    use_gate_in_kernel,
+    A_log,
+    dt_bias,
+):
+    if q.dim() != 4:
+        raise _Decline("expected [B, T, H, K]")
+    B, T, H, _ = q.shape
+    HV = v.shape[2]
+    HO = max(H, HV)
+
+    if cu_seqlens is None:
+        cu = torch.arange(0, (B + 1) * T, T, dtype=torch.int32, device=q.device)
+    else:
+        if B != 1:
+            raise _Decline("varlen requires B==1 (FLA contract)")
+        cu = cu_seqlens.to(torch.int32)
+
+    # Fuse the gate in-kernel via safe_gate: pass raw g logits + a_log/dt_bias and
+    # let the kernel compute -exp(a_log)*softplus(g+dt_bias). safe_gate requires
+    # both; FLA may omit dt_bias, so synthesize a zero bias.
+    a_log_t = dt_bias_t = None
+    if use_gate_in_kernel:
+        if A_log is None:
+            raise _Decline("use_gate_in_kernel requires A_log")
+        a_log_t = A_log.float().reshape(-1)
+        dt_bias_t = dt_bias.float().reshape(-1) if dt_bias is not None else torch.zeros_like(a_log_t)
+        if a_log_t.shape[0] != HO or dt_bias_t.shape[0] != HO:
+            raise _Decline("a_log/dt_bias head count does not match HO=max(H,HV)")
+
+    # g is always fp32 (raw logits under safe_gate, else FLA's precomputed log decay).
+    # beta is io-dtype logits when the kernel applies the sigmoid, else fp32 post-activation.
+    g = g.float()
+    beta = beta.to(q.dtype) if use_beta_sigmoid_in_kernel else beta.float()
+
+    def thd(t):
+        return t.reshape(-1, *t.shape[2:])
+
+    g2, beta2 = thd(g), thd(beta)
+    if g2.shape[-1] != HO or beta2.shape[-1] != HO:
+        raise _Decline("g/beta head count does not match HO=max(H,HV)")
+
+    h0 = None if initial_state is None else initial_state.float().contiguous()
+    o, fs = gated_delta_net(
+        thd(q),
+        thd(k),
+        thd(v),
+        g2,
+        beta2,
+        cu,
+        scale=scale,
+        initial_state=h0,
+        output_final_state=output_final_state,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
+        safe_gate=use_gate_in_kernel,
+        a_log=a_log_t,
+        dt_bias=dt_bias_t,
+    )
+    o = o.reshape(B, T, *o.shape[1:])  # native o is shaped like v (THD) -> [B,T,HV,V]
+    return o, (fs if output_final_state else None)
+
+
+def make_chunk_gated_delta_rule(real_fn):
+    """Wrap FLA's ``chunk_gated_delta_rule`` with a cuDNN fast path + FLA fallback."""
+
+    def chunk_gated_delta_rule(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        scale=None,
+        initial_state=None,
+        output_final_state=False,
+        use_qk_l2norm_in_kernel=False,
+        use_beta_sigmoid_in_kernel=False,
+        allow_neg_eigval=False,
+        state_v_first=False,
+        cu_seqlens=None,
+        cu_seqlens_cpu=None,
+        cp_context=None,
+        **kwargs,
+    ):
+        A_log = kwargs.get("A_log")
+        dt_bias = kwargs.get("dt_bias")
+        use_gate_in_kernel = kwargs.get("use_gate_in_kernel", False)
+
+        def fallback(reason):
+            _LAST["path"] = f"fallback:{reason}"
+            return real_fn(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                scale=scale,
+                initial_state=initial_state,
+                output_final_state=output_final_state,
+                use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+                use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
+                allow_neg_eigval=allow_neg_eigval,
+                state_v_first=state_v_first,
+                cu_seqlens=cu_seqlens,
+                cu_seqlens_cpu=cu_seqlens_cpu,
+                cp_context=cp_context,
+                **kwargs,
+            )
+
+        # Variants the native op does not model -> incumbent.
+        if allow_neg_eigval or cp_context is not None:
+            return fallback("variant")
+        # state_v_first only changes the recurrent-state layout, so it is a no-op
+        # for a stateless (training) call; decline only when a state is exchanged.
+        if state_v_first and (initial_state is not None or output_final_state):
+            return fallback("state_v_first")
+        if not (q.is_cuda and torch.cuda.get_device_capability(q.device)[0] >= 10):
+            return fallback("pre-Blackwell")
+        try:
+            out = _to_native(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                scale,
+                initial_state,
+                output_final_state,
+                cu_seqlens,
+                use_qk_l2norm_in_kernel,
+                use_beta_sigmoid_in_kernel,
+                use_gate_in_kernel,
+                A_log,
+                dt_bias,
+            )
+        except (_Decline, *_DECLINE) as e:
+            return fallback(type(e).__name__)
+        _LAST["path"] = "native"
+        return out
+
+    chunk_gated_delta_rule.__wrapped__ = real_fn
+    return chunk_gated_delta_rule