feat: vendor MXFP8 MoE runtime

README.md

@@ -9,10 +9,11 @@ CUDA kernels for Prime Intellect training stacks, shipped as one wheel, `prime-k
     ├── kernels.toml          # the manifest: one table per kernel
     ├── __init__.py           # registry: is_available / load / status
     ├── _spec.py              # manifest parser (build time + runtime)
-    ├── flash_moe/            # one folder per kernel
+    ├── flash_moe/            # compiled kernel
     │   ├── __init__.py       # Python surface: op wrappers, fake tensors
     │   ├── mxfp8.py
     │   └── csrc/             # the C++/CUDA sources compiled into prime_kernels.flash_moe._C
+    ├── mxfp8_moe/            # Python-only MXFP8 MoE runtime kernels
     └── rmsnorm/
         ├── __init__.py
         ├── csrc/             # the torch binding
@@ -22,7 +23,8 @@ CUDA kernels for Prime Intellect training stacks, shipped as one wheel, `prime-k
 
 The repo root is the wheel: `setup.py` and `pyproject.toml` sit here, and `prime_kernels/`
 is the package you import. A kernel folder holds both halves of one kernel — its Python
-surface and, under `csrc/`, the sources compiled into `prime_kernels.<name>._C`.
+surface and, for compiled kernels, the sources under `csrc/` compiled into
+`prime_kernels.<name>._C`.
 
 This repo is consumed as a git submodule at `deps/prime-kernels/` in
 [prime-rl](https://github.com/PrimeIntellect-ai/prime-rl), which builds and publishes the
@@ -48,10 +50,9 @@ kernel, and returns the scales already in the blocked layout a tensor core GEMM
 its sources are committed for now — its table in `kernels.toml` is commented out, so it is
 neither built nor shipped in the wheel, and the registry does not list it.
 
-`flash_moe` is used by prime-rl's MoE layers under `model.moe_fused_kernel=true`, which
-resolves the kernel during model setup so an unusable install fails before training starts.
-It picks `fused_moe_mxfp8` when the run also quantizes the experts to MXFP8 and
-`fused_moe_bf16` otherwise.
+`flash_moe` is currently dormant in prime-rl. `mxfp8_moe` provides differentiable MXFP8
+grouped GEMM and MXFP8 expert-parallel transport. It is registered as Python-only because
+it orchestrates PyTorch and torchao kernels rather than compiling a `_C` extension here.
 
 ## Installing
 
@@ -94,6 +95,10 @@ cxx-std = 20
    treats the op as non-differentiable. `flash_moe` is the exception — it is forward only,
    and prime-rl wraps it in its own `autograd.Function`.
 
+For a Python-only kernel, set `python-only = true`, omit `ops` and `sources`, and expose
+the differentiable Python surface from `__init__.py`. Optional import requirements belong
+in the manifest's `requires` list so `is_available()` fails during setup.
+
 Whatever the kernel requires of its inputs — block sizes, alignments, layouts — belongs
 here, not in the caller: `TORCH_CHECK` it in the binding, and export the constants
 (e.g. `flash_moe.BLOCK_M`) and any setup-time predicate (`unsupported_shape_reason`) from

prime_kernels/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib
+import importlib.util
 from importlib.machinery import EXTENSION_SUFFIXES
 from pathlib import Path
 from types import ModuleType
@@ -31,7 +32,10 @@ def spec(name: str) -> KernelSpec:
 
 
 def is_built(name: str) -> bool:
-    directory = spec(name).path
+    kernel = spec(name)
+    directory = kernel.path
+    if kernel.python_only:
+        return (directory / "__init__.py").is_file()
     return any((directory / f"_C{suffix}").exists() for suffix in EXTENSION_SUFFIXES)
 
 
@@ -41,6 +45,9 @@ def unavailable_reason(name: str, device: int | None = None) -> str | None:
     kernel = spec(name)
     if not is_built(name):
         return f"{name} was not compiled into this install of prime-kernels"
+    missing = [requirement for requirement in kernel.requires if importlib.util.find_spec(requirement) is None]
+    if missing:
+        return f"{name} requires {', '.join(missing)}, which is not installed"
     if not torch.cuda.is_available():
         return f"{name} requires a CUDA device ({kernel.sm_list}), none is available"
     capability = torch.cuda.get_device_capability(device)

prime_kernels/_spec.py

@@ -57,6 +57,8 @@ class KernelSpec:
     cxx_std: int
     cxx_flags: tuple[str, ...]
     nvcc_flags: tuple[str, ...]
+    python_only: bool
+    requires: tuple[str, ...]
 
     @property
     def module(self) -> str:
@@ -82,11 +84,13 @@ def _kernel(name: str, path: Path, table: dict) -> KernelSpec:
         name=name,
         path=path,
         description=table["description"],
-        ops=table["ops"],
-        sources=tuple(path / source for source in table["sources"]),
+        ops=table.get("ops", ""),
+        sources=tuple(path / source for source in table.get("sources", [])),
         include_dirs=tuple(path / directory for directory in table.get("include-dirs", [])),
         archs=tuple(Arch.parse(arch) for arch in table["arch"]),
         cxx_std=table.get("cxx-std", 20),
         cxx_flags=tuple(table.get("cxx-flags", [])),
         nvcc_flags=tuple(table.get("nvcc-flags", [])),
+        python_only=table.get("python-only", False),
+        requires=tuple(table.get("requires", [])),
     )

prime_kernels/kernels.toml

@@ -8,6 +8,9 @@
 #   sources      C++/CUDA files to compile into prime_kernels.<name>._C
 #   arch         compute capabilities to compile for; matched exactly against the device
 #   cxx-std      C++ standard
+#
+# Vendored Python kernels use `python-only = true`, may list import checks in `requires`,
+# and omit `ops`, `sources`, and compilation settings.
 
 [flash_moe]
 description = "Fused MoE forward (bf16 + mxfp8) on Blackwell tcgen05"
@@ -21,6 +24,12 @@ include-dirs = ["csrc/kernels"]
 arch = ["10.0a"]
 cxx-std = 20
 
+[mxfp8_moe]
+description = "MXFP8 grouped GEMM and expert-parallel transport"
+python-only = true
+requires = ["torchao"]
+arch = ["10.0"]
+
 # rmsnorm is not built yet: only its sources are committed. Uncomment the table below to
 # put it back into the build (and to make the registry report on it).
 #

prime_kernels/mxfp8_moe/LICENSE.torchao

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

prime_kernels/mxfp8_moe/__init__.py

@@ -0,0 +1,9 @@
+from .grouped_gemm import TOKEN_GROUP_ALIGNMENT, grouped_gemm
+from .transport import all_to_all_combine, all_to_all_dispatch
+
+__all__ = [
+    "TOKEN_GROUP_ALIGNMENT",
+    "all_to_all_combine",
+    "all_to_all_dispatch",
+    "grouped_gemm",
+]

prime_kernels/mxfp8_moe/grouped_gemm.py

@@ -0,0 +1,201 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license in
+# LICENSE.torchao. Derived from torchao commit 02105d46c.
+
+from __future__ import annotations
+
+import torch
+from torchao.prototype.moe_training.kernels.mxfp8 import (
+    mx_block_rearrange_2d_M_groups_cuda,
+    mxfp8_quantize_cuda_3d,
+    triton_mx_block_rearrange_2d_K_groups,
+    triton_mx_block_rearrange_2d_M_groups,
+    triton_mx_block_rearrange_per_group_3d,
+)
+from torchao.prototype.mx_formats.config import MXFP8Dim1CastKernelChoice, ScaleCalculationMode
+from torchao.prototype.mx_formats.kernels import triton_mxfp8_dequant_dim0, triton_to_mxfp8_dim0
+from torchao.prototype.mx_formats.utils import _to_mxfp8_dim1_kernel_wrapper
+from torchao.quantization.quantize_.common import KernelPreference
+
+TOKEN_GROUP_ALIGNMENT = 32
+_CUDA_REARRANGE_MAX_GROUPS = 32
+_QUANT_NUMEL_LIMIT = 1 << 31
+_QUANT_CHUNK_NUMEL = 1 << 30
+_SCALING_MODE = ScaleCalculationMode.RCEIL
+
+
+def _quantize_rows(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if x.ndim != 2 or x.numel() < _QUANT_NUMEL_LIMIT:
+        return triton_to_mxfp8_dim0(
+            x,
+            inner_block_size=TOKEN_GROUP_ALIGNMENT,
+            scaling_mode=_SCALING_MODE.value.lower(),
+        )
+
+    rows_per_chunk = max(1, _QUANT_CHUNK_NUMEL // x.shape[1])
+    quantized = [
+        triton_to_mxfp8_dim0(
+            chunk,
+            inner_block_size=TOKEN_GROUP_ALIGNMENT,
+            scaling_mode=_SCALING_MODE.value.lower(),
+        )
+        for chunk in x.split(rows_per_chunk, dim=0)
+    ]
+    return torch.cat([data for data, _ in quantized]), torch.cat([scales for _, scales in quantized])
+
+
+def quantize_rows(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if x.ndim != 2:
+        raise ValueError(f"MXFP8 row quantization expects a 2D tensor, got shape {tuple(x.shape)}")
+    if x.shape[1] % TOKEN_GROUP_ALIGNMENT:
+        raise ValueError(
+            f"MXFP8 row quantization requires columns divisible by {TOKEN_GROUP_ALIGNMENT}, got {x.shape[1]}"
+        )
+    if x.shape[0] == 0:
+        data = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+        scales = torch.empty(
+            (0, x.shape[1] // TOKEN_GROUP_ALIGNMENT),
+            dtype=torch.float8_e8m0fnu,
+            device=x.device,
+        )
+        return data, scales
+    return _quantize_rows(x)
+
+
+def dequantize_rows(data: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    if data.shape[0] == 0:
+        return torch.empty(data.shape, dtype=dtype, device=data.device)
+    return triton_mxfp8_dequant_dim0(
+        data,
+        scales.view(torch.uint8),
+        out_dtype=dtype,
+        scale_block_size=TOKEN_GROUP_ALIGNMENT,
+    )
+
+
+def _rearrange_token_scales(scales: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
+    if offsets.numel() > _CUDA_REARRANGE_MAX_GROUPS:
+        return triton_mx_block_rearrange_2d_M_groups(scales, offsets)
+    return mx_block_rearrange_2d_M_groups_cuda(scales, offsets)
+
+
+def _quantize_dim1(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    mx = _to_mxfp8_dim1_kernel_wrapper(
+        x,
+        TOKEN_GROUP_ALIGNMENT,
+        elem_dtype=torch.float8_e4m3fn,
+        hp_dtype=x.dtype,
+        kernel_preference=KernelPreference.AUTO,
+        cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA,
+        scale_calculation_mode=_SCALING_MODE,
+    )
+    return mx.qdata, mx.scale
+
+
+def _forward(input_act: torch.Tensor, weight_t: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
+    input_data, input_scales = quantize_rows(input_act)
+    weight_data, weight_scales = _quantize_rows(weight_t.transpose(-2, -1))
+    return torch._scaled_grouped_mm(
+        input_data,
+        weight_data.transpose(-2, -1),
+        _rearrange_token_scales(input_scales, offsets),
+        triton_mx_block_rearrange_per_group_3d(weight_scales),
+        offs=offsets,
+        out_dtype=torch.bfloat16,
+    )
+
+
+def _dgrad(grad_output: torch.Tensor, weight_t: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
+    grad_data, grad_scales = quantize_rows(grad_output)
+    weight_data, weight_scales = mxfp8_quantize_cuda_3d(
+        weight_t.transpose(-2, -1),
+        TOKEN_GROUP_ALIGNMENT,
+        scaling_mode=_SCALING_MODE.value.lower(),
+    )
+    return torch._scaled_grouped_mm(
+        grad_data,
+        weight_data,
+        _rearrange_token_scales(grad_scales, offsets),
+        weight_scales,
+        offs=offsets,
+        out_dtype=torch.bfloat16,
+    )
+
+
+def _wgrad(
+    grad_output: torch.Tensor,
+    input_act: torch.Tensor,
+    offsets: torch.Tensor,
+    *,
+    high_precision: bool,
+) -> torch.Tensor:
+    if high_precision:
+        grad_weight = torch._grouped_mm(
+            grad_output.transpose(-2, -1),
+            input_act,
+            offs=offsets,
+            out_dtype=torch.bfloat16,
+        )
+        return grad_weight.transpose(-2, -1)
+
+    grad_data, grad_scales = _quantize_dim1(grad_output)
+    input_data, input_scales = _quantize_dim1(input_act)
+    scale_offsets = offsets // TOKEN_GROUP_ALIGNMENT
+    grad_weight = torch._scaled_grouped_mm(
+        grad_data,
+        input_data.transpose(-2, -1),
+        triton_mx_block_rearrange_2d_K_groups(grad_scales, scale_offsets),
+        triton_mx_block_rearrange_2d_K_groups(input_scales, scale_offsets),
+        offs=offsets,
+        out_dtype=torch.bfloat16,
+    )
+    return grad_weight.transpose(-2, -1)
+
+
+class _GroupedGemm(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_act: torch.Tensor,
+        weight_t: torch.Tensor,
+        offsets: torch.Tensor,
+        high_precision_wgrad: bool,
+    ) -> torch.Tensor:
+        ctx.save_for_backward(input_act, weight_t, offsets)
+        ctx.high_precision_wgrad = high_precision_wgrad
+        return _forward(input_act, weight_t, offsets)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        input_act, weight_t, offsets = ctx.saved_tensors
+        grad_input = _dgrad(grad_output, weight_t, offsets) if ctx.needs_input_grad[0] else None
+        grad_weight = (
+            _wgrad(
+                grad_output,
+                input_act,
+                offsets,
+                high_precision=ctx.high_precision_wgrad,
+            )
+            if ctx.needs_input_grad[1]
+            else None
+        )
+        return grad_input, grad_weight, None, None
+
+
+def grouped_gemm(
+    input_act: torch.Tensor,
+    weight_t: torch.Tensor,
+    offsets: torch.Tensor,
+    *,
+    high_precision_wgrad: bool = False,
+) -> torch.Tensor:
+    """Run differentiable MXFP8 grouped GEMM on already aligned token groups."""
+    if input_act.ndim != 2 or weight_t.ndim != 3:
+        raise ValueError(
+            f"MXFP8 grouped GEMM expects 2D activations and 3D weights, got {input_act.ndim}D and {weight_t.ndim}D"
+        )
+    if offsets.dtype != torch.int32:
+        raise ValueError(f"MXFP8 grouped GEMM offsets must be int32, got {offsets.dtype}")
+    return _GroupedGemm.apply(input_act, weight_t, offsets, high_precision_wgrad)