Implement naive three-stage scaled dot-product attention in Triton. Given query, key, and value matrices Q, K, V all of shape (N,D):
out=softmax(DQK⊤)VUnlike FlashAttention, this implementation materializes the full N×N score matrix in a scratch tensor and runs three separate kernels: a tiled matmul for QK⊤, a row-wise softmax over the score matrix, and a second tiled matmul attn⋅V. Each kernel reads and writes the score tensor to HBM, which is exactly the bandwidth cost that FlashAttention avoids by fusing the three stages. Same math, very different memory footprint.
Input: Q = [[1.0, 0.0], [0.0, 1.0]], K = [[1.0, 0.0], [0.0, 1.0]], V = [[1.0, 2.0], [3.0, 4.0]]
Output: [[1.6604, 2.6604], [2.3396, 3.3396]]
Input: Q = [[0.0, 0.0]], K = [[1.0, 1.0], [2.0, 2.0]], V = [[5.0, 5.0], [10.0, 10.0]]
Output: [[7.5, 7.5]]
Stage 1 is a tiled QK⊤ scaled by 1/D. Each program writes one tile of the score matrix to HBM.
Stage 2 is row-wise softmax on the score matrix. One program per row: subtract the row max, exponentiate, divide by the sum. Same structure as the fused-softmax problem.
Stage 3 is attn⋅V, structurally identical to stage 1 but with the materialized score matrix as the left operand. Three kernels, three HBM round-trips - that is the bandwidth cost.
torch.float32 on cuda. Single head, no batch dimension.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥64 due to two sequential matmuls accumulating drift.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement naive three-stage scaled dot-product attention in Triton. Given query, key, and value matrices Q, K, V all of shape (N,D):
out=softmax(DQK⊤)VUnlike FlashAttention, this implementation materializes the full N×N score matrix in a scratch tensor and runs three separate kernels: a tiled matmul for QK⊤, a row-wise softmax over the score matrix, and a second tiled matmul attn⋅V. Each kernel reads and writes the score tensor to HBM, which is exactly the bandwidth cost that FlashAttention avoids by fusing the three stages. Same math, very different memory footprint.
Input: Q = [[1.0, 0.0], [0.0, 1.0]], K = [[1.0, 0.0], [0.0, 1.0]], V = [[1.0, 2.0], [3.0, 4.0]]
Output: [[1.6604, 2.6604], [2.3396, 3.3396]]
Input: Q = [[0.0, 0.0]], K = [[1.0, 1.0], [2.0, 2.0]], V = [[5.0, 5.0], [10.0, 10.0]]
Output: [[7.5, 7.5]]
Stage 1 is a tiled QK⊤ scaled by 1/D. Each program writes one tile of the score matrix to HBM.
Stage 2 is row-wise softmax on the score matrix. One program per row: subtract the row max, exponentiate, divide by the sum. Same structure as the fused-softmax problem.
Stage 3 is attn⋅V, structurally identical to stage 1 but with the materialized score matrix as the left operand. Three kernels, three HBM round-trips - that is the bandwidth cost.
torch.float32 on cuda. Single head, no batch dimension.atol=1e-2 + rtol=1e-2 * abs(expected) for N≥64 due to two sequential matmuls accumulating drift.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.