16d2ad1d38 · tduncan/Nvidia-Optimization-Stack-VLLM

.pre-commit-config.yaml

+7

@@ -127,6 +127,13 @@ repos:

127 language: python

128 types: [python]

129 additional_dependencies: [regex]

130 + # prevent use torch.cuda APIs

131 + - id: check-torch-cuda-call

132 + name: "Prevent new 'torch.cuda' APIs call"

133 + entry: python tools/pre_commit/check_torch_cuda.py

134 + language: python

135 + types: [python]

136 + additional_dependencies: [regex]

137 - id: validate-config

138 name: Validate configuration has default values and that each field has a docstring

139 entry: python tools/pre_commit/validate_config.py

benchmarks/benchmark_topk_topp.py

+1 -1

@@ -102,7 +102,7 @@ def reset_memory_stats():

102 """Reset peak memory statistics."""

103 reset_buffer_cache()

104 torch.cuda.reset_peak_memory_stats()

105 - torch.cuda.empty_cache()

105 + torch.accelerator.empty_cache()

106 gc.collect()

107

108

benchmarks/kernels/benchmark_moe.py

+1 -1

@@ -54,7 +54,7 @@ def clear_triton_cache():

54

55 # Clear CUDA memory cache

56 if torch.cuda.is_available():

57 - torch.cuda.empty_cache()

57 + torch.accelerator.empty_cache()

58

59 # Try to clear Triton's runtime cache

60 try:

benchmarks/kernels/benchmark_reshape_and_cache.py

+1 -1

@@ -104,7 +104,7 @@ def run_benchmark(

104

105 # free tensors to mitigate OOM when sweeping

106 del key, value, key_cache, value_cache, slot_mapping

107 - torch.cuda.empty_cache()

107 + torch.accelerator.empty_cache()

108

109 return lat

110

benchmarks/kernels/benchmark_reshape_and_cache_flash.py

+1 -1

@@ -129,7 +129,7 @@ def run_benchmark(

129

130 # free tensors to mitigate OOM when sweeping

131 del key, value, key_cache, value_cache, slot_mapping

132 - torch.cuda.empty_cache()

132 + torch.accelerator.empty_cache()

133

134 return lat

135

examples/offline_inference/lora_with_quantization_inference.py

+1 -1

@@ -120,7 +120,7 @@ def main():

120 # Clean up the GPU memory for the next test

121 del engine

122 gc.collect()

123 - torch.cuda.empty_cache()

123 + torch.accelerator.empty_cache()

124

125

126 if __name__ == "__main__":

examples/offline_inference/rlhf_colocate.py

+1 -1

@@ -159,7 +159,7 @@ class RayTrainingActor:

159 s.close()

160 del buffer

161 gc.collect()

162 - torch.cuda.empty_cache()

162 + torch.accelerator.empty_cache()

163

164

165 # Ray manages four GPUs.

examples/offline_inference/rlhf_utils.py

+1 -1

@@ -150,7 +150,7 @@ class ColocateWorkerExtension:

150 socket.close()

151 del buffer

152 gc.collect()

153 - torch.cuda.empty_cache()

153 + torch.accelerator.empty_cache()

154

155 def report_device_id(self) -> str:

156 from vllm.platforms import current_platform

tests/compile/test_dynamic_shapes_compilation.py

+1 -1

@@ -99,7 +99,7 @@ def test_dynamic_shapes_compilation(

99 # Clean up GPU memory

100 del model

101 gc.collect()

102 - torch.cuda.empty_cache()

102 + torch.accelerator.empty_cache()

103 torch.cuda.synchronize()

104 print("GPU memory cleared")

105

tests/conftest.py

+1 -1

@@ -1533,7 +1533,7 @@ def clean_gpu_memory_between_tests():

1533

1534 # Clean up GPU memory after the test

1535 if torch.cuda.is_available():

1536 - torch.cuda.empty_cache()

1536 + torch.accelerator.empty_cache()

1537 gc.collect()

1538

1539

tests/entrypoints/openai/test_tensorizer_entrypoint.py

+1 -1

@@ -24,7 +24,7 @@ LORA_PATH = "davzoku/finqa_adapter_1b"

24

25 def _cleanup():

26 gc.collect()

27 - torch.cuda.empty_cache()

27 + torch.accelerator.empty_cache()

28

29

30 @pytest.fixture(autouse=True)

tests/kernels/mamba/test_causal_conv1d.py

+1 -1

@@ -273,7 +273,7 @@ def test_causal_conv1d_varlen(

273 batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype

274 ):

275 device = "cuda"

276 - torch.cuda.empty_cache()

276 + torch.accelerator.empty_cache()

277 rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)

278 if itype == torch.bfloat16:

279 rtol, atol = 1e-2, 5e-2

tests/kernels/moe/test_moe.py

+1 -1

@@ -769,7 +769,7 @@ def test_mixtral_moe(

769 requires_grad=False,

770 )

771 torch.cuda.synchronize()

772 - torch.cuda.empty_cache()

772 + torch.accelerator.empty_cache()

773

774 # FIXME (zyongye) fix this after we move self.kernel

775 # assignment in FusedMoE.__init__

tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py

+3 -3

@@ -178,7 +178,7 @@ def test_load_without_tensorizer_load_format(vllm_runner, capfd, model_ref):

178 finally:

179 del model

180 gc.collect()

181 - torch.cuda.empty_cache()

181 + torch.accelerator.empty_cache()

182

183

184 def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd, model_ref):

@@ -200,7 +200,7 @@ def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd, model_ref)

200 finally:

201 del model

202 gc.collect()

203 - torch.cuda.empty_cache()

203 + torch.accelerator.empty_cache()

204

205

206 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")

@@ -283,7 +283,7 @@ def test_vllm_tensorized_model_has_same_outputs(

283 model_ref, vllm_runner, tmp_path, model_path

284 ):

285 gc.collect()

286 - torch.cuda.empty_cache()

286 + torch.accelerator.empty_cache()

287 config = TensorizerConfig(tensorizer_uri=str(model_path))

288 args = EngineArgs(model=model_ref)

289

tests/test_regression.py

+1 -1

@@ -49,7 +49,7 @@ def test_gc():

49 del llm

50

51 gc.collect()

52 - torch.cuda.empty_cache()

52 + torch.accelerator.empty_cache()

53

54 # The memory allocated for model and KV cache should be released.

55 # The memory allocated for PyTorch and others should be less than 50MB.

tests/v1/e2e/test_async_spec_decode.py

+1 -1

@@ -125,7 +125,7 @@ def test_no_sync_with_spec_decode(

125 assert len(outputs[0].outputs[0].text) > 0

126

127 del llm

128 - torch.cuda.empty_cache()

128 + torch.accelerator.empty_cache()

129 cleanup_dist_env_and_memory()

130

131 sync_tracker.assert_no_sync()

tests/v1/e2e/test_lora_with_spec_decode.py

+2 -2

@@ -95,7 +95,7 @@ def test_batch_inference_correctness(

95 prompts, sampling_params, lora_request=lora_request

96 )

97 del ref_llm

98 - torch.cuda.empty_cache()

98 + torch.accelerator.empty_cache()

99 cleanup_dist_env_and_memory()

100

101 lora_spec_llm = LLM(

@@ -135,5 +135,5 @@ def test_batch_inference_correctness(

135 print(f"match ratio: {matches}/{len(ref_outputs)}")

136 assert matches > int(0.90 * len(ref_outputs))

137 del lora_spec_llm

138 - torch.cuda.empty_cache()

138 + torch.accelerator.empty_cache()

139 cleanup_dist_env_and_memory()

tests/v1/e2e/test_mamba_prefix_cache.py

+2 -2

@@ -440,7 +440,7 @@ def _run_ref_mamba_state_worker():

440 torch.save(cpu_state_ref, "mamba_kv_cache_dict_ref.pth")

441 mamba_kv_cache_dict.clear()

442 del engine

443 - torch.cuda.empty_cache()

443 + torch.accelerator.empty_cache()

444 cleanup_dist_env_and_memory()

445 except Exception:

446 traceback.print_exc()

@@ -805,5 +805,5 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):

805 check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, keys_to_check)

806 mamba_kv_cache_dict.clear()

807 del engine

808 - torch.cuda.empty_cache()

808 + torch.accelerator.empty_cache()

809 cleanup_dist_env_and_memory()

tests/v1/e2e/test_spec_decode.py

+9 -9

@@ -179,7 +179,7 @@ def test_ngram_and_suffix_correctness(

179 )

180 evaluate_llm_for_gsm8k(spec_llm)

181 del spec_llm

182 - torch.cuda.empty_cache()

182 + torch.accelerator.empty_cache()

183 cleanup_dist_env_and_memory()

184

185

@@ -240,7 +240,7 @@ def test_suffix_decoding_acceptance(

240 assert last_accept_rate > 0.80

241

242 del spec_llm

243 - torch.cuda.empty_cache()

243 + torch.accelerator.empty_cache()

244 cleanup_dist_env_and_memory()

245

246

@@ -307,14 +307,14 @@ def test_speculators_model_integration(

307 verifier_model = spec_llm.llm_engine.vllm_config.model_config.model

308

309 del spec_llm

310 - torch.cuda.empty_cache()

310 + torch.accelerator.empty_cache()

311 cleanup_dist_env_and_memory()

312

313 # Second run: Reference without speculative decoding

314 ref_llm = LLM(model=verifier_model, max_model_len=4096)

315 ref_outputs = ref_llm.chat(test_prompts, sampling_config)

316 del ref_llm

317 - torch.cuda.empty_cache()

317 + torch.accelerator.empty_cache()

318 cleanup_dist_env_and_memory()

319

320 # Compare outputs

@@ -410,7 +410,7 @@ def _run_eagle_correctness(

410 )

411 ref_outputs = ref_llm.chat(test_prompts, sampling_config)

412 del ref_llm

413 - torch.cuda.empty_cache()

413 + torch.accelerator.empty_cache()

414 cleanup_dist_env_and_memory()

415

416 spec_llm = LLM(

@@ -445,7 +445,7 @@ def _run_eagle_correctness(

445

446 assert matches > int(0.6 * len(ref_outputs))

447 del spec_llm

448 - torch.cuda.empty_cache()

448 + torch.accelerator.empty_cache()

449 cleanup_dist_env_and_memory()

450

451

@@ -715,7 +715,7 @@ def test_mtp_correctness(

715 ref_llm, expected_accuracy_threshold=expected_accuracy_threshold

716 )

717 del ref_llm

718 - torch.cuda.empty_cache()

718 + torch.accelerator.empty_cache()

719 cleanup_dist_env_and_memory()

720

721 spec_llm = LLM(

@@ -747,7 +747,7 @@ def test_mtp_correctness(

747 # Upon failure, inspect the outputs to check for inaccuracy.

748 assert matches > int(MTP_SIMILARITY_RATE * len(ref_outputs))

749 del spec_llm

750 - torch.cuda.empty_cache()

750 + torch.accelerator.empty_cache()

751 cleanup_dist_env_and_memory()

752

753

@@ -952,7 +952,7 @@ def assert_draft_model_correctness(args: ArgsTest):

952 )

953

954 del spec_llm # CLEANUP

955 - torch.cuda.empty_cache()

955 + torch.accelerator.empty_cache()

956 cleanup_dist_env_and_memory()

957

958 print(

tests/v1/entrypoints/llm/test_struct_output_generate.py

+1 -1

@@ -857,7 +857,7 @@ def test_structured_output_batched_with_non_structured_outputs_requests(

857 # Free memory as soon as possible as failed assertions

858 # will short circuit and not free up memory

859 del llm

860 - torch.cuda.empty_cache()

860 + torch.accelerator.empty_cache()

861 cleanup_dist_env_and_memory()

862

863 for index, output in enumerate(outputs):

tests/v1/sample/test_logprobs.py

+3 -3

@@ -530,7 +530,7 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):

530 assert positive_values > 0

531 finally:

532 del llm

533 - torch.cuda.empty_cache()

533 + torch.accelerator.empty_cache()

534 cleanup_dist_env_and_memory()

535

536

@@ -1065,7 +1065,7 @@ def test_spec_decode_logprobs(

1065 for logprobs in output.logprobs:

1066 ref_logprobs.extend(logprobs.values())

1067 del ref_llm

1068 - torch.cuda.empty_cache()

1068 + torch.accelerator.empty_cache()

1069 cleanup_dist_env_and_memory()

1070

1071 # Run spec decode LLM.

@@ -1095,7 +1095,7 @@ def test_spec_decode_logprobs(

1095 for logprobs in output.logprobs:

1096 spec_logprobs.extend(logprobs.values())

1097 del spec_llm

1098 - torch.cuda.empty_cache()

1098 + torch.accelerator.empty_cache()

1099 cleanup_dist_env_and_memory()

1100

1101 # Per-token logprobs are expected to be the same.

tools/pre_commit/check_torch_cuda.py new

+43

@@ -0,0 +1,43 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +import sys

4 +

5 +import regex as re

6 +

7 +# --------------------------------------------------------------------------- #

8 +# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`

9 +# --------------------------------------------------------------------------- #

10 +_TORCH_CUDA_PATTERNS = [

11 + r"\btorch\.cuda\.empty_cache\b",

12 +]

13 +

14 +ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}

15 +

16 +

17 +def scan_file(path: str) -> int:

18 + with open(path, encoding="utf-8") as f:

19 + content = f.read()

20 + for pattern in _TORCH_CUDA_PATTERNS:

21 + for match in re.finditer(pattern, content, re.MULTILINE):

22 + # Calculate line number from match position

23 + line_num = content[: match.start() + 1].count("\n") + 1

24 + print(

25 + f"{path}:{line_num}: "

26 + "\033[91merror:\033[0m " # red color

27 + "Found torch.cuda API call"

28 + )

29 + return 1

30 + return 0

31 +

32 +

33 +def main():

34 + returncode = 0

35 + for filename in sys.argv[1:]:

36 + if any(filename.startswith(prefix) for prefix in ALLOWED_FILES):

37 + continue

38 + returncode |= scan_file(filename)

39 + return returncode

40 +

41 +

42 +if __name__ == "__main__":

43 + sys.exit(main())

vllm/compilation/cuda_graph.py

+3 -1

@@ -260,7 +260,9 @@ class CUDAGraphWrapper:

260 # therefore, we only run gc for the first graph,

261 # and disable gc for the rest of the graphs.

262 stack.enter_context(patch("gc.collect", lambda: None))

263 - stack.enter_context(patch("torch.cuda.empty_cache", lambda: None))

263 + stack.enter_context(

264 + patch("torch.accelerator.empty_cache", lambda: None)

265 + )

266

267 if self.graph_pool is not None:

268 set_graph_pool_id(self.graph_pool)

vllm/distributed/elastic_ep/elastic_execute.py

+1 -1

@@ -408,7 +408,7 @@ class ElasticEPScalingExecutor:

408

409 gc.collect()

410 torch.cuda.synchronize()

411 - torch.cuda.empty_cache()

411 + torch.accelerator.empty_cache()

412 unlock_workspace()

413 self.worker.compile_or_warm_up_model()

414 lock_workspace()

vllm/distributed/parallel_state.py

+7 -7

@@ -1916,14 +1916,14 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):

1916 gc.collect()

1917 from vllm.platforms import current_platform

1918

1919 - empty_cache = current_platform.empty_cache

1920 - if empty_cache is not None:

1921 - empty_cache()

1922 - try:

1923 - if not current_platform.is_cpu():

1919 + if not current_platform.is_cpu():

1920 + torch.accelerator.empty_cache()

1921 + try:

1922 torch._C._host_emptyCache()

1925 - except AttributeError:

1926 - logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")

1923 + except AttributeError:

1924 + logger.warning(

1925 + "torch._C._host_emptyCache() only available in Pytorch >=2.5"

1926 + )

1927

1928

1929 def in_the_same_node_as(

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+1 -1

@@ -200,7 +200,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

200 ):

201 num_pad = 256 // weight.element_size()

202 weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]

203 - torch.cuda.empty_cache()

203 + torch.accelerator.empty_cache()

204

205 return weight

206

vllm/model_executor/layers/quantization/quark/quark_moe.py

+3 -3

@@ -961,7 +961,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

961

962 # secondly, process mxfp weights

963 if self.emulate:

964 - torch.cuda.empty_cache()

964 + torch.accelerator.empty_cache()

965 return

966

967 from aiter.utility.fp4_utils import e8m0_shuffle

@@ -995,7 +995,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

995 layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)

996 layer.w13_weight.is_shuffled = True

997 layer.w2_weight.is_shuffled = True

998 - torch.cuda.empty_cache()

998 + torch.accelerator.empty_cache()

999

1000 def get_fused_moe_quant_config(

1001 self, layer: torch.nn.Module

@@ -1116,7 +1116,7 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):

1116 del layer.w2_weight

1117 layer.w13_weight = None

1118 layer.w2_weight = None

1119 - torch.cuda.empty_cache()

1119 + torch.accelerator.empty_cache()

1120

1121 if self.static_input_scales:

1122 if layer.w13_input_scale is None or layer.w2_input_scale is None:

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+1 -1

@@ -1407,7 +1407,7 @@ def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor:

1407 import torch.nn.functional as F

1408

1409 weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]

1410 - torch.cuda.empty_cache()

1410 + torch.accelerator.empty_cache()

1411 return weight

1412

1413

vllm/model_executor/model_loader/bitsandbytes_loader.py

+1 -1

@@ -811,7 +811,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

811 **stacked_quant_state_dict,

812 }

813 self._bind_quant_states_to_params(model, stacked_quant_state_dict)

814 - torch.cuda.empty_cache()

814 + torch.accelerator.empty_cache()

815

816 def download_model(self, model_config: ModelConfig) -> None:

817 self._prepare_weights(model_config.model, model_config.revision)

vllm/utils/mem_utils.py

+3 -3

@@ -96,7 +96,7 @@ class MemorySnapshot:

96 # rather than `torch.cuda.memory_reserved()` .

97 # After `torch.cuda.reset_peak_memory_stats()`,

98 # `torch.cuda.memory_reserved()` will keep growing, and only shrink

99 - # when we call `torch.cuda.empty_cache()` or OOM happens.

99 + # when we call `torch.accelerator.empty_cache()` or OOM happens.

100 self.torch_peak = current_platform.memory_stats(device).get(

101 "allocated_bytes.all.peak", 0

102 )

@@ -250,7 +250,7 @@ def memory_profiling(

250 until after profiling to get (c.).

251 """

252 gc.collect()

253 - current_platform.empty_cache()

253 + torch.accelerator.empty_cache()

254 current_platform.reset_peak_memory_stats(baseline_snapshot.device_)

255

256 result = MemoryProfilingResult(

@@ -264,7 +264,7 @@ def memory_profiling(

264 yield result

265

266 gc.collect()

267 - current_platform.empty_cache()

267 + torch.accelerator.empty_cache()

268

269 result.after_profile.measure()

270

vllm/v1/sample/ops/topk_topp_triton.py

+1 -1

@@ -1036,4 +1036,4 @@ def apply_top_k_top_p_triton(

1036 def reset_buffer_cache():

1037 _TRITON_BUFFER_CACHE.clear()

1038 _TRITON_TABLE_CACHE.clear()

1039 - torch.cuda.empty_cache()

1039 + torch.accelerator.empty_cache()

vllm/v1/worker/gpu/model_runner.py

+1 -1

@@ -496,7 +496,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

496

497 start_time = time.perf_counter()

498 gc.collect()

499 - torch.cuda.empty_cache()

499 + torch.accelerator.empty_cache()

500 start_free_gpu_memory = torch.cuda.mem_get_info()[0]

501

502 with self.maybe_setup_dummy_loras(self.lora_config):

vllm/v1/worker/gpu_worker.py

+2 -2

@@ -278,7 +278,7 @@ class Worker(WorkerBase):

278

279 # Now take memory snapshot after NCCL is initialized

280 gc.collect()

281 - torch.cuda.empty_cache()

281 + torch.accelerator.empty_cache()

282

283 # take current memory snapshot

284 self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)

@@ -585,7 +585,7 @@ class Worker(WorkerBase):

585 # sampling related tensors of max possible shape to avoid memory

586 # fragmentation issue.

587 # NOTE: This is called after `capture_model` on purpose to prevent

588 - # memory buffers from being cleared by `torch.cuda.empty_cache`.

588 + # memory buffers from being cleared by `torch.accelerator.empty_cache`.

589 max_num_reqs = min(

590 self.scheduler_config.max_num_seqs,

591 self.scheduler_config.max_num_batched_tokens,

vllm/v1/worker/xpu_model_runner.py

-1

@@ -46,7 +46,6 @@ def _torch_cuda_wrapper():

46 if supports_xpu_graph():

47 torch.cuda.graph = torch.xpu.graph

48 torch.cuda.CUDAGraph = torch.xpu.XPUGraph

49 - torch.cuda.empty_cache = torch.xpu.empty_cache

49 yield

50 finally:

51 pass

vllm/v1/worker/xpu_worker.py

+2 -2

@@ -62,7 +62,7 @@ class XPUWorker(Worker):

62 self.device = torch.device(f"xpu:{self.local_rank}")

63 current_platform.set_device(self.device)

64 current_platform.check_if_supports_dtype(self.model_config.dtype)

65 - torch.xpu.empty_cache()

65 + torch.accelerator.empty_cache()

66 self.init_gpu_memory = torch.xpu.get_device_properties(

67 self.local_rank

68 ).total_memory

@@ -90,7 +90,7 @@ class XPUWorker(Worker):

90

91 # Now take memory snapshot after NCCL is initialized

92 gc.collect()

93 - torch.xpu.empty_cache()

93 + torch.accelerator.empty_cache()

94

95 # take current memory snapshot

96 self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)

[Hardware] Replace `torch.cuda.empty_cache` with `torch.accelerator.empty_cache` (#30681)