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 12, 14:45 UTC
[XPU]fix Kimi-VL-A3B-thinking on xpu (#29309)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Nov 25, 2025 at 05:02 UTC
3cfa63ad991665b2440155cd29352342024072fd
1 file changed
+14
-6
vllm/model_executor/models/moonvit.py
+14
-6
@@ -56,10 +56,13 @@ from transformers.utils import is_flash_attn_2_available
56
from vllm.model_executor.layers.conv import Conv2dLayer
57
from vllm.model_executor.layers.linear import ReplicatedLinear
58
from vllm.model_executor.models.utils import maybe_prefix
59
+from vllm.platforms import current_platform
60
from vllm.transformers_utils.configs.moonvit import MoonViTConfig
61
62
if is_flash_attn_2_available():
63
from flash_attn import flash_attn_varlen_func
64
+elif current_platform.is_xpu():
65
+ from vllm.attention.utils.fa_utils import flash_attn_varlen_func
66
else:
67
flash_attn_varlen_func = None
68
@@ -106,10 +109,10 @@ def multihead_attention(
109
q,
110
k,
111
v,
109
- q_cu_seqlens,
110
- k_cu_seqlens,
111
- max_seqlen_q,
112
- max_seqlen_k,
112
+ cu_seqlens_q=q_cu_seqlens,
113
+ cu_seqlens_k=k_cu_seqlens,
114
+ max_seqlen_q=max_seqlen_q,
115
+ max_seqlen_k=max_seqlen_k,
116
causal=False,
117
)
118
attn_out = attn_out.flatten(start_dim=-2)
@@ -291,7 +294,12 @@ class Rope2DPosEmb(nn.Module):
294
"""
295
296
def __init__(
294
- self, dim: int, max_height: int, max_width: int, theta_base=10000, device="cuda"
297
+ self,
298
+ dim: int,
299
+ max_height: int,
300
+ max_width: int,
301
+ theta_base=10000,
302
+ device=current_platform.device_type,
303
):
304
super().__init__()
305
self.dim = dim
@@ -437,7 +445,7 @@ class MoonVitEncoderLayer(nn.Module):
445
self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
446
self.attn_implementation = attn_implementation
447
# use fa2 in vllm by default
440
- if is_flash_attn_2_available():
448
+ if is_flash_attn_2_available() or current_platform.is_xpu():
449
self.attn_implementation = "flash_attention_2"
450
451
self.norm0 = nn.LayerNorm(hidden_dim)