MSLK Kernel Library — Complete Reference

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 →

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 →

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 →

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

B: batch
M: sequence
G: head groups
H: heads
K: head dim

Variable-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

M: tokens
K: reduction dim
N: output features

Grouped 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

T: tokens
E: experts

Routing 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.

Backend	Pin it with `op=`	Dtypes	Bwd	Dropout	Varlen / paged	Reach for it when
CUTLASSNVIDIA · any compiled arch	cutlass.FwOpcutlass.BwOp	FP32 · FP16 · BF16	Yes	Yes	Mask-dependent	You need an unusual head dimension, or FP32.
CUTLASS BlackwellNVIDIA · SM100	cutlass_blackwell.FwOp…FwOpDecode · …BwOp	FP16 · BF16	Yes	No	Varlen only	On Blackwell, for the tuned prefill and decode pair.
FlashNVIDIA · SM80	flash.FwOpflash.BwOp	FP16 · BF16	Yes	Yes	Varlen · paged fwd	Default fast path for ordinary training and inference.
Flash3NVIDIA · SM80–SM90	flash3.FwOp…BwOp · …FwOp_KVSplit	FP16 · BF16 · FP8	Yes	No	Varlen · paged (fwd)	Long-context forward passes that want split-KV.
CuTe HopperNVIDIA · SM90	cute_hopper.FwOpcute_hopper.BwOp	FP16 · BF16	Yes	No	Varlen only	You want the CuTe DSL kernels on Hopper.
CuTe BlackwellNVIDIA · SM100	cute_blackwell.FwOp…FwOpDecode · …BwOp	FP16 · BF16 · FP8	Yes	No	Varlen · paged	Decoding on Blackwell against a paged KV cache.
CKAMD · supported gfx	ck.FwOpck.BwOp	FP16 · BF16	Yes	Yes	Bias-dependent	General ROCm path; the only one with bias gradients.
CK decoder / split-KAMD · supported gfx	ck_decoder.FwOpck_splitk.FwOp_S1 … _S128	FP16 · BF16 · FP32	Fwd	No	Varlen · paged	ROCm decode, or you want to fix the split count yourself.
Triton split-KNVIDIA + AMD · Triton	triton_splitk.FwOp…FwOp_S1 … _S128	FP16 · BF16 · FP8 qquantized KV	Fwd	No	Varlen · paged	Your KV cache is INT4 or FP8 — this backend reads it.
Flash MTIAMTIA · build-dependent	flash_mtia.FwOpflash_mtia.BwOp	FP16 · BF16	Yes	Yes	Varlen only	You 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.

Call	In → out	NVIDIA	AMD	Scales you must supply
f8f8bf16_rowwise…_batched · …_grouped_stacked	FP8 × FP8 → BF16	SM90–SM100 tested	gfx942, gfx950	One per row of x and of w — from `quantize_fp8_row`.
f8f8bf16_blockwise	FP8 × FP8 → BF16	SM90–SM100 tested	gfx942, gfx950	One per Bm×Bk tile; block dims are arguments.
f8f8bf16_groupwise	FP8 × FP8 → BF16	SM90–SM100 tested	gfx942, gfx950	Fixed groups of 128 along K.
f8f8f16_rowwise…_preshuffle	FP8 × FP8 → FP16	ROCm only	gfx942, gfx950	Rowwise. The FP16-output twin of the op above.
bf16bf16bf16_grouped_stacked…_cat · …_dynamic	BF16 × BF16 → BF16	SM90+ tested	gfx942 tested	None — but pass concatenated x, w`[G,N,K]` and `M_sizes`.
i8i8bf16i8i8bf16_dynamic	INT8 × INT8 → BF16	SM80+	gfx942, gfx950	One scalar (static) or a tensor scale (dynamic).
bf16i4bf16_rowwise…_batched	BF16 × INT4 → BF16	SM90 native	ROCm Triton	Packed w`[N,K/2]` plus group scale and zero point.
bf16i4bf16_shuffledf8i4bf16_shuffled	BF16 / FP8 × INT4 → BF16	SM90 exactly	Not exposed	As above, after running `preshuffle_i4` on the weights once.
f4f4bf16…_grouped_mm · …_grouped_stacked	FP4 × FP4 → BF16	SM100+	gfx950	One 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_mm	FP4 × FP4 → BF16	SM10.3+, CUDA 13+	No	Offset-grouped NVFP4, with separate global scales per operand.
mx8mx4bf16mx8mx4/mx8mx8…_grouped_mm	MXFP8 × MXFP4 → BF16	SM100+	gfx950	E8M0 block exponents. Layout differs by platform; ROCm MX8×MX4 is hybrid.
mx8mx6bf16mx6mx6bf16	MXFP8 / MXFP6 × MXFP6 → BF16	SM100+	No	Block exponents, with four E2M3 values packed into three bytes.
bf16x9_gemm	FP32 × FP32 → FP32	CUDA 13+	No	None — cuBLAS emulates FP32 with nine BF16 products.
mixed_input_gemmmslk.gemm.blackwell_mixed_input_gemm	INT4 / INT8 × BF16 / FP16	SM100	No	CuTe 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 here

Included: 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.

f8f8bf16: FP8 activations × FP8 weights with BF16 output.
bf16i4bf16: BF16 activations × packed INT4 weights with BF16 output.
rowwise: One quantization scale per logical row, commonly M for activations and N for weights.
blockwise: Scales cover rectangular M×K or N×K tiles; block dimensions are explicit parameters.
groupwise: Scales cover fixed-width groups along K, often 128 values.
MXFP4 / MXFP8: Microscaling formats: small element encodings paired with shared E8M0-style block exponents.
grouped_stacked: Groups share stacked storage and a tensor of per-group row counts (M_sizes).
grouped_mm: Groups are described by offsets into concatenated token storage, usually with a leading expert axis for weights.
preshuffle: Weights and/or scales are reordered once into the layout expected by a specialized kernel.
varlen: Variable-length sequences packed together with cumulative sequence-length metadata.
paged KV: K/V cache rows are addressed through a block table and fixed page size instead of contiguous sequence storage.
split-K / split-KV: Parallelize 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.

MSLK	PyTorch	Python	CUDA	Compiled CUDA archs	ROCm	Compiled ROCm archs
1.3.0	2.13.x	3.10–3.14	13.0, 13.2	8.0, 9.0a, 10.0a, 12.0a	7.1, 7.2	gfx942
1.2.0	2.12.x	3.10–3.14	13.0, 13.2	8.0, 9.0a, 10.0a, 12.0a	7.1, 7.2	gfx942
1.1.0	2.11.x	3.10–3.14	12.6–13.0	8.0, 9.0a, 10.0a, 12.0a	7.0, 7.1	gfx908, gfx90a, gfx942, gfx950
1.0.0	2.10.x	3.10–3.14	12.6–13.0	8.0, 9.0a, 10.0a, 12.0a	7.1, 7.2	gfx908, gfx90a, gfx942, gfx950

Metadata noteThe 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.