[Core][LoRA] Support fp32 lm_head (head_dtype) on the LoRA path (#48525)

Signed-off-by: Karthik Kothuri <[email protected]>
Co-authored-by: Claude <[email protected]>

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
         )

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.

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)