@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5de1add806

[feature]Add int4 quantization support for emulation moe backend (#48451)

Signed-off-by: Qiang Li <qiang.li2@amd.com>

qli88 committed Jul 15, 2026 at 21:00 UTC 5de1add806ce524ddd8c13286500ab8089e9654e
9 files changed +1653 -10
tests/kernels/quantization/test_int4_emulation_moe.py new
+1189
@@ -0,0 +1,1189 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""Correctness tests for Int4EmulationTritonExperts MoE backend.
5 +
6 +Tests the weight dequantization helpers (_unpack_and_dequant_int4_gptq,
7 +_unpack_and_dequant_int4_awq) and full MoE forward pass
8 +(_process_weights_emulation_gptq, _process_weights_emulation_awq)
9 +for both symmetric and asymmetric zero-point cases.
10 +
11 +Run `pytest tests/kernels/quantization/test_int4_emulation_moe.py`.
12 +"""
13 +
14 +import numpy
15 +import pytest
16 +import torch
17 +import torch.nn.functional as F
18 +
19 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
20 +from vllm.model_executor.layers.fused_moe.config import (
21 + FusedMoEConfig,
22 + FusedMoEParallelConfig,
23 + RoutingMethodType,
24 + int4_w4a16_moe_quant_config,
25 +)
26 +from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
27 + Int4EmulationTritonExperts,
28 +)
29 +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
30 + _process_weights_emulation_awq,
31 + _process_weights_emulation_gptq,
32 + _unpack_and_dequant_int4_awq,
33 + _unpack_and_dequant_int4_gptq,
34 +)
35 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
36 + awq_pack,
37 + gptq_pack,
38 +)
39 +from vllm.platforms import current_platform
40 +
41 +pytestmark = pytest.mark.skipif(
42 + not current_platform.is_cuda_alike(),
43 + reason="Int4EmulationTritonExperts requires CUDA.",
44 +)
45 +
46 +device = "cuda"
47 +
48 +# (E, K, N, group_size)
49 +SHAPES = [
50 + pytest.param(2, 64, 32, 32, id="tiny-gs32"),
51 + pytest.param(4, 128, 64, 64, id="small-gs64"),
52 + pytest.param(4, 256, 128, 128, id="medium-gs128"),
53 +]
54 +
55 +# (E, K, N, top_k, group_size, num_tokens)
56 +E2E_CONFIGS = [
57 + pytest.param(4, 64, 32, 2, 32, 8, id="tiny"),
58 + pytest.param(8, 128, 64, 2, 64, 16, id="small"),
59 +]
60 +
61 +
62 +# ---------------------------------------------------------------------------
63 +# Helpers
64 +# ---------------------------------------------------------------------------
65 +
66 +
67 +def _quantize_sym(w_fp: torch.Tensor, group_size: int):
68 + """Quantize [K, N] float to int4 symmetric (uint4b8), return q and scale."""
69 + K, N = w_fp.shape
70 + assert K % group_size == 0
71 + n_groups = K // group_size
72 + w_grouped = w_fp.reshape(n_groups, group_size, N)
73 + scale = w_grouped.abs().amax(dim=1) / 7.0
74 + scale = scale.clamp(min=1e-6)
75 + w_quant = (w_grouped / scale.unsqueeze(1)).round().clamp(-8, 7)
76 + q = (w_quant + 8).to(torch.int32).reshape(K, N)
77 + return q, scale
78 +
79 +
80 +def _quantize_asym(w_fp: torch.Tensor, group_size: int):
81 + """Quantize [K, N] float to uint4 asymmetric, return q, scale, zero."""
82 + K, N = w_fp.shape
83 + assert K % group_size == 0
84 + n_groups = K // group_size
85 + w_grouped = w_fp.reshape(n_groups, group_size, N)
86 + wmin = w_grouped.amin(dim=1)
87 + wmax = w_grouped.amax(dim=1)
88 + scale = (wmax - wmin) / 15.0
89 + scale = scale.clamp(min=1e-6)
90 + zero = (-wmin / scale).round().clamp(0, 15).to(torch.int32)
91 + w_quant = ((w_grouped - wmin.unsqueeze(1)) / scale.unsqueeze(1)).round()
92 + q = w_quant.clamp(0, 15).to(torch.int32).reshape(K, N)
93 + return q, scale, zero
94 +
95 +
96 +def _dequantize_ref(
97 + w_uint: torch.Tensor,
98 + scale: torch.Tensor,
99 + zero=None,
100 + output_dtype: torch.dtype = torch.bfloat16,
101 +):
102 + """Reference dequant for a single [K, N] slice."""
103 + K, N = w_uint.shape
104 + n_groups = scale.shape[0]
105 + group_size = K // n_groups
106 + w = w_uint.reshape(n_groups, group_size, N).to(output_dtype)
107 + s = scale.unsqueeze(1).to(output_dtype)
108 + if zero is None:
109 + return ((w - 8) * s).reshape(K, N)
110 + z = zero.unsqueeze(1).to(output_dtype)
111 + return ((w - z) * s).reshape(K, N)
112 +
113 +
114 +def _pack_gptq_zeros(zero: torch.Tensor, N: int) -> torch.Tensor:
115 + """Pack [n_groups, N] zeros into GPTQ format [n_groups, N//8] int32."""
116 + n_groups, _ = zero.shape
117 + z = zero.to(torch.int32).cpu().numpy().astype(numpy.uint32)
118 + packed = numpy.zeros((n_groups, N // 8), dtype=numpy.uint32)
119 + for i in range(8):
120 + packed |= z[:, i::8] << (i * 4)
121 + return torch.from_numpy(packed.astype(numpy.int32)).to(device)
122 +
123 +
124 +def _pack_awq_zeros(zero: torch.Tensor, N: int) -> torch.Tensor:
125 + """Pack [n_groups, N] zeros into AWQ column format [n_groups, N//8] int32."""
126 + n_groups, _ = zero.shape
127 + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7])
128 + z = zero.to(torch.int32).cpu().numpy().astype(numpy.uint32)
129 + z_interleaved = z.reshape(-1, 8)[:, interleave].reshape(n_groups, N)
130 + packed = numpy.zeros((n_groups, N // 8), dtype=numpy.uint32)
131 + for i in range(8):
132 + packed |= z_interleaved[:, i::8] << (i * 4)
133 + return torch.from_numpy(packed.astype(numpy.int32)).to(device)
134 +
135 +
136 +def _make_moe_config(E, K, N):
137 + return FusedMoEConfig(
138 + num_experts=E,
139 + experts_per_token=2,
140 + hidden_dim=K,
141 + intermediate_size=N,
142 + num_local_experts=E,
143 + num_logical_experts=E,
144 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
145 + activation=MoEActivation.SILU,
146 + in_dtype=torch.bfloat16,
147 + device=device,
148 + routing_method=RoutingMethodType.TopK,
149 + max_num_tokens=512,
150 + )
151 +
152 +
153 +def _make_gptq_moe_weights(E, K, N, group_size, asym=False):
154 + """Build GPTQ MoE weight tensors and per-expert float references."""
155 + torch.manual_seed(7)
156 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device)
157 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device)
158 +
159 + w13_list, w13s_list, w13z_list, w13_ref_list = [], [], [], []
160 + w2_list, w2s_list, w2z_list, w2_ref_list = [], [], [], []
161 +
162 + for e in range(E):
163 + if asym:
164 + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size)
165 + w13_list.append(gptq_pack(q13, 4, K, 2 * N))
166 + w13s_list.append(s13)
167 + w13z_list.append(_pack_gptq_zeros(z13, 2 * N))
168 + w13_ref_list.append(_dequantize_ref(q13, s13, z13))
169 + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size)
170 + w2_list.append(gptq_pack(q2, 4, N, K))
171 + w2s_list.append(s2)
172 + w2z_list.append(_pack_gptq_zeros(z2, K))
173 + w2_ref_list.append(_dequantize_ref(q2, s2, z2))
174 + else:
175 + q13, s13 = _quantize_sym(w13_fp[e], group_size)
176 + w13_list.append(gptq_pack(q13, 4, K, 2 * N))
177 + w13s_list.append(s13)
178 + w13z_list.append(None)
179 + w13_ref_list.append(_dequantize_ref(q13, s13))
180 + q2, s2 = _quantize_sym(w2_fp[e], group_size)
181 + w2_list.append(gptq_pack(q2, 4, N, K))
182 + w2s_list.append(s2)
183 + w2z_list.append(None)
184 + w2_ref_list.append(_dequantize_ref(q2, s2))
185 +
186 + return (
187 + torch.stack(w13_list),
188 + torch.stack(w13s_list),
189 + torch.stack(w13z_list) if asym else None,
190 + torch.stack(w13_ref_list), # [E, K, 2N]
191 + torch.stack(w2_list),
192 + torch.stack(w2s_list),
193 + torch.stack(w2z_list) if asym else None,
194 + torch.stack(w2_ref_list), # [E, N, K]
195 + )
196 +
197 +
198 +def _make_awq_moe_weights(E, K, N, group_size, asym=False):
199 + """Build AWQ MoE weight tensors and per-expert float references."""
200 + torch.manual_seed(8)
201 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device)
202 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device)
203 +
204 + w13_list, w13s_list, w13z_list, w13_ref_list = [], [], [], []
205 + w2_list, w2s_list, w2z_list, w2_ref_list = [], [], [], []
206 +
207 + for e in range(E):
208 + if asym:
209 + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size)
210 + w13_list.append(awq_pack(q13, 4, K, 2 * N))
211 + w13s_list.append(s13)
212 + w13z_list.append(_pack_awq_zeros(z13, 2 * N))
213 + w13_ref_list.append(_dequantize_ref(q13, s13, z13))
214 + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size)
215 + w2_list.append(awq_pack(q2, 4, N, K))
216 + w2s_list.append(s2)
217 + w2z_list.append(_pack_awq_zeros(z2, K))
218 + w2_ref_list.append(_dequantize_ref(q2, s2, z2))
219 + else:
220 + q13, s13 = _quantize_sym(w13_fp[e], group_size)
221 + w13_list.append(awq_pack(q13, 4, K, 2 * N))
222 + w13s_list.append(s13)
223 + w13z_list.append(None)
224 + w13_ref_list.append(_dequantize_ref(q13, s13))
225 + q2, s2 = _quantize_sym(w2_fp[e], group_size)
226 + w2_list.append(awq_pack(q2, 4, N, K))
227 + w2s_list.append(s2)
228 + w2z_list.append(None)
229 + w2_ref_list.append(_dequantize_ref(q2, s2))
230 +
231 + return (
232 + torch.stack(w13_list),
233 + torch.stack(w13s_list),
234 + torch.stack(w13z_list) if asym else None,
235 + torch.stack(w13_ref_list), # [E, K, 2N]
236 + torch.stack(w2_list),
237 + torch.stack(w2s_list),
238 + torch.stack(w2z_list) if asym else None,
239 + torch.stack(w2_ref_list), # [E, N, K]
240 + )
241 +
242 +
243 +def _run_emulation_forward(
244 + experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N
245 +):
246 + ws13_size = hidden_states.shape[0] * topk_ids.shape[1] * max(N, K)
247 + ws2_size = hidden_states.shape[0] * topk_ids.shape[1] * max(2 * N, K)
248 + workspace13 = torch.zeros(ws13_size, dtype=hidden_states.dtype, device=device)
249 + workspace2 = torch.zeros(ws2_size, dtype=hidden_states.dtype, device=device)
250 + output = torch.zeros(
251 + hidden_states.shape[0], K, dtype=hidden_states.dtype, device=device
252 + )
253 + experts.apply(
254 + output=output,
255 + hidden_states=hidden_states,
256 + w1=w13_bf16,
257 + w2=w2_bf16,
258 + topk_weights=topk_weights,
259 + topk_ids=topk_ids,
260 + activation=MoEActivation.SILU,
261 + global_num_experts=E,
262 + expert_map=None,
263 + a1q_scale=None,
264 + a2_scale=None,
265 + workspace13=workspace13,
266 + workspace2=workspace2,
267 + expert_tokens_meta=None,
268 + apply_router_weight_on_input=False,
269 + )
270 + return output
271 +
272 +
273 +# ---------------------------------------------------------------------------
274 +# Tests: _unpack_and_dequant_int4_gptq
275 +# ---------------------------------------------------------------------------
276 +
277 +
278 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
279 +def test_gptq_unpack_symmetric(E, K, N, group_size):
280 + """GPTQ symmetric unpacker matches reference."""
281 + torch.manual_seed(0)
282 + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device)
283 +
284 + packed_list, scale_list, ref_list = [], [], []
285 + for e in range(E):
286 + q, s = _quantize_sym(w_fp[e], group_size)
287 + packed_list.append(gptq_pack(q, 4, K, N))
288 + scale_list.append(s)
289 + ref_list.append(_dequantize_ref(q, s, output_dtype=torch.float32))
290 +
291 + w_packed = torch.stack(packed_list).to(device)
292 + scale = torch.stack(scale_list).to(device)
293 + ref = torch.stack(ref_list).to(device)
294 +
295 + out = _unpack_and_dequant_int4_gptq(
296 + w_packed, scale, None, transpose_output=False, output_dtype=torch.float32
297 + )
298 +
299 + assert out.shape == (E, K, N)
300 + assert torch.allclose(out, ref, atol=0), (
301 + f"max diff: {(out - ref).abs().max().item()}"
302 + )
303 +
304 +
305 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
306 +def test_gptq_unpack_asymmetric(E, K, N, group_size):
307 + """GPTQ asymmetric unpacker matches reference."""
308 + torch.manual_seed(1)
309 + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device)
310 +
311 + packed_list, scale_list, zero_list, ref_list = [], [], [], []
312 + for e in range(E):
313 + q, s, z = _quantize_asym(w_fp[e], group_size)
314 + packed_list.append(gptq_pack(q, 4, K, N))
315 + scale_list.append(s)
316 + zero_list.append(_pack_gptq_zeros(z, N))
317 + ref_list.append(_dequantize_ref(q, s, z, output_dtype=torch.float32))
318 +
319 + w_packed = torch.stack(packed_list).to(device)
320 + scale = torch.stack(scale_list).to(device)
321 + qzeros = torch.stack(zero_list).to(device)
322 + ref = torch.stack(ref_list).to(device)
323 +
324 + out = _unpack_and_dequant_int4_gptq(
325 + w_packed, scale, qzeros, transpose_output=False, output_dtype=torch.float32
326 + )
327 +
328 + assert out.shape == (E, K, N)
329 + assert torch.allclose(out, ref, atol=0), (
330 + f"max diff: {(out - ref).abs().max().item()}"
331 + )
332 +
333 +
334 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
335 +def test_gptq_unpack_transpose(E, K, N, group_size):
336 + """GPTQ transpose_output=True gives [E, N, K]."""
337 + torch.manual_seed(2)
338 + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device)
339 +
340 + packed_list, scale_list = [], []
341 + for e in range(E):
342 + q, s = _quantize_sym(w_fp[e], group_size)
343 + packed_list.append(gptq_pack(q, 4, K, N))
344 + scale_list.append(s)
345 +
346 + w_packed = torch.stack(packed_list).to(device)
347 + scale = torch.stack(scale_list).to(device)
348 +
349 + out_normal = _unpack_and_dequant_int4_gptq(w_packed, scale, None, False)
350 + out_transposed = _unpack_and_dequant_int4_gptq(w_packed, scale, None, True)
351 +
352 + assert out_transposed.shape == (E, N, K)
353 + assert torch.allclose(
354 + out_transposed, out_normal.permute(0, 2, 1).contiguous(), atol=0
355 + )
356 +
357 +
358 +# ---------------------------------------------------------------------------
359 +# Tests: _unpack_and_dequant_int4_awq
360 +# ---------------------------------------------------------------------------
361 +
362 +
363 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
364 +def test_awq_unpack_symmetric(E, K, N, group_size):
365 + """AWQ symmetric unpacker matches reference."""
366 + torch.manual_seed(3)
367 + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device)
368 +
369 + packed_list, scale_list, ref_list = [], [], []
370 + for e in range(E):
371 + q, s = _quantize_sym(w_fp[e], group_size)
372 + packed_list.append(awq_pack(q, 4, K, N))
373 + scale_list.append(s)
374 + ref_list.append(_dequantize_ref(q, s, output_dtype=torch.float32))
375 +
376 + w_packed = torch.stack(packed_list).to(device)
377 + scale = torch.stack(scale_list).to(device)
378 + ref = torch.stack(ref_list).to(device)
379 +
380 + out = _unpack_and_dequant_int4_awq(
381 + w_packed, scale, None, transpose_output=False, output_dtype=torch.float32
382 + )
383 +
384 + assert out.shape == (E, K, N)
385 + assert torch.allclose(out, ref, atol=0), (
386 + f"max diff: {(out - ref).abs().max().item()}"
387 + )
388 +
389 +
390 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
391 +def test_awq_unpack_asymmetric(E, K, N, group_size):
392 + """AWQ asymmetric unpacker matches reference."""
393 + torch.manual_seed(4)
394 + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device)
395 +
396 + packed_list, scale_list, zero_list, ref_list = [], [], [], []
397 + for e in range(E):
398 + q, s, z = _quantize_asym(w_fp[e], group_size)
399 + packed_list.append(awq_pack(q, 4, K, N))
400 + scale_list.append(s)
401 + zero_list.append(_pack_awq_zeros(z, N))
402 + ref_list.append(_dequantize_ref(q, s, z, output_dtype=torch.float32))
403 +
404 + w_packed = torch.stack(packed_list).to(device)
405 + scale = torch.stack(scale_list).to(device)
406 + qzeros = torch.stack(zero_list).to(device)
407 + ref = torch.stack(ref_list).to(device)
408 +
409 + out = _unpack_and_dequant_int4_awq(
410 + w_packed, scale, qzeros, transpose_output=False, output_dtype=torch.float32
411 + )
412 +
413 + assert out.shape == (E, K, N)
414 + assert torch.allclose(out, ref, atol=0), (
415 + f"max diff: {(out - ref).abs().max().item()}"
416 + )
417 +
418 +
419 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
420 +def test_awq_unpack_transpose(E, K, N, group_size):
421 + """AWQ transpose_output=True gives [E, N, K]."""
422 + torch.manual_seed(5)
423 + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device)
424 +
425 + packed_list, scale_list = [], []
426 + for e in range(E):
427 + q, s = _quantize_sym(w_fp[e], group_size)
428 + packed_list.append(awq_pack(q, 4, K, N))
429 + scale_list.append(s)
430 +
431 + w_packed = torch.stack(packed_list).to(device)
432 + scale = torch.stack(scale_list).to(device)
433 +
434 + out_normal = _unpack_and_dequant_int4_awq(w_packed, scale, None, False)
435 + out_transposed = _unpack_and_dequant_int4_awq(w_packed, scale, None, True)
436 +
437 + assert out_transposed.shape == (E, N, K)
438 + assert torch.allclose(
439 + out_transposed, out_normal.permute(0, 2, 1).contiguous(), atol=0
440 + )
441 +
442 +
443 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
444 +def test_awq_gptq_unpack_agree(E, K, N, group_size):
445 + """AWQ and GPTQ unpackers produce identical values for the same weights."""
446 + torch.manual_seed(6)
447 + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device)
448 +
449 + gptq_list, awq_list, scale_list = [], [], []
450 + for e in range(E):
451 + q, s = _quantize_sym(w_fp[e], group_size)
452 + gptq_list.append(gptq_pack(q, 4, K, N))
453 + awq_list.append(awq_pack(q, 4, K, N))
454 + scale_list.append(s)
455 +
456 + scale = torch.stack(scale_list).to(device)
457 + out_gptq = _unpack_and_dequant_int4_gptq(
458 + torch.stack(gptq_list).to(device), scale, None, False, torch.float32
459 + )
460 + out_awq = _unpack_and_dequant_int4_awq(
461 + torch.stack(awq_list).to(device), scale, None, False, torch.float32
462 + )
463 +
464 + assert torch.allclose(out_gptq, out_awq, atol=0), (
465 + f"max diff: {(out_gptq - out_awq).abs().max().item()}"
466 + )
467 +
468 +
469 +# ---------------------------------------------------------------------------
470 +# Tests: _process_weights_emulation_{gptq,awq}
471 +# ---------------------------------------------------------------------------
472 +
473 +
474 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
475 +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"])
476 +def test_gptq_process_weights_shapes_and_values(E, K, N, group_size, asym):
477 + """_process_weights_emulation_gptq shapes and values match reference."""
478 + w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_gptq_moe_weights(
479 + E, K, N, group_size, asym
480 + )
481 + result = _process_weights_emulation_gptq(w13, w2, w13s, w2s, w13z, w2z)
482 + w13_out, w2_out = result[0], result[1]
483 +
484 + assert w13_out.shape == (E, 2 * N, K)
485 + assert w2_out.shape == (E, K, N)
486 + assert w13_out.dtype == torch.bfloat16
487 + assert w2_out.dtype == torch.bfloat16
488 +
489 + expected_w13 = w13_ref.permute(0, 2, 1)
490 + expected_w2 = w2_ref.permute(0, 2, 1)
491 +
492 + assert torch.allclose(w13_out.float(), expected_w13.float(), atol=0), (
493 + f"w13 max diff: {(w13_out.float() - expected_w13.float()).abs().max().item()}"
494 + )
495 + assert torch.allclose(w2_out.float(), expected_w2.float(), atol=0), (
496 + f"w2 max diff: {(w2_out.float() - expected_w2.float()).abs().max().item()}"
497 + )
498 +
499 +
500 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
501 +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"])
502 +def test_awq_process_weights_shapes_and_values(E, K, N, group_size, asym):
503 + """_process_weights_emulation_awq shapes and values match reference."""
504 + w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_awq_moe_weights(
505 + E, K, N, group_size, asym
506 + )
507 + result = _process_weights_emulation_awq(w13, w2, w13s, w2s, w13z, w2z)
508 + w13_out, w2_out = result[0], result[1]
509 +
510 + assert w13_out.shape == (E, 2 * N, K)
511 + assert w2_out.shape == (E, K, N)
512 + assert w13_out.dtype == torch.bfloat16
513 + assert w2_out.dtype == torch.bfloat16
514 +
515 + expected_w13 = w13_ref.permute(0, 2, 1)
516 + expected_w2 = w2_ref.permute(0, 2, 1)
517 +
518 + assert torch.allclose(w13_out.float(), expected_w13.float(), atol=0), (
519 + f"w13 max diff: {(w13_out.float() - expected_w13.float()).abs().max().item()}"
520 + )
521 + assert torch.allclose(w2_out.float(), expected_w2.float(), atol=0), (
522 + f"w2 max diff: {(w2_out.float() - expected_w2.float()).abs().max().item()}"
523 + )
524 +
525 +
526 +@pytest.mark.parametrize("E, K, N, group_size", SHAPES)
527 +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"])
528 +def test_gptq_awq_process_weights_agree(E, K, N, group_size, asym):
529 + """AWQ and GPTQ process_weights produce identical dequantized tensors."""
530 + torch.manual_seed(9)
531 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device)
532 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device)
533 +
534 + g13_list, g13s_list, g13z_list = [], [], []
535 + a13_list, a13s_list, a13z_list = [], [], []
536 + g2_list, g2s_list, g2z_list = [], [], []
537 + a2_list, a2s_list, a2z_list = [], [], []
538 +
539 + for e in range(E):
540 + if asym:
541 + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size)
542 + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size)
543 + g13z_list.append(_pack_gptq_zeros(z13, 2 * N))
544 + a13z_list.append(_pack_awq_zeros(z13, 2 * N))
545 + g2z_list.append(_pack_gptq_zeros(z2, K))
546 + a2z_list.append(_pack_awq_zeros(z2, K))
547 + else:
548 + q13, s13 = _quantize_sym(w13_fp[e], group_size)
549 + q2, s2 = _quantize_sym(w2_fp[e], group_size)
550 + g13z_list.append(None)
551 + a13z_list.append(None)
552 + g2z_list.append(None)
553 + a2z_list.append(None)
554 +
555 + g13_list.append(gptq_pack(q13, 4, K, 2 * N))
556 + a13_list.append(awq_pack(q13, 4, K, 2 * N))
557 + g13s_list.append(s13)
558 + a13s_list.append(s13)
559 + g2_list.append(gptq_pack(q2, 4, N, K))
560 + a2_list.append(awq_pack(q2, 4, N, K))
561 + g2s_list.append(s2)
562 + a2s_list.append(s2)
563 +
564 + gptq_res = _process_weights_emulation_gptq(
565 + torch.stack(g13_list),
566 + torch.stack(g2_list),
567 + torch.stack(g13s_list),
568 + torch.stack(g2s_list),
569 + torch.stack(g13z_list) if asym else None,
570 + torch.stack(g2z_list) if asym else None,
571 + )
572 + awq_res = _process_weights_emulation_awq(
573 + torch.stack(a13_list),
574 + torch.stack(a2_list),
575 + torch.stack(a13s_list),
576 + torch.stack(a2s_list),
577 + torch.stack(a13z_list) if asym else None,
578 + torch.stack(a2z_list) if asym else None,
579 + )
580 +
581 + assert torch.allclose(gptq_res[0].float(), awq_res[0].float(), atol=1e-3), (
582 + f"w13 max diff: {(gptq_res[0] - awq_res[0]).float().abs().max().item()}"
583 + )
584 + assert torch.allclose(gptq_res[1].float(), awq_res[1].float(), atol=1e-3), (
585 + f"w2 max diff: {(gptq_res[1] - awq_res[1]).float().abs().max().item()}"
586 + )
587 +
588 +
589 +# ---------------------------------------------------------------------------
590 +# End-to-end MoE forward pass tests
591 +# ---------------------------------------------------------------------------
592 +
593 +
594 +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens", E2E_CONFIGS)
595 +def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens):
596 + """GPTQ and AWQ emulation backends produce bit-identical forward outputs."""
597 + torch.manual_seed(42)
598 + moe_config = _make_moe_config(E, K, N)
599 +
600 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device)
601 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device)
602 +
603 + g13_list, g13s_list, a13_list, a13s_list = [], [], [], []
604 + g2_list, g2s_list, a2_list, a2s_list = [], [], [], []
605 +
606 + for e in range(E):
607 + q13, s13 = _quantize_sym(w13_fp[e], group_size)
608 + q2, s2 = _quantize_sym(w2_fp[e], group_size)
609 + g13_list.append(gptq_pack(q13, 4, K, 2 * N))
610 + a13_list.append(awq_pack(q13, 4, K, 2 * N))
611 + g13s_list.append(s13)
612 + a13s_list.append(s13.clone())
613 + g2_list.append(gptq_pack(q2, 4, N, K))
614 + a2_list.append(awq_pack(q2, 4, N, K))
615 + g2s_list.append(s2)
616 + a2s_list.append(s2.clone())
617 +
618 + gptq_res = _process_weights_emulation_gptq(
619 + torch.stack(g13_list),
620 + torch.stack(g2_list),
621 + torch.stack(g13s_list),
622 + torch.stack(g2s_list),
623 + None,
624 + None,
625 + )
626 + awq_res = _process_weights_emulation_awq(
627 + torch.stack(a13_list),
628 + torch.stack(a2_list),
629 + torch.stack(a13s_list),
630 + torch.stack(a2s_list),
631 + None,
632 + None,
633 + )
634 + w13_gptq, w2_gptq = gptq_res[0], gptq_res[1]
635 + w13_awq, w2_awq = awq_res[0], awq_res[1]
636 +
637 + dummy_scale = torch.ones(1, dtype=torch.float16, device=device)
638 + experts_gptq = Int4EmulationTritonExperts(
639 + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
640 + )
641 + experts_awq = Int4EmulationTritonExperts(
642 + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
643 + )
644 +
645 + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
646 + topk_weights = torch.softmax(
647 + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1
648 + )
649 + topk_ids = torch.stack(
650 + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)]
651 + ).to(torch.int32)
652 +
653 + out_gptq = _run_emulation_forward(
654 + experts_gptq, w13_gptq, w2_gptq, hidden_states, topk_weights, topk_ids, E, K, N
655 + )
656 + out_awq = _run_emulation_forward(
657 + experts_awq, w13_awq, w2_awq, hidden_states, topk_weights, topk_ids, E, K, N
658 + )
659 +
660 + assert torch.allclose(out_gptq, out_awq, atol=0), (
661 + f"max diff: {(out_gptq - out_awq).abs().max().item()}"
662 + )
663 +
664 +
665 +# ---------------------------------------------------------------------------
666 +# EP (Expert Parallelism) tests
667 +# ---------------------------------------------------------------------------
668 +
669 +# (E, K, N, top_k, group_size, num_tokens, ep_size)
670 +EP_CONFIGS = [
671 + pytest.param(4, 64, 32, 2, 32, 8, 2, id="E4-ep2"),
672 + pytest.param(8, 64, 32, 2, 32, 16, 4, id="E8-ep4"),
673 + pytest.param(8, 128, 64, 2, 64, 16, 2, id="E8-ep2"),
674 +]
675 +
676 +
677 +def _make_expert_map(global_num_experts: int, start: int, end: int) -> torch.Tensor:
678 + """Build expert_map for a rank that owns experts [start, end)."""
679 + expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32, device=device)
680 + expert_map[start:end] = torch.arange(end - start, dtype=torch.int32, device=device)
681 + return expert_map
682 +
683 +
684 +def _run_emulation_forward_ep(
685 + experts,
686 + w13_bf16,
687 + w2_bf16,
688 + hidden_states,
689 + topk_weights,
690 + topk_ids,
691 + global_num_experts,
692 + expert_map,
693 +):
694 + """Run forward with EP expert_map; returns output tensor."""
695 + T, K = hidden_states.shape
696 + N = w2_bf16.shape[2]
697 + ws13_size = T * topk_ids.shape[1] * max(N, K)
698 + ws2_size = T * topk_ids.shape[1] * max(2 * N, K)
699 + workspace13 = torch.zeros(ws13_size, dtype=hidden_states.dtype, device=device)
700 + workspace2 = torch.zeros(ws2_size, dtype=hidden_states.dtype, device=device)
701 + output = torch.zeros(T, K, dtype=hidden_states.dtype, device=device)
702 + experts.apply(
703 + output=output,
704 + hidden_states=hidden_states,
705 + w1=w13_bf16,
706 + w2=w2_bf16,
707 + topk_weights=topk_weights,
708 + topk_ids=topk_ids,
709 + activation=MoEActivation.SILU,
710 + global_num_experts=global_num_experts,
711 + expert_map=expert_map,
712 + a1q_scale=None,
713 + a2_scale=None,
714 + workspace13=workspace13,
715 + workspace2=workspace2,
716 + expert_tokens_meta=None,
717 + apply_router_weight_on_input=False,
718 + )
719 + return output
720 +
721 +
722 +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS)
723 +@pytest.mark.parametrize("fmt", ["gptq", "awq"])
724 +def test_ep_output_matches_no_ep(E, K, N, top_k, group_size, num_tokens, ep_size, fmt):
725 + """EP simulation: sum of per-rank outputs equals the no-EP forward pass."""
726 + assert E % ep_size == 0
727 + num_local = E // ep_size
728 +
729 + torch.manual_seed(20)
730 +
731 + # Build all expert weights in BF16 (no-EP reference)
732 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02
733 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02
734 +
735 + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], []
736 + for e in range(E):
737 + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size)
738 + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size)
739 + if fmt == "gptq":
740 + packed13_list.append(gptq_pack(q13, 4, K, 2 * N))
741 + packed2_list.append(gptq_pack(q2, 4, N, K))
742 + else:
743 + packed13_list.append(awq_pack(q13, 4, K, 2 * N))
744 + packed2_list.append(awq_pack(q2, 4, N, K))
745 + scales13_list.append(s13)
746 + scales2_list.append(s2)
747 +
748 + process_fn = (
749 + _process_weights_emulation_gptq
750 + if fmt == "gptq"
751 + else _process_weights_emulation_awq
752 + )
753 + res = process_fn(
754 + torch.stack(packed13_list),
755 + torch.stack(packed2_list),
756 + torch.stack(scales13_list),
757 + torch.stack(scales2_list),
758 + None,
759 + None,
760 + )
761 + w13_all, w2_all = res[0], res[1] # [E, 2N, K], [E, K, N]
762 +
763 + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
764 + topk_weights = torch.softmax(
765 + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1
766 + )
767 + topk_ids = torch.stack(
768 + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)]
769 + ).to(torch.int32)
770 +
771 + # No-EP reference
772 + dummy_scale = torch.ones(1, dtype=torch.float16, device=device)
773 + moe_config_full = FusedMoEConfig(
774 + num_experts=E,
775 + experts_per_token=top_k,
776 + hidden_dim=K,
777 + intermediate_size=N,
778 + num_local_experts=E,
779 + num_logical_experts=E,
780 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
781 + activation=MoEActivation.SILU,
782 + in_dtype=torch.bfloat16,
783 + device=device,
784 + routing_method=RoutingMethodType.TopK,
785 + max_num_tokens=512,
786 + )
787 + experts_ref = Int4EmulationTritonExperts(
788 + moe_config_full, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
789 + )
790 + out_no_ep = _run_emulation_forward(
791 + experts_ref, w13_all, w2_all, hidden_states, topk_weights, topk_ids, E, K, N
792 + )
793 +
794 + # EP simulation: sum contributions from each rank
795 + out_ep_sum = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device)
796 + for rank in range(ep_size):
797 + start = rank * num_local
798 + end = start + num_local
799 + expert_map = _make_expert_map(E, start, end)
800 + moe_config_ep = FusedMoEConfig(
801 + num_experts=E,
802 + experts_per_token=top_k,
803 + hidden_dim=K,
804 + intermediate_size=N,
805 + num_local_experts=num_local,
806 + num_logical_experts=E,
807 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
808 + activation=MoEActivation.SILU,
809 + in_dtype=torch.bfloat16,
810 + device=device,
811 + routing_method=RoutingMethodType.TopK,
812 + max_num_tokens=512,
813 + )
814 + experts_ep = Int4EmulationTritonExperts(
815 + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
816 + )
817 + out_rank = _run_emulation_forward_ep(
818 + experts_ep,
819 + w13_all[start:end],
820 + w2_all[start:end],
821 + hidden_states,
822 + topk_weights,
823 + topk_ids,
824 + global_num_experts=E,
825 + expert_map=expert_map,
826 + )
827 + out_ep_sum = out_ep_sum + out_rank
828 +
829 + assert torch.allclose(out_ep_sum, out_no_ep, atol=1e-3), (
830 + f"[{fmt}] EP sum max diff: {(out_ep_sum - out_no_ep).abs().max().item():.6f}"
831 + )
832 +
833 +
834 +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS)
835 +def test_ep_gptq_awq_agree(E, K, N, top_k, group_size, num_tokens, ep_size):
836 + """With EP, GPTQ and AWQ emulation produce the same outputs per rank."""
837 + assert E % ep_size == 0
838 + num_local = E // ep_size
839 +
840 + torch.manual_seed(21)
841 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02
842 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02
843 +
844 + g13_list, g13s_list, g2_list, g2s_list = [], [], [], []
845 + a13_list, a13s_list, a2_list, a2s_list = [], [], [], []
846 + for e in range(E):
847 + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size)
848 + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size)
849 + g13_list.append(gptq_pack(q13, 4, K, 2 * N))
850 + a13_list.append(awq_pack(q13, 4, K, 2 * N))
851 + g13s_list.append(s13)
852 + a13s_list.append(s13.clone())
853 + g2_list.append(gptq_pack(q2, 4, N, K))
854 + a2_list.append(awq_pack(q2, 4, N, K))
855 + g2s_list.append(s2)
856 + a2s_list.append(s2.clone())
857 +
858 + gptq_res = _process_weights_emulation_gptq(
859 + torch.stack(g13_list),
860 + torch.stack(g2_list),
861 + torch.stack(g13s_list),
862 + torch.stack(g2s_list),
863 + None,
864 + None,
865 + )
866 + awq_res = _process_weights_emulation_awq(
867 + torch.stack(a13_list),
868 + torch.stack(a2_list),
869 + torch.stack(a13s_list),
870 + torch.stack(a2s_list),
871 + None,
872 + None,
873 + )
874 + w13_gptq, w2_gptq = gptq_res[0], gptq_res[1]
875 + w13_awq, w2_awq = awq_res[0], awq_res[1]
876 +
877 + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
878 + topk_weights = torch.softmax(
879 + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1
880 + )
881 + topk_ids = torch.stack(
882 + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)]
883 + ).to(torch.int32)
884 +
885 + dummy_scale = torch.ones(1, dtype=torch.float16, device=device)
886 +
887 + for rank in range(ep_size):
888 + start = rank * num_local
889 + end = start + num_local
890 + expert_map = _make_expert_map(E, start, end)
891 + moe_config_ep = FusedMoEConfig(
892 + num_experts=E,
893 + experts_per_token=top_k,
894 + hidden_dim=K,
895 + intermediate_size=N,
896 + num_local_experts=num_local,
897 + num_logical_experts=E,
898 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
899 + activation=MoEActivation.SILU,
900 + in_dtype=torch.bfloat16,
901 + device=device,
902 + routing_method=RoutingMethodType.TopK,
903 + max_num_tokens=512,
904 + )
905 + experts_gptq = Int4EmulationTritonExperts(
906 + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
907 + )
908 + experts_awq = Int4EmulationTritonExperts(
909 + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
910 + )
911 + out_gptq = _run_emulation_forward_ep(
912 + experts_gptq,
913 + w13_gptq[start:end],
914 + w2_gptq[start:end],
915 + hidden_states,
916 + topk_weights,
917 + topk_ids,
918 + E,
919 + expert_map,
920 + )
921 + out_awq = _run_emulation_forward_ep(
922 + experts_awq,
923 + w13_awq[start:end],
924 + w2_awq[start:end],
925 + hidden_states,
926 + topk_weights,
927 + topk_ids,
928 + E,
929 + expert_map,
930 + )
931 + assert torch.allclose(out_gptq, out_awq, atol=0), (
932 + f"rank={rank} max diff: {(out_gptq - out_awq).abs().max().item()}"
933 + )
934 +
935 +
936 +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS)
937 +def test_ep_partial_rank_no_active_experts(
938 + E, K, N, top_k, group_size, num_tokens, ep_size
939 +):
940 + """A rank that owns no token-selected experts produces an all-zero output."""
941 + assert E % ep_size == 0
942 + num_local = E // ep_size
943 +
944 + torch.manual_seed(22)
945 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02
946 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02
947 +
948 + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], []
949 + for e in range(E):
950 + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size)
951 + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size)
952 + packed13_list.append(gptq_pack(q13, 4, K, 2 * N))
953 + packed2_list.append(gptq_pack(q2, 4, N, K))
954 + scales13_list.append(s13)
955 + scales2_list.append(s2)
956 +
957 + res = _process_weights_emulation_gptq(
958 + torch.stack(packed13_list),
959 + torch.stack(packed2_list),
960 + torch.stack(scales13_list),
961 + torch.stack(scales2_list),
962 + None,
963 + None,
964 + )
965 + w13_all, w2_all = res[0], res[1]
966 +
967 + # Force topk_ids to only use experts in [0, num_local) — rank 0's slice
968 + topk_ids = torch.zeros(num_tokens, top_k, dtype=torch.int32, device=device)
969 + topk_weights = torch.softmax(
970 + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1
971 + )
972 + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
973 +
974 + dummy_scale = torch.ones(1, dtype=torch.float16, device=device)
975 +
976 + # Last rank owns experts [E-num_local, E), tokens only route to [0, num_local)
977 + last_rank = ep_size - 1
978 + start = last_rank * num_local
979 + end = E
980 + expert_map = _make_expert_map(E, start, end)
981 + moe_config_ep = FusedMoEConfig(
982 + num_experts=E,
983 + experts_per_token=top_k,
984 + hidden_dim=K,
985 + intermediate_size=N,
986 + num_local_experts=num_local,
987 + num_logical_experts=E,
988 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
989 + activation=MoEActivation.SILU,
990 + in_dtype=torch.bfloat16,
991 + device=device,
992 + routing_method=RoutingMethodType.TopK,
993 + max_num_tokens=512,
994 + )
995 + experts_ep = Int4EmulationTritonExperts(
996 + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
997 + )
998 + out = _run_emulation_forward_ep(
999 + experts_ep,
1000 + w13_all[start:end],
1001 + w2_all[start:end],
1002 + hidden_states,
1003 + topk_weights,
1004 + topk_ids,
1005 + E,
1006 + expert_map,
1007 + )
1008 + assert torch.all(out == 0), (
1009 + f"Expected zeros for inactive rank, got max={out.abs().max().item()}"
1010 + )
1011 +
1012 +
1013 +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS)
1014 +def test_ep_sum_equals_full_forward(E, K, N, top_k, group_size, num_tokens, ep_size):
1015 + """With fixed routing, EP rank outputs sum to the single-rank full forward."""
1016 + assert E % ep_size == 0
1017 + num_local = E // ep_size
1018 +
1019 + torch.manual_seed(23)
1020 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02
1021 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02
1022 +
1023 + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], []
1024 + for e in range(E):
1025 + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size)
1026 + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size)
1027 + packed13_list.append(gptq_pack(q13, 4, K, 2 * N))
1028 + packed2_list.append(gptq_pack(q2, 4, N, K))
1029 + scales13_list.append(s13)
1030 + scales2_list.append(s2)
1031 +
1032 + res = _process_weights_emulation_gptq(
1033 + torch.stack(packed13_list),
1034 + torch.stack(packed2_list),
1035 + torch.stack(scales13_list),
1036 + torch.stack(scales2_list),
1037 + None,
1038 + None,
1039 + )
1040 + w13_all, w2_all = res[0], res[1]
1041 +
1042 + # Fix routing so every token uses exactly 2 consecutive experts (round-robin)
1043 + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
1044 + topk_ids = torch.stack(
1045 + [
1046 + torch.tensor([(t * top_k + k) % E for k in range(top_k)], dtype=torch.int32)
1047 + for t in range(num_tokens)
1048 + ]
1049 + ).to(device)
1050 + topk_weights = torch.full((num_tokens, top_k), 1.0 / top_k, device=device)
1051 +
1052 + dummy_scale = torch.ones(1, dtype=torch.float16, device=device)
1053 +
1054 + # Full (no-EP) reference
1055 + moe_config_full = FusedMoEConfig(
1056 + num_experts=E,
1057 + experts_per_token=top_k,
1058 + hidden_dim=K,
1059 + intermediate_size=N,
1060 + num_local_experts=E,
1061 + num_logical_experts=E,
1062 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1063 + activation=MoEActivation.SILU,
1064 + in_dtype=torch.bfloat16,
1065 + device=device,
1066 + routing_method=RoutingMethodType.TopK,
1067 + max_num_tokens=512,
1068 + )
1069 + experts_full = Int4EmulationTritonExperts(
1070 + moe_config_full, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
1071 + )
1072 + out_full = _run_emulation_forward(
1073 + experts_full, w13_all, w2_all, hidden_states, topk_weights, topk_ids, E, K, N
1074 + )
1075 +
1076 + # EP sum
1077 + out_ep_sum = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device)
1078 + for rank in range(ep_size):
1079 + start = rank * num_local
1080 + end = start + num_local
1081 + expert_map = _make_expert_map(E, start, end)
1082 + moe_config_ep = FusedMoEConfig(
1083 + num_experts=E,
1084 + experts_per_token=top_k,
1085 + hidden_dim=K,
1086 + intermediate_size=N,
1087 + num_local_experts=num_local,
1088 + num_logical_experts=E,
1089 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1090 + activation=MoEActivation.SILU,
1091 + in_dtype=torch.bfloat16,
1092 + device=device,
1093 + routing_method=RoutingMethodType.TopK,
1094 + max_num_tokens=512,
1095 + )
1096 + experts_ep = Int4EmulationTritonExperts(
1097 + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
1098 + )
1099 + out_rank = _run_emulation_forward_ep(
1100 + experts_ep,
1101 + w13_all[start:end],
1102 + w2_all[start:end],
1103 + hidden_states,
1104 + topk_weights,
1105 + topk_ids,
1106 + E,
1107 + expert_map,
1108 + )
1109 + out_ep_sum = out_ep_sum + out_rank
1110 +
1111 + assert torch.allclose(out_ep_sum, out_full, atol=1e-3), (
1112 + f"EP sum max diff: {(out_ep_sum - out_full).abs().max().item():.6f}"
1113 + )
1114 +
1115 +
1116 +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens", E2E_CONFIGS)
1117 +@pytest.mark.parametrize("fmt", ["gptq", "awq"])
1118 +def test_emulation_output_close_to_bf16_reference(
1119 + E, K, N, top_k, group_size, num_tokens, fmt
1120 +):
1121 + """Emulation output is close to a direct BF16 MoE forward."""
1122 + torch.manual_seed(11)
1123 + moe_config = _make_moe_config(E, K, N)
1124 +
1125 + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.bfloat16, device=device) * 0.02
1126 + w2_fp = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.02
1127 +
1128 + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], []
1129 + for e in range(E):
1130 + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size)
1131 + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size)
1132 + if fmt == "gptq":
1133 + packed13_list.append(gptq_pack(q13, 4, K, 2 * N))
1134 + packed2_list.append(gptq_pack(q2, 4, N, K))
1135 + else:
1136 + packed13_list.append(awq_pack(q13, 4, K, 2 * N))
1137 + packed2_list.append(awq_pack(q2, 4, N, K))
1138 + scales13_list.append(s13)
1139 + scales2_list.append(s2)
1140 +
1141 + process_fn = (
1142 + _process_weights_emulation_gptq
1143 + if fmt == "gptq"
1144 + else _process_weights_emulation_awq
1145 + )
1146 + res = process_fn(
1147 + torch.stack(packed13_list),
1148 + torch.stack(packed2_list),
1149 + torch.stack(scales13_list),
1150 + torch.stack(scales2_list),
1151 + None,
1152 + None,
1153 + )
1154 + w13_bf16, w2_bf16 = res[0], res[1]
1155 +
1156 + dummy_scale = torch.ones(1, dtype=torch.float16, device=device)
1157 + experts = Int4EmulationTritonExperts(
1158 + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale)
1159 + )
1160 +
1161 + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device)
1162 + topk_weights = torch.softmax(
1163 + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1
1164 + )
1165 + topk_ids = torch.stack(
1166 + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)]
1167 + ).to(torch.int32)
1168 +
1169 + out_emulation = _run_emulation_forward(
1170 + experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N
1171 + )
1172 +
1173 + ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device)
1174 + for m in range(num_tokens):
1175 + acc = torch.zeros(K, dtype=torch.float32, device=device)
1176 + for k in range(top_k):
1177 + e = topk_ids[m, k].item()
1178 + w = topk_weights[m, k].item()
1179 + gate_up = hidden_states[m] @ w13_bf16[e].T
1180 + gate, up = gate_up.chunk(2)
1181 + act = F.silu(gate) * up
1182 + acc += w * (act @ w2_bf16[e].T).float()
1183 + ref[m] = acc.bfloat16()
1184 +
1185 + rel_l2 = (
1186 + torch.norm(out_emulation.float() - ref.float())
1187 + / torch.norm(ref.float()).clamp(min=1e-6)
1188 + ).item()
1189 + assert rel_l2 < 0.15, f"[{fmt}] relative L2 = {rel_l2:.4f} (threshold 0.15)"
vllm/model_executor/layers/fused_moe/config.py
+12
@@ -887,6 +887,9 @@ def int4_w4a16_moe_quant_config(
887 block_shape: list[int] | None = None,
888 a1_gscale: torch.Tensor | None = None,
889 a2_gscale: torch.Tensor | None = None,
890 + gemm1_clamp_limit: float | None = None,
891 + gemm1_alpha: float | None = None,
892 + gemm1_beta: float | None = None,
893 ) -> FusedMoEQuantConfig:
894 """
895 Construct a quant config for 16-bit float activations and int4 weights.
@@ -897,6 +900,9 @@ def int4_w4a16_moe_quant_config(
900 _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale),
901 _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias),
902 _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias),
903 + gemm1_clamp_limit=gemm1_clamp_limit,
904 + gemm1_alpha=gemm1_alpha,
905 + gemm1_beta=gemm1_beta,
906 )
907
908
@@ -950,6 +956,9 @@ def int8_w8a16_moe_quant_config(
956 block_shape: list[int] | None = None,
957 a1_gscale: torch.Tensor | None = None,
958 a2_gscale: torch.Tensor | None = None,
959 + gemm1_clamp_limit: float | None = None,
960 + gemm1_alpha: float | None = None,
961 + gemm1_beta: float | None = None,
962 ) -> FusedMoEQuantConfig:
963 """
964 Construct a quant config for 16-bit float activations and int8 weights.
@@ -960,6 +969,9 @@ def int8_w8a16_moe_quant_config(
969 _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale),
970 _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias),
971 _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias),
972 + gemm1_clamp_limit=gemm1_clamp_limit,
973 + gemm1_alpha=gemm1_alpha,
974 + gemm1_beta=gemm1_beta,
975 )
976
977
vllm/model_executor/layers/fused_moe/experts/int4_emulation_moe.py new
+129
@@ -0,0 +1,129 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Int4 weight-only quantization emulation for MoE.
4 +
5 +Weights are dequantized from packed int4 to BF16 once at load time;
6 +the forward pass then runs plain TritonExperts in BF16.
7 +"""
8 +
9 +import torch
10 +
11 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
14 +from vllm.model_executor.layers.fused_moe.config import (
15 + FusedMoEConfig,
16 + FusedMoEQuantConfig,
17 +)
18 +from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
19 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
20 + QuantKey,
21 + kInt4Static,
22 + kInt4Static32,
23 + kInt4Static32Asym,
24 + kInt4StaticAsym,
25 +)
26 +from vllm.platforms import current_platform
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +class Int4EmulationTritonExperts(TritonExperts):
32 + """Int4 W-only MoE that dequantizes weights to BF16 at load time.
33 +
34 + Weights arrive already dequantized (convert_to_wna16_moe_kernel_format
35 + does the unpacking); apply() simply forwards to TritonExperts.
36 + """
37 +
38 + def __init__(
39 + self,
40 + moe_config: FusedMoEConfig,
41 + quant_config: FusedMoEQuantConfig,
42 + ):
43 + super().__init__(moe_config, quant_config)
44 + logger.warning_once(
45 + "Using Int4EmulationTritonExperts MoE backend. Int4 weights are "
46 + "dequantized to BF16 at load time "
47 + )
48 + # Weights are dequantized to BF16 before apply() is called, so
49 + # TritonExperts must see them as plain float — clear the int4 dtype
50 + # and scales so the hidden-size assertion and kernel dispatch treat
51 + # them as unquantized.
52 + self.quant_config._w1.dtype = None
53 + self.quant_config._w2.dtype = None
54 + self.quant_config._w1.scale = None
55 + self.quant_config._w2.scale = None
56 +
57 + @staticmethod
58 + def _supports_current_device() -> bool:
59 + return current_platform.is_cuda_alike()
60 +
61 + @staticmethod
62 + def _supports_quant_scheme(
63 + weight_key: QuantKey | None,
64 + activation_key: QuantKey | None,
65 + ) -> bool:
66 + return (
67 + weight_key
68 + in (
69 + kInt4Static,
70 + kInt4Static32,
71 + kInt4StaticAsym,
72 + kInt4Static32Asym,
73 + )
74 + and activation_key is None
75 + )
76 +
77 + @property
78 + def quant_dtype(self) -> torch.dtype | str | None:
79 + return None
80 +
81 + @property
82 + def block_shape(self) -> list[int] | None:
83 + return None
84 +
85 + @property
86 + def expects_unquantized_inputs(self) -> bool:
87 + return True
88 +
89 + def apply(
90 + self,
91 + output: torch.Tensor,
92 + hidden_states: torch.Tensor,
93 + w1: torch.Tensor,
94 + w2: torch.Tensor,
95 + topk_weights: torch.Tensor,
96 + topk_ids: torch.Tensor,
97 + activation: MoEActivation,
98 + global_num_experts: int,
99 + expert_map: torch.Tensor | None,
100 + a1q_scale: torch.Tensor | None,
101 + a2_scale: torch.Tensor | None,
102 + workspace13: torch.Tensor,
103 + workspace2: torch.Tensor,
104 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
105 + apply_router_weight_on_input: bool,
106 + ):
107 + if w1.element_size() < 2:
108 + raise RuntimeError(
109 + "Int4EmulationTritonExperts.apply() received packed int4 weights "
110 + "(element_size < 2). Weights must be dequantized to BF16 before "
111 + "the forward pass via convert_to_wna16_moe_kernel_format."
112 + )
113 + return super().apply(
114 + output=output,
115 + hidden_states=hidden_states,
116 + w1=w1,
117 + w2=w2,
118 + topk_weights=topk_weights,
119 + topk_ids=topk_ids,
120 + activation=activation,
121 + global_num_experts=global_num_experts,
122 + expert_map=expert_map,
123 + a1q_scale=None,
124 + a2_scale=None,
125 + workspace13=workspace13,
126 + workspace2=workspace2,
127 + expert_tokens_meta=expert_tokens_meta,
128 + apply_router_weight_on_input=apply_router_weight_on_input,
129 + )
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+281 -2
@@ -51,6 +51,7 @@ class WNA16MoEBackend(Enum):
51 CPU = "CPU"
52 FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
53 XPU = "XPU"
54 + EMULATION = "EMULATION"
55
56
57 def backend_to_kernel_cls(
@@ -87,6 +88,12 @@ def backend_to_kernel_cls(
88 )
89
90 return [CPUExpertsInt4]
91 + elif backend == WNA16MoEBackend.EMULATION:
92 + from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
93 + Int4EmulationTritonExperts,
94 + )
95 +
96 + return [Int4EmulationTritonExperts]
97 else:
98 raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
99
@@ -105,6 +112,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
112 WNA16MoEBackend.MARLIN,
113 WNA16MoEBackend.BATCHED_MARLIN,
114 WNA16MoEBackend.HUMMING,
115 + WNA16MoEBackend.EMULATION,
116 ]
117 return _AVAILABLE_BACKENDS
118
@@ -115,6 +123,7 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend:
123 "marlin": WNA16MoEBackend.MARLIN,
124 "humming": WNA16MoEBackend.HUMMING,
125 "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM,
126 + "emulation": WNA16MoEBackend.EMULATION,
127 }
128 if backend := mapping.get(runner_backend):
129 return backend
@@ -215,6 +224,9 @@ def make_wna16_moe_quant_config(
224 w2_bias: torch.Tensor | None = None,
225 a1_gscale: torch.Tensor | None = None,
226 a2_gscale: torch.Tensor | None = None,
227 + gemm1_clamp_limit: float | None = None,
228 + gemm1_alpha: float | None = None,
229 + gemm1_beta: float | None = None,
230 ) -> FusedMoEQuantConfig:
231 """Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE."""
232 if num_bits == 4:
@@ -228,6 +240,9 @@ def make_wna16_moe_quant_config(
240 block_shape=[0, group_size],
241 a1_gscale=a1_gscale,
242 a2_gscale=a2_gscale,
243 + gemm1_clamp_limit=gemm1_clamp_limit,
244 + gemm1_alpha=gemm1_alpha,
245 + gemm1_beta=gemm1_beta,
246 )
247 else:
248 assert num_bits == 8
@@ -241,6 +256,9 @@ def make_wna16_moe_quant_config(
256 block_shape=[0, group_size],
257 a1_gscale=a1_gscale,
258 a2_gscale=a2_gscale,
259 + gemm1_clamp_limit=gemm1_clamp_limit,
260 + gemm1_alpha=gemm1_alpha,
261 + gemm1_beta=gemm1_beta,
262 )
263
264
@@ -263,19 +281,23 @@ def make_wna16_moe_kernel(
281 from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
282 CPUExpertsInt4,
283 )
284 + from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
285 + Int4EmulationTritonExperts,
286 + )
287 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
288 XPUExpertsWNA16,
289 )
290
291 # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts,
271 - # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, and the Humming
272 - # grouped/indexed experts.
292 + # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, the Humming
293 + # grouped/indexed experts, and Int4EmulationTritonExperts
294 allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = (
295 MarlinExperts,
296 BatchedMarlinExperts,
297 TrtLlmMxint4ExpertsMonolithic,
298 XPUExpertsWNA16,
299 CPUExpertsInt4,
300 + Int4EmulationTritonExperts,
301 )
302 if backend == WNA16MoEBackend.HUMMING:
303 allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING))
@@ -1017,6 +1039,243 @@ def _humming_wna16_weight_schema(
1039 )
1040
1041
1042 +def _unpack_and_dequant_int4_gptq(
1043 + w_int32: torch.Tensor,
1044 + scale: torch.Tensor,
1045 + qzeros: torch.Tensor | None,
1046 + transpose_output: bool,
1047 + output_dtype: torch.dtype = torch.bfloat16,
1048 +) -> torch.Tensor:
1049 + """Unpack GPTQ-packed int4 weights and dequantize to output_dtype.
1050 +
1051 + Args:
1052 + w_int32: packed weights, shape [E, K_packed, N] where K_packed = K//8
1053 + (8 nibbles per int32, LSB-first in the K dimension).
1054 + scale: per-group scales, shape [E, K//group_size, N], float16.
1055 + qzeros: optional asymmetric zero-points, shape [E, K//gs, N//8], int32.
1056 + None for symmetric (uint4b8 with implicit bias 8).
1057 + transpose_output: if True return [E, N, K]; if False return [E, K, N].
1058 + output_dtype: target floating-point dtype (bfloat16 or float16).
1059 +
1060 + Returns:
1061 + Dequantized weight tensor in the requested layout.
1062 + """
1063 + E, K_packed, N = w_int32.shape
1064 + K = K_packed * 8
1065 +
1066 + # Unpack: [E, K_packed, N] -> [E, K_packed, N, 8] via bit-shifts.
1067 + # The nibble index (last dim) enumerates K rows within each packed column,
1068 + # so we must fuse K_packed and the nibble dim, not N and the nibble dim.
1069 + # Permute to [E, K_packed, 8, N] before reshaping to [E, K, N].
1070 + shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4
1071 + nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF # [E, K_packed, N, 8]
1072 +
1073 + # Reshape to [E, K, N]: fuse K_packed and nibble index (dim 1 and 3)
1074 + w = nibbles.permute(0, 1, 3, 2).reshape(E, K, N).to(torch.int16)
1075 +
1076 + if qzeros is None:
1077 + # Symmetric uint4b8: subtract bias so the range is [-8, 7]
1078 + w = w - 8
1079 + else:
1080 + # Asymmetric: unpack zero-points (same 8-nibble packing) and subtract
1081 + # qzeros shape: [E, K//gs, N//8] int32
1082 + gs = K // scale.shape[1]
1083 + n_gs = scale.shape[1]
1084 + zp_shifts = torch.arange(8, device=qzeros.device, dtype=torch.int32) * 4
1085 + zp_nibbles = (qzeros.unsqueeze(-1) >> zp_shifts) & 0xF # [E, n_gs, N//8, 8]
1086 + zp = zp_nibbles.reshape(E, n_gs, N).to(torch.int16) # [E, n_gs, N]
1087 + zp = zp.repeat_interleave(gs, dim=1) # [E, K, N]
1088 + w = w - zp
1089 +
1090 + # Broadcast scale [E, K//gs, N] -> [E, K, N]
1091 + gs = K // scale.shape[1]
1092 + scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype)
1093 +
1094 + w_dequant = w.to(output_dtype) * scale_broadcast # [E, K, N]
1095 +
1096 + if transpose_output:
1097 + return w_dequant.permute(0, 2, 1).contiguous() # [E, N, K]
1098 + return w_dequant.contiguous() # [E, K, N]
1099 +
1100 +
1101 +def _unpack_and_dequant_int4_awq(
1102 + w_int32: torch.Tensor,
1103 + scale: torch.Tensor,
1104 + qzeros: torch.Tensor | None,
1105 + transpose_output: bool,
1106 + output_dtype: torch.dtype = torch.bfloat16,
1107 +) -> torch.Tensor:
1108 + """Unpack AWQ-packed int4 weights and dequantize to output_dtype.
1109 +
1110 + AWQ packs along the N (column) dimension with an interleave permutation
1111 + [0,2,4,6,1,3,5,7] applied before packing, so unpacking must undo that.
1112 +
1113 + Args:
1114 + w_int32: packed weights, shape [E, K, N_packed] where N_packed = N//8
1115 + (8 nibbles per int32, packed along N with AWQ interleaving).
1116 + scale: per-group scales, shape [E, K//group_size, N], float16.
1117 + qzeros: asymmetric zero-points, shape [E, K//gs, N_packed], int32.
1118 + None for symmetric (uint4b8 with implicit bias 8).
1119 + transpose_output: if True return [E, N, K]; if False return [E, K, N].
1120 + output_dtype: target floating-point dtype (bfloat16 or float16).
1121 +
1122 + Returns:
1123 + Dequantized weight tensor in the requested layout.
1124 + """
1125 + E, K, N_packed = w_int32.shape
1126 + N = N_packed * 8
1127 +
1128 + # Unpack 8 nibbles per int32 along the N dimension (LSB-first)
1129 + shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4
1130 + # [E, K, N_packed, 8] -> [E, K, N_packed*8] = [E, K, N_interleaved]
1131 + nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF
1132 + w_interleaved = nibbles.reshape(E, K, N) # [E, K, N] but column-interleaved
1133 +
1134 + # Undo AWQ interleave: packed order is [0,2,4,6,1,3,5,7] within each group
1135 + # of 8. Inverse: position i in packed -> original column interleave[i].
1136 + # To reverse: we need the inverse permutation so that
1137 + # w[:, :, inv_interleave] = w_interleaved gives the natural column order.
1138 + interleave = torch.tensor([0, 2, 4, 6, 1, 3, 5, 7], device=w_int32.device)
1139 + inv_interleave = torch.empty_like(interleave)
1140 + inv_interleave[interleave] = torch.arange(8, device=w_int32.device)
1141 +
1142 + # Apply inverse interleave within each group of 8 columns
1143 + w_reshaped = w_interleaved.reshape(E, K, N // 8, 8) # [E, K, groups, 8]
1144 + w_reordered = w_reshaped[:, :, :, inv_interleave] # undo interleave
1145 + w = w_reordered.reshape(E, K, N).to(torch.int16) # [E, K, N]
1146 +
1147 + if qzeros is None:
1148 + w = w - 8
1149 + else:
1150 + # qzeros: [E, K//gs, N_packed] int32, same AWQ column packing
1151 + gs = K // scale.shape[1]
1152 + n_gs = scale.shape[1]
1153 + zp_nibbles = (qzeros.unsqueeze(-1) >> shifts) & 0xF # [E, n_gs, N_packed, 8]
1154 + zp_interleaved = zp_nibbles.reshape(E, n_gs, N)
1155 + zp_reshaped = zp_interleaved.reshape(E, n_gs, N // 8, 8)
1156 + zp_reordered = zp_reshaped[:, :, :, inv_interleave]
1157 + zp = zp_reordered.reshape(E, n_gs, N).to(torch.int16) # [E, n_gs, N]
1158 + zp = zp.repeat_interleave(gs, dim=1) # [E, K, N]
1159 + w = w - zp
1160 +
1161 + gs = K // scale.shape[1]
1162 + scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype) # [E, K, N]
1163 +
1164 + w_dequant = w.to(output_dtype) * scale_broadcast # [E, K, N]
1165 +
1166 + if transpose_output:
1167 + return w_dequant.permute(0, 2, 1).contiguous() # [E, N, K]
1168 + return w_dequant.contiguous() # [E, K, N]
1169 +
1170 +
1171 +def _process_weights_emulation_gptq(
1172 + w13: torch.Tensor,
1173 + w2: torch.Tensor,
1174 + w13_scale: torch.Tensor,
1175 + w2_scale: torch.Tensor,
1176 + w13_qzeros: torch.Tensor | None,
1177 + w2_qzeros: torch.Tensor | None,
1178 +) -> tuple:
1179 + """Dequantize int4 weights to BF16 for the emulation backend.
1180 +
1181 + Inputs are in GPTQ packed format:
1182 + w13: [E, K//8, 2*N] int32 (gate+up proj stacked on dim 2)
1183 + w2: [E, N//8, K] int32
1184 + w13_scale: [E, K//gs, 2*N] float16
1185 + w2_scale: [E, N//gs, K] float16
1186 +
1187 + Outputs (what TritonExperts expects):
1188 + w13_out: [E, 2*N, K] bfloat16
1189 + w2_out: [E, K, N] bfloat16
1190 + """
1191 + # w13: packed along K (dim 1), output cols are 2*N (dim 2)
1192 + # transpose_output=True yields [E, 2*N, K]
1193 + w13_bf16 = _unpack_and_dequant_int4_gptq(
1194 + w13, w13_scale, w13_qzeros, transpose_output=True
1195 + )
1196 +
1197 + # w2: packed along N (dim 1 is N//8), output cols are K (dim 2)
1198 + # After unpacking we get [E, N, K]; we want [E, K, N] for TritonExperts
1199 + # transpose_output=False gives [E, N, K], then we permute once more
1200 + w2_unpacked = _unpack_and_dequant_int4_gptq(
1201 + w2, w2_scale, w2_qzeros, transpose_output=False
1202 + ) # [E, N, K]
1203 + w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous() # [E, K, N]
1204 +
1205 + dummy = torch.ones(1, dtype=torch.float16, device=w13.device)
1206 + return (
1207 + w13_bf16, # w13_qweight (now bf16, not int32)
1208 + w2_bf16, # w2_qweight (now bf16, not int32)
1209 + dummy, # w13_scales (unused; nulled out in Int4EmulationTritonExperts)
1210 + dummy, # w2_scales (unused)
1211 + None, # w13_g_idx
1212 + None, # w2_g_idx
1213 + None, # w13_g_idx_sort_indices
1214 + None, # w2_g_idx_sort_indices
1215 + None, # w13_qzeros
1216 + None, # w2_qzeros
1217 + None, # w13_input_global_scale
1218 + None, # w2_input_global_scale
1219 + None, # w13_bias
1220 + None, # w2_bias
1221 + )
1222 +
1223 +
1224 +def _process_weights_emulation_awq(
1225 + w13: torch.Tensor,
1226 + w2: torch.Tensor,
1227 + w13_scale: torch.Tensor,
1228 + w2_scale: torch.Tensor,
1229 + w13_qzeros: torch.Tensor | None,
1230 + w2_qzeros: torch.Tensor | None,
1231 +) -> tuple:
1232 + """Dequantize AWQ int4 weights to BF16 for the emulation backend.
1233 +
1234 + AWQ inputs:
1235 + w13: [E, K, 2*N//8] int32 (packed along N, gate+up on dim 2)
1236 + w2: [E, N, K//8] int32 (packed along K)
1237 + w13_scale: [E, K//gs, 2*N] float16
1238 + w2_scale: [E, N//gs, K] float16
1239 +
1240 + Outputs (what TritonExperts expects):
1241 + w13_out: [E, 2*N, K] bfloat16
1242 + w2_out: [E, K, N] bfloat16
1243 + """
1244 + # w13: AWQ-packed along N (dim 2), K is unpacked in dim 1
1245 + # _unpack_and_dequant_int4_awq with transpose_output=True yields [E, 2*N, K]
1246 + w13_bf16 = _unpack_and_dequant_int4_awq(
1247 + w13, w13_scale, w13_qzeros, transpose_output=True
1248 + )
1249 +
1250 + # w2: AWQ packs along K (dim 2 is K//8), N is unpacked in dim 1.
1251 + # AWQ w2 is [E, N, K//8] — same column-pack format applied to the K dim.
1252 + # _unpack_and_dequant_int4_awq expects [E, rows, N_packed] where the
1253 + # packed dim is columns. Treat dim 1 as rows and dim 2 as N_packed:
1254 + # unpacking gives [E, N, K]. Then permute to [E, K, N].
1255 + w2_unpacked = _unpack_and_dequant_int4_awq(
1256 + w2, w2_scale, w2_qzeros, transpose_output=False
1257 + ) # [E, N, K]
1258 + w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous() # [E, K, N]
1259 +
1260 + dummy = torch.ones(1, dtype=torch.float16, device=w13.device)
1261 + return (
1262 + w13_bf16,
1263 + w2_bf16,
1264 + dummy,
1265 + dummy,
1266 + None,
1267 + None,
1268 + None,
1269 + None,
1270 + None,
1271 + None,
1272 + None,
1273 + None,
1274 + None,
1275 + None,
1276 + )
1277 +
1278 +
1279 def convert_to_wna16_moe_kernel_format(
1280 backend: WNA16MoEBackend,
1281 layer: torch.nn.Module,
@@ -1203,5 +1462,25 @@ def convert_to_wna16_moe_kernel_format(
1462 w13_bias_out,
1463 w2_bias_out,
1464 )
1465 + elif backend == WNA16MoEBackend.EMULATION:
1466 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
1467 +
1468 + if isinstance(quant_config, AutoAWQConfig):
1469 + return _process_weights_emulation_awq(
1470 + w13,
1471 + w2,
1472 + w13_scale,
1473 + w2_scale,
1474 + w13_qzeros,
1475 + w2_qzeros,
1476 + )
1477 + return _process_weights_emulation_gptq(
1478 + w13,
1479 + w2,
1480 + w13_scale,
1481 + w2_scale,
1482 + w13_qzeros,
1483 + w2_qzeros,
1484 + )
1485 else:
1486 raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}")
vllm/model_executor/layers/quantization/auto_awq.py
+3
@@ -774,6 +774,9 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
774 w2_bias=getattr(layer, "w2_bias", None),
775 a1_gscale=getattr(layer, "w13_input_global_scale", None),
776 a2_gscale=getattr(layer, "w2_input_global_scale", None),
777 + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
778 + gemm1_alpha=getattr(layer, "swiglu_alpha", None),
779 + gemm1_beta=getattr(layer, "swiglu_beta", None),
780 )
781
782 def select_gemm_impl(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+17
@@ -142,6 +142,23 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
142 return CompressedTensorsW4A16FlydslMoEMethod(
143 weight_quant, input_quant, layer.moe_config
144 )
145 + elif moe_backend == "emulation":
146 + # Although this is called 'Marlin', actually it selects
147 + # emulation backend by calling select_wna16_moe_backend.
148 + # TODO: we need to update CompressedTensorsWNA16MoeMethod
149 + # to honor "--moe-backend" option
150 + from .compressed_tensors_moe_wna16_marlin import (
151 + CompressedTensorsWNA16MarlinMoEMethod,
152 + )
153 +
154 + logger.info_once(
155 + "Using CompressedTensorsWNA16MarlinMoEMethod "
156 + "(emulation backend requested)"
157 + )
158 + return CompressedTensorsWNA16MarlinMoEMethod(
159 + weight_quant, input_quant, layer.moe_config, layer_name
160 + )
161 +
162 from .compressed_tensors_moe_wna16 import (
163 CompressedTensorsWNA16MoEMethod,
164 )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py
+2 -1
@@ -108,7 +108,8 @@ class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod):
108 # grouped actorder isn't supported by this kernel
109 assert weight_quant.actorder != "group"
110 assert weight_quant.symmetric, (
111 - "Only symmetric quantization is supported for MoE"
111 + "Only symmetric quantization is supported for MoE. "
112 + "Try --moe-backend emulation."
113 )
114
115 def create_weights(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+2 -1
@@ -48,7 +48,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
48 # grouped actorder isn't supported by this kernel
49 assert weight_quant.actorder != "group"
50 assert weight_quant.symmetric, (
51 - "Only symmetric quantization is supported for MoE"
51 + "Only symmetric quantization is supported for MoE. "
52 + "Try --moe-backend emulation."
53 )
54
55 def create_weights(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+18 -6
@@ -467,10 +467,16 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
467
468 # Marlin-specific parameters (not needed for Flashinfer)
469 if not is_flashinfer:
470 - replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
471 - replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
472 - replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
473 - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
470 + if w13_g_idx_processed is not None:
471 + replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
472 + if w2_g_idx_processed is not None:
473 + replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
474 + if w13_g_idx_sort_indices is not None:
475 + replace_parameter(
476 + layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices
477 + )
478 + if w2_g_idx_sort_indices is not None:
479 + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
480
481 # Register input global scales if present
482 if w13_input_global_scale is not None:
@@ -484,8 +490,11 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
490 torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
491 )
492
487 - if self.experts_cls is not None and issubclass(
488 - self.experts_cls, FusedMoEExpertsModular
493 + # Marlin workspace — only needed for Marlin-family backends, not emulation.
494 + if (
495 + self.experts_cls is not None
496 + and issubclass(self.experts_cls, FusedMoEExpertsModular)
497 + and self.wna16_backend != WNA16MoEBackend.EMULATION
498 ):
499 layer.workspace = marlin_make_workspace_new(
500 layer.w13_weight_g_idx.device, 4
@@ -528,6 +537,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
537 num_bits=self.num_bits,
538 w1_zp=getattr(layer, "w13_weight_zero_point", None),
539 w2_zp=getattr(layer, "w2_weight_zero_point", None),
540 + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
541 + gemm1_alpha=getattr(layer, "swiglu_alpha", None),
542 + gemm1_beta=getattr(layer, "swiglu_beta", None),
543 )
544
545 def apply_monolithic(