fix: restore kv_cache.append that was accidentally removed

2026-05-31 03:56:58 +00:00
parent 22a89b5a45
commit b519108cab
1 changed files with 1 additions and 0 deletions
--- 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)