Skip to content

Commit 6db7724

Browse files
coderfeliclaude
andcommitted
refactor(mxfp4-moe): consolidate layout-API gemm into utils/moegemm/moe_dispatcher
Collapse the 5 partly-duplicated MXFP4 MoE layout-API modules into 3: * utils.py -- basics (shape constants, K-size formulas, pointer/LDS helpers, e8m0/SwiGLU quant math); no mma/load * moegemm.py -- gemm1 + gemm2 device bodies in one file, plus the shared B/B-scale fx.copy + fx.gemm layout primitives, the A-LDS loaders, and the atomic bf16 epilogue * moe_dispatcher.py-- compile_gemm1/2_a4w4_port + launch dispatch (mxfp4_moe_gemm1/2, caches, gemm1_grid) Removes the helper duplication (gemm1 previously re-defined _raw/_lds_ptr3/ _silu_mul_batch/... already in common) and makes the stack self-contained (no imports from the v1 modules). Pure code-motion -- no behavioral change. Deletes mxfp4_moe_{common,layout,gemm1,gemm2,gemm_2stage}.py and rewires the test import (kernels.mxfp4_moe_gemm_2stage -> kernels.moe_dispatcher). Validated on gfx950: test_moe_gemm_2stage layout-API path (gemm1->gemm2, opus-sort), a4w4 + a8w4, S/M/L -- numeric gate (mean_row_cos > 0.85) passes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 392f41c commit 6db7724

9 files changed

Lines changed: 1706 additions & 1862 deletions

File tree

kernels/moe_dispatcher.py

Lines changed: 550 additions & 0 deletions
Large diffs are not rendered by default.

kernels/moegemm.py

Lines changed: 996 additions & 0 deletions
Large diffs are not rendered by default.

kernels/mxfp4_moe_common.py

Lines changed: 0 additions & 217 deletions
This file was deleted.

0 commit comments

Comments
 (0)