← 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.

Faster generation made the runtime more useful, but the draft model took memory away from the conversation cache. I returned to the deployment question: how much context can fit when everything shares 16 GB?

The KV cache saves information about earlier tokens. Keys help decide which information matters; values hold the information to retrieve. The current token supplies the query.

In the full-attention layers, a longer conversation needs a larger cache.

The rough 16 GB budget allowed 7.1 GiB for weights, 1.3 GiB for runtime overhead and, when enabled, 2.1 GiB for the FP8 DFlash drafter. The model’s 16 full-attention layers each use 4 KV heads of dimension 256: 64 KiB per token of history in BF16.

That leaves roughly 110K token slots without speculation, or 60K with it—short of the model’s 262K context. These are estimates; actual allocations on the target card still need measuring.

The 7.1 GiB figure is the allowance used in this estimate. Keep it separate from the approximately 7.25 GB packed model size in part 1 and the 6.80 GB of active projection weights read per token in part 2. Here, GB means 10⁹ bytes and GiB means 2³⁰ bytes.

Reducing cache precision offered a way to retain more context. This chapter follows three parts of that work:

• wasted cache capacity recovered before changing its number format;
• a bug that corrupted output in the stock INT8 path;
• a new attention kernel modelled on llama.cpp's.

On the development card, the end state is 1.8–1.9× the KV-cache capacity, similar first-token latency to BF16 FA2 and faster plain decode at every measured length. The quality checks below found small logit changes and no statistically detectable task-score difference in the tested cohorts. Capacity counts describe the cache pool; a single request is also bounded by the model's supported context length.

Measurement note. Measurements use one RTX 3090 Ti at 300 W with a 12 GiB KV-cache budget unless stated. The 16 GB deployment capacities are estimates, not measurements on the A4000. A cache pool can serve several requests; it is not the context limit of one request.

For a simple capacity calculation, 1 GiB divided by 64 KiB per token holds 16,384 token slots before cache-manager padding and other state costs. Halving the bytes per token nearly doubles that part of the capacity. It does not double total VRAM: weights, runtime allocations, recurrent state and the draft still occupy their own space.

Recover wasted cache capacity first

The cache manager reserves memory in groups. Padding is space reserved to make those groups the same size, even when some of the slots are empty.

Before compressing anything, I checked why adding DFlash cost so much cache memory. Adding the drafter raised the cost per token of context from 64 KiB to about 89 KiB. The drafter's own attention is a 2,048-token sliding window, so it should cost almost nothing at long context.

I dumped the KV-cache groups vLLM actually built:

vLLM grouped the layers into equal-sized memory allocations. The helper model's five layers made the group size five, so each group containing four full-attention layers reserved space for an unused fifth layer. That empty space was repeated for every token of context. A FIXME comment in the code already pointed to the issue.

The cache-group planning fix chooses the largest group size that leaves the full-attention layers unpadded. Here that is 4, so the 5 drafter layers are the ones padded, to 8, and they cost only about nine fixed-size blocks. With DFlash, the same 12 GiB now holds 171,868 tokens instead of 141,631 (+21%). Step time is unchanged.

An 8-bit cache was already available

FP8 is the usual way to halve a KV cache, but the available execution paths constrained this Ampere deployment:

• FA2 and FlashInfer have no SM86 FP8 path in the versions measured here.
• Triton's e4m3 path needs SM89.
• Triton's e5m2 ran, but with about 7.5% error.

vLLM already has an int8_per_token_head cache format in its Triton attention implementation. It stores each cached number as an 8-bit integer, with a shared 32-bit scale to recover its approximate original value. An attention head is one set of calculations looking back through the saved history.

Including the scales, keys and values take 520 bytes per token per head, against 1,024 bytes in 16-bit BF16: 0.508× the space. See vLLM's KV-cache documentation for cache types and configuration.

The first run produced this:

The capital of France is Paris. The capital of Germany is Berlin ** to ( . within... have not (... ( (

With BF16 on the same Triton backend, the output was correct. With INT8, the first token was correct and subsequent output was corrupted. I instrumented one decode step, comparing the cache contents and the kernel output against a PyTorch reference:

Quality check: numerical agreement

Attention was correct, so the damage was somewhere else. It turned out to be a single line. When the backend first builds its views of the scale slots, it initialises them with fill_(1.0).

The model uses two kinds of memory: a growing attention cache and a fixed-size running summary maintained by Gated DeltaNet. vLLM lets both use parts of the same reserved memory pool. The scale-initialisation code ran after some layers had already written their summaries, and accidentally overwrote part of them. The first token was calculated before the damage; later tokens used the corrupted summaries.

The fixed scale-view setup explains why the fill was unnecessary: every scale that is read is either a slot that was written, or a masked load that defaults to 1.0. Removing it made INT8 output identical to BF16, word for word.

Is 8 bits enough?

llama.cpp's Q8_0 KV format uses one scale per 32 values. vLLM uses one per 256, which is cheaper but more exposed to outlier channels in RoPE-rotated keys. I measured before deciding whether that mattered.

I fed four 32K-token documents to each version, giving both the same sequence of tokens. This teacher-forced comparison asks how each version scores the same text, rather than letting different generated tokens change the next input. The prompt_logprobs option exposes those scores. The attention program reads from the compressed cache, so the comparison includes its rounding error.

In the table, ΔNLL measures the change in how well the model predicts the supplied text; close to zero means little change. Mean |Δ logprob| measures changes in individual token scores. Top-1 agreement is how often both versions prefer the same next token. FA2 means FlashAttention 2, the reference attention implementation here.

Quality check: prediction loss

INT8 moves the logits about as much as switching BF16 from FA2 to Triton does, and the error doesn't grow with position in these documents. That supported keeping per-256 scales for this implementation, with the task-level follow-up below. More varied inputs or longer reasoning could still expose sensitivity these tests missed.

The stock path was slow

Reading half as many cache bytes ought to help long-context decode. The stock implementation did the opposite. With INT8 working correctly, these measurements made access patterns and parallelism the next problem:

Request performance

Three separate problems:

1. Decode got slower, despite reading half the bytes. vLLM's kernel loads the int8 rows narrowly and fetches each scale with a separate scattered load.
2. DFlash throughput fell sharply. The Triton kernel doesn't split the KV sequence across thread blocks for multi-token queries. That's the same problem part 7 fixed for FA2, and this backend had none of that work.
3. Prefill doubled. Triton's prefill kernel is about 2× slower than FA2 even in BF16, and INT8 added another 17%.

Read the compressed cache without copying it first

llama.cpp has a mature quantised-KV path, and its CUDA attention implementation makes a clear split.

• One query token (the VEC kernel, fattn-vec): Q is quantised to int8 once, Q·K runs as DP4A integer dot products directly on the int8 keys, and V is dequantised in registers. No converted copy of the cache is ever written.
• Several query tokens: the whole visible K/V is first converted to an FP16 buffer (ggml_get_to_fp16_cuda), then the FP16 tensor-core kernel runs.

The first is the model for decode. The second would add substantial overhead to speculative verification, which re-reads 128K tokens of context each step with 6 query tokens. I therefore adapted the decode approach and used a different strategy for verification.

I wrote one GPU program for generating a token or checking several proposed tokens. It divides the saved history into sections, lets different workers read those sections together, then combines their results. This is called split-KV attention.

Where the data starts in memory also matters. Each head stores 256 data bytes followed by a four-byte scale, so the next section starts at byte 260. The compiler could not always establish that larger reads were safe and used smaller ones. Expressing the address differently let it prove four-byte alignment and generate wider reads.

Integer Q·K didn't make the default. Following llama.cpp, quantising Q to int8 and running Q·K on int8 tensor cores is faster at 128K: by 17% for decode and 31% for verification. But its error against an exact reference is 1.4–2.1%, against 0.22% for BF16 Q·K, and real queries have outlier channels that random test data doesn't. It's an opt-in switch (GGUF_PQ2_INT8_KV_INT_QK=1), off by default.

Kernel time per layer on a 3090 Ti, in µs, against both vLLM's INT8 kernel and BF16 FA2:

Component test

A geometry sweep then gave the tuned table:

• Decode wants small tiles and many segments (16-token tile, 64 segments).
• Verification wants wide tiles and few segments (64-token tile, 16 segments).

Tuning cut verify time at 128K by a third. One gap remains. A verify step at 128K takes 1,026 µs on a 3090, against 822 µs for part 7's FA2 path. With 36 real query rows padded to Triton's power-of-two 64, the kernel is partly compute-bound: about 34 GFLOP per layer per step.

A hand-written MMA kernel with 8-row granularity could reduce that padding. At this stage, it was the main unresolved performance issue in this attention path.

Read the prompt using a small temporary buffer

The right strategy changes for prefill. One decode query would gain little from converting the entire cache before reading it. Many prompt tokens can reuse a converted range, provided that range fits in a small temporary buffer.

For prefill, llama.cpp's convert-to-FP16 approach is the right one. A 2,048-token chunk does enough work per key that a conversion pass costs little. The catch is memory: converting all of a 128K context at once needs a 512 MiB scratch per layer, too much on a 16 GB card.

Instead, the prompt-processing program converts one section at a time and reuses the same temporary storage:

• Convert earlier cached keys and values in sections of at most 32K tokens into about 128 MiB of reusable storage, then run FlashAttention 2 on each section.
• Handle the current prompt section separately, so each token can see only the tokens before it and itself.
• Combine the partial results with the correct attention weights, as in the program for generating tokens.

Converting costs about 1% of prefill time:

Request performance

When one request is reading and another is writing

With 1–8 concurrent requests and short prompts, INT8 matches BF16 within 0.97–1.04×. Those differences follow DFlash acceptance, which varies from run to run.

For concurrent serving, the important case is a long prompt arriving while another request is decoding. vLLM's chunked prefill then puts the newcomer's prefill chunk and the other request's decode token in the same step. My hook only handled uniform steps, so these mixed steps fell back to vLLM's slow Triton kernel. In the test, a 32K prompt (B) arrives while a 64K request (A) decodes:

Request performance

Request performance

The mixed-request implementation lets the CPU separate the two kinds of work. Requests generating or checking up to 16 tokens use the new program; each section of a new prompt uses FlashAttention 2. This adds a small coordination cost for mixed steps, but restores the faster path for both requests.

Where it landed

INT8 KV against BF16 on FA2, same card, same 12 GiB:

Capacity, numerical agreement and request performance

Cache-pool capacity, including slots that may serve multiple requests. This is not a measured A4000 capacity or the context limit of one request. SVG version.

Plain decode is faster than BF16 FA2, by up to 18% at 128K. Decode attention is memory-bound and now reads half the bytes. DFlash is level with FA2 up to 64K and 10% slower per step at 128K, because of the verify padding described above.

For the 16 GB deployment target, the component budget suggests about 190K token slots without speculative decoding and 90–120K with DFlash, against roughly 110K and 60K in BF16. These estimates still need validation against the target card's actual allocations and workload.

Task-level quality

These tests ask whether the versions still answer questions correctly. In the table, lost / gained counts questions answered correctly by only one version. The p-value describes the evidence for a difference in that paired comparison; it does not prove equal quality.

The teacher-forced test covers prefill. Decode runs through the new kernel, so I also ran the paired task comparison from part 7: BF16 against INT8 KV on the same GPU, greedy, with the same fixed questions.

Quality check: task scores

The GSM8K scores are unchanged. The paired GPQA comparisons did not detect a statistically significant difference, but 115–118 of 198 answers reached the 2,048-token budget before finishing. That limit makes the scores sensitive to which reasoning chains complete. Among questions both versions finished, the scores match without speculation and differ by two with it; this selected subset does not establish equivalent quality.

DFlash acceptance was 61%, the same as with BF16 KV. These are limited checks, and a realistic reasoning budget remains important.

What I took from this

Reducing cache bytes improved capacity, while faster attention required better reads, enough parallel work and a dispatch path that handled mixed requests. Concurrent serving brought all three together: a memory-saving feature was useful only if a new prompt could arrive without stalling someone else's generation.

1. Check what already exists before building. vLLM had an INT8 KV mode all along. It was one bad line away from working, and the work that mattered was making it fast for this model, not inventing a format.
2. Locate corrupted output by comparing with a reference at each boundary. Comparing the cache write and the attention output against PyTorch cleared both in one run, which pointed straight at shared memory. In hybrid models, "my tensor" and "your tensor" can be the same bytes.
3. Borrow selectively from llama.cpp. Its decode kernel was exactly right. Its multi-token strategy was right for prefill, with a bounded buffer, and exactly wrong for speculative verification.
4. Pin benchmark inputs, and let harnesses refuse to run. One long-context driver built its prompts from every report on disk, so writing a report changed the next run's prompts; pinning the source list fixed that. The quality harness checksums its sources and stopped itself twice mid-run: once when the vLLM checkout was rebased under it, and once when I edited plugin files it was using. Both stops were right.

The speed and memory work now form a complete serving path. Chapter 9 compares the current implementation with llama.cpp. To try the model yourself, the installation guide is the practical companion.

Test 29 records the standalone compressed-KV experiment. The later BF16 recurrent-state option changes the fixed DeltaNet state, not the INT8 KV cache studied here.

Sources

• Quantised DFlash context projection, cache-group planning and scale-view fix: changes in my vLLM fork.

• INT8 attention, prompt processing and mixed requests: implementation in my PQ2 plugin fork.

• vLLM quantised KV-cache documentation and vLLM source.

• Triton, FlashAttention-2, and llama.cpp CUDA attention source: backend references discussed above.

• llama.cpp reference: fattn.cu selects attention kernels; fattn-common.cuh implements vec_dot_fattn_vec_KQ_q8_0, dequantize_V_q8_0 and the FP16 conversion for multiple query tokens.