Summary
This issue serves as a tracker for Week 2 (Online Softmax) submissions.
Participants are expected to contribute softmax kernels in CuTe DSL, along with minimal correctness and performance evidence. Multiple PRs can link back to this issue for coordination and review.
Scope
Softmax over the last dimension (row-wise).
Implement forward + backward in CuTe DSL, using an online softmax approach.
Forward
y = softmax(x, dim=-1) with numerical stability.
Backward
Given y and upstream dy, compute dx for softmax backward (row-wise reduction pattern). Validate against PyTorch autograd baseline (same shapes/dtypes).
Motivation / Use Case
Softmax is a critical building block in attention mechanisms and transformer architectures. The naive two-pass approach (compute max, then compute exp/sum) requires reading the input twice, while the online (single-pass) algorithm tracks a running max and rescales partial sums on the fly, halving memory traffic.
Week 2 focuses on exploring CuTe DSL implementations of online softmax and building intuition around numerical stability, warp-level reductions, and memory-bound kernel optimization.
This tracker centralizes:
- Submission PRs
- Optimization discussions
Proposed Solution
- Supported shapes: 2D tensors,
dim=-1 (row-wise)
- Supported dtypes: fp16 / bf16 / fp32 (as applicable) - recall the actual dtype required for softmax
- Forward: online softmax with running max for numerical stability
- Backward:
dx = y * (dy - dot) where dot = (dy * y).sum(dim=-1)
- Correctness validated against PyTorch
torch.softmax and torch.autograd
- Optional benchmark + pytest additions for each submission
Contributions that go beyond the baseline scope (e.g., fused softmax+dropout, multi-dimensional softmax, or FlashAttention-style tiling) are very welcome. Please highlight any such additions in your submission.
Scope Alignment
v0.1 scope (Weeks 0-2)
Summary
This issue serves as a tracker for Week 2 (Online Softmax) submissions.
Participants are expected to contribute softmax kernels in CuTe DSL, along with minimal correctness and performance evidence. Multiple PRs can link back to this issue for coordination and review.
Scope
Softmax over the last dimension (row-wise).
Implement forward + backward in CuTe DSL, using an online softmax approach.
Forward
y = softmax(x, dim=-1)with numerical stability.Backward
Given
yand upstreamdy, computedxfor softmax backward (row-wise reduction pattern). Validate against PyTorch autograd baseline (same shapes/dtypes).Motivation / Use Case
Softmax is a critical building block in attention mechanisms and transformer architectures. The naive two-pass approach (compute max, then compute exp/sum) requires reading the input twice, while the online (single-pass) algorithm tracks a running max and rescales partial sums on the fly, halving memory traffic.
Week 2 focuses on exploring CuTe DSL implementations of online softmax and building intuition around numerical stability, warp-level reductions, and memory-bound kernel optimization.
This tracker centralizes:
Proposed Solution
dim=-1(row-wise)dx = y * (dy - dot)wheredot = (dy * y).sum(dim=-1)torch.softmaxandtorch.autogradContributions that go beyond the baseline scope (e.g., fused softmax+dropout, multi-dimensional softmax, or FlashAttention-style tiling) are very welcome. Please highlight any such additions in your submission.
Scope Alignment
v0.1 scope (Weeks 0-2)