07f9baf756 · tduncan/Nvidia-Optimization-Stack-VLLM

benchmarks/benchmark_topk_topp.py

+4 -2

@@ -132,8 +132,10 @@ def benchmark_function(

132 reset_memory_stats()

133

134 # Benchmark

135 - start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]

136 - end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]

135 + start_events = [

136 + torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)

137 + ]

138 + end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]

139

140 for i in range(benchmark_iters):

141 logits_copy = logits.clone()

benchmarks/kernels/benchmark_moe_defaults.py

+2 -2

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

134 torch.accelerator.synchronize()

135

136 # Benchmark

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

138 - end = torch.Event(enable_timing=True)

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

138 + end = torch.cuda.Event(enable_timing=True)

139 start.record()

140 for _ in range(num_iters):

141 with override_config(config):

benchmarks/kernels/benchmark_selective_state_update.py

+2 -2

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

170 graph.replay()

171 torch.accelerator.synchronize()

172

173 - start = torch.Event(enable_timing=True)

174 - end = torch.Event(enable_timing=True)

173 + start = torch.cuda.Event(enable_timing=True)

174 + end = torch.cuda.Event(enable_timing=True)

175 latencies: list[float] = []

176 for _ in range(num_iters):

177 start.record()

tests/v1/kv_connector/unit/test_hf3fs_connector.py

+1 -1

@@ -33,7 +33,7 @@ def hf3fs_stats():

33 def _make_cuda_event():

34 """Return a real CUDA event when available, otherwise a MagicMock."""

35 if torch.cuda.is_available():

36 - return torch.Event()

36 + return torch.cuda.Event()

37 return MagicMock()

38

39

tools/pre_commit/check_torch_cuda.py

+1 -9

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

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

10 _TORCH_CUDA_PATTERNS = [

12 - r"\btorch\.cuda\.(manual_seed|manual_seed_all|Event)\b",

12 + r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b",

13 r"\bwith\storch\.cuda\.device\b",

14 # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally

15 r"\bcuda_device_count_stateless\(\)\b",

@@ -21,7 +21,6 @@ ALLOWED_FILES = {

21 "vllm/device_allocator/",

22 "vllm/distributed/weight_transfer/ipc_engine.py",

23 "tests/distributed/test_packed_tensor.py",

24 - "tools/pre_commit/check_torch_cuda.py",

24 }

25

26

@@ -40,13 +39,6 @@ def scan_file(path: str) -> int:

39 f"Found {matched_text} API call. Use set_random_seed instead."

40 )

41 return 1

43 - if matched_text == "torch.cuda.Event":

44 - print(

45 - f"{path}:{line_num}: "

46 - "\033[91merror:\033[0m "

47 - "Found torch.cuda.Event API call. Use torch.Event instead."

48 - )

49 - return 1

42 print(

43 f"{path}:{line_num}: "

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

vllm/distributed/eplb/eplb_state.py

+2 -2

@@ -742,8 +742,8 @@ class EplbState:

742 is_main_rank = ep_rank == 0

743 if is_main_rank:

744 if not self.is_async or is_profile:

745 - start_event = torch.Event(enable_timing=True)

746 - end_event = torch.Event(enable_timing=True)

745 + start_event = torch.cuda.Event(enable_timing=True)

746 + end_event = torch.cuda.Event(enable_timing=True)

747 start_event.record()

748 logger.info(

749 "Rearranging experts %s %s...",

vllm/distributed/eplb/eplb_utils.py

+2 -2

@@ -31,7 +31,7 @@ class CpuGpuEvent:

31 """

32

33 def __init__(self):

34 - self._event = torch.Event()

34 + self._event = torch.cuda.Event()

35 self._recorded = threading.Event()

36

37 def wait(self, stream: torch.cuda.Stream | None = None):

@@ -56,7 +56,7 @@ class CpuGpuEvent:

56 "CpuGpuEvent.record() called before the previous event was "

57 "consumed by wait()"

58 )

59 - self._event = torch.Event()

59 + self._event = torch.cuda.Event()

60 self._event.record(stream)

61 self._recorded.set()

62

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

+4 -4

@@ -239,7 +239,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):

239 # this event is complete the request is considered "done sending"

240 # by get_finished; clients block on the per-file flock to wait for

241 # the disk write itself.

242 - self._req_copy_events: dict[str, torch.Event] = {}

242 + self._req_copy_events: dict[str, torch.cuda.Event] = {}

243 # req_ids reported as finished-generating by the scheduler,

244 # accumulated across get_finished calls.

245 self._accumulated_finished_req_ids: set[str] = set()

@@ -320,7 +320,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):

320 @staticmethod

321 def _write_tensors(

322 tensors: dict[str, torch.Tensor],

323 - event: torch.Event,

323 + event: torch.cuda.Event,

324 filename: str,

325 lock_fd: int | None,

326 ) -> None:

@@ -375,7 +375,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):

375 copy_stream = self._get_copy_stream()

376

377 # Ensure the copy stream sees all prior writes on the default stream.

378 - ready_event = torch.Event()

378 + ready_event = torch.cuda.Event()

379 ready_event.record()

380 copy_stream.wait_event(ready_event)

381

@@ -396,7 +396,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):

396 pinned_hs.copy_(hidden_states_gpu, non_blocking=True)

397

398 # Record completion of this copy on the copy stream.

399 - copy_done = torch.Event()

399 + copy_done = torch.cuda.Event()

400 copy_done.record(copy_stream)

401

402 # token_ids is already on CPU (created in request_finished).

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py

+1 -1

@@ -221,7 +221,7 @@ class Hf3fsClient:

221

222 @wsynchronized()

223 def batch_write(

224 - self, offsets: list[int], tensors: list[torch.Tensor], event: torch.Event

224 + self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event

225 ) -> list[int]:

226 """Write data from tensors to the file at specified offsets.

227

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py

+3 -3

@@ -133,7 +133,7 @@ class AsyncOperationManager:

133 # CUDA streams for async operations

134 self._save_stream = torch.cuda.Stream()

135 self._load_stream = torch.cuda.Stream()

136 - self._save_event = torch.Event()

136 + self._save_event = torch.cuda.Event()

137

138 # Buffer allocators for data copying

139 self._save_buffer_allocator = CopyBufferAllocator(

@@ -171,7 +171,7 @@ class AsyncOperationManager:

171 def submit_save_operation(self, request_id: str, block_ids, block_hashes) -> Future:

172 """Submit a save operation for async execution."""

173 future: Future[Any] = Future()

174 - main_stream_event = torch.Event()

174 + main_stream_event = torch.cuda.Event()

175 main_stream_event.record()

176 task = (request_id, block_ids, block_hashes, future, main_stream_event)

177 self._save_queue.put(task)

@@ -304,7 +304,7 @@ class AsyncOperationManager:

304 block_ids, buffers, "gather"

305 )

306

307 - save_stream_event = torch.Event()

307 + save_stream_event = torch.cuda.Event()

308 save_stream_event.record(self._save_stream) # Record gather completion

309

310 # Step3: Write data in batches

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py

+1 -1

@@ -75,7 +75,7 @@ class Hf3fsClient:

75 return torch.frombuffer(buffer_data, dtype=dtype)

76

77 def batch_write(

78 - self, offsets: list[int], tensors: list[torch.Tensor], event: torch.Event

78 + self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event

79 ) -> list[int]:

80 """Write data from tensors to file at specified offsets."""

81 results = []

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py

+4 -4

@@ -430,7 +430,7 @@ class LMCacheMPWorkerAdapter:

430

431 @_lmcache_nvtx_annotate

432 def submit_store_request(

433 - self, request_id: str, op: LoadStoreOp, event: torch.Event

433 + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event

434 ):

435 """

436 Submit a KV cache store request to LMCache

@@ -464,7 +464,7 @@ class LMCacheMPWorkerAdapter:

464

465 @_lmcache_nvtx_annotate

466 def submit_retrieve_request(

467 - self, request_id: str, op: LoadStoreOp, event: torch.Event

467 + self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event

468 ):

469 """

470 Submit a KV cache retrieve request to LMCache

@@ -501,7 +501,7 @@ class LMCacheMPWorkerAdapter:

501 self,

502 request_ids: list[str],

503 ops: list[LoadStoreOp],

504 - event: torch.Event,

504 + event: torch.cuda.Event,

505 ):

506 """

507 Submit a batched store request to LMCache

@@ -550,7 +550,7 @@ class LMCacheMPWorkerAdapter:

550 self,

551 request_ids: list[str],

552 ops: list[LoadStoreOp],

553 - event: torch.Event,

553 + event: torch.cuda.Event,

554 ):

555 """

556 Submit a batched retrieve request to LMCache

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py

+2 -2

@@ -589,7 +589,7 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1):

589 return

590

591 with torch.cuda.stream(torch.cuda.current_stream()):

592 - event = torch.Event(interprocess=True)

592 + event = torch.cuda.Event(interprocess=True)

593 event.record()

594

595 self.worker_adapter.batched_submit_retrieve_requests(

@@ -663,7 +663,7 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1):

663 return

664

665 with torch.cuda.stream(torch.cuda.current_stream()):

666 - event = torch.Event(interprocess=True)

666 + event = torch.cuda.Event(interprocess=True)

667 event.record()

668

669 self.worker_adapter.batched_submit_store_requests(

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

+1 -1

@@ -323,7 +323,7 @@ class ReqMeta:

323 can_save: bool | None = None

324 load_spec: LoadSpec | None = None

325 is_last_chunk: bool | None = None

326 - current_event: torch.Event | None = None

326 + current_event: torch.cuda.Event | None = None

327

328 token_ids: list[int] | None = None

329 num_prompt_tokens: int | None = None

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+1 -1

@@ -1357,7 +1357,7 @@ class MooncakeStoreWorker:

1357 current_event = None

1358 for request in meta.requests:

1359 if request.can_save:

1360 - current_event = torch.Event()

1360 + current_event = torch.cuda.Event()

1361 current_event.record()

1362 break

1363

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py

+1 -1

@@ -56,7 +56,7 @@ class WriteTask:

56 local_block_ids: list[int]

57 remote_block_ids_hint: list[int] | None

58 layer_name: str

59 - event: torch.Event

59 + event: torch.cuda.Event

60 remote_notify_port: int

61 remote_ip: str

62 enqueue_time: float = field(default_factory=time.perf_counter)

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

+1 -1

@@ -1061,7 +1061,7 @@ class MoRIIOConnectorWorker:

1061 # when mori-io supports ibgda functionality

1062

1063 stream = torch.cuda.current_stream()

1064 - event = torch.Event()

1064 + event = torch.cuda.Event()

1065 event.record(stream)

1066

1067 task = WriteTask(

vllm/lora/layers/base_linear.py

+1 -1

@@ -89,7 +89,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

89 vllm_config = get_current_vllm_config()

90 self._lora_stream = _get_lora_aux_cuda_stream()

91 assert current_platform.is_cuda_alike()

92 - self._events = [torch.Event(), torch.Event()]

92 + self._events = [torch.cuda.Event(), torch.cuda.Event()]

93 # lora_linear avoids prefix conflicts with the base layer

94 self.layer_name = self.base_layer.prefix + ".lora_linear_async"

95 compilation_config = vllm_config.compilation_config

vllm/lora/layers/fused_moe.py

+2 -2

@@ -118,7 +118,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

118

119 def _init_lora_stream_context(self) -> None:

120 self._lora_stream: torch.cuda.Stream | None = None

121 - self._events: tuple[torch.Event, ...] | None = None

121 + self._events: tuple[torch.cuda.Event, ...] | None = None

122 if not self._enable_aux_cuda_stream:

123 return

124 if not current_platform.is_cuda_alike():

@@ -127,7 +127,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

127 # 4 events: 2 per (base GEMM, LoRA) pair so w13 and w2 don't reuse

128 # the same event objects; reuse-within-a-pair is fine because the

129 # second pair starts only after intermediate_cache1.add_() has joined.

130 - self._events = tuple(torch.Event() for _ in range(4))

130 + self._events = tuple(torch.cuda.Event() for _ in range(4))

131

132 def _build_lora_context(self):

133 use_dual_stream = (

vllm/model_executor/layers/fused_moe/experts/lora_context.py

+1 -1

@@ -51,7 +51,7 @@ class MoELoRAContext:

51 # Events are paired one-per-overlap-pair: events[0,1] for w13,

52 # events[2,3] for w2, so the two pairs do not race on the same event.

53 aux_stream: torch.cuda.Stream | None = None

54 - events: tuple[torch.Event, ...] | None = None

54 + events: tuple[torch.cuda.Event, ...] | None = None

55

56 # Per-rank token→LoRA mapping after EP dispatch. Set by

57 # FusedMoEPrepareAndFinalizeModular.prepare() when EP+LoRA is active, read

vllm/model_executor/offloader/prefetch.py

+2 -2

@@ -388,7 +388,7 @@ class _ModuleOffloader:

388

389 # Event to signal when H2D copy to static buffer is complete.

390 # Used for per-layer synchronization (both eager and capture modes).

391 - self._copy_done_event = torch.Event()

391 + self._copy_done_event = torch.cuda.Event()

392

393 # Track whether _copy_done_event is valid for eager-mode wait_event.

394 # False when: (1) never recorded, or (2) last recorded during a

@@ -518,7 +518,7 @@ class _ModuleOffloader:

518

519 # Fork: record event on compute stream, copy_stream waits on it

520 # This joins copy_stream to any active CUDA graph capture

521 - fork_event = torch.Event()

521 + fork_event = torch.cuda.Event()

522 torch.cuda.current_stream().record_event(fork_event)

523 self.copy_stream.wait_event(fork_event)

524

vllm/models/deepseek_v4/attention.py

+5 -2

@@ -272,7 +272,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

272 # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;

273 # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins

274 # before post-GEMM starts.

275 - self.ln_events = [torch.Event() for _ in range(4)]

275 + self.ln_events = [torch.cuda.Event() for _ in range(4)]

276

277 assert cache_config is not None, "DeepseekV4 attention requires cache_config"

278 # ---- Attention / KV-cache setup ----

@@ -760,7 +760,10 @@ class DeepseekV4Indexer(nn.Module):

760

761 # None on ROCm — maybe_execute_in_parallel falls back to sequential.

762 self.aux_stream = aux_stream

763 - self.ln_events: list[torch.Event] = [torch.Event(), torch.Event()]

763 + self.ln_events: list[torch.cuda.Event] = [

764 + torch.cuda.Event(),

765 + torch.cuda.Event(),

766 + ]

767

768 def forward(

769 self,

vllm/models/deepseek_v4/xpu/xpu_sparse.py

+11

@@ -44,6 +44,17 @@ class DeepseekV4XPUAttention(DeepseekV4Attention):

44 backend_cls = DeepseekV4XPUSparseBackend

45 use_flashmla_fp8_layout = True

46

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

48 + # torch.cuda.Event() raises RuntimeError on XPU ("dummy base class").

49 + # The Base and DeepseekV4Indexer both create cuda Events in __init__, so

50 + # we temporarily redirect torch.cuda.Event → torch.xpu.Event.

51 + _orig_event = torch.cuda.Event

52 + torch.cuda.Event = torch.xpu.Event # type: ignore[misc]

53 + try:

54 + super().__init__(*args, **kwargs)

55 + finally:

56 + torch.cuda.Event = _orig_event # type: ignore[misc]

57 +

58 def _fused_qnorm_rope_kv_insert(self, q, kv, positions, attn_metadata):

59 from typing import cast

60

vllm/utils/multi_stream_utils.py

+5 -5

@@ -20,8 +20,8 @@ class EventType(Enum):

20 def maybe_execute_in_parallel(

21 fn0: Callable[[], Any],

22 fn1: Callable[[], Any],

23 - event0: torch.Event,

24 - event1: torch.Event,

23 + event0: torch.cuda.Event,

24 + event1: torch.cuda.Event,

25 aux_stream: torch.cuda.Stream | None = None,

26 ) -> tuple[Any, Any]:

27 """Run two functions potentially in parallel on separate CUDA streams.

@@ -61,8 +61,8 @@ def maybe_execute_in_parallel(

61 def execute_in_parallel(

62 default_fn: Callable[[], Any],

63 aux_fns: list[Callable[[], Any] | None],

64 - start_event: torch.Event,

65 - done_events: list[torch.Event],

64 + start_event: torch.cuda.Event,

65 + done_events: list[torch.cuda.Event],

66 aux_streams: list[torch.cuda.Stream] | None = None,

67 enable: bool = False,

68 ) -> tuple[Any, list[Any]]:

@@ -108,7 +108,7 @@ def execute_in_parallel(

108 )

109

110 aux_results = [None] * len(aux_fns)

111 - pending: list[torch.Event] = []

111 + pending: list[torch.cuda.Event] = []

112

113 start_event.record()

114 for i, fn in enumerate(aux_fns):

vllm/v1/spec_decode/ngram_proposer_gpu.py

+2 -2

@@ -468,7 +468,7 @@ class NgramProposerGPU:

468

469

470 def update_scheduler_for_invalid_drafts(

471 - num_valid_draft_tokens_event: torch.Event,

471 + num_valid_draft_tokens_event: torch.cuda.Event,

472 num_valid_draft_tokens_cpu: torch.Tensor,

473 scheduler_output: "SchedulerOutput",

474 req_id_to_index: dict[str, int],

@@ -643,7 +643,7 @@ def _sync_num_tokens(

643 def copy_num_valid_draft_tokens(

644 num_valid_draft_tokens_cpu: torch.Tensor,

645 num_valid_draft_tokens_copy_stream: torch.cuda.Stream,

646 - num_valid_draft_tokens_event: torch.Event,

646 + num_valid_draft_tokens_event: torch.cuda.Event,

647 num_valid_draft_tokens: torch.Tensor | None,

648 batch_size: int,

649 ) -> None:

vllm/v1/worker/cpu/shm.py

+1

@@ -48,6 +48,7 @@ def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]:

48

49

50 torch.Event = _EventPlaceholder

51 +torch.cuda.Event = _EventPlaceholder

52 torch.cuda.Stream = _StreamPlaceholder

53 torch.cuda.set_stream = noop

54 torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()

vllm/v1/worker/gpu/async_utils.py

+2 -2

@@ -24,7 +24,7 @@ class AsyncOutput(AsyncModelRunnerOutput):

24 self.model_runner_output = model_runner_output

25 self.sampler_output = sampler_output

26 self.num_sampled_tokens = num_sampled_tokens

27 - self.copy_event = torch.Event()

27 + self.copy_event = torch.cuda.Event()

28

29 with stream(copy_stream, main_stream):

30 copy_stream.wait_stream(main_stream)

@@ -81,7 +81,7 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):

81 self.model_runner_output = model_runner_output

82 self.pooler_output = pooler_output

83 self.is_valid = is_valid

84 - self.copy_event = torch.Event()

84 + self.copy_event = torch.cuda.Event()

85

86 with stream(copy_stream, main_stream):

87 copy_stream.wait_stream(main_stream)

vllm/v1/worker/gpu/pp_utils.py

+1 -1

@@ -17,7 +17,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch

17 class PendingRecv:

18 """Per-step slot data for a deferred postprocess on the main stream."""

19

20 - event: torch.Event

20 + event: torch.cuda.Event

21

22 sampled_tokens: torch.Tensor # [num_reqs, max_sample_len]

23 num_sampled: torch.Tensor # [num_reqs]

vllm/v1/worker/gpu/spec_decode/utils.py

+1 -1

@@ -12,7 +12,7 @@ class DraftTokensHandler:

12 def __init__(self, device: torch.device | None = None):

13 self.device = device

14 self.copy_stream = torch.cuda.Stream(device)

15 - self.copy_event = torch.Event()

15 + self.copy_event = torch.cuda.Event()

16

17 self.req_ids: list[str] = []

18 self.draft_tokens_np: np.ndarray | None = None

vllm/v1/worker/gpu_model_runner.py

+2 -2

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

860 # N-gram GPU path: async D2H buffer/event for per-request valid draft counts.

861 self._num_valid_draft_tokens: torch.Tensor | None = None

862 self._num_valid_draft_tokens_cpu: torch.Tensor | None = None

863 - self._num_valid_draft_tokens_event: torch.Event | None = None

863 + self._num_valid_draft_tokens_event: torch.cuda.Event | None = None

864 self._num_valid_draft_tokens_copy_stream: torch.cuda.Stream | None = None

865 if (

866 self.speculative_config is not None

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

869 self._num_valid_draft_tokens_cpu = torch.empty(

870 self.max_num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY

871 )

872 - self._num_valid_draft_tokens_event = torch.Event()

872 + self._num_valid_draft_tokens_event = torch.cuda.Event()

873 self._num_valid_draft_tokens_copy_stream = torch.cuda.Stream()

874

875 self._draft_token_req_ids: list[str] | None = None

vllm/v1/worker/xpu_model_runner.py

+1

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

49 torch.cuda.current_stream = partial(torch.xpu.current_stream)

50 torch.cuda.stream = partial(torch.xpu.stream)

51 torch.cuda.set_stream = partial(torch.xpu.set_stream)

52 + torch.cuda.Event = partial(torch.xpu.Event)

53 if supports_xpu_graph():

54 torch.cuda.graph = partial(torch.xpu.graph)

55 torch.cuda.CUDAGraph = torch.xpu.XPUGraph

Revert "[Platform] Replace `torch.cuda.Event` with `torch.Event` (#47140)" (#47668)