Skip to content

Rename fused_marlin_moe kernels by precision scheme and split benchmarks - #5437

Draft
zeroherolin wants to merge 4 commits into
flagos-ai:masterfrom
zeroherolin:zl_marlin_polish
Draft

Rename fused_marlin_moe kernels by precision scheme and split benchmarks#5437
zeroherolin wants to merge 4 commits into
flagos-ai:masterfrom
zeroherolin:zl_marlin_polish

Conversation

@zeroherolin

Copy link
Copy Markdown
Contributor

PR Category

Operator | OP Test | Benchmark

Type of Change

Refactor

Description

Standardize fused_marlin_moe naming to fused_marlin_moe_<weight-act scheme>_<format> and split the benchmark by precision. Pure rename — no logic change, no new comments.

Kernels / entries / tune keys (src/flag_gems/fused/fused_marlin_moe.py + runtime registrations):

Before After
fused_moe_w4a16_gptq fused_marlin_moe_w4a16_int4
fused_moe_mxfp4 fused_marlin_moe_w4a16_mxfp4
_w4a16_* / _mxfp4_* kernels & invokes _w4a16_int4_* / _w4a16_mxfp4_*
w4a16_pack / mxfp4_pack w4a16_int4_pack / w4a16_mxfp4_pack
tune keys fused_marlin_moe_mxfp4[_gemm_silu] fused_marlin_moe_w4a16_mxfp4[_gemm_silu]

INT8 (w8a16_int8) reuses the generic wna16 fallback and has no dedicated kernel; w8a16_fp8 is reserved for future work.

Tests (tests/test_fused_marlin_moe.py): test functions renamed to test_fused_marlin_moe_{w4a16_int4, w8a16_int8, w4a16_mxfp4}.

Benchmark: benchmark/test_fused_marlin_moe.py split into three self-contained per-precision files (test_fused_marlin_moe_{w4a16_int4, w8a16_int8, w4a16_mxfp4}.py), following the test_fused_moe_<precision>.py convention; pytest marker unified to fused_marlin_moe.

Issue

N/A

Progress

  • Change is properly reviewed (1 reviewer required, 2 recommended).
  • Change is responded to an issue.
  • Change is fully covered by a UT.

Performance

No perf impact expected (identifier-only rename). Verified on H20 (sm90): 119 unit tests collected, per-precision subset + guard tests all pass; renamed tune keys load correctly via libtuner.

@Vincent-Xiao
Vincent-Xiao marked this pull request as draft August 13, 2026 08:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants