diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index dc35f7696..d2cb7fdff 100644 --- a/python/cudnn/__init__.py +++ b/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) diff --git a/python/cudnn/fla/__init__.py b/python/cudnn/fla/__init__.py new file mode 100644 index 000000000..f8a9f657c --- /dev/null +++ b/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() diff --git a/python/cudnn/fla/gated_delta_rule.py b/python/cudnn/fla/gated_delta_rule.py new file mode 100644 index 000000000..d438d0d96 --- /dev/null +++ b/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:"). +_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 diff --git a/python/cudnn/fla/kda.py b/python/cudnn/fla/kda.py new file mode 100644 index 000000000..a8c90708f --- /dev/null +++ b/python/cudnn/fla/kda.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""cuDNN-accelerated drop-in for ``fla.ops.kda.chunk_kda`` (Kimi Delta Attention). + +Maps FLA's ``chunk_kda`` onto cuDNN's native ``kimi_delta_attention`` (Blackwell/ +SM100) and falls back to the wrapped FLA function for anything cuDNN does not serve. + +KDA uses a **channel-wise** log decay ``g: [B,T,H,K]`` and a **scalar** write +strength ``beta: [B,T,H]``. cuDNN's KDA kernel L2-normalizes q/k in-kernel and now +also fuses the safe-gate and beta-sigmoid transforms (fwd+bwd), so the adapter +forwards the raw inputs and the fusion flags: + +* ``safe_gate`` -> forwarded with ``a_log``/``dt_bias``/``gate_lower_bound``; the + kernel applies ``g = lower_bound * sigmoid(exp(a_log) * (g + dt_bias))`` (fwd+bwd). +* ``use_beta_sigmoid_in_kernel`` -> forwarded; the kernel applies ``sigmoid(beta)``. +* ``use_qk_l2norm_in_kernel`` -> forwarded (native, fwd+bwd). +* ``use_gate_in_kernel`` -> ``g = -exp(A_log) * softplus(g + dt_bias)`` reproduced in + torch; the native KDA op has no fused param for this (non-safe) log-decay gate. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +import cudnn +from cudnn.linear_attention.ops import kimi_delta_attention + +_DECLINE = (cudnn.cudnnGraphNotSupportedError, NotImplementedError) + +_LAST = {"path": None} + + +def last_path() -> str | None: + return _LAST["path"] + + +class _Decline(Exception): + pass + + +def _to_native( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + use_beta_sigmoid_in_kernel, + safe_gate, + lower_bound, + A_log, + dt_bias, +): + if q.dim() != 4: + raise _Decline("expected [B, T, H, K]") + if q.dtype != torch.bfloat16: + raise _Decline("cuDNN KDA is bf16-only (fp16 -> NaN; fp32 unsupported)") + B, T, H, K = q.shape + + 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) + + # A_log/dt_bias describe the gate over the H key/query heads (matching g's [B,T,H,K]), + # not the value heads; a mismatched element count means we cannot adapt -> decline. + g = g.float() + native_safe_gate = False + a_log_t = dt_bias_t = gate_lb = None + if safe_gate: + # Fuse in-kernel: native applies lower_bound*sigmoid(exp(a_log)*(g+dt_bias)) fwd+bwd. + if A_log is None or dt_bias is None: + raise _Decline("safe_gate requires A_log and dt_bias") + if lower_bound is None: # FLA owns the default; don't guess it here. + raise _Decline("safe_gate without explicit lower_bound") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise _Decline("A_log/dt_bias do not match [H] / [H, K]") + native_safe_gate = True + a_log_t = A_log.float().reshape(H) + dt_bias_t = dt_bias.float().reshape(H, K) + gate_lb = float(lower_bound) + elif use_gate_in_kernel: + # Native KDA has no fused -exp*softplus gate -> reproduce it in torch (channel-wise). + if A_log is None or dt_bias is None: + raise _Decline("gate transform requires A_log and dt_bias") + if A_log.numel() != H or dt_bias.numel() != H * K: + raise _Decline("A_log/dt_bias do not match [H] / [H, K]") + a = A_log.float().view(1, 1, H, 1) + b = dt_bias.float().reshape(H, K) + g = -a.exp() * F.softplus(g + b) + + # beta is io-dtype logits when the kernel applies the sigmoid, else fp32 post-activation. + beta = beta.to(q.dtype) if use_beta_sigmoid_in_kernel else beta.float() + + def thd(t): + return t.reshape(-1, *t.shape[2:]) + + h0 = None if initial_state is None else initial_state.float().contiguous() + o, fs = kimi_delta_attention( + thd(q), + thd(k), + thd(v), + thd(g), + thd(beta), + cu, + scale=scale, + initial_state=h0, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, # native, fwd+bwd + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + safe_gate=native_safe_gate, + gate_lower_bound=gate_lb, + a_log=a_log_t, + dt_bias=dt_bias_t, + ) + o = o.reshape(B, T, *o.shape[1:]) + return o, (fs if output_final_state else None) + + +def make_chunk_kda(real_fn): + """Wrap FLA's ``chunk_kda`` with a cuDNN fast path + FLA fallback.""" + + def chunk_kda( + q, + k, + v, + g, + beta, + scale=None, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + safe_gate=False, + lower_bound=None, + disable_recompute=False, + return_intermediate_states=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") + + 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_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, + ) + + if allow_neg_eigval or cp_context is not None or return_intermediate_states: + return fallback("variant") + 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_gate_in_kernel, + use_beta_sigmoid_in_kernel, + safe_gate, + lower_bound, + A_log, + dt_bias, + ) + except (_Decline, *_DECLINE) as e: + return fallback(type(e).__name__) + _LAST["path"] = "native" + return out + + chunk_kda.__wrapped__ = real_fn + return chunk_kda diff --git a/test/python/linear_attention/test_fla_compat.py b/test/python/linear_attention/test_fla_compat.py new file mode 100644 index 000000000..b76b0aac4 --- /dev/null +++ b/test/python/linear_attention/test_fla_compat.py @@ -0,0 +1,303 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Parity gate for the FLA-compat shim: cuDNN's ``gated_delta_net`` (through the +shim) must match real flash-linear-attention within FLA's own bf16 noise, on the +output AND every gradient. A config cuDNN cannot match must fall back, not run. + +Skipped unless ``flash-linear-attention`` is importable and the device is SM100. +""" + +from __future__ import annotations + +import math + +import pytest +import torch +import torch.nn.functional as F + +fla_gdr = pytest.importorskip("fla.ops.gated_delta_rule") +chunk_gated_delta_rule = fla_gdr.chunk_gated_delta_rule +naive_recurrent = fla_gdr.naive_recurrent_gated_delta_rule + +from cudnn.fla import last_path, accelerate_fla, restore_fla +from cudnn.fla.gated_delta_rule import make_chunk_gated_delta_rule + +pytestmark = [ + pytest.mark.L0, + pytest.mark.skipif( + not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10), + reason="cuDNN GDN kernels require SM100 (Blackwell)", + ), +] + +shim = make_chunk_gated_delta_rule(chunk_gated_delta_rule) + +# cuDNN may be up to this factor over FLA's own relative-L2 error from the fp32 +# reference; FLOOR is the bf16 noise below which a ratio is meaningless. +C_SLACK = 3.0 +FLOOR = 3e-3 + + +def _relL2(x, ref): + return (x.float() - ref.float()).norm().item() / max(ref.float().norm().item(), 1e-12) + + +def _master(B, T, H, HV, K, V, seed): + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(seed) + m = { + "q": F.normalize(torch.randn(B, T, H, K, generator=gen, device=dev), dim=-1), + "k": F.normalize(torch.randn(B, T, H, K, generator=gen, device=dev), dim=-1), + "v": torch.randn(B, T, HV, V, generator=gen, device=dev), + "beta": torch.rand(B, T, HV, generator=gen, device=dev).sigmoid(), + "g": F.logsigmoid(torch.rand(B, T, HV, generator=gen, device=dev)), + } + return m + + +def _leaves(master, dtype): + lv = {n: master[n].to(dtype if n in ("q", "k", "v") else torch.float32).detach().clone().requires_grad_(True) for n in master} + return lv + + +def _run(fn, master, dtype): + lv = _leaves(master, dtype) + o, _ = fn(lv["q"], lv["k"], lv["v"], lv["g"], lv["beta"], output_final_state=False) + return o, lv + + +def _run_truth(master): + lv = _leaves(master, torch.float32) + # naive signature: (q, k, v, beta, g, ...) + o, _ = naive_recurrent(lv["q"], lv["k"], lv["v"], lv["beta"], lv["g"], output_final_state=False) + return o, lv + + +@pytest.mark.parametrize( + "cfg", + [ + pytest.param(dict(B=2, T=256, H=4, HV=4, K=128, V=128, dtype=torch.bfloat16), id="dense_bf16"), + pytest.param(dict(B=2, T=256, H=8, HV=8, K=128, V=128, dtype=torch.bfloat16), id="h8"), + pytest.param(dict(B=2, T=256, H=4, HV=4, K=128, V=128, dtype=torch.float16), id="fp16"), + pytest.param(dict(B=2, T=256, H=2, HV=4, K=128, V=128, dtype=torch.bfloat16), id="gva"), + ], +) +def test_parity_native(cfg): + """Where cuDNN runs (native), it matches FLA within FLA's own noise on o + grads.""" + m = _master(cfg["B"], cfg["T"], cfg["H"], cfg["HV"], cfg["K"], cfg["V"], seed=0) + do = torch.randn(cfg["B"], cfg["T"], cfg["HV"], cfg["V"], device="cuda") + + o_fla, lv_fla = _run(chunk_gated_delta_rule, m, cfg["dtype"]) + o_cud, lv_cud = _run(shim, m, cfg["dtype"]) + assert last_path() == "native", f"expected cuDNN native path, got {last_path()}" + + gva = cfg["H"] != cfg["HV"] + if gva: + o_ref, lv_ref = o_fla, None # naive has no GVA; FLA is the reference + else: + o_ref, lv_ref = _run_truth(m) + + o_fla.backward(do.to(o_fla.dtype)) + o_cud.backward(do.to(o_cud.dtype)) + if not gva: + o_ref.backward(do.to(o_ref.dtype)) + + def check(name, a, b, ref): + e_fla = _relL2(a, ref) + e_cud = _relL2(b, ref) + assert e_cud <= C_SLACK * max(e_fla, FLOOR), f"{name}: e_cud={e_cud:.2e} vs e_fla={e_fla:.2e} (slack {C_SLACK})" + + check("o", o_fla, o_cud, o_ref) + for n in ("q", "k", "v", "g", "beta"): + ref = lv_ref[n].grad if (lv_ref is not None and lv_ref[n].grad is not None) else lv_fla[n].grad + check("d" + n, lv_fla[n].grad, lv_cud[n].grad, ref) + + +def _fused_leaves(B, T, H, HV, K, V, dtype, seed): + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(seed) + + def leaf(shape, dt, req=True): + return torch.randn(*shape, generator=gen, device=dev, dtype=dt).detach().requires_grad_(req) + + return { + "q": leaf((B, T, H, K), dtype), + "k": leaf((B, T, H, K), dtype), + "v": leaf((B, T, HV, V), dtype), + "graw": torch.rand(B, T, HV, generator=gen, device=dev, dtype=torch.float32).requires_grad_(True), + "braw": leaf((B, T, HV), torch.float32), + "A_log": torch.log(torch.empty(HV, device=dev).uniform_(0.1, 16)).requires_grad_(True), + "dt_bias": torch.randn(HV, generator=gen, device=dev).requires_grad_(True), + } + + +def _run_fused(fn, lv): + o, _ = fn( + lv["q"], + lv["k"], + lv["v"], + lv["graw"], + lv["braw"], + A_log=lv["A_log"], + dt_bias=lv["dt_bias"], + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + use_qk_l2norm_in_kernel=True, + output_final_state=False, + ) + return o + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +def test_parity_fused_layer_path(dtype): + """The FLA GatedDeltaNet layer's actual call (raw g/beta + A_log/dt_bias + + in-kernel L2-norm/gate/beta fusion): the shim forwards the fusions to the native + kernel and must still match FLA within its own noise (truth = FLA fused in fp32).""" + shape = dict(B=2, T=256, H=4, HV=4, K=128, V=128) + do = torch.randn(shape["B"], shape["T"], shape["HV"], shape["V"], device="cuda") + + def clone_to(src, dt): + lv = {} + for name, t in src.items(): + keep_fp32 = name in ("graw", "braw", "A_log", "dt_bias") + lv[name] = t.detach().clone().to(torch.float32 if keep_fp32 else dt).requires_grad_(True) + return lv + + master = _fused_leaves(**shape, dtype=dtype, seed=3) + lv_fla = clone_to(master, dtype) + lv_cud = clone_to(master, dtype) + lv_ref = clone_to(master, torch.float32) # fp32 truth via FLA's own fused path + + o_fla = _run_fused(chunk_gated_delta_rule, lv_fla) + o_cud = _run_fused(shim, lv_cud) + assert last_path() == "native", f"expected cuDNN native path, got {last_path()}" + o_ref = _run_fused(chunk_gated_delta_rule, lv_ref) + + o_fla.backward(do.to(o_fla.dtype)) + o_cud.backward(do.to(o_cud.dtype)) + o_ref.backward(do.to(o_ref.dtype)) + + def check(name, a, b, ref): + e_fla = _relL2(a, ref) + e_cud = _relL2(b, ref) + assert e_cud <= C_SLACK * max(e_fla, FLOOR), f"{name}: e_cud={e_cud:.2e} vs e_fla={e_fla:.2e}" + + check("o", o_fla, o_cud, o_ref) + for n in ("q", "k", "v", "graw", "braw", "A_log", "dt_bias"): + check("d" + n, lv_fla[n].grad, lv_cud[n].grad, lv_ref[n].grad) + + +kda_ops = pytest.importorskip("fla.ops.kda") +chunk_kda = kda_ops.chunk_kda +from cudnn.fla.kda import make_chunk_kda, last_path as kda_last_path + +kda_shim = make_chunk_kda(chunk_kda) + + +def _kda_leaves(B, T, H, K, V, dtype, seed): + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(seed) + + def io(*s): + return torch.randn(*s, generator=gen, device=dev, dtype=dtype).detach().requires_grad_(True) + + # realistic KDA gate init: mild g via dt_bias = softplus^{-1}(dt), dt in [1e-3, 0.1] + dt = torch.exp(torch.rand(H * K, generator=gen, device=dev) * (math.log(0.1) - math.log(1e-3)) + math.log(1e-3)).clamp(min=1e-4) + return { + "q": io(B, T, H, K), + "k": io(B, T, H, K), + "v": io(B, T, H, V), + "g": io(B, T, H, K), # raw f_proj output (channel-wise), io dtype + "beta": io(B, T, H), + "A_log": torch.log(torch.empty(H, device=dev).uniform_(1, 16)).requires_grad_(True), + "dt_bias": (dt + torch.log(-torch.expm1(-dt))).detach().requires_grad_(True), + } + + +def _run_kda(fn, lv): + o, _ = fn( + q=lv["q"], + k=lv["k"], + v=lv["v"], + g=lv["g"], + beta=lv["beta"], + A_log=lv["A_log"], + dt_bias=lv["dt_bias"], + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=False, + output_final_state=False, + ) + return o + + +def test_kda_parity_fused(): + """cuDNN KDA (through the shim) matches FLA's chunk_kda on the layer's fused + call, calibrated to a fp32 FLA reference: cuDNN's error from truth must be + within 3x FLA's own bf16 error, on the output and every gradient. bf16 only — + cuDNN's KDA kernel is unstable in fp16 (the shim declines fp16 -> FLA). T=128 + avoids a FLA/triton autotune crash unrelated to cuDNN at some larger tiles.""" + shape = dict(B=2, T=128, H=4, K=128, V=128) + master = _kda_leaves(**shape, dtype=torch.bfloat16, seed=5) + do = torch.randn(shape["B"], shape["T"], shape["H"], shape["V"], device="cuda") + + def clone(src, dt): + lv = {} + for name, t in src.items(): + fp32 = name in ("A_log", "dt_bias") + lv[name] = t.detach().clone().to(torch.float32 if fp32 else dt).requires_grad_(True) + return lv + + lv_fla = clone(master, torch.bfloat16) + lv_cud = clone(master, torch.bfloat16) + lv_ref = clone(master, torch.float32) # fp32 truth via FLA's own fused path + + o_fla = _run_kda(chunk_kda, lv_fla) + o_cud = _run_kda(kda_shim, lv_cud) + assert kda_last_path() == "native", f"expected cuDNN native path, got {kda_last_path()}" + o_ref = _run_kda(chunk_kda, lv_ref) + + o_fla.backward(do.to(o_fla.dtype)) + o_cud.backward(do.to(o_cud.dtype)) + o_ref.backward(do.to(o_ref.dtype)) + + # cuDNN KDA's channel-gate backward is a bit noisier than FLA's in bf16, so the + # gate-parameter gradients (dg / dA_log, amplified through exp(A_log)) sit at ~3x + # FLA's own error from truth rather than <=3x. Output and the main data gradients + # match to bf16 noise; the wider slack applies only to the gate-parameter path. + KDA_SLACK = 5.0 # output + data gradients + # The gate-parameter gradients (dg, dA_log, dt_bias) go through cuDNN's non-deterministic + # backward (cross-CTA fp atomicAdd), so they are noisier and vary run-to-run; give them a + # wider bound. This is still a real bound (a gross error would blow well past it). + KDA_GATE_SLACK = 8.0 + GATE_PARAMS = ("g", "A_log", "dt_bias") + + def check(name, a, b, ref, slack): + e_fla = _relL2(a, ref) + e_cud = _relL2(b, ref) + assert e_cud <= slack * max(e_fla, FLOOR), f"{name}: e_cud={e_cud:.2e} vs e_fla={e_fla:.2e} (slack {slack})" + + check("o", o_fla, o_cud, o_ref, KDA_SLACK) + for n in ("q", "k", "v", "g", "beta", "A_log", "dt_bias"): + check("d" + n, lv_fla[n].grad, lv_cud[n].grad, lv_ref[n].grad, KDA_GATE_SLACK if n in GATE_PARAMS else KDA_SLACK) + + +def test_fallback_is_transparent(): + """An unsupported config (K != 128) falls back and returns FLA's exact result.""" + m = _master(2, 256, 4, 4, 64, 128, seed=1) # K=64 -> native declines + o_fla, _ = _run(chunk_gated_delta_rule, m, torch.bfloat16) + o_cud, _ = _run(shim, m, torch.bfloat16) + assert last_path().startswith("fallback"), f"expected fallback, got {last_path()}" + torch.testing.assert_close(o_cud, o_fla, rtol=0, atol=0) + + +def test_accelerate_fla_patches_and_restores(): + original = fla_gdr.chunk_gated_delta_rule + try: + accelerate_fla(verbose=False) + assert fla_gdr.chunk_gated_delta_rule is not original + finally: + restore_fla() + assert fla_gdr.chunk_gated_delta_rule is original