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 15, 08:30 UTC
[Misc] Replace `cuda` hard code with `current_platform` (#16983)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
May 23, 2025 at 12:38 UTC
9c1baa5bc6caedabeac1a6da57ec79b41e13056d
3 files changed
+11
-7
vllm/distributed/parallel_state.py
+3
-2
@@ -1221,8 +1221,9 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
1221
ray.shutdown()
1222
gc.collect()
1223
from vllm.platforms import current_platform
1224
- if not current_platform.is_cpu():
1225
- torch.cuda.empty_cache()
1224
+ empty_cache = current_platform.empty_cache
1225
+ if empty_cache is not None:
1226
+ empty_cache()
1227
try:
1228
torch._C._host_emptyCache()
1229
except AttributeError:
vllm/forward_context.py
+4
-1
@@ -120,7 +120,10 @@ def set_forward_context(attn_metadata: Any,
120
# we use synchronous scheduling right now,
121
# adding a sync point here should not affect
122
# scheduling of the next batch
123
- torch.cuda.synchronize()
123
+ from vllm.platforms import current_platform
124
+ synchronize = current_platform.synchronize
125
+ if synchronize is not None:
126
+ synchronize()
127
now = time.perf_counter()
128
# time measurement is in milliseconds
129
batchsize_forward_time[batchsize].append(
vllm/spec_decode/metrics.py
+4
-4
@@ -126,12 +126,12 @@ class AsyncMetricsCollector:
126
"""Copy rejection/typical-acceptance sampling metrics
127
(number of accepted tokens, etc) to CPU asynchronously.
128
129
- Returns a CUDA event recording when the copy is complete.
129
+ Returns a device event recording when the copy is complete.
130
"""
131
assert self._copy_stream is not None
132
- self._copy_stream.wait_stream(torch.cuda.current_stream())
132
+ self._copy_stream.wait_stream(current_platform.current_stream())
133
134
- with torch.cuda.stream(self._copy_stream):
134
+ with current_platform.stream(self._copy_stream):
135
self._aggregate_num_accepted_tokens.copy_(
136
self.spec_decode_sampler.num_accepted_tokens,
137
non_blocking=True)
@@ -142,7 +142,7 @@ class AsyncMetricsCollector:
142
self._aggregate_num_draft_tokens = (
143
self.spec_decode_sampler.num_draft_tokens)
144
145
- aggregate_metrics_ready = torch.cuda.Event()
145
+ aggregate_metrics_ready = current_platform.Event()
146
aggregate_metrics_ready.record(self._copy_stream)
147
148
return aggregate_metrics_ready