dsv4.1: compression, KV I/O, and metadata kernels

[hnyls2002]

Summary

• Add KVLayout (kernels/ops/attention/dsv4/kv_layout.py, include/sgl_kernel/deepseek_v4/kv_layout.cuh): the paged FlashMLA main-KV formats as one enum with their per-token data / scale bytes, tile size and page alignment. V4 is the existing 584-byte layout (448 fp8 nope + 64 bf16 rope + ue8m0 scales); V41 (528 B, every dim fp8 with one ue8m0 scale per 32) and V41_FP4 (288 B, e2m1 codes with one e4m3 scale per 16) are the DeepSeek-V4.1 formats. PagedKV<layout, page_bits>::row gives every writer the same address arithmetic, and v41::store_row is the single quantizer for the V4.1 rows.
• Thread layout through the existing cache writers and reader: fused_store_cache (with an optional in-kernel RoPE for the V4.1 layouts), compress_norm_rope_store, fused_k_norm_rope_flashmla, dequantize_k_cache_paged. The default stays V4, so current callers are unchanged.
• Add the ratio-1 / ratio-2 decode compressors (csrc/deepseek_v4/c1.cuh, c2.cuh, ops/attention/dsv4/low_ratio_compress.py): pair-pooling (ratio 2), RMSNorm, RoPE and the main-KV write in one launch, returning the pre-RoPE latent; out_loc == 0 marks a padded graph row.
• Add the V4.1 dequantizers (dequant_k_cache.py), fp4_utils.cuh (e2m1 fake-quant helpers next to the existing fp8_utils.cuh), the pure-torch reference quantizer (torch_quant.py), and two small-batch metadata builders (build_page_table_positions_small, build_low_ratio_metadata).
• Model call sites (pool layout selection, the low-ratio compressor path) are in dsv4.1: remaining model and runtime integration #38798.

Changes to existing kernels

• csrc/deepseek_v4/store.cuh, fused_norm_rope_v2.cuh, main_norm_rope.cuh: the V4 path now addresses pages through PagedKV; same bytes as before. The kernel templates gain a KVLayout parameter.
• csrc/deepseek_v4/c_plan.cuh: ragged_id is a zero-based uint16, so a 65536-token batch is valid; batch_size must stay below 65535 because pack_w(65535, 65535) is the invalid-write sentinel. test_deepseek_v4_compress_plan_bounds.py covers both bounds and the sentinel collision.
• get_paged_mqa_logits_metadata accepts page size 128 as well as 64 (the schedule depends only on the sequence lengths).

Verification

• JIT builds of every layout variant of the touched kernels (c1, c2, store, compress norm-rope, main-K norm-rope) succeed.
• Existing tests on the modified kernels are rerun on this PR: test_c4_v2.py, test_c128_v2.py, test_deepseek_v4_compress_state_runtime_shapes.py, test_fp4_indexer.py, test_paged_mqa_metadata.py, plus test_deepseek_v4_compress_plan_bounds.py.

---

CI States

Latest PR Test (Base): 🚫 Run #35147198293
Latest PR Test (Extra): ❌ Run #35147197934
Latest PR Test (AMD ROCm 10): ⏳ Run #35147198240

[BBuf]

Read the whole diff. The kernel work is high quality and the header comments are unusually good — a few things I checked and can confirm rather than ask about:

• store_row_fp8's claim that max((bits >> 23) - 8 + (mantissa > 0x600000), 114) reproduces ceil(log2(max(amax / 448, 1e-4))) exactly holds. I ran it against torch_quant.cast_scale_inv_to_ue8m0 on GB300 over every positive bf16 value, 134M random fp32 values, the ±64 ulp neighbourhood of every 448 * 2^k, and zero/denormals: 0 mismatches. Nice trick, and worth the comment it has.
• clear_negative_zero / the pack_w sentinel reasoning in c_plan.cuh (batch_id < 65535 keeps pack_w(65535, 65535, ...) unreachable while ragged_id may reach 65535) both check out.

Two things on the diff.

---

1. This PR deletes the draft-pad regression tests and replaces them with tests for something else

test_deepseek_v4_compress_plan_draft_pad.py is rewritten from scratch (+51 / -196). The five tests removed are:

• test_c4_ring_residency, test_c128_ring_residency
• test_cpu_and_gpu_planner_agree
• test_plain_prefill_write_set
• test_over_capacity_under_writes

They do not move anywhere else — I grepped the whole branch for _assert_ring_residency, _written_positions, _max_draft_tokens and each test name; nothing. The three new tests cover the uint16 index boundary this PR changes, which is good, but they are a different invariant in a file that still says draft_pad in its name.

The logic they pinned is still live and unchanged on this branch:

// c_plan.cuh:524
const auto mtp_pad = ring_size > window_size ? ring_size - window_size + 2 : 0;

used at c_plan.cuh:207, :224 and :590, and pool_configurator.py still carries _assert_ring_serves_draft_tokens. The deleted file's own docstring describes exactly why it existed:

c_plan.cuh used to cap that pad at 4 (kMaxMTPDraftTokens), silently under-writing the ring for larger draft counts — no IMA, no NaN, just wrong compressed state.

test_cpu_and_gpu_planner_agree in particular was the only thing pinning the host loop against plan_compress_prefill_kernel0; those two planners are edited independently and now nothing catches them drifting.

Could the new boundary tests be added to the file rather than replacing it? If the old ones had to go for runtime or flakiness reasons, could that be said explicitly so the trade is on the record?

2. None of the new entry points have a test

Every new public surface in this PR is untested on the branch:

added	test
c1_decode_norm_rope_store	none
c2_decode_or_verify_norm_rope_store	none
dequantize_k_cache_paged_v41 (fp8 + fp4 Triton)	none
quantize_k_cache_v41 / dequantize_k_cache_v41	none
fake_quant_compressed_kv / fake_quant_fp4	none
low_ratio_metadata, page_table_positions_small, fill_all_compressed_indices	none

This is the cheapest possible test to write, because the PR ships both sides: torch_quant.py is an explicit byte-for-byte reference ("byte for byte what the store / dequant kernels are checked against") and dequantize_k_cache_paged_v41 is documented as "bit-exact with the pure-torch dequantizer". A single round-trip test — store_row through the kernel, quantize_k_cache_v41 in torch, compare bytes — would lock the whole V41/V41_FP4 layout down. Is that coming in a later PR of the stack, or was the intent to rely on e2e accuracy only?

---

Pushed one cleanup

f33acae on this branch: flash_c2_decode_kernel loaded freqs_cis twice for the same lane and position —

// c2.cuh:141
if (tx >= kNopeThreads) freq.load(params.freqs_cis + (pos - 1) * kRopeDim, tx - kNopeThreads);
...
// c2.cuh:202, inside the branch that consumes it
freq.load(params.freqs_cis + (pos - 1) * kRopeDim, tx - kNopeThreads);

freq is never read between them, so the first is dead. I kept the early one (it is issued before the softmax and the RMSNorm reduction, so its latency overlaps that work) and removed the reload. Verified bit-exact on GB300: identical sha256 of out, the paged cache and the pair-state ring for V4 / V41 / V41_FP4 decode plus a draft_len=4 verify batch, JIT cache cleared between runs, with a deliberate eps perturbation as a negative control to confirm the harness actually recompiles. Revert freely if you would rather keep the diff minimal.

Note this branch is the base of the stack, so the commit is not in dsv4.1-candidate and its descendants — and I saw dsv4.1-metadata get force-pushed while I was reviewing, so drop it if it is in your way.

(Review by Claude Opus 5, run by @BBuf.)

[hnyls2002]

/tag-and-rerun-ci

[hnyls2002]

/rerun-test test/registered/kernels/ops/attention/test_c4_v2.py test/registered/kernels/ops/attention/test_c128_v2.py test/registered/kernel/attention/test_deepseek_v4_compress_plan_bounds.py test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py test/registered/kernels/ops/attention/test_fp4_indexer.py test/registered/kernels/ops/attention/test_paged_mqa_metadata.py

[github-actions]

Results for /rerun-test test/registered/kernels/ops/attention/test_c4_v2.py test/registered/kernels/ops/attention/test_c128_v2.py test/registered/kernel/attention/test_deepseek_v4_compress_plan_bounds.py test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py test/registered/kernels/ops/attention/test_fp4_indexer.py test/registered/kernels/ops/attention/test_paged_mqa_metadata.py:

🚀 1-gpu-h100 (6 tests): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_c4_v2.py
cd test/ && python3 registered/kernels/ops/attention/test_c128_v2.py
cd test/ && python3 registered/kernel/attention/test_deepseek_v4_compress_plan_bounds.py
cd test/ && python3 registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py
cd test/ && python3 registered/kernels/ops/attention/test_fp4_indexer.py
cd test/ && python3 registered/kernels/ops/attention/test_paged_mqa_metadata.py

[hnyls2002]

/rerun-test test/registered/attention/unittests/dsv4/test_deepseek_v4.py test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200.py test/registered/attention/unittests/dsa/test_dsa.py

[github-actions]

Results for /rerun-test test/registered/attention/unittests/dsv4/test_deepseek_v4.py test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200.py test/registered/attention/unittests/dsa/test_dsa.py:

🚀 4-gpu-b200 (3 tests): ❌ View workflow run

cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py
cd test/ && python3 registered/e2e/models/test_deepseek_v4_flash_fp4_b200.py
cd test/ && python3 registered/attention/unittests/dsa/test_dsa.py

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py
cd test/ && python3 registered/attention/unittests/dsa/test_dsa.py