Commit 6db7724
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
Large diffs are not rendered by default.
Large diffs are not rendered by default.
This file was deleted.
0 commit comments