fix(gdn): use block-end decay for SM100 state updates

[guangyunh-nv]

Summary

• restore the block-end cumulative decay when updating the recurrent SM100 GDN state
• remove the thread-local coordinate materialization used to select that scalar
• add a shared regression test covering full, partial, and recurrent blocks on all supported GDN architectures

Root cause

PR #4133 replaced the fixed block-end lookup with an index derived from the last coordinate in each thread's partitioned accumulator fragment:

max_coord = tTR_tCcShared[cute.size(tTR_tCcShared) - 1]
cumprod_total = sCumprod[max_coord[1], 0, gate_handle.index]

That max_coord coordinate is thread-dependent. Its second coordinate is:

tid % 4 == 0: coord[1] = 57
tid % 4 == 1: coord[1] = 59
tid % 4 == 2: coord[1] = 61
tid % 4 == 3: coord[1] = 63

Representative output:

tid=0   coord=(24, 57)
tid=1   coord=(24, 59)
tid=2   coord=(24, 61)
tid=3   coord=(24, 63)
...
tid=124 coord=(127, 57)
tid=125 coord=(127, 59)
tid=126 coord=(127, 61)
tid=127 coord=(127, 63)

The fix directly loads the block-end scalar:

cumprod_total = sCumprod[self.b_t - 1, 0, gate_handle.index]

This remains correct for a partial block. The gate loader predicates out-of-bounds elements after initializing them to the multiplicative identity 1.0. Their log2 contribution is therefore zero, so the inclusive prefix product remains constant after the final valid token and slot BT - 1 contains the product over exactly the valid block.

Reproduction

The regression test uses FP16, sequence lengths [64, 111, 192], one head, normalized K, and alpha/beta sampled from [0.99, 1.0). These lengths cover one full block, a non-multiple tail with recurrent carry, and three aligned blocks.

Revision	Output max error	Output mismatches	State max error	State mismatches
Before #4133 (6258e522)	5.80e-4	0 / 46976	2.28e-4	0 / 49152
After #4133 (f057e15b)	1.59e-2	1660 / 46976	1.01e-2	6361 / 49152
This fix	5.80e-4	0 / 46976	2.28e-4	0 / 49152

Validation

PYTHONPATH=. python -m pytest -q \
  tests/gdn/test_prefill_delta_rule.py::test_prefill_block_end_decay

Result: 1 passed on SM100 with CUDA 13. The test now collects on SM90, SM100, and SM12x so architecture CI covers the shared contract.

Summary by CodeRabbit

• Bug Fixes

  • Improved gated delta network prefill behavior for partial chunks and block-boundary sequence lengths.
  • Ensured cumulative decay uses neutral padding values when processing incomplete chunks.

• Tests

  • Added GPU coverage comparing prefill results and final states against the blockwise reference implementation, including decay factors near one.

[coderabbitai]

📝 Walkthrough

Walkthrough

The Blackwell chunked kernel now reads cumulative decay from the final physical chunk slot. A regression test covers prefill sequences that end at and across block boundaries, comparing outputs and final states with the blockwise reference.

Changes

Chunked decay indexing

Layer / File(s)	Summary
Final-slot decay indexing and regression coverage flashinfer/gdn_kernels/blackwell/gated_delta_net_chunked.py, tests/gdn/test_prefill_delta_rule.py	The kernel removes coordinate-derived indexing and reads the final physical chunk slot. The regression test validates outputs and transposed final states for multi-sequence prefill inputs.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: bkryu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the primary fix to block-end decay handling for SM100 GDN state updates.
Description check	✅ Passed	The description explains the fix, root cause, regression coverage, and validation results, but omits template checklist items and related issues.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix_sm100_gdn_block_end_decay

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[cloudhan]

Other archs should also be checked. Relax it.

[jiahanc]

/bot run tests/gdn

[jiahanc]

Thanks for the fix!

[flashinfer-bot]

GitLab MR !1104 has been created, and the CI pipeline #60794041 is currently running. I'll report back once the pipeline job completes.