From 3c32a77df7ce6dfa94b93bd04512cf2eb5456cd6 Mon Sep 17 00:00:00 2001
From: Ronald Tse <ronald.tse@gmail.com>
Date: Fri, 28 Aug 2026 13:14:11 +0800
Subject: [PATCH] fix(imf): cut greedy-KV generation when the token window
 cycles
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Live failure: short inputs (كتاب) made the int8 students echo the
same phrase until max_len. The guard stops generation when the last
24 tokens repeat verbatim earlier in the output; goldens and parity
unaffected (90/90).
---
 src/imf/export.py          | 13 +++++++
 tests/test_decode_guard.py | 76 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 tests/test_decode_guard.py

diff --git a/src/imf/export.py b/src/imf/export.py
index 152a856..e6775c1 100644
--- a/src/imf/export.py
+++ b/src/imf/export.py
@@ -393,6 +393,12 @@ def onnx_greedy_kv(encoder_sess, kv_sess, text: str, max_len: int = 256) -> list
     pasts = _zero_pasts(kv_sess)
     cur = np.array([[0]], dtype=np.int64)
     generated: list[int] = []
+    # Flat-byte students can loop on short inputs (greedy KV without
+    # sampling): cut generation when the recent token window repeats
+    # verbatim — the live failure was diacritization echoing the same
+    # phrase until max_len.
+    window = 24
+    joined = ""
     for _ in range(max_len):
         out = kv_sess.run(None, {"input_ids": cur, "encoder_hidden_states": hidden, **pasts})
         results = dict(zip(out_names, out, strict=True))
@@ -405,6 +411,13 @@ def onnx_greedy_kv(encoder_sess, kv_sess, text: str, max_len: int = 256) -> list
             for name in pasts
         }
         cur = np.array([[nxt]], dtype=np.int64)
+        if len(generated) >= 2 * window:
+            joined += f",{nxt}"
+            needle = ",".join(str(t) for t in generated[-window:])
+            if joined.find(needle) < len(joined) - len(needle):
+                break
+        else:
+            joined = ",".join(str(t) for t in generated)
     return generated
 
 
diff --git a/tests/test_decode_guard.py b/tests/test_decode_guard.py
new file mode 100644
index 0000000..db969b4
--- /dev/null
+++ b/tests/test_decode_guard.py
@@ -0,0 +1,76 @@
+"""Repetition guard in the greedy KV decode: flat-byte students loop
+on short inputs (live: كتاب -> كَتَابٍ: كَتَابٍ: ... until max_len).
+The guard must cut generation when the recent token window cycles."""
+
+from src.imf.export import onnx_greedy_kv
+
+
+class _FakeEncoder:
+    def run(self, _, feeds):
+        import numpy as np
+
+        return [np.zeros((1, feeds["input_ids"].shape[1], 4), dtype=np.float32)]
+
+
+class _FakeKV:
+    """Cycles two tokens forever, never EOS — the pathological loop."""
+
+    def __init__(self):
+        import numpy as np
+
+        self.np = np
+        self.step = 0
+
+    def get_outputs(self):
+        class O:
+            def __init__(self, name):
+                self.name = name
+
+        return [O("logits"), O("present_k"), O("present_v")]
+
+    def get_inputs(self):
+        class I:
+            def __init__(self, name, shape, typ="tensor(float)"):
+                self.name = name
+                self.shape = shape
+                self.type = typ
+
+        return [
+            I("input_ids", [1, "seq"]),
+            I("encoder_hidden_states", [1, "seq", 4]),
+            I("past_k", [1, 4, "past", 8]),
+            I("past_v", [1, 4, "past", 8]),
+        ]
+
+    def run(self, _, feeds):
+        import numpy as np
+
+        logits = np.full((1, 1, 260), -1e9)
+        # visible cycle: token 10, 11, 10, 11, ...
+        logits[0, -1, 10 if self.step % 2 == 0 else 11] = 1e9
+        self.step += 1
+        return [logits, np.zeros((1,)), np.zeros((1,))]
+
+
+def test_looping_model_is_cut_before_max_len():
+    out = onnx_greedy_kv(_FakeEncoder(), _FakeKV(), "x", max_len=4096)
+    # 2-token cycle caught by the window guard well before max_len
+    assert len(out) < 200
+    assert out[:4] == [10, 11, 10, 11]
+
+
+def test_normal_generation_unaffected():
+    class _StopKV(_FakeKV):
+        def run(self, _, feeds):
+            import numpy as np
+
+            logits = np.full((1, 1, 260), -1e9)
+            if self.step >= 5:
+                logits[0, -1, 1] = 1e9  # EOS_ID
+            else:
+                logits[0, -1, 20 + self.step] = 1e9
+            self.step += 1
+            return [logits, np.zeros((1,)), np.zeros((1,))]
+
+    out = onnx_greedy_kv(_FakeEncoder(), _StopKV(), "x", max_len=256)
+    assert out == [20, 21, 22, 23, 24]