Fix RoPE cache overflow for long prompts with KV cache

[dipeshbabu]

When using KV cache during generation/inference, RoPE cos/sin buffers are sliced with an absolute offset (T0 = kv_cache.get_pos()), but the code only validated T <= cache_len. This can crash once T0 + T exceeds the cached RoPE length (even when T is small), matching the failure reported in #514.

Root cause:
RoPE cache bounds were validated against T instead of T0 + T, so long contexts / long generation runs can hit an out-of-bounds slice:
cos[:, T0:T0+T], sin[:, T0:T0+T].

Change:

• Ensure the RoPE cache covers T0 + T before slicing.
• Grow the cache to the next power of two for amortized behavior.
• Overwrite existing buffers instead of re-registering them.

[dipeshbabu]

@karpathy what do you think about this? could you review it?

[svlandeg]

Hi @dipeshbabu, as per the contribution guidelines, can you please declare any parts that had substantial LLM contribution, and whether there are any parts that you have not written and that you do not fully understand?

[svlandeg]

Is unbounded growth a good idea?

[dipeshbabu]

Is unbounded growth a good idea?

Unbounded growth is not ideal long-term. It fixes the crash, but without a cap it can silently keep allocating memory during very long generation.

[dipeshbabu]

@svlandeg I updated the stale RoPE cache comment, removed the problematic blanket exception pattern, and changed the lazy RoPE cache growth to be bounded max_rotary_seq_len to avoid unbounded memory growth in long generation runs. I also simplified _ensure_rope_cache() to infer device from the existing buffer and kept the forward() fix based on T0 + T.

[svlandeg]

@dipeshbabu:

Thanks! Per the contribution guidelines, can you please declare any parts that had substantial LLM contribution, and whether there are any parts that you have not written and that you do not fully understand?

[dipeshbabu]

Hi @dipeshbabu, as per the contribution guidelines, can you please declare any parts that had substantial LLM contribution, and whether there are any parts that you have not written and that you do not fully understand?

I used an LLM in a limited way to help me think through and explain the RoPE/KV-cache indexing issue more clearly specifically around why the cache bound needs to be checked against T0 + T during cached generation, not just T. It was used for explanation/reasoning support and wording, not for code I copied directly.

I wrote the final code changes myself and I understand the fix and the affected code path.

[svlandeg]

Thanks for the quick follow-up. To be perfectly honest with you, there were a little bit too many small errors / issues with the original implementation, which felt like vibe-coded to me. When I asked Claude, it gave a very similar implementation, code structure & comments.

I think it's good to look into this, as there currently is a "future TODO" in the code on master to dynamically grow the cache when we reach the limit here, and as #514 demonstrated this happens with max_seq_len at 256, which I can replicate.

But maybe this is something that Andrej should look into himself...

[svlandeg]

I don't like the fact that now there's one more magic number...