Hacktoberfest 2026：维护者为十月标记出来的 issue，仍然开放、适合新手。 浏览 Hacktoberfest issue

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

未关闭
  • 哪里不对？
#3,724 0 条评论 1 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 2 天内回复

@VikingDeng 已经在做这个了。

开始于 2026年10月21日。

  • #3726 来自 @VikingDeng —— 未关闭

评估

难度
5/5
预计耗时
一周以上
新手友好度
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 小时
30 天内合并 PR
59

环境准备

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

NVIDIA/TransformerEngine 的其他 Issue

查看 NVIDIA/TransformerEngine 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。