diff --git a/attn_gym/linear/kda/api.py b/attn_gym/linear/kda/api.py index 58aaae06..f27d797b 100644 --- a/attn_gym/linear/kda/api.py +++ b/attn_gym/linear/kda/api.py @@ -38,6 +38,7 @@ def chunk_kda( output_final_state: bool = False, fastmath: bool = False, autotune: bool = True, + fuse_q_l2norm: bool = False, impl: Impl | str = Impl.FUSED, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Apply chunk-parallel KDA for training and prefill. @@ -62,6 +63,10 @@ def chunk_kda( autotune: Benchmark candidate kernel configurations when true (winners are cached and reused); use fixed heuristics when false for repeatable selection across machines and cache states. + fuse_q_l2norm: Accept unnormalized q and L2-normalize each row inside + the fused core (default ``eps=1e-6``), skipping the standalone + normalization pass; k must still be pre-normalized. Forward-only, + eager-only, and rejected with ``"reference"``. impl: ``"fused"`` uses the Blackwell kernels with first-order autograd; ``"reference"`` uses differentiable eager PyTorch in FP32, with no automatic fallback. @@ -73,6 +78,8 @@ def chunk_kda( selected_impl = resolve_impl(impl) if selected_impl is Impl.REFERENCE and fastmath: raise ValueError("fastmath applies only to impl='fused'") + if selected_impl is Impl.REFERENCE and fuse_q_l2norm: + raise ValueError("fuse_q_l2norm applies only to impl='fused'") validate_kda_inputs( q, k, @@ -96,6 +103,7 @@ def chunk_kda( output_final_state=output_final_state, fastmath=fastmath, autotune=autotune, + fuse_q_l2norm=fuse_q_l2norm, ) return reference_kda( partial(naive_chunk_kda_from_cumulative, chunk_size=_CHUNK_SIZE), diff --git a/attn_gym/linear/kda/fwd/cute/chunk_kda_fwd.py b/attn_gym/linear/kda/fwd/cute/chunk_kda_fwd.py index 0a90c628..c2101fcd 100644 --- a/attn_gym/linear/kda/fwd/cute/chunk_kda_fwd.py +++ b/attn_gym/linear/kda/fwd/cute/chunk_kda_fwd.py @@ -87,8 +87,15 @@ def _chunk_kda_fwd( *, output_final_state: bool, autotune: bool, + fuse_q_l2norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor]: - """Run the optimized KDA core using an already selected chunk schedule.""" + """Run the optimized KDA core using an already selected chunk schedule. + + With ``fuse_q_l2norm`` the caller passes raw (unnormalized) q: the grams + stay in the raw-q basis and the output stage computes and applies each + row's inverse norm in-kernel, skipping the standalone q normalization + pass. Forward-only; the returned Aqk stays in the raw-q basis. + """ scale = _HEAD_DIM**-0.5 with profiler_range("kda/fused/chunk_kda_fwd_intra"): @@ -127,6 +134,7 @@ def _chunk_kda_fwd( chunk_size=_CHUNK_SIZE, metadata=metadata, autotune=autotune, + fuse_q_l2norm=fuse_q_l2norm, ) return output, final_state, Aqk, Akk @@ -139,6 +147,7 @@ def _chunk_kda_fwd_shared( beta: torch.Tensor, initial_state: torch.Tensor | None, autotune: bool, + fuse_q_l2norm: bool, output_final_state: bool, ): """Keep the complete composed forward behind one compiler-opaque boundary.""" @@ -154,6 +163,7 @@ def _chunk_kda_fwd_shared( None, output_final_state=output_final_state, autotune=autotune, + fuse_q_l2norm=fuse_q_l2norm, ) @@ -165,6 +175,7 @@ def _chunk_kda_fwd_cuda( beta: torch.Tensor, initial_state: torch.Tensor | None, autotune: bool, + fuse_q_l2norm: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: output, _final_state, Aqk, Akk = _chunk_kda_fwd_shared( q, @@ -174,6 +185,7 @@ def _chunk_kda_fwd_cuda( beta, initial_state, autotune, + fuse_q_l2norm, False, ) return output, Aqk, Akk @@ -187,6 +199,7 @@ def _chunk_kda_fwd_with_state_cuda( beta: torch.Tensor, initial_state: torch.Tensor | None, autotune: bool, + fuse_q_l2norm: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: return _chunk_kda_fwd_shared( q, @@ -196,6 +209,7 @@ def _chunk_kda_fwd_with_state_cuda( beta, initial_state, autotune, + fuse_q_l2norm, True, ) @@ -210,6 +224,7 @@ def _chunk_kda_fwd_ragged_shared( cu_seqlens: torch.Tensor, chunk_offsets: torch.Tensor, autotune: bool, + fuse_q_l2norm: bool, output_final_state: bool, ): """Run ragged forward with caller-prepared routing and fixed-schema tape.""" @@ -231,6 +246,7 @@ def _chunk_kda_fwd_ragged_shared( metadata, output_final_state=output_final_state, autotune=autotune, + fuse_q_l2norm=fuse_q_l2norm, ) @@ -244,6 +260,7 @@ def _chunk_kda_fwd_ragged_cuda( cu_seqlens: torch.Tensor, chunk_offsets: torch.Tensor, autotune: bool, + fuse_q_l2norm: bool, ): output, _state, Aqk, Akk = _chunk_kda_fwd_ragged_shared( q, @@ -255,6 +272,7 @@ def _chunk_kda_fwd_ragged_cuda( cu_seqlens, chunk_offsets, autotune, + fuse_q_l2norm, False, ) return output, Aqk, Akk @@ -270,6 +288,7 @@ def _chunk_kda_fwd_ragged_with_state_cuda( cu_seqlens: torch.Tensor, chunk_offsets: torch.Tensor, autotune: bool, + fuse_q_l2norm: bool, ): return _chunk_kda_fwd_ragged_shared( q, @@ -281,6 +300,7 @@ def _chunk_kda_fwd_ragged_with_state_cuda( cu_seqlens, chunk_offsets, autotune, + fuse_q_l2norm, True, ) diff --git a/attn_gym/linear/kda/fwd/triton/chunk_gla_fwd_o.py b/attn_gym/linear/kda/fwd/triton/chunk_gla_fwd_o.py index 24bb08c8..ebccec00 100644 --- a/attn_gym/linear/kda/fwd/triton/chunk_gla_fwd_o.py +++ b/attn_gym/linear/kda/fwd/triton/chunk_gla_fwd_o.py @@ -25,6 +25,9 @@ from attn_gym.linear.kda.chunk_scheduler import RaggedChunkMetadata, load_ragged_chunk_work from attn_gym.linear.kda.utils import autotune_cache_kwargs, exp, exp2 +# Must match l2norm's default eps so the fused route reproduces the standalone pass. +_L2NORM_EPS = tl.constexpr(1e-6) + @triton.heuristics( { @@ -50,7 +53,7 @@ triton.Config({"BK": 64, "BV": 64}, num_warps=2, num_stages=4), triton.Config({"BK": 64, "BV": 64}, num_warps=8, num_stages=4), ], - key=["H", "K", "V", "T", "BT"], + key=["H", "K", "V", "T", "BT", "FUSE_Q_L2NORM"], **autotune_cache_kwargs, ) @triton.jit( @@ -84,6 +87,7 @@ def chunk_gla_fwd_kernel_o( USE_EXP2: tl.constexpr, IS_VARLEN: tl.constexpr, USE_INT64_OFFSETS: tl.constexpr, + FUSE_Q_L2NORM: tl.constexpr = False, ): i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) if USE_INT64_OFFSETS: @@ -134,6 +138,7 @@ def chunk_gla_fwd_kernel_o( A += ptr_offset((bos, i_h), (H * BT, BT)) b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_ssq = tl.zeros([BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K @@ -144,6 +149,9 @@ def chunk_gla_fwd_kernel_o( # [BT, BK] b_q = tl.load(p_q, mask=m_qg, other=0.0) + if FUSE_Q_L2NORM: + b_q32 = b_q.to(tl.float32) + b_ssq += tl.sum(b_q32 * b_q32, 1) # [BT, BK] b_g = tl.load(p_g, mask=m_qg, other=0.0).to(tl.float32) # [BT, BK] @@ -167,6 +175,10 @@ def chunk_gla_fwd_kernel_o( b_A = tl.load(p_A, mask=m_t[:, None], other=0.0) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v) + if FUSE_Q_L2NORM: + # Both output terms are linear in the raw q row (Aqk carries the raw-q + # gram), so one row scale by 1/||q|| completes the deferred L2 norm. + b_o *= (1 / tl.sqrt(b_ssq + _L2NORM_EPS))[:, None] tl.store(p_o, b_o.to(o.dtype.element_ty), mask=m_tv) @@ -188,13 +200,18 @@ def _compose_output_tma( BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + FUSE_Q_L2NORM: tl.constexpr, ): """Compose one complete output tile with TMA-backed tensor descriptors.""" b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_ssq = tl.zeros([BT], dtype=tl.float32) for key_tile in range(tl.cdiv(K, BK)): key_start = key_tile * BK b_q = q_desc.load([batch, token_start, head, key_start]) b_q = tl.reshape(b_q, [BT, BK]) + if FUSE_Q_L2NORM: + b_q32 = b_q.to(tl.float32) + b_ssq += tl.sum(b_q32 * b_q32, 1) b_g = g_desc.load([batch, token_start, head, key_start]) b_g = tl.reshape(b_g, [BT, BK]).to(tl.float32) b_qg = (b_q * exp2(b_g)).to(b_q.dtype) @@ -211,6 +228,10 @@ def _compose_output_tma( b_v = v_desc.load([batch, token_start, head, value_tile * BV]) b_v = tl.reshape(b_v, [BT, BV]) b_o += tl.dot(b_A.to(b_v.dtype), b_v) + if FUSE_Q_L2NORM: + # Both accumulated terms are linear in the raw q row (Aqk carries the + # raw-q gram), so one row scale by 1/||q|| completes the deferred norm. + b_o *= (1 / tl.sqrt(b_ssq + _L2NORM_EPS))[:, None] o_desc.store( [batch, token_start, head, value_tile * BV], tl.reshape(b_o.to(b_v.dtype), [1, BT, 1, BV]), @@ -232,6 +253,7 @@ def chunk_gla_fwd_kernel_o_tma( BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + FUSE_Q_L2NORM: tl.constexpr = False, ): """Compose fixed KDA output tiles with TMA-backed tensor descriptors.""" value_tile, chunk, batch_head = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -253,6 +275,7 @@ def chunk_gla_fwd_kernel_o_tma( BT, BK, BV, + FUSE_Q_L2NORM, ) @@ -282,6 +305,7 @@ def chunk_gla_fwd_kernel_o_ragged_tma( BK: tl.constexpr, BV: tl.constexpr, num_sequences, + FUSE_Q_L2NORM: tl.constexpr = False, ): """Compose full ragged chunks with TMA and partial tails with masked pointers.""" i_v, global_chunk, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) @@ -312,6 +336,7 @@ def chunk_gla_fwd_kernel_o_ragged_tma( BT, BK, BV, + FUSE_Q_L2NORM, ) else: o_i = tl.arange(0, BT) @@ -321,6 +346,7 @@ def chunk_gla_fwd_kernel_o_ragged_tma( m_v = o_v < V m_tv = m_t[:, None] & m_v[None, :] b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_ssq = tl.zeros([BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K @@ -332,6 +358,9 @@ def chunk_gla_fwd_kernel_o_ragged_tma( (H * K * V, K * V, V, 1), ) b_q = tl.load(p_q, mask=m_qg, other=0.0) + if FUSE_Q_L2NORM: + b_q32 = b_q.to(tl.float32) + b_ssq += tl.sum(b_q32 * b_q32, 1) b_g = tl.load(p_g, mask=m_qg, other=0.0).to(tl.float32) b_qg = (b_q * exp2(b_g)).to(b_q.dtype) b_h = tl.load(p_h, mask=m_k[:, None] & m_v[None, :], other=0.0) @@ -344,6 +373,8 @@ def chunk_gla_fwd_kernel_o_ragged_tma( b_A = tl.load(p_A, mask=m_t[:, None], other=0.0) b_A = tl.where(o_i[:, None] >= o_i[None, :], b_A, 0.0).to(b_v.dtype) b_o += tl.dot(b_A, b_v) + if FUSE_Q_L2NORM: + b_o *= (1 / tl.sqrt(b_ssq + _L2NORM_EPS))[:, None] p_o = o + ptr_offset((o_t[:, None], i_h, o_v[None, :]), (H * V, V, 1)) tl.store(p_o, b_o.to(o.dtype.element_ty), mask=m_tv) @@ -367,8 +398,14 @@ def chunk_gla_fwd_o_gk( chunk_size: int = 64, metadata: RaggedChunkMetadata | None = None, autotune: bool = True, + fuse_q_l2norm: bool = False, ) -> torch.Tensor: - """Compose fixed-length or packed KDA intra- and inter-chunk output terms.""" + """Compose fixed-length or packed KDA intra- and inter-chunk output terms. + + ``fuse_q_l2norm`` applies the q L2 norm as an output row scale computed + in-kernel from the raw q rows; ``q`` and ``A`` must then carry the raw-q + basis. + """ if metadata is not None: metadata.validate_chunk_size(chunk_size) batch, tokens, heads, key_dim = q.shape @@ -422,6 +459,7 @@ def chunk_gla_fwd_o_gk( BT=chunk_size, BK=block_key_dim, BV=block_value_dim, + FUSE_Q_L2NORM=fuse_q_l2norm, num_warps=2, num_stages=3, ) @@ -456,12 +494,14 @@ def chunk_gla_fwd_o_gk( BK=block_key_dim, BV=block_value_dim, num_sequences=metadata.cu_seqlens.shape[0] - 1, + FUSE_Q_L2NORM=fuse_q_l2norm, num_warps=2, num_stages=3, # The rarely taken partial-chunk pointer branch pushes this - # kernel to 180 registers vs the dense kernel's 134; cap at the - # dense budget so only tail CTAs pay with spills. - maxnreg=136, + # kernel to 180 registers vs the dense kernel's 134; cap near + # the dense budget so only tail CTAs pay with spills. The + # fused q-norm accumulator needs a little extra headroom. + maxnreg=152 if fuse_q_l2norm else 136, ) return output @@ -488,6 +528,7 @@ def grid(meta): BT=chunk_size, num_sequences=(0 if metadata is None else metadata.cu_seqlens.shape[0] - 1), USE_EXP2=True, + FUSE_Q_L2NORM=fuse_q_l2norm, ) return output diff --git a/attn_gym/linear/kda/impl/fused.py b/attn_gym/linear/kda/impl/fused.py index 9016d04f..1230bec3 100644 --- a/attn_gym/linear/kda/impl/fused.py +++ b/attn_gym/linear/kda/impl/fused.py @@ -49,16 +49,17 @@ def forward( output_final_state, fastmath, autotune, + fuse_q_l2norm, ): if cu_seqlens is None: assert chunk_offsets is None if output_final_state: output, state, aqk, akk = chunk_fwd_with_state_op( - q, k, v, cumulative_gate, beta, initial_state, autotune + q, k, v, cumulative_gate, beta, initial_state, autotune, fuse_q_l2norm ) else: output, aqk, akk = chunk_fwd_op( - q, k, v, cumulative_gate, beta, initial_state, autotune + q, k, v, cumulative_gate, beta, initial_state, autotune, fuse_q_l2norm ) elif output_final_state: assert chunk_offsets is not None @@ -72,6 +73,7 @@ def forward( cu_seqlens, chunk_offsets, autotune, + fuse_q_l2norm, ) else: assert chunk_offsets is not None @@ -85,6 +87,7 @@ def forward( cu_seqlens, chunk_offsets, autotune, + fuse_q_l2norm, ) ctx.save_for_backward( q, @@ -141,7 +144,7 @@ def backward(ctx, d_output, d_final_state=None): else: dq, dk, dv, dg, db = chunk_bwd_op(*args) d_initial_state = None - return dq, dk, dv, dg, db, d_initial_state, None, None, None, None, None + return dq, dk, dv, dg, db, d_initial_state, None, None, None, None, None, None def chunk_forward( @@ -157,8 +160,32 @@ def chunk_forward( output_final_state: bool = False, fastmath: bool = False, autotune: bool = True, + fuse_q_l2norm: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Normalize inputs and invoke the registered fused chunk operators.""" + """Normalize inputs and invoke the registered fused chunk operators. + + ``fuse_q_l2norm`` accepts unnormalized q and L2-normalizes each row + inside the fused core (default ``eps=1e-6``), skipping the standalone q + normalization pass; k must still be pre-normalized by the caller. It is + forward-only (gradient-tracking inputs are rejected while grad mode is + enabled) and eager-only. The intermediate grams then carry + raw-q magnitudes (about ``sqrt(K)`` larger), so this flag assumes the + bounded-gate contract (``bounded_gate_cumsum``); unbounded synthetic gates + can overflow the BF16 gram range. + """ + if fuse_q_l2norm: + if torch.compiler.is_compiling(): + raise NotImplementedError( + "fuse_q_l2norm is eager-only; it bypasses the compiler-opaque op" + ) + needs_grad = torch.is_grad_enabled() and any( + tensor is not None and tensor.requires_grad + for tensor in (q, k, v, cumulative_gate, beta, initial_state) + ) + if needs_grad: + raise NotImplementedError( + "fuse_q_l2norm is forward-only; normalize q outside the kernel to keep gradients" + ) if metadata is not None: assert cu_seqlens is None metadata.validate_chunk_size(_CHUNK_SIZE) @@ -197,6 +224,7 @@ def chunk_forward( True, fastmath, autotune, + fuse_q_l2norm, ) else: output = _ChunkKDA.apply( @@ -211,6 +239,7 @@ def chunk_forward( False, fastmath, autotune, + fuse_q_l2norm, ) state = None return output.reshape(output_shape).to(output_dtype), state diff --git a/attn_gym/linear/kda/ops.py b/attn_gym/linear/kda/ops.py index 82a75c9e..cf414d2b 100644 --- a/attn_gym/linear/kda/ops.py +++ b/attn_gym/linear/kda/ops.py @@ -17,7 +17,7 @@ # Fixed-arity schema pairs avoid optional outputs on hot paths. _CHUNK_FWD_ARGS = ( "(Tensor q, Tensor k, Tensor v, Tensor cumulative_gate, Tensor beta, Tensor? initial_state," - " bool autotune)" + " bool autotune, bool fuse_q_l2norm)" ) torch.library.define( "attn_gym::kda_chunk_fwd", @@ -30,7 +30,8 @@ _CHUNK_RAGGED_FWD_ARGS = ( "(Tensor q, Tensor k, Tensor v, Tensor cumulative_gate, Tensor beta, " - "Tensor? initial_state, Tensor cu_seqlens, Tensor chunk_offsets, bool autotune)" + "Tensor? initial_state, Tensor cu_seqlens, Tensor chunk_offsets, bool autotune, " + "bool fuse_q_l2norm)" ) torch.library.define( "attn_gym::kda_chunk_fwd_ragged", @@ -213,8 +214,9 @@ def _chunk_fwd_fake( beta: torch.Tensor, initial_state: torch.Tensor | None, autotune: bool, + fuse_q_l2norm: bool, ): - del k, cumulative_gate, beta, initial_state, autotune + del k, cumulative_gate, beta, initial_state, autotune, fuse_q_l2norm return _chunk_fwd_fake_common(q, v) @@ -227,8 +229,9 @@ def _chunk_fwd_with_state_fake( beta: torch.Tensor, initial_state: torch.Tensor | None, autotune: bool, + fuse_q_l2norm: bool, ): - del k, cumulative_gate, beta, initial_state, autotune + del k, cumulative_gate, beta, initial_state, autotune, fuse_q_l2norm output, aqk, akk = _chunk_fwd_fake_common(q, v) state = q.new_empty( (q.shape[0], q.shape[2], q.shape[3], v.shape[-1]), @@ -248,8 +251,10 @@ def _chunk_fwd_ragged_fake( cu_seqlens: torch.Tensor, chunk_offsets: torch.Tensor, autotune: bool, + fuse_q_l2norm: bool, ): del k, cumulative_gate, beta, initial_state, cu_seqlens, chunk_offsets, autotune + del fuse_q_l2norm return _chunk_fwd_fake_common(q, v) @@ -264,8 +269,9 @@ def _chunk_fwd_ragged_with_state_fake( cu_seqlens: torch.Tensor, chunk_offsets: torch.Tensor, autotune: bool, + fuse_q_l2norm: bool, ): - del k, cumulative_gate, beta, initial_state, chunk_offsets, autotune + del k, cumulative_gate, beta, initial_state, chunk_offsets, autotune, fuse_q_l2norm output, aqk, akk = _chunk_fwd_fake_common(q, v) state = q.new_empty( (cu_seqlens.shape[0] - 1, q.shape[2], q.shape[3], v.shape[-1]), diff --git a/test/test_kda_cute_forward.py b/test/test_kda_cute_forward.py index db39da50..639ab693 100644 --- a/test/test_kda_cute_forward.py +++ b/test/test_kda_cute_forward.py @@ -339,12 +339,12 @@ def test_chunk_kda_selects_direct_dense_or_ragged_route( module = importlib.import_module("attn_gym.linear.kda.impl.fused") routes = [] - def dense_forward(q, _k, v, _gate, _beta, _state, _tune): + def dense_forward(q, _k, v, _gate, _beta, _state, _tune, _fuse): routes.append("dense") tape = q.new_empty((*q.shape[:3], 64)) return torch.empty_like(v), tape, tape - def ragged_forward(q, _k, v, _gate, _beta, _state, _cu_seqlens, _chunk_offsets, _tune): + def ragged_forward(q, _k, v, _gate, _beta, _state, _cu_seqlens, _chunk_offsets, _tune, _fuse): routes.append("ragged") tape = q.new_empty((*q.shape[:3], 64)) return torch.empty_like(v), tape, tape @@ -628,6 +628,7 @@ def test_chunk_kda_op_registration(): beta.detach(), initial_state.detach(), True, + False, ) torch.library.opcheck(_chunk_kda_fwd_op, args, rtol=2e-2, atol=2e-3) torch.library.opcheck(_chunk_kda_fwd_with_state_op, args, rtol=2e-2, atol=2e-3) @@ -646,6 +647,7 @@ def test_chunk_kda_backward_op_registration(): beta, initial_state, True, + False, ) torch.library.opcheck( _chunk_kda_bwd_op, diff --git a/test/test_kda_fuse_q_l2norm.py b/test/test_kda_fuse_q_l2norm.py new file mode 100644 index 00000000..963ef632 --- /dev/null +++ b/test/test_kda_fuse_q_l2norm.py @@ -0,0 +1,128 @@ +# 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. + +"""The forward-only fused q L2 norm must match the explicit normalize-then-run path.""" + +import pytest +import torch + +from attn_gym.linear.kda import bounded_gate_cumsum, l2norm + +pytest.importorskip("cutlass", reason="fuse_q_l2norm requires the CuTeDSL backend") +from attn_gym.linear.kda import chunk_kda # noqa: E402 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), + reason="the optimized KDA core requires CUDA capability 10.0", +) + +_GRAD_OPERANDS = ("q", "k", "v", "cumulative_gate", "beta", "initial_state") + + +def _inputs(seq_lens: list[int], heads: int): + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(0) + total, dim = sum(seq_lens), 128 + + def rand(*shape: int) -> torch.Tensor: + return torch.randn(shape, generator=generator, device=device).bfloat16() + + q, k, v = rand(1, total, heads, dim), rand(1, total, heads, dim), rand(1, total, heads, dim) + # Stress the in-kernel norm: widely varying row scales plus an all-zero row + # (which must hit the eps floor identically to the standalone l2norm pass). + scales = torch.logspace(-3, 3, total, device=device).view(1, total, 1, 1) + q = (q.float() * scales).bfloat16() + q[0, total // 2] = 0.0 + raw_gate = rand(1, total, heads, dim) + a_log = torch.randn(heads, generator=generator, device=device) + dt_bias = torch.randn(heads, dim, generator=generator, device=device) + beta = torch.rand(1, total, heads, generator=generator, device=device, dtype=torch.float32) + cu_seqlens = None + if len(seq_lens) > 1: + cu_seqlens = torch.tensor( + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], device=device, dtype=torch.int32 + ) + cumulative_gate = bounded_gate_cumsum(raw_gate, a_log, dt_bias, cu_seqlens=cu_seqlens) + initial_state = torch.randn( + len(seq_lens), heads, dim, dim, generator=generator, device=device, dtype=torch.float32 + ) + return q, k, v, cumulative_gate, beta, cu_seqlens, initial_state + + +@pytest.mark.parametrize("seq_lens", [[512], [192, 65, 255]], ids=["dense", "ragged"]) +def test_fuse_q_l2norm_matches_explicit_normalization(seq_lens: list[int]) -> None: + q, k, v, cumulative_gate, beta, cu_seqlens, initial_state = _inputs(seq_lens, heads=4) + kn = l2norm(k) + expected, expected_state = chunk_kda( + l2norm(q), + kn, + v, + cumulative_gate, + beta, + initial_state, + cu_seqlens=cu_seqlens, + output_final_state=True, + ) + actual, actual_state = chunk_kda( + q, + kn, + v, + cumulative_gate, + beta, + initial_state, + cu_seqlens=cu_seqlens, + output_final_state=True, + fuse_q_l2norm=True, + ) + # The fused path defers the row scale past the bf16 gram rounding, so + # outputs agree to bf16 resolution; the state path never touches q's norm. + torch.testing.assert_close(actual, expected, atol=2e-3, rtol=2e-2) + torch.testing.assert_close(actual_state, expected_state, atol=0.0, rtol=0.0) + + +def test_fuse_q_l2norm_rejects_compile() -> None: + q, k, v, cumulative_gate, beta, _, _ = _inputs([256], heads=2) + kn = l2norm(k) + + compiled = torch.compile( + lambda: chunk_kda(q, kn, v, cumulative_gate, beta, fuse_q_l2norm=True), + fullgraph=True, + ) + # NotImplementedError and dynamo's trace-time wrapper both derive from + # RuntimeError; only the eager-only contract message matters. + with pytest.raises(RuntimeError, match="eager-only"): + compiled() + + +@pytest.mark.parametrize("operand", _GRAD_OPERANDS) +def test_fuse_q_l2norm_rejects_gradients(operand: str) -> None: + q, k, v, cumulative_gate, beta, _, initial_state = _inputs([256], heads=2) + tensors = { + "q": q, + "k": l2norm(k), + "v": v, + "cumulative_gate": cumulative_gate, + "beta": beta, + "initial_state": initial_state, + } + tensors[operand] = tensors[operand].clone().requires_grad_(True) + + def run(): + return chunk_kda( + tensors["q"], + tensors["k"], + tensors["v"], + tensors["cumulative_gate"], + tensors["beta"], + tensors["initial_state"], + fuse_q_l2norm=True, + ) + + with pytest.raises(RuntimeError, match="forward-only"): + run() + with torch.no_grad(): + output, _ = run() + assert not output.requires_grad diff --git a/test/test_kda_impl_dispatch.py b/test/test_kda_impl_dispatch.py index 952b1744..1ca34f12 100644 --- a/test/test_kda_impl_dispatch.py +++ b/test/test_kda_impl_dispatch.py @@ -204,6 +204,13 @@ def test_chunk_reference_rejects_fastmath(): chunk_kda(q, k, v, chunk_cumsum_ref(gate, 64), beta, fastmath=True, impl="reference") +def test_chunk_reference_rejects_fuse_q_l2norm(): + """Keep fused-only knobs from silently changing meaning.""" + q, k, v, gate, beta = _inputs(tokens=8) + with pytest.raises(ValueError, match="fuse_q_l2norm"): + chunk_kda(q, k, v, chunk_cumsum_ref(gate, 64), beta, fuse_q_l2norm=True, impl="reference") + + @pytest.mark.skipif(not BLACKWELL, reason="fused chunk_kda requires CUDA capability 10.0") def test_chunk_autotune_flag_is_deterministic_and_accurate(): """autotune=False pins fixed heuristic configs without changing the math contract.""" diff --git a/test/test_kda_ragged_autograd_ops.py b/test/test_kda_ragged_autograd_ops.py index 241c61ea..66f67659 100644 --- a/test/test_kda_ragged_autograd_ops.py +++ b/test/test_kda_ragged_autograd_ops.py @@ -33,6 +33,7 @@ def test_ragged_custom_op_registrations(): cu_seqlens, metadata.chunk_offsets, True, + False, ) torch.library.opcheck( _chunk_kda_fwd_ragged_with_state_op, @@ -77,6 +78,7 @@ def test_ragged_custom_op_registrations(): cu_seqlens, metadata.chunk_offsets, True, + False, ) with torch.no_grad(): output, Aqk, Akk = _chunk_kda_fwd_ragged_op(*no_state_args)