From e24788138b2e5a9f2813f1a416a8b51f6b7d888b Mon Sep 17 00:00:00 2001 From: Emma Lien Date: Mon, 30 Mar 2026 08:27:28 +0000 Subject: [PATCH 1/3] fix(tpu): fix LoRA crashes, engine interface mismatch, and loader stability --- vllm/lora/layers/logits_processor.py | 3 ++ vllm/lora/layers/utils.py | 42 ++++++++++--------- vllm/lora/layers/vocal_parallel_embedding.py | 4 +- .../layers/quantization/mxfp4.py | 1 + .../model_loader/default_loader.py | 9 ++-- vllm/v1/engine/input_processor.py | 7 +++- 6 files changed, 41 insertions(+), 25 deletions(-) diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index 237a61eace1e..b79e7ad58e50 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/vllm/lora/layers/logits_processor.py @@ -158,6 +158,9 @@ def _get_logits( if logits is None: return None + + if current_platform.is_tpu(): + return logits[:, : self.base_layer.vocab_size] if self.sharded_to_full_mapping_gpu is not None: # Reindex full logits tensor to ensure 1:1 mapping between diff --git a/vllm/lora/layers/utils.py b/vllm/lora/layers/utils.py index c19b097586f5..0b6f3471e726 100644 --- a/vllm/lora/layers/utils.py +++ b/vllm/lora/layers/utils.py @@ -32,26 +32,28 @@ def __post_init__(self): def _get_lora_device(base_layer: nn.Module) -> torch.device: # code borrowed from https://github.com/fmmoret/vllm/blob/fm-support-lora-on-quantized-models/vllm/lora/layers.py#L34 """Returns the device for where to place the LoRA tensors.""" - # unquantizedLinear - if hasattr(base_layer, "weight"): - return base_layer.weight.device - # Compressed Tensor - elif hasattr(base_layer, "weight_packed"): - return base_layer.weight_packed.device - # GPTQ/AWQ - elif hasattr(base_layer, "qweight"): - return base_layer.qweight.device - # MoE layer - elif hasattr(base_layer, "w2_weight"): - return base_layer.w2_weight.device - # MoE Compressed Tensor - elif hasattr(base_layer, "w2_weight_packed"): - return base_layer.w2_weight_packed.device - # MoE GPTQ/AWQ/GGUF - elif hasattr(base_layer, "w2_qweight"): - return base_layer.w2_qweight.device - else: - raise ValueError(f"Unsupported base layer: {base_layer}") + + def get_dev(obj): + if obj is None: + return None + if isinstance(obj, (nn.ParameterList, list, tuple)): + return obj[0].device if len(obj) > 0 else None + if hasattr(obj, "device"): + return obj.device + return None + + attr_names = ["weight", "weight_packed", "qweight", + "w2_weight", "w2_weight_packed", "w2_qweight"] + for attr in attr_names: + target = getattr(base_layer, attr, None) + dev = get_dev(target) + if dev is not None: + return dev + + try: + return next(base_layer.parameters()).device + except StopIteration: + return torch.device("cpu") def _not_fully_sharded_can_replace(can_replace): diff --git a/vllm/lora/layers/vocal_parallel_embedding.py b/vllm/lora/layers/vocal_parallel_embedding.py index 05e7cfa06c85..81cf9273676c 100644 --- a/vllm/lora/layers/vocal_parallel_embedding.py +++ b/vllm/lora/layers/vocal_parallel_embedding.py @@ -94,6 +94,9 @@ def set_lora( ) def forward(self, x: torch.Tensor) -> torch.Tensor: + full_output = self.base_layer.forward(x) + if current_platform.is_tpu(): + return full_output # NB: Don't use torch.narrow here. torch.narrow triggers some # Dynamic Shape specialization in torch.compile num_tokens = x.shape[0] @@ -103,7 +106,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x + indices_1, self.lora_a_stacked_2d, ) - full_output = self.base_layer.forward(x) full_output_org = full_output if full_output.ndim == 3: diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 8ce5432fed83..4edc17267460 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -24,6 +24,7 @@ mxfp4_round_up_hidden_size_and_intermediate_size, select_mxfp4_moe_backend, ) +Mxfp4Backend = Mxfp4MoeBackend from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import QuantizationMethods from vllm.model_executor.layers.quantization.base_config import ( diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 5c9c97f4b64a..4984d58d00af 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -308,9 +308,12 @@ def _init_ep_weight_filter(self, model_config: ModelConfig) -> None: expert weights. By computing the set upfront we can skip non-local expert tensors *before* reading them from disk. """ - from vllm.config import get_current_vllm_config - - vllm_config = get_current_vllm_config() + try: + from vllm.config import get_current_vllm_config + vllm_config = get_current_vllm_config() + except AssertionError: + logger.warning("vLLM config not set during weight filtering, skipping EP filter.") + return parallel_config = vllm_config.parallel_config if not ( diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index b77b9277a48d..629047077197 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -248,7 +248,12 @@ def process_inputs( tokenization_kwargs=tokenization_kwargs, ) - current_platform.validate_request(processed_inputs, params) + if not current_platform.is_tpu(): + current_platform.validate_request(processed_inputs, params) + else: + # todo(TPU): TpuPlatform.validate_request interface is currently + # inconsistent with V1 InputProcessor. Skipping for now. + logger.debug("Skipping validate_request on TPU due to interface mismatch.") encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs) self._validate_model_inputs(encoder_inputs, decoder_inputs) From 636ca8ae80f818db72af806ec7eba2646d8450f1 Mon Sep 17 00:00:00 2001 From: Emma Lien Date: Tue, 31 Mar 2026 02:00:16 +0000 Subject: [PATCH 2/3] chore: revert temporary TPU compatibility patches and sync with latest tpu-inference --- vllm/model_executor/layers/quantization/mxfp4.py | 1 - vllm/v1/engine/input_processor.py | 7 +------ 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 4edc17267460..8ce5432fed83 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -24,7 +24,6 @@ mxfp4_round_up_hidden_size_and_intermediate_size, select_mxfp4_moe_backend, ) -Mxfp4Backend = Mxfp4MoeBackend from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import QuantizationMethods from vllm.model_executor.layers.quantization.base_config import ( diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 629047077197..b77b9277a48d 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -248,12 +248,7 @@ def process_inputs( tokenization_kwargs=tokenization_kwargs, ) - if not current_platform.is_tpu(): - current_platform.validate_request(processed_inputs, params) - else: - # todo(TPU): TpuPlatform.validate_request interface is currently - # inconsistent with V1 InputProcessor. Skipping for now. - logger.debug("Skipping validate_request on TPU due to interface mismatch.") + current_platform.validate_request(processed_inputs, params) encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs) self._validate_model_inputs(encoder_inputs, decoder_inputs) From 024604843aac1d36308b8fa64479142eeb4e8109 Mon Sep 17 00:00:00 2001 From: Emma Lien Date: Mon, 13 Apr 2026 01:48:47 +0000 Subject: [PATCH 3/3] fix(lora): refactor lora device retrieval --- vllm/lora/layers/logits_processor.py | 3 - vllm/lora/layers/utils.py | 1 - vllm/lora/layers/vocal_parallel_embedding.py | 4 +- vllm/lora/model_manager.py | 184 +++++++++--------- .../model_loader/default_loader.py | 9 +- 5 files changed, 97 insertions(+), 104 deletions(-) diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index b79e7ad58e50..237a61eace1e 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/vllm/lora/layers/logits_processor.py @@ -158,9 +158,6 @@ def _get_logits( if logits is None: return None - - if current_platform.is_tpu(): - return logits[:, : self.base_layer.vocab_size] if self.sharded_to_full_mapping_gpu is not None: # Reindex full logits tensor to ensure 1:1 mapping between diff --git a/vllm/lora/layers/utils.py b/vllm/lora/layers/utils.py index 0b6f3471e726..d942228ddd1c 100644 --- a/vllm/lora/layers/utils.py +++ b/vllm/lora/layers/utils.py @@ -32,7 +32,6 @@ def __post_init__(self): def _get_lora_device(base_layer: nn.Module) -> torch.device: # code borrowed from https://github.com/fmmoret/vllm/blob/fm-support-lora-on-quantized-models/vllm/lora/layers.py#L34 """Returns the device for where to place the LoRA tensors.""" - def get_dev(obj): if obj is None: return None diff --git a/vllm/lora/layers/vocal_parallel_embedding.py b/vllm/lora/layers/vocal_parallel_embedding.py index 81cf9273676c..05e7cfa06c85 100644 --- a/vllm/lora/layers/vocal_parallel_embedding.py +++ b/vllm/lora/layers/vocal_parallel_embedding.py @@ -94,9 +94,6 @@ def set_lora( ) def forward(self, x: torch.Tensor) -> torch.Tensor: - full_output = self.base_layer.forward(x) - if current_platform.is_tpu(): - return full_output # NB: Don't use torch.narrow here. torch.narrow triggers some # Dynamic Shape specialization in torch.compile num_tokens = x.shape[0] @@ -106,6 +103,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x + indices_1, self.lora_a_stacked_2d, ) + full_output = self.base_layer.forward(x) full_output_org = full_output if full_output.ndim == 3: diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 9d3772560433..97fc6fdaa7b9 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -6,6 +6,7 @@ from typing import TypeVar import torch +import torchax from torch import nn from vllm.config import VllmConfig @@ -481,101 +482,102 @@ def create_dummy_lora( ) -> LoRAModel: """Create zero-initialized LoRAModel for warmup.""" model = LoRAModel(lora_id, rank, {}) - for module_name, module in self.model.named_modules(): - if ( - not self._match_target_modules(module_name) - or not isinstance(module, BaseLayerWithLoRA) - or self._get_punica_wrapper(module_name) is None - ): - continue - parts = module_name.split(".") - if module_name not in self.packed_modules: - assert embedding_modules is not None - if parts[-1] in embedding_modules: - # Special-case lm_head: wrapped by LogitsProcessorWithLoRA. - # LoRA input dim is hidden_size, output dim is vocab size. - # LogitsProcessorWithLoRA handles extra vocab size directly. - if parts[-1] == "lm_head": - input_dim = module.lora_a_stacked[0].shape[-1] - output_dim = module.lora_b_stacked[0].shape[-2] - else: - input_dim = ( - module.base_layer.org_vocab_size - if hasattr(module.base_layer, "org_vocab_size") - else module.base_layer.weight.shape[1] + with torchax.default_env(): + for module_name, module in self.model.named_modules(): + if ( + not self._match_target_modules(module_name) + or not isinstance(module, BaseLayerWithLoRA) + or self._get_punica_wrapper(module_name) is None + ): + continue + parts = module_name.split(".") + if module_name not in self.packed_modules: + assert embedding_modules is not None + if parts[-1] in embedding_modules: + # Special-case lm_head: wrapped by LogitsProcessorWithLoRA. + # LoRA input dim is hidden_size, output dim is vocab size. + # LogitsProcessorWithLoRA handles extra vocab size directly. + if parts[-1] == "lm_head": + input_dim = module.lora_a_stacked[0].shape[-1] + output_dim = module.lora_b_stacked[0].shape[-2] + else: + input_dim = ( + module.base_layer.org_vocab_size + if hasattr(module.base_layer, "org_vocab_size") + else module.base_layer.weight.shape[1] + ) + output_dim = ( + module.base_layer.embedding_dim + if hasattr(module.base_layer, "embedding_dim") + else module.base_layer.weight.shape[0] + ) + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + input_dim, + output_dim, + rank, + module.lora_a_stacked[0].dtype, + "cpu", ) - output_dim = ( - module.base_layer.embedding_dim - if hasattr(module.base_layer, "embedding_dim") - else module.base_layer.weight.shape[0] + model.loras[module_name] = lora + elif module.__class__.__name__ == "FusedMoE3DWithLoRA": + # Case for 3D moe model + # w2 + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + module.w2_input_size, + module.w2_output_size, + rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts + module.w2_lora_a_stacked[0].dtype, + "cpu", ) - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - input_dim, - output_dim, - rank, - module.lora_a_stacked[0].dtype, - "cpu", - ) - model.loras[module_name] = lora - elif module.__class__.__name__ == "FusedMoE3DWithLoRA": - # Case for 3D moe model - # w2 - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - module.w2_input_size, - module.w2_output_size, - rank * module.w2_lora_a_stacked[0].shape[1], # rank*num_experts - module.w2_lora_a_stacked[0].dtype, - "cpu", - ) - model.loras[module_name] = lora - # w13 - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - module.w13_input_size, - module.w13_output_size, - rank - * module.w13_lora_a_stacked[0].shape[1], # rank*num_experts - module.w13_lora_a_stacked[0].dtype, - "cpu", - ) - model.loras[module_name + ".base_layer"] = lora + model.loras[module_name] = lora + # w13 + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + module.w13_input_size, + module.w13_output_size, + rank + * module.w13_lora_a_stacked[0].shape[1], # rank*num_experts + module.w13_lora_a_stacked[0].dtype, + "cpu", + ) + model.loras[module_name + ".base_layer"] = lora + else: + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name, + module.lora_a_stacked[0].shape[-1], + module.lora_b_stacked[0].shape[-2], + rank, + module.lora_a_stacked[0].dtype, + "cpu", + ) + model.loras[module_name] = lora else: - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name, - module.lora_a_stacked[0].shape[-1], - module.lora_b_stacked[0].shape[-2], - rank, - module.lora_a_stacked[0].dtype, - "cpu", - ) + parts = module_name.split(".") + replacements = self.packed_modules_mapping[parts[-1]] + subloras: list[LoRALayerWeights | None] = [] + for i, r in enumerate(replacements): + lora = LoRALayerWeights.create_dummy_lora_weights( + module_name + "." + r, + module.lora_a_stacked[i].shape[-1], + module.lora_b_stacked[i].shape[-2], + rank, + module.lora_a_stacked[i].dtype, + "cpu", + ) + subloras.append(lora) + if module.__class__.__name__ == "FusedMoEWithLoRA": + # For non-gated MoE, pad subloras to 3 elements per expert + # to match pack_moe expectations (w1, w2, None for w3) + if self._is_non_gated_moe and len(subloras) > 0: + subloras = self._pad_lora_pairs_to_triplets(subloras) + lora = PackedLoRALayerWeights.pack_moe( + subloras, module_name, is_non_gated_moe=self._is_non_gated_moe + ) + else: + lora = PackedLoRALayerWeights.pack(subloras) model.loras[module_name] = lora - else: - parts = module_name.split(".") - replacements = self.packed_modules_mapping[parts[-1]] - subloras: list[LoRALayerWeights | None] = [] - for i, r in enumerate(replacements): - lora = LoRALayerWeights.create_dummy_lora_weights( - module_name + "." + r, - module.lora_a_stacked[i].shape[-1], - module.lora_b_stacked[i].shape[-2], - rank, - module.lora_a_stacked[i].dtype, - "cpu", - ) - subloras.append(lora) - if module.__class__.__name__ == "FusedMoEWithLoRA": - # For non-gated MoE, pad subloras to 3 elements per expert - # to match pack_moe expectations (w1, w2, None for w3) - if self._is_non_gated_moe and len(subloras) > 0: - subloras = self._pad_lora_pairs_to_triplets(subloras) - lora = PackedLoRALayerWeights.pack_moe( - subloras, module_name, is_non_gated_moe=self._is_non_gated_moe - ) - else: - lora = PackedLoRALayerWeights.pack(subloras) - model.loras[module_name] = lora return model def _match_target_modules(self, module_name: str) -> bool: diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 4984d58d00af..5c9c97f4b64a 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -308,12 +308,9 @@ def _init_ep_weight_filter(self, model_config: ModelConfig) -> None: expert weights. By computing the set upfront we can skip non-local expert tensors *before* reading them from disk. """ - try: - from vllm.config import get_current_vllm_config - vllm_config = get_current_vllm_config() - except AssertionError: - logger.warning("vLLM config not set during weight filtering, skipping EP filter.") - return + from vllm.config import get_current_vllm_config + + vllm_config = get_current_vllm_config() parallel_config = vllm_config.parallel_config if not (