Hi @Chelsi-create , thank you for the PR!
There are still a few tests failing related to your changes, could you check that the tests pass?
https://github.com/Lightning-AI/litgpt/actions/runs/19088536192/job/54534011522?pr=2153

Hi @Chelsi-create, thanks very much for working on this! I am the original filer of the bug and I really appreciate you taking up the issue. I have several questions related to your testing and findings.

(1) It looks like you are focused on an issue with sliding window attention in Gemma, but in my original runlitgpt.py script (using LLM.generate), I also see the same problem of repetition when running Llama. Were you not able to replicate the problem with Llama using my runlitgpt.py script? And if you were able to reproduce the problem with Llama, did this problem get somehow fixed as a side effect of the fix you made for sliding window KV cache management in Gemma? I realize that Llama only uses Global Attention, so that's why I am confused how the problem with Llama was also apparently fixed (you mention Llama passing all three "long form" generation tests in your testing summary).

Another thing worth mentioning here ... The problem with Gemma (before your fix) occurs both when using LLM.generate (as in my runlitgpt.py script) and also when running Gemma using litgpt chat (for the same prompts and parameter settings). However, when running Llama, I only see the repetition problem when running Llama with LLM.generate and NOT when I use the same prompt and parameters in litgpt chat. This suggests there are two different bugs. Gemma probably fails in both cases because of the sliding window context issue you are addressing with this PR. However, why would Llama fail when using LLM.generate (with very similar symptoms as the Gemma failure) but not fail when using litgpt chat? This is all quite puzzling to me. Of course, my comments here assume you are able to reproduce what I just described: Llama has the repetitive behavior (just like Gemma) before your fix when using LLM.generate() but not when using litgpt chat.

(2) You make a distinction between LLM.generate and chat_generate(), but your linked files for runlitgpt.py and runlitgpt_v1.py both seem to be using LLM.generate. The only difference seems to be in the prompts. So, I am a little confused about the difference between LLM.generate and chat_generate() and why this distinction matters in the context of this bug.

(3) You mention this in your testing summary: "After fix: coherent outputs with limited repetition (7-9 tokens)." I'm somewhat concerned to hear that there is still any repetition at all. Can you share an example of the "limited repetition" you are seeing?

I'd prefer we not close the original issue until these questions are addressed. Please let me know if there is any way I can help in resolving any of these questions. I'd really like to use litgpt but I'm uneasy about that until I understand the answers to the questions I raised above.

Thanks!

Hi @drwslacy47 ,

Thank you so much for the detailed and thoughtful questions. I have tried to address all your questions below:

(1) Llama Behavior: Why it Failed with LLM.generate() but Works with chat_generate()

This is an excellent observation, and I should clarify what I discovered.

The Core Finding
After careful investigation, both Gemma and Llama fail with the high-level LLM.generate() API (as in your original runlitgpt_v1.py), but the low-level chat_generate() based runlitgpt.py works for both models.

Your original runlitgpt_v1.py using llm.generate() → Shows repetition for Llama
The runlitgpt.py using chat_generate() → Works coherently for Llama

Why does this happen?
The difference is two separate issues:

For Gemma:
Primary issue: Sliding window KV cache bug (what this PR fixes)
When using chat_generate() + correct stop tokens → Works perfectly
When using LLM.generate() → Still fails due to cache management in api.py

For Llama:
No sliding window issue (uses global attention only)
Problem: LLM.generate() in api.py line 762 uses only EOS token for stopping
When using chat_generate() + proper model stop tokens (prompt_style.stop_tokens()) -> this works
When using LLM.generate() + only EOS token → Model keeps generating past natural stopping point

Why Works in litgpt chat but not in LLM.generate()?

• The litgpt chat CLI uses the same chat_generate() function as runlitgpt.py script:
• It correctly uses prompt_style.stop_tokens()
• Each user turn resets the KV cache
• No long-form generation in single call

Core problem with the LLM.generate() API:

• Uses incomplete stop token handling (only EOS)
• Tries to manage cache internally without proper reset
• Has different code path in api.py that bypasses my model.py fix partially

The Key Difference Between the Two Scripts
runlitgpt_v1.py (BROKEN - uses LLM.generate()):

python

response = llm.generate(
    prompt,
    max_new_tokens=1500,
    temperature=0.6,
    top_k=64,
    top_p=0.9
)

runlitgpt.py (WORKING - uses chat_generate()):

python

from litgpt.chat.base import generate as chat_generate

stop_tokens = prompt_style.stop_tokens(tokenizer)  # ← Gets ALL stop tokens!
stream = chat_generate(
    model=llm.model,
    prompt=input_ids,
    max_returned_tokens=max_returned_tokens,
    temperature=0.6,
    top_k=64,
    top_p=0.9,
    stop_tokens=stop_tokens  # ← Uses proper stop tokens
)

Short summary:
The runlitgpt.py script works for BOTH models because:

• Properly pre-sizes KV cache with llm.model.set_kv_cache()
• Uses model-specific stop tokens: prompt_style.stop_tokens(tokenizer)
• Uses low-level chat_generate() directly
• The sliding window fix in model.py is properly utilized

(2) Clarifying LLM.generate() vs chat_generate()

LLM.generate() - High-Level Python API
Location: litgpt/generate/api.py
Advantages: Simple, one-line API
Disadvantages: Limited stop token control, internal cache management issues

chat_generate() - Low-Level Function
Advantages: Complete control, explicit stop tokens, cleaner cache handling
Disadvantages: More setup required

Why This Matters for the Bug
Your original script uses LLM.generate():

python

llm = LLM.load(model_path)
response = llm.generate(prompt, ...)  # High-level API

python

from litgpt.chat.base import generate as chat_generate

stop_tokens = prompt_style.stop_tokens(tokenizer)  # ← Explicitly get all stop tokens
stream = chat_generate(..., stop_tokens=stop_tokens)

So both your runlitgpt_v1.py and runlitgpt.py conceptually do generation, but runlitgpt.py actually uses the lower-level chat_generate() underneath, which gives us full control.

(3) The "Limited Repetition (7-9 tokens)" Explanation

My earlier statement about “limited repetition” was imprecise. What I actually measured was word frequency in the last 200 tokens, which isn’t a good indicator of coherence. Natural repetition (eg “the”, “and”) can appear frequently without being an issue.