← All notes and experiments

Making a ternary reasoning model fast in vLLM

Earlier development stage. Implementation speeds here describe this stage of the work. See the latest published comparison for the current checkpoint.

One-token generation was now much faster. Could the model advance several tokens in one step? That would help concurrent users. It could also help one user, if a small model proposed tokens for the main model to check together.

The work had two steps: make small batches efficient, then make drafting and checking worthwhile.

The first lifted eight-request throughput from 204.1 to 350.3 tok/s. The speculative-decoding experiments then roughly doubled single-request speed, before I extended the gain to longer contexts.

Measurement note. Most comparisons below use an RTX 3090 Ti at 300 W. The draft-length sweep uses an RTX 3090 at 275 W and is labelled separately. Each before/after pair stays on the same card and power limit. These are successive implementation stages, so rates from different tables should be read with their own workloads.

First, make small batches efficient

The batch-eight puzzle

Part 5's batch-eight trace contained a number I'd skimmed past. With eight concurrent requests, the integer projections took about 39 ms per decode step. With one request they take about 10 ms.

A batch-eight step reads exactly the same 6.8 GB of weights as a batch-one step; the activations are a few kilobytes. If decode were bandwidth-bound, eight tokens should cost almost the same as one. Instead they cost four times as much. The batch-eight projections were moving about 200 GB/s, a fifth of the card.

The integer program from Part 3 shared weight reads across tokens, but still did separate arithmetic for each token. It also had to keep many more partial results in fast temporary storage. Eight tokens meant eight times the integer calculations, even though the weight data was shared.

The temporary array had shape (tokens × rows × groups × 32). The eight-token case needed 16,384 32-bit values per program, about 128 registers per thread for that array alone. That storage pressure limited how much work could run at once.

Put the GPU's matrix hardware to work

Ampere's tensor cores multiply small tables of numbers together. They offer much more integer arithmetic throughput than the four-pair DP4A instruction, but only for particular table sizes. Eight tokens did not fit the original arrangement efficiently.

The fix was to put the weight rows on one side of the multiplication and the token inputs on the other. That fits the hardware better. Marlin and Machete use this kind of rearrangement too. MMA is the instruction's name for matrix multiplication followed by addition.

I also arranged the input numbers in the same order as the packed weight codes. That let the GPU extract the codes with fewer operations. The integer sums stayed the same; rounding could still differ when the scaled sums were combined.

On inputs saved from the real model, the Triton version of this kernel cut the batch-eight projection pass from 35.5 to 18.5 ms. In full serving with eight concurrent requests:

Request performance

Single-request throughput was unchanged.

Reading each weight once was only part of the answer. The kernel also needed enough arithmetic throughput and room for its temporary values.

These rates add the output from eight requests. The same saving could help one user if those rows instead held draft tokens from their request.

Then, use those batches to check draft tokens

One verification step can advance several tokens. Its benefit depends on how many tokens are accepted and the combined cost of drafting and verification.

Efficient small batches improve concurrent serving. They can also reduce the wait for an individual user through speculative decoding.

A small draft model proposes k tokens. The main model checks them together and keeps the longest matching prefix. With greedy decoding, this gives the main model’s own output, up to floating-point near-ties.

The check processes k + 1 token rows from one request: M = 4 at k = 3, or M = 6 at k = 5. If checking four tokens costs about the same as generating one, accepted draft tokens become cheap.

Consider a simple example. Plain generation takes 12 ms per token. Drafting and checking take 24 ms and advance by three tokens on average: 8 ms per token, a 1.5× speedup.

If that same step advances by only 1.5 tokens, it costs 16 ms per token and loses. Both acceptance and the cost of a complete step matter.

That made the kernel work above a prerequisite. In the production path at the time, M = 3, 6, 7 and 16 fell back to floating-point kernels with no weight reuse, and cost 80–96 ms per projection pass. Tuned for every M from 2 to 16, the tensor-core kernel costs an almost flat 21–24 ms whatever M is.

The helper, or draft model, is DFlash2, trained by ProCreations for Bonsai. It has five layers and 1.9 billion parameters. It uses intermediate results from five layers of the main model to suggest up to eight tokens together, then ranks its guesses. The DFlash paper describes the method.

Make checking the suggestions cheap

A speculative step yielded about 2.9 tokens but cost 32 ms, against 12.7 ms for a normal step. Verification at M = 4 was the problem: about 2× a single-token step.

Nsight Compute showed why the Triton tensor-core kernel was stuck at 40–50% of bandwidth:

GPU profile

Triton's version repeatedly copied the extracted weights into shared memory, a small storage area shared by a group of GPU workers. The workers then had to wait for one another before continuing. Changing how the work was divided stopped helping at around 500 GB/s.

I wrote a CUDA version that loaded the weights directly into each worker's fast temporary storage. It could issue the matrix instruction without that extra shared copy or wait. The instruction, mma.sync.m16n8k32, has an eight-column side that fits the token group.

At M = 4 the projections now cost 1.08–1.18× a single-token step. On real activations the CUDA kernel beats the Triton one by 1.05–1.43× at M = 2–8, and even beats DP4A at M = 2.

One more cost was hiding outside the five "integer" projection shapes. The floating 5120 × 6144 output projection at M ≤ 4 used a GEMV that re-reads all its weights for every token: 169 µs per layer at M = 4, 10.8 ms of every verification step.

An 8-row tiled kernel brings that to 48 µs. Routing it through the INT8 kernel too brings it to 19 µs, but that changes these outputs by about 0.7%. It's the same Q8 treatment production already applies to this projection at M = 1, so I made it a separate opt-in.

Choose a draft length

Here k is how many tokens the helper proposes. Acceptance is the fraction of proposed tokens the main model accepts. Tokens per step counts how far generation advances after each complete round of drafting and checking.

Single request, 90-token prompt, 512 greedy tokens, three repeats.

RTX 3090 (275 W), BF16 draft:

Request performance

k = 5 gave the best rate in this sweep. At k = 7, verifying 8 tokens costs more and acceptance falls faster than tokens per step rises.

Make the helper model smaller

Even after speeding up the checks, the helper still read 3.85 GB of weights per step. I tried storing its weights in FP8, an 8-bit floating-point format, instead of 16-bit BF16. vLLM can do this using Marlin on Ampere, but startup crashed. DFlash's shortcut for preparing the helper's saved context assumed the weights were in their original format. Compression changed their storage, breaking that shortcut.

The fix in the context projection, about 30 lines in vllm-prism, asks each layer to calculate its own saved context when the weights are compressed. It produces the same arrangement of results as the original combined calculation; the existing unit test checks that the two match.

Here FP8 describes the draft's weight storage. Ampere has no native FP8 tensor-core arithmetic; the Marlin path makes those compressed weights usable on the card. The benefit is fewer bytes to store and read, which also matters for the 16 GB deployment budget.

RTX 3090 Ti (300 W), k = 5 + INT8 output:

Request performance

Acceptance was almost unchanged in this comparison, while step time fell from 22.1 to 19.9 ms. That is consistent with reducing the draft's weight-read cost. Acceptance on this workload is only one check of the compressed draft.

What this checkpoint establishes

• Text differs from plain decoding at near-ties. Verification scores tokens at M = k + 1 rather than M = 1, so floating-point order differs, and greedy decoding flips where two tokens are almost tied. In these runs every speculative configuration diverged at the same spot as other batch-size changes (around character 300). Each configuration was deterministic across its own repeats.
• Under heavy load, speculation still loses. At eight concurrent requests it is 0.76× plain decoding; see below.
• Opt-in throughout. The small-batch kernels (GGUF_PQ2_SMALL_MMQ=1) and especially the INT8 output projection (GGUF_PQ2_SMALL_MMQ_OUTPUT=1) are not defaults.

Keep the gain as the workload changes

More concurrent users

The prediction changes with concurrency. At k = 5, one request needs six target query rows for verification. Four simultaneous requests need 24 rows. If the fast kernel covers only 16, adding users can cross a dispatch boundary and erase the single-user gain.

Speculation's gain depends on staying inside the fast kernels' M range. With k = 5 and the FP8 draft, identical 512-token requests gave:

Request performance

Acceptance barely moves (46–52%). What changes is M: six verification tokens per request means M = 24 at four requests and M = 48 at eight. That's past the M ≤ 16 the tensor-core kernels covered, so verification fell back to the old floating GEMM.

Making the Triton kernel's token tile a parameter (16, 32 or 64 rows) fixed most of it: 2.1–5.3× faster projections at M = 24–64. A fresh paired comparison gave:

Request performance

Four users now gain. Eight still lose, mainly to the floating output projection above M = 16 and to the draft's own cost at batch 8. Plain decoding, meanwhile, reaches about 400 tok/s aggregate at eight requests.

Longer prompts

The same comparison with natural-text prompts (not a repeated seed, which a drafter could copy), one request:

Request performance

Acceptance holds up, and is even higher than on the short prompt. But a speculative step grew by about 33 ms from 4K to 32K, while a plain step grew by about 2.4 ms.

Nsight put 34.5 of those 37 ms on one kernel: FA2's attention with a grid of 1×1×24. That's 24 thread blocks, one per query head, each scanning all 32K keys on an 84-SM GPU. FA2 splits long contexts across blocks for single-token decode, but its paged varlen path refuses to split when there's more than one query token, which is exactly the verification case.

The verification attention implementation uses only existing kernels. Verification attention for n tokens is the shared prefix (all keys but the last n) plus a small causal tail.

• The prefix runs through FA2's single-token split path, with the n tokens packed into the head dimension so the grouped-query mapping still lines up.
• The tail uses a small direct computation.
• The two halves merge exactly by log-sum-exp.

Per layer at 32K that's 2,436 → 281 µs, at the same error as FA2:

Request performance

A speculative step now costs about 22–24 ms at every context length up to 32K. It also holds further out: at 64K DFlash runs at 114.0 tok/s against 55.8 plain (2.04×), and at 128K at 101.4 against 42.5 (2.39×). There a speculative step costs 32 ms, against 23.5 ms per plain token. The ceiling is memory: with the draft's own KV cache, 128K needs 23.1 of the card's 24 GB.

Check whether the answers remain useful

I ran a paired comparison on one GPU: plain decoding against DFlash k = 5 with the full opt-in stack (tensor-core verification and the INT8 output projection), on the same fixed question sets as part 6.

Quality check

The GSM8K score is unchanged, although the lost/gained counts show that the two modes did not answer exactly the same questions correctly. On GPQA, DFlash scored 68 against 75 for plain decoding; this test did not detect a statistically significant difference.

About 60% of answers reached the 2,048-token budget before finishing, making the score sensitive to completion status. Both modes scored 62 on the 66 questions they both finished, but that selected subset does not establish the cause of the overall gap.

A larger evaluation with a realistic reasoning budget is still needed to assess quality.

What I took from this

Speculation paid off after the runtime could verify several tokens efficiently. Acceptance alone did not predict the result: dispatch boundaries and attention parallelism could erase the gain even when the helper kept proposing useful tokens. For deployment, the draft also has to leave enough memory for the conversation.

1. Divide bytes by time before believing "memory-bound". The batch-eight path was reading the same bytes as batch one at a fifth of the speed. A one-line calculation would have flagged it in part 5.
2. The fastest kernel isn't always tunable from Triton. Every Triton knob plateaued about 25% below where the hand-written MMA kernel landed. When the operand layout forces shared-memory staging, occupancy, not tile size, is the limit.
3. Look outside the shapes you optimised. The output projection's per-token GEMV cost more than all the CUDA savings put together, and it wasn't in any of my screens because it wasn't an "integer" shape.
4. Speculative decoding is a kernel problem as much as a model problem. On the same GPU with the same BF16 draft at k = 3, the gain went from +15% to +59% purely by making verification cheap. The longer drafts and the FP8 draft only paid off after that.

The faster path now needs space for both a target model and a helper. Part 8 follows that memory cost: first recovering wasted allocation, then compressing the cache without giving back the speed.

Later measurement. These figures describe the kernels and small prompt samples of this stage. Chapter 9 repeats speculation on the later implementation with twelve prompts per length: about 210 tok/s at 1K and 181–194 at 4K. The component-only small-batch evidence is in tests 21–23.

Sources

• csrc/pq2/pq2_mma_small.cu: Small-batch CUDA matrix calculation, with its shared MMA helper; Triton implementation and dispatch.
• DFlash context projection: the quantised-draft fix in my vLLM fork.
• DFlash2 drafter model card and DFlash paper.
• NVIDIA PTX matrix multiply instruction reference.