Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

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

Đang mở
  • Có gì không ổn?
#3,724 0 bình luận 1 reaction 0 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 2 ngày

@VikingDeng đang làm issue này rồi.

Từ ngày 21/10/2026.

  • #3726 của @VikingDeng — đang mở

Đánh giá

Độ khó
5/5
Thời gian dự kiến
Hơn một tuần
Mức phù hợp với người mới
35/100
Loại issue
Tính năng
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Sôi nổi
Công nghệ
python, pytorch
Lĩnh vực
backend, machine-learning

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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?
Ngôn ngữ chính
Python
Star
3.6k
Fork
872
Merge trung bình
4 ngày 11 giờ
Pull request đã merge (30 ngày)
59

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của NVIDIA/TransformerEngine

Tất cả issue của NVIDIA/TransformerEngine

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.