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

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

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

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

@VikingDeng がすでに取り組んでいます。

2026年10月21日 から。

  • #3726 @VikingDeng による — オープン

評価

難易度
5/5
見積もり時間
1週間以上
初心者へのやさしさ
35/100
issue の種類
機能追加
明瞭さ
おおむね明確
活発さ
活発
技術スタック
python, pytorch
領域
backend, machine-learning

調査の方向性

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.

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

説明

enhancement

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?
主要言語
Python
スター
3.6k
フォーク
872
平均マージ
4日 11時間
マージ済み PR（30日）
59

環境構築

はじめの一歩

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

NVIDIA/TransformerEngine のほかの issue

NVIDIA/TransformerEngine の issue をすべて見る

似ている issue

Python の issue をもっと見る

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

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