[Bugfix][CI] Fix test_head_dtype quant_method test on ROCm (#48654)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Jul 14, 2026 at 18:31 UTC
0f0f28b537ac1a3d4bbeb74c706402d26dd27c76
1 file changed
+13
-1
tests/v1/sample/test_head_dtype.py
+13
-1
@@ -72,14 +72,26 @@ def test_non_fp32_head_dtype_uses_cast_path(default_vllm_config):
72
73
74
def test_head_dtype_equal_to_model_dtype_uses_quant_method(default_vllm_config):
75
+ from unittest import mock
76
+
77
vocab_size, hidden_size = 64, 16
78
lp = _build_processor(vocab_size)
79
lp.head_dtype = torch.bfloat16
80
81
hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
82
weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
83
+ lm_head = _FakeLmHead(weight)
84
82
- logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)
85
+ with mock.patch.object(
86
+ lm_head.quant_method,
87
+ "apply",
88
+ side_effect=lambda layer, x, bias=None: torch.nn.functional.linear(
89
+ x, layer.weight, bias
90
+ ),
91
+ ) as apply_mock:
92
+ logits = lp._get_logits(hidden_states, lm_head, None)
93
+
94
+ apply_mock.assert_called_once()
95
assert logits.dtype == torch.bfloat16
96
97