Skip to content

[MoE Refactor] Migrate MoeWNA16Method quantization method over to using the new MK oracle scheme. - #44120

Merged
vllm-bot merged 79 commits into
vllm-project:mainfrom
neuralmagic:fix-moe-wna16-oracle
Jul 22, 2026
Merged

vllm-bot merged 79 commits into
vllm-project:mainfrom
neuralmagic:fix-moe-wna16-oracle

Conversation

@bnellnm

@bnellnm bnellnm commented May 31, 2026

Copy link
Copy Markdown
Collaborator

Purpose

This PR is an fixed + updated version of #42647

Migrate MoeWNA16Method quantization method over to using the new MK oracle scheme.

Additional fixes

  • Fix Triton backend weight conversion when original weights are from auto_gptq.
  • Fix marlin zero point conversion code when original ZPs come from auto_gptq.

Also fixes broken quantization import when humming module is not installed. This was broken prior to #42647

cc @bedeks @yzong-rh

Test Plan

  • pytest tests/compile/h100/test_startup.py::test_model_startup[kimi_k2.5]
  • pytest tests/evals/gsm8k/test_gsm8k_correctness.py::test_gsm8k_correctness[Qwen1.5-MoE-W4A16-CT] --config-list-file=configs/models-small.txt
  • lm-eval run --model vllm --tasks gsm8k --model_args="pretrained=Qwen/Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4" --num_fewshot 5 --batch_size 2048
  • lm-eval run --model vllm --tasks gsm8k --model_args="pretrained=TheBloke/Mixtral_7Bx2_MoE-AWQ" --num_fewshot 5 --batch_size 2048

Test Result

pytest tests pass.

lm-eval results:

Qwen baseline

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4640|±  |0.0137|
|     |       |strict-match    |     5|exact_match|↑  |0.3283|±  |0.0129|

Qwen PR

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|                                                                                   
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|                                                                                   
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4549|±  |0.0137|                                                                                   
|     |       |strict-match    |     5|exact_match|↑  |0.3237|±  |0.0129|                                                                                   

Mixtral baseline

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6929|±  |0.0127|
|     |       |strict-match    |     5|exact_match|↑  |0.6899|±  |0.0127|

Mixtral PR

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6952|±  |0.0127|
|     |       |strict-match    |     5|exact_match|↑  |0.6922|±  |0.0127|

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

claude and others added 30 commits May 13, 2026 20:36
…ion oracle

Converts the WNA16 Marlin MoE method to the new oracle-based kernel
abstraction pattern, matching the structure of CompressedTensorsW8A8Fp8MoEMethod.

Key changes:
- Add oracle/wna16.py with select_wna16_moe_backend(), make_wna16_moe_quant_config(),
  and make_wna16_moe_kernel() following the same pattern as oracle/fp8.py
- Add experts/trtllm_mxint4_moe.py with TrtLlmMxint4ExpertsMonolithic wrapping
  flashinfer_trtllm_mxint4_moe() as a FusedMoEExpertsMonolithic subclass
- Update CompressedTensorsWNA16MarlinMoEMethod.__init__ to call
  select_wna16_moe_backend() instead of inline backend selection
- Update process_weights_after_loading to build moe_quant_config and moe_kernel
  for 4-bit paths (both Flashinfer and Marlin); non-4-bit Marlin retains
  the direct fused_marlin_moe() fallback in apply()
- Remove GPTQMarlinState enum and select_gemm_impl() override (base class raises
  ValueError; moe_kernel is set so the old path is never reached for 4-bit)
- apply() and apply_monolithic() delegate to moe_kernel when available

Co-authored-by: Claude

https://claude.ai/code/session_01FLcv14ytXvwBUCsKZqpze1
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
@mergify

mergify Bot commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @bnellnm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Jul 6, 2026
@mergify mergify Bot removed the needs-rebase label Jul 14, 2026
@mergify

mergify Bot commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @bnellnm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Jul 16, 2026
bnellnm added 3 commits July 16, 2026 18:00
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
@mergify mergify Bot removed the needs-rebase label Jul 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance Performance-related issues ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants