[Triton][GDN] Support V-major (hvk) state layout in decode kernel

[hsthe29]

Summary

fused_sigmoid_gating_delta_rule_update (the GDN / gated-delta-net decode kernel) hardcoded the recurrent-state pointer arithmetic to a K-major [N, HV, K, V] contiguous layout:

p_h0 = h0_source + idx*HV*K*V + i_hv*K*V + o_k[:, None]*V + o_v[None, :]
#                                           └ K stride = V     └ V stride = 1  => [K, V]

A framework whose recurrent-state pool is V-major [N, HV, V, K] cannot use this kernel: reading the same bytes with the opposite inner-axis order gives a transposed (wrong) state (cos ≈ 0.02), and because the strides are baked in, it can't be fixed with a .transpose() view.

This is exactly SGLang's situation — its MambaPool recurrent state is allocated [num_slots, HV, V, K] (V-major), and its own sgl-fla decode kernel indexes it as o_v*K + o_k. So aiter's GDN decode kernel is currently unusable as a drop-in for SGLang.

Change

Add explicit (slot, head, K, V) strides to the kernel and a state_layout selector on the wrapper:

• "hkv" (default) — [N, HV, K, V], K-major. Strides (HV*K*V, K*V, V, 1), i.e. byte-identical to the previous hardcoded arithmetic (integer pointer offsets unchanged → no behavior change for existing callers).
• "hvk" — [N, HV, V, K], V-major. Strides (HV*V*K, V*K, 1, K). Lets a V-major pool call the kernel in place (via h0_indices scatter) with no transpose-copy.

Both the initial-state read and the final-state write-back use the passed strides.

Correctness

New test op_tests/test_gated_delta_rule.py::test_fused_sigmoid_gating_delta_rule_update_state_layout asserts "hvk" == "hkv" for the same logical state (same recurrence, only storage differs), across K==V and K!=V, fp32/bf16, with a non-identity h0_indices scatter. Output and final state match to FP-noise.

check	cos	max abs err
aiter hvk vs fp32 oracle (out)	1.000000	~1e-8 (fp32) / ~1e-4 (bf16)
aiter hvk vs aiter hkv (out+state)	1.000000	≤1.2e-4 (bf16), FP-noise
aiter hvk vs SGLang sgl-fla decode on [V,K] pool (out)	1.000000	0.0 in several cases

• Existing decode test (default hkv path) unchanged: 11/11 pass.
• New layout test: 5/5 pass. black 26.x / ruff 0.15.x clean.

Scope / follow-up

This covers the decode kernel. The prefill/extend path (chunk_gated_delta_rule) has a separate gap for the same consumers: its public API takes no initial_state_indices (no in-place scatter) and returns no per-chunk h (needed for SGLang's extra_buffer SSM-state tracking). That is a larger, separate change and is intentionally out of scope here.

[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 4057 --add-label <label>

[hsthe29]

Hi @azaidy - this PR extends the GDN decode kernel (fused_sigmoid_gating_delta_rule_update) to support a V-major hvk [N, HV, V, K] recurrent-state layout, so SGLang's MambaPool (allocated V-major) can call it in place with no transpose-copy.

It's fully backward-compatible: the default "hkv" path is byte-identical to the previous hardcoded arithmetic, and the existing decode test still passes 11/11. A new test asserts hvk == hkv for the same logical state (K==V and K!=V, fp32/bf16), and matches SGLang's sgl-fla decode on a [V,K] pool.

Could you please approve the CI workflows for this PR?