A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 01:15 UTC
[ROCm][Test] Fix beam search determinism failures from batch-size-dependent FP divergence and removed wrong marker (#34878)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Feb 19, 2026 at 02:25 UTC
f6220f98779463705e578562f307b2becea8b8b3
3 files changed
+45
-8
.buildkite/test-amd.yaml
+2
-2
@@ -549,7 +549,7 @@ steps:
549
- tests/samplers
550
- tests/conftest.py
551
commands:
552
- - pytest -v -s -m samplers
552
+ - pytest -v -s samplers
553
554
- label: LoRA Test %N # 20min each
555
timeout_in_minutes: 30
@@ -2177,7 +2177,7 @@ steps:
2177
- tests/samplers
2178
- tests/conftest.py
2179
commands:
2180
- - pytest -v -s -m samplers
2180
+ - pytest -v -s samplers
2181
2182
- label: LoRA Test %N # 20min each
2183
timeout_in_minutes: 30
.buildkite/test_areas/samplers.yaml
+1
-1
@@ -18,4 +18,4 @@ steps:
18
depends_on:
19
- image-build-amd
20
commands:
21
- - pytest -v -s -m samplers
21
+ - pytest -v -s samplers
tests/samplers/test_beam_search.py
+42
-5
@@ -9,6 +9,26 @@ import pytest
9
from transformers import AutoModelForSeq2SeqLM
10
11
from vllm.assets.audio import AudioAsset
12
+from vllm.platforms import current_platform
13
+
14
+# Extra engine kwargs needed for numerically deterministic beam search.
15
+# On ROCm, floating-point reductions in attention and GEMM kernels are
16
+# non-associative and sensitive to batch geometry, so we:
17
+# async_scheduling=False – deterministic batch composition
18
+# enforce_eager=True – no CUDA-graph padding changing effective size
19
+# enable_prefix_caching=False – avoid prefix-sharing side effects
20
+# max_num_seqs=1 – fixed batch size across runs
21
+# On other platforms these are not needed and the dict is empty.
22
+EXTRA_ENGINE_KWARGS: dict = (
23
+ dict(
24
+ async_scheduling=False,
25
+ enforce_eager=True,
26
+ enable_prefix_caching=False,
27
+ max_num_seqs=1,
28
+ )
29
+ if current_platform.is_rocm()
30
+ else {}
31
+)
32
33
# FIXME(zhuohan): The test can not pass if we:
34
# 1. Increase max_tokens to 256.
@@ -25,6 +45,7 @@ MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
45
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
46
@pytest.mark.parametrize("beam_width", BEAM_WIDTHS)
47
def test_beam_search_single_input(
48
+ monkeypatch,
49
hf_runner,
50
vllm_runner,
51
example_prompts,
@@ -33,13 +54,16 @@ def test_beam_search_single_input(
54
max_tokens: int,
55
beam_width: int,
56
) -> None:
57
+ if current_platform.is_rocm():
58
+ monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
59
+
60
example_prompts = example_prompts[:1]
61
with hf_runner(model, dtype=dtype) as hf_model:
62
hf_outputs = hf_model.generate_beam_search(
63
example_prompts, beam_width, max_tokens
64
)
65
42
- with vllm_runner(model, dtype=dtype) as vllm_model:
66
+ with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
67
vllm_outputs = vllm_model.generate_beam_search(
68
example_prompts, beam_width, max_tokens
69
)
@@ -66,6 +90,7 @@ def test_beam_search_single_input(
90
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
91
@pytest.mark.parametrize("beam_width", BEAM_WIDTHS)
92
def test_beam_search_with_concurrency_limit(
93
+ monkeypatch,
94
hf_runner,
95
vllm_runner,
96
example_prompts,
@@ -74,21 +99,29 @@ def test_beam_search_with_concurrency_limit(
99
max_tokens: int,
100
beam_width: int,
101
) -> None:
102
+ if current_platform.is_rocm():
103
+ monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
104
+
105
# example_prompts[1]&[3]&[7] fails due to unknown reason even without
106
# concurrency limit. skip them for now.
107
example_prompts = example_prompts[:8]
108
concurrency_limit = 2
109
assert len(example_prompts) > concurrency_limit
82
- with vllm_runner(model, dtype=dtype) as vllm_model:
110
+ with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
111
outputs_with_limit = vllm_model.generate_beam_search(
84
- example_prompts, beam_width, max_tokens, concurrency_limit=concurrency_limit
112
+ example_prompts,
113
+ beam_width,
114
+ max_tokens,
115
+ concurrency_limit=concurrency_limit,
116
)
117
outputs_without_limit = []
118
119
for i in range(0, len(example_prompts), concurrency_limit):
120
outputs_without_limit.extend(
121
vllm_model.generate_beam_search(
91
- example_prompts[i : i + concurrency_limit], beam_width, max_tokens
122
+ example_prompts[i : i + concurrency_limit],
123
+ beam_width,
124
+ max_tokens,
125
)
126
)
127
@@ -118,6 +151,7 @@ def test_beam_search_with_concurrency_limit(
151
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
152
@pytest.mark.parametrize("beam_width", MM_BEAM_WIDTHS)
153
def test_beam_search_passes_multimodal_data(
154
+ monkeypatch,
155
hf_runner,
156
vllm_runner,
157
dtype: str,
@@ -125,6 +159,9 @@ def test_beam_search_passes_multimodal_data(
159
beam_width: int,
160
) -> None:
161
"""Ensure that beam search passes multimodal data through correctly."""
162
+ if current_platform.is_rocm():
163
+ monkeypatch.setenv("VLLM_ROCM_USE_SKINNY_GEMM", "0")
164
+
165
# NOTE - this test is primarily to check that mm data is passed to beams
166
# correctly. As such, we just need to check one extra modality to make
167
# sure things pass through properly.
@@ -145,7 +182,7 @@ def test_beam_search_passes_multimodal_data(
182
audios=audios,
183
)
184
148
- with vllm_runner(model, dtype=dtype) as vllm_model:
185
+ with vllm_runner(model, dtype=dtype, **EXTRA_ENGINE_KWARGS) as vllm_model:
186
vllm_outputs = vllm_model.generate_beam_search(
187
prompts,
188
beam_width=beam_width,