From 7b55be079e9e31cd908c6f7562f3dde49e0280c7 Mon Sep 17 00:00:00 2001 From: drisspg Date: Sun, 16 Aug 2026 21:31:58 +0000 Subject: [PATCH] Add fused recurrent_kda for decode and inference prefill This is the missing public entrypoint that torchtitan's KDA layer (pytorch/torchtitan#4164) expects for its "recurrent" backend. `recurrent_kda(q, k, v, gate, beta, initial_state=None, *, cu_seqlens=None, output_final_state=False)` mirrors `chunk_kda`'s signature and packed contract: device-resident int32 offsets, repeated offsets as empty padding slots whose state passes through bitwise, a terminal offset below physical capacity leaving the tail outside the contract, and FP32 recurrent states with one leading entry per logical sequence. The distinction from the chunked core is encoded in the argument name: `gate` is the per-token log2 decay (`bounded_gate_cumsum` with `chunk_size=1`), not the chunk-local cumulative gate. Queries scale by 1/sqrt(K) internally, matching `chunk_kda` and the `naive_recurrent_kda` default. The Triton kernel scans tokens sequentially per (sequence, head, value-block) program with the [K, BV] FP32 state held in registers, computing in FP32 regardless of input dtype. It is inference-only: gradient-requiring inputs are rejected with a pointer at `chunk_kda` rather than silently detaching. The op uses the repo's define/impl + fake registration with a fixed-arity (output, state) schema, and joins the lazy `KDA_OPS` surface in `attn_gym.linear`. ```bash gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_recurrent.py -q ``` stack-info: PR: https://github.com/meta-pytorch/attention-gym/pull/317, branch: drisspg/stack/79 --- attn_gym/linear/__init__.py | 2 + attn_gym/linear/kda/__init__.py | 2 + attn_gym/linear/kda/fwd/triton/recurrent.py | 301 ++++++++++++++++++ attn_gym/linear/kda/naive.py | 35 ++- attn_gym/linear/kda/validation.py | 75 +++++ test/test_kda_recurrent.py | 324 ++++++++++++++++++++ 6 files changed, 727 insertions(+), 12 deletions(-) create mode 100644 attn_gym/linear/kda/fwd/triton/recurrent.py create mode 100644 attn_gym/linear/kda/validation.py create mode 100644 test/test_kda_recurrent.py diff --git a/attn_gym/linear/__init__.py b/attn_gym/linear/__init__.py index 0e7186aa..810fb40c 100644 --- a/attn_gym/linear/__init__.py +++ b/attn_gym/linear/__init__.py @@ -19,6 +19,7 @@ naive_chunk_kda, naive_chunk_kda_from_cumulative, naive_recurrent_kda, + recurrent_kda, ) # Note: Lazy Imports @@ -36,6 +37,7 @@ "naive_chunk_kda", "naive_chunk_kda_from_cumulative", "naive_recurrent_kda", + "recurrent_kda", ] GDN_OPS = [ diff --git a/attn_gym/linear/kda/__init__.py b/attn_gym/linear/kda/__init__.py index e691ab8d..e893c669 100644 --- a/attn_gym/linear/kda/__init__.py +++ b/attn_gym/linear/kda/__init__.py @@ -10,6 +10,7 @@ from attn_gym.linear.kda.fwd.triton.gate_fwd import bounded_gate_cumsum from attn_gym.linear.kda.fwd.triton.l2norm_fwd import l2norm +from attn_gym.linear.kda.fwd.triton.recurrent import recurrent_kda from attn_gym.linear.kda.masking import ( active_token_mask, mask_inactive_token_gradients, @@ -52,6 +53,7 @@ def __getattr__(name: str): "naive_chunk_kda", "naive_chunk_kda_from_cumulative", "naive_recurrent_kda", + "recurrent_kda", *_CUTEDSL_EXPORTS, ] ) diff --git a/attn_gym/linear/kda/fwd/triton/recurrent.py b/attn_gym/linear/kda/fwd/triton/recurrent.py new file mode 100644 index 00000000..44741227 --- /dev/null +++ b/attn_gym/linear/kda/fwd/triton/recurrent.py @@ -0,0 +1,301 @@ +# Copyright (c) 2025 Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Fused O(T) KDA recurrence for decode and inference prefill. + +The kernel scans tokens sequentially per (sequence, head, value block), holding the +FP32 recurrent state in registers. It mirrors :func:`naive_recurrent_kda` exactly: +per step the state decays by ``exp2(gate)`` per key channel, a beta-scaled delta +writes the new value, and the query reads the updated state. The operation is +inference-only; use ``chunk_kda`` for training. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from attn_gym.linear.kda.validation import validate_kda_inputs + +_MAX_KEY_DIM = 256 + + +@triton.jit +def kda_recurrent_fwd_kernel( + q, + k, + v, + gate, + beta, + output, + h0, + ht, + cu_seqlens, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v = pid % NV + i_nh = pid // NV + i_n, i_h = i_nh // H, i_nh % H + + if IS_VARLEN: + # Assumption: seqlens have been validated prior to call + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + else: + bos = i_n * T + eos = bos + T + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < K + m_v = o_v < V + m_kv = m_k[:, None] & m_v[None, :] + + p_state = i_n * H * K * V + i_h * K * V + o_k[:, None] * V + o_v[None, :] + if USE_INITIAL_STATE: + b_state = tl.load(h0 + p_state, mask=m_kv, other=0.0).to(tl.float32) + else: + b_state = tl.zeros([BK, BV], dtype=tl.float32) + + for t in range(bos, eos): + row = t * H + i_h + b_q = tl.load(q + row * K + o_k, mask=m_k, other=0.0).to(tl.float32) * scale + b_k = tl.load(k + row * K + o_k, mask=m_k, other=0.0).to(tl.float32) + b_g = tl.load(gate + row * K + o_k, mask=m_k, other=0.0).to(tl.float32) + b_beta = tl.load(beta + row).to(tl.float32) + b_v = tl.load(v + row * V + o_v, mask=m_v, other=0.0).to(tl.float32) + + b_state *= tl.exp2(b_g)[:, None] + b_delta = (b_v - tl.sum(b_k[:, None] * b_state, 0)) * b_beta + b_state += b_k[:, None] * b_delta[None, :] + b_o = tl.sum(b_q[:, None] * b_state, 0) + tl.store(output + row * V + o_v, b_o.to(output.dtype.element_ty), mask=m_v) + + if STORE_FINAL_STATE: + tl.store(ht + p_state, b_state, mask=m_kv) + + +def _launch_recurrent_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + store_final_state: bool, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Allocate outputs and launch the sequential scan over token spans.""" + batch, tokens, heads, key_dim = q.shape + value_dim = v.shape[-1] + num_sequences = batch if cu_seqlens is None else cu_seqlens.shape[0] - 1 + output = torch.empty_like(v, dtype=q.dtype) + final_state = ( + q.new_empty(num_sequences, heads, key_dim, value_dim, dtype=torch.float32) + if store_final_state + else None + ) + # BV=32 measured faster than 64 on B200 for both decode and prefill + # (smaller state tiles schedule better; state traffic is identical). + block_v = min(triton.next_power_of_2(value_dim), 32) + # One flat launch dimension: sequence-head counts can exceed the 65,535 + # grid-Y limit, while grid-X is effectively unbounded. + grid = (triton.cdiv(value_dim, block_v) * num_sequences * heads,) + kda_recurrent_fwd_kernel[grid]( + q, + k, + v, + gate, + beta, + output, + initial_state, + final_state, + cu_seqlens, + scale=key_dim**-0.5, + T=tokens, + H=heads, + K=key_dim, + V=value_dim, + BK=triton.next_power_of_2(key_dim), + BV=block_v, + USE_INITIAL_STATE=initial_state is not None, + STORE_FINAL_STATE=store_final_state, + IS_VARLEN=cu_seqlens is not None, + num_warps=4, + ) + return output, final_state + + +_RECURRENT_FWD_ARGS = ( + "(Tensor q, Tensor k, Tensor v, Tensor gate, Tensor beta," + " Tensor? initial_state, Tensor? cu_seqlens)" +) +torch.library.define("attn_gym::kda_recurrent_fwd", _RECURRENT_FWD_ARGS + " -> (Tensor, Tensor)") +torch.library.define("attn_gym::kda_recurrent_fwd_no_state", _RECURRENT_FWD_ARGS + " -> Tensor") + + +def _kda_recurrent_fwd_cuda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + output, final_state = _launch_recurrent_fwd( + q, k, v, gate, beta, initial_state, cu_seqlens, store_final_state=True + ) + assert final_state is not None + return output, final_state + + +def _kda_recurrent_fwd_no_state_cuda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> torch.Tensor: + return _launch_recurrent_fwd( + q, k, v, gate, beta, initial_state, cu_seqlens, store_final_state=False + )[0] + + +torch.library.impl("attn_gym::kda_recurrent_fwd", "CUDA", _kda_recurrent_fwd_cuda) +torch.library.impl( + "attn_gym::kda_recurrent_fwd_no_state", "CUDA", _kda_recurrent_fwd_no_state_cuda +) + + +@torch.library.register_fake("attn_gym::kda_recurrent_fwd") +def _kda_recurrent_fwd_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + del gate, beta, initial_state + num_sequences = q.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1 + final_state = q.new_empty( + num_sequences, q.shape[2], q.shape[3], v.shape[-1], dtype=torch.float32 + ) + return torch.empty_like(v, dtype=q.dtype), final_state + + +@torch.library.register_fake("attn_gym::kda_recurrent_fwd_no_state") +def _kda_recurrent_fwd_no_state_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> torch.Tensor: + del gate, beta, initial_state, cu_seqlens + return torch.empty_like(v, dtype=q.dtype) + + +_recurrent_fwd_op = torch.ops.attn_gym.kda_recurrent_fwd.default +_recurrent_fwd_no_state_op = torch.ops.attn_gym.kda_recurrent_fwd_no_state.default + + +def _validate_recurrent_kda_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> None: + """Validate the shared contract plus the fused scan's own constraints.""" + validate_kda_inputs( + q, k, v, gate, beta, initial_state, cu_seqlens, op_name="recurrent_kda", gate_name="gate" + ) + if q.shape[-1] > _MAX_KEY_DIM: + raise ValueError(f"recurrent_kda requires K in [1, {_MAX_KEY_DIM}], got {q.shape[-1]}") + if not q.is_cuda: + raise ValueError("recurrent_kda requires CUDA tensors") + data_tensors = (q, k, v, gate, beta) + if initial_state is not None: + data_tensors += (initial_state,) + if torch.is_grad_enabled() and any(tensor.requires_grad for tensor in data_tensors): + raise RuntimeError( + "recurrent_kda is inference-only and has no backward; use chunk_kda for " + "training or call under torch.no_grad() / torch.inference_mode()" + ) + + +def recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None = None, + *, + cu_seqlens: torch.Tensor | None = None, + output_final_state: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Apply the fused O(T) KDA delta rule for decode and inference prefill. + + Args: + q: Queries with shape ``[B, T, H, K]``; scaled by ``1/sqrt(K)`` internally. + k: Keys with the same shape as ``q``. + v: Values with shape ``[B, T, H, V]``. + gate: Per-token log2 decay with the same shape as ``q``, as produced by + ``bounded_gate_cumsum(chunk_size=1)`` — not the chunk-local + cumulative gate that ``chunk_kda`` consumes. + beta: Per-token write gate with shape ``[B, T, H]``. + initial_state: Optional recurrent state with one ``[H, K, V]`` entry per + logical sequence. + cu_seqlens: Optional device-resident int32 offsets selecting packed + ``[1, T, H, D]`` execution. Repeated offsets are empty padding slots + whose state passes through unchanged, and the terminal offset may sit + below the physical token capacity; values past it are outside the + operation's contract, so fixed-shape CUDA graphs can replay with + different boundaries and active lengths. + output_final_state: Also return the final recurrent state. When false, + the state is neither allocated nor written. + + Returns: + The output in ``q.dtype`` and, when requested, the FP32 recurrent state. + + The scan computes in FP32 regardless of input dtype. The fused scan is + inference-only: when autograd is enabled, calls whose data inputs require + gradients are rejected instead of silently detaching. + """ + _validate_recurrent_kda_inputs(q, k, v, gate, beta, initial_state, cu_seqlens) + # The kernel loads every operand through an FP32 register cast, so only the + # layout needs normalizing here; recurrent states are always produced in FP32. + q, k, v, gate, beta = (tensor.contiguous() for tensor in (q, k, v, gate, beta)) + if initial_state is not None: + initial_state = initial_state.contiguous() + if output_final_state: + return _recurrent_fwd_op(q, k, v, gate, beta, initial_state, cu_seqlens) + return _recurrent_fwd_no_state_op(q, k, v, gate, beta, initial_state, cu_seqlens), None + + +__all__ = ["recurrent_kda"] diff --git a/attn_gym/linear/kda/naive.py b/attn_gym/linear/kda/naive.py index 9d62d025..fa76b92c 100644 --- a/attn_gym/linear/kda/naive.py +++ b/attn_gym/linear/kda/naive.py @@ -38,17 +38,31 @@ def naive_recurrent_kda( scale: query scale for q k^T (optional; default 1/sqrt(K)) initial_state: initial recurrent state (B, H, K, V) (optional) output_final_state: also return the final state (optional; in the compute dtype) - cu_seqlens: optional int32 offsets (varlen mode); the recurrence restarts at - each document boundary ``[cu_seqlens[i]:cu_seqlens[i + 1]]``. + cu_seqlens: optional int32 offsets (varlen mode) with the public packed + contract: the recurrence restarts at each document boundary, empty + documents pass their state through, and output rows past the + terminal offset stay zero. """ b, t, h, k_dim = q.shape if cu_seqlens is not None: if b != 1: raise ValueError(f"varlen mode packs documents into one row, got batch {b}") - outputs, final_states = [], [] - for doc, (bos, eos) in enumerate(pairwise(cu_seqlens.tolist())): - output, final_state = naive_recurrent_kda( + offsets = cu_seqlens.tolist() + num_documents = len(offsets) - 1 + compute_dtype = torch.promote_types(q.dtype, torch.float32) + output = torch.zeros(1, t, h, v.shape[-1], dtype=q.dtype, device=q.device) + final_state = ( + initial_state.to(compute_dtype).clone() + if initial_state is not None + else torch.zeros( + num_documents, h, k_dim, v.shape[-1], dtype=compute_dtype, device=q.device + ) + ) + for doc, (bos, eos) in enumerate(pairwise(offsets)): + if bos == eos: + continue + doc_output, doc_state = naive_recurrent_kda( q[:, bos:eos], k[:, bos:eos], v[:, bos:eos], @@ -56,14 +70,11 @@ def naive_recurrent_kda( beta[:, bos:eos], scale, initial_state[doc : doc + 1] if initial_state is not None else None, - output_final_state, + output_final_state=True, ) - outputs.append(output) - final_states.append(final_state) - return ( - torch.cat(outputs, dim=1), - torch.cat(final_states) if output_final_state else None, - ) + output[:, bos:eos] = doc_output + final_state[doc] = doc_state[0] + return output, (final_state if output_final_state else None) output_dtype = q.dtype optional = () if initial_state is None else (initial_state,) diff --git a/attn_gym/linear/kda/validation.py b/attn_gym/linear/kda/validation.py new file mode 100644 index 00000000..7f0520ac --- /dev/null +++ b/attn_gym/linear/kda/validation.py @@ -0,0 +1,75 @@ +# Copyright (c) 2025 Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Shared validation for the public KDA delta-rule contract. + +Both ops and both implementations share one tensor/packed contract, stated +here once; implementation-specific constraints stay with the implementation. +""" + +from __future__ import annotations + +import torch + +SUPPORTED_INPUT_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +def validate_kda_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + *, + op_name: str, + gate_name: str, +) -> None: + """Validate the shared KDA operation contract before normalizing inputs.""" + if q.ndim != 4: + raise ValueError(f"q must have shape [B, T, H, K], got {tuple(q.shape)}") + batch, tokens, heads, key_dim = q.shape + if batch == 0 or tokens == 0 or heads == 0 or key_dim == 0: + raise ValueError(f"q must have nonempty dimensions, got {tuple(q.shape)}") + if k.shape != q.shape: + raise ValueError(f"k must have shape {tuple(q.shape)}, got {tuple(k.shape)}") + if gate.shape != q.shape: + raise ValueError(f"{gate_name} must have shape {tuple(q.shape)}, got {tuple(gate.shape)}") + if v.ndim != 4 or v.shape[:3] != (batch, tokens, heads) or v.shape[-1] < 1: + raise ValueError( + f"v must have shape [{batch}, {tokens}, {heads}, V], got {tuple(v.shape)}" + ) + if beta.shape != (batch, tokens, heads): + raise ValueError(f"beta must have shape {(batch, tokens, heads)}, got {tuple(beta.shape)}") + if cu_seqlens is not None: + if batch != 1: + raise ValueError("packed cu_seqlens require q to have batch size one") + if cu_seqlens.ndim != 1 or cu_seqlens.shape[0] < 2: + raise ValueError("cu_seqlens must have shape [num_sequences + 1]") + if ( + cu_seqlens.dtype != torch.int32 + or not cu_seqlens.is_contiguous() + or cu_seqlens.device != q.device + ): + raise ValueError("cu_seqlens must be contiguous int32 on q.device") + state_batch = batch if cu_seqlens is None else cu_seqlens.shape[0] - 1 + expected_state = (state_batch, heads, key_dim, v.shape[-1]) + if initial_state is not None and initial_state.shape != expected_state: + raise ValueError( + f"initial_state must have shape {expected_state}, got {tuple(initial_state.shape)}" + ) + data_tensors = (q, k, v, gate, beta) + if initial_state is not None: + data_tensors += (initial_state,) + if not all(tensor.device == q.device for tensor in data_tensors): + raise ValueError(f"all {op_name} inputs must be on the same device") + if any(tensor.dtype not in SUPPORTED_INPUT_DTYPES for tensor in data_tensors): + supported = ", ".join(str(dtype) for dtype in SUPPORTED_INPUT_DTYPES) + raise TypeError(f"{op_name} inputs must use one of {supported}") + + +__all__ = ["SUPPORTED_INPUT_DTYPES", "validate_kda_inputs"] diff --git a/test/test_kda_recurrent.py b/test/test_kda_recurrent.py new file mode 100644 index 00000000..25ef9292 --- /dev/null +++ b/test/test_kda_recurrent.py @@ -0,0 +1,324 @@ +# Copyright (c) 2025 Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Correctness and integration tests for the fused KDA recurrence.""" + +from itertools import pairwise + +import pytest +import torch +import torch.nn.functional as F + +pytest.importorskip("triton") + +from attn_gym.linear import naive_recurrent_kda, recurrent_kda +from attn_gym.linear.kda.fwd.triton.recurrent import ( + _recurrent_fwd_no_state_op, + _recurrent_fwd_op, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="recurrent_kda requires CUDA" +) + +BLACKWELL = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) + + +def _inputs( + batch: int = 2, + tokens: int = 37, + heads: int = 2, + key_dim: int = 64, + value_dim: int = 64, + dtype: torch.dtype = torch.float32, + initial_state: bool = False, + seed: int = 0, +): + torch.manual_seed(seed) + # KDA L2-normalizes q and k before the core; unnormalized keys make the + # delta-rule recurrence exponentially unstable and useless for comparison. + q = F.normalize(torch.randn(batch, tokens, heads, key_dim, device="cuda"), dim=-1).to(dtype) + k = F.normalize(torch.randn(batch, tokens, heads, key_dim, device="cuda"), dim=-1).to(dtype) + v = torch.randn(batch, tokens, heads, value_dim, device="cuda", dtype=dtype) + # Realistic bounded log2 decays keep the recurrence stable over the scan. + gate = -torch.rand(batch, tokens, heads, key_dim, device="cuda") * 3.0 + beta = torch.rand(batch, tokens, heads, device="cuda") + state = torch.randn(batch, heads, key_dim, value_dim, device="cuda") if initial_state else None + return q, k, v, gate, beta, state + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("use_initial_state", [False, True]) +@pytest.mark.parametrize("tokens", [1, 37]) +def test_recurrent_matches_naive_dense(dtype: torch.dtype, use_initial_state: bool, tokens: int): + """Match the reference scan on dense batches, including single-token decode.""" + q, k, v, gate, beta, state = _inputs( + tokens=tokens, dtype=dtype, initial_state=use_initial_state + ) + + output, final_state = recurrent_kda(q, k, v, gate, beta, state, output_final_state=True) + expected, expected_state = naive_recurrent_kda( + q.float(), + k.float(), + v.float(), + gate, + beta, + initial_state=state, + output_final_state=True, + ) + + tolerance = 1e-5 if dtype == torch.float32 else 2e-2 + assert output.dtype == q.dtype + assert final_state is not None and final_state.dtype == torch.float32 + torch.testing.assert_close(output.float(), expected, rtol=tolerance, atol=tolerance) + torch.testing.assert_close(final_state, expected_state, rtol=tolerance, atol=tolerance) + + +@pytest.mark.parametrize(("key_dim", "value_dim"), [(80, 48), (128, 128)]) +def test_recurrent_matches_naive_non_power_of_two(key_dim: int, value_dim: int): + """Mask partial key and value blocks correctly.""" + q, k, v, gate, beta, _ = _inputs(key_dim=key_dim, value_dim=value_dim, seed=1) + + output, final_state = recurrent_kda(q, k, v, gate, beta, output_final_state=True) + expected, expected_state = naive_recurrent_kda(q, k, v, gate, beta, output_final_state=True) + torch.testing.assert_close(output, expected, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(final_state, expected_state, rtol=1e-5, atol=1e-5) + + +def test_recurrent_packed_capacity_and_empty_slots(): + """Pass empty-slot state through and ignore rows past the terminal offset.""" + q, k, v, gate, beta, _ = _inputs(batch=1, tokens=32, seed=3) + cu_seqlens = torch.tensor([0, 0, 11, 27, 27], device="cuda", dtype=torch.int32) + initial_state = torch.randn(4, q.shape[2], q.shape[3], v.shape[-1], device="cuda") + + output, final_state = recurrent_kda( + q, + k, + v, + gate, + beta, + initial_state, + cu_seqlens=cu_seqlens, + output_final_state=True, + ) + assert final_state is not None + # Empty padding slots preserve their incoming state bitwise. + torch.testing.assert_close(final_state[0], initial_state[0], rtol=0, atol=0) + torch.testing.assert_close(final_state[3], initial_state[3], rtol=0, atol=0) + for sequence, (start, end) in enumerate(pairwise(cu_seqlens.cpu().tolist())): + if start == end: + continue + expected, expected_state = naive_recurrent_kda( + q[:, start:end], + k[:, start:end], + v[:, start:end], + gate[:, start:end], + beta[:, start:end], + initial_state=initial_state[sequence : sequence + 1], + output_final_state=True, + ) + torch.testing.assert_close(output[:, start:end], expected, rtol=1e-5, atol=1e-5) + torch.testing.assert_close( + final_state[sequence : sequence + 1], expected_state, rtol=1e-5, atol=1e-5 + ) + + +@pytest.mark.skipif(not BLACKWELL, reason="chunk_kda requires CUDA capability 10.0") +def test_recurrent_agrees_with_chunked_core(): + """Cross-check the decode scan against the training core on shared inputs.""" + from attn_gym.linear import bounded_gate_cumsum, chunk_kda + + batch, tokens, heads, head_dim = 2, 128, 2, 128 + q, k, v, _, beta = _inputs( + batch=batch, + tokens=tokens, + heads=heads, + key_dim=head_dim, + value_dim=head_dim, + dtype=torch.bfloat16, + seed=4, + )[:5] + raw_gate = torch.randn(batch, tokens, heads, head_dim, device="cuda", dtype=torch.bfloat16) + a_log = torch.zeros(heads, device="cuda") + dt_bias = torch.zeros(heads, head_dim, device="cuda") + + per_token_gate = bounded_gate_cumsum(raw_gate, a_log, dt_bias, chunk_size=1) + cumulative_gate = bounded_gate_cumsum(raw_gate, a_log, dt_bias, chunk_size=64) + + recurrent_output, recurrent_state = recurrent_kda( + q, k, v, per_token_gate, beta, output_final_state=True + ) + chunked_output, chunked_state = chunk_kda( + q, k, v, cumulative_gate, beta, output_final_state=True + ) + torch.testing.assert_close( + recurrent_output.float(), chunked_output.float(), rtol=5e-2, atol=5e-2 + ) + torch.testing.assert_close(recurrent_state, chunked_state, rtol=5e-2, atol=5e-2) + + +def test_recurrent_validates_public_contract(): + """Reject malformed inputs at the public boundary before a kernel launch.""" + q, k, v, gate, beta, _ = _inputs(tokens=4) + with pytest.raises(ValueError, match="k must have shape"): + recurrent_kda(q, k[:, :-1], v, gate, beta) + with pytest.raises(ValueError, match="gate must have shape"): + recurrent_kda(q, k, v, gate[..., :-1], beta) + with pytest.raises(ValueError, match="beta must have shape"): + recurrent_kda(q, k, v, gate, beta[:, :, :-1]) + with pytest.raises(ValueError, match="initial_state must have shape"): + recurrent_kda(q, k, v, gate, beta, q.new_zeros(1, 1, 1, 1)) + with pytest.raises(ValueError, match="batch size one"): + recurrent_kda( + q, + k, + v, + gate, + beta, + cu_seqlens=torch.tensor([0, 4], device="cuda", dtype=torch.int32), + ) + with pytest.raises(ValueError, match="num_sequences"): + recurrent_kda( + q[:1], + k[:1], + v[:1], + gate[:1], + beta[:1], + cu_seqlens=torch.tensor([0], device="cuda", dtype=torch.int32), + ) + with pytest.raises(ValueError, match="contiguous int32"): + recurrent_kda( + q[:1], + k[:1], + v[:1], + gate[:1], + beta[:1], + cu_seqlens=torch.tensor([0, 4], device="cuda", dtype=torch.int64), + ) + with pytest.raises(ValueError, match="requires K in"): + big = q.new_zeros(1, 4, 2, 512) + recurrent_kda(big, big, big, big, q.new_zeros(1, 4, 2)) + with pytest.raises(TypeError, match="inputs must use one of"): + recurrent_kda(q.double(), k.double(), v.double(), gate.double(), beta.double()) + + +@pytest.mark.parametrize("operand", range(6)) +def test_recurrent_rejects_gradient_tracking(operand: int): + """State a clear inference-only contract for every gradient-tracking operand.""" + tensors = list(_inputs(tokens=4, initial_state=True)) + tensors[operand] = tensors[operand].float().requires_grad_() + with pytest.raises(RuntimeError, match="inference-only"): + recurrent_kda(*tensors) + with torch.no_grad(): + output, _ = recurrent_kda(*tensors) + assert not output.requires_grad + + +@pytest.mark.parametrize("packed", [False, True]) +def test_recurrent_custom_op_registration(packed: bool): + """Exercise the schema and fake implementation for both modes.""" + batch = 1 if packed else 2 + q, k, v, gate, beta, _ = _inputs(batch=batch, tokens=17) + cu_seqlens = torch.tensor([0, 2, 7, 17], device="cuda", dtype=torch.int32) if packed else None + num_sequences = 3 if packed else batch + state = torch.randn(num_sequences, q.shape[2], q.shape[3], v.shape[-1], device="cuda") + torch.library.opcheck(_recurrent_fwd_op, (q, k, v, gate, beta, state, cu_seqlens)) + torch.library.opcheck(_recurrent_fwd_no_state_op, (q, k, v, gate, beta, state, cu_seqlens)) + + +@pytest.mark.parametrize("output_final_state", [False, True]) +def test_recurrent_fullgraph_compile(output_final_state: bool): + """Compile both optional-state branches of the public operation.""" + q, k, v, gate, beta, state = _inputs(initial_state=True) + + expected, expected_state = recurrent_kda( + q, k, v, gate, beta, state, output_final_state=output_final_state + ) + compiled = torch.compile(recurrent_kda, fullgraph=True) + output, final_state = compiled( + q, k, v, gate, beta, state, output_final_state=output_final_state + ) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + if output_final_state: + torch.testing.assert_close(final_state, expected_state, rtol=0, atol=0) + else: + assert final_state is None and expected_state is None + + +def test_recurrent_cuda_graph_replay(): + """Replay fixed shapes with mutated boundaries, values, and history.""" + q, k, v, gate, beta, _ = _inputs(batch=1, tokens=32, seed=5) + cu_seqlens = torch.tensor([0, 11, 27, 32], device="cuda", dtype=torch.int32) + initial_state = torch.randn(3, q.shape[2], q.shape[3], v.shape[-1], device="cuda") + _recurrent_fwd_op(q, k, v, gate, beta, initial_state, cu_seqlens) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output, captured_state = _recurrent_fwd_op( + q, k, v, gate, beta, initial_state, cu_seqlens + ) + + active_tokens = 23 + with torch.no_grad(): + initial_state.add_(0.25) + cu_seqlens.copy_( + torch.tensor([0, 8, active_tokens, active_tokens], device="cuda", dtype=torch.int32) + ) + q[:, active_tokens:].fill_(float("nan")) + v[:, active_tokens:].fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + + expected_output, expected_state = _recurrent_fwd_op( + q, k, v, gate, beta, initial_state, cu_seqlens + ) + torch.testing.assert_close( + captured_output[:, :active_tokens], + expected_output[:, :active_tokens], + rtol=0, + atol=0, + ) + torch.testing.assert_close(captured_state, expected_state, rtol=0, atol=0) + + +def test_recurrent_launches_beyond_grid_y_limit(): + """Flat 1-D launches must survive sequence-head counts above 65,535.""" + batch, heads, head_dim = 2200, 32, 16 + assert batch * heads > 65_535 + q, k, v, gate, beta, _ = _inputs( + batch=batch, tokens=1, heads=heads, key_dim=head_dim, value_dim=head_dim, seed=6 + ) + output, _ = recurrent_kda(q, k, v, gate, beta) + expected, _ = naive_recurrent_kda(q, k, v, gate, beta) + torch.testing.assert_close(output, expected, rtol=1e-5, atol=1e-5) + + +def test_naive_recurrent_packed_matches_public_contract(): + """The pure reference honors packed semantics: empty slots and capacity tails.""" + q, k, v, gate, beta, _ = _inputs(batch=1, tokens=32, seed=7) + cu_seqlens = torch.tensor([0, 0, 11, 27, 27], device="cuda", dtype=torch.int32) + initial_state = torch.randn(4, q.shape[2], q.shape[3], v.shape[-1], device="cuda") + + output, final_state = naive_recurrent_kda( + q, + k, + v, + gate, + beta, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=cu_seqlens, + ) + assert output.shape == v.shape and final_state is not None + torch.testing.assert_close(final_state[0], initial_state[0], rtol=0, atol=0) + torch.testing.assert_close(final_state[3], initial_state[3], rtol=0, atol=0) + torch.testing.assert_close(output[:, 27:], torch.zeros_like(output[:, 27:]), rtol=0, atol=0) + fused, fused_state = recurrent_kda( + q, k, v, gate, beta, initial_state, cu_seqlens=cu_seqlens, output_final_state=True + ) + torch.testing.assert_close(output[:, :27], fused[:, :27], rtol=1e-5, atol=1e-5) + torch.testing.assert_close(final_state, fused_state, rtol=1e-5, atol=1e-5)