[DSV4.1] Four BS=1 critical-path fixes for DSpark decode (887.10 -> 905.40 tok/s on 4xGB300)

[BBuf]

Motivation

A BS=1 latency round on DeepSeek-V4.1-Flash / 4xGB300 (attention TP4, MoE TP4
--ep-size 1, DSpark static verify block 5). Profiling the verify cycle turned up four
places where the critical path paid for work it did not need. Each is independent, small,
and individually revertible.

Result

Protocol: 4xGB300 (sm103a), random 4096-in / 1024-out, BS=1 serial client,
SGLANG_SIMULATE_ACC_LEN=5.5 + match-expected, --speculative-algorithm DSPARK --speculative-dspark-block-size 5, no profiler attached. Metric is the client-side
output_tps median. Base is this PR's merge base, 0ea9492; each launch is one warmup
plus 7 measured runs, and the two arms were launched alternately on the same box in the
same session.

	launch A	launch B	launch C	combined median	n
0ea9492	883.30	873.50	888.10	887.10	21
+ this PR	905.40	887.80	909.50	905.40	21

887.10 -> 905.40 tok/s (+2.06%). All three launch pairs improve (+2.50 / +1.64 /
+2.41%). Three launches per arm because the baseline's own launch-to-launch spread is
~2% -- FlashInfer selects GEMM/MoE tactics per process, so a single A/B launch pair
cannot separate a few percent from restart noise.

Accepted length is unchanged (5.505 median in both arms) -- these are cycle-time changes,
not acceptance ones.

Every change here is one a production workload sees; nothing in this PR is specific to
the benchmark mode the measurement runs in.

What changed

#	change	cycle effect
1	split the verify argmax over the vocabulary	-46 us
2	vision router emits the FlashInfer packed top-k ids	-64 us
3	MXFP8 quantize as an epilogue of the mHC combine+RMSNorm	-56 us
4	FlashMLA's split-KV schedule computed by a wide kernel	-83 us

-249 us per verify cycle in total.

Attribution is from the GPU trace (kernel counts and per-cycle kernel time) plus the
acceptance-length-independent cycle_us = (elapsed_s - ttft_s) / verify_steps;
output_tps alone has a +-5% spread at this sample count, wider than any single item.

1. torch.argmax on [6, 129280] fp32 takes 23.6 us -- at::native::reduce_kernel
   gives each row one block, so six very wide rows leave the machine idle for a reduction
   that reads 3.1 MB. A flat two-stage split does it in 3.5 us. The accept sits between
   the target verify and the KV commit, so the whole cost was exposed.
2. vision_topk dropped the packed_out argument the text path passes. It is the
   router every MoE layer of a checkpoint with a vision tower goes through, so each layer
   paid a PackTopkIds launch to rebuild what the gate already had in registers.
3. The mHC combine+norm output was written BF16 and immediately re-read by a
   standalone mxfp8_quantize inside the first projection. At BS=1 the two launches cost
   about the same.
4. FlashMLA builds its tile-scheduler metadata in a <<<1, 32>>> kernel whose entire
   partition loop runs on thread 0, storing each 32-byte entry to global memory. The loop
   is over num_sm_parts = num_sms / s_q, so a BS=1 step walks 152 partitions at
   ~156 ns each -- 28.05 us in the decode graph, fully exposed. It cannot be hoisted out:
   on_after_cuda_graph_warmup clears the scheduler before capture on purpose, because
   the schedule depends on the step's topk_length and a frozen one would leave the tail
   of a longer row unattended. The new decoding_sched_meta runs the same algorithm over
   a whole block -- request state in registers instead of a chain of dependent shared
   loads, and the walk stops once the last request is consumed, leaving the whole idle
   tail (over 100 of the 152 partitions at BS=1) to be filled in parallel. 28.05 -> 7.44
   us in situ, same 80 launches per 20 steps, and no get_mla_metadata_kernel left in
   the trace.

Correctness

Changes 1, 3 and 4 are bitwise identical to what they replace, with new registered tests:

• test/registered/kernel/speculative/test_dspark_fast_argmax.py -- vs torch.argmax
  over random rows, exact ties, +-Inf, an all -Inf row and a strided row. Ties resolve to
  the lowest index, matching ArgMaxOps' strict >.
• test/registered/kernel/hyperconnection/test_hc_combine_norm_mxfp8.py -- vs
  hc_combine_norm followed by flashinfer_mxfp8_quantize(..., is_sf_swizzled_layout= True) on both its cuda and cute-dsl backends, across five magnitude decades:
  the BF16 output, the FP8 values and the swizzled UE8M0 scale block all compare equal.
• test/registered/kernel/attention/test_decoding_sched_meta.py -- the schedule itself
  over batch 1..64, s_q 1 and 6, top-k 512 and 2048 and five topk_length shapes
  including all-zero and all-one rows; the attention output compared as raw bits (a
  random fp8 cache decodes to NaN in places); and the extra-cache schedule, which rounds
  the main length up to a block before adding the extra one. FlashMLA's
  DecodingSchedMeta ends in a _pad word it never writes, so the comparison covers the
  seven defined fields.

Change 2 passes the same ids through a second output buffer the gate already fills.

No dataset evaluation was rerun, because nothing here is an approximation.

Measured and deliberately left out

These were built and measured on the same box in the same session and are not in this
PR, recorded so nobody repeats them:

• Letting SGLANG_SIMULATE_ACC_LEN take the in-graph accept is worth -104 us/cycle
  in that mode and nothing at all outside it. fold_eligible refuses the fold whenever
  the variable is set, so a simulated run takes the eager accept while a real greedy run
  takes the folded one, and the graph's gated copy runs on top of it -- about 40 tiny
  launches, two full-vocabulary argmaxes, three NCCL broadcasts and a second KV-inject
  chain, every step. Staging the drawn length in a device buffer fixes it. Left out
  because it cannot help a production workload: with the variable at its default the
  fold_eligible expression is already equivalent, so the change is inert. It is a
  benchmark-fidelity fix, not a speedup, and belongs on its own if anyone wants it.
• A fused router (the logits GEMM and the gating in one launch) is not worth it here.
  The chain is tiny_gemm_bf16 2.68 us + moe_fused_gate 2.85 us = 5.53 us per layer,
  and the router chain exposes only 64 us/cycle because 87% of it is already hidden on a
  side stream. tiny_gemm_bf16 is not the weak link either -- a split-N Triton GEMV over
  seven configurations came in at 10.8-13.4 us against its 2.68. So a fusion buys one
  launch, ~0.2-0.4% end to end, against a bitwise reimplementation of sqrtsoftplus,
  noaux_tc, renormalize and the in-register pack. An earlier Triton prototype was bitwise
  correct but 21-45 us: a grid-wide release/acquire dominates at this size.
• DeepGEMM mega_mhc is a regression at BS=1: 10.7 us at T=6 against the 5.25 us of
  critical path it replaces, because SGLang hides the mHC statistics and Sinkhorn on a
  side stream and the fused kernel drags them back onto the main one. +434 us/cycle.
  mega_gate is capped at the same 64 us the router chain exposes. Both are worth
  revisiting at BS>=64, where side streams have less to hide.
• FlashMLA's fused norm-rope-attn-rope-cast is at parity at BS=1: 18.05 us at
  s_q=6 h_q=64 topk=512 against ~18.4 us for sglang's sparse decode + q_rope_store +
  the O-RoPE and O-quant it also absorbs. It also does not build under nvcc 13.0 without
  widening two non-dependent static_asserts that sit in discarded if constexpr
  branches (core_attn/kernel.cuh:327 and :536; at h_q=64, FOLD_FACTOR is 2 and both
  branches are dead). That is the same build failure [DeepSeek-V4.1] Bump FlashMLA to the fork's rebase head (v4.1 kernels) #39171 works around by skipping the
  TU group.
• Candidate block selection via DSV4 top-k v2 instead of torch.topk is worth
  another -62 us/cycle (candidate_block_logits goes 82.9 -> 8.1 us), but it cannot keep
  the torch.topk tie contract that candidate_block_logits documents and
  test_dsv4_indexer_postprocess enforces: top-k v2 partitions the row, so on exactly
  tied scores it keeps a different (equally-scoring) set, and it ranks NaN last where
  torch.topk ranks it first. Left out rather than weakening that test.
• Fusing the q-LoRA RMSNorm with its MXFP8 quantize is bitwise identical but slower
  (6109 vs 6028 us/cycle): FlashInfer's two launches spread a [<=8, 1280] row over more
  CTAs than one-CTA-per-row does, and at that shape the parallelism beats the launch and
  the round trip the fusion removes.
• Fusing the mHC post-split into combine+norm, also slower in a dependency chain.

Unrelated, but worth knowing

lmsysorg/sglang:dev-dsv41 ships libcuda.so only under /usr/local/cuda/lib64/stubs,
while the sgl_kernel_jit_cuda_ipc JIT link line uses -L/usr/local/cuda/lib64 -lcuda.
The link fails with cannot find -lcuda, SGLang logs Setup Custom allreduce failed and
falls back to NCCL for every TP all-reduce in the model. On this box that is worth
+28% at BS=1 (681.54 -> 872.50 tok/s) and the only symptom is one warning line.
export LIBRARY_PATH=/usr/local/cuda/lib64/stubs fixes it. Worth checking on any box
where DSV4.1 BS=1 numbers come in low; I have not touched the image or the link line in
this PR.

Checklist

• Format with the pinned pre-commit (pre-commit run --from-ref ... --to-ref ...).
  The check-registered-tests hook is red on dsv4.1 for 23 pre-existing files; the
  three files added here are not among them.
• New registered unit tests for all three bitwise-identical kernels.
• No documentation change needed.

---

CI States

Latest PR Test (Base): ❌ Run #34830084240
Latest PR Test (Extra): ❌ Run #34830083868
Latest PR Test (AMD ROCm 10): ❌ Run #34830084204