[GH-ISSUE #412] How to use the 'EventOverlap'? Is there a way to convert it to a 'torch.cuda.Event()'? #277

New issue

Closed

opened 2026-10-26 09:57:51 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 09:57:51 +00:00

Owner

Originally created by @arminzhu on GitHub (Sep 16, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/412

I want to use the 'EventOverlap' like this

torch.cuda.current_stream().wait_event(_EventOverlap)

I also noticed that there is a method named current_stream_wait in the EventOverlap class, but I’m not sure whether it is what I want.

class EventOverlap:
    """
    A wrapper class to manage CUDA events, also for better overlapping convenience.

    Attributes:
        event: the CUDA event captured.
        extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
    """

    def __init__(self, event: Optional[EventHandle] = None,
                 extra_tensors: Optional[Tuple[torch.Tensor]] = None) -> None:
        """
        Initialize the class.

        Arguments:
            event: the CUDA event captured.
            extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph.
        """
        self.event = event

        # NOTES: we use extra tensors to achieve stream recording, otherwise,
        # stream recording will be incompatible with CUDA graph.
        self.extra_tensors = extra_tensors

    def current_stream_wait(self) -> None:
        """
        The current stream `torch.cuda.current_stream()` waits for the event to be finished.
        """
        assert self.event is not None
        self.event.current_stream_wait()

Originally created by @arminzhu on GitHub (Sep 16, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepEP/issues/412 I want to use the 'EventOverlap' like this `torch.cuda.current_stream().wait_event(_EventOverlap) ` I also noticed that there is a method named `current_stream_wait` in the `EventOverlap` class, but I’m not sure whether it is what I want. > ``` > class EventOverlap: > """ > A wrapper class to manage CUDA events, also for better overlapping convenience. > > Attributes: > event: the CUDA event captured. > extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph. > """ > > def __init__(self, event: Optional[EventHandle] = None, > extra_tensors: Optional[Tuple[torch.Tensor]] = None) -> None: > """ > Initialize the class. > > Arguments: > event: the CUDA event captured. > extra_tensors: an easier way to simulate PyTorch tensor `record_stream`, may be useful with CUDA graph. > """ > self.event = event > > # NOTES: we use extra tensors to achieve stream recording, otherwise, > # stream recording will be incompatible with CUDA graph. > self.extra_tensors = extra_tensors > > def current_stream_wait(self) -> None: > """ > The current stream `torch.cuda.current_stream()` waits for the event to be finished. > """ > assert self.event is not None > self.event.current_stream_wait() > ```

bsakaguchi closed this issue

2026-10-26 09:57:51 +00:00

bsakaguchi referenced this issue

2026-10-27 21:03:15 +00:00

[PR #277] Allow using few SMs for low-latency mode #498