Skip to content

feat: vendor MXFP8 MoE runtime - #9

Open
S1ro1 wants to merge 1 commit into
mainfrom
feat/mxfp8-moe-runtime
Open

feat: vendor MXFP8 MoE runtime#9
S1ro1 wants to merge 1 commit into
mainfrom
feat/mxfp8-moe-runtime

Conversation

@S1ro1

@S1ro1 S1ro1 commented Aug 30, 2026

Copy link
Copy Markdown

Summary

  • Add a Python-only mxfp8_moe kernel package with differentiable grouped GEMM and explicit MXFP8 expert-parallel transport.
  • Keep BF16 parameters and tensors at the public boundary; quantized data and scales remain internal.
  • Vendor the orchestration from torchao commit 02105d46c with its BSD license and wide-MoE/large-tensor handling.
  • Extend the kernel manifest, registry, and wheel builder to support Python-only kernels and runtime requirements.

Validation

  • Built a wheel with PRIME_KERNELS=mxfp8_moe and PRIME_KERNELS_REQUIRE=1 and verified the module, manifest, and license are included.
  • Imported the public module against the pinned torchao build.
  • Ruff check, Ruff format check, and diff checks pass.
  • The registry rejects the kernel on H200 before import as expected. Numerical execution still requires an SM100 runner.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant