6f91edf96d · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 -import gc

4 import tempfile

5 from contextlib import contextmanager

6

@@ -9,8 +8,7 @@ import pytest

8 import torch

9

10 from tests.models.utils import check_logprobs_close

12 -from tests.utils import wait_for_rocm_memory_to_settle

13 -from vllm import LLM, SamplingParams

11 +from vllm import SamplingParams

12 from vllm.compilation.decorators import support_torch_compile

13 from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config

14 from vllm.config.compilation import (

@@ -49,6 +47,7 @@ def get_test_models():

47 @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")

48 def test_dynamic_shapes_compilation(

49 monkeypatch,

50 + vllm_runner,

51 model_name,

52 shapes_type,

53 use_aot_compile,

@@ -79,9 +78,13 @@ def test_dynamic_shapes_compilation(

78

79 print(f"Testing {shapes_type.name} dynamic shapes...")

80

82 - # Initialize the model with specific dynamic shapes configuration

83 - model = LLM(

84 - model=model_name,

81 + sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10)

82 + test_prompts = [prompt, "The capital of France is"]

83 +

84 + # VllmRunner shuts down the engine core on exit, so the eager model

85 + # below never races a lingering compiled engine for GPU memory.

86 + with vllm_runner(

87 + model_name,

88 compilation_config={

89 "mode": CompilationMode.VLLM_COMPILE,

90 "dynamic_shapes_config": {

@@ -90,33 +93,25 @@ def test_dynamic_shapes_compilation(

93 },

94 },

95 max_model_len=1024,

93 - )

94 -

95 - sampling_params = SamplingParams(max_tokens=5, temperature=0, logprobs=10)

96 - test_prompts = [prompt, "The capital of France is"]

97 -

98 - compiled_outputs = []

99 - for p in test_prompts:

100 - output = model.generate(p, sampling_params)[0].outputs[0]

101 - assert len(output.text.strip()) > 0, "Compiled model produced empty output"

102 - compiled_outputs.append((output.token_ids, output.text, output.logprobs))

103 -

104 - del model

105 - gc.collect()

106 - torch.accelerator.empty_cache()

107 - torch.accelerator.synchronize()

108 - wait_for_rocm_memory_to_settle()

109 -

110 - eager_model = LLM(model=model_name, enforce_eager=True, max_model_len=1024)

111 - eager_outputs = []

112 - for p in test_prompts:

113 - output = eager_model.generate(p, sampling_params)[0].outputs[0]

114 - assert len(output.text.strip()) > 0, "Eager model produced empty output"

115 - eager_outputs.append((output.token_ids, output.text, output.logprobs))

116 - del eager_model

117 - gc.collect()

118 - torch.accelerator.empty_cache()

119 - torch.accelerator.synchronize()

96 + enable_chunked_prefill=None,

97 + ) as vllm_model:

98 + compiled_outputs = []

99 + for p in test_prompts:

100 + output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0]

101 + assert len(output.text.strip()) > 0, "Compiled model produced empty output"

102 + compiled_outputs.append((output.token_ids, output.text, output.logprobs))

103 +

104 + with vllm_runner(

105 + model_name,

106 + enforce_eager=True,

107 + max_model_len=1024,

108 + enable_chunked_prefill=None,

109 + ) as vllm_model:

110 + eager_outputs = []

111 + for p in test_prompts:

112 + output = vllm_model.llm.generate(p, sampling_params)[0].outputs[0]

113 + assert len(output.text.strip()) > 0, "Eager model produced empty output"

114 + eager_outputs.append((output.token_ids, output.text, output.logprobs))

115

116 check_logprobs_close(

117 outputs_0_lst=eager_outputs,

@@ -241,44 +236,39 @@ def test_model_specialization_with_evaluate_guards(

236

237

238 @pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")

244 -def test_piecewise_backend_empty_sym_shape_indices():

239 +def test_piecewise_backend_empty_sym_shape_indices(vllm_runner):

240 """Test that PiecewiseBackend handles empty sym_shape_indices correctly.

241

242 When all inputs have static shapes (no torch.SymInt), sym_shape_indices

243 will be empty. The fix in PiecewiseBackend.__call__ handles this case

244 by using the first compiled range_entry.

245 """

251 - gc.collect()

252 - torch.accelerator.empty_cache()

253 - torch.accelerator.synchronize()

254 -

246 # Use small max_model_len and max_num_batched_tokens to encourage

247 # static shape compilation with empty sym_shape_indices

257 - llm = LLM(

258 - model="Qwen/Qwen3-0.6B",

248 + with vllm_runner(

249 + "Qwen/Qwen3-0.6B",

250 max_model_len=512,

251 max_num_batched_tokens=1,

252 + enable_chunked_prefill=None,

253 compilation_config={

254 "mode": CompilationMode.VLLM_COMPILE,

255 "dynamic_shapes_config": {

256 "type": DynamicShapesType.BACKED.value,

257 },

258 },

267 - )

268 -

269 - sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)

259 + ) as vllm_model:

260 + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)

261

271 - # Generate with static shape inputs

272 - output = llm.generate("Hello, my name is", sampling_params=sampling_params)

273 - result = output[0].outputs[0].text

274 - assert len(result) > 0, "Should generate non-empty output"

275 -

276 - # Generate again to verify compilation works with empty sym_shape_indices

277 - output = llm.generate("The capital of France is", sampling_params=sampling_params)

278 - result = output[0].outputs[0].text

279 - assert len(result) > 0, "Should generate non-empty output on second run"

262 + # Generate with static shape inputs

263 + output = vllm_model.llm.generate(

264 + "Hello, my name is", sampling_params=sampling_params

265 + )

266 + result = output[0].outputs[0].text

267 + assert len(result) > 0, "Should generate non-empty output"

268

281 - del llm

282 - gc.collect()

283 - torch.accelerator.empty_cache()

284 - torch.accelerator.synchronize()

269 + # Generate again to verify compilation works with empty sym_shape_indices

270 + output = vllm_model.llm.generate(

271 + "The capital of France is", sampling_params=sampling_params

272 + )

273 + result = output[0].outputs[0].text

274 + assert len(result) > 0, "Should generate non-empty output on second run"

[Test] dynamic_shapes_compilation (#49974)