MSLK Kernel Reference (Website)

MSLK Kernel Reference organizes GPU kernel APIs and usage paths across 7 domains, including Transformer attention, linear layers, and sparse models. Automatic dispatch is the default, but the guide recommends specifying a backend only when specific architecture, paged KV, split-K, or deterministic execution is required.

Installation examples provide wheels for CUDA 13.0 and ROCm 7.1. After registering operators with import mslk, you can load domain modules such as mslk.attention, mslk.gemm, and mslk.moe to use torch.ops.mslk entries.

Key examples include the following.

• Causal attention using memory_efficient_attention
• BF16 GEMM execution after FP8 row-wise quantization
• Partial attention across multiple K/V chunks with exact softmax merging
• MoE routing based on index_shuffling and gather/scatter

GEMM generally takes activations [M,K] and weights [N,K] to compute X @ W.T, and the quantization scale layout must match the quantization method used to generate it. Throughout the documentation, along with the common rule that the last dimension's stride must be 1, attention, GEMM, and MoE tensor shapes are explained, providing reference information that can expand to roughly 400 API entries.