Skip to main content

Module kernels

Module kernels 

Source
Expand description

The lane’s kernel library: one CUDA kernel per kind of op a step runs, each templated over the forms the fusion rules ask of it (a residual added before a norm, an activation quantized in a producer’s epilogue, a result written into the row of a cache a position names). The planner instantiates them; no kernel here knows a model.

Every kernel computes what the language’s sugar means, the formula the VM runs: an RMSNorm divides by sqrt(mean + 1e-6) then scales, an attention scores (q·k)/√hd and normalizes by a max-subtracted softmax, an arg-max keeps the first of equal maxima. Accumulation is f32 in a fixed order (the same inputs give the same bits on every run); the projections are the one place the numerics are a choice — Q8MatKind picks the activation precision of every Q8_0 GEMV:

• Dequant — the activation stays f32, the int8 weight is scaled in the kernel, the dot accumulates in f64: the precise lane, logits a hair from the f64 reference; • Int8Act — the activation quantized to int8 per 32-block and multiplied by __dp4a against the weight (llama.cpp’s MMVQ shape): the fast lane; • TwoPassInt8 — the activation split into a high and a low int8 word, two __dp4a per block: ~14-bit activations at the fast lane’s memory cost.

A quantized activation is a group of 32 values sharing one scale s = max|v| / 127; the high word is round(v / s), the low word round((v − hi·s) / (s/128)) with v − hi·s one fused multiply-add, as the f64 reference computes it.

A Q8_0 matrix is uploaded split, its int8 values in one array and its block scales in another, so the GEMV reads 16-byte words. A matrix in any other GGUF format keeps its blocks as the checkpoint stores them, and BLOCKS decodes a 32-weight sub-block of a row into int8 levels with a scale and an offset per 16 — weight = scale·level + offset, exactly the CPU dequantizer’s arithmetic — so one row dot serves every format under every numerics: the offset term multiplies the activation’s sum, which the int8 lanes take with a __dp4a against ones.

Enums§

Kernel
One library kernel: its name and its source.

Constants§

ATTN_BLOCK_VALUES
The key and value elements an attention block reads a position, across its heads: four latent-attention heads’ worth ((576 + 512) · 4), which scored best from a thousand positions on (the 3060’s attention measurements), against two heads’ at a hundred.
ATTN_FOLD_PAIRS
The value pairs a block of lk_attention_fold folds, a thread each.
ATTN_HELD_BYTES
The most shared memory an attention block (lk_attention) keeps its chunks’ partials in, so a row it takes whole (a batch’s sequence, a chunk’s row) weighs each chunk once: a block of 256 threads keeps three to a streaming multiprocessor.
ATTN_I8_QUERIES
The query vectors a block of the two-word attention (lk_attention_i8) takes, a warp 16 of them.
ATTN_MIN_SPLIT
The fewest positions a split of an attention’s span takes: below it a split’s partial costs more than its scores.
ATTN_MMA_QUERIES
The query vectors a block of the f16 tensor-core attention (lk_attention_mma) takes, a warp 16 of them.
ATTN_SHARED_HEADS
The most query heads lk_attention_shared takes in its one block: four warps of four heads.
ATTN_SHARED_STEP
The cached rows lk_attention_shared stages a step (LK_ST in its source).
ATTN_TILE
The positions the attention scores at a time (LK_TILE in its source).
ATTN_TILE_Q
The query vectors a block of the tiled attention takes: heads of one key/value group over consecutive rows.
AXPY_R
AXPY_SUMS
The bytes of the staged per-head matmul’s own sums in shared memory (red): an f32 for each lane of each of its AXPY_R · AXPY_W warps — it stages only under the quantized numerics, which sum in f32.
AXPY_W
The sub-blocks of outputs and the row groups a block of the staged per-head matmul takes.
BLOCKS
A weight matrix as the kernels read it, and the sub-block decoders of every format besides Q8_0 (template argument F, the GGUF type id): a Q8_0 matrix is its int8 values q and block scales s; any other is its blocks b, a row every rb bytes. lk_sub<F> decodes sub-block sb (weights 32·sb …) of a row into 32 int8 levels and, per half, a scale and an offset; a weight is scale·level + offset rounded as the CPU dequantizer rounds it (a product, then a sum); lk_level<F> decodes one element of a sub-block alike. lk_row_dot<F, Q, M> is a row’s dot with M activation rows (ld apart) under numerics Q, one warp; lk_dot<F, Q, M> the same for any format, a Q8_0 matrix’s split arrays included; lk_gemv_dot<F, Q, M> a row’s dot with M activation rows of its length as the plain GEMV of its format takes it, the one dot of a format every kernel shares so a fused launch computes what its parts computed.
DELTA_CHUNK
The rows of a chunk of the chunked delta rule (LK_DC in its source), the rows of a value head’s memory a block of its walk takes, and the threads of that block: two warps.
DELTA_PREP_SMEM
The dynamic shared memory of lk_delta_prep: a run of 32 of a chunk’s keys’ and queries’ values (then T), A, and the sums T’s lower blocks take — rows padded a value.
DELTA_STATE_ROWS
DELTA_STATE_THREADS
GEMM_BLK_TILE_M
The activation rows of a dense block-format tensor-core GEMM block’s tile (its outputs: GEMM_TILE_N).
GEMM_PARTS
The matrices a launch in parts (lk_parts: lk_gemv_parts, lk_gemv_mma, lk_gemm_blk‘s G 4) multiplies one activation by at most — the kernels’ LK_PARTS: six, so a recurrent layer’s query, key, value, output gate and two gates project its activation in one launch, the gates no launch of their own (at Qwen3.5-0.8B’s batch of 64 a single block on the 3060, 16.8 µs a layer).
GEMM_TILE_N
The outputs a block of the tensor-core block GEMM computes (lk_gemm_blk’s tile).
GEMV_F32_ROWS
The activation rows a block of lk_gemv_f32 takes, together (its LK_F32_ROWS): few enough that a wide chunk’s rows spread over the chip, every row’s dots summed as one row’s are.
GEMV_MMA_MIN_ROWS
The fewest rows a batch’s Q8_0 projection takes to the tensor cores (lk_gemv_mma); fewer run on the GEMV.
GEMV_MMA_ROWS
The most rows lk_gemv_mma takes in one launch: four m-tiles of 16.
HEAD_TILE_M
The rows of a tiled per-head matmul block (lk_head_tiled), each block 64 of a head’s outputs.
PRELUDE
The helpers every kernel shares: the storage types, row addressing, warp and block reductions, the activation quantizer.
ROPE_HEAD
The widest head a rotation turns: its block holds the head’s values in shared memory (lk_rope_head’s hv).
STALL
A stream held back for a span of the global timer. Under --stagger <seed> one runs before each launch of a step, its span drawn from the seed, the form and the launch, so launches its streams leave unordered run in another order under each seed.

Functions§

accumulator
The bytes of the lane’s sums under the numerics of kernel mode mode (mode), and their CUDA type: f64 for the precise numerics (mode 0) — every sum of every kernel, a norm’s, a convolution’s, the delta rule’s, the router’s softmax and the experts’ — f32 for the others.
argmax_blocks
The blocks an arg-max over n values spreads across: a run of at least 2048 values a block, at most 64 blocks.
attn_blocks_a_multiprocessor
The blocks of an attention (lk_attention‘s MB) a multiprocessor holds at once, by the registers it leaves them: three for a launch of more blocks than the chip holds at two (blocks over sms), which come in waves, so occupancy bounds them — on the 3060, 5.7% faster for Llama-3.2-1B’s batch of 64 and 3.4% for Gemma-3-1B’s; two for a launch under one wave, whose blocks’ latency bounds it, and for the f64 sums (acc 8 bytes), which need the registers.
attn_flash_queries
The query vectors a block of the flash attention (lk_attention_flash) takes: 64 for heads and values to 128 wide, 32 past them, so a thread’s outputs stay 64 registers and the block’s shared memory within a block’s.
attn_flash_smem_bytes
The flash attention block’s dynamic shared memory: its queries, two stages of a tile’s keys 32 values of the head at a time and of its values 16 keys at a time, and the tile’s weights — rows padded four values.
attn_flash_threads
The threads of a flash attention block: 256 over 64 queries (a thread 4 queries by 4 keys), 128 over 32 — where the 3060 measured each fastest.
attn_heads_per_block
The query heads an attention block takes together: the largest divisor of group, at most 16, whose keys and values a position stay within ATTN_BLOCK_VALUES (one head whatever its width), whose accumulators (g heads’ value pairs a thread, twice over in a block that takes every chunk) stay within 32 registers, and whose queries, scores and chunk sums fit smem bytes — None when not even one head fits.
attn_held_bytes
The bytes of a block’s g heads’ partials over splits chunks — each chunk’s maximum, sum and vd weighed values, in the accumulator’s acc bytes — and whether the block holds them (held): a row cut into chunks, the partials within ATTN_HELD_BYTES.
attn_holds
attn_i8_keys
The keys a tile of the two-word attention takes: 64, or 32 where heads or values pass 128 wide.
attn_i8_rows
The key rows a key/value group’s two words take for a cache of len positions: whole tiles of 64.
attn_i8_scratch
The two-word attention’s scratch for a cache of len positions over groups key/value groups: the keys’ high and low words, the values’, and each row’s two scales — in bytes.
attn_i8_smem_bytes
The two-word attention block’s dynamic shared memory: a tile’s keys’ and values’ two words, rows padded 16 bytes, and both its scales in two buffers.
attn_mma_keys
The keys a tile of the f16 tensor-core attention takes: 64, or 32 where heads or values pass 192 wide.
attn_mma_smem_bytes
The f16 tensor-core attention block’s dynamic shared memory: a tile’s keys (attn_mma_keys) and their values as f16, rows padded eight values.
attn_path
The path an attention launch takes past a row of one chunk (lk_attention‘s PATH, lk_attn_path), fixed where it launches so its kernel compiles that path alone and ptxas allots its registers to it: the chunks’ partials of a grid of z split blocks a row folded by the last to arrive, held side by side by the row’s one block (side chunks at a time), weighed one at a time by it, or none where a row is one chunk (splits 1).
attn_run_width
The lanes a key’s dot takes in an attention of heads hd wide (lk_attn_run_width, lk_attention’s RW): a half warp for a head of 9 to 16 fours, the whole warp otherwise.
attn_shared_padded_heads
The heads lk_attention_shared pads its g heads to: whole warps of four.
attn_shared_smem_bytes
The dynamic shared memory of a block of lk_attention_shared over g heads of hd: two steps’ rows, and a step’s scores of every padded head.
attn_shared_splits
The most chunks lk_attention_shared cuts a row’s span of span positions into, the decode step’s in every form: one block a streaming multiprocessor (sms), its shared memory a block’s, and no more than give each chunk ATTN_MIN_SPLIT positions.
attn_side_bytes
The shared memory a block weighing side chunks side by side takes past attn_smem_bytes’ one tile of scores: a tile of its g heads’ scores for each other chunk.
attn_side_chunks
The chunks of its row a block that holds them weighs side by side (lk_attn_side): as many as its threads’ value slots hold — each chunk vd / 2 pairs a head, a thread attn_values_per_thread a head — and no more than the row’s splits; one where a chunk fills them.
attn_smem_bytes
The dynamic shared memory an attention block takes: its g query heads of hd (f32), then — in accumulators of acc bytes — their scores over a tile and each head’s maximum, sum and fold factor over each of splits chunks.
attn_splits
The most chunks an attention cuts a row’s span into: enough blocks for two a streaming multiprocessor (sms) over the blocks the head blocks of one row give, and no more than give each chunk ATTN_MIN_SPLIT positions of span.
attn_tile_heads
The heads of one group a tiled attention block takes: the largest power of two dividing group, at most 32 — the block’s rows ATTN_TILE_Q over it.
attn_tile_smem_bytes
The tiled attention block’s dynamic shared memory: a 32-wide chunk of its queries and of a tile’s 64 keys, the tile’s scores, and eight keys’ values.
attn_tile_values
The groups of four values, 128 apart, a thread of the tiled attention accumulates for each of its queries.
attn_values_per_thread
The pairs of each head’s values a thread of an attention block (256 threads) accumulates: lk_attention weighs at most g times as many of its block’s g heads’ value slots a thread.
axpy_pitch
The bytes of a row the staged per-head matmul holds in shared memory (lk_axpy_pitch): the units of AXPY_W sub-blocks and one more, 15 bytes of alignment slack, rounded to an odd number of words.
axpy_staged
Whether a per-head matmul whose rows are its inputs stages them (lk_head_axpy_staged, AXPY_W sub-blocks and AXPY_R row groups a block): under the quantized numerics in every block format but Q6_K and Q8_0 — 1.7–2.8 times the per-lane kernel on the 3060 there (19.9 µs to 9.7 for DeepSeek-V2-Lite’s Q4_K absorption); the per-lane kernel keeps the precise numerics’ f64 sums, Q6_K’s heavier decode and Q8_0’s one activation row. The choice reads no count of activation rows: a prompt’s chunk and a decode step sum alike, so a prompt continues alike in chunks or a token at a time.
block_holds_its_norm
Whether a block of threads runs a norm of a row of cols as the norm’s own block does, each of its threads one of the norm’s (norm_threads). A narrower block plays several of the norm’s warps a warp, one after another, and the row’s norm runs on that one block alone: DeepSeek-V2-Lite’s 320-thread down blocks playing the norm’s 1024 decoded at 142.0 tokens a second on the 3060, 146.7 with the norm’s tail cut and 146.3 with the norm its own launch.
decodes
Whether the block library decodes ty: the formats with a sub-block decoder below.
delta_chunk_scratch
The floats of the chunked delta rule’s scratch over chunks chunks of hv value heads d wide: each chunk’s decayed keys and queries (a value of the width at a time), its T·diag(β) and masked scores, and its rows’ decays to its end with its own.
delta_chunks_take
Whether the chunked delta rule takes heads d wide: at most 128, in whole runs of 32 (a lane of the walk’s warp d / 32 values of the width, the preparation’s runs of 32).
format_arg
The template argument that names a weight format in the block library: its GGUF type id.
gemm_blk_expert_rows
The activation rows of a grouped GEMM block over an expert mixture’s choices: 128 when the average expert has more than 64 rows to multiply (choices / experts), 64 otherwise — a taller tile reads each expert’s decoded weights for more rows, a shorter one leaves fewer of its rows empty.
gemm_blk_launch_smem
The dynamic shared memory one launch of the block-format tensor-core GEMM over k takes: a Q8_0 matrix’s streamed path over a k of whole 128-value scale groups, the block path otherwise — so the streamed path’s blocks are not held to the block path’s larger stages, and two of them share a multiprocessor where its registers allow.
gemm_blk_launch_threads
The threads one launch of the block-format tensor-core GEMM over k takes: a Q8_0 matrix’s streamed path over 128 rows, eight warps of 64 rows by 32 outputs — two blocks of them, their registers and stages, share a multiprocessor — and a warp each 32 × 32 of the tile otherwise (gemm_blk_threads).
gemm_blk_rows
The activation rows a dense tensor-core GEMM of a matrix in format f (its GGUF type id) takes a block: 64 for a Q8_0 matrix under two words, so two blocks of its streamed path, both words of each stage, share a multiprocessor; GEMM_BLK_TILE_M otherwise.
gemm_blk_smem_bytes
The dynamic shared memory a block-format tensor-core GEMM block in format f (its GGUF type id) over bm activation rows may take, whichever path a launch takes (gemm_blk_launch_smem over every k): what its kernel is allowed.
gemm_blk_tail
How a dense block GEMM of tiles tiles and stages stages of k fills its last wave of workers blocks: (whole, split) — the first whole tiles a block each, every later tile cut into split parts of its stages, as many as fill the wave. split 1 when the tiles fill whole waves, or a wave’s blocks cannot cut a remaining tile in two.
gemm_blk_threads
The threads of a block-format tensor-core GEMM block over bm activation rows: a warp each 32 × 32 of the tile.
gemm_f32_splits
The parts an f32 router GEMM (lk_gemm_f32) of tiles 64 × 64 tiles cuts each tile’s chunks runs of 32 k into: as many as give the chip’s sms two blocks each, whole runs of chunks every one.
gemv_blk_norm_threads
The threads of a block-format projection that norms its row in its last block (lk_gemv_blk_norm): eight warps, four where a warp takes per_warp rows past one.
gemv_f32_parts
The warps an f32 GEMV (lk_gemv_f32) cuts each output row’s k into: the most, to eight, that keep 128 values a warp and give its n outputs no more than 16 warps a streaming multiprocessor of the chip’s sms.
gemv_mma_launch
lk_gemv_mma’s launch over m rows (at most GEMV_MMA_ROWS) of a projection of n outputs and k on a chip of sms multiprocessors, (mt, nt, mw, nw): the widest tile its rows fill (gemv_mma_shape), except that a narrow projection of long rows over four m-tiles — its blocks of 128 outputs leaving more than half the chip idle, each running long (k at least 2048) — takes blocks of 64 outputs, twice as many. On the 3060 at 64 rows: Gemma-3-1B’s down projection (6912 by 1152) 94 µs to 54 (62 as the 32-row tile in two slabs of grid layers), Qwen3.5-0.8B’s (3584 by 1024) 54 to 31 (35) and its output projection (2048 by 1024) 33 to 20 (22); wider or shorter projections are faster in blocks of 128.
gemv_mma_shape
lk_gemv_mma’s shape over mt m-tiles, (nt, mw, nw): four n-tiles of 8 outputs a block over one m-tile, a warp each; eight over two, a warp each n-tile over both; sixteen over four, a warp each two n-tiles over two m-tiles — so a block stages its activation rows once for 32·mt outputs, and over four m-tiles each fragment a warp loads feeds two products. On the 3060 at 64 rows the two-by-two warps beat a warp each n-tile over all four m-tiles on Llama-3.2-1B’s down projection (8192 by 2048) 117 µs to 102, its output projection 33.8 to 30.9, its query, key and value 35.6 to 32.9 and its gate and up 171.6 to 168.8; at 16 and 32 rows two n-tiles a warp gain nothing.
gemv_mma_smem_bytes
lk_gemv_mma‘s shared memory (lk_mma_smem): its stages (three, two under the two-word numerics), a stage four blocks of each of its 8·nt matrix rows and 16·mt activation rows (both words’ under two words), rows 144 bytes apart, and the stage’s scales.
gemv_mma_threads
lk_gemv_mma’s threads: a warp each nw n-tiles over each mw of its mt m-tiles.
gemv_mma_tiles
lk_gemv_mma’s m-tiles for m rows (at most GEMV_MMA_ROWS): 1, 2 or 4.
head_mode
The numerics a per-head matmul — a latent’s expansion half — runs under: its activation f32 under every numerics, with f64 sums under the precise one (0) and f32 sums under the quantized ones (3), so the absorbed latent attention computes what the expanded one does.
mode
The template argument that names a numerics mode in the library: 0 f32 activations, 1 int8, 2 two-pass int8.
moe_down_launch
The blocks and threads of a decode step’s down projection (lk_moe_down, lk_moe_down_norm) over a row of d: moe_down_rows outputs a block, a warp each of the k kept experts and the shared expert’s warps where its intermediate sdff is not 0.
moe_down_rows
The outputs a block of the decode step’s expert down projection takes (W of lk_moe_down), and the warps its shared expert’s rows take (SW): four outputs where d divides by four, so each kept expert’s warp has four rows’ weight loads in flight over one load of the hidden row — and a warp a row of the shared expert’s where its rows are wider than an expert’s, so no warp holds the block back. Against a block of one output on the 3060, DeepSeek-V2-Lite’s down projections (Q8_0 beside a shared Q6_K, Q5_0 beside a shared Q4_K) drop 13–31%, gpt-oss-20b’s (MXFP4) 23–36% and Qwen3-30B-A3B’s shape (Q4_K, Q6_K) 16–33%; eight outputs lost but in MXFP4.
moe_group_smem_bytes
The shared memory of lk_moe_group at threads: the running offsets and each warp’s counts.
moe_group_threads
The threads of the block that groups a block’s expert choices (lk_moe_group): a warp a 32 choices, as many warps as 32 and as each warp’s count of every one of experts experts, beside the running offsets, fits 48 KB.
norm_threads
The threads of a block that norms a row of cols (lk_rmsnorm, and the norm inside a rotation): the row’s width in whole warps, at most 1024 — each thread then holds at most cols / 1024 values.
route_warps
The warps an expert router’s row takes (WPE of lk_moe_route): four for an f32 router whose row splits into four parts of whole 16-byte words, so its experts fill four times the blocks; one for a quantized router, whose row is its format’s one dot.
rows_a_warp
The output rows a warp of a GEMV takes (W of lk_row_dot): four in a block format when one activation row meets a matrix of at least 2048 rows of at least 64 sub-blocks — four rows’ weight loads in flight over one load of the activation outran one row a warp by 12–20% on the 3060 there, and lost where fewer warps would cover the chip or a row leaves lanes idle — one otherwise. A Q8_0 row already moves sixteen bytes a lane at a time.