[DSV4.1] Multi-stream prepare for ratio-1/2 layers, fused ratio-1 verify compression, PDL on _q_rope_store

[DarkSharpness]

Motivation

On DeepSeek-V4.1-Flash, 38 of the 40 attention layers have compress ratio 1 or 2. Their decode / verify prepare kept the fused KV-cache write (fused_k_norm_rope_flashmla) on the main stream right behind the Q chain (wq_b GEMM + _q_rope_store), so attention waited for both; only the compressor / indexer went to a side stream, and only in the eight layers that have them. The multi-stream prepare (_forward_prepare_multi_stream) was excluded for these layers, and the special case lived inside _forward_prepare as early_sources.

Two smaller things showed up in the same trace: the ratio-1 KV-source layer still ran the torch compression chain in target-verify (the fused verify path was gated on ratio 2 although c1 is stateless), and _q_rope_store sat behind a launch bubble on the critical path.

Changes

1. _forward_prepare_low_ratio_multi_stream (deepseek_v4.py): the ratio-1/2 decode / target-verify prepare on three streams: compressor + indexer (forward_low_ratio_sources) on one side stream, the KV write on another, the Q chain on the current stream. Each fork is followed by main-stream work and the side blocks are recorded right before their join, so CUDA-graph replay keeps the main chain on one lane ([DSV4.1] Record side-stream work right before its join to keep CUDA-graph replay on one stream #39420). _forward_prepare loses the early_sources branches; dispatch is the former early_sources condition, now for every ratio-1/2 layer.
2. Fused verify compression for ratio 1 (deepseek_v4_backend.py): _low_ratio_compress takes the fused path for compress_ratio in (1, 2) in target-verify. c1 has no pair state, its kernel takes any row count and suppresses padded rows, and c1_out_loc is built for verify. New test test_static_verify_dispatch_ratio_1 (dispatch + pool bytes equal to the unfused chain on a verify-shaped batch). The int64 conversion of positions at the sources entry is dropped (every consumer takes int32 or int64).
3. PDL on _q_rope_store (q_rope_store.py): gdc_wait before the first load, launch_dependents after the store, launch_pdl where supported. Output bit-identical; in a captured GEMM -> kernel chain the kernel-attributable time drops from ~2.0 us to 1.2-1.5 us per call (m = 1..8).

Validation (4x GB300, TP4/EP4, DSpark, bs=1, 64K context, TP0 trace, one verify cycle)

	before	after (1)+(2)
fused_k_norm_rope_flashmla on the _q_rope_store lane	38 / 40 layers	0 / 40 (finishes ~6 us before the Q RoPE starts)
attention-layer time (median, _q_rope_store spacing)	122.9 us	121.0 us
distinct CUDA streams per verify cycle	19	12 (main lane 240 -> 914 kernels; the eight per-layer lanes are gone)
_rmsnorm_fp32 / _rope_fake_quant_pack_indexer / _rope_tail_fake_quant_fp4 per cycle	2 / 1 / 1	0 / 0 / 0
flash_c1_decode_kernel per cycle	0	1

gsm8k (1319 questions, 64 threads) with (1)+(2): 0.9697, stop rate 1.0, no runaways; per-question agreement with the unchanged tree is at the launch-to-launch level (sampling). Graph capture succeeds; smoke outputs byte-identical. Unit tests: test_dsv41_fused_compress 4/4 (incl. the new case), test_dsv41_indexer_prefill 2/2, test_dsv4_q_rope_store 2/2; the split methods were also checked in isolation (eager, side stream, graph replay). A gsm8k run of the full branch including (3) is in flight and will be posted as a comment.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34875916558
Latest PR Test (Extra): ❌ Run #34875915808
Latest PR Test (AMD ROCm 10): ❌ Run #34875916470

[DarkSharpness]

Full-branch check (all three commits) on 4x GB300, TP4/EP4, DSpark, SGLANG_DSV41_DEEP_GEMM_CANDIDATE_INDEXER=1:

• gsm8k, 1319 questions: 0.9704, stop rate 1.0, no runaways; graph capture succeeds; smoke outputs byte-identical to the unchanged tree.
• Trace (bs=1, 64K, TP0, one verify cycle): 12 distinct streams; the KV write is off the Q lane in 40/40 layers and finishes ~5 us before _q_rope_store starts; attention-layer median 119.7 us (122.9 before the branch).
• PDL on _q_rope_store: its predecessor on the stream is the wq_b blockscaled GEMM in 40/40 layers. Gap from GEMM end to kernel start went from +0.48 us (median) to -0.32 us, i.e. the kernel now launches before the GEMM finishes in 38/40 layers; kernel duration 1.82 -> 1.90 us (the prologue now carries the grid-dependency wait). Net ~0.8 us per layer.