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 7, 22:45 UTC
[ROCm][CI] Fix flaky embedding chat test by using tolerance-based comparison (#35050)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Feb 22, 2026 at 03:03 UTC
a8a47c17b68fbd4229a86cc1d4202ebc94bdb9fe
1 file changed
+17
-7
tests/entrypoints/pooling/embed/test_online.py
+17
-7
@@ -58,13 +58,19 @@ if current_platform.is_rocm():
58
torch.backends.cuda.enable_mem_efficient_sdp(False)
59
torch.backends.cuda.enable_math_sdp(True)
60
61
+# On ROCm, floating-point reductions in attention and GEMM kernels are
62
+# non-associative and sensitive to batch geometry. Force LLM instances
63
+# into an identical, deterministic execution mode:
64
+ROCM_DETERMINISM_ARGS: list[str] = (
65
+ ["--max-num-seqs", "1"] if current_platform.is_rocm() else []
66
+)
67
+
68
69
@pytest.fixture(scope="module")
70
def server():
71
args = [
72
"--runner",
73
"pooling",
67
- # use half precision for speed and memory savings in CI environment
74
"--dtype",
75
DTYPE,
76
"--enforce-eager",
@@ -72,12 +78,9 @@ def server():
78
"512",
79
"--chat-template",
80
DUMMY_CHAT_TEMPLATE,
81
+ *ROCM_DETERMINISM_ARGS,
82
]
83
77
- # ROCm: Use Flex Attention to support encoder-only self-attention.
78
- if current_platform.is_rocm():
79
- args.extend(["--attention-backend", "FLEX_ATTENTION"])
80
-
84
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
85
yield remote_server
86
@@ -343,8 +346,15 @@ async def test_chat_request(
346
assert chat_embeddings.id is not None
347
assert completion_embeddings.id is not None
348
assert chat_embeddings.created <= completion_embeddings.created
346
- assert chat_embeddings.model_dump(exclude={"id", "created"}) == (
347
- completion_embeddings.model_dump(exclude={"id", "created"})
349
+ # Use tolerance-based comparison for embeddings
350
+ check_embeddings_close(
351
+ embeddings_0_lst=[d.embedding for d in chat_embeddings.data],
352
+ embeddings_1_lst=[d.embedding for d in completion_embeddings.data],
353
+ name_0="chat",
354
+ name_1="completion",
355
+ )
356
+ assert chat_embeddings.model_dump(exclude={"id", "created", "data"}) == (
357
+ completion_embeddings.model_dump(exclude={"id", "created", "data"})
358
)
359
360
# test add_generation_prompt