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 5, 03:30 UTC
[XPU][MXFP4] add mxfp4 quant op for XPU (#39857)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
zofia committed
Apr 15, 2026 at 20:28 UTC
fc701c80588c215f84af0b745edcf4d127e276bc
2 files changed
+50
vllm/_xpu_ops.py
+46
@@ -144,6 +144,46 @@ def _xpu_mxfp8_quantize_fake(
144
return x.to(dtype), x_s.to(torch.float8_e8m0fnu)
145
146
147
+def _xpu_mxfp4_quantize_impl(
148
+ x: torch.Tensor,
149
+) -> tuple[torch.Tensor, torch.Tensor]:
150
+ MXFP4_BLOCK_SIZE = 32
151
+ eps = 1e-10
152
+ assert x.ndim == 2, "input must be 2-D"
153
+ assert x.shape[-1] % MXFP4_BLOCK_SIZE == 0, (
154
+ f"last dimension {x.shape[-1]} must be divisible by group_size "
155
+ f"{MXFP4_BLOCK_SIZE}"
156
+ )
157
+ assert x.is_contiguous(), "input groups must be contiguous"
158
+
159
+ M, N = x.shape
160
+
161
+ # Packed FP4 output: two nibbles per byte
162
+ x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
163
+ x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
164
+
165
+ torch.ops._C.per_token_group_quant_mxfp4(x, x_q, x_s, MXFP4_BLOCK_SIZE, eps)
166
+
167
+ x_q = x_q.view(torch.float4_e2m1fn_x2)
168
+ x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
169
+ return x_q, x_s
170
+
171
+
172
+def _xpu_mxfp4_quantize_fake(
173
+ x: torch.Tensor,
174
+) -> tuple[torch.Tensor, torch.Tensor]:
175
+ MXFP4_BLOCK_SIZE = 32
176
+ M, N = x.shape
177
+
178
+ # Packed FP4 output: two nibbles per byte
179
+ x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
180
+ x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
181
+
182
+ x_q = x_q.view(torch.float4_e2m1fn_x2)
183
+ x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
184
+ return x_q, x_s
185
+
186
+
187
# Global flag to ensure ops are registered only once
188
_OPS_REGISTERED = False
189
@@ -555,6 +595,12 @@ class xpu_ops:
595
fake_impl=_xpu_mxfp8_quantize_fake,
596
)
597
598
+ direct_register_custom_op(
599
+ op_name="xpu_mxfp4_quantize",
600
+ op_func=_xpu_mxfp4_quantize_impl,
601
+ fake_impl=_xpu_mxfp4_quantize_fake,
602
+ )
603
+
604
_OPS_REGISTERED = True
605
606
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+4
@@ -162,3 +162,7 @@ try:
162
quant_dequant_mxfp4 = torch.ops.vllm.quant_dequant_mxfp4
163
except AttributeError as error:
164
raise error
165
+
166
+
167
+def xpu_mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
168
+ return torch.ops.vllm.xpu_mxfp4_quantize(x)