Triton and Mosaic for linear_softmax_cross_entropy_loss#801

Open

captainpete wants to merge 21 commits into

openxla:mainfrom

captainpete:linear-xentropy-gpu

Commits on Mar 27, 2026

Add Pallas/Triton forward kernel for linear softmax cross-entropy loss
captainpete
committed
Add Pallas/Triton backward kernel for linear softmax cross-entropy loss
captainpete
committed
Add Pallas/Triton Op wiring for linear softmax cross-entropy loss
captainpete
committed
Add GPU benchmark harness and update README for linear_softmax_cross_entropy_loss
captainpete
committed
Rewrite backward to O(3BVH) single kernel with zero-init aliasing
captainpete
committed
Fuse dout scaling into backward kernel
captainpete
committed
Doc: add triton to api.py docstring, retire stale backend='triton' gotcha
captainpete
committed
Add Pallas/Mosaic-GPU SM90 Op for linear softmax cross-entropy loss
captainpete
committed
Register LCE loss benchmark in CI; add mosaic_gpu to benchmark implementations
captainpete
committed
Fix benchmark EXAMPLES: rename 'w' key to 'weights' to match public API
captainpete
committed
Switch mosaic_gpu and triton backward to padded-chunk cuBLAS scan
captainpete
committed
Clarify Triton exclusion: forward segfault remains, backward is resolved
captainpete
committed
Fix stale docstrings and PR.md backend selection order
captainpete
committed
Remove dead backward kernels (Triton atomic_add bwd, SM90 WGMMA bwd)
captainpete
committed
Remove mosaic_gpu from default backend chain; make it explicit opt-in
captainpete
committed
Switch Triton to heuristics config; drop autotuning
captainpete
committed
Improve Triton heuristics config; clean up PR.md
captainpete
committed
Triton: v-padding, heuristic overhaul, memory story in PR.md
captainpete
committed
Another pass on the PR.md
captainpete
committed
Remove PR doc
captainpete
committed
Remove uv.lock
captainpete
committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Triton and Mosaic for linear_softmax_cross_entropy_loss#801

Triton and Mosaic for linear_softmax_cross_entropy_loss#801
captainpete wants to merge 21 commits into
openxla:mainfrom
captainpete:linear-xentropy-gpu

Commits on Mar 27, 2026

Add Pallas/Triton forward kernel for linear softmax cross-entropy loss

Add Pallas/Triton backward kernel for linear softmax cross-entropy loss

Add Pallas/Triton Op wiring for linear softmax cross-entropy loss

Add GPU benchmark harness and update README for linear_softmax_cross_entropy_loss

Rewrite backward to O(3BVH) single kernel with zero-init aliasing

Fuse dout scaling into backward kernel

Doc: add triton to api.py docstring, retire stale backend='triton' gotcha

Add Pallas/Mosaic-GPU SM90 Op for linear softmax cross-entropy loss

Register LCE loss benchmark in CI; add mosaic_gpu to benchmark implementations

Fix benchmark EXAMPLES: rename 'w' key to 'weights' to match public API

Switch mosaic_gpu and triton backward to padded-chunk cuBLAS scan

Clarify Triton exclusion: forward segfault remains, backward is resolved

Fix stale docstrings and PR.md backend selection order

Remove dead backward kernels (Triton atomic_add bwd, SM90 WGMMA bwd)

Remove mosaic_gpu from default backend chain; make it explicit opt-in

Switch Triton to heuristics config; drop autotuning

Improve Triton heuristics config; clean up PR.md

Triton: v-padding, heuristic overhaul, memory story in PR.md

Another pass on the PR.md

Remove PR doc

Remove uv.lock