Keeping the Draft Walk on Chip: One Launch for DSpark’s Markov Head

TL;DR: DSpark drafts a block of γ tokens in one parallel pass, then fixes it up with a small serial walk. Every step of that walk multiplies by a matrix as wide as the vocabulary. We run the whole walk as one persistent kernel that keeps the matrix on chip.

• H100, the walk alone: 331 → 37 µs at γ = 7 and 741 → 62 µs at γ = 16 (bs = 1).
• H100, end to end in SGLang: at temperature 0 and concurrency 1–16, +5.3–5.9% per decoding round at γ = 7 and +10.8–12.1% at γ = 16. Greedy outputs are identical, 605/605.
• Tensor parallelism: splitting the matrix over two GPUs lets it stay on chip in bf16, with no quantization. That wins from batch 8 up.
• B200: a third layer of on-chip memory, TMEM, lets the chip hold the whole bf16 matrix, together with registers and shared memory: 4.5× faster at γ = 16. For a vocabulary too big to fit, we keep 62% on chip and stream the rest: 4.4× faster.
• Other serial heads (Domino, DeLS-Spec, DSpark’s gated and RNN heads) gain 6.6–9.0× over their exact reference paths, from the same kernel template. Domino gains +16.8% per round end to end.

Code: github.com/tie-pilot-qxw/dspark-markov-walk · SGLang PR: sgl-project/sglang#42855

The serial part of a parallel drafter

A block drafter such as DSpark guesses a whole block of tokens in one forward pass. That pass is parallel, so it is cheap. The catch is that every position is guessed blind: position 3 does not know what positions 1 and 2 turned out to be.

DSpark wins part of that back with a small Markov head. After the parallel pass, it walks the block from left to right. At each step it adds a bias that depends on the token it just picked:

prev = anchor
for k in range(γ):
    u       = W1[prev]                 # [256]
    logit_k = bf16(base_k + W2 @ u)    # [V]
    tok_k   = argmax(logit_k)          # or a sample of softmax(logit_k / T)
    prev    = tok_k

The sketch omits one detail: SGLang’s default walk also rounds W2 @ u to bf16 before the add.

Each step needs the token from the step before, so the steps cannot run in parallel.

SGLang’s default walk runs them as γ separate steps. Each step is a full-vocabulary GEMV, then an add, an argmax and a sampler. W2 is V × 256 in bf16: 78 MB at Qwen3’s 151,936-token vocabulary, and 127 MB at 248,320. So every round reads the same tens of megabytes from HBM γ times, just to pick one token per step.

That adds up. On an H100 at bs = 1, the walk costs 331 µs per round at γ = 7 and 741 µs at γ = 16. That is 5.8% and 12% of a whole Qwen3-4B speculative-decoding round (5.72 ms and 6.06 ms). The big target model verifies the draft in parallel. This little serial tail is what is left.

The idea: make the walk resident

Two facts about the walk point to the design:

• What changes between steps is tiny: one token id.
• What each step reads never changes: the same W2, every step, every round.

Persistent RNN kernels used exactly this shape ten years ago. They kept the recurrent weights on chip and kept one kernel alive across all time steps (Diamos et al., 2016). We apply the same idea to the draft walk.

First, the matrix has to fit. An H100 has 65.4 MB of on-chip storage across its 132 SMs (registers plus shared memory). bf16 W2 is 78 MB, so it does not fit. An int8 W2 with one fp32 scale per row is 38.9 MB, and that does. (Why int8 rather than fp8 comes later.)

So the kernel works like this:

• One CTA per SM, one launch per round. Each CTA owns a slice of vocabulary rows, 1,152 rows at V = 151,936. It loads its slice of int8 W2 into registers and shared memory once. Then it runs all γ steps without leaving the kernel. In the bs ≤ 4 kernels, W2 is read from memory once per round instead of γ times. The larger-batch kernel keeps most of it resident and streams the rest (see below).
• Exact integer math. u = W1[prev] is stored as two int8 planes, u ≈ s_hi·q_hi + s_lo·q_lo with s_lo = s_hi/256. Both dot products run as exact int32 sums on the int8 tensor cores, so nothing is dequantized inside the loop.
• One global argmax per step, using an atomic rendezvous instead of grid.sync().
  • Each CTA reduces its rows to one 64-bit key: the order-preserving bits of the fp32 logit, then the inverted row index, so ties go to the smallest token id.
  • Every CTA posts its key to one 16-byte {key, count} pair, with an atomic max on the key and an atomic add on the count. Then it polls the pair until the count says every CTA has arrived.
  • There is no memset between CUDA-graph replays, because the per-round state alternates by round parity.
• Sampling inside the kernel.
  • Philox noise is keyed by (seed, row, step, round counter), so every graph replay draws fresh noise.
  • Small batches use per-row Gumbel-max. Large batches use an inverse CDF inside 8-row groups and Gumbel-max across groups.
  • Temperatures are read per request on the device, so one captured graph serves greedy, sampling and mixed batches.
• Lossless verification. The logits are rounded to bf16 once. The token is picked from exactly those values, and the same bits are written out for SGLang’s verifier, which rebuilds the draft distribution from them. So the rejection sampler sees the distribution the draft was really drawn from, up to the sampler’s floating-point error.

Three kernels cover the batch sizes:

Results on H100

All numbers in this section: H100 80GB SXM, CUDA 13.1.

The walk alone

CUDA-graph replay, cold L2, µs per call. The baseline is SGLang’s default walk.

At V = 248,320 the int8 table is 63.6 MB. That is just under the 65.4 MB total, but each SM also needs room for working buffers, so it no longer fits. The wgmma kernel keeps as much as fits in shared memory and streams the rest from L2 at every step. That is why bs = 1 costs 211 µs there (γ = 15), against 62 µs on the Qwen3 vocabulary (γ = 16).

End to end in SGLang

How we measured:

• Workload: Qwen3-4B target, 605 prompts from 9 datasets (gsm8k, math500, aime25, humaneval, mbpp, livecodebench, mt-bench, alpaca, arena-hard-v2), up to 2,048 new tokens.
• Metric: cost per round from the scheduler’s decode log (accept len / gen throughput), paired by batch size.
• Protocol: two interleaved passes per arm, with the order reversed in the second pass, and default server flags.

Per-round throughput gain:

• At bs = 1 a round goes from 5.721 to 5.424 ms at γ = 7, and from 6.059 to 5.467 ms at γ = 16.
• A bigger target: with Qwen3-14B and deepseek-ai/dspark_qwen3_14b_block7, the gain is +2.2% at concurrency 1. The kernel saves the same 0.29 ms per round as with the 4B target, but a 14B round takes 13.7 ms, so the share is smaller.

This is the general pattern: the kernel saves a fixed amount of time per round, so the percentage depends on how long the rest of the round takes. Longer blocks have longer walks and gain more. Bigger targets have longer rounds and gain less.

Correctness

• Greedy outputs are identical. With --enable-deterministic-inference at temperature 0, the generated text matches the default walk on 605/605 prompts for both drafters, at concurrency 1, 16 and 64.
• Acceptance length is tied. Change in acceptance length, ours vs default, with 95% prompt-cluster bootstrap intervals:
  • γ = 7: −0.09% [−0.36, +0.17] at temperature 0, and +0.11% [−0.39, +0.62] at temperature 1.
  • γ = 16: +0.09% [−0.34, +0.50] at temperature 0, and +0.47% [−0.20, +1.07] at temperature 1.
• The drafts themselves can differ. W2 is int8, with 0.58% relative quantization error (see below), and our kernel rounds once where the default walk rounds twice. So near-ties can go the other way. That only changes how many draft tokens get accepted. With greedy decoding, the outputs matched on every tested prompt. With sampling, rejection sampling keeps the target’s output distribution, though the same seed need not give the same text.

Why int8 and not fp8?

fp8 would fit just as well: it is also one byte per weight, and on H100 the fp8 and int8 tensor cores run at the same rate. We still use int8, for two reasons.

int8 has about 4.6× lower error than e4m3. We quantized W2 of the two public 4B drafters used above, each with one scale per row. The two drafters agree to the digits shown:

The errors are relative Frobenius norms. The bias is W2·u for u = W1 of 4,096 random tokens. The fp8 rows get the same free fp32 scale as int8.

The reason is how each format spends its 8 bits:

• After scaling, each row of W2 is a bell-shaped spread of values. int8 covers it with 255 evenly spaced levels.
• e4m3 spends 4 bits on the exponent, so it has only 3 mantissa bits. The large entries, which dominate the dot product, get steps of up to 1/16 of their size.
• fp8’s extra range only helps the tiny entries, and those barely move the sum.

int8 sums are exact. int8 × int8 products add up exactly in int32, so the result does not depend on summation order. That is why all three of our kernels, with their different tilings, match one integer reference bit for bit.

The two-plane trick for u relies on this too. It gives u about 16 bits of precision for the price of two int8 MMAs, and the integer sums stay exact. With fp8 you either put u in fp8 as well, with 3 mantissa bits, or convert W2 to fp16 inside the loop. Our first fp8 prototype did the latter.

In an offline acceptance check on one of our own 4B drafters, int8 showed no measurable loss. e4m3 and e5m2 lost very little, but measurably.

This choice depends on the GPU. On H100 and B200, int8 and fp8 tensor cores run at the same rate. Blackwell Ultra (B300) cuts dense int8 throughput to a small fraction of fp8, so on B300 the walk will probably have to use fp8 and live with the larger error.

Tensor parallelism: bf16 with no quantization

int8 is what lets W2 fit on one H100. With two GPUs there is another option: split W2 by vocabulary. Each GPU then holds half of the bf16 table, 38.9 MB, which is the same size as the whole int8 table. So the walk can stay on chip with the head’s own bf16 weights and fp32 accumulation.

The price is one more exchange. Each GPU finds the best token in its own half, and the two halves have to agree before the next step can start. We do this inside the kernel through peer-mapped mailboxes. The cross-GPU hop costs about 1.1–1.6 µs per step.

Round time at γ = 7, device clock, µs, greedy | T = 1:

• Small batches: int8 on one GPU wins. At bs = 1 the hop costs ~8 µs per round, more than the bf16 version saves.
• From batch 8 up: bf16 on two GPUs wins, by 5–24% greedy and 11–28% at T = 1.

B200: the bf16 table fits on chip

Blackwell adds a third kind of on-chip memory. Each SM gets 256 KB of tensor memory (TMEM) next to its registers and shared memory, and only the new tcgen05 tensor-core instructions can read it. Across 148 SMs that gives:

The bf16 W2 for Qwen3’s vocabulary is 77.8 MB. That does not fit on an H100, but it does fit on a B200. So on B200 we keep the whole table in bf16 with no quantization:

• 147 CTAs × 1,040 rows;
• 46% of the rows sit in TMEM, 22% in shared memory, and 32% in registers;
• tcgen05 reads the TMEM part, and mma.sync handles the rest.

It uses the same rounding points as SGLang’s walk: fp32 accumulation, the bias rounded to bf16, then the sum rounded to bf16. On our test data the drafts and corrected logits matched the default walk bit for bit. The accumulation order differs, though, so that is not a guarantee.

A caveat about the numbers. We ran on rented B200s (Modal). We could not lock clocks, and our kernel and the baseline ran in different containers rather than interleaved. Treat the B200 ratios as approximate.

V = 151,936: the whole bf16 table on chip

Walk alone, µs per call, against the fastest SGLang walk for each cell:

The ratios are smaller than on H100. That is because the baseline gets faster too: the whole bf16 table fits in B200’s L2, so the default walk drops from 44.0 µs per step on H100 to 19.2 µs on B200. Our int8 kernel, recompiled for B200, barely moves: 2.73 µs per step on H100, 2.82 on B200.

Keeping the table on chip still matters. A probe that reads the whole bf16 table from L2 every step, inside a single launch, needs 6.76 µs per step. The resident kernel needs 3.34 µs, about 2× faster.

Where one step goes, at bs = 1:

• the GEMV takes 1.43 µs;
• the cross-SM argmax exchange takes about 1.3 µs. Even with only 8 CTAs it takes 0.97 µs, so most of it is a fixed synchronization cost, not contention. Every alternative we tried was slower: grid.sync 2.17 µs, mailboxes 1.67–2.10 µs.

Batching is almost free for the tensor cores, because request n is column n of the MMA’s N = 8 operand. What grows with batch size is everything around the GEMV.

V = 248,320: 62% on chip, the rest streamed

At 248,320 tokens, the bf16 table is 127 MB, and even B200 can only keep 78.8 MB (62%) on chip. Every step, the other 48.5 MB has to come from L2:

• One thread of warp 0 copies the streamed rows with TMA into a 64 KB shared-memory ring, 16 KB at a time.
• The same thread issues tcgen05 MMAs straight from shared memory.
• Warps 1–7 work through the resident tiles in parallel.
• W2 does not depend on the token, so the ring keeps refilling across steps. The next step’s first stages are already in flight during the current step’s argmax.

bs = 1, µs per call:

One lesson from this kernel: issue order matters more than you would expect.

• In our first version, each step began with the resident MMAs, then the streamed ones.
• tcgen05.commit waits for all MMAs that thread issued before it. So no ring slot could be freed, and no refill could be issued, until all 64 resident MMAs had finished. Meanwhile the copy engine sat idle.
• Consuming the stages that had already arrived first, and slipping 4 resident MMAs in after each refill, cut the step from 6.86 to 6.27 µs (−8.6%, same container).

Two more things we measured about B200’s tensor cores:

• A shared-memory operand holds the tensor pipe for the whole read. At M = 128, N = 8, one MMA with A in shared memory takes 39 cycles per 4 KB of A. With A in TMEM it takes about 8 cycles. While the slow one runs, nothing else can use the tensor pipe, including mma.sync from other warps.
• tcgen05.cp (shared memory → TMEM) is no way around this. It runs at 64 cycles per 4 KB and holds the pipe too. In our kernels, TMEM paid off only for data that is written once and reused, like the resident table. It did not help for data that is used once, like the streamed part.

End to end on B200

We have not run an end-to-end A/B on B200 yet. We have only measured how much of a round the default walk takes. With Qwen3-4B and the γ = 16 drafter at bs = 1, it is about 6.4% of a B200 round, against 11% on H100.

Beyond DSpark: other serial heads

DSpark’s Markov head is not the only serial head. Other drafters correct their block with a small recurrent cell that also ends in a vocabulary-wide matrix:

• Domino: a GRU cell that also reads the drafter’s hidden states.
• DeLS-Spec: a GRU.
• DSpark’s own gated and RNN heads.

They all have the same shape: a tiny state, and a big matrix that every step reads. So we turned the walk into a template, walk_kernel<Cell>, where each head supplies only its state update.

The question for each head is whether its recurrent matrix fits in the spare shared memory of every CTA:

• The gated head’s 256 × 256 matrix does fit in int8, so every CTA computes the update redundantly and needs no extra exchange.
• The RNN head and the GRUs do not fit. Their state is spread across CTAs and gathered once or twice per step, at about 1.7 µs per gather.

SGLang runs Domino in one of two modes, and we compare against both:

• Exact (pool 0) scores the full vocabulary at every step.
• Default (pool 2,048) first takes the top 2,048 tokens of the base logits, as one candidate pool shared by the whole block. Every step then scores only those candidates.
  • The steps are still serial, but each one is about 74× narrower.
  • It is approximate: a token outside the pool can never be picked.
  • Each step costs about 35 µs, against 73 µs in exact mode. The narrower GEMV does not remove the fixed cost of the small kernels each step launches.

Our kernel scores the full vocabulary, like the exact mode, at 7.2 µs per step.

H100, bs = 1, µs per round, against each head’s own reference implementation. “Head only” means the round minus the lm_head GEMM, which both sides run identically:

• DFlash2 is the exception. It already restricts its serial part to a 16-token candidate lattice, which costs only ~0.6 µs per step. Our full-vocabulary walk costs ~3.2 µs per step, so it only wins on the fixed per-round cost, and the two cross near γ ≈ 21.
• No public checkpoint uses DSpark’s gated or RNN heads, so those two rows use synthetic weights and measure time only.

Domino end to end. We hooked the kernel into SGLang’s Domino path at bs = 1 (Qwen3-4B or Qwen3-8B with the released Huang2020 Domino b16 drafters, temperature 0, concurrency 1, the same 605 prompts):

• The saving per round equals the walk saving. For example, at 4B against the exact path, the round saves 0.96 ms and the walk microbenchmark saves 0.945 ms.
• Greedy outputs are identical. With deterministic inference, the text is 605/605 identical to SGLang’s exact path, at both sizes.
• Acceptance is tied: within 0.06%.

Limitations

• The upstream PR covers only part of this. It targets H100 (sm_90a), vanilla Markov heads (rank 256), and bs ≤ 64. Under TP > 1 it falls back to the default walk.
• The rest is not released yet. The B200 kernels and the serial-head template are not public. Their end-to-end numbers are limited to what is reported above.
• The gain shrinks as the target grows. The saving per round is fixed, so a longer round means a smaller share.

Code

• Repository: github.com/tie-pilot-qxw/dspark-markov-walk. It contains the H100 kernels, the SGLang reference hook, the bf16 tensor-parallel kernels (experimental/bf16_tp/), tests, and the end-to-end A/B harness.
• SGLang PR: sgl-project/sglang#42855, “[Spec][DSpark] Opt-in fused int8 Markov walk kernels (Hopper)”. It is opt-in behind SGLANG_DSPARK_OPT_INT8_MARKOV_WALK=1.