2a2d5d2780 · tduncan/Nvidia-Optimization-Stack-VLLM

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

+2 -2

@@ -255,8 +255,8 @@ def bench_run(

255 torch.cuda.synchronize()

256

257 # Timing

258 - start_event = torch.cuda.Event(enable_timing=True)

259 - end_event = torch.cuda.Event(enable_timing=True)

258 + start_event = torch.Event(enable_timing=True)

259 + end_event = torch.Event(enable_timing=True)

260

261 latencies = []

262 for _ in range(num_iters):

benchmarks/kernels/benchmark_moe.py

+2 -2

@@ -185,8 +185,8 @@ def benchmark_config(

185 graph.replay()

186 torch.cuda.synchronize()

187

188 - start_event = torch.cuda.Event(enable_timing=True)

189 - end_event = torch.cuda.Event(enable_timing=True)

188 + start_event = torch.Event(enable_timing=True)

189 + end_event = torch.Event(enable_timing=True)

190

191 latencies: list[float] = []

192 for i in range(num_iters):

benchmarks/kernels/benchmark_moe_permute_unpermute.py

+4 -4

@@ -105,8 +105,8 @@ def benchmark_permute(

105 graph.replay()

106 torch.cuda.synchronize()

107

108 - start_event = torch.cuda.Event(enable_timing=True)

109 - end_event = torch.cuda.Event(enable_timing=True)

108 + start_event = torch.Event(enable_timing=True)

109 + end_event = torch.Event(enable_timing=True)

110

111 latencies: list[float] = []

112 for i in range(num_iters):

@@ -241,8 +241,8 @@ def benchmark_unpermute(

241 graph.replay()

242 torch.cuda.synchronize()

243

244 - start_event = torch.cuda.Event(enable_timing=True)

245 - end_event = torch.cuda.Event(enable_timing=True)

244 + start_event = torch.Event(enable_timing=True)

245 + end_event = torch.Event(enable_timing=True)

246

247 latencies: list[float] = []

248 for i in range(num_iters):

benchmarks/kernels/benchmark_per_token_group_quant.py

+2 -2

@@ -30,8 +30,8 @@ def _time_cuda(

30 fn()

31 torch.cuda.synchronize()

32

33 - start = torch.cuda.Event(enable_timing=True)

34 - end = torch.cuda.Event(enable_timing=True)

33 + start = torch.Event(enable_timing=True)

34 + end = torch.Event(enable_timing=True)

35

36 start.record()

37 for _ in range(bench_iters):

benchmarks/kernels/benchmark_silu_mul_fp8_quant.py

+2 -2

@@ -253,8 +253,8 @@ def benchmark(

253 )

254 torch.cuda.synchronize()

255

256 - start_event = torch.cuda.Event(enable_timing=True)

257 - end_event = torch.cuda.Event(enable_timing=True)

256 + start_event = torch.Event(enable_timing=True)

257 + end_event = torch.Event(enable_timing=True)

258

259 # Benchmark

260 latencies: list[float] = []

benchmarks/kernels/benchmark_trtllm_decode_attention.py

+2 -2

@@ -127,8 +127,8 @@ def benchmark_decode(

127

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

129 torch.cuda.synchronize()

130 - start = torch.cuda.Event(enable_timing=True)

131 - end = torch.cuda.Event(enable_timing=True)

130 + start = torch.Event(enable_timing=True)

131 + end = torch.Event(enable_timing=True)

132 times = []

133 for i in range(warmup):

134 fn()

benchmarks/kernels/benchmark_trtllm_prefill_attention.py

+2 -2

@@ -139,8 +139,8 @@ def benchmark_prefill(

139

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

141 torch.cuda.synchronize()

142 - start = torch.cuda.Event(enable_timing=True)

143 - end = torch.cuda.Event(enable_timing=True)

142 + start = torch.Event(enable_timing=True)

143 + end = torch.Event(enable_timing=True)

144 times = []

145 for i in range(warmup):

146 fn()

benchmarks/kernels/benchmark_w8a8_block_fp8.py

+2 -2

@@ -183,8 +183,8 @@ def benchmark_config(

183 run()

184 torch.cuda.synchronize()

185

186 - start_event = torch.cuda.Event(enable_timing=True)

187 - end_event = torch.cuda.Event(enable_timing=True)

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):

tests/kernels/attention/test_merge_attn_states.py

+4 -4

@@ -150,8 +150,8 @@ def test_merge_attn_states(

150 output_torch = output.clone()

151 output_lse_torch = output_lse.clone()

152 total_time_torch_kernel = 0

153 - start = torch.cuda.Event(enable_timing=True)

154 - end = torch.cuda.Event(enable_timing=True)

153 + start = torch.Event(enable_timing=True)

154 + end = torch.Event(enable_timing=True)

155

156 # 0. Run the Torch kernel

157 prefix_lse_torch = prefix_lse.clone()

@@ -188,8 +188,8 @@ def test_merge_attn_states(

188 output_lse_ref_triton = output_lse.clone()

189

190 total_time_triton_kernel = 0

191 - start = torch.cuda.Event(enable_timing=True)

192 - end = torch.cuda.Event(enable_timing=True)

191 + start = torch.Event(enable_timing=True)

192 + end = torch.Event(enable_timing=True)

193

194 for _ in range(warmup_times):

195 merge_attn_states_triton(

vllm/v1/kv_offload/worker/cpu_gpu.py

+3 -3

@@ -68,9 +68,9 @@ class CpuGpuOffloadingHandler(OffloadingHandler):

68 self.h2d_stream = torch.cuda.Stream()

69

70 # job_id -> transfer cuda event

71 - self.transfer_events: dict[int, torch.cuda.Event] = {}

71 + self.transfer_events: dict[int, torch.Event] = {}

72 # list of cuda events available for re-use

73 - self.events_pool: list[torch.cuda.Event] = []

73 + self.events_pool: list[torch.Event] = []

74

75 pin_memory = is_pin_memory_available()

76

@@ -153,7 +153,7 @@ class CpuGpuOffloadingHandler(OffloadingHandler):

153 )

154 src_to_dst_tensor = torch.from_numpy(src_to_dst)

155

156 - event = self.events_pool.pop() if self.events_pool else torch.cuda.Event()

156 + event = self.events_pool.pop() if self.events_pool else torch.Event()

157 with torch.cuda.stream(stream):

158 for src_tensor, dst_tensor, kv_dim in zip(

159 src_tensors, dst_tensors, self.kv_dim_before_num_blocks

vllm/v1/worker/cpu_model_runner.py

+3 -3

@@ -96,14 +96,14 @@ def _torch_cuda_wrapper():

96 def __init__(self, *args, **kwargs) -> None:

97 pass

98

99 - cuda_event = torch.cuda.Event

99 + cuda_event = torch.Event

100 cuda_stream = torch.cuda.Stream

101 try:

102 - torch.cuda.Event = _EventPlaceholder

102 + torch.Event = _EventPlaceholder

103 torch.cuda.Stream = _StreamPlaceholder

104 yield

105 finally:

106 - torch.cuda.Event = cuda_event

106 + torch.Event = cuda_event

107 torch.cuda.Stream = cuda_stream

108

109

vllm/v1/worker/gpu_input_batch.py

+2 -2

@@ -265,7 +265,7 @@ class InputBatch:

265 # ids from prior step, if required by current sampling params

266 # (e.g. penalties).

267 self.sampled_token_ids_cpu: torch.Tensor | None = None

268 - self.async_copy_ready_event: torch.cuda.Event | None = None

268 + self.async_copy_ready_event: torch.Event | None = None

269

270 @property

271 def req_ids(self) -> list[str]:

@@ -891,7 +891,7 @@ class InputBatch:

891 def set_async_sampled_token_ids(

892 self,

893 sampled_token_ids_cpu: torch.Tensor,

894 - async_copy_ready_event: torch.cuda.Event,

894 + async_copy_ready_event: torch.Event,

895 ) -> None:

896 """

897 In async scheduling case, store ref to sampled_token_ids_cpu

vllm/v1/worker/gpu_model_runner.py

+6 -6

@@ -185,7 +185,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

185 self._invalid_req_indices = invalid_req_indices

186

187 # Event on the copy stream so we can synchronize the non-blocking copy.

188 - self.async_copy_ready_event = torch.cuda.Event()

188 + self.async_copy_ready_event = torch.Event()

189

190 # Keep a reference to the device tensor to avoid it being

191 # deallocated until we finish copying it to the host.

@@ -435,10 +435,10 @@ class GPUModelRunner(

435 self.async_output_copy_stream: torch.cuda.Stream | None = None

436 # cuda event to synchronize use of reused CPU tensors between steps

437 # when async scheduling is enabled.

438 - self.prepare_inputs_event: torch.cuda.Event | None = None

438 + self.prepare_inputs_event: torch.Event | None = None

439 if self.use_async_scheduling:

440 self.async_output_copy_stream = torch.cuda.Stream()

441 - self.prepare_inputs_event = torch.cuda.Event()

441 + self.prepare_inputs_event = torch.Event()

442

443 # self.cudagraph_batch_sizes sorts in ascending order.

444 if (

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

549

550 # Cached outputs.

551 self._draft_token_ids: list[list[int]] | torch.Tensor | None = None

552 - self.transfer_event = torch.cuda.Event()

552 + self.transfer_event = torch.Event()

553 self.sampled_token_ids_pinned_cpu = torch.empty(

554 (self.max_num_reqs, 1),

555 dtype=torch.int64,

@@ -559,10 +559,10 @@ class GPUModelRunner(

559

560 # Pre-allocated tensor for copying valid sampled token counts to CPU,

561 # with dedicated stream for overlapping and event for coordination.

562 - self.valid_sampled_token_count_event: torch.cuda.Event | None = None

562 + self.valid_sampled_token_count_event: torch.Event | None = None

563 self.valid_sampled_token_count_copy_stream: torch.cuda.Stream | None = None

564 if self.use_async_scheduling and self.num_spec_tokens:

565 - self.valid_sampled_token_count_event = torch.cuda.Event()

565 + self.valid_sampled_token_count_event = torch.Event()

566 self.valid_sampled_token_count_copy_stream = torch.cuda.Stream()

567 self.valid_sampled_token_count_cpu = torch.empty(

568 self.max_num_reqs,

vllm/v1/worker/ubatching.py

+4 -4

@@ -27,8 +27,8 @@ class UBatchContext:

27 ready_barrier: threading.Barrier,

28 cpu_wait_event: threading.Event,

29 cpu_signal_event: threading.Event,

30 - gpu_comm_done_event: torch.cuda.Event,

31 - gpu_compute_done_event: torch.cuda.Event,

30 + gpu_comm_done_event: torch.Event,

31 + gpu_compute_done_event: torch.Event,

32 schedule: str = "default",

33 ):

34 self.id = id

@@ -207,8 +207,8 @@ def make_ubatch_contexts(

207 Create a context manager for micro-batching synchronization.

208 """

209 cpu_events = [threading.Event() for _ in range(num_micro_batches)]

210 - gpu_comm_done_events = [torch.cuda.Event() for _ in range(num_micro_batches)]

211 - gpu_compute_done_events = [torch.cuda.Event() for _ in range(num_micro_batches)]

210 + gpu_comm_done_events = [torch.Event() for _ in range(num_micro_batches)]

211 + gpu_compute_done_events = [torch.Event() for _ in range(num_micro_batches)]

212

213 assert len(forward_contexts) == 2

214

vllm/v1/worker/xpu_model_runner.py

+1 -8

@@ -37,19 +37,12 @@ class XPUModelRunner(GPUModelRunner):

37

38 @contextmanager

39 def _torch_cuda_wrapper():

40 - class _EventPlaceholder:

41 - def __init__(self, *args, **kwargs) -> None:

42 - self.record = lambda: None

43 - self.synchronize = lambda: None

44 -

40 try:

41 # replace cuda APIs with xpu APIs, this should work by default

47 - torch.cuda.Event = torch.xpu.Event

42 torch.cuda.Stream = torch.xpu.Stream

43 torch.cuda.default_stream = torch.xpu.current_stream

44 torch.cuda.current_stream = torch.xpu.current_stream

45 torch.cuda.stream = torch.xpu.stream

46 yield

47 finally:

54 - # if anything goes wrong, just patch it with a placeholder

55 - torch.cuda.Event = _EventPlaceholder

48 + pass

Replace `torch.cuda.Event` with `torch.Event` for better hardware compatibility (#26985)