Skip to content

Commit 2fe9d05

Browse files
ganeshnjganeshnj
committed
attention
1 parent cac9233 commit 2fe9d05

4 files changed

Lines changed: 452 additions & 2 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ Each operation is a single kernel, tested independently against NumPy.
9191
- [x] `rmsnorm` — composed: `square` → `reduce_sum` → `rsqrt` → `scale` (4 launches)
9292
- [x] `linear` — matvec using dot (one output per block)
9393
- [x] `cross_entropy` — composed: `softmax` → `gather` → `-log`
94-
- [ ] `attention` — composed: linear projections + dot + softmax + weighted sum
94+
- [x] `attention` — composed: linear projections + dot + softmax + weighted sum
9595

9696
### Stage 2 — Wire into microgpt
9797

‎docs/attention.md‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
# attention (Scaled Dot-Product Attention) — Design & Implementation
2+
3+
## What it does
4+
5+
Single-head scaled dot-product attention computes a context-weighted sum of value vectors, where the weights come from softmax-normalized dot products between a query and stored keys.
6+
7+
In microgpt:
8+
9+
```python
10+
attn_logits = [sum(q[j] * k_t[j] for j in range(dim)) / dim**0.5
11+
for t in range(seq_len)]
12+
attn_weights = softmax(attn_logits)
13+
out = [sum(attn_weights[t] * v_t[j] for t in range(seq_len))
14+
for j in range(dim)]
15+
```
16+
17+
With Q/K/V projections and output projection, attention is the core building block of every transformer layer.
18+
19+
## 12-kernel decomposition
20+
21+
```
22+
Step 1: q = Wq @ x → linear_kernel (1 launch)
23+
Step 2: k = Wk @ x → linear_kernel (1 launch)
24+
Step 3: v = Wv @ x → linear_kernel (1 launch)
25+
Step 4: scores = K @ q → matvec_kernel (1 launch)
26+
Step 5: scores /= sqrt(n_embd) → kern_div_scalar (1 launch)
27+
Steps 6-10: weights = softmax(scores) → 5 launches
28+
Step 11: attn_out = V^T @ weights → matvec_kernel (1 launch)
29+
Step 12: output = Wo @ attn_out → linear_kernel (1 launch)
30+
```
31+
32+
All 12 steps use existing kernels — no new kernels needed.
33+
34+
## Key design decisions
35+
36+
**Single head only.** Multi-head adds slicing/concatenation complexity without exercising new kernels. For Stage 1, single-head (head_dim = n_embd) is sufficient.
37+
38+
**V^T transpose on host.** The weighted sum `out[j] = sum_t(weights[t] * V[t,j])` is the matrix-vector product `V^T @ weights`. We transpose V on the host to `(n_embd, seq_len)` row-major before passing to `matvec_kernel`.
39+
40+
**seq_len <= 64.** Softmax requires the full scores vector in one block. microgpt processes tokens one at a time, so seq_len grows with context. For the test, seq_len = 4.
41+
42+
**sqrt(n_embd) as a 1-element f32 array.** Reuses the scalar broadcast `kern_div_scalar` pattern from softmax.
43+
44+
## Python orchestrator
45+
46+
```python
47+
def attention(x, Wq, Wk, Wv, Wo, K_cache, V_cache, n_embd):
48+
q = np.zeros(n_embd, dtype=np.float32)
49+
k = np.zeros(n_embd, dtype=np.float32)
50+
v = np.zeros(n_embd, dtype=np.float32)
51+
52+
linear_kernel[(n_embd,)](Wq, x, q, n_embd)
53+
linear_kernel[(n_embd,)](Wk, x, k, n_embd)
54+
linear_kernel[(n_embd,)](Wv, x, v, n_embd)
55+
56+
K_cache.append(k.copy())
57+
V_cache.append(v.copy())
58+
K = np.vstack(K_cache)
59+
V = np.vstack(V_cache)
60+
seq_len = len(K_cache)
61+
62+
scores = np.zeros(seq_len, dtype=np.float32)
63+
matvec_kernel[(seq_len,)](K.flatten(), q, scores, n_embd)
64+
65+
sqrt_d = np.array([np.sqrt(float(n_embd))], dtype=np.float32)
66+
scores_scaled = np.zeros(seq_len, dtype=np.float32)
67+
kern_div_scalar[(1,)](scores, sqrt_d, scores_scaled, seq_len)
68+
69+
weights = np.zeros(seq_len, dtype=np.float32)
70+
softmax(scores_scaled, weights, seq_len)
71+
72+
V_T = np.ascontiguousarray(V.T)
73+
attn_out = np.zeros(n_embd, dtype=np.float32)
74+
matvec_kernel[(n_embd,)](V_T.flatten(), weights, attn_out, seq_len)
75+
76+
output = np.zeros(n_embd, dtype=np.float32)
77+
linear_kernel[(n_embd,)](Wo, attn_out, output, n_embd)
78+
return output
79+
```
80+
81+
## How it maps to microgpt
82+
83+
microgpt's transformer block runs attention + residual + MLP:
84+
85+
```python
86+
q = linear(x, wq)
87+
k = linear(x, wk)
88+
v = linear(x, wv)
89+
keys[li].append(k)
90+
values[li].append(v)
91+
# per-head attention scores, softmax, weighted sum
92+
x = linear(x_attn, wo)
93+
x = [a + b for a, b in zip(x, x_residual)] # residual
94+
```
95+
96+
With tiny-ton, each of these calls maps to the kernel launches above.
97+
98+
## How it flows through the stack
99+
100+
```
101+
Python orchestrator: 12 x kernel[grid](...)
102+
│
103+
▼
104+
JIT (for each kernel): existing builtins
105+
│ linear (3x Q/K/V + 1x Wo) → tinyton.load + tinyton.mul + tinyton.reduce_sum
106+
│ matvec (scores, weighted) → same as linear
107+
│ div_scalar → tinyton.load + tinyton.div
108+
│ softmax (5x) → reduce_max, sub, exp, reduce_sum, div
109+
│
110+
▼
111+
GPU lowering (existing patterns)
112+
│ all ops already lowered
113+
│
114+
▼
115+
PTX: 12 separate kernel launches via cuLaunchKernel
116+
```
117+
118+
## Files changed
119+
120+
| File | What |
121+
|------|------|
122+
| `docs/attention.md` | This design doc |
123+
| `examples/attention_test.py` | Standalone test: 12-launch attention vs NumPy |
124+
125+
No C++ files changed. No new builtins in `jit.py`. All kernels are user-written.
126+
127+
## Testing strategy
128+
129+
`examples/attention_test.py` creates random weight matrices (Wq, Wk, Wv, Wo) of shape (16, 16), processes 4 input vectors to build a KV cache, then runs the full 12-launch attention for the last position. Compares against a NumPy reference that performs the same Q/K/V projections, scaled dot-product attention, and output projection. Tolerance: `atol=1e-3` (12 launches of float accumulation).

‎examples/attention_test.py‎

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
"""Single-head attention (12 kernel launches) -- verified against NumPy.
2+
3+
Run: PYTHONPATH="build/bindings:python" python3 examples/attention_test.py
4+
"""
5+
6+
import numpy as np
7+
import tiny_ton as tt
8+
9+
10+
# --- linear / matvec kernel (same kernel, different grid) -------------------
11+
12+
@tt.jit
13+
def linear_kernel(W_ptr, x_ptr, y_ptr, in_features):
14+
pid = tt.program_id(0)
15+
tid = tt.arange(0, 64)
16+
mask = tid < in_features
17+
w = tt.load(W_ptr + pid * in_features + tid, mask=mask)
18+
x = tt.load(x_ptr + tid, mask=mask)
19+
dot = tt.reduce_sum(w * x)
20+
tt.store(y_ptr + pid, dot)
21+
22+
23+
matvec_kernel = linear_kernel
24+
25+
26+
# --- softmax kernels (reused from softmax_test) -----------------------------
27+
28+
@tt.jit
29+
def kern_reduce_max(src, dst, N):
30+
pid = tt.program_id(0)
31+
off = pid * 64 + tt.arange(0, 64)
32+
mask = off < N
33+
x = tt.load(src + off, mask=mask)
34+
mx = tt.reduce_max(x)
35+
tt.store(dst + pid, mx)
36+
37+
38+
@tt.jit
39+
def kern_sub_scalar(src, scalar_ptr, dst, N):
40+
pid = tt.program_id(0)
41+
off = pid * 64 + tt.arange(0, 64)
42+
mask = off < N
43+
x = tt.load(src + off, mask=mask)
44+
s = tt.load(scalar_ptr)
45+
tt.store(dst + off, x - s, mask=mask)
46+
47+
48+
@tt.jit
49+
def kern_exp(src, dst, N):
50+
pid = tt.program_id(0)
51+
off = pid * 64 + tt.arange(0, 64)
52+
mask = off < N
53+
x = tt.load(src + off, mask=mask)
54+
tt.store(dst + off, tt.exp(x), mask=mask)
55+
56+
57+
@tt.jit
58+
def kern_reduce_sum(src, dst, N):
59+
pid = tt.program_id(0)
60+
off = pid * 64 + tt.arange(0, 64)
61+
mask = off < N
62+
x = tt.load(src + off, mask=mask)
63+
total = tt.reduce_sum(x)
64+
tt.store(dst + pid, total)
65+
66+
67+
@tt.jit
68+
def kern_div_scalar(src, scalar_ptr, dst, N):
69+
pid = tt.program_id(0)
70+
off = pid * 64 + tt.arange(0, 64)
71+
mask = off < N
72+
x = tt.load(src + off, mask=mask)
73+
s = tt.load(scalar_ptr)
74+
tt.store(dst + off, x / s, mask=mask)
75+
76+
77+
def softmax(x, out, N):
78+
"""Host-side softmax: 5 kernel launches."""
79+
grid = (max(1, (N + 63) // 64),)
80+
tmp_max = np.zeros(1, dtype=x.dtype)
81+
tmp_exp = np.zeros(N, dtype=x.dtype)
82+
tmp_sum = np.zeros(1, dtype=x.dtype)
83+
84+
kern_reduce_max[(1,)](x, tmp_max, N)
85+
kern_sub_scalar[grid](x, tmp_max, tmp_exp, N)
86+
kern_exp[grid](tmp_exp, tmp_exp, N)
87+
kern_reduce_sum[(1,)](tmp_exp, tmp_sum, N)
88+
kern_div_scalar[grid](tmp_exp, tmp_sum, out, N)
89+
90+
91+
# --- attention orchestrator (12 launches) -----------------------------------
92+
93+
def attention(x, Wq, Wk, Wv, Wo, K_cache, V_cache, n_embd):
94+
"""Single-head scaled dot-product attention with Q/K/V/O projections."""
95+
q = np.zeros(n_embd, dtype=np.float32)
96+
k = np.zeros(n_embd, dtype=np.float32)
97+
v = np.zeros(n_embd, dtype=np.float32)
98+
99+
linear_kernel[(n_embd,)](Wq, x, q, n_embd)
100+
linear_kernel[(n_embd,)](Wk, x, k, n_embd)
101+
linear_kernel[(n_embd,)](Wv, x, v, n_embd)
102+
103+
K_cache.append(k.copy())
104+
V_cache.append(v.copy())
105+
K = np.ascontiguousarray(np.vstack(K_cache))
106+
V = np.vstack(V_cache)
107+
seq_len = len(K_cache)
108+
109+
scores = np.zeros(seq_len, dtype=np.float32)
110+
matvec_kernel[(seq_len,)](K.flatten(), q, scores, n_embd)
111+
112+
sqrt_d = np.array([np.sqrt(float(n_embd))], dtype=np.float32)
113+
scores_scaled = np.zeros(seq_len, dtype=np.float32)
114+
kern_div_scalar[(1,)](scores, sqrt_d, scores_scaled, seq_len)
115+
116+
weights = np.zeros(seq_len, dtype=np.float32)
117+
softmax(scores_scaled, weights, seq_len)
118+
119+
V_T = np.ascontiguousarray(V.T)
120+
attn_out = np.zeros(n_embd, dtype=np.float32)
121+
matvec_kernel[(n_embd,)](V_T.flatten(), weights, attn_out, seq_len)
122+
123+
output = np.zeros(n_embd, dtype=np.float32)
124+
linear_kernel[(n_embd,)](Wo, attn_out, output, n_embd)
125+
return output
126+
127+
128+
# --- NumPy reference --------------------------------------------------------
129+
130+
def attention_numpy(x, Wq, Wk, Wv, Wo, K_cache, V_cache, n_embd):
131+
"""NumPy reference for single-head attention."""
132+
q = Wq @ x
133+
k = Wk @ x
134+
v = Wv @ x
135+
136+
K_cache.append(k.copy())
137+
V_cache.append(v.copy())
138+
K = np.vstack(K_cache)
139+
V = np.vstack(V_cache)
140+
141+
scores = K @ q / np.sqrt(float(n_embd))
142+
shifted = scores - np.max(scores)
143+
w = np.exp(shifted) / np.sum(np.exp(shifted))
144+
attn_out = V.T @ w
145+
return Wo @ attn_out
146+
147+
148+
# --- test -------------------------------------------------------------------
149+
150+
def main():
151+
np.random.seed(42)
152+
n_embd = 16
153+
n_tokens = 4
154+
155+
Wq = np.random.randn(n_embd, n_embd).astype(np.float32) * 0.1
156+
Wk = np.random.randn(n_embd, n_embd).astype(np.float32) * 0.1
157+
Wv = np.random.randn(n_embd, n_embd).astype(np.float32) * 0.1
158+
Wo = np.random.randn(n_embd, n_embd).astype(np.float32) * 0.1
159+
160+
tokens = [np.random.randn(n_embd).astype(np.float32) for _ in range(n_tokens)]
161+
162+
K_gpu, V_gpu = [], []
163+
K_ref, V_ref = [], []
164+
165+
all_ok = True
166+
for t in range(n_tokens):
167+
x = tokens[t]
168+
gpu_out = attention(x.copy(),
169+
Wq.flatten().copy(), Wk.flatten().copy(),
170+
Wv.flatten().copy(), Wo.flatten().copy(),
171+
K_gpu, V_gpu, n_embd)
172+
ref_out = attention_numpy(x.copy(), Wq, Wk, Wv, Wo,
173+
K_ref, V_ref, n_embd)
174+
175+
ok = np.allclose(gpu_out, ref_out, atol=1e-3)
176+
print(f"attention pos={t} (seq_len={t+1}): {'PASS' if ok else 'FAIL'}")
177+
if not ok:
178+
for i in range(n_embd):
179+
print(f" [{i}] got={gpu_out[i]:.6f} expected={ref_out[i]:.6f}")
180+
all_ok = False
181+
182+
assert all_ok, "attention test failed"
183+
print("All attention tests passed.")
184+
185+
186+
if __name__ == "__main__":
187+
main()

0 commit comments

Comments
 (0)