fix(generate): only feed the new token to forward() after the first step The loop was passing the full growing input_ids to forward() every step while also maintaining a KV cache, so at step i it reprocessed all i tokens on top of a cache that already held them — O(n^2) memory in the number of generated tokens. Short generations that hit EOS early survived; anything longer OOM'd (25.8 GiB allocated for a 158 MB model). Pass only ids[:, -1:] after the first step; the RoPE positions already account for the cache length via `start`. Verified: model.generate() now runs to full length (120+ tokens) with no OOM.

#1
by Compactbot - opened
Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment