35231d4ff4
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.
78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
"""Batched steady-state throughput bench through ``FastLanguageModel`` +
|
|
``UNSLOTH_FAST_INFERENCE=1``. First ``generate`` call primes CUDA graphs;
|
|
subsequent calls report steady state. Compare against April CLI-only
|
|
numbers for the same workload."""
|
|
|
|
import os, sys, time
|
|
from pathlib import Path
|
|
_REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(_REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_REPO_ROOT))
|
|
|
|
import argparse
|
|
import torch
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--model", default="unsloth/Qwen3-4B-Base")
|
|
p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
|
|
p.add_argument("--n_prompts", type=int, default=8)
|
|
p.add_argument("--max_new_tokens", type=int, default=64)
|
|
p.add_argument("--max_batch_size", type=int, default=16)
|
|
p.add_argument("--max_seq_length", type=int, default=1024)
|
|
p.add_argument("--n_rounds", type=int, default=3)
|
|
args = p.parse_args()
|
|
|
|
os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1")
|
|
import unsloth
|
|
from unsloth import FastLanguageModel
|
|
|
|
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
|
|
|
|
model, tok = FastLanguageModel.from_pretrained(
|
|
model_name=args.model,
|
|
max_seq_length=args.max_seq_length,
|
|
dtype=dtype,
|
|
load_in_4bit=False,
|
|
fast_inference=True,
|
|
max_batch_size=args.max_batch_size,
|
|
)
|
|
print(f"[bench] model={args.model} dtype={args.dtype}")
|
|
prompts = [f"In one sentence, a fact about {t} is" for t in [
|
|
"the moon", "gravity", "the ocean", "the sun", "honey",
|
|
"rain", "trees", "mountains"
|
|
][:args.n_prompts]]
|
|
|
|
class _SP:
|
|
max_tokens = args.max_new_tokens
|
|
temperature = 0.0
|
|
|
|
# Warmup round — captures CUDA graphs.
|
|
print("[bench] warmup (CUDA graph capture)...")
|
|
t0 = time.perf_counter()
|
|
_ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
|
|
print(f"[bench] warmup wall: {time.perf_counter() - t0:.2f}s")
|
|
|
|
walls = []
|
|
tok_counts = []
|
|
for r in range(args.n_rounds):
|
|
torch.cuda.synchronize()
|
|
t1 = time.perf_counter()
|
|
outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
|
|
torch.cuda.synchronize()
|
|
dt = time.perf_counter() - t1
|
|
n_tok = sum(len(o.outputs[0].token_ids) for o in outs)
|
|
walls.append(dt)
|
|
tok_counts.append(n_tok)
|
|
print(f"[bench] round {r}: {n_tok} toks in {dt:.2f}s -> {n_tok/dt:.1f} tok/s")
|
|
|
|
if walls:
|
|
wall_med = sorted(walls)[len(walls) // 2]
|
|
tok_med = tok_counts[len(walls) // 2]
|
|
print(f"[bench] median: {tok_med} toks in {wall_med:.2f}s "
|
|
f"=> {tok_med / wall_med:.1f} tok/s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|