{"owner":"linkedin","repo":"Liger-Kernel","hasSkills":true,"totalSkillsCount":33,"totalTokensCount":44015,"categories":["anthropic-skill","root-instruction"],"hasMcp":false,"mcpConfig":null,"found":[".agents/skills/liger-autopatch/SKILL.md",".agents/skills/liger-autopatch/code-generator.md",".agents/skills/liger-autopatch/decision-matrix.md",".agents/skills/liger-autopatch/examples/gemma-profile.md",".agents/skills/liger-autopatch/examples/llama-profile.md",".agents/skills/liger-autopatch/model-analyzer.md",".agents/skills/liger-autopatch/templates/lce-forward-dense.md",".agents/skills/liger-autopatch/templates/lce-forward-moe.md",".agents/skills/liger-autopatch/templates/monkey-patch-fn.md",".agents/skills/liger-autopatch/templates/test-convergence.md",".agents/skills/liger-autopatch/templates/test-instance-patch.md",".agents/skills/liger-autopatch/validator.md",".agents/skills/liger-kernel-dev/SKILL.md",".agents/skills/liger-kernel-dev/analyzer.md",".agents/skills/liger-kernel-dev/examples/cross-entropy-profile.md",".agents/skills/liger-kernel-dev/examples/rms-norm-profile.md",".agents/skills/liger-kernel-dev/examples/swiglu-profile.md",".agents/skills/liger-kernel-dev/generator.md",".agents/skills/liger-kernel-dev/kernel-profile-format.md",".agents/skills/liger-kernel-dev/templates/benchmark.md",".agents/skills/liger-kernel-dev/templates/functional-api.md",".agents/skills/liger-kernel-dev/templates/module-wrapper.md",".agents/skills/liger-kernel-dev/templates/ops-kernel.md",".agents/skills/liger-kernel-dev/templates/unit-test.md",".agents/skills/liger-kernel-dev/validator.md",".agents/skills/liger-kernel-perf/SKILL.md",".agents/skills/liger-kernel-perf/finalizer.md",".agents/skills/liger-kernel-perf/optimization-strategies.md",".agents/skills/liger-kernel-perf/optimizer.md",".agents/skills/liger-kernel-perf/profiler.md",".agents/skills/liger-kernel-perf/templates/optimization-profile.md",".agents/skills/liger-kernel-perf/templates/variant-notes.md","AGENTS.md"],"skills":{".agents/skills/liger-autopatch/SKILL.md":"---\nname: liger-autopatch\ndescription: \"Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies existing monkey-patching. Generates lce_forward, monkey-patch function, tests, and README entry. Use when adding a new model to Liger Kernel, when a user asks to patch an unsupported model, when extending MODEL_TYPE_TO_APPLY_LIGER_FN, or when modifying/updating/fixing an existing monkey-patch (e.g., adding a new kernel to an already-supported model, fixing instance patching, updating a patch for upstream HF changes).\"\n---\n\n# Liger Auto-Patch\n\nAdds Liger Kernel optimization support for a new HuggingFace model, or modifies existing monkey-patching, through a staged pipeline with human review between stages. Supports creating new model patches and modifying existing ones.\n\n## Mode Detection\n\n- **Create mode**: User asks to add/patch/support a new model → full pipeline (Analyze → Generate → Validate)\n- **Modify mode**: User asks to update/fix/change/extend an existing monkey-patch → lighter pipeline (Change Impact Analysis → Apply Changes → Validate)\n\nKeywords that suggest modify mode: update, fix, change, add [kernel] to [existing model], extend, modify, new activation, new norm, bug in patch, upstream changed\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Model Analyzer** workflow in [model-analyzer.md](model-analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nThis stage reads the HF `modeling_*.py` source and produces a **model profile** answering 12 architectural questions from [decision-matrix.md](decision-matrix.md).\n\n**Human checkpoint:** Present the profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md).\n\nGenerates/modifies up to 13 files:\n\n1. `src/liger_kernel/transformers/model/{model}.py` — NEW lce_forward\n2. `src/liger_kernel/transformers/monkey_patch.py` — MODIFY\n3. `src/liger_kernel/transformers/__init__.py` — MODIFY\n4. `src/liger_kernel/transformers/model/output_classes.py` — MODIFY if needed\n5. `test/transformers/test_monkey_patch.py` — MODIFY\n6. `test/convergence/bf16/test_mini_models.py` — MODIFY (FLCE path)\n7. `test/convergence/bf16/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n8. `test/convergence/fp32/test_mini_models.py` — MODIFY (FLCE path)\n9. `test/convergence/fp32/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n10. `test/convergence/bf16/test_mini_models_multimodal.py` — MODIFY if VL model\n11. `test/convergence/fp32/test_mini_models_multimodal.py` — MODIFY if VL model\n12. `test/utils.py` — MODIFY\n13. `README.md` — MODIFY\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns instance patching test, convergence test, and lint check. Retries up to 3 times on failure.\n\n**Human checkpoint:** Report final test results.\n\n## Pipeline (Modify Mode)\n\n### Stage 1: Change Impact Analysis\n\nRead the existing `apply_liger_kernel_to_{model_type}` function in `monkey_patch.py` and the relevant section of the upstream HF `modeling_{model_type}.py`. Produce a short change plan:\n\n- What is being added/changed/fixed\n- Which Liger kernel(s) are involved\n- Which files need modification (subset of the 13 files from create mode)\n- What the expected behavior should be after the change\n\n**Human checkpoint:** Present the change plan. Confirm before proceeding.\n\n### Stage 2: Apply Changes\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md) in **modify mode**.\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md). This stage is **mandatory** — do not skip it. At minimum, run:\n\n1. Instance patching test: `pytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs`\n2. All convergence tests for the model:\n   - `pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, bf16)\n   - `pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, bf16)\n   - `pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, fp32)\n   - `pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, fp32)\n   - If VL (multimodal) model, also run:\n     - `pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n     - `pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n3. Checkstyle: `make checkstyle`\n\n**Human checkpoint:** Report final test results.\n\n## Reference Files\n\n- [decision-matrix.md](decision-matrix.md) — 12 architectural decisions to resolve per model\n- [examples/llama-profile.md](examples/llama-profile.md) — Reference profile for standard dense model\n- [examples/gemma-profile.md](examples/gemma-profile.md) — Reference profile showing GeGLU + offset variant\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n",".agents/skills/liger-autopatch/code-generator.md":"# Code Generator Agent\n\nTakes a confirmed model profile (create mode) or a change plan (modify mode) and generates or modifies files for Liger Kernel support.\n\n## Mode\n\n- **Create mode** (default): Generating all files for a new model. Follow the full \"Files to Generate\" list below.\n- **Modify mode**: Making targeted changes to an existing monkey-patch. Follow the \"Modification Checklist\" section instead.\n\n## Pre-Requisites\n\nBefore generating, read the reference implementation closest to this model:\n- Dense → `src/liger_kernel/transformers/model/llama.py`\n- MoE → `src/liger_kernel/transformers/model/mixtral.py`\n- Vision-Language → `src/liger_kernel/transformers/model/qwen2_vl.py`\n- Gemma-family → `src/liger_kernel/transformers/model/gemma.py`\n\nAlso read the corresponding patching function in `monkey_patch.py` and the templates in [templates/](templates/).\n\n## Files to Generate\n\n### 1. `src/liger_kernel/transformers/model/{model_type}.py` (NEW)\n\nThe `lce_forward` function. See [templates/lce-forward-dense.md](templates/lce-forward-dense.md) or [templates/lce-forward-moe.md](templates/lce-forward-moe.md).\n\nKey rules:\n- Match the exact forward signature from HF's `ForCausalLM.forward`\n- Use `lce_maybe_trainable_lm_head` from `llama.py` (shared PEFT/FSDP utility)\n- If model needs custom loss args (e.g., softcapping), write a local helper instead\n\n### 2. `src/liger_kernel/transformers/monkey_patch.py` (MODIFY)\n\nThree changes — see [templates/monkey-patch-fn.md](templates/monkey-patch-fn.md):\n\n**A.** Add lce_forward import (~line 18-28):\n```python\nfrom liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n```\n\n**B.** Add `apply_liger_kernel_to_{model_type}` function with both class-level and instance-level patching paths.\n\n**C.** Add entry to `MODEL_TYPE_TO_APPLY_LIGER_FN` dict (~line 3067).\n\n### 3. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the function in three locations (maintain alphabetical order):\n- `TYPE_CHECKING` block\n- `__getattr__` monkey_patch_symbols set\n- `__all__` list extension\n\n### 4. `src/liger_kernel/transformers/model/output_classes.py` (MODIFY if needed)\n\nOnly for models needing custom output (MoE with `aux_loss`, VL with `rope_deltas`). Follow the existing guarded-import pattern in the file.\n\n### 5. `test/transformers/test_monkey_patch.py` (MODIFY)\n\nSee [templates/test-instance-patch.md](templates/test-instance-patch.md). Add availability checker + skipif-decorated test function using `inspect.getsource()` assertions.\n\n### 6. Convergence tests (MODIFY multiple files)\n\nSee [templates/test-convergence.md](templates/test-convergence.md). Every model needs entries in multiple convergence test files:\n\n**All text models (dense + MoE)** — add to these 4 files:\n- `test/convergence/bf16/test_mini_models.py` — FLCE path, bf16\n- `test/convergence/bf16/test_mini_models_with_logits.py` — non-FLCE path (tests RMSNorm/SwiGLU/RoPE only), bf16\n- `test/convergence/fp32/test_mini_models.py` — FLCE path, fp32\n- `test/convergence/fp32/test_mini_models_with_logits.py` — non-FLCE path, fp32\n\n**Vision-language models** — also add to these 2:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\nEach file needs: imports, availability guard, `MiniModelConfig` entry in `MINI_MODEL_SETUPS` dict, and a `pytest.param` entry in the parametrize block. The `MiniModelConfig` entry is identical across all files for the same model. The `pytest.param` tolerances differ — use bf16 tolerances (looser) for bf16 files and fp32 tolerances (tighter) for fp32 files. Copy tolerance values from a similar existing model (e.g., Llama for dense, Mixtral for MoE).\n\n### 7. `test/utils.py` (MODIFY)\n\nAdd `revert_liger_kernel_to_{model_type}` function that reloads the modeling module.\n\n### 8. `README.md` (MODIFY)\n\nAdd row to the Patching table under \"### Patching\":\n```\n| {ModelName} | `liger_kernel.transformers.apply_liger_kernel_to_{model_type}` | {Supported Operations} |\n```\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order** alongside similar existing entries. Never append to the end of a section — find the correct alphabetical position.\n- After generating all files, run `make checkstyle` to verify formatting. If it fails, run `ruff check . --fix && ruff format .` to auto-fix, then verify with `make checkstyle` again.\n\n## Modification Checklist (Modify Mode)\n\nBefore making changes, read the existing implementation:\n1. Read `apply_liger_kernel_to_{model_type}` in `monkey_patch.py`\n2. Read the existing test in `test_monkey_patch.py` for this model\n3. Read the relevant HF modeling source for context\n\n### Rules for All Modifications\n\n**R1. Both patching levels.** If adding a new kernel, it must appear in BOTH:\n  - Class-level patching (the main body of `apply_liger_kernel_to_{model_type}`)\n  - Instance-level patching (the `if model is not None` block)\n\n  Omitting one is the most common mistake.\n\n**R2. New parameter with default.** Every new kernel gets a bool parameter on the\n  apply function signature (e.g., `relu_squared: bool = True`). Default should be `True`\n  for kernels that are safe to enable by default, `False` otherwise.\n\n**R3. Update docstring.** Update the function's docstring to:\n  - Add an `Args` entry for the new parameter\n  - Remove any stale notes that the new kernel invalidates\n    (e.g., \"squared ReLU is not supported\" → remove if you're adding it)\n\n**R4. Update tests.** In the existing `test_apply_liger_kernel_to_instance_for_{model_type}`:\n  - Add import for the new Liger kernel class\n  - Add \"not yet patched\" assertion before `_apply_liger_kernel_to_instance`\n  - Add \"correctly patched\" assertion after\n  - Follow the exact pattern of existing assertions in the same test\n\n**R5. Check revert function.** Read `revert_liger_kernel_to_{model_type}` in `test/utils.py`.\n  The revert function uses `importlib.reload(modeling_{model_type})` to undo all patches.\n  This handles most cases automatically, but check if the new kernel requires additional\n  revert logic (e.g., if the kernel patches something outside the modeling module, or\n  replaces a global like `ACT2FN` that `importlib.reload` won't fully restore). Update\n  the revert function if needed.\n\n**R6. Run convergence tests.** Don't modify convergence test files unless the change\n  requires it (e.g., new mini model config fields). But DO run existing convergence\n  tests in the Validate stage to verify no regression. This is critical — the Validator\n  agent (Stage 3) handles this, but if you are generating code without a separate\n  Validate stage, run these yourself:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  ```\n  For VL (multimodal) models, also run:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  ```\n\n**R7. Update README.md.** If the change adds a visibly new capability to the model's\n  row in the patching table (e.g., a new operation), update the supported operations list.\n\n### Common Modification Patterns\n\n**Adding an activation kernel (e.g., relu_squared for nemotron):**\n- Import the Liger kernel class at the top of `monkey_patch.py`\n- Add bool parameter to apply function signature\n- Class-level: replace in `ACT2FN` dict or replace the MLP class\n- Instance-level: patch each `decoder_layer`'s activation/MLP\n- Test: `assert isinstance` checks on the activation/MLP\n\n**Adding a norm variant:**\n- Add bool parameter to apply function\n- Class-level: replace the norm class\n- Instance-level: use `_patch_rms_norm_module` or `_patch_layer_norm_module` on all norm attrs\n- Test: `assert isinstance` checks on norm modules\n\n**Fixing missing instance patching:**\n- Read the class-level patching to see what's patched\n- Add corresponding instance-level patches in the `if model is not None` block\n- Test: add assertions that were missing\n\n**Updating for upstream HF changes:**\n- Compare the current HF modeling file against what the patch assumes\n- Update class names, attribute names, forward signatures as needed\n- May require updating `lce_forward` if the base model's forward changed\n",".agents/skills/liger-autopatch/decision-matrix.md":"# Decision Matrix\n\nWhen analyzing a HuggingFace model for Liger Kernel support, you must resolve these 12 architectural decisions by reading the model's `modeling_*.py` source code.\n\n## 1. Norm Type\n\n**Question:** Does the model use RMSNorm, LayerNorm, or both?\n\n**How to detect:**\n- Search for `class *RMSNorm` in the modeling file → RMSNorm\n- Search for `nn.LayerNorm` usage → LayerNorm\n- Multimodal models often use both (RMSNorm for text, LayerNorm for vision)\n\n**Liger mapping:** `LigerRMSNorm` or `LigerLayerNorm`\n\n## 2. RMSNorm Casting Mode\n\n**Question:** How does the model handle dtype casting during normalization?\n\n**How to detect:** Read the RMSNorm forward method:\n- Casts input to fp32, computes variance, casts back → `\"gemma\"`\n- Computes variance in fp32 only (input stays original dtype) → `\"llama\"`\n- No casting at all → `\"none\"`\n\n**Default:** `\"llama\"` (most common)\n\n## 3. RMSNorm Offset\n\n**Question:** Does the weight have a +1.0 offset?\n\n**How to detect:** In the RMSNorm forward, look for `(1 + self.weight)` or `self.weight + 1`:\n- Present → `offset=1.0` (Gemma family)\n- Absent → `offset=0.0` (most models)\n\n## 4. RMSNorm In-Place\n\n**Question:** Can the backward pass modify dY in-place?\n\n**How to detect:** Check if the model has two sequential norm layers with a residual connection between them (like Gemma2's `pre_feedforward_layernorm` + `post_feedforward_layernorm`):\n- Sequential norms with residual → `in_place=False`\n- Otherwise → `in_place=True`\n\n## 5. MLP Activation Type\n\n**Question:** What activation function does the gated MLP use?\n\n**How to detect:** Read the MLP class forward method:\n- `silu` or `F.silu` → SwiGLU → `LigerSwiGLUMLP`\n- `gelu` or `gelu_new` or `gelu_fast` → GeGLU → `LigerGEGLUMLP`\n- Phi3-style (single gate+up projection split) → `LigerPhi3SwiGLUMLP`\n\n**Also check:** The config's `hidden_act` field.\n\n## 6. Dense vs MoE\n\n**Question:** Is the model dense, MoE, or hybrid (some layers dense, some MoE)?\n\n**How to detect:**\n- Search for `Expert`, `MoE`, `SparseMoe`, `TopK` routing classes\n- Check if decoder layers have a `block_sparse_moe` or `experts` attribute\n- Hybrid: check for `is_moe_layer` or conditional MoE per-layer\n\n**Liger mapping:**\n- Dense → standard patching\n- MoE (transformers v5) → `LigerExperts`\n- MoE (transformers v4) → `LigerBlockSparseTop2MLP`\n- Qwen3-style MoE → `LigerQwen3MoeSwiGLUMLP`\n\n## 7. Vision Components\n\n**Question:** Does the model have a vision encoder?\n\n**How to detect:**\n- Check for `pixel_values` in the `forward` signature\n- Look for a separate vision model class (e.g., `*VisionModel`)\n- Check config for `vision_config` or `text_config` sub-configs\n\n**If yes:**\n- Vision encoder norms are usually `nn.LayerNorm` → patch with `LigerLayerNorm`\n- Text and vision must be patched separately\n\n## 8. RoPE Variant\n\n**Question:** What type of positional embedding does the model use?\n\n**How to detect:** Search for the `apply_rotary_pos_emb` function:\n- Standard (q, k, cos, sin) → `liger_rotary_pos_emb` (rope=True)\n- Llama4-style → `liger_llama4_text_rotary_pos_emb`\n- Qwen2VL MRoPE → `liger_multimodal_rotary_pos_emb`\n- No rotary embedding or custom variant → `rope=False`\n\n## 9. Output Class\n\n**Question:** What return type does the model's ForCausalLM.forward use?\n\n**How to detect:** Read the return statement and type annotation:\n- Standard → `LigerCausalLMOutputWithPast`\n- MoE (has `aux_loss`) → `LigerMoeCausalLMOutputWithPast`\n- Custom VL output → create model-specific output class in `output_classes.py`\n\n## 10. Hidden State Access\n\n**Question:** How does the model access hidden states from base model output?\n\n**How to detect:** In the ForCausalLM.forward, after calling `self.model(...)`:\n- `outputs[0]` → most models (Llama, Mistral, Gemma, etc.)\n- `outputs.last_hidden_state` → Phi3, Qwen3.5 MoE, some newer models\n\n## 11. Logit Softcapping\n\n**Question:** Does the model apply softcapping to logits before loss?\n\n**How to detect:** Check config for `final_logit_softcapping`:\n- Present → pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`\n- Absent → no softcapping (most models)\n- **VL models:** Config path may be `self.config.text_config.final_logit_softcapping` instead of `self.config.final_logit_softcapping`. Check whether the model uses a composite config with `text_config` sub-config.\n\n**Models with softcapping:** Gemma2, Gemma3\n\n## 12. Decoder Layer Norm Names\n\n**Question:** What are the attribute names of norm layers in each decoder layer?\n\n**How to detect:** Read the decoder layer class `__init__`:\n- Standard: `input_layernorm`, `post_attention_layernorm`\n- Gemma2 extra: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- GLM4: `post_self_attn_layernorm`, `post_mlp_layernorm`\n- Some models: `q_norm`, `k_norm` on self_attn\n\nAlso check the final norm on the base model (usually `model.norm` or `model.final_layernorm`).\n",".agents/skills/liger-autopatch/examples/gemma-profile.md":"# Model Profile: Gemma\n\nThis profile demonstrates the key differences from Llama: GeGLU activation, RMSNorm offset, and Gemma casting mode.\n\n## Identity\n- model_type: gemma\n- causal_lm_class: GemmaForCausalLM\n- base_model_class: GemmaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.gemma.modeling_gemma\n- config_module: transformers.models.gemma.configuration_gemma\n\n## Normalization\n- norm_class: GemmaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: gemma (everything cast to fp32, then computed, then cast back)\n- offset: 1.0 (weight uses `1 + self.weight` pattern)\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: GemmaMLP\n- activation: gelu (uses GELU activation, not SiLU)\n- liger_mlp_class: LigerGEGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nSame as Llama — no extra parameters.\n\n## Mini Model Config\n```python\nGemmaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-6,\n    hidden_activation=\"gelu_pytorch_tanh\",\n)\n```\n\n## Key Differences from Llama\n\n1. **Activation**: Uses `geglu` parameter (not `swiglu`) in the patch function\n2. **RMSNorm**: Requires `offset=1.0` and `casting_mode=\"gemma\"`\n3. **MLP class**: `LigerGEGLUMLP` instead of `LigerSwiGLUMLP`\n4. **Patching uses partial**: `_patch_rms_norm_module_for_gemma = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)`\n\n## Gemma2 Additional Differences\n- `in_place: false` (residual between sequential norms)\n- Extra norm layers: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- Has `final_logit_softcapping` in config\n- Uses `LigerRMSNormForGemma2` variant\n",".agents/skills/liger-autopatch/examples/llama-profile.md":"# Model Profile: Llama\n\nThis is the reference profile for the simplest and most common model type.\n\n## Identity\n- model_type: llama\n- causal_lm_class: LlamaForCausalLM\n- base_model_class: LlamaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.llama.modeling_llama\n- config_module: transformers.models.llama.configuration_llama\n\n## Normalization\n- norm_class: LlamaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: llama\n- offset: 0.0\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: LlamaMLP\n- activation: silu\n- liger_mlp_class: LigerSwiGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nStandard parameters only:\n- input_ids, attention_mask, position_ids, past_key_values\n- inputs_embeds, labels, use_cache\n- output_attentions, output_hidden_states, return_dict\n- cache_position, logits_to_keep, skip_logits\n- **kwargs\n\nNo model-specific extra parameters.\n\n## Mini Model Config\n```python\nLlamaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-5,\n)\n```\n\n## Notes\n- Llama is the base reference for all dense LLaMA-family models\n- Most models (Mistral, Qwen2, OLMo, etc.) follow this exact pattern\n- The `lce_forward` for Llama reuses `lce_maybe_trainable_lm_head` for PEFT/FSDP support\n- Has `pretraining_tp` check (raises error if > 1)\n",".agents/skills/liger-autopatch/model-analyzer.md":"# Model Analyzer Agent\n\nReads a HuggingFace model's source code and produces a structured model profile.\n\n## Input\n\nA model identifier: model_type string (e.g., `\"nemotron\"`), class name (e.g., `\"NemotronForCausalLM\"`), or Hub model name.\n\n## Steps\n\n1. Locate `transformers/models/{model_type}/modeling_{model_type}.py` and `configuration_{model_type}.py` on disk\n2. Read [decision-matrix.md](decision-matrix.md) for the 12 decisions to resolve\n3. Read the modeling file and extract all architectural details\n4. Produce the profile in the format below\n\n## What to Extract\n\nFrom the modeling file, identify:\n\n- `*ForCausalLM`, `*Model`, `*DecoderLayer`, `*RMSNorm`/norm, `*MLP`, `*Expert`/`*MoE` classes\n- `ForCausalLM.forward` signature — note all parameters\n- Norm class `forward` — determine casting mode and offset\n- All norm attribute names on decoder layer and base model\n- MLP activation function and gate/up/down projection names\n- Config's `model_type` string, `final_logit_softcapping`, `pretraining_tp`\n\n## Output Format\n\n```markdown\n# Model Profile: {ModelName}\n\n## Identity\n- model_type: {from config}\n- causal_lm_class: {e.g., NemotronForCausalLM}\n- base_model_class: {e.g., NemotronModel}\n- base_model_prefix: {e.g., \"model\"}\n- modeling_module: transformers.models.{model_type}.modeling_{model_type}\n\n## Normalization\n- norm_type: RMSNorm | LayerNorm | both\n- casting_mode: llama | gemma | none\n- offset: 0.0 | 1.0\n- in_place: true | false\n- final_norm_attr: {e.g., \"model.norm\"}\n- decoder_norm_attrs: {list}\n- attn_norm_attrs: {e.g., \"self_attn.q_norm\" or \"none\"}\n\n## MLP\n- activation: silu | gelu | gelu_new\n- liger_mlp_class: LigerSwiGLUMLP | LigerGEGLUMLP | LigerPhi3SwiGLUMLP\n\n## Structure\n- type: dense | moe | hybrid\n- moe_expert_class: {if MoE}\n- shared_expert: true | false\n\n## Positional Embedding\n- rope_type: standard | llama4 | qwen2vl_mrope | none\n\n## Output\n- output_class: LigerCausalLMOutputWithPast | LigerMoeCausalLMOutputWithPast | custom\n- hidden_state_access: outputs[0] | outputs.last_hidden_state\n- has_logit_softcapping: true | false\n\n## Vision\n- has_vision: true | false\n- vision_norm_type: {if applicable}\n\n## Forward Signature\n{List model-specific parameters not in standard Llama signature}\n\n## Mini Model Config\n{Suggest minimal config values for testing — 2 layers, 32 hidden, 64 intermediate}\n```\n\n## Rules\n\n- Never guess — always read the source to determine each value\n- Read [examples/llama-profile.md](examples/llama-profile.md) and [examples/gemma-profile.md](examples/gemma-profile.md) as reference\n- Compare against existing supported models to find the closest match\n",".agents/skills/liger-autopatch/templates/lce-forward-dense.md":"# Template: lce_forward for Dense Models\n\nReference implementation: `src/liger_kernel/transformers/model/llama.py`. Read it first.\n\n## Structure\n\nThe file needs `lce_forward` and reuses `lce_maybe_trainable_lm_head` from `llama.py`.\n\n### Imports\n\n```python\nfrom typing import TYPE_CHECKING, List, Optional, Tuple, Union\nimport torch\nfrom liger_kernel.transformers.model.llama import lce_maybe_trainable_lm_head\nfrom liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss, unpack_cross_entropy_result\nfrom liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast\nif TYPE_CHECKING:\n    from transformers.cache_utils import Cache\n```\n\n### lce_forward body pattern\n\n1. Set defaults from config (output_attentions, output_hidden_states, return_dict)\n2. Call `self.model(...)` with all forwarded params\n3. Extract hidden states: `outputs[0]` or `outputs.last_hidden_state`\n4. Slice for `logits_to_keep`\n5. Extract `shift_labels = kwargs.pop(\"shift_labels\", None)`\n6. `skip_logits` logic — default True during training with labels\n7. Compute loss via `lce_maybe_trainable_lm_head` (skip_logits) or `self.lm_head` (else)\n8. Return `LigerCausalLMOutputWithPast`\n\n### Variations\n\n- **Logit softcapping**: Pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`. Requires a local helper instead of reusing `lce_maybe_trainable_lm_head`.\n- **Hidden state access**: Use `outputs.last_hidden_state` if HF model uses that\n- **pretraining_tp**: Remove check if model doesn't have this config field\n\n### When NOT to reuse llama's lce_maybe_trainable_lm_head\n\nIf the model passes extra kwargs to `LigerForCausalLMLoss` (like softcapping), write a local `_liger_for_causal_lm_loss` and inline the PEFT/FSDP handling from llama.py.\n",".agents/skills/liger-autopatch/templates/lce-forward-moe.md":"# Template: lce_forward for MoE Models\n\nReference implementation: `src/liger_kernel/transformers/model/mixtral.py`. Read it first.\n\n## Differences from Dense\n\n1. **Output class**: `LigerMoeCausalLMOutputWithPast` (includes `aux_loss`)\n2. **Extra param**: `output_router_logits` in forward signature\n3. **Auxiliary loss**: Compute `load_balancing_loss_func` from router logits\n4. **Return**: Include `aux_loss` and `router_logits` in output\n\n## Auxiliary Loss Pattern\n\n```python\naux_loss = None\nif output_router_logits:\n    router_logits = outputs.router_logits if return_dict else outputs[-1]\n    aux_loss = load_balancing_loss_func(\n        router_logits, self.num_experts, self.num_experts_per_tok, attention_mask,\n    )\n    if labels is not None and loss is not None:\n        loss += self.router_aux_loss_coef * aux_loss\n```\n\n## Monkey Patch Differences for MoE\n\nExpert patching is version-aware:\n\n```python\n# Transformers v5\nmodeling_module.{ModelExperts} = LigerExperts\n# Instance: _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\n\n# Transformers v4\nmodeling_module.{ModelExpertMLP} = LigerBlockSparseTop2MLP\n# Instance: for expert in decoder_layer.block_sparse_moe.experts: _patch_swiglu_module(...)\n```\n\nCheck exact attribute names in the HF decoder layer class.\n",".agents/skills/liger-autopatch/templates/monkey-patch-fn.md":"# Template: apply_liger_kernel_to_{model_type}\n\nReference: `apply_liger_kernel_to_llama` (dense) or `apply_liger_kernel_to_mixtral` (MoE) in `monkey_patch.py`.\n\n## Function Signature\n\n```python\ndef apply_liger_kernel_to_{model_type}(\n    rope: bool = True,              # False if model lacks standard RoPE\n    cross_entropy: bool = False,\n    fused_linear_cross_entropy: bool = True,\n    rms_norm: bool = True,          # Or layer_norm for LayerNorm models\n    swiglu: bool = True,            # Or geglu for Gemma-family\n    model: PreTrainedModel = None,\n) -> None:\n```\n\n## Body Pattern\n\nClass-level patching always runs (even when `model` is provided). Instance-level patching runs additionally when `model is not None`. They are NOT mutually exclusive — this matches the existing Llama/Mistral/Gemma patterns.\n\nThe only exception is `fused_linear_cross_entropy`: use `MethodType` binding when instance is given, class replacement otherwise.\n\n```python\nassert not (cross_entropy and fused_linear_cross_entropy), \"...\"\n\nfrom transformers.models.{model_type} import modeling_{model_type}\n\n# Class-level patching (always runs)\nif rope: modeling_{model_type}.apply_rotary_pos_emb = liger_rotary_pos_emb\nif rms_norm: modeling_{model_type}.{ModelRMSNorm} = LigerRMSNorm\nif swiglu: modeling_{model_type}.{ModelMLP} = LigerSwiGLUMLP\nif cross_entropy: modeling_{model_type}.CrossEntropyLoss = LigerCrossEntropyLoss\nif fused_linear_cross_entropy:\n    if model is not None:\n        model.forward = MethodType({model_type}_lce_forward, model)\n    else:\n        modeling_{model_type}.{ModelForCausalLM}.forward = {model_type}_lce_forward\n\n# Instance-level patching (only when model instance provided)\nif model is not None:\n    base_model = getattr(model, model.base_model_prefix, model)\n    if rms_norm: _patch_rms_norm_module(base_model.norm)\n    for decoder_layer in base_model.layers:\n        if swiglu: _patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)\n        if rms_norm:\n            _patch_rms_norm_module(decoder_layer.input_layernorm)\n            _patch_rms_norm_module(decoder_layer.post_attention_layernorm)\n```\n\n## Common Variations\n\n**Gemma-family** (offset + casting):\n```python\n_patch_rms_norm_module_for_model = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)\n```\n\n**MoE experts** (version-aware):\n```python\nif IS_TRANSFORMERS_V5_OR_LATER:\n    _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\nelse:\n    for expert in decoder_layer.block_sparse_moe.experts:\n        _patch_swiglu_module(expert, LigerBlockSparseTop2MLP)\n```\n\n**Multimodal VL instance patching:**\n\nFor VL models, extract text_model and vision_model from the composite instance using `isinstance` checks. Each VL model has different accessor paths — read the existing VL patches in `monkey_patch.py` for the exact pattern. Example from Llama4:\n\n```python\nif model is not None:\n    if isinstance(model, {Model}ForConditionalGeneration):\n        language_model = model.language_model          # or model.model.language_model\n        vision_model = model.vision_model              # or model.model.visual\n        text_model = language_model.model\n    elif isinstance(model, {Model}ForCausalLM):\n        text_model = model.model\n        vision_model = None\n    elif isinstance(model, {Model}TextModel):\n        text_model = model\n        vision_model = None\n    else:\n        raise ValueError(f\"Unsupported model type: {type(model)}\")\n\n    # Patch text layers\n    if text_model:\n        if rms_norm: _patch_rms_norm_module(text_model.norm)\n        for decoder_layer in text_model.layers:\n            ...\n\n    # Patch vision layers\n    if vision_model:\n        for layer in vision_model.encoder.layers:  # accessor varies by model\n            if layer_norm:\n                _patch_layer_norm_module(layer.layer_norm1)\n                _patch_layer_norm_module(layer.layer_norm2)\n```\n\nThe exact accessor paths (`model.language_model`, `model.model.language_model`, `model.model.visual`) vary per VL model. Always read the HF model's `__init__` to find the correct paths.\n",".agents/skills/liger-autopatch/templates/test-convergence.md":"# Template: Convergence Tests\n\nReference: existing entries across `test/convergence/` files. Read the Llama entries as baseline.\n\n## Which Files to Modify\n\nEvery model needs entries in **at least 4 files**:\n\n| File | What it tests | dtype |\n|------|-------------|-------|\n| `test/convergence/bf16/test_mini_models.py` | FLCE path (fused linear cross entropy) | bfloat16 |\n| `test/convergence/bf16/test_mini_models_with_logits.py` | Non-FLCE path (tests RMSNorm/SwiGLU/RoPE only, verifies logits) | bfloat16 |\n| `test/convergence/fp32/test_mini_models.py` | FLCE path | float32 |\n| `test/convergence/fp32/test_mini_models_with_logits.py` | Non-FLCE path | float32 |\n\n**Vision-language models** also need:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\n## Placement Rule\n\nWhen adding entries to existing files, always insert in **alphabetical order** alongside similar existing entries. Never append to the end of a section. Specifically:\n- Imports: alphabetical among other `apply_liger_kernel_to_*` / `revert_liger_kernel_to_*` imports\n- Availability guards: alphabetical among other try/except model import blocks\n- `MINI_MODEL_SETUPS` entries: alphabetical among other model setup blocks\n- `pytest.param` entries: alphabetical among other model param entries\n\n## What to Add in Each File\n\nEach file has a `MINI_MODEL_SETUPS` dict and a `pytest.param` block. Add to both.\n\n### 1. Imports (same across all files, insert in alphabetical order)\n\n```python\nfrom liger_kernel.transformers import apply_liger_kernel_to_{model_type}\nfrom test.utils import revert_liger_kernel_to_{model_type}\n```\n\n### 2. Availability Guard (same across all files)\n\nThe `try/except ImportError` handles version compatibility automatically: if the model doesn't exist in the installed transformers version (e.g., 4.52.0), the import fails and the tests skip. No explicit version check needed.\n\n```python\ntry:\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n    from transformers.models.{model_type}.modeling_{model_type} import {ModelForCausalLM}\n    {MODEL_UPPER}_AVAILABLE = True\nexcept ImportError:\n    {MODEL_UPPER}_AVAILABLE = False\n```\n\n### 3. MINI_MODEL_SETUPS Entry (identical across all files)\n\n```python\nif {MODEL_UPPER}_AVAILABLE:\n    MINI_MODEL_SETUPS[\"mini_{model_type}\"] = MiniModelConfig(\n        liger_kernel_patch_func=apply_liger_kernel_to_{model_type},\n        liger_kernel_patch_revert_func=revert_liger_kernel_to_{model_type},\n        model_class={ModelForCausalLM},\n        mini_model_config={ModelConfig}(\n            hidden_size=32, intermediate_size=64, num_hidden_layers=2,\n            num_attention_heads=2, num_key_value_heads=2, vocab_size=1024,\n        ),\n    )\n```\n\n### 4. pytest.param Entry (tolerances differ by dtype)\n\n**bf16 files** — looser tolerances (copy from similar existing model):\n```python\npytest.param(\n    \"mini_{model_type}\", 32, 1e-5, torch.bfloat16,\n    1e-2, 5e-2, 1e-1, 1e-2, 1e-2, 1e-2,\n    marks=pytest.mark.skipif(not supports_bfloat16(), reason=\"bfloat16 not supported\"),\n),\n```\n\n**fp32 files** — tighter tolerances:\n```python\n(\"mini_{model_type}\", 32, 1e-4, torch.float32, 1e-8, 2e-5, 5e-3, 1e-5, 5e-3, 1e-5),\n```\n\n### 5. Revert Function in test/utils.py (shared by all files)\n\n```python\ndef revert_liger_kernel_to_{model_type}(model_config):\n    from transformers.models.{model_type} import modeling_{model_type}\n    importlib.reload(modeling_{model_type})\n    model_config.model_class = modeling_{model_type}.{ModelForCausalLM}\n```\n\n## Key Differences Between Test Variants\n\n**`with_logits`**: Sets `fused_linear_cross_entropy=False, cross_entropy=False` and checks `output.logits` match. Tests that non-loss kernels are numerically correct on their own. Same `MINI_MODEL_SETUPS` entry, just add a `pytest.param`.\n\n**`multimodal`**: Uses `multimodal_collate_fn`, image datasets, and model configs with vision sub-config. Only for models with `has_vision: true`. Multimodal tests may also require fake tokenizer/processor configs in `test/resources/fake_configs/`. Check existing VL model entries (e.g., Qwen2-VL, Llava) for the pattern — if the model needs a tokenizer/processor that can't be created programmatically, add a fake config JSON file.\n\n## Mini Config Guidelines\n\nUse minimal values: hidden_size=32, intermediate_size=64, num_hidden_layers=2, num_attention_heads=2, vocab_size=1024. Add model-specific required fields (e.g., `num_local_experts` for MoE). Copy config structure from the most similar existing model.\n",".agents/skills/liger-autopatch/templates/test-instance-patch.md":"# Template: Instance Patching Test\n\nReference: `test_apply_liger_kernel_to_instance_for_llama` in `test/transformers/test_monkey_patch.py`.\n\n## Add Availability Checker (at top of file, near other `is_*_available` functions ~lines 59-240, in alphabetical order)\n\n```python\ndef is_{model_type}_available():\n    try:\n        import transformers.models.{model_type}\n        return True\n    except ImportError:\n        return False\n```\n\n## Add Test Function (after existing `test_apply_liger_kernel_to_instance_for_*` functions, in alphabetical order)\n\n```python\n@pytest.mark.skipif(not is_{model_type}_available(), reason=\"{model_type} not available\")\ndef test_apply_liger_kernel_to_instance_for_{model_type}():\n    from liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n\n    config = {ModelConfig}(hidden_size=32, intermediate_size=64, num_hidden_layers=2)\n    dummy_model_instance = AutoModelForCausalLM.from_config(config)\n\n    # Verify NOT patched\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) != inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)\n\n    # Apply patch\n    _apply_liger_kernel_to_instance(model=dummy_model_instance)\n\n    # Verify patched\n    assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource({model_type}_lce_forward)\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) == inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)\n```\n\n## Adapt For\n\n- **GeGLU models**: Use `LigerGEGLUMLP` in assertions\n- **MoE models**: Check `decoder_layer.mlp.experts.forward` against `LigerExperts`\n- **Extra norms**: Add assertions for `pre_feedforward_layernorm` etc.\n- **Vision**: Add assertions for `LigerLayerNorm` on vision layers\n",".agents/skills/liger-autopatch/validator.md":"# Validator Agent\n\nRuns tests on generated code and fixes issues iteratively.\n\n## Pre-Requisites\n\nVerify the project is installed in development mode: `python -c \"import liger_kernel\"` must succeed. If not, ask the user to run `pip install -e \".[dev]\"`.\n\n## Validation Steps\n\n### Step 1: Syntax Check\n\n```bash\npython -c \"from liger_kernel.transformers import apply_liger_kernel_to_{model_type}\"\n```\n\n### Step 2: Instance Patching Test\n\n```bash\npytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `ImportError: cannot import name` | Check exact class name in HF source |\n| `AttributeError: has no attribute` | Read decoder layer `__init__` for correct attr names |\n| `AssertionError: getsource mismatch` | Check class being replaced in monkey_patch |\n| `ModuleNotFoundError` | Add `@pytest.mark.skipif` with availability check |\n\n### Step 3: Convergence Tests (requires GPU)\n\nRun **all applicable** convergence tests. For text models, there are 4 files. For VL models, 6.\n\n**FLCE path (fused linear cross entropy):**\n```bash\npytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Non-FLCE path (logits verification — tests RMSNorm/SwiGLU/RoPE without loss fusion):**\n```bash\npytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Multimodal (only for VL models):**\n```bash\npytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `CUDA out of memory` | Reduce hidden_size/num_layers in mini config |\n| `RuntimeError: shape mismatch` | Check `self.config.hidden_size` usage |\n| `AssertionError: allclose` | Check casting_mode, offset, in_place settings |\n| `KeyError in config` | Add missing field to mini model config |\n\n### Step 4: Checkstyle\n\nRun the project's full lint and format check:\n\n```bash\nmake checkstyle\n```\n\nIf it fails, auto-fix with `ruff check . --fix && ruff format .`, then re-run `make checkstyle`.\n\n## Retry Logic\n\n**Hard failures** (Steps 1, 2, 4 — import, patching, checkstyle): These indicate code bugs. Read traceback, identify root cause, fix, re-run. Max 3 retries per step. After 3 failures, report the exact error and stop.\n\n**Soft failures** (Step 3 — convergence): If a convergence test fails after 3 retries with `allclose` tolerance errors, mark it as `SKIP (tolerance tuning needed)` in the report rather than `FAIL`. This means the patching is correct but tolerances need manual adjustment. Note this in the PR description so reviewers are aware.\n\nConvergence failures from other errors (shape mismatch, CUDA OOM, KeyError) are still hard failures — fix and retry.\n\n## Success Report\n\n```\n| Check                        | Status |\n|------------------------------|--------|\n| Import                       | PASS   |\n| Instance Patching            | PASS   |\n| Convergence bf16 FLCE        | PASS   |\n| Convergence bf16 with_logits | PASS   |\n| Convergence fp32 FLCE        | PASS   |\n| Convergence fp32 with_logits | PASS   |\n| Convergence multimodal       | SKIP (text-only model) |\n| Checkstyle                   | PASS   |\n```\n\nInclude any `SKIP (tolerance tuning needed)` entries in the PR description under a \"Known issues\" section.\n",".agents/skills/liger-kernel-dev/SKILL.md":"---\nname: liger-kernel-dev\ndescription: \"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.\"\n---\n\n# Liger Kernel Dev\n\nDevelops 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.\n\n## Mode Detection\n\n- **Create mode**: User asks to create/add/generate/write/build a new kernel → full pipeline\n- **Modify mode**: User asks to update/fix/change/extend an existing kernel → skip Analyze, modify files, then Validate\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Analyzer** workflow in [analyzer.md](analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nAccepts any input: local file, URL, code snippet, natural language description, or model component reference. Produces a standalone PyTorch reference implementation and a kernel profile.\n\n**Human checkpoint:** Present PyTorch reference + kernel profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Generator** workflow in [generator.md](generator.md).\n\nGenerates/modifies up to 8 files:\n\n1. `src/liger_kernel/ops/{kernel}.py` — NEW Triton kernels + autograd Function\n2. `src/liger_kernel/transformers/{kernel}.py` — NEW nn.Module wrapper\n3. `src/liger_kernel/transformers/functional.py` — MODIFY add functional API\n4. `src/liger_kernel/ops/__init__.py` — MODIFY export Function class\n5. `src/liger_kernel/transformers/__init__.py` — MODIFY export Module + `__all__`\n6. `test/transformers/test_{kernel}.py` — NEW unit tests\n7. `benchmark/scripts/benchmark_{kernel}.py` — NEW benchmark script\n8. `benchmark/data/all_benchmark_data.csv` — MODIFY (after benchmarks run)\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns checkstyle, unit tests (hard gate — stops on persistent failure), benchmarks, and generates plots. Optionally runs ncu profiling.\n\n**Human checkpoint:** Report final results with benchmark numbers and plots.\n\n## Pipeline (Modify Mode)\n\n1. Read existing kernel files to understand current implementation\n2. Understand the requested modification\n3. Make targeted changes (Generator handles this)\n4. Run full Validate stage (same as create mode)\n\n## Reference Files\n\n- [kernel-profile-format.md](kernel-profile-format.md) — Kernel profile schema and field descriptions\n- [examples/swiglu-profile.md](examples/swiglu-profile.md) — Tier 1 (element-wise) reference\n- [examples/rms-norm-profile.md](examples/rms-norm-profile.md) — Tier 2 (reduction) reference\n- [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md) — Tier 3 (fused/complex) reference\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n",".agents/skills/liger-kernel-dev/analyzer.md":"# Analyzer Agent\n\nUnderstands a PyTorch operation from any input form and produces a standalone PyTorch reference implementation + kernel profile.\n\n## Input Handling\n\nThe user may provide the operation in any form:\n\n1. **Local file path** → Read the file directly\n2. **URL** (GitHub, HuggingFace, etc.) → Fetch via WebFetch tool\n3. **Code snippet** → Pasted in the conversation\n4. **Natural language** → Mathematical description (e.g., \"element-wise SiLU(x) * y\")\n5. **Model component** → e.g., \"the MLP in Phi-4\" — locate in transformers source and extract\n\n## Steps\n\n### 1. Understand the Operation\n\n- Read/fetch the source code from whatever input was provided\n- Identify the mathematical operation (forward pass)\n- Derive the backward pass (gradient computation)\n- Identify all inputs, outputs, and their expected shapes/dtypes\n- Note any precision-sensitive operations that need float32 upcasting (sigmoid, rsqrt, exp, log, tanh)\n\n### 2. Write PyTorch Reference\n\nCreate a standalone implementation that:\n- Depends only on `torch` (no external libraries)\n- Implements both forward and backward behavior (either as an `nn.Module` or a plain function)\n- Will serve as the correctness baseline for testing\n- Is clean, readable, and well-named\n\n### 3. Classify Into Tier\n\nRead [kernel-profile-format.md](kernel-profile-format.md) for the full schema.\n\n**Tier 1 — Element-wise**: No reductions across dimensions. One row per program. Examples: SwiGLU, GeGLU, DyT.\n- Read reference: `src/liger_kernel/ops/swiglu.py`\n\n**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.\n- Read reference: `src/liger_kernel/ops/rms_norm.py`\n\n**Tier 3 — Fused/Complex**: Multi-pass algorithms, gradient-in-forward tricks, multiple outputs. Examples: CrossEntropy, FusedLinearCrossEntropy.\n- Read reference: `src/liger_kernel/ops/cross_entropy.py`\n\nAlso read the closest example profile:\n- Tier 1 → [examples/swiglu-profile.md](examples/swiglu-profile.md)\n- Tier 2 → [examples/rms-norm-profile.md](examples/rms-norm-profile.md)\n- Tier 3 → [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md)\n\n### 4. Produce Kernel Profile\n\nFill in all fields from [kernel-profile-format.md](kernel-profile-format.md).\n\n### 5. Present to User\n\nShow:\n1. The PyTorch reference implementation (full code)\n2. The kernel profile (all fields)\n3. Which existing kernel is closest (for the Generator to use as reference)\n\nWait for user confirmation before proceeding to Stage 2.\n",".agents/skills/liger-kernel-dev/examples/cross-entropy-profile.md":"# Kernel Profile: CrossEntropy (Tier 3 — Fused/Complex)\n\n## Identity\n- operation_name: cross_entropy\n- function_class_name: LigerCrossEntropyFunction\n- module_class_name: LigerCrossEntropyLoss\n- functional_name: liger_cross_entropy\n\n## Classification\n- tier: 3\n- tier_description: fused/complex\n- closest_existing_kernel: fused_linear_cross_entropy (extends this with linear layer fusion)\n\n## Forward Pass\n- forward_inputs:\n  - _input: shape (B*T, V), logits\n  - target: shape (B*T,), label indices\n  - weight: shape (V,) or None, class weights\n  - ignore_index: int, label to ignore\n  - label_smoothing: float\n  - reduction: str, \"mean\" | \"sum\" | \"none\"\n  - softcap: float or None\n- forward_outputs:\n  - loss: scalar or shape (B*T,) if reduction=\"none\"\n  - z_loss: optional auxiliary loss\n  - token_accuracy: optional accuracy metric\n  - predicted_tokens: optional argmax tokens\n- forward_computation: Two-pass online softmax + cross entropy loss with optional smoothing and softcapping\n- precision_sensitive_ops: [exp, log]\n\n## Backward Pass\n- backward_saved_tensors: [_input] — gradient is computed during forward and stored in-place\n- backward_recompute: none (gradient-in-forward trick)\n- gradient_formulas:\n  - d_input: already computed in forward pass and stored in _input tensor. Backward just scales by grad_output.\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row (one row = one token's logits over vocab)\n- block_size_source: custom — iterates over vocab in chunks of BLOCK_SIZE\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - weight: optional class weights\n  - ignore_index: int = -100\n  - lse_square_scale: float = 0.0\n  - label_smoothing: float = 0.0\n  - reduction: str = \"mean\"\n  - softcap: float or None = None\n  - return_z_loss: bool = False\n- learnable_params: none\n\n## Benchmarking\n- benchmark_variable: vocab_size (V)\n- benchmark_x_label: \"V\"\n- benchmark_x_values_suggestion: [4096, 8192, 16384, 32768, 65536, 131072]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {B: 8, T: 512, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Online softmax**: Two-pass algorithm. Pass 1: compute running max and logsumexp. Pass 2: compute softmax and gradients. Avoids materializing the full softmax vector.\n- **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.\n- **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.\n- **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.\n- **Multiple loss components**: Combines original CE loss, label smoothing loss, and z-loss (for training stability). Each component contributes to both loss and gradient.\n",".agents/skills/liger-kernel-dev/examples/rms-norm-profile.md":"# Kernel Profile: RMSNorm (Tier 2 — Reduction)\n\n## Identity\n- operation_name: rms_norm\n- function_class_name: LigerRMSNormFunction\n- module_class_name: LigerRMSNorm\n- functional_name: liger_rms_norm\n\n## Classification\n- tier: 2\n- tier_description: reduction\n- closest_existing_kernel: layer_norm (similar pattern with different normalization)\n\n## Forward Pass\n- forward_inputs:\n  - X: shape (B, T, H), input tensor\n  - W: shape (H,), weight tensor\n  - eps: float, epsilon for numerical stability\n  - offset: float, weight offset (0.0 for Llama, 1.0 for Gemma)\n  - casting_mode: str, \"llama\" | \"gemma\" | \"none\"\n- forward_outputs:\n  - Y: shape (B, T, H), normalized output\n- forward_computation: Y = (X / RMS(X)) * (W + offset), RMS = sqrt(mean(X^2) + eps)\n- precision_sensitive_ops: [rsqrt]\n\n## Backward Pass\n- backward_saved_tensors: [X, W, RSTD] — RSTD cached from forward to avoid recomputation\n- backward_recompute: none (RSTD is expensive to recompute)\n- gradient_formulas:\n  - dX: rstd * (dY*(W+offset) - (1/N) * rstd^2 * dot(dY*(W+offset), X) * X)\n  - dW: sum over (B,T) of dY * (X * rstd)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: forward uses one program per row `(n_rows,)`, backward uses SM-based partitioning `(sm_count,)` with `rows_per_program`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: true (for dW reduction — each SM accumulates partial dW, then summed)\n\n## Module Parameters\n- module_init_params:\n  - hidden_size: int\n  - eps: float = 1e-6\n  - offset: float = 0.0\n  - casting_mode: str = \"llama\"\n  - init_fn: str = \"ones\"\n  - in_place: bool = True\n  - elementwise_affine: bool = True\n- learnable_params:\n  - weight: shape (hidden_size,), init ones or zeros\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {M: 4096, eps: 1e-6, dtype: torch.float32}\n\n## Key Patterns\n\n- **RSTD caching**: Forward computes and stores `rstd = rsqrt(mean(X^2) + eps)` — 1 value per row, tiny memory cost, saves 4 ops in backward\n- **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)`\n- **Casting modes as constexpr**: `casting_mode` is `tl.constexpr` so the compiler eliminates dead branches\n- **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)\n- **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)\n",".agents/skills/liger-kernel-dev/examples/swiglu-profile.md":"# Kernel Profile: SwiGLU (Tier 1 — Element-wise)\n\n## Identity\n- operation_name: swiglu\n- function_class_name: LigerSiLUMulFunction\n- module_class_name: LigerSwiGLUMLP\n- functional_name: liger_swiglu\n\n## Classification\n- tier: 1\n- tier_description: element-wise\n- closest_existing_kernel: geglu (same structure, different activation)\n\n## Forward Pass\n- forward_inputs:\n  - a: shape (B, T, H), gate projection output\n  - b: shape (B, T, H), up projection output\n- forward_outputs:\n  - c: shape (B, T, H), silu(a) * b\n- forward_computation: c = silu(a) * b, where silu(x) = x * sigmoid(x)\n- precision_sensitive_ops: [sigmoid]\n\n## Backward Pass\n- backward_saved_tensors: [a, b] (reshaped to 2D in forward wrapper)\n- backward_recompute: recompute silu(a) and sigmoid(a) in backward\n- gradient_formulas:\n  - da: dc * (silu(a) * (1 - sigmoid(a)) + sigmoid(a)) * b\n  - db: dc * silu(a)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row, `(n_rows,)`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - config: HuggingFace model config object\n- learnable_params:\n  - gate_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - up_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - down_proj: Linear(intermediate_size, hidden_size, bias=False)\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"torch\", \"torch_compile\"]\n- benchmark_fixed_config: {BT: 4096, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Recomputation over saving**: Forward saves `a, b` but backward recomputes `sigmoid(a)` and `silu(a)` — saves memory, sigmoid is cheap\n- **In-place backward**: Writes gradients directly to `a_ptr` and `b_ptr` (the saved tensors) — saves allocation\n- **Float32 for sigmoid**: `a_row` cast to `tl.float32` before `tl.sigmoid`, result cast back via `.cast(b_row.dtype)`\n- **No intermediate allocations**: Forward kernel writes directly to output `c`; backward kernel overwrites saved `a, b`\n",".agents/skills/liger-kernel-dev/generator.md":"# Generator Agent\n\nTakes a confirmed kernel profile and PyTorch reference, generates all files for a production-ready Liger kernel.\n\n## Pre-Requisites\n\nBefore generating, read the existing implementation closest to this kernel's tier:\n- Tier 1 → `src/liger_kernel/ops/swiglu.py` + `src/liger_kernel/transformers/swiglu.py`\n- Tier 2 → `src/liger_kernel/ops/rms_norm.py` + `src/liger_kernel/transformers/rms_norm.py`\n- Tier 3 → `src/liger_kernel/ops/cross_entropy.py` + `src/liger_kernel/transformers/cross_entropy.py`\n\nAlso read the corresponding test and benchmark files for the reference kernel.\n\n## Files to Generate\n\n### 1. `src/liger_kernel/ops/{kernel}.py` (NEW)\n\nThe core Triton kernel. See [templates/ops-kernel.md](templates/ops-kernel.md).\n\nKey rules:\n- `@triton.jit` for forward and backward kernels\n- `torch.autograd.Function` subclass with `@ensure_contiguous` decorator\n- Use `calculate_settings(n_cols)` for BLOCK_SIZE and num_warps\n- Cast to `tl.float32` for precision-sensitive ops (sigmoid, rsqrt, exp, log, tanh)\n- Use `tl.program_id(0).to(tl.int64)` to avoid overflow\n- Use stride parameters for memory access, not hardcoded offsets\n- Prefer recomputation over saving for backward when computation is cheap\n\n### 2. `src/liger_kernel/transformers/{kernel}.py` (NEW)\n\nThe nn.Module wrapper. See [templates/module-wrapper.md](templates/module-wrapper.md).\n\nKey rules:\n- Import from `liger_kernel.ops` (not from the submodule directly)\n- Include `extra_repr` method\n- Match the PyTorch API as closely as possible\n\n### 3. `src/liger_kernel/transformers/functional.py` (MODIFY)\n\nAdd functional API. See [templates/functional-api.md](templates/functional-api.md).\n\nTwo changes:\n- Add import of the Function class at the top (alphabetical order)\n- Add `liger_{kernel}()` function (alphabetical order among existing functions)\n\n### 4. `src/liger_kernel/ops/__init__.py` (MODIFY)\n\nAdd import of the Function class. Insert in alphabetical order among existing imports. Use `# noqa: F401` suffix.\n\n### 5. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the Module class in three locations (alphabetical order in each):\n- Direct import at top (among \"Always-safe imports\")\n- `__all__` list\n- If the Module has a `TYPE_CHECKING` import or `__getattr__` entry needed, add there too\n\n### 6. `test/transformers/test_{kernel}.py` (NEW)\n\nUnit tests. See [templates/unit-test.md](templates/unit-test.md).\n\nKey rules:\n- Include a PyTorch reference implementation (copy from Analyzer output)\n- Test both forward and backward correctness\n- Parametrize over shapes (include non-power-of-2 and small/large sizes)\n- Parametrize over dtypes with appropriate tolerances\n- Use `set_seed()`, `assert_verbose_allclose()`, `infer_device()`, `supports_bfloat16()`\n- Import from `test.utils`, NOT `liger_kernel.test.utils`\n\n### 7. `benchmark/scripts/benchmark_{kernel}.py` (NEW)\n\nBenchmarks. See [templates/benchmark.md](templates/benchmark.md).\n\nKey rules:\n- Define `_setup_{kernel}()` function returning `(input_tensor, layer)`\n- Define `bench_speed_{kernel}()` using `run_speed_benchmark`\n- Define `bench_memory_{kernel}()` using `run_memory_benchmark`\n- Compare \"liger\" vs \"torch\" (and optionally \"torch_compile\") providers\n- Use `get_benchmark_model_config`, `estimate_kernel_peak_memory`, and sweep config utilities\n- Run both speed and memory benchmarks for forward, backward, and full modes\n- **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\n\n## Modify Mode\n\nWhen modifying an existing kernel:\n1. Read the existing kernel files thoroughly\n2. Understand the current implementation patterns\n3. Make targeted changes — do not rewrite entire files\n4. Add new test cases for new functionality\n5. Ensure existing tests still cover original behavior\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order**\n- Import style: one import per line, no multi-imports\n",".agents/skills/liger-kernel-dev/kernel-profile-format.md":"# Kernel Profile Format\n\nWhen analyzing an operation, produce a profile with all of these fields.\n\n## Fields\n\n```yaml\n# Identity\noperation_name: snake_case name (e.g., \"rms_norm\", \"dyt\", \"sparsemax\")\nfunction_class_name: \"Liger{PascalCase}Function\" (e.g., \"LigerRMSNormFunction\")\nmodule_class_name: \"Liger{PascalCase}\" (e.g., \"LigerRMSNorm\")\nfunctional_name: \"liger_{snake_case}\" (e.g., \"liger_rms_norm\")\n\n# Classification\ntier: 1 | 2 | 3\ntier_description: \"element-wise\" | \"reduction\" | \"fused/complex\"\nclosest_existing_kernel: which existing Liger kernel is most similar\n\n# Forward Pass\nforward_inputs:\n  - name: X\n    shape: \"(B, T, H)\"\n    dtype: \"input dtype\"\n  # ... list all inputs with shapes\nforward_outputs:\n  - name: Y\n    shape: \"(B, T, H)\"\n    dtype: \"same as input\"\nforward_computation: brief math description (e.g., \"Y = X * rsqrt(mean(X^2) + eps) * W\")\nprecision_sensitive_ops: list ops needing float32 (e.g., [rsqrt, sigmoid, exp])\n\n# Backward Pass\nbackward_saved_tensors: what to save_for_backward (e.g., [X, W, RSTD])\nbackward_recompute: what to recompute instead of saving (e.g., \"recompute silu in backward\")\ngradient_formulas:\n  dX: \"formula or description\"\n  dW: \"formula or description\"\n\n# Tiling Strategy\ngrid_dimensions: 1D | 2D\ngrid_description: \"one program per row\" | \"one program per (row_block, col_block)\" | etc.\nblock_size_source: \"calculate_settings(n_cols)\" | \"custom\" with explanation\nneeds_sm_parallelism: true | false (for weight gradient reduction in backward)\n\n# Module Parameters\nmodule_init_params:\n  - name: hidden_size\n    type: int\n  - name: eps\n    type: float\n    default: 1e-6\n  # ... list all __init__ parameters\nlearnable_params:\n  - name: weight\n    shape: \"(hidden_size,)\"\n    init: \"ones\"\n  # ... list learnable parameters\n\n# Benchmarking\nbenchmark_variable: \"hidden_size\" | \"seq_len\" | \"vocab_size\" | etc.\nbenchmark_x_label: human-readable label for x-axis\nbenchmark_x_values_suggestion: suggested range to sweep\nbenchmark_providers: [\"liger\", \"torch\"] (add \"torch_compile\" if applicable)\nbenchmark_fixed_config: dict of fixed params (e.g., {BT: 4096, dtype: torch.bfloat16})\n```\n\n## Naming Conventions\n\n| Component | Pattern | Example |\n|-----------|---------|---------|\n| Ops file | `src/liger_kernel/ops/{operation_name}.py` | `ops/dyt.py` |\n| Transformer file | `src/liger_kernel/transformers/{operation_name}.py` | `transformers/dyt.py` |\n| Test file | `test/transformers/test_{operation_name}.py` | `test_dyt.py` |\n| Benchmark file | `benchmark/scripts/benchmark_{operation_name}.py` | `benchmark_dyt.py` |\n| Function class | `Liger{PascalCase}Function` | `LigerDyTFunction` |\n| Module class | `Liger{PascalCase}` | `LigerDyT` |\n| Functional | `liger_{snake_case}` | `liger_dyt` |\n| Forward wrapper | `{snake_case}_forward` or `liger_{snake_case}_fwd` | `liger_dyt_fwd` |\n| Backward wrapper | `{snake_case}_backward` or `liger_{snake_case}_bwd` | `liger_dyt_bwd` |\n",".agents/skills/liger-kernel-dev/templates/benchmark.md":"# Benchmark Template\n\n## File: `benchmark/scripts/benchmark_{kernel}.py`\n\n```python\nimport os\nimport sys\n\nimport torch\n\nfrom benchmark_model_configs import compute_hidden_size_sweep_config\nfrom benchmark_model_configs import estimate_kernel_peak_memory\nfrom benchmark_model_configs import get_benchmark_model_config\nfrom utils import SingleBenchmarkRunInput\nfrom utils import SingleBenchmarkRunOutput\nfrom utils import parse_benchmark_script_args\nfrom utils import run_benchmarks\nfrom utils import run_memory_benchmark\nfrom utils import run_speed_benchmark\n\nfrom liger_kernel.utils import infer_device\n\ndevice = infer_device()\n\nsys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), \"../..\")))\n\n\ndef _setup_{kernel}(input: SingleBenchmarkRunInput):\n    \"\"\"Create input tensor and layer from benchmark config.\"\"\"\n    # Import both Liger and PyTorch reference implementations\n    from test.transformers.test_{kernel} import Liger{Kernel}\n    from test.transformers.test_{kernel} import Torch{Kernel}\n\n    cfg = input.extra_benchmark_config\n    hidden_size = input.x  # Or whatever the sweep variable is\n\n    x = torch.randn(cfg[\"BT\"], hidden_size, device=device, dtype=cfg[\"dtype\"], requires_grad=True)\n\n    if input.kernel_provider == \"liger\":\n        layer = Liger{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch\":\n        layer = Torch{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch_compile\":\n        layer = torch.compile(Torch{Kernel}(hidden_size=hidden_size).to(device))\n    else:\n        raise ValueError(f\"Invalid provider: {input.kernel_provider} for {kernel}\")\n\n    return x, layer\n\n\ndef bench_speed_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x])\n\n\ndef bench_memory_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode)\n\n\nBT = 4096  # Fixed batch*seq for sweeping hidden_size\n\n\nif __name__ == \"__main__\":\n    args = parse_benchmark_script_args()\n    model = get_benchmark_model_config(args.model)\n\n    def _probe():\n        probe_input = SingleBenchmarkRunInput(\n            x=model.hidden_size,\n            kernel_provider=\"torch\",\n            extra_benchmark_config={\"BT\": BT, \"dtype\": model.dtype},\n        )\n        x, layer = _setup_{kernel}(probe_input)\n        return layer(x)\n\n    peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)\n    sweep_config = compute_hidden_size_sweep_config(model, peak_bytes, bt=BT)\n    x_values = [1024 * i for i in range(1, 17) if 1024 * i <= sweep_config.max_hidden_size] or [model.hidden_size]\n\n    common_configs = {\n        \"kernel_name\": \"{kernel}\",\n        \"x_name\": \"hidden_size\",\n        \"x_label\": \"hidden_size\",\n        \"x_values\": x_values,\n        \"kernel_providers\": [\"liger\", \"torch\", \"torch_compile\"],\n        \"extra_benchmark_configs\": [{\"BT\": sweep_config.bt, \"dtype\": model.dtype}],\n        \"overwrite\": args.overwrite,\n    }\n\n    run_benchmarks(\n        bench_test_fn=bench_speed_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"speed\",\n        metric_unit=\"ms\",\n        **common_configs,\n    )\n    run_benchmarks(\n        bench_test_fn=bench_memory_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"memory\",\n        metric_unit=\"MB\",\n        **common_configs,\n    )\n```\n\n### Key Rules\n\n1. **Import both implementations from test file** — avoids duplicating reference code\n2. **Support \"liger\", \"torch\", and \"torch_compile\" providers**\n3. **Use `run_speed_benchmark` and `run_memory_benchmark`** from utils — don't write custom measurement\n4. **Use `get_benchmark_model_config`** and sweep utilities for standard model sizes\n5. **Use `estimate_kernel_peak_memory`** to auto-compute safe x_values range\n6. **Run both speed AND memory** for all three modes (forward, backward, full)\n7. **Use `BT = 4096`** as default batch*seq when sweeping hidden_size\n8. **Adapt sweep variable** — some kernels sweep `seq_len` or `vocab_size` instead of `hidden_size`. Use the appropriate `compute_*_sweep_config` utility.\n9. **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.\n",".agents/skills/liger-kernel-dev/templates/functional-api.md":"# Functional API Template\n\n## File: `src/liger_kernel/transformers/functional.py` (MODIFY)\n\n### Step 1: Add Import\n\nAdd at the top of the file, in alphabetical order among existing imports:\n\n```python\nfrom liger_kernel.ops import Liger{Kernel}Function\n```\n\n### Step 2: Add Function\n\nAdd in alphabetical order among existing functions:\n\n```python\ndef liger_{kernel}(x, ...params):\n    return Liger{Kernel}Function.apply(x, ...params)\n```\n\nFor functions with complex return types (like CrossEntropy), wrap the output:\n\n```python\ndef liger_{kernel}(input, target, ...params):\n    result = Liger{Kernel}Function.apply(input, target, ...params)\n    return result\n```\n\n### Key Rules\n\n- Function signature should match the PyTorch equivalent where applicable\n- Use positional args for tensor inputs, keyword args for config params\n- Keep the function body minimal — just call `.apply()` and return\n",".agents/skills/liger-kernel-dev/templates/module-wrapper.md":"# Module Wrapper Template\n\n## File: `src/liger_kernel/transformers/{kernel}.py`\n\n```python\nimport torch\nimport torch.nn as nn\n\nfrom liger_kernel.ops import Liger{Kernel}Function\n\n\nclass Liger{Kernel}(nn.Module):\n    def __init__(self, ...params):\n        super().__init__()\n        # Store config\n        self.param = param\n        # Create learnable parameters\n        self.weight = nn.Parameter(torch.ones(hidden_size))\n\n    def forward(self, x):\n        return Liger{Kernel}Function.apply(x, self.weight, ...self.params)\n\n    def extra_repr(self):\n        return f\"param={self.param}, ...\"\n```\n\n### Key Rules\n\n- Import from `liger_kernel.ops` (vendor-replaceable), not from the submodule\n- Match the PyTorch reference API as closely as possible (same __init__ params, same forward signature)\n- Include `extra_repr` for debugging\n- For loss modules, name the class `Liger{Kernel}Loss` and the forward takes `(input, target)`\n- For model-specific variants, create subclasses (see `LigerRMSNormForGemma` in `rms_norm.py`)\n",".agents/skills/liger-kernel-dev/templates/ops-kernel.md":"# Ops Kernel Template\n\n## File: `src/liger_kernel/ops/{kernel}.py`\n\nBefore writing, read the reference kernel for this tier in the actual codebase.\n\n### Import Block\n\n```python\nimport torch\nimport triton\nimport triton.language as tl\n\nfrom liger_kernel.ops.utils import calculate_settings\nfrom liger_kernel.ops.utils import ensure_contiguous\n```\n\nAdd additional imports only as needed:\n- `math` — if using `math.ceil` for row partitioning\n- `from liger_kernel.ops.utils import compare_version` — if checking triton version\n- `from liger_kernel.ops.utils import torch_to_triton_dtype` — if passing dtype as kernel arg\n- `rsqrt` from triton libdevice — if using inverse square root (see `rms_norm.py` for the import pattern with version guards)\n- `tanh` from triton libdevice — if using tanh (see `dyt.py` for the import pattern)\n\n### Forward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_forward_kernel(\n    # Output pointers first, then input pointers\n    Y_ptr, Y_stride,\n    X_ptr, X_stride,\n    # ... other input pointers and strides\n    n_cols,\n    # constexpr params for compile-time optimization\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Calculate base pointers using strides\n    X_ptr += row_idx * X_stride\n    Y_ptr += row_idx * Y_stride\n\n    # Load with mask and float32 upcast for precision\n    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0).to(tl.float32)\n\n    # === FORWARD COMPUTATION ===\n\n    tl.store(Y_ptr + col_offsets, result, mask=mask)\n```\n\n### Backward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_backward_kernel(\n    # Gradient output, then gradient inputs, then saved tensors\n    dY_ptr, dY_stride,\n    dX_ptr, dX_stride,\n    # ... saved tensors from forward\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Load gradient and saved tensors\n    # === BACKWARD COMPUTATION ===\n\n    tl.store(dX_ptr + row_idx * dX_stride + col_offsets, dX_row, mask=mask)\n```\n\n### Python Wrappers\n\n```python\ndef {kernel}_forward(X, ...):\n    ori_shape = X.shape\n    n_cols = ori_shape[-1]\n    X = X.view(-1, n_cols)\n    n_rows = X.shape[0]\n\n    Y = torch.empty_like(X)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_forward_kernel[(n_rows,)](\n        Y, Y.stride(-2),\n        X, X.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return Y.view(*ori_shape)\n\n\ndef {kernel}_backward(dY, ...saved):\n    ori_shape = dY.shape\n    n_cols = ori_shape[-1]\n    dY = dY.view(-1, n_cols)\n    n_rows = dY.shape[0]\n\n    dX = torch.empty_like(dY)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_backward_kernel[(n_rows,)](\n        dY, dY.stride(-2),\n        dX, dX.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return dX.view(*ori_shape)\n```\n\n### Autograd Function\n\n```python\nclass Liger{Kernel}Function(torch.autograd.Function):\n    @staticmethod\n    @ensure_contiguous\n    def forward(ctx, X, ...params):\n        Y = {kernel}_forward(X, ...params)\n        ctx.save_for_backward(...)  # Save only what backward needs\n        return Y\n\n    @staticmethod\n    @ensure_contiguous\n    def backward(ctx, dY):\n        ...saved = ctx.saved_tensors\n        dX = {kernel}_backward(dY, ...saved)\n        # Return one gradient per forward input, None for non-tensor params\n        return dX, None, None, ...\n```\n\n### Tier-Specific Patterns\n\n**Tier 1 (element-wise):**\n- Grid: `(n_rows,)` — one program per row\n- Simple 1D indexing with `col_offsets` and `mask`\n- Prefer recomputation in backward over saving activations\n- See `swiglu.py`: saves `a, b` but recomputes `silu(a)` in backward\n\n**Tier 2 (reduction):**\n- Forward may cache intermediate state (e.g., RSTD in RMSNorm)\n- Backward may use SM-based parallelism for weight gradient reduction:\n  ```python\n  sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count\n  _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)\n  grid = (sm_count,)\n  rows_per_program = math.ceil(n_rows / sm_count)\n  ```\n- See `rms_norm.py` for the full pattern\n\n**Tier 3 (fused/complex):**\n- Multi-pass algorithms (e.g., online softmax: pass 1 finds max+sum, pass 2 computes gradients)\n- Gradient-in-forward trick: compute gradients during forward, store in input tensor\n- Multiple `tl.constexpr` flags for compile-time dead code elimination\n- See `cross_entropy.py` for the full pattern\n\n### Key Rules\n\n1. Always use `tl.program_id(0).to(tl.int64)` (prevents overflow on large tensors)\n2. Always use `mask = col_offsets < n_cols` for boundary handling\n3. Use `tl.constexpr` for parameters that enable compile-time optimization\n4. Use stride parameters from tensors, never hardcode memory layout\n5. Cast precision-sensitive ops to `tl.float32`, cast back for storage\n6. Return `None` for each non-tensor parameter in backward\n",".agents/skills/liger-kernel-dev/templates/unit-test.md":"# Unit Test Template\n\n## File: `test/transformers/test_{kernel}.py`\n\n```python\nimport pytest\nimport torch\n\nfrom test.utils import assert_verbose_allclose\nfrom test.utils import infer_device\nfrom test.utils import set_seed\nfrom test.utils import supports_bfloat16\n\nfrom liger_kernel.ops import Liger{Kernel}Function\nfrom liger_kernel.transformers.{kernel} import Liger{Kernel}\nfrom liger_kernel.transformers.functional import liger_{kernel}\n\n\n# ---- PyTorch Reference ----\n# Paste the standalone PyTorch reference from the Analyzer stage here.\n\nclass Torch{Kernel}(torch.nn.Module):\n    def __init__(self, ...):\n        super().__init__()\n        ...\n\n    def forward(self, x):\n        ...\n\n\nset_seed(42)\ndevice = infer_device()\n\n\n# ---- Correctness Test (Module) ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (4, 16, 2048),\n        (1, 1, 1023),   # Non-power-of-2\n        (3, 7, 256),    # Small + prime numbers\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-1,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_correctness(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Initialize both with same weights\n    torch_layer = Torch{Kernel}(...).to(device).to(dtype)\n    liger_layer = Liger{Kernel}(...).to(device).to(dtype)\n    # Copy weights from torch to liger\n    # liger_layer.weight.data = torch_layer.weight.data.clone()\n\n    # Forward\n    torch_output = torch_layer(x1)\n    liger_output = liger_layer(x2)\n    assert_verbose_allclose(torch_output, liger_output, rtol=rtol, atol=atol, extra_info=\"[output]\")\n\n    # Backward\n    grad_output = torch.randn_like(_input)\n    torch_output.backward(grad_output)\n    liger_output.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol, extra_info=\"[input.grad]\")\n\n    # Check weight gradients if applicable\n    # assert_verbose_allclose(\n    #     torch_layer.weight.grad, liger_layer.weight.grad,\n    #     rtol=rtol, atol=atol, extra_info=\"[weight.grad]\"\n    # )\n\n\n# ---- Functional API Test ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (1, 1, 1023),\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-8,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_functional(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Test that functional and Function.apply produce identical results\n    output1 = liger_{kernel}(x1, ...)\n    output2 = Liger{Kernel}Function.apply(x2, ...)\n\n    assert_verbose_allclose(output1, output2, rtol=rtol, atol=atol)\n\n    grad_output = torch.randn_like(_input)\n    output1.backward(grad_output)\n    output2.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol)\n```\n\n### Key Testing Rules\n\n1. **Always test both forward AND backward** — backward bugs are common\n2. **Include non-power-of-2 shapes** — tests the mask boundary handling\n3. **Parametrize dtypes with per-dtype tolerances**:\n   - `float32`: `atol=1e-5, rtol=1e-5` (tight)\n   - `bfloat16`: `atol=1e-1, rtol=5e-2` (loose — bf16 has limited precision)\n4. **Use `set_seed(42)` at module level** for reproducibility\n5. **Use `infer_device()`** not `\"cuda\"` — supports multi-vendor\n6. **Guard bfloat16 with `supports_bfloat16()`** — skips on older GPUs\n7. **Use `assert_verbose_allclose`** not `torch.testing.assert_close` — gives better error output\n8. **Test weight gradients** when the kernel has learnable parameters\n9. **Add `@pytest.mark.flaky(reruns=3)`** if the test is sensitive to random seeds\n10. **Add extra test functions** for edge cases specific to the kernel\n",".agents/skills/liger-kernel-dev/validator.md":"# Validator Agent\n\nValidates the generated kernel through checkstyle, testing, benchmarking, and optional profiling.\n\n## Validation Flow\n\n### Step 1: Checkstyle (Fast Gate)\n\n```bash\nmake checkstyle\n```\n\nIf it fails:\n```bash\nruff check . --fix && ruff format .\nmake checkstyle\n```\n\nIf it still fails after auto-fix, identify and fix the remaining issues manually, then re-run.\n\n### Step 2: Unit Tests (Hard Gate)\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\nIf tests fail:\n1. Read the error output carefully\n2. Identify root cause: kernel logic bug, shape mismatch, dtype issue, or tolerance too tight\n3. Fix the issue in the relevant file(s)\n4. Re-run `make checkstyle` (in case fix introduced style issues)\n5. Re-run the failing test\n\n**Retry up to 3 times total.** If tests still fail after 3 attempts:\n- **STOP IMMEDIATELY. Do NOT proceed to benchmarks.**\n- Report to the user:\n  - Which tests are failing and the exact error messages\n  - What fixes were attempted\n  - The current state of the code\n  - Suggest what might be wrong and ask for guidance\n\n### Step 3: Benchmarks\n\nRun the benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nThis produces speed and memory measurements for forward, backward, and full modes, comparing Liger vs PyTorch baseline.\n\n### Step 4: Generate Plots\n\n```bash\npython benchmark/benchmarks_visualizer.py\n```\n\nPlots are saved to `benchmark/visualizations/`. If the visualizer fails or is not available, this step is non-blocking — report the raw numbers instead.\n\n### Step 5: Optional ncu Profiling\n\nOnly if the user explicitly requests it. Requires NVIDIA Nsight Compute (`ncu`) to be installed.\n\nProfile both the Liger kernel and PyTorch reference:\n```bash\nncu --set full --target-processes all -o liger_profile python -c \"\nimport torch\nfrom liger_kernel.ops.{kernel} import Liger{Kernel}Function\nx = torch.randn(..., device='cuda', dtype=torch.float32, requires_grad=True)\nout = Liger{Kernel}Function.apply(x, ...)\nout.backward(torch.randn_like(out))\n\"\n```\n\nReport key metrics: SM occupancy, memory throughput, compute throughput.\n\nIf ncu is not available, report this and skip.\n\n### Step 6: Report Results\n\nPresent a clear summary to the user:\n\n```\n## Results\n\n**Checkstyle:** PASS\n**Unit Tests:** PASS (X tests passed)\n\n**Speed Benchmarks:**\n| Mode     | Liger (ms) | PyTorch (ms) | Speedup |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...x    |\n| Backward | ...       | ...         | ...x    |\n| Full     | ...       | ...         | ...x    |\n\n**Memory Benchmarks:**\n| Mode     | Liger (MB) | PyTorch (MB) | Savings |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...%    |\n| Backward | ...       | ...         | ...%    |\n| Full     | ...       | ...         | ...%    |\n\n**Assessment:** [Which dimension improved — speed, memory, or both.\nFlag if either metric is catastrophically worse.]\n\n**Plots:** benchmark/visualizations/\n\n**ncu Profiling:** [Results if requested, or \"Not requested\"]\n```\n",".agents/skills/liger-kernel-perf/SKILL.md":"---\nname: liger-kernel-perf\ndescription: \"Optimizes the performance of existing Liger Kernel Triton kernels. Profiles kernels, diagnoses bottlenecks (memory-bound vs compute-bound), generates multiple optimization variants with benchmarking, and applies the best variant while maintaining correctness. Supports GPU architecture-specific optimization (Ampere, Hopper, Blackwell). Use when a user asks to optimize, speed up, tune, profile, or reduce memory of an existing Liger kernel.\"\n---\n\n# Liger Kernel Perf\n\nOptimizes existing Liger Kernel Triton kernels through a 3-stage pipeline: Profile, Optimize, Finalize. Supports interactive mode (human checkpoints between stages) and autonomous mode (runs end-to-end). NVIDIA GPUs only.\n\n## Mode Detection\n\n- **Interactive mode** (default): Human checkpoints between each stage\n- **Autonomous mode**: User says \"just optimize it\", \"run without asking me\", \"optimize autonomously\" → all stages run end-to-end, user sees only the final report\n\n## Input Parsing\n\nExtract from the user's request:\n\n| Field | Description | Default |\n|-------|-------------|---------|\n| `target_kernel` | Which kernel to optimize (e.g., \"rms_norm\", \"cross_entropy\") | **Required** |\n| `optimization_goal` | speed / memory / balanced | balanced |\n| `scope` | Specific pass (forward/backward), input regime, or general | general |\n| `target_gpu` | Ampere / Hopper / Blackwell / auto-detect | auto-detect |\n| `autonomy` | interactive / autonomous | interactive |\n| `max_variants` | Max optimization variants to try | 8 |\n| `target_metric` | Optional concrete target (e.g., \"forward under 0.3ms at hidden_size=4096\") | none |\n\n## Pre-Flight Validation\n\nBefore starting the pipeline, validate:\n\n1. Kernel file exists: `src/liger_kernel/ops/{kernel}.py`\n2. Benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`\n3. Test file exists: `test/transformers/test_{kernel}.py`\n4. GPU is available and CUDA works\n5. Project is installed in dev mode (`pip install -e \".[dev]\"`)\n\nIf any validation fails, report clearly and stop.\n\n## Pipeline\n\n### Stage 1: Profile\n\nFollow the **Profiler** workflow in [profiler.md](profiler.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nThis stage:\n1. Creates the workspace directory `optimization/{kernel}/`\n2. Copies the original kernel as a snapshot\n3. Runs baseline benchmarks using the existing benchmark script\n4. Detects GPU architecture (or uses user-specified target)\n5. Optionally runs NCU profiling (if `ncu` is available)\n6. Analyzes the kernel code (tier classification, patterns, optimization opportunities)\n7. Classifies the bottleneck: memory-bound vs compute-bound\n8. Produces an optimization profile with a recommended strategy order\n9. Saves profile to `optimization/{kernel}/profile.md`\n\n**Human checkpoint (interactive mode):** Present the optimization profile with bottleneck diagnosis and proposed strategy order. Confirm before proceeding.\n\n### Stage 2: Optimize\n\nFollow the **Optimizer** workflow in [optimizer.md](optimizer.md).\n\nThis stage runs an autonomous optimization loop:\n\n1. Read the optimization profile and original kernel\n2. **Always try parameter tuning first** (BLOCK_SIZE, num_warps, num_stages manual sweep -- NOT @triton.autotune)\n3. Then apply diagnosis-driven techniques from [optimization-strategies.md](optimization-strategies.md)\n4. For each variant:\n   a. Generate the variant code → `optimization/{kernel}/{kernel}_vN.py`\n   b. Write the variant lab notebook → `optimization/{kernel}/{kernel}_vN_notes.md`\n   c. Run quick smoke test (single shape, float32, forward+backward) → discard on failure\n   d. Run the **full existing benchmark script** → `optimization/{kernel}/benchmarks/vN_results.csv`\n   e. Check guardrails (no catastrophic regressions)\n   f. Update the variant notes with actual results\n5. Read all prior variant notes before generating the next variant\n6. **Stop when:** budget exhausted, 2 consecutive variants with <1% improvement, or target metric met\n7. Produce a comparison table of ALL variants\n\n**Human checkpoint (interactive mode):** Present the comparison table across all variants. User approves the winner (or skill picks best if autonomous).\n\n### Stage 3: Finalize\n\nFollow the **Finalizer** workflow in [finalizer.md](finalizer.md).\n\nThis stage:\n1. Applies the winning variant in-place to `src/liger_kernel/ops/{kernel}.py`\n2. Runs the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs` (hard gate)\n3. Runs checkstyle: `make checkstyle` (auto-fix with `ruff check . --fix && ruff format .`)\n4. Generates 3-way comparison plots (original liger vs optimized liger vs huggingface baseline) using `benchmarks_visualizer.py`\n5. Generates the final optimization report → `optimization/{kernel}/report.md`\n6. Creates a PR with only the kernel code changes (no plots or optimization workspace files)\n7. Presents the before/after summary with plots\n\n**Human checkpoint (interactive mode):** Present the final report with before/after numbers, comparison plots, and test results.\n\n## Guardrails\n\nThese apply to EVERY variant, regardless of mode:\n\n| Guardrail | Threshold | Action |\n|-----------|-----------|--------|\n| Non-target metric regression | >5% worse | Reject variant |\n| Cross-pass regression | >10% on one pass to marginally improve other | Reject variant |\n| Smoke test failure | Any correctness failure | Discard variant immediately |\n| Full test suite failure | Any | Do NOT apply winner, report failure, stop |\n| Checkstyle failure | Any | Auto-fix with ruff, retry once |\n\n## Reference Files\n\n- [profiler.md](profiler.md) -- Profiler Agent specification\n- [optimizer.md](optimizer.md) -- Optimizer Agent specification\n- [finalizer.md](finalizer.md) -- Finalizer Agent specification\n- [optimization-strategies.md](optimization-strategies.md) -- Catalog of optimization techniques\n- [templates/optimization-profile.md](templates/optimization-profile.md) -- Profiling output format (cross-stage contract)\n- [templates/variant-notes.md](templates/variant-notes.md) -- Per-variant lab notebook format\n",".agents/skills/liger-kernel-perf/finalizer.md":"# Finalizer Agent\n\nApplies the winning optimization variant to the production kernel, validates correctness through the full test suite and checkstyle, and produces the final optimization report.\n\nThis is Stage 3 of the 3-stage pipeline. It receives the winning variant from Stage 2 (Optimizer) and is responsible for ensuring the optimization is safe to ship.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `winning_variant`: variant number chosen by user (interactive) or by the Optimizer (autonomous)\n- `optimization_workspace`: path to `optimization/{kernel}/` containing all variant files, benchmarks, and the profile\n\n## Expected Workspace State\n\nBefore the Finalizer starts, the following files MUST exist (created by earlier stages):\n\n```\noptimization/{kernel}/\n  original_{kernel}.py          # Snapshot of the original kernel (from Profiler)\n  profile.md                    # Optimization profile (from Profiler)\n  {kernel}_v1.py                # Variant source files (from Optimizer)\n  {kernel}_v1_notes.md          # Variant lab notebooks (from Optimizer)\n  {kernel}_v2.py                # ... more variants\n  {kernel}_v2_notes.md\n  benchmarks/\n    v0_baseline.csv             # Baseline benchmark data\n    v1_results.csv              # Per-variant benchmark data\n    v2_results.csv\n```\n\nIf any required files are missing, report the missing files and stop.\n\n## Steps\n\n### Step 1: Validate Inputs and Read Context\n\n1. Confirm the winning variant file exists: `optimization/{kernel}/{kernel}_v{N}.py`\n2. Read the winning variant code\n3. Read the original kernel: `optimization/{kernel}/original_{kernel}.py`\n4. Read the optimization profile: `optimization/{kernel}/profile.md`\n5. Read ALL variant notes files (`{kernel}_v*_notes.md`) to collect data for the report\n6. Read ALL benchmark CSVs (`benchmarks/v*_results.csv`) to build the comparison table\n7. Read the current production kernel: `src/liger_kernel/ops/{kernel}.py`\n\nVerify the original kernel snapshot matches the current production kernel. If they differ (someone else modified the file during the optimization), warn the user and confirm before proceeding.\n\n### Step 2: Apply the Winning Variant\n\nReplace the production kernel with the winning variant:\n\n```bash\ncp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis is a destructive operation. The original is preserved at `optimization/{kernel}/original_{kernel}.py`.\n\n### Step 3: Run the Full Test Suite (HARD GATE)\n\nRun the complete test file for the kernel:\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\n**This is a HARD GATE.** The test suite must pass completely.\n\n#### If tests PASS:\nProceed to Step 4.\n\n#### If tests FAIL:\nExecute the test failure recovery procedure:\n\n**Attempt 1: Diagnose and fix**\n1. Read the full test error output carefully\n2. Identify the root cause (numerical tolerance issue, shape mismatch, missing import, API change, etc.)\n3. Apply a targeted fix to `src/liger_kernel/ops/{kernel}.py`\n4. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**Attempt 2: Second fix attempt (if Attempt 1 still fails)**\n1. Read the new error output\n2. Try a different fix approach\n3. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**If both attempts fail: REVERT**\n1. Restore the original kernel:\n   ```bash\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n2. Verify the revert by running the test suite again:\n   ```bash\n   python -m pytest test/transformers/test_{kernel}.py -xvs\n   ```\n3. Set `optimization_applied = false`\n4. Record the failure reason and all error outputs\n5. Proceed to Step 4 to generate the report (the report documents what was attempted even if it could not be applied)\n\n**Do NOT proceed to finalize (Step 5) if the optimization was reverted. The report must clearly state the optimization was not applied.**\n\n### Step 4: Run Checkstyle\n\nRun the project's checkstyle validation:\n\n```bash\nmake checkstyle\n```\n\n#### If checkstyle PASSES:\nProceed to Step 5.\n\n#### If checkstyle FAILS:\nAuto-fix with ruff:\n\n```bash\nruff check . --fix && ruff format .\n```\n\nRetry checkstyle once:\n\n```bash\nmake checkstyle\n```\n\nIf it still fails after auto-fix, report the remaining checkstyle issues but do NOT revert the kernel (checkstyle is not a correctness gate). Note the issues in the report so they can be fixed manually.\n\n### Step 5: Generate Comparison Plots\n\nGenerate 3-way comparison plots showing the original Liger kernel, the optimized kernel, and the HuggingFace/PyTorch baseline. This gives the user a visual demonstration of the improvement.\n\nThe procedure has 3 phases: (1) benchmark both kernels for plots, (2) generate plots, (3) clean the CSV so only the optimized kernel's data remains permanently.\n\n#### Step 5a: Benchmark Both Kernels for Plotting\n\n1. **Benchmark the original kernel as \"liger_original\"**: Temporarily restore the original kernel and run benchmarks with a custom provider name:\n   ```bash\n   # Swap in original kernel\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Edit benchmark script: replace \"liger\" with \"liger_original\" in kernel_providers\n   sed -i 's/\"liger\"/\"liger_original\"/g' benchmark/scripts/benchmark_{kernel}.py\n   \n   # Run benchmarks (appends \"liger_original\" rows to CSV)\n   cd benchmark/scripts && python benchmark_{kernel}.py\n   \n   # Revert the benchmark script change\n   git checkout benchmark/scripts/benchmark_{kernel}.py\n   ```\n\n2. **Benchmark the optimized kernel as \"liger\"**: Re-apply the optimized kernel and run benchmarks normally. Use `--overwrite` so the \"liger\" rows reflect the optimized kernel:\n   ```bash\n   # Swap in optimized kernel\n   cp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Run benchmarks (writes \"liger\" rows with optimized kernel)\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   ```\n\nAt this point the CSV has 3 providers: `\"liger_original\"`, `\"liger\"` (optimized), and `\"huggingface\"`/`\"torch\"`. This is what we need for the 3-way comparison plots.\n\n#### Step 5b: Generate Plots\n\nGenerate plots while all 3 providers are in the CSV:\n\n```bash\n# Speed plots for all modes (forward, backward, full)\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name speed \\\n  --overwrite\n\n# Memory plot\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name memory \\\n  --overwrite\n```\n\nCopy plots to the optimization workspace:\n\n```bash\ncp benchmark/visualizations/{kernel}_speed_*.png optimization/{kernel}/\ncp benchmark/visualizations/{kernel}_memory_*.png optimization/{kernel}/\n```\n\n#### Step 5c: Clean CSV for Permanent State\n\nNow remove the temporary `\"liger_original\"` rows and ensure the CSV only contains clean data for the optimized kernel. Delete ALL rows for this kernel, then re-run the benchmark to append fresh data:\n\n```bash\n# Delete all existing rows for this kernel from CSV (keep header + other kernels)\nhead -1 benchmark/data/all_benchmark_data.csv > /tmp/benchmark_header.csv\ngrep -v \"^{kernel},\" benchmark/data/all_benchmark_data.csv > /tmp/benchmark_other.csv || true\ncat /tmp/benchmark_header.csv <(tail -n +2 /tmp/benchmark_other.csv) > benchmark/data/all_benchmark_data.csv\n\n# Re-run benchmarks with optimized kernel (appends fresh rows)\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nAfter this, `all_benchmark_data.csv` contains:\n- `\"liger\"` — the **optimized** kernel (clean fresh data)\n- `\"huggingface\"` / `\"torch\"` — the baseline (also fresh)\n- NO `\"liger_original\"` — that was only needed for the plots\n\n#### Step 5c: Generate Additional Custom Plots (Optional)\n\nIf the optimization shows interesting size-dependent behavior (e.g., bigger gains at larger sizes), generate targeted plots highlighting this. Use matplotlib directly if the standard visualizer does not cover the comparison you want to show.\n\nSave any additional plots to `optimization/{kernel}/`.\n\n### Step 6: Generate the Final Optimization Report\n\nWrite the report to `optimization/{kernel}/report.md`. The report is the permanent record of the entire optimization effort.\n\nThe report must include these sections:\n\n1. **Summary**: kernel name, GPU, goal, variants tried, best variant, improvement %, and whether the optimization was applied (`Yes` / `No — reverted due to test failures`)\n2. **Baseline Profile**: bottleneck classification with evidence, NCU highlights if available\n3. **Variant Comparison**: include the table from `optimization/{kernel}/comparison.md` (produced by the Optimizer)\n4. **Winning Changes**: what was changed, why it worked, and a key code diff. Run `diff optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py` — do NOT write the diff from memory.\n5. **Rejected Approaches**: for each rejected variant, what was tried and why it failed\n6. **Comparison Plots**: reference the plots generated in Step 5 with file paths\n7. **Correctness Verification**: smoke test results, full test suite results, checkstyle results\n8. **Recommendations**: unexplored optimization opportunities, suggested follow-ups\n9. **Reproduction** section:\n   ```bash\n   # Reproduce baseline benchmarks:\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   # Revert to original:\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n\nInclude the variant comparison data from `optimization/{kernel}/comparison.md` (produced by the Optimizer).\n\n### Report Quality Checklist\n\nBefore saving the report, verify:\n\n- [ ] All variants are listed in the comparison table (including rejected ones)\n- [ ] Speed ratios are calculated consistently (higher = faster)\n- [ ] Memory ratios are calculated consistently (higher = less memory used)\n- [ ] The winning variant is clearly marked\n- [ ] The bottleneck classification matches the profile from Stage 1\n- [ ] If optimization was reverted, this is stated in the Summary, the header, and Section 6\n- [ ] Reproduction commands use the correct file paths\n- [ ] The code diff in Section 4 is accurate (diff the actual files, do not write from memory)\n- [ ] Comparison plots are generated and referenced in the report\n\n### Step 7: Create Pull Request\n\nCreate a PR with the kernel code changes and the updated benchmark data. Do NOT include plots or optimization workspace files in the PR.\n\n#### Step 7a: Stage Only Relevant Files\n\n```bash\ngit add src/liger_kernel/ops/{kernel}.py\ngit add benchmark/data/all_benchmark_data.csv\n```\n\nDo NOT add:\n- `optimization/` directory (workspace files, notes, plots)\n- `benchmark/visualizations/` (generated plots -- these are for local review only)\n\n#### Step 7b: Create a Descriptive Commit\n\n```bash\ngit commit -m \"$(cat <<'EOF'\n[Perf] Optimize {kernel} kernel: {headline improvement}\n\n{2-3 sentence summary of the optimization technique and its impact.}\n\nKey changes:\n- {Change 1}: {what and why}\n- {Change 2}: {what and why}\n\nBenchmark results ({gpu_name}):\n- Speed (forward): {delta}% {faster/slower}\n- Speed (backward): {delta}% {faster/slower}\n- Speed (full): {delta}% {faster/slower}\n- Memory: {delta}% {reduction/increase}\nEOF\n)\"\n```\n\nIf the host AI runtime has its own commit-attribution convention (e.g., a `Co-Authored-By` trailer), append it per that runtime's guidelines.\n\n#### Step 7c: Push and Create PR\n\n```bash\ngit push -u origin HEAD\n```\n\nCreate the PR using `gh pr create`. The PR body must clearly explain:\n\n1. **What was optimized** -- which kernel, which pass, what bottleneck\n2. **What changes were made** -- specific code changes with rationale for each\n3. **How they help** -- tie back to the bottleneck diagnosis (e.g., \"reduced register pressure from 115 to 95 registers/thread, increasing occupancy from 12.5% to 25%\")\n4. **Benchmark results** -- before/after table with speed and memory numbers at representative sizes\n5. **Correctness** -- confirm all tests pass\n\nUse this PR template:\n\n```bash\ngh pr create --title \"[Perf] Optimize {kernel}: {headline}\" --body \"$(cat <<'EOF'\n## Summary\n\n{2-3 sentences: what was optimized, the key technique, and the headline result.}\n\n## Bottleneck Diagnosis\n\n- **Kernel**: `{kernel}` ({tier} — {tier_description})\n- **GPU**: {gpu_name} ({architecture})\n- **Classification**: {memory-bound / compute-bound / balanced}\n- **Root cause**: {specific bottleneck explanation}\n\n## Changes\n\n{For each change, explain WHAT was changed and WHY it helps:}\n\n### 1. {Change title}\n{Description of the code change and its performance rationale.}\n\n### 2. {Change title}\n{Description of the code change and its performance rationale.}\n\n## Benchmark Results\n\nTested on {gpu_name}. Values are median ms (speed) or MB (memory).\n\n### Speed (ms, lower is better)\n\n| x_value | Mode | Before | After | Improvement |\n|---------|------|--------|-------|-------------|\n| {x1} | forward | {before} | {after} | {delta}% |\n| {x1} | backward | {before} | {after} | {delta}% |\n| {x1} | full | {before} | {after} | {delta}% |\n| ... | ... | ... | ... | ... |\n\n### Memory (MB, lower is better)\n\n| x_value | Before | After | Improvement |\n|---------|--------|-------|-------------|\n| {x1} | {before} | {after} | {delta}% |\n| ... | ... | ... | ... |\n\n## Correctness\n\n- Full test suite: **PASSED** (`python -m pytest test/transformers/test_{kernel}.py -xvs`)\n- Checkstyle: **PASSED** (`make checkstyle`)\n\n## Test Plan\n\n- [x] All existing unit tests pass\n- [x] Benchmarks show improvement across all input sizes\n- [x] No regression on non-target metrics (speed/memory balance maintained)\nEOF\n)\"\n```\n\nIf the host AI runtime appends a \"Generated with\" footer (e.g., Claude Code, Cursor, Copilot), include it per that runtime's guidelines.\n\n**Important**: Do NOT include plots as image attachments in the PR. Plots are for local review only and live in the optimization workspace.\n\n### Step 8: Present the Before/After Summary\n\nInclude the plot file paths in the summary so the user can view them.\n\nDisplay a concise summary to the user:\n\n```\n=== Optimization Complete: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: {APPLIED / REVERTED}\nWinner: v{N} — {strategy}\n\nBefore → After (at largest x_value):\n  Speed (forward):  {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (backward): {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (full):     {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Memory:           {baseline_mb} MB → {optimized_mb} MB ({improvement}% reduction)\n\nVariants tried: {N}\nTest suite: PASSED ({N} tests)\nCheckstyle: PASSED\n\nReport: optimization/{kernel}/report.md\n```\n\nIf the optimization was reverted:\n\n```\n=== Optimization REVERTED: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: REVERTED — test suite failed after applying winner\n\nWinner candidate: v{N} — {strategy}\nFailure reason: {brief error description}\nFix attempts: 2 (both failed)\n\nThe original kernel has been restored. No changes were made to production code.\n\nVariants tried: {N} (see report for full comparison)\nReport: optimization/{kernel}/report.md\n```\n\n## Error Handling\n\n| Error | Action |\n|-------|--------|\n| Winning variant file not found | Report missing file, stop |\n| Original kernel snapshot missing | Report missing file, stop (cannot safely revert) |\n| Production kernel differs from snapshot | Warn user, ask for confirmation (interactive) or stop (autonomous) |\n| Test suite fails (after 2 fix attempts) | Revert to original, generate report documenting failure |\n| Checkstyle fails (after auto-fix) | Note in report, do NOT revert (not a correctness issue) |\n| Benchmark CSVs missing | Generate report with available data, note missing data |\n| Cannot write report | Print report contents to stdout as fallback |\n\n## Key Principles\n\n1. **Correctness is non-negotiable.** The full test suite is a hard gate. Never skip it, never ignore failures, never mark a failing optimization as successful.\n2. **Always preserve the original.** The snapshot at `optimization/{kernel}/original_{kernel}.py` is the safety net. Never delete or overwrite it.\n3. **The report is permanent.** Even if the optimization is reverted, the report documents everything that was tried. This prevents future engineers from repeating failed approaches.\n4. **Be honest in the report.** Report actual measured numbers. Do not round favorably. If a variant was only 0.3% faster, say 0.3%, not \"marginal improvement.\" If the optimization was reverted, say so prominently.\n5. **Diff, don't recall.** When writing the code diff in the report, actually diff the files. Do not write the diff from memory -- it will contain errors.\n",".agents/skills/liger-kernel-perf/optimization-strategies.md":"# Optimization Strategies Catalog\n\nReference catalog for the Optimizer Agent. Techniques are ordered by recommended trial order within each category. The agent should **always start with parameter tuning**, then apply diagnosis-driven techniques based on the bottleneck classification from the Profiler stage.\n\n---\n\n## Decision Framework\n\nAfter reading the optimization profile, follow this priority:\n\n1. **Always first**: Parameter Tuning (Section 1)\n2. **Memory-bound kernel**: Section 2 techniques, then Section 5\n3. **Compute-bound kernel**: Section 3 techniques, then Section 5\n4. **Balanced kernel**: Interleave Section 2 and Section 3, then Section 5\n5. **Architecture-specific**: Layer Section 4 techniques on top of the best variant when targeting a specific GPU\n\n---\n\n## 1. Always First: Parameter Tuning\n\n**Why first**: Parameter tuning is the cheapest technique (minimal code changes) and frequently delivers the largest single improvement. The current codebase uses `calculate_settings()` which picks BLOCK_SIZE as the next power of 2 and num_warps via a fixed heuristic. This is a reasonable default but leaves significant performance on the table for specific input shapes and GPU architectures.\n\n### Why NOT @triton.autotune\n\nLiger production kernels **cannot use `@triton.autotune`** for three architectural reasons:\n\n1. **Forward-backward context coupling**: Liger kernels use `torch.autograd.Function` where `forward()` computes `BLOCK_SIZE`/`num_warps` via `calculate_settings()`, stores them in `ctx`, and `backward()` retrieves them to launch its kernel with identical config. With `@triton.autotune`, these become compile-time constants managed by Triton -- they cannot be stored in `ctx`, and the backward kernel has no way to know which config the forward kernel auto-selected for a given input.\n\n2. **Multi-backend incompatibility**: NPU/Ascend does not support `num_warps` or `num_stages` at all (commit `8b965b9` explicitly removed them). `@triton.autotune` relies on these parameters.\n\n3. **Inference warmup latency**: Each new input shape triggers autotuning (benchmarks all configs). Adds 100ms+ first-call penalty per unique shape, unacceptable for LLM inference with variable sequence lengths.\n\nThe commented-out `@triton.autotune` decorators in `dyt.py` and `grpo_loss.py` are evidence of prior exploration that was abandoned for these reasons. Do NOT re-introduce `@triton.autotune` on production kernels.\n\n### Parameters to Tune\n\n| Parameter | Description | Typical Range |\n|-----------|-------------|---------------|\n| `BLOCK_SIZE` | Elements processed per program instance | 256, 512, 1024, 2048, 4096, 8192, 16384 |\n| `num_warps` | Warps per program instance (32 threads each) | 2, 4, 8, 16, 32 |\n| `num_stages` | Software pipelining depth | 1, 2, 3, 4, 5 |\n| `BLOCK_ROW` | Rows per program (block-row mode only) | 1, 2, 4, 8, 16, 32 |\n\n### Method: Manual Sweep with Improved Heuristic\n\nReplace `calculate_settings()` with a better heuristic derived from benchmarking. The variant must still pass `BLOCK_SIZE` and `num_warps` through `ctx` to the backward kernel.\n\n**Step 1**: Identify the current config for the benchmark's input sizes:\n```python\n# calculate_settings() returns (next_power_of_2(n), num_warps_heuristic)\n# For n_cols=4096: BLOCK_SIZE=4096, num_warps=8\n# For n_cols=8192: BLOCK_SIZE=8192, num_warps=16\n# Note: num_stages is never set (Triton defaults to 1)\n```\n\n**Step 2**: Create a variant with an improved selection function:\n```python\ndef optimized_settings(n):\n    \"\"\"Improved heuristic derived from benchmarking on {gpu_arch}.\"\"\"\n    BLOCK_SIZE = triton.next_power_of_2(n)\n    if BLOCK_SIZE > 65536:\n        raise RuntimeError(...)\n    # Tuned values (from benchmark sweep):\n    if BLOCK_SIZE >= 8192:\n        num_warps = 16\n        num_stages = 3   # NEW: calculate_settings doesn't set this\n    elif BLOCK_SIZE >= 2048:\n        num_warps = 8\n        num_stages = 2\n    else:\n        num_warps = 4\n        num_stages = 2\n    return BLOCK_SIZE, num_warps, num_stages\n```\n\n**Step 3**: Update the kernel launcher to pass `num_stages`:\n```python\n# The kernel launch gains num_stages (previously unset, defaulted to 1):\nkernel[(grid,)](..., BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages)\n```\n\n**Step 4**: Update `ctx` to also store `num_stages` and pass it to the backward kernel.\n\n**Key insight**: `num_stages` is the most commonly overlooked parameter. `calculate_settings()` does not set it at all, so Triton defaults to `num_stages=1`. Simply adding `num_stages=2` or `num_stages=3` to existing kernel launches is often the single biggest win, especially on Hopper.\n\n### Architecture-Specific Guidance\n\n- **Ampere (SM 8.x)**: `num_stages=2` is usually optimal. Higher stages rarely help and can increase register pressure. `num_warps` sweet spot is 4-16.\n- **Hopper (SM 9.0)**: Benefits significantly from higher `num_stages` (3-5) due to TMA and async pipeline support. Also tolerates `num_warps=32` better due to deeper warp scheduling.\n- **Blackwell (SM 10.0)**: Similar to Hopper but with even more pipeline depth tolerance. Start with `num_stages=4-5`.\n\n### When to Apply\n\nAlways. This is the mandatory first optimization attempt for every kernel.\n\n### Expected Impact\n\n5-40% speedup depending on how far the `calculate_settings()` heuristic is from optimal for the target shapes and GPU. Adding `num_stages` alone can yield 10-25% on Hopper.\n\n### Risks/Gotchas\n\n- Configs are shape-dependent. A config optimal for `n_cols=4096` may be suboptimal for `n_cols=128`. Always benchmark across ALL x_values from the existing benchmark suite.\n- If BLOCK_SIZE exceeds `n_cols`, the kernel still works (masked loads) but wastes warps. Prune configs where `BLOCK_SIZE >> n_cols`.\n- The `calculate_settings()` function enforces `MAX_FUSED_SIZE = 65536`. This limit exists to prevent register spilling. Respect this unless verified otherwise.\n- When adding `num_stages` to existing kernels, the backward kernel must also receive the same `num_stages` value via `ctx`.\n- On NPU/Ascend, `num_warps` and `num_stages` are ignored. Optimizations targeting these parameters are NVIDIA-specific.\n\n### Kernel Tiers\n\nAll tiers benefit. Tier 1 kernels often see the largest relative gains because their simple structure makes the config choice proportionally more impactful.\n\n---\n\n## 2. Memory-Bound Optimizations\n\nApply these when the profiler classifies the kernel as memory-bound (high memory throughput, low compute utilization, simple operations per load/store).\n\n### 2.1 Memory Access Coalescing\n\n**When to apply**: NCU shows low memory throughput efficiency, or kernel uses non-contiguous access patterns (e.g., strided access across rows, column-major access on row-major data).\n\n**How to implement**:\n\nEnsure that threads in a warp access adjacent memory addresses. In Triton, this means the innermost offset computation should produce contiguous addresses across the `tl.arange` dimension.\n\n```python\n# GOOD: contiguous access within a warp\ncol_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_idx * stride + col_offsets, mask=col_offsets < n_cols)\n\n# BAD: strided access (each thread hits a different cache line)\nrow_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_offsets * stride + col_idx)\n```\n\nIn Liger kernels, most access is already contiguous via `col_offsets = tl.arange(0, BLOCK_SIZE)`. Check for cases where stride-based pointer arithmetic breaks coalescing, especially in backward kernels or multi-dimensional grids.\n\n**Expected impact**: 10-50% when fixing uncoalesced access. Minimal gain if access is already coalesced.\n\n**Risks**: Changing access patterns may require transposing data or restructuring the grid, which can be invasive.\n\n**Kernel tiers**: All tiers, but Tier 2/3 backward passes most often have coalescing issues due to accumulation patterns.\n\n### 2.2 Cache Modifiers\n\n**When to apply**: Kernel reads data that is used exactly once (stream through), or reads data that is reused many times (should be cached aggressively). NCU shows poor L2 hit rate.\n\n**How to implement**:\n\nTriton supports cache modifiers via the `eviction_policy` parameter on `tl.load` and `tl.store`:\n\n```python\n# Data read once, don't pollute cache (streaming access)\nX_block = tl.load(X_ptr + offsets, eviction_policy=\"evict_first\")\n\n# Data reused across iterations, keep in cache\nW_row = tl.load(W_ptr + offsets, eviction_policy=\"evict_last\")\n\n# Default (let hardware decide)\nX_block = tl.load(X_ptr + offsets)  # no modifier\n```\n\n| Modifier | Triton API | When to Use |\n|----------|------------|-------------|\n| `.ca` (cache all) | `eviction_policy=\"evict_last\"` | Data reused across loop iterations or by other programs |\n| `.cs` (cache streaming) | `eviction_policy=\"evict_first\"` | Data used once; large tensors that would thrash cache |\n| `.wb` (write-back) | default store behavior | Standard write-back through cache hierarchy |\n| `.wt` (write-through) | not directly exposed | Write-through (rarely needed in Triton) |\n\n**Expected impact**: 5-20% on memory-bound kernels with predictable access patterns.\n\n**Risks**: Wrong eviction policy can degrade performance. `.evict_first` on reused data will cause redundant memory fetches. Always benchmark both policies against the default.\n\n**Kernel tiers**: Tier 2 and Tier 3 benefit most (weight reuse in backward passes, multi-pass algorithms).\n\n### 2.3 Reduce Memory Traffic\n\n**When to apply**: Kernel reads or writes data that could be computed on the fly, or writes intermediate results that are only consumed once.\n\n**How to implement**:\n\nLiger already uses several memory-saving patterns. Look for additional opportunities:\n\n```python\n# Pattern 1: Recompute instead of saving to memory\n# Instead of: save silu(a) to a buffer, load in backward\n# Do: recompute silu(a) in backward from saved a\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a  # recomputed in backward\n\n# Pattern 2: In-place gradient storage (gradient-in-forward trick)\n# Cross entropy already does this: stores gradient in the input tensor\n# Look for other kernels that allocate separate grad tensors unnecessarily\n\n# Pattern 3: Fuse producer-consumer pairs to avoid materializing intermediates\n# Instead of: Y = kernel1(X); Z = kernel2(Y)  (Y written to HBM, then read)\n# Do: Z = fused_kernel(X)  (Y lives only in registers/SRAM)\n\n# Pattern 4: Write in-place to dY for dX when shapes match\nif in_place:\n    dX = dY  # reuse memory, zero additional allocation\n```\n\n**Expected impact**: 10-40% from eliminating unnecessary memory round-trips. Memory savings also reduce peak memory.\n\n**Risks**: Recomputation trades memory for compute. Only beneficial when the kernel is memory-bound. Can also complicate backward pass correctness -- always verify gradient accuracy.\n\n**Kernel tiers**: All tiers. Tier 3 kernels (fused operations) have the most opportunity for intermediate elimination.\n\n### 2.4 Data Reuse (Pre-Load Outside Loops)\n\n**When to apply**: A value is loaded inside a loop but does not change across iterations. Common with weight tensors and scaling factors.\n\n**How to implement**:\n\n```python\n# BEFORE: weight loaded inside the loop (redundant loads if compiler doesn't hoist)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    W_block = tl.load(W_ptr + offsets)  # loaded every iteration if W is large\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_block\n\n# AFTER: hoist the invariant load\nW_row = tl.load(W_ptr + col_offsets, mask=col_offsets < n_cols)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_row  # reuse from registers\n```\n\nNote: Triton's compiler often hoists loads automatically, but it is not guaranteed for all patterns, especially when the loop body is complex.\n\n**Expected impact**: 5-15% when applicable.\n\n**Risks**: Hoisting large loads outside loops increases register pressure. If the weight vector is very large relative to BLOCK_SIZE, it may cause register spilling that offsets the benefit.\n\n**Kernel tiers**: Tier 2 (norm kernels with weight vectors), Tier 3 (cross-entropy with weight tensors).\n\n### 2.5 Block-Row Mode (Multiple Rows per Program)\n\n**When to apply**: The kernel processes one row per program instance, but each row is small (e.g., `n_cols <= 256`). This under-utilizes each SM because the program does very little work.\n\n**How to implement**:\n\nLiger's `rms_norm.py` already implements this as `_block_rms_norm_forward_kernel`. The pattern:\n\n```python\n@triton.jit\ndef kernel_block_row(\n    ...,\n    n_rows,\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n    BLOCK_ROW: tl.constexpr,  # number of rows per program\n):\n    # Each program processes BLOCK_ROW rows\n    row_idx = tl.program_id(0) * BLOCK_ROW + tl.arange(0, BLOCK_ROW)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    row_mask = row_idx < n_rows\n    col_mask = col_offsets < n_cols\n\n    # 2D load: [BLOCK_ROW, BLOCK_SIZE]\n    X = tl.load(\n        X_ptr + row_idx[:, None] * stride + col_offsets[None, :],\n        mask=row_mask[:, None] & col_mask[None, :],\n        other=0,\n    )\n    # Process all rows at once...\n```\n\nGrid: `(triton.cdiv(n_rows, BLOCK_ROW),)` instead of `(n_rows,)`.\n\nUse conditional dispatch to choose between row mode and block-row mode:\n\n```python\nif BLOCK_SIZE > 256 or n_rows < 4096 * 8:\n    # Row mode: one row per program\n    kernel_row[(n_rows,)](...)\nelse:\n    # Block-row mode: multiple rows per program\n    BLOCK_ROW = 16\n    kernel_block_row[(triton.cdiv(n_rows, BLOCK_ROW),)](\n        ..., BLOCK_ROW=BLOCK_ROW\n    )\n```\n\n**Expected impact**: 10-30% for kernels with small `n_cols` and large `n_rows`.\n\n**Risks**: Increases register pressure and shared memory usage. The 2D load pattern is less efficient when `BLOCK_ROW` causes register spilling. Also adds code complexity (two kernel variants). Does not help when `n_cols` is already large.\n\n**Kernel tiers**: Tier 1 and Tier 2 primarily. Many Liger Tier 1 kernels (swiglu, geglu, relu_squared) do not yet have block-row mode.\n\n### 2.6 Vectorized Loads\n\n**When to apply**: Kernel loads/stores many small elements individually. Memory bandwidth is limited by the number of transactions rather than total bytes.\n\n**How to implement**:\n\nTriton generally handles vectorization automatically when loads are contiguous and aligned. However, you can help by:\n\n1. Ensuring BLOCK_SIZE is a multiple of 8 (for fp16/bf16) or 4 (for fp32) to enable 128-bit loads.\n2. Ensuring pointer alignment (tensors created by PyTorch are typically aligned).\n3. Avoiding masked loads at boundaries where possible -- pad the input instead.\n\n```python\n# Triton auto-vectorizes contiguous loads. Just ensure alignment:\n# BLOCK_SIZE should be power of 2 (already enforced by calculate_settings)\n# Data should be contiguous (already enforced by ensure_contiguous decorator)\n\n# For explicit vectorization (rare, advanced):\n# Cast pointer to wider type, load fewer elements\n# This is rarely needed as Triton handles it, but can help in edge cases\n```\n\n**Expected impact**: 5-10% if loads are currently unvectorized. Usually Triton already handles this.\n\n**Risks**: Minimal. Padding inputs wastes a small amount of memory.\n\n**Kernel tiers**: All tiers, but gains are usually small since Triton auto-vectorizes well.\n\n---\n\n## 3. Compute-Bound Optimizations\n\nApply these when the profiler classifies the kernel as compute-bound (high compute utilization, low memory throughput utilization, complex math per element).\n\n### 3.1 Algorithmic Improvements\n\n**When to apply**: The kernel uses a naive algorithm where a more efficient one exists. Examples: naive softmax vs online softmax, explicit matrix inverse vs iterative method.\n\n**How to implement**:\n\nLiger already uses online softmax in cross_entropy. Look for similar opportunities:\n\n```python\n# Online algorithm pattern: single-pass with running statistics\n# Instead of: pass 1 = compute max, pass 2 = compute sum, pass 3 = compute output\n# Do: single pass maintaining running max and sum\n\nm = float(\"-inf\")\nd = 0.0\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    block_max = tl.max(X_block)\n    m_new = tl.maximum(m, block_max)\n    d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new))\n    m = m_new\n# Result: lse = m + tl.log(d)\n```\n\nOther algorithmic opportunities:\n- Welford's algorithm for variance (numerically stable single-pass)\n- Kahan summation for high-precision reduction\n- Fast approximations for transcendentals when full precision is not needed\n\n**Expected impact**: 20-60% when applicable (reducing O(n) passes).\n\n**Risks**: Numerical stability can differ from reference implementation. Always validate against PyTorch reference with tight tolerances. Online algorithms may accumulate floating-point errors differently.\n\n**Kernel tiers**: Tier 2 and Tier 3 primarily. Tier 1 kernels rarely have algorithmic improvement opportunities.\n\n### 3.2 Reduce Redundant Computation\n\n**When to apply**: The same sub-expression is computed multiple times within a kernel, or across forward and backward passes.\n\n**How to implement**:\n\n```python\n# BEFORE: sigmoid computed twice\nsilu_a = a_row * tl.sigmoid(a_row)\ngrad = dc_row * (silu_a * (1 - tl.sigmoid(a_row)) + tl.sigmoid(a_row)) * b_row\n\n# AFTER: compute once, reuse\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a\ngrad = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row\n```\n\nLiger already does this well (see swiglu backward). Look for:\n- Repeated `tl.exp()`, `tl.sigmoid()`, `tl.log()` calls with the same argument\n- Reduction results (sum, max) computed multiple times\n- Softmax probabilities recomputed instead of cached in registers\n\n**Expected impact**: 5-20% depending on the cost of the duplicated operation.\n\n**Risks**: Caching sub-expressions increases register pressure. Trade-off with register spilling.\n\n**Kernel tiers**: Tier 2 and Tier 3. Tier 1 kernels are typically simple enough that redundancy is already eliminated.\n\n### 3.3 Tensor Core Utilization\n\n**When to apply**: Kernel performs matrix-multiply-like operations (dot products, outer products) that could leverage tensor cores. This is most relevant for Tier 3 kernels with linear layers or attention-like patterns.\n\n**How to implement**:\n\nTensor cores require specific alignment:\n- Block dimensions must be multiples of 16 (for fp16/bf16) or 8 (for tf32)\n- Input/output matrices must be in supported layouts\n- Use `tl.dot()` which maps to tensor core instructions when dimensions align\n\n```python\n# Ensure BLOCK_M, BLOCK_N, BLOCK_K are multiples of 16\nBLOCK_M: tl.constexpr = 64  # multiple of 16\nBLOCK_N: tl.constexpr = 64  # multiple of 16\nBLOCK_K: tl.constexpr = 32  # multiple of 16\n\n# tl.dot maps to tensor cores automatically\nacc = tl.dot(A_block, B_block)  # [BLOCK_M, BLOCK_K] x [BLOCK_K, BLOCK_N]\n```\n\nFor `tl.dot` to use tensor cores:\n- Both operands must be 2D\n- Operand dtypes must be fp16, bf16, or tf32\n- Inner dimension K must be >= 16\n\n**Expected impact**: 2-10x for matrix-multiply-heavy kernels. Not applicable to element-wise or simple reduction kernels.\n\n**Risks**: Tensor core operations have lower precision than FP32 scalar operations. The `tf32` mode on Ampere+ truncates FP32 mantissa to 10 bits. Always validate numerical accuracy. Not all kernels have operations that map to tensor cores.\n\n**Kernel tiers**: Tier 3 only (fused_linear_cross_entropy, attention-like kernels). Not applicable to Tier 1 or most Tier 2 kernels.\n\n### 3.4 Register Pressure Reduction\n\n**When to apply**: NCU shows high register usage per thread (>128 registers) or occupancy is limited by register count. Kernel body has many live variables.\n\n**How to implement**:\n\n```python\n# Strategy 1: Use constexpr branches to eliminate dead code paths\n# Triton compiles out branches on constexpr values, reducing register usage\ncasting_mode: tl.constexpr  # compiler eliminates unused casting branches\nHAS_WEIGHT: tl.constexpr    # compiler eliminates weight-related registers when False\n\n# Strategy 2: Split large kernels into smaller passes\n# Instead of one kernel doing everything:\n# Kernel 1: forward + save intermediates\n# Kernel 2: backward using intermediates\n# This is a last resort as it adds kernel launch overhead\n\n# Strategy 3: Reduce live range of temporary variables\n# Reorder operations so temporary values are consumed immediately\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    # Use X_block immediately, don't hold it across iterations\n    result = compute(X_block)\n    tl.store(..., result)\n    # X_block and result are now dead, registers freed\n\n# Strategy 4: Use smaller data types for intermediates where precision allows\n# fp16 uses half the registers of fp32\nintermediate = X_block.to(tl.float16)  # only if precision loss is acceptable\n```\n\n**Expected impact**: 10-25% when occupancy is register-limited. Minimal impact when occupancy is already high.\n\n**Risks**: Splitting kernels adds launch overhead and memory traffic for intermediates. Reducing precision can cause numerical issues. Constexpr branches are already used extensively in Liger -- check for remaining opportunities.\n\n**Kernel tiers**: Tier 3 (most register pressure due to complex logic). Tier 2 backward passes (accumulation variables).\n\n### 3.5 2D Grid Tiling for Better SM Utilization\n\n**When to apply**: Kernel uses a 1D grid but the problem has two natural dimensions (e.g., batch and feature). A 1D grid may leave SMs idle when the number of programs is not a multiple of SM count.\n\n**How to implement**:\n\n```python\n# 1D grid (current Liger pattern for most kernels):\ngrid = (n_rows,)\n# If n_rows = 100 and SM count = 108, 8 SMs are idle in the last wave\n\n# 2D grid for better utilization:\ngrid = (triton.cdiv(n_rows, BLOCK_ROW), triton.cdiv(n_cols, BLOCK_COL))\n\n@triton.jit\ndef kernel_2d(\n    ...,\n    BLOCK_ROW: tl.constexpr,\n    BLOCK_COL: tl.constexpr,\n):\n    row_block = tl.program_id(0)\n    col_block = tl.program_id(1)\n    row_start = row_block * BLOCK_ROW\n    col_start = col_block * BLOCK_COL\n    # Process a tile of [BLOCK_ROW, BLOCK_COL]\n```\n\nA 2D grid is most useful when the problem is large in both dimensions and each tile does enough work to amortize launch overhead.\n\n**Expected impact**: 5-15% for kernels with suboptimal SM utilization. More impactful when `n_rows` is small relative to SM count.\n\n**Risks**: 2D tiling for reduction kernels requires careful handling of partial results and synchronization. More complex code. Not beneficial when `n_cols` fits in a single block.\n\n**Kernel tiers**: Tier 2 (reductions with large feature dimensions), Tier 3 (large fused operations).\n\n### 3.6 Loop Unrolling Hints\n\n**When to apply**: Kernel has a hot loop with a known iteration count or small trip count. Triton's compiler unrolls some loops automatically, but explicit hints can help.\n\n**How to implement**:\n\n```python\n# Triton unrolls loops with constexpr bounds automatically.\n# For dynamic bounds, you can help by:\n\n# Strategy 1: Make the loop bound a constexpr when possible\n# If n_cols is known at compile time (e.g., passed as constexpr):\nfor i in range(0, N_COLS, BLOCK_SIZE):  # N_COLS: tl.constexpr -> unrolled\n    ...\n\n# Strategy 2: Use tl.static_range for explicit unrolling\nfor i in tl.static_range(4):  # always unrolled\n    X_block = tl.load(X_ptr + i * BLOCK_SIZE + offsets)\n    ...\n\n# Strategy 3: Break a large loop into a known-count outer loop + remainder\nCHUNKS: tl.constexpr = n_cols // BLOCK_SIZE  # constexpr division\nfor i in tl.static_range(CHUNKS):\n    ...\n# Handle remainder separately\n```\n\n**Expected impact**: 5-10% for tight loops with few iterations. Diminishing returns for large trip counts.\n\n**Risks**: Excessive unrolling increases instruction cache pressure and compile time. Can also increase register pressure if the loop body is large.\n\n**Kernel tiers**: Tier 2 (reduction loops), Tier 3 (multi-pass loops).\n\n---\n\n## 4. Architecture-Specific Optimizations\n\nApply these on top of the best general-purpose variant when targeting a specific GPU architecture.\n\n### 4.1 Ampere (SM 8.x: A100, A10, A30)\n\n**Async Copy (cp.async)**:\n- Triton uses `num_stages` to control software pipelining, which maps to `cp.async` on Ampere.\n- `num_stages=2` is the sweet spot for most Ampere kernels. Going higher increases shared memory usage without benefit.\n- Triton handles `cp.async` automatically when `num_stages > 1`.\n\n**L2 Cache Persistence**:\n- A100 has 40MB L2 cache. Tensors smaller than this can benefit from persistence hints.\n- Use `eviction_policy=\"evict_last\"` on reused data (weights, biases).\n- For cross-entropy, the target tensor and weight tensor are good candidates.\n\n**Optimal Warp Counts**:\n- Ampere excels with 4-16 warps per SM. 32 warps often increases scheduling overhead.\n- For memory-bound kernels: fewer warps (4-8) with larger BLOCK_SIZE.\n- For compute-bound kernels: more warps (8-16) to hide latency.\n\n**Implementation**:\n```python\n# Ampere-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=2),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=2),\n```\n\n### 4.2 Hopper (SM 9.0: H100, H200)\n\n**TMA (Tensor Memory Accelerator)**:\n- Hopper's TMA offloads address computation and data movement to a dedicated unit.\n- Triton >= 3.0 exposes TMA via `tl.make_block_ptr` and experimental APIs.\n- TMA is most beneficial for 2D tiled access patterns (matmul-like kernels).\n\n```python\n# TMA-friendly block pointer (Triton >= 3.0)\nblock_ptr = tl.make_block_ptr(\n    base=X_ptr,\n    shape=(M, N),\n    strides=(stride_m, stride_n),\n    offsets=(row_start, col_start),\n    block_shape=(BLOCK_M, BLOCK_N),\n    order=(1, 0),  # row-major\n)\nX_block = tl.load(block_ptr, boundary_check=[0, 1])\n```\n\n**Higher num_stages**:\n- Hopper's async pipeline supports deeper prefetching. `num_stages=3-5` often outperforms 2.\n- The deeper pipeline hides memory latency more effectively on H100's higher-bandwidth HBM3.\n\n**WGMMA (Warp Group Matrix Multiply Accumulate)**:\n- H100 tensor cores are organized in warp groups (4 warps = 128 threads).\n- `num_warps` should be a multiple of 4 for optimal tensor core usage.\n- Triton maps `tl.dot` to WGMMA automatically when conditions are met.\n\n**Cluster Launch**:\n- Hopper supports thread block clusters (groups of SMs that can share SRAM).\n- Currently experimental in Triton. Potentially useful for large reductions.\n\n**Implementation**:\n```python\n# Hopper-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=4),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=32, num_stages=4),\n```\n\n### 4.3 Blackwell (SM 10.0: B200, GB200)\n\n**5th-Gen Tensor Cores**:\n- Higher throughput for FP8, FP16, BF16, and TF32 operations.\n- Block sizes aligned to 16 remain optimal. FP8 support enables even higher throughput.\n- Triton support is evolving; check Triton release notes for Blackwell-specific features.\n\n**Larger Shared Memory**:\n- Blackwell supports up to 256KB shared memory per SM (vs 228KB on Hopper).\n- Larger BLOCK_SIZE values become viable without spilling to registers or HBM.\n- Can enable larger `BLOCK_ROW` in block-row mode.\n\n**Higher Memory Bandwidth**:\n- HBM3E provides more bandwidth than HBM3. Memory-bound kernels see direct benefit.\n- The compute-to-memory ratio shifts, making some previously memory-bound kernels balanced or compute-bound.\n\n**Implementation**:\n```python\n# Blackwell-tuned config (aggressive block sizes)\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 16384}, num_warps=32, num_stages=5),\n```\n\n**Note**: Blackwell-specific Triton features are still maturing. Focus on general optimizations first, then add Blackwell-specific configs as Triton support stabilizes.\n\n---\n\n## 5. Balanced Optimizations\n\nThese techniques help both memory-bound and compute-bound kernels.\n\n### 5.1 Kernel Fusion Opportunities\n\n**When to apply**: Two or more kernels execute sequentially on the same data, writing intermediate results to HBM between them. The fused version keeps intermediates in registers/SRAM.\n\n**How to implement**:\n\nLiger already excels at fusion (cross_entropy fuses softmax+loss+gradient; fused_linear_cross_entropy fuses matmul+CE). Look for remaining opportunities:\n\n```python\n# Pattern: operation followed by elementwise scaling\n# BEFORE: two kernel launches\ny = rms_norm_kernel(x, w)        # writes y to HBM\nz = elementwise_mul_kernel(y, s)  # reads y from HBM, writes z\n\n# AFTER: fused kernel\n@triton.jit\ndef fused_rms_norm_mul(X_ptr, W_ptr, S_ptr, Z_ptr, ...):\n    # RMS norm computation\n    y = x_row * rstd * (offset + w_row)\n    # Elementwise mul (no HBM round-trip)\n    z = y * s\n    tl.store(Z_ptr + offsets, z)\n```\n\nCommon fusion candidates in Liger:\n- Norm + activation (e.g., rms_norm + swiglu in some architectures)\n- Loss + metric computation (already done: cross_entropy + accuracy)\n- Dropout + residual add + norm\n\n**Expected impact**: 20-50% from eliminating HBM round-trips between fused operations.\n\n**Risks**: Fused kernels are harder to maintain, test, and debug. Register pressure increases. The fused kernel may have worse autotuning characteristics. Only fuse when the intermediate tensor is large relative to cache.\n\n**Kernel tiers**: Creates new Tier 3 kernels from combinations of Tier 1/2 kernels. Existing Tier 3 kernels may have remaining fusion opportunities.\n\n### 5.2 Better Grid Sizing\n\n**When to apply**: The number of programs (grid size) is much smaller or much larger than the SM count, leading to load imbalance or insufficient parallelism.\n\n**How to implement**:\n\n```python\nsm_count = torch.cuda.get_device_properties(device).multi_processor_count\n\n# For backward passes using SM-based parallelism (like rms_norm backward):\ngrid = (sm_count,)\nrows_per_program = math.ceil(n_rows / sm_count)\n\n# For forward passes: ensure grid is a small multiple of sm_count\n# If n_rows >> sm_count, this is automatic\n# If n_rows < sm_count, consider block-row mode or 2D grid to increase occupancy\n\n# For persistent kernels (advanced):\n# Launch exactly sm_count programs, each processes multiple work items\ngrid = (sm_count,)\nfor work_item in range(pid, total_work, NUM_SMS):\n    process(work_item)\n```\n\nThe rms_norm backward kernel already uses this SM-based grid pattern. Look for other backward kernels that use `(n_rows,)` grids but could benefit from SM-based parallelism.\n\n**Expected impact**: 5-20% from better load balancing and reduced tail effects.\n\n**Risks**: SM-based grids require careful work distribution and can introduce non-determinism if work items have unequal cost.\n\n**Kernel tiers**: All tiers. Most impactful for Tier 2 backward passes.\n\n### 5.3 Reduce Kernel Launch Overhead\n\n**When to apply**: Kernel is called many times with very short execution time (< 10 microseconds). Launch overhead dominates.\n\n**How to implement**:\n\n```python\n# Strategy 1: Batch multiple operations into one kernel\n# Instead of launching N small kernels, launch one kernel with N work items\n\n# Strategy 2: Use CUDA graphs (from PyTorch side, not inside Triton)\n# Captures a sequence of kernel launches and replays them with minimal overhead\ngraph = torch.cuda.CUDAGraph()\nwith torch.cuda.graph(graph):\n    output = liger_kernel(input)\n# Replay: graph.replay()\n\n# Strategy 3: Increase work per kernel (larger BLOCK_SIZE, block-row mode)\n# Already covered in autotuning and block-row mode sections\n```\n\n**Expected impact**: 5-15% for very short kernels. Negligible for kernels that run > 100 microseconds.\n\n**Risks**: CUDA graphs have limitations (fixed input shapes, no dynamic control flow). Batching work items increases kernel complexity.\n\n**Kernel tiers**: Tier 1 primarily (short element-wise kernels). Tier 2/3 kernels usually run long enough that launch overhead is negligible.\n\n### 5.4 Mixed Precision Strategies\n\n**When to apply**: Kernel uses FP32 throughout but only needs FP32 for specific operations (reductions, accumulations). Or kernel uses FP16/BF16 throughout but suffers from numerical issues.\n\n**How to implement**:\n\nLiger already uses mixed precision extensively (e.g., rms_norm with casting_mode). Look for additional opportunities:\n\n```python\n# Pattern 1: Compute in fp32 only where needed (Llama-style)\nX_row = tl.load(...)              # fp16/bf16\nX_fp32 = X_row.to(tl.float32)     # upcast for reduction\nmean_sq = tl.sum(X_fp32 * X_fp32) / n_cols  # fp32 reduction\nrstd = rsqrt(mean_sq + eps)        # fp32\nresult = (X_row * rstd.to(X_row.dtype)) * W  # back to original dtype\n\n# Pattern 2: Accumulate in fp32, store in original dtype\nacc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)\nfor ...:\n    acc += (dY_row * X_row).to(tl.float32)  # accumulate in fp32\ntl.store(dW_ptr, acc.to(original_dtype))     # store in original dtype\n\n# Pattern 3: Use BF16 instead of FP16 for better dynamic range\n# BF16 has same exponent bits as FP32, so less overflow/underflow risk\n# But lower mantissa precision than FP16\n```\n\n**Expected impact**: 5-20% from reducing unnecessary FP32 computation. Can also improve accuracy if precision is currently insufficient.\n\n**Risks**: Under-casting causes numerical errors. Over-casting wastes compute. The casting_mode pattern in Liger (llama vs gemma vs none) shows that different models need different precision strategies. Always validate against reference implementation.\n\n**Kernel tiers**: All tiers. Tier 2 reductions are most sensitive to precision. Tier 1 kernels with transcendentals (sigmoid, exp) benefit from targeted FP32 upcast.\n\n---\n\n## Quick Reference: Technique Selection by Symptom\n\n| Symptom | Likely Bottleneck | Try These Techniques |\n|---------|------------------|---------------------|\n| NCU memory throughput > 60% | Memory-bound | 2.1, 2.2, 2.3, 2.5 |\n| NCU compute throughput > 60% | Compute-bound | 3.1, 3.2, 3.3, 3.4 |\n| Low SM occupancy | Register pressure or grid size | 3.4, 5.2, 1 (num_warps) |\n| Small n_cols, large n_rows | Under-utilized programs | 2.5, 5.2, 1 (BLOCK_SIZE) |\n| Large n_cols, small n_rows | Poor SM utilization | 3.5, 5.2 |\n| Multiple sequential kernel launches | Fusion opportunity | 5.1 |\n| Forward fast, backward slow | Backward-specific issue | 2.3, 2.4, 5.2 (SM-based grid) |\n| Performance varies wildly by shape | Bad fixed config | 1 (autotune with key) |\n| High register count in NCU | Register pressure | 3.4, 1 (smaller BLOCK_SIZE) |\n| Poor L2 cache hit rate | Cache misuse | 2.2, 2.4 |\n\n## Quick Reference: Technique Applicability by Kernel Tier\n\n| Technique | Tier 1 (element-wise) | Tier 2 (reduction) | Tier 3 (fused/complex) |\n|-----------|----------------------|--------------------|-----------------------|\n| 1. Autotuning | High | High | High |\n| 2.1 Coalescing | Low (usually ok) | Medium | Medium |\n| 2.2 Cache modifiers | Low | Medium | High |\n| 2.3 Reduce traffic | Medium | Medium | High |\n| 2.4 Data reuse | Low | Medium | Medium |\n| 2.5 Block-row mode | High | High | Low |\n| 2.6 Vectorized loads | Low | Low | Low |\n| 3.1 Algorithmic | Low | High | High |\n| 3.2 Reduce redundancy | Low | Medium | High |\n| 3.3 Tensor cores | N/A | N/A | High |\n| 3.4 Register pressure | Low | Medium | High |\n| 3.5 2D grid tiling | Low | Medium | Medium |\n| 3.6 Loop unrolling | Low | Medium | Medium |\n| 4.x Arch-specific | Medium | Medium | Medium |\n| 5.1 Kernel fusion | Medium (fuse into Tier 3) | Medium | High |\n| 5.2 Grid sizing | Low | High | Medium |\n| 5.3 Launch overhead | High | Low | Low |\n| 5.4 Mixed precision | Medium | High | High |\n",".agents/skills/liger-kernel-perf/optimizer.md":"# Optimizer Agent\n\nRuns an autonomous optimization loop over an existing Liger kernel, generating multiple variants, benchmarking each against the baseline, and selecting the best performer while maintaining correctness.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\", \"cross_entropy\", \"swiglu\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `max_variants`: maximum number of variants to try (default 8)\n- `target_metric`: optional concrete target (e.g., \"forward under 0.3ms at hidden_size=4096\")\n- `optimization_profile`: the profile produced by the Profiler Agent at `optimization/{kernel}/profile.md`\n- `original_kernel`: the original kernel snapshot at `optimization/{kernel}/original_{kernel}.py`\n\n## Pre-Requisites\n\nBefore starting the optimization loop:\n\n1. Verify the workspace exists: `optimization/{kernel}/` with `original_{kernel}.py` and `profile.md`\n2. Verify the benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`\n3. Verify the test file exists: `test/transformers/test_{kernel}.py`\n4. Read the optimization profile thoroughly — extract bottleneck classification, tier, current config, and recommended strategy order\n5. Read the original kernel code thoroughly — understand every `@triton.jit` function, the `torch.autograd.Function`, and the launcher logic\n6. Read `optimization-strategies.md` for the full catalog of available techniques\n\n## Variant Numbering\n\n- v0 = baseline (original kernel, already benchmarked by the Profiler)\n- v1, v2, ... = optimization variants generated by this agent\n- Each variant gets two files:\n  - `optimization/{kernel}/{kernel}_vN.py` — the modified kernel code\n  - `optimization/{kernel}/{kernel}_vN_notes.md` — the lab notebook entry\n\n## The Optimization Loop\n\n```\nvariant_number = 1\nbest_variant = 0  (baseline)\nconsecutive_no_improvement = 0\n\nwhile variant_number <= max_variants:\n    1. Read ALL prior variant notes (v1 through v{variant_number-1})\n    2. Decide what to try next (see Strategy Selection below)\n    3. Generate variant code\n    4. Write variant notes (hypothesis section)\n    5. Run smoke test → if FAIL, mark REJECTED, increment variant_number, continue\n    6. Run full benchmark → record results\n    7. Check guardrails → if FAIL, mark REJECTED, increment variant_number, continue\n    8. Update variant notes with actual results and verdict\n    9. If improvement over best_variant: update best_variant, reset consecutive_no_improvement\n   10. Else: increment consecutive_no_improvement\n   11. If consecutive_no_improvement >= 2: STOP (diminishing returns)\n   12. If target_metric is set and met: STOP (goal achieved)\n   13. Increment variant_number\n```\n\n## Strategy Selection\n\n### Always Try Parameter Tuning First (v1)\n\nThe first variant MUST be a parameter sweep over `BLOCK_SIZE`, `num_warps`, and `num_stages`. This is the cheapest optimization with the highest expected payoff.\n\n**Why NOT @triton.autotune**: Liger kernels use `torch.autograd.Function` where the forward pass stores `BLOCK_SIZE` and `num_warps` in `ctx` and passes them to the backward kernel. With `@triton.autotune`, these become compile-time constants that cannot be stored in `ctx` — the backward kernel would not know which config the forward used. Additionally, NPU/Ascend backends do not support `num_warps` or `num_stages` at all. For these reasons, `@triton.autotune` is intentionally not used in Liger production kernels.\n\n**Instead, do a manual parameter sweep**: Replace the `calculate_settings(n_cols)` call with hardcoded values and benchmark each configuration:\n\n```python\n# In the variant, replace:\n#   BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n# With a specific config to test:\nBLOCK_SIZE = 2048   # Try: 256, 512, 1024, 2048, 4096, 8192\nnum_warps = 8       # Try: 2, 4, 8, 16, 32\nnum_stages = 2      # Try: 1, 2, 3, 4\n```\n\n**Sweep procedure for v1**:\n1. Read the current `calculate_settings(n_cols)` output for the benchmark's representative input sizes\n2. Create a variant that replaces `calculate_settings` with a custom function that returns a better config\n3. Test 3-5 promising configs (not the full cartesian product — use the architecture guidance below to prune)\n4. Pick the best config across ALL benchmark x_values (not just one size)\n5. The variant should implement a smarter `calculate_settings` replacement, not just a single hardcoded value\n\n**Architecture-specific guidance for config selection**:\n- **Ampere**: `num_stages=2`, `num_warps` up to 16, BLOCK_SIZE 1024-4096\n- **Hopper**: `num_stages=3-5`, `num_warps` up to 32, BLOCK_SIZE 2048-8192\n- **Blackwell**: `num_stages=4-5`, similar to Hopper\n\n**Important**: Keep the same forward/backward ctx pattern. Both kernels must use the same `BLOCK_SIZE` and `num_warps` for a given input. The variant must still pass these through `ctx` just like the original.\n\n### Subsequent Variants: Diagnosis-Driven\n\nAfter autotuning, choose the next technique based on:\n\n1. **The bottleneck classification from the profile** (memory-bound vs compute-bound)\n2. **Learnings from all prior variants** (what worked, what did not, what was surprising)\n3. **The recommended strategy order from the profile**\n\nConsult [optimization-strategies.md](optimization-strategies.md) for the full technique catalog. Common patterns by bottleneck type:\n\n**Memory-bound kernels** (Tier 1, simple element-wise):\n- Fuse multiple memory-bound operations into a single kernel\n- Add cache eviction hints (`tl.store(..., eviction_policy=\"evict_first\")`)\n- Increase BLOCK_SIZE to amortize memory latency\n- Use `tl.load(..., eviction_policy=\"evict_last\")` for reused data\n- Reduce unnecessary memory traffic (eliminate redundant loads/stores)\n- In-place operations where safe\n\n**Compute-bound kernels** (Tier 3, heavy math):\n- Reduce precision of intermediate computations where safe\n- Use fast-math intrinsics (`tl.math.fast_expf`, etc.)\n- Restructure computation to reduce total FLOPs\n- Improve instruction-level parallelism via loop unrolling\n- Increase occupancy by reducing register pressure\n\n**Balanced / Tier 2 kernels** (reductions):\n- Optimize the reduction pattern (two-pass vs single-pass)\n- Tune BLOCK_ROW for multi-row processing\n- Balance between register usage and parallelism\n- Consider split-K or hierarchical reductions for large dimensions\n- Block-row mode threshold tuning (the `n_rows < 4096 * 8` style checks)\n\n### Combining Techniques\n\nA single variant may combine multiple techniques if they are complementary (e.g., autotuning + cache hints). However, prefer isolated changes so that the lab notebook clearly attributes improvements to specific techniques. When combining, explain which part of the improvement is expected from each technique.\n\n### When All Strategies Exhausted\n\nIf all recommended strategies have been tried and the consecutive_no_improvement counter has not triggered, the agent may:\n- Try creative combinations of previous partial improvements\n- Attempt architecture-specific tricks (Hopper warp specialization, Blackwell TMA)\n- Revisit a rejected variant with a modified approach based on learnings\n\n## Generating a Variant\n\nFor each variant, produce a complete, self-contained kernel file at `optimization/{kernel}/{kernel}_vN.py`.\n\nRules:\n1. Start from the original kernel (not a previous variant) unless explicitly building on a prior variant's approach. If building on a prior variant, note which one and why.\n2. The variant file must be a drop-in replacement for `src/liger_kernel/ops/{kernel}.py` — same public API, same class names, same function signatures.\n3. Do not change the `torch.autograd.Function` class name or the public API.\n4. Do not remove safety checks, dtype handling, or multi-device support (XPU, NPU paths) unless the change is specifically about those areas.\n5. **Preserve ALL existing comments** — function docstrings, mathematical formula explanations, inline rationale comments, and section headers. Do NOT delete or shorten comments that explain what the code does or why.\n6. **Add explanatory comments for every optimization change** — each modified section should have a comment explaining what was changed and why it improves performance. For example:\n   ```python\n   # Optimization: compute dW before dX to allow register reuse.\n   # The original order kept 8 vectors live simultaneously (115 regs/thread).\n   # This reordering reduces peak liveness to 5-6 vectors (95 regs/thread).\n   ```\n7. **Preserve function-level docstrings** — if the original kernel functions have docstrings or parameter descriptions, keep them. Update them if the function signature changes.\n8. Add a comment block at the top of the file documenting what was changed:\n\n```python\n\"\"\"\nOptimization Variant {N}: {Short Title}\nBased on: original (or v{M} if building on a prior variant)\nChanges: {brief description of modifications}\n\"\"\"\n```\n\n## Variant Lab Notebook\n\nFor each variant, write `optimization/{kernel}/{kernel}_vN_notes.md` using the format in [templates/variant-notes.md](templates/variant-notes.md).\n\nKey requirements:\n- The Hypothesis must reference the profile's bottleneck classification and learnings from prior variants\n- Changes must include rationale for each modification with code snippets\n- Expected Outcome must be specific and quantitative (\"10-20% forward speed improvement\" not \"should be faster\")\n- Actual Results must include ALL x_values if improvement varies across sizes\n- Verdict: ACCEPTED (improves target without violating guardrails), REJECTED (fails smoke test/guardrails/no improvement), or MIXED (improves some metrics, regresses others within limits)\n- Learnings must include insights to carry forward to the next variant\n\n### Writing Notes Across the Loop\n\n- **Before smoke test**: Write the Hypothesis, Changes, and Expected Outcome sections\n- **After benchmark**: Fill in Actual Results, Verdict, and Learnings\n- **On smoke test failure**: Fill Actual Results with \"SMOKE TEST FAILED: {error}\", set Verdict to REJECTED, and write Learnings about what went wrong\n\n## Smoke Test Procedure\n\nThe smoke test validates basic correctness before investing time in a full benchmark. Execute it by temporarily swapping the variant into the source tree.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run a Minimal Correctness Check\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80\n```\n\nThis runs only float32 tests (the most numerically stable) with a timeout to catch hangs. The `-x` flag stops at the first failure.\n\n### Step 3: Interpret Results\n\n- **All tests pass**: Smoke test PASSES. Proceed to full benchmark.\n- **Any test fails**: Smoke test FAILS. Read the error carefully.\n  - If the failure is a **tolerance issue** (allclose fails but values are close): Note this in the variant notes. The optimization may have changed numerical behavior slightly. Try adjusting the variant to preserve precision, or note the tolerance difference.\n  - If the failure is a **crash or shape error**: The variant has a bug. Mark REJECTED immediately. Record the error in the notes for learning.\n  - If the failure is a **hang/timeout**: The variant has a deadlock or infinite loop. Mark REJECTED immediately.\n\n### Step 4: Restore the Original\n\n**CRITICAL**: Always restore the original kernel after the smoke test, regardless of outcome:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis must happen even if the smoke test crashes. Use a pattern that ensures restoration:\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py && \\\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80; \\\nSMOKE_EXIT=$?; \\\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py; \\\nexit $SMOKE_EXIT\n```\n\n## Full Benchmark Procedure\n\nAfter a variant passes the smoke test, run the full existing benchmark suite to get precise measurements.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run the Full Benchmark\n\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nThis runs speed benchmarks (forward, backward, full) and memory benchmarks across all configured x_values, comparing the \"liger\" provider against the baseline (\"huggingface\" or \"torch\").\n\n### Step 3: Capture Results\n\nThe benchmark writes results to `benchmark/data/all_benchmark_data.csv`. Copy the relevant rows:\n\n```bash\ncp benchmark/data/all_benchmark_data.csv optimization/{kernel}/benchmarks/vN_results.csv\n```\n\n### Step 4: Parse Results\n\nRead the CSV and extract:\n- Speed (ms) at the median (p50) for forward, backward, and full modes\n- Memory (MB) at the median (p50) for full mode\n- Values at each x_value to detect size-dependent behavior\n\nCompare against:\n- v0 baseline: the absolute improvement\n- Current best variant: whether this is the new best\n\n### Step 5: Restore the Original\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n## Guardrail Checks\n\nAfter benchmarking each variant, apply these guardrails. All comparisons are against the v0 baseline.\n\n### Guardrail 1: Non-Target Metric Regression (>5%)\n\nIf the optimization goal is \"speed\", check that memory does not regress by more than 5%. If the goal is \"memory\", check that speed does not regress by more than 5%. If the goal is \"balanced\", check both.\n\n```\nif goal == \"speed\":\n    reject if memory_vN > memory_v0 * 1.05\nelif goal == \"memory\":\n    reject if speed_vN > speed_v0 * 1.05  (higher ms = slower = worse)\nelif goal == \"balanced\":\n    reject if memory_vN > memory_v0 * 1.05 OR speed_vN > speed_v0 * 1.05\n```\n\nCheck this at EVERY x_value. If the regression occurs only at certain sizes, note this in the variant notes but still reject — partial regressions are not acceptable.\n\n### Guardrail 2: Cross-Pass Regression (>10%)\n\nDo not sacrifice one pass (forward or backward) to improve the other beyond a 10% threshold:\n\n```\nreject if speed_fwd_vN > speed_fwd_v0 * 1.10 AND speed_bwd_vN < speed_bwd_v0\nreject if speed_bwd_vN > speed_bwd_v0 * 1.10 AND speed_fwd_vN < speed_fwd_v0\n```\n\nAn optimization that makes forward 15% slower to make backward 5% faster is rejected.\n\n### Guardrail 3: Smoke Test Failure\n\nAlready handled in the smoke test procedure. Any correctness failure means immediate rejection.\n\n### Handling Guardrail Failures\n\nWhen a variant is rejected by guardrails:\n1. Set Verdict to REJECTED in the variant notes\n2. Record exactly which guardrail was violated and by how much\n3. Record the Learnings — this information is valuable for future variants\n4. Do NOT increment consecutive_no_improvement (guardrail failures are not \"stagnation\")\n5. Proceed to the next variant\n\n## Stopping Conditions\n\nThe loop terminates when ANY of these conditions is met:\n\n### Condition 1: Budget Exhausted\n\n`variant_number > max_variants`\n\nReport: \"Optimization budget of {max_variants} variants exhausted.\"\n\n### Condition 2: Diminishing Returns\n\n`consecutive_no_improvement >= 2`\n\nTwo consecutive variants that pass guardrails but improve the target metric by less than 1% over the current best. This indicates the kernel is near its optimization ceiling for the current approach.\n\nReport: \"Stopping due to diminishing returns — last 2 accepted variants showed <1% improvement.\"\n\n### Condition 3: Target Metric Met\n\nThe user-specified `target_metric` has been achieved.\n\nReport: \"Target metric achieved by variant v{N}: {metric_description}.\"\n\n## Reading Prior Variant Notes\n\nBefore generating each new variant, the agent MUST read ALL prior variant notes. This is the accumulated learning mechanism that makes later variants smarter.\n\nWhen reading prior notes, extract:\n1. **What techniques were tried**: Avoid repeating the exact same approach\n2. **What worked**: Build on successful directions\n3. **What failed and why**: Avoid repeating mistakes\n4. **Surprising results**: These often reveal the true bottleneck\n5. **Size-dependent behavior**: Some optimizations help at large sizes but hurt at small sizes\n\nUse this information to form a more informed hypothesis for the next variant.\n\n## Final Comparison Table\n\nAfter the loop terminates (for any reason), produce a comparison table of ALL variants and save it to `optimization/{kernel}/comparison.md`.\n\nThe comparison must include:\n- A master table with ALL variants (accepted and rejected): Variant | Strategy | Speed Fwd/Bwd/Full (ms) | Memory (MB) | Verdict\n- Values at the representative x_value (typically the largest or most common) with percentage delta from baseline\n- Winner declaration with justification\n- Cumulative learnings summary (what worked, what didn't, why)\n\n## Error Recovery\n\n### Benchmark Script Crashes\n\nIf the benchmark script crashes (not a test failure — an actual crash):\n1. Restore the original kernel immediately\n2. Read the crash output\n3. If it is a Triton compilation error: the variant has a syntax or semantic error in the kernel. Mark REJECTED with the error.\n4. If it is a CUDA OOM: the variant uses too much memory for the GPU. Mark REJECTED. Note the memory constraint.\n5. If it is an unrelated infrastructure error: retry once. If it fails again, report to the user.\n\n### Workspace Corruption\n\nIf at any point `src/liger_kernel/ops/{kernel}.py` is found to differ from `optimization/{kernel}/original_{kernel}.py` when it should be the original:\n\n```bash\ndiff src/liger_kernel/ops/{kernel}.py optimization/{kernel}/original_{kernel}.py\n```\n\nIf they differ unexpectedly, restore immediately:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Agent Failure Mid-Loop\n\nIf the agent encounters an unrecoverable error:\n1. Restore the original kernel\n2. Save all variant notes written so far\n3. Write a partial comparison table with whatever data is available\n4. Report to the user with the current state and what went wrong\n\n## Important Constraints\n\n1. **Never leave a variant installed in the source tree.** The original kernel must always be restored after each smoke test and benchmark. This is the most critical invariant.\n2. **Never modify test files.** The existing tests are the correctness oracle. If a variant fails tests, the variant is wrong, not the tests.\n3. **Never modify benchmark files.** The existing benchmarks are the measurement standard. Variants must be measured by the same benchmark to enable fair comparison.\n4. **Always write notes before running experiments.** The hypothesis must be written before seeing results to avoid post-hoc rationalization.\n5. **Always read all prior notes before generating a new variant.** This is what makes the optimization loop learn from itself.\n6. **Prefer the simplest change that achieves the goal.** A 5-line autotuning change that yields 15% speedup is better than a 200-line rewrite that yields 20%.\n",".agents/skills/liger-kernel-perf/profiler.md":"# Profiler Agent\n\nProfiles an existing Liger kernel to diagnose performance bottlenecks and produce an optimization strategy.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `target_gpu`: Ampere / Hopper / Blackwell / auto-detect\n\n## Steps\n\n### Step 1: Set Up Workspace\n\n1. Create directory `optimization/{kernel}/` at the repo root\n2. Copy `src/liger_kernel/ops/{kernel}.py` to `optimization/{kernel}/original_{kernel}.py`\n3. Create `optimization/{kernel}/benchmarks/` subdirectory\n\n### Step 2: Detect GPU Architecture\n\nRun:\n```bash\npython -c \"import torch; props = torch.cuda.get_device_properties(0); print(f'GPU: {props.name}, SM: {props.major}.{props.minor}, SMs: {props.multi_processor_count}, Mem: {props.total_mem // (1024**3)}GB')\"\n```\n\nClassify architecture:\n- SM 8.x → Ampere (A100, A10, etc.)\n- SM 9.0 → Hopper (H100, H200, etc.)\n- SM 10.0 → Blackwell (B200, etc.)\n\nIf user specified a target_gpu, note the mismatch (if any) but proceed with the actual hardware.\n\n### Step 3: Run Baseline Benchmarks\n\nRun the existing benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nRecord the baseline results. Copy the relevant rows from `benchmark/data/all_benchmark_data.csv` to `optimization/{kernel}/benchmarks/v0_baseline.csv`.\n\nParse and summarize the baseline numbers:\n- Speed: forward, backward, full (median ms at each x_value)\n- Memory: full (median MB at each x_value)\n- Identify the slowest operation mode and largest x_values\n\n### Step 4: Check for NCU Availability\n\n```bash\nwhich ncu 2>/dev/null && echo \"NCU_AVAILABLE\" || echo \"NCU_NOT_AVAILABLE\"\n```\n\nIf available, offer to run profiling (interactive mode) or run it (autonomous mode).\n\nTo determine the correct import and invocation, read the kernel's `torch.autograd.Function` class from `src/liger_kernel/ops/{kernel}.py` and construct a minimal reproduction script. For example, for `rms_norm`:\n\n```bash\nncu --set full --target-processes all -o optimization/{kernel}/ncu_profile \\\n    python -c \"\nimport torch\nfrom liger_kernel.ops.rms_norm import LigerRMSNormFunction\nx = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16, requires_grad=True)\nw = torch.ones(4096, device='cuda', dtype=torch.bfloat16)\ny = LigerRMSNormFunction.apply(x, w, 1e-6, 0.0, 'llama', False)\ny.sum().backward()\n\"\n```\n\nAdapt the import path, function arguments, and tensor shapes for the target kernel.\n\nExtract key metrics from NCU output:\n- **SM Occupancy** (%)\n- **Compute Throughput** (% of peak)\n- **Memory Throughput** (% of peak)\n- **L1/L2 Cache Hit Rates**\n- **Warp Stall Reasons** (top 3)\n\n### Step 5: Analyze Kernel Code\n\nRead `src/liger_kernel/ops/{kernel}.py` thoroughly. Extract:\n\n1. **Tier classification**:\n   - Tier 1 (element-wise): No reductions. One row per program.\n   - Tier 2 (reduction): Cross-column reductions. May need SM-based parallelism.\n   - Tier 3 (fused/complex): Multi-pass algorithms, gradient-in-forward tricks.\n\n2. **Current configuration**:\n   - How BLOCK_SIZE is set (calculate_settings? hardcoded? autotune?)\n   - num_warps value\n   - num_stages value (if specified)\n   - Grid dimensions (1D, 2D, block-row mode?)\n\n3. **Memory access patterns**:\n   - Stride-based or contiguous?\n   - Cache modifiers used?\n   - In-place operations?\n   - What is saved_for_backward vs recomputed?\n\n4. **Compute patterns**:\n   - Precision-sensitive operations (sigmoid, rsqrt, exp, log)?\n   - Casting modes?\n   - Online/chunked algorithms?\n   - constexpr parameters?\n\n5. **Existing optimization opportunities**:\n   - Any TODO/FIXME comments about performance?\n   - Missing block-row mode variant?\n   - Suboptimal BLOCK_SIZE heuristic?\n   - Unnecessary memory traffic?\n\n### Step 6: Classify Bottleneck\n\nUse NCU data (if available) or heuristics:\n\n**Memory-bound indicators**:\n- NCU: Memory throughput > 60% of peak, compute throughput < 40%\n- Heuristic: Simple element-wise ops, large tensors, few arithmetic operations per load\n- Heuristic: Tier 1 kernels are almost always memory-bound\n\n**Compute-bound indicators**:\n- NCU: Compute throughput > 60% of peak, memory throughput < 40%\n- Heuristic: Complex math per element (exp, log, sigmoid chains), reductions, many operations per load\n- Heuristic: Tier 3 kernels with heavy computation\n\n**Balanced (both bottlenecks)**:\n- NCU: Both throughputs 30-60%\n- Heuristic: Tier 2 kernels often fall here\n\n### Step 7: Produce Optimization Profile\n\nWrite `optimization/{kernel}/profile.md` using the [optimization-profile template](templates/optimization-profile.md).\n\nThe profile must include:\n- GPU architecture and properties\n- Baseline benchmark summary (table)\n- NCU summary (if available)\n- Kernel tier classification\n- Bottleneck classification with evidence\n- Current configuration analysis\n- Recommended strategy order (from [optimization-strategies.md](optimization-strategies.md))\n- Specific opportunities identified in the code\n\n### Step 8: Present to User (Interactive Mode)\n\nShow:\n1. GPU detected: {name} ({architecture})\n2. Baseline performance summary table\n3. Bottleneck classification: {memory-bound / compute-bound / balanced}\n4. Evidence for classification\n5. Proposed optimization strategy order\n6. Estimated iteration count\n\nWait for user confirmation before proceeding to Stage 2.\n\nIn autonomous mode, skip this step and proceed directly.\n",".agents/skills/liger-kernel-perf/templates/optimization-profile.md":"# Optimization Profile: {kernel}\n\n> Generated by Profiler Agent (Stage 1)\n> Date: {date}\n\n## Kernel Identity\n\n| Field | Value |\n|-------|-------|\n| Kernel Name | `{kernel}` |\n| Source File | `src/liger_kernel/ops/{kernel}.py` |\n| Benchmark Script | `benchmark/scripts/benchmark_{kernel}.py` |\n| Test File | `test/transformers/test_{kernel}.py` |\n| Tier Classification | {tier} ({tier_description}) |\n\n## GPU Information\n\n| Property | Value |\n|----------|-------|\n| GPU Name | {gpu_name} |\n| Architecture | {gpu_arch} (SM {sm_version}) |\n| SM Count | {sm_count} |\n| Total Memory | {gpu_memory_gb} GB |\n| Target GPU (user-specified) | {target_gpu} |\n| Hardware Match | {hardware_match} |\n\n## Baseline Benchmarks\n\n### Speed (median ms)\n\n| x_value | Mode | {kernel} (Liger) | {baseline_provider} | Speedup |\n|---------|------|-------------------|----------------------|---------|\n| {x_val_1} | forward | {liger_fwd_1} | {baseline_fwd_1} | {speedup_fwd_1}x |\n| {x_val_1} | backward | {liger_bwd_1} | {baseline_bwd_1} | {speedup_bwd_1}x |\n| {x_val_1} | full | {liger_full_1} | {baseline_full_1} | {speedup_full_1}x |\n| {x_val_2} | forward | {liger_fwd_2} | {baseline_fwd_2} | {speedup_fwd_2}x |\n| {x_val_2} | backward | {liger_bwd_2} | {baseline_bwd_2} | {speedup_bwd_2}x |\n| {x_val_2} | full | {liger_full_2} | {baseline_full_2} | {speedup_full_2}x |\n| ... | ... | ... | ... | ... |\n\n### Memory (median MB)\n\n| x_value | {kernel} (Liger) | {baseline_provider} | Savings |\n|---------|-------------------|----------------------|---------|\n| {x_val_1} | {liger_mem_1} | {baseline_mem_1} | {mem_savings_1}% |\n| {x_val_2} | {liger_mem_2} | {baseline_mem_2} | {mem_savings_2}% |\n| ... | ... | ... | ... |\n\n### Observations\n\n- Slowest mode: {slowest_mode} ({slowest_mode_reason})\n- Largest x_value tested: {largest_x_val}\n- Performance scaling trend: {scaling_trend}\n\n## NCU Profiling Summary\n\n> Status: {ncu_status}\n\n{ncu_available_section_start}\n\n| Metric | Value | Assessment |\n|--------|-------|------------|\n| SM Occupancy | {sm_occupancy}% | {occupancy_assessment} |\n| Compute Throughput | {compute_throughput}% of peak | {compute_assessment} |\n| Memory Throughput | {memory_throughput}% of peak | {memory_assessment} |\n| L1 Cache Hit Rate | {l1_hit_rate}% | {l1_assessment} |\n| L2 Cache Hit Rate | {l2_hit_rate}% | {l2_assessment} |\n\n### Top Warp Stall Reasons\n\n1. {warp_stall_1} ({warp_stall_pct_1}%)\n2. {warp_stall_2} ({warp_stall_pct_2}%)\n3. {warp_stall_3} ({warp_stall_pct_3}%)\n\n{ncu_available_section_end}\n\n{ncu_unavailable_section_start}\n\nNCU is not available on this system. Bottleneck classification is based on code heuristics only.\n\n{ncu_unavailable_section_end}\n\n## Bottleneck Classification\n\n**Classification: {bottleneck_class}**\n\n### Evidence\n\n| Signal | Observation | Supports |\n|--------|-------------|----------|\n| {evidence_signal_1} | {evidence_obs_1} | {evidence_supports_1} |\n| {evidence_signal_2} | {evidence_obs_2} | {evidence_supports_2} |\n| {evidence_signal_3} | {evidence_obs_3} | {evidence_supports_3} |\n| ... | ... | ... |\n\n### Reasoning\n\n{bottleneck_reasoning}\n\n## Current Configuration Analysis\n\n| Parameter | Current Value | How Set | Notes |\n|-----------|--------------|---------|-------|\n| BLOCK_SIZE | {block_size} | {block_size_source} | {block_size_notes} |\n| num_warps | {num_warps} | {num_warps_source} | {num_warps_notes} |\n| num_stages | {num_stages} | {num_stages_source} | {num_stages_notes} |\n| Grid Dimensions | {grid_dims} | {grid_dims_source} | {grid_dims_notes} |\n\n### Memory Access Patterns\n\n- Stride pattern: {stride_pattern}\n- Cache modifiers: {cache_modifiers}\n- In-place operations: {in_place_ops}\n- Saved for backward vs recomputed: {save_vs_recompute}\n\n### Compute Patterns\n\n- Precision-sensitive ops: {precision_ops}\n- Casting mode: {casting_mode}\n- Algorithm type: {algorithm_type}\n- constexpr parameters: {constexpr_params}\n\n## Identified Optimization Opportunities\n\n{opportunity_list}\n\n## Recommended Strategy Order\n\nBased on the bottleneck classification ({bottleneck_class}) and kernel analysis:\n\n{strategy_order}\n\n### Estimated Iterations\n\n- Expected variants: {estimated_variants}\n- Rationale: {variants_rationale}\n\n---\n\n*Profile saved to `optimization/{kernel}/profile.md`*\n",".agents/skills/liger-kernel-perf/templates/variant-notes.md":"# Variant {variant_number}: {variant_title}\n\n> Kernel: `{kernel}`\n> Date: {date}\n> Parent variant: {parent_variant}\n\n## Hypothesis\n\n**Target bottleneck:** {target_bottleneck}\n\n**Why this change should help:**\n\n{hypothesis_reasoning}\n\n**References from prior variants:**\n\n{prior_variant_references}\n\n## Changes\n\n{changes_list}\n\n## Expected Outcome\n\n**Predicted improvement:** {predicted_improvement}\n\n**Reasoning:** {prediction_reasoning}\n\n## Actual Results\n\n### Speed (median ms)\n\n| Metric | x_value | Baseline (v0) | This Variant (v{variant_number}) | Delta % |\n|--------|---------|---------------|----------------------------------|---------|\n| Speed forward | {x_val_1} | {baseline_fwd_1} | {variant_fwd_1} | {delta_fwd_1}% |\n| Speed forward | {x_val_2} | {baseline_fwd_2} | {variant_fwd_2} | {delta_fwd_2}% |\n| Speed backward | {x_val_1} | {baseline_bwd_1} | {variant_bwd_1} | {delta_bwd_1}% |\n| Speed backward | {x_val_2} | {baseline_bwd_2} | {variant_bwd_2} | {delta_bwd_2}% |\n| Speed full | {x_val_1} | {baseline_full_1} | {variant_full_1} | {delta_full_1}% |\n| Speed full | {x_val_2} | {baseline_full_2} | {variant_full_2} | {delta_full_2}% |\n\n### Memory (median MB)\n\n| Metric | x_value | Baseline (v0) | This Variant (v{variant_number}) | Delta % |\n|--------|---------|---------------|----------------------------------|---------|\n| Memory full | {x_val_1} | {baseline_mem_1} | {variant_mem_1} | {delta_mem_1}% |\n| Memory full | {x_val_2} | {baseline_mem_2} | {variant_mem_2} | {delta_mem_2}% |\n\n## Guardrail Checks\n\n| Check | Threshold | Result | Status |\n|-------|-----------|--------|--------|\n| Non-target metric regression | <5% worse | {nontarget_result} | {nontarget_status} |\n| Cross-pass regression | <10% on one pass to marginally improve other | {crosspass_result} | {crosspass_status} |\n| Smoke test (single shape, float32, fwd+bwd) | Pass | {smoke_result} | {smoke_status} |\n\n## Verdict\n\n**{verdict}**\n\n{verdict_explanation}\n\n## Learnings\n\n{learnings}\n\n---\n\n*Variant code: `optimization/{kernel}/{kernel}_v{variant_number}.py`*\n*Benchmark data: `optimization/{kernel}/benchmarks/v{variant_number}_results.csv`*\n","AGENTS.md":"# AGENTS.md\n\nGuidance for AI coding assistants working in this repository.\n\n## Skills\n\nRepository-specific workflow guides (\"skills\") live in [`.agents/skills/`](.agents/skills/). Each subdirectory is a self-contained guide for a multi-stage workflow. Read the `SKILL.md` at the top of each subdirectory for an overview; the other files in that directory are referenced from `SKILL.md` and should be read on demand.\n\n| Skill | What it does |\n|-------|--------------|\n| [`liger-kernel-dev`](.agents/skills/liger-kernel-dev/SKILL.md) | Develops new Triton kernels from a PyTorch reference (or modifies existing kernels). 3-stage pipeline: Analyze → Generate → Validate. NVIDIA GPUs only. |\n| [`liger-autopatch`](.agents/skills/liger-autopatch/SKILL.md) | Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies an existing monkey-patch. 3-stage pipeline: Analyze → Generate → Validate. |\n| [`liger-kernel-perf`](.agents/skills/liger-kernel-perf/SKILL.md) | Optimizes the performance of an existing Liger Triton kernel. 3-stage pipeline: Profile → Optimize → Finalize. NVIDIA GPUs only. |\n\nThe skills are written to be runtime-agnostic — they describe the workflow as a sequence of stages a competent agent (or human) can follow. Where a stage says \"Follow the X workflow in `x.md`\", that's a directive to read and execute that file's instructions; runtimes that support parallel subagents may delegate the stage, but it is not required.\n\n## Vendor-specific shortcuts\n\nFor convenience, some assistants auto-discover skills from vendor-specific paths. These point at the canonical `.agents/skills/` directory:\n\n- `.claude/skills` → symlink → `.agents/skills` (for Claude Code)\n\nIf you're adding support for another assistant, add a symlink (or your tool's preferred adapter) pointing to `.agents/skills/`. Do not duplicate the content.\n\n## Repo conventions\n\n- Source layout: `src/liger_kernel/{ops,transformers}/` for Triton ops and `nn.Module` / HF wrappers respectively\n- Tests: `test/transformers/` (unit) and `test/convergence/{bf16,fp32}/` (model convergence)\n- Benchmarks: `benchmark/scripts/` (scripts) and `benchmark/data/all_benchmark_data.csv` (results)\n- Lint/format: `make checkstyle` (uses `ruff`)\n- Install dev mode: `pip install -e \".[dev]\"`\n\nSee `README.md` for the project overview and contribution guide.\n"},"files":{".agents/skills/liger-autopatch/SKILL.md":"---\nname: liger-autopatch\ndescription: \"Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies existing monkey-patching. Generates lce_forward, monkey-patch function, tests, and README entry. Use when adding a new model to Liger Kernel, when a user asks to patch an unsupported model, when extending MODEL_TYPE_TO_APPLY_LIGER_FN, or when modifying/updating/fixing an existing monkey-patch (e.g., adding a new kernel to an already-supported model, fixing instance patching, updating a patch for upstream HF changes).\"\n---\n\n# Liger Auto-Patch\n\nAdds Liger Kernel optimization support for a new HuggingFace model, or modifies existing monkey-patching, through a staged pipeline with human review between stages. Supports creating new model patches and modifying existing ones.\n\n## Mode Detection\n\n- **Create mode**: User asks to add/patch/support a new model → full pipeline (Analyze → Generate → Validate)\n- **Modify mode**: User asks to update/fix/change/extend an existing monkey-patch → lighter pipeline (Change Impact Analysis → Apply Changes → Validate)\n\nKeywords that suggest modify mode: update, fix, change, add [kernel] to [existing model], extend, modify, new activation, new norm, bug in patch, upstream changed\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Model Analyzer** workflow in [model-analyzer.md](model-analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nThis stage reads the HF `modeling_*.py` source and produces a **model profile** answering 12 architectural questions from [decision-matrix.md](decision-matrix.md).\n\n**Human checkpoint:** Present the profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md).\n\nGenerates/modifies up to 13 files:\n\n1. `src/liger_kernel/transformers/model/{model}.py` — NEW lce_forward\n2. `src/liger_kernel/transformers/monkey_patch.py` — MODIFY\n3. `src/liger_kernel/transformers/__init__.py` — MODIFY\n4. `src/liger_kernel/transformers/model/output_classes.py` — MODIFY if needed\n5. `test/transformers/test_monkey_patch.py` — MODIFY\n6. `test/convergence/bf16/test_mini_models.py` — MODIFY (FLCE path)\n7. `test/convergence/bf16/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n8. `test/convergence/fp32/test_mini_models.py` — MODIFY (FLCE path)\n9. `test/convergence/fp32/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n10. `test/convergence/bf16/test_mini_models_multimodal.py` — MODIFY if VL model\n11. `test/convergence/fp32/test_mini_models_multimodal.py` — MODIFY if VL model\n12. `test/utils.py` — MODIFY\n13. `README.md` — MODIFY\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns instance patching test, convergence test, and lint check. Retries up to 3 times on failure.\n\n**Human checkpoint:** Report final test results.\n\n## Pipeline (Modify Mode)\n\n### Stage 1: Change Impact Analysis\n\nRead the existing `apply_liger_kernel_to_{model_type}` function in `monkey_patch.py` and the relevant section of the upstream HF `modeling_{model_type}.py`. Produce a short change plan:\n\n- What is being added/changed/fixed\n- Which Liger kernel(s) are involved\n- Which files need modification (subset of the 13 files from create mode)\n- What the expected behavior should be after the change\n\n**Human checkpoint:** Present the change plan. Confirm before proceeding.\n\n### Stage 2: Apply Changes\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md) in **modify mode**.\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md). This stage is **mandatory** — do not skip it. At minimum, run:\n\n1. Instance patching test: `pytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs`\n2. All convergence tests for the model:\n   - `pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, bf16)\n   - `pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, bf16)\n   - `pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, fp32)\n   - `pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, fp32)\n   - If VL (multimodal) model, also run:\n     - `pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n     - `pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n3. Checkstyle: `make checkstyle`\n\n**Human checkpoint:** Report final test results.\n\n## Reference Files\n\n- [decision-matrix.md](decision-matrix.md) — 12 architectural decisions to resolve per model\n- [examples/llama-profile.md](examples/llama-profile.md) — Reference profile for standard dense model\n- [examples/gemma-profile.md](examples/gemma-profile.md) — Reference profile showing GeGLU + offset variant\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n",".agents/skills/liger-autopatch/code-generator.md":"# Code Generator Agent\n\nTakes a confirmed model profile (create mode) or a change plan (modify mode) and generates or modifies files for Liger Kernel support.\n\n## Mode\n\n- **Create mode** (default): Generating all files for a new model. Follow the full \"Files to Generate\" list below.\n- **Modify mode**: Making targeted changes to an existing monkey-patch. Follow the \"Modification Checklist\" section instead.\n\n## Pre-Requisites\n\nBefore generating, read the reference implementation closest to this model:\n- Dense → `src/liger_kernel/transformers/model/llama.py`\n- MoE → `src/liger_kernel/transformers/model/mixtral.py`\n- Vision-Language → `src/liger_kernel/transformers/model/qwen2_vl.py`\n- Gemma-family → `src/liger_kernel/transformers/model/gemma.py`\n\nAlso read the corresponding patching function in `monkey_patch.py` and the templates in [templates/](templates/).\n\n## Files to Generate\n\n### 1. `src/liger_kernel/transformers/model/{model_type}.py` (NEW)\n\nThe `lce_forward` function. See [templates/lce-forward-dense.md](templates/lce-forward-dense.md) or [templates/lce-forward-moe.md](templates/lce-forward-moe.md).\n\nKey rules:\n- Match the exact forward signature from HF's `ForCausalLM.forward`\n- Use `lce_maybe_trainable_lm_head` from `llama.py` (shared PEFT/FSDP utility)\n- If model needs custom loss args (e.g., softcapping), write a local helper instead\n\n### 2. `src/liger_kernel/transformers/monkey_patch.py` (MODIFY)\n\nThree changes — see [templates/monkey-patch-fn.md](templates/monkey-patch-fn.md):\n\n**A.** Add lce_forward import (~line 18-28):\n```python\nfrom liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n```\n\n**B.** Add `apply_liger_kernel_to_{model_type}` function with both class-level and instance-level patching paths.\n\n**C.** Add entry to `MODEL_TYPE_TO_APPLY_LIGER_FN` dict (~line 3067).\n\n### 3. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the function in three locations (maintain alphabetical order):\n- `TYPE_CHECKING` block\n- `__getattr__` monkey_patch_symbols set\n- `__all__` list extension\n\n### 4. `src/liger_kernel/transformers/model/output_classes.py` (MODIFY if needed)\n\nOnly for models needing custom output (MoE with `aux_loss`, VL with `rope_deltas`). Follow the existing guarded-import pattern in the file.\n\n### 5. `test/transformers/test_monkey_patch.py` (MODIFY)\n\nSee [templates/test-instance-patch.md](templates/test-instance-patch.md). Add availability checker + skipif-decorated test function using `inspect.getsource()` assertions.\n\n### 6. Convergence tests (MODIFY multiple files)\n\nSee [templates/test-convergence.md](templates/test-convergence.md). Every model needs entries in multiple convergence test files:\n\n**All text models (dense + MoE)** — add to these 4 files:\n- `test/convergence/bf16/test_mini_models.py` — FLCE path, bf16\n- `test/convergence/bf16/test_mini_models_with_logits.py` — non-FLCE path (tests RMSNorm/SwiGLU/RoPE only), bf16\n- `test/convergence/fp32/test_mini_models.py` — FLCE path, fp32\n- `test/convergence/fp32/test_mini_models_with_logits.py` — non-FLCE path, fp32\n\n**Vision-language models** — also add to these 2:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\nEach file needs: imports, availability guard, `MiniModelConfig` entry in `MINI_MODEL_SETUPS` dict, and a `pytest.param` entry in the parametrize block. The `MiniModelConfig` entry is identical across all files for the same model. The `pytest.param` tolerances differ — use bf16 tolerances (looser) for bf16 files and fp32 tolerances (tighter) for fp32 files. Copy tolerance values from a similar existing model (e.g., Llama for dense, Mixtral for MoE).\n\n### 7. `test/utils.py` (MODIFY)\n\nAdd `revert_liger_kernel_to_{model_type}` function that reloads the modeling module.\n\n### 8. `README.md` (MODIFY)\n\nAdd row to the Patching table under \"### Patching\":\n```\n| {ModelName} | `liger_kernel.transformers.apply_liger_kernel_to_{model_type}` | {Supported Operations} |\n```\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order** alongside similar existing entries. Never append to the end of a section — find the correct alphabetical position.\n- After generating all files, run `make checkstyle` to verify formatting. If it fails, run `ruff check . --fix && ruff format .` to auto-fix, then verify with `make checkstyle` again.\n\n## Modification Checklist (Modify Mode)\n\nBefore making changes, read the existing implementation:\n1. Read `apply_liger_kernel_to_{model_type}` in `monkey_patch.py`\n2. Read the existing test in `test_monkey_patch.py` for this model\n3. Read the relevant HF modeling source for context\n\n### Rules for All Modifications\n\n**R1. Both patching levels.** If adding a new kernel, it must appear in BOTH:\n  - Class-level patching (the main body of `apply_liger_kernel_to_{model_type}`)\n  - Instance-level patching (the `if model is not None` block)\n\n  Omitting one is the most common mistake.\n\n**R2. New parameter with default.** Every new kernel gets a bool parameter on the\n  apply function signature (e.g., `relu_squared: bool = True`). Default should be `True`\n  for kernels that are safe to enable by default, `False` otherwise.\n\n**R3. Update docstring.** Update the function's docstring to:\n  - Add an `Args` entry for the new parameter\n  - Remove any stale notes that the new kernel invalidates\n    (e.g., \"squared ReLU is not supported\" → remove if you're adding it)\n\n**R4. Update tests.** In the existing `test_apply_liger_kernel_to_instance_for_{model_type}`:\n  - Add import for the new Liger kernel class\n  - Add \"not yet patched\" assertion before `_apply_liger_kernel_to_instance`\n  - Add \"correctly patched\" assertion after\n  - Follow the exact pattern of existing assertions in the same test\n\n**R5. Check revert function.** Read `revert_liger_kernel_to_{model_type}` in `test/utils.py`.\n  The revert function uses `importlib.reload(modeling_{model_type})` to undo all patches.\n  This handles most cases automatically, but check if the new kernel requires additional\n  revert logic (e.g., if the kernel patches something outside the modeling module, or\n  replaces a global like `ACT2FN` that `importlib.reload` won't fully restore). Update\n  the revert function if needed.\n\n**R6. Run convergence tests.** Don't modify convergence test files unless the change\n  requires it (e.g., new mini model config fields). But DO run existing convergence\n  tests in the Validate stage to verify no regression. This is critical — the Validator\n  agent (Stage 3) handles this, but if you are generating code without a separate\n  Validate stage, run these yourself:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  ```\n  For VL (multimodal) models, also run:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  ```\n\n**R7. Update README.md.** If the change adds a visibly new capability to the model's\n  row in the patching table (e.g., a new operation), update the supported operations list.\n\n### Common Modification Patterns\n\n**Adding an activation kernel (e.g., relu_squared for nemotron):**\n- Import the Liger kernel class at the top of `monkey_patch.py`\n- Add bool parameter to apply function signature\n- Class-level: replace in `ACT2FN` dict or replace the MLP class\n- Instance-level: patch each `decoder_layer`'s activation/MLP\n- Test: `assert isinstance` checks on the activation/MLP\n\n**Adding a norm variant:**\n- Add bool parameter to apply function\n- Class-level: replace the norm class\n- Instance-level: use `_patch_rms_norm_module` or `_patch_layer_norm_module` on all norm attrs\n- Test: `assert isinstance` checks on norm modules\n\n**Fixing missing instance patching:**\n- Read the class-level patching to see what's patched\n- Add corresponding instance-level patches in the `if model is not None` block\n- Test: add assertions that were missing\n\n**Updating for upstream HF changes:**\n- Compare the current HF modeling file against what the patch assumes\n- Update class names, attribute names, forward signatures as needed\n- May require updating `lce_forward` if the base model's forward changed\n",".agents/skills/liger-autopatch/decision-matrix.md":"# Decision Matrix\n\nWhen analyzing a HuggingFace model for Liger Kernel support, you must resolve these 12 architectural decisions by reading the model's `modeling_*.py` source code.\n\n## 1. Norm Type\n\n**Question:** Does the model use RMSNorm, LayerNorm, or both?\n\n**How to detect:**\n- Search for `class *RMSNorm` in the modeling file → RMSNorm\n- Search for `nn.LayerNorm` usage → LayerNorm\n- Multimodal models often use both (RMSNorm for text, LayerNorm for vision)\n\n**Liger mapping:** `LigerRMSNorm` or `LigerLayerNorm`\n\n## 2. RMSNorm Casting Mode\n\n**Question:** How does the model handle dtype casting during normalization?\n\n**How to detect:** Read the RMSNorm forward method:\n- Casts input to fp32, computes variance, casts back → `\"gemma\"`\n- Computes variance in fp32 only (input stays original dtype) → `\"llama\"`\n- No casting at all → `\"none\"`\n\n**Default:** `\"llama\"` (most common)\n\n## 3. RMSNorm Offset\n\n**Question:** Does the weight have a +1.0 offset?\n\n**How to detect:** In the RMSNorm forward, look for `(1 + self.weight)` or `self.weight + 1`:\n- Present → `offset=1.0` (Gemma family)\n- Absent → `offset=0.0` (most models)\n\n## 4. RMSNorm In-Place\n\n**Question:** Can the backward pass modify dY in-place?\n\n**How to detect:** Check if the model has two sequential norm layers with a residual connection between them (like Gemma2's `pre_feedforward_layernorm` + `post_feedforward_layernorm`):\n- Sequential norms with residual → `in_place=False`\n- Otherwise → `in_place=True`\n\n## 5. MLP Activation Type\n\n**Question:** What activation function does the gated MLP use?\n\n**How to detect:** Read the MLP class forward method:\n- `silu` or `F.silu` → SwiGLU → `LigerSwiGLUMLP`\n- `gelu` or `gelu_new` or `gelu_fast` → GeGLU → `LigerGEGLUMLP`\n- Phi3-style (single gate+up projection split) → `LigerPhi3SwiGLUMLP`\n\n**Also check:** The config's `hidden_act` field.\n\n## 6. Dense vs MoE\n\n**Question:** Is the model dense, MoE, or hybrid (some layers dense, some MoE)?\n\n**How to detect:**\n- Search for `Expert`, `MoE`, `SparseMoe`, `TopK` routing classes\n- Check if decoder layers have a `block_sparse_moe` or `experts` attribute\n- Hybrid: check for `is_moe_layer` or conditional MoE per-layer\n\n**Liger mapping:**\n- Dense → standard patching\n- MoE (transformers v5) → `LigerExperts`\n- MoE (transformers v4) → `LigerBlockSparseTop2MLP`\n- Qwen3-style MoE → `LigerQwen3MoeSwiGLUMLP`\n\n## 7. Vision Components\n\n**Question:** Does the model have a vision encoder?\n\n**How to detect:**\n- Check for `pixel_values` in the `forward` signature\n- Look for a separate vision model class (e.g., `*VisionModel`)\n- Check config for `vision_config` or `text_config` sub-configs\n\n**If yes:**\n- Vision encoder norms are usually `nn.LayerNorm` → patch with `LigerLayerNorm`\n- Text and vision must be patched separately\n\n## 8. RoPE Variant\n\n**Question:** What type of positional embedding does the model use?\n\n**How to detect:** Search for the `apply_rotary_pos_emb` function:\n- Standard (q, k, cos, sin) → `liger_rotary_pos_emb` (rope=True)\n- Llama4-style → `liger_llama4_text_rotary_pos_emb`\n- Qwen2VL MRoPE → `liger_multimodal_rotary_pos_emb`\n- No rotary embedding or custom variant → `rope=False`\n\n## 9. Output Class\n\n**Question:** What return type does the model's ForCausalLM.forward use?\n\n**How to detect:** Read the return statement and type annotation:\n- Standard → `LigerCausalLMOutputWithPast`\n- MoE (has `aux_loss`) → `LigerMoeCausalLMOutputWithPast`\n- Custom VL output → create model-specific output class in `output_classes.py`\n\n## 10. Hidden State Access\n\n**Question:** How does the model access hidden states from base model output?\n\n**How to detect:** In the ForCausalLM.forward, after calling `self.model(...)`:\n- `outputs[0]` → most models (Llama, Mistral, Gemma, etc.)\n- `outputs.last_hidden_state` → Phi3, Qwen3.5 MoE, some newer models\n\n## 11. Logit Softcapping\n\n**Question:** Does the model apply softcapping to logits before loss?\n\n**How to detect:** Check config for `final_logit_softcapping`:\n- Present → pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`\n- Absent → no softcapping (most models)\n- **VL models:** Config path may be `self.config.text_config.final_logit_softcapping` instead of `self.config.final_logit_softcapping`. Check whether the model uses a composite config with `text_config` sub-config.\n\n**Models with softcapping:** Gemma2, Gemma3\n\n## 12. Decoder Layer Norm Names\n\n**Question:** What are the attribute names of norm layers in each decoder layer?\n\n**How to detect:** Read the decoder layer class `__init__`:\n- Standard: `input_layernorm`, `post_attention_layernorm`\n- Gemma2 extra: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- GLM4: `post_self_attn_layernorm`, `post_mlp_layernorm`\n- Some models: `q_norm`, `k_norm` on self_attn\n\nAlso check the final norm on the base model (usually `model.norm` or `model.final_layernorm`).\n",".agents/skills/liger-autopatch/examples/gemma-profile.md":"# Model Profile: Gemma\n\nThis profile demonstrates the key differences from Llama: GeGLU activation, RMSNorm offset, and Gemma casting mode.\n\n## Identity\n- model_type: gemma\n- causal_lm_class: GemmaForCausalLM\n- base_model_class: GemmaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.gemma.modeling_gemma\n- config_module: transformers.models.gemma.configuration_gemma\n\n## Normalization\n- norm_class: GemmaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: gemma (everything cast to fp32, then computed, then cast back)\n- offset: 1.0 (weight uses `1 + self.weight` pattern)\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: GemmaMLP\n- activation: gelu (uses GELU activation, not SiLU)\n- liger_mlp_class: LigerGEGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nSame as Llama — no extra parameters.\n\n## Mini Model Config\n```python\nGemmaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-6,\n    hidden_activation=\"gelu_pytorch_tanh\",\n)\n```\n\n## Key Differences from Llama\n\n1. **Activation**: Uses `geglu` parameter (not `swiglu`) in the patch function\n2. **RMSNorm**: Requires `offset=1.0` and `casting_mode=\"gemma\"`\n3. **MLP class**: `LigerGEGLUMLP` instead of `LigerSwiGLUMLP`\n4. **Patching uses partial**: `_patch_rms_norm_module_for_gemma = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)`\n\n## Gemma2 Additional Differences\n- `in_place: false` (residual between sequential norms)\n- Extra norm layers: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- Has `final_logit_softcapping` in config\n- Uses `LigerRMSNormForGemma2` variant\n",".agents/skills/liger-autopatch/examples/llama-profile.md":"# Model Profile: Llama\n\nThis is the reference profile for the simplest and most common model type.\n\n## Identity\n- model_type: llama\n- causal_lm_class: LlamaForCausalLM\n- base_model_class: LlamaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.llama.modeling_llama\n- config_module: transformers.models.llama.configuration_llama\n\n## Normalization\n- norm_class: LlamaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: llama\n- offset: 0.0\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: LlamaMLP\n- activation: silu\n- liger_mlp_class: LigerSwiGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nStandard parameters only:\n- input_ids, attention_mask, position_ids, past_key_values\n- inputs_embeds, labels, use_cache\n- output_attentions, output_hidden_states, return_dict\n- cache_position, logits_to_keep, skip_logits\n- **kwargs\n\nNo model-specific extra parameters.\n\n## Mini Model Config\n```python\nLlamaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-5,\n)\n```\n\n## Notes\n- Llama is the base reference for all dense LLaMA-family models\n- Most models (Mistral, Qwen2, OLMo, etc.) follow this exact pattern\n- The `lce_forward` for Llama reuses `lce_maybe_trainable_lm_head` for PEFT/FSDP support\n- Has `pretraining_tp` check (raises error if > 1)\n",".agents/skills/liger-autopatch/model-analyzer.md":"# Model Analyzer Agent\n\nReads a HuggingFace model's source code and produces a structured model profile.\n\n## Input\n\nA model identifier: model_type string (e.g., `\"nemotron\"`), class name (e.g., `\"NemotronForCausalLM\"`), or Hub model name.\n\n## Steps\n\n1. Locate `transformers/models/{model_type}/modeling_{model_type}.py` and `configuration_{model_type}.py` on disk\n2. Read [decision-matrix.md](decision-matrix.md) for the 12 decisions to resolve\n3. Read the modeling file and extract all architectural details\n4. Produce the profile in the format below\n\n## What to Extract\n\nFrom the modeling file, identify:\n\n- `*ForCausalLM`, `*Model`, `*DecoderLayer`, `*RMSNorm`/norm, `*MLP`, `*Expert`/`*MoE` classes\n- `ForCausalLM.forward` signature — note all parameters\n- Norm class `forward` — determine casting mode and offset\n- All norm attribute names on decoder layer and base model\n- MLP activation function and gate/up/down projection names\n- Config's `model_type` string, `final_logit_softcapping`, `pretraining_tp`\n\n## Output Format\n\n```markdown\n# Model Profile: {ModelName}\n\n## Identity\n- model_type: {from config}\n- causal_lm_class: {e.g., NemotronForCausalLM}\n- base_model_class: {e.g., NemotronModel}\n- base_model_prefix: {e.g., \"model\"}\n- modeling_module: transformers.models.{model_type}.modeling_{model_type}\n\n## Normalization\n- norm_type: RMSNorm | LayerNorm | both\n- casting_mode: llama | gemma | none\n- offset: 0.0 | 1.0\n- in_place: true | false\n- final_norm_attr: {e.g., \"model.norm\"}\n- decoder_norm_attrs: {list}\n- attn_norm_attrs: {e.g., \"self_attn.q_norm\" or \"none\"}\n\n## MLP\n- activation: silu | gelu | gelu_new\n- liger_mlp_class: LigerSwiGLUMLP | LigerGEGLUMLP | LigerPhi3SwiGLUMLP\n\n## Structure\n- type: dense | moe | hybrid\n- moe_expert_class: {if MoE}\n- shared_expert: true | false\n\n## Positional Embedding\n- rope_type: standard | llama4 | qwen2vl_mrope | none\n\n## Output\n- output_class: LigerCausalLMOutputWithPast | LigerMoeCausalLMOutputWithPast | custom\n- hidden_state_access: outputs[0] | outputs.last_hidden_state\n- has_logit_softcapping: true | false\n\n## Vision\n- has_vision: true | false\n- vision_norm_type: {if applicable}\n\n## Forward Signature\n{List model-specific parameters not in standard Llama signature}\n\n## Mini Model Config\n{Suggest minimal config values for testing — 2 layers, 32 hidden, 64 intermediate}\n```\n\n## Rules\n\n- Never guess — always read the source to determine each value\n- Read [examples/llama-profile.md](examples/llama-profile.md) and [examples/gemma-profile.md](examples/gemma-profile.md) as reference\n- Compare against existing supported models to find the closest match\n",".agents/skills/liger-autopatch/templates/lce-forward-dense.md":"# Template: lce_forward for Dense Models\n\nReference implementation: `src/liger_kernel/transformers/model/llama.py`. Read it first.\n\n## Structure\n\nThe file needs `lce_forward` and reuses `lce_maybe_trainable_lm_head` from `llama.py`.\n\n### Imports\n\n```python\nfrom typing import TYPE_CHECKING, List, Optional, Tuple, Union\nimport torch\nfrom liger_kernel.transformers.model.llama import lce_maybe_trainable_lm_head\nfrom liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss, unpack_cross_entropy_result\nfrom liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast\nif TYPE_CHECKING:\n    from transformers.cache_utils import Cache\n```\n\n### lce_forward body pattern\n\n1. Set defaults from config (output_attentions, output_hidden_states, return_dict)\n2. Call `self.model(...)` with all forwarded params\n3. Extract hidden states: `outputs[0]` or `outputs.last_hidden_state`\n4. Slice for `logits_to_keep`\n5. Extract `shift_labels = kwargs.pop(\"shift_labels\", None)`\n6. `skip_logits` logic — default True during training with labels\n7. Compute loss via `lce_maybe_trainable_lm_head` (skip_logits) or `self.lm_head` (else)\n8. Return `LigerCausalLMOutputWithPast`\n\n### Variations\n\n- **Logit softcapping**: Pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`. Requires a local helper instead of reusing `lce_maybe_trainable_lm_head`.\n- **Hidden state access**: Use `outputs.last_hidden_state` if HF model uses that\n- **pretraining_tp**: Remove check if model doesn't have this config field\n\n### When NOT to reuse llama's lce_maybe_trainable_lm_head\n\nIf the model passes extra kwargs to `LigerForCausalLMLoss` (like softcapping), write a local `_liger_for_causal_lm_loss` and inline the PEFT/FSDP handling from llama.py.\n",".agents/skills/liger-autopatch/templates/lce-forward-moe.md":"# Template: lce_forward for MoE Models\n\nReference implementation: `src/liger_kernel/transformers/model/mixtral.py`. Read it first.\n\n## Differences from Dense\n\n1. **Output class**: `LigerMoeCausalLMOutputWithPast` (includes `aux_loss`)\n2. **Extra param**: `output_router_logits` in forward signature\n3. **Auxiliary loss**: Compute `load_balancing_loss_func` from router logits\n4. **Return**: Include `aux_loss` and `router_logits` in output\n\n## Auxiliary Loss Pattern\n\n```python\naux_loss = None\nif output_router_logits:\n    router_logits = outputs.router_logits if return_dict else outputs[-1]\n    aux_loss = load_balancing_loss_func(\n        router_logits, self.num_experts, self.num_experts_per_tok, attention_mask,\n    )\n    if labels is not None and loss is not None:\n        loss += self.router_aux_loss_coef * aux_loss\n```\n\n## Monkey Patch Differences for MoE\n\nExpert patching is version-aware:\n\n```python\n# Transformers v5\nmodeling_module.{ModelExperts} = LigerExperts\n# Instance: _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\n\n# Transformers v4\nmodeling_module.{ModelExpertMLP} = LigerBlockSparseTop2MLP\n# Instance: for expert in decoder_layer.block_sparse_moe.experts: _patch_swiglu_module(...)\n```\n\nCheck exact attribute names in the HF decoder layer class.\n",".agents/skills/liger-autopatch/templates/monkey-patch-fn.md":"# Template: apply_liger_kernel_to_{model_type}\n\nReference: `apply_liger_kernel_to_llama` (dense) or `apply_liger_kernel_to_mixtral` (MoE) in `monkey_patch.py`.\n\n## Function Signature\n\n```python\ndef apply_liger_kernel_to_{model_type}(\n    rope: bool = True,              # False if model lacks standard RoPE\n    cross_entropy: bool = False,\n    fused_linear_cross_entropy: bool = True,\n    rms_norm: bool = True,          # Or layer_norm for LayerNorm models\n    swiglu: bool = True,            # Or geglu for Gemma-family\n    model: PreTrainedModel = None,\n) -> None:\n```\n\n## Body Pattern\n\nClass-level patching always runs (even when `model` is provided). Instance-level patching runs additionally when `model is not None`. They are NOT mutually exclusive — this matches the existing Llama/Mistral/Gemma patterns.\n\nThe only exception is `fused_linear_cross_entropy`: use `MethodType` binding when instance is given, class replacement otherwise.\n\n```python\nassert not (cross_entropy and fused_linear_cross_entropy), \"...\"\n\nfrom transformers.models.{model_type} import modeling_{model_type}\n\n# Class-level patching (always runs)\nif rope: modeling_{model_type}.apply_rotary_pos_emb = liger_rotary_pos_emb\nif rms_norm: modeling_{model_type}.{ModelRMSNorm} = LigerRMSNorm\nif swiglu: modeling_{model_type}.{ModelMLP} = LigerSwiGLUMLP\nif cross_entropy: modeling_{model_type}.CrossEntropyLoss = LigerCrossEntropyLoss\nif fused_linear_cross_entropy:\n    if model is not None:\n        model.forward = MethodType({model_type}_lce_forward, model)\n    else:\n        modeling_{model_type}.{ModelForCausalLM}.forward = {model_type}_lce_forward\n\n# Instance-level patching (only when model instance provided)\nif model is not None:\n    base_model = getattr(model, model.base_model_prefix, model)\n    if rms_norm: _patch_rms_norm_module(base_model.norm)\n    for decoder_layer in base_model.layers:\n        if swiglu: _patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)\n        if rms_norm:\n            _patch_rms_norm_module(decoder_layer.input_layernorm)\n            _patch_rms_norm_module(decoder_layer.post_attention_layernorm)\n```\n\n## Common Variations\n\n**Gemma-family** (offset + casting):\n```python\n_patch_rms_norm_module_for_model = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)\n```\n\n**MoE experts** (version-aware):\n```python\nif IS_TRANSFORMERS_V5_OR_LATER:\n    _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\nelse:\n    for expert in decoder_layer.block_sparse_moe.experts:\n        _patch_swiglu_module(expert, LigerBlockSparseTop2MLP)\n```\n\n**Multimodal VL instance patching:**\n\nFor VL models, extract text_model and vision_model from the composite instance using `isinstance` checks. Each VL model has different accessor paths — read the existing VL patches in `monkey_patch.py` for the exact pattern. Example from Llama4:\n\n```python\nif model is not None:\n    if isinstance(model, {Model}ForConditionalGeneration):\n        language_model = model.language_model          # or model.model.language_model\n        vision_model = model.vision_model              # or model.model.visual\n        text_model = language_model.model\n    elif isinstance(model, {Model}ForCausalLM):\n        text_model = model.model\n        vision_model = None\n    elif isinstance(model, {Model}TextModel):\n        text_model = model\n        vision_model = None\n    else:\n        raise ValueError(f\"Unsupported model type: {type(model)}\")\n\n    # Patch text layers\n    if text_model:\n        if rms_norm: _patch_rms_norm_module(text_model.norm)\n        for decoder_layer in text_model.layers:\n            ...\n\n    # Patch vision layers\n    if vision_model:\n        for layer in vision_model.encoder.layers:  # accessor varies by model\n            if layer_norm:\n                _patch_layer_norm_module(layer.layer_norm1)\n                _patch_layer_norm_module(layer.layer_norm2)\n```\n\nThe exact accessor paths (`model.language_model`, `model.model.language_model`, `model.model.visual`) vary per VL model. Always read the HF model's `__init__` to find the correct paths.\n",".agents/skills/liger-autopatch/templates/test-convergence.md":"# Template: Convergence Tests\n\nReference: existing entries across `test/convergence/` files. Read the Llama entries as baseline.\n\n## Which Files to Modify\n\nEvery model needs entries in **at least 4 files**:\n\n| File | What it tests | dtype |\n|------|-------------|-------|\n| `test/convergence/bf16/test_mini_models.py` | FLCE path (fused linear cross entropy) | bfloat16 |\n| `test/convergence/bf16/test_mini_models_with_logits.py` | Non-FLCE path (tests RMSNorm/SwiGLU/RoPE only, verifies logits) | bfloat16 |\n| `test/convergence/fp32/test_mini_models.py` | FLCE path | float32 |\n| `test/convergence/fp32/test_mini_models_with_logits.py` | Non-FLCE path | float32 |\n\n**Vision-language models** also need:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\n## Placement Rule\n\nWhen adding entries to existing files, always insert in **alphabetical order** alongside similar existing entries. Never append to the end of a section. Specifically:\n- Imports: alphabetical among other `apply_liger_kernel_to_*` / `revert_liger_kernel_to_*` imports\n- Availability guards: alphabetical among other try/except model import blocks\n- `MINI_MODEL_SETUPS` entries: alphabetical among other model setup blocks\n- `pytest.param` entries: alphabetical among other model param entries\n\n## What to Add in Each File\n\nEach file has a `MINI_MODEL_SETUPS` dict and a `pytest.param` block. Add to both.\n\n### 1. Imports (same across all files, insert in alphabetical order)\n\n```python\nfrom liger_kernel.transformers import apply_liger_kernel_to_{model_type}\nfrom test.utils import revert_liger_kernel_to_{model_type}\n```\n\n### 2. Availability Guard (same across all files)\n\nThe `try/except ImportError` handles version compatibility automatically: if the model doesn't exist in the installed transformers version (e.g., 4.52.0), the import fails and the tests skip. No explicit version check needed.\n\n```python\ntry:\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n    from transformers.models.{model_type}.modeling_{model_type} import {ModelForCausalLM}\n    {MODEL_UPPER}_AVAILABLE = True\nexcept ImportError:\n    {MODEL_UPPER}_AVAILABLE = False\n```\n\n### 3. MINI_MODEL_SETUPS Entry (identical across all files)\n\n```python\nif {MODEL_UPPER}_AVAILABLE:\n    MINI_MODEL_SETUPS[\"mini_{model_type}\"] = MiniModelConfig(\n        liger_kernel_patch_func=apply_liger_kernel_to_{model_type},\n        liger_kernel_patch_revert_func=revert_liger_kernel_to_{model_type},\n        model_class={ModelForCausalLM},\n        mini_model_config={ModelConfig}(\n            hidden_size=32, intermediate_size=64, num_hidden_layers=2,\n            num_attention_heads=2, num_key_value_heads=2, vocab_size=1024,\n        ),\n    )\n```\n\n### 4. pytest.param Entry (tolerances differ by dtype)\n\n**bf16 files** — looser tolerances (copy from similar existing model):\n```python\npytest.param(\n    \"mini_{model_type}\", 32, 1e-5, torch.bfloat16,\n    1e-2, 5e-2, 1e-1, 1e-2, 1e-2, 1e-2,\n    marks=pytest.mark.skipif(not supports_bfloat16(), reason=\"bfloat16 not supported\"),\n),\n```\n\n**fp32 files** — tighter tolerances:\n```python\n(\"mini_{model_type}\", 32, 1e-4, torch.float32, 1e-8, 2e-5, 5e-3, 1e-5, 5e-3, 1e-5),\n```\n\n### 5. Revert Function in test/utils.py (shared by all files)\n\n```python\ndef revert_liger_kernel_to_{model_type}(model_config):\n    from transformers.models.{model_type} import modeling_{model_type}\n    importlib.reload(modeling_{model_type})\n    model_config.model_class = modeling_{model_type}.{ModelForCausalLM}\n```\n\n## Key Differences Between Test Variants\n\n**`with_logits`**: Sets `fused_linear_cross_entropy=False, cross_entropy=False` and checks `output.logits` match. Tests that non-loss kernels are numerically correct on their own. Same `MINI_MODEL_SETUPS` entry, just add a `pytest.param`.\n\n**`multimodal`**: Uses `multimodal_collate_fn`, image datasets, and model configs with vision sub-config. Only for models with `has_vision: true`. Multimodal tests may also require fake tokenizer/processor configs in `test/resources/fake_configs/`. Check existing VL model entries (e.g., Qwen2-VL, Llava) for the pattern — if the model needs a tokenizer/processor that can't be created programmatically, add a fake config JSON file.\n\n## Mini Config Guidelines\n\nUse minimal values: hidden_size=32, intermediate_size=64, num_hidden_layers=2, num_attention_heads=2, vocab_size=1024. Add model-specific required fields (e.g., `num_local_experts` for MoE). Copy config structure from the most similar existing model.\n",".agents/skills/liger-autopatch/templates/test-instance-patch.md":"# Template: Instance Patching Test\n\nReference: `test_apply_liger_kernel_to_instance_for_llama` in `test/transformers/test_monkey_patch.py`.\n\n## Add Availability Checker (at top of file, near other `is_*_available` functions ~lines 59-240, in alphabetical order)\n\n```python\ndef is_{model_type}_available():\n    try:\n        import transformers.models.{model_type}\n        return True\n    except ImportError:\n        return False\n```\n\n## Add Test Function (after existing `test_apply_liger_kernel_to_instance_for_*` functions, in alphabetical order)\n\n```python\n@pytest.mark.skipif(not is_{model_type}_available(), reason=\"{model_type} not available\")\ndef test_apply_liger_kernel_to_instance_for_{model_type}():\n    from liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n\n    config = {ModelConfig}(hidden_size=32, intermediate_size=64, num_hidden_layers=2)\n    dummy_model_instance = AutoModelForCausalLM.from_config(config)\n\n    # Verify NOT patched\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) != inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)\n\n    # Apply patch\n    _apply_liger_kernel_to_instance(model=dummy_model_instance)\n\n    # Verify patched\n    assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource({model_type}_lce_forward)\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) == inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)\n```\n\n## Adapt For\n\n- **GeGLU models**: Use `LigerGEGLUMLP` in assertions\n- **MoE models**: Check `decoder_layer.mlp.experts.forward` against `LigerExperts`\n- **Extra norms**: Add assertions for `pre_feedforward_layernorm` etc.\n- **Vision**: Add assertions for `LigerLayerNorm` on vision layers\n",".agents/skills/liger-autopatch/validator.md":"# Validator Agent\n\nRuns tests on generated code and fixes issues iteratively.\n\n## Pre-Requisites\n\nVerify the project is installed in development mode: `python -c \"import liger_kernel\"` must succeed. If not, ask the user to run `pip install -e \".[dev]\"`.\n\n## Validation Steps\n\n### Step 1: Syntax Check\n\n```bash\npython -c \"from liger_kernel.transformers import apply_liger_kernel_to_{model_type}\"\n```\n\n### Step 2: Instance Patching Test\n\n```bash\npytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `ImportError: cannot import name` | Check exact class name in HF source |\n| `AttributeError: has no attribute` | Read decoder layer `__init__` for correct attr names |\n| `AssertionError: getsource mismatch` | Check class being replaced in monkey_patch |\n| `ModuleNotFoundError` | Add `@pytest.mark.skipif` with availability check |\n\n### Step 3: Convergence Tests (requires GPU)\n\nRun **all applicable** convergence tests. For text models, there are 4 files. For VL models, 6.\n\n**FLCE path (fused linear cross entropy):**\n```bash\npytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Non-FLCE path (logits verification — tests RMSNorm/SwiGLU/RoPE without loss fusion):**\n```bash\npytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Multimodal (only for VL models):**\n```bash\npytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `CUDA out of memory` | Reduce hidden_size/num_layers in mini config |\n| `RuntimeError: shape mismatch` | Check `self.config.hidden_size` usage |\n| `AssertionError: allclose` | Check casting_mode, offset, in_place settings |\n| `KeyError in config` | Add missing field to mini model config |\n\n### Step 4: Checkstyle\n\nRun the project's full lint and format check:\n\n```bash\nmake checkstyle\n```\n\nIf it fails, auto-fix with `ruff check . --fix && ruff format .`, then re-run `make checkstyle`.\n\n## Retry Logic\n\n**Hard failures** (Steps 1, 2, 4 — import, patching, checkstyle): These indicate code bugs. Read traceback, identify root cause, fix, re-run. Max 3 retries per step. After 3 failures, report the exact error and stop.\n\n**Soft failures** (Step 3 — convergence): If a convergence test fails after 3 retries with `allclose` tolerance errors, mark it as `SKIP (tolerance tuning needed)` in the report rather than `FAIL`. This means the patching is correct but tolerances need manual adjustment. Note this in the PR description so reviewers are aware.\n\nConvergence failures from other errors (shape mismatch, CUDA OOM, KeyError) are still hard failures — fix and retry.\n\n## Success Report\n\n```\n| Check                        | Status |\n|------------------------------|--------|\n| Import                       | PASS   |\n| Instance Patching            | PASS   |\n| Convergence bf16 FLCE        | PASS   |\n| Convergence bf16 with_logits | PASS   |\n| Convergence fp32 FLCE        | PASS   |\n| Convergence fp32 with_logits | PASS   |\n| Convergence multimodal       | SKIP (text-only model) |\n| Checkstyle                   | PASS   |\n```\n\nInclude any `SKIP (tolerance tuning needed)` entries in the PR description under a \"Known issues\" section.\n",".agents/skills/liger-kernel-dev/SKILL.md":"---\nname: liger-kernel-dev\ndescription: \"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.\"\n---\n\n# Liger Kernel Dev\n\nDevelops 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.\n\n## Mode Detection\n\n- **Create mode**: User asks to create/add/generate/write/build a new kernel → full pipeline\n- **Modify mode**: User asks to update/fix/change/extend an existing kernel → skip Analyze, modify files, then Validate\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Analyzer** workflow in [analyzer.md](analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nAccepts any input: local file, URL, code snippet, natural language description, or model component reference. Produces a standalone PyTorch reference implementation and a kernel profile.\n\n**Human checkpoint:** Present PyTorch reference + kernel profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Generator** workflow in [generator.md](generator.md).\n\nGenerates/modifies up to 8 files:\n\n1. `src/liger_kernel/ops/{kernel}.py` — NEW Triton kernels + autograd Function\n2. `src/liger_kernel/transformers/{kernel}.py` — NEW nn.Module wrapper\n3. `src/liger_kernel/transformers/functional.py` — MODIFY add functional API\n4. `src/liger_kernel/ops/__init__.py` — MODIFY export Function class\n5. `src/liger_kernel/transformers/__init__.py` — MODIFY export Module + `__all__`\n6. `test/transformers/test_{kernel}.py` — NEW unit tests\n7. `benchmark/scripts/benchmark_{kernel}.py` — NEW benchmark script\n8. `benchmark/data/all_benchmark_data.csv` — MODIFY (after benchmarks run)\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns checkstyle, unit tests (hard gate — stops on persistent failure), benchmarks, and generates plots. Optionally runs ncu profiling.\n\n**Human checkpoint:** Report final results with benchmark numbers and plots.\n\n## Pipeline (Modify Mode)\n\n1. Read existing kernel files to understand current implementation\n2. Understand the requested modification\n3. Make targeted changes (Generator handles this)\n4. Run full Validate stage (same as create mode)\n\n## Reference Files\n\n- [kernel-profile-format.md](kernel-profile-format.md) — Kernel profile schema and field descriptions\n- [examples/swiglu-profile.md](examples/swiglu-profile.md) — Tier 1 (element-wise) reference\n- [examples/rms-norm-profile.md](examples/rms-norm-profile.md) — Tier 2 (reduction) reference\n- [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md) — Tier 3 (fused/complex) reference\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n",".agents/skills/liger-kernel-dev/analyzer.md":"# Analyzer Agent\n\nUnderstands a PyTorch operation from any input form and produces a standalone PyTorch reference implementation + kernel profile.\n\n## Input Handling\n\nThe user may provide the operation in any form:\n\n1. **Local file path** → Read the file directly\n2. **URL** (GitHub, HuggingFace, etc.) → Fetch via WebFetch tool\n3. **Code snippet** → Pasted in the conversation\n4. **Natural language** → Mathematical description (e.g., \"element-wise SiLU(x) * y\")\n5. **Model component** → e.g., \"the MLP in Phi-4\" — locate in transformers source and extract\n\n## Steps\n\n### 1. Understand the Operation\n\n- Read/fetch the source code from whatever input was provided\n- Identify the mathematical operation (forward pass)\n- Derive the backward pass (gradient computation)\n- Identify all inputs, outputs, and their expected shapes/dtypes\n- Note any precision-sensitive operations that need float32 upcasting (sigmoid, rsqrt, exp, log, tanh)\n\n### 2. Write PyTorch Reference\n\nCreate a standalone implementation that:\n- Depends only on `torch` (no external libraries)\n- Implements both forward and backward behavior (either as an `nn.Module` or a plain function)\n- Will serve as the correctness baseline for testing\n- Is clean, readable, and well-named\n\n### 3. Classify Into Tier\n\nRead [kernel-profile-format.md](kernel-profile-format.md) for the full schema.\n\n**Tier 1 — Element-wise**: No reductions across dimensions. One row per program. Examples: SwiGLU, GeGLU, DyT.\n- Read reference: `src/liger_kernel/ops/swiglu.py`\n\n**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.\n- Read reference: `src/liger_kernel/ops/rms_norm.py`\n\n**Tier 3 — Fused/Complex**: Multi-pass algorithms, gradient-in-forward tricks, multiple outputs. Examples: CrossEntropy, FusedLinearCrossEntropy.\n- Read reference: `src/liger_kernel/ops/cross_entropy.py`\n\nAlso read the closest example profile:\n- Tier 1 → [examples/swiglu-profile.md](examples/swiglu-profile.md)\n- Tier 2 → [examples/rms-norm-profile.md](examples/rms-norm-profile.md)\n- Tier 3 → [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md)\n\n### 4. Produce Kernel Profile\n\nFill in all fields from [kernel-profile-format.md](kernel-profile-format.md).\n\n### 5. Present to User\n\nShow:\n1. The PyTorch reference implementation (full code)\n2. The kernel profile (all fields)\n3. Which existing kernel is closest (for the Generator to use as reference)\n\nWait for user confirmation before proceeding to Stage 2.\n",".agents/skills/liger-kernel-dev/examples/cross-entropy-profile.md":"# Kernel Profile: CrossEntropy (Tier 3 — Fused/Complex)\n\n## Identity\n- operation_name: cross_entropy\n- function_class_name: LigerCrossEntropyFunction\n- module_class_name: LigerCrossEntropyLoss\n- functional_name: liger_cross_entropy\n\n## Classification\n- tier: 3\n- tier_description: fused/complex\n- closest_existing_kernel: fused_linear_cross_entropy (extends this with linear layer fusion)\n\n## Forward Pass\n- forward_inputs:\n  - _input: shape (B*T, V), logits\n  - target: shape (B*T,), label indices\n  - weight: shape (V,) or None, class weights\n  - ignore_index: int, label to ignore\n  - label_smoothing: float\n  - reduction: str, \"mean\" | \"sum\" | \"none\"\n  - softcap: float or None\n- forward_outputs:\n  - loss: scalar or shape (B*T,) if reduction=\"none\"\n  - z_loss: optional auxiliary loss\n  - token_accuracy: optional accuracy metric\n  - predicted_tokens: optional argmax tokens\n- forward_computation: Two-pass online softmax + cross entropy loss with optional smoothing and softcapping\n- precision_sensitive_ops: [exp, log]\n\n## Backward Pass\n- backward_saved_tensors: [_input] — gradient is computed during forward and stored in-place\n- backward_recompute: none (gradient-in-forward trick)\n- gradient_formulas:\n  - d_input: already computed in forward pass and stored in _input tensor. Backward just scales by grad_output.\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row (one row = one token's logits over vocab)\n- block_size_source: custom — iterates over vocab in chunks of BLOCK_SIZE\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - weight: optional class weights\n  - ignore_index: int = -100\n  - lse_square_scale: float = 0.0\n  - label_smoothing: float = 0.0\n  - reduction: str = \"mean\"\n  - softcap: float or None = None\n  - return_z_loss: bool = False\n- learnable_params: none\n\n## Benchmarking\n- benchmark_variable: vocab_size (V)\n- benchmark_x_label: \"V\"\n- benchmark_x_values_suggestion: [4096, 8192, 16384, 32768, 65536, 131072]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {B: 8, T: 512, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Online softmax**: Two-pass algorithm. Pass 1: compute running max and logsumexp. Pass 2: compute softmax and gradients. Avoids materializing the full softmax vector.\n- **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.\n- **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.\n- **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.\n- **Multiple loss components**: Combines original CE loss, label smoothing loss, and z-loss (for training stability). Each component contributes to both loss and gradient.\n",".agents/skills/liger-kernel-dev/examples/rms-norm-profile.md":"# Kernel Profile: RMSNorm (Tier 2 — Reduction)\n\n## Identity\n- operation_name: rms_norm\n- function_class_name: LigerRMSNormFunction\n- module_class_name: LigerRMSNorm\n- functional_name: liger_rms_norm\n\n## Classification\n- tier: 2\n- tier_description: reduction\n- closest_existing_kernel: layer_norm (similar pattern with different normalization)\n\n## Forward Pass\n- forward_inputs:\n  - X: shape (B, T, H), input tensor\n  - W: shape (H,), weight tensor\n  - eps: float, epsilon for numerical stability\n  - offset: float, weight offset (0.0 for Llama, 1.0 for Gemma)\n  - casting_mode: str, \"llama\" | \"gemma\" | \"none\"\n- forward_outputs:\n  - Y: shape (B, T, H), normalized output\n- forward_computation: Y = (X / RMS(X)) * (W + offset), RMS = sqrt(mean(X^2) + eps)\n- precision_sensitive_ops: [rsqrt]\n\n## Backward Pass\n- backward_saved_tensors: [X, W, RSTD] — RSTD cached from forward to avoid recomputation\n- backward_recompute: none (RSTD is expensive to recompute)\n- gradient_formulas:\n  - dX: rstd * (dY*(W+offset) - (1/N) * rstd^2 * dot(dY*(W+offset), X) * X)\n  - dW: sum over (B,T) of dY * (X * rstd)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: forward uses one program per row `(n_rows,)`, backward uses SM-based partitioning `(sm_count,)` with `rows_per_program`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: true (for dW reduction — each SM accumulates partial dW, then summed)\n\n## Module Parameters\n- module_init_params:\n  - hidden_size: int\n  - eps: float = 1e-6\n  - offset: float = 0.0\n  - casting_mode: str = \"llama\"\n  - init_fn: str = \"ones\"\n  - in_place: bool = True\n  - elementwise_affine: bool = True\n- learnable_params:\n  - weight: shape (hidden_size,), init ones or zeros\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {M: 4096, eps: 1e-6, dtype: torch.float32}\n\n## Key Patterns\n\n- **RSTD caching**: Forward computes and stores `rstd = rsqrt(mean(X^2) + eps)` — 1 value per row, tiny memory cost, saves 4 ops in backward\n- **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)`\n- **Casting modes as constexpr**: `casting_mode` is `tl.constexpr` so the compiler eliminates dead branches\n- **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)\n- **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)\n",".agents/skills/liger-kernel-dev/examples/swiglu-profile.md":"# Kernel Profile: SwiGLU (Tier 1 — Element-wise)\n\n## Identity\n- operation_name: swiglu\n- function_class_name: LigerSiLUMulFunction\n- module_class_name: LigerSwiGLUMLP\n- functional_name: liger_swiglu\n\n## Classification\n- tier: 1\n- tier_description: element-wise\n- closest_existing_kernel: geglu (same structure, different activation)\n\n## Forward Pass\n- forward_inputs:\n  - a: shape (B, T, H), gate projection output\n  - b: shape (B, T, H), up projection output\n- forward_outputs:\n  - c: shape (B, T, H), silu(a) * b\n- forward_computation: c = silu(a) * b, where silu(x) = x * sigmoid(x)\n- precision_sensitive_ops: [sigmoid]\n\n## Backward Pass\n- backward_saved_tensors: [a, b] (reshaped to 2D in forward wrapper)\n- backward_recompute: recompute silu(a) and sigmoid(a) in backward\n- gradient_formulas:\n  - da: dc * (silu(a) * (1 - sigmoid(a)) + sigmoid(a)) * b\n  - db: dc * silu(a)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row, `(n_rows,)`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - config: HuggingFace model config object\n- learnable_params:\n  - gate_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - up_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - down_proj: Linear(intermediate_size, hidden_size, bias=False)\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"torch\", \"torch_compile\"]\n- benchmark_fixed_config: {BT: 4096, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Recomputation over saving**: Forward saves `a, b` but backward recomputes `sigmoid(a)` and `silu(a)` — saves memory, sigmoid is cheap\n- **In-place backward**: Writes gradients directly to `a_ptr` and `b_ptr` (the saved tensors) — saves allocation\n- **Float32 for sigmoid**: `a_row` cast to `tl.float32` before `tl.sigmoid`, result cast back via `.cast(b_row.dtype)`\n- **No intermediate allocations**: Forward kernel writes directly to output `c`; backward kernel overwrites saved `a, b`\n",".agents/skills/liger-kernel-dev/generator.md":"# Generator Agent\n\nTakes a confirmed kernel profile and PyTorch reference, generates all files for a production-ready Liger kernel.\n\n## Pre-Requisites\n\nBefore generating, read the existing implementation closest to this kernel's tier:\n- Tier 1 → `src/liger_kernel/ops/swiglu.py` + `src/liger_kernel/transformers/swiglu.py`\n- Tier 2 → `src/liger_kernel/ops/rms_norm.py` + `src/liger_kernel/transformers/rms_norm.py`\n- Tier 3 → `src/liger_kernel/ops/cross_entropy.py` + `src/liger_kernel/transformers/cross_entropy.py`\n\nAlso read the corresponding test and benchmark files for the reference kernel.\n\n## Files to Generate\n\n### 1. `src/liger_kernel/ops/{kernel}.py` (NEW)\n\nThe core Triton kernel. See [templates/ops-kernel.md](templates/ops-kernel.md).\n\nKey rules:\n- `@triton.jit` for forward and backward kernels\n- `torch.autograd.Function` subclass with `@ensure_contiguous` decorator\n- Use `calculate_settings(n_cols)` for BLOCK_SIZE and num_warps\n- Cast to `tl.float32` for precision-sensitive ops (sigmoid, rsqrt, exp, log, tanh)\n- Use `tl.program_id(0).to(tl.int64)` to avoid overflow\n- Use stride parameters for memory access, not hardcoded offsets\n- Prefer recomputation over saving for backward when computation is cheap\n\n### 2. `src/liger_kernel/transformers/{kernel}.py` (NEW)\n\nThe nn.Module wrapper. See [templates/module-wrapper.md](templates/module-wrapper.md).\n\nKey rules:\n- Import from `liger_kernel.ops` (not from the submodule directly)\n- Include `extra_repr` method\n- Match the PyTorch API as closely as possible\n\n### 3. `src/liger_kernel/transformers/functional.py` (MODIFY)\n\nAdd functional API. See [templates/functional-api.md](templates/functional-api.md).\n\nTwo changes:\n- Add import of the Function class at the top (alphabetical order)\n- Add `liger_{kernel}()` function (alphabetical order among existing functions)\n\n### 4. `src/liger_kernel/ops/__init__.py` (MODIFY)\n\nAdd import of the Function class. Insert in alphabetical order among existing imports. Use `# noqa: F401` suffix.\n\n### 5. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the Module class in three locations (alphabetical order in each):\n- Direct import at top (among \"Always-safe imports\")\n- `__all__` list\n- If the Module has a `TYPE_CHECKING` import or `__getattr__` entry needed, add there too\n\n### 6. `test/transformers/test_{kernel}.py` (NEW)\n\nUnit tests. See [templates/unit-test.md](templates/unit-test.md).\n\nKey rules:\n- Include a PyTorch reference implementation (copy from Analyzer output)\n- Test both forward and backward correctness\n- Parametrize over shapes (include non-power-of-2 and small/large sizes)\n- Parametrize over dtypes with appropriate tolerances\n- Use `set_seed()`, `assert_verbose_allclose()`, `infer_device()`, `supports_bfloat16()`\n- Import from `test.utils`, NOT `liger_kernel.test.utils`\n\n### 7. `benchmark/scripts/benchmark_{kernel}.py` (NEW)\n\nBenchmarks. See [templates/benchmark.md](templates/benchmark.md).\n\nKey rules:\n- Define `_setup_{kernel}()` function returning `(input_tensor, layer)`\n- Define `bench_speed_{kernel}()` using `run_speed_benchmark`\n- Define `bench_memory_{kernel}()` using `run_memory_benchmark`\n- Compare \"liger\" vs \"torch\" (and optionally \"torch_compile\") providers\n- Use `get_benchmark_model_config`, `estimate_kernel_peak_memory`, and sweep config utilities\n- Run both speed and memory benchmarks for forward, backward, and full modes\n- **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\n\n## Modify Mode\n\nWhen modifying an existing kernel:\n1. Read the existing kernel files thoroughly\n2. Understand the current implementation patterns\n3. Make targeted changes — do not rewrite entire files\n4. Add new test cases for new functionality\n5. Ensure existing tests still cover original behavior\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order**\n- Import style: one import per line, no multi-imports\n",".agents/skills/liger-kernel-dev/kernel-profile-format.md":"# Kernel Profile Format\n\nWhen analyzing an operation, produce a profile with all of these fields.\n\n## Fields\n\n```yaml\n# Identity\noperation_name: snake_case name (e.g., \"rms_norm\", \"dyt\", \"sparsemax\")\nfunction_class_name: \"Liger{PascalCase}Function\" (e.g., \"LigerRMSNormFunction\")\nmodule_class_name: \"Liger{PascalCase}\" (e.g., \"LigerRMSNorm\")\nfunctional_name: \"liger_{snake_case}\" (e.g., \"liger_rms_norm\")\n\n# Classification\ntier: 1 | 2 | 3\ntier_description: \"element-wise\" | \"reduction\" | \"fused/complex\"\nclosest_existing_kernel: which existing Liger kernel is most similar\n\n# Forward Pass\nforward_inputs:\n  - name: X\n    shape: \"(B, T, H)\"\n    dtype: \"input dtype\"\n  # ... list all inputs with shapes\nforward_outputs:\n  - name: Y\n    shape: \"(B, T, H)\"\n    dtype: \"same as input\"\nforward_computation: brief math description (e.g., \"Y = X * rsqrt(mean(X^2) + eps) * W\")\nprecision_sensitive_ops: list ops needing float32 (e.g., [rsqrt, sigmoid, exp])\n\n# Backward Pass\nbackward_saved_tensors: what to save_for_backward (e.g., [X, W, RSTD])\nbackward_recompute: what to recompute instead of saving (e.g., \"recompute silu in backward\")\ngradient_formulas:\n  dX: \"formula or description\"\n  dW: \"formula or description\"\n\n# Tiling Strategy\ngrid_dimensions: 1D | 2D\ngrid_description: \"one program per row\" | \"one program per (row_block, col_block)\" | etc.\nblock_size_source: \"calculate_settings(n_cols)\" | \"custom\" with explanation\nneeds_sm_parallelism: true | false (for weight gradient reduction in backward)\n\n# Module Parameters\nmodule_init_params:\n  - name: hidden_size\n    type: int\n  - name: eps\n    type: float\n    default: 1e-6\n  # ... list all __init__ parameters\nlearnable_params:\n  - name: weight\n    shape: \"(hidden_size,)\"\n    init: \"ones\"\n  # ... list learnable parameters\n\n# Benchmarking\nbenchmark_variable: \"hidden_size\" | \"seq_len\" | \"vocab_size\" | etc.\nbenchmark_x_label: human-readable label for x-axis\nbenchmark_x_values_suggestion: suggested range to sweep\nbenchmark_providers: [\"liger\", \"torch\"] (add \"torch_compile\" if applicable)\nbenchmark_fixed_config: dict of fixed params (e.g., {BT: 4096, dtype: torch.bfloat16})\n```\n\n## Naming Conventions\n\n| Component | Pattern | Example |\n|-----------|---------|---------|\n| Ops file | `src/liger_kernel/ops/{operation_name}.py` | `ops/dyt.py` |\n| Transformer file | `src/liger_kernel/transformers/{operation_name}.py` | `transformers/dyt.py` |\n| Test file | `test/transformers/test_{operation_name}.py` | `test_dyt.py` |\n| Benchmark file | `benchmark/scripts/benchmark_{operation_name}.py` | `benchmark_dyt.py` |\n| Function class | `Liger{PascalCase}Function` | `LigerDyTFunction` |\n| Module class | `Liger{PascalCase}` | `LigerDyT` |\n| Functional | `liger_{snake_case}` | `liger_dyt` |\n| Forward wrapper | `{snake_case}_forward` or `liger_{snake_case}_fwd` | `liger_dyt_fwd` |\n| Backward wrapper | `{snake_case}_backward` or `liger_{snake_case}_bwd` | `liger_dyt_bwd` |\n",".agents/skills/liger-kernel-dev/templates/benchmark.md":"# Benchmark Template\n\n## File: `benchmark/scripts/benchmark_{kernel}.py`\n\n```python\nimport os\nimport sys\n\nimport torch\n\nfrom benchmark_model_configs import compute_hidden_size_sweep_config\nfrom benchmark_model_configs import estimate_kernel_peak_memory\nfrom benchmark_model_configs import get_benchmark_model_config\nfrom utils import SingleBenchmarkRunInput\nfrom utils import SingleBenchmarkRunOutput\nfrom utils import parse_benchmark_script_args\nfrom utils import run_benchmarks\nfrom utils import run_memory_benchmark\nfrom utils import run_speed_benchmark\n\nfrom liger_kernel.utils import infer_device\n\ndevice = infer_device()\n\nsys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), \"../..\")))\n\n\ndef _setup_{kernel}(input: SingleBenchmarkRunInput):\n    \"\"\"Create input tensor and layer from benchmark config.\"\"\"\n    # Import both Liger and PyTorch reference implementations\n    from test.transformers.test_{kernel} import Liger{Kernel}\n    from test.transformers.test_{kernel} import Torch{Kernel}\n\n    cfg = input.extra_benchmark_config\n    hidden_size = input.x  # Or whatever the sweep variable is\n\n    x = torch.randn(cfg[\"BT\"], hidden_size, device=device, dtype=cfg[\"dtype\"], requires_grad=True)\n\n    if input.kernel_provider == \"liger\":\n        layer = Liger{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch\":\n        layer = Torch{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch_compile\":\n        layer = torch.compile(Torch{Kernel}(hidden_size=hidden_size).to(device))\n    else:\n        raise ValueError(f\"Invalid provider: {input.kernel_provider} for {kernel}\")\n\n    return x, layer\n\n\ndef bench_speed_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x])\n\n\ndef bench_memory_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode)\n\n\nBT = 4096  # Fixed batch*seq for sweeping hidden_size\n\n\nif __name__ == \"__main__\":\n    args = parse_benchmark_script_args()\n    model = get_benchmark_model_config(args.model)\n\n    def _probe():\n        probe_input = SingleBenchmarkRunInput(\n            x=model.hidden_size,\n            kernel_provider=\"torch\",\n            extra_benchmark_config={\"BT\": BT, \"dtype\": model.dtype},\n        )\n        x, layer = _setup_{kernel}(probe_input)\n        return layer(x)\n\n    peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)\n    sweep_config = compute_hidden_size_sweep_config(model, peak_bytes, bt=BT)\n    x_values = [1024 * i for i in range(1, 17) if 1024 * i <= sweep_config.max_hidden_size] or [model.hidden_size]\n\n    common_configs = {\n        \"kernel_name\": \"{kernel}\",\n        \"x_name\": \"hidden_size\",\n        \"x_label\": \"hidden_size\",\n        \"x_values\": x_values,\n        \"kernel_providers\": [\"liger\", \"torch\", \"torch_compile\"],\n        \"extra_benchmark_configs\": [{\"BT\": sweep_config.bt, \"dtype\": model.dtype}],\n        \"overwrite\": args.overwrite,\n    }\n\n    run_benchmarks(\n        bench_test_fn=bench_speed_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"speed\",\n        metric_unit=\"ms\",\n        **common_configs,\n    )\n    run_benchmarks(\n        bench_test_fn=bench_memory_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"memory\",\n        metric_unit=\"MB\",\n        **common_configs,\n    )\n```\n\n### Key Rules\n\n1. **Import both implementations from test file** — avoids duplicating reference code\n2. **Support \"liger\", \"torch\", and \"torch_compile\" providers**\n3. **Use `run_speed_benchmark` and `run_memory_benchmark`** from utils — don't write custom measurement\n4. **Use `get_benchmark_model_config`** and sweep utilities for standard model sizes\n5. **Use `estimate_kernel_peak_memory`** to auto-compute safe x_values range\n6. **Run both speed AND memory** for all three modes (forward, backward, full)\n7. **Use `BT = 4096`** as default batch*seq when sweeping hidden_size\n8. **Adapt sweep variable** — some kernels sweep `seq_len` or `vocab_size` instead of `hidden_size`. Use the appropriate `compute_*_sweep_config` utility.\n9. **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.\n",".agents/skills/liger-kernel-dev/templates/functional-api.md":"# Functional API Template\n\n## File: `src/liger_kernel/transformers/functional.py` (MODIFY)\n\n### Step 1: Add Import\n\nAdd at the top of the file, in alphabetical order among existing imports:\n\n```python\nfrom liger_kernel.ops import Liger{Kernel}Function\n```\n\n### Step 2: Add Function\n\nAdd in alphabetical order among existing functions:\n\n```python\ndef liger_{kernel}(x, ...params):\n    return Liger{Kernel}Function.apply(x, ...params)\n```\n\nFor functions with complex return types (like CrossEntropy), wrap the output:\n\n```python\ndef liger_{kernel}(input, target, ...params):\n    result = Liger{Kernel}Function.apply(input, target, ...params)\n    return result\n```\n\n### Key Rules\n\n- Function signature should match the PyTorch equivalent where applicable\n- Use positional args for tensor inputs, keyword args for config params\n- Keep the function body minimal — just call `.apply()` and return\n",".agents/skills/liger-kernel-dev/templates/module-wrapper.md":"# Module Wrapper Template\n\n## File: `src/liger_kernel/transformers/{kernel}.py`\n\n```python\nimport torch\nimport torch.nn as nn\n\nfrom liger_kernel.ops import Liger{Kernel}Function\n\n\nclass Liger{Kernel}(nn.Module):\n    def __init__(self, ...params):\n        super().__init__()\n        # Store config\n        self.param = param\n        # Create learnable parameters\n        self.weight = nn.Parameter(torch.ones(hidden_size))\n\n    def forward(self, x):\n        return Liger{Kernel}Function.apply(x, self.weight, ...self.params)\n\n    def extra_repr(self):\n        return f\"param={self.param}, ...\"\n```\n\n### Key Rules\n\n- Import from `liger_kernel.ops` (vendor-replaceable), not from the submodule\n- Match the PyTorch reference API as closely as possible (same __init__ params, same forward signature)\n- Include `extra_repr` for debugging\n- For loss modules, name the class `Liger{Kernel}Loss` and the forward takes `(input, target)`\n- For model-specific variants, create subclasses (see `LigerRMSNormForGemma` in `rms_norm.py`)\n",".agents/skills/liger-kernel-dev/templates/ops-kernel.md":"# Ops Kernel Template\n\n## File: `src/liger_kernel/ops/{kernel}.py`\n\nBefore writing, read the reference kernel for this tier in the actual codebase.\n\n### Import Block\n\n```python\nimport torch\nimport triton\nimport triton.language as tl\n\nfrom liger_kernel.ops.utils import calculate_settings\nfrom liger_kernel.ops.utils import ensure_contiguous\n```\n\nAdd additional imports only as needed:\n- `math` — if using `math.ceil` for row partitioning\n- `from liger_kernel.ops.utils import compare_version` — if checking triton version\n- `from liger_kernel.ops.utils import torch_to_triton_dtype` — if passing dtype as kernel arg\n- `rsqrt` from triton libdevice — if using inverse square root (see `rms_norm.py` for the import pattern with version guards)\n- `tanh` from triton libdevice — if using tanh (see `dyt.py` for the import pattern)\n\n### Forward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_forward_kernel(\n    # Output pointers first, then input pointers\n    Y_ptr, Y_stride,\n    X_ptr, X_stride,\n    # ... other input pointers and strides\n    n_cols,\n    # constexpr params for compile-time optimization\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Calculate base pointers using strides\n    X_ptr += row_idx * X_stride\n    Y_ptr += row_idx * Y_stride\n\n    # Load with mask and float32 upcast for precision\n    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0).to(tl.float32)\n\n    # === FORWARD COMPUTATION ===\n\n    tl.store(Y_ptr + col_offsets, result, mask=mask)\n```\n\n### Backward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_backward_kernel(\n    # Gradient output, then gradient inputs, then saved tensors\n    dY_ptr, dY_stride,\n    dX_ptr, dX_stride,\n    # ... saved tensors from forward\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Load gradient and saved tensors\n    # === BACKWARD COMPUTATION ===\n\n    tl.store(dX_ptr + row_idx * dX_stride + col_offsets, dX_row, mask=mask)\n```\n\n### Python Wrappers\n\n```python\ndef {kernel}_forward(X, ...):\n    ori_shape = X.shape\n    n_cols = ori_shape[-1]\n    X = X.view(-1, n_cols)\n    n_rows = X.shape[0]\n\n    Y = torch.empty_like(X)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_forward_kernel[(n_rows,)](\n        Y, Y.stride(-2),\n        X, X.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return Y.view(*ori_shape)\n\n\ndef {kernel}_backward(dY, ...saved):\n    ori_shape = dY.shape\n    n_cols = ori_shape[-1]\n    dY = dY.view(-1, n_cols)\n    n_rows = dY.shape[0]\n\n    dX = torch.empty_like(dY)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_backward_kernel[(n_rows,)](\n        dY, dY.stride(-2),\n        dX, dX.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return dX.view(*ori_shape)\n```\n\n### Autograd Function\n\n```python\nclass Liger{Kernel}Function(torch.autograd.Function):\n    @staticmethod\n    @ensure_contiguous\n    def forward(ctx, X, ...params):\n        Y = {kernel}_forward(X, ...params)\n        ctx.save_for_backward(...)  # Save only what backward needs\n        return Y\n\n    @staticmethod\n    @ensure_contiguous\n    def backward(ctx, dY):\n        ...saved = ctx.saved_tensors\n        dX = {kernel}_backward(dY, ...saved)\n        # Return one gradient per forward input, None for non-tensor params\n        return dX, None, None, ...\n```\n\n### Tier-Specific Patterns\n\n**Tier 1 (element-wise):**\n- Grid: `(n_rows,)` — one program per row\n- Simple 1D indexing with `col_offsets` and `mask`\n- Prefer recomputation in backward over saving activations\n- See `swiglu.py`: saves `a, b` but recomputes `silu(a)` in backward\n\n**Tier 2 (reduction):**\n- Forward may cache intermediate state (e.g., RSTD in RMSNorm)\n- Backward may use SM-based parallelism for weight gradient reduction:\n  ```python\n  sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count\n  _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)\n  grid = (sm_count,)\n  rows_per_program = math.ceil(n_rows / sm_count)\n  ```\n- See `rms_norm.py` for the full pattern\n\n**Tier 3 (fused/complex):**\n- Multi-pass algorithms (e.g., online softmax: pass 1 finds max+sum, pass 2 computes gradients)\n- Gradient-in-forward trick: compute gradients during forward, store in input tensor\n- Multiple `tl.constexpr` flags for compile-time dead code elimination\n- See `cross_entropy.py` for the full pattern\n\n### Key Rules\n\n1. Always use `tl.program_id(0).to(tl.int64)` (prevents overflow on large tensors)\n2. Always use `mask = col_offsets < n_cols` for boundary handling\n3. Use `tl.constexpr` for parameters that enable compile-time optimization\n4. Use stride parameters from tensors, never hardcode memory layout\n5. Cast precision-sensitive ops to `tl.float32`, cast back for storage\n6. Return `None` for each non-tensor parameter in backward\n",".agents/skills/liger-kernel-dev/templates/unit-test.md":"# Unit Test Template\n\n## File: `test/transformers/test_{kernel}.py`\n\n```python\nimport pytest\nimport torch\n\nfrom test.utils import assert_verbose_allclose\nfrom test.utils import infer_device\nfrom test.utils import set_seed\nfrom test.utils import supports_bfloat16\n\nfrom liger_kernel.ops import Liger{Kernel}Function\nfrom liger_kernel.transformers.{kernel} import Liger{Kernel}\nfrom liger_kernel.transformers.functional import liger_{kernel}\n\n\n# ---- PyTorch Reference ----\n# Paste the standalone PyTorch reference from the Analyzer stage here.\n\nclass Torch{Kernel}(torch.nn.Module):\n    def __init__(self, ...):\n        super().__init__()\n        ...\n\n    def forward(self, x):\n        ...\n\n\nset_seed(42)\ndevice = infer_device()\n\n\n# ---- Correctness Test (Module) ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (4, 16, 2048),\n        (1, 1, 1023),   # Non-power-of-2\n        (3, 7, 256),    # Small + prime numbers\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-1,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_correctness(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Initialize both with same weights\n    torch_layer = Torch{Kernel}(...).to(device).to(dtype)\n    liger_layer = Liger{Kernel}(...).to(device).to(dtype)\n    # Copy weights from torch to liger\n    # liger_layer.weight.data = torch_layer.weight.data.clone()\n\n    # Forward\n    torch_output = torch_layer(x1)\n    liger_output = liger_layer(x2)\n    assert_verbose_allclose(torch_output, liger_output, rtol=rtol, atol=atol, extra_info=\"[output]\")\n\n    # Backward\n    grad_output = torch.randn_like(_input)\n    torch_output.backward(grad_output)\n    liger_output.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol, extra_info=\"[input.grad]\")\n\n    # Check weight gradients if applicable\n    # assert_verbose_allclose(\n    #     torch_layer.weight.grad, liger_layer.weight.grad,\n    #     rtol=rtol, atol=atol, extra_info=\"[weight.grad]\"\n    # )\n\n\n# ---- Functional API Test ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (1, 1, 1023),\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-8,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_functional(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Test that functional and Function.apply produce identical results\n    output1 = liger_{kernel}(x1, ...)\n    output2 = Liger{Kernel}Function.apply(x2, ...)\n\n    assert_verbose_allclose(output1, output2, rtol=rtol, atol=atol)\n\n    grad_output = torch.randn_like(_input)\n    output1.backward(grad_output)\n    output2.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol)\n```\n\n### Key Testing Rules\n\n1. **Always test both forward AND backward** — backward bugs are common\n2. **Include non-power-of-2 shapes** — tests the mask boundary handling\n3. **Parametrize dtypes with per-dtype tolerances**:\n   - `float32`: `atol=1e-5, rtol=1e-5` (tight)\n   - `bfloat16`: `atol=1e-1, rtol=5e-2` (loose — bf16 has limited precision)\n4. **Use `set_seed(42)` at module level** for reproducibility\n5. **Use `infer_device()`** not `\"cuda\"` — supports multi-vendor\n6. **Guard bfloat16 with `supports_bfloat16()`** — skips on older GPUs\n7. **Use `assert_verbose_allclose`** not `torch.testing.assert_close` — gives better error output\n8. **Test weight gradients** when the kernel has learnable parameters\n9. **Add `@pytest.mark.flaky(reruns=3)`** if the test is sensitive to random seeds\n10. **Add extra test functions** for edge cases specific to the kernel\n",".agents/skills/liger-kernel-dev/validator.md":"# Validator Agent\n\nValidates the generated kernel through checkstyle, testing, benchmarking, and optional profiling.\n\n## Validation Flow\n\n### Step 1: Checkstyle (Fast Gate)\n\n```bash\nmake checkstyle\n```\n\nIf it fails:\n```bash\nruff check . --fix && ruff format .\nmake checkstyle\n```\n\nIf it still fails after auto-fix, identify and fix the remaining issues manually, then re-run.\n\n### Step 2: Unit Tests (Hard Gate)\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\nIf tests fail:\n1. Read the error output carefully\n2. Identify root cause: kernel logic bug, shape mismatch, dtype issue, or tolerance too tight\n3. Fix the issue in the relevant file(s)\n4. Re-run `make checkstyle` (in case fix introduced style issues)\n5. Re-run the failing test\n\n**Retry up to 3 times total.** If tests still fail after 3 attempts:\n- **STOP IMMEDIATELY. Do NOT proceed to benchmarks.**\n- Report to the user:\n  - Which tests are failing and the exact error messages\n  - What fixes were attempted\n  - The current state of the code\n  - Suggest what might be wrong and ask for guidance\n\n### Step 3: Benchmarks\n\nRun the benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nThis produces speed and memory measurements for forward, backward, and full modes, comparing Liger vs PyTorch baseline.\n\n### Step 4: Generate Plots\n\n```bash\npython benchmark/benchmarks_visualizer.py\n```\n\nPlots are saved to `benchmark/visualizations/`. If the visualizer fails or is not available, this step is non-blocking — report the raw numbers instead.\n\n### Step 5: Optional ncu Profiling\n\nOnly if the user explicitly requests it. Requires NVIDIA Nsight Compute (`ncu`) to be installed.\n\nProfile both the Liger kernel and PyTorch reference:\n```bash\nncu --set full --target-processes all -o liger_profile python -c \"\nimport torch\nfrom liger_kernel.ops.{kernel} import Liger{Kernel}Function\nx = torch.randn(..., device='cuda', dtype=torch.float32, requires_grad=True)\nout = Liger{Kernel}Function.apply(x, ...)\nout.backward(torch.randn_like(out))\n\"\n```\n\nReport key metrics: SM occupancy, memory throughput, compute throughput.\n\nIf ncu is not available, report this and skip.\n\n### Step 6: Report Results\n\nPresent a clear summary to the user:\n\n```\n## Results\n\n**Checkstyle:** PASS\n**Unit Tests:** PASS (X tests passed)\n\n**Speed Benchmarks:**\n| Mode     | Liger (ms) | PyTorch (ms) | Speedup |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...x    |\n| Backward | ...       | ...         | ...x    |\n| Full     | ...       | ...         | ...x    |\n\n**Memory Benchmarks:**\n| Mode     | Liger (MB) | PyTorch (MB) | Savings |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...%    |\n| Backward | ...       | ...         | ...%    |\n| Full     | ...       | ...         | ...%    |\n\n**Assessment:** [Which dimension improved — speed, memory, or both.\nFlag if either metric is catastrophically worse.]\n\n**Plots:** benchmark/visualizations/\n\n**ncu Profiling:** [Results if requested, or \"Not requested\"]\n```\n",".agents/skills/liger-kernel-perf/SKILL.md":"---\nname: liger-kernel-perf\ndescription: \"Optimizes the performance of existing Liger Kernel Triton kernels. Profiles kernels, diagnoses bottlenecks (memory-bound vs compute-bound), generates multiple optimization variants with benchmarking, and applies the best variant while maintaining correctness. Supports GPU architecture-specific optimization (Ampere, Hopper, Blackwell). Use when a user asks to optimize, speed up, tune, profile, or reduce memory of an existing Liger kernel.\"\n---\n\n# Liger Kernel Perf\n\nOptimizes existing Liger Kernel Triton kernels through a 3-stage pipeline: Profile, Optimize, Finalize. Supports interactive mode (human checkpoints between stages) and autonomous mode (runs end-to-end). NVIDIA GPUs only.\n\n## Mode Detection\n\n- **Interactive mode** (default): Human checkpoints between each stage\n- **Autonomous mode**: User says \"just optimize it\", \"run without asking me\", \"optimize autonomously\" → all stages run end-to-end, user sees only the final report\n\n## Input Parsing\n\nExtract from the user's request:\n\n| Field | Description | Default |\n|-------|-------------|---------|\n| `target_kernel` | Which kernel to optimize (e.g., \"rms_norm\", \"cross_entropy\") | **Required** |\n| `optimization_goal` | speed / memory / balanced | balanced |\n| `scope` | Specific pass (forward/backward), input regime, or general | general |\n| `target_gpu` | Ampere / Hopper / Blackwell / auto-detect | auto-detect |\n| `autonomy` | interactive / autonomous | interactive |\n| `max_variants` | Max optimization variants to try | 8 |\n| `target_metric` | Optional concrete target (e.g., \"forward under 0.3ms at hidden_size=4096\") | none |\n\n## Pre-Flight Validation\n\nBefore starting the pipeline, validate:\n\n1. Kernel file exists: `src/liger_kernel/ops/{kernel}.py`\n2. Benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`\n3. Test file exists: `test/transformers/test_{kernel}.py`\n4. GPU is available and CUDA works\n5. Project is installed in dev mode (`pip install -e \".[dev]\"`)\n\nIf any validation fails, report clearly and stop.\n\n## Pipeline\n\n### Stage 1: Profile\n\nFollow the **Profiler** workflow in [profiler.md](profiler.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nThis stage:\n1. Creates the workspace directory `optimization/{kernel}/`\n2. Copies the original kernel as a snapshot\n3. Runs baseline benchmarks using the existing benchmark script\n4. Detects GPU architecture (or uses user-specified target)\n5. Optionally runs NCU profiling (if `ncu` is available)\n6. Analyzes the kernel code (tier classification, patterns, optimization opportunities)\n7. Classifies the bottleneck: memory-bound vs compute-bound\n8. Produces an optimization profile with a recommended strategy order\n9. Saves profile to `optimization/{kernel}/profile.md`\n\n**Human checkpoint (interactive mode):** Present the optimization profile with bottleneck diagnosis and proposed strategy order. Confirm before proceeding.\n\n### Stage 2: Optimize\n\nFollow the **Optimizer** workflow in [optimizer.md](optimizer.md).\n\nThis stage runs an autonomous optimization loop:\n\n1. Read the optimization profile and original kernel\n2. **Always try parameter tuning first** (BLOCK_SIZE, num_warps, num_stages manual sweep -- NOT @triton.autotune)\n3. Then apply diagnosis-driven techniques from [optimization-strategies.md](optimization-strategies.md)\n4. For each variant:\n   a. Generate the variant code → `optimization/{kernel}/{kernel}_vN.py`\n   b. Write the variant lab notebook → `optimization/{kernel}/{kernel}_vN_notes.md`\n   c. Run quick smoke test (single shape, float32, forward+backward) → discard on failure\n   d. Run the **full existing benchmark script** → `optimization/{kernel}/benchmarks/vN_results.csv`\n   e. Check guardrails (no catastrophic regressions)\n   f. Update the variant notes with actual results\n5. Read all prior variant notes before generating the next variant\n6. **Stop when:** budget exhausted, 2 consecutive variants with <1% improvement, or target metric met\n7. Produce a comparison table of ALL variants\n\n**Human checkpoint (interactive mode):** Present the comparison table across all variants. User approves the winner (or skill picks best if autonomous).\n\n### Stage 3: Finalize\n\nFollow the **Finalizer** workflow in [finalizer.md](finalizer.md).\n\nThis stage:\n1. Applies the winning variant in-place to `src/liger_kernel/ops/{kernel}.py`\n2. Runs the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs` (hard gate)\n3. Runs checkstyle: `make checkstyle` (auto-fix with `ruff check . --fix && ruff format .`)\n4. Generates 3-way comparison plots (original liger vs optimized liger vs huggingface baseline) using `benchmarks_visualizer.py`\n5. Generates the final optimization report → `optimization/{kernel}/report.md`\n6. Creates a PR with only the kernel code changes (no plots or optimization workspace files)\n7. Presents the before/after summary with plots\n\n**Human checkpoint (interactive mode):** Present the final report with before/after numbers, comparison plots, and test results.\n\n## Guardrails\n\nThese apply to EVERY variant, regardless of mode:\n\n| Guardrail | Threshold | Action |\n|-----------|-----------|--------|\n| Non-target metric regression | >5% worse | Reject variant |\n| Cross-pass regression | >10% on one pass to marginally improve other | Reject variant |\n| Smoke test failure | Any correctness failure | Discard variant immediately |\n| Full test suite failure | Any | Do NOT apply winner, report failure, stop |\n| Checkstyle failure | Any | Auto-fix with ruff, retry once |\n\n## Reference Files\n\n- [profiler.md](profiler.md) -- Profiler Agent specification\n- [optimizer.md](optimizer.md) -- Optimizer Agent specification\n- [finalizer.md](finalizer.md) -- Finalizer Agent specification\n- [optimization-strategies.md](optimization-strategies.md) -- Catalog of optimization techniques\n- [templates/optimization-profile.md](templates/optimization-profile.md) -- Profiling output format (cross-stage contract)\n- [templates/variant-notes.md](templates/variant-notes.md) -- Per-variant lab notebook format\n",".agents/skills/liger-kernel-perf/finalizer.md":"# Finalizer Agent\n\nApplies the winning optimization variant to the production kernel, validates correctness through the full test suite and checkstyle, and produces the final optimization report.\n\nThis is Stage 3 of the 3-stage pipeline. It receives the winning variant from Stage 2 (Optimizer) and is responsible for ensuring the optimization is safe to ship.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `winning_variant`: variant number chosen by user (interactive) or by the Optimizer (autonomous)\n- `optimization_workspace`: path to `optimization/{kernel}/` containing all variant files, benchmarks, and the profile\n\n## Expected Workspace State\n\nBefore the Finalizer starts, the following files MUST exist (created by earlier stages):\n\n```\noptimization/{kernel}/\n  original_{kernel}.py          # Snapshot of the original kernel (from Profiler)\n  profile.md                    # Optimization profile (from Profiler)\n  {kernel}_v1.py                # Variant source files (from Optimizer)\n  {kernel}_v1_notes.md          # Variant lab notebooks (from Optimizer)\n  {kernel}_v2.py                # ... more variants\n  {kernel}_v2_notes.md\n  benchmarks/\n    v0_baseline.csv             # Baseline benchmark data\n    v1_results.csv              # Per-variant benchmark data\n    v2_results.csv\n```\n\nIf any required files are missing, report the missing files and stop.\n\n## Steps\n\n### Step 1: Validate Inputs and Read Context\n\n1. Confirm the winning variant file exists: `optimization/{kernel}/{kernel}_v{N}.py`\n2. Read the winning variant code\n3. Read the original kernel: `optimization/{kernel}/original_{kernel}.py`\n4. Read the optimization profile: `optimization/{kernel}/profile.md`\n5. Read ALL variant notes files (`{kernel}_v*_notes.md`) to collect data for the report\n6. Read ALL benchmark CSVs (`benchmarks/v*_results.csv`) to build the comparison table\n7. Read the current production kernel: `src/liger_kernel/ops/{kernel}.py`\n\nVerify the original kernel snapshot matches the current production kernel. If they differ (someone else modified the file during the optimization), warn the user and confirm before proceeding.\n\n### Step 2: Apply the Winning Variant\n\nReplace the production kernel with the winning variant:\n\n```bash\ncp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis is a destructive operation. The original is preserved at `optimization/{kernel}/original_{kernel}.py`.\n\n### Step 3: Run the Full Test Suite (HARD GATE)\n\nRun the complete test file for the kernel:\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\n**This is a HARD GATE.** The test suite must pass completely.\n\n#### If tests PASS:\nProceed to Step 4.\n\n#### If tests FAIL:\nExecute the test failure recovery procedure:\n\n**Attempt 1: Diagnose and fix**\n1. Read the full test error output carefully\n2. Identify the root cause (numerical tolerance issue, shape mismatch, missing import, API change, etc.)\n3. Apply a targeted fix to `src/liger_kernel/ops/{kernel}.py`\n4. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**Attempt 2: Second fix attempt (if Attempt 1 still fails)**\n1. Read the new error output\n2. Try a different fix approach\n3. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**If both attempts fail: REVERT**\n1. Restore the original kernel:\n   ```bash\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n2. Verify the revert by running the test suite again:\n   ```bash\n   python -m pytest test/transformers/test_{kernel}.py -xvs\n   ```\n3. Set `optimization_applied = false`\n4. Record the failure reason and all error outputs\n5. Proceed to Step 4 to generate the report (the report documents what was attempted even if it could not be applied)\n\n**Do NOT proceed to finalize (Step 5) if the optimization was reverted. The report must clearly state the optimization was not applied.**\n\n### Step 4: Run Checkstyle\n\nRun the project's checkstyle validation:\n\n```bash\nmake checkstyle\n```\n\n#### If checkstyle PASSES:\nProceed to Step 5.\n\n#### If checkstyle FAILS:\nAuto-fix with ruff:\n\n```bash\nruff check . --fix && ruff format .\n```\n\nRetry checkstyle once:\n\n```bash\nmake checkstyle\n```\n\nIf it still fails after auto-fix, report the remaining checkstyle issues but do NOT revert the kernel (checkstyle is not a correctness gate). Note the issues in the report so they can be fixed manually.\n\n### Step 5: Generate Comparison Plots\n\nGenerate 3-way comparison plots showing the original Liger kernel, the optimized kernel, and the HuggingFace/PyTorch baseline. This gives the user a visual demonstration of the improvement.\n\nThe procedure has 3 phases: (1) benchmark both kernels for plots, (2) generate plots, (3) clean the CSV so only the optimized kernel's data remains permanently.\n\n#### Step 5a: Benchmark Both Kernels for Plotting\n\n1. **Benchmark the original kernel as \"liger_original\"**: Temporarily restore the original kernel and run benchmarks with a custom provider name:\n   ```bash\n   # Swap in original kernel\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Edit benchmark script: replace \"liger\" with \"liger_original\" in kernel_providers\n   sed -i 's/\"liger\"/\"liger_original\"/g' benchmark/scripts/benchmark_{kernel}.py\n   \n   # Run benchmarks (appends \"liger_original\" rows to CSV)\n   cd benchmark/scripts && python benchmark_{kernel}.py\n   \n   # Revert the benchmark script change\n   git checkout benchmark/scripts/benchmark_{kernel}.py\n   ```\n\n2. **Benchmark the optimized kernel as \"liger\"**: Re-apply the optimized kernel and run benchmarks normally. Use `--overwrite` so the \"liger\" rows reflect the optimized kernel:\n   ```bash\n   # Swap in optimized kernel\n   cp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Run benchmarks (writes \"liger\" rows with optimized kernel)\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   ```\n\nAt this point the CSV has 3 providers: `\"liger_original\"`, `\"liger\"` (optimized), and `\"huggingface\"`/`\"torch\"`. This is what we need for the 3-way comparison plots.\n\n#### Step 5b: Generate Plots\n\nGenerate plots while all 3 providers are in the CSV:\n\n```bash\n# Speed plots for all modes (forward, backward, full)\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name speed \\\n  --overwrite\n\n# Memory plot\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name memory \\\n  --overwrite\n```\n\nCopy plots to the optimization workspace:\n\n```bash\ncp benchmark/visualizations/{kernel}_speed_*.png optimization/{kernel}/\ncp benchmark/visualizations/{kernel}_memory_*.png optimization/{kernel}/\n```\n\n#### Step 5c: Clean CSV for Permanent State\n\nNow remove the temporary `\"liger_original\"` rows and ensure the CSV only contains clean data for the optimized kernel. Delete ALL rows for this kernel, then re-run the benchmark to append fresh data:\n\n```bash\n# Delete all existing rows for this kernel from CSV (keep header + other kernels)\nhead -1 benchmark/data/all_benchmark_data.csv > /tmp/benchmark_header.csv\ngrep -v \"^{kernel},\" benchmark/data/all_benchmark_data.csv > /tmp/benchmark_other.csv || true\ncat /tmp/benchmark_header.csv <(tail -n +2 /tmp/benchmark_other.csv) > benchmark/data/all_benchmark_data.csv\n\n# Re-run benchmarks with optimized kernel (appends fresh rows)\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nAfter this, `all_benchmark_data.csv` contains:\n- `\"liger\"` — the **optimized** kernel (clean fresh data)\n- `\"huggingface\"` / `\"torch\"` — the baseline (also fresh)\n- NO `\"liger_original\"` — that was only needed for the plots\n\n#### Step 5c: Generate Additional Custom Plots (Optional)\n\nIf the optimization shows interesting size-dependent behavior (e.g., bigger gains at larger sizes), generate targeted plots highlighting this. Use matplotlib directly if the standard visualizer does not cover the comparison you want to show.\n\nSave any additional plots to `optimization/{kernel}/`.\n\n### Step 6: Generate the Final Optimization Report\n\nWrite the report to `optimization/{kernel}/report.md`. The report is the permanent record of the entire optimization effort.\n\nThe report must include these sections:\n\n1. **Summary**: kernel name, GPU, goal, variants tried, best variant, improvement %, and whether the optimization was applied (`Yes` / `No — reverted due to test failures`)\n2. **Baseline Profile**: bottleneck classification with evidence, NCU highlights if available\n3. **Variant Comparison**: include the table from `optimization/{kernel}/comparison.md` (produced by the Optimizer)\n4. **Winning Changes**: what was changed, why it worked, and a key code diff. Run `diff optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py` — do NOT write the diff from memory.\n5. **Rejected Approaches**: for each rejected variant, what was tried and why it failed\n6. **Comparison Plots**: reference the plots generated in Step 5 with file paths\n7. **Correctness Verification**: smoke test results, full test suite results, checkstyle results\n8. **Recommendations**: unexplored optimization opportunities, suggested follow-ups\n9. **Reproduction** section:\n   ```bash\n   # Reproduce baseline benchmarks:\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   # Revert to original:\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n\nInclude the variant comparison data from `optimization/{kernel}/comparison.md` (produced by the Optimizer).\n\n### Report Quality Checklist\n\nBefore saving the report, verify:\n\n- [ ] All variants are listed in the comparison table (including rejected ones)\n- [ ] Speed ratios are calculated consistently (higher = faster)\n- [ ] Memory ratios are calculated consistently (higher = less memory used)\n- [ ] The winning variant is clearly marked\n- [ ] The bottleneck classification matches the profile from Stage 1\n- [ ] If optimization was reverted, this is stated in the Summary, the header, and Section 6\n- [ ] Reproduction commands use the correct file paths\n- [ ] The code diff in Section 4 is accurate (diff the actual files, do not write from memory)\n- [ ] Comparison plots are generated and referenced in the report\n\n### Step 7: Create Pull Request\n\nCreate a PR with the kernel code changes and the updated benchmark data. Do NOT include plots or optimization workspace files in the PR.\n\n#### Step 7a: Stage Only Relevant Files\n\n```bash\ngit add src/liger_kernel/ops/{kernel}.py\ngit add benchmark/data/all_benchmark_data.csv\n```\n\nDo NOT add:\n- `optimization/` directory (workspace files, notes, plots)\n- `benchmark/visualizations/` (generated plots -- these are for local review only)\n\n#### Step 7b: Create a Descriptive Commit\n\n```bash\ngit commit -m \"$(cat <<'EOF'\n[Perf] Optimize {kernel} kernel: {headline improvement}\n\n{2-3 sentence summary of the optimization technique and its impact.}\n\nKey changes:\n- {Change 1}: {what and why}\n- {Change 2}: {what and why}\n\nBenchmark results ({gpu_name}):\n- Speed (forward): {delta}% {faster/slower}\n- Speed (backward): {delta}% {faster/slower}\n- Speed (full): {delta}% {faster/slower}\n- Memory: {delta}% {reduction/increase}\nEOF\n)\"\n```\n\nIf the host AI runtime has its own commit-attribution convention (e.g., a `Co-Authored-By` trailer), append it per that runtime's guidelines.\n\n#### Step 7c: Push and Create PR\n\n```bash\ngit push -u origin HEAD\n```\n\nCreate the PR using `gh pr create`. The PR body must clearly explain:\n\n1. **What was optimized** -- which kernel, which pass, what bottleneck\n2. **What changes were made** -- specific code changes with rationale for each\n3. **How they help** -- tie back to the bottleneck diagnosis (e.g., \"reduced register pressure from 115 to 95 registers/thread, increasing occupancy from 12.5% to 25%\")\n4. **Benchmark results** -- before/after table with speed and memory numbers at representative sizes\n5. **Correctness** -- confirm all tests pass\n\nUse this PR template:\n\n```bash\ngh pr create --title \"[Perf] Optimize {kernel}: {headline}\" --body \"$(cat <<'EOF'\n## Summary\n\n{2-3 sentences: what was optimized, the key technique, and the headline result.}\n\n## Bottleneck Diagnosis\n\n- **Kernel**: `{kernel}` ({tier} — {tier_description})\n- **GPU**: {gpu_name} ({architecture})\n- **Classification**: {memory-bound / compute-bound / balanced}\n- **Root cause**: {specific bottleneck explanation}\n\n## Changes\n\n{For each change, explain WHAT was changed and WHY it helps:}\n\n### 1. {Change title}\n{Description of the code change and its performance rationale.}\n\n### 2. {Change title}\n{Description of the code change and its performance rationale.}\n\n## Benchmark Results\n\nTested on {gpu_name}. Values are median ms (speed) or MB (memory).\n\n### Speed (ms, lower is better)\n\n| x_value | Mode | Before | After | Improvement |\n|---------|------|--------|-------|-------------|\n| {x1} | forward | {before} | {after} | {delta}% |\n| {x1} | backward | {before} | {after} | {delta}% |\n| {x1} | full | {before} | {after} | {delta}% |\n| ... | ... | ... | ... | ... |\n\n### Memory (MB, lower is better)\n\n| x_value | Before | After | Improvement |\n|---------|--------|-------|-------------|\n| {x1} | {before} | {after} | {delta}% |\n| ... | ... | ... | ... |\n\n## Correctness\n\n- Full test suite: **PASSED** (`python -m pytest test/transformers/test_{kernel}.py -xvs`)\n- Checkstyle: **PASSED** (`make checkstyle`)\n\n## Test Plan\n\n- [x] All existing unit tests pass\n- [x] Benchmarks show improvement across all input sizes\n- [x] No regression on non-target metrics (speed/memory balance maintained)\nEOF\n)\"\n```\n\nIf the host AI runtime appends a \"Generated with\" footer (e.g., Claude Code, Cursor, Copilot), include it per that runtime's guidelines.\n\n**Important**: Do NOT include plots as image attachments in the PR. Plots are for local review only and live in the optimization workspace.\n\n### Step 8: Present the Before/After Summary\n\nInclude the plot file paths in the summary so the user can view them.\n\nDisplay a concise summary to the user:\n\n```\n=== Optimization Complete: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: {APPLIED / REVERTED}\nWinner: v{N} — {strategy}\n\nBefore → After (at largest x_value):\n  Speed (forward):  {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (backward): {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (full):     {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Memory:           {baseline_mb} MB → {optimized_mb} MB ({improvement}% reduction)\n\nVariants tried: {N}\nTest suite: PASSED ({N} tests)\nCheckstyle: PASSED\n\nReport: optimization/{kernel}/report.md\n```\n\nIf the optimization was reverted:\n\n```\n=== Optimization REVERTED: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: REVERTED — test suite failed after applying winner\n\nWinner candidate: v{N} — {strategy}\nFailure reason: {brief error description}\nFix attempts: 2 (both failed)\n\nThe original kernel has been restored. No changes were made to production code.\n\nVariants tried: {N} (see report for full comparison)\nReport: optimization/{kernel}/report.md\n```\n\n## Error Handling\n\n| Error | Action |\n|-------|--------|\n| Winning variant file not found | Report missing file, stop |\n| Original kernel snapshot missing | Report missing file, stop (cannot safely revert) |\n| Production kernel differs from snapshot | Warn user, ask for confirmation (interactive) or stop (autonomous) |\n| Test suite fails (after 2 fix attempts) | Revert to original, generate report documenting failure |\n| Checkstyle fails (after auto-fix) | Note in report, do NOT revert (not a correctness issue) |\n| Benchmark CSVs missing | Generate report with available data, note missing data |\n| Cannot write report | Print report contents to stdout as fallback |\n\n## Key Principles\n\n1. **Correctness is non-negotiable.** The full test suite is a hard gate. Never skip it, never ignore failures, never mark a failing optimization as successful.\n2. **Always preserve the original.** The snapshot at `optimization/{kernel}/original_{kernel}.py` is the safety net. Never delete or overwrite it.\n3. **The report is permanent.** Even if the optimization is reverted, the report documents everything that was tried. This prevents future engineers from repeating failed approaches.\n4. **Be honest in the report.** Report actual measured numbers. Do not round favorably. If a variant was only 0.3% faster, say 0.3%, not \"marginal improvement.\" If the optimization was reverted, say so prominently.\n5. **Diff, don't recall.** When writing the code diff in the report, actually diff the files. Do not write the diff from memory -- it will contain errors.\n",".agents/skills/liger-kernel-perf/optimization-strategies.md":"# Optimization Strategies Catalog\n\nReference catalog for the Optimizer Agent. Techniques are ordered by recommended trial order within each category. The agent should **always start with parameter tuning**, then apply diagnosis-driven techniques based on the bottleneck classification from the Profiler stage.\n\n---\n\n## Decision Framework\n\nAfter reading the optimization profile, follow this priority:\n\n1. **Always first**: Parameter Tuning (Section 1)\n2. **Memory-bound kernel**: Section 2 techniques, then Section 5\n3. **Compute-bound kernel**: Section 3 techniques, then Section 5\n4. **Balanced kernel**: Interleave Section 2 and Section 3, then Section 5\n5. **Architecture-specific**: Layer Section 4 techniques on top of the best variant when targeting a specific GPU\n\n---\n\n## 1. Always First: Parameter Tuning\n\n**Why first**: Parameter tuning is the cheapest technique (minimal code changes) and frequently delivers the largest single improvement. The current codebase uses `calculate_settings()` which picks BLOCK_SIZE as the next power of 2 and num_warps via a fixed heuristic. This is a reasonable default but leaves significant performance on the table for specific input shapes and GPU architectures.\n\n### Why NOT @triton.autotune\n\nLiger production kernels **cannot use `@triton.autotune`** for three architectural reasons:\n\n1. **Forward-backward context coupling**: Liger kernels use `torch.autograd.Function` where `forward()` computes `BLOCK_SIZE`/`num_warps` via `calculate_settings()`, stores them in `ctx`, and `backward()` retrieves them to launch its kernel with identical config. With `@triton.autotune`, these become compile-time constants managed by Triton -- they cannot be stored in `ctx`, and the backward kernel has no way to know which config the forward kernel auto-selected for a given input.\n\n2. **Multi-backend incompatibility**: NPU/Ascend does not support `num_warps` or `num_stages` at all (commit `8b965b9` explicitly removed them). `@triton.autotune` relies on these parameters.\n\n3. **Inference warmup latency**: Each new input shape triggers autotuning (benchmarks all configs). Adds 100ms+ first-call penalty per unique shape, unacceptable for LLM inference with variable sequence lengths.\n\nThe commented-out `@triton.autotune` decorators in `dyt.py` and `grpo_loss.py` are evidence of prior exploration that was abandoned for these reasons. Do NOT re-introduce `@triton.autotune` on production kernels.\n\n### Parameters to Tune\n\n| Parameter | Description | Typical Range |\n|-----------|-------------|---------------|\n| `BLOCK_SIZE` | Elements processed per program instance | 256, 512, 1024, 2048, 4096, 8192, 16384 |\n| `num_warps` | Warps per program instance (32 threads each) | 2, 4, 8, 16, 32 |\n| `num_stages` | Software pipelining depth | 1, 2, 3, 4, 5 |\n| `BLOCK_ROW` | Rows per program (block-row mode only) | 1, 2, 4, 8, 16, 32 |\n\n### Method: Manual Sweep with Improved Heuristic\n\nReplace `calculate_settings()` with a better heuristic derived from benchmarking. The variant must still pass `BLOCK_SIZE` and `num_warps` through `ctx` to the backward kernel.\n\n**Step 1**: Identify the current config for the benchmark's input sizes:\n```python\n# calculate_settings() returns (next_power_of_2(n), num_warps_heuristic)\n# For n_cols=4096: BLOCK_SIZE=4096, num_warps=8\n# For n_cols=8192: BLOCK_SIZE=8192, num_warps=16\n# Note: num_stages is never set (Triton defaults to 1)\n```\n\n**Step 2**: Create a variant with an improved selection function:\n```python\ndef optimized_settings(n):\n    \"\"\"Improved heuristic derived from benchmarking on {gpu_arch}.\"\"\"\n    BLOCK_SIZE = triton.next_power_of_2(n)\n    if BLOCK_SIZE > 65536:\n        raise RuntimeError(...)\n    # Tuned values (from benchmark sweep):\n    if BLOCK_SIZE >= 8192:\n        num_warps = 16\n        num_stages = 3   # NEW: calculate_settings doesn't set this\n    elif BLOCK_SIZE >= 2048:\n        num_warps = 8\n        num_stages = 2\n    else:\n        num_warps = 4\n        num_stages = 2\n    return BLOCK_SIZE, num_warps, num_stages\n```\n\n**Step 3**: Update the kernel launcher to pass `num_stages`:\n```python\n# The kernel launch gains num_stages (previously unset, defaulted to 1):\nkernel[(grid,)](..., BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages)\n```\n\n**Step 4**: Update `ctx` to also store `num_stages` and pass it to the backward kernel.\n\n**Key insight**: `num_stages` is the most commonly overlooked parameter. `calculate_settings()` does not set it at all, so Triton defaults to `num_stages=1`. Simply adding `num_stages=2` or `num_stages=3` to existing kernel launches is often the single biggest win, especially on Hopper.\n\n### Architecture-Specific Guidance\n\n- **Ampere (SM 8.x)**: `num_stages=2` is usually optimal. Higher stages rarely help and can increase register pressure. `num_warps` sweet spot is 4-16.\n- **Hopper (SM 9.0)**: Benefits significantly from higher `num_stages` (3-5) due to TMA and async pipeline support. Also tolerates `num_warps=32` better due to deeper warp scheduling.\n- **Blackwell (SM 10.0)**: Similar to Hopper but with even more pipeline depth tolerance. Start with `num_stages=4-5`.\n\n### When to Apply\n\nAlways. This is the mandatory first optimization attempt for every kernel.\n\n### Expected Impact\n\n5-40% speedup depending on how far the `calculate_settings()` heuristic is from optimal for the target shapes and GPU. Adding `num_stages` alone can yield 10-25% on Hopper.\n\n### Risks/Gotchas\n\n- Configs are shape-dependent. A config optimal for `n_cols=4096` may be suboptimal for `n_cols=128`. Always benchmark across ALL x_values from the existing benchmark suite.\n- If BLOCK_SIZE exceeds `n_cols`, the kernel still works (masked loads) but wastes warps. Prune configs where `BLOCK_SIZE >> n_cols`.\n- The `calculate_settings()` function enforces `MAX_FUSED_SIZE = 65536`. This limit exists to prevent register spilling. Respect this unless verified otherwise.\n- When adding `num_stages` to existing kernels, the backward kernel must also receive the same `num_stages` value via `ctx`.\n- On NPU/Ascend, `num_warps` and `num_stages` are ignored. Optimizations targeting these parameters are NVIDIA-specific.\n\n### Kernel Tiers\n\nAll tiers benefit. Tier 1 kernels often see the largest relative gains because their simple structure makes the config choice proportionally more impactful.\n\n---\n\n## 2. Memory-Bound Optimizations\n\nApply these when the profiler classifies the kernel as memory-bound (high memory throughput, low compute utilization, simple operations per load/store).\n\n### 2.1 Memory Access Coalescing\n\n**When to apply**: NCU shows low memory throughput efficiency, or kernel uses non-contiguous access patterns (e.g., strided access across rows, column-major access on row-major data).\n\n**How to implement**:\n\nEnsure that threads in a warp access adjacent memory addresses. In Triton, this means the innermost offset computation should produce contiguous addresses across the `tl.arange` dimension.\n\n```python\n# GOOD: contiguous access within a warp\ncol_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_idx * stride + col_offsets, mask=col_offsets < n_cols)\n\n# BAD: strided access (each thread hits a different cache line)\nrow_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_offsets * stride + col_idx)\n```\n\nIn Liger kernels, most access is already contiguous via `col_offsets = tl.arange(0, BLOCK_SIZE)`. Check for cases where stride-based pointer arithmetic breaks coalescing, especially in backward kernels or multi-dimensional grids.\n\n**Expected impact**: 10-50% when fixing uncoalesced access. Minimal gain if access is already coalesced.\n\n**Risks**: Changing access patterns may require transposing data or restructuring the grid, which can be invasive.\n\n**Kernel tiers**: All tiers, but Tier 2/3 backward passes most often have coalescing issues due to accumulation patterns.\n\n### 2.2 Cache Modifiers\n\n**When to apply**: Kernel reads data that is used exactly once (stream through), or reads data that is reused many times (should be cached aggressively). NCU shows poor L2 hit rate.\n\n**How to implement**:\n\nTriton supports cache modifiers via the `eviction_policy` parameter on `tl.load` and `tl.store`:\n\n```python\n# Data read once, don't pollute cache (streaming access)\nX_block = tl.load(X_ptr + offsets, eviction_policy=\"evict_first\")\n\n# Data reused across iterations, keep in cache\nW_row = tl.load(W_ptr + offsets, eviction_policy=\"evict_last\")\n\n# Default (let hardware decide)\nX_block = tl.load(X_ptr + offsets)  # no modifier\n```\n\n| Modifier | Triton API | When to Use |\n|----------|------------|-------------|\n| `.ca` (cache all) | `eviction_policy=\"evict_last\"` | Data reused across loop iterations or by other programs |\n| `.cs` (cache streaming) | `eviction_policy=\"evict_first\"` | Data used once; large tensors that would thrash cache |\n| `.wb` (write-back) | default store behavior | Standard write-back through cache hierarchy |\n| `.wt` (write-through) | not directly exposed | Write-through (rarely needed in Triton) |\n\n**Expected impact**: 5-20% on memory-bound kernels with predictable access patterns.\n\n**Risks**: Wrong eviction policy can degrade performance. `.evict_first` on reused data will cause redundant memory fetches. Always benchmark both policies against the default.\n\n**Kernel tiers**: Tier 2 and Tier 3 benefit most (weight reuse in backward passes, multi-pass algorithms).\n\n### 2.3 Reduce Memory Traffic\n\n**When to apply**: Kernel reads or writes data that could be computed on the fly, or writes intermediate results that are only consumed once.\n\n**How to implement**:\n\nLiger already uses several memory-saving patterns. Look for additional opportunities:\n\n```python\n# Pattern 1: Recompute instead of saving to memory\n# Instead of: save silu(a) to a buffer, load in backward\n# Do: recompute silu(a) in backward from saved a\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a  # recomputed in backward\n\n# Pattern 2: In-place gradient storage (gradient-in-forward trick)\n# Cross entropy already does this: stores gradient in the input tensor\n# Look for other kernels that allocate separate grad tensors unnecessarily\n\n# Pattern 3: Fuse producer-consumer pairs to avoid materializing intermediates\n# Instead of: Y = kernel1(X); Z = kernel2(Y)  (Y written to HBM, then read)\n# Do: Z = fused_kernel(X)  (Y lives only in registers/SRAM)\n\n# Pattern 4: Write in-place to dY for dX when shapes match\nif in_place:\n    dX = dY  # reuse memory, zero additional allocation\n```\n\n**Expected impact**: 10-40% from eliminating unnecessary memory round-trips. Memory savings also reduce peak memory.\n\n**Risks**: Recomputation trades memory for compute. Only beneficial when the kernel is memory-bound. Can also complicate backward pass correctness -- always verify gradient accuracy.\n\n**Kernel tiers**: All tiers. Tier 3 kernels (fused operations) have the most opportunity for intermediate elimination.\n\n### 2.4 Data Reuse (Pre-Load Outside Loops)\n\n**When to apply**: A value is loaded inside a loop but does not change across iterations. Common with weight tensors and scaling factors.\n\n**How to implement**:\n\n```python\n# BEFORE: weight loaded inside the loop (redundant loads if compiler doesn't hoist)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    W_block = tl.load(W_ptr + offsets)  # loaded every iteration if W is large\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_block\n\n# AFTER: hoist the invariant load\nW_row = tl.load(W_ptr + col_offsets, mask=col_offsets < n_cols)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_row  # reuse from registers\n```\n\nNote: Triton's compiler often hoists loads automatically, but it is not guaranteed for all patterns, especially when the loop body is complex.\n\n**Expected impact**: 5-15% when applicable.\n\n**Risks**: Hoisting large loads outside loops increases register pressure. If the weight vector is very large relative to BLOCK_SIZE, it may cause register spilling that offsets the benefit.\n\n**Kernel tiers**: Tier 2 (norm kernels with weight vectors), Tier 3 (cross-entropy with weight tensors).\n\n### 2.5 Block-Row Mode (Multiple Rows per Program)\n\n**When to apply**: The kernel processes one row per program instance, but each row is small (e.g., `n_cols <= 256`). This under-utilizes each SM because the program does very little work.\n\n**How to implement**:\n\nLiger's `rms_norm.py` already implements this as `_block_rms_norm_forward_kernel`. The pattern:\n\n```python\n@triton.jit\ndef kernel_block_row(\n    ...,\n    n_rows,\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n    BLOCK_ROW: tl.constexpr,  # number of rows per program\n):\n    # Each program processes BLOCK_ROW rows\n    row_idx = tl.program_id(0) * BLOCK_ROW + tl.arange(0, BLOCK_ROW)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    row_mask = row_idx < n_rows\n    col_mask = col_offsets < n_cols\n\n    # 2D load: [BLOCK_ROW, BLOCK_SIZE]\n    X = tl.load(\n        X_ptr + row_idx[:, None] * stride + col_offsets[None, :],\n        mask=row_mask[:, None] & col_mask[None, :],\n        other=0,\n    )\n    # Process all rows at once...\n```\n\nGrid: `(triton.cdiv(n_rows, BLOCK_ROW),)` instead of `(n_rows,)`.\n\nUse conditional dispatch to choose between row mode and block-row mode:\n\n```python\nif BLOCK_SIZE > 256 or n_rows < 4096 * 8:\n    # Row mode: one row per program\n    kernel_row[(n_rows,)](...)\nelse:\n    # Block-row mode: multiple rows per program\n    BLOCK_ROW = 16\n    kernel_block_row[(triton.cdiv(n_rows, BLOCK_ROW),)](\n        ..., BLOCK_ROW=BLOCK_ROW\n    )\n```\n\n**Expected impact**: 10-30% for kernels with small `n_cols` and large `n_rows`.\n\n**Risks**: Increases register pressure and shared memory usage. The 2D load pattern is less efficient when `BLOCK_ROW` causes register spilling. Also adds code complexity (two kernel variants). Does not help when `n_cols` is already large.\n\n**Kernel tiers**: Tier 1 and Tier 2 primarily. Many Liger Tier 1 kernels (swiglu, geglu, relu_squared) do not yet have block-row mode.\n\n### 2.6 Vectorized Loads\n\n**When to apply**: Kernel loads/stores many small elements individually. Memory bandwidth is limited by the number of transactions rather than total bytes.\n\n**How to implement**:\n\nTriton generally handles vectorization automatically when loads are contiguous and aligned. However, you can help by:\n\n1. Ensuring BLOCK_SIZE is a multiple of 8 (for fp16/bf16) or 4 (for fp32) to enable 128-bit loads.\n2. Ensuring pointer alignment (tensors created by PyTorch are typically aligned).\n3. Avoiding masked loads at boundaries where possible -- pad the input instead.\n\n```python\n# Triton auto-vectorizes contiguous loads. Just ensure alignment:\n# BLOCK_SIZE should be power of 2 (already enforced by calculate_settings)\n# Data should be contiguous (already enforced by ensure_contiguous decorator)\n\n# For explicit vectorization (rare, advanced):\n# Cast pointer to wider type, load fewer elements\n# This is rarely needed as Triton handles it, but can help in edge cases\n```\n\n**Expected impact**: 5-10% if loads are currently unvectorized. Usually Triton already handles this.\n\n**Risks**: Minimal. Padding inputs wastes a small amount of memory.\n\n**Kernel tiers**: All tiers, but gains are usually small since Triton auto-vectorizes well.\n\n---\n\n## 3. Compute-Bound Optimizations\n\nApply these when the profiler classifies the kernel as compute-bound (high compute utilization, low memory throughput utilization, complex math per element).\n\n### 3.1 Algorithmic Improvements\n\n**When to apply**: The kernel uses a naive algorithm where a more efficient one exists. Examples: naive softmax vs online softmax, explicit matrix inverse vs iterative method.\n\n**How to implement**:\n\nLiger already uses online softmax in cross_entropy. Look for similar opportunities:\n\n```python\n# Online algorithm pattern: single-pass with running statistics\n# Instead of: pass 1 = compute max, pass 2 = compute sum, pass 3 = compute output\n# Do: single pass maintaining running max and sum\n\nm = float(\"-inf\")\nd = 0.0\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    block_max = tl.max(X_block)\n    m_new = tl.maximum(m, block_max)\n    d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new))\n    m = m_new\n# Result: lse = m + tl.log(d)\n```\n\nOther algorithmic opportunities:\n- Welford's algorithm for variance (numerically stable single-pass)\n- Kahan summation for high-precision reduction\n- Fast approximations for transcendentals when full precision is not needed\n\n**Expected impact**: 20-60% when applicable (reducing O(n) passes).\n\n**Risks**: Numerical stability can differ from reference implementation. Always validate against PyTorch reference with tight tolerances. Online algorithms may accumulate floating-point errors differently.\n\n**Kernel tiers**: Tier 2 and Tier 3 primarily. Tier 1 kernels rarely have algorithmic improvement opportunities.\n\n### 3.2 Reduce Redundant Computation\n\n**When to apply**: The same sub-expression is computed multiple times within a kernel, or across forward and backward passes.\n\n**How to implement**:\n\n```python\n# BEFORE: sigmoid computed twice\nsilu_a = a_row * tl.sigmoid(a_row)\ngrad = dc_row * (silu_a * (1 - tl.sigmoid(a_row)) + tl.sigmoid(a_row)) * b_row\n\n# AFTER: compute once, reuse\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a\ngrad = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row\n```\n\nLiger already does this well (see swiglu backward). Look for:\n- Repeated `tl.exp()`, `tl.sigmoid()`, `tl.log()` calls with the same argument\n- Reduction results (sum, max) computed multiple times\n- Softmax probabilities recomputed instead of cached in registers\n\n**Expected impact**: 5-20% depending on the cost of the duplicated operation.\n\n**Risks**: Caching sub-expressions increases register pressure. Trade-off with register spilling.\n\n**Kernel tiers**: Tier 2 and Tier 3. Tier 1 kernels are typically simple enough that redundancy is already eliminated.\n\n### 3.3 Tensor Core Utilization\n\n**When to apply**: Kernel performs matrix-multiply-like operations (dot products, outer products) that could leverage tensor cores. This is most relevant for Tier 3 kernels with linear layers or attention-like patterns.\n\n**How to implement**:\n\nTensor cores require specific alignment:\n- Block dimensions must be multiples of 16 (for fp16/bf16) or 8 (for tf32)\n- Input/output matrices must be in supported layouts\n- Use `tl.dot()` which maps to tensor core instructions when dimensions align\n\n```python\n# Ensure BLOCK_M, BLOCK_N, BLOCK_K are multiples of 16\nBLOCK_M: tl.constexpr = 64  # multiple of 16\nBLOCK_N: tl.constexpr = 64  # multiple of 16\nBLOCK_K: tl.constexpr = 32  # multiple of 16\n\n# tl.dot maps to tensor cores automatically\nacc = tl.dot(A_block, B_block)  # [BLOCK_M, BLOCK_K] x [BLOCK_K, BLOCK_N]\n```\n\nFor `tl.dot` to use tensor cores:\n- Both operands must be 2D\n- Operand dtypes must be fp16, bf16, or tf32\n- Inner dimension K must be >= 16\n\n**Expected impact**: 2-10x for matrix-multiply-heavy kernels. Not applicable to element-wise or simple reduction kernels.\n\n**Risks**: Tensor core operations have lower precision than FP32 scalar operations. The `tf32` mode on Ampere+ truncates FP32 mantissa to 10 bits. Always validate numerical accuracy. Not all kernels have operations that map to tensor cores.\n\n**Kernel tiers**: Tier 3 only (fused_linear_cross_entropy, attention-like kernels). Not applicable to Tier 1 or most Tier 2 kernels.\n\n### 3.4 Register Pressure Reduction\n\n**When to apply**: NCU shows high register usage per thread (>128 registers) or occupancy is limited by register count. Kernel body has many live variables.\n\n**How to implement**:\n\n```python\n# Strategy 1: Use constexpr branches to eliminate dead code paths\n# Triton compiles out branches on constexpr values, reducing register usage\ncasting_mode: tl.constexpr  # compiler eliminates unused casting branches\nHAS_WEIGHT: tl.constexpr    # compiler eliminates weight-related registers when False\n\n# Strategy 2: Split large kernels into smaller passes\n# Instead of one kernel doing everything:\n# Kernel 1: forward + save intermediates\n# Kernel 2: backward using intermediates\n# This is a last resort as it adds kernel launch overhead\n\n# Strategy 3: Reduce live range of temporary variables\n# Reorder operations so temporary values are consumed immediately\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    # Use X_block immediately, don't hold it across iterations\n    result = compute(X_block)\n    tl.store(..., result)\n    # X_block and result are now dead, registers freed\n\n# Strategy 4: Use smaller data types for intermediates where precision allows\n# fp16 uses half the registers of fp32\nintermediate = X_block.to(tl.float16)  # only if precision loss is acceptable\n```\n\n**Expected impact**: 10-25% when occupancy is register-limited. Minimal impact when occupancy is already high.\n\n**Risks**: Splitting kernels adds launch overhead and memory traffic for intermediates. Reducing precision can cause numerical issues. Constexpr branches are already used extensively in Liger -- check for remaining opportunities.\n\n**Kernel tiers**: Tier 3 (most register pressure due to complex logic). Tier 2 backward passes (accumulation variables).\n\n### 3.5 2D Grid Tiling for Better SM Utilization\n\n**When to apply**: Kernel uses a 1D grid but the problem has two natural dimensions (e.g., batch and feature). A 1D grid may leave SMs idle when the number of programs is not a multiple of SM count.\n\n**How to implement**:\n\n```python\n# 1D grid (current Liger pattern for most kernels):\ngrid = (n_rows,)\n# If n_rows = 100 and SM count = 108, 8 SMs are idle in the last wave\n\n# 2D grid for better utilization:\ngrid = (triton.cdiv(n_rows, BLOCK_ROW), triton.cdiv(n_cols, BLOCK_COL))\n\n@triton.jit\ndef kernel_2d(\n    ...,\n    BLOCK_ROW: tl.constexpr,\n    BLOCK_COL: tl.constexpr,\n):\n    row_block = tl.program_id(0)\n    col_block = tl.program_id(1)\n    row_start = row_block * BLOCK_ROW\n    col_start = col_block * BLOCK_COL\n    # Process a tile of [BLOCK_ROW, BLOCK_COL]\n```\n\nA 2D grid is most useful when the problem is large in both dimensions and each tile does enough work to amortize launch overhead.\n\n**Expected impact**: 5-15% for kernels with suboptimal SM utilization. More impactful when `n_rows` is small relative to SM count.\n\n**Risks**: 2D tiling for reduction kernels requires careful handling of partial results and synchronization. More complex code. Not beneficial when `n_cols` fits in a single block.\n\n**Kernel tiers**: Tier 2 (reductions with large feature dimensions), Tier 3 (large fused operations).\n\n### 3.6 Loop Unrolling Hints\n\n**When to apply**: Kernel has a hot loop with a known iteration count or small trip count. Triton's compiler unrolls some loops automatically, but explicit hints can help.\n\n**How to implement**:\n\n```python\n# Triton unrolls loops with constexpr bounds automatically.\n# For dynamic bounds, you can help by:\n\n# Strategy 1: Make the loop bound a constexpr when possible\n# If n_cols is known at compile time (e.g., passed as constexpr):\nfor i in range(0, N_COLS, BLOCK_SIZE):  # N_COLS: tl.constexpr -> unrolled\n    ...\n\n# Strategy 2: Use tl.static_range for explicit unrolling\nfor i in tl.static_range(4):  # always unrolled\n    X_block = tl.load(X_ptr + i * BLOCK_SIZE + offsets)\n    ...\n\n# Strategy 3: Break a large loop into a known-count outer loop + remainder\nCHUNKS: tl.constexpr = n_cols // BLOCK_SIZE  # constexpr division\nfor i in tl.static_range(CHUNKS):\n    ...\n# Handle remainder separately\n```\n\n**Expected impact**: 5-10% for tight loops with few iterations. Diminishing returns for large trip counts.\n\n**Risks**: Excessive unrolling increases instruction cache pressure and compile time. Can also increase register pressure if the loop body is large.\n\n**Kernel tiers**: Tier 2 (reduction loops), Tier 3 (multi-pass loops).\n\n---\n\n## 4. Architecture-Specific Optimizations\n\nApply these on top of the best general-purpose variant when targeting a specific GPU architecture.\n\n### 4.1 Ampere (SM 8.x: A100, A10, A30)\n\n**Async Copy (cp.async)**:\n- Triton uses `num_stages` to control software pipelining, which maps to `cp.async` on Ampere.\n- `num_stages=2` is the sweet spot for most Ampere kernels. Going higher increases shared memory usage without benefit.\n- Triton handles `cp.async` automatically when `num_stages > 1`.\n\n**L2 Cache Persistence**:\n- A100 has 40MB L2 cache. Tensors smaller than this can benefit from persistence hints.\n- Use `eviction_policy=\"evict_last\"` on reused data (weights, biases).\n- For cross-entropy, the target tensor and weight tensor are good candidates.\n\n**Optimal Warp Counts**:\n- Ampere excels with 4-16 warps per SM. 32 warps often increases scheduling overhead.\n- For memory-bound kernels: fewer warps (4-8) with larger BLOCK_SIZE.\n- For compute-bound kernels: more warps (8-16) to hide latency.\n\n**Implementation**:\n```python\n# Ampere-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=2),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=2),\n```\n\n### 4.2 Hopper (SM 9.0: H100, H200)\n\n**TMA (Tensor Memory Accelerator)**:\n- Hopper's TMA offloads address computation and data movement to a dedicated unit.\n- Triton >= 3.0 exposes TMA via `tl.make_block_ptr` and experimental APIs.\n- TMA is most beneficial for 2D tiled access patterns (matmul-like kernels).\n\n```python\n# TMA-friendly block pointer (Triton >= 3.0)\nblock_ptr = tl.make_block_ptr(\n    base=X_ptr,\n    shape=(M, N),\n    strides=(stride_m, stride_n),\n    offsets=(row_start, col_start),\n    block_shape=(BLOCK_M, BLOCK_N),\n    order=(1, 0),  # row-major\n)\nX_block = tl.load(block_ptr, boundary_check=[0, 1])\n```\n\n**Higher num_stages**:\n- Hopper's async pipeline supports deeper prefetching. `num_stages=3-5` often outperforms 2.\n- The deeper pipeline hides memory latency more effectively on H100's higher-bandwidth HBM3.\n\n**WGMMA (Warp Group Matrix Multiply Accumulate)**:\n- H100 tensor cores are organized in warp groups (4 warps = 128 threads).\n- `num_warps` should be a multiple of 4 for optimal tensor core usage.\n- Triton maps `tl.dot` to WGMMA automatically when conditions are met.\n\n**Cluster Launch**:\n- Hopper supports thread block clusters (groups of SMs that can share SRAM).\n- Currently experimental in Triton. Potentially useful for large reductions.\n\n**Implementation**:\n```python\n# Hopper-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=4),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=32, num_stages=4),\n```\n\n### 4.3 Blackwell (SM 10.0: B200, GB200)\n\n**5th-Gen Tensor Cores**:\n- Higher throughput for FP8, FP16, BF16, and TF32 operations.\n- Block sizes aligned to 16 remain optimal. FP8 support enables even higher throughput.\n- Triton support is evolving; check Triton release notes for Blackwell-specific features.\n\n**Larger Shared Memory**:\n- Blackwell supports up to 256KB shared memory per SM (vs 228KB on Hopper).\n- Larger BLOCK_SIZE values become viable without spilling to registers or HBM.\n- Can enable larger `BLOCK_ROW` in block-row mode.\n\n**Higher Memory Bandwidth**:\n- HBM3E provides more bandwidth than HBM3. Memory-bound kernels see direct benefit.\n- The compute-to-memory ratio shifts, making some previously memory-bound kernels balanced or compute-bound.\n\n**Implementation**:\n```python\n# Blackwell-tuned config (aggressive block sizes)\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 16384}, num_warps=32, num_stages=5),\n```\n\n**Note**: Blackwell-specific Triton features are still maturing. Focus on general optimizations first, then add Blackwell-specific configs as Triton support stabilizes.\n\n---\n\n## 5. Balanced Optimizations\n\nThese techniques help both memory-bound and compute-bound kernels.\n\n### 5.1 Kernel Fusion Opportunities\n\n**When to apply**: Two or more kernels execute sequentially on the same data, writing intermediate results to HBM between them. The fused version keeps intermediates in registers/SRAM.\n\n**How to implement**:\n\nLiger already excels at fusion (cross_entropy fuses softmax+loss+gradient; fused_linear_cross_entropy fuses matmul+CE). Look for remaining opportunities:\n\n```python\n# Pattern: operation followed by elementwise scaling\n# BEFORE: two kernel launches\ny = rms_norm_kernel(x, w)        # writes y to HBM\nz = elementwise_mul_kernel(y, s)  # reads y from HBM, writes z\n\n# AFTER: fused kernel\n@triton.jit\ndef fused_rms_norm_mul(X_ptr, W_ptr, S_ptr, Z_ptr, ...):\n    # RMS norm computation\n    y = x_row * rstd * (offset + w_row)\n    # Elementwise mul (no HBM round-trip)\n    z = y * s\n    tl.store(Z_ptr + offsets, z)\n```\n\nCommon fusion candidates in Liger:\n- Norm + activation (e.g., rms_norm + swiglu in some architectures)\n- Loss + metric computation (already done: cross_entropy + accuracy)\n- Dropout + residual add + norm\n\n**Expected impact**: 20-50% from eliminating HBM round-trips between fused operations.\n\n**Risks**: Fused kernels are harder to maintain, test, and debug. Register pressure increases. The fused kernel may have worse autotuning characteristics. Only fuse when the intermediate tensor is large relative to cache.\n\n**Kernel tiers**: Creates new Tier 3 kernels from combinations of Tier 1/2 kernels. Existing Tier 3 kernels may have remaining fusion opportunities.\n\n### 5.2 Better Grid Sizing\n\n**When to apply**: The number of programs (grid size) is much smaller or much larger than the SM count, leading to load imbalance or insufficient parallelism.\n\n**How to implement**:\n\n```python\nsm_count = torch.cuda.get_device_properties(device).multi_processor_count\n\n# For backward passes using SM-based parallelism (like rms_norm backward):\ngrid = (sm_count,)\nrows_per_program = math.ceil(n_rows / sm_count)\n\n# For forward passes: ensure grid is a small multiple of sm_count\n# If n_rows >> sm_count, this is automatic\n# If n_rows < sm_count, consider block-row mode or 2D grid to increase occupancy\n\n# For persistent kernels (advanced):\n# Launch exactly sm_count programs, each processes multiple work items\ngrid = (sm_count,)\nfor work_item in range(pid, total_work, NUM_SMS):\n    process(work_item)\n```\n\nThe rms_norm backward kernel already uses this SM-based grid pattern. Look for other backward kernels that use `(n_rows,)` grids but could benefit from SM-based parallelism.\n\n**Expected impact**: 5-20% from better load balancing and reduced tail effects.\n\n**Risks**: SM-based grids require careful work distribution and can introduce non-determinism if work items have unequal cost.\n\n**Kernel tiers**: All tiers. Most impactful for Tier 2 backward passes.\n\n### 5.3 Reduce Kernel Launch Overhead\n\n**When to apply**: Kernel is called many times with very short execution time (< 10 microseconds). Launch overhead dominates.\n\n**How to implement**:\n\n```python\n# Strategy 1: Batch multiple operations into one kernel\n# Instead of launching N small kernels, launch one kernel with N work items\n\n# Strategy 2: Use CUDA graphs (from PyTorch side, not inside Triton)\n# Captures a sequence of kernel launches and replays them with minimal overhead\ngraph = torch.cuda.CUDAGraph()\nwith torch.cuda.graph(graph):\n    output = liger_kernel(input)\n# Replay: graph.replay()\n\n# Strategy 3: Increase work per kernel (larger BLOCK_SIZE, block-row mode)\n# Already covered in autotuning and block-row mode sections\n```\n\n**Expected impact**: 5-15% for very short kernels. Negligible for kernels that run > 100 microseconds.\n\n**Risks**: CUDA graphs have limitations (fixed input shapes, no dynamic control flow). Batching work items increases kernel complexity.\n\n**Kernel tiers**: Tier 1 primarily (short element-wise kernels). Tier 2/3 kernels usually run long enough that launch overhead is negligible.\n\n### 5.4 Mixed Precision Strategies\n\n**When to apply**: Kernel uses FP32 throughout but only needs FP32 for specific operations (reductions, accumulations). Or kernel uses FP16/BF16 throughout but suffers from numerical issues.\n\n**How to implement**:\n\nLiger already uses mixed precision extensively (e.g., rms_norm with casting_mode). Look for additional opportunities:\n\n```python\n# Pattern 1: Compute in fp32 only where needed (Llama-style)\nX_row = tl.load(...)              # fp16/bf16\nX_fp32 = X_row.to(tl.float32)     # upcast for reduction\nmean_sq = tl.sum(X_fp32 * X_fp32) / n_cols  # fp32 reduction\nrstd = rsqrt(mean_sq + eps)        # fp32\nresult = (X_row * rstd.to(X_row.dtype)) * W  # back to original dtype\n\n# Pattern 2: Accumulate in fp32, store in original dtype\nacc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)\nfor ...:\n    acc += (dY_row * X_row).to(tl.float32)  # accumulate in fp32\ntl.store(dW_ptr, acc.to(original_dtype))     # store in original dtype\n\n# Pattern 3: Use BF16 instead of FP16 for better dynamic range\n# BF16 has same exponent bits as FP32, so less overflow/underflow risk\n# But lower mantissa precision than FP16\n```\n\n**Expected impact**: 5-20% from reducing unnecessary FP32 computation. Can also improve accuracy if precision is currently insufficient.\n\n**Risks**: Under-casting causes numerical errors. Over-casting wastes compute. The casting_mode pattern in Liger (llama vs gemma vs none) shows that different models need different precision strategies. Always validate against reference implementation.\n\n**Kernel tiers**: All tiers. Tier 2 reductions are most sensitive to precision. Tier 1 kernels with transcendentals (sigmoid, exp) benefit from targeted FP32 upcast.\n\n---\n\n## Quick Reference: Technique Selection by Symptom\n\n| Symptom | Likely Bottleneck | Try These Techniques |\n|---------|------------------|---------------------|\n| NCU memory throughput > 60% | Memory-bound | 2.1, 2.2, 2.3, 2.5 |\n| NCU compute throughput > 60% | Compute-bound | 3.1, 3.2, 3.3, 3.4 |\n| Low SM occupancy | Register pressure or grid size | 3.4, 5.2, 1 (num_warps) |\n| Small n_cols, large n_rows | Under-utilized programs | 2.5, 5.2, 1 (BLOCK_SIZE) |\n| Large n_cols, small n_rows | Poor SM utilization | 3.5, 5.2 |\n| Multiple sequential kernel launches | Fusion opportunity | 5.1 |\n| Forward fast, backward slow | Backward-specific issue | 2.3, 2.4, 5.2 (SM-based grid) |\n| Performance varies wildly by shape | Bad fixed config | 1 (autotune with key) |\n| High register count in NCU | Register pressure | 3.4, 1 (smaller BLOCK_SIZE) |\n| Poor L2 cache hit rate | Cache misuse | 2.2, 2.4 |\n\n## Quick Reference: Technique Applicability by Kernel Tier\n\n| Technique | Tier 1 (element-wise) | Tier 2 (reduction) | Tier 3 (fused/complex) |\n|-----------|----------------------|--------------------|-----------------------|\n| 1. Autotuning | High | High | High |\n| 2.1 Coalescing | Low (usually ok) | Medium | Medium |\n| 2.2 Cache modifiers | Low | Medium | High |\n| 2.3 Reduce traffic | Medium | Medium | High |\n| 2.4 Data reuse | Low | Medium | Medium |\n| 2.5 Block-row mode | High | High | Low |\n| 2.6 Vectorized loads | Low | Low | Low |\n| 3.1 Algorithmic | Low | High | High |\n| 3.2 Reduce redundancy | Low | Medium | High |\n| 3.3 Tensor cores | N/A | N/A | High |\n| 3.4 Register pressure | Low | Medium | High |\n| 3.5 2D grid tiling | Low | Medium | Medium |\n| 3.6 Loop unrolling | Low | Medium | Medium |\n| 4.x Arch-specific | Medium | Medium | Medium |\n| 5.1 Kernel fusion | Medium (fuse into Tier 3) | Medium | High |\n| 5.2 Grid sizing | Low | High | Medium |\n| 5.3 Launch overhead | High | Low | Low |\n| 5.4 Mixed precision | Medium | High | High |\n",".agents/skills/liger-kernel-perf/optimizer.md":"# Optimizer Agent\n\nRuns an autonomous optimization loop over an existing Liger kernel, generating multiple variants, benchmarking each against the baseline, and selecting the best performer while maintaining correctness.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\", \"cross_entropy\", \"swiglu\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `max_variants`: maximum number of variants to try (default 8)\n- `target_metric`: optional concrete target (e.g., \"forward under 0.3ms at hidden_size=4096\")\n- `optimization_profile`: the profile produced by the Profiler Agent at `optimization/{kernel}/profile.md`\n- `original_kernel`: the original kernel snapshot at `optimization/{kernel}/original_{kernel}.py`\n\n## Pre-Requisites\n\nBefore starting the optimization loop:\n\n1. Verify the workspace exists: `optimization/{kernel}/` with `original_{kernel}.py` and `profile.md`\n2. Verify the benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`\n3. Verify the test file exists: `test/transformers/test_{kernel}.py`\n4. Read the optimization profile thoroughly — extract bottleneck classification, tier, current config, and recommended strategy order\n5. Read the original kernel code thoroughly — understand every `@triton.jit` function, the `torch.autograd.Function`, and the launcher logic\n6. Read `optimization-strategies.md` for the full catalog of available techniques\n\n## Variant Numbering\n\n- v0 = baseline (original kernel, already benchmarked by the Profiler)\n- v1, v2, ... = optimization variants generated by this agent\n- Each variant gets two files:\n  - `optimization/{kernel}/{kernel}_vN.py` — the modified kernel code\n  - `optimization/{kernel}/{kernel}_vN_notes.md` — the lab notebook entry\n\n## The Optimization Loop\n\n```\nvariant_number = 1\nbest_variant = 0  (baseline)\nconsecutive_no_improvement = 0\n\nwhile variant_number <= max_variants:\n    1. Read ALL prior variant notes (v1 through v{variant_number-1})\n    2. Decide what to try next (see Strategy Selection below)\n    3. Generate variant code\n    4. Write variant notes (hypothesis section)\n    5. Run smoke test → if FAIL, mark REJECTED, increment variant_number, continue\n    6. Run full benchmark → record results\n    7. Check guardrails → if FAIL, mark REJECTED, increment variant_number, continue\n    8. Update variant notes with actual results and verdict\n    9. If improvement over best_variant: update best_variant, reset consecutive_no_improvement\n   10. Else: increment consecutive_no_improvement\n   11. If consecutive_no_improvement >= 2: STOP (diminishing returns)\n   12. If target_metric is set and met: STOP (goal achieved)\n   13. Increment variant_number\n```\n\n## Strategy Selection\n\n### Always Try Parameter Tuning First (v1)\n\nThe first variant MUST be a parameter sweep over `BLOCK_SIZE`, `num_warps`, and `num_stages`. This is the cheapest optimization with the highest expected payoff.\n\n**Why NOT @triton.autotune**: Liger kernels use `torch.autograd.Function` where the forward pass stores `BLOCK_SIZE` and `num_warps` in `ctx` and passes them to the backward kernel. With `@triton.autotune`, these become compile-time constants that cannot be stored in `ctx` — the backward kernel would not know which config the forward used. Additionally, NPU/Ascend backends do not support `num_warps` or `num_stages` at all. For these reasons, `@triton.autotune` is intentionally not used in Liger production kernels.\n\n**Instead, do a manual parameter sweep**: Replace the `calculate_settings(n_cols)` call with hardcoded values and benchmark each configuration:\n\n```python\n# In the variant, replace:\n#   BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n# With a specific config to test:\nBLOCK_SIZE = 2048   # Try: 256, 512, 1024, 2048, 4096, 8192\nnum_warps = 8       # Try: 2, 4, 8, 16, 32\nnum_stages = 2      # Try: 1, 2, 3, 4\n```\n\n**Sweep procedure for v1**:\n1. Read the current `calculate_settings(n_cols)` output for the benchmark's representative input sizes\n2. Create a variant that replaces `calculate_settings` with a custom function that returns a better config\n3. Test 3-5 promising configs (not the full cartesian product — use the architecture guidance below to prune)\n4. Pick the best config across ALL benchmark x_values (not just one size)\n5. The variant should implement a smarter `calculate_settings` replacement, not just a single hardcoded value\n\n**Architecture-specific guidance for config selection**:\n- **Ampere**: `num_stages=2`, `num_warps` up to 16, BLOCK_SIZE 1024-4096\n- **Hopper**: `num_stages=3-5`, `num_warps` up to 32, BLOCK_SIZE 2048-8192\n- **Blackwell**: `num_stages=4-5`, similar to Hopper\n\n**Important**: Keep the same forward/backward ctx pattern. Both kernels must use the same `BLOCK_SIZE` and `num_warps` for a given input. The variant must still pass these through `ctx` just like the original.\n\n### Subsequent Variants: Diagnosis-Driven\n\nAfter autotuning, choose the next technique based on:\n\n1. **The bottleneck classification from the profile** (memory-bound vs compute-bound)\n2. **Learnings from all prior variants** (what worked, what did not, what was surprising)\n3. **The recommended strategy order from the profile**\n\nConsult [optimization-strategies.md](optimization-strategies.md) for the full technique catalog. Common patterns by bottleneck type:\n\n**Memory-bound kernels** (Tier 1, simple element-wise):\n- Fuse multiple memory-bound operations into a single kernel\n- Add cache eviction hints (`tl.store(..., eviction_policy=\"evict_first\")`)\n- Increase BLOCK_SIZE to amortize memory latency\n- Use `tl.load(..., eviction_policy=\"evict_last\")` for reused data\n- Reduce unnecessary memory traffic (eliminate redundant loads/stores)\n- In-place operations where safe\n\n**Compute-bound kernels** (Tier 3, heavy math):\n- Reduce precision of intermediate computations where safe\n- Use fast-math intrinsics (`tl.math.fast_expf`, etc.)\n- Restructure computation to reduce total FLOPs\n- Improve instruction-level parallelism via loop unrolling\n- Increase occupancy by reducing register pressure\n\n**Balanced / Tier 2 kernels** (reductions):\n- Optimize the reduction pattern (two-pass vs single-pass)\n- Tune BLOCK_ROW for multi-row processing\n- Balance between register usage and parallelism\n- Consider split-K or hierarchical reductions for large dimensions\n- Block-row mode threshold tuning (the `n_rows < 4096 * 8` style checks)\n\n### Combining Techniques\n\nA single variant may combine multiple techniques if they are complementary (e.g., autotuning + cache hints). However, prefer isolated changes so that the lab notebook clearly attributes improvements to specific techniques. When combining, explain which part of the improvement is expected from each technique.\n\n### When All Strategies Exhausted\n\nIf all recommended strategies have been tried and the consecutive_no_improvement counter has not triggered, the agent may:\n- Try creative combinations of previous partial improvements\n- Attempt architecture-specific tricks (Hopper warp specialization, Blackwell TMA)\n- Revisit a rejected variant with a modified approach based on learnings\n\n## Generating a Variant\n\nFor each variant, produce a complete, self-contained kernel file at `optimization/{kernel}/{kernel}_vN.py`.\n\nRules:\n1. Start from the original kernel (not a previous variant) unless explicitly building on a prior variant's approach. If building on a prior variant, note which one and why.\n2. The variant file must be a drop-in replacement for `src/liger_kernel/ops/{kernel}.py` — same public API, same class names, same function signatures.\n3. Do not change the `torch.autograd.Function` class name or the public API.\n4. Do not remove safety checks, dtype handling, or multi-device support (XPU, NPU paths) unless the change is specifically about those areas.\n5. **Preserve ALL existing comments** — function docstrings, mathematical formula explanations, inline rationale comments, and section headers. Do NOT delete or shorten comments that explain what the code does or why.\n6. **Add explanatory comments for every optimization change** — each modified section should have a comment explaining what was changed and why it improves performance. For example:\n   ```python\n   # Optimization: compute dW before dX to allow register reuse.\n   # The original order kept 8 vectors live simultaneously (115 regs/thread).\n   # This reordering reduces peak liveness to 5-6 vectors (95 regs/thread).\n   ```\n7. **Preserve function-level docstrings** — if the original kernel functions have docstrings or parameter descriptions, keep them. Update them if the function signature changes.\n8. Add a comment block at the top of the file documenting what was changed:\n\n```python\n\"\"\"\nOptimization Variant {N}: {Short Title}\nBased on: original (or v{M} if building on a prior variant)\nChanges: {brief description of modifications}\n\"\"\"\n```\n\n## Variant Lab Notebook\n\nFor each variant, write `optimization/{kernel}/{kernel}_vN_notes.md` using the format in [templates/variant-notes.md](templates/variant-notes.md).\n\nKey requirements:\n- The Hypothesis must reference the profile's bottleneck classification and learnings from prior variants\n- Changes must include rationale for each modification with code snippets\n- Expected Outcome must be specific and quantitative (\"10-20% forward speed improvement\" not \"should be faster\")\n- Actual Results must include ALL x_values if improvement varies across sizes\n- Verdict: ACCEPTED (improves target without violating guardrails), REJECTED (fails smoke test/guardrails/no improvement), or MIXED (improves some metrics, regresses others within limits)\n- Learnings must include insights to carry forward to the next variant\n\n### Writing Notes Across the Loop\n\n- **Before smoke test**: Write the Hypothesis, Changes, and Expected Outcome sections\n- **After benchmark**: Fill in Actual Results, Verdict, and Learnings\n- **On smoke test failure**: Fill Actual Results with \"SMOKE TEST FAILED: {error}\", set Verdict to REJECTED, and write Learnings about what went wrong\n\n## Smoke Test Procedure\n\nThe smoke test validates basic correctness before investing time in a full benchmark. Execute it by temporarily swapping the variant into the source tree.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run a Minimal Correctness Check\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80\n```\n\nThis runs only float32 tests (the most numerically stable) with a timeout to catch hangs. The `-x` flag stops at the first failure.\n\n### Step 3: Interpret Results\n\n- **All tests pass**: Smoke test PASSES. Proceed to full benchmark.\n- **Any test fails**: Smoke test FAILS. Read the error carefully.\n  - If the failure is a **tolerance issue** (allclose fails but values are close): Note this in the variant notes. The optimization may have changed numerical behavior slightly. Try adjusting the variant to preserve precision, or note the tolerance difference.\n  - If the failure is a **crash or shape error**: The variant has a bug. Mark REJECTED immediately. Record the error in the notes for learning.\n  - If the failure is a **hang/timeout**: The variant has a deadlock or infinite loop. Mark REJECTED immediately.\n\n### Step 4: Restore the Original\n\n**CRITICAL**: Always restore the original kernel after the smoke test, regardless of outcome:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis must happen even if the smoke test crashes. Use a pattern that ensures restoration:\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py && \\\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80; \\\nSMOKE_EXIT=$?; \\\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py; \\\nexit $SMOKE_EXIT\n```\n\n## Full Benchmark Procedure\n\nAfter a variant passes the smoke test, run the full existing benchmark suite to get precise measurements.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run the Full Benchmark\n\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nThis runs speed benchmarks (forward, backward, full) and memory benchmarks across all configured x_values, comparing the \"liger\" provider against the baseline (\"huggingface\" or \"torch\").\n\n### Step 3: Capture Results\n\nThe benchmark writes results to `benchmark/data/all_benchmark_data.csv`. Copy the relevant rows:\n\n```bash\ncp benchmark/data/all_benchmark_data.csv optimization/{kernel}/benchmarks/vN_results.csv\n```\n\n### Step 4: Parse Results\n\nRead the CSV and extract:\n- Speed (ms) at the median (p50) for forward, backward, and full modes\n- Memory (MB) at the median (p50) for full mode\n- Values at each x_value to detect size-dependent behavior\n\nCompare against:\n- v0 baseline: the absolute improvement\n- Current best variant: whether this is the new best\n\n### Step 5: Restore the Original\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n## Guardrail Checks\n\nAfter benchmarking each variant, apply these guardrails. All comparisons are against the v0 baseline.\n\n### Guardrail 1: Non-Target Metric Regression (>5%)\n\nIf the optimization goal is \"speed\", check that memory does not regress by more than 5%. If the goal is \"memory\", check that speed does not regress by more than 5%. If the goal is \"balanced\", check both.\n\n```\nif goal == \"speed\":\n    reject if memory_vN > memory_v0 * 1.05\nelif goal == \"memory\":\n    reject if speed_vN > speed_v0 * 1.05  (higher ms = slower = worse)\nelif goal == \"balanced\":\n    reject if memory_vN > memory_v0 * 1.05 OR speed_vN > speed_v0 * 1.05\n```\n\nCheck this at EVERY x_value. If the regression occurs only at certain sizes, note this in the variant notes but still reject — partial regressions are not acceptable.\n\n### Guardrail 2: Cross-Pass Regression (>10%)\n\nDo not sacrifice one pass (forward or backward) to improve the other beyond a 10% threshold:\n\n```\nreject if speed_fwd_vN > speed_fwd_v0 * 1.10 AND speed_bwd_vN < speed_bwd_v0\nreject if speed_bwd_vN > speed_bwd_v0 * 1.10 AND speed_fwd_vN < speed_fwd_v0\n```\n\nAn optimization that makes forward 15% slower to make backward 5% faster is rejected.\n\n### Guardrail 3: Smoke Test Failure\n\nAlready handled in the smoke test procedure. Any correctness failure means immediate rejection.\n\n### Handling Guardrail Failures\n\nWhen a variant is rejected by guardrails:\n1. Set Verdict to REJECTED in the variant notes\n2. Record exactly which guardrail was violated and by how much\n3. Record the Learnings — this information is valuable for future variants\n4. Do NOT increment consecutive_no_improvement (guardrail failures are not \"stagnation\")\n5. Proceed to the next variant\n\n## Stopping Conditions\n\nThe loop terminates when ANY of these conditions is met:\n\n### Condition 1: Budget Exhausted\n\n`variant_number > max_variants`\n\nReport: \"Optimization budget of {max_variants} variants exhausted.\"\n\n### Condition 2: Diminishing Returns\n\n`consecutive_no_improvement >= 2`\n\nTwo consecutive variants that pass guardrails but improve the target metric by less than 1% over the current best. This indicates the kernel is near its optimization ceiling for the current approach.\n\nReport: \"Stopping due to diminishing returns — last 2 accepted variants showed <1% improvement.\"\n\n### Condition 3: Target Metric Met\n\nThe user-specified `target_metric` has been achieved.\n\nReport: \"Target metric achieved by variant v{N}: {metric_description}.\"\n\n## Reading Prior Variant Notes\n\nBefore generating each new variant, the agent MUST read ALL prior variant notes. This is the accumulated learning mechanism that makes later variants smarter.\n\nWhen reading prior notes, extract:\n1. **What techniques were tried**: Avoid repeating the exact same approach\n2. **What worked**: Build on successful directions\n3. **What failed and why**: Avoid repeating mistakes\n4. **Surprising results**: These often reveal the true bottleneck\n5. **Size-dependent behavior**: Some optimizations help at large sizes but hurt at small sizes\n\nUse this information to form a more informed hypothesis for the next variant.\n\n## Final Comparison Table\n\nAfter the loop terminates (for any reason), produce a comparison table of ALL variants and save it to `optimization/{kernel}/comparison.md`.\n\nThe comparison must include:\n- A master table with ALL variants (accepted and rejected): Variant | Strategy | Speed Fwd/Bwd/Full (ms) | Memory (MB) | Verdict\n- Values at the representative x_value (typically the largest or most common) with percentage delta from baseline\n- Winner declaration with justification\n- Cumulative learnings summary (what worked, what didn't, why)\n\n## Error Recovery\n\n### Benchmark Script Crashes\n\nIf the benchmark script crashes (not a test failure — an actual crash):\n1. Restore the original kernel immediately\n2. Read the crash output\n3. If it is a Triton compilation error: the variant has a syntax or semantic error in the kernel. Mark REJECTED with the error.\n4. If it is a CUDA OOM: the variant uses too much memory for the GPU. Mark REJECTED. Note the memory constraint.\n5. If it is an unrelated infrastructure error: retry once. If it fails again, report to the user.\n\n### Workspace Corruption\n\nIf at any point `src/liger_kernel/ops/{kernel}.py` is found to differ from `optimization/{kernel}/original_{kernel}.py` when it should be the original:\n\n```bash\ndiff src/liger_kernel/ops/{kernel}.py optimization/{kernel}/original_{kernel}.py\n```\n\nIf they differ unexpectedly, restore immediately:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Agent Failure Mid-Loop\n\nIf the agent encounters an unrecoverable error:\n1. Restore the original kernel\n2. Save all variant notes written so far\n3. Write a partial comparison table with whatever data is available\n4. Report to the user with the current state and what went wrong\n\n## Important Constraints\n\n1. **Never leave a variant installed in the source tree.** The original kernel must always be restored after each smoke test and benchmark. This is the most critical invariant.\n2. **Never modify test files.** The existing tests are the correctness oracle. If a variant fails tests, the variant is wrong, not the tests.\n3. **Never modify benchmark files.** The existing benchmarks are the measurement standard. Variants must be measured by the same benchmark to enable fair comparison.\n4. **Always write notes before running experiments.** The hypothesis must be written before seeing results to avoid post-hoc rationalization.\n5. **Always read all prior notes before generating a new variant.** This is what makes the optimization loop learn from itself.\n6. **Prefer the simplest change that achieves the goal.** A 5-line autotuning change that yields 15% speedup is better than a 200-line rewrite that yields 20%.\n",".agents/skills/liger-kernel-perf/profiler.md":"# Profiler Agent\n\nProfiles an existing Liger kernel to diagnose performance bottlenecks and produce an optimization strategy.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `target_gpu`: Ampere / Hopper / Blackwell / auto-detect\n\n## Steps\n\n### Step 1: Set Up Workspace\n\n1. Create directory `optimization/{kernel}/` at the repo root\n2. Copy `src/liger_kernel/ops/{kernel}.py` to `optimization/{kernel}/original_{kernel}.py`\n3. Create `optimization/{kernel}/benchmarks/` subdirectory\n\n### Step 2: Detect GPU Architecture\n\nRun:\n```bash\npython -c \"import torch; props = torch.cuda.get_device_properties(0); print(f'GPU: {props.name}, SM: {props.major}.{props.minor}, SMs: {props.multi_processor_count}, Mem: {props.total_mem // (1024**3)}GB')\"\n```\n\nClassify architecture:\n- SM 8.x → Ampere (A100, A10, etc.)\n- SM 9.0 → Hopper (H100, H200, etc.)\n- SM 10.0 → Blackwell (B200, etc.)\n\nIf user specified a target_gpu, note the mismatch (if any) but proceed with the actual hardware.\n\n### Step 3: Run Baseline Benchmarks\n\nRun the existing benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nRecord the baseline results. Copy the relevant rows from `benchmark/data/all_benchmark_data.csv` to `optimization/{kernel}/benchmarks/v0_baseline.csv`.\n\nParse and summarize the baseline numbers:\n- Speed: forward, backward, full (median ms at each x_value)\n- Memory: full (median MB at each x_value)\n- Identify the slowest operation mode and largest x_values\n\n### Step 4: Check for NCU Availability\n\n```bash\nwhich ncu 2>/dev/null && echo \"NCU_AVAILABLE\" || echo \"NCU_NOT_AVAILABLE\"\n```\n\nIf available, offer to run profiling (interactive mode) or run it (autonomous mode).\n\nTo determine the correct import and invocation, read the kernel's `torch.autograd.Function` class from `src/liger_kernel/ops/{kernel}.py` and construct a minimal reproduction script. For example, for `rms_norm`:\n\n```bash\nncu --set full --target-processes all -o optimization/{kernel}/ncu_profile \\\n    python -c \"\nimport torch\nfrom liger_kernel.ops.rms_norm import LigerRMSNormFunction\nx = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16, requires_grad=True)\nw = torch.ones(4096, device='cuda', dtype=torch.bfloat16)\ny = LigerRMSNormFunction.apply(x, w, 1e-6, 0.0, 'llama', False)\ny.sum().backward()\n\"\n```\n\nAdapt the import path, function arguments, and tensor shapes for the target kernel.\n\nExtract key metrics from NCU output:\n- **SM Occupancy** (%)\n- **Compute Throughput** (% of peak)\n- **Memory Throughput** (% of peak)\n- **L1/L2 Cache Hit Rates**\n- **Warp Stall Reasons** (top 3)\n\n### Step 5: Analyze Kernel Code\n\nRead `src/liger_kernel/ops/{kernel}.py` thoroughly. Extract:\n\n1. **Tier classification**:\n   - Tier 1 (element-wise): No reductions. One row per program.\n   - Tier 2 (reduction): Cross-column reductions. May need SM-based parallelism.\n   - Tier 3 (fused/complex): Multi-pass algorithms, gradient-in-forward tricks.\n\n2. **Current configuration**:\n   - How BLOCK_SIZE is set (calculate_settings? hardcoded? autotune?)\n   - num_warps value\n   - num_stages value (if specified)\n   - Grid dimensions (1D, 2D, block-row mode?)\n\n3. **Memory access patterns**:\n   - Stride-based or contiguous?\n   - Cache modifiers used?\n   - In-place operations?\n   - What is saved_for_backward vs recomputed?\n\n4. **Compute patterns**:\n   - Precision-sensitive operations (sigmoid, rsqrt, exp, log)?\n   - Casting modes?\n   - Online/chunked algorithms?\n   - constexpr parameters?\n\n5. **Existing optimization opportunities**:\n   - Any TODO/FIXME comments about performance?\n   - Missing block-row mode variant?\n   - Suboptimal BLOCK_SIZE heuristic?\n   - Unnecessary memory traffic?\n\n### Step 6: Classify Bottleneck\n\nUse NCU data (if available) or heuristics:\n\n**Memory-bound indicators**:\n- NCU: Memory throughput > 60% of peak, compute throughput < 40%\n- Heuristic: Simple element-wise ops, large tensors, few arithmetic operations per load\n- Heuristic: Tier 1 kernels are almost always memory-bound\n\n**Compute-bound indicators**:\n- NCU: Compute throughput > 60% of peak, memory throughput < 40%\n- Heuristic: Complex math per element (exp, log, sigmoid chains), reductions, many operations per load\n- Heuristic: Tier 3 kernels with heavy computation\n\n**Balanced (both bottlenecks)**:\n- NCU: Both throughputs 30-60%\n- Heuristic: Tier 2 kernels often fall here\n\n### Step 7: Produce Optimization Profile\n\nWrite `optimization/{kernel}/profile.md` using the [optimization-profile template](templates/optimization-profile.md).\n\nThe profile must include:\n- GPU architecture and properties\n- Baseline benchmark summary (table)\n- NCU summary (if available)\n- Kernel tier classification\n- Bottleneck classification with evidence\n- Current configuration analysis\n- Recommended strategy order (from [optimization-strategies.md](optimization-strategies.md))\n- Specific opportunities identified in the code\n\n### Step 8: Present to User (Interactive Mode)\n\nShow:\n1. GPU detected: {name} ({architecture})\n2. Baseline performance summary table\n3. Bottleneck classification: {memory-bound / compute-bound / balanced}\n4. Evidence for classification\n5. Proposed optimization strategy order\n6. Estimated iteration count\n\nWait for user confirmation before proceeding to Stage 2.\n\nIn autonomous mode, skip this step and proceed directly.\n",".agents/skills/liger-kernel-perf/templates/optimization-profile.md":"# Optimization Profile: {kernel}\n\n> Generated by Profiler Agent (Stage 1)\n> Date: {date}\n\n## Kernel Identity\n\n| Field | Value |\n|-------|-------|\n| Kernel Name | `{kernel}` |\n| Source File | `src/liger_kernel/ops/{kernel}.py` |\n| Benchmark Script | `benchmark/scripts/benchmark_{kernel}.py` |\n| Test File | `test/transformers/test_{kernel}.py` |\n| Tier Classification | {tier} ({tier_description}) |\n\n## GPU Information\n\n| Property | Value |\n|----------|-------|\n| GPU Name | {gpu_name} |\n| Architecture | {gpu_arch} (SM {sm_version}) |\n| SM Count | {sm_count} |\n| Total Memory | {gpu_memory_gb} GB |\n| Target GPU (user-specified) | {target_gpu} |\n| Hardware Match | {hardware_match} |\n\n## Baseline Benchmarks\n\n### Speed (median ms)\n\n| x_value | Mode | {kernel} (Liger) | {baseline_provider} | Speedup |\n|---------|------|-------------------|----------------------|---------|\n| {x_val_1} | forward | {liger_fwd_1} | {baseline_fwd_1} | {speedup_fwd_1}x |\n| {x_val_1} | backward | {liger_bwd_1} | {baseline_bwd_1} | {speedup_bwd_1}x |\n| {x_val_1} | full | {liger_full_1} | {baseline_full_1} | {speedup_full_1}x |\n| {x_val_2} | forward | {liger_fwd_2} | {baseline_fwd_2} | {speedup_fwd_2}x |\n| {x_val_2} | backward | {liger_bwd_2} | {baseline_bwd_2} | {speedup_bwd_2}x |\n| {x_val_2} | full | {liger_full_2} | {baseline_full_2} | {speedup_full_2}x |\n| ... | ... | ... | ... | ... |\n\n### Memory (median MB)\n\n| x_value | {kernel} (Liger) | {baseline_provider} | Savings |\n|---------|-------------------|----------------------|---------|\n| {x_val_1} | {liger_mem_1} | {baseline_mem_1} | {mem_savings_1}% |\n| {x_val_2} | {liger_mem_2} | {baseline_mem_2} | {mem_savings_2}% |\n| ... | ... | ... | ... |\n\n### Observations\n\n- Slowest mode: {slowest_mode} ({slowest_mode_reason})\n- Largest x_value tested: {largest_x_val}\n- Performance scaling trend: {scaling_trend}\n\n## NCU Profiling Summary\n\n> Status: {ncu_status}\n\n{ncu_available_section_start}\n\n| Metric | Value | Assessment |\n|--------|-------|------------|\n| SM Occupancy | {sm_occupancy}% | {occupancy_assessment} |\n| Compute Throughput | {compute_throughput}% of peak | {compute_assessment} |\n| Memory Throughput | {memory_throughput}% of peak | {memory_assessment} |\n| L1 Cache Hit Rate | {l1_hit_rate}% | {l1_assessment} |\n| L2 Cache Hit Rate | {l2_hit_rate}% | {l2_assessment} |\n\n### Top Warp Stall Reasons\n\n1. {warp_stall_1} ({warp_stall_pct_1}%)\n2. {warp_stall_2} ({warp_stall_pct_2}%)\n3. {warp_stall_3} ({warp_stall_pct_3}%)\n\n{ncu_available_section_end}\n\n{ncu_unavailable_section_start}\n\nNCU is not available on this system. Bottleneck classification is based on code heuristics only.\n\n{ncu_unavailable_section_end}\n\n## Bottleneck Classification\n\n**Classification: {bottleneck_class}**\n\n### Evidence\n\n| Signal | Observation | Supports |\n|--------|-------------|----------|\n| {evidence_signal_1} | {evidence_obs_1} | {evidence_supports_1} |\n| {evidence_signal_2} | {evidence_obs_2} | {evidence_supports_2} |\n| {evidence_signal_3} | {evidence_obs_3} | {evidence_supports_3} |\n| ... | ... | ... |\n\n### Reasoning\n\n{bottleneck_reasoning}\n\n## Current Configuration Analysis\n\n| Parameter | Current Value | How Set | Notes |\n|-----------|--------------|---------|-------|\n| BLOCK_SIZE | {block_size} | {block_size_source} | {block_size_notes} |\n| num_warps | {num_warps} | {num_warps_source} | {num_warps_notes} |\n| num_stages | {num_stages} | {num_stages_source} | {num_stages_notes} |\n| Grid Dimensions | {grid_dims} | {grid_dims_source} | {grid_dims_notes} |\n\n### Memory Access Patterns\n\n- Stride pattern: {stride_pattern}\n- Cache modifiers: {cache_modifiers}\n- In-place operations: {in_place_ops}\n- Saved for backward vs recomputed: {save_vs_recompute}\n\n### Compute Patterns\n\n- Precision-sensitive ops: {precision_ops}\n- Casting mode: {casting_mode}\n- Algorithm type: {algorithm_type}\n- constexpr parameters: {constexpr_params}\n\n## Identified Optimization Opportunities\n\n{opportunity_list}\n\n## Recommended Strategy Order\n\nBased on the bottleneck classification ({bottleneck_class}) and kernel analysis:\n\n{strategy_order}\n\n### Estimated Iterations\n\n- Expected variants: {estimated_variants}\n- Rationale: {variants_rationale}\n\n---\n\n*Profile saved to `optimization/{kernel}/profile.md`*\n",".agents/skills/liger-kernel-perf/templates/variant-notes.md":"# Variant {variant_number}: {variant_title}\n\n> Kernel: `{kernel}`\n> Date: {date}\n> Parent variant: {parent_variant}\n\n## Hypothesis\n\n**Target bottleneck:** {target_bottleneck}\n\n**Why this change should help:**\n\n{hypothesis_reasoning}\n\n**References from prior variants:**\n\n{prior_variant_references}\n\n## Changes\n\n{changes_list}\n\n## Expected Outcome\n\n**Predicted improvement:** {predicted_improvement}\n\n**Reasoning:** {prediction_reasoning}\n\n## Actual Results\n\n### Speed (median ms)\n\n| Metric | x_value | Baseline (v0) | This Variant (v{variant_number}) | Delta % |\n|--------|---------|---------------|----------------------------------|---------|\n| Speed forward | {x_val_1} | {baseline_fwd_1} | {variant_fwd_1} | {delta_fwd_1}% |\n| Speed forward | {x_val_2} | {baseline_fwd_2} | {variant_fwd_2} | {delta_fwd_2}% |\n| Speed backward | {x_val_1} | {baseline_bwd_1} | {variant_bwd_1} | {delta_bwd_1}% |\n| Speed backward | {x_val_2} | {baseline_bwd_2} | {variant_bwd_2} | {delta_bwd_2}% |\n| Speed full | {x_val_1} | {baseline_full_1} | {variant_full_1} | {delta_full_1}% |\n| Speed full | {x_val_2} | {baseline_full_2} | {variant_full_2} | {delta_full_2}% |\n\n### Memory (median MB)\n\n| Metric | x_value | Baseline (v0) | This Variant (v{variant_number}) | Delta % |\n|--------|---------|---------------|----------------------------------|---------|\n| Memory full | {x_val_1} | {baseline_mem_1} | {variant_mem_1} | {delta_mem_1}% |\n| Memory full | {x_val_2} | {baseline_mem_2} | {variant_mem_2} | {delta_mem_2}% |\n\n## Guardrail Checks\n\n| Check | Threshold | Result | Status |\n|-------|-----------|--------|--------|\n| Non-target metric regression | <5% worse | {nontarget_result} | {nontarget_status} |\n| Cross-pass regression | <10% on one pass to marginally improve other | {crosspass_result} | {crosspass_status} |\n| Smoke test (single shape, float32, fwd+bwd) | Pass | {smoke_result} | {smoke_status} |\n\n## Verdict\n\n**{verdict}**\n\n{verdict_explanation}\n\n## Learnings\n\n{learnings}\n\n---\n\n*Variant code: `optimization/{kernel}/{kernel}_v{variant_number}.py`*\n*Benchmark data: `optimization/{kernel}/benchmarks/v{variant_number}_results.csv`*\n","AGENTS.md":"# AGENTS.md\n\nGuidance for AI coding assistants working in this repository.\n\n## Skills\n\nRepository-specific workflow guides (\"skills\") live in [`.agents/skills/`](.agents/skills/). Each subdirectory is a self-contained guide for a multi-stage workflow. Read the `SKILL.md` at the top of each subdirectory for an overview; the other files in that directory are referenced from `SKILL.md` and should be read on demand.\n\n| Skill | What it does |\n|-------|--------------|\n| [`liger-kernel-dev`](.agents/skills/liger-kernel-dev/SKILL.md) | Develops new Triton kernels from a PyTorch reference (or modifies existing kernels). 3-stage pipeline: Analyze → Generate → Validate. NVIDIA GPUs only. |\n| [`liger-autopatch`](.agents/skills/liger-autopatch/SKILL.md) | Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies an existing monkey-patch. 3-stage pipeline: Analyze → Generate → Validate. |\n| [`liger-kernel-perf`](.agents/skills/liger-kernel-perf/SKILL.md) | Optimizes the performance of an existing Liger Triton kernel. 3-stage pipeline: Profile → Optimize → Finalize. NVIDIA GPUs only. |\n\nThe skills are written to be runtime-agnostic — they describe the workflow as a sequence of stages a competent agent (or human) can follow. Where a stage says \"Follow the X workflow in `x.md`\", that's a directive to read and execute that file's instructions; runtimes that support parallel subagents may delegate the stage, but it is not required.\n\n## Vendor-specific shortcuts\n\nFor convenience, some assistants auto-discover skills from vendor-specific paths. These point at the canonical `.agents/skills/` directory:\n\n- `.claude/skills` → symlink → `.agents/skills` (for Claude Code)\n\nIf you're adding support for another assistant, add a symlink (or your tool's preferred adapter) pointing to `.agents/skills/`. Do not duplicate the content.\n\n## Repo conventions\n\n- Source layout: `src/liger_kernel/{ops,transformers}/` for Triton ops and `nn.Module` / HF wrappers respectively\n- Tests: `test/transformers/` (unit) and `test/convergence/{bf16,fp32}/` (model convergence)\n- Benchmarks: `benchmark/scripts/` (scripts) and `benchmark/data/all_benchmark_data.csv` (results)\n- Lint/format: `make checkstyle` (uses `ruff`)\n- Install dev mode: `pip install -e \".[dev]\"`\n\nSee `README.md` for the project overview and contribution guide.\n"},"items":[{"name":"AGENTS.md","path":"AGENTS.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/AGENTS.md","title":"AI Agent Protocol & Instructions","category":"root-instruction","format":"markdown","content":"# AGENTS.md\n\nGuidance for AI coding assistants working in this repository.\n\n## Skills\n\nRepository-specific workflow guides (\"skills\") live in [`.agents/skills/`](.agents/skills/). Each subdirectory is a self-contained guide for a multi-stage workflow. Read the `SKILL.md` at the top of each subdirectory for an overview; the other files in that directory are referenced from `SKILL.md` and should be read on demand.\n\n| Skill | What it does |\n|-------|--------------|\n| [`liger-kernel-dev`](.agents/skills/liger-kernel-dev/SKILL.md) | Develops new Triton kernels from a PyTorch reference (or modifies existing kernels). 3-stage pipeline: Analyze → Generate → Validate. NVIDIA GPUs only. |\n| [`liger-autopatch`](.agents/skills/liger-autopatch/SKILL.md) | Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies an existing monkey-patch. 3-stage pipeline: Analyze → Generate → Validate. |\n| [`liger-kernel-perf`](.agents/skills/liger-kernel-perf/SKILL.md) | Optimizes the performance of an existing Liger Triton kernel. 3-stage pipeline: Profile → Optimize → Finalize. NVIDIA GPUs only. |\n\nThe skills are written to be runtime-agnostic — they describe the workflow as a sequence of stages a competent agent (or human) can follow. Where a stage says \"Follow the X workflow in `x.md`\", that's a directive to read and execute that file's instructions; runtimes that support parallel subagents may delegate the stage, but it is not required.\n\n## Vendor-specific shortcuts\n\nFor convenience, some assistants auto-discover skills from vendor-specific paths. These point at the canonical `.agents/skills/` directory:\n\n- `.claude/skills` → symlink → `.agents/skills` (for Claude Code)\n\nIf you're adding support for another assistant, add a symlink (or your tool's preferred adapter) pointing to `.agents/skills/`. Do not duplicate the content.\n\n## Repo conventions\n\n- Source layout: `src/liger_kernel/{ops,transformers}/` for Triton ops and `nn.Module` / HF wrappers respectively\n- Tests: `test/transformers/` (unit) and `test/convergence/{bf16,fp32}/` (model convergence)\n- Benchmarks: `benchmark/scripts/` (scripts) and `benchmark/data/all_benchmark_data.csv` (results)\n- Lint/format: `make checkstyle` (uses `ruff`)\n- Install dev mode: `pip install -e \".[dev]\"`\n\nSee `README.md` for the project overview and contribution guide.\n","isInternal":false,"tokens":572,"sizeBytes":2361},{"name":"code-generator.md","path":".agents/skills/liger-autopatch/code-generator.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/code-generator.md","title":"Liger-autopatch Skill","category":"anthropic-skill","format":"markdown","content":"# Code Generator Agent\n\nTakes a confirmed model profile (create mode) or a change plan (modify mode) and generates or modifies files for Liger Kernel support.\n\n## Mode\n\n- **Create mode** (default): Generating all files for a new model. Follow the full \"Files to Generate\" list below.\n- **Modify mode**: Making targeted changes to an existing monkey-patch. Follow the \"Modification Checklist\" section instead.\n\n## Pre-Requisites\n\nBefore generating, read the reference implementation closest to this model:\n- Dense → `src/liger_kernel/transformers/model/llama.py`\n- MoE → `src/liger_kernel/transformers/model/mixtral.py`\n- Vision-Language → `src/liger_kernel/transformers/model/qwen2_vl.py`\n- Gemma-family → `src/liger_kernel/transformers/model/gemma.py`\n\nAlso read the corresponding patching function in `monkey_patch.py` and the templates in [templates/](templates/).\n\n## Files to Generate\n\n### 1. `src/liger_kernel/transformers/model/{model_type}.py` (NEW)\n\nThe `lce_forward` function. See [templates/lce-forward-dense.md](templates/lce-forward-dense.md) or [templates/lce-forward-moe.md](templates/lce-forward-moe.md).\n\nKey rules:\n- Match the exact forward signature from HF's `ForCausalLM.forward`\n- Use `lce_maybe_trainable_lm_head` from `llama.py` (shared PEFT/FSDP utility)\n- If model needs custom loss args (e.g., softcapping), write a local helper instead\n\n### 2. `src/liger_kernel/transformers/monkey_patch.py` (MODIFY)\n\nThree changes — see [templates/monkey-patch-fn.md](templates/monkey-patch-fn.md):\n\n**A.** Add lce_forward import (~line 18-28):\n```python\nfrom liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n```\n\n**B.** Add `apply_liger_kernel_to_{model_type}` function with both class-level and instance-level patching paths.\n\n**C.** Add entry to `MODEL_TYPE_TO_APPLY_LIGER_FN` dict (~line 3067).\n\n### 3. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the function in three locations (maintain alphabetical order):\n- `TYPE_CHECKING` block\n- `__getattr__` monkey_patch_symbols set\n- `__all__` list extension\n\n### 4. `src/liger_kernel/transformers/model/output_classes.py` (MODIFY if needed)\n\nOnly for models needing custom output (MoE with `aux_loss`, VL with `rope_deltas`). Follow the existing guarded-import pattern in the file.\n\n### 5. `test/transformers/test_monkey_patch.py` (MODIFY)\n\nSee [templates/test-instance-patch.md](templates/test-instance-patch.md). Add availability checker + skipif-decorated test function using `inspect.getsource()` assertions.\n\n### 6. Convergence tests (MODIFY multiple files)\n\nSee [templates/test-convergence.md](templates/test-convergence.md). Every model needs entries in multiple convergence test files:\n\n**All text models (dense + MoE)** — add to these 4 files:\n- `test/convergence/bf16/test_mini_models.py` — FLCE path, bf16\n- `test/convergence/bf16/test_mini_models_with_logits.py` — non-FLCE path (tests RMSNorm/SwiGLU/RoPE only), bf16\n- `test/convergence/fp32/test_mini_models.py` — FLCE path, fp32\n- `test/convergence/fp32/test_mini_models_with_logits.py` — non-FLCE path, fp32\n\n**Vision-language models** — also add to these 2:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\nEach file needs: imports, availability guard, `MiniModelConfig` entry in `MINI_MODEL_SETUPS` dict, and a `pytest.param` entry in the parametrize block. The `MiniModelConfig` entry is identical across all files for the same model. The `pytest.param` tolerances differ — use bf16 tolerances (looser) for bf16 files and fp32 tolerances (tighter) for fp32 files. Copy tolerance values from a similar existing model (e.g., Llama for dense, Mixtral for MoE).\n\n### 7. `test/utils.py` (MODIFY)\n\nAdd `revert_liger_kernel_to_{model_type}` function that reloads the modeling module.\n\n### 8. `README.md` (MODIFY)\n\nAdd row to the Patching table under \"### Patching\":\n```\n| {ModelName} | `liger_kernel.transformers.apply_liger_kernel_to_{model_type}` | {Supported Operations} |\n```\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order** alongside similar existing entries. Never append to the end of a section — find the correct alphabetical position.\n- After generating all files, run `make checkstyle` to verify formatting. If it fails, run `ruff check . --fix && ruff format .` to auto-fix, then verify with `make checkstyle` again.\n\n## Modification Checklist (Modify Mode)\n\nBefore making changes, read the existing implementation:\n1. Read `apply_liger_kernel_to_{model_type}` in `monkey_patch.py`\n2. Read the existing test in `test_monkey_patch.py` for this model\n3. Read the relevant HF modeling source for context\n\n### Rules for All Modifications\n\n**R1. Both patching levels.** If adding a new kernel, it must appear in BOTH:\n  - Class-level patching (the main body of `apply_liger_kernel_to_{model_type}`)\n  - Instance-level patching (the `if model is not None` block)\n\n  Omitting one is the most common mistake.\n\n**R2. New parameter with default.** Every new kernel gets a bool parameter on the\n  apply function signature (e.g., `relu_squared: bool = True`). Default should be `True`\n  for kernels that are safe to enable by default, `False` otherwise.\n\n**R3. Update docstring.** Update the function's docstring to:\n  - Add an `Args` entry for the new parameter\n  - Remove any stale notes that the new kernel invalidates\n    (e.g., \"squared ReLU is not supported\" → remove if you're adding it)\n\n**R4. Update tests.** In the existing `test_apply_liger_kernel_to_instance_for_{model_type}`:\n  - Add import for the new Liger kernel class\n  - Add \"not yet patched\" assertion before `_apply_liger_kernel_to_instance`\n  - Add \"correctly patched\" assertion after\n  - Follow the exact pattern of existing assertions in the same test\n\n**R5. Check revert function.** Read `revert_liger_kernel_to_{model_type}` in `test/utils.py`.\n  The revert function uses `importlib.reload(modeling_{model_type})` to undo all patches.\n  This handles most cases automatically, but check if the new kernel requires additional\n  revert logic (e.g., if the kernel patches something outside the modeling module, or\n  replaces a global like `ACT2FN` that `importlib.reload` won't fully restore). Update\n  the revert function if needed.\n\n**R6. Run convergence tests.** Don't modify convergence test files unless the change\n  requires it (e.g., new mini model config fields). But DO run existing convergence\n  tests in the Validate stage to verify no regression. This is critical — the Validator\n  agent (Stage 3) handles this, but if you are generating code without a separate\n  Validate stage, run these yourself:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  ```\n  For VL (multimodal) models, also run:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  ```\n\n**R7. Update README.md.** If the change adds a visibly new capability to the model's\n  row in the patching table (e.g., a new operation), update the supported operations list.\n\n### Common Modification Patterns\n\n**Adding an activation kernel (e.g., relu_squared for nemotron):**\n- Import the Liger kernel class at the top of `monkey_patch.py`\n- Add bool parameter to apply function signature\n- Class-level: replace in `ACT2FN` dict or replace the MLP class\n- Instance-level: patch each `decoder_layer`'s activation/MLP\n- Test: `assert isinstance` checks on the activation/MLP\n\n**Adding a norm variant:**\n- Add bool parameter to apply function\n- Class-level: replace the norm class\n- Instance-level: use `_patch_rms_norm_module` or `_patch_layer_norm_module` on all norm attrs\n- Test: `assert isinstance` checks on norm modules\n\n**Fixing missing instance patching:**\n- Read the class-level patching to see what's patched\n- Add corresponding instance-level patches in the `if model is not None` block\n- Test: add assertions that were missing\n\n**Updating for upstream HF changes:**\n- Compare the current HF modeling file against what the patch assumes\n- Update class names, attribute names, forward signatures as needed\n- May require updating `lce_forward` if the base model's forward changed\n","isInternal":false,"tokens":2184,"sizeBytes":8695},{"name":"decision-matrix.md","path":".agents/skills/liger-autopatch/decision-matrix.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/decision-matrix.md","title":"Liger-autopatch Skill","category":"anthropic-skill","format":"markdown","content":"# Decision Matrix\n\nWhen analyzing a HuggingFace model for Liger Kernel support, you must resolve these 12 architectural decisions by reading the model's `modeling_*.py` source code.\n\n## 1. Norm Type\n\n**Question:** Does the model use RMSNorm, LayerNorm, or both?\n\n**How to detect:**\n- Search for `class *RMSNorm` in the modeling file → RMSNorm\n- Search for `nn.LayerNorm` usage → LayerNorm\n- Multimodal models often use both (RMSNorm for text, LayerNorm for vision)\n\n**Liger mapping:** `LigerRMSNorm` or `LigerLayerNorm`\n\n## 2. RMSNorm Casting Mode\n\n**Question:** How does the model handle dtype casting during normalization?\n\n**How to detect:** Read the RMSNorm forward method:\n- Casts input to fp32, computes variance, casts back → `\"gemma\"`\n- Computes variance in fp32 only (input stays original dtype) → `\"llama\"`\n- No casting at all → `\"none\"`\n\n**Default:** `\"llama\"` (most common)\n\n## 3. RMSNorm Offset\n\n**Question:** Does the weight have a +1.0 offset?\n\n**How to detect:** In the RMSNorm forward, look for `(1 + self.weight)` or `self.weight + 1`:\n- Present → `offset=1.0` (Gemma family)\n- Absent → `offset=0.0` (most models)\n\n## 4. RMSNorm In-Place\n\n**Question:** Can the backward pass modify dY in-place?\n\n**How to detect:** Check if the model has two sequential norm layers with a residual connection between them (like Gemma2's `pre_feedforward_layernorm` + `post_feedforward_layernorm`):\n- Sequential norms with residual → `in_place=False`\n- Otherwise → `in_place=True`\n\n## 5. MLP Activation Type\n\n**Question:** What activation function does the gated MLP use?\n\n**How to detect:** Read the MLP class forward method:\n- `silu` or `F.silu` → SwiGLU → `LigerSwiGLUMLP`\n- `gelu` or `gelu_new` or `gelu_fast` → GeGLU → `LigerGEGLUMLP`\n- Phi3-style (single gate+up projection split) → `LigerPhi3SwiGLUMLP`\n\n**Also check:** The config's `hidden_act` field.\n\n## 6. Dense vs MoE\n\n**Question:** Is the model dense, MoE, or hybrid (some layers dense, some MoE)?\n\n**How to detect:**\n- Search for `Expert`, `MoE`, `SparseMoe`, `TopK` routing classes\n- Check if decoder layers have a `block_sparse_moe` or `experts` attribute\n- Hybrid: check for `is_moe_layer` or conditional MoE per-layer\n\n**Liger mapping:**\n- Dense → standard patching\n- MoE (transformers v5) → `LigerExperts`\n- MoE (transformers v4) → `LigerBlockSparseTop2MLP`\n- Qwen3-style MoE → `LigerQwen3MoeSwiGLUMLP`\n\n## 7. Vision Components\n\n**Question:** Does the model have a vision encoder?\n\n**How to detect:**\n- Check for `pixel_values` in the `forward` signature\n- Look for a separate vision model class (e.g., `*VisionModel`)\n- Check config for `vision_config` or `text_config` sub-configs\n\n**If yes:**\n- Vision encoder norms are usually `nn.LayerNorm` → patch with `LigerLayerNorm`\n- Text and vision must be patched separately\n\n## 8. RoPE Variant\n\n**Question:** What type of positional embedding does the model use?\n\n**How to detect:** Search for the `apply_rotary_pos_emb` function:\n- Standard (q, k, cos, sin) → `liger_rotary_pos_emb` (rope=True)\n- Llama4-style → `liger_llama4_text_rotary_pos_emb`\n- Qwen2VL MRoPE → `liger_multimodal_rotary_pos_emb`\n- No rotary embedding or custom variant → `rope=False`\n\n## 9. Output Class\n\n**Question:** What return type does the model's ForCausalLM.forward use?\n\n**How to detect:** Read the return statement and type annotation:\n- Standard → `LigerCausalLMOutputWithPast`\n- MoE (has `aux_loss`) → `LigerMoeCausalLMOutputWithPast`\n- Custom VL output → create model-specific output class in `output_classes.py`\n\n## 10. Hidden State Access\n\n**Question:** How does the model access hidden states from base model output?\n\n**How to detect:** In the ForCausalLM.forward, after calling `self.model(...)`:\n- `outputs[0]` → most models (Llama, Mistral, Gemma, etc.)\n- `outputs.last_hidden_state` → Phi3, Qwen3.5 MoE, some newer models\n\n## 11. Logit Softcapping\n\n**Question:** Does the model apply softcapping to logits before loss?\n\n**How to detect:** Check config for `final_logit_softcapping`:\n- Present → pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`\n- Absent → no softcapping (most models)\n- **VL models:** Config path may be `self.config.text_config.final_logit_softcapping` instead of `self.config.final_logit_softcapping`. Check whether the model uses a composite config with `text_config` sub-config.\n\n**Models with softcapping:** Gemma2, Gemma3\n\n## 12. Decoder Layer Norm Names\n\n**Question:** What are the attribute names of norm layers in each decoder layer?\n\n**How to detect:** Read the decoder layer class `__init__`:\n- Standard: `input_layernorm`, `post_attention_layernorm`\n- Gemma2 extra: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- GLM4: `post_self_attn_layernorm`, `post_mlp_layernorm`\n- Some models: `q_norm`, `k_norm` on self_attn\n\nAlso check the final norm on the base model (usually `model.norm` or `model.final_layernorm`).\n","isInternal":false,"tokens":1384,"sizeBytes":4962},{"name":"gemma-profile.md","path":".agents/skills/liger-autopatch/examples/gemma-profile.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/examples/gemma-profile.md","title":"Examples Skill","category":"anthropic-skill","format":"markdown","content":"# Model Profile: Gemma\n\nThis profile demonstrates the key differences from Llama: GeGLU activation, RMSNorm offset, and Gemma casting mode.\n\n## Identity\n- model_type: gemma\n- causal_lm_class: GemmaForCausalLM\n- base_model_class: GemmaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.gemma.modeling_gemma\n- config_module: transformers.models.gemma.configuration_gemma\n\n## Normalization\n- norm_class: GemmaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: gemma (everything cast to fp32, then computed, then cast back)\n- offset: 1.0 (weight uses `1 + self.weight` pattern)\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: GemmaMLP\n- activation: gelu (uses GELU activation, not SiLU)\n- liger_mlp_class: LigerGEGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nSame as Llama — no extra parameters.\n\n## Mini Model Config\n```python\nGemmaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-6,\n    hidden_activation=\"gelu_pytorch_tanh\",\n)\n```\n\n## Key Differences from Llama\n\n1. **Activation**: Uses `geglu` parameter (not `swiglu`) in the patch function\n2. **RMSNorm**: Requires `offset=1.0` and `casting_mode=\"gemma\"`\n3. **MLP class**: `LigerGEGLUMLP` instead of `LigerSwiGLUMLP`\n4. **Patching uses partial**: `_patch_rms_norm_module_for_gemma = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)`\n\n## Gemma2 Additional Differences\n- `in_place: false` (residual between sequential norms)\n- Extra norm layers: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- Has `final_logit_softcapping` in config\n- Uses `LigerRMSNormForGemma2` variant\n","isInternal":false,"tokens":657,"sizeBytes":2240},{"name":"llama-profile.md","path":".agents/skills/liger-autopatch/examples/llama-profile.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/examples/llama-profile.md","title":"Examples Skill","category":"anthropic-skill","format":"markdown","content":"# Model Profile: Llama\n\nThis is the reference profile for the simplest and most common model type.\n\n## Identity\n- model_type: llama\n- causal_lm_class: LlamaForCausalLM\n- base_model_class: LlamaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.llama.modeling_llama\n- config_module: transformers.models.llama.configuration_llama\n\n## Normalization\n- norm_class: LlamaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: llama\n- offset: 0.0\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: LlamaMLP\n- activation: silu\n- liger_mlp_class: LigerSwiGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nStandard parameters only:\n- input_ids, attention_mask, position_ids, past_key_values\n- inputs_embeds, labels, use_cache\n- output_attentions, output_hidden_states, return_dict\n- cache_position, logits_to_keep, skip_logits\n- **kwargs\n\nNo model-specific extra parameters.\n\n## Mini Model Config\n```python\nLlamaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-5,\n)\n```\n\n## Notes\n- Llama is the base reference for all dense LLaMA-family models\n- Most models (Mistral, Qwen2, OLMo, etc.) follow this exact pattern\n- The `lce_forward` for Llama reuses `lce_maybe_trainable_lm_head` for PEFT/FSDP support\n- Has `pretraining_tp` check (raises error if > 1)\n","isInternal":false,"tokens":549,"sizeBytes":1919},{"name":"model-analyzer.md","path":".agents/skills/liger-autopatch/model-analyzer.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/model-analyzer.md","title":"Liger-autopatch Skill","category":"anthropic-skill","format":"markdown","content":"# Model Analyzer Agent\n\nReads a HuggingFace model's source code and produces a structured model profile.\n\n## Input\n\nA model identifier: model_type string (e.g., `\"nemotron\"`), class name (e.g., `\"NemotronForCausalLM\"`), or Hub model name.\n\n## Steps\n\n1. Locate `transformers/models/{model_type}/modeling_{model_type}.py` and `configuration_{model_type}.py` on disk\n2. Read [decision-matrix.md](decision-matrix.md) for the 12 decisions to resolve\n3. Read the modeling file and extract all architectural details\n4. Produce the profile in the format below\n\n## What to Extract\n\nFrom the modeling file, identify:\n\n- `*ForCausalLM`, `*Model`, `*DecoderLayer`, `*RMSNorm`/norm, `*MLP`, `*Expert`/`*MoE` classes\n- `ForCausalLM.forward` signature — note all parameters\n- Norm class `forward` — determine casting mode and offset\n- All norm attribute names on decoder layer and base model\n- MLP activation function and gate/up/down projection names\n- Config's `model_type` string, `final_logit_softcapping`, `pretraining_tp`\n\n## Output Format\n\n```markdown\n# Model Profile: {ModelName}\n\n## Identity\n- model_type: {from config}\n- causal_lm_class: {e.g., NemotronForCausalLM}\n- base_model_class: {e.g., NemotronModel}\n- base_model_prefix: {e.g., \"model\"}\n- modeling_module: transformers.models.{model_type}.modeling_{model_type}\n\n## Normalization\n- norm_type: RMSNorm | LayerNorm | both\n- casting_mode: llama | gemma | none\n- offset: 0.0 | 1.0\n- in_place: true | false\n- final_norm_attr: {e.g., \"model.norm\"}\n- decoder_norm_attrs: {list}\n- attn_norm_attrs: {e.g., \"self_attn.q_norm\" or \"none\"}\n\n## MLP\n- activation: silu | gelu | gelu_new\n- liger_mlp_class: LigerSwiGLUMLP | LigerGEGLUMLP | LigerPhi3SwiGLUMLP\n\n## Structure\n- type: dense | moe | hybrid\n- moe_expert_class: {if MoE}\n- shared_expert: true | false\n\n## Positional Embedding\n- rope_type: standard | llama4 | qwen2vl_mrope | none\n\n## Output\n- output_class: LigerCausalLMOutputWithPast | LigerMoeCausalLMOutputWithPast | custom\n- hidden_state_access: outputs[0] | outputs.last_hidden_state\n- has_logit_softcapping: true | false\n\n## Vision\n- has_vision: true | false\n- vision_norm_type: {if applicable}\n\n## Forward Signature\n{List model-specific parameters not in standard Llama signature}\n\n## Mini Model Config\n{Suggest minimal config values for testing — 2 layers, 32 hidden, 64 intermediate}\n```\n\n## Rules\n\n- Never guess — always read the source to determine each value\n- Read [examples/llama-profile.md](examples/llama-profile.md) and [examples/gemma-profile.md](examples/gemma-profile.md) as reference\n- Compare against existing supported models to find the closest match\n","isInternal":false,"tokens":731,"sizeBytes":2629},{"name":"SKILL.md","path":".agents/skills/liger-autopatch/SKILL.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/SKILL.md","title":"liger-autopatch","category":"anthropic-skill","format":"markdown","content":"---\nname: liger-autopatch\ndescription: \"Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies existing monkey-patching. Generates lce_forward, monkey-patch function, tests, and README entry. Use when adding a new model to Liger Kernel, when a user asks to patch an unsupported model, when extending MODEL_TYPE_TO_APPLY_LIGER_FN, or when modifying/updating/fixing an existing monkey-patch (e.g., adding a new kernel to an already-supported model, fixing instance patching, updating a patch for upstream HF changes).\"\n---\n\n# Liger Auto-Patch\n\nAdds Liger Kernel optimization support for a new HuggingFace model, or modifies existing monkey-patching, through a staged pipeline with human review between stages. Supports creating new model patches and modifying existing ones.\n\n## Mode Detection\n\n- **Create mode**: User asks to add/patch/support a new model → full pipeline (Analyze → Generate → Validate)\n- **Modify mode**: User asks to update/fix/change/extend an existing monkey-patch → lighter pipeline (Change Impact Analysis → Apply Changes → Validate)\n\nKeywords that suggest modify mode: update, fix, change, add [kernel] to [existing model], extend, modify, new activation, new norm, bug in patch, upstream changed\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Model Analyzer** workflow in [model-analyzer.md](model-analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nThis stage reads the HF `modeling_*.py` source and produces a **model profile** answering 12 architectural questions from [decision-matrix.md](decision-matrix.md).\n\n**Human checkpoint:** Present the profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md).\n\nGenerates/modifies up to 13 files:\n\n1. `src/liger_kernel/transformers/model/{model}.py` — NEW lce_forward\n2. `src/liger_kernel/transformers/monkey_patch.py` — MODIFY\n3. `src/liger_kernel/transformers/__init__.py` — MODIFY\n4. `src/liger_kernel/transformers/model/output_classes.py` — MODIFY if needed\n5. `test/transformers/test_monkey_patch.py` — MODIFY\n6. `test/convergence/bf16/test_mini_models.py` — MODIFY (FLCE path)\n7. `test/convergence/bf16/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n8. `test/convergence/fp32/test_mini_models.py` — MODIFY (FLCE path)\n9. `test/convergence/fp32/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n10. `test/convergence/bf16/test_mini_models_multimodal.py` — MODIFY if VL model\n11. `test/convergence/fp32/test_mini_models_multimodal.py` — MODIFY if VL model\n12. `test/utils.py` — MODIFY\n13. `README.md` — MODIFY\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns instance patching test, convergence test, and lint check. Retries up to 3 times on failure.\n\n**Human checkpoint:** Report final test results.\n\n## Pipeline (Modify Mode)\n\n### Stage 1: Change Impact Analysis\n\nRead the existing `apply_liger_kernel_to_{model_type}` function in `monkey_patch.py` and the relevant section of the upstream HF `modeling_{model_type}.py`. Produce a short change plan:\n\n- What is being added/changed/fixed\n- Which Liger kernel(s) are involved\n- Which files need modification (subset of the 13 files from create mode)\n- What the expected behavior should be after the change\n\n**Human checkpoint:** Present the change plan. Confirm before proceeding.\n\n### Stage 2: Apply Changes\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md) in **modify mode**.\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md). This stage is **mandatory** — do not skip it. At minimum, run:\n\n1. Instance patching test: `pytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs`\n2. All convergence tests for the model:\n   - `pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, bf16)\n   - `pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, bf16)\n   - `pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, fp32)\n   - `pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, fp32)\n   - If VL (multimodal) model, also run:\n     - `pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n     - `pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n3. Checkstyle: `make checkstyle`\n\n**Human checkpoint:** Report final test results.\n\n## Reference Files\n\n- [decision-matrix.md](decision-matrix.md) — 12 architectural decisions to resolve per model\n- [examples/llama-profile.md](examples/llama-profile.md) — Reference profile for standard dense model\n- [examples/gemma-profile.md](examples/gemma-profile.md) — Reference profile showing GeGLU + offset variant\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n","frontmatter":{"name":"liger-autopatch","description":"Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies existing monkey-patching. Generates lce_forward, monkey-patch function, tests, and README entry. Use when adding a new model to Liger Kernel, when a user asks to patch an unsupported model, when extending MODEL_TYPE_TO_APPLY_LIGER_FN, or when modifying/updating/fixing an existing monkey-patch (e.g., adding a new kernel to an already-supported model, fixing instance patching, updating a patch for upstream HF changes)."},"isInternal":false,"tokens":1269,"sizeBytes":5177},{"name":"lce-forward-dense.md","path":".agents/skills/liger-autopatch/templates/lce-forward-dense.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/templates/lce-forward-dense.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Template: lce_forward for Dense Models\n\nReference implementation: `src/liger_kernel/transformers/model/llama.py`. Read it first.\n\n## Structure\n\nThe file needs `lce_forward` and reuses `lce_maybe_trainable_lm_head` from `llama.py`.\n\n### Imports\n\n```python\nfrom typing import TYPE_CHECKING, List, Optional, Tuple, Union\nimport torch\nfrom liger_kernel.transformers.model.llama import lce_maybe_trainable_lm_head\nfrom liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss, unpack_cross_entropy_result\nfrom liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast\nif TYPE_CHECKING:\n    from transformers.cache_utils import Cache\n```\n\n### lce_forward body pattern\n\n1. Set defaults from config (output_attentions, output_hidden_states, return_dict)\n2. Call `self.model(...)` with all forwarded params\n3. Extract hidden states: `outputs[0]` or `outputs.last_hidden_state`\n4. Slice for `logits_to_keep`\n5. Extract `shift_labels = kwargs.pop(\"shift_labels\", None)`\n6. `skip_logits` logic — default True during training with labels\n7. Compute loss via `lce_maybe_trainable_lm_head` (skip_logits) or `self.lm_head` (else)\n8. Return `LigerCausalLMOutputWithPast`\n\n### Variations\n\n- **Logit softcapping**: Pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`. Requires a local helper instead of reusing `lce_maybe_trainable_lm_head`.\n- **Hidden state access**: Use `outputs.last_hidden_state` if HF model uses that\n- **pretraining_tp**: Remove check if model doesn't have this config field\n\n### When NOT to reuse llama's lce_maybe_trainable_lm_head\n\nIf the model passes extra kwargs to `LigerForCausalLMLoss` (like softcapping), write a local `_liger_for_causal_lm_loss` and inline the PEFT/FSDP handling from llama.py.\n","isInternal":false,"tokens":457,"sizeBytes":1792},{"name":"lce-forward-moe.md","path":".agents/skills/liger-autopatch/templates/lce-forward-moe.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/templates/lce-forward-moe.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Template: lce_forward for MoE Models\n\nReference implementation: `src/liger_kernel/transformers/model/mixtral.py`. Read it first.\n\n## Differences from Dense\n\n1. **Output class**: `LigerMoeCausalLMOutputWithPast` (includes `aux_loss`)\n2. **Extra param**: `output_router_logits` in forward signature\n3. **Auxiliary loss**: Compute `load_balancing_loss_func` from router logits\n4. **Return**: Include `aux_loss` and `router_logits` in output\n\n## Auxiliary Loss Pattern\n\n```python\naux_loss = None\nif output_router_logits:\n    router_logits = outputs.router_logits if return_dict else outputs[-1]\n    aux_loss = load_balancing_loss_func(\n        router_logits, self.num_experts, self.num_experts_per_tok, attention_mask,\n    )\n    if labels is not None and loss is not None:\n        loss += self.router_aux_loss_coef * aux_loss\n```\n\n## Monkey Patch Differences for MoE\n\nExpert patching is version-aware:\n\n```python\n# Transformers v5\nmodeling_module.{ModelExperts} = LigerExperts\n# Instance: _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\n\n# Transformers v4\nmodeling_module.{ModelExpertMLP} = LigerBlockSparseTop2MLP\n# Instance: for expert in decoder_layer.block_sparse_moe.experts: _patch_swiglu_module(...)\n```\n\nCheck exact attribute names in the HF decoder layer class.\n","isInternal":false,"tokens":319,"sizeBytes":1283},{"name":"monkey-patch-fn.md","path":".agents/skills/liger-autopatch/templates/monkey-patch-fn.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/templates/monkey-patch-fn.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Template: apply_liger_kernel_to_{model_type}\n\nReference: `apply_liger_kernel_to_llama` (dense) or `apply_liger_kernel_to_mixtral` (MoE) in `monkey_patch.py`.\n\n## Function Signature\n\n```python\ndef apply_liger_kernel_to_{model_type}(\n    rope: bool = True,              # False if model lacks standard RoPE\n    cross_entropy: bool = False,\n    fused_linear_cross_entropy: bool = True,\n    rms_norm: bool = True,          # Or layer_norm for LayerNorm models\n    swiglu: bool = True,            # Or geglu for Gemma-family\n    model: PreTrainedModel = None,\n) -> None:\n```\n\n## Body Pattern\n\nClass-level patching always runs (even when `model` is provided). Instance-level patching runs additionally when `model is not None`. They are NOT mutually exclusive — this matches the existing Llama/Mistral/Gemma patterns.\n\nThe only exception is `fused_linear_cross_entropy`: use `MethodType` binding when instance is given, class replacement otherwise.\n\n```python\nassert not (cross_entropy and fused_linear_cross_entropy), \"...\"\n\nfrom transformers.models.{model_type} import modeling_{model_type}\n\n# Class-level patching (always runs)\nif rope: modeling_{model_type}.apply_rotary_pos_emb = liger_rotary_pos_emb\nif rms_norm: modeling_{model_type}.{ModelRMSNorm} = LigerRMSNorm\nif swiglu: modeling_{model_type}.{ModelMLP} = LigerSwiGLUMLP\nif cross_entropy: modeling_{model_type}.CrossEntropyLoss = LigerCrossEntropyLoss\nif fused_linear_cross_entropy:\n    if model is not None:\n        model.forward = MethodType({model_type}_lce_forward, model)\n    else:\n        modeling_{model_type}.{ModelForCausalLM}.forward = {model_type}_lce_forward\n\n# Instance-level patching (only when model instance provided)\nif model is not None:\n    base_model = getattr(model, model.base_model_prefix, model)\n    if rms_norm: _patch_rms_norm_module(base_model.norm)\n    for decoder_layer in base_model.layers:\n        if swiglu: _patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)\n        if rms_norm:\n            _patch_rms_norm_module(decoder_layer.input_layernorm)\n            _patch_rms_norm_module(decoder_layer.post_attention_layernorm)\n```\n\n## Common Variations\n\n**Gemma-family** (offset + casting):\n```python\n_patch_rms_norm_module_for_model = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)\n```\n\n**MoE experts** (version-aware):\n```python\nif IS_TRANSFORMERS_V5_OR_LATER:\n    _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\nelse:\n    for expert in decoder_layer.block_sparse_moe.experts:\n        _patch_swiglu_module(expert, LigerBlockSparseTop2MLP)\n```\n\n**Multimodal VL instance patching:**\n\nFor VL models, extract text_model and vision_model from the composite instance using `isinstance` checks. Each VL model has different accessor paths — read the existing VL patches in `monkey_patch.py` for the exact pattern. Example from Llama4:\n\n```python\nif model is not None:\n    if isinstance(model, {Model}ForConditionalGeneration):\n        language_model = model.language_model          # or model.model.language_model\n        vision_model = model.vision_model              # or model.model.visual\n        text_model = language_model.model\n    elif isinstance(model, {Model}ForCausalLM):\n        text_model = model.model\n        vision_model = None\n    elif isinstance(model, {Model}TextModel):\n        text_model = model\n        vision_model = None\n    else:\n        raise ValueError(f\"Unsupported model type: {type(model)}\")\n\n    # Patch text layers\n    if text_model:\n        if rms_norm: _patch_rms_norm_module(text_model.norm)\n        for decoder_layer in text_model.layers:\n            ...\n\n    # Patch vision layers\n    if vision_model:\n        for layer in vision_model.encoder.layers:  # accessor varies by model\n            if layer_norm:\n                _patch_layer_norm_module(layer.layer_norm1)\n                _patch_layer_norm_module(layer.layer_norm2)\n```\n\nThe exact accessor paths (`model.language_model`, `model.model.language_model`, `model.model.visual`) vary per VL model. Always read the HF model's `__init__` to find the correct paths.\n","isInternal":false,"tokens":976,"sizeBytes":4074},{"name":"test-convergence.md","path":".agents/skills/liger-autopatch/templates/test-convergence.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/templates/test-convergence.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Template: Convergence Tests\n\nReference: existing entries across `test/convergence/` files. Read the Llama entries as baseline.\n\n## Which Files to Modify\n\nEvery model needs entries in **at least 4 files**:\n\n| File | What it tests | dtype |\n|------|-------------|-------|\n| `test/convergence/bf16/test_mini_models.py` | FLCE path (fused linear cross entropy) | bfloat16 |\n| `test/convergence/bf16/test_mini_models_with_logits.py` | Non-FLCE path (tests RMSNorm/SwiGLU/RoPE only, verifies logits) | bfloat16 |\n| `test/convergence/fp32/test_mini_models.py` | FLCE path | float32 |\n| `test/convergence/fp32/test_mini_models_with_logits.py` | Non-FLCE path | float32 |\n\n**Vision-language models** also need:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\n## Placement Rule\n\nWhen adding entries to existing files, always insert in **alphabetical order** alongside similar existing entries. Never append to the end of a section. Specifically:\n- Imports: alphabetical among other `apply_liger_kernel_to_*` / `revert_liger_kernel_to_*` imports\n- Availability guards: alphabetical among other try/except model import blocks\n- `MINI_MODEL_SETUPS` entries: alphabetical among other model setup blocks\n- `pytest.param` entries: alphabetical among other model param entries\n\n## What to Add in Each File\n\nEach file has a `MINI_MODEL_SETUPS` dict and a `pytest.param` block. Add to both.\n\n### 1. Imports (same across all files, insert in alphabetical order)\n\n```python\nfrom liger_kernel.transformers import apply_liger_kernel_to_{model_type}\nfrom test.utils import revert_liger_kernel_to_{model_type}\n```\n\n### 2. Availability Guard (same across all files)\n\nThe `try/except ImportError` handles version compatibility automatically: if the model doesn't exist in the installed transformers version (e.g., 4.52.0), the import fails and the tests skip. No explicit version check needed.\n\n```python\ntry:\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n    from transformers.models.{model_type}.modeling_{model_type} import {ModelForCausalLM}\n    {MODEL_UPPER}_AVAILABLE = True\nexcept ImportError:\n    {MODEL_UPPER}_AVAILABLE = False\n```\n\n### 3. MINI_MODEL_SETUPS Entry (identical across all files)\n\n```python\nif {MODEL_UPPER}_AVAILABLE:\n    MINI_MODEL_SETUPS[\"mini_{model_type}\"] = MiniModelConfig(\n        liger_kernel_patch_func=apply_liger_kernel_to_{model_type},\n        liger_kernel_patch_revert_func=revert_liger_kernel_to_{model_type},\n        model_class={ModelForCausalLM},\n        mini_model_config={ModelConfig}(\n            hidden_size=32, intermediate_size=64, num_hidden_layers=2,\n            num_attention_heads=2, num_key_value_heads=2, vocab_size=1024,\n        ),\n    )\n```\n\n### 4. pytest.param Entry (tolerances differ by dtype)\n\n**bf16 files** — looser tolerances (copy from similar existing model):\n```python\npytest.param(\n    \"mini_{model_type}\", 32, 1e-5, torch.bfloat16,\n    1e-2, 5e-2, 1e-1, 1e-2, 1e-2, 1e-2,\n    marks=pytest.mark.skipif(not supports_bfloat16(), reason=\"bfloat16 not supported\"),\n),\n```\n\n**fp32 files** — tighter tolerances:\n```python\n(\"mini_{model_type}\", 32, 1e-4, torch.float32, 1e-8, 2e-5, 5e-3, 1e-5, 5e-3, 1e-5),\n```\n\n### 5. Revert Function in test/utils.py (shared by all files)\n\n```python\ndef revert_liger_kernel_to_{model_type}(model_config):\n    from transformers.models.{model_type} import modeling_{model_type}\n    importlib.reload(modeling_{model_type})\n    model_config.model_class = modeling_{model_type}.{ModelForCausalLM}\n```\n\n## Key Differences Between Test Variants\n\n**`with_logits`**: Sets `fused_linear_cross_entropy=False, cross_entropy=False` and checks `output.logits` match. Tests that non-loss kernels are numerically correct on their own. Same `MINI_MODEL_SETUPS` entry, just add a `pytest.param`.\n\n**`multimodal`**: Uses `multimodal_collate_fn`, image datasets, and model configs with vision sub-config. Only for models with `has_vision: true`. Multimodal tests may also require fake tokenizer/processor configs in `test/resources/fake_configs/`. Check existing VL model entries (e.g., Qwen2-VL, Llava) for the pattern — if the model needs a tokenizer/processor that can't be created programmatically, add a fake config JSON file.\n\n## Mini Config Guidelines\n\nUse minimal values: hidden_size=32, intermediate_size=64, num_hidden_layers=2, num_attention_heads=2, vocab_size=1024. Add model-specific required fields (e.g., `num_local_experts` for MoE). Copy config structure from the most similar existing model.\n","isInternal":false,"tokens":1180,"sizeBytes":4576},{"name":"test-instance-patch.md","path":".agents/skills/liger-autopatch/templates/test-instance-patch.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/templates/test-instance-patch.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Template: Instance Patching Test\n\nReference: `test_apply_liger_kernel_to_instance_for_llama` in `test/transformers/test_monkey_patch.py`.\n\n## Add Availability Checker (at top of file, near other `is_*_available` functions ~lines 59-240, in alphabetical order)\n\n```python\ndef is_{model_type}_available():\n    try:\n        import transformers.models.{model_type}\n        return True\n    except ImportError:\n        return False\n```\n\n## Add Test Function (after existing `test_apply_liger_kernel_to_instance_for_*` functions, in alphabetical order)\n\n```python\n@pytest.mark.skipif(not is_{model_type}_available(), reason=\"{model_type} not available\")\ndef test_apply_liger_kernel_to_instance_for_{model_type}():\n    from liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n\n    config = {ModelConfig}(hidden_size=32, intermediate_size=64, num_hidden_layers=2)\n    dummy_model_instance = AutoModelForCausalLM.from_config(config)\n\n    # Verify NOT patched\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) != inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)\n\n    # Apply patch\n    _apply_liger_kernel_to_instance(model=dummy_model_instance)\n\n    # Verify patched\n    assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource({model_type}_lce_forward)\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) == inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)\n```\n\n## Adapt For\n\n- **GeGLU models**: Use `LigerGEGLUMLP` in assertions\n- **MoE models**: Check `decoder_layer.mlp.experts.forward` against `LigerExperts`\n- **Extra norms**: Add assertions for `pre_feedforward_layernorm` etc.\n- **Vision**: Add assertions for `LigerLayerNorm` on vision layers\n","isInternal":false,"tokens":538,"sizeBytes":2326},{"name":"validator.md","path":".agents/skills/liger-autopatch/validator.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-autopatch/validator.md","title":"Liger-autopatch Skill","category":"anthropic-skill","format":"markdown","content":"# Validator Agent\n\nRuns tests on generated code and fixes issues iteratively.\n\n## Pre-Requisites\n\nVerify the project is installed in development mode: `python -c \"import liger_kernel\"` must succeed. If not, ask the user to run `pip install -e \".[dev]\"`.\n\n## Validation Steps\n\n### Step 1: Syntax Check\n\n```bash\npython -c \"from liger_kernel.transformers import apply_liger_kernel_to_{model_type}\"\n```\n\n### Step 2: Instance Patching Test\n\n```bash\npytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `ImportError: cannot import name` | Check exact class name in HF source |\n| `AttributeError: has no attribute` | Read decoder layer `__init__` for correct attr names |\n| `AssertionError: getsource mismatch` | Check class being replaced in monkey_patch |\n| `ModuleNotFoundError` | Add `@pytest.mark.skipif` with availability check |\n\n### Step 3: Convergence Tests (requires GPU)\n\nRun **all applicable** convergence tests. For text models, there are 4 files. For VL models, 6.\n\n**FLCE path (fused linear cross entropy):**\n```bash\npytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Non-FLCE path (logits verification — tests RMSNorm/SwiGLU/RoPE without loss fusion):**\n```bash\npytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Multimodal (only for VL models):**\n```bash\npytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `CUDA out of memory` | Reduce hidden_size/num_layers in mini config |\n| `RuntimeError: shape mismatch` | Check `self.config.hidden_size` usage |\n| `AssertionError: allclose` | Check casting_mode, offset, in_place settings |\n| `KeyError in config` | Add missing field to mini model config |\n\n### Step 4: Checkstyle\n\nRun the project's full lint and format check:\n\n```bash\nmake checkstyle\n```\n\nIf it fails, auto-fix with `ruff check . --fix && ruff format .`, then re-run `make checkstyle`.\n\n## Retry Logic\n\n**Hard failures** (Steps 1, 2, 4 — import, patching, checkstyle): These indicate code bugs. Read traceback, identify root cause, fix, re-run. Max 3 retries per step. After 3 failures, report the exact error and stop.\n\n**Soft failures** (Step 3 — convergence): If a convergence test fails after 3 retries with `allclose` tolerance errors, mark it as `SKIP (tolerance tuning needed)` in the report rather than `FAIL`. This means the patching is correct but tolerances need manual adjustment. Note this in the PR description so reviewers are aware.\n\nConvergence failures from other errors (shape mismatch, CUDA OOM, KeyError) are still hard failures — fix and retry.\n\n## Success Report\n\n```\n| Check                        | Status |\n|------------------------------|--------|\n| Import                       | PASS   |\n| Instance Patching            | PASS   |\n| Convergence bf16 FLCE        | PASS   |\n| Convergence bf16 with_logits | PASS   |\n| Convergence fp32 FLCE        | PASS   |\n| Convergence fp32 with_logits | PASS   |\n| Convergence multimodal       | SKIP (text-only model) |\n| Checkstyle                   | PASS   |\n```\n\nInclude any `SKIP (tolerance tuning needed)` entries in the PR description under a \"Known issues\" section.\n","isInternal":false,"tokens":942,"sizeBytes":3646},{"name":"analyzer.md","path":".agents/skills/liger-kernel-dev/analyzer.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/analyzer.md","title":"Liger-kernel-dev Skill","category":"anthropic-skill","format":"markdown","content":"# Analyzer Agent\n\nUnderstands a PyTorch operation from any input form and produces a standalone PyTorch reference implementation + kernel profile.\n\n## Input Handling\n\nThe user may provide the operation in any form:\n\n1. **Local file path** → Read the file directly\n2. **URL** (GitHub, HuggingFace, etc.) → Fetch via WebFetch tool\n3. **Code snippet** → Pasted in the conversation\n4. **Natural language** → Mathematical description (e.g., \"element-wise SiLU(x) * y\")\n5. **Model component** → e.g., \"the MLP in Phi-4\" — locate in transformers source and extract\n\n## Steps\n\n### 1. Understand the Operation\n\n- Read/fetch the source code from whatever input was provided\n- Identify the mathematical operation (forward pass)\n- Derive the backward pass (gradient computation)\n- Identify all inputs, outputs, and their expected shapes/dtypes\n- Note any precision-sensitive operations that need float32 upcasting (sigmoid, rsqrt, exp, log, tanh)\n\n### 2. Write PyTorch Reference\n\nCreate a standalone implementation that:\n- Depends only on `torch` (no external libraries)\n- Implements both forward and backward behavior (either as an `nn.Module` or a plain function)\n- Will serve as the correctness baseline for testing\n- Is clean, readable, and well-named\n\n### 3. Classify Into Tier\n\nRead [kernel-profile-format.md](kernel-profile-format.md) for the full schema.\n\n**Tier 1 — Element-wise**: No reductions across dimensions. One row per program. Examples: SwiGLU, GeGLU, DyT.\n- Read reference: `src/liger_kernel/ops/swiglu.py`\n\n**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.\n- Read reference: `src/liger_kernel/ops/rms_norm.py`\n\n**Tier 3 — Fused/Complex**: Multi-pass algorithms, gradient-in-forward tricks, multiple outputs. Examples: CrossEntropy, FusedLinearCrossEntropy.\n- Read reference: `src/liger_kernel/ops/cross_entropy.py`\n\nAlso read the closest example profile:\n- Tier 1 → [examples/swiglu-profile.md](examples/swiglu-profile.md)\n- Tier 2 → [examples/rms-norm-profile.md](examples/rms-norm-profile.md)\n- Tier 3 → [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md)\n\n### 4. Produce Kernel Profile\n\nFill in all fields from [kernel-profile-format.md](kernel-profile-format.md).\n\n### 5. Present to User\n\nShow:\n1. The PyTorch reference implementation (full code)\n2. The kernel profile (all fields)\n3. Which existing kernel is closest (for the Generator to use as reference)\n\nWait for user confirmation before proceeding to Stage 2.\n","isInternal":false,"tokens":632,"sizeBytes":2648},{"name":"cross-entropy-profile.md","path":".agents/skills/liger-kernel-dev/examples/cross-entropy-profile.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/examples/cross-entropy-profile.md","title":"Examples Skill","category":"anthropic-skill","format":"markdown","content":"# Kernel Profile: CrossEntropy (Tier 3 — Fused/Complex)\n\n## Identity\n- operation_name: cross_entropy\n- function_class_name: LigerCrossEntropyFunction\n- module_class_name: LigerCrossEntropyLoss\n- functional_name: liger_cross_entropy\n\n## Classification\n- tier: 3\n- tier_description: fused/complex\n- closest_existing_kernel: fused_linear_cross_entropy (extends this with linear layer fusion)\n\n## Forward Pass\n- forward_inputs:\n  - _input: shape (B*T, V), logits\n  - target: shape (B*T,), label indices\n  - weight: shape (V,) or None, class weights\n  - ignore_index: int, label to ignore\n  - label_smoothing: float\n  - reduction: str, \"mean\" | \"sum\" | \"none\"\n  - softcap: float or None\n- forward_outputs:\n  - loss: scalar or shape (B*T,) if reduction=\"none\"\n  - z_loss: optional auxiliary loss\n  - token_accuracy: optional accuracy metric\n  - predicted_tokens: optional argmax tokens\n- forward_computation: Two-pass online softmax + cross entropy loss with optional smoothing and softcapping\n- precision_sensitive_ops: [exp, log]\n\n## Backward Pass\n- backward_saved_tensors: [_input] — gradient is computed during forward and stored in-place\n- backward_recompute: none (gradient-in-forward trick)\n- gradient_formulas:\n  - d_input: already computed in forward pass and stored in _input tensor. Backward just scales by grad_output.\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row (one row = one token's logits over vocab)\n- block_size_source: custom — iterates over vocab in chunks of BLOCK_SIZE\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - weight: optional class weights\n  - ignore_index: int = -100\n  - lse_square_scale: float = 0.0\n  - label_smoothing: float = 0.0\n  - reduction: str = \"mean\"\n  - softcap: float or None = None\n  - return_z_loss: bool = False\n- learnable_params: none\n\n## Benchmarking\n- benchmark_variable: vocab_size (V)\n- benchmark_x_label: \"V\"\n- benchmark_x_values_suggestion: [4096, 8192, 16384, 32768, 65536, 131072]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {B: 8, T: 512, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Online softmax**: Two-pass algorithm. Pass 1: compute running max and logsumexp. Pass 2: compute softmax and gradients. Avoids materializing the full softmax vector.\n- **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.\n- **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.\n- **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.\n- **Multiple loss components**: Combines original CE loss, label smoothing loss, and z-loss (for training stability). Each component contributes to both loss and gradient.\n","isInternal":false,"tokens":797,"sizeBytes":3161},{"name":"rms-norm-profile.md","path":".agents/skills/liger-kernel-dev/examples/rms-norm-profile.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/examples/rms-norm-profile.md","title":"Examples Skill","category":"anthropic-skill","format":"markdown","content":"# Kernel Profile: RMSNorm (Tier 2 — Reduction)\n\n## Identity\n- operation_name: rms_norm\n- function_class_name: LigerRMSNormFunction\n- module_class_name: LigerRMSNorm\n- functional_name: liger_rms_norm\n\n## Classification\n- tier: 2\n- tier_description: reduction\n- closest_existing_kernel: layer_norm (similar pattern with different normalization)\n\n## Forward Pass\n- forward_inputs:\n  - X: shape (B, T, H), input tensor\n  - W: shape (H,), weight tensor\n  - eps: float, epsilon for numerical stability\n  - offset: float, weight offset (0.0 for Llama, 1.0 for Gemma)\n  - casting_mode: str, \"llama\" | \"gemma\" | \"none\"\n- forward_outputs:\n  - Y: shape (B, T, H), normalized output\n- forward_computation: Y = (X / RMS(X)) * (W + offset), RMS = sqrt(mean(X^2) + eps)\n- precision_sensitive_ops: [rsqrt]\n\n## Backward Pass\n- backward_saved_tensors: [X, W, RSTD] — RSTD cached from forward to avoid recomputation\n- backward_recompute: none (RSTD is expensive to recompute)\n- gradient_formulas:\n  - dX: rstd * (dY*(W+offset) - (1/N) * rstd^2 * dot(dY*(W+offset), X) * X)\n  - dW: sum over (B,T) of dY * (X * rstd)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: forward uses one program per row `(n_rows,)`, backward uses SM-based partitioning `(sm_count,)` with `rows_per_program`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: true (for dW reduction — each SM accumulates partial dW, then summed)\n\n## Module Parameters\n- module_init_params:\n  - hidden_size: int\n  - eps: float = 1e-6\n  - offset: float = 0.0\n  - casting_mode: str = \"llama\"\n  - init_fn: str = \"ones\"\n  - in_place: bool = True\n  - elementwise_affine: bool = True\n- learnable_params:\n  - weight: shape (hidden_size,), init ones or zeros\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {M: 4096, eps: 1e-6, dtype: torch.float32}\n\n## Key Patterns\n\n- **RSTD caching**: Forward computes and stores `rstd = rsqrt(mean(X^2) + eps)` — 1 value per row, tiny memory cost, saves 4 ops in backward\n- **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)`\n- **Casting modes as constexpr**: `casting_mode` is `tl.constexpr` so the compiler eliminates dead branches\n- **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)\n- **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)\n","isInternal":false,"tokens":812,"sizeBytes":2773},{"name":"swiglu-profile.md","path":".agents/skills/liger-kernel-dev/examples/swiglu-profile.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/examples/swiglu-profile.md","title":"Examples Skill","category":"anthropic-skill","format":"markdown","content":"# Kernel Profile: SwiGLU (Tier 1 — Element-wise)\n\n## Identity\n- operation_name: swiglu\n- function_class_name: LigerSiLUMulFunction\n- module_class_name: LigerSwiGLUMLP\n- functional_name: liger_swiglu\n\n## Classification\n- tier: 1\n- tier_description: element-wise\n- closest_existing_kernel: geglu (same structure, different activation)\n\n## Forward Pass\n- forward_inputs:\n  - a: shape (B, T, H), gate projection output\n  - b: shape (B, T, H), up projection output\n- forward_outputs:\n  - c: shape (B, T, H), silu(a) * b\n- forward_computation: c = silu(a) * b, where silu(x) = x * sigmoid(x)\n- precision_sensitive_ops: [sigmoid]\n\n## Backward Pass\n- backward_saved_tensors: [a, b] (reshaped to 2D in forward wrapper)\n- backward_recompute: recompute silu(a) and sigmoid(a) in backward\n- gradient_formulas:\n  - da: dc * (silu(a) * (1 - sigmoid(a)) + sigmoid(a)) * b\n  - db: dc * silu(a)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row, `(n_rows,)`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - config: HuggingFace model config object\n- learnable_params:\n  - gate_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - up_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - down_proj: Linear(intermediate_size, hidden_size, bias=False)\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"torch\", \"torch_compile\"]\n- benchmark_fixed_config: {BT: 4096, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Recomputation over saving**: Forward saves `a, b` but backward recomputes `sigmoid(a)` and `silu(a)` — saves memory, sigmoid is cheap\n- **In-place backward**: Writes gradients directly to `a_ptr` and `b_ptr` (the saved tensors) — saves allocation\n- **Float32 for sigmoid**: `a_row` cast to `tl.float32` before `tl.sigmoid`, result cast back via `.cast(b_row.dtype)`\n- **No intermediate allocations**: Forward kernel writes directly to output `c`; backward kernel overwrites saved `a, b`\n","isInternal":false,"tokens":600,"sizeBytes":2140},{"name":"generator.md","path":".agents/skills/liger-kernel-dev/generator.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/generator.md","title":"Liger-kernel-dev Skill","category":"anthropic-skill","format":"markdown","content":"# Generator Agent\n\nTakes a confirmed kernel profile and PyTorch reference, generates all files for a production-ready Liger kernel.\n\n## Pre-Requisites\n\nBefore generating, read the existing implementation closest to this kernel's tier:\n- Tier 1 → `src/liger_kernel/ops/swiglu.py` + `src/liger_kernel/transformers/swiglu.py`\n- Tier 2 → `src/liger_kernel/ops/rms_norm.py` + `src/liger_kernel/transformers/rms_norm.py`\n- Tier 3 → `src/liger_kernel/ops/cross_entropy.py` + `src/liger_kernel/transformers/cross_entropy.py`\n\nAlso read the corresponding test and benchmark files for the reference kernel.\n\n## Files to Generate\n\n### 1. `src/liger_kernel/ops/{kernel}.py` (NEW)\n\nThe core Triton kernel. See [templates/ops-kernel.md](templates/ops-kernel.md).\n\nKey rules:\n- `@triton.jit` for forward and backward kernels\n- `torch.autograd.Function` subclass with `@ensure_contiguous` decorator\n- Use `calculate_settings(n_cols)` for BLOCK_SIZE and num_warps\n- Cast to `tl.float32` for precision-sensitive ops (sigmoid, rsqrt, exp, log, tanh)\n- Use `tl.program_id(0).to(tl.int64)` to avoid overflow\n- Use stride parameters for memory access, not hardcoded offsets\n- Prefer recomputation over saving for backward when computation is cheap\n\n### 2. `src/liger_kernel/transformers/{kernel}.py` (NEW)\n\nThe nn.Module wrapper. See [templates/module-wrapper.md](templates/module-wrapper.md).\n\nKey rules:\n- Import from `liger_kernel.ops` (not from the submodule directly)\n- Include `extra_repr` method\n- Match the PyTorch API as closely as possible\n\n### 3. `src/liger_kernel/transformers/functional.py` (MODIFY)\n\nAdd functional API. See [templates/functional-api.md](templates/functional-api.md).\n\nTwo changes:\n- Add import of the Function class at the top (alphabetical order)\n- Add `liger_{kernel}()` function (alphabetical order among existing functions)\n\n### 4. `src/liger_kernel/ops/__init__.py` (MODIFY)\n\nAdd import of the Function class. Insert in alphabetical order among existing imports. Use `# noqa: F401` suffix.\n\n### 5. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the Module class in three locations (alphabetical order in each):\n- Direct import at top (among \"Always-safe imports\")\n- `__all__` list\n- If the Module has a `TYPE_CHECKING` import or `__getattr__` entry needed, add there too\n\n### 6. `test/transformers/test_{kernel}.py` (NEW)\n\nUnit tests. See [templates/unit-test.md](templates/unit-test.md).\n\nKey rules:\n- Include a PyTorch reference implementation (copy from Analyzer output)\n- Test both forward and backward correctness\n- Parametrize over shapes (include non-power-of-2 and small/large sizes)\n- Parametrize over dtypes with appropriate tolerances\n- Use `set_seed()`, `assert_verbose_allclose()`, `infer_device()`, `supports_bfloat16()`\n- Import from `test.utils`, NOT `liger_kernel.test.utils`\n\n### 7. `benchmark/scripts/benchmark_{kernel}.py` (NEW)\n\nBenchmarks. See [templates/benchmark.md](templates/benchmark.md).\n\nKey rules:\n- Define `_setup_{kernel}()` function returning `(input_tensor, layer)`\n- Define `bench_speed_{kernel}()` using `run_speed_benchmark`\n- Define `bench_memory_{kernel}()` using `run_memory_benchmark`\n- Compare \"liger\" vs \"torch\" (and optionally \"torch_compile\") providers\n- Use `get_benchmark_model_config`, `estimate_kernel_peak_memory`, and sweep config utilities\n- Run both speed and memory benchmarks for forward, backward, and full modes\n- **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\n\n## Modify Mode\n\nWhen modifying an existing kernel:\n1. Read the existing kernel files thoroughly\n2. Understand the current implementation patterns\n3. Make targeted changes — do not rewrite entire files\n4. Add new test cases for new functionality\n5. Ensure existing tests still cover original behavior\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order**\n- Import style: one import per line, no multi-imports\n","isInternal":false,"tokens":1022,"sizeBytes":4178},{"name":"kernel-profile-format.md","path":".agents/skills/liger-kernel-dev/kernel-profile-format.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/kernel-profile-format.md","title":"Liger-kernel-dev Skill","category":"anthropic-skill","format":"markdown","content":"# Kernel Profile Format\n\nWhen analyzing an operation, produce a profile with all of these fields.\n\n## Fields\n\n```yaml\n# Identity\noperation_name: snake_case name (e.g., \"rms_norm\", \"dyt\", \"sparsemax\")\nfunction_class_name: \"Liger{PascalCase}Function\" (e.g., \"LigerRMSNormFunction\")\nmodule_class_name: \"Liger{PascalCase}\" (e.g., \"LigerRMSNorm\")\nfunctional_name: \"liger_{snake_case}\" (e.g., \"liger_rms_norm\")\n\n# Classification\ntier: 1 | 2 | 3\ntier_description: \"element-wise\" | \"reduction\" | \"fused/complex\"\nclosest_existing_kernel: which existing Liger kernel is most similar\n\n# Forward Pass\nforward_inputs:\n  - name: X\n    shape: \"(B, T, H)\"\n    dtype: \"input dtype\"\n  # ... list all inputs with shapes\nforward_outputs:\n  - name: Y\n    shape: \"(B, T, H)\"\n    dtype: \"same as input\"\nforward_computation: brief math description (e.g., \"Y = X * rsqrt(mean(X^2) + eps) * W\")\nprecision_sensitive_ops: list ops needing float32 (e.g., [rsqrt, sigmoid, exp])\n\n# Backward Pass\nbackward_saved_tensors: what to save_for_backward (e.g., [X, W, RSTD])\nbackward_recompute: what to recompute instead of saving (e.g., \"recompute silu in backward\")\ngradient_formulas:\n  dX: \"formula or description\"\n  dW: \"formula or description\"\n\n# Tiling Strategy\ngrid_dimensions: 1D | 2D\ngrid_description: \"one program per row\" | \"one program per (row_block, col_block)\" | etc.\nblock_size_source: \"calculate_settings(n_cols)\" | \"custom\" with explanation\nneeds_sm_parallelism: true | false (for weight gradient reduction in backward)\n\n# Module Parameters\nmodule_init_params:\n  - name: hidden_size\n    type: int\n  - name: eps\n    type: float\n    default: 1e-6\n  # ... list all __init__ parameters\nlearnable_params:\n  - name: weight\n    shape: \"(hidden_size,)\"\n    init: \"ones\"\n  # ... list learnable parameters\n\n# Benchmarking\nbenchmark_variable: \"hidden_size\" | \"seq_len\" | \"vocab_size\" | etc.\nbenchmark_x_label: human-readable label for x-axis\nbenchmark_x_values_suggestion: suggested range to sweep\nbenchmark_providers: [\"liger\", \"torch\"] (add \"torch_compile\" if applicable)\nbenchmark_fixed_config: dict of fixed params (e.g., {BT: 4096, dtype: torch.bfloat16})\n```\n\n## Naming Conventions\n\n| Component | Pattern | Example |\n|-----------|---------|---------|\n| Ops file | `src/liger_kernel/ops/{operation_name}.py` | `ops/dyt.py` |\n| Transformer file | `src/liger_kernel/transformers/{operation_name}.py` | `transformers/dyt.py` |\n| Test file | `test/transformers/test_{operation_name}.py` | `test_dyt.py` |\n| Benchmark file | `benchmark/scripts/benchmark_{operation_name}.py` | `benchmark_dyt.py` |\n| Function class | `Liger{PascalCase}Function` | `LigerDyTFunction` |\n| Module class | `Liger{PascalCase}` | `LigerDyT` |\n| Functional | `liger_{snake_case}` | `liger_dyt` |\n| Forward wrapper | `{snake_case}_forward` or `liger_{snake_case}_fwd` | `liger_dyt_fwd` |\n| Backward wrapper | `{snake_case}_backward` or `liger_{snake_case}_bwd` | `liger_dyt_bwd` |\n","isInternal":false,"tokens":845,"sizeBytes":2925},{"name":"SKILL.md","path":".agents/skills/liger-kernel-dev/SKILL.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/SKILL.md","title":"liger-kernel-dev","category":"anthropic-skill","format":"markdown","content":"---\nname: liger-kernel-dev\ndescription: \"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.\"\n---\n\n# Liger Kernel Dev\n\nDevelops 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.\n\n## Mode Detection\n\n- **Create mode**: User asks to create/add/generate/write/build a new kernel → full pipeline\n- **Modify mode**: User asks to update/fix/change/extend an existing kernel → skip Analyze, modify files, then Validate\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Analyzer** workflow in [analyzer.md](analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nAccepts any input: local file, URL, code snippet, natural language description, or model component reference. Produces a standalone PyTorch reference implementation and a kernel profile.\n\n**Human checkpoint:** Present PyTorch reference + kernel profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Generator** workflow in [generator.md](generator.md).\n\nGenerates/modifies up to 8 files:\n\n1. `src/liger_kernel/ops/{kernel}.py` — NEW Triton kernels + autograd Function\n2. `src/liger_kernel/transformers/{kernel}.py` — NEW nn.Module wrapper\n3. `src/liger_kernel/transformers/functional.py` — MODIFY add functional API\n4. `src/liger_kernel/ops/__init__.py` — MODIFY export Function class\n5. `src/liger_kernel/transformers/__init__.py` — MODIFY export Module + `__all__`\n6. `test/transformers/test_{kernel}.py` — NEW unit tests\n7. `benchmark/scripts/benchmark_{kernel}.py` — NEW benchmark script\n8. `benchmark/data/all_benchmark_data.csv` — MODIFY (after benchmarks run)\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns checkstyle, unit tests (hard gate — stops on persistent failure), benchmarks, and generates plots. Optionally runs ncu profiling.\n\n**Human checkpoint:** Report final results with benchmark numbers and plots.\n\n## Pipeline (Modify Mode)\n\n1. Read existing kernel files to understand current implementation\n2. Understand the requested modification\n3. Make targeted changes (Generator handles this)\n4. Run full Validate stage (same as create mode)\n\n## Reference Files\n\n- [kernel-profile-format.md](kernel-profile-format.md) — Kernel profile schema and field descriptions\n- [examples/swiglu-profile.md](examples/swiglu-profile.md) — Tier 1 (element-wise) reference\n- [examples/rms-norm-profile.md](examples/rms-norm-profile.md) — Tier 2 (reduction) reference\n- [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md) — Tier 3 (fused/complex) reference\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n","frontmatter":{"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."},"isInternal":false,"tokens":732,"sizeBytes":3226},{"name":"benchmark.md","path":".agents/skills/liger-kernel-dev/templates/benchmark.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/templates/benchmark.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Benchmark Template\n\n## File: `benchmark/scripts/benchmark_{kernel}.py`\n\n```python\nimport os\nimport sys\n\nimport torch\n\nfrom benchmark_model_configs import compute_hidden_size_sweep_config\nfrom benchmark_model_configs import estimate_kernel_peak_memory\nfrom benchmark_model_configs import get_benchmark_model_config\nfrom utils import SingleBenchmarkRunInput\nfrom utils import SingleBenchmarkRunOutput\nfrom utils import parse_benchmark_script_args\nfrom utils import run_benchmarks\nfrom utils import run_memory_benchmark\nfrom utils import run_speed_benchmark\n\nfrom liger_kernel.utils import infer_device\n\ndevice = infer_device()\n\nsys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), \"../..\")))\n\n\ndef _setup_{kernel}(input: SingleBenchmarkRunInput):\n    \"\"\"Create input tensor and layer from benchmark config.\"\"\"\n    # Import both Liger and PyTorch reference implementations\n    from test.transformers.test_{kernel} import Liger{Kernel}\n    from test.transformers.test_{kernel} import Torch{Kernel}\n\n    cfg = input.extra_benchmark_config\n    hidden_size = input.x  # Or whatever the sweep variable is\n\n    x = torch.randn(cfg[\"BT\"], hidden_size, device=device, dtype=cfg[\"dtype\"], requires_grad=True)\n\n    if input.kernel_provider == \"liger\":\n        layer = Liger{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch\":\n        layer = Torch{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch_compile\":\n        layer = torch.compile(Torch{Kernel}(hidden_size=hidden_size).to(device))\n    else:\n        raise ValueError(f\"Invalid provider: {input.kernel_provider} for {kernel}\")\n\n    return x, layer\n\n\ndef bench_speed_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x])\n\n\ndef bench_memory_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode)\n\n\nBT = 4096  # Fixed batch*seq for sweeping hidden_size\n\n\nif __name__ == \"__main__\":\n    args = parse_benchmark_script_args()\n    model = get_benchmark_model_config(args.model)\n\n    def _probe():\n        probe_input = SingleBenchmarkRunInput(\n            x=model.hidden_size,\n            kernel_provider=\"torch\",\n            extra_benchmark_config={\"BT\": BT, \"dtype\": model.dtype},\n        )\n        x, layer = _setup_{kernel}(probe_input)\n        return layer(x)\n\n    peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)\n    sweep_config = compute_hidden_size_sweep_config(model, peak_bytes, bt=BT)\n    x_values = [1024 * i for i in range(1, 17) if 1024 * i <= sweep_config.max_hidden_size] or [model.hidden_size]\n\n    common_configs = {\n        \"kernel_name\": \"{kernel}\",\n        \"x_name\": \"hidden_size\",\n        \"x_label\": \"hidden_size\",\n        \"x_values\": x_values,\n        \"kernel_providers\": [\"liger\", \"torch\", \"torch_compile\"],\n        \"extra_benchmark_configs\": [{\"BT\": sweep_config.bt, \"dtype\": model.dtype}],\n        \"overwrite\": args.overwrite,\n    }\n\n    run_benchmarks(\n        bench_test_fn=bench_speed_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"speed\",\n        metric_unit=\"ms\",\n        **common_configs,\n    )\n    run_benchmarks(\n        bench_test_fn=bench_memory_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"memory\",\n        metric_unit=\"MB\",\n        **common_configs,\n    )\n```\n\n### Key Rules\n\n1. **Import both implementations from test file** — avoids duplicating reference code\n2. **Support \"liger\", \"torch\", and \"torch_compile\" providers**\n3. **Use `run_speed_benchmark` and `run_memory_benchmark`** from utils — don't write custom measurement\n4. **Use `get_benchmark_model_config`** and sweep utilities for standard model sizes\n5. **Use `estimate_kernel_peak_memory`** to auto-compute safe x_values range\n6. **Run both speed AND memory** for all three modes (forward, backward, full)\n7. **Use `BT = 4096`** as default batch*seq when sweeping hidden_size\n8. **Adapt sweep variable** — some kernels sweep `seq_len` or `vocab_size` instead of `hidden_size`. Use the appropriate `compute_*_sweep_config` utility.\n9. **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.\n","isInternal":false,"tokens":1066,"sizeBytes":4579},{"name":"functional-api.md","path":".agents/skills/liger-kernel-dev/templates/functional-api.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/templates/functional-api.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Functional API Template\n\n## File: `src/liger_kernel/transformers/functional.py` (MODIFY)\n\n### Step 1: Add Import\n\nAdd at the top of the file, in alphabetical order among existing imports:\n\n```python\nfrom liger_kernel.ops import Liger{Kernel}Function\n```\n\n### Step 2: Add Function\n\nAdd in alphabetical order among existing functions:\n\n```python\ndef liger_{kernel}(x, ...params):\n    return Liger{Kernel}Function.apply(x, ...params)\n```\n\nFor functions with complex return types (like CrossEntropy), wrap the output:\n\n```python\ndef liger_{kernel}(input, target, ...params):\n    result = Liger{Kernel}Function.apply(input, target, ...params)\n    return result\n```\n\n### Key Rules\n\n- Function signature should match the PyTorch equivalent where applicable\n- Use positional args for tensor inputs, keyword args for config params\n- Keep the function body minimal — just call `.apply()` and return\n","isInternal":false,"tokens":211,"sizeBytes":893},{"name":"module-wrapper.md","path":".agents/skills/liger-kernel-dev/templates/module-wrapper.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/templates/module-wrapper.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Module Wrapper Template\n\n## File: `src/liger_kernel/transformers/{kernel}.py`\n\n```python\nimport torch\nimport torch.nn as nn\n\nfrom liger_kernel.ops import Liger{Kernel}Function\n\n\nclass Liger{Kernel}(nn.Module):\n    def __init__(self, ...params):\n        super().__init__()\n        # Store config\n        self.param = param\n        # Create learnable parameters\n        self.weight = nn.Parameter(torch.ones(hidden_size))\n\n    def forward(self, x):\n        return Liger{Kernel}Function.apply(x, self.weight, ...self.params)\n\n    def extra_repr(self):\n        return f\"param={self.param}, ...\"\n```\n\n### Key Rules\n\n- Import from `liger_kernel.ops` (vendor-replaceable), not from the submodule\n- Match the PyTorch reference API as closely as possible (same __init__ params, same forward signature)\n- Include `extra_repr` for debugging\n- For loss modules, name the class `Liger{Kernel}Loss` and the forward takes `(input, target)`\n- For model-specific variants, create subclasses (see `LigerRMSNormForGemma` in `rms_norm.py`)\n","isInternal":false,"tokens":249,"sizeBytes":1022},{"name":"ops-kernel.md","path":".agents/skills/liger-kernel-dev/templates/ops-kernel.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/templates/ops-kernel.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Ops Kernel Template\n\n## File: `src/liger_kernel/ops/{kernel}.py`\n\nBefore writing, read the reference kernel for this tier in the actual codebase.\n\n### Import Block\n\n```python\nimport torch\nimport triton\nimport triton.language as tl\n\nfrom liger_kernel.ops.utils import calculate_settings\nfrom liger_kernel.ops.utils import ensure_contiguous\n```\n\nAdd additional imports only as needed:\n- `math` — if using `math.ceil` for row partitioning\n- `from liger_kernel.ops.utils import compare_version` — if checking triton version\n- `from liger_kernel.ops.utils import torch_to_triton_dtype` — if passing dtype as kernel arg\n- `rsqrt` from triton libdevice — if using inverse square root (see `rms_norm.py` for the import pattern with version guards)\n- `tanh` from triton libdevice — if using tanh (see `dyt.py` for the import pattern)\n\n### Forward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_forward_kernel(\n    # Output pointers first, then input pointers\n    Y_ptr, Y_stride,\n    X_ptr, X_stride,\n    # ... other input pointers and strides\n    n_cols,\n    # constexpr params for compile-time optimization\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Calculate base pointers using strides\n    X_ptr += row_idx * X_stride\n    Y_ptr += row_idx * Y_stride\n\n    # Load with mask and float32 upcast for precision\n    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0).to(tl.float32)\n\n    # === FORWARD COMPUTATION ===\n\n    tl.store(Y_ptr + col_offsets, result, mask=mask)\n```\n\n### Backward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_backward_kernel(\n    # Gradient output, then gradient inputs, then saved tensors\n    dY_ptr, dY_stride,\n    dX_ptr, dX_stride,\n    # ... saved tensors from forward\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Load gradient and saved tensors\n    # === BACKWARD COMPUTATION ===\n\n    tl.store(dX_ptr + row_idx * dX_stride + col_offsets, dX_row, mask=mask)\n```\n\n### Python Wrappers\n\n```python\ndef {kernel}_forward(X, ...):\n    ori_shape = X.shape\n    n_cols = ori_shape[-1]\n    X = X.view(-1, n_cols)\n    n_rows = X.shape[0]\n\n    Y = torch.empty_like(X)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_forward_kernel[(n_rows,)](\n        Y, Y.stride(-2),\n        X, X.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return Y.view(*ori_shape)\n\n\ndef {kernel}_backward(dY, ...saved):\n    ori_shape = dY.shape\n    n_cols = ori_shape[-1]\n    dY = dY.view(-1, n_cols)\n    n_rows = dY.shape[0]\n\n    dX = torch.empty_like(dY)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_backward_kernel[(n_rows,)](\n        dY, dY.stride(-2),\n        dX, dX.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return dX.view(*ori_shape)\n```\n\n### Autograd Function\n\n```python\nclass Liger{Kernel}Function(torch.autograd.Function):\n    @staticmethod\n    @ensure_contiguous\n    def forward(ctx, X, ...params):\n        Y = {kernel}_forward(X, ...params)\n        ctx.save_for_backward(...)  # Save only what backward needs\n        return Y\n\n    @staticmethod\n    @ensure_contiguous\n    def backward(ctx, dY):\n        ...saved = ctx.saved_tensors\n        dX = {kernel}_backward(dY, ...saved)\n        # Return one gradient per forward input, None for non-tensor params\n        return dX, None, None, ...\n```\n\n### Tier-Specific Patterns\n\n**Tier 1 (element-wise):**\n- Grid: `(n_rows,)` — one program per row\n- Simple 1D indexing with `col_offsets` and `mask`\n- Prefer recomputation in backward over saving activations\n- See `swiglu.py`: saves `a, b` but recomputes `silu(a)` in backward\n\n**Tier 2 (reduction):**\n- Forward may cache intermediate state (e.g., RSTD in RMSNorm)\n- Backward may use SM-based parallelism for weight gradient reduction:\n  ```python\n  sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count\n  _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)\n  grid = (sm_count,)\n  rows_per_program = math.ceil(n_rows / sm_count)\n  ```\n- See `rms_norm.py` for the full pattern\n\n**Tier 3 (fused/complex):**\n- Multi-pass algorithms (e.g., online softmax: pass 1 finds max+sum, pass 2 computes gradients)\n- Gradient-in-forward trick: compute gradients during forward, store in input tensor\n- Multiple `tl.constexpr` flags for compile-time dead code elimination\n- See `cross_entropy.py` for the full pattern\n\n### Key Rules\n\n1. Always use `tl.program_id(0).to(tl.int64)` (prevents overflow on large tensors)\n2. Always use `mask = col_offsets < n_cols` for boundary handling\n3. Use `tl.constexpr` for parameters that enable compile-time optimization\n4. Use stride parameters from tensors, never hardcode memory layout\n5. Cast precision-sensitive ops to `tl.float32`, cast back for storage\n6. Return `None` for each non-tensor parameter in backward\n","isInternal":false,"tokens":1359,"sizeBytes":5094},{"name":"unit-test.md","path":".agents/skills/liger-kernel-dev/templates/unit-test.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/templates/unit-test.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Unit Test Template\n\n## File: `test/transformers/test_{kernel}.py`\n\n```python\nimport pytest\nimport torch\n\nfrom test.utils import assert_verbose_allclose\nfrom test.utils import infer_device\nfrom test.utils import set_seed\nfrom test.utils import supports_bfloat16\n\nfrom liger_kernel.ops import Liger{Kernel}Function\nfrom liger_kernel.transformers.{kernel} import Liger{Kernel}\nfrom liger_kernel.transformers.functional import liger_{kernel}\n\n\n# ---- PyTorch Reference ----\n# Paste the standalone PyTorch reference from the Analyzer stage here.\n\nclass Torch{Kernel}(torch.nn.Module):\n    def __init__(self, ...):\n        super().__init__()\n        ...\n\n    def forward(self, x):\n        ...\n\n\nset_seed(42)\ndevice = infer_device()\n\n\n# ---- Correctness Test (Module) ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (4, 16, 2048),\n        (1, 1, 1023),   # Non-power-of-2\n        (3, 7, 256),    # Small + prime numbers\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-1,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_correctness(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Initialize both with same weights\n    torch_layer = Torch{Kernel}(...).to(device).to(dtype)\n    liger_layer = Liger{Kernel}(...).to(device).to(dtype)\n    # Copy weights from torch to liger\n    # liger_layer.weight.data = torch_layer.weight.data.clone()\n\n    # Forward\n    torch_output = torch_layer(x1)\n    liger_output = liger_layer(x2)\n    assert_verbose_allclose(torch_output, liger_output, rtol=rtol, atol=atol, extra_info=\"[output]\")\n\n    # Backward\n    grad_output = torch.randn_like(_input)\n    torch_output.backward(grad_output)\n    liger_output.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol, extra_info=\"[input.grad]\")\n\n    # Check weight gradients if applicable\n    # assert_verbose_allclose(\n    #     torch_layer.weight.grad, liger_layer.weight.grad,\n    #     rtol=rtol, atol=atol, extra_info=\"[weight.grad]\"\n    # )\n\n\n# ---- Functional API Test ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (1, 1, 1023),\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-8,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_functional(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Test that functional and Function.apply produce identical results\n    output1 = liger_{kernel}(x1, ...)\n    output2 = Liger{Kernel}Function.apply(x2, ...)\n\n    assert_verbose_allclose(output1, output2, rtol=rtol, atol=atol)\n\n    grad_output = torch.randn_like(_input)\n    output1.backward(grad_output)\n    output2.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol)\n```\n\n### Key Testing Rules\n\n1. **Always test both forward AND backward** — backward bugs are common\n2. **Include non-power-of-2 shapes** — tests the mask boundary handling\n3. **Parametrize dtypes with per-dtype tolerances**:\n   - `float32`: `atol=1e-5, rtol=1e-5` (tight)\n   - `bfloat16`: `atol=1e-1, rtol=5e-2` (loose — bf16 has limited precision)\n4. **Use `set_seed(42)` at module level** for reproducibility\n5. **Use `infer_device()`** not `\"cuda\"` — supports multi-vendor\n6. **Guard bfloat16 with `supports_bfloat16()`** — skips on older GPUs\n7. **Use `assert_verbose_allclose`** not `torch.testing.assert_close` — gives better error output\n8. **Test weight gradients** when the kernel has learnable parameters\n9. **Add `@pytest.mark.flaky(reruns=3)`** if the test is sensitive to random seeds\n10. **Add extra test functions** for edge cases specific to the kernel\n","isInternal":false,"tokens":1219,"sizeBytes":4471},{"name":"validator.md","path":".agents/skills/liger-kernel-dev/validator.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-dev/validator.md","title":"Liger-kernel-dev Skill","category":"anthropic-skill","format":"markdown","content":"# Validator Agent\n\nValidates the generated kernel through checkstyle, testing, benchmarking, and optional profiling.\n\n## Validation Flow\n\n### Step 1: Checkstyle (Fast Gate)\n\n```bash\nmake checkstyle\n```\n\nIf it fails:\n```bash\nruff check . --fix && ruff format .\nmake checkstyle\n```\n\nIf it still fails after auto-fix, identify and fix the remaining issues manually, then re-run.\n\n### Step 2: Unit Tests (Hard Gate)\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\nIf tests fail:\n1. Read the error output carefully\n2. Identify root cause: kernel logic bug, shape mismatch, dtype issue, or tolerance too tight\n3. Fix the issue in the relevant file(s)\n4. Re-run `make checkstyle` (in case fix introduced style issues)\n5. Re-run the failing test\n\n**Retry up to 3 times total.** If tests still fail after 3 attempts:\n- **STOP IMMEDIATELY. Do NOT proceed to benchmarks.**\n- Report to the user:\n  - Which tests are failing and the exact error messages\n  - What fixes were attempted\n  - The current state of the code\n  - Suggest what might be wrong and ask for guidance\n\n### Step 3: Benchmarks\n\nRun the benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nThis produces speed and memory measurements for forward, backward, and full modes, comparing Liger vs PyTorch baseline.\n\n### Step 4: Generate Plots\n\n```bash\npython benchmark/benchmarks_visualizer.py\n```\n\nPlots are saved to `benchmark/visualizations/`. If the visualizer fails or is not available, this step is non-blocking — report the raw numbers instead.\n\n### Step 5: Optional ncu Profiling\n\nOnly if the user explicitly requests it. Requires NVIDIA Nsight Compute (`ncu`) to be installed.\n\nProfile both the Liger kernel and PyTorch reference:\n```bash\nncu --set full --target-processes all -o liger_profile python -c \"\nimport torch\nfrom liger_kernel.ops.{kernel} import Liger{Kernel}Function\nx = torch.randn(..., device='cuda', dtype=torch.float32, requires_grad=True)\nout = Liger{Kernel}Function.apply(x, ...)\nout.backward(torch.randn_like(out))\n\"\n```\n\nReport key metrics: SM occupancy, memory throughput, compute throughput.\n\nIf ncu is not available, report this and skip.\n\n### Step 6: Report Results\n\nPresent a clear summary to the user:\n\n```\n## Results\n\n**Checkstyle:** PASS\n**Unit Tests:** PASS (X tests passed)\n\n**Speed Benchmarks:**\n| Mode     | Liger (ms) | PyTorch (ms) | Speedup |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...x    |\n| Backward | ...       | ...         | ...x    |\n| Full     | ...       | ...         | ...x    |\n\n**Memory Benchmarks:**\n| Mode     | Liger (MB) | PyTorch (MB) | Savings |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...%    |\n| Backward | ...       | ...         | ...%    |\n| Full     | ...       | ...         | ...%    |\n\n**Assessment:** [Which dimension improved — speed, memory, or both.\nFlag if either metric is catastrophically worse.]\n\n**Plots:** benchmark/visualizations/\n\n**ncu Profiling:** [Results if requested, or \"Not requested\"]\n```\n","isInternal":false,"tokens":777,"sizeBytes":3085},{"name":"finalizer.md","path":".agents/skills/liger-kernel-perf/finalizer.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-perf/finalizer.md","title":"Liger-kernel-perf Skill","category":"anthropic-skill","format":"markdown","content":"# Finalizer Agent\n\nApplies the winning optimization variant to the production kernel, validates correctness through the full test suite and checkstyle, and produces the final optimization report.\n\nThis is Stage 3 of the 3-stage pipeline. It receives the winning variant from Stage 2 (Optimizer) and is responsible for ensuring the optimization is safe to ship.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `winning_variant`: variant number chosen by user (interactive) or by the Optimizer (autonomous)\n- `optimization_workspace`: path to `optimization/{kernel}/` containing all variant files, benchmarks, and the profile\n\n## Expected Workspace State\n\nBefore the Finalizer starts, the following files MUST exist (created by earlier stages):\n\n```\noptimization/{kernel}/\n  original_{kernel}.py          # Snapshot of the original kernel (from Profiler)\n  profile.md                    # Optimization profile (from Profiler)\n  {kernel}_v1.py                # Variant source files (from Optimizer)\n  {kernel}_v1_notes.md          # Variant lab notebooks (from Optimizer)\n  {kernel}_v2.py                # ... more variants\n  {kernel}_v2_notes.md\n  benchmarks/\n    v0_baseline.csv             # Baseline benchmark data\n    v1_results.csv              # Per-variant benchmark data\n    v2_results.csv\n```\n\nIf any required files are missing, report the missing files and stop.\n\n## Steps\n\n### Step 1: Validate Inputs and Read Context\n\n1. Confirm the winning variant file exists: `optimization/{kernel}/{kernel}_v{N}.py`\n2. Read the winning variant code\n3. Read the original kernel: `optimization/{kernel}/original_{kernel}.py`\n4. Read the optimization profile: `optimization/{kernel}/profile.md`\n5. Read ALL variant notes files (`{kernel}_v*_notes.md`) to collect data for the report\n6. Read ALL benchmark CSVs (`benchmarks/v*_results.csv`) to build the comparison table\n7. Read the current production kernel: `src/liger_kernel/ops/{kernel}.py`\n\nVerify the original kernel snapshot matches the current production kernel. If they differ (someone else modified the file during the optimization), warn the user and confirm before proceeding.\n\n### Step 2: Apply the Winning Variant\n\nReplace the production kernel with the winning variant:\n\n```bash\ncp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis is a destructive operation. The original is preserved at `optimization/{kernel}/original_{kernel}.py`.\n\n### Step 3: Run the Full Test Suite (HARD GATE)\n\nRun the complete test file for the kernel:\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\n**This is a HARD GATE.** The test suite must pass completely.\n\n#### If tests PASS:\nProceed to Step 4.\n\n#### If tests FAIL:\nExecute the test failure recovery procedure:\n\n**Attempt 1: Diagnose and fix**\n1. Read the full test error output carefully\n2. Identify the root cause (numerical tolerance issue, shape mismatch, missing import, API change, etc.)\n3. Apply a targeted fix to `src/liger_kernel/ops/{kernel}.py`\n4. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**Attempt 2: Second fix attempt (if Attempt 1 still fails)**\n1. Read the new error output\n2. Try a different fix approach\n3. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**If both attempts fail: REVERT**\n1. Restore the original kernel:\n   ```bash\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n2. Verify the revert by running the test suite again:\n   ```bash\n   python -m pytest test/transformers/test_{kernel}.py -xvs\n   ```\n3. Set `optimization_applied = false`\n4. Record the failure reason and all error outputs\n5. Proceed to Step 4 to generate the report (the report documents what was attempted even if it could not be applied)\n\n**Do NOT proceed to finalize (Step 5) if the optimization was reverted. The report must clearly state the optimization was not applied.**\n\n### Step 4: Run Checkstyle\n\nRun the project's checkstyle validation:\n\n```bash\nmake checkstyle\n```\n\n#### If checkstyle PASSES:\nProceed to Step 5.\n\n#### If checkstyle FAILS:\nAuto-fix with ruff:\n\n```bash\nruff check . --fix && ruff format .\n```\n\nRetry checkstyle once:\n\n```bash\nmake checkstyle\n```\n\nIf it still fails after auto-fix, report the remaining checkstyle issues but do NOT revert the kernel (checkstyle is not a correctness gate). Note the issues in the report so they can be fixed manually.\n\n### Step 5: Generate Comparison Plots\n\nGenerate 3-way comparison plots showing the original Liger kernel, the optimized kernel, and the HuggingFace/PyTorch baseline. This gives the user a visual demonstration of the improvement.\n\nThe procedure has 3 phases: (1) benchmark both kernels for plots, (2) generate plots, (3) clean the CSV so only the optimized kernel's data remains permanently.\n\n#### Step 5a: Benchmark Both Kernels for Plotting\n\n1. **Benchmark the original kernel as \"liger_original\"**: Temporarily restore the original kernel and run benchmarks with a custom provider name:\n   ```bash\n   # Swap in original kernel\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Edit benchmark script: replace \"liger\" with \"liger_original\" in kernel_providers\n   sed -i 's/\"liger\"/\"liger_original\"/g' benchmark/scripts/benchmark_{kernel}.py\n   \n   # Run benchmarks (appends \"liger_original\" rows to CSV)\n   cd benchmark/scripts && python benchmark_{kernel}.py\n   \n   # Revert the benchmark script change\n   git checkout benchmark/scripts/benchmark_{kernel}.py\n   ```\n\n2. **Benchmark the optimized kernel as \"liger\"**: Re-apply the optimized kernel and run benchmarks normally. Use `--overwrite` so the \"liger\" rows reflect the optimized kernel:\n   ```bash\n   # Swap in optimized kernel\n   cp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Run benchmarks (writes \"liger\" rows with optimized kernel)\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   ```\n\nAt this point the CSV has 3 providers: `\"liger_original\"`, `\"liger\"` (optimized), and `\"huggingface\"`/`\"torch\"`. This is what we need for the 3-way comparison plots.\n\n#### Step 5b: Generate Plots\n\nGenerate plots while all 3 providers are in the CSV:\n\n```bash\n# Speed plots for all modes (forward, backward, full)\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name speed \\\n  --overwrite\n\n# Memory plot\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name memory \\\n  --overwrite\n```\n\nCopy plots to the optimization workspace:\n\n```bash\ncp benchmark/visualizations/{kernel}_speed_*.png optimization/{kernel}/\ncp benchmark/visualizations/{kernel}_memory_*.png optimization/{kernel}/\n```\n\n#### Step 5c: Clean CSV for Permanent State\n\nNow remove the temporary `\"liger_original\"` rows and ensure the CSV only contains clean data for the optimized kernel. Delete ALL rows for this kernel, then re-run the benchmark to append fresh data:\n\n```bash\n# Delete all existing rows for this kernel from CSV (keep header + other kernels)\nhead -1 benchmark/data/all_benchmark_data.csv > /tmp/benchmark_header.csv\ngrep -v \"^{kernel},\" benchmark/data/all_benchmark_data.csv > /tmp/benchmark_other.csv || true\ncat /tmp/benchmark_header.csv <(tail -n +2 /tmp/benchmark_other.csv) > benchmark/data/all_benchmark_data.csv\n\n# Re-run benchmarks with optimized kernel (appends fresh rows)\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nAfter this, `all_benchmark_data.csv` contains:\n- `\"liger\"` — the **optimized** kernel (clean fresh data)\n- `\"huggingface\"` / `\"torch\"` — the baseline (also fresh)\n- NO `\"liger_original\"` — that was only needed for the plots\n\n#### Step 5c: Generate Additional Custom Plots (Optional)\n\nIf the optimization shows interesting size-dependent behavior (e.g., bigger gains at larger sizes), generate targeted plots highlighting this. Use matplotlib directly if the standard visualizer does not cover the comparison you want to show.\n\nSave any additional plots to `optimization/{kernel}/`.\n\n### Step 6: Generate the Final Optimization Report\n\nWrite the report to `optimization/{kernel}/report.md`. The report is the permanent record of the entire optimization effort.\n\nThe report must include these sections:\n\n1. **Summary**: kernel name, GPU, goal, variants tried, best variant, improvement %, and whether the optimization was applied (`Yes` / `No — reverted due to test failures`)\n2. **Baseline Profile**: bottleneck classification with evidence, NCU highlights if available\n3. **Variant Comparison**: include the table from `optimization/{kernel}/comparison.md` (produced by the Optimizer)\n4. **Winning Changes**: what was changed, why it worked, and a key code diff. Run `diff optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py` — do NOT write the diff from memory.\n5. **Rejected Approaches**: for each rejected variant, what was tried and why it failed\n6. **Comparison Plots**: reference the plots generated in Step 5 with file paths\n7. **Correctness Verification**: smoke test results, full test suite results, checkstyle results\n8. **Recommendations**: unexplored optimization opportunities, suggested follow-ups\n9. **Reproduction** section:\n   ```bash\n   # Reproduce baseline benchmarks:\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   # Revert to original:\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n\nInclude the variant comparison data from `optimization/{kernel}/comparison.md` (produced by the Optimizer).\n\n### Report Quality Checklist\n\nBefore saving the report, verify:\n\n- [ ] All variants are listed in the comparison table (including rejected ones)\n- [ ] Speed ratios are calculated consistently (higher = faster)\n- [ ] Memory ratios are calculated consistently (higher = less memory used)\n- [ ] The winning variant is clearly marked\n- [ ] The bottleneck classification matches the profile from Stage 1\n- [ ] If optimization was reverted, this is stated in the Summary, the header, and Section 6\n- [ ] Reproduction commands use the correct file paths\n- [ ] The code diff in Section 4 is accurate (diff the actual files, do not write from memory)\n- [ ] Comparison plots are generated and referenced in the report\n\n### Step 7: Create Pull Request\n\nCreate a PR with the kernel code changes and the updated benchmark data. Do NOT include plots or optimization workspace files in the PR.\n\n#### Step 7a: Stage Only Relevant Files\n\n```bash\ngit add src/liger_kernel/ops/{kernel}.py\ngit add benchmark/data/all_benchmark_data.csv\n```\n\nDo NOT add:\n- `optimization/` directory (workspace files, notes, plots)\n- `benchmark/visualizations/` (generated plots -- these are for local review only)\n\n#### Step 7b: Create a Descriptive Commit\n\n```bash\ngit commit -m \"$(cat <<'EOF'\n[Perf] Optimize {kernel} kernel: {headline improvement}\n\n{2-3 sentence summary of the optimization technique and its impact.}\n\nKey changes:\n- {Change 1}: {what and why}\n- {Change 2}: {what and why}\n\nBenchmark results ({gpu_name}):\n- Speed (forward): {delta}% {faster/slower}\n- Speed (backward): {delta}% {faster/slower}\n- Speed (full): {delta}% {faster/slower}\n- Memory: {delta}% {reduction/increase}\nEOF\n)\"\n```\n\nIf the host AI runtime has its own commit-attribution convention (e.g., a `Co-Authored-By` trailer), append it per that runtime's guidelines.\n\n#### Step 7c: Push and Create PR\n\n```bash\ngit push -u origin HEAD\n```\n\nCreate the PR using `gh pr create`. The PR body must clearly explain:\n\n1. **What was optimized** -- which kernel, which pass, what bottleneck\n2. **What changes were made** -- specific code changes with rationale for each\n3. **How they help** -- tie back to the bottleneck diagnosis (e.g., \"reduced register pressure from 115 to 95 registers/thread, increasing occupancy from 12.5% to 25%\")\n4. **Benchmark results** -- before/after table with speed and memory numbers at representative sizes\n5. **Correctness** -- confirm all tests pass\n\nUse this PR template:\n\n```bash\ngh pr create --title \"[Perf] Optimize {kernel}: {headline}\" --body \"$(cat <<'EOF'\n## Summary\n\n{2-3 sentences: what was optimized, the key technique, and the headline result.}\n\n## Bottleneck Diagnosis\n\n- **Kernel**: `{kernel}` ({tier} — {tier_description})\n- **GPU**: {gpu_name} ({architecture})\n- **Classification**: {memory-bound / compute-bound / balanced}\n- **Root cause**: {specific bottleneck explanation}\n\n## Changes\n\n{For each change, explain WHAT was changed and WHY it helps:}\n\n### 1. {Change title}\n{Description of the code change and its performance rationale.}\n\n### 2. {Change title}\n{Description of the code change and its performance rationale.}\n\n## Benchmark Results\n\nTested on {gpu_name}. Values are median ms (speed) or MB (memory).\n\n### Speed (ms, lower is better)\n\n| x_value | Mode | Before | After | Improvement |\n|---------|------|--------|-------|-------------|\n| {x1} | forward | {before} | {after} | {delta}% |\n| {x1} | backward | {before} | {after} | {delta}% |\n| {x1} | full | {before} | {after} | {delta}% |\n| ... | ... | ... | ... | ... |\n\n### Memory (MB, lower is better)\n\n| x_value | Before | After | Improvement |\n|---------|--------|-------|-------------|\n| {x1} | {before} | {after} | {delta}% |\n| ... | ... | ... | ... |\n\n## Correctness\n\n- Full test suite: **PASSED** (`python -m pytest test/transformers/test_{kernel}.py -xvs`)\n- Checkstyle: **PASSED** (`make checkstyle`)\n\n## Test Plan\n\n- [x] All existing unit tests pass\n- [x] Benchmarks show improvement across all input sizes\n- [x] No regression on non-target metrics (speed/memory balance maintained)\nEOF\n)\"\n```\n\nIf the host AI runtime appends a \"Generated with\" footer (e.g., Claude Code, Cursor, Copilot), include it per that runtime's guidelines.\n\n**Important**: Do NOT include plots as image attachments in the PR. Plots are for local review only and live in the optimization workspace.\n\n### Step 8: Present the Before/After Summary\n\nInclude the plot file paths in the summary so the user can view them.\n\nDisplay a concise summary to the user:\n\n```\n=== Optimization Complete: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: {APPLIED / REVERTED}\nWinner: v{N} — {strategy}\n\nBefore → After (at largest x_value):\n  Speed (forward):  {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (backward): {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (full):     {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Memory:           {baseline_mb} MB → {optimized_mb} MB ({improvement}% reduction)\n\nVariants tried: {N}\nTest suite: PASSED ({N} tests)\nCheckstyle: PASSED\n\nReport: optimization/{kernel}/report.md\n```\n\nIf the optimization was reverted:\n\n```\n=== Optimization REVERTED: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: REVERTED — test suite failed after applying winner\n\nWinner candidate: v{N} — {strategy}\nFailure reason: {brief error description}\nFix attempts: 2 (both failed)\n\nThe original kernel has been restored. No changes were made to production code.\n\nVariants tried: {N} (see report for full comparison)\nReport: optimization/{kernel}/report.md\n```\n\n## Error Handling\n\n| Error | Action |\n|-------|--------|\n| Winning variant file not found | Report missing file, stop |\n| Original kernel snapshot missing | Report missing file, stop (cannot safely revert) |\n| Production kernel differs from snapshot | Warn user, ask for confirmation (interactive) or stop (autonomous) |\n| Test suite fails (after 2 fix attempts) | Revert to original, generate report documenting failure |\n| Checkstyle fails (after auto-fix) | Note in report, do NOT revert (not a correctness issue) |\n| Benchmark CSVs missing | Generate report with available data, note missing data |\n| Cannot write report | Print report contents to stdout as fallback |\n\n## Key Principles\n\n1. **Correctness is non-negotiable.** The full test suite is a hard gate. Never skip it, never ignore failures, never mark a failing optimization as successful.\n2. **Always preserve the original.** The snapshot at `optimization/{kernel}/original_{kernel}.py` is the safety net. Never delete or overwrite it.\n3. **The report is permanent.** Even if the optimization is reverted, the report documents everything that was tried. This prevents future engineers from repeating failed approaches.\n4. **Be honest in the report.** Report actual measured numbers. Do not round favorably. If a variant was only 0.3% faster, say 0.3%, not \"marginal improvement.\" If the optimization was reverted, say so prominently.\n5. **Diff, don't recall.** When writing the code diff in the report, actually diff the files. Do not write the diff from memory -- it will contain errors.\n","isInternal":false,"tokens":4036,"sizeBytes":16757},{"name":"optimization-strategies.md","path":".agents/skills/liger-kernel-perf/optimization-strategies.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-perf/optimization-strategies.md","title":"Liger-kernel-perf Skill","category":"anthropic-skill","format":"markdown","content":"# Optimization Strategies Catalog\n\nReference catalog for the Optimizer Agent. Techniques are ordered by recommended trial order within each category. The agent should **always start with parameter tuning**, then apply diagnosis-driven techniques based on the bottleneck classification from the Profiler stage.\n\n---\n\n## Decision Framework\n\nAfter reading the optimization profile, follow this priority:\n\n1. **Always first**: Parameter Tuning (Section 1)\n2. **Memory-bound kernel**: Section 2 techniques, then Section 5\n3. **Compute-bound kernel**: Section 3 techniques, then Section 5\n4. **Balanced kernel**: Interleave Section 2 and Section 3, then Section 5\n5. **Architecture-specific**: Layer Section 4 techniques on top of the best variant when targeting a specific GPU\n\n---\n\n## 1. Always First: Parameter Tuning\n\n**Why first**: Parameter tuning is the cheapest technique (minimal code changes) and frequently delivers the largest single improvement. The current codebase uses `calculate_settings()` which picks BLOCK_SIZE as the next power of 2 and num_warps via a fixed heuristic. This is a reasonable default but leaves significant performance on the table for specific input shapes and GPU architectures.\n\n### Why NOT @triton.autotune\n\nLiger production kernels **cannot use `@triton.autotune`** for three architectural reasons:\n\n1. **Forward-backward context coupling**: Liger kernels use `torch.autograd.Function` where `forward()` computes `BLOCK_SIZE`/`num_warps` via `calculate_settings()`, stores them in `ctx`, and `backward()` retrieves them to launch its kernel with identical config. With `@triton.autotune`, these become compile-time constants managed by Triton -- they cannot be stored in `ctx`, and the backward kernel has no way to know which config the forward kernel auto-selected for a given input.\n\n2. **Multi-backend incompatibility**: NPU/Ascend does not support `num_warps` or `num_stages` at all (commit `8b965b9` explicitly removed them). `@triton.autotune` relies on these parameters.\n\n3. **Inference warmup latency**: Each new input shape triggers autotuning (benchmarks all configs). Adds 100ms+ first-call penalty per unique shape, unacceptable for LLM inference with variable sequence lengths.\n\nThe commented-out `@triton.autotune` decorators in `dyt.py` and `grpo_loss.py` are evidence of prior exploration that was abandoned for these reasons. Do NOT re-introduce `@triton.autotune` on production kernels.\n\n### Parameters to Tune\n\n| Parameter | Description | Typical Range |\n|-----------|-------------|---------------|\n| `BLOCK_SIZE` | Elements processed per program instance | 256, 512, 1024, 2048, 4096, 8192, 16384 |\n| `num_warps` | Warps per program instance (32 threads each) | 2, 4, 8, 16, 32 |\n| `num_stages` | Software pipelining depth | 1, 2, 3, 4, 5 |\n| `BLOCK_ROW` | Rows per program (block-row mode only) | 1, 2, 4, 8, 16, 32 |\n\n### Method: Manual Sweep with Improved Heuristic\n\nReplace `calculate_settings()` with a better heuristic derived from benchmarking. The variant must still pass `BLOCK_SIZE` and `num_warps` through `ctx` to the backward kernel.\n\n**Step 1**: Identify the current config for the benchmark's input sizes:\n```python\n# calculate_settings() returns (next_power_of_2(n), num_warps_heuristic)\n# For n_cols=4096: BLOCK_SIZE=4096, num_warps=8\n# For n_cols=8192: BLOCK_SIZE=8192, num_warps=16\n# Note: num_stages is never set (Triton defaults to 1)\n```\n\n**Step 2**: Create a variant with an improved selection function:\n```python\ndef optimized_settings(n):\n    \"\"\"Improved heuristic derived from benchmarking on {gpu_arch}.\"\"\"\n    BLOCK_SIZE = triton.next_power_of_2(n)\n    if BLOCK_SIZE > 65536:\n        raise RuntimeError(...)\n    # Tuned values (from benchmark sweep):\n    if BLOCK_SIZE >= 8192:\n        num_warps = 16\n        num_stages = 3   # NEW: calculate_settings doesn't set this\n    elif BLOCK_SIZE >= 2048:\n        num_warps = 8\n        num_stages = 2\n    else:\n        num_warps = 4\n        num_stages = 2\n    return BLOCK_SIZE, num_warps, num_stages\n```\n\n**Step 3**: Update the kernel launcher to pass `num_stages`:\n```python\n# The kernel launch gains num_stages (previously unset, defaulted to 1):\nkernel[(grid,)](..., BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages)\n```\n\n**Step 4**: Update `ctx` to also store `num_stages` and pass it to the backward kernel.\n\n**Key insight**: `num_stages` is the most commonly overlooked parameter. `calculate_settings()` does not set it at all, so Triton defaults to `num_stages=1`. Simply adding `num_stages=2` or `num_stages=3` to existing kernel launches is often the single biggest win, especially on Hopper.\n\n### Architecture-Specific Guidance\n\n- **Ampere (SM 8.x)**: `num_stages=2` is usually optimal. Higher stages rarely help and can increase register pressure. `num_warps` sweet spot is 4-16.\n- **Hopper (SM 9.0)**: Benefits significantly from higher `num_stages` (3-5) due to TMA and async pipeline support. Also tolerates `num_warps=32` better due to deeper warp scheduling.\n- **Blackwell (SM 10.0)**: Similar to Hopper but with even more pipeline depth tolerance. Start with `num_stages=4-5`.\n\n### When to Apply\n\nAlways. This is the mandatory first optimization attempt for every kernel.\n\n### Expected Impact\n\n5-40% speedup depending on how far the `calculate_settings()` heuristic is from optimal for the target shapes and GPU. Adding `num_stages` alone can yield 10-25% on Hopper.\n\n### Risks/Gotchas\n\n- Configs are shape-dependent. A config optimal for `n_cols=4096` may be suboptimal for `n_cols=128`. Always benchmark across ALL x_values from the existing benchmark suite.\n- If BLOCK_SIZE exceeds `n_cols`, the kernel still works (masked loads) but wastes warps. Prune configs where `BLOCK_SIZE >> n_cols`.\n- The `calculate_settings()` function enforces `MAX_FUSED_SIZE = 65536`. This limit exists to prevent register spilling. Respect this unless verified otherwise.\n- When adding `num_stages` to existing kernels, the backward kernel must also receive the same `num_stages` value via `ctx`.\n- On NPU/Ascend, `num_warps` and `num_stages` are ignored. Optimizations targeting these parameters are NVIDIA-specific.\n\n### Kernel Tiers\n\nAll tiers benefit. Tier 1 kernels often see the largest relative gains because their simple structure makes the config choice proportionally more impactful.\n\n---\n\n## 2. Memory-Bound Optimizations\n\nApply these when the profiler classifies the kernel as memory-bound (high memory throughput, low compute utilization, simple operations per load/store).\n\n### 2.1 Memory Access Coalescing\n\n**When to apply**: NCU shows low memory throughput efficiency, or kernel uses non-contiguous access patterns (e.g., strided access across rows, column-major access on row-major data).\n\n**How to implement**:\n\nEnsure that threads in a warp access adjacent memory addresses. In Triton, this means the innermost offset computation should produce contiguous addresses across the `tl.arange` dimension.\n\n```python\n# GOOD: contiguous access within a warp\ncol_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_idx * stride + col_offsets, mask=col_offsets < n_cols)\n\n# BAD: strided access (each thread hits a different cache line)\nrow_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_offsets * stride + col_idx)\n```\n\nIn Liger kernels, most access is already contiguous via `col_offsets = tl.arange(0, BLOCK_SIZE)`. Check for cases where stride-based pointer arithmetic breaks coalescing, especially in backward kernels or multi-dimensional grids.\n\n**Expected impact**: 10-50% when fixing uncoalesced access. Minimal gain if access is already coalesced.\n\n**Risks**: Changing access patterns may require transposing data or restructuring the grid, which can be invasive.\n\n**Kernel tiers**: All tiers, but Tier 2/3 backward passes most often have coalescing issues due to accumulation patterns.\n\n### 2.2 Cache Modifiers\n\n**When to apply**: Kernel reads data that is used exactly once (stream through), or reads data that is reused many times (should be cached aggressively). NCU shows poor L2 hit rate.\n\n**How to implement**:\n\nTriton supports cache modifiers via the `eviction_policy` parameter on `tl.load` and `tl.store`:\n\n```python\n# Data read once, don't pollute cache (streaming access)\nX_block = tl.load(X_ptr + offsets, eviction_policy=\"evict_first\")\n\n# Data reused across iterations, keep in cache\nW_row = tl.load(W_ptr + offsets, eviction_policy=\"evict_last\")\n\n# Default (let hardware decide)\nX_block = tl.load(X_ptr + offsets)  # no modifier\n```\n\n| Modifier | Triton API | When to Use |\n|----------|------------|-------------|\n| `.ca` (cache all) | `eviction_policy=\"evict_last\"` | Data reused across loop iterations or by other programs |\n| `.cs` (cache streaming) | `eviction_policy=\"evict_first\"` | Data used once; large tensors that would thrash cache |\n| `.wb` (write-back) | default store behavior | Standard write-back through cache hierarchy |\n| `.wt` (write-through) | not directly exposed | Write-through (rarely needed in Triton) |\n\n**Expected impact**: 5-20% on memory-bound kernels with predictable access patterns.\n\n**Risks**: Wrong eviction policy can degrade performance. `.evict_first` on reused data will cause redundant memory fetches. Always benchmark both policies against the default.\n\n**Kernel tiers**: Tier 2 and Tier 3 benefit most (weight reuse in backward passes, multi-pass algorithms).\n\n### 2.3 Reduce Memory Traffic\n\n**When to apply**: Kernel reads or writes data that could be computed on the fly, or writes intermediate results that are only consumed once.\n\n**How to implement**:\n\nLiger already uses several memory-saving patterns. Look for additional opportunities:\n\n```python\n# Pattern 1: Recompute instead of saving to memory\n# Instead of: save silu(a) to a buffer, load in backward\n# Do: recompute silu(a) in backward from saved a\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a  # recomputed in backward\n\n# Pattern 2: In-place gradient storage (gradient-in-forward trick)\n# Cross entropy already does this: stores gradient in the input tensor\n# Look for other kernels that allocate separate grad tensors unnecessarily\n\n# Pattern 3: Fuse producer-consumer pairs to avoid materializing intermediates\n# Instead of: Y = kernel1(X); Z = kernel2(Y)  (Y written to HBM, then read)\n# Do: Z = fused_kernel(X)  (Y lives only in registers/SRAM)\n\n# Pattern 4: Write in-place to dY for dX when shapes match\nif in_place:\n    dX = dY  # reuse memory, zero additional allocation\n```\n\n**Expected impact**: 10-40% from eliminating unnecessary memory round-trips. Memory savings also reduce peak memory.\n\n**Risks**: Recomputation trades memory for compute. Only beneficial when the kernel is memory-bound. Can also complicate backward pass correctness -- always verify gradient accuracy.\n\n**Kernel tiers**: All tiers. Tier 3 kernels (fused operations) have the most opportunity for intermediate elimination.\n\n### 2.4 Data Reuse (Pre-Load Outside Loops)\n\n**When to apply**: A value is loaded inside a loop but does not change across iterations. Common with weight tensors and scaling factors.\n\n**How to implement**:\n\n```python\n# BEFORE: weight loaded inside the loop (redundant loads if compiler doesn't hoist)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    W_block = tl.load(W_ptr + offsets)  # loaded every iteration if W is large\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_block\n\n# AFTER: hoist the invariant load\nW_row = tl.load(W_ptr + col_offsets, mask=col_offsets < n_cols)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_row  # reuse from registers\n```\n\nNote: Triton's compiler often hoists loads automatically, but it is not guaranteed for all patterns, especially when the loop body is complex.\n\n**Expected impact**: 5-15% when applicable.\n\n**Risks**: Hoisting large loads outside loops increases register pressure. If the weight vector is very large relative to BLOCK_SIZE, it may cause register spilling that offsets the benefit.\n\n**Kernel tiers**: Tier 2 (norm kernels with weight vectors), Tier 3 (cross-entropy with weight tensors).\n\n### 2.5 Block-Row Mode (Multiple Rows per Program)\n\n**When to apply**: The kernel processes one row per program instance, but each row is small (e.g., `n_cols <= 256`). This under-utilizes each SM because the program does very little work.\n\n**How to implement**:\n\nLiger's `rms_norm.py` already implements this as `_block_rms_norm_forward_kernel`. The pattern:\n\n```python\n@triton.jit\ndef kernel_block_row(\n    ...,\n    n_rows,\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n    BLOCK_ROW: tl.constexpr,  # number of rows per program\n):\n    # Each program processes BLOCK_ROW rows\n    row_idx = tl.program_id(0) * BLOCK_ROW + tl.arange(0, BLOCK_ROW)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    row_mask = row_idx < n_rows\n    col_mask = col_offsets < n_cols\n\n    # 2D load: [BLOCK_ROW, BLOCK_SIZE]\n    X = tl.load(\n        X_ptr + row_idx[:, None] * stride + col_offsets[None, :],\n        mask=row_mask[:, None] & col_mask[None, :],\n        other=0,\n    )\n    # Process all rows at once...\n```\n\nGrid: `(triton.cdiv(n_rows, BLOCK_ROW),)` instead of `(n_rows,)`.\n\nUse conditional dispatch to choose between row mode and block-row mode:\n\n```python\nif BLOCK_SIZE > 256 or n_rows < 4096 * 8:\n    # Row mode: one row per program\n    kernel_row[(n_rows,)](...)\nelse:\n    # Block-row mode: multiple rows per program\n    BLOCK_ROW = 16\n    kernel_block_row[(triton.cdiv(n_rows, BLOCK_ROW),)](\n        ..., BLOCK_ROW=BLOCK_ROW\n    )\n```\n\n**Expected impact**: 10-30% for kernels with small `n_cols` and large `n_rows`.\n\n**Risks**: Increases register pressure and shared memory usage. The 2D load pattern is less efficient when `BLOCK_ROW` causes register spilling. Also adds code complexity (two kernel variants). Does not help when `n_cols` is already large.\n\n**Kernel tiers**: Tier 1 and Tier 2 primarily. Many Liger Tier 1 kernels (swiglu, geglu, relu_squared) do not yet have block-row mode.\n\n### 2.6 Vectorized Loads\n\n**When to apply**: Kernel loads/stores many small elements individually. Memory bandwidth is limited by the number of transactions rather than total bytes.\n\n**How to implement**:\n\nTriton generally handles vectorization automatically when loads are contiguous and aligned. However, you can help by:\n\n1. Ensuring BLOCK_SIZE is a multiple of 8 (for fp16/bf16) or 4 (for fp32) to enable 128-bit loads.\n2. Ensuring pointer alignment (tensors created by PyTorch are typically aligned).\n3. Avoiding masked loads at boundaries where possible -- pad the input instead.\n\n```python\n# Triton auto-vectorizes contiguous loads. Just ensure alignment:\n# BLOCK_SIZE should be power of 2 (already enforced by calculate_settings)\n# Data should be contiguous (already enforced by ensure_contiguous decorator)\n\n# For explicit vectorization (rare, advanced):\n# Cast pointer to wider type, load fewer elements\n# This is rarely needed as Triton handles it, but can help in edge cases\n```\n\n**Expected impact**: 5-10% if loads are currently unvectorized. Usually Triton already handles this.\n\n**Risks**: Minimal. Padding inputs wastes a small amount of memory.\n\n**Kernel tiers**: All tiers, but gains are usually small since Triton auto-vectorizes well.\n\n---\n\n## 3. Compute-Bound Optimizations\n\nApply these when the profiler classifies the kernel as compute-bound (high compute utilization, low memory throughput utilization, complex math per element).\n\n### 3.1 Algorithmic Improvements\n\n**When to apply**: The kernel uses a naive algorithm where a more efficient one exists. Examples: naive softmax vs online softmax, explicit matrix inverse vs iterative method.\n\n**How to implement**:\n\nLiger already uses online softmax in cross_entropy. Look for similar opportunities:\n\n```python\n# Online algorithm pattern: single-pass with running statistics\n# Instead of: pass 1 = compute max, pass 2 = compute sum, pass 3 = compute output\n# Do: single pass maintaining running max and sum\n\nm = float(\"-inf\")\nd = 0.0\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    block_max = tl.max(X_block)\n    m_new = tl.maximum(m, block_max)\n    d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new))\n    m = m_new\n# Result: lse = m + tl.log(d)\n```\n\nOther algorithmic opportunities:\n- Welford's algorithm for variance (numerically stable single-pass)\n- Kahan summation for high-precision reduction\n- Fast approximations for transcendentals when full precision is not needed\n\n**Expected impact**: 20-60% when applicable (reducing O(n) passes).\n\n**Risks**: Numerical stability can differ from reference implementation. Always validate against PyTorch reference with tight tolerances. Online algorithms may accumulate floating-point errors differently.\n\n**Kernel tiers**: Tier 2 and Tier 3 primarily. Tier 1 kernels rarely have algorithmic improvement opportunities.\n\n### 3.2 Reduce Redundant Computation\n\n**When to apply**: The same sub-expression is computed multiple times within a kernel, or across forward and backward passes.\n\n**How to implement**:\n\n```python\n# BEFORE: sigmoid computed twice\nsilu_a = a_row * tl.sigmoid(a_row)\ngrad = dc_row * (silu_a * (1 - tl.sigmoid(a_row)) + tl.sigmoid(a_row)) * b_row\n\n# AFTER: compute once, reuse\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a\ngrad = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row\n```\n\nLiger already does this well (see swiglu backward). Look for:\n- Repeated `tl.exp()`, `tl.sigmoid()`, `tl.log()` calls with the same argument\n- Reduction results (sum, max) computed multiple times\n- Softmax probabilities recomputed instead of cached in registers\n\n**Expected impact**: 5-20% depending on the cost of the duplicated operation.\n\n**Risks**: Caching sub-expressions increases register pressure. Trade-off with register spilling.\n\n**Kernel tiers**: Tier 2 and Tier 3. Tier 1 kernels are typically simple enough that redundancy is already eliminated.\n\n### 3.3 Tensor Core Utilization\n\n**When to apply**: Kernel performs matrix-multiply-like operations (dot products, outer products) that could leverage tensor cores. This is most relevant for Tier 3 kernels with linear layers or attention-like patterns.\n\n**How to implement**:\n\nTensor cores require specific alignment:\n- Block dimensions must be multiples of 16 (for fp16/bf16) or 8 (for tf32)\n- Input/output matrices must be in supported layouts\n- Use `tl.dot()` which maps to tensor core instructions when dimensions align\n\n```python\n# Ensure BLOCK_M, BLOCK_N, BLOCK_K are multiples of 16\nBLOCK_M: tl.constexpr = 64  # multiple of 16\nBLOCK_N: tl.constexpr = 64  # multiple of 16\nBLOCK_K: tl.constexpr = 32  # multiple of 16\n\n# tl.dot maps to tensor cores automatically\nacc = tl.dot(A_block, B_block)  # [BLOCK_M, BLOCK_K] x [BLOCK_K, BLOCK_N]\n```\n\nFor `tl.dot` to use tensor cores:\n- Both operands must be 2D\n- Operand dtypes must be fp16, bf16, or tf32\n- Inner dimension K must be >= 16\n\n**Expected impact**: 2-10x for matrix-multiply-heavy kernels. Not applicable to element-wise or simple reduction kernels.\n\n**Risks**: Tensor core operations have lower precision than FP32 scalar operations. The `tf32` mode on Ampere+ truncates FP32 mantissa to 10 bits. Always validate numerical accuracy. Not all kernels have operations that map to tensor cores.\n\n**Kernel tiers**: Tier 3 only (fused_linear_cross_entropy, attention-like kernels). Not applicable to Tier 1 or most Tier 2 kernels.\n\n### 3.4 Register Pressure Reduction\n\n**When to apply**: NCU shows high register usage per thread (>128 registers) or occupancy is limited by register count. Kernel body has many live variables.\n\n**How to implement**:\n\n```python\n# Strategy 1: Use constexpr branches to eliminate dead code paths\n# Triton compiles out branches on constexpr values, reducing register usage\ncasting_mode: tl.constexpr  # compiler eliminates unused casting branches\nHAS_WEIGHT: tl.constexpr    # compiler eliminates weight-related registers when False\n\n# Strategy 2: Split large kernels into smaller passes\n# Instead of one kernel doing everything:\n# Kernel 1: forward + save intermediates\n# Kernel 2: backward using intermediates\n# This is a last resort as it adds kernel launch overhead\n\n# Strategy 3: Reduce live range of temporary variables\n# Reorder operations so temporary values are consumed immediately\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    # Use X_block immediately, don't hold it across iterations\n    result = compute(X_block)\n    tl.store(..., result)\n    # X_block and result are now dead, registers freed\n\n# Strategy 4: Use smaller data types for intermediates where precision allows\n# fp16 uses half the registers of fp32\nintermediate = X_block.to(tl.float16)  # only if precision loss is acceptable\n```\n\n**Expected impact**: 10-25% when occupancy is register-limited. Minimal impact when occupancy is already high.\n\n**Risks**: Splitting kernels adds launch overhead and memory traffic for intermediates. Reducing precision can cause numerical issues. Constexpr branches are already used extensively in Liger -- check for remaining opportunities.\n\n**Kernel tiers**: Tier 3 (most register pressure due to complex logic). Tier 2 backward passes (accumulation variables).\n\n### 3.5 2D Grid Tiling for Better SM Utilization\n\n**When to apply**: Kernel uses a 1D grid but the problem has two natural dimensions (e.g., batch and feature). A 1D grid may leave SMs idle when the number of programs is not a multiple of SM count.\n\n**How to implement**:\n\n```python\n# 1D grid (current Liger pattern for most kernels):\ngrid = (n_rows,)\n# If n_rows = 100 and SM count = 108, 8 SMs are idle in the last wave\n\n# 2D grid for better utilization:\ngrid = (triton.cdiv(n_rows, BLOCK_ROW), triton.cdiv(n_cols, BLOCK_COL))\n\n@triton.jit\ndef kernel_2d(\n    ...,\n    BLOCK_ROW: tl.constexpr,\n    BLOCK_COL: tl.constexpr,\n):\n    row_block = tl.program_id(0)\n    col_block = tl.program_id(1)\n    row_start = row_block * BLOCK_ROW\n    col_start = col_block * BLOCK_COL\n    # Process a tile of [BLOCK_ROW, BLOCK_COL]\n```\n\nA 2D grid is most useful when the problem is large in both dimensions and each tile does enough work to amortize launch overhead.\n\n**Expected impact**: 5-15% for kernels with suboptimal SM utilization. More impactful when `n_rows` is small relative to SM count.\n\n**Risks**: 2D tiling for reduction kernels requires careful handling of partial results and synchronization. More complex code. Not beneficial when `n_cols` fits in a single block.\n\n**Kernel tiers**: Tier 2 (reductions with large feature dimensions), Tier 3 (large fused operations).\n\n### 3.6 Loop Unrolling Hints\n\n**When to apply**: Kernel has a hot loop with a known iteration count or small trip count. Triton's compiler unrolls some loops automatically, but explicit hints can help.\n\n**How to implement**:\n\n```python\n# Triton unrolls loops with constexpr bounds automatically.\n# For dynamic bounds, you can help by:\n\n# Strategy 1: Make the loop bound a constexpr when possible\n# If n_cols is known at compile time (e.g., passed as constexpr):\nfor i in range(0, N_COLS, BLOCK_SIZE):  # N_COLS: tl.constexpr -> unrolled\n    ...\n\n# Strategy 2: Use tl.static_range for explicit unrolling\nfor i in tl.static_range(4):  # always unrolled\n    X_block = tl.load(X_ptr + i * BLOCK_SIZE + offsets)\n    ...\n\n# Strategy 3: Break a large loop into a known-count outer loop + remainder\nCHUNKS: tl.constexpr = n_cols // BLOCK_SIZE  # constexpr division\nfor i in tl.static_range(CHUNKS):\n    ...\n# Handle remainder separately\n```\n\n**Expected impact**: 5-10% for tight loops with few iterations. Diminishing returns for large trip counts.\n\n**Risks**: Excessive unrolling increases instruction cache pressure and compile time. Can also increase register pressure if the loop body is large.\n\n**Kernel tiers**: Tier 2 (reduction loops), Tier 3 (multi-pass loops).\n\n---\n\n## 4. Architecture-Specific Optimizations\n\nApply these on top of the best general-purpose variant when targeting a specific GPU architecture.\n\n### 4.1 Ampere (SM 8.x: A100, A10, A30)\n\n**Async Copy (cp.async)**:\n- Triton uses `num_stages` to control software pipelining, which maps to `cp.async` on Ampere.\n- `num_stages=2` is the sweet spot for most Ampere kernels. Going higher increases shared memory usage without benefit.\n- Triton handles `cp.async` automatically when `num_stages > 1`.\n\n**L2 Cache Persistence**:\n- A100 has 40MB L2 cache. Tensors smaller than this can benefit from persistence hints.\n- Use `eviction_policy=\"evict_last\"` on reused data (weights, biases).\n- For cross-entropy, the target tensor and weight tensor are good candidates.\n\n**Optimal Warp Counts**:\n- Ampere excels with 4-16 warps per SM. 32 warps often increases scheduling overhead.\n- For memory-bound kernels: fewer warps (4-8) with larger BLOCK_SIZE.\n- For compute-bound kernels: more warps (8-16) to hide latency.\n\n**Implementation**:\n```python\n# Ampere-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=2),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=2),\n```\n\n### 4.2 Hopper (SM 9.0: H100, H200)\n\n**TMA (Tensor Memory Accelerator)**:\n- Hopper's TMA offloads address computation and data movement to a dedicated unit.\n- Triton >= 3.0 exposes TMA via `tl.make_block_ptr` and experimental APIs.\n- TMA is most beneficial for 2D tiled access patterns (matmul-like kernels).\n\n```python\n# TMA-friendly block pointer (Triton >= 3.0)\nblock_ptr = tl.make_block_ptr(\n    base=X_ptr,\n    shape=(M, N),\n    strides=(stride_m, stride_n),\n    offsets=(row_start, col_start),\n    block_shape=(BLOCK_M, BLOCK_N),\n    order=(1, 0),  # row-major\n)\nX_block = tl.load(block_ptr, boundary_check=[0, 1])\n```\n\n**Higher num_stages**:\n- Hopper's async pipeline supports deeper prefetching. `num_stages=3-5` often outperforms 2.\n- The deeper pipeline hides memory latency more effectively on H100's higher-bandwidth HBM3.\n\n**WGMMA (Warp Group Matrix Multiply Accumulate)**:\n- H100 tensor cores are organized in warp groups (4 warps = 128 threads).\n- `num_warps` should be a multiple of 4 for optimal tensor core usage.\n- Triton maps `tl.dot` to WGMMA automatically when conditions are met.\n\n**Cluster Launch**:\n- Hopper supports thread block clusters (groups of SMs that can share SRAM).\n- Currently experimental in Triton. Potentially useful for large reductions.\n\n**Implementation**:\n```python\n# Hopper-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=4),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=32, num_stages=4),\n```\n\n### 4.3 Blackwell (SM 10.0: B200, GB200)\n\n**5th-Gen Tensor Cores**:\n- Higher throughput for FP8, FP16, BF16, and TF32 operations.\n- Block sizes aligned to 16 remain optimal. FP8 support enables even higher throughput.\n- Triton support is evolving; check Triton release notes for Blackwell-specific features.\n\n**Larger Shared Memory**:\n- Blackwell supports up to 256KB shared memory per SM (vs 228KB on Hopper).\n- Larger BLOCK_SIZE values become viable without spilling to registers or HBM.\n- Can enable larger `BLOCK_ROW` in block-row mode.\n\n**Higher Memory Bandwidth**:\n- HBM3E provides more bandwidth than HBM3. Memory-bound kernels see direct benefit.\n- The compute-to-memory ratio shifts, making some previously memory-bound kernels balanced or compute-bound.\n\n**Implementation**:\n```python\n# Blackwell-tuned config (aggressive block sizes)\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 16384}, num_warps=32, num_stages=5),\n```\n\n**Note**: Blackwell-specific Triton features are still maturing. Focus on general optimizations first, then add Blackwell-specific configs as Triton support stabilizes.\n\n---\n\n## 5. Balanced Optimizations\n\nThese techniques help both memory-bound and compute-bound kernels.\n\n### 5.1 Kernel Fusion Opportunities\n\n**When to apply**: Two or more kernels execute sequentially on the same data, writing intermediate results to HBM between them. The fused version keeps intermediates in registers/SRAM.\n\n**How to implement**:\n\nLiger already excels at fusion (cross_entropy fuses softmax+loss+gradient; fused_linear_cross_entropy fuses matmul+CE). Look for remaining opportunities:\n\n```python\n# Pattern: operation followed by elementwise scaling\n# BEFORE: two kernel launches\ny = rms_norm_kernel(x, w)        # writes y to HBM\nz = elementwise_mul_kernel(y, s)  # reads y from HBM, writes z\n\n# AFTER: fused kernel\n@triton.jit\ndef fused_rms_norm_mul(X_ptr, W_ptr, S_ptr, Z_ptr, ...):\n    # RMS norm computation\n    y = x_row * rstd * (offset + w_row)\n    # Elementwise mul (no HBM round-trip)\n    z = y * s\n    tl.store(Z_ptr + offsets, z)\n```\n\nCommon fusion candidates in Liger:\n- Norm + activation (e.g., rms_norm + swiglu in some architectures)\n- Loss + metric computation (already done: cross_entropy + accuracy)\n- Dropout + residual add + norm\n\n**Expected impact**: 20-50% from eliminating HBM round-trips between fused operations.\n\n**Risks**: Fused kernels are harder to maintain, test, and debug. Register pressure increases. The fused kernel may have worse autotuning characteristics. Only fuse when the intermediate tensor is large relative to cache.\n\n**Kernel tiers**: Creates new Tier 3 kernels from combinations of Tier 1/2 kernels. Existing Tier 3 kernels may have remaining fusion opportunities.\n\n### 5.2 Better Grid Sizing\n\n**When to apply**: The number of programs (grid size) is much smaller or much larger than the SM count, leading to load imbalance or insufficient parallelism.\n\n**How to implement**:\n\n```python\nsm_count = torch.cuda.get_device_properties(device).multi_processor_count\n\n# For backward passes using SM-based parallelism (like rms_norm backward):\ngrid = (sm_count,)\nrows_per_program = math.ceil(n_rows / sm_count)\n\n# For forward passes: ensure grid is a small multiple of sm_count\n# If n_rows >> sm_count, this is automatic\n# If n_rows < sm_count, consider block-row mode or 2D grid to increase occupancy\n\n# For persistent kernels (advanced):\n# Launch exactly sm_count programs, each processes multiple work items\ngrid = (sm_count,)\nfor work_item in range(pid, total_work, NUM_SMS):\n    process(work_item)\n```\n\nThe rms_norm backward kernel already uses this SM-based grid pattern. Look for other backward kernels that use `(n_rows,)` grids but could benefit from SM-based parallelism.\n\n**Expected impact**: 5-20% from better load balancing and reduced tail effects.\n\n**Risks**: SM-based grids require careful work distribution and can introduce non-determinism if work items have unequal cost.\n\n**Kernel tiers**: All tiers. Most impactful for Tier 2 backward passes.\n\n### 5.3 Reduce Kernel Launch Overhead\n\n**When to apply**: Kernel is called many times with very short execution time (< 10 microseconds). Launch overhead dominates.\n\n**How to implement**:\n\n```python\n# Strategy 1: Batch multiple operations into one kernel\n# Instead of launching N small kernels, launch one kernel with N work items\n\n# Strategy 2: Use CUDA graphs (from PyTorch side, not inside Triton)\n# Captures a sequence of kernel launches and replays them with minimal overhead\ngraph = torch.cuda.CUDAGraph()\nwith torch.cuda.graph(graph):\n    output = liger_kernel(input)\n# Replay: graph.replay()\n\n# Strategy 3: Increase work per kernel (larger BLOCK_SIZE, block-row mode)\n# Already covered in autotuning and block-row mode sections\n```\n\n**Expected impact**: 5-15% for very short kernels. Negligible for kernels that run > 100 microseconds.\n\n**Risks**: CUDA graphs have limitations (fixed input shapes, no dynamic control flow). Batching work items increases kernel complexity.\n\n**Kernel tiers**: Tier 1 primarily (short element-wise kernels). Tier 2/3 kernels usually run long enough that launch overhead is negligible.\n\n### 5.4 Mixed Precision Strategies\n\n**When to apply**: Kernel uses FP32 throughout but only needs FP32 for specific operations (reductions, accumulations). Or kernel uses FP16/BF16 throughout but suffers from numerical issues.\n\n**How to implement**:\n\nLiger already uses mixed precision extensively (e.g., rms_norm with casting_mode). Look for additional opportunities:\n\n```python\n# Pattern 1: Compute in fp32 only where needed (Llama-style)\nX_row = tl.load(...)              # fp16/bf16\nX_fp32 = X_row.to(tl.float32)     # upcast for reduction\nmean_sq = tl.sum(X_fp32 * X_fp32) / n_cols  # fp32 reduction\nrstd = rsqrt(mean_sq + eps)        # fp32\nresult = (X_row * rstd.to(X_row.dtype)) * W  # back to original dtype\n\n# Pattern 2: Accumulate in fp32, store in original dtype\nacc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)\nfor ...:\n    acc += (dY_row * X_row).to(tl.float32)  # accumulate in fp32\ntl.store(dW_ptr, acc.to(original_dtype))     # store in original dtype\n\n# Pattern 3: Use BF16 instead of FP16 for better dynamic range\n# BF16 has same exponent bits as FP32, so less overflow/underflow risk\n# But lower mantissa precision than FP16\n```\n\n**Expected impact**: 5-20% from reducing unnecessary FP32 computation. Can also improve accuracy if precision is currently insufficient.\n\n**Risks**: Under-casting causes numerical errors. Over-casting wastes compute. The casting_mode pattern in Liger (llama vs gemma vs none) shows that different models need different precision strategies. Always validate against reference implementation.\n\n**Kernel tiers**: All tiers. Tier 2 reductions are most sensitive to precision. Tier 1 kernels with transcendentals (sigmoid, exp) benefit from targeted FP32 upcast.\n\n---\n\n## Quick Reference: Technique Selection by Symptom\n\n| Symptom | Likely Bottleneck | Try These Techniques |\n|---------|------------------|---------------------|\n| NCU memory throughput > 60% | Memory-bound | 2.1, 2.2, 2.3, 2.5 |\n| NCU compute throughput > 60% | Compute-bound | 3.1, 3.2, 3.3, 3.4 |\n| Low SM occupancy | Register pressure or grid size | 3.4, 5.2, 1 (num_warps) |\n| Small n_cols, large n_rows | Under-utilized programs | 2.5, 5.2, 1 (BLOCK_SIZE) |\n| Large n_cols, small n_rows | Poor SM utilization | 3.5, 5.2 |\n| Multiple sequential kernel launches | Fusion opportunity | 5.1 |\n| Forward fast, backward slow | Backward-specific issue | 2.3, 2.4, 5.2 (SM-based grid) |\n| Performance varies wildly by shape | Bad fixed config | 1 (autotune with key) |\n| High register count in NCU | Register pressure | 3.4, 1 (smaller BLOCK_SIZE) |\n| Poor L2 cache hit rate | Cache misuse | 2.2, 2.4 |\n\n## Quick Reference: Technique Applicability by Kernel Tier\n\n| Technique | Tier 1 (element-wise) | Tier 2 (reduction) | Tier 3 (fused/complex) |\n|-----------|----------------------|--------------------|-----------------------|\n| 1. Autotuning | High | High | High |\n| 2.1 Coalescing | Low (usually ok) | Medium | Medium |\n| 2.2 Cache modifiers | Low | Medium | High |\n| 2.3 Reduce traffic | Medium | Medium | High |\n| 2.4 Data reuse | Low | Medium | Medium |\n| 2.5 Block-row mode | High | High | Low |\n| 2.6 Vectorized loads | Low | Low | Low |\n| 3.1 Algorithmic | Low | High | High |\n| 3.2 Reduce redundancy | Low | Medium | High |\n| 3.3 Tensor cores | N/A | N/A | High |\n| 3.4 Register pressure | Low | Medium | High |\n| 3.5 2D grid tiling | Low | Medium | Medium |\n| 3.6 Loop unrolling | Low | Medium | Medium |\n| 4.x Arch-specific | Medium | Medium | Medium |\n| 5.1 Kernel fusion | Medium (fuse into Tier 3) | Medium | High |\n| 5.2 Grid sizing | Low | High | Medium |\n| 5.3 Launch overhead | High | Low | Low |\n| 5.4 Mixed precision | Medium | High | High |\n","isInternal":false,"tokens":8828,"sizeBytes":35446},{"name":"optimizer.md","path":".agents/skills/liger-kernel-perf/optimizer.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-perf/optimizer.md","title":"Liger-kernel-perf Skill","category":"anthropic-skill","format":"markdown","content":"# Optimizer Agent\n\nRuns an autonomous optimization loop over an existing Liger kernel, generating multiple variants, benchmarking each against the baseline, and selecting the best performer while maintaining correctness.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\", \"cross_entropy\", \"swiglu\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `max_variants`: maximum number of variants to try (default 8)\n- `target_metric`: optional concrete target (e.g., \"forward under 0.3ms at hidden_size=4096\")\n- `optimization_profile`: the profile produced by the Profiler Agent at `optimization/{kernel}/profile.md`\n- `original_kernel`: the original kernel snapshot at `optimization/{kernel}/original_{kernel}.py`\n\n## Pre-Requisites\n\nBefore starting the optimization loop:\n\n1. Verify the workspace exists: `optimization/{kernel}/` with `original_{kernel}.py` and `profile.md`\n2. Verify the benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`\n3. Verify the test file exists: `test/transformers/test_{kernel}.py`\n4. Read the optimization profile thoroughly — extract bottleneck classification, tier, current config, and recommended strategy order\n5. Read the original kernel code thoroughly — understand every `@triton.jit` function, the `torch.autograd.Function`, and the launcher logic\n6. Read `optimization-strategies.md` for the full catalog of available techniques\n\n## Variant Numbering\n\n- v0 = baseline (original kernel, already benchmarked by the Profiler)\n- v1, v2, ... = optimization variants generated by this agent\n- Each variant gets two files:\n  - `optimization/{kernel}/{kernel}_vN.py` — the modified kernel code\n  - `optimization/{kernel}/{kernel}_vN_notes.md` — the lab notebook entry\n\n## The Optimization Loop\n\n```\nvariant_number = 1\nbest_variant = 0  (baseline)\nconsecutive_no_improvement = 0\n\nwhile variant_number <= max_variants:\n    1. Read ALL prior variant notes (v1 through v{variant_number-1})\n    2. Decide what to try next (see Strategy Selection below)\n    3. Generate variant code\n    4. Write variant notes (hypothesis section)\n    5. Run smoke test → if FAIL, mark REJECTED, increment variant_number, continue\n    6. Run full benchmark → record results\n    7. Check guardrails → if FAIL, mark REJECTED, increment variant_number, continue\n    8. Update variant notes with actual results and verdict\n    9. If improvement over best_variant: update best_variant, reset consecutive_no_improvement\n   10. Else: increment consecutive_no_improvement\n   11. If consecutive_no_improvement >= 2: STOP (diminishing returns)\n   12. If target_metric is set and met: STOP (goal achieved)\n   13. Increment variant_number\n```\n\n## Strategy Selection\n\n### Always Try Parameter Tuning First (v1)\n\nThe first variant MUST be a parameter sweep over `BLOCK_SIZE`, `num_warps`, and `num_stages`. This is the cheapest optimization with the highest expected payoff.\n\n**Why NOT @triton.autotune**: Liger kernels use `torch.autograd.Function` where the forward pass stores `BLOCK_SIZE` and `num_warps` in `ctx` and passes them to the backward kernel. With `@triton.autotune`, these become compile-time constants that cannot be stored in `ctx` — the backward kernel would not know which config the forward used. Additionally, NPU/Ascend backends do not support `num_warps` or `num_stages` at all. For these reasons, `@triton.autotune` is intentionally not used in Liger production kernels.\n\n**Instead, do a manual parameter sweep**: Replace the `calculate_settings(n_cols)` call with hardcoded values and benchmark each configuration:\n\n```python\n# In the variant, replace:\n#   BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n# With a specific config to test:\nBLOCK_SIZE = 2048   # Try: 256, 512, 1024, 2048, 4096, 8192\nnum_warps = 8       # Try: 2, 4, 8, 16, 32\nnum_stages = 2      # Try: 1, 2, 3, 4\n```\n\n**Sweep procedure for v1**:\n1. Read the current `calculate_settings(n_cols)` output for the benchmark's representative input sizes\n2. Create a variant that replaces `calculate_settings` with a custom function that returns a better config\n3. Test 3-5 promising configs (not the full cartesian product — use the architecture guidance below to prune)\n4. Pick the best config across ALL benchmark x_values (not just one size)\n5. The variant should implement a smarter `calculate_settings` replacement, not just a single hardcoded value\n\n**Architecture-specific guidance for config selection**:\n- **Ampere**: `num_stages=2`, `num_warps` up to 16, BLOCK_SIZE 1024-4096\n- **Hopper**: `num_stages=3-5`, `num_warps` up to 32, BLOCK_SIZE 2048-8192\n- **Blackwell**: `num_stages=4-5`, similar to Hopper\n\n**Important**: Keep the same forward/backward ctx pattern. Both kernels must use the same `BLOCK_SIZE` and `num_warps` for a given input. The variant must still pass these through `ctx` just like the original.\n\n### Subsequent Variants: Diagnosis-Driven\n\nAfter autotuning, choose the next technique based on:\n\n1. **The bottleneck classification from the profile** (memory-bound vs compute-bound)\n2. **Learnings from all prior variants** (what worked, what did not, what was surprising)\n3. **The recommended strategy order from the profile**\n\nConsult [optimization-strategies.md](optimization-strategies.md) for the full technique catalog. Common patterns by bottleneck type:\n\n**Memory-bound kernels** (Tier 1, simple element-wise):\n- Fuse multiple memory-bound operations into a single kernel\n- Add cache eviction hints (`tl.store(..., eviction_policy=\"evict_first\")`)\n- Increase BLOCK_SIZE to amortize memory latency\n- Use `tl.load(..., eviction_policy=\"evict_last\")` for reused data\n- Reduce unnecessary memory traffic (eliminate redundant loads/stores)\n- In-place operations where safe\n\n**Compute-bound kernels** (Tier 3, heavy math):\n- Reduce precision of intermediate computations where safe\n- Use fast-math intrinsics (`tl.math.fast_expf`, etc.)\n- Restructure computation to reduce total FLOPs\n- Improve instruction-level parallelism via loop unrolling\n- Increase occupancy by reducing register pressure\n\n**Balanced / Tier 2 kernels** (reductions):\n- Optimize the reduction pattern (two-pass vs single-pass)\n- Tune BLOCK_ROW for multi-row processing\n- Balance between register usage and parallelism\n- Consider split-K or hierarchical reductions for large dimensions\n- Block-row mode threshold tuning (the `n_rows < 4096 * 8` style checks)\n\n### Combining Techniques\n\nA single variant may combine multiple techniques if they are complementary (e.g., autotuning + cache hints). However, prefer isolated changes so that the lab notebook clearly attributes improvements to specific techniques. When combining, explain which part of the improvement is expected from each technique.\n\n### When All Strategies Exhausted\n\nIf all recommended strategies have been tried and the consecutive_no_improvement counter has not triggered, the agent may:\n- Try creative combinations of previous partial improvements\n- Attempt architecture-specific tricks (Hopper warp specialization, Blackwell TMA)\n- Revisit a rejected variant with a modified approach based on learnings\n\n## Generating a Variant\n\nFor each variant, produce a complete, self-contained kernel file at `optimization/{kernel}/{kernel}_vN.py`.\n\nRules:\n1. Start from the original kernel (not a previous variant) unless explicitly building on a prior variant's approach. If building on a prior variant, note which one and why.\n2. The variant file must be a drop-in replacement for `src/liger_kernel/ops/{kernel}.py` — same public API, same class names, same function signatures.\n3. Do not change the `torch.autograd.Function` class name or the public API.\n4. Do not remove safety checks, dtype handling, or multi-device support (XPU, NPU paths) unless the change is specifically about those areas.\n5. **Preserve ALL existing comments** — function docstrings, mathematical formula explanations, inline rationale comments, and section headers. Do NOT delete or shorten comments that explain what the code does or why.\n6. **Add explanatory comments for every optimization change** — each modified section should have a comment explaining what was changed and why it improves performance. For example:\n   ```python\n   # Optimization: compute dW before dX to allow register reuse.\n   # The original order kept 8 vectors live simultaneously (115 regs/thread).\n   # This reordering reduces peak liveness to 5-6 vectors (95 regs/thread).\n   ```\n7. **Preserve function-level docstrings** — if the original kernel functions have docstrings or parameter descriptions, keep them. Update them if the function signature changes.\n8. Add a comment block at the top of the file documenting what was changed:\n\n```python\n\"\"\"\nOptimization Variant {N}: {Short Title}\nBased on: original (or v{M} if building on a prior variant)\nChanges: {brief description of modifications}\n\"\"\"\n```\n\n## Variant Lab Notebook\n\nFor each variant, write `optimization/{kernel}/{kernel}_vN_notes.md` using the format in [templates/variant-notes.md](templates/variant-notes.md).\n\nKey requirements:\n- The Hypothesis must reference the profile's bottleneck classification and learnings from prior variants\n- Changes must include rationale for each modification with code snippets\n- Expected Outcome must be specific and quantitative (\"10-20% forward speed improvement\" not \"should be faster\")\n- Actual Results must include ALL x_values if improvement varies across sizes\n- Verdict: ACCEPTED (improves target without violating guardrails), REJECTED (fails smoke test/guardrails/no improvement), or MIXED (improves some metrics, regresses others within limits)\n- Learnings must include insights to carry forward to the next variant\n\n### Writing Notes Across the Loop\n\n- **Before smoke test**: Write the Hypothesis, Changes, and Expected Outcome sections\n- **After benchmark**: Fill in Actual Results, Verdict, and Learnings\n- **On smoke test failure**: Fill Actual Results with \"SMOKE TEST FAILED: {error}\", set Verdict to REJECTED, and write Learnings about what went wrong\n\n## Smoke Test Procedure\n\nThe smoke test validates basic correctness before investing time in a full benchmark. Execute it by temporarily swapping the variant into the source tree.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run a Minimal Correctness Check\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80\n```\n\nThis runs only float32 tests (the most numerically stable) with a timeout to catch hangs. The `-x` flag stops at the first failure.\n\n### Step 3: Interpret Results\n\n- **All tests pass**: Smoke test PASSES. Proceed to full benchmark.\n- **Any test fails**: Smoke test FAILS. Read the error carefully.\n  - If the failure is a **tolerance issue** (allclose fails but values are close): Note this in the variant notes. The optimization may have changed numerical behavior slightly. Try adjusting the variant to preserve precision, or note the tolerance difference.\n  - If the failure is a **crash or shape error**: The variant has a bug. Mark REJECTED immediately. Record the error in the notes for learning.\n  - If the failure is a **hang/timeout**: The variant has a deadlock or infinite loop. Mark REJECTED immediately.\n\n### Step 4: Restore the Original\n\n**CRITICAL**: Always restore the original kernel after the smoke test, regardless of outcome:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis must happen even if the smoke test crashes. Use a pattern that ensures restoration:\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py && \\\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80; \\\nSMOKE_EXIT=$?; \\\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py; \\\nexit $SMOKE_EXIT\n```\n\n## Full Benchmark Procedure\n\nAfter a variant passes the smoke test, run the full existing benchmark suite to get precise measurements.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run the Full Benchmark\n\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nThis runs speed benchmarks (forward, backward, full) and memory benchmarks across all configured x_values, comparing the \"liger\" provider against the baseline (\"huggingface\" or \"torch\").\n\n### Step 3: Capture Results\n\nThe benchmark writes results to `benchmark/data/all_benchmark_data.csv`. Copy the relevant rows:\n\n```bash\ncp benchmark/data/all_benchmark_data.csv optimization/{kernel}/benchmarks/vN_results.csv\n```\n\n### Step 4: Parse Results\n\nRead the CSV and extract:\n- Speed (ms) at the median (p50) for forward, backward, and full modes\n- Memory (MB) at the median (p50) for full mode\n- Values at each x_value to detect size-dependent behavior\n\nCompare against:\n- v0 baseline: the absolute improvement\n- Current best variant: whether this is the new best\n\n### Step 5: Restore the Original\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n## Guardrail Checks\n\nAfter benchmarking each variant, apply these guardrails. All comparisons are against the v0 baseline.\n\n### Guardrail 1: Non-Target Metric Regression (>5%)\n\nIf the optimization goal is \"speed\", check that memory does not regress by more than 5%. If the goal is \"memory\", check that speed does not regress by more than 5%. If the goal is \"balanced\", check both.\n\n```\nif goal == \"speed\":\n    reject if memory_vN > memory_v0 * 1.05\nelif goal == \"memory\":\n    reject if speed_vN > speed_v0 * 1.05  (higher ms = slower = worse)\nelif goal == \"balanced\":\n    reject if memory_vN > memory_v0 * 1.05 OR speed_vN > speed_v0 * 1.05\n```\n\nCheck this at EVERY x_value. If the regression occurs only at certain sizes, note this in the variant notes but still reject — partial regressions are not acceptable.\n\n### Guardrail 2: Cross-Pass Regression (>10%)\n\nDo not sacrifice one pass (forward or backward) to improve the other beyond a 10% threshold:\n\n```\nreject if speed_fwd_vN > speed_fwd_v0 * 1.10 AND speed_bwd_vN < speed_bwd_v0\nreject if speed_bwd_vN > speed_bwd_v0 * 1.10 AND speed_fwd_vN < speed_fwd_v0\n```\n\nAn optimization that makes forward 15% slower to make backward 5% faster is rejected.\n\n### Guardrail 3: Smoke Test Failure\n\nAlready handled in the smoke test procedure. Any correctness failure means immediate rejection.\n\n### Handling Guardrail Failures\n\nWhen a variant is rejected by guardrails:\n1. Set Verdict to REJECTED in the variant notes\n2. Record exactly which guardrail was violated and by how much\n3. Record the Learnings — this information is valuable for future variants\n4. Do NOT increment consecutive_no_improvement (guardrail failures are not \"stagnation\")\n5. Proceed to the next variant\n\n## Stopping Conditions\n\nThe loop terminates when ANY of these conditions is met:\n\n### Condition 1: Budget Exhausted\n\n`variant_number > max_variants`\n\nReport: \"Optimization budget of {max_variants} variants exhausted.\"\n\n### Condition 2: Diminishing Returns\n\n`consecutive_no_improvement >= 2`\n\nTwo consecutive variants that pass guardrails but improve the target metric by less than 1% over the current best. This indicates the kernel is near its optimization ceiling for the current approach.\n\nReport: \"Stopping due to diminishing returns — last 2 accepted variants showed <1% improvement.\"\n\n### Condition 3: Target Metric Met\n\nThe user-specified `target_metric` has been achieved.\n\nReport: \"Target metric achieved by variant v{N}: {metric_description}.\"\n\n## Reading Prior Variant Notes\n\nBefore generating each new variant, the agent MUST read ALL prior variant notes. This is the accumulated learning mechanism that makes later variants smarter.\n\nWhen reading prior notes, extract:\n1. **What techniques were tried**: Avoid repeating the exact same approach\n2. **What worked**: Build on successful directions\n3. **What failed and why**: Avoid repeating mistakes\n4. **Surprising results**: These often reveal the true bottleneck\n5. **Size-dependent behavior**: Some optimizations help at large sizes but hurt at small sizes\n\nUse this information to form a more informed hypothesis for the next variant.\n\n## Final Comparison Table\n\nAfter the loop terminates (for any reason), produce a comparison table of ALL variants and save it to `optimization/{kernel}/comparison.md`.\n\nThe comparison must include:\n- A master table with ALL variants (accepted and rejected): Variant | Strategy | Speed Fwd/Bwd/Full (ms) | Memory (MB) | Verdict\n- Values at the representative x_value (typically the largest or most common) with percentage delta from baseline\n- Winner declaration with justification\n- Cumulative learnings summary (what worked, what didn't, why)\n\n## Error Recovery\n\n### Benchmark Script Crashes\n\nIf the benchmark script crashes (not a test failure — an actual crash):\n1. Restore the original kernel immediately\n2. Read the crash output\n3. If it is a Triton compilation error: the variant has a syntax or semantic error in the kernel. Mark REJECTED with the error.\n4. If it is a CUDA OOM: the variant uses too much memory for the GPU. Mark REJECTED. Note the memory constraint.\n5. If it is an unrelated infrastructure error: retry once. If it fails again, report to the user.\n\n### Workspace Corruption\n\nIf at any point `src/liger_kernel/ops/{kernel}.py` is found to differ from `optimization/{kernel}/original_{kernel}.py` when it should be the original:\n\n```bash\ndiff src/liger_kernel/ops/{kernel}.py optimization/{kernel}/original_{kernel}.py\n```\n\nIf they differ unexpectedly, restore immediately:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Agent Failure Mid-Loop\n\nIf the agent encounters an unrecoverable error:\n1. Restore the original kernel\n2. Save all variant notes written so far\n3. Write a partial comparison table with whatever data is available\n4. Report to the user with the current state and what went wrong\n\n## Important Constraints\n\n1. **Never leave a variant installed in the source tree.** The original kernel must always be restored after each smoke test and benchmark. This is the most critical invariant.\n2. **Never modify test files.** The existing tests are the correctness oracle. If a variant fails tests, the variant is wrong, not the tests.\n3. **Never modify benchmark files.** The existing benchmarks are the measurement standard. Variants must be measured by the same benchmark to enable fair comparison.\n4. **Always write notes before running experiments.** The hypothesis must be written before seeing results to avoid post-hoc rationalization.\n5. **Always read all prior notes before generating a new variant.** This is what makes the optimization loop learn from itself.\n6. **Prefer the simplest change that achieves the goal.** A 5-line autotuning change that yields 15% speedup is better than a 200-line rewrite that yields 20%.\n","isInternal":false,"tokens":4424,"sizeBytes":19076},{"name":"profiler.md","path":".agents/skills/liger-kernel-perf/profiler.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-perf/profiler.md","title":"Liger-kernel-perf Skill","category":"anthropic-skill","format":"markdown","content":"# Profiler Agent\n\nProfiles an existing Liger kernel to diagnose performance bottlenecks and produce an optimization strategy.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `target_gpu`: Ampere / Hopper / Blackwell / auto-detect\n\n## Steps\n\n### Step 1: Set Up Workspace\n\n1. Create directory `optimization/{kernel}/` at the repo root\n2. Copy `src/liger_kernel/ops/{kernel}.py` to `optimization/{kernel}/original_{kernel}.py`\n3. Create `optimization/{kernel}/benchmarks/` subdirectory\n\n### Step 2: Detect GPU Architecture\n\nRun:\n```bash\npython -c \"import torch; props = torch.cuda.get_device_properties(0); print(f'GPU: {props.name}, SM: {props.major}.{props.minor}, SMs: {props.multi_processor_count}, Mem: {props.total_mem // (1024**3)}GB')\"\n```\n\nClassify architecture:\n- SM 8.x → Ampere (A100, A10, etc.)\n- SM 9.0 → Hopper (H100, H200, etc.)\n- SM 10.0 → Blackwell (B200, etc.)\n\nIf user specified a target_gpu, note the mismatch (if any) but proceed with the actual hardware.\n\n### Step 3: Run Baseline Benchmarks\n\nRun the existing benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nRecord the baseline results. Copy the relevant rows from `benchmark/data/all_benchmark_data.csv` to `optimization/{kernel}/benchmarks/v0_baseline.csv`.\n\nParse and summarize the baseline numbers:\n- Speed: forward, backward, full (median ms at each x_value)\n- Memory: full (median MB at each x_value)\n- Identify the slowest operation mode and largest x_values\n\n### Step 4: Check for NCU Availability\n\n```bash\nwhich ncu 2>/dev/null && echo \"NCU_AVAILABLE\" || echo \"NCU_NOT_AVAILABLE\"\n```\n\nIf available, offer to run profiling (interactive mode) or run it (autonomous mode).\n\nTo determine the correct import and invocation, read the kernel's `torch.autograd.Function` class from `src/liger_kernel/ops/{kernel}.py` and construct a minimal reproduction script. For example, for `rms_norm`:\n\n```bash\nncu --set full --target-processes all -o optimization/{kernel}/ncu_profile \\\n    python -c \"\nimport torch\nfrom liger_kernel.ops.rms_norm import LigerRMSNormFunction\nx = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16, requires_grad=True)\nw = torch.ones(4096, device='cuda', dtype=torch.bfloat16)\ny = LigerRMSNormFunction.apply(x, w, 1e-6, 0.0, 'llama', False)\ny.sum().backward()\n\"\n```\n\nAdapt the import path, function arguments, and tensor shapes for the target kernel.\n\nExtract key metrics from NCU output:\n- **SM Occupancy** (%)\n- **Compute Throughput** (% of peak)\n- **Memory Throughput** (% of peak)\n- **L1/L2 Cache Hit Rates**\n- **Warp Stall Reasons** (top 3)\n\n### Step 5: Analyze Kernel Code\n\nRead `src/liger_kernel/ops/{kernel}.py` thoroughly. Extract:\n\n1. **Tier classification**:\n   - Tier 1 (element-wise): No reductions. One row per program.\n   - Tier 2 (reduction): Cross-column reductions. May need SM-based parallelism.\n   - Tier 3 (fused/complex): Multi-pass algorithms, gradient-in-forward tricks.\n\n2. **Current configuration**:\n   - How BLOCK_SIZE is set (calculate_settings? hardcoded? autotune?)\n   - num_warps value\n   - num_stages value (if specified)\n   - Grid dimensions (1D, 2D, block-row mode?)\n\n3. **Memory access patterns**:\n   - Stride-based or contiguous?\n   - Cache modifiers used?\n   - In-place operations?\n   - What is saved_for_backward vs recomputed?\n\n4. **Compute patterns**:\n   - Precision-sensitive operations (sigmoid, rsqrt, exp, log)?\n   - Casting modes?\n   - Online/chunked algorithms?\n   - constexpr parameters?\n\n5. **Existing optimization opportunities**:\n   - Any TODO/FIXME comments about performance?\n   - Missing block-row mode variant?\n   - Suboptimal BLOCK_SIZE heuristic?\n   - Unnecessary memory traffic?\n\n### Step 6: Classify Bottleneck\n\nUse NCU data (if available) or heuristics:\n\n**Memory-bound indicators**:\n- NCU: Memory throughput > 60% of peak, compute throughput < 40%\n- Heuristic: Simple element-wise ops, large tensors, few arithmetic operations per load\n- Heuristic: Tier 1 kernels are almost always memory-bound\n\n**Compute-bound indicators**:\n- NCU: Compute throughput > 60% of peak, memory throughput < 40%\n- Heuristic: Complex math per element (exp, log, sigmoid chains), reductions, many operations per load\n- Heuristic: Tier 3 kernels with heavy computation\n\n**Balanced (both bottlenecks)**:\n- NCU: Both throughputs 30-60%\n- Heuristic: Tier 2 kernels often fall here\n\n### Step 7: Produce Optimization Profile\n\nWrite `optimization/{kernel}/profile.md` using the [optimization-profile template](templates/optimization-profile.md).\n\nThe profile must include:\n- GPU architecture and properties\n- Baseline benchmark summary (table)\n- NCU summary (if available)\n- Kernel tier classification\n- Bottleneck classification with evidence\n- Current configuration analysis\n- Recommended strategy order (from [optimization-strategies.md](optimization-strategies.md))\n- Specific opportunities identified in the code\n\n### Step 8: Present to User (Interactive Mode)\n\nShow:\n1. GPU detected: {name} ({architecture})\n2. Baseline performance summary table\n3. Bottleneck classification: {memory-bound / compute-bound / balanced}\n4. Evidence for classification\n5. Proposed optimization strategy order\n6. Estimated iteration count\n\nWait for user confirmation before proceeding to Stage 2.\n\nIn autonomous mode, skip this step and proceed directly.\n","isInternal":false,"tokens":1371,"sizeBytes":5455},{"name":"SKILL.md","path":".agents/skills/liger-kernel-perf/SKILL.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-perf/SKILL.md","title":"liger-kernel-perf","category":"anthropic-skill","format":"markdown","content":"---\nname: liger-kernel-perf\ndescription: \"Optimizes the performance of existing Liger Kernel Triton kernels. Profiles kernels, diagnoses bottlenecks (memory-bound vs compute-bound), generates multiple optimization variants with benchmarking, and applies the best variant while maintaining correctness. Supports GPU architecture-specific optimization (Ampere, Hopper, Blackwell). Use when a user asks to optimize, speed up, tune, profile, or reduce memory of an existing Liger kernel.\"\n---\n\n# Liger Kernel Perf\n\nOptimizes existing Liger Kernel Triton kernels through a 3-stage pipeline: Profile, Optimize, Finalize. Supports interactive mode (human checkpoints between stages) and autonomous mode (runs end-to-end). NVIDIA GPUs only.\n\n## Mode Detection\n\n- **Interactive mode** (default): Human checkpoints between each stage\n- **Autonomous mode**: User says \"just optimize it\", \"run without asking me\", \"optimize autonomously\" → all stages run end-to-end, user sees only the final report\n\n## Input Parsing\n\nExtract from the user's request:\n\n| Field | Description | Default |\n|-------|-------------|---------|\n| `target_kernel` | Which kernel to optimize (e.g., \"rms_norm\", \"cross_entropy\") | **Required** |\n| `optimization_goal` | speed / memory / balanced | balanced |\n| `scope` | Specific pass (forward/backward), input regime, or general | general |\n| `target_gpu` | Ampere / Hopper / Blackwell / auto-detect | auto-detect |\n| `autonomy` | interactive / autonomous | interactive |\n| `max_variants` | Max optimization variants to try | 8 |\n| `target_metric` | Optional concrete target (e.g., \"forward under 0.3ms at hidden_size=4096\") | none |\n\n## Pre-Flight Validation\n\nBefore starting the pipeline, validate:\n\n1. Kernel file exists: `src/liger_kernel/ops/{kernel}.py`\n2. Benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`\n3. Test file exists: `test/transformers/test_{kernel}.py`\n4. GPU is available and CUDA works\n5. Project is installed in dev mode (`pip install -e \".[dev]\"`)\n\nIf any validation fails, report clearly and stop.\n\n## Pipeline\n\n### Stage 1: Profile\n\nFollow the **Profiler** workflow in [profiler.md](profiler.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nThis stage:\n1. Creates the workspace directory `optimization/{kernel}/`\n2. Copies the original kernel as a snapshot\n3. Runs baseline benchmarks using the existing benchmark script\n4. Detects GPU architecture (or uses user-specified target)\n5. Optionally runs NCU profiling (if `ncu` is available)\n6. Analyzes the kernel code (tier classification, patterns, optimization opportunities)\n7. Classifies the bottleneck: memory-bound vs compute-bound\n8. Produces an optimization profile with a recommended strategy order\n9. Saves profile to `optimization/{kernel}/profile.md`\n\n**Human checkpoint (interactive mode):** Present the optimization profile with bottleneck diagnosis and proposed strategy order. Confirm before proceeding.\n\n### Stage 2: Optimize\n\nFollow the **Optimizer** workflow in [optimizer.md](optimizer.md).\n\nThis stage runs an autonomous optimization loop:\n\n1. Read the optimization profile and original kernel\n2. **Always try parameter tuning first** (BLOCK_SIZE, num_warps, num_stages manual sweep -- NOT @triton.autotune)\n3. Then apply diagnosis-driven techniques from [optimization-strategies.md](optimization-strategies.md)\n4. For each variant:\n   a. Generate the variant code → `optimization/{kernel}/{kernel}_vN.py`\n   b. Write the variant lab notebook → `optimization/{kernel}/{kernel}_vN_notes.md`\n   c. Run quick smoke test (single shape, float32, forward+backward) → discard on failure\n   d. Run the **full existing benchmark script** → `optimization/{kernel}/benchmarks/vN_results.csv`\n   e. Check guardrails (no catastrophic regressions)\n   f. Update the variant notes with actual results\n5. Read all prior variant notes before generating the next variant\n6. **Stop when:** budget exhausted, 2 consecutive variants with <1% improvement, or target metric met\n7. Produce a comparison table of ALL variants\n\n**Human checkpoint (interactive mode):** Present the comparison table across all variants. User approves the winner (or skill picks best if autonomous).\n\n### Stage 3: Finalize\n\nFollow the **Finalizer** workflow in [finalizer.md](finalizer.md).\n\nThis stage:\n1. Applies the winning variant in-place to `src/liger_kernel/ops/{kernel}.py`\n2. Runs the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs` (hard gate)\n3. Runs checkstyle: `make checkstyle` (auto-fix with `ruff check . --fix && ruff format .`)\n4. Generates 3-way comparison plots (original liger vs optimized liger vs huggingface baseline) using `benchmarks_visualizer.py`\n5. Generates the final optimization report → `optimization/{kernel}/report.md`\n6. Creates a PR with only the kernel code changes (no plots or optimization workspace files)\n7. Presents the before/after summary with plots\n\n**Human checkpoint (interactive mode):** Present the final report with before/after numbers, comparison plots, and test results.\n\n## Guardrails\n\nThese apply to EVERY variant, regardless of mode:\n\n| Guardrail | Threshold | Action |\n|-----------|-----------|--------|\n| Non-target metric regression | >5% worse | Reject variant |\n| Cross-pass regression | >10% on one pass to marginally improve other | Reject variant |\n| Smoke test failure | Any correctness failure | Discard variant immediately |\n| Full test suite failure | Any | Do NOT apply winner, report failure, stop |\n| Checkstyle failure | Any | Auto-fix with ruff, retry once |\n\n## Reference Files\n\n- [profiler.md](profiler.md) -- Profiler Agent specification\n- [optimizer.md](optimizer.md) -- Optimizer Agent specification\n- [finalizer.md](finalizer.md) -- Finalizer Agent specification\n- [optimization-strategies.md](optimization-strategies.md) -- Catalog of optimization techniques\n- [templates/optimization-profile.md](templates/optimization-profile.md) -- Profiling output format (cross-stage contract)\n- [templates/variant-notes.md](templates/variant-notes.md) -- Per-variant lab notebook format\n","frontmatter":{"name":"liger-kernel-perf","description":"Optimizes the performance of existing Liger Kernel Triton kernels. Profiles kernels, diagnoses bottlenecks (memory-bound vs compute-bound), generates multiple optimization variants with benchmarking, and applies the best variant while maintaining correctness. Supports GPU architecture-specific optimization (Ampere, Hopper, Blackwell). Use when a user asks to optimize, speed up, tune, profile, or reduce memory of an existing Liger kernel."},"isInternal":false,"tokens":1401,"sizeBytes":6155},{"name":"optimization-profile.md","path":".agents/skills/liger-kernel-perf/templates/optimization-profile.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-perf/templates/optimization-profile.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Optimization Profile: {kernel}\n\n> Generated by Profiler Agent (Stage 1)\n> Date: {date}\n\n## Kernel Identity\n\n| Field | Value |\n|-------|-------|\n| Kernel Name | `{kernel}` |\n| Source File | `src/liger_kernel/ops/{kernel}.py` |\n| Benchmark Script | `benchmark/scripts/benchmark_{kernel}.py` |\n| Test File | `test/transformers/test_{kernel}.py` |\n| Tier Classification | {tier} ({tier_description}) |\n\n## GPU Information\n\n| Property | Value |\n|----------|-------|\n| GPU Name | {gpu_name} |\n| Architecture | {gpu_arch} (SM {sm_version}) |\n| SM Count | {sm_count} |\n| Total Memory | {gpu_memory_gb} GB |\n| Target GPU (user-specified) | {target_gpu} |\n| Hardware Match | {hardware_match} |\n\n## Baseline Benchmarks\n\n### Speed (median ms)\n\n| x_value | Mode | {kernel} (Liger) | {baseline_provider} | Speedup |\n|---------|------|-------------------|----------------------|---------|\n| {x_val_1} | forward | {liger_fwd_1} | {baseline_fwd_1} | {speedup_fwd_1}x |\n| {x_val_1} | backward | {liger_bwd_1} | {baseline_bwd_1} | {speedup_bwd_1}x |\n| {x_val_1} | full | {liger_full_1} | {baseline_full_1} | {speedup_full_1}x |\n| {x_val_2} | forward | {liger_fwd_2} | {baseline_fwd_2} | {speedup_fwd_2}x |\n| {x_val_2} | backward | {liger_bwd_2} | {baseline_bwd_2} | {speedup_bwd_2}x |\n| {x_val_2} | full | {liger_full_2} | {baseline_full_2} | {speedup_full_2}x |\n| ... | ... | ... | ... | ... |\n\n### Memory (median MB)\n\n| x_value | {kernel} (Liger) | {baseline_provider} | Savings |\n|---------|-------------------|----------------------|---------|\n| {x_val_1} | {liger_mem_1} | {baseline_mem_1} | {mem_savings_1}% |\n| {x_val_2} | {liger_mem_2} | {baseline_mem_2} | {mem_savings_2}% |\n| ... | ... | ... | ... |\n\n### Observations\n\n- Slowest mode: {slowest_mode} ({slowest_mode_reason})\n- Largest x_value tested: {largest_x_val}\n- Performance scaling trend: {scaling_trend}\n\n## NCU Profiling Summary\n\n> Status: {ncu_status}\n\n{ncu_available_section_start}\n\n| Metric | Value | Assessment |\n|--------|-------|------------|\n| SM Occupancy | {sm_occupancy}% | {occupancy_assessment} |\n| Compute Throughput | {compute_throughput}% of peak | {compute_assessment} |\n| Memory Throughput | {memory_throughput}% of peak | {memory_assessment} |\n| L1 Cache Hit Rate | {l1_hit_rate}% | {l1_assessment} |\n| L2 Cache Hit Rate | {l2_hit_rate}% | {l2_assessment} |\n\n### Top Warp Stall Reasons\n\n1. {warp_stall_1} ({warp_stall_pct_1}%)\n2. {warp_stall_2} ({warp_stall_pct_2}%)\n3. {warp_stall_3} ({warp_stall_pct_3}%)\n\n{ncu_available_section_end}\n\n{ncu_unavailable_section_start}\n\nNCU is not available on this system. Bottleneck classification is based on code heuristics only.\n\n{ncu_unavailable_section_end}\n\n## Bottleneck Classification\n\n**Classification: {bottleneck_class}**\n\n### Evidence\n\n| Signal | Observation | Supports |\n|--------|-------------|----------|\n| {evidence_signal_1} | {evidence_obs_1} | {evidence_supports_1} |\n| {evidence_signal_2} | {evidence_obs_2} | {evidence_supports_2} |\n| {evidence_signal_3} | {evidence_obs_3} | {evidence_supports_3} |\n| ... | ... | ... |\n\n### Reasoning\n\n{bottleneck_reasoning}\n\n## Current Configuration Analysis\n\n| Parameter | Current Value | How Set | Notes |\n|-----------|--------------|---------|-------|\n| BLOCK_SIZE | {block_size} | {block_size_source} | {block_size_notes} |\n| num_warps | {num_warps} | {num_warps_source} | {num_warps_notes} |\n| num_stages | {num_stages} | {num_stages_source} | {num_stages_notes} |\n| Grid Dimensions | {grid_dims} | {grid_dims_source} | {grid_dims_notes} |\n\n### Memory Access Patterns\n\n- Stride pattern: {stride_pattern}\n- Cache modifiers: {cache_modifiers}\n- In-place operations: {in_place_ops}\n- Saved for backward vs recomputed: {save_vs_recompute}\n\n### Compute Patterns\n\n- Precision-sensitive ops: {precision_ops}\n- Casting mode: {casting_mode}\n- Algorithm type: {algorithm_type}\n- constexpr parameters: {constexpr_params}\n\n## Identified Optimization Opportunities\n\n{opportunity_list}\n\n## Recommended Strategy Order\n\nBased on the bottleneck classification ({bottleneck_class}) and kernel analysis:\n\n{strategy_order}\n\n### Estimated Iterations\n\n- Expected variants: {estimated_variants}\n- Rationale: {variants_rationale}\n\n---\n\n*Profile saved to `optimization/{kernel}/profile.md`*\n","isInternal":false,"tokens":1249,"sizeBytes":4231},{"name":"variant-notes.md","path":".agents/skills/liger-kernel-perf/templates/variant-notes.md","rawUrl":"https://raw.githubusercontent.com/linkedin/Liger-Kernel/HEAD/.agents/skills/liger-kernel-perf/templates/variant-notes.md","title":"Templates Skill","category":"anthropic-skill","format":"markdown","content":"# Variant {variant_number}: {variant_title}\n\n> Kernel: `{kernel}`\n> Date: {date}\n> Parent variant: {parent_variant}\n\n## Hypothesis\n\n**Target bottleneck:** {target_bottleneck}\n\n**Why this change should help:**\n\n{hypothesis_reasoning}\n\n**References from prior variants:**\n\n{prior_variant_references}\n\n## Changes\n\n{changes_list}\n\n## Expected Outcome\n\n**Predicted improvement:** {predicted_improvement}\n\n**Reasoning:** {prediction_reasoning}\n\n## Actual Results\n\n### Speed (median ms)\n\n| Metric | x_value | Baseline (v0) | This Variant (v{variant_number}) | Delta % |\n|--------|---------|---------------|----------------------------------|---------|\n| Speed forward | {x_val_1} | {baseline_fwd_1} | {variant_fwd_1} | {delta_fwd_1}% |\n| Speed forward | {x_val_2} | {baseline_fwd_2} | {variant_fwd_2} | {delta_fwd_2}% |\n| Speed backward | {x_val_1} | {baseline_bwd_1} | {variant_bwd_1} | {delta_bwd_1}% |\n| Speed backward | {x_val_2} | {baseline_bwd_2} | {variant_bwd_2} | {delta_bwd_2}% |\n| Speed full | {x_val_1} | {baseline_full_1} | {variant_full_1} | {delta_full_1}% |\n| Speed full | {x_val_2} | {baseline_full_2} | {variant_full_2} | {delta_full_2}% |\n\n### Memory (median MB)\n\n| Metric | x_value | Baseline (v0) | This Variant (v{variant_number}) | Delta % |\n|--------|---------|---------------|----------------------------------|---------|\n| Memory full | {x_val_1} | {baseline_mem_1} | {variant_mem_1} | {delta_mem_1}% |\n| Memory full | {x_val_2} | {baseline_mem_2} | {variant_mem_2} | {delta_mem_2}% |\n\n## Guardrail Checks\n\n| Check | Threshold | Result | Status |\n|-------|-----------|--------|--------|\n| Non-target metric regression | <5% worse | {nontarget_result} | {nontarget_status} |\n| Cross-pass regression | <10% on one pass to marginally improve other | {crosspass_result} | {crosspass_status} |\n| Smoke test (single shape, float32, fwd+bwd) | Pass | {smoke_result} | {smoke_status} |\n\n## Verdict\n\n**{verdict}**\n\n{verdict_explanation}\n\n## Learnings\n\n{learnings}\n\n---\n\n*Variant code: `optimization/{kernel}/{kernel}_v{variant_number}.py`*\n*Benchmark data: `optimization/{kernel}/benchmarks/v{variant_number}_results.csv`*\n","isInternal":false,"tokens":627,"sizeBytes":2133}],"systemPromptSnippet":"<yakaai_skills repo=\"linkedin/Liger-Kernel\">\n<!-- File: AGENTS.md (Tokens: ~572 | Category: root-instruction) -->\n# AGENTS.md\n\nGuidance for AI coding assistants working in this repository.\n\n## Skills\n\nRepository-specific workflow guides (\"skills\") live in [`.agents/skills/`](.agents/skills/). Each subdirectory is a self-contained guide for a multi-stage workflow. Read the `SKILL.md` at the top of each subdirectory for an overview; the other files in that directory are referenced from `SKILL.md` and should be read on demand.\n\n| Skill | What it does |\n|-------|--------------|\n| [`liger-kernel-dev`](.agents/skills/liger-kernel-dev/SKILL.md) | Develops new Triton kernels from a PyTorch reference (or modifies existing kernels). 3-stage pipeline: Analyze → Generate → Validate. NVIDIA GPUs only. |\n| [`liger-autopatch`](.agents/skills/liger-autopatch/SKILL.md) | Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies an existing monkey-patch. 3-stage pipeline: Analyze → Generate → Validate. |\n| [`liger-kernel-perf`](.agents/skills/liger-kernel-perf/SKILL.md) | Optimizes the performance of an existing Liger Triton kernel. 3-stage pipeline: Profile → Optimize → Finalize. NVIDIA GPUs only. |\n\nThe skills are written to be runtime-agnostic — they describe the workflow as a sequence of stages a competent agent (or human) can follow. Where a stage says \"Follow the X workflow in `x.md`\", that's a directive to read and execute that file's instructions; runtimes that support parallel subagents may delegate the stage, but it is not required.\n\n## Vendor-specific shortcuts\n\nFor convenience, some assistants auto-discover skills from vendor-specific paths. These point at the canonical `.agents/skills/` directory:\n\n- `.claude/skills` → symlink → `.agents/skills` (for Claude Code)\n\nIf you're adding support for another assistant, add a symlink (or your tool's preferred adapter) pointing to `.agents/skills/`. Do not duplicate the content.\n\n## Repo conventions\n\n- Source layout: `src/liger_kernel/{ops,transformers}/` for Triton ops and `nn.Module` / HF wrappers respectively\n- Tests: `test/transformers/` (unit) and `test/convergence/{bf16,fp32}/` (model convergence)\n- Benchmarks: `benchmark/scripts/` (scripts) and `benchmark/data/all_benchmark_data.csv` (results)\n- Lint/format: `make checkstyle` (uses `ruff`)\n- Install dev mode: `pip install -e \".[dev]\"`\n\nSee `README.md` for the project overview and contribution guide.\n\n<!-- File: .agents/skills/liger-autopatch/code-generator.md (Tokens: ~2184 | Category: anthropic-skill) -->\n# Code Generator Agent\n\nTakes a confirmed model profile (create mode) or a change plan (modify mode) and generates or modifies files for Liger Kernel support.\n\n## Mode\n\n- **Create mode** (default): Generating all files for a new model. Follow the full \"Files to Generate\" list below.\n- **Modify mode**: Making targeted changes to an existing monkey-patch. Follow the \"Modification Checklist\" section instead.\n\n## Pre-Requisites\n\nBefore generating, read the reference implementation closest to this model:\n- Dense → `src/liger_kernel/transformers/model/llama.py`\n- MoE → `src/liger_kernel/transformers/model/mixtral.py`\n- Vision-Language → `src/liger_kernel/transformers/model/qwen2_vl.py`\n- Gemma-family → `src/liger_kernel/transformers/model/gemma.py`\n\nAlso read the corresponding patching function in `monkey_patch.py` and the templates in [templates/](templates/).\n\n## Files to Generate\n\n### 1. `src/liger_kernel/transformers/model/{model_type}.py` (NEW)\n\nThe `lce_forward` function. See [templates/lce-forward-dense.md](templates/lce-forward-dense.md) or [templates/lce-forward-moe.md](templates/lce-forward-moe.md).\n\nKey rules:\n- Match the exact forward signature from HF's `ForCausalLM.forward`\n- Use `lce_maybe_trainable_lm_head` from `llama.py` (shared PEFT/FSDP utility)\n- If model needs custom loss args (e.g., softcapping), write a local helper instead\n\n### 2. `src/liger_kernel/transformers/monkey_patch.py` (MODIFY)\n\nThree changes — see [templates/monkey-patch-fn.md](templates/monkey-patch-fn.md):\n\n**A.** Add lce_forward import (~line 18-28):\n```python\nfrom liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n```\n\n**B.** Add `apply_liger_kernel_to_{model_type}` function with both class-level and instance-level patching paths.\n\n**C.** Add entry to `MODEL_TYPE_TO_APPLY_LIGER_FN` dict (~line 3067).\n\n### 3. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the function in three locations (maintain alphabetical order):\n- `TYPE_CHECKING` block\n- `__getattr__` monkey_patch_symbols set\n- `__all__` list extension\n\n### 4. `src/liger_kernel/transformers/model/output_classes.py` (MODIFY if needed)\n\nOnly for models needing custom output (MoE with `aux_loss`, VL with `rope_deltas`). Follow the existing guarded-import pattern in the file.\n\n### 5. `test/transformers/test_monkey_patch.py` (MODIFY)\n\nSee [templates/test-instance-patch.md](templates/test-instance-patch.md). Add availability checker + skipif-decorated test function using `inspect.getsource()` assertions.\n\n### 6. Convergence tests (MODIFY multiple files)\n\nSee [templates/test-convergence.md](templates/test-convergence.md). Every model needs entries in multiple convergence test files:\n\n**All text models (dense + MoE)** — add to these 4 files:\n- `test/convergence/bf16/test_mini_models.py` — FLCE path, bf16\n- `test/convergence/bf16/test_mini_models_with_logits.py` — non-FLCE path (tests RMSNorm/SwiGLU/RoPE only), bf16\n- `test/convergence/fp32/test_mini_models.py` — FLCE path, fp32\n- `test/convergence/fp32/test_mini_models_with_logits.py` — non-FLCE path, fp32\n\n**Vision-language models** — also add to these 2:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\nEach file needs: imports, availability guard, `MiniModelConfig` entry in `MINI_MODEL_SETUPS` dict, and a `pytest.param` entry in the parametrize block. The `MiniModelConfig` entry is identical across all files for the same model. The `pytest.param` tolerances differ — use bf16 tolerances (looser) for bf16 files and fp32 tolerances (tighter) for fp32 files. Copy tolerance values from a similar existing model (e.g., Llama for dense, Mixtral for MoE).\n\n### 7. `test/utils.py` (MODIFY)\n\nAdd `revert_liger_kernel_to_{model_type}` function that reloads the modeling module.\n\n### 8. `README.md` (MODIFY)\n\nAdd row to the Patching table under \"### Patching\":\n```\n| {ModelName} | `liger_kernel.transformers.apply_liger_kernel_to_{model_type}` | {Supported Operations} |\n```\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order** alongside similar existing entries. Never append to the end of a section — find the correct alphabetical position.\n- After generating all files, run `make checkstyle` to verify formatting. If it fails, run `ruff check . --fix && ruff format .` to auto-fix, then verify with `make checkstyle` again.\n\n## Modification Checklist (Modify Mode)\n\nBefore making changes, read the existing implementation:\n1. Read `apply_liger_kernel_to_{model_type}` in `monkey_patch.py`\n2. Read the existing test in `test_monkey_patch.py` for this model\n3. Read the relevant HF modeling source for context\n\n### Rules for All Modifications\n\n**R1. Both patching levels.** If adding a new kernel, it must appear in BOTH:\n  - Class-level patching (the main body of `apply_liger_kernel_to_{model_type}`)\n  - Instance-level patching (the `if model is not None` block)\n\n  Omitting one is the most common mistake.\n\n**R2. New parameter with default.** Every new kernel gets a bool parameter on the\n  apply function signature (e.g., `relu_squared: bool = True`). Default should be `True`\n  for kernels that are safe to enable by default, `False` otherwise.\n\n**R3. Update docstring.** Update the function's docstring to:\n  - Add an `Args` entry for the new parameter\n  - Remove any stale notes that the new kernel invalidates\n    (e.g., \"squared ReLU is not supported\" → remove if you're adding it)\n\n**R4. Update tests.** In the existing `test_apply_liger_kernel_to_instance_for_{model_type}`:\n  - Add import for the new Liger kernel class\n  - Add \"not yet patched\" assertion before `_apply_liger_kernel_to_instance`\n  - Add \"correctly patched\" assertion after\n  - Follow the exact pattern of existing assertions in the same test\n\n**R5. Check revert function.** Read `revert_liger_kernel_to_{model_type}` in `test/utils.py`.\n  The revert function uses `importlib.reload(modeling_{model_type})` to undo all patches.\n  This handles most cases automatically, but check if the new kernel requires additional\n  revert logic (e.g., if the kernel patches something outside the modeling module, or\n  replaces a global like `ACT2FN` that `importlib.reload` won't fully restore). Update\n  the revert function if needed.\n\n**R6. Run convergence tests.** Don't modify convergence test files unless the change\n  requires it (e.g., new mini model config fields). But DO run existing convergence\n  tests in the Validate stage to verify no regression. This is critical — the Validator\n  agent (Stage 3) handles this, but if you are generating code without a separate\n  Validate stage, run these yourself:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs\n  ```\n  For VL (multimodal) models, also run:\n  ```bash\n  pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs\n  ```\n\n**R7. Update README.md.** If the change adds a visibly new capability to the model's\n  row in the patching table (e.g., a new operation), update the supported operations list.\n\n### Common Modification Patterns\n\n**Adding an activation kernel (e.g., relu_squared for nemotron):**\n- Import the Liger kernel class at the top of `monkey_patch.py`\n- Add bool parameter to apply function signature\n- Class-level: replace in `ACT2FN` dict or replace the MLP class\n- Instance-level: patch each `decoder_layer`'s activation/MLP\n- Test: `assert isinstance` checks on the activation/MLP\n\n**Adding a norm variant:**\n- Add bool parameter to apply function\n- Class-level: replace the norm class\n- Instance-level: use `_patch_rms_norm_module` or `_patch_layer_norm_module` on all norm attrs\n- Test: `assert isinstance` checks on norm modules\n\n**Fixing missing instance patching:**\n- Read the class-level patching to see what's patched\n- Add corresponding instance-level patches in the `if model is not None` block\n- Test: add assertions that were missing\n\n**Updating for upstream HF changes:**\n- Compare the current HF modeling file against what the patch assumes\n- Update class names, attribute names, forward signatures as needed\n- May require updating `lce_forward` if the base model's forward changed\n\n<!-- File: .agents/skills/liger-autopatch/decision-matrix.md (Tokens: ~1384 | Category: anthropic-skill) -->\n# Decision Matrix\n\nWhen analyzing a HuggingFace model for Liger Kernel support, you must resolve these 12 architectural decisions by reading the model's `modeling_*.py` source code.\n\n## 1. Norm Type\n\n**Question:** Does the model use RMSNorm, LayerNorm, or both?\n\n**How to detect:**\n- Search for `class *RMSNorm` in the modeling file → RMSNorm\n- Search for `nn.LayerNorm` usage → LayerNorm\n- Multimodal models often use both (RMSNorm for text, LayerNorm for vision)\n\n**Liger mapping:** `LigerRMSNorm` or `LigerLayerNorm`\n\n## 2. RMSNorm Casting Mode\n\n**Question:** How does the model handle dtype casting during normalization?\n\n**How to detect:** Read the RMSNorm forward method:\n- Casts input to fp32, computes variance, casts back → `\"gemma\"`\n- Computes variance in fp32 only (input stays original dtype) → `\"llama\"`\n- No casting at all → `\"none\"`\n\n**Default:** `\"llama\"` (most common)\n\n## 3. RMSNorm Offset\n\n**Question:** Does the weight have a +1.0 offset?\n\n**How to detect:** In the RMSNorm forward, look for `(1 + self.weight)` or `self.weight + 1`:\n- Present → `offset=1.0` (Gemma family)\n- Absent → `offset=0.0` (most models)\n\n## 4. RMSNorm In-Place\n\n**Question:** Can the backward pass modify dY in-place?\n\n**How to detect:** Check if the model has two sequential norm layers with a residual connection between them (like Gemma2's `pre_feedforward_layernorm` + `post_feedforward_layernorm`):\n- Sequential norms with residual → `in_place=False`\n- Otherwise → `in_place=True`\n\n## 5. MLP Activation Type\n\n**Question:** What activation function does the gated MLP use?\n\n**How to detect:** Read the MLP class forward method:\n- `silu` or `F.silu` → SwiGLU → `LigerSwiGLUMLP`\n- `gelu` or `gelu_new` or `gelu_fast` → GeGLU → `LigerGEGLUMLP`\n- Phi3-style (single gate+up projection split) → `LigerPhi3SwiGLUMLP`\n\n**Also check:** The config's `hidden_act` field.\n\n## 6. Dense vs MoE\n\n**Question:** Is the model dense, MoE, or hybrid (some layers dense, some MoE)?\n\n**How to detect:**\n- Search for `Expert`, `MoE`, `SparseMoe`, `TopK` routing classes\n- Check if decoder layers have a `block_sparse_moe` or `experts` attribute\n- Hybrid: check for `is_moe_layer` or conditional MoE per-layer\n\n**Liger mapping:**\n- Dense → standard patching\n- MoE (transformers v5) → `LigerExperts`\n- MoE (transformers v4) → `LigerBlockSparseTop2MLP`\n- Qwen3-style MoE → `LigerQwen3MoeSwiGLUMLP`\n\n## 7. Vision Components\n\n**Question:** Does the model have a vision encoder?\n\n**How to detect:**\n- Check for `pixel_values` in the `forward` signature\n- Look for a separate vision model class (e.g., `*VisionModel`)\n- Check config for `vision_config` or `text_config` sub-configs\n\n**If yes:**\n- Vision encoder norms are usually `nn.LayerNorm` → patch with `LigerLayerNorm`\n- Text and vision must be patched separately\n\n## 8. RoPE Variant\n\n**Question:** What type of positional embedding does the model use?\n\n**How to detect:** Search for the `apply_rotary_pos_emb` function:\n- Standard (q, k, cos, sin) → `liger_rotary_pos_emb` (rope=True)\n- Llama4-style → `liger_llama4_text_rotary_pos_emb`\n- Qwen2VL MRoPE → `liger_multimodal_rotary_pos_emb`\n- No rotary embedding or custom variant → `rope=False`\n\n## 9. Output Class\n\n**Question:** What return type does the model's ForCausalLM.forward use?\n\n**How to detect:** Read the return statement and type annotation:\n- Standard → `LigerCausalLMOutputWithPast`\n- MoE (has `aux_loss`) → `LigerMoeCausalLMOutputWithPast`\n- Custom VL output → create model-specific output class in `output_classes.py`\n\n## 10. Hidden State Access\n\n**Question:** How does the model access hidden states from base model output?\n\n**How to detect:** In the ForCausalLM.forward, after calling `self.model(...)`:\n- `outputs[0]` → most models (Llama, Mistral, Gemma, etc.)\n- `outputs.last_hidden_state` → Phi3, Qwen3.5 MoE, some newer models\n\n## 11. Logit Softcapping\n\n**Question:** Does the model apply softcapping to logits before loss?\n\n**How to detect:** Check config for `final_logit_softcapping`:\n- Present → pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`\n- Absent → no softcapping (most models)\n- **VL models:** Config path may be `self.config.text_config.final_logit_softcapping` instead of `self.config.final_logit_softcapping`. Check whether the model uses a composite config with `text_config` sub-config.\n\n**Models with softcapping:** Gemma2, Gemma3\n\n## 12. Decoder Layer Norm Names\n\n**Question:** What are the attribute names of norm layers in each decoder layer?\n\n**How to detect:** Read the decoder layer class `__init__`:\n- Standard: `input_layernorm`, `post_attention_layernorm`\n- Gemma2 extra: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- GLM4: `post_self_attn_layernorm`, `post_mlp_layernorm`\n- Some models: `q_norm`, `k_norm` on self_attn\n\nAlso check the final norm on the base model (usually `model.norm` or `model.final_layernorm`).\n\n<!-- File: .agents/skills/liger-autopatch/examples/gemma-profile.md (Tokens: ~657 | Category: anthropic-skill) -->\n# Model Profile: Gemma\n\nThis profile demonstrates the key differences from Llama: GeGLU activation, RMSNorm offset, and Gemma casting mode.\n\n## Identity\n- model_type: gemma\n- causal_lm_class: GemmaForCausalLM\n- base_model_class: GemmaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.gemma.modeling_gemma\n- config_module: transformers.models.gemma.configuration_gemma\n\n## Normalization\n- norm_class: GemmaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: gemma (everything cast to fp32, then computed, then cast back)\n- offset: 1.0 (weight uses `1 + self.weight` pattern)\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: GemmaMLP\n- activation: gelu (uses GELU activation, not SiLU)\n- liger_mlp_class: LigerGEGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nSame as Llama — no extra parameters.\n\n## Mini Model Config\n```python\nGemmaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-6,\n    hidden_activation=\"gelu_pytorch_tanh\",\n)\n```\n\n## Key Differences from Llama\n\n1. **Activation**: Uses `geglu` parameter (not `swiglu`) in the patch function\n2. **RMSNorm**: Requires `offset=1.0` and `casting_mode=\"gemma\"`\n3. **MLP class**: `LigerGEGLUMLP` instead of `LigerSwiGLUMLP`\n4. **Patching uses partial**: `_patch_rms_norm_module_for_gemma = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)`\n\n## Gemma2 Additional Differences\n- `in_place: false` (residual between sequential norms)\n- Extra norm layers: `pre_feedforward_layernorm`, `post_feedforward_layernorm`\n- Has `final_logit_softcapping` in config\n- Uses `LigerRMSNormForGemma2` variant\n\n<!-- File: .agents/skills/liger-autopatch/examples/llama-profile.md (Tokens: ~549 | Category: anthropic-skill) -->\n# Model Profile: Llama\n\nThis is the reference profile for the simplest and most common model type.\n\n## Identity\n- model_type: llama\n- causal_lm_class: LlamaForCausalLM\n- base_model_class: LlamaModel\n- base_model_prefix: \"model\"\n- modeling_module: transformers.models.llama.modeling_llama\n- config_module: transformers.models.llama.configuration_llama\n\n## Normalization\n- norm_class: LlamaRMSNorm\n- norm_type: RMSNorm\n- casting_mode: llama\n- offset: 0.0\n- in_place: true\n- final_norm_attr: model.norm\n- decoder_norm_attrs:\n  - input_layernorm\n  - post_attention_layernorm\n- attn_norm_attrs: none\n\n## MLP\n- mlp_class: LlamaMLP\n- activation: silu\n- liger_mlp_class: LigerSwiGLUMLP\n- gate_proj_attr: gate_proj\n- up_proj_attr: up_proj\n- down_proj_attr: down_proj\n\n## Structure\n- type: dense\n- moe_expert_class: n/a\n- moe_router_class: n/a\n- shared_expert: false\n\n## Positional Embedding\n- rope_type: standard\n- rope_function: apply_rotary_pos_emb\n\n## Output\n- output_class: LigerCausalLMOutputWithPast\n- hidden_state_access: outputs[0]\n- has_logit_softcapping: false\n- softcapping_config_attr: none\n\n## Vision\n- has_vision: false\n\n## Forward Signature\nStandard parameters only:\n- input_ids, attention_mask, position_ids, past_key_values\n- inputs_embeds, labels, use_cache\n- output_attentions, output_hidden_states, return_dict\n- cache_position, logits_to_keep, skip_logits\n- **kwargs\n\nNo model-specific extra parameters.\n\n## Mini Model Config\n```python\nLlamaConfig(\n    hidden_size=32,\n    intermediate_size=64,\n    num_hidden_layers=2,\n    num_attention_heads=2,\n    num_key_value_heads=2,\n    vocab_size=1024,\n    rms_norm_eps=1e-5,\n)\n```\n\n## Notes\n- Llama is the base reference for all dense LLaMA-family models\n- Most models (Mistral, Qwen2, OLMo, etc.) follow this exact pattern\n- The `lce_forward` for Llama reuses `lce_maybe_trainable_lm_head` for PEFT/FSDP support\n- Has `pretraining_tp` check (raises error if > 1)\n\n<!-- File: .agents/skills/liger-autopatch/model-analyzer.md (Tokens: ~731 | Category: anthropic-skill) -->\n# Model Analyzer Agent\n\nReads a HuggingFace model's source code and produces a structured model profile.\n\n## Input\n\nA model identifier: model_type string (e.g., `\"nemotron\"`), class name (e.g., `\"NemotronForCausalLM\"`), or Hub model name.\n\n## Steps\n\n1. Locate `transformers/models/{model_type}/modeling_{model_type}.py` and `configuration_{model_type}.py` on disk\n2. Read [decision-matrix.md](decision-matrix.md) for the 12 decisions to resolve\n3. Read the modeling file and extract all architectural details\n4. Produce the profile in the format below\n\n## What to Extract\n\nFrom the modeling file, identify:\n\n- `*ForCausalLM`, `*Model`, `*DecoderLayer`, `*RMSNorm`/norm, `*MLP`, `*Expert`/`*MoE` classes\n- `ForCausalLM.forward` signature — note all parameters\n- Norm class `forward` — determine casting mode and offset\n- All norm attribute names on decoder layer and base model\n- MLP activation function and gate/up/down projection names\n- Config's `model_type` string, `final_logit_softcapping`, `pretraining_tp`\n\n## Output Format\n\n```markdown\n# Model Profile: {ModelName}\n\n## Identity\n- model_type: {from config}\n- causal_lm_class: {e.g., NemotronForCausalLM}\n- base_model_class: {e.g., NemotronModel}\n- base_model_prefix: {e.g., \"model\"}\n- modeling_module: transformers.models.{model_type}.modeling_{model_type}\n\n## Normalization\n- norm_type: RMSNorm | LayerNorm | both\n- casting_mode: llama | gemma | none\n- offset: 0.0 | 1.0\n- in_place: true | false\n- final_norm_attr: {e.g., \"model.norm\"}\n- decoder_norm_attrs: {list}\n- attn_norm_attrs: {e.g., \"self_attn.q_norm\" or \"none\"}\n\n## MLP\n- activation: silu | gelu | gelu_new\n- liger_mlp_class: LigerSwiGLUMLP | LigerGEGLUMLP | LigerPhi3SwiGLUMLP\n\n## Structure\n- type: dense | moe | hybrid\n- moe_expert_class: {if MoE}\n- shared_expert: true | false\n\n## Positional Embedding\n- rope_type: standard | llama4 | qwen2vl_mrope | none\n\n## Output\n- output_class: LigerCausalLMOutputWithPast | LigerMoeCausalLMOutputWithPast | custom\n- hidden_state_access: outputs[0] | outputs.last_hidden_state\n- has_logit_softcapping: true | false\n\n## Vision\n- has_vision: true | false\n- vision_norm_type: {if applicable}\n\n## Forward Signature\n{List model-specific parameters not in standard Llama signature}\n\n## Mini Model Config\n{Suggest minimal config values for testing — 2 layers, 32 hidden, 64 intermediate}\n```\n\n## Rules\n\n- Never guess — always read the source to determine each value\n- Read [examples/llama-profile.md](examples/llama-profile.md) and [examples/gemma-profile.md](examples/gemma-profile.md) as reference\n- Compare against existing supported models to find the closest match\n\n<!-- File: .agents/skills/liger-autopatch/SKILL.md (Tokens: ~1269 | Category: anthropic-skill) -->\n---\nname: liger-autopatch\ndescription: \"Adds Liger Kernel support for a new HuggingFace Transformers model, or modifies existing monkey-patching. Generates lce_forward, monkey-patch function, tests, and README entry. Use when adding a new model to Liger Kernel, when a user asks to patch an unsupported model, when extending MODEL_TYPE_TO_APPLY_LIGER_FN, or when modifying/updating/fixing an existing monkey-patch (e.g., adding a new kernel to an already-supported model, fixing instance patching, updating a patch for upstream HF changes).\"\n---\n\n# Liger Auto-Patch\n\nAdds Liger Kernel optimization support for a new HuggingFace model, or modifies existing monkey-patching, through a staged pipeline with human review between stages. Supports creating new model patches and modifying existing ones.\n\n## Mode Detection\n\n- **Create mode**: User asks to add/patch/support a new model → full pipeline (Analyze → Generate → Validate)\n- **Modify mode**: User asks to update/fix/change/extend an existing monkey-patch → lighter pipeline (Change Impact Analysis → Apply Changes → Validate)\n\nKeywords that suggest modify mode: update, fix, change, add [kernel] to [existing model], extend, modify, new activation, new norm, bug in patch, upstream changed\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Model Analyzer** workflow in [model-analyzer.md](model-analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nThis stage reads the HF `modeling_*.py` source and produces a **model profile** answering 12 architectural questions from [decision-matrix.md](decision-matrix.md).\n\n**Human checkpoint:** Present the profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md).\n\nGenerates/modifies up to 13 files:\n\n1. `src/liger_kernel/transformers/model/{model}.py` — NEW lce_forward\n2. `src/liger_kernel/transformers/monkey_patch.py` — MODIFY\n3. `src/liger_kernel/transformers/__init__.py` — MODIFY\n4. `src/liger_kernel/transformers/model/output_classes.py` — MODIFY if needed\n5. `test/transformers/test_monkey_patch.py` — MODIFY\n6. `test/convergence/bf16/test_mini_models.py` — MODIFY (FLCE path)\n7. `test/convergence/bf16/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n8. `test/convergence/fp32/test_mini_models.py` — MODIFY (FLCE path)\n9. `test/convergence/fp32/test_mini_models_with_logits.py` — MODIFY (non-FLCE path)\n10. `test/convergence/bf16/test_mini_models_multimodal.py` — MODIFY if VL model\n11. `test/convergence/fp32/test_mini_models_multimodal.py` — MODIFY if VL model\n12. `test/utils.py` — MODIFY\n13. `README.md` — MODIFY\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns instance patching test, convergence test, and lint check. Retries up to 3 times on failure.\n\n**Human checkpoint:** Report final test results.\n\n## Pipeline (Modify Mode)\n\n### Stage 1: Change Impact Analysis\n\nRead the existing `apply_liger_kernel_to_{model_type}` function in `monkey_patch.py` and the relevant section of the upstream HF `modeling_{model_type}.py`. Produce a short change plan:\n\n- What is being added/changed/fixed\n- Which Liger kernel(s) are involved\n- Which files need modification (subset of the 13 files from create mode)\n- What the expected behavior should be after the change\n\n**Human checkpoint:** Present the change plan. Confirm before proceeding.\n\n### Stage 2: Apply Changes\n\nFollow the **Code Generator** workflow in [code-generator.md](code-generator.md) in **modify mode**.\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md). This stage is **mandatory** — do not skip it. At minimum, run:\n\n1. Instance patching test: `pytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs`\n2. All convergence tests for the model:\n   - `pytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, bf16)\n   - `pytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, bf16)\n   - `pytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs` (FLCE, fp32)\n   - `pytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs` (non-FLCE, fp32)\n   - If VL (multimodal) model, also run:\n     - `pytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n     - `pytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs`\n3. Checkstyle: `make checkstyle`\n\n**Human checkpoint:** Report final test results.\n\n## Reference Files\n\n- [decision-matrix.md](decision-matrix.md) — 12 architectural decisions to resolve per model\n- [examples/llama-profile.md](examples/llama-profile.md) — Reference profile for standard dense model\n- [examples/gemma-profile.md](examples/gemma-profile.md) — Reference profile showing GeGLU + offset variant\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n\n<!-- File: .agents/skills/liger-autopatch/templates/lce-forward-dense.md (Tokens: ~457 | Category: anthropic-skill) -->\n# Template: lce_forward for Dense Models\n\nReference implementation: `src/liger_kernel/transformers/model/llama.py`. Read it first.\n\n## Structure\n\nThe file needs `lce_forward` and reuses `lce_maybe_trainable_lm_head` from `llama.py`.\n\n### Imports\n\n```python\nfrom typing import TYPE_CHECKING, List, Optional, Tuple, Union\nimport torch\nfrom liger_kernel.transformers.model.llama import lce_maybe_trainable_lm_head\nfrom liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss, unpack_cross_entropy_result\nfrom liger_kernel.transformers.model.output_classes import LigerCausalLMOutputWithPast\nif TYPE_CHECKING:\n    from transformers.cache_utils import Cache\n```\n\n### lce_forward body pattern\n\n1. Set defaults from config (output_attentions, output_hidden_states, return_dict)\n2. Call `self.model(...)` with all forwarded params\n3. Extract hidden states: `outputs[0]` or `outputs.last_hidden_state`\n4. Slice for `logits_to_keep`\n5. Extract `shift_labels = kwargs.pop(\"shift_labels\", None)`\n6. `skip_logits` logic — default True during training with labels\n7. Compute loss via `lce_maybe_trainable_lm_head` (skip_logits) or `self.lm_head` (else)\n8. Return `LigerCausalLMOutputWithPast`\n\n### Variations\n\n- **Logit softcapping**: Pass `final_logit_softcapping=self.config.final_logit_softcapping` to `LigerForCausalLMLoss`. Requires a local helper instead of reusing `lce_maybe_trainable_lm_head`.\n- **Hidden state access**: Use `outputs.last_hidden_state` if HF model uses that\n- **pretraining_tp**: Remove check if model doesn't have this config field\n\n### When NOT to reuse llama's lce_maybe_trainable_lm_head\n\nIf the model passes extra kwargs to `LigerForCausalLMLoss` (like softcapping), write a local `_liger_for_causal_lm_loss` and inline the PEFT/FSDP handling from llama.py.\n\n<!-- File: .agents/skills/liger-autopatch/templates/lce-forward-moe.md (Tokens: ~319 | Category: anthropic-skill) -->\n# Template: lce_forward for MoE Models\n\nReference implementation: `src/liger_kernel/transformers/model/mixtral.py`. Read it first.\n\n## Differences from Dense\n\n1. **Output class**: `LigerMoeCausalLMOutputWithPast` (includes `aux_loss`)\n2. **Extra param**: `output_router_logits` in forward signature\n3. **Auxiliary loss**: Compute `load_balancing_loss_func` from router logits\n4. **Return**: Include `aux_loss` and `router_logits` in output\n\n## Auxiliary Loss Pattern\n\n```python\naux_loss = None\nif output_router_logits:\n    router_logits = outputs.router_logits if return_dict else outputs[-1]\n    aux_loss = load_balancing_loss_func(\n        router_logits, self.num_experts, self.num_experts_per_tok, attention_mask,\n    )\n    if labels is not None and loss is not None:\n        loss += self.router_aux_loss_coef * aux_loss\n```\n\n## Monkey Patch Differences for MoE\n\nExpert patching is version-aware:\n\n```python\n# Transformers v5\nmodeling_module.{ModelExperts} = LigerExperts\n# Instance: _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\n\n# Transformers v4\nmodeling_module.{ModelExpertMLP} = LigerBlockSparseTop2MLP\n# Instance: for expert in decoder_layer.block_sparse_moe.experts: _patch_swiglu_module(...)\n```\n\nCheck exact attribute names in the HF decoder layer class.\n\n<!-- File: .agents/skills/liger-autopatch/templates/monkey-patch-fn.md (Tokens: ~976 | Category: anthropic-skill) -->\n# Template: apply_liger_kernel_to_{model_type}\n\nReference: `apply_liger_kernel_to_llama` (dense) or `apply_liger_kernel_to_mixtral` (MoE) in `monkey_patch.py`.\n\n## Function Signature\n\n```python\ndef apply_liger_kernel_to_{model_type}(\n    rope: bool = True,              # False if model lacks standard RoPE\n    cross_entropy: bool = False,\n    fused_linear_cross_entropy: bool = True,\n    rms_norm: bool = True,          # Or layer_norm for LayerNorm models\n    swiglu: bool = True,            # Or geglu for Gemma-family\n    model: PreTrainedModel = None,\n) -> None:\n```\n\n## Body Pattern\n\nClass-level patching always runs (even when `model` is provided). Instance-level patching runs additionally when `model is not None`. They are NOT mutually exclusive — this matches the existing Llama/Mistral/Gemma patterns.\n\nThe only exception is `fused_linear_cross_entropy`: use `MethodType` binding when instance is given, class replacement otherwise.\n\n```python\nassert not (cross_entropy and fused_linear_cross_entropy), \"...\"\n\nfrom transformers.models.{model_type} import modeling_{model_type}\n\n# Class-level patching (always runs)\nif rope: modeling_{model_type}.apply_rotary_pos_emb = liger_rotary_pos_emb\nif rms_norm: modeling_{model_type}.{ModelRMSNorm} = LigerRMSNorm\nif swiglu: modeling_{model_type}.{ModelMLP} = LigerSwiGLUMLP\nif cross_entropy: modeling_{model_type}.CrossEntropyLoss = LigerCrossEntropyLoss\nif fused_linear_cross_entropy:\n    if model is not None:\n        model.forward = MethodType({model_type}_lce_forward, model)\n    else:\n        modeling_{model_type}.{ModelForCausalLM}.forward = {model_type}_lce_forward\n\n# Instance-level patching (only when model instance provided)\nif model is not None:\n    base_model = getattr(model, model.base_model_prefix, model)\n    if rms_norm: _patch_rms_norm_module(base_model.norm)\n    for decoder_layer in base_model.layers:\n        if swiglu: _patch_swiglu_module(decoder_layer.mlp, LigerSwiGLUMLP)\n        if rms_norm:\n            _patch_rms_norm_module(decoder_layer.input_layernorm)\n            _patch_rms_norm_module(decoder_layer.post_attention_layernorm)\n```\n\n## Common Variations\n\n**Gemma-family** (offset + casting):\n```python\n_patch_rms_norm_module_for_model = partial(_patch_rms_norm_module, casting_mode=\"gemma\", offset=1.0)\n```\n\n**MoE experts** (version-aware):\n```python\nif IS_TRANSFORMERS_V5_OR_LATER:\n    _patch_swiglu_module(decoder_layer.mlp.experts, LigerExperts)\nelse:\n    for expert in decoder_layer.block_sparse_moe.experts:\n        _patch_swiglu_module(expert, LigerBlockSparseTop2MLP)\n```\n\n**Multimodal VL instance patching:**\n\nFor VL models, extract text_model and vision_model from the composite instance using `isinstance` checks. Each VL model has different accessor paths — read the existing VL patches in `monkey_patch.py` for the exact pattern. Example from Llama4:\n\n```python\nif model is not None:\n    if isinstance(model, {Model}ForConditionalGeneration):\n        language_model = model.language_model          # or model.model.language_model\n        vision_model = model.vision_model              # or model.model.visual\n        text_model = language_model.model\n    elif isinstance(model, {Model}ForCausalLM):\n        text_model = model.model\n        vision_model = None\n    elif isinstance(model, {Model}TextModel):\n        text_model = model\n        vision_model = None\n    else:\n        raise ValueError(f\"Unsupported model type: {type(model)}\")\n\n    # Patch text layers\n    if text_model:\n        if rms_norm: _patch_rms_norm_module(text_model.norm)\n        for decoder_layer in text_model.layers:\n            ...\n\n    # Patch vision layers\n    if vision_model:\n        for layer in vision_model.encoder.layers:  # accessor varies by model\n            if layer_norm:\n                _patch_layer_norm_module(layer.layer_norm1)\n                _patch_layer_norm_module(layer.layer_norm2)\n```\n\nThe exact accessor paths (`model.language_model`, `model.model.language_model`, `model.model.visual`) vary per VL model. Always read the HF model's `__init__` to find the correct paths.\n\n<!-- File: .agents/skills/liger-autopatch/templates/test-convergence.md (Tokens: ~1180 | Category: anthropic-skill) -->\n# Template: Convergence Tests\n\nReference: existing entries across `test/convergence/` files. Read the Llama entries as baseline.\n\n## Which Files to Modify\n\nEvery model needs entries in **at least 4 files**:\n\n| File | What it tests | dtype |\n|------|-------------|-------|\n| `test/convergence/bf16/test_mini_models.py` | FLCE path (fused linear cross entropy) | bfloat16 |\n| `test/convergence/bf16/test_mini_models_with_logits.py` | Non-FLCE path (tests RMSNorm/SwiGLU/RoPE only, verifies logits) | bfloat16 |\n| `test/convergence/fp32/test_mini_models.py` | FLCE path | float32 |\n| `test/convergence/fp32/test_mini_models_with_logits.py` | Non-FLCE path | float32 |\n\n**Vision-language models** also need:\n- `test/convergence/bf16/test_mini_models_multimodal.py`\n- `test/convergence/fp32/test_mini_models_multimodal.py`\n\n## Placement Rule\n\nWhen adding entries to existing files, always insert in **alphabetical order** alongside similar existing entries. Never append to the end of a section. Specifically:\n- Imports: alphabetical among other `apply_liger_kernel_to_*` / `revert_liger_kernel_to_*` imports\n- Availability guards: alphabetical among other try/except model import blocks\n- `MINI_MODEL_SETUPS` entries: alphabetical among other model setup blocks\n- `pytest.param` entries: alphabetical among other model param entries\n\n## What to Add in Each File\n\nEach file has a `MINI_MODEL_SETUPS` dict and a `pytest.param` block. Add to both.\n\n### 1. Imports (same across all files, insert in alphabetical order)\n\n```python\nfrom liger_kernel.transformers import apply_liger_kernel_to_{model_type}\nfrom test.utils import revert_liger_kernel_to_{model_type}\n```\n\n### 2. Availability Guard (same across all files)\n\nThe `try/except ImportError` handles version compatibility automatically: if the model doesn't exist in the installed transformers version (e.g., 4.52.0), the import fails and the tests skip. No explicit version check needed.\n\n```python\ntry:\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n    from transformers.models.{model_type}.modeling_{model_type} import {ModelForCausalLM}\n    {MODEL_UPPER}_AVAILABLE = True\nexcept ImportError:\n    {MODEL_UPPER}_AVAILABLE = False\n```\n\n### 3. MINI_MODEL_SETUPS Entry (identical across all files)\n\n```python\nif {MODEL_UPPER}_AVAILABLE:\n    MINI_MODEL_SETUPS[\"mini_{model_type}\"] = MiniModelConfig(\n        liger_kernel_patch_func=apply_liger_kernel_to_{model_type},\n        liger_kernel_patch_revert_func=revert_liger_kernel_to_{model_type},\n        model_class={ModelForCausalLM},\n        mini_model_config={ModelConfig}(\n            hidden_size=32, intermediate_size=64, num_hidden_layers=2,\n            num_attention_heads=2, num_key_value_heads=2, vocab_size=1024,\n        ),\n    )\n```\n\n### 4. pytest.param Entry (tolerances differ by dtype)\n\n**bf16 files** — looser tolerances (copy from similar existing model):\n```python\npytest.param(\n    \"mini_{model_type}\", 32, 1e-5, torch.bfloat16,\n    1e-2, 5e-2, 1e-1, 1e-2, 1e-2, 1e-2,\n    marks=pytest.mark.skipif(not supports_bfloat16(), reason=\"bfloat16 not supported\"),\n),\n```\n\n**fp32 files** — tighter tolerances:\n```python\n(\"mini_{model_type}\", 32, 1e-4, torch.float32, 1e-8, 2e-5, 5e-3, 1e-5, 5e-3, 1e-5),\n```\n\n### 5. Revert Function in test/utils.py (shared by all files)\n\n```python\ndef revert_liger_kernel_to_{model_type}(model_config):\n    from transformers.models.{model_type} import modeling_{model_type}\n    importlib.reload(modeling_{model_type})\n    model_config.model_class = modeling_{model_type}.{ModelForCausalLM}\n```\n\n## Key Differences Between Test Variants\n\n**`with_logits`**: Sets `fused_linear_cross_entropy=False, cross_entropy=False` and checks `output.logits` match. Tests that non-loss kernels are numerically correct on their own. Same `MINI_MODEL_SETUPS` entry, just add a `pytest.param`.\n\n**`multimodal`**: Uses `multimodal_collate_fn`, image datasets, and model configs with vision sub-config. Only for models with `has_vision: true`. Multimodal tests may also require fake tokenizer/processor configs in `test/resources/fake_configs/`. Check existing VL model entries (e.g., Qwen2-VL, Llava) for the pattern — if the model needs a tokenizer/processor that can't be created programmatically, add a fake config JSON file.\n\n## Mini Config Guidelines\n\nUse minimal values: hidden_size=32, intermediate_size=64, num_hidden_layers=2, num_attention_heads=2, vocab_size=1024. Add model-specific required fields (e.g., `num_local_experts` for MoE). Copy config structure from the most similar existing model.\n\n<!-- File: .agents/skills/liger-autopatch/templates/test-instance-patch.md (Tokens: ~538 | Category: anthropic-skill) -->\n# Template: Instance Patching Test\n\nReference: `test_apply_liger_kernel_to_instance_for_llama` in `test/transformers/test_monkey_patch.py`.\n\n## Add Availability Checker (at top of file, near other `is_*_available` functions ~lines 59-240, in alphabetical order)\n\n```python\ndef is_{model_type}_available():\n    try:\n        import transformers.models.{model_type}\n        return True\n    except ImportError:\n        return False\n```\n\n## Add Test Function (after existing `test_apply_liger_kernel_to_instance_for_*` functions, in alphabetical order)\n\n```python\n@pytest.mark.skipif(not is_{model_type}_available(), reason=\"{model_type} not available\")\ndef test_apply_liger_kernel_to_instance_for_{model_type}():\n    from liger_kernel.transformers.model.{model_type} import lce_forward as {model_type}_lce_forward\n    from transformers.models.{model_type}.configuration_{model_type} import {ModelConfig}\n\n    config = {ModelConfig}(hidden_size=32, intermediate_size=64, num_hidden_layers=2)\n    dummy_model_instance = AutoModelForCausalLM.from_config(config)\n\n    # Verify NOT patched\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) != inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) != inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) != inspect.getsource(LigerRMSNorm.forward)\n\n    # Apply patch\n    _apply_liger_kernel_to_instance(model=dummy_model_instance)\n\n    # Verify patched\n    assert inspect.getsource(dummy_model_instance.forward) == inspect.getsource({model_type}_lce_forward)\n    assert inspect.getsource(dummy_model_instance.model.norm.forward) == inspect.getsource(LigerRMSNorm.forward)\n    for layer in dummy_model_instance.model.layers:\n        assert inspect.getsource(layer.mlp.forward) == inspect.getsource({LigerMLPClass}.forward)\n        assert inspect.getsource(layer.input_layernorm.forward) == inspect.getsource(LigerRMSNorm.forward)\n```\n\n## Adapt For\n\n- **GeGLU models**: Use `LigerGEGLUMLP` in assertions\n- **MoE models**: Check `decoder_layer.mlp.experts.forward` against `LigerExperts`\n- **Extra norms**: Add assertions for `pre_feedforward_layernorm` etc.\n- **Vision**: Add assertions for `LigerLayerNorm` on vision layers\n\n<!-- File: .agents/skills/liger-autopatch/validator.md (Tokens: ~942 | Category: anthropic-skill) -->\n# Validator Agent\n\nRuns tests on generated code and fixes issues iteratively.\n\n## Pre-Requisites\n\nVerify the project is installed in development mode: `python -c \"import liger_kernel\"` must succeed. If not, ask the user to run `pip install -e \".[dev]\"`.\n\n## Validation Steps\n\n### Step 1: Syntax Check\n\n```bash\npython -c \"from liger_kernel.transformers import apply_liger_kernel_to_{model_type}\"\n```\n\n### Step 2: Instance Patching Test\n\n```bash\npytest test/transformers/test_monkey_patch.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `ImportError: cannot import name` | Check exact class name in HF source |\n| `AttributeError: has no attribute` | Read decoder layer `__init__` for correct attr names |\n| `AssertionError: getsource mismatch` | Check class being replaced in monkey_patch |\n| `ModuleNotFoundError` | Add `@pytest.mark.skipif` with availability check |\n\n### Step 3: Convergence Tests (requires GPU)\n\nRun **all applicable** convergence tests. For text models, there are 4 files. For VL models, 6.\n\n**FLCE path (fused linear cross entropy):**\n```bash\npytest test/convergence/bf16/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Non-FLCE path (logits verification — tests RMSNorm/SwiGLU/RoPE without loss fusion):**\n```bash\npytest test/convergence/bf16/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_with_logits.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\n**Multimodal (only for VL models):**\n```bash\npytest test/convergence/bf16/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\npytest test/convergence/fp32/test_mini_models_multimodal.py -k \"{model_type}\" -xvs 2>&1 | tail -80\n```\n\nCommon failures:\n\n| Error | Fix |\n|-------|-----|\n| `CUDA out of memory` | Reduce hidden_size/num_layers in mini config |\n| `RuntimeError: shape mismatch` | Check `self.config.hidden_size` usage |\n| `AssertionError: allclose` | Check casting_mode, offset, in_place settings |\n| `KeyError in config` | Add missing field to mini model config |\n\n### Step 4: Checkstyle\n\nRun the project's full lint and format check:\n\n```bash\nmake checkstyle\n```\n\nIf it fails, auto-fix with `ruff check . --fix && ruff format .`, then re-run `make checkstyle`.\n\n## Retry Logic\n\n**Hard failures** (Steps 1, 2, 4 — import, patching, checkstyle): These indicate code bugs. Read traceback, identify root cause, fix, re-run. Max 3 retries per step. After 3 failures, report the exact error and stop.\n\n**Soft failures** (Step 3 — convergence): If a convergence test fails after 3 retries with `allclose` tolerance errors, mark it as `SKIP (tolerance tuning needed)` in the report rather than `FAIL`. This means the patching is correct but tolerances need manual adjustment. Note this in the PR description so reviewers are aware.\n\nConvergence failures from other errors (shape mismatch, CUDA OOM, KeyError) are still hard failures — fix and retry.\n\n## Success Report\n\n```\n| Check                        | Status |\n|------------------------------|--------|\n| Import                       | PASS   |\n| Instance Patching            | PASS   |\n| Convergence bf16 FLCE        | PASS   |\n| Convergence bf16 with_logits | PASS   |\n| Convergence fp32 FLCE        | PASS   |\n| Convergence fp32 with_logits | PASS   |\n| Convergence multimodal       | SKIP (text-only model) |\n| Checkstyle                   | PASS   |\n```\n\nInclude any `SKIP (tolerance tuning needed)` entries in the PR description under a \"Known issues\" section.\n\n<!-- File: .agents/skills/liger-kernel-dev/analyzer.md (Tokens: ~632 | Category: anthropic-skill) -->\n# Analyzer Agent\n\nUnderstands a PyTorch operation from any input form and produces a standalone PyTorch reference implementation + kernel profile.\n\n## Input Handling\n\nThe user may provide the operation in any form:\n\n1. **Local file path** → Read the file directly\n2. **URL** (GitHub, HuggingFace, etc.) → Fetch via WebFetch tool\n3. **Code snippet** → Pasted in the conversation\n4. **Natural language** → Mathematical description (e.g., \"element-wise SiLU(x) * y\")\n5. **Model component** → e.g., \"the MLP in Phi-4\" — locate in transformers source and extract\n\n## Steps\n\n### 1. Understand the Operation\n\n- Read/fetch the source code from whatever input was provided\n- Identify the mathematical operation (forward pass)\n- Derive the backward pass (gradient computation)\n- Identify all inputs, outputs, and their expected shapes/dtypes\n- Note any precision-sensitive operations that need float32 upcasting (sigmoid, rsqrt, exp, log, tanh)\n\n### 2. Write PyTorch Reference\n\nCreate a standalone implementation that:\n- Depends only on `torch` (no external libraries)\n- Implements both forward and backward behavior (either as an `nn.Module` or a plain function)\n- Will serve as the correctness baseline for testing\n- Is clean, readable, and well-named\n\n### 3. Classify Into Tier\n\nRead [kernel-profile-format.md](kernel-profile-format.md) for the full schema.\n\n**Tier 1 — Element-wise**: No reductions across dimensions. One row per program. Examples: SwiGLU, GeGLU, DyT.\n- Read reference: `src/liger_kernel/ops/swiglu.py`\n\n**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.\n- Read reference: `src/liger_kernel/ops/rms_norm.py`\n\n**Tier 3 — Fused/Complex**: Multi-pass algorithms, gradient-in-forward tricks, multiple outputs. Examples: CrossEntropy, FusedLinearCrossEntropy.\n- Read reference: `src/liger_kernel/ops/cross_entropy.py`\n\nAlso read the closest example profile:\n- Tier 1 → [examples/swiglu-profile.md](examples/swiglu-profile.md)\n- Tier 2 → [examples/rms-norm-profile.md](examples/rms-norm-profile.md)\n- Tier 3 → [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md)\n\n### 4. Produce Kernel Profile\n\nFill in all fields from [kernel-profile-format.md](kernel-profile-format.md).\n\n### 5. Present to User\n\nShow:\n1. The PyTorch reference implementation (full code)\n2. The kernel profile (all fields)\n3. Which existing kernel is closest (for the Generator to use as reference)\n\nWait for user confirmation before proceeding to Stage 2.\n\n<!-- File: .agents/skills/liger-kernel-dev/examples/cross-entropy-profile.md (Tokens: ~797 | Category: anthropic-skill) -->\n# Kernel Profile: CrossEntropy (Tier 3 — Fused/Complex)\n\n## Identity\n- operation_name: cross_entropy\n- function_class_name: LigerCrossEntropyFunction\n- module_class_name: LigerCrossEntropyLoss\n- functional_name: liger_cross_entropy\n\n## Classification\n- tier: 3\n- tier_description: fused/complex\n- closest_existing_kernel: fused_linear_cross_entropy (extends this with linear layer fusion)\n\n## Forward Pass\n- forward_inputs:\n  - _input: shape (B*T, V), logits\n  - target: shape (B*T,), label indices\n  - weight: shape (V,) or None, class weights\n  - ignore_index: int, label to ignore\n  - label_smoothing: float\n  - reduction: str, \"mean\" | \"sum\" | \"none\"\n  - softcap: float or None\n- forward_outputs:\n  - loss: scalar or shape (B*T,) if reduction=\"none\"\n  - z_loss: optional auxiliary loss\n  - token_accuracy: optional accuracy metric\n  - predicted_tokens: optional argmax tokens\n- forward_computation: Two-pass online softmax + cross entropy loss with optional smoothing and softcapping\n- precision_sensitive_ops: [exp, log]\n\n## Backward Pass\n- backward_saved_tensors: [_input] — gradient is computed during forward and stored in-place\n- backward_recompute: none (gradient-in-forward trick)\n- gradient_formulas:\n  - d_input: already computed in forward pass and stored in _input tensor. Backward just scales by grad_output.\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row (one row = one token's logits over vocab)\n- block_size_source: custom — iterates over vocab in chunks of BLOCK_SIZE\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - weight: optional class weights\n  - ignore_index: int = -100\n  - lse_square_scale: float = 0.0\n  - label_smoothing: float = 0.0\n  - reduction: str = \"mean\"\n  - softcap: float or None = None\n  - return_z_loss: bool = False\n- learnable_params: none\n\n## Benchmarking\n- benchmark_variable: vocab_size (V)\n- benchmark_x_label: \"V\"\n- benchmark_x_values_suggestion: [4096, 8192, 16384, 32768, 65536, 131072]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {B: 8, T: 512, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Online softmax**: Two-pass algorithm. Pass 1: compute running max and logsumexp. Pass 2: compute softmax and gradients. Avoids materializing the full softmax vector.\n- **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.\n- **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.\n- **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.\n- **Multiple loss components**: Combines original CE loss, label smoothing loss, and z-loss (for training stability). Each component contributes to both loss and gradient.\n\n<!-- File: .agents/skills/liger-kernel-dev/examples/rms-norm-profile.md (Tokens: ~812 | Category: anthropic-skill) -->\n# Kernel Profile: RMSNorm (Tier 2 — Reduction)\n\n## Identity\n- operation_name: rms_norm\n- function_class_name: LigerRMSNormFunction\n- module_class_name: LigerRMSNorm\n- functional_name: liger_rms_norm\n\n## Classification\n- tier: 2\n- tier_description: reduction\n- closest_existing_kernel: layer_norm (similar pattern with different normalization)\n\n## Forward Pass\n- forward_inputs:\n  - X: shape (B, T, H), input tensor\n  - W: shape (H,), weight tensor\n  - eps: float, epsilon for numerical stability\n  - offset: float, weight offset (0.0 for Llama, 1.0 for Gemma)\n  - casting_mode: str, \"llama\" | \"gemma\" | \"none\"\n- forward_outputs:\n  - Y: shape (B, T, H), normalized output\n- forward_computation: Y = (X / RMS(X)) * (W + offset), RMS = sqrt(mean(X^2) + eps)\n- precision_sensitive_ops: [rsqrt]\n\n## Backward Pass\n- backward_saved_tensors: [X, W, RSTD] — RSTD cached from forward to avoid recomputation\n- backward_recompute: none (RSTD is expensive to recompute)\n- gradient_formulas:\n  - dX: rstd * (dY*(W+offset) - (1/N) * rstd^2 * dot(dY*(W+offset), X) * X)\n  - dW: sum over (B,T) of dY * (X * rstd)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: forward uses one program per row `(n_rows,)`, backward uses SM-based partitioning `(sm_count,)` with `rows_per_program`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: true (for dW reduction — each SM accumulates partial dW, then summed)\n\n## Module Parameters\n- module_init_params:\n  - hidden_size: int\n  - eps: float = 1e-6\n  - offset: float = 0.0\n  - casting_mode: str = \"llama\"\n  - init_fn: str = \"ones\"\n  - in_place: bool = True\n  - elementwise_affine: bool = True\n- learnable_params:\n  - weight: shape (hidden_size,), init ones or zeros\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"huggingface\"]\n- benchmark_fixed_config: {M: 4096, eps: 1e-6, dtype: torch.float32}\n\n## Key Patterns\n\n- **RSTD caching**: Forward computes and stores `rstd = rsqrt(mean(X^2) + eps)` — 1 value per row, tiny memory cost, saves 4 ops in backward\n- **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)`\n- **Casting modes as constexpr**: `casting_mode` is `tl.constexpr` so the compiler eliminates dead branches\n- **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)\n- **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)\n\n<!-- File: .agents/skills/liger-kernel-dev/examples/swiglu-profile.md (Tokens: ~600 | Category: anthropic-skill) -->\n# Kernel Profile: SwiGLU (Tier 1 — Element-wise)\n\n## Identity\n- operation_name: swiglu\n- function_class_name: LigerSiLUMulFunction\n- module_class_name: LigerSwiGLUMLP\n- functional_name: liger_swiglu\n\n## Classification\n- tier: 1\n- tier_description: element-wise\n- closest_existing_kernel: geglu (same structure, different activation)\n\n## Forward Pass\n- forward_inputs:\n  - a: shape (B, T, H), gate projection output\n  - b: shape (B, T, H), up projection output\n- forward_outputs:\n  - c: shape (B, T, H), silu(a) * b\n- forward_computation: c = silu(a) * b, where silu(x) = x * sigmoid(x)\n- precision_sensitive_ops: [sigmoid]\n\n## Backward Pass\n- backward_saved_tensors: [a, b] (reshaped to 2D in forward wrapper)\n- backward_recompute: recompute silu(a) and sigmoid(a) in backward\n- gradient_formulas:\n  - da: dc * (silu(a) * (1 - sigmoid(a)) + sigmoid(a)) * b\n  - db: dc * silu(a)\n\n## Tiling Strategy\n- grid_dimensions: 1D\n- grid_description: one program per row, `(n_rows,)`\n- block_size_source: calculate_settings(n_cols)\n- needs_sm_parallelism: false\n\n## Module Parameters\n- module_init_params:\n  - config: HuggingFace model config object\n- learnable_params:\n  - gate_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - up_proj: Linear(hidden_size, intermediate_size, bias=False)\n  - down_proj: Linear(intermediate_size, hidden_size, bias=False)\n\n## Benchmarking\n- benchmark_variable: hidden_size\n- benchmark_x_label: \"hidden_size\"\n- benchmark_x_values_suggestion: [1024, 2048, 4096, 8192, 16384]\n- benchmark_providers: [\"liger\", \"torch\", \"torch_compile\"]\n- benchmark_fixed_config: {BT: 4096, dtype: torch.bfloat16}\n\n## Key Patterns\n\n- **Recomputation over saving**: Forward saves `a, b` but backward recomputes `sigmoid(a)` and `silu(a)` — saves memory, sigmoid is cheap\n- **In-place backward**: Writes gradients directly to `a_ptr` and `b_ptr` (the saved tensors) — saves allocation\n- **Float32 for sigmoid**: `a_row` cast to `tl.float32` before `tl.sigmoid`, result cast back via `.cast(b_row.dtype)`\n- **No intermediate allocations**: Forward kernel writes directly to output `c`; backward kernel overwrites saved `a, b`\n\n<!-- File: .agents/skills/liger-kernel-dev/generator.md (Tokens: ~1022 | Category: anthropic-skill) -->\n# Generator Agent\n\nTakes a confirmed kernel profile and PyTorch reference, generates all files for a production-ready Liger kernel.\n\n## Pre-Requisites\n\nBefore generating, read the existing implementation closest to this kernel's tier:\n- Tier 1 → `src/liger_kernel/ops/swiglu.py` + `src/liger_kernel/transformers/swiglu.py`\n- Tier 2 → `src/liger_kernel/ops/rms_norm.py` + `src/liger_kernel/transformers/rms_norm.py`\n- Tier 3 → `src/liger_kernel/ops/cross_entropy.py` + `src/liger_kernel/transformers/cross_entropy.py`\n\nAlso read the corresponding test and benchmark files for the reference kernel.\n\n## Files to Generate\n\n### 1. `src/liger_kernel/ops/{kernel}.py` (NEW)\n\nThe core Triton kernel. See [templates/ops-kernel.md](templates/ops-kernel.md).\n\nKey rules:\n- `@triton.jit` for forward and backward kernels\n- `torch.autograd.Function` subclass with `@ensure_contiguous` decorator\n- Use `calculate_settings(n_cols)` for BLOCK_SIZE and num_warps\n- Cast to `tl.float32` for precision-sensitive ops (sigmoid, rsqrt, exp, log, tanh)\n- Use `tl.program_id(0).to(tl.int64)` to avoid overflow\n- Use stride parameters for memory access, not hardcoded offsets\n- Prefer recomputation over saving for backward when computation is cheap\n\n### 2. `src/liger_kernel/transformers/{kernel}.py` (NEW)\n\nThe nn.Module wrapper. See [templates/module-wrapper.md](templates/module-wrapper.md).\n\nKey rules:\n- Import from `liger_kernel.ops` (not from the submodule directly)\n- Include `extra_repr` method\n- Match the PyTorch API as closely as possible\n\n### 3. `src/liger_kernel/transformers/functional.py` (MODIFY)\n\nAdd functional API. See [templates/functional-api.md](templates/functional-api.md).\n\nTwo changes:\n- Add import of the Function class at the top (alphabetical order)\n- Add `liger_{kernel}()` function (alphabetical order among existing functions)\n\n### 4. `src/liger_kernel/ops/__init__.py` (MODIFY)\n\nAdd import of the Function class. Insert in alphabetical order among existing imports. Use `# noqa: F401` suffix.\n\n### 5. `src/liger_kernel/transformers/__init__.py` (MODIFY)\n\nAdd the Module class in three locations (alphabetical order in each):\n- Direct import at top (among \"Always-safe imports\")\n- `__all__` list\n- If the Module has a `TYPE_CHECKING` import or `__getattr__` entry needed, add there too\n\n### 6. `test/transformers/test_{kernel}.py` (NEW)\n\nUnit tests. See [templates/unit-test.md](templates/unit-test.md).\n\nKey rules:\n- Include a PyTorch reference implementation (copy from Analyzer output)\n- Test both forward and backward correctness\n- Parametrize over shapes (include non-power-of-2 and small/large sizes)\n- Parametrize over dtypes with appropriate tolerances\n- Use `set_seed()`, `assert_verbose_allclose()`, `infer_device()`, `supports_bfloat16()`\n- Import from `test.utils`, NOT `liger_kernel.test.utils`\n\n### 7. `benchmark/scripts/benchmark_{kernel}.py` (NEW)\n\nBenchmarks. See [templates/benchmark.md](templates/benchmark.md).\n\nKey rules:\n- Define `_setup_{kernel}()` function returning `(input_tensor, layer)`\n- Define `bench_speed_{kernel}()` using `run_speed_benchmark`\n- Define `bench_memory_{kernel}()` using `run_memory_benchmark`\n- Compare \"liger\" vs \"torch\" (and optionally \"torch_compile\") providers\n- Use `get_benchmark_model_config`, `estimate_kernel_peak_memory`, and sweep config utilities\n- Run both speed and memory benchmarks for forward, backward, and full modes\n- **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\n\n## Modify Mode\n\nWhen modifying an existing kernel:\n1. Read the existing kernel files thoroughly\n2. Understand the current implementation patterns\n3. Make targeted changes — do not rewrite entire files\n4. Add new test cases for new functionality\n5. Ensure existing tests still cover original behavior\n\n## Code Style\n\n- Line length 120, double quotes, single imports sorted with isort\n- Follow exact patterns from existing code — do not innovate on style\n- When modifying existing files, insert new entries in **alphabetical order**\n- Import style: one import per line, no multi-imports\n\n<!-- File: .agents/skills/liger-kernel-dev/kernel-profile-format.md (Tokens: ~845 | Category: anthropic-skill) -->\n# Kernel Profile Format\n\nWhen analyzing an operation, produce a profile with all of these fields.\n\n## Fields\n\n```yaml\n# Identity\noperation_name: snake_case name (e.g., \"rms_norm\", \"dyt\", \"sparsemax\")\nfunction_class_name: \"Liger{PascalCase}Function\" (e.g., \"LigerRMSNormFunction\")\nmodule_class_name: \"Liger{PascalCase}\" (e.g., \"LigerRMSNorm\")\nfunctional_name: \"liger_{snake_case}\" (e.g., \"liger_rms_norm\")\n\n# Classification\ntier: 1 | 2 | 3\ntier_description: \"element-wise\" | \"reduction\" | \"fused/complex\"\nclosest_existing_kernel: which existing Liger kernel is most similar\n\n# Forward Pass\nforward_inputs:\n  - name: X\n    shape: \"(B, T, H)\"\n    dtype: \"input dtype\"\n  # ... list all inputs with shapes\nforward_outputs:\n  - name: Y\n    shape: \"(B, T, H)\"\n    dtype: \"same as input\"\nforward_computation: brief math description (e.g., \"Y = X * rsqrt(mean(X^2) + eps) * W\")\nprecision_sensitive_ops: list ops needing float32 (e.g., [rsqrt, sigmoid, exp])\n\n# Backward Pass\nbackward_saved_tensors: what to save_for_backward (e.g., [X, W, RSTD])\nbackward_recompute: what to recompute instead of saving (e.g., \"recompute silu in backward\")\ngradient_formulas:\n  dX: \"formula or description\"\n  dW: \"formula or description\"\n\n# Tiling Strategy\ngrid_dimensions: 1D | 2D\ngrid_description: \"one program per row\" | \"one program per (row_block, col_block)\" | etc.\nblock_size_source: \"calculate_settings(n_cols)\" | \"custom\" with explanation\nneeds_sm_parallelism: true | false (for weight gradient reduction in backward)\n\n# Module Parameters\nmodule_init_params:\n  - name: hidden_size\n    type: int\n  - name: eps\n    type: float\n    default: 1e-6\n  # ... list all __init__ parameters\nlearnable_params:\n  - name: weight\n    shape: \"(hidden_size,)\"\n    init: \"ones\"\n  # ... list learnable parameters\n\n# Benchmarking\nbenchmark_variable: \"hidden_size\" | \"seq_len\" | \"vocab_size\" | etc.\nbenchmark_x_label: human-readable label for x-axis\nbenchmark_x_values_suggestion: suggested range to sweep\nbenchmark_providers: [\"liger\", \"torch\"] (add \"torch_compile\" if applicable)\nbenchmark_fixed_config: dict of fixed params (e.g., {BT: 4096, dtype: torch.bfloat16})\n```\n\n## Naming Conventions\n\n| Component | Pattern | Example |\n|-----------|---------|---------|\n| Ops file | `src/liger_kernel/ops/{operation_name}.py` | `ops/dyt.py` |\n| Transformer file | `src/liger_kernel/transformers/{operation_name}.py` | `transformers/dyt.py` |\n| Test file | `test/transformers/test_{operation_name}.py` | `test_dyt.py` |\n| Benchmark file | `benchmark/scripts/benchmark_{operation_name}.py` | `benchmark_dyt.py` |\n| Function class | `Liger{PascalCase}Function` | `LigerDyTFunction` |\n| Module class | `Liger{PascalCase}` | `LigerDyT` |\n| Functional | `liger_{snake_case}` | `liger_dyt` |\n| Forward wrapper | `{snake_case}_forward` or `liger_{snake_case}_fwd` | `liger_dyt_fwd` |\n| Backward wrapper | `{snake_case}_backward` or `liger_{snake_case}_bwd` | `liger_dyt_bwd` |\n\n<!-- File: .agents/skills/liger-kernel-dev/SKILL.md (Tokens: ~732 | Category: anthropic-skill) -->\n---\nname: liger-kernel-dev\ndescription: \"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.\"\n---\n\n# Liger Kernel Dev\n\nDevelops 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.\n\n## Mode Detection\n\n- **Create mode**: User asks to create/add/generate/write/build a new kernel → full pipeline\n- **Modify mode**: User asks to update/fix/change/extend an existing kernel → skip Analyze, modify files, then Validate\n\n## Pipeline (Create Mode)\n\n### Stage 1: Analyze\n\nFollow the **Analyzer** workflow in [analyzer.md](analyzer.md). If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.\n\nAccepts any input: local file, URL, code snippet, natural language description, or model component reference. Produces a standalone PyTorch reference implementation and a kernel profile.\n\n**Human checkpoint:** Present PyTorch reference + kernel profile. Confirm before proceeding.\n\n### Stage 2: Generate\n\nFollow the **Generator** workflow in [generator.md](generator.md).\n\nGenerates/modifies up to 8 files:\n\n1. `src/liger_kernel/ops/{kernel}.py` — NEW Triton kernels + autograd Function\n2. `src/liger_kernel/transformers/{kernel}.py` — NEW nn.Module wrapper\n3. `src/liger_kernel/transformers/functional.py` — MODIFY add functional API\n4. `src/liger_kernel/ops/__init__.py` — MODIFY export Function class\n5. `src/liger_kernel/transformers/__init__.py` — MODIFY export Module + `__all__`\n6. `test/transformers/test_{kernel}.py` — NEW unit tests\n7. `benchmark/scripts/benchmark_{kernel}.py` — NEW benchmark script\n8. `benchmark/data/all_benchmark_data.csv` — MODIFY (after benchmarks run)\n\n**Human checkpoint:** Present changes for review.\n\n### Stage 3: Validate\n\nFollow the **Validator** workflow in [validator.md](validator.md).\n\nRuns checkstyle, unit tests (hard gate — stops on persistent failure), benchmarks, and generates plots. Optionally runs ncu profiling.\n\n**Human checkpoint:** Report final results with benchmark numbers and plots.\n\n## Pipeline (Modify Mode)\n\n1. Read existing kernel files to understand current implementation\n2. Understand the requested modification\n3. Make targeted changes (Generator handles this)\n4. Run full Validate stage (same as create mode)\n\n## Reference Files\n\n- [kernel-profile-format.md](kernel-profile-format.md) — Kernel profile schema and field descriptions\n- [examples/swiglu-profile.md](examples/swiglu-profile.md) — Tier 1 (element-wise) reference\n- [examples/rms-norm-profile.md](examples/rms-norm-profile.md) — Tier 2 (reduction) reference\n- [examples/cross-entropy-profile.md](examples/cross-entropy-profile.md) — Tier 3 (fused/complex) reference\n- Templates in [templates/](templates/) — Code generation patterns for each file type\n\n<!-- File: .agents/skills/liger-kernel-dev/templates/benchmark.md (Tokens: ~1066 | Category: anthropic-skill) -->\n# Benchmark Template\n\n## File: `benchmark/scripts/benchmark_{kernel}.py`\n\n```python\nimport os\nimport sys\n\nimport torch\n\nfrom benchmark_model_configs import compute_hidden_size_sweep_config\nfrom benchmark_model_configs import estimate_kernel_peak_memory\nfrom benchmark_model_configs import get_benchmark_model_config\nfrom utils import SingleBenchmarkRunInput\nfrom utils import SingleBenchmarkRunOutput\nfrom utils import parse_benchmark_script_args\nfrom utils import run_benchmarks\nfrom utils import run_memory_benchmark\nfrom utils import run_speed_benchmark\n\nfrom liger_kernel.utils import infer_device\n\ndevice = infer_device()\n\nsys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), \"../..\")))\n\n\ndef _setup_{kernel}(input: SingleBenchmarkRunInput):\n    \"\"\"Create input tensor and layer from benchmark config.\"\"\"\n    # Import both Liger and PyTorch reference implementations\n    from test.transformers.test_{kernel} import Liger{Kernel}\n    from test.transformers.test_{kernel} import Torch{Kernel}\n\n    cfg = input.extra_benchmark_config\n    hidden_size = input.x  # Or whatever the sweep variable is\n\n    x = torch.randn(cfg[\"BT\"], hidden_size, device=device, dtype=cfg[\"dtype\"], requires_grad=True)\n\n    if input.kernel_provider == \"liger\":\n        layer = Liger{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch\":\n        layer = Torch{Kernel}(hidden_size=hidden_size).to(device)\n    elif input.kernel_provider == \"torch_compile\":\n        layer = torch.compile(Torch{Kernel}(hidden_size=hidden_size).to(device))\n    else:\n        raise ValueError(f\"Invalid provider: {input.kernel_provider} for {kernel}\")\n\n    return x, layer\n\n\ndef bench_speed_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_speed_benchmark(lambda: layer(x), input.kernel_operation_mode, [x])\n\n\ndef bench_memory_{kernel}(input: SingleBenchmarkRunInput) -> SingleBenchmarkRunOutput:\n    x, layer = _setup_{kernel}(input)\n    return run_memory_benchmark(lambda: layer(x), input.kernel_operation_mode)\n\n\nBT = 4096  # Fixed batch*seq for sweeping hidden_size\n\n\nif __name__ == \"__main__\":\n    args = parse_benchmark_script_args()\n    model = get_benchmark_model_config(args.model)\n\n    def _probe():\n        probe_input = SingleBenchmarkRunInput(\n            x=model.hidden_size,\n            kernel_provider=\"torch\",\n            extra_benchmark_config={\"BT\": BT, \"dtype\": model.dtype},\n        )\n        x, layer = _setup_{kernel}(probe_input)\n        return layer(x)\n\n    peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)\n    sweep_config = compute_hidden_size_sweep_config(model, peak_bytes, bt=BT)\n    x_values = [1024 * i for i in range(1, 17) if 1024 * i <= sweep_config.max_hidden_size] or [model.hidden_size]\n\n    common_configs = {\n        \"kernel_name\": \"{kernel}\",\n        \"x_name\": \"hidden_size\",\n        \"x_label\": \"hidden_size\",\n        \"x_values\": x_values,\n        \"kernel_providers\": [\"liger\", \"torch\", \"torch_compile\"],\n        \"extra_benchmark_configs\": [{\"BT\": sweep_config.bt, \"dtype\": model.dtype}],\n        \"overwrite\": args.overwrite,\n    }\n\n    run_benchmarks(\n        bench_test_fn=bench_speed_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"speed\",\n        metric_unit=\"ms\",\n        **common_configs,\n    )\n    run_benchmarks(\n        bench_test_fn=bench_memory_{kernel},\n        kernel_operation_modes=[\"full\", \"forward\", \"backward\"],\n        metric_name=\"memory\",\n        metric_unit=\"MB\",\n        **common_configs,\n    )\n```\n\n### Key Rules\n\n1. **Import both implementations from test file** — avoids duplicating reference code\n2. **Support \"liger\", \"torch\", and \"torch_compile\" providers**\n3. **Use `run_speed_benchmark` and `run_memory_benchmark`** from utils — don't write custom measurement\n4. **Use `get_benchmark_model_config`** and sweep utilities for standard model sizes\n5. **Use `estimate_kernel_peak_memory`** to auto-compute safe x_values range\n6. **Run both speed AND memory** for all three modes (forward, backward, full)\n7. **Use `BT = 4096`** as default batch*seq when sweeping hidden_size\n8. **Adapt sweep variable** — some kernels sweep `seq_len` or `vocab_size` instead of `hidden_size`. Use the appropriate `compute_*_sweep_config` utility.\n9. **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.\n\n<!-- File: .agents/skills/liger-kernel-dev/templates/functional-api.md (Tokens: ~211 | Category: anthropic-skill) -->\n# Functional API Template\n\n## File: `src/liger_kernel/transformers/functional.py` (MODIFY)\n\n### Step 1: Add Import\n\nAdd at the top of the file, in alphabetical order among existing imports:\n\n```python\nfrom liger_kernel.ops import Liger{Kernel}Function\n```\n\n### Step 2: Add Function\n\nAdd in alphabetical order among existing functions:\n\n```python\ndef liger_{kernel}(x, ...params):\n    return Liger{Kernel}Function.apply(x, ...params)\n```\n\nFor functions with complex return types (like CrossEntropy), wrap the output:\n\n```python\ndef liger_{kernel}(input, target, ...params):\n    result = Liger{Kernel}Function.apply(input, target, ...params)\n    return result\n```\n\n### Key Rules\n\n- Function signature should match the PyTorch equivalent where applicable\n- Use positional args for tensor inputs, keyword args for config params\n- Keep the function body minimal — just call `.apply()` and return\n\n<!-- File: .agents/skills/liger-kernel-dev/templates/module-wrapper.md (Tokens: ~249 | Category: anthropic-skill) -->\n# Module Wrapper Template\n\n## File: `src/liger_kernel/transformers/{kernel}.py`\n\n```python\nimport torch\nimport torch.nn as nn\n\nfrom liger_kernel.ops import Liger{Kernel}Function\n\n\nclass Liger{Kernel}(nn.Module):\n    def __init__(self, ...params):\n        super().__init__()\n        # Store config\n        self.param = param\n        # Create learnable parameters\n        self.weight = nn.Parameter(torch.ones(hidden_size))\n\n    def forward(self, x):\n        return Liger{Kernel}Function.apply(x, self.weight, ...self.params)\n\n    def extra_repr(self):\n        return f\"param={self.param}, ...\"\n```\n\n### Key Rules\n\n- Import from `liger_kernel.ops` (vendor-replaceable), not from the submodule\n- Match the PyTorch reference API as closely as possible (same __init__ params, same forward signature)\n- Include `extra_repr` for debugging\n- For loss modules, name the class `Liger{Kernel}Loss` and the forward takes `(input, target)`\n- For model-specific variants, create subclasses (see `LigerRMSNormForGemma` in `rms_norm.py`)\n\n<!-- File: .agents/skills/liger-kernel-dev/templates/ops-kernel.md (Tokens: ~1359 | Category: anthropic-skill) -->\n# Ops Kernel Template\n\n## File: `src/liger_kernel/ops/{kernel}.py`\n\nBefore writing, read the reference kernel for this tier in the actual codebase.\n\n### Import Block\n\n```python\nimport torch\nimport triton\nimport triton.language as tl\n\nfrom liger_kernel.ops.utils import calculate_settings\nfrom liger_kernel.ops.utils import ensure_contiguous\n```\n\nAdd additional imports only as needed:\n- `math` — if using `math.ceil` for row partitioning\n- `from liger_kernel.ops.utils import compare_version` — if checking triton version\n- `from liger_kernel.ops.utils import torch_to_triton_dtype` — if passing dtype as kernel arg\n- `rsqrt` from triton libdevice — if using inverse square root (see `rms_norm.py` for the import pattern with version guards)\n- `tanh` from triton libdevice — if using tanh (see `dyt.py` for the import pattern)\n\n### Forward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_forward_kernel(\n    # Output pointers first, then input pointers\n    Y_ptr, Y_stride,\n    X_ptr, X_stride,\n    # ... other input pointers and strides\n    n_cols,\n    # constexpr params for compile-time optimization\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Calculate base pointers using strides\n    X_ptr += row_idx * X_stride\n    Y_ptr += row_idx * Y_stride\n\n    # Load with mask and float32 upcast for precision\n    X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0).to(tl.float32)\n\n    # === FORWARD COMPUTATION ===\n\n    tl.store(Y_ptr + col_offsets, result, mask=mask)\n```\n\n### Backward Kernel\n\n```python\n@triton.jit\ndef _{kernel}_backward_kernel(\n    # Gradient output, then gradient inputs, then saved tensors\n    dY_ptr, dY_stride,\n    dX_ptr, dX_stride,\n    # ... saved tensors from forward\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n):\n    row_idx = tl.program_id(0).to(tl.int64)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    mask = col_offsets < n_cols\n\n    # Load gradient and saved tensors\n    # === BACKWARD COMPUTATION ===\n\n    tl.store(dX_ptr + row_idx * dX_stride + col_offsets, dX_row, mask=mask)\n```\n\n### Python Wrappers\n\n```python\ndef {kernel}_forward(X, ...):\n    ori_shape = X.shape\n    n_cols = ori_shape[-1]\n    X = X.view(-1, n_cols)\n    n_rows = X.shape[0]\n\n    Y = torch.empty_like(X)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_forward_kernel[(n_rows,)](\n        Y, Y.stride(-2),\n        X, X.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return Y.view(*ori_shape)\n\n\ndef {kernel}_backward(dY, ...saved):\n    ori_shape = dY.shape\n    n_cols = ori_shape[-1]\n    dY = dY.view(-1, n_cols)\n    n_rows = dY.shape[0]\n\n    dX = torch.empty_like(dY)\n    BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n\n    _{kernel}_backward_kernel[(n_rows,)](\n        dY, dY.stride(-2),\n        dX, dX.stride(-2),\n        n_cols=n_cols,\n        BLOCK_SIZE=BLOCK_SIZE,\n        num_warps=num_warps,\n    )\n    return dX.view(*ori_shape)\n```\n\n### Autograd Function\n\n```python\nclass Liger{Kernel}Function(torch.autograd.Function):\n    @staticmethod\n    @ensure_contiguous\n    def forward(ctx, X, ...params):\n        Y = {kernel}_forward(X, ...params)\n        ctx.save_for_backward(...)  # Save only what backward needs\n        return Y\n\n    @staticmethod\n    @ensure_contiguous\n    def backward(ctx, dY):\n        ...saved = ctx.saved_tensors\n        dX = {kernel}_backward(dY, ...saved)\n        # Return one gradient per forward input, None for non-tensor params\n        return dX, None, None, ...\n```\n\n### Tier-Specific Patterns\n\n**Tier 1 (element-wise):**\n- Grid: `(n_rows,)` — one program per row\n- Simple 1D indexing with `col_offsets` and `mask`\n- Prefer recomputation in backward over saving activations\n- See `swiglu.py`: saves `a, b` but recomputes `silu(a)` in backward\n\n**Tier 2 (reduction):**\n- Forward may cache intermediate state (e.g., RSTD in RMSNorm)\n- Backward may use SM-based parallelism for weight gradient reduction:\n  ```python\n  sm_count = torch.cuda.get_device_properties(X.device).multi_processor_count\n  _dW = torch.empty((sm_count, n_cols), dtype=torch.float32, device=W.device)\n  grid = (sm_count,)\n  rows_per_program = math.ceil(n_rows / sm_count)\n  ```\n- See `rms_norm.py` for the full pattern\n\n**Tier 3 (fused/complex):**\n- Multi-pass algorithms (e.g., online softmax: pass 1 finds max+sum, pass 2 computes gradients)\n- Gradient-in-forward trick: compute gradients during forward, store in input tensor\n- Multiple `tl.constexpr` flags for compile-time dead code elimination\n- See `cross_entropy.py` for the full pattern\n\n### Key Rules\n\n1. Always use `tl.program_id(0).to(tl.int64)` (prevents overflow on large tensors)\n2. Always use `mask = col_offsets < n_cols` for boundary handling\n3. Use `tl.constexpr` for parameters that enable compile-time optimization\n4. Use stride parameters from tensors, never hardcode memory layout\n5. Cast precision-sensitive ops to `tl.float32`, cast back for storage\n6. Return `None` for each non-tensor parameter in backward\n\n<!-- File: .agents/skills/liger-kernel-dev/templates/unit-test.md (Tokens: ~1219 | Category: anthropic-skill) -->\n# Unit Test Template\n\n## File: `test/transformers/test_{kernel}.py`\n\n```python\nimport pytest\nimport torch\n\nfrom test.utils import assert_verbose_allclose\nfrom test.utils import infer_device\nfrom test.utils import set_seed\nfrom test.utils import supports_bfloat16\n\nfrom liger_kernel.ops import Liger{Kernel}Function\nfrom liger_kernel.transformers.{kernel} import Liger{Kernel}\nfrom liger_kernel.transformers.functional import liger_{kernel}\n\n\n# ---- PyTorch Reference ----\n# Paste the standalone PyTorch reference from the Analyzer stage here.\n\nclass Torch{Kernel}(torch.nn.Module):\n    def __init__(self, ...):\n        super().__init__()\n        ...\n\n    def forward(self, x):\n        ...\n\n\nset_seed(42)\ndevice = infer_device()\n\n\n# ---- Correctness Test (Module) ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (4, 16, 2048),\n        (1, 1, 1023),   # Non-power-of-2\n        (3, 7, 256),    # Small + prime numbers\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-1,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_correctness(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Initialize both with same weights\n    torch_layer = Torch{Kernel}(...).to(device).to(dtype)\n    liger_layer = Liger{Kernel}(...).to(device).to(dtype)\n    # Copy weights from torch to liger\n    # liger_layer.weight.data = torch_layer.weight.data.clone()\n\n    # Forward\n    torch_output = torch_layer(x1)\n    liger_output = liger_layer(x2)\n    assert_verbose_allclose(torch_output, liger_output, rtol=rtol, atol=atol, extra_info=\"[output]\")\n\n    # Backward\n    grad_output = torch.randn_like(_input)\n    torch_output.backward(grad_output)\n    liger_output.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol, extra_info=\"[input.grad]\")\n\n    # Check weight gradients if applicable\n    # assert_verbose_allclose(\n    #     torch_layer.weight.grad, liger_layer.weight.grad,\n    #     rtol=rtol, atol=atol, extra_info=\"[weight.grad]\"\n    # )\n\n\n# ---- Functional API Test ----\n\n@pytest.mark.parametrize(\n    \"B, T, hidden_size\",\n    [\n        (2, 8, 4096),\n        (1, 1, 1023),\n    ],\n)\n@pytest.mark.parametrize(\n    \"dtype, atol, rtol\",\n    [\n        (torch.float32, 1e-5, 1e-5),\n        pytest.param(\n            torch.bfloat16,\n            1e-8,\n            5e-2,\n            marks=pytest.mark.skipif(\n                not supports_bfloat16(),\n                reason=\"bfloat16 not supported on this GPU\",\n            ),\n        ),\n    ],\n)\ndef test_liger_{kernel}_functional(B, T, hidden_size, dtype, atol, rtol):\n    _input = torch.randn(B, T, hidden_size, device=device, dtype=dtype)\n\n    x1 = _input.clone().requires_grad_(True)\n    x2 = _input.clone().requires_grad_(True)\n\n    # Test that functional and Function.apply produce identical results\n    output1 = liger_{kernel}(x1, ...)\n    output2 = Liger{Kernel}Function.apply(x2, ...)\n\n    assert_verbose_allclose(output1, output2, rtol=rtol, atol=atol)\n\n    grad_output = torch.randn_like(_input)\n    output1.backward(grad_output)\n    output2.backward(grad_output)\n    assert_verbose_allclose(x1.grad, x2.grad, rtol=rtol, atol=atol)\n```\n\n### Key Testing Rules\n\n1. **Always test both forward AND backward** — backward bugs are common\n2. **Include non-power-of-2 shapes** — tests the mask boundary handling\n3. **Parametrize dtypes with per-dtype tolerances**:\n   - `float32`: `atol=1e-5, rtol=1e-5` (tight)\n   - `bfloat16`: `atol=1e-1, rtol=5e-2` (loose — bf16 has limited precision)\n4. **Use `set_seed(42)` at module level** for reproducibility\n5. **Use `infer_device()`** not `\"cuda\"` — supports multi-vendor\n6. **Guard bfloat16 with `supports_bfloat16()`** — skips on older GPUs\n7. **Use `assert_verbose_allclose`** not `torch.testing.assert_close` — gives better error output\n8. **Test weight gradients** when the kernel has learnable parameters\n9. **Add `@pytest.mark.flaky(reruns=3)`** if the test is sensitive to random seeds\n10. **Add extra test functions** for edge cases specific to the kernel\n\n<!-- File: .agents/skills/liger-kernel-dev/validator.md (Tokens: ~777 | Category: anthropic-skill) -->\n# Validator Agent\n\nValidates the generated kernel through checkstyle, testing, benchmarking, and optional profiling.\n\n## Validation Flow\n\n### Step 1: Checkstyle (Fast Gate)\n\n```bash\nmake checkstyle\n```\n\nIf it fails:\n```bash\nruff check . --fix && ruff format .\nmake checkstyle\n```\n\nIf it still fails after auto-fix, identify and fix the remaining issues manually, then re-run.\n\n### Step 2: Unit Tests (Hard Gate)\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\nIf tests fail:\n1. Read the error output carefully\n2. Identify root cause: kernel logic bug, shape mismatch, dtype issue, or tolerance too tight\n3. Fix the issue in the relevant file(s)\n4. Re-run `make checkstyle` (in case fix introduced style issues)\n5. Re-run the failing test\n\n**Retry up to 3 times total.** If tests still fail after 3 attempts:\n- **STOP IMMEDIATELY. Do NOT proceed to benchmarks.**\n- Report to the user:\n  - Which tests are failing and the exact error messages\n  - What fixes were attempted\n  - The current state of the code\n  - Suggest what might be wrong and ask for guidance\n\n### Step 3: Benchmarks\n\nRun the benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nThis produces speed and memory measurements for forward, backward, and full modes, comparing Liger vs PyTorch baseline.\n\n### Step 4: Generate Plots\n\n```bash\npython benchmark/benchmarks_visualizer.py\n```\n\nPlots are saved to `benchmark/visualizations/`. If the visualizer fails or is not available, this step is non-blocking — report the raw numbers instead.\n\n### Step 5: Optional ncu Profiling\n\nOnly if the user explicitly requests it. Requires NVIDIA Nsight Compute (`ncu`) to be installed.\n\nProfile both the Liger kernel and PyTorch reference:\n```bash\nncu --set full --target-processes all -o liger_profile python -c \"\nimport torch\nfrom liger_kernel.ops.{kernel} import Liger{Kernel}Function\nx = torch.randn(..., device='cuda', dtype=torch.float32, requires_grad=True)\nout = Liger{Kernel}Function.apply(x, ...)\nout.backward(torch.randn_like(out))\n\"\n```\n\nReport key metrics: SM occupancy, memory throughput, compute throughput.\n\nIf ncu is not available, report this and skip.\n\n### Step 6: Report Results\n\nPresent a clear summary to the user:\n\n```\n## Results\n\n**Checkstyle:** PASS\n**Unit Tests:** PASS (X tests passed)\n\n**Speed Benchmarks:**\n| Mode     | Liger (ms) | PyTorch (ms) | Speedup |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...x    |\n| Backward | ...       | ...         | ...x    |\n| Full     | ...       | ...         | ...x    |\n\n**Memory Benchmarks:**\n| Mode     | Liger (MB) | PyTorch (MB) | Savings |\n|----------|-----------|-------------|---------|\n| Forward  | ...       | ...         | ...%    |\n| Backward | ...       | ...         | ...%    |\n| Full     | ...       | ...         | ...%    |\n\n**Assessment:** [Which dimension improved — speed, memory, or both.\nFlag if either metric is catastrophically worse.]\n\n**Plots:** benchmark/visualizations/\n\n**ncu Profiling:** [Results if requested, or \"Not requested\"]\n```\n\n<!-- File: .agents/skills/liger-kernel-perf/finalizer.md (Tokens: ~4036 | Category: anthropic-skill) -->\n# Finalizer Agent\n\nApplies the winning optimization variant to the production kernel, validates correctness through the full test suite and checkstyle, and produces the final optimization report.\n\nThis is Stage 3 of the 3-stage pipeline. It receives the winning variant from Stage 2 (Optimizer) and is responsible for ensuring the optimization is safe to ship.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `winning_variant`: variant number chosen by user (interactive) or by the Optimizer (autonomous)\n- `optimization_workspace`: path to `optimization/{kernel}/` containing all variant files, benchmarks, and the profile\n\n## Expected Workspace State\n\nBefore the Finalizer starts, the following files MUST exist (created by earlier stages):\n\n```\noptimization/{kernel}/\n  original_{kernel}.py          # Snapshot of the original kernel (from Profiler)\n  profile.md                    # Optimization profile (from Profiler)\n  {kernel}_v1.py                # Variant source files (from Optimizer)\n  {kernel}_v1_notes.md          # Variant lab notebooks (from Optimizer)\n  {kernel}_v2.py                # ... more variants\n  {kernel}_v2_notes.md\n  benchmarks/\n    v0_baseline.csv             # Baseline benchmark data\n    v1_results.csv              # Per-variant benchmark data\n    v2_results.csv\n```\n\nIf any required files are missing, report the missing files and stop.\n\n## Steps\n\n### Step 1: Validate Inputs and Read Context\n\n1. Confirm the winning variant file exists: `optimization/{kernel}/{kernel}_v{N}.py`\n2. Read the winning variant code\n3. Read the original kernel: `optimization/{kernel}/original_{kernel}.py`\n4. Read the optimization profile: `optimization/{kernel}/profile.md`\n5. Read ALL variant notes files (`{kernel}_v*_notes.md`) to collect data for the report\n6. Read ALL benchmark CSVs (`benchmarks/v*_results.csv`) to build the comparison table\n7. Read the current production kernel: `src/liger_kernel/ops/{kernel}.py`\n\nVerify the original kernel snapshot matches the current production kernel. If they differ (someone else modified the file during the optimization), warn the user and confirm before proceeding.\n\n### Step 2: Apply the Winning Variant\n\nReplace the production kernel with the winning variant:\n\n```bash\ncp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis is a destructive operation. The original is preserved at `optimization/{kernel}/original_{kernel}.py`.\n\n### Step 3: Run the Full Test Suite (HARD GATE)\n\nRun the complete test file for the kernel:\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -xvs\n```\n\n**This is a HARD GATE.** The test suite must pass completely.\n\n#### If tests PASS:\nProceed to Step 4.\n\n#### If tests FAIL:\nExecute the test failure recovery procedure:\n\n**Attempt 1: Diagnose and fix**\n1. Read the full test error output carefully\n2. Identify the root cause (numerical tolerance issue, shape mismatch, missing import, API change, etc.)\n3. Apply a targeted fix to `src/liger_kernel/ops/{kernel}.py`\n4. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**Attempt 2: Second fix attempt (if Attempt 1 still fails)**\n1. Read the new error output\n2. Try a different fix approach\n3. Re-run the full test suite: `python -m pytest test/transformers/test_{kernel}.py -xvs`\n\n**If both attempts fail: REVERT**\n1. Restore the original kernel:\n   ```bash\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n2. Verify the revert by running the test suite again:\n   ```bash\n   python -m pytest test/transformers/test_{kernel}.py -xvs\n   ```\n3. Set `optimization_applied = false`\n4. Record the failure reason and all error outputs\n5. Proceed to Step 4 to generate the report (the report documents what was attempted even if it could not be applied)\n\n**Do NOT proceed to finalize (Step 5) if the optimization was reverted. The report must clearly state the optimization was not applied.**\n\n### Step 4: Run Checkstyle\n\nRun the project's checkstyle validation:\n\n```bash\nmake checkstyle\n```\n\n#### If checkstyle PASSES:\nProceed to Step 5.\n\n#### If checkstyle FAILS:\nAuto-fix with ruff:\n\n```bash\nruff check . --fix && ruff format .\n```\n\nRetry checkstyle once:\n\n```bash\nmake checkstyle\n```\n\nIf it still fails after auto-fix, report the remaining checkstyle issues but do NOT revert the kernel (checkstyle is not a correctness gate). Note the issues in the report so they can be fixed manually.\n\n### Step 5: Generate Comparison Plots\n\nGenerate 3-way comparison plots showing the original Liger kernel, the optimized kernel, and the HuggingFace/PyTorch baseline. This gives the user a visual demonstration of the improvement.\n\nThe procedure has 3 phases: (1) benchmark both kernels for plots, (2) generate plots, (3) clean the CSV so only the optimized kernel's data remains permanently.\n\n#### Step 5a: Benchmark Both Kernels for Plotting\n\n1. **Benchmark the original kernel as \"liger_original\"**: Temporarily restore the original kernel and run benchmarks with a custom provider name:\n   ```bash\n   # Swap in original kernel\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Edit benchmark script: replace \"liger\" with \"liger_original\" in kernel_providers\n   sed -i 's/\"liger\"/\"liger_original\"/g' benchmark/scripts/benchmark_{kernel}.py\n   \n   # Run benchmarks (appends \"liger_original\" rows to CSV)\n   cd benchmark/scripts && python benchmark_{kernel}.py\n   \n   # Revert the benchmark script change\n   git checkout benchmark/scripts/benchmark_{kernel}.py\n   ```\n\n2. **Benchmark the optimized kernel as \"liger\"**: Re-apply the optimized kernel and run benchmarks normally. Use `--overwrite` so the \"liger\" rows reflect the optimized kernel:\n   ```bash\n   # Swap in optimized kernel\n   cp optimization/{kernel}/{kernel}_v{N}.py src/liger_kernel/ops/{kernel}.py\n   \n   # Run benchmarks (writes \"liger\" rows with optimized kernel)\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   ```\n\nAt this point the CSV has 3 providers: `\"liger_original\"`, `\"liger\"` (optimized), and `\"huggingface\"`/`\"torch\"`. This is what we need for the 3-way comparison plots.\n\n#### Step 5b: Generate Plots\n\nGenerate plots while all 3 providers are in the CSV:\n\n```bash\n# Speed plots for all modes (forward, backward, full)\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name speed \\\n  --overwrite\n\n# Memory plot\npython benchmark/benchmarks_visualizer.py \\\n  --kernel-name {kernel} \\\n  --metric-name memory \\\n  --overwrite\n```\n\nCopy plots to the optimization workspace:\n\n```bash\ncp benchmark/visualizations/{kernel}_speed_*.png optimization/{kernel}/\ncp benchmark/visualizations/{kernel}_memory_*.png optimization/{kernel}/\n```\n\n#### Step 5c: Clean CSV for Permanent State\n\nNow remove the temporary `\"liger_original\"` rows and ensure the CSV only contains clean data for the optimized kernel. Delete ALL rows for this kernel, then re-run the benchmark to append fresh data:\n\n```bash\n# Delete all existing rows for this kernel from CSV (keep header + other kernels)\nhead -1 benchmark/data/all_benchmark_data.csv > /tmp/benchmark_header.csv\ngrep -v \"^{kernel},\" benchmark/data/all_benchmark_data.csv > /tmp/benchmark_other.csv || true\ncat /tmp/benchmark_header.csv <(tail -n +2 /tmp/benchmark_other.csv) > benchmark/data/all_benchmark_data.csv\n\n# Re-run benchmarks with optimized kernel (appends fresh rows)\ncd benchmark/scripts && python benchmark_{kernel}.py\n```\n\nAfter this, `all_benchmark_data.csv` contains:\n- `\"liger\"` — the **optimized** kernel (clean fresh data)\n- `\"huggingface\"` / `\"torch\"` — the baseline (also fresh)\n- NO `\"liger_original\"` — that was only needed for the plots\n\n#### Step 5c: Generate Additional Custom Plots (Optional)\n\nIf the optimization shows interesting size-dependent behavior (e.g., bigger gains at larger sizes), generate targeted plots highlighting this. Use matplotlib directly if the standard visualizer does not cover the comparison you want to show.\n\nSave any additional plots to `optimization/{kernel}/`.\n\n### Step 6: Generate the Final Optimization Report\n\nWrite the report to `optimization/{kernel}/report.md`. The report is the permanent record of the entire optimization effort.\n\nThe report must include these sections:\n\n1. **Summary**: kernel name, GPU, goal, variants tried, best variant, improvement %, and whether the optimization was applied (`Yes` / `No — reverted due to test failures`)\n2. **Baseline Profile**: bottleneck classification with evidence, NCU highlights if available\n3. **Variant Comparison**: include the table from `optimization/{kernel}/comparison.md` (produced by the Optimizer)\n4. **Winning Changes**: what was changed, why it worked, and a key code diff. Run `diff optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py` — do NOT write the diff from memory.\n5. **Rejected Approaches**: for each rejected variant, what was tried and why it failed\n6. **Comparison Plots**: reference the plots generated in Step 5 with file paths\n7. **Correctness Verification**: smoke test results, full test suite results, checkstyle results\n8. **Recommendations**: unexplored optimization opportunities, suggested follow-ups\n9. **Reproduction** section:\n   ```bash\n   # Reproduce baseline benchmarks:\n   cd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n   # Revert to original:\n   cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n   ```\n\nInclude the variant comparison data from `optimization/{kernel}/comparison.md` (produced by the Optimizer).\n\n### Report Quality Checklist\n\nBefore saving the report, verify:\n\n- [ ] All variants are listed in the comparison table (including rejected ones)\n- [ ] Speed ratios are calculated consistently (higher = faster)\n- [ ] Memory ratios are calculated consistently (higher = less memory used)\n- [ ] The winning variant is clearly marked\n- [ ] The bottleneck classification matches the profile from Stage 1\n- [ ] If optimization was reverted, this is stated in the Summary, the header, and Section 6\n- [ ] Reproduction commands use the correct file paths\n- [ ] The code diff in Section 4 is accurate (diff the actual files, do not write from memory)\n- [ ] Comparison plots are generated and referenced in the report\n\n### Step 7: Create Pull Request\n\nCreate a PR with the kernel code changes and the updated benchmark data. Do NOT include plots or optimization workspace files in the PR.\n\n#### Step 7a: Stage Only Relevant Files\n\n```bash\ngit add src/liger_kernel/ops/{kernel}.py\ngit add benchmark/data/all_benchmark_data.csv\n```\n\nDo NOT add:\n- `optimization/` directory (workspace files, notes, plots)\n- `benchmark/visualizations/` (generated plots -- these are for local review only)\n\n#### Step 7b: Create a Descriptive Commit\n\n```bash\ngit commit -m \"$(cat <<'EOF'\n[Perf] Optimize {kernel} kernel: {headline improvement}\n\n{2-3 sentence summary of the optimization technique and its impact.}\n\nKey changes:\n- {Change 1}: {what and why}\n- {Change 2}: {what and why}\n\nBenchmark results ({gpu_name}):\n- Speed (forward): {delta}% {faster/slower}\n- Speed (backward): {delta}% {faster/slower}\n- Speed (full): {delta}% {faster/slower}\n- Memory: {delta}% {reduction/increase}\nEOF\n)\"\n```\n\nIf the host AI runtime has its own commit-attribution convention (e.g., a `Co-Authored-By` trailer), append it per that runtime's guidelines.\n\n#### Step 7c: Push and Create PR\n\n```bash\ngit push -u origin HEAD\n```\n\nCreate the PR using `gh pr create`. The PR body must clearly explain:\n\n1. **What was optimized** -- which kernel, which pass, what bottleneck\n2. **What changes were made** -- specific code changes with rationale for each\n3. **How they help** -- tie back to the bottleneck diagnosis (e.g., \"reduced register pressure from 115 to 95 registers/thread, increasing occupancy from 12.5% to 25%\")\n4. **Benchmark results** -- before/after table with speed and memory numbers at representative sizes\n5. **Correctness** -- confirm all tests pass\n\nUse this PR template:\n\n```bash\ngh pr create --title \"[Perf] Optimize {kernel}: {headline}\" --body \"$(cat <<'EOF'\n## Summary\n\n{2-3 sentences: what was optimized, the key technique, and the headline result.}\n\n## Bottleneck Diagnosis\n\n- **Kernel**: `{kernel}` ({tier} — {tier_description})\n- **GPU**: {gpu_name} ({architecture})\n- **Classification**: {memory-bound / compute-bound / balanced}\n- **Root cause**: {specific bottleneck explanation}\n\n## Changes\n\n{For each change, explain WHAT was changed and WHY it helps:}\n\n### 1. {Change title}\n{Description of the code change and its performance rationale.}\n\n### 2. {Change title}\n{Description of the code change and its performance rationale.}\n\n## Benchmark Results\n\nTested on {gpu_name}. Values are median ms (speed) or MB (memory).\n\n### Speed (ms, lower is better)\n\n| x_value | Mode | Before | After | Improvement |\n|---------|------|--------|-------|-------------|\n| {x1} | forward | {before} | {after} | {delta}% |\n| {x1} | backward | {before} | {after} | {delta}% |\n| {x1} | full | {before} | {after} | {delta}% |\n| ... | ... | ... | ... | ... |\n\n### Memory (MB, lower is better)\n\n| x_value | Before | After | Improvement |\n|---------|--------|-------|-------------|\n| {x1} | {before} | {after} | {delta}% |\n| ... | ... | ... | ... |\n\n## Correctness\n\n- Full test suite: **PASSED** (`python -m pytest test/transformers/test_{kernel}.py -xvs`)\n- Checkstyle: **PASSED** (`make checkstyle`)\n\n## Test Plan\n\n- [x] All existing unit tests pass\n- [x] Benchmarks show improvement across all input sizes\n- [x] No regression on non-target metrics (speed/memory balance maintained)\nEOF\n)\"\n```\n\nIf the host AI runtime appends a \"Generated with\" footer (e.g., Claude Code, Cursor, Copilot), include it per that runtime's guidelines.\n\n**Important**: Do NOT include plots as image attachments in the PR. Plots are for local review only and live in the optimization workspace.\n\n### Step 8: Present the Before/After Summary\n\nInclude the plot file paths in the summary so the user can view them.\n\nDisplay a concise summary to the user:\n\n```\n=== Optimization Complete: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: {APPLIED / REVERTED}\nWinner: v{N} — {strategy}\n\nBefore → After (at largest x_value):\n  Speed (forward):  {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (backward): {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Speed (full):     {baseline_ms} ms → {optimized_ms} ms ({improvement}% faster)\n  Memory:           {baseline_mb} MB → {optimized_mb} MB ({improvement}% reduction)\n\nVariants tried: {N}\nTest suite: PASSED ({N} tests)\nCheckstyle: PASSED\n\nReport: optimization/{kernel}/report.md\n```\n\nIf the optimization was reverted:\n\n```\n=== Optimization REVERTED: {kernel} ===\n\nGPU: {gpu_name} ({architecture})\nStatus: REVERTED — test suite failed after applying winner\n\nWinner candidate: v{N} — {strategy}\nFailure reason: {brief error description}\nFix attempts: 2 (both failed)\n\nThe original kernel has been restored. No changes were made to production code.\n\nVariants tried: {N} (see report for full comparison)\nReport: optimization/{kernel}/report.md\n```\n\n## Error Handling\n\n| Error | Action |\n|-------|--------|\n| Winning variant file not found | Report missing file, stop |\n| Original kernel snapshot missing | Report missing file, stop (cannot safely revert) |\n| Production kernel differs from snapshot | Warn user, ask for confirmation (interactive) or stop (autonomous) |\n| Test suite fails (after 2 fix attempts) | Revert to original, generate report documenting failure |\n| Checkstyle fails (after auto-fix) | Note in report, do NOT revert (not a correctness issue) |\n| Benchmark CSVs missing | Generate report with available data, note missing data |\n| Cannot write report | Print report contents to stdout as fallback |\n\n## Key Principles\n\n1. **Correctness is non-negotiable.** The full test suite is a hard gate. Never skip it, never ignore failures, never mark a failing optimization as successful.\n2. **Always preserve the original.** The snapshot at `optimization/{kernel}/original_{kernel}.py` is the safety net. Never delete or overwrite it.\n3. **The report is permanent.** Even if the optimization is reverted, the report documents everything that was tried. This prevents future engineers from repeating failed approaches.\n4. **Be honest in the report.** Report actual measured numbers. Do not round favorably. If a variant was only 0.3% faster, say 0.3%, not \"marginal improvement.\" If the optimization was reverted, say so prominently.\n5. **Diff, don't recall.** When writing the code diff in the report, actually diff the files. Do not write the diff from memory -- it will contain errors.\n\n<!-- File: .agents/skills/liger-kernel-perf/optimization-strategies.md (Tokens: ~8828 | Category: anthropic-skill) -->\n# Optimization Strategies Catalog\n\nReference catalog for the Optimizer Agent. Techniques are ordered by recommended trial order within each category. The agent should **always start with parameter tuning**, then apply diagnosis-driven techniques based on the bottleneck classification from the Profiler stage.\n\n---\n\n## Decision Framework\n\nAfter reading the optimization profile, follow this priority:\n\n1. **Always first**: Parameter Tuning (Section 1)\n2. **Memory-bound kernel**: Section 2 techniques, then Section 5\n3. **Compute-bound kernel**: Section 3 techniques, then Section 5\n4. **Balanced kernel**: Interleave Section 2 and Section 3, then Section 5\n5. **Architecture-specific**: Layer Section 4 techniques on top of the best variant when targeting a specific GPU\n\n---\n\n## 1. Always First: Parameter Tuning\n\n**Why first**: Parameter tuning is the cheapest technique (minimal code changes) and frequently delivers the largest single improvement. The current codebase uses `calculate_settings()` which picks BLOCK_SIZE as the next power of 2 and num_warps via a fixed heuristic. This is a reasonable default but leaves significant performance on the table for specific input shapes and GPU architectures.\n\n### Why NOT @triton.autotune\n\nLiger production kernels **cannot use `@triton.autotune`** for three architectural reasons:\n\n1. **Forward-backward context coupling**: Liger kernels use `torch.autograd.Function` where `forward()` computes `BLOCK_SIZE`/`num_warps` via `calculate_settings()`, stores them in `ctx`, and `backward()` retrieves them to launch its kernel with identical config. With `@triton.autotune`, these become compile-time constants managed by Triton -- they cannot be stored in `ctx`, and the backward kernel has no way to know which config the forward kernel auto-selected for a given input.\n\n2. **Multi-backend incompatibility**: NPU/Ascend does not support `num_warps` or `num_stages` at all (commit `8b965b9` explicitly removed them). `@triton.autotune` relies on these parameters.\n\n3. **Inference warmup latency**: Each new input shape triggers autotuning (benchmarks all configs). Adds 100ms+ first-call penalty per unique shape, unacceptable for LLM inference with variable sequence lengths.\n\nThe commented-out `@triton.autotune` decorators in `dyt.py` and `grpo_loss.py` are evidence of prior exploration that was abandoned for these reasons. Do NOT re-introduce `@triton.autotune` on production kernels.\n\n### Parameters to Tune\n\n| Parameter | Description | Typical Range |\n|-----------|-------------|---------------|\n| `BLOCK_SIZE` | Elements processed per program instance | 256, 512, 1024, 2048, 4096, 8192, 16384 |\n| `num_warps` | Warps per program instance (32 threads each) | 2, 4, 8, 16, 32 |\n| `num_stages` | Software pipelining depth | 1, 2, 3, 4, 5 |\n| `BLOCK_ROW` | Rows per program (block-row mode only) | 1, 2, 4, 8, 16, 32 |\n\n### Method: Manual Sweep with Improved Heuristic\n\nReplace `calculate_settings()` with a better heuristic derived from benchmarking. The variant must still pass `BLOCK_SIZE` and `num_warps` through `ctx` to the backward kernel.\n\n**Step 1**: Identify the current config for the benchmark's input sizes:\n```python\n# calculate_settings() returns (next_power_of_2(n), num_warps_heuristic)\n# For n_cols=4096: BLOCK_SIZE=4096, num_warps=8\n# For n_cols=8192: BLOCK_SIZE=8192, num_warps=16\n# Note: num_stages is never set (Triton defaults to 1)\n```\n\n**Step 2**: Create a variant with an improved selection function:\n```python\ndef optimized_settings(n):\n    \"\"\"Improved heuristic derived from benchmarking on {gpu_arch}.\"\"\"\n    BLOCK_SIZE = triton.next_power_of_2(n)\n    if BLOCK_SIZE > 65536:\n        raise RuntimeError(...)\n    # Tuned values (from benchmark sweep):\n    if BLOCK_SIZE >= 8192:\n        num_warps = 16\n        num_stages = 3   # NEW: calculate_settings doesn't set this\n    elif BLOCK_SIZE >= 2048:\n        num_warps = 8\n        num_stages = 2\n    else:\n        num_warps = 4\n        num_stages = 2\n    return BLOCK_SIZE, num_warps, num_stages\n```\n\n**Step 3**: Update the kernel launcher to pass `num_stages`:\n```python\n# The kernel launch gains num_stages (previously unset, defaulted to 1):\nkernel[(grid,)](..., BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages)\n```\n\n**Step 4**: Update `ctx` to also store `num_stages` and pass it to the backward kernel.\n\n**Key insight**: `num_stages` is the most commonly overlooked parameter. `calculate_settings()` does not set it at all, so Triton defaults to `num_stages=1`. Simply adding `num_stages=2` or `num_stages=3` to existing kernel launches is often the single biggest win, especially on Hopper.\n\n### Architecture-Specific Guidance\n\n- **Ampere (SM 8.x)**: `num_stages=2` is usually optimal. Higher stages rarely help and can increase register pressure. `num_warps` sweet spot is 4-16.\n- **Hopper (SM 9.0)**: Benefits significantly from higher `num_stages` (3-5) due to TMA and async pipeline support. Also tolerates `num_warps=32` better due to deeper warp scheduling.\n- **Blackwell (SM 10.0)**: Similar to Hopper but with even more pipeline depth tolerance. Start with `num_stages=4-5`.\n\n### When to Apply\n\nAlways. This is the mandatory first optimization attempt for every kernel.\n\n### Expected Impact\n\n5-40% speedup depending on how far the `calculate_settings()` heuristic is from optimal for the target shapes and GPU. Adding `num_stages` alone can yield 10-25% on Hopper.\n\n### Risks/Gotchas\n\n- Configs are shape-dependent. A config optimal for `n_cols=4096` may be suboptimal for `n_cols=128`. Always benchmark across ALL x_values from the existing benchmark suite.\n- If BLOCK_SIZE exceeds `n_cols`, the kernel still works (masked loads) but wastes warps. Prune configs where `BLOCK_SIZE >> n_cols`.\n- The `calculate_settings()` function enforces `MAX_FUSED_SIZE = 65536`. This limit exists to prevent register spilling. Respect this unless verified otherwise.\n- When adding `num_stages` to existing kernels, the backward kernel must also receive the same `num_stages` value via `ctx`.\n- On NPU/Ascend, `num_warps` and `num_stages` are ignored. Optimizations targeting these parameters are NVIDIA-specific.\n\n### Kernel Tiers\n\nAll tiers benefit. Tier 1 kernels often see the largest relative gains because their simple structure makes the config choice proportionally more impactful.\n\n---\n\n## 2. Memory-Bound Optimizations\n\nApply these when the profiler classifies the kernel as memory-bound (high memory throughput, low compute utilization, simple operations per load/store).\n\n### 2.1 Memory Access Coalescing\n\n**When to apply**: NCU shows low memory throughput efficiency, or kernel uses non-contiguous access patterns (e.g., strided access across rows, column-major access on row-major data).\n\n**How to implement**:\n\nEnsure that threads in a warp access adjacent memory addresses. In Triton, this means the innermost offset computation should produce contiguous addresses across the `tl.arange` dimension.\n\n```python\n# GOOD: contiguous access within a warp\ncol_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_idx * stride + col_offsets, mask=col_offsets < n_cols)\n\n# BAD: strided access (each thread hits a different cache line)\nrow_offsets = tl.arange(0, BLOCK_SIZE)\nX_block = tl.load(X_ptr + row_offsets * stride + col_idx)\n```\n\nIn Liger kernels, most access is already contiguous via `col_offsets = tl.arange(0, BLOCK_SIZE)`. Check for cases where stride-based pointer arithmetic breaks coalescing, especially in backward kernels or multi-dimensional grids.\n\n**Expected impact**: 10-50% when fixing uncoalesced access. Minimal gain if access is already coalesced.\n\n**Risks**: Changing access patterns may require transposing data or restructuring the grid, which can be invasive.\n\n**Kernel tiers**: All tiers, but Tier 2/3 backward passes most often have coalescing issues due to accumulation patterns.\n\n### 2.2 Cache Modifiers\n\n**When to apply**: Kernel reads data that is used exactly once (stream through), or reads data that is reused many times (should be cached aggressively). NCU shows poor L2 hit rate.\n\n**How to implement**:\n\nTriton supports cache modifiers via the `eviction_policy` parameter on `tl.load` and `tl.store`:\n\n```python\n# Data read once, don't pollute cache (streaming access)\nX_block = tl.load(X_ptr + offsets, eviction_policy=\"evict_first\")\n\n# Data reused across iterations, keep in cache\nW_row = tl.load(W_ptr + offsets, eviction_policy=\"evict_last\")\n\n# Default (let hardware decide)\nX_block = tl.load(X_ptr + offsets)  # no modifier\n```\n\n| Modifier | Triton API | When to Use |\n|----------|------------|-------------|\n| `.ca` (cache all) | `eviction_policy=\"evict_last\"` | Data reused across loop iterations or by other programs |\n| `.cs` (cache streaming) | `eviction_policy=\"evict_first\"` | Data used once; large tensors that would thrash cache |\n| `.wb` (write-back) | default store behavior | Standard write-back through cache hierarchy |\n| `.wt` (write-through) | not directly exposed | Write-through (rarely needed in Triton) |\n\n**Expected impact**: 5-20% on memory-bound kernels with predictable access patterns.\n\n**Risks**: Wrong eviction policy can degrade performance. `.evict_first` on reused data will cause redundant memory fetches. Always benchmark both policies against the default.\n\n**Kernel tiers**: Tier 2 and Tier 3 benefit most (weight reuse in backward passes, multi-pass algorithms).\n\n### 2.3 Reduce Memory Traffic\n\n**When to apply**: Kernel reads or writes data that could be computed on the fly, or writes intermediate results that are only consumed once.\n\n**How to implement**:\n\nLiger already uses several memory-saving patterns. Look for additional opportunities:\n\n```python\n# Pattern 1: Recompute instead of saving to memory\n# Instead of: save silu(a) to a buffer, load in backward\n# Do: recompute silu(a) in backward from saved a\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a  # recomputed in backward\n\n# Pattern 2: In-place gradient storage (gradient-in-forward trick)\n# Cross entropy already does this: stores gradient in the input tensor\n# Look for other kernels that allocate separate grad tensors unnecessarily\n\n# Pattern 3: Fuse producer-consumer pairs to avoid materializing intermediates\n# Instead of: Y = kernel1(X); Z = kernel2(Y)  (Y written to HBM, then read)\n# Do: Z = fused_kernel(X)  (Y lives only in registers/SRAM)\n\n# Pattern 4: Write in-place to dY for dX when shapes match\nif in_place:\n    dX = dY  # reuse memory, zero additional allocation\n```\n\n**Expected impact**: 10-40% from eliminating unnecessary memory round-trips. Memory savings also reduce peak memory.\n\n**Risks**: Recomputation trades memory for compute. Only beneficial when the kernel is memory-bound. Can also complicate backward pass correctness -- always verify gradient accuracy.\n\n**Kernel tiers**: All tiers. Tier 3 kernels (fused operations) have the most opportunity for intermediate elimination.\n\n### 2.4 Data Reuse (Pre-Load Outside Loops)\n\n**When to apply**: A value is loaded inside a loop but does not change across iterations. Common with weight tensors and scaling factors.\n\n**How to implement**:\n\n```python\n# BEFORE: weight loaded inside the loop (redundant loads if compiler doesn't hoist)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    W_block = tl.load(W_ptr + offsets)  # loaded every iteration if W is large\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_block\n\n# AFTER: hoist the invariant load\nW_row = tl.load(W_ptr + col_offsets, mask=col_offsets < n_cols)\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(X_ptr + offsets)\n    Y_block = X_block * W_row  # reuse from registers\n```\n\nNote: Triton's compiler often hoists loads automatically, but it is not guaranteed for all patterns, especially when the loop body is complex.\n\n**Expected impact**: 5-15% when applicable.\n\n**Risks**: Hoisting large loads outside loops increases register pressure. If the weight vector is very large relative to BLOCK_SIZE, it may cause register spilling that offsets the benefit.\n\n**Kernel tiers**: Tier 2 (norm kernels with weight vectors), Tier 3 (cross-entropy with weight tensors).\n\n### 2.5 Block-Row Mode (Multiple Rows per Program)\n\n**When to apply**: The kernel processes one row per program instance, but each row is small (e.g., `n_cols <= 256`). This under-utilizes each SM because the program does very little work.\n\n**How to implement**:\n\nLiger's `rms_norm.py` already implements this as `_block_rms_norm_forward_kernel`. The pattern:\n\n```python\n@triton.jit\ndef kernel_block_row(\n    ...,\n    n_rows,\n    n_cols,\n    BLOCK_SIZE: tl.constexpr,\n    BLOCK_ROW: tl.constexpr,  # number of rows per program\n):\n    # Each program processes BLOCK_ROW rows\n    row_idx = tl.program_id(0) * BLOCK_ROW + tl.arange(0, BLOCK_ROW)\n    col_offsets = tl.arange(0, BLOCK_SIZE)\n    row_mask = row_idx < n_rows\n    col_mask = col_offsets < n_cols\n\n    # 2D load: [BLOCK_ROW, BLOCK_SIZE]\n    X = tl.load(\n        X_ptr + row_idx[:, None] * stride + col_offsets[None, :],\n        mask=row_mask[:, None] & col_mask[None, :],\n        other=0,\n    )\n    # Process all rows at once...\n```\n\nGrid: `(triton.cdiv(n_rows, BLOCK_ROW),)` instead of `(n_rows,)`.\n\nUse conditional dispatch to choose between row mode and block-row mode:\n\n```python\nif BLOCK_SIZE > 256 or n_rows < 4096 * 8:\n    # Row mode: one row per program\n    kernel_row[(n_rows,)](...)\nelse:\n    # Block-row mode: multiple rows per program\n    BLOCK_ROW = 16\n    kernel_block_row[(triton.cdiv(n_rows, BLOCK_ROW),)](\n        ..., BLOCK_ROW=BLOCK_ROW\n    )\n```\n\n**Expected impact**: 10-30% for kernels with small `n_cols` and large `n_rows`.\n\n**Risks**: Increases register pressure and shared memory usage. The 2D load pattern is less efficient when `BLOCK_ROW` causes register spilling. Also adds code complexity (two kernel variants). Does not help when `n_cols` is already large.\n\n**Kernel tiers**: Tier 1 and Tier 2 primarily. Many Liger Tier 1 kernels (swiglu, geglu, relu_squared) do not yet have block-row mode.\n\n### 2.6 Vectorized Loads\n\n**When to apply**: Kernel loads/stores many small elements individually. Memory bandwidth is limited by the number of transactions rather than total bytes.\n\n**How to implement**:\n\nTriton generally handles vectorization automatically when loads are contiguous and aligned. However, you can help by:\n\n1. Ensuring BLOCK_SIZE is a multiple of 8 (for fp16/bf16) or 4 (for fp32) to enable 128-bit loads.\n2. Ensuring pointer alignment (tensors created by PyTorch are typically aligned).\n3. Avoiding masked loads at boundaries where possible -- pad the input instead.\n\n```python\n# Triton auto-vectorizes contiguous loads. Just ensure alignment:\n# BLOCK_SIZE should be power of 2 (already enforced by calculate_settings)\n# Data should be contiguous (already enforced by ensure_contiguous decorator)\n\n# For explicit vectorization (rare, advanced):\n# Cast pointer to wider type, load fewer elements\n# This is rarely needed as Triton handles it, but can help in edge cases\n```\n\n**Expected impact**: 5-10% if loads are currently unvectorized. Usually Triton already handles this.\n\n**Risks**: Minimal. Padding inputs wastes a small amount of memory.\n\n**Kernel tiers**: All tiers, but gains are usually small since Triton auto-vectorizes well.\n\n---\n\n## 3. Compute-Bound Optimizations\n\nApply these when the profiler classifies the kernel as compute-bound (high compute utilization, low memory throughput utilization, complex math per element).\n\n### 3.1 Algorithmic Improvements\n\n**When to apply**: The kernel uses a naive algorithm where a more efficient one exists. Examples: naive softmax vs online softmax, explicit matrix inverse vs iterative method.\n\n**How to implement**:\n\nLiger already uses online softmax in cross_entropy. Look for similar opportunities:\n\n```python\n# Online algorithm pattern: single-pass with running statistics\n# Instead of: pass 1 = compute max, pass 2 = compute sum, pass 3 = compute output\n# Do: single pass maintaining running max and sum\n\nm = float(\"-inf\")\nd = 0.0\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    block_max = tl.max(X_block)\n    m_new = tl.maximum(m, block_max)\n    d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new))\n    m = m_new\n# Result: lse = m + tl.log(d)\n```\n\nOther algorithmic opportunities:\n- Welford's algorithm for variance (numerically stable single-pass)\n- Kahan summation for high-precision reduction\n- Fast approximations for transcendentals when full precision is not needed\n\n**Expected impact**: 20-60% when applicable (reducing O(n) passes).\n\n**Risks**: Numerical stability can differ from reference implementation. Always validate against PyTorch reference with tight tolerances. Online algorithms may accumulate floating-point errors differently.\n\n**Kernel tiers**: Tier 2 and Tier 3 primarily. Tier 1 kernels rarely have algorithmic improvement opportunities.\n\n### 3.2 Reduce Redundant Computation\n\n**When to apply**: The same sub-expression is computed multiple times within a kernel, or across forward and backward passes.\n\n**How to implement**:\n\n```python\n# BEFORE: sigmoid computed twice\nsilu_a = a_row * tl.sigmoid(a_row)\ngrad = dc_row * (silu_a * (1 - tl.sigmoid(a_row)) + tl.sigmoid(a_row)) * b_row\n\n# AFTER: compute once, reuse\nsig_a = tl.sigmoid(a_row)\nsilu_a = a_row * sig_a\ngrad = dc_row * (silu_a * (1 - sig_a) + sig_a) * b_row\n```\n\nLiger already does this well (see swiglu backward). Look for:\n- Repeated `tl.exp()`, `tl.sigmoid()`, `tl.log()` calls with the same argument\n- Reduction results (sum, max) computed multiple times\n- Softmax probabilities recomputed instead of cached in registers\n\n**Expected impact**: 5-20% depending on the cost of the duplicated operation.\n\n**Risks**: Caching sub-expressions increases register pressure. Trade-off with register spilling.\n\n**Kernel tiers**: Tier 2 and Tier 3. Tier 1 kernels are typically simple enough that redundancy is already eliminated.\n\n### 3.3 Tensor Core Utilization\n\n**When to apply**: Kernel performs matrix-multiply-like operations (dot products, outer products) that could leverage tensor cores. This is most relevant for Tier 3 kernels with linear layers or attention-like patterns.\n\n**How to implement**:\n\nTensor cores require specific alignment:\n- Block dimensions must be multiples of 16 (for fp16/bf16) or 8 (for tf32)\n- Input/output matrices must be in supported layouts\n- Use `tl.dot()` which maps to tensor core instructions when dimensions align\n\n```python\n# Ensure BLOCK_M, BLOCK_N, BLOCK_K are multiples of 16\nBLOCK_M: tl.constexpr = 64  # multiple of 16\nBLOCK_N: tl.constexpr = 64  # multiple of 16\nBLOCK_K: tl.constexpr = 32  # multiple of 16\n\n# tl.dot maps to tensor cores automatically\nacc = tl.dot(A_block, B_block)  # [BLOCK_M, BLOCK_K] x [BLOCK_K, BLOCK_N]\n```\n\nFor `tl.dot` to use tensor cores:\n- Both operands must be 2D\n- Operand dtypes must be fp16, bf16, or tf32\n- Inner dimension K must be >= 16\n\n**Expected impact**: 2-10x for matrix-multiply-heavy kernels. Not applicable to element-wise or simple reduction kernels.\n\n**Risks**: Tensor core operations have lower precision than FP32 scalar operations. The `tf32` mode on Ampere+ truncates FP32 mantissa to 10 bits. Always validate numerical accuracy. Not all kernels have operations that map to tensor cores.\n\n**Kernel tiers**: Tier 3 only (fused_linear_cross_entropy, attention-like kernels). Not applicable to Tier 1 or most Tier 2 kernels.\n\n### 3.4 Register Pressure Reduction\n\n**When to apply**: NCU shows high register usage per thread (>128 registers) or occupancy is limited by register count. Kernel body has many live variables.\n\n**How to implement**:\n\n```python\n# Strategy 1: Use constexpr branches to eliminate dead code paths\n# Triton compiles out branches on constexpr values, reducing register usage\ncasting_mode: tl.constexpr  # compiler eliminates unused casting branches\nHAS_WEIGHT: tl.constexpr    # compiler eliminates weight-related registers when False\n\n# Strategy 2: Split large kernels into smaller passes\n# Instead of one kernel doing everything:\n# Kernel 1: forward + save intermediates\n# Kernel 2: backward using intermediates\n# This is a last resort as it adds kernel launch overhead\n\n# Strategy 3: Reduce live range of temporary variables\n# Reorder operations so temporary values are consumed immediately\nfor i in range(0, n_cols, BLOCK_SIZE):\n    X_block = tl.load(...)\n    # Use X_block immediately, don't hold it across iterations\n    result = compute(X_block)\n    tl.store(..., result)\n    # X_block and result are now dead, registers freed\n\n# Strategy 4: Use smaller data types for intermediates where precision allows\n# fp16 uses half the registers of fp32\nintermediate = X_block.to(tl.float16)  # only if precision loss is acceptable\n```\n\n**Expected impact**: 10-25% when occupancy is register-limited. Minimal impact when occupancy is already high.\n\n**Risks**: Splitting kernels adds launch overhead and memory traffic for intermediates. Reducing precision can cause numerical issues. Constexpr branches are already used extensively in Liger -- check for remaining opportunities.\n\n**Kernel tiers**: Tier 3 (most register pressure due to complex logic). Tier 2 backward passes (accumulation variables).\n\n### 3.5 2D Grid Tiling for Better SM Utilization\n\n**When to apply**: Kernel uses a 1D grid but the problem has two natural dimensions (e.g., batch and feature). A 1D grid may leave SMs idle when the number of programs is not a multiple of SM count.\n\n**How to implement**:\n\n```python\n# 1D grid (current Liger pattern for most kernels):\ngrid = (n_rows,)\n# If n_rows = 100 and SM count = 108, 8 SMs are idle in the last wave\n\n# 2D grid for better utilization:\ngrid = (triton.cdiv(n_rows, BLOCK_ROW), triton.cdiv(n_cols, BLOCK_COL))\n\n@triton.jit\ndef kernel_2d(\n    ...,\n    BLOCK_ROW: tl.constexpr,\n    BLOCK_COL: tl.constexpr,\n):\n    row_block = tl.program_id(0)\n    col_block = tl.program_id(1)\n    row_start = row_block * BLOCK_ROW\n    col_start = col_block * BLOCK_COL\n    # Process a tile of [BLOCK_ROW, BLOCK_COL]\n```\n\nA 2D grid is most useful when the problem is large in both dimensions and each tile does enough work to amortize launch overhead.\n\n**Expected impact**: 5-15% for kernels with suboptimal SM utilization. More impactful when `n_rows` is small relative to SM count.\n\n**Risks**: 2D tiling for reduction kernels requires careful handling of partial results and synchronization. More complex code. Not beneficial when `n_cols` fits in a single block.\n\n**Kernel tiers**: Tier 2 (reductions with large feature dimensions), Tier 3 (large fused operations).\n\n### 3.6 Loop Unrolling Hints\n\n**When to apply**: Kernel has a hot loop with a known iteration count or small trip count. Triton's compiler unrolls some loops automatically, but explicit hints can help.\n\n**How to implement**:\n\n```python\n# Triton unrolls loops with constexpr bounds automatically.\n# For dynamic bounds, you can help by:\n\n# Strategy 1: Make the loop bound a constexpr when possible\n# If n_cols is known at compile time (e.g., passed as constexpr):\nfor i in range(0, N_COLS, BLOCK_SIZE):  # N_COLS: tl.constexpr -> unrolled\n    ...\n\n# Strategy 2: Use tl.static_range for explicit unrolling\nfor i in tl.static_range(4):  # always unrolled\n    X_block = tl.load(X_ptr + i * BLOCK_SIZE + offsets)\n    ...\n\n# Strategy 3: Break a large loop into a known-count outer loop + remainder\nCHUNKS: tl.constexpr = n_cols // BLOCK_SIZE  # constexpr division\nfor i in tl.static_range(CHUNKS):\n    ...\n# Handle remainder separately\n```\n\n**Expected impact**: 5-10% for tight loops with few iterations. Diminishing returns for large trip counts.\n\n**Risks**: Excessive unrolling increases instruction cache pressure and compile time. Can also increase register pressure if the loop body is large.\n\n**Kernel tiers**: Tier 2 (reduction loops), Tier 3 (multi-pass loops).\n\n---\n\n## 4. Architecture-Specific Optimizations\n\nApply these on top of the best general-purpose variant when targeting a specific GPU architecture.\n\n### 4.1 Ampere (SM 8.x: A100, A10, A30)\n\n**Async Copy (cp.async)**:\n- Triton uses `num_stages` to control software pipelining, which maps to `cp.async` on Ampere.\n- `num_stages=2` is the sweet spot for most Ampere kernels. Going higher increases shared memory usage without benefit.\n- Triton handles `cp.async` automatically when `num_stages > 1`.\n\n**L2 Cache Persistence**:\n- A100 has 40MB L2 cache. Tensors smaller than this can benefit from persistence hints.\n- Use `eviction_policy=\"evict_last\"` on reused data (weights, biases).\n- For cross-entropy, the target tensor and weight tensor are good candidates.\n\n**Optimal Warp Counts**:\n- Ampere excels with 4-16 warps per SM. 32 warps often increases scheduling overhead.\n- For memory-bound kernels: fewer warps (4-8) with larger BLOCK_SIZE.\n- For compute-bound kernels: more warps (8-16) to hide latency.\n\n**Implementation**:\n```python\n# Ampere-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=2),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=2),\n```\n\n### 4.2 Hopper (SM 9.0: H100, H200)\n\n**TMA (Tensor Memory Accelerator)**:\n- Hopper's TMA offloads address computation and data movement to a dedicated unit.\n- Triton >= 3.0 exposes TMA via `tl.make_block_ptr` and experimental APIs.\n- TMA is most beneficial for 2D tiled access patterns (matmul-like kernels).\n\n```python\n# TMA-friendly block pointer (Triton >= 3.0)\nblock_ptr = tl.make_block_ptr(\n    base=X_ptr,\n    shape=(M, N),\n    strides=(stride_m, stride_n),\n    offsets=(row_start, col_start),\n    block_shape=(BLOCK_M, BLOCK_N),\n    order=(1, 0),  # row-major\n)\nX_block = tl.load(block_ptr, boundary_check=[0, 1])\n```\n\n**Higher num_stages**:\n- Hopper's async pipeline supports deeper prefetching. `num_stages=3-5` often outperforms 2.\n- The deeper pipeline hides memory latency more effectively on H100's higher-bandwidth HBM3.\n\n**WGMMA (Warp Group Matrix Multiply Accumulate)**:\n- H100 tensor cores are organized in warp groups (4 warps = 128 threads).\n- `num_warps` should be a multiple of 4 for optimal tensor core usage.\n- Triton maps `tl.dot` to WGMMA automatically when conditions are met.\n\n**Cluster Launch**:\n- Hopper supports thread block clusters (groups of SMs that can share SRAM).\n- Currently experimental in Triton. Potentially useful for large reductions.\n\n**Implementation**:\n```python\n# Hopper-tuned config\ntriton.Config({\"BLOCK_SIZE\": 4096}, num_warps=8, num_stages=4),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=32, num_stages=4),\n```\n\n### 4.3 Blackwell (SM 10.0: B200, GB200)\n\n**5th-Gen Tensor Cores**:\n- Higher throughput for FP8, FP16, BF16, and TF32 operations.\n- Block sizes aligned to 16 remain optimal. FP8 support enables even higher throughput.\n- Triton support is evolving; check Triton release notes for Blackwell-specific features.\n\n**Larger Shared Memory**:\n- Blackwell supports up to 256KB shared memory per SM (vs 228KB on Hopper).\n- Larger BLOCK_SIZE values become viable without spilling to registers or HBM.\n- Can enable larger `BLOCK_ROW` in block-row mode.\n\n**Higher Memory Bandwidth**:\n- HBM3E provides more bandwidth than HBM3. Memory-bound kernels see direct benefit.\n- The compute-to-memory ratio shifts, making some previously memory-bound kernels balanced or compute-bound.\n\n**Implementation**:\n```python\n# Blackwell-tuned config (aggressive block sizes)\ntriton.Config({\"BLOCK_SIZE\": 8192}, num_warps=16, num_stages=5),\ntriton.Config({\"BLOCK_SIZE\": 16384}, num_warps=32, num_stages=5),\n```\n\n**Note**: Blackwell-specific Triton features are still maturing. Focus on general optimizations first, then add Blackwell-specific configs as Triton support stabilizes.\n\n---\n\n## 5. Balanced Optimizations\n\nThese techniques help both memory-bound and compute-bound kernels.\n\n### 5.1 Kernel Fusion Opportunities\n\n**When to apply**: Two or more kernels execute sequentially on the same data, writing intermediate results to HBM between them. The fused version keeps intermediates in registers/SRAM.\n\n**How to implement**:\n\nLiger already excels at fusion (cross_entropy fuses softmax+loss+gradient; fused_linear_cross_entropy fuses matmul+CE). Look for remaining opportunities:\n\n```python\n# Pattern: operation followed by elementwise scaling\n# BEFORE: two kernel launches\ny = rms_norm_kernel(x, w)        # writes y to HBM\nz = elementwise_mul_kernel(y, s)  # reads y from HBM, writes z\n\n# AFTER: fused kernel\n@triton.jit\ndef fused_rms_norm_mul(X_ptr, W_ptr, S_ptr, Z_ptr, ...):\n    # RMS norm computation\n    y = x_row * rstd * (offset + w_row)\n    # Elementwise mul (no HBM round-trip)\n    z = y * s\n    tl.store(Z_ptr + offsets, z)\n```\n\nCommon fusion candidates in Liger:\n- Norm + activation (e.g., rms_norm + swiglu in some architectures)\n- Loss + metric computation (already done: cross_entropy + accuracy)\n- Dropout + residual add + norm\n\n**Expected impact**: 20-50% from eliminating HBM round-trips between fused operations.\n\n**Risks**: Fused kernels are harder to maintain, test, and debug. Register pressure increases. The fused kernel may have worse autotuning characteristics. Only fuse when the intermediate tensor is large relative to cache.\n\n**Kernel tiers**: Creates new Tier 3 kernels from combinations of Tier 1/2 kernels. Existing Tier 3 kernels may have remaining fusion opportunities.\n\n### 5.2 Better Grid Sizing\n\n**When to apply**: The number of programs (grid size) is much smaller or much larger than the SM count, leading to load imbalance or insufficient parallelism.\n\n**How to implement**:\n\n```python\nsm_count = torch.cuda.get_device_properties(device).multi_processor_count\n\n# For backward passes using SM-based parallelism (like rms_norm backward):\ngrid = (sm_count,)\nrows_per_program = math.ceil(n_rows / sm_count)\n\n# For forward passes: ensure grid is a small multiple of sm_count\n# If n_rows >> sm_count, this is automatic\n# If n_rows < sm_count, consider block-row mode or 2D grid to increase occupancy\n\n# For persistent kernels (advanced):\n# Launch exactly sm_count programs, each processes multiple work items\ngrid = (sm_count,)\nfor work_item in range(pid, total_work, NUM_SMS):\n    process(work_item)\n```\n\nThe rms_norm backward kernel already uses this SM-based grid pattern. Look for other backward kernels that use `(n_rows,)` grids but could benefit from SM-based parallelism.\n\n**Expected impact**: 5-20% from better load balancing and reduced tail effects.\n\n**Risks**: SM-based grids require careful work distribution and can introduce non-determinism if work items have unequal cost.\n\n**Kernel tiers**: All tiers. Most impactful for Tier 2 backward passes.\n\n### 5.3 Reduce Kernel Launch Overhead\n\n**When to apply**: Kernel is called many times with very short execution time (< 10 microseconds). Launch overhead dominates.\n\n**How to implement**:\n\n```python\n# Strategy 1: Batch multiple operations into one kernel\n# Instead of launching N small kernels, launch one kernel with N work items\n\n# Strategy 2: Use CUDA graphs (from PyTorch side, not inside Triton)\n# Captures a sequence of kernel launches and replays them with minimal overhead\ngraph = torch.cuda.CUDAGraph()\nwith torch.cuda.graph(graph):\n    output = liger_kernel(input)\n# Replay: graph.replay()\n\n# Strategy 3: Increase work per kernel (larger BLOCK_SIZE, block-row mode)\n# Already covered in autotuning and block-row mode sections\n```\n\n**Expected impact**: 5-15% for very short kernels. Negligible for kernels that run > 100 microseconds.\n\n**Risks**: CUDA graphs have limitations (fixed input shapes, no dynamic control flow). Batching work items increases kernel complexity.\n\n**Kernel tiers**: Tier 1 primarily (short element-wise kernels). Tier 2/3 kernels usually run long enough that launch overhead is negligible.\n\n### 5.4 Mixed Precision Strategies\n\n**When to apply**: Kernel uses FP32 throughout but only needs FP32 for specific operations (reductions, accumulations). Or kernel uses FP16/BF16 throughout but suffers from numerical issues.\n\n**How to implement**:\n\nLiger already uses mixed precision extensively (e.g., rms_norm with casting_mode). Look for additional opportunities:\n\n```python\n# Pattern 1: Compute in fp32 only where needed (Llama-style)\nX_row = tl.load(...)              # fp16/bf16\nX_fp32 = X_row.to(tl.float32)     # upcast for reduction\nmean_sq = tl.sum(X_fp32 * X_fp32) / n_cols  # fp32 reduction\nrstd = rsqrt(mean_sq + eps)        # fp32\nresult = (X_row * rstd.to(X_row.dtype)) * W  # back to original dtype\n\n# Pattern 2: Accumulate in fp32, store in original dtype\nacc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)\nfor ...:\n    acc += (dY_row * X_row).to(tl.float32)  # accumulate in fp32\ntl.store(dW_ptr, acc.to(original_dtype))     # store in original dtype\n\n# Pattern 3: Use BF16 instead of FP16 for better dynamic range\n# BF16 has same exponent bits as FP32, so less overflow/underflow risk\n# But lower mantissa precision than FP16\n```\n\n**Expected impact**: 5-20% from reducing unnecessary FP32 computation. Can also improve accuracy if precision is currently insufficient.\n\n**Risks**: Under-casting causes numerical errors. Over-casting wastes compute. The casting_mode pattern in Liger (llama vs gemma vs none) shows that different models need different precision strategies. Always validate against reference implementation.\n\n**Kernel tiers**: All tiers. Tier 2 reductions are most sensitive to precision. Tier 1 kernels with transcendentals (sigmoid, exp) benefit from targeted FP32 upcast.\n\n---\n\n## Quick Reference: Technique Selection by Symptom\n\n| Symptom | Likely Bottleneck | Try These Techniques |\n|---------|------------------|---------------------|\n| NCU memory throughput > 60% | Memory-bound | 2.1, 2.2, 2.3, 2.5 |\n| NCU compute throughput > 60% | Compute-bound | 3.1, 3.2, 3.3, 3.4 |\n| Low SM occupancy | Register pressure or grid size | 3.4, 5.2, 1 (num_warps) |\n| Small n_cols, large n_rows | Under-utilized programs | 2.5, 5.2, 1 (BLOCK_SIZE) |\n| Large n_cols, small n_rows | Poor SM utilization | 3.5, 5.2 |\n| Multiple sequential kernel launches | Fusion opportunity | 5.1 |\n| Forward fast, backward slow | Backward-specific issue | 2.3, 2.4, 5.2 (SM-based grid) |\n| Performance varies wildly by shape | Bad fixed config | 1 (autotune with key) |\n| High register count in NCU | Register pressure | 3.4, 1 (smaller BLOCK_SIZE) |\n| Poor L2 cache hit rate | Cache misuse | 2.2, 2.4 |\n\n## Quick Reference: Technique Applicability by Kernel Tier\n\n| Technique | Tier 1 (element-wise) | Tier 2 (reduction) | Tier 3 (fused/complex) |\n|-----------|----------------------|--------------------|-----------------------|\n| 1. Autotuning | High | High | High |\n| 2.1 Coalescing | Low (usually ok) | Medium | Medium |\n| 2.2 Cache modifiers | Low | Medium | High |\n| 2.3 Reduce traffic | Medium | Medium | High |\n| 2.4 Data reuse | Low | Medium | Medium |\n| 2.5 Block-row mode | High | High | Low |\n| 2.6 Vectorized loads | Low | Low | Low |\n| 3.1 Algorithmic | Low | High | High |\n| 3.2 Reduce redundancy | Low | Medium | High |\n| 3.3 Tensor cores | N/A | N/A | High |\n| 3.4 Register pressure | Low | Medium | High |\n| 3.5 2D grid tiling | Low | Medium | Medium |\n| 3.6 Loop unrolling | Low | Medium | Medium |\n| 4.x Arch-specific | Medium | Medium | Medium |\n| 5.1 Kernel fusion | Medium (fuse into Tier 3) | Medium | High |\n| 5.2 Grid sizing | Low | High | Medium |\n| 5.3 Launch overhead | High | Low | Low |\n| 5.4 Mixed precision | Medium | High | High |\n\n<!-- File: .agents/skills/liger-kernel-perf/optimizer.md (Tokens: ~4424 | Category: anthropic-skill) -->\n# Optimizer Agent\n\nRuns an autonomous optimization loop over an existing Liger kernel, generating multiple variants, benchmarking each against the baseline, and selecting the best performer while maintaining correctness.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\", \"cross_entropy\", \"swiglu\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `max_variants`: maximum number of variants to try (default 8)\n- `target_metric`: optional concrete target (e.g., \"forward under 0.3ms at hidden_size=4096\")\n- `optimization_profile`: the profile produced by the Profiler Agent at `optimization/{kernel}/profile.md`\n- `original_kernel`: the original kernel snapshot at `optimization/{kernel}/original_{kernel}.py`\n\n## Pre-Requisites\n\nBefore starting the optimization loop:\n\n1. Verify the workspace exists: `optimization/{kernel}/` with `original_{kernel}.py` and `profile.md`\n2. Verify the benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`\n3. Verify the test file exists: `test/transformers/test_{kernel}.py`\n4. Read the optimization profile thoroughly — extract bottleneck classification, tier, current config, and recommended strategy order\n5. Read the original kernel code thoroughly — understand every `@triton.jit` function, the `torch.autograd.Function`, and the launcher logic\n6. Read `optimization-strategies.md` for the full catalog of available techniques\n\n## Variant Numbering\n\n- v0 = baseline (original kernel, already benchmarked by the Profiler)\n- v1, v2, ... = optimization variants generated by this agent\n- Each variant gets two files:\n  - `optimization/{kernel}/{kernel}_vN.py` — the modified kernel code\n  - `optimization/{kernel}/{kernel}_vN_notes.md` — the lab notebook entry\n\n## The Optimization Loop\n\n```\nvariant_number = 1\nbest_variant = 0  (baseline)\nconsecutive_no_improvement = 0\n\nwhile variant_number <= max_variants:\n    1. Read ALL prior variant notes (v1 through v{variant_number-1})\n    2. Decide what to try next (see Strategy Selection below)\n    3. Generate variant code\n    4. Write variant notes (hypothesis section)\n    5. Run smoke test → if FAIL, mark REJECTED, increment variant_number, continue\n    6. Run full benchmark → record results\n    7. Check guardrails → if FAIL, mark REJECTED, increment variant_number, continue\n    8. Update variant notes with actual results and verdict\n    9. If improvement over best_variant: update best_variant, reset consecutive_no_improvement\n   10. Else: increment consecutive_no_improvement\n   11. If consecutive_no_improvement >= 2: STOP (diminishing returns)\n   12. If target_metric is set and met: STOP (goal achieved)\n   13. Increment variant_number\n```\n\n## Strategy Selection\n\n### Always Try Parameter Tuning First (v1)\n\nThe first variant MUST be a parameter sweep over `BLOCK_SIZE`, `num_warps`, and `num_stages`. This is the cheapest optimization with the highest expected payoff.\n\n**Why NOT @triton.autotune**: Liger kernels use `torch.autograd.Function` where the forward pass stores `BLOCK_SIZE` and `num_warps` in `ctx` and passes them to the backward kernel. With `@triton.autotune`, these become compile-time constants that cannot be stored in `ctx` — the backward kernel would not know which config the forward used. Additionally, NPU/Ascend backends do not support `num_warps` or `num_stages` at all. For these reasons, `@triton.autotune` is intentionally not used in Liger production kernels.\n\n**Instead, do a manual parameter sweep**: Replace the `calculate_settings(n_cols)` call with hardcoded values and benchmark each configuration:\n\n```python\n# In the variant, replace:\n#   BLOCK_SIZE, num_warps = calculate_settings(n_cols)\n# With a specific config to test:\nBLOCK_SIZE = 2048   # Try: 256, 512, 1024, 2048, 4096, 8192\nnum_warps = 8       # Try: 2, 4, 8, 16, 32\nnum_stages = 2      # Try: 1, 2, 3, 4\n```\n\n**Sweep procedure for v1**:\n1. Read the current `calculate_settings(n_cols)` output for the benchmark's representative input sizes\n2. Create a variant that replaces `calculate_settings` with a custom function that returns a better config\n3. Test 3-5 promising configs (not the full cartesian product — use the architecture guidance below to prune)\n4. Pick the best config across ALL benchmark x_values (not just one size)\n5. The variant should implement a smarter `calculate_settings` replacement, not just a single hardcoded value\n\n**Architecture-specific guidance for config selection**:\n- **Ampere**: `num_stages=2`, `num_warps` up to 16, BLOCK_SIZE 1024-4096\n- **Hopper**: `num_stages=3-5`, `num_warps` up to 32, BLOCK_SIZE 2048-8192\n- **Blackwell**: `num_stages=4-5`, similar to Hopper\n\n**Important**: Keep the same forward/backward ctx pattern. Both kernels must use the same `BLOCK_SIZE` and `num_warps` for a given input. The variant must still pass these through `ctx` just like the original.\n\n### Subsequent Variants: Diagnosis-Driven\n\nAfter autotuning, choose the next technique based on:\n\n1. **The bottleneck classification from the profile** (memory-bound vs compute-bound)\n2. **Learnings from all prior variants** (what worked, what did not, what was surprising)\n3. **The recommended strategy order from the profile**\n\nConsult [optimization-strategies.md](optimization-strategies.md) for the full technique catalog. Common patterns by bottleneck type:\n\n**Memory-bound kernels** (Tier 1, simple element-wise):\n- Fuse multiple memory-bound operations into a single kernel\n- Add cache eviction hints (`tl.store(..., eviction_policy=\"evict_first\")`)\n- Increase BLOCK_SIZE to amortize memory latency\n- Use `tl.load(..., eviction_policy=\"evict_last\")` for reused data\n- Reduce unnecessary memory traffic (eliminate redundant loads/stores)\n- In-place operations where safe\n\n**Compute-bound kernels** (Tier 3, heavy math):\n- Reduce precision of intermediate computations where safe\n- Use fast-math intrinsics (`tl.math.fast_expf`, etc.)\n- Restructure computation to reduce total FLOPs\n- Improve instruction-level parallelism via loop unrolling\n- Increase occupancy by reducing register pressure\n\n**Balanced / Tier 2 kernels** (reductions):\n- Optimize the reduction pattern (two-pass vs single-pass)\n- Tune BLOCK_ROW for multi-row processing\n- Balance between register usage and parallelism\n- Consider split-K or hierarchical reductions for large dimensions\n- Block-row mode threshold tuning (the `n_rows < 4096 * 8` style checks)\n\n### Combining Techniques\n\nA single variant may combine multiple techniques if they are complementary (e.g., autotuning + cache hints). However, prefer isolated changes so that the lab notebook clearly attributes improvements to specific techniques. When combining, explain which part of the improvement is expected from each technique.\n\n### When All Strategies Exhausted\n\nIf all recommended strategies have been tried and the consecutive_no_improvement counter has not triggered, the agent may:\n- Try creative combinations of previous partial improvements\n- Attempt architecture-specific tricks (Hopper warp specialization, Blackwell TMA)\n- Revisit a rejected variant with a modified approach based on learnings\n\n## Generating a Variant\n\nFor each variant, produce a complete, self-contained kernel file at `optimization/{kernel}/{kernel}_vN.py`.\n\nRules:\n1. Start from the original kernel (not a previous variant) unless explicitly building on a prior variant's approach. If building on a prior variant, note which one and why.\n2. The variant file must be a drop-in replacement for `src/liger_kernel/ops/{kernel}.py` — same public API, same class names, same function signatures.\n3. Do not change the `torch.autograd.Function` class name or the public API.\n4. Do not remove safety checks, dtype handling, or multi-device support (XPU, NPU paths) unless the change is specifically about those areas.\n5. **Preserve ALL existing comments** — function docstrings, mathematical formula explanations, inline rationale comments, and section headers. Do NOT delete or shorten comments that explain what the code does or why.\n6. **Add explanatory comments for every optimization change** — each modified section should have a comment explaining what was changed and why it improves performance. For example:\n   ```python\n   # Optimization: compute dW before dX to allow register reuse.\n   # The original order kept 8 vectors live simultaneously (115 regs/thread).\n   # This reordering reduces peak liveness to 5-6 vectors (95 regs/thread).\n   ```\n7. **Preserve function-level docstrings** — if the original kernel functions have docstrings or parameter descriptions, keep them. Update them if the function signature changes.\n8. Add a comment block at the top of the file documenting what was changed:\n\n```python\n\"\"\"\nOptimization Variant {N}: {Short Title}\nBased on: original (or v{M} if building on a prior variant)\nChanges: {brief description of modifications}\n\"\"\"\n```\n\n## Variant Lab Notebook\n\nFor each variant, write `optimization/{kernel}/{kernel}_vN_notes.md` using the format in [templates/variant-notes.md](templates/variant-notes.md).\n\nKey requirements:\n- The Hypothesis must reference the profile's bottleneck classification and learnings from prior variants\n- Changes must include rationale for each modification with code snippets\n- Expected Outcome must be specific and quantitative (\"10-20% forward speed improvement\" not \"should be faster\")\n- Actual Results must include ALL x_values if improvement varies across sizes\n- Verdict: ACCEPTED (improves target without violating guardrails), REJECTED (fails smoke test/guardrails/no improvement), or MIXED (improves some metrics, regresses others within limits)\n- Learnings must include insights to carry forward to the next variant\n\n### Writing Notes Across the Loop\n\n- **Before smoke test**: Write the Hypothesis, Changes, and Expected Outcome sections\n- **After benchmark**: Fill in Actual Results, Verdict, and Learnings\n- **On smoke test failure**: Fill Actual Results with \"SMOKE TEST FAILED: {error}\", set Verdict to REJECTED, and write Learnings about what went wrong\n\n## Smoke Test Procedure\n\nThe smoke test validates basic correctness before investing time in a full benchmark. Execute it by temporarily swapping the variant into the source tree.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run a Minimal Correctness Check\n\n```bash\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80\n```\n\nThis runs only float32 tests (the most numerically stable) with a timeout to catch hangs. The `-x` flag stops at the first failure.\n\n### Step 3: Interpret Results\n\n- **All tests pass**: Smoke test PASSES. Proceed to full benchmark.\n- **Any test fails**: Smoke test FAILS. Read the error carefully.\n  - If the failure is a **tolerance issue** (allclose fails but values are close): Note this in the variant notes. The optimization may have changed numerical behavior slightly. Try adjusting the variant to preserve precision, or note the tolerance difference.\n  - If the failure is a **crash or shape error**: The variant has a bug. Mark REJECTED immediately. Record the error in the notes for learning.\n  - If the failure is a **hang/timeout**: The variant has a deadlock or infinite loop. Mark REJECTED immediately.\n\n### Step 4: Restore the Original\n\n**CRITICAL**: Always restore the original kernel after the smoke test, regardless of outcome:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\nThis must happen even if the smoke test crashes. Use a pattern that ensures restoration:\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py && \\\npython -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k \"float32\" 2>&1 | tail -80; \\\nSMOKE_EXIT=$?; \\\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py; \\\nexit $SMOKE_EXIT\n```\n\n## Full Benchmark Procedure\n\nAfter a variant passes the smoke test, run the full existing benchmark suite to get precise measurements.\n\n### Step 1: Swap In the Variant\n\n```bash\ncp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Step 2: Run the Full Benchmark\n\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nThis runs speed benchmarks (forward, backward, full) and memory benchmarks across all configured x_values, comparing the \"liger\" provider against the baseline (\"huggingface\" or \"torch\").\n\n### Step 3: Capture Results\n\nThe benchmark writes results to `benchmark/data/all_benchmark_data.csv`. Copy the relevant rows:\n\n```bash\ncp benchmark/data/all_benchmark_data.csv optimization/{kernel}/benchmarks/vN_results.csv\n```\n\n### Step 4: Parse Results\n\nRead the CSV and extract:\n- Speed (ms) at the median (p50) for forward, backward, and full modes\n- Memory (MB) at the median (p50) for full mode\n- Values at each x_value to detect size-dependent behavior\n\nCompare against:\n- v0 baseline: the absolute improvement\n- Current best variant: whether this is the new best\n\n### Step 5: Restore the Original\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n## Guardrail Checks\n\nAfter benchmarking each variant, apply these guardrails. All comparisons are against the v0 baseline.\n\n### Guardrail 1: Non-Target Metric Regression (>5%)\n\nIf the optimization goal is \"speed\", check that memory does not regress by more than 5%. If the goal is \"memory\", check that speed does not regress by more than 5%. If the goal is \"balanced\", check both.\n\n```\nif goal == \"speed\":\n    reject if memory_vN > memory_v0 * 1.05\nelif goal == \"memory\":\n    reject if speed_vN > speed_v0 * 1.05  (higher ms = slower = worse)\nelif goal == \"balanced\":\n    reject if memory_vN > memory_v0 * 1.05 OR speed_vN > speed_v0 * 1.05\n```\n\nCheck this at EVERY x_value. If the regression occurs only at certain sizes, note this in the variant notes but still reject — partial regressions are not acceptable.\n\n### Guardrail 2: Cross-Pass Regression (>10%)\n\nDo not sacrifice one pass (forward or backward) to improve the other beyond a 10% threshold:\n\n```\nreject if speed_fwd_vN > speed_fwd_v0 * 1.10 AND speed_bwd_vN < speed_bwd_v0\nreject if speed_bwd_vN > speed_bwd_v0 * 1.10 AND speed_fwd_vN < speed_fwd_v0\n```\n\nAn optimization that makes forward 15% slower to make backward 5% faster is rejected.\n\n### Guardrail 3: Smoke Test Failure\n\nAlready handled in the smoke test procedure. Any correctness failure means immediate rejection.\n\n### Handling Guardrail Failures\n\nWhen a variant is rejected by guardrails:\n1. Set Verdict to REJECTED in the variant notes\n2. Record exactly which guardrail was violated and by how much\n3. Record the Learnings — this information is valuable for future variants\n4. Do NOT increment consecutive_no_improvement (guardrail failures are not \"stagnation\")\n5. Proceed to the next variant\n\n## Stopping Conditions\n\nThe loop terminates when ANY of these conditions is met:\n\n### Condition 1: Budget Exhausted\n\n`variant_number > max_variants`\n\nReport: \"Optimization budget of {max_variants} variants exhausted.\"\n\n### Condition 2: Diminishing Returns\n\n`consecutive_no_improvement >= 2`\n\nTwo consecutive variants that pass guardrails but improve the target metric by less than 1% over the current best. This indicates the kernel is near its optimization ceiling for the current approach.\n\nReport: \"Stopping due to diminishing returns — last 2 accepted variants showed <1% improvement.\"\n\n### Condition 3: Target Metric Met\n\nThe user-specified `target_metric` has been achieved.\n\nReport: \"Target metric achieved by variant v{N}: {metric_description}.\"\n\n## Reading Prior Variant Notes\n\nBefore generating each new variant, the agent MUST read ALL prior variant notes. This is the accumulated learning mechanism that makes later variants smarter.\n\nWhen reading prior notes, extract:\n1. **What techniques were tried**: Avoid repeating the exact same approach\n2. **What worked**: Build on successful directions\n3. **What failed and why**: Avoid repeating mistakes\n4. **Surprising results**: These often reveal the true bottleneck\n5. **Size-dependent behavior**: Some optimizations help at large sizes but hurt at small sizes\n\nUse this information to form a more informed hypothesis for the next variant.\n\n## Final Comparison Table\n\nAfter the loop terminates (for any reason), produce a comparison table of ALL variants and save it to `optimization/{kernel}/comparison.md`.\n\nThe comparison must include:\n- A master table with ALL variants (accepted and rejected): Variant | Strategy | Speed Fwd/Bwd/Full (ms) | Memory (MB) | Verdict\n- Values at the representative x_value (typically the largest or most common) with percentage delta from baseline\n- Winner declaration with justification\n- Cumulative learnings summary (what worked, what didn't, why)\n\n## Error Recovery\n\n### Benchmark Script Crashes\n\nIf the benchmark script crashes (not a test failure — an actual crash):\n1. Restore the original kernel immediately\n2. Read the crash output\n3. If it is a Triton compilation error: the variant has a syntax or semantic error in the kernel. Mark REJECTED with the error.\n4. If it is a CUDA OOM: the variant uses too much memory for the GPU. Mark REJECTED. Note the memory constraint.\n5. If it is an unrelated infrastructure error: retry once. If it fails again, report to the user.\n\n### Workspace Corruption\n\nIf at any point `src/liger_kernel/ops/{kernel}.py` is found to differ from `optimization/{kernel}/original_{kernel}.py` when it should be the original:\n\n```bash\ndiff src/liger_kernel/ops/{kernel}.py optimization/{kernel}/original_{kernel}.py\n```\n\nIf they differ unexpectedly, restore immediately:\n\n```bash\ncp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py\n```\n\n### Agent Failure Mid-Loop\n\nIf the agent encounters an unrecoverable error:\n1. Restore the original kernel\n2. Save all variant notes written so far\n3. Write a partial comparison table with whatever data is available\n4. Report to the user with the current state and what went wrong\n\n## Important Constraints\n\n1. **Never leave a variant installed in the source tree.** The original kernel must always be restored after each smoke test and benchmark. This is the most critical invariant.\n2. **Never modify test files.** The existing tests are the correctness oracle. If a variant fails tests, the variant is wrong, not the tests.\n3. **Never modify benchmark files.** The existing benchmarks are the measurement standard. Variants must be measured by the same benchmark to enable fair comparison.\n4. **Always write notes before running experiments.** The hypothesis must be written before seeing results to avoid post-hoc rationalization.\n5. **Always read all prior notes before generating a new variant.** This is what makes the optimization loop learn from itself.\n6. **Prefer the simplest change that achieves the goal.** A 5-line autotuning change that yields 15% speedup is better than a 200-line rewrite that yields 20%.\n\n<!-- File: .agents/skills/liger-kernel-perf/profiler.md (Tokens: ~1371 | Category: anthropic-skill) -->\n# Profiler Agent\n\nProfiles an existing Liger kernel to diagnose performance bottlenecks and produce an optimization strategy.\n\n## Inputs\n\n- `target_kernel`: kernel name (e.g., \"rms_norm\")\n- `optimization_goal`: speed / memory / balanced\n- `scope`: general / forward / backward / specific input regime\n- `target_gpu`: Ampere / Hopper / Blackwell / auto-detect\n\n## Steps\n\n### Step 1: Set Up Workspace\n\n1. Create directory `optimization/{kernel}/` at the repo root\n2. Copy `src/liger_kernel/ops/{kernel}.py` to `optimization/{kernel}/original_{kernel}.py`\n3. Create `optimization/{kernel}/benchmarks/` subdirectory\n\n### Step 2: Detect GPU Architecture\n\nRun:\n```bash\npython -c \"import torch; props = torch.cuda.get_device_properties(0); print(f'GPU: {props.name}, SM: {props.major}.{props.minor}, SMs: {props.multi_processor_count}, Mem: {props.total_mem // (1024**3)}GB')\"\n```\n\nClassify architecture:\n- SM 8.x → Ampere (A100, A10, etc.)\n- SM 9.0 → Hopper (H100, H200, etc.)\n- SM 10.0 → Blackwell (B200, etc.)\n\nIf user specified a target_gpu, note the mismatch (if any) but proceed with the actual hardware.\n\n### Step 3: Run Baseline Benchmarks\n\nRun the existing benchmark script:\n```bash\ncd benchmark/scripts && python benchmark_{kernel}.py --overwrite\n```\n\nRecord the baseline results. Copy the relevant rows from `benchmark/data/all_benchmark_data.csv` to `optimization/{kernel}/benchmarks/v0_baseline.csv`.\n\nParse and summarize the baseline numbers:\n- Speed: forward, backward, full (median ms at each x_value)\n- Memory: full (median MB at each x_value)\n- Identify the slowest operation mode and largest x_values\n\n### Step 4: Check for NCU Availability\n\n```bash\nwhich ncu 2>/dev/null && echo \"NCU_AVAILABLE\" || echo \"NCU_NOT_AVAILABLE\"\n```\n\nIf available, offer to run profiling (interactive mode) or run it (autonomous mode).\n\nTo determine the correct import and invocation, read the kernel's `torch.autograd.Function` class from `src/liger_kernel/ops/{kernel}.py` and construct a minimal reproduction script. For example, for `rms_norm`:\n\n```bash\nncu --set full --target-processes all -o optimization/{kernel}/ncu_profile \\\n    python -c \"\nimport torch\nfrom liger_kernel.ops.rms_norm import LigerRMSNormFunction\nx = torch.randn(4096, 4096, device='cuda', dtype=torch.bfloat16, requires_grad=True)\nw = torch.ones(4096, device='cuda', dtype=torch.bfloat16)\ny = LigerRMSNormFunction.apply(x, w, 1e-6, 0.0, 'llama', False)\ny.sum().backward()\n\"\n```\n\nAdapt the import path, function arguments, and tensor shapes for the target kernel.\n\nExtract key metrics from NCU output:\n- **SM Occupancy** (%)\n- **Compute Throughput** (% of peak)\n- **Memory Throughput** (% of peak)\n- **L1/L2 Cache Hit Rates**\n- **Warp Stall Reasons** (top 3)\n\n### Step 5: Analyze Kernel Code\n\nRead `src/liger_kernel/ops/{kernel}.py` thoroughly. Extract:\n\n1. **Tier classification**:\n   - Tier 1 (element-wise): No reductions. One row per program.\n   - Tier 2 (reduction): Cross-column reductions. May need SM-based parallelism.\n   - Tier 3 (fused/complex): Multi-pass algorithms, gradient-in-forward tricks.\n\n2. **Current configuration**:\n   - How BLOCK_SIZE is set (calculate_settings? hardcoded? autotune?)\n   - num_warps value\n   - num_stages value (if specified)\n   - Grid dimensions (1D, 2D, block-row mode?)\n\n3. **Memory access patterns**:\n   - Stride-based or contiguous?\n   - Cache modifiers used?\n   - In-place operations?\n   - What is saved_for_backward vs recomputed?\n\n4. **Compute patterns**:\n   - Precision-sensitive operations (sigmoid, rsqrt, exp, log)?\n   - Casting modes?\n   - Online/chunked algorithms?\n   - constexpr parameters?\n\n5. **Existing optimization opportunities**:\n   - Any TODO/FIXME comments about performance?\n   - Missing block-row mode variant?\n   - Suboptimal BLOCK_SIZE heuristic?\n   - Unnecessary memory traffic?\n\n### Step 6: Classify Bottleneck\n\nUse NCU data (if available) or heuristics:\n\n**Memory-bound indicators**:\n- NCU: Memory throughput > 60% of peak, compute throughput < 40%\n- Heuristic: Simple element-wise ops, large tensors, few arithmetic operations per load\n- Heuristic: Tier 1 kernels are almost always memory-bound\n\n**Compute-bound indicators**:\n- NCU: Compute throughput > 60% of peak, memory throughput < 40%\n- Heuristic: Complex math per element (exp, log, sigmoid chains), reductions, many operations per load\n- Heuristic: Tier 3 kernels with heavy computation\n\n**Balanced (both bottlenecks)**:\n- NCU: Both throughputs 30-60%\n- Heuristic: Tier 2 kernels often fall here\n\n### Step 7: Produce Optimization Profile\n\nWrite `optimization/{kernel}/profile.md` using the [optimization-profile template](templates/optimization-profile.md).\n\nThe profile must include:\n- GPU architecture and properties\n- Baseline benchmark summary (table)\n- NCU summary (if available)\n- Kernel tier classification\n- Bottleneck classification with evidence\n- Current configuration analysis\n- Recommended strategy order (from [optimization-strategies.md](optimization-strategies.md))\n- Specific opportunities identified in the code\n\n### Step 8: Present to User (Interactive Mode)\n\nShow:\n1. GPU detected: {name} ({architecture})\n2. Baseline performance summary table\n3. Bottleneck classification: {memory-bound / compute-bound / balanced}\n4. Evidence for classification\n5. Proposed optimization strategy order\n6. Estimated iteration count\n\nWait for user confirmation before proceeding to Stage 2.\n\nIn autonomous mode, skip this step and proceed directly.\n\n<!-- ... +3 more modular skills available via /api/skills/linkedin/Liger-Kernel?path={path} -->\n</yakaai_skills>"}