DSP & GPU Coursework
CUDA Attention Kernels
GPU kernels that move attention tensors between layouts. They do almost no arithmetic, so speed comes down to how threads touch memory.
4 tokens, 2 heads, 3 features each, in memory order.
Before: [token][head][feature]
After: [head][token][feature]
Token 2, head 1, feature 0: input address (2·2 + 1)·3 + 0 = 15, output address (1·4 + 2)·3 + 0 = 18.
Head 1 is shaded. After the permutation each head's values sit together.
What I did
I implemented permute and unpermute kernels for Q, K and V, choosing thread mappings for coalesced access, staging through shared memory and sizing blocks with occupancy in mind.
Result
- Permute and unpermute kernels implemented with coalescing, tiling and warp divergence considered.
- No GPU model, tensor shapes, timings or speedup are given, so none are claimed.
- Layout work around attention, not a fused or FlashAttention-style kernel; no Tensor Cores.
Design decision
Coalesce the expensive side
One side of a permutation is always strided. Mapping threads so a warp's accesses share transactions, and staging through shared memory, keeps both sides contiguous.
- Q, K, V[seq][head][feat]
- Permutemy kernel
- Per-head attentionnot part of this work
- Unpermutemy kernel
- Next layer[seq][head][feat]
More detail
Why permute
Attention splits each token's features into heads. Q, K and V move from token-major to head-major layout for the per-head work, and the output moves back. Values stay the same; addresses change.
Index arithmetic
A [sequence][head][feature] array becomes [head][sequence][feature]: input index (s·H + h)·D + d, output index (h·S + s)·D + d.
Other design decisions
Occupancy is a means
Size blocks against register and shared-memory use. Instead of: maximising threads per block.
More resident warps only help if memory traffic is the bottleneck.