[Perf] Optimize rmsnorm/layernorm

docs/prebuilt_kernels_guide.md

@@ -34,29 +34,28 @@ executor = build_layernorm_module(N=8192, dtype_str="bf16")
 | Constant | Value | Description |
 |---|---|---|
 | `BLOCK_THREADS` | 256 | Threads per block |
-| `WARP_SIZE` | 64 | AMD wavefront size |
+| `WARP_SIZE` | 64 on CDNA, 32 on RDNA | Wavefront size, resolved from the target arch |
 | `VEC_WIDTH` | 8 | Vector load/store width |
-| `VEC_ALIGN` | 16 | Alignment for vector ops (bytes) |
 | `EPS` | 1e-5 | Numerical stability epsilon |
-| `USE_NONTEMPORAL` | True | Non-temporal stores for output |
 
 **Algorithm:**
 - **Two-pass normalization**: Pass 1 computes mean and variance, Pass 2 applies affine transform
-- **Fast path**: When `N == BLOCK_THREADS * VEC_WIDTH * 4` (for example, N=8192), uses fully register-resident computation with no scalar tail
-- **Generic path**: Handles arbitrary N with vector body + scalar tail
+- **Vectorized path**: When the element type is 16-bit and `N % VEC_WIDTH == 0`, the row is covered by `N / VEC_WIDTH` vector tiles with no scalar tail
+- **Scalar path**: FP32, or any `N` not divisible by `VEC_WIDTH`, falls back to a fully scalar two-pass implementation
 - **bf16 handling**: Software round-to-nearest-even (RNE) pack on gfx942; hardware `cvt_pk_bf16_f32` on gfx950+
 - **Warp reduction**: XOR-shuffle-based intra-wave reduction (shifts: 32, 16, 8, 4, 2, 1), then LDS-based cross-wave synchronization
 
-**Kernel signature** (using `@flyc.kernel` API):
-```
-GPU_MODULE_NAME = "layernorm_module"
-
-@kernel
-layernorm_kernel(self, Input, Gamma, Beta, Output, m_in)
+**Kernel signature:**
+```python
+@flyc.kernel
+layernorm_kernel(Input, Gamma, Beta, Output, Mean, Rstd)
 
-@jit
-__call__(self, Input, Gamma, Beta, Output, m_in)
+@flyc.jit
+launch_layernorm(Input, Gamma, Beta, Output, m_in, stream=...)
+# store_stats=True inserts Mean and Rstd before m_in
 ```
+The builder returns the `launch_layernorm` closure. The row count `m_in` is a
+runtime launch argument, not a kernel parameter.
 
 ### 1.2 RMSNorm (`kernels/norm/rmsnorm_kernel.py`)
 
@@ -70,11 +69,19 @@ executor = build_rmsnorm_module(N=8192, dtype_str="bf16", store_rstd=False)
 ```
 
 `build_rmsnorm_module(N, dtype_str, store_rstd=False, eps=EPS,
-BLOCK_THREADS=BLOCK_THREADS, weight_dtype_str=None)` optionally writes the
+BLOCK_THREADS=None, weight_dtype_str=None)` optionally writes the
 per-row reciprocal std (`rstd`) for use by the backward pass.
 `weight_dtype_str` defaults to `dtype_str`; FP16/BF16 activations additionally
 support FP32 weights.
 
+**Quantized variants:** The DynamicQuant and SmoothQuant builders emit int8
+`Output` and fp32 per-row `YScale`. `Input` must use the element dtype named by
+the builder's `dtype_str`, and every other operand — `Gamma`, the fused-add
+`ResidualIn`/`ResidualOut`, and SmoothQuant `XScale` — must match it. The
+launchers raise `ValueError` on a mismatch, so a wrong dtype fails at compile
+time rather than silently producing corrupted scales. Unlike the plain forward,
+the quantized builders do not accept FP32 weights with FP16/BF16 activations.
+
 **Backward:** `build_rmsnorm_bwd_module(N, dtype_str,
 weight_dtype_str=None)` builds the fused RMSNorm backward kernel (grid `(M,)`,
 one block per row). Kernel signature
@@ -84,19 +91,33 @@ grad). The forward bakes `eps` into `Rstd`, so the backward does not need it.
 The public plain and fused-add training wrappers return `dweight` in the
 original weight dtype.
 
-**Configuration constants:** Same as LayerNorm (BLOCK_THREADS=256, VEC_WIDTH=8, etc.)
+**Configuration constants:**
+| Constant | Value | Description |
+|---|---|---|
+| `BLOCK_THREADS` | 256; 512 on gfx95x when `N >= 8192` | Resolved by `default_block_threads(N, arch)` when the builder argument is left at `None` |
+| `WARP_SIZE` | 64 on CDNA, 32 on RDNA | Wavefront size, resolved from the target arch |
+| `VEC_WIDTH` | 8 | Vector load/store width |
+| `EPS` | 1e-5 | Numerical stability epsilon |
+
+**Algorithm (2-pass, row cached in registers):**
+1. **Pass 1**: One vectorized global read per row; the input stays in registers
+   and the sum of squares is accumulated in the same pass. A scalar tail covers
+   the `N % VEC_WIDTH` leftover elements.
+2. **Pass 2**: Normalize, multiply by gamma, and store, reusing the registers
+   from pass 1. `Gamma` is preloaded during pass 1 on the gfx942 BF16 fast path.
 
-**Algorithm (3-pass with LDS caching):**
-1. **Pass 0**: Global → LDS row cache (one-pass global read, vectorized)
-2. **Pass 1**: Sum-of-squares computation from LDS row cache
-3. **Pass 2**: Normalize + gamma multiply + store with software pipeline for Gamma prefetch
+LDS holds only the cross-wave reduction slots, sized by the wave count rather
+than by `N`; the row itself never passes through shared memory. The quantized
+builders add a third pass that applies the per-row scale.
 
 **Kernel signature:**
-```
-GPU_MODULE_NAME = "rmsnorm_module"
+```python
+@flyc.kernel
+rmsnorm_kernel(Input, Gamma, Rstd, Output)
 
-@kernel
-rmsnorm_kernel(self, Input, Gamma, Output, m_in)
+@flyc.jit
+launch_rmsnorm(Input, Gamma, Output, m_in, stream=...)
+# store_rstd=True inserts Rstd between Output and m_in
 ```
 
 ---