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.

---

## Inputs the agent should gather first

Before writing any code, confirm:

1. **Input dtype** — fp16, bf16, or fp32 for `x` and `residual`. Both should match. The summed residual write-back keeps the input dtype; the normalized `y` typically keeps the same dtype as well.
2. **Hidden dimension H** — fixed or dynamic? Power of 2? Always divisible by `BLOCK_SIZE`? Does it fit in a single `BLOCK_SIZE` (single-pass) or require a loop (two-pass)?
3. **Batch shape** — (N, H), (B, T, H), or higher rank. Confirm normalization is over the last dimension only.
4. **Epsilon** — typically 1e-5 or 1e-6. Confirm placement (inside vs outside `sqrt`); PyTorch RMSNorm convention is inside.
5. **Weight (gamma) dtype** — most checkpoints store gamma in fp16/bf16. Multiply in fp32 then cast back.
6. **In-place vs out-of-place residual write** — does the caller want `x` overwritten with `x + residual`, or written to a separate output buffer? Aliasing rules matter.
7. **Backward needed?** — if yes, you must save `rstd` and the summed activation, and design a backward that splits the gradient into the residual path and the rmsnorm path.
8. **Eps double-counting** — confirm the downstream kernel (e.g., the next block's QKV) does not also add eps if it consumes the rstd you produce.

---

## Required reasoning process

1. **Assign one program per row.** `row_idx = tl.program_id(0)`. Grid is `(N_rows,)` where `N_rows = prod(batch_dims)`. Compute the row base pointer as `x_ptr + row_idx * row_stride` and similarly for `residual_ptr`, `y_ptr`, and the summed-output pointer.

2. **Load x and residual together with a mask.**
   ```python
   col_offsets = tl.arange(0, BLOCK_SIZE)
   mask = col_offsets < H
   x   = tl.load(x_ptr   + row_idx * stride + col_offsets, mask=mask, other=0.0)
   res = tl.load(res_ptr + row_idx * stride + col_offsets, mask=mask, other=0.0)
   ```
   `other=0.0` is correct: out-of-bounds positions contribute 0 to the sum and 0 to the variance.

3. **Cast to fp32 and compute the sum.**
   ```python
   x   = x.to(tl.float32)
   res = res.to(tl.float32)
   h   = x + res
   ```
   This is the single most important precision choice. Residual streams in deep transformers grow over 32+ layers; performing the add in fp16/bf16 loses bits before the squaring step and corrupts the variance estimate.

4. **Write the summed residual back immediately.**
   ```python
   tl.store(sum_ptr + row_idx * stride + col_offsets, h.to(input_dtype), mask=mask)
   ```
   Cast back to the input dtype on the way out. This is the value the next block consumes as its residual; it is not optional and must be written in the same kernel pass to realize the fusion benefit.

5. **Compute the mean square in fp32.**
   ```python
   var = tl.sum(h * h, axis=0) / H
   rstd = 1.0 / tl.sqrt(var + eps)
   ```
   Divide by `H`, not by the count of unmasked elements. Out-of-bounds positions loaded as 0.0 contribute 0 to the sum-of-squares; dividing by `H` is correct.

6. **Normalize and apply gamma.**
   ```python
   gamma = tl.load(g_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32)
   y     = h * rstd * gamma
   ```
   Gamma is a 1D tensor of shape `(H,)`. Use `col_offsets` only — no row offset.

7. **Store the normalized output.**
   ```python
   tl.store(y_ptr + row_idx * stride + col_offsets, y.to(input_dtype), mask=mask)
   ```
   Same boundary mask as the load.

8. **For H > BLOCK_SIZE,** loop over chunks. Two paths:
   - **Two-pass:** pass 1 loads `x`, `res`, computes `h`, stores `h` to the sum buffer, accumulates `sum(h*h)`. Pass 2 reloads `h` from the sum buffer, normalizes, stores `y`. Requires reading `h` twice but avoids any extra storage.
   - **Single-pass with online accumulation:** keep a running `sum_sq` across chunks; on the second pass over chunks, recompute `h = x + res` (since `h` was not stashed) and normalize. This re-reads `x` and `res` and is usually worse than two-pass.
   Pick two-pass. It writes the residual once (which you must do anyway) and reads it back once.

9. **For training, save `rstd` and the summed residual.** The backward kernel needs both. The summed residual is already written by step 4; just store `rstd` to a `(N_rows,)` buffer.

10. **Backward (if required).** The forward output is `y = h * rstd * gamma` where `h = x + res`. The gradient of the loss with respect to `x` is `dL/dh` (since `dh/dx = 1`); the gradient with respect to `res` is also `dL/dh`. So `dL/dx = dL/dres = drmsnorm_dh`, and gradients on the residual stream pass through unchanged. Implement the rmsnorm backward to produce `drmsnorm_dh` once, then write it to both `dx` and `dres` (or accumulate into `dx` if `res` aliases `x`).

---

## Kernel design rules

- `BLOCK_SIZE` must be a power of 2 and declared `tl.constexpr`. Choose the smallest power of 2 >= `H` for the single-pass case; otherwise pick a fixed tile (e.g., 1024 or 2048) and loop.
- All arithmetic from the add through the variance must run in fp32. Cast inputs immediately after loading. Cast back to the input dtype only at the two stores.
- The summed residual write (`x + residual`) must happen in the same kernel pass. If you skip it, the fusion is meaningless — the next block will recompute the same sum.
- Epsilon placement is `1.0 / sqrt(var + eps)` (inside the sqrt). This matches the PyTorch RMSNorm convention. Document the choice if you deviate.
- Gamma is 1D, shape `(H,)`, shared across all rows. Pointer arithmetic uses `col_offsets` only.
- Row stride for `x`, `residual`, `y`, and the summed output must be passed as kernel arguments. Do not hardcode `stride = H`. They are usually equal but the contract should not assume it.
- If `x` and the summed-output pointer alias (in-place residual write), make sure you finish reading `x` before writing the summed residual. In Triton, since the load is a register copy, the in-place store after the add is safe.
- Do not add a bias term — RMSNorm has gamma but no beta. Adding beta silently turns the kernel into a malformed LayerNorm.

---

## Correctness requirements

- The summed residual must be written to memory in the same kernel pass. Skipping this write produces silently wrong outputs starting at layer 2, because every subsequent block reads a stale residual.
- The variance denominator is `H` (full row length), not the count of unmasked positions. Out-of-bounds positions contribute 0 because they are loaded with `other=0.0` and `0 * 0 = 0`.
- The add `h = x + residual` must execute in fp32. Performing the add in fp16/bf16 truncates the residual stream and produces a measurably wrong variance — visible as small but consistent drift versus a reference implementation, especially at layer counts > 16.
- Out-of-bounds load uses `other=0.0`. Using `other=-inf` or any nonzero value corrupts the sum-of-squares.
- The output store mask must match the input load mask exactly (same `col_offsets < H`).
- Gamma loads use `col_offsets` only — no row offset. Adding `row_idx * H` is a common bug that reads garbage.
- Eps must be strictly positive. Reject or warn on `eps <= 0`.
- If `x` and the summed-output buffer point to the same memory (in-place residual update), the kernel must read `x` into registers before writing the sum back. Do not interleave per-element loads and stores between `x` and the sum buffer — load the full tile first.
- The summed residual must be cast back to the input dtype on store. Writing fp32 into an fp16 buffer (or vice versa without a cast) is a memory-corruption bug.

---

## Performance requirements

The agent must reason through the following before finalizing:

- **Memory traffic.** Naive: read `x`, read `residual`, write `h` (the add kernel), read `h`, read `gamma`, write `y` (the rmsnorm kernel) — about `5 * N_rows * H * sizeof(dtype)` bytes plus one gamma read per row. Fused: read `x`, read `residual`, read `gamma`, write `h`, write `y` — about `4 * N_rows * H * sizeof(dtype)` bytes. The saving is one full read pass over the activation tensor. Quote the saving as `~20%` of the unfused traffic for the most common case, not as a speedup multiplier.
- **Single-pass vs two-pass.** Two-pass is correct for `H > BLOCK_SIZE` but adds one read of the summed residual. For `H <= BLOCK_SIZE` (the common case for `H` up to 8192 on modern hardware), single-pass keeps `h` in registers and avoids the reload.
- **Register pressure.** Holding fp32 versions of both `x` and `residual` plus their fp32 sum doubles register use compared to LayerNorm without fusion. For `BLOCK_SIZE = 8192` fp16 inputs, that is roughly 96 KB of register data per program if you keep all three live; you cannot. Reuse the registers: cast and add immediately, drop `x` and `res`, keep only `h`.
- **Gamma reuse.** Gamma is read per row but resides in L2 after the first few programs. Do not preload to shared memory.
- **Bandwidth ceiling.** A well-implemented fused kernel should approach HBM peak for moderate `H` (1024–8192). State the achieved bandwidth in benchmark output, not a speedup multiplier alone.
- **Launch overhead.** For very small workloads (e.g., decoding with `N_rows = 1`), the absolute saving is one launch — about 5–10 microseconds. Below that scale the fusion still helps but the relative speedup is dominated by Python and PyTorch overhead, not the kernel.

---

## Output format

The agent should produce:

1. **The Triton kernel function** with `@triton.jit`, taking `x_ptr`, `residual_ptr`, `gamma_ptr`, `y_ptr`, `sum_ptr` (the summed-residual output, may alias `x_ptr`), `rstd_ptr` (optional, training only), `stride`, `H`, `eps`, and `BLOCK_SIZE: tl.constexpr` as arguments.
2. **A Python launcher** that flattens the leading dimensions, computes the grid as `(N_rows,)`, asserts contiguity on the last dim, extracts strides, picks `BLOCK_SIZE` based on `H`, and exposes an in-place vs out-of-place flag for the residual.
3. **A correctness test** comparing against a reference `(x + residual).pow(2).mean(-1, keepdim=True).add(eps).rsqrt() * (x + residual) * gamma`, plus a check that the summed residual buffer matches `x + residual` exactly. Use `torch.allclose(atol=1e-3, rtol=1e-3)` for fp16/bf16 and `atol=1e-5` for fp32.
4. **A bandwidth benchmark** comparing the fused kernel against an unfused baseline (PyTorch `add` + reference `rmsnorm`) using `triton.testing.do_bench` or `torch.utils.benchmark.Timer`. Report achieved GB/s, not just speedup.
5. **A note on backward** — if forward-only, state explicitly that this kernel is forward-only. If backward is included, state how the residual gradient is split.
6. **Comment block** documenting epsilon placement, dtype contract for gamma, and the aliasing assumption between `x` and the summed-output buffer.

---

## Common failure modes

- **Forgetting to write the summed residual.** The kernel produces the correct `y` for the current layer but the next layer reads a stale residual. The model output is silently wrong from layer 2 onward. Tests on a single layer pass; tests on the full stack fail. Always include an explicit check that the summed buffer equals `x + residual`.
- **fp16 add before squaring.** Performing `h = x + res` in fp16/bf16 before the cast loses precision in the residual stream. The variance estimate drifts and per-token logits diverge over many layers. Always cast to fp32 first, then add, then square.
- **Aliasing bug between `x` and `sum_ptr`.** If the caller passes `sum_ptr == x_ptr` for in-place residual update and the kernel writes the sum before fully reading `x`, the loaded `x` is corrupted. In Triton this is fine if you load the full tile first then add then store, but is a real bug if you tile within a single program.
- **Double-counting eps.** If a downstream kernel (e.g., a second normalization or a reused rstd) also adds eps to the variance, the effective eps is `2 * eps`. Document who owns eps in the contract.
- **Gamma cast at the wrong stage.** Loading gamma in fp16 and multiplying by fp32 `h * rstd` produces an implicit narrowing in some Triton versions. Cast gamma to fp32 explicitly before multiplying.
- **Wrong denominator for variance.** Dividing by `tl.sum(mask.to(tl.float32))` instead of `H` gives a different (incorrect) variance when `H` is not a multiple of `BLOCK_SIZE`.
- **Storing fp32 into an fp16 output buffer.** A forgotten `.to(input_dtype)` on either the summed-residual store or the normalized-output store either crashes (Triton type check) or corrupts memory (silent in some versions).
- **Backward forgetting the residual path.** A forward-fused kernel paired with a backward that only handles the rmsnorm path produces wrong gradients for `x` and `residual`. Both must receive `drmsnorm_dh`.
- **Using this for LayerNorm by accident.** RMSNorm has no mean subtraction and no beta. Adding either turns the kernel into a malformed LayerNorm and wastes compute.

---

## Review checklist

- [ ] `BLOCK_SIZE` is a power of 2 and declared `tl.constexpr`.
- [ ] Both `x` and `residual` are cast to fp32 immediately after loading, before the add.
- [ ] The add `h = x + residual` runs in fp32.
- [ ] The summed residual `h` is stored back to memory in the same kernel pass, cast to the input dtype.
- [ ] Out-of-bounds load uses `other=0.0` for both `x` and `residual`.
- [ ] Variance denominator is `H` (full row length), not count of valid elements.
- [ ] Variance and rstd are computed in fp32.
- [ ] Eps is placed inside the sqrt: `1.0 / sqrt(var + eps)`.
- [ ] Gamma loads use `col_offsets` only — no row offset added.
- [ ] Gamma is cast to fp32 before multiplying with `h * rstd`.
- [ ] No beta term and no mean subtraction (this is RMSNorm, not LayerNorm).
- [ ] Output store applies the same boundary mask as the input load.
- [ ] Row strides for all four buffers are passed as kernel arguments.
- [ ] Aliasing contract between `x` and the summed-output buffer is documented and respected (load full tile before storing).
- [ ] For training: `rstd` is stored to a `(N_rows,)` buffer; the summed residual is already in memory.
- [ ] Backward (if implemented) routes `drmsnorm_dh` to both `dx` and `dresidual`.
- [ ] Correctness test verifies both the normalized output and the summed residual buffer.
- [ ] Benchmark reports achieved bandwidth, not only a speedup multiplier, and compares against an unfused baseline.
- [ ] No claim that fusion is faster without a measured number on the target hardware.