From 9f1ea24470a201755313114a537751d393e49ec5 Mon Sep 17 00:00:00 2001 From: Yeonsil Yoon Date: Wed, 16 Sep 2026 10:25:00 -0700 Subject: [PATCH] [XPU] Keep MoE expert weights offloaded during transpose The XPU branch of UnquantizedFusedMoEMethod.process_weights_after_loading runs w13/w2.transpose(-1, -2).contiguous(). When experts are CPU-offloaded via UVAOffloader they present as accelerator tensors backed by host memory (marked _vllm_is_uva_offloaded); .contiguous() then materializes a real on-device copy, silently de-offloading every expert and OOMing the card at load time, before KV profiling. Transpose offloaded weights on the host and rebuild the UVA view so they stay off-device; resident (non-offloaded) weights keep the original device path and are numerically unchanged. Re-assert the offload marker after _setup_kernel, since replace_parameter registers a fresh Parameter and drops the old one's attributes, so a later reload (e.g. RL weight update) still takes the host-transpose path. Scoped to UnquantizedMoeBackend.XPU only; a no-op for all non-offloaded and non-XPU paths. Signed-off-by: Yeonsil Yoon Co-authored-by: Claude Opus 4.8 --- .../fused_moe/unquantized_fused_moe_method.py | 35 +++++++++++++++++-- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 523d1df0d80e..d0d6b5c3c7e0 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -205,17 +205,46 @@ def process_weights_after_loading(self, layer: "RoutedExperts") -> None: return elif self.unquantized_backend == UnquantizedMoeBackend.XPU: + from vllm.utils.torch_utils import ( + get_accelerator_view_from_cpu_tensor, + ) + w13 = layer.w13_weight w2 = layer.w2_weight - - w13.data = w13.transpose(-1, -2).contiguous() - w2.data = w2.transpose(-1, -2).contiguous() + w13_offloaded = getattr(w13, "_vllm_is_uva_offloaded", False) + w2_offloaded = getattr(w2, "_vllm_is_uva_offloaded", False) + + def _xpu_transpose(w, offloaded): + # CPU-offloaded expert weights (UVAOffloader) present as + # accelerator tensors backed by host memory. A plain + # .transpose().contiguous() materializes a real on-device + # copy, silently de-offloading every expert and OOMing the + # card. Transpose on the host and rebuild the UVA view to keep + # the offload; resident weights take the device path. + if offloaded: + cpu_t = w.data.to("cpu").transpose(-1, -2).contiguous() + w.data = get_accelerator_view_from_cpu_tensor(cpu_t) + w._vllm_is_uva_offloaded = True + else: + w.data = w.transpose(-1, -2).contiguous() + + _xpu_transpose(w13, w13_offloaded) + _xpu_transpose(w2, w2_offloaded) self._setup_kernel( layer=layer, w13=w13, w2=w2, ) + + # _setup_kernel -> replace_parameter registers a fresh Parameter + # and drops attributes of the old one; re-assert the offload marker + # so a later reload (e.g. RL weight update) still detects the + # offload and takes the host-transpose path above. + if w13_offloaded: + layer.w13_weight._vllm_is_uva_offloaded = True + if w2_offloaded: + layer.w2_weight._vllm_is_uva_offloaded = True else: self._setup_kernel( layer=layer,