There is a bug when both packing and context parallelism (CP) are enabled at the same time.

I created a sequence of length 7, where the first 4 tokens are set to 1 and the next 3 tokens are set to 2. Then I padded the sequence with a value of 200 to reach a total length of 10. I enabled context parallelism (CP) with cp=2, but found that the computed results are inconsistent compared to when CP is disabled.

The core issue appears to be that with CP enabled, the sequence length is implicitly rounded down. That is, a sequence of original length 7 is treated as if its length were 7 // 2 * 2 = 6 during computation, ignoring the last valid token.

import torch
import torch.distributed as dist
from transformer_engine.pytorch import DotProductAttention
import os

def init_dist():
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank())
    print(f"[Rank {dist.get_rank()}] Initialized process group")

def main():
    init_dist()
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    assert world_size == 2, "This test assumes 2 GPUs."

    torch.manual_seed(42)

    # 配置参数
    batch_size = 1
    num_heads = 1
    head_dim = 8
    dtype = torch.float16

    if rank == 0:
        q_len = 4
        kv_len = 4
        # 构造不同 rank 的数据
        q = torch.full((5, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
        k = torch.full((5, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
        v = torch.full((5, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)

        q[q_len:] = 2
        k[kv_len:] = 2
        v[kv_len:] = 2


        # 拼接 cu_seqlens
        cu_seqlens_q = torch.tensor([0, 7], dtype=torch.int32, device="cuda")
        cu_seqlens_kv = torch.tensor([0, 7], dtype=torch.int32, device="cuda")
        cu_seqlens_q_padded = torch.tensor([0, 10], dtype=torch.int32, device="cuda")
        cu_seqlens_kv_padded = torch.tensor([0, 10], dtype=torch.int32, device="cuda")
    else:
        q_len = 3
        kv_len = 3
        # 构造不同 rank 的数据
        q = torch.full((5, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
        k = torch.full((5, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
        v = torch.full((5, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)

        q[2:] = 100
        k[2:] = 100
        v[2:] = 100

        # 拼接 cu_seqlens
        cu_seqlens_q = torch.tensor([0, 7], dtype=torch.int32, device="cuda")
        cu_seqlens_kv = torch.tensor([0, 7], dtype=torch.int32, device="cuda")
        cu_seqlens_q_padded = torch.tensor([0, 10], dtype=torch.int32, device="cuda")
        cu_seqlens_kv_padded = torch.tensor([0, 10], dtype=torch.int32, device="cuda")

    # 初始化 Attention
    attn = DotProductAttention(
        num_attention_heads=num_heads,
        kv_channels=head_dim,
        attention_type="self",
        qkv_format="thd",
        attn_mask_type="padding"
    )

    cp_group = dist.group.WORLD
    cp_ranks = list(range(world_size))
    cp_stream = torch.cuda.Stream()

    attn.set_context_parallel_group(
        cp_group=cp_group,
        cp_global_ranks=cp_ranks,
        cp_stream=cp_stream,
        cp_comm_type="p2p"
    )

    torch.cuda.synchronize()

    # 执行 attention（每个 rank 上）
    out_cp = attn(
        query_layer=q,
        key_layer=k,
        value_layer=v,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_kv=cu_seqlens_kv,
        cu_seqlens_q_padded=cu_seqlens_q_padded,
        cu_seqlens_kv_padded=cu_seqlens_kv_padded,
        max_seqlen_q=10,
        max_seqlen_kv=10,
        qkv_format="thd",
        attn_mask_type="padding"
    )

    # 收集所有 rank 的输出
    out_all = [torch.empty_like(out_cp) for _ in range(world_size)]
    dist.all_gather(out_all, out_cp)
    torch.cuda.synchronize()
    # ====== Rank 0 上验证 full attention ======
    if rank == 0:
        # 构造全量输入：q,k,v 合并
        full_q = torch.cat([
            torch.full((4, num_heads, head_dim), 1.0, device="cuda", dtype=dtype),
            torch.full((3, num_heads, head_dim), 2.0, device="cuda", dtype=dtype)
        ], dim=0)
        full_k = torch.cat([
            torch.full((4, num_heads, head_dim), 1.0, device="cuda", dtype=dtype),
            torch.full((3, num_heads, head_dim), 2.0, device="cuda", dtype=dtype)
        ], dim=0)
        full_v = full_k.clone()

        full_attn = DotProductAttention(
            num_attention_heads=num_heads,
            kv_channels=head_dim,
            attention_type="cross",
            qkv_format="sbhd",
            attn_mask_type="no_mask"
        )

        # reshape thd -> sbhd
        full_q = full_q.view(7, 1, num_heads, head_dim)
        full_k = full_k.view(7, 1, num_heads, head_dim)
        full_v = full_v.view(7, 1, num_heads, head_dim)

        out_ref = full_attn(
            query_layer=full_q,
            key_layer=full_k,
            value_layer=full_v,
            max_seqlen_q=7,
            max_seqlen_kv=7
        )

        out_ref = out_ref.squeeze(1)  # [q_len*2, num_heads, head_dim]

        print("cp out", out_all)
        print("out_ref", out_ref)

if __name__ == "__main__":
    main()

cp out [tensor([[100., 100., 100., 100., 100., 100., 100., 100.],
[100., 100., 100., 100., 100., 100., 100., 100.],
[100., 100., 100., 100., 100., 100., 100., 100.],
[ 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0.]], device='cuda:0',
dtype=torch.float16), tensor([[100., 100., 100., 100., 100., 100., 100., 100.],
[100., 100., 100., 100., 100., 100., 100., 100.],
[100., 100., 100., 100., 100., 100., 100., 100.],
[ 0., 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 0., 0., 0., 0., 0., 0.]], device='cuda:0',
dtype=torch.float16)]
out_ref tensor([[1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268],
[1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268],
[1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268],
[1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268, 1.9268],
[1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951],
[1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951],
[1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951, 1.9951]],
device='cuda:0', dtype=torch.float16)