@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c01368fd64

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

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>

HDCharles committed Sep 28, 2026 at 13:22 UTC c01368fd64eaecc0c7242e4d933ec80fe8e047f5

6 files changed +223 -210

tests/kernels/quantization/test_int4_emulation_moe.py

+43 -8

@@ -240,6 +240,29 @@ def _make_awq_moe_weights(E, K, N, group_size, asym=False):

240 )

241

242

243 +def _canonicalize_gptq_moe_inputs(

244 + w13: torch.Tensor,

245 + w2: torch.Tensor,

246 + w13_scale: torch.Tensor,

247 + w2_scale: torch.Tensor,

248 + w13_qzeros: torch.Tensor | None,

249 + w2_qzeros: torch.Tensor | None,

250 +):

251 + """Convert GPTQ K-first tensors to the oracle's canonical N-first layout."""

252 +

253 + def transpose(value: torch.Tensor | None) -> torch.Tensor | None:

254 + return value.transpose(1, 2).contiguous() if value is not None else None

255 +

256 + return (

257 + w13.transpose(1, 2).contiguous(),

258 + w2.transpose(1, 2).contiguous(),

259 + w13_scale.transpose(1, 2).contiguous(),

260 + w2_scale.transpose(1, 2).contiguous(),

261 + transpose(w13_qzeros),

262 + transpose(w2_qzeros),

263 + )

264 +

265 +

266 def _run_emulation_forward(

267 experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N

268 ):

@@ -478,7 +501,8 @@ def test_gptq_process_weights_shapes_and_values(E, K, N, group_size, asym):

501 w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_gptq_moe_weights(

502 E, K, N, group_size, asym

503 )

481 - result = _process_weights_emulation_gptq(w13, w2, w13s, w2s, w13z, w2z)

504 + gptq_inputs = _canonicalize_gptq_moe_inputs(w13, w2, w13s, w2s, w13z, w2z)

505 + result = _process_weights_emulation_gptq(*gptq_inputs)

506 w13_out, w2_out = result[0], result[1]

507

508 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):

585 g2s_list.append(s2)

586 a2s_list.append(s2)

587

564 - gptq_res = _process_weights_emulation_gptq(

588 + gptq_inputs = _canonicalize_gptq_moe_inputs(

589 torch.stack(g13_list),

590 torch.stack(g2_list),

591 torch.stack(g13s_list),

@@ -569,6 +593,7 @@ def test_gptq_awq_process_weights_agree(E, K, N, group_size, asym):

593 torch.stack(g13z_list) if asym else None,

594 torch.stack(g2z_list) if asym else None,

595 )

596 + gptq_res = _process_weights_emulation_gptq(*gptq_inputs)

597 awq_res = _process_weights_emulation_awq(

598 torch.stack(a13_list),

599 torch.stack(a2_list),

@@ -615,7 +640,7 @@ def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens):

640 g2s_list.append(s2)

641 a2s_list.append(s2.clone())

642

618 - gptq_res = _process_weights_emulation_gptq(

643 + gptq_inputs = _canonicalize_gptq_moe_inputs(

644 torch.stack(g13_list),

645 torch.stack(g2_list),

646 torch.stack(g13s_list),

@@ -623,6 +648,7 @@ def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens):

648 None,

649 None,

650 )

651 + gptq_res = _process_weights_emulation_gptq(*gptq_inputs)

652 awq_res = _process_weights_emulation_awq(

653 torch.stack(a13_list),

654 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

776 if fmt == "gptq"

777 else _process_weights_emulation_awq

778 )

753 - res = process_fn(

779 + process_inputs = (

780 torch.stack(packed13_list),

781 torch.stack(packed2_list),

782 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

784 None,

785 None,

786 )

787 + if fmt == "gptq":

788 + process_inputs = _canonicalize_gptq_moe_inputs(*process_inputs)

789 + res = process_fn(*process_inputs)

790 w13_all, w2_all = res[0], res[1] # [E, 2N, K], [E, K, N]

791

792 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):

884 g2s_list.append(s2)

885 a2s_list.append(s2.clone())

886

858 - gptq_res = _process_weights_emulation_gptq(

887 + gptq_inputs = _canonicalize_gptq_moe_inputs(

888 torch.stack(g13_list),

889 torch.stack(g2_list),

890 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):

892 None,

893 None,

894 )

895 + gptq_res = _process_weights_emulation_gptq(*gptq_inputs)

896 awq_res = _process_weights_emulation_awq(

897 torch.stack(a13_list),

898 torch.stack(a2_list),

@@ -954,7 +984,7 @@ def test_ep_partial_rank_no_active_experts(

984 scales13_list.append(s13)

985 scales2_list.append(s2)

986

957 - res = _process_weights_emulation_gptq(

987 + gptq_inputs = _canonicalize_gptq_moe_inputs(

988 torch.stack(packed13_list),

989 torch.stack(packed2_list),

990 torch.stack(scales13_list),

@@ -962,6 +992,7 @@ def test_ep_partial_rank_no_active_experts(

992 None,

993 None,

994 )

995 + res = _process_weights_emulation_gptq(*gptq_inputs)

996 w13_all, w2_all = res[0], res[1]

997

998 # 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

1060 scales13_list.append(s13)

1061 scales2_list.append(s2)

1062

1032 - res = _process_weights_emulation_gptq(

1063 + gptq_inputs = _canonicalize_gptq_moe_inputs(

1064 torch.stack(packed13_list),

1065 torch.stack(packed2_list),

1066 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

1068 None,

1069 None,

1070 )

1071 + res = _process_weights_emulation_gptq(*gptq_inputs)

1072 w13_all, w2_all = res[0], res[1]

1073

1074 # Fix routing so every token uses exactly 2 consecutive experts (round-robin)

@@ -1143,7 +1175,7 @@ def test_emulation_output_close_to_bf16_reference(

1175 if fmt == "gptq"

1176 else _process_weights_emulation_awq

1177 )

1146 - res = process_fn(

1178 + process_inputs = (

1179 torch.stack(packed13_list),

1180 torch.stack(packed2_list),

1181 torch.stack(scales13_list),

@@ -1151,6 +1183,9 @@ def test_emulation_output_close_to_bf16_reference(

1183 None,

1184 None,

1185 )

1186 + if fmt == "gptq":

1187 + process_inputs = _canonicalize_gptq_moe_inputs(*process_inputs)

1188 + res = process_fn(*process_inputs)

1189 w13_bf16, w2_bf16 = res[0], res[1]

1190

1191 dummy_scale = torch.ones(1, dtype=torch.float16, device=device)

tests/quantization/test_moe_wna16.py

+44 -9

@@ -149,7 +149,11 @@ def test_wna16_oracle_rejects_incompatible_quant_structures(

149 assert expected in reason

150

151

152 -def test_compressed_tensors_weights_are_transposed_for_triton():

152 +def test_compressed_tensors_nfirst_weights_pass_through_for_triton():

153 + """CT provides N-first [E, N, K_packed] int32; Triton expects N-first uint8.

154 +

155 + No transpose needed — just a view from int32 to uint8.

156 + """

157 quant_config = QuantizationArgs(

158 num_bits=4,

159 type=QuantizationType.INT,

@@ -158,10 +162,41 @@ def test_compressed_tensors_weights_are_transposed_for_triton():

162 dynamic=False,

163 group_size=32,

164 )

161 - w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)

162 - w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)

163 - w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)

164 - w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)

165 + # N-first layout: [E, N, K_packed]

166 + w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2)

167 + w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2)

168 + w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4)

169 + w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3)

170 +

171 + converted = convert_to_wna16_moe_kernel_format(

172 + backend=WNA16MoEBackend.TRITON,

173 + layer=torch.nn.Module(),

174 + quant_config=quant_config,

175 + input_dtype=None,

176 + w13=w13,

177 + w2=w2,

178 + w13_scale=w13_scale,

179 + w2_scale=w2_scale,

180 + )

181 +

182 + assert converted is not None

183 + assert torch.equal(converted[0], w13.contiguous().view(torch.uint8))

184 + assert torch.equal(converted[1], w2.contiguous().view(torch.uint8))

185 + assert torch.equal(converted[2], w13_scale)

186 + assert torch.equal(converted[3], w2_scale)

187 +

188 +

189 +def test_gptq_nfirst_weights_pass_through_for_triton():

190 + """GPTQ frontend normalizes to N-first before calling oracle.

191 +

192 + Verifies the oracle handles N-first GPTQ inputs the same way as CT.

193 + """

194 + quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})

195 + # N-first layout (after GPTQ frontend transpose): [E, N, K_packed]

196 + w13 = torch.arange(16, dtype=torch.int32).reshape(1, 8, 2)

197 + w2 = torch.arange(12, dtype=torch.int32).reshape(1, 6, 2)

198 + w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 8, 4)

199 + w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 6, 3)

200

201 converted = convert_to_wna16_moe_kernel_format(

202 backend=WNA16MoEBackend.TRITON,

@@ -175,10 +210,10 @@ def test_compressed_tensors_weights_are_transposed_for_triton():

210 )

211

212 assert converted is not None

178 - assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))

179 - assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))

180 - assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())

181 - assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())

213 + assert torch.equal(converted[0], w13.contiguous().view(torch.uint8))

214 + assert torch.equal(converted[1], w2.contiguous().view(torch.uint8))

215 + assert torch.equal(converted[2], w13_scale)

216 + assert torch.equal(converted[3], w2_scale)

217

218

219 def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):

vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

+88 -103

@@ -567,13 +567,26 @@ def _process_weights_marlin(

567 """Standard Marlin weight post-processing shared by MARLIN and

568 BATCHED_MARLIN backends.

569

570 + Inputs arrive in canonical N-first layout ``[E, N, K_packed]``.

571 +

572 Steps

573 -----

572 - 1. Optional FP8 preprocessing of packed weights / scales.

573 - 2. Repack weights via ``gptq_marlin_moe_repack``.

574 - 3. Permute scales (and optionally extract INT8 global scales).

575 - 4. Permute bias tensors.

574 + 1. Transpose canonical N-first to K-first for Marlin.

575 + 2. Optional FP8 preprocessing of packed weights / scales.

576 + 3. Repack weights via ``gptq_marlin_moe_repack``.

577 + 4. Permute scales (and optionally extract INT8 global scales).

578 + 5. Permute bias tensors.

579 """

580 + # Canonical N-first → K-first for Marlin.

581 + w13_qweight = w13_qweight.transpose(1, 2).contiguous()

582 + w2_qweight = w2_qweight.transpose(1, 2).contiguous()

583 + w13_scales = w13_scales.transpose(1, 2).contiguous()

584 + w2_scales = w2_scales.transpose(1, 2).contiguous()

585 + if w13_qzeros is not None:

586 + w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()

587 + if w2_qzeros is not None:

588 + w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()

589 +

590 is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1

591

592 marlin_w13_qweight: torch.Tensor

@@ -934,7 +947,12 @@ def _process_weights_cpu(

947 torch.Tensor | None, # w13_bias

948 torch.Tensor | None, # w2_bias

949 ]:

937 - """CPU INT4 W4A16 weight post-processing."""

950 + """CPU INT4 W4A16 weight post-processing.

951 +

952 + Non-AWQ inputs arrive in canonical N-first layout and are transposed

953 + to K-first internally. AWQ uses a different packing axis and arrives

954 + in its native format.

955 + """

956 from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (

957 prepare_int4_moe_layer_for_cpu,

958 )

@@ -945,15 +963,9 @@ def _process_weights_cpu(

963 AutoGPTQConfig,

964 )

965

948 - # Detect packing format.

949 - # AWQ: qweight is [E, K, 2*N//8] (packed along output/N dim).

950 - # GPTQ: qweight is [E, K//8, 2*N] (packed along input/K dim).

951 - # compressed-tensors: qweight is [E, K//8, 2*N] (packed along input/K dim).

966 if isinstance(quant_config, AutoAWQConfig):

953 - # AWQ: K is stored unpacked in dim 1.

967 cpu_quant_algo = ops.CPUQuantAlgo.AWQ

968 elif isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):

956 - # GPTQ / compressed-tensors: K//8 is stored packed in dim 1.

969 if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act:

970 raise NotImplementedError(

971 "CPU WNA16 MoE backend does not support GPTQ with "

@@ -961,6 +973,15 @@ def _process_weights_cpu(

973 "reordering support."

974 )

975 cpu_quant_algo = ops.CPUQuantAlgo.GPTQ

976 + # Canonical N-first → K-first for CPU kernel.

977 + w13 = w13.transpose(1, 2).contiguous()

978 + w2 = w2.transpose(1, 2).contiguous()

979 + w13_scale = w13_scale.transpose(1, 2).contiguous()

980 + w2_scale = w2_scale.transpose(1, 2).contiguous()

981 + if w13_qzeros is not None:

982 + w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()

983 + if w2_qzeros is not None:

984 + w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()

985 else:

986 raise TypeError(

987 "CPU WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig "

@@ -1030,41 +1051,28 @@ def _process_weights_xpu(

1051 torch.Tensor | None, # w13_bias

1052 torch.Tensor | None, # w2_bias

1053 ]:

1033 - """Repack GPTQ-format INT4 MoE weights into the layout

1034 - `vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True)` expects:

1035 -

1036 - w13: [E, 2*N, K] int4 (uint8 storage [E, 2*N, K // 2])

1037 - w13_scales: [E, 2*N, K // group_size] params_dtype

1038 - w2: [E, K, N] int4 (uint8 storage [E, K, N // 2])

1039 - w2_scales: [E, K, N // group_size] params_dtype

1040 -

1041 - Input GPTQ layout from MoERunner.weight_loader:

1042 - w13: [E, K // 8, 2*N] int32 (8 nibbles per int32 along the input dim)

1043 - w13_scales: [E, K // group_size, 2*N] params_dtype

1044 - w2: [E, N // 8, K] int32

1045 - w2_scales: [E, N // group_size, K] params_dtype

1046 -

1047 - Transpose dim 1 ↔ dim 2 then view int32 → uint8 to recover sequential

1048 - int4-packed bytes along the input dim. Each packed int32 holds 8 nibbles

1049 - `(n7<<28)|(n6<<24)|...|(n1<<4)|n0` in ascending K order; on a

1050 - little-endian host the int32→uint8 view exposes them as bytes

1051 - `[n1<<4|n0, n3<<4|n2, n5<<4|n4, n7<<4|n6]`, i.e. two nibbles per byte

1052 - with the lower nibble = lower input-K index. xpu_fused_moe(is_int4=True)

1053 - expects this convention; on a big-endian host the byte order reverses

1054 - and the kernel would silently miscompute, so we hard-fail.

1054 + """Repack INT4 MoE weights into the layout

1055 + `vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True)` expects.

1056 +

1057 + Inputs arrive in canonical N-first layout ``[E, N, K_packed]`` int32.

1058 + The int32 → uint8 view recovers sequential int4-packed bytes along

1059 + the input dim. Each packed int32 holds 8 nibbles in ascending K order;

1060 + on a little-endian host the view exposes them as two nibbles per byte

1061 + with the lower nibble = lower input-K index. xpu_fused_moe(is_int4=True)

1062 + expects this convention; big-endian hosts are unsupported.

1063 """

1064 del layer, quant_config # unused — kept for parity with the marlin helper

1065

1066 if sys.byteorder != "little":

1067 raise NotImplementedError(

1060 - "_process_weights_xpu requires a little-endian host: the GPTQ "

1068 + "_process_weights_xpu requires a little-endian host: the "

1069 "int32 → uint8 nibble repack relies on LE byte ordering."

1070 )

1071

1064 - w13_xpu = w13_qweight.transpose(1, 2).contiguous().view(torch.uint8)

1065 - w2_xpu = w2_qweight.transpose(1, 2).contiguous().view(torch.uint8)

1066 - w13_scales_xpu = w13_scales.transpose(1, 2).contiguous()

1067 - w2_scales_xpu = w2_scales.transpose(1, 2).contiguous()

1072 + w13_xpu = w13_qweight.contiguous().view(torch.uint8)

1073 + w2_xpu = w2_qweight.contiguous().view(torch.uint8)

1074 + w13_scales_xpu = w13_scales.contiguous()

1075 + w2_scales_xpu = w2_scales.contiguous()

1076

1077 return (

1078 w13_xpu,

@@ -1315,16 +1323,23 @@ def _process_weights_emulation_gptq(

1323 ) -> tuple:

1324 """Dequantize int4 weights to BF16 for the emulation backend.

1325

1318 - Inputs are in GPTQ packed format:

1319 - w13: [E, K//8, 2*N] int32 (gate+up proj stacked on dim 2)

1320 - w2: [E, N//8, K] int32

1321 - w13_scale: [E, K//gs, 2*N] float16

1322 - w2_scale: [E, N//gs, K] float16

1326 + Inputs arrive in canonical N-first layout and are transposed to

1327 + K-first internally for the GPTQ dequantization routines.

1328

1329 Outputs (what TritonExperts expects):

1330 w13_out: [E, 2*N, K] bfloat16

1331 w2_out: [E, K, N] bfloat16

1332 """

1333 + # Canonical N-first → K-first for dequantization.

1334 + w13 = w13.transpose(1, 2).contiguous()

1335 + w2 = w2.transpose(1, 2).contiguous()

1336 + w13_scale = w13_scale.transpose(1, 2).contiguous()

1337 + w2_scale = w2_scale.transpose(1, 2).contiguous()

1338 + if w13_qzeros is not None:

1339 + w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()

1340 + if w2_qzeros is not None:

1341 + w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()

1342 +

1343 # w13: packed along K (dim 1), output cols are 2*N (dim 2)

1344 # transpose_output=True yields [E, 2*N, K]

1345 w13_bf16 = _unpack_and_dequant_int4_gptq(

@@ -1435,6 +1450,10 @@ def convert_to_wna16_moe_kernel_format(

1450 ):

1451 """Dispatch weight post-processing to the appropriate per-backend handler.

1452

1453 + All non-AWQ frontends must supply weights and scales in N-first

1454 + canonical layout ``[E, N_out, K_packed]``. AWQ uses a different

1455 + packing axis and is handled by dedicated per-backend helpers.

1456 +

1457 To add a new backend, implement a ``_process_weights_<name>`` helper and

1458 add a branch here. Backends that rewrite the layer's parameters in place

1459 (e.g. Humming) return ``None``; the caller then skips the param scatter.

@@ -1626,67 +1645,33 @@ def convert_to_wna16_moe_kernel_format(

1645 w2_qzeros,

1646 )

1647 elif backend == WNA16MoEBackend.TRITON:

1629 - # Two possible input layouts depending on the quantization source:

1630 - #

1631 - # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first

1632 - # → just view as uint8 (no-op)

1633 - #

1634 - # AutoGPTQ/compressed-tensors (int32, K-first):

1635 - # (E, K // pack32, N_out)

1636 - # → transpose to N-first, then view as uint8 to get

1637 - # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s]

1638 - # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs)

1639 - from vllm.model_executor.layers.quantization.auto_gptq import (

1640 - AutoGPTQConfig,

1641 - )

1648 + # All inputs arrive in canonical N-first format; view as uint8.

1649 + w13_uint8 = w13.contiguous().view(torch.uint8)

1650 + w2_uint8 = w2.contiguous().view(torch.uint8)

1651 +

1652 + # Repack N-first int32 zero-points to uint8:

1653 + # [E, N//pf, K//gs] int32 → [E, N*4//pf, K//gs] uint8

1654 + if w13_qzeros is not None and w13_qzeros.dtype == torch.int32:

1655 + E13, Np13, Kg13 = w13_qzeros.shape

1656 + w13_qzeros = (

1657 + w13_qzeros.contiguous()

1658 + .view(torch.uint8)

1659 + .reshape(E13, Np13, Kg13, 4)

1660 + .permute(0, 1, 3, 2)

1661 + .reshape(E13, Np13 * 4, Kg13)

1662 + .contiguous()

1663 + )

1664 + if w2_qzeros is not None and w2_qzeros.dtype == torch.int32:

1665 + E2, Np2, Kg2 = w2_qzeros.shape

1666 + w2_qzeros = (

1667 + w2_qzeros.contiguous()

1668 + .view(torch.uint8)

1669 + .reshape(E2, Np2, Kg2, 4)

1670 + .permute(0, 1, 3, 2)

1671 + .reshape(E2, Np2 * 4, Kg2)

1672 + .contiguous()

1673 + )

1674

1643 - if isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):

1644 - # These integrations build in K-first format even when the Triton

1645 - # backend is selected. Transpose to N-first first.

1646 - w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8)

1647 - w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8)

1648 - w13_scale = w13_scale.transpose(1, 2).contiguous()

1649 - w2_scale = w2_scale.transpose(1, 2).contiguous()

1650 - # Zero points from compressed-tensors checkpoints are K-first int32

1651 - # with 8 int4 ZPs packed per element: shape (E, K//gs, N//8).

1652 - # fused_moe_kernel_gptq_awq expects N-first uint8 with 2 int4 ZPs

1653 - # per byte: shape (E, N//2, K//gs), indexed as

1654 - # (offs_bn // 2) * stride_bzn + offs_k_group * stride_bzk.

1655 - # Conversion steps:

1656 - # (E, K//gs, N//8) int32

1657 - # → transpose(1,2) → (E, N//8, K//gs) int32

1658 - # → view(uint8) → (E, N//8, K//gs*4) [each int32 → 4 bytes]

1659 - # → reshape(…,4) → (E, N//8, K//gs, 4) [isolate byte index]

1660 - # → permute(0,1,3,2) → (E, N//8, 4, K//gs) [byte index before K]

1661 - # → reshape → (E, N//2, K//gs) [kernel expected layout]

1662 - # After this, element [e, offs_bn//2, k_group] is the uint8 byte

1663 - # holding the two int4 ZPs for output channels offs_bn and offs_bn+1.

1664 - if w13_qzeros is not None:

1665 - E13, Kg13, Np13 = w13_qzeros.shape

1666 - w13_qzeros = (

1667 - w13_qzeros.transpose(1, 2)

1668 - .contiguous()

1669 - .view(torch.uint8)

1670 - .reshape(E13, Np13, Kg13, 4)

1671 - .permute(0, 1, 3, 2)

1672 - .reshape(E13, Np13 * 4, Kg13)

1673 - .contiguous()

1674 - )

1675 - if w2_qzeros is not None:

1676 - E2, Kg2, Np2 = w2_qzeros.shape

1677 - w2_qzeros = (

1678 - w2_qzeros.transpose(1, 2)

1679 - .contiguous()

1680 - .view(torch.uint8)

1681 - .reshape(E2, Np2, Kg2, 4)

1682 - .permute(0, 1, 3, 2)

1683 - .reshape(E2, Np2 * 4, Kg2)

1684 - .contiguous()

1685 - )

1686 - else:

1687 - # MoeWNA16 uses N-first uint8 weights and scales.

1688 - w13_uint8 = w13.view(torch.uint8)

1689 - w2_uint8 = w2.view(torch.uint8)

1675 return (

1676 w13_uint8,

1677 w2_uint8,

vllm/model_executor/layers/fused_moe/routed_experts.py

+1 -1

@@ -676,7 +676,7 @@ class RoutedExperts(PluggableLayer):

676 # Hereafter, `expert_id` is local physical id

677

678 # is_transposed: if the dim to shard the weight

679 - # should be flipped. Required by GPTQ, compressed-tensors

679 + # should be flipped. Required by GPTQ/AWQ (K-first format).

680 # should be whatever dimension intermediate_size_per_partition is

681 is_transposed = getattr(param, "is_transposed", False)

682

vllm/model_executor/layers/quantization/auto_gptq.py

+18 -6

@@ -626,19 +626,31 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

626 layer.register_parameter("w2_bias", None)

627 w2_bias = None

628

629 + # Normalize GPTQ K-first layout to canonical N-first format.

630 + w13 = layer.w13_qweight.data.transpose(1, 2).contiguous()

631 + w2 = layer.w2_qweight.data.transpose(1, 2).contiguous()

632 + w13_scale = layer.w13_scales.data.transpose(1, 2).contiguous()

633 + w2_scale = layer.w2_scales.data.transpose(1, 2).contiguous()

634 + w13_qzeros = getattr(layer, "w13_qzeros", None)

635 + w2_qzeros = getattr(layer, "w2_qzeros", None)

636 + if w13_qzeros is not None:

637 + w13_qzeros = w13_qzeros.data.transpose(1, 2).contiguous()

638 + if w2_qzeros is not None:

639 + w2_qzeros = w2_qzeros.data.transpose(1, 2).contiguous()

640 +

641 converted = convert_to_wna16_moe_kernel_format(

642 backend=self.wna16_moe_backend,

643 layer=layer,

644 quant_config=self.quant_config,

645 input_dtype=self.input_dtype,

634 - w13=layer.w13_qweight,

635 - w2=layer.w2_qweight,

636 - w13_scale=layer.w13_scales,

637 - w2_scale=layer.w2_scales,

646 + w13=w13,

647 + w2=w2,

648 + w13_scale=w13_scale,

649 + w2_scale=w2_scale,

650 w13_bias=w13_bias,

651 w2_bias=w2_bias,

640 - w13_qzeros=getattr(layer, "w13_qzeros", None),

641 - w2_qzeros=getattr(layer, "w2_qzeros", None),

652 + w13_qzeros=w13_qzeros,

653 + w2_qzeros=w2_qzeros,

654 )

655

656 if converted is None:

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py

+29 -83

@@ -112,10 +112,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

112 WNA16MoEBackend.MARLIN,

113 WNA16MoEBackend.BATCHED_MARLIN,

114 ]

115 - self.is_transposed = self.wna16_backend not in (

116 - WNA16MoEBackend.FLASHINFER_TRTLLM,

117 - WNA16MoEBackend.HUMMING,

118 - )

115 + # CT checkpoints always load in N-first format [E, N, K_packed].

116 + # Backend-specific transposition is handled in

117 + # convert_to_wna16_moe_kernel_format.

118 + self.is_transposed = False

119

120 if self.is_marlin:

121 assert check_moe_marlin_supports_config(

@@ -141,80 +141,31 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

141 num_groups_w2: int | None = None,

142 num_groups_w13: int | None = None,

143 ) -> tuple[int, int, int]:

144 - """Get the shape of the weight based on the weight name, number of experts

145 - hidden size, intermediate size per partition, number of groups for w2,

146 - and number of groups for w13. Pass in num_groups_w2 and num_groups_w13

147 - for weight scales/zero_points.

144 + """Return N-first buffer shape for the given weight/scale/zp tensor.

145 +

146 + All CT MoE weights are allocated in the checkpoint's native N-first

147 + layout ``[E, N, K_packed]``. Backend-specific transposition (e.g. to

148 + Marlin's K-first layout) is deferred to

149 + ``convert_to_wna16_moe_kernel_format``.

150 """

149 - if weight_name in ("w13_scale", "w13_zp"):

150 - assert num_groups_w13 is not None, (

151 - "num_groups_w13 must be provided for weight scales/zero_points"

152 - )

153 - if weight_name in ("w2_scale", "w2_zp"):

154 - assert num_groups_w2 is not None, (

155 - "num_groups_w2 must be provided for weight scales/zero_points"

156 - )

151 w13_num_shards = 2 if self.moe.is_act_and_mul else 1

158 - shape_map: dict[str, dict[str, tuple[int, int | None, int | None]]] = {

159 - "w13_weight": {

160 - "Flashinfer": (

161 - num_experts,

162 - w13_num_shards * intermediate_size_per_partition,

163 - self._packed_dim(hidden_size),

164 - ),

165 - "Marlin": (

166 - num_experts,

167 - self._packed_dim(hidden_size),

168 - w13_num_shards * intermediate_size_per_partition,

169 - ),

170 - },

171 - "w13_scale": {

172 - "Flashinfer": (

173 - num_experts,

174 - w13_num_shards * intermediate_size_per_partition,

175 - num_groups_w13,

176 - ),

177 - "Marlin": (

178 - num_experts,

179 - num_groups_w13,

180 - w13_num_shards * intermediate_size_per_partition,

181 - ),

182 - },

183 - "w13_zp": {

184 - "Marlin": (

185 - num_experts,

186 - num_groups_w13,

187 - self._packed_dim(w13_num_shards * intermediate_size_per_partition),

188 - ),

189 - },

190 - "w2_weight": {

191 - "Flashinfer": (

192 - num_experts,

193 - hidden_size,

194 - self._packed_dim(intermediate_size_per_partition),

195 - ),

196 - "Marlin": (

197 - num_experts,

198 - self._packed_dim(intermediate_size_per_partition),

199 - hidden_size,

200 - ),

201 - },

202 - "w2_scale": {

203 - "Flashinfer": (num_experts, hidden_size, num_groups_w2),

204 - "Marlin": (num_experts, num_groups_w2, hidden_size),

205 - },

206 - "w2_zp": {

207 - "Marlin": (

208 - num_experts,

209 - num_groups_w2,

210 - self._packed_dim(hidden_size),

211 - ),

212 - },

152 + w13_n = w13_num_shards * intermediate_size_per_partition

153 + pf = self.packed_factor

154 + shape_map = {

155 + "w13_weight": (num_experts, w13_n, hidden_size // pf),

156 + "w13_scale": (num_experts, w13_n, num_groups_w13),

157 + "w13_zp": (num_experts, w13_n // pf, num_groups_w13),

158 + "w2_weight": (

159 + num_experts,

160 + hidden_size,

161 + intermediate_size_per_partition // pf,

162 + ),

163 + "w2_scale": (num_experts, hidden_size, num_groups_w2),

164 + "w2_zp": (num_experts, hidden_size // pf, num_groups_w2),

165 }

214 - backend_key = "Marlin" if self.is_transposed else "Flashinfer"

215 - shape = shape_map[weight_name][backend_key]

216 - assert shape[1] is not None and shape[2] is not None

217 - return shape[0], shape[1], shape[2]

166 + shape = shape_map[weight_name]

167 + assert isinstance(shape, tuple) and all(isinstance(d, int) for d in shape)

168 + return shape # type: ignore[return-value]

169

170 def create_weights(

171 self,

@@ -225,11 +176,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

176 params_dtype: torch.dtype,

177 **extra_weight_attrs,

178 ):

228 - # Will transpose the loaded weight along the

229 - # intermediate and hidden dim sizes. Will

230 - # shard for TP along the transposed dims

179 extra_weight_attrs.update(

232 - {"is_transposed": self.is_transposed, "quant_method": self.strategy}

180 + {"is_transposed": False, "quant_method": self.strategy}

181 )

182

183 w13_weight = torch.nn.Parameter(

@@ -440,12 +388,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

388 replace_parameter(layer, "w13_weight_scale", w13_scales)

389 replace_parameter(layer, "w2_weight_scale", w2_scales)

390

443 - # CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points

444 - # even for symmetric quant (the oracle synthesizes them).

391 + # CPU fused_experts_cpu requires zero points even for symmetric quant.

392 # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None.

393 if (

447 - not self.symmetric

448 - or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3)

394 + not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU

395 ) and self.wna16_backend != WNA16MoEBackend.EMULATION:

396 assert w13_qzeros is not None and w2_qzeros is not None

397 replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)