diff --git a/flashinfer/parallel_attention/__init__.py b/flashinfer/parallel_attention/__init__.py
new file mode 100644
index 0000000000..7cb1d30070
--- /dev/null
+++ b/flashinfer/parallel_attention/__init__.py
@@ -0,0 +1,19 @@
+from .parallel_attention import ParallelAttention as ParallelAttention
+from .parallel_config import UnevenCPConfig as UnevenCPConfig
+from .parallel_config import VarlenCPConfig as VarlenCPConfig
+from .utils import split_varlen_input as split_varlen_input
+from .utils import ulysses_varlen_config as ulysses_varlen_config
+from .utils import ring_varlen_config as ring_varlen_config
+from .utils import uneven_cp_config as uneven_cp_config
+from .utils import get_parallel_groups as get_parallel_groups
+
+__all__ = [
+    "ParallelAttention",
+    "UnevenCPConfig",
+    "VarlenCPConfig",
+    "split_varlen_input",
+    "ulysses_varlen_config",
+    "ring_varlen_config",
+    "uneven_cp_config",
+    "get_parallel_groups",
+]
diff --git a/flashinfer/parallel_attention/attention_ops.py b/flashinfer/parallel_attention/attention_ops.py
new file mode 100644
index 0000000000..35b1a6c578
--- /dev/null
+++ b/flashinfer/parallel_attention/attention_ops.py
@@ -0,0 +1,253 @@
+import logging
+
+import math
+import torch
+
+from .utils import (
+    convert_output_layout,
+    convert_qkv_layout,
+)
+
+logger = logging.getLogger(__name__)
+
+try:
+    import flash_attn_interface
+except ImportError:
+    flash_attn_interface = None
+
+from flashinfer.prefill import fmha_varlen
+
+
+class AttentionOpManager:
+    _attn_registry: dict[str, type] = {}
+
+    @classmethod
+    def op_type(cls):
+        return "attention"
+
+    @classmethod
+    def set_attn_config(cls, **kwargs):
+        for key, value in kwargs.items():
+            if hasattr(cls, key):
+                setattr(cls, key, value)
+            else:
+                raise AttributeError(f"'{cls.__name__}' has no attribute '{key}'")
+
+    @classmethod
+    def register_attn(cls, attn_type):
+        def decorator(attn_class):
+            # Register the attention class
+            cls._attn_registry[attn_type] = attn_class
+            return attn_class
+
+        return decorator
+
+    @classmethod
+    def get_impl(cls, name=None):
+        if name is None:
+            name = cls.attn_type
+        attn_class = cls._attn_registry.get(name)
+        if attn_class is None:
+            raise ValueError(f"Attention function {name} not found in registry")
+        return attn_class()  # Create and return an instance
+
+    @classmethod
+    def get_registered_types(cls):
+        return list(cls._attn_registry.keys())
+
+
+@AttentionOpManager.register_attn("flash-attn3")
+class FlashAttn3:
+    def __call__(
+        self,
+        query,
+        key,
+        value,
+        attn_mask=None,
+        is_causal=False,
+        return_lse=False,
+        tensor_layout="HND",
+        cur_rank_cu_seqlens_q=None,
+        cur_rank_cu_seqlens_k=None,
+        cur_rank_max_seqlen_q=0,
+        cur_rank_max_seqlen_k=0,
+        **kwargs,
+    ):
+        if flash_attn_interface is None:
+            raise ImportError("FlashAttn3 is not installed")
+
+        if tensor_layout not in ["HND", "NHD"]:
+            raise NotImplementedError("Tensor layout not supported for FlashAttn3")
+
+        if tensor_layout == "HND":
+            query, key, value = convert_qkv_layout(
+                query, key, value, src_layout="HND", dst_layout="NHD"
+            )
+
+        if attn_mask is not None:
+            raise NotImplementedError("FlashAttn3 does not support attn_mask yet")
+
+        # FA3 only supports float16 and bfloat16
+        origin_dtype = query.dtype
+        if query.dtype not in [torch.float16, torch.bfloat16]:
+            query = query.to(torch.float16)
+            key = key.to(torch.float16)
+            value = value.to(torch.float16)
+
+        if cur_rank_cu_seqlens_q is None:
+            query = torch.unsqueeze(query, dim=0)
+            key = torch.unsqueeze(key, dim=0)
+            value = torch.unsqueeze(value, dim=0)
+            output = flash_attn_interface.flash_attn_func(
+                q=query,
+                k=key,
+                v=value,
+                softmax_scale=None,
+                causal=is_causal,
+                qv=None,
+                q_descale=None,
+                k_descale=None,
+                v_descale=None,
+                window_size=(-1, -1),
+                attention_chunk=0,
+                softcap=0.0,
+                num_splits=1,
+                pack_gqa=None,
+                deterministic=False,
+                sm_margin=0,
+                return_attn_probs=return_lse,
+            )
+
+            if isinstance(output, tuple):
+                lse = torch.squeeze(output[1], dim=0)
+                output = torch.squeeze(output[0], dim=0)
+                output = (output, lse)
+            else:
+                output = torch.squeeze(output, dim=0)
+
+        else:
+            output = flash_attn_interface.flash_attn_varlen_func(
+                q=query,
+                k=key,
+                v=value,
+                cu_seqlens_q=cur_rank_cu_seqlens_q,
+                cu_seqlens_k=cur_rank_cu_seqlens_k,
+                max_seqlen_q=cur_rank_max_seqlen_q,
+                max_seqlen_k=cur_rank_max_seqlen_k,
+                seqused_q=None,
+                seqused_k=None,
+                softmax_scale=None,
+                causal=is_causal,
+                qv=None,
+                q_descale=None,
+                k_descale=None,
+                v_descale=None,
+                window_size=(-1, -1),
+                attention_chunk=0,
+                softcap=0.0,
+                num_splits=1,
+                pack_gqa=None,
+                deterministic=False,
+                sm_margin=0,
+                return_attn_probs=return_lse,
+            )
+
+        lse = None
+        if isinstance(output, tuple):
+            lse = output[1]
+            output = output[0]
+
+        if tensor_layout == "HND":
+            output = convert_output_layout(output, src_layout="NHD", dst_layout="HND")
+
+        if tensor_layout == "NHD" and lse is not None:
+            lse = lse.permute(1, 0)
+
+        if output.dtype != origin_dtype:
+            output = output.to(origin_dtype)
+
+        if return_lse:
+            assert lse is not None, "lse is not returned by FlashAttn3"
+            return output, lse
+        else:
+            return output
+
+
+@AttentionOpManager.register_attn("cutlass")
+class CutlassFmha:
+    def __call__(
+        self,
+        query,
+        key,
+        value,
+        attn_mask=None,
+        is_causal=False,
+        return_lse=False,
+        tensor_layout="HND",
+        cur_rank_cu_seqlens_q=None,
+        cur_rank_cu_seqlens_k=None,
+        cur_rank_max_seqlen_q=0,
+        cur_rank_max_seqlen_k=0,
+        **kwargs,
+    ):
+        if tensor_layout not in ["HND", "NHD"]:
+            raise NotImplementedError("Tensor layout not supported for CutlassFmha")
+
+        if tensor_layout == "HND":
+            query, key, value = convert_qkv_layout(
+                query, key, value, src_layout="HND", dst_layout="NHD"
+            )
+
+        if attn_mask is not None:
+            raise NotImplementedError("CutlassFmha does not support attn_mask yet")
+
+        # CutlassFmha only supports float16 and bfloat16
+        origin_dtype = query.dtype
+        if query.dtype not in [torch.float16, torch.bfloat16]:
+            query = query.to(torch.float16)
+            key = key.to(torch.float16)
+            value = value.to(torch.float16)
+
+        if cur_rank_cu_seqlens_q is None:
+            qo_segment_offsets = torch.tensor(
+                [0, query.shape[0]], device=query.device, dtype=torch.int32
+            )
+            kv_segment_offsets = torch.tensor(
+                [0, key.shape[0]], device=key.device, dtype=torch.int32
+            )
+            max_qo_len = query.shape[0]
+        else:
+            qo_segment_offsets = cur_rank_cu_seqlens_q
+            kv_segment_offsets = cur_rank_cu_seqlens_k
+            max_qo_len = cur_rank_max_seqlen_q
+
+        output = fmha_varlen(
+            query,
+            key,
+            value,
+            qo_segment_offsets=qo_segment_offsets,
+            kv_segment_offsets=kv_segment_offsets,
+            max_qo_len=max_qo_len,
+            causal=is_causal,
+            sm_scale=1.0 / math.sqrt(query.size(-1)),
+            return_lse=return_lse,
+        )
+
+        lse = None
+        if isinstance(output, tuple):
+            lse = output[1]
+            output = output[0]
+
+        if tensor_layout == "HND":
+            output = convert_output_layout(output, src_layout="NHD", dst_layout="HND")
+            if lse is not None:
+                lse = lse.permute(1, 0)
+
+        if output.dtype != origin_dtype:
+            output = output.to(origin_dtype)
+
+        if return_lse:
+            assert lse is not None, "lse is not returned by cutlass fmha"
+            return output, lse
+        else:
+            return output
diff --git a/flashinfer/parallel_attention/parallel_attention.py b/flashinfer/parallel_attention/parallel_attention.py
new file mode 100644
index 0000000000..eb00ede17b
--- /dev/null
+++ b/flashinfer/parallel_attention/parallel_attention.py
@@ -0,0 +1,117 @@
+import logging
+
+import torch
+
+from .attention_ops import AttentionOpManager
+from .parallel_config import UnevenCPConfig, VarlenCPConfig
+from .parallel_wrapper import ring_wrapper, ulysses_wrapper
+
+logger = logging.getLogger(__name__)
+
+
+class ParallelAttention:
+    """Runs an attention backend with Ulysses and/or Ring parallelism.
+
+    Wraps any registered attention implementation (see :class:`AttentionOpManager`)
+    and transparently applies Ulysses (all-to-all head splitting) and Ring
+    (P2P KV exchange with online softmax merging) parallelism via decorators.
+
+    Args:
+        attn_type: Name of the registered attention backend (e.g. ``"flash-attn3"``).
+        ulysses_group: Ulysses process group.
+        ring_group: Ring process group.
+        uneven_cp_config: Configuration for uneven context parallelism where
+            sequence lengths are not evenly divisible across ranks.
+        varlen_cp_config: Configuration for variable-length context parallelism
+            where multiple sequences of different lengths are packed together.
+        fuse_qkv: If ``True``, fuse Q/K/V into a single all-to-all communication
+            in Ulysses parallelism (reduces 3 NCCL calls to 1).
+
+    Example::
+
+        config = AttnParallelConfig()
+        config.set_config(ulysses_size=2, ring_size=2)
+        attn = ParallelAttention(
+            attn_type="flash-attn3",
+            ulysses_group=ulysses_group,
+            ring_group=ring_group,
+        )
+        output = attn.run(query, key, value, tensor_layout="HND")
+    """
+
+    def __init__(
+        self,
+        attn_type: str,
+        ulysses_group: torch.distributed.ProcessGroup,
+        ring_group: torch.distributed.ProcessGroup,
+        uneven_cp_config: UnevenCPConfig = None,
+        varlen_cp_config: VarlenCPConfig = None,
+        fuse_qkv: bool = False,
+    ):
+        self.attn_type = attn_type
+        self.attn_impl = AttentionOpManager.get_impl(attn_type)
+        self.ulysses_group = ulysses_group
+        self.ring_group = ring_group
+        self.uneven_cp_config = uneven_cp_config
+        self.varlen_cp_config = varlen_cp_config
+        self.fuse_qkv = fuse_qkv
+
+    @ulysses_wrapper
+    @ring_wrapper
+    def run(
+        self,
+        query,
+        key,
+        value,
+        tensor_layout,
+        attn_mask=None,
+        is_causal=False,
+        return_lse=False,
+        cur_rank_cu_seqlens_q=None,
+        cur_rank_cu_seqlens_k=None,
+        cur_rank_max_seqlen_q=0,
+        cur_rank_max_seqlen_k=0,
+        **kwargs,
+    ):
+        """Run parallel attention on the local rank's portion of Q/K/V.
+
+        The Ulysses and Ring wrappers transparently handle communication
+        before and after this method is called.
+
+        Args:
+            query: Query tensor, shape ``[H, S, D]`` (HND) or ``[S, H, D]`` (NHD).
+            key: Key tensor, same layout as query.
+            value: Value tensor, same layout as query.
+            tensor_layout: ``"HND"`` or ``"NHD"``.
+            attn_mask: Optional attention mask (not yet supported).
+            is_causal: Whether to apply causal masking (not yet supported).
+            return_lse: Must be ``False``; internally managed by ring wrapper.
+            cur_rank_cu_seqlens_q/ cur_rank_cu_seqlens_k/
+            cur_rank_max_seqlen_q/ cur_rank_max_seqlen_k:
+            please do not set this manually. This will be set by the parallel wrapper.
+            The sequence lengths should be set in the uneven_cp_config or varlen_cp_config.
+            **kwargs: Additional arguments forwarded to the attention backend.
+
+        Returns:
+            torch.Tensor: Attention output for the local rank, same layout as input.
+        """
+        if is_causal:
+            raise NotImplementedError(
+                "parallel attention does not support causal attention right now"
+            )
+
+        attn_inputs = {
+            "query": query,
+            "key": key,
+            "value": value,
+            "tensor_layout": tensor_layout,
+            "attn_mask": attn_mask,
+            "is_causal": is_causal,
+            "return_lse": return_lse,
+            "cur_rank_cu_seqlens_q": cur_rank_cu_seqlens_q,
+            "cur_rank_cu_seqlens_k": cur_rank_cu_seqlens_k,
+            "cur_rank_max_seqlen_q": cur_rank_max_seqlen_q,
+            "cur_rank_max_seqlen_k": cur_rank_max_seqlen_k,
+        }
+
+        return self.attn_impl(**attn_inputs, **kwargs)
diff --git a/flashinfer/parallel_attention/parallel_config.py b/flashinfer/parallel_attention/parallel_config.py
new file mode 100644
index 0000000000..71eac4c6c8
--- /dev/null
+++ b/flashinfer/parallel_attention/parallel_config.py
@@ -0,0 +1,156 @@
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class UnevenCPConfig:
+    """Configuration for uneven context parallelism.
+
+    Handles the case where the total sequence length is not evenly divisible
+    by the number of ranks. Each rank may hold a different number of tokens,
+    and the last rank typically gets fewer tokens (the remainder).
+
+    The parallel wrappers use this information to truncate padding and zero
+    out extra output positions on the last rank.
+
+    Use the :func:`uneven_cp_config` utility function to compute
+    ``seq_len_cur_ring_group`` via ``all_gather``, then pass the result
+    to this dataclass.
+
+    Attributes:
+        seq_len: Actual (unpadded) total sequence length.
+        seq_len_padded: Padded total sequence length (divisible by
+            ``world_size``).
+        seq_len_cur_ring_group: Tensor of per-rank sequence lengths within
+            the current ring group, shape ``[ring_size]``. ``None`` when
+            ``ring_size == 1`` (no ring parallelism).
+
+    Example::
+
+        # Total sequence length 1023, world_size = 8
+        # Each rank gets ceil(1023/8) = 128 tokens, except last rank gets 127
+        ring_group, ulysses_group = get_parallel_groups(ulysses_size=2, ring_size=4)
+        seq_len_cur_ring_group = uneven_cp_config(
+            seq_len=1023,
+            seq_len_padded=1024,
+            seq_len_cur_rank=128 if rank < 7 else 127,
+            ulysses_group=ulysses_group,
+            ring_group=ring_group,
+        )
+        config = UnevenCPConfig(
+            seq_len=1023,
+            seq_len_padded=1024,
+            seq_len_cur_ring_group=seq_len_cur_ring_group,
+        )
+    """
+
+    seq_len: Optional[int] = None
+    seq_len_padded: Optional[int] = None
+    seq_len_cur_ring_group: Optional[torch.Tensor] = None
+
+    def reset(self):
+        self.seq_len = None
+        self.seq_len_padded = None
+        self.seq_len_cur_ring_group = None
+
+
+@dataclass
+class VarlenCPConfig:
+    """Configuration for variable-length context parallelism.
+
+    Handles the case where multiple sequences of different lengths are packed
+    together (varlen). Cumulative sequence length arrays
+    (``cu_seqlens``) are computed so that the attention kernel can correctly
+    identify sequence boundaries.
+
+    Supports two modes (mutually exclusive):
+
+    - **Ulysses-only** (``ring_size == 1``): The packed sequences are treated
+      as a whole and split across heads via all-to-all. No per-sequence
+      splitting is needed — only the overall ``cu_seqlens`` are stored so the
+      attention kernel knows where each sequence starts and ends.
+    - **Ring-only** (``ulysses_size == 1``): Each individual sequence is split
+      across ranks along the sequence dimension. ``cu_seqlens`` are stored as
+      a 2D tensor of shape ``[ring_size, num_seqs + 1]``, one row per rank,
+      because each rank holds a different slice of every sequence.
+
+    Attributes:
+        cu_seqlens_q_cur_ulysses_group: Cumulative query sequence lengths for
+            the current ulysses group (shared across all ulysses ranks).
+        cu_seqlens_kv_cur_ulysses_group: Cumulative key/value sequence lengths
+            for the current ulysses group.
+        max_seq_len_q_cur_ulysses_group: Max query sequence length in the
+            current ulysses group.
+        max_seq_len_kv_cur_ulysses_group: Max key/value sequence length in the
+            current ulysses group.
+        cu_seqlens_q_cur_ring_group: Cumulative query sequence lengths for all
+            ranks in the current ring group, shape ``[ring_size, num_seqs + 1]``.
+        cu_seqlens_kv_cur_ring_group: Cumulative key/value sequence lengths for
+            all ranks in the current ring group.
+        max_seq_len_q_cur_ring_group: Max query sequence length across all
+            ranks in the ring group (per-rank padded length).
+        max_seq_len_kv_cur_ring_group: Max key/value sequence length across all
+            ranks in the ring group (per-rank padded length).
+    """
+
+    cu_seqlens_q_cur_ulysses_group: Optional[torch.Tensor] = None
+    cu_seqlens_kv_cur_ulysses_group: Optional[torch.Tensor] = None
+    max_seq_len_q_cur_ulysses_group: Optional[int] = None
+    max_seq_len_kv_cur_ulysses_group: Optional[int] = None
+    cu_seqlens_q_cur_ring_group: Optional[torch.Tensor] = None
+    cu_seqlens_kv_cur_ring_group: Optional[torch.Tensor] = None
+    max_seq_len_q_cur_ring_group: Optional[int] = None
+    max_seq_len_kv_cur_ring_group: Optional[int] = None
+
+    def set_varlen_cp_config(
+        self,
+        cu_seqlens_q_all_ranks,
+        cu_seqlens_kv_all_ranks,
+        max_seq_len_q,
+        max_seq_len_kv,
+        ulysses_group,
+        ring_group,
+    ):
+        ring_size = (
+            torch.distributed.get_world_size(ring_group)
+            if ring_group is not None
+            else 1
+        )
+        ulysses_size = (
+            torch.distributed.get_world_size(ulysses_group)
+            if ulysses_group is not None
+            else 1
+        )
+
+        if ring_size == 1:
+            self.cu_seqlens_q_cur_ulysses_group = cu_seqlens_q_all_ranks
+            self.cu_seqlens_kv_cur_ulysses_group = cu_seqlens_kv_all_ranks
+            self.max_seq_len_q_cur_ulysses_group = max_seq_len_q
+            self.max_seq_len_kv_cur_ulysses_group = max_seq_len_kv
+            return
+
+        if ulysses_size == 1:
+            self.cu_seqlens_q_cur_ring_group = cu_seqlens_q_all_ranks
+            self.cu_seqlens_kv_cur_ring_group = cu_seqlens_kv_all_ranks
+            self.max_seq_len_q_cur_ring_group = max_seq_len_q
+            self.max_seq_len_kv_cur_ring_group = max_seq_len_kv
+            return
+
+        raise NotImplementedError(
+            "Varlen CP only supported when ulysses_size == 1 or ring_size == 1"
+        )
+
+    def reset(self):
+        self.cu_seqlens_q_cur_ulysses_group = None
+        self.cu_seqlens_kv_cur_ulysses_group = None
+        self.max_seq_len_q_cur_ulysses_group = None
+        self.max_seq_len_kv_cur_ulysses_group = None
+        self.cu_seqlens_q_cur_ring_group = None
+        self.cu_seqlens_kv_cur_ring_group = None
+        self.max_seq_len_q_cur_ring_group = None
+        self.max_seq_len_kv_cur_ring_group = None
diff --git a/flashinfer/parallel_attention/parallel_wrapper.py b/flashinfer/parallel_attention/parallel_wrapper.py
new file mode 100644
index 0000000000..5bb0c88d40
--- /dev/null
+++ b/flashinfer/parallel_attention/parallel_wrapper.py
@@ -0,0 +1,527 @@
+import logging
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+def all_to_all(tensor, scatter_idx, gather_idx, tensor_layout, group=None):
+    """Perform all-to-all communication on a tensor.
+
+    Args:
+        tensor (torch.Tensor): Input tensor for all-to-all communication
+        scatter_idx (int): Dimension to scatter, will split along this dimension
+            and then scatter to all processes
+        gather_idx (int): Dimension to gather, will gather from all processes
+            and then concatenate along this dimension
+        group (ProcessGroup, optional): Process group to use for communication
+
+    Returns:
+        torch.Tensor
+    """
+    if not dist.is_initialized():
+        return tensor
+
+    world_size = dist.get_world_size(group)
+    if world_size == 1:
+        return tensor
+
+    if scatter_idx == gather_idx:
+        raise ValueError("scatter_idx and gather_idx must be different")
+
+    def chunk_tensor(tensor, scatter_idx):
+        t_shape = list(tensor.shape)
+        if t_shape[scatter_idx] % world_size != 0:
+            raise ValueError(
+                f"Dimension {scatter_idx} of tensor {tensor.shape} "
+                f"must be divisible by world size {world_size}"
+            )
+        chunk_size = t_shape[scatter_idx] // world_size
+        new_shape = []
+        for i in range(len(t_shape)):
+            if i != scatter_idx:
+                new_shape.append(t_shape[i])
+            else:
+                new_shape.extend([world_size, chunk_size])
+        tensor = tensor.reshape(*new_shape)
+        # move scatter_idx to front
+        tensor = tensor.permute(
+            scatter_idx,
+            *[i for i in range(len(new_shape)) if i != scatter_idx],
+        ).contiguous()
+        return tensor
+
+    # chunk tensor for all_to_all
+    tensor = chunk_tensor(tensor, scatter_idx)
+
+    # Perform all2all
+    output = torch.empty_like(tensor)
+    dist.all_to_all_single(output, tensor, group=group)
+
+    # output: e.g., [world_size, chunked_H, chunked_S, D]
+    # if scatter_idx == 0, gather_idx == 1 -> [chunked_H, S, D]
+    def reorder_tensor(tensor, gather_idx):
+        t_shape = list(tensor.shape)
+        world_size = t_shape[0]
+        # insert front to gather_idx + 1
+        permute_idx = []
+        for i in range(1, len(t_shape)):
+            if i != gather_idx + 1:
+                permute_idx.append(i)
+            else:
+                permute_idx.extend([0, i])
+        tensor = tensor.permute(*permute_idx).contiguous()
+
+        # reshape tensor
+        new_shape = []
+        for i in range(1, len(t_shape)):
+            if i != gather_idx + 1:
+                new_shape.append(t_shape[i])
+            else:
+                new_shape.append(world_size * t_shape[i])
+
+        tensor = tensor.reshape(*new_shape)
+
+        return tensor
+
+    output = reorder_tensor(output, gather_idx)
+
+    return output
+
+
+def ulysses_a2a_in(
+    query,
+    key,
+    value,
+    attn_mask,
+    tensor_layout,
+    ulysses_size=1,
+    ulysses_rank=0,
+    ulysses_group=None,
+    fuse_qkv=False,
+):
+    if ulysses_size == 1:
+        return query, key, value, attn_mask
+
+    if attn_mask is not None:
+        raise NotImplementedError("Attn mask not supported for ulysses_a2a_in")
+
+    if tensor_layout == "HND":
+        scatter_idx = 0
+        gather_idx = 1
+    elif tensor_layout == "NHD":
+        scatter_idx = 1
+        gather_idx = 0
+    else:
+        raise ValueError(f"Invalid tensor layout: {tensor_layout}")
+
+    # [H, S/N, D] -> [H/N, S, D]
+    if fuse_qkv:
+        # Fused communication: concatenate q/k/v into [3, H, S/N, D],
+        # single all-to-all, then split.
+        # This reduces 3 NCCL calls to 1, improving efficiency.
+        query = torch.unsqueeze(query, 0)
+        key = torch.unsqueeze(key, 0)
+        value = torch.unsqueeze(value, 0)
+        qkv = torch.cat([query, key, value], dim=0)
+        qkv = all_to_all(
+            qkv,
+            scatter_idx=scatter_idx + 1,
+            gather_idx=gather_idx + 1,
+            group=ulysses_group,
+            tensor_layout=tensor_layout,
+        )
+        query, key, value = torch.chunk(qkv, 3, dim=0)
+        query = query.squeeze(0)
+        key = key.squeeze(0)
+        value = value.squeeze(0)
+    else:
+        # Independent communication: 3 separate all-to-all operations (default, safe)
+        query = all_to_all(
+            query,
+            scatter_idx=scatter_idx,
+            gather_idx=gather_idx,
+            group=ulysses_group,
+            tensor_layout=tensor_layout,
+        )
+        key = all_to_all(
+            key,
+            scatter_idx=scatter_idx,
+            gather_idx=gather_idx,
+            group=ulysses_group,
+            tensor_layout=tensor_layout,
+        )
+        value = all_to_all(
+            value,
+            scatter_idx=scatter_idx,
+            gather_idx=gather_idx,
+            group=ulysses_group,
+            tensor_layout=tensor_layout,
+        )
+
+    return query, key, value, attn_mask
+
+
+def ulysses_a2a_out(output, tensor_layout, ulysses_size=1, ulysses_group=None):
+    if ulysses_size == 1:
+        return output
+
+    assert tensor_layout in ["NHD", "HND"], (
+        f"tensor_layout must be NHD or HND, but got {tensor_layout}"
+    )
+    if tensor_layout == "HND":
+        scatter_idx = 1
+        gather_idx = 0
+    elif tensor_layout == "NHD":
+        scatter_idx = 0
+        gather_idx = 1
+    else:
+        raise ValueError(f"Invalid tensor layout: {tensor_layout}")
+    # [H/N, S, D] -> [H, S/N, D]
+    output = all_to_all(
+        output,
+        scatter_idx=scatter_idx,
+        gather_idx=gather_idx,
+        tensor_layout=tensor_layout,
+        group=ulysses_group,
+    )
+    return output
+
+
+def ring_fwd_out_correction(
+    out: torch.Tensor,
+    out_per_step: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    softmax_lse_per_step: torch.Tensor,
+):
+    """Merge partial outputs of each step in ring attention"""
+    new_out = out - F.sigmoid(
+        softmax_lse_per_step.unsqueeze(-1) - softmax_lse.unsqueeze(-1)
+    ) * (out - out_per_step)
+    out.copy_(new_out)
+
+
+def ring_fwd_softmax_lse_correction(
+    softmax_lse: torch.Tensor,
+    softmax_lse_per_step: torch.Tensor,
+):
+    """Merge softmax stats of each step in ring attention"""
+    new_lse = softmax_lse - F.logsigmoid(softmax_lse - softmax_lse_per_step)
+    softmax_lse.copy_(new_lse)
+
+
+def ring_attn_p2p_communicate(
+    rank, send_tensor, send_dst, recv_tensor, recv_src, ring_group
+):
+    """Point-to-point communications of KV and dKV in ring attention"""
+    send_recv_ops = []
+    if rank % 2 == 0:
+        send_op = torch.distributed.P2POp(
+            torch.distributed.isend,
+            send_tensor,
+            group_peer=send_dst,
+            group=ring_group,
+        )
+        recv_op = torch.distributed.P2POp(
+            torch.distributed.irecv,
+            recv_tensor,
+            group_peer=recv_src,
+            group=ring_group,
+        )
+        send_recv_ops.append(send_op)
+        send_recv_ops.append(recv_op)
+    else:
+        recv_op = torch.distributed.P2POp(
+            torch.distributed.irecv,
+            recv_tensor,
+            group_peer=recv_src,
+            group=ring_group,
+        )
+        send_op = torch.distributed.P2POp(
+            torch.distributed.isend,
+            send_tensor,
+            group_peer=send_dst,
+            group=ring_group,
+        )
+        send_recv_ops.append(recv_op)
+        send_recv_ops.append(send_op)
+    send_recv_reqs = torch.distributed.batch_isend_irecv(send_recv_ops)
+
+    return send_recv_reqs
+
+
+def ulysses_wrapper(func):
+    def wrapper(self, query, key, value, tensor_layout, attn_mask=None, **kwargs):
+        ulysses_group = self.ulysses_group
+        ring_group = self.ring_group
+        uneven_cp_config = self.uneven_cp_config
+        varlen_cp_config = self.varlen_cp_config
+
+        ulysses_size = (
+            dist.get_world_size(ulysses_group) if ulysses_group is not None else 1
+        )
+        ring_size = dist.get_world_size(ring_group) if ring_group is not None else 1
+
+        if kwargs.get("return_lse", False):
+            raise ValueError("return_lse=True is not supported in parallel attention")
+
+        if ulysses_size == 1:
+            return func(self, query, key, value, tensor_layout, attn_mask, **kwargs)
+
+        ulysses_rank = dist.get_rank(ulysses_group)
+
+        assert tensor_layout in ["NHD", "HND"], (
+            f"tensor_layout must be NHD or HND, but got {tensor_layout}"
+        )
+        if tensor_layout == "HND":
+            seq_dim = 1
+            head_dim = 0
+        elif tensor_layout == "NHD":
+            seq_dim = 0
+            head_dim = 1
+        else:
+            raise ValueError(f"Invalid tensor layout: {tensor_layout}")
+
+        if query.shape[head_dim] % ulysses_size != 0:
+            raise ValueError(
+                f"Head dim {head_dim} of query {query.shape} "
+                f"must be divisible by ulysses size {ulysses_size}"
+            )
+        if key.shape[head_dim] % ulysses_size != 0:
+            raise ValueError(
+                f"Head dim {head_dim} of key {key.shape} "
+                f"must be divisible by ulysses size {ulysses_size}"
+            )
+        if value.shape[head_dim] % ulysses_size != 0:
+            raise ValueError(
+                f"Head dim {head_dim} of value {value.shape} "
+                f"must be divisible by ulysses size {ulysses_size}"
+            )
+
+        # Apply ulysses_a2a_in before the function call
+        query, key, value, attn_mask = ulysses_a2a_in(
+            query,
+            key,
+            value,
+            attn_mask,
+            tensor_layout,
+            ulysses_size=ulysses_size,
+            ulysses_rank=ulysses_rank,
+            ulysses_group=ulysses_group,
+            fuse_qkv=self.fuse_qkv,
+        )
+
+        # truncate and pad if cp is uneven
+        truncate_and_pad = uneven_cp_config is not None
+
+        if ring_size == 1 and truncate_and_pad:
+            # there is no ring, so we can use uneven_cp_config.seq_len
+            # to do truncate and pad
+            seq_len = uneven_cp_config.seq_len
+
+            # Truncate key and value tensors using torch.narrow
+            key = torch.narrow(key, seq_dim, 0, seq_len).contiguous()
+            value = torch.narrow(value, seq_dim, 0, seq_len).contiguous()
+
+        if ring_size == 1 and varlen_cp_config is not None:
+            cu_seqlens_q = varlen_cp_config.cu_seqlens_q_cur_ulysses_group
+            cu_seqlens_kv = varlen_cp_config.cu_seqlens_kv_cur_ulysses_group
+            kwargs["cur_rank_cu_seqlens_q"] = cu_seqlens_q
+            kwargs["cur_rank_cu_seqlens_k"] = cu_seqlens_kv
+            kwargs["cur_rank_max_seqlen_q"] = (
+                varlen_cp_config.max_seq_len_q_cur_ulysses_group
+            )
+            kwargs["cur_rank_max_seqlen_k"] = (
+                varlen_cp_config.max_seq_len_kv_cur_ulysses_group
+            )
+
+            if key.shape[seq_dim] != cu_seqlens_kv[-1]:
+                # Truncate kv_inputs using torch.narrow
+                key = torch.narrow(key, seq_dim, 0, cu_seqlens_kv[-1])
+                value = torch.narrow(value, seq_dim, 0, cu_seqlens_kv[-1])
+
+        # Call the original function
+        result = func(self, query, key, value, tensor_layout, attn_mask, **kwargs)
+
+        # if ring size is 1, return_lse is false, result only has output.
+        if ring_size == 1 and truncate_and_pad and result.shape[seq_dim] > seq_len:
+            # Zero out padding using torch.narrow
+            padding_part = torch.narrow(
+                result, seq_dim, seq_len, result.shape[seq_dim] - seq_len
+            )
+            padding_part.zero_()
+
+        if (
+            ring_size == 1
+            and varlen_cp_config is not None
+            and result.shape[seq_dim]
+            > varlen_cp_config.cu_seqlens_q_cur_ulysses_group[-1]
+        ):
+            # Zero out padding using torch.narrow
+            cu_end = varlen_cp_config.cu_seqlens_q_cur_ulysses_group[-1]
+            padding_part = torch.narrow(
+                result, seq_dim, cu_end, result.shape[seq_dim] - cu_end
+            )
+            padding_part.zero_()
+
+        result = ulysses_a2a_out(
+            result,
+            tensor_layout,
+            ulysses_size=ulysses_size,
+            ulysses_group=ulysses_group,
+        )
+
+        return result
+
+    return wrapper
+
+
+def get_kv_rank(ring_size, ring_rank, cur_iter):
+    # get the the source rank of kv tensor in current iter
+    return (ring_size + ring_rank - cur_iter) % ring_size
+
+
+def ring_wrapper(func):
+    def wrapper(self, query, key, value, tensor_layout, attn_mask=None, **kwargs):
+        ring_group = self.ring_group
+        uneven_cp_config = self.uneven_cp_config
+        varlen_cp_config = self.varlen_cp_config
+
+        ring_size = dist.get_world_size(ring_group) if ring_group is not None else 1
+
+        if ring_size == 1:
+            return func(self, query, key, value, tensor_layout, attn_mask, **kwargs)
+
+        rank = dist.get_rank(ring_group)
+        send_dst = (rank + 1) % ring_size
+        recv_src = (rank - 1) % ring_size
+
+        # Determine sequence dimension based on tensor layout
+        if tensor_layout == "HND":
+            seq_dim = 1  # query, key, value shape are: [H, S, D], so seq_dim is 1
+        elif tensor_layout == "NHD":
+            seq_dim = 0  # query, key, value shape are: [S, H, D], so seq_dim is 0
+        else:
+            raise ValueError(f"Invalid tensor layout: {tensor_layout}")
+
+        p2p_comm_buffers = [None, None]
+        p2p_comm_buffers[0] = torch.cat((key.unsqueeze(0), value.unsqueeze(0)), dim=0)
+        send_recv_reqs = [[], []]
+
+        out = None
+        softmax_lse = None
+        for i in range(ring_size):
+            kv_rank = get_kv_rank(ring_size, rank, i)
+            # wait until KV is received
+            for req in send_recv_reqs[(i + 1) % 2]:
+                req.wait()
+
+            if i < (ring_size - 1):
+                p2p_comm_buffers[(i + 1) % 2] = torch.empty_like(
+                    p2p_comm_buffers[i % 2]
+                )
+                send_recv_reqs[i % 2] = ring_attn_p2p_communicate(
+                    rank,
+                    p2p_comm_buffers[i % 2],
+                    send_dst,
+                    p2p_comm_buffers[(i + 1) % 2],
+                    recv_src,
+                    ring_group,
+                )
+            kv_inputs = p2p_comm_buffers[i % 2]
+
+            # do truncate and pad if cp is uneven,
+            if uneven_cp_config is not None:
+                # seq_dim+1 because kv_inputs is concated to
+                # [2, H, S, D] or [2, S, H, D]
+                if (
+                    kv_inputs.shape[seq_dim + 1]
+                    != uneven_cp_config.seq_len_cur_ring_group[kv_rank]
+                ):
+                    # Truncate kv_inputs using torch.narrow
+                    kv_inputs = torch.narrow(
+                        kv_inputs,
+                        seq_dim + 1,
+                        0,
+                        uneven_cp_config.seq_len_cur_ring_group[kv_rank],
+                    )
+
+            if varlen_cp_config is not None:
+                cu_seqlens_q = varlen_cp_config.cu_seqlens_q_cur_ring_group[rank]
+                cu_seqlens_kv = varlen_cp_config.cu_seqlens_kv_cur_ring_group[kv_rank]
+                kwargs["cur_rank_cu_seqlens_q"] = cu_seqlens_q
+                kwargs["cur_rank_cu_seqlens_k"] = cu_seqlens_kv
+
+                if kv_inputs.shape[seq_dim + 1] != cu_seqlens_kv[-1]:
+                    # Truncate kv_inputs using torch.narrow
+                    kv_inputs = torch.narrow(
+                        kv_inputs, seq_dim + 1, 0, cu_seqlens_kv[-1]
+                    )
+
+                kwargs["cur_rank_max_seqlen_q"] = (
+                    varlen_cp_config.max_seq_len_q_cur_ring_group
+                )
+                kwargs["cur_rank_max_seqlen_k"] = (
+                    varlen_cp_config.max_seq_len_kv_cur_ring_group
+                )
+
+            kwargs["return_lse"] = True
+            # we need this line because a bug in flash-attn4
+            # https://github.com/Dao-AILab/flash-attention/pull/1793
+            with torch.cuda.device(query.device.index):
+                block_out = func(
+                    self,
+                    query,
+                    kv_inputs[0],
+                    kv_inputs[1],
+                    tensor_layout,
+                    attn_mask,
+                    **kwargs,
+                )
+
+            out_per_step = block_out[0]
+            softmax_lse_per_step = block_out[1]
+
+            if i == 0:
+                softmax_lse = torch.clone(softmax_lse_per_step).to(torch.float)
+                out = torch.clone(out_per_step)
+            else:
+                ring_fwd_out_correction(
+                    out, out_per_step, softmax_lse, softmax_lse_per_step
+                )
+                ring_fwd_softmax_lse_correction(softmax_lse, softmax_lse_per_step)
+
+        # Determine output sequence dimension based on tensor layout
+        # (for output tensor)
+        if tensor_layout == "HND":
+            out_seq_dim = 1  # out is [H, S, D], so seq_dim is 1
+        elif tensor_layout == "NHD":
+            out_seq_dim = 0  # out is [S, H, D], so seq_dim is 0
+        else:
+            raise ValueError(f"Invalid tensor layout: {tensor_layout}")
+
+        start_pos = out.shape[out_seq_dim]
+
+        if (
+            uneven_cp_config is not None
+            and out.shape[out_seq_dim] > uneven_cp_config.seq_len_cur_ring_group[rank]
+        ):
+            start_pos = uneven_cp_config.seq_len_cur_ring_group[rank]
+
+        if (
+            varlen_cp_config is not None
+            and out.shape[out_seq_dim]
+            > varlen_cp_config.cu_seqlens_q_cur_ring_group[rank][-1]
+        ):
+            start_pos = varlen_cp_config.cu_seqlens_q_cur_ring_group[rank][-1]
+
+        if start_pos < out.shape[out_seq_dim]:
+            padding_length = out.shape[out_seq_dim] - start_pos
+            padding_part = torch.narrow(out, out_seq_dim, start_pos, padding_length)
+            padding_part.zero_()
+
+        return out
+
+    return wrapper
diff --git a/flashinfer/parallel_attention/utils.py b/flashinfer/parallel_attention/utils.py
new file mode 100644
index 0000000000..4e603112b6
--- /dev/null
+++ b/flashinfer/parallel_attention/utils.py
@@ -0,0 +1,445 @@
+import logging
+import torch
+from torch.distributed.device_mesh import init_device_mesh
+
+logger = logging.getLogger(__name__)
+
+
+def convert_qkv_layout(q, k, v, src_layout, dst_layout):
+    if src_layout == "HND" and dst_layout == "NHD":
+        # [H, S, D] -> [S, H, D]
+        q = q.permute(1, 0, 2).contiguous()
+        k = k.permute(1, 0, 2).contiguous()
+        v = v.permute(1, 0, 2).contiguous()
+    elif src_layout == "NHD" and dst_layout == "HND":
+        # [S, H, D] -> [H, S, D]
+        q = q.permute(1, 0, 2).contiguous()
+        k = k.permute(1, 0, 2).contiguous()
+        v = v.permute(1, 0, 2).contiguous()
+    else:
+        raise NotImplementedError(
+            f"Unsupported tensor layout conversion: {src_layout} -> {dst_layout}"
+        )
+    return q, k, v
+
+
+def convert_output_layout(out, src_layout, dst_layout):
+    if src_layout == "HND" and dst_layout == "NHD":
+        # [S, H, D] -> [H, S, D]
+        out = out.permute(1, 0, 2).contiguous()
+    elif src_layout == "NHD" and dst_layout == "HND":
+        # [H, S, D] -> [S, H, D]
+        out = out.permute(1, 0, 2).contiguous()
+    else:
+        raise NotImplementedError(
+            f"Unsupported tensor layout conversion: {src_layout} -> {dst_layout}"
+        )
+    return out
+
+
+def split_varlen_input(tensor, seq_len_list, world_size, rank, tensor_layout="HND"):
+    """Split a packed variable-length tensor across ranks for context parallelism.
+
+    Given a tensor whose sequence dimension is the concatenation of multiple
+    sub-sequences, split each sub-sequence into ``world_size`` chunks and return
+    the ``rank``-th chunk concatenated together. The first ``world_size - 1``
+    ranks each get ``ceil(seq_len / world_size)`` tokens per sub-sequence;
+    the last rank gets the remainder. The result is zero-padded so that all
+    ranks have the same total sequence length.
+
+    Args:
+        tensor: Input tensor of shape ``[H, total_seq_len, D]`` (HND) or
+            ``[total_seq_len, H, D]`` (NHD).
+        seq_len_list: Individual sequence lengths that sum to ``total_seq_len``,
+            e.g. ``[1021, 1024, 1027]``. Can be a list, tuple, or torch.Tensor.
+        world_size: Number of ranks to split across.
+        rank: Which rank's chunk to return (0-indexed).
+        tensor_layout: ``"HND"`` or ``"NHD"``.
+
+    Returns:
+        torch.Tensor: The rank's chunk, zero-padded to uniform length across ranks.
+    """
+    if not isinstance(seq_len_list, torch.Tensor):
+        seq_len_list = torch.tensor(seq_len_list, dtype=torch.int32)
+
+    if tensor_layout == "NHD":
+        chunk_dim = 0
+    elif tensor_layout == "HND":
+        chunk_dim = 1
+    else:
+        raise ValueError(f"Invalid tensor layout: {tensor_layout}")
+
+    seq_len_padded = (seq_len_list + world_size - 1) // world_size * world_size
+    total_seq_len_padded = sum(seq_len_padded)
+    seq_len_padded_cur_rank = (
+        (total_seq_len_padded + world_size - 1) // world_size
+    ).to(torch.int32)
+
+    chunks = []
+    offset = 0
+    for seq_len in seq_len_list:
+        seq_len = int(seq_len)
+        # First (world_size - 1) ranks get ceil(seq_len / world_size),
+        # last rank gets whatever is left.
+        base = (seq_len + world_size - 1) // world_size
+        if rank < world_size - 1:
+            chunk_len = base
+            start = offset + base * rank
+        else:
+            # Last rank gets the remainder
+            start = offset + base * (world_size - 1)
+            chunk_len = seq_len - base * (world_size - 1)
+
+        chunks.append(tensor.narrow(chunk_dim, start, chunk_len))
+        offset += seq_len
+
+    res = torch.cat(chunks, dim=chunk_dim)
+
+    if res.shape[chunk_dim] < seq_len_padded_cur_rank:
+        pad_len = seq_len_padded_cur_rank - res.shape[chunk_dim]
+        pad_shape = list(res.shape)
+        pad_shape[chunk_dim] = pad_len
+        res = torch.cat(
+            [res, torch.zeros(pad_shape, device=res.device, dtype=res.dtype)],
+            dim=chunk_dim,
+        )
+
+    return res
+
+
+def ulysses_varlen_config(seq_lens_q, seq_lens_kv):
+    """Compute cumulative sequence lengths for Ulysses-only variable-length parallelism.
+
+    In Ulysses-only mode (``ring_size == 1``), the packed sequences are treated
+    as a whole and split across heads via all-to-all. This function builds the
+    ``cu_seqlens`` arrays that the attention kernel needs to locate sequence
+    boundaries within the packed input.
+
+    Args:
+        seq_lens_q: Per-sequence query lengths, e.g. ``[1021, 2048, 512]``.
+            Can be a list, tuple, or torch.Tensor.
+        seq_lens_kv: Per-sequence key/value lengths, same format as
+            ``seq_lens_q``.
+
+    Returns:
+        Tuple of four elements:
+
+        - **cu_seqlens_q** (*torch.Tensor*): Cumulative query sequence lengths
+          of shape ``[num_seqs + 1]``, starting with 0.
+        - **cu_seqlens_kv** (*torch.Tensor*): Cumulative key/value sequence
+          lengths, same shape.
+        - **max_seqlen_q** (*int*): Maximum query sequence length.
+        - **max_seqlen_kv** (*int*): Maximum key/value sequence length.
+    """
+    rank = torch.distributed.get_rank()
+    device = torch.device(f"cuda:{rank}")
+
+    cu_seqlens_q = (
+        torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32),
+                torch.cumsum(torch.tensor(seq_lens_q, dtype=torch.int32), dim=0),
+            ]
+        )
+        .to(device)
+        .to(torch.int32)
+    )
+
+    cu_seqlens_kv = (
+        torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32),
+                torch.cumsum(torch.tensor(seq_lens_kv, dtype=torch.int32), dim=0),
+            ]
+        )
+        .to(device)
+        .to(torch.int32)
+    )
+
+    max_seqlen_q = max(seq_lens_q)
+    max_seqlen_k = max(seq_lens_kv)
+
+    return cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_k
+
+
+def ring_varlen_config(seq_lens_q, seq_lens_kv, ring_group):
+    """Compute per-rank cumulative sequence lengths for Ring-only variable-length parallelism.
+
+    In Ring-only mode (``ulysses_size == 1``), each individual sequence is
+    split along the sequence dimension across ring ranks. Each rank holds a
+    chunk of every sequence, so ``cu_seqlens`` are stored as a 2-D tensor of
+    shape ``[ring_size, num_seqs + 1]`` — one row per rank — because each
+    rank's chunk has different per-sequence lengths (the last rank gets the
+    remainder after padding).
+
+    Sequences are padded to be divisible by ``ring_size`` so that the first
+    ``ring_size - 1`` ranks each get ``ceil(seq_len / ring_size)`` tokens per
+    sequence, and the last rank gets the remainder.
+
+    Args:
+        seq_lens_q: Per-sequence query lengths, e.g. ``[1021, 1024, 1027]``.
+            Can be a list, tuple, or torch.Tensor.
+        seq_lens_kv: Per-sequence key/value lengths, same format as
+            ``seq_lens_q``.
+        ring_group: The ring attention process group, used to determine
+            ``ring_size`` via ``torch.distributed.get_world_size(ring_group)``.
+
+    Returns:
+        Tuple of four elements:
+
+        - **cu_seqlens_q_all_ranks** (*torch.Tensor*): Cumulative query
+          sequence lengths for all ranks, shape ``[ring_size, num_seqs + 1]``.
+        - **cu_seqlens_kv_all_ranks** (*torch.Tensor*): Cumulative key/value
+          sequence lengths, same shape.
+        - **max_seq_len_q** (*torch.Tensor*): Maximum per-rank padded query
+          sequence length.
+        - **max_seq_len_kv** (*torch.Tensor*): Maximum per-rank padded
+          key/value sequence length.
+
+    Example::
+
+        # seq_lens_q = [1021, 1024, 1027], ring_size = 4
+        #
+        # Padding: 1021 -> 1024, 1024 -> 1024, 1027 -> 1028
+        # Per-rank: 256,          256,          257
+        # Last rank gets remainder for each sequence:
+        #   1021: 256 - (1024-1021) = 253
+        #   1024: 256 - (1024-1024) = 256
+        #   1027: 257 - (1028-1027) = 256
+        #
+        # cu_seqlens_q_all_ranks (shape [4, 4]):
+        #   rank 0: [0, 256, 512, 769]
+        #   rank 1: [0, 256, 512, 769]
+        #   rank 2: [0, 256, 512, 769]
+        #   rank 3: [0, 253, 509, 765]  (last rank, shorter chunks)
+    """
+    if not isinstance(seq_lens_q, torch.Tensor):
+        seq_lens_q = torch.tensor(seq_lens_q, dtype=torch.int32)
+    if not isinstance(seq_lens_kv, torch.Tensor):
+        seq_lens_kv = torch.tensor(seq_lens_kv, dtype=torch.int32)
+
+    world_size = (
+        torch.distributed.get_world_size(ring_group) if ring_group is not None else 1
+    )
+
+    rank = torch.distributed.get_rank()
+    device = torch.device(f"cuda:{rank}")
+
+    padded_seq_lens_q = (seq_lens_q + world_size - 1) // world_size * world_size
+    padded_seq_lens_kv = (seq_lens_kv + world_size - 1) // world_size * world_size
+
+    padded_seq_len_q_cur_rank = padded_seq_lens_q // world_size
+    padded_seq_len_kv_cur_rank = padded_seq_lens_kv // world_size
+
+    max_seq_len_q = padded_seq_len_q_cur_rank.max()
+    max_seq_len_kv = padded_seq_len_kv_cur_rank.max()
+
+    cu_seqlens_q_all_ranks = []
+    cu_seqlens_kv_all_ranks = []
+
+    for i in range(world_size):
+        if i == world_size - 1:
+            seq_len_q_cur_rank = padded_seq_len_q_cur_rank - (
+                padded_seq_lens_q - seq_lens_q
+            )
+            seq_len_kv_cur_rank = padded_seq_len_kv_cur_rank - (
+                padded_seq_lens_kv - seq_lens_kv
+            )
+        else:
+            seq_len_q_cur_rank = padded_seq_len_q_cur_rank
+            seq_len_kv_cur_rank = padded_seq_len_kv_cur_rank
+
+        cu_seqlens_q = (
+            torch.cat(
+                [
+                    torch.zeros(1),
+                    torch.cumsum(seq_len_q_cur_rank, dim=0),
+                ]
+            )
+            .to(device)
+            .to(torch.int32)
+        )
+        cu_seqlens_q_all_ranks.append(cu_seqlens_q)
+
+        cu_seqlens_kv = (
+            torch.cat(
+                [
+                    torch.zeros(1),
+                    torch.cumsum(seq_len_kv_cur_rank, dim=0),
+                ]
+            )
+            .to(device)
+            .to(torch.int32)
+        )
+        cu_seqlens_kv_all_ranks.append(cu_seqlens_kv)
+
+    return (
+        torch.stack(cu_seqlens_q_all_ranks),
+        torch.stack(cu_seqlens_kv_all_ranks),
+        max_seq_len_q.item(),
+        max_seq_len_kv.item(),
+    )
+
+
+def uneven_cp_config(
+    seq_len,
+    seq_len_padded,
+    seq_len_cur_rank,
+    ulysses_group=None,
+    ring_group=None,
+):
+    """Gather per-rank sequence lengths and compute the current ring group sequence length.
+
+    Args:
+        seq_len: Actual (unpadded) total sequence length.
+        seq_len_padded: Padded total sequence length (divisible by world_size).
+        seq_len_cur_rank: Number of real (non-padding) tokens on this rank. for example, if the total sequence
+        length is 1023 and the world size is 8, and the rank is 0, then seq_len_cur_rank is 128. If the rank is 7, then seq_len_cur_rank is 127.
+        ulysses_group: Ulysses process group.
+        ring_group: Ring process group.
+    """
+
+    rank = torch.distributed.get_rank()
+    device = torch.device(f"cuda:{rank}")
+
+    seq_len_cur_rank = torch.tensor(
+        [seq_len_cur_rank], dtype=torch.int32, device=device
+    )
+    gather_list = [
+        torch.empty_like(seq_len_cur_rank)
+        for _ in range(torch.distributed.get_world_size())
+    ]
+    torch.distributed.all_gather(gather_list, seq_len_cur_rank)
+    seq_len_all_ranks = torch.cat(gather_list, dim=0).cpu()
+
+    ring_size = (
+        torch.distributed.get_world_size(ring_group) if ring_group is not None else 1
+    )
+    ulysses_size = (
+        torch.distributed.get_world_size(ulysses_group)
+        if ulysses_group is not None
+        else 1
+    )
+
+    if ring_size == 1:
+        return
+
+    if ulysses_size == 1:
+        ring_ranks = torch.distributed.get_process_group_ranks(ring_group)
+        seq_len_cur_ring_group = seq_len_all_ranks[torch.tensor(ring_ranks)]
+        return seq_len_cur_ring_group
+
+    ulysses_ranks = torch.distributed.get_process_group_ranks(ulysses_group)
+    seq_len_cur_ulysses_group = seq_len_all_ranks[torch.tensor(ulysses_ranks)]
+    ring_seq_cur_ring_rank = torch.sum(seq_len_cur_ulysses_group, dtype=torch.int32).to(
+        device
+    )
+    gather_list = [
+        torch.empty(1, dtype=torch.int32, device=device) for _ in range(ring_size)
+    ]
+
+    torch.distributed.all_gather(
+        gather_list,
+        ring_seq_cur_ring_rank,
+        group=ring_group,
+    )
+    seq_len_cur_ring_group = torch.cat(gather_list, dim=0)
+
+    return seq_len_cur_ring_group
+
+
+def get_parallel_groups(
+    ulysses_size: int,
+    ring_size: int,
+    device_type: str = "cuda",
+):
+    """Create a device mesh and return the Ring and Ulysses process groups.
+
+    Builds a ``DeviceMesh`` with up to three dimensions — ``redundant``
+    (when ``world_size > ulysses_size * ring_size``), ``ring``, and
+    ``ulysses`` — and extracts the corresponding process groups for use
+    in :class:`ParallelAttention`.
+
+    Args:
+        ulysses_size: Ulysses parallel degree (number of ranks that
+            participate in all-to-all head splitting).
+        ring_size: Ring attention parallel degree (number of ranks that
+            exchange KV chunks in a ring).
+        device_type: Device type for the mesh, defaults to ``"cuda"``.
+
+    Returns:
+        Tuple of two elements:
+
+        - **ring_group** (*Optional[ProcessGroup]*): The ring attention
+          process group, or ``None`` if ``ring_size == 1``.
+        - **ulysses_group** (*Optional[ProcessGroup]*): The Ulysses
+          process group, or ``None`` if ``ulysses_size == 1``.
+
+    Raises:
+        ValueError: If ``world_size`` is not divisible by
+            ``ulysses_size * ring_size``.
+
+    Note:
+        The device mesh dimensions are created in the order:
+        ``redundant`` (if needed) → ``ring`` → ``ulysses``.
+    """
+
+    total_parallel_size = ulysses_size * ring_size
+    world_size = torch.distributed.get_world_size()
+    if world_size % total_parallel_size != 0:
+        raise ValueError(
+            f"World size ({world_size}) is not divisible by "
+            f"total parallel size ({total_parallel_size})"
+        )
+
+    logger.debug(
+        f"Setting up device mesh with total parallel size: {total_parallel_size} "
+        f"ulysses_size: {ulysses_size}, ring_size: {ring_size}"
+    )
+
+    if total_parallel_size == 1:
+        logger.debug("No parallelism needed, skipping device mesh setup")
+        return None, None
+
+    mesh_dims = []
+    mesh_sizes = []
+
+    if world_size != total_parallel_size:
+        mesh_dims.append("redundant")
+        mesh_sizes.append(world_size // total_parallel_size)
+        logger.debug(
+            f"Added redundant dimension: "
+            f"{world_size // total_parallel_size}, "
+            f"world_size={world_size}, "
+            f"total_parallel_size={total_parallel_size}"
+        )
+
+    if ring_size > 1:
+        mesh_dims.append("ring")
+        mesh_sizes.append(ring_size)
+        logger.debug(f"Added Ring dimension: {ring_size}")
+    if ulysses_size > 1:
+        mesh_dims.append("ulysses")
+        mesh_sizes.append(ulysses_size)
+        logger.debug(f"Added Ulysses dimension: {ulysses_size}")
+
+    if not mesh_dims:
+        logger.debug("No mesh dimensions needed")
+        return None, None
+    else:
+        logger.info(f"Creating device mesh: dims={mesh_dims}, sizes={mesh_sizes}")
+        device_mesh = init_device_mesh(
+            device_type,
+            tuple(mesh_sizes),
+            mesh_dim_names=tuple(mesh_dims),
+        )
+        logger.info("Device mesh created successfully")
+
+    ring_group = None
+    ulysses_group = None
+    if ring_size > 1:
+        ring_group = device_mesh.get_group("ring")
+    if ulysses_size > 1:
+        ulysses_group = device_mesh.get_group("ulysses")
+
+    return ring_group, ulysses_group
diff --git a/scripts/task_test_multi_gpu_comm_kernels.sh b/scripts/task_test_multi_gpu_comm_kernels.sh
index 94f2761dc8..e2e261c129 100755
--- a/scripts/task_test_multi_gpu_comm_kernels.sh
+++ b/scripts/task_test_multi_gpu_comm_kernels.sh
@@ -24,6 +24,10 @@ source "${SCRIPT_DIR}/test_utils.sh"
 # Add others back once they are fixed
 TEST_FILES="tests/comm/test_allreduce_unified_api.py"
 
+# Tests that require torchrun instead of mpirun
+TORCHRUN_TEST_FILES="tests/attention/test_parallel_attention.py"
+: "${TORCHRUN_PREFIX:=torchrun --nproc_per_node=4}"
+
 # Main execution
 main() {
     # Parse command line arguments
@@ -49,6 +53,28 @@ main() {
         execute_tests "$TEST_FILES"
     fi
 
+    # Execute torchrun tests (torchrun requires -m pytest, not direct pytest invocation)
+    echo "Multi-GPU torchrun test files:"
+    for test_file in $TORCHRUN_TEST_FILES; do
+        echo "  $test_file"
+    done
+    echo ""
+
+    for test_file in $TORCHRUN_TEST_FILES; do
+        echo "=========================================="
+        echo "Running: ${TORCHRUN_PREFIX} -m pytest ${test_file} -v"
+        echo "=========================================="
+        if [ "$DRY_RUN" != "true" ]; then
+            if ${TORCHRUN_PREFIX} -m pytest "${test_file}" -v; then
+                echo "PASSED: $test_file"
+            else
+                echo "FAILED: $test_file"
+                EXIT_CODE=1
+            fi
+        fi
+        echo ""
+    done
+
     exit "$EXIT_CODE"
 }
 
diff --git a/tests/attention/test_parallel_attention.py b/tests/attention/test_parallel_attention.py
new file mode 100644
index 0000000000..acd21db5bc
--- /dev/null
+++ b/tests/attention/test_parallel_attention.py
@@ -0,0 +1,394 @@
+"""Tests for parallel attention (Ulysses + Ring).
+
+Launch with:
+    torchrun --nproc_per_node=4 -m pytest tests/attention/test_parallel_attention.py -v
+"""
+
+import math
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from flashinfer.utils import (
+    get_compute_capability,
+    is_sm90a_supported,
+    is_sm100a_supported,
+)
+from flashinfer.parallel_attention import (
+    UnevenCPConfig,
+    VarlenCPConfig,
+    ParallelAttention,
+    split_varlen_input,
+    get_parallel_groups,
+    uneven_cp_config,
+    ulysses_varlen_config,
+    ring_varlen_config,
+)
+
+# Skip all tests when not launched via torchrun / torch.distributed.launch
+pytestmark = pytest.mark.skipif(
+    "RANK" not in os.environ,
+    reason="Must be launched with torchrun (RANK env var not set)",
+)
+
+
+# ── Fixtures ──────────────────────────────────────────────────────────────
+
+
+@pytest.fixture(scope="session", autouse=True)
+def dist_setup():
+    """Initialize and tear down the distributed process group once per session."""
+    if not dist.is_initialized():
+        dist.init_process_group("nccl")
+    yield
+    dist.destroy_process_group()
+
+
+@pytest.fixture
+def world_size():
+    return dist.get_world_size()
+
+
+@pytest.fixture
+def rank():
+    return dist.get_rank()
+
+
+@pytest.fixture
+def device(rank):
+    return torch.device(f"cuda:{rank}")
+
+
+@pytest.fixture(autouse=True)
+def skip_if_unsupported(request):
+    """Skip test if the attention backend requires unsupported hardware."""
+    attn_type = request.node.callspec.params.get("attn_type", None)
+    if attn_type == "flash-attn3" and not is_sm90a_supported(torch.device("cuda")):
+        cc = get_compute_capability(torch.device("cuda"))
+        pytest.skip(f"flash-attn3 requires SM90a+, got {cc}")
+
+    if attn_type == "cutlass" and not is_sm100a_supported(torch.device("cuda")):
+        cc = get_compute_capability(torch.device("cuda"))
+        pytest.skip(f"cutlass requires SM100a+, got {cc}")
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────
+
+
+def _sample_tensors(num_heads, seq_len, head_dim, world_size):
+    """Create sample tensors for attention testing."""
+    shape = (num_heads, seq_len, head_dim)
+    rank = dist.get_rank()
+    device = torch.device(f"cuda:{rank}")
+
+    q = torch.randn(shape, device=device, dtype=torch.bfloat16)
+    k = torch.randn(shape, device=device, dtype=torch.bfloat16)
+    v = torch.randn(shape, device=device, dtype=torch.bfloat16)
+
+    dist.broadcast(q, src=0)
+    dist.broadcast(k, src=0)
+    dist.broadcast(v, src=0)
+
+    local_q = q.chunk(world_size, dim=1)[rank]
+    local_k = k.chunk(world_size, dim=1)[rank]
+    local_v = v.chunk(world_size, dim=1)[rank]
+    return q, k, v, local_q, local_k, local_v
+
+
+def _sample_ring_varlen_tensors(num_heads, head_dim, world_size, seq_len_list):
+    rank = dist.get_rank()
+    device = torch.device(f"cuda:{rank}")
+
+    total_seq_len = sum(seq_len_list)
+    shape = (num_heads, total_seq_len, head_dim)
+
+    q = torch.randn(shape, device=device, dtype=torch.bfloat16)
+    k = torch.randn(shape, device=device, dtype=torch.bfloat16)
+    v = torch.randn(shape, device=device, dtype=torch.bfloat16)
+
+    dist.broadcast(q, src=0)
+    dist.broadcast(k, src=0)
+    dist.broadcast(v, src=0)
+
+    local_q = split_varlen_input(q, seq_len_list, world_size, rank)
+    local_k = split_varlen_input(k, seq_len_list, world_size, rank)
+    local_v = split_varlen_input(v, seq_len_list, world_size, rank)
+
+    return q, k, v, local_q, local_k, local_v
+
+
+def _assert_cos_similarity(output, ref_output, threshold=0.99):
+    cos_sim = torch.nn.CosineSimilarity(dim=0, eps=1e-6)
+    similarity = cos_sim(
+        output.reshape(-1).to(torch.float32),
+        ref_output.reshape(-1).to(torch.float32),
+    )
+    assert similarity >= threshold, f"Cosine similarity {similarity:.6f} < {threshold}"
+
+
+# ── Tests ─────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("num_heads", [24])
+@pytest.mark.parametrize("seq_len", [6 * 8 * 1024])
+@pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("ulysses_size,ring_size", [(4, 1), (1, 4), (2, 2)])
+@pytest.mark.parametrize("attn_type", ["flash-attn3", "cutlass"])
+@pytest.mark.parametrize("tensor_layout", ["HND", "NHD"])
+def test_attn_parallel(
+    num_heads,
+    seq_len,
+    head_dim,
+    ulysses_size,
+    ring_size,
+    attn_type,
+    tensor_layout,
+    world_size,
+):
+    query, key, value, local_query, local_key, local_value = _sample_tensors(
+        num_heads, seq_len, head_dim, world_size
+    )
+
+    if tensor_layout == "NHD":
+        local_query = local_query.permute(1, 0, 2).contiguous()
+        local_key = local_key.permute(1, 0, 2).contiguous()
+        local_value = local_value.permute(1, 0, 2).contiguous()
+
+    ring_group, ulysses_group = get_parallel_groups(
+        ulysses_size=ulysses_size, ring_size=ring_size
+    )
+    attn = ParallelAttention(
+        attn_type=attn_type,
+        ulysses_group=ulysses_group,
+        ring_group=ring_group,
+    )
+
+    local_output = attn.run(local_query, local_key, local_value, tensor_layout)
+
+    if tensor_layout == "NHD":
+        local_output = local_output.permute(1, 0, 2)
+
+    ref_output = F.scaled_dot_product_attention(
+        query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0), is_causal=False
+    )
+    local_ref_output = ref_output.chunk(world_size, dim=2)[dist.get_rank()]
+
+    _assert_cos_similarity(local_output, local_ref_output)
+
+
+@pytest.mark.parametrize("num_heads", [24])
+@pytest.mark.parametrize("seq_len_padded", [6 * 8 * 1024])
+@pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("ulysses_size,ring_size", [(2, 2)])
+@pytest.mark.parametrize("attn_type", ["flash-attn3", "cutlass"])
+@pytest.mark.parametrize("tensor_layout", ["HND", "NHD"])
+def test_uneven_attn_parallel(
+    num_heads,
+    seq_len_padded,
+    head_dim,
+    ulysses_size,
+    ring_size,
+    attn_type,
+    tensor_layout,
+    world_size,
+    rank,
+    device,
+):
+    # _sample_tensors returns HND layout
+    query, key, value, local_query, local_key, local_value = _sample_tensors(
+        num_heads, seq_len_padded, head_dim, world_size
+    )
+
+    uneven_number = world_size - 1
+    seq_len_cur_rank = local_query.shape[1]
+    if rank == world_size - 1:
+        seq_len_cur_rank = seq_len_cur_rank - uneven_number
+
+    if tensor_layout == "NHD":
+        local_query = local_query.permute(1, 0, 2).contiguous()
+        local_key = local_key.permute(1, 0, 2).contiguous()
+        local_value = local_value.permute(1, 0, 2).contiguous()
+
+    ring_group, ulysses_group = get_parallel_groups(
+        ulysses_size=ulysses_size, ring_size=ring_size
+    )
+    seq_len_cur_ring_group = uneven_cp_config(
+        seq_len=seq_len_padded - uneven_number,
+        seq_len_padded=seq_len_padded,
+        seq_len_cur_rank=seq_len_cur_rank,
+        ulysses_group=ulysses_group,
+        ring_group=ring_group,
+    )
+    ucp_config = UnevenCPConfig(
+        seq_len=seq_len_padded - uneven_number,
+        seq_len_padded=seq_len_padded,
+        seq_len_cur_ring_group=seq_len_cur_ring_group,
+    )
+    attn = ParallelAttention(
+        attn_type=attn_type,
+        ulysses_group=ulysses_group,
+        ring_group=ring_group,
+        uneven_cp_config=ucp_config,
+    )
+
+    local_output = attn.run(
+        local_query, local_key, local_value, tensor_layout=tensor_layout
+    )
+
+    if tensor_layout == "NHD":
+        local_output = local_output.permute(1, 0, 2)
+
+    query = query[:, :-uneven_number, :]
+    key = key[:, :-uneven_number, :]
+    value = value[:, :-uneven_number, :]
+
+    ref_output = F.scaled_dot_product_attention(
+        query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0), is_causal=False
+    )
+    local_ref_output = ref_output.chunk(world_size, dim=2)[rank]
+
+    if rank == world_size - 1:
+        local_output = local_output[:, :-uneven_number, :]
+
+    _assert_cos_similarity(local_output, local_ref_output)
+
+
+@pytest.mark.parametrize("num_heads", [24])
+@pytest.mark.parametrize("seq_len_list", [[1 * 8 * 1024 - 1, 3 * 8 * 1024]])
+@pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("attn_type", ["flash-attn3", "cutlass"])
+@pytest.mark.parametrize("tensor_layout", ["HND", "NHD"])
+def test_ulysses_varlen_attn_parallel(
+    num_heads, seq_len_list, head_dim, attn_type, tensor_layout, world_size, rank
+):
+    ulysses_size = world_size
+    ring_size = 1
+
+    total_seq_len = sum(seq_len_list)
+    seq_len_padded = math.ceil(total_seq_len / world_size) * world_size
+    uneven_number = seq_len_padded - total_seq_len
+
+    # _sample_tensors returns HND layout
+    query, key, value, local_query, local_key, local_value = _sample_tensors(
+        num_heads, seq_len_padded, head_dim, world_size
+    )
+
+    if tensor_layout == "NHD":
+        local_query = local_query.permute(1, 0, 2).contiguous()
+        local_key = local_key.permute(1, 0, 2).contiguous()
+        local_value = local_value.permute(1, 0, 2).contiguous()
+
+    ring_group, ulysses_group = get_parallel_groups(
+        ulysses_size=ulysses_size, ring_size=ring_size
+    )
+    cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = ulysses_varlen_config(
+        seq_len_list, seq_len_list
+    )
+    vcp_config = VarlenCPConfig(
+        cu_seqlens_q_cur_ulysses_group=cu_seqlens_q,
+        cu_seqlens_kv_cur_ulysses_group=cu_seqlens_kv,
+        max_seq_len_q_cur_ulysses_group=max_seqlen_q,
+        max_seq_len_kv_cur_ulysses_group=max_seqlen_kv,
+    )
+    attn = ParallelAttention(
+        attn_type=attn_type,
+        ulysses_group=ulysses_group,
+        ring_group=ring_group,
+        varlen_cp_config=vcp_config,
+    )
+
+    local_output = attn.run(
+        local_query, local_key, local_value, tensor_layout=tensor_layout
+    )
+
+    if tensor_layout == "NHD":
+        local_output = local_output.permute(1, 0, 2)
+
+    cu_seqlens_q = cu_seqlens_q.cpu()
+    cu_seqlens_kv = cu_seqlens_kv.cpu()
+    local_ref_output_list = []
+    for i in range(len(seq_len_list)):
+        q_tmp = query[:, cu_seqlens_q[i] : cu_seqlens_q[i + 1], :]
+        k_tmp = key[:, cu_seqlens_kv[i] : cu_seqlens_kv[i + 1], :]
+        v_tmp = value[:, cu_seqlens_kv[i] : cu_seqlens_kv[i + 1], :]
+        tmp_output = F.scaled_dot_product_attention(
+            q_tmp.unsqueeze(0), k_tmp.unsqueeze(0), v_tmp.unsqueeze(0), is_causal=False
+        )
+        local_ref_output_list.append(tmp_output)
+
+    ref_output = torch.cat(local_ref_output_list, dim=2)
+    local_ref_output = ref_output.chunk(world_size, dim=2)[rank]
+
+    if rank == world_size - 1 and seq_len_padded > total_seq_len:
+        local_output = local_output[:, :-uneven_number, :]
+
+    _assert_cos_similarity(local_output, local_ref_output)
+
+
+@pytest.mark.parametrize("num_heads", [24])
+@pytest.mark.parametrize(
+    "seq_len_list",
+    [torch.tensor([1021, 1024, 1027, 750, 826], dtype=torch.int32)],
+)
+@pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("attn_type", ["flash-attn3", "cutlass"])
+@pytest.mark.parametrize("tensor_layout", ["HND", "NHD"])
+def test_ring_varlen_attn_parallel(
+    num_heads, seq_len_list, head_dim, attn_type, tensor_layout, world_size, rank
+):
+    ring_size = world_size
+
+    full_cu_seqlens = [0]
+    for seq_len in seq_len_list:
+        full_cu_seqlens.append(full_cu_seqlens[-1] + seq_len)
+
+    # _sample_ring_varlen_tensors returns HND layout
+    query, key, value, local_query, local_key, local_value = (
+        _sample_ring_varlen_tensors(num_heads, head_dim, world_size, seq_len_list)
+    )
+
+    if tensor_layout == "NHD":
+        local_query = local_query.permute(1, 0, 2).contiguous()
+        local_key = local_key.permute(1, 0, 2).contiguous()
+        local_value = local_value.permute(1, 0, 2).contiguous()
+
+    ring_group, ulysses_group = get_parallel_groups(ulysses_size=1, ring_size=ring_size)
+    cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv = ring_varlen_config(
+        seq_len_list, seq_len_list, ring_group
+    )
+    vcp_config = VarlenCPConfig(
+        cu_seqlens_q_cur_ring_group=cu_seqlens_q,
+        cu_seqlens_kv_cur_ring_group=cu_seqlens_kv,
+        max_seq_len_q_cur_ring_group=max_seqlen_q,
+        max_seq_len_kv_cur_ring_group=max_seqlen_kv,
+    )
+    attn = ParallelAttention(
+        attn_type=attn_type,
+        ulysses_group=ulysses_group,
+        ring_group=ring_group,
+        varlen_cp_config=vcp_config,
+    )
+    local_output = attn.run(
+        local_query, local_key, local_value, tensor_layout=tensor_layout
+    )
+
+    if tensor_layout == "NHD":
+        local_output = local_output.permute(1, 0, 2)
+
+    local_ref_output_list = []
+    for i in range(len(seq_len_list)):
+        q_tmp = query[:, full_cu_seqlens[i] : full_cu_seqlens[i + 1], :]
+        k_tmp = key[:, full_cu_seqlens[i] : full_cu_seqlens[i + 1], :]
+        v_tmp = value[:, full_cu_seqlens[i] : full_cu_seqlens[i + 1], :]
+        tmp_output = F.scaled_dot_product_attention(
+            q_tmp.unsqueeze(0), k_tmp.unsqueeze(0), v_tmp.unsqueeze(0), is_causal=False
+        )
+        local_ref_output_list.append(tmp_output)
+
+    ref_output = torch.cat(local_ref_output_list, dim=2).squeeze(0)
+    local_ref_output = split_varlen_input(ref_output, seq_len_list, world_size, rank)
+
+    _assert_cos_similarity(local_output, local_ref_output)