Setting UNSLOTH_FAST_INFERENCE=1 makes
FastLanguageModel.from_pretrained(..., fast_inference=True) return the
flex attention + paged KV + CUDA-graph backend instead of vLLM for
Qwen3, Llama-3, and Gemma-4-E2B-it. Default "0" keeps the existing vLLM
path.
The flex engine lives in a new unsloth.inference subpackage so the
FastLanguageModel wiring and any external caller can reach it without
going through scripts/. The user-facing GRPO notebooks for Qwen3,
Llama-3, and Gemma-4 run end-to-end with no code edits, just the env
var.
New subpackage
- unsloth/inference/flex_paged_attention.py: paged KV cache and
block-mask helpers.
- unsloth/inference/flex_qwen3_llama.py: FlexInference for Qwen3 and
Llama-3.2. tokenize() skips pre-populated input_ids so the engine
can accept vLLM-style token-id prompts.
- unsloth/inference/flex_gemma4.py: FlexGemma4Inference. Text-shell
extraction from Gemma4ForConditionalGeneration, per-layer-type
sliding-window masks, KV sharing for layers 15-34.
- unsloth/inference/flex_engine.py: FlexEngine with the vLLM LLM
surface. Arch dispatch, deep-copy of the HF model for colocate
rollout so training forward stays intact, lazy pristine-base copy
on first bind_peft_model, torch.amp.autocast("cuda", dtype=...)
wrapping for the whole generate path, auto-tune for Triton block
sizes per GPU capability and head_dim band. FA4 default OFF
because the sm_100 FLASH backend crashes on short prompts; users
can opt in with fa4_prefill=True.
- unsloth/inference/vllm_shim.py: LoRARequest, RequestOutput,
CompletionOutput dataclasses plus save_lora / load_lora that
mirror unsloth_zoo.vllm_utils line for line so the TRL GRPO patch
sees the same attribute shape.
Wiring
- unsloth/models/loader.py: UNSLOTH_FAST_INFERENCE=1 bypasses the
vLLM import guard in FastLanguageModel.from_pretrained (L356) and
FastModel.from_pretrained (L982).
- unsloth/models/llama.py: the elif not fast_inference load branch
now also handles the flex case, pops max_batch_size from kwargs
(FlexEngine-only), snapshots the model into an inference copy
BEFORE Unsloth's post-patch so flex-attention patching lands on a
clean HF layout, and after the tokenizer loads constructs
FlexEngine and attaches .vllm_engine / .fast_generate /
.fast_generate_batches. FastLlamaModel.get_peft_model calls
engine.bind_peft_model(peft) right after patch_peft_fast_inference
so state_dict() reads the training LoRA tensors.
- unsloth/models/vision.py: same pattern for FastBaseModel. Gemma-4
is allowed through the vision-model gate when the flex backend is
selected; the engine extracts the text shell from
Gemma4ForConditionalGeneration internally.
- unsloth/models/_utils.py: patch_peft_fast_inference picks
unsloth.inference.vllm_shim.save_lora / load_lora when the engine
is a FlexEngine, so the flex path never imports
vllm.lora.request. fast_inference_setup skips patch_vllm() when
UNSLOTH_FAST_INFERENCE=1.
Smoke tests (B200 sm_100)
- Qwen3-4B-Base bf16 (no LoRA and with get_peft_model LoRA)
- Qwen3-4B-Base fp16
- Llama-3.2-3B-Instruct bf16
- gemma-4-E2B-it bf16
all load through FastLanguageModel + UNSLOTH_FAST_INFERENCE=1 and
generate coherent completions.
Batched steady-state throughput via FastLanguageModel (autocast ON,
n_prompts=8 max_new_tokens=64 max_batch_size=16 max_seq_length=1024):
| Model | Flex CLI | Integration | Delta |
|----------------------------|-----------|-------------|-------|
| Qwen3-4B-Base | 819 tok/s | 595 tok/s | -27% |
| Llama-3.2-3B-Instruct | 1163 | 1357 | +17% |
Qwen3's regression is outside the 10% budget; most likely cause is
q_norm / k_norm being promoted to fp32 under autocast. Tracked for
a follow-up; the integration is functional and GRPO-usable today.
Note
The flex engine source files live here in unsloth/inference/. The
CLI benchmark drivers in scripts/benchmarks/ land in a separate PR
(#5108) which also carries the Gemma-4 and Llama-3 regression runs.
This PR depends only on what is already on main.
Known follow-ups
- FA4 on sm_100 needs a more robust enable rule before it can be
default-on.
- Memory: 2x base-model VRAM without LoRA, 3x with LoRA.
- FlexEngine.sleep() / .wake_up() are no-op stubs; real CPU-offload
parity with UNSLOTH_VLLM_STANDBY is a separate PR.
Daniel Han