Add forward-only fuse_q_l2norm to chunk_kda by drisspg · Pull Request #345 · meta-pytorch/attention-gym

drisspg · 2026-08-18T20:08:40Z

Add forward-only fuse_q_l2norm to chunk_kda

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.

pytest test/test_kda_fuse_q_l2norm.py test/test_kda*.py
python agent_space/bench_three_way.py

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

chatgpt-codex-connector · 2026-08-18T20:12:24Z

💡 Codex Review

attention-gym/attn_gym/linear/kda/fwd/triton/chunk_delta_h.py

Lines 422 to 423 in 81db694

    
           h, v_new, final_state = torch.ops.attn_gym.kda_delta_h_with_state( 
        
               k, w, u, gk, initial_state, cu_seqlens, chunk_offsets, chunks

Avoid nesting the delta-H operator in the chunk operator

For ordinary public chunk_kda calls, _ChunkKDA.forward has already entered one of the compiler-opaque kda_chunk_fwd* operators before execution reaches this call, so dispatching through kda_delta_h* adds a second registered-operator boundary to every hot forward pass. This directly offsets part of the recurrence's latency improvement; call the shared _delta_h_launch from the already-opaque chunk implementation and retain the registered delta-H operator only for standalone compiled callers.

AGENTS.md reference: AGENTS.md:L78-L84

ℹ️ 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".

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

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8aaba8e214

ℹ️ 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".

chatgpt-codex-connector · 2026-08-18T20:22:56Z

+        output, state, _, _ = _chunk_kda_fwd(
+            q,


Restrict fused normalization to BF16 queries

When q is an accepted FP16/FP32 input, this branch receives the tensor only after line 192 has converted the raw values to BF16, whereas the advertised explicit path computes l2norm(q) in the original dtype and converts the normalized result afterward. Inputs whose small components are rounded or underflow during the early BF16 cast can therefore produce a different direction—or an all-zero output—rather than matching the skipped normalization pass. Either reject non-BF16 q when this flag is enabled or preserve the original-precision normalization semantics.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-08-18T20:22:56Z

+        cu_seqlens = torch.tensor(
+            [0, *torch.tensor(seq_lens).cumsum(0).tolist()], device=device, dtype=torch.int32
+        )


Reuse the shared packed-offset builder

Replace this local cu_seqlens construction with attn_gym.testing.kda.cumulative_sequence_offsets(seq_lens). The scoped test guidance explicitly requires checking and reusing that helper for packed offsets, so keeping a second builder here duplicates the canonical dtype/device and boundary conventions.

AGENTS.md reference: test/AGENTS.md:L5-L10

Useful? React with 👍 / 👎.

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

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a10335be2f

ℹ️ 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".

chatgpt-codex-connector · 2026-08-18T20:28:01Z

+    )
+    # 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)


Calibrate the fused-output tolerance against a reference

Replace this fixed atol/rtol assertion with assert_matches_low_precision_reference, using the reference implementation as the high-precision result and the explicit-normalization fused run as the low-precision result. With the current random fixture, the absolute 2e-3 allowance can hide errors in small expected values and is not tied to the BF16 rounding error that the comment intends to permit; the scoped test convention specifically requires data-derived error budgets rather than hand-tuned tolerances.

AGENTS.md reference: test/AGENTS.md:L5-L13

Useful? React with 👍 / 👎.

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

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 38f56140c1

ℹ️ 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".

chatgpt-codex-connector · 2026-08-18T21:48:18Z

+    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]


Prevent raw-query overflow before applying the fused norm

For BF16 rows with large but finite values, the raw-q dot products can overflow before this final scale is applied, producing inf * 0 = NaN. For example, an all-3e38 query row and positive normalized keys make Aqk overflow, while b_ssq is also infinite; the explicit l2norm(q) path instead reduces that row to zero and returns a finite result. The bounded-gate requirement does not prevent this query-magnitude overflow, so the fused path needs overflow-safe scaling or an explicit magnitude restriction.

Useful? React with 👍 / 👎.

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 · 2026-08-19T13:50:17Z

Dropping this: a dedicated forward-only prefill path is planned and the flag belongs there rather than on the training API. Work preserved on branch archive/kda-fuse-q-l2norm.

drisspg force-pushed the drisspg/stack/83 branch from a120ff9 to 3f3f755 Compare August 18, 2026 20:08

drisspg force-pushed the drisspg/stack/84 branch from 81db694 to 3e0d7a3 Compare August 18, 2026 20:08

This was referenced Aug 18, 2026

Cap ragged KDA output-kernel registers at the dense budget #343

Merged

Replace the KDA inter-chunk recurrence with a warp-specialized kernel #344

Merged

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 18, 2026

drisspg marked this pull request as draft August 18, 2026 20:20

drisspg changed the base branch from drisspg/stack/83 to main August 18, 2026 20:20

drisspg force-pushed the drisspg/stack/84 branch from 3e0d7a3 to 8aaba8e Compare August 18, 2026 20:20

drisspg changed the base branch from main to drisspg/stack/83 August 18, 2026 20:20

drisspg marked this pull request as ready for review August 18, 2026 20:20

chatgpt-codex-connector Bot reviewed Aug 18, 2026

View reviewed changes

drisspg marked this pull request as draft August 18, 2026 20:25

drisspg changed the base branch from drisspg/stack/83 to main August 18, 2026 20:25

drisspg force-pushed the drisspg/stack/84 branch from 8aaba8e to a10335b Compare August 18, 2026 20:25

drisspg changed the base branch from main to drisspg/stack/83 August 18, 2026 20:25

drisspg marked this pull request as ready for review August 18, 2026 20:25

chatgpt-codex-connector Bot reviewed Aug 18, 2026

View reviewed changes

drisspg marked this pull request as draft August 18, 2026 21:40

drisspg changed the base branch from drisspg/stack/83 to main August 18, 2026 21:40

drisspg force-pushed the drisspg/stack/84 branch from a10335b to 2871c81 Compare August 18, 2026 21:40

drisspg changed the base branch from main to drisspg/stack/83 August 18, 2026 21:40

drisspg marked this pull request as ready for review August 18, 2026 21:40

drisspg marked this pull request as draft August 18, 2026 21:45

drisspg changed the base branch from drisspg/stack/83 to main August 18, 2026 21:45

drisspg force-pushed the drisspg/stack/84 branch from 2871c81 to 38f5614 Compare August 18, 2026 21:45

drisspg changed the base branch from main to drisspg/stack/83 August 18, 2026 21:45

drisspg marked this pull request as ready for review August 18, 2026 21:45

chatgpt-codex-connector Bot reviewed Aug 18, 2026

View reviewed changes

drisspg marked this pull request as draft August 18, 2026 22:31

drisspg changed the base branch from drisspg/stack/83 to main August 18, 2026 22:31

drisspg force-pushed the drisspg/stack/84 branch from 38f5614 to afcbcee Compare August 18, 2026 22:31

drisspg changed the base branch from main to drisspg/stack/83 August 18, 2026 22:31

drisspg marked this pull request as ready for review August 18, 2026 22:31

drisspg marked this pull request as draft August 19, 2026 03:27

drisspg changed the base branch from drisspg/stack/83 to main August 19, 2026 03:27

drisspg force-pushed the drisspg/stack/84 branch from afcbcee to cd6d6ff Compare August 19, 2026 03:27

drisspg changed the base branch from main to drisspg/stack/83 August 19, 2026 03:27

drisspg marked this pull request as ready for review August 19, 2026 03:27

drisspg marked this pull request as draft August 19, 2026 13:42

drisspg changed the base branch from drisspg/stack/83 to main August 19, 2026 13:43

drisspg force-pushed the drisspg/stack/84 branch from cd6d6ff to ede717d Compare August 19, 2026 13:43

drisspg marked this pull request as ready for review August 19, 2026 13:43

drisspg closed this Aug 19, 2026

Navigation Menu

Uh oh!

Add forward-only fuse_q_l2norm to chunk_kda - #345

Add forward-only fuse_q_l2norm to chunk_kda#345
drisspg wants to merge 1 commit into
mainfrom
drisspg/stack/84

drisspg commented Aug 18, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot commented Aug 18, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Uh oh!

drisspg commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

drisspg commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chatgpt-codex-connector Bot commented Aug 18, 2026

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

drisspg commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

drisspg commented Aug 18, 2026 •

edited

Loading