# MSLK - complete kernel reference for LLM agents

> Every documented MSLK symbol with its signature, behavior, constraints and source location. For orientation (shape conventions, which backend or GEMM op to call, sharp edges) read llms.txt first.

Source: https://github.com/meta-pytorch/MSLK @ 69ae1b897f2d546d72acab129e1ae4bc2924900f. Generated from docs/index.html by docs/build-llms-txt.mjs - do not edit by hand.

Contents: 54 groups, 424 symbols.

## Attention - Core fMHA

### memory_efficient_attention

```python
memory_efficient_attention(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None) -> Tensor
```

Autograd-enabled fused attention with automatic backend dispatch.

Accepts [B,M,K], [B,M,H,K], or experimental [B,M,G,H,K] inputs. It computes scaled QK^T, 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.

**Methods / arguments:**
- `op=(FwOp, BwOp)` - Explicitly pins forward and backward operator classes.
- `attn_bias=AttentionBias | Tensor | None` - Uses optimized masks when possible; arbitrary dense tensors are generally slower.

**Constraints:**
- 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.

Kind: function - Stability: Public - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: fMHA, dispatch, BMK, BMHK, BMGHK
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### memory_efficient_attention_forward

```python
memory_efficient_attention_forward(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None) -> Tensor
```

Forward-only attention path that does not retain manual-backward context.

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.

**Constraints:**
- 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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: fMHA, dispatch, BMK, BMHK, BMGHK
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### memory_efficient_attention_forward_requires_grad

```python
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]
```

Manual-backward forward pass returning output and log-sum-exp.

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

**Constraints:**
- Nonzero dropout is rejected by this manual API.
- The forward and backward operators must agree on variable-length LSE layout.

Kind: function - Stability: Public - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: fMHA, dispatch, BMK, BMHK, BMGHK
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### memory_efficient_attention_backward

```python
memory_efficient_attention_backward(grad, output, lse, query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None) -> tuple[Tensor, Tensor, Tensor]
```

Explicit backward pass for an output/LSE pair.

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.

**Constraints:**
- 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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: fMHA, dispatch, BMK, BMHK, BMGHK
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### memory_efficient_attention_partial

```python
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]
```

Computes an output and LSE for one disjoint K/V shard.

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

**Constraints:**
- 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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: fMHA, dispatch, BMK, BMHK, BMGHK
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### merge_attentions

```python
merge_attentions(attn_split, lse_split, write_lse=True, output_dtype=None) -> tuple[Tensor, Optional[Tensor]]
```

Log-sum-exp-correct merge of attention computed over K/V chunks.

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

**Constraints:**
- Direct autograd through this function is unsupported; use merge_attentions_autograd.
- Inputs requiring gradients require write_lse=True.

Kind: function - Stability: Public - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: fMHA, dispatch, BMK, BMHK, BMGHK
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

## Attention - Compile-friendly wrappers

### torch.ops.mslk.memory_efficient_attention_forward

```python
torch.ops.mslk.memory_efficient_attention_forward(q, k, v, b=None, p=0.0, scale=None) -> Tensor
```

Compile-friendly forward attention with optional bias.

These custom-library wrappers keep dispatch outside a torch.compile trace and expose a reduced forward-only contract.

**Returns:** Attention output tensor matching q.

**Constraints:**
- 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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, Meta / fake - Topics: torch.compile, custom op, forward
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### torch.ops.mslk.memory_efficient_attention_forward_with_bias

```python
torch.ops.mslk.memory_efficient_attention_forward_with_bias(q, k, v, b, p=0.0, scale=None) -> Tensor
```

Compile-friendly forward attention requiring an explicit bias tensor.

These custom-library wrappers keep dispatch outside a torch.compile trace and expose a reduced forward-only contract.

**Returns:** Attention output tensor matching q.

**Constraints:**
- 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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, Meta / fake - Topics: torch.compile, custom op, forward
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

## Attention - Trainable partial attention

### Partial

```python
Partial(attn: Tensor, lse: Tensor, placeholder: Tensor)
```

Autograd carrier for a partial attention output and its merge context.

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 / arguments:**
- `is_bmghk() -> bool` - Reports whether the stored output uses the grouped-head BMGHK layout.
- `apply(fn: Callable[[Tensor], Tensor]) -> Partial` - Applies a layout-preserving transform that must not manipulate the embedding dimension.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.merge_training` - Platforms: NVIDIA, AMD - Topics: autograd, split KV, LSE
Source: [mslk/attention/fmha/merge_training.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/merge_training.py)

### memory_efficient_attention_partial_autograd

```python
memory_efficient_attention_partial_autograd(query, key, value, attn_bias=None, p=0.0, scale=None, *, op=None, output_dtype=None) -> Partial
```

Autograd-safe wrapper around partial attention.

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.

**Constraints:**
- Arguments mirror memory_efficient_attention_partial.
- Dropout remains unsupported by the underlying partial path.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.merge_training` - Platforms: NVIDIA, AMD - Topics: autograd, split KV, LSE
Source: [mslk/attention/fmha/merge_training.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/merge_training.py)

### merge_attentions_autograd

```python
merge_attentions_autograd(*partials: Partial) -> Tensor
```

Merges one or more Partial objects with correct gradients.

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.

**Constraints:**
- Raises ValueError when called without partials.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.merge_training` - Platforms: NVIDIA, AMD - Topics: autograd, split KV, LSE
Source: [mslk/attention/fmha/merge_training.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/merge_training.py)

## Attention - Biases & masks

### AttentionBias

```python
AttentionBias()
```

Abstract base for optimized attention-bias descriptions.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

**Constraints:**
- Subclasses define materialize and, where state is stored, to(device).

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### LowerTriangularMask

```python
LowerTriangularMask(device: Optional[torch.device] = None)
```

Top-left-aligned causal mask.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `add_bias(bias: Tensor) -> LowerTriangularMaskWithTensorBias` - Combines causal masking with an arbitrary additive tensor.
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Produces zeros on/below the diagonal and -inf above it.

**Constraints:**
- The device constructor argument is retained only for compatibility and is ignored.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### LowerTriangularMaskWithTensorBias

```python
LowerTriangularMaskWithTensorBias(bias: Tensor)
```

Top-left causal mask plus a dense additive tensor bias.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

**Constraints:**
- Backend support is narrower than for a pure causal mask.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### LowerTriangularFromBottomRightMask

```python
LowerTriangularFromBottomRightMask()
```

Bottom-right-aligned causal mask for Mq != Mkv.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `make_local_attention(window_size: int) -> LowerTriangularFromBottomRightLocalAttentionMask` - Restricts the causal region to a trailing window.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### LowerTriangularFromBottomRightLocalAttentionMask

```python
LowerTriangularFromBottomRightLocalAttentionMask(_window_size: int)
```

Bottom-right causal sliding-window mask.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### LocalAttentionFromBottomRightMask

```python
LocalAttentionFromBottomRightMask(window_left: int, window_right: int)
```

Non-causal local window aligned from the bottom right.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalMask

```python
BlockDiagonalMask(q_seqinfo, k_seqinfo, _batch_sizes=None)
```

Packed independent sequences where each Q block sees its matching K/V block.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens(q_seqlen, kv_seqlen=None, *, device=None) -> BlockDiagonalMask` - Builds sequence metadata from lengths.
- `from_tensor_list(tensors) -> tuple[BlockDiagonalMask, Tensor]` - Concatenates a list and returns its matching block mask.
- `from_tensor_lists_qkv(tensors_q, tensors_k, tensors_v=None) -> tuple[BlockDiagonalMask, Tensor, Tensor, Optional[Tensor]]` - Packs Q/K/(V) lists and returns all tensors plus a mask.
- `split_queries(tensor) / split_kv(tensor) / split(tensor)` - Restores packed outputs to the original logical sequences.
- `make_causal() / make_causal_from_bottomright()` - Returns the corresponding causal subclass.
- `make_local_attention(window_size) / make_local_attention_from_bottomright(window_size)` - Returns a local block-diagonal variant.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalCausalMask

```python
BlockDiagonalCausalMask(q_seqinfo, k_seqinfo, _batch_sizes=None)
```

Top-left causal mask inside every packed sequence block.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalCausalFromBottomRightMask

```python
BlockDiagonalCausalFromBottomRightMask(q_seqinfo, k_seqinfo, _batch_sizes=None)
```

Bottom-right causal mask inside every packed block.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

**Constraints:**
- Each K/V block must be at least as long as its Q block.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalCausalLocalAttentionMask

```python
BlockDiagonalCausalLocalAttentionMask(q_seqinfo, k_seqinfo, _batch_sizes=None, _window_size=0)
```

Top-left causal local window inside packed blocks.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

**Constraints:**
- Although the generated dataclass signature defaults _window_size to 0, construction rejects non-positive values.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalCausalLocalAttentionFromBottomRightMask

```python
BlockDiagonalCausalLocalAttentionFromBottomRightMask(q_seqinfo, k_seqinfo, _batch_sizes=None, _window_size=0)
```

Bottom-right causal local window inside packed blocks.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

**Constraints:**
- Although the generated dataclass signature defaults _window_size to 0, construction rejects non-positive values.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalPaddedKeysMask

```python
BlockDiagonalPaddedKeysMask(q_seqinfo, k_seqinfo)
```

Packed Q blocks attending fixed-capacity padded K/V lanes.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens(q_seqlen, kv_padding, kv_seqlen, causal_diagonal=None, *, device=None)` - Creates padded-lane metadata.
- `make_paged(block_tables, page_size, paged_type)` - Maps logical lanes to physical KV pages.
- `make_local_attention(window_left, window_right)` - Returns a padded local-attention variant.

**Constraints:**
- causal_diagonal currently exists for compatibility and is unused.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalCausalWithOffsetPaddedKeysMask

```python
BlockDiagonalCausalWithOffsetPaddedKeysMask(q_seqinfo, k_seqinfo, causal_diagonal=None)
```

Bottom-right causal masking over padded K/V lanes.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens(q_seqlen, kv_padding, kv_seqlen, causal_diagonal=None, *, device=None)` - Creates bottom-right causal padded metadata.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalLocalAttentionPaddedKeysMask

```python
BlockDiagonalLocalAttentionPaddedKeysMask(q_seqinfo, k_seqinfo, window_left, window_right)
```

Non-causal local windows over padded K/V lanes.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens_local(q_seqlen, kv_padding, kv_seqlen, window_left, window_right)` - Constructs the local padded mask.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalCausalLocalAttentionPaddedKeysMask

```python
BlockDiagonalCausalLocalAttentionPaddedKeysMask(q_seqinfo, k_seqinfo, _window_size)
```

Bottom-right causal local windows over padded K/V lanes.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens_local(q_seqlen, kv_padding, kv_seqlen, window_size)` - Constructs the causal-local padded mask.
- `make_paged(block_tables, page_size, paged_type)` - Returns the paged counterpart; paged_type is required.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### PagedBlockDiagonalPaddedKeysMask

```python
PagedBlockDiagonalPaddedKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)
```

Padded logical K/V lanes mapped to physical cache pages.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens(q_seqlen, kv_seqlen, block_tables, page_size, *, device=None)` - Builds page-aware metadata.

**Constraints:**
- Physical K/V uses [1, physical_pages x page_size, H, D] or its grouped-head equivalent.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### PagedBlockDiagonalCausalWithOffsetPaddedKeysMask

```python
PagedBlockDiagonalCausalWithOffsetPaddedKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)
```

Bottom-right causal paged padded mask.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### PagedBlockDiagonalCausalLocalPaddedKeysMask

```python
PagedBlockDiagonalCausalLocalPaddedKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size, _window_size)
```

Bottom-right causal local paged padded mask.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens_local(q_seqlen, kv_seqlen, block_tables, page_size, window_size, *, device=None)` - Constructs the paged causal-local mask.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalGappyKeysMask

```python
BlockDiagonalGappyKeysMask(q_seqinfo, k_seqinfo)
```

K/V lanes with arbitrary starts and independently used lengths.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens(q_seqlen, kv_seqstarts, kv_seqlen, *, device=None)` - Describes gappy lanes.
- `make_paged(block_tables, page_size, notional_padding, paged_type)` - Maps the gappy layout to physical pages.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalLocalAttentionFromBottomRightGappyKeysMask

```python
BlockDiagonalLocalAttentionFromBottomRightGappyKeysMask(q_seqinfo, k_seqinfo, window_left, window_right)
```

Bottom-right-aligned local windows over gappy K/V lanes.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens_local_gappy(q_seqlen, kv_seqstarts, kv_seqlen, window_left, window_right, device)` - Constructs the local gappy mask.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### BlockDiagonalCausalWithOffsetGappyKeysMask

```python
BlockDiagonalCausalWithOffsetGappyKeysMask(q_seqinfo, k_seqinfo)
```

Bottom-right causal masking over gappy K/V lanes.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### PagedBlockDiagonalGappyKeysMask

```python
PagedBlockDiagonalGappyKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)
```

Gappy logical K/V lanes mapped through a page table.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `from_seqlens(q_seqlen, kv_seqstarts, kv_seqlen, block_tables, page_size, *, device=None)` - Builds paged-gappy metadata.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### PagedBlockDiagonalCausalWithOffsetGappyKeysMask

```python
PagedBlockDiagonalCausalWithOffsetGappyKeysMask(q_seqinfo, k_seqinfo, block_tables, page_size)
```

Bottom-right causal paged-gappy mask.

AttentionBias objects describe sparsity or additive bias without materializing the full Q x K matrix, allowing compatible kernels to encode the mask directly.

**Returns:** An object accepted as attn_bias by fMHA.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

### VARLEN_BIASES

```python
VARLEN_BIASES: tuple[type[AttentionBias], ...]
```

Dispatch grouping for packed, padded, gappy, and paged variable-length masks.

Used by backend support checks and manual-backward LSE layout detection.

**Returns:** Tuple of bias base classes.

**Methods / arguments:**
- `materialize(shape, dtype=torch.float32, device='cpu') -> Tensor` - Builds the dense bias for debugging/reference checks; intentionally slow.
- `to(device)` - Moves stored sequence metadata or tensor bias to a device.

Kind: constant - Stability: Low-level - Module: `mslk.attention.fmha.attn_bias` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: attention bias, mask
Source: [mslk/attention/fmha/attn_bias.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/attn_bias.py)

## Attention - Backend operators

### cutlass.FwOp

```python
cutlass.FwOp
```

Classic CUTLASS forward attention with the broadest NVIDIA dtype/head-dim range.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Maximum Q/K head dimension is 65,536.
- Supports different V head dimension and grouped-head forward.
- Does not support partial attention.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cutlass` - Platforms: NVIDIA SM60-SM90 - Topics: FP32, FP16, BF16, dropout, tensor bias, BMGHK
Source: [mslk/attention/fmha/cutlass.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass.py)

### cutlass.BwOp

```python
cutlass.BwOp
```

Classic CUTLASS attention backward, including supported tensor-bias gradients.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cutlass` - Platforms: NVIDIA SM60-SM90 - Topics: FP32, FP16, BF16, dropout, bias grad
Source: [mslk/attention/fmha/cutlass.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass.py)

### cutlass_blackwell.FwOp

```python
cutlass_blackwell.FwOp
```

Blackwell CUTLASS prefill/varlen forward backend.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- 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.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cutlass_blackwell` - Platforms: NVIDIA SM100+ - Topics: FP16, BF16, varlen
Source: [mslk/attention/fmha/cutlass_blackwell.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass_blackwell.py)

### cutlass_blackwell.FwOpDecode

```python
cutlass_blackwell.FwOpDecode
```

Blackwell inference/decode forward backend.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Inference-only; Q heads per KV head are limited to 16.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cutlass_blackwell` - Platforms: NVIDIA SM100+ - Topics: BF16, decode, GQA
Source: [mslk/attention/fmha/cutlass_blackwell.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass_blackwell.py)

### cutlass_blackwell.BwOp

```python
cutlass_blackwell.BwOp
```

Blackwell CUTLASS backward backend.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- BMHK only; head dimensions 64 or 128.
- No dropout or different V dimension.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cutlass_blackwell` - Platforms: NVIDIA SM100+ - Topics: FP16, BF16, BMHK
Source: [mslk/attention/fmha/cutlass_blackwell.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cutlass_blackwell.py)

### flash.FwOp

```python
flash.FwOp
```

FlashAttention 2 forward backend for the general NVIDIA fast path.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- 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.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.flash` - Platforms: NVIDIA SM80+ - Topics: FP16, BF16, dropout, partial, varlen, paged, BMGHK
Source: [mslk/attention/fmha/flash.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py)

### flash.BwOp

```python
flash.BwOp
```

FlashAttention 2 backward backend.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Grouped-head BMGHK is forward-only.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.flash` - Platforms: NVIDIA SM80+ - Topics: FP16, BF16, dropout, varlen
Source: [mslk/attention/fmha/flash.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py)

### flash3.FwOp

```python
flash3.FwOp
```

FlashAttention 3 forward with FP8 and flexible V-dimension support.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Head dimension must be 64, 128, 192, or 256.
- No dropout.
- Ampere requires Q/K and V dimensions to match.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.flash3` - Platforms: NVIDIA SM80-SM90 - Topics: FP16, BF16, FP8 E4M3, partial, paged, gappy, BMGHK
Source: [mslk/attention/fmha/flash3.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py)

### flash3.BwOp

```python
flash3.BwOp
```

FlashAttention 3 backward backend.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Padded, gappy, paged, and BMGHK paths are forward-only.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.flash3` - Platforms: NVIDIA SM80-SM90 - Topics: FP16, BF16, FP8 E4M3
Source: [mslk/attention/fmha/flash3.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py)

### flash3.FwOp_KVSplit

```python
flash3.FwOp_KVSplit
```

FlashAttention 3 split-KV forward variant.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.flash3` - Platforms: NVIDIA SM80-SM90 - Topics: split KV, decode, FP8
Source: [mslk/attention/fmha/flash3.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py)

### ck.FwOp

```python
ck.FwOp
```

Composable Kernel attention forward for ROCm.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Maximum head dimension 256.
- PyTorch still reports the HIP device type as cuda.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.ck` - Platforms: AMD ROCm - Topics: FP16, BF16, dropout, partial, BMGHK, different V dim
Source: [mslk/attention/fmha/ck.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck.py)

### ck.BwOp

```python
ck.BwOp
```

Composable Kernel attention backward, including supported bias gradients.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.ck` - Platforms: AMD ROCm - Topics: FP16, BF16, dropout, bias grad
Source: [mslk/attention/fmha/ck.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck.py)

### ck_decoder.FwOp

```python
ck_decoder.FwOp
```

ROCm inference decoder attention backend.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Inference-only; fixed per-lane query length.
- Key padding is limited to 8,192.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.ck_decoder` - Platforms: AMD ROCm / MI250X - Topics: FP16, BF16, FP32, decode
Source: [mslk/attention/fmha/ck_decoder.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_decoder.py)

### ck_splitk.FwOp

```python
ck_splitk.FwOp
```

ROCm split-K decoder/prefix attention with optional quantized K/V.

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 / arguments:**
- `get_split_k(B, H, Mk) -> int` - Selects a split count from batch, heads, and KV length.

**Constraints:**
- No dropout or backward.
- Maximum head dimension 256; causal inputs have query-length restrictions.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: FP16, BF16, FP32, INT32 KV, split K
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### triton_splitk.FwOp

```python
triton_splitk.FwOp
```

Portable split-K forward with quantized KV-cache support.

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 / arguments:**
- `get_operator(splitk, *, block_m=None, block_n=None, num_warps=None, num_stages=None, split_k_early_exit=None)` - Builds/selects a specialized operator class.
- `get_split_k(B, G, H, Mk, Mq, page_size, is_paged=False) -> int` - Heuristic split selector.

**Constraints:**
- 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.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm / MI300X - Topics: FP16, BF16, FP8, INT4 KV, FP8 KV, partial, paged
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### cute_hopper.FwOp

```python
cute_hopper.FwOp
```

CuTe DSL Hopper forward attention.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Head dimension 64 or 128.
- No dropout, partial attention, or different V dimension.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cute_hopper` - Platforms: NVIDIA SM90-SM99 - Topics: FP16, BF16, CuTe, BMHK, BMGHK
Source: [mslk/attention/fmha/cute_hopper.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_hopper.py)

### cute_hopper.BwOp

```python
cute_hopper.BwOp
```

CuTe DSL Hopper backward attention.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cute_hopper` - Platforms: NVIDIA SM90-SM99 - Topics: FP16, BF16, CuTe
Source: [mslk/attention/fmha/cute_hopper.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_hopper.py)

### cute_blackwell.FwOp

```python
cute_blackwell.FwOp
```

CuTe DSL Blackwell prefill forward attention.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Prefill head dimensions: 64, 96, 128.
- No dropout, partial attention, or different V dimension.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cute_blackwell` - Platforms: NVIDIA SM100+ - Topics: FP16, BF16, FP8 E4M3, paged, CuTe
Source: [mslk/attention/fmha/cute_blackwell.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_blackwell.py)

### cute_blackwell.FwOpDecode

```python
cute_blackwell.FwOpDecode
```

CuTe DSL Blackwell decode/paged forward attention.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Decode head dimension 64 or 128.
- Paged-gappy support is limited to this decode path.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cute_blackwell` - Platforms: NVIDIA SM100+ - Topics: FP16, BF16, FP8 E4M3, decode, paged, gappy
Source: [mslk/attention/fmha/cute_blackwell.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_blackwell.py)

### cute_blackwell.BwOp

```python
cute_blackwell.BwOp
```

CuTe DSL Blackwell backward attention.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

**Constraints:**
- Backward head dimensions: 64, 96, 128.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.cute_blackwell` - Platforms: NVIDIA SM100+ - Topics: FP16, BF16, FP8 E4M3, CuTe
Source: [mslk/attention/fmha/cute_blackwell.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/cute_blackwell.py)

### flash_mtia.FwOp

```python
flash_mtia.FwOp
```

MTIA counterpart of the FlashAttention 2 forward protocol.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.flash_mtia` - Platforms: MTIA - Topics: FP16, BF16, dropout, partial
Source: [mslk/attention/fmha/flash_mtia.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py)

### flash_mtia.BwOp

```python
flash_mtia.BwOp
```

MTIA counterpart of the FlashAttention 2 backward protocol.

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 / arguments:**
- `is_available() -> bool` - Checks whether the compiled/JIT implementation is present.
- `supports(inputs) -> bool` - Returns whether the operator accepts a normalized Inputs object.
- `not_supported_reasons(inputs) -> list[str]` - Explains rejected dtype, layout, shape, mask, dropout, or architecture constraints.

Kind: backend - Stability: Advanced - Module: `mslk.attention.fmha.flash_mtia` - Platforms: MTIA - Topics: FP16, BF16, dropout
Source: [mslk/attention/fmha/flash_mtia.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py)

## Attention - Fixed split-K backends

### ck_splitk.FwOp_S1

```python
ck_splitk.FwOp_S1
```

CK split-K forward fixed to 1 split.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 1
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### ck_splitk.FwOp_S2

```python
ck_splitk.FwOp_S2
```

CK split-K forward fixed to 2 splits.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 2
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### ck_splitk.FwOp_S4

```python
ck_splitk.FwOp_S4
```

CK split-K forward fixed to 4 splits.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 4
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### ck_splitk.FwOp_S8

```python
ck_splitk.FwOp_S8
```

CK split-K forward fixed to 8 splits.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 8
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### ck_splitk.FwOp_S16

```python
ck_splitk.FwOp_S16
```

CK split-K forward fixed to 16 splits.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 16
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### ck_splitk.FwOp_S32

```python
ck_splitk.FwOp_S32
```

CK split-K forward fixed to 32 splits.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 32
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### ck_splitk.FwOp_S64

```python
ck_splitk.FwOp_S64
```

CK split-K forward fixed to 64 splits.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 64
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### ck_splitk.FwOp_S128

```python
ck_splitk.FwOp_S128
```

CK split-K forward fixed to 128 splits.

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 - Module: `mslk.attention.fmha.ck_splitk` - Platforms: AMD ROCm - Topics: CK, split 128
Source: [mslk/attention/fmha/ck_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/ck_splitk.py)

### triton_splitk.FwOp_S1

```python
triton_splitk.FwOp_S1
```

Triton split-K forward fixed to 1 split.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 1
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### triton_splitk.FwOp_S2

```python
triton_splitk.FwOp_S2
```

Triton split-K forward fixed to 2 splits.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 2
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### triton_splitk.FwOp_S4

```python
triton_splitk.FwOp_S4
```

Triton split-K forward fixed to 4 splits.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 4
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### triton_splitk.FwOp_S8

```python
triton_splitk.FwOp_S8
```

Triton split-K forward fixed to 8 splits.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 8
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### triton_splitk.FwOp_S16

```python
triton_splitk.FwOp_S16
```

Triton split-K forward fixed to 16 splits.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 16
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### triton_splitk.FwOp_S32

```python
triton_splitk.FwOp_S32
```

Triton split-K forward fixed to 32 splits.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 32
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### triton_splitk.FwOp_S64

```python
triton_splitk.FwOp_S64
```

Triton split-K forward fixed to 64 splits.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 64
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### triton_splitk.FwOp_S128

```python
triton_splitk.FwOp_S128
```

Triton split-K forward fixed to 128 splits.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: Triton, split 128
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

## Attention - Preassembled operator pairs

### MemoryEfficientAttentionCutlassOp

```python
MemoryEfficientAttentionCutlassOp = (cutlass.FwOp, cutlass.BwOp)
```

Classic CUTLASS forward and backward.

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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionCutlassBlackwellOp

```python
MemoryEfficientAttentionCutlassBlackwellOp = (cutlass_blackwell.FwOp, cutlass_blackwell.BwOp)
```

Blackwell CUTLASS forward and backward.

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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionCutlassFwdFlashBwOp

```python
MemoryEfficientAttentionCutlassFwdFlashBwOp = (cutlass.FwOp, flash.BwOp)
```

Classic CUTLASS forward with FlashAttention 2 backward.

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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionFlashAttentionOp

```python
MemoryEfficientAttentionFlashAttentionOp = (flash.FwOp, flash.BwOp)
```

FlashAttention 2 pair.

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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionFlashMtiaAttentionOp

```python
MemoryEfficientAttentionFlashMtiaAttentionOp = (flash_mtia.FwOp, flash_mtia.BwOp)
```

MTIA Flash pair.

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 - Module: `mslk.attention.fmha` - Platforms: MTIA - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionCkOp

```python
MemoryEfficientAttentionCkOp = (ck.FwOp, ck.BwOp)
```

ROCm CK pair.

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 - Module: `mslk.attention.fmha` - Platforms: AMD ROCm - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionCkDecoderOp

```python
MemoryEfficientAttentionCkDecoderOp = (ck_decoder.FwOp, ck.BwOp)
```

Inference-only convenience tuple using CK decoder forward.

A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.

**Returns:** AttentionOp tuple.

**Constraints:**
- The forward class produces no backward context; the paired ck.BwOp does not make this tuple valid for autograd training.

Kind: constant - Stability: Advanced - Module: `mslk.attention.fmha` - Platforms: AMD ROCm - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionSplitKCkOp

```python
MemoryEfficientAttentionSplitKCkOp = (ck_splitk.FwOp, ck.BwOp)
```

Inference-only convenience tuple using CK split-K forward.

A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.

**Returns:** AttentionOp tuple.

**Constraints:**
- The forward class produces no backward context; the paired ck.BwOp does not make this tuple valid for autograd training.

Kind: constant - Stability: Advanced - Module: `mslk.attention.fmha` - Platforms: AMD ROCm - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### MemoryEfficientAttentionCuteFlashAttentionOp

```python
MemoryEfficientAttentionCuteFlashAttentionOp = (cute_blackwell.FwOp, cute_blackwell.BwOp)
```

CuTe Blackwell forward/backward pair.

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 - Module: `mslk.attention.fmha` - Platforms: NVIDIA SM100+ - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### ALL_FW_OPS

```python
ALL_FW_OPS: list[type[AttentionFwOpBase]]
```

Explicit forward-backend enumeration used by tests and introspection.

A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.

**Returns:** AttentionOp tuple.

**Constraints:**
- This is not the automatic dispatch priority list.

Kind: constant - Stability: Advanced - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### ALL_BW_OPS

```python
ALL_BW_OPS: list[type[AttentionBwOpBase]]
```

Explicit backward-backend enumeration used by tests and introspection.

A ready-made (forward class, backward class) tuple for the op= argument. Availability still depends on platform and build.

**Returns:** AttentionOp tuple.

**Constraints:**
- This is not the automatic dispatch priority list.

Kind: constant - Stability: Advanced - Module: `mslk.attention.fmha` - Platforms: NVIDIA, AMD, MTIA - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/__init__.py)

### triton_splitk.FwOp_Map

```python
triton_splitk.FwOp_Map: dict[int, type[triton_splitk.FwOp]]
```

Maps supported split counts to specialized Triton split-K forward classes.

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 - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA SM80+, AMD ROCm - Topics: AttentionOp, explicit dispatch
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

## Attention - Operator protocol & inputs

### AttentionOp

```python
AttentionOp = tuple[Optional[type[AttentionFwOpBase]], Optional[type[AttentionBwOpBase]]]
```

Forward/backward operator pair type accepted by high-level fMHA.

Kind: type alias - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### AttentionOpBase

```python
AttentionOpBase
```

Shared backend metadata, availability, and support-check protocol.

**Methods / arguments:**
- `supports(inputs: Inputs) -> bool` - Whether the backend accepts an input descriptor.
- `not_supported_reasons(inputs: Inputs) -> list[str]` - Human-readable rejection reasons.
- `shape_not_supported_reasons(Mq, Mkv, K, Kv) -> list[str]` - Shape-only compatibility diagnostics.
- `is_available() -> bool` - Whether the implementation is loadable.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### AttentionFwOpBase

```python
AttentionFwOpBase(AttentionOpBase)
```

Base protocol for forward attention operators.

**Methods / arguments:**
- `apply(inp: Inputs, needs_gradient: bool) -> tuple[Tensor, Optional[Context]]` - Runs the backend and optionally returns backward context.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### AttentionBwOpBase

```python
AttentionBwOpBase(AttentionOpBase)
```

Base protocol for backward attention operators.

**Methods / arguments:**
- `apply(ctx: Context, inp: Inputs, grad: Tensor) -> Gradients` - Computes normalized Q/K/V/(bias) gradients.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### Inputs

```python
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)
```

Normalized fMHA input descriptor used by dispatch and backends.

Validates dtypes, layouts, shapes, masks, scale and dropout; can normalize BMK/BMHK/BMGHK logical layouts.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### InputsFp8

```python
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)
```

Input descriptor for row/group quantized FP8-style K/V caches.

Carries packed scale/shift metadata used by direct quantized split-K backends.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### InputsMXFp8

```python
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)
```

Input descriptor for MXFP8 K/V with E8M0 block scales.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### Context

```python
Context(lse: Tensor, out: Tensor, op_bw: Optional[type[AttentionBwOpBase]]=None, rng_state=None, qkv_share_storage: bool=False)
```

Forward context consumed by an attention backward operator.

**Methods / arguments:**
- `get_padded_lse(pad_to: int, force_pad_inf: bool=False) -> Tensor` - Returns LSE padded for a backend's alignment contract.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### Gradients

```python
Gradients(dq: Tensor, dk: Tensor, dv: Tensor, db: Optional[Tensor]=None)
```

Normalized attention-gradient bundle.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### ScaledTensor

```python
ScaledTensor(data: Tensor, scale: Tensor, dequant_func: Callable[[Tensor, Tensor], Tensor], original_dtype: torch.dtype, require_grad: bool=False)
```

Tensor subclass carrying separate quantization-scale metadata.

Kind: class - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### pack_fp8_tensorwise_per_head

```python
pack_fp8_tensorwise_per_head(x: Tensor, scale: Tensor | float, original_dtype) -> ScaledTensor
```

Wraps tensorwise/per-head FP8 storage and scale metadata for backend input.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### bmghk2bmhk

```python
bmghk2bmhk(x: Tensor, x_scale: Optional[Tensor]=None, handle_rep_heads: bool=False) -> tuple[Tensor, Optional[Tensor]]
```

Flattens grouped heads G x H into a single H axis, with optional scale handling.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### bmhk2bhk

```python
bmhk2bhk(x: Tensor, x_scale: Optional[Tensor]=None, handle_mqa: bool=True) -> tuple[Tensor, Optional[Tensor]]
```

Converts BMHK storage to the backend BHK-style view for single-query/decode paths.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### bmk2bmhk

```python
bmk2bmhk(tensor: Tensor, num_heads: int) -> Tensor
```

Adds/reshapes a head axis for BMK tensors.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

### check_lastdim_alignment_stride1

```python
check_lastdim_alignment_stride1(reasons, name, x, alignment) -> None
```

Appends backend support failures for last-dimension alignment and stride.

Kind: function - Stability: Developer - Module: `mslk.attention.fmha.common` - Platforms: NVIDIA, AMD, MTIA, CPU / Python - Topics: dispatch, introspection, normalized inputs
Source: [mslk/attention/fmha/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/common.py)

## Attention - Dispatch controls

### _get_use_fa3

```python
_get_use_fa3() -> bool
```

Returns whether CUDA automatic dispatch is permitted to select FlashAttention 3.

**Constraints:**
- Underscore-prefixed but intentionally included in fmha.__all__.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.dispatch` - Platforms: NVIDIA, CPU / Python - Topics: FlashAttention 3, dispatch, compatibility
Source: [mslk/attention/fmha/dispatch.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/dispatch.py)

### _set_use_fa3

```python
_set_use_fa3(use_flash_attention3: bool) -> None
```

Enables or disables FlashAttention 3 in automatic dispatch.

**Constraints:**
- Process-global dispatch control.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.dispatch` - Platforms: NVIDIA, CPU / Python - Topics: FlashAttention 3, dispatch, compatibility
Source: [mslk/attention/fmha/dispatch.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/dispatch.py)

### fa3_available

```python
fa3_available() -> bool
```

Reports whether a compatible FlashAttention 3 implementation is available.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.dispatch` - Platforms: NVIDIA, CPU / Python - Topics: FlashAttention 3, dispatch, compatibility
Source: [mslk/attention/fmha/dispatch.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/dispatch.py)

### is_pt_cutlass_compatible

```python
is_pt_cutlass_compatible(force=False) -> bool
```

Checks PyTorch's bundled CUTLASS attention compatibility.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.torch_attention_compat` - Platforms: NVIDIA, CPU / Python - Topics: FlashAttention 3, dispatch, compatibility
Source: [mslk/attention/fmha/torch_attention_compat.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/torch_attention_compat.py)

### ensure_pt_flash_ok

```python
ensure_pt_flash_ok() -> None
```

Raises when the active PyTorch FlashAttention integration is incompatible.

Kind: function - Stability: Advanced - Module: `mslk.attention.fmha.torch_attention_compat` - Platforms: NVIDIA, CPU / Python - Topics: FlashAttention 3, dispatch, compatibility
Source: [mslk/attention/fmha/torch_attention_compat.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/torch_attention_compat.py)

## Attention - Tree & speculative attention

### TreeAttnMetadata

```python
TreeAttnMetadata.from_tree_choices(tree_choices, dtype=None, device=None) -> TreeAttnMetadata
```

Builds explicit mask and traversal metadata for a speculative token tree.

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 / arguments:**
- `from_tree_choices(tree_choices, dtype=None, device=None)` - Builds metadata from node paths sorted by depth.
- `from_tree_choices_cached(tree_choices: tuple[tuple[int,...],...], dtype=None, device=None)` - LRU-cached constructor for hashable choices.

Kind: class - Stability: Public - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

### tree_attention

```python
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
```

Computes prefix-cache and speculative-tree attention, then LSE-merges them.

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.

**Constraints:**
- Supports paged prefix masks.
- Quantized int32 KV uses Triton split-K; FP8 Q/cache paths accept explicit scales.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

### use_triton_splitk_for_prefix

```python
use_triton_splitk_for_prefix(B: int, G: int, tree_size: int) -> bool
```

Heuristic for choosing Triton split-K on the tree prefix.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

### select_prefix_op

```python
select_prefix_op(B, G, tree_size, autotune, attn_bias, kv_cache_dtype) -> Optional[type[AttentionFwOpBase]]
```

Selects the prefix forward operator for tree attention.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

### SplitKAutotune

```python
SplitKAutotune(triton_splitk.FwOp)
```

Tree-attention split-K operator variant with runtime autotuning.

Kind: backend - Stability: Experimental - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

### construct_full_tree_choices

```python
construct_full_tree_choices(tree_depth: int, branching: int) -> list[tuple[int, ...]]
```

Generates node choices for a regular full tree.

**Constraints:**
- Its depth convention counts non-root levels; get_full_tree_size uses a different range convention.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

### construct_tree_choices

```python
construct_tree_choices(branching: list[int]) -> list[tuple[int, ...]]
```

Generates node choices for per-level branching factors.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

### get_full_tree_size

```python
get_full_tree_size(tree_depth: int, branching: int) -> int
```

Returns sum of branching^i over range(tree_depth).

**Constraints:**
- The tree_depth convention does not directly match construct_full_tree_choices.

Kind: function - Stability: Public - Module: `mslk.attention.fmha.tree_attention` - Platforms: NVIDIA, AMD - Topics: Medusa, EAGLE, Hydra, speculative decoding, tree
Source: [mslk/attention/fmha/tree_attention.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/tree_attention.py)

## Attention - iRoPE & page transforms

### split_blocks_for_decoding_gpu_part

```python
split_blocks_for_decoding_gpu_part(input_bias, batchify_len, block_tables=None, page_size=None) -> Optional[tuple[Tensor, Tensor]]
```

Computes GPU-side sequence-start/length data for decoding block splits.

Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.

Kind: function - Stability: Experimental - Module: `mslk.attention.fmha.split_blocks_fairinternal` - Platforms: NVIDIA, AMD - Topics: iRoPE, paged KV, gappy keys, advanced
Source: [mslk/attention/fmha/split_blocks_fairinternal.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py)

### split_blocks_for_decoding

```python
split_blocks_for_decoding(input_bias, batchify_len, block_tables=None, page_size=None, gpu_data=None) -> Optional[BlockDiagonalGappyKeysMask | PagedBlockDiagonalGappyKeysMask]
```

Transforms padded decoding lanes into gappy or paged-gappy masks.

Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.

Kind: function - Stability: Experimental - Module: `mslk.attention.fmha.split_blocks_fairinternal` - Platforms: NVIDIA, AMD - Topics: iRoPE, paged KV, gappy keys, advanced
Source: [mslk/attention/fmha/split_blocks_fairinternal.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py)

### split_blocks_for_prefill

```python
split_blocks_for_prefill(input_bias, batchify_len) -> Optional[BlockDiagonalPaddedKeysMask]
```

Re-batches padded prefill lanes into a split block description.

Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.

Kind: function - Stability: Experimental - Module: `mslk.attention.fmha.split_blocks_fairinternal` - Platforms: NVIDIA, AMD - Topics: iRoPE, paged KV, gappy keys, advanced
Source: [mslk/attention/fmha/split_blocks_fairinternal.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py)

### maybe_make_paged

```python
maybe_make_paged(attn_bias, block_tables, page_size, notional_padding) -> Optional[AttentionBias]
```

Converts supported padded/gappy masks to paged counterparts when a block table is supplied.

Advanced transformations that convert padded iRoPE lanes into chunked, gappy, or paged bias descriptions.

Kind: function - Stability: Experimental - Module: `mslk.attention.fmha.split_blocks_fairinternal` - Platforms: NVIDIA, AMD - Topics: iRoPE, paged KV, gappy keys, advanced
Source: [mslk/attention/fmha/split_blocks_fairinternal.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/split_blocks_fairinternal.py)

## Attention - ROCm MLA

### mla_decode_fwd

```python
mla_decode_fwd(query, kv_buffer, block_tables, cu_seqlens_q, seqused_k, softmax_scale=None) -> Tensor
```

Paged MLA decode forward for one query position per sequence.

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

**Constraints:**
- Query shape is [B,128,576].
- Input may be BF16 or compatible ROCm FP8 FNUZ.

Kind: function - Stability: Public - Module: `mslk.attention.mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

### mla_prefill_fwd

```python
mla_prefill_fwd(query, kv_buffer, block_tables, cu_seqlens_q, seqused_k, softmax_scale=None) -> Tensor
```

Packed variable-length MLA prefill forward.

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

**Constraints:**
- Query shape is [total_tokens,128,576].

Kind: function - Stability: Public - Module: `mslk.attention.mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

### MLA_NUM_HEADS

```python
MLA_NUM_HEADS = 128
```

Fixed MLA query-head count.

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 - Module: `mslk.attention.mla.triton_mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

### MLA_NUM_KV_HEADS

```python
MLA_NUM_KV_HEADS = 1
```

Fixed latent KV-head count.

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 - Module: `mslk.attention.mla.triton_mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

### MLA_KV_LORA_RANK

```python
MLA_KV_LORA_RANK = 512
```

Latent KV LoRA rank.

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 - Module: `mslk.attention.mla.triton_mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

### MLA_QK_ROPE_HEAD_DIM

```python
MLA_QK_ROPE_HEAD_DIM = 64
```

RoPE portion of the Q/K head.

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 - Module: `mslk.attention.mla.triton_mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

### MLA_QK_HEAD_DIM

```python
MLA_QK_HEAD_DIM = 576
```

Full Q/K head dimension after absorption.

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 - Module: `mslk.attention.mla.triton_mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

### MLA_V_HEAD_DIM

```python
MLA_V_HEAD_DIM = 512
```

MLA value/output head dimension.

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 - Module: `mslk.attention.mla.triton_mla` - Platforms: AMD gfx942, AMD gfx950 - Topics: MLA, DeepSeek V3, paged KV, Triton, BF16, FP8 FNUZ
Source: [mslk/attention/mla/triton_mla.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/mla/triton_mla.py)

## Attention - FlyDSL Flash Attention

### flydsl_flash_attn_func

```python
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)
```

ROCm FlyDSL flash-attention forward across dense, varlen, GQA/MQA, and paged KV.

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.

**Constraints:**
- 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 - Module: `mslk.attention.flydsl` - Platforms: AMD gfx942, AMD gfx950 - Topics: FlyDSL, dense, varlen, paged, GQA, MQA, split K
Source: [mslk/attention/flydsl/flash_attn_interface.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/flydsl/flash_attn_interface.py)

### dualwave_splitk_workspace_elems

```python
dualwave_splitk_workspace_elems(batch_size, num_heads, seq_len, num_kv_splits, head_dim=128) -> int
```

Returns the FP32 element count needed by the gfx950 dual-wave split-K workspace.

Kind: function - Stability: Public - Module: `mslk.attention.flydsl.flash_attn_interface` - Platforms: AMD gfx942, AMD gfx950 - Topics: FlyDSL, dense, varlen, paged, GQA, MQA, split K
Source: [mslk/attention/flydsl/flash_attn_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/flydsl/flash_attn_utils.py)

## Attention - Standalone facades

### flash_attn_func

```python
flash_attn_func(...) -> dependency-defined
```

Optional standalone CuTe Flash-Attention facade.

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.

**Constraints:**
- 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 - Module: `mslk.attention.flash_attn` - Platforms: NVIDIA / dependency-defined - Topics: CuTe, optional dependency, standalone
Source: [mslk/attention/flash_attn/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/flash_attn/__init__.py)

## Attention - Standalone Blackwell FMHA

### cutlass_blackwell_fmha_func

```python
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)
```

Autograd-capable standalone Blackwell FMHA facade.

Handles dense or variable-length inputs and routes sequence length sq=1 through the generated decode path.

**Returns:** Attention output tensor.

**Constraints:**
- The generation path leaves its split dimension unmerged and has no backward implementation.

Kind: function - Stability: Public - Module: `mslk.attention.cutlass_blackwell_fmha` - Platforms: NVIDIA SM100+ - Topics: CUTLASS, Blackwell, varlen, paged, decode
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py)

### cutlass_blackwell_fmha_decode_forward

```python
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]
```

Inference-only split-K decode forward with raw partial outputs.

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 - Module: `mslk.attention.cutlass_blackwell_fmha` - Platforms: NVIDIA SM100+ - Topics: CUTLASS, Blackwell, varlen, paged, decode
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py)

### _cutlass_blackwell_fmha_forward

```python
_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]
```

Testing/implementation forward wrapper returning output and LSE.

Kind: function - Stability: Low-level - Module: `mslk.attention.cutlass_blackwell_fmha` - Platforms: NVIDIA SM100+ - Topics: CUTLASS, Blackwell, varlen, paged, decode
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py)

### cutlass_blackwell_fmha_custom_op

```python
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)
```

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 - Module: `mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_custom_op` - Platforms: NVIDIA SM100+ - Topics: CUTLASS, Blackwell, varlen, paged, decode
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py)

### get_splitk_heuristic

```python
get_splitk_heuristic(batch: int, seqlen_kv: int, kv_heads: int=1, tile_n: int=256, sm_count: int | None=None) -> int
```

Chooses a decode split size from cache length, KV heads, and available SMs.

**Constraints:**
- Returns 0 to disable split-K when a single split would cover the sequence.

Kind: function - Stability: Public - Module: `mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_interface` - Platforms: NVIDIA SM100+ - Topics: CUTLASS, Blackwell, varlen, paged, decode
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py)

### GenKernelType

```python
GenKernelType.UMMA_I | GenKernelType.UMMA_P
```

Selects the generated Blackwell UMMA kernel family.

Kind: enum - Stability: Low-level - Module: `mslk.attention.cutlass_blackwell_fmha.cutlass_blackwell_fmha_interface` - Platforms: NVIDIA SM100+ - Topics: CUTLASS, Blackwell, varlen, paged, decode
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py)

## Attention - Raw Blackwell torch ops

### torch.ops.mslk.fmha_fwd

```python
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]
```

Standalone Blackwell forward implementation op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+, Meta / fake - Topics: native op, CUTLASS, implementation
Source: [csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd.cu](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd.cu)

### torch.ops.mslk.fmha_bwd

```python
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]
```

Standalone Blackwell backward implementation op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+, Meta / fake - Topics: native op, CUTLASS, implementation
Source: [csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_bwd.cu](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_bwd.cu)

### torch.ops.mslk.fmha_gen_fwd

```python
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]
```

Generated UMMA Blackwell forward implementation op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+, Meta / fake - Topics: native op, CUTLASS, implementation
Source: [csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_gen_impl.cu](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/cuda/cutlass_blackwell_fmha/blackwell_gen_impl.cu)

### torch.ops.mslk.cutlass_blackwell_fmha_fwd

```python
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]
```

fMHA Blackwell forward native op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+, Meta / fake - Topics: native op, CUTLASS, implementation
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py)

### torch.ops.mslk.cutlass_blackwell_fmha_bwd

```python
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]
```

fMHA Blackwell backward native op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+, Meta / fake - Topics: native op, CUTLASS, implementation
Source: [mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_custom_op.py)

## Attention - Merge & storage utilities

### triton_splitk.merge_attentions

```python
triton_splitk.merge_attentions(attn_out, lse_out, attn_split, lse_split) -> None
```

In-place low-level reduction of split attention into preallocated outputs.

This is the allocation-free implementation layer beneath the high-level merge.

Kind: function - Stability: Low-level - Module: `mslk.attention.fmha.triton_splitk` - Platforms: NVIDIA, AMD - Topics: merge, storage views, autograd
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### torch.ops.mslk.fmha_merge_attentions_varargs

```python
torch.ops.mslk.fmha_merge_attentions_varargs(attn_split: Tensor[], lse_split: Tensor[], write_lse, output_dtype, B, M, G, H, Kq) -> Tensor[]
```

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: merge, storage views, autograd
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### torch.ops.mslk.merge_attentions_varargs_backward

```python
torch.ops.mslk.merge_attentions_varargs_backward(attn_split, lse_split, attn_out, lse_out, grad_attn, grad_lse) -> tuple[Tensor[], Tensor[]]
```

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: merge, storage views, autograd
Source: [mslk/attention/fmha/triton_splitk.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/triton_splitk.py)

### get_stack_strides

```python
get_stack_strides(tensors: Sequence[Tensor], dim: int) -> Optional[tuple[int, ...]]
```

Detects whether tensors are views of one common stacked storage layout.

Kind: function - Stability: Low-level - Module: `mslk.attention.fmha.unbind` - Platforms: NVIDIA, AMD - Topics: merge, storage views, autograd
Source: [mslk/attention/fmha/unbind.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/unbind.py)

### unbind

```python
unbind(x: Tensor, dim: int) -> tuple[Tensor, ...]
```

Autograd-aware unbind preserving shared-storage information.

Kind: function - Stability: Low-level - Module: `mslk.attention.fmha.unbind` - Platforms: NVIDIA, AMD - Topics: merge, storage views, autograd
Source: [mslk/attention/fmha/unbind.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/unbind.py)

### stack_or_none

```python
stack_or_none(tensors: Sequence[Tensor], dim: int) -> Optional[Tensor]
```

Returns a zero-copy/common-storage stack view when possible, else None.

Kind: function - Stability: Low-level - Module: `mslk.attention.fmha.unbind` - Platforms: NVIDIA, AMD - Topics: merge, storage views, autograd
Source: [mslk/attention/fmha/unbind.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/unbind.py)

## Attention - Raw backend ops

### torch.ops.xformers.efficient_attention_forward_ck

```python
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]
```

Raw CK forward attention op.

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 - Module: `torch.ops` - Platforms: AMD ROCm - Topics: raw op, backend implementation, unstable
Source: [csrc/attention/ck/fmha/hip_fmha/attention_forward_generic_ck_tiled.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_fmha/attention_forward_generic_ck_tiled.cpp)

### torch.ops.xformers.efficient_attention_backward_ck

```python
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]
```

Raw CK backward attention op.

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 - Module: `torch.ops` - Platforms: AMD ROCm - Topics: raw op, backend implementation, unstable
Source: [csrc/attention/ck/fmha/hip_fmha/attention_backward_generic_ck_tiled.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_fmha/attention_backward_generic_ck_tiled.cpp)

### torch.ops.xformers.efficient_attention_forward_decoder_ck

```python
torch.ops.xformers.efficient_attention_forward_decoder_ck(query, key, value, seq_positions, scale) -> Tensor
```

Raw CK decoder forward op.

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 - Module: `torch.ops` - Platforms: AMD ROCm - Topics: raw op, backend implementation, unstable
Source: [csrc/attention/ck/fmha/hip_decoder/attention_forward_decoder.hip](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_decoder/attention_forward_decoder.hip)

### torch.ops.xformers.efficient_attention_forward_decoder_splitk_ck

```python
torch.ops.xformers.efficient_attention_forward_decoder_splitk_ck(query, key, value, seq_positions, scale, split_k) -> Tensor
```

Raw CK split-K decoder forward op.

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 - Module: `torch.ops` - Platforms: AMD ROCm - Topics: raw op, backend implementation, unstable
Source: [csrc/attention/ck/fmha/hip_decoder/attention_forward_splitk.hip](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_decoder/attention_forward_splitk.hip)

### torch.ops.xformers._ck_rand_uniform

```python
torch.ops.xformers._ck_rand_uniform(p: float, out: Tensor) -> Tensor
```

CK random-uniform helper used by dropout implementation.

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 - Module: `torch.ops` - Platforms: AMD ROCm - Topics: raw op, backend implementation, unstable
Source: [csrc/attention/ck/fmha/hip_fmha/attention_ck_rand_uniform.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/attention/ck/fmha/hip_fmha/attention_ck_rand_uniform.cpp)

### torch.ops.mslk_flash.flash_fwd

```python
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]
```

Conditional bundled FlashAttention 2 forward op returning output, LSE, and RNG state.

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 - Module: `torch.ops` - Platforms: NVIDIA SM80+ - Topics: raw op, backend implementation, unstable
Source: [mslk/attention/fmha/flash.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py)

### torch.ops.mslk_flash.flash_bwd

```python
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]
```

Conditional bundled FlashAttention 2 backward op.

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 - Module: `torch.ops` - Platforms: NVIDIA SM80+ - Topics: raw op, backend implementation, unstable
Source: [mslk/attention/fmha/flash.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash.py)

### torch.ops.mslk_flash3.flash_fwd

```python
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]
```

Conditional bundled FlashAttention 3 forward op returning output and LSE.

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 - Module: `torch.ops` - Platforms: NVIDIA SM80-SM90 - Topics: raw op, backend implementation, unstable
Source: [mslk/attention/fmha/flash3.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py)

### torch.ops.mslk_flash3.flash_bwd

```python
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]
```

Conditional bundled FlashAttention 3 backward op.

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 - Module: `torch.ops` - Platforms: NVIDIA SM80-SM90 - Topics: raw op, backend implementation, unstable
Source: [mslk/attention/fmha/flash3.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash3.py)

### torch.ops.mslk_flash_mtia.flash_fwd

```python
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]
```

Conditional MTIA Flash forward op returning output, LSE, and RNG state.

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 - Module: `torch.ops` - Platforms: MTIA - Topics: raw op, backend implementation, unstable
Source: [mslk/attention/fmha/flash_mtia.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py)

### torch.ops.mslk_flash_mtia.flash_bwd

```python
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]
```

Conditional MTIA Flash backward op.

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 - Module: `torch.ops` - Platforms: MTIA - Topics: raw op, backend implementation, unstable
Source: [mslk/attention/fmha/flash_mtia.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/attention/fmha/flash_mtia.py)

## GEMM - BF16 grouped torch ops

### torch.ops.mslk.bf16bf16bf16_grouped

```python
torch.ops.mslk.bf16bf16bf16_grouped(X: Tensor[], W: Tensor[]) -> Tensor[]
```

List-of-tensors BF16 grouped GEMM.

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.

**Constraints:**
- Legacy TensorList path; deprecated/unsupported by the ROCm implementation. Prefer grouped_stacked.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: BF16, grouped GEMM, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16bf16bf16_grouped_cat

```python
torch.ops.mslk.bf16bf16bf16_grouped_cat(X: Tensor[], W: Tensor[]) -> Tensor
```

List-of-tensors BF16 grouped GEMM with concatenated output.

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

**Constraints:**
- Legacy path; prefer grouped_stacked.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: BF16, grouped GEMM, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16bf16bf16_grouped_dynamic

```python
torch.ops.mslk.bf16bf16bf16_grouped_dynamic(X: Tensor, W: Tensor, zero_start_index_M: Tensor) -> Tensor
```

Dynamic-row BF16 grouped GEMM using per-group zero starts.

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

**Constraints:**
- Legacy dynamic path; deprecated/unsupported by the ROCm implementation.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: BF16, grouped GEMM, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16bf16bf16_grouped_stacked

```python
torch.ops.mslk.bf16bf16bf16_grouped_stacked(X, W, M_sizes, out=None, num_sms=None) -> Tensor
```

Preferred stacked BF16 grouped forward GEMM.

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.

**Constraints:**
- M_sizes gives the valid row count for each W[g].
- num_sms optionally limits participating SMs.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: BF16, grouped GEMM, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16bf16bf16_grouped_grad

```python
torch.ops.mslk.bf16bf16bf16_grouped_grad(X, W, M_sizes, out=None, num_sms=None) -> Tensor
```

Grouped BF16 data-gradient matrix product.

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

**Constraints:**
- CUDA uses CUTLASS; ROCm registers the Triton implementation when mslk.gemm is imported.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: BF16, grouped GEMM, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16bf16bf16_grouped_wgrad

```python
torch.ops.mslk.bf16bf16bf16_grouped_wgrad(X, W, M_sizes, output=None, output_accum=False, num_sms=None) -> Tensor
```

Grouped BF16 weight-gradient product.

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.

**Constraints:**
- output_accum=True requires a preallocated FP32 output tensor.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: BF16, grouped GEMM, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

## GEMM - FP8 row/block/group torch ops

### torch.ops.mslk.f8f8bf16_blockwise

```python
torch.ops.mslk.f8f8bf16_blockwise(XQ, WQ, x_scale, w_scale, block_m=128, block_n=128, block_k=128) -> Tensor
```

Block-scaled FP8 x FP8 GEMM.

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

**Constraints:**
- Scale grids correspond to the declared M/N/K block sizes.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise

```python
torch.ops.mslk.f8f8bf16_rowwise(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor
```

Primary rowwise FP8 x FP8 -> BF16 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_out

```python
torch.ops.mslk.f8f8bf16_rowwise_out(XQ, WQ, x_scale, w_scale, output, bias=None, use_fast_accum=True) -> None
```

In-place/out rowwise FP8 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_batched

```python
torch.ops.mslk.f8f8bf16_rowwise_batched(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True, output=None) -> Tensor
```

Batched rowwise FP8 GEMM with optional preallocated output.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_grouped

```python
torch.ops.mslk.f8f8bf16_rowwise_grouped(XQ: Tensor[], WQ: Tensor[], x_scale: Tensor[], w_scale: Tensor[]) -> Tensor[]
```

TensorList grouped rowwise FP8 GEMM.

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

**Constraints:**
- Legacy/deprecated on ROCm; prefer grouped_stacked or grouped_mm.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_grouped_cat

```python
torch.ops.mslk.f8f8bf16_rowwise_grouped_cat(XQ: Tensor[], WQ: Tensor[], x_scale: Tensor[], w_scale: Tensor[]) -> Tensor
```

TensorList grouped rowwise FP8 GEMM with concatenated output.

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

**Constraints:**
- Legacy/deprecated on ROCm.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_grouped_stacked

```python
torch.ops.mslk.f8f8bf16_rowwise_grouped_stacked(XQ, WQ, x_scale, w_scale, M_sizes) -> Tensor
```

Stacked grouped rowwise FP8 GEMM.

Uses concatenated activation rows and a group/expert axis on weights, with M_sizes selecting each valid segment.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_grouped_dynamic

```python
torch.ops.mslk.f8f8bf16_rowwise_grouped_dynamic(XQ, WQ, x_scale, w_scale, zero_start_index_M, zeroing_output_tensor=True) -> Tensor
```

Dynamic-row grouped FP8 GEMM.

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

**Constraints:**
- Legacy/deprecated on ROCm.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_groupwise

```python
torch.ops.mslk.f8f8bf16_groupwise(XQ, WQ, x_scale, w_scale) -> Tensor
```

FP8 GEMM with fixed K-group scale granularity 128.

XQ [M,K], WQ [N,K], x_scale [K/128,M], and w_scale [K/128,N/128].

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_groupwise_grouped

```python
torch.ops.mslk.f8f8bf16_groupwise_grouped(XQ, WQ, x_scale, w_scale, M_sizes) -> Tensor
```

Stacked grouped FP8 GEMM with fixed K-group scales.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: FP8, BF16 output, reciprocal scale, X @ W.T
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

## GEMM - ROCm FP8 torch ops

### torch.ops.mslk.f8f8f16_rowwise

```python
torch.ops.mslk.f8f8f16_rowwise(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor
```

ROCm rowwise FP8 GEMM with FP16 output.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: AMD ROCm - Topics: FP8, ROCm, CK, preshuffle
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_preshuffle

```python
torch.ops.mslk.f8f8bf16_rowwise_preshuffle(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor
```

ROCm preshuffled rowwise FP8 GEMM with BF16 output.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: AMD ROCm - Topics: FP8, ROCm, CK, preshuffle
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8f16_rowwise_preshuffle

```python
torch.ops.mslk.f8f8f16_rowwise_preshuffle(XQ, WQ, x_scale, w_scale, bias=None, use_fast_accum=True) -> Tensor
```

ROCm preshuffled rowwise FP8 GEMM with FP16 output.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: AMD ROCm - Topics: FP8, ROCm, CK, preshuffle
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8f8bf16_rowwise_grouped_mm

```python
torch.ops.mslk.f8f8bf16_rowwise_grouped_mm(XQ, WQ, x_scale, w_scale, offsets=None, output=None) -> Tensor
```

Generic ROCm grouped FP8 GEMM over 2D/3D layout combinations.

Supports 2D x 3D expert grouping, 3D x 2D output grouping, batched 3D x 3D without offsets, and K-grouped 2D x 2D with offsets.

**Returns:** The mutated output tensor.

**Constraints:**
- Scales are FP32; N generally must be divisible by 8.
- offsets must be omitted for the 3D x 3D layout.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: AMD ROCm - Topics: FP8, ROCm, CK, preshuffle
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

## GEMM - FP4 & microscaling torch ops

### torch.ops.mslk.f4f4bf16

```python
torch.ops.mslk.f4f4bf16(XQ, WQ, x_scale, w_scale, output=None, global_scale=None, mxfp4_block_size=32) -> Tensor
```

Packed FP4 x FP4 GEMM selecting MXFP4, MXFP4-16, or NVFP4 mode.

global_scale selects NVFP4. Without it, block size 32 selects standard MXFP4 and 16 selects MXFP4-16.

**Constraints:**
- For NVFP4 GEMM, global_scale is reciprocal(a_global_scale x 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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD gfx950 - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f4f4bf16_grouped_mm

```python
torch.ops.mslk.f4f4bf16_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None, global_scale=None) -> Tensor
```

Offset-described grouped FP4 GEMM.

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.

**Constraints:**
- On ROCm gfx950 this supports only 2D XQ x 3D WQ standard MXFP4; global_scale is accepted for schema compatibility but ignored.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD gfx950 - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f4f4bf16_grouped_stacked

```python
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
```

Stacked grouped FP4 GEMM with optional per-segment padding metadata.

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.

**Constraints:**
- 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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD gfx950 - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f4f4bf16_ultra_grouped_mm

```python
torch.ops.mslk.f4f4bf16_ultra_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, x_global_scale, w_global_scale, output=None) -> Tensor
```

Ultra grouped FP4 GEMM with separate activation/weight global scales.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM10.3+ / CUDA 13+ - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.mx8mx4bf16

```python
torch.ops.mslk.mx8mx4bf16(XQ, WQ, x_scale, w_scale, output=None) -> Tensor
```

MXFP8 activation x MXFP4 weight GEMM.

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.

**Constraints:**
- On ROCm, activation scales are blocked/swizzled while weight scales remain plain [N,K/32].

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD gfx950 - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.mx8mx4bf16_grouped_mm

```python
torch.ops.mslk.mx8mx4bf16_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None) -> Tensor
```

Grouped MXFP8 x MXFP4 GEMM described by offsets.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD gfx950 - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.mx8mx8bf16_grouped_mm

```python
torch.ops.mslk.mx8mx8bf16_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None, actual_num_tokens=None) -> Tensor
```

Grouped MXFP8 x MXFP8 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD gfx950 - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.mx8mx6bf16

```python
torch.ops.mslk.mx8mx6bf16(XQ, WQ, x_scale, w_scale, output=None) -> Tensor
```

MXFP8 x packed MXFP6 E2M3 GEMM.

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.

**Constraints:**
- Pack quantize_bf16_to_mx6_e2m3 output with pack_fp6_e2m3 first.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+ - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.mx6mx6bf16

```python
torch.ops.mslk.mx6mx6bf16(XQ, WQ, x_scale, w_scale, output=None, splits=0) -> Tensor
```

Packed MXFP6 E2M3 x MXFP6 E2M3 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+ - Topics: FP4, MXFP4, MXFP8, MXFP6, BF16 output, packed
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

## GEMM - INT and mixed torch ops

### torch.ops.mslk.i8i8bf16

```python
torch.ops.mslk.i8i8bf16(XQ, WQ, scale: float, split_k=1) -> Tensor
```

INT8 x INT8 GEMM with a static scalar scale.

Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].

**Constraints:**
- 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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.i8i8bf16_dynamic

```python
torch.ops.mslk.i8i8bf16_dynamic(XQ, WQ, scale: Tensor, split_k=1) -> Tensor
```

INT8 x INT8 GEMM with a tensor-valued dynamic scale.

Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].

**Constraints:**
- ROCm registration requires importing mslk.gemm.triton.int8_gemm.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16i4bf16_rowwise

```python
torch.ops.mslk.bf16i4bf16_rowwise(X, W, w_scale_group, w_zero_group) -> Tensor
```

BF16 activation x row/group-quantized INT4 weight GEMM.

Integer and mixed-precision matrix products. INT4 weights generally use [N,K/2] packed storage and produce BF16 [M,N].

**Constraints:**
- ROCm registration requires importing mslk.gemm.triton.int4_gemm.

Kind: torch op - Stability: Public - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16i4bf16_rowwise_batched

```python
torch.ops.mslk.bf16i4bf16_rowwise_batched(X, WQ, w_scale, w_zp) -> Tensor
```

Batched BF16 x rowwise INT4 weight GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16i4bf16_shuffled

```python
torch.ops.mslk.bf16i4bf16_shuffled(X, W, w_scale_group, w_zero_group) -> Tensor
```

BF16 x preshuffled INT4 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM90 - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16i4bf16_shuffled_batched

```python
torch.ops.mslk.bf16i4bf16_shuffled_batched(X, WQ, w_scale, w_zp) -> Tensor
```

Batched BF16 x preshuffled INT4 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM90 - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16i4bf16_shuffled_grouped

```python
torch.ops.mslk.bf16i4bf16_shuffled_grouped(X, WQ, w_scale_group, w_zero_group, M_sizes) -> Tensor
```

Stacked grouped BF16 x preshuffled INT4 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM90 - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8i4bf16_rowwise

```python
torch.ops.mslk.f8i4bf16_rowwise(XQ, WQ, x_scale, w_scale, w_zp) -> Tensor
```

Rowwise FP8 activation x INT4 weight GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM90 - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8i4bf16_shuffled

```python
torch.ops.mslk.f8i4bf16_shuffled(XQ, WQ, x_scale, w_scale, w_scale_group) -> Tensor
```

FP8 activation x preshuffled INT4 weight GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM90 - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.f8i4bf16_shuffled_grouped

```python
torch.ops.mslk.f8i4bf16_shuffled_grouped(XQ, WQ, x_scale, w_scale, w_scale_group, M_sizes) -> Tensor
```

Stacked grouped FP8 x preshuffled INT4 GEMM.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM90 - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.preshuffle_i4

```python
torch.ops.mslk.preshuffle_i4(WQ, w_scale) -> tuple[Tensor, Tensor]
```

Preprocesses packed INT4 weights and scales for shuffled CUDA kernels.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

### torch.ops.mslk.bf16x9_gemm

```python
torch.ops.mslk.bf16x9_gemm(A, B, output=None) -> Tensor
```

cuBLAS BF16x9-emulation GEMM over FP32 inputs and output.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA / CUDA 13+ - Topics: INT8, INT4, mixed input, BF16 output, packed weights
Source: [csrc/gemm/gemm_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/gemm/gemm_ops.cpp)

## GEMM - Triton FP8 GEMM

### matmul_fp8_row

```python
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
```

Feature-rich Triton rowwise FP8 matrix product.

Computes a @ b.T for 2D or higher-rank a, applies reciprocal dequantization scales and optional bias, and returns BF16 by default.

**Constraints:**
- 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 - Module: `mslk.gemm.triton.fp8_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, FP8, rowwise, blockwise, X @ W.T
Source: [mslk/gemm/triton/fp8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py)

### matmul_fp8_block

```python
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
```

Triton block-scaled FP8 matrix product.

Computes a @ b.T using an explicit 3D scale-block contract.

**Constraints:**
- Accelerator-only.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.fp8_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, FP8, rowwise, blockwise, X @ W.T
Source: [mslk/gemm/triton/fp8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py)

### to_mxfp8

```python
to_mxfp8(data_hp: Tensor, block_size=32) -> tuple[Tensor, Tensor]
```

Converts high-precision data to MXFP8 with E8M0 block scales.

**Returns:** (scale_e8m0_biased, data_fp8).

**Constraints:**
- Return order is scales first, data second-the reverse of most quantizers.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.fp8_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, FP8, rowwise, blockwise, X @ W.T
Source: [mslk/gemm/triton/fp8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py)

### torch.ops.triton.matmul_fp8_row

```python
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
```

Custom-op registration for matmul_fp8_row.

**Constraints:**
- 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 - Module: `torch.ops.triton` - Platforms: NVIDIA, AMD - Topics: Triton, FP8, rowwise, blockwise, X @ W.T
Source: [mslk/gemm/triton/fp8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py)

### torch.ops.triton.matmul_fp8_block

```python
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
```

Custom-op registration for matmul_fp8_block.

**Constraints:**
- 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 - Module: `torch.ops.triton` - Platforms: NVIDIA, AMD - Topics: Triton, FP8, rowwise, blockwise, X @ W.T
Source: [mslk/gemm/triton/fp8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_gemm.py)

## GEMM - Triton grouped GEMM

### grouped_gemm

```python
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
```

General BF16 grouped GEMM with optional fused bias/router scaling.

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 - Module: `mslk.gemm.triton.grouped_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, grouped, MoE, M_sizes
Source: [mslk/gemm/triton/grouped_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py)

### grouped_gemm_fp8_rowwise

```python
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
```

Rowwise FP8 grouped GEMM for expert layers.

Concatenated-token grouped GEMMs with one weight matrix per group/expert. M_sizes partitions x along its first dimension.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.grouped_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, grouped, MoE, M_sizes
Source: [mslk/gemm/triton/grouped_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py)

### grouped_gemm_dgrad

```python
grouped_gemm_dgrad(x, w, m_sizes, out=None, num_sms=None) -> Tensor
```

Triton grouped BF16 data-gradient product.

Concatenated-token grouped GEMMs with one weight matrix per group/expert. M_sizes partitions x along its first dimension.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.grouped_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, grouped, MoE, M_sizes
Source: [mslk/gemm/triton/grouped_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py)

### grouped_gemm_wgrad

```python
grouped_gemm_wgrad(x, w, m_sizes, output=None, output_accum=False, num_sms=None) -> Tensor
```

Triton grouped BF16 weight-gradient product.

Concatenated-token grouped GEMMs with one weight matrix per group/expert. M_sizes partitions x along its first dimension.

**Constraints:**
- FP32 preallocated output is required when accumulating.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.grouped_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, grouped, MoE, M_sizes
Source: [mslk/gemm/triton/grouped_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/grouped_gemm.py)

## GEMM - Triton groupwise & microscaling

### matmul_f8f8bf16_groupwise

```python
matmul_f8f8bf16_groupwise(XQ, WQ, x_scale, w_scale, output=None) -> Tensor
```

Triton FP8 GEMM with K-group scale granularity 128.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.fp8_groupwise_gemm` - Platforms: NVIDIA, AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/fp8_groupwise_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_groupwise_gemm.py)

### matmul_f8f8bf16_groupwise_grouped

```python
matmul_f8f8bf16_groupwise_grouped(XQ, WQ, x_scale, w_scale, M_sizes, output=None) -> Tensor
```

Triton stacked grouped FP8 groupwise GEMM.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.fp8_groupwise_grouped_gemm` - Platforms: NVIDIA, AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/fp8_groupwise_grouped_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/fp8_groupwise_grouped_gemm.py)

### matmul_mx8mx4bf16

```python
matmul_mx8mx4bf16(XQ, WQ, x_scale, w_scale, output=None) -> Tensor
```

Triton MXFP8 x MXFP4 GEMM.

**Constraints:**
- ROCm gfx950 native block-scaled MFMA.
- ROCm activation scales are blocked/swizzled; weight scales are row-major [N,K/32].

Kind: function - Stability: Public - Module: `mslk.gemm.triton.mx8mx4_gemm` - Platforms: AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/mx8mx4_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/mx8mx4_gemm.py)

### matmul_mx8mx4bf16_grouped

```python
matmul_mx8mx4bf16_grouped(XQ, WQ, x_scale, w_scale, offsets, output=None) -> Tensor
```

Triton grouped MXFP8 x MXFP4 GEMM.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.mx8mx4_gemm` - Platforms: AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/mx8mx4_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/mx8mx4_gemm.py)

### matmul_mx8mx8bf16_grouped

```python
matmul_mx8mx8bf16_grouped(XQ, WQ, x_scale, w_scale, offsets, output=None, actual_num_tokens=None) -> Tensor
```

Triton grouped MXFP8 x MXFP8 GEMM.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.mx8mx8_gemm` - Platforms: AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/mx8mx8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/mx8mx8_gemm.py)

### mxfp4_gemm

```python
mxfp4_gemm(XQ, WQ, x_scale, w_scale, output=None) -> Tensor
```

Triton packed MXFP4 x MXFP4 GEMM.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.f4f4bf16` - Platforms: AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/f4f4bf16.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py)

### mxfp4_grouped_mm

```python
mxfp4_grouped_mm(XQ, WQ, x_scale, w_scale, offsets, output=None, global_scale=None) -> Tensor
```

ROCm gfx950 offset-grouped standard MXFP4 GEMM.

**Constraints:**
- Supports only XQ [total_M,K/2] with transposed WQ [G,K/2,N]; the CUDA 2D x 2D K-grouped layout is unsupported.
- global_scale is accepted for schema compatibility but ignored; leave it None.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.f4f4bf16` - Platforms: AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/f4f4bf16.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py)

### mxfp4_grouped_stacked_gemm

```python
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
```

ROCm gfx950 stacked grouped standard MXFP4 GEMM.

**Constraints:**
- 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 - Module: `mslk.gemm.triton.f4f4bf16` - Platforms: AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/f4f4bf16.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py)

### f4f4bf16

```python
f4f4bf16(XQ, WQ, x_scale, w_scale, output=None, global_scale=None, mxfp4_block_size=32) -> Tensor
```

ROCm compatibility wrapper for standard MXFP4 block-size-32 GEMM.

**Constraints:**
- 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 - Module: `mslk.gemm.triton.f4f4bf16` - Platforms: AMD gfx950 - Topics: Triton, groupwise, MXFP4, MXFP8, gfx950
Source: [mslk/gemm/triton/f4f4bf16.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/f4f4bf16.py)

## GEMM - Triton integer GEMM

### matmul_bf16i4_rowwise

```python
matmul_bf16i4_rowwise(X, W, w_scale_group, w_zero_group) -> Tensor
```

Triton BF16 activation x row/group INT4 weight GEMM.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.int4_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, INT4, INT8, BF16
Source: [mslk/gemm/triton/int4_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int4_gemm.py)

### matmul_bf16i4_rowwise_batched

```python
matmul_bf16i4_rowwise_batched(X, W, w_scale, w_zp) -> Tensor
```

Triton batched BF16 x INT4 GEMM.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.int4_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, INT4, INT8, BF16
Source: [mslk/gemm/triton/int4_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int4_gemm.py)

### i8i8bf16_triton

```python
i8i8bf16_triton(XQ, WQ, scale: float, split_k=1) -> Tensor
```

Triton INT8 x INT8 GEMM with scalar scale.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.int8_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, INT4, INT8, BF16
Source: [mslk/gemm/triton/int8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int8_gemm.py)

### i8i8bf16_dynamic_triton

```python
i8i8bf16_dynamic_triton(XQ, WQ, scale: Tensor, split_k=1) -> Tensor
```

Triton INT8 x INT8 GEMM with tensor scale.

Kind: function - Stability: Public - Module: `mslk.gemm.triton.int8_gemm` - Platforms: NVIDIA, AMD - Topics: Triton, INT4, INT8, BF16
Source: [mslk/gemm/triton/int8_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/triton/int8_gemm.py)

## GEMM - Blackwell mixed-input CuTeDSL

### MixedInputGemmKernel

```python
MixedInputGemmKernel(scale_granularity_m, scale_granularity_k, acc_dtype, use_2cta_instrs, mma_tiler_mnk, cluster_shape_mn, use_tma_store)
```

Configurable compiled Blackwell mixed-input GEMM kernel object.

CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.

**Methods / arguments:**
- `__call__(a, a_scale, b, c, max_active_clusters, stream)` - Launches a configured kernel.
- `can_implement(...)` - Checks alignment, tiler, cluster, layout, and epilogue constraints.

Kind: class - Stability: Advanced - Module: `mslk.gemm.blackwell_mixed_input_gemm` - Platforms: NVIDIA SM100 - Topics: CuTeDSL, Blackwell, INT4, INT8, BF16, mixed input
Source: [mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py)

### mixed_input_gemm

```python
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
```

General Blackwell mixed-input GEMM frontend.

CuTeDSL kernels for a narrow integer operand and a wide BF16/FP16 operand, with optional block scaling and TMA output stores.

**Constraints:**
- Without TMA store, output dimensions must not leave out-of-bounds epilogue tiles.

Kind: function - Stability: Public - Module: `mslk.gemm.blackwell_mixed_input_gemm` - Platforms: NVIDIA SM100 - Topics: CuTeDSL, Blackwell, INT4, INT8, BF16, mixed input
Source: [mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py)

### int4bf16bf16_gemm

```python
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
```

Convenience INT4 narrow operand x BF16 wide operand GEMM.

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 - Module: `mslk.gemm.blackwell_mixed_input_gemm` - Platforms: NVIDIA SM100 - Topics: CuTeDSL, Blackwell, INT4, INT8, BF16, mixed input
Source: [mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py)

### int8bf16bf16_gemm

```python
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
```

Convenience INT8 narrow operand x BF16 wide operand GEMM.

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 - Module: `mslk.gemm.blackwell_mixed_input_gemm` - Platforms: NVIDIA SM100 - Topics: CuTeDSL, Blackwell, INT4, INT8, BF16, mixed input
Source: [mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py)

### create_tensors

```python
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)
```

Creates benchmark/reference tensors and layouts for mixed-input GEMM.

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 - Module: `mslk.gemm.blackwell_mixed_input_gemm` - Platforms: NVIDIA SM100 - Topics: CuTeDSL, Blackwell, INT4, INT8, BF16, mixed input
Source: [mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py)

### compare

```python
compare(a_torch_cpu, b_torch_cpu, a_scale_torch_cpu, c_torch_gpu, c_dtype, tolerance) -> None
```

Compares kernel output with a CPU reference.

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 - Module: `mslk.gemm.blackwell_mixed_input_gemm` - Platforms: NVIDIA SM100 - Topics: CuTeDSL, Blackwell, INT4, INT8, BF16, mixed input
Source: [mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py)

### run

```python
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)
```

Benchmark/validation driver for mixed-input configurations.

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 - Module: `mslk.gemm.blackwell_mixed_input_gemm` - Platforms: NVIDIA SM100 - Topics: CuTeDSL, Blackwell, INT4, INT8, BF16, mixed input
Source: [mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/gemm/blackwell_mixed_input_gemm/mixed_input_gemm.py)

## Quantization - FP8 quantization

### triton_quantize_fp8_row

```python
triton_quantize_fp8_row(a, scale_ub=None, zero_start_index_M=None, align_rows_to=None, eps_opt=1/512) -> tuple[Tensor, Tensor]
```

Triton rowwise FP8 quantizer.

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

**Constraints:**
- 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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### quantize_fp8_row

```python
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]
```

Public rowwise FP8 adapter selecting the Triton kernel or a PyTorch fallback.

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

**Constraints:**
- 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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### scale_fp8_row

```python
scale_fp8_row(a, x_scale, w_scale) -> Tensor
```

Applies activation and weight row scales to an accumulator/result tensor.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### triton_quantize_fp8_block

```python
triton_quantize_fp8_block(x, block_m=256, block_k=256, scale_ub=None, k_major=True) -> tuple[Tensor, Tensor]
```

Triton 2D blockwise FP8 quantizer.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### quantize_fp8_block

```python
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]
```

Public blockwise FP8 adapter.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### triton_quantize_fp8_group

```python
triton_quantize_fp8_group(x, group_size=128, scale_ub=None, m_sizes=None, k_major=True) -> tuple[Tensor, Tensor]
```

Triton K-group FP8 quantizer, optionally aware of grouped valid row counts.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### quantize_fp8_group

```python
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]
```

Public groupwise FP8 adapter.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### triton_quantize_fp8_tensor

```python
triton_quantize_fp8_tensor(a) -> tuple[Tensor, Tensor]
```

Triton tensorwise FP8 quantizer.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### quantize_fp8_tensor

```python
quantize_fp8_tensor(a, use_triton=True) -> tuple[Tensor, Tensor]
```

Public tensorwise FP8 adapter.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### dequantize_fp8_row

```python
dequantize_fp8_row(xq, x_scale) -> Tensor
```

Dequantizes rowwise FP8 storage to BF16.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### dequantize_fp8_block

```python
dequantize_fp8_block(xq, x_scale, block_m=256, block_k=256) -> Tensor
```

Dequantizes blockwise FP8 storage to BF16.

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 - Module: `mslk.quantize.triton.fp8_quantize` - Platforms: NVIDIA, AMD - Topics: FP8, E4M3, reciprocal scale, Triton
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

## Quantization - FP8 custom ops & dtype helpers

### torch.ops.triton.quantize_fp8_row

```python
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]
```

Custom-op registration for rowwise FP8 quantization.

**Constraints:**
- The Meta implementation currently omits the runtime eps_opt argument.

Kind: torch op - Stability: Low-level - Module: `torch.ops.triton` - Platforms: NVIDIA, AMD, Meta / fake - Topics: torch.compile, custom op, FP8 dtype
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### torch.ops.triton.quantize_fp8_block

```python
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]
```

Custom-op registration for blockwise FP8 quantization.

Kind: torch op - Stability: Low-level - Module: `torch.ops.triton` - Platforms: NVIDIA, AMD, Meta / fake - Topics: torch.compile, custom op, FP8 dtype
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### torch.ops.triton.quantize_fp8_tensor

```python
torch.ops.triton.quantize_fp8_tensor(a, use_triton=True) -> (Tensor, Tensor)
```

Custom-op registration for tensorwise FP8 quantization.

Kind: torch op - Stability: Low-level - Module: `torch.ops.triton` - Platforms: NVIDIA, AMD, Meta / fake - Topics: torch.compile, custom op, FP8 dtype
Source: [mslk/quantize/triton/fp8_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp8_quantize.py)

### get_fp8_constants

```python
get_fp8_constants() -> tuple[torch.dtype, tl.dtype, float, float]
```

Returns the platform PyTorch dtype, Triton dtype, maximum finite value, and epsilon.

Kind: function - Stability: Public - Module: `mslk.utils.triton.fp8_utils` - Platforms: NVIDIA, AMD, Meta / fake - Topics: torch.compile, custom op, FP8 dtype
Source: [mslk/utils/triton/fp8_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/triton/fp8_utils.py)

### reinterpret_fp8_type

```python
reinterpret_fp8_type(tensor: Tensor, dtype: tl.dtype) -> TensorWrapper
```

Reinterprets FP8 storage for a Triton dtype without numeric conversion.

Kind: function - Stability: Advanced - Module: `mslk.utils.triton.fp8_utils` - Platforms: NVIDIA, AMD, Meta / fake - Topics: torch.compile, custom op, FP8 dtype
Source: [mslk/utils/triton/fp8_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/triton/fp8_utils.py)

## Quantization - MXFP4 quantization

### triton_quantize_mx4

```python
triton_quantize_mx4(input, *, rounding_mode=RoundingMode.ceil, seed=None) -> tuple[Tensor, Tensor]
```

Preferred public MXFP4 group-32 quantizer.

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.

**Constraints:**
- CUDA scale storage is flattened 128 x 4 blocked/swizzled int8.
- ROCm scale storage is plain [...,K/32] uint8.
- Stochastic rounding is CUDA-only.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+, AMD gfx950 - Topics: MXFP4, E2M1, E8M0, packed, rounding
Source: [mslk/quantize/triton/fp4_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py)

### triton_quantize_mx4_unpack

```python
triton_quantize_mx4_unpack(input, group_size=32, ebits=2, mbits=1, rounding_mode=RoundingMode.ceil, stochastic_casting=False, *, seed=None) -> tuple[Tensor, Tensor]
```

Compatibility adapter supporting group size 16/32 and legacy rounding arguments.

Quantizes BF16/FP16 values into packed E2M1 elements with shared E8M0 scales. The last dimension must be divisible by the group size.

**Constraints:**
- Only E2M1 (ebits=2, mbits=1) is implemented.
- stochastic_casting=True forces stochastic rounding.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+, AMD gfx950 - Topics: MXFP4, E2M1, E8M0, packed, rounding
Source: [mslk/quantize/triton/fp4_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py)

### quantize_mx4

```python
quantize_mx4(x, group_size=32, rounding_mode=RoundingMode.ceil, *, seed=None) -> tuple[Tensor, Tensor]
```

Direct MXFP4 kernel frontend underlying the public adapter.

Quantizes BF16/FP16 values into packed E2M1 elements with shared E8M0 scales. The last dimension must be divisible by the group size.

**Constraints:**
- For a 1D input, preserves an extra row dimension; the public adapter restores the expected shape.

Kind: function - Stability: Advanced - Module: `mslk.quantize.triton.quantize_kernels.mx4` - Platforms: NVIDIA SM100+, AMD gfx950 - Topics: MXFP4, E2M1, E8M0, packed, rounding
Source: [mslk/quantize/triton/quantize_kernels/mx4.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/quantize_kernels/mx4.py)

### quantize_mx4_stacked

```python
quantize_mx4_stacked(m_sizes, x, group_size=32, rounding_mode=RoundingMode.ceil, *, seed=None) -> tuple[Tensor, Tensor]
```

Segment-aware MXFP4 quantization for stacked grouped inputs.

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 - Module: `mslk.quantize.triton.quantize_kernels.mx4_stacked` - Platforms: NVIDIA SM100+, AMD gfx950 - Topics: MXFP4, E2M1, E8M0, packed, rounding
Source: [mslk/quantize/triton/quantize_kernels/mx4_stacked.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/quantize_kernels/mx4_stacked.py)

## Quantization - NVFP4 & legacy FP4

### triton_quantize_nvfp4

```python
triton_quantize_nvfp4(x, global_scale, use_e8m0_scale=False, use_precise_math=True) -> tuple[Tensor, Tensor]
```

NVFP4 quantizer with per-16-value E4M3 scale factors.

**Returns:** (packed FP4, swizzled scales).

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/quantize.py)

### triton_fake_quantize_nvfp4_per_tensor

```python
triton_fake_quantize_nvfp4_per_tensor(input, static_scales=None, scale_ub=None) -> tuple[Tensor, Tensor]
```

Per-tensor NVFP4 fake quantization returning BF16 values.

**Returns:** (BF16 fake-quantized tensor, FP32 scale or amax).

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/fake_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fake_quantize.py)

### nvfp4_quantize_stacked

```python
nvfp4_quantize_stacked(m_sizes, input, global_scale) -> tuple[Tensor, Tensor]
```

Segment-aware stacked NVFP4 quantization.

**Returns:** (packed FP4, padded two-dimensional swizzled scale buffer).

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/quantize_stacked.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/quantize_stacked.py)

### nvfp4_quantize_stacked_with_token_scale

```python
nvfp4_quantize_stacked_with_token_scale(m_sizes, input) -> tuple[Tensor, Tensor, Tensor]
```

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 - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/quantize_stacked.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/quantize_stacked.py)

### calculate_group_max

```python
calculate_group_max(input, m_sizes) -> tuple[Tensor, Tensor]
```

Computes per-segment NVFP4 global scales and a row-to-segment map.

The name is misleading: the first output is effectively 448 x 6/max_abs, not the raw maximum.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/global_scale.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/global_scale.py)

### cal_global_scale_mx4_as_nvfp4

```python
cal_global_scale_mx4_as_nvfp4(x) -> Tensor
```

Derives an NVFP4-style global scale for MXFP4 data.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/primitives.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py)

### global_scale_nvfp4

```python
global_scale_nvfp4(x) -> Tensor
```

Computes the standard global NVFP4 scale from input range.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_utils` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/fp4_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py)

### fp4_to_float

```python
fp4_to_float(x) -> Tensor
```

Decodes FP4 nibbles to FP32 values.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_utils` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/fp4_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py)

### scale_nvfp4

```python
scale_nvfp4(x, scale, global_scale, group_size=16) -> Tensor
```

Applies NVFP4 local and global scaling to decoded values.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_utils` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/fp4_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py)

### dequantize_nvfp4

```python
dequantize_nvfp4(input_quantized, scale, global_scale, group_size=16) -> Tensor
```

Dequantizes packed NVFP4 to BF16.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_utils` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/fp4_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py)

### dequantize_mx4

```python
dequantize_mx4(input_quantized, scale, group_size=32) -> Tensor
```

Dequantizes packed MXFP4 with NVIDIA blocked scales to BF16.

**Constraints:**
- Not compatible with the plain row-major scale output produced on ROCm.

Kind: function - Stability: Public - Module: `mslk.quantize.triton.fp4_utils` - Platforms: NVIDIA SM100+ - Topics: NVFP4, E4M3 scales, global scale, packed
Source: [mslk/quantize/triton/legacy/fp4_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/fp4_utils.py)

## Quantization - FP4 constants & primitives

### RoundingMode

```python
RoundingMode.nearest=0 | floor=1 | even=2 | stochastic=3 | ceil=4
```

Rounding policy enum used by MXFP4 quantizers.

**Constraints:**
- Stochastic mode is rejected on ROCm.

Kind: enum - Stability: Advanced - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py)

### get_mx4_exp_bias

```python
get_mx4_exp_bias(ebits) -> int
```

Returns the MX exponent bias for a format width.

Kind: function - Stability: Advanced - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/legacy/primitives.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py)

### FP4_E2M1_MAX

```python
FP4_E2M1_MAX = 6.0
```

Largest finite E2M1 magnitude.

Kind: constant - Stability: Advanced - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py)

### FP8_E4M3_MAX

```python
FP8_E4M3_MAX = 448
```

Largest finite OCP E4M3 magnitude.

Kind: constant - Stability: Advanced - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py)

### FP4_EBITS

```python
FP4_EBITS = 2
```

E2M1 exponent-bit count.

Kind: constant - Stability: Advanced - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py)

### FP4_MBITS

```python
FP4_MBITS = 1
```

E2M1 mantissa-bit count.

Kind: constant - Stability: Advanced - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_quantize.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_quantize.py)

### E8M0_EXPONENT_BIAS

```python
E8M0_EXPONENT_BIAS = tl.constexpr(127)
```

Shared exponent bias used to encode MX block scales.

**Constraints:**
- Host/Python consumers can read the wrapped scalar through E8M0_EXPONENT_BIAS.value.

Kind: constant - Stability: Advanced - Module: `mslk.quantize.triton.fp4_primitives` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_primitives/constants.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/constants.py)

### BF16_MIN_NORMAL

```python
BF16_MIN_NORMAL = tl.constexpr(2**-126)
```

Smallest normal BF16 magnitude used by safe scale math.

**Constraints:**
- Host/Python consumers can read the wrapped scalar through BF16_MIN_NORMAL.value.

Kind: constant - Stability: Advanced - Module: `mslk.quantize.triton.fp4_primitives` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_primitives/constants.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/constants.py)

### blocked_scale_offset

```python
blocked_scale_offset(logical_row, logical_col, n_col_blocks, num_cols)
```

Maps a logical scale coordinate into Blackwell's 128 x 4 blocked layout.

Kind: function - Stability: Advanced - Module: `mslk.quantize.triton.fp4_primitives` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_primitives/layout.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/layout.py)

### stacked_segment_map

```python
stacked_segment_map(m_sizes_ptr, pid_m, M_PER_BLOCK, NUM_SEGMENTS, PREFIX_NUM, BSEARCH_ITERS)
```

Triton helper mapping a row tile to a stacked segment/expert.

Kind: function - Stability: Advanced - Module: `mslk.quantize.triton.fp4_primitives` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_primitives/layout.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/layout.py)

### mx4_scale_normalize_encode

```python
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)
```

Triton primitive that normalizes an MX block and encodes its E8M0 scale.

Kind: function - Stability: Advanced - Module: `mslk.quantize.triton.fp4_primitives` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_primitives/scale.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/scale.py)

### convert_fp32_to_fp4_packed

```python
convert_fp32_to_fp4_packed(x_pairs, IS_GFX950, IS_ROCM)
```

Triton primitive converting pairs of FP32 values into packed E2M1 nibbles.

Kind: function - Stability: Advanced - Module: `mslk.quantize.triton.fp4_primitives` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/fp4_primitives/packing.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/fp4_primitives/packing.py)

### unsigned_fp32_to_e8m0

```python
unsigned_fp32_to_e8m0(tensor, mbits, scale_round_mode)
```

Compatibility Triton primitive encoding positive FP32 magnitudes as E8M0 scales.

Kind: function - Stability: Low-level - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/legacy/primitives.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py)

### nvfp4_scale_swizzle

```python
nvfp4_scale_swizzle(offs_m)
```

Compatibility Triton helper mapping rows into the NVFP4 scale-swizzle layout.

Kind: function - Stability: Low-level - Module: `mslk.quantize.triton.fp4_quantize` - Platforms: NVIDIA, AMD - Topics: FP4, Triton primitive, rounding, constants
Source: [mslk/quantize/triton/legacy/primitives.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/triton/legacy/primitives.py)

## Quantization - INT4 preprocessing

### pack_int4

```python
pack_int4(x: Tensor) -> Tensor
```

Packs two logical 4-bit integer values into each int8 byte.

**Returns:** Packed int8 with last dimension halved.

Kind: function - Stability: Public - Module: `mslk.quantize.shuffle` - Platforms: NVIDIA, AMD - Topics: INT4, packing, preshuffle, zero point
Source: [mslk/quantize/shuffle.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py)

### int4_row_quantize_zp

```python
int4_row_quantize_zp(x, group_size=128) -> tuple[Tensor, Tensor, Tensor]
```

Groupwise asymmetric INT4 quantization with zero points.

**Returns:** (unpacked int8 values [N,K], scales, zero_points).

**Constraints:**
- Values are not packed; call pack_int4 before a packed-weight GEMM.

Kind: function - Stability: Public - Module: `mslk.quantize.shuffle` - Platforms: NVIDIA, AMD - Topics: INT4, packing, preshuffle, zero point
Source: [mslk/quantize/shuffle.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py)

### int4_row_quantize

```python
int4_row_quantize(x, group_size=128) -> tuple[Tensor, Tensor]
```

Groupwise symmetric INT4 quantization.

**Returns:** (unpacked int8 values [N,K], scales).

**Constraints:**
- 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 - Module: `mslk.quantize.shuffle` - Platforms: NVIDIA, AMD - Topics: INT4, packing, preshuffle, zero point
Source: [mslk/quantize/shuffle.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py)

### quantize_int4_preshuffle

```python
quantize_int4_preshuffle(w, group_size=128, dtype='fp8', use_zp=True) -> tuple[Tensor, tuple[Tensor, Tensor]]
```

Quantizes, packs, and CUDA-preshuffles INT4 weights and scale terms.

**Returns:** (preshuffled packed weights, (scale_a, scale_b)).

Kind: function - Stability: Public - Module: `mslk.quantize.shuffle` - Platforms: NVIDIA - Topics: INT4, packing, preshuffle, zero point
Source: [mslk/quantize/shuffle.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py)

### ck_preshuffle

```python
ck_preshuffle(src, NXdl=16) -> Tensor
```

Reorders data into the AMD Composable Kernel XDL weight layout.

**Constraints:**
- Source must use float8_e4m3fnuz.

Kind: function - Stability: Public - Module: `mslk.quantize.shuffle` - Platforms: AMD ROCm - Topics: INT4, packing, preshuffle, zero point
Source: [mslk/quantize/shuffle.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/shuffle.py)

## Quantization - MXFP6 helpers

### quantize_bf16_to_mx6_e2m3

```python
quantize_bf16_to_mx6_e2m3(x, block_size=32) -> tuple[Tensor, Tensor]
```

Quantizes BF16 to unpacked six-bit E2M3 codes plus E8M0 block scales.

**Returns:** (unpacked uint8 codes, E8M0 scales).

**Constraints:**
- The codes are not GEMM-ready until pack_fp6_e2m3 is applied.

Kind: function - Stability: Public - Module: `mslk.quantize.mx_mixed_dtype_utils` - Platforms: NVIDIA SM100+ - Topics: MXFP6, E2M3, E8M0, packing
Source: [mslk/quantize/mx_mixed_dtype_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py)

### pack_fp6_e2m3

```python
pack_fp6_e2m3(unpacked: Tensor) -> Tensor
```

Bit-packs four E2M3 values into three uint8 bytes.

Kind: function - Stability: Public - Module: `mslk.quantize.mx_mixed_dtype_utils` - Platforms: NVIDIA SM100+ - Topics: MXFP6, E2M3, E8M0, packing
Source: [mslk/quantize/mx_mixed_dtype_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py)

### E2M3_DECODE

```python
E2M3_DECODE: tuple[float, ...]
```

Lookup table mapping 6-bit E2M3 codes to decoded values.

Kind: constant - Stability: Public - Module: `mslk.quantize.mx_mixed_dtype_utils` - Platforms: NVIDIA SM100+ - Topics: MXFP6, E2M3, E8M0, packing
Source: [mslk/quantize/mx_mixed_dtype_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py)

### E2M3_MAX

```python
E2M3_MAX = 7.5
```

Largest finite E2M3 magnitude.

Kind: constant - Stability: Public - Module: `mslk.quantize.mx_mixed_dtype_utils` - Platforms: NVIDIA SM100+ - Topics: MXFP6, E2M3, E8M0, packing
Source: [mslk/quantize/mx_mixed_dtype_utils.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/quantize/mx_mixed_dtype_utils.py)

## MoE - Routing & token movement

### index_shuffling

```python
index_shuffling(routing_scores, expert_index_start=None, expert_index_end=None, valid_token_count=None, top_k=1) -> tuple[Tensor, Tensor, Tensor]
```

Selects top-k experts and groups token indices by expert.

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 x top_k], token_indices [T x top_k].

**Constraints:**
- 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 - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [csrc/moe/moe_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/moe/moe_ops.cpp)

### gather_scale_dense_tokens

```python
gather_scale_dense_tokens(x, token_indices, expert_indices, scores, valid_token_count=None) -> Tensor
```

Gathers routed tokens and multiplies each by its router score.

Output row i is x[token_indices[i]] x scores[token_indices[i], expert_indices[i]]. x is [T,D] and the result is [selected,D].

**Constraints:**
- 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 - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### gather_scale_quant_dense_tokens

```python
gather_scale_quant_dense_tokens(x, token_indices, expert_indices, scores, scale_ub=None, valid_token_count=None) -> tuple[Tensor, Tensor]
```

Gathers/scales routed tokens and rowwise-quantizes them to FP8.

**Returns:** (FP8 [selected,D], float32 reciprocal scales [selected]).

**Constraints:**
- 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 - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### scatter_add_dense_tokens

```python
scatter_add_dense_tokens(out_tokens, in_tokens, token_indices, valid_token_count=None) -> None
```

Atomically accumulates routed expert rows back to dense token order.

Mutates out_tokens [T,D] by adding in_tokens[i] into row token_indices[i].

**Constraints:**
- Contiguous tensors required.
- NVIDIA requires CUDA toolkit 12.4+; ROCm is supported.

Kind: function - Stability: Public - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### scatter_add_padded_tokens

```python
scatter_add_padded_tokens(in_tokens, token_counts, token_indices, out_tokens) -> None
```

Scatters expert/rank-padded token blocks back into dense token storage.

Consumes in_tokens [EP,T_K,D], per-expert token counts, and token indices, then mutates out_tokens [T,D].

**Constraints:**
- T_K must be divisible by T and E by EP.
- NVIDIA requires CUDA toolkit 12.4+; ROCm is supported.

Kind: function - Stability: Public - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### combine_shuffling

```python
combine_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False) -> tuple[Tensor, Tensor]
```

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 - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/shuffling.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py)

### split_shuffling

```python
split_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False, init_with_zeros=False) -> Tensor
```

Inverse of combine_shuffling for expert communication results.

**Constraints:**
- 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 - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/shuffling.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py)

### silu_mul

```python
silu_mul(x0, x1, valid_token_count=None) -> Tensor
```

Fused SwiGLU activation x0-sigmoid(x0)-x1.

**Returns:** Same shape and dtype as [T,D] inputs.

**Constraints:**
- 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 - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/activation.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py)

### silu_mul_quant

```python
silu_mul_quant(x0, x1, scale_ub=None, valid_token_count=None) -> tuple[Tensor, Tensor]
```

Fused SwiGLU plus rowwise FP8 quantization.

**Returns:** (FP8 [T,D], float32 reciprocal scales [T]).

**Constraints:**
- 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 - Module: `mslk.moe` - Platforms: NVIDIA, AMD - Topics: routing, tokens, experts, Triton
Source: [mslk/moe/activation.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py)

## MoE - Registered MoE ops

### torch.ops.mslk.index_shuffling

```python
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)
```

Native top-k expert grouping op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [csrc/moe/moe_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/moe/moe_ops.cpp)

### torch.ops.mslk.scatter_add_along_first_dim

```python
torch.ops.mslk.scatter_add_along_first_dim(Dst, Src, Index) -> None
```

In-place first-dimension scatter-add with a fast BF16 TMA path.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA - Topics: dispatcher, custom op, MoE
Source: [csrc/moe/moe_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/moe/moe_ops.cpp)

### torch.ops.mslk.silu_mul

```python
torch.ops.mslk.silu_mul(x0, x1, valid_token_count=None) -> Tensor
```

Registered fused SwiGLU activation.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/activation.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py)

### torch.ops.mslk.silu_mul_quant

```python
torch.ops.mslk.silu_mul_quant(x0, x1, scale_ub=None, valid_token_count=None) -> (Tensor, Tensor)
```

Registered fused SwiGLU plus FP8 quantization.

Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.

**Constraints:**
- Current fake/Meta implementation returns only one tensor despite this two-output schema.

Kind: torch op - Stability: Low-level - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/activation.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/activation.py)

### torch.ops.mslk.gather_scale_dense_tokens

```python
torch.ops.mslk.gather_scale_dense_tokens(x, token_indices, expert_indices, scores, valid_token_count=None) -> Tensor
```

Registered routed-token gather/scale op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### torch.ops.mslk.gather_scale_quant_dense_tokens

```python
torch.ops.mslk.gather_scale_quant_dense_tokens(x, token_indices, expert_indices, scores, scale_ub=None, valid_token_count=None) -> Tensor
```

Registered quantizing routed-token gather with a known return-schema mismatch.

Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.

**Constraints:**
- 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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### torch.ops.mslk.scatter_add_dense_tokens

```python
torch.ops.mslk.scatter_add_dense_tokens(out_tokens, in_tokens, token_indices, valid_token_count=None) -> None
```

Registered dense-token scatter-add op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### torch.ops.mslk.scatter_add_padded_tokens

```python
torch.ops.mslk.scatter_add_padded_tokens(in_tokens, token_counts, token_indices, out_tokens) -> None
```

Registered padded expert-token scatter op.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/gather_scatter.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/gather_scatter.py)

### torch.ops.mslk.combine_shuffling

```python
torch.ops.mslk.combine_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False) -> (Tensor, Tensor)
```

Registered rank-major to expert-major reorder.

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 - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/shuffling.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py)

### torch.ops.mslk.split_shuffling

```python
torch.ops.mslk.split_shuffling(tokens, token_counts, expert_start=None, expert_end=None, is_padded=False, init_with_zeros=False) -> Tensor
```

Registered inverse expert-shuffling reorder.

Dispatcher-level schemas behind the Python wrappers. Import mslk.moe before use so native, custom, and fake implementations are registered.

**Constraints:**
- Current registered implementation signatures omit init_with_zeros.

Kind: torch op - Stability: Low-level - Module: `torch.ops.mslk` - Platforms: NVIDIA, AMD, Meta / fake - Topics: dispatcher, custom op, MoE
Source: [mslk/moe/shuffling.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/shuffling.py)

## MoE - Composed MoE layers

### MoEArgs

```python
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)
```

Frozen configuration dataclass for the provided MoE modules.

**Methods / arguments:**
- `num_local_experts -> int` - Cached num_experts // ep_size property.

**Constraints:**
- No constructor defaults.
- precision is not interpreted directly; initialization methods typically select BF16 versus FP8 rowwise storage.

Kind: class - Stability: Public - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

### BaselineMoE

```python
BaselineMoE(ep_group: ProcessGroup, ep_mp_group: ProcessGroup, moe_args: MoEArgs)
```

Dense/reference distributed SwiGLU MoE implementation.

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

**Methods / arguments:**
- `build(init_methods=None) -> BaselineMoE` - Initializes router/shared/routed parameters; required before forward.
- `forward(x, use_static_shape)` - Computes shared and routed experts; Baseline accepts but does not use use_static_shape.
- `router_DE / E / EG / K` - Router and geometry properties.
- `is_shared_fp8_rowwise / is_routed_fp8_rowwise` - Reports scaled FP8 parameter storage.

**Constraints:**
- Caller performs model-parallel all-reduce externally.
- Imports require fairscale and pyre_extensions.

Kind: class - Stability: Public - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

### MetaShufflingMoE

```python
MetaShufflingMoE(ep_group: ProcessGroup, ep_mp_group: ProcessGroup, moe_args: MoEArgs)
```

Optimized inference MoE using shuffle/gather/grouped-GEMM primitives and overlapped communication.

**Returns:** forward(x [B,T,D], use_static_shape) -> [B,T,D].

**Constraints:**
- 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 - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

### ScaledParameter

```python
ScaledParameter(data: Tensor, scale: Optional[Tensor]=None)
```

Non-trainable Parameter carrying optional quantization scales.

**Methods / arguments:**
- `weights -> Tensor` - Underlying parameter data.
- `scales -> Tensor` - Scale tensor; asserts that one exists.
- `is_scaled -> bool` - Whether scales are attached.

Kind: class - Stability: Advanced - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

### Experts

```python
Experts(dim: int, hidden_dim: int)
```

Abstract shared base for expert weight containers.

**Methods / arguments:**
- `build(init_methods=None) -> Experts` - Initializes and packs expert parameters.
- `w13 / w2 -> ScaledParameter` - Packed SwiGLU input/gate and output weights.
- `is_fp8_rowwise -> bool` - Detects FP8 rowwise storage from w13 dtype.

**Constraints:**
- Model-parallel world size must divide dim and hidden_dim.

Kind: class - Stability: Advanced - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

### RoutedExperts

```python
RoutedExperts(num_local_experts: int, dim: int, hidden_dim: int)
```

Per-local-expert stacked weight container.

Kind: class - Stability: Advanced - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

### SharedExperts

```python
SharedExperts(dim: int, hidden_dim: int)
```

Shared SwiGLU expert weight container.

Kind: class - Stability: Advanced - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

### init_params

```python
init_params(key: str, param: ScaledParameter, init_methods: Mapping) -> None
```

Applies a named initialization/quantization callback or Kaiming-uniform fallback.

**Constraints:**
- Recognized keys cover routed w_in/w_out/w_swiglu, shared equivalents, and router_DE.

Kind: function - Stability: Advanced - Module: `mslk.moe.layers` - Platforms: NVIDIA, AMD, Distributed - Topics: PyTorch module, expert parallel, model parallel, inference
Source: [mslk/moe/layers.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/moe/layers.py)

## Convolution - FP8 3D convolution

### torch.ops.mslk.f8f8bf16_conv

```python
torch.ops.mslk.f8f8bf16_conv(activation, filter, scale, padding: int[3], stride: int[3], dilation: int[3]) -> Tensor
```

FP8 3D cross-correlation/convolution with combined scale and BF16 output.

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.

**Constraints:**
- 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 - Module: `torch.ops.mslk` - Platforms: NVIDIA SM100+, Meta / fake - Topics: FP8, BF16, 3D convolution, NDHWC, channels_last_3d, CUTLASS
Source: [csrc/conv/conv_ops.cpp](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/csrc/conv/conv_ops.cpp)

## Runtime - FlyDSL runtime & AOT

### is_flydsl_available

```python
is_flydsl_available() -> bool
```

Cached check for FlyDSL importability and current-architecture support.

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.

**Methods / arguments:**
- `is_flydsl_available.cache_clear()` - Invalidates the lru_cache after environment changes.

**Constraints:**
- mslk.flydsl.__init__ is empty; import the explicit common, jit, or aot submodule.

Kind: function - Stability: Public - Module: `mslk.flydsl.common` - Platforms: AMD ROCm, CPU / Python - Topics: FlyDSL, JIT, AOT, cache
Source: [mslk/flydsl/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/common.py)

### require_flydsl

```python
require_flydsl() -> None
```

Raises RuntimeError with an installation hint when FlyDSL is unavailable.

**Constraints:**
- mslk.flydsl.__init__ is empty; import the explicit common, jit, or aot submodule.

Kind: function - Stability: Public - Module: `mslk.flydsl.common` - Platforms: AMD ROCm, CPU / Python - Topics: FlyDSL, JIT, AOT, cache
Source: [mslk/flydsl/common.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/common.py)

### configure_runtime_cache

```python
configure_runtime_cache() -> None
```

Points FlyDSL at bundled AOT artifacts and applies MSLK's JIT-disable switch.

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.

**Constraints:**
- Called automatically when mslk.flydsl.jit is imported.

Kind: function - Stability: Public - Module: `mslk.flydsl.jit` - Platforms: AMD ROCm, CPU / Python - Topics: FlyDSL, JIT, AOT, cache
Source: [mslk/flydsl/jit.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/jit.py)

### run_compiled

```python
run_compiled(launcher: Callable[..., Any], *args: Any) -> None
```

Compiles a FlyDSL launcher once, then invokes the cached compiled function.

The first compiler call executes the launch itself; only subsequent calls invoke launcher._mslk_cf separately.

**Constraints:**
- mslk.flydsl.__init__ is empty; import the explicit common, jit, or aot submodule.

Kind: function - Stability: Public - Module: `mslk.flydsl.jit` - Platforms: AMD ROCm, CPU / Python - Topics: FlyDSL, JIT, AOT, cache
Source: [mslk/flydsl/jit.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/jit.py)

### collect_aot_jobs

```python
collect_aot_jobs() -> list[tuple[str, dict[str, Any], str]]
```

Builds the Cartesian product of registered kernel configs and architectures.

**Constraints:**
- The shipped _AOT_KERNEL_MODULES list is currently empty, so the default result contains no jobs.

Kind: function - Stability: Public - Module: `mslk.flydsl.aot` - Platforms: AMD ROCm, CPU / Python - Topics: FlyDSL, JIT, AOT, cache
Source: [mslk/flydsl/aot.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/aot.py)

### compile_aot

```python
compile_aot(cache_dir: str) -> None
```

Multiprocess-compiles all collected FlyDSL AOT jobs into a runtime cache.

Uses spawn workers under COMPILE_ONLY=1, aggregates failures, and restores environment state.

**Constraints:**
- 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 - Module: `mslk.flydsl.aot` - Platforms: AMD ROCm, CPU / Python - Topics: FlyDSL, JIT, AOT, cache
Source: [mslk/flydsl/aot.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/aot.py)

## Runtime - Device detection

### is_cuda

```python
is_cuda() -> bool
```

True for an NVIDIA CUDA build with an available device.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### is_rocm

```python
is_rocm() -> bool
```

True for a PyTorch HIP build with an available AMD device.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### compute_capability_in

```python
compute_capability_in(major_min: int, major_max: int | None=None) -> bool
```

Checks current device compute-capability major against an inclusive range.

**Constraints:**
- Returns false without a device; implementation also reads the tuple on ROCm.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### compute_capability_at_least

```python
compute_capability_at_least(major_min: int, minor_min: int=0) -> bool
```

Checks the current capability tuple against a minimum.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### cuda_version_at_least

```python
cuda_version_at_least(major_min: int) -> bool
```

Checks the PyTorch build's CUDA toolkit major, not driver or device capability.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### get_gfx_arch_name

```python
get_gfx_arch_name() -> str
```

Returns current ROCm gcnArchName, or an empty string on unavailable/error.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### gfx_arch_in

```python
gfx_arch_in(arch_list: Iterable[str]) -> bool
```

Substring-matches any requested gfx name against gcnArchName.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### is_gfx942

```python
is_gfx942() -> bool
```

True on AMD MI300X/CDNA3 gfx942.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### is_gfx950

```python
is_gfx950() -> bool
```

True on AMD MI350/CDNA4 gfx950.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

### supports_float8_fnuz

```python
supports_float8_fnuz(throw_on_hip_incompatibility=True) -> bool
```

Reports whether the active ROCm target uses the FNUZ FP8 flavor.

Returns true for gfx942. gfx950 uses OCP FP8 and returns false unless MSLK_ROCM_FORCE_FP8FNUZ_TYPE is set before import.

**Constraints:**
- The current function docstring claims gfx950 FNUZ support, but implementation and architecture comments say OCP FP8; this reference follows implementation.

Kind: function - Stability: Public - Module: `mslk.utils.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: device, architecture, capability, FP8
Source: [mslk/utils/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/device.py)

## Runtime - Test decorators

### skipUnlessCuda

```python
skipUnlessCuda() -> Callable
```

Skips unless running on a strict NVIDIA CUDA device.

**Returns:** A unittest-style skip decorator.

Kind: function - Stability: Public - Module: `mslk.testing.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: unittest, pytest, skip decorator
Source: [mslk/testing/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py)

### skipUnlessRocm

```python
skipUnlessRocm() -> Callable
```

Skips unless running on a strict AMD ROCm device.

**Returns:** A unittest-style skip decorator.

Kind: function - Stability: Public - Module: `mslk.testing.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: unittest, pytest, skip decorator
Source: [mslk/testing/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py)

### skipUnlessCudaCapability

```python
skipUnlessCudaCapability(major_min, major_max=None, *, minor_min=0) -> Callable
```

Narrows CUDA tests to a capability range and remains transparent on ROCm.

**Returns:** A unittest-style skip decorator.

Kind: function - Stability: Public - Module: `mslk.testing.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: unittest, pytest, skip decorator
Source: [mslk/testing/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py)

### skipUnlessCudaVersion

```python
skipUnlessCudaVersion(major_min: int) -> Callable
```

Narrows CUDA tests to a minimum toolkit major and remains transparent on ROCm.

**Returns:** A unittest-style skip decorator.

Kind: function - Stability: Public - Module: `mslk.testing.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: unittest, pytest, skip decorator
Source: [mslk/testing/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py)

### skipUnlessGfxArch

```python
skipUnlessGfxArch(arch: str, *more_archs: str) -> Callable
```

Narrows ROCm tests to selected gfx targets and remains transparent on CUDA.

**Returns:** A unittest-style skip decorator.

Kind: function - Stability: Public - Module: `mslk.testing.device` - Platforms: NVIDIA, AMD, CPU / Python - Topics: unittest, pytest, skip decorator
Source: [mslk/testing/device.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/testing/device.py)

## Runtime - Package loading & runtime knobs

### open_source

```python
mslk.open_source = True
```

Signals the OSS package variant to tests and integration code.

Kind: constant - Stability: Advanced - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [mslk/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py)

### __version__ / __target__ / __variant__

```python
mslk.__version__: str; mslk.__target__: str; mslk.__variant__: str
```

Generated package-build identity values, with internal/default fallbacks.

Kind: constants - Stability: Advanced - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [mslk/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py)

### load_library_buck

```python
load_library_buck(buck_target: str) -> None
```

Internal/OSS bridge used by domain imports to load split Buck native libraries.

**Constraints:**
- No-ops in Python-only mode and suppresses OSError in OSS.

Kind: function - Stability: Advanced - Module: `mslk.utils.torch.library` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [mslk/utils/torch/library.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/utils/torch/library.py)

### MSLK_PYTHON_ONLY

```python
MSLK_PYTHON_ONLY=1
```

Skips mslk.so native loading/compilation for Python-only development.

**Constraints:**
- Does not provide CPU implementations of GPU kernels.

Kind: environment variable - Stability: Advanced - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [mslk/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py)

### MSLK runtime environment variables

```python
MSLK_CACHE_DIR - MSLK_AUTOTUNE_USE_CUDA_GRAPH - MSLK_AUTOTUNE_COLLECT_STATS - MSLK_ROCM_FORCE_FP8FNUZ_TYPE
```

Runtime cache, autotune, timing-stat, and ROCm FP8 override environment variables.

Cache defaults to $HOME/.mslk. Presence of the autotune variables enables graph benchmarking or timing-stat collection.

Kind: configuration - Stability: Advanced - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [mslk/__init__.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/__init__.py)

### MSLK diagnostic compile-time defines

```python
MSLK_MEMCHECK - MSLK_ISOLATE_KERNEL_LAUNCH - MSLK_TENSORCHECK
```

Preprocessor defines enabling accessor bounds checks, isolated launches, and tensor value checks.

**Constraints:**
- These are compile-time defines, not environment variables read by a running process.

Kind: configuration - Stability: Advanced - Module: `native build` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [include/mslk/utils/kernel_launcher.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh)

### FlyDSL environment variables

```python
FLYDSL_RUNTIME_CACHE_DIR - MSLK_FLYDSL_DISABLE_JIT - FLYDSL_RUNTIME_RUN_ONLY - MSLK_FLYDSL_AOT_WORKERS - COMPILE_ONLY
```

Runtime cache, run-only, worker-count, and compile-only controls for FlyDSL.

Kind: configuration - Stability: Advanced - Module: `mslk.flydsl` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [mslk/flydsl/jit.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/mslk/flydsl/jit.py)

### setup.py build CLI

```python
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]
```

Source-build command-line controls.

**Constraints:**
- 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 - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [setup.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/setup.py)

### MSLK build environment variables

```python
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
```

Package identity, Python-only, toolkit, architecture, and toolchain inputs used by setup.py.

Kind: configuration - Stability: Advanced - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [setup.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/setup.py)

### MSLK CMake knobs

```python
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
```

Direct CMake configuration variables for native builds.

**Constraints:**
- Native code is built as C++20.

Kind: configuration - Stability: Advanced - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [CMakeLists.txt](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/CMakeLists.txt)

### Package dependencies

```python
install_requires=['numpy']; extra: mslk[flash3]; ROCm wheels add pinned FlyDSL
```

Declared runtime dependency boundary and optional kernel packages.

**Constraints:**
- 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 - Module: `mslk` - Platforms: NVIDIA, AMD, CPU / Python - Topics: import, shared library, environment, build
Source: [setup.py](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/setup.py)

## C++ - Numeric & device utilities

### nextPowerOf2

```python
constexpr int64_t nextPowerOf2(int64_t num)
```

Rounds a positive integer up to the next power of two.

Kind: C++ function - Stability: Developer - Module: `namespace mslk` - Platforms: NVIDIA, AMD, CPU - Topics: header, constexpr, device properties
Source: [include/mslk/utils/utils.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils.h)

### roundUp

```python
constexpr int64_t roundUp(int64_t num, int64_t multiple)
```

Rounds num up to a requested multiple.

Kind: C++ function - Stability: Developer - Module: `namespace mslk` - Platforms: NVIDIA, AMD, CPU - Topics: header, constexpr, device properties
Source: [include/mslk/utils/utils.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils.h)

### nextPowerOf2OrRoundUp

```python
constexpr int64_t nextPowerOf2OrRoundUp(int64_t num, int64_t roundUpTo, int64_t threshold)
```

Uses power-of-two rounding below a threshold and fixed-multiple rounding above it.

Kind: C++ function - Stability: Developer - Module: `namespace mslk` - Platforms: NVIDIA, AMD, CPU - Topics: header, constexpr, device properties
Source: [include/mslk/utils/utils.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils.h)

### getDeviceArch

```python
inline int getDeviceArch()
```

Returns cached current-device compute-capability major.

**Constraints:**
- Enforces CUDA runtime 12.8+ when the major is 10 or newer.

Kind: C++ function - Stability: Developer - Module: `namespace mslk` - Platforms: NVIDIA - Topics: header, constexpr, device properties
Source: [include/mslk/utils/utils_gpu.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils_gpu.h)

### getSMCount

```python
inline int64_t getSMCount(int device_index, std::optional<int64_t> num_sms)
```

Returns an override or the available SM count after PyTorch carveout.

**Constraints:**
- Current cached implementation obtains device-0 properties and does not use device_index for that lookup.

Kind: C++ function - Stability: Developer - Module: `namespace mslk` - Platforms: NVIDIA - Topics: header, constexpr, device properties
Source: [include/mslk/utils/utils_gpu.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/utils_gpu.h)

## C++ - CUDA stream & shared memory

### get_device_for_stream

```python
get_device_for_stream(const cudaStream_t& stream) -> device index
```

Resolves the CUDA device associated with a raw stream.

Kind: C++ function - Stability: Developer - Module: `mslk::utils::device` - Platforms: NVIDIA - Topics: CUDA stream, device, dynamic shared memory
Source: [include/mslk/utils/device/cuda_utilities.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/device/cuda_utilities.cuh)

### to_cuda_stream

```python
to_cuda_stream(stream, device_index=-1) -> c10::cuda::CUDAStream
```

Normalizes raw cudaStream_t or c10 stream input to CUDAStream.

Kind: C++ function - Stability: Developer - Module: `mslk::utils::device` - Platforms: NVIDIA - Topics: CUDA stream, device, dynamic shared memory
Source: [include/mslk/utils/device/cuda_utilities.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/device/cuda_utilities.cuh)

### set_gpu_max_dynamic_shared_memory

```python
set_gpu_max_dynamic_shared_memory(kernel, smem_bytes, device=current) -> void
```

Opts a kernel into the requested dynamic shared-memory limit after availability checks.

Kind: C++ function - Stability: Developer - Module: `mslk::utils::device` - Platforms: NVIDIA - Topics: CUDA stream, device, dynamic shared memory
Source: [include/mslk/utils/device/cuda_utilities.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/device/cuda_utilities.cuh)

## C++ - Source context & timing

### source_location

```python
using source_location = std::source_location  // or backport selected by build
```

Portable source-location type used by SourceContext.

Kind: C++ alias - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: diagnostics, source_location, events, benchmark
Source: [include/mslk/utils/source_context.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/source_context.h)

### SourceContext

```python
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
```

Carries source location and a human-readable label through kernel diagnostics.

**Methods / arguments:**
- `description() const -> std::string` - Formats the contextual description.
- `withSummary(summary) -> SourceContext` - Copies context with a new summary.

Kind: C++ type - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: diagnostics, source_location, events, benchmark
Source: [include/mslk/utils/source_context.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/source_context.h)

### SOURCE_CONTEXT_CURRENT

```python
SOURCE_CONTEXT_CURRENT(label)
```

Captures the current file/function/line into a SourceContext.

Kind: C++ macro - Stability: Developer - Module: `global macro` - Platforms: NVIDIA, AMD, CPU - Topics: diagnostics, source_location, events, benchmark
Source: [include/mslk/utils/source_context.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/source_context.h)

### KernelExecutionTimer

```python
KernelExecutionTimer(c10::cuda::CUDAStream stream)
```

Single-use CUDA-event kernel timer with explicit state checks.

**Methods / arguments:**
- `start()` - Records the start event.
- `stop()` - Records the stop event.
- `elapsedMillis() -> float` - Synchronizes the recorded stop event and returns elapsed milliseconds.

Kind: C++ type - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA - Topics: diagnostics, source_location, events, benchmark
Source: [include/mslk/utils/bench/kernel_execution_timer.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/bench/kernel_execution_timer.cuh)

## C++ - Tensor accessors

### TensorAccessor

```python
TensorAccessor<T, N, PtrTraits=DefaultPtrTraits, index_t=int64_t>
```

ATen-compatible tensor accessor with optional named/contextual device bounds assertions.

**Methods / arguments:**
- `using PtrType = typename PtrTraits<T>::PtrType` - Resolved pointer type for the selected pointer trait.
- `TensorAccessor(PtrType data, const index_t* sizes, const index_t* strides, const char* name, const char* context)` - Constructs a contextual accessor over external size/stride storage.
- `operator[](index)` - Indexes a dimension or scalar element.
- `numel() -> size_t` - Returns the addressed storage extent.
- `at(index) -> T&` - Bounds-checked flat access in the enhanced implementation.

Kind: C++ type - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### PackedTensorAccessor

```python
PackedTensorAccessor<T, N, PtrTraits=DefaultPtrTraits, index_t=int64_t>
```

By-value size/stride accessor suitable for passing into GPU kernels.

**Methods / arguments:**
- `using PtrType = typename PtrTraits<T>::PtrType` - Resolved pointer type for the selected pointer trait.
- `PackedTensorAccessor(PtrType data, const index_t* sizes, const index_t* strides, const char* name, const char* context)` - Constructs from native index_t size/stride arrays.
- `template <typename source_index_t=int64_t> PackedTensorAccessor(PtrType data, const source_index_t* sizes, const source_index_t* strides, const char* name, const char* context)` - Converts int64 size/stride arrays when the accessor uses a narrower index type.
- `operator[] / numel() / at()` - Same indexing contract as TensorAccessor.
- `transpose(dim1, dim2)` - Returns a host-side accessor with exchanged sizes/strides.

Kind: C++ type - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### PackedTensorAccessor32

```python
PackedTensorAccessor32<T, N, PtrTraits> = PackedTensorAccessor<..., int32_t>
```

Packed accessor using 32-bit indices.

Kind: C++ alias - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### PackedTensorAccessor64

```python
PackedTensorAccessor64<T, N, PtrTraits> = PackedTensorAccessor<..., int64_t>
```

Packed accessor using 64-bit indices.

Kind: C++ alias - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### DefaultPtrTraits

```python
DefaultPtrTraits<T> = at::DefaultPtrTraits<T>
```

Default accessor pointer trait alias.

Kind: C++ alias - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### RestrictPtrTraits

```python
RestrictPtrTraits<T> = at::RestrictPtrTraits<T>
```

CUDA/HIP restricted-pointer trait alias.

**Constraints:**
- Declared only while compiling with __CUDACC__ or __HIPCC__; it is absent from ordinary host-only C++ compilation.

Kind: C++ alias - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### pta

```python
namespace pta = mslk::utils  // MSLK_MEMCHECK; otherwise at
```

Compile-time selector between enhanced MSLK and standard ATen accessors.

Kind: C++ namespace alias - Stability: Developer - Module: `global namespace` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### overflow_safe_int_t

```python
using overflow_safe_int_t = int64_t
```

Global index type reserved for overflow-safe size arithmetic.

Kind: C++ alias - Stability: Developer - Module: `global namespace` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### ::PackedTensorAccessor

```python
template <typename T, size_t N, template <typename U> class PtrTraits, typename index_t> using PackedTensorAccessor = at::GenericPackedTensorAccessor<T, N, PtrTraits, index_t>
```

Host-build global alias used when MSLK_MEMCHECK is not defined.

**Constraints:**
- This is distinct from mslk::utils::PackedTensorAccessor and exists only in the non-MSLK_MEMCHECK branch.

Kind: C++ alias - Stability: Developer - Module: `global namespace` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

### NAME_MAX_LEN / CONTEXT_MAX_LEN

```python
NAME_MAX_LEN = 32; CONTEXT_MAX_LEN = 256
```

Fixed diagnostic string capacities stored in enhanced accessors.

Kind: C++ constants - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD, CPU - Topics: TensorAccessor, bounds checking, MSLK_MEMCHECK, header
Source: [include/mslk/utils/torch/tensor_accessor.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor.h)

## C++ - Accessor builder & launch

### scalar_type_for

```python
template <typename T> at::ScalarType scalar_type_for()
```

Maps a C++ scalar type to the corresponding ATen ScalarType.

Kind: C++ function - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/torch/tensor_accessor_builder.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h)

### TensorAccessorBuilder

```python
TensorAccessorBuilder<T, N, index_nbits=64, packed=true, PtrTraits=at::DefaultPtrTraits>
```

Defers tensor validation and accessor construction until kernel launch context is known.

**Methods / arguments:**
- `using index_t = conditional_t<index_nbits == 64, int64_t, int32_t>` - Index type selected from index_nbits.
- `using accessor_t = conditional_t<packed, pta::PackedTensorAccessor<...>, pta::TensorAccessor<...>>` - Concrete accessor selected from the packed template flag.
- `TensorAccessorBuilder(const std::string_view& name, const at::Tensor& tensor) noexcept` - Captures a tensor and diagnostic name without copying the tensor.
- `validate_tensor(context)` - Checks rank, scalar type, and 32-bit index bounds.
- `build_ta(context) / build_pta(context)` - Constructs standard or packed accessor.
- `build(context)` - Chooses accessor form from the packed template parameter.
- `checkValues(context)` - Checks the wrapped tensor for NaN/Inf in diagnostic mode.

Kind: C++ type - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/torch/tensor_accessor_builder.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h)

### TA_B / PTA_B

```python
TA_B(tensor, T, N, index_nbits) - PTA_B(tensor, T, N, index_nbits)
```

Creates named non-packed/packed accessor builders while capturing the variable name.

Kind: C++ macros - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/torch/tensor_accessor_builder.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h)

### mslk::utils::pta

```python
namespace mslk::utils { namespace pta = mslk::utils; }  // MSLK_MEMCHECK
namespace mslk::utils { namespace pta = at; }           // otherwise
```

Builder-local selector for enhanced MSLK versus standard ATen accessors.

Kind: C++ namespace alias - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/torch/tensor_accessor_builder.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h)

### MAKE_PTA_WITH_NAME

```python
MAKE_PTA_WITH_NAME(...)
```

Legacy named packed-accessor construction macro.

Kind: C++ macro - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/torch/tensor_accessor_builder.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/torch/tensor_accessor_builder.h)

### is_tensor_accessor_builder

```python
is_tensor_accessor_builder<T>; is_tensor_accessor_builder_v<T>
```

Detects TensorAccessorBuilder template instances.

Kind: C++ trait - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/kernel_launcher.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh)

### transform_kernel_arg

```python
transform_kernel_arg(const SourceContext&, T&& arg)
```

Builds accessor-builder arguments with launch context; forwards other arguments.

Kind: C++ function - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/kernel_launcher.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh)

### check_kernel_arg

```python
check_kernel_arg(const SourceContext&, T&& arg)
```

Runs configured tensor value checks on accessor-builder arguments.

Kind: C++ function - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/kernel_launcher.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh)

### KernelLauncher

```python
KernelLauncher<LaunchCooperative=false, EnableDSA=false, EnableBarrierIsolation=false, EnableNaNChecks=false, EnableExecutionTimer=false>(const SourceContext& context) noexcept
```

Validated normal/cooperative GPU launch wrapper with optional diagnostics and timing.

**Methods / arguments:**
- `launch_kernel(kernel, grid, block, shared_mem, stream, args...)` - Validates dimensions/smem, transforms args, launches, checks errors, and optionally returns elapsed ms.
- `checkGridSizesInRange / checkBlockSizesInRange / checkThreadCountNotExceeded` - Launch-shape validation helpers.
- `checkSharedMemoryPerBlockNotExceeded` - Validates opt-in CUDA or HIP block shared-memory limit.
- `kernelLaunchCheck()` - Raises a contextual launch/DSA failure.

Kind: C++ type - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/kernel_launcher.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh)

### MSLK_LAUNCH_KERNEL

```python
MSLK_LAUNCH_KERNEL(kernel, grid, block, smem, stream, ...)
```

Instantiates KernelLauncher with build-time diagnostic toggles and launches a normal kernel.

Kind: C++ macro - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/kernel_launcher.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh)

### MSLK_LAUNCH_COOPERATIVE_KERNEL

```python
MSLK_LAUNCH_COOPERATIVE_KERNEL(kernel, grid, block, smem, stream, ...)
```

Launches through the cooperative KernelLauncher specialization.

Kind: C++ macro - Stability: Developer - Module: `mslk::utils` - Platforms: NVIDIA, AMD - Topics: kernel launch, validation, NaN check, DSA
Source: [include/mslk/utils/kernel_launcher.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/kernel_launcher.cuh)

## C++ - Autotuning & CUTLASS helpers

### TuningCache

```python
TuningCache(const std::string& kernelName)
```

Persistent per-shape best-kernel cache with optional CUDA-graph benchmarking.

**Methods / arguments:**
- `findBestKernelMaybeAutotune(cache_key, kernel_map, args...) -> Kernel` - Returns a cached winner or benchmarks candidates and persists the result.

**Constraints:**
- 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.

Kind: C++ type - Stability: Developer - Module: `mslk` - Platforms: NVIDIA, AMD - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/utils/tuning_cache.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/utils/tuning_cache.cuh)

### KernelMode

```python
enum class KernelMode { Small, Medium, Large, Default }
```

Coarse GEMM-shape mode used to select CUTLASS kernels.

Kind: C++ enum - Stability: Developer - Module: `mslk` - Platforms: NVIDIA, AMD - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/gemm/cutlass/kernel_mode.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/kernel_mode.h)

### get_kernel_mode

```python
KernelMode get_kernel_mode(at::Tensor XQ, at::Tensor WQ)
```

Classifies a non-batched GEMM into a KernelMode.

Kind: C++ function - Stability: Developer - Module: `mslk` - Platforms: NVIDIA, AMD - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/gemm/cutlass/kernel_mode.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/kernel_mode.h)

### get_batched_kernel_mode

```python
KernelMode get_batched_kernel_mode(at::Tensor XQ, at::Tensor WQ)
```

Classifies a batched GEMM into a KernelMode.

Kind: C++ function - Stability: Developer - Module: `mslk` - Platforms: NVIDIA, AMD - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/gemm/cutlass/kernel_mode.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/kernel_mode.h)

### LinearCombinationOnDevice

```python
template <typename ElementOutput_, int Count, typename ElementAccumulator_=ElementOutput_, typename ElementCompute_=ElementOutput_, ScaleType::Kind Scale=ScaleType::Default, FloatRoundStyle Round=FloatRoundStyle::round_to_nearest> class LinearCombinationOnDevice
```

CUTLASS epilogue linear-combination operator whose scale data lives on device.

**Methods / arguments:**
- `using ElementOutput / ElementAccumulator / ElementCompute` - Public scalar-type aliases for output, accumulator, and compute types.
- `using FragmentOutput / FragmentAccumulator / ComputeFragment` - Count-element CUTLASS Array aliases for each processing stage.
- `using ParamsBase = LinearCombinationParams` - Base parameter-storage alias.
- `LinearCombinationOnDevice(Params const& params)` - Constructs the epilogue functor and retains the device alpha pointer.
- `bool is_source_needed() const` - Reports whether the selected scale mode needs a source fragment.
- `void set_k_partition(int k_partition, int k_partition_count)` - Adjusts beta for serial K-partition reduction.
- `FragmentOutput operator()(FragmentAccumulator const&, FragmentOutput const&) const` - Computes D = alpha x accumulator + beta x source.
- `FragmentOutput operator()(FragmentAccumulator const&) const` - Computes D = alpha x accumulator.

Kind: C++ type - Stability: Developer - Module: `cutlass::epilogue::thread` - Platforms: NVIDIA - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/gemm/cutlass/threadblock.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/threadblock.h)

### LinearCombinationOnDevice::Params

```python
struct LinearCombinationOnDevice<ElementOutput_, Count, ElementAccumulator_, ElementCompute_, Scale, Round>::Params : LinearCombinationParams
```

Host/device-constructable alpha/beta parameter carrier for LinearCombinationOnDevice.

**Methods / arguments:**
- `Params()` - Defaults to alpha=1 and beta=0.
- `Params(ElementCompute alpha, ElementCompute beta)` - Stores scalar alpha and beta values.
- `Params(ElementCompute alpha)` - Stores alpha with beta=0.
- `Params(const ElementCompute* alpha_ptr, const ElementCompute* beta_ptr)` - Uses device pointers for both scales.
- `Params(const ElementCompute* alpha_ptr)` - Uses a device alpha pointer and no beta pointer.
- `Params(ParamsBase const& base)` - Converts CUTLASS base parameter storage.

Kind: C++ type - Stability: Developer - Module: `cutlass::epilogue::thread` - Platforms: NVIDIA - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/gemm/cutlass/threadblock.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/threadblock.h)

### GroupedGemmInputType

```python
enum GroupedGemmInputType { _2D2D, _2D3D }
```

Selects grouped GEMM activation/weight rank layout.

Kind: C++ enum - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/gemm/cutlass/grouped_common.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/grouped_common.cuh)

### set_grouped_gemm_args_kernel

```python
template <typename ProblemShape, typename ElementA, typename ElementB, typename ElementC, typename ScaleDtype, typename StrideA, typename StrideB, typename StrideC, typename LayoutSFA, typename LayoutSFB, typename Sm1xxBlkScaledConfig, typename ElementGlobalScale=float>
__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)
```

Builds per-group CUTLASS pointer/shape/stride argument arrays on device.

Kind: CUDA kernel - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: autotune, cache, CUTLASS, grouped GEMM
Source: [include/mslk/gemm/cutlass/grouped_common.cuh](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/cutlass/grouped_common.cuh)

## C++ - Native GEMM declarations - FP8 & BF16

### mslk::gemm::f8f8bf16_blockwise

```python
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)
```

Native rectangular-block-scaled FP8 GEMM returning BF16.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise_preshuffle

```python
at::Tensor mslk::gemm::f8f8bf16_rowwise_preshuffle(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> bias=std::nullopt, bool use_fast_accum=true)
```

ROCm preshuffled rowwise FP8 GEMM returning BF16.

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 - Module: `mslk::gemm` - Platforms: AMD ROCm - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8f16_rowwise_preshuffle

```python
at::Tensor mslk::gemm::f8f8f16_rowwise_preshuffle(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> bias=std::nullopt, bool use_fast_accum=true)
```

ROCm preshuffled rowwise FP8 GEMM declared with FP16 output.

Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.

**Constraints:**
- The current dispatcher registration mistakenly binds this schema to f8f8bf16_rowwise_preshuffle; treat that as a source defect.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: AMD ROCm - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise_grouped

```python
std::vector<at::Tensor> mslk::gemm::f8f8bf16_rowwise_grouped(at::TensorList XQ, at::TensorList WQ, at::TensorList x_scale, at::TensorList w_scale)
```

Native list-based grouped rowwise FP8 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise_grouped_cat

```python
at::Tensor mslk::gemm::f8f8bf16_rowwise_grouped_cat(at::TensorList XQ, at::TensorList WQ, at::TensorList x_scale, at::TensorList w_scale)
```

Native list-based grouped rowwise FP8 GEMM with concatenated output.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise_grouped_stacked

```python
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)
```

Native stacked grouped rowwise FP8 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise_grouped_dynamic

```python
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)
```

Native dynamic-row grouped FP8 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise_batched

```python
at::Tensor mslk::gemm::f8f8bf16_rowwise_batched(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> bias=std::nullopt, bool use_fast_accum=true, std::optional<at::Tensor> output=std::nullopt)
```

Native batched rowwise FP8 GEMM with optional output reuse.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16bf16bf16_grouped

```python
std::vector<at::Tensor> mslk::gemm::bf16bf16bf16_grouped(at::TensorList X, at::TensorList W)
```

Native list-based grouped BF16 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16bf16bf16_grouped_cat

```python
at::Tensor mslk::gemm::bf16bf16bf16_grouped_cat(at::TensorList X, at::TensorList W)
```

Native list-based grouped BF16 GEMM with concatenated output.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16bf16bf16_grouped_dynamic

```python
at::Tensor mslk::gemm::bf16bf16bf16_grouped_dynamic(at::Tensor X, at::Tensor W, at::Tensor zero_start_index_M)
```

Native dynamic-row grouped BF16 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16bf16bf16_grouped_stacked

```python
at::Tensor mslk::gemm::bf16bf16bf16_grouped_stacked(at::Tensor X, at::Tensor W, at::Tensor M_sizes, std::optional<at::Tensor> out=std::nullopt, std::optional<int64_t> num_sms=std::nullopt)
```

Native stacked grouped BF16 GEMM with output/SM overrides.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise

```python
at::Tensor mslk::gemm::f8f8bf16_rowwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> bias=std::nullopt, bool use_fast_accum=true)
```

Native rowwise FP8 GEMM returning BF16.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_rowwise_out

```python
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<at::Tensor> bias=std::nullopt, bool use_fast_accum=true)
```

Native rowwise FP8 GEMM writing a caller-owned output tensor.

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 - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8f16_rowwise

```python
at::Tensor mslk::gemm::f8f8f16_rowwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> bias=std::nullopt, bool use_fast_accum=true)
```

ROCm native rowwise FP8 GEMM returning FP16.

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 - Module: `mslk::gemm` - Platforms: AMD ROCm - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_groupwise

```python
at::Tensor mslk::gemm::f8f8bf16_groupwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale)
```

Native declaration for K-group-scaled FP8 GEMM.

Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.

**Constraints:**
- The ROCm schema implementation is supplied by Python/Triton.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8f8bf16_groupwise_grouped

```python
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)
```

Native declaration for stacked grouped K-group FP8 GEMM.

Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.

**Constraints:**
- The ROCm schema implementation is supplied by Python/Triton.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::i8i8bf16

```python
at::Tensor mslk::gemm::i8i8bf16(at::Tensor XQ, at::Tensor WQ, double scale, int64_t split_k)
```

Native INT8 GEMM with scalar dequantization scale.

Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.

**Constraints:**
- The dispatcher schema supplies split_k=1; the C++ declaration itself has no default.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::i8i8bf16_dynamic

```python
at::Tensor mslk::gemm::i8i8bf16_dynamic(at::Tensor XQ, at::Tensor WQ, at::Tensor scale, int64_t split_k=1)
```

Native INT8 GEMM with tensor-valued dynamic scale.

Published ATen-native declarations behind the same-name torch.ops.mslk schemas. Prefer the dispatcher contract from Python; these declarations are for native extensions.

**Constraints:**
- The ROCm schema implementation is supplied by Python/Triton.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, GEMM
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

## C++ - Native GEMM declarations - FP4, INT4 & MX

### mslk::gemm::f4f4bf16_grouped_stacked

```python
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<at::Tensor> global_scale=std::nullopt, std::optional<at::Tensor> starting_row_after_padding=std::nullopt, bool use_mx=true)
```

Native stacked grouped MXFP4/NVFP4 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16x9_gemm

```python
at::Tensor mslk::gemm::bf16x9_gemm(at::Tensor A, at::Tensor B, std::optional<at::Tensor> output=std::nullopt)
```

Native cuBLAS BF16x9 emulation GEMM producing FP32.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16i4bf16_shuffled

```python
at::Tensor mslk::gemm::bf16i4bf16_shuffled(at::Tensor X, at::Tensor W, at::Tensor w_scale_group, at::Tensor w_zero_group)
```

Native BF16 x preshuffled INT4 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8i4bf16_shuffled_grouped

```python
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)
```

Native stacked grouped FP8 x preshuffled INT4 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16i4bf16_shuffled_grouped

```python
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)
```

Native stacked grouped BF16 x preshuffled INT4 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16i4bf16_shuffled_batched

```python
at::Tensor mslk::gemm::bf16i4bf16_shuffled_batched(at::Tensor X, at::Tensor WQ, at::Tensor w_scale, at::Tensor w_zp)
```

Native batched BF16 x preshuffled INT4 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16i4bf16_rowwise_batched

```python
at::Tensor mslk::gemm::bf16i4bf16_rowwise_batched(at::Tensor X, at::Tensor WQ, at::Tensor w_scale, at::Tensor w_zp)
```

Native declaration for batched BF16 x rowwise INT4 GEMM.

Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.

**Constraints:**
- The ROCm schema implementation is supplied by Python/Triton.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::bf16i4bf16_rowwise

```python
at::Tensor mslk::gemm::bf16i4bf16_rowwise(at::Tensor X, at::Tensor W, at::Tensor w_scale_group, at::Tensor w_zero_group)
```

Native declaration for BF16 x row/group-quantized INT4 GEMM.

Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.

**Constraints:**
- The ROCm schema implementation is supplied by Python/Triton.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8i4bf16_rowwise

```python
at::Tensor mslk::gemm::f8i4bf16_rowwise(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, at::Tensor w_zp)
```

Native rowwise FP8 activation x INT4 weight GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::f8i4bf16_shuffled

```python
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)
```

Native FP8 activation x preshuffled INT4 weight GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::preshuffle_i4

```python
std::tuple<at::Tensor, at::Tensor> mslk::gemm::preshuffle_i4(at::Tensor WQ, at::Tensor w_scale)
```

Native one-time INT4 weight/scale preprocessing entry.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::mx8mx4bf16

```python
at::Tensor mslk::gemm::mx8mx4bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> output=std::nullopt)
```

Native declaration for MXFP8 x MXFP4 GEMM.

Published ATen-native declarations for packed and microscaled GEMM families. Same-name dispatcher entries above carry the Python-facing shape and scale contracts.

**Constraints:**
- The ROCm schema implementation is supplied by Python/Triton.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::mx8mx6bf16

```python
at::Tensor mslk::gemm::mx8mx6bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> output=std::nullopt)
```

Native MXFP8 x MXFP6 GEMM.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

### mslk::gemm::mx6mx6bf16

```python
at::Tensor mslk::gemm::mx6mx6bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> output=std::nullopt, int64_t splits=0)
```

Native MXFP6 x MXFP6 GEMM with optional split reduction.

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 - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, packed
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)

## C++ - Native offset-grouped GEMM declarations

### mslk::gemm::f8f8bf16_rowwise_grouped_mm

```python
at::Tensor mslk::gemm::f8f8bf16_rowwise_grouped_mm(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> offsets, at::Tensor& output)
```

ROCm native layout-polymorphic grouped rowwise FP8 GEMM.

Published native declarations for offset-described or layout-polymorphic grouped GEMMs.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: AMD ROCm - Topics: ATen, native entrypoint, dispatcher mirror, grouped MM
Source: [include/mslk/gemm/gemm_torch.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h)

### mslk::gemm::mx8mx8bf16_grouped_mm

```python
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<at::Tensor> output=std::nullopt, std::optional<int64_t> actual_num_tokens=std::nullopt)
```

Native offset-grouped MXFP8 x MXFP8 GEMM.

Published native declarations for offset-described or layout-polymorphic grouped GEMMs.

**Constraints:**
- The ROCm dispatcher implementation is supplied by Python/Triton rather than this declaration.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, grouped MM
Source: [include/mslk/gemm/gemm_torch.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h)

### mslk::gemm::f4f4bf16_grouped_mm

```python
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<at::Tensor> output=std::nullopt, std::optional<at::Tensor> global_scale=std::nullopt)
```

Native offset-grouped MXFP4/NVFP4 GEMM.

Published native declarations for offset-described or layout-polymorphic grouped GEMMs.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, grouped MM
Source: [include/mslk/gemm/gemm_torch.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h)

### mslk::gemm::f4f4bf16_ultra_grouped_mm

```python
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<at::Tensor> output=std::nullopt)
```

Native ultra grouped FP4 GEMM with separate global scales.

Published native declarations for offset-described or layout-polymorphic grouped GEMMs.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, grouped MM
Source: [include/mslk/gemm/gemm_torch.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h)

### mslk::gemm::mx8mx4bf16_grouped_mm

```python
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<at::Tensor> output=std::nullopt)
```

Native offset-grouped MXFP8 x MXFP4 GEMM.

Published native declarations for offset-described or layout-polymorphic grouped GEMMs.

**Constraints:**
- The ROCm dispatcher implementation is supplied by Python/Triton rather than this declaration.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, grouped MM
Source: [include/mslk/gemm/gemm_torch.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h)

### mslk::gemm::f4f4bf16

```python
at::Tensor mslk::gemm::f4f4bf16(at::Tensor XQ, at::Tensor WQ, at::Tensor x_scale, at::Tensor w_scale, std::optional<at::Tensor> output=std::nullopt, std::optional<at::Tensor> global_scale=std::nullopt, int64_t mxfp4_block_size=32)
```

Native packed FP4 GEMM selecting MXFP4, MXFP4-16, or NVFP4.

Published native declarations for offset-described or layout-polymorphic grouped GEMMs.

Kind: C++ function - Stability: Developer - Module: `mslk::gemm` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror, grouped MM
Source: [include/mslk/gemm/gemm_torch.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm_torch.h)

## C++ - Native Conv & MoE declarations

### mslk::conv::f8f8bf16_conv

```python
at::Tensor mslk::conv::f8f8bf16_conv(at::Tensor activation, at::Tensor filter, at::Tensor scale, std::vector<int64_t> padding, std::vector<int64_t> stride, std::vector<int64_t> dilation)
```

Native FP8 3D convolution/cross-correlation entry.

Kind: C++ function - Stability: Developer - Module: `mslk::conv` - Platforms: NVIDIA SM100+ - Topics: ATen, native entrypoint, dispatcher mirror
Source: [include/mslk/conv/conv.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/conv/conv.h)

### mslk::moe::index_shuffling_torch

```python
std::tuple<at::Tensor, at::Tensor, at::Tensor> mslk::moe::index_shuffling_torch(const at::Tensor& routing_scores, const std::optional<int64_t>& expert_index_start, const std::optional<int64_t>& expert_index_end, const std::optional<at::Tensor>& valid_token_count, int64_t top_k)
```

Native implementation entry for the index_shuffling dispatcher schema.

Kind: C++ function - Stability: Developer - Module: `mslk::moe` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, dispatcher mirror
Source: [include/mslk/moe/moe.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/moe/moe.h)

### mslk::moe::scatter_add_along_first_dim

```python
void mslk::moe::scatter_add_along_first_dim(at::Tensor dst, at::Tensor src, at::Tensor index)
```

Native in-place first-dimension scatter-add entry.

Kind: C++ function - Stability: Developer - Module: `mslk::moe` - Platforms: NVIDIA - Topics: ATen, native entrypoint, dispatcher mirror
Source: [include/mslk/moe/moe.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/moe/moe.h)

## C++ - Native entry headers

### Native GEMM / Conv / MoE declarations

```python
#include <mslk/gemm/gemm.h>
#include <mslk/gemm/gemm_torch.h>
#include <mslk/conv/conv.h>
#include <mslk/moe/moe.h>
```

ATen-native declarations corresponding to the documented torch.ops schemas.

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 / arguments:**
- `gemm.h / gemm_torch.h` - Native GEMM declarations for every platform-compiled family.
- `conv.h::f8f8bf16_conv(...)` - Native FP8 convolution entry.
- `moe.h::index_shuffling_torch(...)` - Native routing entry.
- `moe.h::scatter_add_along_first_dim(...)` - Native in-place scatter-add entry.

Kind: C++ API family - Stability: Developer - Module: `mslk C++ headers` - Platforms: NVIDIA, AMD - Topics: ATen, native entrypoint, headers
Source: [include/mslk/gemm/gemm.h](https://github.com/meta-pytorch/MSLK/blob/69ae1b897f2d546d72acab129e1ae4bc2924900f/include/mslk/gemm/gemm.h)