PR #39136: [ROCm][Quantization][2/N] Refactor quark_moe w4a8 w/ oracle

• Repository: vllm-project/vllm
• Author: BowenBao
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/39136

Description (problem / solution / changelog)

1. Remove QuarkOCP_MX_MoEMethod_OSS and add aiter w4a8 backend.
2. Add unittest cases for rocm w4a16, w4a8 fused moe.
3. Validated locally with

pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
    --config-list-file=tests/evals/gpt_oss/configs/models-gfx950.txt

pytest -s -v tests/evals/gsm8k/test_gsm8k_correctness.py \
    --config-list-file=tests/evals/gsm8k/configs/models-qwen35-mi355.txt

Changed files

• tests/kernels/moe/test_ocp_mx_moe.py (modified, +388/-8)
• vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py (modified, +174/-26)
• vllm/model_executor/layers/fused_moe/oracle/mxfp4.py (modified, +126/-13)
• vllm/model_executor/layers/quantization/mxfp4.py (modified, +3/-3)
• vllm/model_executor/layers/quantization/quark/quark_moe.py (modified, +66/-200)

PR #41317: [Refactor] Extract shared helpers from MXFP4 MoE backend selectors

• Repository: vllm-project/vllm
• Author: demian-overflow
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/41317

Description (problem / solution / changelog)

Summary

Addresses #41291 by extracting the duplicated machinery shared between select_gpt_oss_mxfp4_moe_backend and select_mxfp4_moe_backend into private module-level helpers in vllm/model_executor/layers/fused_moe/oracle/mxfp4.py.

What changes

Both selectors carried four nested closures each (_make_log_backend, _make_log_unsupported, _return_or_raise) plus near-identical copies of the priority-iteration loop and the explicit-runner-backend branch. This PR lifts those into six private module-level helpers:

• _activation_format_for_config(config) — pulls the BatchedExperts vs Standard decision out of both functions.
• _make_log_backend(backend) / _make_log_unsupported(backend, reason) — the two log-message templates, deduplicated.
• _try_kernel_classes(backend, config, ..., *, log_failures=False) — probes every kernel class registered for backend. The log_failures switch preserves the difference between the priority-loop path (which logs each unsupported kernel via debug_once) and the explicit-backend path (which only surfaces failures by raising).
• _return_or_raise(backend, config, ...) — resolves a single backend or raises ValueError, used by every direct-pick site (env vars, runner-backend, XPU fallback).
• _select_explicit_runner_backend(config, activation_format) — honors config.moe_backend when set, returns None to signal "fall through to priority list".
• _select_first_supported_backend(config, priority_backends, activation_format) — walks a priority list and returns the first backend whose kernel matches.

select_gpt_oss_mxfp4_moe_backend and select_mxfp4_moe_backend are now thin orchestration that compose those helpers. Both keep their public signatures and observable behavior, so the three callers in vllm/model_executor/layers/quantization/mxfp4.py and vllm/model_executor/layers/quantization/quark/quark_moe.py need no edit.

Behavioral preservation

The refactor preserves observable behavior 1:1, with two intentional no-op cleanups bundled in:

1. Removed the redundant logger.info_once(_make_log_backend(backend)) immediately before _return_or_raise(Mxfp4MoeBackend.XPU, ...) in the GPT-OSS XPU branch — _return_or_raise already emits the same info_once message on success, so the original would have logged twice and info_once deduplicated it.
2. Removed the explicit scope="local" kwargs in select_mxfp4_moe_backend's log calls. "local" is the default for info_once / debug_once per vllm/logger.py, so the behavior is identical and the calls match the rest of the file.

The GPT-OSS path's behavior matrix is fully preserved:

• LoRA branch (raises on non-CUDA, picks TRITON_UNFUSED or MARLIN).
• All four env-var override branches (VLLM_USE_FLASHINFER_MOE_MXFP4_BF16, VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8, VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS, VLLM_MXFP4_USE_MARLIN) — kept in select_gpt_oss_mxfp4_moe_backend, since they don't apply to the generic MXFP4 path.
• The BATCHED_MARLIN coercion when the format is BatchedExperts.
• The XPU fallback after the priority loop.
• The CUDA/ROCm NotImplementedError and the final (NONE, None) return.
• The per-kernel-class debug_once for unsupported configs in the priority loop (preserved via log_failures=True).

Why this shape and not the issue's literal ask

The issue proposes "a single select_mxfp4_moe_backend() function … with activation_key parameter or quantization_config parameter". Reading both functions, the divergence is bigger than that:

A single function with a mode='gpt_oss' flag (or several boolean kwargs) would be harder to read than two named entry points and would lock in a shape that the larger class-hierarchy refactor in PR #37776 will rework anyway. This PR achieves the spirit of the dedup ask — eliminate the four copies of the helper closures and the duplicated priority loop — while keeping both entry points distinct.

Why it isn't a duplicate

Searched for existing PRs:

gh pr list --repo vllm-project/vllm --state all --search "41291 in:body"        → no PR
gh pr list --repo vllm-project/vllm --state open --search "select_mxfp4_moe_backend"  → no overlap
gh pr list --repo vllm-project/vllm --state open --search "merge mxfp4 backend"       → no overlap

The closest open PR is #37776 ([MoE] Unify MoE oracles with class structure, needs-rebase), which restructures all MoE oracles into a class hierarchy. That PR predates the GPT-OSS split that #39604 introduced and #40860 reshaped, so it does not touch select_gpt_oss_mxfp4_moe_backend at all. The two PRs don't conflict; whichever lands first, the other rebases cleanly.

Tracking issue: #41291.

Tests run

This is a Python-only refactor with no observable behavior change, so I ran the lint/typecheck commands instead of building the full vLLM wheel:

• ruff check vllm/model_executor/layers/fused_moe/oracle/mxfp4.py — clean.
• ruff format --check vllm/model_executor/layers/fused_moe/oracle/mxfp4.py — already formatted.
• mypy --python-version 3.10 --follow-imports=skip --ignore-missing-imports vllm/model_executor/layers/fused_moe/oracle/mxfp4.py — Success: no issues found in 1 source file.
• typos vllm/model_executor/layers/fused_moe/oracle/mxfp4.py — clean.
• python -c "import ast; ast.parse(open(...).read())" — OK.
• grep -rn "select_gpt_oss_mxfp4_moe_backend\|select_mxfp4_moe_backend" across the repo — three call sites in quantization/mxfp4.py and quantization/quark/quark_moe.py, none need editing because the public signatures are unchanged.

I did not run the full pytest suite or build the wheel — the PR makes no behavior changes and the helpers are private (no new public API surface). Happy to run any specific kernel/MoE test the reviewer thinks is relevant.

AI assistance disclosure

Claude (Anthropic) assisted with: reading both selector functions and the issue thread, comparing against the in-flight PR #37776, designing the helper extraction shape, applying the edit, and running the lint/mypy/typos commands. I (Demian Havdun) reviewed every changed line, checked that the helpers preserve the original behavior path-by-path (especially the per-kernel-class debug logging and the XPU branch double-log dedup), and signed off as committer per DCO. Co-author trailer added per AGENTS.md.

Changed files

• vllm/model_executor/layers/fused_moe/oracle/mxfp4.py (modified, +141/-126)