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 6, 00:45 UTC
[XPU] support MLA model on Intel GPU (#37143)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Mar 25, 2026 at 17:43 UTC
14771f715085f5026919d30048329c3e822d4b3c
4 files changed
+15
-12
vllm/_xpu_ops.py
+1
@@ -170,6 +170,7 @@ class xpu_ops:
170
num_splits=0,
171
return_softmax_lse: bool | None = False,
172
s_aux: torch.Tensor | None = None,
173
+ return_attn_probs: bool | None = False,
174
):
175
assert cu_seqlens_k is not None or seqused_k is not None, (
176
"cu_seqlens_k or seqused_k must be provided"
vllm/model_executor/layers/attention/mla_attention.py
+4
@@ -1059,6 +1059,10 @@ except ImportError:
1059
"MLA models using TRITON_MLA will require flash_attn. "
1060
"AITER_MLA backends use aiter kernels instead."
1061
)
1062
+ elif current_platform.is_xpu():
1063
+ from vllm._xpu_ops import xpu_ops as ops
1064
+
1065
+ flash_attn_varlen_func = ops.flash_attn_varlen_func # type: ignore[no-redef]
1066
1067
1068
def dynamic_per_batched_tensor_quant(
vllm/model_executor/layers/quantization/input_quant_fp8.py
+10
@@ -165,6 +165,16 @@ class QuantFP8(CustomOp):
165
# Fallback to CUDA implementation
166
return self.forward_cuda(x, scale, scale_ub)
167
168
+ def forward_xpu(
169
+ self,
170
+ x: torch.Tensor,
171
+ scale: torch.Tensor | None = None,
172
+ scale_ub: torch.Tensor | None = None,
173
+ use_triton: bool = False,
174
+ ) -> tuple[torch.Tensor, torch.Tensor]:
175
+ # XPU can use same code path as CUDA.
176
+ return self.forward_cuda(x, scale, scale_ub, use_triton)
177
+
178
def forward_native(
179
self,
180
x: torch.Tensor,
vllm/platforms/xpu.py
-12
@@ -160,7 +160,6 @@ class XPUPlatform(Platform):
160
@classmethod
161
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
162
cache_config = vllm_config.cache_config
163
- model_config = vllm_config.model_config
163
parallel_config = vllm_config.parallel_config
164
# in V1(or with chunked prefill) block_size is 64
165
if cache_config and not cache_config.user_specified_block_size:
@@ -209,17 +208,6 @@ class XPUPlatform(Platform):
208
if vllm_config.kv_transfer_config is not None:
209
vllm_config.kv_transfer_config.enable_permute_local_kv = True
210
212
- if model_config and model_config.use_mla:
213
- logger.info(
214
- "MLA is enabled on a non-GPU platform; forcing chunked "
215
- "prefill and prefix caching to be disabled."
216
- )
217
- vllm_config.scheduler_config.enable_chunked_prefill = False
218
- vllm_config.scheduler_config.max_num_batched_tokens = max(
219
- vllm_config.model_config.max_model_len,
220
- vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,
221
- )
222
-
211
# In some cases, the internal memory type cache can misdetect GPU
212
# memory as host memory, also leading to invalid memory access.
213
# This cache can be disabled by setting UCX_MEMTYPE_CACHE=n.