979 Commits
Author SHA1 Message Date
Yulong CaoandJari Kolehmainen aed0dfbb58 refactor(kl_div): route fused-linear KL inner kernel through the dispatcher (#1476)
## Summary

Follow-up promised during the #1423 review (@kolehma8): expose the
fused-linear KL's inner computation as a dispatchable primitive,
mirroring what `jsd_loss_and_grad` already does for `fused_linear_jsd`.

`LigerFusedLinearKLDivFunction` used to launch `_kl_div_kernel` directly
from its chunk loop, bypassing impl selection — the same class of bug
that the `jsd_loss_and_grad` registration guards against (composed ops
silently bypass the dispatcher; cf. #1228). After this PR:

- `kl_loss_and_grad` is a registered non-autograd primitive in
`ops/backends/_triton/kl_div.py`, wrapping `_kl_div_kernel` (the kernel
stays in `ops/fused_linear_kl_div.py`, mirroring the FLCE backend
layout);
- the composed op dispatches the primitive per chunk and **rebinds the
returned dx**, so a future impl that returns a fresh tensor instead of
writing in place needs no changes to the composed op;
- `functional.py` declares the primitive's op locations.

No user-facing API or numerical behavior changes: same kernel, same
chunk geometry, same GEMMs.

## Why this matters beyond tidiness

- The elementwise `kl_div` op already ships a CuTe DSL backend (#1418).
A CuTe DSL / cuTile `kl_loss_and_grad` primitive is the natural next
step (cuTile registrations are actively landing, e.g. #1456, tracked in
#1205) — and after this PR it only needs a `@register_op` under the same
primitive name; the composed op picks it up via auto-select.
- Small contract detail surfaced while wiring this: the kernel
early-returns for `ignore_index` rows without writing the loss, so the
primitive pre-zeroes its per-row loss buffer (the composed op's old
`loss_1d` used to provide that zeroing implicitly).

## Validation

- New `test/ops/test_fused_linear_kl_div.py` (14 cases): primitive
numerics vs a PyTorch reference (2 shapes × 3 dtypes × with/without
`ignore_index`), `available_backends` sanity, and a tracking-dispatch
test asserting the composed op's default path routes through
`dispatch("kl_loss_and_grad")`.
- Regression (RTX 4060 Laptop):
`test/transformers/test_fused_linear_kl_div.py` (full behavioral net
incl. AMP and `detect_anomaly` paths) + `test/ops/test_kl_div.py` +
`test/ops/test_fused_linear_jsd.py` + `test/backends` → **184 passed, 19
skipped** (env-conditional skips).
- `ruff check` + `ruff format --check` clean on all touched files.

GPU CI workflows need maintainer approval to run on this fork PR (same
as #1423/#1424) — happy to provide additional numbers if useful.

Co-authored-by: Jari Kolehmainen <[email protected]>
2026-09-24 21:43:17 +00:00
Daniel LentonandJari Kolehmainen ec10acbbe6 perf(ppo): skip dW/dB calculation when lm_head parameters are frozen (#1488)
## Summary

`LigerFusedLinearPPOBase`, which backs `LigerFusedLinearGRPOLoss` and
its variants (GRPO, DAPO, Dr. GRPO, BNPO, CISPO, SAPO, ...), always
makes the `lm_head` weight (and bias) differentiable
(`weight.detach().requires_grad_(True)`). As a result,
`_ChunkedSelectiveLogProbFunction.backward` allocates a float32 `[V, H]`
gradient buffer and runs a `grad_logits.T @ hidden` GEMM on every
sequence × vocab chunk, even when the weight does not require grad. With
LoRA/PEFT the `lm_head` is frozen and that gradient is discarded.

This PR differentiates the weight and bias only when they require grad,
as #1441 did for the preference losses.
`_ChunkedSelectiveLogProbFunction` now computes only the gradients
autograd asks for (`ctx.needs_input_grad`). The loss and the input
gradients are unchanged, and trainable parameters take the same path as
before.

## Details

Forward + backward in bf16, `LigerFusedLinearGRPOLoss(beta=0.0,
use_ref_model=False)` with its other defaults (compiled), B = 8, T =
1024, H = 4096. Each figure is the median of 10 timed iterations after 3
warm-ups, then the median over separate processes, with main and this
branch run alternately.

| V | lm_head | main | this PR |
|---|---|---|---|
| 128,256 | frozen | 471.6 ms, 6.14 GiB | **263.7 ms, 3.25 GiB (−44%)**
|
| 151,936 | frozen | 557.4 ms, 7.22 GiB | **311.7 ms, 3.79 GiB (−44%)**
|
| 128,256 | trainable | 472.2 ms | 471.6 ms |
| 151,936 | trainable | 558.8 ms | 558.0 ms |

#1294 also edits `fused_linear_ppo.py`, changing the GEMM operand dtypes
in the same backward. The two changes are independent, but whichever
lands second needs a rebase.

## Testing Done

- Hardware Type: NVIDIA H100 80GB HBM3 (torch 2.14.0+cu130, triton
3.8.0)
- [ ] run `make test` to ensure correctness: I ran the test file for the
changed base rather than the full suite.
`test/chunked_loss/test_grpo_loss.py`: 994 passed and 614 skipped on
this branch (976 passed and 614 skipped on main, plus 18 new).
- [x] run `make checkstyle` to ensure code style (ruff 0.15.22 `check`
and `format --check` on the changed files)
- [ ] run `make test-convergence` to ensure convergence: not run; the
convergence tests do not exercise the chunked losses.

New tests:
- A frozen vs a trainable `lm_head`, with and without bias, with and
without a reference model, and at a vocab size above the 4096 vocab
chunk, so both chunk loops run.
- A trainable weight with a frozen bias.
- A check that `_selective_logprob_backward` is asked for the weight and
bias gradients only when they require grad.

The frozen-head test calls `torch.compiler.reset()` first, as
`test_correctness_large_seq_exercises_chunking` does. Without it,
starting from a cold compile cache on CPU, two of its compiled cases hit
a dynamo guard error ("sources must not be empty for symbol"). The error
reproduces identically on main with the same test, so it is not related
to this change.

Co-authored-by: Jari Kolehmainen <[email protected]>
2026-09-24 21:36:46 +00:00
Rain fbdea0ff3b fix(npu): Replace Ascend A5 comments with 950 (#1484)
## Summary

Follow
https://github.com/linkedin/Liger-Kernel/issues/969#issuecomment-5773453873

Replace the machine-type name `A5` with `950` in Ascend comments and
docstrings, matching the `A2` / `A3` / `950` naming rule.
- Touched files: `ub_manager.py`, `group_norm.py`, and `grpo_loss.py`.
- Kernel logic, dispatch, and numerics are unchanged.

## Testing Done

- Hardware Type: N/A (comment-only)
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
2026-09-24 18:21:47 +00:00
Rain 09f7febf1a fix(npu): Ascend A3/A5 UB tile cap and A5 workarounds for group norm, GRPO, fused MoE (#1464)
## Summary

This PR lands two commits that unblock Ascend **A5** regression coverage
for several transformers ops while **triton-ascend** fixes are pending
upstream.

1. **`ub_manager`**: Hard-cap block tiles using detected UB and the
default **auto-multi-buffer (~4/3)** live-set factor, so margin-only
sizing cannot pick tiles that only fit mis-reported UB (notably some
**A5 GLU** paths). Uses **A3 ≈ 192 KiB / A5 ≈ 256 KiB** from runtime
detection—not SoC-specific branches in kernels.
2. **Kernel workarounds (A5)**: Targeted changes in **group norm**,
**GRPO / CISPO / SAPO**, and **fused MoE** so the same Triton semantics
pass on A5. Several paths are **temporary** until the linked
**triton-ascend** issues are fixed; we intend to revert or simplify once
upstream ships.

## Upstream triton-ascend issues (filed)

These are **compiler/backend bugs on the A5 pipeline**, not application
logic errors. This PR works around them; long-term fix belongs in
[triton-lang/triton-ascend](https://github.com/triton-lang/triton-ascend/issues).

| Issue | Symptom on A5 | Workaround in this PR |
|-------|----------------|------------------------|
| [#2239](https://github.com/triton-lang/triton-ascend/issues/2239) |
BiShengHIR `memref.extract_strided_metadata` when **dynamic `task_base`
+ `tl.sum(..., axis=1)`** on padded 2D tiles (e.g. group norm
`(1,1,1,3)`) | `_group_norm_forward_single_task_kernel`:
**`tl.sum(X_vals)`** / **`tl.sum(X_vals * X_vals)`** (equivalent for
`cpg=1`) |
| [#2246](https://github.com/triton-lang/triton-ascend/issues/2246) |
Compile-time **UB overflow** for large-vocab scans with
**`tl.static_range`** + bf16→fp32 online softmax (`V=151936`) | Inner
vocab loops: **`tl.static_range` → `tl.range(0, N, BLOCK_N)`** in
`grpo_loss.py` (fwd / softmax / bwd) |
| [#2247](https://github.com/triton-lang/triton-ascend/issues/2247) |
**2-D reduction** over `[TPT, K_POW2]` drops lanes when **`K_POW2 < 8`**
(MoE router histogram undercount) | `_moe_router_histogram_kernel`: pad
inner dim to **`HIST_K_POW2 = max(8, K_POW2)`** with unchanged `k < K`
mask |

Additional **fused MoE** A5 fixes in the same commit (gather **`TILE_H`
padding**, backward **one row per program** for `tl.dot`, **H on grid**
for `dX_expanded`, bf16 K-reduce via mask weights) are documented
internally and may get separate upstream issues; they are **not**
covered by #2239/#2246/#2247.

## Details

### `ub_manager` (`8ba7d90`)

- Cap: `max_block_size = min(margin_solve, ub_capacity //
(bits_per_block * 4/3))`.
- Does **not** use `multiplier * 2` (would unnecessarily shrink A5 GLU
tiles that legitimately fit 256 KiB UB).
- On A3, plain CE/LCE fwd vocab tile may step down (e.g. 16384→8192)
where 4/3 straddles a power of two; SwiGLU / RMSNorm / RoPE tiles
unchanged in our A/B checks.

### Group norm (`d30508b`)

- Single-task forward: avoid **`axis=1`** reduce on A5 (see #2239).

### GRPO loss (`d30508b`)

- Replace inner **`tl.static_range`** with **`tl.range`** for vocab
dimension (see #2246).
- Keeps **`multibuffer=True`**; shrinking `BLOCK_N` alone does not fix
A5 compile UB for `static_range`.

### Fused MoE (`d30508b`)

- Histogram: **`HIST_K_POW2`** (see #2247).
- Token gather: **`TILE_H = BLOCK_H + BLOCK_H // 2`**; K-reduction uses
explicit weights where bf16 **`tl.sum(axis=0)`** was unsafe on A5.
- Backward down-proj / dX: row-expanded grid, **`pid_base`**,
**`row_mask = (m_offs == 0) & …`**; dX launch **`(num_rows,
n_h_tiles)`** instead of nested H+I **`tl.range`** in one program.

## Testing done

- **Hardware**: 

Ascend **A3** (192 KiB UB)
<img width="1039" height="115" alt="A3 第二次执行"
src="https://github.com/user-attachments/assets/6a8208ae-ff70-49c4-a245-597e6cc16ec1"
/>

Ascend **A5** (256 KiB UB)
<img width="925" height="109" alt="A5 全量测试用例执行情况"
src="https://github.com/user-attachments/assets/3c9ccb90-7b85-4214-a045-7876c0c3a603"
/>

The difference in the execution results of the 76 test cases between A3
and A5 is due to different `transformers` versions.

- **`make test`** (transformers): full suite green on A5 after above
(run **`test_fused_moe.py` in a subprocess** if combining with other
tests—MoE correctness can poison the device in-process).
- **`make checkstyle`**
- **`ub_manager` A/B** (same machine, old vs new): SwiGLU / RMSNorm /
RoPE tile sizes and latency unchanged where expected; CE flat on
V=32k/128k; FLCE spot check per prior notes on this PR.

## Follow-up

- Drop or narrow workarounds when
[#2239](https://github.com/triton-lang/triton-ascend/issues/2239),
[#2246](https://github.com/triton-lang/triton-ascend/issues/2246), and
[#2247](https://github.com/triton-lang/triton-ascend/issues/2247) are
resolved in a released **triton-ascend** build.
- Track remaining MoE A5 items against future upstream issues as filed.
2026-09-21 13:29:41 +00:00
c5b9f31d3e Fix CUTEDSL FLCE functional wrapper arguments (#1475)
## Summary

- forward `chunk_mem_const` only when the selected FLCE autograd
Function advertises the extended inner-dispatch signature
- preserve the standalone CUTEDSL and Ascend Function argument contracts
- restore the intended SM100 hardware skip for CUTEDSL FLCE
functional-wrapper tests on H200 instead of raising an argument-count
`TypeError`

## Root cause

`chunk_mem_const` was appended unconditionally by
`liger_fused_linear_cross_entropy`. The standalone CUTEDSL Function
accepts the original 15 user arguments, so the wrapper supplied one
extra positional argument before the Function could reach its SM100
capability gate.

## Testing

Run in pod `f45151a224e134137869-n0-0` on H200/SM90:

- `make checkstyle`
- `make test`: 5067 passed, 915 skipped, 3 xfailed
- `make test-cutedsl`: 488 passed, 246 skipped
- `make test-cutile`: 211 passed, 113 skipped
- `make test-cute`: 4 passed, 10 skipped
- `make test-convergence`: 151 passed, 15 skipped, 3 xfailed, 1 xpassed,
8 failed

The eight convergence failures are unrelated CUTEDSL LayerNorm failures
at `src/liger_kernel/ops/backends/_cutedsl/layer_norm.py:615`, where
`view()` rejects non-contiguous multimodal inputs. A targeted rerun
against untouched `origin/main` reproduced all eight failures,
confirming this change introduces no additional regression.

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-18 19:47:57 +00:00
19d92e7309 Publish LCK with PyPI Trusted Publishing (#1474)
## Summary

- publish `liger-cute-kernels` through PyPI Trusted Publishing instead
of long-lived username/password secrets
- bind the publish job to the `liger-cute-kernels` GitHub environment
configured in the pending publisher
- grant `id-token: write` only to the publish job
- use the PyPA-supported `pypa/gh-action-pypi-publish` action
- retain the guarded manual `publish=true` retry added by #1473

## Pending publisher configuration

The PyPI pending publisher must exactly match:

```text
Project: liger-cute-kernels
GitHub owner: linkedin
Repository: Liger-Kernel
Workflow: publish-liger-cute-kernels.yml
Environment: liger-cute-kernels
```

A pending publisher intentionally returns 404 for the project until its
first successful OIDC upload. That first upload creates the project and
associates it with the PyPI account that configured the pending
publisher.

## Publishing 0.8.3 after merge

Run **Build and Publish Liger Cute Kernels** manually on `main` with
`publish=true`. The workflow rebuilds and validates both cores, packages
the combined wheel, obtains a short-lived token through GitHub OIDC, and
uploads it to PyPI.

`PYPI_LCK_USERNAME` and `PYPI_LCK_PASSWORD` are not used by this
workflow and can be removed after the OIDC path succeeds.

## Validation

- workflow YAML parses successfully
- publish job uses environment `liger-cute-kernels`
- `id-token: write` is scoped to the publish job
- no Twine username/password secrets are referenced
- nightly and default manual builds remain non-publishing

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-17 01:25:14 +00:00
3899871085 Use a dedicated PyPI credential for LCK (#1473)
## Summary

- publish `liger-cute-kernels` with dedicated `PYPI_LCK_USERNAME` and
`PYPI_LCK_PASSWORD` repository secrets
- add an explicit `workflow_dispatch` `publish` input, defaulting to
`false`
- allow a manual publish retry after credentials are configured, while
preserving non-publishing nightly and default manual builds
- validate that both credentials exist before invoking Twine
- run Twine in verbose mode so future PyPI authorization errors include
the server response

## Background

The v0.8.3 release built and packaged the dual-architecture wheel
successfully, but public PyPI upload failed with HTTP 403. The existing
credentials can publish `liger-kernel` but are not authorized for the
separate `liger-cute-kernels` project.

Failed release run:
https://github.com/linkedin/Liger-Kernel/actions/runs/35137571890

## Required repository configuration

Create and verify the separate PyPI account, enable 2FA, and create an
**entire-account API token** for its first project upload. Add these
secrets to the main repository:

```text
PYPI_LCK_USERNAME=__token__
PYPI_LCK_PASSWORD=pypi-...
```

Do not include either value in source, PR comments, or workflow inputs.

After the first upload creates `liger-cute-kernels`, `PYPI_LCK_PASSWORD`
can be replaced with a project-scoped token from that account while
retaining the same username secret.

## Publishing 0.8.3 after merge

Run **Build and Publish Liger Cute Kernels** manually on `main` with
`publish=true`. The workflow rebuilds and validates both architecture
cores before the dedicated publish job runs.

## Validation

- workflow YAML parses successfully
- manual `publish` input is boolean and defaults to `false`
- release events still publish automatically
- main pushes and default manual dispatches still skip publication
- both dedicated credentials are checked before upload

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-16 22:36:26 +00:00
bad6c3ec42 Keep Triton FLCE C=1 by default and use C=8 for GPT-OSS (#1472)
## Summary

- keep the public Triton fused linear cross entropy chunk-memory default
at `C=1`
- add an explicit per-call `chunk_mem_const` override while preserving
backward compatibility for existing autograd callers
- use `C=8` only in the GPT-OSS patched training path, where C=1 causes
a deterministic BF16 convergence mismatch
- keep every other model and direct/public FLCE caller at C=1
- retain the reusable-logits, direct-dX, and direct-dW optimizations
from #1468
- lock the default and override chunk geometry in regression coverage

## Why the override is scoped

The exact mini GPT-OSS convergence case fails with C=1 and passes with
C=4, C=8, and C=16. However, setting C=8 globally causes the BF16
Qwen3.5-MoE convergence case to fail. The model-scoped override
preserves both trajectories:

- public/default FLCE: C=1
- GPT-OSS patched loss: C=8
- fused linear KL: unchanged independently tuned policy

## H200 memory and latency

BF16, `BT=8192`, `H=4096`, `V=128256`:

| Policy | Chunk size | Chunks | Time | Incremental peak allocation |
|---|---:|---:|---:|---:|
| Public default C=1 | 256 | 32 | 53.039 ms | 1128.66 MiB |
| GPT-OSS override C=8 | 2048 | 4 | 37.748 ms | 1568.04 MiB |
| Previous global C=16 | 4096 | 2 | 35.074 ms | 2068.04 MiB |

The default therefore retains the minimum memory budget, while GPT-OSS
uses the smallest requested release override with stable convergence.

## Testing done

On H200 with PyTorch 2.9.1.6+cu129 and Triton 3.6.0:

- **144 passed**: FLCE ops suite
- **144 passed**: FLCE ops suite with `LIGER_KERNEL_STRICT=1`
- **141 passed**: transformer-facing FLCE suite
- **71 passed, 3 dependency-gated skips**: FP32 + BF16 text-model
convergence suites
- exact BF16 mini GPT-OSS convergence passes with the scoped C=8
override
- BF16 mini Qwen3.5-MoE convergence passes because its default remains
C=1
- Ruff lint/format and `git diff --check` pass

The broader six-file convergence attempt also exposed four PaliGemma
multimodal failures in the ad hoc Transformers 5.15 environment. Those
configurations explicitly set `fused_linear_cross_entropy=False`, so
they are unrelated to this change.

No owned B300 pod was available for an additional Blackwell run.

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
v0.8.3
2026-09-16 18:36:55 +00:00
d5df610cbf Bump Liger packages to 0.8.3 (#1471)
## Summary

- bump the shared public version from `0.8.2` to `0.8.3`
- update both `liger-kernel` and `liger-cute-kernels`, since the native
package resolves its version from the root `pyproject.toml`
- add `liger-kernel[lck]` as an optional install path
- pin the extra to `liger-cute-kernels==0.8.3` so the Python and native
distributions cannot drift across releases
- document stable and editable `lck` extra installation

## Resulting installs

```bash
pip install liger-kernel
pip install "liger-kernel[lck]"
```

The second command installs the normal Liger dependencies plus the
exactly matching standalone native wheel.

## Validation

- root wheel: `liger_kernel-0.8.3-py3-none-any.whl`
- root wheel metadata includes `Provides-Extra: lck`
- root wheel metadata includes `liger-cute-kernels==0.8.3; extra ==
"lck"`
- standalone LCK package metadata resolves to `0.8.3`
- targeted LCK version tests pass
- Ruff check and format checks pass for `setup.py`

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-16 16:17:49 +00:00
cb6462b60d Refresh LigerCute README figures and results (#1469)
## Summary

- replace the outdated Hopper-only benchmark figures and result claims
in `liger_cute_kernels/README.md`
- add the current H200 and B300 GPU-scaling and token-scaling figures
from the Liger Communication Kernels manuscript
- add the current forward/backward pipeline diagrams and Qwen3/Qwen3.5
end-to-end training comparison
- update the design description and benchmark table to match the current
generated figures and training statistics
- retain the existing fused scaled linear cross-entropy build and
runtime documentation without carrying forward its obsolete benchmark
claims

## Figure updates

Removed:

- `hopper_gpu_scaling.png`
- `hopper_token_scaling.png`
- `hopper_training.png`

Added:

- H200 and B300 MoE GPU-scaling composites
- H200 and B300 token-scaling composites
- forward and backward LigerMoE pipeline diagrams
- H200 Qwen3 / two-host B200 Qwen3.5 training comparison

## Validation

- all seven README image references resolve to committed assets
- no legacy figure filenames or obsolete result claims remain
- generated PNG assets total approximately 1.4 MiB
- `git diff --check` passes

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-16 00:03:17 +00:00
b53fc76341 Fix remaining LCK container build dependencies (#1467)
## Summary

- replace GitHub toolcache Python with Ubuntu 22.04's native Python 3.10
inside the CUDA build containers
- install `libibverbs-dev` so the PyPI NVSHMEM IBGDA headers can include
`infiniband/mlx5dv.h`
- retain GitHub-hosted Python 3.12 and pip caching for the non-container
wheel packaging job

## Failures addressed

PR #1466 fixed the first `setup-python` cache-hook failure, but the
subsequent pre-merge workflow runs exposed two additional container
dependencies:

1. GitHub's Python 3.12 toolcache binary requires glibc 2.38, while the
digest-pinned CUDA Ubuntu 22.04 image provides glibc 2.35.
2. Both native architectures require `infiniband/mlx5dv.h` through
NVSHMEM's IBGDA headers; Ubuntu supplies it in `libibverbs-dev`.

The native core is Python-ABI-independent, and `apache-tvm-ffi==0.1.14`
provides a CPython 3.10 manylinux wheel.

## Validation

Pre-merge fork dispatch:
https://github.com/kolehma8/Liger-Kernel/actions/runs/35032379309

- dependency installation and CUDA 12.9 configuration pass for both
architectures
- SM100f core build and artifact upload pass in 6m34s
- SM90a all-config core/cubin build is still running
- the manual dispatch does not execute the Twine publication job

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-16 00:01:59 +00:00
Vaibhav JindalandCopilot de77856a01 perf(flce): reuse logits buffers and write input gradients directly (#1468)
## Summary

Reduce peak allocated memory in fused linear cross entropy by reusing
chunk-logits storage and writing input gradients directly into their
destination.

On B300 with BF16/FP16, the reference 8,192-token workload uses **33.3%
less incremental peak memory**; the 65,536-token workload uses **46.5%
less**. Latency is effectively unchanged. This is a buffer-management
optimization, not a new kernel or chunk-size policy.

## Details

- Allocate one reusable logits buffer on eligible NVIDIA CUDA paths, and
use `torch.mm(..., out=...)` to fill each chunk.
- Write dX directly into its preallocated slice instead of allocating a
GEMM result and copying it.
- Release old logits, gradient/transpose views, CE argument tuples, and
token-scaling temporaries once consumed. Release reusable workspace
before the final gradient-dtype conversion.
- Preserve the original allocating matmul and bias-promotion behavior
when direct output would change dtype or autocast semantics.

The direct-output path requires no-grad execution, matching
FP16/BF16/FP32 operands, and disabled or dtype-matching CUDA autocast.
Logits reuse additionally requires absent/same-dtype bias; direct dX
requires a matching destination dtype. Other cases retain the original
matmul behavior. The older autocast-query API remains supported.

## B300 measurements

Full forward + backward, mean reduction, no bias, default accumulation,
H=4096 and V=128256. Memory is incremental peak allocated MiB above the
resident input/parameter baseline, not reserved or total process memory.

| Dtype | Tokens | Before ms | After ms | Before peak MiB | After peak
MiB | Memory reduction |
|---|---:|---:|---:|---:|---:|---:|
| BF16 | 8,192 | 19.193 | 19.067 | 3,102.04 | 2,068.04 | 33.3% |
| BF16 | 16,384 | 38.301 | 38.420 | 5,202.08 | 3,134.08 | 39.8% |
| BF16 | 65,536 | 146.720 | 148.472 | 17,802.31 | 9,530.31 | 46.5% |
| FP16 | 8,192 | 16.120 | 16.164 | 3,102.04 | 2,068.04 | 33.3% |
| FP16 | 16,384 | 31.688 | 31.556 | 5,202.08 | 3,134.08 | 39.8% |
| FP16 | 65,536 | 124.956 | 123.641 | 17,802.31 | 9,530.31 | 46.5% |

Across 23 cases covering BF16, FP16, FP32, FP32 accumulation, AMP, and
optional features, observed latency changes ranged from **2.2% faster to
1.4% slower**. No material speedup is claimed. The measured losses and
gradients matched the baseline with zero numerical difference in these
cases.

Methodology:

- NVIDIA B300 SM103; PyTorch 2.11.0.1+cu130; Triton 3.6.0.
- Baseline: FLCE source at `f2ac965f619aa4aa090c34252f87349a968b7597`,
unchanged by the two subsequent main commits used as this PR's base.
- Baseline and candidate loaded in the same process; identical seeded
tensors, chunk sizes, precision, and upstream gradient 1.0.
- Fresh forward/backward graph per iteration; both leaf gradients
cleared. No optimizer step.
- Compilation excluded; eager CUDA-event `triton.testing.do_bench`, 100
ms warmup, 400 ms repetition budget, four alternating-order rounds;
median of round medians.
- Memory: median of three separate warmed peak-allocation probes.

The improvement removes approximately one overlapping logits slab plus a
temporary dX chunk. It does **not** make memory independent of token
count: the unchanged C=16 policy still grows the chunk size with total
tokens.

## Hopper and other GPUs

There are no new SM100/SM103-specific instructions or kernels. The
optimization uses ordinary PyTorch/cuBLAS output-buffer operations, so
the same allocation-lifetime benefit is expected on eligible Hopper and
other NVIDIA CUDA workloads.

**Only B300 has been measured.** H100/H200/Ampere latency, exact peak
allocations, and real older-PyTorch behavior still need hardware
coverage; identical performance percentages are not claimed. ROCm and
non-CUDA devices do not enter the new direct-output path. The HIP
eligibility simulation below is not a ROCm hardware run.

## Testing Done

- **322 passed:** transformer FLCE suite, ops FLCE suite, and CuTe
shape-fallback contracts.
- **170 passed:** ops FLCE suite with `LIGER_KERNEL_STRICT=1` (overlaps
the preceding coverage).
- **58 passed / 19 NVIDIA-only skips:** targeted tests with only the HIP
eligibility flag simulated on B300.
- Regression coverage includes actual aten dispatch, storage lifetime
through views, ragged/noncontiguous chunks, AMP and bias promotion, FP32
accumulation, token scaling, forward-only execution, and deferred
`reduction="none"` backward.
- Ruff lint, format checks, and `git diff --check` pass.
- Full-repository `make test` and convergence suites were not run.

Only the FLCE implementation and its existing regression-test file are
included. Local dependency pins, benchmark artifacts, and unrelated
files are excluded.

Co-authored-by: Copilot <[email protected]>
2026-09-15 23:50:07 +00:00
dc1512630d Fix Python setup in LCK container builds (#1466)
## Summary

- disable `actions/setup-python`'s pip cache hook in the CUDA container
matrix jobs
- retain Python 3.12 setup and all dependency installation through
`python -m pip`
- keep pip caching enabled for the non-container wheel-packaging job

## Root cause

Post-merge run
[35030952199](https://github.com/linkedin/Liger-Kernel/actions/runs/35030952199)
failed before compilation in both architecture jobs. `setup-python`
installed CPython successfully, then its cache hook tried to execute
`/__t/Python/3.12.14/x64/bin/pip` inside the job container and failed
with `ENOENT`.

The cache is only an optimization and is incompatible with this
container/toolcache path layout. Removing `cache: pip` avoids the
failing hook; subsequent steps already use `python -m pip`.

## Validation

- workflow YAML parses with the container cache disabled and the host
packaging cache retained
- fork workflow run:
https://github.com/kolehma8/Liger-Kernel/actions/runs/35031510256

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-15 22:34:54 +00:00
d72addc575 Publish multi-architecture Liger Cute wheel on releases (#1465)
## Summary

- add an independent release workflow for `liger-cute-kernels` using the
pinned CUDA 12.9.1 Ubuntu 22.04 build image and CUTLASS v4.4.1
- build SM90a and SM100f cores in parallel with four compiler jobs, then
package both into one Python-ABI-independent wheel
- use the same public version as `liger-kernel` and publish only on
release events; manual dispatches build artifacts without publishing
- stop bundling NVSHMEM shared libraries and depend on
`nvidia-nvshmem-cu12==3.6.5`
- select the SM90a or SM100f core at runtime from the active GPU
capability while retaining legacy single-core wheel compatibility
- validate FSLCE forward/backward exports, wheel contents, version,
dependency metadata, manylinux tag, and the 100 MiB PyPI limit

## Testing done

- `ruff check` and `ruff format --check` on the changed Python files
- 17 targeted packaging/loader tests passed
- workflow YAML parsed successfully and `git diff --check` passed
- built both native cores on H200 pods with `LIGER_CUTE_BUILD_JOBS=4`:
  - SM90a: 45m09s, peak compiler RSS 8.69 GiB
  - SM100f: 2m42s, peak compiler RSS 7.42 GiB
- confirmed both cores export fused linear scaled cross-entropy forward
and backward APIs
- packaged one dual-architecture wheel: 87,726,555 bytes (83.66 MiB),
below PyPI's 100 MiB limit
- installed the wheel in an isolated target directory with
`nvidia-nvshmem-cu12==3.6.5` and verified on H200:
  - selected `libliger_cute_kernels_sm90a.so`
  - no NVSHMEM libraries were bundled in the LCK package
- NVSHMEM unique-ID initialization reported rank 0 / world size 1 and
finalized successfully
- `twine check` passed without warnings

The H200 measurement pod uses glibc 2.38, so its validation artifact was
tagged `py3-none-manylinux_2_38_x86_64`. The release workflow builds
inside the digest-pinned Ubuntu 22.04 image, asserts glibc 2.35, and
emits `py3-none-manylinux_2_35_x86_64`.

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-15 22:25:19 +00:00
f2ac965f61 Add SM100 tensor-parallel fused scaled cross entropy (#1458)
## Summary

Adds a native Blackwell-family (`sm_100f`, including B200/B300)
tensor-parallel fused scaled linear cross-entropy implementation.

- Fuses the classifier GEMM with NLL/LSE and optional entropy.
- Adds a persistent SM100 backward that executes dZ, dX, dW, token-wave
scheduling, node-local NVLS, and two-host reduction in one cluster-2
launch.
- Supports TP1/2/4/8 node-local execution and two-host hierarchical TP
with 1/2/4/8 selected GPUs per host, including TP subgroups of a larger
NVSHMEM world.
- Uses the shared forward-verified QP put-signal ring for inter-host
communication.
- Distributes the remote FP32 dX merge over warp 1 from every resident
CTA, synchronized with an HBM atomic completion counter.
- Adds reusable workspace configuration/query APIs, public API
documentation, correctness validators, phase benchmarks, and topology
contract tests.

## Code entry points

| Layer | Forward | Backward |
|---|---|---|
| PyTorch autograd |
[`LigerFusedLinearScaledCrossEntropyNativeTPFunction.forward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/src/liger_kernel/ops/cute/fused_linear_scaled_cross_entropy_tp.py#L139)
|
[`LigerFusedLinearScaledCrossEntropyNativeTPFunction.backward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/src/liger_kernel/ops/cute/fused_linear_scaled_cross_entropy_tp.py#L186)
|
| TVM FFI |
[`fused_linear_scaled_cross_entropy_forward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp#L535)
|
[`fused_linear_scaled_cross_entropy_backward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp#L610)
|
| SM100 host launcher |
[`fused_linear_scaled_cross_entropy_forward_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward_sm100.cu#L703)
|
[`fused_linear_scaled_cross_entropy_backward_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward_sm100.cu#L445)
|
| Device kernel |
[`forward_gemm_tp_kernel_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_kernel_sm100.cuh#L26)
|
[`backward_gemm_tp_kernel_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm100.cuh#L1109)
|
| Role implementations |
[`forward_gemm_roles_sm100.cuh`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_roles_sm100.cuh)
|
[`backward_gemm_mainloop_sm100.cuh`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm100.cuh)
|

## Multi-host subgroup topology

Multi-host setup is relative to the configured TP team rather than the
full NVSHMEM world. `nvshmem_team_split_2d` derives:

- one local NVLS row containing only the selected GPUs on the current
host;
- one matching-rank remote column containing the corresponding selected
GPU from each host.

For a 16-rank world on two eight-GPU hosts, TP group `{0,4,8,12}`
becomes local rows `{0,4}` / `{8,12}` and remote pairs `{0,8}` /
`{4,12}`. SM100 remote kernels use a parent-team stream barrier followed
by a cooperative clustered CUDA launch, so disjoint TP groups do not
require world-synchronous kernel entry.

Validated concurrently as four TP4 groups `{0,4,8,12}`, `{1,5,9,13}`,
`{2,6,10,14}`, and `{3,7,11,15}`:

- 28/28 numerical cases pass across the four groups.
- Covers entropy, ignored targets, ragged vocabulary, full waves, and
two-wave tails.
- Cross-rank dX equality remains exact.
- Deliberately staggering group entry by 500 ms also passes for forward
and backward, confirming subgroup-scoped progress.
- Two-host local widths 1, 2, 4, and 8 GPUs per host pass.
- `NVSHMEM_IB_ENABLE_IBGDA=1` was forced for the final campaign. Runtime
initialization reported `Successfully initialized the transport: IBGDA.
It will be used for device-side APIs over IB`, created DCI/RC QPs, and
selected each configured `mlx5_*` device.

Random BF16, `M=4096`, `H=4096`, global `V=131072`, 10 warmups, and 30
measured iterations:

| TP group root | Native forward ms | Native forward TF/GPU | Forward
speedup | Native F+B ms | Native F+B TF/GPU | F+B speedup |
|---:|---:|---:|---:|---:|---:|---:|
| 0 | 1.080 | 1017.9 | 2.82x | 5.193 | 846.9 | 1.69x |
| 1 | 1.072 | 1025.8 | 2.85x | 5.166 | 851.3 | 1.69x |
| 2 | 1.058 | 1039.5 | 2.86x | 5.128 | 857.6 | 1.69x |
| 3 | 1.090 | 1009.2 | 2.78x | 5.175 | 849.9 | 1.68x |

Native versus fallback forward maximum absolute difference was
`5.45e-4`.

### SM90 regression validation

Validated the shared topology, workspace, and binding changes on two
eight-GPU
H200 hosts with a clean `sm_90a` build:

- Complete native core build passed.
- 12/12 shared reduction and parent-relative topology tests passed.
- 8/8 SM90 FSLCE configuration and online-softmax contract tests passed.
- 12/12 standalone SM90 forward GEMM numerical cases passed.
- Public FFI world-team and process-group subteam tests passed,
including
  forward, dX, dW, entropy, and CUDA graph replay.
- A two-host TP16 world team passed forward, dX, dW, and entropy
validation.
- Four concurrent strided TP4 groups `{0,4,8,12}`, `{1,5,9,13}`,
  `{2,6,10,14}`, and `{3,7,11,15}` all passed on H200.
- NVSHMEM reported IBGDA as the active device-side InfiniBand transport
on all
  16 ranks.

## Implementation

### Forward

- 2SM UMMA, `M256 x N256 x K64`, cluster `2 x 1`.
- Five-stage TMA mainloop.
- Four source slots and a default wave width of 64 N256 vocabulary
tiles.
- Warp 0 performs local reduction, warp 1 runs the matching-rank host
ring, warp 2 owns TMA, warp 3 issues UMMA, and warps 4-11 drain
TMEM/epilogue work.
- The `__global__` entry point is isolated in
`forward_gemm_kernel_sm100.cuh` and dispatches those roles explicitly;
templated implementations live in `forward_gemm_roles_sm100.cuh`.

### Backward

- One persistent 384-thread cluster-2 kernel.
- dZ: five-stage `M256 x N256 x K64`.
- dX: four-stage 2SM N512 schedule, publishing two M128xN256
communication tiles.
- dW: six-stage paired-N schedule with TMA store for wave 0 and TMA
reduce-add for later waves.
- A single 512-column `Allocator2Sm` TMEM allocation is reused across
all phases.
- TP16:
  1. every CTA's warp 0 performs node-local NVLS reduce-scatter;
  2. CTA 0 warp 1 performs the existing inter-host ring transfer;
3. warp 1 from all resident CTAs merges a global-strided portion of the
remote FP32 shard;
  4. one HBM atomic arrival per CTA releases CTA 0;
  5. warp 0 performs node-local all-gather and BF16 scatter.

The measured TP16 remote stage is 0.285-0.328 ms:

| Component | Latency |
|---|---:|
| IBRC transfer, 8 MiB/GPU | 0.119-0.156 ms |
| All-CTA FP32 merge | 0.160-0.176 ms |
| Total remote stage | 0.285-0.328 ms |
| In-kernel dW overlap window | ~2.41 ms |
| Exposed remote tail | 0 ms |

## Correctness

- TP1/2/4/8/16: 35/35 distributed numerical cases pass.
- Covers aligned and ragged shapes, entropy, ignored targets, 4K waves,
4097-token boundaries, and 5000-token two-wave tails.
- Cross-rank dX equality is exact.
- Maximum observed relative error:
  - dX: 0.262%
  - dW: 0.498%
- Forward native versus Verl maximum absolute difference: `1.09e-3`.
- C++ tests:
  - shared reduction/topology: 11/11
  - SM100 backward contracts: 12/12
- `make checkstyle` passes.

## Performance methodology

- GPU: NVIDIA B300.
- BF16 inputs, FP32 accumulation.
- Hidden size: 4096.
- Fixed global vocabulary: 131072.
- Tokens: 4096, 8192, 16384.
- TP: 1, 2, 4, 8, 16.
- 10 warmups, 30 measured iterations, median CUDA-event latency.
- Verl baseline: repository 512-token chunked TP fallback.
- Forward+backward throughput counts four classifier GEMMs: forward
logits,
  backward logits recomputation for dZ, dX, and dW.
- Efficiency is per-GPU TFLOP/s relative to TP1 for the same shape.
- TP16 uses two eight-GPU hosts and explicit one-HCA-per-local-GPU
mapping.

### M=4096 forward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 2.852 | 1542.0 | 100.0% | 8.489 | 518.1 | 100.0% | 2.98x |
| 2 | 1.457 | 1508.9 | 97.9% | 4.653 | 472.6 | 91.2% | 3.19x |
| 4 | 0.764 | 1438.3 | 93.3% | 2.856 | 384.9 | 74.3% | 3.74x |
| 8 | 0.428 | 1283.6 | 83.2% | 2.134 | 257.6 | 49.7% | 4.98x |
| 16 | 0.465 | 591.0 | 38.3% | 2.096 | 131.1 | 25.3% | 4.51x |

### M=4096 forward + backward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 10.950 | 1606.5 | 100.0% | 29.898 | 588.4 | 100.0% | 2.73x |
| 2 | 6.170 | 1425.6 | 88.7% | 15.570 | 564.9 | 96.0% | 2.52x |
| 4 | 3.562 | 1234.7 | 76.9% | 8.451 | 520.4 | 88.5% | 2.37x |
| 8 | 2.253 | 976.0 | 60.8% | 5.426 | 405.2 | 68.9% | 2.41x |
| 16 | 2.171 | 506.5 | 31.5% | 4.472 | 245.9 | 41.8% | 2.06x |

### M=8192 forward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 5.744 | 1531.4 | 100.0% | 17.107 | 514.2 | 100.0% | 2.98x |
| 2 | 2.867 | 1533.8 | 100.2% | 9.376 | 469.1 | 91.2% | 3.27x |
| 4 | 1.469 | 1496.9 | 97.7% | 5.638 | 390.1 | 75.9% | 3.84x |
| 8 | 0.793 | 1385.8 | 90.5% | 4.185 | 262.7 | 51.1% | 5.28x |
| 16 | 0.762 | 721.6 | 47.1% | 4.077 | 134.8 | 26.2% | 5.35x |

### M=8192 forward + backward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 22.264 | 1580.3 | 100.0% | 58.375 | 602.8 | 100.0% | 2.62x |
| 2 | 11.600 | 1516.5 | 96.0% | 31.618 | 556.4 | 92.3% | 2.73x |
| 4 | 6.463 | 1360.9 | 86.1% | 17.348 | 507.1 | 84.1% | 2.68x |
| 8 | 4.219 | 1042.5 | 66.0% | 10.653 | 412.8 | 68.5% | 2.52x |
| 16 | 4.020 | 546.9 | 34.6% | 8.659 | 254.0 | 42.1% | 2.15x |

### M=16384 forward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 12.645 | 1391.2 | 100.0% | 35.020 | 502.4 | 100.0% | 2.77x |
| 2 | 6.399 | 1374.5 | 98.8% | 19.560 | 449.7 | 89.5% | 3.06x |
| 4 | 3.206 | 1371.9 | 98.6% | 11.426 | 384.9 | 76.6% | 3.56x |
| 8 | 1.674 | 1313.3 | 94.4% | 8.197 | 268.3 | 53.4% | 4.90x |
| 16 | 1.449 | 759.0 | 54.6% | 9.809 | 112.1 | 22.3% | 6.77x |

### M=16384 forward + backward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 50.387 | 1396.5 | 100.0% | 115.214 | 610.8 | 100.0% | 2.29x |
| 2 | 23.270 | 1512.0 | 108.3% | 62.505 | 562.9 | 92.2% | 2.69x |
| 4 | 12.579 | 1398.5 | 100.1% | 35.498 | 495.6 | 81.1% | 2.82x |
| 8 | 8.281 | 1062.1 | 76.1% | 21.188 | 415.2 | 68.0% | 2.56x |
| 16 | 7.898 | 556.8 | 39.9% | 16.706 | 263.2 | 43.1% | 2.12x |

## Memory

Inputs, weights, and final outputs are excluded. Native reports its
exact reusable C++/NVSHMEM forward+backward workspace. Verl reports
incremental PyTorch peak memory.

### M=4096

| TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB
|
|---:|---:|---:|---:|
| 1 | 1810.40 | 1024.10 | 2720.11 |
| 2 | 1202.42 | 512.10 | 1376.11 |
| 4 | 898.46 | 256.10 | 704.11 |
| 8 | 746.53 | 128.10 | 368.11 |
| 16 | 682.67 | 64.10 | 200.11 |

### M=8192

| TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB
|
|---:|---:|---:|---:|
| 1 | 2004.79 | 1024.19 | 2752.21 |
| 2 | 1300.81 | 512.19 | 1408.21 |
| 4 | 948.84 | 256.19 | 736.21 |
| 8 | 772.91 | 128.19 | 400.21 |
| 16 | 709.06 | 64.19 | 232.21 |

### M=16384

| TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB
|
|---:|---:|---:|---:|
| 1 | 2393.56 | 1024.35 | 2816.40 |
| 2 | 1497.58 | 512.35 | 1472.40 |
| 4 | 1049.61 | 256.35 | 800.40 |
| 8 | 825.68 | 128.35 | 464.40 |
| 16 | 761.83 | 64.35 | 296.40 |

## Additional weak-scaling production result

With fixed `V_local=65536`, `M=4096`, and `H=4096`:

| TP | Native forward TF/GPU | Verl forward TF/GPU | Native F+B TF/GPU |
Verl F+B TF/GPU |
|---:|---:|---:|---:|---:|
| 1 | 1524.5 | 488.6 | 1607.1 | 570.3 |
| 2 | 1505.3 | 471.8 | 1424.4 | 560.7 |
| 4 | 1513.3 | 459.2 | 1541.9 | 543.2 |
| 8 | 1490.8 | 446.8 | 1580.4 | 537.9 |
| 16 | 1323.0 | 412.2 | 1434.5 | 533.7 |

## Model-shape campaign

The following production model shapes use `M=8192`, random BF16 inputs,
10 warmups, and 30 measured iterations. TP16 spans the same two
eight-GPU hosts as the fixed-shape campaign.

### Qwen3-30B-A3B (`H=2048`, `V=151936`)

#### Forward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 3.692 | 1380.8 | 100.0% | 15.885 | 320.9 | 100.0% | 4.30x |
| 2 | 1.937 | 1315.9 | 95.3% | 9.192 | 277.3 | 86.4% | 4.75x |
| 4 | 1.040 | 1225.4 | 88.7% | 5.581 | 228.4 | 71.2% | 5.37x |
| 8 | 0.592 | 1076.2 | 77.9% | 4.136 | 154.1 | 48.0% | 6.99x |
| 16 | 0.664 | 479.8 | 34.7% | 4.012 | 79.4 | 24.7% | 6.04x |

#### Forward + backward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 13.692 | 1489.3 | 100.0% | 47.689 | 427.6 | 100.0% | 3.48x |
| 2 | 7.480 | 1363.2 | 91.5% | 26.363 | 386.8 | 90.4% | 3.52x |
| 4 | 4.100 | 1243.3 | 83.5% | 14.834 | 343.7 | 80.4% | 3.62x |
| 8 | 2.507 | 1016.5 | 68.3% | 9.720 | 262.3 | 61.3% | 3.88x |
| 16 | 2.545 | 500.8 | 33.6% | 8.346 | 152.7 | 35.7% | 3.28x |

#### Memory

| TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB
|
|---:|---:|---:|---:|
| 1 | 1975.79 | 1187.19 | 1962.08 |
| 2 | 1286.31 | 593.69 | 996.65 |
| 4 | 941.84 | 296.94 | 515.33 |
| 8 | 769.41 | 152.19 | 276.94 |
| 16 | 695.56 | 74.37 | 154.58 |

### Llama 4 (`H=5120`, `V=202048`)

#### Forward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 11.727 | 1445.3 | 100.0% | 28.506 | 594.6 | 100.0% | 2.43x |
| 2 | 5.795 | 1462.3 | 101.2% | 15.228 | 556.5 | 93.6% | 2.63x |
| 4 | 2.880 | 1471.2 | 101.8% | 8.489 | 499.1 | 83.9% | 2.95x |
| 8 | 1.479 | 1433.0 | 99.1% | 5.339 | 396.8 | 66.7% | 3.61x |
| 16 | 1.103 | 960.3 | 66.4% | 7.578 | 139.8 | 23.5% | 6.87x |

#### Forward + backward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 50.180 | 1351.1 | 100.0% | 100.330 | 675.7 | 100.0% | 2.00x |
| 2 | 22.677 | 1494.8 | 110.6% | 54.966 | 616.7 | 91.3% | 2.42x |
| 4 | 11.785 | 1438.3 | 106.4% | 29.663 | 571.5 | 84.6% | 2.52x |
| 8 | 6.939 | 1221.3 | 90.4% | 16.230 | 522.1 | 77.3% | 2.34x |
| 16 | 6.026 | 703.2 | 52.0% | 23.633 | 179.3 | 26.5% | 3.92x |

#### Memory

| TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB
|
|---:|---:|---:|---:|
| 1 | 2655.29 | 1578.69 | 5015.46 |
| 2 | 1626.31 | 792.19 | 2547.99 |
| 4 | 1111.84 | 394.81 | 1315.52 |
| 8 | 854.41 | 200.19 | 698.15 |
| 16 | 756.06 | 98.84 | 389.87 |

### DeepSeek-V3 (`H=7168`, `V=129280`)

#### Forward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 10.923 | 1390.0 | 100.0% | 22.352 | 679.2 | 100.0% | 2.05x |
| 2 | 5.317 | 1427.6 | 102.7% | 11.923 | 636.7 | 93.7% | 2.24x |
| 4 | 2.638 | 1439.0 | 103.5% | 6.943 | 546.7 | 80.5% | 2.63x |
| 8 | 1.356 | 1399.8 | 100.7% | 4.556 | 416.5 | 61.3% | 3.36x |
| 16 | 1.041 | 911.7 | 65.6% | 4.862 | 195.2 | 28.7% | 4.67x |

#### Forward + backward

| TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU |
Efficiency | Speedup |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 1 | 43.588 | 1393.3 | 100.0% | 81.751 | 742.9 | 100.0% | 1.88x |
| 2 | 22.167 | 1369.9 | 98.3% | 44.494 | 682.4 | 91.9% | 2.01x |
| 4 | 12.411 | 1223.3 | 87.8% | 24.582 | 617.6 | 83.1% | 1.98x |
| 8 | 7.600 | 998.8 | 71.7% | 13.475 | 563.3 | 75.8% | 1.77x |
| 16 | 7.024 | 540.4 | 38.8% | 10.843 | 350.0 | 47.1% | 1.54x |

#### Memory

| TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB
|
|---:|---:|---:|---:|
| 1 | 2278.79 | 1010.19 | 4279.46 |
| 2 | 1437.81 | 505.19 | 2196.71 |
| 4 | 1017.34 | 256.19 | 1156.21 |
| 8 | 807.41 | 128.19 | 634.08 |
| 16 | 744.56 | 64.19 | 373.04 |

## Validation commands

```bash
make checkstyle
```

On B300 (`-DLIGER_CUTE_CUDA_ARCH=100f`):

```bash
cmake --build build --target \
  libliger_cute_kernels.so \
  test_fslce_shared_reduction \
  test_fused_scaled_linear_cross_entropy_backward_sm100

./build/tests/cpp/test_fslce_shared_reduction
./build/tests/cpp/test_fused_scaled_linear_cross_entropy_backward_sm100
```

Distributed numerical validation:

```bash
python -m torch.distributed.run --standalone --nproc_per_node=<1|2|4|8> \
  liger_cute_kernels/tests/python/validate_fused_scaled_linear_cross_entropy_backward_sm100.py
```

TP16 uses two eight-GPU hosts with:

```bash
export NVSHMEM_REMOTE_TRANSPORT=ibrc
export NVSHMEM_IB_ENABLE_IBGDA=1
export NVSHMEM_ENABLE_NIC_PE_MAPPING=1
unset NVSHMEM_HCA_LIST
export NVSHMEM_HCA_PE_MAPPING='mlx5_0:1:1,mlx5_2:1:1,mlx5_3:1:1,mlx5_4:1:1,mlx5_5:1:1,mlx5_6:1:1,mlx5_8:1:1,mlx5_9:1:1'
```

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-14 22:18:34 +00:00
5786fd4c76 Specialize SM90 MoE communication and fix backward synchronization (#1463)
## Summary

- Specialize SM90 MoE communication for local-only and inter-host
execution so local kernels do not carry unnecessary IBGDA code/register
pressure, while retaining IBGDA for inter-host communication.
- Fix the SM90 backward hangs: publish the split MLP3 loop bounds and
perform the producer/consumer rendezvous before the blocking producer
loop; advance shared cursors even when an empty expert range or K slice
takes an early `continue`.
- Add the external non-RDC SM90 module build, packaging, dispatch and
compatibility checks, including the all-configuration module and
1/2/4/8/16-PE coverage.
- Extend TMA/routing coverage and ensure a rejected or incompletely
initialized external module is unloaded through scoped ownership.

Communication semaphore counts and acquire/release semantics are
preserved. The existing SM100-family GPU implementation is unchanged.

## MoEBenchmark results: H200 and B300

**Protocol:** CUDA 12.9, BF16 routed-expert MoE MLP, eight microbatches,
250 ms per-timed-launch cooldown. The summary below uses **T=8192 per
rank** and reports throughput in **TFLOP/s per GPU**. Sixteen GPUs means
two hosts with eight GPUs each. Shared experts and attention are not
included in the benchmark FLOP count.

| GPU | GPUs | Llama 4 Scout forward | Llama 4 Scout backward |
Qwen3-30B-A3B forward | Qwen3-30B-A3B backward |
|---|---:|---:|---:|---:|---:|
| H200 | 1 | 511.36 | 299.54 | 354.38 | 199.28 |
| H200 | 2 | 461.73 | 319.30 | 291.08 | 222.21 |
| H200 | 4 | 473.09 | 337.02 | 255.68 | 217.06 |
| H200 | 8 | 461.04 | 340.12 | 232.64 | 209.07 |
| H200 | 16 | 318.45 | 278.54 | 50.64 | 74.71 |
| B300 | 1 | 977.74 | 723.75 | 599.60 | 376.85 |
| B300 | 2 | 939.54 | 835.58 | 547.99 | 396.22 |
| B300 | 4 | 1004.79 | 955.76 | 517.30 | 406.17 |
| B300 | 8 | 886.71 | 714.04 | 467.82 | 402.50 |
| B300 | 16 | 696.53 | 647.42 | 172.13 | 215.10 |

### Representative backend comparisons

Each cell below is **forward / backward TFLOP/s per GPU** for
Qwen3-30B-A3B at T=8192.

| GPU | GPUs | TE | DeepEP + TE | Comet | FlashMoE | Liger |
|---|---:|---:|---:|---:|---:|---:|
| H200 | 8 | 111.37 / 137.38 | 100.17 / 158.89 | 156.37 / N/A | 102.37 /
N/A | **232.64 / 209.07** |
| H200 | 16 | 12.70 / 30.63 | N/A | N/A | 27.44 / N/A | **50.64 /
74.71** |
| B300 | 8 | 192.26 / 182.09 | 94.53 / 178.85 | N/A | N/A | **467.82 /
402.50** |
| B300 | 16 | 71.46 / 126.43 | N/A | N/A | N/A | **172.13 / 215.10** |

This is not a claim of a uniform win across models. For example, at
eight GPUs, Scout backward is 340.12 versus TE's 443.46 TFLOP/s on H200,
and 714.04 versus 805.15 on B300. At sixteen B300 GPUs, Scout backward
is close to TE: 647.42 versus 653.07 TFLOP/s.

### Coverage and limitations

- Seven model families: Llama 4 Scout, Mixtral 8x7B/8x22B,
Qwen3-30B-A3B/235B-A22B, and Qwen3.5-35B-A3B/122B-A10B.
- H200 named-model campaign: T=8192, 1/2/4/8/16 GPUs, 238 attempts: 228
numeric results and ten expected expert-divisibility skips.
- B300 named-model campaign: T=1024/2048/4096/8192/16384, 1/2/4/8/16
GPUs, 910 attempts: 870 numeric results and forty expected
expert-divisibility skips.
- Mixtral's eight experts cannot be evenly distributed over sixteen
expert-parallel ranks.
- DeepEP is included at 2/4/8 GPUs; its sixteen-GPU transport path
failed in the preserved mainline baseline and is excluded. Comet is
forward-only and single-host in this setup. FlashMoE is forward-only.
- B300 Comet/Flux is unavailable in the pinned source's architecture
support. B300 FlashMoE compilation fails against the pinned CUDA
12.9/CCCL interfaces. These cells are N/A, not substituted measurements.
- The final sixteen-GPU Liger results use IBGDA with GPU-local HCA
affinity: four usable compute NICs per H200 host and eight per B300
host. Local runs use NVLink; the NIC correction does not change their
protocol.
- Kernel-managed RDMA QP counters do not cover IBGDA's direct QPs.
Per-rank HCA assignment and transport initialization are retained in
provenance, but **per-rail traffic volume and aggregate network
bandwidth were not measured**.

The named-model campaigns above are separate from the completed fixed
eight-GPU sweeps below.

### Additional eight-GPU sweeps

- **H200 token sweep:** seven named-model families at
T=1024/2048/4096/8192/16384. The final dataset contains 224 new attempts
plus
56 accepted T=8192 rows from the immutable named-model campaign: 280
rows,
  ten LaTeX tables, 63 PNG/PDF/EPS plots, and zero missing cases.
- **H200 parameter grid:** T=8192, D=I in {2048,4096},
E in {8,16,32,64,128}, and K in {1,2,4,8} with K<E: 38 configurations
and
  304/304 finite-positive attempts. Forward includes TE, DeepEP, Comet,
FlashMoE and Liger; backward includes TE, DeepEP and Liger. Liger is
fastest
in 38/38 forward rows with a 1.388x geometric-mean speedup over the
fastest
  comparison backend. It is fastest in 29/38 backward rows with a 1.103x
  geometric-mean speedup over the fastest comparison backend.
- **B300 parameter grid:** the same 38 configurations, with TE, DeepEP
and
Liger in both directions: 228/228 finite-positive attempts. Liger is
fastest
in 38/38 forward rows and 38/38 backward rows, with geometric-mean
speedups
of 1.880x forward and 1.406x backward over the fastest comparison
backend.
Both generated tables contain only the three deployed B300 backends;
Comet
and FlashMoE are omitted for the availability reasons documented above.

### SM100 regression comparison

A separate, protocol-matched historical comparison uses **one microbatch
and 250 ms cooldown** on both mainline and candidate. Across 330 numeric
Liger comparisons, GPU-count/direction geometric-mean throughput ratios
range from approximately **0.9955 to 1.0090**. The largest individual
decrease is **4.81%** for eight-GPU Qwen3-30B-A3B at T8192; no
individual comparison is more than 5% lower. These are single-run
comparisons, not confidence intervals or a claim that every measurement
improved.

The candidate's per-host benchmark/core hashes match the SM100 artifacts
used for the final campaign. Both sides of the historical sixteen-GPU
comparison used the same older four-NIC mapping; those historical values
are **not** mixed into the corrected-affinity, eight-microbatch table
above. The comparison shows no broad SM100 throughput loss, while
keeping protocol changes and run-to-run variation explicit.

## Testing Done

Hardware: NVIDIA H200 (SM90) and NVIDIA B300 SXM6 (SM100-family,
`sm_103a`).

- Complete SM90 native core and all-configuration external cubin built
with CUDA 12.9.
- Local and actual two-host forward/backward reference runs completed,
including the formerly hanging backward path.
- Named-model campaigns, the H200 token sweep, and the H200/B300 fixed
  eight-GPU parameter grids completed with the coverage above.
- `test_cute_build.py` and `test_export_nonrdc_nvshmem_ptx.py`: thirteen
tests passed,
including duplicate-symbol, missing-SETMAXNREG, and identical-path
rejection coverage.
- `test_moe_bindings.py`: six tests passed and the explicit
NVSHMEM-runtime placeholder
remained skipped; changed topology is rejected before either schedule is
rewritten.
- `test_moe_fwd_bwd_cuda_graph`: dense and sparse routing passed with
TMA enabled and
with the IB-capable getmem fallback forced (four cases total, two H200
GPUs).
- Ruff lint and format checks passed for all seven Python files in this
change.
- The review follow-up all-configuration SM90 wheel compiled with the
existing CUDA
12.9 Release flags in 44m16s. It removes the dead same-host forward copy
path,
validates cached topology before schedule updates, and fingerprints the
selected
  specialization menu.

Repository-wide checklist:
- [ ] `make test`: not rerun; targeted native/CUDA/reference coverage is
listed above.
- [ ] `make checkstyle`: targeted non-mutating Ruff checks were used
instead of the repository-wide auto-fixing target.
- [ ] `make test-convergence`: not run; this change concerns the native
MoE implementation, not the Transformers integration.

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-14 16:16:11 +00:00
sunyi0505 bcb5962dc0 [NPU] Fix Ascend JSD UT failures and register LayerNorm dispatch (#1452)
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Unit tests for JSD / FusedLinearJSD on Ascend failed, mainly due to
signature misalignment with the NVIDIA frontend, and invalid kernel
launches when all samples are ignored. The ascend-triton dispatcher
registration for LayerNorm is also added so that the NPU can
automatically select the UB-aware implementation.

JSD / FusedLinearJSD signature alignment: NVIDIA's LigerJSDFunction /
LigerFusedLinearJSDFunction accepts jsd_impl / jsd_mode (apply(...,
None, None) is used in tests). The Ascend kernel is self-contained and
does not use the inner dispatcher, so it simply accepts and ignores
these two parameters, and the number of backward return values is
supplemented accordingly; supports_inner_impl_dispatch is set to False
for JSD.

Short-circuit for all-ignored case: When n_non_ignore == 0, skip kernel
launch and directly return zero loss and zero gradients. This covers
cases such as test_correctness_all_ignored to prevent division-by-zero /
invalid grid.

LayerNorm registration: Add ops/backends/_ascend/layer_norm.py and
register it in declare_op_locations("layer_norm") within functional.py.
Otherwise, auto-selection only detects the NVIDIA-triton implementation
for CUDA and cannot use _ascend.ops.layer_norm.
<!---
## 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. --->
pytest tests/transformers
<img width="1255" height="383" alt="aa9edfdb912726cb88d8905e1b40f870"
src="https://github.com/user-attachments/assets/db596c84-8a94-480c-82a9-2a98351b4510"
/>

<!-- 
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 800T A3(X86)
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
2026-09-14 12:57:46 +00:00
Yulong Cao f6385e909b feat: add FusedLinearKLDivLoss (fused linear + KL divergence for distillation) (#1423)
## Summary

Add `FusedLinearKLDivLoss`, which fuses the vocabulary-head linear
projection with
KL(Q || P) loss against an explicit target distribution, for knowledge
distillation
workflows where teacher probabilities are available directly (e.g.
cached / offline
soft labels), complementing `LigerKLDIVLoss` (element-wise, no fusion)
and
`LigerFusedLinearJSD` (teacher side needs its own projection).

The kernel follows the chunked gradient-in-forward structure of
`fused_linear_jsd.py`,
so the full `BT x V` logits tensor is never materialized, and adopts the
`_CHUNK_MEM_CONST = 16` token-chunk geometry of
`fused_linear_cross_entropy.py`.

## Details

- Supports `reduction` (`"batchmean"` / `"mean"` / `"sum"`, same
semantics as
`torch.nn.KLDivLoss`), `ignore_index` via `shift_labels`, softmax
`temperature`,
  and an `eps` clamp so that `0 * log(0)` contributes 0.
- Supports `accum_dtype` for fp32 weight-gradient accumulation across
chunked
  low-precision GEMMs (same semantics as `LigerFusedLinearJSD` /
`LigerFusedLinearCrossEntropyLoss`), including the
`torch.addmm(out_dtype=fp32)`
  fast path on torch >= 2.8 / sm80+ and a fallback elsewhere.
- The log-softmax backward is computed in place on the saved log-probs
and the
`chunk_size x V` fp32 temporaries are freed eagerly, so the peak during
the
  trailing GEMMs stays low.
- `reduction="none"` is intentionally unsupported (gradient-in-forward
requires a
  scalar loss); a `ValueError` points users to `LigerKLDIVLoss` instead.
- bf16 + `reduction="sum"` test tolerances follow the existing
convention in
  `test_fused_linear_cross_entropy.py`.
- Peak-memory improvement over the unfused PyTorch baseline (full
fwd+bwd,
`max_memory_allocated`, bf16, llama-3 head shape H=4096 / V=128256, RTX
4060
Laptop GPU): BT=2048 6537 -> 4310 MB (1.52x), BT=4096 12072 -> 6594 MB
(1.83x);
at BT=1024 the benchmark harness reports 4793 -> 4244 MB. Runtime is at
parity
  with the baseline at BT=1024 (0.89x).

## Testing Done

- Hardware Type: NVIDIA GeForce RTX 4060 Laptop GPU (sm89); NVIDIA Tesla
T4 (sm75)

- [x] run `make test` to ensure correctness (32 pre-existing failures
unrelated to
this change: 18 `test_mlp.py` shared-memory OOR on sm89 and 14
`test_grpo_loss.py`
  LUSPO numerics; both reproduce on `main` without this change)
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence (does not
exercise the new
  standalone loss)

<details>
<summary>Test logs (RTX 4060 Laptop, torch 2.11.0+cu128, triton
3.7.1)</summary>

```
make test: 32 failed, 4152 passed, 914 skipped, 15 xfailed in 1050.50s (0:17:30)
  test/transformers/test_fused_linear_kl_div.py: 172 passed, 0 failed, 0 skipped
  failures (all pre-existing on this GPU, reproduced on main without this change):
    18 x test/transformers/test_mlp.py
        triton.runtime.errors.OutOfResources: shared memory
        Required: 131072, Hardware limit: 101376 (sm89 has 100 KB usable smem)
    14 x test/chunked_loss/test_grpo_loss.py
        chunked GRPO LUSPO numerics (also seen by other contributors, see #1397)
make checkstyle: ruff check and ruff format --check pass on all files in this PR
```

</details>

<details>
<summary>Test logs (Kaggle Tesla T4, torch 2.10.0+cu128, triton
3.6.0)</summary>

```
test/transformers/test_fused_linear_kl_div.py: 51 passed, 25 skipped in 10.29s
(bf16 cases skipped by design: T4 is sm75 without native bf16 support)
```

</details>
2026-09-14 02:34:33 +00:00
xvxuopop c64878355b Add DTensor support for Ascend SwiGLU (#1448)
## Summary

Add DTensor support to the Ascend SwiGLU implementation.

SwiGLU now runs on local DTensor shards and preserves the original
device
mesh and placements in forward and backward. Also fix
`infer_comm_backend()` to return `hccl` on Ascend NPU.

## Testing Done

Environment: 4 x Ascend 910 NPU, CANN 9.1.0, PyTorch/torch-npu 2.9.0.

- Non-DTensor SwiGLU tests: `67 passed, 18 skipped`
- Four-device FP32/BF16 SwiGLU DTensor tests: passed
- Four-device RMSNorm DTensor test for HCCL regression: passed
- `make checkstyle`: passed

- Hardware Type: 4 x Ascend 910 NPU
- [ ] run `make test` to ensure correctness (targeted tests run instead)
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence (not applicable
to this DTensor change)
2026-09-13 16:47:03 +00:00
Vaibhav JindalandCopilot 95b01e9402 perf(flce): accumulate low-precision Triton dW with addmm (#1454)
## Summary

Optimize Liger's Triton fused linear cross entropy (FLCE)
weight-gradient accumulation for BF16/FP16 parameters on NVIDIA SM80+
GPUs.

The previous default-accumulator path computes a parameter-sized matrix,
converts it to FP32, then adds it into the BF16/FP16 running gradient
buffer for **every token chunk**:

```python
grad_weight += torch.mm(grad_logits_chunk.t(), input_chunk).float()
```

This PR instead accumulates directly into the existing dtype-matched
gradient buffer:

```python
torch.addmm(grad_weight, grad_logits_t, input_chunk, out=grad_weight)
```

This removes the separate parameter-sized product/conversion temporary
and extra addition pass. The stored gradient-buffer dtype is unchanged;
the old `.float()` temporary was not an FP32 running accumulator. The
fused accumulation can change rounding relative to the old standalone
product, so bitwise equivalence is not claimed.

The existing guarded FP32 accumulation fast path, AMP operand alignment,
and legacy dtype/device fallbacks are preserved. The CE kernel, its
launch configuration, the default token-chunk heuristic, and the
forward/backward scheduling are unchanged.

**Scope:** this is the first stack layer only: Triton accumulation plus
regressions in the existing FLCE test file. It does **not** add a public
`chunk_size` argument, change cuTile, or include the later three-way
benchmark implementation.

## Performance improvements attributable to this PR

### Unchanged production defaults

BF16 inputs/weights, `accum_dtype=None`, and the existing
`_CHUNK_MEM_CONST=16` heuristic on both sides. The number of token
chunks is identical before and after. At the Llama-3-sized reference
shape, this is **1.26x faster** with **38.9% lower incremental peak
allocated memory** without changing any user settings.

| Tokens / hidden / vocab | Chunk rows | Before ms | After ms | Speedup
| Before / after peak MiB | Memory reduction |
|---|---:|---:|---:|---:|---:|---:|
| 8,192 / 4,096 / 128,256 | 4,096 | 25.002 | 19.897 | **1.26x** |
5,074.04 / 3,102.04 | **38.9%** |
| 4,096 / 2,048 / 32,000 | 4,096 | 1.774 | 1.366 | **1.30x** | 766.02 /
407.02 | **46.9%** |

### Controlled smaller chunks: benefit per accumulation step

These are additional controlled experiments, **not the production
defaults**. For the reference shape, the existing private chunk-budget
constant was set to 4 (1,024-row chunks) or 1 (256-row chunks) on
**both** versions. No new chunk-size API or later stack-layer code was
needed.

| Tokens / hidden / vocab | Chunk rows | Before ms | After ms | Speedup
| Before / after peak MiB | Memory reduction |
|---|---:|---:|---:|---:|---:|---:|
| 8,192 / 4,096 / 128,256 | 1,024 | 43.167 | 21.836 | **1.98x** |
4,322.54 / 1,575.04 | **63.6%** |
| 8,192 / 4,096 / 128,256 | 256 | 123.443 | 31.686 | **3.90x** |
4,134.66 / 1,193.29 | **71.1%** |

The larger gain at smaller chunks is expected: the old implementation
repeatedly allocated/converted/added a full `[vocab, hidden]` matrix,
and this PR removes that work on each iteration. At the reference shape,
one such FP32 matrix alone is approximately 2,004 MiB.

### FP32-accumulator control

The existing `accum_dtype=torch.float32` fast path is intentionally
preserved. At 8,192 tokens / hidden 4,096 / vocab 128,256, with the
unchanged default 4,096-row chunk:

- Before: **20.473 ms**, **4,104.04 MiB** incremental peak.
- After: **20.471 ms**, **4,104.04 MiB** incremental peak.

This control is effectively unchanged; the improvement targets the
default low-precision gradient-buffer path.

### Measurement methodology

- **Hardware:** NVIDIA B200, SM100, driver 580.105.08.
- **Software:** PyTorch 2.11.0.1+cu130, Triton 3.6.0.
- **Before / after source:** base FLCE implementation `96233e3` versus
`6a5077f`, loaded against the same installed dependencies. Current head
`d5f3593` changes only a test comment relative to `6a5077f`; its kernel
implementation is identical.
- Identical seeded BF16 inputs, weights, targets, and token-chunk
geometry; both inputs and weights require gradients; no bias or optional
loss features in the timing cases.
- Measure **full forward + backward**, mean reduction, fixed upstream
gradient 1.0. Each iteration uses a fresh graph and clears both leaf
gradients. No optimizer step or model-training-throughput claim.
- Compilation plus five warmup iterations excluded. Each run uses 30
CUDA-event timed iterations; tables report the median of three runs,
alternating before/after order in the same process.
- Memory is **incremental peak allocated MiB above the live input/weight
baseline**, with both gradients cleared before reset and a dedicated
memory-probe iteration. It is not total process memory or reserved
memory.
- These compare the same Triton-backed implementation before/after the
portable cuBLAS accumulation improvement. They are not
CuTe-versus-Triton or DSL/compiler speedup claims.

## Testing Done

- **Hardware:** NVIDIA B200 / SM100.
- Latest relevant suites: **470 passed, 1 environment-gated skipped**
across `test/ops/test_fused_linear_cross_entropy.py`,
`test/transformers/test_fused_linear_cross_entropy.py`, and
`test/cutedsl/test_fused_linear_cross_entropy.py`.
- All **70 addmm regression cases** now live in the existing ops FLCE
test file alongside its existing 23 cases. Their chunk-budget fixture is
class-scoped and does not alter the pre-existing tests.
- Coverage includes real aten dispatch checks for direct `addmm.out` and
no parameter-shaped upcasts, preserved FP32 `dtype_out`, AMP and
FP32-bias promotion, BF16/FP16/FP32, mean/sum/nonuniform per-token
gradients, ignored chunks, frozen weights, noncontiguous tensors, and
old-Torch/SM capability fallbacks.
- Review requests addressed: restored the original inline FP32 guard,
removed the extra boolean aliases, consolidated tests, and renamed the
marker to `requires_sm80_or_higher`.
- [x] `make checkstyle` passes locally; GitHub Checkstyle is also green.
- [ ] Full-repository `make test` was not run; the targeted suites above
were run.
- [ ] `make test-convergence` was not run; no convergence claim is made.

---------

Co-authored-by: Copilot <[email protected]>
2026-09-10 18:30:11 +00:00
Arup DeandCopilot 96233e3d9b feat(cutile): cuTile backends for RMSNorm, LayerNorm, Softmax, JSD, FusedLinearJSD (#1420)
# feat(cutile): cuTile backends for RMSNorm, LayerNorm, Softmax, JSD,
FusedLinearJSD

**Head:** `arde171:arde/liger-E-cutile` → **Base:** `main`. Final PR of
the multi-backend series; **#1419 (CuTeDSL CE) is now merged**, so this
diff is **cuTile-only** — no rebase caveat anymore.

## What this adds
Registers a **cuTile** backend family for five ops, gated strictly to
Blackwell. All impls declare **`min_cc=(10, 0)`** → only reachable on
sm_100+ (B200/B300); Hopper/Ampere transparently keep using
Triton/CuTeDSL.

| file | + |
|---|---|
| `ops/backends/_cutile/rms_norm.py` | 1221 |
| `ops/backends/_cutile/layer_norm.py` | 626 |
| `ops/backends/_cutile/softmax.py` | 448 |
| `ops/backends/_cutile/jsd.py` | 319 |
| `ops/backends/_cutile/fused_linear_jsd.py` | 64 |
| `ops/backends/_cutile/__init__.py` | 11 |
| `functional.py` | 6 (append `_cutile.*` to 5 op declarations) |
| `ops/fused_linear_jsd.py` | 19 (`_max_capability(device=None)`) |
| `testing/_op_helpers.py` | 5 (`strict_grad` flag) |
| `test/backends/test_cutile_select_mode.py` | 83 |
| `test/cutile/test_cutile_backends_parity.py` | 799 (113-case parity
matrix) |
| `test/ops/{test_rms_norm,test_layer_norm}.py` | 44 |

No CuTeDSL file touched; purely additive registration behind a Blackwell
gate.

## Default-selection is benchmark-justified (nothing defaults to a
slower kernel)
Ranks are chosen so cuTile is **only ever opt-in** — it is never
auto-selected over a faster existing path. Measured median latency
(bf16, fwd+bwd) on **B200 (sm_100)**:

| op | rank | Triton | cuTile | CuTeDSL | auto-selected |
|---|---|---|---|---|---|
| `rms_norm`   | 80 | 0.262 ms | 0.473 ms | **0.190 ms** | CuTeDSL |
| `layer_norm` | 80 | 0.271 ms | 0.623 ms | **0.219 ms** | CuTeDSL |
| `softmax`    | 80 | **0.151 ms** | 0.188 ms | 0.156 ms | Triton |
| `jsd`        | 25 | 0.802 ms | 0.661 ms | **0.583 ms** | CuTeDSL |
| `fused_linear_jsd` | 25 | — | — | — | CuTeDSL |

- **Norms → `preference_rank=80`** (above Triton's 50): cuTile is
0.44–0.55× of Triton on B200, so it must stay opt-in.
- **JSD / FusedLinearJSD → `preference_rank=25`**: both cuTile (1.21×)
and CuTeDSL (1.38×) beat Triton, but CuTeDSL wins, so CuTeDSL stays the
default and cuTile is opt-in below it.

## Testing (deterministic `-p no:randomly`)
Full `test/backends test/ops test/cutile` + the 113-case
CuTeDSL↔cuTile↔Torch parity matrix:
- **B300 (cc 10.3):** `1149 passed, 48 skipped, 0 failed` + `113/113`
parity.
- **B200 (cc 10.0):** `1149 passed, 48 skipped, 0 failed` + `113/113`
parity.
- **H200 (cc 9.0):** green (cuTile `min_cc=(10,0)` correctly skips on
Hopper).
- ruff check + format clean. Mixed-precision affine (dW/dB in their own
dtypes), device-safe launches (all kernels use the input tensor's
device), and `strict_grad` elementwise dW/dB assertions added.

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: Copilot <[email protected]>
2026-09-10 02:57:03 +00:00
d2b5474de7 Add SM100 fused linear SwiGLU (#1340)
## Summary

Adds a unified CuTe DSL SwiGLU implementation on top of the current
`main`
SM100 GEMM:

- a differentiable elementwise forward/backward replacement for
  `LigerSiLUMulFunction`;
- an exact-SM100 fused-linear forward path that computes gate and up
projections
in one grouped GEMM and applies SwiGLU in the epilogue without
materializing
either projection, reading a plainly **stacked** (or **separate**)
gate/up
  weight directly with no pre-packing and no extra weight copy;
- safe fallback to two linear projections plus the elementwise CuTe DSL
kernel
  for unsupported architectures, dtypes, and K alignment.

This PR is rebased directly onto `main`; the former dependency on #1339
has
been removed.

## Implementation

- The fused path consumes a plainly **stacked** `[2N, K]` weight (gate
rows
`[0:N]` then up rows `[N:2N]`) directly:
`run_grouped_epilogue_gemm_stacked`
rebuilds the B tensor with a nested `((32, 2, T):(K, N·K, 32·K))`
N-layout so
gate/up for the same 32 output features still land as adjacent
accumulator
fragments local to the same SM100 CTA — **without any physical repack or
a
  second weight copy**. The MMA mainloop and the FP32 TMEM epilogue are
  unchanged from the interleaved-pack version.
- The grouped custom-epilogue mode is integrated into the current
reusable
`_sm100_gemm.py` without restoring the retired FLCE row-reduction paths.
- The existing single-fragment FLCE configuration remains unchanged.
- Native fused execution requires exact compute capability `(10, 0)`,
  FP16/BF16 operands, and K divisible by 64.
- Unsupported inputs run two `torch.nn.functional.linear` calls followed
by the
CuTe DSL elementwise kernel (separate tensors are used directly here, no
copy).
- Elementwise and fused kernels launch under the input tensor's CUDA
device
guard and on PyTorch's active CUDA stream for both TVM-FFI and DLPack
calling
  paths.

## Weight API

```python
# Stacked (zero-copy): gate rows [0:N] then up rows [N:2N]
# e.g. a model's fused gate_up_proj.weight, or torch.cat([gate, up], dim=0)
h = fused_swiglu(x, gate_up_weight)

# Separate tensors: concatenated once internally for the native path
h = fused_swiglu(x, gate_weight, up_weight)
```

- The **stacked** form is read as a re-strided view of a single buffer,
so there
is **no pre-pack step and no transient 2× weight copy** (previously
incurred by
the 32-row interleave). This resolves the review feedback on the
weight-packing
  memory spike.
- **Separate** `gate`/`up` tensors are accepted for convenience; the
native path
  concatenates them once (a strided view cannot span two independent
allocations), while the fallback path consumes them with no copy. Prefer
the
  stacked form on hot paths.
- Because the weight stays a standard `[2N, K]` (or two `[N, K]`)
tensor, it is
directly usable as an `nn.Parameter` and is checkpoint-/shard-friendly.
- `output_features` is **keyword-only** and inferred as `N` when
omitted;
trailing padded columns (when `N % 32 != 0`) are trimmed automatically.
- `pack_swiglu_weights` is retained as a thin `torch.cat` convenience
helper.

## B200 optimization

NCU at `M=8192, H=4096, N=14336` classified the grouped kernel as
compute-bound:

- compute throughput: **96.07%**
- memory throughput: **62.11%**
- registers/thread: **96**
- dynamic shared memory/block: **213.25 KiB**
- achieved occupancy: **8.63%**

A manual pipeline-stage/swizzle sweep selected:

- five A/B stages at 1K tokens;
- six stages from 2K through 16K;
- four stages with the existing eight-way swizzle from 32K upward.

The large-M specialization raises normalized speedup from
`1.015x`/`0.981x`
at 32K/64K to `1.053x`/`1.052x`.

## BF16 benchmarks

NVIDIA B200, `H=4096`, `N=14336`, seven warmups, 21 interleaved samples.
The
stacked weight is read as a zero-copy view, so nothing is excluded from
latency
on the fast path (separate tensors incur a one-time `torch.cat`). Peak
memory is
measured in isolated provider processes and includes input, one weight
representation, output, and temporary PyTorch CUDA allocations.

| Tokens | Two GEMMs + SwiGLU (ms) | Fused SwiGLU (ms) | Speedup |
Baseline peak (MiB) | Fused peak (MiB) | Memory ratio |
|---:|---:|---:|---:|---:|---:|---:|
| 1,024 | 0.214208 | 0.223136 | 0.960x | 348 | 260 | 1.34x |
| 2,048 | 0.355776 | 0.356896 | 0.997x | 440 | 296 | 1.49x |
| 4,096 | 0.691104 | 0.647360 | 1.068x | 624 | 368 | 1.70x |
| 8,192 | 1.285664 | 1.217984 | 1.056x | 992 | 512 | 1.94x |
| 16,384 | 2.565120 | 2.412224 | 1.063x | 1,728 | 800 | 2.16x |
| 32,768 | 5.130048 | 4.873984 | 1.053x | 3,200 | 1,376 | 2.33x |
| 65,536 | 10.174336 | 9.672512 | 1.052x | 6,144 | 2,528 | 2.43x |

## Correctness

- CuTe DSL SwiGLU, upstream elementwise SwiGLU, SM100 GEMM, and FLCE
regression
  suites: **301 passed, 1 skipped**
- Stacked zero-copy path, separate-weight path (byte-identical to
stacked), and
weight-API validation are covered in `test/cutedsl/test_swiglu.py`
(verified
  alongside `test/cutedsl/test_sm100_gemm.py` on B200 SM100)
- Current-stream coverage: elementwise forward, elementwise backward,
and
  fused-linear forward
- FP16/BF16 fused correctness, padded output tails, exact-SM100
dispatch, FP32
fallback, arbitrary elementwise shapes, multipliers, and PyTorch/Triton
parity
- `make checkstyle`: **passed**

Validation environment:

- NVIDIA B200, SM100, driver `580.105.08`
- PyTorch `2.13.0.1+cu130`
- Triton `3.7.1`
- `nvidia-cutlass-dsl` `4.6.0`
- `apache-tvm-ffi` `0.1.12`

## Scope

- `src/liger_kernel/ops/cutedsl/ops/swiglu.py`
- `src/liger_kernel/ops/cutedsl/ops/_sm100_gemm.py`
- `test/cutedsl/test_swiglu.py`

---------

Co-authored-by: Stella Dey <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>
2026-09-09 23:11:17 +00:00
luca-888 ad2436e80e fix(cutedsl): use the current stream for RMSNorm backward (#1450)
## Summary

CuTe DSL RMSNorm backward does not pass the current CUDA stream to the
kernel launch. During CUDA Graph capture, the core backward kernel is
omitted while subsequent gradient reductions are captured, causing
incorrect gradients on replay. FusedAddRMSNorm shares the same backward
helper.

Forward the TVM FFI environment stream through the backward launcher to
the kernel launch. Add regression tests in the CuTe DSL GPU suite that
replay backward with changed upstream gradients and compare against
eager execution.

## Testing Done

- Hardware Type: NVIDIA H100 80GB and NVIDIA B200 180GB
- Environment: PyTorch 2.9.1 / CUDA 12.8 / CuTe DSL 4.6.0 / TVM FFI
0.1.13.post3
- Before the fix on H100: all 8 CuTe DSL regression cases failed; all 8
Triton controls passed.
- After the fix on H100: all 16 CUDA Graph regression cases and 97
existing CuTe DSL RMSNorm cases passed, with no skips (separate runs).
- After the fix on B200: all 16 CUDA Graph regression cases and 97
existing CuTe DSL RMSNorm cases passed, with no skips (113 cases in one
run).
- Regression cases cover FP32/BF16, RMSNorm in-place/out-of-place
backward, and FusedAddRMSNorm with and without an explicit
residual-output gradient.
- The CUDA Graph regression cases passed with
`LIGER_KERNEL_IMPL=cutedsl` on both GPUs.

Command for the combined targeted suite:

```bash
LIGER_KERNEL_IMPL=cutedsl python -m pytest \
  test/cutedsl/test_rms_norm_stream.py test/cutedsl/test_rms_norm.py
```

- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

Full-repository correctness and model convergence suites were not run.
Validation was focused on the affected kernels and CUDA Graph replay
behavior.
2026-09-09 15:55:13 +00:00
sunyi0505 8c4d822408 [NPU] Optimize Ascend CE/FLCE, RMSNorm, and RoPE kernels (#1426)
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
The core idea of the fusion operator in torch_npu is to combine multiple
scattered computation operations (operators) into one larger operation.
This aims to reduce unnecessary overhead, make full use of the hardware
features of Ascend NPU, and thus significantly improve the performance
of model training and inference. Typically, we directly call the
corresponding interfaces in torch_npu, for example:
torch_npu.npu_cross_entropy_loss/torch_npu.npu_rms_norm/torch_npu.npu_rotary_mul/torch_npu.npu_swiglu.

Tune the Ascend Triton backends for CrossEntropy,
FusedLinearCrossEntropy, RMSNorm, and RoPE. The common theme is fewer
HBM round‑trips, constexpr‑DCE'd masks, software‑pipelined tl.range, and
in‑place buffers so last‑layer training does not peak at two full BT×V
tensors, so as to achieve performance on par with or even exceeding the
fused operators in torch_npu. Refer to
[torch_npu_baselines.py](https://github.com/user-attachments/files/31633905/torch_npu_baselines.py)
for the complete implementation script


<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
### CrossEntropy
- Plain forward processes contiguous row chunks per vector core
(sequential MTE2) instead of one row per program.
- Full tiles are unmasked; `HAS_TAIL` is constexpr so the remainder mask
is DCE'd when `V` divides `BLOCK_SIZE`.
- `tl.range` pipelines MTE2 / VEC / MTE3; OOB labels are clamped before
DMA to avoid a fault ahead of the host assert.
- Dedicated `get_plain_ce_block_size` for the smaller live set on the
no-weight path.
- Last-layer backward (`grad_output == 1`) overwrites the logits buffer
in-place.
### FusedLinearCrossEntropy
- One Cube GEMM for the full `(BT, V)` logits instead of many token
chunks.
- Reuses the optimized CE kernels; writes `dlogits` in-place into the
logits buffer (no extra BT×V).
- Optionally retains forward logits (capped at `min(4 GiB, device_mem /
8)`) to skip the backward logits GEMM.
- `torch.mm(..., out=)` for `dX` / `dW` when dtypes match, avoiding an
extra V×H workspace.
### RMSNorm
- Full-width single-pass fused kernels when a row fits UB (same pattern
as Ascend LayerNorm), avoiding a second load of X after the rstd
reduction.
- 2D fused path (`ROWS_PER_BLOCK × n_cols`) when UB allows; column
tiling only when a full row does not fit.
- On 910B (`BT=8192, H=4096, bf16`) the previous two-pass tiled forward
was ~350 µs (`aiv_mte2_ratio=0.69`) vs `npu_rms_norm` at ~107 µs.
### RoPE
- Native BNSD layout; the old BSND transpose round-trip was ~40% of
Liger RoPE time.
- Work maps to `(batch, head, s_tile)` so 48 AIV cores stay occupied.
- Fast path: contiguous `(BLOCK_S, D)` DMA, `rotate_half` in UB via
`extract_slice` / `insert_slice`.
- Sequential q-then-k launches (dual-stream fights over the same 48 AIV
cores). Short aligned shapes use one two-phase launch to hide the second
host launch.


## 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. 
-->
pytest test/transformers/test_cross_entropy.py
test/transformers/test_rms_norm.py test/transformers/test_rope.py
test/transformers/test_fused_linear_cross_entropy.py
<img width="1215" height="250" alt="image"
src="https://github.com/user-attachments/assets/29f8e01a-a690-4d43-84df-00a80a38db11"
/>


- Hardware Type:  Atls 800T A3(X86)
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence

| Kernel name | Performance |
|---------------|----------------|
| rope forward | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/e95a09cc-9ae4-41b7-a425-cf2b97967f22"
/>|
| rope backward | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/63dca63b-d2cd-4c3f-9ec2-3896cc15eeba"
/> |
| rope full | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/cfd901fb-ee2b-42b2-88dd-12a89d44eb0d"
/>|
| cross_entropy forward | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/9960806f-0da4-4215-b7ee-1820c5130187"
/>|
| cross_entropy backward | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/d631bb0e-6d13-40da-97fa-bffdd9c06ea6"
/>|
| cross_entropy full | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/fd90e4b7-506e-4229-b6d1-74ce435b52d5"
/> |
| cross_entropy no grad forward | <img width="1000" height="600"
alt="image"
src="https://github.com/user-attachments/assets/1eaa515b-adef-4570-9359-13108969dd1c"
/>|
| rms_norm forward | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/a677522c-4d42-4ea4-8a70-1a8970b2d097"
/> |
| rms_norm backward | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/c9b3e04c-11cd-43c5-ad6d-49d376a41375"
/>|
| rms_norm full | <img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/101ccd94-0a6a-4012-90f9-b5828d28a6a2"
/>|
| fused_linear_cross_entropy forward |<img width="1000" height="600"
alt="image"
src="https://github.com/user-attachments/assets/cfdc390a-aff9-4397-95a3-b7a8862a55e8"
/>|
| fused_linear_cross_entropy backward | <img width="1000" height="600"
alt="image"
src="https://github.com/user-attachments/assets/ddc65a5d-eb3f-4efb-8893-09330ff86076"
/>|
| fused_linear_cross_entropy full |<img width="1000" height="600"
alt="image"
src="https://github.com/user-attachments/assets/470090ed-db0d-4aa1-b8c7-861228fed38b"
/>|
| swiglu forward |<img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/83474768-c545-4ee5-8d0d-b596e5eddd48"
/>|
| swiglu backward |<img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/a198dbd3-7f6f-4886-8b46-75f374817fe7"
/>|
| swiglu full |<img width="1000" height="600" alt="image"
src="https://github.com/user-attachments/assets/112527a4-7226-45ac-920e-43387b03d725"
/>|
2026-09-09 12:17:05 +00:00
Piyush UmateandVaibhav Jindal 77206f839d perf(preference): skip dW/dB calculation when parameters are frozen (#1441)
## Description
In parameter-efficient fine-tuning (LoRA / PEFT) or frozen-backbone/head
alignment, model weights `weight` are frozen (`weight.requires_grad =
False`).

Previously, `LigerFusedLinearPreferenceBase` (powering **DPO**,
**SimPO**, **CPO**, **ORPO**) and
`LigerFusedLinearUnpairedPreferenceBase` (**KTO**) unconditionally:
1. Allocated a full weight gradient buffer `grad_weight =
torch.zeros_like(weight)` (consuming ~1.25 GB to 2.5 GB of VRAM for
vocab size 152K).
2. Traced and computed the `d_weight = d_logits^T @ input_chunk` GEMM on
**every single chunk** via `torch.func.grad_and_value(..., argnums=(0,
1), ...)`.
3. Incurred memory bandwidth accumulating
`grad_weight.add_(chunk_grad_weight)` across all chunks.

### Changes
- **Dynamic Autodiff `argnums`**: Dynamically inspect
`weight.requires_grad`, `_input.requires_grad`, and
`bias.requires_grad`. If `weight.requires_grad` is `False`, argument 1
(`weight`) is excluded from `argnums`, completely skipping reverse-mode
autodiff GEMM calculations for `weight`.
- **Zero-Allocation Weight Buffer**: Set `grad_weight = None` when
`weight.requires_grad = False`, eliminating static memory allocation and
chunk-level accumulation.
- **Fixed Latent Broadcasting Bug**: In
`fused_linear_unpaired_preference.py`, fixed an upstream broadcasting
error where `preference_labels_chunk.unsqueeze(1)` failed for
`chunk_size > 1` against 3D `logits_chunk`.
- **Backward Return Compatibility**: When `weight.requires_grad =
False`, `backward()` safely returns `None` for `grad_weight`, matching
PyTorch autograd convention.

---

## Hardware Benchmarks (NVIDIA A100-SXM4-40GB GPU)

Benchmarks measured on **NVIDIA A100-SXM4-40GB** (`bfloat16`, PyTorch
2.14, Triton 3.8.0):

| Workload Configuration | Batch / Context | Vocab Size | Full FT
Latency | LoRA Skip dW Latency | Speedup |
| :--- | :---: | :---: | :---: | :---: | :---: |
| **Qwen / Mistral Alignment** | $B=4, T=512$ | 32,000 | 17.32 ms |
**11.88 ms** | **$1.46\times$ (31.4% faster)** |
| **Long Context DPO** | $B=4, T=1024$ | 32,000 | 30.86 ms | **21.39
ms** | **$1.44\times$ (30.7% faster)** |
| **Large Vocab Alignment** | $B=4, T=512$ | 64,000 | 33.73 ms | **24.39
ms** | **$1.38\times$ (27.7% faster)** |
| **Llama-3 Architecture** | $B=2, T=1024$ | 128,256 | 64.91 ms |
**42.55 ms** | **$1.53\times$ (34.4% faster)** |

---

## Testing
- Added `test/chunked_loss/test_preference_frozen_weight.py` covering
DPO, SimPO, CPO, ORPO, and KTO under:
  - `compiled=True` and `compiled=False`
  - `has_bias=True` and `has_bias=False`
  - mixed `requires_grad` (weight frozen, bias trainable)
  - `20 passed, 0 failed`
- Existing regression test suites (`test_dpo_loss.py`,
`test_simpo_loss.py`) passed 100%.

Co-authored-by: Vaibhav Jindal <[email protected]>
2026-09-09 02:03:32 +00:00
David 2b4f1a413e fix(monkey_patch): report the Liger name on patched gemma3/gemma4/exaone4 MLPs (#1444)
## Summary

`_patch_geglu_module` and `_patch_swiglu_module` rebind both `forward`
and `_get_name`, so a patched MLP prints as `LigerGEGLUMLP` /
`LigerSwiGLUMLP`. Three sites call `_bind_method_to_module` directly and
rebind only `forward`:

- `apply_liger_kernel_to_gemma3_text`
- `apply_liger_kernel_to_gemma4_text`
- `apply_liger_kernel_to_exaone4`

Their MLPs run the Liger kernel but still report the stock name. So
`print(model)`, which is the usual way of checking that patching landed,
says `Gemma3MLP` where gemma2 in the same session says `LigerGEGLUMLP`:

```
gemma2      mlp repr name : LigerGEGLUMLP
gemma3_text mlp repr name : Gemma3MLP
forward actually swapped?  gemma2 True / gemma3_text True
```

Removing exactly this discrepancy was point 2 of #834; these three sites
just never went through the helpers. Routing them through the same
helpers is the whole change. Nothing outside the reported name moves.

## Testing Done

- Added the `_get_name` assertion next to the six existing GeGLU
instance checks. On `main` four of them fail (`gemma3_text`,
`gemma3_conditional_generation`, `gemma4_text`,
`gemma4_conditional_generation`) and three pass (`gemma`, `gemma2`,
`paligemma`, which already go through the helper), so the split matches
the three sites above. All pass with the fix.
- `exaone4` had no instance test at all, so this adds one covering the
lce forward, RMSNorm and SwiGLU MLP patching. It fails on `main` on the
`_get_name` assertion.
- `pytest test/transformers/test_monkey_patch.py`: 63 passed / 1 skipped
on `main`, 64 passed / 1 skipped here (the new exaone4 test). The skip
is `muse_glimmer` missing from the installed transformers.
- `ruff check .` and `ruff format --check .` with the v0.14.11 pinned in
`.pre-commit-config.yaml`: clean.

Hardware Type: CPU. The change is monkey-patch bookkeeping and
`test_monkey_patch.py` has no CUDA gating. Ran on transformers 5.9.0 /
torch 2.11.0.
2026-09-08 22:08:49 +00:00
Arup DeandCopilot 64594266bb feat(cutedsl): CuTeDSL backends for CrossEntropy, FusedLinearCE, FusedLinearJSD (#1419)
# feat(cutedsl): register CuTeDSL CrossEntropy / FusedLinearCE /
FusedLinearJSD backends

**Head:** `arde171:arde/liger-D-cutedsl-ce` → **Base:** `main`. Stacks
on #1418. PR **4/5**.

## What this PR is (and is not)

This is a **dispatcher-wiring + one correctness fix** PR. The CuTeDSL
cross-entropy-family **kernels already exist on `main`** (added by
#1334/#1355/#1370/#1380 etc.). This PR only **registers** them with the
multi-backend dispatcher from #1416 and adds the composed-op adapters.
**It does not modify any existing kernel implementation** —
`fused_scaled_cross_entropy_sm90.py` (fragment forward, #1370),
`_sm100_gemm.py` (selector, #1380), and `cross_entropy.py` `num_warps`
are all **byte-identical to `upstream/main`**.

Diff: **532 insertions / 9 deletions**, 6 files.

## What's added
-
`ops/backends/_cutedsl/{cross_entropy,fused_linear_cross_entropy,fused_linear_jsd}.py`
— `register_op` adapters, plus the shared inner primitives
`cross_entropy_loss_and_grad` / `jsd_loss_and_grad`.
- `functional.py` — `declare_op_locations` for the above.
- `ops/cutedsl/ops/cross_entropy.py` — **+33/-9**, the fp32 correctness
fix only (see below).
- `test/cutedsl/test_cross_entropy.py` — **+46**, a focused
confident-prediction precision test.

## Default-selection safety (addresses the perf review)
Large-vocab CE is **HBM-bandwidth-bound**. Measured on B200 (T=8192,
V=128256, fwd+bwd): standalone cutedsl CE is **tied** in bf16 (1.185 vs
Triton 1.179 ms) and **~5% slower** in fp32 (1.894 vs 1.803). So:
- **standalone `cross_entropy` cutedsl → `preference_rank=60`** (below
Triton's 50): **Triton stays the auto-default**; cutedsl is opt-in via
`impl="nvidia-cutedsl"`.
- **`cross_entropy_loss_and_grad` inner primitive stays
`preference_rank=10`** — it's the kernel the fused-linear CE path
composes.

## Correctness fix (CuTeDSL analogue of Triton #1329)
The cutedsl CE backward now **recomputes the true-class gradient `dx_y`
fully in fp32 and overwrites the target slot with a single store** (was
a read-modify-write of the already-low-precision value). `dx_y =
softmax(x_y) − (1−ls)` cancels catastrophically as `softmax(x_y) → 1`,
so the `−(1−ls)` term must be folded in fp32 before the round. Covers
weighted / label-smoothing / z-loss / softcap / ignore-index; new
`test_ce_true_class_grad_confident_predictions` locks it.

## Testing (Blackwell sm_100 ×2 + Hopper sm_90), deterministic `-p
no:randomly`
| suite | B200 sm_100 | H200 sm_90 |
|---|---|---|
| `test/cutedsl/test_cross_entropy.py` (incl. new precision test) |
**166 / 0** | ✓ |
| `test/transformers/test_cross_entropy.py` | **185 / 0** | 185 / 0 |
| `test/backends` + `test/ops` (D-only) | **634 / 0** | ✓ |

ruff check + format clean.

Co-authored-by: Copilot <[email protected]>
---
**Multi-backend series — review/merge in order:** #1416 (A) → #1417 (B)
→ #1418 (C) → **#1419 (D)** → #1420 (E · cuTile). Stacked; diffs
auto-narrow as earlier PRs land.

Co-authored-by: Copilot <[email protected]>
2026-09-03 15:55:23 +00:00
Vaibhav Jindal 0a5e9a55f8 Update tests.py to use B200 instead of B200! (#1437)
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
"B200!" is not a defined GPU type unlike "H100!". Need to use "B200"
instead of "B200!" at all locations.

Refer: https://modal.com/docs/guide/gpu#b200-gpus

<!---
## 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
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
2026-09-02 17:45:50 +00:00
71293bebcc fix(mlp): correct SwiGLU backward_dI on Blackwell (SM100/B200) (#1435)
## Summary

`test/transformers/test_mlp.py` (the Triton SwiGLU MLP correctness
suite) fails on B200 (SM100) for **every** dtype. Root-caused to a
single kernel — `_swiglu_kernel_backward_dI` — and fixed with a minimal
restructuring of its accumulation loop.

## Symptoms on B200

| dtype | failure |
|-------|---------|
| fp32 | `triton.runtime.errors.OutOfResources: shared memory, Required:
327744, Hardware limit: 232448` |
| bf16 / fp16 | wrong `x.grad` — `calc_diff ≈ 0.12`, relative error ≈
**45%** (forward + weight-grads pass; only `dI` is wrong) |

## Root cause

`dI = dG @ Wg + dU @ Wu` was computed by interleaving **two `tl.dot(...,
acc=acc)` calls into the same accumulator within one K-loop iteration**:

```python
for k in range(0, hidden_dim, BLOCK_K):
    acc = tl.dot(dG_block, gate_weight_block, acc=acc)
    acc = tl.dot(dU_block, up_weight_block, acc=acc)   # same acc
```

On Blackwell (tcgen05) this miscompiles. Minimal repro on B200 (TMA
descriptors, `M,K,N = 256,512,256`, bf16):

| pattern | relative error |
|---------|-----|
| single accumulating dot in a K-loop | ✅ 1.7e-3 |
| **two dots interleaved into one accumulator** | ❌ **4.6e-1** |
| two dots via `+=` into one accumulator | ❌ 4.6e-1 |
| two separate accumulators, added at end | ✅ 1.7e-3 |
| **two sequential K-loops, one accumulator** | ✅ 1.7e-3 |

The same interleaving also pipelines all four operands (`dG, Wg, dU,
Wu`) at once, so for fp32 with `num_stages≥2` the smallest config needs
≥327744 bytes of shared memory — more than SM100 exposes (232448).
Because **every** config in the autotune grid OOMs, the autotuner falls
back to the first (also-OOM) config and the launch crashes.

## Fix

Accumulate the two matmuls in **two sequential K-loops sharing one
accumulator**. This:
- is correct on SM100 (and unchanged behaviour on SM90/H100), and
- pipelines only two operands at a time, so the `BLOCK_K=64` configs fit
in smem and the autotuner selects a valid config for fp32.

The `forward` kernel (two separate accumulators) and `backward_dGU`
kernel (single dot) already use safe patterns and are untouched. The
`fixA` alternative (two separate accumulators) was rejected: it keeps
the 4-operand smem footprint (fp32 still OOMs) and hit `PTXAS Internal
Triton PTX codegen` errors on B200 for `BLOCK_M=128` configs.

## Verification

On B200: `test/transformers/test_mlp.py` → **59 passed, 1 xfailed** (was
28 failed before). fp32/bf16/fp16 across all parametrised shapes now
pass.

Co-authored-by: GitHub Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-02 00:37:28 +00:00
1621f42d7e ci(nvi): add cuTile tests on B200 to the daily run (#1436)
## Summary

Adds a B200 (SM100) counterpart to the daily H100 cuTile CI job, so the
cuTile suite runs on **both** archs — mirroring cuTeDSL, which already
runs on H100 and B200. Arch-specific cuTile regressions (e.g.
SM100/tcgen05 codegen issues) would otherwise go uncaught by CI.

## Changes

- **`dev/modal/tests.py`** — new `liger_cutile_tests_b200` Modal
function, identical to `liger_cutile_tests` but on B200. Installs
`.[dev,cutile-tileiras]` (so `import cuda.tile` resolves and the tests
run rather than `importorskip`) and runs `make test-cutile` (which sets
`LIGER_KERNEL_IMPL=cutile`).
- **`.github/workflows/nvi-ci.yml`** — adds a `cuTile (B200)` entry to
the `daily` matrix and refreshes the surrounding comments.

## Strict `B200!` GPU spec

Both B200 jobs now use the strict `gpu="B200!"` form, matching the H100
jobs' `H100!`. The `!` tells Modal not to substitute a different GPU
family. This matters for arch-specific SM100 tests: without it, a
fallback/substituted GPU could produce a **misleading pass**. This PR
also flips the pre-existing `liger_cutedsl_tests_b200` from `"B200"` →
`"B200!"` for consistency (small, related change — easy to drop if
undesired).

`max-parallel: 3` and the workflow-level concurrency group are
unchanged, so the 5th daily matrix entry still runs at most 3 GPU jobs
at once.

## Verification

Locally on a B200: `LIGER_KERNEL_IMPL=cutile pytest test/cutile/` →
**199 passed, 12 skipped**. The new job is meaningful and currently
green on B200.

Co-authored-by: GitHub Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-01 23:00:26 +00:00
0245181a11 Use native CUTLASS + NVSHMEM package terminology (#1434)
## Summary

- Remove the internal "LCK" nickname from Liger frontend classes,
helpers, messages, tests, and user-facing documentation.
- Use `native` terminology for the CUTLASS + NVSHMEM tensor-parallel
adapter.
- Describe `liger_cute_kernels` as the native CUTLASS + NVSHMEM kernel
package rather than a MoE-only package.
- Document both current operator families: expert-parallel MoE and
tensor-parallel fused scaled linear cross entropy.
- Replace all remaining `LigerCommKernels` and `liger_comm_kernels`
references with `LigerCuteKernels`.
- Keep existing internal build identifiers unchanged; this is a
terminology/documentation update with no runtime behavior change.

## Validation

- `make checkstyle`
- H200 frontend tests: 28 passed

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-01 21:54:38 +00:00
cc7f68fc9c [CUDA] Add tensor-parallel fused scaled cross entropy frontend (#1428)
## Summary

- Add `LigerFusedLinearScaledCrossEntropyTPFunction` and
`liger_fused_linear_scaled_cross_entropy_tp`, with `tp_group` passed as
a call-time argument.
- Route Hopper BF16 inputs to the optional `liger_cute_kernels`
implementation when its FSLCE runtime surface is available.
- Keep the fallback entirely in Liger using a chunked tensor-parallel
implementation adapted from Verl's fused PPO formulas; there is no
runtime dependency on the Verl package.
- Match both paths on global per-token NLL/entropy, ignored targets,
local `dW`, and TP-summed `dX`.
- Leave NVSHMEM bootstrap and team preparation application-owned. FSLCE
resolves only the call-time `tp_group`, so it coexists with LigerMoE
using a different `ep_group` without frontend-global process-group
state.
- Remove the native NCCL dependency and all Torch communicator internals
from the LCK path. Forward and backward communicate exclusively through
the prepared NVSHMEM team, and the TVM-FFI core remains
Torch-ABI-independent.
- Extend the reusable local reduction API with FP32 `MAX`. The forward
WGMMA kernel performs split last-arrival merging, local MAX, correction,
and local SUM in its epilogue.
- For two-host execution, each local GPU sends one unique contiguous
token shard over IBRC; the follow-up kernel merges the remote
online-softmax state, gathers shards with NVLS, and writes
NLL/LSE/entropy.
- Bind native FFI calls to the input CUDA device so autograd worker
threads have a valid context for TMA descriptor construction.

## Validation

- `make checkstyle`
- `python -m pytest
test/transformers/test_fused_linear_scaled_cross_entropy.py -q` - 15
passed, 1 skipped
- H200 distributed frontend matrix using `torch.multiprocessing.spawn` -
10 passed in 77.45s:
  - WORLD process groups at TP1, TP2, TP4, and TP8
  - 4-GPU contiguous subgroups `[0,1]` / `[2,3]`
  - 4-GPU strided subgroups `[0,2]` / `[1,3]`
  - Built-in Liger fallback at TP2 WORLD and strided TP4
- LigerMoE on `ep_group=WORLD` followed by FSLCE on independent strided
TP groups in the same four processes
  - DirectPeer fallback with NVLS disabled
- Native forward self-test: aligned, entropy, ragged M/H/V, temperature,
and large split-N cases passed.
- Raw TVM-FFI world/subgroup tests, including forward and backward CUDA
Graph replay: 4 passed.
- Two-host TP16 IBRC random-data correctness: passed.
- `libliger_cute_kernels.so` has no NCCL dynamic dependency.

### Forward performance

M=8192, H=4096, Vglobal=131072, entropy disabled:

| TP | NVSHMEM p50 | Previous p50 | Delta |
|---:|---:|---:|---:|
| 1 | 11.227 ms | 11.148 ms | +0.7% |
| 2 | 5.639 ms | 5.594 ms | +0.8% |
| 4 | 2.832 ms | 2.821 ms | +0.4% |
| 8 | 1.442 ms | 1.453 ms | -0.8% |
| 16 (two-host IBRC) | 0.829 ms | 0.784 ms | +5.7% |

TP1-8 are within the agreed 1% run-to-run noise threshold; the TP16 IBRC
gap is accepted.

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-01 21:22:46 +00:00
6b7e3dd6fa fix(cutedsl): self-skip SM100-only FLCE tests on SM90 (H100 CI) (#1431)
## Summary

Fixes the **H100 (SM90)** half of the `NVIDIA GPU` daily CI job `cuTeDSL
SM90 (H100)`.

`make test-cutedsl` (`LIGER_KERNEL_IMPL=cutedsl pytest test/cutedsl/`)
runs the whole `test/cutedsl/` folder on **both** H100 (SM90) and B200
(SM100) — per `nvi-ci.yml`, each arch is meant to *self-skip the other
half*. But the native FLCE op
`liger_kernel.ops.cutedsl.ops.fused_linear_cross_entropy` is
**SM100/Blackwell-only** and raises

```
RuntimeError: Native CuTe DSL FLCE requires exact SM100 hardware, FP16/BF16 input and weight, and mean/sum reduction.
```

on non-SM100 GPUs. On the H100 job this made **214 FLCE tests fail
instead of skip**, because the `_run_or_skip` helper only caught
`NotImplementedError`.

## Root cause

The file mixes two skip mechanisms. The 6 `@sm100_required` tests skip
correctly, but the ~34 parity tests are `@cuda_required` and route
through `_run_or_skip`, which only swallowed `NotImplementedError` — not
the SM100 `RuntimeError`. A few tests also call the native kernel
directly (outside `_run_or_skip`), and
`test_flce_native_support_requires_exact_sm100` implicitly assumes
Blackwell hardware.

## Changes (test-only)

- **Broaden `_run_or_skip`** to also `pytest.skip` on the SM100-hardware
`RuntimeError` (message contains `SM100`), re-raising every other
`RuntimeError` so genuine SM100 failures still surface on B200.
- **Route the 3 direct-call tests** through `_run_or_skip`:
`test_flce_native_z_loss_preserves_input_dtype`,
`test_flce_independent_requires_grad_matches_torch`,
`test_flce_all_ignored_targets_are_valid`.
- **Mark `test_flce_native_support_requires_exact_sm100`
`@sm100_required`**: on SM90 `_native_sm100_supported` short-circuits at
`infer_device_arch() != "blackwell"`, so `get_device_capability` is
never called and its `assert seen == [tensor.device]` can't hold.

Guard/validation tests that use `pytest.raises` (`rejects_fp32`,
`rejects_non_sm100`, `reduction_none`, OOB targets, `invalid_inputs`)
still **run and pass** on SM90 — coverage preserved.

## Verification (on an H100 node)

- FLCE file alone: **214 failed → 0 failed** (`12 passed, 225 skipped`,
6.9s).
- Full `test/cutedsl/` (with #1430 also applied): **462 passed, 229
skipped, 0 failed**.
- `ruff check` + `ruff format --check`: clean.

## Relationship to #1430

The other 28 SM90 failures (26 CE-fp32 `assumed_align` + 2 routing) are
fixed by #1430. This PR is complementary and independent (touches only
`test/cutedsl/test_fused_linear_cross_entropy.py`); together they make
the H100 job fully green.

---------

Co-authored-by: svc-kingkong_LinkedIn <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-01 20:26:39 +00:00
083c92342b fix(cutedsl): fp32 CE loss precision on B200 + restore LigerCrossEntropyFunction export (#1430)
## Summary

Two bugs fixed, both caught by `make test-cutedsl` on B200 (SM100).

---

### Bug 1 — `assumed_align=2` causes bf16-precision stores for fp32 loss
tensors (pre-existing)

**Failing tests (26):** all fp32 CE tests with feature flags — z_loss,
softcap, label_smoothing — when `requires_grad=True`.

**Symptom:** Forward loss quantised to bf16 precision (`1.46875` instead
of `1.47212`, `~2.8e-2` error) even though `loss_1d` is fp32. Error only
appears with `HAS_GRAD=True`.

**Root cause:** `to_cute_tensor(loss, assumed_align=2)` passed a 2-byte
alignment hint to `from_dlpack` for a 4-byte (fp32) tensor. When
compiled with `HAS_GRAD=True`, the CuTe DSL compiler generated 16-bit
(bf16) stores for `mLoss[row]` instead of correct 32-bit stores. Same
issue affected the backward `_scale_in_place` gradient scaler.

**Fix:** Replace hardcoded `assumed_align=2` with
`assumed_align=<tensor>.element_size()` at three call sites (loss
buffer, z_loss buffer, scale buffer). Keeps 2-byte alignment for
bf16/fp16, uses 4-byte for fp32.

---

### Bug 2 — `LigerCrossEntropyFunction` not exported from
`transformers.cross_entropy` (since #1416)

**Failing tests (2):**
`test_backend.py::test_liger_kernel_impl_cutedsl_routes` and
`test_cross_entropy.py::test_liger_kernel_impl_cutedsl_selects_cutedsl_ce`.

**Root cause:** PR #1416 refactored
`liger_kernel/transformers/cross_entropy.py` to use `dispatch()` from
the new multi-DSL backends system, removing the explicit
`LigerCrossEntropyFunction` import. The cutedsl routing tests expect
this attribute to be importable and to resolve to the cutedsl class when
`LIGER_KERNEL_IMPL=cutedsl`.

**Fix:** Re-add the import as a re-export (`# noqa: F401`). Because
`liger_kernel.ops` applies backend replacement at import time, the
attribute correctly resolves to the cutedsl implementation when
`LIGER_KERNEL_IMPL=cutedsl`.

---

## Verification (B200 / SM100)

```
641 passed, 50 skipped, 0 failed  (50 skipped = SM90-only tests self-skip on B200)
```

Co-authored-by: GitHub Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-09-01 17:19:17 +00:00
Buffett LiuandZhipeng Wang e6a81bb0c3 Add Liger Kernel support for Muse Glimmer (#1390)
## Summary
Adds Liger Kernel support for Muse Glimmer (`model_type=muse_glimmer`),
covering both capabilities requested in #1378:

1. FusedLinearCrossEntropy / skip-logits loss path, avoiding the `[B, S,
V]` logits tensor (Muse Glimmer's `vocab_size` is 202,048)
2. `return_token_accuracy=True` → `outputs.token_accuracy`, so TRL's
`SFTTrainer` can log `mean_token_accuracy` without a trainer-side
fallback

Registered in `MODEL_TYPE_TO_APPLY_LIGER_FN`, so `use_liger_kernel=True`
picks it up. Patches RoPE, RMSNorm, SwiGLU, vision LayerNorm (instance
path only), CrossEntropy and FusedLinearCrossEntropy.

## Details
**Softcap.** HF computes `T * tanh(logits * m / T)` where `m =
output_multiplier` and `T = final_logit_softcapping`, but Liger's fused
softcap only implements `T * tanh(logits / T)`. `m` is folded into the
hidden states instead, since `(m * h) @ W.T == m * (h @ W.T)` — scaling
`[batch, seq, hidden]` rather than `[vocab, hidden]`.

**Two RMSNorm variants.** `MuseGlimmerRMSNorm` scales by `weight`
directly (gemma-style fp32 compute, offset 0);
`MuseGlimmerTextCenteredRMSNorm` scales by `(1 + weight)` from zero-init
(offset 1). The `with_scale=False` instances, `qk_norm`, `embed_norm`,
`perception_emb_norm`, have no weight parameter and fall back to a torch
implementation matching HF exactly; an assert guards against that
fallback silently dropping a scale.

**`in_place=False`** on the TextCentered norms is required, not
conservative: each decoder layer hands `pre_feedforward_layernorm` the
same tensor it retains as the residual.

Note that `token_accuracy` is populated on the fused path only; with
`skip_logits=False` (eval) it stays `None`, matching existing patched
models such as Qwen3.

## Testing Done
- `test/transformers/test_muse_glimmer.py` (new) — instance-path
numerical parity against the unpatched model: loss plus embedding /
SwiGLU / RMSNorm gradients in fp32 and bf16, the non-fused softcap
branch, vision LayerNorm, and the `token_accuracy` / `predicted_tokens`
contract. The convergence suite patches at module level, so this is what
covers `_apply_liger_kernel_to_instance` — the path TRL actually takes.
- `test/transformers/test_monkey_patch.py` — instance patching
- Convergence in bf16 and fp32: `test_mini_models`,
`test_mini_models_multimodal`, `test_mini_models_with_logits`

### Memory

H100, bf16, forward + backward, no optimizer step. Randomly initialised
(no public checkpoint), real `hidden_size=6656` / `vocab_size=202048`,
`num_hidden_layers` reduced from 52 so a 30B-class model fits on one
GPU.

| seq_len | peak mem (HF) | peak mem (Liger) | saved |
|--------:|--------------:|-----------------:|------:|
| 1024 | 18.4 GB | 20.0 GB | -9% |
| 2048 | 19.6 GB | 21.2 GB | -8% |
| 4096 | 30.4 GB | 23.5 GB | 23% |
| 8192 | 52.0 GB | 28.2 GB | 46% |
| 16384 | OOM | 38.2 GB | — |

The crossover sits near 3K: FLCE accumulates `grad_weight` in fp32, a
fixed 5.4 GB for `[202048, 6656]` against bf16's 2.7 GB. Below ~3K that
outweighs the logits it avoids; above it the saving compounds, and 16K
only runs at all with Liger.

### Throughput

FLCE trades speed for memory by design. End to end that penalty is
diluted by the transformer body, and the reduced depth above overstates
it:

| layers | lm_head share | tok/s (HF) | tok/s (Liger) | ratio |
|-------:|--------------:|-----------:|--------------:|------:|
| 4 | 41% | 22,405 | 13,433 | 0.60x |
| 8 | 26% | 13,685 | 10,396 | 0.76x |
| 16 | 15% | 7,705 | 7,063 | 0.92x |

`lm_head share` is measured from the model's parameter split, which
tracks the matmul FLOP split. The shipping model's 52 layers put it at
~5%, below the last row here.

- 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: Zhipeng Wang <[email protected]>
2026-09-01 05:32:27 +00:00
Arup DeandCopilot 2798d08546 feat(cutedsl): CuTeDSL backends for SwiGLU, GeGLU, Softmax, RoPE, KL-div, JSD (#1418)
# feat(cutedsl): CuTeDSL backends for SwiGLU, GeGLU, Softmax, RoPE,
KL-div, JSD

**Head:** `arde171:arde/liger-C-cutedsl-elementwise` → **Base:** `main`
— stacks on PR 2. PR **3/5**.

## Summary

Adds the **CuTeDSL elementwise/reduction** op backends, registered via
`declare_op_locations` (Blackwell auto-select, Triton fallback).
Includes several **correctness fixes** surfaced by an independent
GPT-5.6 review of these kernels.

## What's added

- `ops/backends/_cutedsl/{swiglu,geglu,softmax,rope,kl_div,jsd}.py` +
`ops/cutedsl/ops/rope.py`.
- `functional.py` — appends the `_cutedsl` paths for these six ops.
- `test/cutedsl/test_rope.py`,
`test/ops/test_jsd_chunk_ignore_regression.py`.

## Correctness fixes included

- **SwiGLU / DTensor:** the CuTeDSL SwiGLU now **preserves DTensor
(tensor-parallel) semantics** by delegating DTensor inputs to Triton
(the previous path materialized the full tensor and returned plain
tensors). Dense fast path unchanged. fp32 fwd tolerance calibrated for
the fast sigmoid.
- **Softmax fp32 stability:** compute `exp2((x - max) * log2e)` instead
of `exp2(x*log2e - max*log2e)` — shift-invariant, avoids fp32
cancellation at large offsets.
- **JSD chunked ignore-labels (P0):** decide the label-stream per chunk
(was comparing a global non-ignore count with the per-chunk size), which
had let ignored rows leak loss/gradient. Regression test added (asserts
ignored rows get exactly zero input-grad across chunked shapes).
- **RoPE TMA guard:** `_tma_supported` rejects shapes whose launch block
would exceed 1024 threads (e.g. bf16 head_dim=2048) → falls back to the
portable token kernel.
- **GeGLU:** fp32 tolerance calibrated for the fast tanh-gelu.

## Correctness & testing (B300 sm_103 + H200 sm_90)

- Registration: `swiglu:[triton,cutedsl]`, `jsd:[triton,cutedsl]`,
`jsd_loss_and_grad:[triton,cutedsl]` (inner op declared here, co-located
with its `_cutedsl/jsd.py` impl).
- Full backend suite `test/backends` + `test/ops`, deterministic (`-p
no:randomly`): **B300 609 passed / 0 failed**, **H200 609 passed / 0
failed** (39 skipped). The chunked-FLJSD ignore-label regression test
exercises the CuTeDSL JSD path on both archs.
- ruff clean.

Co-authored-by: Copilot <[email protected]>
---
**Multi-backend series — review/merge in order:** #1416 (A · dispatcher
foundation) → #1417 (B · CuTeDSL norms) → #1418 (C · CuTeDSL
elementwise) → #1419 (D · CuTeDSL CE) → #1420 (E · cuTile). Each PR is
stacked on the previous; the diff auto-narrows as earlier PRs land. GPU
tests are cron-only (per #1409) so per-PR CI = checkstyle + CPU tests;
multi-backend suites verified out-of-band on **B300 (sm_103)** + **H200
(sm_90)**.

Co-authored-by: Copilot <[email protected]>
2026-09-01 00:45:13 +00:00
Arup DeandCopilot 8dedc9703a feat(cutedsl): CuTeDSL backends for RMSNorm, LayerNorm, fused_add_rms_norm (#1417)
# feat(cutedsl): CuTeDSL backends for RMSNorm, LayerNorm,
fused_add_rms_norm

**Head:** `arde171:arde/liger-B-cutedsl-norms` → **Base:** `main` —
stacks on PR 1 (dispatcher). PR **2/5**.

## Summary

Adds the first **CuTeDSL** (NVIDIA CUTLASS Python DSL) backends — the
three normalization ops — plus the shared `_cute_lib/` helpers.
Registered into the dispatcher via `declare_op_locations`;
**auto-selected on Blackwell** (where `cutlass.cute` is usable), with
**Triton as the fallback** on Hopper/Ampere.

## What's added

- `ops/backends/_cutedsl/{rms_norm,layer_norm,fused_add_rms_norm}.py` —
the backend impls.
- `ops/backends/_cutedsl/_cute_lib/*` — shared
compile/copy/reduce/layout utilities used by all CuTeDSL kernels.
- `ops/cutedsl/ops/{rms_norm,rms_norm_fastpath}.py` — the low-level
CuTeDSL kernels.
- `functional.py` — appends the
`_cutedsl.{rms_norm,layer_norm,fused_add_rms_norm}` paths to those ops'
declarations.

## Design

`min_cc=(9,0)` (Hopper+). The two-pass RMSNorm keeps a per-SM fp32 `dW`
partial reduced across SMs; forward/backward mirror Quack's SM-count
heuristic. The dispatcher auto-selects CuTeDSL on Blackwell and defers
to Triton where `cutlass.cute` isn't usable.

## Correctness & testing (B300 sm_103 + H200 sm_90, torch 2.13/cu130)

- Registration: `rms_norm: [nvidia-triton, nvidia-cutedsl]`.
- Full backend suite `test/backends` + `test/ops`, deterministic (`-p
no:randomly`): **B300 474 passed / 0 failed**, **H200 474 passed / 0
failed** (49 skipped); also confirmed on **H100 474 / 0**.
- Bf16 weight-grad (`dw`) parity uses a reduction-appropriate tolerance
(it's a batch sum-reduction; the strict elementwise tol is for fwd/dx).
- ruff clean.

## Scope

CuTeDSL norm backends only; additive registration. No change to Triton
behavior; non-Blackwell falls back to Triton.

Co-authored-by: Copilot <[email protected]>
---
**Multi-backend series — review/merge in order:** #1416 (A · dispatcher
foundation) → #1417 (B · CuTeDSL norms) → #1418 (C · CuTeDSL
elementwise) → #1419 (D · CuTeDSL CE) → #1420 (E · cuTile). Each PR is
stacked on the previous; the diff auto-narrows as earlier PRs land. GPU
tests are cron-only (per #1409) so per-PR CI = checkstyle + CPU tests;
multi-backend suites verified out-of-band on **B300 (sm_103)** + **H200
(sm_90)**.

Co-authored-by: Copilot <[email protected]>
2026-08-31 22:50:45 +00:00
Arup DeandCopilot 90f807a321 feat(backends): multi-DSL dispatcher foundation + Triton backend layer (#1416)
# feat(backends): multi-DSL dispatcher foundation + Triton backend layer

**Head:** `arde171:arde/liger-A-slim` → **Base:**
`linkedin/Liger-Kernel:main` — 61 files, PR **1/5** of the multi-backend
series.

## Summary

Introduces a **per-(op, impl) backend dispatcher** so a single Liger op
can have multiple kernel implementations (Triton, CuTeDSL, cuTile) and
pick one by **preference-rank + capability**, with **Triton as the
universal fallback**. This PR ships *only* the foundation + the Triton
backend layer; on its own it routes every op through Triton exactly as
before (no behavior change). The CuTeDSL and cuTile backends register
into this in follow-up PRs 2–5.

## What's added

- `src/liger_kernel/backends/`
- `dispatch.py` — selection: explicit `impl=` → env
(`LIGER_KERNEL_IMPL[_OP]`) → global `set_impl` → capability-ranked
auto-pick → Triton fallback. Raises a precise `ImplNotAvailableError`
when a requested impl isn't registered.
- `registry.py` — `register_impl`/`declare_op_locations`; lazy per-op
discovery (imports a backend module only on first use of its op).
- `capability.py` — declarative `Capability(min_cc, max_cc, modules,
requires_cuda)`; resolves compute capability from the **input tensor's
device** (`torch.cuda.get_device_capability(device)`), threaded through
`dispatch`→`resolve_impl`→`available_impls`→`is_satisfied`, so selection
is correct even for heterogeneous multi-GPU-per-process. Memoized per
`(op, device, device-cc, cuda-visible)`.
- `ops/backends/_triton/*` — the existing Triton kernels as first-class
registered impls.
- `functional.py` — `declare_op_locations` for every op (this PR:
**Triton path only**; later PRs append their own `_cutedsl`/`_cutile`
paths — additive, so no PR references a backend it doesn't ship).
- `ops/_nvidia_shared.py`, `testing/_op_helpers.py` (shared
correctness/tolerance harness), `test/backends/*`.
- A **test-isolation fixture** (`test/conftest.py`) that resets the
global impl selection after each test, so backend-selecting tests don't
leak into later tests.

## Design — preference-rank + Triton fallback

Each impl declares a rank and a `Capability`. Auto-selection picks the
lowest-rank impl whose capability holds on the input tensor's device; if
none (or the module can't import), it falls back to Triton. Discovery is
lazy and idempotent.

## Correctness & testing (B300 sm_103 + H200 sm_90, torch 2.13/cu130)

- `import liger_kernel` — **warning-clean** (declares only shipped
Triton paths).
- Full backend suite `test/backends` + `test/ops`, deterministic (`-p
no:randomly`): **B300 348 passed / 0 failed**, **H200 348 passed / 0
failed** (44 skipped).
- ruff check + format: clean across the repo.

## Scope

Foundation only. No CuTeDSL/cuTile code, no benchmark data, no
dev-scaffold, no doc essays (those are in later PRs / stay internal).
Behavior is Triton-identical.

Co-authored-by: Copilot <[email protected]>
---
**Multi-backend series — review/merge in order:** #1416 (A · dispatcher
foundation) → #1417 (B · CuTeDSL norms) → #1418 (C · CuTeDSL
elementwise) → #1419 (D · CuTeDSL CE) → #1420 (E · cuTile). Each PR is
stacked on the previous; the diff auto-narrows as earlier PRs land. GPU
tests are cron-only (per #1409) so per-PR CI = checkstyle + CPU tests;
multi-backend suites verified out-of-band on **B300 (sm_103)** + **H200
(sm_90)**.

Co-authored-by: Copilot <[email protected]>
2026-08-31 22:37:35 +00:00
Vaibhav JindalandCopilot c02362cb5b fix(cutile): pin cuda-tile>=1.5.0 so CI gets CompilerOptions.num_worker_warps (#1427)
## Summary

Fixes the scheduled **cuTile (H100)** job, which failed at
import/collection with:

```
src/liger_kernel/ops/cutile/ops/dyt.py:144: in <module>
    _dyt_bwd_kernel_nww8 = _dyt_bwd_kernel.replace_hints(num_worker_warps=8)
...
E   TypeError: CompilerOptions.__init__() got an unexpected keyword argument 'num_worker_warps'
```

## Root cause

The cuTile ops use `num_worker_warps` both as a `@ct.kernel(...)`
argument and via
`replace_hints(num_worker_warps=...)` (in `dyt.py`,
`fused_add_rms_norm.py`,
`poly_norm.py`, `softmax.py`, `fused_linear_cross_entropy.py`,
`swiglu.py`).

That `CompilerOptions` field was **only added in cuda-tile 1.4.0** —
it's absent in
1.3.0 and 1.0.0 (verified by inspecting the wheels). Because the
`cutile` /
`cutile-tileiras` extras pinned only `"cuda-tile"` (unpinned),
dependency resolution
backtracked to an older cuda-tile whose `[tileiras]` extra constraints
were
satisfiable, landing below 1.4.0. The field was then missing, so
importing
`liger_kernel.ops.cutile.ops` (triggered under
`LIGER_KERNEL_IMPL=cutile`) crashed and
took down the whole cuTile suite at collection.

Since this runs at import via the backend-replacement in
`liger_kernel/ops/__init__.py`,
it is not recoverable at runtime — the version floor must guarantee the
field exists.

## Fix

Pin the floor to `cuda-tile>=1.5.0` (the validated version) in both
cuTile extras in
`setup.py`, so the resolver can no longer fall back to a version without
`num_worker_warps`.

## Notes

- cuTeDSL (SM90/SM100) jobs also failed in the same run, but that's a
**separate**
  library (`nvidia-cutlass-dsl`) and will be addressed separately.
- Verified locally: `num_worker_warps` present in cuda-tile 1.4.0 &
1.5.0, absent in
1.3.0 & 1.0.0; `setup.py` parses and is ruff-clean. The full cuTile
suite needs an
  H100 to confirm end-to-end (no GPU in the dev sandbox).

Co-authored-by: Copilot <[email protected]>
2026-08-31 21:43:55 +00:00
99747c53ba [CUDA] Add Hopper tensor-parallel fused linear cross entropy (#1422)
## Summary

Adds Hopper tensor-parallel fused linear scaled cross entropy
forward/backward with NVLS, direct-peer, and hierarchical two-host
paths.

Production choices:
- zero-spill cluster-2 M-fast dZ handoff with compact precomputed
coefficients;
- zero-spill four-stage combined dX+dW with one reused R24 accumulator;
- H2048 dX split-K2 when the local vocabulary has an even K64 tile
count,
  otherwise split-K1; H4096 uses split-K1;
- N32 FP32 TMA dX publication, N64/S2 dW epilogue, final-only
publication fence;
- post-compute local/hierarchical FP32-to-BF16 finalization;
- device-resident graph-safe launch epochs with one stream-ordered tail
barrier;
- non-RDC WGMMA translation units with remote NVSHMEM code isolated in
RDC sources.

## Tensor-parallel reduction API

Reduction transport is now reusable infrastructure under
`liger_cute::detail`:

- `TpReducePlan` selects the local NVLS or DirectPeer backend on the
host.
- Kernels receive only the selected compact device view; the full host
plan is never passed to device code.
- `RequiresRemote` is a launcher template argument, so the local path
carries no runtime remote branch or remote fields.
- Two-host execution is local NVLS reduce-scatter, explicit
`launch_remote_reduce(...)`, then the common local all-gather/scatter.
- DirectPeer remains supported behind the same local reduction API.

The refactor removes the unused pre-cluster dZ/role framework and old
non-cluster hierarchical implementation. Net change: 1,907 fewer lines
while preserving the validated kernel schedules.

## Kernel structure

### Warp functions and roles

**Forward `forward_gemm_kernel_sm90` (12 warps, cluster-2)**
- Warp 0 is the TMA producer for X and weight panels.
- Warps 1-3 are reserved/inactive.
- Warp groups 1-2 (8 warps) execute the N320 logical GEMM as two N160
accumulator banks, then perform the N80/N40 online-softmax and entropy
fold.
- Cluster peers share the token-side input while owning adjacent
vocabulary work.

**Backward dZ handoff (12 warps, cluster-2)**
- Warp 3 drives the three-stage TMA pipeline for X/W panels.
- Warps 0-2 consume handed-off FP16 logits, apply precomputed row
coefficients, target correction, BF16 conversion, and TMA output stores.
- Warp groups 1-2 execute M128xN256xK64 WGMMA and hand the rounded FP16
logits through shared memory.
- The handoff separates accumulator and epilogue liveness, eliminating
the old dZ spill path.

**Backward combined dX+dW (12 warps, cluster-2)**
- Warp 0 drives the continuous four-stage TMA producer pipeline.
- Two communication warps consume the FP32 dX ring, perform local NVLS
reduction/publication, and overlap with dW.
- Warp groups 1-2 execute dX then dW using one explicitly reused
M128xN256 FP32 accumulator bank at R24.
- dX uses W multicast and N32 FP32 TMA publication. dW uses dZ
multicast, independent X loads, and N64/S2 TMA store/reduce-add.
- H2048 uses split-K2 with atomic accumulation into a stream-cleared
durable
  buffer when its K64 tile count is even; odd tile counts and H4096 use
  split-K1.

### Forward launch sequence

1. `forward_gemm_kernel_sm90`: local GEMM plus online
max/sum/entropy/target partials.
2. `forward_split_reduce_kernel_sm90`: reduce local vocabulary splits.
3. Tensor-parallel MAX reduction.
4. `pack_forward_reduction_kernel`: correct/pack sum fields.
5. Tensor-parallel SUM reduction.
6. `finalize_forward_kernel`: produce log-probability, LSE, and optional
entropy.

### Backward launch sequence

For each token wave:
1. `backward_dz_handoff_wave_kernel_sm90`: logits GEMM, compact dZ math,
BF16 dZ workspace.
2. `backward_dx_cluster2_gemm_wave_kernel_sm90`: continuous four-stage
dX then dW compute with overlapped local communication.

After all waves:
3. Single host: local FP32 reduction result to BF16 scatter.
4. Hierarchical: separate-RDC inter-host reduction, local NVLS
broadcast/all-gather, then BF16 scatter.
5. Direct-peer systems use the post-compute fallback finalizer.
6. One stream-ordered tensor-parallel tail barrier prevents staging
reuse before peers finish.

### Hoisted versus repeated work

**Hoisted**
- Backend capability/team selection, workspace reservation, rank
mapping, and TMA descriptor construction occur once on the host.
- TMA descriptors are prefetched once per CTA role before persistent
tile loops.
- Shape counts, raster bounds, and split-K selection are computed
outside inner K loops.
- dZ stores per-row `exp_bias`, `correction_bias`, and
`correction_slope`; row coefficients/target are hoisted outside the
inner column loop.
- dX and dW share one accumulator object and aliased A/B shared-memory
arenas.
- IB communication and final BF16 materialization remain outside the hot
combined WGMMA kernel.

**Not hoisted**
- Per-wave dZ and combined dX+dW launches.
- Per-tile TMA loads, WGMMA issue, accumulator-to-SMEM conversion, and
dX publication.
- Per-element dZ `exp2`, entropy correction, and target subtraction.
- Per-stage NVLS ring synchronization/reduction and split-K2 atomic
accumulation.
- dW N64 TMA store/reduce-add for each output subtile and wave.

## Validation

- `make checkstyle`
- clean CMake shared-library build
- C++ tests: 8/8 passed
- Python WORLD and strided-subteam forward/backward tests: 2/2 passed,
including two CUDA Graph replays
- backward TP1/2/4/8 selftests: entropy on/off, ragged shapes,
multi-wave/full-grid, TPR1/2/4
- poison staging/dZ and delayed communication warp tests
- forward TP1/2/4/8 selftests: entropy on/off, ragged, temperature,
large split
- hierarchical TP16 forward/backward
- benchmarked kernel schedule before the transport-only API refactor:
92/92 unique matched points passed in both entropy modes; TP1/2/4/8/16
- post-refactor: checkstyle and focused Python facade regression passed;
GPU compilation remains delegated to PR validation

## Benchmark methodology

H200 SXM, CUDA 12.9, 8 warmups and 30 measured iterations, p50 shown
below. TFLOP/s are per GPU unless labeled aggregate. Liger entropy-off
is compared with Verl's no-entropy configuration. Memory factor is `Verl
GiB / Liger GiB`; values above 1 mean Liger uses less memory.
Allocation-free means no hot-path allocation; Liger's reported footprint
includes pre-reserved symmetric communication workspaces. H100 and H200
absolute results are kept separate.

## Overall Liger vs Verl

| Pass | Matched points | Geomean speedup | Range | Slower points |
Geomean mem V/L |
|---|---:|---:|---:|---:|---:|
| backward | 46 | 1.54x | 0.95-1.87x | 1 | 0.82x |
| forward | 46 | 2.11x | 1.69-3.79x | 0 | 0.47x |

Forward is faster at every matched point. Backward is faster at every
matched point except M8192/H2048/V131072 at TP16.

## Strong-scaling efficiency

Efficiency is per-GPU TFLOP/s at TPn divided by TP1 per-GPU TFLOP/s for
the same global workload.

| Pass | M | H | V | TP1 TF/GPU | TP2 eff. | TP4 eff. | TP8 eff. | TP16
eff. |
|---|---:|---:|---:|---:|---:|---:|---:|---:|
| backward | 2048 | 4096 | 131072 | 767.9 | 98.1% | 92.7% | 82.5% |
50.1% |
| backward | 8192 | 2048 | 131072 | 694.9 | 91.4% | 82.8% | 68.6% |
37.8% |
| backward | 8192 | 4096 | 131072 | 762.8 | 97.8% | 93.7% | 84.0% |
54.5% |
| backward | 16384 | 4096 | 131072 | 761.1 | 97.3% | 93.8% | 83.7% |
55.3% |
| backward | 2048 | 4096 | 262144 | 783.1 | 98.2% | 97.0% | 91.1% |
66.8% |
| backward | 8192 | 2048 | 262144 | 711.3 | 96.7% | 92.3% | 85.0% |
58.2% |
| backward | 8192 | 4096 | 262144 | 777.4 | 98.4% | 96.6% | 91.8% |
70.6% |
| backward | 16384 | 4096 | 262144 | 775.8 | 98.5% | 95.8% | 91.9% |
71.2% |
| forward | 2048 | 4096 | 131072 | 787.7 | 99.9% | 97.7% | 83.5% | 58.4%
|
| forward | 8192 | 2048 | 131072 | 747.7 | 99.4% | 96.7% | 90.4% | 79.1%
|
| forward | 8192 | 4096 | 131072 | 789.0 | 99.7% | 98.8% | 95.9% | 88.9%
|
| forward | 16384 | 4096 | 131072 | 757.4 | 99.7% | 98.3% | 95.0% |
91.4% |
| forward | 2048 | 4096 | 262144 | 791.1 | 99.0% | 100.0% | 95.4% |
77.0% |
| forward | 8192 | 2048 | 262144 | 747.5 | 99.7% | 99.1% | 95.6% | 86.8%
|
| forward | 8192 | 4096 | 262144 | 788.3 | 99.9% | 100.1% | 98.3% |
94.1% |
| forward | 16384 | 4096 | 262144 | 755.3 | 100.0% | 99.8% | 98.3% |
93.9% |

## Weak-scaling efficiency

Global vocabulary grows with TP while local vocabulary remains 65,536.
Efficiency is per-GPU TFLOP/s relative to TP1.

| Pass | H | TP1 TF/GPU | TP2 eff. | TP4 eff. | TP8 eff. | TP16 eff. |
|---|---:|---:|---:|---:|---:|---:|
| backward | 2048 | 661.4 | 96.1% | 99.2% | 99.2% | 94.7% |
| backward | 4096 | 757.6 | 98.4% | 99.1% | 98.8% | 95.8% |
| forward | 2048 | 748.2 | 99.4% | 99.0% | 98.5% | 97.4% |
| forward | 4096 | 788.5 | 99.7% | 100.1% | 99.3% | 98.4% |

## Full Liger vs Verl results

| Pass | M | H | V | TP | Liger TF/GPU | Verl TF/GPU | Speedup | Liger
GiB | Verl GiB | Mem V/L |
|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
| backward | 8192 | 2048 | 65536 | 1 | 661.4 | 375.2 | 1.76x | 2.26 |
1.15 | 0.51x |
| backward | 8192 | 4096 | 65536 | 1 | 757.6 | 507.3 | 1.49x | 2.82 |
1.99 | 0.71x |
| backward | 2048 | 4096 | 131072 | 1 | 767.9 | 445.2 | 1.72x | 3.85 |
3.14 | 0.82x |
| backward | 2048 | 4096 | 131072 | 2 | 753.3 | 436.4 | 1.73x | 2.72 |
1.65 | 0.61x |
| backward | 2048 | 4096 | 131072 | 4 | 711.6 | 385.0 | 1.85x | 2.16 |
0.88 | 0.41x |
| backward | 2048 | 4096 | 131072 | 8 | 633.3 | 362.6 | 1.75x | 1.88 |
0.52 | 0.28x |
| backward | 2048 | 4096 | 131072 | 16 | 385.0 | 311.3 | 1.24x | 0.17 |
0.33 | 1.94x |
| backward | 8192 | 2048 | 131072 | 1 | 694.9 | 378.9 | 1.83x | 2.88 |
1.89 | 0.66x |
| backward | 8192 | 2048 | 131072 | 2 | 635.4 | 372.0 | 1.71x | 2.26 |
1.15 | 0.51x |
| backward | 8192 | 2048 | 131072 | 4 | 575.2 | 329.3 | 1.75x | 1.94 |
0.71 | 0.37x |
| backward | 8192 | 2048 | 131072 | 8 | 476.9 | 313.8 | 1.52x | 1.79 |
0.56 | 0.31x |
| backward | 8192 | 2048 | 131072 | 16 | 262.4 | 277.1 | 0.95x | 0.14 |
0.49 | 3.50x |
| backward | 8192 | 4096 | 131072 | 1 | 762.8 | 510.6 | 1.49x | 3.94 |
3.48 | 0.88x |
| backward | 8192 | 4096 | 131072 | 2 | 745.7 | 501.2 | 1.49x | 2.82 |
1.99 | 0.71x |
| backward | 8192 | 4096 | 131072 | 4 | 714.7 | 446.2 | 1.60x | 2.26 |
1.18 | 0.52x |
| backward | 8192 | 4096 | 131072 | 8 | 640.6 | 431.4 | 1.48x | 1.97 |
0.84 | 0.43x |
| backward | 8192 | 4096 | 131072 | 16 | 415.7 | 387.3 | 1.07x | 0.27 |
0.68 | 2.52x |
| backward | 16384 | 4096 | 131072 | 1 | 761.1 | 519.8 | 1.46x | 4.07 |
3.93 | 0.97x |
| backward | 16384 | 4096 | 131072 | 2 | 740.2 | 510.8 | 1.45x | 2.94 |
2.46 | 0.84x |
| backward | 16384 | 4096 | 131072 | 4 | 714.1 | 458.3 | 1.56x | 2.38 |
1.58 | 0.66x |
| backward | 16384 | 4096 | 131072 | 8 | 637.1 | 444.4 | 1.43x | 2.10 |
1.28 | 0.61x |
| backward | 16384 | 4096 | 131072 | 16 | 421.0 | 404.6 | 1.04x | 0.39 |
1.13 | 2.90x |
| backward | 2048 | 4096 | 262144 | 1 | 783.1 | 449.4 | 1.74x | 6.10 |
6.14 | 1.01x |
| backward | 2048 | 4096 | 262144 | 2 | 768.8 | 444.1 | 1.73x | 3.85 |
3.14 | 0.82x |
| backward | 2048 | 4096 | 262144 | 4 | 759.7 | 435.4 | 1.74x | 2.72 |
1.65 | 0.61x |
| backward | 2048 | 4096 | 262144 | 8 | 713.6 | 383.2 | 1.86x | 2.16 |
0.88 | 0.41x |
| backward | 2048 | 4096 | 262144 | 16 | 523.4 | 351.2 | 1.49x | 0.31 |
0.52 | 1.68x |
| backward | 8192 | 2048 | 262144 | 1 | 711.3 | 381.3 | 1.87x | 4.13 |
3.35 | 0.81x |
| backward | 8192 | 2048 | 262144 | 2 | 687.8 | 376.7 | 1.83x | 2.88 |
1.89 | 0.66x |
| backward | 8192 | 2048 | 262144 | 4 | 656.4 | 371.3 | 1.77x | 2.26 |
1.15 | 0.51x |
| backward | 8192 | 2048 | 262144 | 8 | 604.6 | 329.0 | 1.84x | 1.94 |
0.71 | 0.37x |
| backward | 8192 | 2048 | 262144 | 16 | 414.0 | 306.1 | 1.35x | 0.22 |
0.56 | 2.55x |
| backward | 8192 | 4096 | 262144 | 1 | 777.4 | 512.1 | 1.52x | 6.19 |
6.45 | 1.04x |
| backward | 8192 | 4096 | 262144 | 2 | 765.1 | 505.7 | 1.51x | 3.94 |
3.48 | 0.88x |
| backward | 8192 | 4096 | 262144 | 4 | 751.1 | 500.6 | 1.50x | 2.82 |
1.99 | 0.71x |
| backward | 8192 | 4096 | 262144 | 8 | 714.0 | 447.2 | 1.60x | 2.26 |
1.18 | 0.52x |
| backward | 8192 | 4096 | 262144 | 16 | 548.6 | 421.2 | 1.30x | 0.41 |
0.84 | 2.05x |
| backward | 16384 | 4096 | 262144 | 1 | 775.8 | 520.5 | 1.49x | 6.32 |
6.87 | 1.09x |
| backward | 16384 | 4096 | 262144 | 2 | 763.9 | 516.4 | 1.48x | 4.07 |
3.93 | 0.97x |
| backward | 16384 | 4096 | 262144 | 4 | 743.2 | 509.9 | 1.46x | 2.94 |
2.46 | 0.84x |
| backward | 16384 | 4096 | 262144 | 8 | 713.2 | 458.0 | 1.56x | 2.38 |
1.58 | 0.66x |
| backward | 16384 | 4096 | 262144 | 16 | 552.6 | 429.2 | 1.29x | 0.53 |
1.28 | 2.42x |
| backward | 8192 | 2048 | 524288 | 8 | 655.9 | 369.7 | 1.77x | 2.26 |
1.15 | 0.51x |
| backward | 8192 | 4096 | 524288 | 8 | 748.4 | 499.0 | 1.50x | 2.82 |
1.99 | 0.71x |
| backward | 8192 | 2048 | 1048576 | 16 | 626.5 | 366.9 | 1.71x | 0.69 |
1.15 | 1.67x |
| backward | 8192 | 4096 | 1048576 | 16 | 726.1 | 494.5 | 1.47x | 1.25 |
1.99 | 1.59x |
| forward | 8192 | 2048 | 65536 | 1 | 748.2 | 361.7 | 2.07x | 0.28 |
0.29 | 1.04x |
| forward | 8192 | 4096 | 65536 | 1 | 788.5 | 417.4 | 1.89x | 0.56 |
0.57 | 1.02x |
| forward | 2048 | 4096 | 131072 | 1 | 787.7 | 460.7 | 1.71x | 1.02 |
1.02 | 1.00x |
| forward | 2048 | 4096 | 131072 | 2 | 786.9 | 439.2 | 1.79x | 0.98 |
0.52 | 0.53x |
| forward | 2048 | 4096 | 131072 | 4 | 769.3 | 408.7 | 1.88x | 0.73 |
0.27 | 0.37x |
| forward | 2048 | 4096 | 131072 | 8 | 657.7 | 351.8 | 1.87x | 0.60 |
0.14 | 0.23x |
| forward | 2048 | 4096 | 131072 | 16 | 460.3 | 253.4 | 1.82x | 0.46 |
0.08 | 0.17x |
| forward | 8192 | 2048 | 131072 | 1 | 747.7 | 226.1 | 3.31x | 0.53 |
0.55 | 1.04x |
| forward | 8192 | 2048 | 131072 | 2 | 743.4 | 199.3 | 3.73x | 0.74 |
0.29 | 0.39x |
| forward | 8192 | 2048 | 131072 | 4 | 722.7 | 343.3 | 2.11x | 0.62 |
0.16 | 0.26x |
| forward | 8192 | 2048 | 131072 | 8 | 676.0 | 321.2 | 2.10x | 0.56 |
0.10 | 0.18x |
| forward | 8192 | 2048 | 131072 | 16 | 591.1 | 251.8 | 2.35x | 0.45 |
0.06 | 0.13x |
| forward | 8192 | 4096 | 131072 | 1 | 789.0 | 224.1 | 3.52x | 1.06 |
1.08 | 1.02x |
| forward | 8192 | 4096 | 131072 | 2 | 786.3 | 411.8 | 1.91x | 1.03 |
0.57 | 0.55x |
| forward | 8192 | 4096 | 131072 | 4 | 779.4 | 401.0 | 1.94x | 0.78 |
0.32 | 0.41x |
| forward | 8192 | 4096 | 131072 | 8 | 756.8 | 388.6 | 1.95x | 0.65 |
0.19 | 0.29x |
| forward | 8192 | 4096 | 131072 | 16 | 701.1 | 333.9 | 2.10x | 0.51 |
0.13 | 0.25x |
| forward | 16384 | 4096 | 131072 | 1 | 757.4 | 372.8 | 2.03x | 1.13 |
1.16 | 1.03x |
| forward | 16384 | 4096 | 131072 | 2 | 754.9 | 365.0 | 2.07x | 1.09 |
0.64 | 0.59x |
| forward | 16384 | 4096 | 131072 | 4 | 744.7 | 358.3 | 2.08x | 0.84 |
0.38 | 0.45x |
| forward | 16384 | 4096 | 131072 | 8 | 719.9 | 331.1 | 2.17x | 0.72 |
0.25 | 0.35x |
| forward | 16384 | 4096 | 131072 | 16 | 692.2 | 327.0 | 2.12x | 0.57 |
0.19 | 0.33x |
| forward | 2048 | 4096 | 262144 | 1 | 791.1 | 468.7 | 1.69x | 2.02 |
2.02 | 1.00x |
| forward | 2048 | 4096 | 262144 | 2 | 783.1 | 422.3 | 1.85x | 1.48 |
1.02 | 0.69x |
| forward | 2048 | 4096 | 262144 | 4 | 790.9 | 437.6 | 1.81x | 0.98 |
0.52 | 0.53x |
| forward | 2048 | 4096 | 262144 | 8 | 754.4 | 404.2 | 1.87x | 0.73 |
0.27 | 0.37x |
| forward | 2048 | 4096 | 262144 | 16 | 609.5 | 327.5 | 1.86x | 0.52 |
0.14 | 0.27x |
| forward | 8192 | 2048 | 262144 | 1 | 747.5 | 372.3 | 2.01x | 1.03 |
1.06 | 1.03x |
| forward | 8192 | 2048 | 262144 | 2 | 745.5 | 360.5 | 2.07x | 0.99 |
0.55 | 0.56x |
| forward | 8192 | 2048 | 262144 | 4 | 740.7 | 353.1 | 2.10x | 0.74 |
0.29 | 0.39x |
| forward | 8192 | 2048 | 262144 | 8 | 714.4 | 341.1 | 2.09x | 0.62 |
0.16 | 0.26x |
| forward | 8192 | 2048 | 262144 | 16 | 648.8 | 309.6 | 2.10x | 0.48 |
0.10 | 0.21x |
| forward | 8192 | 4096 | 262144 | 1 | 788.3 | 208.0 | 3.79x | 2.06 |
2.09 | 1.01x |
| forward | 8192 | 4096 | 262144 | 2 | 787.6 | 233.2 | 3.38x | 1.53 |
1.08 | 0.71x |
| forward | 8192 | 4096 | 262144 | 4 | 789.0 | 411.4 | 1.92x | 1.03 |
0.57 | 0.55x |
| forward | 8192 | 4096 | 262144 | 8 | 775.0 | 401.4 | 1.93x | 0.78 |
0.32 | 0.41x |
| forward | 8192 | 4096 | 262144 | 16 | 741.9 | 379.7 | 1.95x | 0.57 |
0.19 | 0.33x |
| forward | 16384 | 4096 | 262144 | 1 | 755.3 | 374.1 | 2.02x | 2.13 |
2.19 | 1.03x |
| forward | 16384 | 4096 | 262144 | 2 | 755.0 | 367.4 | 2.05x | 1.59 |
1.16 | 0.73x |
| forward | 16384 | 4096 | 262144 | 4 | 753.8 | 364.0 | 2.07x | 1.09 |
0.64 | 0.59x |
| forward | 16384 | 4096 | 262144 | 8 | 742.1 | 357.1 | 2.08x | 0.84 |
0.38 | 0.45x |
| forward | 16384 | 4096 | 262144 | 16 | 709.3 | 344.8 | 2.06x | 0.64 |
0.25 | 0.39x |
| forward | 8192 | 2048 | 524288 | 8 | 736.9 | 351.5 | 2.10x | 0.74 |
0.29 | 0.39x |
| forward | 8192 | 4096 | 524288 | 8 | 782.6 | 410.2 | 1.91x | 1.03 |
0.57 | 0.55x |
| forward | 8192 | 2048 | 1048576 | 16 | 729.1 | 347.2 | 2.10x | 0.66 |
0.29 | 0.44x |
| forward | 8192 | 4096 | 1048576 | 16 | 776.0 | 407.4 | 1.90x | 0.95 |
0.57 | 0.60x |

## Liger entropy-on results

| Pass | M | H | V | TP | p50 ms | TF/GPU | Aggregate TF | GiB | Backend
| split-K |
|---|---:|---:|---:|---:|---:|---:|---:|---:|---|---:|
| backward | 2048 | 4096 | 131072 | 1 | 8.589 | 768.1 | 768.1 | 3.85 |
nvls | 1 |
| backward | 2048 | 4096 | 131072 | 2 | 4.395 | 750.4 | 1500.9 | 2.72 |
nvls | 1 |
| backward | 2048 | 4096 | 131072 | 4 | 2.311 | 713.7 | 2854.6 | 2.16 |
nvls | 1 |
| backward | 2048 | 4096 | 131072 | 8 | 1.291 | 638.5 | 5108.2 | 1.88 |
nvls | 1 |
| backward | 2048 | 4096 | 131072 | 16 | 1.066 | 386.9 | 6190.2 | 0.17 |
hierarchical | 1 |
| backward | 2048 | 4096 | 262144 | 1 | 16.842 | 783.4 | 783.4 | 6.10 |
nvls | 1 |
| backward | 2048 | 4096 | 262144 | 2 | 8.570 | 769.8 | 1539.6 | 3.85 |
nvls | 1 |
| backward | 2048 | 4096 | 262144 | 4 | 4.366 | 755.6 | 3022.2 | 2.72 |
nvls | 1 |
| backward | 2048 | 4096 | 262144 | 8 | 2.306 | 715.1 | 5720.5 | 2.16 |
nvls | 1 |
| backward | 2048 | 4096 | 262144 | 16 | 1.574 | 523.8 | 8380.4 | 0.31 |
hierarchical | 1 |
| backward | 8192 | 2048 | 65536 | 1 | 10.008 | 659.2 | 659.2 | 2.26 |
nvls | 2 |
| backward | 8192 | 2048 | 131072 | 1 | 19.086 | 691.3 | 691.3 | 2.88 |
nvls | 2 |
| backward | 8192 | 2048 | 131072 | 2 | 10.421 | 633.0 | 1266.1 | 2.26 |
nvls | 2 |
| backward | 8192 | 2048 | 131072 | 4 | 5.759 | 572.8 | 2291.2 | 1.94 |
nvls | 2 |
| backward | 8192 | 2048 | 131072 | 8 | 3.467 | 475.7 | 3805.4 | 1.79 |
nvls | 2 |
| backward | 8192 | 2048 | 131072 | 16 | 3.141 | 262.6 | 4201.1 | 0.14 |
hierarchical | 2 |
| backward | 8192 | 2048 | 262144 | 1 | 37.342 | 706.7 | 706.7 | 4.13 |
nvls | 2 |
| backward | 8192 | 2048 | 262144 | 2 | 19.290 | 684.0 | 1368.0 | 2.88 |
nvls | 2 |
| backward | 8192 | 2048 | 262144 | 4 | 10.105 | 652.8 | 2611.3 | 2.26 |
nvls | 2 |
| backward | 8192 | 2048 | 262144 | 8 | 5.454 | 604.8 | 4838.6 | 1.94 |
nvls | 2 |
| backward | 8192 | 2048 | 262144 | 16 | 3.999 | 412.4 | 6598.0 | 0.22 |
hierarchical | 2 |
| backward | 8192 | 2048 | 524288 | 8 | 10.070 | 655.1 | 5241.0 | 2.26 |
nvls | 2 |
| backward | 8192 | 2048 | 1048576 | 16 | 10.573 | 624.0 | 9983.4 | 0.69
| hierarchical | 2 |
| backward | 8192 | 4096 | 65536 | 1 | 17.420 | 757.4 | 757.4 | 2.82 |
nvls | 1 |
| backward | 8192 | 4096 | 131072 | 1 | 34.601 | 762.6 | 762.6 | 3.94 |
nvls | 1 |
| backward | 8192 | 4096 | 131072 | 2 | 17.774 | 742.3 | 1484.7 | 2.82 |
nvls | 1 |
| backward | 8192 | 4096 | 131072 | 4 | 9.214 | 716.0 | 2863.9 | 2.26 |
nvls | 1 |
| backward | 8192 | 4096 | 131072 | 8 | 5.153 | 640.2 | 5121.4 | 1.97 |
nvls | 1 |
| backward | 8192 | 4096 | 131072 | 16 | 3.963 | 416.2 | 6658.9 | 0.27 |
hierarchical | 1 |
| backward | 8192 | 4096 | 262144 | 1 | 67.858 | 777.8 | 777.8 | 6.19 |
nvls | 1 |
| backward | 8192 | 4096 | 262144 | 2 | 34.470 | 765.5 | 1531.1 | 3.94 |
nvls | 1 |
| backward | 8192 | 4096 | 262144 | 4 | 17.560 | 751.4 | 3005.5 | 2.82 |
nvls | 1 |
| backward | 8192 | 4096 | 262144 | 8 | 9.209 | 716.4 | 5731.1 | 2.26 |
nvls | 1 |
| backward | 8192 | 4096 | 262144 | 16 | 6.014 | 548.5 | 8775.6 | 0.41 |
hierarchical | 1 |
| backward | 8192 | 4096 | 524288 | 8 | 17.647 | 747.7 | 5981.2 | 2.82 |
nvls | 1 |
| backward | 8192 | 4096 | 1048576 | 16 | 18.181 | 725.7 | 11611.2 |
1.25 | hierarchical | 1 |
| backward | 16384 | 4096 | 131072 | 1 | 69.235 | 762.3 | 762.3 | 4.07 |
nvls | 1 |
| backward | 16384 | 4096 | 131072 | 2 | 35.679 | 739.6 | 1479.2 | 2.94
| nvls | 1 |
| backward | 16384 | 4096 | 131072 | 4 | 18.499 | 713.2 | 2853.0 | 2.38
| nvls | 1 |
| backward | 16384 | 4096 | 131072 | 8 | 10.382 | 635.4 | 5083.3 | 2.10
| nvls | 1 |
| backward | 16384 | 4096 | 131072 | 16 | 7.832 | 421.1 | 6738.3 | 0.39
| hierarchical | 1 |
| backward | 16384 | 4096 | 262144 | 1 | 136.168 | 775.2 | 775.2 | 6.32
| nvls | 1 |
| backward | 16384 | 4096 | 262144 | 2 | 69.172 | 763.0 | 1526.0 | 4.07
| nvls | 1 |
| backward | 16384 | 4096 | 262144 | 4 | 35.475 | 743.9 | 2975.4 | 2.94
| nvls | 1 |
| backward | 16384 | 4096 | 262144 | 8 | 18.444 | 715.4 | 5722.9 | 2.38
| nvls | 1 |
| backward | 16384 | 4096 | 262144 | 16 | 11.933 | 552.8 | 8845.5 | 0.53
| hierarchical | 1 |
| forward | 2048 | 4096 | 131072 | 1 | 2.842 | 773.9 | 773.9 | 1.02 |
nvls | |
| forward | 2048 | 4096 | 131072 | 2 | 1.392 | 789.8 | 1579.6 | 0.98 |
nvls | |
| forward | 2048 | 4096 | 131072 | 4 | 0.720 | 763.8 | 3055.3 | 0.73 |
nvls | |
| forward | 2048 | 4096 | 131072 | 8 | 0.424 | 648.0 | 5183.8 | 0.60 |
nvls | |
| forward | 2048 | 4096 | 131072 | 16 | 0.303 | 453.9 | 7262.3 | 0.46 |
hierarchical | |
| forward | 2048 | 4096 | 262144 | 1 | 5.612 | 783.6 | 783.6 | 2.02 |
nvls | |
| forward | 2048 | 4096 | 262144 | 2 | 2.820 | 779.7 | 1559.3 | 1.48 |
nvls | |
| forward | 2048 | 4096 | 262144 | 4 | 1.395 | 788.4 | 3153.5 | 0.98 |
nvls | |
| forward | 2048 | 4096 | 262144 | 8 | 0.736 | 746.5 | 5972.1 | 0.73 |
nvls | |
| forward | 2048 | 4096 | 262144 | 16 | 0.456 | 603.4 | 9653.6 | 0.52 |
hierarchical | |
| forward | 8192 | 2048 | 65536 | 1 | 2.957 | 743.6 | 743.6 | 0.28 |
nvls | |
| forward | 8192 | 2048 | 131072 | 1 | 5.899 | 745.5 | 745.5 | 0.53 |
nvls | |
| forward | 8192 | 2048 | 131072 | 2 | 2.973 | 739.7 | 1479.4 | 0.74 |
nvls | |
| forward | 8192 | 2048 | 131072 | 4 | 1.533 | 717.4 | 2869.6 | 0.62 |
nvls | |
| forward | 8192 | 2048 | 131072 | 8 | 0.815 | 674.2 | 5393.4 | 0.56 |
nvls | |
| forward | 8192 | 2048 | 131072 | 16 | 0.473 | 580.9 | 9294.3 | 0.45 |
hierarchical | |
| forward | 8192 | 2048 | 262144 | 1 | 11.777 | 746.9 | 746.9 | 1.03 |
nvls | |
| forward | 8192 | 2048 | 262144 | 2 | 5.918 | 743.1 | 1486.2 | 0.99 |
nvls | |
| forward | 8192 | 2048 | 262144 | 4 | 2.988 | 736.0 | 2944.0 | 0.74 |
nvls | |
| forward | 8192 | 2048 | 262144 | 8 | 1.548 | 710.5 | 5683.6 | 0.62 |
nvls | |
| forward | 8192 | 2048 | 262144 | 16 | 0.854 | 643.8 | 10300.6 | 0.48 |
hierarchical | |
| forward | 8192 | 2048 | 524288 | 8 | 3.002 | 732.6 | 5860.9 | 0.74 |
nvls | |
| forward | 8192 | 2048 | 1048576 | 16 | 3.045 | 722.1 | 11553.5 | 0.66
| hierarchical | |
| forward | 8192 | 4096 | 65536 | 1 | 5.596 | 785.9 | 785.9 | 0.56 |
nvls | |
| forward | 8192 | 4096 | 131072 | 1 | 11.182 | 786.7 | 786.7 | 1.06 |
nvls | |
| forward | 8192 | 4096 | 131072 | 2 | 5.575 | 788.9 | 1577.7 | 1.03 |
nvls | |
| forward | 8192 | 4096 | 131072 | 4 | 2.828 | 777.7 | 3110.7 | 0.78 |
nvls | |
| forward | 8192 | 4096 | 131072 | 8 | 1.458 | 754.1 | 6033.1 | 0.65 |
nvls | |
| forward | 8192 | 4096 | 131072 | 16 | 0.789 | 697.0 | 11151.2 | 0.51 |
hierarchical | |
| forward | 8192 | 4096 | 262144 | 1 | 22.375 | 786.2 | 786.2 | 2.06 |
nvls | |
| forward | 8192 | 4096 | 262144 | 2 | 11.141 | 789.5 | 1579.1 | 1.53 |
nvls | |
| forward | 8192 | 4096 | 262144 | 4 | 5.594 | 786.2 | 3144.9 | 1.03 |
nvls | |
| forward | 8192 | 4096 | 262144 | 8 | 2.833 | 776.2 | 6209.9 | 0.78 |
nvls | |
| forward | 8192 | 4096 | 262144 | 16 | 1.492 | 737.0 | 11791.4 | 0.57 |
hierarchical | |
| forward | 8192 | 4096 | 524288 | 8 | 5.645 | 779.2 | 6233.2 | 1.03 |
nvls | |
| forward | 8192 | 4096 | 1048576 | 16 | 5.677 | 774.7 | 12395.4 | 0.95
| hierarchical | |
| forward | 16384 | 4096 | 131072 | 1 | 23.312 | 754.6 | 754.6 | 1.13 |
nvls | |
| forward | 16384 | 4096 | 131072 | 2 | 11.685 | 752.8 | 1505.6 | 1.09 |
nvls | |
| forward | 16384 | 4096 | 131072 | 4 | 5.932 | 741.4 | 2965.4 | 0.84 |
nvls | |
| forward | 16384 | 4096 | 131072 | 8 | 3.067 | 716.9 | 5735.6 | 0.72 |
nvls | |
| forward | 16384 | 4096 | 131072 | 16 | 1.600 | 687.0 | 10992.6 | 0.57
| hierarchical | |
| forward | 16384 | 4096 | 262144 | 1 | 46.752 | 752.6 | 752.6 | 2.13 |
nvls | |
| forward | 16384 | 4096 | 262144 | 2 | 23.374 | 752.6 | 1505.3 | 1.59 |
nvls | |
| forward | 16384 | 4096 | 262144 | 4 | 11.702 | 751.7 | 3006.6 | 1.09 |
nvls | |
| forward | 16384 | 4096 | 262144 | 8 | 5.946 | 739.6 | 5917.2 | 0.84 |
nvls | |
| forward | 16384 | 4096 | 262144 | 16 | 3.116 | 705.7 | 11291.1 | 0.64
| hierarchical | |

## Resource results

- compact dZ handoff: 168 registers, zero spills, 11 barriers; dZ-only
p50 11.839 ms entropy-off and 11.898 ms entropy-on in the isolated
validation build.
- combined split-K1 and split-K2: 168 registers, zero spills, 7
barriers, 230,400 B dynamic shared memory; TPR1/2/4 validated.
- dX/dW mainloops are within 10% of tuned cuBLASLt; dX exceeds cuBLAS on
multicast-friendly shapes and dW is approximately 4% behind.

---------

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-08-31 20:15:17 +00:00
jiahao su f65ae6576f Update the ascend npu test repository link (#1411)
### Summary

Transfer the ascend npu test of liger-kernel from the individual
warehouse to the Ascend organization.
2026-08-31 18:57:35 +00:00
Rain 777799588a fix(mlp): keep liger_kernel.ops importable without TensorDescriptor (#1421)
## Summary

- `LigerMLP` (#1357) does a hard `from triton.tools.tensor_descriptor
import TensorDescriptor`. That module first shipped in official Triton
**3.4.0**, so any earlier / vendor build fails while importing
`liger_kernel.ops`.
- `ops/__init__.py` always loads default kernels before backend overlay,
so one missing import takes down every op (RMSNorm, Softmax, RoPE, …),
not just MLP.
- Make the import optional. `LigerMLPFunction.forward` raises a clear
`RuntimeError` if the API is missing. GPU / Triton >= 3.4 is unchanged.
- Skip `test/transformers/test_mlp.py` when
`triton.tools.tensor_descriptor` is not installed.

## Testing Done

- `pytest test/transformers/ --collect-only` on a Triton 3.2 build:
collection succeeds (no longer 33 collection errors).
- `test_mlp.py` is skipped when the API is absent; other transformers
tests collect and run as before.
2026-08-30 13:02:20 +00:00
Vaibhav JindalandCopilot 00e72b86ca ci: schedule NVIDIA tests by cadence/GPU with a hard 3-job concurrency cap (#1415)
## Summary

Reorganizes the **NVIDIA GPU** workflow into a cadence- and GPU-aware
schedule, and adds the Modal functions to back it. Follows up #1410
(which added the per-DSL test folders and `make test-cutile` /
`test-cutedsl` / `test-cute` targets).

### Schedule

**Daily** (cron `24 9 * * *`):
| Suite | GPU | Modal fn |
|---|---|---|
| general correctness | H100 | `liger_correctness_tests` |
| cuTile | H100 | `liger_cutile_tests` (new) |
| cuTeDSL — SM90 half | H100 | `liger_cutedsl_tests_h100` (new) |
| cuTeDSL — SM100 half | B200 | `liger_cutedsl_tests_b200` (new) |

cuTeDSL runs on **both** archs on purpose: its SM90 tests only run on
H100 and its SM100 tests only run on B200 — each arch self-skips the
other half, so both are needed for full coverage.

**Weekly** (cron `24 13 * * 0`, Sundays, H100): `convergence` (latest
tf) + oldest-supported-tf **correctness** and **convergence**.

**Not scheduled:** `cute` (fused-MoE) tests, for now.

### The ≤3-concurrent fix

A 4th simultaneous Modal GPU call was being rejected (one job died
instantly at ~0.7 min). This is now capped two ways:
- **`strategy.max-parallel: 3`** on each matrix job (daily's 4 entries
run 3-then-1; weekly's 3 run together).
- **Workflow-level `concurrency: { group: <workflow>,
cancel-in-progress: false }`** serializes whole runs, so the daily and
weekly runs never overlap.

Only **one** cadence job activates per trigger (the job-level `if`
selects daily vs weekly from `github.event.schedule` / the
`workflow_dispatch` `suite` input), so at most **3** GPU jobs ever run
together. The `if` deliberately uses only the `github` context — **not**
the `matrix` context, which isn't available in a job-level `if`.

### Manual runs
`workflow_dispatch` gains a `suite` input (`daily` / `weekly`, default
`daily`).


## Notes / caveats
- **Convergence is currently red** on `main`'s scheduled runs
(pre-existing; correctness passes). Weekly scheduling is fine but expect
failures until that's fixed separately.
- **DSL suites run in CI for the first time** — the cuTeDSL-B200 (SM100)
path and the cuTile routing test may surface latent failures on first
run.
- **B200 access**: uses `gpu="B200"` on Modal — confirm the token/plan
has B200 capacity.
- `cuTile` installs the `cutile-tileiras` extra so `import cuda.tile`
resolves (otherwise those tests importorskip).
- Validated locally with a YAML parser + a gating simulation (exactly
one job per trigger) + modal-fn-name cross-check; `actionlint` wasn't
installable in the dev sandbox, and GPU suites can't run there.

Co-authored-by: Copilot <[email protected]>
2026-08-28 04:53:37 +00:00
Arup DeandCopilot aae4bef01e perf(flce): widen the transient-logits chunk-memory budget constant to C=16 (#1414)
# perf(flce): widen the transient-logits chunk-memory budget constant to
C=16

**Head:** `arde171:arde/optos-liger-flce-chunk` → **Base:**
`linkedin/Liger-Kernel:main` (`4310065`)
One file, `+18 / -1`:
`src/liger_kernel/ops/fused_linear_cross_entropy.py`.

## Summary

The chunked `fused_linear_cross_entropy` sizes its token chunk from a
memory budget constant. The historical formula `inc_factor = cdiv(V, H)`
corresponds to a budget of **C=1** — the tightest possible
transient-logits footprint (~`BT x H`). At real LLM vocab shapes this
fractures the batch into many tiny chunks and makes the op
**launch-bound**. Widening the budget to **C=16 x BT x H** collapses the
chunk loop with a still-bounded, modest memory increase.

```diff
+_CHUNK_MEM_CONST = 16
...
-    inc_factor = triton.cdiv(V, H)                       # C=1: a memory floor, not a perf target
+    inc_factor = triton.cdiv(V, _CHUNK_MEM_CONST * H)    # widen transient-logits budget to C x BT x H
     chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))
+    chunk_size = min(chunk_size, BT)                     # a single chunk covers BT when the budget allows
```

At llama-3 shapes (`V=128256, H=4096`), C=1 gives `inc_factor=32` ->
`chunk_size=BT/32` -> 32 loop iterations; C=16 collapses that.

## Performance (B200, full fwd+bwd, bf16)

| shape | C=1 | C=16 | speedup |
|---|---|---|---|
| V=128256 H=4096 BT=8192 | 123.4 ms | 24.7 ms | **5.0x** |
| V=128256 H=4096 BT=4096 | 119.3 ms | 15.1 ms | **7.9x** |
| V=128256 H=4096 BT=1024 | 116.9 ms | 9.07 ms | **12.9x** |
| V=32000 H=4096 BT=8192 | 11.0 ms | 5.4 ms | 2.0x (flat for C>=8) |

C=16 is the knee -- gains flatten beyond it at BT>=4096. Memory stays
bounded at `O(C x BT x H)`.

## Correctness

Numerics are **partition-invariant per row**: the loss is
**bit-identical** across chunk count C; the input-gradient differs only
at the ~1e-5 GEMM reduction-order level, and the weight-grad only in
`addmm` chunk-accumulation order (the same class as the existing
`fused_linear_jsd` chunk-memory idiom). Same constant on all devices.

**Tests** -- `test/transformers/test_fused_linear_cross_entropy.py`:
- **B300 (sm_103): 141 passed**
- **H200 (sm_90): 141 passed**

## Scope

One constant + a `min(chunk_size, BT)` clamp in the Triton composition.
No kernel or API change; no new backend. Purely a chunk-geometry perf
tuning.

Co-authored-by: Copilot <[email protected]>

Co-authored-by: Copilot <[email protected]>
2026-08-27 21:37:38 +00:00
Arup DeandCopilot 50607471ea fix(fused_moe): correct input-gradient (dx) on Blackwell/B300 (#1413)
# PR: Fix fused-MoE input gradient (dx) on Blackwell/B300

**Head:** `arde171:arde/optos-liger-moe-dx-fix` → **Base:**
`linkedin/Liger-Kernel:main` (`4310065`)

---

## Summary

`LigerFusedMoEFunction` returns an **incorrect gradient with respect to
its input** (`x.grad`) on **B300 / sm_103 (Blackwell)**. The error is
~50–64% relative-mean and **silently corrupts MoE training** (no NaN, no
error) — the input gradient feeds every layer below the MoE block.
Forward output, `gate_up_proj.grad` (dW1), and `down_proj.grad` (dW2)
are all correct; only `dx` is affected.

The bug is **Blackwell-specific**: the identical code is correct on H200
/ sm_90, which is why it isn't caught by upstream CI running on
Hopper/Ampere.

## Root cause

`_moe_bwd_dX_expanded_kernel` computes `dx_expanded = d_gate @ W1_gateᵀ
+ d_up @ W1_upᵀ` by folding **both** matmuls into **one shared
accumulator**, with two `tl.dot(..., acc=acc)` calls per loop iteration:

```python
acc = tl.zeros(...)
for k in range(0, I, BLOCK_K):
    acc = tl.dot(d_gate, w_gate, acc=acc)   # gate half
    acc = tl.dot(d_up,   w_up,   acc=acc)   # up half  → same acc, same iteration
```

On B300/sm_103 (Triton in torch 2.13/cu130) this pattern miscompiles and
drops/corrupts one of the two contributions. Because SwiGLU's gate and
up terms are comparable in magnitude, losing one produces the observed
~50–64% error. The error is:
- **top-k independent** (the router weight is already folded into
`d_pre_act` upstream — even K=1 is wrong),
- **only visible with ≥2 experts** (with a single expert the failure
cancels in the single contiguous tile),
- **config-dependent** (magnitude shifts with `BLOCK_N`/`BLOCK_K`),
which is the fingerprint of a tiling/accumulation miscompile rather than
numerical tolerance.

The forward counterpart `_fused_up_proj_swiglu_kernel` already computes
the gate/up split with **two independent accumulators** (`acc_gate`,
`acc_up`) and is correct.

## Fix

Accumulate the gate and up contributions into **separate fp32 tiles**
and sum once after the loop — mirroring the forward kernel.
Mathematically identical; the two independent recurrences stay
simultaneously live, so the backend cannot reassociate them back into
the hazardous single-accumulator form.

```python
acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
acc_up   = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, I, BLOCK_K):
    ...
    acc_gate = tl.dot(d_gate, w_gate, acc=acc_gate)
    ...
    acc_up   = tl.dot(d_up,   w_up,   acc=acc_up)
acc = acc_gate + acc_up
```

One file, `+15 / −4`.

## Validation

Measured on **B300 SXM6 (sm_103)** and **H200 (sm_90)**, torch
2.13.0+cu130.

**dx microbench** (trusted-autograd + `transformers` MixtralExperts
references, fp32, K=1/2/4 × 3 seeds):

| | before | after |
|---|---|---|
| `dx` rel-mean (worst) | 52–64% | **0.22%** (tf32 level, matches
fwd/dW1/dW2) |

**Test suite (B300), before → after:**

| test | before | after |
|---|---|---|
| `test/transformers/test_fused_moe.py` (fp32+bf16) | 5+ fp32 fail |
**19 passed** |
| `test/transformers/test_swiglu.py::test_correctness_mixtralexperts` |
3 fail on `[x.grad]` (`max_abs≈1.5e4`) | **4 passed** |
| H200 (both suites) | already passed | passes (Blackwell-specific) |

Pre-existing, unrelated: `mini_qwen3_moe`/`qwen3_5_moe` bf16 convergence
fails identically with and without this change — a known
machine-dependent flaky test documented in the test file (`# TODO(tcc):
... diverges on ci test (A10G) ... diverges on h100`).

## Notes

- A single contiguous-`2I` GEMM (one dot per iteration) is an equivalent
alternative and a good **follow-up performance** change; it should be
validated separately with B300 perf + correctness. This PR intentionally
makes the minimal correctness fix that mirrors the proven forward
pattern.
- Upstream CI cannot reproduce this without Blackwell runners; the fix
is validated on B300 hardware and is correct by construction (identical
to the forward accumulator structure).

Co-authored-by: Copilot <[email protected]>

Co-authored-by: Copilot <[email protected]>
2026-08-27 21:37:20 +00:00
df40cf97d4 Isolate SM90 and SM100 MoE kernel paths (#1412)
## Summary

- duplicate the complete MoE GEMM dependency graph for SM90 and SM100
behind compile-time entrypoint selectors
- isolate architecture-specific math, TMEM, barriers, dispatch menus,
and tuning-table wrappers so neither build depends on the opposite
architecture
- preserve mainline M128/NC4 communication geometry and add tests
enforcing the boundary

## Validation

- `make checkstyle`
- architecture boundary tests: 3 passed
- B200 distributed MoE suite: 8 passed
- H100 distributed MoE suite: 8 passed
- SM100: full SASS and resource report identical to mainline
- SM90: normalized SASS and resource report identical to mainline; only
compiler-generated internal NVSHMEM symbol hashes differ
- paired performance (difference vs mainline):
  - SM100 Qwen3-30B-A3B: +0.11% total
  - SM100 Mixtral-8x7B: -0.16% total
  - SM90 Qwen3-30B-A3B: -0.20% total
  - SM90 Mixtral-8x7B: -0.22% total

Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-08-27 17:44:06 +00:00
Vaibhav JindalandCopilot 38c6a455ea test: organize DSL tests into per-DSL folders (cutedsl/cutile/cute) (#1410)
## Summary

The three opt-in GPU DSL backends — **cutedsl**, **cutile**, and
**cute** (fused-MoE) — had their tests inconsistently placed: a partial
`test/cutedsl/` folder existed, but most DSL tests were scattered in
`test/transformers/`. This PR gives each DSL its **own test folder**,
mirroring the existing `test/cutedsl/` convention (drop the DSL prefix
from filenames), and wires up matching `make` targets.

One minor change in cron schedule for nvidia-ci.

## Changes

### cutedsl → `test/cutedsl/`
- Moved `test_cutedsl_fused_linear_cross_entropy_sm90` →
`test_fused_linear_cross_entropy_sm90`
- Moved `test_cutedsl_rms_norm_fastpath` → `test_rms_norm_fastpath`
- Moved `test_cutedsl_rope` → `test_rope`
- Moved `test_swiglu_cutedsl` → `test_swiglu`
- **Deleted** `test_cutedsl_rms_norm` — it was a duplicate *subset* of
the canonical `test/cutedsl/test_rms_norm.py` (same two test-function
names; the canonical file is a strict superset with the extra Blackwell
fast-path tests and stricter mode-specific bf16 dW tolerances). Deleting
avoids a filename + test-name collision; the canonical file is left
untouched.
- Fixed a stale docstring path in `test/cutedsl/test_backend.py`.

### cutile → new `test/cutile/`
- Moved `test_cutile_backend` → `test_backend`,
`test_cutile_fused_linear_cross_entropy_sm90` →
`test_fused_linear_cross_entropy_sm90`, `test_cutile_rms_norm` →
`test_rms_norm`, `test_cutile_rope` → `test_rope`, `test_cutile_swiglu`
→ `test_swiglu`.
- **Split** the 7 env-gated `test_cutile_*` tests out of the shared
`test/transformers/test_fused_linear_scaled_cross_entropy.py` into
`test/cutile/test_fused_linear_scaled_cross_entropy.py`. The
backend-agnostic frontend/dispatcher tests stay in `test/transformers/`.

### cute (fused-MoE) → new `test/cute/`
- Moved `test_cute_moe_autograd` → `test_moe_autograd` and `test_moe` →
`test_moe`.
- `test_fused_moe.py` is the general Triton MoE test and **stays** in
`test/transformers/`.

### Makefile
- `test:` and `test-verbose:` now `--ignore=test/cutile` and
`--ignore=test/cute` (and `test-verbose` also ignores `test/cutedsl` for
consistency), so DSL tests run only via their dedicated targets.
- New **`test-cutile`** (`LIGER_KERNEL_IMPL=cutile`, mirrors
`test-cutedsl`).
- New **`test-cute`** — intentionally does **not** set
`LIGER_KERNEL_IMPL=cute`, because selecting it would force `import
liger_kernel.ops` to load the separate `liger_cute_kernels` (lck) wheel
and break the always-on registration/import-guard tests.

## Notes / decisions
- All moves use `git mv` to preserve history. The `test.utils` import
blocks in the moved rope/swiglu files were re-sorted by ruff because
`test` is now first-party inside the packaged DSL folders.
- **Coverage shift (intended):** a few always-on pure-Python tests
(cutedsl rms_norm fastpath, cute registration/import-guard) previously
ran under `make test`; they now run via the dedicated `test-*` targets,
consistent with the one-folder-per-DSL model.
- No new basename collisions: the DSL folders all have `__init__.py`
(unique package-qualified module names) and are ignored by `make test`.

## Validation
- `ruff check` and `ruff format --check` pass across the whole `test/`
tree.
- All affected files pass `py_compile`; `make -n` dry-runs verified for
the new targets.
- ⚠️ GPU/Modal test execution (`make test-cutedsl` / `test-cutile` /
`test-cute`) was **not** run in this environment (no CUDA/torch
available). These should be validated on a GPU host with the
corresponding extras installed.

---------

Co-authored-by: Copilot <[email protected]>
2026-08-27 05:37:31 +00:00
0515c5aa2c Fix broken star history chart (#1385)
The star history chart on the README and docs was broken due to GitHub
stargazer API restrictions. Update the chart to use a working data
source that requires no API token so the chart renders correctly again.

Co-authored-by: Dessalines39394 <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>
2026-08-27 01:16:10 +00:00
Vaibhav JindalandCopilot 0406e4ddc0 ci: run NVIDIA GPU tests on manual dispatch and daily cron only (#1409)
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->

Remove push, pull_request, and merge_group triggers from the NVIDIA GPU
workflow so it no longer runs on every PR/merge. It now runs only via
workflow_dispatch (manual) and a daily schedule (00:00 UTC). Also drop
the now-dead merge_group if-guards on all jobs and simplify the
concurrency key.

<!---
## 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. --->
CI test to be done

<!-- 
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 convergence

Co-authored-by: Copilot <[email protected]>
2026-08-26 21:48:05 +00:00
b28a0ca4f4 [cuTeDSL] Make SwiGLU CUDA-graph capturable + harmonize plain SiLU-Mul tests (#1408)
## Summary
Two related changes to the plain SiLU-Mul (SwiGLU:
`silu(a*gate)*b*down`) kernels — a CUDA-graph capturability fix for the
CuTeDSL backend, and harmonized test coverage across all three backends
(Triton / CuTeDSL / cuTile). No change to the Triton or cuTile kernels;
the fused-GEMM MLP kernels are untouched.

## 1. CuTeDSL SwiGLU: CUDA-graph capturability fix
The CuTeDSL SwiGLU launched all four kernels (vectorized fast path +
predicated fallback, forward and backward) without threading PyTorch's
current CUDA stream, so they ran on the default/null stream. As a result
the op was **not CUDA-graph capturable** — graph replay produced
stale/empty results (verified: graph-vs-eager max diff ~10-15, should be
0).

Fix mirrors `rms_norm.py` (the existing TVM-FFI + stream reference):
- add `_cute_stream()` (a cached `cuda.CUstream` over
`torch.cuda.current_stream()`),
- add `stream: cuda.CUstream = None` to each `@cute.jit` launcher and
pass `stream=stream` into every `.launch(...)`,
- for the TVM-FFI-compiled paths, compile against
`make_fake_stream(use_tvm_ffi_env_stream=True)` so the kernel tracks the
caller's env stream (the non-TVM-FFI `_dyn` path passes a real
`_cute_stream()`).

Numerics and gate/down semantics are unchanged. After the fix,
CUDA-graph capture + replay is **bit-identical to eager** on all paths
(fast/predicated, fwd/bwd), and the existing CuTeDSL suite still passes.

## 2. Harmonized plain SiLU-Mul tests across the three backends
Each backend test file (`test_swiglu.py`, `test_swiglu_cutedsl.py`,
`test_cutile_swiglu.py`) now defines a **byte-identical shared block**
(`_SILU_SHAPES`, `_SILU_MULTIPLIERS`, dtype lists, an fp32 reference,
tolerances — verified same md5 across files) and a **single parallel
set** of tests exercising each backend's own SiLU-Mul Function:
- correctness vs PyTorch (fwd+bwd) over aligned / predicated / tiny /
odd / non-power-of-2 "cliff" shapes (incl. Qwen2.5 11008 / 13824 /
18944) x fp32/fp16/bf16,
- gate/down multipliers,
- default-multiplier backward-compat,
- shape flexibility (1D-4D),
- CUDA-graph capturability (guards the fix in part 1),
- parity vs the Triton kernel (CuTeDSL and cuTile files).

The set is intentionally non-redundant: the cliff sizes live in the
shared `_SILU_SHAPES` (so every backend covers them once), the earlier
cuTile-specific cliff-sweep tests were folded into that shared list, and
Triton's multiplier coverage stays in its existing
`test_correctness_silumul_with_multipliers` rather than a duplicate.
Existing backend-specific tests are preserved: Triton's HF MLP-module +
DTensor tests, and CuTeDSL's low-level `swiglu_forward` functional test.

## Results (B200)
- cuTile: **68 passed**
- CuTeDSL: **70 passed**
- Triton: plain SiLU-Mul suite **48 passed** (adds
correctness-vs-PyTorch, shape-flexibility, and CUDA-graph tests;
multipliers already covered by the existing test). The only Triton
failures are the **pre-existing fused-GEMM `MixtralExperts`** tests —
unchanged from `main` and out of scope for plain SiLU-Mul.
- The new CUDA-graph tests pass on all three backends; lint (`ruff
check` + `format --check`) clean; the shared `_SILU_SHAPES` block is
byte-identical across the three files.

Co-authored-by: Copilot <[email protected]>

Co-authored-by: vaibhavjindal <[email protected]>
Co-authored-by: Copilot <[email protected]>
2026-08-26 17:14:47 +00:00