← All notes and experiments

Making a ternary reasoning model fast in vLLM

Latest published comparison. The opening tables report the current checkpoint; later sections explain the earlier stages that led to it.

The first implementation got Bonsai running in vLLM. I then improved its arithmetic, weight reads, speculative checks and conversation cache.

This chapter brings the results up to date with three more rounds of kernel work.

The latest measurements reach 94.4 tok/s for one request and 451 tok/s across eight requests, with FP32 recurrent state and no speculative decoding. Prompt processing is 1.52–1.60× llama.cpp at 1,024–16,384 tokens, and close to level at 128.

BF16 recurrent state makes plain generation faster still. Its long-reasoning validation remains incomplete.

Measurement note. Both engines ran sequentially on one RTX 3090 Ti at 300 W, with the same PQ2 model, real-text prompts and 8-bit KV caches. Prefix/prompt caching was off. The updated vLLM columns were measured at the end of the session; the llama.cpp column is the earlier run on the same card.

These are development measurements, not A4000 results. The records were integrated into this article on 16 October; no GPU benchmarks were rerun for the website.

For a quick guide to the labels and precision choices, see how to read the results.

Current checkpoint

Plain decoding

The main vLLM column uses the FP32 recurrent state selected by auto for this model. The BF16 column changes that state with --mamba-ssm-cache-dtype bfloat16; it does not change the INT8 KV cache. Rates for concurrent requests are aggregate output across users.

Request performance

The reported per-request rates at four requests are 67.1, 70.3 and 32.3 tok/s respectively; at eight they are 58.7, 63.2 and 20.4. The comparison uses the recorded aggregate rates rather than reconstructing them from these per-request summaries. The vLLM refresh uses medians of three runs for one request and two for four and eight requests.

Prompt processing

The llama.cpp rates disable recurrent-state checkpoints. The vLLM configuration includes the output-projection INT8 path and the later prefill kernel changes described below.

Request performance

These request-based prefill rates use time to first token as a proxy for prompt-processing time. They include effects outside the matrix kernels. Earlier intermediate runs still showed a short-prompt deficit of about 5%; the final refresh is within about 1% in the FP32 column. Each table belongs to its own measurement stage.

-ngl 999 -fa on -ctk q8_0 -ctv q8_0 -b 2048 -ub 512 --ctx-checkpoints 0

Establish a fair starting comparison

Earlier checkpoint. The sections below trace the stages leading to the opening tables. Their intermediate rates describe those stages, rather than the final configuration.

The first short-prompt result appeared to favour vLLM by 1.82×: 1,129 versus 622 tok/s at 128 tokens. llama-server was saving recurrent-state checkpoints while processing the prompt. On this hybrid model, that added roughly 100 ms per request: about 205 ms rather than 100 ms in the server's prompt timing.

With --ctx-checkpoints 0, llama.cpp reached 1,210 tok/s, close to llama-bench. The initial long-prompt prefill rates were also effectively level. The original short-prompt advantage therefore did not establish an engine difference. The comparison tables above use checkpoints off.

The initial decode profile told a different story. At a 4,096-token prompt followed by 256 generated tokens, the ternary matrix-vector kernels took 2.27 seconds in vLLM and 2.28 in llama.cpp. llama.cpp also accumulated about 0.76 seconds of idle gaps, roughly 3 ms between tokens, while vLLM kept the GPU busy for about 98% of decode time. The early single-request advantage came from the complete execution path; later kernel work increased it further.

Round one: improve the existing prefill path

The first changes addressed a floating-point fallback, tile selection and input preparation:

Request performance and numerical behaviour

The output projections accounted for 128 floating calls and about 16% of prefill. Isolated timings predicted a 0.29-second saving; the request measurement saved about 0.15 seconds. Once those projections moved to MMQ, that kernel accounted for 85% of prefill, making its tile geometry the more useful target.

Larger tiles spread unpacking work across more token rows. At 512 rows or more, 256 × 64 tiles with eight warps were 1.3–1.45× faster in the measured cases. Batched row preparation gave a smaller gain of about 1.5%, without the operator changes needed for the planned fusion. Plain decode remained about 84–85 tok/s at this stage.

Quality evidence for the output-projection change

The earlier article left task-level evaluation pending. The updated records now include it:

Quality check

GPQA changed 33 of 198 answers in either direction, and more than half reached the 2,048-token budget without a parsed answer. The study did not detect a clear score difference, but it is a limited test of reasoning quality. DFlash acceptance was unchanged in the recorded check.

Round two: CUDA prefill and activation reuse

The small-batch kernel from chapter 7 calculated directly from raw code planes. I extended that approach to large prefill tiles, using a hand-written INT8 tensor-core kernel. Each plane contributes a matrix product; the ternary −1 offset is applied once per quantisation group. The implementation uses a cp.async pipeline and arranges each lane's activation fragments for its matrix instructions.

It matched the Triton path bit for bit on 110 test cases, with identical end-to-end prompt log-probabilities in the recorded comparison.

Component test

Work per clock improved more than sustained throughput: the new kernel ran around 1.2 GHz rather than 1.58 GHz under the power cap. This repeated the lesson from chapter 5. Short timings and fixed-clock comparisons were insufficient to choose the serving configuration; sustained throughput at the operating limit mattered.

For small-batch decode, several warps were re-reading the same activation data for their own weight rows. Allowing a warp to handle two or four row tiles per activation load was bit-exact and improved wide projections by 9–23%. In that paired serving experiment, aggregate rates rose from 241.5 to 251.5 tok/s at four requests, and 375.7 to 406.8 at eight.

Round three: improve decode on real inputs

After those changes, a new profile put ternary matrix operations at 7.99 ms of a 10.78 ms single-request step, and 10.7 ms of a 16.77 ms eight-request step. Gated DeltaNet recurrence had become 2.52 ms, about 15% of the eight-request step. Hadamard and activation preparation took 0.93 and 1.16 ms respectively. These are the profile conditions before the next changes, not a breakdown of the final table.

Stage narrow projections through a ring

The down and output projections offered too few output tiles to occupy the GPU well. A read-only probe showed that wider contiguous row segments could approach 940 GB/s, compared with about 420 for the earlier pattern on those shapes.

I kept the prepared row-major weights and staged reads through a shared-memory cp.async ring, splitting the reduction across thread blocks. On captured eight-row inputs, the selected path reduced the 257-projection pass from 13.70 to 12.43 ms, a 9.3% latency reduction.

Down improved by 34% and output by 23% in that replay. In a separate serving comparison, eight-request aggregate throughput rose from 412 to 453 tok/s, four-request throughput improved by about 5%, and one request was unchanged because the dispatch applies at 2–8 rows.

The split-K down entries change FP32 summation order. They are deterministic, with relative RMS difference at most about 1.1e-4, but are not bit-identical. GGUF_PQ2_SMALL_MMQ_RING_EXACT retains only the bit-identical selections. A speculative GSM8K check scored 123 before and 121 after, with 126 of 128 answers identical. That small cohort does not establish equivalent task quality.

The component investigation also tested more extensive tile-major and window-interleaved layouts. They improved some reads while slowing other consumers of the weights. Tests 24–25 and 30–36 explain why the integrated path retained row-major storage.

Use raw code planes for one token

The dedicated Triton kernel still expanded packed codes through PRMT. The new single-token CUDA path uses the raw two-bit planes with DP4A instead. Activations are quantised once into plane order, the ternary offset is applied once per group, and the kernel stages activations and prefetches weight windows for the one-row case.

Using real weights changed the expected benefit. The Triton gate/up timing was 60.9 µs on random bytes and 67.7 on real ternary weights. The plane-DP4A timings were almost unchanged, at 55.6 and 56.0 µs. The recorded experiment establishes content sensitivity; it does not isolate its hardware cause.

Across captured real inputs, the 257-projection pass fell from 11.29 to 8.81 ms, a 21.9% latency reduction. All 771 outputs were finite and deterministic, but only 278 matched the previous bits; maximum relative RMS difference was about 1.4e-4. See tests 37–39 for the component evidence.

In the paired serving run, one request improved from 82.4 to 94.4 tok/s, about 14.5%. Four and eight requests, and prefill, were unchanged. GSM8K scored 121 in both configurations, with 126 of 128 answers identical. The numerical and short-task checks support further use of the path, while leaving broader task validation open. GGUF_PQ2_M1_CUDA=0 restores the Triton implementation.

Reducing recurrent-state traffic

For this model, --mamba-ssm-cache-dtype auto resolved to FP32. DeltaNet reads and writes its recurrent state each step, independently of the growing KV cache. BF16 halves that state per sequence. In a controlled A/B it improved plain decoding by 2.3% at one request, 4.7% at four and 8.7% at eight. The final refresh above shows comparable gains.

The long-window check reported changes in prediction loss of −0.00025 / −0.00023 / −0.00029 / −0.00046 / +0.00002 nats over the 0–16K / 16–32K / 32–48K / 0–32K / 20–52K windows. GSM8K scored 120 in both modes, with all 128 answers identical.

This is a useful option with an explicit validation limit. GPQA and multi-thousand-token reasoning were not evaluated for this flag. Small long-window loss changes and short-answer agreement do not establish behaviour on long reasoning chains.

Speculative decoding across more prompts

The earlier speculative results depended on small prompt samples. I repeated the comparison on 12 different real-text prompts per length, generating 512 tokens per prompt, with one request and one server at a time. Each state setting ran twice in the order FP32, BF16, FP32, BF16. The target used the later kernels, DFlash and k = 5.

Request performance

Relative to plain decoding around 94–97 tok/s, these mean rates are roughly 2.2× at 1K and 1.9–2.0× at 4K. The comparison is specific to these prompt samples.

Repeated runs of the same setting agreed within 0.5% on the mean and had identical acceptance counts. Variation between prompts was much larger: 140–300 tok/s at 1K and 138–311 at 4K. Five-prompt samples could therefore produce substantially different means without a kernel change. These twelve-prompt results replace the earlier small-sample rates as the current speculative summary; earlier paired rows remain useful within their own experiments.

Paired per-prompt BF16 gains averaged +0.7% at 1K and +7.8% at 4K, with standard deviations of 12–17% and standard errors of 3.6 and 4.8 percentage points. BF16 was faster on 9 and 8 of 12 prompts respectively. Twelve prompts do not establish a general state-precision advantage for speculation, even though the plain-decode A/B shows a gain.

Check the measurement process as well as the kernel

Two follow-ups corrected earlier interpretations.

Compile-cache state affected numerical comparisons. Small cross-launch prediction-loss differences initially attributed to noise were reproducible on a quiet system. Checks found bit-exact kernels on all 384 live calls, identical outputs in eager mode, and identical outputs after fresh compilation of both configurations. The differences followed cached compile artefacts. A later A/B showed the same issue. For bit-level comparisons, the recorded procedure disables vLLM's compile cache:

VLLM_DISABLE_COMPILE_CACHE=1

This corrects the earlier account that treated the small differences as ordinary launch-to-launch noise. Bit-exact component kernels alone do not guarantee matching compiled graphs and caches.

Background workloads affected timings. Interrupted scripts left a benchmark and a server running on the same GPU, producing apparent single-request rates of 55–66 rather than about 82 tok/s. The accepted measurements exclude those runs. Checking for remaining GPU workloads became part of the timing procedure.

The real-versus-random weight experiment adds a third lesson: synthetic data can change a kernel's ranking. Component timing, real-input replay and complete requests remain separate checks.

Limits of the results

The matrix kernels are only part of the complete step. Normalisation, recurrent state, attention and input preparation also contribute; at 128K, attention accounts for about 38% of prefill in the recorded profile. Long-reasoning quality has not been established by these checks, particularly for the BF16 state option.

Several tested changes were not integrated: a universal storage-layout switch regressed other readers; a warp-private ring was 8–27% slower on some wide projections; and fused Hadamard/plane quantisation saved only about 0.5–0.6 µs per call, around 1.4% of a single-request step. Their component evidence remains available in the experiment reference.

The current results were tuned on the development card. Tile, power and memory behaviour on the 16 GB deployment target has not been measured. The earlier 7.1 GiB model allowance in chapter 8 remains a capacity estimate pending confirmation of the actual allocation.

The GPU concepts companion explains the layouts, workers and numerical checks. The experiment reference records the 39 component investigations. Implementation changes are in the public PQ2 plugin fork and the vLLM fork.