# MSLK - kernel library reference for LLM agents > Fused GPU kernels for transformer workloads: attention, low-precision GEMM, quantization, MoE routing, convolution. This file is the whole public surface in one page: read the orientation, then find your symbol in the index. Full per-symbol detail (behavior, constraints, methods) is in llms-full.txt; the same catalog as data is in api.json. Source: https://github.com/meta-pytorch/MSLK @ 69ae1b897f2d546d72acab129e1ae4bc2924900f. Generated from docs/index.html by docs/build-llms-txt.mjs - do not edit by hand. ## What this library is 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. - Latest: MSLK 1.3.0 - PyTorch 2.13 - NVIDIA: CUDA 13.0 / 13.2 - SM80 - 90a - 100a - 120a - AMD: ROCm 7.1 / 7.2 - gfx942 - Python: 3.10 - 3.14 ## Getting a kernel to run ### Install ```bash # 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. ### Attention ```python 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. ### FP8 GEMM ```python 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 x 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. ### Split KV ```python 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. ### MoE primitives ```python 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. ## Shape conventions Every signature below uses these axis letters. 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 Axes: 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 x K, not K x N - `out` `[M, N]` - computed as x @ w.T Axes: 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` `[ceil(M/Bm), ceil(K/Bk)]` - one per Bm x 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 x top_k]` - token and expert index pairs - `counts` `[E + 2]` - tokens routed per expert Axes: 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. ## Choosing an implementation ### 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 | | --- | --- | --- | --- | --- | --- | --- | | CUTLASS (NVIDIA - any compiled arch) | cutlass.FwOp (cutlass.BwOp) | FP32 - FP16 - BF16 | Yes | Yes | Mask-dependent | You need an unusual head dimension, or FP32. | | CUTLASS Blackwell (NVIDIA - SM100) | cutlass_blackwell.FwOp (...FwOpDecode - ...BwOp) | FP16 - BF16 | Yes | No | Varlen only | On Blackwell, for the tuned prefill and decode pair. | | Flash (NVIDIA - SM80) | flash.FwOp (flash.BwOp) | FP16 - BF16 | Yes | Yes | Varlen - paged fwd | Default fast path for ordinary training and inference. | | Flash3 (NVIDIA - SM80-SM90) | flash3.FwOp (...BwOp - ...FwOp_KVSplit) | FP16 - BF16 - FP8 | Yes | No | Varlen - paged (fwd) | Long-context forward passes that want split-KV. | | CuTe Hopper (NVIDIA - SM90) | cute_hopper.FwOp (cute_hopper.BwOp) | FP16 - BF16 | Yes | No | Varlen only | You want the CuTe DSL kernels on Hopper. | | CuTe Blackwell (NVIDIA - SM100) | cute_blackwell.FwOp (...FwOpDecode - ...BwOp) | FP16 - BF16 - FP8 | Yes | No | Varlen - paged | Decoding on Blackwell against a paged KV cache. | | CK (AMD - supported gfx) | ck.FwOp (ck.BwOp) | FP16 - BF16 | Yes | Yes | Bias-dependent | General ROCm path; the only one with bias gradients. | | CK decoder / split-K (AMD - supported gfx) | ck_decoder.FwOp (ck_splitk.FwOp_S1 ... _S128) | FP16 - BF16 - FP32 | Fwd | No | Varlen - paged | ROCm decode, or you want to fix the split count yourself. | | Triton split-K (NVIDIA + AMD - Triton) | triton_splitk.FwOp (...FwOp_S1 ... _S128) | FP16 - BF16 - FP8 q (quantized KV) | Fwd | No | Varlen - paged | Your KV cache is INT4 or FP8 - this backend reads it. | | Flash MTIA (MTIA - build-dependent) | flash_mtia.FwOp (flash_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 x 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 x FP8 -> BF16 | SM90-SM100 tested | gfx942, gfx950 | One per row of x and of w - from quantize_fp8_row. | | f8f8bf16_blockwise | FP8 x FP8 -> BF16 | SM90-SM100 tested | gfx942, gfx950 | One per Bm x Bk tile; block dims are arguments. | | f8f8bf16_groupwise | FP8 x FP8 -> BF16 | SM90-SM100 tested | gfx942, gfx950 | Fixed groups of 128 along K. | | f8f8f16_rowwise (..._preshuffle) | FP8 x FP8 -> FP16 | ROCm only | gfx942, gfx950 | Rowwise. The FP16-output twin of the op above. | | bf16bf16bf16_grouped_stacked (..._cat - ..._dynamic) | BF16 x BF16 -> BF16 | SM90+ tested | gfx942 tested | None - but pass concatenated x, w[G,N,K] and M_sizes. | | i8i8bf16 (i8i8bf16_dynamic) | INT8 x INT8 -> BF16 | SM80+ | gfx942, gfx950 | One scalar (static) or a tensor scale (dynamic). | | bf16i4bf16_rowwise (..._batched) | BF16 x INT4 -> BF16 | SM90 native | ROCm Triton | Packed w[N,K/2] plus group scale and zero point. | | bf16i4bf16_shuffled (f8i4bf16_shuffled) | BF16 / FP8 x INT4 -> BF16 | SM90 exactly | Not exposed | As above, after running preshuffle_i4 on the weights once. | | f4f4bf16 (..._grouped_mm - ..._grouped_stacked) | FP4 x 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 x FP4 -> BF16 | SM10.3+, CUDA 13+ | No | Offset-grouped NVFP4, with separate global scales per operand. | | mx8mx4bf16 (mx8mx4/mx8mx8..._grouped_mm) | MXFP8 x MXFP4 -> BF16 | SM100+ | gfx950 | E8M0 block exponents. Layout differs by platform; ROCm MX8 x MX4 is hybrid. | | mx8mx6bf16 (mx6mx6bf16) | MXFP8 / MXFP6 x MXFP6 -> BF16 | SM100+ | No | Block exponents, with four E2M3 values packed into three bytes. | | bf16x9_gemm | FP32 x FP32 -> FP32 | CUDA 13+ | No | None - cuBLAS emulates FP32 with nine BF16 products. | | mixed_input_gemm (mslk.gemm.blackwell_mixed_input_gemm) | INT4 / INT8 x BF16 / FP16 | SM100 | No | CuTe DSL kernel: one narrow operand against one wide operand. | ## Sharp edges - **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. ## Naming - **f8f8bf16** - FP8 activations x FP8 weights with BF16 output. - **bf16i4bf16** - BF16 activations x 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 x K or N x 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. ## Release compatibility | 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 | ## Symbol index All 424 documented symbols, grouped. Format: name - summary (module). Platforms are shown only where support is restricted; everything else runs on both CUDA and ROCm. Signatures, constraints and source links are in llms-full.txt under the same headings. ### Attention - Core fMHA Module: `mslk.attention.fmha` - memory_efficient_attention - Autograd-enabled fused attention with automatic backend dispatch. [NVIDIA, AMD, MTIA] - memory_efficient_attention_forward - Forward-only attention path that does not retain manual-backward context. [NVIDIA, AMD, MTIA] - memory_efficient_attention_forward_requires_grad - Manual-backward forward pass returning output and log-sum-exp. [NVIDIA, AMD, MTIA] - memory_efficient_attention_backward - Explicit backward pass for an output/LSE pair. [NVIDIA, AMD, MTIA] - memory_efficient_attention_partial - Computes an output and LSE for one disjoint K/V shard. [NVIDIA, AMD, MTIA] - merge_attentions - Log-sum-exp-correct merge of attention computed over K/V chunks. [NVIDIA, AMD, MTIA] ### Attention - Compile-friendly wrappers Module: `torch.ops.mslk` - torch.ops.mslk.memory_efficient_attention_forward - Compile-friendly forward attention with optional bias. [NVIDIA, Meta / fake] - torch.ops.mslk.memory_efficient_attention_forward_with_bias - Compile-friendly forward attention requiring an explicit bias tensor. [NVIDIA, Meta / fake] ### Attention - Trainable partial attention Module: `mslk.attention.fmha.merge_training` - Partial - Autograd carrier for a partial attention output and its merge context. - memory_efficient_attention_partial_autograd - Autograd-safe wrapper around partial attention. - merge_attentions_autograd - Merges one or more Partial objects with correct gradients. ### Attention - Biases & masks Module: `mslk.attention.fmha.attn_bias` - AttentionBias - Abstract base for optimized attention-bias descriptions. [NVIDIA, AMD, MTIA, CPU / Python] - LowerTriangularMask - Top-left-aligned causal mask. [NVIDIA, AMD, MTIA, CPU / Python] - LowerTriangularMaskWithTensorBias - Top-left causal mask plus a dense additive tensor bias. [NVIDIA, AMD, MTIA, CPU / Python] - LowerTriangularFromBottomRightMask - Bottom-right-aligned causal mask for Mq != Mkv. [NVIDIA, AMD, MTIA, CPU / Python] - LowerTriangularFromBottomRightLocalAttentionMask - Bottom-right causal sliding-window mask. [NVIDIA, AMD, MTIA, CPU / Python] - LocalAttentionFromBottomRightMask - Non-causal local window aligned from the bottom right. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalMask - Packed independent sequences where each Q block sees its matching K/V block. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalCausalMask - Top-left causal mask inside every packed sequence block. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalCausalFromBottomRightMask - Bottom-right causal mask inside every packed block. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalCausalLocalAttentionMask - Top-left causal local window inside packed blocks. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalCausalLocalAttentionFromBottomRightMask - Bottom-right causal local window inside packed blocks. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalPaddedKeysMask - Packed Q blocks attending fixed-capacity padded K/V lanes. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalCausalWithOffsetPaddedKeysMask - Bottom-right causal masking over padded K/V lanes. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalLocalAttentionPaddedKeysMask - Non-causal local windows over padded K/V lanes. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalCausalLocalAttentionPaddedKeysMask - Bottom-right causal local windows over padded K/V lanes. [NVIDIA, AMD, MTIA, CPU / Python] - PagedBlockDiagonalPaddedKeysMask - Padded logical K/V lanes mapped to physical cache pages. [NVIDIA, AMD, MTIA, CPU / Python] - PagedBlockDiagonalCausalWithOffsetPaddedKeysMask - Bottom-right causal paged padded mask. [NVIDIA, AMD, MTIA, CPU / Python] - PagedBlockDiagonalCausalLocalPaddedKeysMask - Bottom-right causal local paged padded mask. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalGappyKeysMask - K/V lanes with arbitrary starts and independently used lengths. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalLocalAttentionFromBottomRightGappyKeysMask - Bottom-right-aligned local windows over gappy K/V lanes. [NVIDIA, AMD, MTIA, CPU / Python] - BlockDiagonalCausalWithOffsetGappyKeysMask - Bottom-right causal masking over gappy K/V lanes. [NVIDIA, AMD, MTIA, CPU / Python] - PagedBlockDiagonalGappyKeysMask - Gappy logical K/V lanes mapped through a page table. [NVIDIA, AMD, MTIA, CPU / Python] - PagedBlockDiagonalCausalWithOffsetGappyKeysMask - Bottom-right causal paged-gappy mask. [NVIDIA, AMD, MTIA, CPU / Python] - VARLEN_BIASES - Dispatch grouping for packed, padded, gappy, and paged variable-length masks. [NVIDIA, AMD, MTIA, CPU / Python] ### Attention - Backend operators - cutlass.FwOp - Classic CUTLASS forward attention with the broadest NVIDIA dtype/head-dim range. (mslk.attention.fmha.cutlass) [NVIDIA SM60-SM90] - cutlass.BwOp - Classic CUTLASS attention backward, including supported tensor-bias gradients. (mslk.attention.fmha.cutlass) [NVIDIA SM60-SM90] - cutlass_blackwell.FwOp - Blackwell CUTLASS prefill/varlen forward backend. (mslk.attention.fmha.cutlass_blackwell) [NVIDIA SM100+] - cutlass_blackwell.FwOpDecode - Blackwell inference/decode forward backend. (mslk.attention.fmha.cutlass_blackwell) [NVIDIA SM100+] - cutlass_blackwell.BwOp - Blackwell CUTLASS backward backend. (mslk.attention.fmha.cutlass_blackwell) [NVIDIA SM100+] - flash.FwOp - FlashAttention 2 forward backend for the general NVIDIA fast path. (mslk.attention.fmha.flash) [NVIDIA SM80+] - flash.BwOp - FlashAttention 2 backward backend. (mslk.attention.fmha.flash) [NVIDIA SM80+] - flash3.FwOp - FlashAttention 3 forward with FP8 and flexible V-dimension support. (mslk.attention.fmha.flash3) [NVIDIA SM80-SM90] - flash3.BwOp - FlashAttention 3 backward backend. (mslk.attention.fmha.flash3) [NVIDIA SM80-SM90] - flash3.FwOp_KVSplit - FlashAttention 3 split-KV forward variant. (mslk.attention.fmha.flash3) [NVIDIA SM80-SM90] - ck.FwOp - Composable Kernel attention forward for ROCm. (mslk.attention.fmha.ck) [AMD ROCm] - ck.BwOp - Composable Kernel attention backward, including supported bias gradients. (mslk.attention.fmha.ck) [AMD ROCm] - ck_decoder.FwOp - ROCm inference decoder attention backend. (mslk.attention.fmha.ck_decoder) [AMD ROCm / MI250X] - ck_splitk.FwOp - ROCm split-K decoder/prefix attention with optional quantized K/V. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - triton_splitk.FwOp - Portable split-K forward with quantized KV-cache support. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm / MI300X] - cute_hopper.FwOp - CuTe DSL Hopper forward attention. (mslk.attention.fmha.cute_hopper) [NVIDIA SM90-SM99] - cute_hopper.BwOp - CuTe DSL Hopper backward attention. (mslk.attention.fmha.cute_hopper) [NVIDIA SM90-SM99] - cute_blackwell.FwOp - CuTe DSL Blackwell prefill forward attention. (mslk.attention.fmha.cute_blackwell) [NVIDIA SM100+] - cute_blackwell.FwOpDecode - CuTe DSL Blackwell decode/paged forward attention. (mslk.attention.fmha.cute_blackwell) [NVIDIA SM100+] - cute_blackwell.BwOp - CuTe DSL Blackwell backward attention. (mslk.attention.fmha.cute_blackwell) [NVIDIA SM100+] - flash_mtia.FwOp - MTIA counterpart of the FlashAttention 2 forward protocol. (mslk.attention.fmha.flash_mtia) [MTIA] - flash_mtia.BwOp - MTIA counterpart of the FlashAttention 2 backward protocol. (mslk.attention.fmha.flash_mtia) [MTIA] ### Attention - Fixed split-K backends - ck_splitk.FwOp_S1 - CK split-K forward fixed to 1 split. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - ck_splitk.FwOp_S2 - CK split-K forward fixed to 2 splits. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - ck_splitk.FwOp_S4 - CK split-K forward fixed to 4 splits. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - ck_splitk.FwOp_S8 - CK split-K forward fixed to 8 splits. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - ck_splitk.FwOp_S16 - CK split-K forward fixed to 16 splits. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - ck_splitk.FwOp_S32 - CK split-K forward fixed to 32 splits. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - ck_splitk.FwOp_S64 - CK split-K forward fixed to 64 splits. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - ck_splitk.FwOp_S128 - CK split-K forward fixed to 128 splits. (mslk.attention.fmha.ck_splitk) [AMD ROCm] - triton_splitk.FwOp_S1 - Triton split-K forward fixed to 1 split. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] - triton_splitk.FwOp_S2 - Triton split-K forward fixed to 2 splits. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] - triton_splitk.FwOp_S4 - Triton split-K forward fixed to 4 splits. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] - triton_splitk.FwOp_S8 - Triton split-K forward fixed to 8 splits. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] - triton_splitk.FwOp_S16 - Triton split-K forward fixed to 16 splits. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] - triton_splitk.FwOp_S32 - Triton split-K forward fixed to 32 splits. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] - triton_splitk.FwOp_S64 - Triton split-K forward fixed to 64 splits. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] - triton_splitk.FwOp_S128 - Triton split-K forward fixed to 128 splits. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] ### Attention - Preassembled operator pairs - MemoryEfficientAttentionCutlassOp - Classic CUTLASS forward and backward. (mslk.attention.fmha) [NVIDIA, AMD, MTIA] - MemoryEfficientAttentionCutlassBlackwellOp - Blackwell CUTLASS forward and backward. (mslk.attention.fmha) [NVIDIA, AMD, MTIA] - MemoryEfficientAttentionCutlassFwdFlashBwOp - Classic CUTLASS forward with FlashAttention 2 backward. (mslk.attention.fmha) [NVIDIA, AMD, MTIA] - MemoryEfficientAttentionFlashAttentionOp - FlashAttention 2 pair. (mslk.attention.fmha) [NVIDIA, AMD, MTIA] - MemoryEfficientAttentionFlashMtiaAttentionOp - MTIA Flash pair. (mslk.attention.fmha) [MTIA] - MemoryEfficientAttentionCkOp - ROCm CK pair. (mslk.attention.fmha) [AMD ROCm] - MemoryEfficientAttentionCkDecoderOp - Inference-only convenience tuple using CK decoder forward. (mslk.attention.fmha) [AMD ROCm] - MemoryEfficientAttentionSplitKCkOp - Inference-only convenience tuple using CK split-K forward. (mslk.attention.fmha) [AMD ROCm] - MemoryEfficientAttentionCuteFlashAttentionOp - CuTe Blackwell forward/backward pair. (mslk.attention.fmha) [NVIDIA SM100+] - ALL_FW_OPS - Explicit forward-backend enumeration used by tests and introspection. (mslk.attention.fmha) [NVIDIA, AMD, MTIA] - ALL_BW_OPS - Explicit backward-backend enumeration used by tests and introspection. (mslk.attention.fmha) [NVIDIA, AMD, MTIA] - triton_splitk.FwOp_Map - Maps supported split counts to specialized Triton split-K forward classes. (mslk.attention.fmha.triton_splitk) [NVIDIA SM80+, AMD ROCm] ### Attention - Operator protocol & inputs Module: `mslk.attention.fmha.common` - AttentionOp - Forward/backward operator pair type accepted by high-level fMHA. [NVIDIA, AMD, MTIA, CPU / Python] - AttentionOpBase - Shared backend metadata, availability, and support-check protocol. [NVIDIA, AMD, MTIA, CPU / Python] - AttentionFwOpBase - Base protocol for forward attention operators. [NVIDIA, AMD, MTIA, CPU / Python] - AttentionBwOpBase - Base protocol for backward attention operators. [NVIDIA, AMD, MTIA, CPU / Python] - Inputs - Normalized fMHA input descriptor used by dispatch and backends. [NVIDIA, AMD, MTIA, CPU / Python] - InputsFp8 - Input descriptor for row/group quantized FP8-style K/V caches. [NVIDIA, AMD, MTIA, CPU / Python] - InputsMXFp8 - Input descriptor for MXFP8 K/V with E8M0 block scales. [NVIDIA, AMD, MTIA, CPU / Python] - Context - Forward context consumed by an attention backward operator. [NVIDIA, AMD, MTIA, CPU / Python] - Gradients - Normalized attention-gradient bundle. [NVIDIA, AMD, MTIA, CPU / Python] - ScaledTensor - Tensor subclass carrying separate quantization-scale metadata. [NVIDIA, AMD, MTIA, CPU / Python] - pack_fp8_tensorwise_per_head - Wraps tensorwise/per-head FP8 storage and scale metadata for backend input. [NVIDIA, AMD, MTIA, CPU / Python] - bmghk2bmhk - Flattens grouped heads G x H into a single H axis, with optional scale handling. [NVIDIA, AMD, MTIA, CPU / Python] - bmhk2bhk - Converts BMHK storage to the backend BHK-style view for single-query/decode paths. [NVIDIA, AMD, MTIA, CPU / Python] - bmk2bmhk - Adds/reshapes a head axis for BMK tensors. [NVIDIA, AMD, MTIA, CPU / Python] - check_lastdim_alignment_stride1 - Appends backend support failures for last-dimension alignment and stride. [NVIDIA, AMD, MTIA, CPU / Python] ### Attention - Dispatch controls - _get_use_fa3 - Returns whether CUDA automatic dispatch is permitted to select FlashAttention 3. (mslk.attention.fmha.dispatch) [NVIDIA, CPU / Python] - _set_use_fa3 - Enables or disables FlashAttention 3 in automatic dispatch. (mslk.attention.fmha.dispatch) [NVIDIA, CPU / Python] - fa3_available - Reports whether a compatible FlashAttention 3 implementation is available. (mslk.attention.fmha.dispatch) [NVIDIA, CPU / Python] - is_pt_cutlass_compatible - Checks PyTorch's bundled CUTLASS attention compatibility. (mslk.attention.fmha.torch_attention_compat) [NVIDIA, CPU / Python] - ensure_pt_flash_ok - Raises when the active PyTorch FlashAttention integration is incompatible. (mslk.attention.fmha.torch_attention_compat) [NVIDIA, CPU / Python] ### Attention - Tree & speculative attention Module: `mslk.attention.fmha.tree_attention` - TreeAttnMetadata - Builds explicit mask and traversal metadata for a speculative token tree. - tree_attention - Computes prefix-cache and speculative-tree attention, then LSE-merges them. - use_triton_splitk_for_prefix - Heuristic for choosing Triton split-K on the tree prefix. - select_prefix_op - Selects the prefix forward operator for tree attention. - SplitKAutotune - Tree-attention split-K operator variant with runtime autotuning. - construct_full_tree_choices - Generates node choices for a regular full tree. - construct_tree_choices - Generates node choices for per-level branching factors. - get_full_tree_size - Returns sum of branching^i over range(tree_depth). ### Attention - iRoPE & page transforms Module: `mslk.attention.fmha.split_blocks_fairinternal` - split_blocks_for_decoding_gpu_part - Computes GPU-side sequence-start/length data for decoding block splits. - split_blocks_for_decoding - Transforms padded decoding lanes into gappy or paged-gappy masks. - split_blocks_for_prefill - Re-batches padded prefill lanes into a split block description. - maybe_make_paged - Converts supported padded/gappy masks to paged counterparts when a block table is supplied. ### Attention - ROCm MLA - mla_decode_fwd - Paged MLA decode forward for one query position per sequence. (mslk.attention.mla) [AMD gfx942, AMD gfx950] - mla_prefill_fwd - Packed variable-length MLA prefill forward. (mslk.attention.mla) [AMD gfx942, AMD gfx950] - MLA_NUM_HEADS - Fixed MLA query-head count. (mslk.attention.mla.triton_mla) [AMD gfx942, AMD gfx950] - MLA_NUM_KV_HEADS - Fixed latent KV-head count. (mslk.attention.mla.triton_mla) [AMD gfx942, AMD gfx950] - MLA_KV_LORA_RANK - Latent KV LoRA rank. (mslk.attention.mla.triton_mla) [AMD gfx942, AMD gfx950] - MLA_QK_ROPE_HEAD_DIM - RoPE portion of the Q/K head. (mslk.attention.mla.triton_mla) [AMD gfx942, AMD gfx950] - MLA_QK_HEAD_DIM - Full Q/K head dimension after absorption. (mslk.attention.mla.triton_mla) [AMD gfx942, AMD gfx950] - MLA_V_HEAD_DIM - MLA value/output head dimension. (mslk.attention.mla.triton_mla) [AMD gfx942, AMD gfx950] ### Attention - FlyDSL Flash Attention - flydsl_flash_attn_func - ROCm FlyDSL flash-attention forward across dense, varlen, GQA/MQA, and paged KV. (mslk.attention.flydsl) [AMD gfx942, AMD gfx950] - dualwave_splitk_workspace_elems - Returns the FP32 element count needed by the gfx950 dual-wave split-K workspace. (mslk.attention.flydsl.flash_attn_interface) [AMD gfx942, AMD gfx950] ### Attention - Standalone facades Module: `mslk.attention.flash_attn` - flash_attn_func - Optional standalone CuTe Flash-Attention facade. [NVIDIA / dependency-defined] ### Attention - Standalone Blackwell FMHA - cutlass_blackwell_fmha_func - Autograd-capable standalone Blackwell FMHA facade. (mslk.attention.cutlass_blackwell_fmha) [NVIDIA SM100+] - cutlass_blackwell_fmha_decode_forward - Inference-only split-K decode forward with raw partial outputs. (mslk.attention.cutlass_blackwell_fmha) [NVIDIA SM100+] - _cutlass_blackwell_fmha_forward - Testing/implementation forward wrapper returning output and LSE. (mslk.attention.cutlass_blackwell_fmha) [NVIDIA SM100+] - cutlass_blackwell_fmha_custom_op - torch.library custom-op facade for standalone Blackwell FMHA. (mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_custom_op) [NVIDIA SM100+] - get_splitk_heuristic - Chooses a decode split size from cache length, KV heads, and available SMs. (mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_interface) [NVIDIA SM100+] - GenKernelType - Selects the generated Blackwell UMMA kernel family. (mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_interface) [NVIDIA SM100+] ### Attention - Raw Blackwell torch ops Module: `torch.ops.mslk` - torch.ops.mslk.fmha_fwd - Standalone Blackwell forward implementation op. [NVIDIA SM100+, Meta / fake] - torch.ops.mslk.fmha_bwd - Standalone Blackwell backward implementation op. [NVIDIA SM100+, Meta / fake] - torch.ops.mslk.fmha_gen_fwd - Generated UMMA Blackwell forward implementation op. [NVIDIA SM100+, Meta / fake] - torch.ops.mslk.cutlass_blackwell_fmha_fwd - fMHA Blackwell forward native op. [NVIDIA SM100+, Meta / fake] - torch.ops.mslk.cutlass_blackwell_fmha_bwd - fMHA Blackwell backward native op. [NVIDIA SM100+, Meta / fake] ### Attention - Merge & storage utilities - triton_splitk.merge_attentions - In-place low-level reduction of split attention into preallocated outputs. (mslk.attention.fmha.triton_splitk) - torch.ops.mslk.fmha_merge_attentions_varargs - Custom op for variable-argument partial-attention merge. (torch.ops.mslk) - torch.ops.mslk.merge_attentions_varargs_backward - Custom backward op for variable-argument attention merge. (torch.ops.mslk) - get_stack_strides - Detects whether tensors are views of one common stacked storage layout. (mslk.attention.fmha.unbind) - unbind - Autograd-aware unbind preserving shared-storage information. (mslk.attention.fmha.unbind) - stack_or_none - Returns a zero-copy/common-storage stack view when possible, else None. (mslk.attention.fmha.unbind) ### Attention - Raw backend ops Module: `torch.ops` - torch.ops.xformers.efficient_attention_forward_ck - Raw CK forward attention op. [AMD ROCm] - torch.ops.xformers.efficient_attention_backward_ck - Raw CK backward attention op. [AMD ROCm] - torch.ops.xformers.efficient_attention_forward_decoder_ck - Raw CK decoder forward op. [AMD ROCm] - torch.ops.xformers.efficient_attention_forward_decoder_splitk_ck - Raw CK split-K decoder forward op. [AMD ROCm] - torch.ops.xformers._ck_rand_uniform - CK random-uniform helper used by dropout implementation. [AMD ROCm] - torch.ops.mslk_flash.flash_fwd - Conditional bundled FlashAttention 2 forward op returning output, LSE, and RNG state. [NVIDIA SM80+] - torch.ops.mslk_flash.flash_bwd - Conditional bundled FlashAttention 2 backward op. [NVIDIA SM80+] - torch.ops.mslk_flash3.flash_fwd - Conditional bundled FlashAttention 3 forward op returning output and LSE. [NVIDIA SM80-SM90] - torch.ops.mslk_flash3.flash_bwd - Conditional bundled FlashAttention 3 backward op. [NVIDIA SM80-SM90] - torch.ops.mslk_flash_mtia.flash_fwd - Conditional MTIA Flash forward op returning output, LSE, and RNG state. [MTIA] - torch.ops.mslk_flash_mtia.flash_bwd - Conditional MTIA Flash backward op. [MTIA] ### GEMM - BF16 grouped torch ops Module: `torch.ops.mslk` - torch.ops.mslk.bf16bf16bf16_grouped - List-of-tensors BF16 grouped GEMM. - torch.ops.mslk.bf16bf16bf16_grouped_cat - List-of-tensors BF16 grouped GEMM with concatenated output. - torch.ops.mslk.bf16bf16bf16_grouped_dynamic - Dynamic-row BF16 grouped GEMM using per-group zero starts. - torch.ops.mslk.bf16bf16bf16_grouped_stacked - Preferred stacked BF16 grouped forward GEMM. - torch.ops.mslk.bf16bf16bf16_grouped_grad - Grouped BF16 data-gradient matrix product. - torch.ops.mslk.bf16bf16bf16_grouped_wgrad - Grouped BF16 weight-gradient product. ### GEMM - FP8 row/block/group torch ops Module: `torch.ops.mslk` - torch.ops.mslk.f8f8bf16_blockwise - Block-scaled FP8 x FP8 GEMM. - torch.ops.mslk.f8f8bf16_rowwise - Primary rowwise FP8 x FP8 -> BF16 GEMM. - torch.ops.mslk.f8f8bf16_rowwise_out - In-place/out rowwise FP8 GEMM. - torch.ops.mslk.f8f8bf16_rowwise_batched - Batched rowwise FP8 GEMM with optional preallocated output. - torch.ops.mslk.f8f8bf16_rowwise_grouped - TensorList grouped rowwise FP8 GEMM. - torch.ops.mslk.f8f8bf16_rowwise_grouped_cat - TensorList grouped rowwise FP8 GEMM with concatenated output. - torch.ops.mslk.f8f8bf16_rowwise_grouped_stacked - Stacked grouped rowwise FP8 GEMM. - torch.ops.mslk.f8f8bf16_rowwise_grouped_dynamic - Dynamic-row grouped FP8 GEMM. - torch.ops.mslk.f8f8bf16_groupwise - FP8 GEMM with fixed K-group scale granularity 128. - torch.ops.mslk.f8f8bf16_groupwise_grouped - Stacked grouped FP8 GEMM with fixed K-group scales. ### GEMM - ROCm FP8 torch ops Module: `torch.ops.mslk` - torch.ops.mslk.f8f8f16_rowwise - ROCm rowwise FP8 GEMM with FP16 output. [AMD ROCm] - torch.ops.mslk.f8f8bf16_rowwise_preshuffle - ROCm preshuffled rowwise FP8 GEMM with BF16 output. [AMD ROCm] - torch.ops.mslk.f8f8f16_rowwise_preshuffle - ROCm preshuffled rowwise FP8 GEMM with FP16 output. [AMD ROCm] - torch.ops.mslk.f8f8bf16_rowwise_grouped_mm - Generic ROCm grouped FP8 GEMM over 2D/3D layout combinations. [AMD ROCm] ### GEMM - FP4 & microscaling torch ops Module: `torch.ops.mslk` - torch.ops.mslk.f4f4bf16 - Packed FP4 x FP4 GEMM selecting MXFP4, MXFP4-16, or NVFP4 mode. [NVIDIA, AMD gfx950] - torch.ops.mslk.f4f4bf16_grouped_mm - Offset-described grouped FP4 GEMM. [NVIDIA, AMD gfx950] - torch.ops.mslk.f4f4bf16_grouped_stacked - Stacked grouped FP4 GEMM with optional per-segment padding metadata. [NVIDIA, AMD gfx950] - torch.ops.mslk.f4f4bf16_ultra_grouped_mm - Ultra grouped FP4 GEMM with separate activation/weight global scales. [NVIDIA SM10.3+ / CUDA 13+] - torch.ops.mslk.mx8mx4bf16 - MXFP8 activation x MXFP4 weight GEMM. [NVIDIA, AMD gfx950] - torch.ops.mslk.mx8mx4bf16_grouped_mm - Grouped MXFP8 x MXFP4 GEMM described by offsets. [NVIDIA, AMD gfx950] - torch.ops.mslk.mx8mx8bf16_grouped_mm - Grouped MXFP8 x MXFP8 GEMM. [NVIDIA, AMD gfx950] - torch.ops.mslk.mx8mx6bf16 - MXFP8 x packed MXFP6 E2M3 GEMM. [NVIDIA SM100+] - torch.ops.mslk.mx6mx6bf16 - Packed MXFP6 E2M3 x MXFP6 E2M3 GEMM. [NVIDIA SM100+] ### GEMM - INT and mixed torch ops Module: `torch.ops.mslk` - torch.ops.mslk.i8i8bf16 - INT8 x INT8 GEMM with a static scalar scale. - torch.ops.mslk.i8i8bf16_dynamic - INT8 x INT8 GEMM with a tensor-valued dynamic scale. - torch.ops.mslk.bf16i4bf16_rowwise - BF16 activation x row/group-quantized INT4 weight GEMM. - torch.ops.mslk.bf16i4bf16_rowwise_batched - Batched BF16 x rowwise INT4 weight GEMM. - torch.ops.mslk.bf16i4bf16_shuffled - BF16 x preshuffled INT4 GEMM. [NVIDIA SM90] - torch.ops.mslk.bf16i4bf16_shuffled_batched - Batched BF16 x preshuffled INT4 GEMM. [NVIDIA SM90] - torch.ops.mslk.bf16i4bf16_shuffled_grouped - Stacked grouped BF16 x preshuffled INT4 GEMM. [NVIDIA SM90] - torch.ops.mslk.f8i4bf16_rowwise - Rowwise FP8 activation x INT4 weight GEMM. [NVIDIA SM90] - torch.ops.mslk.f8i4bf16_shuffled - FP8 activation x preshuffled INT4 weight GEMM. [NVIDIA SM90] - torch.ops.mslk.f8i4bf16_shuffled_grouped - Stacked grouped FP8 x preshuffled INT4 GEMM. [NVIDIA SM90] - torch.ops.mslk.preshuffle_i4 - Preprocesses packed INT4 weights and scales for shuffled CUDA kernels. [NVIDIA] - torch.ops.mslk.bf16x9_gemm - cuBLAS BF16x9-emulation GEMM over FP32 inputs and output. [NVIDIA / CUDA 13+] ### GEMM - Triton FP8 GEMM - matmul_fp8_row - Feature-rich Triton rowwise FP8 matrix product. (mslk.gemm.triton.fp8_gemm) - matmul_fp8_block - Triton block-scaled FP8 matrix product. (mslk.gemm.triton.fp8_gemm) - to_mxfp8 - Converts high-precision data to MXFP8 with E8M0 block scales. (mslk.gemm.triton.fp8_gemm) - torch.ops.triton.matmul_fp8_row - Custom-op registration for matmul_fp8_row. (torch.ops.triton) - torch.ops.triton.matmul_fp8_block - Custom-op registration for matmul_fp8_block. (torch.ops.triton) ### GEMM - Triton grouped GEMM Module: `mslk.gemm.triton.grouped_gemm` - grouped_gemm - General BF16 grouped GEMM with optional fused bias/router scaling. - grouped_gemm_fp8_rowwise - Rowwise FP8 grouped GEMM for expert layers. - grouped_gemm_dgrad - Triton grouped BF16 data-gradient product. - grouped_gemm_wgrad - Triton grouped BF16 weight-gradient product. ### GEMM - Triton groupwise & microscaling - matmul_f8f8bf16_groupwise - Triton FP8 GEMM with K-group scale granularity 128. (mslk.gemm.triton.fp8_groupwise_gemm) [NVIDIA, AMD gfx950] - matmul_f8f8bf16_groupwise_grouped - Triton stacked grouped FP8 groupwise GEMM. (mslk.gemm.triton.fp8_groupwise_grouped_gemm) [NVIDIA, AMD gfx950] - matmul_mx8mx4bf16 - Triton MXFP8 x MXFP4 GEMM. (mslk.gemm.triton.mx8mx4_gemm) [AMD gfx950] - matmul_mx8mx4bf16_grouped - Triton grouped MXFP8 x MXFP4 GEMM. (mslk.gemm.triton.mx8mx4_gemm) [AMD gfx950] - matmul_mx8mx8bf16_grouped - Triton grouped MXFP8 x MXFP8 GEMM. (mslk.gemm.triton.mx8mx8_gemm) [AMD gfx950] - mxfp4_gemm - Triton packed MXFP4 x MXFP4 GEMM. (mslk.gemm.triton.f4f4bf16) [AMD gfx950] - mxfp4_grouped_mm - ROCm gfx950 offset-grouped standard MXFP4 GEMM. (mslk.gemm.triton.f4f4bf16) [AMD gfx950] - mxfp4_grouped_stacked_gemm - ROCm gfx950 stacked grouped standard MXFP4 GEMM. (mslk.gemm.triton.f4f4bf16) [AMD gfx950] - f4f4bf16 - ROCm compatibility wrapper for standard MXFP4 block-size-32 GEMM. (mslk.gemm.triton.f4f4bf16) [AMD gfx950] ### GEMM - Triton integer GEMM - matmul_bf16i4_rowwise - Triton BF16 activation x row/group INT4 weight GEMM. (mslk.gemm.triton.int4_gemm) - matmul_bf16i4_rowwise_batched - Triton batched BF16 x INT4 GEMM. (mslk.gemm.triton.int4_gemm) - i8i8bf16_triton - Triton INT8 x INT8 GEMM with scalar scale. (mslk.gemm.triton.int8_gemm) - i8i8bf16_dynamic_triton - Triton INT8 x INT8 GEMM with tensor scale. (mslk.gemm.triton.int8_gemm) ### GEMM - Blackwell mixed-input CuTeDSL Module: `mslk.gemm.blackwell_mixed_input_gemm` - MixedInputGemmKernel - Configurable compiled Blackwell mixed-input GEMM kernel object. [NVIDIA SM100] - mixed_input_gemm - General Blackwell mixed-input GEMM frontend. [NVIDIA SM100] - int4bf16bf16_gemm - Convenience INT4 narrow operand x BF16 wide operand GEMM. [NVIDIA SM100] - int8bf16bf16_gemm - Convenience INT8 narrow operand x BF16 wide operand GEMM. [NVIDIA SM100] - create_tensors - Creates benchmark/reference tensors and layouts for mixed-input GEMM. [NVIDIA SM100] - compare - Compares kernel output with a CPU reference. [NVIDIA SM100] - run - Benchmark/validation driver for mixed-input configurations. [NVIDIA SM100] ### Quantization - FP8 quantization Module: `mslk.quantize.triton.fp8_quantize` - triton_quantize_fp8_row - Triton rowwise FP8 quantizer. - quantize_fp8_row - Public rowwise FP8 adapter selecting the Triton kernel or a PyTorch fallback. - scale_fp8_row - Applies activation and weight row scales to an accumulator/result tensor. - triton_quantize_fp8_block - Triton 2D blockwise FP8 quantizer. - quantize_fp8_block - Public blockwise FP8 adapter. - triton_quantize_fp8_group - Triton K-group FP8 quantizer, optionally aware of grouped valid row counts. - quantize_fp8_group - Public groupwise FP8 adapter. - triton_quantize_fp8_tensor - Triton tensorwise FP8 quantizer. - quantize_fp8_tensor - Public tensorwise FP8 adapter. - dequantize_fp8_row - Dequantizes rowwise FP8 storage to BF16. - dequantize_fp8_block - Dequantizes blockwise FP8 storage to BF16. ### Quantization - FP8 custom ops & dtype helpers - torch.ops.triton.quantize_fp8_row - Custom-op registration for rowwise FP8 quantization. (torch.ops.triton) [NVIDIA, AMD, Meta / fake] - torch.ops.triton.quantize_fp8_block - Custom-op registration for blockwise FP8 quantization. (torch.ops.triton) [NVIDIA, AMD, Meta / fake] - torch.ops.triton.quantize_fp8_tensor - Custom-op registration for tensorwise FP8 quantization. (torch.ops.triton) [NVIDIA, AMD, Meta / fake] - get_fp8_constants - Returns the platform PyTorch dtype, Triton dtype, maximum finite value, and epsilon. (mslk.utils.triton.fp8_utils) [NVIDIA, AMD, Meta / fake] - reinterpret_fp8_type - Reinterprets FP8 storage for a Triton dtype without numeric conversion. (mslk.utils.triton.fp8_utils) [NVIDIA, AMD, Meta / fake] ### Quantization - MXFP4 quantization - triton_quantize_mx4 - Preferred public MXFP4 group-32 quantizer. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+, AMD gfx950] - triton_quantize_mx4_unpack - Compatibility adapter supporting group size 16/32 and legacy rounding arguments. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+, AMD gfx950] - quantize_mx4 - Direct MXFP4 kernel frontend underlying the public adapter. (mslk.quantize.triton.quantize_kernels.mx4) [NVIDIA SM100+, AMD gfx950] - quantize_mx4_stacked - Segment-aware MXFP4 quantization for stacked grouped inputs. (mslk.quantize.triton.quantize_kernels.mx4_stacked) [NVIDIA SM100+, AMD gfx950] ### Quantization - NVFP4 & legacy FP4 - triton_quantize_nvfp4 - NVFP4 quantizer with per-16-value E4M3 scale factors. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+] - triton_fake_quantize_nvfp4_per_tensor - Per-tensor NVFP4 fake quantization returning BF16 values. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+] - nvfp4_quantize_stacked - Segment-aware stacked NVFP4 quantization. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+] - nvfp4_quantize_stacked_with_token_scale - Stacked NVFP4 quantization deriving an inverse scale per token. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+] - calculate_group_max - Computes per-segment NVFP4 global scales and a row-to-segment map. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+] - cal_global_scale_mx4_as_nvfp4 - Derives an NVFP4-style global scale for MXFP4 data. (mslk.quantize.triton.fp4_quantize) [NVIDIA SM100+] - global_scale_nvfp4 - Computes the standard global NVFP4 scale from input range. (mslk.quantize.triton.fp4_utils) [NVIDIA SM100+] - fp4_to_float - Decodes FP4 nibbles to FP32 values. (mslk.quantize.triton.fp4_utils) [NVIDIA SM100+] - scale_nvfp4 - Applies NVFP4 local and global scaling to decoded values. (mslk.quantize.triton.fp4_utils) [NVIDIA SM100+] - dequantize_nvfp4 - Dequantizes packed NVFP4 to BF16. (mslk.quantize.triton.fp4_utils) [NVIDIA SM100+] - dequantize_mx4 - Dequantizes packed MXFP4 with NVIDIA blocked scales to BF16. (mslk.quantize.triton.fp4_utils) [NVIDIA SM100+] ### Quantization - FP4 constants & primitives - RoundingMode - Rounding policy enum used by MXFP4 quantizers. (mslk.quantize.triton.fp4_quantize) - get_mx4_exp_bias - Returns the MX exponent bias for a format width. (mslk.quantize.triton.fp4_quantize) - FP4_E2M1_MAX - Largest finite E2M1 magnitude. (mslk.quantize.triton.fp4_quantize) - FP8_E4M3_MAX - Largest finite OCP E4M3 magnitude. (mslk.quantize.triton.fp4_quantize) - FP4_EBITS - E2M1 exponent-bit count. (mslk.quantize.triton.fp4_quantize) - FP4_MBITS - E2M1 mantissa-bit count. (mslk.quantize.triton.fp4_quantize) - E8M0_EXPONENT_BIAS - Shared exponent bias used to encode MX block scales. (mslk.quantize.triton.fp4_primitives) - BF16_MIN_NORMAL - Smallest normal BF16 magnitude used by safe scale math. (mslk.quantize.triton.fp4_primitives) - blocked_scale_offset - Maps a logical scale coordinate into Blackwell's 128 x 4 blocked layout. (mslk.quantize.triton.fp4_primitives) - stacked_segment_map - Triton helper mapping a row tile to a stacked segment/expert. (mslk.quantize.triton.fp4_primitives) - mx4_scale_normalize_encode - Triton primitive that normalizes an MX block and encodes its E8M0 scale. (mslk.quantize.triton.fp4_primitives) - convert_fp32_to_fp4_packed - Triton primitive converting pairs of FP32 values into packed E2M1 nibbles. (mslk.quantize.triton.fp4_primitives) - unsigned_fp32_to_e8m0 - Compatibility Triton primitive encoding positive FP32 magnitudes as E8M0 scales. (mslk.quantize.triton.fp4_quantize) - nvfp4_scale_swizzle - Compatibility Triton helper mapping rows into the NVFP4 scale-swizzle layout. (mslk.quantize.triton.fp4_quantize) ### Quantization - INT4 preprocessing Module: `mslk.quantize.shuffle` - pack_int4 - Packs two logical 4-bit integer values into each int8 byte. - int4_row_quantize_zp - Groupwise asymmetric INT4 quantization with zero points. - int4_row_quantize - Groupwise symmetric INT4 quantization. - quantize_int4_preshuffle - Quantizes, packs, and CUDA-preshuffles INT4 weights and scale terms. [NVIDIA] - ck_preshuffle - Reorders data into the AMD Composable Kernel XDL weight layout. [AMD ROCm] ### Quantization - MXFP6 helpers Module: `mslk.quantize.mx_mixed_dtype_utils` - quantize_bf16_to_mx6_e2m3 - Quantizes BF16 to unpacked six-bit E2M3 codes plus E8M0 block scales. [NVIDIA SM100+] - pack_fp6_e2m3 - Bit-packs four E2M3 values into three uint8 bytes. [NVIDIA SM100+] - E2M3_DECODE - Lookup table mapping 6-bit E2M3 codes to decoded values. [NVIDIA SM100+] - E2M3_MAX - Largest finite E2M3 magnitude. [NVIDIA SM100+] ### MoE - Routing & token movement Module: `mslk.moe` - index_shuffling - Selects top-k experts and groups token indices by expert. - gather_scale_dense_tokens - Gathers routed tokens and multiplies each by its router score. - gather_scale_quant_dense_tokens - Gathers/scales routed tokens and rowwise-quantizes them to FP8. - scatter_add_dense_tokens - Atomically accumulates routed expert rows back to dense token order. - scatter_add_padded_tokens - Scatters expert/rank-padded token blocks back into dense token storage. - combine_shuffling - Reorders rank-major expert-token blocks to expert-major/rank-minor. - split_shuffling - Inverse of combine_shuffling for expert communication results. - silu_mul - Fused SwiGLU activation x0-sigmoid(x0)-x1. - silu_mul_quant - Fused SwiGLU plus rowwise FP8 quantization. ### MoE - Registered MoE ops Module: `torch.ops.mslk` - torch.ops.mslk.index_shuffling - Native top-k expert grouping op. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.scatter_add_along_first_dim - In-place first-dimension scatter-add with a fast BF16 TMA path. [NVIDIA] - torch.ops.mslk.silu_mul - Registered fused SwiGLU activation. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.silu_mul_quant - Registered fused SwiGLU plus FP8 quantization. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.gather_scale_dense_tokens - Registered routed-token gather/scale op. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.gather_scale_quant_dense_tokens - Registered quantizing routed-token gather with a known return-schema mismatch. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.scatter_add_dense_tokens - Registered dense-token scatter-add op. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.scatter_add_padded_tokens - Registered padded expert-token scatter op. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.combine_shuffling - Registered rank-major to expert-major reorder. [NVIDIA, AMD, Meta / fake] - torch.ops.mslk.split_shuffling - Registered inverse expert-shuffling reorder. [NVIDIA, AMD, Meta / fake] ### MoE - Composed MoE layers Module: `mslk.moe.layers` - MoEArgs - Frozen configuration dataclass for the provided MoE modules. [NVIDIA, AMD, Distributed] - BaselineMoE - Dense/reference distributed SwiGLU MoE implementation. [NVIDIA, AMD, Distributed] - MetaShufflingMoE - Optimized inference MoE using shuffle/gather/grouped-GEMM primitives and overlapped communication. [NVIDIA, AMD, Distributed] - ScaledParameter - Non-trainable Parameter carrying optional quantization scales. [NVIDIA, AMD, Distributed] - Experts - Abstract shared base for expert weight containers. [NVIDIA, AMD, Distributed] - RoutedExperts - Per-local-expert stacked weight container. [NVIDIA, AMD, Distributed] - SharedExperts - Shared SwiGLU expert weight container. [NVIDIA, AMD, Distributed] - init_params - Applies a named initialization/quantization callback or Kaiming-uniform fallback. [NVIDIA, AMD, Distributed] ### Convolution - FP8 3D convolution Module: `torch.ops.mslk` - torch.ops.mslk.f8f8bf16_conv - FP8 3D cross-correlation/convolution with combined scale and BF16 output. [NVIDIA SM100+, Meta / fake] ### Runtime - FlyDSL runtime & AOT - is_flydsl_available - Cached check for FlyDSL importability and current-architecture support. (mslk.flydsl.common) [AMD ROCm, CPU / Python] - require_flydsl - Raises RuntimeError with an installation hint when FlyDSL is unavailable. (mslk.flydsl.common) [AMD ROCm, CPU / Python] - configure_runtime_cache - Points FlyDSL at bundled AOT artifacts and applies MSLK's JIT-disable switch. (mslk.flydsl.jit) [AMD ROCm, CPU / Python] - run_compiled - Compiles a FlyDSL launcher once, then invokes the cached compiled function. (mslk.flydsl.jit) [AMD ROCm, CPU / Python] - collect_aot_jobs - Builds the Cartesian product of registered kernel configs and architectures. (mslk.flydsl.aot) [AMD ROCm, CPU / Python] - compile_aot - Multiprocess-compiles all collected FlyDSL AOT jobs into a runtime cache. (mslk.flydsl.aot) [AMD ROCm, CPU / Python] ### Runtime - Device detection Module: `mslk.utils.device` - is_cuda - True for an NVIDIA CUDA build with an available device. [NVIDIA, AMD, CPU / Python] - is_rocm - True for a PyTorch HIP build with an available AMD device. [NVIDIA, AMD, CPU / Python] - compute_capability_in - Checks current device compute-capability major against an inclusive range. [NVIDIA, AMD, CPU / Python] - compute_capability_at_least - Checks the current capability tuple against a minimum. [NVIDIA, AMD, CPU / Python] - cuda_version_at_least - Checks the PyTorch build's CUDA toolkit major, not driver or device capability. [NVIDIA, AMD, CPU / Python] - get_gfx_arch_name - Returns current ROCm gcnArchName, or an empty string on unavailable/error. [NVIDIA, AMD, CPU / Python] - gfx_arch_in - Substring-matches any requested gfx name against gcnArchName. [NVIDIA, AMD, CPU / Python] - is_gfx942 - True on AMD MI300X/CDNA3 gfx942. [NVIDIA, AMD, CPU / Python] - is_gfx950 - True on AMD MI350/CDNA4 gfx950. [NVIDIA, AMD, CPU / Python] - supports_float8_fnuz - Reports whether the active ROCm target uses the FNUZ FP8 flavor. [NVIDIA, AMD, CPU / Python] ### Runtime - Test decorators Module: `mslk.testing.device` - skipUnlessCuda - Skips unless running on a strict NVIDIA CUDA device. [NVIDIA, AMD, CPU / Python] - skipUnlessRocm - Skips unless running on a strict AMD ROCm device. [NVIDIA, AMD, CPU / Python] - skipUnlessCudaCapability - Narrows CUDA tests to a capability range and remains transparent on ROCm. [NVIDIA, AMD, CPU / Python] - skipUnlessCudaVersion - Narrows CUDA tests to a minimum toolkit major and remains transparent on ROCm. [NVIDIA, AMD, CPU / Python] - skipUnlessGfxArch - Narrows ROCm tests to selected gfx targets and remains transparent on CUDA. [NVIDIA, AMD, CPU / Python] ### Runtime - Package loading & runtime knobs - open_source - Signals the OSS package variant to tests and integration code. (mslk) [NVIDIA, AMD, CPU / Python] - __version__ / __target__ / __variant__ - Generated package-build identity values, with internal/default fallbacks. (mslk) [NVIDIA, AMD, CPU / Python] - load_library_buck - Internal/OSS bridge used by domain imports to load split Buck native libraries. (mslk.utils.torch.library) [NVIDIA, AMD, CPU / Python] - MSLK_PYTHON_ONLY - Skips mslk.so native loading/compilation for Python-only development. (mslk) [NVIDIA, AMD, CPU / Python] - MSLK runtime environment variables - Runtime cache, autotune, timing-stat, and ROCm FP8 override environment variables. (mslk) [NVIDIA, AMD, CPU / Python] - MSLK diagnostic compile-time defines - Preprocessor defines enabling accessor bounds checks, isolated launches, and tensor value checks. (native build) [NVIDIA, AMD, CPU / Python] - FlyDSL environment variables - Runtime cache, run-only, worker-count, and compile-only controls for FlyDSL. (mslk.flydsl) [NVIDIA, AMD, CPU / Python] - setup.py build CLI - Source-build command-line controls. (mslk) [NVIDIA, AMD, CPU / Python] - MSLK build environment variables - Package identity, Python-only, toolkit, architecture, and toolchain inputs used by setup.py. (mslk) [NVIDIA, AMD, CPU / Python] - MSLK CMake knobs - Direct CMake configuration variables for native builds. (mslk) [NVIDIA, AMD, CPU / Python] - Package dependencies - Declared runtime dependency boundary and optional kernel packages. (mslk) [NVIDIA, AMD, CPU / Python] ### C++ - Numeric & device utilities Module: `namespace mslk` - nextPowerOf2 - Rounds a positive integer up to the next power of two. [NVIDIA, AMD, CPU] - roundUp - Rounds num up to a requested multiple. [NVIDIA, AMD, CPU] - nextPowerOf2OrRoundUp - Uses power-of-two rounding below a threshold and fixed-multiple rounding above it. [NVIDIA, AMD, CPU] - getDeviceArch - Returns cached current-device compute-capability major. [NVIDIA] - getSMCount - Returns an override or the available SM count after PyTorch carveout. [NVIDIA] ### C++ - CUDA stream & shared memory Module: `mslk::utils::device` - get_device_for_stream - Resolves the CUDA device associated with a raw stream. [NVIDIA] - to_cuda_stream - Normalizes raw cudaStream_t or c10 stream input to CUDAStream. [NVIDIA] - set_gpu_max_dynamic_shared_memory - Opts a kernel into the requested dynamic shared-memory limit after availability checks. [NVIDIA] ### C++ - Source context & timing - source_location - Portable source-location type used by SourceContext. (mslk::utils) [NVIDIA, AMD, CPU] - SourceContext - Carries source location and a human-readable label through kernel diagnostics. (mslk::utils) [NVIDIA, AMD, CPU] - SOURCE_CONTEXT_CURRENT - Captures the current file/function/line into a SourceContext. (global macro) [NVIDIA, AMD, CPU] - KernelExecutionTimer - Single-use CUDA-event kernel timer with explicit state checks. (mslk::utils) [NVIDIA] ### C++ - Tensor accessors - TensorAccessor - ATen-compatible tensor accessor with optional named/contextual device bounds assertions. (mslk::utils) [NVIDIA, AMD, CPU] - PackedTensorAccessor - By-value size/stride accessor suitable for passing into GPU kernels. (mslk::utils) [NVIDIA, AMD, CPU] - PackedTensorAccessor32 - Packed accessor using 32-bit indices. (mslk::utils) [NVIDIA, AMD, CPU] - PackedTensorAccessor64 - Packed accessor using 64-bit indices. (mslk::utils) [NVIDIA, AMD, CPU] - DefaultPtrTraits - Default accessor pointer trait alias. (mslk::utils) [NVIDIA, AMD, CPU] - RestrictPtrTraits - CUDA/HIP restricted-pointer trait alias. (mslk::utils) - pta - Compile-time selector between enhanced MSLK and standard ATen accessors. (global namespace) [NVIDIA, AMD, CPU] - overflow_safe_int_t - Global index type reserved for overflow-safe size arithmetic. (global namespace) [NVIDIA, AMD, CPU] - ::PackedTensorAccessor - Host-build global alias used when MSLK_MEMCHECK is not defined. (global namespace) [NVIDIA, AMD, CPU] - NAME_MAX_LEN / CONTEXT_MAX_LEN - Fixed diagnostic string capacities stored in enhanced accessors. (mslk::utils) [NVIDIA, AMD, CPU] ### C++ - Accessor builder & launch Module: `mslk::utils` - scalar_type_for - Maps a C++ scalar type to the corresponding ATen ScalarType. - TensorAccessorBuilder - Defers tensor validation and accessor construction until kernel launch context is known. - TA_B / PTA_B - Creates named non-packed/packed accessor builders while capturing the variable name. - mslk::utils::pta - Builder-local selector for enhanced MSLK versus standard ATen accessors. - MAKE_PTA_WITH_NAME - Legacy named packed-accessor construction macro. - is_tensor_accessor_builder - Detects TensorAccessorBuilder template instances. - transform_kernel_arg - Builds accessor-builder arguments with launch context; forwards other arguments. - check_kernel_arg - Runs configured tensor value checks on accessor-builder arguments. - KernelLauncher - Validated normal/cooperative GPU launch wrapper with optional diagnostics and timing. - MSLK_LAUNCH_KERNEL - Instantiates KernelLauncher with build-time diagnostic toggles and launches a normal kernel. - MSLK_LAUNCH_COOPERATIVE_KERNEL - Launches through the cooperative KernelLauncher specialization. ### C++ - Autotuning & CUTLASS helpers - TuningCache - Persistent per-shape best-kernel cache with optional CUDA-graph benchmarking. (mslk) - KernelMode - Coarse GEMM-shape mode used to select CUTLASS kernels. (mslk) - get_kernel_mode - Classifies a non-batched GEMM into a KernelMode. (mslk) - get_batched_kernel_mode - Classifies a batched GEMM into a KernelMode. (mslk) - LinearCombinationOnDevice - CUTLASS epilogue linear-combination operator whose scale data lives on device. (cutlass::epilogue::thread) [NVIDIA] - LinearCombinationOnDevice::Params - Host/device-constructable alpha/beta parameter carrier for LinearCombinationOnDevice. (cutlass::epilogue::thread) [NVIDIA] - GroupedGemmInputType - Selects grouped GEMM activation/weight rank layout. (mslk::gemm) - set_grouped_gemm_args_kernel - Builds per-group CUTLASS pointer/shape/stride argument arrays on device. (mslk::gemm) [NVIDIA] ### C++ - Native GEMM declarations - FP8 & BF16 Module: `mslk::gemm` - mslk::gemm::f8f8bf16_blockwise - Native rectangular-block-scaled FP8 GEMM returning BF16. - mslk::gemm::f8f8bf16_rowwise_preshuffle - ROCm preshuffled rowwise FP8 GEMM returning BF16. [AMD ROCm] - mslk::gemm::f8f8f16_rowwise_preshuffle - ROCm preshuffled rowwise FP8 GEMM declared with FP16 output. [AMD ROCm] - mslk::gemm::f8f8bf16_rowwise_grouped - Native list-based grouped rowwise FP8 GEMM. - mslk::gemm::f8f8bf16_rowwise_grouped_cat - Native list-based grouped rowwise FP8 GEMM with concatenated output. - mslk::gemm::f8f8bf16_rowwise_grouped_stacked - Native stacked grouped rowwise FP8 GEMM. - mslk::gemm::f8f8bf16_rowwise_grouped_dynamic - Native dynamic-row grouped FP8 GEMM. - mslk::gemm::f8f8bf16_rowwise_batched - Native batched rowwise FP8 GEMM with optional output reuse. - mslk::gemm::bf16bf16bf16_grouped - Native list-based grouped BF16 GEMM. - mslk::gemm::bf16bf16bf16_grouped_cat - Native list-based grouped BF16 GEMM with concatenated output. - mslk::gemm::bf16bf16bf16_grouped_dynamic - Native dynamic-row grouped BF16 GEMM. - mslk::gemm::bf16bf16bf16_grouped_stacked - Native stacked grouped BF16 GEMM with output/SM overrides. - mslk::gemm::f8f8bf16_rowwise - Native rowwise FP8 GEMM returning BF16. - mslk::gemm::f8f8bf16_rowwise_out - Native rowwise FP8 GEMM writing a caller-owned output tensor. - mslk::gemm::f8f8f16_rowwise - ROCm native rowwise FP8 GEMM returning FP16. [AMD ROCm] - mslk::gemm::f8f8bf16_groupwise - Native declaration for K-group-scaled FP8 GEMM. - mslk::gemm::f8f8bf16_groupwise_grouped - Native declaration for stacked grouped K-group FP8 GEMM. - mslk::gemm::i8i8bf16 - Native INT8 GEMM with scalar dequantization scale. - mslk::gemm::i8i8bf16_dynamic - Native INT8 GEMM with tensor-valued dynamic scale. ### C++ - Native GEMM declarations - FP4, INT4 & MX Module: `mslk::gemm` - mslk::gemm::f4f4bf16_grouped_stacked - Native stacked grouped MXFP4/NVFP4 GEMM. [NVIDIA] - mslk::gemm::bf16x9_gemm - Native cuBLAS BF16x9 emulation GEMM producing FP32. [NVIDIA] - mslk::gemm::bf16i4bf16_shuffled - Native BF16 x preshuffled INT4 GEMM. [NVIDIA] - mslk::gemm::f8i4bf16_shuffled_grouped - Native stacked grouped FP8 x preshuffled INT4 GEMM. [NVIDIA] - mslk::gemm::bf16i4bf16_shuffled_grouped - Native stacked grouped BF16 x preshuffled INT4 GEMM. [NVIDIA] - mslk::gemm::bf16i4bf16_shuffled_batched - Native batched BF16 x preshuffled INT4 GEMM. [NVIDIA] - mslk::gemm::bf16i4bf16_rowwise_batched - Native declaration for batched BF16 x rowwise INT4 GEMM. - mslk::gemm::bf16i4bf16_rowwise - Native declaration for BF16 x row/group-quantized INT4 GEMM. - mslk::gemm::f8i4bf16_rowwise - Native rowwise FP8 activation x INT4 weight GEMM. [NVIDIA] - mslk::gemm::f8i4bf16_shuffled - Native FP8 activation x preshuffled INT4 weight GEMM. [NVIDIA] - mslk::gemm::preshuffle_i4 - Native one-time INT4 weight/scale preprocessing entry. [NVIDIA] - mslk::gemm::mx8mx4bf16 - Native declaration for MXFP8 x MXFP4 GEMM. - mslk::gemm::mx8mx6bf16 - Native MXFP8 x MXFP6 GEMM. [NVIDIA] - mslk::gemm::mx6mx6bf16 - Native MXFP6 x MXFP6 GEMM with optional split reduction. [NVIDIA] ### C++ - Native offset-grouped GEMM declarations Module: `mslk::gemm` - mslk::gemm::f8f8bf16_rowwise_grouped_mm - ROCm native layout-polymorphic grouped rowwise FP8 GEMM. [AMD ROCm] - mslk::gemm::mx8mx8bf16_grouped_mm - Native offset-grouped MXFP8 x MXFP8 GEMM. [NVIDIA] - mslk::gemm::f4f4bf16_grouped_mm - Native offset-grouped MXFP4/NVFP4 GEMM. [NVIDIA] - mslk::gemm::f4f4bf16_ultra_grouped_mm - Native ultra grouped FP4 GEMM with separate global scales. [NVIDIA] - mslk::gemm::mx8mx4bf16_grouped_mm - Native offset-grouped MXFP8 x MXFP4 GEMM. [NVIDIA] - mslk::gemm::f4f4bf16 - Native packed FP4 GEMM selecting MXFP4, MXFP4-16, or NVFP4. [NVIDIA] ### C++ - Native Conv & MoE declarations - mslk::conv::f8f8bf16_conv - Native FP8 3D convolution/cross-correlation entry. (mslk::conv) [NVIDIA SM100+] - mslk::moe::index_shuffling_torch - Native implementation entry for the index_shuffling dispatcher schema. (mslk::moe) - mslk::moe::scatter_add_along_first_dim - Native in-place first-dimension scatter-add entry. (mslk::moe) [NVIDIA] ### C++ - Native entry headers Module: `mslk C++ headers` - Native GEMM / Conv / MoE declarations - ATen-native declarations corresponding to the documented torch.ops schemas.