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 10, 00:15 UTC
[ROCm][CI] Add ROCm attention backend support for EAGLE DP tests (#32363)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jan 19, 2026 at 03:57 UTC
c0a350ca73fa8f10bd718fdcec47b075dc1d102f
1 file changed
+28
-6
tests/v1/distributed/test_eagle_dp.py
+28
-6
@@ -9,18 +9,40 @@ import pytest
9
10
from vllm import SamplingParams
11
from vllm.engine.arg_utils import AsyncEngineArgs
12
+from vllm.platforms import current_platform
13
from vllm.sampling_params import RequestOutputKind
14
from vllm.v1.engine.async_llm import AsyncLLM
15
16
DP_SIZE = int(os.getenv("DP_SIZE", 2))
17
18
+if current_platform.is_rocm():
19
+ ATTN_BACKENDS = ["ROCM_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"]
20
+else:
21
+ ATTN_BACKENDS = ["FLASH_ATTN"]
22
+
23
24
@pytest.mark.asyncio
19
-async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch):
20
- # This test checks that running a model with and without eagle
21
- # leads to identical tokens. This is only true in batch invariant mode
22
- # (because the target model verifies all draft tokens in one big forward pass)
23
- monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
25
+@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
26
+@pytest.mark.xfail(
27
+ current_platform.is_rocm(),
28
+ reason="Test may fail on ROCm until batch invariance is enabled."
29
+ "See: https://github.com/vllm-project/vllm/issues/27433",
30
+ strict=False,
31
+)
32
+async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str):
33
+ if not current_platform.is_rocm():
34
+ # This test checks that running a model with and without eagle
35
+ # leads to identical tokens.
36
+ #
37
+ # NOTE: This is only true in batch invariant mode
38
+ # (because the target model verifies all draft tokens in one big
39
+ # forward pass)
40
+ #
41
+ # TODO[ROCm]: Test is passing on ROCm CI but may break in future.
42
+ # Enable batch invariance for ROCm when possible. See:
43
+ # https://github.com/vllm-project/vllm/issues/27433
44
+
45
+ monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
46
47
target_model = "meta-llama/Llama-3.1-8B-Instruct"
48
draft_model = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
@@ -34,7 +56,7 @@ async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch):
56
data_parallel_backend="mp", # ray takes more time
57
trust_remote_code=True,
58
max_model_len=16384,
37
- attention_config={"backend": "FLASH_ATTN"},
59
+ attention_config={"backend": attn_backend},
60
)
61
62
eagle_engine_args = replace(