[ROCm][Bugfix] Enable the fp32 head_dtype torch.mm fast path on ROCm

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Turner Jabbour <doubleujabbour@gmail.com>

Author: wjabbour

tests/v1/sample/test_head_dtype.py

@@ -95,6 +95,37 @@ def test_head_dtype_equal_to_model_dtype_uses_quant_method(default_vllm_config):
     assert logits.dtype == torch.bfloat16
 
 
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Exercises the torch.mm(out_dtype=...) device fast path, "
+    "available on CUDA and ROCm.",
+)
+def test_fp32_head_uses_mm_fast_path_on_device(default_vllm_config):
+    # On ROCm, current_platform.is_cuda() is False, so this previously fell
+    # through to the cast path (F.linear) instead of torch.mm(out_dtype=...),
+    # even though ROCm supports the out_dtype mm via its non-Lt GEMM path.
+    from unittest import mock
+
+    vocab_size, hidden_size, num_tokens = 64, 16, 4
+    lp = _build_processor(vocab_size)
+    lp.head_dtype = torch.float32
+
+    hidden_states = torch.randn(
+        num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda"
+    )
+    weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16, device="cuda")
+
+    with mock.patch(
+        "vllm.model_executor.layers.logits_processor.F.linear"
+    ) as linear_mock:
+        logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)
+
+    linear_mock.assert_not_called()
+    assert logits.dtype == torch.float32
+    expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
+    torch.testing.assert_close(logits, expected)
+
+
 def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
     lp = _build_processor(64)
     lp.head_dtype = torch.float32

vllm/model_executor/layers/logits_processor.py

@@ -115,14 +115,15 @@ def _apply_head(
             )
         if (
             self.head_dtype == torch.float32
-            and current_platform.is_cuda()
+            and (current_platform.is_cuda() or current_platform.is_rocm())
             and hidden_states.is_cuda
         ):
             # Accumulate the projection directly into fp32. This avoids
             # materializing an fp32 copy of the lm_head weight on every step,
-            # unlike casting both operands. `torch.mm(out_dtype=...)` is
-            # CUDA-only and only supports fp32 output for fp16/bf16 inputs, so
-            # other cases fall back to the cast path below.
+            # unlike casting both operands. `torch.mm(out_dtype=...)` only
+            # supports fp32 output for fp16/bf16 inputs, and is only
+            # implemented for CUDA and ROCm (the latter via the non-Lt GEMM
+            # path); other platforms fall back to the cast path below.
             flat = hidden_states.reshape(-1, hidden_states.shape[-1])
             logits = torch.mm(flat, lm_head.weight.t(), out_dtype=self.head_dtype)
             if embedding_bias is not None: