Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

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

Aperta
  • Qualcosa non va?
#3,724 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

@VikingDeng ci sta già lavorando.

Dal 21/10/2026.

  • #3726 di @VikingDeng — aperta

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
35/100
Tipo di issue
Funzionalità
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
backend, machine-learning

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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?
Lingua principale
Python
Stelle
3.6k
Fork
872
Merge medio
4g 12h
PR unite (30g)
60

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/TransformerEngine

Tutte le issue di NVIDIA/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.