"""Main-thread synchronous driver for `ContinuousBatchProcessor`. `ContinuousBatchingManager.start()` spawns a background thread that owns the decode loop. That thread conflicts with two things we want to enable here: 1. `torch.compile(mode="reduce-overhead")` which uses `cudagraph_trees` and requires main-thread TLS. 2. Raw `torch.cuda.CUDAGraph` capture / replay on the decode forward. The manager's dead `warmup()` path suggests CB was supposed to grow CUDA graph support upstream but `init_continuous_batching` raises `NotImplementedError` on `use_cuda_graph=True`. This driver side-steps the whole manager thread, so the cudagraph integration point is now available. Key fixed-shape invariant: with `slice_inputs=False`, the full pre-allocated tensor buffers (input_ids, position_ids, cu_seq_lens_*, attention_mask, read_index / write_index) are returned as views of the same storage every step, so their shapes are constant across iterations. That is the precondition for graph replay / cudagraph_trees to be safe. Greedy sampling only (`do_sample=False`). `torch.multinomial` is not graph-friendly. Usage: driver = SyncCBDriver(model, gen_config, CBSyncConfig(compile_mode="reduce-overhead")) # Reuse across many rollouts (cache / compiled forward stay warm): for batch in batches: driver.add_requests(batch) out = driver.drive_until_empty() """ from __future__ import annotations import queue import time from dataclasses import dataclass, field from typing import Optional import torch from transformers.generation.configuration_utils import GenerationConfig from transformers.generation.continuous_batching import ( PagedAttentionCache, RequestStatus, ) from transformers.generation.continuous_batching.continuous_api import ( ContinuousBatchProcessor, ContinuousBatchingManager, ) from transformers.generation.continuous_batching.scheduler import FIFOScheduler @dataclass class CBSyncConfig: """Tunables for the sync driver.""" max_new_tokens: int = 512 # torch.compile mode for the model forward. None = eager. # "reduce-overhead" triggers cudagraph_trees, which captures a CUDA # graph per unique input shape and replays it afterwards. compile_mode: Optional[str] = None do_sample: bool = False # greedy only (graph-safe) eos_token_id: Optional[int] = None pad_token_id: Optional[int] = None # `slice_inputs=False` would give fixed shapes every step (graph-friendly) # but forces every decode step to do `max_batch_tokens` tokens of work, # which at max_batch_tokens=8192 is ~256x more than the real decode batch. # Prefer `slice_inputs=True` (natural shapes) and let torch.compile bucket # per shape. Steady-state decode has one shape so most steps replay the # same graph anyway. slice_inputs: bool = True # Paged cache upper bounds. max_batch_tokens: int = 8192 num_blocks: int = 8192 # `torch._dynamo.config.cache_size_limit`: raise when varying shapes. dynamo_cache_size_limit: int = 256 on_step: Optional[callable] = field(default = None) class SyncCBDriver: """Main-thread driver that owns the PagedAttentionCache, ContinuousBatchProcessor, and optionally a `torch.compile`-compiled forward. Reusable across multiple `drive_until_empty` calls -- the cache and compiled forward stay warm between rounds. """ def __init__( self, model: torch.nn.Module, generation_config: GenerationConfig, cfg: CBSyncConfig, ): self.model = model.eval() self.cfg = cfg gc = GenerationConfig.from_dict(generation_config.to_dict()) gc.do_sample = cfg.do_sample if cfg.max_new_tokens: gc.max_new_tokens = cfg.max_new_tokens if cfg.eos_token_id is not None: gc.eos_token_id = cfg.eos_token_id if cfg.pad_token_id is not None: gc.pad_token_id = cfg.pad_token_id gc.max_batch_tokens = cfg.max_batch_tokens gc.num_blocks = cfg.num_blocks self.generation_config = gc self.manager = ContinuousBatchingManager( model = self.model, generation_config = gc, manual_eviction = False, streaming = False, slice_inputs = cfg.slice_inputs, ) self.cache = PagedAttentionCache( self.model.config, gc, self.model.device, self.model.dtype, tp_size = getattr(self.model, "_tp_size", None), ) self.batch_processor = ContinuousBatchProcessor( self.cache, self.model.config, gc, self.manager.input_queue, self.manager.output_queue, self.manager.stop_event, self.model.device, self.model.dtype, FIFOScheduler(self.cache), streaming = False, manual_eviction = False, slice_inputs = cfg.slice_inputs, ) self.manager.batch_processor = self.batch_processor # torch.compile on the model forward. With slice_inputs=True the # forward sees varying shapes (prefill bursts + decode steady state); # `dynamic=True` lets Inductor bucket per shape without re-tracing # every call, and `mode="reduce-overhead"` wraps each bucket in a # CUDA graph replay path. if cfg.compile_mode: import torch._dynamo torch._dynamo.config.cache_size_limit = cfg.dynamo_cache_size_limit # GRPO's `requires_grad_` issue doesn't apply here (eval mode). try: torch._dynamo.config.allow_unspec_int_on_nn_module = True except AttributeError: pass print( f"[cb_sync] torch.compile(model, mode='{cfg.compile_mode}', " f"dynamic=True)" ) self.model.forward = torch.compile( self.model.forward, mode = cfg.compile_mode, dynamic = True, fullgraph = False, ) self._step_count = 0 def add_requests(self, prompt_ids_list: list[list[int]]) -> list[str]: return [self.manager.add_request(ids) for ids in prompt_ids_list] def drive_until_empty(self) -> dict[str, list[int]]: """Run the decode loop until every request finishes. Returns a dict {request_id: generated_token_ids}. Reusable across calls on the same driver -- the paged cache and the compiled forward stay warm. """ results: dict[str, list[int]] = {} while True: if ( self.manager.input_queue.empty() and not self.batch_processor.has_pending_requests() ): break if not self.batch_processor.prepare_next_batch(): break # With `reduce-overhead`, Inductor captures a CUDA graph on the # first call with a given shape signature and replays it after. # Our shapes are constant (slice_inputs=False), so the second call # is already replaying. No per-step torch.cuda.synchronize() -- # the replay itself fences appropriately. self.manager._generation_step(self.batch_processor) self.batch_processor.update_batch() self._step_count += 1 if self.cfg.on_step is not None: self.cfg.on_step(self._step_count, 0) # Drain output_queue as requests finish. while True: try: out = self.manager.output_queue.get_nowait() except queue.Empty: break if out.status == RequestStatus.FINISHED: results[out.request_id] = out.generated_tokens while True: try: out = self.manager.output_queue.get_nowait() except queue.Empty: break if out.status == RequestStatus.FINISHED: results[out.request_id] = out.generated_tokens return results def close(self): self.cache = None self.batch_processor = None self.manager.batch_processor = None # Simple microbench harness so the file is runnable standalone. if __name__ == "__main__": import argparse import json import os import sys from pathlib import Path HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) import flash_attn_fa4_shim # noqa: E402 flash_attn_fa4_shim.apply() parser = argparse.ArgumentParser() parser.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") parser.add_argument("--n_prompts", type = int, default = 32) parser.add_argument("--n_rounds", type = int, default = 2) parser.add_argument("--max_new_tokens", type = int, default = 512) parser.add_argument("--attn_impl", default = "paged_attention") parser.add_argument( "--compile_mode", default = None, choices = [ None, "default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs", ], ) parser.add_argument("--max_batch_tokens", type = int, default = 8192) parser.add_argument("--num_blocks", type = int, default = 8192) parser.add_argument("--lora_adapter", default = None) parser.add_argument("--stats_path", required = True) args = parser.parse_args() from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig tok = AutoTokenizer.from_pretrained(args.model_name) if tok.pad_token is None: tok.pad_token = tok.eos_token model = AutoModelForCausalLM.from_pretrained( args.model_name, dtype = torch.bfloat16, attn_implementation = args.attn_impl, ).to("cuda") model.eval() if args.lora_adapter: from peft import PeftModel model = PeftModel.from_pretrained( model, str(Path(args.lora_adapter).resolve()), is_trainable = False ) model.eval() from unsloth_grpo_common import ( SYSTEM_PROMPT, apply_chat_template_to_tokenizer, ) from datasets import load_dataset apply_chat_template_to_tokenizer(tok) ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train") ds = ds.shuffle(seed = 3407).select(range(args.n_prompts)) messages = [ [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": x["prompt"]}, ] for x in ds ] prompt_ids = [ tok.apply_chat_template(m, add_generation_prompt = True, tokenize = True) for m in messages ] gc_cfg = GenerationConfig( max_new_tokens = args.max_new_tokens, do_sample = False, pad_token_id = tok.pad_token_id, bos_token_id = tok.bos_token_id, eos_token_id = tok.eos_token_id, use_cache = True, ) cfg = CBSyncConfig( max_new_tokens = args.max_new_tokens, compile_mode = args.compile_mode, max_batch_tokens = args.max_batch_tokens, num_blocks = args.num_blocks, eos_token_id = tok.eos_token_id, pad_token_id = tok.pad_token_id or tok.eos_token_id, ) torch.cuda.reset_peak_memory_stats() # One driver, multiple rounds -- cache + compiled forward stay warm. driver = SyncCBDriver(model, gc_cfg, cfg) # Warmup (first 16 prompts). With compile, this amortizes the capture. print("[cb_sync] warmup...") driver.add_requests(prompt_ids[:16]) _ = driver.drive_until_empty() torch.cuda.synchronize() wall_times = [] total_decoded = 0 for r in range(args.n_rounds): torch.cuda.synchronize() t0 = time.perf_counter() driver.add_requests(prompt_ids) results = driver.drive_until_empty() torch.cuda.synchronize() wall_times.append(time.perf_counter() - t0) total_decoded = sum(len(v) for v in results.values()) print( f"[cb_sync] round {r}: {wall_times[-1]:.2f}s, {total_decoded} tokens, " f"{total_decoded / wall_times[-1]:.1f} tok/s" ) med = sorted(wall_times)[len(wall_times) // 2] out = { "backend": "cb_sync_driver", "compile_mode": args.compile_mode, "attn_impl": args.attn_impl, "lora_adapter": args.lora_adapter, "n_prompts": args.n_prompts, "n_decoded_tokens": total_decoded, "wall_times_s": wall_times, "median_wall_s": med, "decode_tps": total_decoded / med if med else 0, "max_new_tokens": args.max_new_tokens, "peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3, } os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True) with open(args.stats_path, "w") as f: json.dump(out, f, indent = 2) print(json.dumps(out, indent = 2)) driver.close() os._exit(0)