[XPU][LoRA] Fix torch.compile DEVICE_LOST by avoiding view-mutation in LoRA shrink (#47944)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
Jul 9, 2026 at 14:08 UTC
b8c7c86533d52e70433d9718ddafdd26542191f1
1 file changed
+15
-2
vllm/lora/punica_wrapper/punica_xpu.py
+15
-2
@@ -98,7 +98,14 @@ class PunicaWrapperXPU(PunicaWrapperBase):
98
w_t_all: torch.Tensor,
99
scale: float,
100
):
101
- bgmv_shrink(x, w_t_all, y, self._get_token_lora_indices(x), scale)
101
+ buf = torch.zeros(
102
+ x.size(0),
103
+ w_t_all.size(-2),
104
+ dtype=x.dtype,
105
+ device=x.device,
106
+ )
107
+ bgmv_shrink(x, w_t_all, buf, self._get_token_lora_indices(x), scale)
108
+ y.copy_(buf)
109
110
def _apply_expand(
111
self,
@@ -448,11 +455,17 @@ class PunicaWrapperXPU(PunicaWrapperBase):
455
_,
456
_,
457
lora_ids,
451
- _,
458
+ no_lora_flag,
459
num_active_loras,
460
) = self.token_mapping_meta.meta_args(
461
x.size(0), self.lora_config.specialize_active_lora
462
)
463
+
464
+ assert no_lora_flag.numel() == 1
465
+ if no_lora_flag.item():
466
+ # None of the inputs require LoRA.
467
+ return
468
+
469
if token_lora_mapping is None:
470
token_lora_mapping = token_lora_mapping_meta
471
fused_moe_lora(