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

[Feature] Fused gated RMSNorm (RMSNorm(x) * SiLU(gate)) for GDN output norm

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

Los mantenedores suelen responder en 2 días

@VikingDeng ya está trabajando en esto.

Desde el 21/10/2026.

  • #3726 de @VikingDeng — abierto

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
35/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
backend, machine-learning

Línea de trabajo

Start by reading the existing rmsnorm_fwd and rmsnorm_bwd paths, tests/pytorch/test_numerics.py (test_rmsnorm_accuracy), and the op checks in tests/pytorch/test_fusible_ops.py. The proposed work has two stages: establish the PyTorch RMSNormGated API and numerics, then add native nvte_rmsnorm_gated_fwd/bwd kernels, including strided gates. Done means y, dx, dgamma, and dgate match an FP64 reference for both gamma modes and contiguous and strided gates; kernel performance work spans GPU architectures the author cannot test.

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

Descripción

Is your feature request related to a problem? Please describe.

Gated DeltaNet layers (Qwen3.5 / Qwen3-Next style) normalize the attention output and then gate it: y = RMSNorm(x) * silu(gate). TE's RMSNorm has no gate input (neither nvte_rmsnorm_fwd nor the PyTorch module/op), so Megatron-LM runs this as out_norm(x) * act_fn(gate) in megatron/core/ssm/gated_delta_net/gdn.py (_apply_gated_norm). There is also an opt-in Triton kernel (megatron/core/fusions/fused_gated_norm.py, behind gdn_gated_output_norm_fusion). This is the "Gated RMSNorm implementation in TE" item in the Megatron-LM MoE roadmap (NVIDIA/Megatron-LM#6757); @sbhavani suggested opening a design issue here.

Describe the solution you'd like

A separate op rather than a new argument on RMSNorm, following the precedent of backward_add_rmsnorm / nvte_rmsnorm_bwd_add:

norm = te.ops.RMSNormGated(hidden_size, eps=1e-6, zero_centered_gamma=False)
y = norm(x, gate)          # y = rmsnorm(x) * silu(gate), dgate returned in backward

I'd do it in two steps:

  1. A PyTorch op that reuses the existing rmsnorm_fwd / rmsnorm_bwd, with the gating as pointwise work. Backward is dn = dy * silu(g), dg = dy * n * silu'(g), then rmsnorm_bwd(dn). This settles the API and the numerics.
  2. Native nvte_rmsnorm_gated_fwd/bwd kernels so x, gate and y are each read or written once, starting with contiguous BF16/FP16, then strided gate (GDN slices the gate out of a projection output).

For numerics I'd match Megatron's Triton kernel: FP32 accumulation, n and dn rounded to x's dtype before the gate multiply, and dgate in the gate's dtype. That way switching Megatron from the unfused or Triton path to TE doesn't change results beyond rounding.

Tests would follow test_rmsnorm_accuracy in tests/pytorch/test_numerics.py and the op-level checks in tests/pytorch/test_fusible_ops.py: y, dx, dgamma and dgate against an FP64 reference, with zero-centered gamma on and off, and contiguous and strided gate.

Describe alternatives you've considered

  • An optional gate argument on RMSNorm. That touches the C API and every existing RMSNorm path (cuDNN backend, FP8/MXFP8 outputs, JAX), and backward would need an extra output, so a separate op seems cleaner.
  • Keeping the Triton kernel in Megatron only. That works, but other frameworks can't use it, and it doesn't compose with TE's quantized outputs.

Additional context

I can develop and check correctness on SM86 (RTX 3090, NVTE_CUDA_ARCHS=86) and SM120 (RTX PRO 6000). I don't have SM90/SM100 access, so performance numbers there would need someone else's help. Questions:

  • Is a separate RMSNormGated op the shape you'd want, or would you rather fold it into the ops fuser?
  • Should step 1 land on its own, or only together with the native kernel?
  • Should the first version support FP8/MXFP8 output, or only BF16/FP16?
Lenguaje dominante
Python
Estrellas
3.6k
Forks
872
Merge medio
4 d 12 h
PR fusionados (30 d)
60

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.