feat: add NVFP4 MoE kernels with optional 4over6 quantization

MANIFEST.in

@@ -0,0 +1,2 @@
+include build_backend.py
+recursive-include prime_kernels/nvfp4_moe *.py *.md *.cpp *.cu *.cuh *.h MSLK_LICENSE

build_backend.py

@@ -0,0 +1,31 @@
+"""Setuptools builds with dependencies selected from the kernel manifest."""
+
+import os
+import tomllib
+from pathlib import Path
+
+from setuptools.build_meta import build_editable as build_editable
+from setuptools.build_meta import build_sdist as build_sdist
+from setuptools.build_meta import build_wheel as build_wheel
+from setuptools.build_meta import prepare_metadata_for_build_editable as prepare_metadata_for_build_editable
+from setuptools.build_meta import prepare_metadata_for_build_wheel as prepare_metadata_for_build_wheel
+
+
+def get_requires_for_build_wheel(config_settings=None):
+    manifest = tomllib.loads((Path(__file__).parent / "prime_kernels/kernels.toml").read_text())
+    selected = {name for name in os.environ.get("PRIME_KERNELS", "").split(",") if name}
+    unknown = selected - manifest.keys()
+    if unknown:
+        raise ValueError(f"PRIME_KERNELS contains unknown kernel(s): {', '.join(sorted(unknown))}")
+    return sorted(
+        {
+            requirement
+            for name, kernel in manifest.items()
+            if not selected or name in selected
+            for requirement in kernel.get("build-requires", [])
+        }
+    )
+
+
+get_requires_for_build_editable = get_requires_for_build_wheel
+get_requires_for_build_sdist = get_requires_for_build_wheel

prime_kernels/_spec.py

@@ -59,6 +59,8 @@ class KernelSpec:
     nvcc_flags: tuple[str, ...]
     python_only: bool
     requires: tuple[str, ...]
+    include_packages: dict[str, tuple[str, ...]]
+    build_requires: tuple[str, ...]
 
     @property
     def module(self) -> str:
@@ -93,4 +95,6 @@ def _kernel(name: str, path: Path, table: dict) -> KernelSpec:
         nvcc_flags=tuple(table.get("nvcc-flags", [])),
         python_only=table.get("python-only", False),
         requires=tuple(table.get("requires", [])),
+        include_packages={name: tuple(paths) for name, paths in table.get("include-packages", {}).items()},
+        build_requires=tuple(table.get("build-requires", [])),
     )

prime_kernels/kernels.toml

@@ -8,6 +8,7 @@
 #   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
+#   build-requires  Python distributions needed only to build this kernel
 #
 # Vendored Python kernels use `python-only = true`, may list import checks in `requires`,
 # and omit `ops`, `sources`, and compilation settings.
@@ -30,6 +31,26 @@ python-only = true
 requires = ["torchao"]
 arch = ["10.0"]
 
+[nvfp4_moe]
+description = "Per-token NVFP4 grouped GEMM with BF16 backward"
+build-requires = ["nvidia-cutlass==4.2.0.0"]
+ops = "prime_rl"
+sources = [
+    "csrc/torch_interface.cpp",
+    "csrc/quantize/bindings.cpp",
+    "csrc/quantize/quantize.cu",
+    "csrc/grouped_gemm/bindings.cpp",
+    "csrc/grouped_gemm/f4f4bf16_ultra_grouped.cu",
+    "csrc/grouped_gemm/f4f4bf16_ultra_grouped_256_128_256_2_1_1.cu",
+    "csrc/grouped_gemm/f4f4bf16_ultra_grouped_256_256_256_2_1_1.cu",
+]
+include-dirs = ["csrc/grouped_gemm"]
+include-packages = { cutlass_library = ["source/include", "source/tools/util/include"] }
+arch = ["10.0a"]
+cxx-std = 17
+cxx-flags = ["-O3"]
+nvcc-flags = ["-O3", "--expt-relaxed-constexpr", "--expt-extended-lambda", "--threads=2"]
+
 # 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/nvfp4_moe/MSLK_LICENSE

@@ -0,0 +1,30 @@
+BSD License
+
+For MSLK software
+
+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:
+
+ * Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+ * 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.
+
+ * Neither the name Facebook 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/nvfp4_moe/README.md

@@ -0,0 +1,43 @@
+# NVFP4 routed-expert GEMM
+
+This SM100 kernel quantizes BF16 inputs and expert weights to E2M1 with
+block-16 E4M3 scales. Activations have one FP32 outer scale per token; weights
+have one per expert. Gated experts must pass a fused gate/up weight to share
+the outer weight scale.
+
+`grouped_gemm(x, weight_t, offs=offsets, backward="dequant_bf16")` accepts
+`x[M, K]`, `weight_t[E, K, N]`, and cumulative INT32 row offsets. The token
+dispatcher aligns expert row counts to `TOKEN_GROUP_ALIGNMENT`. Expert hidden
+and intermediate dimensions must be positive multiples of 32.
+
+The default backward reconstructs the packed forward operands and computes
+BF16 dgrad and wgrad. `backward="bf16"` uses the original BF16 operands.
+Master parameters and optimizer precision are controlled by the trainer.
+Use `fullgraph=False` with the trainer.
+
+`four_over_six=True` enables adaptive map-to-4/map-to-6 scaling for both
+operands, matching FlashInfer's default 4/6 recipe: E4M3 normalization bound
+448, MAE error selection, strict error scoring, and the default approximate
+candidate arithmetic. The option defaults to `False`. The same option is
+available on `quantize_activations` and `quantize_weights`.
+
+On inference, enable `FLASHINFER_NVFP4_4OVER6=1` and leave the other FlashInfer
+NVFP4 options at their defaults. Changing the normalization bound, error metric,
+error-scoring mode, or `FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH` would select a
+different recipe. The quantized layout and backward operands are unchanged.
+
+Build with a CUDA development toolkit matching PyTorch's CUDA major.
+Kernel-specific dependencies are declared in `kernels.toml` under
+`build-requires`; the build backend installs the dependencies for the selected
+kernels into the isolated build environment:
+
+```sh
+PRIME_KERNELS=nvfp4_moe PRIME_KERNELS_REQUIRE=1 MAX_JOBS=2 \
+  UV_TORCH_BACKEND=cu130 uv build --wheel .
+uv run pytest tests/test_nvfp4_moe.py
+```
+
+The grouped GEMM derives from MSLK and retains `MSLK_LICENSE`. The quantizer
+retains Transformer Engine and FlashInfer attribution in its CUDA source. The build
+uses packaged CUTLASS headers; no compiler or CUTLASS package is needed by a
+finished wheel at runtime.

prime_kernels/nvfp4_moe/__init__.py

@@ -0,0 +1,19 @@
+from . import _C  # noqa: F401
+from .grouped_gemm import grouped_gemm
+from .quantize import quantize_activations, quantize_weights
+
+__all__ = [
+    "TOKEN_GROUP_ALIGNMENT",
+    "grouped_gemm",
+    "quantize_activations",
+    "quantize_weights",
+    "unsupported_shape_reason",
+]
+
+TOKEN_GROUP_ALIGNMENT = 32
+
+
+def unsupported_shape_reason(dim: int, hidden_dim: int) -> str | None:
+    if dim <= 0 or hidden_dim <= 0 or dim % 32 or hidden_dim % 32:
+        return "NVFP4 expert dimensions must be positive multiples of 32"
+    return None

prime_kernels/nvfp4_moe/_ops.py

@@ -0,0 +1,107 @@
+import torch
+
+
+def _round_up(value: int, multiple: int) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+def _quantize_activations_fake(
+    matrix: torch.Tensor,
+    offsets: torch.Tensor,
+    four_over_six: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    rows, contraction_size = matrix.shape
+    groups = offsets.shape[0]
+    scale_columns = _round_up(contraction_size // 16, 4)
+    padded_scale_rows = _round_up(rows + groups * 127, 128)
+    return (
+        matrix.new_empty((rows, contraction_size // 2), dtype=torch.uint8),
+        matrix.new_empty(
+            (padded_scale_rows, scale_columns),
+            dtype=torch.float8_e4m3fn,
+        ),
+        matrix.new_empty((rows,), dtype=torch.float32),
+    )
+
+
+def _quantize_weights_fake(
+    weight_rows: torch.Tensor,
+    four_over_six: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    groups, output_size, contraction_size = weight_rows.shape
+    padded_output_size = _round_up(output_size, 128)
+    scale_columns = _round_up(contraction_size // 16, 4)
+    return (
+        weight_rows.new_empty(
+            (groups, output_size, contraction_size // 2),
+            dtype=torch.uint8,
+        ),
+        weight_rows.new_empty(
+            (groups, padded_output_size * scale_columns),
+            dtype=torch.float8_e4m3fn,
+        ),
+        weight_rows.new_empty((groups,), dtype=torch.float32),
+    )
+
+
+def _dequantize_activations_fake(
+    packed: torch.Tensor,
+    block_scales: torch.Tensor,
+    global_scales: torch.Tensor,
+    offsets: torch.Tensor,
+) -> torch.Tensor:
+    del block_scales, global_scales, offsets
+    return packed.new_empty(
+        (packed.shape[0], packed.shape[1] * 2),
+        dtype=torch.bfloat16,
+    )
+
+
+def _dequantize_weights_fake(
+    packed: torch.Tensor,
+    block_scales: torch.Tensor,
+    global_scales: torch.Tensor,
+) -> torch.Tensor:
+    del block_scales, global_scales
+    return packed.new_empty(
+        (packed.shape[0], packed.shape[1], packed.shape[2] * 2),
+        dtype=torch.bfloat16,
+    )
+
+
+def _grouped_mm_fake(
+    activations: torch.Tensor,
+    weight: torch.Tensor,
+    activation_block_scales: torch.Tensor,
+    weight_block_scales: torch.Tensor,
+    offsets: torch.Tensor,
+    activation_token_scales: torch.Tensor,
+    weight_expert_scales: torch.Tensor,
+) -> torch.Tensor:
+    del (
+        activation_block_scales,
+        weight_block_scales,
+        offsets,
+        activation_token_scales,
+        weight_expert_scales,
+    )
+    return activations.new_empty(
+        (activations.shape[0], weight.shape[-1]),
+        dtype=torch.bfloat16,
+    )
+
+
+torch.library.register_fake("prime_rl::quantize_nvfp4_activations", _quantize_activations_fake)
+_quantize_activations = torch.ops.prime_rl.quantize_nvfp4_activations.default
+
+torch.library.register_fake("prime_rl::quantize_nvfp4_weights", _quantize_weights_fake)
+_quantize_weights = torch.ops.prime_rl.quantize_nvfp4_weights.default
+
+torch.library.register_fake("prime_rl::dequantize_nvfp4_activations", _dequantize_activations_fake)
+_dequantize_activations = torch.ops.prime_rl.dequantize_nvfp4_activations.default
+
+torch.library.register_fake("prime_rl::dequantize_nvfp4_weights", _dequantize_weights_fake)
+_dequantize_weights = torch.ops.prime_rl.dequantize_nvfp4_weights.default
+
+torch.library.register_fake("prime_rl::grouped_nvfp4_gemm", _grouped_mm_fake)
+_grouped_mm = torch.ops.prime_rl.grouped_nvfp4_gemm.default

prime_kernels/nvfp4_moe/csrc/grouped_gemm/bindings.cpp

@@ -0,0 +1,37 @@
+#include "nvfp4_grouped.h"
+
+#include <torch/library.h>
+
+namespace prime_kernels::nvfp4 {
+
+at::Tensor grouped_mm(
+    const at::Tensor& activations,
+    const at::Tensor& weight,
+    const at::Tensor& activation_block_scales,
+    const at::Tensor& weight_block_scales,
+    const at::Tensor& offsets,
+    const at::Tensor& activation_token_scales,
+    const at::Tensor& weight_expert_scales) {
+  return f4f4bf16_ultra_grouped_mm(
+      activations,
+      weight,
+      activation_block_scales,
+      weight_block_scales,
+      offsets,
+      activation_token_scales,
+      weight_expert_scales);
+}
+
+TORCH_LIBRARY_FRAGMENT(prime_rl, m) {
+  m.def(
+      "grouped_nvfp4_gemm(Tensor activations, Tensor weight, "
+      "Tensor activation_block_scales, Tensor weight_block_scales, "
+      "Tensor offsets, Tensor activation_token_scales, "
+      "Tensor weight_expert_scales) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(prime_rl, CUDA, m) {
+  m.impl("grouped_nvfp4_gemm", grouped_mm);
+}
+
+} // namespace prime_kernels::nvfp4