[BugFix] route awq_dequantize through the upstream custom op

[Lfan-ke]

Purpose

vllm_metax/quant_config/auto_awq.py has two paths, selected by group_size % 32:

group_size % 32	weights	compute
== 0	repacked to GPTQ-4bit (awq_to_gptq_4bit)	fused awq_gemm
!= 0	left in AWQ layout	awq_dequantize + torch.matmul

The second path is dead. It called mx_ops.awq_dequantize (auto_awq.py:140), but vllm_metax/_custom_ops.py never defined that name, so the call raises:

AttributeError: module 'vllm_metax._custom_ops' has no attribute 'awq_dequantize'

The per-channel case reaches it: group_size = -1 gives -1 % 32 == 31, so every AWQ checkpoint whose group_size is not a multiple of 32 fails to load.

Changes

awq_dequantize is not modified on MetaX. csrc/libtorch_stable/torch_bindings.cpp:680 registers it under _C through STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops), which is the namespace vLLM's own wrapper resolves against, so upstream's vllm._custom_ops.awq_dequantize already reaches the MetaX kernel.

The note at the top of vllm_metax/_custom_ops.py says that file only carries ops that are modified or newly added in vllm_metax compared to vllm, and to check upstream first to avoid duplicates. So the call goes to vllm._custom_ops directly, as it already does in models/deepseek_v2.py, model_executor/layers/attention/mla_attention.py and the three sparse_attn_indexer backends.

awq_gemm and awq_to_gptq_4bit stay on mx_ops: the former takes two extra MetaX arguments (temp_space, dtype_bf16), the latter has no upstream counterpart.

Test Plan

On a MetaX C500 (MACA 3.5.3.20), resolve the op through both modules and call it on a 4-bit AWQ weight.

Test Result

>>> import vllm_metax._custom_ops as mx
>>> mx.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
AttributeError: module 'vllm_metax._custom_ops' has no attribute 'awq_dequantize'

>>> from vllm import _custom_ops as ops
>>> ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0).shape
torch.Size([128, 64])          # float16, on maca

torch.ops._C.awq_dequantize is present on the MetaX build, and vllm._custom_ops.awq_dequantize dispatches to it.

[gemini-code-assist]

Code Review

This pull request introduces the awq_dequantize function in vllm_metax/_custom_ops.py, which supports dequantization using either Triton AWQ or a custom C operator depending on the environment configuration. There are no review comments, and I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Lfan-ke]

@codex review please. read code at pr and repositories.