MSLK、Transformer 向け GPU カーネルライブラリを公開

Generated from source · CUDA + ROCm

MSLK is a library of fused GPU kernels for transformer workloads: attention, low-precision GEMM, quantization, MoE routing, and convolution. Most of it is reached through torch.ops.mslk.* after import mslk. This page documents every public surface and, more importantly, tells you which one to call.

LatestMSLK 1.3.0 · PyTorch 2.13

NVIDIACUDA 13.0 / 13.2 · SM80 · 90a · 100a · 120a

AMDROCm 7.1 / 7.2 · gfx942

Python3.10 – 3.14

Kernel map

Seven domains.

Pick the workload family you are working on — each card opens the reference filtered to that domain.

[QK

Attention

fMHA dispatch, masks, architecture backends, tree attention, MLA, and paged KV.

156 symbols →](#api-reference)

[AB

GEMM

BF16, FP8, FP4, MX, INT8, INT4, grouped, batched, and mixed-input matrix products.

70 symbols →](#api-reference)

[Q

Quantization

FP8, MXFP4, NVFP4, INT4, and MXFP6 conversion, packing, scaling, and dequantization.

54 symbols →](#api-reference)

[E×

MoE

Routing, shuffling, token gather/scatter, fused activations, and composed expert layers.

27 symbols →](#api-reference)

[3D

Convolution

The low-precision FP8 3D convolution operator.

1 symbol →](#api-reference)

[RT

Runtime

FlyDSL JIT/AOT, device detection, architecture checks, and test decorators.

32 symbols →](#api-reference)

[++

C++

Kernel launch, accessors, tuning, timing, source context, and CUTLASS helpers.

84 symbols →](#api-reference)

Choose a route

The shortest path to the right kernel.

01 · TRANSFORMER CORE

Fused attention

Begin with automatic dispatch. Reach for explicit backends only when you need a specific architecture, paged KV layout, split-K, or deterministic behavior.

Start with memory_efficient_attention →

02 · LINEAR LAYERS

Quantize, then GEMM

Choose a scale granularity that matches the GEMM family: tensor, row, block, group, MXFP4, NVFP4, or packed INT4.

See the FP8 pipeline →

03 · SPARSE MODELS

Route, gather, compute, scatter

MSLK exposes the routing pieces independently and also includes baseline and Meta-shuffling MoE layers for composed execution.

Open the routing API →

Quick start

From install to output.

These are deliberately small, copyable paths through the major public surfaces.

# CUDA 13.0 wheel
pip install mslk --index-url https://download.pytorch.org/whl/cu130

# ROCm 7.1 wheel
pip install mslk \
  --index-url https://download.pytorch.org/whl/rocm7.1/ \
  --extra-index-url https://pypi.org/simple

Import registers operators

import mslk loads mslk.so. Import a domain such as mslk.gemm or mslk.moe before calling its torch.ops.mslk entries so Python-side registrations are installed.

import torch
from mslk.attention import fmha

B, M, H, K = 2, 2048, 32, 128
q = torch.randn(B, M, H, K, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)

out = fmha.memory_efficient_attention(
    q, k, v,
    attn_bias=fmha.LowerTriangularMask(),
)
# out: [B, M, H, K]

Let dispatch work

Automatic dispatch evaluates the input dtype, head dimension, mask, dropout, gradient requirements, and hardware. Supply op=(FwOp, BwOp) only when deliberately pinning a backend.

import torch
import mslk.gemm
from mslk.quantize.triton.fp8_quantize import quantize_fp8_row

x = torch.randn(1024, 4096, device="cuda", dtype=torch.bfloat16)
w = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
xq, x_scale = quantize_fp8_row(x)
wq, w_scale = quantize_fp8_row(w)

out = torch.ops.mslk.f8f8bf16_rowwise(
    xq, wq, x_scale, w_scale
)
# Conceptually: dequant(xq) @ dequant(wq).T → BF16

Weights are N × K

Most MSLK GEMMs take activations [M,K] and weights [N,K], then compute X @ W.T. Keep scale layout paired with the quantizer that produced it.

import torch
from mslk.attention.fmha.merge_training import (
    memory_efficient_attention_partial_autograd,
    merge_attentions_autograd,
)

B, Mq, Mkv, H, K = 1, 128, 1024, 16, 128
q = torch.randn(B, Mq, H, K, device="cuda", dtype=torch.bfloat16)
k = torch.randn(B, Mkv, H, K, device="cuda", dtype=torch.bfloat16)
v = torch.randn_like(k)
k0, k1 = k.chunk(2, dim=1)
v0, v1 = v.chunk(2, dim=1)
p0 = memory_efficient_attention_partial_autograd(q, k0, v0)
p1 = memory_efficient_attention_partial_autograd(q, k1, v1)
out = merge_attentions_autograd(p0, p1)

Exact softmax merge

Each partial carries its output and log-sum-exp. The merge reweights chunks mathematically, so it is equivalent to attention over the concatenated K/V sequence.

import torch
import mslk.moe

T, D, E = 256, 512, 8
x = torch.randn(T, D, device="cuda", dtype=torch.bfloat16)
routing_scores = torch.softmax(
    torch.randn(T, E, device="cuda"), dim=-1
)
counts, experts, tokens = torch.ops.mslk.index_shuffling(
    routing_scores, top_k=1
)
expert_x = mslk.moe.gather_scale_dense_tokens(
    x, tokens, experts, routing_scores
)
# Replace this identity with grouped expert GEMMs + silu_mul.
expert_y = expert_x
out = torch.zeros_like(x)
mslk.moe.scatter_add_dense_tokens(out, expert_y, tokens)

Composable routing pieces

The low-level API makes data movement explicit. For a composed module, use BaselineMoE or the top-1-only MetaShufflingMoE.

Conventions

How to read a shape in these docs.

Every signature on this page describes tensors with the same axis letters. Learn them once and the 400 entries below stop needing individual explanation. One rule holds nearly everywhere: the last dimension must have stride 1, even when the others are non-contiguous.

Attention

q, k, v*[B, M, H, K]**the usual case*

q, k, v*[B, M, G, H, K]**GQA/MQA, experimental — you expand K/V yourself*

out*[B, M, H, Kv]**same layout as q, last axis from V*

Bbatch

Msequence

Ghead groups

Hheads

Khead dim

Variable-length batches are packed into B=1 with sequence metadata carried by the mask instead.

GEMM

x*[M, K]**activations*

w*[N, K]**weights — N×K, not K×N*

out*[M, N]**computed as x @ w.T*

Mtokens

Kreduction dim

Noutput features

Grouped variants keep this layout and add a group description: a list of tensors, a leading expert axis, M_sizes alongside concatenated tokens, or offsets. Output is BF16 unless the op name says f16.

Quantization scales

tensorwise*[1]**one scale for the whole tensor*

rowwise*[M]**one per row — the common FP8 path*

blockwise*[⌈M/Bm⌉, ⌈K/Bk⌉]**one per Bm×Bk tile*

A quantized tensor is the packed data *plus* its scales — a GEMM only accepts the granularity it was written for, so keep each scale tensor with the quantizer that produced it. MX formats add E8M0 block exponents whose layout differs between CUDA and ROCm; those buffers are not interchangeable.

MoE routing

scores*[T, E]**router output*

indices*[T × top_k]**token and expert index pairs*

counts*[E + 2]**tokens routed per expert*

Ttokens

Eexperts

Routing order stays explicit rather than hidden inside a fused layer, which is what lets the expert GEMM run as one grouped call over contiguous segments.

Support matrices

What runs where.

Two orientation maps for the choices that actually block you: which attention backend can serve your case, and which GEMM op matches the dtypes you already have. Both are read from source, and neither replaces the runtime checks — exact shapes, masks, and the archs compiled into your wheel still decide.

Attention — backends behind memory_efficient_attention

This table is only about attention. Every row is a forward/backward operator class under mslk.attention.fmha. Leave op=None and dispatch picks one for you; pass op=(FwOp, BwOp) when you need a specific one. GEMM, MoE and quantization ops do not dispatch through this.

GEMM — picking a low-precision op

Nothing here is auto-selected: you call the op that matches the dtypes you already hold, and it is your job to hand it scales in the exact granularity it expects. Names encode the contract — f8f8bf16_rowwise is FP8 in × FP8 in → BF16 out, one scale per row. An op always exists after import; it raises at call time if your wheel has no kernel for the arch.

Complete reference

Filter by symbol, concept (“paged”, “rowwise”), module, dtype or platform. Open an entry for its signature, behavior, support contract, caveats, and source.

What counts as documented here

Included: exported and directly callable Python APIs, registered dispatcher schemas, selectable backend classes, integration-level raw ops, and published C++ headers. Excluded: underscore-only kernel bodies, Meta/fake implementations, benchmarks, and test-only reference routines — unless they expose a documented integration contract.

Showing 424 of 424 documented symbols.

memory_efficient_attention#

mslk.attention.fmha

Autograd-enabled fused attention with automatic backend dispatch.

PublicfunctionNVIDIAAMD

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

Accepts [B,M,K], [B,M,H,K], or experimental [B,M,G,H,K] inputs. It computes scaled QKᵀ, applies an optimized bias or dense tensor bias, softmax/dropout, then multiplies V without materializing the full attention matrix.

Returns: Attention output in the query layout, with its last dimension replaced by the V head dimension.

Methods / constructors

op=(FwOp, BwOp)Explicitly pins forward and backward operator classes.

attn_bias=AttentionBias | Tensor | NoneUses optimized masks when possible; arbitrary dense tensors are generally slower.

Notes & 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.

GroupAttention / Core fMHA

SupportNVIDIA · AMD · MTIA

TopicsfMHA · dispatch · BMK · BMHK · BMGHK

Sourcemslk/attention/fmha/__init__.py

memory_efficient_attention_forward#

mslk.attention.fmha

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

PublicfunctionNVIDIAAMD

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

Validates and normalizes the same inputs as the high-level API, selects a forward backend, and returns only the attention output.

Returns: Attention output tensor.

Notes & constraints

• Pass a forward operator class, not an operator pair, through op.
• Use this path when no gradients or LSE are needed.

GroupAttention / Core fMHA

SupportNVIDIA · AMD · MTIA

TopicsfMHA · dispatch · BMK · BMHK · BMGHK

Sourcemslk/attention/fmha/__init__.py

memory_efficient_attention_forward_requires_grad#

mslk.attention.fmha

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

PublicfunctionNVIDIAAMD

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]

Despite its name, this is the explicit non-autograd API. The returned LSE is the context consumed by memory_efficient_attention_backward.

Returns: (output, lse).

Notes & constraints

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

GroupAttention / Core fMHA

SupportNVIDIA · AMD · MTIA

TopicsfMHA · dispatch · BMK · BMHK · BMGHK

Sourcemslk/attention/fmha/__init__.py

memory_efficient_attention_backward#

mslk.attention.fmha

Explicit backward pass for an output/LSE pair.

PublicfunctionNVIDIAAMD

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

Consumes the original Q/K/V, forward output, LSE, and upstream gradient. Dispatches a compatible backward kernel unless one is supplied.

Returns: (dq, dk, dv) reshaped to the original inputs.

Notes & constraints

• Dropout is not supported on the manual forward/backward API.
• This return contract does not include a dense tensor-bias gradient.

GroupAttention / Core fMHA

SupportNVIDIA · AMD · MTIA

TopicsfMHA · dispatch · BMK · BMHK · BMGHK

Sourcemslk/attention/fmha/__init__.py

memory_efficient_attention_partial#

mslk.attention.fmha

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

PublicfunctionNVIDIAAMD

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]

Call it repeatedly with the same query and separate key/value chunks, then combine the results with merge_attentions. This enables exact split-KV attention.

Returns: (partial_output, partial_lse).

Notes & 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.

GroupAttention / Core fMHA

SupportNVIDIA · AMD · MTIA

TopicsfMHA · dispatch · BMK · BMHK · BMGHK

Sourcemslk/attention/fmha/__init__.py

merge_attentions#

mslk.attention.fmha

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

PublicfunctionNVIDIAAMD

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

Accepts sequences of chunk tensors or tensors stacked along a leading chunk axis. It computes the exact normalized output over the union of chunks rather than averaging partial outputs.

Returns: (merged_output, merged_lse) when write_lse=True; otherwise (merged_output, None).

Notes & constraints

• Direct autograd through this function is unsupported; use merge_attentions_autograd.
• Inputs requiring gradients require write_lse=True.

GroupAttention / Core fMHA

SupportNVIDIA · AMD · MTIA

TopicsfMHA · dispatch · BMK · BMHK · BMGHK

Sourcemslk/attention/fmha/__init__.py

torch.ops.mslk.memory_efficient_attention_forward#

torch.ops.mslk

Compile-friendly forward attention with optional bias.

Low-leveltorch opNVIDIAMeta / fake

torch.ops.mslk.memory_efficient_attention_forward(q, k, v, b=None, p=0.0, scale=None) -> 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.

Notes & constraints

• No explicit op or output_dtype argument.
• K and V head dimensions must match.
• Only selected bias forms are trace-safe.

GroupAttention / Compile-friendly wrappers

SupportNVIDIA · Meta / fake

Topicstorch.compile · custom op · forward

Sourcemslk/attention/fmha/__init__.py

torch.ops.mslk.memory_efficient_attention_forward_with_bias#

torch.ops.mslk

Compile-friendly forward attention requiring an explicit bias tensor.

Low-leveltorch opNVIDIAMeta / fake

torch.ops.mslk.memory_efficient_attention_forward_with_bias(q, k, v, b, p=0.0, scale=None) -> 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.

Notes & constraints

• No explicit op or output_dtype argument.
• K and V head dimensions must match.
• Only selected bias forms are trace-safe.

GroupAttention / Compile-friendly wrappers

SupportNVIDIA · Meta / fake

Topicstorch.compile · custom op · forward

Sourcemslk/attention/fmha/__init__.py

Partial#

mslk.attention.fmha.merge_training

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

PublicclassNVIDIAAMD

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

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 / constructors

is_bmghk() -> boolReports whether the stored output uses the grouped-head BMGHK layout.

apply(fn: Callable[[Tensor], Tensor]) -> PartialApplies a layout-preserving transform that must not manipulate the embedding dimension.

GroupAttention / Trainable partial attention

SupportNVIDIA · AMD

Topicsautograd · split KV · LSE

Sourcemslk/attention/fmha/merge_training.py

memory_efficient_attention_partial_autograd#

mslk.attention.fmha.merge_training

Autograd-safe wrapper around partial attention.

PublicfunctionNVIDIAAMD

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

Records Q/K/V and uses hidden placeholder gradients to route the merged output and full LSE into each shard’s backward pass.

Returns: Partial.

Notes & constraints

• Arguments mirror memory_efficient_attention_partial.
• Dropout remains unsupported by the underlying partial path.

GroupAttention / Trainable partial attention

SupportNVIDIA · AMD

Topicsautograd · split KV · LSE

Sourcemslk/attention/fmha/merge_training.py

merge_attentions_autograd#

mslk.attention.fmha.merge_training

Merges one or more Partial objects with correct gradients.

PublicfunctionNVIDIAAMD

merge_attentions_autograd(*partials: Partial) -> Tensor

For multiple shards it performs the exact LSE merge. Passing a single Partial returns its output while retaining the autograd plumbing.

Returns: Merged attention output.

Notes & constraints

• Raises ValueError when called without partials.

GroupAttention / Trainable partial attention

SupportNVIDIA · AMD

Topicsautograd · split KV · LSE

Sourcemslk/attention/fmha/merge_training.py

AttentionBias#

mslk.attention.fmha.attn_bias

Abstract base for optimized attention-bias descriptions.

PublicclassNVIDIAAMD

AttentionBias()

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

Returns: An object accepted as attn_bias by fMHA.

Methods / constructors

materialize(shape, dtype=torch.float32, device='cpu') -> TensorBuilds the dense bias for debugging/reference checks; intentionally slow.

to(device)Moves stored sequence metadata or tensor bias to a device.

Notes & constraints

• Subclasses define materialize and, where state is stored, to(device).

GroupAttention / Biases & masks

SupportNVIDIA · AMD · MTIA · CPU / Python

Topicsattention bias · mask

Sourcemslk/attention/fmha/attn_bias.py

LowerTriangularMask#

mslk.attention.fmha.attn_bias

Top-left-aligned causal mask.

PublicclassNVIDIAAMD

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

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

Returns: An object accepted as attn_bias by fMHA.

Methods / constructors

add_bias(bias: Tensor) -> LowerTriangularMaskWithTensorBiasCombines causal masking with an arbitrary additive tensor.

materialize(shape, dtype=torch.float32, device='cpu') -> TensorProduces zeros on/below the diagonal and −∞ above it.

Notes & constraints

• The device constructor argument is retained only for compatibility and is ignored.

GroupAttention / Biases & masks

SupportNVIDIA · AMD · MTIA · CPU / Python

Topicsattention bias · mask

Sourcemslk/attention/fmha/attn_bias.py

LowerTriangularMaskWithTensorBias#

mslk.attention.fmha.attn_bias

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

PublicclassNVIDIAAMD

LowerTriangularMaskWithTensorBias(bias: Tensor)

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

Returns: An object accepted as attn_bias by fMHA.

Methods / constructors

materialize(shape, dtype=torch.float32, device='cpu') -> TensorBuilds the dense bias for debugging/reference checks; intentionally slow.

to(device)Moves stored sequence metadata or tensor bias to a device.

Notes & constraints

• Backend support is narrower than for a pure causal mask.

GroupAttention / Biases & masks

SupportNVIDIA · AMD · MTIA · CPU / Python

Topicsattention bias · mask

Sourcemslk/attention/fmha/attn_bias.py

LowerTriangularFromBottomRightMask#

mslk.attention.fmha.attn_bias

Bottom-right-aligned causal mask for Mq ≠ Mkv.

PublicclassNVIDIAAMD

LowerTriangularFromBottomRightMask()

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

Returns: An object accepted as attn_bias by fMHA.

Methods / constructors

make_local_attention(window_size: int) -> LowerTriangularFromBottomRightLocalAttentionMaskRestricts the causal region to a trailing window.

GroupAttention / Biases & masks

SupportNVIDIA · AMD · MTIA · CPU / Python

Topicsattention bias · mask

Sourcemslk/attention/fmha/attn_bias.py

LowerTriangularFromBottomRightLocalAttentionMask#

mslk.attention.fmha.attn_bias

Bottom-right causal sliding-window mask.

PublicclassNVIDIAAMD

LowerTriangularFromBottomRightLocalAttentionMask(_window_size: int)

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

Returns: An object accepted as attn_bias by fMHA.

Methods / constructors

materialize(shape, dtype=torch.float32, device='cpu') -> TensorBuilds the dense bias for debugging/reference checks; intentionally slow.

to(device)Moves stored sequence metadata or tensor bias to a device.

GroupAttention / Biases & masks

SupportNVIDIA · AMD