[QUESTION]：How to pad sequences that cannot be evenly split in context parallelism (CP)?

For example, suppose a sequence has length 7 and I want to apply self-attention using CP with 2 devices. Since the sequence length isn’t divisible by 2 * cp, I pad it to 8. At this point, I want to mask out the padded token.

I tried the following: setting qkv_format="thd", attn_mask_type="padding", and using cu_seqlens_q = cu_seqlens_kv = [0, 7]. However, this doesn’t work as expected — it internally behaves as if cu_seqlens_q = cu_seqlens_kv = [0, 6], which causes issues. Specifically, it treats only the first 3 tokens in each of the 4-token chunks (on each device) as valid.

Question:
In this scenario, how should I properly handle padding and masking for sequences that cannot be evenly split across context-parallel devices?

I designed a test case where the first 4 tokens are all set to 1, tokens 5 to 7 are set to 2, and the 8th token is a padded token, set to 100. Then, I compare the output of context parallel (CP) attention and the output of full (non-CP) attention.

The results are:
```
cp out [tensor([[1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443],
[1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443],
[1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443, 1.9443],
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],
device='cuda:0', dtype=torch.float16), tensor([[1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971],
[1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971],
[1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971, 1.9971],
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],
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)
```
```
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((4, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
k = torch.full((4, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
v = torch.full((4, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)

# 拼接 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")
else:
q_len = 3
kv_len = 3
# 构造不同 rank 的数据
q = torch.full((4, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
k = torch.full((4, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)
v = torch.full((4, num_heads, head_dim), float(rank + 1), device="cuda", dtype=dtype)

q[q_len:] = 100
k[kv_len:] = 100
v[kv_len:] = 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")

# 初始化 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,
max_seqlen_q=8,
max_seqlen_kv=8,
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()
```
I would really appreciate any relevant suggestions.