[Triton][GDN] Add in-place state scatter + h output to VK chunk

[hsthe29]

Summary

chunk_gated_delta_rule_opt_vk already computes the gated-delta-rule chunk with a V-major [N, H, V, K] hidden state and GVA (Hg k-heads / H v-heads) — but it could not back a paged / radix recurrent-state pool, for two reasons:

1. It took a dense, per-sequence initial_state ([N, H, V, K]) and produced a separate dense final_state. A pool keyed by per-request slot indices could not be read/written in place.
2. It discarded the per-chunk hidden-state snapshots h (produced internally by the VK h-scan), which a state-tracking cache needs to snapshot SSM state at chunk boundaries.

Both are exactly what a serving framework's linear-attention state cache needs (e.g. SGLang's MambaPool + radix/extra-buffer tracking, whose own sgl-fla VK h-scan does precisely this).

Change

Two new optional args; defaults preserve existing behavior byte-for-byte:

• initial_state_indices — in-place scatter. When given, the h-scan reads the initial state from, and writes the final state back to, the same pool at slot initial_state_indices[i_n] ((index * H + i_h) * V * K) rather than a dense [N, …] buffer. Only the Triton VK path supports it (use_chunk_hip / use_chunk_flydsl raise a clear error). Mirrors the sgl-fla VK h-scan (h0 == ht == pool + index*stride).
• return_h — expose the per-chunk snapshots h [B, NT, H, V, K] (already computed) for chunk-boundary state tracking.

Threaded through: kernel chunk_gated_delta_rule_fwd_kernel_h_opt_vk (+initial_state_indices, +USE_STATE_INDICES constexpr) → wrapper chunk_gated_delta_rule_fwd_h_opt_vk → chunk_gated_delta_rule_fwd_opt_vk (+return_h) → public chunk_gated_delta_rule_opt_vk.

Correctness

New op_test test_chunk_opt_vk_state_indices: the in-place-scatter path equals the trusted dense path for the same logical state, and return_h has the documented shape. Coverage: GVA (H=3·Hg) + non-GVA, K==V + K!=V, fp16 + bf16, varlen with aligned and unaligned chunk lengths, non-identity slot permutation.

check	cos	max abs err
o : indices(in-place) vs dense	1.000000	0.0 (bit-exact)
final state : indices(in-place) vs dense	1.000000	0.0 (bit-exact)
h shape	[B, NT, H, V, K]	✓

• New op_test: 5/5 pass. Existing test_chunk_opt_vk / _varlen: 20/20 pass (dense path unchanged). black / ruff clean.
• Drop-in cross-check (external, mamba_v2_integration/task_5/test_aiter_chunk_vs_sglfla.py): vs SGLang's own sgl-fla chunk_gated_delta_rule on the same [V,K] pool + indices, o / h / in-place final state all cos = 1.000000 (max abs err at bf16/fp16 noise, ~1e-4…4e-3), same h shape [1, NT, H, V, K].

Benchmark (kernel, MI350X gfx950, GPU idle, min-of-N)

Qwen3.6-like GDN shapes (Hg=16, H=48 GVA×3, K=V=128, bf16), on the same [V,K] pool + indices. aiter_triton_idx = this PR's Triton VK path (indices + return_h); sglfla_idx = SGLang's own sgl-fla chunk (baseline it runs today). Reproducible over 2 runs.

seqlens	aiter_triton_idx µs	aiter_hip µs	sglfla_idx µs	aiter/sglfla
[2048]	210.8	228.0	321.7	0.66×
[1024,1024]	204.7	205.2	297.0	0.69×
[512×4]	219.3	198.2	306.2	0.72×
[200,64,130,777] (ragged)	165.6	164.0	259.0	0.64×
[4096]	399.9	414.4	520.9	0.77×

• aiter's opt_vk chunk is ~1.30–1.63× faster than sgl-fla at the kernel level.
• The PR's initial_state_indices + return_h add no overhead (aiter_triton_idx ≈ the dense aiter_hip), so the in-place scatter path is as fast as aiter's fastest dense path.

Notes

• HIP / FlyDSL h-scan variants don't yet support initial_state_indices (raise). A follow-up can add it there.
• Together with the decode PR, this lets a V-major, paged recurrent-state pool use aiter's GDN kernels end-to-end (extend + decode) with no transpose-copy.

Related PR: #4057

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4058 --add-label <label>

[hsthe29]

Hi @vgokhale @cagrikymk @azaidy 👋

This is my first contribution from a fork, so all CI workflows are currently sitting at
action_required and need a maintainer to Approve and run workflows before they can
execute. Could one of you kindly kick them off when you have a moment?

Since this PR specifically backs SGLang's MambaPool / radix state-cache integration, it
would also be great to run the SGLang path - could you add the ci:sglang label (and
ci:triton-300x if you'd like MI300X coverage)? I don't have permission to add labels
myself.

For context, everything is green on my side locally:

• new op_test test_chunk_opt_vk_state_indices 5/5 pass; existing test_chunk_opt_vk[_varlen] 20/20 pass (dense path byte-for-byte unchanged)
• black / ruff clean
• cross-checked against SGLang's own sgl-fla chunk kernel: o / h / in-place final state all cos = 1.000000, and initial_state_indices + return_h add no measurable overhead

This pairs with #4057 (V-major decode kernel) - together they enable an end-to-end V-major
paged recurrent-state pool. Happy to split, rebase, or adjust review order however works best.

Thanks a lot for taking a look! 🙏

[cagrikymk]

Could you add description of the new args you added?

[hsthe29]

Done — added descriptions for both new args (initial_state_indices and
return_h) to the chunk_gated_delta_rule_fwd_opt_vk docstring, and mirrored
them on the public chunk_gated_delta_rule_opt_vk and the
chunk_gated_delta_rule_fwd_h_opt_vk h-scan wrapper. Also documented the extra
h value returned when return_h=True. Pushed in 3079003 (docstring-only, no
functional change). Thanks for catching it!

[hsthe29]

Pushed 3079003 to address @cagrikymk's review - the new initial_state_indices and return_h args are now documented across chunk_gated_delta_rule_fwd_opt_vk, the public chunk_gated_delta_rule_opt_vk, and the
chunk_gated_delta_rule_fwd_h_opt_vk wrapper (docstring-only, no functional change).

Still green locally on gfx950:

• new test_chunk_opt_vk_state_indices 5/5, existing test_chunk_opt_vk[_varlen] 20/20 (dense path unchanged)
• black / ruff clean