A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 13:30 UTC
[XPU][MoE] Add WNA16 oracle backend for GPTQ sym-int4 (xpu_fused_moe) (#41426)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Jason Elie Bou Kheir committed
May 28, 2026 at 09:30 UTC
3207e7680e52853db757aeb14489612704687cd6
2 files changed
+179
-7
vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
+48
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
QuantKey,
17
kFp8DynamicTensorSym,
18
kFp8StaticTensorSym,
19
+ kInt4Static,
20
kMxfp4Static,
21
kMxfp8Dynamic,
22
kMxfp8Static,
@@ -48,6 +49,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
49
num_dispatchers,
50
)
51
self.is_fp8 = False
52
+ self.is_int4 = False
53
self.is_mxfp4 = False
54
self.is_mxfp8 = False
55
self.fused_moe_impl: XpuFusedMoe | None = None
@@ -133,6 +135,15 @@ class XPUExperts(mk.FusedMoEExpertsModular):
135
expert_tokens_meta: mk.ExpertTokensMetadata | None,
136
apply_router_weight_on_input: bool,
137
):
138
+ # The kernel takes is_fp8/is_int4/is_mxfp4 as independent booleans.
139
+ # In this hierarchy each subclass flips exactly one to True; assert
140
+ # the invariant so a future subclass that sets two doesn't silently
141
+ # miscompute (kernel-side priority is undocumented).
142
+ assert sum([self.is_fp8, self.is_int4, self.is_mxfp4]) <= 1, (
143
+ "XPUExperts: at most one of is_fp8, is_int4, is_mxfp4 may be True; "
144
+ f"got is_fp8={self.is_fp8}, is_int4={self.is_int4}, "
145
+ f"is_mxfp4={self.is_mxfp4}."
146
+ )
147
if self.fused_moe_impl is None:
148
topk = topk_ids.size(-1)
149
self.fused_moe_impl = XpuFusedMoe(
@@ -148,6 +159,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
159
ep_rank=self.moe_config.ep_rank,
160
ep_size=self.moe_config.ep_size,
161
is_fp8=self.is_fp8,
162
+ is_int4=self.is_int4,
163
is_mxfp4=self.is_mxfp4,
164
is_mxfp8=self.is_mxfp8,
165
)
@@ -217,6 +229,42 @@ class XPUExpertsMxfp8(XPUExpertsFp8):
229
return (weight_key, activation_key) in SUPPORTED_W_A
230
231
232
+class XPUExpertsWNA16(XPUExperts):
233
+ """W4A16 INT4-symmetric MoE backed by `xpu_fused_moe(is_int4=True)`.
234
+
235
+ Weight layout when `is_int4=True` (per `xpu_fused_moe` docstring):
236
+ w13: [num_experts, 2*inter_size, hidden_size] contiguous int4-packed
237
+ w13_scales: [num_experts, 2*inter_size, hidden_size // group_size]
238
+ w2: [num_experts, hidden_size, inter_size] contiguous int4-packed
239
+ w2_scales: [num_experts, hidden_size, inter_size // group_size]
240
+
241
+ Pairs with `INCXPULinearMethod` for the linear layers; together they
242
+ cover full-attn + MoE on Intel XPU end-to-end without IPEX.
243
+ """
244
+
245
+ def __init__(
246
+ self,
247
+ moe_config: FusedMoEConfig,
248
+ quant_config: FusedMoEQuantConfig,
249
+ max_num_tokens: int | None = None,
250
+ num_dispatchers: int | None = None,
251
+ ):
252
+ super().__init__(
253
+ moe_config,
254
+ quant_config,
255
+ max_num_tokens,
256
+ num_dispatchers,
257
+ )
258
+ self.is_int4 = True
259
+
260
+ @staticmethod
261
+ def _supports_quant_scheme(
262
+ weight_key: QuantKey | None,
263
+ activation_key: QuantKey | None,
264
+ ) -> bool:
265
+ return (weight_key, activation_key) == (kInt4Static, None)
266
+
267
+
268
class XPUExpertsMXFp4(XPUExperts):
269
def __init__(
270
self,
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+131
-7
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import sys
4
from enum import Enum
5
from typing import TYPE_CHECKING
6
@@ -26,6 +27,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
QuantKey,
29
)
30
+from vllm.platforms import current_platform
31
32
if TYPE_CHECKING:
33
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
@@ -37,6 +39,7 @@ logger = init_logger(__name__)
39
class WNA16MoEBackend(Enum):
40
MARLIN = "MARLIN"
41
BATCHED_MARLIN = "BATCHED_MARLIN"
42
+ XPU = "XPU"
43
44
45
def backend_to_kernel_cls(
@@ -57,6 +60,13 @@ def backend_to_kernel_cls(
60
61
return [BatchedMarlinExperts]
62
63
+ elif backend == WNA16MoEBackend.XPU:
64
+ from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
65
+ XPUExpertsWNA16,
66
+ )
67
+
68
+ return [XPUExpertsWNA16]
69
+
70
else:
71
raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
72
@@ -65,11 +75,12 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
75
"""
76
Get available backends in priority order based on platform and config.
77
"""
68
- _AVAILABLE_BACKENDS = [
78
+ if current_platform.is_xpu():
79
+ return [WNA16MoEBackend.XPU]
80
+ return [
81
WNA16MoEBackend.MARLIN,
82
WNA16MoEBackend.BATCHED_MARLIN,
83
]
72
- return _AVAILABLE_BACKENDS
84
85
86
def select_wna16_moe_backend(
@@ -156,12 +167,14 @@ def make_wna16_moe_kernel(
167
w2_g_idx_sort_indices: torch.Tensor | None,
168
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
169
) -> mk.FusedMoEKernel:
159
- # Currently, we only support MarlinExperts and BatchedMarlinExperts
160
- assert experts_cls in (MarlinExperts, BatchedMarlinExperts)
161
-
170
from vllm.model_executor.layers.fused_moe.all2all_utils import (
171
maybe_make_prepare_finalize,
172
)
173
+ from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
174
+ XPUExpertsWNA16,
175
+ )
176
+
177
+ assert experts_cls in (MarlinExperts, BatchedMarlinExperts, XPUExpertsWNA16)
178
179
prepare_finalize = maybe_make_prepare_finalize(
180
moe=moe_config,
@@ -172,11 +185,24 @@ def make_wna16_moe_kernel(
185
assert prepare_finalize is not None
186
assert isinstance(prepare_finalize, mk.FusedMoEPrepareAndFinalizeModular)
187
175
- if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
188
+ if experts_cls is XPUExpertsWNA16:
189
+ assert (
190
+ prepare_finalize.activation_format == mk.FusedMoEActivationFormat.Standard
191
+ ), (
192
+ "XPUExpertsWNA16 only supports the Standard activation format; "
193
+ "xpu_fused_moe(is_int4=True) does not implement BatchedExperts."
194
+ )
195
+ experts: mk.FusedMoEExperts = XPUExpertsWNA16(
196
+ moe_config=moe_config,
197
+ quant_config=moe_quant_config,
198
+ )
199
+ elif (
200
+ prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts
201
+ ):
202
assert experts_cls == BatchedMarlinExperts
203
max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
204
assert max_num_tokens is not None
179
- experts: mk.FusedMoEExperts = BatchedMarlinExperts(
205
+ experts = BatchedMarlinExperts(
206
max_num_tokens=max_num_tokens,
207
num_dispatchers=prepare_finalize.num_dispatchers(),
208
moe_config=moe_config,
@@ -524,6 +550,69 @@ def _process_awq_weights_marlin(
550
)
551
552
553
+def _process_weights_xpu(
554
+ layer: torch.nn.Module,
555
+ quant_config: QuantizationConfig,
556
+ w13_qweight: torch.Tensor,
557
+ w2_qweight: torch.Tensor,
558
+ w13_scales: torch.Tensor,
559
+ w2_scales: torch.Tensor,
560
+ w13_bias: torch.Tensor | None = None,
561
+ w2_bias: torch.Tensor | None = None,
562
+) -> tuple[
563
+ torch.Tensor, # w13_qweight
564
+ torch.Tensor, # w2_qweight
565
+ torch.Tensor, # w13_scales
566
+ torch.Tensor, # w2_scales
567
+ torch.Tensor | None, # w13_bias
568
+ torch.Tensor | None, # w2_bias
569
+]:
570
+ """Repack GPTQ-format INT4 MoE weights into the layout
571
+ `vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True)` expects:
572
+
573
+ w13: [E, 2*N, K] int4 (uint8 storage [E, 2*N, K // 2])
574
+ w13_scales: [E, 2*N, K // group_size] params_dtype
575
+ w2: [E, K, N] int4 (uint8 storage [E, K, N // 2])
576
+ w2_scales: [E, K, N // group_size] params_dtype
577
+
578
+ Input GPTQ layout from FusedMoE.weight_loader:
579
+ w13: [E, K // 8, 2*N] int32 (8 nibbles per int32 along the input dim)
580
+ w13_scales: [E, K // group_size, 2*N] params_dtype
581
+ w2: [E, N // 8, K] int32
582
+ w2_scales: [E, N // group_size, K] params_dtype
583
+
584
+ Transpose dim 1 ↔ dim 2 then view int32 → uint8 to recover sequential
585
+ int4-packed bytes along the input dim. Each packed int32 holds 8 nibbles
586
+ `(n7<<28)|(n6<<24)|...|(n1<<4)|n0` in ascending K order; on a
587
+ little-endian host the int32→uint8 view exposes them as bytes
588
+ `[n1<<4|n0, n3<<4|n2, n5<<4|n4, n7<<4|n6]`, i.e. two nibbles per byte
589
+ with the lower nibble = lower input-K index. xpu_fused_moe(is_int4=True)
590
+ expects this convention; on a big-endian host the byte order reverses
591
+ and the kernel would silently miscompute, so we hard-fail.
592
+ """
593
+ del layer, quant_config # unused — kept for parity with the marlin helper
594
+
595
+ if sys.byteorder != "little":
596
+ raise NotImplementedError(
597
+ "_process_weights_xpu requires a little-endian host: the GPTQ "
598
+ "int32 → uint8 nibble repack relies on LE byte ordering."
599
+ )
600
+
601
+ w13_xpu = w13_qweight.transpose(1, 2).contiguous().view(torch.uint8)
602
+ w2_xpu = w2_qweight.transpose(1, 2).contiguous().view(torch.uint8)
603
+ w13_scales_xpu = w13_scales.transpose(1, 2).contiguous()
604
+ w2_scales_xpu = w2_scales.transpose(1, 2).contiguous()
605
+
606
+ return (
607
+ w13_xpu,
608
+ w2_xpu,
609
+ w13_scales_xpu,
610
+ w2_scales_xpu,
611
+ w13_bias,
612
+ w2_bias,
613
+ )
614
+
615
+
616
def convert_to_wna16_moe_kernel_format(
617
backend: WNA16MoEBackend,
618
layer: torch.nn.Module,
@@ -617,5 +706,40 @@ def convert_to_wna16_moe_kernel_format(
706
w13_bias,
707
w2_bias,
708
)
709
+ elif backend == WNA16MoEBackend.XPU:
710
+ (
711
+ w13_xpu,
712
+ w2_xpu,
713
+ w13_scale_xpu,
714
+ w2_scale_xpu,
715
+ w13_bias_out,
716
+ w2_bias_out,
717
+ ) = _process_weights_xpu(
718
+ layer,
719
+ quant_config,
720
+ w13,
721
+ w2,
722
+ w13_scale,
723
+ w2_scale,
724
+ w13_bias,
725
+ w2_bias,
726
+ )
727
+ empty = torch.empty((0,), dtype=torch.int32, device=w13.device)
728
+ return (
729
+ w13_xpu,
730
+ w2_xpu,
731
+ w13_scale_xpu,
732
+ w2_scale_xpu,
733
+ empty, # w13_g_idx
734
+ empty, # w2_g_idx
735
+ empty, # w13_g_idx_sort_indices
736
+ empty, # w2_g_idx_sort_indices
737
+ None, # w13_qzeros — sym int4 on XPU has none; kernel does uint4b8→s4
738
+ None, # w2_qzeros
739
+ None, # w13_input_global_scale
740
+ None, # w2_input_global_scale
741
+ w13_bias_out,
742
+ w2_bias_out,
743
+ )
744
else:
745
raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}")