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

Base of the NVFP4 stack: #9#12. Rebased onto current main; the package/registry implementation is unchanged. Indexed-attention PR #8 remains a separate child of this PR.

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.

@S1ro1
S1ro1 force-pushed the feat/mxfp8-moe-runtime branch from 3fb83eb to 0c4c878 Compare September 10, 2026 13:35
@S1ro1
S1ro1 removed this pull request from stack #10 September 10, 2026 13:36
@S1ro1
S1ro1 added this pull request to stack #13 September 10, 2026 13:36
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