[Feature][DSpark] Support Speculators-format checkpoints and ragged multimodal CUDA graphs

[qq1060]

Status: implementation, public-checkpoint compatibility, and paired
fixed-output performance tests are complete. Greedy target-output equivalence
remains unresolved, so this PR does not claim lossless output.

Motivation

This PR extends SGLang's DSpark path to run native Speculators checkpoints end
to end in eager mode, CUDA Graph replay, and ragged multimodal verification. It
covers the checkpoint config and weight formats, both DSpark anchor layouts,
independent draft vocabularies, and the runtime paths needed to verify them.

Native Speculators checkpoints are not limited to the shared-vocabulary layout
that DSpark originally supported in SGLang. They may contain an independent
reduced-vocabulary LM head plus a draft-to-target (d2t) mapping. For example,
the public
RedHatAI/Qwen3.6-35B-A3B-speculator.dspark
checkpoint has:

• target vocabulary: 248,320 tokens;
• draft lm_head: 32,000 tokens;
• markov_w1: target-vocabulary input;
• markov_w2: draft-vocabulary output;
• d2t: a 32,000-entry mapping from sampled draft IDs to target IDs.

Treating these tensors as a shared-vocabulary model either fails during loading
or feeds/scatters IDs in the wrong vocabulary space.

Speculators checkpoints can also use either DSpark draft-block layout:

• sample_from_anchor=false: the anchor is only a conditioning token. The draft
  forward has gamma + 1 slots, and slot 0 is excluded from sampling and
  verification.
• sample_from_anchor=true: the anchor is a trained prediction. The draft
  forward has gamma slots.

Using the wrong width does not always crash; it can shift real draft positions
and reduce the acceptance length toward 1.

After checkpoint loading and draft geometry were working, mixed image/text
validation exposed separate problems in ragged mRoPE and linear attention. The
ragged-path changes make multimodal verification work under CUDA Graphs.

During the same validation, I found a pre-existing interaction between
/flush_cache and DSpark draft CUDA Graphs: emptying the allocator cache can
invalidate addresses retained by a captured graph. This is included as a
separate DSpark graph-lifecycle fix.

Modifications

• Normalize the native Speculators DSpark config format, whose backbone lives
  under transformer_layer_config, without importing checkpoint-side remote
  config code. This path is gated on speculators_model_type=dspark and reuses
  the existing config-dictionary fetch.
• Load an independent tensor-parallel ParallelLMHead when
  draft_vocab_size != target_vocab_size; preserve the original shared-head
  path for full-vocabulary DSpark checkpoints.
• Support the asymmetric Markov geometry: markov_w1 consumes target IDs while
  markov_w2 and the draft LM head produce draft-vocabulary logits.
• Map sampled draft IDs through d2t before storing tokens or feeding the next
  Markov step. For probabilistic acceptance, scatter draft probabilities into
  target-vocabulary columns and leave unmapped target tokens at -inf.
• Validate mapping shape/range and fail early for unsupported reduced-vocabulary
  checkpoints without d2t. Load t2d as checkpoint metadata without using it
  in the decode path.
• Read authoritative gamma and sample_from_anchor metadata, reject
  inconsistent checkpoint geometry at startup, and propagate the same draft
  width through eager execution and CUDA Graph replay.
• Keep verification, KV commit, and acceptance accounting at gamma real draft
  tokens. Only the draft forward width changes.
• Build speculative mRoPE positions from each request's actual verify length;
  CUDA Graph padding uses zero-delta ghost positions.
• Add packed-to-dense scatter/gather support for ragged linear attention and
  declare ragged CUDA Graph support per attention backend.
• Preserve captured DSpark draft graph allocations across /flush_cache while
  still clearing logical request, KV, radix, grammar, and metrics state.
• Use an 8-token SPS M bin for compact scheduling and recompute the simulated
  bonus token at the acceptance boundary.
• Add regression coverage for full/reduced vocabularies, configuration
  normalization, tensor-parallel LM heads, Markov shapes, d2t mapping,
  corrected-logit scatter, both checkpoint layouts, CUDA Graphs, ragged mRoPE,
  linear attention, cache flushing, and SPS fitting.

Status

• Bonus-anchor and dense anchor-first checkpoint layouts
• Full-vocabulary DSpark checkpoint regression
• Independent reduced-vocabulary LM head
• Draft-to-target vocabulary mapping and probabilistic corrected logits
• Eager and CUDA Graph draft-width propagation
• TP=1 and TP=2 public RedHat checkpoint inference
• Greedy and non-greedy sampling
• Ragged multimodal mRoPE and linear-attention layouts
• DSpark CUDA Graph behavior after /flush_cache
• Greedy token-level equivalence with target-only decoding
• Same-GPU paired performance benchmark at concurrency 1 and 8

Correctness Tests

On the PR branch rebased onto upstream main on 2026-08-13, the focused suite
reports 212 passed, 54 subtests passed; the reduced-vocabulary file reports
26 passed, 2 subtests passed. Python compilation, formatting, focused Ruff
checks, and git diff --check also pass.

The public-checkpoint validation uses
RedHatAI/Qwen3.6-35B-A3B-NVFP4
as the target and
RedHatAI/Qwen3.6-35B-A3B-speculator.dspark
as the draft. All reported runs used NVIDIA RTX PRO 5000 72GB Blackwell GPUs
(driver 580.126.20); each run used one GPU, and paired performance measurements
reused the same physical GPU. The environment has PyTorch 2.13.0+cu130, Transformers 5.12.1,
and FlashInfer 0.6.17. The official workloads are MT-Bench (80 two-turn
conversations, 160 requests) and the MMSpec test split (600 conversations, 723
turns). Correctness runs compare target-only autoregressive (AR) decoding with
DSpark under greedy decoding at concurrency 1.

Workload	AR repeat	DSpark completed	Exact first-turn AR/DSpark token sequences
MT-Bench	160/160	160/160	8/80
MMSpec	723/723	723/723	18/600

The AR repeat is exact on both workloads, confirming a deterministic
baseline. DSpark completes every request, but target-output equivalence is not
established. A separate 20-prompt eager probe also mismatches on all 20 prompts,
so disabling decode CUDA Graph does not remove the divergence; its cause remains
unresolved.

Acceptance is aggregated from raw server counters. The acceptance length below
includes the target bonus token.

Workload	Draft acceptance rate	Acceptance length including bonus
MT-Bench	32.863%	3.6257
MMSpec	24.956%	2.9925

Separate focused lanes also pass TP=1 and TP=2 inference, greedy and non-greedy
sampling, ragged multimodal CUDA Graph replay, and a post-/flush_cache request.

Fixed-output Performance

The formal fixed-output benchmark consists of three same-GPU paired repeats
with service restart and unmeasured warm-up before each mode:

• MT-Bench first turns with exactly 256 output tokens per request;
• a deterministic 20-per-topic MMSpec subset with exactly 128 output tokens per
  request;
• closed-loop concurrency 1 and 8;
• pair order target -> DSpark, DSpark -> target, target -> DSpark.

All 2,400 measured requests completed with the required output length. Values
below are medians across the three runs; the interval is the range of the three
paired throughput ratios. The raw service logs and result manifests bind these
runs to the reported source commit and execution-environment freeze; strict
reanalyzing after the provenance audit produced an identical summary.

Workload	C	Target tok/s	DSpark tok/s	Paired throughput	p50 E2E target -> DSpark
MT-Bench	1	151.64	262.32	1.730x [1.726, 1.741]	1.687s -> 1.008s
MT-Bench	8	715.30	942.30	1.317x [1.300, 1.327]	2.853s -> 2.076s
MMSpec	1	132.89	178.49	1.343x [1.343, 1.346]	0.937s -> 0.715s
MMSpec	8	467.98	519.74	1.111x [1.110, 1.124]	2.086s -> 1.754s

The gain is not uniform across the latency distribution. MT-Bench C=8 p99 E2E
regresses in all three repeats (paired target/DSpark ratios 0.906-0.920).
MMSpec C=8 has about 9% higher p50 TPOT and mixed p99 E2E results (paired
ratios 0.938-1.092, median 0.961). This supports a stable fixed-output
throughput and p50 E2E improvement under the reported configuration, not a
universal tail-latency improvement.

SSE chunk gaps are not reported as token-level ITL because one speculative
event can expose multiple tokens. TPOT is client-observed and amortized from
the first to last token-bearing event.

Known limitation

Greedy token-level equivalence with target-only decoding has not been established
for the tested public checkpoint. The reported results demonstrate checkpoint
compatibility, request completion, online acceptance, and fixed-output
performance, but do not support a lossless-output claim.

Related work

• Original SGLang DSpark implementation: #30261
• Bonus-anchor checkpoint layout: #30982
• vLLM Speculators checkpoint support: vLLM #47093

Thanks to the authors and reviewers of these changes for the checkpoint-layout
analysis and reference behavior.

Contributors

Engine Architecture Group 5, Engine Infrastructure Department, Xiaohongshu
(RedNote): Su Zhan, Fei Ziyu, Hong Chenchen, Luo Zhaokai, Jin Huayi.

---

CI States

Latest PR Test (Base): ❌ Run #32136288277
Latest PR Test (Extra): ❌ Run #32136288017

[qq1060]

/tag-and-rerun-ci