MSLK kernel reference (Website)

# MSLK Kernel Library — Complete Reference Source: [https://lucasb.eyer.be/lab/mslk/](https://lucasb.eyer.be/lab/mslk/) Generated from source · CUDA \+ ROCm MSLK is a library of fused GPU kernels for transformer workloads: attention, low\-precision GEMM, quantization, MoE routing, and convolution\. Most of it is reached through`torch\.ops\.mslk\.\*`after`import mslk`\. This page documents every public surface and, more importantly, tells you which one to call\. LatestMSLK 1\.3\.0 · PyTorch 2\.13NVIDIACUDA 13\.0 / 13\.2 · SM80 · 90a · 100a · 120aAMDROCm 7\.1 / 7\.2 · gfx942Python3\.10 – 3\.14 Kernel map ## Seven domains\. Pick the workload family you are working on — each card opens the reference filtered to that domain\. Choose a route ## The shortest path to the right kernel\. 01 · TRANSFORMER CORE### Fused attention Begin with automatic dispatch\. Reach for explicit backends only when you need a specific architecture, paged KV layout, split\-K, or deterministic behavior\. [Start with memory\_efficient\_attention →](https://lucasb.eyer.be/lab/mslk/#api-memory-efficient-attention) 02 · LINEAR LAYERS### Quantize, then GEMM Choose a scale granularity that matches the GEMM family: tensor, row, block, group, MXFP4, NVFP4, or packed INT4\. [See the FP8 pipeline →](https://lucasb.eyer.be/lab/mslk/#quickstart) 03 · SPARSE MODELS### Route, gather, compute, scatter MSLK exposes the routing pieces independently and also includes baseline and Meta\-shuffling MoE layers for composed execution\. [Open the routing API →](https://lucasb.eyer.be/lab/mslk/#api-index-shuffling) Quick start ## From install to output\. These are deliberately small, copyable paths through the major public surfaces\. ``` # CUDA 13.0 wheel pip install mslk --index-url https://download.pytorch.org/whl/cu130 # ROCm 7.1 wheel pip install mslk \ --index-url https://download.pytorch.org/whl/rocm7.1/ \ --extra-index-url https://pypi.org/simple ``` ### Import registers operators `import mslk`loads`mslk\.so`\. Import a domain such as`mslk\.gemm`or`mslk\.moe`before calling its`torch\.ops\.mslk`entries so Python\-side registrations are installed\. ``` import torch from mslk.attention import fmha B, M, H, K = 2, 2048, 32, 128 q = torch.randn(B, M, H, K, device="cuda", dtype=torch.bfloat16) k = torch.randn_like(q) v = torch.randn_like(q) out = fmha.memory_efficient_attention( q, k, v, attn_bias=fmha.LowerTriangularMask(), ) # out: [B, M, H, K] ``` ### Let dispatch work Automatic dispatch evaluates the input dtype, head dimension, mask, dropout, gradient requirements, and hardware\. Supply`op=\(FwOp, BwOp\)`only when deliberately pinning a backend\. ``` import torch import mslk.gemm from mslk.quantize.triton.fp8_quantize import quantize_fp8_row x = torch.randn(1024, 4096, device="cuda", dtype=torch.bfloat16) w = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16) xq, x_scale = quantize_fp8_row(x) wq, w_scale = quantize_fp8_row(w) out = torch.ops.mslk.f8f8bf16_rowwise( xq, wq, x_scale, w_scale ) # Conceptually: dequant(xq) @ dequant(wq).T → BF16 ``` ### Weights are N × K Most MSLK GEMMs take activations`\[M,K\]`and weights`\[N,K\]`, then compute`X @ W\.T`\. Keep scale layout paired with the quantizer that produced it\. ``` import torch from mslk.attention.fmha.merge_training import ( memory_efficient_attention_partial_autograd, merge_attentions_autograd, ) B, Mq, Mkv, H, K = 1, 128, 1024, 16, 128 q = torch.randn(B, Mq, H, K, device="cuda", dtype=torch.bfloat16) k = torch.randn(B, Mkv, H, K, device="cuda", dtype=torch.bfloat16) v = torch.randn_like(k) k0, k1 = k.chunk(2, dim=1) v0, v1 = v.chunk(2, dim=1) p0 = memory_efficient_attention_partial_autograd(q, k0, v0) p1 = memory_efficient_attention_partial_autograd(q, k1, v1) out = merge_attentions_autograd(p0, p1) ``` ### Exact softmax merge Each partial carries its output and log\-sum\-exp\. The merge reweights chunks mathematically, so it is equivalent to attention over the concatenated K/V sequence\. ``` import torch import mslk.moe T, D, E = 256, 512, 8 x = torch.randn(T, D, device="cuda", dtype=torch.bfloat16) routing_scores = torch.softmax( torch.randn(T, E, device="cuda"), dim=-1 ) counts, experts, tokens = torch.ops.mslk.index_shuffling( routing_scores, top_k=1 ) expert_x = mslk.moe.gather_scale_dense_tokens( x, tokens, experts, routing_scores ) # Replace this identity with grouped expert GEMMs + silu_mul. expert_y = expert_x out = torch.zeros_like(x) mslk.moe.scatter_add_dense_tokens(out, expert_y, tokens) ``` ### Composable routing pieces The low\-level API makes data movement explicit\. For a composed module, use`BaselineMoE`or the top\-1\-only`MetaShufflingMoE`\. Conventions ## How to read a shape in these docs\. Every signature on this page describes tensors with the same axis letters\. Learn them once and the 400 entries below stop needing individual explanation\. One rule holds nearly everywhere:**the last dimension must have stride 1**, even when the others are non\-contiguous\. ### Attention **q, k, v***\[B, M, H, K\]**the usual case* **q, k, v***\[B, M, G, H, K\]**GQA/MQA, experimental — you expand K/V yourself* **out***\[B, M, H, Kv\]**same layout as q, last axis from V* BbatchMsequenceGhead groupsHheadsKhead dimVariable\-length batches are packed into`B=1`with sequence metadata carried by the mask instead\. ### GEMM **x***\[M, K\]**activations* **w***\[N, K\]**weights — N×K, not K×N* **out***\[M, N\]**computed as x @ w\.T* MtokensKreduction dimNoutput featuresGrouped variants keep this layout and add a group description: a list of tensors, a leading expert axis,`M\_sizes`alongside concatenated tokens, or offsets\. Output is BF16 unless the op name says`f16`\. ### Quantization scales **tensorwise***\[1\]**one scale for the whole tensor* **rowwise***\[M\]**one per row — the common FP8 path* **blockwise***\[⌈M/Bm⌉, ⌈K/Bk⌉\]**one per Bm×Bk tile* A quantized tensor is the packed data*plus*its scales — a GEMM only accepts the granularity it was written for, so keep each scale tensor with the quantizer that produced it\. MX formats add E8M0 block exponents whose layout differs between CUDA and ROCm; those buffers are not interchangeable\. ### MoE routing **scores***\[T, E\]**router output* **indices***\[T × top\_k\]**token and expert index pairs* **counts***\[E \+ 2\]**tokens routed per expert* TtokensEexpertsRouting order stays explicit rather than hidden inside a fused layer, which is what lets the expert GEMM run as one grouped call over contiguous segments\. Support matrices ## What runs where\. Two orientation maps for the choices that actually block you: which attention backend can serve your case, and which GEMM op matches the dtypes you already have\. Both are read from source, and neither replaces the runtime checks — exact shapes, masks, and the archs compiled into your wheel still decide\. ### Attention — backends behind`memory\_efficient\_attention` This table is**only about attention**\. Every row is a forward/backward operator class under`mslk\.attention\.fmha`\. Leave`op=None`and dispatch picks one for you; pass`op=\(FwOp, BwOp\)`when you need a specific one\. GEMM, MoE and quantization ops do not dispatch through this\. BackendPin it with`op=`DtypesBwdDropoutVarlen / pagedReach for it whenCUTLASSNVIDIA · any compiled archcutlass\.FwOpcutlass\.BwOpFP32 · FP16 · BF16YesYesMask\-dependentYou need an unusual head dimension, or FP32\.CUTLASS BlackwellNVIDIA · SM100cutlass\_blackwell\.FwOp…FwOpDecode · …BwOpFP16 · BF16YesNoVarlen onlyOn Blackwell, for the tuned prefill and decode pair\.FlashNVIDIA · SM80flash\.FwOpflash\.BwOpFP16 · BF16YesYesVarlen · paged fwdDefault fast path for ordinary training and inference\.Flash3NVIDIA · SM80–SM90flash3\.FwOp…BwOp · …FwOp\_KVSplitFP16 · BF16 · FP8YesNoVarlen · paged \(fwd\)Long\-context forward passes that want split\-KV\.CuTe HopperNVIDIA · SM90cute\_hopper\.FwOpcute\_hopper\.BwOpFP16 · BF16YesNoVarlen onlyYou want the CuTe DSL kernels on Hopper\.CuTe BlackwellNVIDIA · SM100cute\_blackwell\.FwOp…FwOpDecode · …BwOpFP16 · BF16 · FP8YesNoVarlen · pagedDecoding on Blackwell against a paged KV cache\.CKAMD · supported gfxck\.FwOpck\.BwOpFP16 · BF16YesYesBias\-dependentGeneral ROCm path; the only one with bias gradients\.CK decoder / split\-KAMD · supported gfxck\_decoder\.FwOpck\_splitk\.FwOp\_S1 … \_S128FP16 · BF16 · FP32FwdNoVarlen · pagedROCm decode, or you want to fix the split count yourself\.Triton split\-KNVIDIA \+ AMD · Tritontriton\_splitk\.FwOp…FwOp\_S1 … \_S128FP16 · BF16 · FP8 qquantized KVFwdNoVarlen · pagedYour KV cache is INT4 or FP8 — this backend reads it\.Flash MTIAMTIA · build\-dependentflash\_mtia\.FwOpflash\_mtia\.BwOpFP16 · BF16YesYesVarlen onlyYou are running on MTIA\. ### GEMM — picking a low\-precision op Nothing here is auto\-selected:**you call the op that matches the dtypes you already hold**, and it is your job to hand it scales in the exact granularity it expects\. Names encode the contract —`f8f8bf16\_rowwise`is FP8 in × FP8 in → BF16 out, one scale per row\. An op always exists after import; it raises at call time if your wheel has no kernel for the arch\. CallIn → outNVIDIAAMDScales you must supplyf8f8bf16\_rowwise…\_batched · …\_grouped\_stackedFP8 × FP8 → BF16SM90–SM100 testedgfx942, gfx950One per row of x and of w — from`quantize\_fp8\_row`\.f8f8bf16\_blockwiseFP8 × FP8 → BF16SM90–SM100 testedgfx942, gfx950One per Bm×Bk tile; block dims are arguments\.f8f8bf16\_groupwiseFP8 × FP8 → BF16SM90–SM100 testedgfx942, gfx950Fixed groups of 128 along K\.f8f8f16\_rowwise…\_preshuffleFP8 × FP8 → FP16ROCm onlygfx942, gfx950Rowwise\. The FP16\-output twin of the op above\.bf16bf16bf16\_grouped\_stacked…\_cat · …\_dynamicBF16 × BF16 → BF16SM90\+ testedgfx942 testedNone — but pass concatenated x, w`\[G,N,K\]`and`M\_sizes`\.i8i8bf16i8i8bf16\_dynamicINT8 × INT8 → BF16SM80\+gfx942, gfx950One scalar \(static\) or a tensor scale \(dynamic\)\.bf16i4bf16\_rowwise…\_batchedBF16 × INT4 → BF16SM90 nativeROCm TritonPacked w`\[N,K/2\]`plus group scale and zero point\.bf16i4bf16\_shuffledf8i4bf16\_shuffledBF16 / FP8 × INT4 → BF16SM90 exactlyNot exposedAs above, after running`preshuffle\_i4`on the weights once\.f4f4bf16…\_grouped\_mm · …\_grouped\_stackedFP4 × FP4 → BF16SM100\+gfx950One op for three formats — NVFP4, MXFP4 or MXFP4\-16 is selected by the scales you pass\. MXFP4\-16 and NVFP4 are CUDA\-only\.f4f4bf16\_ultra\_grouped\_mmFP4 × FP4 → BF16SM10\.3\+, CUDA 13\+NoOffset\-grouped NVFP4, with separate global scales per operand\.mx8mx4bf16mx8mx4/mx8mx8…\_grouped\_mmMXFP8 × MXFP4 → BF16SM100\+gfx950E8M0 block exponents\. Layout differs by platform; ROCm MX8×MX4 is hybrid\.mx8mx6bf16mx6mx6bf16MXFP8 / MXFP6 × MXFP6 → BF16SM100\+NoBlock exponents, with four E2M3 values packed into three bytes\.bf16x9\_gemmFP32 × FP32 → FP32CUDA 13\+NoNone — cuBLAS emulates FP32 with nine BF16 products\.mixed\_input\_gemmmslk\.gemm\.blackwell\_mixed\_input\_gemmINT4 / INT8 × BF16 / FP16SM100NoCuTe DSL kernel: one narrow operand against one wide operand\. ## Complete reference Filter by symbol, concept \(“paged”, “rowwise”\), module, dtype or platform\. Open an entry for its signature, behavior, support contract, caveats, and source\. What counts as documented hereIncluded: exported and directly callable Python APIs, registered dispatcher schemas, selectable backend classes, integration\-level raw ops, and published C\+\+ headers\. Excluded: underscore\-only kernel bodies, Meta/fake implementations, benchmarks, and test\-only reference routines — unless they expose a documented integration contract\. Read before shipping ## Sharp edges worth knowing\. These are current source\-level caveats, not generic GPU advice\. ### Registration is import\-driven `import mslk`loads the consolidated native library\. Domain Meta and Python implementations appear only after their modules are imported\. On ROCm specifically, import`mslk\.gemm\.triton\.int4\_gemm`or`int8\_gemm`before the matching`torch\.ops\.mslk`calls;`mslk\.gemm`does not currently import those two for you\. ### Two MoE schemas deserve verification The current gather/quant and fused SiLU/quant paths have multi\-output implementations whose Python\-side registration has historically diverged from declared return schemas\. Validate`gather\_scale\_quant\_dense\_tokens`and`silu\_mul\_quant`against your installed build before tracing or exporting\. ### Flash\-attn varlen export gap Tests reference`flash\_attn\_varlen\_func`, but`mslk\.attention\.flash\_attn\.\_\_init\_\_`currently exports only`flash\_attn\_func`; the varlen export is commented out\. Prefer fMHA masks or the FlyDSL varlen entry point unless your build adds it\. ### Python\-only mode is not a CPU kernel build `MSLK\_PYTHON\_ONLY=1`skips native compilation so Python and Triton code can be inspected or tested\. It does not make CUDA/ROCm native kernels available on CPU\. ### GQA broadcasting is explicit The high\-level attention API accepts`\[B,M,G,H,K\]`, but does not automatically broadcast K/V heads\. Reshape and expand K/V yourself; backward support for 5D and partial paths is more restricted\. ### Architecture names are contracts Blackwell means SM100\-class compiled targets; Hopper CuTe means SM90; FlyDSL flash attention targets ROCm with architecture\-sensitive fast paths\. A matching dtype is not enough if the binary or DSL target is absent\. Glossary ## Decode the names\. f8f8bf16FP8 activations × FP8 weights with BF16 output\.bf16i4bf16BF16 activations × packed INT4 weights with BF16 output\.rowwiseOne quantization scale per logical row, commonly M for activations and N for weights\.blockwiseScales cover rectangular M×K or N×K tiles; block dimensions are explicit parameters\.groupwiseScales cover fixed\-width groups along K, often 128 values\.MXFP4 / MXFP8Microscaling formats: small element encodings paired with shared E8M0\-style block exponents\.grouped\_stackedGroups share stacked storage and a tensor of per\-group row counts \(`M\_sizes`\)\.grouped\_mmGroups are described by offsets into concatenated token storage, usually with a leading expert axis for weights\.preshuffleWeights and/or scales are reordered once into the layout expected by a specialized kernel\.varlenVariable\-length sequences packed together with cumulative sequence\-length metadata\.paged KVK/V cache rows are addressed through a block table and fixed page size instead of contiguous sequence storage\.split\-K / split\-KVParallelize reduction across K or the KV sequence, then combine partial outputs\. Environment ## Install and build matrix\. Wheel compatibility follows PyTorch releases\. Native kernel availability still depends on the architecture compiled into that wheel\. MSLKPyTorchPythonCUDACompiled CUDA archsROCmCompiled ROCm archs1\.3\.02\.13\.x3\.10–3\.1413\.0, 13\.28\.0, 9\.0a, 10\.0a, 12\.0a7\.1, 7\.2gfx9421\.2\.02\.12\.x3\.10–3\.1413\.0, 13\.28\.0, 9\.0a, 10\.0a, 12\.0a7\.1, 7\.2gfx9421\.1\.02\.11\.x3\.10–3\.1412\.6–13\.08\.0, 9\.0a, 10\.0a, 12\.0a7\.0, 7\.1gfx908, gfx90a, gfx942, gfx9501\.0\.02\.10\.x3\.10–3\.1412\.6–13\.08\.0, 9\.0a, 10\.0a, 12\.0a7\.1, 7\.2gfx908, gfx90a, gfx942, gfx950 **Metadata note**The table above follows the current README release contract\.`setup\.py`classifiers still list Python 3\.9–3\.13 while the README lists 3\.10–3\.14, and the setup URL still names the older`pytorch/MSLK`path\. Treat release wheels and the README as the practical compatibility authority\. ### Build CUDA `\./ci/integration/mslk\_oss\_build\.bash`creates a conda environment\. Activate it, then iterate with`python setup\.py install`\. ### Build ROCm Set`BUILD\_VARIANT=rocm`, a matching`BUILD\_ROCM\_VERSION`, and`PYTORCH\_ROCM\_ARCH`\(for example`gfx942`\) when invoking the build\.