TopolearnGeneration latency and memory are dominated by the KV cache; streaming is essential for responsive UIs.
During generation, each new token attends to all previous tokens. Recomputing their keys and values every step would be quadratic, so models store them in a KV cache: the first forward pass (prefill) processes the prompt and fills the cache; each decode step then processes just one new token and appends its keys and values.
The cache grows with layers x heads x sequence length x batch size, and for long contexts it can exceed the weights in memory. Transformers offers several cache classes:
DynamicCache(default): grows as needed.StaticCache(cache_implementation="static"): preallocates a fixed maximum length so shapes never change, which enablestorch.compileto produce a much faster decode step.- Offloaded and quantized caches trade speed for memory on long contexts; sliding-window models keep only recent tokens.
You can reuse a cache across calls: pass past_key_values= to continue from a shared prefix such as a long system prompt instead of reprocessing it.
Streaming shows tokens as they are generated. TextStreamer(tok) prints to stdout; TextIteratorStreamer yields text chunks in your code, but generate blocks, so run it in a background thread and iterate the streamer in the main thread, for example to send server-sent events.
For high-throughput serving with many users, dedicated engines (vLLM, TGI, SGLang) or transformers serve with continuous batching and paged attention are more efficient than calling generate per request.
Example
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
repo = "HuggingFaceTB/SmolLM2-135M-Instruct"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo)
msgs = [{"role": "user", "content": "Name three planets."}]
inputs = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt")
streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
thread = Thread(target=model.generate,
kwargs=dict(**inputs, max_new_tokens=40, streamer=streamer))
thread.start()
for chunk in streamer: # text arrives as it is generated
print(chunk, end="", flush=True)
thread.join()
# static cache: fixed shapes, compile-friendly decode step
out = model.generate(**inputs, max_new_tokens=40, cache_implementation="static")Remember
- Prefill fills the KV cache; each decode step adds one token's keys and values
- StaticCache gives fixed shapes for torch.compile; DynamicCache is the flexible default
- Stream with TextIteratorStreamer and generate in a background thread
Watch out for
- Iterating a TextIteratorStreamer in the same thread that calls generate, which blocks until generation is over.