Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

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

Abierto
  • ¿Algo no está bien?
#2,868 3 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 2 días

Nadie ha tomado este issue todavía.

  • #2901 de @beccohov — cerrado sin fusionar

Evaluación

Dificultad
3/5
Tiempo estimado
1-2 días
Aptitud para principiantes
68/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Tranquilo
Stack tecnológico
python, pytorch
Área
distributed-systems, machine-learning

Línea de trabajo

Comienza con transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py, especialmente con la lógica de enable_mla cerca de la línea 1316 y el reshape de la salida cerca de la línea 1854. Ejecuta la reproducción distribuida de CUDA proporcionada con dimensiones K/V asimétricas y GQA, y compara después la ruta CP con la ruta no-CP. Se considera terminado cuando el paso forward de CP se completa y produce una salida con la forma esperada de los attention heads.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
3.6k
Forks
862
Merge medio
3 d 17 h
PR fusionados (30 d)
55

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/TransformerEngine

Todos los issues de NVIDIA/TransformerEngine

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.