dsv4.1-amd: KV cache layouts, FP4 indexer, compressor and router kernels

[kevin-mii]

This PR was "stack 2/4" of the DeepSeek-V4.1 AMD series (#41018 to #41021). It now follows the CUDA dsv4.1 layout (#39646, #39652, #39653, #39656, #39664, then #38798): kernel PRs by domain, each on main with no callers, then one integration PR. This is the second kernel PR; it does not depend on #41018.

Summary

• Build the V4.1 KV layouts on ROCm (include/sgl_kernel/deepseek_v4/kv_layout.cuh, fp4_utils.cuh): v41::store_row quantizes V41 (fp8, ue8m0 per 32) and V41_FP4 (e2m1, e4m3 per 16) rows on gfx950 with round-to-nearest-even as on CUDA, so fused_store_cache, compress_norm_rope_store, fused_k_norm_rope_flashmla and dequantize_k_cache_paged_v41 serve both layouts there.
• FP4 indexer on ROCm: the RoPE pack kernels build with wave64-aware shuffles (fp4_indexer_rope.cuh); fp4_indexer_rope_hip.cuh writes the split FlyDSL index-K layout; fp4_indexer_hip.py adds the split writer and reader, the FlyDSL query packer, a one-launch index-Q pack + head weights and a page-table bucket width.
• In-tree FP4 indexer decode schedule (fp4_indexer_schedule_hip.py, next to the existing in-tree prefill schedule): build_decode_schedule writes aiter compute_varctx_schedule's cta_info in two dispatches whose registers do not grow with the row count. aiter's single kernel spills past ~1K rows and could not reserve its scratch at DSpark batch 256 (1,536 rows), which crashed decode graph capture.
• ROCm decode router: a batch-invariant split-K bf16 GEMV (ops/gemm/router_gemv_hip.py, beside main's router_gemv.py) and a sqrtsoftplus top-k gate over its partials bitwise equal to aiter's topk_gating_kernel_opt (ops/moe/rocm_router_gate.py). It is here because indexer_head_weights runs the GEMV.
• The low-ratio compressors (c1.cuh, c2.cuh, c2_decode_pool.py) run on ROCm, taking the fp8 pool as bytes there (low_ratio_compress.py); fused_k_norm_rope_flashmla(q=...) ropes the query heads in the K launch on ROCm; sglang/test/kernels/deepseek_v4/dsv41_kv_quant_reference.py is an independent torch oracle for V41_FP4 pages.
• Model call sites are in the integration PR of this stack.

Changes to existing kernels

• CUDA-visible, same outputs: main_norm_rope.cuh gains a defaulted kRopeQ template flag; fp4_indexer.py and fp4_rope_fake_quant.py move their row math into @triton.jit helpers (positions=None keeps the old call); c1.cuh and fp4_indexer_rope.cuh match kDLGPU, which is kDLCUDA on CUDA; c2_decode_pool.py divides with tl.div_rn instead of libdevice.div_rn (both IEEE round-to-nearest; Triton's HIP libdevice has no div_rn).
• ROCm numerics: with SGLANG_USE_AITER, linear_bf16_fp32 now takes the default fp32 torch.mm path instead of aiter's bf16-rounded tgemm (V4 compressor and router logits), and the aiter branch is removed. The V4.1 paths are asserted gfx950-only, and gfx942 code is unchanged.

Verification

• New tests: test_v41_kv_store.py (byte-exact stores and dequant for both layouts, query rope in the K launch), test_rocm_router_gate.py (bitwise to aiter on ties and non-finite logits, the fused split-K gate), test_router_gemv_hip.py (the batch-invariant GEMV, and linear_bf16_fp32 staying in fp32, which fails on the removed aiter route); test_fp4_indexer_hip.py is extended (split index-K writer bitwise to the Triton RMSNorm + RoPE + FP4 chain, decode schedule bitwise to aiter).
• MI350X (ROCm 10.0.0, torch 2.11, triton 3.8, aiter at the Dockerfile pin): these tests and the existing tests on the modified kernels (test_c4_v2.py, test_c128_v2.py, test_deepseek_v4_compress_state_runtime_shapes.py, test_dsv4_unified_fp8_compress_store.py, test_dsv4_fp8_cast.py) pass on this branch, as do the AMD-registered DSV4 attention, MXFP8, mHC and HiCache kernel tests.
• With the integration PR's callers, DeepSeek-V4.1-Flash at TP4 / EP4 with DSpark captures decode graphs up to batch 256 and scores 0.885-0.900 on GSM8K (200 questions, three runs).
• Pending on a CUDA runner: test_v41_kv_store.py (SM100), test_fp4_indexer.py, test_deepseek_v4.py.

Stack

• Base: main. Independent of dsv4.1-amd: gfx950 MXFP8 matmul kernels and fp8-grid producers #41018 (gfx950 MXFP8 matmul kernels).

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #36227764374
Latest PR Test (Extra): ✅ Run #36227764031
Latest PR Test (AMD ROCm 10): ❌ Run #36227764383

[kevin-mii]

/tag-and-rerun-ci

[kevin-mii]

/rerun-failed-ci