diff --git a/single_shot_inference.py b/single_shot_inference.py
index 77deee94..c8e6e8b8 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -484,6 +484,7 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
 
     # -- Get full KV from cache (already RoPE'd) --
     k_full, v_full = kv_cache.get()  # (1, seq_len, hd) each — RoPE'd, K=V
+    seq_len = k_full.shape[1]
 
     # -- FMHA: (n_h, T, hd) × (1, seq_len, hd) → (n_h, T, hd) --
     q_input = q_heads.permute(1, 0, 2)  # (n_h, T, hd)