Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

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

オープン
  • 何か問題がありますか？
#2,868 コメント 3 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 2 日以内に返信

まだ誰も着手していません。

  • #2901 @beccohov による — マージされずにクローズ

評価

難易度
3/5
見積もり時間
1〜2日
初心者へのやさしさ
68/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
静か
技術スタック
python, pytorch
領域
distributed-systems, machine-learning

調査の方向性

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py から始め、特に 1316 行付近の enable_mla ロジックと、1854 行付近の出力の reshape を確認します。非対称な K/V 次元と GQA を使って、提供されている分散 CUDA 再現コードを実行し、その後 CP パスと非 CP パスを比較します。CP の forward pass が完了し、期待される attention-head の形状を持つ出力が生成されれば完了です。

索引モデルが issue の本文から書いたものです。

説明

bug

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.

主要言語
Python
スター
3.6k
フォーク
862
平均マージ
4日 15時間
マージ済み PR（30日）
58

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

NVIDIA/TransformerEngine のほかの issue

NVIDIA/TransformerEngine の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。