fix(dspark): reduce V2 working VRAM

[alexbi29]

Summary

Reduces DeepSeek V4 / DSpark V2 working VRAM and increases usable KV cache on the 2x SM120 DSpark serving config.

Main changes:

• Skip the target MTP hidden-state staging buffer unless the speculative method is actually mtp.
• Reuse one C128A top-k backing buffer for decode + prefill metadata rows.
• Let DSpark V2 skip DFlash's persistent context hidden-state staging tensor.
• Avoid reserving the packed-cache padded-Q scratch when n_local_heads == padded_heads; in that case the fused qnorm/RoPE/KV-insert op can write Q in place.
• For SM120 packed C128 prefill, map local top-k indices to global KV slots in place once per step and reuse the global view across layers. This removes a full-context runtime temporary/conversion launch; it does not change the KV profile line.
• Fuse final MHC post-processing paths so DSpark/target finalization no longer materializes [T, hc, H] just to produce hc_head output and Eagle aux means.
• Reuse the previous MHC residual buffer for large-token fused post/pre transitions when the old residual is dead, avoiding an additional [T, hc, H] allocation.
• Reuse the dead MHC layer input as the next fused post/pre output buffer, removing one 64 MiB BF16 [T, H] profile allocation per rank without changing kernel math or launch count.
• Keep DSpark probabilistic rejection logits in a persistent graph-written buffer. A current-step dense handoff was tested but backed out because DSpark drafting is CUDA-graph replayed; Python reassignment does not run per replay and caused lower acceptance/TPS.

Live memory result

Measured on 2x RTX PRO 6000 Blackwell / SM120, TP=2, DSpark V2, GMU 0.975, max model len 524288.

State	KV memory	KV tokens	Max concurrency
baseline before this cleanup	10.54 GiB	1,095,149	2.09x
after MTP-buffer cleanup	10.79 GiB	1,121,196	2.14x
after in-place Q cleanup	11.05 GiB	1,147,140	2.19x
after C128 SM120 prefill in-place global indices	11.05 GiB	1,147,140	2.19x
after fused MHC final/aux captures	11.30 GiB	1,173,083	2.24x
after MHC residual-buffer reuse	11.55 GiB	1,199,131	2.29x
after MHC input-buffer reuse	11.61 GiB	1,205,591	2.30x

Net from baseline: +1.07 GiB available KV and +110,442 KV tokens.

Validation

• git diff --check, targeted py_compile passed.
• Prototype checkout: tests/v1/spec_decode/test_dspark_config.py, test_rejection_sampler_utils.py, test_dspark.py passed (48 passed).
• Clean PR27 fold worktree: config + rejection sampler tests passed (30 passed).
• Direct DSpark serve smoke: 7*8 -> 56.
• Live reload on the PR27 head confirmed the 11.55 GiB available KV, 1,199,131 KV tokens, 2.29x concurrency profile.
• Concurrency probe, 12 requests, temp=0, thinking=false: 5,986 output tokens in 9.30s, 643.5 tok/s aggregate, draft accept rate 0.4683, mean AL 3.341.
• 2026-07-09 current-step draft-logits A/B: persistent graph-written draft logits restored acceptance/TPS. Warm decode run: avg 309.2 tok/s, median 310.9, range 304.7-311.4, versus current-step post-JIT baseline median 297.7.
• Same A/B prefill run: 6106 / 8004 / 7710 / 6369 tok/s at 1K / 4K / 16K / 65K prompt sizes.
• Coherency smoke: 17 + 25 -> 42.
• Clean PR worktree fixup tests after restoring persistent draft logits:
  • tests/v1/spec_decode/test_dspark_config.py -> 2 passed
  • test_dense_draft_logits_index_mapping_matches_state_indexed and test_gumbel_sample_stores_processed_logits_inplace -> 2 passed
• 2026-07-10 MHC input-reuse validation: focused CUDA gate 20 passed; deployed profile 11.61 GiB / 1,205,591 tokens / 2.30x; warm decode A/B 307.6 vs 307.5 tok/s; final deployed warm median 310.0 tok/s; coherency 17 + 25 -> 42.
• Deployed concurrent decode (max_tokens=500, temperature=0, thinking off): c=2 275.7 aggregate tok/s (52.95% draft accept), c=4 393.2 tok/s (55.07%), c=8 577.9 tok/s (56.92%).
• Deployed long-context decode-only TPS (2048 output tokens, temperature 1.0): 32,783-token prompt median 286.5 tok/s; 65,539-token prompt median 287.9 tok/s; both ~7% below the 310.0 tok/s short-context median with no further 64K degradation.

[jasl]

Sorry, I missed this one

[jasl]

Merged — thanks, this was a substantial piece of work and it holds up well under review.

The seven fix(dspark)/test/chore commits went in as 1a2e039039, and ship on
codex/ds4-sm120-min-enable + ds4-sm120-preview-dev as part of 1ee1ec7930
(which also carries 34 upstream commits merged on top).

Validation on GB10 (2× SM121a, TP=2, MTP2, mml 49152, util 0.85, fp8 KV)

gate	result
GSM8K	0.9568 — identical to the pre-merge baseline
arthur long-context coherence, c=1	2/2
arthur long-context coherence, c=12	23/24 (our known 22–23/24 band)
toolcall-15 (en × 3 thinking modes, rounds=3)	86% — unchanged
instruction-following JSON-only	PASS
IMA / assert in serve log	0

Unit suite green: 795 kv_offload, 6 DSv4 sparse-SWA/ubatch/indexer, 35 rejection sampler.
A DSpark serve answers correctly (17×23 → 391, capital of France, first three primes).

Re-gated after merging 34 upstream commits on top: arthur c=1 2/2, IMA 0, and GSM8K three times
— 0.9469 / 0.9507 / 0.9583. Worth stating plainly since it bears on the numbers above: that
1.1pp spread on one unchanged build is wider than our entire historical build-to-build range, so
GSM8K here discriminates gross correctness breakage, not sub-1pp quality deltas.

Kernel-test A/B at 520c25949c vs the merge, same node and environment:
test_mhc_kernels.py 13 → 11 failures, test_rejection_sampler_utils.py 29 → 31 passed
(your two new tests pass). The one AFTER-only failure,
test_mhc_fused_post_pre_reuses_dead_buffers[128], is the [128] parametrization that already
fails for the pre-existing test_mhc_fused_post_pre[*-128] without the reuse variant — it
inherits an existing failure rather than introducing one. No new failure modes.

Review notes

The parts that needed the most scrutiny were the in-place mutations, since geometry- and
aliasing-invariant bugs have bitten this branch repeatedly. All of them check out:

• Rewriting the C128A prefill indices in place local→global and caching them across layers
  is safe: the Triton kernel reads and writes the same element of each row with no cross-row
  reads, and the super()._forward_prefill fallback — whose parent reads that field expecting
  local indices — is gated on the global VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL, so
  the whole model takes one path and layers never mix readers.
• Sharing one c128a_topk_buffer between decode and prefill partitions exactly on
  is_decode = token_idx < num_decode_tokens; the stale tail of the shared lens buffer is
  harmless because the consumer stores topk_lens[token_idx] unconditionally.
• Making _mtp_hidden_buffer conditional is fine for DeepSeek-V4-Flash: its
  inference/config.json has both compress_ratios and hc_mult (4), so the MTP path still
  allocates it.

Two deltas relative to your current branch head

1. rejection_sampler.py — we bind the trailing optionals by keyword. When we merged, the
draft_logits_index_mapping threading was not yet present, so self.synthetic_conditional_rates
was passed positionally into the slot rejection_sample() reserves for
draft_logits_index_mapping — the two are adjacent in the signature, so it bound silently with
no conflict and no error. You have since fixed the same thing by adding the missing argument;
ours is the same semantics with keyword binding plus a comment, which is more robust if that
signature grows again. Functionally equivalent — just flagging that we kept ours.

2. fix(config): keep sequence-parallel MoE enabled at DP=1 — we did not take this one.
Not a disagreement with your measurements, but it doesn't fit our branch:

• It's inert for us. use_sequence_parallel_moe also requires enable_expert_parallel, and
  expert parallel is off in our standard GB10 serve, so the property is False either way.
• Reverting an upstream commit means carrying a conflict point through every future upstream
  merge, and silently re-breaking the pure-TP memory case Fix: Restore data_parallel_size > 1 for use_sequence_parallel_moe vllm-project/vllm#48849 was written for, for everyone
  using this fork.
• If we later enable EP at DP=1, flipping use_sequence_parallel_moe to True also flips
  pad_shared_expert to False in both dspark.py and model.py — the DSpark shared-expert
  weight-loading path that broke a 4-node deployment a week ago.

Your own commit message points at the right answer: make the gate config-aware rather than
requiring DP > 1. That's an upstream conversation, and your numbers are good ammunition for it —
worth taking to vllm-project#48849 directly.

Two things worth knowing

• mhc_fused_post_pre_tilelang_reuse_residual is not direct_register_custom_op'd and is
  absent from __all__, unlike the mhc_fused_post_pre_tilelang it replaces on the hot path.
  We measured no effect (0 graph breaks, CUDA-graph capture 24 s), and it is arguably the safer
  side — an unregistered plain function leaves its in-place mutation visible to the tracer.
  Please don't "fix" this by registering it under the file's uniform mutates_args=[]:
  that would declare a mutation that does happen as absent.
• We could not confirm the VRAM saving on our hardware, and we're not disputing it. The
  serve-reported "available KV memory / GPU KV cache size" turned out to swing ~1 GiB / 9%
  between runs of near-identical builds on our setup, which is larger than the effect we were
  trying to resolve, so single-run comparisons of that figure are not meaningful here. Our
  default path is MTP, which keeps _mtp_hidden_buffer and so doesn't collect the DSpark win
  anyway.

Closing as merged. Thanks again.