unsloth/scripts/benchmarks/cb_vs_vllm_generation.py

"""Standalone generation microbenchmark: vLLM vs transformers continuous batching.

Measures prompt-tokens/s, decode tokens/s, and end-to-end wall-clock on
`N` prompts sampled from DAPO-Math-17k with the GRPO chat template applied.

Run:
    CUDA_VISIBLE_DEVICES=2 python scripts/cb_vs_vllm_generation.py \
        --backend vllm   --stats_path logs/vllm_gen.json
    CUDA_VISIBLE_DEVICES=2 python scripts/cb_vs_vllm_generation.py \
        --backend tpaged --stats_path logs/cb_gen.json

One backend per process (both engines are GPU-greedy). Results are then
combined offline.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import torch  # noqa: E402


def build_prompts(tokenizer, n_prompts):
    from unsloth_grpo_common import (
        apply_chat_template_to_tokenizer,
        SYSTEM_PROMPT,
    )
    from datasets import load_dataset

    apply_chat_template_to_tokenizer(tokenizer)
    ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
    ds = ds.shuffle(seed = 3407).select(range(n_prompts))
    messages = [
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": x["prompt"]},
        ]
        for x in ds
    ]
    prompts_text = [
        tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = False)
        for m in messages
    ]
    prompt_ids = [
        tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = True)
        for m in messages
    ]
    return prompts_text, prompt_ids


def run_vllm(args):
    import os as _os

    _os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
    from unsloth import FastLanguageModel

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = args.model_name,
        max_seq_length = args.max_seq_length,
        load_in_4bit = False,
        fast_inference = True,
        max_lora_rank = 32,
        gpu_memory_utilization = args.gpu_memory_utilization,
    )
    prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)

    from vllm import SamplingParams

    sp = SamplingParams(
        temperature = 1.0,
        min_p = 0.1,
        top_p = 1.0,
        top_k = -1,
        seed = 3407,
        max_tokens = args.max_new_tokens,
        stop = [tokenizer.eos_token],
        include_stop_str_in_output = True,
    )

    # Warmup on 16 prompts then discard.
    warmup_text = prompts_text[:16]
    _ = model.fast_generate(warmup_text, sampling_params = sp, lora_request = None)
    torch.cuda.synchronize()

    # Three measured rounds on the full batch.
    n_prompt_tokens = sum(len(p) for p in prompt_ids)
    wall_times = []
    total_decoded = None
    for _ in range(args.n_rounds):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        outputs = model.fast_generate(
            prompts_text, sampling_params = sp, lora_request = None
        )
        torch.cuda.synchronize()
        wall_times.append(time.perf_counter() - t0)
        decoded = sum(len(o.outputs[0].token_ids) for o in outputs)
        total_decoded = decoded

    med = sorted(wall_times)[len(wall_times) // 2]
    return {
        "backend": "vllm",
        "n_prompts": args.n_prompts,
        "n_prompt_tokens": n_prompt_tokens,
        "n_decoded_tokens": total_decoded,
        "wall_times_s": wall_times,
        "median_wall_s": med,
        "prompt_tps": n_prompt_tokens / med,
        "decode_tps": (total_decoded or 0) / med,
        "max_new_tokens": args.max_new_tokens,
    }


def run_tpaged(args):
    # Vanilla HF load. Unsloth's Qwen3Attention monkey-patch does not
    # compose with the `paged|<impl>` functional attention interface.
    from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name,
        dtype = torch.bfloat16,
        attn_implementation = args.attn_impl,
    ).to("cuda")
    model.eval()
    prompts_text, prompt_ids = build_prompts(tokenizer, args.n_prompts)

    gen_config = GenerationConfig(
        max_new_tokens = args.max_new_tokens,
        do_sample = True,
        temperature = 1.0,
        top_p = 1.0,
        min_p = 0.1,
        pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id,
        bos_token_id = tokenizer.bos_token_id,
        eos_token_id = tokenizer.eos_token_id,
        use_cache = True,
    )
    # Raise the paged-cache upper bounds; defaults (256 / 4096) throttle CB.
    gen_config.max_batch_tokens = args.max_batch_tokens
    gen_config.num_blocks = args.num_blocks

    # Warmup on 16 prompts.
    warmup_ids = prompt_ids[:16]
    with torch.inference_mode():
        _ = model.generate_batch(
            warmup_ids, generation_config = gen_config, progress_bar = False
        )
    torch.cuda.synchronize()

    n_prompt_tokens = sum(len(p) for p in prompt_ids)
    wall_times = []
    total_decoded = None
    for _ in range(args.n_rounds):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.inference_mode():
            outputs = model.generate_batch(
                prompt_ids, generation_config = gen_config, progress_bar = False
            )
        torch.cuda.synchronize()
        wall_times.append(time.perf_counter() - t0)
        decoded = sum(len(v.generated_tokens) for v in outputs.values())
        total_decoded = decoded

    med = sorted(wall_times)[len(wall_times) // 2]
    return {
        "backend": "tpaged",
        "attn_impl": args.attn_impl,
        "n_prompts": args.n_prompts,
        "n_prompt_tokens": n_prompt_tokens,
        "n_decoded_tokens": total_decoded,
        "wall_times_s": wall_times,
        "median_wall_s": med,
        "prompt_tps": n_prompt_tokens / med,
        "decode_tps": (total_decoded or 0) / med,
        "max_new_tokens": args.max_new_tokens,
    }


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--backend", choices = ["vllm", "tpaged"], required = True)
    p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
    p.add_argument("--max_seq_length", type = int, default = 2048)
    p.add_argument("--n_prompts", type = int, default = 64)
    p.add_argument("--n_rounds", type = int, default = 3)
    p.add_argument("--max_new_tokens", type = int, default = 1024)
    p.add_argument("--gpu_memory_utilization", type = float, default = 0.8)
    p.add_argument("--attn_impl", default = "sdpa")
    p.add_argument("--max_batch_tokens", type = int, default = 8192)
    p.add_argument("--num_blocks", type = int, default = 16384)
    p.add_argument("--stats_path", required = True)
    return p.parse_args()


def main():
    args = parse_args()
    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)

    torch.cuda.reset_peak_memory_stats()
    if args.backend == "vllm":
        out = run_vllm(args)
    else:
        out = run_tpaged(args)

    out["peak_memory_gb"] = torch.cuda.max_memory_allocated() / 1024**3
    with open(args.stats_path, "w") as f:
        json.dump(out, f, indent = 2)
    print(json.dumps(out, indent = 2))


if __name__ == "__main__":
    main()