Claude skill for Triton kernel development (liger-kernel-dev) (#1170)

## Summary This PR adds a **Claude Code skill** (`liger-kernel-dev`) that automates the full lifecycle of Triton kernel development for Liger Kernel. Instead of manually writing 8+ files per kernel (ops, transformers wrapper, functional API, exports, tests, benchmarks), a contributor can now describe a PyTorch operation and the skill handles analysis, code generation, and validation through a 3-stage agentic pipeline. ### How it works 1. **Analyze** — An Analyzer agent reads the PyTorch operation (from a local file, URL, code snippet, natural language description, or model component reference), writes a standalone PyTorch reference implementation, and produces a structured kernel profile classifying the operation into one of 3 complexity tiers 2. **Generate** — A Generator agent takes the profile and creates/modifies up to 8 files: Triton forward+backward kernels with `torch.autograd.Function`, `nn.Module` wrapper, functional API, `__init__.py` exports, unit tests (parametrized over shapes and dtypes), and benchmark scripts 3. **Validate** — A Validator agent runs `make checkstyle`, unit tests (hard gate with 3 retries — stops on persistent failure), benchmarks (speed + memory for fwd/bwd/full), and generates plots. Optionally runs `ncu` profiling ### Key design decisions - **Progressive disclosure**: SKILL.md is a concise entrypoint (~63 lines); detailed instructions live in separate agent files and templates that are loaded only when needed - **Three complexity tiers**: Templates cover element-wise ops (SwiGLU-like), reduction ops (RMSNorm-like), and fused/complex ops (CrossEntropy-like) with tier-specific patterns - **Correctness is a hard gate**: If tests fail after 3 retries, the pipeline stops and reports to the user rather than pushing forward with broken code - **Single dtype benchmarks**: Benchmarks use `model.dtype` (typically bfloat16). Multi-dtype coverage is handled by unit tests - **Supports both create and modify modes**: Detects intent automatically — creating new kernels goes through the full pipeline, modifying existing kernels skips analysis ### Skill file structure ``` .claude/skills/liger-kernel-dev/ ├── SKILL.md # Main entry point (63 lines) ├── analyzer.md # Stage 1: understand operation, produce profile ├── generator.md # Stage 2: generate all 8 files ├── validator.md # Stage 3: checkstyle → tests → benchmarks → plots ├── kernel-profile-format.md # Kernel profile schema + naming conventions ├── templates/ │ ├── ops-kernel.md # Triton kernel patterns by tier │ ├── module-wrapper.md # nn.Module wrapper pattern │ ├── functional-api.md # functional.py modification pattern │ ├── unit-test.md # Test file pattern with testing rules │ └── benchmark.md # Benchmark script pattern └── examples/ ├── swiglu-profile.md # Tier 1 (element-wise) reference ├── rms-norm-profile.md # Tier 2 (reduction) reference └── cross-entropy-profile.md # Tier 3 (fused/complex) reference ``` ## Testing Done - Tested end-to-end on **ReLU Squared** (Tier 1, element-wise activation used in Nemotron models) → #1171 - Skill generated all 8 files: Triton kernel, module wrapper, functional API, exports, 18 unit tests, and benchmarks - All tests pass on H100 (float32 + bfloat16) - Benchmarks show **1.9-3.3x speedup** and **37.5% memory savings** vs PyTorch - Full PR including benchmark plots was generated end-to-end by the skill - Iteratively refined the skill based on issues found during testing (single-dtype benchmarks, template accuracy) --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-10-12 02:17:47 +00:00 · 2026-03-27 19:56:37 +00:00
parent 07dd9beebd
commit 0169c4def2
13 changed files with 1094 additions and 0 deletions
@@ -0,0 +1,63 @@
+---
+name: liger-kernel-dev
+description: "Develops production-ready Triton kernels for Liger Kernel. Creates new kernels from PyTorch operations (local files, URLs, code snippets, or natural language) with ops, module wrappers, functional APIs, unit tests, benchmarks, and plots. Also modifies existing Liger kernels. Use when adding a new Triton kernel, converting a PyTorch operation to Triton, or updating an existing Liger kernel."
+---
+
+# Liger Kernel Dev
+
+Develops Triton kernels for Liger Kernel through a 3-stage pipeline with human review between stages. Supports creating new kernels and modifying existing ones. NVIDIA GPUs only.
+
+## Mode Detection
+
+- **Create mode**: User asks to create/add/generate/write/build a new kernel → full pipeline
+- **Modify mode**: User asks to update/fix/change/extend an existing kernel → skip Analyze, modify files, then Validate
+
+## Pipeline (Create Mode)
+
+### Stage 1: Analyze
+
+Spawn an **Analyzer** agent (read [analyzer.md](analyzer.md)).
+
+Accepts any input: local file, URL, code snippet, natural language description, or model component reference. Produces a standalone PyTorch reference implementation and a kernel profile.
+
+**Human checkpoint:** Present PyTorch reference + kernel profile. Confirm before proceeding.
+
+### Stage 2: Generate
+
+Spawn a **Generator** agent (read [generator.md](generator.md)).
+
+Generates/modifies up to 8 files:
+
+1. `src/liger_kernel/ops/{kernel}.py` — NEW Triton kernels + autograd Function
+2. `src/liger_kernel/transformers/{kernel}.py` — NEW nn.Module wrapper
+3. `src/liger_kernel/transformers/functional.py` — MODIFY add functional API
+4. `src/liger_kernel/ops/__init__.py` — MODIFY export Function class
+5. `src/liger_kernel/transformers/__init__.py` — MODIFY export Module + `__all__`
+6. `test/transformers/test_{kernel}.py` — NEW unit tests
+7. `benchmark/scripts/benchmark_{kernel}.py` — NEW benchmark script
+8. `benchmark/data/all_benchmark_data.csv` — MODIFY (after benchmarks run)
+
+**Human checkpoint:** Present changes for review.
+
+### Stage 3: Validate
+
+Spawn a **Validator** agent (read [validator.md](validator.md)).
+
+Runs checkstyle, unit tests (hard gate — stops on persistent failure), benchmarks, and generates plots. Optionally runs ncu profiling.
+
+**Human checkpoint:** Report final results with benchmark numbers and plots.
+
+## Pipeline (Modify Mode)
+
+1. Read existing kernel files to understand current implementation
+2. Understand the requested modification
+3. Make targeted changes (Generator handles this)
+4. Run full Validate stage (same as create mode)
+
+## Reference Files
+
+- [kernel-profile-format.md](kernel-profile-format.md) — Kernel profile schema and field descriptions
+- [examples/swiglu-profile.md](examples/swiglu-profile.md) — Tier 1 (element-wise) reference
+- [examples/rms-norm-profile.md](examples/rms-norm-profile.md) — Tier 2 (reduction) reference
+- [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md) — Tier 3 (fused/complex) reference
+- Templates in [templates/](templates/) — Code generation patterns for each file type
@@ -0,0 +1,62 @@
+# Analyzer Agent
+
+Understands a PyTorch operation from any input form and produces a standalone PyTorch reference implementation + kernel profile.
+
+## Input Handling
+
+The user may provide the operation in any form:
+
+1. **Local file path** → Read the file directly
+2. **URL** (GitHub, HuggingFace, etc.) → Fetch via WebFetch tool
+3. **Code snippet** → Pasted in the conversation
+4. **Natural language** → Mathematical description (e.g., "element-wise SiLU(x) * y")
+5. **Model component** → e.g., "the MLP in Phi-4" — locate in transformers source and extract
+
+## Steps
+
+### 1. Understand the Operation
+
+- Read/fetch the source code from whatever input was provided
+- Identify the mathematical operation (forward pass)
+- Derive the backward pass (gradient computation)
+- Identify all inputs, outputs, and their expected shapes/dtypes
+- Note any precision-sensitive operations that need float32 upcasting (sigmoid, rsqrt, exp, log, tanh)
+
+### 2. Write PyTorch Reference
+
+Create a standalone implementation that:
+- Depends only on `torch` (no external libraries)
+- Implements both forward and backward behavior (either as an `nn.Module` or a plain function)
+- Will serve as the correctness baseline for testing
+- Is clean, readable, and well-named
+
+### 3. Classify Into Tier
+
+Read [kernel-profile-format.md](kernel-profile-format.md) for the full schema.
+
+**Tier 1 — Element-wise**: No reductions across dimensions. One row per program. Examples: SwiGLU, GeGLU, DyT.
+- Read reference: `src/liger_kernel/ops/swiglu.py`
+
+**Tier 2 — Reduction**: Cross-column reductions (tl.sum, tl.max). May need to save intermediate state for backward. May need SM-based parallelism for weight gradient reduction. Examples: RMSNorm, LayerNorm, Softmax, Sparsemax.
+- Read reference: `src/liger_kernel/ops/rms_norm.py`
+
+**Tier 3 — Fused/Complex**: Multi-pass algorithms, gradient-in-forward tricks, multiple outputs. Examples: CrossEntropy, FusedLinearCrossEntropy.
+- Read reference: `src/liger_kernel/ops/cross_entropy.py`
+
+Also read the closest example profile:
+- Tier 1 → [examples/swiglu-profile.md](examples/swiglu-profile.md)
+- Tier 2 → [examples/rms-norm-profile.md](examples/rms-norm-profile.md)
+- Tier 3 → [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md)
+
+### 4. Produce Kernel Profile
+
+Fill in all fields from [kernel-profile-format.md](kernel-profile-format.md).
+
+### 5. Present to User
+
+Show:
+1. The PyTorch reference implementation (full code)
+2. The kernel profile (all fields)
+3. Which existing kernel is closest (for the Generator to use as reference)
+
+Wait for user confirmation before proceeding to Stage 2.
@@ -0,0 +1,67 @@
+# Kernel Profile: CrossEntropy (Tier 3 — Fused/Complex)
+
+## Identity
+- operation_name: cross_entropy
+- function_class_name: LigerCrossEntropyFunction
+- module_class_name: LigerCrossEntropyLoss
+- functional_name: liger_cross_entropy
+
+## Classification
+- tier: 3
+- tier_description: fused/complex
+- closest_existing_kernel: fused_linear_cross_entropy (extends this with linear layer fusion)
+
+## Forward Pass
+- forward_inputs:
+  - _input: shape (B*T, V), logits
+  - target: shape (B*T,), label indices
+  - weight: shape (V,) or None, class weights
+  - ignore_index: int, label to ignore
+  - label_smoothing: float
+  - reduction: str, "mean" | "sum" | "none"
+  - softcap: float or None
+- forward_outputs:
+  - loss: scalar or shape (B*T,) if reduction="none"
+  - z_loss: optional auxiliary loss
+  - token_accuracy: optional accuracy metric
+  - predicted_tokens: optional argmax tokens
+- forward_computation: Two-pass online softmax + cross entropy loss with optional smoothing and softcapping
+- precision_sensitive_ops: [exp, log]
+
+## Backward Pass
+- backward_saved_tensors: [_input] — gradient is computed during forward and stored in-place
+- backward_recompute: none (gradient-in-forward trick)
+- gradient_formulas:
+  - d_input: already computed in forward pass and stored in _input tensor. Backward just scales by grad_output.
+
+## Tiling Strategy
+- grid_dimensions: 1D
+- grid_description: one program per row (one row = one token's logits over vocab)
+- block_size_source: custom — iterates over vocab in chunks of BLOCK_SIZE
+- needs_sm_parallelism: false
+
+## Module Parameters
+- module_init_params:
+  - weight: optional class weights
+  - ignore_index: int = -100
+  - lse_square_scale: float = 0.0
+  - label_smoothing: float = 0.0
+  - reduction: str = "mean"
+  - softcap: float or None = None
+  - return_z_loss: bool = False
+- learnable_params: none
+
+## Benchmarking
+- benchmark_variable: vocab_size (V)
+- benchmark_x_label: "V"
+- benchmark_x_values_suggestion: [4096, 8192, 16384, 32768, 65536, 131072]
+- benchmark_providers: ["liger", "huggingface"]
+- benchmark_fixed_config: {B: 8, T: 512, dtype: torch.bfloat16}
+
+## Key Patterns
+
+- **Online softmax**: Two-pass algorithm. Pass 1: compute running max and logsumexp. Pass 2: compute softmax and gradients. Avoids materializing the full softmax vector.
+- **Gradient-in-forward trick**: The forward kernel computes the gradient and stores it directly in `_input` (overwriting logits). The backward pass just retrieves this and multiplies by `grad_output`. This saves having to recompute softmax in backward.
+- **Constexpr flags for code elimination**: `HAS_WEIGHT`, `HAS_SOFTCAPPING`, `HAS_GRADIENTS`, `RETURN_Z_LOSS`, `RETURN_TOKEN_ACCURACY` — each is `tl.constexpr`, so the compiler removes unused code paths entirely.
+- **Chunked vocab iteration**: The kernel loops over vocabulary in `BLOCK_SIZE` chunks: `for i in range(0, n_cols, BLOCK_SIZE)`. This handles arbitrarily large vocabularies without requiring BLOCK_SIZE >= n_cols.
+- **Multiple loss components**: Combines original CE loss, label smoothing loss, and z-loss (for training stability). Each component contributes to both loss and gradient.
@@ -0,0 +1,64 @@
+# Kernel Profile: RMSNorm (Tier 2 — Reduction)
+
+## Identity
+- operation_name: rms_norm
+- function_class_name: LigerRMSNormFunction
+- module_class_name: LigerRMSNorm
+- functional_name: liger_rms_norm
+
+## Classification
+- tier: 2
+- tier_description: reduction
+- closest_existing_kernel: layer_norm (similar pattern with different normalization)
+
+## Forward Pass
+- forward_inputs:
+  - X: shape (B, T, H), input tensor
+  - W: shape (H,), weight tensor
+  - eps: float, epsilon for numerical stability
+  - offset: float, weight offset (0.0 for Llama, 1.0 for Gemma)
+  - casting_mode: str, "llama" | "gemma" | "none"
+- forward_outputs:
+  - Y: shape (B, T, H), normalized output
+- forward_computation: Y = (X / RMS(X)) * (W + offset), RMS = sqrt(mean(X^2) + eps)
+- precision_sensitive_ops: [rsqrt]
+
+## Backward Pass
+- backward_saved_tensors: [X, W, RSTD] — RSTD cached from forward to avoid recomputation
+- backward_recompute: none (RSTD is expensive to recompute)
+- gradient_formulas:
+  - dX: rstd * (dY*(W+offset) - (1/N) * rstd^2 * dot(dY*(W+offset), X) * X)
+  - dW: sum over (B,T) of dY * (X * rstd)
+
+## Tiling Strategy
+- grid_dimensions: 1D
+- grid_description: forward uses one program per row `(n_rows,)`, backward uses SM-based partitioning `(sm_count,)` with `rows_per_program`
+- block_size_source: calculate_settings(n_cols)
+- needs_sm_parallelism: true (for dW reduction — each SM accumulates partial dW, then summed)
+
+## Module Parameters
+- module_init_params:
+  - hidden_size: int
+  - eps: float = 1e-6
+  - offset: float = 0.0
+  - casting_mode: str = "llama"
+  - init_fn: str = "ones"
+  - in_place: bool = True
+  - elementwise_affine: bool = True
+- learnable_params:
+  - weight: shape (hidden_size,), init ones or zeros
+
+## Benchmarking
+- benchmark_variable: hidden_size
+- benchmark_x_label: "hidden_size"
+- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]
+- benchmark_providers: ["liger", "huggingface"]
+- benchmark_fixed_config: {M: 4096, eps: 1e-6, dtype: torch.float32}
+
+## Key Patterns
+
+- **RSTD caching**: Forward computes and stores `rstd = rsqrt(mean(X^2) + eps)` — 1 value per row, tiny memory cost, saves 4 ops in backward
+- **SM-based backward**: Weight gradient needs reduction across all rows. Each SM processes `rows_per_program` rows and accumulates into `_dW[sm_id, :]`. Final `dW = _dW.sum(dim=0)`
+- **Casting modes as constexpr**: `casting_mode` is `tl.constexpr` so the compiler eliminates dead branches
+- **In-place backward option**: `in_place=True` writes dX into dY tensor to save memory. Set `False` when dY is needed elsewhere (e.g., Gemma2 residual)
+- **Two kernel variants**: Row-wise for `BLOCK_SIZE > 256 or n_rows < 4096*8`, block-wise otherwise (processes `BLOCK_ROW=16` rows per program for better GPU utilization)
@@ -0,0 +1,56 @@
+# Kernel Profile: SwiGLU (Tier 1 — Element-wise)
+
+## Identity
+- operation_name: swiglu
+- function_class_name: LigerSiLUMulFunction
+- module_class_name: LigerSwiGLUMLP
+- functional_name: liger_swiglu
+
+## Classification
+- tier: 1
+- tier_description: element-wise
+- closest_existing_kernel: geglu (same structure, different activation)
+
+## Forward Pass
+- forward_inputs:
+  - a: shape (B, T, H), gate projection output
+  - b: shape (B, T, H), up projection output
+- forward_outputs:
+  - c: shape (B, T, H), silu(a) * b
+- forward_computation: c = silu(a) * b, where silu(x) = x * sigmoid(x)
+- precision_sensitive_ops: [sigmoid]
+
+## Backward Pass
+- backward_saved_tensors: [a, b] (reshaped to 2D in forward wrapper)
+- backward_recompute: recompute silu(a) and sigmoid(a) in backward
+- gradient_formulas:
+  - da: dc * (silu(a) * (1 - sigmoid(a)) + sigmoid(a)) * b
+  - db: dc * silu(a)
+
+## Tiling Strategy
+- grid_dimensions: 1D
+- grid_description: one program per row, `(n_rows,)`
+- block_size_source: calculate_settings(n_cols)
+- needs_sm_parallelism: false
+
+## Module Parameters
+- module_init_params:
+  - config: HuggingFace model config object
+- learnable_params:
+  - gate_proj: Linear(hidden_size, intermediate_size, bias=False)
+  - up_proj: Linear(hidden_size, intermediate_size, bias=False)
+  - down_proj: Linear(intermediate_size, hidden_size, bias=False)
+
+## Benchmarking
+- benchmark_variable: hidden_size
+- benchmark_x_label: "hidden_size"
+- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]
+- benchmark_providers: ["liger", "torch", "torch_compile"]
+- benchmark_fixed_config: {BT: 4096, dtype: torch.bfloat16}
+
+## Key Patterns
+
+- **Recomputation over saving**: Forward saves `a, b` but backward recomputes `sigmoid(a)` and `silu(a)` — saves memory, sigmoid is cheap
+- **In-place backward**: Writes gradients directly to `a_ptr` and `b_ptr` (the saved tensors) — saves allocation
+- **Float32 for sigmoid**: `a_row` cast to `tl.float32` before `tl.sigmoid`, result cast back via `.cast(b_row.dtype)`
+- **No intermediate allocations**: Forward kernel writes directly to output `c`; backward kernel overwrites saved `a, b`
@@ -0,0 +1,96 @@
+# Generator Agent
+
+Takes a confirmed kernel profile and PyTorch reference, generates all files for a production-ready Liger kernel.
+
+## Pre-Requisites
+
+Before generating, read the existing implementation closest to this kernel's tier:
+- Tier 1 → `src/liger_kernel/ops/swiglu.py` + `src/liger_kernel/transformers/swiglu.py`
+- Tier 2 → `src/liger_kernel/ops/rms_norm.py` + `src/liger_kernel/transformers/rms_norm.py`
+- Tier 3 → `src/liger_kernel/ops/cross_entropy.py` + `src/liger_kernel/transformers/cross_entropy.py`
+
+Also read the corresponding test and benchmark files for the reference kernel.
+
+## Files to Generate
+
+### 1. `src/liger_kernel/ops/{kernel}.py` (NEW)
+
+The core Triton kernel. See [templates/ops-kernel.md](templates/ops-kernel.md).
+
+Key rules:
+- `@triton.jit` for forward and backward kernels
+- `torch.autograd.Function` subclass with `@ensure_contiguous` decorator
+- Use `calculate_settings(n_cols)` for BLOCK_SIZE and num_warps
+- Cast to `tl.float32` for precision-sensitive ops (sigmoid, rsqrt, exp, log, tanh)
+- Use `tl.program_id(0).to(tl.int64)` to avoid overflow
+- Use stride parameters for memory access, not hardcoded offsets
+- Prefer recomputation over saving for backward when computation is cheap
+
+### 2. `src/liger_kernel/transformers/{kernel}.py` (NEW)
+
+The nn.Module wrapper. See [templates/module-wrapper.md](templates/module-wrapper.md).
+
+Key rules:
+- Import from `liger_kernel.ops` (not from the submodule directly)
+- Include `extra_repr` method
+- Match the PyTorch API as closely as possible
+
+### 3. `src/liger_kernel/transformers/functional.py` (MODIFY)
+
+Add functional API. See [templates/functional-api.md](templates/functional-api.md).
+
+Two changes:
+- Add import of the Function class at the top (alphabetical order)
+- Add `liger_{kernel}()` function (alphabetical order among existing functions)
+
+### 4. `src/liger_kernel/ops/__init__.py` (MODIFY)
+
+Add import of the Function class. Insert in alphabetical order among existing imports. Use `# noqa: F401` suffix.
+
+### 5. `src/liger_kernel/transformers/__init__.py` (MODIFY)
+
+Add the Module class in three locations (alphabetical order in each):
+- Direct import at top (among "Always-safe imports")
+- `__all__` list
+- If the Module has a `TYPE_CHECKING` import or `__getattr__` entry needed, add there too
+
+### 6. `test/transformers/test_{kernel}.py` (NEW)
+
+Unit tests. See [templates/unit-test.md](templates/unit-test.md).
+
+Key rules:
+- Include a PyTorch reference implementation (copy from Analyzer output)
+- Test both forward and backward correctness
+- Parametrize over shapes (include non-power-of-2 and small/large sizes)
+- Parametrize over dtypes with appropriate tolerances
+- Use `set_seed()`, `assert_verbose_allclose()`, `infer_device()`, `supports_bfloat16()`
+- Import from `test.utils`, NOT `liger_kernel.test.utils`
+
+### 7. `benchmark/scripts/benchmark_{kernel}.py` (NEW)
+
+Benchmarks. See [templates/benchmark.md](templates/benchmark.md).
+
+Key rules:
+- Define `_setup_{kernel}()` function returning `(input_tensor, layer)`
+- Define `bench_speed_{kernel}()` using `run_speed_benchmark`
+- Define `bench_memory_{kernel}()` using `run_memory_benchmark`
+- Compare "liger" vs "torch" (and optionally "torch_compile") providers
+- Use `get_benchmark_model_config`, `estimate_kernel_peak_memory`, and sweep config utilities
+- Run both speed and memory benchmarks for forward, backward, and full modes
+- **Use a single dtype** from `model.dtype` (typically bfloat16). Do NOT generate separate benchmarks for multiple dtypes — multi-dtype coverage is handled by unit tests, not benchmarks
+
+## Modify Mode
+
+When modifying an existing kernel:
+1. Read the existing kernel files thoroughly
+2. Understand the current implementation patterns
+3. Make targeted changes — do not rewrite entire files
+4. Add new test cases for new functionality
+5. Ensure existing tests still cover original behavior
+
+## Code Style
+
+- Line length 120, double quotes, single imports sorted with isort
+- Follow exact patterns from existing code — do not innovate on style
+- When modifying existing files, insert new entries in **alphabetical order**
+- Import style: one import per line, no multi-imports
@@ -0,0 +1,79 @@
+# Kernel Profile Format
+
+When analyzing an operation, produce a profile with all of these fields.
+
+## Fields
+
+```yaml
+# Identity
+operation_name: snake_case name (e.g., "rms_norm", "dyt", "sparsemax")
+function_class_name: "Liger{PascalCase}Function" (e.g., "LigerRMSNormFunction")
+module_class_name: "Liger{PascalCase}" (e.g., "LigerRMSNorm")
+functional_name: "liger_{snake_case}" (e.g., "liger_rms_norm")
+
+# Classification
+tier: 1 | 2 | 3
+tier_description: "element-wise" | "reduction" | "fused/complex"
+closest_existing_kernel: which existing Liger kernel is most similar
+
+# Forward Pass
+forward_inputs:
+  - name: X
+    shape: "(B, T, H)"
+    dtype: "input dtype"
+  # ... list all inputs with shapes
+forward_outputs:
+  - name: Y
+    shape: "(B, T, H)"
+    dtype: "same as input"
+forward_computation: brief math description (e.g., "Y = X * rsqrt(mean(X^2) + eps) * W")
+precision_sensitive_ops: list ops needing float32 (e.g., [rsqrt, sigmoid, exp])
+
+# Backward Pass
+backward_saved_tensors: what to save_for_backward (e.g., [X, W, RSTD])
+backward_recompute: what to recompute instead of saving (e.g., "recompute silu in backward")
+gradient_formulas:
+  dX: "formula or description"
+  dW: "formula or description"
+
+# Tiling Strategy
+grid_dimensions: 1D | 2D
+grid_description: "one program per row" | "one program per (row_block, col_block)" | etc.
+block_size_source: "calculate_settings(n_cols)" | "custom" with explanation
+needs_sm_parallelism: true | false (for weight gradient reduction in backward)
+
+# Module Parameters
+module_init_params:
+  - name: hidden_size
+    type: int
+  - name: eps
+    type: float
+    default: 1e-6
+  # ... list all __init__ parameters
+learnable_params:
+  - name: weight
+    shape: "(hidden_size,)"
+    init: "ones"
+  # ... list learnable parameters
+
+# Benchmarking
+benchmark_variable: "hidden_size" | "seq_len" | "vocab_size" | etc.
+benchmark_x_label: human-readable label for x-axis
+benchmark_x_values_suggestion: suggested range to sweep
+benchmark_providers: ["liger", "torch"] (add "torch_compile" if applicable)
+benchmark_fixed_config: dict of fixed params (e.g., {BT: 4096, dtype: torch.bfloat16})
+```
+
+## Naming Conventions
+
+| Component | Pattern | Example |
+|-----------|---------|---------|
+| Ops file | `src/liger_kernel/ops/{operation_name}.py` | `ops/dyt.py` |
+| Transformer file | `src/liger_kernel/transformers/{operation_name}.py` | `transformers/dyt.py` |
+| Test file | `test/transformers/test_{operation_name}.py` | `test_dyt.py` |
+| Benchmark file | `benchmark/scripts/benchmark_{operation_name}.py` | `benchmark_dyt.py` |
+| Function class | `Liger{PascalCase}Function` | `LigerDyTFunction` |
+| Module class | `Liger{PascalCase}` | `LigerDyT` |
+| Functional | `liger_{snake_case}` | `liger_dyt` |
+| Forward wrapper | `{snake_case}_forward` or `liger_{snake_case}_fwd` | `liger_dyt_fwd` |
+| Backward wrapper | `{snake_case}_backward` or `liger_{snake_case}_bwd` | `liger_dyt_bwd` |
@@ -0,0 +1,117 @@
+# Benchmark Template
+
+## File: `benchmark/scripts/benchmark_{kernel}.py`
+
+```python
+import os
+import sys
+
+import torch
+
+from benchmark_model_configs import compute_hidden_size_sweep_config
+from benchmark_model_configs import estimate_kernel_peak_memory
+from benchmark_model_configs import get_benchmark_model_config
+from utils import SingleBenchmarkRunInput
+from utils import SingleBenchmarkRunOutput
+from utils import parse_benchmark_script_args
+from utils import run_benchmarks
+from utils import run_memory_benchmark
+from utils import run_speed_benchmark
+
+from liger_kernel.utils import infer_device
+
+device = infer_device()
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+
+def _setup_{kernel}(input: SingleBenchmarkRunInput):
+    """Create input tensor and layer from benchmark config."""
+    # Import both Liger and PyTorch reference implementations
+    from test.transformers.test_{kernel} import Liger{Kernel}
+    from test.transformers.test_{kernel} import Torch{Kernel}
+
+    cfg = input.extra_benchmark_config
+    hidden_size = input.x  # Or whatever the sweep variable is
+
+    x = torch.randn(cfg["BT"], hidden_size, device=device, dtype=cfg["dtype"], requires_grad=True)
+
+    if input.kernel_provider == "liger":
+        layer = Liger{Kernel}(hidden_size=hidden_size).to(device)
+    elif input.kernel_provider == "torch":
+        layer = Torch{Kernel}(hidden_size=hidden_size).to(device)
+    elif input.kernel_provider == "torch_compile":
+        layer = torch.compile(Torch{Kernel}(hidden_size=hidden_size).to(device))
+    else:
+        raise ValueError(f"Invalid provider: {input.kernel_provider} for {kernel}")
+
+    return x, layer
+
+
+def bench_speed_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    x, layer = _setup_{kernel}(input)
+    return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x])
+
+
+def bench_memory_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:
+    x, layer = _setup_{kernel}(input)
+    return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode)
+
+
+BT = 4096  # Fixed batch*seq for sweeping hidden_size
+
+
+if __name__ == "__main__":
+    args = parse_benchmark_script_args()
+    model = get_benchmark_model_config(args.model)
+
+    def _probe():
+        probe_input = SingleBenchmarkRunInput(
+            x=model.hidden_size,
+            kernel_provider="torch",
+            extra_benchmark_config={"BT": BT, "dtype": model.dtype},
+        )
+        x, layer = _setup_{kernel}(probe_input)
+        return layer(x)
+
+    peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
+    sweep_config = compute_hidden_size_sweep_config(model, peak_bytes, bt=BT)
+    x_values = [1024 * i for i in range(1, 17) if 1024 * i <= sweep_config.max_hidden_size] or [model.hidden_size]
+
+    common_configs = {
+        "kernel_name": "{kernel}",
+        "x_name": "hidden_size",
+        "x_label": "hidden_size",
+        "x_values": x_values,
+        "kernel_providers": ["liger", "torch", "torch_compile"],
+        "extra_benchmark_configs": [{"BT": sweep_config.bt, "dtype": model.dtype}],
+        "overwrite": args.overwrite,
+    }
+
+    run_benchmarks(
+        bench_test_fn=bench_speed_{kernel},
+        kernel_operation_modes=["full", "forward", "backward"],
+        metric_name="speed",
+        metric_unit="ms",
+        **common_configs,
+    )
+    run_benchmarks(
+        bench_test_fn=bench_memory_{kernel},
+        kernel_operation_modes=["full", "forward", "backward"],
+        metric_name="memory",
+        metric_unit="MB",
+        **common_configs,
+    )
+```
+
+### Key Rules
+
+1. **Import both implementations from test file** — avoids duplicating reference code
+2. **Support "liger", "torch", and "torch_compile" providers**
+3. **Use `run_speed_benchmark` and `run_memory_benchmark`** from utils — don't write custom measurement
+4. **Use `get_benchmark_model_config`** and sweep utilities for standard model sizes
+5. **Use `estimate_kernel_peak_memory`** to auto-compute safe x_values range
+6. **Run both speed AND memory** for all three modes (forward, backward, full)
+7. **Use `BT = 4096`** as default batch*seq when sweeping hidden_size
+8. **Adapt sweep variable** — some kernels sweep `seq_len` or `vocab_size` instead of `hidden_size`. Use the appropriate `compute_*_sweep_config` utility.
+9. **Use a single dtype** — use `model.dtype` from the benchmark model config (typically `torch.bfloat16`). Do NOT generate separate benchmarks for multiple dtypes. Multi-dtype testing is handled by the unit tests, not benchmarks.
@@ -0,0 +1,34 @@
+# Functional API Template
+
+## File: `src/liger_kernel/transformers/functional.py` (MODIFY)
+
+### Step 1: Add Import
+
+Add at the top of the file, in alphabetical order among existing imports:
+
+```python
+from liger_kernel.ops import Liger{Kernel}Function
+```
+
+### Step 2: Add Function
+
+Add in alphabetical order among existing functions:
+
+```python
+def liger_{kernel}(x, ...params):
+    return Liger{Kernel}Function.apply(x, ...params)
+```
+
+For functions with complex return types (like CrossEntropy), wrap the output:
+
+```python
+def liger_{kernel}(input, target, ...params):
+    result = Liger{Kernel}Function.apply(input, target, ...params)
+    return result
+```
+
+### Key Rules
+
+- Function signature should match the PyTorch equivalent where applicable
+- Use positional args for tensor inputs, keyword args for config params
+- Keep the function body minimal — just call `.apply()` and return
@@ -0,0 +1,33 @@
+# Module Wrapper Template
+
+## File: `src/liger_kernel/transformers/{kernel}.py`
+
+```python
+import torch
+import torch.nn as nn
+
+from liger_kernel.ops import Liger{Kernel}Function
+
+
+class Liger{Kernel}(nn.Module):
+    def __init__(self, ...params):
+        super().__init__()
+        # Store config
+        self.param = param
+        # Create learnable parameters
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, x):
+        return Liger{Kernel}Function.apply(x, self.weight, ...self.params)
+
+    def extra_repr(self):
+        return f"param={self.param}, ..."
+```
+
+### Key Rules
+
+- Import from `liger_kernel.ops` (vendor-replaceable), not from the submodule
+- Match the PyTorch reference API as closely as possible (same __init__ params, same forward signature)
+- Include `extra_repr` for debugging
+- For loss modules, name the class `Liger{Kernel}Loss` and the forward takes `(input, target)`
+- For model-specific variants, create subclasses (see `LigerRMSNormForGemma` in `rms_norm.py`)
@@ -0,0 +1,169 @@
+# Ops Kernel Template
+
+## File: `src/liger_kernel/ops/{kernel}.py`
+
+Before writing, read the reference kernel for this tier in the actual codebase.
+
+### Import Block
+
+```python
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.utils import calculate_settings
+from liger_kernel.ops.utils import ensure_contiguous
+```
+
+Add additional imports only as needed:
+- `math` — if using `math.ceil` for row partitioning
+- `from liger_kernel.ops.utils import compare_version` — if checking triton version
+- `from liger_kernel.ops.utils import torch_to_triton_dtype` — if passing dtype as kernel arg
+- `rsqrt` from triton libdevice — if using inverse square root (see `rms_norm.py` for the import pattern with version guards)
+- `tanh` from triton libdevice — if using tanh (see `dyt.py` for the import pattern)
+
+### Forward Kernel
+
+```python
+@triton.jit
+def _{kernel}_forward_kernel(
+    # Output pointers first, then input pointers
+    Y_ptr, Y_stride,
+    X_ptr, X_stride,
+    # ... other input pointers and strides
+    n_cols,
+    # constexpr params for compile-time optimization
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_idx = tl.program_id(0).to(tl.int64)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    # Calculate base pointers using strides
+    X_ptr += row_idx * X_stride
+    Y_ptr += row_idx * Y_stride
+
+    # Load with mask and float32 upcast for precision
+    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0).to(tl.float32)
+
+    # === FORWARD COMPUTATION ===
+
+    tl.store(Y_ptr + col_offsets, result, mask=mask)
+```
+
+### Backward Kernel
+
+```python
+@triton.jit
+def _{kernel}_backward_kernel(
+    # Gradient output, then gradient inputs, then saved tensors
+    dY_ptr, dY_stride,
+    dX_ptr, dX_stride,
+    # ... saved tensors from forward
+    n_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_idx = tl.program_id(0).to(tl.int64)
+    col_offsets = tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < n_cols
+
+    # Load gradient and saved tensors
+    # === BACKWARD COMPUTATION ===
+
+    tl.store(dX_ptr + row_idx * dX_stride + col_offsets, dX_row, mask=mask)
+```
+
+### Python Wrappers
+
+```python
+def {kernel}_forward(X, ...):
+    ori_shape = X.shape
+    n_cols = ori_shape[-1]
+    X = X.view(-1, n_cols)
+    n_rows = X.shape[0]
+
+    Y = torch.empty_like(X)
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    _{kernel}_forward_kernel[(n_rows,)](
+        Y, Y.stride(-2),
+        X, X.stride(-2),
+        n_cols=n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return Y.view(*ori_shape)
+
+
+def {kernel}_backward(dY, ...saved):
+    ori_shape = dY.shape
+    n_cols = ori_shape[-1]
+    dY = dY.view(-1, n_cols)
+    n_rows = dY.shape[0]
+
+    dX = torch.empty_like(dY)
+    BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+
+    _{kernel}_backward_kernel[(n_rows,)](
+        dY, dY.stride(-2),
+        dX, dX.stride(-2),
+        n_cols=n_cols,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return dX.view(*ori_shape)
+```
+
+### Autograd Function
+
+```python
+class Liger{Kernel}Function(torch.autograd.Function):
+    @staticmethod
+    @ensure_contiguous
+    def forward(ctx, X, ...params):
+        Y = {kernel}_forward(X, ...params)
+        ctx.save_for_backward(...)  # Save only what backward needs
+        return Y
+
+    @staticmethod
+    @ensure_contiguous
+    def backward(ctx, dY):
+        ...saved = ctx.saved_tensors
+        dX = {kernel}_backward(dY, ...saved)
+        # Return one gradient per forward input, None for non-tensor params
+        return dX, None, None, ...
+```
+
+### Tier-Specific Patterns
+
+**Tier 1 (element-wise):**
+- Grid: `(n_rows,)` — one program per row
+- Simple 1D indexing with `col_offsets` and `mask`
+- Prefer recomputation in backward over saving activations
+- See `swiglu.py`: saves `a, b` but recomputes `silu(a)` in backward
+
+**Tier 2 (reduction):**
+- Forward may cache intermediate state (e.g., RSTD in RMSNorm)
+- Backward may use SM-based parallelism for weight gradient reduction:
+  ```python
+  sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count
+  _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)
+  grid = (sm_count,)
+  rows_per_program = math.ceil(n_rows / sm_count)
+  ```
+- See `rms_norm.py` for the full pattern
+
+**Tier 3 (fused/complex):**
+- Multi-pass algorithms (e.g., online softmax: pass 1 finds max+sum, pass 2 computes gradients)
+- Gradient-in-forward trick: compute gradients during forward, store in input tensor
+- Multiple `tl.constexpr` flags for compile-time dead code elimination
+- See `cross_entropy.py` for the full pattern
+
+### Key Rules
+
+1. Always use `tl.program_id(0).to(tl.int64)` (prevents overflow on large tensors)
+2. Always use `mask = col_offsets < n_cols` for boundary handling
+3. Use `tl.constexpr` for parameters that enable compile-time optimization
+4. Use stride parameters from tensors, never hardcode memory layout
+5. Cast precision-sensitive ops to `tl.float32`, cast back for storage
+6. Return `None` for each non-tensor parameter in backward
@@ -0,0 +1,146 @@
+# Unit Test Template
+
+## File: `test/transformers/test_{kernel}.py`
+
+```python
+import pytest
+import torch
+
+from test.utils import assert_verbose_allclose
+from test.utils import infer_device
+from test.utils import set_seed
+from test.utils import supports_bfloat16
+
+from liger_kernel.ops import Liger{Kernel}Function
+from liger_kernel.transformers.{kernel} import Liger{Kernel}
+from liger_kernel.transformers.functional import liger_{kernel}
+
+
+# ---- PyTorch Reference ----
+# Paste the standalone PyTorch reference from the Analyzer stage here.
+
+class Torch{Kernel}(torch.nn.Module):
+    def __init__(self, ...):
+        super().__init__()
+        ...
+
+    def forward(self, x):
+        ...
+
+
+set_seed(42)
+device = infer_device()
+
+
+# ---- Correctness Test (Module) ----
+
+@pytest.mark.parametrize(
+    "B, T, hidden_size",
+    [
+        (2, 8, 4096),
+        (4, 16, 2048),
+        (1, 1, 1023),   # Non-power-of-2
+        (3, 7, 256),    # Small + prime numbers
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        (torch.float32, 1e-5, 1e-5),
+        pytest.param(
+            torch.bfloat16,
+            1e-1,
+            5e-2,
+            marks=pytest.mark.skipif(
+                not supports_bfloat16(),
+                reason="bfloat16 not supported on this GPU",
+            ),
+        ),
+    ],
+)
+def test_liger_{kernel}_correctness(B, T, hidden_size, dtype, atol, rtol):
+    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)
+
+    x1 = _input.clone().requires_grad_(True)
+    x2 = _input.clone().requires_grad_(True)
+
+    # Initialize both with same weights
+    torch_layer = Torch{Kernel}(...).to(device).to(dtype)
+    liger_layer = Liger{Kernel}(...).to(device).to(dtype)
+    # Copy weights from torch to liger
+    # liger_layer.weight.data = torch_layer.weight.data.clone()
+
+    # Forward
+    torch_output = torch_layer(x1)
+    liger_output = liger_layer(x2)
+    assert_verbose_allclose(torch_output, liger_output, rtol=rtol, atol=atol, extra_info="[output]")
+
+    # Backward
+    grad_output = torch.randn_like(_input)
+    torch_output.backward(grad_output)
+    liger_output.backward(grad_output)
+    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol, extra_info="[input.grad]")
+
+    # Check weight gradients if applicable
+    # assert_verbose_allclose(
+    #     torch_layer.weight.grad, liger_layer.weight.grad,
+    #     rtol=rtol, atol=atol, extra_info="[weight.grad]"
+    # )
+
+
+# ---- Functional API Test ----
+
+@pytest.mark.parametrize(
+    "B, T, hidden_size",
+    [
+        (2, 8, 4096),
+        (1, 1, 1023),
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype, atol, rtol",
+    [
+        (torch.float32, 1e-5, 1e-5),
+        pytest.param(
+            torch.bfloat16,
+            1e-8,
+            5e-2,
+            marks=pytest.mark.skipif(
+                not supports_bfloat16(),
+                reason="bfloat16 not supported on this GPU",
+            ),
+        ),
+    ],
+)
+def test_liger_{kernel}_functional(B, T, hidden_size, dtype, atol, rtol):
+    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)
+
+    x1 = _input.clone().requires_grad_(True)
+    x2 = _input.clone().requires_grad_(True)
+
+    # Test that functional and Function.apply produce identical results
+    output1 = liger_{kernel}(x1, ...)
+    output2 = Liger{Kernel}Function.apply(x2, ...)
+
+    assert_verbose_allclose(output1, output2, rtol=rtol, atol=atol)
+
+    grad_output = torch.randn_like(_input)
+    output1.backward(grad_output)
+    output2.backward(grad_output)
+    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol)
+```
+
+### Key Testing Rules
+
+1. **Always test both forward AND backward** — backward bugs are common
+2. **Include non-power-of-2 shapes** — tests the mask boundary handling
+3. **Parametrize dtypes with per-dtype tolerances**:
+   - `float32`: `atol=1e-5, rtol=1e-5` (tight)
+   - `bfloat16`: `atol=1e-1, rtol=5e-2` (loose — bf16 has limited precision)
+4. **Use `set_seed(42)` at module level** for reproducibility
+5. **Use `infer_device()`** not `"cuda"` — supports multi-vendor
+6. **Guard bfloat16 with `supports_bfloat16()`** — skips on older GPUs
+7. **Use `assert_verbose_allclose`** not `torch.testing.assert_close` — gives better error output
+8. **Test weight gradients** when the kernel has learnable parameters
+9. **Add `@pytest.mark.flaky(reruns=3)`** if the test is sensitive to random seeds
+10. **Add extra test functions** for edge cases specific to the kernel
@@ -0,0 +1,108 @@
+# Validator Agent
+
+Validates the generated kernel through checkstyle, testing, benchmarking, and optional profiling.
+
+## Validation Flow
+
+### Step 1: Checkstyle (Fast Gate)
+
+```bash
+make checkstyle
+```
+
+If it fails:
+```bash
+ruff check . --fix && ruff format .
+make checkstyle
+```
+
+If it still fails after auto-fix, identify and fix the remaining issues manually, then re-run.
+
+### Step 2: Unit Tests (Hard Gate)
+
+```bash
+python -m pytest test/transformers/test_{kernel}.py -xvs
+```
+
+If tests fail:
+1. Read the error output carefully
+2. Identify root cause: kernel logic bug, shape mismatch, dtype issue, or tolerance too tight
+3. Fix the issue in the relevant file(s)
+4. Re-run `make checkstyle` (in case fix introduced style issues)
+5. Re-run the failing test
+
+**Retry up to 3 times total.** If tests still fail after 3 attempts:
+- **STOP IMMEDIATELY. Do NOT proceed to benchmarks.**
+- Report to the user:
+  - Which tests are failing and the exact error messages
+  - What fixes were attempted
+  - The current state of the code
+  - Suggest what might be wrong and ask for guidance
+
+### Step 3: Benchmarks
+
+Run the benchmark script:
+```bash
+cd benchmark/scripts && python benchmark_{kernel}.py
+```
+
+This produces speed and memory measurements for forward, backward, and full modes, comparing Liger vs PyTorch baseline.
+
+### Step 4: Generate Plots
+
+```bash
+python benchmark/benchmarks_visualizer.py
+```
+
+Plots are saved to `benchmark/visualizations/`. If the visualizer fails or is not available, this step is non-blocking — report the raw numbers instead.
+
+### Step 5: Optional ncu Profiling
+
+Only if the user explicitly requests it. Requires NVIDIA Nsight Compute (`ncu`) to be installed.
+
+Profile both the Liger kernel and PyTorch reference:
+```bash
+ncu --set full --target-processes all -o liger_profile python -c "
+import torch
+from liger_kernel.ops.{kernel} import Liger{Kernel}Function
+x = torch.randn(..., device='cuda', dtype=torch.float32, requires_grad=True)
+out = Liger{Kernel}Function.apply(x, ...)
+out.backward(torch.randn_like(out))
+"
+```
+
+Report key metrics: SM occupancy, memory throughput, compute throughput.
+
+If ncu is not available, report this and skip.
+
+### Step 6: Report Results
+
+Present a clear summary to the user:
+
+```
+## Results
+
+**Checkstyle:** PASS
+**Unit Tests:** PASS (X tests passed)
+
+**Speed Benchmarks:**
+| Mode     | Liger (ms) | PyTorch (ms) | Speedup |
+|----------|-----------|-------------|---------|
+| Forward  | ...       | ...         | ...x    |
+| Backward | ...       | ...         | ...x    |
+| Full     | ...       | ...         | ...x    |
+
+**Memory Benchmarks:**
+| Mode     | Liger (MB) | PyTorch (MB) | Savings |
+|----------|-----------|-------------|---------|
+| Forward  | ...       | ...         | ...%    |
+| Backward | ...       | ...         | ...%    |
+| Full     | ...       | ...         | ...%    |
+
+**Assessment:** [Which dimension improved — speed, memory, or both.
+Flag if either metric is catastrophically worse.]
+
+**Plots:** benchmark/visualizations/
+
+**ncu Profiling:** [Results if requested, or "Not requested"]
+```