{ "name": "MSLK", "description": "Fused GPU kernels for transformer workloads: attention, low-precision GEMM, quantization, MoE routing, convolution.", "repository": "https://github.com/meta-pytorch/MSLK", "commit": "69ae1b897f2d546d72acab129e1ae4bc2924900f", "generatedFrom": "docs/index.html", "categories": { "Attention": { "description": "fMHA dispatch, masks, architecture backends, tree attention, MLA, and paged KV.", "symbols": 156 }, "GEMM": { "description": "BF16, FP8, FP4, MX, INT8, INT4, grouped, batched, and mixed-input matrix products.", "symbols": 70 }, "Quantization": { "description": "FP8, MXFP4, NVFP4, INT4, and MXFP6 conversion, packing, scaling, and dequantization.", "symbols": 54 }, "MoE": { "description": "Routing, shuffling, token gather/scatter, fused activations, and composed expert layers.", "symbols": 27 }, "Convolution": { "description": "The low-precision FP8 3D convolution operator.", "symbols": 1 }, "Runtime": { "description": "FlyDSL JIT/AOT, device detection, architecture checks, and test decorators.", "symbols": 32 }, "C++": { "description": "Kernel launch, accessors, tuning, timing, source context, and CUTLASS helpers.", "symbols": 84 } }, "symbols": [ { "name": "memory_efficient_attention", "signature": "memory_efficient_attention(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None) -> Tensor", "summary": "Autograd-enabled fused attention with automatic backend dispatch.", "details": "Accepts [B,M,K], [B,M,H,K], or experimental [B,M,G,H,K] inputs. It computes scaled QKᵀ, applies an optimized bias or dense tensor bias, softmax/dropout, then multiplies V without materializing the full attention matrix.", "returns": "Attention output in the query layout, with its last dimension replaced by the V head dimension.", "notes": [ "The last dimension must have stride 1; other dimensions may be non-contiguous.", "For GQA/MQA, reshape and explicitly expand K/V heads; MSLK does not broadcast them automatically.", "Leave op=None for recommended dispatch; pass a (forward, backward) class tuple to pin a backend." ], "methods": [ { "signature": "op=(FwOp, BwOp)", "summary": "Explicitly pins forward and backward operator classes." }, { "signature": "attn_bias=AttentionBias | Tensor | None", "summary": "Uses optimized masks when possible; arbitrary dense tensors are generally slower." } ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Core fMHA", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "fMHA", "dispatch", "BMK", "BMHK", "BMGHK" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention" }, { "name": "memory_efficient_attention_forward", "signature": "memory_efficient_attention_forward(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None) -> Tensor", "summary": "Forward-only attention path that does not retain manual-backward context.", "details": "Validates and normalizes the same inputs as the high-level API, selects a forward backend, and returns only the attention output.", "returns": "Attention output tensor.", "notes": [ "Pass a forward operator class, not an operator pair, through op.", "Use this path when no gradients or LSE are needed." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Core fMHA", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "fMHA", "dispatch", "BMK", "BMHK", "BMGHK" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-forward" }, { "name": "memory_efficient_attention_forward_requires_grad", "signature": "memory_efficient_attention_forward_requires_grad(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None) -> tuple[Tensor, Tensor]", "summary": "Manual-backward forward pass returning output and log-sum-exp.", "details": "Despite its name, this is the explicit non-autograd API. The returned LSE is the context consumed by memory_efficient_attention_backward.", "returns": "(output, lse).", "notes": [ "Nonzero dropout is rejected by this manual API.", "The forward and backward operators must agree on variable-length LSE layout." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Core fMHA", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "fMHA", "dispatch", "BMK", "BMHK", "BMGHK" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-forward-requires-grad" }, { "name": "memory_efficient_attention_backward", "signature": "memory_efficient_attention_backward(grad, output, lse, query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None) -> tuple[Tensor, Tensor, Tensor]", "summary": "Explicit backward pass for an output/LSE pair.", "details": "Consumes the original Q/K/V, forward output, LSE, and upstream gradient. Dispatches a compatible backward kernel unless one is supplied.", "returns": "(dq, dk, dv) reshaped to the original inputs.", "notes": [ "Dropout is not supported on the manual forward/backward API.", "This return contract does not include a dense tensor-bias gradient." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Core fMHA", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "fMHA", "dispatch", "BMK", "BMHK", "BMGHK" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-backward" }, { "name": "memory_efficient_attention_partial", "signature": "memory_efficient_attention_partial(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None, _allow_backward=False) -> tuple[Tensor, Tensor]", "summary": "Computes an output and LSE for one disjoint K/V shard.", "details": "Call it repeatedly with the same query and separate key/value chunks, then combine the results with merge_attentions. This enables exact split-KV attention.", "returns": "(partial_output, partial_lse).", "notes": [ "Dropout is unsupported.", "The public path is forward-only. _allow_backward is private and deliberately restricted.", "Use the merge_training wrappers for safe training." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Core fMHA", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "fMHA", "dispatch", "BMK", "BMHK", "BMGHK" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-partial" }, { "name": "merge_attentions", "signature": "merge_attentions(attn_split, lse_split, write_lse=True, output_dtype=None) -> tuple[Tensor, Optional[Tensor]]", "summary": "Log-sum-exp-correct merge of attention computed over K/V chunks.", "details": "Accepts sequences of chunk tensors or tensors stacked along a leading chunk axis. It computes the exact normalized output over the union of chunks rather than averaging partial outputs.", "returns": "(merged_output, merged_lse) when write_lse=True; otherwise (merged_output, None).", "notes": [ "Direct autograd through this function is unsupported; use merge_attentions_autograd.", "Inputs requiring gradients require write_lse=True." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Core fMHA", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "fMHA", "dispatch", "BMK", "BMHK", "BMGHK" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-merge-attentions" }, { "name": "torch.ops.mslk.memory_efficient_attention_forward", "signature": "torch.ops.mslk.memory_efficient_attention_forward(q, k, v, b=None, p=0.0, scale=None) -> Tensor", "summary": "Compile-friendly forward attention with optional bias.", "details": "These custom-library wrappers keep dispatch outside a torch.compile trace and expose a reduced forward-only contract.", "returns": "Attention output tensor matching q.", "notes": [ "No explicit op or output_dtype argument.", "K and V head dimensions must match.", "Only selected bias forms are trace-safe." ], "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Compile-friendly wrappers", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "Meta / fake" ], "tags": [ "torch.compile", "custom op", "forward" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-mslk-memory-efficient-attention-forward" }, { "name": "torch.ops.mslk.memory_efficient_attention_forward_with_bias", "signature": "torch.ops.mslk.memory_efficient_attention_forward_with_bias(q, k, v, b, p=0.0, scale=None) -> Tensor", "summary": "Compile-friendly forward attention requiring an explicit bias tensor.", "details": "These custom-library wrappers keep dispatch outside a torch.compile trace and expose a reduced forward-only contract.", "returns": "Attention output tensor matching q.", "notes": [ "No explicit op or output_dtype argument.", "K and V head dimensions must match.", "Only selected bias forms are trace-safe." ], "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Compile-friendly wrappers", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "Meta / fake" ], "tags": [ "torch.compile", "custom op", "forward" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-mslk-memory-efficient-attention-forward-with-bias" }, { "name": "Partial", "signature": "Partial(attn: Tensor, lse: Tensor, placeholder: Tensor)", "summary": "Autograd carrier for a partial attention output and its merge context.", "details": "Users normally receive Partial from memory_efficient_attention_partial_autograd rather than constructing it directly.", "returns": "A lightweight object consumed by merge_attentions_autograd.", "methods": [ { "signature": "is_bmghk() -> bool", "summary": "Reports whether the stored output uses the grouped-head BMGHK layout." }, { "signature": "apply(fn: Callable[[Tensor], Tensor]) -> Partial", "summary": "Applies a layout-preserving transform that must not manipulate the embedding dimension." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Trainable partial attention", "module": "mslk.attention.fmha.merge_training", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autograd", "split KV", "LSE" ], "source": "mslk/attention/fmha/merge_training.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/merge_training.py", "href": "#api-partial" }, { "name": "memory_efficient_attention_partial_autograd", "signature": "memory_efficient_attention_partial_autograd(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None) -> Partial", "summary": "Autograd-safe wrapper around partial attention.", "details": "Records Q/K/V and uses hidden placeholder gradients to route the merged output and full LSE into each shard’s backward pass.", "returns": "Partial.", "notes": [ "Arguments mirror memory_efficient_attention_partial.", "Dropout remains unsupported by the underlying partial path." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Trainable partial attention", "module": "mslk.attention.fmha.merge_training", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autograd", "split KV", "LSE" ], "source": "mslk/attention/fmha/merge_training.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/merge_training.py", "href": "#api-memory-efficient-attention-partial-autograd" }, { "name": "merge_attentions_autograd", "signature": "merge_attentions_autograd(*partials: Partial) -> Tensor", "summary": "Merges one or more Partial objects with correct gradients.", "details": "For multiple shards it performs the exact LSE merge. Passing a single Partial returns its output while retaining the autograd plumbing.", "returns": "Merged attention output.", "notes": [ "Raises ValueError when called without partials." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Trainable partial attention", "module": "mslk.attention.fmha.merge_training", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autograd", "split KV", "LSE" ], "source": "mslk/attention/fmha/merge_training.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/merge_training.py", "href": "#api-merge-attentions-autograd" }, { "name": "AttentionBias", "signature": "AttentionBias()", "summary": "Abstract base for optimized attention-bias descriptions.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "Subclasses define materialize and, where state is stored, to(device)." ], "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-attention-bias" }, { "name": "LowerTriangularMask", "signature": "LowerTriangularMask(device: Optional[torch.device] = None)", "summary": "Top-left-aligned causal mask.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "The device constructor argument is retained only for compatibility and is ignored." ], "methods": [ { "signature": "add_bias(bias: Tensor) -> LowerTriangularMaskWithTensorBias", "summary": "Combines causal masking with an arbitrary additive tensor." }, { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Produces zeros on/below the diagonal and −∞ above it." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-lower-triangular-mask" }, { "name": "LowerTriangularMaskWithTensorBias", "signature": "LowerTriangularMaskWithTensorBias(bias: Tensor)", "summary": "Top-left causal mask plus a dense additive tensor bias.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "Backend support is narrower than for a pure causal mask." ], "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-lower-triangular-mask-with-tensor-bias" }, { "name": "LowerTriangularFromBottomRightMask", "signature": "LowerTriangularFromBottomRightMask()", "summary": "Bottom-right-aligned causal mask for Mq ≠ Mkv.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "make_local_attention(window_size: int) -> LowerTriangularFromBottomRightLocalAttentionMask", "summary": "Restricts the causal region to a trailing window." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-lower-triangular-from-bottom-right-mask" }, { "name": "LowerTriangularFromBottomRightLocalAttentionMask", "signature": "LowerTriangularFromBottomRightLocalAttentionMask(_window_size: int)", "summary": "Bottom-right causal sliding-window mask.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-lower-triangular-from-bottom-right-local-attention-mask" }, { "name": "LocalAttentionFromBottomRightMask", "signature": "LocalAttentionFromBottomRightMask(window_left: int, window_right: int)", "summary": "Non-causal local window aligned from the bottom right.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-local-attention-from-bottom-right-mask" }, { "name": "BlockDiagonalMask", "signature": "BlockDiagonalMask(q_seqinfo, k_seqinfo, _batch_sizes=None)", "summary": "Packed independent sequences where each Q block sees its matching K/V block.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens(q_seqlen, kv_seqlen=None, *, device=None) -> BlockDiagonalMask", "summary": "Builds sequence metadata from lengths." }, { "signature": "from_tensor_list(tensors) -> tuple[BlockDiagonalMask, Tensor]", "summary": "Concatenates a list and returns its matching block mask." }, { "signature": "from_tensor_lists_qkv(tensors_q, tensors_k, tensors_v=None) -> tuple[BlockDiagonalMask, Tensor, Tensor, Optional[Tensor]]", "summary": "Packs Q/K/(V) lists and returns all tensors plus a mask." }, { "signature": "split_queries(tensor) / split_kv(tensor) / split(tensor)", "summary": "Restores packed outputs to the original logical sequences." }, { "signature": "make_causal() / make_causal_from_bottomright()", "summary": "Returns the corresponding causal subclass." }, { "signature": "make_local_attention(window_size) / make_local_attention_from_bottomright(window_size)", "summary": "Returns a local block-diagonal variant." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-mask" }, { "name": "BlockDiagonalCausalMask", "signature": "BlockDiagonalCausalMask(q_seqinfo, k_seqinfo, _batch_sizes=None)", "summary": "Top-left causal mask inside every packed sequence block.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-causal-mask" }, { "name": "BlockDiagonalCausalFromBottomRightMask", "signature": "BlockDiagonalCausalFromBottomRightMask(q_seqinfo, k_seqinfo, _batch_sizes=None)", "summary": "Bottom-right causal mask inside every packed block.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "Each K/V block must be at least as long as its Q block." ], "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-causal-from-bottom-right-mask" }, { "name": "BlockDiagonalCausalLocalAttentionMask", "signature": "BlockDiagonalCausalLocalAttentionMask(q_seqinfo, k_seqinfo, _batch_sizes=None, _window_size=0)", "summary": "Top-left causal local window inside packed blocks.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "Although the generated dataclass signature defaults _window_size to 0, construction rejects non-positive values." ], "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-causal-local-attention-mask" }, { "name": "BlockDiagonalCausalLocalAttentionFromBottomRightMask", "signature": "BlockDiagonalCausalLocalAttentionFromBottomRightMask(q_seqinfo, k_seqinfo, _batch_sizes=None, _window_size=0)", "summary": "Bottom-right causal local window inside packed blocks.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "Although the generated dataclass signature defaults _window_size to 0, construction rejects non-positive values." ], "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-causal-local-attention-from-bottom-right-mask" }, { "name": "BlockDiagonalPaddedKeysMask", "signature": "BlockDiagonalPaddedKeysMask(q_seqinfo, k_seqinfo)", "summary": "Packed Q blocks attending fixed-capacity padded K/V lanes.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "causal_diagonal currently exists for compatibility and is unused." ], "methods": [ { "signature": "from_seqlens(q_seqlen, kv_padding, kv_seqlen, causal_diagonal=None, *, device=None)", "summary": "Creates padded-lane metadata." }, { "signature": "make_paged(block_tables, page_size, paged_type)", "summary": "Maps logical lanes to physical KV pages." }, { "signature": "make_local_attention(window_left, window_right)", "summary": "Returns a padded local-attention variant." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-padded-keys-mask" }, { "name": "BlockDiagonalCausalWithOffsetPaddedKeysMask", "signature": "BlockDiagonalCausalWithOffsetPaddedKeysMask(q_seqinfo, k_seqinfo, causal_diagonal=None)", "summary": "Bottom-right causal masking over padded K/V lanes.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens(q_seqlen, kv_padding, kv_seqlen, causal_diagonal=None, *, device=None)", "summary": "Creates bottom-right causal padded metadata." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-causal-with-offset-padded-keys-mask" }, { "name": "BlockDiagonalLocalAttentionPaddedKeysMask", "signature": "BlockDiagonalLocalAttentionPaddedKeysMask(q_seqinfo, k_seqinfo, window_left, window_right)", "summary": "Non-causal local windows over padded K/V lanes.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens_local(q_seqlen, kv_padding, kv_seqlen, window_left, window_right)", "summary": "Constructs the local padded mask." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-local-attention-padded-keys-mask" }, { "name": "BlockDiagonalCausalLocalAttentionPaddedKeysMask", "signature": "BlockDiagonalCausalLocalAttentionPaddedKeysMask(q_seqinfo, k_seqinfo, _window_size)", "summary": "Bottom-right causal local windows over padded K/V lanes.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens_local(q_seqlen, kv_padding, kv_seqlen, window_size)", "summary": "Constructs the causal-local padded mask." }, { "signature": "make_paged(block_tables, page_size, paged_type)", "summary": "Returns the paged counterpart; paged_type is required." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-causal-local-attention-padded-keys-mask" }, { "name": "PagedBlockDiagonalPaddedKeysMask", "signature": "PagedBlockDiagonalPaddedKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)", "summary": "Padded logical K/V lanes mapped to physical cache pages.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "notes": [ "Physical K/V uses [1, physical_pages × page_size, H, D] or its grouped-head equivalent." ], "methods": [ { "signature": "from_seqlens(q_seqlen, kv_seqlen, block_tables, page_size, *, device=None)", "summary": "Builds page-aware metadata." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-paged-block-diagonal-padded-keys-mask" }, { "name": "PagedBlockDiagonalCausalWithOffsetPaddedKeysMask", "signature": "PagedBlockDiagonalCausalWithOffsetPaddedKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)", "summary": "Bottom-right causal paged padded mask.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-paged-block-diagonal-causal-with-offset-padded-keys-mask" }, { "name": "PagedBlockDiagonalCausalLocalPaddedKeysMask", "signature": "PagedBlockDiagonalCausalLocalPaddedKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size, _window_size)", "summary": "Bottom-right causal local paged padded mask.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens_local(q_seqlen, kv_seqlen, block_tables, page_size, window_size, *, device=None)", "summary": "Constructs the paged causal-local mask." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-paged-block-diagonal-causal-local-padded-keys-mask" }, { "name": "BlockDiagonalGappyKeysMask", "signature": "BlockDiagonalGappyKeysMask(q_seqinfo, k_seqinfo)", "summary": "K/V lanes with arbitrary starts and independently used lengths.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens(q_seqlen, kv_seqstarts, kv_seqlen, *, device=None)", "summary": "Describes gappy lanes." }, { "signature": "make_paged(block_tables, page_size, notional_padding, paged_type)", "summary": "Maps the gappy layout to physical pages." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-gappy-keys-mask" }, { "name": "BlockDiagonalLocalAttentionFromBottomRightGappyKeysMask", "signature": "BlockDiagonalLocalAttentionFromBottomRightGappyKeysMask(q_seqinfo, k_seqinfo, window_left, window_right)", "summary": "Bottom-right-aligned local windows over gappy K/V lanes.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens_local_gappy(q_seqlen, kv_seqstarts, kv_seqlen, window_left, window_right, device)", "summary": "Constructs the local gappy mask." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-local-attention-from-bottom-right-gappy-keys-mask" }, { "name": "BlockDiagonalCausalWithOffsetGappyKeysMask", "signature": "BlockDiagonalCausalWithOffsetGappyKeysMask(q_seqinfo, k_seqinfo)", "summary": "Bottom-right causal masking over gappy K/V lanes.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-block-diagonal-causal-with-offset-gappy-keys-mask" }, { "name": "PagedBlockDiagonalGappyKeysMask", "signature": "PagedBlockDiagonalGappyKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)", "summary": "Gappy logical K/V lanes mapped through a page table.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "from_seqlens(q_seqlen, kv_seqstarts, kv_seqlen, block_tables, page_size, *, device=None)", "summary": "Builds paged-gappy metadata." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-paged-block-diagonal-gappy-keys-mask" }, { "name": "PagedBlockDiagonalCausalWithOffsetGappyKeysMask", "signature": "PagedBlockDiagonalCausalWithOffsetGappyKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)", "summary": "Bottom-right causal paged-gappy mask.", "details": "AttentionBias objects describe sparsity or additive bias without materializing the full Q×K matrix, allowing compatible kernels to encode the mask directly.", "returns": "An object accepted as attn_bias by fMHA.", "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-paged-block-diagonal-causal-with-offset-gappy-keys-mask" }, { "name": "VARLEN_BIASES", "signature": "VARLEN_BIASES: tuple[type[AttentionBias], ...]", "summary": "Dispatch grouping for packed, padded, gappy, and paged variable-length masks.", "details": "Used by backend support checks and manual-backward LSE layout detection.", "returns": "Tuple of bias base classes.", "methods": [ { "signature": "materialize(shape, dtype=torch.float32, device='cpu') -> Tensor", "summary": "Builds the dense bias for debugging/reference checks; intentionally slow." }, { "signature": "to(device)", "summary": "Moves stored sequence metadata or tensor bias to a device." } ], "kind": "constant", "stability": "Low-level", "category": "Attention", "group": "Biases & masks", "module": "mslk.attention.fmha.attn_bias", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "attention bias", "mask" ], "source": "mslk/attention/fmha/attn_bias.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py", "href": "#api-varlen-biases" }, { "name": "cutlass.FwOp", "signature": "cutlass.FwOp", "summary": "Classic CUTLASS forward attention with the broadest NVIDIA dtype/head-dim range.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Maximum Q/K head dimension is 65,536.", "Supports different V head dimension and grouped-head forward.", "Does not support partial attention." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cutlass", "platforms": [ "NVIDIA SM60–SM90" ], "tags": [ "FP32", "FP16", "BF16", "dropout", "tensor bias", "BMGHK" ], "source": "mslk/attention/fmha/cutlass.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass.py", "href": "#api-cutlass-fw-op" }, { "name": "cutlass.BwOp", "signature": "cutlass.BwOp", "summary": "Classic CUTLASS attention backward, including supported tensor-bias gradients.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cutlass", "platforms": [ "NVIDIA SM60–SM90" ], "tags": [ "FP32", "FP16", "BF16", "dropout", "bias grad" ], "source": "mslk/attention/fmha/cutlass.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass.py", "href": "#api-cutlass-bw-op" }, { "name": "cutlass_blackwell.FwOp", "signature": "cutlass_blackwell.FwOp", "summary": "Blackwell CUTLASS prefill/varlen forward backend.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "High-level backend supports head dimensions 64 or 128.", "No dropout, partial attention, or different V head dimension.", "Automatic CUDA dispatch does not currently prioritize this class; pin it explicitly." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cutlass_blackwell", "platforms": [ "NVIDIA SM100+" ], "tags": [ "FP16", "BF16", "varlen" ], "source": "mslk/attention/fmha/cutlass_blackwell.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass_blackwell.py", "href": "#api-cutlass-blackwell-fw-op" }, { "name": "cutlass_blackwell.FwOpDecode", "signature": "cutlass_blackwell.FwOpDecode", "summary": "Blackwell inference/decode forward backend.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Inference-only; Q heads per KV head are limited to 16." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cutlass_blackwell", "platforms": [ "NVIDIA SM100+" ], "tags": [ "BF16", "decode", "GQA" ], "source": "mslk/attention/fmha/cutlass_blackwell.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass_blackwell.py", "href": "#api-cutlass-blackwell-fw-op-decode" }, { "name": "cutlass_blackwell.BwOp", "signature": "cutlass_blackwell.BwOp", "summary": "Blackwell CUTLASS backward backend.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "BMHK only; head dimensions 64 or 128.", "No dropout or different V dimension." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cutlass_blackwell", "platforms": [ "NVIDIA SM100+" ], "tags": [ "FP16", "BF16", "BMHK" ], "source": "mslk/attention/fmha/cutlass_blackwell.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass_blackwell.py", "href": "#api-cutlass-blackwell-bw-op" }, { "name": "flash.FwOp", "signature": "flash.FwOp", "summary": "FlashAttention 2 forward backend for the general NVIDIA fast path.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Maximum head dimension 256; Q/K and V head dimensions must match.", "Paged support depends on the bundled versus PyTorch Flash implementation; enabled paths require page size divisible by 256." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.flash", "platforms": [ "NVIDIA SM80+" ], "tags": [ "FP16", "BF16", "dropout", "partial", "varlen", "paged", "BMGHK" ], "source": "mslk/attention/fmha/flash.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py", "href": "#api-flash-fw-op" }, { "name": "flash.BwOp", "signature": "flash.BwOp", "summary": "FlashAttention 2 backward backend.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Grouped-head BMGHK is forward-only." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.flash", "platforms": [ "NVIDIA SM80+" ], "tags": [ "FP16", "BF16", "dropout", "varlen" ], "source": "mslk/attention/fmha/flash.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py", "href": "#api-flash-bw-op" }, { "name": "flash3.FwOp", "signature": "flash3.FwOp", "summary": "FlashAttention 3 forward with FP8 and flexible V-dimension support.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Head dimension must be 64, 128, 192, or 256.", "No dropout.", "Ampere requires Q/K and V dimensions to match." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.flash3", "platforms": [ "NVIDIA SM80–SM90" ], "tags": [ "FP16", "BF16", "FP8 E4M3", "partial", "paged", "gappy", "BMGHK" ], "source": "mslk/attention/fmha/flash3.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py", "href": "#api-flash3-fw-op" }, { "name": "flash3.BwOp", "signature": "flash3.BwOp", "summary": "FlashAttention 3 backward backend.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Padded, gappy, paged, and BMGHK paths are forward-only." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.flash3", "platforms": [ "NVIDIA SM80–SM90" ], "tags": [ "FP16", "BF16", "FP8 E4M3" ], "source": "mslk/attention/fmha/flash3.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py", "href": "#api-flash3-bw-op" }, { "name": "flash3.FwOp_KVSplit", "signature": "flash3.FwOp_KVSplit", "summary": "FlashAttention 3 split-KV forward variant.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.flash3", "platforms": [ "NVIDIA SM80–SM90" ], "tags": [ "split KV", "decode", "FP8" ], "source": "mslk/attention/fmha/flash3.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py", "href": "#api-flash3-fw-op-kvsplit" }, { "name": "ck.FwOp", "signature": "ck.FwOp", "summary": "Composable Kernel attention forward for ROCm.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Maximum head dimension 256.", "PyTorch still reports the HIP device type as cuda." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.ck", "platforms": [ "AMD ROCm" ], "tags": [ "FP16", "BF16", "dropout", "partial", "BMGHK", "different V dim" ], "source": "mslk/attention/fmha/ck.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck.py", "href": "#api-ck-fw-op" }, { "name": "ck.BwOp", "signature": "ck.BwOp", "summary": "Composable Kernel attention backward, including supported bias gradients.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.ck", "platforms": [ "AMD ROCm" ], "tags": [ "FP16", "BF16", "dropout", "bias grad" ], "source": "mslk/attention/fmha/ck.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck.py", "href": "#api-ck-bw-op" }, { "name": "ck_decoder.FwOp", "signature": "ck_decoder.FwOp", "summary": "ROCm inference decoder attention backend.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Inference-only; fixed per-lane query length.", "Key padding is limited to 8,192." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.ck_decoder", "platforms": [ "AMD ROCm / MI250X" ], "tags": [ "FP16", "BF16", "FP32", "decode" ], "source": "mslk/attention/fmha/ck_decoder.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_decoder.py", "href": "#api-ck-decoder-fw-op" }, { "name": "ck_splitk.FwOp", "signature": "ck_splitk.FwOp", "summary": "ROCm split-K decoder/prefix attention with optional quantized K/V.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "No dropout or backward.", "Maximum head dimension 256; causal inputs have query-length restrictions." ], "methods": [ { "signature": "get_split_k(B, H, Mk) -> int", "summary": "Selects a split count from batch, heads, and KV length." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "FP16", "BF16", "FP32", "INT32 KV", "split K" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op" }, { "name": "triton_splitk.FwOp", "signature": "triton_splitk.FwOp", "summary": "Portable split-K forward with quantized KV-cache support.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Supported head dimensions: 16, 32, 64, 128, 256, 512.", "No dropout or backward.", "Supports fused row/group INT4 and rowwise FP8 caches stored through int32 representations.", "FwOp_Map also contains split counts 48, 72, 80, 96, and 112 in addition to the named power-of-two aliases." ], "methods": [ { "signature": "get_operator(splitk, *, block_m=None, block_n=None, num_warps=None, num_stages=None, split_k_early_exit=None)", "summary": "Builds/selects a specialized operator class." }, { "signature": "get_split_k(B, G, H, Mk, Mq, page_size, is_paged=False) -> int", "summary": "Heuristic split selector." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm / MI300X" ], "tags": [ "FP16", "BF16", "FP8", "INT4 KV", "FP8 KV", "partial", "paged" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op" }, { "name": "cute_hopper.FwOp", "signature": "cute_hopper.FwOp", "summary": "CuTe DSL Hopper forward attention.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Head dimension 64 or 128.", "No dropout, partial attention, or different V dimension." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cute_hopper", "platforms": [ "NVIDIA SM90–SM99" ], "tags": [ "FP16", "BF16", "CuTe", "BMHK", "BMGHK" ], "source": "mslk/attention/fmha/cute_hopper.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_hopper.py", "href": "#api-cute-hopper-fw-op" }, { "name": "cute_hopper.BwOp", "signature": "cute_hopper.BwOp", "summary": "CuTe DSL Hopper backward attention.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cute_hopper", "platforms": [ "NVIDIA SM90–SM99" ], "tags": [ "FP16", "BF16", "CuTe" ], "source": "mslk/attention/fmha/cute_hopper.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_hopper.py", "href": "#api-cute-hopper-bw-op" }, { "name": "cute_blackwell.FwOp", "signature": "cute_blackwell.FwOp", "summary": "CuTe DSL Blackwell prefill forward attention.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Prefill head dimensions: 64, 96, 128.", "No dropout, partial attention, or different V dimension." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cute_blackwell", "platforms": [ "NVIDIA SM100+" ], "tags": [ "FP16", "BF16", "FP8 E4M3", "paged", "CuTe" ], "source": "mslk/attention/fmha/cute_blackwell.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_blackwell.py", "href": "#api-cute-blackwell-fw-op" }, { "name": "cute_blackwell.FwOpDecode", "signature": "cute_blackwell.FwOpDecode", "summary": "CuTe DSL Blackwell decode/paged forward attention.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Decode head dimension 64 or 128.", "Paged-gappy support is limited to this decode path." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cute_blackwell", "platforms": [ "NVIDIA SM100+" ], "tags": [ "FP16", "BF16", "FP8 E4M3", "decode", "paged", "gappy" ], "source": "mslk/attention/fmha/cute_blackwell.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_blackwell.py", "href": "#api-cute-blackwell-fw-op-decode" }, { "name": "cute_blackwell.BwOp", "signature": "cute_blackwell.BwOp", "summary": "CuTe DSL Blackwell backward attention.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "notes": [ "Backward head dimensions: 64, 96, 128." ], "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.cute_blackwell", "platforms": [ "NVIDIA SM100+" ], "tags": [ "FP16", "BF16", "FP8 E4M3", "CuTe" ], "source": "mslk/attention/fmha/cute_blackwell.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_blackwell.py", "href": "#api-cute-blackwell-bw-op" }, { "name": "flash_mtia.FwOp", "signature": "flash_mtia.FwOp", "summary": "MTIA counterpart of the FlashAttention 2 forward protocol.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.flash_mtia", "platforms": [ "MTIA" ], "tags": [ "FP16", "BF16", "dropout", "partial" ], "source": "mslk/attention/fmha/flash_mtia.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py", "href": "#api-flash-mtia-fw-op" }, { "name": "flash_mtia.BwOp", "signature": "flash_mtia.BwOp", "summary": "MTIA counterpart of the FlashAttention 2 backward protocol.", "details": "Backend classes implement the fMHA operator protocol. Pass a forward class to forward-only APIs or an (FwOp, BwOp) tuple to memory_efficient_attention. Each class exposes availability and input-support checks before launch.", "returns": "An operator class used for dispatch; call through the fMHA frontends.", "methods": [ { "signature": "is_available() -> bool", "summary": "Checks whether the compiled/JIT implementation is present." }, { "signature": "supports(inputs) -> bool", "summary": "Returns whether the operator accepts a normalized Inputs object." }, { "signature": "not_supported_reasons(inputs) -> list[str]", "summary": "Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints." } ], "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Backend operators", "module": "mslk.attention.fmha.flash_mtia", "platforms": [ "MTIA" ], "tags": [ "FP16", "BF16", "dropout" ], "source": "mslk/attention/fmha/flash_mtia.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py", "href": "#api-flash-mtia-bw-op" }, { "name": "ck_splitk.FwOp_S1", "signature": "ck_splitk.FwOp_S1", "summary": "CK split-K forward fixed to 1 split.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 1" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s1" }, { "name": "ck_splitk.FwOp_S2", "signature": "ck_splitk.FwOp_S2", "summary": "CK split-K forward fixed to 2 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 2" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s2" }, { "name": "ck_splitk.FwOp_S4", "signature": "ck_splitk.FwOp_S4", "summary": "CK split-K forward fixed to 4 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 4" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s4" }, { "name": "ck_splitk.FwOp_S8", "signature": "ck_splitk.FwOp_S8", "summary": "CK split-K forward fixed to 8 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 8" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s8" }, { "name": "ck_splitk.FwOp_S16", "signature": "ck_splitk.FwOp_S16", "summary": "CK split-K forward fixed to 16 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 16" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s16" }, { "name": "ck_splitk.FwOp_S32", "signature": "ck_splitk.FwOp_S32", "summary": "CK split-K forward fixed to 32 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 32" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s32" }, { "name": "ck_splitk.FwOp_S64", "signature": "ck_splitk.FwOp_S64", "summary": "CK split-K forward fixed to 64 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 64" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s64" }, { "name": "ck_splitk.FwOp_S128", "signature": "ck_splitk.FwOp_S128", "summary": "CK split-K forward fixed to 128 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.ck_splitk", "platforms": [ "AMD ROCm" ], "tags": [ "CK", "split 128" ], "source": "mslk/attention/fmha/ck_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py", "href": "#api-ck-splitk-fw-op-s128" }, { "name": "triton_splitk.FwOp_S1", "signature": "triton_splitk.FwOp_S1", "summary": "Triton split-K forward fixed to 1 split.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 1" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s1" }, { "name": "triton_splitk.FwOp_S2", "signature": "triton_splitk.FwOp_S2", "summary": "Triton split-K forward fixed to 2 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 2" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s2" }, { "name": "triton_splitk.FwOp_S4", "signature": "triton_splitk.FwOp_S4", "summary": "Triton split-K forward fixed to 4 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 4" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s4" }, { "name": "triton_splitk.FwOp_S8", "signature": "triton_splitk.FwOp_S8", "summary": "Triton split-K forward fixed to 8 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 8" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s8" }, { "name": "triton_splitk.FwOp_S16", "signature": "triton_splitk.FwOp_S16", "summary": "Triton split-K forward fixed to 16 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 16" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s16" }, { "name": "triton_splitk.FwOp_S32", "signature": "triton_splitk.FwOp_S32", "summary": "Triton split-K forward fixed to 32 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 32" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s32" }, { "name": "triton_splitk.FwOp_S64", "signature": "triton_splitk.FwOp_S64", "summary": "Triton split-K forward fixed to 64 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 64" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s64" }, { "name": "triton_splitk.FwOp_S128", "signature": "triton_splitk.FwOp_S128", "summary": "Triton split-K forward fixed to 128 splits.", "details": "Fixed subclasses force a compile-time/launch-time split count. They otherwise follow their family’s support contract.", "returns": "A forward operator class.", "kind": "backend", "stability": "Advanced", "category": "Attention", "group": "Fixed split-K backends", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "Triton", "split 128" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-s128" }, { "name": "MemoryEfficientAttentionCutlassOp", "signature": "MemoryEfficientAttentionCutlassOp = (cutlass.FwOp, cutlass.BwOp)", "summary": "Classic CUTLASS forward and backward.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-cutlass-op" }, { "name": "MemoryEfficientAttentionCutlassBlackwellOp", "signature": "MemoryEfficientAttentionCutlassBlackwellOp = (cutlass_blackwell.FwOp, cutlass_blackwell.BwOp)", "summary": "Blackwell CUTLASS forward and backward.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-cutlass-blackwell-op" }, { "name": "MemoryEfficientAttentionCutlassFwdFlashBwOp", "signature": "MemoryEfficientAttentionCutlassFwdFlashBwOp = (cutlass.FwOp, flash.BwOp)", "summary": "Classic CUTLASS forward with FlashAttention 2 backward.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-cutlass-fwd-flash-bw-op" }, { "name": "MemoryEfficientAttentionFlashAttentionOp", "signature": "MemoryEfficientAttentionFlashAttentionOp = (flash.FwOp, flash.BwOp)", "summary": "FlashAttention 2 pair.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-flash-attention-op" }, { "name": "MemoryEfficientAttentionFlashMtiaAttentionOp", "signature": "MemoryEfficientAttentionFlashMtiaAttentionOp = (flash_mtia.FwOp, flash_mtia.BwOp)", "summary": "MTIA Flash pair.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "MTIA" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-flash-mtia-attention-op" }, { "name": "MemoryEfficientAttentionCkOp", "signature": "MemoryEfficientAttentionCkOp = (ck.FwOp, ck.BwOp)", "summary": "ROCm CK pair.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "AMD ROCm" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-ck-op" }, { "name": "MemoryEfficientAttentionCkDecoderOp", "signature": "MemoryEfficientAttentionCkDecoderOp = (ck_decoder.FwOp, ck.BwOp)", "summary": "Inference-only convenience tuple using CK decoder forward.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "notes": [ "The forward class produces no backward context; the paired ck.BwOp does not make this tuple valid for autograd training." ], "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "AMD ROCm" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-ck-decoder-op" }, { "name": "MemoryEfficientAttentionSplitKCkOp", "signature": "MemoryEfficientAttentionSplitKCkOp = (ck_splitk.FwOp, ck.BwOp)", "summary": "Inference-only convenience tuple using CK split-K forward.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "notes": [ "The forward class produces no backward context; the paired ck.BwOp does not make this tuple valid for autograd training." ], "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "AMD ROCm" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-split-kck-op" }, { "name": "MemoryEfficientAttentionCuteFlashAttentionOp", "signature": "MemoryEfficientAttentionCuteFlashAttentionOp = (cute_blackwell.FwOp, cute_blackwell.BwOp)", "summary": "CuTe Blackwell forward/backward pair.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA SM100+" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-memory-efficient-attention-cute-flash-attention-op" }, { "name": "ALL_FW_OPS", "signature": "ALL_FW_OPS: list[type[AttentionFwOpBase]]", "summary": "Explicit forward-backend enumeration used by tests and introspection.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "notes": [ "This is not the automatic dispatch priority list." ], "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-all-fw-ops" }, { "name": "ALL_BW_OPS", "signature": "ALL_BW_OPS: list[type[AttentionBwOpBase]]", "summary": "Explicit backward-backend enumeration used by tests and introspection.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "notes": [ "This is not the automatic dispatch priority list." ], "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha", "platforms": [ "NVIDIA", "AMD", "MTIA" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py", "href": "#api-all-bw-ops" }, { "name": "triton_splitk.FwOp_Map", "signature": "triton_splitk.FwOp_Map: dict[int, type[triton_splitk.FwOp]]", "summary": "Maps supported split counts to specialized Triton split-K forward classes.", "details": "A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.", "returns": "AttentionOp tuple.", "kind": "constant", "stability": "Advanced", "category": "Attention", "group": "Preassembled operator pairs", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA SM80+", "AMD ROCm" ], "tags": [ "AttentionOp", "explicit dispatch" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-fw-op-map" }, { "name": "AttentionOp", "signature": "AttentionOp = tuple[Optional[type[AttentionFwOpBase]], Optional[type[AttentionBwOpBase]]]", "summary": "Forward/backward operator pair type accepted by high-level fMHA.", "kind": "type alias", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-attention-op" }, { "name": "AttentionOpBase", "signature": "AttentionOpBase", "summary": "Shared backend metadata, availability, and support-check protocol.", "methods": [ { "signature": "supports(inputs: Inputs) -> bool", "summary": "Whether the backend accepts an input descriptor." }, { "signature": "not_supported_reasons(inputs: Inputs) -> list[str]", "summary": "Human-readable rejection reasons." }, { "signature": "shape_not_supported_reasons(Mq, Mkv, K, Kv) -> list[str]", "summary": "Shape-only compatibility diagnostics." }, { "signature": "is_available() -> bool", "summary": "Whether the implementation is loadable." } ], "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-attention-op-base" }, { "name": "AttentionFwOpBase", "signature": "AttentionFwOpBase(AttentionOpBase)", "summary": "Base protocol for forward attention operators.", "methods": [ { "signature": "apply(inp: Inputs, needs_gradient: bool) -> tuple[Tensor, Optional[Context]]", "summary": "Runs the backend and optionally returns backward context." } ], "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-attention-fw-op-base" }, { "name": "AttentionBwOpBase", "signature": "AttentionBwOpBase(AttentionOpBase)", "summary": "Base protocol for backward attention operators.", "methods": [ { "signature": "apply(ctx: Context, inp: Inputs, grad: Tensor) -> Gradients", "summary": "Computes normalized Q/K/V/(bias) gradients." } ], "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-attention-bw-op-base" }, { "name": "Inputs", "signature": "Inputs(query: Tensor, key: Tensor, value: Tensor, attn_bias: Optional[Tensor | AttentionBias]=None, p: float=0.0, scale: Optional[float]=None, output_dtype: Optional[torch.dtype]=None, is_partial: bool=False, quantize_pv_to_fp8: bool=False, quantize_qk_to_fp8: bool=False, use_fp32_scales: bool=False, num_splits: int=0)", "summary": "Normalized fMHA input descriptor used by dispatch and backends.", "details": "Validates dtypes, layouts, shapes, masks, scale and dropout; can normalize BMK/BMHK/BMGHK logical layouts.", "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-inputs" }, { "name": "InputsFp8", "signature": "InputsFp8(query, key, value, attn_bias=None, p=0.0, scale=None, output_dtype=None, is_partial=False, quantize_pv_to_fp8=False, quantize_qk_to_fp8=False, use_fp32_scales=False, num_splits=0, k_fp8_scale_shift=None, v_fp8_scale_shift=None, q_fp8_scale_shift=None)", "summary": "Input descriptor for row/group quantized FP8-style K/V caches.", "details": "Carries packed scale/shift metadata used by direct quantized split-K backends.", "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-inputs-fp8" }, { "name": "InputsMXFp8", "signature": "InputsMXFp8(query, key, value, attn_bias=None, p=0.0, scale=None, output_dtype=None, is_partial=False, quantize_pv_to_fp8=False, quantize_qk_to_fp8=False, use_fp32_scales=False, num_splits=0, k_fp8_scale=None, v_fp8_scale=None, q_fp8_scale=None)", "summary": "Input descriptor for MXFP8 K/V with E8M0 block scales.", "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-inputs-mxfp8" }, { "name": "Context", "signature": "Context(lse: Tensor, out: Tensor, op_bw: Optional[type[AttentionBwOpBase]]=None, rng_state=None, qkv_share_storage: bool=False)", "summary": "Forward context consumed by an attention backward operator.", "methods": [ { "signature": "get_padded_lse(pad_to: int, force_pad_inf: bool=False) -> Tensor", "summary": "Returns LSE padded for a backend’s alignment contract." } ], "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-context" }, { "name": "Gradients", "signature": "Gradients(dq: Tensor, dk: Tensor, dv: Tensor, db: Optional[Tensor]=None)", "summary": "Normalized attention-gradient bundle.", "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-gradients" }, { "name": "ScaledTensor", "signature": "ScaledTensor(data: Tensor, scale: Tensor, dequant_func: Callable[[Tensor, Tensor], Tensor], original_dtype: torch.dtype, require_grad: bool=False)", "summary": "Tensor subclass carrying separate quantization-scale metadata.", "kind": "class", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-scaled-tensor" }, { "name": "pack_fp8_tensorwise_per_head", "signature": "pack_fp8_tensorwise_per_head(x: Tensor, scale: Tensor | float, original_dtype) -> ScaledTensor", "summary": "Wraps tensorwise/per-head FP8 storage and scale metadata for backend input.", "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-pack-fp8-tensorwise-per-head" }, { "name": "bmghk2bmhk", "signature": "bmghk2bmhk(x: Tensor, x_scale: Optional[Tensor]=None, handle_rep_heads: bool=False) -> tuple[Tensor, Optional[Tensor]]", "summary": "Flattens grouped heads G×H into a single H axis, with optional scale handling.", "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-bmghk2bmhk" }, { "name": "bmhk2bhk", "signature": "bmhk2bhk(x: Tensor, x_scale: Optional[Tensor]=None, handle_mqa: bool=True) -> tuple[Tensor, Optional[Tensor]]", "summary": "Converts BMHK storage to the backend BHK-style view for single-query/decode paths.", "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-bmhk2bhk" }, { "name": "bmk2bmhk", "signature": "bmk2bmhk(tensor: Tensor, num_heads: int) -> Tensor", "summary": "Adds/reshapes a head axis for BMK tensors.", "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-bmk2bmhk" }, { "name": "check_lastdim_alignment_stride1", "signature": "check_lastdim_alignment_stride1(reasons, name, x, alignment) -> None", "summary": "Appends backend support failures for last-dimension alignment and stride.", "kind": "function", "stability": "Developer", "category": "Attention", "group": "Operator protocol & inputs", "module": "mslk.attention.fmha.common", "platforms": [ "NVIDIA", "AMD", "MTIA", "CPU / Python" ], "tags": [ "dispatch", "introspection", "normalized inputs" ], "source": "mslk/attention/fmha/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py", "href": "#api-check-lastdim-alignment-stride1" }, { "name": "_get_use_fa3", "signature": "_get_use_fa3() -> bool", "summary": "Returns whether CUDA automatic dispatch is permitted to select FlashAttention 3.", "notes": [ "Underscore-prefixed but intentionally included in fmha.__all__." ], "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Dispatch controls", "module": "mslk.attention.fmha.dispatch", "platforms": [ "NVIDIA", "CPU / Python" ], "tags": [ "FlashAttention 3", "dispatch", "compatibility" ], "source": "mslk/attention/fmha/dispatch.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/dispatch.py", "href": "#api-get-use-fa3" }, { "name": "_set_use_fa3", "signature": "_set_use_fa3(use_flash_attention3: bool) -> None", "summary": "Enables or disables FlashAttention 3 in automatic dispatch.", "notes": [ "Process-global dispatch control." ], "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Dispatch controls", "module": "mslk.attention.fmha.dispatch", "platforms": [ "NVIDIA", "CPU / Python" ], "tags": [ "FlashAttention 3", "dispatch", "compatibility" ], "source": "mslk/attention/fmha/dispatch.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/dispatch.py", "href": "#api-set-use-fa3" }, { "name": "fa3_available", "signature": "fa3_available() -> bool", "summary": "Reports whether a compatible FlashAttention 3 implementation is available.", "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Dispatch controls", "module": "mslk.attention.fmha.dispatch", "platforms": [ "NVIDIA", "CPU / Python" ], "tags": [ "FlashAttention 3", "dispatch", "compatibility" ], "source": "mslk/attention/fmha/dispatch.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/dispatch.py", "href": "#api-fa3-available" }, { "name": "is_pt_cutlass_compatible", "signature": "is_pt_cutlass_compatible(force=False) -> bool", "summary": "Checks PyTorch’s bundled CUTLASS attention compatibility.", "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Dispatch controls", "module": "mslk.attention.fmha.torch_attention_compat", "platforms": [ "NVIDIA", "CPU / Python" ], "tags": [ "FlashAttention 3", "dispatch", "compatibility" ], "source": "mslk/attention/fmha/torch_attention_compat.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/torch_attention_compat.py", "href": "#api-is-pt-cutlass-compatible" }, { "name": "ensure_pt_flash_ok", "signature": "ensure_pt_flash_ok() -> None", "summary": "Raises when the active PyTorch FlashAttention integration is incompatible.", "kind": "function", "stability": "Advanced", "category": "Attention", "group": "Dispatch controls", "module": "mslk.attention.fmha.torch_attention_compat", "platforms": [ "NVIDIA", "CPU / Python" ], "tags": [ "FlashAttention 3", "dispatch", "compatibility" ], "source": "mslk/attention/fmha/torch_attention_compat.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/torch_attention_compat.py", "href": "#api-ensure-pt-flash-ok" }, { "name": "TreeAttnMetadata", "signature": "TreeAttnMetadata.from_tree_choices(tree_choices, dtype=None, device=None) -> TreeAttnMetadata", "summary": "Builds explicit mask and traversal metadata for a speculative token tree.", "details": "Stores the tree mask, draft-token mapping, leaf retrieval paths, path lengths, position IDs, parent/child tables, sibling candidate indices, per-level counts, and subtree sizes.", "methods": [ { "signature": "from_tree_choices(tree_choices, dtype=None, device=None)", "summary": "Builds metadata from node paths sorted by depth." }, { "signature": "from_tree_choices_cached(tree_choices: tuple[tuple[int,...],...], dtype=None, device=None)", "summary": "LRU-cached constructor for hashable choices." } ], "kind": "class", "stability": "Public", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-tree-attn-metadata" }, { "name": "tree_attention", "signature": "tree_attention(q, spec_k, spec_v, cache_k, cache_v, spec_attn_bias, prefix_attn_bias, prefix_op=None, suffix_op=None, autotune=False, quantized_kv_scales=None, q_fp8=None) -> Tensor", "summary": "Computes prefix-cache and speculative-tree attention, then LSE-merges them.", "details": "Implements Medusa/EAGLE/Hydra-style inference as two partial attention calls: a long prefix over cache K/V and a tree-masked suffix over speculative K/V.", "returns": "Merged attention output in BMHK or BMGHK layout.", "notes": [ "Supports paged prefix masks.", "Quantized int32 KV uses Triton split-K; FP8 Q/cache paths accept explicit scales." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-tree-attention" }, { "name": "use_triton_splitk_for_prefix", "signature": "use_triton_splitk_for_prefix(B: int, G: int, tree_size: int) -> bool", "summary": "Heuristic for choosing Triton split-K on the tree prefix.", "kind": "function", "stability": "Public", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-use-triton-splitk-for-prefix" }, { "name": "select_prefix_op", "signature": "select_prefix_op(B, G, tree_size, autotune, attn_bias, kv_cache_dtype) -> Optional[type[AttentionFwOpBase]]", "summary": "Selects the prefix forward operator for tree attention.", "kind": "function", "stability": "Public", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-select-prefix-op" }, { "name": "SplitKAutotune", "signature": "SplitKAutotune(triton_splitk.FwOp)", "summary": "Tree-attention split-K operator variant with runtime autotuning.", "kind": "backend", "stability": "Experimental", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-split-kautotune" }, { "name": "construct_full_tree_choices", "signature": "construct_full_tree_choices(tree_depth: int, branching: int) -> list[tuple[int, ...]]", "summary": "Generates node choices for a regular full tree.", "notes": [ "Its depth convention counts non-root levels; get_full_tree_size uses a different range convention." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-construct-full-tree-choices" }, { "name": "construct_tree_choices", "signature": "construct_tree_choices(branching: list[int]) -> list[tuple[int, ...]]", "summary": "Generates node choices for per-level branching factors.", "kind": "function", "stability": "Public", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-construct-tree-choices" }, { "name": "get_full_tree_size", "signature": "get_full_tree_size(tree_depth: int, branching: int) -> int", "summary": "Returns Σ branchingⁱ over range(tree_depth).", "notes": [ "The tree_depth convention does not directly match construct_full_tree_choices." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Tree & speculative attention", "module": "mslk.attention.fmha.tree_attention", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Medusa", "EAGLE", "Hydra", "speculative decoding", "tree" ], "source": "mslk/attention/fmha/tree_attention.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py", "href": "#api-get-full-tree-size" }, { "name": "split_blocks_for_decoding_gpu_part", "signature": "split_blocks_for_decoding_gpu_part(input_bias, batchify_len, block_tables=None, page_size=None) -> Optional[tuple[Tensor, Tensor]]", "summary": "Computes GPU-side sequence-start/length data for decoding block splits.", "details": "Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.", "kind": "function", "stability": "Experimental", "category": "Attention", "group": "iRoPE & page transforms", "module": "mslk.attention.fmha.split_blocks_fairinternal", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "iRoPE", "paged KV", "gappy keys", "advanced" ], "source": "mslk/attention/fmha/split_blocks_fairinternal.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py", "href": "#api-split-blocks-for-decoding-gpu-part" }, { "name": "split_blocks_for_decoding", "signature": "split_blocks_for_decoding(input_bias, batchify_len, block_tables=None, page_size=None, gpu_data=None) -> Optional[BlockDiagonalGappyKeysMask | PagedBlockDiagonalGappyKeysMask]", "summary": "Transforms padded decoding lanes into gappy or paged-gappy masks.", "details": "Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.", "kind": "function", "stability": "Experimental", "category": "Attention", "group": "iRoPE & page transforms", "module": "mslk.attention.fmha.split_blocks_fairinternal", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "iRoPE", "paged KV", "gappy keys", "advanced" ], "source": "mslk/attention/fmha/split_blocks_fairinternal.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py", "href": "#api-split-blocks-for-decoding" }, { "name": "split_blocks_for_prefill", "signature": "split_blocks_for_prefill(input_bias, batchify_len) -> Optional[BlockDiagonalPaddedKeysMask]", "summary": "Re-batches padded prefill lanes into a split block description.", "details": "Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.", "kind": "function", "stability": "Experimental", "category": "Attention", "group": "iRoPE & page transforms", "module": "mslk.attention.fmha.split_blocks_fairinternal", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "iRoPE", "paged KV", "gappy keys", "advanced" ], "source": "mslk/attention/fmha/split_blocks_fairinternal.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py", "href": "#api-split-blocks-for-prefill" }, { "name": "maybe_make_paged", "signature": "maybe_make_paged(attn_bias, block_tables, page_size, notional_padding) -> Optional[AttentionBias]", "summary": "Converts supported padded/gappy masks to paged counterparts when a block table is supplied.", "details": "Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.", "kind": "function", "stability": "Experimental", "category": "Attention", "group": "iRoPE & page transforms", "module": "mslk.attention.fmha.split_blocks_fairinternal", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "iRoPE", "paged KV", "gappy keys", "advanced" ], "source": "mslk/attention/fmha/split_blocks_fairinternal.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py", "href": "#api-maybe-make-paged" }, { "name": "mla_decode_fwd", "signature": "mla_decode_fwd(query, kv_buffer, block_tables, cu_seqlens_q, seqused_k, softmax_scale=None) -> Tensor", "summary": "Paged MLA decode forward for one query position per sequence.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "returns": "BF16 [B, 128, 512].", "notes": [ "Query shape is [B,128,576].", "Input may be BF16 or compatible ROCm FP8 FNUZ." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-decode-fwd" }, { "name": "mla_prefill_fwd", "signature": "mla_prefill_fwd(query, kv_buffer, block_tables, cu_seqlens_q, seqused_k, softmax_scale=None) -> Tensor", "summary": "Packed variable-length MLA prefill forward.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "returns": "BF16 [total_tokens, 128, 512].", "notes": [ "Query shape is [total_tokens,128,576]." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-prefill-fwd" }, { "name": "MLA_NUM_HEADS", "signature": "MLA_NUM_HEADS = 128", "summary": "Fixed MLA query-head count.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "kind": "constant", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla.triton_mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-num-heads" }, { "name": "MLA_NUM_KV_HEADS", "signature": "MLA_NUM_KV_HEADS = 1", "summary": "Fixed latent KV-head count.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "kind": "constant", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla.triton_mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-num-kv-heads" }, { "name": "MLA_KV_LORA_RANK", "signature": "MLA_KV_LORA_RANK = 512", "summary": "Latent KV LoRA rank.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "kind": "constant", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla.triton_mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-kv-lora-rank" }, { "name": "MLA_QK_ROPE_HEAD_DIM", "signature": "MLA_QK_ROPE_HEAD_DIM = 64", "summary": "RoPE portion of the Q/K head.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "kind": "constant", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla.triton_mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-qk-rope-head-dim" }, { "name": "MLA_QK_HEAD_DIM", "signature": "MLA_QK_HEAD_DIM = 576", "summary": "Full Q/K head dimension after absorption.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "kind": "constant", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla.triton_mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-qk-head-dim" }, { "name": "MLA_V_HEAD_DIM", "signature": "MLA_V_HEAD_DIM = 512", "summary": "MLA value/output head dimension.", "details": "Post-weight-absorption Multi-head Latent Attention for the fixed DeepSeek-V3-style geometry: 128 Q heads, one latent KV head, QK dimension 576, and V dimension 512.", "kind": "constant", "stability": "Public", "category": "Attention", "group": "ROCm MLA", "module": "mslk.attention.mla.triton_mla", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "MLA", "DeepSeek V3", "paged KV", "Triton", "BF16", "FP8 FNUZ" ], "source": "mslk/attention/mla/triton_mla.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py", "href": "#api-mla-v-head-dim" }, { "name": "flydsl_flash_attn_func", "signature": "flydsl_flash_attn_func(q, k, v, *, causal=True, num_kv_heads=None, cu_seqlens_q=None, cu_seqlens_kv=None, max_seqlen_q=None, max_seqlen_kv=None, cross_seqlen=None, block_table=None, seqlen_k=None, kv_cache_layout='linear', num_kv_splits=1, q_descale=None, k_descale=None, v_descale=None, out=None, waves_per_eu=2, daz=True, dualwave_swp_lazy_rescale=True, dualwave_swp_setprio=True, dualwave_swp_enable_stagger=True, debug_counts=None, return_lse=False, stream=None)", "summary": "ROCm FlyDSL flash-attention forward across dense, varlen, GQA/MQA, and paged KV.", "details": "gfx942 uses a generic path; gfx950 uses dual-wave kernels. Dense FP8 is gfx950-only and requires scalar FP32 q/k/v descales.", "returns": "Output tensor, or (output, LSE) when return_lse=True on supported non-FP8/non-paged paths.", "notes": [ "Head dimension must be at least 64 and divisible by 32.", "Split-K supports D=64/128, BF16/FP16, Q length ≥384, and is incompatible with varlen.", "Paged runtime accepts kv_cache_layout 'linear' or 'vectorized', page size 64, and D=64/128.", "The docstring’s 'linear3d' layout is not accepted by current validation." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "FlyDSL Flash Attention", "module": "mslk.attention.flydsl", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "FlyDSL", "dense", "varlen", "paged", "GQA", "MQA", "split K" ], "source": "mslk/attention/flydsl/flash_attn_interface.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/flydsl/flash_attn_interface.py", "href": "#api-flydsl-flash-attn-func" }, { "name": "dualwave_splitk_workspace_elems", "signature": "dualwave_splitk_workspace_elems(batch_size, num_heads, seq_len, num_kv_splits, head_dim=128) -> int", "summary": "Returns the FP32 element count needed by the gfx950 dual-wave split-K workspace.", "kind": "function", "stability": "Public", "category": "Attention", "group": "FlyDSL Flash Attention", "module": "mslk.attention.flydsl.flash_attn_interface", "platforms": [ "AMD gfx942", "AMD gfx950" ], "tags": [ "FlyDSL", "dense", "varlen", "paged", "GQA", "MQA", "split K" ], "source": "mslk/attention/flydsl/flash_attn_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/flydsl/flash_attn_utils.py", "href": "#api-dualwave-splitk-workspace-elems" }, { "name": "flash_attn_func", "signature": "flash_attn_func(...) -> dependency-defined", "summary": "Optional standalone CuTe Flash-Attention facade.", "details": "Aliases an internal implementation or flash_attn.cute.interface.flash_attn_func, so the exact signature follows the installed dependency.", "returns": "Dependency-defined attention result.", "notes": [ "The current OSS package exports only flash_attn_func.", "Repository tests mention flash_attn_varlen_func and other helpers that are absent from this OSS directory; do not rely on them." ], "kind": "function", "stability": "Experimental", "category": "Attention", "group": "Standalone facades", "module": "mslk.attention.flash_attn", "platforms": [ "NVIDIA / dependency-defined" ], "tags": [ "CuTe", "optional dependency", "standalone" ], "source": "mslk/attention/flash_attn/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/flash_attn/__init__.py", "href": "#api-flash-attn-func" }, { "name": "cutlass_blackwell_fmha_func", "signature": "cutlass_blackwell_fmha_func(q, k, v, softmax_scale=None, causal=False, cu_seqlens_q=None, cu_seqlens_k=None, max_seq_len_q=None, max_seq_len_k=None, seqlen_kv=None, page_table=None, seqlen_k=None, window_size=(-1, -1), bottom_right=True, deterministic=False)", "summary": "Autograd-capable standalone Blackwell FMHA facade.", "details": "Handles dense or variable-length inputs and routes sequence length sq=1 through the generated decode path.", "returns": "Attention output tensor.", "notes": [ "The generation path leaves its split dimension unmerged and has no backward implementation." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Standalone Blackwell FMHA", "module": "mslk.attention.cutlass_blackwell_fmha", "platforms": [ "NVIDIA SM100+" ], "tags": [ "CUTLASS", "Blackwell", "varlen", "paged", "decode" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "href": "#api-cutlass-blackwell-fmha-func" }, { "name": "cutlass_blackwell_fmha_decode_forward", "signature": "cutlass_blackwell_fmha_decode_forward(q, k, v, seqlen_kv=None, cu_seqlens_q=None, cu_seqlens_k=None, max_seq_len_q=None, max_seq_len_k=None, softmax_scale=None, causal=False, window_left=-1, window_right=-1, bottom_right=True, split_k_size=0, use_heuristic=True) -> tuple[Tensor, Tensor]", "summary": "Inference-only split-K decode forward with raw partial outputs.", "details": "Accepts q [B,H,D] or [B,1,H,D] and requires seqlen_kv. Unlike the fMHA backend class, this low-level function leaves the split dimension unmerged.", "returns": "Output [B,1,H,num_splits,D] and LSE [B,num_splits,H,1].", "kind": "function", "stability": "Public", "category": "Attention", "group": "Standalone Blackwell FMHA", "module": "mslk.attention.cutlass_blackwell_fmha", "platforms": [ "NVIDIA SM100+" ], "tags": [ "CUTLASS", "Blackwell", "varlen", "paged", "decode" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "href": "#api-cutlass-blackwell-fmha-decode-forward" }, { "name": "_cutlass_blackwell_fmha_forward", "signature": "_cutlass_blackwell_fmha_forward(q, k, v, cu_seqlens_q=None, cu_seqlens_k=None, max_seq_len_q=None, max_seq_len_k=None, softmax_scale=None, causal=False, seqlen_kv=None, page_table=None, seqlen_k=None, window_left=-1, window_right=-1, bottom_right=True) -> tuple[Tensor, Tensor]", "summary": "Testing/implementation forward wrapper returning output and LSE.", "kind": "function", "stability": "Low-level", "category": "Attention", "group": "Standalone Blackwell FMHA", "module": "mslk.attention.cutlass_blackwell_fmha", "platforms": [ "NVIDIA SM100+" ], "tags": [ "CUTLASS", "Blackwell", "varlen", "paged", "decode" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "href": "#api-cutlass-blackwell-fmha-forward" }, { "name": "cutlass_blackwell_fmha_custom_op", "signature": "cutlass_blackwell_fmha_custom_op(q, k, v, softmax_scale=None, causal=False, cu_seqlens_q=None, cu_seqlens_k=None, max_seq_len_q=None, max_seq_len_k=None, seqlen_kv=None, page_table=None, seqlen_k=-1, window_size_left=-1, window_size_right=-1, bottom_right=True)", "summary": "torch.library custom-op facade for standalone Blackwell FMHA.", "returns": "A single output tensor; the underlying native op’s LSE result is discarded.", "kind": "function", "stability": "Low-level", "category": "Attention", "group": "Standalone Blackwell FMHA", "module": "mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_custom_op", "platforms": [ "NVIDIA SM100+" ], "tags": [ "CUTLASS", "Blackwell", "varlen", "paged", "decode" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py", "href": "#api-cutlass-blackwell-fmha-custom-op" }, { "name": "get_splitk_heuristic", "signature": "get_splitk_heuristic(batch: int, seqlen_kv: int, kv_heads: int=1, tile_n: int=256, sm_count: int | None=None) -> int", "summary": "Chooses a decode split size from cache length, KV heads, and available SMs.", "notes": [ "Returns 0 to disable split-K when a single split would cover the sequence." ], "kind": "function", "stability": "Public", "category": "Attention", "group": "Standalone Blackwell FMHA", "module": "mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_interface", "platforms": [ "NVIDIA SM100+" ], "tags": [ "CUTLASS", "Blackwell", "varlen", "paged", "decode" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "href": "#api-get-splitk-heuristic" }, { "name": "GenKernelType", "signature": "GenKernelType.UMMA_I | GenKernelType.UMMA_P", "summary": "Selects the generated Blackwell UMMA kernel family.", "kind": "enum", "stability": "Low-level", "category": "Attention", "group": "Standalone Blackwell FMHA", "module": "mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_interface", "platforms": [ "NVIDIA SM100+" ], "tags": [ "CUTLASS", "Blackwell", "varlen", "paged", "decode" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py", "href": "#api-gen-kernel-type" }, { "name": "torch.ops.mslk.fmha_fwd", "signature": "torch.ops.mslk.fmha_fwd(query, key, value, cu_seqlens_q=None, cu_seqlens_k=None, max_seq_len_q=None, max_seq_len_k=None, softmax_scale=None, causal=False, seqlen_kv=None, page_table=None, seqlen_k=None, window_size_left=-1, window_size_right=-1, bottom_right=True) -> tuple[Tensor, Tensor]", "summary": "Standalone Blackwell forward implementation op.", "details": "Native implementation schemas used by the standalone and fMHA wrappers. Prefer the Python facades unless integrating at dispatcher level.", "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Raw Blackwell torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+", "Meta / fake" ], "tags": [ "native op", "CUTLASS", "implementation" ], "source": "csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd.cu", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd.cu", "href": "#api-mslk-fmha-fwd" }, { "name": "torch.ops.mslk.fmha_bwd", "signature": "torch.ops.mslk.fmha_bwd(dOutput, query, key, value, output, softmax_lse, cu_seqlens_q=None, cu_seqlens_k=None, max_seq_len_q=None, max_seq_len_k=None, softmax_scale=None, causal=False, window_size_left=-1, window_size_right=-1, bottom_right=True, deterministic=False) -> tuple[Tensor, Tensor, Tensor]", "summary": "Standalone Blackwell backward implementation op.", "details": "Native implementation schemas used by the standalone and fMHA wrappers. Prefer the Python facades unless integrating at dispatcher level.", "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Raw Blackwell torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+", "Meta / fake" ], "tags": [ "native op", "CUTLASS", "implementation" ], "source": "csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_bwd.cu", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_bwd.cu", "href": "#api-mslk-fmha-bwd" }, { "name": "torch.ops.mslk.fmha_gen_fwd", "signature": "torch.ops.mslk.fmha_gen_fwd(query, key, value, seqlen_kv, batch_idx=None, kernel_type=0, window_left=-1, window_right=-1, split_k_size=1024) -> tuple[Tensor, Tensor]", "summary": "Generated UMMA Blackwell forward implementation op.", "details": "Native implementation schemas used by the standalone and fMHA wrappers. Prefer the Python facades unless integrating at dispatcher level.", "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Raw Blackwell torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+", "Meta / fake" ], "tags": [ "native op", "CUTLASS", "implementation" ], "source": "csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_gen_impl.cu", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_gen_impl.cu", "href": "#api-mslk-fmha-gen-fwd" }, { "name": "torch.ops.mslk.cutlass_blackwell_fmha_fwd", "signature": "torch.ops.mslk.cutlass_blackwell_fmha_fwd(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seq_len_q, max_seq_len_k, softmax_scale, causal, seqlen_kv, page_table, seqlen_k=-1, window_size_left=-1, window_size_right=-1, bottom_right=True) -> tuple[Tensor, Tensor]", "summary": "fMHA Blackwell forward native op.", "details": "Native implementation schemas used by the standalone and fMHA wrappers. Prefer the Python facades unless integrating at dispatcher level.", "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Raw Blackwell torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+", "Meta / fake" ], "tags": [ "native op", "CUTLASS", "implementation" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py", "href": "#api-mslk-cutlass-blackwell-fmha-fwd" }, { "name": "torch.ops.mslk.cutlass_blackwell_fmha_bwd", "signature": "torch.ops.mslk.cutlass_blackwell_fmha_bwd(dout, q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, max_seq_len_q, max_seq_len_k, softmax_scale, causal, window_size_left=-1, window_size_right=-1, bottom_right=True, deterministic=False) -> tuple[Tensor, Tensor, Tensor]", "summary": "fMHA Blackwell backward native op.", "details": "Native implementation schemas used by the standalone and fMHA wrappers. Prefer the Python facades unless integrating at dispatcher level.", "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Raw Blackwell torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+", "Meta / fake" ], "tags": [ "native op", "CUTLASS", "implementation" ], "source": "mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py", "href": "#api-mslk-cutlass-blackwell-fmha-bwd" }, { "name": "triton_splitk.merge_attentions", "signature": "triton_splitk.merge_attentions(attn_out, lse_out, attn_split, lse_split) -> None", "summary": "In-place low-level reduction of split attention into preallocated outputs.", "details": "This is the allocation-free implementation layer beneath the high-level merge.", "kind": "function", "stability": "Low-level", "category": "Attention", "group": "Merge & storage utilities", "module": "mslk.attention.fmha.triton_splitk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "merge", "storage views", "autograd" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-triton-splitk-merge-attentions" }, { "name": "torch.ops.mslk.fmha_merge_attentions_varargs", "signature": "torch.ops.mslk.fmha_merge_attentions_varargs(attn_split: Tensor[], lse_split: Tensor[], write_lse, output_dtype, B, M, G, H, Kq) -> Tensor[]", "summary": "Custom op for variable-argument partial-attention merge.", "returns": "A one-element [attention] list, or [attention, LSE] when write_lse is true.", "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Merge & storage utilities", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "merge", "storage views", "autograd" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-mslk-fmha-merge-attentions-varargs" }, { "name": "torch.ops.mslk.merge_attentions_varargs_backward", "signature": "torch.ops.mslk.merge_attentions_varargs_backward(attn_split, lse_split, attn_out, lse_out, grad_attn, grad_lse) -> tuple[Tensor[], Tensor[]]", "summary": "Custom backward op for variable-argument attention merge.", "returns": "Separate lists of gradients for attention shards and LSE shards.", "kind": "torch op", "stability": "Low-level", "category": "Attention", "group": "Merge & storage utilities", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "merge", "storage views", "autograd" ], "source": "mslk/attention/fmha/triton_splitk.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py", "href": "#api-mslk-merge-attentions-varargs-backward" }, { "name": "get_stack_strides", "signature": "get_stack_strides(tensors: Sequence[Tensor], dim: int) -> Optional[tuple[int, ...]]", "summary": "Detects whether tensors are views of one common stacked storage layout.", "kind": "function", "stability": "Low-level", "category": "Attention", "group": "Merge & storage utilities", "module": "mslk.attention.fmha.unbind", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "merge", "storage views", "autograd" ], "source": "mslk/attention/fmha/unbind.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/unbind.py", "href": "#api-get-stack-strides" }, { "name": "unbind", "signature": "unbind(x: Tensor, dim: int) -> tuple[Tensor, ...]", "summary": "Autograd-aware unbind preserving shared-storage information.", "kind": "function", "stability": "Low-level", "category": "Attention", "group": "Merge & storage utilities", "module": "mslk.attention.fmha.unbind", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "merge", "storage views", "autograd" ], "source": "mslk/attention/fmha/unbind.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/unbind.py", "href": "#api-unbind" }, { "name": "stack_or_none", "signature": "stack_or_none(tensors: Sequence[Tensor], dim: int) -> Optional[Tensor]", "summary": "Returns a zero-copy/common-storage stack view when possible, else None.", "kind": "function", "stability": "Low-level", "category": "Attention", "group": "Merge & storage utilities", "module": "mslk.attention.fmha.unbind", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "merge", "storage views", "autograd" ], "source": "mslk/attention/fmha/unbind.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/unbind.py", "href": "#api-stack-or-none" }, { "name": "torch.ops.xformers.efficient_attention_forward_ck", "signature": "torch.ops.xformers.efficient_attention_forward_ck(query, key, value, attn_bias, seqstart_q, seqstart_k, max_seqlen_q, dropout_p, compute_logsumexp, custom_mask_type, scale, seqlen_k, window_size, block_tables, page_size) -> tuple[Tensor, Optional[Tensor], int, int]", "summary": "Raw CK forward attention op.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "AMD ROCm" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "csrc/attention/ck/fmha/hip_fmha/attention_forward_generic_ck_tiled.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_fmha/attention_forward_generic_ck_tiled.cpp", "href": "#api-torch-ops-xformers-efficient-attention-forward-ck" }, { "name": "torch.ops.xformers.efficient_attention_backward_ck", "signature": "torch.ops.xformers.efficient_attention_backward_ck(grad_out, query, key, value, attn_bias, seqstart_q, seqstart_k, max_seqlen_q, max_seqlen_k, seqlen_k, logsumexp, output, dropout_p, rng_seed, rng_offset, custom_mask_type, scale, window_size) -> tuple[Tensor, Tensor, Tensor, Tensor]", "summary": "Raw CK backward attention op.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "AMD ROCm" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "csrc/attention/ck/fmha/hip_fmha/attention_backward_generic_ck_tiled.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_fmha/attention_backward_generic_ck_tiled.cpp", "href": "#api-torch-ops-xformers-efficient-attention-backward-ck" }, { "name": "torch.ops.xformers.efficient_attention_forward_decoder_ck", "signature": "torch.ops.xformers.efficient_attention_forward_decoder_ck(query, key, value, seq_positions, scale) -> Tensor", "summary": "Raw CK decoder forward op.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "AMD ROCm" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "csrc/attention/ck/fmha/hip_decoder/attention_forward_decoder.hip", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_decoder/attention_forward_decoder.hip", "href": "#api-torch-ops-xformers-efficient-attention-forward-decoder-ck" }, { "name": "torch.ops.xformers.efficient_attention_forward_decoder_splitk_ck", "signature": "torch.ops.xformers.efficient_attention_forward_decoder_splitk_ck(query, key, value, seq_positions, scale, split_k) -> Tensor", "summary": "Raw CK split-K decoder forward op.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "AMD ROCm" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "csrc/attention/ck/fmha/hip_decoder/attention_forward_splitk.hip", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_decoder/attention_forward_splitk.hip", "href": "#api-torch-ops-xformers-efficient-attention-forward-decoder-splitk-ck" }, { "name": "torch.ops.xformers._ck_rand_uniform", "signature": "torch.ops.xformers._ck_rand_uniform(p: float, out: Tensor) -> Tensor", "summary": "CK random-uniform helper used by dropout implementation.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "AMD ROCm" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "csrc/attention/ck/fmha/hip_fmha/attention_ck_rand_uniform.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_fmha/attention_ck_rand_uniform.cpp", "href": "#api-torch-ops-xformers-ck-rand-uniform" }, { "name": "torch.ops.mslk_flash.flash_fwd", "signature": "torch.ops.mslk_flash.flash_fwd(query, key, value, cu_seqlens_q, cu_seqlens_k, seqused_k, max_seqlen_q, max_seqlen_k, p, softmax_scale, is_causal, window_left, window_right, return_softmax, block_tables) -> tuple[Tensor, Tensor, Tensor]", "summary": "Conditional bundled FlashAttention 2 forward op returning output, LSE, and RNG state.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "NVIDIA SM80+" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "mslk/attention/fmha/flash.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py", "href": "#api-torch-ops-mslk-flash-flash-fwd" }, { "name": "torch.ops.mslk_flash.flash_bwd", "signature": "torch.ops.mslk_flash.flash_bwd(grads_share_storage, grad, query, key, value, out, lse, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, p, softmax_scale, is_causal, window_left, window_right, rng_state) -> tuple[Tensor, Tensor, Tensor]", "summary": "Conditional bundled FlashAttention 2 backward op.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "NVIDIA SM80+" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "mslk/attention/fmha/flash.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py", "href": "#api-torch-ops-mslk-flash-flash-bwd" }, { "name": "torch.ops.mslk_flash3.flash_fwd", "signature": "torch.ops.mslk_flash3.flash_fwd(query, key, value, cu_seqlens_q, cu_seqlens_k, seqused_k, leftpad_k, max_seqlen_q, max_seqlen_k, p, softmax_scale, is_causal, descale_q=None, descale_k=None, descale_v=None, block_table=None, use_kvsplit=False, window_left=-1, window_right=-1) -> tuple[Tensor, Tensor]", "summary": "Conditional bundled FlashAttention 3 forward op returning output and LSE.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "NVIDIA SM80–SM90" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "mslk/attention/fmha/flash3.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py", "href": "#api-torch-ops-mslk-flash3-flash-fwd" }, { "name": "torch.ops.mslk_flash3.flash_bwd", "signature": "torch.ops.mslk_flash3.flash_bwd(grads_share_storage, dout, query, key, value, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, softmax_scale, is_causal, window_left, window_right) -> tuple[Tensor, Tensor, Tensor]", "summary": "Conditional bundled FlashAttention 3 backward op.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "NVIDIA SM80–SM90" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "mslk/attention/fmha/flash3.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py", "href": "#api-torch-ops-mslk-flash3-flash-bwd" }, { "name": "torch.ops.mslk_flash_mtia.flash_fwd", "signature": "torch.ops.mslk_flash_mtia.flash_fwd(query, key, value, cu_seqlens_q, cu_seqlens_k, seqused_k, max_seqlen_q, max_seqlen_k, p, softmax_scale, is_causal, window_left, window_right, return_softmax, block_tables) -> tuple[Tensor, Tensor, Tensor]", "summary": "Conditional MTIA Flash forward op returning output, LSE, and RNG state.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "MTIA" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "mslk/attention/fmha/flash_mtia.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py", "href": "#api-torch-ops-mslk-flash-mtia-flash-fwd" }, { "name": "torch.ops.mslk_flash_mtia.flash_bwd", "signature": "torch.ops.mslk_flash_mtia.flash_bwd(grads_share_storage, grad, query, key, value, out, lse, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, p, softmax_scale, is_causal, window_left, window_right, rng_state) -> tuple[Tensor, Tensor, Tensor]", "summary": "Conditional MTIA Flash backward op.", "details": "Backend implementation interfaces exposed through PyTorch’s dispatcher. Their schemas track bundled kernel code and are not stable user-level contracts; prefer fMHA operator classes.", "kind": "torch op", "stability": "Implementation", "category": "Attention", "group": "Raw backend ops", "module": "torch.ops", "platforms": [ "MTIA" ], "tags": [ "raw op", "backend implementation", "unstable" ], "source": "mslk/attention/fmha/flash_mtia.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py", "href": "#api-torch-ops-mslk-flash-mtia-flash-bwd" }, { "name": "torch.ops.mslk.bf16bf16bf16_grouped", "signature": "torch.ops.mslk.bf16bf16bf16_grouped(X: Tensor[], W: Tensor[]) -> Tensor[]", "summary": "List-of-tensors BF16 grouped GEMM.", "details": "Grouped BF16 matrix products. The preferred stacked layout uses X [total_M,K], W [G,N,K], and M_sizes [G], producing [total_M,N].", "returns": "One [Mi,N] BF16 result per group.", "notes": [ "Legacy TensorList path; deprecated/unsupported by the ROCm implementation. Prefer grouped_stacked." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "BF16 grouped torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "BF16", "grouped GEMM", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16bf16bf16-grouped" }, { "name": "torch.ops.mslk.bf16bf16bf16_grouped_cat", "signature": "torch.ops.mslk.bf16bf16bf16_grouped_cat(X: Tensor[], W: Tensor[]) -> Tensor", "summary": "List-of-tensors BF16 grouped GEMM with concatenated output.", "details": "Grouped BF16 matrix products. The preferred stacked layout uses X [total_M,K], W [G,N,K], and M_sizes [G], producing [total_M,N].", "returns": "BF16 [sum(Mi),N].", "notes": [ "Legacy path; prefer grouped_stacked." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "BF16 grouped torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "BF16", "grouped GEMM", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16bf16bf16-grouped-cat" }, { "name": "torch.ops.mslk.bf16bf16bf16_grouped_dynamic", "signature": "torch.ops.mslk.bf16bf16bf16_grouped_dynamic(X: Tensor, W: Tensor, zero_start_index_M: Tensor) -> Tensor", "summary": "Dynamic-row BF16 grouped GEMM using per-group zero starts.", "details": "Grouped BF16 matrix products. The preferred stacked layout uses X [total_M,K], W [G,N,K], and M_sizes [G], producing [total_M,N].", "notes": [ "Legacy dynamic path; deprecated/unsupported by the ROCm implementation." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "BF16 grouped torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "BF16", "grouped GEMM", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16bf16bf16-grouped-dynamic" }, { "name": "torch.ops.mslk.bf16bf16bf16_grouped_stacked", "signature": "torch.ops.mslk.bf16bf16bf16_grouped_stacked(X, W, M_sizes, out=None, num_sms=None) -> Tensor", "summary": "Preferred stacked BF16 grouped forward GEMM.", "details": "Grouped BF16 matrix products. The preferred stacked layout uses X [total_M,K], W [G,N,K], and M_sizes [G], producing [total_M,N].", "returns": "BF16 [total_M,N], optionally written into out.", "notes": [ "M_sizes gives the valid row count for each W[g].", "num_sms optionally limits participating SMs." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "BF16 grouped torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "BF16", "grouped GEMM", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16bf16bf16-grouped-stacked" }, { "name": "torch.ops.mslk.bf16bf16bf16_grouped_grad", "signature": "torch.ops.mslk.bf16bf16bf16_grouped_grad(X, W, M_sizes, out=None, num_sms=None) -> Tensor", "summary": "Grouped BF16 data-gradient matrix product.", "details": "Grouped BF16 matrix products. The preferred stacked layout uses X [total_M,K], W [G,N,K], and M_sizes [G], producing [total_M,N].", "returns": "BF16 [total_M,N].", "notes": [ "CUDA uses CUTLASS; ROCm registers the Triton implementation when mslk.gemm is imported." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "BF16 grouped torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "BF16", "grouped GEMM", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16bf16bf16-grouped-grad" }, { "name": "torch.ops.mslk.bf16bf16bf16_grouped_wgrad", "signature": "torch.ops.mslk.bf16bf16bf16_grouped_wgrad(X, W, M_sizes, output=None, output_accum=False, num_sms=None) -> Tensor", "summary": "Grouped BF16 weight-gradient product.", "details": "With X [total_M,N] and W/dY [total_M,K], produces a per-group [G,N,K] gradient.", "returns": "BF16 output, or accumulated FP32 when output_accum=True.", "notes": [ "output_accum=True requires a preallocated FP32 output tensor." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "BF16 grouped torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "BF16", "grouped GEMM", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16bf16bf16-grouped-wgrad" }, { "name": "torch.ops.mslk.f8f8bf16_blockwise", "signature": "torch.ops.mslk.f8f8bf16_blockwise(XQ, WQ, x_scale, w_scale, block_m=128, block_n=128, block_k=128) -> Tensor", "summary": "Block-scaled FP8 × FP8 GEMM.", "details": "Quantized FP8 activation/weight matrix products. Rowwise operands conventionally use XQ [...,M,K], WQ [N,K], reciprocal scales [...,M] and [N], and return BF16 [...,M,N].", "notes": [ "Scale grids correspond to the declared M/N/K block sizes." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-blockwise" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise", "signature": "torch.ops.mslk.f8f8bf16_rowwise(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor", "summary": "Primary rowwise FP8 × FP8 → BF16 GEMM.", "details": "Quantized FP8 activation/weight matrix products. Rowwise operands conventionally use XQ [...,M,K], WQ [N,K], reciprocal scales [...,M] and [N], and return BF16 [...,M,N].", "returns": "BF16 [...,M,N], with optional bias.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_out", "signature": "torch.ops.mslk.f8f8bf16_rowwise_out(XQ, WQ, x_scale, w_scale, output, bias=None, use_fast_accum=True) -> None", "summary": "In-place/out rowwise FP8 GEMM.", "details": "Quantized FP8 activation/weight matrix products. Rowwise operands conventionally use XQ [...,M,K], WQ [N,K], reciprocal scales [...,M] and [N], and return BF16 [...,M,N].", "returns": "None; mutates output.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-out" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_batched", "signature": "torch.ops.mslk.f8f8bf16_rowwise_batched(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True, output=None) -> Tensor", "summary": "Batched rowwise FP8 GEMM with optional preallocated output.", "details": "Quantized FP8 activation/weight matrix products. Rowwise operands conventionally use XQ [...,M,K], WQ [N,K], reciprocal scales [...,M] and [N], and return BF16 [...,M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-batched" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_grouped", "signature": "torch.ops.mslk.f8f8bf16_rowwise_grouped(XQ: Tensor[], WQ: Tensor[], x_scale: Tensor[], w_scale: Tensor[]) -> Tensor[]", "summary": "TensorList grouped rowwise FP8 GEMM.", "details": "Quantized FP8 activation/weight matrix products. Rowwise operands conventionally use XQ [...,M,K], WQ [N,K], reciprocal scales [...,M] and [N], and return BF16 [...,M,N].", "notes": [ "Legacy/deprecated on ROCm; prefer grouped_stacked or grouped_mm." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-grouped" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_grouped_cat", "signature": "torch.ops.mslk.f8f8bf16_rowwise_grouped_cat(XQ: Tensor[], WQ: Tensor[], x_scale: Tensor[], w_scale: Tensor[]) -> Tensor", "summary": "TensorList grouped rowwise FP8 GEMM with concatenated output.", "details": "Quantized FP8 activation/weight matrix products. Rowwise operands conventionally use XQ [...,M,K], WQ [N,K], reciprocal scales [...,M] and [N], and return BF16 [...,M,N].", "notes": [ "Legacy/deprecated on ROCm." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-grouped-cat" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_grouped_stacked", "signature": "torch.ops.mslk.f8f8bf16_rowwise_grouped_stacked(XQ, WQ, x_scale, w_scale, M_sizes) -> Tensor", "summary": "Stacked grouped rowwise FP8 GEMM.", "details": "Uses concatenated activation rows and a group/expert axis on weights, with M_sizes selecting each valid segment.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-grouped-stacked" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_grouped_dynamic", "signature": "torch.ops.mslk.f8f8bf16_rowwise_grouped_dynamic(XQ, WQ, x_scale, w_scale, zero_start_index_M, zeroing_output_tensor=True) -> Tensor", "summary": "Dynamic-row grouped FP8 GEMM.", "details": "Quantized FP8 activation/weight matrix products. Rowwise operands conventionally use XQ [...,M,K], WQ [N,K], reciprocal scales [...,M] and [N], and return BF16 [...,M,N].", "notes": [ "Legacy/deprecated on ROCm." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-grouped-dynamic" }, { "name": "torch.ops.mslk.f8f8bf16_groupwise", "signature": "torch.ops.mslk.f8f8bf16_groupwise(XQ, WQ, x_scale, w_scale) -> Tensor", "summary": "FP8 GEMM with fixed K-group scale granularity 128.", "details": "XQ [M,K], WQ [N,K], x_scale [K/128,M], and w_scale [K/128,N/128].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-groupwise" }, { "name": "torch.ops.mslk.f8f8bf16_groupwise_grouped", "signature": "torch.ops.mslk.f8f8bf16_groupwise_grouped(XQ, WQ, x_scale, w_scale, M_sizes) -> Tensor", "summary": "Stacked grouped FP8 GEMM with fixed K-group scales.", "details": "XQ [total_M,K], WQ [G,N,K], x_scale [total_M,K/128], w_scale [G,K/128,N/128].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP8 row/block/group torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "BF16 output", "reciprocal scale", "X @ W.T" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-groupwise-grouped" }, { "name": "torch.ops.mslk.f8f8f16_rowwise", "signature": "torch.ops.mslk.f8f8f16_rowwise(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor", "summary": "ROCm rowwise FP8 GEMM with FP16 output.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "ROCm FP8 torch ops", "module": "torch.ops.mslk", "platforms": [ "AMD ROCm" ], "tags": [ "FP8", "ROCm", "CK", "preshuffle" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8f16-rowwise" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_preshuffle", "signature": "torch.ops.mslk.f8f8bf16_rowwise_preshuffle(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor", "summary": "ROCm preshuffled rowwise FP8 GEMM with BF16 output.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "ROCm FP8 torch ops", "module": "torch.ops.mslk", "platforms": [ "AMD ROCm" ], "tags": [ "FP8", "ROCm", "CK", "preshuffle" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-preshuffle" }, { "name": "torch.ops.mslk.f8f8f16_rowwise_preshuffle", "signature": "torch.ops.mslk.f8f8f16_rowwise_preshuffle(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor", "summary": "ROCm preshuffled rowwise FP8 GEMM with FP16 output.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "ROCm FP8 torch ops", "module": "torch.ops.mslk", "platforms": [ "AMD ROCm" ], "tags": [ "FP8", "ROCm", "CK", "preshuffle" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8f16-rowwise-preshuffle" }, { "name": "torch.ops.mslk.f8f8bf16_rowwise_grouped_mm", "signature": "torch.ops.mslk.f8f8bf16_rowwise_grouped_mm(XQ, WQ, x_scale, w_scale, offsets=None, output=None) -> Tensor", "summary": "Generic ROCm grouped FP8 GEMM over 2D/3D layout combinations.", "details": "Supports 2D×3D expert grouping, 3D×2D output grouping, batched 3D×3D without offsets, and K-grouped 2D×2D with offsets.", "returns": "The mutated output tensor.", "notes": [ "Scales are FP32; N generally must be divisible by 8.", "offsets must be omitted for the 3D×3D layout." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "ROCm FP8 torch ops", "module": "torch.ops.mslk", "platforms": [ "AMD ROCm" ], "tags": [ "FP8", "ROCm", "CK", "preshuffle" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8f8bf16-rowwise-grouped-mm" }, { "name": "torch.ops.mslk.f4f4bf16", "signature": "torch.ops.mslk.f4f4bf16(XQ, WQ, x_scale, w_scale, output=None, global_scale=None, mxfp4_block_size=32) -> Tensor", "summary": "Packed FP4 × FP4 GEMM selecting MXFP4, MXFP4-16, or NVFP4 mode.", "details": "global_scale selects NVFP4. Without it, block size 32 selects standard MXFP4 and 16 selects MXFP4-16.", "notes": [ "For NVFP4 GEMM, global_scale is reciprocal(a_global_scale × b_global_scale).", "ROCm gfx950 supports only standard MXFP4 with block size 32: global_scale must be None and MXFP4-16/NVFP4 are unsupported." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f4f4bf16" }, { "name": "torch.ops.mslk.f4f4bf16_grouped_mm", "signature": "torch.ops.mslk.f4f4bf16_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None, global_scale=None) -> Tensor", "summary": "Offset-described grouped FP4 GEMM.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "notes": [ "On ROCm gfx950 this supports only 2D XQ × 3D WQ standard MXFP4; global_scale is accepted for schema compatibility but ignored." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f4f4bf16-grouped-mm" }, { "name": "torch.ops.mslk.f4f4bf16_grouped_stacked", "signature": "torch.ops.mslk.f4f4bf16_grouped_stacked(XQ, WQ, x_scale, w_scale, M_sizes, global_scale=None, starting_row_after_padding=None, use_mx=True) -> Tensor", "summary": "Stacked grouped FP4 GEMM with optional per-segment padding metadata.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "notes": [ "ROCm gfx950 requires use_mx=True. global_scale and starting_row_after_padding are compatibility arguments and are ignored there." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f4f4bf16-grouped-stacked" }, { "name": "torch.ops.mslk.f4f4bf16_ultra_grouped_mm", "signature": "torch.ops.mslk.f4f4bf16_ultra_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, x_global_scale, w_global_scale, output=None) -> Tensor", "summary": "Ultra grouped FP4 GEMM with separate activation/weight global scales.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM10.3+ / CUDA 13+" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f4f4bf16-ultra-grouped-mm" }, { "name": "torch.ops.mslk.mx8mx4bf16", "signature": "torch.ops.mslk.mx8mx4bf16(XQ, WQ, x_scale, w_scale, output=None) -> Tensor", "summary": "MXFP8 activation × MXFP4 weight GEMM.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "notes": [ "On ROCm, activation scales are blocked/swizzled while weight scales remain plain [N,K/32]." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-mx8mx4bf16" }, { "name": "torch.ops.mslk.mx8mx4bf16_grouped_mm", "signature": "torch.ops.mslk.mx8mx4bf16_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None) -> Tensor", "summary": "Grouped MXFP8 × MXFP4 GEMM described by offsets.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-mx8mx4bf16-grouped-mm" }, { "name": "torch.ops.mslk.mx8mx8bf16_grouped_mm", "signature": "torch.ops.mslk.mx8mx8bf16_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None, actual_num_tokens=None) -> Tensor", "summary": "Grouped MXFP8 × MXFP8 GEMM.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-mx8mx8bf16-grouped-mm" }, { "name": "torch.ops.mslk.mx8mx6bf16", "signature": "torch.ops.mslk.mx8mx6bf16(XQ, WQ, x_scale, w_scale, output=None) -> Tensor", "summary": "MXFP8 × packed MXFP6 E2M3 GEMM.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "notes": [ "Pack quantize_bf16_to_mx6_e2m3 output with pack_fp6_e2m3 first." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-mx8mx6bf16" }, { "name": "torch.ops.mslk.mx6mx6bf16", "signature": "torch.ops.mslk.mx6mx6bf16(XQ, WQ, x_scale, w_scale, output=None, splits=0) -> Tensor", "summary": "Packed MXFP6 E2M3 × MXFP6 E2M3 GEMM.", "details": "Block/microscaled matrix products. FP4 stores two values per byte; MXFP6 stores four 6-bit values per three bytes. Scale layout is part of each operator contract.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "FP4 & microscaling torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+" ], "tags": [ "FP4", "MXFP4", "MXFP8", "MXFP6", "BF16 output", "packed" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-mx6mx6bf16" }, { "name": "torch.ops.mslk.i8i8bf16", "signature": "torch.ops.mslk.i8i8bf16(XQ, WQ, scale: float, split_k=1) -> Tensor", "summary": "INT8 × INT8 GEMM with a static scalar scale.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "notes": [ "ROCm registration requires importing mslk.gemm.triton.int8_gemm.", "ROCm accepts split_k for parity but ignores non-1 values with a warning." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-i8i8bf16" }, { "name": "torch.ops.mslk.i8i8bf16_dynamic", "signature": "torch.ops.mslk.i8i8bf16_dynamic(XQ, WQ, scale: Tensor, split_k=1) -> Tensor", "summary": "INT8 × INT8 GEMM with a tensor-valued dynamic scale.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "notes": [ "ROCm registration requires importing mslk.gemm.triton.int8_gemm." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-i8i8bf16-dynamic" }, { "name": "torch.ops.mslk.bf16i4bf16_rowwise", "signature": "torch.ops.mslk.bf16i4bf16_rowwise(X, W, w_scale_group, w_zero_group) -> Tensor", "summary": "BF16 activation × row/group-quantized INT4 weight GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "notes": [ "ROCm registration requires importing mslk.gemm.triton.int4_gemm." ], "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16i4bf16-rowwise" }, { "name": "torch.ops.mslk.bf16i4bf16_rowwise_batched", "signature": "torch.ops.mslk.bf16i4bf16_rowwise_batched(X, WQ, w_scale, w_zp) -> Tensor", "summary": "Batched BF16 × rowwise INT4 weight GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16i4bf16-rowwise-batched" }, { "name": "torch.ops.mslk.bf16i4bf16_shuffled", "signature": "torch.ops.mslk.bf16i4bf16_shuffled(X, W, w_scale_group, w_zero_group) -> Tensor", "summary": "BF16 × preshuffled INT4 GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM90" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16i4bf16-shuffled" }, { "name": "torch.ops.mslk.bf16i4bf16_shuffled_batched", "signature": "torch.ops.mslk.bf16i4bf16_shuffled_batched(X, WQ, w_scale, w_zp) -> Tensor", "summary": "Batched BF16 × preshuffled INT4 GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM90" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16i4bf16-shuffled-batched" }, { "name": "torch.ops.mslk.bf16i4bf16_shuffled_grouped", "signature": "torch.ops.mslk.bf16i4bf16_shuffled_grouped(X, WQ, w_scale_group, w_zero_group, M_sizes) -> Tensor", "summary": "Stacked grouped BF16 × preshuffled INT4 GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM90" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16i4bf16-shuffled-grouped" }, { "name": "torch.ops.mslk.f8i4bf16_rowwise", "signature": "torch.ops.mslk.f8i4bf16_rowwise(XQ, WQ, x_scale, w_scale, w_zp) -> Tensor", "summary": "Rowwise FP8 activation × INT4 weight GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM90" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8i4bf16-rowwise" }, { "name": "torch.ops.mslk.f8i4bf16_shuffled", "signature": "torch.ops.mslk.f8i4bf16_shuffled(XQ, WQ, x_scale, w_scale, w_scale_group) -> Tensor", "summary": "FP8 activation × preshuffled INT4 weight GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM90" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8i4bf16-shuffled" }, { "name": "torch.ops.mslk.f8i4bf16_shuffled_grouped", "signature": "torch.ops.mslk.f8i4bf16_shuffled_grouped(XQ, WQ, x_scale, w_scale, w_scale_group, M_sizes) -> Tensor", "summary": "Stacked grouped FP8 × preshuffled INT4 GEMM.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM90" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-f8i4bf16-shuffled-grouped" }, { "name": "torch.ops.mslk.preshuffle_i4", "signature": "torch.ops.mslk.preshuffle_i4(WQ, w_scale) -> tuple[Tensor, Tensor]", "summary": "Preprocesses packed INT4 weights and scales for shuffled CUDA kernels.", "details": "Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-preshuffle-i4" }, { "name": "torch.ops.mslk.bf16x9_gemm", "signature": "torch.ops.mslk.bf16x9_gemm(A, B, output=None) -> Tensor", "summary": "cuBLAS BF16x9-emulation GEMM over FP32 inputs and output.", "details": "Despite its name, A [M,K] and B [N,K] are contiguous FP32 and the result [M,N] is FP32.", "kind": "torch op", "stability": "Public", "category": "GEMM", "group": "INT and mixed torch ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA / CUDA 13+" ], "tags": [ "INT8", "INT4", "mixed input", "BF16 output", "packed weights" ], "source": "csrc/gemm/gemm_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp", "href": "#api-mslk-bf16x9-gemm" }, { "name": "matmul_fp8_row", "signature": "matmul_fp8_row(a, b, a_scale, b_scale, bias=None, dot_out_dtype=None, allow_tf32=True, fp8_fast_accum=True, imprecise_acc=False, tma_persistent=True, no_use_persistent=None, use_warp_specialization=False) -> Tensor", "summary": "Feature-rich Triton rowwise FP8 matrix product.", "details": "Computes a @ b.T for 2D or higher-rank a, applies reciprocal dequantization scales and optional bias, and returns BF16 by default.", "notes": [ "The source docstring’s division formula is stale; implementation/dequantization convention multiplies reciprocal scales.", "Persistent/TMA/warp-specialized paths are architecture-dependent." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton FP8 GEMM", "module": "mslk.gemm.triton.fp8_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "FP8", "rowwise", "blockwise", "X @ W.T" ], "source": "mslk/gemm/triton/fp8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py", "href": "#api-matmul-fp8-row" }, { "name": "matmul_fp8_block", "signature": "matmul_fp8_block(a, b, a_scale, b_scale, scale_block_m=256, scale_block_n=256, scale_block_k=256, dot_out_dtype=None, allow_tf32=True, fp8_fast_accum=True) -> Tensor", "summary": "Triton block-scaled FP8 matrix product.", "details": "Computes a @ b.T using an explicit 3D scale-block contract.", "notes": [ "Accelerator-only." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton FP8 GEMM", "module": "mslk.gemm.triton.fp8_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "FP8", "rowwise", "blockwise", "X @ W.T" ], "source": "mslk/gemm/triton/fp8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py", "href": "#api-matmul-fp8-block" }, { "name": "to_mxfp8", "signature": "to_mxfp8(data_hp: Tensor, block_size=32) -> tuple[Tensor, Tensor]", "summary": "Converts high-precision data to MXFP8 with E8M0 block scales.", "returns": "(scale_e8m0_biased, data_fp8).", "notes": [ "Return order is scales first, data second—the reverse of most quantizers." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton FP8 GEMM", "module": "mslk.gemm.triton.fp8_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "FP8", "rowwise", "blockwise", "X @ W.T" ], "source": "mslk/gemm/triton/fp8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py", "href": "#api-to-mxfp8" }, { "name": "torch.ops.triton.matmul_fp8_row", "signature": "torch.ops.triton.matmul_fp8_row(a, b, a_scale, b_scale, bias=None, dot_out_dtype=None, allow_tf32=True, fp8_fast_accum=True, imprecise_acc=False, tma_persistent=True, no_use_persistent=None, use_warp_specialization=False) -> Tensor", "summary": "Custom-op registration for matmul_fp8_row.", "notes": [ "Runtime accepts higher-rank a by flattening its leading dimensions, but the registered fake/Meta implementation assumes 2D a and b; torch.compile/export may fail for higher-rank inputs." ], "kind": "torch op", "stability": "Low-level", "category": "GEMM", "group": "Triton FP8 GEMM", "module": "torch.ops.triton", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "FP8", "rowwise", "blockwise", "X @ W.T" ], "source": "mslk/gemm/triton/fp8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py", "href": "#api-triton-matmul-fp8-row" }, { "name": "torch.ops.triton.matmul_fp8_block", "signature": "torch.ops.triton.matmul_fp8_block(a, b, a_scale, b_scale, scale_block_m=256, scale_block_n=256, scale_block_k=256, dot_out_dtype=None, allow_tf32=True, fp8_fast_accum=True) -> Tensor", "summary": "Custom-op registration for matmul_fp8_block.", "notes": [ "Runtime accepts higher-rank a by flattening its leading dimensions, but the registered fake/Meta implementation assumes 2D a and b.", "The fake/Meta implementation always reports BF16 output even when dot_out_dtype requests another dtype." ], "kind": "torch op", "stability": "Low-level", "category": "GEMM", "group": "Triton FP8 GEMM", "module": "torch.ops.triton", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "FP8", "rowwise", "blockwise", "X @ W.T" ], "source": "mslk/gemm/triton/fp8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py", "href": "#api-triton-matmul-fp8-block" }, { "name": "grouped_gemm", "signature": "grouped_gemm(x, w, m_sizes, bias=None, token_weights=None, use_fast_accum=True, *, _use_warp_specialization=True, _output_tensor=None, _scatter_add_indices=None) -> Tensor", "summary": "General BF16 grouped GEMM with optional fused bias/router scaling.", "details": "Supports per-expert bias, per-token router weights, preallocated output, and an internal fused scatter-add path used by MoE.", "returns": "[total_M,N].", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton grouped GEMM", "module": "mslk.gemm.triton.grouped_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "grouped", "MoE", "M_sizes" ], "source": "mslk/gemm/triton/grouped_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py", "href": "#api-grouped-gemm" }, { "name": "grouped_gemm_fp8_rowwise", "signature": "grouped_gemm_fp8_rowwise(x, w, m_sizes, x_scale, w_scale, use_fast_accum=True, *, _use_warp_specialization=True, _output_tensor=None, _scatter_add_indices=None) -> Tensor", "summary": "Rowwise FP8 grouped GEMM for expert layers.", "details": "Concatenated-token grouped GEMMs with one weight matrix per group/expert. M_sizes partitions x along its first dimension.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton grouped GEMM", "module": "mslk.gemm.triton.grouped_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "grouped", "MoE", "M_sizes" ], "source": "mslk/gemm/triton/grouped_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py", "href": "#api-grouped-gemm-fp8-rowwise" }, { "name": "grouped_gemm_dgrad", "signature": "grouped_gemm_dgrad(x, w, m_sizes, out=None, num_sms=None) -> Tensor", "summary": "Triton grouped BF16 data-gradient product.", "details": "Concatenated-token grouped GEMMs with one weight matrix per group/expert. M_sizes partitions x along its first dimension.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton grouped GEMM", "module": "mslk.gemm.triton.grouped_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "grouped", "MoE", "M_sizes" ], "source": "mslk/gemm/triton/grouped_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py", "href": "#api-grouped-gemm-dgrad" }, { "name": "grouped_gemm_wgrad", "signature": "grouped_gemm_wgrad(x, w, m_sizes, output=None, output_accum=False, num_sms=None) -> Tensor", "summary": "Triton grouped BF16 weight-gradient product.", "details": "Concatenated-token grouped GEMMs with one weight matrix per group/expert. M_sizes partitions x along its first dimension.", "notes": [ "FP32 preallocated output is required when accumulating." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton grouped GEMM", "module": "mslk.gemm.triton.grouped_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "grouped", "MoE", "M_sizes" ], "source": "mslk/gemm/triton/grouped_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py", "href": "#api-grouped-gemm-wgrad" }, { "name": "matmul_f8f8bf16_groupwise", "signature": "matmul_f8f8bf16_groupwise(XQ, WQ, x_scale, w_scale, output=None) -> Tensor", "summary": "Triton FP8 GEMM with K-group scale granularity 128.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.fp8_groupwise_gemm", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/fp8_groupwise_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_groupwise_gemm.py", "href": "#api-matmul-f8f8bf16-groupwise" }, { "name": "matmul_f8f8bf16_groupwise_grouped", "signature": "matmul_f8f8bf16_groupwise_grouped(XQ, WQ, x_scale, w_scale, M_sizes, output=None) -> Tensor", "summary": "Triton stacked grouped FP8 groupwise GEMM.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.fp8_groupwise_grouped_gemm", "platforms": [ "NVIDIA", "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/fp8_groupwise_grouped_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_groupwise_grouped_gemm.py", "href": "#api-matmul-f8f8bf16-groupwise-grouped" }, { "name": "matmul_mx8mx4bf16", "signature": "matmul_mx8mx4bf16(XQ, WQ, x_scale, w_scale, output=None) -> Tensor", "summary": "Triton MXFP8 × MXFP4 GEMM.", "notes": [ "ROCm gfx950 native block-scaled MFMA.", "ROCm activation scales are blocked/swizzled; weight scales are row-major [N,K/32]." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.mx8mx4_gemm", "platforms": [ "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/mx8mx4_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/mx8mx4_gemm.py", "href": "#api-matmul-mx8mx4bf16" }, { "name": "matmul_mx8mx4bf16_grouped", "signature": "matmul_mx8mx4bf16_grouped(XQ, WQ, x_scale, w_scale, offsets, output=None) -> Tensor", "summary": "Triton grouped MXFP8 × MXFP4 GEMM.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.mx8mx4_gemm", "platforms": [ "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/mx8mx4_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/mx8mx4_gemm.py", "href": "#api-matmul-mx8mx4bf16-grouped" }, { "name": "matmul_mx8mx8bf16_grouped", "signature": "matmul_mx8mx8bf16_grouped(XQ, WQ, x_scale, w_scale, offsets, output=None, actual_num_tokens=None) -> Tensor", "summary": "Triton grouped MXFP8 × MXFP8 GEMM.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.mx8mx8_gemm", "platforms": [ "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/mx8mx8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/mx8mx8_gemm.py", "href": "#api-matmul-mx8mx8bf16-grouped" }, { "name": "mxfp4_gemm", "signature": "mxfp4_gemm(XQ, WQ, x_scale, w_scale, output=None) -> Tensor", "summary": "Triton packed MXFP4 × MXFP4 GEMM.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.f4f4bf16", "platforms": [ "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/f4f4bf16.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py", "href": "#api-mxfp4-gemm" }, { "name": "mxfp4_grouped_mm", "signature": "mxfp4_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None, global_scale=None) -> Tensor", "summary": "ROCm gfx950 offset-grouped standard MXFP4 GEMM.", "notes": [ "Supports only XQ [total_M,K/2] with transposed WQ [G,K/2,N]; the CUDA 2D×2D K-grouped layout is unsupported.", "global_scale is accepted for schema compatibility but ignored; leave it None." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.f4f4bf16", "platforms": [ "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/f4f4bf16.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py", "href": "#api-mxfp4-grouped-mm" }, { "name": "mxfp4_grouped_stacked_gemm", "signature": "mxfp4_grouped_stacked_gemm(XQ, WQ, x_scale, w_scale, M_sizes, output=None, global_scale=None, starting_row_after_padding=None, use_mx=True, offsets_override=None) -> Tensor", "summary": "ROCm gfx950 stacked grouped standard MXFP4 GEMM.", "notes": [ "use_mx must be True; NVFP4 is unsupported.", "global_scale and starting_row_after_padding are accepted for compatibility but ignored.", "offsets_override is an AMD-only optional cumulative-offset fast path." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.f4f4bf16", "platforms": [ "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/f4f4bf16.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py", "href": "#api-mxfp4-grouped-stacked-gemm" }, { "name": "f4f4bf16", "signature": "f4f4bf16(XQ, WQ, x_scale, w_scale, output=None, global_scale=None, mxfp4_block_size=32) -> Tensor", "summary": "ROCm compatibility wrapper for standard MXFP4 block-size-32 GEMM.", "notes": [ "Requires gfx950.", "global_scale must be None; NVFP4 is unsupported.", "mxfp4_block_size must be 32; MXFP4-16 is CUDA-only." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton groupwise & microscaling", "module": "mslk.gemm.triton.f4f4bf16", "platforms": [ "AMD gfx950" ], "tags": [ "Triton", "groupwise", "MXFP4", "MXFP8", "gfx950" ], "source": "mslk/gemm/triton/f4f4bf16.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py", "href": "#api-triton-f4f4bf16" }, { "name": "matmul_bf16i4_rowwise", "signature": "matmul_bf16i4_rowwise(X, W, w_scale_group, w_zero_group) -> Tensor", "summary": "Triton BF16 activation × row/group INT4 weight GEMM.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton integer GEMM", "module": "mslk.gemm.triton.int4_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "INT4", "INT8", "BF16" ], "source": "mslk/gemm/triton/int4_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int4_gemm.py", "href": "#api-matmul-bf16i4-rowwise" }, { "name": "matmul_bf16i4_rowwise_batched", "signature": "matmul_bf16i4_rowwise_batched(X, W, w_scale, w_zp) -> Tensor", "summary": "Triton batched BF16 × INT4 GEMM.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton integer GEMM", "module": "mslk.gemm.triton.int4_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "INT4", "INT8", "BF16" ], "source": "mslk/gemm/triton/int4_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int4_gemm.py", "href": "#api-matmul-bf16i4-rowwise-batched" }, { "name": "i8i8bf16_triton", "signature": "i8i8bf16_triton(XQ, WQ, scale: float, split_k=1) -> Tensor", "summary": "Triton INT8 × INT8 GEMM with scalar scale.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton integer GEMM", "module": "mslk.gemm.triton.int8_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "INT4", "INT8", "BF16" ], "source": "mslk/gemm/triton/int8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int8_gemm.py", "href": "#api-i8i8bf16-triton" }, { "name": "i8i8bf16_dynamic_triton", "signature": "i8i8bf16_dynamic_triton(XQ, WQ, scale: Tensor, split_k=1) -> Tensor", "summary": "Triton INT8 × INT8 GEMM with tensor scale.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Triton integer GEMM", "module": "mslk.gemm.triton.int8_gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "Triton", "INT4", "INT8", "BF16" ], "source": "mslk/gemm/triton/int8_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int8_gemm.py", "href": "#api-i8i8bf16-dynamic-triton" }, { "name": "MixedInputGemmKernel", "signature": "MixedInputGemmKernel(scale_granularity_m, scale_granularity_k, acc_dtype, use_2cta_instrs, mma_tiler_mnk, cluster_shape_mn, use_tma_store)", "summary": "Configurable compiled Blackwell mixed-input GEMM kernel object.", "details": "CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.", "methods": [ { "signature": "__call__(a, a_scale, b, c, max_active_clusters, stream)", "summary": "Launches a configured kernel." }, { "signature": "can_implement(...)", "summary": "Checks alignment, tiler, cluster, layout, and epilogue constraints." } ], "kind": "class", "stability": "Advanced", "category": "GEMM", "group": "Blackwell mixed-input CuTeDSL", "module": "mslk.gemm.blackwell_mixed_input_gemm", "platforms": [ "NVIDIA SM100" ], "tags": [ "CuTeDSL", "Blackwell", "INT4", "INT8", "BF16", "mixed input" ], "source": "mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "href": "#api-mixed-input-gemm-kernel" }, { "name": "mixed_input_gemm", "signature": "mixed_input_gemm(A, B, A_scale=None, C=None, scale_granularity_m=1, scale_granularity_k=128, acc_dtype=None, mma_tiler_mnk=(128,128,128), cluster_shape_mn=(1,1), use_2cta_instrs=False, use_tma_store=False) -> Tensor", "summary": "General Blackwell mixed-input GEMM frontend.", "details": "CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.", "notes": [ "Without TMA store, output dimensions must not leave out-of-bounds epilogue tiles." ], "kind": "function", "stability": "Public", "category": "GEMM", "group": "Blackwell mixed-input CuTeDSL", "module": "mslk.gemm.blackwell_mixed_input_gemm", "platforms": [ "NVIDIA SM100" ], "tags": [ "CuTeDSL", "Blackwell", "INT4", "INT8", "BF16", "mixed input" ], "source": "mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "href": "#api-mixed-input-gemm" }, { "name": "int4bf16bf16_gemm", "signature": "int4bf16bf16_gemm(A, B, A_scale, C=None, scale_granularity_m=1, scale_granularity_k=128, acc_dtype=None, mma_tiler_mnk=(128,128,128), cluster_shape_mn=(1,1), use_2cta_instrs=False, use_tma_store=False) -> Tensor", "summary": "Convenience INT4 narrow operand × BF16 wide operand GEMM.", "details": "CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Blackwell mixed-input CuTeDSL", "module": "mslk.gemm.blackwell_mixed_input_gemm", "platforms": [ "NVIDIA SM100" ], "tags": [ "CuTeDSL", "Blackwell", "INT4", "INT8", "BF16", "mixed input" ], "source": "mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "href": "#api-int4bf16bf16-gemm" }, { "name": "int8bf16bf16_gemm", "signature": "int8bf16bf16_gemm(A, B, C=None, acc_dtype=None, mma_tiler_mnk=(128,128,128), cluster_shape_mn=(1,1), use_2cta_instrs=False, use_tma_store=False) -> Tensor", "summary": "Convenience INT8 narrow operand × BF16 wide operand GEMM.", "details": "CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.", "kind": "function", "stability": "Public", "category": "GEMM", "group": "Blackwell mixed-input CuTeDSL", "module": "mslk.gemm.blackwell_mixed_input_gemm", "platforms": [ "NVIDIA SM100" ], "tags": [ "CuTeDSL", "Blackwell", "INT4", "INT8", "BF16", "mixed input" ], "source": "mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "href": "#api-int8bf16bf16-gemm" }, { "name": "create_tensors", "signature": "create_tensors(l, m, n, k, a_major, b_major, c_major, a_dtype, b_dtype, c_dtype, scale_granularity_m=0, scale_granularity_k=0)", "summary": "Creates benchmark/reference tensors and layouts for mixed-input GEMM.", "details": "CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.", "kind": "function", "stability": "Developer", "category": "GEMM", "group": "Blackwell mixed-input CuTeDSL", "module": "mslk.gemm.blackwell_mixed_input_gemm", "platforms": [ "NVIDIA SM100" ], "tags": [ "CuTeDSL", "Blackwell", "INT4", "INT8", "BF16", "mixed input" ], "source": "mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "href": "#api-create-tensors" }, { "name": "compare", "signature": "compare(a_torch_cpu, b_torch_cpu, a_scale_torch_cpu, c_torch_gpu, c_dtype, tolerance) -> None", "summary": "Compares kernel output with a CPU reference.", "details": "CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.", "kind": "function", "stability": "Developer", "category": "GEMM", "group": "Blackwell mixed-input CuTeDSL", "module": "mslk.gemm.blackwell_mixed_input_gemm", "platforms": [ "NVIDIA SM100" ], "tags": [ "CuTeDSL", "Blackwell", "INT4", "INT8", "BF16", "mixed input" ], "source": "mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "href": "#api-compare" }, { "name": "run", "signature": "run(mnkl, scale_granularity_m, scale_granularity_k, a_dtype, b_dtype, c_dtype, acc_dtype, a_major, b_major, c_major, mma_tiler_mnk, cluster_shape_mn, use_2cta_instrs, use_tma_store, tolerance, warmup_iterations=0, iterations=1, skip_ref_check=False, use_cold_l2=False, **kwargs)", "summary": "Benchmark/validation driver for mixed-input configurations.", "details": "CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.", "kind": "function", "stability": "Developer", "category": "GEMM", "group": "Blackwell mixed-input CuTeDSL", "module": "mslk.gemm.blackwell_mixed_input_gemm", "platforms": [ "NVIDIA SM100" ], "tags": [ "CuTeDSL", "Blackwell", "INT4", "INT8", "BF16", "mixed input" ], "source": "mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py", "href": "#api-run" }, { "name": "triton_quantize_fp8_row", "signature": "triton_quantize_fp8_row(a, scale_ub=None, zero_start_index_M=None, align_rows_to=None, eps_opt=1/512) -> tuple[Tensor, Tensor]", "summary": "Triton rowwise FP8 quantizer.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "returns": "(FP8 data, reciprocal row scales).", "notes": [ "zero_start_index_M gives the number of valid rows for each logical matrix; suffix rows at or beyond that count are emitted as zeros.", "align_rows_to rounds up the quantized output’s last dimension; the reciprocal-scale tensor retains shape a.shape[:-1]." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-triton-quantize-fp8-row" }, { "name": "quantize_fp8_row", "signature": "quantize_fp8_row(a, scale_ub=None, zero_start_index_M=None, use_triton=True, output_device=None, align_rows_to=None, eps_opt=1/512) -> tuple[Tensor, Tensor]", "summary": "Public rowwise FP8 adapter selecting the Triton kernel or a PyTorch fallback.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "returns": "(FP8 data, reciprocal row scales).", "notes": [ "The Triton path honors zero_start_index_M and align_rows_to as described above.", "The PyTorch fallback honors scale_ub and output_device but currently ignores zero_start_index_M, align_rows_to, and eps_opt." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-quantize-fp8-row" }, { "name": "scale_fp8_row", "signature": "scale_fp8_row(a, x_scale, w_scale) -> Tensor", "summary": "Applies activation and weight row scales to an accumulator/result tensor.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-scale-fp8-row" }, { "name": "triton_quantize_fp8_block", "signature": "triton_quantize_fp8_block(x, block_m=256, block_k=256, scale_ub=None, k_major=True) -> tuple[Tensor, Tensor]", "summary": "Triton 2D blockwise FP8 quantizer.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "returns": "(FP8 data, reciprocal block-scale grid).", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-triton-quantize-fp8-block" }, { "name": "quantize_fp8_block", "signature": "quantize_fp8_block(x, block_m=256, block_k=256, scale_ub=None, use_triton=True, output_device=None, k_major=True) -> tuple[Tensor, Tensor]", "summary": "Public blockwise FP8 adapter.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-quantize-fp8-block" }, { "name": "triton_quantize_fp8_group", "signature": "triton_quantize_fp8_group(x, group_size=128, scale_ub=None, m_sizes=None, k_major=True) -> tuple[Tensor, Tensor]", "summary": "Triton K-group FP8 quantizer, optionally aware of grouped valid row counts.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-triton-quantize-fp8-group" }, { "name": "quantize_fp8_group", "signature": "quantize_fp8_group(x, group_size=128, scale_ub=None, m_sizes=None, k_major=True, use_triton=True, output_device=None) -> tuple[Tensor, Tensor]", "summary": "Public groupwise FP8 adapter.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-quantize-fp8-group" }, { "name": "triton_quantize_fp8_tensor", "signature": "triton_quantize_fp8_tensor(a) -> tuple[Tensor, Tensor]", "summary": "Triton tensorwise FP8 quantizer.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "returns": "(FP8 data, scalar reciprocal scale).", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-triton-quantize-fp8-tensor" }, { "name": "quantize_fp8_tensor", "signature": "quantize_fp8_tensor(a, use_triton=True) -> tuple[Tensor, Tensor]", "summary": "Public tensorwise FP8 adapter.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-quantize-fp8-tensor" }, { "name": "dequantize_fp8_row", "signature": "dequantize_fp8_row(xq, x_scale) -> Tensor", "summary": "Dequantizes rowwise FP8 storage to BF16.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-dequantize-fp8-row" }, { "name": "dequantize_fp8_block", "signature": "dequantize_fp8_block(xq, x_scale, block_m=256, block_k=256) -> Tensor", "summary": "Dequantizes blockwise FP8 storage to BF16.", "details": "Quantizes to the platform FP8 flavor and returns reciprocal scales consumed by MSLK GEMMs. CUDA uses E4M3FN; compatible ROCm targets use E4M3FNUZ, while gfx950 uses OCP E4M3FN.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 quantization", "module": "mslk.quantize.triton.fp8_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP8", "E4M3", "reciprocal scale", "Triton" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-dequantize-fp8-block" }, { "name": "torch.ops.triton.quantize_fp8_row", "signature": "torch.ops.triton.quantize_fp8_row(a, scale_ub=None, zero_start_index_M=None, use_triton=True, output_device=None, align_rows_to=None, eps_opt=1/512) -> tuple[Tensor, Tensor]", "summary": "Custom-op registration for rowwise FP8 quantization.", "notes": [ "The Meta implementation currently omits the runtime eps_opt argument." ], "kind": "torch op", "stability": "Low-level", "category": "Quantization", "group": "FP8 custom ops & dtype helpers", "module": "torch.ops.triton", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "torch.compile", "custom op", "FP8 dtype" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-torch-ops-triton-quantize-fp8-row" }, { "name": "torch.ops.triton.quantize_fp8_block", "signature": "torch.ops.triton.quantize_fp8_block(x, block_m=256, block_k=256, scale_ub=None, use_triton=True, output_device=None, k_major=True) -> tuple[Tensor, Tensor]", "summary": "Custom-op registration for blockwise FP8 quantization.", "kind": "torch op", "stability": "Low-level", "category": "Quantization", "group": "FP8 custom ops & dtype helpers", "module": "torch.ops.triton", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "torch.compile", "custom op", "FP8 dtype" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-torch-ops-triton-quantize-fp8-block" }, { "name": "torch.ops.triton.quantize_fp8_tensor", "signature": "torch.ops.triton.quantize_fp8_tensor(a, use_triton=True) -> (Tensor, Tensor)", "summary": "Custom-op registration for tensorwise FP8 quantization.", "kind": "torch op", "stability": "Low-level", "category": "Quantization", "group": "FP8 custom ops & dtype helpers", "module": "torch.ops.triton", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "torch.compile", "custom op", "FP8 dtype" ], "source": "mslk/quantize/triton/fp8_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py", "href": "#api-torch-ops-triton-quantize-fp8-tensor" }, { "name": "get_fp8_constants", "signature": "get_fp8_constants() -> tuple[torch.dtype, tl.dtype, float, float]", "summary": "Returns the platform PyTorch dtype, Triton dtype, maximum finite value, and epsilon.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "FP8 custom ops & dtype helpers", "module": "mslk.utils.triton.fp8_utils", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "torch.compile", "custom op", "FP8 dtype" ], "source": "mslk/utils/triton/fp8_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/triton/fp8_utils.py", "href": "#api-get-fp8-constants" }, { "name": "reinterpret_fp8_type", "signature": "reinterpret_fp8_type(tensor: Tensor, dtype: tl.dtype) -> TensorWrapper", "summary": "Reinterprets FP8 storage for a Triton dtype without numeric conversion.", "kind": "function", "stability": "Advanced", "category": "Quantization", "group": "FP8 custom ops & dtype helpers", "module": "mslk.utils.triton.fp8_utils", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "torch.compile", "custom op", "FP8 dtype" ], "source": "mslk/utils/triton/fp8_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/triton/fp8_utils.py", "href": "#api-reinterpret-fp8-type" }, { "name": "triton_quantize_mx4", "signature": "triton_quantize_mx4(input, *, rounding_mode=RoundingMode.ceil, seed=None) -> tuple[Tensor, Tensor]", "summary": "Preferred public MXFP4 group-32 quantizer.", "details": "Quantizes BF16/FP16 values into packed E2M1 elements with shared E8M0 scales. The last dimension must be divisible by the group size.", "returns": "Packed uint8 [...,K/2] and E8M0 scales.", "notes": [ "CUDA scale storage is flattened 128×4 blocked/swizzled int8.", "ROCm scale storage is plain [...,K/32] uint8.", "Stochastic rounding is CUDA-only." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "MXFP4 quantization", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+", "AMD gfx950" ], "tags": [ "MXFP4", "E2M1", "E8M0", "packed", "rounding" ], "source": "mslk/quantize/triton/fp4_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py", "href": "#api-triton-quantize-mx4" }, { "name": "triton_quantize_mx4_unpack", "signature": "triton_quantize_mx4_unpack(input, group_size=32, ebits=2, mbits=1, rounding_mode=RoundingMode.ceil, stochastic_casting=False, *, seed=None) -> tuple[Tensor, Tensor]", "summary": "Compatibility adapter supporting group size 16/32 and legacy rounding arguments.", "details": "Quantizes BF16/FP16 values into packed E2M1 elements with shared E8M0 scales. The last dimension must be divisible by the group size.", "notes": [ "Only E2M1 (ebits=2, mbits=1) is implemented.", "stochastic_casting=True forces stochastic rounding." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "MXFP4 quantization", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+", "AMD gfx950" ], "tags": [ "MXFP4", "E2M1", "E8M0", "packed", "rounding" ], "source": "mslk/quantize/triton/fp4_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py", "href": "#api-triton-quantize-mx4-unpack" }, { "name": "quantize_mx4", "signature": "quantize_mx4(x, group_size=32, rounding_mode=RoundingMode.ceil, *, seed=None) -> tuple[Tensor, Tensor]", "summary": "Direct MXFP4 kernel frontend underlying the public adapter.", "details": "Quantizes BF16/FP16 values into packed E2M1 elements with shared E8M0 scales. The last dimension must be divisible by the group size.", "notes": [ "For a 1D input, preserves an extra row dimension; the public adapter restores the expected shape." ], "kind": "function", "stability": "Advanced", "category": "Quantization", "group": "MXFP4 quantization", "module": "mslk.quantize.triton.quantize_kernels.mx4", "platforms": [ "NVIDIA SM100+", "AMD gfx950" ], "tags": [ "MXFP4", "E2M1", "E8M0", "packed", "rounding" ], "source": "mslk/quantize/triton/quantize_kernels/mx4.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/quantize_kernels/mx4.py", "href": "#api-quantize-mx4" }, { "name": "quantize_mx4_stacked", "signature": "quantize_mx4_stacked(m_sizes, x, group_size=32, rounding_mode=RoundingMode.ceil, *, seed=None) -> tuple[Tensor, Tensor]", "summary": "Segment-aware MXFP4 quantization for stacked grouped inputs.", "details": "Quantizes BF16/FP16 values into packed E2M1 elements with shared E8M0 scales. The last dimension must be divisible by the group size.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "MXFP4 quantization", "module": "mslk.quantize.triton.quantize_kernels.mx4_stacked", "platforms": [ "NVIDIA SM100+", "AMD gfx950" ], "tags": [ "MXFP4", "E2M1", "E8M0", "packed", "rounding" ], "source": "mslk/quantize/triton/quantize_kernels/mx4_stacked.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/quantize_kernels/mx4_stacked.py", "href": "#api-quantize-mx4-stacked" }, { "name": "triton_quantize_nvfp4", "signature": "triton_quantize_nvfp4(x, global_scale, use_e8m0_scale=False, use_precise_math=True) -> tuple[Tensor, Tensor]", "summary": "NVFP4 quantizer with per-16-value E4M3 scale factors.", "returns": "(packed FP4, swizzled scales).", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/quantize.py", "href": "#api-triton-quantize-nvfp4" }, { "name": "triton_fake_quantize_nvfp4_per_tensor", "signature": "triton_fake_quantize_nvfp4_per_tensor(input, static_scales=None, scale_ub=None) -> tuple[Tensor, Tensor]", "summary": "Per-tensor NVFP4 fake quantization returning BF16 values.", "returns": "(BF16 fake-quantized tensor, FP32 scale or amax).", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/fake_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fake_quantize.py", "href": "#api-triton-fake-quantize-nvfp4-per-tensor" }, { "name": "nvfp4_quantize_stacked", "signature": "nvfp4_quantize_stacked(m_sizes, input, global_scale) -> tuple[Tensor, Tensor]", "summary": "Segment-aware stacked NVFP4 quantization.", "returns": "(packed FP4, padded two-dimensional swizzled scale buffer).", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/quantize_stacked.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/quantize_stacked.py", "href": "#api-nvfp4-quantize-stacked" }, { "name": "nvfp4_quantize_stacked_with_token_scale", "signature": "nvfp4_quantize_stacked_with_token_scale(m_sizes, input) -> tuple[Tensor, Tensor, Tensor]", "summary": "Stacked NVFP4 quantization deriving an inverse scale per token.", "returns": "(packed FP4, padded two-dimensional swizzled scale buffer, token_scale_inv).", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/quantize_stacked.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/quantize_stacked.py", "href": "#api-nvfp4-quantize-stacked-with-token-scale" }, { "name": "calculate_group_max", "signature": "calculate_group_max(input, m_sizes) -> tuple[Tensor, Tensor]", "summary": "Computes per-segment NVFP4 global scales and a row-to-segment map.", "details": "The name is misleading: the first output is effectively 448×6/max_abs, not the raw maximum.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/global_scale.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/global_scale.py", "href": "#api-calculate-group-max" }, { "name": "cal_global_scale_mx4_as_nvfp4", "signature": "cal_global_scale_mx4_as_nvfp4(x) -> Tensor", "summary": "Derives an NVFP4-style global scale for MXFP4 data.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/primitives.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py", "href": "#api-cal-global-scale-mx4-as-nvfp4" }, { "name": "global_scale_nvfp4", "signature": "global_scale_nvfp4(x) -> Tensor", "summary": "Computes the standard global NVFP4 scale from input range.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/fp4_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py", "href": "#api-global-scale-nvfp4" }, { "name": "fp4_to_float", "signature": "fp4_to_float(x) -> Tensor", "summary": "Decodes FP4 nibbles to FP32 values.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/fp4_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py", "href": "#api-fp4-to-float" }, { "name": "scale_nvfp4", "signature": "scale_nvfp4(x, scale, global_scale, group_size=16) -> Tensor", "summary": "Applies NVFP4 local and global scaling to decoded values.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/fp4_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py", "href": "#api-scale-nvfp4" }, { "name": "dequantize_nvfp4", "signature": "dequantize_nvfp4(input_quantized, scale, global_scale, group_size=16) -> Tensor", "summary": "Dequantizes packed NVFP4 to BF16.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/fp4_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py", "href": "#api-dequantize-nvfp4" }, { "name": "dequantize_mx4", "signature": "dequantize_mx4(input_quantized, scale, group_size=32) -> Tensor", "summary": "Dequantizes packed MXFP4 with NVIDIA blocked scales to BF16.", "notes": [ "Not compatible with the plain row-major scale output produced on ROCm." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "NVFP4 & legacy FP4", "module": "mslk.quantize.triton.fp4_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "NVFP4", "E4M3 scales", "global scale", "packed" ], "source": "mslk/quantize/triton/legacy/fp4_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py", "href": "#api-dequantize-mx4" }, { "name": "RoundingMode", "signature": "RoundingMode.nearest=0 | floor=1 | even=2 | stochastic=3 | ceil=4", "summary": "Rounding policy enum used by MXFP4 quantizers.", "notes": [ "Stochastic mode is rejected on ROCm." ], "kind": "enum", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py", "href": "#api-rounding-mode" }, { "name": "get_mx4_exp_bias", "signature": "get_mx4_exp_bias(ebits) -> int", "summary": "Returns the MX exponent bias for a format width.", "kind": "function", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/legacy/primitives.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py", "href": "#api-get-mx4-exp-bias" }, { "name": "FP4_E2M1_MAX", "signature": "FP4_E2M1_MAX = 6.0", "summary": "Largest finite E2M1 magnitude.", "kind": "constant", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py", "href": "#api-fp4-e2-m1-max" }, { "name": "FP8_E4M3_MAX", "signature": "FP8_E4M3_MAX = 448", "summary": "Largest finite OCP E4M3 magnitude.", "kind": "constant", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py", "href": "#api-fp8-e4-m3-max" }, { "name": "FP4_EBITS", "signature": "FP4_EBITS = 2", "summary": "E2M1 exponent-bit count.", "kind": "constant", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py", "href": "#api-fp4-ebits" }, { "name": "FP4_MBITS", "signature": "FP4_MBITS = 1", "summary": "E2M1 mantissa-bit count.", "kind": "constant", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_quantize.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py", "href": "#api-fp4-mbits" }, { "name": "E8M0_EXPONENT_BIAS", "signature": "E8M0_EXPONENT_BIAS = tl.constexpr(127)", "summary": "Shared exponent bias used to encode MX block scales.", "notes": [ "Host/Python consumers can read the wrapped scalar through E8M0_EXPONENT_BIAS.value." ], "kind": "constant", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_primitives", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_primitives/constants.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/constants.py", "href": "#api-e8-m0-exponent-bias" }, { "name": "BF16_MIN_NORMAL", "signature": "BF16_MIN_NORMAL = tl.constexpr(2**-126)", "summary": "Smallest normal BF16 magnitude used by safe scale math.", "notes": [ "Host/Python consumers can read the wrapped scalar through BF16_MIN_NORMAL.value." ], "kind": "constant", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_primitives", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_primitives/constants.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/constants.py", "href": "#api-bf16-min-normal" }, { "name": "blocked_scale_offset", "signature": "blocked_scale_offset(logical_row, logical_col, n_col_blocks, num_cols)", "summary": "Maps a logical scale coordinate into Blackwell’s 128×4 blocked layout.", "kind": "function", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_primitives", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_primitives/layout.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/layout.py", "href": "#api-blocked-scale-offset" }, { "name": "stacked_segment_map", "signature": "stacked_segment_map(m_sizes_ptr, pid_m, M_PER_BLOCK, NUM_SEGMENTS, PREFIX_NUM, BSEARCH_ITERS)", "summary": "Triton helper mapping a row tile to a stacked segment/expert.", "kind": "function", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_primitives", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_primitives/layout.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/layout.py", "href": "#api-stacked-segment-map" }, { "name": "mx4_scale_normalize_encode", "signature": "mx4_scale_normalize_encode(x_blocks, block_amax, pid_m, pid_n, seed, N, M_PER_BLOCK, NUM_GROUPS, GROUP_SIZE, ROUNDING_MODE, EBITS, MBITS, STOCHASTIC)", "summary": "Triton primitive that normalizes an MX block and encodes its E8M0 scale.", "kind": "function", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_primitives", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_primitives/scale.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/scale.py", "href": "#api-mx4-scale-normalize-encode" }, { "name": "convert_fp32_to_fp4_packed", "signature": "convert_fp32_to_fp4_packed(x_pairs, IS_GFX950, IS_ROCM)", "summary": "Triton primitive converting pairs of FP32 values into packed E2M1 nibbles.", "kind": "function", "stability": "Advanced", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_primitives", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/fp4_primitives/packing.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/packing.py", "href": "#api-convert-fp32-to-fp4-packed" }, { "name": "unsigned_fp32_to_e8m0", "signature": "unsigned_fp32_to_e8m0(tensor, mbits, scale_round_mode)", "summary": "Compatibility Triton primitive encoding positive FP32 magnitudes as E8M0 scales.", "kind": "function", "stability": "Low-level", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/legacy/primitives.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py", "href": "#api-unsigned-fp32-to-e8m0" }, { "name": "nvfp4_scale_swizzle", "signature": "nvfp4_scale_swizzle(offs_m)", "summary": "Compatibility Triton helper mapping rows into the NVFP4 scale-swizzle layout.", "kind": "function", "stability": "Low-level", "category": "Quantization", "group": "FP4 constants & primitives", "module": "mslk.quantize.triton.fp4_quantize", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "FP4", "Triton primitive", "rounding", "constants" ], "source": "mslk/quantize/triton/legacy/primitives.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py", "href": "#api-nvfp4-scale-swizzle" }, { "name": "pack_int4", "signature": "pack_int4(x: Tensor) -> Tensor", "summary": "Packs two logical 4-bit integer values into each int8 byte.", "returns": "Packed int8 with last dimension halved.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "INT4 preprocessing", "module": "mslk.quantize.shuffle", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "INT4", "packing", "preshuffle", "zero point" ], "source": "mslk/quantize/shuffle.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py", "href": "#api-pack-int4" }, { "name": "int4_row_quantize_zp", "signature": "int4_row_quantize_zp(x, group_size=128) -> tuple[Tensor, Tensor, Tensor]", "summary": "Groupwise asymmetric INT4 quantization with zero points.", "returns": "(unpacked int8 values [N,K], scales, zero_points).", "notes": [ "Values are not packed; call pack_int4 before a packed-weight GEMM." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "INT4 preprocessing", "module": "mslk.quantize.shuffle", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "INT4", "packing", "preshuffle", "zero point" ], "source": "mslk/quantize/shuffle.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py", "href": "#api-int4-row-quantize-zp" }, { "name": "int4_row_quantize", "signature": "int4_row_quantize(x, group_size=128) -> tuple[Tensor, Tensor]", "summary": "Groupwise symmetric INT4 quantization.", "returns": "(unpacked int8 values [N,K], scales).", "notes": [ "The source docstring incorrectly says [N,K/2]; each logical INT4 value occupies one int8 element until pack_int4 is called." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "INT4 preprocessing", "module": "mslk.quantize.shuffle", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "INT4", "packing", "preshuffle", "zero point" ], "source": "mslk/quantize/shuffle.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py", "href": "#api-int4-row-quantize" }, { "name": "quantize_int4_preshuffle", "signature": "quantize_int4_preshuffle(w, group_size=128, dtype='fp8', use_zp=True) -> tuple[Tensor, tuple[Tensor, Tensor]]", "summary": "Quantizes, packs, and CUDA-preshuffles INT4 weights and scale terms.", "returns": "(preshuffled packed weights, (scale_a, scale_b)).", "kind": "function", "stability": "Public", "category": "Quantization", "group": "INT4 preprocessing", "module": "mslk.quantize.shuffle", "platforms": [ "NVIDIA" ], "tags": [ "INT4", "packing", "preshuffle", "zero point" ], "source": "mslk/quantize/shuffle.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py", "href": "#api-quantize-int4-preshuffle" }, { "name": "ck_preshuffle", "signature": "ck_preshuffle(src, NXdl=16) -> Tensor", "summary": "Reorders data into the AMD Composable Kernel XDL weight layout.", "notes": [ "Source must use float8_e4m3fnuz." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "INT4 preprocessing", "module": "mslk.quantize.shuffle", "platforms": [ "AMD ROCm" ], "tags": [ "INT4", "packing", "preshuffle", "zero point" ], "source": "mslk/quantize/shuffle.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py", "href": "#api-ck-preshuffle" }, { "name": "quantize_bf16_to_mx6_e2m3", "signature": "quantize_bf16_to_mx6_e2m3(x, block_size=32) -> tuple[Tensor, Tensor]", "summary": "Quantizes BF16 to unpacked six-bit E2M3 codes plus E8M0 block scales.", "returns": "(unpacked uint8 codes, E8M0 scales).", "notes": [ "The codes are not GEMM-ready until pack_fp6_e2m3 is applied." ], "kind": "function", "stability": "Public", "category": "Quantization", "group": "MXFP6 helpers", "module": "mslk.quantize.mx_mixed_dtype_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "MXFP6", "E2M3", "E8M0", "packing" ], "source": "mslk/quantize/mx_mixed_dtype_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py", "href": "#api-quantize-bf16-to-mx6-e2m3" }, { "name": "pack_fp6_e2m3", "signature": "pack_fp6_e2m3(unpacked: Tensor) -> Tensor", "summary": "Bit-packs four E2M3 values into three uint8 bytes.", "kind": "function", "stability": "Public", "category": "Quantization", "group": "MXFP6 helpers", "module": "mslk.quantize.mx_mixed_dtype_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "MXFP6", "E2M3", "E8M0", "packing" ], "source": "mslk/quantize/mx_mixed_dtype_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py", "href": "#api-pack-fp6-e2m3" }, { "name": "E2M3_DECODE", "signature": "E2M3_DECODE: tuple[float, ...]", "summary": "Lookup table mapping 6-bit E2M3 codes to decoded values.", "kind": "constant", "stability": "Public", "category": "Quantization", "group": "MXFP6 helpers", "module": "mslk.quantize.mx_mixed_dtype_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "MXFP6", "E2M3", "E8M0", "packing" ], "source": "mslk/quantize/mx_mixed_dtype_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py", "href": "#api-e2-m3-decode" }, { "name": "E2M3_MAX", "signature": "E2M3_MAX = 7.5", "summary": "Largest finite E2M3 magnitude.", "kind": "constant", "stability": "Public", "category": "Quantization", "group": "MXFP6 helpers", "module": "mslk.quantize.mx_mixed_dtype_utils", "platforms": [ "NVIDIA SM100+" ], "tags": [ "MXFP6", "E2M3", "E8M0", "packing" ], "source": "mslk/quantize/mx_mixed_dtype_utils.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py", "href": "#api-e2-m3-max" }, { "name": "index_shuffling", "signature": "index_shuffling(routing_scores, expert_index_start=None, expert_index_end=None, valid_token_count=None, top_k=1) -> tuple[Tensor, Tensor, Tensor]", "summary": "Selects top-k experts and groups token indices by expert.", "details": "routing_scores is [T,E] BF16/FP32. Optional half-open expert bounds select a local expert range and returned expert IDs are localized by subtracting expert_index_start.", "returns": "int32 token_counts [E+2], expert_indices [T×top_k], token_indices [T×top_k].", "notes": [ "token_counts[:E] holds per-expert selected counts; token_counts[-2] is the total input-token count; token_counts[-1] is the selected-token count and bounds the valid index-array prefix.", "E must be 16, 32, 128, or 320.", "CUDA top_k is 1, 2, or 4; ROCm supports top_k=1.", "The alias is None when torch.cuda.is_available() is false." ], "kind": "torch op alias", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "csrc/moe/moe_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/moe/moe_ops.cpp", "href": "#api-index-shuffling" }, { "name": "gather_scale_dense_tokens", "signature": "gather_scale_dense_tokens(x, token_indices, expert_indices, scores, valid_token_count=None) -> Tensor", "summary": "Gathers routed tokens and multiplies each by its router score.", "details": "Output row i is x[token_indices[i]] × scores[token_indices[i], expert_indices[i]]. x is [T,D] and the result is [selected,D].", "notes": [ "x and index arrays must be contiguous; scores may be strided.", "D must be divisible by 1024 for large-T occupancy, otherwise 512.", "valid_token_count gates the written prefix; remaining rows are uninitialized." ], "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-gather-scale-dense-tokens" }, { "name": "gather_scale_quant_dense_tokens", "signature": "gather_scale_quant_dense_tokens(x, token_indices, expert_indices, scores, scale_ub=None, valid_token_count=None) -> tuple[Tensor, Tensor]", "summary": "Gathers/scales routed tokens and rowwise-quantizes them to FP8.", "returns": "(FP8 [selected,D], float32 reciprocal scales [selected]).", "notes": [ "Prefer this direct Python wrapper: the current torch op schema declares only one Tensor even though the implementation returns two.", "Rows outside the valid output prefix are left uninitialized." ], "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-gather-scale-quant-dense-tokens" }, { "name": "scatter_add_dense_tokens", "signature": "scatter_add_dense_tokens(out_tokens, in_tokens, token_indices, valid_token_count=None) -> None", "summary": "Atomically accumulates routed expert rows back to dense token order.", "details": "Mutates out_tokens [T,D] by adding in_tokens[i] into row token_indices[i].", "notes": [ "Contiguous tensors required.", "NVIDIA requires CUDA toolkit 12.4+; ROCm is supported." ], "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-scatter-add-dense-tokens" }, { "name": "scatter_add_padded_tokens", "signature": "scatter_add_padded_tokens(in_tokens, token_counts, token_indices, out_tokens) -> None", "summary": "Scatters expert/rank-padded token blocks back into dense token storage.", "details": "Consumes in_tokens [EP,T_K,D], per-expert token counts, and token indices, then mutates out_tokens [T,D].", "notes": [ "T_K must be divisible by T and E by EP.", "NVIDIA requires CUDA toolkit 12.4+; ROCm is supported." ], "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-scatter-add-padded-tokens" }, { "name": "combine_shuffling", "signature": "combine_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False) -> tuple[Tensor, Tensor]", "summary": "Reorders rank-major expert-token blocks to expert-major/rank-minor.", "returns": "A same-shape reordered buffer and selected-expert counts [EG+1], whose final element is total valid rows.", "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/shuffling.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py", "href": "#api-combine-shuffling" }, { "name": "split_shuffling", "signature": "split_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False, init_with_zeros=False) -> Tensor", "summary": "Inverse of combine_shuffling for expert communication results.", "notes": [ "init_with_zeros selects zero-filled versus uninitialized allocation where holes are possible.", "The current registered Meta/CUDA implementation functions omit this sixth argument; direct wrapper behavior is the intended contract." ], "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/shuffling.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py", "href": "#api-split-shuffling" }, { "name": "silu_mul", "signature": "silu_mul(x0, x1, valid_token_count=None) -> Tensor", "summary": "Fused SwiGLU activation x0·sigmoid(x0)·x1.", "returns": "Same shape and dtype as [T,D] inputs.", "notes": [ "Inputs must match and be row-contiguous.", "D divisibility is 1024 for T at least the SM count, otherwise 512.", "Rows beyond valid_token_count are left uninitialized." ], "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/activation.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py", "href": "#api-silu-mul" }, { "name": "silu_mul_quant", "signature": "silu_mul_quant(x0, x1, scale_ub=None, valid_token_count=None) -> tuple[Tensor, Tensor]", "summary": "Fused SwiGLU plus rowwise FP8 quantization.", "returns": "(FP8 [T,D], float32 reciprocal scales [T]).", "notes": [ "Rows beyond valid_token_count are uninitialized.", "The torch schema returns two tensors, but the current Meta implementation returns only the FP8 tensor; verify compile/export flows." ], "kind": "function", "stability": "Public", "category": "MoE", "group": "Routing & token movement", "module": "mslk.moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "routing", "tokens", "experts", "Triton" ], "source": "mslk/moe/activation.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py", "href": "#api-silu-mul-quant" }, { "name": "torch.ops.mslk.index_shuffling", "signature": "torch.ops.mslk.index_shuffling(routing_scores, expert_index_start=None, expert_index_end=None, valid_token_count=None, top_k=1) -> (Tensor, Tensor, Tensor)", "summary": "Native top-k expert grouping op.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "csrc/moe/moe_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/moe/moe_ops.cpp", "href": "#api-mslk-index-shuffling" }, { "name": "torch.ops.mslk.scatter_add_along_first_dim", "signature": "torch.ops.mslk.scatter_add_along_first_dim(Dst, Src, Index) -> None", "summary": "In-place first-dimension scatter-add with a fast BF16 TMA path.", "details": "Equivalent to Dst.scatter_add_(0, Index[:,None].expand(-1,K), Src). The fast contiguous BF16 path accepts int32/int64 index and K divisible by 256; other cases fall back to ATen.", "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "csrc/moe/moe_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/moe/moe_ops.cpp", "href": "#api-mslk-scatter-add-along-first-dim" }, { "name": "torch.ops.mslk.silu_mul", "signature": "torch.ops.mslk.silu_mul(x0, x1, valid_token_count=None) -> Tensor", "summary": "Registered fused SwiGLU activation.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/activation.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py", "href": "#api-mslk-silu-mul" }, { "name": "torch.ops.mslk.silu_mul_quant", "signature": "torch.ops.mslk.silu_mul_quant(x0, x1, scale_ub=None, valid_token_count=None) -> (Tensor, Tensor)", "summary": "Registered fused SwiGLU plus FP8 quantization.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "notes": [ "Current fake/Meta implementation returns only one tensor despite this two-output schema." ], "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/activation.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py", "href": "#api-mslk-silu-mul-quant" }, { "name": "torch.ops.mslk.gather_scale_dense_tokens", "signature": "torch.ops.mslk.gather_scale_dense_tokens(x, token_indices, expert_indices, scores, valid_token_count=None) -> Tensor", "summary": "Registered routed-token gather/scale op.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-mslk-gather-scale-dense-tokens" }, { "name": "torch.ops.mslk.gather_scale_quant_dense_tokens", "signature": "torch.ops.mslk.gather_scale_quant_dense_tokens(x, token_indices, expert_indices, scores, scale_ub=None, valid_token_count=None) -> Tensor", "summary": "Registered quantizing routed-token gather with a known return-schema mismatch.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "notes": [ "Schema says Tensor; CUDA and Meta implementations return (FP8, reciprocal_scale). Prefer the direct Python wrapper and validate your installed build." ], "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-mslk-gather-scale-quant-dense-tokens" }, { "name": "torch.ops.mslk.scatter_add_dense_tokens", "signature": "torch.ops.mslk.scatter_add_dense_tokens(out_tokens, in_tokens, token_indices, valid_token_count=None) -> None", "summary": "Registered dense-token scatter-add op.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-mslk-scatter-add-dense-tokens" }, { "name": "torch.ops.mslk.scatter_add_padded_tokens", "signature": "torch.ops.mslk.scatter_add_padded_tokens(in_tokens, token_counts, token_indices, out_tokens) -> None", "summary": "Registered padded expert-token scatter op.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/gather_scatter.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py", "href": "#api-mslk-scatter-add-padded-tokens" }, { "name": "torch.ops.mslk.combine_shuffling", "signature": "torch.ops.mslk.combine_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False) -> (Tensor, Tensor)", "summary": "Registered rank-major to expert-major reorder.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/shuffling.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py", "href": "#api-mslk-combine-shuffling" }, { "name": "torch.ops.mslk.split_shuffling", "signature": "torch.ops.mslk.split_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False, init_with_zeros=False) -> Tensor", "summary": "Registered inverse expert-shuffling reorder.", "details": "Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.", "notes": [ "Current registered implementation signatures omit init_with_zeros." ], "kind": "torch op", "stability": "Low-level", "category": "MoE", "group": "Registered MoE ops", "module": "torch.ops.mslk", "platforms": [ "NVIDIA", "AMD", "Meta / fake" ], "tags": [ "dispatcher", "custom op", "MoE" ], "source": "mslk/moe/shuffling.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py", "href": "#api-mslk-split-shuffling" }, { "name": "MoEArgs", "signature": "MoEArgs(precision: str, dim: int, hidden_dim: int, num_experts: int, top_k: int, mp_size: int, ep_size: int, mp_size_for_routed_experts: Optional[int], use_fast_accum: bool, dedup_comm: bool)", "summary": "Frozen configuration dataclass for the provided MoE modules.", "notes": [ "No constructor defaults.", "precision is not interpreted directly; initialization methods typically select BF16 versus FP8 rowwise storage." ], "methods": [ { "signature": "num_local_experts -> int", "summary": "Cached num_experts // ep_size property." } ], "kind": "class", "stability": "Public", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-mo-eargs" }, { "name": "BaselineMoE", "signature": "BaselineMoE(ep_group: ProcessGroup, ep_mp_group: ProcessGroup, moe_args: MoEArgs)", "summary": "Dense/reference distributed SwiGLU MoE implementation.", "details": "Uses sigmoid routing plus top-k, one shared SwiGLU expert path, routed experts, and expert all-to-all when EP>1. Forward runs under no_grad.", "returns": "forward(x [B,T,D], use_static_shape) -> [B,T,D].", "notes": [ "Caller performs model-parallel all-reduce externally.", "Imports require fairscale and pyre_extensions." ], "methods": [ { "signature": "build(init_methods=None) -> BaselineMoE", "summary": "Initializes router/shared/routed parameters; required before forward." }, { "signature": "forward(x, use_static_shape)", "summary": "Computes shared and routed experts; Baseline accepts but does not use use_static_shape." }, { "signature": "router_DE / E / EG / K", "summary": "Router and geometry properties." }, { "signature": "is_shared_fp8_rowwise / is_routed_fp8_rowwise", "summary": "Reports scaled FP8 parameter storage." } ], "kind": "class", "stability": "Public", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-baseline-mo-e" }, { "name": "MetaShufflingMoE", "signature": "MetaShufflingMoE(ep_group: ProcessGroup, ep_mp_group: ProcessGroup, moe_args: MoEArgs)", "summary": "Optimized inference MoE using shuffle/gather/grouped-GEMM primitives and overlapped communication.", "returns": "forward(x [B,T,D], use_static_shape) -> [B,T,D].", "notes": [ "top_k must equal 1.", "mp_size must equal the effective routed-expert MP size.", "Allocates CUDA streams/events in the constructor.", "use_static_shape selects static overlap/all-gather paths versus dynamic all-to-all behavior." ], "kind": "class", "stability": "Public", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-meta-shuffling-mo-e" }, { "name": "ScaledParameter", "signature": "ScaledParameter(data: Tensor, scale: Optional[Tensor]=None)", "summary": "Non-trainable Parameter carrying optional quantization scales.", "methods": [ { "signature": "weights -> Tensor", "summary": "Underlying parameter data." }, { "signature": "scales -> Tensor", "summary": "Scale tensor; asserts that one exists." }, { "signature": "is_scaled -> bool", "summary": "Whether scales are attached." } ], "kind": "class", "stability": "Advanced", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-scaled-parameter" }, { "name": "Experts", "signature": "Experts(dim: int, hidden_dim: int)", "summary": "Abstract shared base for expert weight containers.", "notes": [ "Model-parallel world size must divide dim and hidden_dim." ], "methods": [ { "signature": "build(init_methods=None) -> Experts", "summary": "Initializes and packs expert parameters." }, { "signature": "w13 / w2 -> ScaledParameter", "summary": "Packed SwiGLU input/gate and output weights." }, { "signature": "is_fp8_rowwise -> bool", "summary": "Detects FP8 rowwise storage from w13 dtype." } ], "kind": "class", "stability": "Advanced", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-experts" }, { "name": "RoutedExperts", "signature": "RoutedExperts(num_local_experts: int, dim: int, hidden_dim: int)", "summary": "Per-local-expert stacked weight container.", "kind": "class", "stability": "Advanced", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-routed-experts" }, { "name": "SharedExperts", "signature": "SharedExperts(dim: int, hidden_dim: int)", "summary": "Shared SwiGLU expert weight container.", "kind": "class", "stability": "Advanced", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-shared-experts" }, { "name": "init_params", "signature": "init_params(key: str, param: ScaledParameter, init_methods: Mapping) -> None", "summary": "Applies a named initialization/quantization callback or Kaiming-uniform fallback.", "notes": [ "Recognized keys cover routed w_in/w_out/w_swiglu, shared equivalents, and router_DE." ], "kind": "function", "stability": "Advanced", "category": "MoE", "group": "Composed MoE layers", "module": "mslk.moe.layers", "platforms": [ "NVIDIA", "AMD", "Distributed" ], "tags": [ "PyTorch module", "expert parallel", "model parallel", "inference" ], "source": "mslk/moe/layers.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py", "href": "#api-init-params" }, { "name": "torch.ops.mslk.f8f8bf16_conv", "signature": "torch.ops.mslk.f8f8bf16_conv(activation, filter, scale, padding: int[3], stride: int[3], dilation: int[3]) -> Tensor", "summary": "FP8 3D cross-correlation/convolution with combined scale and BF16 output.", "details": "Accepts rank-5 activation as NDHWC [N,D,H,W,C] or logical NCDHW backed by channels_last_3d; filter as KTRSC [K,T,R,S,C] or logical KCTRS backed by channels_last_3d. Spatial lists use D,H,W order.", "returns": "BF16 [N,Z,P,Q,K] for NDHWC, or [N,K,Z,P,Q] in channels_last_3d layout for NCDHW.", "notes": [ "Import mslk.conv to register the op and fake implementation; there is no Python convenience wrapper.", "Activation/filter channels must match and their channel dimension must be contiguous.", "Requires the SM100 Blackwell CUTLASS build; the schema is excluded on ROCm.", "padding, stride, and dilation must each have length 3." ], "kind": "torch op", "stability": "Public", "category": "Convolution", "group": "FP8 3D convolution", "module": "torch.ops.mslk", "platforms": [ "NVIDIA SM100+", "Meta / fake" ], "tags": [ "FP8", "BF16", "3D convolution", "NDHWC", "channels_last_3d", "CUTLASS" ], "source": "csrc/conv/conv_ops.cpp", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/conv/conv_ops.cpp", "href": "#api-mslk-f8f8bf16-conv" }, { "name": "is_flydsl_available", "signature": "is_flydsl_available() -> bool", "summary": "Cached check for FlyDSL importability and current-architecture support.", "details": "Returns true only if the package is discoverable and the current ROCm arch appears in FlyDSL’s shared-memory capacity map. Catches errors and returns false.", "notes": [ "mslk.flydsl.__init__ is empty; import the explicit common, jit, or aot submodule." ], "methods": [ { "signature": "is_flydsl_available.cache_clear()", "summary": "Invalidates the lru_cache after environment changes." } ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "FlyDSL runtime & AOT", "module": "mslk.flydsl.common", "platforms": [ "AMD ROCm", "CPU / Python" ], "tags": [ "FlyDSL", "JIT", "AOT", "cache" ], "source": "mslk/flydsl/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/common.py", "href": "#api-is-flydsl-available" }, { "name": "require_flydsl", "signature": "require_flydsl() -> None", "summary": "Raises RuntimeError with an installation hint when FlyDSL is unavailable.", "notes": [ "mslk.flydsl.__init__ is empty; import the explicit common, jit, or aot submodule." ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "FlyDSL runtime & AOT", "module": "mslk.flydsl.common", "platforms": [ "AMD ROCm", "CPU / Python" ], "tags": [ "FlyDSL", "JIT", "AOT", "cache" ], "source": "mslk/flydsl/common.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/common.py", "href": "#api-require-flydsl" }, { "name": "configure_runtime_cache", "signature": "configure_runtime_cache() -> None", "summary": "Points FlyDSL at bundled AOT artifacts and applies MSLK’s JIT-disable switch.", "details": "If bundled aot_artifacts exists and FLYDSL_RUNTIME_CACHE_DIR is unset, it selects that directory. MSLK_FLYDSL_DISABLE_JIT=1 maps to FLYDSL_RUNTIME_RUN_ONLY=1.", "notes": [ "Called automatically when mslk.flydsl.jit is imported." ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "FlyDSL runtime & AOT", "module": "mslk.flydsl.jit", "platforms": [ "AMD ROCm", "CPU / Python" ], "tags": [ "FlyDSL", "JIT", "AOT", "cache" ], "source": "mslk/flydsl/jit.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/jit.py", "href": "#api-configure-runtime-cache" }, { "name": "run_compiled", "signature": "run_compiled(launcher: Callable[..., Any], *args: Any) -> None", "summary": "Compiles a FlyDSL launcher once, then invokes the cached compiled function.", "details": "The first compiler call executes the launch itself; only subsequent calls invoke launcher._mslk_cf separately.", "notes": [ "mslk.flydsl.__init__ is empty; import the explicit common, jit, or aot submodule." ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "FlyDSL runtime & AOT", "module": "mslk.flydsl.jit", "platforms": [ "AMD ROCm", "CPU / Python" ], "tags": [ "FlyDSL", "JIT", "AOT", "cache" ], "source": "mslk/flydsl/jit.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/jit.py", "href": "#api-run-compiled" }, { "name": "collect_aot_jobs", "signature": "collect_aot_jobs() -> list[tuple[str, dict[str, Any], str]]", "summary": "Builds the Cartesian product of registered kernel configs and architectures.", "notes": [ "The shipped _AOT_KERNEL_MODULES list is currently empty, so the default result contains no jobs." ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "FlyDSL runtime & AOT", "module": "mslk.flydsl.aot", "platforms": [ "AMD ROCm", "CPU / Python" ], "tags": [ "FlyDSL", "JIT", "AOT", "cache" ], "source": "mslk/flydsl/aot.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/aot.py", "href": "#api-collect-aot-jobs" }, { "name": "compile_aot", "signature": "compile_aot(cache_dir: str) -> None", "summary": "Multiprocess-compiles all collected FlyDSL AOT jobs into a runtime cache.", "details": "Uses spawn workers under COMPILE_ONLY=1, aggregates failures, and restores environment state.", "notes": [ "With the current empty registry, prints that there are no kernels and skips.", "Worker count uses MSLK_FLYDSL_AOT_WORKERS or a bounded CPU-affinity default." ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "FlyDSL runtime & AOT", "module": "mslk.flydsl.aot", "platforms": [ "AMD ROCm", "CPU / Python" ], "tags": [ "FlyDSL", "JIT", "AOT", "cache" ], "source": "mslk/flydsl/aot.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/aot.py", "href": "#api-compile-aot" }, { "name": "is_cuda", "signature": "is_cuda() -> bool", "summary": "True for an NVIDIA CUDA build with an available device.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-is-cuda" }, { "name": "is_rocm", "signature": "is_rocm() -> bool", "summary": "True for a PyTorch HIP build with an available AMD device.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-is-rocm" }, { "name": "compute_capability_in", "signature": "compute_capability_in(major_min: int, major_max: int | None=None) -> bool", "summary": "Checks current device compute-capability major against an inclusive range.", "notes": [ "Returns false without a device; implementation also reads the tuple on ROCm." ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-compute-capability-in" }, { "name": "compute_capability_at_least", "signature": "compute_capability_at_least(major_min: int, minor_min: int=0) -> bool", "summary": "Checks the current capability tuple against a minimum.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-compute-capability-at-least" }, { "name": "cuda_version_at_least", "signature": "cuda_version_at_least(major_min: int) -> bool", "summary": "Checks the PyTorch build’s CUDA toolkit major, not driver or device capability.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-cuda-version-at-least" }, { "name": "get_gfx_arch_name", "signature": "get_gfx_arch_name() -> str", "summary": "Returns current ROCm gcnArchName, or an empty string on unavailable/error.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-get-gfx-arch-name" }, { "name": "gfx_arch_in", "signature": "gfx_arch_in(arch_list: Iterable[str]) -> bool", "summary": "Substring-matches any requested gfx name against gcnArchName.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-gfx-arch-in" }, { "name": "is_gfx942", "signature": "is_gfx942() -> bool", "summary": "True on AMD MI300X/CDNA3 gfx942.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-is-gfx942" }, { "name": "is_gfx950", "signature": "is_gfx950() -> bool", "summary": "True on AMD MI350/CDNA4 gfx950.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-is-gfx950" }, { "name": "supports_float8_fnuz", "signature": "supports_float8_fnuz(throw_on_hip_incompatibility=True) -> bool", "summary": "Reports whether the active ROCm target uses the FNUZ FP8 flavor.", "details": "Returns true for gfx942. gfx950 uses OCP FP8 and returns false unless MSLK_ROCM_FORCE_FP8FNUZ_TYPE is set before import.", "notes": [ "The current function docstring claims gfx950 FNUZ support, but implementation and architecture comments say OCP FP8; this reference follows implementation." ], "kind": "function", "stability": "Public", "category": "Runtime", "group": "Device detection", "module": "mslk.utils.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "device", "architecture", "capability", "FP8" ], "source": "mslk/utils/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py", "href": "#api-supports-float8-fnuz" }, { "name": "skipUnlessCuda", "signature": "skipUnlessCuda() -> Callable", "summary": "Skips unless running on a strict NVIDIA CUDA device.", "returns": "A unittest-style skip decorator.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Test decorators", "module": "mslk.testing.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "unittest", "pytest", "skip decorator" ], "source": "mslk/testing/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py", "href": "#api-skip-unless-cuda" }, { "name": "skipUnlessRocm", "signature": "skipUnlessRocm() -> Callable", "summary": "Skips unless running on a strict AMD ROCm device.", "returns": "A unittest-style skip decorator.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Test decorators", "module": "mslk.testing.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "unittest", "pytest", "skip decorator" ], "source": "mslk/testing/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py", "href": "#api-skip-unless-rocm" }, { "name": "skipUnlessCudaCapability", "signature": "skipUnlessCudaCapability(major_min, major_max=None, *, minor_min=0) -> Callable", "summary": "Narrows CUDA tests to a capability range and remains transparent on ROCm.", "returns": "A unittest-style skip decorator.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Test decorators", "module": "mslk.testing.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "unittest", "pytest", "skip decorator" ], "source": "mslk/testing/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py", "href": "#api-skip-unless-cuda-capability" }, { "name": "skipUnlessCudaVersion", "signature": "skipUnlessCudaVersion(major_min: int) -> Callable", "summary": "Narrows CUDA tests to a minimum toolkit major and remains transparent on ROCm.", "returns": "A unittest-style skip decorator.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Test decorators", "module": "mslk.testing.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "unittest", "pytest", "skip decorator" ], "source": "mslk/testing/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py", "href": "#api-skip-unless-cuda-version" }, { "name": "skipUnlessGfxArch", "signature": "skipUnlessGfxArch(arch: str, *more_archs: str) -> Callable", "summary": "Narrows ROCm tests to selected gfx targets and remains transparent on CUDA.", "returns": "A unittest-style skip decorator.", "kind": "function", "stability": "Public", "category": "Runtime", "group": "Test decorators", "module": "mslk.testing.device", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "unittest", "pytest", "skip decorator" ], "source": "mslk/testing/device.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py", "href": "#api-skip-unless-gfx-arch" }, { "name": "open_source", "signature": "mslk.open_source = True", "summary": "Signals the OSS package variant to tests and integration code.", "kind": "constant", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "mslk/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py", "href": "#api-open-source" }, { "name": "__version__ / __target__ / __variant__", "signature": "mslk.__version__: str; mslk.__target__: str; mslk.__variant__: str", "summary": "Generated package-build identity values, with internal/default fallbacks.", "kind": "constants", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "mslk/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py", "href": "#api-version-target-variant" }, { "name": "load_library_buck", "signature": "load_library_buck(buck_target: str) -> None", "summary": "Internal/OSS bridge used by domain imports to load split Buck native libraries.", "notes": [ "No-ops in Python-only mode and suppresses OSError in OSS." ], "kind": "function", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk.utils.torch.library", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "mslk/utils/torch/library.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/torch/library.py", "href": "#api-load-library-buck" }, { "name": "MSLK_PYTHON_ONLY", "signature": "MSLK_PYTHON_ONLY=1", "summary": "Skips mslk.so native loading/compilation for Python-only development.", "notes": [ "Does not provide CPU implementations of GPU kernels." ], "kind": "environment variable", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "mslk/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py", "href": "#api-mslk-python-only" }, { "name": "MSLK runtime environment variables", "signature": "MSLK_CACHE_DIR · MSLK_AUTOTUNE_USE_CUDA_GRAPH · MSLK_AUTOTUNE_COLLECT_STATS · MSLK_ROCM_FORCE_FP8FNUZ_TYPE", "summary": "Runtime cache, autotune, timing-stat, and ROCm FP8 override environment variables.", "details": "Cache defaults to $HOME/.mslk. Presence of the autotune variables enables graph benchmarking or timing-stat collection.", "kind": "configuration", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "mslk/__init__.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py", "href": "#api-mslk-runtime-environment-variables" }, { "name": "MSLK diagnostic compile-time defines", "signature": "MSLK_MEMCHECK · MSLK_ISOLATE_KERNEL_LAUNCH · MSLK_TENSORCHECK", "summary": "Preprocessor defines enabling accessor bounds checks, isolated launches, and tensor value checks.", "notes": [ "These are compile-time defines, not environment variables read by a running process." ], "kind": "configuration", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "native build", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "include/mslk/utils/kernel_launcher.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh", "href": "#api-mslk-diagnostic-compile-time-defines" }, { "name": "FlyDSL environment variables", "signature": "FLYDSL_RUNTIME_CACHE_DIR · MSLK_FLYDSL_DISABLE_JIT · FLYDSL_RUNTIME_RUN_ONLY · MSLK_FLYDSL_AOT_WORKERS · COMPILE_ONLY", "summary": "Runtime cache, run-only, worker-count, and compile-only controls for FlyDSL.", "kind": "configuration", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk.flydsl", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "mslk/flydsl/jit.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/jit.py", "href": "#api-fly-dsl-environment-variables" }, { "name": "setup.py build CLI", "signature": "python setup.py install [--verbose] [--debug {0,1,2}] [--dryrun] [--build-target default] [--build-variant {cpu,cuda,rocm}] [--package_channel {nightly,test,release}] [--nvml_lib_path PATH] [--nccl_lib_path PATH] [--build_fb_code] [--cxxprefix PATH]", "summary": "Source-build command-line controls.", "notes": [ "Debug level 1 enables device-side assertion defines; level 2 additionally builds unoptimized debug code.", "Although the parser accepts cpu, CMake rejects a native CPU build; use Python-only mode for no-native installs.", "Source builds are supported on Linux." ], "kind": "configuration", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "setup.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/setup.py", "href": "#api-setup-py-build-cli" }, { "name": "MSLK build environment variables", "signature": "MSLK_PYTHON_ONLY · MSLK_PYTHON_ONLY_PLAT · MSLK_PACKAGE_NAME · MSLK_VERSION_OVERRIDE · MSLK_BUILD_FB_CODE · CHANNEL · CU_VERSION · BUILD_ROCM_VERSION · PYTORCH_ROCM_ARCH · ROCM_PATH · CUDA_BIN_PATH · CUDACXX · CUB_DIR", "summary": "Package identity, Python-only, toolkit, architecture, and toolchain inputs used by setup.py.", "kind": "configuration", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "setup.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/setup.py", "href": "#api-mslk-build-environment-variables" }, { "name": "MSLK CMake knobs", "signature": "MSLK_BUILD_TARGET · MSLK_BUILD_VARIANT · BUILD_FB_CODE · MSLK_FBPKG_BUILD · TORCH_CUDA_ARCH_LIST · AMDGPU_TARGETS · HIP_ROOT_DIR · NVML_LIB_PATH · NCCL_INCLUDE_DIRS · NCCL_LIBRARIES", "summary": "Direct CMake configuration variables for native builds.", "notes": [ "Native code is built as C++20." ], "kind": "configuration", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "CMakeLists.txt", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/CMakeLists.txt", "href": "#api-mslk-cmake-knobs" }, { "name": "Package dependencies", "signature": "install_requires=['numpy']; extra: mslk[flash3]; ROCm wheels add pinned FlyDSL", "summary": "Declared runtime dependency boundary and optional kernel packages.", "notes": [ "PyTorch is intentionally omitted from install_requires.", "MoE layers additionally import fairscale and pyre_extensions; ensure they are present in environments using mslk.moe.layers." ], "kind": "configuration", "stability": "Advanced", "category": "Runtime", "group": "Package loading & runtime knobs", "module": "mslk", "platforms": [ "NVIDIA", "AMD", "CPU / Python" ], "tags": [ "import", "shared library", "environment", "build" ], "source": "setup.py", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/setup.py", "href": "#api-package-dependencies" }, { "name": "nextPowerOf2", "signature": "constexpr int64_t nextPowerOf2(int64_t num)", "summary": "Rounds a positive integer up to the next power of two.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Numeric & device utilities", "module": "namespace mslk", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "header", "constexpr", "device properties" ], "source": "include/mslk/utils/utils.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils.h", "href": "#api-next-power-of2" }, { "name": "roundUp", "signature": "constexpr int64_t roundUp(int64_t num, int64_t multiple)", "summary": "Rounds num up to a requested multiple.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Numeric & device utilities", "module": "namespace mslk", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "header", "constexpr", "device properties" ], "source": "include/mslk/utils/utils.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils.h", "href": "#api-round-up" }, { "name": "nextPowerOf2OrRoundUp", "signature": "constexpr int64_t nextPowerOf2OrRoundUp(int64_t num, int64_t roundUpTo, int64_t threshold)", "summary": "Uses power-of-two rounding below a threshold and fixed-multiple rounding above it.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Numeric & device utilities", "module": "namespace mslk", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "header", "constexpr", "device properties" ], "source": "include/mslk/utils/utils.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils.h", "href": "#api-next-power-of2-or-round-up" }, { "name": "getDeviceArch", "signature": "inline int getDeviceArch()", "summary": "Returns cached current-device compute-capability major.", "notes": [ "Enforces CUDA runtime 12.8+ when the major is 10 or newer." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Numeric & device utilities", "module": "namespace mslk", "platforms": [ "NVIDIA" ], "tags": [ "header", "constexpr", "device properties" ], "source": "include/mslk/utils/utils_gpu.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils_gpu.h", "href": "#api-get-device-arch" }, { "name": "getSMCount", "signature": "inline int64_t getSMCount(int device_index, std::optional num_sms)", "summary": "Returns an override or the available SM count after PyTorch carveout.", "notes": [ "Current cached implementation obtains device-0 properties and does not use device_index for that lookup." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Numeric & device utilities", "module": "namespace mslk", "platforms": [ "NVIDIA" ], "tags": [ "header", "constexpr", "device properties" ], "source": "include/mslk/utils/utils_gpu.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils_gpu.h", "href": "#api-get-smcount" }, { "name": "get_device_for_stream", "signature": "get_device_for_stream(const cudaStream_t& stream) -> device index", "summary": "Resolves the CUDA device associated with a raw stream.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "CUDA stream & shared memory", "module": "mslk::utils::device", "platforms": [ "NVIDIA" ], "tags": [ "CUDA stream", "device", "dynamic shared memory" ], "source": "include/mslk/utils/device/cuda_utilities.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/device/cuda_utilities.cuh", "href": "#api-get-device-for-stream" }, { "name": "to_cuda_stream", "signature": "to_cuda_stream(stream, device_index=-1) -> c10::cuda::CUDAStream", "summary": "Normalizes raw cudaStream_t or c10 stream input to CUDAStream.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "CUDA stream & shared memory", "module": "mslk::utils::device", "platforms": [ "NVIDIA" ], "tags": [ "CUDA stream", "device", "dynamic shared memory" ], "source": "include/mslk/utils/device/cuda_utilities.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/device/cuda_utilities.cuh", "href": "#api-to-cuda-stream" }, { "name": "set_gpu_max_dynamic_shared_memory", "signature": "set_gpu_max_dynamic_shared_memory(kernel, smem_bytes, device=current) -> void", "summary": "Opts a kernel into the requested dynamic shared-memory limit after availability checks.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "CUDA stream & shared memory", "module": "mslk::utils::device", "platforms": [ "NVIDIA" ], "tags": [ "CUDA stream", "device", "dynamic shared memory" ], "source": "include/mslk/utils/device/cuda_utilities.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/device/cuda_utilities.cuh", "href": "#api-set-gpu-max-dynamic-shared-memory" }, { "name": "source_location", "signature": "using source_location = std::source_location // or backport selected by build", "summary": "Portable source-location type used by SourceContext.", "kind": "C++ alias", "stability": "Developer", "category": "C++", "group": "Source context & timing", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "diagnostics", "source_location", "events", "benchmark" ], "source": "include/mslk/utils/source_context.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/source_context.h", "href": "#api-source-location" }, { "name": "SourceContext", "signature": "constexpr mslk::utils::SourceContext(const source_location& location, const std::string_view& summary, const std::string_view& template_filepath, const std::string_view& dsa_file_descriptor) noexcept", "summary": "Carries source location and a human-readable label through kernel diagnostics.", "methods": [ { "signature": "description() const -> std::string", "summary": "Formats the contextual description." }, { "signature": "withSummary(summary) -> SourceContext", "summary": "Copies context with a new summary." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Source context & timing", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "diagnostics", "source_location", "events", "benchmark" ], "source": "include/mslk/utils/source_context.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/source_context.h", "href": "#api-source-context" }, { "name": "SOURCE_CONTEXT_CURRENT", "signature": "SOURCE_CONTEXT_CURRENT(label)", "summary": "Captures the current file/function/line into a SourceContext.", "kind": "C++ macro", "stability": "Developer", "category": "C++", "group": "Source context & timing", "module": "global macro", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "diagnostics", "source_location", "events", "benchmark" ], "source": "include/mslk/utils/source_context.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/source_context.h", "href": "#api-source-context-current" }, { "name": "KernelExecutionTimer", "signature": "KernelExecutionTimer(c10::cuda::CUDAStream stream)", "summary": "Single-use CUDA-event kernel timer with explicit state checks.", "methods": [ { "signature": "start()", "summary": "Records the start event." }, { "signature": "stop()", "summary": "Records the stop event." }, { "signature": "elapsedMillis() -> float", "summary": "Synchronizes the recorded stop event and returns elapsed milliseconds." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Source context & timing", "module": "mslk::utils", "platforms": [ "NVIDIA" ], "tags": [ "diagnostics", "source_location", "events", "benchmark" ], "source": "include/mslk/utils/bench/kernel_execution_timer.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/bench/kernel_execution_timer.cuh", "href": "#api-kernel-execution-timer" }, { "name": "TensorAccessor", "signature": "TensorAccessor", "summary": "ATen-compatible tensor accessor with optional named/contextual device bounds assertions.", "methods": [ { "signature": "using PtrType = typename PtrTraits::PtrType", "summary": "Resolved pointer type for the selected pointer trait." }, { "signature": "TensorAccessor(PtrType data, const index_t* sizes, const index_t* strides, const char* name, const char* context)", "summary": "Constructs a contextual accessor over external size/stride storage." }, { "signature": "operator[](index)", "summary": "Indexes a dimension or scalar element." }, { "signature": "numel() -> size_t", "summary": "Returns the addressed storage extent." }, { "signature": "at(index) -> T&", "summary": "Bounds-checked flat access in the enhanced implementation." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-tensor-accessor" }, { "name": "PackedTensorAccessor", "signature": "PackedTensorAccessor", "summary": "By-value size/stride accessor suitable for passing into GPU kernels.", "methods": [ { "signature": "using PtrType = typename PtrTraits::PtrType", "summary": "Resolved pointer type for the selected pointer trait." }, { "signature": "PackedTensorAccessor(PtrType data, const index_t* sizes, const index_t* strides, const char* name, const char* context)", "summary": "Constructs from native index_t size/stride arrays." }, { "signature": "template PackedTensorAccessor(PtrType data, const source_index_t* sizes, const source_index_t* strides, const char* name, const char* context)", "summary": "Converts int64 size/stride arrays when the accessor uses a narrower index type." }, { "signature": "operator[] / numel() / at()", "summary": "Same indexing contract as TensorAccessor." }, { "signature": "transpose(dim1, dim2)", "summary": "Returns a host-side accessor with exchanged sizes/strides." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-packed-tensor-accessor" }, { "name": "PackedTensorAccessor32", "signature": "PackedTensorAccessor32 = PackedTensorAccessor<..., int32_t>", "summary": "Packed accessor using 32-bit indices.", "kind": "C++ alias", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-packed-tensor-accessor32" }, { "name": "PackedTensorAccessor64", "signature": "PackedTensorAccessor64 = PackedTensorAccessor<..., int64_t>", "summary": "Packed accessor using 64-bit indices.", "kind": "C++ alias", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-packed-tensor-accessor64" }, { "name": "DefaultPtrTraits", "signature": "DefaultPtrTraits = at::DefaultPtrTraits", "summary": "Default accessor pointer trait alias.", "kind": "C++ alias", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-default-ptr-traits" }, { "name": "RestrictPtrTraits", "signature": "RestrictPtrTraits = at::RestrictPtrTraits", "summary": "CUDA/HIP restricted-pointer trait alias.", "notes": [ "Declared only while compiling with __CUDACC__ or __HIPCC__; it is absent from ordinary host-only C++ compilation." ], "kind": "C++ alias", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-restrict-ptr-traits" }, { "name": "pta", "signature": "namespace pta = mslk::utils // MSLK_MEMCHECK; otherwise at", "summary": "Compile-time selector between enhanced MSLK and standard ATen accessors.", "kind": "C++ namespace alias", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "global namespace", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-pta" }, { "name": "overflow_safe_int_t", "signature": "using overflow_safe_int_t = int64_t", "summary": "Global index type reserved for overflow-safe size arithmetic.", "kind": "C++ alias", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "global namespace", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-overflow-safe-int-t" }, { "name": "::PackedTensorAccessor", "signature": "template class PtrTraits, typename index_t> using PackedTensorAccessor = at::GenericPackedTensorAccessor", "summary": "Host-build global alias used when MSLK_MEMCHECK is not defined.", "notes": [ "This is distinct from mslk::utils::PackedTensorAccessor and exists only in the non-MSLK_MEMCHECK branch." ], "kind": "C++ alias", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "global namespace", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-global-packed-tensor-accessor" }, { "name": "NAME_MAX_LEN / CONTEXT_MAX_LEN", "signature": "NAME_MAX_LEN = 32; CONTEXT_MAX_LEN = 256", "summary": "Fixed diagnostic string capacities stored in enhanced accessors.", "kind": "C++ constants", "stability": "Developer", "category": "C++", "group": "Tensor accessors", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD", "CPU" ], "tags": [ "TensorAccessor", "bounds checking", "MSLK_MEMCHECK", "header" ], "source": "include/mslk/utils/torch/tensor_accessor.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h", "href": "#api-name-max-len-context-max-len" }, { "name": "scalar_type_for", "signature": "template at::ScalarType scalar_type_for()", "summary": "Maps a C++ scalar type to the corresponding ATen ScalarType.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/torch/tensor_accessor_builder.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h", "href": "#api-scalar-type-for" }, { "name": "TensorAccessorBuilder", "signature": "TensorAccessorBuilder", "summary": "Defers tensor validation and accessor construction until kernel launch context is known.", "methods": [ { "signature": "using index_t = conditional_t", "summary": "Index type selected from index_nbits." }, { "signature": "using accessor_t = conditional_t, pta::TensorAccessor<...>>", "summary": "Concrete accessor selected from the packed template flag." }, { "signature": "TensorAccessorBuilder(const std::string_view& name, const at::Tensor& tensor) noexcept", "summary": "Captures a tensor and diagnostic name without copying the tensor." }, { "signature": "validate_tensor(context)", "summary": "Checks rank, scalar type, and 32-bit index bounds." }, { "signature": "build_ta(context) / build_pta(context)", "summary": "Constructs standard or packed accessor." }, { "signature": "build(context)", "summary": "Chooses accessor form from the packed template parameter." }, { "signature": "checkValues(context)", "summary": "Checks the wrapped tensor for NaN/Inf in diagnostic mode." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/torch/tensor_accessor_builder.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h", "href": "#api-tensor-accessor-builder" }, { "name": "TA_B / PTA_B", "signature": "TA_B(tensor, T, N, index_nbits) · PTA_B(tensor, T, N, index_nbits)", "summary": "Creates named non-packed/packed accessor builders while capturing the variable name.", "kind": "C++ macros", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/torch/tensor_accessor_builder.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h", "href": "#api-ta-b-pta-b" }, { "name": "mslk::utils::pta", "signature": "namespace mslk::utils { namespace pta = mslk::utils; } // MSLK_MEMCHECK\nnamespace mslk::utils { namespace pta = at; } // otherwise", "summary": "Builder-local selector for enhanced MSLK versus standard ATen accessors.", "kind": "C++ namespace alias", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/torch/tensor_accessor_builder.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h", "href": "#api-mslk-utils-pta" }, { "name": "MAKE_PTA_WITH_NAME", "signature": "MAKE_PTA_WITH_NAME(...)", "summary": "Legacy named packed-accessor construction macro.", "kind": "C++ macro", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/torch/tensor_accessor_builder.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h", "href": "#api-make-pta-with-name" }, { "name": "is_tensor_accessor_builder", "signature": "is_tensor_accessor_builder; is_tensor_accessor_builder_v", "summary": "Detects TensorAccessorBuilder template instances.", "kind": "C++ trait", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/kernel_launcher.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh", "href": "#api-is-tensor-accessor-builder" }, { "name": "transform_kernel_arg", "signature": "transform_kernel_arg(const SourceContext&, T&& arg)", "summary": "Builds accessor-builder arguments with launch context; forwards other arguments.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/kernel_launcher.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh", "href": "#api-transform-kernel-arg" }, { "name": "check_kernel_arg", "signature": "check_kernel_arg(const SourceContext&, T&& arg)", "summary": "Runs configured tensor value checks on accessor-builder arguments.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/kernel_launcher.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh", "href": "#api-check-kernel-arg" }, { "name": "KernelLauncher", "signature": "KernelLauncher(const SourceContext& context) noexcept", "summary": "Validated normal/cooperative GPU launch wrapper with optional diagnostics and timing.", "methods": [ { "signature": "launch_kernel(kernel, grid, block, shared_mem, stream, args...)", "summary": "Validates dimensions/smem, transforms args, launches, checks errors, and optionally returns elapsed ms." }, { "signature": "checkGridSizesInRange / checkBlockSizesInRange / checkThreadCountNotExceeded", "summary": "Launch-shape validation helpers." }, { "signature": "checkSharedMemoryPerBlockNotExceeded", "summary": "Validates opt-in CUDA or HIP block shared-memory limit." }, { "signature": "kernelLaunchCheck()", "summary": "Raises a contextual launch/DSA failure." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/kernel_launcher.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh", "href": "#api-kernel-launcher" }, { "name": "MSLK_LAUNCH_KERNEL", "signature": "MSLK_LAUNCH_KERNEL(kernel, grid, block, smem, stream, ...)", "summary": "Instantiates KernelLauncher with build-time diagnostic toggles and launches a normal kernel.", "kind": "C++ macro", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/kernel_launcher.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh", "href": "#api-mslk-launch-kernel" }, { "name": "MSLK_LAUNCH_COOPERATIVE_KERNEL", "signature": "MSLK_LAUNCH_COOPERATIVE_KERNEL(kernel, grid, block, smem, stream, ...)", "summary": "Launches through the cooperative KernelLauncher specialization.", "kind": "C++ macro", "stability": "Developer", "category": "C++", "group": "Accessor builder & launch", "module": "mslk::utils", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "kernel launch", "validation", "NaN check", "DSA" ], "source": "include/mslk/utils/kernel_launcher.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh", "href": "#api-mslk-launch-cooperative-kernel" }, { "name": "TuningCache", "signature": "TuningCache(const std::string& kernelName)", "summary": "Persistent per-shape best-kernel cache with optional CUDA-graph benchmarking.", "notes": [ "Cache directory is MSLK_CACHE_DIR or $HOME/.mslk.", "Only device 0 writes cache files.", "MSLK_AUTOTUNE_USE_CUDA_GRAPH and MSLK_AUTOTUNE_COLLECT_STATS are presence-based switches." ], "methods": [ { "signature": "findBestKernelMaybeAutotune(cache_key, kernel_map, args...) -> Kernel", "summary": "Returns a cached winner or benchmarks candidates and persists the result." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/utils/tuning_cache.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/tuning_cache.cuh", "href": "#api-tuning-cache" }, { "name": "KernelMode", "signature": "enum class KernelMode { Small, Medium, Large, Default }", "summary": "Coarse GEMM-shape mode used to select CUTLASS kernels.", "kind": "C++ enum", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/gemm/cutlass/kernel_mode.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/kernel_mode.h", "href": "#api-kernel-mode" }, { "name": "get_kernel_mode", "signature": "KernelMode get_kernel_mode(at::Tensor XQ, at::Tensor WQ)", "summary": "Classifies a non-batched GEMM into a KernelMode.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/gemm/cutlass/kernel_mode.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/kernel_mode.h", "href": "#api-get-kernel-mode" }, { "name": "get_batched_kernel_mode", "signature": "KernelMode get_batched_kernel_mode(at::Tensor XQ, at::Tensor WQ)", "summary": "Classifies a batched GEMM into a KernelMode.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "mslk", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/gemm/cutlass/kernel_mode.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/kernel_mode.h", "href": "#api-get-batched-kernel-mode" }, { "name": "LinearCombinationOnDevice", "signature": "template class LinearCombinationOnDevice", "summary": "CUTLASS epilogue linear-combination operator whose scale data lives on device.", "methods": [ { "signature": "using ElementOutput / ElementAccumulator / ElementCompute", "summary": "Public scalar-type aliases for output, accumulator, and compute types." }, { "signature": "using FragmentOutput / FragmentAccumulator / ComputeFragment", "summary": "Count-element CUTLASS Array aliases for each processing stage." }, { "signature": "using ParamsBase = LinearCombinationParams", "summary": "Base parameter-storage alias." }, { "signature": "LinearCombinationOnDevice(Params const& params)", "summary": "Constructs the epilogue functor and retains the device alpha pointer." }, { "signature": "bool is_source_needed() const", "summary": "Reports whether the selected scale mode needs a source fragment." }, { "signature": "void set_k_partition(int k_partition, int k_partition_count)", "summary": "Adjusts beta for serial K-partition reduction." }, { "signature": "FragmentOutput operator()(FragmentAccumulator const&, FragmentOutput const&) const", "summary": "Computes D = alpha × accumulator + beta × source." }, { "signature": "FragmentOutput operator()(FragmentAccumulator const&) const", "summary": "Computes D = alpha × accumulator." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "cutlass::epilogue::thread", "platforms": [ "NVIDIA" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/gemm/cutlass/threadblock.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/threadblock.h", "href": "#api-linear-combination-on-device" }, { "name": "LinearCombinationOnDevice::Params", "signature": "struct LinearCombinationOnDevice::Params : LinearCombinationParams", "summary": "Host/device-constructable alpha/beta parameter carrier for LinearCombinationOnDevice.", "methods": [ { "signature": "Params()", "summary": "Defaults to alpha=1 and beta=0." }, { "signature": "Params(ElementCompute alpha, ElementCompute beta)", "summary": "Stores scalar alpha and beta values." }, { "signature": "Params(ElementCompute alpha)", "summary": "Stores alpha with beta=0." }, { "signature": "Params(const ElementCompute* alpha_ptr, const ElementCompute* beta_ptr)", "summary": "Uses device pointers for both scales." }, { "signature": "Params(const ElementCompute* alpha_ptr)", "summary": "Uses a device alpha pointer and no beta pointer." }, { "signature": "Params(ParamsBase const& base)", "summary": "Converts CUTLASS base parameter storage." } ], "kind": "C++ type", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "cutlass::epilogue::thread", "platforms": [ "NVIDIA" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/gemm/cutlass/threadblock.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/threadblock.h", "href": "#api-linear-combination-on-device-params" }, { "name": "GroupedGemmInputType", "signature": "enum GroupedGemmInputType { _2D2D, _2D3D }", "summary": "Selects grouped GEMM activation/weight rank layout.", "kind": "C++ enum", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/gemm/cutlass/grouped_common.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/grouped_common.cuh", "href": "#api-grouped-gemm-input-type" }, { "name": "set_grouped_gemm_args_kernel", "signature": "template \n__global__ void set_grouped_gemm_args_kernel(int64_t G, int64_t M, int64_t N, int64_t K, ProblemShape* problem_shape_ptr, ElementA* xq, const ElementA** xq_ptr, ElementB* wq, const ElementB** wq_ptr, ScaleDtype* x_scale, const ScaleDtype** x_scale_ptr, int32_t x_scale_size, ScaleDtype* w_scale, const ScaleDtype** w_scale_ptr, ElementC* output, ElementC** output_ptr, StrideA* stride_a_ptr, StrideB* stride_b_ptr, StrideC* stride_c_ptr, int32_t* offsets, LayoutSFA* layout_SFA, LayoutSFB* layout_SFB, GroupedGemmInputType gemm_type, bool is_transposeAB, ElementGlobalScale* global_scale=nullptr, const ElementGlobalScale** global_scale_ptr=nullptr)", "summary": "Builds per-group CUTLASS pointer/shape/stride argument arrays on device.", "kind": "CUDA kernel", "stability": "Developer", "category": "C++", "group": "Autotuning & CUTLASS helpers", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "autotune", "cache", "CUTLASS", "grouped GEMM" ], "source": "include/mslk/gemm/cutlass/grouped_common.cuh", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/grouped_common.cuh", "href": "#api-set-grouped-gemm-args-kernel" }, { "name": "mslk::gemm::f8f8bf16_blockwise", "signature": "at::Tensor mslk::gemm::f8f8bf16_blockwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, int64_t block_m=128, int64_t block_n=128, int64_t block_k=128)", "summary": "Native rectangular-block-scaled FP8 GEMM returning BF16.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-blockwise" }, { "name": "mslk::gemm::f8f8bf16_rowwise_preshuffle", "signature": "at::Tensor mslk::gemm::f8f8bf16_rowwise_preshuffle(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional bias=std::nullopt, bool use_fast_accum=true)", "summary": "ROCm preshuffled rowwise FP8 GEMM returning BF16.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "AMD ROCm" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-preshuffle" }, { "name": "mslk::gemm::f8f8f16_rowwise_preshuffle", "signature": "at::Tensor mslk::gemm::f8f8f16_rowwise_preshuffle(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional bias=std::nullopt, bool use_fast_accum=true)", "summary": "ROCm preshuffled rowwise FP8 GEMM declared with FP16 output.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "notes": [ "The current dispatcher registration mistakenly binds this schema to f8f8bf16_rowwise_preshuffle; treat that as a source defect." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "AMD ROCm" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8f16-rowwise-preshuffle" }, { "name": "mslk::gemm::f8f8bf16_rowwise_grouped", "signature": "std::vector mslk::gemm::f8f8bf16_rowwise_grouped(at::TensorList XQ, at::TensorList WQ, at::TensorList x_scale, at::TensorList w_scale)", "summary": "Native list-based grouped rowwise FP8 GEMM.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-grouped" }, { "name": "mslk::gemm::f8f8bf16_rowwise_grouped_cat", "signature": "at::Tensor mslk::gemm::f8f8bf16_rowwise_grouped_cat(at::TensorList XQ, at::TensorList WQ, at::TensorList x_scale, at::TensorList w_scale)", "summary": "Native list-based grouped rowwise FP8 GEMM with concatenated output.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-grouped-cat" }, { "name": "mslk::gemm::f8f8bf16_rowwise_grouped_stacked", "signature": "at::Tensor mslk::gemm::f8f8bf16_rowwise_grouped_stacked(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor M_sizes)", "summary": "Native stacked grouped rowwise FP8 GEMM.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-grouped-stacked" }, { "name": "mslk::gemm::f8f8bf16_rowwise_grouped_dynamic", "signature": "at::Tensor mslk::gemm::f8f8bf16_rowwise_grouped_dynamic(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor zero_start_index_M, bool zeroing_output_tensor=true)", "summary": "Native dynamic-row grouped FP8 GEMM.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-grouped-dynamic" }, { "name": "mslk::gemm::f8f8bf16_rowwise_batched", "signature": "at::Tensor mslk::gemm::f8f8bf16_rowwise_batched(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional bias=std::nullopt, bool use_fast_accum=true, std::optional output=std::nullopt)", "summary": "Native batched rowwise FP8 GEMM with optional output reuse.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-batched" }, { "name": "mslk::gemm::bf16bf16bf16_grouped", "signature": "std::vector mslk::gemm::bf16bf16bf16_grouped(at::TensorList X, at::TensorList W)", "summary": "Native list-based grouped BF16 GEMM.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16bf16bf16-grouped" }, { "name": "mslk::gemm::bf16bf16bf16_grouped_cat", "signature": "at::Tensor mslk::gemm::bf16bf16bf16_grouped_cat(at::TensorList X, at::TensorList W)", "summary": "Native list-based grouped BF16 GEMM with concatenated output.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16bf16bf16-grouped-cat" }, { "name": "mslk::gemm::bf16bf16bf16_grouped_dynamic", "signature": "at::Tensor mslk::gemm::bf16bf16bf16_grouped_dynamic(at::Tensor X, at::Tensor W, at::Tensor zero_start_index_M)", "summary": "Native dynamic-row grouped BF16 GEMM.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16bf16bf16-grouped-dynamic" }, { "name": "mslk::gemm::bf16bf16bf16_grouped_stacked", "signature": "at::Tensor mslk::gemm::bf16bf16bf16_grouped_stacked(at::Tensor X, at::Tensor W, at::Tensor M_sizes, std::optional out=std::nullopt, std::optional num_sms=std::nullopt)", "summary": "Native stacked grouped BF16 GEMM with output/SM overrides.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16bf16bf16-grouped-stacked" }, { "name": "mslk::gemm::f8f8bf16_rowwise", "signature": "at::Tensor mslk::gemm::f8f8bf16_rowwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional bias=std::nullopt, bool use_fast_accum=true)", "summary": "Native rowwise FP8 GEMM returning BF16.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise" }, { "name": "mslk::gemm::f8f8bf16_rowwise_out", "signature": "void mslk::gemm::f8f8bf16_rowwise_out(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor output, std::optional bias=std::nullopt, bool use_fast_accum=true)", "summary": "Native rowwise FP8 GEMM writing a caller-owned output tensor.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-out" }, { "name": "mslk::gemm::f8f8f16_rowwise", "signature": "at::Tensor mslk::gemm::f8f8f16_rowwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional bias=std::nullopt, bool use_fast_accum=true)", "summary": "ROCm native rowwise FP8 GEMM returning FP16.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "AMD ROCm" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8f16-rowwise" }, { "name": "mslk::gemm::f8f8bf16_groupwise", "signature": "at::Tensor mslk::gemm::f8f8bf16_groupwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale)", "summary": "Native declaration for K-group-scaled FP8 GEMM.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "notes": [ "The ROCm schema implementation is supplied by Python/Triton." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-groupwise" }, { "name": "mslk::gemm::f8f8bf16_groupwise_grouped", "signature": "at::Tensor mslk::gemm::f8f8bf16_groupwise_grouped(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor M_sizes)", "summary": "Native declaration for stacked grouped K-group FP8 GEMM.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "notes": [ "The ROCm schema implementation is supplied by Python/Triton." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8f8bf16-groupwise-grouped" }, { "name": "mslk::gemm::i8i8bf16", "signature": "at::Tensor mslk::gemm::i8i8bf16(at::Tensor XQ, at::Tensor WQ, double scale, int64_t split_k)", "summary": "Native INT8 GEMM with scalar dequantization scale.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "notes": [ "The dispatcher schema supplies split_k=1; the C++ declaration itself has no default." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-i8i8bf16" }, { "name": "mslk::gemm::i8i8bf16_dynamic", "signature": "at::Tensor mslk::gemm::i8i8bf16_dynamic(at::Tensor XQ, at::Tensor WQ, at::Tensor scale, int64_t split_k=1)", "summary": "Native INT8 GEMM with tensor-valued dynamic scale.", "details": "Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.", "notes": [ "The ROCm schema implementation is supplied by Python/Triton." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP8 & BF16", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "GEMM" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-i8i8bf16-dynamic" }, { "name": "mslk::gemm::f4f4bf16_grouped_stacked", "signature": "at::Tensor mslk::gemm::f4f4bf16_grouped_stacked(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor M_sizes, std::optional global_scale=std::nullopt, std::optional starting_row_after_padding=std::nullopt, bool use_mx=true)", "summary": "Native stacked grouped MXFP4/NVFP4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f4f4bf16-grouped-stacked" }, { "name": "mslk::gemm::bf16x9_gemm", "signature": "at::Tensor mslk::gemm::bf16x9_gemm(at::Tensor A, at::Tensor B, std::optional output=std::nullopt)", "summary": "Native cuBLAS BF16x9 emulation GEMM producing FP32.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16x9-gemm" }, { "name": "mslk::gemm::bf16i4bf16_shuffled", "signature": "at::Tensor mslk::gemm::bf16i4bf16_shuffled(at::Tensor X, at::Tensor W, at::Tensor w_scale_group, at::Tensor w_zero_group)", "summary": "Native BF16 × preshuffled INT4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16i4bf16-shuffled" }, { "name": "mslk::gemm::f8i4bf16_shuffled_grouped", "signature": "at::Tensor mslk::gemm::f8i4bf16_shuffled_grouped(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor w_scale_group, at::Tensor M_sizes)", "summary": "Native stacked grouped FP8 × preshuffled INT4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8i4bf16-shuffled-grouped" }, { "name": "mslk::gemm::bf16i4bf16_shuffled_grouped", "signature": "at::Tensor mslk::gemm::bf16i4bf16_shuffled_grouped(at::Tensor X, at::Tensor WQ, at::Tensor w_scale_group, at::Tensor w_zero_group, at::Tensor M_sizes)", "summary": "Native stacked grouped BF16 × preshuffled INT4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16i4bf16-shuffled-grouped" }, { "name": "mslk::gemm::bf16i4bf16_shuffled_batched", "signature": "at::Tensor mslk::gemm::bf16i4bf16_shuffled_batched(at::Tensor X, at::Tensor WQ, at::Tensor w_scale, at::Tensor w_zp)", "summary": "Native batched BF16 × preshuffled INT4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16i4bf16-shuffled-batched" }, { "name": "mslk::gemm::bf16i4bf16_rowwise_batched", "signature": "at::Tensor mslk::gemm::bf16i4bf16_rowwise_batched(at::Tensor X, at::Tensor WQ, at::Tensor w_scale, at::Tensor w_zp)", "summary": "Native declaration for batched BF16 × rowwise INT4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "notes": [ "The ROCm schema implementation is supplied by Python/Triton." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16i4bf16-rowwise-batched" }, { "name": "mslk::gemm::bf16i4bf16_rowwise", "signature": "at::Tensor mslk::gemm::bf16i4bf16_rowwise(at::Tensor X, at::Tensor W, at::Tensor w_scale_group, at::Tensor w_zero_group)", "summary": "Native declaration for BF16 × row/group-quantized INT4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "notes": [ "The ROCm schema implementation is supplied by Python/Triton." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-bf16i4bf16-rowwise" }, { "name": "mslk::gemm::f8i4bf16_rowwise", "signature": "at::Tensor mslk::gemm::f8i4bf16_rowwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor w_zp)", "summary": "Native rowwise FP8 activation × INT4 weight GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8i4bf16-rowwise" }, { "name": "mslk::gemm::f8i4bf16_shuffled", "signature": "at::Tensor mslk::gemm::f8i4bf16_shuffled(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor w_scale_group)", "summary": "Native FP8 activation × preshuffled INT4 weight GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-f8i4bf16-shuffled" }, { "name": "mslk::gemm::preshuffle_i4", "signature": "std::tuple mslk::gemm::preshuffle_i4(at::Tensor WQ, at::Tensor w_scale)", "summary": "Native one-time INT4 weight/scale preprocessing entry.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-preshuffle-i4" }, { "name": "mslk::gemm::mx8mx4bf16", "signature": "at::Tensor mslk::gemm::mx8mx4bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional output=std::nullopt)", "summary": "Native declaration for MXFP8 × MXFP4 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "notes": [ "The ROCm schema implementation is supplied by Python/Triton." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-mx8mx4bf16" }, { "name": "mslk::gemm::mx8mx6bf16", "signature": "at::Tensor mslk::gemm::mx8mx6bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional output=std::nullopt)", "summary": "Native MXFP8 × MXFP6 GEMM.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-mx8mx6bf16" }, { "name": "mslk::gemm::mx6mx6bf16", "signature": "at::Tensor mslk::gemm::mx6mx6bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional output=std::nullopt, int64_t splits=0)", "summary": "Native MXFP6 × MXFP6 GEMM with optional split reduction.", "details": "Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native GEMM declarations · FP4, INT4 & MX", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "packed" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-mslk-gemm-mx6mx6bf16" }, { "name": "mslk::gemm::f8f8bf16_rowwise_grouped_mm", "signature": "at::Tensor mslk::gemm::f8f8bf16_rowwise_grouped_mm(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional offsets, at::Tensor& output)", "summary": "ROCm native layout-polymorphic grouped rowwise FP8 GEMM.", "details": "Published native declarations for offset-described or layout-polymorphic grouped GEMMs.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native offset-grouped GEMM declarations", "module": "mslk::gemm", "platforms": [ "AMD ROCm" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "grouped MM" ], "source": "include/mslk/gemm/gemm_torch.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h", "href": "#api-mslk-gemm-f8f8bf16-rowwise-grouped-mm" }, { "name": "mslk::gemm::mx8mx8bf16_grouped_mm", "signature": "at::Tensor mslk::gemm::mx8mx8bf16_grouped_mm(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor offsets, std::optional output=std::nullopt, std::optional actual_num_tokens=std::nullopt)", "summary": "Native offset-grouped MXFP8 × MXFP8 GEMM.", "details": "Published native declarations for offset-described or layout-polymorphic grouped GEMMs.", "notes": [ "The ROCm dispatcher implementation is supplied by Python/Triton rather than this declaration." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native offset-grouped GEMM declarations", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "grouped MM" ], "source": "include/mslk/gemm/gemm_torch.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h", "href": "#api-mslk-gemm-mx8mx8bf16-grouped-mm" }, { "name": "mslk::gemm::f4f4bf16_grouped_mm", "signature": "at::Tensor mslk::gemm::f4f4bf16_grouped_mm(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor offsets, std::optional output=std::nullopt, std::optional global_scale=std::nullopt)", "summary": "Native offset-grouped MXFP4/NVFP4 GEMM.", "details": "Published native declarations for offset-described or layout-polymorphic grouped GEMMs.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native offset-grouped GEMM declarations", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "grouped MM" ], "source": "include/mslk/gemm/gemm_torch.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h", "href": "#api-mslk-gemm-f4f4bf16-grouped-mm" }, { "name": "mslk::gemm::f4f4bf16_ultra_grouped_mm", "signature": "at::Tensor mslk::gemm::f4f4bf16_ultra_grouped_mm(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor offsets, at::Tensor x_global_scale, at::Tensor w_global_scale, std::optional output=std::nullopt)", "summary": "Native ultra grouped FP4 GEMM with separate global scales.", "details": "Published native declarations for offset-described or layout-polymorphic grouped GEMMs.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native offset-grouped GEMM declarations", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "grouped MM" ], "source": "include/mslk/gemm/gemm_torch.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h", "href": "#api-mslk-gemm-f4f4bf16-ultra-grouped-mm" }, { "name": "mslk::gemm::mx8mx4bf16_grouped_mm", "signature": "at::Tensor mslk::gemm::mx8mx4bf16_grouped_mm(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor offsets, std::optional output=std::nullopt)", "summary": "Native offset-grouped MXFP8 × MXFP4 GEMM.", "details": "Published native declarations for offset-described or layout-polymorphic grouped GEMMs.", "notes": [ "The ROCm dispatcher implementation is supplied by Python/Triton rather than this declaration." ], "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native offset-grouped GEMM declarations", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "grouped MM" ], "source": "include/mslk/gemm/gemm_torch.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h", "href": "#api-mslk-gemm-mx8mx4bf16-grouped-mm" }, { "name": "mslk::gemm::f4f4bf16", "signature": "at::Tensor mslk::gemm::f4f4bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional output=std::nullopt, std::optional global_scale=std::nullopt, int64_t mxfp4_block_size=32)", "summary": "Native packed FP4 GEMM selecting MXFP4, MXFP4-16, or NVFP4.", "details": "Published native declarations for offset-described or layout-polymorphic grouped GEMMs.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native offset-grouped GEMM declarations", "module": "mslk::gemm", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror", "grouped MM" ], "source": "include/mslk/gemm/gemm_torch.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h", "href": "#api-mslk-gemm-f4f4bf16" }, { "name": "mslk::conv::f8f8bf16_conv", "signature": "at::Tensor mslk::conv::f8f8bf16_conv(at::Tensor activation, at::Tensor filter, at::Tensor scale, std::vector padding, std::vector stride, std::vector dilation)", "summary": "Native FP8 3D convolution/cross-correlation entry.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native Conv & MoE declarations", "module": "mslk::conv", "platforms": [ "NVIDIA SM100+" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror" ], "source": "include/mslk/conv/conv.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/conv/conv.h", "href": "#api-mslk-conv-f8f8bf16-conv" }, { "name": "mslk::moe::index_shuffling_torch", "signature": "std::tuple mslk::moe::index_shuffling_torch(const at::Tensor& routing_scores, const std::optional& expert_index_start, const std::optional& expert_index_end, const std::optional& valid_token_count, int64_t top_k)", "summary": "Native implementation entry for the index_shuffling dispatcher schema.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native Conv & MoE declarations", "module": "mslk::moe", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror" ], "source": "include/mslk/moe/moe.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/moe/moe.h", "href": "#api-mslk-moe-index-shuffling-torch" }, { "name": "mslk::moe::scatter_add_along_first_dim", "signature": "void mslk::moe::scatter_add_along_first_dim(at::Tensor dst, at::Tensor src, at::Tensor index)", "summary": "Native in-place first-dimension scatter-add entry.", "kind": "C++ function", "stability": "Developer", "category": "C++", "group": "Native Conv & MoE declarations", "module": "mslk::moe", "platforms": [ "NVIDIA" ], "tags": [ "ATen", "native entrypoint", "dispatcher mirror" ], "source": "include/mslk/moe/moe.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/moe/moe.h", "href": "#api-mslk-moe-scatter-add-along-first-dim" }, { "name": "Native GEMM / Conv / MoE declarations", "signature": "#include \n#include \n#include \n#include ", "summary": "ATen-native declarations corresponding to the documented torch.ops schemas.", "details": "gemm.h and gemm_torch.h declare quantized/grouped GEMMs; conv.h declares f8f8bf16_conv; moe.h declares index_shuffling_torch and scatter_add_along_first_dim. The dispatcher schemas are the more stable Python-facing contract.", "methods": [ { "signature": "gemm.h / gemm_torch.h", "summary": "Native GEMM declarations for every platform-compiled family." }, { "signature": "conv.h::f8f8bf16_conv(...)", "summary": "Native FP8 convolution entry." }, { "signature": "moe.h::index_shuffling_torch(...)", "summary": "Native routing entry." }, { "signature": "moe.h::scatter_add_along_first_dim(...)", "summary": "Native in-place scatter-add entry." } ], "kind": "C++ API family", "stability": "Developer", "category": "C++", "group": "Native entry headers", "module": "mslk C++ headers", "platforms": [ "NVIDIA", "AMD" ], "tags": [ "ATen", "native entrypoint", "headers" ], "source": "include/mslk/gemm/gemm.h", "sourceUrl": "https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h", "href": "#api-native-gemm-conv-mo-e-declarations" } ] }