Shader shelf
WGSL kernels ported verbatim from the old llm/src/backend/wgpu/shaders/.
None of them are wired into mr's op dispatch yet — they sit here as
ready-to-use source so the implementation knowledge ships forward.
Wire-up prerequisites
Every kernel in this directory expects:
FrameAllocatorfor buffer lifetime (present:mr/src/wgpu_rs/alloc.rs).dispatch_2d(n)for workgroup grid sizing whenN > 65535(present:mr/src/wgpu_rs/alloc.rs).- A graph executor that dispatches one op per compute pass
(absent:
mr/src/ir/executor.rsis a stub).
Once the executor lands, each shader moves one directory up (out of shelf/)
and gets a Rust kernels/<name>.rs binding.
Inventory
Matmul kernels
q4_matmul.wgsl— Q4_0 (32-value block, 18 bytes/block) matmul. Baseline. Sub-group reduce inside workgroup.q4k_matmul.wgsl— Q4_K (256-value superblock, 144 bytes/block) matmul. Reads raw superblocks directly on GPU, no CPU-side conversion. Inline decodes the packed 12-byte scale region using llama.cpp'sget_scale_min_k4bit extraction. Known needs-opt: superblock decode inside workgroup loop loses parallelism (audit: 13 tok/s vs Q4_0's 29).q6k_matmul.wgsl— Q6_K (256-value superblock, 210 bytes/block) matmul. Used for lm_head in most qwen models.ternary_matmul.wgsl— 2-bit packed BitNet weights ({-1, 0, +1}). Multiply-free inner loop: branches on encoding to fadd / fsub / skip. 8× less bandwidth than Q4 → memory-bound decode wins proportionally. Delivered 102 tok/s end-to-end on BitNet-2B in the old runtime.
Attention kernel
attention.wgsl— single-pass fused SDPA decode. One workgroup per head, 256 threads. Online softmax viasubgroupMax→ 8-slot shared reduction →subgroupAddfor exp sum. All in registers; no second pass over the sequence. Requiresenable subgroups;(WGSL spec-optional).
Fused ops
fused_norm_q4.wgsl— RMSNorm inlined into Q4 matmul. Normalizes into shared memory, Q4 dequant reads from shared instead of DRAM. Saves 1 dispatch + 1 global write + 1 global read per linear layer. Limited tohidden_size ≤ 4096(hardcoded shared-memory cap).fused_skip_norm.wgsl— Add + RMSNorm in one workgroup. Produces two outputs (normalized and raw residual) from one pass.
What did not make the shelf
From the old shaders/ dir, intentionally omitted:
turbo_quant.wgsl— non-atomicf32add masquerading asatomicAdd; correct only because each token gets its own workgroup. Needs a proper atomic path before shipping.q4k_embed.wgsl,q2k/q3k/q5kmatmul — less commonly used quants; copy when a model needs them.activations.wgsl,elementwise.wgsl,argmax.wgsl,rms_norm.wgsl,rope.wgsl,layernorm.wgsl,norm.wgsl,conv.wgsl,spatial.wgsl,special.wgsl,adapter.wgsl,kv_cache.wgsl,f16_matmul.wgsl,f32_matmul.wgsl— straightforward ops whose mr/ native kernels already exist or are trivial to write fresh when needed.