From fe30c86908e9eb4e41bc08479b963b7c565b7211 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 5 Aug 2026 15:06:43 +0000 Subject: [PATCH] llm: restore master form in TransformerBlock._attention non-quantized path is byte-identical to master again; the only change in this block is the stacked_kv hoist and the quantized_attention branch --- tinygrad/llm/model.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 45e8a28fa5..f9e0221b07 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -173,10 +173,18 @@ class TransformerBlock(FFNBlock): attn = quantized_attention(q, stacked_kv, self.cache_kv, self.cache_kv_scale, start_pos) else: assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(stacked_kv.uop))) - k, v = assigned_kv[0, :, :, 0:start_pos+T, :], assigned_kv[1, :, :, 0:start_pos+T, :] - mask = Tensor.full((1, 1, T, k.shape[-2]), float("-inf"), dtype=q.dtype, device=q.device, buffer=False).triu(start_pos+1) \ + k = assigned_kv[0, :, :, 0:start_pos+T, :] + v = assigned_kv[1, :, :, 0:start_pos+T, :] + + #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(stacked_kv) + #k = self.cache_kv[0, :, :, 0:start_pos+T, :] + #v = self.cache_kv[1, :, :, 0:start_pos+T, :] + + # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True + # TODO: this if statement should be removed and it shouldn't generate extra kernels + mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \ if resolve(T != 1) else None - attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) + attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) # (B,H,T,Hd) attn = attn.transpose(1, 2).reshape(B, T, -1) # back to (B,T,D) return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid()))