@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 109 lines 3.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import asyncio
4 import os
5 from contextlib import AsyncExitStack
6 from dataclasses import replace
7
8 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 # On SM<90 (e.g., L4), batch invariance does not support CUDA graphs.
24 # See https://github.com/vllm-project/vllm/pull/30018 and
25 # tests/v1/determinism/utils.py for the documented limitation.
26 IS_DEVICE_CAPABILITY_BELOW_90 = not current_platform.has_device_capability(90)
27
28
29 @pytest.mark.asyncio
30 @pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
31 @pytest.mark.xfail(
32 current_platform.is_rocm(),
33 reason="Test may fail on ROCm until batch invariance is enabled. "
34 "See: https://github.com/vllm-project/vllm/issues/27433",
35 strict=False,
36 )
37 async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str):
38 if not current_platform.is_rocm() and not current_platform.is_xpu():
39 # This test checks that running a model with and without eagle
40 # leads to identical tokens.
41 #
42 # NOTE: This is only true in batch invariant mode
43 # (because the target model verifies all draft tokens in one big
44 # forward pass)
45 #
46 # TODO[ROCm]: Test is passing on ROCm CI but may break in future.
47 # Enable batch invariance for ROCm when possible. See:
48 # https://github.com/vllm-project/vllm/issues/27433
49
50 monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
51
52 target_model = "meta-llama/Llama-3.1-8B-Instruct"
53 draft_model = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
54
55 engine_args = AsyncEngineArgs(
56 model=target_model,
57 tokenizer_mode="auto",
58 enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
59 tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
60 data_parallel_size=DP_SIZE,
61 data_parallel_backend="mp", # ray takes more time
62 trust_remote_code=True,
63 max_model_len=16384,
64 attention_config={"backend": attn_backend},
65 )
66
67 eagle_engine_args = replace(
68 engine_args,
69 speculative_config={
70 "model": draft_model,
71 "method": "eagle",
72 "num_speculative_tokens": 3,
73 },
74 )
75
76 prompt = "This is a test of data parallel with eagle"
77 num_expected_tokens = 100
78 sampling_params = SamplingParams(
79 max_tokens=num_expected_tokens,
80 ignore_eos=True,
81 output_kind=RequestOutputKind.FINAL_ONLY,
82 temperature=0,
83 )
84
85 async def generate_with_timeout(given_engine: AsyncLLM):
86 async for out in given_engine.generate(
87 request_id="test-eagle-dp", prompt=prompt, sampling_params=sampling_params
88 ):
89 token_ids = out.outputs[0].token_ids
90 assert len(token_ids) == num_expected_tokens
91 return token_ids
92
93 async def engine_create_and_generate(engine_args: AsyncEngineArgs):
94 async with AsyncExitStack() as after:
95 engine = AsyncLLM.from_engine_args(engine_args)
96 after.callback(engine.shutdown)
97
98 token_ids = await asyncio.wait_for(
99 generate_with_timeout(engine), timeout=30
100 )
101
102 assert not engine.output_processor.has_unfinished_requests()
103 return token_ids
104
105 token_ids_with_eagle = await engine_create_and_generate(eagle_engine_args)
106 token_ids_no_eagle = await engine_create_and_generate(engine_args)
107
108 # Test for correctness
109 assert token_ids_with_eagle == token_ids_no_eagle