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, 22:45 UTC
[Bugfix] Fix eagle dp tests on A100 (#31241)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Dec 24, 2025 at 19:05 UTC
254f6b986720c92ddf97fbb1a6a6465da8e87e29
1 file changed
+7
-1
tests/v1/distributed/test_eagle_dp.py
+7
-1
@@ -16,7 +16,12 @@ DP_SIZE = int(os.getenv("DP_SIZE", 2))
16
17
18
@pytest.mark.asyncio
19
-async def test_run_eagle_dp():
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")
24
+
25
target_model = "meta-llama/Llama-3.1-8B-Instruct"
26
draft_model = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
27
@@ -29,6 +34,7 @@ async def test_run_eagle_dp():
34
data_parallel_backend="mp", # ray takes more time
35
trust_remote_code=True,
36
max_model_len=16384,
37
+ attention_config={"backend": "FLASH_ATTN"},
38
)
39
40
eagle_engine_args = replace(