diff --git a/test/unit/test_llm_amd.py b/test/unit/test_llm_amd.py
index a5afd1863a..ba0ba16bd3 100644
--- a/test/unit/test_llm_amd.py
+++ b/test/unit/test_llm_amd.py
@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, dtypes, nn
+from tinygrad import Tensor, UOp, dtypes, nn
 from tinygrad.llm.kernels import Linear
 from tinygrad.llm.kernels.amd import q8_quantize, quantized_attention
 from tinygrad.llm.gguf import ggml_data_to_tensor
@@ -34,4 +34,22 @@ class TestQ8Quantize(unittest.TestCase):
     out = quantized_attention(q, Tensor.stack(k, v), cache, scale, 0).realize()
     np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
 
+  def test_prefill_attention_unaligned_start(self):
+    if not str(Tensor.empty(1).device).startswith("AMD"): self.skipTest("AMD required")
+    rng = np.random.default_rng(42)
+    start_pos = 31
+    q = Tensor.zeros(1, 8, 32, 128)
+    old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
+    new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
+    cache = Tensor.empty(2, 1, 1, 256, 128, dtype=dtypes.int8).contiguous()
+    scale = Tensor.empty(2, 1, 1, 256, dtype=dtypes.float16).contiguous()
+    old_scale = np.maximum(np.max(np.abs(old_kv), axis=-1, keepdims=True) / 127, 1e-8).astype(np.float16)
+    Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(np.rint(old_kv / old_scale).astype(np.int8))),
+                   scale[:, :, :, :start_pos].assign(Tensor(old_scale.squeeze(-1))))
+    out = quantized_attention(q, Tensor(new_kv), cache, scale, UOp.variable("start_pos", 0, 255).bind(start_pos)).realize()
+    values = cache[1, 0, 0, :start_pos+32].numpy().astype(np.float32) * \
+      scale[1, 0, 0, :start_pos+32].numpy().astype(np.float32)[:, None]
+    expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
+    np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
+
 if __name__ == "__main__": unittest.main()
diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py
index e1e7df254e..eb84d6fa6c 100644
--- a/tinygrad/llm/kernels/amd.py
+++ b/tinygrad/llm/kernels/amd.py
@@ -134,7 +134,8 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, kv_scale:UOp, valid_kv_len:int
   m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
   ri_w = UOp.range(TM, 270)
   m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce(m_ij[ri_w], lane, maximum=True)).end(ri_w))
-  S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
+  tile_max = m_ij.reshape(TM, 1).expand(TM, TN).maximum(-1e30)
+  S_reg = S_reg.after(S_reg.store(((S_reg - tile_max) * LOG2E).exp2()))
   p_local, ri_ws = _reg((TM,), 8, 0, n_tile), UOp.range(TM, 295, AxisType.WEAK)
   p_sum = p_local.after(p_local[ri_ws].store(sum((warp_reduce(S_reg[ri_ws, rn], lane) for rn in range(TN)), S_reg.const_like(0))).end(ri_ws))
   P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)