mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
540575bc2b
## Summary This PR optimizes the backward pass performance of the LayerNorm operator by implementing a Persistent Kernel with Partial Reduction strategy. The previous implementation relied on `atomic_add` to accumulate gradients `(dw, db)` across the batch/sequence dimension. For input shapes with a large number of tokens (e.g., large Batch Size or Sequence Length), this caused severe atomic lock contention, leading to significant performance degradation (kernel serialization). This change replaces the atomic operations with a deterministic reduction approach, achieving 5x-12x speedups on large-scale inputs while maintaining numerical correctness. ## Testing Done - Hardware Type: A100 80GB SXM4 - [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: niyunsheng <[email protected]> Co-authored-by: Steven Shimizu <[email protected]>