Compile a model

Quantization saves memory, but dequantizing weights inside every matmul can make generation slower than in bf16. torch.compile fuses that work into optimized kernels and recovers much of the speed.

Compile automatically in generate

generate() compiles the model’s forward pass automatically when all of these hold.

• The quantization method supports compilation (compressed-tensors, Fine-grained FP8, GGUF, NVFP4, and torchao).
• The model uses a cache with a fixed size, like the static cache.
• The model runs on CUDA, XPU, Neuron, or TPU.
• No part of the model is offloaded to the CPU or disk.

Pass cache_implementation="static" to generate(). The example below uses torchao int8 weight-only quantization.

import torch
from torchao.quantization import Int8WeightOnlyConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig

model_id = "Qwen/Qwen3-0.6B"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=TorchAoConfig(quant_type=Int8WeightOnlyConfig()),
    device_map="cuda",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=32, cache_implementation="static")
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

The first calls are slow while the model compiles. Later calls with the same shapes reuse the compiled kernels.

Pass disable_compile=True to generate() to turn automatic compilation off.