diff --git a/README.md b/README.md index f92b69f..1420a28 100644 --- a/README.md +++ b/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.._C`. +surface and, for compiled kernels, the sources under `csrc/` compiled into +`prime_kernels.._C`. This repo builds and publishes its own prebuilt wheels via [`build_kernels.yaml`](.github/workflows/build_kernels.yaml). It is also consumed as a git @@ -50,10 +52,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 @@ -96,6 +97,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 diff --git a/prime_kernels/__init__.py b/prime_kernels/__init__.py index 31dd207..b48d3d4 100644 --- a/prime_kernels/__init__.py +++ b/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) diff --git a/prime_kernels/_spec.py b/prime_kernels/_spec.py index 3964a7e..75b6190 100644 --- a/prime_kernels/_spec.py +++ b/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", [])), ) diff --git a/prime_kernels/kernels.toml b/prime_kernels/kernels.toml index 4edf105..2c6f5ec 100644 --- a/prime_kernels/kernels.toml +++ b/prime_kernels/kernels.toml @@ -8,6 +8,9 @@ # sources C++/CUDA files to compile into prime_kernels.._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). # diff --git a/prime_kernels/mxfp8_moe/LICENSE.torchao b/prime_kernels/mxfp8_moe/LICENSE.torchao new file mode 100644 index 0000000..767f40b --- /dev/null +++ b/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. diff --git a/prime_kernels/mxfp8_moe/__init__.py b/prime_kernels/mxfp8_moe/__init__.py new file mode 100644 index 0000000..8ca978f --- /dev/null +++ b/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", +] diff --git a/prime_kernels/mxfp8_moe/grouped_gemm.py b/prime_kernels/mxfp8_moe/grouped_gemm.py new file mode 100644 index 0000000..8c629ae --- /dev/null +++ b/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) diff --git a/prime_kernels/mxfp8_moe/transport.py b/prime_kernels/mxfp8_moe/transport.py new file mode 100644 index 0000000..59c130e --- /dev/null +++ b/prime_kernels/mxfp8_moe/transport.py @@ -0,0 +1,96 @@ +# 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 +import torch.distributed as dist +from torch.distributed import ProcessGroup + +from .grouped_gemm import dequantize_rows, quantize_rows + + +def _all_to_all( + x: torch.Tensor, + output_splits: list[int], + input_splits: list[int], + group: ProcessGroup, +) -> torch.Tensor: + output = x.new_empty((sum(output_splits), *x.shape[1:])) + dist.all_to_all_single(output, x.contiguous(), output_splits, input_splits, group=group) + return output + + +def _quantized_all_to_all( + x: torch.Tensor, + output_splits: list[int], + input_splits: list[int], + group: ProcessGroup, +) -> torch.Tensor: + data, scales = quantize_rows(x) + output_data = _all_to_all(data, output_splits, input_splits, group) + output_scales = _all_to_all(scales.view(torch.uint8), output_splits, input_splits, group) + return dequantize_rows(output_data, output_scales.view(torch.float8_e8m0fnu), x.dtype) + + +class _Dispatch(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: torch.Tensor, + output_splits: list[int], + input_splits: list[int], + group: ProcessGroup, + ) -> torch.Tensor: + ctx.output_splits = output_splits + ctx.input_splits = input_splits + ctx.group = group + return _quantized_all_to_all(x, output_splits, input_splits, group) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_input = _all_to_all(grad_output, ctx.input_splits, ctx.output_splits, ctx.group) + return grad_input, None, None, None + + +class _Combine(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: torch.Tensor, + output_splits: list[int], + input_splits: list[int], + group: ProcessGroup, + ) -> torch.Tensor: + ctx.output_splits = output_splits + ctx.input_splits = input_splits + ctx.group = group + return _all_to_all(x, output_splits, input_splits, group) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_input = _quantized_all_to_all(grad_output, ctx.input_splits, ctx.output_splits, ctx.group) + return grad_input, None, None, None + + +def all_to_all_dispatch( + x: torch.Tensor, + output_splits: list[int], + input_splits: list[int], + group: ProcessGroup, +) -> torch.Tensor: + """Send expert inputs as MXFP8 and return explicit high-precision tensors.""" + return _Dispatch.apply(x, output_splits, input_splits, group) + + +def all_to_all_combine( + x: torch.Tensor, + output_splits: list[int], + input_splits: list[int], + group: ProcessGroup, +) -> torch.Tensor: + """Combine expert outputs in high precision and send their gradients as MXFP8.""" + return _Combine.apply(x, output_splits, input_splits, group) diff --git a/setup.py b/setup.py index 4230da4..06c81a4 100644 --- a/setup.py +++ b/setup.py @@ -70,9 +70,14 @@ def _extension(kernel) -> CUDAExtension: selection = {name for name in os.environ.get("PRIME_KERNELS", "").split(",") if name} require = os.environ.get("PRIME_KERNELS_REQUIRE") == "1" cuda = _nvcc_version() +unknown = selection - kernels.keys() +if unknown: + raise SystemExit(f"PRIME_KERNELS contains unknown kernel(s): {', '.join(sorted(unknown))}") extensions, skipped = [], {} for name, kernel in kernels.items(): + if kernel.python_only: + continue if selection and name not in selection: skipped[name] = "not selected by PRIME_KERNELS" continue @@ -84,15 +89,19 @@ def _extension(kernel) -> CUDAExtension: for name, reason in skipped.items(): print(f"prime-kernels: skipping {name}: {reason}", file=sys.stderr) -if require and skipped: - raise SystemExit(f"PRIME_KERNELS_REQUIRE=1 but {len(skipped)} kernel(s) were skipped") +required_skips = {name: reason for name, reason in skipped.items() if not selection or name in selection} +if require and required_skips: + raise SystemExit(f"PRIME_KERNELS_REQUIRE=1 but {len(required_skips)} requested kernel(s) were skipped") print(f"prime-kernels: building {', '.join(ext.name for ext in extensions) or ''}", file=sys.stderr) setuptools.setup( # Listed explicitly: the kernel folders carry C++/CUDA sources next to their Python, and # only the Python surface plus the compiled extension belongs in the wheel. packages=["prime_kernels", *(f"prime_kernels.{name}" for name in kernels)], - package_data={"prime_kernels": ["kernels.toml"]}, + package_data={ + "prime_kernels": ["kernels.toml"], + "prime_kernels.mxfp8_moe": ["LICENSE.torchao"], + }, ext_modules=extensions, cmdclass={"build_ext": BuildExtension}, )