Add fused recurrent_kda for decode and inference prefill - #317
Merged
Conversation
drisspg
added a commit
that referenced
this pull request
Aug 16, 2026
## Human Note ## Agent note 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`. ## Test Plan ```bash gpu-run auto -- ~/.venvs/ag-linear/bin/python -m pytest test/test_kda_recurrent.py -q # 14 passed on B200: dense/packed vs naive_recurrent_kda (fp32 exact-tolerance, bf16), # non-power-of-two K/V masking, empty-slot state passthrough, capacity tails, # cross-check vs chunk_kda via bounded_gate_cumsum(chunk_size=1 vs 64), gradient # rejection, opcheck, fullgraph compile, and CUDA-graph replay with mutated # boundaries and history. ``` stack-info: PR: #317, branch: drisspg/stack/79
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 16, 2026 18:02
73b0e0e to
2953ab4
Compare
meta-cla
Bot
added
the
CLA Signed
drisspg
marked this pull request as draft
August 16, 2026 18:54
drisspg
marked this pull request as ready for review
August 16, 2026 18:54
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2953ab441e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+52
to
+53
| bos = tl.load(cu_seqlens + i_n).to(tl.int64) | ||
| eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) |
There was a problem hiding this comment.
When cu_seqlens contains a negative, decreasing, nonzero initial, or over-capacity boundary, these unchecked values drive row outside the Q/K/V/output allocations, potentially causing an illegal CUDA memory access or corrupting adjacent output memory. The public validation checks only the tensor's metadata, whereas the existing ragged chunk scheduler device-validates 0 <= begin <= end <= tokens and that the first offset is zero; enforce the same invariants before this scan, including during CUDA Graph replay.
Useful? React with 👍 / 👎.
| output = torch.empty_like(v, dtype=q.dtype) | ||
| final_state = q.new_empty(num_sequences, heads, key_dim, value_dim, dtype=torch.float32) | ||
| block_v = min(triton.next_power_of_2(value_dim), 64) | ||
| grid = (triton.cdiv(value_dim, block_v), num_sequences * heads) |
There was a problem hiding this comment.
For otherwise valid dense or packed inputs where num_sequences * heads > 65,535 (for example, 4,096 packed sequences with 32 heads), this places the product in CUDA's grid-Y dimension and the launch fails because grid Y is limited to 65,535. No validation imposes a corresponding limit, so flatten the work into grid X or otherwise split the sequence and head dimensions without exceeding the grid-Y bound.
Useful? React with 👍 / 👎.
drisspg
marked this pull request as draft
August 16, 2026 21:39
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 16, 2026 21:39
2953ab4 to
1fc481c
Compare
drisspg
marked this pull request as ready for review
August 16, 2026 21:39
drisspg
marked this pull request as draft
August 16, 2026 23:06
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 16, 2026 23:06
1fc481c to
a0aec00
Compare
drisspg
marked this pull request as ready for review
August 16, 2026 23:07
drisspg
marked this pull request as draft
August 16, 2026 23:37
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 16, 2026 23:37
a0aec00 to
35d2e42
Compare
drisspg
marked this pull request as ready for review
August 16, 2026 23:37
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 35d2e42cda
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 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() | ||
| output, final_state = _recurrent_fwd_op(q, k, v, gate, beta, initial_state, cu_seqlens) |
There was a problem hiding this comment.
When output_final_state=False (the default), this still invokes the two-output operator, whose launcher allocates a num_sequences * heads * K * V FP32 tensor and whose kernel stores the entire final state, only for the next line to discard it. For batched inference with large head dimensions this can consume hundreds of megabytes or cause an OOM even though the caller requested only the token output; use a fixed-arity no-state operator/launcher variant that compiles out the allocation and ht store, as the existing chunk_kda path does.
AGENTS.md reference: AGENTS.md:L69-L73
Useful? React with 👍 / 👎.
drisspg
marked this pull request as draft
August 16, 2026 23:48
drisspg
marked this pull request as ready for review
August 16, 2026 23:49
drisspg
marked this pull request as draft
August 17, 2026 00:40
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 17, 2026 00:40
c5f58b2 to
f063927
Compare
drisspg
marked this pull request as ready for review
August 17, 2026 00:40
drisspg
marked this pull request as draft
August 17, 2026 01:06
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 17, 2026 01:06
f063927 to
3826b81
Compare
drisspg
marked this pull request as ready for review
August 17, 2026 01:06
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3826b81e91
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| Shapes: q, k, g (B, T, H, K); v (B, T, H, V); beta (B, T, H); state (B, H, K, V), | ||
| or (N, H, K, V) over the N documents of a packed row in varlen mode. | ||
| Shapes: q, k, g (B, T, H, K); v (B, T, H, V); beta (B, T, H); state (B, H, K, V). |
There was a problem hiding this comment.
Calls to naive_recurrent_kda(..., cu_seqlens=offsets) now raise an unexpected-keyword TypeError because this commit removes the public argument and segmented execution that were added in the parent history. Packed callers therefore lose the reference implementation needed to validate the newly introduced fused packed path; retain the cu_seqlens parameter and its per-sequence recurrence behavior.
Useful? React with 👍 / 👎.
drisspg
marked this pull request as draft
August 17, 2026 01:31
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 17, 2026 01:32
3826b81 to
37c1048
Compare
drisspg
marked this pull request as ready for review
August 17, 2026 01:32
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 37c10488ee
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+58
to
+61
| tl.device_assert( | ||
| (bos >= 0) & (bos <= eos) & (eos <= T) & ((i_n != 0) | (bos == 0)), | ||
| "invalid packed cu_seqlens", | ||
| ) |
There was a problem hiding this comment.
The earlier offset-validation issue remains: the fresh evidence is that the added guard uses tl.device_assert, which Triton disables unless TRITON_DEBUG=1. In normal production runs, negative, decreasing, nonzero-initial, or over-capacity offsets can therefore still make the loop access outside the Q/K/V/output buffers; validate these values through an always-enabled device-side path that remains compatible with CUDA Graph replay.
Useful? React with 👍 / 👎.
Comment on lines
+55
to
+59
| 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 |
There was a problem hiding this comment.
When packed execution uses output_final_state=False, this branch still allocates an FP32 [num_documents, H, K, V] tensor—or clones the entire supplied initial state—and then fills it document by document before discarding it at return. For many packed documents or large head dimensions, the reference call can now consume gigabytes or OOM even though the previous implementation retained only each document's transient recurrence state; allocate and accumulate this tensor only when the caller requests it.
Useful? React with 👍 / 👎.
drisspg
marked this pull request as draft
August 17, 2026 01:40
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 17, 2026 01:40
37c1048 to
4c5cbb1
Compare
drisspg
marked this pull request as ready for review
August 17, 2026 01:40
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: #317, branch: drisspg/stack/79
drisspg
marked this pull request as draft
August 17, 2026 01:42
drisspg
force-pushed
the
drisspg/stack/79
branch
from
August 17, 2026 01:42
4c5cbb1 to
7b55be0
Compare
drisspg
marked this pull request as ready for review
August 17, 2026 01:42
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b55be079e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| output, final_state = naive_recurrent_kda( | ||
| offsets = cu_seqlens.tolist() | ||
| num_documents = len(offsets) - 1 | ||
| compute_dtype = torch.promote_types(q.dtype, torch.float32) |
There was a problem hiding this comment.
When q is FP32 but another operand or initial_state is FP64, each recursive document recurrence promotes its state to FP64, but this buffer is derived from q alone, so assigning doc_state silently truncates the returned packed final states to FP32; empty documents also immediately lose precision through the FP32 clone. The previous packed implementation preserved the recursively promoted dtype when concatenating states, so compute this dtype from k, v, g, beta, and the optional initial state just as the dense branch does.
Useful? React with 👍 / 👎.
drisspg
added a commit
that referenced
this pull request
Aug 18, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
drisspg
added a commit
that referenced
this pull request
Aug 18, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
drisspg
added a commit
that referenced
this pull request
Aug 18, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
drisspg
added a commit
that referenced
this pull request
Aug 18, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
drisspg
added a commit
that referenced
this pull request
Aug 18, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
drisspg
added a commit
that referenced
this pull request
Aug 18, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
drisspg
added a commit
that referenced
this pull request
Aug 19, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
drisspg
added a commit
that referenced
this pull request
Aug 19, 2026
Inference callers currently normalize q with a standalone l2norm pass before chunk_kda, which reads and writes the full [T, H, 128] tensor just to apply a per-row scalar. Every q-dependent term in the forward is linear in q's row and lands in the same output row, so the norm can be deferred: the grams stay in the raw-q basis (intra and K3b unchanged), the output kernel accumulates each row's sum of squares inside its existing k-tile loop, and one 1/||q|| row scale on the accumulated output completes the normalization. An earlier variant that emitted an rstd tensor from the intra kernel cost +154us from register pressure in that 1-warp kernel and was rejected. The fused ragged output launch needs a slightly larger register cap (152 vs 136) for the extra accumulator. The flag is opt-in and forward-only: it raises when any input requires grad (the saved Aqk stays in the raw-q basis, so the existing backward would be wrong) and when compiling (the fused route bypasses the compiler-opaque autograd op). This mirrors the precedent elsewhere: FLA's use_qk_l2norm_in_kernel supports backward only because it launches the standalone l2norm kernels inside the op rather than fusing them, while every true fusion (Helion varlen, FlashInfer decode) is forward-only. k keeps its explicit l2norm; this complements the decode-side preprocessing fusion in the recurrent path (#317, #334) on the chunked prefill side. Raw-q grams run about sqrt(K) hotter, so the flag assumes the bounded-gate contract; the eps floor matches l2norm's default and the tests cover zero rows and 1e-3..1e3 row scales against the explicit path. GB300, six frozen contract shapes, output-only forward, warm CUDA events: removes the 38-60us standalone q pass per call (l2norm kernel eliminated; +5us sumsq cost inside the o kernel). This flipped the h96/h64 mixed-shape comparisons against Helion's pretuned linear-attention kernels; final six-shape geomean time ratio ours/helion = 0.952. ```bash pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py python agent_space/bench_three_way.py ``` stack-info: PR: #345, branch: drisspg/stack/84
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked PRs:
Add fused recurrent_kda for decode and inference prefill
Human Note
Agent note
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)mirrorschunk_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:
gateis the per-token log2 decay (bounded_gate_cumsumwithchunk_size=1), not the chunk-local cumulative gate. Queries scale by 1/sqrt(K) internally,matching
chunk_kdaand thenaive_recurrent_kdadefault.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_kdarather thansilently detaching. The op uses the repo's define/impl + fake registration with a fixed-arity
(output, state) schema, and joins the lazy
KDA_OPSsurface inattn_gym.linear.Test Plan