Optimizing Loops for Softmax and FlashAttention in Triton
Puzzle 8: Long Softmax This puzzle implements a batched softmax operation. Given a batch of logits of shape [N0, T], compute the softmax for each row. The spec is: def softmax_spec(x: Float32[4, 200]) -> Float32[4, 200]: x_max = x.max(1, keepdim=True)[0] x = x - x_max x_exp = x.exp() return x_exp...