perf(ce): hardware exp2 (ex2.approx) for softmax by justinhh4 · Pull Request #1266 · linkedin/Liger-Kernel

justinhh4 · 2026-06-25T21:05:31Z

PR1: CE `exp2` — B200 + H100 benchmarks

Optimization measured against the frozen original baseline (v0_base, == upstream HEAD). Correctness: full CE suite (177 cases, bf16+fp32, weight/softcap/smoothing/z_loss/ignore_index combos) passes.

What & why

Replace tl.exp with the hardware ex2.approx instruction via e^x = 2^(x·log₂e), in both the online-softmax accumulation and the gradient pass.

Why: NVIDIA GPUs have no native e^x — tl.exp lowers to a multi-instruction accurate sequence, while tl.exp2 emits a single ex2.approx.f32. This is the most frequent FP op in the kernel (once per vocab element × 2 passes), so it cuts total issued instructions.
Accuracy: ex2.approx has ~2.4e-7 relative error — within the suite's fp32 rtol=1e-6 and bf16 rtol=5e-2, and softmax normalization further cancels the (correlated) error. Helps every dtype.

Headline shape (BT=8192, V=128256, full fwd+bwd)

GPU	dtype	baseline	this PR	latency	throughput
H100	bf16	2.583 ms	2.225 ms	-13.9%	+16.1%
H100	fp32	4.489 ms	4.475 ms	-0.3%	+0.3%
B200	bf16	1.847 ms	1.687 ms	-8.6%	+9.5%
B200	fp32	2.085 ms	1.887 ms	-9.5%	+10.5%

Across context length (V=128256, BT 1024–65536)

GPU	dtype	min	avg	max
H100	bf16	-10%	-13%	-15%
H100	fp32	-0%	-0%	-1%
B200	bf16	-4%	-8%	-10%
B200	fp32	-5%	-8%	-10%

Figures

Latency vs context length and vs vocab, baseline vs this PR, on H100 and B200.

vaibhavjindal · 2026-06-25T22:44:12Z

Good find! Added a minor comment.

@Tcc0403 @Mecoli1219 do you see any issues with this change?

vaibhavjindal · 2026-06-25T22:43:29Z

We should define 1.4426950408889634 as a constant and use it instead of explicitly writing it multiple times.

Changed it to a constant.

Address review feedback on linkedin#1266: define 1.4426950408889634 once as the module-level LOG2_E constant and reference it from the three exp2 softmax sites instead of repeating the literal. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Tcc0403

LGTM

Address review feedback on linkedin#1266: define 1.4426950408889634 once as the module-level LOG2_E constant and reference it from the three exp2 softmax sites instead of repeating the literal. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

yueyiming2009 · 2026-06-27T18:44:32Z

A late review but should we have a wrapper to improve readability? @vaibhavjindal @Tcc0403 ?

@triton.jit def fast_exp(x): """ Hardware-accelerated base-e exponentiation. Uses ex2.approx via the identity e^x = 2^(x * log2(e)). """ LOG2_E = 1.4426950408889634 return tl.exp2(x * LOG2_E)

Tcc0403 · 2026-06-27T19:16:08Z

I prefer plain exp2 call, leaving developers to decide ordering for numerical concerns. E.g. exp2((a - b) - log2_e) or exp2(a * log2_e - b * log2_e)

…inkedin#1275) Follow up to linkedin#1266 the cross entropy kernel reads the module level `LOG2_E` constant from inside the `@triton.jit` kernel. Newer Triton no longer allows accessing a plain (non-constexpr) global from within a kernel, so the kernel fails to compile with a NameError and the cross entropy correctness tests error out. Wrapping the constant in `tl.constexpr` is what the Triton error itself suggests, and it keeps the value identical. The cross entropy tests pass again with this change.

## CE `num_warps`: dtype- & arch-aware (Blackwell-gated) Rebased onto current `main` (post-linkedin#1266). All numbers below use **linkedin#1266 (`v1_exp2`) as the baseline**, so they isolate *this* PR's contribution and exclude linkedin#1266's exp2 gains. Measured on **B200**, which is what this gate targets. Correctness: full CE suite (bf16+fp32, with weight/softcap/label_smoothing/z_loss/ignore_index combos) passes. ### What & why Replace the hardcoded `num_warps=32` in CE forward with a dtype- and arch-aware choice: - **Blackwell (sm_100+) bf16/fp16 → 8**; everything else (Hopper & earlier, all fp32) → **32**; AMD → **16**. - Adds a small reusable helper `get_gpu_arch()` in `ops/utils.py` that returns `"blackwell"` / `"hopper"` / `"ampere"` / `""`. It uses `infer_device()` as the CUDA guard and excludes AMD via `is_hip()` (since `infer_device()` reports AMD as `"cuda"` too). The gate reads `get_gpu_arch() == "blackwell"`. - **Why 8 on Blackwell bf16:** CE there is *instruction-issue-bound* (ncu: ALU is the top pipe; the dominant stall is "Not Selected" from a surplus of warps) — fewer warps cut issue contention. Hopper/earlier and fp32 are *bandwidth-bound* and keep 32 to hide memory latency. - **Scoped to exactly bf16/fp16** (`element_size() == 2`) on sm_100+; fp8 (1 B) and fp32 keep 32 warps (unmeasured → safe default). `num_warps` is purely a scheduling parameter → no correctness impact. ### Headline (BT=8192, V=128256, full fwd+bwd; baseline = linkedin#1266) | dtype | baseline (linkedin#1266) | this PR | latency | throughput | |---|--:|--:|--:|--:| | bf16 | 1.686 ms | 1.438 ms | **−14.7%** | +17.2% | | fp32 | 1.892 ms | 1.892 ms | +0.0% (unchanged) | +0.0% | Only bf16/fp16 fires the gate; fp32 keeps 32 warps and is flat — shown to confirm no regression. ### bf16 across context length (V=128256; baseline = linkedin#1266) | BT | baseline | this PR | latency | |--:|--:|--:|--:| | 1024 | 0.405 ms | 0.391 ms | −3.4% | | 2048 | 0.583 ms | 0.542 ms | −7.0% | | 4096 | 0.950 ms | 0.839 ms | −11.6% | | 8192 | 1.686 ms | 1.438 ms | −14.7% | | 16384 | 3.184 ms | 2.694 ms | −15.4% | | 32768 | 6.092 ms | 5.101 ms | −16.3% | | 65536 | 11.958 ms | 9.924 ms | −17.0% | Gain grows with batch size (−3% → −17%; avg −12% over the sweep). ### bf16 across vocab (BT=8192; baseline = linkedin#1266) −17.1% @ V=32k → −8.6% @ V=262k (avg −13.0%) — larger relative win at smaller vocab. ### Figures <img width="818" height="566" alt="pr2_bf16_bt" src="https://github.com/user-attachments/assets/fda19ba4-aab2-4dbc-98de-8fc0ad060f7b" /> <img width="818" height="566" alt="pr2_bf16_vocab" src="https://github.com/user-attachments/assets/537b6af0-689c-4be0-97cd-f8b8c4953a08" /> <img width="818" height="566" alt="pr2_fp32_bt" src="https://github.com/user-attachments/assets/600c4e4f-c8b5-4791-a595-8bc5ea73419a" /> <img width="818" height="566" alt="pr2_fp32_vocab" src="https://github.com/user-attachments/assets/fc541ce2-1f62-44c4-9417-b0d039953299" /> --------- Co-authored-by: Justin Hu <181588904+justinhh4@users.noreply.github.com> Co-authored-by: Vaibhav Jindal <vaibhav.jndl@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

This was referenced Jun 25, 2026

perf(ce): dtype-aware num_warps (Blackwell-gated) #1267

Merged

perf(ce): single-shot dx_y gradient correction #1268

Open

vaibhavjindal requested review from Mecoli1219 and Tcc0403 June 25, 2026 22:40

vaibhavjindal reviewed Jun 25, 2026

View reviewed changes

Tcc0403 approved these changes Jun 26, 2026

View reviewed changes

justinhh4 and others added 2 commits June 26, 2026 10:47

perf(ce): hardware exp2 (ex2.approx) for softmax exp

4f4e707

justinhh4 force-pushed the ce-opt-1-exp2 branch from eb7ef00 to e538f80 Compare June 26, 2026 17:48

vaibhavjindal approved these changes Jun 26, 2026

View reviewed changes

vaibhavjindal added this pull request to the merge queue Jun 26, 2026

Merged via the queue into linkedin:main with commit 8832b2e Jun 26, 2026
5 of 7 checks passed

kashif mentioned this pull request Jun 29, 2026

Make LOG2_E a tl.constexpr so the CE kernel compiles on newer Triton #1275

Merged

This was referenced Jul 6, 2026

[CuteDSL] Add benchmarking, CE memory optimizations, and match Triton num_warps justinhh4/Liger-Kernel#1

Open

[CuteDSL] Add benchmarking, CE memory optimizations, and match Triton num_warps #1290

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

perf(ce): hardware exp2 (ex2.approx) for softmax#1266

perf(ce): hardware exp2 (ex2.approx) for softmax#1266
vaibhavjindal merged 2 commits into
linkedin:mainfrom
justinhh4:ce-opt-1-exp2

justinhh4 commented Jun 25, 2026 •

edited

Loading

Uh oh!

vaibhavjindal commented Jun 25, 2026 •

edited

Loading

Uh oh!

vaibhavjindal Jun 25, 2026

Uh oh!

justinhh4 Jun 25, 2026

Uh oh!

Tcc0403 left a comment

Uh oh!

Uh oh!

yueyiming2009 commented Jun 27, 2026

Uh oh!

Tcc0403 commented Jun 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

justinhh4 commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR1: CE exp2 — B200 + H100 benchmarks

What & why

Headline shape (BT=8192, V=128256, full fwd+bwd)

Across context length (V=128256, BT 1024–65536)

Figures

Uh oh!

vaibhavjindal commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vaibhavjindal Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

justinhh4 Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

Tcc0403 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yueyiming2009 commented Jun 27, 2026

Uh oh!

Tcc0403 commented Jun 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

justinhh4 commented Jun 25, 2026 •

edited

Loading

PR1: CE `exp2` — B200 + H100 benchmarks

vaibhavjindal commented Jun 25, 2026 •

edited

Loading