[Fix] Fold chunk state launch grids to fit the 65535-block limit

[xy200303]

Summary

Avoid invalid CUDA launches when packing many short sequences. Several chunk-state kernels put N * H on grid axis z, and vector global cumsum puts it on axis y. Both CUDA axes are limited to 65535 blocks, where N = len(cu_seqlens) - 1 is the number of packed sequences.

For example, 2050 sequences of 16 tokens with 32 heads require 65600 blocks and fail at launch with only 32800 input tokens. The trigger is the number of sequences times heads, independent of dtype; a large token count alone does not trigger it.

Flatten these grids onto CUDA's x axis and recover the logical 2D/3D indices through a shared unflatten_program_id(X, Y) Triton helper. This covers the shared chunk-state kernels used by simple GLA, GLA, GSA, RWKV6 and MesaNet, the separate IPLR/DPLR state kernels, and vector global cumsum used by normalized linear attention and other callers. Per-block computations, dtypes and public APIs are unchanged.

Test plan

• Add one FP16 extreme case to the existing simple-GLA varlen test: 2050 sequences × 32 heads, comparing outputs, final states and all gradients with the fused recurrent reference at the existing tolerances. Add a vector-cumsum case with the same sequence/head count and two feature tiles.
• Local NVIDIA GB200, PyTorch 2.11.0.post1+cu132.msh, strict FLA_CI_ENV=0, default Triton implementations and normal NaN memory poisoning: 259 selected tests passed across simple GLA (32), GLA (30), linear attention (54), RWKV6 (13), IPLR (14), DPLR (35), cumsum (33), GSA (33) and MesaNet (15). All listed test files pass in full except simple GLA, which selects -k 'test_chunk or test_fused_recurrent'. DPLR uses Triton 3.6.0; the other suites use 3.5.1. MesaNet's reference required loading the image's BLAS/LAPACK libraries.
• The broader simple-GLA suite exposes an existing fused_chunk failure on this environment. Independent base/head runs load the correct worktrees, have identical fused_chunk source hashes, and produce the same o diff=60.875, ratio=0.999282. This path does not use the changed helper. Triton 3.5.1 also fails to compile an unchanged DPLR wy_fast_bwd.py kernel at D=100 on both base and head; the DPLR suite passes with 3.6.0.
• Pre-commit, header and diff checks pass. Extracting the shared helper emits identical PTX instructions to the inline flattened mapping for all seven affected kernels, in dense and varlen modes on sm90 (14 comparisons).

Benchmark / NCU (kernel changes only)

Each before/after pair runs on the same GB200, comparing e27e1aaf → 8e8024fc. The grid implementation is unchanged in final head 3c94fc43; linear attention was revalidated and remeasured against its updated base (d07d7ad8 → 3c94fc43). PyTorch 2.11.0.post1+cu132.msh / Triton 3.5.1, BF16, H=8, D=64. Dense: B=2, T=2048. Varlen: B=1, T=4096, 32 sequences of 128 tokens. Values are median before → after, µs.

Registry-generated inputs and benchmarks.ops.run.benchmark_op, with 6 warmup iterations, 40 ms timing warmup and 200 ms repetitions (100/500 ms for the final linear-attention rerun):

Op	Dense fwd	Dense fwd+bwd	Varlen fwd	Varlen fwd+bwd
simple_gla	213.6 → 214.4	915.9 → 895.8	293.0 → 285.5	1061.5 → 1104.4
gla	332.1 → 319.3	1096.0 → 1056.0	421.8 → 399.2	1316.2 → 1272.8
linear_attn	336.2 → 356.1	1127.8 → 1112.6	542.6 → 536.7	1564.9 → 1536.3
rwkv6	314.6 → 311.2	1095.2 → 1062.7	395.3 → 374.0	1258.3 → 1246.9

IPLR and DPLR additionally use triton.testing.do_bench_cudagraph with 500 ms repetitions to isolate GPU execution. Eager DPLR timings varied substantially between repeated processes, so those timings do not support a stable regression conclusion. IPLR has no implemented backward API.

Op	Dense fwd	Dense fwd+bwd	Varlen fwd	Varlen fwd+bwd
iplr_delta_rule	112.8 → 114.3	N/A	49.3 → 47.8	N/A
dplr_delta_rule	348.5 → 343.5	1062.8 → 1050.5	105.3 → 105.3	419.6 → 416.5

NCU --set full / --set source on the shared state-forward kernel: duration 12.99 → 12.99 µs, 32 → 32 registers/thread, memory throughput 646.21 → 646.60 GB/s. The profiled state kernel has unchanged latency; IPLR/DPLR graph timings vary by at most 3% on these workloads.

Breaking changes

None.

Checklist

• I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
• I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
• Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
• Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
• This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.

[xy200303]

The failing smoke test is pre-existing on main and unrelated to this PR.

tests/layers/test_based.py::test_forward[B1-T1-H2-K16-V16-dtype1-chunk] already fails the
same way on main before this branch existed. The scheduled nvidia-h100-ci run on
cfaac246 (the current main tip) reports the identical error:

FAILED tests/layers/test_based.py::test_forward[B1-T1-H2-K16-V16-dtype1-chunk]
  - torch.AcceleratorError: CUDA error: an illegal memory access was encountered

https://github.com/fla-org/flash-linear-attention/actions/runs/34407286027

The dependent-test list is run in order and stops at the first failing file, so the smoke job
never reached the tests for the files changed here (test_simple_gla.py, test_gla.py,
test_linear_attn.py, test_rwkv6.py, test_iplr_delta.py, test_dplr_delta.py,
tests/ops/utils/test_cumsum.py all appear later in the same list). The new
tests/ops/test_simple_gla.py::test_chunk_varlen_many_documents did not get a chance to run
in this job.

For reference, the changed paths were verified locally on RTX 5090 D (sm_120, torch 2.8.0+cu128,
triton 3.4.0) — the six affected ops launch cleanly at N*H = 65600 after the change, output is
bitwise identical to before with the config pinned, and test_cumsum / test_iplr_delta /
test_rwkv6 / test_linear_attn pass (89 tests, 0 failures). sm_120 is not in the CI matrix, so
that evidence is local only.

Update: the base branch has since been updated and now includes #1241
([Fix] Reject odd K/V for 16-bit dtypes in chunk_bwd_dv), which lands on the same path this
test exercises — K=16 expands to an odd 153-wide feature map in fp16, and the backward goes
through chunk_bwd_dv. CI has been re-triggered on the updated head.

[xy200303]

pytest tests/ops/test_simple_gla.py tests/ops/test_linear_attn.py tests/ops/utils/test_cumsum.py -q

[xy200303]

The Smoke job is red from the 2-hour job cap, not from a test failure.

gh api .../actions/jobs/103132039378/logs shows no FAILED lines at all — every test that ran
passed. The job walks the dependent-test list alphabetically and was cancelled at
The job has exceeded the maximum execution time of 2h0m0s right after
tests/ops/test_rwkv6.py (one of the files this PR changes — it passed).

Touching fla/ops/common/chunk_h.py and fla/ops/utils/cumsum.py makes
scripts/find_dependent_tests.py expand to ~50 GPU test files, and the alphabetical tail of that
list — test_rwkv7.py, test_simple_gla.py, test_wall_attn.py, tests/ops/utils/test_cumsum.py —
never gets a chance to run inside the 2h budget. In particular the new
tests/ops/test_simple_gla.py::test_chunk_varlen_many_documents did not run here.

Two things worth noting for whoever picks this up:

• The tests that did run are green, including tests/ops/test_rwkv6.py and
  tests/ops/test_linear_attn.py, both of which exercise kernels this PR touches.
• Earlier, the same file tests/layers/test_based.py::test_forward[B1-T1-H2-K16-V16-dtype1-chunk]
  was failing here; that failure is gone from this run's list, so it is no longer masking anything.

I could not get H100 coverage for the changed files through the pytest <args> comment route —
both Issue / PR pytest command runs I triggered came back skipped. If there is a preferred way
to run a targeted subset (or if you would rather this PR not touch cumsum.py to keep the
dependent-test set smaller), let me know and I will adjust.

[xy200303]

For anyone looking at the red Smoke (PyTorch 2.12) check: this run is a budget timeout, not a test failure or a hang.

The run I inspected (job 103132039378, run 34557126450, 2026-09-11) sat at the job's timeout-minutes: 120 and was cancelled. Reconstructing the log:

• find_dependent_tests.py selected 48 test files for this PR's changed files.
• The job runs each file in its own pytest process. Summing the 48 invocations gives 119.6 min — essentially the entire 120-minute budget.
• There is no single long hang: the largest gap between consecutive log lines is 7.1 min, and 123 gaps exceed 5 s. The budget is consumed by breadth, spread across many files (slowest single invocation ≈ 20 min).

The reason the fanout is this wide is that fla/ops/common/chunk_h.py is imported by essentially every chunk-state kernel, so touching it pulls in nearly the whole ops suite. That is a property of the module, not of this PR.

Two other data points that point the same way:

• The same branch's earlier smoke runs (34468307083, 34467553463, 09-10) failed in 1.6 min and 6.7 min respectively — the check is not deterministically red on this change.
• Recent runs on other branches show the same squeeze: 34745550725 on fix/l2norm-autotune-key ran 81.2 min before failing, against the same 120-minute ceiling.

Net: the code in this PR has not been shown to fail any test; the smoke tier just no longer fits the suite it now selects. Closing this out needs a main-side change — raising timeout-minutes for the smoke tier, or narrowing what it selects for core shared modules — rather than a change in this branch. I've left the fix as-is and not pushed anything to work around the timeout.