def attention(
        q,
        k,
        v,
        sp_group,
        dropout_p: float = 0.0,
        softmax_scale: Optional[float] = None,
        deterministic: bool = False,
        return_softmax: bool = False,
    ):
        """Ring attention

        Args:
            q (torch.Tensor): shape [B, S, N, D]
            k (torch.Tensor): shape [B, S, N, D]
            v (torch.Tensor): shape [B, S, N, D]
            sp_group (dist.ProcessGroup): sequence parallel group
            dropout_p (float, optional): dropout prob. Defaults to 0.0.
            softmax_scale (Optional[float], optional): softmax scale. Defaults to None.
            deterministic (Optional[bool], optional): backward deterministic mode. Defaults to False.
            return_softmax (bool, optional): return softmax or not. Defaults to False.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: output and log sum exp. Output's shape should be [B, S, N, D]. LSE's shape should be [B, N, S].
        """
        if RingAttention.ATTN_DONE is None:
            RingAttention.ATTN_DONE = torch.cuda.Event()
        if RingAttention.SP_STREAM is None:
            RingAttention.SP_STREAM = torch.cuda.Stream()
        out, softmax_lse = RingAttention.apply(
            q, k, v, sp_group, RingAttention.SP_STREAM, dropout_p, softmax_scale, deterministic
        )
        if return_softmax:
            return out, softmax_lse
        return out
Ring attention

        Args:
            q (torch.Tensor): shape [B, S, N, D]
            k (torch.Tensor): shape [B, S, N, D]
            v (torch.Tensor): shape [B, S, N, D]
            sp_group (dist.ProcessGroup): sequence parallel group
            dropout_p (float, optional): dropout prob. Defaults to 0.0.
            softmax_scale (Optional[float], optional): softmax scale. Defaults to None.
            deterministic (Optional[bool], optional): backward deterministic mode. Defaults to False.
            return_softmax (bool, optional): return softmax or not. Defaults to False.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: output and log sum exp. Output's shape should be [B, S, N, D]. LSE's shape should be [B, N, S].
        
opensora/models/mmdit/distributed.py