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

By the end of part 3, integer arithmetic and a better weight layout had doubled generation to about 50 tokens per second on a 3090 Ti. llama.cpp was still faster for one request.

Concurrent serving was the goal, but each user could still wait for thousands of reasoning tokens. Single-request speed mattered too.

On 5–6 October I traced the remaining gap, built a more reliable way to measure it and tailored the kernel to one token at a time.

Measurement note. Serving results in this chapter use one RTX 3090 Ti under the early 275 W cap. Workload lengths and repeat counts are stated with each comparison. The latest matched comparison is recorded separately.

How big is the gap, and where is it?

The matched comparison used a 90-token prompt, 2048 output tokens and the median of three runs. This measures the software gap under the same conditions:

Request performance

Two other results from that stage suggested the single-request path deserved attention:

• Prefill: vLLM with MMQ was ahead, at 1093 vs 249 tok/s for a 128-token prompt and 1111 vs 1059 at 4096.
• Eight concurrent requests: vLLM was ahead too, at 21.4 vs 18.7 tok/s per request.

The short-prompt prefill comparison needs a qualification: chapter 9 later showed that llama-server's checkpoint defaults can dominate short-prompt timings. That finding makes these earlier prefill figures insufficient to establish an engine advantage without checking the old run's settings.

So the gap was specifically generating an answer for one request. TTFT was about equal (194 vs 202 ms), so the whole gap was in sustained generation: about 2.6 ms per token.

It's not the CPU

The usual suspect for "vLLM is slower at batch 1" is Python and scheduler overhead. Nsight said otherwise:

GPU profile

vLLM's GPU was busy 99.2% of the time; CUDA graphs were limiting launch overhead. llama.cpp actually had more idle gaps, but its kernels finished sooner. Removing every vLLM gap would have saved under 1%. In this trace, the gap was kernel time.

GPU profile

The 11% floating fallback was 64 calls per generated token. That is exactly 48 linear-attention output projections plus 16 full-attention o_projs, all of shape 5120 × 6144, which the integer path didn't yet cover.

Reducing integer GEMV duration by 20% would save about 5.6 s in this trace, roughly the size of the gap.

A 20% increase in a kernel's processing rate is a different claim from a 20% reduction in its duration. If it accounts for 69% of runtime and its rate rises by 20%, its duration falls to 1/1.2 of the original. Total runtime becomes 0.31 + 0.69/1.2 = 0.885, for about a 13% overall speedup. Speeding up a small kernel can never save more than its original share of the step.

Reading llama.cpp's kernel

Both engines use DP4A and 128-weight groups, so I compared the two decode kernels line by line:

The first comparison did not point to excessive register use. Registers are the small, fast storage each GPU worker uses for temporary numbers. I tried:

• dividing each matrix into pieces suited to its size,
• recovering packed weights with a byte-selection instruction, PRMT,
• adding several integer dot products into one running total,
• combining the MLP's gating step with the weighted calculation,
• combining input mixing and rounding into one GPU program,
• using the integer program for the output matrices still using floating-point arithmetic.

Measure with inputs taken from the real model

Each level adds conditions missing from the previous measurement. The last step is a deployment check, beyond the 3090 experiments reported here.

Running complete requests takes minutes for each candidate, and timings vary. Testing one GPU program with made-up inputs is quick, but had already misled me. I built a replay harness: a tool that saves inputs from the real model and runs those same inputs through different GPU programs.

1. Run the real model once with instrumentation and capture the real inputs and outputs of every PQ2 projection at decode positions 1, 32 and 128. That is 257 projection layers, 771 samples and 64 linked MLP gate→down pairs.
2. Replay those captured activations through any candidate kernel under CUDA graphs. Check bit-exactness against the saved outputs (all 771 baseline comparisons were exact), or check against numerical thresholds for candidates that legitimately change rounding.
3. Time candidates in adjacent pairs with the baseline, alternating order, so clock drift hits both equally.

A comparison of nine options ran in 22 seconds.

Each level tests a different prediction. Synthetic inputs expose instruction and tile costs. Captured activations expose real values, strides and linked layer effects. Serving adds the scheduler, other kernels, graph execution and power behaviour. A candidate has to survive that progression before its microbenchmark win is useful.

Here is how the candidates fared at each level of the measurement ladder:

Component test and request performance

Three things stand out:

• Replay got the ranking right. It put the three winners 10.3% ahead of "everything on", and serving measured 11.6% in the reference run. It did not get the size right: 1.20× in replay became 1.24× in serving.
• "Enable everything" was worse than picking. Chained DP4A hurt at M = 1 (−13% in replay), and the two fusions cost slightly more than they saved.
• The reference implementation needed to be tested in this layout. llama.cpp uses chained accumulation, but it was slower here. Its effect depends on the thread layout and batch size; part 5 shows a case where it helps.

Recover four weights with one lookup

The original unpacking code needed several operations on bits to recover the weights. The GPU's PRMT instruction can instead select four bytes from a small lookup table in one instruction. Its low-level name is prmt.b32.

The constant 0x020100ff holds the four possible decoded bytes: 0xff (−1), 0x00 (0), 0x01 (+1) and 0x02 (+2). Build a selector from the four 2-bit codes, one per nibble, and a single prmt produces four signed INT8 weights:

For codes [0, 1, 2, 1], the selector is 0x1210. PRMT selects lookup bytes [0xff, 0x00, 0x01, 0x00], which represent [−1, 0, +1, 0]. The represented weights are unchanged. The hypothesis is that selecting bytes costs less than constructing them through several arithmetic instructions.

@triton.jit
def _decode_codes(packed, PRMT: tl.constexpr):
    if PRMT:
        selector = (packed & 3) | ((packed >> 2 & 3) << 4)
        selector |= ((packed >> 4 & 3) << 8) | ((packed >> 6 & 3) << 12)
        return tl.inline_asm_elementwise(
            "prmt.b32 $0, 0x020100ff, 0x020100ff, $1;",
            constraints="=r,r", args=[selector], dtype=tl.int32,
            is_pure=True, pack=1,
        )
    ...

From pq2_int_gemv.py.

Result: three winners

With 4096 input tokens and 2048 output tokens, the three changes together improved generation from 50.30 to 62.17 tok/s (+23.6%). They became the defaults within the integer path; chained DP4A and the two combined programs were rejected for single-token generation.

Try different ways of dividing the work

The first program divided every matrix the same way. A next-token scoring table with 248,320 output rows and a layer with 5120 output rows need different amounts of parallel work. I tested 180 combinations of how many rows and weight groups each worker handles, and how many workers cooperate.

Two lessons came out of this:

1. Noise is large relative to the effects. Repeated identical controls spread by up to about 12%. Anything under a few percent needed confirmation in a fresh process.
2. Replay winners didn't always survive. A "geometry winners" set scored 1.217× in replay versus 1.199× for the three winners. Longer paired runs then rejected its gate/up and output-projection changes.

The tiles that survived became shape-tuned-v2:

Component test

In serving that was +3.2% in serving. The clock telemetry showed something odd: the median SM clock was 1395 MHz for the reference and 1425 MHz for the candidate, both pinned at about 273 W against the 275 W cap. A more efficient kernel was letting the governor clock higher. That clock effect gets its own section in part 5.

After v2, _int_gemv was still 74.7% of decode kernel time. FA2 split-KV attention was 5.0%, Hadamard 4.4%, Gated DeltaNet 3.2% and quantisation 2.6%. The GPU was busy 99.15% of the time.

The win: a program just for one token

The remaining opportunity came from the program's generality. I wanted to remove the dimensions and branches a single token never used, then check whether the compiler needed less temporary storage and could keep more work running at once.

The integer GEMV kernel was written to handle 1–16 tokens, both weight layouts and several decode modes. At M = 1 with prepared weights, those additional cases are unnecessary, but the compiler still has to allocate resources for them.

So I wrote a dedicated single-token kernel:

• No token dimension at all. The activation is one row.
• The scale plane is passed as its own pointer, instead of computing its offset inside the combined allocation.
• Each DP4A takes an immediate zero accumulator (dp4a ... $1, $2, 0), so every DP4A is independent and the sum happens afterwards.
• A one-stage loop (num_stages=1). No software pipelining for the compiler to over-allocate for.

@triton.jit
def _single(Q, S, W, WS, Y, N: tl.constexpr, K: tl.constexpr,
            BN: tl.constexpr, BC: tl.constexpr):
    rows = tl.program_id(0) * BN + tl.arange(0, BN)
    offset = tl.arange(0, BC)
    acc = tl.full((BN, BC), 0, tl.float32)
    for base in tl.range(0, tl.cdiv(K // 128, BC), num_stages=1):
        blocks = base * BC + offset
        packed, ws = _load_tile(W, WS, rows, blocks, N, K)
        acc += _consume(Q, S, packed, ws, blocks, K)
    tl.store(Y + rows, tl.sum(acc, 1), rows < N)

From pq2_single_gemv.py.

On the dominant gate/up shape, compiler resources dropped from 128 to 93 registers per thread and from 8192 to 1024 bytes of shared memory, with no spills.

In replay with real activations, the single-token kernel alone gained +11.8%, and +14.1% combined with v2 tiles. On its own it was bit-exact against the previous kernel.

I tested several further changes in the same comparison:

Component test

Deeper pipelining and explicit prefetch made things worse. That pattern repeats through the rest of the series.

Serving

Full model, 4096-token prompt, 2048 generated tokens, median of three runs:

Request performance

TTFT was unchanged. Numerical validation covered 1,542 projections, 384 linked MLP outputs and 12 CUDA-graph input-mutation checks, with a maximum relative RMS of 9.5 × 10⁻⁵.

Clocks rose again, from 1395 to 1455 MHz in this run. Some of that +18% is the power governor rewarding a leaner kernel, and I'll come back to that.

Both changes became the defaults inside the integer path (GGUF_PQ2_SINGLE_GEMV=1, GGUF_PQ2_INT_GEMV_VARIANT=shape-tuned-v2). The re-measured promoted default was 72.7 tok/s.

Checkpoint: faster through specialisation

Single-request decode on a 3090 Ti went from about 50 to about 72 tok/s over these two days, and from 25 at the start of the series.

The 72.5 tok/s result is encouraging, but the llama.cpp figure of 58.6 tok/s used a 90-token prompt, while this run used 4,096 tokens. The changed workload prevents a matched speedup claim. The same-GPU comparison in chapter 9 supplies the later checkpoint.

A short-prompt check from the eight-request work used a 90-token prompt and 512 output tokens. The selected configuration ran at 62.3 tok/s in that run, compared with llama.cpp's 58.6 tok/s from 5 October. The output lengths and test dates differ, so this is supporting evidence rather than a matched final comparison. A same-day comparison with identical requests remains the stronger test.

Against the ceiling from part 2, 72 tok/s is about 49% of the 148 tok/s weight-bandwidth bound. That is an improvement from the initial 17%, while leaving a substantial gap to the idealised limit. The next chapter examines the practical constraints behind that gap.

The useful change was removing work the common case did not need. The tile screens and compiler reports helped explain it; complete requests decided whether it mattered. Next, those same measurement tools would produce a puzzle: a faster replay that barely changed serving speed.

The supporting replay records are tests 12–14. Later tests 37–39 revisit the single-token path with raw code planes and real weights; chapter 9 reports their separate serving follow-up.

Sources

Implementation and instruction references:

• triton/pq2_single_gemv.py: Single-token calculation; integer decode and PRMT unpacking.
• NVIDIA PTX ISA: prmt.b32 and dp4a instruction references.
• Nsight Compute documentation: GPU kernel profiling and hardware-counter reference.