Kernel experiments: recorded evidence
Making a ternary reasoning model fast in vLLM
Published measurements are snapshots of the stated hardware, settings and revisions.
This companion reference records 39 standalone component experiments, grouped by the concept they test. Some use synthetic data and others replay inputs captured from the model. These timings do not measure a complete vLLM or llama.cpp request, and numerical agreement at one projection does not establish unchanged answer quality.
The source extract was prepared on 15 October 2026 and includes the later same-day follow-ups. We have not rerun its GPU experiments for this article. The recorded qualifications are retained; test numbers follow the order on this page. The GPU concepts guide explains the terms and data flow; chapter 9 covers the separate serving and task-quality measurements.
Measurement type: component test. These records examine individual operations or linked replays, not complete requests. Numerical checks compare outputs; they do not score answers.
Each test starts with the idea and its measured result. Expand the explanation for the reasoning, limitations, checks and code. The snippets include current implementations and archived experiment versions as labelled in the source; they are explanatory extracts, not a standalone installation recipe.
Browse the experiments
Each section introduces a GPU concept with a diagram, then presents the experiments that test it. Tests are numbered 1–39 in reading order. Use the index below to jump to a concept. The diagrams illustrate the data flow; they are not measured traces. This site renders them as static SVGs, so no Mermaid software is needed to read them.
A. Packed weights: storage, loads and decoding
Takeaway. Keep weights compressed, but make them easier to read and unpack. The useful layout depends on the kernel that consumes it.
Terms for this section: Row-major / tile-major · Stride and coalescing · Groups, planes and fragments.
Diagram 1: The same compressed values, arranged differently
Both layouts store 32 code bytes plus a 2-byte scale per 128 weights. The prepared allocation puts all code bytes first, followed by all scales; it does not expand weights to dense values. Separate planes change addresses, not values. Test 1 explores preparation and decoding; test 2 expands each loaded byte once; test 3 confirms simpler decoding; test 4 checks the layout across readers.
Test 1. Packed loading, unpacking and scheduling (Oct 1)
Visual reference: Diagram 1 · Topic overview.
The idea: We try different ways to store, load and decode the same compressed weights. One candidate separates the packed weight bytes from their scale values; another replaces several decoding steps with a shorter sequence of integer operations.
The result: Aligned code/scale buffers plus simpler unpacking give 1.098–1.326× the speed of the tuned baseline. Unpacking alone gives 1.025–1.140×. Confirmation outputs match bit for bit. Manual wider-lane loads and pipelines usually lose. Preparation is outside the timing and keeps an extra packed copy in this experiment.
Explanation, checks and code
Why it matters on the GPU: The original blocks mix a 2-byte scale with 32 code bytes, giving a 34-byte stride. That makes some wide aligned loads awkward. Separate code and scale planes give the compiler a more regular layout for wide loads. But making each thread load a wider word can require extra exchanges between threads to put values where DP4A expects them. Wider source loads therefore don’t automatically mean less work.
What to take from it: Aligned planes and simpler decoding won in the measured comparisons. The machine-instruction listings show fewer unpacking operations and different loads. They don’t directly measure memory transactions. Preparing an extra copy also costs memory and needs enough repeated uses to pay for that extra work.
Setup: We try 78 initial variants and 30 further variants that preserve the dot-product lanes. Confirmation uses a new random seed, both dtypes and strides, and both timing orders. Separate checks cover all packed byte values and saved model inputs.
Relevant code:
The tested add/XOR byte expansion replaces individual signed-code subtractions.
Source: pq2-dev/benchmarks/plugin-archive/pq2_packed_experiment.py:13–27.
@triton.jit
def _expand_byte(p, FAST: tl.constexpr):
if FAST:
# Each byte is 0..3. Adding 127 cannot carry between bytes; XOR maps
# 127/128/129/130 to signed -1/0/1/2 without four separate subtractions.
spread = (p & 3) | ((p << 6) & 0x300)
spread |= ((p << 12) & 0x30000) | ((p << 18) & 0x3000000)
word = ((spread.to(tl.uint32) + 0x7F7F7F7F) ^ 0x80808080).to(tl.int32)
else:
c0 = ((p & 3) - 1) & 255
c1 = (((p >> 2) & 3) - 1) & 255
c2 = (((p >> 4) & 3) - 1) & 255
c3 = (((p >> 6) & 3) - 1) & 255
word = c0 | (c1 << 8) | (c2 << 16) | (c3 << 24)
return word
Preparation separates aligned code and scale buffers; preparation cost is outside repeated-call timing.
Source: pq2-dev/benchmarks/plugin-archive/pq2_packed_experiment.py:142–181.
def prepare_pq2(weight, *, load_bytes=1, separate=False):
"""Prepare reusable buffers. Separate layout allocates a new packed copy.
Raw 16-bit loads require an even base offset and row stride. No dense
dequantization; separate buffers still use exactly 34 bytes per block.
"""
if (
weight.ndim != 2
or weight.dtype != torch.uint8
or not weight.is_cuda
or weight.shape[1] == 0
or weight.shape[1] % 34
):
raise ValueError("Expected a CUDA uint8 PQ2 matrix")
if load_bytes not in (1, 2, 4) or (not separate and load_bytes == 4):
raise ValueError("Raw PQ2 supports 1/2-byte loads; separate supports 1/2/4")
if weight.stride(1) != 1:
weight = weight.contiguous()
n, width = weight.shape
k = width // 34 * 128
if separate:
blocks = weight.view(n, width // 34, 34)
scales = (
blocks[:, :, :2]
.clone(memory_format=torch.contiguous_format)
.view(torch.float16)
.reshape(n, width // 34)
)
codes = (
blocks[:, :, 2:]
.clone(memory_format=torch.contiguous_format)
.reshape(n, width // 34 * 32)
)
else:
codes = weight
scales = weight # Unused pointer for this specialization; no extra storage.
if load_bytes == 2 and (weight.data_ptr() % 2 or weight.stride(0) % 2):
raise ValueError("Raw 16-bit loads require an even base and row stride")
dtype = {1: torch.uint8, 2: torch.int16, 4: torch.int32}[load_bytes]
return PreparedPQ2(codes.view(dtype), scales, k, load_bytes, separate)
Test 2. MMQ byte expansion (Oct 2)
Visual reference: Diagram 1 · Topic overview.
The idea: We change the unpacking step: fetch each packed byte once, then expand its four two-bit codes into the 128 values needed for a group. We compare the new kernel against a saved older version and require exactly the same output bytes.
The result: All 24 cases match bit for bit. Raw weights give a median speedup of 1.190×, ranging from 1.125–1.319×. Prepared weights give a median of 2.290×, ranging from 1.898–2.464×. These ratios are against the previous MMQ implementation.
Explanation, checks and code
Why it matters on the GPU: The mathematical dot product and Tensor Core work stay the same. The change gives the compiler a different load-and-expand layout instead of addressing the same packed bytes through a 128-element grid. That can reduce load/addressing and rearrangement overhead, especially with separate prepared code and scale planes. Repeated addresses in the source don’t prove that the old kernel fetched each byte from DRAM four times.
What to take from it: Exact equality makes this an implementation efficiency change, rather than a different quantisation approximation. The larger gain for prepared weights shows that unpacking strategy interacts strongly with storage layout.
Setup: We compare six shapes at M = 128 and 512, using raw and prepared weights in BF16. Both method orders are timed with 100 ms CUDA-graph measurements against a saved previous kernel.
Relevant code:
Saved byte-expanded MMQ version: load 32 packed bytes, expand four planes, reshape into the 128-element dot.
Source: pq2-dev/benchmarks/results/pq2-mmq-groups-20261002/reference-group128.py:31–65.
@triton.jit(do_not_specialize=["M"])
def _mmq(
Q,
S,
W,
Y,
M,
N: tl.constexpr,
K: tl.constexpr,
STRIDE: tl.constexpr,
BM: tl.constexpr,
BN: tl.constexpr,
PREPARED: tl.constexpr = False,
):
m = tl.program_id(0) * BM + tl.arange(0, BM)
n = tl.program_id(1) * BN + tl.arange(0, BN)
k = tl.arange(0, 128)
acc = tl.full((BM, BN), 0, tl.float32)
for block in range(K // 128):
a = tl.load(
Q + m[:, None] * K + block * 128 + k[None, :], m[:, None] < M, other=0
)
ws = _pq2_scale(W, n, block, N, K, STRIDE, n < N, PREPARED)
# Load each packed byte once before expanding its four signed codes.
byte = tl.arange(0, 32)
packed = _pq2_codes(
W, n[None, :], block, byte[:, None], K, STRIDE, n[None, :] < N, PREPARED
)
shifts = 2 * tl.arange(0, 4)
codes = ((packed.to(tl.int32)[:, None, :] >> shifts[None, :, None]) & 3) - 1
b = tl.reshape(codes.to(tl.int8), (128, BN))
partial = tl.dot(a, b).to(tl.float32)
xs = tl.load(S + m * (K // 128) + block, m < M, other=0)
acc += partial * xs[:, None] * ws[None, :]
tl.store(Y + m[:, None] * N + n[None, :], acc, (m[:, None] < M) & (n[None, :] < N))
Complete-call old/new comparison, byte equality, and both timing orders.
Source: pq2-dev/benchmarks/experiments/integer-decode/pq2_mmq_unpack_bench.py:47–59.
methods = {
"reference": partial(reference.pq2_mmq, x, w, prepared=prepared),
"candidate": partial(pq2_mmq, x, w, prepared=prepared),
}
expected = methods["reference"]()
actual = methods["candidate"]()
torch.testing.assert_close(
actual.view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0
)
timings = [
{mode: do_bench_cudagraph(methods[mode], rep=100) for mode in order}
for order in (list(methods), list(reversed(methods)))
]
Test 3. Integrated unpack arithmetic confirmation (Oct 1)
Visual reference: Diagram 1 · Topic overview.
The idea: We check the shorter packed-byte decoding sequence after it has been connected to the normal integer projection API. It compares the integrated result bit for bit with a frozen previous kernel across shapes, token counts, dtypes and row strides.
The result: All 100 paired outputs and all 54 saved outputs match bit for bit. Median speedups are 1.040× for gate/up, 1.049× for down, 1.077× for GDN, 1.061× for QKV and 1.038× for the head. Timings include input quantisation.
Explanation, checks and code
Why it matters on the GPU: The compressed bytes, scales, tile shape and DP4A products stay the same. The optimisation changes the integer instructions used to expand codes into signed 8-bit lanes. Fewer masks/subtractions can leave more instruction capacity for useful work and reduce dependencies, even when the number of bytes read from memory is unchanged.
What to take from it: Exact output equality isolates a decoding implementation change. Its speedup is measured against the original fixed tile, so it should not be added directly to percentages from the separately tuned packing experiment.
Setup: We compare 100 full-shape pairs across five projections, M = 1, 2, 4, 5 and 8, both dtypes and both layouts. Timings use both orders and 100 ms graphs; another check replays 54 saved outputs.
Relevant code:
Simplified byte expansion whose integrated implementation is being checked.
Source: pq2-dev/benchmarks/plugin-archive/pq2_packed_experiment.py:13–27.
@triton.jit
def _expand_byte(p, FAST: tl.constexpr):
if FAST:
# Each byte is 0..3. Adding 127 cannot carry between bytes; XOR maps
# 127/128/129/130 to signed -1/0/1/2 without four separate subtractions.
spread = (p & 3) | ((p << 6) & 0x300)
spread |= ((p << 12) & 0x30000) | ((p << 18) & 0x3000000)
word = ((spread.to(tl.uint32) + 0x7F7F7F7F) ^ 0x80808080).to(tl.int32)
else:
c0 = ((p & 3) - 1) & 255
c1 = (((p >> 2) & 3) - 1) & 255
c2 = (((p >> 4) & 3) - 1) & 255
c3 = (((p >> 6) & 3) - 1) & 255
word = c0 | (c1 << 8) | (c2 << 16) | (c3 << 24)
return word
Frozen previous/current API outputs must match as integer bits, then both method orders are timed.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/confirm_pq2_unpack.py:51–74.
methods = {
"previous": partial(previous.pq2_int_gemv, x, weight),
"integrated": partial(pq2_int_gemv, x, weight),
}
expected = methods["previous"]()
actual = methods["integrated"]()
torch.testing.assert_close(
actual.view(torch.int16),
expected.view(torch.int16),
atol=0,
rtol=0,
)
timings = []
for order in (list(methods), list(reversed(methods))):
timings.append(
{
name: do_bench_cudagraph(methods[name], rep=100)
for name in order
}
)
row = dict(
projection=projection,
m=m,
dtype=str(dtype),
Test 4. Shared prepared-layout kernel matrix (Oct 2)
Visual reference: Diagram 1 · Topic overview.
The idea: We give all projection readers the same prepared representation: one allocation containing a code plane followed by a scale plane. It compares that representation against the original interleaved blocks, measuring the repeated projection call after preparation.
The result: Median raw/prepared decode speedups are 1.118× for gate/up, 1.083× for down, 1.077× for GDN, 1.031× for QKV, 1.109× for the head and 1.059× for output. Prefill improves gate/up by 1.058× and down by 1.030×, but other shapes slow down by about 3–10%. All 96 synthetic and 108 saved-output strict checks show no differences.
Explanation, checks and code
Why it matters on the GPU: Separating the planes makes addresses more regular and can permit wider aligned loads without storing dense weights. The total packed payload stays 34 bytes per 128 weights. Different kernels traverse that payload differently: a layout that suits small-row decode can be less convenient for a large tiled prefill kernel. Keeping one common allocation also avoids requiring separate resident weight formats for every reader.
What to take from it: All measured decode shapes improve, while several prefill shapes regress. That shows why a memory-layout optimisation must be tested across its readers. These timings exclude one-time conversion and do not by themselves establish load-time peak memory.
Setup: We check 96 synthetic cases across six shapes in BF16 and FP16, using M = 1, 4, 8, 128 and 512. Both timing orders are measured, with preparation outside the timing.
Relevant code:
The prepared matrix retains equal packed storage size, compares raw/prepared reader outputs and times complete calls.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/benchmark_pq2_prepared.py:51–91.
torch.uint8
).reshape(n, k // 128, 2)
prepared = prepare_pq2_layout(weight)
assert prepared.untyped_storage().nbytes() == weight.numel()
x = torch.randn(m, k, device="cuda", dtype=dtype)
method = (
pq2_int_gemv
if m <= 8 and projection != "attention_gdn_output"
else pq2_matmul
)
methods = {
"raw": partial(method, x, weight),
"prepared": partial(method, x, prepared, prepared=True),
}
expected = methods["raw"]()
torch.testing.assert_close(
methods["prepared"]().view(torch.uint8),
expected.view(torch.uint8),
atol=0,
rtol=0,
)
timings = (
[]
if args.validate_only
else [
{
name: do_bench_cudagraph(methods[name], rep=100)
for name in order
}
for order in (list(methods), list(reversed(methods)))
]
)
row = dict(
projection=projection,
m=m,
dtype=str(dtype),
layout=layout,
method=method.__name__,
different_outputs=0,
resident_prepared_bytes=prepared.numel(),
timings_ms=timings,
B. Reusing weights and choosing compute instructions
Takeaway. Match the arithmetic to the batch size: small integers helped decode; Tensor Cores became useful when more token rows were available.
Terms for this section: Tiles · Groups, planes and fragments.
Diagram 2: One weight tile, several tokens, different arithmetic
The arithmetic choices are alternatives, not consecutive stages. Floating GEMV uses floating inputs; DP4A and INT8 Tensor Core paths first quantise the inputs. Test 5 reuses weights with floating arithmetic; test 6 uses DP4A; test 7 uses Tensor Cores. Reuse reduces repeated work, but larger live tiles consume registers.
Test 5. Floating batched GEMV (Oct 1)
Visual reference: Diagram 2 · Topic overview.
The idea: We multiply the same compressed weights by several token vectors together. Instead of loading and decoding a weight separately for each token, a program loads a small weight tile and uses it for several tokens before moving on.
The result: At M = 4, the candidate is 1.70–1.90× as fast as native across the projections. At M = 1, it manages only 0.53–0.56×, and most cases at M ≥ 5 slow down. All 280 outputs pass the FP32-reference check.
Explanation, checks and code
Why it matters on the GPU: Weights arrive from GPU memory in packed two-bit form, then become floating-point values for multiplication on the ordinary arithmetic units. Reuse saves repeated weight loads and decoding. However, keeping partial answers for several tokens uses more registers, the small storage attached to running threads. Too much live data can reduce how much work fits on each GPU multiprocessor at once.
What to take from it: Four tokens benefit from reuse in this particular design. One token has no cross-token reuse to pay for the extra structure, while larger batches need a more suitable tiled matrix-multiply design. The test doesn’t pin down the hardware cause of every slowdown.
Setup: We try five projections at M = 1, 2, 4, 5, 8 and 16, using BF16 and FP16 with both contiguous and strided rows. There are 240 initial records and 40 repeat records.
Relevant code:
Archived worktree implementation of the floating batch-reuse kernel; its exact identity with the October 1 timed revision is not established.
Source: pq2-dev/benchmarks/worktrees/fox-moe-safety/vllm_gguf_plugin/triton/pq2_gemv.py:11–58.
@triton.jit(do_not_specialize=["M"])
def _batched_gemv(
X,
W,
Y,
M,
N: tl.constexpr,
K: tl.constexpr,
STRIDE: tl.constexpr,
BT: tl.constexpr,
BN: tl.constexpr,
BK: tl.constexpr,
PREPARED: tl.constexpr = False,
):
rows = tl.program_id(0) * BN + tl.arange(0, BN)
tokens = tl.arange(0, BT)
offsets = tl.arange(0, BK)
acc = tl.full((BT, BN, BK), 0, tl.float32)
for base in range(tl.cdiv(K, BK)):
k = base * BK + offsets
mask = (rows[:, None] < N) & (k[None, :] < K)
scale = _pq2_scale(
W, rows[:, None], k[None, :] // 128, N, K, STRIDE, mask, PREPARED
)
packed = _pq2_codes(
W,
rows[:, None],
k[None, :] // 128,
(k[None, :] % 128) // 4,
K,
STRIDE,
mask,
PREPARED,
).to(tl.int32)
code = ((packed >> (2 * (k[None, :] % 4))) & 3) - 1
a = tl.load(
X + tokens[:, None] * K + k[None, :],
(tokens[:, None] < M) & (k[None, :] < K),
other=0,
)
b = (code.to(tl.float32) * scale).to(a.dtype).to(tl.float32)
acc += a.to(tl.float32)[:, None, :] * b[None, :, :]
values = tl.sum(acc, 2)
tl.store(
Y + tokens[:, None] * N + rows[None, :],
values,
(tokens[:, None] < M) & (rows[None, :] < N),
)
The benchmark validates the batched floating result against the dense reference and selects complete-call functions.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/benchmark_prism.py:136–159.
gemv_first_call_ms = None
if "batched_gemv" in methods:
torch.cuda.synchronize()
start = time.perf_counter()
candidate = pq2_batched_gemv(x, w)
torch.cuda.synchronize()
gemv_first_call_ms = (time.perf_counter() - start) * 1000
torch.testing.assert_close(
candidate,
expected,
atol=0.02 if dtype == torch.bfloat16 else 0.005,
rtol=0.02 if dtype == torch.bfloat16 else 0.005,
)
del expected, actual
functions = {
"int_gemv": lambda: pq2_int_gemv(x, w),
"batched_gemv": lambda: pq2_batched_gemv(x, w),
"native": lambda: pq2_matmul(x, w),
"mmq": lambda: pq2_mmq(x, w),
"dequant_dense": lambda: x @ dequantize().T,
"dense": lambda: x @ dense.T,
}
rows = []
for method in methods:
Test 6. Integer DP4A GEMV (Oct 1)
Visual reference: Diagram 2 · Topic overview.
The idea: We round each group of 128 input values to small signed integers, multiplies them by the decoded compressed weights, then uses scale factors to put the answer back into floating-point units. It measures the whole operation, including rounding the inputs.
The result: Against native, speedups range from 1.49–1.70× at M = 1 and 4.09–4.90× at M = 4. None of the 120 cases is slower than native. The largest synthetic relative RMS error is 0.008666. In the repeat, speedups are 1.58–1.77× at M = 1 and 4.53–4.77× at M = 4.
Explanation, checks and code
Why it matters on the GPU: DP4A is an integer dot-product instruction: one instruction multiplies four pairs of 8-bit integers and adds their products to a 32-bit sum. It uses integer execution hardware, rather than Tensor Core matrix instructions. The two-bit weights remain compressed in memory; only the currently needed values are decoded. Input quantisation is performed once and reused across output rows, reducing repeated preparation work.
What to take from it: The speedups show that compressed loads and integer dot products work well for these small matrices. Rounding the inputs changes the numerical problem slightly, so a close projection result is not proof of unchanged model quality.
Setup: We quantise inputs in groups of 128, calculate DP4A integer dots, then apply scales in FP32. Complete-call graph timings cover five shapes, six M values, two dtypes and two layouts. A repeat reverses the method order.
Relevant code:
Saved group128 quantizer: signed INT8 values with nearest-even rounding and FP32 scales.
Source: pq2-dev/benchmarks/results/pq2-mmq-groups-20261002/reference-group128.py:16–28.
@triton.jit
def _quantize(X, Q, S, K: tl.constexpr):
block = tl.program_id(0)
i = tl.arange(0, 128)
a = tl.load(X + block * 128 + i).to(tl.float32)
scale = tl.max(tl.abs(a), 0) / 127.0
normalized = tl.div_rn(a, tl.where(scale > 0, scale, 1.0))
# Round nearest, ties to even, then clamp before the narrowing conversion.
q = tl.minimum(
127.0, tl.maximum(-127.0, tl.extra.cuda.libdevice.nearbyint(normalized))
)
tl.store(Q + block * 128 + i, q.to(tl.int8))
tl.store(S + block, scale)
Current implementation of the DP4A group dot; retained to show the arithmetic, rather than claiming this is the original October 1 revision.
Source: vllm-gguf-plugin-prism/vllm_gguf_plugin/triton/pq2_int_gemv.py:33–54.
@triton.jit
def _dot_groups(Q, W, tokens, rows, block, M, N: tl.constexpr, K: tl.constexpr,
STRIDE: tl.constexpr, PREPARED: tl.constexpr, CHAINED: tl.constexpr):
valid = block < K // 128
mask = (rows[:, None] < N) & valid[None, :]
words = tl.arange(0, 4 if CHAINED else 32)
dots = tl.full((tokens.shape[0], rows.shape[0], block.shape[0], words.shape[0]), 0, tl.int32)
for j in tl.static_range(8 if CHAINED else 1):
groups = words + j * 4
q = tl.load(Q + tokens[:, None, None] * (K // 4)
+ block[None, :, None] * 32 + groups[None, None, :],
(tokens[:, None, None] < M) & valid[None, :, None], other=0)
packed = _pq2_codes(W, rows[:, None, None], block[None, :, None],
groups[None, None, :], K, STRIDE,
mask[:, :, None], PREPARED).to(tl.int32)
w = _decode_codes(packed)
dots = tl.inline_asm_elementwise(
"dp4a.s32.s32 $0, $1, $2, $3;", constraints="=r,r,r,r",
args=[q[:, None, :, :], w[None, :, :, :], dots],
dtype=tl.int32, is_pure=True, pack=1,
)
return tl.sum(dots, 3).to(tl.float32)
Original benchmark numerical comparison for the integer API.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/benchmark_prism.py:113–135.
)
if not math.isfinite(mmq_relative_rms) or mmq_relative_rms > 0.03:
raise AssertionError(f"MMQ relative RMS {mmq_relative_rms} exceeds 0.03")
int_gemv_first_call_ms = None
int_gemv_relative_rms = None
if "int_gemv" in methods:
torch.cuda.synchronize()
start = time.perf_counter()
candidate = pq2_int_gemv(x, w)
torch.cuda.synchronize()
int_gemv_first_call_ms = (time.perf_counter() - start) * 1000
int_gemv_relative_rms = (
(
(candidate.float() - expected.float()).square().sum()
/ expected.float().square().sum()
)
.sqrt()
.item()
)
if not math.isfinite(int_gemv_relative_rms) or int_gemv_relative_rms > 0.03:
raise AssertionError(
f"Integer GEMV relative RMS {int_gemv_relative_rms} exceeds 0.03"
)
Test 7. INT8 tensor-core MMQ prefill and full shape coverage (Oct 1)
Visual reference: Diagram 2 · Topic overview.
The idea: We round input blocks to 8-bit integers and uses the GPU's matrix-multiply hardware to process a tile of tokens and output rows together. It checks whether this works better than the floating packed-weight kernel as the number of tokens grows.
The result: At larger M, gate/up reaches about 1.5× native, down reaches 1.42×, and the head reaches 1.59× at M = 1,024. All single-token shapes slow down. Relative RMS error in the full matrix is 0.457–0.954% against the rounded floating reference. The prototype reports 2–4 spills.
Explanation, checks and code
Why it matters on the GPU: Tensor Cores execute small matrix products in one instruction sequence. They need enough rows and columns to do useful work; tiny batches can leave much of a tile unused. The kernel decodes only the required packed weights, accumulates each group's integer products, then applies that group's input and weight scales.
More tokens allow more reuse of each weight tile, but large tiles put pressure on registers and shared memory.
What to take from it: The larger-token wins and single-token losses identify a workload crossover. The prototype's spills show a resource cost as well. The result supports a selective Tensor Core path, rather than sending every matrix size through it.
Setup: The kernel uses group-128 Q8 inputs, an INT32 Tensor Core dot for each group, and FP32 scaling. An initial FFN tile search is followed by 400 records—200 pairs—covering five shapes and M from 1 to 1,024. Timings include input quantisation.
Relevant code:
Saved earlier MMQ implementation: one INT8 tensor-core dot and scale application per group.
Source: pq2-dev/benchmarks/results/pq2-mmq-unpack-20261002/reference-kernel.py:31–61.
@triton.jit(do_not_specialize=["M"])
def _mmq(
Q,
S,
W,
Y,
M,
N: tl.constexpr,
K: tl.constexpr,
STRIDE: tl.constexpr,
BM: tl.constexpr,
BN: tl.constexpr,
PREPARED: tl.constexpr = False,
):
m = tl.program_id(0) * BM + tl.arange(0, BM)
n = tl.program_id(1) * BN + tl.arange(0, BN)
k = tl.arange(0, 128)
acc = tl.full((BM, BN), 0, tl.float32)
for block in range(K // 128):
a = tl.load(
Q + m[:, None] * K + block * 128 + k[None, :], m[:, None] < M, other=0
)
ws = _pq2_scale(W, n, block, N, K, STRIDE, n < N, PREPARED)
packed = _pq2_codes(
W, n[None, :], block, k[:, None] // 4, K, STRIDE, n[None, :] < N, PREPARED
)
b = (((packed.to(tl.int32) >> (2 * (k[:, None] % 4))) & 3) - 1).to(tl.int8)
partial = tl.dot(a, b).to(tl.float32)
xs = tl.load(S + m * (K // 128) + block, m < M, other=0)
acc += partial * xs[:, None] * ws[None, :]
tl.store(Y + m[:, None] * N + n[None, :], acc, (m[:, None] < M) & (n[None, :] < N))
The standalone tile sweep launches the candidate MMQ geometry and times it.
Source: pq2-dev/benchmarks/results/pq2-mmq-20261001/tile_sweep.py:14–33.
for bm, bn, warps in [
(32, 64, 4),
(64, 64, 4),
(64, 128, 4),
(32, 128, 4),
(64, 64, 8),
(64, 128, 8),
(128, 64, 8),
]:
def fn():
_mmq[(triton.cdiv(m, bm), triton.cdiv(n, bn))](
q, s, w, y, m, n, k, w.stride(0), bm, bn, num_warps=warps
)
try:
fn()
print(n, k, bm, bn, warps, do_bench_cudagraph(fn, rep=100), flush=True)
except Exception as e:
print(type(e).__name__, str(e)[:100], flush=True)
C. Activation groups: numerical error versus overhead
Takeaway. Smaller activation groups reduced numerical error, but added latency. Better projection accuracy alone does not establish better answers.
Terms for this section: Groups, planes and fragments.
Diagram 3: Changing how many input values share a scale
Smaller groups give local values a more suitable quantisation scale, but require more scale storage and separately scaled work. Test 8 isolates numerical error on the CPU; test 9 measures GPU error and latency. The weight format stays unchanged.
Test 8. CPU activation-group numerical control (Oct 1)
Visual reference: Diagram 3 · Topic overview.
The idea: We do an offline calculation on the CPU using saved inputs. It reconstructs the same compressed weights and changes only how many activation values share a quantisation scale, so the numerical effect of group size can be studied in isolation.
The result: The group-128 reconstruction matches the saved quantised FP64 reference within 1e-12. Mean activation relative RMS falls from 0.00540395 to 0.00460424, improving all 54 samples. Mean projection relative RMS falls from 0.00368201 to 0.00315170, improving 50/54. This test has no GPU latency result.
Explanation, checks and code
Why it matters on the GPU: It deliberately does not measure GPU memory layout, occupancy or compute-core speed. The scale and code reconstruction mirrors the formats used by the GPU, while 64-bit accumulation reduces confounding from floating-point reduction order. Smaller groups would require more scales and separately scaled work on the GPU, but those costs must come from the separate GPU timing test.
What to take from it: The error reduction supports the reason for trying smaller groups. It does not establish that smaller groups run faster, or that the model's final answers improve.
Setup: We replay 54 saved samples on the CPU. Rounding, FP32 activation scales, FP16 weight scales and FP64 accumulation stay fixed; only activation group size changes between 128 and 32.
Relevant code:
The diagnostic uses relative RMS for both activation and projection errors.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/compare_pq2_q8_groups.py:11–12.
def relative_rms(actual, expected):
return ((actual - expected).square().sum() / expected.square().sum()).sqrt().item()
CPU replay holds the weights and scale rounding fixed while changing activation group size.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/compare_pq2_q8_groups.py:22–51.
for probe in json.loads((args.captures / "projection-comparison.json").read_text()):
saved = torch.load(args.captures / probe["artifact"], weights_only=True)
x = saved["x"].float()
m, k = x.shape
blocks = saved["packed_rows"].reshape(-1, k // 128, 34)
scales = blocks[:, :, :2].contiguous().view(torch.float16).double()
codes = ((blocks[:, :, 2:, None].int() >> (2 * torch.arange(4))) & 3) - 1
weights = (codes.reshape(-1, k // 128, 128).double() * scales).reshape(-1, k)
reference = x.double() @ weights.T
metrics = {}
for group in (128, 32):
activations = x.reshape(m, k // group, group)
# Triton lowers division by this constant to an FP32 reciprocal
# multiply. Preserve that scale rounding in both group variants.
scale = activations.abs().amax(-1, keepdim=True) * (1.0 / 127)
quantized = (
(activations / torch.where(scale > 0, scale, 1))
.round()
.clamp(-127, 127)
)
reconstructed = (quantized.double() * scale.double()).reshape(m, k)
projected = reconstructed @ weights.T
if group == 128:
torch.testing.assert_close(
projected, saved["quantized_fp64"], atol=1e-12, rtol=1e-12
)
metrics[str(group)] = dict(
activation_relative_rms=relative_rms(reconstructed, x.double()),
projection_relative_rms=relative_rms(projected, reference),
)
Test 9. Smaller Q8 activation groups (Oct 2)
Visual reference: Diagram 3 · Topic overview.
The idea: We give smaller groups of input values their own rounding scale. It asks how much accuracy improves when a large value can affect only 32 or 64 neighbours instead of 128, and how much slower that becomes.
The result: Mean RMSE falls from 0.004343 with group 128 to 0.004134 with group 64 and 0.003897 with group 32. Groups 64 and 32 improve 38/54 and 50/54 probes respectively. Their median latency costs are +10.2% and +52.4%.
Explanation, checks and code
Why it matters on the GPU: The compressed weights and their 128-value weight scales stay the same. Smaller activation groups add more input-scale values and more separately scaled partial dot products. Tensor Cores can still perform the integer products, but the program must handle more group boundaries, scale loads and floating-point combinations. The activation buffer also needs more scale storage.
What to take from it: The saved-input comparisons show a modest error reduction alongside a measurable latency cost. This is an accuracy-versus-overhead trade-off; it does not show that the extra scale traffic alone caused the slowdown.
Setup: We replay 54 saved projections against FP64 references. The separate timing comparison has 20 complete-call cases, with groups of 128, 64 and 32, raw/prepared weights, and M = 128 or 512. Creating the original captures is outside the test.
Relevant code:
Group-size variants use identical activation tensors and compare both timing orders.
Source: pq2-dev/benchmarks/experiments/integer-decode/pq2_mmq_groups.py:50–73.
for prepared in (False, True):
w = prepare_pq2_layout(weight) if prepared else weight
for m in (128, 512):
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
methods = {
str(g): partial(
pq2_mmq, x, w, prepared=prepared, activation_group=g
)
for g in (128, 64, 32)
}
if reference:
torch.testing.assert_close(
methods["128"]().view(torch.uint8),
reference.pq2_mmq(x, w, prepared=prepared).view(torch.uint8),
atol=0,
rtol=0,
)
# Errors are measured separately on saved real activations.
for method in methods.values():
method()
timings = [
{g: do_bench_cudagraph(methods[g], rep=100) for g in order}
for order in (list(methods), list(reversed(methods)))
]
The separate saved-activation replay measures FP64-reference errors for each group size.
Source: pq2-dev/benchmarks/experiments/integer-decode/pq2_mmq_groups.py:87–109.
errors = []
if args.captures:
for path in sorted(args.captures.glob("projection-*.pt")):
saved = torch.load(path, weights_only=True)
x, w = saved["x"].cuda(), saved["packed_rows"].cuda()
fp64 = saved["fp64"].double()
row = {
"artifact": path.name,
"m": x.shape[0],
"n": w.shape[0],
"k": x.shape[1],
"groups": {},
}
for group in (128, 64, 32):
y = pq2_mmq(x, w, activation_group=group).cpu().double()
delta = y - fp64
row["groups"][str(group)] = {
"rmse": delta.square().mean().sqrt().item(),
"relative_l2": (delta.norm() / fp64.norm().clamp_min(1e-30)).item(),
"max_abs": delta.abs().max().item(),
}
errors.append(row)
del x, w
D. Tiles, warps and specialisation
Takeaway. Divide work to suit the actual matrix and batch size. Confirm promising changes on real inputs; hand-written CUDA is not automatically faster.
Terms for this section: Tiles · Warps and reductions · Butterflies and shuffles.
Diagram 4: Dividing a projection into work that fits the GPU
This is a conceptual scheduling view, not a literal mapping of one matrix cell to one thread. Test 10 changes compile-time knowledge; tests 11, 12–14 and 15 change geometry or specialisation; test 16 fits floating tiles to eight rows; test 17 changes warp reductions and persistent scheduling.
Test 10. Runtime-M specialisation and native matrix (Oct 1)
Visual reference: Diagram 4 · Topic overview.
The idea: We check whether one compiled matrix-multiply program can handle several numbers of input tokens, instead of compiling a new program for every token count. The test feeds both versions exactly the same weights and inputs, checks their answers, and measures first-use delay separately from repeated execution.
The result: All 14 paired outputs match bit for bit. After the first runtime-M compilation, subsequent first calls take 0.68–9.71 ms; compiling a new fixed-M variant takes 338–435 ms. Large-M steady-state changes range from −1.1% to +0.4%, while small FFN-down cases slow down by 8.2–8.9%.
Explanation, checks and code
Why it matters on the GPU: The weights stay in the same places in GPU memory and the basic arithmetic stays the same. The change is what the compiler knows in advance. With a fixed token count it can remove some boundary checks and simplify work; with a runtime count the GPU must handle valid and padded rows more generally.
This is primarily a compilation-versus-execution trade-off, rather than a new memory format or a move to different compute cores.
What to take from it: It substantially reduces repeated compilation, but the small-matrix slowdowns show that compiler flexibility can cost execution efficiency. Exact output agreement means this trade-off does not change the tested answers.
Setup: We use synthetic matrices for 240 cases, then compare runtime-M and fixed-M versions directly in 14 paired cases. Timings use warmed CUDA graphs on an RTX 3090 Ti.
Relevant code:
The same inputs are passed to constexpr-M and runtime-M implementations; cold calls, warmed graph samples and exact parity are checked.
Source: pq2-dev/benchmarks/results/pq2-kernels-20261001/compare_runtime_m.py:20–38.
w = torch.randint(0, 256, (n, k // 128, 34), device="cuda", dtype=torch.uint8)
scales = torch.rand(n, k // 128, device="cuda", dtype=torch.float16) * 0.02
w[:, :, :2] = scales.view(torch.uint8).reshape(n, k // 128, 2)
w = w.reshape(n, -1)
for m in (5, 8, 13, 128, 256, 258, 1024):
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
outputs = []
for name, fn in (("constexpr", reference.pq2_matmul), ("runtime", pq2_matmul)):
torch.cuda.synchronize()
start = time.perf_counter()
outputs.append(fn(x, w))
torch.cuda.synchronize()
first_ms = (time.perf_counter() - start) * 1000
for _ in range(3):
fn(x, w)
samples = [do_bench_cudagraph(lambda: fn(x, w), rep=100) for _ in range(3)]
writer.writerow([n, k, m, name, first_ms, sorted(samples)[1]])
stream.flush()
torch.testing.assert_close(outputs[0], outputs[1], atol=0, rtol=0)
Test 11. Integer tile tuning (Oct 1)
Visual reference: Diagram 4 · Topic overview.
The idea: We change how many output rows and 128-value groups one program handles, and how many groups of threads work on it. It searches for the most efficient way to divide the same integer matrix multiply across the GPU, then confirms the winners with input preparation included.
The result: Every confirmed case improves, with speedups of 1.068–1.158×. The selected settings report no spills. The largest relative RMS difference from the existing quantised kernel is 0.000025065.
Explanation, checks and code
Why it matters on the GPU: A tile is a chunk of the matrix assigned to a program. Larger tiles can reuse inputs and reduce scheduling overhead, but need more registers and possibly shared memory. More warps can expose more parallel work, but they also share each multiprocessor's fixed resources. The test balances reuse, resource use and parallelism without changing the stored weight format or DP4A instruction family.
What to take from it: These settings help in the confirmed cases. Zero spills means the compiler did not report overflowing register storage; it doesn’t mean that the GPU spent no time waiting for memory.
Setup: We screen 24 tile settings across two FFN shapes at M = 1, 4 and 8. We then check complete calls in 24 dtype/stride cases, timing both execution orders.
Relevant code:
The candidate tile is passed directly into the integer GEMV launch.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/tune_pq2_int_gemv.py:29–43.
def launch(q, s, w, y, m, n, k, bn, bc, warps):
return _int_gemv[(triton.cdiv(n, bn),)](
q,
s,
w,
y,
m,
n,
k,
w.stride(0),
triton.next_power_of_2(m),
bn,
bc,
num_warps=warps,
)
Confirmation includes activation quantisation and scratch allocation.
Source: pq2-dev/benchmarks/plugin-archive/benchmarks/confirm_pq2_int_tiles.py:18–26.
def complete(x, w, tile):
m, k = x.shape
n = w.shape[0]
q = torch.empty((m, k), dtype=torch.int8, device=x.device)
scales = torch.empty((m, k // 128), dtype=torch.float32, device=x.device)
output = torch.empty((m, n), dtype=x.dtype, device=x.device)
_quantize[(m * (k // 128),)](x, q, scales, k)
launch(q.view(torch.int32), scales, w, output, m, n, k, *tile)
return output
Test 12. Single-token geometry, PRMT, chaining and fusion screens (Oct 5)
Visual reference: Diagram 4 · Topic overview.
The idea: We change several parts of a one-token projection: the tile shape, how packed codes are decoded, whether integer dot products share a running accumulator, and whether neighbouring preparation operations are fused into the projection.
The result: Gate/up takes 92.56 µs at baseline, 84.18 µs with the shape change, 85.47 µs with PRMT, and 103.34 µs with chaining. Down takes 50.96 µs at baseline, then 46.85, 48.01 and 60.25 µs respectively. Chaining is slower in this one-token test.
Explanation, checks and code
Why it matters on the GPU: PRMT rearranges bytes in registers to turn packed codes into the signed values expected by DP4A. Chaining DP4A instructions can reduce separate additions, but each instruction then depends on the previous answer and may expose less independent work.
Fusion can avoid writing an intermediate tensor to memory and launching another kernel, at the cost of keeping more work and data in one program. Geometry determines how the remaining work is spread across multiprocessors.
What to take from it: The screen separates these alternatives instead of assuming that fewer instructions or fewer launches always wins. Chaining was slower at one token here; that does not contradict its later benefit with eight tokens.
Setup: We use synthetic BF16 inputs and prepared weights, timing complete calls including quantisation over three passes. Fusion candidates get their own equivalent unfused baselines.
Relevant code:
Candidate geometry, PRMT, chaining and fusion calls are timed against the corresponding standalone baselines.
Source: pq2-dev/benchmarks/experiments/single-request/pq2_all_decode_microbench.py:27–60.
s=torch.randn(n,k//128,device='cuda',dtype=torch.float16)*0.01
w[:,:,:2]=s.view(torch.uint8).reshape(n,k//128,2)
w=prepare_pq2_layout(w.reshape(n,-1))
x=torch.randn(1,k,device='cuda',dtype=torch.bfloat16)
def integer(z,variant='default',decode='shift',chained=False,gated=False,quantized=None):
return pq2_int_gemv(z,w,prepared=True,variant=variant,decode=decode,chained=chained,gated=gated,quantized=quantized)
methods={'default':lambda:integer(x),'shape':lambda:integer(x,variant='shape-tuned'),
'prmt':lambda:integer(x,decode='prmt'),'chained':lambda:integer(x,chained=True),
'prmt-chained':lambda:integer(x,decode='prmt',chained=True),
'shape-prmt-chained':lambda:integer(x,variant='shape-tuned',decode='prmt',chained=True)}
expected=methods['default']()
for name,method in methods.items():
actual=method();torch.testing.assert_close(actual,expected,atol=0.02,rtol=0.01)
if label=='attention_gdn_output':methods['existing_floating']=lambda:pq2_matmul(x,w,prepared=True)
def activated(z):
gate,up=z.chunk(2,-1)
return (torch.nn.functional.silu(gate.float())*up.float()).to(x.dtype)
if label=='gate_up':
methods['gate_baseline']=lambda:activated(integer(x))
methods['gate_fused']=lambda:integer(x,gated=True)
torch.testing.assert_close(methods['gate_fused'](),methods['gate_baseline'](),atol=0.03,rtol=0.01)
methods['hadamard_baseline']=lambda:integer(hadamard(x,None,1024))
methods['hadamard_fused']=lambda:integer(x,quantized=hadamard_quantize(x,None,1024))
torch.testing.assert_close(methods['hadamard_fused'](),methods['hadamard_baseline'](),atol=0,rtol=0)
if label=='gate_up':
methods['all_baseline']=lambda:activated(integer(hadamard(x,None,1024)))
methods['all']=lambda:integer(x,variant='shape-tuned',decode='prmt',chained=True,gated=True,quantized=hadamard_quantize(x,None,1024))
else:
methods['all_baseline']=methods['hadamard_baseline']
methods['all']=lambda:integer(x,variant='shape-tuned',decode='prmt',chained=True,quantized=hadamard_quantize(x,None,1024))
torch.testing.assert_close(methods['all'](),methods['all_baseline'](),atol=0.03,rtol=0.01)
timings={name:[] for name in methods}
for order in (list(methods),list(reversed(methods)),list(methods)):
for name in order:timings[name].append(do_bench_cudagraph(methods[name],rep=100))
Test 13. Real-weight geometry confirmation (Oct 5)
Visual reference: Diagram 4 · Topic overview.
The idea: We repeat the geometry search using actual model weights and saved token inputs. It rotates through different layers rather than repeatedly multiplying one small matrix, and uses longer alternating reference/candidate runs to reject apparent winners caused by timing noise.
The result: Confirmed projection gains are +7.83% for QKV, +6.82% for GDN input and +11.78% for the head. The apparent gate/up and output gains do not survive confirmation. Direct linked-replay gains are +1.62%, +1.72% and +2.41%. An earlier 1.217× replay result uses an older baseline and mustn't be added to these gains.
Explanation, checks and code
Why it matters on the GPU: A repeatedly reused matrix can remain in the GPU's L2 cache, which is much faster to access than board memory. Rotating through 6.33 GiB of distinct weights reduces that artificial advantage and better tests weight streaming. Tile choices also change how many thread blocks exist, how well they fill the multiprocessors, and how much register state each block needs.
What to take from it: Several initial wins disappear under confirmation. The surviving settings are more credible for these real-weight projection workloads, although this replay still omits the rest of a decode step.
Setup: We screen 180 geometries across six shapes, then run longer paired confirmations. The replay has 257 projections captured at decode steps 1, 32 and 128, 64 linked MLPs, and 6.33 GiB of distinct weights.
Relevant code:
Rotating per-shape projection calls, alternating timing pairs and all-position numerical checks determine which geometry survives.
Source: pq2-dev/benchmarks/experiments/single-request/pq2_geometry_confirm.py:25–56.
for ranking in screen['ranking']:
n,k=ranking['shape'];names=[name for name in r.order if [r.layers[name]['n'],r.layers[name]['k']]==[n,k]]
def run(config):
y=None
for name in names:
for _ in range(r.manifest['calls_per_decode_step'][name]):y=r.projection(name,r.samples[32,name]['x'],config)
return y
candidates=sorted([x for x in screen['records'] if x['shape']==[n,k] and x.get('speedup',0)>1.01],key=lambda x:-x['speedup'])[:3]
if not candidates:continue
control,out=graph_for(lambda:run(base))
for item in candidates:
config=base|{'geometry':{f'{n},{k}':item['geometry']}}
candidate,candidate_out=graph_for(lambda:run(config))
# Settle clock/cache state after compilation, then alternate order.
for _ in range(5):time_graph(control,100);time_graph(candidate,100)
ratios=[];pairs=[]
for repeat in range(7):
if repeat%2:
cv=time_graph(candidate,150);bv=time_graph(control,150)
else:
bv=time_graph(control,150);cv=time_graph(candidate,150)
ratios.append(cv/bv);pairs.append([bv,cv])
metrics=[]
for step in r.manifest['steps']:
for name in names:
x=r.samples[step,name]['x'];metrics.append(error_metrics(r.projection(name,x,config),r.projection(name,x,base)))
assert all(x['finite'] and x['relative_rms']<0.002 for x in metrics)
row={'shape':[n,k],'geometry':item['geometry'],'speedup':1/statistics.median(ratios),'min_speedup':1/max(ratios),'max_speedup':1/min(ratios),'pairs_ms':pairs,'checked_outputs':len(metrics),'max_relative_rms':max(x['relative_rms'] for x in metrics),'max_abs':max(x['max_abs'] for x in metrics)}
rows.append(row);write_json(a.output/'progress.json',rows);print(json.dumps(row),flush=True)
del candidate,candidate_out
eligible=[x for x in rows if x['shape']==[n,k] and x['speedup']>1.02 and x['min_speedup']>1.01]
if eligible:chosen[f'{n},{k}']=max(eligible,key=lambda x:x['speedup'])['geometry']
Test 14. Single-token structural specialisation (Oct 6)
Visual reference: Diagram 4 · Topic overview.
The idea: We use a simpler kernel specifically for one token. It removes the bookkeeping needed to handle several tokens and compares byte loads, wider loads, alternative packed layouts and deeper pipelining.
The result: The byte-load specialisation improves linked replay by 11.83%; v2 geometry gives 2.05%; together they give 14.08%. The specialisation matches exactly. Gate/up drops from 128 to 93 registers and from 8,192 to 1,024 shared-memory bytes, with no spills. Wider loads and deeper staging don't beat it.
Explanation, checks and code
Why it matters on the GPU: A single-token kernel needs fewer live dimensions and temporary values. That can reduce registers and shared memory per block, leaving more room for concurrent work on a multiprocessor. Independent DP4A operations also give the compiler opportunities to overlap instructions. Wider loads may save load instructions but introduce rearrangement or larger live tiles, so they can lose overall.
What to take from it: The exact-output specialisation wins and uses fewer measured compiler resources. Those observations are consistent with reduced resource pressure, but the experiment changes several things together and doesn’t isolate one hardware cause.
Setup: We rotate through real weights and one-row inputs captured at decode steps 1, 32 and 128. Comparisons use seven alternating graph-timing pairs, followed by confirmation in a fresh process.
Relevant code:
The single-row specialisation consumes packed tiles with independent DP4A operations.
Source: pq2-dev/benchmarks/harness/pq2_single_structural.py:39–48.
@triton.jit
def _consume(Q,S,packed,ws,blocks,K:tl.constexpr):
words=tl.arange(0,32)
q=tl.load(Q+blocks[:,None]*32+words[None,:],blocks[:,None]<K//128,other=0)
w=_decode_codes(packed,True)
dots=tl.inline_asm_elementwise('dp4a.s32.s32 $0, $1, $2, 0;',
constraints='=r,r,r',args=[q[None,:,:],w],dtype=tl.int32,is_pure=True,pack=1)
partial=tl.sum(dots,2).to(tl.float32)
xs=tl.load(S+blocks,blocks<K//128,other=0)
return partial*xs[None,:]*ws
The specialised loop removes the generic token dimension and tests load/pipeline choices.
Source: pq2-dev/benchmarks/harness/pq2_single_structural.py:51–73.
@triton.jit
def _single(Q,S,W,WS,Y,N:tl.constexpr,K:tl.constexpr,BN:tl.constexpr,BC:tl.constexpr,
WIDTH:tl.constexpr,ROW_TILE:tl.constexpr,STAGES:tl.constexpr,
PREFETCH: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)
count=tl.cdiv(K//128,BC)
if PREFETCH:
# Keep the next packed tile live while consuming the current tile.
blocks=offset
packed,ws=_load_tile(W,WS,rows,blocks,N,K,WIDTH,ROW_TILE)
for base in range(count):
next_blocks=(base+1)*BC+offset
next_packed,next_ws=_load_tile(W,WS,rows,next_blocks,N,K,WIDTH,ROW_TILE)
acc+=_consume(Q,S,packed,ws,base*BC+offset,K)
packed,ws=next_packed,next_ws
else:
for base in tl.range(0,count,num_stages=STAGES):
blocks=base*BC+offset
packed,ws=_load_tile(W,WS,rows,blocks,N,K,WIDTH,ROW_TILE)
acc+=_consume(Q,S,packed,ws,blocks,K)
tl.store(Y+rows,tl.sum(acc,1),rows<N)
Paired CUDA graph timing controls the comparison.
Source: pq2-dev/benchmarks/harness/pq2_single_structural.py:142–151.
def paired(candidate,control,ms=100):
for _ in range(3):time_graph(control,ms);time_graph(candidate,ms)
pairs=[]
for i in range(7):
if i%2:cv=time_graph(candidate,ms);bv=time_graph(control,ms)
else:bv=time_graph(control,ms);cv=time_graph(candidate,ms)
pairs.append([bv,cv])
ratios=[a/b for a,b in pairs]
return {'speedup':statistics.median(ratios),'min_speedup':min(ratios),'max_speedup':max(ratios),
'control_ms':statistics.median(x[0] for x in pairs),'candidate_ms':statistics.median(x[1] for x in pairs),'pairs_ms':pairs}
Test 15. Batch-eight geometry and chained DP4A (Oct 6)
Visual reference: Diagram 4 · Topic overview.
The idea: We run projections for eight saved token rows and varies tile sizes, token grouping and chained integer accumulation. A separate variant also changes the floating output projection to quantised integer arithmetic.
The result: No geometry passes the improvement gate. Chaining matches all 771 saved outputs bit for bit and improves linked replay by 3.60–4.00%. Integer output gives 10.57–10.99%, but changes the approximation: its largest relative RMS difference is 0.008532.
Explanation, checks and code
Why it matters on the GPU: Reusing a weight across eight tokens avoids rereading it, but the kernel must retain partial sums for all eight. That can inflate register use and integer instruction count. Smaller token groups reduce live data but can reread weights; placing related groups close together tries to reuse weights through L2 cache. Chained DP4A changes the accumulation dependency pattern without changing the packed weight bytes.
What to take from it: The geometry search found no reliable win. Chaining preserved the saved outputs and improved linked replay modestly. Integer output gained more but changed the approximation, so its speedup is a different kind of result.
Setup: We use 771 saved outputs from real eight-row inputs at three decode steps. The screen covers 260 combinations of row, reduction, warp and token-tile settings, followed by linked replay.
Relevant code:
The experimental token-tiled batched DP4A kernel.
Source: pq2-dev/benchmarks/harness/pq2_batch8_replay.py:12–39.
@triton.jit
def _token_tiled_gemv(Q, S, W, Y, M: tl.constexpr, N: tl.constexpr,
K: tl.constexpr, STRIDE: tl.constexpr, BT: tl.constexpr,
BN: tl.constexpr, BC: tl.constexpr, PREPARED: tl.constexpr,
INTERLEAVED: tl.constexpr = True):
if INTERLEAVED:
# Adjacent programs reuse one row tile across token groups through L2.
token_tiles = tl.cdiv(M, BT)
tokens = (tl.program_id(0) % token_tiles) * BT + tl.arange(0, BT)
rows = (tl.program_id(0) // token_tiles) * BN + tl.arange(0, BN)
else:
tokens = tl.program_id(1) * BT + tl.arange(0, BT)
rows = tl.program_id(0) * BN + tl.arange(0, BN)
offsets = tl.arange(0, BC)
acc = tl.full((BT, BN, BC), 0, tl.float32)
for base in range(tl.cdiv(K // 128, BC)):
block = base * BC + offsets
valid = block < K // 128
mask = (rows[:, None] < N) & valid[None, :]
ws = _pq2_scale(W, rows[:, None], block[None, :], N, K,
STRIDE, mask, PREPARED)
partial = _dot_groups(Q, W, tokens, rows, block, M, N, K,
STRIDE, PREPARED, True, False)
xs = tl.load(S + tokens[:, None] * (K // 128) + block[None, :],
(tokens[:, None] < M) & valid[None, :], other=0)
acc += partial * xs[:, None, :] * ws[None, :, :]
tl.store(Y + tokens[:, None] * N + rows[None, :], tl.sum(acc, 2),
(tokens[:, None] < M) & (rows[None, :] < N))
The follow-up switches chaining and integer output while retaining the equivalent transforms.
Source: pq2-dev/benchmarks/experiments/batch8/pq2_batch8_followups.py:20–40.
def projection(self, name, x, config, gated=False):
layer = self.layers[name]
n, k = layer['n'], layer['k']
integer = (n, k) in prism._BATCHED_GEMV_SHAPES or (config.get('integer_output') and (n, k) == (5120, 6144))
if not config.get('chained', False) or not integer:
return super().projection(name, x, config, gated=gated)
if layer['hadamard_block']:
x = prism.hadamard(x, layer['signs'], layer['hadamard_block'])
if layer['layout_object'] is not None:
x = layer['layout_object'].input_to_gguf(x)
m = x.shape[0]
q = torch.empty((m, k), device=x.device, dtype=torch.int8)
s = torch.empty((m, k // 128), device=x.device, dtype=torch.float32)
_quantize[(m * (k // 128),)](x.contiguous(), q, s, k)
y = torch.empty((m, n), device=x.device, dtype=x.dtype)
w = layer['weight']
_int_gemv[(triton.cdiv(n, 16),)](q.view(torch.int32), s, w, y, m, n, k, w.stride(0),
8, 16, 4, layer['prepared'], True, True, False, num_warps=4)
if layer['bias'] is not None:
y.add_(layer['bias'])
return y
Test 16. Eight-row structural CUDA/Triton replay (Oct 6)
Visual reference: Diagram 4 · Topic overview.
The idea: We compare several independently written CUDA and Triton ways of performing the same eight-row projections. It also tries smaller floating-point output tiles, different activation-group formats, and storing intermediate results in 32-bit floating point.
The result: The float-bm8 variant gives +9.22% linked replay and a 2.179× output-projection speedup, with exact outputs. Chained Triton gives +3.06%. CUDA native-r2-w2 gives −55.26% against the reference. Native Q8-interleaved gives +69.38% against its own slower group-32 control. All 14,649 projection checks pass, but one of 3,648 linked checks exceeds the error threshold.
Explanation, checks and code
Why it matters on the GPU: A floating tile sized for 64 token rows wastes work when only eight exist; using an eight-row tile fits the actual batch better. The CUDA prototypes divide reduction work among warps and combine their partial answers through shared memory. Interleaving quantised inputs with scales changes where each thread finds its operands. Each choice trades regular loads, register use, shared-memory traffic and synchronisation.
What to take from it: The smaller floating tile is a strong exact-output improvement. The large interleaved-packing percentage is relative to a slower CUDA prototype, not the best Triton baseline. One linked numerical failure also prevents treating every screened variant as accepted.
Setup: We compare 19 settings on real eight-row captures, checking 771 projections per setting. Timings use seven alternating pairs, with layout conversion left outside the measurement.
Relevant code:
The independently written CUDA projection kernel tests packing, scale groups, warp reductions and chained versus independent DP4A.
Source: pq2-dev/benchmarks/experiments/single-request/pq2_structural.cu:5–78.
// Independent benchmark kernel: two/four warps split K, each lane handles 32 weights.
template<int R, int WARPS, bool CHAIN, int GROUP>
__global__ void project(const int8_t* q, const void* scales, const uint8_t* w,
void* out, int n, int k, bool prepared, int group,
bool half_scale, bool packed_q, bool float_out) {
int lane = threadIdx.x, warp = threadIdx.y;
int tid = warp * 32 + lane;
float acc[8][R] = {};
for (int chunk = tid; chunk < k / 32; chunk += WARPS * 32) {
int wb = chunk / 4, sub = chunk % 4;
#pragma unroll
for (int r = 0; r < R; ++r) {
int row = blockIdx.x * R + r;
if (row >= n) continue;
const uint8_t* codes = prepared ? w + row * (k / 4) + wb * 32
: w + row * (k / 128 * 34) + wb * 34 + 2;
const half* ws = prepared ? reinterpret_cast<const half*>(w + n * (k / 4)) + row * (k / 128) + wb
: reinterpret_cast<const half*>(w + row * (k / 128 * 34) + wb * 34);
float weight_scale = __half2float(*ws);
int decoded[8];
#pragma unroll
for (int j = 0; j < 8; ++j) {
uint32_t b = codes[sub * 8 + j];
uint32_t selector = (b & 3) | (((b >> 2) & 3) << 4) |
(((b >> 4) & 3) << 8) | (((b >> 6) & 3) << 12);
decoded[j] = __byte_perm(0x020100ff, 0x020100ff, selector);
}
#pragma unroll
for (int token = 0; token < 8; ++token) {
const int* words = packed_q ? reinterpret_cast<const int*>(q + (token * (k / 32) + chunk) * 36 + 4)
: reinterpret_cast<const int*>(q + token * k + chunk * 32);
int sum = 0;
#pragma unroll
for (int j = 0; j < 8; ++j) {
if constexpr (CHAIN) sum = __dp4a(words[j], decoded[j], sum);
else sum += __dp4a(words[j], decoded[j], 0);
}
int si = token * (k / GROUP) + chunk / (GROUP / 32);
float xs = packed_q ? __half2float(*reinterpret_cast<const half*>(q + (token * (k / 32) + chunk) * 36))
: half_scale ? __half2float(static_cast<const half*>(scales)[si])
: static_cast<const float*>(scales)[si];
acc[token][r] += (float(sum) * xs) * weight_scale;
}
}
}
__shared__ float shared[WARPS][8][R];
#pragma unroll
for (int t = 0; t < 8; ++t) {
#pragma unroll
for (int r = 0; r < R; ++r) {
float v = acc[t][r];
for (int delta = 16; delta; delta >>= 1) v += __shfl_down_sync(0xffffffff, v, delta);
if (lane == 0) shared[warp][t][r] = v;
}
}
__syncthreads();
if (warp == 0 && lane == 0) {
#pragma unroll
for (int t = 0; t < 8; ++t) {
#pragma unroll
for (int r = 0; r < R; ++r) {
int row = blockIdx.x * R + r;
if (row < n) {
float v = 0;
#pragma unroll
for (int z = 0; z < WARPS; ++z) v += shared[z][t][r];
if (float_out) static_cast<float*>(out)[t * n + row] = v;
else static_cast<__nv_bfloat16*>(out)[t * n + row] = __float2bfloat16(v);
}
}
}
}
}
The structural variants are compared with alternating paired CUDA graph timings.
Source: pq2-dev/benchmarks/experiments/single-request/pq2_structural_replay.py:136–148.
def paired(candidate, control, passes=7, ms=100):
for _ in range(2):
time_graph(control,ms);time_graph(candidate,ms)
pairs=[]
for i in range(passes):
if i%2:
cv=time_graph(candidate,ms);bv=time_graph(control,ms)
else:
bv=time_graph(control,ms);cv=time_graph(candidate,ms)
pairs.append([bv,cv])
ratios=[bv/cv for bv,cv in pairs]
return {'control_ms':statistics.median(p[0] for p in pairs),'candidate_ms':statistics.median(p[1] for p in pairs),
'speedup':statistics.median(ratios),'min_speedup':min(ratios),'max_speedup':max(ratios),'pairs_ms':pairs}
Test 17. CUDA warp reductions and persistent Triton (Oct 6)
Visual reference: Diagram 4 · Topic overview.
The idea: We rewrite the projection using explicit CUDA warp-level operations, with two different ways of summing partial answers. It also tests a persistent Triton schedule in which a fixed pool of blocks repeatedly takes more tiles.
The result: The best CUDA/reference ratios are 0.664–0.770×: every tested shape is slower. Persistent Triton is roughly 0.940–1.030×. No CUDA shape passes the performance gate. Selected profiles and controls pass 3,084 projection, 768 linked and 48 graph checks; the losing candidates don't all receive that full validation.
Explanation, checks and code
Why it matters on the GPU: A warp is 32 threads executing together. Shuffles exchange values between its threads to combine dot products; the arrangement determines how many exchanges and dependencies are needed. Persistent blocks change how work is assigned to multiprocessors and can reduce poorly filled final waves. But too few persistent blocks can reduce concurrency, while too much per-warp work raises register demand.
What to take from it: The tested CUDA designs lose despite offering explicit control over the hardware. A handwritten kernel still has to beat the compiler's load scheduling, instruction layout and resource balance. The small persistent gains are shape-specific.
Setup: We screen 384 shape/setting pairs. CUDA variants try integer reduction or lane-local accumulation, 1, 2 or 4 groups and rows per warp, and 4 or 8 warps. Persistent grids use 1, 2, 4 or 8 blocks per SM.
Relevant code:
Standalone CUDA integer-reduction warp kernel.
Source: pq2-dev/benchmarks/cuda/pq2_warp_gemv.cu:7–49.
__device__ __forceinline__ int decode(unsigned x) {
unsigned selector=(x&3)|((x>>2&3)<<4)|((x>>4&3)<<8)|((x>>6&3)<<12);
return (int)__byte_perm(0x020100ff,0x020100ff,selector);
}
template<class Out,int G,int R,int W,bool P>
__global__ void warp_gemv(const int* __restrict__ q,const float* __restrict__ xs,
const uint8_t* __restrict__ codes,const half* __restrict__ ws,
Out* __restrict__ y,int n,int groups) {
int lane=threadIdx.x&31,warp=threadIdx.x>>5;
constexpr int L=32/G;
int sub=lane/L,local=lane%L;
int tiles=(n+W*R-1)/(W*R);
for(int tile=blockIdx.x;tile<tiles;tile+=P?gridDim.x:tiles) {
int row=(tile*W+warp)*R;
float acc[R]={};
for(int base=0;base<groups;base+=G) {
int group=base+sub;
int activation[4];
#pragma unroll
for(int j=0;j<G;j++)activation[j]=group<groups?q[group*32+local+j*L]:0;
#pragma unroll
for(int r=0;r<R;r++) {
int dot=0;
if(row+r<n && group<groups) {
const uint8_t* block=codes+(size_t)(row+r)*groups*32+group*32;
#pragma unroll
for(int j=0;j<G;j++)dot=__dp4a(activation[j],decode(block[local+j*L]),dot);
}
#pragma unroll
for(int d=L/2;d>0;d/=2)dot+=__shfl_down_sync(0xffffffff,dot,d,L);
if(local==0 && row+r<n && group<groups)
acc[r]+=((float)dot*xs[group])*__half2float(ws[(size_t)(row+r)*groups+group]);
}
}
#pragma unroll
for(int r=0;r<R;r++) {
float sum=acc[r];
#pragma unroll
for(int d=16;d>0;d/=2)sum+=__shfl_down_sync(0xffffffff,sum,d);
if(lane==0 && row+r<n)y[row+r]=Out(sum);
}
}
}
Lane-local scaled accumulation variant.
Source: pq2-dev/benchmarks/cuda/pq2_float_warp_gemv.cu:7–49.
__device__ __forceinline__ int decode(unsigned x) {
unsigned selector=(x&3)|((x>>2&3)<<4)|((x>>4&3)<<8)|((x>>6&3)<<12);
return (int)__byte_perm(0x020100ff,0x020100ff,selector);
}
template<class Out,int G,int R,int W,bool P>
__global__ void warp_gemv(const int* __restrict__ q,const float* __restrict__ xs,
const uint8_t* __restrict__ codes,const half* __restrict__ ws,
Out* __restrict__ y,int n,int groups) {
int lane=threadIdx.x&31,warp=threadIdx.x>>5;
constexpr int L=32/G;
int sub=lane/L,local=lane%L;
int tiles=(n+W*R-1)/(W*R);
for(int tile=blockIdx.x;tile<tiles;tile+=P?gridDim.x:tiles) {
int row=(tile*W+warp)*R;
float acc[R]={};
for(int base=0;base<groups;base+=G) {
int group=base+sub;
int activation[4];
#pragma unroll
for(int j=0;j<G;j++)activation[j]=group<groups?q[group*32+local*G+j]:0;
#pragma unroll
for(int r=0;r<R;r++) {
int dot=0;
if(row+r<n && group<groups) {
const uint8_t* block=codes+(size_t)(row+r)*groups*32+group*32;
unsigned packed;
if constexpr(G==4)packed=((const unsigned*)block)[local];
else if constexpr(G==2)packed=((const unsigned short*)block)[local];
else packed=block[local];
#pragma unroll
for(int j=0;j<G;j++)dot=__dp4a(activation[j],decode((packed>>(8*j))&255),dot);
}
if(row+r<n && group<groups)
acc[r]+=((float)dot*xs[group])*__half2float(ws[(size_t)(row+r)*groups+group]);
}
}
#pragma unroll
for(int r=0;r<R;r++) {
float sum=acc[r];
#pragma unroll
for(int d=16;d>0;d/=2)sum+=__shfl_down_sync(0xffffffff,sum,d);
if(lane==0 && row+r<n)y[row+r]=Out(sum);
}
Persistent Triton launch uses a fixed number of blocks per SM.
Source: pq2-dev/benchmarks/experiments/kernel-latency/pq2_warp_standalone.py:19–27.
@triton.jit
def persistent(Q,S,W,WS,Y,N:tl.constexpr,K:tl.constexpr,BN:tl.constexpr,BC:tl.constexpr):
for tile in range(tl.program_id(0),tl.cdiv(N,BN),tl.num_programs(0)):
rows=tile*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)
E. Keeping preparation data close to compute
Takeaway. Combining preparation steps saved work, but the gain became much smaller once the projection was included.
Terms for this section: Butterflies and shuffles · Occupancy and pipelining.
Diagram 5: Separate preparation kernels versus fused preparation
The fused path reduces intermediate memory round trips and launches. Test 18 also has to preserve reduction order and rounding: a tiny floating-point difference can change an INT8 code. The diagram does not claim that every intermediate fits only in registers.
Test 18. RMSNorm + Hadamard + Q8 fusion (Oct 6)
Visual reference: Diagram 5 · Topic overview.
The idea: We combine normalisation, the Hadamard mixing transform and input quantisation so that preparation can pass values directly between steps. It checks that combining them preserves the same rounded values, quantised codes and residual output.
The result: The first 12 variants fail the numerical checks. Accepted variants keep0 and keep1 give preparation speedups of 1.3529× and 1.2644×; including GEMV, these become 1.0331× and 1.0261×. All 768 BF16 checks match codes, projections and residuals exactly, with another 16 graph checks.
Explanation, checks and code
Why it matters on the GPU: Separate kernels normally write an intermediate array to GPU memory and a later kernel reads it back. Fusion can keep intermediate values in registers or shared memory and remove launches. But normalisation requires a sum across a row, and floating-point sums depend slightly on their order.
A small rounding change can push a value across an INT8 rounding boundary. Matching the reference reduction structure is therefore part of the optimisation.
What to take from it: Many initially faster-looking structures fail the numerical gate. The accepted version improves preparation substantially, while preparation plus the much larger GEMV improves only a few percent.
Setup: We use 384 raw normalisation samples from 128 boundaries at three decode positions. The accepted variant matches the compiled reference's 2,048-element reduction order using 16 warps. Timings rotate weights and alternate graph pairs.
Relevant code:
Matched normalisation reduction order preserves the compiled-reference rounding.
Source: pq2-dev/benchmarks/experiments/kernel-latency/pq2_norm_quant_kernels.py:27–38.
@triton.jit
def _compiled_stats(X,R,ROUT,INV,K:tl.constexpr,EPS:tl.constexpr,HAS_R:tl.constexpr):
i=tl.arange(0,2048);acc=tl.full((2048,),0,tl.float32)
for offset in tl.range(0,K,2048):
j=offset+i;x=tl.load(X+j,j<K,other=0).to(tl.float32)
if HAS_R:x+=tl.load(R+j,j<K,other=0).to(tl.float32)
acc=tl.where(j<K,acc+x*x,acc)
variance=tl.sum(acc,0)/K;tl.store(INV,tl.extra.cuda.libdevice.rsqrt(variance+EPS))
if HAS_R:
for offset in tl.range(0,K,2048):
j=offset+i;x=tl.load(X+j,j<K,other=0).to(tl.float32)+tl.load(R+j,j<K,other=0).to(tl.float32)
tl.store(ROUT+j,x,j<K)
The fused preparation tile applies norm, Hadamard, dtype rounding and Q8 conversion.
Source: pq2-dev/benchmarks/experiments/kernel-latency/pq2_norm_quant_kernels.py:40–49.
@triton.jit
def _tile(X,R,W,SIGNS,INV,Q,S,NORM,K:tl.constexpr,HAS_R:tl.constexpr,KEEP_NORM:tl.constexpr):
i=tl.program_id(0)*1024+tl.arange(0,1024);x=tl.load(X+i).to(tl.float32)
if HAS_R:x+=tl.load(R+i).to(tl.float32)
w=tl.load(W+i).to(tl.float32)+1.
normalized=((x*tl.load(INV))*w).to(X.dtype.element_ty)
if KEEP_NORM:tl.store(NORM+i,normalized)
v=normalized.to(tl.float32)*tl.load(SIGNS+i).to(tl.float32)
v=_fwht(v,1024).to(X.dtype.element_ty).to(tl.float32)
_q8(v,Q,S,tl.program_id(0)*1024,K,1024)
The final quantisation retains group-128 arithmetic.
Source: pq2-dev/benchmarks/experiments/kernel-latency/pq2_norm_quant_kernels.py:5–11.
@triton.jit
def _q8(v,Q,S,offset,K:tl.constexpr,B:tl.constexpr):
groups=tl.reshape(v,(B//128,128));scale=tl.max(tl.abs(groups),1)/127.
normalized=tl.div_rn(groups,tl.where(scale[:,None]>0,scale[:,None],1.))
q=tl.minimum(127.,tl.maximum(-127.,tl.extra.cuda.libdevice.nearbyint(normalized)))
i=offset+tl.arange(0,B);tl.store(Q+i,tl.reshape(q,(B,)).to(tl.int8),i<K)
gi=offset//128+tl.arange(0,B//128);tl.store(S+gi,scale,gi<K//128)
F. Diagnosing and hiding memory waits
Takeaway. Waiting for data does not mean every attempt to prefetch it will help. Extra buffers and registers can leave fewer warps ready to run.
Terms for this section: Occupancy and pipelining.
Diagram 6: Why starting a load earlier can help or hurt
The diagram shows the overlap case. Without prefetch, the sequence is load A → compute A → load B → compute B. Arrows show dependencies, not a measured time scale. Test 19 measures stalls and resource use; test 20 tries overlap, unrolling and register caps. Overlap helps only if its storage and scheduling costs do not erase the saved waiting time.
Test 19. GEMV resource and stall diagnostics (Oct 6)
Visual reference: Diagram 6 · Topic overview.
The idea: We ask the GPU profiler what happens while each representative projection runs: how often instructions wait for loads, how many warps are active, and how busy memory and execution hardware are. It is a diagnosis of the current kernel, rather than a replacement kernel.
The result: Long-scoreboard stalls account for 46–67% of the reported warp-cycle measure. The scheduler has no eligible warp in 42–58% of cycles, and achieved occupancy is 30–65%. No local load/store instructions appear in this collection. The head reaches 90.2% of peak DRAM throughput. These are isolated profiling observations.
Explanation, checks and code
Why it matters on the GPU: A multiprocessor can switch to another ready warp while one waits for data. Registers and shared memory limit how many warps can reside there. A long-scoreboard stall indicates waiting for an outstanding memory-related dependency; a cycle with no eligible warp means the scheduler has no ready warp to issue. The head's high DRAM utilisation suggests less remaining bandwidth headroom than some other shapes.
What to take from it: The counters give us a reason to try earlier loads and different layouts. They do not show that raising occupancy alone will help: more resident warps may require smaller tiles, more instructions or register spills.
Setup: We profile isolated launches for six real-weight shapes using captured decode-step-32 activations and the Hadamard transform. Nsight collects separate passes with uncontrolled caches.
Relevant code:
The profiler target runs one representative real-weight GEMV per shape with NVTX labels and finite-output checks.
Source: pq2-dev/benchmarks/experiments/single-request/pq2_gemv_counters.py:15–25.
r=Replay(ROOT/'pq2-dev/benchmarks/results/pq2-real-decode-replay-20261005/capture');rows=[]
names=list({(x['n'],x['k']):name for name,x in r.layers.items()}.values())
for name in names:
layer=r.layers[name];x=r.samples[32,name]['x']
if layer['hadamard_block']:x=prism.hadamard(x,layer['signs'],layer['hadamard_block'])
if layer['layout_object']:x=layer['layout_object'].input_to_gguf(x)
torch.cuda.nvtx.range_push(f"PQ2:{layer['n']},{layer['k']}")
y=pq2_int_gemv(x,layer['weight'],prepared=True)
torch.cuda.synchronize();torch.cuda.nvtx.range_pop()
rows.append({'name':name,'n':layer['n'],'k':layer['k'],'finite':bool(y.isfinite().all())})
write_json(a.output,{'complete':True,'shapes':rows})
Resource inspection launches the kernel and queries compiler/driver occupancy information; these are distinct from measured stall counters.
Source: pq2-dev/benchmarks/experiments/single-request/pq2_gemv_resources.py:20–45.
def attr(code):
v=ctypes.c_int();assert lib.cuDeviceGetAttribute(ctypes.byref(v),code,0)==0;return v.value
with gpu_lock(0):
r=Replay(ROOT/'pq2-dev/benchmarks/results/pq2-real-decode-replay-20261005/capture');entries={};rows=[]
for name,layer in r.layers.items():
n,k=layer['n'],layer['k'];x=prism.hadamard(r.samples[32,name]['x'],layer['signs'],layer['hadamard_block'])
q=torch.empty_like(x,dtype=torch.int8);s=torch.empty((1,k//128),device=x.device,dtype=torch.float32);_quantize[(k//128,)](x,q,s,k)
y=torch.empty((1,n),device=x.device,dtype=x.dtype)
bn,bc=(32,8) if n==34816 else ((16,16) if k in (6144,17408) else ((64,4) if n==248320 else (16,8)))
entries[name]=(q,s,layer['weight'],y,n,k,bn,bc)
for n,k in dict.fromkeys((l['n'],l['k']) for l in r.layers.values()):
names=[name for name,l in r.layers.items() if (l['n'],l['k'])==(n,k)]
kernel=launch_single(*entries[names[0]]);blocks=ctypes.c_int();assert lib.cuOccupancyMaxActiveBlocksPerMultiprocessor(ctypes.byref(blocks),kernel.function,128,kernel.metadata.shared)==0
static=ctypes.c_int();assert lib.cuFuncGetAttribute(ctypes.byref(static),1,kernel.function)==0
def run():
for name in names:
for _ in range(r.manifest['calls_per_decode_step'][name]):launch_single(*entries[name])
return entries[names[-1]][3]
graph,_=graph_for(run)
for _ in range(3):time_graph(graph,100)
times=[time_graph(graph,150) for _ in range(7)];calls=sum(r.manifest['calls_per_decode_step'][name] for name in names)
bytes_per_call=n*k//4+n*(k//128)*2
ms=statistics.median(times);row={'shape':[n,k],'calls':calls,'registers':kernel.n_regs,'spills':kernel.n_spills,'static_shared_bytes':static.value,'dynamic_shared_bytes':kernel.metadata.shared,'blocks_per_sm':blocks.value,'theoretical_active_warps':blocks.value*4,'theoretical_occupancy_pct':blocks.value*128/attr(39)*100,'bn':entries[names[0]][6],'bc':entries[names[0]][7],'median_ms':ms,'samples_ms':times,'logical_weight_bytes':bytes_per_call*calls,'logical_weight_GB_s':bytes_per_call*calls/(ms*1e6),'scope':'Logical code+scale bytes / elapsed time, not measured DRAM traffic.'}
rows.append(row);write_json(a.output/'resources.json',{'complete':False,'rows':rows,'device':{'max_threads_per_sm':attr(39),'registers_per_sm':attr(82),'shared_bytes_per_sm':attr(81)}})
print(json.dumps(row),flush=True);del graph
write_json(a.output/'resources.json',{'complete':True,'rows':rows,'device':{'max_threads_per_sm':attr(39),'registers_per_sm':attr(82),'shared_bytes_per_sm':attr(81)},'hardware_counters':'Unavailable: ERR_NVGPUCTRPERM, driver restricts access to administrators.'})
Test 20. Latency hiding: pipelines, prefetch, unroll and register caps (Oct 6)
Visual reference: Diagram 6 · Topic overview.
The idea: We try to request upcoming weights earlier and keep more independent work in flight. It varies loop unrolling, pipeline stages, explicit prefetching, tile sizes and compiler register limits.
The result: Selected shape speedups are 1.0747× for QKV, 1.0474× for GDN, 1.0266× for the head, 1.0447× for down and 1.0600× for output. Gate/up keeps the reference. Linked replay gives 1.0261–1.0301×. All 771 projection, 192 linked and 12 graph checks pass. Extra staging and explicit prefetch generally lose.
Explanation, checks and code
Why it matters on the GPU: A load from board memory takes time. Starting it before the current arithmetic finishes can hide some of that delay. However, prefetched values occupy registers or shared memory until used. Register caps can allow more warps to fit, but may force temporary data into slower spill storage. Unrolling exposes independent instructions while increasing the live working set.
What to take from it: Some limited unrolling and geometry changes help. More staging and explicit prefetch usually lose, showing that additional overlap can cost more resources than it saves. The modest linked gain is the measured result; no single stall percentage explains it by itself.
Setup: We screen 126 candidate/shape pairs and six unchanged controls, rotating through real weights. Each comparison uses seven 40 ms pairs. Selection checks spills and improvement against the control before linked confirmation.
Relevant code:
Pipelining and loop-unrolling knobs change when packed tiles are consumed.
Source: pq2-dev/benchmarks/harness/pq2_latency_kernels.py:7–15.
@triton.jit
def _pipeline(Q,S,W,WS,Y,N:tl.constexpr,K:tl.constexpr,BN:tl.constexpr,BC:tl.constexpr,
STAGES:tl.constexpr,UNROLL:tl.constexpr):
rows=tl.program_id(0)*BN+tl.arange(0,BN);off=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=STAGES,loop_unroll_factor=UNROLL):
blocks=base*BC+off;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)
Explicit next-tile prefetch variant.
Source: pq2-dev/benchmarks/harness/pq2_latency_kernels.py:17–26.
@triton.jit
def _prefetch(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);off=tl.arange(0,BC)
acc=tl.full((BN,BC),0,tl.float32)
packed,ws=_load_tile(W,WS,rows,off,N,K)
for base in tl.range(0,tl.cdiv(K//128,BC),num_stages=1):
nxt,nws=_load_tile(W,WS,rows,(base+1)*BC+off,N,K)
acc+=_consume(Q,S,packed,ws,base*BC+off,K)
packed,ws=nxt,nws
tl.store(Y+rows,tl.sum(acc,1),rows<N)
Geometry, warp count and register caps are passed into each candidate launch.
Source: pq2-dev/benchmarks/harness/pq2_latency_kernels.py:31–37.
def launch(q,s,weight,y,n,k,cfg):
bn,bc=geometry(n,k);bn=cfg.get('bn',bn);bc=cfg.get('bc',bc)
flat=weight.reshape(-1);args=(q.view(torch.int32),s,flat[:n*k//4],flat[n*k//4:].view(torch.float16),y,n,k,bn,bc)
options={'num_warps':cfg.get('warps',4)}
if cfg.get('maxnreg'):options['maxnreg']=cfg['maxnreg']
if cfg.get('mode')=='prefetch':return _prefetch[(triton.cdiv(n,bn),)](*args,**options)
return _pipeline[(triton.cdiv(n,bn),)](*args,cfg.get('stages',1),cfg.get('unroll',1),**options)
G. Preparing operands for Tensor Cores
Takeaway. Small batches need operands arranged for the matrix instructions. Direct register fragments can avoid costly intermediate staging.
Terms for this section: Groups, planes and fragments · Tiles.
Diagram 7: Four code planes matched to activation slices
The minus-one correction accounts for PQ2 values being code minus one. Test 21 uses masked-plane Triton MMA; test 22 checks different token counts. This diagram shows the shared mathematical idea; exact shifts and operand layouts differ between implementations.
Diagram 8: Shared-memory staging versus direct register fragments
A fragment is the part of a hardware matrix tile held by one thread. Test 23 aims to remove shared-memory staging and barriers from the multiply path. More direct control still leaves register limits and work distribution to solve.
Test 21. Masked-plane Triton small-batch MMA (Oct 7)
Visual reference: Diagram 7, Diagram 8 · Topic overview.
The idea: We replace eight-token DP4A work with small integer matrix multiplies. It rearranges the rounded inputs into bit-plane order so that each plane of the compressed weights can multiply a contiguous slice of inputs, then corrects for the weight code's minus-one offset.
The result: The whole projection pass falls from 35.53 to 18.53 ms, a 1.93× speedup. Selected shapes gain 2.00–2.43×. All 579 output checks pass, with maximum relative RMS error of 1.1e-4. The output projection keeps its floating path.
Explanation, checks and code
Why it matters on the GPU: A packed byte holds four two-bit codes. Plane ordering makes those codes usable by Tensor Core dot products with less rearrangement. The kernel uses masks and shifts, applies one activation-sum correction per group, and uses shared-memory staging as generated by Triton. Split-K gives the smaller down-projection more independent blocks, with a second kernel combining partial sums in a fixed order.
What to take from it: The near-twofold whole-projection improvement shows that matching the arithmetic and operand layout to Tensor Cores helps this eight-row workload. The floating output projection remains outside the accelerated set.
Setup: We use real eight-row captures for 257 projections at three decode steps. The search varies BN = 32, 64 or 128; 1 or 2 groups; 4 or 8 warps; and 2–4 stages. Down also gets deterministic split-K settings.
Relevant code:
Current plane-order quantisation and per-group activation sum, matching the algorithm described in the recorded test.
Source: vllm-gguf-plugin-prism/vllm_gguf_plugin/triton/pq2_small_mmq.py:214–225.
@triton.jit
def _quantize_planes_sum(X, Q, S, T, K: tl.constexpr):
# As _quantize_planes, plus the integer sum of each group's codes (for the -1 offset).
block = tl.program_id(0)
i = tl.arange(0, 128)
a = tl.load(X + block * 128 + i).to(tl.float32)
scale = tl.max(tl.abs(a), 0) / 127.0
normalized = tl.div_rn(a, tl.where(scale > 0, scale, 1.0))
q = tl.minimum(127.0, tl.maximum(-127.0, tl.extra.cuda.libdevice.nearbyint(normalized)))
tl.store(Q + block * 128 + (i % 4) * 32 + i // 4, q.to(tl.int8))
tl.store(S + block, scale)
tl.store(T + block, tl.sum(q.to(tl.int32), 0))
Current masked-plane tensor-core kernel, including exact group correction.
Source: vllm-gguf-plugin-prism/vllm_gguf_plugin/triton/pq2_small_mmq.py:246–281.
@triton.jit(do_not_specialize=["M"])
def _small_mmq_masked(Q, S, T, W, WS, Y, M, N: tl.constexpr, K: tl.constexpr,
BN: tl.constexpr, G: tl.constexpr, STAGES: tl.constexpr,
SPLIT: tl.constexpr = 1, BT: tl.constexpr = 16):
"""Codes stay in place: plane s is `byte & (3 << 2s)` = code * 4^s (plane 3 pre-shifted
by 2 to stay below 128). The exact INT32 product is shifted back, and the ternary -1
offset is applied once per group as -sum(q). Integer result equals the DP4A kernel's."""
GROUPS: tl.constexpr = K // 128
rows = tl.program_id(0) * BN + tl.arange(0, BN)
tokens = tl.arange(0, BT)
byte = tl.arange(0, 32)
row_ok = rows < N
tok_ok = tokens < M
acc = tl.zeros((BN, BT), tl.float32)
PER: tl.constexpr = GROUPS // SPLIT
first = tl.program_id(1) * PER
for base in tl.range(first, first + PER, G, num_stages=STAGES):
for j in tl.static_range(G):
g = base + j
packed = tl.load(W + rows[:, None] * (K // 4) + g * 32 + byte[None, :],
row_ok[:, None], other=0)
q0 = tl.load(Q + tokens[None, :] * K + g * 128 + byte[:, None], tok_ok[None, :], other=0)
q1 = tl.load(Q + tokens[None, :] * K + g * 128 + 32 + byte[:, None], tok_ok[None, :], other=0)
q2 = tl.load(Q + tokens[None, :] * K + g * 128 + 64 + byte[:, None], tok_ok[None, :], other=0)
q3 = tl.load(Q + tokens[None, :] * K + g * 128 + 96 + byte[:, None], tok_ok[None, :], other=0)
part = tl.dot((packed & 0x03).to(tl.int8), q0)
part += tl.dot((packed & 0x0C).to(tl.int8), q1) >> 2
part += tl.dot((packed & 0x30).to(tl.int8), q2) >> 4
part += tl.dot(((packed >> 2) & 0x30).to(tl.int8), q3) >> 4
part -= tl.load(T + tokens * GROUPS + g, tok_ok, other=0)[None, :]
ws = tl.load(WS + rows * GROUPS + g, row_ok, other=0).to(tl.float32)
xs = tl.load(S + tokens * GROUPS + g, tok_ok, other=0)
acc += part.to(tl.float32) * ws[:, None] * xs[None, :]
# SPLIT == 1 stores the output directly; otherwise Y is an FP32 (SPLIT, M, N) workspace.
out = Y + tl.program_id(1) * M * N
tl.store(out + tokens[None, :] * N + rows[:, None], acc, row_ok[:, None] & tok_ok[None, :])
Standalone alternating confirmation timings.
Source: pq2-dev/benchmarks/experiments/batch8/pq2_small_mmq_screen.py:45–53.
def paired(control,candidate,pairs,ms):
for _ in range(2):time_graph(control,ms);time_graph(candidate,ms)
rows,ratios=[],[]
for r in range(pairs):
if r%2:cv,bv=time_graph(candidate,ms),time_graph(control,ms)
else:bv,cv=time_graph(control,ms),time_graph(candidate,ms)
rows.append([bv,cv]);ratios.append(cv/bv)
return {'reference_ms':statistics.median(p[0] for p in rows),'candidate_ms':statistics.median(p[1] for p in rows),
'speedup':1/statistics.median(ratios),'min_speedup':1/max(ratios),'max_speedup':1/min(ratios),'pairs_ms':rows}
Test 22. Small-M MMA extension and floating output tile (Oct 7)
Visual reference: Diagram 7, Diagram 8 · Topic overview.
The idea: We repeat the small Tensor Core projection search for different numbers of token rows, including the sizes used during speculative verification. It also tests floating output tiles sized closer to the actual row count.
The result: Whole-projection speedups are 0.99× at M = 2, 3.85× at 3, 1.25× at 4, 1.82× at 5, 4.19× at 6, 4.18× at 7, and 3.95× at 16. The production baseline changes with M. Maximum relative RMS is 1.4e-4 against Q8, and 0.0036–0.0069 against floating output. These timings predate the output-tile improvement.
Explanation, checks and code
Why it matters on the GPU: Tensor Core tiles have fixed hardware dimensions, so two, four or sixteen rows can use the same tile very differently. At small M, padding and launch costs may outweigh reuse. The old dispatch also selected different kernel families for different M, making some baselines much weaker than others. Smaller floating tiles avoid processing or allocating execution resources for many padded rows.
What to take from it: The large speedups at some M reflect both a better kernel and a weak previous dispatch for that M. The numerical reference is quantised integer arithmetic; matching it doesn’t mean exact agreement with the old floating path.
Setup: We try real-row inputs at M = 2–7 and 16. Smaller M takes rows from a saved capture; M = 16 joins two steps. Winners get seven confirmation pairs and 1,035 output checks.
Relevant code:
Build the tested M-row inputs from saved eight-row captures.
Source: pq2-dev/benchmarks/experiments/batch8/pq2_small_mmq_mscreen.py:27–30.
def rows_for(replay,name,m,step):
if m<=8:return replay.samples[step,name]['x'][:m].contiguous()
other={1:32,32:128,128:1}[step]
return torch.cat([replay.samples[step,name]['x'],replay.samples[other,name]['x']])[:m].contiguous()
The numerical reference is the integer path at the same M.
Source: pq2-dev/benchmarks/experiments/batch8/pq2_small_mmq_mscreen.py:37–40.
def integer_reference(replay,name,x):
layer=replay.layers[name];y=pq2_int_gemv(rotate(layer,x),layer['weight'],prepared=True,chained=False)
if layer['bias'] is not None:y.add_(layer['bias'])
return y
Geometry screen, reference-error gate, seven-pair confirmation and all-step checks.
Source: pq2-dev/benchmarks/experiments/batch8/pq2_small_mmq_mscreen.py:63–89.
control,_c=graph_for(lambda:run(base))
splits=(1,2,4) if n<=8192 else (1,)
for bn,g,stages,split in itertools.product((32,64,128),(1,2),(2,3,4),splits):
if (k//128)%split or (k//128//split)%g:continue
geometry=[bn,g,4,stages,split];row={'m':m,'shape':[n,k],'geometry':geometry}
write_json(out/'status.json',{'complete':False,'stage':'screen','m':m,'shape':[n,k],'geometry':geometry})
try:
config={f'{n},{k}':geometry};x=inputs[names[0]]
cand=replay.projection(names[0],x,config)
metrics=error_metrics(cand,integer_reference(replay,names[0],x))
if not metrics['finite'] or metrics['relative_rms']>REL_RMS_LIMIT:raise ValueError(f'numerics {metrics}')
cg,_o=graph_for(lambda:run(config))
row|=paired(control,cg,3,60);row['validation_vs_integer']=metrics
row['resources']=replay.resources.get(f'{n},{k},{bn},{g},4,{stages},{split}');del cg,_o
except Exception as error:row['error']=str(error)[:300]
records.setdefault(key,[]).append(row);write_json(out/'progress.json',records)
ok=[r for r in records[key] if 'speedup' in r]
if not ok:raise RuntimeError(f'No valid candidate for {key}')
top=max(ok,key=lambda r:r['speedup']);config={f'{n},{k}':top['geometry']}
cg,_o=graph_for(lambda:run(config));confirm=paired(control,cg,7,150);del cg,_o
vs_int,vs_prod=[],[]
for step in replay.manifest['steps']:
for name in names:
x=rows_for(replay,name,m,step);cand=replay.projection(name,x,config)
mi=error_metrics(cand,integer_reference(replay,name,x))
if not mi['finite'] or mi['relative_rms']>REL_RMS_LIMIT:raise RuntimeError(f'{key} {name} step {step} failed: {mi}')
vs_int.append(mi['relative_rms']);vs_prod.append(error_metrics(cand,replay.projection(name,x,base))['relative_rms'])
Test 23. Direct-fragment CUDA tensor-core MMA (Oct 7)
Visual reference: Diagram 7, Diagram 8 · Topic overview.
The idea: We supply the Tensor Core instruction's operands directly in thread registers using CUDA. It tries to avoid the shared-memory staging that limited the earlier Triton small-batch kernel, and loads several groups ahead before consuming them.
The result: 33 of 35 M/shape entries qualify against the current Triton/DP4A selection. Whole-pass gains are 1.23× at M = 2, 1.19× at 3, 1.15× at 4 and 5, 1.11× at 6, 1.07× at 7, and 1.05× at 8. Maximum relative RMS is 1.7e-4. Down keeps Triton at M = 7 and 8.
Explanation, checks and code
Why it matters on the GPU: Tensor Core matrix instructions expect specific pieces of a tile in specific threads and registers; these pieces are called fragments. Arranging loads and input quantisation to match those fragments removes a shared-memory round trip and block barriers from the core multiply. Holding more groups in flight can hide memory latency, but uses more registers. At larger M the increased register footprint makes this approach less attractive.
What to take from it: The direct-fragment design improves most small-M entries. The down projection at M = 7 or 8 keeps the Triton split-K route, illustrating that enough parallel blocks can matter more than eliminating shared-memory staging.
Setup: The CUDA kernel uses direct register fragments for INT8 Tensor Core MMA, group-128 Q8 and an activation-sum correction. Synthetic and real-row screens try 2 or 4 warps and 2, 4 or 8 groups in flight.
Relevant code:
Current direct-fragment CUDA MMA instruction wrapper.
Source: vllm-gguf-plugin-prism/vllm_gguf_plugin/csrc/pq2/pq2_mma_small.cu:22–28.
__device__ __forceinline__ void mma_s8(int (&c)[4], const uint32_t (&a)[4], uint32_t b0, uint32_t b1) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
: "+r"(c[0]), "+r"(c[1]), "+r"(c[2]), "+r"(c[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
}
Current register-fragment plane extraction and scaled accumulation; the current source may include later RT-tiling refinements.
Source: vllm-gguf-plugin-prism/vllm_gguf_plugin/csrc/pq2/pq2_mma_small.cu:85–130.
#pragma unroll
for (int u = 0; u < UNROLL; ++u) {
const int g = g0 + u < g_end ? g0 + u : g_end - 1;
#pragma unroll
for (int r = 0; r < RT; ++r) {
w0[u][r] = ok0[r] ? __ldg(reinterpret_cast<const uint2*>(c0[r] + g * 32)) : make_uint2(0x55555555u, 0x55555555u);
w1[u][r] = ok1[r] ? __ldg(reinterpret_cast<const uint2*>(c1[r] + g * 32)) : make_uint2(0x55555555u, 0x55555555u);
ws0[u][r] = ok0[r] ? __half2float(__ldg(wscale + (size_t)row0[r] * groups + g)) : 0.f;
ws1[u][r] = ok1[r] ? __half2float(__ldg(wscale + (size_t)row1[r] * groups + g)) : 0.f;
}
#pragma unroll
for (int t = 0; t < NT; ++t) {
#pragma unroll
for (int s = 0; s < 4; ++s) {
b[u][t][2 * s] = tok_ok[t] ? __ldg(reinterpret_cast<const uint32_t*>(qb[t] + g * 128 + s * 32)) : 0u;
b[u][t][2 * s + 1] = tok_ok[t] ? __ldg(reinterpret_cast<const uint32_t*>(qb[t] + g * 128 + s * 32 + 16)) : 0u;
}
#pragma unroll
for (int j = 0; j < 2; ++j) {
const int tk = c_ok[t][j] ? t * 8 + tid * 2 + j : 0;
x[u][t][j] = c_ok[t][j] ? __ldg(xs + (size_t)tk * groups + g) : 0.f;
qs[u][t][j] = c_ok[t][j] ? __ldg(qsum + (size_t)tk * groups + g) : 0;
}
}
}
#pragma unroll
for (int u = 0; u < UNROLL; ++u) {
if (g0 + u >= g_end) break;
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int r = 0; r < RT; ++r) {
int c[4] = {0, 0, 0, 0};
#pragma unroll
for (int s = 0; s < 4; ++s) {
const uint32_t a[4] = {(w0[u][r].x >> (2 * s)) & 0x03030303u, (w1[u][r].x >> (2 * s)) & 0x03030303u,
(w0[u][r].y >> (2 * s)) & 0x03030303u, (w1[u][r].y >> (2 * s)) & 0x03030303u};
mma_s8(c, a, b[u][t][2 * s], b[u][t][2 * s + 1]);
}
#pragma unroll
for (int j = 0; j < 2; ++j) {
acc[r][t][j] += (float)(c[j] - qs[u][t][j]) * ws0[u][r] * x[u][t][j];
acc[r][t][2 + j] += (float)(c[2 + j] - qs[u][t][j]) * ws1[u][r] * x[u][t][j];
}
}
}
The standalone PyTorch/ctypes launcher quantises fragments and calls the CUDA shared library.
Source: pq2-dev/benchmarks/experiments/batch8/pq2_mma_small.py:45–62.
def mma_small(x,weight,*,split=1,warps=4,unroll=2):
m,k=x.shape;n=weight.shape[0];groups=k//128
if not 1<=m<=16 or k%128 or weight.shape[1]!=groups*34 or groups%split:raise ValueError('bad shape')
if not weight.is_contiguous():raise ValueError('prepared contiguous weight required')
x=x.contiguous()
q=torch.empty((m,k),dtype=torch.int8,device=x.device)
s=torch.empty((m,groups),dtype=torch.float32,device=x.device)
t=torch.empty((m,groups),dtype=torch.int32,device=x.device)
_quantize_fragments[(m*groups,)](x,q,s,t,k)
y=torch.empty((m,n),dtype=x.dtype,device=x.device)
partial=torch.empty((split,m,n),dtype=torch.float32,device=x.device) if split>1 else None
codes=weight.data_ptr();scales=codes+n*k//4
err=_library().pq2_mma_small_launch(q.data_ptr(),s.data_ptr(),t.data_ptr(),codes,scales,y.data_ptr(),
partial.data_ptr() if partial is not None else None,m,n,k,int(x.dtype==torch.bfloat16),split,warps,unroll,
torch.cuda.current_stream(x.device).cuda_stream)
if err:raise RuntimeError(f'pq2_mma_small launch failed: {err}')
if split>1:y.copy_(partial.sum(0))
return y
H. Memory access patterns and tile-major storage
Takeaway. Putting nearby reads together improved the load-only probe. A complete projection also needs enough parallel work and efficient arithmetic.
Terms for this section: Row-major / tile-major · Stride and coalescing.
Diagram 9: Put the bytes consumed together next to each other
The row-major case illustrates the tested current pattern; row-major storage can also be accessed efficiently with larger contiguous chunks. A 512-byte contiguous warp region still uses multiple hardware transactions. Test 24 isolates loads; test 25 checks whether the advantage survives real compute.
Test 24. Bandwidth access-pattern probes (Oct 15)
Visual reference: Diagram 9 · Topic overview.
The idea: We remove the projection arithmetic and asks how quickly each load pattern can read compressed code bytes. It compares a linear stream, the current row pattern, larger contiguous row chunks, an asynchronous staging ring and a tile-major arrangement.
The result: For K = 5,120, the current pattern reaches 417–447 GB/s on gate/up and 320–354 GB/s on the head. Tile-major reaches 951 and 971 GB/s respectively. Row-major chunks of at least 128 bytes and cp.async also approach the streaming limit. With arithmetic removed, this points to a load-pattern limit.
Explanation, checks and code
Why it matters on the GPU: GPU memory serves neighbouring addresses efficiently when the threads in a warp request them together. Small pieces from widely separated rows can use the load hardware less efficiently even when the total byte count is unchanged.
Tile-major storage puts the pieces consumed by a warp next to each other; larger row chunks achieve a similar goal in a different way. The cp.async variant stages data into shared memory while other work proceeds.
What to take from it: Rotating buffers larger than L2 reduces cache-only results. With computation removed, the large gap on K=5120 shapes is strong evidence that the current access pattern itself leaves bandwidth unused. These are code-byte throughput figures, not full-projection performance.
Setup: We run copy-only CUDA probes through ctypes and PyTorch on an RTX 3090 Ti. Rotating at least 160 MB of weights reduces L2 reuse. Patterns include linear reads, row-major chunks, cp.async and tile-major reads.
Relevant code:
Current-pattern and contiguous-row load probes isolate memory-access behaviour.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bw_probe.cu:23–73.
template <int U>
__global__ void k_cur(const uint8_t* __restrict__ codes, int K, uint32_t* out) {
const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, gid = lane >> 2, tid = lane & 3;
const int tile = blockIdx.x * (blockDim.x >> 5) + warp;
const uint8_t* c0 = codes + (size_t)(tile * 16 + gid) * (K >> 2) + tid * 8;
const uint8_t* c1 = c0 + (size_t)8 * (K >> 2);
const int groups = K >> 7;
uint32_t a = 0;
for (int g0 = 0; g0 < groups; g0 += U) {
uint2 w0[U], w1[U];
#pragma unroll
for (int u = 0; u < U; ++u) {
const int g = g0 + u < groups ? g0 + u : groups - 1;
w0[u] = __ldg(reinterpret_cast<const uint2*>(c0 + g * 32));
w1[u] = __ldg(reinterpret_cast<const uint2*>(c1 + g * 32));
}
#pragma unroll
for (int u = 0; u < U; ++u) a ^= w0[u].x ^ w0[u].y ^ w1[u].x ^ w1[u].y;
}
out[blockIdx.x * blockDim.x + threadIdx.x] = a;
}
// 2. segment-length study: each warp owns 16 rows, reads ROWB contiguous bytes per row per iteration with LDG.128,
// U iterations in flight.
template <int ROWB, int U>
__global__ void k_seg(const uint8_t* __restrict__ codes, int K, uint32_t* out) {
constexpr int LPR = ROWB / 16; // lanes per row
constexpr int RPI = 32 / LPR; // rows per warp instruction
constexpr int INSTR = 16 / RPI; // instructions per 16-row tile per iteration
const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
const int tile = blockIdx.x * (blockDim.x >> 5) + warp;
const int rl = lane / LPR, part = lane % LPR;
const int nseg = (K >> 2) / ROWB;
const uint8_t* base = codes + (size_t)(tile * 16 + rl) * (K >> 2) + part * 16;
uint32_t a = 0;
for (int s0 = 0; s0 < nseg; s0 += U) {
uint4 v[U][INSTR];
#pragma unroll
for (int u = 0; u < U; ++u) {
const int s = s0 + u < nseg ? s0 + u : nseg - 1;
#pragma unroll
for (int i = 0; i < INSTR; ++i)
v[u][i] = __ldg(reinterpret_cast<const uint4*>(base + (size_t)i * RPI * (K >> 2) + (size_t)s * ROWB));
}
#pragma unroll
for (int u = 0; u < U; ++u)
#pragma unroll
for (int i = 0; i < INSTR; ++i) a ^= v[u][i].x ^ v[u][i].y ^ v[u][i].z ^ v[u][i].w;
}
out[blockIdx.x * blockDim.x + threadIdx.x] = a;
}
Tile-major probe uses contiguous lane loads.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bw_probe.cu:108–130.
// WS=1 also reads one 4-byte scale word (rows gid, gid+8) per lane per group from a [tile][group][8 lanes] stream.
template <int U, int WS>
__global__ void k_tile(const uint8_t* __restrict__ codes, const uint8_t* __restrict__ scales, int K, uint32_t* out) {
const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, gid = lane >> 2;
const int tile = blockIdx.x * (blockDim.x >> 5) + warp;
const int groups = K >> 7;
const uint4* c = reinterpret_cast<const uint4*>(codes) + (size_t)tile * groups * 32 + lane;
const uint32_t* sc = reinterpret_cast<const uint32_t*>(scales) + (size_t)tile * groups * 8 + gid;
uint32_t a = 0;
for (int g0 = 0; g0 < groups; g0 += U) {
uint4 v[U];
uint32_t w[U];
#pragma unroll
for (int u = 0; u < U; ++u) {
const int g = g0 + u < groups ? g0 + u : groups - 1;
v[u] = __ldg(c + (size_t)g * 32);
if (WS) w[u] = __ldg(sc + (size_t)g * 8);
}
#pragma unroll
for (int u = 0; u < U; ++u) { a ^= v[u].x ^ v[u].y ^ v[u].z ^ v[u].w; if (WS) a ^= w[u]; }
}
out[blockIdx.x * blockDim.x + threadIdx.x] = a;
}
The PyTorch driver rotates buffers inside CUDA graphs and reports effective code-byte bandwidth.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/run_probe.py:10–29.
def bench(launch, copies, bytes_each, reps=20):
st=torch.cuda.Stream()
with torch.cuda.stream(st):
for c in copies:
r=launch(c)
if r not in (0,None): raise RuntimeError(f"launch rc={r}")
torch.cuda.synchronize()
err=torch.cuda.current_stream() # surfaces async errors
torch.cuda.synchronize()
g=torch.cuda.CUDAGraph()
with torch.cuda.graph(g,stream=st):
for c in copies: launch(c)
for _ in range(3): g.replay()
torch.cuda.synchronize()
e0=torch.cuda.Event(enable_timing=True); e1=torch.cuda.Event(enable_timing=True)
e0.record()
for _ in range(reps): g.replay()
e1.record(); torch.cuda.synchronize()
ms=e0.elapsed_time(e1)/reps
return bytes_each*len(copies)/(ms*1e-3)/1e9, ms*1e3/len(copies)
Test 25. Tile-major MMA and M=1 headroom (Oct 15)
Later follow-up: this is the initial kernel without the successful down split-K setting. Tests 30–31 address that limitation; tests 32–39 compare layouts and the newer integrated kernels.
Visual reference: Diagram 9 · Topic overview.
The idea: We take the promising tile-major layout from the read-only probes and uses it in a real integer Tensor Core projection. It compares answers and timings against the current projection, then separately measures how close the existing one-token path is to the load-only ceiling.
The result: Gate/up falls from 71.2 to 54.5 µs, a 23% latency reduction; input falls from 34.8 to 30.9, QKV from 29.4 to 29.1, and the head from 385 to 354 µs. These outputs match bit for bit. Down rises from 50.4 to 65.3 µs and output from 21.1 to 25.0 µs. M = 1 bandwidth is 625–924 GB/s. The estimated 1.3 ms step saving is an extrapolation, not measured model throughput.
Explanation, checks and code
Why it matters on the GPU: Tile-major weights place each warp's 16-byte-per-thread loads next to each other, feeding register fragments efficiently. The kernel keeps several row tiles or groups in flight to reuse inputs and overlap loads. Small-output matrices create fewer row tiles, however, so some multiprocessors can sit idle unless the reduction is also split across blocks. A good memory layout can’t compensate for too little parallel work.
What to take from it: The layout helps gate/up and the head but loses on down and output, which lack the needed parallelism in this variant. The estimated whole-step saving is an extrapolation from component timings, not an engine measurement.
Setup: We compare eight-row tile-major CUDA MMA against production, including its existing row-tile (RT) settings, and check output bits. A separate probe measures the production one-row path. Both rotate at least 160 MB of weights.
Relevant code:
The dated standalone test explicitly converts prepared weights into tile-major layout.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bench_tm.py:7–12.
def tile_major(w,n,k,rt):
G=k//128; ST=n//(16*rt)
flat=w.reshape(-1)
c=flat[:n*k//4].view(ST,rt,2,8,G,4,8).permute(0,4,1,3,5,2,6).contiguous()
s=flat[n*k//4:].view(torch.float16).view(ST,rt,2,8,G).permute(0,4,1,3,2).contiguous()
return c,s
The dated tile-major CUDA kernel loads tiles and scales ahead, executes masked-plane MMA and accumulates scaled results.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/mma_tm.cu:63–117.
const uint4* cbase = codes + (size_t)st * groups * RT * 32 + lane;
const uint32_t* sbase = wscale + (size_t)st * groups * RT * 8 + gid;
for (int g0 = 0; g0 < groups; g0 += UNROLL) {
uint4 wv[UNROLL][RT];
uint32_t wsv[UNROLL][RT];
uint32_t b[UNROLL][NT][8];
float x[UNROLL][NT][2];
int qs[UNROLL][NT][2];
#pragma unroll
for (int u = 0; u < UNROLL; ++u) {
const int g = g0 + u < groups ? g0 + u : groups - 1;
#pragma unroll
for (int r = 0; r < RT; ++r) {
wv[u][r] = __ldg(cbase + ((size_t)g * RT + r) * 32);
wsv[u][r] = __ldg(sbase + ((size_t)g * RT + r) * 8);
}
#pragma unroll
for (int t = 0; t < NT; ++t) {
#pragma unroll
for (int s = 0; s < 4; ++s) {
b[u][t][2 * s] = tok_ok[t] ? __ldg(reinterpret_cast<const uint32_t*>(qb[t] + g * 128 + s * 32)) : 0u;
b[u][t][2 * s + 1] = tok_ok[t] ? __ldg(reinterpret_cast<const uint32_t*>(qb[t] + g * 128 + s * 32 + 16)) : 0u;
}
#pragma unroll
for (int j = 0; j < 2; ++j) {
const int tk = c_ok[t][j] ? t * 8 + tid * 2 + j : 0;
x[u][t][j] = c_ok[t][j] ? __ldg(xs + (size_t)tk * groups + g) : 0.f;
qs[u][t][j] = c_ok[t][j] ? __ldg(qsum + (size_t)tk * groups + g) : 0;
}
}
}
#pragma unroll
for (int u = 0; u < UNROLL; ++u) {
if (g0 + u >= groups) break;
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int r = 0; r < RT; ++r) {
int c[4] = {0, 0, 0, 0};
#pragma unroll
for (int s = 0; s < 4; ++s) {
const uint32_t a[4] = {(wv[u][r].x >> (2 * s)) & 0x03030303u, (wv[u][r].z >> (2 * s)) & 0x03030303u,
(wv[u][r].y >> (2 * s)) & 0x03030303u, (wv[u][r].w >> (2 * s)) & 0x03030303u};
mma_s8(c, a, b[u][t][2 * s], b[u][t][2 * s + 1]);
}
const float ws0 = __half2float(__ushort_as_half((unsigned short)(wsv[u][r] & 0xffffu)));
const float ws1 = __half2float(__ushort_as_half((unsigned short)(wsv[u][r] >> 16)));
#pragma unroll
for (int j = 0; j < 2; ++j) {
acc[r][t][j] += (float)(c[j] - qs[u][t][j]) * ws0 * x[u][t][j];
acc[r][t][2 + j] += (float)(c[2 + j] - qs[u][t][j]) * ws1 * x[u][t][j];
}
}
}
}
The standalone M=1 probe measures the existing production projection path without a model engine.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/m1.py:1–18.
import torch,sys
sys.path.insert(0,"/tmp/claude-1000/-home-amorley-vllm/9d014f21-8765-4ef9-9067-3fa7be39352d/scratchpad/tm")
exec(open("/tmp/claude-1000/-home-amorley-vllm/9d014f21-8765-4ef9-9067-3fa7be39352d/scratchpad/tm/bench_tm.py").read().split("Ms=[int")[0])
from vllm_gguf_plugin.triton import prism
prism._EXPERIMENTAL_INT_GEMV=True
import os
print("M=1 production decode path (integer GEMV / single), rotating weight copies:")
for n,k,name in [(34816,5120,"gate_up"),(16384,5120,"in_proj"),(14336,5120,"qkv"),(5120,17408,"down"),(5120,6144,"out"),(248320,5120,"lm_head")]:
nbytes=n*(k//128)*34; nc=max(1,int(160e6//nbytes))
ws=[]
for i in range(nc):
raw=torch.randint(0,256,(n,k//128,34),device="cuda",dtype=torch.uint8)
raw[:,:,:2]=(torch.rand(n,k//128,device="cuda")*0.04+0.005).half().view(torch.uint8).view(n,k//128,2)
ws.append(prepare_pq2_layout(raw.view(n,-1).contiguous()))
x=torch.randn(1,k,device="cuda",dtype=torch.bfloat16)
t=timed([(lambda w=w: prism.pq2_matmul(x,w,prepared=True)) for w in ws])
print(f"{name:8s} {n}x{k}: {t:6.1f} us {nbytes/1e9/t*1e6:4.0f} GB/s")
del ws
I. Attention: parallel histories and compressed KV
Takeaway. Attention needs a kernel and split count suited to the query length and cache format. More splits or fewer cache bytes do not guarantee a faster request.
Terms for this section: Warps and reductions · Stride and coalescing.
Diagram 10: Split one query over its history, then merge softmax results
Splits create more parallel blocks but add partial storage and merge work. Partial attention outputs can’t simply be averaged. Tests 26 and 27 vary FA2 splitting; test 28 compares backends and segment counts on real pages; test 29 adds KV compression and arithmetic choices.
Only two history parts are drawn to keep the diagram compact; the tests use more. KV can use BF16 storage or smaller INT8 values with scales; this storage choice is separate from the split count.
Test 26. Synthetic FA2 attention split sweep (Oct 6)
Visual reference: Diagram 10 · Topic overview.
The idea: We create one attention query and a synthetic history of keys and values. It divides the history into different numbers of pieces, computes attention on those pieces in parallel, and combines the partial answers.
The result: The best 16-split setting gives 1.272× at 4,096 tokens, 1.234× at 6,144 and 1.203× at 8,192. At 128 tokens, the best result is 1.001×—effectively unchanged. The test calls the installed FA2 operation directly, without an engine.
Explanation, checks and code
Why it matters on the GPU: One query can expose too little parallel work to fill the GPU. Splitting its long history creates more thread blocks so more multiprocessors can work simultaneously. Each split needs partial softmax statistics and output storage, followed by a combine step. Too many splits add memory traffic and launch/reduction overhead; short histories provide too little work to pay for that overhead.
What to take from it: Sixteen splits helps the tested long histories, while the short case is effectively flat. This is primarily a work-distribution experiment within the same FA2 kernel family, not a change to how the model weights are stored.
Setup: We create synthetic BF16 queries, keys and values with 24 query heads, 4 KV heads, head dimension 256 and 16-token pages. Contexts are 128, 4,096, 6,144 and 8,192 tokens. We try automatic splitting and 1, 2, 4, 8, 16 or 32 splits, with FP32 and changing-query checks.
Relevant code:
Synthetic paged Q/K/V, direct FA2 split override, independent FP32 attention, and changing-input graph validation.
Source: pq2-dev/benchmarks/experiments/attention/pq2_single_attention.py:39–75.
cuq=torch.tensor([0,1],device='cuda',dtype=torch.int32)
cuk=torch.zeros_like(cuq)
for length in (128,4096,6144,8192):
write_json(a.output/'status.json',{'complete':False,'stage':'attention split sweep','context':length})
q=torch.randn(1,24,256,device='cuda',dtype=torch.bfloat16)
original_q=q.clone()
k=torch.randn(length//16,16,4,256,device='cuda',dtype=torch.bfloat16)
v=torch.randn_like(k)
table=torch.arange(length//16,device='cuda',dtype=torch.int32).unsqueeze(0)
used=torch.tensor([length],device='cuda',dtype=torch.int32)
def run(splits):
return torch.ops._vllm_fa2_C.varlen_fwd(q,k,v,None,cuq,cuk,used,None,table,None,
1,length,0.,1/math.sqrt(256),False,True,-1,-1,0.,False,splits,None)[0]
# Independent float32 reference, repeating each KV head six times.
keys=k.reshape(length,4,256).permute(1,0,2).repeat_interleave(6,0).float()
vals=v.reshape(length,4,256).permute(1,0,2).repeat_interleave(6,0).float()
scores=torch.einsum('hd,hld->hl',q[0].float(),keys)/math.sqrt(256)
truth=torch.einsum('hl,hld->hd',scores.softmax(-1),vals).unsqueeze(0).to(q.dtype)
control,baseline=graph_for(lambda:run(0))
control.replay();torch.cuda.synchronize()
metrics=error_metrics(baseline,truth)
if not metrics['finite'] or metrics['relative_rms']>0.01:raise RuntimeError(f'FA2 reference mismatch: {metrics}')
checks.append({'context':length,'splits':0,'float32_reference':metrics})
for splits in (1,2,4,8,16,32):
try:
q.copy_(original_q)
actual=run(splits);metrics=error_metrics(actual,truth)
if not metrics['finite'] or metrics['relative_rms']>0.01:raise RuntimeError(f'Numerical mismatch: {metrics}')
candidate,y=graph_for(lambda:run(splits))
candidate.replay();torch.cuda.synchronize();old=y.clone()
q.mul_(0.8).add_(0.04);expected=run(splits)
candidate.replay();torch.cuda.synchronize()
torch.testing.assert_close(y,expected,atol=0,rtol=0)
if torch.equal(y,old):raise RuntimeError('Graph did not update')
q.copy_(original_q)
row={'context':length,'splits':splits,'validation':metrics,'graph_mutation_exact':True,
**paired(candidate,control,100)}
Test 27. Real-KV FA2 split-count sweep (Oct 6)
Visual reference: Diagram 10 · Topic overview.
The idea: We revisit attention splitting using the real saved query and KV pages instead of synthetic arrays. It holds those inputs fixed, changes how many pieces of the history are processed in parallel, and checks both the captured answers and an independent floating-point reference.
The result: At 4K, 16K and 32K, 16 splits improve attention by 14.18%, 7.16% and 4.04%. Forty splits give −6.00%, +6.10% and +12.86%. All 48 cases pass the 16-layer and graph-mutation checks, and nine FP32 references pass. Worst relative RMS error is 0.00147142.
Explanation, checks and code
Why it matters on the GPU: The actual 784-token pages and interleaved K/V strides determine the addresses threads read. Increasing splits spreads a query over more blocks and multiprocessors, but creates more partial softmax/output data to store and combine. The best split count depends on history length, page access and the GPU's number of multiprocessors.
What to take from it: Sixteen splits helps across these captured contexts; forty loses at 4K but helps at 32K. The different result from a synthetic split test is expected because the memory layout, bounds and work distribution differ.
Setup: We try 16 split settings at three contexts, using 16 layers captured at decode step 32. Pages hold 784 tokens and the maximum K bound is 34,816. Seven graph pairs include an automatic/automatic control; the KV cache is read-only.
Relevant code:
The split count goes directly to the installed FA2 kernel over saved layer tensors.
Source: pq2-dev/benchmarks/experiments/attention/pq2_attention_split_worker.py:60–61.
def run(count):
return [torch.ops._vllm_fa2_C.varlen_fwd(a['q'],a['k'],a['v'],None,a['cu_seqlens_q'],a['cu_seqlens_k'],a['seqused_k'],None,a['block_table'],None,1,max_k,0.,a['softmax_scale'],False,True,-1,-1,0.,False,count,None)[0] for a in layers]
Captured-output and FP32-reference checks precede numerical, changing-input graph and paired timing checks.
Source: pq2-dev/benchmarks/experiments/attention/pq2_attention_split_worker.py:62–108.
reference = run(0)
checks=[]
for i,(a,actual) in enumerate(zip(layers,reference)):
metrics = error_metrics(actual,a['expected'])
assert metrics['finite'] and metrics['relative_rms']<.01, ('Captured cache/query mismatch',i,metrics)
checks.append({'layer':i,'eager_attention_comparison':metrics})
# Independent grouped-query FP32 reference, without replicating KV heads.
for i in (0,7,15):
a=layers[i];length=int(a['seqused_k'][0]);page=a['k'].shape[1]
pages=a['block_table'][0,:(length+page-1)//page].long()
keys=a['k'].index_select(0,pages).reshape(-1,4,256)[:length].permute(1,2,0).float()
values=a['v'].index_select(0,pages).reshape(-1,4,256)[:length].permute(1,0,2).float()
query=a['q'][0].reshape(4,6,256).float()
truth=((query@keys*a['softmax_scale']).softmax(-1)@values).reshape(1,24,256)
metrics=error_metrics(reference[i],truth)
assert metrics['finite'] and metrics['relative_rms']<.01, ('Independent FP32 mismatch',i,metrics)
checks[i]['fp32_comparison']=metrics
del keys,values,query,truth
control,reference = graph_for(lambda:run(0))
control.replay();torch.cuda.synchronize()
records=[];failures=[];saved_q=[a['q'].clone() for a in layers]
try:
for count in splits:
try:
actual=run(count)
validation=[error_metrics(x,y) for x,y in zip(actual,reference)]
assert all(v['finite'] and v['relative_rms']<.01 for v in validation), ('Split numerical mismatch',count,validation)
candidate,outputs=graph_for(lambda:run(count))
candidate.replay();torch.cuda.synchronize()
old=outputs[0].clone()
layers[0]['q'].mul_(.8).add_(.04)
changed=run(count)
candidate.replay();torch.cuda.synchronize()
assert not torch.equal(outputs[0],old), 'Graph ignored modified query'
for x,y in zip(outputs,changed):torch.testing.assert_close(x,y,atol=0,rtol=0)
for a,q in zip(layers,saved_q):a['q'].copy_(q)
row={'splits':count,'validation_max_relative_rms':max(v['relative_rms'] for v in validation),'graph_mutation_exact':True,**paired(candidate,control,100)}
records.append(row);print('ATTENTION_SPLIT',metadata[0]['used_k'],count,row['speedup'],flush=True)
del candidate,outputs
except Exception as error:
failures.append({'splits':count,'error':str(error)})
for a,q in zip(layers,saved_q):a['q'].copy_(q)
print('ATTENTION_REJECTED',count,str(error),flush=True)
write_json(root/'progress.json',records)
summary={'complete':True,'actual_context':metadata[0]['used_k'],'layers':16,'max_k':max_k,'page_size':layers[0]['k'].shape[1],'checks':checks,'records':records,'failures':failures,'scope':'Real attention queries and live paged KV from one single-request decode step across all 16 full-attention layers. Each graph replays all 16 layers in order, excluding cache updates and other model operations. Maximum K bound matches the serving graph.'}
write_json(root/'summary.json',summary)
assert len(records)==len(splits) and not failures, ('Sweep incomplete',failures)
Test 28. Real-KV attention backend screen (Oct 6)
Visual reference: Diagram 10 · Topic overview.
The idea: We give several attention implementations the same real query and saved paged key/value cache. It times only the attention calculations, trying FA2 split choices, Triton segment counts, and FlashInfer's ordinary-core and Tensor Core routes.
The result: On GPU 1, Triton with 32 segments gives 1.517×, 1.291× and 1.267× against FA2 automatic splitting at 4K, 16K and 32K. With 64 segments, it gives 1.437×, 1.335× and 1.424×. FlashInfer Tensor Cores give 1.336×, 1.397× and 1.305×. The final screen passes 480 layer, 9 FP32-reference and 30 poisoned-output graph checks; it replaces earlier screens with insufficient mutation checks.
Explanation, checks and code
Why it matters on the GPU: Keys and values occupy pages with the actual model's strides, rather than a convenient new contiguous array. That layout affects how adjacent threads combine their memory requests. Segmentation determines how many blocks can work in parallel. Tensor Core routes use matrix-multiply hardware, while ordinary-core routes use other arithmetic units; planning and temporary workspace also affect practical integration, even though planning is excluded from these timings.
What to take from it: The winners depend on context length and implementation. The numerical checks and poisoned-output graph checks make the component comparison stronger, but it does not measure cache updates or request throughput.
Setup: We use real decode-step-32 queries and paged KV from 16 layers. Contexts are 4K, 16K and 32K, with 24 query heads, 4 KV heads, head dimension 256 and 784-token pages. Timings use seven pairs and exclude KV updates and planning.
Relevant code:
Direct FA2 and segmented Triton calls over captured layer tensors; this excerpt excludes capture/engine setup.
Source: pq2-dev/benchmarks/experiments/attention/pq2_attention_kernel_worker.py:67–82.
def run(count):
if isinstance(count, str):
segments = int(count.split('-')[-1])
outputs = []
for a, (out, partial, maximum, expsum) in zip(layers,buffers[segments]):
unified_attention(a['q'],a['k'],a['v'],out,
a['cu_seqlens_q'],1,a['seqused_k'],max_k,
a['softmax_scale'],True,(-1,-1),a['block_table'],0.,
None,None,None,seq_threshold_3D=1 if segments else None,
num_par_softmax_segments=segments or None,
softmax_segm_output=partial if segments else None,
softmax_segm_max=maximum if segments else None,
softmax_segm_expsum=expsum if segments else None)
outputs.append(out)
return outputs
return [torch.ops._vllm_fa2_C.varlen_fwd(a['q'],a['k'],a['v'],None,a['cu_seqlens_q'],a['cu_seqlens_k'],a['seqused_k'],None,a['block_table'],None,1,max_k,0.,a['softmax_scale'],False,True,-1,-1,0.,False,count,None)[0] for a in layers]
FlashInfer workspace/plan setup and direct captured-KV execution; planning is outside replay timing.
Source: pq2-dev/benchmarks/experiments/attention/pq2_flashinfer_kernel_worker.py:69–105.
a = layers[0]
page = a['k'].shape[1]
length = int(a['seqused_k'][0])
indices = a['block_table'][0,:(length+page-1)//page].contiguous()
indptr = torch.tensor([0,indices.numel()],device=a['q'].device,dtype=torch.int32)
last = torch.tensor([(length-1)%page+1],device=a['q'].device,dtype=torch.int32)
workspace = torch.empty(128*1024*1024,device=a['q'].device,dtype=torch.uint8)
for tc in (False,True):
wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(workspace,
use_cuda_graph=True,use_tensor_cores=tc,
paged_kv_indptr_buffer=indptr,paged_kv_indices_buffer=indices,
paged_kv_last_page_len_buffer=last)
wrapper.plan(indptr.cpu(),indices.cpu(),last.cpu(),24,4,256,page,
q_data_type=torch.bfloat16,kv_data_type=torch.bfloat16,
sm_scale=a['softmax_scale'])
wrappers['flashinfer-'+str(int(tc))] = (wrapper,[torch.empty_like(a['q']) for a in layers])
def run(count):
if isinstance(count,str) and count.startswith('flashinfer'):
wrapper, outputs = wrappers[count]
for a,out in zip(layers,outputs):
wrapper.run(a['q'],(a['k'],a['v']),out=out)
return outputs
if isinstance(count, str):
segments = int(count.split('-')[-1])
outputs = []
for a, (out, partial, maximum, expsum) in zip(layers,buffers[segments]):
unified_attention(a['q'],a['k'],a['v'],out,
a['cu_seqlens_q'],1,a['seqused_k'],max_k,
a['softmax_scale'],True,(-1,-1),a['block_table'],0.,
None,None,None,seq_threshold_3D=1 if segments else None,
num_par_softmax_segments=segments or None,
softmax_segm_output=partial if segments else None,
softmax_segm_max=maximum if segments else None,
softmax_segm_expsum=expsum if segments else None)
outputs.append(out)
return outputs
return [torch.ops._vllm_fa2_C.varlen_fwd(a['q'],a['k'],a['v'],None,a['cu_seqlens_q'],a['cu_seqlens_k'],a['seqused_k'],None,a['block_table'],None,1,max_k,0.,a['softmax_scale'],False,True,-1,-1,0.,False,count,None)[0] for a in layers]
Graph validation poisons outputs before replay to verify that modified queries are consumed.
Source: pq2-dev/benchmarks/experiments/attention/pq2_attention_kernel_worker.py:112–124.
old=outputs[0].clone()
layers[0]['q'].mul_(.8).add_(.04)
changed=[x.clone() for x in run(count)]
for x in outputs:x.fill_(float('nan'))
candidate.replay();torch.cuda.synchronize()
assert not torch.equal(outputs[0],old), 'Graph ignored modified query'
for x,y in zip(outputs,changed):torch.testing.assert_close(x,y,atol=0,rtol=0)
for a,q in zip(layers,saved_q):a['q'].copy_(q)
row={'splits':count,'validation_max_relative_rms':max(v['relative_rms'] for v in validation),'graph_mutation_exact':True,**paired(candidate,control,100)}
records.append(row);print('ATTENTION_SPLIT',metadata[0]['used_k'],count,row['speedup'],flush=True)
del candidate,outputs
except Exception as error:
failures.append({'splits':count,'error':str(error)})
Test 29. INT8 KV attention microbenchmark (Oct 7)
Visual reference: Diagram 10 · Topic overview.
The idea: We store the attention history as 8-bit keys and values, then compares several ways to use that compressed history. One custom mode also quantises the query/key dot-product arithmetic; another keeps that part in BF16 to reduce added error.
The result: With one query at 4K, FA2 BF16 takes 43.91 µs and unified INT8 takes 52.61 µs. The custom 32-segment BF16-QK path takes 24.22 µs with relative error 0.002169; INT-QK takes 22.90 µs with error 0.02058. At 32K, BF16-QK takes 124.74 µs against FA2's 188.21 µs. These errors are the recorded rel_err values, not task accuracy.
Explanation, checks and code
Why it matters on the GPU: Eight-bit KV storage reduces the bytes needed for the history compared with 16-bit storage, but introduces scale reads and conversion work. The arithmetic mode determines whether the query/key product uses integer or BF16 operands. Segments divide a long history across the GPU before a combine step. The test measures these actual kernel paths; halving stored bytes doesn’t automatically halve latency.
What to take from it: The INT-QK mode is slightly faster in the quoted 4 K case but has substantially larger recorded relative error than BF16-QK. This exposes a concrete speed-versus-numerical-error choice in attention, separate from weight compression.
Setup: We create standalone PyTorch tensors for one decode query or six verification queries, at 4K, 32K and 64K contexts. CUDA graphs time direct kernel calls; no vLLM engine is launched.
Relevant code:
CUDA event timing around repeated captured graph execution.
Source: pq2-dev/benchmarks/suites/pq2_int8_kv_attn_bench.py:21–34.
def timed(fn, iters=50):
fn()
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
fn()
g.replay()
torch.cuda.synchronize()
start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iters):
g.replay()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) / iters * 1000 # us
The standalone matrix compares custom INT8-QK/BF16-QK with direct unified-attention and FA2 component calls.
Source: pq2-dev/benchmarks/suites/pq2_int8_kv_attn_bench.py:68–91.
def main():
rows_out = []
for rows in (1, 6):
for used in (4096, 32768, 65536, 131072):
a = t._inputs([used], [rows], seed=used + rows)
ref = t._reference(a)
out = torch.empty_like(a['q'])
row = {'query_tokens': rows, 'context': used}
row['fa2_bf16_us'] = timed(fa2_bf16(used, rows))
row['vllm_int8_us'] = timed(lambda: vllm_int8(a, out, used))
row['vllm_int8_rel_err'] = t._rel(out, ref)
for segments in (32, 64):
for int_qk in (True, False):
key = f'ours_s{segments}_{"intqk" if int_qk else "bf16qk"}'
row[key + '_us'] = timed(lambda: int8_kv_attention(**a, out=out, segments=segments,
int_qk=int_qk))
row[key + '_rel_err'] = t._rel(out, ref)
rows_out.append(row)
print(json.dumps({k: (round(v, 4) if isinstance(v, float) else v) for k, v in row.items()}),
flush=True)
del a, ref
torch.cuda.empty_cache()
with open(sys.argv[1], 'w') as f:
json.dump(rows_out, f, indent=1)
J. Following the layout idea through to real inputs
Takeaway. Follow a promising design through real weights, other readers and the production projection API. These component checks remain separate from request speed and answer quality.
From the blog: these are further component checks related to chapter 9. They also revisit the one-token and small-batch ideas in posts 4 and 7.
Start with tests 24 and 25 for the initial load probes and tile-major kernel. The follow-ups below ask three questions: can split-K fix the small-output cases, can we keep one weight layout for every reader, and can a simpler one-token arithmetic path avoid a layout change altogether?
New terms: a completion counter counts finished blocks so the last one can combine their partial sums. A thread fence makes memory writes visible in the required order; it is not a whole-GPU barrier. Window interleaving rearranges the small pieces inside a row’s 128-byte code window.
Shared-memory banks serve different addresses in parallel; a bank conflict makes some simultaneous accesses wait. An L2 sector is a chunk counted by the cache traffic counters. RB counts independent row slabs handled by a warp in these one-token prototypes; U counts windows kept in flight.
Neither is a universal hardware unit.
Diagram 11: Improve the reads without changing every weight reader
The ring changes the load schedule rather than the stored model format. A fixed combination order gives repeatability; it only gives the same bits as the old path when its grouping also matches.
Diagram 12: One-token DP4A without expanding every signed code
The input-sum correction handles the weight code’s minus-one offset. This keeps the shared row-major weight layout. Different floating-point summation order can still change the final output bits.
| Follow-up | Main question |
|---|---|
| 30, 31 | Does split-K turn the tile-major probe into a useful real-input projection? |
| 32, 34, 35 | Ring buffering or a new window layout—and what happens to the other readers? |
| 33, 37, 38 | What works for one token, and does random weight data mislead us? |
| 36, 39 | Do the selected kernels help when called through the production projection API? |
Reading the percentages: in this follow-up group, a latency reduction is 1 − candidate time / reference time. It is not the speedup ratio used elsewhere. For example, 13.52 → 10.88 ms is 19.5% less time and about 1.24× speed. Each comparison keeps its own baseline; do not combine the savings across rows.
Test 30. Split-K repairs the small-output tile-major cases (15 October)
Visual reference: Diagram 11 · Follow-up overview.
The idea: We split the long input reduction across several blocks, then let the last block combine their partial answers. This gives the GPU more independent work without launching a separate final reduction kernel.
The result: Down falls from 50.3 to 35.9 µs, a 28.6% latency reduction, with the same output bits when split 4 reproduces the production order. Output falls from 21.2 to 16.4 µs, a 22.6% reduction, but differs by up to about 1e-3 in the synthetic sweep. Gate/up, input and the head also improve; QKV is exact at split 1.
Explanation, checks and code
Why it matters on the GPU: Tile-major loads helped large matrices in test 25, but smaller output dimensions produced too few blocks. Split-K creates more blocks along K. A completion counter tells the last block when all partials are ready; it sums split 0, then split 1, and so on. That order is deterministic, but only matches the reference bits if the reference uses the same grouping.
What to take from it: This follow-up fixes the earlier down-projection regression. It does not make every output exact, and the completion protocol still needs temporary storage and synchronisation.
Setup: Rotating-weight M = 8 sweeps include the separate activation-quantisation launch on both sides. The candidate uses tile-major weights, row tiling and configurable split-K. Percentages here are calculated as 1 − candidate time / reference time.
Relevant code:
The split path publishes partial sums, elects the last block, reduces in fixed order and resets its counter.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/mma_tmk_fixed.cu:130–166.
y[(size_t)tk * N + row0[r]] = from_float<T>(acc[r][t][j]);
y[(size_t)tk * N + row1[r]] = from_float<T>(acc[r][t][2 + j]);
}
return;
}
// Split-K: write this split's FP32 partials, then the last CTA to finish a row block sums the splits in
// fixed order (z = 0..split-1) and writes the output; deterministic regardless of arrival order.
#pragma unroll
for (int r = 0; r < RT; ++r)
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int j = 0; j < 2; ++j) {
const int tk = t * 8 + tid * 2 + j;
if (tk >= M) continue;
float* p = partial + ((size_t)blockIdx.y * M + tk) * N;
p[row0[r]] = acc[r][t][j];
p[row1[r]] = acc[r][t][2 + j];
}
__threadfence();
__syncthreads();
__shared__ unsigned is_last;
if (threadIdx.x == 0) is_last = (atomicAdd(&counters[blockIdx.x], 1u) == (unsigned)(split - 1));
__syncthreads();
if (!is_last) return;
__threadfence();
const int rows_per_block = (blockDim.x >> 5) * RT * 16;
const int base = blockIdx.x * rows_per_block;
for (int i = threadIdx.x; i < rows_per_block * M; i += blockDim.x) {
const int tk = i / rows_per_block, row = base + i % rows_per_block;
if (row >= N) continue;
float a = __ldcg(partial + (size_t)tk * N + row);
for (int z = 1; z < split; ++z) a += __ldcg(partial + ((size_t)z * M + tk) * N + row);
y[(size_t)tk * N + row] = from_float<T>(a);
}
if (threadIdx.x == 0) counters[blockIdx.x] = 0;
}
The driver tries split counts and checks exactness before timing the complete call.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bench_tmk.py:44–65.
y=torch.empty(M,n,dtype=torch.bfloat16,device="cuda")
part=torch.empty(8*M*n,dtype=torch.float32,device="cuda"); cnt=torch.zeros(1<<16,dtype=torch.int32,device="cuda")
best=None; allexact=True; bests={}
ST=torch.cuda.current_stream
for rt in (1,2,4):
if n%(16*rt): continue
tm=[tile_major(w,n,k,rt) for w in ws]
for split,wp,un in itertools.product((1,2,4,8),(2,4),(2,4,8)):
if (k//128)%split: continue
args=lambda c,sc: (q.data_ptr(),s.data_ptr(),t.data_ptr(),c.data_ptr(),sc.data_ptr(),y.data_ptr(),part.data_ptr(),cnt.data_ptr(),M,n,k,split,wp,un,rt,torch.cuda.current_stream().cuda_stream)
if L(*args(tm[0][0],tm[0][1]))!=0: continue
torch.cuda.synchronize()
ex=torch.equal(y,ref); err=float((y.float()-ref.float()).abs().max())
def mk(c,sc): return lambda: (sm._quantize_fragments[(M*(k//128),)](x,q,s,t,k), L(*args(c,sc)))
tt=timed([mk(c,sc) for c,sc in tm])
key=split
if key not in bests or tt<bests[key][0]: bests[key]=(tt,rt,wp,un,ex,err)
if best is None or tt<best[0]: best=(tt,rt,wp,un,ex,err,split)
gb=nbytes*n*0+ (n*(k//128)*34)/1e9
print(f"{name:8s} {n}x{k}: prod {tprod:6.1f}us ({gb/tprod*1e6:4.0f} GB/s) | best {best[0]:6.1f}us ({gb/best[0]*1e6:4.0f} GB/s) split={best[6]} RT={best[1]} w{best[2]} u{best[3]} {100*(1-best[0]/tprod):+3.0f}% exact={best[4]} maxerr={best[5]:.1e}",flush=True)
print(" per split: "+" ".join(f"s{k_}: {v[0]:.1f}us (RT{v[1]} w{v[2]} u{v[3]}, exact={v[4]})" for k_,v in sorted(bests.items())),flush=True)
del ws
Test 31. Tile-major split-K on real eight-row captures (15 October)
Visual reference: Diagram 11 · Follow-up overview.
The idea: We check whether the synthetic split-K improvements survive real model weights and token inputs, rather than repeatedly using random matrices.
The result: All 257 projections take 13.52 → 10.88 ms in the captured replay: 2.64 ms saved, or a 19.5% latency reduction. Of 771 samples, 635 match production bit for bit. All 136 differing samples are output projections; maximum absolute difference is about 3.9e-3 and relative RMS is about 1.3e-4. All outputs are finite.
Explanation, checks and code
Why it matters on the GPU: The replay streams weights from different layers, so a small cached matrix cannot hide load costs. Split-K helps fill the GPU on the smaller output shapes, while tile-major storage supplies each warp with adjacent weight pieces.
What to take from it: The full projection pass benefits, but this is still a component replay. Output-projection error against the older captured baseline stays about 8.5e-3 for both paths; that does not mean the candidate and current production have identical outputs.
Setup: 771 samples cover 257 real projections at decode steps 1, 32 and 128, with eight rows each. CUDA-graph timing uses rotating real weights and interleaved comparisons. Layout conversion is outside replay timing. The saved replay_tm.json and replay_tm.log retain the numerical and timing evidence.
Relevant code:
The candidate quantises each input and calls the selected split-K kernel.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_tm.py:36–40.
def mine(l,xr,q,s,t,y):
n,k=l["n"],l["k"]; sp,rt,wp,un=l["cfg"]; c,sc=l["tm"]
sm._quantize_fragments[(M*(k//128),)](xr,q,s,t,k)
r=L(q.data_ptr(),s.data_ptr(),t.data_ptr(),c.data_ptr(),sc.data_ptr(),y.data_ptr(),part.data_ptr(),cnt.data_ptr(),M,n,k,sp,wp,un,rt,torch.cuda.current_stream().cuda_stream)
assert r==0,r
Checks distinguish candidate-versus-production error from error against the older captured output.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_tm.py:41–56.
# --- numerical comparison on real data ---
stats={}
for (step,name),smp in sorted(samples.items()):
l=layers[name]; n,k=l["n"],l["k"]
x=smp["x"].cuda(); assert x.shape==(M,k),(x.shape,)
xr=prism.hadamard(x,l["signs"],1024)
prod=prism.pq2_matmul(xr,l["w"],prepared=True)
q=torch.empty(M,k,dtype=torch.int8,device="cuda"); s=torch.empty(M,k//128,dtype=torch.float32,device="cuda"); t=torch.empty(M,k//128,dtype=torch.int32,device="cuda")
y=torch.empty(M,n,dtype=prod.dtype,device="cuda"); mine(l,xr,q,s,t,y); torch.cuda.synchronize()
base=smp["y"].cuda()
d=(y.float()-prod.float())
rec=dict(equal=bool(torch.equal(y,prod)),ndiff=int((y!=prod).sum()),maxabs=float(d.abs().max()),
rel_mine_prod=float(d.norm()/prod.float().norm()),
rel_prod_base=float((prod.float()-base.float()).norm()/base.float().norm()),
rel_mine_base=float((y.float()-base.float()).norm()/base.float().norm()),finite=bool(torch.isfinite(y).all()))
stats.setdefault((n,k),[]).append(rec)
Graph timing repeats the projection call list over real weights.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_tm.py:71–87.
def timed(fns,reps=20):
st=torch.cuda.Stream()
with torch.cuda.stream(st):
for f in fns: f()
torch.cuda.synchronize()
g=torch.cuda.CUDAGraph()
with torch.cuda.graph(g,stream=st):
for f in fns: f()
for _ in range(3): g.replay()
torch.cuda.synchronize()
out=[]
for _ in range(5):
e0=torch.cuda.Event(enable_timing=True); e1=torch.cuda.Event(enable_timing=True)
e0.record()
for _ in range(reps): g.replay()
e1.record(); torch.cuda.synchronize(); out.append(e0.elapsed_time(e1)/reps)
return out
Test 32. Keep row-major weights and stage them through a ring (15 October)
Visual reference: Diagram 11 · Follow-up overview.
The idea: We try to get better loads without changing the stored weight layout. Several shared-memory buffers form a ring: while one tile is being used, another can be filled by asynchronous copies.
The result: In the recorded real-data session, the exact down split-4 ring gives 13.13 ms and saves 0.64 ms (4.6%) against its paired reference. Down split 2 gives 12.22 ms and saves 1.63 ms (11.8%), but only 630/771 outputs match bits because down is summed differently. The tile-major comparison remains 10.88 ms with 635/771 exact samples.
Explanation, checks and code
Why it matters on the GPU: The cp.async copies gather larger contiguous chunks of each existing row into shared memory. The arithmetic then reads a layout arranged for its Tensor Core operands. This spends shared-memory capacity and barrier work to improve the load/compute schedule, but avoids keeping a second full weight format.
What to take from it: The faster split-2 setting changes rounding order; the exact setting gives a smaller gain. Baselines in this session range from 13.5–13.9 ms, so compare the reported paired savings rather than subtracting one universal reference. Later dispatch-level results in test 36 use a newer selection.
Setup: The original layout comparison uses synthetic rotating-weight sweeps plus 257 real projections and 771 samples. The ring variants differ in the down-projection split order. Tests 32, 34 and 36 are successive experiments, not gains to add together.
Relevant code:
Each stage asynchronously copies weight bytes, activation bytes and group metadata into shared memory.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/ring.cu:45–73.
auto load_stage = [&](int s, int it) {
uint8_t* Wc = smem + s * C::STAGE;
uint8_t* Q = Wc + C::WC;
uint8_t* Ws = Q + C::QB;
uint8_t* Xs = Ws + C::WSB;
uint8_t* Qs = Xs + C::XSB;
const int g0 = g_begin + it * SG;
constexpr int WCH = SG * 32 / 16;
for (int i = threadIdx.x; i < C::BMW * WCH; i += C::NTH) {
const int r = i / WCH, c = i % WCH;
cp_async<16>(Wc + r * C::WSTRIDE + c * 16, codes + (size_t)min(row0 + r, N - 1) * (K >> 2) + g0 * 32 + c * 16);
}
constexpr int QCH = SG * 8;
for (int i = threadIdx.x; i < C::BT * QCH; i += C::NTH) {
const int t = i / QCH, c = i % QCH;
cp_async<16>(Q + t * C::QSTRIDE + c * 16, q + (size_t)min(t, M - 1) * K + g0 * 128 + c * 16);
}
for (int i = threadIdx.x; i < C::BMW; i += C::NTH)
cp_async<SG * 2>(Ws + i * SG * 2, wscale + (size_t)min(row0 + i, N - 1) * G + g0);
for (int i = threadIdx.x; i < C::BT; i += C::NTH) {
cp_async<SG * 4>(Xs + i * SG * 4, xs + (size_t)min(i, M - 1) * G + g0);
cp_async<SG * 4>(Qs + i * SG * 4, qsum + (size_t)min(i, M - 1) * G + g0);
}
};
float acc[RT][NT][4];
#pragma unroll
for (int r = 0; r < RT; ++r)
The main loop waits for the needed stage and schedules the next stage ahead.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/ring.cu:85–105.
const uint8_t* Wc = smem + (it % ST) * C::STAGE;
const uint8_t* Q = Wc + C::WC;
const __half* Ws = reinterpret_cast<const __half*>(Q + C::QB);
const float* Xs = reinterpret_cast<const float*>(reinterpret_cast<const uint8_t*>(Ws) + C::WSB);
const int* Qs = reinterpret_cast<const int*>(reinterpret_cast<const uint8_t*>(Xs) + C::XSB);
#pragma unroll
for (int gl = 0; gl < SG; ++gl) {
uint32_t b[NT][8];
float x0[NT], x1[NT];
int q0[NT], q1[NT];
#pragma unroll
for (int n = 0; n < NT; ++n) {
const uint4* bp = reinterpret_cast<const uint4*>(Q + (n * 8 + gid) * C::QSTRIDE + gl * 128 + tid * 32);
const uint4 u0 = bp[0], u1 = bp[1];
b[n][0] = u0.x; b[n][1] = u0.y; b[n][2] = u0.z; b[n][3] = u0.w;
b[n][4] = u1.x; b[n][5] = u1.y; b[n][6] = u1.z; b[n][7] = u1.w;
const int tc = n * 8 + tid * 2;
x0[n] = Xs[tc * SG + gl]; x1[n] = Xs[(tc + 1) * SG + gl];
q0[n] = 0x4B400000 - Qs[tc * SG + gl]; q1[n] = 0x4B400000 - Qs[(tc + 1) * SG + gl];
}
#pragma unroll
The replay chooses the ring or the reference on a per-shape basis.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_ring2.py:37–44.
def mine(l,xr,q,s,t,y):
n,k=l["n"],l["k"]
if l["cfg"] is None:
y.copy_(prism.pq2_matmul(xr,l["w"],prepared=True)); return
sp,cfg=l["cfg"]; w=l["w"]
quantize_prefill(xr,q,s,t)
r=L(q.data_ptr(),s.data_ptr(),t.data_ptr(),w.data_ptr(),w.data_ptr()+n*k//4,y.data_ptr(),part.data_ptr(),cnt.data_ptr(),M,n,k,sp,cfg,torch.cuda.current_stream().cuda_stream)
assert r==0,r
Test 33. Does a tile-major layout help one token? (15 October)
Visual reference: Diagram 12 · Follow-up overview.
The idea: We reuse the alternative layout at M = 1 to see whether a memory arrangement that helped eight rows also helps a single token.
The result: The naïve Triton tile-major port is about twice as slow on gate/up: roughly 97–115 µs against 56 µs. The tile-major MMA alternative gives reported latency reductions of about 6% on gate/up and input, 8% on QKV, 3% on down and 4% on the head; output is about 1% slower. The source estimates roughly 4% across the weighted projection step.
Explanation, checks and code
Why it matters on the GPU: The Triton port needs wide integer loads followed by byte extraction and operand rearrangement. With one row, there is much less reuse to pay for that work. The MMA alternative changes the arithmetic path as well as the layout. A good load-only result cannot predict either complete kernel by itself.
What to take from it: These results argue against a universal layout switch for single-row decode. The source attributes the Triton problem to integer instruction work; treat that as its diagnosis rather than a separately measured instruction-pipeline limit. Test 38 later shows another complication: real weight content changes timings.
Setup: The standalone M = 1 comparison rotates weights and includes input quantisation. single_tm.py implements the Triton port; bench_tmk_m1.py tries the Tensor Core alternative. No full-model speed measurement is included here.
Relevant code:
The alternative Triton kernel reads tile-major words and reconstructs bytewise dot operands.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/single_tm.py:8–38.
@triton.jit
def _single_tm(Q, S, W, WS, Y, N: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr):
"""M = 1 DP4A GEMV on tile-major planes. A program handles BT 16-row tiles. Per (tile, group) the 512 B
are 64 contiguous int64 pieces l = (gid * 4 + tid) * 2 + h (8 B of row h*8 + gid, bytes tid*8..tid*8+7)."""
G: tl.constexpr = K // 128
W64 = W.to(tl.pointer_type(tl.int64))
t = tl.program_id(0) * BT + tl.arange(0, BT)
lane = tl.arange(0, 64)
j = tl.arange(0, 8)
offset = tl.arange(0, BC)
okt = t < N // 16
acc = tl.full((BT, 64), 0, tl.float32)
for base in tl.range(0, tl.cdiv(G, BC), num_stages=1):
blocks = base * BC + offset
okb = blocks < G
ok = okt[:, None, None] & okb[None, :, None]
tg = t[:, None, None] * G + blocks[None, :, None]
v = tl.load(W64 + tg * 64 + lane[None, None, :], ok, other=0) # [BT, BC, 64] int64
ws = tl.load(WS + (tg * 8 + (lane // 8)[None, None, :]) * 2 + (lane % 2)[None, None, :], ok, other=0).to(tl.float32)
b = ((v[:, :, :, None] >> (j * 8)) & 0xFF).to(tl.int32) # [BT, BC, 64, 8]
w = _decode_codes(b)
q = tl.load(Q + blocks[:, None, None] * 32 + ((lane // 2) % 4)[None, :, None] * 8 + j[None, None, :],
okb[:, None, None], other=0) # [BC, 64, 8]
dots = tl.inline_asm_elementwise('dp4a.s32.s32 $0, $1, $2, 0;', constraints='=r,r,r',
args=[q[None, :, :, :], w], dtype=tl.int32, is_pure=True, pack=1)
partial = tl.sum(dots, 3).to(tl.float32) # [BT, BC, 64]
xs = tl.load(S + blocks, okb, other=0)
acc += tl.sum(partial * xs[None, :, None] * ws, 1)
per = tl.sum(tl.reshape(acc, (BT, 8, 4, 2)), 2) # [BT, gid, h]
rows = t[:, None, None] * 16 + tl.arange(0, 2)[None, None, :] * 8 + tl.arange(0, 8)[None, :, None]
tl.store(Y + rows, per, okt[:, None, None] & (rows < N))
The M = 1 driver searches split, row-tile and unroll settings and records output agreement.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bench_tmk_m1.py:45–64.
tprod=timed([ (lambda w=w: prodf(w)) for w in ws])
y=torch.empty(M,n,dtype=torch.bfloat16,device="cuda")
part=torch.empty(8*M*n,dtype=torch.float32,device="cuda"); cnt=torch.zeros(1<<16,dtype=torch.int32,device="cuda")
best=None; allexact=True; bests={}
ST=torch.cuda.current_stream
for rt in (1,2,4):
if n%(16*rt): continue
tm=[tile_major(w,n,k,rt) for w in ws]
for split,wp,un in itertools.product((1,2,4,8),(2,4),(2,4,8)):
if (k//128)%split: continue
args=lambda c,sc: (q.data_ptr(),s.data_ptr(),t.data_ptr(),c.data_ptr(),sc.data_ptr(),y.data_ptr(),part.data_ptr(),cnt.data_ptr(),M,n,k,split,wp,un,rt,torch.cuda.current_stream().cuda_stream)
if L(*args(tm[0][0],tm[0][1]))!=0: continue
torch.cuda.synchronize()
ex=torch.equal(y,ref); err=float((y.float()-ref.float()).abs().max())
def mk(c,sc): return lambda: (sm._quantize_fragments[(M*(k//128),)](x,q,s,t,k), L(*args(c,sc)))
tt=timed([mk(c,sc) for c,sc in tm])
key=split
if key not in bests or tt<bests[key][0]: bests[key]=(tt,rt,wp,un,ex,err)
if best is None or tt<best[0]: best=(tt,rt,wp,un,ex,err,split)
gb=nbytes*n*0+ (n*(k//128)*34)/1e9
Test 34. Window-interleaved weights on real eight-row inputs (15 October)
Visual reference: Diagram 11 · Follow-up overview.
The idea: We keep each row together but reorder the four groups inside each 128-byte window. That puts the pieces a lane needs closer together without moving whole row tiles across the matrix.
The result: The 257-projection pass falls from 13.64 to 11.80 ms: 1.84 ms saved, or a 13.5% latency reduction. Gate/up falls from 5.540 to 4.509 ms, down from 4.265 to 2.951 ms, and output from 1.344 to 1.207 ms. Of 771 samples, 643 match bit for bit; all differences are output projections, with maximum relative RMS about 1.3e-4.
Explanation, checks and code
Why it matters on the GPU: Within a row, the code bytes are stored as window → lane piece → group → eight-byte piece. Scales are unchanged. Direct loads feed register fragments without the ring’s shared-memory staging. The kernel still needs split-K for some shapes, and every other reader must understand the new addresses.
What to take from it: This is a middle ground between unchanged row-major storage and full tile-major storage. The earlier layout screen had no window replay yet; this later run supplies it. Test 35 checks the cost to other readers before drawing an integration conclusion.
Setup: Real weights and eight-row activations cover 257 projections at three decode steps. The selected window kernel leaves QKV and the head on production. Comparing with the captured baseline gives unchanged output-projection relative error of about 8.53e-3.
Relevant code:
The physical rearrangement changes code bytes within each window and leaves scales separate.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bench_win.py:7–11.
def win_layout(w,n,k,rt):
G=k//128; flat=w.reshape(-1)
c=flat[:n*k//4].view(n,G//4,4,4,8).permute(0,1,3,2,4).contiguous() # [row][window][tid][gl][j]
s=flat[n*k//4:].contiguous() # row-major halves, untouched
return c,s
Selected shapes use the window kernel; unselected shapes retain the existing operation.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_win.py:36–43.
def mine(l,xr,q,s,t,y):
n,k=l["n"],l["k"]
if l["cfg"] is None:
y.copy_(prism.pq2_matmul(xr,l["w"],prepared=True)); return
sp,rt,wp,un=l["cfg"]; c,sc=l["tm"]
sm._quantize_fragments[(M*(k//128),)](xr,q,s,t,k)
r=L(q.data_ptr(),s.data_ptr(),t.data_ptr(),c.data_ptr(),sc.data_ptr(),y.data_ptr(),part.data_ptr(),cnt.data_ptr(),M,n,k,sp,wp,un,rt,torch.cuda.current_stream().cuda_stream)
assert r==0,r
The replay records bit equality and both production and captured-baseline error.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_win.py:44–59.
# --- numerical comparison on real data ---
stats={}
for (step,name),smp in sorted(samples.items()):
l=layers[name]; n,k=l["n"],l["k"]
x=smp["x"].cuda(); assert x.shape==(M,k),(x.shape,)
xr=prism.hadamard(x,l["signs"],1024)
prod=prism.pq2_matmul(xr,l["w"],prepared=True)
q=torch.empty(M,k,dtype=torch.int8,device="cuda"); s=torch.empty(M,k//128,dtype=torch.float32,device="cuda"); t=torch.empty(M,k//128,dtype=torch.int32,device="cuda")
y=torch.empty(M,n,dtype=prod.dtype,device="cuda"); mine(l,xr,q,s,t,y); torch.cuda.synchronize()
base=smp["y"].cuda()
d=(y.float()-prod.float())
rec=dict(equal=bool(torch.equal(y,prod)),ndiff=int((y!=prod).sum()),maxabs=float(d.abs().max()),
rel_mine_prod=float(d.norm()/prod.float().norm()),
rel_prod_base=float((prod.float()-base.float()).norm()/base.float().norm()),
rel_mine_base=float((y.float()-base.float()).norm()/base.float().norm()),finite=bool(torch.isfinite(y).all()))
stats.setdefault((n,k),[]).append(rec)
Test 35. Check the window layout against every reader (15 October)
Visual reference: Diagram 11 · Follow-up overview.
The idea: We change only the weight addresses in the other kernels, then check whether the proposed layout helps or hurts single-token decode, small batches, prefill and embedding lookup.
The result: All tested reader outputs remain bit-identical. However, M = 1 is about 7% slower on K = 5,120 shapes and 23.8% slower on the head. The M = 2 and M = 4 integer head paths slow by 91% and 61%. At M = 2,048, Triton prefill slows by roughly 18–29%; CUDA prefill with four-group windows slows gate/up by 16.0% and input by 14.3%. Embedding output is unchanged.
Explanation, checks and code
Why it matters on the GPU: The same compressed bytes now arrive in a different order. A layout that suits direct small-batch fragments can force more awkward loads or keep larger shared-memory read results live in a prefill kernel. Recorded CUDA-prefill counters show unchanged DRAM bytes (789 MB), fewer shared-memory bank conflicts (13.5M → 2.4M), but more registers (191 → 255).
What to take from it: Fewer bank conflicts did not make this kernel faster. The timing penalty tracks register growth, which supports the report’s diagnosis but does not prove a single cause. Window size 2 removes the extra L2-sector requests without removing the timing penalty. The extra eight-row saving over the ring is too small to ignore these wider regressions.
Setup: Triton readers change only their code-byte address expression. Comparisons use each reader’s row-major version on rotating real-shape weights. CUDA-prefill variants test windows of 4 and 2 groups; embedding checks 256 token IDs. Positive percentages in source section 10 mean more time, hence slower.
Relevant code:
The conversion preserves the original scale plane.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/triton_win.py:10–13.
def win_layout_full(w,n,k):
G=k//128; flat=w.reshape(-1).clone(); nb=n*k//4
flat[:nb]=flat[:nb].view(n,G//_WIN,_WIN,4,8).permute(0,1,3,2,4).reshape(-1)
return flat.view(w.shape)
The single-token reader changes its addresses while preserving its computation.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/triton_win.py:21–29.
@triton.jit
def _load_tile_win(W,WS,rows,blocks,N:tl.constexpr,K:tl.constexpr):
valid=(rows[:,None]<N)&(blocks[None,:]<K//128)
words=tl.arange(0,32)
address=(rows[:,None,None]*(K//4)+(blocks[None,:,None]//4)*128+(words[None,None,:]//8)*32
+(blocks[None,:,None]%4)*8+(words[None,None,:]%8))
scales=tl.load(WS+rows[:,None]*(K//128)+blocks[None,:],valid,other=0).to(tl.float32)
values=tl.load(W+address,valid[:,:,None],other=0).to(tl.int32)
return values,scales
The prefill check compares exact outputs and times row-major versus window loads.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bench_prefill_win.py:22–29.
y=torch.empty(M,n,dtype=torch.bfloat16,device="cuda")
st=lambda: torch.cuda.current_stream().cuda_stream
cfg=4 if n>=14336 else 6
def call(L,w): return L(q.data_ptr(),s.data_ptr(),t.data_ptr(),w.data_ptr(),w.data_ptr()+n*k//4,y.data_ptr(),M,n,k,cfg,st())
assert call(Lrm,rm[0])==0; torch.cuda.synchronize(); e_rm=torch.equal(y,ref)
y.zero_(); assert call(Lwin,wn[0])==0; torch.cuda.synchronize(); e_win=torch.equal(y,ref)
tr=timed([(lambda w=w: call(Lrm,w)) for w in rm],5); tw=timed([(lambda w=w: call(Lwin,w)) for w in wn],5)
print(f" {name:8s} {n}x{k} cfg{cfg}: row-major {tr:7.0f}us window {tw:7.0f}us {100*(tw/tr-1):+5.1f}% time | exact vs Triton: row-major={e_rm} window={e_win}",flush=True)
Test 36. The selected ring through the production projection API (15 October)
Visual reference: Diagram 11 · Follow-up overview.
The idea: We test the integrated ring selection through the ordinary projection function, using saved inputs rather than an inference server. The selection changes only shapes where the ring helps.
The result: At M = 8, all 257 projections fall from 13.70 to 12.43 ms: 1.27 ms saved, a 9.3% latency reduction. Down falls from 4.179 to 2.755 ms; output from 1.323 to 1.018 ms; QKV from 0.586 to 0.542 ms. Down matches 51/192 samples, with maximum relative RMS 1.07e-4. Output matches 192/192 and QKV 48/48. All results are finite and deterministic.
Explanation, checks and code
Why it matters on the GPU: Weights stay in their shared row-major format. The integrated dispatch sends down to split 2, selects suitable output splits by M, and uses the ring for QKV at M = 7 and 8. It keeps the direct-fragment path elsewhere. This avoids imposing a new layout on prefill, embeddings and one-token decode.
What to take from it: The clean synthetic sweep’s input-projection wins did not transfer to real data, so they were dropped. The exact-only mode retains only bit-identical entries: output at M = 6–8 and QKV at M = 7–8. Deterministic down results still differ from the old bits because the sum order changed. Serving and quality results in the same source section are excluded from this companion.
Setup: The clean sweep uses an idle GPU and rotating weights; an earlier overlapping run is discarded. The reported selected sweep has down reductions of about 17–37% and output about 14–32% at M = 2–8. Real replay then checks 771 samples through pq2_matmul with the ring flag off and on, and interleaves graph timings. Source integration is identified as commits 535f571 and 87b43cd.
Relevant code:
The clean sweep checks both the fastest setting and the fastest exactly matching setting.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bench_ring3.py:42–57.
best=None; best_exact=None; allrows=[]
for cfg,split in itertools.product(range(13),(1,2,4,8)):
nt2=cfg in (10,11,12)
if (M>8)!=nt2: continue
def run(w): return L(q.data_ptr(),s.data_ptr(),t.data_ptr(),w.data_ptr(),w.data_ptr()+n*k//4,y.data_ptr(),part.data_ptr(),cnt.data_ptr(),M,n,k,split,cfg,torch.cuda.current_stream().cuda_stream)
quantize_prefill(x,q,s,t)
if run(ws[0])!=0: continue
torch.cuda.synchronize()
ex=torch.equal(y,ref); err=float((y.float()-ref.float()).abs().max())
tt=timed([(lambda w=w: (quantize_prefill(x,q,s,t),run(w))) for w in ws])
allrows.append((tt,cfg,split,ex,err))
if best is None or tt<best[0]: best=(tt,cfg,split,ex,err)
if ex and (best_exact is None or tt<best_exact[0]): best_exact=(tt,cfg,split,ex,err)
gb=nbytes/1e9
RES.setdefault(str(M),{})[name]={"prod_us":tprod,"best":best,"best_exact":best_exact,"top5":sorted(allrows)[:5]}
json.dump(RES,open("/tmp/claude-1000/-home-amorley-vllm/9d014f21-8765-4ef9-9067-3fa7be39352d/scratchpad/ring/sweep_ring_clean.json","w"))
The same public projection path is used with the ring selection toggled.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_prod.py:14–17.
def run(name,x,ring):
small._RING=ring
l=layers[name]; xr=prism.hadamard(x,l["signs"],1024)
return prism.pq2_matmul(xr,l["w"],prepared=True)
Each saved sample is checked for equality, relative error, determinism and finite output.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_prod.py:18–24.
stats=collections.defaultdict(list)
for (step,name),smp in sorted(samples.items()):
x=smp["x"].cuda(); l=layers[name]
a=run(name,x,False); b=run(name,x,True); c=run(name,x,True)
d=(a.float()-b.float())
stats[(l["n"],l["k"])].append(dict(eq=bool(torch.equal(a,b)),rel=float(d.norm()/a.float().norm()),det=bool(torch.equal(b,c)),fin=bool(torch.isfinite(b).all())))
print("Real captured data (771 samples), ring ON vs OFF through pq2_matmul:")
Timing uses pre-transformed inputs and direct projection calls, not a serving request.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_prod.py:45–47.
def graph_time(nms,ring):
small._RING=ring
return timed([(lambda nm=nm: prism.pq2_matmul(inputs[nm],layers[nm]["w"],prepared=True)) for nm in nms])
Test 37. Plane-DP4A single-token prototypes and tuning (15 October)
Visual reference: Diagram 12 · Follow-up overview.
The idea: We give the one-token integer kernel the raw code planes directly, rather than expanding every code through PRMT. We then try prefetch, shared scales and fusing the activation quantiser into each block.
The result: The first prototype improves gate/up by about 3%, input 8%, QKV 12%, output 4% and the head 3%, but down is about 4% slower. The tuned v2 log gives weighted projection time 8.63 → 8.10 ms, a 6.2% latency reduction. Fusing the quantiser into the kernel is bit-identical to the Triton quantiser but is reported 6–21% slower.
Explanation, checks and code
Why it matters on the GPU: Plane-order inputs let DP4A multiply unsigned code lanes by the matching signed activation lanes. Subtracting the group’s activation sum supplies the code-minus-one correction. Adjacent 16-byte weight loads feed several rows, while activation values and scales are staged once per block. The fused variant repeats input quantisation in every output block, so removing a launch also duplicates a lot of work.
What to take from it: The v2 improvements are from synthetic rotating-weight tuning and should not be confused with the larger real-input gains in test 39. The plan’s operation-count estimates explain the design; they are not measured pipeline utilisation. Ideas still listed as future work in the original plan are not counted as completed tests.
Setup: bench_m1.py screens the first prototype; bench_m1v2.py varies row slabs (RB), windows in flight (U), warp count, shared-scale staging and prefetch. bench_m1v3.py tests quantiser fusion. Quantisation launches are included where separate, and numerical checks precede timing.
Relevant code:
Four raw planes feed DP4A, then a lane exchange and group-sum correction produce scaled partials.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/m1v2.cu:95–108.
const uint32_t wd[4] = {wv[u][r].x, wv[u][r].y, wv[u][r].z, wv[u][r].w};
const uint32_t qa[4] = {q0.x, q0.y, q0.z, q0.w}, qb[4] = {q1.x, q1.y, q1.z, q1.w};
const uint32_t qc[4] = {q2.x, q2.y, q2.z, q2.w}, qd[4] = {q3.x, q3.y, q3.z, q3.w};
int dot = 0;
#pragma unroll
for (int j = 0; j < 4; ++j) {
dot = __dp4a((int)( wd[j] & 0x03030303u), (int)qa[j], dot);
dot = __dp4a((int)((wd[j] >> 2) & 0x03030303u), (int)qb[j], dot);
dot = __dp4a((int)((wd[j] >> 4) & 0x03030303u), (int)qc[j], dot);
dot = __dp4a((int)((wd[j] >> 6) & 0x03030303u), (int)qd[j], dot);
}
dot += __shfl_xor_sync(0xffffffffu, dot, 1);
const float wsv = SSM ? __half2float(sws[(srow0 + r * 4) * G + g]) : ws[u][r];
if (half == 0) acc[r] += (float)(dot - qs) * x * wsv;
The sweep checks relative error and times complete candidate calls.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/bench_m1v2.py:23–36.
y=torch.empty(1,n,dtype=torch.bfloat16,device="cuda"); res=collections.defaultdict(lambda:None); best=None
for (rb,u,wp),ssm,pre in itertools.product(CFG,(0,1),(0,1)):
def run(w): return L(q.data_ptr(),s.data_ptr(),t.data_ptr(),w.data_ptr(),w.data_ptr()+n*k//4,y.data_ptr(),n,k,rb,u,wp,ssm,pre,torch.cuda.current_stream().cuda_stream)
sm._quantize_planes_sum[(G,)](x,q,s,t,k)
if run(rm[0])!=0: continue
torch.cuda.synchronize()
rel=float((y.float()-ref.float()).norm()/ref.float().norm())
if rel>5e-4: print(" BAD",name,(rb,u,wp,ssm,pre),rel); continue
tt=timed([(lambda w=w: (sm._quantize_planes_sum[(G,)](x,q,s,t,k),run(w))) for w in rm],10)
key=(ssm,pre)
if res[key] is None or tt<res[key][0]: res[key]=(tt,rb,u,wp,rel)
if best is None or tt<best[0]: best=(tt,rb,u,wp,ssm,pre,rel)
gb=nb/1e9
combos=" ".join(f"[ssm{k_[0]} pre{k_[1]}: {v[0]:5.1f}us RB{v[1]}U{v[2]}w{v[3]}]" for k_,v in sorted(res.items()) if v)
The fused variant quantises the input vector inside each output block.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/m1v3.cu:50–75.
if (QF) {
// One warp per 128-value group: lane owns elements 4*lane..4*lane+3 = byte `lane`, planes s = 0..3.
// Arithmetic identical to _quantize_planes_sum: scale = amax * (1/127) (FP32 reciprocal multiply), q = rint(a / scale).
for (int g = warp; g < G; g += WARPS) {
const uint2 raw = *reinterpret_cast<const uint2*>(xin + (size_t)g * 128 + lane * 4);
const T* e = reinterpret_cast<const T*>(&raw);
float a[4];
#pragma unroll
for (int i = 0; i < 4; ++i) a[i] = static_cast<float>(e[i]);
float m = fmaxf(fmaxf(fabsf(a[0]), fabsf(a[1])), fmaxf(fabsf(a[2]), fabsf(a[3])));
#pragma unroll
for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o));
const float scale = m * (1.0f / 127.0f);
const float div = scale > 0.f ? scale : 1.0f;
int sum = 0;
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float v = fminf(127.0f, fmaxf(-127.0f, rintf(__fdiv_rn(a[i], div))));
const int qi = (int)v;
sq[g * 128 + i * 32 + lane] = (int8_t)qi;
sum += qi;
}
#pragma unroll
for (int o = 16; o > 0; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
if (lane == 0) { sxs[g] = scale; sqs[g] = sum; }
}
Test 38. Real ternary weights versus random packed bytes (15 October)
Visual reference: Diagram 12 · Follow-up overview.
The idea: We keep the matrix shape and input vector fixed, but swap random packed weights for captured ternary weights. This checks whether the synthetic data itself changes the comparison.
The result: For gate/up, the Triton PRMT path takes about 60.9 µs on random bytes and 67.7 µs on real ternary weights—roughly 11% slower. The plane-DP4A path is almost unchanged at 55.6 and 56.0 µs. Real code values are about one third each of 0, 1 and 2, with no code 3.
Explanation, checks and code
Why it matters on the GPU: Two-bit storage permits four code values, but the model’s ternary weights only use three. The test changes the packed bit patterns seen by each instruction sequence, while preserving the overall matrix dimensions. The observed content sensitivity means a random-byte benchmark can understate the candidate’s advantage.
What to take from it: The measured difference is useful evidence, but this test does not isolate the hardware cause of the PRMT path’s sensitivity. It is a reason to confirm with real captured weights, not a general claim that random data always runs faster.
Setup: The script loads up to 24 distinct real gate/up weight tensors, creates matching random tensors, prints code frequencies, and times both kernels with the same one-row activation vector and graph-timing helper.
Relevant code:
The driver samples real code frequencies and runs the same two methods against real and random weight lists.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/real_vs_random.py:10–30.
L.restype=ctypes.c_int; L.argtypes=[ctypes.c_void_p]*6+[ctypes.c_int]*7+[ctypes.c_void_p]
n,k=34816,5120
real=[]
seen=set()
for r in man["layers"]:
if (r["n"],r["k"])==(n,k) and r["weight_file"] not in seen:
seen.add(r["weight_file"]); real.append(torch.load(CAP/r["weight_file"],weights_only=True,map_location="cpu").cuda())
if len(real)==24: break
rand,_=make(n,k,24)
w=real[0].reshape(-1)[:n*k//4]
print("real code bytes: fraction of 2-bit codes ==0/1/2/3:",[round(float(((w.view(torch.uint8)>>(2*s))&3).eq(v).float().mean()),3) for s in (0,) for v in range(4)])
wr=rand[0].reshape(-1)[:n*k//4]
print("random code bytes: ",[round(float(((wr>>(2*s))&3).eq(v).float().mean()),3) for s in (0,) for v in range(4)])
x=torch.randn(1,k,device="cuda",dtype=torch.bfloat16); G=k//128
q=torch.empty(1,k,dtype=torch.int8,device="cuda"); s=torch.empty(1,G,dtype=torch.float32,device="cuda"); t=torch.empty(1,G,dtype=torch.int32,device="cuda"); y=torch.empty(1,n,dtype=torch.bfloat16,device="cuda")
def mine(w):
sm._quantize_planes_sum[(G,)](x,q,s,t,k)
L(q.data_ptr(),s.data_ptr(),t.data_ptr(),w.data_ptr(),w.data_ptr()+n*k//4,y.data_ptr(),n,k,2,1,4,1,1,torch.cuda.current_stream().cuda_stream)
for label,ws in (("random weights",rand),("real weights ",real)):
tp=timed([(lambda w=w: pq2_int_gemv(x,w,prepared=True)) for w in ws]); tm=timed([(lambda w=w: mine(w)) for w in ws])
print(f"{label}: production {tp:6.1f} us plane-DP4A {tm:6.1f} us ({100*(1-tm/tp):+.0f}%)")
Test 39. The single-token CUDA path on real captured inputs (15 October)
Visual reference: Diagram 12 · Follow-up overview.
The idea: We enable the integrated one-token CUDA kernel through the normal projection API and check it against the previous path on all the saved single-row samples.
The result: All 257 projections fall from 11.29 to 8.81 ms: 2.48 ms saved, a 21.9% latency reduction. Per-shape latency reductions are 22.7% for gate/up, 23.4% for input, 15.1% for QKV, 15.9% for down, 7.9% for output and 37.8% for the head. All 771 outputs are finite and deterministic, but only 278 match the previous bits; maximum relative RMS is about 1.4e-4.
Explanation, checks and code
Why it matters on the GPU: This uses plane-DP4A on the existing prepared row-major weights. The input is quantised once into plane order, then the CUDA kernel shares those values across output rows and stages work to suit the one-token case. It changes FP32 group accumulation order, which explains why exact integer group products can still lead to different rounded outputs.
What to take from it: This is a projection-only flag comparison, not an engine benchmark. Its larger gain than the synthetic v2 estimate reflects a different real-data comparison, including the content sensitivity in test 38. The recorded results do not establish unchanged task answers. The source’s serving and GSM8K comparisons are deliberately excluded.
Setup: The replay loads 257 projections at steps 1, 32 and 128. It compares the public pq2_matmul path with _M1_CUDA disabled/enabled, repeats the candidate for determinism, and interleaves graph timings on pre-transformed inputs. The integrated kernel is identified as commit 3e2538b, gated to M = 1, prepared BF16/FP16 weights, SM86 and the supported shapes with K divisible by 512.
Relevant code:
The existing projection API is toggled between its old and new one-row paths.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_m1_prod.py:13–15.
def run(name,xr,flag):
ig._M1_CUDA=flag
return prism.pq2_matmul(xr,layers[name]["w"],prepared=True)
Saved samples are checked for bit equality, relative error, repeatability and finite values.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_m1_prod.py:16–21.
stats=collections.defaultdict(list)
for (step,name),smp in sorted(samples.items()):
l=layers[name]; xr=prism.hadamard(smp["x"].cuda(),l["signs"],1024)
a=run(name,xr,False).clone(); b=run(name,xr,True).clone(); c=run(name,xr,True).clone()
d=a.float()-b.float()
stats[(l["n"],l["k"])].append(dict(eq=bool(torch.equal(a,b)),rel=float(d.norm()/a.float().norm()),det=bool(torch.equal(b,c)),fin=bool(torch.isfinite(b).all())))
Graph timing uses each shape’s real weights and prepared inputs.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/replay_m1_prod.py:43–45.
def gt(nms,flag):
ig._M1_CUDA=flag
return timed([(lambda nm=nm: prism.pq2_matmul(inputs[nm],layers[nm]["w"],prepared=True)) for nm in nms])
The final lane reduction shows where the floating-point accumulation order differs.
Source: pq2-dev/benchmarks/results/pq2-decode-kernel-design-20261015/m1v2.cu:128–134.
#pragma unroll
for (int r = 0; r < RB; ++r) {
float v = acc[r];
v += __shfl_xor_sync(0xffffffffu, v, 2);
v += __shfl_xor_sync(0xffffffffu, v, 4);
if (part == 0 && ok[r]) y[row[r]] = static_cast<T>(v);
}
Reproduction and scope
The snippets preserve the selected experiment code. In the supplied extract, 44 of 104 source snapshots listed in its manifest are present and match their recorded hashes; 60 listed snapshots are absent. The inline examples remain useful for explaining the recorded tests, but this is not a complete reproduction archive.
Some drivers rely on temporary shared-library paths, helper modules, captured tensors and rebuilt CUDA libraries, so they need adaptation before use elsewhere. Follow the installation guide to run the model. Code paths shown beside snippets identify their origin; raw benchmark reports and source archives are not published as website links.
The records distinguish completed experiments from future ideas and superseded screens. Where an original report mixes component timings with serving or task-quality results, this reference keeps only the component evidence. In particular, tests 36 and 39 do not by themselves establish the final generation rates or task scores; those are reported separately in chapter 9.