import pickle import gc import torch import torch.distributed as dist from multiprocessing.synchronize import Event from multiprocessing.shared_memory import SharedMemory import sys from ..config import Config from .sequence import Sequence from ..layers.sampler import Sampler from ..utils.context import set_context, get_context, reset_context import socket def find_available_port(start_port: int = 2333, max_attempts: int = 100) -> int: """Find an available port starting from start_port. Args: start_port: The starting port number to check max_attempts: Maximum number of ports to try Returns: An available port number Raises: RuntimeError: If no available port is found within max_attempts """ for i in range(max_attempts): port = start_port + i try: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) s.bind(('localhost', port)) return port except OSError: # Port is in use, try next one continue raise RuntimeError(f"Could not find an available port starting from {start_port} after {max_attempts} attempts") class ModelRunner: def __init__(self, config: Config, rank: int, event: Event | list[Event], model_object=None, graph_pool_handle=None): # Enable capturing scalar outputs to avoid graph breaks from Tensor.item() calls torch._dynamo.config.capture_scalar_outputs = True self.config = config hf_config = config.hf_config self.block_size = config.kvcache_block_size self.enforce_eager = config.enforce_eager self.world_size = config.tensor_parallel_size self.rank = rank self.event = event if self.world_size > 1: dist_port = find_available_port() print(f"[debug]dist_port: {dist_port}") # Use gloo backend on Windows, nccl on Linux/other platforms backend = "gloo" if sys.platform == "win32" else "nccl" dist.init_process_group(backend, f"tcp://127.0.0.1:{dist_port}", world_size=self.world_size, rank=rank) torch.cuda.set_device(rank) else: if torch.cuda.is_available(): torch.cuda.set_device(0) default_dtype = torch.get_default_dtype() # Use dtype instead of deprecated torch_dtype config_dtype = getattr(hf_config, 'dtype', getattr(hf_config, 'torch_dtype', None)) # Validate and convert config_dtype to a valid torch floating-point dtype # Default to bfloat16 for CUDA (required for Flash Attention 2) if config_dtype is None: config_dtype = torch.bfloat16 elif isinstance(config_dtype, str): # Convert string dtype to torch dtype dtype_map = { 'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16, 'float64': torch.float64, 'torch.float32': torch.float32, 'torch.float16': torch.float16, 'torch.bfloat16': torch.bfloat16, 'torch.float64': torch.float64, } config_dtype = dtype_map.get(config_dtype.lower(), torch.bfloat16) elif not isinstance(config_dtype, torch.dtype) or not config_dtype.is_floating_point: # If not a valid floating-point torch dtype, default to bfloat16 config_dtype = torch.bfloat16 self.dtype = config_dtype # Save for later use self._runtime_ready = False self._graph_cache = {} self._graph_cache_order = [] self._logits_bias_cache = {} self._sampling_generator = None self._runtime_signature = None self._graph_pool_seed = graph_pool_handle self._guard_counts = {} self._guard_seen_details = set() torch.set_default_dtype(config_dtype) if model_object is None: raise RuntimeError( "nanovllm now requires a preloaded MMGP model object. " "Pass model_object=... when creating LLM." ) self.model = model_object self.sampler = Sampler() # Pre-allocate buffers for sampling (optimization: avoid repeated tensor creation) # Must be called before model execution paths that use these buffers. self._allocate_sample_buffers() torch.set_default_device("cpu") torch.set_default_dtype(default_dtype) if self.world_size > 1: if rank == 0: self.shm = SharedMemory(name="nanovllm", create=True, size=2**20) dist.barrier() else: dist.barrier() self.shm = SharedMemory(name="nanovllm") self.loop() def ensure_runtime_ready(self): if self._runtime_ready: # In eager mode MMGP may move parameter storage between CPU/GPU after prefill. # That changes data_ptr() without invalidating the live KV cache or sequence state. # Resetting here drops the prompt context and makes legacy decode drift off-topic. if self.enforce_eager: return current_sig = self._get_graph_capture_signature() if current_sig == self._runtime_signature: return self._note_guard("runtime_reprepare_signature_change") self.reset_runtime_state() if self.model is None: raise RuntimeError("LLM model object is not available.") self._tie_word_embeddings_if_needed() self.allocate_kv_cache() self._prepare_model_sequence_state() if not self.enforce_eager: self.capture_cudagraph() self._runtime_ready = True self._runtime_signature = self._get_graph_capture_signature() def reset_generation_state(self): if self.model is None: return try: for module in self.model.modules(): reset_sequence_state = getattr(module, "reset_sequence_state", None) if callable(reset_sequence_state): reset_sequence_state() continue if hasattr(module, "conv_state"): module.conv_state = None if hasattr(module, "recurrent_state"): module.recurrent_state = None except Exception: pass self._logits_bias_cache.clear() reset_context() def _prepare_model_sequence_state(self): if self.model is None: return model_device = self._get_model_device() for module in self.model.modules(): prepare = getattr(module, "prepare_sequence_state", None) if callable(prepare): prepare(self.config.max_num_seqs, model_device, self.dtype) def _get_tied_embeddings(self): if self.model is None: return None hf_config = getattr(self.config, "hf_config", None) if not bool(getattr(hf_config, "tie_word_embeddings", False)): return None lm_head = getattr(self.model, "lm_head", None) embed = getattr(self.model, "embed_tokens", None) if lm_head is None or embed is None: return None lm_w = getattr(lm_head, "weight", None) emb_w = getattr(embed, "weight", None) if lm_w is None or emb_w is None: return None return lm_head, embed def _tie_word_embeddings_if_needed(self): tied = self._get_tied_embeddings() if tied is None: return lm_head, embed = tied if lm_head.weight is not embed.weight: lm_head.register_parameter("weight", embed.weight) if lm_head.weight is not embed.weight: raise RuntimeError("Failed to retie lm_head.weight with embed_tokens.weight.") def reset_runtime_state(self): if not self._runtime_ready: return # Clear attention KV cache refs so we don't write into freed storage later. try: if self.model is not None: for module in self.model.modules(): if hasattr(module, "k_cache") and hasattr(module, "v_cache"): module.k_cache = module.v_cache = torch.tensor([]) release_sequence_state = getattr(module, "release_sequence_state", None) if callable(release_sequence_state): release_sequence_state() continue reset_sequence_state = getattr(module, "reset_sequence_state", None) if callable(reset_sequence_state): reset_sequence_state() else: if hasattr(module, "conv_state"): module.conv_state = None if hasattr(module, "recurrent_state"): module.recurrent_state = None except Exception: pass if hasattr(self, "kv_cache"): try: del self.kv_cache except Exception: pass # CUDA graphs captured against previous model/KV pointers are unsafe after runtime reset. # Force recapture on next prepare to avoid stale-pointer illegal memory access. try: self.clear_graph_cache() except Exception: pass try: for attr_name in ("graphs", "graph_vars", "graph_bs", "graph_pool"): if hasattr(self, attr_name): delattr(self, attr_name) except Exception: pass try: torch.cuda.synchronize() except Exception: pass try: torch.cuda.empty_cache() except Exception: pass try: torch.cuda.ipc_collect() except Exception: pass self._logits_bias_cache.clear() self._sampling_generator = None self._runtime_signature = None self._runtime_ready = False gc.collect() def _get_graph_capture_signature(self): model_ptr = -1 kv_ptr = -1 try: first_param = next(self.model.parameters()) if first_param.is_cuda: model_ptr = int(first_param.data_ptr()) except Exception: pass try: if hasattr(self, "kv_cache") and torch.is_tensor(self.kv_cache) and self.kv_cache.is_cuda: kv_ptr = int(self.kv_cache.data_ptr()) except Exception: pass return (model_ptr, kv_ptr, int(self.config.max_model_len), int(self.config.max_num_seqs)) def _get_model_device(self) -> torch.device: try: return next(self.model.parameters()).device except Exception: return torch.device("cpu") def _get_runtime_device(self) -> torch.device: if torch.cuda.is_available(): return torch.device("cuda", torch.cuda.current_device()) if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return torch.device("mps") return self._get_model_device() def _pin_memory_enabled(self) -> bool: return self._get_runtime_device().type == "cuda" def _to_runtime_device(self, tensor: torch.Tensor) -> torch.Tensor: device = self._get_runtime_device() return tensor.to(device=device, non_blocking=device.type == "cuda") def _drop_graph_cache_entry(self, cache_key): entry = self._graph_cache.pop(cache_key, None) if cache_key in self._graph_cache_order: self._graph_cache_order.remove(cache_key) if entry is None: return try: del entry["graphs"] del entry["pool"] del entry["vars"] del entry["bs"] except Exception: pass def clear_graph_cache(self): if self._graph_cache: for key in list(self._graph_cache.keys()): self._drop_graph_cache_entry(key) self._graph_cache.clear() self._graph_cache_order.clear() def _note_guard(self, name: str, detail: str | None = None): count = self._guard_counts.get(name, 0) + 1 self._guard_counts[name] = count if detail: detail_key = (name, detail) if detail_key not in self._guard_seen_details: print(f"[nanovllm][guard] {name}: {detail}") self._guard_seen_details.add(detail_key) return if count == 1: print(f"[nanovllm][guard] {name}") def _get_logits_bias(self, seq: Sequence, logits: torch.Tensor): bias = getattr(seq, "logits_bias", None) if bias is None or not torch.is_tensor(bias): return None key = (id(bias), logits.device, logits.dtype) cached = self._logits_bias_cache.get(key) if cached is not None: return cached cached = bias.to(device=logits.device, dtype=logits.dtype) self._logits_bias_cache[key] = cached return cached def set_sampling_seed(self, seed: int | None): if seed is None: self._sampling_generator = None return device = self._get_runtime_device() try: generator = torch.Generator(device=device) generator.manual_seed(int(seed)) self._sampling_generator = generator except Exception: self._sampling_generator = None @staticmethod def _apply_logits_bias(logits_row: torch.Tensor, bias: torch.Tensor): logits_row.add_(bias) def _allocate_sample_buffers(self): """Pre-allocate reusable buffers for sampling to avoid repeated tensor creation.""" max_bs = self.config.max_num_seqs max_tokens = self.config.max_num_batched_tokens max_num_blocks = (self.config.max_model_len + self.block_size - 1) // self.block_size pin_memory = self._pin_memory_enabled() # Pre-allocate pinned memory buffers on CPU for fast transfer # Must explicitly specify device="cpu" since default device may be "cuda" self._cpu_temperatures = torch.zeros(max_bs, dtype=torch.float32, device="cpu", pin_memory=pin_memory) self._cpu_cfg_scales = torch.zeros(max_bs, dtype=torch.float32, device="cpu", pin_memory=pin_memory) self._cpu_top_ks = torch.zeros(max_bs, dtype=torch.int32, device="cpu", pin_memory=pin_memory) self._cpu_top_ps = torch.zeros(max_bs, dtype=torch.float32, device="cpu", pin_memory=pin_memory) self._cpu_min_ps = torch.zeros(max_bs, dtype=torch.float32, device="cpu", pin_memory=pin_memory) self._cpu_repetition_penalties = torch.zeros(max_bs, dtype=torch.float32, device="cpu", pin_memory=pin_memory) # Pre-allocate decode buffers on CPU with pinned memory self._cpu_input_ids = torch.zeros(max_bs, dtype=torch.int64, device="cpu", pin_memory=pin_memory) self._cpu_positions = torch.zeros(max_bs, dtype=torch.int64, device="cpu", pin_memory=pin_memory) self._cpu_slot_mapping = torch.zeros(max_bs, dtype=torch.int32, device="cpu", pin_memory=pin_memory) self._cpu_context_lens = torch.zeros(max_bs, dtype=torch.int32, device="cpu", pin_memory=pin_memory) # Pre-allocate prefill buffers on CPU with pinned memory (optimization to avoid repeated tensor creation) self._cpu_prefill_input_ids = torch.zeros(max_tokens, dtype=torch.int64, device="cpu", pin_memory=pin_memory) self._cpu_prefill_positions = torch.zeros(max_tokens, dtype=torch.int64, device="cpu", pin_memory=pin_memory) self._cpu_prefill_cu_seqlens = torch.zeros(max_bs + 1, dtype=torch.int32, device="cpu", pin_memory=pin_memory) self._cpu_prefill_slot_mapping = torch.zeros(max_tokens, dtype=torch.int32, device="cpu", pin_memory=pin_memory) # Pre-allocate block tables buffer (shared by both decode and prefill) self._cpu_block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32, device="cpu", pin_memory=pin_memory) # Pre-allocate buffer for sequence token IDs (used in logits processor and sampler) # Max length is max_model_len since sequences can be that long self._seq_token_ids_buffer = torch.zeros(max_bs, self.config.max_model_len, dtype=torch.int64, device="cpu", pin_memory=pin_memory) def _release_sample_buffers(self): buffer_names = [ "_cpu_temperatures", "_cpu_cfg_scales", "_cpu_top_ks", "_cpu_top_ps", "_cpu_min_ps", "_cpu_repetition_penalties", "_cpu_input_ids", "_cpu_positions", "_cpu_slot_mapping", "_cpu_context_lens", "_cpu_prefill_input_ids", "_cpu_prefill_positions", "_cpu_prefill_cu_seqlens", "_cpu_prefill_slot_mapping", "_cpu_block_tables", "_seq_token_ids_buffer", ] for name in buffer_names: if hasattr(self, name): try: delattr(self, name) except Exception: pass def exit(self): try: self.reset_runtime_state() except Exception: pass self._release_sample_buffers() self._logits_bias_cache.clear() self._guard_counts.clear() self._guard_seen_details.clear() if hasattr(self, "sampler"): self.sampler = None if hasattr(self, "model"): self.model = None if self.world_size > 1: self.shm.close() dist.barrier() if self.rank == 0: self.shm.unlink() if not self.enforce_eager: if hasattr(self, "graphs"): del self.graphs if hasattr(self, "graph_vars"): del self.graph_vars if hasattr(self, "graph_bs"): del self.graph_bs if hasattr(self, "graph_pool"): del self.graph_pool try: torch.cuda.synchronize() except Exception: pass try: torch.cuda.empty_cache() except Exception: pass try: torch.cuda.ipc_collect() except Exception: pass if dist.is_initialized(): dist.destroy_process_group() def loop(self): while True: method_name, args = self.read_shm() self.call(method_name, *args) if method_name == "exit": break def read_shm(self): assert self.world_size > 1 and self.rank > 0 self.event.wait() n = int.from_bytes(self.shm.buf[0:4], "little") method_name, *args = pickle.loads(self.shm.buf[4:n+4]) self.event.clear() return method_name, args def write_shm(self, method_name, *args): assert self.world_size > 1 and self.rank == 0 data = pickle.dumps([method_name, *args]) n = len(data) self.shm.buf[0:4] = n.to_bytes(4, "little") self.shm.buf[4:n+4] = data for event in self.event: event.set() def call(self, method_name, *args): if self.world_size > 1 and self.rank == 0: self.write_shm(method_name, *args) method = getattr(self, method_name, None) return method(*args) def _get_kv_cache_modules(self): if self.model is None: return [] return [module for module in self.model.modules() if hasattr(module, "k_cache") and hasattr(module, "v_cache")] def allocate_kv_cache(self): config = self.config hf_config = config.hf_config runtime_device = self._get_runtime_device() is_cuda_runtime = runtime_device.type == "cuda" num_kv_heads = hf_config.num_key_value_heads // self.world_size head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads) kv_cache_modules = self._get_kv_cache_modules() kv_cache_layer_count = len(kv_cache_modules) block_bytes = 2 * kv_cache_layer_count * self.block_size * num_kv_heads * head_dim * self.dtype.itemsize # Strict policy: allocate exactly the blocks required by requested runtime limits. required_blocks_per_seq = (config.max_model_len + self.block_size - 1) // self.block_size required_total_blocks = required_blocks_per_seq * max(1, config.max_num_seqs) config.num_kvcache_blocks = max(1, int(required_total_blocks)) required_kv_bytes = config.num_kvcache_blocks * block_bytes if is_cuda_runtime: free, total = torch.cuda.mem_get_info() current = torch.cuda.memory_stats()["allocated_bytes.all.current"] target_total_usage = total * config.gpu_memory_utilization allowed_kv_bytes = max(0, target_total_usage - current) if required_kv_bytes > allowed_kv_bytes: raise RuntimeError( f"Insufficient GPU memory for strict KV cache sizing under gpu_memory_utilization={config.gpu_memory_utilization:.2f}. " f"Required KV: {required_kv_bytes / 1024**3:.2f} GB " f"(blocks={config.num_kvcache_blocks}, block={block_bytes / 1024**2:.2f} MB), " f"Allowed by limit: {allowed_kv_bytes / 1024**3:.2f} GB, " f"Free: {free / 1024**3:.2f} GB, Current: {current / 1024**3:.2f} GB, " f"Requested max_model_len={config.max_model_len}, max_num_seqs={config.max_num_seqs}." ) try: self.kv_cache = torch.empty( 2, kv_cache_layer_count, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim, device=runtime_device, dtype=self.dtype, ) except RuntimeError as exc: if "out of memory" in str(exc).lower(): extra = "" if is_cuda_runtime: free_now, total_now = torch.cuda.mem_get_info() extra = f" Current free VRAM: {free_now / 1024**3:.2f} GB / {total_now / 1024**3:.2f} GB total." raise RuntimeError( f"Failed to allocate strict KV cache ({required_kv_bytes / 1024**3:.2f} GB) for " f"max_model_len={config.max_model_len}, max_num_seqs={config.max_num_seqs} on {runtime_device}.{extra}" ) from exc raise for layer_id, module in enumerate(kv_cache_modules): module.k_cache = self.kv_cache[0, layer_id] module.v_cache = self.kv_cache[1, layer_id] def prepare_block_tables(self, seqs: list[Sequence]): bs = len(seqs) max_len = max(len(seq.block_table) for seq in seqs) block_tables = self._cpu_block_tables[:bs, :max_len] block_tables.fill_(-1) for row, seq in enumerate(seqs): if not seq.block_table: continue block_tables[row, :len(seq.block_table)] = torch.tensor(seq.block_table, dtype=torch.int32, device="cpu") return self._to_runtime_device(block_tables) def prepare_prefill(self, seqs: list[Sequence]): use_prompt_embeds = any(getattr(seq, "prompt_embeds", None) is not None for seq in seqs) if use_prompt_embeds and not all(getattr(seq, "prompt_embeds", None) is not None for seq in seqs): raise RuntimeError("Mixed embedded/non-embedded prefill batches are not supported.") input_ids = [] positions = [] prompt_embeds = [] prompt_position_ids = [] cu_seqlens_q = [0] cu_seqlens_k = [0] max_seqlen_q = 0 max_seqlen_k = 0 slot_mapping = [] block_tables = None has_previous_state = any(int(getattr(seq, "num_cached_tokens", 0) or 0) > 0 for seq in seqs) for seq in seqs: seqlen = len(seq) input_ids.extend(seq[seq.num_cached_tokens:]) seqlen_q = seqlen - seq.num_cached_tokens seqlen_k = seqlen if use_prompt_embeds: seq_prompt_embeds = getattr(seq, "prompt_embeds", None) seq_position_ids = getattr(seq, "prompt_position_ids", None) if seq_prompt_embeds is None or seq_position_ids is None: raise RuntimeError("Embedded prefill requires both prompt_embeds and prompt_position_ids.") seq_prompt_embeds = seq_prompt_embeds[seq.num_cached_tokens:seqlen] if seq_prompt_embeds.ndim != 2 or seq_prompt_embeds.shape[0] != seqlen_q: raise RuntimeError("Embedded prefill prompt_embeds shape does not match uncached prompt length.") if seq_position_ids.ndim == 3: seq_position_ids = seq_position_ids[:, 0] if seq_position_ids.ndim != 2 or seq_position_ids.shape[0] != 3: raise RuntimeError("Embedded prefill prompt_position_ids must have shape [3, seq_len].") seq_position_ids = seq_position_ids[:, seq.num_cached_tokens:seqlen] if seq_position_ids.shape[1] != seqlen_q: raise RuntimeError("Embedded prefill prompt_position_ids shape does not match uncached prompt length.") prompt_embeds.append(seq_prompt_embeds) prompt_position_ids.append(seq_position_ids) else: positions.extend(list(range(seq.num_cached_tokens, seqlen))) cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) max_seqlen_q = max(seqlen_q, max_seqlen_q) max_seqlen_k = max(seqlen_k, max_seqlen_k) if not seq.block_table: # warmup: no blocks allocated yet slot_mapping.extend([-1] * seqlen_q) continue cached_tokens = max(0, int(seq.num_cached_tokens or 0)) cached_partial_tokens = cached_tokens % self.block_size for i in range(seq.num_cached_blocks, seq.num_blocks): start = seq.block_table[i] * self.block_size if i == seq.num_cached_blocks and cached_partial_tokens > 0: start += cached_partial_tokens if i != seq.num_blocks - 1: end = seq.block_table[i] * self.block_size + self.block_size else: end = seq.block_table[i] * self.block_size + seq.last_block_num_tokens if end > start: slot_mapping.extend(list(range(start, end))) if cu_seqlens_k[-1] > cu_seqlens_q[-1]: # prefix cache block_tables = self.prepare_block_tables(seqs) if use_prompt_embeds: input_ids = None positions = self._to_runtime_device(torch.cat(prompt_position_ids, dim=1).unsqueeze(1).contiguous()) inputs_embeds = self._to_runtime_device(torch.cat(prompt_embeds, dim=0).unsqueeze(0).contiguous()) else: pin_memory = self._pin_memory_enabled() input_ids = self._to_runtime_device(torch.tensor(input_ids, dtype=torch.int64, pin_memory=pin_memory)) positions = self._to_runtime_device(torch.tensor(positions, dtype=torch.int64, pin_memory=pin_memory)) inputs_embeds = None pin_memory = self._pin_memory_enabled() cu_seqlens_q = self._to_runtime_device(torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=pin_memory)) cu_seqlens_k = self._to_runtime_device(torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=pin_memory)) slot_mapping = self._to_runtime_device(torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=pin_memory)) set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables, has_previous_state=has_previous_state) return input_ids, positions, inputs_embeds @torch.inference_mode() def prefill_only(self, seqs: list[Sequence]) -> None: self.ensure_runtime_ready() input_ids, positions, inputs_embeds = self.prepare_prefill(seqs) self.run_model(input_ids, positions, True, inputs_embeds=inputs_embeds) for seq in seqs: seq.clear_prompt_data() reset_context() def prepare_decode(self, seqs: list[Sequence]): """Optimized decode preparation using pre-allocated buffers.""" bs = len(seqs) # Use pre-allocated CPU buffers for i, seq in enumerate(seqs): self._cpu_input_ids[i] = seq.last_token self._cpu_positions[i] = len(seq) - 1 + int(getattr(seq, "position_offset", 0) or 0) self._cpu_context_lens[i] = len(seq) self._cpu_slot_mapping[i] = seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1 # Transfer to the runtime device using sliced views input_ids = self._to_runtime_device(self._cpu_input_ids[:bs]) positions = self._to_runtime_device(self._cpu_positions[:bs]) slot_mapping = self._to_runtime_device(self._cpu_slot_mapping[:bs]) context_lens = self._to_runtime_device(self._cpu_context_lens[:bs]) block_tables = self.prepare_block_tables(seqs) set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables) return input_ids, positions def prepare_sample(self, seqs: list[Sequence], is_cfg_batch: bool = False): """Optimized sample preparation using pre-allocated buffers.""" if is_cfg_batch: num_seqs = len(seqs) // 2 target_seqs = seqs[:num_seqs] else: num_seqs = len(seqs) target_seqs = seqs # Fill pre-allocated CPU buffers top_ks_is_zero = True top_ps_is_one = True min_ps_is_zero = True repetition_penalties_is_one = True for i, seq in enumerate(target_seqs): self._cpu_temperatures[i] = seq.temperature self._cpu_cfg_scales[i] = seq.cfg_scale self._cpu_top_ks[i] = seq.top_k if seq.top_k is not None else 0 if seq.top_k is not None and seq.top_k > 0: top_ks_is_zero = False self._cpu_top_ps[i] = seq.top_p if seq.top_p is not None else 1.0 if seq.top_p is not None and seq.top_p != 1.0: top_ps_is_one = False self._cpu_min_ps[i] = seq.min_p if seq.min_p is not None else 0.0 if seq.min_p is not None and seq.min_p > 0.0: min_ps_is_zero = False self._cpu_repetition_penalties[i] = seq.repetition_penalty if seq.repetition_penalty is not None else 1.0 if seq.repetition_penalty is not None and seq.repetition_penalty != 1.0: repetition_penalties_is_one = False # Transfer to the runtime device using sliced views (single batched transfer) temperatures = self._to_runtime_device(self._cpu_temperatures[:num_seqs]) cfg_scales = self._to_runtime_device(self._cpu_cfg_scales[:num_seqs]) top_ks = self._to_runtime_device(self._cpu_top_ks[:num_seqs]) if not top_ks_is_zero else None top_ps = self._to_runtime_device(self._cpu_top_ps[:num_seqs]) if not top_ps_is_one else None min_ps = self._to_runtime_device(self._cpu_min_ps[:num_seqs]) if not min_ps_is_zero else None repetition_penalties = self._to_runtime_device(self._cpu_repetition_penalties[:num_seqs]) if not repetition_penalties_is_one else None return temperatures, cfg_scales, top_ks, top_ps, min_ps, repetition_penalties @torch.inference_mode() def run_model(self, input_ids: torch.Tensor | None, positions: torch.Tensor, is_prefill: bool, inputs_embeds: torch.Tensor | None = None): decode_batch_size = input_ids.size(0) if input_ids is not None else int(inputs_embeds.shape[0]) model_kwargs = {"input_ids": input_ids, "positions": positions} if inputs_embeds is not None: model_kwargs["inputs_embeds"] = inputs_embeds if is_prefill or self.enforce_eager or decode_batch_size > 512: return self.model.compute_logits(self.model(**model_kwargs)) else: bs = decode_batch_size context = get_context() # Check if block_tables size exceeds pre-allocated buffer size # This can happen when conditional and unconditional sequences have different lengths # in CFG mode, causing block_tables to have more columns than expected max_num_blocks = self.graph_vars["block_tables"].size(1) if context.block_tables.size(1) > max_num_blocks: # Fall back to eager mode when block_tables is too large for CUDA graph self._note_guard( "cudagraph_fallback_block_table_cols", f"requested={context.block_tables.size(1)} max={max_num_blocks}", ) return self.model.compute_logits(self.model(**model_kwargs)) # Fix: Also check if block_tables row count matches batch size # Dimension mismatch can cause CUDA illegal memory access during graph replay if context.block_tables.size(0) != bs: # Fall back to eager mode when block_tables row count doesn't match batch size self._note_guard( "cudagraph_fallback_block_table_rows", f"rows={context.block_tables.size(0)} bs={bs}", ) return self.model.compute_logits(self.model(**model_kwargs)) # Fix: Verify slot_mapping and context_lens dimensions match batch size if context.slot_mapping.size(0) != bs or context.context_lens.size(0) != bs: # Fall back to eager mode when dimensions don't match self._note_guard( "cudagraph_fallback_context_shape", f"slot={context.slot_mapping.size(0)} ctx={context.context_lens.size(0)} bs={bs}", ) return self.model.compute_logits(self.model(**model_kwargs)) graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] graph_vars = self.graph_vars graph_vars["input_ids"][:bs] = input_ids graph_vars["positions"][:bs] = positions graph_vars["slot_mapping"].fill_(-1) graph_vars["slot_mapping"][:bs] = context.slot_mapping graph_vars["context_lens"].zero_() graph_vars["context_lens"][:bs] = context.context_lens # Clear block_tables first to ensure no stale data from previous runs graph_vars["block_tables"][:bs].fill_(-1) graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables graph.replay() return self.model.compute_logits(graph_vars["outputs"][:bs]) def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]: """Run model forward and sampling. For CFG sequences, batch is structured as: [cond_seq1, cond_seq2, ..., uncond_seq1, uncond_seq2, ...] where uncond_seqi is the paired unconditional sequence of cond_seqi.""" self.ensure_runtime_ready() # Check if this is a CFG batch (contains paired conditional and unconditional sequences) is_cfg_batch = seqs[0].cfg_scale > 1.0 and seqs[0].paired_seq is not None if is_cfg_batch: # CFG batch: seqs = [cond_seq1, cond_seq2, ..., uncond_seq1, uncond_seq2, ...] num_cond = len(seqs) // 2 cond_seqs = seqs[:num_cond] # uncond_seqs = seqs[num_cond:] # Prepare inputs for both conditional and unconditional (they're already in the batch) if is_prefill: input_ids, positions, inputs_embeds = self.prepare_prefill(seqs) else: input_ids, positions = self.prepare_decode(seqs) inputs_embeds = None sample_params = self.prepare_sample(seqs, is_cfg_batch=True) if self.rank == 0 else None if sample_params is not None: temperatures, cfg_scales, top_ks, top_ps, min_ps, repetition_penalties = sample_params else: temperatures = cfg_scales = top_ks = top_ps = min_ps = repetition_penalties = None # Run model forward (processes entire batch: cond + uncond) logits_all = self.run_model(input_ids, positions, is_prefill, inputs_embeds=inputs_embeds) if is_prefill: for seq in seqs: seq.clear_prompt_data() reset_context() if self.rank == 0: # Split logits: first half is conditional, second half is unconditional logits_cond = logits_all[:num_cond] logits_uncond = logits_all[num_cond:] # Apply repetition penalty to conditional logits (before CFG) if repetition_penalties is not None: for i, seq in enumerate(cond_seqs): penalty = repetition_penalties[i].item() if penalty != 1.0: # Only penalize completion tokens (not prompt tokens) completion_tokens = torch.tensor(seq.completion_token_ids, device=logits_cond.device) if len(completion_tokens) > 0: # Create token mask: mark tokens that appeared in completion token_mask = torch.zeros(logits_cond.shape[1], dtype=torch.bool, device=logits_cond.device) token_mask[completion_tokens] = True # Apply standard repetition penalty formula (matching transformers implementation): # For tokens in completion: if score < 0 then score * penalty, else score / penalty penalty_scores = torch.where( logits_cond[i] < 0, logits_cond[i] * penalty, logits_cond[i] / penalty ) # Only apply penalty to tokens that appeared in completion logits_cond[i] = torch.where(token_mask, penalty_scores, logits_cond[i]) # Apply CFG formula: logits_cfg = logits_uncond + cfg_scale * (logits_cond - logits_uncond) cfg_scales_tensor = cfg_scales.unsqueeze(1) # [num_cond, 1] logits_cfg = logits_uncond + cfg_scales_tensor * (logits_cond - logits_uncond) # Apply optional per-sequence logits bias before processors/sampling. for i, seq in enumerate(cond_seqs): bias = self._get_logits_bias(seq, logits_cfg) if bias is not None: self._apply_logits_bias(logits_cfg[i], bias) # Apply logits processor for constrained decoding (if any sequence has one) for i, seq in enumerate(cond_seqs): if seq.logits_processor is not None: # Create input_ids tensor for this sequence seq_input_ids = torch.tensor([seq.token_ids], device=logits_cfg.device) # Apply processor to this sequence's logits logits_cfg[i:i+1] = seq.logits_processor(seq_input_ids, logits_cfg[i:i+1]) # Prepare input_ids for sampler (for repetition penalty, though we already applied it) # cond_input_ids = torch.tensor([seq.token_ids for seq in cond_seqs], device=logits_cfg.device) # Sample from CFG logits token_ids_cfg = self.sampler( logits_cfg, temperatures, top_ks=top_ks if top_ks is not None else None, top_ps=top_ps if top_ps is not None else None, min_ps=min_ps if min_ps is not None else None, repetition_penalties=None, # Already applied above generator=self._sampling_generator, # input_ids=cond_input_ids, ).tolist() # Update logits processor state after sampling # NOTE: Only update for the first sequence since all sequences share the same processor # Updating multiple times would cause duplicate state updates (e.g., codes_count += N instead of += 1) if cond_seqs and cond_seqs[0].logits_processor_update_state is not None: cond_seqs[0].logits_processor_update_state(token_ids_cfg[0]) # Return token_ids (will be applied to both conditional and unconditional sequences) return token_ids_cfg else: return None else: # Normal batch (non-CFG) if is_prefill: input_ids, positions, inputs_embeds = self.prepare_prefill(seqs) else: input_ids, positions = self.prepare_decode(seqs) inputs_embeds = None sample_params = self.prepare_sample(seqs, is_cfg_batch=False) if self.rank == 0 else None if sample_params is not None: temperatures, cfg_scales, top_ks, top_ps, min_ps, repetition_penalties = sample_params else: temperatures = cfg_scales = top_ks = top_ps = min_ps = repetition_penalties = None logits = self.run_model(input_ids, positions, is_prefill, inputs_embeds=inputs_embeds) if is_prefill: for seq in seqs: seq.clear_prompt_data() reset_context() if self.rank == 0: # Apply repetition penalty to logits if repetition_penalties is not None: for i, seq in enumerate(seqs): penalty = repetition_penalties[i].item() if penalty != 1.0: # Only penalize completion tokens (not prompt tokens) completion_tokens = torch.tensor(seq.completion_token_ids, device=logits.device) if len(completion_tokens) > 0: # Create token mask: mark tokens that appeared in completion token_mask = torch.zeros(logits.shape[1], dtype=torch.bool, device=logits.device) token_mask[completion_tokens] = True # Apply standard repetition penalty formula (matching transformers implementation): # For tokens in completion: if score < 0 then score * penalty, else score / penalty penalty_scores = torch.where( logits[i] < 0, logits[i] * penalty, logits[i] / penalty ) # Only apply penalty to tokens that appeared in completion logits[i] = torch.where(token_mask, penalty_scores, logits[i]) # Apply logits processor for constrained decoding (if any sequence has one) # Clone logits to avoid in-place update issues in inference mode logits = logits.clone() for i, seq in enumerate(seqs): bias = self._get_logits_bias(seq, logits) if bias is not None: self._apply_logits_bias(logits[i], bias) for i, seq in enumerate(seqs): if seq.logits_processor is not None: # Create input_ids tensor for this sequence seq_input_ids = torch.tensor([seq.token_ids], device=logits.device) # Apply processor to this sequence's logits (clone to avoid inference mode issues) processed = seq.logits_processor(seq_input_ids, logits[i:i+1].clone()) logits[i] = processed[0] # Prepare input_ids for sampler # seq_input_ids = torch.tensor([seq.token_ids for seq in seqs], device=logits.device) token_ids = self.sampler( logits, temperatures, top_ks=top_ks if top_ks is not None else None, top_ps=top_ps if top_ps is not None else None, min_ps=min_ps if min_ps is not None else None, repetition_penalties=None, # Already applied above generator=self._sampling_generator, # input_ids=seq_input_ids, ).tolist() # Update logits processor state after sampling # NOTE: Only update for the first sequence since all sequences may share the same processor # (when using a single SamplingParams for batch generation) # Updating multiple times would cause duplicate state updates (e.g., codes_count += N instead of += 1) if seqs and seqs[0].logits_processor_update_state is not None: seqs[0].logits_processor_update_state(token_ids[0]) return token_ids else: return None @torch.inference_mode() def capture_cudagraph(self): if self._get_runtime_device().type != "cuda": self.enforce_eager = True return config = self.config cache_key = (config.max_model_len, config.max_num_seqs) model_device = torch.device("cuda") if torch.cuda.is_available() else self._get_model_device() cached = self._graph_cache.get(cache_key) if cached is not None: current_sig = self._get_graph_capture_signature() if cached.get("sig") == current_sig: self.graphs = cached["graphs"] self.graph_pool = cached["pool"] self.graph_vars = cached["vars"] self.graph_bs = cached["bs"] if cache_key in self._graph_cache_order: self._graph_cache_order.remove(cache_key) self._graph_cache_order.append(cache_key) return self._drop_graph_cache_entry(cache_key) hf_config = config.hf_config max_bs = min(self.config.max_num_seqs, 512) max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size input_ids = torch.zeros(max_bs, dtype=torch.int64, device=model_device) positions = torch.zeros(max_bs, dtype=torch.int64, device=model_device) slot_mapping = torch.zeros(max_bs, dtype=torch.int32, device=model_device) context_lens = torch.zeros(max_bs, dtype=torch.int32, device=model_device) block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32, device=model_device) outputs = torch.zeros(max_bs, hf_config.hidden_size, device=model_device, dtype=self.dtype) base_graph_bs = [1, 2, 4, 8] self.graph_bs = [bs for bs in base_graph_bs if bs <= max_bs] if max_bs > 8: self.graph_bs.extend(range(16, max_bs + 1, 16)) if not self.graph_bs: self.graph_bs = [max_bs] self.graphs = {} self.graph_pool = self._graph_pool_seed for bs in reversed(self.graph_bs): graph = torch.cuda.CUDAGraph() set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs]) outputs[:bs] = self.model(input_ids[:bs], positions[:bs]) # warmup with torch.cuda.graph(graph, self.graph_pool): outputs[:bs] = self.model(input_ids[:bs], positions[:bs]) # capture if self.graph_pool is None: self.graph_pool = graph.pool() self.graphs[bs] = graph torch.cuda.synchronize() reset_context() self.graph_vars = dict( input_ids=input_ids, positions=positions, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables, outputs=outputs, ) self._graph_cache[cache_key] = { "graphs": self.graphs, "pool": self.graph_pool, "vars": self.graph_vars, "bs": self.graph_bs, "sig": self._get_graph_capture_signature(), } if cache_key in self._graph_cache_order: self._graph_cache_order.remove(cache_key) self._graph_cache_order.append(cache_key) while len(self._graph_cache_order) > 5: old_key = self._graph_cache_order.pop(0) self._drop_graph_cache_entry(old_key)