Skip to content

Repository files navigation

crucible

CI

Training small language models from scratch and building a custom inference engine, targeting NVIDIA Blackwell (sm_120) hardware.

Status: 120M model trained (1.44B tokens, val loss 3.28). Rust engine decodes at ~1460 tok/s (1.7x llama.cpp at batch 1, token-identical greedy output) and prefills at ~39,700 tok/s at seq 512 on a hand-written tensor-core GEMM — 2.4x the scalar kernel it replaced. Remaining gap to llama.cpp on prefill is 2.8x, down from 7.2x. A paged KV cache and a continuous-batching scheduler now keep multiple requests resident and decode them together, graph-replayed: 8,366 tok/s aggregate at batch 16 and 1,618 tok/s at batch 1, against ~1,640 for one single-request stream. 93-98% of a decode step is now GPU kernel execution. An HTTP service with SSE streaming exposes it, and a Ratatui terminal client uses it: llm-engine serve then llm-engine tui.

Why Blackwell specifically

sm_120 is recent silicon, and the surrounding software stack is still maturing — llama.cpp, vLLM, and Triton all have less-optimised Blackwell paths than they do for Ampere or Ada. That makes it a useful target: FP8 tensor-core work here is measurable against baselines that have not yet been tuned to death.

Hardware

GPU NVIDIA RTX PRO 4000 Blackwell Laptop, 16 GB GDDR7 ECC
Compute capability sm_120 (12.0)
Driver 596.86 (CUDA 13.2 capable)
CPU Intel Core Ultra 9 285HX
RAM 64 GB DDR5
Host OS Windows 11 + WSL2 (Ubuntu 26.04 LTS)
Power limit 134 W (requires Dell Optimizer "Ultra Performance")

Setup

bash scripts/setup_wsl.sh

All GPU work happens inside WSL2. This is a deliberate choice, not convenience:

  • Triton and the CUDA toolchain are first-class here. (An earlier version of this list also cited vLLM's lack of Windows support; that reason was withdrawn once vLLM turned out not to run under WSL2 either — see the comparison section.)
  • Triton — which torch.compile uses to generate fused kernels — is first-class on Linux and unreliable on Windows.
  • CUDA C++ on Windows requires the MSVC toolchain; on Linux gcc suffices.

The one place native Windows would be better is profiling: Nsight Compute can hit restrictions collecting some hardware counters under WSL2 virtualisation.

Environment traps this handles

  1. Ubuntu 26.04 ships only Python 3.14, which PyTorch does not yet publish wheels for. The setup script installs a standalone Python 3.12 via uv rather than fighting the system interpreter.
  2. Never install an NVIDIA driver inside WSL. The Windows driver already exposes the GPU through /dev/dxg. The script installs the CUDA toolkit only; installing a driver in the guest breaks passthrough.
  3. sm_120 requires CUDA ≥ 12.8. Earlier toolkits compile but fall back to PTX JIT, which is slower and easy to miss. verify_gpu.py checks the arch list explicitly rather than trusting that CUDA "works".

Measured hardware baseline

PyTorch 2.13.0+cu130 reports arch list sm_75 sm_80 sm_86 sm_90 sm_100 sm_120, so sm_120 kernels are native rather than PTX-JIT'd.

Median of 15 trials × 50 iterations, 4096³ matmul, via scripts/bench.py:

Precision TFLOP/s spread
FP32 17.6 21.7%
TF32 50.2 8.6%
FP16 78.8 22.2%
BF16 76.4 23.2%
FP8 e4m3 179.7 13.6%
FP8 e4m3 (fast accum) 186.0 4.6%

What these numbers support

  • FP8 is worth targeting — roughly 2.4× BF16, the headline path for the inference engine.
  • BF16 and FP16 are equivalent in speed. An early single-shot run suggested BF16 was ~20% faster; repeated measurement showed that was an artifact. BF16 remains the training default for numerical range, not throughput.
  • FP8 requires torch._scaled_mm. Plain a @ b raises NotImplementedError for FP8 dtypes regardless of hardware support, and the right-hand operand must be column-major.
  • fast_accum is not established as faster. It led in two consecutive runs but by less than the noise floor. Recorded as suggestive, not as a result.

Measurement caveats

The GPU idles at ~700 MHz and ramps to ~1700 MHz under load (49 °C → 67 °C, 24 W → 134 W within a single run). Workloads execute in fixed order, so the first one is measured coldest, which inflates every ratio computed against it. Reported FP32 speedups reached 11.1× but are realistically closer to 10×.

bench.py now warms the GPU before the first workload; the table above predates that fix and is kept for the caveat it documents. Absolute per-precision medians are unaffected by ordering — only cross-workload ratios are.

Training

# tokenise a corpus (streams from the Hub, writes uint16 shards)
python data/prepare.py --tokens 1e8 --out data/tiny

# train
python train.py --preset 30m --data data/tiny --steps 2000
Preset Total params Non-embedding Embedding share
30m 28.5 M 9.1 M 68%
120m 113.0 M 74.4 M 34%
350m 317.4 M 265.9 M 16%

Parameter counts are reported non-embedding by default, following scaling-law convention. This matters at small scale: with a 50k vocabulary the embedding table is 19 M parameters, so two-thirds of the "30M" preset is embeddings. Comparing presets on total parameters would mostly compare embedding tables.

Choosing a model size for a fixed GPU budget

Training the largest model that fits is the wrong instinct on one GPU. Measured throughput per preset (scripts/probe_batch.py) against an 8-hour budget:

preset non-embedding tok/s 8h tokens tokens/param vs Chinchilla
30m 9.4M 250,000 2.00B 211.8 1059%
120m 74.3M 50,827 1.46B 19.7 98%
350m 265.9M 17,995 0.52B 1.9 10%

At 8 hours the 350M preset reaches under 2 tokens per parameter against the ~20 that compute-optimal training calls for. It is so far undertrained that the 120M, trained properly on the same GPU-hours, should reach a lower loss. scripts/compute_budget.py reproduces this table from measured throughput.

Trained model

The 120M preset, trained on 1.44B tokens (22,000 steps × 65,536):

30M (ablation control) 120M
non-embedding params 9.1 M 74.3 M
tokens seen 0.13 B 1.44 B
tokens / param 14 19.4
best val loss 4.3733 3.2839

A 1.09 improvement in validation loss, and the first run in this repo where GQA is genuinely GQA (n_kv_head=3) rather than the degenerate MQA the 30M sweep silently used.

VRAM: this platform does not raise OutOfMemoryError

Probing batch sizes at the 350M preset:

 batch   ms/step   peak GB      tok/s
     4     227.6     12.61     17,995
     6     419.4     16.57     14,651
     8    2778.5     20.61      2,948
    12   44883.2     28.57        274

Peak allocation runs past the card's 16.3 GB to 28 GB without failing, because the WSL2 driver spills to host RAM over PCIe rather than aborting. Batch 12 is 200× slower than batch 4 and still reports success. Any batch-size search relying on torch.cuda.OutOfMemoryError will therefore recommend a configuration that technically runs and is unusable. probe_batch.py detects the throughput cliff and the VRAM overshoot directly instead.

Throughput

Measured against the 76.4 TFLOP/s BF16 figure above rather than a datasheet number:

preset MFU tokens/s ms/step
30m 27% 250,000 282
120m 19% → 52% 26,000 → 70,600 2,520 → 929

The 120M figure is two numbers because the run had two regimes. Steps 0–16,000 held ~2,520 ms/step; from step ~16,000 to the end it held ~929 ms/step. A clean 2.7× shift, stable on both sides, with nothing in the training code changing at that point.

The cause is almost certainly the laptop power profile being switched mid-run — this machine toggles between a ~55 W eco cap and a ~135 W performance cap, and 2.7× is the right magnitude for that change. It cannot be confirmed from the log, because train.py recorded step time and MFU but never recorded the power limit. That omission is what made a one-line explanation look like a mystery, and it is why every GPU benchmark in this repo now prints enforced.power.limit alongside its result.

The consequence for the schedule was large: 12.5 hours instead of the 6.1 projected, entirely from the slow regime. Two lessons are now baked into the repo: throughput must be monitored during long runs rather than sampled at the start, and a projection built from the first 20 steps is worthless.

At the fast regime, 52% against 27% at 30M confirms the low figure at 30M is a small-model artifact — kernel launch overhead dominating a model too small to saturate the GPU — not a defect in the training loop.

Replacing the naive data loader with a vectorised gather plus background prefetch thread changed 30M MFU by 0.3 points — nothing, which is what ruled the input pipeline out as the cause. The prefetch loader is kept because it is not worse and matters at larger batch sizes, but it was not the fix.

Architecture ablations

python sweep.py --steps 2000 --seeds 3      # run the sweep
python sweep.py --report                    # aggregate results

Five axes, each varied independently against a fixed control (gqa / rope / swiglu / pre-rmsnorm):

Axis Variants
attention mha, gqa, mqa
pos_encoding rope, alibi, learned, none
activation swiglu, gelu
norm rmsnorm, layernorm
norm_placement pre, post

SwiGLU's hidden dimension is scaled to 2/3 × 4d so its parameter count stays comparable to a GeLU MLP — otherwise the comparison would measure extra parameters rather than the activation function.

Methodology

Every configuration runs across multiple seeds, and results are classified into three outcomes rather than two:

  1. Diverged. A seed whose best validation loss exceeds the control by more than 1.0 has failed to train, not merely done worse. It is excluded from the mean and SD, and its variant is flagged UNSTABLE regardless of how the surviving seeds did.
  2. A difference, if |Δ| exceeds 2× the larger seed-to-seed standard deviation.
  3. Within noise otherwise. A single-seed result is unknown, never a finding.

Step 1 is not hypothetical — it was added because the first version got norm_placement=post wrong. Averaging its diverged seed in produced an SD of 1.78 against the control's 0.016, and that inflated SD then swallowed the very difference it should have exposed, reporting a training failure as "within noise".

The two-sigma rule has likewise already overturned a result here: a single-shot benchmark showed BF16 outperforming FP16 by 20%, which repeated measurement revealed to be pure noise. Architecture differences at this scale are small enough that the same trap applies throughout.

Each run executes in a separate process so CUDA state, compilation caches, and RNG do not leak between configurations. Completed runs are skipped, so a sweep is resumable.

Results

27 runs (9 configs × 3 seeds), 2000 steps each, 5.6 h on one GPU. Full table in RESULTS.md; figures in figures/.

Control gqa/rope/swiglu/pre-rmsnorm: 4.3733 ± 0.0157 best validation loss.

Axis Variant Δ vs control Verdict
pos_encoding none +0.2579 worse
pos_encoding learned +0.1843 worse
pos_encoding alibi −0.0188 within noise
attention mha −0.0031 within noise
attention mqa −0.0070 within noise
activation gelu −0.0005 within noise
norm layernorm −0.0049 within noise
norm_placement post +0.1821 UNSTABLE — 1/3 seeds diverged

At this scale, positional encoding is the only architecture choice that measurably affects loss. Activation and norm type sit inside seed noise, so they can be chosen on efficiency grounds rather than quality — swiglu vs gelu differs by 0.0005 with parameter count held constant.

Post-norm diverged on one seed of three. A single-seed experiment would more likely than not have reported it as merely "somewhat worse" and missed the instability entirely.

Correction — the attention axis in this sweep is partly void. GPTConfig resolved gqa via max(1, n_head // 4), which collapses to n_kv_head = 1 for any n_head <= 7. The 30M preset has 6 heads, so the gqa control and the mqa variant were the same architecture, and their null result is vacuous rather than informative. The row labelled control throughout this sweep is therefore MQA, not GQA.

What survives: mha (6 KV heads) and mqa (1 KV head) were genuinely different configurations, and the difference between them was inside seed noise. So the real finding is that cutting from 6 KV heads to 1 costs nothing measurable in loss at this scale — a 6× KV-cache reduction, not the 12× originally claimed here, which came from the default 12-head config rather than the preset actually trained.

Every other axis is unaffected: all variants shared the same attention baseline, so those comparisons remain valid. GQA proper is untested at 30M. Fixed in GPTConfig.__post_init__ (floor of 2, plus a hard failure when GQA would degenerate) and pinned by test_gqa_never_degenerates.

Throughput observations, and why the MFU column proved unusable

Three effects make MFU unsafe to compare across configurations:

  1. torch.compile is bistable. Two layernorm seeds ran at 759 ms/step and the third at 296 ms — identical code and config, 2.6× apart, because compilation landed on a different kernel schedule. Loss was unaffected (SD 0.0022, the tightest of any variant).
  2. Sustained thermal drift. Runs slowed from 9.9 to ~20 min over 5.6 hours.
  3. Within-run consistency is not evidence. The sweep showed MHA at 270–282 ms against GQA at 287–288 ms, consistent across all three seeds, which looked like a real ~5% effect and an obvious culprit: repeat_interleave materialising the KV-head expansion.

Point 3 was wrong, and scripts/bench_attention.py was written to check it:

variant                    ms/step   spread   peak GB
MHA                           72.6     3.9%     6.78
GQA repeat_interleave         70.3     2.3%     6.76
GQA fused                     72.6     3.4%     6.64
MQA repeat_interleave         71.3     7.5%     6.76
MQA fused                     72.0     9.8%     6.64

Total range across variants is 3.2% against a 9.8% noise floor. No attention implementation is faster than another at this size. The apparent MHA advantage was a compiled-schedule artifact that happened to be stable within one sweep — the same failure mode as the LayerNorm bistability, and precisely why seed-consistency alone cannot establish a throughput result.

The fused-GQA path (enable_gqa) is kept anyway, on the one measurable difference: 6.64 GB vs 6.76 GB peak memory, in the expected direction for both GQA and MQA. That gap grows with KV-cache length, which matters for inference rather than training. tests/test_attention.py pins it to produce outputs identical to the materialised path.

Known limitation

At 2000 steps the 30M preset consumes 131 M tokens against a 90 M-token corpus, so data repeats. Conclusions transfer to larger scale only loosely. Expanding the corpus (prepare.py --tokens 5e8) removes this.

Inference engine (Rust)

cd engine && cargo build --release
./target/release/llm-engine inspect ../export/120m
./target/release/llm-engine logits ../export/120m --tokens 464,2159,318,1719

Checkpoints cross the Python/Rust boundary as safetensors (export.py), which is flat, zero-copy and memory-mappable — unlike .pt, which is a Python pickle. config.json carries the architecture, because head counts, norm type and positional encoding all change the compute graph and the engine refuses to infer them from tensor shapes.

Validating the forward pass

The CPU path is written to be obviously correct rather than fast: scalar kernels, f64 accumulation, no SIMD or threading. It exists so that when a GPU kernel disagrees with it, the bug is known to be in the GPU kernel.

Correctness is established by comparing logits against PyTorch on identical input, not by inspecting output text:

Rust PyTorch
sum over vocabulary −334002.5076 −334002.4851
min −12.244563 −12.244562
max 7.800590 7.800590

Top-10 token ids match in order, values to ~6 decimal places. The remaining difference is f32 summation order across 50,304 logits — 6.7e-8 relative.

This matters more than it might look. A transposed projection, the other RoPE pairing convention, or query heads mapped to the wrong KV heads under GQA all produce a model that loads, runs, and generates fluent-looking text. None of them announce themselves, and none would be caught by reading samples. scripts/reference_logits.py regenerates the PyTorch side.

Current CPU throughput: 445 ms to load, 261 ms for a 4-token forward pass — single-threaded scalar code, and the baseline every optimisation is measured against.

Prefill and decode are different problems

The engine has two paths, and the reason is arithmetic rather than tidiness.

Decode has one token in flight. Every matmul is a matrix-vector product, no weight is reused, and the work is bandwidth-bound — which is why int8 and CUDA graphs are what move it.

Prefill has the whole prompt at once with no sequential dependency between its tokens, so one weight matrix serves every row. That is matrix-matrix: compute-bound, and it wants tiling and tensor cores instead.

Using the decode path for prefill, as this engine did at first, applies bandwidth-bound tooling to a compute-bound problem and pays ~150 kernel launches per prompt token. GpuModel::forward now routes any multi-token input to prefill; CRUCIBLE_PREFILL=serial forces the old path, which is how the two were compared.

Both paths write into one KV cache layout, so a prompt can be prefilled and then extended token by token with no conversion between them.

Tokenizer

GPT-2 BPE, reimplemented in Rust. The vocabulary is exported from the same tiktoken encoding the training data was built with (scripts/export_tokenizer.py), in a trivial binary format so the engine needs no JSON or base64 dependency.

python scripts/export_tokenizer.py --out export/gpt2.tok
./target/release/llm-engine tokenize export/gpt2.tok "The World is a stage"

The pre-tokenizer pattern needs fancy-regex rather than the standard regex crate, because \s+(?!\S) is a negative lookahead and regex has no lookaround support.

A bug worth recording. The first merge loop removed the merged element before recomputing neighbouring ranks, so lookahead indexed past the wrong boundary and merging stopped early: "The" encoded as [817, 68] ("Th", "e") instead of [464], and the probe string produced 29 tokens instead of 14. Decoding round-tripped perfectly the whole time — the text came back byte-identical, so nothing looked wrong. The only symptom would have been a model fed ids it was never trained on, which presents as degraded output and reads like a bad model rather than a bad tokenizer.

It was caught by comparing ids against tiktoken, and matches_tiktoken_ids now pins them.

Generation

./target/release/llm-engine generate export/120m \
    --tokenizer export/gpt2.tok \
    --prompt "The capital of France is" --max-tokens 30 --temperature 0.7

Output from the 120M model (val loss 3.28):

The capital of France is in the possession of the royal treasury. The French government is in the possession of the French and French governments of the entire territory of the world.

Photosynthesis is the process by which plants convert light into energy. Plants can use photosynthesis to generate heat. The resulting heat is used to power the plant's turbines

The three branches of government are the three branches of government, the executive branch and the judiciary, and the judiciary.

Fluent and locally coherent, with the failure modes expected at this scale: factual drift, repetition, and word-sense confusion (plant the organism versus plant the factory).

Sampling is top-k with temperature, seeded by a small xorshift64* so runs reproduce exactly without a rand dependency.

KV cache

Without a cache, generating token N re-runs attention over all N positions, so a sequence costs O(N²) and the entire prompt is recomputed every step. The cache retains each position's projected keys and values, making each new token O(N).

Measured on the same prompt and seed, 30 tokens:

no cache with cache
decode time 35.22 s 3.10 s
throughput 0.85 tok/s 9.66 tok/s
per token 1,174 ms 103 ms

11.4× faster, byte-identical output. Identical output is the check that matters: a broken cache still generates fluent text, just different text, so matching the uncached path token-for-token under a fixed seed is what proves it correct.

The advantage grows with length, since only the uncached path is quadratic. At 150 tokens throughput holds at 9.42 tok/s — essentially flat — while the uncached path would need roughly 25× its 30-token time. That is ~55× at 150 tokens, widening further from there.

This required restructuring the forward pass from layer-major to token-major: each token now flows through every layer before the next token begins. By the time position p reaches layer L, every earlier position has already written its layer-L keys and values, so attention reads them instead of recomputing. The logits still match PyTorch exactly (sum -334002.5076), which is how the restructure was verified.

Cache layout is [layer][position][kv_head * head_dim], contiguous in the last dimension — attention reads one position's keys for one head at a time, so those values sit adjacent in memory. At 1024 context the 120M model's cache is 18.9 MB; with MHA instead of GQA it would be 75.5 MB.

CUDA kernels

cargo build --release --features cuda
./target/release/llm-engine gpu-validate
./target/release/llm-engine gpu-bench --rows 8192 --cols 4096

Kernels are compiled at runtime with NVRTC, not offline with nvcc. Two reasons, one forced and one earned:

  • CUDA 13's headers conflict with glibc 2.43 on Ubuntu 26.04 — both declare rsqrt/rsqrtf with incompatible exception specifications, and nvcc injects host headers even for --ptx, so every compile fails including an empty kernel. No compiler flag fixes it; __GLIBC_USE(IEC_60559_FUNCS_EXT_C23) is not overridable. NVRTC never includes host headers, so the conflict cannot arise.
  • As a side effect the engine builds without the CUDA toolkit present, and targets the exact GPU at runtime.

Validation

Every kernel has a scalar CPU twin in ops.rs, and gpu-validate compares them rather than assuming they agree:

kernel max relative difference
gemv 0 (bit-exact)
rmsnorm 1.1e-7
softmax 3.7e-7
silu_mul 1.8e-7
rope 9.6e-7
gemm (scalar, f32 / int8) 3.3e-7 / 3.3e-7
gemm (tensor core, f32 / int8) 6.4e-5 / 3.4e-5

Exact equality is not the bar — the GPU reduces in a different order and use_fast_math trades accuracy for speed. Rounding-level agreement is.

The tensor-core rows are held to a looser bound than the scalar ones because they convert activations to half, but a looser bound is not no bound: half carries 11 mantissa bits, so with K=768 accumulating in f32 the error must stay under 2^-11 ≈ 4.9e-4. Measured at 6.4e-5, comfortably inside. Anything above 1e-2 would mean the tiling or the fragment layout is wrong, not that half is imprecise.

Two ways this test passed for the wrong reason

The generators feeding gpu-validate are not arbitrary, and both properties were learned by getting them wrong.

The first version used (i % 89 - 44) / 128 — small integers over a power of two. half represents those exactly, so the tensor-core path and the scalar path agreed bit-for-bit and the test reported 0.000e0. A kernel that dropped precision catastrophically would have passed identically. Test data has to be mantissa-dense before a precision test means anything.

The fix after that used raw sin/cos, which is mantissa-dense but signed. A 768-term dot product of random signs cancels down to near zero, and relative error against a near-zero result is meaningless — it reported ~2.0, which is what a sign flip on noise looks like, not a broken kernel. Shifting both operands positive keeps the mantissas full while making the sum accumulate monotonically.

Both failures produced a confident-looking number. The first said the kernel was perfect, the second said it was broken, and the kernel was the same kernel.

Why decode is bandwidth-bound

Generating one token at a time makes every matmul a matrix-vector product: weights are read once and reused for a single output element. Measured with scripts/bench_bandwidth.py at a 135 W enforced limit:

757 GB/s median (spread 13.1%, best 759). Against the 120M model's weights:

precision GB/token ceiling
f32 0.45 1,674 tok/s
f16 / bf16 0.23 3,348 tok/s
int8 0.11 6,696 tok/s

Two consequences:

FP8 tensor cores are close to irrelevant for single-stream decode. The 186 TFLOP/s figure only pays off in batched prefill, where weights are reused across many tokens. For decode, quantisation helps because it moves fewer bytes — not because it multiplies faster.

The GEMV kernel is already done. It reaches 723–770 GB/s, at the measured device ceiling. No amount of kernel tuning raises decode throughput further in f32; only reading fewer bytes does.

Comparing kernels requires paired trials

float4 vectorised loads were expected to beat the scalar kernel. Measuring that took three attempts, and the first two were the interesting ones.

The first design ran all scalar trials, then all float4 trials, and reported a 120% difference against a 168% spread. The second, after correcting the power profile, still showed 76% spread. Both were unusable, for the same structural reason: running kernel A to completion then kernel B cannot separate a real difference from clock drift between the two phases, and on a laptop the clocks move far more than any kernel effect.

The fix is a paired design — both kernels run back to back inside each trial, with the order alternating between trials, and the reported statistic is the per-trial ratio. Drift then affects both kernels equally and cancels:

  scalar      723 GB/s   [462-745]  spread 39.1%
  float4      738 GB/s   [492-770]  spread 37.6%

  paired ratio float4/scalar: 1.034   [0.910-1.128]  spread 21.1%
  -> 3.4% difference against 21.1% paired spread: not distinguishable

Absolute throughput still swings ~39%, but the paired comparison narrows to 21% and gives a usable answer: float4 is not faster. The scalar kernel already saturates memory bandwidth, so there is nothing for wider loads to recover.

Every measurement must record its power envelope

Getting a stable number took longer than getting a fast one. An early run spanned 167–505 GB/s, a 168% spread, with clocks sampled at 180–300 MHz during 100% utilisation while boosting to 1867 MHz when idle:

SM MHz  temp  watts  util
  1095    46   63.0   100
   180    45   45.4   100     <- 180 MHz at full utilisation
  1867    44   25.0     0     <- 1867 MHz once idle

That looked like a clock-governor fault and was written up here as one. It was not. The laptop was in a 55 W eco power profile, and the GPU was doing exactly what it had been told. Temperature never exceeded 46 °C because the cap was the binding constraint, not heat.

The lesson is methodological, and it now applies to every GPU number in this repo: this machine's power limit is user-switchable between roughly 55 W and 175 W, so a measurement that does not record its envelope is not reproducible. gpu-bench and gpu-validate now print enforced.power.limit and the maximum SM clock alongside every result, and gpu-bench warns when spread exceeds 30% — the signature of a throttled or capped run.

The size of the error is worth stating: the eco-profile baseline measured 479 GB/s, against 757 GB/s at 135 W. Every derived figure — the decode ceiling, the value of quantisation, how close the GEMV kernel sits to peak — was computed from a number that was 37% low.

GPU forward pass

cargo build --release --features cuda
./target/release/llm-engine gpu-logits export/120m --decode 64
./target/release/llm-engine generate export/120m --tokenizer export/gpt2.tok     --prompt "The capital of France is" --gpu

Weights and the KV cache stay resident on the device; only the token id goes in and the logits come out. Kernels are queued without synchronising between them, so the host does not stall on each of the ~170 launches a token requires.

CPU (scalar reference) GPU
decode 9.66 tok/s 471–501 tok/s
per token 103 ms 2.00 ms
resident 18.9 MB cache 485 MB weights + cache

~49× faster, and byte-identical output under the same seed — which is the check that matters, since a wrong kernel still produces fluent text. Against the CPU reference, logits agree to 2.0e-4 maximum relative difference with the top-10 identical in order; the residual comes from use_fast_math (__expf, rsqrtf) and f32 rather than f64 accumulation.

Attention is fused into a single kernel per head — scores, softmax and the weighted sum of values, with scores held in dynamic shared memory. Splitting those into three kernels would mean 432 launches per token at 12 heads × 12 layers, and launch overhead alone would dominate a model this size.

What limits it now

2.00 ms/token moves 0.45 GB of weights, an effective 225 GB/s against the measured 757 GB/s ceiling — about 30%. The gap is not the kernels: GEMV in isolation reaches 723–770 GB/s. It is that ~170 separate launches per token leave the GPU idle between small kernels.

Two levers, in order of expected value:

  1. Quantisation. int8 weights cut bytes-per-token by 4×, raising the ceiling to ~6,700 tok/s. (Implemented — and it delivered 4× on memory but only 1.06× on speed. See the int8 section below: the ceiling rose, but the engine was never near it.)
  2. CUDA graphs, capturing the per-token launch sequence once and replaying it, to remove per-launch overhead. (Implemented — 1.39x on its own, and it raised int8's contribution from 1.05x to 1.21x. See the CUDA graphs section below.)

int8 quantisation

./target/release/llm-engine gpu-eval export/120m --data data/fineweb-2b/val.bin     --tokens 1024 --quant f32,int8

Weight-only, symmetric, per-output-row scales. Activations stay f32 — they are a negligible share of the bytes moved during decode, so quantising them would cost accuracy for nothing. Per-row rather than per-tensor because one scale across a whole matrix is set by its largest outlier, crushing the resolution of every other row.

Measured on 1024 held-out tokens:

f32 int8
cross-entropy 3.720299 3.720334
perplexity 41.2767 41.2782
weights resident 452 MB 114 MB
decode 405 tok/s 431 tok/s

Quality cost is +0.001% cross-entropy — free, within any reasonable tolerance. Memory drops 4×. But speed rises only 1.06×, not the 4× the bandwidth argument predicted.

That gap is the interesting part, and it corrects an earlier claim in this README. int8 moves 0.11 GB per token, which at 757 GB/s is 0.15 ms — yet a token takes 2.32 ms. Roughly 2.2 ms is fixed overhead, about 13 µs across the ~170 kernel launches a token requires. Decode at this model size is launch-bound, not bandwidth-bound, so removing bytes barely moves the wall clock.

The bandwidth ceiling reasoning was not wrong, it was premature: the ceiling did rise from ~1,674 to ~6,700 tok/s, but the engine sits at 431, nowhere near either. Quantisation cashes in only once per-token overhead is gone.

Revised priority. CUDA graphs — capturing the per-token launch sequence once and replaying it — now comes before further quantisation work, because it is what makes quantisation pay. int8 is worth keeping regardless for the 4× memory reduction, which is what determines how large a model fits in 16 GB.

CUDA graphs

./target/release/llm-engine gpu-logits export/120m --quant int8 --decode 256 --graph
./target/release/llm-engine generate export/120m --tokenizer export/gpt2.tok \
    --prompt "The capital of France is" --graph

Decoding one token issues ~170 kernel launches, each costing microseconds of driver work. A CUDA graph captures that sequence once and replays it as a single submission.

eager graph graph gain
f32 509 tok/s 707 tok/s 1.39x
int8 533 tok/s 852 tok/s 1.60x
int8 gain 1.05x 1.21x

Output is byte-identical to eager and to the CPU reference under the same seed, and cross-entropy is unchanged to six decimals.

The two optimisations compose, and one enables the other. int8 alone was worth 1.05x; once graphs removed the launch overhead it became 1.21x, because only then was there enough bandwidth-bound time left for fewer bytes to matter. Combined: 509 to 852 tok/s, and 88x over the 9.66 tok/s CPU reference.

Making capture work

Two failures, neither obvious from the error text, both worth recording.

A captured graph freezes its kernel arguments. token, pos, seq_len and the KV cache slot offset all change every step, so passing them by value would bake step 0 into the graph and every replay would recompute the same token. Every per-step scalar therefore lives in a small device buffer that kernels index into, updated with one host-to-device copy of 5 ints per token -- one transfer replacing ~170 launches. Slot 0 holds a permanent zero so call sites needing a constant offset take the same code path instead of requiring a second kernel variant. Shared memory for attention is sized for maximum context rather than current length, for the same reason: a graph fixes its allocation at capture.

CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED. CUDA forbids capturing the legacy default stream, which is what ctx.default_stream() returns. Fixed by creating a dedicated stream.

CUDA_ERROR_STREAM_CAPTURE_ISOLATION. cudarc records a CUDA event per buffer and inserts cuStreamWaitEvent on every kernel touching it, to order work across streams. A captured launch waiting on an event recorded by uncaptured work is exactly what capture isolation forbids -- so the safety mechanism made capture impossible. Creating a dedicated stream is also what activated that tracking, since cudarc only engages it in multi-stream mode.

Resolved with disable_event_tracking(), which is sound here because the engine uses a single stream and issues everything in program order, so the stream itself provides the ordering those events exist to guarantee. It must be called before any allocation: the flag is read when a buffer is created, so buffers allocated earlier keep their events and still poison capture.

What limits it now

1.17 ms/token at int8 moves 0.11 GB of weights, which at 757 GB/s is 0.15 ms. So roughly 1.0 ms per token is still unaccounted for, and decode sits at about 13% of the bandwidth ceiling.

Launch overhead is no longer the main suspect -- graphs removed most of it. The candidates were:

  • The attention kernel's inner loop is uncoalesced. (Confirmed by profiling and fixed -- see below. It was the largest stage at 32%.)
  • Small-matrix inefficiency. (Confirmed and partly fixed: warp-per-row GEMV for int8.)
  • Logits transfer. (Measured at 4.6% of decode -- real but minor, and still unaddressed.)

Profiling, and getting the profiler right

./target/release/llm-engine gpu-profile export/120m --quant int8 --warm 256

Nsight Systems reports no CUDA kernel data under WSL2 virtualisation, so attribution is done in-engine: the stream is synchronised between stages and each is timed on the host.

The first version of that produced a nonsense ranking. A 768-element rmsnorm appeared to cost more than a 50304x768 matmul, because rmsnorm runs 24 times per token and lm_head once, so it absorbed 24 syncs to the other's one. The profiler was measuring itself.

The correction needs a per-block overhead estimate. Timing syncs on an idle stream gave 70.8 us, which cannot be right: 111 blocks would then cost 7.9 ms against a 3.9 ms measured total. Syncing an idle stream is simply a different operation from syncing after queued work. The estimate now comes from the cheapest stage that still launches a kernel -- embed, one 768-element row copy -- giving ~21 us and an adjusted total of 1.37 ms against 0.87 ms of real decode, which is close enough for ranking.

stage adjusted ms share
mlp 0.390 28.5%
attention 0.362 26.4%
lm_head 0.169 12.3%
qkv_proj 0.168 12.2%
rope 0.117 8.5%
logits_copy 0.062 4.6%
rmsnorm 0.060 4.4%
o_proj 0.035 2.6%
residual 0.006 0.4%

Two kernel fixes the profile found

Attention: uncoalesced reads. The original scoring loop gave each thread a cached position and walked that key vector sequentially, so neighbouring threads read addresses cache_stride floats apart -- 768 bytes here -- and every lane issued its own memory transaction. Rewritten to one warp per position with lanes striding across the key, reads become contiguous and coalesce. The value accumulation was already coalesced but left three quarters of each block idle at head_dim 64; each warp now sums a slice of positions into its own partial vector in shared memory.

Attention fell from 0.531 to 0.362 ms, and decode went 852 to 1021 tok/s.

GEMV: starved blocks. A block per output row suits a long row and wastes a short one. gate_proj is 2048x768, which as int8 with char4 loads is 192 elements against 256 threads: a quarter of the block idle, one load per active thread, then a full eight-warp block reduction to combine them. One warp per row removes the idle threads and replaces the block reduction with shuffles alone.

int8 measured 1170-1299 tok/s against 852, a 1.42x gain well outside the spread. The same change did not help f32 -- 737/736/864 against 825, a possible regression sitting inside its own 17% spread. Unproven in either direction, so f32 keeps block-per-row and the switch is applied only where it was measured to help. An f32 row carries four times the bytes, so its block is far less starved.

Where decode stands

path tok/s ms/token
CPU reference (scalar) 9.7 103
GPU eager, f32 509 1.96
GPU graph, f32 825 1.21
GPU graph, int8 ~1149 0.87

119x the CPU reference, with cross-entropy unchanged at every step (3.720334 for int8, before and after both kernel rewrites).

Still well under the bandwidth ceiling. The MLP was the largest stage and its three projections are most of the model's weights, so the next gain looked like it needed fusing the per-layer sequence rather than tuning kernels further -- which is what happened next.

Kernel fusion

Graphs removed the CPU cost of launching, but each kernel still pays GPU-side dispatch, so kernel count keeps mattering. Three fusions, all int8:

  • SwiGLU: silu(gate . x) * (up . x) in one kernel instead of three, with no hidden-sized intermediates. MLP time halved, 0.349 to 0.174 ms.
  • Residual into projection: o_proj and down_proj accumulate straight into the residual stream, removing the separate add_inplace at both sites.

Decode went 1149 to 1320 tok/s (median of five), 0.76 ms/token. Cross-entropy unchanged at 3.720334.

The fusion shipped broken first, and speed alone would have passed it. gemv_i8_at routes to warp-per-row only when cols/4 < 256; down_proj has 2048 columns, so it took the block-per-row kernel, which had no accumulate parameter. It overwrote the residual stream instead of adding to it. The engine ran faster while doing this, generated fluent text, and perplexity went from 41.28 to 64,300. Only the held-out cross-entropy check caught it. The f32 path now fails loudly rather than silently overwriting.

The profiler broke in the same way, more quietly: profile_step duplicates the forward pass so it can sync between stages, and only queue_token was fused -- so it spent a round reporting a residual stage that no longer existed. Both now mirror each other, and the duplication is called out in the code as something that must be kept in step.

Where decode stands

path tok/s ms/token
CPU reference (scalar) 9.7 103
GPU eager, f32 509 1.96
GPU graph, f32 825 1.21
GPU graph, int8 1149 0.87
GPU graph, int8, fused 1320 0.76

137x the CPU reference, cross-entropy unchanged at every step.

Current stage breakdown (gpu-profile, position ~256):

stage adjusted ms share
attention 0.427 38.8%
mlp 0.174 15.8%
qkv_proj 0.173 15.7%
rmsnorm 0.100 9.1%
rope 0.073 6.7%
lm_head 0.061 5.5%
logits_copy 0.060 5.5%
o_proj 0.031 2.8%

Split-position attention: implemented, exact, and it did not help

Attention dominated at 38.8%, and not because of kernel quality. At position 256 it reads roughly 4.7 MB of KV cache per token, which at 757 GB/s should take about 6 us; it took 427. The kernel launches one block per head -- 12 blocks on a GPU with dozens of SMs, so most of the machine sat idle. Coalescing the reads (an earlier fix) helped, but no per-thread tuning fixes a grid that cannot fill the device.

So the sequence was split across blocks too: grid (n_head, n_chunks), each block reducing one chunk into a partial softmax, with a second kernel rescaling by exp(m_chunk - m_global) and merging -- flash-decoding. n_chunks is fixed at capacity rather than current length, because a captured CUDA graph freezes grid dimensions.

It is exact: cross-entropy is identical to the single-block path to six decimals. It is also not faster.

decode tok/s, int8 + graph 256 tokens 900 tokens
single block per head 1484 1424
split positions 1305 1389

Median of three. Splitting costs a second kernel dispatch per layer, 12 more per token, which at this size is about what the extra parallelism saves. Clearly worse at short context, a wash at long.

Kept behind CRUCIBLE_ATTN=split rather than deleted: the trade should invert with more heads, a larger head_dim, or context well beyond 1024, where attention work grows while the extra dispatch does not. It is off by default because on this model it loses.

A limit of the profiler, found the hard way

The profiler said attention had halved -- 0.427 to 0.206 ms adjusted -- while end-to-end decode did not move. Both cannot be true.

The profiler subtracts one launch-plus-sync overhead per timed block. The split path runs two kernels inside that one block, so it was credited with one subtraction where it should have had two, and looked better than it was. The comparison it was built for -- ranking stages within one configuration -- is still valid. Comparing configurations whose stages contain different numbers of kernels is not something it can do.

This is the third time a measurement here has produced a confident wrong answer: first the sync overhead attributed to whoever called most often, then the idle-stream probe that contradicted its own arithmetic, now this. The pattern is consistent -- the tool is fine for what it was built for and silently wrong just outside it, and only an end-to-end number catches the difference.

Tensor-core GEMM

Prefill multiplies a [seq, n_embd] activation block by every weight matrix, so unlike decode it is compute-bound rather than waiting on memory. The scalar 16x16 tiled kernel sustained ~3.4 TFLOP/s — 4.5% of this GPU's BF16 tensor-core peak — which was the entire remaining gap to llama.cpp.

Replacing it with a wmma kernel:

seq scalar tiled tensor core speedup
128 15,684 18,353 1.17x
256 17,203 28,560 1.66x
512 16,858 39,691 2.35x
1024 15,016 34,900 2.32x

The tensor-core column is the shipped default (9 interleaved trials, 170.91 W enforced / 3090 MHz). The scalar column is from an earlier 5-trial run at 156 W; it is paired here rather than re-measured because the scalar kernel reproduces across runs to within 0.3%, which the tensor-core paths do not.

Throughput falls off past 512 on both paths because attention is O(n^2): its share of the work grows with sequence length while the GEMM's shrinks.

Two design choices carry the int8 path. Weights stay int8 in global memory and convert to half in shared memory — keeping a half copy of the model would cost 226 MB and give back most of what quantisation bought, and int8 -> half is exact, so that conversion loses nothing. Activations are f32 and convert to half on load, which does lose mantissa bits; that one is a real numerical change and is priced below rather than assumed away.

GEMM tuning is complete, and neither remaining technique helped

Two techniques were on the roadmap. Both are now resolved, neither ships.

ldmatrix was already in use. SASS for gemm_i8_wmma contains 4 LDSM.16.M88.4 alongside 4 HMMA.16816.F32, and the big tile 8 and 16 — counts that match the fragment math exactly. wmma::load_matrix_sync lowers to ldmatrix on sm_120 already, so writing it by hand with raw mma.sync PTX would emit the same instructions. This cost nothing to establish and would have cost days to "implement".

Double buffering was implemented and rejected. Prefetching the next K tile into registers while the current tile's mma runs, then staging it into an alternate shared buffer. cp.async is not usable here because the int8 -> half conversion happens during staging and cp.async copies raw bytes.

seq small prod small dbuf big prod big dbuf
128 18,374 19,777 14,983 7,100
256 28,572 31,217 27,022 13,596
512 37,835 35,270 39,249 21,269
1024 30,482 31,175 35,856 26,398

Shared memory doubles by construction, which costs the small tile 8 -> 5 blocks/SM and the big tile 6 -> 4. On the big tile that is fatal. On the small tile the result is genuinely mixed, and reproducibly so: through the production wmma-auto policy it gained 5.9% then 6.9% at seq 256 across two independent runs, and lost 8.3% then 8.7% at seq 512. Seq 512 is both a common prompt length and where this engine performs best, so a reproducible regression there is not a shippable trade.

A sequence-length-dependent rule could capture the short-prompt gain, but that is a dispatch-heuristic change and belongs in its own measured experiment.

Neither tile size wins everywhere

The obvious next move was a larger block tile. A 64x64 tile produces four times the output per block for twice the loads — ~64 FLOP per element loaded against ~25 — so it should win outright. It does not.

9 interleaved trials per point, 170.91 W enforced / 3090 MHz. auto picks per launch, described below:

seq 16x64 64x64 auto
128 18,371 14,971 18,353
256 26,983 25,408 28,560
512 35,483 36,821 39,691
1024 29,453 34,515 34,900

The 64x64 tile gives up ~19% at seq 128 and takes ~17% at seq 1024. The intensity argument is correct and was not the binding constraint: M is the sequence length, so a 64-row tile at seq=128 with n_embd=768 launches ceil(128/64) * ceil(768/64) = 24 blocks on a 60-SM GPU where the 16-row tile launches 96. Arithmetic intensity only starts paying once there is enough work to fill the machine.

Both tiles are generated from one templated device function so they cannot drift apart. The default (wmma-auto) picks per launch, using the big tile only when a launch would still produce at least two blocks per SM — read from the device, not hard-coded. It beats both fixed tiles because the projections in a single forward pass have very different N — 768, 2048, and 50304 for the lm_head — and want different tiles: it ties the small tile at seq 128 and wins by 5.8%, 11.9% and 18.5% at 256/512/1024.

CRUCIBLE_GEMM=wmma-small and wmma-big pin a fixed tile, for benchmarking and debugging; tiled selects the scalar kernel.

Measurement notes, because a 2x claim was wrong here once already. An earlier 5-trial run reported the 64x64 tile at 9,950 tok/s at seq 128 — a 2x deficit rather than the real ~19% — with a 9.9% spread that should have disqualified the median on the spot. Two later runs agree with each other and not with it. The 9-trial numbers above were reproduced by an independent 5-trial run at a different power envelope (148.86 W) to within 1% on every median, which is the check that earlier number never got.

What the speedup costs in accuracy

The tensor-core path converts activations to half, so it is a different computation and not just a faster one. Two measurements, because neither is inferable from the other.

gpu-validate bounds the kernel error against the CPU reference: 6.4e-5 (f32) and 3.4e-5 (int8), against 3.3e-7 for the scalar path. half carries 11 mantissa bits, so with K=768 and an f32 accumulator the error has to stay under 2^-11 ≈ 4.9e-4 — measured well inside that, which makes it rounding rather than a structural defect.

That bounds the kernel but says nothing about the model, and the existing gpu-eval could not help: it fed one token per call, making every matmul a matrix-vector product, so it never executed the GEMM at all and could not have detected a prefill regression. gpu-eval --prefill-ctx N scores each position from a fresh prefill of the preceding N tokens, which puts the batched GEMM on the path that produces the number. Over 511 windows at ctx=256:

scalar tiled tensor core delta
f32 cross-entropy 3.089870 3.089860 -1.0e-5
int8 cross-entropy 3.090417 3.090399 -1.8e-5

The GEMM change moves cross-entropy by ~1e-5, thirty times less than int8 quantisation's own cost of +5.4e-4, and slightly downward — which is noise, not an improvement. Ragged sequence lengths (7, 17, 100, 333) agree to within half rounding, and --prefill 1 is bit-identical because M=1 routes through GEMV, confirming decode is untouched.

The per-launch tile choice adds nothing to this. Both tiles walk K in the same order with the same half conversion, so they are bit-identical, not merely close: wmma-auto reproduces the fixed small tile's logits exactly at every ragged length tested, and lands on the same cross-entropy to six decimals (3.090399). Selecting a tile is a speed decision with no accuracy component, which is why the numbers above did not need re-deriving when it became the default.

Paged KV cache and continuous batching

A contiguous cache makes one decision at load time -- how many tokens a sequence may ever hold -- and charges every sequence that much. Serving several requests from it means reserving the maximum context each (mostly wasted) or recopying the cache whenever a request joins or leaves. Paging removes both.

Layout, and why it was nearly free

Pages are fixed at 16 tokens and hold every layer:

pool: [n_pages][n_layer][PAGE_TOKENS][kv_dim]
offset(page, layer, slot) =
    page * n_layer * PAGE_TOKENS * kv_dim   <- dynamic, from the page table
  + layer * PAGE_TOKENS * kv_dim            <- host constant per launch
  + slot * kv_dim

That split is the reason this was cheap. The layer term is exactly the layer_base argument the projection kernels already took, and the page term is exactly the per-step scalar they already read from device memory. So decode's K/V write and its RoPE needed no kernel changes at all, and CUDA graph capture stays valid, because the only value that varies per step still arrives through the parameter buffer rather than as a kernel argument. Only attention and the prefill scatter needed paged variants.

16 tokens: power of two, so translation is a shift and a mask rather than a division in the attention inner loop. At kv_dim = 192 floats a page is 12 KB of contiguous K per layer, far above the 128-byte transaction granularity, so the coalescing the attention kernel was tuned for is untouched. Internal fragmentation is bounded at 15 tokens per sequence at any context length.

Attention reads the paged representation directly. Gathering pages into a contiguous buffer before attending would work and would defeat the purpose.

Correctness

Paging moves storage, not arithmetic -- the same values are summed in the same order -- so the bar is bit-exact, not close. gpu-paged checks every length that straddles a page boundary, through both the prefill and the decode path, with and without graph replay:

len 1 15 16 17 33 127 128 129 511 512 1023
max diff 0 0 0 0 0 0 0 0 0 0 0

All 64 pages return to the pool on reset.

gpu-batch then checks that a request's output does not depend on who it is batched with, across prompts of 15/16/17/63/64/65/255/511 tokens under both simultaneous and staggered admission. Every request is identical batched and alone, and no pages leak.

A tie-break that looked like a cache bug

The first version of that test compared batched output against the single-request forward path and reported two requests as mismatched. They were not. forward computes its lm_head with a GEMV; the batched path uses a GEMM, and the two sum in different orders for a ~1e-4 difference. Greedy argmax turns that into a different token whenever the top two logits are closer than that, and one different token diverges everything after it.

The evidence is quantitative: the one request that diverged had a closest top-2 logit gap of 6.9e-4, against 1.8e-2 to 1.3e0 for every request that agreed.

The fix was to the test, not the engine. Cross-request isolation is now checked against the same code path with max_batch = 1, where agreement must be exact, and the forward comparison is reported alongside the tie gap rather than asserted.

Scheduler

submit / step / completed. Admission is first-come-first-served, bounded by batch slots and free pages; a request that cannot get pages stays pending, so the pool is backpressure rather than an error surface. Retirement is by swap_remove, which deliberately reorders slots -- every per-request quantity is rebuilt into the metadata arrays each step, so nothing may be tied to a slot index across steps, and this is what would catch it if something were.

Throughput

Heterogeneous prompts (32/128/256/512 cycled), 64 decode steps, 3 trials, 150.31 W enforced:

batch aggregate per-request step ms vs batch 1 pages wasted slots
1 152 t/s 152 t/s 6.56 1.00x 3 15
2 295 t/s 147 t/s 6.78 1.94x 12 30
4 549 t/s 137 t/s 7.28 3.60x 62 60
8 1,055 t/s 132 t/s 7.58 6.92x 124 120
16 1,916 t/s 120 t/s 8.35 12.57x 248 240

Aggregate is total tokens per second; per-request is what one client sees. Batching raises the first and lowers the second. That is a throughput result and not a latency result, and it is worth being precise about which one is being claimed.

Batched GEMV: the projections were the whole problem

Profiling the batched step first, rather than guessing, showed the projections were 83-84% of it -- and that their cost was flat across batch sizes:

stage b1 b4 b16
qkv_proj 2.061 ms 1.828 ms 1.974 ms
down_proj 1.760 ms 1.641 ms 1.767 ms

Sixteen times the work for the same wall time is what an occupancy-starved kernel looks like. At decode M the tiled GEMM launches ceil(768/64) x 1 = 12 blocks on a 60-SM GPU.

The replacement puts one warp on each output row, carrying every request's accumulator at once. The obvious alternative -- one warp per (row, request), grid (rows, batch) -- also restores parallelism but re-reads the whole weight matrix once per request, which at batch 16 is 16x the traffic on the tensor that dominates a bandwidth-bound step. Keeping the batch inside the warp reads each weight row exactly once regardless of batch.

BMAX is a template parameter, not a runtime count: acc[b] indexed by a runtime value is not a register array, it is local memory -- the exact failure the big-tile register investigation documented. Instantiated at 1/2/4/8/16.

Speedup over the tiled GEMM (int8, 200 iterations, median of 3, 158 W):

shape b1 b2 b4 b8 b16
q/o 768x768 4.69x 4.02x 4.49x 3.74x 2.96x
k/v 192x768 3.68x 4.19x 3.73x 2.97x 3.09x
gate/up 2048x768 3.53x 3.50x 3.92x 3.01x 2.41x
down 768x2048 7.87x 10.64x 7.30x 6.35x 5.39x
lm_head 50304x768 7.76x 5.56x 3.57x 1.77x 1.00x

The crossover is shape-dependent because the shapes genuinely differ. The per-layer projections top out at 2048 output rows, so the GEMM is starved and GEMV wins across the range. The lm_head has 50304 rows and 786 blocks, so it is not starved: GEMV wins to batch 8 and reaches parity by 10. The dispatch is therefore rows <= 4096 || batch <= 8.

That lm_head crossover sits on the instantiation boundary rather than anywhere physical -- a batch of 10 runs the BMAX=16 kernel and discards six accumulators -- which is a reason not to read the constant as fundamental.

Residual accumulation is folded back into o_proj and down_proj through the kernel's accumulate flag, the same fusion single-request decode uses, removing 24 kernels per step.

End-to-end through the runtime, GEMV and GEMM trials interleaved:

batch GEMM aggregate GEMV aggregate per-request step ms speedup
1 150 t/s 599 t/s 599 t/s 1.67 4.00x
2 290 t/s 1,137 t/s 568 t/s 1.76 3.92x
4 537 t/s 1,696 t/s 424 t/s 2.36 3.16x
8 1,040 t/s 2,974 t/s 372 t/s 2.69 2.86x
16 1,892 t/s 4,587 t/s 287 t/s 3.49 2.42x

GEMV is also the more accurate path

Worth stating because it was not the goal. The GEMM being replaced is the tensor-core one, which rounds activations to half; GEMV accumulates in f32 throughout. Their outputs differ by up to 8.7e-3 absolute, and most of that is the GEMM's.

The evidence is a test that used to fail. When the batched path used GEMM, one of eight requests diverged from the single-request forward reference on a near-tie -- top-2 logit gap 6.9e-4, against 1.8e-2 to 1.3e0 for the seven that agreed. With GEMV, all eight now match forward exactly, including that one. Switching kernels for speed removed a numerical discrepancy as a side effect.

CUDA graphs for the batched path

After batched GEMV the projections were no longer the bottleneck; issuing 183 kernels eagerly was. Graph replay removes that.

Capture constraints. Everything varying per step falls into three groups. Token ids, positions, sequence lengths and page-table entries already live in persistent device buffers, so they are free -- only their contents change between replays. All scratch, pool and weight addresses are allocated once in enable_paging and never move. That leaves exactly one thing that cannot be absorbed: the active count n, which drives grid dimensions, the batch kernel argument, attention's grid.y, which GEMV BMAX instantiation runs, and the lm_head GEMV/GEMM dispatch.

Two things stay outside the graph: the host-to-device metadata upload (memcpy_htod from a temporary buffer would be captured as a node holding a dangling host pointer) and the final device-to-host logits copy.

Cache keyed by exact batch size, up to max_batch. Bucketing at BMAX would need an active mask threaded through the KV scatter, RoPE and the GEMV store, because padded page-table slots read 0 -- an inactive row would write KV into physical page 0, corrupting whichever request owns it. That is a kernel change plus a safety proof to save eleven graphs that together do not measure on the VRAM counter. Exact keying also preserves the measured GEMV/GEMM dispatch for free, since n is never rounded up to graph capacity.

Nothing captured depends on which request occupies a slot, only on slot position, so the scheduler reordering slots with swap_remove cannot invalidate a graph. Graphs are dropped only when something they baked in could have changed: buffer addresses (enable_paging) or which kernels run (set_force_decode_gemm). CRUCIBLE_BATCH_GRAPH=0 keeps the eager path.

batch eager graph speedup per-request step ms spread
1 606 t/s 1,492 2.46x 1,492 0.67 6.2%
2 1,175 t/s 2,265 1.93x 1,133 0.88 2.6%
4 1,699 t/s 2,568 1.51x 642 1.56 2.9%
8 2,992 t/s 4,067 1.36x 508 1.97 2.3%
16 4,723 t/s 5,764 1.22x 360 2.78 1.9%

Five graphs, 0.4-0.5 ms to capture each (1.5 ms each from a cold context), and graph storage does not move nvidia-smi's MB-resolution counter.

Launch overhead is now solved, and the profiler was over-estimating

Replaying the graph back to back with a single sync gives the GPU's execution time for the kernel sequence with no upload, copy-back or host work -- the floor a full step can approach:

batch pure replay measured step non-kernel kernel share
1 0.561 ms 0.670 ms 0.110 ms 83.6%
2 0.709 ms 0.883 ms 0.174 ms 80.3%
4 1.255 ms 1.558 ms 0.302 ms 80.6%
8 1.478 ms 1.967 ms 0.489 ms 75.2%
16 1.943 ms 2.776 ms 0.833 ms 70.0%

This also corrects the profiler. It had put batch-1 kernel work at 1.015 ms adjusted, and graph replay completes a whole step in 0.670 ms -- so that figure was an over-estimate, and launch overhead was a larger share of the old 1.67 ms step than the 40% it implied. The cause is methodological: the profiler syncs between stages, which suppresses the kernel overlap a replay gets for free. "Adjusted" is an upper bound on kernel time, not a measurement of it.

Device-side argmax

The scheduler needs one token id per request. It was getting that by copying the full logits row back and scanning it on the host -- 3.2 MB per step at batch 16 to find sixteen integers.

The kernel is one block per row, warp-shuffle reduction over (value, index) pairs. What matters is the tie-break, because the host's rule is if v[i] > v[best]: strict, so the lowest index wins a tie, and NaN never displaces anything -- it can only win by starting at index 0. Stated without reference to iteration order that is "the winner is the lowest index i such that no j satisfies v[j] > v[i]", which is order-independent and therefore safe to evaluate as a tree. The merge takes the challenger if it strictly dominates and otherwise keeps the lower index; equal values and NaN comparisons both fall into "neither dominates" and resolve by index. gpu-validate checks all nine cases against the host rule, including exact ties, max at index 0 and at 50303, and NaN at index 0 / mid-row / everywhere.

The kernel is captured as the last node of the decode graph, so token ids are ready the moment replay ends. It runs unconditionally and the full-logit path ignores it, which keeps one graph per batch size serving both paths.

batch logits D2H ids D2H reduction
1 201,216 B 4 B 50,304x
4 804,864 B 16 B 50,304x
16 3,219,456 B 64 B 50,304x

Full-logit and device-argmax trials interleaved, graph replay on for both:

batch full logits device argmax speedup per-request step ms
1 1,481 t/s 1,618 1.09x 1,618 0.62
2 2,271 t/s 2,544 1.12x 1,272 0.79
4 2,644 t/s 3,056 1.16x 764 1.31
8 4,249 t/s 5,268 1.24x 659 1.52
16 6,020 t/s 8,366 1.39x 523 1.91

The gain scales with batch because the transfer it removes does. CRUCIBLE_DEVICE_ARGMAX=0 restores the full-logit path, which is what gpu-logits, gpu-eval and gpu-paged still use -- retrieving complete logits was never removed.

The runtime is now kernel-bound

batch graph replay ids D2H host + scheduler step kernel share
1 0.576 ms 0.023 ms 0.019 ms 0.618 ms 93.2%
4 1.257 ms 0.023 ms 0.029 ms 1.309 ms 96.0%
16 1.882 ms 0.026 ms 0.005 ms 1.912 ms 98.4%

The token-id copy is 0.022-0.026 ms regardless of batch -- latency-bound at 64 bytes, not bandwidth-bound. Host and scheduler work is now under 0.04 ms.

There is no longer a meaningful non-kernel cost to remove, which is the reason to stop optimising the runtime here rather than to look for the next thing. Going further would mean making the transformer itself cheaper, and that is a different project from making the runtime efficient.

What limited the runtime before

The non-kernel column. Host-side argmax is measured at 0.015 ms per request-step -- 2.3% of a batch-1 step, 8.4% at batch 16 -- so it is real but not dominant. The larger part is the device-to-host copy of full logits rows: 3.2 MB at batch 16, which at observed PCIe rates accounts for most of the remaining 0.6 ms.

Both have the same fix, and it is not another kernel optimisation: the scheduler only needs argmax, so a device-side reduction would turn 3.2 MB of transfer plus 0.24 ms of host scanning into 64 bytes. That is the next thing worth measuring, and it is deliberately not done here.

What limited the batched path before

Re-profiling after the change (155 W):

batch 1 batch 16
adjusted step, before 6.895 ms 6.971 ms
adjusted step, after 1.015 ms 2.061 ms
projections, share of step 16% 44%
attention 31.5% 16.6%
logits copy 23.3% 22.9%

The bottleneck is no longer the projections. It is launch overhead: the measured step is 1.67 ms at batch 1 against 1.015 ms of adjusted kernel work, so roughly 40% of the step is the cost of issuing ~183 kernels eagerly. That is what CUDA graphs exist to remove, and it is now the largest single item -- which is the evidence for making graph capture the next task rather than a guess that it would help.

What limited the batched path before

Batch 1 through the batched runtime is 152 tok/s against 1,662 for the single-request path. That looks alarming and is not a paging cost:

single-request path tok/s
contiguous cache + CUDA graph (today's engine) 1,662
contiguous cache, eager 826
paged cache + CUDA graph, single-request path 1,635

Paging costs 1.6% at batch 1. The gap belongs to the batched execution path, which gives up two separate things: graph replay (worth 2.0x) and the GEMV decode kernels. Both eager paths launch ~170 kernels per step, yet the batched one is 5.4x slower, so it is not launch overhead -- it is the kernels. At M=1 the tiled GEMM launches (768/64, 1) = 12 blocks on a 60-SM GPU, where GEMV parallelises across output rows.

That is also why per-request throughput is nearly flat from batch 1 to 16 while aggregate scales 12.6x: the step is dominated by work that does not grow with the batch. A batched GEMV -- grid (rows, batch) -- is the obvious fix and is not attempted here; this task was architecture and correctness, and shipping an unmeasured kernel would be the wrong order.

The single-request path is unchanged. Paging is opt-in via enable_paging; generate, gpu-logits and gpu-eval still use the contiguous cache and graph replay, and their numbers are identical to before.

HTTP inference service

llm-engine serve export/120m --tokenizer export/gpt2.tok --port 8080 --max-batch 16

Binds 127.0.0.1 unless --host says otherwise. There is no authentication, so reaching the network has to be a deliberate act rather than a default.

Concurrency architecture

axum handlers  --jobs-->  inference thread  --tokens-->  per-request channel
                                 |
                         Runtime (scheduler, paged KV, graphs)

The GPU has exactly one owner: a dedicated OS thread. Handlers never touch it. A dedicated thread rather than a Tokio task for two reasons -- the CUDA context and its buffers are not Sync, and a decode step is a blocking GPU call that would stall an async worker for its whole duration. No mutex is ever held across a GPU launch.

That boundary is what keeps batching intact. Had each handler called the model behind a lock, concurrent requests would serialise and the batching engine would have been bypassed by the very layer meant to feed it. Instead every in-flight request goes to one scheduler, which decides how they share a step.

Endpoints

method path purpose
GET /health model, device, max batch, context, sampling mode
GET /metrics active/queued/completed, KV pages, batch size, uptime
POST /v1/generate non-streaming; same scheduler, waits for completion
POST /v1/generate/stream SSE token stream
curl -s localhost:8080/v1/generate   -H 'content-type: application/json'   -d '{"prompt":"The capital of France is","max_tokens":32}'
curl -N localhost:8080/v1/generate/stream   -H 'content-type: application/json'   -d '{"prompt":"The capital of France is","max_tokens":32}'
event: token
data: {"token_id":6342,"text":" Paris"}

event: done
data: {"finish_reason":"length","tokens_generated":32,"text":""}

Sampling is greedy argmax, and /health says so rather than leaving it implied: the schema has no temperature field because the engine has no sampler, not because the field was forgotten.

Streaming text is not per-token decoding

A GPT-2 token is a byte string, not a character. Multi-byte UTF-8 is routinely split across tokens and an emoji spans three or four, so decoding each token alone would emit replacement characters mid-word. IncrementalDecoder buffers bytes until they form something valid, and reproduces decode's lossy handling of genuinely invalid sequences so a stream cannot diverge from a batch decode. Tested by feeding a 4-byte emoji one byte at a time, and end to end by checking the concatenated stream equals the non-streaming response exactly.

Cancellation

A client that disconnects drops its channel. The inference thread notices before its next step, withdraws the request, and hands its KV pages straight back -- an abandoned generation must not keep occupying the batch until max_tokens.

Cancellation takes effect between steps, never inside one: a step is a single fused CUDA graph launch and cannot be interrupted partway. In practice that is one step of latency.

Backpressure

max_queue bounds the waiting queue and max_batch the resident set; over either, the server answers 429 immediately rather than parking the connection. Prompt length, max_tokens and their sum against the model context are all validated before anything is submitted. The page allocator remains the final authority on memory -- these limits stop HTTP turning into unbounded allocation before it gets there.

Service benchmark

Heterogeneous prompts, 128 tokens each, 3 trials, 166 W:

conc e2e tok/s steady tok/s per-request TTFT med TTFT p95 gap med gap p95 batch
1 1,439 1,540 1,439 7.2 ms 7.2 ms 0.65 ms 0.77 ms 1.0
2 2,522 2,887 1,261 10.8 ms 14.3 ms 0.69 ms 0.84 ms 2.0
4 3,891 4,875 973 27.7 ms 27.9 ms 0.82 ms 0.97 ms 4.0
8 5,410 7,506 676 54.3 ms 54.8 ms 1.07 ms 1.27 ms 8.0
16 6,975 10,875 436 106.4 ms 106.9 ms 1.47 ms 1.88 ms 16.0

Two throughput columns because they answer different questions. e2e spans the whole overlapped window including prefill, admission ramp-up and drain -- what a client experiences. steady is n / median inter-token gap, excluding those, and is the number comparable to a decode-only runtime benchmark. Reporting only one would either flatter the service or hide its real cost.

The batch column is the proof of continuous batching over HTTP: mean tokens per decode step is exactly 1.0, 2.0, 4.0, 8.0 and 16.0. Every step carried every in-flight request. Four requests merely completing would have proved nothing.

HTTP overhead is about 5%, measured at concurrency 1 where the comparison is like-for-like: 0.65 ms per token through HTTP and SSE against a 0.62 ms direct runtime step. The higher-concurrency steady figures exceed the direct-runtime table (10,875 against 8,366 at 16) and that is not the service being faster than the engine it wraps -- this benchmark uses short natural prompts while the runtime benchmark cycled 32/128/256/512-token prompts, so attention is cheaper here. Absolute cross-benchmark comparison at those sizes is not valid.

TTFT grows with concurrency because admission prefills each prompt serially, so the sixteenth request waits behind fifteen prefills.

Known limitations

  • Greedy only. No temperature, top-p or beam search.
  • Cancellation is effective at the next step boundary, not immediately.
  • Prompt work is now packed across requests under a token budget; see the cross-request prefill decision record below. Queued requests still wait for resident slots and lifetime KV reservations.
  • No auth, TLS, or multi-model serving; local development scope.
  • /metrics counters are cumulative since startup, which is why the benchmark computes per-run averages from deltas rather than reading them directly.

Sampling

Greedy by default everywhere. Temperature and top-k are opt-in, per request, and deterministic.

curl -s localhost:8080/v1/generate -H 'content-type: application/json'   -d '{"prompt":"Once upon a time","max_tokens":64,
       "temperature":0.8,"top_k":40,"seed":4242}'

Omitting temperature gives greedy, which is what every client written before this feature sends. A default of 0.8 here would have silently changed the output of every existing caller, so absence means greedy rather than "use the usual defaults". top_k or seed without a positive temperature is rejected rather than ignored: silently discarding a parameter the caller set is worse than saying no.

There is no entropy-seeded mode. A sampled reply the user cannot reproduce is worth less than one they can, so an omitted seed uses a fixed documented default (1234) rather than something unrepeatable.

One sampler, not two

Token selection lives in sampling.rs, shared by the CLI and the runtime. It was previously duplicated, which meant the same prompt and seed could produce different text depending on which entry point ran it. Unifying them changed three behaviours, all deliberate:

  • Greedy ties now resolve to the lowest index. The CLI used max_by, which returns the last maximum; the GPU argmax kernel returns the first. Two rules for the same operation is a latent bug, so both now use the kernel's.
  • A NaN logit no longer panics. The old comparator was partial_cmp().unwrap(). In a CLI that is a crash; in a shared inference thread it would take down every other request in the batch. NaN now sorts lowest and never wins, matching the kernel.
  • Candidate order is canonical. The candidate list used to come out of select_nth_unstable_by and be walked in whatever order that left behind: correct as a set, arbitrary as a sequence. Since the walk order decides which token a random threshold lands on, that arbitrary order was load-bearing without being defined anywhere. It is now higher logit first, ties to the lower token id — a strict total order, because token ids are unique. See device-side top-k, which is why it had to be pinned down.

The last of these changed which token a given seed produces. The candidate set is identical and each candidate's probability is identical, so the distribution did not change at all — only which member of it a particular threshold selects. Seeds recorded before the change do not reproduce their old text; seeds recorded after it reproduce everywhere, host or device, which is the property worth having.

Per-request RNG

Each request owns its RNG, keyed on request identity rather than scheduler slot. That matters because retire and cancel use swap_remove, so slot indices are reused by unrelated requests between steps. The invariant:

a request run alone with seed S produces exactly the same tokens when run concurrently with unrelated requests.

Verified by gpu-sampling across seven requests mixing greedy and sampled, prompt lengths 15 to 511, temperatures 0.3 to 1.0, top-k 5 to 500 and five different seeds — identical alone and batched, under simultaneous and staggered admission, and across a cancellation that reuses the same seed. The whole matrix runs twice, once on each selection path, so the fallback is held to the same guarantee and the two paths are pinned to each other.

Greedy consumes no randomness, which is what lets a batch mix both without one perturbing the other.

Mixed batches

The forward pass stays batched regardless of decoding policy; only selection diverges after the logits exist. Greedy rows keep the device-argmax fast path and copy 4 bytes. Sampled rows take their candidates from the device top-k kernel below, or — if they ask for more candidates than it holds — copy their own full logits row.

That choice is made per row, not per batch. One request asking for an unusually wide top-k must not drag the other fifteen onto the slow path with it, and gpu-sampling runs exactly that mixture: seven requests, six taking the device path and one, at top_k 500, not.

Device-side top-k

Sampled decode was the slowest thing in the runtime, and the reason was measurable rather than mysterious. At batch 16, greedy ran at 7,913 tok/s and all-sampled at 3,931 — because every sampled row copied a 50,304-float logits row to the host so the host could pick 40 of them.

D2H per step greedy sampled, full logits sampled, device top-k
batch 1 4 B 201,220 B 1,028 B
batch 4 16 B 804,880 B 4,112 B
batch 8 32 B 1,609,760 B 8,224 B
batch 16 64 B 3,219,520 B 16,448 B

196x less at batch 16. The transfer was only half of it: the host then ran a selection over 50,304 entries per row per step, which measures 80 µs on its own.

topk_rows_f32 returns the k best (value, id) pairs per row instead, in canonical order, and the host samples from those. One block per row; row_k[row] == 0 skips the row, which is how greedy requests pay nothing.

Radix select, not a heap. Extracting the maximum k times is k passes over the row — at k = 40 that is 8 MB per row and it would dominate the step. A per-thread top-k in registers needs k slots per thread to be correct in the worst case, since one thread can own several of the global top-k, and that does not fit. Counting instead finds the exact threshold: one histogram pass per digit, narrowing the key from the top down, then one pass to collect everything at or above it. Three digits (11 + 11 + 10 bits) resolve the float key, and by then the k-th value is unique, so the common case is four passes and only the first touches DRAM.

Ties are where this gets interesting. The key is 64 bits — the order-preserving map of the float in the high half, the complement of the token id in the low half — so descending order on it is the canonical order, and no two keys are ever equal. That is what makes the result independent of thread scheduling: a correct top-k set has exactly one canonical arrangement, so the kernel and the host reference agree by construction rather than by luck. When logits tie exactly, the index half of the key decides, and the kernel keeps narrowing into it rather than taking whichever ids a thread happened to see first.

Two special cases exist only to match the host comparator: NaN maps below everything including -inf, and -0.0 maps to +0.0's key, because IEEE says they are equal and the order must then fall through to the id. Both are done on the bit pattern, so --use_fast_math cannot move a value across a boundary.

The host reference computes the same key, with masks rather than branches on both sides — the sign of a logit is a coin flip, and the branchy version of that function measured 3.7x slower over a 50304-entry row. Having one definition of the order in two languages is the price of computing it in two places; having two definitions would be the bug.

Buffers and graph integration

row_k ([max_batch] i32), cand_vals and cand_ids ([max_batch, 128] each) are allocated once with the rest of the batch scratch — 16 KB at batch 16, nothing per step. The kernel writes from inside the captured decode graph, so candidates are ready the moment replay ends.

Which rows sample, and with what k, lives in row_k — a device buffer, so it can change every step without recapturing anything. Only whether any row samples at all selects a graph, which is one bit, so there are two graphs per batch size instead of one. A single graph carrying the launch unconditionally also works — row_k already makes it a no-op — but it measured a consistent 1% off batch-1 greedy, and greedy not paying for sampling is the point. For the same reason row_k is only uploaded when it changes: an all-greedy run would otherwise pay a host-to-device copy per step that the engine did not make before.

Fallback

top_k above 128 falls back to the full-logit path, which stays as the reference the kernel is validated against and can be forced on with CRUCIBLE_DEVICE_TOPK=0. 128 covers the default of 40 and everything above it with a reason to exist; requests are not clamped to it, because clamping would quietly answer a different question than the one asked.

What it costs, and what it buys

152.44 W, 3090 MHz, prompts 32/128/256/512 cycled, 64 tokens, temperature 0.8, top-k 40, 5 interleaved trials. Every sampled mode ran both ways within the same trial, so drift cannot masquerade as a difference between them:

batch greedy sampled, full logits sampled, device top-k mixed, full mixed, device
1 1,536 t/s 1,224 (0.80x) 1,466 (0.95x)
4 3,031 t/s 2,178 (0.72x) 2,868 (0.95x) 2,537 (0.84x) 2,895 (0.96x)
8 5,228 t/s 3,097 (0.59x) 5,030 (0.96x) 3,915 (0.75x) 5,018 (0.96x)
16 7,913 t/s 3,931 (0.50x) 7,781 (0.98x) 5,180 (0.65x) 7,528 (0.95x)

Sampling used to cost half the throughput at batch 16. It now costs 2%. The ratios in bold are against greedy; against the full-logit path the same rows are 1.20x, 1.32x, 1.62x and 1.98x. At batch 1 the mixed columns are empty because a single row cannot be half sampled.

The selection stage on its own, at k = 40 (gpu-topk-bench, 150.45 W):

sampled rows kernel candidate D2H full-logit D2H host top-k new total speedup
1 24.9 µs 41.6 µs 38.7 µs 79.9 µs 66 µs 1.8x
4 23.1 µs 61.1 µs 153.4 µs 318.2 µs 84 µs 5.6x
8 23.5 µs 36.6 µs 312.4 µs 581.4 µs 60 µs 14.9x
16 24.4 µs 38.8 µs 635.6 µs 1,249.1 µs 63 µs 29.8x

The kernel is flat in both dimensions: one block per row, and 16 blocks fit on 60 SMs, so sixteen rows cost what one row costs. Flat in k too, from 5 to 128 — the work is the passes over the row, not the size of the answer. The candidate transfer is two blocks of rows * 128 regardless of k, small enough that its 37–62 µs is transfer-call overhead rather than bytes.

No dispatch threshold was needed. The plan allowed for one, on the suspicion that the kernel might lose to a single 201 KB copy at small batch. It does not: the device path is ahead at one sampled row — 1.20x end to end, 1.8x in isolation — and pulls further ahead from there. Shipping a crossover the measurements do not show would have been shipping a guess.

Greedy did not regress

Interleaved against the previous commit, five rounds with the order alternating, 149.0–165.6 W, gpu-serve-bench medians:

batch before after
1 1,613 t/s 1,618 t/s
4 3,038 t/s 3,038 t/s
8 5,193 t/s 5,192 t/s
16 7,790 t/s 7,849 t/s

+0.3%, 0.0%, 0.0%, +0.8% — well inside the benchmark's own run-to-run spread, with signs in both directions.

Terminal client

# terminal 1
llm-engine serve export/120m --tokenizer export/gpt2.tok --port 8080 --max-batch 16

# terminal 2
llm-engine tui --server http://127.0.0.1:8080
Crucible  ● connected   120m  greedy  max batch 16  batch 3
+-- conversation ----------------------------------------------+
|  You                                                          |
|  Explain why CUDA graphs help small-model inference.          |
|                                                               |
|  Crucible                                                     |
|  CUDA graphs help because the launch overhead of ...          |
+---------------------------------------------------------------+
+-- generating - Esc to cancel ---------------------------------+
| how does paged attention work_                                |
+---------------------------------------------------------------+
| this request (client-observed) | service                      |
| TTFT            11.3 ms        | active / queued     3 / 0    |
| tok/s              893         | kv pages       178 / 1024    |
| gap median       1.01 ms       | batch (last/avg)   3 / 2.4   |
| generated           87         | tokens / uptime  9812 / 211s |
Enter send   Esc cancel   PgUp/PgDn scroll   F1 help   F2 telemetry   Ctrl+C quit

The TUI is a client. It links crucible::protocol and nothing else from the inference side -- no GpuModel, no Runtime, no PagePool, no CUDA type -- and it builds with --no-default-features --features tui, without the cuda feature at all. That is the boundary being enforced by the compiler rather than asserted in a comment. It never loads the model; start serve first.

Keys

key action
Enter send
Alt+Enter newline
Esc cancel the active generation
Ctrl+U clear input
Left/Right, Home/End cursor
PgUp/PgDn scroll (End returns to the newest text)
F1 help overlay
F2 telemetry panel
F3 generation settings (mode, temperature, top-k, seed)
Ctrl+C quit

Event architecture

keyboard task ─┐
SSE task ──────┤
metrics task ──┼──> AppEvent channel ──> app loop ──> render
health task ───┘

Tasks only send events; one task owns App and applies them. No lock is held across an await, and no task touches a widget.

Rendering is decoupled from token arrival

The backend emits well over a thousand tokens a second and a terminal cannot usefully repaint that fast. Events are applied as they arrive into an 8192-deep channel; the screen redraws on a ~30 FPS tick only when something changed. The coalescing is in the drawing, never in the data -- every token is applied exactly once, and the final text is exact.

Measured: 0.4% CPU idle, 1.7% CPU while receiving ~900 tok/s.

It does not throttle the server

A client that rendered per token would become backpressure on the stream. Driving the same prompt and token budget through the server, one client at a time:

client tokens steps tok/s
plain HTTP reader 512 511 735
plain HTTP reader (repeat) 512 511 1,106
TUI 512 511 998

The TUI's rate sits inside the plain client's own run-to-run range, and both complete 512 tokens in 511 decode steps. No throttling is detectable above that variance — which is the honest claim; the two plain runs differ by 50% from each other, so any smaller effect would be unmeasurable here.

Cancellation

Esc aborts the stream task, which drops the HTTP response. The server sees the disconnect and reclaims the request at its next scheduler boundary. There is no second cancellation protocol, and the TUI adds none. Verified end to end: the server's cancelled_requests increments and every KV page returns to the pool.

A cancelled message is labelled and kept — partial output is not discarded to report that it was interrupted.

Connection handling

● connected, ○ reconnecting, × disconnected. /health is fetched on connect and after a failure, not polled; /metrics polls at 700 ms. One failed poll degrades the indicator rather than ending the session, and reconnect probes run at 1.5 s rather than spinning.

Known limitations

  • Greedy by default; sampling is opt-in per request (F3 in the TUI).
  • No top-p, min-p or repetition penalties.
  • One conversation per process; no persistence, no Markdown, no highlighting.
  • Cancellation takes effect at the next scheduler boundary, not instantly.
  • No model selector, settings screen, or multi-server management.
  • No mouse support; every action has a key.

OpenAI-compatible API

Point an OpenAI client at Crucible and it works:

from openai import OpenAI

client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="not-used")

print(client.chat.completions.create(
    model="crucible-120m",
    messages=[{"role": "user", "content": "Hello"}],
    max_tokens=32,
).choices[0].message.content)

Verified against the official openai Python package (3.8.0) and against the published OpenAPI specification, version 2.3.0. There is no authentication: the SDK requires an api_key argument, any placeholder satisfies it, and the server ignores whatever is sent.

This is a compatibility subset, not an implementation of the OpenAI API. What is implemented is what Crucible actually does. Everything else returns a 4xx naming the parameter, rather than being accepted and ignored -- a server that accepts top_p and then samples with top-k has told the client something false about its own output, and the client has no way to find out.

An adapter, not a second engine

OpenAI request -> DTO -> native job -> [one runtime, one scheduler] -> tokens -> OpenAI chunks

The handlers in openai/ build the same job /v1/generate builds and push it through the same bounded queue to the same inference thread. There is no path from them to GpuModel, no second scheduler, and no separate queue. A compatibility request batches with native requests and with the TUI, and is subject to the same backpressure and the same cancellation.

The native endpoints are unchanged, and the TUI still speaks the native protocol. It is Crucible's own client; making it depend on somebody else's schema would be backwards.

Endpoints

endpoint notes
GET /v1/models one entry, the loaded checkpoint
GET /v1/models/{model} 404 for anything else
POST /v1/completions streaming and non-streaming
POST /v1/chat/completions streaming and non-streaming
curl -s localhost:8080/v1/completions -H 'content-type: application/json' \
  -d '{"model":"crucible-120m","prompt":"The capital of France is","max_tokens":12}'
curl -sN localhost:8080/v1/chat/completions -H 'content-type: application/json' \
  -d '{"model":"crucible-120m","stream":true,"max_tokens":16,
       "messages":[{"role":"user","content":"Hello"}]}'

The model id is crucible-120m, fixed, and appears identically in /v1/models, in every response and in every stream chunk. It is not the model directory: a filesystem path is not something to publish in a field clients paste into config files. Serving a different checkpoint should pass --model-id.

Supported parameters

parameter behaviour
model validated; a different model is 404, never silently served
prompt / messages required
max_tokens, max_completion_tokens both accepted; disagreeing values are an error
temperature absent or 0 is greedy — see below
seed per-request deterministic sampling
stream SSE, terminated with data: [DONE]
stream_options.include_usage adds the final usage-only chunk
top_k a Crucible extension, not an OpenAI field

temperature omitted means greedy, which is a deliberate divergence: OpenAI's default is 1.0. Crucible's native API has always been greedy by default, and matching it means the same prompt produces the same text on every surface — which is what makes the cross-endpoint determinism tests above possible. A client that wants sampling asks for it.

seed without a temperature is accepted rather than refused: greedy decoding already delivers the reproducibility the field is asking for. top_k without a positive temperature is refused, because it is Crucible's own extension and a caller who set it has misunderstood something.

Explicitly unsupported

Each returns 400 with the offending parameter named, unless it is at its no-op value (n: 1, top_p: 1.0, stop: null, penalties at 0, tools: []), which is accepted so ordinary clients that always send defaults still work.

parameter why
top_p Crucible samples with top-k; nucleus sampling is not implemented
frequency_penalty, presence_penalty repetition penalties not implemented
logit_bias not implemented
logprobs, top_logprobs the engine does not return probabilities; zeros would be a fabrication
stop see below
n > 1, best_of one choice per request
tools, tool_choice, functions, function_call this model has no tool-calling training and would never emit a valid call
response_format structured output needs constrained decoding
echo, suffix not implemented
images, audio, modalities text only; a non-text content part is refused, never stringified into the prompt
tool / function message roles representing a tool result would mean inventing a convention the model never saw

stop is refused rather than faked. Implementing it by generating past the stop sequence and trimming afterwards would leave the KV cache, the RNG stream and the token accounting describing text the client never received. It belongs in the scheduler, matching across token boundaries, and that is a runtime feature rather than an adapter one.

The chat template, and what it does not imply

The 120M checkpoint is a base language model. It was trained on FineWeb-Edu, it has never been instruction-tuned, and it ships no trained chat template. Chat Completions here is protocol compatibility, not a capability claim: the model will continue a transcript-shaped prompt with something transcript-shaped, which is not the same as following instructions.

Since no real template exists, this one is defined explicitly rather than guessed:

System: be terse

User: hello

Assistant: hi

User: and again

Assistant:
  • No pseudo-special tokens. <|im_start|> and friends are not in the GPT-2 vocabulary, so they would tokenise into several unrelated pieces the model has never seen arranged that way. Plain English role labels tokenise as ordinary words that do occur in dialogue on the web, which is the only prior this model has.
  • The trailing prime has no space after the colon. GPT-2 merges a leading space into the following token — " hi" is one token, "hi" is another — so ending with "Assistant: " would force a space token followed by a word-without-space token, a pairing the training distribution almost never contains. This is the one detail here that is genuinely about GPT-2 rather than about taste. It is also why replies begin with a space: that space is part of the model's first token, and trimming it would mean the chat endpoint no longer matched the native one for the same prompt.
  • A trailing assistant message is a continuation, not a new turn.
  • One function does this, shared by the streaming and non-streaming handlers, so the two cannot disagree about what was sent to the model.

Finish reasons

length — always, and truthfully. Generation stops when the token budget runs out and for no other reason: this checkpoint has no trained stop token, and stop sequences are not implemented. Reporting stop would assert a natural ending that never happened.

A cancelled request has already lost its connection, so it receives no final chunk and no [DONE]. Fabricating a successful ending for a client that is gone would only corrupt anything replaying the stream from a log.

Usage

prompt_tokens, completion_tokens and total_tokens, from the tokenizer rather than from string lengths. For chat, prompt_tokens counts the serialized transcript actually submitted, which is what the model processed. The prompt_tokens_details and completion_tokens_details sub-objects are omitted rather than zero-filled: reporting cached_tokens: 0 would assert that prompt caching exists here and happened not to help.

On streams, usage appears only when stream_options.include_usage is set, in a final chunk with an empty choices array, exactly as the schema describes.

Errors

The OpenAI envelope, with all four inner fields always present:

{"error": {"message": "...", "type": "invalid_request_error",
           "param": "top_p", "code": "unsupported_parameter"}}
status when
400 malformed JSON, missing field, unsupported parameter, context_length_exceeded
404 unknown model id
429 the bounded queue is full
503 the inference thread is not running
500 inference failed mid-request

Native endpoints keep their own error shape. Rust error text never reaches a client verbatim.

Overhead

Same server, same prompts, same generation config, three surfaces interleaved (168.78 W, 3090 MHz, 64 tokens, median of 3):

clients surface aggregate TTFT ms inter-token ms vs native
1 native 1,419 t/s 7.5 0.61
1 completions 1,420 t/s 7.4 0.60 1.001x
1 chat 1,425 t/s 7.5 0.60 1.004x
4 native 3,376 t/s 27.6 0.78
4 completions 3,340 t/s 27.0 0.78 0.989x
4 chat 3,351 t/s 27.9 0.78 0.993x
8 native 4,434 t/s 52.4 1.02
8 completions 4,350 t/s 53.0 1.02 0.981x
8 chat 4,395 t/s 53.3 1.02 0.991x
16 native 5,336 t/s 104.1 1.42
16 completions 5,266 t/s 103.5 1.42 0.987x
16 chat 5,266 t/s 104.7 1.42 0.987x

0.981x to 1.004x — inside run-to-run noise, with the inter-token interval identical to three significant figures at every concurrency. The adapter is serialisation and nothing else.

Chat's extra cost is prompt length, not protocol, and the two are reported separately for that reason. Its transcript wrapper adds exactly 6 tokens (User: \n\nAssistant:) to every prompt measured:

  5 raw ->  11 chat   'The capital of France is'
 12 raw ->  18 chat   'In a distant galaxy, a small crew of engineers...'
  8 raw ->  14 chat   'The history of the printing press begins in'
 11 raw ->  17 chat   'Q: Why is the sky blue?\nA:'

At these lengths that is invisible in TTFT. On a long conversation it would not be, and it would still be the template's cost rather than HTTP's.

Security boundary

Unchanged from the native service: loopback by default, no authentication, no TLS, bounded request sizes, bounded queue, no user-supplied model path. Local development. The api_key an SDK insists on is ignored.

Anthropic-compatible API

Point an Anthropic client at Crucible and it works:

import anthropic

client = anthropic.Anthropic(base_url="http://127.0.0.1:8080", api_key="not-used")

print(client.messages.create(
    model="crucible-120m",
    max_tokens=64,
    messages=[{"role": "user", "content": "Hello"}],
).content[0].text)

Verified against the official anthropic Python package (1.3.0), whose types are generated from Anthropic's own spec. Supported API version: anthropic-version: 2023-06-01, which is what the SDK sends. An absent header is accepted; a different version is a 400, because answering it with 2023-06-01 semantics would claim to implement changes this adapter has never seen.

A compatibility subset, and not Claude. The model behind it is the same 120M-parameter base LM as everywhere else in this repository — no instruction tuning, no RLHF, no tool training, not a reasoning model. The protocol is Anthropic's; the capabilities are Crucible's.

A sibling adapter, not a wrapper

native handlers    ─┐
OpenAI adapter     ─┼─> one bounded queue ─> one inference thread
Anthropic adapter  ─┘                        one scheduler, one paged KV, one GPU

anthropic/ is not built on openai/. The two protocols disagree about where the system prompt lives, whether max_tokens is required, how content is shaped and what streaming looks like; routing one through the other would make OpenAI's DTOs the interface Anthropic is written against.

What they genuinely share — turning a conversation into prompt text — was extracted into chat_template.rs and is shared there. That is a correctness property rather than tidiness: an Anthropic conversation and the equivalent OpenAI one must reach the same prompt, and with one implementation they cannot drift. The tests check it end to end, and the benchmark prints both surfaces' prompt token counts side by side.

Endpoints

endpoint notes
POST /v1/messages streaming and non-streaming
POST /v1/messages/count_tokens exact, not an estimate
GET /v1/models, GET /v1/models/{id} see the collision below
curl -s localhost:8080/v1/messages \
  -H 'content-type: application/json' -H 'anthropic-version: 2023-06-01' \
  -d '{"model":"crucible-120m","max_tokens":32,
       "messages":[{"role":"user","content":"Explain CUDA graphs."}]}'
curl -sN localhost:8080/v1/messages \
  -H 'content-type: application/json' -H 'anthropic-version: 2023-06-01' \
  -d '{"model":"crucible-120m","max_tokens":16,"stream":true,
       "system":"Be terse.","messages":[{"role":"user","content":"Hello"}]}'

Streaming

Anthropic's protocol is a sequence of typed events with no [DONE] sentinel — message_stop terminates it — so none of the OpenAI chunk framing is reused:

event: message_start        {"message": {...,"usage":{"input_tokens":7,"output_tokens":0}}}
event: content_block_start  {"index":0,"content_block":{"type":"text","text":""}}
event: content_block_delta  {"index":0,"delta":{"type":"text_delta","text":" Hello"}}
...
event: content_block_stop   {"index":0}
event: message_delta        {"delta":{"stop_reason":"max_tokens",...},"usage":{"output_tokens":12}}
event: message_stop         {"type":"message_stop"}

Deltas go through the same IncrementalDecoder the native stream uses, so a character whose UTF-8 bytes span several model tokens arrives whole. Concatenated deltas equal the non-streaming content[0].text byte for byte, and that is tested on CJK output where every character is three tokens.

Sampling: there is no standard knob any more

temperature, top_p and top_k do not appear anywhere in the current SDK's Messages types — not in message_create_params, not under types/ at all. The SDK itself rejects them:

TypeError: Messages.create() got an unexpected keyword argument 'temperature'

So this endpoint takes no standard sampling parameter and decodes greedily. A raw-HTTP request that sends temperature, top_p or top_k gets a 400 pointing at the extensions below, rather than being quietly given greedy output.

Crucible's sampler is still reachable, under names nobody can mistake for Anthropic fields:

extension effect
crucible_temperature positive value enables sampling
crucible_top_k candidate count
crucible_seed deterministic seed

The official SDK never sends these and never needs to. The native API and the TUI remain the places where Crucible's own sampling controls are first-class.

Supported and unsupported

Supported: model (required), max_tokens (required), messages, system as a string or text blocks, stream. Message content may be a string or text blocks. Roles are user and assistant; system inside messages is refused with a pointer at the top-level field. Consecutive same-role messages and a trailing assistant message (prefill) both work. metadata and service_tier are accepted and inert.

Refused with a 400 naming the field, unless empty: stop_sequences, tools, tool_choice, thinking, output_config, container, cache_control, and any non-text content block (images, documents, tool results, thinking blocks).

stop_sequences is refused rather than approximated. Generating past the sequence and trimming afterwards would leave the KV cache, the RNG stream and the output-token count describing text the client never received.

Stop reasons

max_tokens, always, and truthfully. Generation ends when the budget runs out and for no other reason: this checkpoint has no trained end-of-turn token, so end_turn would assert a semantic it never acquired, and stop_sequence, tool_use, pause_turn and refusal all describe machinery that does not exist here. stop_sequence is always null.

Usage and token counting

input_tokens and output_tokens, from the tokenizer. input_tokens counts the serialized transcript — system prompt, role labels, separators, assistant prime — because that is what the model processed. The cache and server-tool counters the schema allows are omitted rather than zero-filled.

count_tokens is exact, not an estimate, because it runs the same conversation builder, the same template and the same tokenizer that /v1/messages runs. There is deliberately no second counting implementation to drift. Three independent paths agree on every tested conversation:

count_tokens               7
usage.input_tokens         7
tokenize "User: Hello\n\nAssistant:"   7

The /v1/models collision

Both protocols claim GET /v1/models with incompatible schemas — OpenAI returns {object, data:[{id, object, created, owned_by}]}, Anthropic returns {data:[{id, type, display_name, created_at}], has_more, first_id, last_id} — so one body cannot satisfy both.

anthropic-version is the discriminator: present means Anthropic, absent keeps the existing OpenAI behaviour byte for byte. The Anthropic SDK sends the header on every request and the OpenAI SDK never does, so the signal is reliable, and both SDKs are tested against this endpoint for exactly that reason. The model id is the same on both sides.

Errors

Anthropic's envelope, built separately from OpenAI's:

{"type": "error",
 "error": {"type": "invalid_request_error", "message": "..."},
 "request_id": "req_..."}
status type when
400 invalid_request_error malformed JSON, missing field, unsupported parameter or content, context overflow
404 not_found_error unknown model id
529 overloaded_error the bounded queue is full
500 api_error inference failed, or the runtime is not running

A full queue is 529 here and 429 on the OpenAI surface: same condition, two vocabularies, each the one its client library retries on. Anthropic's 429 means a quota was exceeded, and this server has no quotas, so using it would be describing a state that does not exist.

Every response and every error carries a request-id header, which the SDK exposes as message._request_id. After a stream has begun there is no status line left, so a mid-stream failure arrives as an Anthropic error event.

Overhead

Same server, same prompts, same budget, four surfaces interleaved (169.88 W, 3090 MHz, 64 tokens, median of 3):

clients native completions chat messages
1 1,427 t/s 1,415 (0.992x) 1,409 (0.987x) 1,374 (0.963x)
4 3,373 t/s 3,329 (0.987x) 3,348 (0.993x) 3,283 (0.973x)
8 4,365 t/s 4,317 (0.989x) 4,362 (0.999x) 4,295 (0.984x)
16 5,288 t/s 5,192 (0.982x) 5,212 (0.986x) 5,163 (0.976x)

The Messages surface runs 1.6–3.7% below native, a little further back than the OpenAI ones. That is per-request framing, not per-token cost: an Anthropic stream sends five framing events around the deltas where an OpenAI stream sends two, and the effect is largest at concurrency 1 where framing is the largest share of a request. The median inter-token interval is identical across all four surfaces at every concurrency (0.60–0.62 / 0.78 / 1.02–1.03 / 1.43 ms).

Template cost is reported separately from protocol cost, and the two conversation surfaces agree exactly — which is the shared serializer, measured:

  5 raw ->  11 chat,  11 messages (+6)   'The capital of France is'
 12 raw ->  18 chat,  18 messages (+6)   'In a distant galaxy, a small crew...'
  8 raw ->  14 chat,  14 messages (+6)   'The history of the printing press begins in'
 11 raw ->  17 chat,  17 messages (+6)   'Q: Why is the sky blue?\nA:'

Security boundary

Unchanged: loopback by default, no authentication, no TLS, bounded request sizes, bounded queue, no user-supplied model path. API keys are ignored — the SDK requires an api_key argument, any placeholder satisfies it, nothing is validated and nothing is logged. Localhost binding is the security boundary.

Prefill scheduling, and why chunking it did not pay

Admission used to prefill a whole prompt inline, so a scheduler step that admitted a long prompt did the entire prefill before decoding anything. Every stream already running felt that as one gap.

Measured, with established streams decoding and one prompt arriving partway through (168.65 W enforced, 3090 MHz max SM clock, medians over three trials, gaps windowed to exclude the established streams' own admission):

existing streams arriving prompt their median gap worst gap stall
1 33 tok 0.86 ms 7.52 ms 8.8x
1 535 tok 0.87 ms 14.59 ms 16.8x
1 941 tok 0.86 ms 28.63 ms 33.3x
4 33 tok 0.93 ms 7.90 ms 8.5x
4 535 tok 0.92 ms 15.16 ms 16.4x
4 941 tok 0.94 ms 29.43 ms 31.4x

The stall tracks prompt length, which is the signature of head-of-line blocking rather than of general load. Sixteen requests arriving together showed the same thing from the other side: a 137.7 ms median time-to-first-token, because sixteen prompts prefilled serially before a single decode step ran.

The scheduler now treats prefill as work, not as part of admission

admit    pending -> prefilling   (page allocation only, no GPU work)
decode   one step for every request that already has a token
prefill  at most one bounded chunk for the oldest prefilling request
retire   whoever finished

Prefilling is an explicit state holding the request's pages, its prompt and how much of it is consumed. Admission no longer touches the GPU, so accepting a 941-token prompt costs the same as accepting a 5-token one. Decode runs before prefill because decode latency is what is being protected, and prefill always runs afterwards, so neither class can starve the other — there is no ratio to tune, the interleave is one to one. Chunks go to the oldest prefilling request until its prompt is done; round-robin would delay every waiting prompt's first token instead of one of them.

Chunks write straight into the request's own pages. Attention for chunk row r covers 0 ..= done + r — the whole cached prefix plus this chunk's causal history — which the paged prefill kernel already supported through its pos_offset argument, so no kernel changed. Nothing is recomputed, nothing is staged and copied, and one SequencePages mapping serves prefill, decode, completion and cancellation alike.

It is correct, and it is off by default

Correctness first, because a policy that is wrong is not worth benchmarking. gpu-prefill-check compares generated token ids, chunked against monolithic:

  • 18 prompt lengths straddling every 16-token page boundary (1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 255, 256, 257, 511, 512), greedy and sampled, at chunk sizes 32/64/128/256 — identical in all 144 comparisons;
  • chunk sizes that align with neither pages nor the prompt (7, 13, 17, 100, 251, 300 against a 251-token prompt) — identical;
  • a request run alone versus prefilled alongside five other prompts — identical, so output does not depend on other requests' timing;
  • cancellation before the first chunk, after one, midway and one chunk from the end — every page returned, the pool reusable, the next request served.

Then the measurement that decides the default. A 941-token prompt arriving into one established stream, each policy on its own server, 159-166 W enforced and 3090 MHz max SM clock throughout:

policy worst gap its TTFT burst-16 TTFT burst-16 aggregate
monolithic 29.1 ms 31.1 ms 143.6 ms 3,733 tok/s
chunk 256 27.7 ms 53.6 ms 101.8 ms 2,940 tok/s
chunk 128 21.1 ms 83.8 ms 166.7 ms 2,329 tok/s
chunk 64 19.5 ms 140.6 ms 237.1 ms 1,693 tok/s
chunk 32 17.5 ms 261.1 ms 236.4 ms 1,083 tok/s

No chunk size wins. The stall shrinks by at most 1.7x, and buying that costs 8.4x the time-to-first-token and two thirds of the throughput. Chunk 256 is the one row with something to show -- the best burst time-to-first-token in the table -- and it barely moves the stall it exists to fix, 29.1 ms to 27.7 ms.

Why: prefill costs 6.8 ms before it looks at a token

TTFT on an idle server, which is prefill plus one decode step (171.76 W, 3090 MHz):

  prompt tokens     1     8    17    33    67   134   201   268   401   535   732   941
  TTFT ms         7.33  7.24  7.32  7.28  7.59  7.99  9.15 10.11 12.46 15.55 22.61 31.01

Flat at ~7.3 ms from 1 token to 134, then linear at about 28.5 us per token. That plateau is a fixed cost paid before a prefill looks at a token — about 6.6 ms of it, net of the one decode step TTFT also contains — so the break-even chunk is around 230 tokens: below that, a chunk's fixed cost exceeds the marginal cost of the tokens it carries.

At this point I attributed that plateau to launch overhead: 208 eager kernel launches through twelve layers, with no captured CUDA graph, the same problem the decode path once had. The next section tested that and it was wrong.

That is enough to explain the whole table. A 941-token prompt at chunk 128 is eight chunks: 8 x 6.6 ms of overhead added to ~27 ms of real work predicts 84 ms, and 83.8 ms is what was measured. At chunk 256 it predicts 58 ms against 53.6 ms measured.

And it explains why no size works. To shrink the stall a chunk has to be small, and small chunks are all fixed cost; to stay efficient a chunk has to be large, and a 256-token chunk already costs 13.9 ms, half of the 29.1 ms monolithic prefill it was meant to break up. The window where chunking both shrinks the stall and keeps the throughput is empty.

What would make it pay is removing that fixed cost, not working around it. The next section does exactly that experiment.

One methodological note, since it nearly produced a wrong number. An earlier version of this curve was taken with the laptop in a power-saving profile and showed 74.3 ms at 941 tokens with a jump between 67 and 134 tokens. The enforced power limit read 160 W throughout, so the cap was not the problem: the GPU was idle between requests and clocked down to 180 MHz, and the curve was measuring boost ramp rather than prefill. The interruption and sweep numbers were unaffected, because there the GPU is already busy. enforced.power.limit alone does not establish an envelope.

What shipped

The scheduler restructuring, because it is correct, it is the necessary foundation, and it fixes real defects on its own: admission no longer blocks on GPU work, a client that disconnects mid-prompt is now cancellable and gets its pages back immediately, and prefill has metrics (prefilling_requests, prefill_chunks, prefill_tokens, last_prefill_chunk_tokens) counted from scheduler state without a single GPU synchronisation.

The chunking policy is off. CRUCIBLE_CHUNKED_PREFILL=1, or --prefill-chunk-tokens N, turns it on for anyone who would rather trade throughput for a smaller worst-case gap, and it doubles as the A/B control every comparison above was paired against.

Established decode is unaffected — decode_active is byte-identical and the step loop adds two calls that return immediately when nothing is prefilling. Measured across three runs at 152–164 W enforced, 3090 MHz: batch 16 at 7,932 / 8,205 / 7,974 tok/s against 7,849 before, batch 1 at 1,521 / 1,543 / 1,576 against 1,618, inside the benchmark's own 3.5–4.9% run-to-run spread.

Limitations

Prompt pages are still reserved in full at admission rather than progressively: the request either fits or waits, which keeps allocator exhaustion a single decision made before anything is written. Prefill remains one request at a time; batching chunks from several requests into one GEMM is a different idea and is not justified while the fixed per-call cost dominates. And the head-of-line stall this milestone set out to remove is still there when chunking is off, which is the honest state of it.

Prefill CUDA graphs: a rejected hypothesis

The previous section ended with a prediction: prefill's ~6.6 ms fixed cost is launch overhead, so capturing a CUDA graph should collapse it, and chunked prefill would then pay. Graphs were built, verified, and measured.

The prediction was wrong. Submission is 0.5 ms of the 6.6.

What the graphs are

One graph per (chunk length, produces logits). Length is in the key because it sets every grid dimension in the sequence; want_logits is in the key because a non-final chunk deliberately skips the lm_head projection, and folding both into one graph would put back work the previous milestone removed.

Exact lengths, no padding buckets — the same standard the batched-decode graphs were held to. Masking rows that must not write KV, must not shift attention and must not affect logits is a correctness proof; exact lengths need no proof.

One value blocked capture. pos_offset, a chunk's offset into its own prompt, was a by-value kernel argument to rope_batch, cache_store_paged and attention_prefill_paged, and a captured graph freezes those. It moved into the existing per-step parameter buffer — the mechanism the decode path already uses for exactly this reason, whose comment already said "anything that changes per token must be read from memory rather than passed by value". Three kernels now read params[PARAM_PREFILL_POS]. Nothing else changed: same GEMM path, same paged attention, same page layout, same pos_offset semantics.

That is what makes the cache small for the case that matters. Because the offset is not in the key, every full chunk of a chunked prefill shares one graph however deep into the prompt it is — a 700-token prompt at chunk 64 replays two graphs eleven times. Monolithic prefill needs one graph per distinct prompt length instead, so the cache is capped at 64 and anything past it runs eager.

Capture is lazy, costs 1.3 ms per graph, and a failure is remembered and falls back to eager rather than taking the server down.

Correctness

gpu-prefill-graph-check compares generated token ids, four paths crossed: eager/graph against monolithic/chunked.

  • 10 prompt lengths x 5 sampling policies (greedy, top-k 5/40/128, and the top-k 500 full-logit fallback) — identical in all 50;
  • the full four-way matrix at chunk sizes 32/64/128/256 and the misaligned 37/73/131 — identical everywhere;
  • request isolation: six requests with different token ids, different pages and different offsets driven through the same two cached graphs, with a cancellation between each — identical to eager, 24 replays, every page returned. That is the test that says nothing request-specific was captured;
  • non-zero offsets: chunk 64 replays one graph at eleven different offsets — identical to eager, which is only possible if the offset was never baked in.

Where the time actually goes

gpu-prefill-bench replays the captured graph in a loop with one sync at the end, which is the kernel sequence executing with no submission cost at all (163.00 W, 3090 MHz, median of 30):

tokens eager graph pure GPU replay submission GPU share
1 6.74 ms 6.23 ms 5.76 ms 0.51 ms 92%
33 6.37 ms 5.88 ms 5.74 ms 0.49 ms 97%
134 6.95 ms 6.45 ms 6.29 ms 0.50 ms 97%
535 14.50 ms 14.12 ms 13.74 ms 0.38 ms 97%
941 30.02 ms 29.62 ms 29.34 ms 0.40 ms 99%

Submission is ~0.5 ms, flat, and does not scale with anything. Prefilling a single token costs 5.76 ms of pure GPU execution. The plateau was never the host talking to the driver; it is the device running ~208 kernels whose grids are far too small to fill it.

Which is a problem this repository has already solved once, elsewhere. The batched GEMV section records the same discovery for decode: "at M=1 it launches 12 blocks on a 60-SM GPU, and profiling showed its cost is flat from batch 1 to batch 16". Prefill runs six of those GEMMs per layer, seventy-two per call, and at small T they are occupancy-starved in exactly the same way. Graphs remove the cost of asking for that work. They cannot make it smaller.

What graphs are still worth

0.4–0.5 ms per prefill, always, for a verified-identical result. That is 7% of a short prefill and 1.3% of a long one, and it shows up where prefills are frequent (161–167 W, 3090 MHz):

eager graph
33-token arrival, worst gap 8.18 ms 7.00 ms
941-token arrival, worst gap 29.03 ms 28.62 ms
burst of 16, median TTFT 139.7 ms 131.9 ms
burst of 16, aggregate 3,799 tok/s 3,931 tok/s

Small, free and never wrong, so graph prefill ships on. CRUCIBLE_PREFILL_GRAPH=0 is the control every number above was paired against.

Chunking, re-decided on the new numbers

Repeating the sweep with graphs on (157–167 W, 3090 MHz, 941-token arrival):

policy worst gap its TTFT burst-16 TTFT burst-16 p95 burst-16 agg
monolithic 28.4 ms 30.6 ms 131.7 ms 132.7 ms 3,905 t/s
chunk 256 27.4 ms 51.4 ms 129.1 ms 213.2 ms 3,069 t/s
chunk 128 20.8 ms 79.8 ms 82.6 ms 307.0 ms 2,431 t/s
chunk 64 19.0 ms 132.7 ms 154.3 ms 465.7 ms 1,790 t/s
chunk 32 17.1 ms 246.1 ms 144.8 ms 813.0 ms 1,138 t/s

Recomputed break-even: fixed cost ~6.3 ms with graphs against ~6.6 ms without, marginal 25.5 us per token, so ~247 tokens against ~230. The model still predicts the table — 941 tokens at chunk 128 is 8 x 6.3 + 24 ms + a decode step = 80 ms, and 79.8 ms was measured — because the term graphs removed was never the one that mattered.

Chunk 128 does now win the burst median time-to-first-token, 82.6 ms against 131.7. It also triples the burst tail and costs 38% of aggregate throughput. So chunking stays off, for the same reason and by nearly the same margin as before.

What this makes the next candidate

Not graphs, and not chunking. The residual is measured GPU compute inefficiency inside small prefills — 5.76 ms to prefill one token, on kernels that launch twelve blocks onto sixty SMs. That is the precondition the previous milestone set for considering cross-request batched prefill: prefill chunks from several waiting requests concatenated into one tensor-core GEMM, so the grids are full. It is the same fix batched GEMV was for decode, and this measurement is what justifies attempting it. It is not attempted here.

Cross-request packed prefill: larger GEMMs, bounded scheduled work

The previous graph experiment identified a roughly 5.76 ms GPU floor for one prompt token. Packing addresses that floor by making independent prompt rows share transformer launches. The result is a partial win: short bursts gain substantially, heterogeneous prompts reach their first token sooner, and small per-request chunks still cost too much to enable by default.

This milestone started at 8f5c745c6309f25f4dae0466a2275598bf333042 with a clean working tree. A separate binary was rebuilt from an archive of that exact commit. The full measurement report, implementation contract, and three-round evidence contain the baseline, parameter sweeps, power envelopes, correctness matrix, and reproduction commands. No weights or GEMM tiles changed.

What actually runs

PrefillBatchPlan carries stable request IDs, prompt offsets, real row counts, and final-slice flags. For four 64-token slices, each projection sees M=256 in one call. Code inspection corrects the earlier estimate: K, V, Q, O, gate, up, and down are seven GEMMs per layer, 84 across twelve layers. Reference and packed execution share one transformer loop.

Persistent int32 owner and absolute-position arrays route RoPE, K/V writes, and attention. Each query reads only its owner's paged history through its own absolute position. The existing 16-token page layout and attention reduction body are shared; there is no dense mask or padding with fake tokens. Only finishing prompt rows reach final normalization and the vocabulary projection. Greedy readback is four bytes per finishing request. Device top-k through 128 and per-row full-logit fallback above 128 retain their existing sampling rules.

Each iteration admits fitting requests, decodes existing streams, then executes one token-budgeted prefill plan. Unfinished contributors rotate to the back. With at most R resident requests, an eligible survivor receives work within R nonempty plans, including under arrivals and cancellation. This is a progress bound, not a latency deadline. Admission remains FCFS and reserves each request's full lifetime page requirement; prompt pages are assigned at admission and later decode growth is lazy. A semaphore now bounds the combined HTTP channel and runtime waiting queue. Draining the channel no longer creates unlimited internal waiting capacity.

Added GPU metadata is 8,256 bytes at token capacity 1024 and batch capacity 16. Existing prefill scratch, page tables, and decode selection buffers are reused. Cancellation is observed between completed GPU calls; pages cannot be recycled while kernels use them. An execution failure stops the inference owner and fails queued/live work instead of continuing with suspect state.

Real service results

These are medians of three alternating baseline/candidate rounds on AC power with the Windows Best Performance overlay. Per-round distributions and loaded clocks are retained in the evidence file; Dell's thermal-mode label could not be verified through its administrator-only CLI. TTFT starts at client request submission, including queueing, prefill, and delivery. The p95 column is the median of each round's request p95, not a percentile inferred from three means.

Workload Baseline TTFT p50 / p95 Packed TTFT p50 / p95 Baseline / packed generation tok/s
16 short prompts (8–64 tokens) 105.90 / 107.01 ms 16.15 / 16.99 ms 5,031 / 8,739
Four 256-token prompts 35.87 / 36.01 ms 23.85 / 24.10 ms 2,553 / 2,846
Four 512-token prompts 54.51 / 54.71 ms 35.17 / 47.28 ms 1,868 / 1,928
Four 941-token prompts 116.47 / 116.63 ms 73.92 / 115.08 ms 1,127 / 1,107
Eight 512-token prompts 109.72 / 109.94 ms 58.31 / 93.59 ms 2,389 / 2,547

The long-prompt tail and aggregate rate matter: lower median TTFT alone does not establish a throughput gain. With an established decoder and four arriving prompts, its median worst gap across rounds fell from 43.03 to 26.64 ms, with material arrival-order variation. A lone 941-token arrival still creates a roughly 27 ms worst gap. The 1024-token budget bounds submitted work; it does not promise sub-millisecond interruption.

One-token HTTP bursts are a further limit: eight and sixteen clients benefit, while two clients had worse TTFT and throughput, and four had lower throughput despite better TTFT. Arrival composition and clock variation are material in that small workload. The dispatch threshold counts actual contributors in a plan; it does not guarantee that concurrent clients arrive in one batch. The full report retains these regressions as well as the successful workloads.

Correctness changed one older decode decision

The expanded HTTP corpus exposed a pre-existing batch-composition bug in the baseline. Above eight active requests, decode switched its final projection to WMMA, rounding activations to half precision. A seeded top-k-500 sequence could change despite identical request input and seed. Int8 decode now retains the same batched GEMV arithmetic at every supported batch size, 1–16. The explicit forced-GEMM comparison remains available. The report separates the cost of this correctness repair from scheduler overhead.

The GPU checks passed 1,084 exact int8 sequences with 24 deterministic fuzz sets and 847 exact f32 sequences, both with zero final-logit difference against independent monolithic prefill. Coverage includes page boundaries, misaligned offsets, permutations, changing neighbors, 16→1→8→2 stale metadata, mixed sampling, cancellation/reuse, malformed metadata, and 16 active decoders shrinking to eight. HTTP added 384 exact comparisons and real page-pressure checks. Native, OpenAI, Anthropic, official SDK, TUI, CPU/mock-server, graph, paged-cache, and held-out CE checks passed; the full report records commands and counts.

An isolated six-request experiment proved that all five client surfaces shared one 94-row packed call after a 941-token singleton blocker. A 60-second mixed load at 35 requests/s completed 1,976 requests plus 124 scheduled disconnects with zero errors and peak queue depth two. A separate 100-request/s saturation test exercised bounded rejection and complete reclamation.

CUDA-event stage attribution puts the 84 transformer GEMMs at about 83% of instrumented time for 16 × 16 rows; attention is about 4%. At 4 × 256 rows, GEMMs account for about 58% and attention 33%, making attention the largest individual stage. Event overhead and clock variation limit absolute timing; the report retains those caveats and the uninstrumented measurements. Nsight GPU activity and Compute Sanitizer instrumentation were unavailable under this WSL setup. No sanitizer pass is claimed.

Production decision

Setting Default
Cross-request packing On; CRUCIBLE_BATCHED_PREFILL=0 preserves reference scheduling
Aggregate prefill budget 1024 real tokens, clamped to model capacity
Contributors per plan max_batch, up to 16
Additional per-request chunking Off; without it, the per-request cap is model capacity
Optional chunk cap 128 when explicitly enabled
Dispatch One contributor uses the existing single-request path; two or more use packed execution
Single-request prefill graphs On, existing bounded exact-length cache
Packed serving graphs None; eager packed work, temporary graphs only for measurement

Budgets 64/128/256/512/1024 and chunk caps 32/64/128/256/none were measured. Small caps reduce individual interruptions but repeat the fixed GPU cost and can worsen long-prompt tails and aggregate throughput. The measured crossover supports the simple two-request dispatch; an adaptive policy and a packed graph cache are not justified by these results. Prefix caching, multi-stream overlap, attention rewrites, and further GEMM tuning remain outside this milestone.

Paged prefill attention: share K/V, preserve the reduction

The attention milestone starts at 135b0953e51801d9748d52672713f19cb48b7ffa. Packing had made the 4 x 256 attention stage the largest individual stage in the prefill profile. The new path keeps the canonical paged layout and shares K/V transport across GQA heads and query positions. The design, complete results, and evidence record the reference, rejected candidates, validation, power envelopes, and reproduction commands.

The reference and the numerical constraint

The old launch assigns a 256-thread block to each query row and query head. Eight warps form dot products, scores occupy a full-history shared buffer, and a block reduction normalizes them. Only 64 threads then accumulate V in ascending history order. Four GQA sibling heads independently load the same KV history, and different query rows share no explicit staging.

The first candidates staged 16/32/64-position K/V tiles and used stable online softmax, retaining a running maximum, normalization sum, and weighted output. For each tile it computes m' = max(m, tile_max), rescales prior state by exp(m - m'), adds exp(score - m') to the denominator and its weighted V to the output, and divides once after the final tile. They passed tensor tolerances with maximum absolute error around 1.1e-6. They nevertheless changed a seeded top-k-500 sequence at its ninth token. Teacher forcing reproduced the failure with identical RNG state and the same 500 candidates: two nearly equal logits swapped ranks 428 and 429. The existing WMMA projections also round attention outputs to half precision; small input differences can cross those rounding boundaries. An unchanged rounded CE value did not make the sequence change acceptable.

Online attention remains an explicit diagnostic experiment. The production kernel preserves the reference's floating-point dot products, exponential normalization sums, value accumulation order, and final reciprocal. No sampling or vocabulary-projection arithmetic changed.

The selected kernel

One block handles four query positions from one request slice and one KV head, hence sixteen independent query distributions. Its 512 threads assign one warp to each distribution, with two output components per lane. A 64-position shared tile first stages K and then V. Each logical page is looked up separately; a tile never assumes consecutive physical page numbers.

This is a two-pass tiled kernel with full per-query score storage on chip, not an online-softmax kernel. To preserve the canonical denominator, each warp emulates the reference's eight virtual warps and both summation trees. Empty virtual warps contribute the same zeros. Each query still has its own scores, normalization, causal bound, and output accumulator.

Descriptors contain packed start, slice length, absolute start, and page-table index. They are uploaded in one bulk transfer. Historical and current K/V both come from the durable pool[page][layer][16][kv_dim]; GQA mapping remains query_head / 4. Query tiles cannot cross request boundaries. Four-row tiles can reuse a fetched component across sixteen distributions, but cache effects mean that source-level reuse is not a measured 16x DRAM-bandwidth reduction.

Shared-score capacity is the largest history in the call rounded up to a 256-position boundary. For this model, the four capacities require 32/48/64/80 KiB of shared storage per block, including the 64-position K or V tile. Driver resource and occupancy-calculator results are reported in the full results; they are not achieved occupancy. The singleton graph key includes this shared capacity as launch topology. Positions, tokens, and page IDs remain dynamic, and the existing 64-entry graph-cache limit remains in force. The selected kernel uses 56 registers per thread and no driver-reported local memory. CUDA's occupancy calculator permits two blocks per SM at 256-score capacity and one at 1024, versus six for the reference. Those limits describe resource residency, not measured execution occupancy.

An additional persistent K buffer was tested so attention could read the historical prefix through pages and the current slice through projection scratch. Both K and V were still stored into the durable pool before attention. That experiment costs exactly 786,432 bytes at 1,024 rows and KV width 192. Its small timing benefit did not justify enabling it. Page-table staging is also disabled. Production adds only the compact segment buffers, 256 bytes each on host and device at sixteen requests, plus small scalar bookkeeping.

Real service, not just kernel timing

These are medians from three freshly warmed, alternating baseline/candidate rounds. The baseline was built from an independently verified archive of the starting commit. Requests generate 64 tokens; p95 is the median of per-round request p95 values. The candidate in these runs explicitly selects the same exact kernel now used by default.

Workload Reference TTFT p50 / p95, ms Exact attention p50 / p95, ms Reference / exact tok/s
16 short prompts 16.54 / 17.05 16.06 / 16.71 8,746 / 8,917
Four 256-token prompts 23.66 / 23.83 22.18 / 22.21 2,849 / 2,913
Four 512-token prompts 35.59 / 48.12 30.04 / 40.42 1,903 / 2,059
Four 941-token prompts 75.19 / 117.37 63.80 / 98.26 1,088 / 1,207
Eight 512-token prompts 61.09 / 98.40 49.06 / 77.06 2,440 / 2,814

The long-prompt gains repeat across rounds. Short cases are mostly unchanged; for example, short-8 throughput is about 0.9% lower, within the variation in these runs. Mixed-8 median TTFT improves while its median p95 is slightly worse, 48.72 to 49.31 ms. These results do not establish universal tail-latency improvement. A lone 941-token arrival reduces the established decoder's median worst gap from 27.49 to 24.40 ms; the four-arrival HOL test has material arrival-order variation and its full distributions remain in the report.

Loaded clock medians differ by less than 1% in each service pair, and maximum SM clock is 3,090 MHz throughout. They are whole-run samples, not counters aligned to every request. Earlier exploratory microbenchmarks include power changes and noisy rounds; their ranges are retained and are not substituted for the service evidence. Tiny singleton attention kernels can lose a few microseconds even though the short service workloads do not materially regress.

Controls and validation

CRUCIBLE_PREFILL_ATTN=reference selects the preserved oracle; CRUCIBLE_PREFILL_ATTN=exact selects the production exact-q4-k64 kernel. Unset uses that exact path for paged prefill with head width 64 and four query heads per KV head. Other geometries use the reference. There is no history threshold. Explicit q*-k* names select the rejected online experiments; the exact-q*-k*, -hybrid, and -cache variants are diagnostic ablations.

The dedicated tensor command compares every candidate with the old kernel and includes shuffled physical pages, nonzero offsets, future poisoning, distinct GQA heads, and persistent descriptor reuse. Exact variants use a bitwise gate. The selected kernel also passed 1,084 complete-model exact sequences with 24 fixed-seed fuzz sets and zero final-logit difference, including mixed sampling, cancellation, and page reuse. Commands and the complete regression record are in the results document:

engine/target/release/llm-engine gpu-prefill-attention-check --fuzz 128
engine/target/release/llm-engine gpu-prefill-attention-bench --variants exact-q4-k64 --iters 50 --trials 5
engine/target/release/llm-engine gpu-packed-prefill-check export/120m --steps 16 --fuzz 24

The final tensor run covers 53 fixed and 128 fuzz cases: 8,688 reference comparisons, 960 future-poison checks and 6,000 composition, permutation and page-remap checks across the variant inventory. All 36 Exact variants match 86,522,880 output elements each bit for bit. The production-default int8 model passes the full 16-step corpus; the separate f32 path passes 847 exact sequences in its four-step corpus. Paged held-out CE matches the reference at printed precision for contexts 32, 256 and 941 (31, 63 and 17 scored positions respectively); the historical decode CE remains 3.720334.

The dedicated microbenchmark covers 25 shapes, including all requested lengths through 1024, ragged groups and 64-row chunks at nonzero offsets. In its five paired rounds, 4 x 256 temporary attention replay measures 576.84 us reference versus 360.98 us selected, with a median within-round ratio of 1.62 (range 1.43–1.74). Power and clocks changed during this run, so these are descriptive kernel observations, not controlled-clock speedup evidence. The independent, closely matched service A/B determines the default.

Scheduler and graph decisions

Three rounds re-evaluated aggregate budgets 256/512/1024 and chunk caps 64/128/256/off under the selected kernel. One resumed round ran at a different clock envelope, so the report compares policies within each round. Chunk 128 reduced the 941-token arrival's worst decoder gap by a median 55.6%, but increased that arrival's TTFT by 184%. Smaller budgets also worsened long-prompt tails. Packing stays on, budget stays 1024, and chunk cap stays off. Singleton graphs remain enabled.

Temporary graphs of the complete packed greedy computation reduced the selected path's synchronized host time relative to full eager calls by 14.1% for 16 x 16 and 17.7% for 4 x 256 in the longer profile runs. These graphs include final selection; replay excludes metadata upload, readback and capture cost. The difference is not pure launch overhead. A production transformer-only graph would keep selection outside and use bounded topology keys for rows, request count, query-grid extent, score capacity and variant. Capture amortization and cache hit rates under ragged service were not measured, so packed serving graphs remain off.

What the final profile and regressions establish

Both the original 20-pass profiles and the repeated 100-pass profiles retain clock differences between baseline and candidate. They are excluded from causal stage-speedup claims. The selected kernel's own longer profile is still useful for locating work: at 4 x 256, attention takes 3.953 ms, gate/up GEMMs 3.801 ms, Q/K/V GEMMs 2.661 ms and down GEMMs 2.586 ms. Attention remains the largest individual category, close to gate/up; this milestone reduces the service bottleneck without claiming to eliminate all attention cost. For 16 x 16, attention is only 0.240 ms of the 8.930 ms instrumented total, and GEMMs dominate. No next-stage optimization is included.

All four required Rust build/check commands and existing GPU suites pass. The HTTP oracle matches 384 sequences. Native, OpenAI, Anthropic and TUI regressions pass 42/125/125/19 checks, including both official SDKs. All five surfaces demonstrably share one 94-row packed batch. The 60-second stable run handles 2,100 submissions with 124 requested cancellations and no overload responses. The 30-second overload run handles 3,000 submissions, returns 1,207 bounded overload responses and never exceeds its queue limit of 64. Both finish with zero runtime failures and all 1,024 pages free; the separate 32-page pressure test also reclaims every page. WSL sanitizer and hardware-counter limitations remain explicit in the full report.

Against llama.cpp and vLLM

python scripts/export_hf.py runs/120m-main --out export/120m-hf --verify
python llama.cpp/convert_hf_to_gguf.py export/120m-hf --outfile export/120m-f32.gguf --outtype f32
llama-quantize export/120m-f32.gguf export/120m-q8_0.gguf Q8_0
python scripts/bench_compare.py --tokens 256 --trials 7

Every other number in this README compares crucible against an earlier version of itself, which answers "did that change help" and not "is this any good".

Decode — 256 tokens, batch 1, greedy, 7 interleaved rounds, 151 W enforced:

engine tok/s spread weights
crucible 1463.6 5.8% int8, per-row scales, 114 MB
llama.cpp (b925e117, CUDA) 862.1 15.6% Q8_0, 122 MB
vLLM 0.28.0 cannot run here, see below

crucible led in all seven rounds with no overlap.

Prefill — 512-token prompt:

engine tok/s
llama.cpp 109,459 ± 21%
crucible (batched) 15,221 7.2x slower
crucible (token at a time) 888 123x slower

The middle row is the point. crucible originally processed prompts one token at a time, so 512 tokens cost ~77,000 kernel launches, each a matrix-vector product — and prefill came out slower than decode, which is backwards and was the tell. Prompt tokens have no sequential dependency on each other, so the whole prompt can go through as a matrix-matrix multiply: compute-bound, with arithmetic intensity that grows with prompt length.

Batching it is 17x faster and costs ~14 launches per layer for the entire prompt instead of ~150 per token. Logits agree with the serial path to 5e-6 relative, and generated text is unchanged.

The prefill gap is now 2.8x, down from 7.2x. That 7.2x was a tensor-core gap: crucible's prefill sustained ~3.4 TFLOP/s, 4.5% of this GPU's BF16 tensor-core peak, on a plain 16x16 tiled f32 kernel, while llama.cpp dispatches to cuBLAS. A hand-written wmma GEMM raised prefill to ~39,700 tok/s at seq 512, roughly 8 TFLOP/s or ~11% of that peak.

What remains is still a kernel-quality gap rather than a mystery. cuBLAS does things this kernel does not: double-buffered shared-memory loads that overlap the next tile's fetch with the current tile's math, ldmatrix for fragment loads, and per-shape tuned tiles. ~11% of peak is a working tensor-core kernel, not a tuned one.

The comparison is valid because the outputs are identical

Greedy decoding, same prompt, both engines:

The capital of France is the capital of the United Kingdom. The capital of the United Kingdom is the capital of the United Kingdom. The capital of the United

Token for token, from two independent implementations using different quantisation schemes — block-wise Q8_0 against per-row int8. That is far stronger evidence that the same model is being measured than any logit tolerance would be. The HF export feeding the GGUF is separately logit-verified against the reference implementation (export_hf.py --verify).

llama.cpp's own trace reports graphs reused = 26, so it is using CUDA graphs as well; neither side has a structural advantage there.

What this result is not

crucible implements one architecture, batch 1, greedy, one quantisation scheme, one context length. llama.cpp supports dozens of architectures, CPU and GPU backends, many quantisation formats, a server, batching and full sampling — and is tuned for models above 7B, where a 120M model at batch 1 is nowhere near its design point. A 1.7x decode margin on this workload says the specialised path is faster on the case it was specialised for; the remaining 7.2x prefill deficit says where tensor cores still earn their keep. It says nothing about the projects.

Three ways this benchmark was wrong before it was right

Each favoured crucible, and each was caught before publication.

The -ngl default. llama-bench does not put all layers on the GPU by default: 583.7 tok/s. With -ngl 99, the same build does 846.5. Publishing the default would have overstated the margin by 45%.

Block ordering. Running all of crucible's trials and then all of llama.cpp's gave llama.cpp 614.6 tok/s at 20% spread against 846.5 standalone — the second engine inherits a hot GPU. This is the same confound already fixed for kernel comparisons and not applied here until it bit. The harness now interleaves, one trial per engine per round, and llama.cpp's number moved back to 862.

An unverified model. Until the greedy outputs were compared, there was no evidence llama.cpp was computing the same thing. A subtly broken GGUF could have been fast for entirely uninteresting reasons.

vLLM cannot run under WSL2

RuntimeError: UVA is not available. WSL2's GPU paravirtualisation does not expose Unified Virtual Addressing, which vLLM's memory management requires. It is a platform limitation, not a configuration problem: native Linux would work, WSL2 and Windows will not.

This also corrects a justification given earlier in this README. The choice of WSL2 was defended partly on the grounds that "vLLM does not officially support Windows" — but vLLM does not run under WSL2 either, so that argument never held. The other reasons stand: Triton is first-class on Linux, and CUDA C++ needs MSVC on Windows where gcc suffices here.

A CUDA 13.3 note

Building llama.cpp needs nvcc, which could not compile any .cu file on this machine: CUDA 13.0's headers conflict with glibc 2.43 over rsqrt. CUDA 13.3 fixes itcuda-nvcc-13-3 compiles cleanly and llama.cpp builds with native sm_120a. crucible stays on NVRTC regardless, since runtime compilation also removes the build-time toolkit dependency.

Layout

model.py           # transformer; ablation axes are config flags
train.py           # training loop, BF16 autocast, MFU accounting
sweep.py           # ablation runner
analyze.py         # figures + significance-aware results table
data/
  prepare.py       # streams FineWeb-Edu into uint16 shards
scripts/
  setup_wsl.sh     # one-shot environment bootstrap
  verify_gpu.py    # device checks, FP8 probe, quick throughput baseline
  bench.py         # repeated-trial harness: median, IQR, thermal state
  bench_prefill.py # GEMM variants: scalar vs tensor-core tiles, interleaved
tests/
  test_attention.py  # GQA path equivalence, causality, param parity
export.py          # checkpoint -> safetensors for the Rust engine
engine/            # Rust inference engine
  src/config.rs      # architecture, refuses to guess
  src/weights.rs     # mmap'd safetensors, bf16/f16/f32
  src/ops.rs         # scalar reference kernels
  src/model.rs       # CPU forward pass
  src/tokenizer.rs   # GPT-2 BPE, pinned against tiktoken
  src/cache.rs       # KV cache for incremental decode
  src/paged.rs       # page pool, per-sequence page tables, CPU-testable
  src/runtime.rs     # continuous-batching scheduler
  src/sampling.rs    # token selection, shared by the CLI and the runtime
  src/chat_template.rs # conversation -> prompt, shared by both adapters
  src/protocol.rs    # native wire types, usable without the engine
  src/server.rs      # axum service, one GPU-owning thread
  src/openai/        # OpenAI-compatible adapter over that service
  src/anthropic/     # Anthropic-compatible Messages adapter, a sibling
  src/tui/           # Ratatui client, native protocol only
  src/gpu.rs         # CUDA backend, NVRTC compilation, validation
  src/gpu_model.rs   # full forward pass on device
  src/quant.rs       # int8 weight quantisation
  kernels/kernels.cu # gemv, rmsnorm, rope, softmax, silu
scripts/
  bench_bandwidth.py # memory bandwidth and the decode ceiling it implies
  export_hf.py       # checkpoint -> HuggingFace Llama, logit-verified
  bench_compare.py   # interleaved comparison against llama.cpp / vLLM
  bench_serve.py     # TTFT and inter-token latency under concurrency
  bench_compat.py    # native vs every compatibility surface, interleaved
  test_serve.py      # native HTTP behaviour against a running server
  test_openai.py     # OpenAI compatibility, raw HTTP and official SDK
  test_anthropic.py  # Anthropic compatibility, raw HTTP and official SDK
  bench_prefill_stall.py # what an established stream feels when a prompt arrives
  smoke_tui.py       # drives the TUI in a pty against a running server
runs/              # per-run log.csv + best.pt checkpoint
figures/           # ablation plots

Tests

python -m pytest tests/ -v

The load-bearing test is test_gqa_paths_match. Grouped-query attention can be computed by materialising KV heads with repeat_interleave or by letting the fused SDPA kernel group them internally, and the two are interchangeable only if their query-head-to-KV-head mapping agrees. If the conventions differed, training would still run and loss would still fall — the model would simply be wrong. The paths are therefore pinned to each other numerically across n_rep ∈ {2, 4, 12}, with and without an attention bias, rather than assumed equivalent.

test_causality verifies that editing token t leaves every output before t bit-identical, which catches a broken mask that a falling loss curve would hide.

The engine's Rust unit and mock-server suites run without a GPU. With cudarc's binding-version override, this also covers host code behind the CUDA feature:

cargo test --manifest-path engine/Cargo.toml --locked
CUDARC_CUDA_VERSION=13020 CUDA_VISIBLE_DEVICES=-1 cargo test --manifest-path engine/Cargo.toml --features cuda --locked

Real kernel checks need a GPU; HTTP and TUI checks also need a provisioned model, tokenizer, and running server:

cd engine
./target/release/llm-engine gpu-validate       # every kernel against the CPU reference
python ../scripts/test_serve.py --port 8080    # native HTTP
python ../scripts/test_openai.py --port 8080 --sdk /path/to/python
python ../scripts/test_anthropic.py --port 8080 --sdk /path/to/python
python ../scripts/smoke_tui.py --binary ./target/release/llm-engine --port 8080

test_openai.py and test_anthropic.py each run twice over: once against raw HTTP, checking payloads field by field against the published schema, and once through the official SDK. An SDK will paper over a wrong object value, a missing logprobs key or a mis-named stream event; raw HTTP will not. Passing --sdk is optional and the SDK half is skipped without it.

Both suites also check what shape tests cannot: streamed text equal to non-streamed across split UTF-8, cross-protocol prompt equivalence, seeded determinism, disconnect cancellation, and that requests from every protocol share decode steps in one scheduler.

Continuous integration

Hosted CI reports existing Rust formatting debt informationally until a separate normalization commit, and gates Clippy correctness and suspicious-code diagnostics, locked CPU tests, feature isolation, and Python/JSON/documentation/repository hygiene. CUDA-feature Rust compilation and host protocol tests run without a toolkit or GPU; kernel execution and NVRTC correctness are covered only by the separate GPU suite.

Optional GPU CI is prepared for manual dispatch on trusted main with local assets and a provisioned NVIDIA runner in a group restricted to this workflow/ref. It offers smoke, full, and native-Linux sanitizer modes. That independent runner access policy must be configured before registration to keep public PR workflows off the GPU machine. The first GitHub runs are pending until commit and push. See CONTRIBUTING.md for exact commands and setup, and the CI report for validation status. Performance measurements remain separate from CI correctness gates.

Roadmap

  • Environment bootstrap and hardware baseline
  • Data pipeline (FineWeb-Edu → uint16 shards)
  • Transformer with switchable architecture axes
  • Training loop with MFU accounting
  • Architecture ablations across seeds (27 runs, 5 axes)
  • Compute-budget analysis; 120M selected as compute-optimal for this GPU
  • Training at 120M (1.44B tokens, val loss 3.2839)
  • Rust inference engine: checkpoint loading, CPU forward pass validated against PyTorch
  • GPT-2 BPE tokenizer in Rust, generation end to end
  • KV cache (11.4x at 30 tokens, ~55x at 150)
  • CUDA kernels, validated against the CPU reference
  • GPU forward pass end to end (49x over CPU, identical output)
  • int8 quantisation: 4x smaller weights, +0.001% cross-entropy
  • CUDA graphs: 1.39x (f32) / 1.60x (int8), and int8's own gain rose 1.05x -> 1.21x
  • Profiler with sync-overhead correction; coalesced attention; warp-per-row int8 GEMV
  • Kernel fusion: SwiGLU in one kernel, residual folded into the projections
  • Split-position attention (flash-decoding) — exact, but slower here; kept opt-in
  • Paged KV cache — 16-token pages, bit-identical to the contiguous cache
  • Batched decode over heterogeneous lengths, no padding to the longest
  • Continuous-batching scheduler — admission, retirement, page reclaim
  • Batched GEMV — 2.4x to 4.0x end-to-end; aggregate 1,916 -> 4,587 tok/s
  • CUDA graphs for the batched path — 1.22x to 2.46x; 70-84% of the step is now pure kernel execution
  • Device-side argmax — 50,304x less D2H; 1.09x to 1.39x end-to-end
  • HTTP inference service — axum, one GPU owner, bounded queues
  • Token streaming over SSE, with cancellation and page reclaim
  • Ratatui TUI client — streaming chat, cancellation, live telemetry
  • Per-request sampling — temperature, top-k, deterministic seeds, through HTTP/SSE and the TUI, with the greedy fast path preserved
  • Device-side top-k — 196x less D2H; sampled decode 1.20x to 1.98x faster and now within 2% of greedy, with greedy itself unchanged
  • OpenAI-compatible API — models, completions and chat completions, streaming and not, over the same scheduler; 0.98x to 1.00x of native
  • Anthropic-compatible Messages API — messages, streaming and not, exact count_tokens, over the same scheduler; 0.96x to 0.98x of native
  • Prefill as scheduled work — explicit prefilling state, non-blocking admission, cancellation mid-prompt, prefill metrics
  • Prefill CUDA graphs — shipped, and the hypothesis behind them rejected: submission was 0.5 ms of the 6.6 ms fixed cost, not the whole of it
  • Cross-request packed prefill — shared transformer GEMMs, token-budgeted scheduling and request isolation; short-burst TTFT 105.9 -> 16.2 ms
  • [~] Chunked prefill — implemented and verified identical to monolithic, but small caps trade smaller decode interruptions for worse prompt tails and throughput, so the additional per-request cap ships off
  • Batched fused SwiGLU (gate/up are still two launches)
  • Throughput comparison against llama.cpp (decode 1.7x faster, prefill 104x slower)
  • Batched prefill — 17x faster, prompt processed as a matrix
  • Tensor-core GEMM — prefill 2.4x faster, llama.cpp gap 7.2x -> 2.8x
  • Tensor-core GEMM tuning closed — ldmatrix already emitted by wmma; double buffering measured and rejected (below)
  • Per-launch tile dispatch — measured, now the default: ties at seq 128, +5.8% / +11.9% / +18.5% at 256/512/1024, bit-identical output
  • vLLM comparison — blocked: WSL2 does not expose UVA, needs native Linux
  • Hosted CPU/source CI and optional manual GPU workflow prepared; first GitHub runs pending

License

Apache-2.0.

About

Training small language models from scratch and building a high-performance Rust/CUDA inference engine for NVIDIA gpus (optimized for Blackwell).

Topics

Resources

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Contributors

Languages