perf: Add fused Triton RMSNormGated kernel for Qwen3.5 GDN layers#697
perf: Add fused Triton RMSNormGated kernel for Qwen3.5 GDN layers#697
Conversation
…el definitions Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
This PR introduces fused Triton implementations to reduce kernel launch overhead and intermediate tensor traffic in normalization-heavy paths (notably Qwen3.5 GDN’s RMSNormGated and Gemma-style RMSNorm), aiming to improve decode-time performance.
Changes:
- Added two Triton kernels that fuse RMSNorm + SiLU gating + weight multiplication for
RMSNormGatedwhenhead_dim == 128. - Switched
RMSNormGated.forward()to prefer the Triton fast path with fallback to the native implementation for unsupported configurations. - Replaced the
GemmaRMSNorm.forward_cudatorch.compile-based path with a dedicated Triton implementation and removed duplicate late imports of Triton.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| z_ptr, | ||
| weight_ptr, | ||
| out_ptr, | ||
| num_rows: tl.constexpr, |
| torch.nn.init.ones_(self.weight) | ||
|
|
||
| def forward_triton(self, x: torch.Tensor, z: torch.Tensor): | ||
| if ( | ||
| z is None | ||
| or x.ndim != 3 | ||
| or self.group_size is not None | ||
| or not self.norm_before_gate | ||
| or x.shape[-1] != 128 | ||
| or not x.is_contiguous() | ||
| or not z.is_contiguous() |
| return self.forward_static(self.weight.data, self.variance_epsilon, x, residual) | ||
|
|
||
| def forward_cuda( | ||
| self, | ||
| x: torch.Tensor, | ||
| residual: torch.Tensor | None = None, | ||
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: | ||
| if torch.compiler.is_compiling(): | ||
| return self.forward_native(x, residual) | ||
| from atom.model_ops.triton_gemma_rmsnorm import gemma_rmsnorm_triton | ||
|
|
||
| if not getattr(self, "_is_compiled", False): | ||
| self.forward_static = torch.compile(self.forward_static) # type: ignore | ||
| self._is_compiled = True | ||
| return self.forward_native(x, residual) | ||
| return gemma_rmsnorm_triton( |
zovonoir
changed the title
…ures Keep the two new Triton kernels but route forward() back to forward_native() so we can verify whether CI failures are pre-existing. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…e-existing Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
I'm an AI assistant (Claude) helping with this PR. I wanted to flag a CI observation: The current Accuracy test failures appear to be pre-existing and unrelated to this PR's changes. To verify this, we temporarily reverted the This confirms that the CI failures are not caused by the new Triton RMSNormGated kernels introduced in this PR. Could you help check if these are known infrastructure or baseline issues? Thanks! |
| z_ptr, | ||
| weight_ptr, | ||
| out_ptr, | ||
| num_rows: tl.constexpr, |
| def forward_triton(self, x: torch.Tensor, z: torch.Tensor): | ||
| if ( | ||
| z is None | ||
| or x.ndim != 3 | ||
| or self.group_size is not None | ||
| or not self.norm_before_gate | ||
| or x.shape[-1] != 128 | ||
| or not x.is_contiguous() | ||
| or not z.is_contiguous() | ||
| ): |
| from atom.model_ops.triton_gemma_rmsnorm import gemma_rmsnorm_triton | ||
|
|
||
| if not getattr(self, "_is_compiled", False): | ||
| self.forward_static = torch.compile(self.forward_static) # type: ignore | ||
| self._is_compiled = True | ||
| return self.forward_native(x, residual) | ||
| return gemma_rmsnorm_triton( | ||
| x, self.weight.data, self.variance_epsilon, residual | ||
| ) |
| x: torch.Tensor, | ||
| residual: torch.Tensor | None = None, | ||
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: | ||
| if torch.compiler.is_compiling(): | ||
| return self.forward_native(x, residual) | ||
| from atom.model_ops.triton_gemma_rmsnorm import gemma_rmsnorm_triton | ||
|
|
||
| if not getattr(self, "_is_compiled", False): | ||
| self.forward_static = torch.compile(self.forward_static) # type: ignore | ||
| self._is_compiled = True | ||
| return self.forward_native(x, residual) | ||
| return gemma_rmsnorm_triton( | ||
| x, self.weight.data, self.variance_epsilon, residual | ||
| ) |
|
Since we already have |
|
@ganyi1996ppo Good point! I looked into the aiter HIP kernel ( Adding a bf16 output path should be straightforward — a compile-time We're evaluating this approach and will follow up with a patch to aiter. In the meantime, let's keep this PR open — the Triton kernels here serve as a working reference and fallback until the aiter HIP path is ready. — Claude (AI assistant) |
|
Superseded by #708, which includes all RMSNormGated changes plus additional MRoPE and GemmaRMSNorm optimizations. |
perf: Add fused Triton RMSNormGated kernel for Qwen3.5 GDN layers
Summary
RMSNormGatedwithhead_dim=128forward_nativepath inRMSNormGated.forward()withforward_triton, which falls back toforward_nativefor unsupported configurationsimport tritonto file top-level and remove duplicate importstorch.compileusage inGemmaRMSNorm.forward_staticwith a dedicated Triton kernelMotivation
Profiling Qwen3.5-27B on MI308X (TP2, ISL=60381, OSL=132) in ATOM SGLang plugin mode revealed that the native
RMSNormGatedimplementation (element-wise ops via PyTorch) contributes significant overhead, especially in the GDN (Gated Delta Network) layers where each layer callsRMSNormGatedon tensors shaped(num_tokens, num_heads, 128).The native path launches multiple small GPU kernels (variance reduction, rsqrt, silu, element-wise multiply) that are memory-bandwidth bound at decode batch sizes. Fusing these into a single Triton kernel eliminates intermediate tensor allocations and reduces kernel launch overhead.
Implementation
Two Triton kernels optimized for different workloads:
_rmsnorm_gated_contiguous_128_kernel(decode path): One thread block per (token, head) pair. Usesnum_warps=1since each block processes only 128 elements. Optimal for small batch sizes (decode)._rmsnorm_gated_contiguous_128_tiled_rows_kernel(prefill path): Each thread block processesblock_rows=32rows of 128 elements. Used whennum_rows >= 65536(e.g., 60K tokens x 24 heads = 1.4M rows). Usesnum_warps=4for better occupancy.forward_tritonmethod safely falls back toforward_nativewhen:zis None (no gating)(tokens, heads, dim)group_sizeis setnorm_before_gateis Falsehead_dim != 128Benchmark Results
Tested on MI308X with Qwen3.5-27B, ATOM SGLang plugin mode, TP2, ISL=60381, OSL=132:
TPOT reduced by 11% (18.7 → 16.7 ms), with no impact on prefill latency (TTFT unchanged).
Test plan
forward_nativeworks for unsupported configurations (group_size != None, non-contiguous tensors, head_dim != 128, etc.)