66a2209645 · tduncan/Nvidia-Optimization-Stack-VLLM

benchmarks/attention_benchmarks/mla_runner.py

+2 -2

@@ -701,7 +701,7 @@ def _run_single_benchmark(

701 # Warmup

702 for _ in range(config.warmup_iters):

703 forward_fn()

704 - torch.cuda.synchronize()

704 + torch.accelerator.synchronize()

705

706 # Benchmark

707 times = []

@@ -714,7 +714,7 @@ def _run_single_benchmark(

714 forward_fn()

715 end.record()

716

717 - torch.cuda.synchronize()

717 + torch.accelerator.synchronize()

718 elapsed_ms = start.elapsed_time(end)

719 times.append(elapsed_ms / 1000.0 / config.num_layers)

720

benchmarks/attention_benchmarks/runner.py

+2 -2

@@ -391,7 +391,7 @@ def _run_single_benchmark(

391 attn_metadata,

392 output=out,

393 )

394 - torch.cuda.synchronize()

394 + torch.accelerator.synchronize()

395

396 # Benchmark

397 times = []

@@ -412,7 +412,7 @@ def _run_single_benchmark(

412 )

413 end.record()

414

415 - torch.cuda.synchronize()

415 + torch.accelerator.synchronize()

416 elapsed_ms = start.elapsed_time(end)

417 times.append(elapsed_ms / 1000.0 / config.num_layers) # seconds per layer

418

benchmarks/benchmark_topk_topp.py

+3 -3

@@ -94,7 +94,7 @@ def create_logits(

94

95 def measure_memory() -> tuple[int, int]:

96 """Return (allocated, reserved) memory in bytes."""

97 - torch.cuda.synchronize()

97 + torch.accelerator.synchronize()

98 return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()

99

100

@@ -123,7 +123,7 @@ def benchmark_function(

123 for _ in range(warmup_iters):

124 logits_copy = logits.clone()

125 func(logits_copy, k, p)

126 - torch.cuda.synchronize()

126 + torch.accelerator.synchronize()

127

128 # Reset memory stats before benchmark

129 reset_memory_stats()

@@ -140,7 +140,7 @@ def benchmark_function(

140 func(logits_copy, k, p)

141 end_events[i].record()

142

143 - torch.cuda.synchronize()

143 + torch.accelerator.synchronize()

144

145 # Calculate timing

146 times = [

benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py

+1 -1

@@ -168,7 +168,7 @@ def bench_impl(

168 # warmup

169 for kwargs in kwargs_list:

170 impl_type.get_impl()(**kwargs)

171 - torch.cuda.synchronize()

171 + torch.accelerator.synchronize()

172

173 # Merge into a single kwargs and qualify arguments as ArgPool

174 kwargs = {k: ArgPool([]) for k in kwargs_list[0]}

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

+4 -4

@@ -171,7 +171,7 @@ def bench_run(

171 activation=MoEActivation.SILU,

172 global_num_experts=num_experts,

173 )

174 - torch.cuda.synchronize()

174 + torch.accelerator.synchronize()

175

176 # Create CUDA graphs for Triton (match benchmark_moe.py pattern exactly)

177 triton_stream = torch.cuda.Stream()

@@ -187,14 +187,14 @@ def bench_run(

187 topk_ids,

188 quant_config=quant_config,

189 )

190 - torch.cuda.synchronize()

190 + torch.accelerator.synchronize()

191

192 def bench_cuda_graph(graph, num_warmup=5, num_iters=100):

193 """Benchmark CUDA graph using events like benchmark_moe.py"""

194 # Warmup

195 for _ in range(num_warmup):

196 graph.replay()

197 - torch.cuda.synchronize()

197 + torch.accelerator.synchronize()

198

199 # Timing

200 start_event = torch.Event(enable_timing=True)

@@ -202,7 +202,7 @@ def bench_run(

202

203 latencies = []

204 for _ in range(num_iters):

205 - torch.cuda.synchronize()

205 + torch.accelerator.synchronize()

206 start_event.record()

207 graph.replay()

208 end_event.record()

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

+3 -3

@@ -307,7 +307,7 @@ def bench_run(

307 def replay_graph(graph, num_repeats):

308 for _ in range(num_repeats):

309 graph.replay()

310 - torch.cuda.synchronize()

310 + torch.accelerator.synchronize()

311

312 cutlass_stream = torch.cuda.Stream()

313 cutlass_graph = torch.cuda.CUDAGraph()

@@ -330,7 +330,7 @@ def bench_run(

330 e=num_experts,

331 device=device,

332 )

333 - torch.cuda.synchronize()

333 + torch.accelerator.synchronize()

334

335 triton_stream = torch.cuda.Stream()

336 triton_graph = torch.cuda.CUDAGraph()

@@ -345,7 +345,7 @@ def bench_run(

345 w2_fp8scale,

346 a_fp8_scale,

347 )

348 - torch.cuda.synchronize()

348 + torch.accelerator.synchronize()

349

350 min_run_time = 5

351 num_warmup = 5

benchmarks/kernels/benchmark_device_communicators.py

+5 -5

@@ -342,7 +342,7 @@ class CommunicatorBenchmark:

342 if not should_use_fn(tensor):

343 return None

344

345 - torch.cuda.synchronize()

345 + torch.accelerator.synchronize()

346 stream = torch.cuda.Stream()

347 with torch.cuda.stream(stream):

348 graph_input = tensor.clone()

@@ -360,17 +360,17 @@ class CommunicatorBenchmark:

360 for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):

361 allreduce_fn(graph_input)

362

363 - torch.cuda.synchronize()

363 + torch.accelerator.synchronize()

364 for _ in range(num_warmup):

365 graph.replay()

366 - torch.cuda.synchronize()

366 + torch.accelerator.synchronize()

367

368 - torch.cuda.synchronize()

368 + torch.accelerator.synchronize()

369 start_time = time.perf_counter()

370

371 for _ in range(num_trials):

372 graph.replay()

373 - torch.cuda.synchronize()

373 + torch.accelerator.synchronize()

374

375 end_time = time.perf_counter()

376

benchmarks/kernels/benchmark_fused_collective.py

+4 -4

@@ -385,7 +385,7 @@ def benchmark_operation(

385 # Warmup before graph capture

386 for _ in range(warmup):

387 operation_func(*args, **kwargs)

388 - torch.cuda.synchronize()

388 + torch.accelerator.synchronize()

389

390 # Create CUDA graph

391 graph = torch.cuda.CUDAGraph()

@@ -398,19 +398,19 @@ def benchmark_operation(

398 operation_func(*args, **kwargs)

399

400 # Graph warmup

401 - torch.cuda.synchronize()

401 + torch.accelerator.synchronize()

402 for _ in range(warmup):

403 graph.replay()

404

405 # Benchmark with CUDA graph

406 - torch.cuda.synchronize()

406 + torch.accelerator.synchronize()

407 start_time = time.perf_counter()

408

409 for _ in range(trials // num_op_per_cudagraph):

410 # operation_func(*args, **kwargs)

411 graph.replay()

412

413 - torch.cuda.synchronize()

413 + torch.accelerator.synchronize()

414 end_time = time.perf_counter()

415

416 avg_time_ms = ((end_time - start_time) / trials) * 1000

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

+3 -3

@@ -224,7 +224,7 @@ def bench_run(

224 def replay_graph(graph, num_repeats):

225 for _ in range(num_repeats):

226 graph.replay()

227 - torch.cuda.synchronize()

227 + torch.accelerator.synchronize()

228

229 cutlass_stream = torch.cuda.Stream()

230 cutlass_graph = torch.cuda.CUDAGraph()

@@ -239,7 +239,7 @@ def bench_run(

239 topk_weights,

240 topk_ids,

241 )

242 - torch.cuda.synchronize()

242 + torch.accelerator.synchronize()

243

244 triton_stream = torch.cuda.Stream()

245 triton_graph = torch.cuda.CUDAGraph()

@@ -254,7 +254,7 @@ def bench_run(

254 w2_scale,

255 a_scale,

256 )

257 - torch.cuda.synchronize()

257 + torch.accelerator.synchronize()

258

259 min_run_time = 5

260 num_warmup = 5

benchmarks/kernels/benchmark_layernorm.py

+2 -2

@@ -34,14 +34,14 @@ def main(

34 residual = torch.randn_like(x) * scale if add_residual else None

35

36 def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:

37 - torch.cuda.synchronize()

37 + torch.accelerator.synchronize()

38 if profile:

39 torch.cuda.cudart().cudaProfilerStart()

40 start_time = time.perf_counter()

41

42 for _ in range(num_iters):

43 layer(x, residual)

44 - torch.cuda.synchronize()

44 + torch.accelerator.synchronize()

45

46 end_time = time.perf_counter()

47 if profile:

benchmarks/kernels/benchmark_lora.py

+1 -1

@@ -1035,7 +1035,7 @@ def bench_optype(

1035 # Run bench function so that _LORA_A_PTR_DICT and _LORA_B_PTR_DICT are set up

1036 for kwargs in kwargs_list:

1037 op_type.bench_fn()(**kwargs)

1038 - torch.cuda.synchronize()

1038 + torch.accelerator.synchronize()

1039

1040 # Merge into a single kwargs and qualify arguments as ArgPool

1041 kwargs = {k: ArgPool([]) for k in kwargs_list[0]}

benchmarks/kernels/benchmark_mla_k_concat.py

+2 -2

@@ -47,13 +47,13 @@ def benchmark_method(

47 # Warmup

48 for _ in range(num_warmup):

49 _ = method(k_nope, k_pe)

50 - torch.cuda.synchronize()

50 + torch.accelerator.synchronize()

51

52 # Benchmark

53 start = time.perf_counter()

54 for _ in range(num_iters):

55 _ = method(k_nope, k_pe)

56 - torch.cuda.synchronize()

56 + torch.accelerator.synchronize()

57 end = time.perf_counter()

58

59 return (end - start) / num_iters * 1000 # Convert to ms

benchmarks/kernels/benchmark_moe.py

+4 -4

@@ -304,19 +304,19 @@ def benchmark_config(

304

305 # JIT compilation & warmup

306 run()

307 - torch.cuda.synchronize()

307 + torch.accelerator.synchronize()

308

309 # Capture 10 invocations with CUDA graph

310 graph = torch.cuda.CUDAGraph()

311 with torch.cuda.graph(graph):

312 for _ in range(10):

313 run()

314 - torch.cuda.synchronize()

314 + torch.accelerator.synchronize()

315

316 # Warmup

317 for _ in range(5):

318 graph.replay()

319 - torch.cuda.synchronize()

319 + torch.accelerator.synchronize()

320

321 start_event = torch.Event(enable_timing=True)

322 end_event = torch.Event(enable_timing=True)

@@ -324,7 +324,7 @@ def benchmark_config(

324 latencies: list[float] = []

325 for i in range(num_iters):

326 prepare(i)

327 - torch.cuda.synchronize()

327 + torch.accelerator.synchronize()

328

329 start_event.record()

330 graph.replay()

benchmarks/kernels/benchmark_moe_defaults.py

+2 -2

@@ -131,7 +131,7 @@ def benchmark_config(

131 topk_ids,

132 quant_config=quant_config,

133 )

134 - torch.cuda.synchronize()

134 + torch.accelerator.synchronize()

135

136 # Benchmark

137 start = torch.cuda.Event(enable_timing=True)

@@ -149,7 +149,7 @@ def benchmark_config(

149 quant_config=quant_config,

150 )

151 end.record()

152 - torch.cuda.synchronize()

152 + torch.accelerator.synchronize()

153 return start.elapsed_time(end) / num_iters * 1000 # ms -> us

154

155

benchmarks/kernels/benchmark_moe_permute_unpermute.py

+8 -8

@@ -69,19 +69,19 @@ def benchmark_permute(

69

70 # JIT compilation & warmup

71 run()

72 - torch.cuda.synchronize()

72 + torch.accelerator.synchronize()

73

74 # Capture 10 invocations with CUDA graph

75 graph = torch.cuda.CUDAGraph()

76 with torch.cuda.graph(graph):

77 for _ in range(10):

78 run()

79 - torch.cuda.synchronize()

79 + torch.accelerator.synchronize()

80

81 # Warmup

82 for _ in range(5):

83 graph.replay()

84 - torch.cuda.synchronize()

84 + torch.accelerator.synchronize()

85

86 start_event = torch.Event(enable_timing=True)

87 end_event = torch.Event(enable_timing=True)

@@ -89,7 +89,7 @@ def benchmark_permute(

89 latencies: list[float] = []

90 for i in range(num_iters):

91 prepare(i)

92 - torch.cuda.synchronize()

92 + torch.accelerator.synchronize()

93

94 start_event.record()

95 graph.replay()

@@ -159,26 +159,26 @@ def benchmark_unpermute(

159 # JIT compilation & warmup

160 input = prepare()

161 run(input)

162 - torch.cuda.synchronize()

162 + torch.accelerator.synchronize()

163

164 # Capture 10 invocations with CUDA graph

165 graph = torch.cuda.CUDAGraph()

166 with torch.cuda.graph(graph):

167 for _ in range(10):

168 run(input)

169 - torch.cuda.synchronize()

169 + torch.accelerator.synchronize()

170

171 # Warmup

172 for _ in range(5):

173 graph.replay()

174 - torch.cuda.synchronize()

174 + torch.accelerator.synchronize()

175

176 start_event = torch.Event(enable_timing=True)

177 end_event = torch.Event(enable_timing=True)

178

179 latencies: list[float] = []

180 for i in range(num_iters):

181 - torch.cuda.synchronize()

181 + torch.accelerator.synchronize()

182 start_event.record()

183 graph.replay()

184 end_event.record()

benchmarks/kernels/benchmark_mrope.py

+5 -5

@@ -135,14 +135,14 @@ def benchmark_mrope(

135 key.clone(),

136 )

137

138 - torch.cuda.synchronize()

138 + torch.accelerator.synchronize()

139

140 # Time reference implementation

141 torch_times = []

142 for _ in range(benchmark_iter):

143 query_clone = query.clone()

144 key_clone = key.clone()

145 - torch.cuda.synchronize()

145 + torch.accelerator.synchronize()

146 start_time = time.time()

147

148 mrope_helper_class.forward_native(

@@ -151,7 +151,7 @@ def benchmark_mrope(

151 key_clone,

152 )

153

154 - torch.cuda.synchronize()

154 + torch.accelerator.synchronize()

155 torch_times.append(time.time() - start_time)

156

157 # Time triton kernel implementation

@@ -159,14 +159,14 @@ def benchmark_mrope(

159 for _ in range(benchmark_iter):

160 query_clone = query.clone()

161 key_clone = key.clone()

162 - torch.cuda.synchronize()

162 + torch.accelerator.synchronize()

163 start_time = time.time()

164 mrope_helper_class.forward_cuda(

165 positions,

166 query_clone,

167 key_clone,

168 )

169 - torch.cuda.synchronize()

169 + torch.accelerator.synchronize()

170 triton_times.append(time.time() - start_time)

171

172 # Calculate statistics

benchmarks/kernels/benchmark_paged_attention.py

+2 -2

@@ -103,7 +103,7 @@ def main(

103 max_logits = torch.empty_like(exp_sums)

104

105 def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:

106 - torch.cuda.synchronize()

106 + torch.accelerator.synchronize()

107 if profile:

108 torch.cuda.cudart().cudaProfilerStart()

109 start_time = time.perf_counter()

@@ -173,7 +173,7 @@ def main(

173 )

174 else:

175 raise ValueError(f"Invalid version: {version}")

176 - torch.cuda.synchronize()

176 + torch.accelerator.synchronize()

177

178 end_time = time.perf_counter()

179 if profile:

benchmarks/kernels/benchmark_per_token_group_quant.py

+2 -2

@@ -28,7 +28,7 @@ def _time_cuda(

28 # warmup

29 for _ in range(warmup_iters):

30 fn()

31 - torch.cuda.synchronize()

31 + torch.accelerator.synchronize()

32

33 start = torch.Event(enable_timing=True)

34 end = torch.Event(enable_timing=True)

@@ -37,7 +37,7 @@ def _time_cuda(

37 for _ in range(bench_iters):

38 fn()

39 end.record()

40 - torch.cuda.synchronize()

40 + torch.accelerator.synchronize()

41

42 return start.elapsed_time(end) / bench_iters # ms/iter

43

benchmarks/kernels/benchmark_quant.py

+2 -2

@@ -29,7 +29,7 @@ def main(

29 scale = torch.randn(1, 1, dtype=torch.float32) if static_scale else None

30

31 def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:

32 - torch.cuda.synchronize()

32 + torch.accelerator.synchronize()

33 if profile:

34 torch.cuda.cudart().cudaProfilerStart()

35 start_time = time.perf_counter()

@@ -39,7 +39,7 @@ def main(

39 ops.scaled_int8_quant(x, scale)

40 else:

41 ops.scaled_fp8_quant(x, scale)

42 - torch.cuda.synchronize()

42 + torch.accelerator.synchronize()

43

44 end_time = time.perf_counter()

45 if profile:

benchmarks/kernels/benchmark_reshape_and_cache.py

+3 -3

@@ -84,16 +84,16 @@ def run_benchmark(

84 g = torch.cuda.CUDAGraph()

85 with torch.cuda.graph(g):

86 function_under_test()

87 - torch.cuda.synchronize()

87 + torch.accelerator.synchronize()

88 function_under_test = lambda: g.replay()

89

90 def run_cuda_benchmark(n_iters: int) -> float:

91 nonlocal key, value, key_cache, value_cache, slot_mapping

92 - torch.cuda.synchronize()

92 + torch.accelerator.synchronize()

93 start = time.perf_counter()

94 for _ in range(n_iters):

95 function_under_test()

96 - torch.cuda.synchronize()

96 + torch.accelerator.synchronize()

97 end = time.perf_counter()

98 return (end - start) / n_iters

99

benchmarks/kernels/benchmark_reshape_and_cache_flash.py

+3 -3

@@ -109,16 +109,16 @@ def run_benchmark(

109 g = torch.cuda.CUDAGraph()

110 with torch.cuda.graph(g):

111 function_under_test()

112 - torch.cuda.synchronize()

112 + torch.accelerator.synchronize()

113 function_under_test = lambda: g.replay()

114

115 def run_cuda_benchmark(n_iters: int) -> float:

116 nonlocal key, value, key_cache, value_cache, slot_mapping

117 - torch.cuda.synchronize()

117 + torch.accelerator.synchronize()

118 start = time.perf_counter()

119 for _ in range(n_iters):

120 function_under_test()

121 - torch.cuda.synchronize()

121 + torch.accelerator.synchronize()

122 end = time.perf_counter()

123 return (end - start) / n_iters

124

benchmarks/kernels/benchmark_silu_mul_fp8_quant.py

+2 -2

@@ -251,7 +251,7 @@ def benchmark(

251 kernel(

252 y, tokens_per_expert, num_parallel_tokens=num_parallel_tokens, group_size=G

253 )

254 - torch.cuda.synchronize()

254 + torch.accelerator.synchronize()

255

256 start_event = torch.Event(enable_timing=True)

257 end_event = torch.Event(enable_timing=True)

@@ -259,7 +259,7 @@ def benchmark(

259 # Benchmark

260 latencies: list[float] = []

261 for _ in range(runs):

262 - torch.cuda.synchronize()

262 + torch.accelerator.synchronize()

263

264 start_event.record()

265 for i in range(iterations_per_run):

benchmarks/kernels/benchmark_trtllm_decode_attention.py

+2 -2

@@ -126,7 +126,7 @@ def benchmark_decode(

126 )

127

128 def time_fn(fn, warmup=10, trials=20):

129 - torch.cuda.synchronize()

129 + torch.accelerator.synchronize()

130 start = torch.Event(enable_timing=True)

131 end = torch.Event(enable_timing=True)

132 times = []

@@ -136,7 +136,7 @@ def benchmark_decode(

136 start.record()

137 fn()

138 end.record()

139 - torch.cuda.synchronize()

139 + torch.accelerator.synchronize()

140 times.append(start.elapsed_time(end)) # ms

141 return sum(times) / len(times), torch.std(torch.tensor(times))

142

benchmarks/kernels/benchmark_trtllm_prefill_attention.py

+2 -2

@@ -138,7 +138,7 @@ def benchmark_prefill(

138 )

139

140 def time_fn(fn, warmup=10, trials=20):

141 - torch.cuda.synchronize()

141 + torch.accelerator.synchronize()

142 start = torch.Event(enable_timing=True)

143 end = torch.Event(enable_timing=True)

144 times = []

@@ -148,7 +148,7 @@ def benchmark_prefill(

148 start.record()

149 fn()

150 end.record()

151 - torch.cuda.synchronize()

151 + torch.accelerator.synchronize()

152 times.append(start.elapsed_time(end)) # ms

153 return sum(times) / len(times), torch.std(torch.tensor(times))

154

benchmarks/kernels/benchmark_w8a8_block_fp8.py

+3 -3

@@ -177,18 +177,18 @@ def benchmark_config(

177 def run():

178 w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype)

179

180 - torch.cuda.synchronize()

180 + torch.accelerator.synchronize()

181 # JIT complication & warmup

182 for _ in range(5):

183 run()

184 - torch.cuda.synchronize()

184 + torch.accelerator.synchronize()

185

186 start_event = torch.Event(enable_timing=True)

187 end_event = torch.Event(enable_timing=True)

188

189 latencies: list[float] = []

190 for i in range(num_iters):

191 - torch.cuda.synchronize()

191 + torch.accelerator.synchronize()

192 start_event.record()

193 run()

194 end_event.record()

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

+4 -4

@@ -35,7 +35,7 @@ def benchmark_shape(

35 B = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)

36

37 # Reference result in BF16

38 - torch.cuda.synchronize()

38 + torch.accelerator.synchronize()

39 C_ref = A @ B.t()

40

41 # Pre-quantize B for all implementations

@@ -121,14 +121,14 @@ def benchmark_shape(

121 # Warmup

122 for _ in range(warmup):

123 func()

124 - torch.cuda.synchronize()

124 + torch.accelerator.synchronize()

125

126 # Timing loop

127 - torch.cuda.synchronize()

127 + torch.accelerator.synchronize()

128 start = time.time()

129 for _ in range(repeat):

130 func()

131 - torch.cuda.synchronize()

131 + torch.accelerator.synchronize()

132 end = time.time()

133

134 # Calculate timing and TFLOPS

docs/design/model_runner_v2.md

+1 -1

@@ -50,7 +50,7 @@ V1 was not originally designed with async scheduling in mind, and support requir

50

51 ## 3. Removing Async Barrier

52

53 -A key requirement for async execution is that CPU operations remain non-blocking. Both explicit sync (for example, `torch.cuda.synchronize`) and implicit sync (for example, unpinned `.to("cuda")`) must be avoided.

53 +A key requirement for async execution is that CPU operations remain non-blocking. Both explicit sync (for example, `torch.accelerator.synchronize`) and implicit sync (for example, unpinned `.to("cuda")`) must be avoided.

54

55 However, async execution can introduce race conditions when CPU and GPU concurrently touch the same memory.

56

docs/usage/troubleshooting.md

+1 -1

@@ -95,7 +95,7 @@ If GPU/CPU communication cannot be established, you can use the following Python

95 torch.cuda.set_device(local_rank)

96 data = torch.FloatTensor([1,] * 128).to("cuda")

97 dist.all_reduce(data, op=dist.ReduceOp.SUM)

98 - torch.cuda.synchronize()

98 + torch.accelerator.synchronize()

99 value = data.mean().item()

100 world_size = dist.get_world_size()

101 assert value == world_size, f"Expected {world_size}, got {value}"

examples/offline_inference/rlhf_colocate.py

+2 -2

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

88 # Zero out all the parameters.

89 for name, p in self.model.named_parameters():

90 p.data.zero_()

91 - torch.cuda.synchronize()

91 + torch.accelerator.synchronize()

92 # The argument for `get_device_uuid` is the index of the GPU in the

93 # list of visible devices.

94 from vllm.platforms import current_platform

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

151 p.data.view(-1).view(dtype=torch.uint8), non_blocking=True

152 )

153 offset += get_size(p)

154 - torch.cuda.synchronize()

154 + torch.accelerator.synchronize()

155 s.send_pyobj(named_tensors)

156 s.recv()

157 s.send_pyobj(None)

examples/offline_inference/rlhf_utils.py

+2 -2

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

120 process_weights_after_loading(

121 self.model_runner.model, self.model_config, self.device

122 )

123 - torch.cuda.synchronize()

123 + torch.accelerator.synchronize()

124 socket.send(b"")

125 break

126 if isinstance(payload, tuple):

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

144 weights.append((item["name"], tensor))

145 self.model_runner.model.load_weights(weights=weights)

146 del weights

147 - torch.cuda.synchronize()

147 + torch.accelerator.synchronize()

148 socket.send(b"")

149

150 socket.close()

tests/compile/test_dynamic_shapes_compilation.py

+1 -1

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

100 del model

101 gc.collect()

102 torch.accelerator.empty_cache()

103 - torch.cuda.synchronize()

103 + torch.accelerator.synchronize()

104 print("GPU memory cleared")

105

106

tests/distributed/test_ca_buffer_sharing.py

+3 -3

@@ -32,7 +32,7 @@ pointers = CustomAllreduce.create_shared_buffer(buffer_size_in_bytes)

32 print(f"Rank {rank} has pointers {pointers}")

33

34 dist.barrier()

35 -torch.cuda.synchronize()

35 +torch.accelerator.synchronize()

36

37 if rank == 0:

38 # the first rank tries to write to all buffers

@@ -41,7 +41,7 @@ if rank == 0:

41 lib.cudaMemset(pointer, byte_value, buffer_size_in_bytes)

42

43 dist.barrier()

44 -torch.cuda.synchronize()

44 +torch.accelerator.synchronize()

45

46 host_data = (ctypes.c_char * buffer_size_in_bytes)()

47

@@ -59,6 +59,6 @@ for p in pointers:

59 print(f"Rank {rank} verified all buffers")

60

61 dist.barrier()

62 -torch.cuda.synchronize()

62 +torch.accelerator.synchronize()

63

64 CustomAllreduce.free_shared_buffer(pointers)

tests/distributed/test_custom_all_reduce.py

+2 -2

@@ -48,7 +48,7 @@ def graph_allreduce(

48 data = torch.zeros(1)

49 data = data.to(device=device)

50 torch.distributed.all_reduce(data, group=group)

51 - torch.cuda.synchronize()

51 + torch.accelerator.synchronize()

52 del data

53

54 # we use the first group to communicate once

@@ -68,7 +68,7 @@ def graph_allreduce(

68 inp2 = torch.randint(

69 1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()

70 )

71 - torch.cuda.synchronize()

71 + torch.accelerator.synchronize()

72 graph = torch.cuda.CUDAGraph()

73 with torch.cuda.graph(graph, stream=graph_capture_context.stream):

74 for i in range(num_communication):

tests/distributed/test_pynccl.py

+15 -15

@@ -68,7 +68,7 @@ def worker_fn():

68 )

69 tensor = torch.ones(16, 1024, 1024, dtype=torch.float32).cuda(pynccl_comm.rank)

70 tensor = pynccl_comm.all_reduce(tensor)

71 - torch.cuda.synchronize()

71 + torch.accelerator.synchronize()

72 assert torch.all(tensor == pynccl_comm.world_size).cpu().item()

73

74

@@ -93,11 +93,11 @@ def multiple_allreduce_worker_fn():

93 if torch.distributed.get_rank() in [0, 1]:

94 tensor = pynccl_comm.all_reduce(tensor)

95 tensor = pynccl_comm.all_reduce(tensor)

96 - torch.cuda.synchronize()

96 + torch.accelerator.synchronize()

97 assert torch.all(tensor == 4).cpu().item()

98 else:

99 tensor = pynccl_comm.all_reduce(tensor)

100 - torch.cuda.synchronize()

100 + torch.accelerator.synchronize()

101 assert torch.all(tensor == 2).cpu().item()

102

103

@@ -121,11 +121,11 @@ def multiple_allreduce_with_vllm_worker_fn():

121 if torch.distributed.get_rank() in [0, 1]:

122 tensor = tensor_model_parallel_all_reduce(tensor)

123 tensor = tensor_model_parallel_all_reduce(tensor)

124 - torch.cuda.synchronize()

124 + torch.accelerator.synchronize()

125 assert torch.all(tensor == 4).cpu().item()

126 else:

127 tensor = tensor_model_parallel_all_reduce(tensor)

128 - torch.cuda.synchronize()

128 + torch.accelerator.synchronize()

129 assert torch.all(tensor == 2).cpu().item()

130

131

@@ -147,12 +147,12 @@ def worker_fn_with_cudagraph():

147 )

148 # run something in the default stream to initialize torch engine

149 a = torch.ones((4, 4), device=f"cuda:{pynccl_comm.rank}")

150 - torch.cuda.synchronize()

150 + torch.accelerator.synchronize()

151 with torch.cuda.graph(graph):

152 a_out = pynccl_comm.all_reduce(a)

153 - torch.cuda.synchronize()

153 + torch.accelerator.synchronize()

154 graph.replay()

155 - torch.cuda.synchronize()

155 + torch.accelerator.synchronize()

156 assert torch.all(a_out == pynccl_comm.world_size).cpu().item()

157

158

@@ -180,7 +180,7 @@ def all_gather_worker_fn():

180 ).to(device)

181

182 pynccl_comm.all_gather(result, tensor)

183 - torch.cuda.synchronize()

183 + torch.accelerator.synchronize()

184 torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)

185

186

@@ -215,7 +215,7 @@ def all_gatherv_worker_fn():

215 ).to(device)

216

217 pynccl_comm.all_gatherv(result, tensor, sizes=sizes)

218 - torch.cuda.synchronize()

218 + torch.accelerator.synchronize()

219 torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)

220

221

@@ -255,7 +255,7 @@ def reduce_scatter_worker_fn():

255 ).to(device)

256

257 pynccl_comm.reduce_scatter(result, tensor)

258 - torch.cuda.synchronize()

258 + torch.accelerator.synchronize()

259 torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)

260

261

@@ -293,7 +293,7 @@ def reduce_scatterv_worker_fn():

293 expected = sum(tensor[start:end] for tensor in all_tensors).to(device)

294

295 pynccl_comm.reduce_scatterv(result, tensor, sizes=sizes)

296 - torch.cuda.synchronize()

296 + torch.accelerator.synchronize()

297 torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)

298

299

@@ -325,7 +325,7 @@ def send_recv_worker_fn():

325 pynccl_comm.send(tensor, dst=(pynccl_comm.rank + 1) % pynccl_comm.world_size)

326 else:

327 pynccl_comm.recv(tensor, src=(pynccl_comm.rank - 1) % pynccl_comm.world_size)

328 - torch.cuda.synchronize()

328 + torch.accelerator.synchronize()

329 assert torch.all(tensor == 1).cpu().item()

330

331

@@ -355,7 +355,7 @@ def multiple_send_recv_worker_fn():

355 pynccl_comm.send(tensor, dst=(pynccl_comm.rank + 1) % pynccl_comm.world_size)

356 else:

357 pynccl_comm.recv(tensor, src=(pynccl_comm.rank - 1) % pynccl_comm.world_size)

358 - torch.cuda.synchronize()

358 + torch.accelerator.synchronize()

359 if torch.distributed.get_rank() in [0, 2]:

360 assert torch.all(tensor == 1).cpu().item()

361 else:

@@ -396,7 +396,7 @@ def broadcast_worker_fn():

396 pynccl_comm.broadcast(recv_tensors[i], src=i)

397 # the broadcast op might be launched in a different stream

398 # need to synchronize to make sure the tensor is ready

399 - torch.cuda.synchronize()

399 + torch.accelerator.synchronize()

400 assert torch.all(recv_tensors[i] == i).cpu().item()

401

402

tests/distributed/test_quick_all_reduce.py

+2 -2

@@ -52,7 +52,7 @@ def graph_quickreduce(

52 data = torch.zeros(1)

53 data = data.to(device=device)

54 torch.distributed.all_reduce(data, group=group)

55 - torch.cuda.synchronize()

55 + torch.accelerator.synchronize()

56 del data

57

58 # we use the first group to communicate once

@@ -71,7 +71,7 @@ def graph_quickreduce(

71 inp2 = torch.randint(

72 -23, 1, (sz,), dtype=dtype, device=torch.cuda.current_device()

73 )

74 - torch.cuda.synchronize()

74 + torch.accelerator.synchronize()

75 graph = torch.cuda.CUDAGraph()

76 with torch.cuda.graph(graph, stream=graph_capture_context.stream):

77 for _ in range(num_communication):

tests/distributed/test_utils.py

+2 -2

@@ -79,11 +79,11 @@ def gpu_worker(rank, WORLD_SIZE, port1, port2):

79 data = torch.tensor([rank]).cuda()

80 pynccl1.all_reduce(data)

81 pg1.barrier()

82 - torch.cuda.synchronize()

82 + torch.accelerator.synchronize()

83 if rank <= 2:

84 pynccl2.all_reduce(data)

85 pg2.barrier()

86 - torch.cuda.synchronize()

86 + torch.accelerator.synchronize()

87 item = data[0].item()

88 print(f"rank: {rank}, item: {item}")

89 if rank == 3:

tests/distributed/test_weight_transfer.py

+4 -4

@@ -251,7 +251,7 @@ def trainer_broadcast_tensor(

251 dtype = getattr(torch, tensor_dtype)

252 tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device="cuda:0")

253 comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream())

254 - torch.cuda.synchronize()

254 + torch.accelerator.synchronize()

255

256 return True

257

@@ -309,7 +309,7 @@ def inference_receive_tensor(

309 shapes=[tensor_shape],

310 )

311 engine.receive_weights(update_info, noop_load_weights)

312 - torch.cuda.synchronize()

312 + torch.accelerator.synchronize()

313

314 # Verify we received the tensor

315 success = False

@@ -630,7 +630,7 @@ class TrainerActor:

630 ipc_handle = reduce_tensor(self.tensor)

631 gpu_uuid = get_physical_gpu_id(0)

632

633 - torch.cuda.synchronize()

633 + torch.accelerator.synchronize()

634

635 self.ipc_handle_dict = {

636 "ipc_handle": ipc_handle,

@@ -704,7 +704,7 @@ def inference_receive_ipc_tensor(

704

705 update_info = engine.parse_update_info(update_dict)

706 engine.receive_weights(update_info, noop_load_weights)

707 - torch.cuda.synchronize()

707 + torch.accelerator.synchronize()

708

709 # Verify we received the tensor

710 success = False

tests/kernels/attention/test_merge_attn_states.py

+6 -6

@@ -165,7 +165,7 @@ def test_merge_attn_states(

165 suffix_lse_torch,

166 output_lse_torch,

167 )

168 - torch.cuda.synchronize()

168 + torch.accelerator.synchronize()

169

170 for _ in range(repeat_times):

171 start.record()

@@ -178,7 +178,7 @@ def test_merge_attn_states(

178 output_lse_torch,

179 )

180 end.record()

181 - torch.cuda.synchronize()

181 + torch.accelerator.synchronize()

182 total_time_torch_kernel += start.elapsed_time(end)

183

184 avg_time_torch_kernel = total_time_torch_kernel / repeat_times

@@ -200,7 +200,7 @@ def test_merge_attn_states(

200 suffix_lse,

201 output_lse_ref_triton,

202 )

203 - torch.cuda.synchronize()

203 + torch.accelerator.synchronize()

204

205 for _ in range(repeat_times):

206 start.record()

@@ -213,7 +213,7 @@ def test_merge_attn_states(

213 output_lse_ref_triton,

214 )

215 end.record()

216 - torch.cuda.synchronize()

216 + torch.accelerator.synchronize()

217 total_time_triton_kernel += start.elapsed_time(end)

218

219 avg_time_triton_kernel = total_time_triton_kernel / repeat_times

@@ -232,7 +232,7 @@ def test_merge_attn_states(

232 suffix_lse,

233 output_lse_cuda,

234 )

235 - torch.cuda.synchronize()

235 + torch.accelerator.synchronize()

236

237 for _ in range(repeat_times):

238 start.record()

@@ -245,7 +245,7 @@ def test_merge_attn_states(

245 output_lse_cuda,

246 )

247 end.record()

248 - torch.cuda.synchronize()

248 + torch.accelerator.synchronize()

249 total_time_cuda_kernel += start.elapsed_time(end)

250

251 avg_time_cuda_kernel = total_time_cuda_kernel / repeat_times

tests/kernels/attention/test_prefix_prefill.py

+8 -8

@@ -239,7 +239,7 @@ def test_contexted_kv_attention(

239 v_scale,

240 sliding_window=sliding_window,

241 )

242 - torch.cuda.synchronize()

242 + torch.accelerator.synchronize()

243 start_time = time.time()

244 op(

245 query,

@@ -258,7 +258,7 @@ def test_contexted_kv_attention(

258 v_scale,

259 sliding_window=sliding_window,

260 )

261 - torch.cuda.synchronize()

261 + torch.accelerator.synchronize()

262 end_time = time.time()

263 print(f"triton Time: {(end_time - start_time) * 1000:.2f} ms")

264

@@ -298,7 +298,7 @@ def test_contexted_kv_attention(

298 dropout_p=0.0,

299 scale=scale,

300 )

301 - torch.cuda.synchronize()

301 + torch.accelerator.synchronize()

302 start_time = time.time()

303 output_ref = F.scaled_dot_product_attention(

304 query_sdpa,

@@ -308,7 +308,7 @@ def test_contexted_kv_attention(

308 dropout_p=0.0,

309 scale=scale,

310 )

311 - torch.cuda.synchronize()

311 + torch.accelerator.synchronize()

312 end_time = time.time()

313 print(f"PyTorch SDPA Time: {(end_time - start_time) * 1000:.2f} ms")

314

@@ -482,7 +482,7 @@ def test_contexted_kv_attention_alibi(

482 v_scale,

483 alibi_slopes=alibi_slopes,

484 )

485 - torch.cuda.synchronize()

485 + torch.accelerator.synchronize()

486 start_time = time.time()

487 op(

488 query,

@@ -501,7 +501,7 @@ def test_contexted_kv_attention_alibi(

501 v_scale,

502 alibi_slopes=alibi_slopes,

503 )

504 - torch.cuda.synchronize()

504 + torch.accelerator.synchronize()

505 end_time = time.time()

506 print(f"triton Time: {(end_time - start_time) * 1000:.2f} ms")

507 scale = float(1.0 / (head_size**0.5))

@@ -517,7 +517,7 @@ def test_contexted_kv_attention_alibi(

517

518 output_ref = torch.empty_like(output)

519

520 - torch.cuda.synchronize()

520 + torch.accelerator.synchronize()

521 start_time = time.time()

522

523 query_start = 0

@@ -572,7 +572,7 @@ def test_contexted_kv_attention_alibi(

572 query_start = query_end

573 key_start = key_end

574

575 - torch.cuda.synchronize()

575 + torch.accelerator.synchronize()

576 end_time = time.time()

577 print(f"PyTorch SDPA Time: {(end_time - start_time) * 1000:.2f} ms")

578 atol = 1e-3 if "fp8" in kv_cache_dtype else 1e-6

tests/kernels/core/test_layernorm.py

+1 -1

@@ -127,7 +127,7 @@ def test_fused_rms_norm_quant(

127 out_quant, x_unfused.contiguous(), quant_scale_t

128 )

129

130 - torch.cuda.synchronize()

130 + torch.accelerator.synchronize()

131 torch.testing.assert_close(residual_fused, residual, atol=1e-2, rtol=1e-2)

132 opcheck(

133 torch.ops._C.fused_add_rms_norm_static_fp8_quant,

tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py

+1 -1

@@ -34,7 +34,7 @@ def do_profile(

34 record_shapes=True,

35 ) as tprof:

36 fn(**fn_kwargs)

37 - torch.cuda.synchronize(torch.cuda.current_device())

37 + torch.accelerator.synchronize(torch.cuda.current_device())

38

39 # TODO (varun): Add a descriptive trace file name

40 tprof.export_chrome_trace(

tests/kernels/moe/test_block_fp8.py

+2 -2

@@ -318,8 +318,8 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)

318 out = deep_gemm_moe_fp8_fn(

319 a, w1, w2, w1_s, w2_s, topk_weights, topk_ids

320 )

321 - torch.cuda.synchronize()

321 + torch.accelerator.synchronize()

322 graph.replay()

323 - torch.cuda.synchronize()

323 + torch.accelerator.synchronize()

324

325 torch.testing.assert_close(out, ref_out, atol=0.035, rtol=0.035)

tests/kernels/moe/test_cutlass_moe.py

+2 -2

@@ -399,9 +399,9 @@ def test_cutlass_moe_8_bit_cuda_graph(

399 mt, topk_weights, topk_ids, per_act_token, per_out_ch

400 )

401

402 - torch.cuda.synchronize()

402 + torch.accelerator.synchronize()

403 graph.replay()

404 - torch.cuda.synchronize()

404 + torch.accelerator.synchronize()

405

406 torch.testing.assert_close(triton_output, cutlass_output, atol=9e-2, rtol=1e-2)

407

tests/kernels/moe/test_moe.py

+3 -3

@@ -272,9 +272,9 @@ def run_moe_test(

272 global_num_experts=global_num_experts,

273 expert_map=expert_map,

274 )

275 - torch.cuda.synchronize()

275 + torch.accelerator.synchronize()

276 graph.replay()

277 - torch.cuda.synchronize()

277 + torch.accelerator.synchronize()

278

279 torch.testing.assert_close(test_output, baseline_output, atol=atol, rtol=rtol)

280

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

768 F.pad(vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],

769 requires_grad=False,

770 )

771 - torch.cuda.synchronize()

771 + torch.accelerator.synchronize()

772 torch.accelerator.empty_cache()

773

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

tests/kernels/quantization/test_allspark_gemm.py

+1 -1

@@ -122,7 +122,7 @@ def test_gptq_allspark_gemm_ampere(mnk_factors, group_size, has_zp, dtype):

122 )

123

124 output_ref = torch.matmul(input, w_ref)

125 - torch.cuda.synchronize()

125 + torch.accelerator.synchronize()

126 max_diff = compute_max_diff(output, output_ref)

127

128 assert max_diff < 0.04

tests/kernels/quantization/test_cutlass_w4a8_moe.py

+1 -1

@@ -269,7 +269,7 @@ def test_cutlass_w4a8_moe_mm_end_to_end(shape, random_zero):

269 setup.c_strides,

270 setup.group_scale_strides,

271 )

272 - torch.cuda.synchronize()

272 + torch.accelerator.synchronize()

273

274 out_ref = compute_moe_reference_output(setup)

275 torch.testing.assert_close(setup.out, out_ref, rtol=1e-2, atol=1e-2)

tests/kernels/quantization/test_marlin_gemm.py

+4 -4

@@ -260,7 +260,7 @@ def test_gptq_marlin_repack(

260 marlin_q_w_2 = ops.gptq_marlin_repack(

261 q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit

262 )

263 - torch.cuda.synchronize()

263 + torch.accelerator.synchronize()

264

265 torch.testing.assert_close(marlin_q_w_1, marlin_q_w_2)

266

@@ -308,7 +308,7 @@ def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, is_a_8bit, nk_factors):

308 marlin_q_w_2 = ops.awq_marlin_repack(

309 q_w_awq, size_k, size_n, quant_type.size_bits, is_a_8bit

310 )

311 - torch.cuda.synchronize()

311 + torch.accelerator.synchronize()

312

313 torch.testing.assert_close(marlin_q_w_1, marlin_q_w_2)

314

@@ -564,7 +564,7 @@ def test_marlin_gemm_subset_input():

564 )

565 output_ref = torch.matmul(a_input, w_ref)

566

567 - torch.cuda.synchronize()

567 + torch.accelerator.synchronize()

568

569 max_diff = compute_max_diff(output, output_ref)

570

@@ -613,7 +613,7 @@ def test_marlin_gemm_with_bias(size_m):

613 )

614 output_ref = torch.matmul(a_input, w_ref) + b_bias.view(1, -1)

615

616 - torch.cuda.synchronize()

616 + torch.accelerator.synchronize()

617

618 max_diff = compute_max_diff(output, output_ref)

619

tests/kernels/test_cache_kernels.py

+1 -1

@@ -57,7 +57,7 @@ def test_gather_cache_oob():

57 seq_starts,

58 )

59

60 - torch.cuda.synchronize()

60 + torch.accelerator.synchronize()

61 assert True

62

63

tests/kernels/test_top_k_per_row.py

+1 -1

@@ -219,7 +219,7 @@ def _run_top_k_per_row_decode_test(

219 top_k,

220 )

221

222 - torch.cuda.synchronize()

222 + torch.accelerator.synchronize()

223

224 # Run reference implementation

225 torch_indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")

tests/models/language/generation/test_common.py

+1 -1

@@ -195,4 +195,4 @@ def test_models(

195 # unit tests. On ROCm, when using AITER

196 # the memory might not be deallocated completely

197 # before running the next test case

198 - torch.cuda.synchronize()

198 + torch.accelerator.synchronize()

tests/quantization/test_compressed_tensors.py

+1 -1

@@ -196,7 +196,7 @@ def test_compressed_tensors_w8a8_logprobs(

196 )

197

198 if current_platform.is_rocm():

199 - torch.cuda.synchronize()

199 + torch.accelerator.synchronize()

200

201

202 def test_compressed_tensors_no_enforce_eager(vllm_runner):

tools/pre_commit/check_torch_cuda.py

+1

@@ -9,6 +9,7 @@ import regex as re

9 # --------------------------------------------------------------------------- #

10 _TORCH_CUDA_PATTERNS = [

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

12 + r"\btorch\.cuda\.synchronize\b",

13 ]

14

15 ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}

vllm/distributed/elastic_ep/elastic_execute.py

+4 -4

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

217 dp_group=standby_dp_group,

218 expert_weights=model.expert_weights,

219 )

220 - torch.cuda.synchronize()

220 + torch.accelerator.synchronize()

221

222 def broadcast_expert_mapping(self) -> None:

223 standby_dp_group = get_standby_dp_group()

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

407 reset_compile_wrapper(self.worker.model_runner.get_model())

408

409 gc.collect()

410 - torch.cuda.synchronize()

410 + torch.accelerator.synchronize()

411 torch.accelerator.empty_cache()

412 unlock_workspace()

413 self.worker.compile_or_warm_up_model()

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

446

447 eplb_state.rearrange(rank_mapping=rank_mapping)

448 # NOTE(yongji): check whether we need to synchronize here

449 - torch.cuda.synchronize()

449 + torch.accelerator.synchronize()

450 # reset expert_rearrangement_step to ensure all ranks are synchronized

451 eplb_state.expert_rearrangement_step = 0

452 eplb_state.num_valid_physical_experts = (

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

491 dp_group=dp_group,

492 expert_weights=model.expert_weights,

493 )

494 - torch.cuda.synchronize()

494 + torch.accelerator.synchronize()

495

496 def receive_expert_mapping(self) -> tuple[torch.Tensor, int, int]:

497 dp_group = get_dp_group()

vllm/distributed/eplb/rebalance_execute.py

+1 -1

@@ -622,7 +622,7 @@ def rearrange_expert_weights_inplace(

622

623 # NOTE(bowen): We need this synchronize to run, but I don't know why.

624 # If you figure out the reason, please let me know -- thank you!

625 - torch.cuda.synchronize()

625 + torch.accelerator.synchronize()

626

627 old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy()

628 new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()

vllm/model_executor/kernels/linear/mixed_precision/cutlass.py

+1 -1

@@ -77,7 +77,7 @@ class CutlassW4A8LinearKernel(MPLinearKernel):

77 def transform_w_q(x):

78 assert isinstance(x, BasevLLMParameter)

79 convert_packed_uint4b8_to_signed_int4_inplace(x.data)

80 - torch.cuda.synchronize()

80 + torch.accelerator.synchronize()

81 permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)

82 x.data = ops.cutlass_encode_and_reorder_int4b(x.data.t().contiguous().t())

83 return x

vllm/v1/worker/gpu/model_runner.py

+2 -2

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

457 else:

458 self._dummy_pooler_run(hidden_states)

459

460 - torch.cuda.synchronize()

460 + torch.accelerator.synchronize()

461 del hidden_states, sample_hidden_states

462 gc.collect()

463

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

525 # to trigger JIT compilation.

526 if all("FLASHINFER" in b.get_name() for b in self.attn_backends.values()):

527 self._dummy_run(self.max_num_tokens, skip_attn=False)

528 - torch.cuda.synchronize()

528 + torch.accelerator.synchronize()

529

530 def finish_requests(self, scheduler_output: SchedulerOutput) -> None:

531 finished_req_ids = scheduler_output.finished_req_ids

vllm/v1/worker/gpu/warmup.py

+1 -1

@@ -102,4 +102,4 @@ def warmup_kernels(model_runner: GPUModelRunner) -> None:

102 cleanup_output.finished_req_ids = set(req_ids)

103 model_runner.execute_model(cleanup_output)

104 model_runner.kv_connector.set_disabled(False)

105 - torch.cuda.synchronize()

105 + torch.accelerator.synchronize()

vllm/v1/worker/gpu_model_runner.py

+4 -4

@@ -928,7 +928,7 @@ class GPUModelRunner(

928

929 # Note: used for model runner override.

930 def _sync_device(self) -> None:

931 - torch.cuda.synchronize()

931 + torch.accelerator.synchronize()

932

933 def _update_states(self, scheduler_output: "SchedulerOutput") -> None:

934 """Update the cached states and the persistent batch with the scheduler

@@ -5345,7 +5345,7 @@ class GPUModelRunner(

5345 cudagraph_runtime_mode=runtime_mode,

5346 )

5347

5348 - torch.cuda.synchronize()

5348 + torch.accelerator.synchronize()

5349 end_free_gpu_memory = torch.cuda.mem_get_info()[0]

5350

5351 # Disable cudagraph capturing globally, so any unexpected cudagraph

@@ -6266,13 +6266,13 @@ class GPUModelRunner(

6266 group_refs = group_lora_refs[current_item_idx : current_item_idx + num_items]

6267 group_request_ids = {req_id for req_id, _ in group_refs}

6268

6269 - torch.cuda.synchronize()

6269 + torch.accelerator.synchronize()

6270 start_time = time.perf_counter()

6271

6272 try:

6273 yield

6274 finally:

6275 - torch.cuda.synchronize()

6275 + torch.accelerator.synchronize()

6276 elapsed = time.perf_counter() - start_time

6277

6278 per_request_time = elapsed / max(len(group_request_ids), 1)

vllm/v1/worker/xpu_model_runner.py

-4

@@ -29,9 +29,6 @@ class XPUModelRunner(GPUModelRunner):

29 # FIXME: To be verified.

30 self.cascade_attn_enabled = False

31

32 - def _sync_device(self) -> None:

33 - torch.xpu.synchronize()

34 -

32

33 @contextmanager

34 def _torch_cuda_wrapper():

@@ -42,7 +39,6 @@ def _torch_cuda_wrapper():

39 torch.cuda.current_stream = torch.xpu.current_stream

40 torch.cuda.stream = torch.xpu.stream

41 torch.cuda.mem_get_info = torch.xpu.mem_get_info

45 - torch.cuda.synchronize = torch.xpu.synchronize

42 if supports_xpu_graph():

43 torch.cuda.graph = torch.xpu.graph

44 torch.cuda.CUDAGraph = torch.xpu.XPUGraph

[Hardware] Replace `torch.cuda.synchronize()` api with `torch.accelerator.synchronize` (#36085)