write-triton-fused-add-rmsnorm-kernel

Skill: Write a Triton Fused Add+RMSNorm Kernel

Purpose

Guide the agent through implementing a single Triton kernel that computes y = rmsnorm(x + residual) while also writing back x + residual for the next transformer block's residual stream. This fusion is the dominant pattern in LLaMA, Mistral, Qwen, and similar decoder blocks: every attention sub-block and every MLP sub-block ends with residual_add -> rmsnorm. Done correctly, the kernel saves one full read+write pass over the activation tensor compared to a naive add kernel followed by an rmsnorm kernel, and removes one launch.

Use this when

• You are implementing a transformer inference path (LLaMA-family, Mistral, Qwen, Gemma, DeepSeek, etc.) and the block structure is h = x + residual; out = rmsnorm(h) with h becoming the residual into the next block.
• Profiling shows the unfused add and rmsnorm kernels each touching the full activation tensor and the pipeline is HBM-bandwidth bound.
• You need a custom Triton implementation because the framework path (PyTorch eager, vendor library) does not fuse these two ops, or torch.compile is unavailable, partial, or breaks the graph.
• You are matching the kernel surface of vLLM (fused_add_rms_norm), FlashInfer, or Liger-Kernel and need a Triton equivalent.
• You want to keep the residual write as part of the same kernel pass to avoid an extra launch and an extra read of x in the next block.

Do not use this when

• The next operation is itself fusable with the normalization output (e.g., RMSNorm immediately followed by a QKV projection where the matmul's loader can ingest unnormalized values from the epilogue of a previous kernel). Prefer a larger fused block over a chain of small ones.
• Hidden dimension is small (< 512). The kernel is launch-overhead bound at that size, and a generic add + framework rmsnorm is fine — the saved bandwidth is dwarfed by launch latency.
• Training is required and you have not designed a backward that handles gradients through both the add and the norm. Forward fusion is straightforward; backward fusion is materially harder (see Common failure modes).
• Mean subtraction is needed (LayerNorm, not RMSNorm). Use the standard Triton LayerNorm skill instead — RMSNorm omits the mean.
• The residual is on a different dtype, layout, or device than x. Resolve the layout mismatch before fusing.