DeepSeek-V4 FP4: fused_compress FP4 scatter + rmsnorm_rope_rotate FP4 KV-cache kernel

[junhaha666]

Summary

FP4 support for the DeepSeek-V4 CSA compressor path (gfx950): FlyDSL fused_compress_attn FP4 scatter + the rmsnorm_rope_rotate_activation_fp4quant_kvcache HIP kernel, plus a KV-scale layout fix.

Changes

• rmsnorm_rope_rotate_activation_fp4quant_kvcache (dsv4_rotate_quant.cu): fused RMSNorm + RoPE + optional Hadamard rotate + FP4 (E2M1) quant + paged KV-cache write.
• fused_compress_attn FP4 quant scatter (legacy + K-split), unified behind a single quant_mode selector (none | per_row_fp8 | group_fp8 | fp4); quant/quant_fp4 are derived from it.
• KV e8m0 scale layout fix: write the scale interleaved ((pos%16)*(kv_block_size/16) + pos/16) so the pa_mqa_logits_fp4 reader's packed-dword load lines up. The linear layout mis-read every scale.

Tests

• Expanded test_dsv4_rotate_quant.py (FP4 kvcache, shuffle scale, rotate) and test_flydsl_compress_attn.py (FP4 shapes).

🤖 Generated with Claude Code

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

[zufayu]

Hi @junhaha666 — reviewed the FP4 changes in depth. Overall the quality is solid: fp4_max is derived via finfo (not hardcoded) with arch-correct comments (448 e4m3fn gfx950 / 240 e4m3fnuz gfx942), the out.dtype dispatch in rotate_activation/rope_rotate_activation is complete (fp4x2 / fp8 / else→bf16, each with a scale assert, no silent fall-through), amax has an all-zero-group eps floor, and there's a torch_ref + tests + green CI. One thing I need confirmed before approving, plus two smaller notes.

1. fp4 auto-K-split: the docstring contradicts the code — and I want to confirm the auto path is tested.

flydsl_fused_compress_attn's docstring (fused_compress_attn.py ~line 2709) says fp4 is "Legacy single-wave kernel only (never auto K-split)". But tracing the actual dispatch, fp4 does auto-engage K-split:

• _is_csa_indexer_fp4 (~line 2973) matches D128 / RD64 / ratio4 / overlap + _mode=="fp4";
• with k_split_num_waves=None and has_bt, it enters the auto branch → nw_eff = csa_ksplit_num_waves(...) (and the comment there notes the FP8 wave-count heuristic carries over to fp4, so nw>1);
• use_ksplit = nw_eff>1 and has_bt and (not nm_asm or _is_csa_main_nm) — fp4 is not nm_asm, so nothing gates it out → it runs compile_...ksplit(quant_mode="fp4").

So in the default case (no explicit num_waves + block_table present), the fp4 indexer runs the K-split kernel. The code clearly intends this (the module docstring lists "FP4 scatter" and there's a dedicated fp4 branch in _build_kernel_ksplit) — that function-level docstring line is just stale.

My concern isn't the contradiction itself, it's coverage: does the test suite exercise this auto path? The fp4 scatter/amax/e8m0 code in _build_kernel_ksplit is largely mirrored from the single-wave _build_kernel, but the wave layout differs (single-wave tid vs K-split lid 0-63, different group-reduce span). If the tests only cover single-wave fp4 (or pass an explicit k_split_num_waves), the auto-triggered ksplit fp4 path is effectively untested.

Could you (a) confirm the fp4 auto-K-split path is covered by a test that does NOT pass an explicit k_split_num_waves, and (b) fix the stale "never auto K-split" line in the docstring?

2. KV e8m0 scale interleave layout — reader side isn't in this diff.

The scale writer moves to the interleaved layout (sflat=(slot%16)*4+(slot//16)), but the consumer pa_mqa_logits_fp4 isn't part of this PR. If the writer/reader layouts don't match, every scale is read wrong (silent accuracy drop). Could you paste the reader-side packed-dword read, or an end-to-end accuracy comparison showing the two sides agree?

3. develop=True on the dsv4_rotate_quant ops bypasses the JIT cache — fine while experimental, but please remove before this lands in a release branch.

Nothing else is blocking — arch handling, dispatch completeness, and boundary guards all look correct. Once the fp4 K-split coverage is confirmed, I'm happy to approve.