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, 11:00 UTC
[XPU] add xpu backend implementation of mxfp8 quant (#38682)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
zofia committed
Apr 8, 2026 at 08:30 UTC
ad3304425b4c3622a3c10f57699d740e6f5f7cf3
2 files changed
+54
vllm/_xpu_ops.py
+48
@@ -102,6 +102,48 @@ def _xpu_ops_deepseek_scaling_rope_fake(
102
return query, key
103
104
105
+def _xpu_mxfp8_quantize_impl(
106
+ x: torch.Tensor, dtype: torch.dtype | None = None
107
+) -> tuple[torch.Tensor, torch.Tensor]:
108
+ MXFP8_BLOCK_SIZE = 32
109
+ assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0
110
+ if dtype is not None:
111
+ assert dtype in (torch.float8_e4m3fn, torch.float8_e5m2), (
112
+ f"Unsupported dtype for xpu_mxfp8_quantize: {dtype}. "
113
+ f"Expected torch.float8_e4m3fn or torch.float8_e5m2."
114
+ )
115
+ else:
116
+ dtype = current_platform.fp8_dtype()
117
+
118
+ finfo = torch.finfo(dtype)
119
+ fp8_min = finfo.min
120
+ fp8_max = finfo.max
121
+ eps = 1e-10
122
+
123
+ x_q = torch.empty_like(x, device=x.device, dtype=dtype)
124
+ shape = x.shape[:-1] + (x.shape[-1] // MXFP8_BLOCK_SIZE,)
125
+ x_s = torch.empty(shape, device=x.device, dtype=torch.float32)
126
+ torch.ops._C.per_token_group_fp8_quant(
127
+ x, x_q, x_s, MXFP8_BLOCK_SIZE, eps, fp8_min, fp8_max, True
128
+ )
129
+ x_s = x_s.to(torch.float8_e8m0fnu)
130
+ return x_q, x_s
131
+
132
+
133
+def _xpu_mxfp8_quantize_fake(
134
+ x: torch.Tensor, dtype: torch.dtype | None = None
135
+) -> tuple[torch.Tensor, torch.Tensor]:
136
+ if dtype is None:
137
+ dtype = current_platform.fp8_dtype()
138
+
139
+ MXFP8_BLOCK_SIZE = 32
140
+
141
+ shape = x.shape[:-1] + (x.shape[-1] // MXFP8_BLOCK_SIZE,)
142
+ x_s = torch.zeros(shape, device=x.device, dtype=torch.float32)
143
+
144
+ return x.to(dtype), x_s.to(torch.float8_e8m0fnu)
145
+
146
+
147
# Global flag to ensure ops are registered only once
148
_OPS_REGISTERED = False
149
@@ -504,6 +546,12 @@ class xpu_ops:
546
dispatch_key=current_platform.dispatch_key,
547
)
548
549
+ direct_register_custom_op(
550
+ op_name="xpu_mxfp8_quantize",
551
+ op_func=_xpu_mxfp8_quantize_impl,
552
+ fake_impl=_xpu_mxfp8_quantize_fake,
553
+ )
554
+
555
_OPS_REGISTERED = True
556
557
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+6
@@ -205,6 +205,12 @@ direct_register_custom_op(
205
)
206
207
208
+def xpu_mxfp8_quantize(
209
+ x: torch.Tensor, dtype: torch.dtype | None = None
210
+) -> tuple[torch.Tensor, torch.Tensor]:
211
+ return torch.ops.vllm.xpu_mxfp8_quantize(x, dtype)
212
+
213
+
214
class Mxfp8LinearOp:
215
def __init__(self):
216
self.backend = select_mxfp8_linear_backend()