Kris Bailey

Inference

Shortlist MTP Mostly Worked Around a Bad Kernel Choice

Cutting Qwen's draft vocabulary from 248,320 tokens to 16,000 looked like a big throughput win. After I fixed the kernel selection, most of that win disappeared.

I was playing with serving Qwen3.8-Flash-Next on a DGX Spark using MTP and had what seemed like an obvious idea: stop making the draft head score all 248,320 tokens in Qwen's vocabulary. A draft does not need to be able to produce every token. It just needs to be right often enough that verifying it saves more time than drafting it costs.

My early tests showed more than a 20% throughput improvement from using a 16,000-token shortlist. Unfortunately, most of that improvement was not because shortlisting was a great MTP trick, it was because the full vocabulary projection was using an inefficient kernel for the weird, skinny matrix shapes involved in drafting. Reducing the vocabulary made the inefficient operation a lot smaller by lowering N. Once I fixed the kernel selection, the shortlist still helped, but only a little.

What I was testing

I was serving primitive-ai/Qwen3.8-Flash-Next-NVFP4 on one NVIDIA GB10 with 128 GB of unified memory. The MTP head generates speculative tokens and the full model checks them. A bad draft wastes time, but it does not change the final output because the verifier still checks the complete vocabulary and remains authoritative.

The draft head normally projects a hidden state with width 2,560 across the complete 248,320-token vocabulary. The M dimension is tiny during decoding. It is usually one for a draft step and three for an MTP2 verification step. That leaves us doing a very wide, bandwidth-heavy matrix operation for only a couple rows at a time.

Full vocabulary248,320 tokens
Shortlist16,000 tokens
Vocabulary retained6.4%
Projection weights read1.27 GB to about 82 MB

I changed the draft head to score only tokens that were plausible drafts. If the token the model really wants is not in the list, the draft is less likely to be accepted. That hurts performance, but it cannot make the final answer wrong because the full model still verifies everything.

Building the 16k shortlist

I built the shortlist from 224 million output tokens from the faunix/Qwen3.8-27B-Distillation-40K dataset. I did not just take the 16,000 most common tokens. That drops structural tokens which are not especially common individually but are necessary for code, formatting, tool calls, and punctuation.

The final 16,000-token list combined three groups:

  • Every token with an ID below 8,192.
  • A 4,152-token structural core covering ASCII, whitespace, punctuation, operators, byte tokens, and special tags.
  • The most frequent remaining content tokens until the list reached 16,000 entries.

The structural core mattered because 3,531 of those tokens had IDs above 8,192 and a simple low-ID cutoff would have dropped them. The frequency distribution made 16,000 look like a reasonable size. Reasoning output needed 9,554 distinct tokens to cover 98% of the output I measured. Non-reasoning output was flatter, but 14,182 tokens still covered 95%.

I also tried finding about 100 high-impact missing tokens that I could add back. That went nowhere. The misses were a flat, prompt-specific tail. On a fresh corpus, adding the 100 tokens lowered the miss rate by only 0.19 percentage points.

At first, the shortlist looked REALLY good

In an earlier concurrency-eight sweep, MTP2 with the 16,000-token shortlist reached about 104.0 tok/s. The full-vocabulary version reached 87.4 tok/s. Depending on the exact run and comparison, I was seeing an improvement around 19% to 21.3%.

Earlier exploratory MTP2 result
Draft projectionSingle stream8 concurrent
16k shortlist26.3 tok/s104.0 tok/s
Full vocabulary26.9 tok/s87.4 tok/s

What was actually happening

The normal projection was being handled poorly for the M=1 and M=3 shapes used by this path. The shortlist looked fast mostly because doing an inefficient operation over 16,000 output rows is a lot cheaper than doing that same inefficient operation over 248,320 rows. I lowered N instead of fixing the real problem.

I enabled a CuTe low-M GEMM path for SM121 and added plans for the single-GPU tensor shapes the model uses. The patch covered the dense BF16 decode operations left after the NVFP4 routed experts, including the full LM head. Fixing that gave me about a 6.9% single-stream step-rate improvement by itself.

I also routed the shortlist projection through the same skinny kernel. I needed both paths to use an appropriate kernel before the comparison meant anything. At one point I had the full head using the better kernel while the shortlist still used ordinary F.linear. The shortlist actually lost at single-stream decoding in that setup even though it read WAY fewer weights.

Of course an inefficient operation over 16,000 rows is faster than one over 248,320 rows. What I actually needed to know was whether the shortlist still helped when both versions used an appropriate kernel.

Re-running the comparison fairly

I reran the comparison with the same memory configuration, the skinny GEMM patch, and local argmax reduction enabled for both variants. Single-stream measurements used 256 generated tokens at temperature zero. Concurrency measurements used eight simultaneous streams and the median of three repetitions after warmup.

Final matched-stack results
MTP depthDraft projectionSingle stream8 concurrent
MTP216k shortlist27.6 tok/s109.7 tok/s
MTP2Full vocabulary26.6 tok/s106.9 tok/s
MTP316k shortlist27.7 tok/s106.0 tok/s
MTP3Full vocabulary25.2 tok/s106.4 tok/s

At MTP2, the shortlist was about 4% faster for one stream. At MTP3 it was about 10% faster. The improvements are real, but 4% is noise and 10% is marginal at best. At concurrency eight, everything saturated around 107 to 110 tok/s. The differences were within the roughly 6% boot-to-boot noise I saw as MTP acceptance changed.

The reliable concurrency gains came from fixing the skinny GEMMs, enabling local argmax reduction, and giving the runtime enough memory for concurrency. Changing MTP depth or using the shortlist did not meaningfully move the saturated result.

Making sure the benchmark was clean

The final differences were small enough that another machine hitting the endpoint could have decided the winner. After every run, I checked how many requests reached the container, made sure they all came from the host, and confirmed that the peak running-request count never went above the concurrency I requested.

In the matched grid I sent 58 requests and the container received 58. I also kept the raw concurrency runs. The MTP2 shortlist results were 114, 110, and 110 tok/s. The full-head results were 92, 107, and 110 tok/s. With that much run-to-run movement, claiming the shortlist provided a meaningful improvement to concurrent throughput would be stretching the truth a bit.

I still kept the shortlist enabled

I kept the 16,000-token shortlist enabled in my MTP2 configuration. It gave me a small, consistent single-stream improvement once both paths used the right kernel. It cannot change the final output, and it was cheap to leave in place. I used MTP2 because it tied MTP3 within the noise, had tighter variance, and left more cache blocks available.

I dropped a more complicated dual shortlist scheme that switched between reasoning and non-reasoning vocabularies. The thought was that reasoning token distributions are not likely to match response token distributions. This checkpoint reasons inline and does not emit a clean closing think marker, so I did not have a reliable place to switch lists though, so it was just using a different reasoning-specific shortlist for everything and doing worse than the 16k list. There was no reason to keep the extra complexity.

What I got out of it

The shortlist experiment was still useful because it exposed a bad execution path. If I had stopped after the first benchmark, I would have walked away thinking draft vocabulary size was the main problem.

Changing a matrix dimension can also change kernel selection, memory traffic, launch behavior, and the library path that gets used. A big benchmark improvement does not necessarily mean the high-level idea is good. Sometimes it just makes the wrong kernel do less work.

That is what happened here. Lowering N was a workaround. Selecting and tuning the right low-M kernel fixed the actual problem. After that, the shortlist became what it really was: a small, measurable optimization worth keeping, but nothing that made a big difference.

(c) 2026 Kris Bailey