mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
main
979
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
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]> |
||
|
|
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]> |
||
|
|
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 |
||
|
|
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. |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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 |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]>
|
||
|
|
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
|
||
|
|
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>
|
||
|
|
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) |
||
|
|
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]> |
||
|
|
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]>
|
||
|
|
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]> |
||
|
|
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. |
||
|
|
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" />| |
||
|
|
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]> |
||
|
|
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. |
||
|
|
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]> |
||
|
|
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 |
||
|
|
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]>
|
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]>
|
||
|
|
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]>
|
||
|
|
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]> |
||
|
|
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]>
|
||
|
|
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]> |
||
|
|
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. |
||
|
|
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. |
||
|
|
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]> |
||
|
|
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]>
|
||
|
|
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]>
|
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |