docs / xpu

One model, start to last token

What happens when dev.cajeta.llm runs Meta-Llama-3.1-8B-Instruct at Q4_K_M, from the first instruction of the process to the last generated token, on the two GPUs the engine is measured on. Every step below answers the same four questions: what happens, what code runs, where the NVIDIA and AMD paths part, and what the optimizer decided and when. The pieces are documented on their own pages; this page is the thread through them.

The run: a 512-token prompt, 128 generated tokens, four sequences at once, the shape every number on the benchmarks is read at. The engine is dev.cajeta.llm at commit 04227de (2026-10-25) on cajeta 0.46.3; the reference machines are Phoenix (RTX 4090, sm_89, CUDA 13.3, under WSL2) and Proton (Radeon 8060S, gfx1151, ROCm). Kernel names in code are the names in the catalog, one row each.

The model, as the engine sees it

“Q4_K_M” is a recipe, not a format: a route is chosen per tensor, so this one model exercises the Q4_K and the Q6_K kernel families on every step, and the Q6_K output head runs on the decode path 128 times per sequence.

1. Execution start

What happens. A program (the cajeta run command, cabra, or a caller of the API) constructs EngineOptions with the model path and calls LlmEngine.load. Nothing has touched a device yet. The first thing load does is decide which device it will be.

What code runs. LlmEngine.load (LlmEngine.cajeta:102). If the options name a device it calls Device.force(name) before any other device call; otherwise the first device touch selects one backend from those the binary bundles, in the fixed order CUDA, HIP, Vulkan, CPU, restricted by CAJETA_XPU_BACKEND when set. The choice is per process and cached; Device.activeBackend() names it. Next, resolveTuneDir (:73) settles where this machine’s measured choices live: the tuneDir option, else CAJETA_LLM_TUNE_DIR (off keeps none), else the platform cache directory with cajeta-llm/tune appended; the directory is handed to every kernel store before a kernel is bound, so no probe can run without a place to remember its answer.

The device is then read, not looked up. Device answers from the driver: multiprocessor count, shared memory per block with and without the opt-in ceiling, registers and threads per multiprocessor, the wave size, free memory, the datasheet bandwidth, what other processes hold (foreignLoadPercent). The kernel registry is the other half of the question: Device.kernelAvailable(name) says whether this backend holds device code for a kernel whose shared memory fits this part, and the library wraps it as KernelCap.usable(name), which also consults the library’s own defect registry and latches the answer per kernel for the process.

NVIDIA against AMD. On Phoenix the selection lands on cuda; the 4090 reports 128 multiprocessors, a 32-lane wave, 48 KB of shared memory per block and 99 KB with opt-in, 100 KB per multiprocessor. On Proton it lands on hip; gfx1151 reports a 32-lane wave for every kernel that declares @Wave(width = 32) (the backend pins +wavefrontsize32 at build), and 64 KB of shared memory per block, which is also the whole per-compute-unit figure. That one difference, 99 KB against 64 KB, is what later decides which prefill tile member survives on each part. A Vulkan run (RADV on Proton, or the NVIDIA ICD on Phoenix) reports its subgroup size and no event clock. Nothing in the library reads the vendor’s name to make a decision; the one place a backend name is tested is backendIsVulkan, and §3 says why.

How it optimizes. Nothing is tuned yet, but two things that will decide every later choice are fixed here: the device key under which the autotune store files its answers (the backend plus the geometry that changes the right answer, so a store copied between machines is ignored rather than believed), and the registry of kernels this backend can run. A kernel the compiler could not lower for this backend was declared so at build time (@Unlowered, tracked against a plan item) and is simply absent; nothing falls through to it.

2. Load

What happens. The file is memory-mapped, its metadata becomes the model configuration, the tokenizer and the chat template are read from it, the model is built and bound to the checkpoint tensor by tensor, and then the engine makes the model device resident and runs every probe it can run without a user’s prompt, so the first request pays for none of them.

What code runs, in the order LlmEngine.load runs it:

1. GgufFile opens the checkpoint; general.architecture = llama and the llama.* keys fill ModelConfig; the tokenizer (tok/) and the end-of-generation set come from the same file. The host holds no second copy of the weights.
2. DecoderModel is built from the configuration and bound by tensor name: 32 DecoderLayers, each an Attention (four Linear projections) and an Mlp (three), plus the embedding, the final norm and the output head. There is no rename table; a tensor the layout does not expect is an error.
3. Residency. Each Linear uploads its packed weight to a KernelBuffer and performs the relayouts its device kernels want: a Q6_K weight is split into a scale prefix and 208-byte bodies (splitScaleKernel, splitPayloadKernel) so the tile kernels read it with aligned 16-byte loads; a Q4_K slab gets a 32-bit word view for the kernels that widen nibbles in registers. Under the default auto prefill-weight policy the formats whose int8 widen route pays get a resident int8 copy (q4kWidenKernel, q6kWidenKernel), inside the widen budget; decode always reads the packed bytes.
4. The K/V cache: on the host, PagedKvCache in 16-token blocks shared by reference count; on the device, KvPool holds one f16 K plane and one V plane per layer, slot s at element s × plane, so one batched attention launch reaches every sequence. The ceiling at four slots is reported as kv-ceiling.
5. prewarmPrefillWeights, then prewarmBatchShapes(maxSeqs × chunk): for every projection, Linear.prewarmTileAt(rows) asks its TileFamily to settle the tile member for the shapes the engine will meet (§3 says how). probeColsFor runs the Cols probe for each format the batched decode will use, and ColsSegProbe.prewarm the fused q|k|v and gate|up variants.
6. The mixture-of-experts residency reserve (none for this dense model), then one throwaway chunk-sized prefill and one decode step, which run the remaining once-per-process probes: matrixCoresPay, tileBeatsDeq, the mat-vec variant probe (probeArms: prefetching, plain, or two-row), and the fused-Cols A/B at the engine’s batch width. A load-phase record breaks the load into bytes moved, packs, prewarm and warm-up.
7. Scheduler(model, chunk) is built with the 150 ms chunk budget; the tokenizer is attached; the block store opens if configured.

NVIDIA against AMD. The same bytes go up on both. What differs is what the prewarm decides: on the 4090 the 128×128 prefill tile opts into 99 KB of shared memory and survives the filter; on gfx1151 it does not fit the 64 KB block limit and the 64×64 member serves (§3). The int8 widen copies are made where tileBeatsDeq says the widen route still pays for a format; the answer is per device and per format and is not assumed from the part’s name. On Proton the K/V planes are allocated from the unified memory the device shares with the host, and the engine counts and reports pool growth because on WSL2 a 2 GiB mapping costs about 30 ms to create and as much to destroy.

How it optimizes. Load is where the store fills. After this step the tune directory holds, for this device key: the tile member per prefill shape (shape.q4kTcTile.<rows>x<out>x<in> and the Q6_K twin, each with a .margin), whether the Cols kernel pays per format at the batch width (colsPays-<type>-b4), the fused segment verdicts (colsSeg-qkv-…, colsSeg-gateup-…), and the mat-vec variant per weight shape (mvVariant-<type>-<rows>x<cols>-c<ncols>), every entry written against the build id of the kernels it measured. A second load on the same machine reads them back and runs no probe. matrixCoresPay and tileBeatsDeq are kept in memory only, once per process.

3. Prefill

What happens. Four requests arrive with 512-token prompts. The scheduler admits them (each one’s whole uncached suffix fits in free K/V blocks, plus one block per layer of decode headroom), and because all four are greedy and the resident batched prefill is ready, it runs one prefill over all their rows: 2048 rows through 32 layers, each projection as one GEMM. The output of the step is one token per sequence, the first generated token.

What code runs. Scheduler.stepOnce collects the prompt chunks (ChunkPolicy.chunkForShare picks the chunk; with nothing decoding it learns the largest that fits, up to 512) and calls CausalLM.stepPrefillBatch (CausalLM.cajeta:826), which checks prefillBatchReady and batchReady(total rows) and otherwise falls back to one stepSeq(seq, tokens, off, n) per request with a prefill-batch-refused record saying why. Inside a chunk:

• A PrefillArena sized to the chunk is carved into borrowed slices for the hidden state and every intermediate (stepSeq:700); nothing survives the call, so peak activation memory is one layer’s at any depth.
• Embedding on the device: the token ids go up (2 KB for 512 tokens) and embedRowQ4kKernel gathers and dequantizes the Q4_K table rows into the hidden state where the layer stack reads it. The host decode of every row plus the upload it replaced cost 45 ms of a 118 ms prefill on sm_89.
• forwardRows(rows) runs the layers. Per layer, on the device throughout:
  1. the input norm fused with the q8_K activation pack (rmsnormPackRowF32 when the width is a multiple of 256, else rmsnormRowsF32 then q8kPackKernel): activations become 256-value q8_K blocks so the integer dot product does the arithmetic;
  2. the query, key and value projections, each through Linear’s batch route (§3.1), the Q4_K ones on the tensor-core tile QkTile_128_128_32_64_128_q4kQ8 or its 64×64 member, the Q6_K value projection (in 16 layers) on Q6Tile_…;
  3. rotary embedding and the K/V write (ropeAppendRowsKernel, kvAppendRangeKernel, or the fused qkPrepBatchKernel), writing f16 keys and values into the layer’s planes of KvPool;
  4. flash attention over the sequence so far, on the f16 matrix cores: attnFlashPrefillF16Kernel, the single-pass online tile at llama.cpp’s precision (64 query-head pairs on one K/V head, the query held in registers, keys and values staged through one shared tile, the softmax kept running and the output rescaled in its fragments). It takes 45 KB of shared memory, so it lowers on gfx1151 where the 88 KB attnFlashPrefillOnlineKernel, the same tile carrying a residual pair for the query and the probabilities, does not; on a backend that holds the f16 kernel back (Vulkan today) the route serves on the residual kernel. Where neither lowers it falls to the two-pass attnFlashPrefillTileKernel, and on a device where a load-time probe finds the tile does not pay, to the wave kernel attnFlashPrefillGqa4Kernel, one wave reading each K/V block once for the four query heads that share it;
  5. the output projection on its tile, accumulated into the residual (addF32);
  6. the post-attention norm and pack; gate and up projections on their tiles; the GLU fused with the next pack (gluFusedPackF32); the down projection (Q6_K in 16 layers) on its tile, accumulated.
• After the last layer only the last row of each sequence goes through the tail: the final norm, the Q6_K output head, and argmaxRowKernel, one integer per sequence read back; the sampler’s temperature is zero in the measured runs, so the argmax is the token.

Every projection announces the arm it took once per route as a batch-route record, so a run’s diagnostic stream names every kernel that served it.

3.1 How a projection chooses its GEMM

Linear.batchRefusal(rows) (Linear.cajeta:2143) runs first and answers with the gate that refused, or null: quantized weight, routing on, a batch GEMM kernel on this backend, the format admitted. For an admitted projection the arms are tried in order (:2600), each gated by a capability the device answered, a measurement, or a fit:

A refused tile says so once with every gate’s value (q4 tcTile refused: avail=… slot=… window=… int8Faster=… ragged=… fits=…), so a route that fell to a slower arm is never silent. Split-K adds workgroups along the reduction when a problem does not fill the device and tcSplitFixupKernel sums the partials.

NVIDIA against AMD. On sm_89 the tile’s int8 operands with int32 accumulators are the native wmma shape, so the fragments are the hardware’s and the multiply is the tensor-core instruction; the 128×128 member needs more than 48 KB of shared memory and opts into the 99 KB ceiling. On gfx1151 the same int8 shape is native WMMA at wave 32, but the block limit is 64 KB with no opt-in, so the filter prunes the 128×128 member under rule R3 and the 64×64 member (128 threads) serves; the fromWords fragment build is a register shuffle on AMD and a shared-memory staging pass on NVIDIA, inside the compiler, so the tile’s source is one. Under Vulkan the engine turns the WMMA, tile and widen routes off (backendIsVulkan, the one backend-name test in the library): the accumulator epilogue demotes to a software tile that trips the display ring’s watchdog and the Q6_K WMMA loses the RADV context, so Vulkan batches through the f16 cooperative GEMM, the mmq tile, or the floor. The bf16 matrix form is native on NVPTX and the software tier on Vulkan; f32 the other way round; f16 native on both, which is why the cooperative f16 GEMM is the portable arm.

How it optimizes. Three moments, three mechanisms.

• Compile time. QkTile is one body over <TM, TN, WM, WN, KS>; @Instantiate built the 128×128 (256 threads) and 64×64 (128 threads) members, @Requires refused any shape the body cannot serve, and @Occupancy(maxThreads = 2 × TN) budgeted registers for the real block (255 on the big member, no spill). The manifest beside each member records registers, spill, shared bytes, the pinned block, and that its matrix verbs lowered to wmma at the native tier. Two gates fail the build rather than ship a slow kernel: a spill, or a backend with no device code that no @Unlowered declaration tracks.
• Bind. TileFamily builds a ShapeFilter over DeviceFacts.current() for the problem (rows, outDim, inDim): R1 prunes a spilling member, R2 one whose threads × registers exceed the register file, R3 one whose shared bytes exceed the per-block limit, R4 one with no resident workgroup; R5 and R6 only record whether the grid fills the device (or the split-K that would) and the padding wasted. Every pruned member keeps its rule and the limit it exceeded; an empty survivor set is an XpuLaunchException naming them, never a silent fallback.
• First use. With two survivors, ShapeChoice times both through the family’s probe, five interleaved rounds, fastest median wins, a lead under three percent recorded as a tie, and remembers the winner under shape.q4kTcTile.2048x4096x4096 against a build id of every survivor’s code hash. On this model the prewarm at load already settled every shape the batch can take, so the first real prefill runs the remembered member. Padding to the grain and the split-K count are decided per launch without timing. Two measurements gate the whole arm once per process: matrixCoresPay admits the matrix cores only if a matrix-core GEMM beats the dot-product kernel by more than ten percent on a synthetic 4096² Q4_K problem (on the CPU backend it does not), and tileBeatsDeq decides between the tile and the widen route at a 512-row gate shape.

3.2 Where prefill stands

On the 4090 the Q4_K and Q6_K prefill routes run at 1.1 to 1.2 of llama.cpp’s; the Q2_K, Q3_K and Q5_K routes at 0.42 to 0.69, because their tiles are younger. The single-sequence prefill of the 30B mixture-of-experts model swings 1.75× between runs a minute apart on Phoenix, the GPU clock, which is why the benchmark reads prefill at four sequences and only after the tune store has settled.

4. Decode

What happens. From the second token on, each sequence advances one token per step, and a token’s cost is almost entirely reading every weight byte once: 4.58 GiB a step, against which the arithmetic is noise. A decode kernel’s only figure of merit is bytes per second against the memory bus. Two regimes exist and the scheduler picks per step: a lone greedy sequence runs the pipelined decode, and several sequences run the batched decode, where one pass over a weight row serves them all.

4.1 One sequence: the pipelined decode

What code runs. Scheduler.stepPipelined (Scheduler.cajeta:555) keeps one step in flight. On the first call it seeds a pinned device token slot with the prompt’s last argmax and calls CausalLM.queueStep(seq, in, out), which returns without waiting. On every later call it queues the next step speculatively, then waits for the previous one’s event and reads back one integer. The host never knows a token before the step that consumes it has been queued. queueStep (CausalLM.cajeta:1074):

1. embedRowQ4kKernel reads the token id from the device slot and writes the embedding into the resident hidden state; no host round trip.
2. forwardRows(rows = 1) takes the resident path (decode-mode: resident, said once): residentLayer per layer:
   • the input norm staged on the device (rmsnormRowF32);
   • the three attention projections in one launch where the load-time A/B said the fused form pays (qkvWaveMatVecKernel), else three wave mat-vecs: q4kQ8WaveMatVecKernel for q and k, and for v the Q4_K or the Q6_K kernel (q6kQ8WaveMatVecKernel) as the layer’s tensor dictates; the prefetching twins (q4kQ8WavePfMatVecKernel, q6kQ8WavePfMatVecKernel) where the variant probe said they pay on this device;
   • DeviceKv.prep: norm, rotary and the K/V write fused in qkPrepKernel (else ropeDecodeKernel and kvAppendKernel), appending this position’s key and value to the layer’s planes;
   • attention split across the context, attnFlashDecodeGqa4Kernel writing partials and attnReducePackQ8Kernel reducing them and re-quantizing the result to q8_K in the same launch, so the output projection reads packed activations directly;
   • the output projection accumulated into the residual on the device;
   • the post-attention norm, gate and up as wave mat-vecs, gluFusedPackF32, the down projection (Q6_K or Q4_K) accumulated back.
3. The fused tail: the final norm, the Q6_K output head as one wave mat-vec over 128256 rows, and argmaxRowKernel writing the token into the pinned output slot; an Event is recorded on the stream and the stream flushed.

One wave owns one output row throughout the mat-vecs: consecutive lanes read consecutive bytes of the row’s 256-value blocks as whole words, dot them against the q8_K activation bytes with the integer dot product, fold the sub-block scales in, and a wave reduction finishes the row. The lane index is tid % Group.width() at every one of the fifty-one sites, so the same source runs at wave 32 and wave 64.

NVIDIA against AMD. On CUDA the decode stream is deferred (KernelStream.setDeferred): the first sight of the step’s launch sequence runs direct, the third is recorded as a graph, and later steps replay it with the changed grids and arguments patched into the nodes, which is how a step of about 1,360 launches is issued in about a millisecond of host time. Deferral is a no-op on HIP and Vulkan, so there the host pays for every launch, and the pipelining (the host’s share of a token overlapping the device’s) matters more, not less. The fused tail, the device argmax and the device embedding are what make the pipeline possible; where any is missing the step runs plainly with a sync. The wave mat-vecs are the same source on both parts at wave 32; on Vulkan a declared 32-lane kernel runs as a virtual wave over the device’s subgroups. The K/V planes are f16 on every backend.

How it optimizes. Decode’s choices were all made before the first token: the wave kernel per weight at bind (Linear latches the format’s wave route when the kernel is usable, the activations are q8_K packed and the width a multiple of 256, else the f32-activation wave, else the scalar kernel, else the host); the prefetching variant per weight shape from mvVariant-…; the fused q|k|v from the load-time A/B. Nothing is timed during decode. The measured standing on the 4090 is at or above llama.cpp’s single-sequence rate at every quantization from Q2_K to Q8_0: the kernels were at the bus, and the gaps between them were the host, which the pipeline removed.

4.2 Four sequences: the batched decode

What code runs. decodeBatchReady (CausalLM.cajeta:3800) is latched once per model: the K/V cache device resident, the attention kernels usable, head dimension at most 128, hidden and query widths multiples of 256, every projection and the output head answering colsReady (no epilogue, packed activations, a Cols kernel for its format in the registry). A refusal is recorded as decode-batch refused: <why> and the sequences decode singly. When it holds, Scheduler.stepOnce collects every decoding row into one batched step (:4000 onward):

• every sequence’s token embedded on the device (embedRowsLaunch), the hidden rows normed and packed (q8kPackRowsLaunch);
• each projection as one Cols launch over the batch’s activation columns: q4kQ8WaveMatVecColsKernel (or the prefetching and two-row variants q4kQ8WavePfMatVecColsKernel, q4kQ8WaveR2MatVecColsKernel where the Cols probe said they pay at width 4), q6kQ8WaveMatVecColsKernel for the Q6_K tensors; the fused q4kQ8WaveMatVecColsSegPlainKernel and …ColsSegPfKernel serving q|k|v and gate|up in one launch where the load-time A/B said so. One pass over a weight row serves four sequences, which is why four cost close to one in weight traffic;
• qkPrepBatchKernel, then one attnFlashDecodeGqa4BatchKernel launch attending every sequence through the whole KvPool buffer (slot s at element s × plane), attnReducePackQ8BatchKernel reducing and packing;
• the tail per row and argmaxRowKernel per row; a decode-batch-phase record splits the step into host embedding, upload, issue and device tail.

Since 2026-10-24 the batched greedy decode also runs one step ahead on HIP, the single-sequence pipeline’s design applied to the batch, rows embedded on the device from the previous step’s argmaxes with one event per step.

NVIDIA against AMD. The Cols kernels exist for exactly the six formats the batched decode admits (Q4_K, Q6_K, Q5_K, Q8_0, Q2_K, Q3_K); this model uses two of them. On the 4090 the four-sequence decode sits at 0.95 of llama.cpp for Q4_K_M, and between 0.89 and 0.97 across the six formats; on Proton’s Vulkan path the decode at four sequences sits at 0.98 of llama.cpp’s Vulkan. The batched pipeline is HIP’s answer to the graph replay it cannot have.

How it optimizes. Per format and batch width, the Cols probe at load timed the Cols kernel against single-row launches after checking each candidate against the host mat-vec (colsPays-q4k-b4); the fused segments were A/B’d in place on the real batched step; the per-launch choice between one Cols launch and several single-row launches reads those answers without timing.

4.3 Longer contexts

At 8k and 32k tokens the attention’s share grows: the decode kernel splits the context across workgroups and the reduce kernel’s partial count grows with it; the K/V pool’s planes grow to the slots a batch needs, each under 2 GiB, and growth is counted because on WSL2 a 2 GiB mapping costs tens of milliseconds to create and destroy, so the engine moves it off the clock rather than making it cheap. The benchmark’s llm/decode-8k and llm/decode-32k rows measure exactly this.

5. What the optimizer did, kernel by kernel

For every kernel this page named, where its configuration came from.

The CPU backend runs the same catalog as a target, not an emulator: the work-item loop vectorizes at the host’s SIMD width, the wave verbs become vector reductions and shuffles, and because the backend reports no register or spill count, the filter prunes nothing there and what differs is which members are fast; matrixCoresPay answers no, and the dot-product floor serves the prefill.

6. What the profiler shows

CAJETA_PROFILER=1 arms the host sampler and the device capture before any backend initializes and writes a .pftrace (into .cajeta/profiles/, one file per run) that the IntelliJ tool window or Perfetto opens. Read against this page it shows:

• the load as a load-phase breakdown: bytes moved, the relayout and widen kernels, the prewarm probes’ bracketed launches (recognizable as twenty-launch bursts of one kernel, five rounds, interleaved), the throwaway prefill and decode step;
• each prefill chunk as the per-layer sequence of §3, the tiles the longest spans, each device span carrying its tier (DEVICE from CUPTI on CUDA, from rocprofiler on HIP; EVENT where only the driver is installed; HOST on Vulkan outside the timestamp seam), with batch-route and tile-choice records beside the launches that took them;
• each decode step as a dense run of about 1,360 launches of a few microseconds each (an armed profiler keeps every launch direct, so the graph replay is off while tracing), the two events per step in the pipeline, and the host’s wait on the previous event as the only host gap; sched-phase records give the step’s admission, K/V reserve, model and total nanoseconds, and decode-batch-phase the batched step’s host embedding, upload, issue and device tail.

The trace this section will link is the first quiet-box run’s, the same execution the benchmarks read, so that the numbers and the picture are one measurement. Until that run exists this section describes the trace from the instrumentation’s own records rather than from a captured file.

Where to go from here

• The kernel families and how one is chosen, in full: Language models on the device.
• The engine around them: Inside the engine, the API, configuration.
• The substrate: One kernel, any device, the scheduler, optimization, and the contract, cajeta.xpu.
• How the numbers are read: Measuring and the benchmark pages.

Source: docs/xpu/walkthrough.md · 32 min read · 7126 words