### Bug description
### Bug
`generate_text_simple_stream(..., context_size=None)` exposes a `context_size` parameter, but never reads it: the initial cache prefill passes the entire `token_ids` tensor to the model. In the same file, non-streaming `generate_text_simple` uses `context_size` to slice the prompt. Consequently, callers cannot cap the streaming prompt through the advertised argument, and an input longer than the model's positional context can fail before producing a token.
Source: [streaming function](https://github.com/rasbt/LLMs-from-scratch/blob/f49cad747931fe6af0a7d2802cd480d1c1a0ec8e/pkg/llms_from_scratch/kv_cache/generate.py#L34-L42), [non-streaming comparison](https://github.com/rasbt/LLMs-from-scratch/blob/f49cad747931fe6af0a7d2802cd480d1c1a0ec8e/pkg/llms_from_scratch/kv_cache/generate.py#L11-L22). This also affects the [single-turn Qwen3 Chainlit example](https://github.com/rasbt/LLMs-from-scratch/blob/f49cad747931fe6af0a7d2802cd480d1c1a0ec8e/ch05/11_qwen3/qwen3-chat-interface/qwen3-chat-interface.py#L116-L130), which passes encoded user input directly to the streaming helper. The multi-turn example explicitly calls `trim_input_tensor` first.
### Minimal reproduction
Using the repository's small GPT-2 test configuration with `context_length=16`:
```python
model = make_model() # from pkg/llms_from_scratch/tests/test_kv_cache_gpt2.py
prompt = torch.randint(0, 96, (1, 20))
list(generate_text_simple_stream(model, prompt, max_new_tokens=1, context_size=8))
```
The streaming function forwards all 20 prompt tokens, despite `context_size=8`; GPT-2's positional embedding has only 16 entries. The non-streaming helper truncates to the requested 8-token context before its first forward pass. This is based on source inspection; I have not run this PyTorch snippet locally.
### Expected behavior and suggested fix
Apply the effective context limit before the streaming prefill, accounting for the requested generation budget because this cache path does not slide its window. The existing `trim_input_tensor` helper can be reused after validating `max_new_tokens < effective_context_length` and capping/validating against `model.cfg["context_length"]`. A regression test should pass a prompt longer than the requested context and assert that the model receives only the retained suffix; add a test for an overlong prompt in the single-turn chat path.
This report concerns the package's generation behavior and its included UI example, not a request for unrelated features.
### What operating system are you using?
None
### Where do you run your code?
None
### Environment
Source-level finding. I have not executed the PyTorch reproduction locally; environment details are not applicable.
Ran your repro at `f49cad7` (CPU, torch 2.2.2, the `TEST_CONFIG` from `tests/test_kv_cache_gpt2.py`) since you flagged that you had not executed it. It reproduces: `IndexError: index out of range in self` from the positional embedding, and `context_size=8`, `context_size=16` and `context_size=None` all fail identically — the argument is inert rather than under-applied.
Worth widening the scope slightly: the non-streaming comparison has the same overflow in its cached branch. `GPTModel.forward` sets `use_cache = use_cache or cache is not None` and then advances `current_pos`, so any cached call breaks once `kept_prompt + max_new_tokens > context_length`, at prefill or at decode. With `context_length=16`:
| call | result |
| --- | --- |
| `generate_text_simple(prompt=20, context_size=8, max_new_tokens=8, use_cache=True)` | ok |
| the same with `max_new_tokens=9` | `IndexError` |
| `generate_text_simple(prompt=20, context_size=None, max_new_tokens=1, use_cache=True)` | `IndexError` |
| `generate_text_simple_stream(prompt=8, max_new_tokens=8)` | ok |
| the same with `max_new_tokens=9` | `IndexError` |
| `use_cache=False`, prompt 20, `max_new_tokens=50` | ok — that branch re-slices every step |
So slicing the streaming prefill to `context_size` alone would move the failure rather than remove it, which I think supports your `trim_input_tensor` suggestion applying to both functions. It does work end to end: trimming a 20-token prompt to `(1, 8)` and then streaming 8 tokens completes.
On the Chainlit example the symptom is different, because Qwen3 is RoPE rather than a positional table. An overlong prompt raises at `kv_cache/qwen3.py:298` (`x * cos`, "The size of tensor a (20) must match the size of tensor b (16) at non-singleton dimension 2"). Overflow during decode is the confusing one: `apply_rope(x, cos, sin, offset=16)` with `seq_len=1` returns shape `(1, 4, 0, 8)` silently, because `cos[16:17]` is an empty slice, so queries and keys lose the new token and the traceback surfaces at `masked_fill` (`qwen3.py:254`, 17 against 16) with no mention of context length. Nothing I tried produced wrong output silently — every overflow raised somewhere — but that message would be hard to trace back to this bug, which may be worth a line in whichever test you add.
(Qwen3 numbers use the `dummy_cfg_base` from `tests/test_qwen3.py` with `context_length` set to 16.)
Thanks @mira687 for running the reproduction and checking both model families. You're right: my non-streaming comparison was too narrow. It applies `context_size` to the initial slice, but its cached branch can still overflow during decoding. I agree the fix should cover both cached helpers.
At `f49cad7`, both loops also forward the newly generated token after the *last* requested output. With the current loop structure, the safe invariant is `prefill_len + max_new_tokens <= effective_context_len`. A focused fix would be:
```python
ctx_len = (model.cfg["context_length"] if context_size is None
else min(context_size, model.cfg["context_length"]))
# Validate ctx_len > 0 and max_new_tokens >= 0; return early for zero tokens.
# For the cached path, reject max_new_tokens >= ctx_len with a clear ValueError.
prefill = trim_input_tensor(input_ids, ctx_len, max_new_tokens)
logits = model(prefill, cache=cache)
```
Apply that prefill to `generate_text_simple_stream` and the `use_cache=True` branch of `generate_text_simple`, while keeping the latter's original `idx` for its returned sequence. The uncached branch can keep re-slicing on each step. Explicitly checking `context_size is None` also avoids treating `0` as the default. Longer generations would need cache eviction/sliding or a clear error; merely trimming the prompt to `context_size` is insufficient.
Suggested regression cases with model context 16: prompt 20 / budget 8 (prefill 8, succeeds), prompt 20 / budget 9 (prefill 7, succeeds), prompt 20 / `context_size=8` / budget 1 (prefill 7, succeeds), and budget 16 (clear validation error). For Qwen3, include both overlong prefill and decode-at-boundary cases, and assert the maximum forward position never exceeds the configured context. I have verified the budget against the source loops, but have not independently run PyTorch locally; your measured results supply the runtime reproduction.
Thanks again for the precise boundary cases. I would keep this as one issue because both paths share the same cache-budget assumption.
Ran your sketch as written (same setup: `f49cad7`, CPU, torch 2.2.2, `TEST_CONFIG` with `context_length=16`), since you flagged you had not run PyTorch locally. **All four of your cases behave exactly as you predicted, including the prefill lengths.**
| case | you predicted | measured (streaming / non-streaming) |
| --- | --- | --- |
| prompt 20, budget 8 | prefill 8, succeeds | prefill 8, 8 tokens yielded / 28 ids returned |
| prompt 20, budget 9 | prefill 7, succeeds | prefill 7, 9 tokens yielded / 29 ids returned |
| prompt 20, `context_size=8`, budget 1 | prefill 7, succeeds | prefill 7, 1 token yielded / 21 ids returned |
| prompt 20, budget 16 | validation error | `ValueError` before any forward pass |
Three things worth adding.
**Your invariant is tight, not conservative.** In the two succeeding cases above the model's `current_pos` finishes at exactly 16, i.e. on the limit. Unpatched, `prefill + max_new_tokens == ctx_len` is the last safe point at both ends of the range — `(prompt 8, budget 8)` and `(prompt 1, budget 15)` both complete, and `(8, 9)` and `(1, 16)` both raise `IndexError`. Your read of the extra forward is also confirmed by counting: `stream(prompt=4, max_new_tokens=8)` issues 9 forwards, shapes `(1,4)` then 8×`(1,1)`, and the last one's logits are discarded when the loop exits. So that wasted forward is what consumes the final slot; skipping it would relax the budget to `prefill + max_new_tokens - 1` and save a forward per call, but that changes the loop structure, and your invariant is correct for the loop as it stands.
**The explicit `ValueError` is worth it rather than redundant.** The rejection already exists today, inside `trim_input_tensor`, but as a bare `assert max_new_tokens < context_len`: at `context_len=16, max_new=16` it raises a completely message-less `AssertionError`. Under `python -O` the assert is stripped, `keep_len = max(1, ...)` clamps the prompt to **1 token**, and you get the same `IndexError` at the end of generation instead — the original bug, just deferred. Raising in the generate functions fixes both.
**`context_size=0` is worse than being treated as the default.** `idx[:, -0:]` returns the *entire* prompt (measured: 20 tokens in, 20 out), so a caller asking for zero context currently gets the maximum, then `IndexError`. Your explicit `is None` check plus rejecting `ctx_len <= 0` turns that into the `ValueError`.
One consequence the fix inherits, which may deserve a test of its own: the multi-turn Chainlit example already calls `trim_input_tensor(..., context_len=MODEL.cfg["context_length"], max_new_tokens=MAX_NEW_TOKENS)` with `MAX_NEW_TOKENS = 38912`, and every shipped Qwen3 config except `30B-A3B` has `context_length = 40960`. So `keep_len` is **2048** — the whole conversation is silently left-truncated to the last 2048 tokens of a 40960 window, and since `build_prompt_from_history` puts the system message first, that is the part dropped. Lowering `context_length` below 38912 to save memory hits the message-less `AssertionError` on the first message. Moving the trim inside the generate functions makes every caller inherit that silent shrink, so a warning when truncation actually happens might be worth it alongside the error.
(`trim_input_tensor` and the config numbers were run directly; I did not download Qwen3 weights or run Chainlit, so the 2048 is arithmetic on their own constants, not an observed session. Agree on keeping it one issue.)
Thanks for running the boundary cases and checking the Qwen3 path. Your measurements confirm the prefill budget, and I agree the fix should include a few safeguards:
- Validate the effective context length and generation budget explicitly, raising a descriptive ValueError before slicing or model execution. The current assertion in trim_input_tensor becomes a message-less AssertionError at the boundary and disappears under python -O.
- Reject context_size=0 as invalid. Python’s idx[:, -0:] keeps the full prompt, which is the opposite of a zero-token limit.
- Avoid silently dropping the system message when reserving a very large generation budget. In the multi-turn Qwen3 example, a 38,912-token output budget in a 40,960-token window leaves only 2,048 prompt tokens. The tensor-level generation helper cannot identify chat roles, so the chat caller should trim message history while preserving the system prompt and newest turns; the helper should also make any raw-token truncation visible to callers.
I’d add regressions for context length 16 with (1) budget 16 and (2) context_size 0, asserting a clear error before the first forward pass, and run the boundary validation with and without -O. For Qwen3, verify that history trimming retains the system message and emits a warning when it removes turns. Thanks again—this makes the report and proposed fix more precise.
Correction to something I got wrong above, before it ends up in the test: **`context_size=0` does not keep the full prompt.** I measured the bare slice instead of the code path. `ctx_len = context_size or model.cfg["context_length"]` short-circuits first, so `0` becomes `16` and the prefill actually fed to the model is `(1, 16)`, not `(1, 20)` — `idx[:, -0:]` is correct Python but unreachable here. Your conclusion survives (`0` is indistinguishable from `None`, and a caller asking for zero silently gets something else), but the guard has to sit before the `or`, and a regression asserting that the full prompt was forwarded would fail.
One thing for the other test, same reason: at `context_length=16, max_new_tokens=16` with asserts on, `trim_input_tensor` already raises **before the first forward** — 0 forwards in, `AssertionError` with `args=()`. So "clear error before the first forward pass" passes today on the timing and fails only on the message. Under `-O` it is the real case: `keep_len = max(1, 0) = 1`, a 20-token prompt is silently cut to 1, and `IndexError: index out of range in self` arrives 17 forwards later. Worth writing that assertion against `-O` specifically.
(`f49cad7`, CPU, torch 2.2.2, Python 3.12.14, `TEST_CONFIG` with `context_length=16`.)
#1090 is merged (`3aae8b9`, tree identical to `7e6219a`), so here is where the cases from this thread stand. Ran the new test file on CPU: **16 passed** — the four `llama3` params error in my environment on `nn.RMSNorm`, which needs torch >= 2.4 and is nothing to do with the change.
The guard is a clear improvement. Same call before and after, `context_length=8`, 6-token prompt, `max_new_tokens=16`, no EOS:
| | before (`f49cad7`) | after |
| --- | --- | --- |
| `Qwen3Model` | `RuntimeError: The size of tensor a (9) must match the size of tensor b (8) at non-singleton dimension 3` | `ValueError: Sequence length 9 exceeds the model's context length of 8 tokens.` |
| `GPTModel` | `IndexError: index out of range in self` | the same `ValueError` |
Two things that case still does, so this issue probably stays open:
**The title case is unchanged.** `generate_text_simple_stream` accepts `context_size` at `kv_cache/generate.py:34` and never references it in the body, so passing `context_size=4` gives the identical result and the message names 8 — the number the caller was trying to override.
**The budget is still not checked on arrival.** The stream yields 3 tokens and *then* raises, before and after. The call is decidable on entry (6 + 16 > 8), and the rule is already written down in the repo: `trim_input_tensor` asserts `max_new_tokens < context_len` and left-truncates the prompt to `context_len - max_new_tokens` (`pkg/llms_from_scratch/generate.py:9`). Its only call site in the repo is `qwen3-chat-interface-multiturn.py:145`; `kv_cache/generate.py:6` and `kv_cache_batched/generate.py:6` re-export it `# noqa: F401` and neither calls it.
Separate from the guard but easier to see now that the message is explicit: on that same model `generate_text_simple(..., max_new_tokens=16, use_cache=False)` returns 22 ids without complaint, because it re-slices `idx[:, -ctx_len:]` every step, while `use_cache=True` raises. Same function, same arguments, two answers about whether the call is legal.
`test_large_generation_budget_can_stop_early_at_eos` covers the branch where EOS arrives first. Its twin is the no-EOS case above, if it is worth a second assertion.
Just also as a heads-up regarding `context_size=None` and never using it. It's been a while, but as far as I recall, it was intentional because the main chapters use a `context_size` parameter. This way, people can use the kv-cache generation functions as a drop-in replacement without breaking the code. I am leaning towards documenting it instead of removing it.
Thanks — the `warnings.warn` text in #1094 covers the part that mattered ("Prompts are not truncated"), so nothing further from me on `context_size` itself.
Two small leftovers from the batch, both measured at `8943bb5`.
**The new warning fires inside the test suite.** `test_streaming_generation_matches_cached_generation` passes `context_size=16` to the stream at `tests/test_kv_cache_gpt2.py:118`, so that file now runs `11 passed, 1 warning`, with `stacklevel=2` pointing at the call. Harmless as configured — there is no `filterwarnings` setting in the repo — and the argument is not what that test checks anyway, since it compares streamed against cached output. Dropping it would keep the run clean.
**#1091 moved the budget boundary and `trim_input_tensor` did not follow.** Largest `prompt + max_new_tokens` that completes, streaming, `context_length=16`, no EOS, measured at each commit:
| commit | largest prompt + budget | forwards for (4, 8) |
| --- | --- | --- |
| `3aae8b9` (#1090) | 16 | 9 |
| `aa32ed1` (#1091) | 17 | 8 |
| `f9313ea` / `8943bb5` | 17 | 8 |
Both ends moved: `(1, 16)` and `(15, 2)` complete now and both raised at `3aae8b9`. That is the slot dropping the final forward freed.
`trim_input_tensor` still encodes the pre-#1091 rule, though — it raises on `max_new_tokens >= context_len` and keeps `context_len - max_new_tokens`, which is now one token tight at both ends. At `context_length=16`, `stream(prompt=1, max_new_tokens=16)` yields 16 tokens while `trim_input_tensor(x, 16, 16)` raises `ValueError`; and at budget 9 the largest prefill that completes is 8 (9 raises) while the helper keeps 7. The one-token over-trim is nothing in the multi-turn Chainlit caller (`qwen3-chat-interface-multiturn.py:145` — 2048 vs 2049 of a 40960 window), but the rejection turns a call that now works into an error. `max_new_tokens <= context_len` with `keep_len = context_len - max_new_tokens + 1` matches the merged loop at both ends.
(`8943bb5`, CPU, torch 2.2.2, Python 3.12.14, `TEST_CONFIG` from `tests/test_kv_cache_gpt2.py` with `context_length=16`. The `11 passed` is the gpt2 test file alone — `test_kv_cache_context_limits.py` does not import in my environment, no `tiktoken`.)