Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Appendix B — The kernel dispatch table

status: draft · path: Muse Glimmer, pinned Muser tree

The decode path on one page: every kernel the engine dispatches for one token, in execution order, with the shader it comes from, the Rust wrapper that binds it, what it reads and writes, and the chapter that explains it. All file:line tags were verified against the pinned tree 6d0807da (see PINNED.md).

B.1 The three kernel sources (legend)

Muser deliberately runs compute from three libraries (Ch 4, crates/muser-engine/src/metal/context.rs:59-131). Every table row below says which source its kernel comes from:

TagSourceBuiltCode
SServing concat library — 27 .metal files concatenated with include_str!, compiled at engine init with fast-math ON (MSL 3.1). Muser-authored muse_reference.metal + nvfp4.metal plus the ferrite-lineage extractionruntime new_library_with_sourcecontext.rs:59-110; 66-name PIPELINES registry metal/encode.rs:21-88
XStrict-f32 cross-vendor library — the same two Muser source files (muse_reference + nvfp4), recompiled with fast-math OFF to match CUDA’s explicit scalar boundaries for the exact/verification lanesruntime new_library_with_sourcecontext.rs:111-121; PSO fields metal/encode.rs:205-277
Lllama.cpp pinned metallib — prebuilt .metallib from llama.cpp commit 89e0aa6f…, loaded from the MUSER_GGML_METALLIB path; supplies the ggml matvec/matmul/norm/rope/unary and flash_attn_ext kernels Muser refuses to re-expressnew_library_with_filecontext.rs:122-131; PSO picks metal/encode.rs:278-293; flash PSO table LlamaFlashAttnPipelines metal/encode.rs:150-167

Rows tagged X run only under MUSER_CROSS_VENDOR_QK (Appendix C); they are the strict arithmetic-ABI mirror, not the serving default.

B.2 Per-layer decode chain (52 layers, execution order)

From encode_token (crates/muser-engine/src/decode.rs:5515-5906), the teacher-forced single-token graph. The serving route packs 1..=4 resident sequences through forward_decode_groupencode_decode_group (decode.rs:4869, :4954), which mirrors this op sequence row for row (every wrapper below appears there too, with rows 1..=4); where the serving route differs, the row says so. Everything from embedding to softcap is one Metal command buffer per token (concurrent dispatch type, explicit barriers, decode.rs:5448-5460).

#StageKernel / functionSrcShader (file:line)Rust dispatch (file:line)Reads → writes (one line)Ch
1Embedding lookupmuser_embedding_q4k (F16 table: muser_embedding_f16)Sshaders/muse_reference.metal:961 (nvfp4.metal:755)encode_embedding_q4k metal/encode/qkv.rs:338 (bind :365); walk decode.rs:5524one u32 token id + one Q4_K embedding row → the [hidden_dim] f32 residual stream11
2Entry norm (weight = ones)llama kernel_rms_norm_mul_f32_4; fallback rms_norm_batchL / Sllama metallib; shaders/ferrite/rmsnorm_batch_tail.metal:1encode_rms_norm_mul metal/encode/norm.rs:244 (bind :270); decode.rs:5535hidden row → normed (RMS, eps 1e-5, × ones)12
3Attention pre-norm — layer 0 only (later layers receive it fused from the previous tail)same as row 2L / Ssame as row 2encode_rms_norm_mul norm.rs:244; decode.rs:5553-5564normedpost_norm with layer.attn_norm12
4Q, K, V, gate projections — one concurrent set of 4 matvecs sharing the input rowllama kernel_mul_mv_q4_K_f32 (Q5_K q5_K, Q6_K q6_K); fallbacks muser_matvec_q4k_4r2s, muser_matvec_q5k_4sg; NVFP4 muser_nvfp4_w4a4_matvec_c1; F16 muser_f16_matvec_c1L / Sllama metallib (pick metal/encode.rs:278-280); muse_reference.metal:735 / :790; nvfp4.metal:312 / :738encode_projection decode.rs:6044encode_quantized_matmul qkv.rs:414 (ggml bind :439, fallback bind :464) / encode_nvfp4_matmul qkv.rs:128 / encode_f16_matmul qkv.rs:68; call decode.rs:5569-5598post_norm row + the four weight matrices → q, k, v, gate activations13
5Per-head QK-norm (parameterless; the ≈3.87 scale is folded into the norm weights)same ggml/rms_norm_batch path as row 2; DFlash’s own route uses rms_norm_per_head (B.6)L / Sllama metallib; rmsnorm_batch_tail.metal:1; ferrite/rms_norm_per_head.metal:15 (registry encode.rs:58)encode_qk_norm norm.rs:286encode_rms_norm_mul norm.rs:244; Q call decode.rs:5600-5608, K call :5609-5617128-wide head slices of q and k, normalized in place14
6RoPE — SWA layers only (uses_rope(), config.rs:68-70); interleaved GPT-J pairs; NoPE full layers skip the dispatch entirelyrope_norm_batch_cached (or llama rope_norm_f32 pinned PSO)S / Lshaders/ferrite/rope.metal:624 (plain rope_batch_cached :566)encode_rope_norm_batch_cached metal/encode/rope.rs:45 (ggml pick :88-137, bind :139); decode.rs:5621-5640cached frequency table + q,k → rotated q,k in place14
7aKV store (token-major ring) + attention — SWA, vec-eligiblemuser_kv_store_f16, then memory barrier, then llama kernel_flash_attn_ext_vec_f16_dk128_dv128 (+ flash_attn_ext_pad when visible % 32 ≠ 0, + flash_attn_ext_vec_reduce)S + Lmuse_reference.metal:979; llama metallib (LlamaFlashAttnPipelines encode.rs:150-167)encode_kv_store_f16 attn.rs:635 (bind :647); encode_llama_flash_attn_decode_vec_f16 attn.rs:437; decode.rs:5660-5693K,V rows → ring slot write_physical; then q + whole ring → attention15, 16
7bAttention — SWA fallback (window not 32-aligned, or no metallib)muser_attention_decode_splitk_f16 + muser_attention_decode_splitk_reduce_f32Smuse_reference.metal:1052 + :1169encode_attention_decode_splitk_f16 attn.rs:708 (binds :748, :776); decode.rs:5695-5723q + ring (split-K partials per workgroup) → attention; geometry splitk_geometry attn.rs:888-89616
7cKV store (head-major plane) + attention — NoPE, vec-eligiblemuser_kv_store_batch_f16 + barrier + llama vec kernel (ns10 = 128)S + Lmuse_reference.metal:1203; llama metallibencode_kv_store_batch_f16 attn.rs:787 (bind :812); attn.rs:437; decode.rs:5728-5770K,V rows → growing plane at position; q + plane → attention15, 16
7dAttention — NoPE fallbackflash_attn_decode_vec_f16_gqa_interleaved + flash_attn_decode_reduce_v2Sshaders/ferrite/flash_attn_decode_vec_contiguous_f16.metal:494; flash_attn_decode_reduce_v2.metal:4encode_ferrite_attention_decode_interleaved_f16 attn.rs:189 (PSOs encode.rs:298-313); decode.rs:5772-5790q + head-major plane (partials) → attention; also re-reads current k,v to dodge a store-load race16
Route ladder predicatesllama_vec_rows = (strict ‖ has_llama_flash_attn_vec) && len>0 && capacity≥32 && (origin_physical==0 ‖ len==capacity); llama_swa = llama_vec_rows && len % 32 == 0decode.rs:5645-565716
8Sigmoid attention-output gatesigmoid_gate_inplaceSshaders/ferrite/sigmoid_gate.metal:7encode_sigmoid_gate metal/encode/gate.rs:7 (bind :17); decode.rs:5793-5799attention × sigmoid(gate) → attention in place17
9o_proj matvecsame stack as row 4L / Ssame as row 4project decode.rs:5800project_tokens :5909-5930encode_projection :6044gated attention + output weights → projected17
10Fused post-attention tail: residual add + post-norm (eps 1e-8) + FFN-norm (eps 1e-5)muser_fused_norm_residual_rms_norm_32sgSshaders/ferrite/rmsnorm_batch_tail.metal:147encode_fused_norm_residual_rms_norm_32sg norm.rs:163 (via …_32sg_batch :190, bind :227); decode.rs:5806-5818normed (residual) + projectedpost_norm; 32 SIMD groups, 1,024 threads, 144 B threadgroup12, 17
11FFN gate+up — fused dual-read when MUSER_FERRITE_FFN_GATE_UP is set and both tensors are Q4_K (the release artifact is); else exact controlffn_q4k_gate_up_silu_4r2s; control path: two row-4 matvecs + muser_silu_mul_inplaceSshaders/ferrite/ffn_fused_tail.metal:496; muse_reference.metal:4encode_ffn_q4k_gate_up_silu_4r2s ffn.rs:10 (bind :25) / encode_silu_mul ffn.rs:38 (bind :49); decode.rs:5819-5862post_norm read once → SiLU(gate)·up written to ffn_gate18
12ffn_down matvecsame stack as row 4L / Ssame as row 4decode.rs:5863-5868:6044ffn_gate + ffn_down weights → projected19
13Fused post-FFN tail: residual + post-FFN-norm (1e-8) + next layer’s attn-norm (last layer: final norm)muser_fused_norm_residual_rms_norm_32sgSrmsnorm_batch_tail.metal:147norm.rs:163; next_norm selected decode.rs:5869-5876, dispatch :5877-5889normed + projected → next layer’s input (or hidden on layer 51)12, 19

Serving-route variants of rows 2/10/13: the packed decode group uses the same …_32sg kernel with rows up to 4 (encode_fused_norm_residual_rms_norm_32sg_batch, norm.rs:190, call decode.rs:5128, :5202); the batch-prefill graph uses the rows-general muser_fused_norm_residual_rms_norm_batch_dual_eps (rmsnorm_batch_tail.metal:250, calls decode.rs:4513, :4656). Under MUSER_CROSS_VENDOR_QK every fused tail decomposes into strict cross-vendor kernels (X) with barriers at each model-dtype boundary (norm.rs:208-225).

B.3 Once-per-token tail

#StageKernel / functionSrcShader (file:line)Rust dispatch (file:line)Reads → writesCh
14Final RMSNormfused into row 13’s tail on the single-token graph (last layer’s next_norm = output_norm, decode.rs:5874-5875); the decode-group path emits a separate normllama kernel_rms_norm_mul_f32_4 / rms_norm_batchL / Sfused: decode.rs:5869-5888; separate: encode_rms_norm_mul decode.rs:5229-5241hidden → normed logits input20
15LM head matvecsame matvec stack as row 4 (kquant); NVFP4 lane: unquantized F16 head muser_f16_matvec_c1 — the ~3.46 ms/token cost that keeps NVFP4 decode at parity, not fasterL / Sllama metallib; nvfp4.metal:738project decode.rs:5892-5897:6044 / qkv.rs:68hidden + vocab projection → logits[vocab]20
16Logit scale + soft capmuser_scale_softcap_inplace; or, to match llama’s graph literally, four ggml unary nodes (scale → tanh → scale)S / Lshaders/muse_reference.metal:15; llama metallib unary (PSO pick encode.rs:289-290)encode_scale_softcap lmhead.rs:163…_count :188 (ggml quartet :230-259); decode.rs:5898-5905logits × logit_scale (= 1/√26 ≈ 0.196116, GGUF metadata, config.rs:190-192) then tanh at softcap 20, in place20
17Sampling read-backargmax / MT19937 sampling on CPU over the read-back distribution; a GPU argmax pair exists for the no-readback benchmark lanes: argmax_f32_phase{1,2} + greedy_argmax_f32_phase{1,2}S (GPU lanes)shaders/ferrite/argmax_f32.metal:7, :41, :77, :125CPU: Session::decode api.rs:696-741 + sampling.rs (distribution buffer retained in place, api.rs:700-703); GPU: encode_argmax_f32_rows lmhead.rs:83, encode_greedy_argmax_f32 lmhead.rs:123full-vocab f32 row read back once per token (4 bytes out on the greedy GPU lane)21

B.4 Prefill batch kernels (the second graph)

Prefill is not decode with more rows: it is a separate batch graph, Session::prefill (api.rs:634) → forward_batch (decode.rs:2857) → forward_batch_hidden (:3788) → encode_batch_hidden_range (:3858-4365), chunked at PREFILL_BATCH_TOKENS = 512 idle / 64 once a decode waits (decode.rs:53-54, :2095-2113). Projections go through encode_batch_projection (decode.rs:5946-5980); attention routes at decode.rs:4090-4365 (Ch 36):

StageKernel / functionSrcShader (file:line)Rust dispatch (file:line)Reads → writesCh
Embedding (per chunk row)muser_embedding_q4kSmuse_reference.metal:961encode_embedding_q4k_from_u32_buffer qkv.rs:376; decode.rs:1756u32 token buffer row → batch hidden36
Projections, NVFP4 M16 route (16 rows, n_in % 64 == 0)muser_nvfp4_w4a4_quantize_m16 + muser_nvfp4_w4a4_prequant_m16_n32 (B.5)Xnvfp4.metal:468 + :504encode_nvfp4_w4a4_prequant_m16 qkv.rs:13; picked decode.rs:5955-5977activations quantized once, then one weight-stationary 32-row tile per projection36
Projections, small batch 4..=8 rowsllama kernel_mul_mv_ext_{q4,q5,q6}_K_f32_r1_{2..5}Lllama metallib (LlamaMulMvExtPipelines encode.rs:174-179)encode_quantized_matmul qkv.rs:482-508the llama-pinned dispatch boundary; changing it breaks logprob parity (qkv.rs:476-481)36
Projections, 16-row K-quant blocksm16_q4k_n32 / m16_q5k_n32 / m16_q6k_n32Sshaders/ferrite/batch_m16_n32.metal:59 / :266 / :163encode_quantized_matmul qkv.rs:556-579DFlash verify/draft blocks; weight-stationary n32 tile, 6 KiB threadgroup36, 33
Projections, aligned Q4_K batchesmatmul_q4k_batch_sgm_alignedSshaders/ferrite/batch_sgm_q4_aligned.metal:69encode_quantized_matmul qkv.rs:584-601Ferrite’s accepted high-occupancy SIMD-group-matrix GEMM36
Projections, general batchllama kernel_mul_mm_q{4,5,6}_K_f32 (aligned/bounds); fallback muser_matmul_q4k/_q5kL / Sllama metallib; muse_reference.metal:912 / :929qkv.rs:604-640the roofline-flipped batch GEMM regime36
KV store, contiguous routemuser_kv_store_batch_f16Smuse_reference.metal:1203attn.rs:787; decode.rs:4135, :4388chunk K,V rows → plane/ring before attention reads them back36
Attention (a): short chunks, < 20 queriesllama vec kernel, one unmasked launch per query rowLllama metallib (attn.rs:437 family)llama_vec_prefill_route_available decode.rs:65; per-row launches decode.rs:4153-4179q row + visible cache → one attention row36
Attention (b): NoPE at llama chunk boundsmuser_fa_causal_mask_f16 + llama flash_attn_ext_blk (once per chunk) + llama kernel_flash_attn_ext_f16_dk128_dv128S + Lmuse_reference.metal:1514; llama metallibencode_llama_fa_prefill_mask_blk attn.rs:266 (binds :283, :307); encode_llama_flash_attn_prefill_f16 attn.rs:328; decode.rs:4197-4208causal f16 mask + skip/partial/dense block bytes, then the masked causal prefill kernel36
Attention (c): local FlashAttention-2flash_attn_v2; one-query GQA specialization muser_flash_attn_decode_gqa_fa2Sshaders/ferrite/flash_attn_v2.metal:59; flash_attn_decode_gqa_fa2.metal:39encode_flash_attention_v2 attn.rs:13 (specialization :61-79); decode.rs:4227, :4329q chunk + f16 KV cache → attention chunk36
SWA ring wrap: staging shadowmuser_stage_swa_prefill_f16 (chunked) / muser_stage_swa_llama_decode_f16 (single-row, llama’s 256-row-padded indices)Smuse_reference.metal:1240 / :1281encode_stage_swa_prefill_f16 attn.rs:103; encode_stage_swa_llama_decode_f16 attn.rs:145; decode.rs:4264, :4281old ring rows + new chunk → detached F16 shadow; ring metadata committed after (append_batch, decode.rs:4348)36, 23
NoPE prefill fallbackmuser_attention_prefill_flash_f16Smuse_reference.metal:1409encode_attention_prefill_f32 attn.rs:830 (bind :861); decode.rs:4359q chunk + current K,V + cache planes → attention chunk (one threadgroup per (head, token))36

Note the last row: the Rust wrapper is named encode_attention_prefill_f32, but at the pin it binds muser_attention_prefill_flash_f16 (attn.rs:861). A sibling kernel muser_attention_prefill_f32 (muse_reference.metal:1326) exists in the registry but is not dispatched by this route (see B.8, conflict 3).

B.5 NVFP4 lane kernels (native 4-bit weights)

All in shaders/nvfp4.metal (Muser-authored); the W4A4 integer contraction and its two scalar epilogue multiplies compile in the no-fast-math library (X) so they match the ModelOpt/MLX integer order (qkv.rs:203-205, Ch 7). Dispatch: encode_nvfp4_matmul (qkv.rs:128) picks by activation scale presence and column count.

KernelSrcShader (file:line)RoleCh
muser_nvfp4_matvec_c{1,2,4,8,16}Xnvfp4.metal:226plain dequantizing NVFP4 matvec; width-1 is the decode kernel, wider calls cover DFlash verification and bounded prefill7
muser_nvfp4_a16_q8_matvecXnvfp4.metal:622weight-only W4A16 route (no input_scale_inv): activations dynamically quantized to Q8 per 16-block, n_in % 2567
muser_nvfp4_w4a4_matvec_c{1,2,4,8,16}Xnvfp4.metal:312W4A4 integer-dot matvec family (weight-stationary 2/4-column specializations)7
muser_nvfp4_w4a4_m16_n32Xnvfp4.metal:33716-column weight-stationary tile, N=32 output rows7, 36
muser_nvfp4_w4a4_quantize_m16 + muser_nvfp4_w4a4_prequant_m16_n32Xnvfp4.metal:468 + :504exact two-pass M=16 route: quantize activations once, then tile — the prefill/verify pair (disabled by MUSER_NO_M16_N32)36
muser_f16_matvec_c{1,2,4,8,16}Xnvfp4.metal:738F16 weights on the NVFP4 lane, incl. the unquantized LM head7
muser_embedding_f16Xnvfp4.metal:755F16 embedding table lookup (row 1’s F16 branch)7
muser_nvfp4_dequant_fixtureXnvfp4.metal:775test fixture: bit-exact dequant of every finite E4M3FN scale7

B.6 DFlash draft kernels (speculative lane)

The five-layer draft runs its own graph in metal/dflash.rs (Ch 8). Prepared-greedy layer loop metal/dflash.rs:1061-1180 (per layer ×5):

StageKernelSrcShader (file:line)Rust dispatch (file:line)Ch
Input normrms_norm_batch (ggml rms_norm when present)L / Srmsnorm_batch_tail.metal:1encode_rms_norm_mul norm.rs:244; metal/dflash.rs:10638
q/k/v projectionsdense f32 matmul_f32_batch_tiled (batch ≥ 4) / matmul_f32_batch; kquant sidecar → the B.4 batch stack incl. m16_*_n32 blocksSbatch_f32_support.metal:45 / :7encode_projection metal/dflash.rs:382encode_dense_f32_batch encode.rs:443 / encode_quantized_matmul qkv.rs:4088
QK-norm (draft is Qwen-style: real weights, not folded scales)rms_norm_per_headSferrite/rms_norm_per_head.metal:15encode_rms_norm_per_head encode.rs:589 (bind :605); metal/dflash.rs:1085-11018
RoPE — NeoX pairing, not the target’s interleaved pairsrope_batch_cachedSferrite/rope.metal:566encode_rope_neox_batch_cached encode.rs:638 (bind :674); metal/dflash.rs:1107-11148
Dual-context attention (64-row sink + sliding window)dflash_dual_attention_f32Sferrite/attention_dflash_dual.metal:15encode_dflash_dual_attention encode.rs:690; metal/dflash.rs:11198
o_proj → fused residual+normfused_residual_rms_norm_batchSrmsnorm_batch_tail.metal:72encode_fused_residual_norm encode.rs:554 (bind :578); metal/dflash.rs:11438
gate/up projections → SwiGLUsilu_hadamard_batchSbatch_ffn_activation_tail.metal:14encode_silu_hadamard_batch encode.rs:482; metal/dflash.rs:11658
down projection → residual addresidual_add_batchSbatch_f32_support.metal:155encode_residual_add_batch encode.rs:500; metal/dflash.rs:11788
Final normrms_norm_batch_inplaceSrmsnorm_batch_tail.metal:42encode_rms_norm_inplace encode.rs:525 (bind :531); metal/dflash.rs:11858
Capture pack (target hidden states → draft input)pack_dflash_layer_major_f32; copy_f32_bufferScopy_f32_buffer.metal:18; :3encode_pack_dflash_layer_major encode.rs:769; metal/dflash.rs:7448
Verify side (target, 16-row blocks)m16_q{4,5,6}k_n32 (B.4) — the L-series tile that took the 16-row verify matmul from ~148 to ~83 ms/cycleSbatch_m16_n32.metal:59/163/266encode_quantized_matmul qkv.rs:556-579 via the mirror-SD suffix begin_dflash_verify_suffix decode.rs:329833
Verify side, opt-in multi-column matvecmuser_matvec_multicol_{q4k,q5k,q6k}_c{1,2,4} (macro)Sferrite/matvec_multicol.metal:400MultiColPipelines multicol.rs:84-103; gate MUSER_MULTI_COL_VERIFY multicol.rs:70-82; also the decode-group route encode_quantized_decode_group qkv.rs:23433

B.7 Strict cross-vendor kernels (X, MUSER_CROSS_VENDOR_QK)

For reference: muser_cross_vendor_{q4k,q5k,q6k} (projections, muse_reference.metal:73/146/221), muser_cross_vendor_rms_per_head (:301), _rms_unweighted (:332) + _mul_weight (:358), _swiglu (:414), _scale (:433), _tanh (:442), _rope (:457) / _rope_neox (:497), _attention_decode (:555), _attention_prefill (:617), _sigmoid_gate (:680), _dual_norm_residual (:690), _residual_add (:724) — all in muse_reference.metal, all compiled fast-math OFF (context.rs:111-121). These replace their S-tagged serving counterparts row-by-row when the flag is set (e.g. decode.rs:5645, norm.rs:208-225, gate.rs:14).

B.8 Conflicts found while building this table

Where the research map, a chapter, and the pinned tree disagreed, the tree won; the disagreements:

  1. QK-norm kernel (row 5). The research map lists rms_norm_per_head as the decode-path QK-norm kernel. At the pin, the target route (encode_qk_norm, norm.rs:286) delegates to the same ggml rms-norm path as row 2 (or the cross-vendor decomposition); rms_norm_per_head lives in the registry (encode.rs:58) and is the DFlash QK-norm kernel (encode.rs:589). Ch 14 already words this correctly (“exists in the pipeline registry”); the table records the tree’s behavior.
  2. Sigmoid gate (row 8). Ch 17 quotes muse_reference.metal:680, which at the pin is muser_cross_vendor_sigmoid_gate (the strict variant). The live serving kernel is sigmoid_gate_inplace (ferrite/sigmoid_gate.metal:7, dispatched gate.rs:17); the cross-vendor kernel is the X mirror of the same math. Both are recorded.
  3. NoPE prefill fallback (B.4 last row). The research map says the fallback dispatches muser_attention_prefill_f32 (muse_reference.metal:1326). At the pin the wrapper encode_attention_prefill_f32 binds muser_attention_prefill_flash_f16 (attn.rs:861, muse_reference.metal:1409); muser_attention_prefill_f32 is in the 66-name registry (encode.rs:42) but has no dispatch on this route.
  4. Chapter in-body line tags. Several chapters cite line numbers inside a kernel body rather than its kernel void line (Ch 11 muse_reference.metal:973 for a kernel that starts :961; Ch 15 :1224 vs kernel start :1203; Ch 16 flash_attn_decode_vec_contiguous_f16.metal:519 vs :494). All resolve within the quoted kernel; this table cites the kernel void lines.

What comes next: the lane matrix and every MUSER_* flag an operator can actually meet — Appendix C.