Parallel speculative drafters use a lightweight serial head to recover token dependencies within a block generated in parallel. Each head step updates only a small state, yet projects that state through the same vocabulary-sized matrix. Step-wise execution repeatedly transfers this large, immutable projection, making the head a meaningful source of latency despite its low arithmetic cost. We present Weights That Stay, an execution design that keeps projection shards on chip for the entire serial draft walk. A persistent megakernel couples weight residency with GPU-side token selection and state advancement, reusing resident weights while exchanging only compact state between steps. The design supports full or partial residency and accommodates both Markov and recurrent heads on H100 and B200. Integrated into SGLang, Weights That Stay improves throughput for Markov drafting on Qwen3-4B and Qwen3-8B, with gains of up to 12.3%. At the head level, the DSpark walk is up to 11.6× faster than the fastest measured SGLang reference path. These results show that aligning weight residency with the full draft walk translates cross-step reuse into end-to-end serving gains.
Figures & tables
Figure 1: Walk-resident execution across GPU memory tiers. Each SM retains a shard of W throughout the walk, exchanging only tokens or recurrent state. Weights are int8 on H100 and bf16 on B200; optional loads stream nonresident tiles.
Figure 2: Serial refinement with a shared vocabulary projection ( V=12 , r=4 , three steps). Each selected token advances the next step. Base-logit addition and explicit state-update nodes are omitted; the full recurrence is given in Equations 1 – 4 .
Figure 3: Memory footprint and latency cost of serial draft heads. (a) Rank-256 projection size versus nominal GPU capacities (MB). (b) One SGLang step at B=1 , V=248 k under CUDA-graph replay. (c) Head microbenchmark latency divided by measured serving-round latency at C=1 , T=0 ; DSpark uses the default folded path, and Domino uses the full-vocabulary head.
Figure 4: Whole-walk fusion removes launch boundaries, while residency removes repeated fetches of W from L2 or HBM ( V =151,936, B =1, greedy). Timings are fitted from the complete-call measurements in Section 5.4 . The SGLang row is truncated after one step, and the within-step segments are illustrative. The resident kernel’s entry-cost excess over streaming is drawn as a one-time shard load.
Figure 5: Operand storage and data paths for GPU matrix instructions ( NVIDIA Corporation, 2026 ) . Thin arrows denote copies; thick arrows denote operand and result access. mma.sync uses registers; wgmma reads A from RF/SMEM and B from SMEM, accumulating in RF; tcgen05.mma reads A from TMEM/SMEM and B from SMEM, accumulating in TMEM.
Figure 6: Cumulative distributions of normalized weight magnitudes (band: two released 4B heads) and representable positive values (quantization levels). Left of the dashed line, E4M3 has finer spacing than int8 but places 75% of its levels in a region containing only 12–17% of the weights.
Figure 7: Vocabulary projection using int8 tensor cores on H100. Two int8 vectors represent the latent u and are multiplied by W together. Their scaled results are summed, added to the base logits, and rounded once to bf16.
Figure 8: Cross-CTA token selection. Each CTA scores its rows ❶, packs its local winner ❷, atomically updates the maximum key and arrival count ❸, waits for all CTAs ❹, and uses the common winner ❺. The {key, count} observation uses the hardware assumption stated in the text.
Figure 9: End-to-end throughput gains for Qwen3-4B and Qwen3-8B over default SGLang: block-16 Markov drafting on H100, using all 3,030 prompts.
Figure 10: Markov-head latency across batch sizes, walk lengths, and vocabularies on H100. Top: T=0 ; bottom: T=1 . Latency uses a log scale; labels give speedup over the fastest measured SGLang path. The 152k path is fully resident at B≤4 and partially resident above; the 248k path is always partially resident.
Figure 11: Separating fusion and residency gains on H100 ( B=1 , V=152 k, greedy). Points measure complete calls; least-squares fits are labeled by their slopes in µs/step. The axis break spans 210–260 µs.
Figure 12: RF/SMEM placement tradeoff on H100 ( B=1 , greedy), with 72 weight tiles per CTA. Points are per-step costs fitted from complete calls at γ=1,7,16 . Shading marks register spills.
Figure 13: Latency and per-step breakdown across serial heads on H100 ( B=1 , greedy). (a) Complete calls: Markov runs 16 steps, recurrent heads 15; labels give speedup over each head’s reference. *DSpark RNN uses random recurrent weights. (b) Disjoint per-step timings, excluding entry cost. The fixed-latent probe repeats projection and selection with an unchanged latent.
Figure 14: B200 bf16 head latency under full (152k) and partial (248k) residency ( B=1 , greedy).
Figure 15: Effect of int8 execution on Qwen3-4B draft acceptance length relative to default SGLang. Estimates pool runs across concurrency settings; lines give 95% prompt-cluster bootstrap intervals on 605 prompts.
LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.
Speculative decoding losslessly accelerates large language model inference by having a lightweight draft model predict future tokens for verification by the target model. Recent block diffusion drafters further reduce drafting latency by predicting multiple tokens in parallel. However, existing block drafters project target hidden states at every input position into a separate drafter-side KV cache, incurring per-request memory and KV-write overhead that grow with concurrency; directly reusing target KVs in place removes this cache but fails to sustain draft quality throughout the block. We propose a hybrid target-context injection method that complements direct target KV reuse with target hidden states only at the last input position, requiring no separate drafter-side KV cache. Building on this design, we propose H-Spec, a hybrid Mamba-attention parallel drafter that consumes the two target-context sources through complementary modules. Mamba modules are initialized with projected last-token target hidden states, while attention modules reuse target KVs in place. Despite its recurrent formulation, Mamba's parallel scan allows H-Spec to preserve block-parallel drafting. Across three target models and diverse tasks, H-Spec improves over the best baseline by 5.0--13.3% in mean accepted length and 5.3--12.6% in batch-size-1 inter-token latency speedup. Under concurrent serving, H-Spec consistently achieves higher throughput while maintaining lower KV cache utilization than baselines across evaluated concurrency levels.
Parallel block drafting accelerates speculative decoding by predicting multiple future positions in a single backbone pass, but on-device deployment constrains both drafter capacity and context storage. Under tight memory budgets, drafters must model long-range dependencies and fine-grained draft-position interactions without the growing KV-cache overhead of full-context attention. We propose HyFlash, a compact parallel drafter that introduces a prefix-memory prior to separate long-range context modeling from fine-grained draft-position interactions with limited capacity. Local attention captures recent context and draft-position interactions, while recurrent memory compresses historical target features into fixed-size states. This inductive bias enables effective local and long-range modeling with bounded draft-side memory. Experiments show that HyFlash achieves higher acceptance than full-context attention with substantially less KV storage and with larger gains for shallow drafters and long contexts.
DSpark-style parallel drafters have made speculative decoding highly effective, yet their draft phase remains serialized on the critical path of every round. Parallel speculative decoding (PSD) overlaps drafting with verification, yet existing methods must guess the accepted prefix and bonus token in advance: a wrong guess reverts the whole batch to serial drafting. We present DPara, a PSD framework that reuses effective parallel drafters yet guarantees backbone--verification overlap in every round, thereby eliminating this probabilistic fallback altogether. While the target verifies, DPara's diffusion backbone precomputes draft representations for every acceptance boundary with the bonus left unspecified; a lightweight autoregressive head then combines the revealed verification outcome with the matching precomputed representation to emit the next round's draft tokens almost instantly---fully parallelizing the dominant backbone forward with verification and leaving only the negligible head cost serial. Experiments on Qwen3-8B and Qwen3-14B across seven math, coding, and chat benchmarks show that DPara achieves average speedups of 3.21× and 3.52× over autoregressive decoding, surpassing the strongest serial and parallel speculative decoding baselines alike.
Fuliang Liu, Xue Li, Kun Qian +4
State Key Laboratory of Novel Software Technology, Nanjing University · Alibaba Group