mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
main
979
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 lowdy1
|
899233a6af |
[Benckmark] Benchmark refactor example (#1199)
## Summary
The current benchmark scripts contain significant boilerplate when
constructing `common_configs` for `run_benchmarks`. Although
`compute_model_config_sweep_config` and `compute_seq_len_sweep_config`
provide the core sweep logic, each script still:
* Reimplements probe logic
* Manually builds `extra_benchmark_config`
* Duplicates `common_configs` assembly
* Defines redundant helpers like `_resolve_*` and `*_model_config`
variants
This PR removes that duplication by introducing higher-level builders
that standardize how benchmarks are defined.
---
## Proposal
### 1. Introduce higher-level sweep builders
Add two unified helper functions in `benchmark_model_configs.py`:
```python
build_model_config_sweep(...)
build_token_length_sweep(...)
```
These functions:
* Wrap existing sweep utilities (`compute_*_sweep_config`)
* Internally handle memory probing via `setup_fn + forward_fn`
* Automatically construct `extra_benchmark_config` from:
* `model_keys` (dynamic model attributes)
* `extra_configs` (static overrides)
* Return a fully-formed `common_configs` dict
So benchmark scripts reduce to:
```python
common_configs = build_*(...)
run_benchmarks(**common_configs)
```
---
### 2. Standardize kernel definition via `setup_fn`
Instead of manually writing `probe_fn`, all kernels now define:
```python
setup_fn: SingleBenchmarkRunInput -> Tuple[Any, ...]
forward_fn: Tuple[Any, ...] -> torch.Tensor (optional)
```
The builders handle:
```python
setup_out = setup_fn(input)
output = forward_fn(*setup_out)
```
A default is provided:
```python
forward_fn = lambda x, layer: layer(x)
```
This removes duplicated forward/probe logic across scripts.
---
### 3. Eliminate redundant helpers
The following patterns are removed across benchmark scripts:
* `probe_fn` definitions
* `extra_config_fn`
* `_resolve_model_config_*`
* `bench_*_model_config` variants
All are now handled centrally by the builders.
---
## New APIs
### `build_model_config_sweep`
* Sweeps across model configurations (x-axis = model name)
* Keeps total tokens (`B * T`) approximately constant
* Uses `setup_fn + forward_fn` to estimate memory per model
```python
build_model_config_sweep(
kernel_name,
all_model_configs,
setup_fn,
model_keys,
forward_fn=...,
probe_provider="torch",
extra_configs=None,
bt=2048,
overwrite=False,
)
```
---
### `build_token_length_sweep`
* Sweeps across sequence length (x-axis = T)
* Automatically adjusts batch size based on memory estimation
* Uses the same `setup_fn + forward_fn` abstraction
```python
build_token_length_sweep(
kernel_name,
probe_seq_len,
model,
setup_fn,
model_keys,
extra_configs=None,
forward_fn=...,
probe_provider="torch",
x_values_fn=...,
overwrite=False,
)
```
---
## Example (after refactor)
```python
common_configs = build_token_length_sweep(
kernel_name="layer_norm",
probe_seq_len=1024,
model=model,
setup_fn=_setup_layer_norm,
model_keys=["hidden_size", "dtype"],
extra_configs={"eps": 1e-6},
probe_provider="huggingface",
)
common_configs["kernel_providers"] = ["liger", "huggingface"]
run_benchmarks(..., **common_configs)
```
---
## **Benchmark Command Examples**
```bash
python ./benchmark/scripts/benchmark_swiglu.py --sweep-mode model_config [--model llama_3_8b]
python ./benchmark/scripts/benchmark_swiglu.py [--sweep-mode token_length] [--bt 2048]
```
---
### Notes
* `model_config`: sweeps across different model configurations (fixed
total tokens)
* `token_length`: sweeps across sequence lengths / batch sizes (fixed
model) and this is the default and can be omitted
---
- Hardware Type: A100-80G-PCIe
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
|
||
|
Chun-Mao (Michael) Lai
|
c4b16d43f9 |
Upgrade version for release (#1212)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Upgrade version for release <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> NA <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergencev0.8.0 |
||
|
|
3dd9e3a2b7 |
Extend LigerExperts patching to qwen3_vl_moe and glm4v_moe (#1192)
## Summary
- Add fused MoE expert (LigerExperts) monkey-patching to
**qwen3_vl_moe** and **glm4v_moe**, the two remaining MoE models that
were missing it
- Fix buggy glm4v_moe instance patching that was outside the `for` loop
with duplicate logic and setting `decoder_layer.mlp = None`
- Enable 3 previously-skipped qwen3_vl_moe convergence tests (bf16/fp32
with_logits, fp32 multimodal)
- All patching gated behind `IS_TRANSFORMERS_V5_OR_LATER` for backward
compatibility with transformers v4
### Models patched
| Model | Class-level patch | Instance-level patch |
|---|---|---|
| `qwen3_vl_moe` | `Qwen3VLMoeTextExperts = LigerExperts` |
Attribute-based detection of `experts` on MoE layers |
| `glm4v_moe` | `Glm4vMoeTextNaiveMoe = LigerExperts` | Attribute-based
detection of `experts` + `shared_experts` |
### Why not llama4?
`Llama4TextExperts.forward(hidden_states)` takes only 1 arg (routing
done externally in MoE block), incompatible with LigerExperts'
`forward(hidden_states,
top_k_index, top_k_weights)`.
### Compatibility
Verified across transformers v4.57.6, v5.0–v5.4, and v5.5.4. Class names
and forward signatures are stable across all v5 releases.
## Test plan
- [x] Instance patching tests (`test_monkey_patch.py`) for all 3
qwen3_vl_moe variants and glm4v_moe
- [x] Convergence tests (bf16 + fp32, FLCE + with_logits) for
qwen3_vl_moe and glm4v_moe
- [x] Enabled 3 previously-skipped qwen3_vl_moe convergence tests
- [x] Verified backward compatibility: on v4,
`IS_TRANSFORMERS_V5_OR_LATER` gates skip all LigerExperts code paths
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>
|
||
|
Shivam Sharma
|
65de95822e |
perf(dyt): enable triton.autotune for DyT forward and backward kernels (#1208)
## Summary
Enables `@triton.autotune` on the DyT (Dynamic Tanh) forward and
backward kernels. The autotune configs were drafted by the original
author but left commented out with a static `BLOCK_N` heuristic in the
Python wrappers (`liger_dyt_fwd` / `liger_dyt_bwd`) as a placeholder.
### Changes
1. **Enable `@triton.autotune`** on both `_dyt_fwd_kernel` and
`_dyt_bwd_kernel` with a pruned config space (18 configs each):
- `BLOCK_N in {1024, 2048, 4096}`
- `num_stages in {1, 2}` (dropped `num_stages=4` from the original stub
— rarely wins for memory-bound elementwise kernels)
- `num_warps in {4, 8, 16}` (dropped `num_warps=32` from fwd stub —
overkill for this kernel)
2. **Remove the manual `BLOCK_N` heuristic** (`if N >= 4096:`) and
manual `kwargs` in both Python wrappers — autotune handles selection.
3. **Add `reset_to_zero=[\"DA\"]` to the backward autotune decorator**.
The `DA` reduction buffer is indexed by `program_id(0)`, so different
`BLOCK_N` configs write to different slot counts per SM. Autotune trials
don't zero output buffers between runs, so stale slots from a prior
config would leak into the final `da.sum()` of the winning config. Same
pattern as used in `fused_moe_kernels.py`.
### Performance (NVIDIA A10G, sm_86, bf16, BT=2048)
Measured via `benchmark/scripts/benchmark_dyt.py --sweep-mode
model_config` across 7 model hidden sizes (2048 -> 8192) x both
`beta=True`/`beta=False` x fwd/bwd/full. Median over `do_bench` trials.
**Forward pass (consistent wins):**
| Model | hidden | main | autotuned | speedup |
|---|---|---|---|---|
| deepseek_v2_lite | 2048 | 0.039ms | 0.039ms | 1.00x |
| qwen2.5_7b | 3584 | 0.066ms | 0.062ms | 1.05x |
| llama_2_7b | 4096 | 0.075ms | 0.072ms | 1.04x |
| llama_3_8b | 4096 | 0.075ms | 0.072ms | 1.04x |
| qwen2.5_14b | 5120 | 0.092ms | 0.089ms | 1.03x |
| deepseek_v3 | 7168 | 0.127ms | 0.124ms | 1.02x |
| qwen2.5_72b | 8192 | 0.144ms | 0.140ms | 1.03x |
**Full (fwd + bwd):**
| Model | main | autotuned | speedup |
|---|---|---|---|
| deepseek_v2_lite | 0.137ms | 0.136ms | 1.01x |
| qwen2.5_7b | 0.219ms | 0.214ms | 1.02x |
| llama_2_7b | 0.248ms | 0.244ms | 1.01x |
| llama_3_8b | 0.249ms | 0.246ms | 1.01x |
| qwen2.5_14b | 0.304ms | 0.298ms | 1.02x |
| deepseek_v3 | 0.419ms | 0.404ms | **1.04x** |
| qwen2.5_72b | 0.466ms | 0.462ms | 1.01x |
**Zero regressions observed across 42 (beta x mode x model)
combinations** (max noise: -1% on 2 configs, within `do_bench`
variance). Memory unchanged.
Backward shows smaller improvements because it's dominated by the per-SM
`rows_per_program` loop over M rather than the column axis.
### Caveats
- Autotune compile cost (~9s per unique `N`) is paid once per process
and absorbed by `triton.testing.do_bench`'s built-in warmup, so measured
speeds above are post-warmup steady-state.
- Config space was chosen based on A10G (Ampere) benchmarks.
Hopper/Blackwell users may benefit from additional configs; Triton's
autotune will still pick the best available from the current set.
## Testing Done
- Hardware: NVIDIA A10G (sm_86, Ampere)
- `pytest test/transformers/test_dyt.py -xvs` -> 24 passed (correctness
+ functional, both `beta=True`/`beta=False`, dtypes fp32 and bf16)
- `python benchmark/scripts/benchmark_dyt.py --sweep-mode model_config
--overwrite`
- [x] run `make test` to ensure correctness *(ran `pytest
test/transformers/test_dyt.py` — scoped to changed kernel)*
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence *(not applicable
— no numerical changes; autotune is a performance-only change)*
|
||
|
sunyi0505
|
dfbc4ddd38 |
[NPU] Modify the grid calculation logic to support larger batch sizes for group norm. (#1211)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
Chunyu
|
ed6b2ffc21 |
[fix] Distinguish the import path for compute_routing_metadata based on device type (#1209)
## Summary As Title. Fix: https://github.com/linkedin/Liger-Kernel/pull/1183#discussion_r3163137729 - Hardware Type: npu - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
Shivam Sharma
|
63598637cc |
Add Falcon H1 SwiGLU multipliers support (#936) (#1201)
## Summary Resolves #936. Adds Liger SwiGLU support for Falcon H1's MLP block, which pre-scales the gate pre-activation and post-scales the down-projection output: y = down_proj(silu(gate_proj(x) * gate_mult) * up_proj(x)) * down_mult Previously `apply_liger_kernel_to_falcon_h1(swiglu=True)` was a no-op that just logged a warning; now it actually patches the MLP. ## Details - `src/liger_kernel/ops/swiglu.py` — forward/backward Triton kernels extended with a `gate_multiplier` scalar applied to the gate pre-activation in fp32 before silu. Backward chain rule pulls an extra `gate_multiplier` factor into `da`. `LigerSiLUMulFunction` accepts `gate_multiplier` and `down_multiplier` as optional args (defaults 1.0), so all existing callers (Llama, Qwen, Granite, Phi3, Olmo, Hunyuan, etc.) are unaffected. `down_multiplier` is applied outside the kernel in the autograd function — a single elementwise multiply on the output, near-zero overhead. - `src/liger_kernel/transformers/swiglu.py` — new `LigerFalconH1SwiGLUMLP` reads `config.mlp_multipliers` and forwards them to `LigerSiLUMulFunction`. Falls back to `self.config.mlp_multipliers` when patched onto an already-instantiated HF module via `_patch_swiglu_module` (which only rebinds `forward`). - `src/liger_kernel/transformers/monkey_patch.py` — flip `apply_liger_kernel_to_falcon_h1`'s `swiglu` default from `False` to `True`, wire it to `LigerFalconH1SwiGLUMLP`, and fix the decoder-layer attribute to `feed_forward` (Falcon H1's actual HF attribute name). - `src/liger_kernel/transformers/__init__.py` — export `LigerFalconH1SwiGLUMLP`. - `test/transformers/test_swiglu.py` — three new parametrized tests (21 total cases): - `test_correctness_silumul_with_multipliers` — 12 cases (3 multiplier combos × 2 shapes × 2 dtypes) validating kernel math against a pure PyTorch reference. - `test_silumul_default_multipliers_backward_compat` — confirms no-arg `LigerSiLUMulFunction.apply(a, b)` is bit-identical to explicit `(1.0, 1.0)`. - `test_correctness_falcon_h1_mlp` — 8 cases end-to-end parity for `LigerFalconH1SwiGLUMLP` against a pure-PyTorch Falcon H1 MLP reference. Out of scope: Qwen3-VL `swiglu=True` no-op (#956) — separate monkey-patch wiring issue, not kernel math. ## Testing Done - Hardware Type: NVIDIA A10G (g5.xlarge) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Notes on `make test`: - All 21 new tests pass. - All pre-existing SwiGLU tests that exercise `LigerSiLUMulFunction` pass (Llama, Phi3, Mixtral fp32, functional, falcon_h1 monkey_patch). - Unrelated pre-existing failures on `main` that also show up here (confirmed by running the same tests on `main`): - `test_correctness_mixtralexperts[bf16]` — Triton 3.2 `atomic_add does not support bf16` in `src/liger_kernel/ops/fused_moe.py`. - `test_apply_liger_kernel_to_instance_for_paligemma` — `SiglipVisionModel` attribute change in transformers v5. - `test_apply_liger_kernel_to_instance_for_gemma3_conditional_generation` — `Gemma3Config` v5 API change. - `test_grpo_loss::test_correctness[sequence-luspo-...*-8-128-1024-4096]` (3 large-shape cases) — pre-existing GRPO numerical sensitivity. Convergence tests were not run locally (single-GPU box). Happy to run them if a maintainer can advise on the expected workflow, or to add a Falcon H1 mini-model entry in a follow-up. |
||
|
|
9e3f2618c3 |
[Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) (#1196)
## Summary Adds Liger-Kernel support for Google's **Gemma 4** family (text, dense variants — primary target `google/gemma-4-31B`). Follows the exact integration pattern of the Gemma 3 port: module-level class swaps (`Gemma4RMSNorm`, `Gemma4TextMLP`), class-forward swap on `Gemma4ForCausalLM`, and `MODEL_TYPE_TO_APPLY_LIGER_FN["gemma4_text"]` registration so `AutoLigerKernelForCausalLM` routes transparently. ### What you get - `LigerRMSNormForGemma4` (ones-init, no `(1+w)` offset, fp32 compute, handles `with_scale=False` for `v_norm`) - `LigerGEGLUMLPForGemma4` (thin wrapper over `LigerGEGLUMLP` that absorbs `Gemma4TextMLP`'s `layer_idx` arg) - Fused-linear-CE `causal_forward` on both `Gemma4ForCausalLM` (stock HF) and `Gemma4TextForCausalLM` (some users extract this subclass to dodge HF issue #45200's mm_token_type_ids check) - Registration + exports from `liger_kernel.transformers` ### What you don't get (explicit non-goals) - Multimodal (`Gemma4ForConditionalGeneration`, `Gemma4VisionModel`, `Gemma4AudioModel`) - MoE (`Gemma4TextExperts`/`Gemma4TextRouter`, used by 26B-A4B) - PLE (Per-Layer Embeddings — used by E2B/E4B) - Double-wide MLP on KV-shared layers (31B doesn't use either) - **Rope kernel swap**: HF Gemma 4's `apply_rotary_pos_emb(x, cos, sin)` is single-tensor; `liger_rotary_pos_emb(q, k, cos, sin)` is dual-tensor. Signatures are incompatible. `rope=True` is a no-op with a warning. The large memory win (LCE) is unaffected. ## Headline: Peak HBM at seq_len=8192 Measured on **AMD MI250X (LUMI)** with a 6-layer Gemma-4-shaped mini model at full 31B vocab (262,144), `bf16`, batch=1, seq=8192, forward+backward: | Peak HBM | Value | |---|---| | HF baseline | **41.47 GB** | | Liger (patched) | **10.89 GB** | | **Saved** | **30.58 GB (73.7% reduction)** | Driver: fused-linear-CE (`skip_logits=True`) eliminates the 262,144 × 8192 × bf16 logits tensor (~4 GB materialized) **plus** its gradient, activation buffers, and softcap intermediates — combined ~30 GB freed. Loss parity on the same configuration: `loss_HF = 12.6841440`, `loss_liger = 12.6840858` (abs diff ~6e-5 — numerically identical). ## Numerical correctness ### fp32 (kernel correctness) Whole-model forward, 6-layer Gemma-4-shaped config, vocab 32000, shape `(2, 256, 32000)`: | statistic | value | |---|---| | max \|logits_liger − logits_hf\| | **2.55e-03** | | mean abs diff | 5.11e-05 | | p99 abs diff | 3.84e-04 | ### bf16 (expected dtype noise) | statistic | value | |---|---| | max abs diff | 7.97e-01 | | mean abs diff | 3.05e-02 | | p99 abs diff | 1.92e-01 | Interpretation: kernels are numerically correct (fp32 max 2.55e-3); bf16 drift is 6-layer dtype-inherent precision noise, well inside industry-standard ranges. ## Discovered during LUMI verification (fixed in-PR) Three plan-level assumptions that only showed up under real model construction — all fixed here: 1. **`v_norm` uses `with_scale=False`** — no weight parameter, so the naive subclass broke class-level swap. Fix: `LigerRMSNormForGemma4` accepts `with_scale` kwarg and delegates forward to a plain-torch path when weight is absent. 2. **`Gemma4TextMLP(config, layer_idx)` signature** (vs `LigerGEGLUMLP(config)`) — crashed at layer construction. Fix: `LigerGEGLUMLPForGemma4` wrapper absorbs the extra arg. 3. **`apply_rotary_pos_emb(x, cos, sin)` single-tensor signature** (vs `liger_rotary_pos_emb(q, k, cos, sin)`) — incompatible. Fix: `rope=True` becomes a no-op + warning for Gemma 4. Plus: `Gemma4RMSNorm` inherits `Gemma3nRMSNorm` (not `Gemma3RMSNorm`) — reusing `LigerRMSNormForGemma3` would silently diverge training via its `(1+w)` offset. New subclass is mandatory, not cosmetic. ## Testing Done - Hardware Type: AMD MI250X (LUMI HPC) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence ### CPU-runnable (included in `make test`) \`\`\` $ python -m pytest test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text -v test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text PASSED [100%] ======================== 1 passed, 1 warning in 21.27s ========================= \`\`\` ### checkstyle \`\`\` $ ruff check --output-format=concise . All checks passed! $ ruff format --check . 264 files already formatted \`\`\` ### LUMI (MI250X, ROCm 6.2.4, PyTorch 2.7.1, transformers 5.5.4) - [x] `pytest test/convergence/bf16/test_mini_models.py -k mini_gemma4_text` — MI250X ROCm - [x] `pytest test/convergence/bf16/test_mini_models_with_logits.py -k mini_gemma4_text` — MI250X ROCm - [x] Gemma 3 regression (3 tests) — all pass, no Gemma 3 behaviour change - [x] Revert helper functional test — identical outputs (`max diff 0.00e+00`) after revert - [x] HBM benchmark at seq_len=8192 — **30.58 GB saved**, **73.7% reduction** Tolerance selection for `mini_gemma4_text`: - `loss_atol=5e-2` (vs gemma3's `1e-2`) — 6-layer mini (minimum for sliding+global cycle) has ~0.05 bf16 drift vs gemma3's 4-layer which fits 1e-2 - `logprobs_atol=5e-1` — accumulated bf16 noise flips 3 of ~20,480 top-k logprob ranks on near-tie tokens --- 🤖 Drafted with [Claude Code](https://claude.com/claude-code) (Claude Opus 4.7). No auto-generation skill was used end-to-end — scaffolding and LUMI-verified fixes were developed and reviewed interactively. --------- Co-authored-by: amrtini <[email protected]> Co-authored-by: Claude Opus 4.7 <[email protected]> |
||
|
Kashif Rasul
|
b8f093adb3 |
[GRPO] add chunked grpo streaming over vocab (#1160)
## Summary This PR fixes the chunked GRPO loss to compute only selected-token log-probs by streaming over the vocab dimension. This reduces peak memory for the fused-linear chunked path and preserves the existing high-level fused-linear API. We also fixes the luspo reduction in the chunked path to match TRL exactly, and tightens the `torch.compile` boundary so we only compile the pure loss computation instead of compiling through the closure that calls `torch.autograd.grad`. <img width="1800" height="900" alt="two_way_grpo_time" src="https://github.com/user-attachments/assets/7b696797-d9a7-4b37-83b0-b408a17c8d6d" /> <img width="1800" height="900" alt="two_way_grpo_memory" src="https://github.com/user-attachments/assets/48cbc8fb-ba5e-497e-9f1d-5439b7499350" /> So now the two implementations are correctly implementing the trade-offs of their design. <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
Chunyu
|
e4831e4d64 |
[NPU] Add NPU Fused MoE kernel (#1183)
## Motivation This pr ports `fused_moe.py` and `fused_moe_kernels.py` to an NPU-affine implementation while preserving the original math. The computational definition is unchanged: forward remains `W1 (gate/up) -> SwiGLU -> W2 -> token-weighted gather`, and backward still follows `dA' = dO @ W2^T` to produce `d_pre_act / dS / dW2 / dX / dW1`. The main changes are execution-strategy optimizations for NPU. ## Note: Use the Skill For this fused_moe kernel migration, we followed the skill document from https://github.com/linkedin/Liger-Kernel/issues/1197. ## Testing Done - Hardware Type: Ascend 910B2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence 🤖 Generated with: [cursor](https://cursor.com/). |
||
|
Chenguang Li
|
dcd404b1ea |
[Fix]: guard LigerQwen3_5CausalLMOutputWithPast import (#1169)
LigerQwen3_5CausalLMOutputWithPast is only defined in output_classes.py
when the installed transformers version includes the Qwen3.5 model. In
older transformers versions, the base class import fails silently andthe
Liger subclass is never defined, causing an ImportError when qwen3_5.py
unconditionally imports it at module level.
This breaks CI（Ascend） environments with a transformers version that
does not yet ship Qwen3.5, as test_monkey_patch.py imports qwen3_5.py at
collection time, causing the entire test run to abort with:
```
ImportError: cannot import name 'LigerQwen3_5CausalLMOutputWithPast'
from 'liger_kernel.transformers.model.output_classes'
```
Fix by wrapping the import in a try/except block (falling back to None)
and converting the return type annotation to a string literal to avoid
evaluation at function definition time.
- Hardware Type: Atlas 800I A2
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
|
||
|
lowdy1
|
3978fe2846 |
[NPU] Add optimized NPU mhc (#1173)
Add Ascend NPU Triton kernels for the three mHC sub-operators: - Fused matmul + RMS normalization (forward/backward) - Sinkhorn routing with split pre/post/residual coefficients (forward/backward) - Pre-aggregate weighted sum (forward/backward) - Post + residual mixing (forward/backward) NPU optimizations applied: - Unified UB tiling via compute_default_tiling_strategy for matrix - Persistent grid-stride loops (tl.range + num_programs) - Adaptive BLOCK_N/BLOCK_M for core utilisation at small seq_len - Fused backward coefficient assembly kernel Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
lowdy1andnoemotiovon
|
2ca3bd0d75 |
[Benchmark] Add ModelConfig Sweep Support and Pre-Probe to Remaining Benchmarks (#1195)
## Summary * Add `model_config` sweep support to training and Rope op benchmarks, enabling evaluation across different model architectures at a fixed sequence length * Introduce a pre-sweep probe step to determine safe configurations and prevent OOM during benchmarking benchmark_attn_res.py benchmark_embedding.py benchmark_fused_neighborhood_attention.py (retains all kernel configs for `model_config` sweep instead of sweeping model architectures) benchmark_mhc.py benchmark_mhc_lm.py benchmark_multi_token_attention.py benchmark_relu_squared.py benchmark_softmax.py benchmark_sparsemax.py benchmark_sparse_multi_token_attention.py benchmark_tiled_mlp.py - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: noemotiovon <[email protected]> |
||
|
lowdy1andnoemotiovon
|
ef85f2469e |
[Benchmark] Benchmark refactor for training and RoPE ops (#1193)
## Summary * Add `model_config` sweep support to training and Rope op benchmarks, enabling evaluation across different model architectures at a fixed sequence length * Introduce a pre-sweep probe step to determine safe configurations and prevent OOM during benchmarking Loss kernels: benchmark_cross_entropy.py benchmark_fused_linear_cross_entropy.py benchmark_jsd.py benchmark_fused_linear_jsd.py benchmark_kl_div.py benchmark_tvd.py benchmark_rope.py benchmark_qwen2vl_mrope.py benchmark_llama4_rope.py `BT` denotes the product of batch size and sequence length (`batch_size` × `sequence_length`). Sweep mode `model_config`: `python benchmark_cross_entropy.py --sweep-mode model_config --bt 2048` Sweep mode `BT`: `python benchmark_cross_entropy.py --model llama_3_8b` - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: noemotiovon <[email protected]> |
||
|
lowdy1andnoemotiovon
|
1796cf0d4f |
[Benchmark] Benchmark refactor for loss ops (#1191)
## Summary * Add `model_config` sweep support to loss ops in Liger-Kernel benchmarks, enabling evaluation across different model architectures at a fixed sequence length * Introduce a pre-sweep probe step to determine safe configurations and prevent OOM during benchmarking Loss kernels: benchmark_cpo_loss.py benchmark_distill_cosine_loss.py benchmark_distill_jsd_loss.py benchmark_dpo_loss.py benchmark_grpo_loss.py benchmark_kto_loss.py benchmark_orpo_loss.py benchmark_simpo_loss.py `BT` denotes the product of batch size and sequence length (`batch_size` × `sequence_length`). Sweep mode `model_config`: `python benchmark_grpo_loss.py --sweep-mode model_config --bt 2048` Sweep mode `BT`: `python benchmark_grpo_loss.py --model llama_3_8b` Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: noemotiovon <[email protected]> |
||
|
lowdy1andnoemotiovon
|
eb6e42383f |
[Benchmark] Add model_config sweep for norm benchmarks (#1189)
## Summary - Add Qwen 2.5 models (7B / 14B / 72B) and DeepSeek models (V2 Lite / V3) to MODEL_REGISTRY, including MoE-specific parameters for DeepSeek models - Add `model_config` sweep support to normalization benchmark scripts, enabling evaluation across different model architectures at a fixed sequence length Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: noemotiovon <[email protected]> |
||
|
luca-888
|
6fe468120c |
[fix]Add logits_to_keep and shift_labels support for Qwen3-VL and Qwen3-VL-MoE (#1181)
# Summary This PR adds `logits_to_keep` and `shift_labels` support for both `Qwen3-VL` and `Qwen3-VL-MoE` in the Liger-patched forward path. The change aligns the patched implementation with the expected Hugging Face interface and enables selective logits materialization for long-context inference. # Testing Done - `make test` - not fully green - observed existing failures in `GRPO`, `fused_neighborhood_attention`, and `gemma3` monkey patch tests - `make test-convergence` - not fully green - observed failure in `test/convergence/bf16/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_llama4-...]` - `make checkstyle` - passed Known limitation: - The failed `make test` / `make test-convergence` cases above do not directly exercise the `Qwen3-VL` or `Qwen3-VL-MoE` `logits_to_keep` / `shift_labels` change in this PR |
||
|
UserChen666andmamba-chen
|
fcaae507af |
[NPU]Improvement performence for grpo_loss (#1174)
## Summary This PR implements Triton kernel-level performance optimizations for the grpo_loss operator. ## Details 1. Block Size (BLOCK_N) & Memory Coefficient Tuning Unified default tiling by increasing BLOCK_N from 2048 to 4096 across all kernels, improving per-batch computation granularity and reducing loop counts and kernel launch overhead. Adjusted memory multipliers downward to fit the NPU’s 192KB Unified Buffer (UB) capacity and enhance memory utilization: Softmax: 6.0 → 3.0 Forward: 10.0 → 4.0 Backward: 12.0 → 8.0 2. Computation Instruction Optimization (Reduce Divisions, Improve Instruction Efficiency) Precompute inv_temp = 1.0 / TEMPERATURE to replace multiple in-loop divisions with single multiplications, reducing floating-point latency. Simplified backward gradient expression: dlogp = dlogp * dloss * inv_temp instead of the original chained division, lowering the number of floating-point operations. 3. Loop & Compilation Optimizations Changed inner kernel loops from range to tl.static_range to provide explicit loop-unrolling hints to the compiler, optimizing instruction scheduling and pipeline efficiency. Explicit index type casting: INPUT_IDS cast to int32 to avoid implicit type conversion overhead on NPU. 4. Masking & Memory Access Optimization Unified use of the cols_mask variable to reuse memory access masks, reducing redundant calculations and improving memory access throughput. Simplified gradient calculation logic: (cols_idx - probs) * dlogp instead of tl.where branching, minimizing branch judgment overhead. ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <img width="1108" height="749" alt="image" src="https://github.com/user-attachments/assets/02c2563f-50cd-427b-b2a9-c3b338fb0bc1" /> <img width="2003" height="828" alt="image" src="https://github.com/user-attachments/assets/64fdf419-bbb3-4a04-8c4e-74fb4c3b4145" /> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: Atlas A800I - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: mamba-chen <[email protected]> |
||
|
Chun-Mao (Michael) Lai
|
1faa0131db |
MoE Triton Kernel (#1179)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Triton implementation of a fused MoE expert kernel for training, inspired by SonicMoE: https://arxiv.org/abs/2512.14080 Related to: #958 ## Details **Known limitation:** The current implementation scales well with number of experts (more parallelism) but not with number of tokens. For models with few experts (e.g., Mixtral-8x7B) at large token counts, the eager per-expert loop is likely faster. **Convergence:** the convergence test uses wider tolerances for MoE models patched with this kernel. ## Benchmarking Results - Model: Qwen3-30B-A3B - #Tokens: 8192 | #Experts: 128 | top-K: 8 | hidden: 2048 | intermediate: 768 | Sweep Dim | #Tokens | #Experts | | :--- | --- | --- | | Memory | <img width="1000" height="600" alt="fused_moe_memory_full_token_length" src="https://github.com/user-attachments/assets/5c55e185-ff15-4d1a-9aab-c0c315f94ea0" /> | <img width="1000" height="600" alt="fused_moe_memory_full_token_length" src="https://github.com/user-attachments/assets/f274f84f-8a7e-4c67-900f-1b96ca5b80d0" /> | | Speed (Full) | <img width="1000" height="600" alt="fused_moe_speed_full_token_length" src="https://github.com/user-attachments/assets/4350afdc-2820-4748-bd0c-7d05812ddc60" /> | <img width="1000" height="600" alt="fused_moe_speed_full_token_length" src="https://github.com/user-attachments/assets/d537c312-af6d-407f-9271-962b54284cd1" /> | | Speed (Forward) | <img width="1000" height="600" alt="fused_moe_speed_forward_token_length" src="https://github.com/user-attachments/assets/8256910e-1910-4117-adcf-44cd91e6125a" /> | <img width="1000" height="600" alt="fused_moe_speed_forward_token_length" src="https://github.com/user-attachments/assets/4ece35dd-215d-4c86-a6b6-f92d60c8c22c" /> | | Speed (Backward) | <img width="1000" height="600" alt="fused_moe_speed_backward_token_length" src="https://github.com/user-attachments/assets/2a12f652-9424-48be-bce8-b5e8c1a8289d" /> | <img width="1000" height="600" alt="fused_moe_speed_backward_token_length" src="https://github.com/user-attachments/assets/69e308bb-51e0-4ec4-bb08-e6cf3f87c0ce" /> | ### Benchmarking with patching into Qwen3-30B-A3B - Parameters - Max sequence length: 32768 - Per device batch size: 1 - Gradient accumulate steps: 8 - GPU: 2 * H100_8 - Result - Processing **8.24x** more tokens per second - Training speed up: 8.19x - Slightly save memory: -1% - Eval speed up: 4.1x <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: H100 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
sunyi0505
|
76a0821780 |
optimize group_norm for ASCEND_NPU (#1154)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
Vaibhav JindalandClaude Opus 4.6
|
b7025a7027 |
Improve bwd kernel for fused_add_rms_norm (#1187)
## Summary This has been generated using the `liger-kernel-perf` skill introduced in #1185 . Restructure the backward kernel of `fused_add_rms_norm` to reduce peak register liveness, dropping register count from **115 → 95 per thread** on H100. This eliminates a performance cliff at large hidden sizes and yields up to **70% backward speedup** at H=16384. ## Root Cause Analysis NCU profiling on H100 revealed the backward kernel was severely underutilizing the GPU: | Metric | Forward Kernel | Backward Kernel | |--------|---------------|-----------------| | Registers/Thread | 40 | **115** | | Theoretical Occupancy | 75% | **25%** | | Achieved Occupancy | 67.87% | **12.49%** | | Blocks/SM (register-limited) | 6 | **2** | | Waves/SM | 5.17 | **0.50** | | Memory Throughput (% peak) | 78.24% | 60.88% | | Compute Throughput (% peak) | 27.64% | 19.87% | The H100 has 65,536 registers per SM. With 256 threads/block (8 warps x 32 threads), each thread using 115 registers consumes 115 x 256 = 29,440 registers per block. This means only `floor(65536 / 29440) = 2` blocks can run concurrently per SM, yielding just 12.5% occupancy. With so few active warps, the GPU cannot overlap memory accesses with computation — stalling 77.5% of scheduler slots. ### Why so many registers? The backward kernel processes one row per loop iteration, maintaining several BLOCK_SIZE-wide vectors simultaneously in registers. With BLOCK_SIZE=4096 and 256 threads, each thread handles 16 elements per vector, requiring 16 registers per vector. The original code had up to **~8 vectors live simultaneously** at peak: | Vector | Lifetime | Purpose | |--------|----------|---------| | `W_row` | Persistent (loop-invariant) | Weight + offset | | `dW_row` | Persistent (accumulator) | Weight gradient | | `col_offsets` | Persistent | Column indices | | `mask` | Persistent | Bounds mask | | `dY_row` | Per-iteration | Upstream gradient | | `X_row` | Per-iteration | Saved residual (fp32) | | `m` | Per-iteration | `dY * W` intermediate | | `dX_row` | Per-iteration | Output gradient | | `dS_out_row` | Per-iteration (conditional) | Residual stream gradient | At 16 registers per vector x 8 vectors = 128 registers, plus scalars/pointers/loop vars, this matches the observed 115 registers (compiler optimizes some away). ## Changes ### 1. Reorder: compute dW before dX **Before:** The original code computed `m -> dX -> dW -> store dX`. This meant `dY_row`, `m`, `X_row`, `dX_row`, and `dW_row` were all live simultaneously during the dW accumulation. **After:** By computing `dW` first (`dW -> m -> dX -> store dX`), the compiler can reuse `dY_row`'s registers when computing `m = dY_row * W_row`, since `dY_row` is no longer needed independently after the dW step. This is safe because both `dW` and `m` read `dY_row` — they don't modify it (except in GEMMA mode where the cast happens after dW). ### 2. Factor the dX formula with a precomputed scalar **Before:** ```python dX_row = rstd_row * m dX_row += (rstd_row) * (-(1 / n_cols) * rstd_row * rstd_row * tl.sum(m * X_row, axis=0) * X_row) ``` This expression keeps `m`, `X_row`, and the intermediate `rstd * (... * X_row)` vector all live while computing the compound expression. **After:** ```python dot = tl.sum(m * X_row, axis=0) # scalar reduction c = -(1.0 / n_cols) * rstd_row * rstd_row * rstd_row * dot # scalar dX_row = rstd_row * m + c * X_row # single fused expression ``` By precomputing `c` as a scalar, the final `dX_row` expression is a simple element-wise combination of two vectors (`m` and `X_row`) scaled by scalars. The compiler doesn't need to keep intermediate vector results alive — it can compute each element of `dX_row` independently, allowing better register reuse. ### 3. Defer dS_out load **Before:** `dS_out_row` was loaded and added to `dX_row` inside the compound correction expression, contributing to peak liveness. **After:** `dS_out_row` is loaded only after the `dX_row = rstd * m + c * X` computation is complete. By this point, `m` registers are freed, so `dS_out_row` reuses them rather than competing for register space. ### 4. Add `num_stages=2` for Hopper software pipelining **Rationale:** The current code never sets `num_stages`, so Triton defaults to `num_stages=1` (no software pipelining). On H100 (Hopper architecture), setting `num_stages >= 2` enables the compiler to overlap memory loads for the next iteration with computation for the current iteration using Hopper's async copy engine (TMA). A parameter sweep confirmed `num_stages=2` is safe across all hidden sizes with no regressions. Higher values (3-4) showed no additional benefit for this kernel because the backward kernel's low occupancy limits the pipeline depth that can be effectively utilized. `num_stages=2` was chosen as the conservative optimal: it enables basic pipelining without increasing register pressure (higher stages require more pipeline buffer registers). The `num_stages` value is passed through `ctx` from forward to backward, following the same pattern as `BLOCK_SIZE` and `num_warps`. ### Combined effect After these changes, the peak register liveness drops to ~5-6 vectors: | Vector | Still live? | |--------|------------| | `W_row` | Yes (persistent) | | `dW_row` | Yes (persistent) | | `col_offsets` | Yes (persistent) | | `mask` | Yes (persistent) | | `dY_row` / `m` | **Shared** (compiler reuses registers) | | `X_row` / `dX_row` | **Shared** (compiler reuses after dot product) | | `dS_out_row` | Loaded late, **reuses freed `m` registers** | NCU confirmed registers dropped from **115 → 95 per thread** (17% reduction). While this doesn't cross the 85-register threshold needed for 3 blocks/SM, the reduced pressure gives the compiler more freedom for instruction scheduling and ILP (instruction-level parallelism), which explains the dramatic speedup especially at H=16384 where the original had extreme register spilling. ## Benchmark Results **NVIDIA H100 80GB HBM3, M=2048, dtype=float32** ### Speed | Mode | H | Before (ms) | After (ms) | Improvement | |------|---|------------|-----------|-------------| | Backward | 2048 | 0.124 | 0.096 | **23% faster** | | Backward | 4096 | 0.124 | 0.101 | **18% faster** | | Backward | 8192 | 0.153 | 0.154 | ~same | | Backward | 16384 | 1.006 | 0.300 | **70% faster (3.35x)** | | Backward | 32768 | 1.390 | 1.119 | **20% faster** | | Full | 16384 | 1.186 | 0.482 | **59% faster** | | Full | 32768 | 1.819 | 1.550 | **15% faster** | | Forward | all sizes | - | - | unchanged | ### Memory No change at any size (~30% savings vs HuggingFace maintained). ### Plots **Backward Speed** (before / after / HuggingFace): <!-- Paste backward speed plot here --> **Full Pass Speed**: <!-- Paste full pass speed plot here --> **Forward Speed** (unchanged): <!-- Paste forward speed plot here --> **Memory** (unchanged): <!-- Paste memory plot here --> ## Test plan - [x] Full test suite passed: `python -m pytest test/transformers/test_fused_add_rms_norm.py -xvs` (40/40 tests, all casting modes + dtypes + shapes) - [x] Checkstyle passed (ruff) - [x] No memory regression at any hidden size - [x] No forward pass regression - [x] Benchmarked on NVIDIA H100 80GB HBM3 - [x] NCU profiling confirmed register reduction (115 → 95) --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
Vaibhav JindalandClaude Opus 4.6
|
2bf7117bd9 |
[Skill] Add liger-kernel-perf skill for Triton kernel performance optimization (#1185)
## Summary Adds a new Claude Code skill (`liger-kernel-perf`) that optimizes existing Liger Kernel Triton kernels through a 3-stage pipeline: Profile -> Optimize -> Finalize. The skill automatically diagnoses bottlenecks, generates multiple optimization variants, benchmarks them against the full suite, and creates a PR with the winning variant -- all while maintaining correctness. Tested on `fused_add_rms_norm` (see PR #1187), which achieved up to 70% backward speedup on H100 via register pressure reduction. ## What the skill does - **Stage 1 (Profiler):** Runs baseline benchmarks, detects GPU architecture (Ampere/Hopper/Blackwell), optionally runs NCU profiling, classifies bottleneck (memory-bound vs compute-bound), produces optimization profile with recommended strategy order - **Stage 2 (Optimizer):** Autonomous optimization loop -- tries parameter tuning first (manual sweep, NOT @triton.autotune), then diagnosis-driven techniques. Each variant gets a versioned file + lab notebook tracking hypothesis/changes/results/learnings. Stops on budget exhaustion, diminishing returns, or target met - **Stage 3 (Finalizer):** Applies winner in-place, runs full test suite (hard gate), generates 3-way comparison plots (original vs optimized vs baseline), updates benchmark CSV, creates PR with descriptive body ## Key design decisions - **No @triton.autotune**: Incompatible with Liger's forward-backward ctx coupling pattern and NPU backends. Uses manual parameter sweeps instead. - **Full benchmark suite every iteration**: No lightweight shortcuts -- ensures optimization is good across ALL input sizes, not just cherry-picked ones. - **Balanced guardrails**: Rejects variants that regress non-target metric >5% or regress one pass >10% to improve the other. - **Comment preservation**: All existing comments kept, explanatory comments added for every optimization change. - **Autonomous + interactive modes**: Runs end-to-end when told "just optimize it", or pauses at human checkpoints between stages. ## Files (7 files, ~2,100 lines) | File | Lines | Purpose | |------|-------|---------| | `SKILL.md` | 116 | Orchestration, input parsing, pipeline flow, guardrails | | `profiler.md` | 156 | Stage 1: baseline benchmarks, GPU detection, bottleneck diagnosis | | `optimizer.md` | 395 | Stage 2: optimization loop with accumulated learning | | `finalizer.md` | 417 | Stage 3: apply, test, plot, update CSV, create PR, report | | `optimization-strategies.md` | 794 | Technique catalog: parameter tuning, memory-bound, compute-bound, architecture-specific | | `templates/optimization-profile.md` | 140 | Cross-stage contract between Profiler and Optimizer | | `templates/variant-notes.md` | 70 | Per-variant lab notebook format for learning accumulation | ## Test plan - [x] Skill triggers correctly on "optimize the X kernel" prompts - [x] Tested end-to-end on `fused_add_rms_norm` (PR #1187) - [x] Verified against Claude Code skill best practices (conciseness, progressive disclosure, SKILL.md under 500 lines) --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
lowdy1
|
969d4ab55a |
[NPU] Add fused_linear_cross_entropy operator (#1164)
## Summary To address the UB overflow issue observed in the benchmark, we introduced an operator with an NPU-friendly implementation of fused linear cross entropy. This fused operator relies on several underlying operations (e.g., large matrix multiplication, softmax, and cross entropy), so its current benchmark performance is not yet optimal. Further optimization may be needed. ## Testing Done Device: Atlas A3 `python -m pytest ./test/transformers/test_fused_linear_cross_entropy.py` <img width="3270" height="499" alt="image" src="https://github.com/user-attachments/assets/7f8a63df-f325-43fe-80b9-6268c5f10e29" /> - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Vaibhav JindalandClaude Opus 4.6
|
1c013e286c |
[Feature] Use Liger's Relu_Squared kernel for Nemotron models (#1176)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Use relu_squared in nemotron. This PR is generated using the liger-autopatch skill and tests the changes in #1177 . <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> Class patching and instance patching of relu_squared function. ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: H100 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
Vaibhav JindalandClaude Opus 4.6
|
f16c9f76f4 |
Improve liger-autopatch skill to enable modifying existing monkey-patches (#1177)
## Summary Extends the `liger-autopatch` Claude Code skill to support **modifying existing monkey-patches**, not just adding new models. Previously, the skill only triggered for new model support, so modification tasks (e.g., adding `LigerReLUSquared` to nemotron in #1172) bypassed the skill entirely and missed its guidelines. Changes to 2 skill files: **SKILL.md:** - Broadened skill description to trigger on modification-related keywords (update, fix, change, extend, etc.) - Added **Mode Detection** section that routes to create vs modify pipelines - Added **Modify Pipeline** with 3 stages: Change Impact Analysis → Apply Changes → Validate - Validate stage lists all convergence test commands explicitly (including multimodal) **code-generator.md:** - Added **Mode** section distinguishing create vs modify mode - Added **Modification Checklist** with 7 rules (R1-R7): - R1: Both patching levels (class-level + instance-level) - R2: New parameter with default - R3: Update docstring - R4: Update tests - R5: Check revert function in test/utils.py - R6: Run convergence tests (all 6 files including multimodal) - R7: Update README.md - Added **Common Modification Patterns** for: adding activation kernels, adding norm variants, fixing missing instance patching, updating for upstream HF changes ## Testing Done Tested by asking Claude to add `relu_squared` to nemotron's existing monkey-patch. Verified the skill triggers in modify mode and follows the modification checklist. - Hardware Type: N/A (skill files only, no code changes) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
Chenguang Li
|
35cf6a0d5c |
[NPU]: optimize cross entropy kernel gradient computation (#1178)
- Hoist loop-invariant scalar computations (z_scale, one_minus_ls, z_deriv) out of the inner loop to avoid redundant recalculation - Fuse softmax, z-loss derivative, and smoothing term into a single vector expression in the non-weighted branch - Guard tl.where with block-range check (y >= i and y < i + BLOCK_SIZE) to skip unnecessary vector operations when target index is not in the current block Hardware Type: Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Chenguang LiandTcc0403
|
d991472e1e |
[Benchmark]: Add --sweep-mode and --bt to benchmark CLI. (#1163)
This PR follows [PR1162](https://github.com/linkedin/Liger-Kernel/pull/1116) and implements Phase 2. - benchmark_model_configs: replace hidden-size sweep with compute_model_config_sweep_config / ModelConfigSweepConfig; probe each registry model to pick safe batch_size and seq_len for discrete sweeps. - benchmark_geglu / benchmark_swiglu: support model_config sweep across MODEL_REGISTRY via _resolve_model_config_* helpers. - benchmark_dyt: default path sweeps B*T with fixed model dimensions (compute_seq_len_sweep_config); optional model_config sweep; setup uses cfg hidden_size and input.x as BT. - utils: allow string x / x_values for model name indices; extend types. - benchmarks_visualizer: forward extra_config_filter to plotting. - BENCHMARK_GUIDELINES: document D1/D2 sweep patterns and model_config flow. Hardware Type: Atlas 800I A2 - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Tcc0403 <[email protected]> |
||
|
kirsten-1andClaude Opus 4.6
|
3f85b7206b |
Add Kimi AttentionResiduals (AttnRes) kernelFeature/add attn res kernel (#1161)
## Summary Implements **Attention Residuals (AttnRes)** from Kimi/Moonshot AI ([arxiv.org/abs/2603.15031](https://arxiv.org/abs/2603.15031)). This PR addresses issue #1158. ### What is AttnRes? AttnRes replaces standard residual connections with softmax attention over depth blocks to solve the **PreNorm dilution** problem where deep layer contributions get diluted: ```python V = stack(blocks) # [N, B, T, D] K = RMSNorm(V) # per-block normalize scores = einsum(w, K) # [N, B, T] — w is [D] learned query alpha = softmax(scores, 0) # over block dim h = einsum(alpha, V) # [B, T, D] — weighted sum ``` ### Implementation - Single fused Triton kernel: RMSNorm + dot product + softmax + weighted sum - Efficient memory usage: Scores stored in registers (supports N≤32 blocks) - Complete autograd support: Forward and backward kernels with @ensure_contiguous decorator - Benchmark script: Compares against PyTorch and torch.compile ### Files Added - src/liger_kernel/ops/attn_res.py - Core kernel implementation (318 lines) - benchmark/scripts/benchmark_attn_res.py - Benchmark and correctness tests (246 lines) - Updated src/liger_kernel/ops/__init__.py to export LigerAttnResFunction ### Test Plan - Run correctness tests: python benchmark/scripts/benchmark_attn_res.py --quick - Run full benchmark: python benchmark/scripts/benchmark_attn_res.py - Test on RTX 5090 (Blackwell architecture) - Verify forward pass correctness (fp16, bf16, fp32) - Verify backward pass correctness ### Benchmark Results Tested on: **NVIDIA GeForce RTX 5090**, CUDA 12.8 > Note: Due to resource constraints, testing was only performed on RTX 5090. Additional testing on datacenter GPUs (A100, H100) would be valuable to validate performance across different architectures. Maintainers are welcome to run benchmarks on other hardware configurations. To reproduce: `python benchmark/scripts/benchmark_attn_res.py` #### Forward Pass Performance | Config | PyTorch | torch.compile | Liger AttnRes | Speedup vs PyTorch | Speedup vs compile | |--------|---------|---------------|---------------|--------------------|--------------------| | N=4, D=4096, fp16 | 5.164 ms | 0.691 ms | **0.206 ms** | **25.11x** ✨ | 3.35x | | N=8, D=4096, fp16 | 10.011 ms | 2.716 ms | **0.394 ms** | **25.39x** ✨ | 6.89x | | N=8, D=8192, fp16 | 20.076 ms | 3.183 ms | **0.780 ms** | **25.74x** ✨ | 4.08x | | N=16, D=4096, fp16 | 19.946 ms | 2.996 ms | **1.004 ms** | **19.86x** ✨ | 2.98x | | N=8, D=4096, bf16 | 10.009 ms | 1.596 ms | **0.393 ms** | **25.47x** ✨ | 4.06x | #### Forward + Backward Performance | Config | PyTorch | Liger AttnRes | Speedup | |--------|---------|---------------|---------| | N=4, D=4096, fp16 | 20.525 ms | **1.003 ms** | **20.46x** ✨ | | N=8, D=4096, fp16 | 39.880 ms | **1.747 ms** | **22.83x** ✨ | | N=8, D=8192, fp16 | 79.723 ms | **3.317 ms** | **24.03x** ✨ | | N=16, D=4096, fp16 | 80.480 ms | **3.555 ms** | **22.64x** ✨ | | N=8, D=4096, bf16 | 39.888 ms | **1.744 ms** | **22.88x** ✨ | #### Visualizations **Speed - Full (Forward + Backward):** <img width="1000" alt="attn_res_speed_full" src="https://github.com/user-attachments/assets/a1bb7710-1d37-4824-8d30-79fd9e14d774" /> **Speed - Forward:** <img width="1000" alt="attn_res_speed_forward" src="https://github.com/user-attachments/assets/a5649c68-fbf4-4d8e-9cb6-8a119ca6f8a3" /> **Speed - Backward:** <img width="1000" alt="attn_res_speed_backward" src="https://github.com/user-attachments/assets/78161703-fea0-405f-a2ec-ae5be44f72ae" /> **Memory - Full:** <img width="1000" alt="attn_res_memory_full" src="https://github.com/user-attachments/assets/d902bf36-ff65-4657-a471-8e92524261e0" /> #### Correctness Tests All tests passed with expected numerical precision: - Forward pass: max diff < 2e-3 (fp16/bf16), < 2e-6 (fp32) - Backward pass: max diff < 4e-3 (fp16/bf16), < 2e-6 (fp32) ### Key Insights 1. Exceptional speedup: 20-25x faster than PyTorch's einsum-based implementation 2. Beats torch.compile: 3-7x faster than torch.compile in forward pass 3. Scales well: Performance maintained across different N (number of blocks) and D (hidden dimension) 4. Memory efficient: Single-pass fused kernel minimizes memory traffic ### The dramatic speedup is achieved by: - Fusing RMSNorm + attention + weighted sum into a single kernel - Storing attention scores in registers (no global memory roundtrip) - Optimized memory access patterns for coalesced reads/writes --- Closes #1158 --------- Co-authored-by: Claude Opus 4.6 <[email protected]> |
||
|
Chunyu
|
cc7d6055b4 |
[bugfix] Remove the pos_ids parameter from apply_rotary_pos_emb (#1162)
## Summary After updating transformers to version 5.2.0, the `pos_ids` parameter was removed from `apply_rotary_pos_emb`. For details, see: https://github.com/huggingface/transformers/blob/v5.2.0/src/transformers/models/llama/modeling_llama.py#L146. The above change causes the error reported in issue https://github.com/linkedin/Liger-Kernel/issues/1155. ## Testing Done <img width="1553" height="577" alt="image" src="https://github.com/user-attachments/assets/7f682d15-e13a-4217-8a2d-39231d3385af" /> - Hardware Type: all - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
Vaibhav JindalandClaude Opus 4.6
|
8e6232a08f |
Add Liger Kernel support for Ministral models (#1166)
This PR was generated automatically by Claude Code with minimal human intervention. Claude uses the skill written in #1167 to generate this PR and test it out. Ministral (model_type: "ministral") is a separate model type in transformers v5+ with its own MinistralForCausalLM class. This adds full Liger Kernel optimization support including RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, and FusedLinearCrossEntropy. ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> - Adds full Liger Kernel optimization for Ministral models (model_type: "ministral") in transformers v5+ - Patches: RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy - All 5 tests (instance patching + 4 convergence) validated on H100 GPU <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
Vaibhav JindalandClaude Opus 4.6
|
eac9ce6ede |
Add Liger Kernel support for nemotron models (#1165)
This PR is generated by Claude code using the skill in #1167. ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> - Adds monkey patching support for NVIDIA Nemotron models with apply_liger_kernel_to_nemotron - Supports FusedLinearCrossEntropy (default) and CrossEntropyLoss optimizations - RoPE, MLP, and LayerNorm patching are intentionally excluded due to Nemotron's unique architecture <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> - Instance patching test passes - All 4 convergence tests pass on H100 GPU (bf16/fp32 x FLCE/with_logits) - Lint passes <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
Vaibhav JindalandClaude Opus 4.6
|
f1b7e47f09 |
Claude skill for automatically monkey patching HF transformers models (liger-autopatch) (#1167)
## Summary This PR adds a **Claude Code skill** (`liger-autopatch`) that automates adding Liger Kernel support for new HuggingFace Transformers models. Instead of manually writing 10+ files per model, a contributor can now say *"Add Liger Kernel support for nemotron"* and the skill handles analysis, code generation, testing, and validation through a 3-stage agentic pipeline. ### How it works 1. **Analyze** — A Model Analyzer agent reads the HF `modeling_*.py` source and produces a structured model profile, resolving 12 architectural decisions (norm type, MLP activation, MoE structure, RoPE variant, casting mode, etc.) 2. **Generate** — A Code Generator agent takes the profile and creates/modifies up to 13 files: `lce_forward`, monkey-patch function, `__init__.py` exports, output classes, instance patching tests, convergence tests (bf16 + fp32, FLCE + with_logits + multimodal), revert utilities, and README entry 3. **Validate** — A Validator agent runs the instance patching test, all applicable convergence tests, and `make checkstyle`, with up to 3 retry attempts per step ### Key design decisions - **Progressive disclosure**: SKILL.md is a concise entrypoint (~55 lines); detailed instructions live in separate agent files and templates - **Version compatibility**: Uses `try/except ImportError` for availability guards — CI on transformers 4.52.0 naturally skips newer models, no explicit version strings needed - **Placement enforcement**: Templates explicitly instruct alphabetical ordering for all insertions into existing files (imports, availability checkers, test functions, MINI_MODEL_SETUPS entries) - **Full convergence coverage**: Generates test entries in all 6 convergence test files (bf16/fp32 x FLCE/with_logits/multimodal) as appropriate for the model type ## Testing Done - Tested end-to-end on **nemotron** (dense, RMSNorm + SwiGLU + RoPE) → #1165 - Tested on **ministral** (dense, Mistral-family) → #1166 - Both PRs generated fully working code including monkey-patch, lce_forward, instance patching tests, and convergence tests across bf16/fp32 FLCE and with_logits variants - Iteratively refined the skill based on issues found during testing (alphabetical placement, convergence test coverage, checkstyle integration, version compatibility) --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
Vaibhav JindalandClaude Opus 4.6
|
39a1f454f3 |
[Feature] Add Triton kernel for ReLU Squared activation (#1171)
This code change was generated end-to-end by Claude Code using the `liger-kernel-dev` skill introduced in #1170. ## Summary Adds triton implementation for relu squared activation function. This is used for NVIDIA's nemotron models. (https://github.com/huggingface/transformers/blob/78bdaf0b39c29737b9ca48a274ef4a34bdafd4d1/src/transformers/activations.py#L206) ### Files added/modified | File | Description | |------|-------------| | `src/liger_kernel/ops/relu_squared.py` | **NEW** — Triton forward/backward kernels + `LigerReLUSquaredFunction` autograd Function | | `src/liger_kernel/transformers/relu_squared.py` | **NEW** — `LigerReLUSquared` nn.Module wrapper | | `src/liger_kernel/transformers/functional.py` | **MODIFIED** — added `liger_relu_squared()` functional API | | `src/liger_kernel/ops/__init__.py` | **MODIFIED** — export `LigerReLUSquaredFunction` + forward/backward | | `src/liger_kernel/transformers/__init__.py` | **MODIFIED** — export `LigerReLUSquared` + added to `__all__` | | `test/transformers/test_relu_squared.py` | **NEW** — 18 unit tests (correctness + functional API, float32 + bf16, multiple shapes) | | `benchmark/scripts/benchmark_relu_squared.py` | **NEW** — speed + memory benchmarks | | `benchmark/data/all_benchmark_data.csv` | **MODIFIED** — benchmark results appended | ### Kernel design - **Tier 1 (element-wise)** — simplest kernel tier, one program per row - **Forward**: fused `relu + square` in a single kernel pass - **Backward**: `dX = dY * 2 * relu(X)`, reads saved X without modifying it - **No float32 upcasting needed** — relu and square are numerically stable - Uses `calculate_settings(n_cols)` for BLOCK_SIZE/num_warps auto-tuning ## Benchmark Results (NVIDIA H100 80GB HBM3, bfloat16, M=4096) ### Speed | Hidden Size | Mode | Liger (ms) | PyTorch (ms) | Speedup | |-------------|------|-----------|-------------|---------| | 4096 | Forward | 0.028 | 0.053 | **1.9x** | | 4096 | Backward | 0.039 | 0.126 | **3.2x** | | 4096 | Full | 0.087 | 0.178 | **2.0x** | | 16384 | Forward | 0.094 | 0.186 | **2.0x** | | 16384 | Backward | 0.137 | 0.453 | **3.3x** | | 16384 | Full | 0.229 | 0.635 | **2.8x** | ### Memory | Hidden Size | Liger (MB) | PyTorch (MB) | Savings | |-------------|-----------|-------------|---------| | 4096 | 0.039 | 0.063 | **37.5%** | | 16384 | 0.156 | 0.250 | **37.5%** | ### Speed — Forward <!-- INSERT relu_squared_speed_forward.png --> <img width="1000" height="600" alt="relu_squared_speed_forward" src="https://github.com/user-attachments/assets/747e0ede-8d97-449c-9505-a1ba21f5470d" />. ### Speed — Backward <!-- INSERT relu_squared_speed_backward.png --> <img width="1000" height="600" alt="relu_squared_speed_backward" src="https://github.com/user-attachments/assets/033f3373-f1f5-44f0-b630-393b82bfd93c" /> ### Speed — Full (Forward + Backward) <!-- INSERT relu_squared_speed_full.png --> <img width="1000" height="600" alt="relu_squared_speed_full" src="https://github.com/user-attachments/assets/cb32cc32-6a39-4b2a-a61e-2df92c88414e" /> ### Memory — Full <!-- INSERT relu_squared_memory_full.png --> <img width="1000" height="600" alt="relu_squared_memory_full" src="https://github.com/user-attachments/assets/0128bade-ec3b-419a-9885-f7e37ca8c400" /> ## Testing Done Hardware Type: NVIDIA H100 80GB HBM3 - [x] 18/18 unit tests pass on H100 (float32 + bfloat16, 6 shape configs including non-power-of-2) - [x] `ruff check` and `ruff format` pass - [x] Speed benchmarks run successfully (forward, backward, full) - [x] Memory benchmarks run successfully --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
Vaibhav JindalandClaude Opus 4.6
|
0169c4def2 |
Claude skill for Triton kernel development (liger-kernel-dev) (#1170)
## Summary
This PR adds a **Claude Code skill** (`liger-kernel-dev`) that automates
the full lifecycle of Triton kernel development for Liger Kernel.
Instead of manually writing 8+ files per kernel (ops, transformers
wrapper, functional API, exports, tests, benchmarks), a contributor can
now describe a PyTorch operation and the skill handles analysis, code
generation, and validation through a 3-stage agentic pipeline.
### How it works
1. **Analyze** — An Analyzer agent reads the PyTorch operation (from a
local file, URL, code snippet, natural language description, or model
component reference), writes a standalone PyTorch reference
implementation, and produces a structured kernel profile classifying the
operation into one of 3 complexity tiers
2. **Generate** — A Generator agent takes the profile and
creates/modifies up to 8 files: Triton forward+backward kernels with
`torch.autograd.Function`, `nn.Module` wrapper, functional API,
`__init__.py` exports, unit tests (parametrized over shapes and dtypes),
and benchmark scripts
3. **Validate** — A Validator agent runs `make checkstyle`, unit tests
(hard gate with 3 retries — stops on persistent failure), benchmarks
(speed + memory for fwd/bwd/full), and generates plots. Optionally runs
`ncu` profiling
### Key design decisions
- **Progressive disclosure**: SKILL.md is a concise entrypoint (~63
lines); detailed instructions live in separate agent files and templates
that are loaded only when needed
- **Three complexity tiers**: Templates cover element-wise ops
(SwiGLU-like), reduction ops (RMSNorm-like), and fused/complex ops
(CrossEntropy-like) with tier-specific patterns
- **Correctness is a hard gate**: If tests fail after 3 retries, the
pipeline stops and reports to the user rather than pushing forward with
broken code
- **Single dtype benchmarks**: Benchmarks use `model.dtype` (typically
bfloat16). Multi-dtype coverage is handled by unit tests
- **Supports both create and modify modes**: Detects intent
automatically — creating new kernels goes through the full pipeline,
modifying existing kernels skips analysis
### Skill file structure
```
.claude/skills/liger-kernel-dev/
├── SKILL.md # Main entry point (63 lines)
├── analyzer.md # Stage 1: understand operation, produce profile
├── generator.md # Stage 2: generate all 8 files
├── validator.md # Stage 3: checkstyle → tests → benchmarks → plots
├── kernel-profile-format.md # Kernel profile schema + naming conventions
├── templates/
│ ├── ops-kernel.md # Triton kernel patterns by tier
│ ├── module-wrapper.md # nn.Module wrapper pattern
│ ├── functional-api.md # functional.py modification pattern
│ ├── unit-test.md # Test file pattern with testing rules
│ └── benchmark.md # Benchmark script pattern
└── examples/
├── swiglu-profile.md # Tier 1 (element-wise) reference
├── rms-norm-profile.md # Tier 2 (reduction) reference
└── cross-entropy-profile.md # Tier 3 (fused/complex) reference
```
## Testing Done
- Tested end-to-end on **ReLU Squared** (Tier 1, element-wise activation
used in Nemotron models) → #1171
- Skill generated all 8 files: Triton kernel, module wrapper, functional
API, exports, 18 unit tests, and benchmarks
- All tests pass on H100 (float32 + bfloat16)
- Benchmarks show **1.9-3.3x speedup** and **37.5% memory savings** vs
PyTorch
- Full PR including benchmark plots was generated end-to-end by the
skill
- Iteratively refined the skill based on issues found during testing
(single-dtype benchmarks, template accuracy)
---------
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
|
||
|
hailey-zhandlowdy1
|
07dd9beebd |
[NPU]support fused neighborhood attention for npu (#1034)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> This PR introduces support for Fused Neighborhood Attention (FNA) optimized specifically for NPU architectures. The implementation focuses on memory efficiency and hardware affinity to prevent performance bottlenecks. Key modifications include: Grid Dimension Refactoring: Adjusted the attention grid to a 1D structure. This change optimizes thread block mapping and prevents User Buffer (UB) overflow, ensuring the workload fits within the NPU's local memory constraints. NPU-Affinity Softmax: Refactored the Softmax tiling and grid dimensions to align with NPU compute unit sizes, maximizing throughput and reducing synchronization overhead. ## Details <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> Hardware Type: NPU（Atlas A2) ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> Tested passed with `python benchmark/scripts/benchmark_fused_neighborhood_attention.py ` `pytest -v test/transformers/test_fused_neighborhood_attention.py` - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: lowdy1 <[email protected]> |
||
|
aryan
|
d8d6630e5c |
fix: guard save_for_backward on grad_bias not bias in fused linear CE forward (#1157)
## Summary Fixes a bug in `LigerFusedLinearCrossEntropyFunction.forward` where `ctx.save_for_backward` checks `bias is not None` instead of `grad_bias is not None` before calling `.detach()`. When `input_requires_grad=False`, `fused_linear_cross_entropy_forward` sets `grad_bias = None` regardless of whether `bias` is provided. The stale check then calls `None.detach()` and raises `AttributeError`. Fixes #1156 ## Details One-line fix in `src/liger_kernel/ops/fused_linear_cross_entropy.py`: Before: `grad_bias.detach() if bias is not None else None,` After: `grad_bias.detach() if grad_bias is not None else None,` ## Testing Done - Hardware Type: CPU (logic-only fix, no GPU kernel change) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
lowdy1
|
167741f5f3 |
[NPU] Add NPU optimized sparsemax (#1104)
## Summary This PR adds a usable Sparsemax implementation on NPU that passes all tests and runs all benchmarks without errors. It includes both tiled and non-tiled kernels for forward and backward. ## Testing Done Tested with `python -m pytest ./test/transformers/test_sparsemax.py -v` Atlas 800I A2(32GB) <img width="3267" height="313" alt="image" src="https://github.com/user-attachments/assets/f9ca7eab-d01f-41ba-a50a-0282d6260dca" /> - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Yutong Wu
|
0438fa8dec |
Add patches for Qwen3_5ForConditionalGeneration to support multimodal. (#1150)
## Summary Add support for Qwen3.5 multimodal patches as follow-up for #1123. Implement extra `lce_forward` for `Qwen3_5ForConditionalGeneration` and add corresponding tests. ### Details Modified files: - `model/qwen3_5.py`: add `lce_forward_for_multimodal` for `Qwen3_5ForConditionalGeneration`. - `model/output_classes.py`: add `LigerQwen3_5CausalLMOutputWithPast` as the return type for qwen3_5. - `test/utils.py`: `revert_liger_kernel_to_qwen3_5` to support `conditional_generation` for qwen3_5. - `test_monkey_patch.py`: add `test_apply_liger_kernel_to_instance_for_qwen3_5_for_conditional_generation`. - `bf16/test_mini_models_multimodal.py`: add `MINI_MODEL_SETUPS["mini_qwen3_5"]` to convergence test. - `fp32/test_mini_models_multimodal.py`: skipped ## Testing Done - Hardware Type: NVIDIA A100 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
Chunyu
|
781083b2cd |
[Optimize, NPU] Remove tl.where from _rms_norm_forward/backward_kernel_tiled() (#1153)
## Summary
When the mask has a large shape, `tl.where` is not NPU-friendly in
`triton-ascend`, leading to low kernel performance. When writing
kernels, it's best to use alternative logic instead. This can result in
significant performance improvements.
**Will these changes affect accuracy?** Since the masking operation has
already been applied when loading X_block, it will not affect the
calculation result during tl.sum.
## Testing Done
### Accuracy first
The shapes in `test_rms_norm.py` are too small to trigger the
`_rms_norm_forward_kernel_tiled` kernel, so we need a new configuration.
```python
@pytest.mark.parametrize(
"bs, sl, hd",
[
(2, 2048, 4096),
(2, 2048, 8192),
(2, 2048, 16384),
(2, 2048, 32768),
],
)
@pytest.mark.parametrize(
"dtype, atol, rtol",
[
(torch.float32, 1e-4, 1e-6),
],
)
@pytest.mark.parametrize(
"reference, offset, casting_mode",
[
(LlamaRMSNorm, 0.0, "llama"),
(GemmaRMSNorm, 1.0, "gemma"),
pytest.param(
BaseRMSNorm,
0.0,
"none",
marks=pytest.mark.skipif(device == "npu", reason="Ascend NPU does not support this test"),
),
],
)
@pytest.mark.parametrize(
"in_place",
[
True,
False,
],
)
@pytest.mark.parametrize(
"elementwise_affine",
[
True,
False,
],
)
```
#### Env
<img width="696" height="251" alt="image"
src="https://github.com/user-attachments/assets/c0cd523f-8a78-4205-8f18-6854477b9d0a"
/>
#### Results after code modification
<img width="1178" height="836" alt="image"
src="https://github.com/user-attachments/assets/b520faa0-f209-4ce0-9574-6000bb612f91"
/>
### Benchmark test
The test cases in `benchmark_rms_norm.py` should keep the same shapes as
those in `test_rms_norm.py`.
```python
common_configs = {
...
"x_values": [2**i for i in range(12, 16)],
...
}
```
#### Before Optimization
forward
<img width="1000" height="600" alt="rms_norm_speed_forward"
src="https://github.com/user-attachments/assets/92a701fc-37a9-4298-a285-95174f97ef98"
/>
backward
<img width="1000" height="600" alt="rms_norm_speed_backward"
src="https://github.com/user-attachments/assets/8cfa175d-d6d3-4472-9b9d-7162e6a9df02"
/>
full
<img width="1000" height="600" alt="rms_norm_speed_full"
src="https://github.com/user-attachments/assets/f8248fc6-52a5-4963-a74e-54f7d85173cd"
/>
memory
<img width="1000" height="600" alt="rms_norm_memory_full"
src="https://github.com/user-attachments/assets/7d4b01bb-d425-4e2b-87ef-3d3b1638324b"
/>
[all_benchmark_data_raw.csv](https://github.com/user-attachments/files/26131388/all_benchmark_data_raw.csv)
#### After Optimization
forward
<img width="1000" height="600" alt="rms_norm_speed_forward"
src="https://github.com/user-attachments/assets/89b1a07f-273a-4d07-a6c2-5eb20c7abeb1"
/>
backward
<img width="1000" height="600" alt="rms_norm_speed_backward"
src="https://github.com/user-attachments/assets/4bd741c9-d401-4ade-ac7c-e3b8a241bc32"
/>
full
<img width="1000" height="600" alt="rms_norm_speed_full"
src="https://github.com/user-attachments/assets/c11f7b8c-a3d9-433e-a24e-ebdae343e2e9"
/>
memory
<img width="1000" height="600" alt="rms_norm_memory_full"
src="https://github.com/user-attachments/assets/e2bfedf2-65b2-440d-96fc-4f389cd5ddda"
/>
[all_benchmark_data_optimized.csv](https://github.com/user-attachments/files/26131403/all_benchmark_data_optimized.csv)
- Hardware Type: Atlas 900 A2 PoD
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
|
||
|
TianHao324
|
830ed7cf11 |
[NPU]:Add support for the cross_entropy operator. (#1148)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> To address the ub overflow issue in the benchmark, we have added an operator with a NPU-friendly implementation of cross_entropy . The current performance is 5-6 times higher than the native code of the GPU, and it is only slightly lower than that of Hugging Face. Further research will be conducted in the future. - Memory-Aware Block Size - NPU Core-Aware Grid Sizing <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <img width="923" height="414" alt="image" src="https://github.com/user-attachments/assets/b95b6209-15b4-48a3-a011-31249a2bc3d0" /> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
UserChen666andmamba-chen
|
ab7fe2ad54 |
[NPU] Add support for the grpo_loss (#1146)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <img width="1838" height="696" alt="image" src="https://github.com/user-attachments/assets/f2315bbb-da4e-4147-87e8-f40a705c24bc" /> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: mamba-chen <[email protected]> |
||
|
Chunyu
|
f762d98202 |
[Fix, NPU] Set soc_info for the NPU device (#1149)
## Summary
With the `teb` update, the `soc version` must be explicitly set in
`set_current_compile_soc_info` for `get_soc_spec` to return the correct
`UB_SIZE`.
```
import torch
from tbe.common.platform import get_soc_spec
from tbe.common.platform import set_current_compile_soc_info
ub_size = get_soc_spec("UB_SIZE")
print(f"UB size: {ub_size}")
```
<img width="646" height="32" alt="image"
src="https://github.com/user-attachments/assets/d57e52e5-e468-4e46-a4db-94351364d218"
/> ❌
## Testing Done
### Atlas 900 A2 PoD
#### Script: ub_size_test.py
```python
import torch
from tbe.common.platform import get_soc_spec
from tbe.common.platform import set_current_compile_soc_info
device = getattr(torch, "npu")
soc_info = device.get_device_name(device.current_device())
print("Current SOC info:", soc_info)
ub_size = get_soc_spec("UB_SIZE")
print(f"UB size: {ub_size}")
```
#### Result
<img width="713" height="304" alt="image"
src="https://github.com/user-attachments/assets/56215e32-1c1b-49ed-9e5b-a859e80729d1"
/>
### Atlas 800I A2
Thanks @TianHao324 for verifying on the Atlas 800I A2 — results are
correct there too.
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
|
||
|
TianHao324
|
05f5878201 |
[NPU]:Add support for the fused_linear_jsd operator. (#1151)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> To address the ub overflow issue in the benchmark, we have added an operator with a NPU-friendly implementation of fused_linear_jsd . - Fix the UB overflow issue in the benchmark - Adaptation of the jsd operator for ascend - Friendly operator writing method for adapting to NPU <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <img width="1194" height="485" alt="image" src="https://github.com/user-attachments/assets/c1bcee50-4754-49c1-850d-1e7e2737d51d" /> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Chenguang Li
|
c2f4c02015 |
[Test]: Refactor benchmark_geglu with standardized model configs (#1116)
Introduce `benchmark_model_configs.py` with canonical LLM architecture profiles (`LLAMA_2_7B`, `LLAMA_3_8B`) and two utilities: - **`estimate_kernel_bytes_per_token`**: runtime probe that measures peak GPU memory per token for a given kernel, providing a safe upper bound. - **`compute_benchmark_shape`**: derives batch size and sequence length from available device memory, the model profile, and the measured bytes-per-token, replacing hardcoded shapes. Add **`run_speed_benchmark`** and **`run_memory_benchmark`** helpers to `utils.py` to eliminate repeated `triton.testing` / `_test_memory` blocks across scripts. Also add a `--model` CLI argument to `parse_benchmark_script_args` so callers can select a model profile at runtime. Refactor `benchmark_dyt.py`, `benchmark_geglu.py`, and `benchmark_swiglu.py` to: - Extract a **setup factory function** to avoid duplicating layer/tensor construction between speed and memory benchmarks. - Delegate to `run_speed_benchmark` / `run_memory_benchmark`. - Compute `x_values`, `bsz`, and related fields dynamically via `compute_benchmark_shape` instead of hardcoded constants. - Expand memory benchmark modes from `["full"]` to `["full", "forward", "backward"]` for consistency with speed benchmarks. Hardware Type: NVIDIA A100-SXM4-80GB - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
orangeH25
|
68a74894b1 |
[NPU] Add group norm support on NPU (#1144)
## Summary This PR introduces a functional GroupNorm operator for Ascend NPU. Key improvements: - Fixes the runtime error `grid should be less than 65536!` and `ub overflow` that occurs when the original GPU-oriented liger-kernel GroupNorm implementation is executed on NPU. - Adjusts the kernel launch and tiling strategy to comply with Ascend NPU execution constraints. - Resolves numerical accuracy issues with PyTorch reference outputs. While the current implementation is still slower than the HuggingFace implementation in end-to-end benchmarks, it provides a stable and functional GroupNorm path for Ascend NPU. This PR mainly focuses on correctness and NPU compatibility. Further kernel-level optimizations will be explored in follow-up work. ## Testing Done <img width="1128" height="122" alt="image" src="https://github.com/user-attachments/assets/75f738d7-162e-4711-a4d9-cb7ede0b43c2" /> - Hardware Type: Atlas 800I A2 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Tcc0403
|
db14ea856f |
Set warn_only for unsupported deterministic algorithms (#1145)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Before the fix (current main branch): ``` ❯ python3 -m pytest test/convergence/bf16/test_mini_models.py -k qwen3_5_moe ========================================================= test session starts ========================================================= platform linux -- Python 3.13.1, pytest-9.0.2, pluggy-1.6.0 rootdir: /home/tcc/Liger-Kernel configfile: pyproject.toml plugins: anyio-4.12.1, rerunfailures-16.1, cov-7.0.0, asyncio-1.3.0, xdist-3.8.0 asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function collected 32 items / 31 deselected / 1 selected test/convergence/bf16/test_mini_models.py::test_mini_model[mini_qwen3_5_moe-32-1e-05-dtype27-0.05-0.2-0.1-0.1-0.01-0.01] FAILED [100%] ``` ``` E RuntimeError: _histc_cuda does not have a deterministic implementation, but you set 'torch.use_deterministic_algorithms(True)'. You can turn off determinism just for this operation, or you can use the 'warn_only=True' option, if that's acceptable for your application. You can also file an issue at https://github.com/pytorch/pytorch/issues to help us prioritize adding deterministic support for this operation. .venv/lib/python3.13/site-packages/transformers/integrations/moe.py:369: RuntimeError ``` After the fix (this PR): ``` ❯ python3 -m pytest test/convergence/bf16/test_mini_models.py -k qwen3_5_moe ========================================================= test session starts ========================================================= platform linux -- Python 3.13.1, pytest-9.0.2, pluggy-1.6.0 rootdir: /home/tcc/Liger-Kernel configfile: pyproject.toml plugins: anyio-4.12.1, rerunfailures-16.1, cov-7.0.0, asyncio-1.3.0, xdist-3.8.0 asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function collected 32 items / 31 deselected / 1 selected test/convergence/bf16/test_mini_models.py::test_mini_model[mini_qwen3_5_moe-32-1e-05-dtype27-0.05-0.2-0.1-0.1-0.01-0.01] PASSED [100%] ``` <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> Root cause: `torch.histc` doesn't support deterministic algorithm https://github.com/huggingface/transformers/blob/adc2f16bf1824f7b57c790b4cf3bc48f95ecec69/src/transformers/integrations/moe.py#L373 ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Signed-off-by: Tcc0403 <[email protected]> |
||
|
lowdy1
|
3ebe3b1bdd |
[NPU] FIX Softmax for loop (#1147)
## Summary The previous fix used an incorrect loop step. This change updates it to `num_prog`, since each grid processes one row out of every `num_prog` rows. ## Testing Done `python -m pytest ./test/transformers/test_softmax.py` Device type: Atlas 800I A2 - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Albert Villanova del Moral
|
e6f84b8431 |
Add mm_token_type_ids to qwen3_5_moe lce_forward (#1140)
## Summary This PR fixes the `lce_forward` function for qwen3_5_moe model, adding support for `mm_token_type_ids` optional parameter related to multimodal processing. Follow-up to: - #1120 - #1109 This fixes a ValueError in `model.generate()` with transformers > 5.2.0, after they merged: - https://github.com/huggingface/transformers/pull/43972 See related issue downstream in TRL: - https://github.com/huggingface/trl/issues/5216 - https://github.com/huggingface/trl/issues/5201 <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
lowdy1
|
8e7641864a |
[NPU] Implement a conservative llama4_rope (#1143)
## Summary In recent versions of triton-ascend, the previous implementation of the `llama4_rope` kernel may fail to compile. The issue appears to be related to the use of `tl.interleave()` in the kernel implementation. So this PR introduces a more conservative implementation of the kernel. Additionally, the current `compute_default_tiling_strategy` is not fully accurate for this kernel due to some two-dimensional tensors. As a temporary workaround, the tiling configuration has been adjusted to a larger and more conservative setting to ensure correct compilation and execution. ## Testing Done `python -m pytest ./test/transformers/test_llama4_rope.py -v` `python ./benchmark/scripts/benchmark_llama4_rope.py` Device: Atlas A2 800I - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
ruilin-gif
|
a01a0020b4 |
Add support for Qwen3.5 dense models (#1123)
## Summary Adds Liger Kernel support for the Qwen3.5 dense model family (0.8B / 2B / 4B / 9B / 27B). Closes #1119 Qwen3.5 uses the same hybrid GDN + full-attention architecture as Qwen3 Next but is a distinct model type (`qwen3_5`) in transformers v5.2+. This PR adds a dedicated `apply_liger_kernel_to_qwen3_5` entry point so auto-patching works correctly. Patched kernels: - **RMSNorm** — Gemma-style offset RMSNorm (`LigerRMSNormForQwen3Next`, `offset=1.0, casting_mode="gemma"`) - **SwiGLU** — Fused SwiGLU MLP (`LigerQwen3MoeSwiGLUMLP`, matching Qwen3.5's `__init__` signature) - **Fused Linear Cross Entropy** — avoids materializing the full logits tensor (vocab_size=248,320) - **RoPE** — disabled (raises `NotImplementedError`); only 1 in 4 layers uses RoPE via full attention ## Details **New file:** - `src/liger_kernel/transformers/model/qwen3_5.py` — `lce_forward` for `Qwen3_5ForCausalLM` **Modified files:** - `monkey_patch.py` — `apply_liger_kernel_to_qwen3_5` with class-level and instance-level patching; registered as `qwen3_5` in `MODEL_TYPE_TO_APPLY_LIGER_FN` - `__init__.py` — export - `test/utils.py` — `revert_liger_kernel_to_qwen3_5` - `test_monkey_patch.py` — instance patching test - `bf16/test_mini_models.py` — convergence test - `bf16/test_mini_models_with_logits.py` — convergence with logits - `fp32/test_mini_models.py` — convergence (skipped; `ChunkGatedDeltaRuleFunction` doesn't support float32) - `fp32/test_mini_models_with_logits.py` — convergence with logits (skipped, same reason) - `README.md` — added Qwen3.5 to supported models table **Validation on real models (Qwen3.5-0.8B through 27B):** | Model | Logits cosine sim | Training loss rel diff | |-------|:-----------------:|:----------------------:| | 0.8B | > 0.999 | 0.32% | | 2B | > 0.999 | 0.20% | | 4B | > 0.999 | 0.28% | | 9B | > 0.999 | 0.39% | | 27B | > 0.999 | 0.41% | ## Testing Done - Hardware Type: NVIDIA H200 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence ``` test_apply_liger_kernel_to_instance_for_qwen3_5 PASSED test_mini_model[mini_qwen3_5 bf16] PASSED (6.47s) test_mini_model[mini_qwen3_5 bf16 with_logits] PASSED (7.23s) test_mini_model[mini_qwen3_5 fp32] SKIPPED (ChunkGatedDeltaRuleFunction no fp32) test_mini_model[mini_qwen3_5 fp32 with_logits] SKIPPED (same) ``` |