diff --git a/tests/kernels/quantization/test_int4_emulation_moe.py b/tests/kernels/quantization/test_int4_emulation_moe.py index 7016860d86..cad0dd94d2 100644 --- a/tests/kernels/quantization/test_int4_emulation_moe.py +++ b/tests/kernels/quantization/test_int4_emulation_moe.py @@ -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) diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 989bf967fa..77d7368ca1 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -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): diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index ad591d9ea7..791dc56766 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -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_`` 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, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 1d411ecc72..3de0aa8dab 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -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) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index ff57f3b8c5..ed430f7445 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -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: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 6eb2455d0e..c03d2bfcf4 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -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)