[Bug] Context parallel crashes with asymmetric K/V head dims (GQA + enable_mla path)

Describe the bug

When using DotProductAttention with context parallelism (CP) and asymmetric K/V head dimensions (kv_channels=(k_dim, v_dim) where k_dim != v_dim), the CP forward pass crashes with a shape mismatch in context_parallel.py.

The root cause: enable_mla = k.shape[-1] != v.shape[-1] (here) triggers the MLA code path, which reshapes the attention output using v_shape (derived from the V input tensor with num_kv_heads). However, with GQA the attention output has num_attention_heads (post-expansion), not
num_kv_heads, causing a size mismatch.

This affects models like MiMo-V2-Flash from Megatron Bridge (I'm currently implementing support for it here) which uses head_dim=192 for Q/K but v_head_dim=128 for V with GQA (num_attention_heads=64, num_kv_heads=4).

Steps/Code to reproduce bug

import os

import torch
import transformer_engine as te
from transformer_engine.pytorch.attention.dot_product_attention import DotProductAttention

print(f"TE version: {te.__version__}")

local_rank = int(os.environ.get("LOCAL_RANK", "0"))
torch.cuda.set_device(local_rank)
torch.distributed.init_process_group("nccl")
device = torch.device(f"cuda:{local_rank}")

B, T, num_heads, num_kv_heads = 1, 32, 8, 2
qk_head_dim, v_head_dim = 64, 48  # asymmetric: k != v

cp_group = torch.distributed.group.WORLD

# CP case doesn't work:
attn = DotProductAttention(
        num_attention_heads=num_heads,
        kv_channels=(qk_head_dim, v_head_dim),
        num_gqa_groups=num_kv_heads,
        attention_dropout=0.0,
        tp_size=1,
        tp_group=None,
        cp_global_ranks=list(range(torch.distributed.get_world_size())),
        cp_group=cp_group,
        cp_stream=torch.cuda.Stream(),
        softmax_type="vanilla", #"learnable" SWA uses attention sink bias
).to(device)

# non-CP works:
# attn = DotProductAttention(
#         num_attention_heads=num_heads,
#         kv_channels=(qk_head_dim, v_head_dim),
#         num_gqa_groups=num_kv_heads,
#         attention_dropout=0.0,
#         softmax_type="vanilla", #"learnable" SWA uses attention sink bias
# ).to(device)

q = torch.randn(T, B, num_heads, qk_head_dim, device=device, dtype=torch.bfloat16)
k = torch.randn(T, B, num_kv_heads, qk_head_dim, device=device, dtype=torch.bfloat16)
v = torch.randn(T, B, num_kv_heads, v_head_dim, device=device, dtype=torch.bfloat16)
out = attn(q, k, v, attn_mask_type="causal")

Expected behavior

CP forward should handle asymmetric K/V head dims with GQA correctly. The output reshape at context_parallel.py should use the attention output's actual shape (num_attention_heads) rather than the V input shape (num_kv_heads).

Environment overview (please complete the following information)

• PyTorch: 2.7
• Transformer Engine: 2.13.0
• CUDA: 13.0

Device details

• GPU model

Additional context

Add any other context about the problem here.

Start with transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py, especially the enable_mla logic near line 1316 and the output reshape near line 1854. Run the provided distributed CUDA reproduction with asymmetric K/V dimensions and GQA, then compare the CP path with the non-CP path. Done means the CP forward pass completes and produces an output with the expected attention-head shape.