[DSV4.1] Record side-stream work right before its join to keep CUDA-graph replay on one stream

[DarkSharpness]

Motivation

A DSpark verify cycle on DeepSeek-V4.1 replays its CUDA graphs across ~144 distinct CUDA streams (TP0, bs=1, 64K context): 85 opened by _hc_mix_stats_partial_kernel, 45 by the routed MXFP8 pre-quant, three new streams per layer, and the main chain moves onto a fresh stream at every mHC join. Traces become hard to read (every layer lands on new stream ids) and the graph does not run the way it was captured.

The rule, established on the running model: at replay, a node that joins several branches continues on the lane of its first-recorded parent. Both offenders record their side-stream block before the main-stream kernels that block is later joined with: the mHC statistics right after the combine (before attention / MoE), the routed pre-quant before the router top-k. The join then follows the side lane and drags the main chain with it. Reordering at the fork does not help (tested: 144 -> 144); recording the side-stream work right before its join does.

Changes

• deepseek_v4.py: _hc_mix_and_combine becomes _hc_combine (the combine on the current stream plus the fork of stats_stream, keeping the tiny-input rule of forking after the combine) and _hc_mix_stats (the statistics on stats_stream). forward_hc_pre_from_prev records _hc_mix_stats after the attention / MoE, right before it joins the stream. The DSpark stage does the same.
• deepseek_v2.py: the routed MXFP8 pre-quant on routed_quant_stream is recorded after the router top-k, right before the routed MoE joins its event. The fork point is unchanged.

Graph topology is unchanged: the fork events stay where they were, so at replay the statistics and the pre-quant still overlap the same main-stream kernels; only the capture order moves.

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

	distinct streams	opened by mix-stats	opened by routed pre-quant
before	144	85	45
pre-quant reorder only	103	85	4
both	19	2	4

The remaining 19 are one per DSpark segment plus a handful of singletons. Overlap is preserved: the mix-stats kernel still starts ~3.4 us after the combine and runs alongside the wqkv GEMM (3.7 us before); 62% of mix-stats launches overlap another kernel in both traces; attention layer time 122.3 -> 122.9 us and verify cycle 6.36 -> 6.41 ms are within single-launch noise. gsm8k, 1319 questions: 0.974 (0.970-0.972 across launches before). Graph capture succeeds; smoke outputs are byte-identical.

The split methods were checked in isolation against the fused kernels called directly: bit-equal in eager mode, with a side stream, and under CUDA-graph replay.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34846601359
Latest PR Test (Extra): ❌ Run #34846600774
Latest PR Test (AMD ROCm 10): ❌ Run #34846601007