[ROCm][Bugfix] Enable the fp32 head_dtype torch.mm fast path on ROCm (#48688)
Signed-off-by: Turner Jabbour <doubleujabbour@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Turner Jabbour committed
Jul 15, 2026 at 02:18 UTC
7aab6e2684894fb1dba17cae117f8e84c5a7bbb3
2 files changed
+36
-4
tests/v1/sample/test_head_dtype.py
+31
@@ -95,6 +95,37 @@ def test_head_dtype_equal_to_model_dtype_uses_quant_method(default_vllm_config):
95
assert logits.dtype == torch.bfloat16
96
97
98
+@pytest.mark.skipif(
99
+ not torch.cuda.is_available(),
100
+ reason="Exercises the torch.mm(out_dtype=...) device fast path, "
101
+ "available on CUDA and ROCm.",
102
+)
103
+def test_fp32_head_uses_mm_fast_path_on_device(default_vllm_config):
104
+ # On ROCm, current_platform.is_cuda() is False, so this previously fell
105
+ # through to the cast path (F.linear) instead of torch.mm(out_dtype=...),
106
+ # even though ROCm supports the out_dtype mm via its non-Lt GEMM path.
107
+ from unittest import mock
108
+
109
+ vocab_size, hidden_size, num_tokens = 64, 16, 4
110
+ lp = _build_processor(vocab_size)
111
+ lp.head_dtype = torch.float32
112
+
113
+ hidden_states = torch.randn(
114
+ num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda"
115
+ )
116
+ weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16, device="cuda")
117
+
118
+ with mock.patch(
119
+ "vllm.model_executor.layers.logits_processor.F.linear"
120
+ ) as linear_mock:
121
+ logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)
122
+
123
+ linear_mock.assert_not_called()
124
+ assert logits.dtype == torch.float32
125
+ expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
126
+ torch.testing.assert_close(logits, expected)
127
+
128
+
129
def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
130
lp = _build_processor(64)
131
lp.head_dtype = torch.float32
vllm/model_executor/layers/logits_processor.py
+5
-4
@@ -115,14 +115,15 @@ class LogitsProcessor(PluggableLayer):
115
)
116
if (
117
self.head_dtype == torch.float32
118
- and current_platform.is_cuda()
118
+ and (current_platform.is_cuda() or current_platform.is_rocm())
119
and hidden_states.is_cuda
120
):
121
# Accumulate the projection directly into fp32. This avoids
122
# materializing an fp32 copy of the lm_head weight on every step,
123
- # unlike casting both operands. `torch.mm(out_dtype=...)` is
124
- # CUDA-only and only supports fp32 output for fp16/bf16 inputs, so
125
- # other cases fall back to the cast path below.
123
+ # unlike casting both operands. `torch.mm(out_dtype=...)` only
124
+ # supports fp32 output for fp16/bf16 inputs, and is only
125
+ # implemented for CUDA and ROCm (the latter via the non-Lt GEMM
126
+ # path); other platforms fall back to the cast path below.
127
flat = hidden_states.reshape(-1, hidden_states.shape[-1])
128
logits = torch.mm(flat, lm_head.weight.t(), out_dtype=self.head_dtype)
129
if embedding_bias is not None: