From 4f7fffb92f6e50fbbee47a89f6ebb4b662aba135 Mon Sep 17 00:00:00 2001 From: Karthik Kothuri Date: Tue, 14 Jul 2026 18:51:09 -0700 Subject: [PATCH] [Core][LoRA] Support fp32 lm_head (head_dtype) on the LoRA path (#48525) Signed-off-by: Karthik Kothuri Co-authored-by: Claude --- tests/lora/test_layers.py | 5 ++++- tests/v1/sample/test_head_dtype.py | 16 ---------------- vllm/lora/layers/logits_processor.py | 19 +++++++------------ 3 files changed, 11 insertions(+), 29 deletions(-) diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index c366b2cf297..95964e9f286 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -368,8 +368,9 @@ def test_embeddings( @pytest.mark.parametrize("device", DEVICES) @pytest.mark.parametrize("vocab_size", [64000, 256512, 258048]) @pytest.mark.parametrize("stage", STAGES) +@pytest.mark.parametrize("head_dtype", [None, torch.float32]) def test_lm_head_logits_processor( - default_vllm_config, dist_init, num_loras, device, vocab_size, stage + default_vllm_config, dist_init, num_loras, device, vocab_size, stage, head_dtype ) -> None: if current_platform.is_cuda_alike() or current_platform.is_xpu(): torch.accelerator.set_device_index(device) @@ -391,6 +392,8 @@ def test_lm_head_logits_processor( linear.weight.data = torch.rand_like(linear.weight.data) linear.weight.data[:, vocab_size:] = 0 logits_processor = LogitsProcessor(vocab_size) + # Exercise an fp32 lm_head (head_dtype) on the LoRA path. + logits_processor.head_dtype = head_dtype lora_logits_processor = LogitsProcessorWithLoRA( logits_processor, 1024, linear.weight.dtype, linear.weight.device, None ) diff --git a/tests/v1/sample/test_head_dtype.py b/tests/v1/sample/test_head_dtype.py index 22276a4c28b..37f5a6e4531 100644 --- a/tests/v1/sample/test_head_dtype.py +++ b/tests/v1/sample/test_head_dtype.py @@ -136,22 +136,6 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config): assert torch.equal(top, expected) -def test_fp32_head_rejected_with_lora(default_vllm_config): - from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA - - base = _build_processor(64) - base.head_dtype = torch.float32 - - with pytest.raises(ValueError, match="not yet supported with LoRA"): - LogitsProcessorWithLoRA( - base, - hidden_size=16, - dtype=torch.bfloat16, - device=torch.device("cpu"), - sharded_to_full_mapping=None, - ) - - @pytest.mark.core_model def test_fp32_head_e2e_no_nan(): """An fp32 head produces finite logprobs end-to-end. diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index 31fa79698fb..496413c21f4 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/vllm/lora/layers/logits_processor.py @@ -45,15 +45,6 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA): self.hidden_size = hidden_size self.dtype = dtype self.device = device - # The fp32 lm_head path lives in the base LogitsProcessor._get_logits, - # which this wrapper bypasses. Rather than silently emit model-dtype - # logits, reject the combination until the LoRA path supports it. - head_dtype = getattr(base_layer, "head_dtype", None) - if head_dtype is not None and head_dtype != dtype: - raise ValueError( - "A head_dtype different from the model dtype (e.g. an fp32 " - "lm_head) is not yet supported with LoRA." - ) self.tp_size = get_tensor_model_parallel_world_size() self.tp_rank = get_tensor_model_parallel_rank() self.sharded_to_full_mapping = sharded_to_full_mapping @@ -158,9 +149,13 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA): actual_lm_head = lm_head.base_layer else: actual_lm_head = lm_head - logits = actual_lm_head.quant_method.apply(actual_lm_head, hidden_states) - if embedding_bias is not None: - logits += embedding_bias + # Run the base projection through the LogitsProcessor so head_dtype + # (e.g. an fp32 lm_head) is honored on the LoRA path too. The LoRA + # delta is accumulated into these logits by add_lora_logits below, + # whose internal buffer is already fp32. + logits = self.base_layer._apply_head( + actual_lm_head, hidden_states, embedding_bias + ) # Gather logits for TP logits = self.base_layer._gather_logits(logits)