iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Quant] Use canonical N-first weight format for CT WNA16 MoE (#52798)

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: HDCharles <{"message":"Not Found","documentation_url":"https://docs.github.com/rest/users/emails#list-email-addresses-for-the-authenticated-user","status":"404"}>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
HDCharles
2026-09-28 11:22:44 -06:00
committed by GitHub
co-authored by Claude Opus 4.6 HDCharles <{"message":"Not Found","documentation_url":"https://docs.github.com/rest/users/emails#list-email-addresses-for-the-authenticated-user","status":"404"}> OpenAI Codex OpenAI Codex
parent 6d3ea3c2c9
commit c01368fd64
6 changed files with 222 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/quantization/test_int4_emulation_moe.py
+43 -8
View File
@@ -240,6 +240,29 @@ def _make_awq_moe_weights(E, K, N, group_size, asym=False):
)
def _canonicalize_gptq_moe_inputs(
w13: torch.Tensor,
w2: torch.Tensor,
w13_scale: torch.Tensor,
w2_scale: torch.Tensor,
w13_qzeros: torch.Tensor | None,
w2_qzeros: torch.Tensor | None,
):
"""Convert GPTQ K-first tensors to the oracle's canonical N-first layout."""
def transpose(value: torch.Tensor | None) -> torch.Tensor | None:
return value.transpose(1, 2).contiguous() if value is not None else None
return (
w13.transpose(1, 2).contiguous(),
w2.transpose(1, 2).contiguous(),
w13_scale.transpose(1, 2).contiguous(),
w2_scale.transpose(1, 2).contiguous(),
transpose(w13_qzeros),
transpose(w2_qzeros),
)
def _run_emulation_forward(
experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N
):
@@ -478,7 +501,8 @@ def test_gptq_process_weights_shapes_and_values(E, K, N, group_size, asym):
w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_gptq_moe_weights(
E, K, N, group_size, asym
)
result = _process_weights_emulation_gptq(w13, w2, w13s, w2s, w13z, w2z)
gptq_inputs = _canonicalize_gptq_moe_inputs(w13, w2, w13s, w2s, w13z, w2z)
result = _process_weights_emulation_gptq(*gptq_inputs)
w13_out, w2_out = result[0], result[1]
assert w13_out.shape == (E, 2 * N, K)
@@ -561,7 +585,7 @@ def test_gptq_awq_process_weights_agree(E, K, N, group_size, asym):
g2s_list.append(s2)
a2s_list.append(s2)
gptq_res = _process_weights_emulation_gptq(
gptq_inputs = _canonicalize_gptq_moe_inputs(
torch.stack(g13_list),
torch.stack(g2_list),
torch.stack(g13s_list),
@@ -569,6 +593,7 @@ def test_gptq_awq_process_weights_agree(E, K, N, group_size, asym):
torch.stack(g13z_list) if asym else None,
torch.stack(g2z_list) if asym else None,
)
gptq_res = _process_weights_emulation_gptq(*gptq_inputs)
awq_res = _process_weights_emulation_awq(
torch.stack(a13_list),
torch.stack(a2_list),
@@ -615,7 +640,7 @@ def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens):
g2s_list.append(s2)
a2s_list.append(s2.clone())
gptq_res = _process_weights_emulation_gptq(
gptq_inputs = _canonicalize_gptq_moe_inputs(
torch.stack(g13_list),
torch.stack(g2_list),
torch.stack(g13s_list),
@@ -623,6 +648,7 @@ def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens):
None,
None,
)
gptq_res = _process_weights_emulation_gptq(*gptq_inputs)
awq_res = _process_weights_emulation_awq(
torch.stack(a13_list),
torch.stack(a2_list),
@@ -750,7 +776,7 @@ def test_ep_output_matches_no_ep(E, K, N, top_k, group_size, num_tokens, ep_size
if fmt == "gptq"
else _process_weights_emulation_awq
)
res = process_fn(
process_inputs = (
torch.stack(packed13_list),
torch.stack(packed2_list),
torch.stack(scales13_list),
@@ -758,6 +784,9 @@ def test_ep_output_matches_no_ep(E, K, N, top_k, group_size, num_tokens, ep_size
None,
None,
)
if fmt == "gptq":
process_inputs = _canonicalize_gptq_moe_inputs(*process_inputs)
res = process_fn(*process_inputs)
w13_all, w2_all = res[0], res[1] # [E, 2N, K], [E, K, N]
hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
@@ -855,7 +884,7 @@ def test_ep_gptq_awq_agree(E, K, N, top_k, group_size, num_tokens, ep_size):
g2s_list.append(s2)
a2s_list.append(s2.clone())
gptq_res = _process_weights_emulation_gptq(
gptq_inputs = _canonicalize_gptq_moe_inputs(
torch.stack(g13_list),
torch.stack(g2_list),
torch.stack(g13s_list),
@@ -863,6 +892,7 @@ def test_ep_gptq_awq_agree(E, K, N, top_k, group_size, num_tokens, ep_size):
None,
None,
)
gptq_res = _process_weights_emulation_gptq(*gptq_inputs)
awq_res = _process_weights_emulation_awq(
torch.stack(a13_list),
torch.stack(a2_list),
@@ -954,7 +984,7 @@ def test_ep_partial_rank_no_active_experts(
scales13_list.append(s13)
scales2_list.append(s2)
res = _process_weights_emulation_gptq(
gptq_inputs = _canonicalize_gptq_moe_inputs(
torch.stack(packed13_list),
torch.stack(packed2_list),
torch.stack(scales13_list),
@@ -962,6 +992,7 @@ def test_ep_partial_rank_no_active_experts(
None,
None,
)
res = _process_weights_emulation_gptq(*gptq_inputs)
w13_all, w2_all = res[0], res[1]
# Force topk_ids to only use experts in [0, num_local) — rank 0's slice
@@ -1029,7 +1060,7 @@ def test_ep_sum_equals_full_forward(E, K, N, top_k, group_size, num_tokens, ep_s
scales13_list.append(s13)
scales2_list.append(s2)
res = _process_weights_emulation_gptq(
gptq_inputs = _canonicalize_gptq_moe_inputs(
torch.stack(packed13_list),
torch.stack(packed2_list),
torch.stack(scales13_list),
@@ -1037,6 +1068,7 @@ def test_ep_sum_equals_full_forward(E, K, N, top_k, group_size, num_tokens, ep_s
None,
None,
)
res = _process_weights_emulation_gptq(*gptq_inputs)
w13_all, w2_all = res[0], res[1]
# Fix routing so every token uses exactly 2 consecutive experts (round-robin)
@@ -1143,7 +1175,7 @@ def test_emulation_output_close_to_bf16_reference(
if fmt == "gptq"
else _process_weights_emulation_awq
)
res = process_fn(
process_inputs = (
torch.stack(packed13_list),
torch.stack(packed2_list),
torch.stack(scales13_list),
@@ -1151,6 +1183,9 @@ def test_emulation_output_close_to_bf16_reference(
None,
None,
)
if fmt == "gptq":
process_inputs = _canonicalize_gptq_moe_inputs(*process_inputs)
res = process_fn(*process_inputs)
w13_bf16, w2_bf16 = res[0], res[1]
dummy_scale = torch.ones(1, dtype=torch.float16, device=device)
tests/quantization/test_moe_wna16.py
+44 -9
View File
@@ -149,7 +149,11 @@ def test_wna16_oracle_rejects_incompatible_quant_structures(
assert expected in reason
def test_compressed_tensors_weights_are_transposed_for_triton():
def test_compressed_tensors_nfirst_weights_pass_through_for_triton():
"""CT provides N-first [E, N, K_packed] int32; Triton expects N-first uint8.
No transpose needed — just a view from int32 to uint8.
"""
quant_config = QuantizationArgs(
num_bits=4,
type=QuantizationType.INT,
@@ -158,10 +162,11 @@ def test_compressed_tensors_weights_are_transposed_for_triton():
dynamic=False,
group_size=32,
)
w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)
w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)
w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)
w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)
# N-first layout: [E, N, K_packed]
w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2)
w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2)
w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4)
w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3)
converted = convert_to_wna16_moe_kernel_format(
backend=WNA16MoEBackend.TRITON,
@@ -175,10 +180,40 @@ def test_compressed_tensors_weights_are_transposed_for_triton():
)
assert converted is not None
assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))
assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))
assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())
assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())
assert torch.equal(converted[0], w13.contiguous().view(torch.uint8))
assert torch.equal(converted[1], w2.contiguous().view(torch.uint8))
assert torch.equal(converted[2], w13_scale)
assert torch.equal(converted[3], w2_scale)
def test_gptq_nfirst_weights_pass_through_for_triton():
"""GPTQ frontend normalizes to N-first before calling oracle.
Verifies the oracle handles N-first GPTQ inputs the same way as CT.
"""
quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
# N-first layout (after GPTQ frontend transpose): [E, N, K_packed]
w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2)
w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2)
w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4)
w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3)
converted = convert_to_wna16_moe_kernel_format(
backend=WNA16MoEBackend.TRITON,
layer=torch.nn.Module(),
quant_config=quant_config,
input_dtype=None,
w13=w13,
w2=w2,
w13_scale=w13_scale,
w2_scale=w2_scale,
)
assert converted is not None
assert torch.equal(converted[0], w13.contiguous().view(torch.uint8))
assert torch.equal(converted[1], w2.contiguous().view(torch.uint8))
assert torch.equal(converted[2], w13_scale)
assert torch.equal(converted[3], w2_scale)
def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+87 -102
View File
@@ -567,13 +567,26 @@ def _process_weights_marlin(
"""Standard Marlin weight post-processing shared by MARLIN and
BATCHED_MARLIN backends.
Inputs arrive in canonical N-first layout ``[E, N, K_packed]``.
Steps
-----
1. Optional FP8 preprocessing of packed weights / scales.
2. Repack weights via ``gptq_marlin_moe_repack``.
3. Permute scales (and optionally extract INT8 global scales).
4. Permute bias tensors.
1. Transpose canonical N-first to K-first for Marlin.
2. Optional FP8 preprocessing of packed weights / scales.
3. Repack weights via ``gptq_marlin_moe_repack``.
4. Permute scales (and optionally extract INT8 global scales).
5. Permute bias tensors.
"""
# Canonical N-first → K-first for Marlin.
w13_qweight = w13_qweight.transpose(1, 2).contiguous()
w2_qweight = w2_qweight.transpose(1, 2).contiguous()
w13_scales = w13_scales.transpose(1, 2).contiguous()
w2_scales = w2_scales.transpose(1, 2).contiguous()
if w13_qzeros is not None:
w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()
if w2_qzeros is not None:
w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
marlin_w13_qweight: torch.Tensor
@@ -934,7 +947,12 @@ def _process_weights_cpu(
torch.Tensor | None, # w13_bias
torch.Tensor | None, # w2_bias
]:
"""CPU INT4 W4A16 weight post-processing."""
"""CPU INT4 W4A16 weight post-processing.
Non-AWQ inputs arrive in canonical N-first layout and are transposed
to K-first internally. AWQ uses a different packing axis and arrives
in its native format.
"""
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
prepare_int4_moe_layer_for_cpu,
)
@@ -945,15 +963,9 @@ def _process_weights_cpu(
AutoGPTQConfig,
)
# Detect packing format.
# AWQ: qweight is [E, K, 2*N//8] (packed along output/N dim).
# GPTQ: qweight is [E, K//8, 2*N] (packed along input/K dim).
# compressed-tensors: qweight is [E, K//8, 2*N] (packed along input/K dim).
if isinstance(quant_config, AutoAWQConfig):
# AWQ: K is stored unpacked in dim 1.
cpu_quant_algo = ops.CPUQuantAlgo.AWQ
elif isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):
# GPTQ / compressed-tensors: K//8 is stored packed in dim 1.
if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act:
raise NotImplementedError(
"CPU WNA16 MoE backend does not support GPTQ with "
@@ -961,6 +973,15 @@ def _process_weights_cpu(
"reordering support."
)
cpu_quant_algo = ops.CPUQuantAlgo.GPTQ
# Canonical N-first → K-first for CPU kernel.
w13 = w13.transpose(1, 2).contiguous()
w2 = w2.transpose(1, 2).contiguous()
w13_scale = w13_scale.transpose(1, 2).contiguous()
w2_scale = w2_scale.transpose(1, 2).contiguous()
if w13_qzeros is not None:
w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()
if w2_qzeros is not None:
w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()
else:
raise TypeError(
"CPU WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig "
@@ -1030,41 +1051,28 @@ def _process_weights_xpu(
torch.Tensor | None, # w13_bias
torch.Tensor | None, # w2_bias
]:
"""Repack GPTQ-format INT4 MoE weights into the layout
`vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True)` expects:
"""Repack INT4 MoE weights into the layout
`vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True)` expects.
w13: [E, 2*N, K] int4 (uint8 storage [E, 2*N, K // 2])
w13_scales: [E, 2*N, K // group_size] params_dtype
w2: [E, K, N] int4 (uint8 storage [E, K, N // 2])
w2_scales: [E, K, N // group_size] params_dtype
Input GPTQ layout from MoERunner.weight_loader:
w13: [E, K // 8, 2*N] int32 (8 nibbles per int32 along the input dim)
w13_scales: [E, K // group_size, 2*N] params_dtype
w2: [E, N // 8, K] int32
w2_scales: [E, N // group_size, K] params_dtype
Transpose dim 1 ↔ dim 2 then view int32 → uint8 to recover sequential
int4-packed bytes along the input dim. Each packed int32 holds 8 nibbles
`(n7<<28)|(n6<<24)|...|(n1<<4)|n0` in ascending K order; on a
little-endian host the int32→uint8 view exposes them as bytes
`[n1<<4|n0, n3<<4|n2, n5<<4|n4, n7<<4|n6]`, i.e. two nibbles per byte
with the lower nibble = lower input-K index. xpu_fused_moe(is_int4=True)
expects this convention; on a big-endian host the byte order reverses
and the kernel would silently miscompute, so we hard-fail.
Inputs arrive in canonical N-first layout ``[E, N, K_packed]`` int32.
The int32 → uint8 view recovers sequential int4-packed bytes along
the input dim. Each packed int32 holds 8 nibbles in ascending K order;
on a little-endian host the view exposes them as two nibbles per byte
with the lower nibble = lower input-K index. xpu_fused_moe(is_int4=True)
expects this convention; big-endian hosts are unsupported.
"""
del layer, quant_config # unused — kept for parity with the marlin helper
if sys.byteorder != "little":
raise NotImplementedError(
"_process_weights_xpu requires a little-endian host: the GPTQ "
"_process_weights_xpu requires a little-endian host: the "
"int32 → uint8 nibble repack relies on LE byte ordering."
)
w13_xpu = w13_qweight.transpose(1, 2).contiguous().view(torch.uint8)
w2_xpu = w2_qweight.transpose(1, 2).contiguous().view(torch.uint8)
w13_scales_xpu = w13_scales.transpose(1, 2).contiguous()
w2_scales_xpu = w2_scales.transpose(1, 2).contiguous()
w13_xpu = w13_qweight.contiguous().view(torch.uint8)
w2_xpu = w2_qweight.contiguous().view(torch.uint8)
w13_scales_xpu = w13_scales.contiguous()
w2_scales_xpu = w2_scales.contiguous()
return (
w13_xpu,
@@ -1315,16 +1323,23 @@ def _process_weights_emulation_gptq(
) -> tuple:
"""Dequantize int4 weights to BF16 for the emulation backend.
Inputs are in GPTQ packed format:
w13: [E, K//8, 2*N] int32 (gate+up proj stacked on dim 2)
w2: [E, N//8, K] int32
w13_scale: [E, K//gs, 2*N] float16
w2_scale: [E, N//gs, K] float16
Inputs arrive in canonical N-first layout and are transposed to
K-first internally for the GPTQ dequantization routines.
Outputs (what TritonExperts expects):
w13_out: [E, 2*N, K] bfloat16
w2_out: [E, K, N] bfloat16
"""
# Canonical N-first → K-first for dequantization.
w13 = w13.transpose(1, 2).contiguous()
w2 = w2.transpose(1, 2).contiguous()
w13_scale = w13_scale.transpose(1, 2).contiguous()
w2_scale = w2_scale.transpose(1, 2).contiguous()
if w13_qzeros is not None:
w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()
if w2_qzeros is not None:
w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()
# w13: packed along K (dim 1), output cols are 2*N (dim 2)
# transpose_output=True yields [E, 2*N, K]
w13_bf16 = _unpack_and_dequant_int4_gptq(
@@ -1435,6 +1450,10 @@ def convert_to_wna16_moe_kernel_format(
):
"""Dispatch weight post-processing to the appropriate per-backend handler.
All non-AWQ frontends must supply weights and scales in N-first
canonical layout ``[E, N_out, K_packed]``. AWQ uses a different
packing axis and is handled by dedicated per-backend helpers.
To add a new backend, implement a ``_process_weights_<name>`` helper and
add a branch here. Backends that rewrite the layer's parameters in place
(e.g. Humming) return ``None``; the caller then skips the param scatter.
@@ -1626,67 +1645,33 @@ def convert_to_wna16_moe_kernel_format(
w2_qzeros,
)
elif backend == WNA16MoEBackend.TRITON:
# Two possible input layouts depending on the quantization source:
#
# MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first
# → just view as uint8 (no-op)
#
# AutoGPTQ/compressed-tensors (int32, K-first):
# (E, K // pack32, N_out)
# → transpose to N-first, then view as uint8 to get
# (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s]
# Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs)
from vllm.model_executor.layers.quantization.auto_gptq import (
AutoGPTQConfig,
)
# All inputs arrive in canonical N-first format; view as uint8.
w13_uint8 = w13.contiguous().view(torch.uint8)
w2_uint8 = w2.contiguous().view(torch.uint8)
# Repack N-first int32 zero-points to uint8:
# [E, N//pf, K//gs] int32 → [E, N*4//pf, K//gs] uint8
if w13_qzeros is not None and w13_qzeros.dtype == torch.int32:
E13, Np13, Kg13 = w13_qzeros.shape
w13_qzeros = (
w13_qzeros.contiguous()
.view(torch.uint8)
.reshape(E13, Np13, Kg13, 4)
.permute(0, 1, 3, 2)
.reshape(E13, Np13 * 4, Kg13)
.contiguous()
)
if w2_qzeros is not None and w2_qzeros.dtype == torch.int32:
E2, Np2, Kg2 = w2_qzeros.shape
w2_qzeros = (
w2_qzeros.contiguous()
.view(torch.uint8)
.reshape(E2, Np2, Kg2, 4)
.permute(0, 1, 3, 2)
.reshape(E2, Np2 * 4, Kg2)
.contiguous()
)
if isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):
# These integrations build in K-first format even when the Triton
# backend is selected. Transpose to N-first first.
w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8)
w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8)
w13_scale = w13_scale.transpose(1, 2).contiguous()
w2_scale = w2_scale.transpose(1, 2).contiguous()
# Zero points from compressed-tensors checkpoints are K-first int32
# with 8 int4 ZPs packed per element: shape (E, K//gs, N//8).
# fused_moe_kernel_gptq_awq expects N-first uint8 with 2 int4 ZPs
# per byte: shape (E, N//2, K//gs), indexed as
# (offs_bn // 2) * stride_bzn + offs_k_group * stride_bzk.
# Conversion steps:
# (E, K//gs, N//8) int32
# → transpose(1,2) → (E, N//8, K//gs) int32
# → view(uint8) → (E, N//8, K//gs*4) [each int32 → 4 bytes]
# → reshape(…,4) → (E, N//8, K//gs, 4) [isolate byte index]
# → permute(0,1,3,2) → (E, N//8, 4, K//gs) [byte index before K]
# → reshape → (E, N//2, K//gs) [kernel expected layout]
# After this, element [e, offs_bn//2, k_group] is the uint8 byte
# holding the two int4 ZPs for output channels offs_bn and offs_bn+1.
if w13_qzeros is not None:
E13, Kg13, Np13 = w13_qzeros.shape
w13_qzeros = (
w13_qzeros.transpose(1, 2)
.contiguous()
.view(torch.uint8)
.reshape(E13, Np13, Kg13, 4)
.permute(0, 1, 3, 2)
.reshape(E13, Np13 * 4, Kg13)
.contiguous()
)
if w2_qzeros is not None:
E2, Kg2, Np2 = w2_qzeros.shape
w2_qzeros = (
w2_qzeros.transpose(1, 2)
.contiguous()
.view(torch.uint8)
.reshape(E2, Np2, Kg2, 4)
.permute(0, 1, 3, 2)
.reshape(E2, Np2 * 4, Kg2)
.contiguous()
)
else:
# MoeWNA16 uses N-first uint8 weights and scales.
w13_uint8 = w13.view(torch.uint8)
w2_uint8 = w2.view(torch.uint8)
return (
w13_uint8,
w2_uint8,
vllm/model_executor/layers/fused_moe/routed_experts.py
+1 -1
View File
@@ -676,7 +676,7 @@ class RoutedExperts(PluggableLayer):
# Hereafter, `expert_id` is local physical id
# is_transposed: if the dim to shard the weight
# should be flipped. Required by GPTQ, compressed-tensors
# should be flipped. Required by GPTQ/AWQ (K-first format).
# should be whatever dimension intermediate_size_per_partition is
is_transposed = getattr(param, "is_transposed", False)
vllm/model_executor/layers/quantization/auto_gptq.py
+18 -6
View File
@@ -626,19 +626,31 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
layer.register_parameter("w2_bias", None)
w2_bias = None
# Normalize GPTQ K-first layout to canonical N-first format.
w13 = layer.w13_qweight.data.transpose(1, 2).contiguous()
w2 = layer.w2_qweight.data.transpose(1, 2).contiguous()
w13_scale = layer.w13_scales.data.transpose(1, 2).contiguous()
w2_scale = layer.w2_scales.data.transpose(1, 2).contiguous()
w13_qzeros = getattr(layer, "w13_qzeros", None)
w2_qzeros = getattr(layer, "w2_qzeros", None)
if w13_qzeros is not None:
w13_qzeros = w13_qzeros.data.transpose(1, 2).contiguous()
if w2_qzeros is not None:
w2_qzeros = w2_qzeros.data.transpose(1, 2).contiguous()
converted = convert_to_wna16_moe_kernel_format(
backend=self.wna16_moe_backend,
layer=layer,
quant_config=self.quant_config,
input_dtype=self.input_dtype,
w13=layer.w13_qweight,
w2=layer.w2_qweight,
w13_scale=layer.w13_scales,
w2_scale=layer.w2_scales,
w13=w13,
w2=w2,
w13_scale=w13_scale,
w2_scale=w2_scale,
w13_bias=w13_bias,
w2_bias=w2_bias,
w13_qzeros=getattr(layer, "w13_qzeros", None),
w2_qzeros=getattr(layer, "w2_qzeros", None),
w13_qzeros=w13_qzeros,
w2_qzeros=w2_qzeros,
)
if converted is None:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+29 -83
View File
@@ -112,10 +112,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
WNA16MoEBackend.MARLIN,
WNA16MoEBackend.BATCHED_MARLIN,
]
self.is_transposed = self.wna16_backend not in (
WNA16MoEBackend.FLASHINFER_TRTLLM,
WNA16MoEBackend.HUMMING,
)
# CT checkpoints always load in N-first format [E, N, K_packed].
# Backend-specific transposition is handled in
# convert_to_wna16_moe_kernel_format.
self.is_transposed = False
if self.is_marlin:
assert check_moe_marlin_supports_config(
@@ -141,80 +141,31 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
num_groups_w2: int | None = None,
num_groups_w13: int | None = None,
) -> tuple[int, int, int]:
"""Get the shape of the weight based on the weight name, number of experts
hidden size, intermediate size per partition, number of groups for w2,
and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
for weight scales/zero_points.
"""Return N-first buffer shape for the given weight/scale/zp tensor.
All CT MoE weights are allocated in the checkpoint's native N-first
layout ``[E, N, K_packed]``. Backend-specific transposition (e.g. to
Marlin's K-first layout) is deferred to
``convert_to_wna16_moe_kernel_format``.
"""
if weight_name in ("w13_scale", "w13_zp"):
assert num_groups_w13 is not None, (
"num_groups_w13 must be provided for weight scales/zero_points"
)
if weight_name in ("w2_scale", "w2_zp"):
assert num_groups_w2 is not None, (
"num_groups_w2 must be provided for weight scales/zero_points"
)
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
shape_map: dict[str, dict[str, tuple[int, int | None, int | None]]] = {
"w13_weight": {
"Flashinfer": (
num_experts,
w13_num_shards * intermediate_size_per_partition,
self._packed_dim(hidden_size),
),
"Marlin": (
num_experts,
self._packed_dim(hidden_size),
w13_num_shards * intermediate_size_per_partition,
),
},
"w13_scale": {
"Flashinfer": (
num_experts,
w13_num_shards * intermediate_size_per_partition,
num_groups_w13,
),
"Marlin": (
num_experts,
num_groups_w13,
w13_num_shards * intermediate_size_per_partition,
),
},
"w13_zp": {
"Marlin": (
num_experts,
num_groups_w13,
self._packed_dim(w13_num_shards * intermediate_size_per_partition),
),
},
"w2_weight": {
"Flashinfer": (
num_experts,
hidden_size,
self._packed_dim(intermediate_size_per_partition),
),
"Marlin": (
num_experts,
self._packed_dim(intermediate_size_per_partition),
hidden_size,
),
},
"w2_scale": {
"Flashinfer": (num_experts, hidden_size, num_groups_w2),
"Marlin": (num_experts, num_groups_w2, hidden_size),
},
"w2_zp": {
"Marlin": (
num_experts,
num_groups_w2,
self._packed_dim(hidden_size),
),
},
w13_n = w13_num_shards * intermediate_size_per_partition
pf = self.packed_factor
shape_map = {
"w13_weight": (num_experts, w13_n, hidden_size // pf),
"w13_scale": (num_experts, w13_n, num_groups_w13),
"w13_zp": (num_experts, w13_n // pf, num_groups_w13),
"w2_weight": (
num_experts,
hidden_size,
intermediate_size_per_partition // pf,
),
"w2_scale": (num_experts, hidden_size, num_groups_w2),
"w2_zp": (num_experts, hidden_size // pf, num_groups_w2),
}
backend_key = "Marlin" if self.is_transposed else "Flashinfer"
shape = shape_map[weight_name][backend_key]
assert shape[1] is not None and shape[2] is not None
return shape[0], shape[1], shape[2]
shape = shape_map[weight_name]
assert isinstance(shape, tuple) and all(isinstance(d, int) for d in shape)
return shape # type: ignore[return-value]
def create_weights(
self,
@@ -225,11 +176,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
params_dtype: torch.dtype,
**extra_weight_attrs,
):
# Will transpose the loaded weight along the
# intermediate and hidden dim sizes. Will
# shard for TP along the transposed dims
extra_weight_attrs.update(
{"is_transposed": self.is_transposed, "quant_method": self.strategy}
{"is_transposed": False, "quant_method": self.strategy}
)
w13_weight = torch.nn.Parameter(
@@ -440,12 +388,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
replace_parameter(layer, "w13_weight_scale", w13_scales)
replace_parameter(layer, "w2_weight_scale", w2_scales)
# CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points
# even for symmetric quant (the oracle synthesizes them).
# CPU fused_experts_cpu requires zero points even for symmetric quant.
# EMULATION bakes ZP into the dequantized bf16 weights — ZP is None.
if (
not self.symmetric
or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3)
not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU
) and self.wna16_backend != WNA16MoEBackend.EMULATION:
assert w13_qzeros is not None and w2_qzeros is not None
replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)