diff --git a/batchgen/batchgen_client.py b/batchgen/batchgen_client.py index 6dea54107..b91adfe3e 100644 --- a/batchgen/batchgen_client.py +++ b/batchgen/batchgen_client.py @@ -6,6 +6,8 @@ import argparse from typing import List, Optional, Dict, Any +from batchgen.deprecation import LegacyInferenceDeprecated + try: import requests _REQUESTS_AVAILABLE = True @@ -208,47 +210,18 @@ def submit_inference( temperature: Optional[float] = None, top_p: Optional[float] = None, ) -> List[str]: - """Submit inference request and get decoded string results. - - Args: - prompts: List of prompt strings - max_input_len: Maximum input sequence length. If None, determined - dynamically from the longest prompt in the batch. - max_output_len: Maximum output/decoding length - ignore_eos: If True, ignore EOS tokens and decode to max_output_len - temperature: Sampling temperature (None = greedy decoding) - top_p: Nucleus sampling threshold (None = disabled) + """DEPRECATED and disabled. Use submit_batch() instead. - Returns: - List of decoded output strings + The method is kept, rather than deleted, so a caller gets the + explanation above instead of an AttributeError telling it only that + something is gone. It raises without any network call: the server + answers /v1/inference with 410 anyway, and failing here keeps the + deprecated request off the wire entirely. Raises: - RuntimeError: If inference fails or returns unexpected format + LegacyInferenceDeprecated: always. """ - payload: Dict[str, Any] = { - "prompts": prompts, - "max_input_len": max_input_len, - "max_output_len": max_output_len, - "ignore_eos": ignore_eos, - } - if temperature is not None: - payload["temperature"] = temperature - if top_p is not None: - payload["top_p"] = top_p - - response = self.post_json("/v1/inference", payload) - - if response.get("status") != "success": - raise RuntimeError(f"Inference failed: {response}") - - results = response.get("results") - if not results: - raise RuntimeError("Server returned empty results.") - - if not isinstance(results, list): - raise RuntimeError(f"Unexpected result format: {type(results)}") - - return results + raise LegacyInferenceDeprecated() # ==================== Batch API Methods ==================== diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 92cf18bb5..cef5bd2f2 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -1,6 +1,7 @@ import concurrent.futures import copy import functools +import json import psutil import logging import math @@ -16,6 +17,7 @@ from tqdm import tqdm from batchgen.config.model_registry import load_config from batchgen.config.tokenizer_registry import load_tokenizer +from batchgen.deprecation import LegacyInferenceDeprecated # Use new wrapper system - Attn_Wrapper/Expert_Wrapper are aliases for backward compatibility from batchgen.models.wrappers import BaseModuleWrapper, AttnWrapperBase, ExpertWrapperBase @@ -274,24 +276,125 @@ class _DualKVLoadPointers: aux_page_counts: torch.Tensor +class QueryBookPoolCapacityError(RuntimeError): + """A QueryBook pool request exceeded the rows/width actually allocated.""" + + +def allocate_node_shared_int64( + name: str, + rows: int, + width: int, + is_creator: bool, + barrier, +) -> Tuple[torch.Tensor, object]: + """Map ONE int64 ``[rows, width]`` CPU tensor per node into every worker. + + The tokenized global batch is identical on every rank (``_tokenize_global_batch`` + all-gathers the results to all of them), so each worker used to hold its own + private copy of the same input-ids table — ``world_size`` duplicates of the + same bytes, which is what OOM-killed the node. + + ``is_creator`` must be true on exactly one rank per node. The creator makes + the segment (POSIX guarantees it is zero-filled, matching the ``torch.zeros`` + it replaces), everyone waits on ``barrier``, then the rest attach. ``barrier`` + is ``dist.barrier`` in the worker and an ``mp.Barrier`` in tests. + + Returns ``(tensor, shm)``. The caller MUST keep ``shm`` alive for as long as + the tensor is reachable: the tensor points straight into the mapping. + + Two operational notes: the segment lands in /dev/shm, so the container's + shm budget has to cover it; and CPython < 3.13 registers a segment with the + resource_tracker on attach as well as on create, so every non-creator rank + prints one "leaked shared_memory objects" warning at shutdown. That warning + is cosmetic — unlink only drops the name, never a live mapping. + """ + from multiprocessing import shared_memory + + nbytes = rows * width * 8 + if is_creator: + try: + # A crashed predecessor can leave the name behind; reusing its + # (possibly smaller) segment would silently truncate. + shared_memory.SharedMemory(name=name).unlink() + except FileNotFoundError: + pass + shm = shared_memory.SharedMemory(name=name, create=True, size=nbytes) + barrier() + if not is_creator: + shm = shared_memory.SharedMemory(name=name) + if shm.size < nbytes: + raise QueryBookPoolCapacityError( + f"shared input_ids segment '{name}' is {shm.size} bytes, " + f"need {nbytes} ({rows} rows x {width} tokens x 8B)" + ) + buf = torch.frombuffer(shm.buf, dtype=torch.int64, count=rows * width).view(rows, width) + # Nobody may unlink/close until every rank has mapped it. + barrier() + return buf, shm + + class QueryBookBufferPool: """Pre-allocated contiguous buffers for query book tensors. Eliminates per-sequence tensor allocation in Phase 3 of _tokenize_global_batch(). With 16 ranks each creating 12K tensors, allocator contention causes ~19 min init. This replaces 24K allocations per rank with 2 large allocations + views. + + ``input_ids_buffer`` may be passed in as a node-shared tensor (see + ``allocate_node_shared_int64``) — its contents are identical on every rank, + so one copy per node is enough. ``decoded_tokens_buffer`` stays PRIVATE: only + the owning rank writes a sequence's decoded tokens, so the ranks' copies + legitimately differ. + + ``input_ids_width`` is the widest ``seq_extended_size`` the pool can serve. + It is sized from the batch that is actually being admitted, NOT from the + model context length: at K3's 1,048,576-token context a 10240-slot pool + would be 80 GiB of zeros per worker. """ - def __init__(self, num_sequences: int, model_context_length: int, max_decoding_length: int, pad_token_id: int = 0): - self.input_ids_buffer = torch.zeros((num_sequences, model_context_length), dtype=torch.long) + def __init__( + self, + num_sequences: int, + input_ids_width: int, + max_decoding_length: int, + pad_token_id: int = 0, + input_ids_buffer: Optional[torch.Tensor] = None, + input_ids_shm: object = None, + ): + if input_ids_buffer is None: + input_ids_buffer = torch.zeros((num_sequences, input_ids_width), dtype=torch.long) + elif tuple(input_ids_buffer.shape) != (num_sequences, input_ids_width): + raise QueryBookPoolCapacityError( + f"shared input_ids buffer has shape {tuple(input_ids_buffer.shape)}, " + f"pool needs ({num_sequences}, {input_ids_width})" + ) + self.input_ids_buffer = input_ids_buffer + self.input_ids_shm = input_ids_shm self.decoded_tokens_buffer = torch.full((num_sequences, max_decoding_length), pad_token_id, dtype=torch.int64) self.pad_token_id = pad_token_id self.num_sequences = num_sequences - self.model_context_length = model_context_length + self.input_ids_width = input_ids_width self.max_decoding_length = max_decoding_length self._free_slots: set = set() self._next_slot: int = 0 + def adopt(self, old: "QueryBookBufferPool") -> None: + """Carry contents and slot bookkeeping over from a superseded pool.""" + rows = min(self.num_sequences, old.num_sequences) + cols = min(self.input_ids_width, old.input_ids_width) + self.input_ids_buffer[:rows, :cols] = old.input_ids_buffer[:rows, :cols] + dec = min(self.max_decoding_length, old.max_decoding_length) + self.decoded_tokens_buffer[:rows, :dec] = old.decoded_tokens_buffer[:rows, :dec] + self._free_slots = set(old._free_slots) + self._next_slot = old._next_slot + + def reset(self) -> None: + """Return the pool to its just-allocated state (legacy per-batch reuse).""" + self._free_slots = set() + self._next_slot = 0 + self.input_ids_buffer.zero_() + self.decoded_tokens_buffer.fill_(self.pad_token_id) + def allocate_slot(self) -> int: if self._free_slots: slot = self._free_slots.pop() @@ -301,7 +404,10 @@ def allocate_slot(self) -> int: return slot slot = self._next_slot if slot >= self.num_sequences: - raise RuntimeError(f"QueryBookBufferPool exhausted: {self.num_sequences} slots used") + raise QueryBookPoolCapacityError( + f"QueryBookBufferPool exhausted: {self.num_sequences} slots used " + f"(raise --max-pool-size)" + ) self._next_slot += 1 return slot @@ -309,6 +415,13 @@ def free_slot(self, slot: int): self._free_slots.add(slot) def get_input_ids_view(self, slot: int, seq_extended_size: int) -> torch.Tensor: + if seq_extended_size > self.input_ids_width: + # Slicing would silently hand back a SHORT view and truncate the + # prompt. The pool must be grown instead (_ensure_buffer_pool). + raise QueryBookPoolCapacityError( + f"input_ids view of {seq_extended_size} tokens requested from a pool " + f"allocated {self.input_ids_width} tokens wide (slot={slot})" + ) return self.input_ids_buffer[slot:slot+1, :seq_extended_size] def get_decoded_tokens_view(self, slot: int) -> torch.Tensor: @@ -758,9 +871,24 @@ def __init__(self, args: BatchGenWorkerArgs): # Request pool: admission queue and response queue for persistent loop self._admission_queue = None # mp.Queue, set via set_admission_queue() self._response_queue = None # mp.Queue, set via set_response_queue() + # global_idx -> decoded text for sequences completed during PREFILL + # (C4). Captured before _report_completion pops the local maps; the + # legacy end-of-generate() gather merges it. Legacy mode only. + self._prefill_completed_results: Dict[int, str] = {} self._shutdown_requested = False self._max_pool_size = args.max_pool_size # 0 = legacy mode + # QueryBook buffer pool. Allocated lazily by _ensure_buffer_pool() once + # the first batch's tokenized lengths are known — its input_ids buffer + # is ONE shared-memory segment per node, so it cannot be sized from + # static config. + self._buffer_pool: Optional[QueryBookBufferPool] = None + self._buffer_pool_generation = 0 + self._shared_buffer_tag: Optional[str] = None + # Superseded pools stay mapped for the process lifetime (see + # _retire_buffer_pool). + self._retired_buffer_pools: List[QueryBookBufferPool] = [] + logging.info(f"Rank {self.rank}: BatchGenWorker __init__ completed.") def Init(self, max_input_length, max_decoding_length, num_queries, max_context_length=None): @@ -1167,6 +1295,17 @@ def _poll_admissions(self) -> bool: elif isinstance(msg, dict) and msg.get("type") == "admit": msg_data = msg has_new = True + elif isinstance(msg, dict) and "prompts" in msg: + # A legacy /v1/inference payload (worker_manager.infer builds + # exactly this shape). It matches no branch above, so it used + # to be dropped right here while the caller sat on + # response_queue.get() and took the next batch's completion. + # The HTTP route now returns 410, so reaching this line means + # some other producer is putting legacy payloads on the queue: + # fail loudly rather than park. Deliberately NOT a catch-all + # for unknown messages -- {"command": "reload"} also lands + # here and must keep its current handling. + raise LegacyInferenceDeprecated() except queue_mod.Empty: pass @@ -1261,7 +1400,9 @@ def _tokenize_admitted_sequences(self, uuids: List[str]) -> None: Reuses the same parallel tokenization + buffer pool fill pattern as _tokenize_global_batch Phase 1 + Phase 3. Key differences: - - Uses existing buffer pool (not creating a new one) + - Allocates the buffer pool on the first admission and grows it when a + later admission is wider (the pool cannot be pre-sized: its widths + come from the requests, not from static config) - Only processes the new sequences, not the full global_batch Optimization: uses padding=False to avoid creating a large padded 2D @@ -1344,6 +1485,35 @@ def _tokenize_admitted_sequences(self, uuids: List[str]) -> None: }) self.global_batch.remove_sequence(uuid) + # Phase 2.75: size the pool for what this admission actually needs. + # COLLECTIVE — every rank runs it with the same numbers: the admission + # message was broadcast and the tokenized lengths were all-gathered above. + required_input_width = 0 + required_decode_width = 0 + for i, seq in enumerate(sequences): + if seq.uuid in rejected_uuids: + continue + prompt_len = tokenized_by_idx[i]["length"] + required_input_width = max( + required_input_width, + min(prompt_len + seq.max_decode_length, self.model_context_length), + ) + required_decode_width = max( + required_decode_width, + min(seq.max_decode_length, self.model_context_length), + ) + if required_input_width > 0: + # Rows keep their --max-pool-size meaning: the pool is NOT widened to + # fit an over-subscribed batch, allocate_slot() still hard-fails. + self._ensure_buffer_pool( + required_rows=( + self._max_pool_size if self._max_pool_size > 0 else len(sequences) + ), + required_input_width=required_input_width, + required_decode_width=required_decode_width, + reason=f"admission of {len(sequences)} sequences", + ) + # Phase 3: Assign buffer pool slots and fill token data # Same pattern as _tokenize_global_batch Phase 3 — allocate slot from # existing buffer pool, write tokens directly into the view. @@ -1484,6 +1654,9 @@ def _report_completion(self, uuid: str, gathered_text: str = None) -> None: if hasattr(self, '_buffer_pool') and self._buffer_pool is not None: if seq._buffer_slot >= 0: self._buffer_pool.free_slot(seq._buffer_slot) + # Guard against double-free / stale reuse: a re-entered report + # for this seq must not free a slot now owned by another seq. + seq._buffer_slot = -1 # Free local index mapping. # DIAGNOSTIC: log the pop on the owning rank so we can correlate @@ -3026,6 +3199,15 @@ def _release_gpu_kv_pages(self, local_sequence_ids: List[int]) -> None: f"Rank {self.rank} Released GPU KV pages for global_idx: {global_sequence_ids}" ) + # Kimi-Linear: release KDA state slots alongside GPU KV pages (no-op for + # other models — slot_manager is None). + try: + from batchgen.models.moonshotai.kimi_linear.wrappers import KimiLinearKDAWrapper + if KimiLinearKDAWrapper.slot_manager is not None: + KimiLinearKDAWrapper.free_sequences(global_sequence_ids) + except ImportError: + pass + # FIX Bug 2: Remove from tracking set and reset gpu_pages_allocated for local_idx in local_sequence_ids: uuid = self._local_to_uuid_map.get(local_idx) @@ -3578,11 +3760,16 @@ def _execute_single_kv_migration(self, uuid: str, from_rank: int, to_rank: int) # are already contiguous (.contiguous() returns same tensor, not a copy) dist.send(tensor=qb.encoded["input_ids"].clone(), dst=to_rank, group=gloo_group) dist.send(tensor=qb.decoded_tokens.clone(), dst=to_rank, group=gloo_group) - # Free buffer slot after send completes - seq_for_slot = self.global_batch.get_sequence(uuid) - if hasattr(seq_for_slot, '_buffer_slot') and seq_for_slot._buffer_slot >= 0: - self._buffer_pool.free_slot(seq_for_slot._buffer_slot) - seq_for_slot._buffer_slot = -1 + # The buffer slot is NOT freed here. slot -> row is GLOBAL + # state: every rank allocates the same slot for the same + # sequence at tokenization and frees it in _report_completion, + # and the destination below reuses this very slot index. + # Freeing it only on the source made this rank's pool disagree + # with every other rank's -- it could hand row S to a new + # admission while everyone else still reads S as this sequence, + # and it left _buffer_slot = -1, so an eviction re-entry wrote + # into row -1 (the LAST row). Now that input_ids is one + # node-shared segment that divergence is data corruption. if BATCHGEN_CB_DEBUG: logging.debug(f"MIGRATION: Rank {self.rank}: Sent query_book for {uuid[:8]}...") else: @@ -3811,9 +3998,19 @@ def _rebalance_host_kv(self) -> None: f"reusing existing_slot={existing_slot}, budget={budget}" ) if existing_slot < 0: - logging.info(f"Rank {self.rank}: Migration receive {uuid[:8]} has no buffer slot (expected for cross-rank migration), allocating new") - existing_slot = self._buffer_pool.allocate_slot() - seq._buffer_slot = existing_slot + # Was: allocate a fresh slot here. That is a rank-LOCAL + # allocation of a globally-agreed index, so this rank would + # then write the sequence into a row every other rank reads + # as somebody else's -- silent corruption of the shared + # input_ids segment. The slot is allocated on every rank at + # tokenization and released on every rank in + # _report_completion, so reaching here means that invariant + # is already broken. + raise QueryBookPoolCapacityError( + f"Rank {self.rank}: migration receive of {uuid[:8]} found no " + f"buffer slot (_buffer_slot={existing_slot}); slot assignment " + f"has diverged from the other ranks" + ) self._buffer_pool.input_ids_buffer[existing_slot, :budget] = pending['input_ids'][0, :budget] self._buffer_pool.decoded_tokens_buffer[existing_slot, :] = pending['decoded_tokens'][0, :] input_ids_view = self._buffer_pool.get_input_ids_view(existing_slot, budget) @@ -3920,6 +4117,7 @@ def process_new_batch( # Step 1: Initialize global batch self.global_batch = SequenceBatch() + self._prefill_completed_results = {} for idx, text in enumerate(global_prompts): max_dec = self.max_decoding_length if per_sequence_max_tokens is not None and idx < len(per_sequence_max_tokens): @@ -4202,6 +4400,128 @@ def _sync_decode_uuids_tensor( self._make_sync_context(), decode_uuids ) + # ============ QueryBook Buffer Pool ============ + + def _node_shared_tag(self) -> str: + """Run-unique tag shared by every rank, for shared-memory segment names.""" + if self._shared_buffer_tag is None: + tag = [os.urandom(6).hex() if self.rank == 0 else None] + dist.broadcast_object_list(tag, src=0) + self._shared_buffer_tag = tag[0] + return self._shared_buffer_tag + + def _ensure_buffer_pool( + self, + required_rows: int, + required_input_width: int, + required_decode_width: int, + reason: str, + ) -> None: + """Allocate — or grow — the QueryBook buffer pool. + + COLLECTIVE: every rank must call this with identical arguments. They do, + because both call sites derive the requirement from the tokenized batch, + which is all-gathered to every rank before this runs. + + ``input_ids_buffer`` is ONE shared-memory segment per node. Sizing is by + actual need: ``required_input_width`` is the widest ``seq_extended_size`` + (prompt + that request's decode budget) the batch will ask for, capped at + the model context length — never the context length itself, and never the + ``--max-pool-size`` flag, which keeps its row-count meaning only. + + A later admission that needs more never silently truncates: it grows the + pool with a WARNING naming both sizes, copies the live rows over and + rebinds every view. ``get_input_ids_view`` hard-fails + (``QueryBookPoolCapacityError``) if a request ever slips past this. + """ + old = self._buffer_pool + if old is not None and ( + required_rows <= old.num_sequences + and required_input_width <= old.input_ids_width + and required_decode_width <= old.max_decoding_length + ): + return + + rows = max(required_rows, old.num_sequences if old is not None else 0) + in_w = max(required_input_width, old.input_ids_width if old is not None else 0) + dec_w = max(required_decode_width, old.max_decoding_length if old is not None else 0) + + self._buffer_pool_generation += 1 + node_id = self.rank // NUM_GPUS_PER_NODE + name = ( + f"batchgen_input_ids_{self._node_shared_tag()}" + f"_n{node_id}_g{self._buffer_pool_generation}" + ) + is_creator = (self.rank % NUM_GPUS_PER_NODE) == 0 + shared_input_ids, shm = allocate_node_shared_int64( + name, rows, in_w, is_creator, dist.barrier + ) + new_pool = QueryBookBufferPool( + num_sequences=rows, + input_ids_width=in_w, + max_decoding_length=dec_w, + pad_token_id=self.pad_token_id, + input_ids_buffer=shared_input_ids, + input_ids_shm=shm, + ) + shared_gib = rows * in_w * 8 / 2**30 + private_gib = rows * dec_w * 8 / 2**30 + if old is None: + logging.info( + f"Rank {self.rank}: QueryBook pool allocated ({reason}): rows={rows}, " + f"input_ids_width={in_w}, decoded_width={dec_w} -> input_ids " + f"{shared_gib:.3f} GiB SHARED per node ('{name}'), decoded_tokens " + f"{private_gib:.3f} GiB per rank" + ) + else: + logging.warning( + f"Rank {self.rank}: QueryBook pool GROWN ({reason}): rows " + f"{old.num_sequences}->{rows}, input_ids_width " + f"{old.input_ids_width}->{in_w}, decoded_width " + f"{old.max_decoding_length}->{dec_w}; new input_ids segment " + f"{shared_gib:.3f} GiB SHARED per node ('{name}')" + ) + new_pool.adopt(old) + self._buffer_pool = new_pool + if old is not None: + self._rebind_buffer_pool_views() + self._retire_buffer_pool(old, is_creator) + + def _retire_buffer_pool(self, old: QueryBookBufferPool, is_creator: bool) -> None: + """Drop a superseded pool's NAME but keep its mapping alive. + + Views handed out before the grow may still be referenced somewhere this + rebind does not reach; unmapping under them would segfault. Unlinking on + the node's creator keeps /dev/shm from accumulating one entry per grow — + POSIX frees the pages once the last mapping goes, i.e. at process exit. + """ + self._retired_buffer_pools.append(old) + if is_creator and old.input_ids_shm is not None: + try: + old.input_ids_shm.unlink() + except FileNotFoundError: + pass + + def _rebind_buffer_pool_views(self) -> None: + """Repoint every live sequence and query-book entry at the current pool.""" + pool = self._buffer_pool + rebound = 0 + for seq in self.global_batch: + slot = getattr(seq, '_buffer_slot', -1) + if slot < 0: + continue + input_ids_view = pool.get_input_ids_view(slot, seq.kv_token_budget) + decoded_view = pool.get_decoded_tokens_view(slot) + seq.input_ids = input_ids_view + seq.decoded_tokens = decoded_view + local_idx = self._uuid_to_local_map.get(seq.uuid) + if local_idx is not None and self.query_book and local_idx in self.query_book: + entry = self.query_book[local_idx] + entry.encoded["input_ids"] = input_ids_view + entry.decoded_tokens = decoded_view + rebound += 1 + logging.warning(f"Rank {self.rank}: rebound {rebound} sequences onto the grown QueryBook pool") + # ============ Tokenization and Assignment ============ def _tokenize_global_batch(self) -> None: @@ -4372,17 +4692,30 @@ def _tokenize_global_batch(self) -> None: # Use max_pool_size for pre-allocation if in pool mode (allows future admissions) pool_capacity = max(num_seqs, self._max_pool_size) if self._max_pool_size > 0 else num_seqs - self._buffer_pool = QueryBookBufferPool( - num_sequences=pool_capacity, - model_context_length=self.model_context_length, - max_decoding_length=self.max_decoding_length, - pad_token_id=self.pad_token_id, + # Width by actual need: the widest seq_extended_size the loop below will + # ask get_input_ids_view() for. Sizing it at model_context_length instead + # costs 8 bytes x pool_capacity x context — 80 GiB per worker at K3's 1M + # context — for a buffer whose rows are only ever read up to their own + # prompt length. + required_width = min( + max_prompt_length + self.max_decoding_length, + self.model_context_length, ) + self._ensure_buffer_pool( + required_rows=pool_capacity, + required_input_width=required_width, + required_decode_width=self.max_decoding_length, + reason="legacy batch tokenization", + ) + # Legacy mode tokenizes a whole new global batch per call, so the slot + # bookkeeping (and buffer contents) must start clean even when the + # existing allocation is reused. + self._buffer_pool.reset() t_alloc = time.perf_counter() - phase3_start logging.info( - f"Rank {self.rank}: Phase 3 buffer pool allocated in {t_alloc:.2f}s " - f"(input_ids: [{num_seqs}, {self.model_context_length}], " - f"decoded_tokens: [{num_seqs}, {self.max_decoding_length}])" + f"Rank {self.rank}: Phase 3 buffer pool ready in {t_alloc:.2f}s " + f"(input_ids: [{pool_capacity}, {self._buffer_pool.input_ids_width}] shared per node, " + f"decoded_tokens: [{pool_capacity}, {self._buffer_pool.max_decoding_length}] per rank)" ) for seq_i, seq in enumerate(self.global_batch): @@ -4809,7 +5142,7 @@ def _check_and_handle_completions( seq._rep_detected = True seq.eos_reached = True logging.warning( - f"Rank {self.rank}: REPETITION (ngram) {seq.uuid[:8]} " + f"Rank {self.rank}: REPETITION (ngram) {seq.uuid} " f"gid={seq.global_idx} at decoded_len={dl}" ) @@ -4889,8 +5222,129 @@ def _submit_completed_to_incremental_writer( for global_idx, tokens, finish_reason in rank_tokens: writer.submit(global_idx, tokens, finish_reason=finish_reason) + def _finish_prefill_completed_sequences(self, prefill_uuids: List[str]) -> List[str]: + """Complete the sequences whose budget is satisfied by the prefill token. + + Prefill already samples the first token and appends it through the + normal decode write path (``query_book[..].decoded_tokens`` at + ``seq.decoded_length``, then ``decoded_length += 1``; see the writeback + loop at the end of ``prefill``/``prefill_prepacked``). For a + ``max_tokens=1`` request that token IS the whole completion, so the + sequence is finished before decode starts. Previously every prefilled + sequence was handed to the decode phase unconditionally, which + - loaded the decode model and configured decoding for nothing, and + - on a prefill-only model (Kimi-K3 ``stream_all_modules``, M-PR-6) + replaced the answer with decode's ``NotImplementedError`` text. + + Only the length budget is checked here: that is exactly the first test + decode's own boundary check makes (``CompletionHandler`` / + ``_check_and_handle_completions``: ``decoded_length >= + max_decode_length``), and it is also the test that wins in + ``get_finish_reason``, so these sequences report the same + length-capped ``finish_reason`` they report today. Sequences that + stop for any other reason (EOS, context limit) are left PREFILLED and + reach decode exactly as before — ``max_tokens > 1`` behaviour is + unchanged. + + Rank alignment: ``decoded_length`` is advanced only on the owning + rank, so the set is derived AFTER ``_sync_sequence_metadata`` + replicates it to every rank. Every rank then computes the identical + set from identical batch-global state — required both for the + collectives below and because the resulting PREFILLED -> COMPLETED + transition is what makes the decode ``while`` loop's + ``has_prefilled()`` false. A rank-divergent set would deadlock the + next collective. + + Returns the list of completed uuids (identical on every rank). + """ + if not prefill_uuids: + return [] + + # Replicate owner-side decoded_length / current_context_length to all + # ranks. Also required by _report_completion, which reads those fields + # on rank 0 for sequences owned elsewhere. + self._sync_sequence_metadata(prefill_uuids) + + completed_uuids = [] + for uuid in prefill_uuids: + seq = self.global_batch.get_sequence(uuid) + if seq is None or seq.status != SequenceStatus.PREFILLED: + continue + if seq.decoded_length >= seq.max_decode_length: + completed_uuids.append(uuid) + + if not completed_uuids: + return [] + + if self.rank == 0: + logging.info( + f"[PREFILL] {len(completed_uuids)}/{len(prefill_uuids)} sequences " + f"completed at prefill (decode budget satisfied by the first " + f"sampled token); they skip the decode phase" + ) + + # Same order as the decode-phase completion handling in generate(): + # writer -> gather text -> release KV -> scalar cleanup -> status -> + # report. _submit_completed_to_incremental_writer and + # _gather_completed_tokens are collectives; every rank calls them with + # the identical uuid list. + self._submit_completed_to_incremental_writer(completed_uuids) + gathered_texts = self._gather_completed_tokens(completed_uuids) + + my_completed = [u for u in completed_uuids if u in self._uuid_to_local_map] + if my_completed: + # prefill_prepacked writes KV straight to host, so most of these + # never registered with the GPU paged manager. + gpu_allocated = [u for u in my_completed if u in self._sequences_with_gpu_kv] + if gpu_allocated: + self._release_gpu_kv_pages(self._get_local_indices_for_uuids(gpu_allocated)) + self._release_host_kv_pages_for_batch(my_completed) + for uuid in completed_uuids: + seq = self.global_batch.get_sequence(uuid) + if seq is not None: + seq.gpu_pages_allocated = 0 + seq.host_pages_allocated = 0 + seq.host_token_capacity = 0 + self._sequences_with_gpu_kv.discard(uuid) + + self._update_batch_status(completed_uuids, SequenceStatus.COMPLETED) + + # Legacy /v1/inference gathers results at the END of generate() by + # iterating _local_to_uuid_map. _report_completion below pops that map + # (release_local_query_slot also drops the query_book entry), so a + # C4-completed sequence would be invisible to that gather and the + # request would return "Results unexpectedly empty after inference". + # Capture the text while the slot still exists. + # Gated on the absence of a response queue: pool/batch mode is fed by + # _report_completion and _submit_completed_to_incremental_writer, so it + # must NOT also accumulate here -- that store is never drained in a + # persistent server and would grow without bound. + if self._response_queue is None: + # getattr, not a plain attribute read: hot reload rebinds methods on a + # LIVE worker and never re-runs __init__, so an attribute introduced in + # __init__ is absent on a reloaded process. _validate_reload + # (server_worker_main_loop.py:68) warns about exactly this and does not + # fix it -- "These will cause AttributeError if accessed." + store = getattr(self, '_prefill_completed_results', None) + if store is None: + store = self._prefill_completed_results = {} + for uuid in completed_uuids: + local_idx = self._uuid_to_local_map.get(uuid) + seq = self.global_batch.get_sequence(uuid) + if local_idx is None or seq is None or local_idx not in self.query_book: + continue + _decoded = self.query_book[local_idx].decoded_tokens[:, :seq.decoded_length] + store[seq.global_idx] = self._decode_tokens_to_string(_decoded) + + # Runs LAST: _report_completion pops the local-index map and frees the + # buffer-pool slot. + for uuid in completed_uuids: + self._report_completion(uuid, gathered_text=gathered_texts.get(uuid)) + + return completed_uuids + def _try_load_new_sequences( - self, + self, current_decode_uuids: List[str], current_local_indices: List[int] ) -> Tuple[List[str], List[int]]: @@ -5114,19 +5568,17 @@ def generate_persistent(self): # Initialize empty global batch (Init may have created one via _reset) self.global_batch = SequenceBatch() - # Pre-allocate buffer pool for max_pool_size. - # Use model_context_length for decoded_tokens buffer (not max_decoding_length) - # because per-request max_completion_tokens can be up to the full context window. - self._buffer_pool = QueryBookBufferPool( - num_sequences=self._max_pool_size, - model_context_length=self.model_context_length, - max_decoding_length=self.model_context_length, - pad_token_id=self.pad_token_id, - ) + # The buffer pool is NOT pre-allocated here. Both of its widths depend on + # the requests: input_ids needs prompt + that request's decode budget, + # decoded_tokens needs that request's max_completion_tokens — and neither + # is known until the first admission is tokenized. Sizing them at + # model_context_length "just in case" is what allocated 2 x 80 GiB per + # worker at K3's 1,048,576-token context. _tokenize_admitted_sequences + # allocates on first admission and grows if a later one needs more. + self._buffer_pool = None logging.info( - f"Rank {self.rank}: Buffer pool pre-allocated for {self._max_pool_size} sequences " - f"(context_length={self.model_context_length}, " - f"max_decoding={self.model_context_length})" + f"Rank {self.rank}: Buffer pool deferred to first admission " + f"(rows={self._max_pool_size}, widths sized per batch)" ) # Initialize index maps @@ -5411,6 +5863,12 @@ def generate(self): container = [msg] dist.broadcast_object_list(container, src=0) self._handle_hot_reload(msg) + elif isinstance(msg, dict) and "prompts" in msg: + # Legacy /v1/inference payload -- see the matching + # guard in _poll_admissions. The `else` below would + # swallow it and the caller would then steal the next + # completion off the shared response queue. + raise LegacyInferenceDeprecated() else: status = torch.tensor([0, 0, 0], dtype=torch.int32, device=self.torch_device) dist.broadcast(status, src=0) @@ -5571,6 +6029,14 @@ def generate(self): self._update_batch_status(prefill_uuids, SequenceStatus.PREFILLED) dist.barrier() + # C4: a request whose whole budget is the prefill-sampled + # token is DONE here. Completing it now (instead of sending + # it into decode to be completed on the first boundary + # check) keeps the decode phase — and its decode-model load + # — off the critical path for max_tokens=1, and is the only + # way a prefill-only model can answer at all. + self._finish_prefill_completed_sequences(prefill_uuids) + # After prefill completes, poll for newly arrived sequences. # If more QUEUEING sequences exist and host KV has capacity, # loop back to prefill instead of entering decode. @@ -5651,6 +6117,11 @@ def generate(self): # Incremental write: submit sequences completed between decode rounds if global_completed: + # Refresh rank-0's sequence replicas first: _report_completion + # reads prompt_length/decoded_length from the local entry, + # which is stale here for sequences owned by other ranks + # (only the completion BIT was all-reduced above). + self._sync_sequence_metadata(list(global_completed)) self._submit_completed_to_incremental_writer(list(global_completed)) # Gather decoded tokens from owning ranks before reporting # (each rank only writes decoded tokens for its own sequences) @@ -5873,6 +6344,9 @@ def generate(self): # With 12K sequences × 1MB tensors = 12GB, all_gather_object OOMs. # Gathering strings (~KB each) instead reduces memory by ~100x. local_results = [] + # Sequences completed during prefill (C4) were reported and popped from + # the local maps back then; their text was captured at that point. + local_results.extend(getattr(self, '_prefill_completed_results', {}).items()) for local_idx, uuid in self._local_to_uuid_map.items(): seq = self.global_batch.get_sequence(uuid) if seq is None: @@ -5998,6 +6472,10 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: ) # STEP 1: Configure model for prefill + # Hand the NCCL communicator to managers that need it during prefill + # (e.g. Kimi-Linear MoE EP all-reduce); harmless no-op for others. + if hasattr(self.parallel_manager, "set_comm"): + self.parallel_manager.set_comm(self.comm) self.model, self.weight_copy_task = self.parallel_manager.configure_prefill() self.set_phase("prefill") @@ -6106,6 +6584,14 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: # Rebuild input_ids with new prompt — reuse buffer pool slot seq_extended_size = seq.kv_token_budget slot = seq._buffer_slot + if slot < 0: + # A negative index silently rewrites the LAST row of the pool, + # which is another sequence's prompt (and, now that input_ids is + # node-shared, every rank's copy of it). + raise QueryBookPoolCapacityError( + f"Rank {self.rank}: re-entry of {uuid[:8]} has no buffer slot " + f"(_buffer_slot={slot}); slot assignment has diverged" + ) self._buffer_pool.input_ids_buffer[slot, :] = 0 self._buffer_pool.input_ids_buffer[slot, :new_prompt_len] = evicted_ids seq.input_ids = self._buffer_pool.get_input_ids_view(slot, seq_extended_size) @@ -6915,6 +7401,9 @@ def prefill_prepacked(self, batch: list[int]): output_tokens = [] + # Pure forward wall time: started here so configure_prefill (already + # reported separately as `Config completed`) is NEVER folded in. + _prefill_forward_t0 = time.perf_counter() with torch.inference_mode(): for batch_idx, (seq_start, seq_end) in tqdm( enumerate(micro_batches), @@ -6974,7 +7463,10 @@ def prefill_prepacked(self, batch: list[int]): ) batch_max_seqlen = max(batch_seq_lengths) - # Set up Attn_Wrapper for this micro-batch + # Set up Attn_Wrapper for this micro-batch. + # These class attrs are the per-step worker->model contract read + # by attention wrappers; see semantics in + # batchgen-context/architecture/PSM_WORKER_CONTRACT.md (§2) Attn_Wrapper.prepack_mode = True Attn_Wrapper.prepack_cu_seqlens = batch_cu_seqlens Attn_Wrapper.prepack_max_seqlen = batch_max_seqlen @@ -6999,17 +7491,63 @@ def prefill_prepacked(self, batch: list[int]): # Reshape to 3D: [1, batch_total_tokens, hidden_dim] hidden_states = inputs_embeds.unsqueeze(0) + # unsqueeze is a VIEW, so `hidden_states` already keeps the + # embedding storage alive for exactly as long as layer 0 needs + # it. Keeping `inputs_embeds` bound as well pins that storage + # for the WHOLE stack instead -- 1.75 GiB dead across layers + # 1-92 at S=131,072 / H=7168 / bf16 + # (batchgen_design/model_support/kimi_k3/ + # PREFILL_MEMORY_AUDIT.md section 7, fix 4). + del inputs_embeds + + # Block Attention Residuals (Kimi-K3): the depth-mix REPLACES the + # classic residual body, so the per-layer state has to be carried + # here. Without it every layer sees a zero-width residual and the + # model runs happily while computing something that is not K3 -- + # wrong text, no error. Mirrors KimiLinearModel.forward + # (kimi_linear/model.py:880-910), which is the eager reference. + use_attn_res = getattr(self.model.model, "use_attn_residuals", False) + block_residual = None + if use_attn_res: + # Zero-column view of a buffer preallocated for ALL the + # stack's block boundaries, so the per-boundary `cat` never + # holds the (S,nb,H) and (S,nb+1,H) tensors at once (12.25 + # GiB at K3's last boundary; PREFILL_MEMORY_AUDIT.md fix 3). + # `block_residual = None` above is load-bearing: it drops + # the previous micro-batch's view before the next buffer is + # allocated. + block_residual = self.model.model._new_block_residual(hidden_states) for layer_idx, decoder_layer in enumerate(self.model.model.layers): - layer_outputs = decoder_layer( - hidden_states, - attention_mask=None, - position_ids=None, - past_key_value=None, - output_attentions=False, - use_cache=False, - ) - hidden_states = layer_outputs[0] + if use_attn_res: + hidden_states, block_residual = decoder_layer( + hidden_states, + attention_mask=None, + position_ids=None, + past_key_value=None, + output_attentions=False, + use_cache=False, + block_residual=block_residual, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=None, + position_ids=None, + past_key_value=None, + output_attentions=False, + use_cache=False, + ) + hidden_states = layer_outputs[0] + + # Output depth-mix, then norm -- that ORDER is load-bearing + # (kimi_linear/model.py:904-913). + if use_attn_res: + hidden_dim = hidden_states.shape[2] + batch_sz, seq_sz = hidden_states.shape[:2] + hidden_states = self.model.model._apply_output_attn_res( + hidden_states.view(-1, hidden_dim), block_residual + ).view(batch_sz, seq_sz, hidden_dim) # Final norm hidden_states = self.model.model.norm(hidden_states) @@ -7045,6 +7583,52 @@ def prefill_prepacked(self, batch: list[int]): ) output_tokens.append(batch_new_tokens) + # The FIRST generated token, straight out of prefill. A + # max_tokens=1 request is now completed right after prefill + # (PREFILL_PLAN C4, _finish_prefill_completed_sequences) and + # the token reaches the client through the normal response + # path, so this log is a cross-check of that path rather than + # the only way to see the token. Rank 0 only; ids only (the + # worker has no tokenizer -- decode them client-side). + if self.rank == 0: + logging.info( + "[PREFILL] first sampled token ids: %s", + batch_new_tokens.reshape(-1).tolist()[:16]) + + _prefill_forward_s = time.perf_counter() - _prefill_forward_t0 + + # Structured prefill record, one JSON line per rank that actually ran a + # prefill (batchgen-benchmark docs/prefill_metrics_proposal.md). The + # report tool prefers this over scraping the tqdm bar, which is + # presentation, not an API. + # + # Deliberately NOT gated on rank 0, departing from the proposal's + # "rank 0 only". prefill_prepacked runs only under + # `if local_prefill_indices:`, and the tqdm bar above is + # disable=(self.rank != 0) -- so when rank 0 owns none of the batch + # there is neither a bar nor a rank-0 line anywhere in the log, and the + # run reports `Prefill: 0.0s` with nothing to scrape. That is exactly + # what the 131,069-token run produced when rank 2 owned the sequence. + # Every participating rank emits its own tagged line; the wall time for + # the batch is the MAX of `prefill_s` over the emitting ranks. + logging.info("[METRICS] %s", json.dumps({ + "phase": "prefill", + "prefill_s": _prefill_forward_s, + "sequences": num_sequences, + "tokens_total": total_tokens_all, + "seq_len_min": min(seq_lengths_list) if seq_lengths_list else 0, + "seq_len_max": max(seq_lengths_list) if seq_lengths_list else 0, + "micro_batches": len(micro_batches), + "max_tokens_per_micro_batch": MAX_TOKENS_PER_MICRO_BATCH, + "world_size": self.world_size, + "rank": self.rank, + # Guarded: the proposal's unguarded output_tokens[0] is an IndexError + # on a rank that ran zero micro-batches. + "first_sampled_token_ids": ( + output_tokens[0].reshape(-1).tolist()[:16] if output_tokens else [] + ), + }, separators=(",", ":"))) + # Reset prepack mode Attn_Wrapper.prepack_mode = False Attn_Wrapper.prepack_cu_seqlens = None @@ -9542,6 +10126,10 @@ def decoding_continuous( # P0: Pre-allocate pinned memory buffer for non-blocking GPU→CPU token transfer _new_tokens_pinned = torch.empty(max(max_batch_size, 1), 1, dtype=torch.long, pin_memory=True) + # Heartbeat state for the rate-limited [DECODE] progress line below + _hb_last_time = time.perf_counter() + _hb_tokens = 0 + # Main decode loop — enable decode watchdog for monitoring self.enable_decode_watchdog() while decode_uuids: @@ -9552,6 +10140,20 @@ def decoding_continuous( self.feed_watchdog() self.feed_decode_watchdog() + # Rate-limited decode heartbeat (rank 0, ~every 30 s) so the log + # monitor sees liveness during long decode phases + _hb_tokens += len(decode_uuids) + if self.rank == 0 and time.perf_counter() - _hb_last_time >= 30.0: + _hb_elapsed = time.perf_counter() - _hb_last_time + _hb_finished = len(self.global_batch.get_sequences_by_status(SequenceStatus.COMPLETED)) + logging.info( + f"[DECODE] step={self._cumulative_decode_iterations} " + f"active={len(decode_uuids)} finished={_hb_finished} " + f"tok/s={_hb_tokens / _hb_elapsed:.2f}" + ) + _hb_last_time = time.perf_counter() + _hb_tokens = 0 + # Page boundary check - use DECISION_INTERVAL (configurable via BATCHGEN_DECISION_FREQUENCY_PAGES) if local_iteration - last_boundary >= self.DECISION_INTERVAL: last_boundary = local_iteration @@ -10464,7 +11066,7 @@ def kv_append_callback_aux(layer_idx: int, k_tensor: torch.Tensor, v_tensor: tor lifespan.dump_lifespan(seq.uuid, seq.global_idx, seq._lifespan_log, "REPETITION") logging.warning( - f"Rank {self.rank}: REPETITION {seq.uuid[:8]} gid={seq.global_idx} " + f"Rank {self.rank}: REPETITION {seq.uuid} gid={seq.global_idx} " f"token={token_id} x{seq._rep_count} at decoded_len={seq.decoded_length}" ) else: @@ -10478,7 +11080,7 @@ def kv_append_callback_aux(layer_idx: int, k_tensor: torch.Tensor, v_tensor: tor seq._rep_detected = True seq.eos_reached = True logging.warning( - f"Rank {self.rank}: REPETITION (ngram) {seq.uuid[:8]} " + f"Rank {self.rank}: REPETITION (ngram) {seq.uuid} " f"gid={seq.global_idx} at decoded_len={_dl}" ) @@ -11816,6 +12418,15 @@ def _unregister_fp8_weights(self): if not hasattr(self.loaded_model_config, 'first_k_dense_replace'): return + # Models whose MoE layers don't expose DeepSeek-style `.mlp` (e.g. + # Kimi-Linear uses `.block_sparse_moe` with BF16 experts) have no FP8 + # weights to unregister. + _fkd = self.loaded_model_config.first_k_dense_replace + if _fkd < len(self.model.model.layers) and not hasattr( + self.model.model.layers[_fkd], 'mlp' + ): + return + for layer_idx in range(len(self.model.model.layers)): attn_module = self.model.model.layers[layer_idx].self_attn if hasattr(attn_module, '_unregister_fp8_weights'): diff --git a/batchgen/ckpt_converter/ckpt_converter.py b/batchgen/ckpt_converter/ckpt_converter.py index cbbae4083..a5da834f4 100644 --- a/batchgen/ckpt_converter/ckpt_converter.py +++ b/batchgen/ckpt_converter/ckpt_converter.py @@ -90,6 +90,22 @@ def _apply_marlin_repack(self, ckpt): packed = ckpt[name] # [N, K//8] int32 or uint8 scale = ckpt[scale_name] # [N, K//32] bf16 + # Refuse MXFP4. This path is uniform-INT4-only: it reinterprets the + # uint8 buffer as packed uint4b8 below, and converts the SCALE + # TENSOR VALUE to bfloat16 — an E8M0 exponent byte of 121 would + # become 121.0 instead of 2**(121-127) = 0.015625. Both corruptions + # are silent. Scoped to tensors this function would actually touch: + # an unscoped check would reject every MXFP4 checkpoint the repack + # never matches, and `batchgen/tools/convert_checkpoint.py:144` + # passes marlin=True BY DEFAULT. + if scale.dtype == torch.uint8: + raise ValueError( + f"{scale_name}: uint8 (E8M0) scale on a tensor the Marlin " + "repack matches. Marlin handles uniform INT4 with BF16 " + "scales only; MXFP4 needs an E2M1 nibble decode and E8M0 " + "scale handling. Re-run with --no-marlin." + ) + # Convert uint8 → int32 if needed if packed.dtype == torch.uint8: N_dim = scale.shape[0] diff --git a/batchgen/config/model_registry.py b/batchgen/config/model_registry.py index f8c817f21..a84f4704d 100644 --- a/batchgen/config/model_registry.py +++ b/batchgen/config/model_registry.py @@ -94,6 +94,10 @@ "GLM-5.1": "glm_moe_dsa", "GLM-5-FP8": "glm_moe_dsa", "GLM-5": "glm_moe_dsa", + # Kimi-Linear (testbed) + Kimi-K3 family (hybrid KDA + NoPE-MLA MoE) + "Kimi-Linear-48B-A3B": "kimi_linear", + "Kimi-Linear": "kimi_linear", + "Kimi-K3": "kimi_k3", } for model_id in KIMI_K25_BACKEND_MODEL_IDS: @@ -217,16 +221,24 @@ def load_config(model_identifier: str) -> "BaseModelConfig": config = None - # Step 1: Try to detect model type from identifier patterns - detected_type = _detect_model_type_from_identifier(model_identifier) - if detected_type and detected_type in CONFIG_REGISTRY: - logger.info(f"Using built-in config for model_type={detected_type}") - config = CONFIG_REGISTRY[detected_type]() - config._name_or_path = model_identifier - return config + # A local checkout's config.json is authoritative — prefer it over the + # name-pattern shortcut (Step 1), which returns curated *defaults* and would + # silently drop data-driven fields (e.g. kimi_linear's `linear_attn_config`) + # for a local dir whose name happens to match a pattern. + _local_config_json = Path(model_identifier) / "config.json" + _is_local_dir = _local_config_json.exists() + + # Step 1: Try to detect model type from identifier patterns (HF model IDs). + if not _is_local_dir: + detected_type = _detect_model_type_from_identifier(model_identifier) + if detected_type and detected_type in CONFIG_REGISTRY: + logger.info(f"Using built-in config for model_type={detected_type}") + config = CONFIG_REGISTRY[detected_type]() + config._name_or_path = model_identifier + return config # Step 2: Check if it's a local directory with config.json - config_path = Path(model_identifier) / "config.json" + config_path = _local_config_json if config_path.exists(): with open(config_path, 'r') as f: data = json.load(f) @@ -313,6 +325,11 @@ def _import_model_configs(): except ImportError: pass + try: + from batchgen.models.moonshotai.kimi_linear import config as _ # noqa: F401 + except ImportError: + pass + try: from batchgen.models.minimax.minimax_m25 import config as _ # noqa: F401 except ImportError: diff --git a/batchgen/config/tokenizer_registry.py b/batchgen/config/tokenizer_registry.py index afc52e243..b0ac2d97a 100644 --- a/batchgen/config/tokenizer_registry.py +++ b/batchgen/config/tokenizer_registry.py @@ -42,6 +42,7 @@ """ from typing import Dict, Type, Optional, TYPE_CHECKING +import importlib import logging from .model_name_utils import KIMI_K25_BACKEND_MODEL_IDS @@ -85,6 +86,17 @@ "GLM-5": "glm_moe_dsa", "MiniMax-M2.5": "minimax_m25", "MiniMaxAI/MiniMax-M2.5": "minimax_m25", + "Kimi-Linear": "kimi_linear", + "kimi-linear": "kimi_linear", + # Kimi-K3 shares the Kimi-Linear ARCHITECTURE but NOT its tokenizer. The two + # ship different added_tokens_decoder tables over a byte-identical BPE merge + # file (163586 is "<|end_of_msg|>" in K3, "<|im_end|>" in the 48B) and K3 has + # no Jinja chat template at all -- its XTML format is Python. Cross-loading + # renders a 12-token K3 fragment as 32 marker-free tokens, silently + # (bug_log.md 2026-07-31). Neither string is a substring of the other, so + # ordering against "Kimi-Linear" does not matter. + "Kimi-K3": "kimi_k3", + "kimi-k3": "kimi_k3", } for model_id in KIMI_K25_BACKEND_MODEL_IDS: @@ -138,6 +150,18 @@ def load_tokenizer(model_identifier: str) -> "BaseTokenizer": logger.info(f"Using registered tokenizer for type={tokenizer_type}") # Tokenizer loads from its own package directory (no path argument) return TOKENIZER_REGISTRY[tokenizer_type]() + # Falling through to a later, less specific pattern means serving + # this model with a DIFFERENT model's tokenizer. That is intended + # and documented for GLM-5.2 (identical vocab, see above); anywhere + # else it is the bug_log.md 2026-07-31 failure mode. Never silent. + logger.warning( + "Model %r matched pattern %r -> tokenizer type %r, which is not " + "registered. Falling through to a less specific pattern; the " + "tokenizer that ends up serving this model is NOT the one its " + "name selected. This is correct only when the two share a vocab " + "AND a chat template -- verify before relying on it.", + model_identifier, pattern, tokenizer_type, + ) raise ValueError( f"No tokenizer registered for model: {model_identifier}. " @@ -158,51 +182,31 @@ def get_registered_tokenizers() -> Dict[str, Type["BaseTokenizer"]]: # Import model-specific tokenizers to register them # These imports trigger the @register_tokenizer decorators def _import_tokenizers(): - """Import all model-specific tokenizer modules to register them.""" - try: - from batchgen.models.deepseek.deepseekv4_flash import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.deepseek.deepseekv3 import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.deepseek.deepseekv2 import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.openai.gpt_oss_120b import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.mixtral import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.moonshotai.kimi_k25 import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.glm.glm5 import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.minimax.minimax_m25 import tokenizer as _ # noqa: F401 - except ImportError: - pass - - try: - from batchgen.models.glm.glm5 import tokenizer as _ # noqa: F401 - except ImportError: - pass + """Import all model-specific tokenizer modules to register them. + + A model package that cannot be imported (optional extra not installed, a + typo in a module, an asset missing from the wheel) is tolerated -- but it is + logged. Swallowing it silently turns a broken tokenizer into the misleading + "No tokenizer registered for model" further down. + """ + for module_path in ( + "batchgen.models.deepseek.deepseekv4_flash.tokenizer", + "batchgen.models.deepseek.deepseekv3.tokenizer", + "batchgen.models.deepseek.deepseekv2.tokenizer", + "batchgen.models.openai.gpt_oss_120b.tokenizer", + "batchgen.models.mixtral.tokenizer", + "batchgen.models.moonshotai.kimi_k25.tokenizer", + "batchgen.models.moonshotai.kimi_linear.tokenizer", + "batchgen.models.moonshotai.kimi_k3.tokenizer", + "batchgen.models.glm.glm5.tokenizer", + "batchgen.models.minimax.minimax_m25.tokenizer", + ): + try: + importlib.import_module(module_path) + except ImportError as exc: + logger.warning( + "Tokenizer module %s could not be imported (%s); any model " + "routed to it will fail to load a tokenizer.", module_path, exc) # Auto-import on module load diff --git a/batchgen/deprecation.py b/batchgen/deprecation.py new file mode 100644 index 000000000..f099220a6 --- /dev/null +++ b/batchgen/deprecation.py @@ -0,0 +1,26 @@ +"""The single place the /v1/inference deprecation is spelled out. + +Three modules must agree on this text and cannot import one another: +`batchgen_client` (imported by `batchgen/__init__`, so it must carry no heavy +dependencies), `server/http_server` (FastAPI + pydantic) and `batchgen_worker` +(torch). A module that imports nothing at all is the only seam the three of +them can share. +""" + +LEGACY_INFERENCE_ERROR_CODE = "legacy_inference_deprecated" + +LEGACY_INFERENCE_MESSAGE = ( + "/v1/inference is deprecated and disabled. Submit inference through the " + "batch API instead: POST /v1/files (purpose=batch), then POST /v1/batches. " + "The legacy path carried no request-id routing: on a pool-mode server it " + "parked in the worker admission loop and then took the next completion off " + "the shared response queue, corrupting a concurrent batch." +) + + +class LegacyInferenceDeprecated(RuntimeError): + """Raised wherever a /v1/inference-shaped request is refused.""" + + def __init__(self, message: str = LEGACY_INFERENCE_MESSAGE) -> None: + super().__init__(message) + self.code = LEGACY_INFERENCE_ERROR_CODE diff --git a/batchgen/kv_cache/host_kv_mananger_config.py b/batchgen/kv_cache/host_kv_mananger_config.py index a8ec16aef..83a5a1a91 100644 --- a/batchgen/kv_cache/host_kv_mananger_config.py +++ b/batchgen/kv_cache/host_kv_mananger_config.py @@ -145,6 +145,33 @@ def bytes_per_page(self) -> int: kv_dtype="bfloat16", ) +# Kimi-K3: MLA latent KV, SAME geometry as kimi-linear (compressed_kv_dim=576) +# but 93 ENGINE layers, not 27. K3 has 24 MLA layers sitting at engine indices +# 3,7,...,87,91,92 — and `wrappers.py::_offload_prepacked_kv` indexes the pool +# by ENGINE layer index, not by a dense MLA counter. Sized at 27 (the +# kimi-linear value it used to alias onto), every MLA layer at index >= 27 +# writes past the end of the pool: silent memory corruption from the very +# first prefill, at layer 28. +_KIMI_K3_MLA_PROFILE = _HostKVModelProfile( + num_layers=93, + num_k_heads=1, + k_head_dim=576, + num_v_heads=0, + v_head_dim=0, + kv_dtype="bfloat16", +) + +# Kimi-Linear: MLA latent KV (compressed_kv_dim=576, 27 engine layers; only +# the 7 MLA layers ever append — KDA layers hold no KV). +_KIMI_LINEAR_MLA_PROFILE = _HostKVModelProfile( + num_layers=27, + num_k_heads=1, + k_head_dim=576, + num_v_heads=0, + v_head_dim=0, + kv_dtype="bfloat16", +) + _PROFILE_REGISTRY: Dict[str, _HostKVModelProfile] = { "deepseek_mla": _DEEPSEEK_MLA_PROFILE, "deepseek_v4_flash": _DEEPSEEK_V4_FLASH_PROFILE, @@ -154,6 +181,8 @@ def bytes_per_page(self) -> int: "minimax_m25_gqa": _MINIMAX_M25_GQA_PROFILE, "glm5_mla": _GLM5_MLA_PROFILE, "glm5_indexer": _GLM5_INDEXER_PROFILE, + "kimi_linear_mla": _KIMI_LINEAR_MLA_PROFILE, + "kimi_k3_mla": _KIMI_K3_MLA_PROFILE, } _PROFILE_ALIASES: Dict[str, str] = {} @@ -197,6 +226,19 @@ def bytes_per_page(self) -> int: "minimax-m2.5", "minimax", ), + "kimi_linear_mla": ( + "moonshotai/kimi-linear-48b-a3b-instruct", + "moonshotai/kimi-linear", + "kimi-linear", + "kimi_linear", + ), + # K3 is NOT an alias of kimi-linear here: same KV geometry, 93 engine + # layers instead of 27. See _KIMI_K3_MLA_PROFILE. + "kimi_k3_mla": ( + "moonshotai/kimi-k3", + "kimi-k3", + "kimi_k3", + ), "glm5_mla": ( "zai-org/glm-5-fp8", "zai-org/glm-5", diff --git a/batchgen/kv_cache/kda_state_gpu_manager.py b/batchgen/kv_cache/kda_state_gpu_manager.py new file mode 100644 index 000000000..712a3f7e6 --- /dev/null +++ b/batchgen/kv_cache/kda_state_gpu_manager.py @@ -0,0 +1,379 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Sequence, Tuple, Union + +import torch + +from batchgen.kv_cache.coordinator_utils import ( + resolve_from_layer_mapping, +) +from batchgen.kv_cache.gpu_paged_kv_manager import ( + _normalize_device, + _normalize_gpu_layer_mapping, + _TensorStack, +) + + +@dataclass(frozen=True) +class KDAStateGPUConfig: + """Geometry for the fixed-size per-sequence KDA state pool. + + Unlike the rolling compressor there is no ring: the fla kernel updates + each sequence's recurrent + short-conv state *in place*, so a sequence + owns exactly one slot per state item for its whole lifetime. + + ``conv_width`` is the conv KERNEL width W; the conv pools store the last + W-1 raw inputs per slot — the ``causal_conv1d.cu`` state contract + (per-layer view shape ``(num_state_items, conv_dim, W-1)``). + """ + + num_kda_layers: int + num_state_items: int + num_heads: int # HV (number of value heads) + head_dim: int = 128 + conv_dim: Optional[int] = None # num_heads * head_dim; derived if None + conv_width: int = 4 # kernel width W; pools store W-1 entries per slot + recurrent_dtype: torch.dtype = torch.float32 + conv_dtype: torch.dtype = torch.bfloat16 + cuda_graph_max_slots: Optional[int] = None + logical_to_physical_layer: Optional[Sequence[int]] = None + + def resolved_conv_dim(self) -> int: + if self.conv_dim is not None: + return int(self.conv_dim) + return int(self.num_heads) * int(self.head_dim) + + +@dataclass(frozen=True) +class KDAStateGPUStats: + num_total_state_items: int + num_free_state_items: int + num_used_state_items: int + num_active_sequences: int + + +class KDAStateGPUManager: + """GPU storage for per-sequence Kimi Delta Attention recurrent + conv state. + + Each active sequence owns one fixed-size state item (one slot) per KDA + layer. The fla kernel mutates ``recurrent_state`` and the three short + causal-conv states in place, so this manager only handles slot + allocation, view export, free-list bookkeeping, and zero-on-alloc. + + M5.1: this manager is the canonical, CUDA-graph-ready home of the KDA + state — the recurrent pool, the three conv pools and the persistent + decode slot-index buffer are each allocated ONCE with a fixed address. + ``KimiLinearKDAWrapper``'s per-layer pools are views of these tensors + and its slot facade delegates all alloc/free/zeroing here. Per-layer + conv views are ``(num_state_items, conv_dim, conv_width-1)`` — the + ``causal_conv1d.cu`` layout (matching the wrapper). + """ + + manager_name = "KDAStateGPUManager" + + def __init__( + self, + *, + config: KDAStateGPUConfig, + device: Union[str, int, torch.device], + ) -> None: + self.config = config + self.device = _normalize_device(device) + if self.config.num_kda_layers <= 0: + raise ValueError("num_kda_layers must be > 0") + if self.config.num_state_items <= 0: + raise ValueError("num_state_items must be > 0") + if self.config.num_heads <= 0: + raise ValueError("num_heads must be > 0") + if self.config.head_dim <= 0: + raise ValueError("head_dim must be > 0") + if self.config.conv_width < 2: + raise ValueError( + "conv_width must be >= 2 (pools store W-1 entries per slot)" + ) + if self.config.resolved_conv_dim() <= 0: + raise ValueError("conv_dim must be > 0") + self._logical_to_physical_layer = _normalize_gpu_layer_mapping( + self.config.logical_to_physical_layer, + self.config.num_kda_layers, + ) + self._reset_runtime_state() + + # ------------------------------------------------------------------ + # lifecycle + # ------------------------------------------------------------------ + def initialize(self) -> None: + if self._is_initialized: + return + if self.device.type == "cuda": + torch.cuda.set_device(self.device) + cfg = self.config + conv_dim = cfg.resolved_conv_dim() + self._recurrent_state = torch.zeros( + ( + cfg.num_kda_layers, + cfg.num_state_items, + cfg.num_heads, + cfg.head_dim, + cfg.head_dim, + ), + dtype=cfg.recurrent_dtype, + device=self.device, + ) + # causal_conv1d.cu contract: a slot holds the last W-1 raw inputs, + # so per-layer views are contiguous (num_state_items, conv_dim, W-1) + # 3-D tensors satisfying the kernel's dim()==3 && size(2)==W-1 check. + # The 4-D allocation keeps ONE fixed base address per q/k/v pool + # (CUDA-graph capture requirement). + conv_shape = ( + cfg.num_kda_layers, + cfg.num_state_items, + conv_dim, + cfg.conv_width - 1, + ) + self._conv_q = torch.zeros( + conv_shape, dtype=cfg.conv_dtype, device=self.device + ) + self._conv_k = torch.zeros( + conv_shape, dtype=cfg.conv_dtype, device=self.device + ) + self._conv_v = torch.zeros( + conv_shape, dtype=cfg.conv_dtype, device=self.device + ) + self._free_state_items = _TensorStack(cfg.num_state_items) + self._ensure_prepared_state_slot_buffer() + self._is_initialized = True + + def destroy(self, *, empty_cuda_cache: bool = False) -> None: + if not self._is_initialized: + return + self._reset_runtime_state() + if empty_cuda_cache and torch.cuda.is_available(): + torch.cuda.empty_cache() + + # ------------------------------------------------------------------ + # slot allocation / free + # ------------------------------------------------------------------ + def allocate_state_item(self, sequence_id: int) -> int: + self._ensure_initialized() + return self._ensure_state_item(int(sequence_id)) + + def allocate_state_items_for_sequences( + self, sequence_ids: Sequence[int] + ) -> dict[int, int]: + self._ensure_initialized() + return { + int(seq_id): self._ensure_state_item(int(seq_id)) + for seq_id in sequence_ids + } + + def release_sequence_states(self, sequence_ids: Sequence[int]) -> None: + self._ensure_initialized() + reclaimed: list[int] = [] + for seq_id in [int(seq_id) for seq_id in sequence_ids]: + state_item_id = self._sequence_state_items.pop(seq_id, None) + if state_item_id is not None: + reclaimed.append(state_item_id) + if reclaimed: + self._free_state_items.push(reclaimed) + + def reset_state_items(self, state_item_ids: Sequence[int]) -> None: + """Zero recurrent + conv (+ aux, if present) state for the slots.""" + self._ensure_initialized() + slots = [int(s) for s in state_item_ids] + if not slots: + return + idx = torch.as_tensor(slots, dtype=torch.long, device=self.device) + for slot in slots: + if slot < 0 or slot >= self.config.num_state_items: + raise IndexError(f"state item id {slot} out of range") + # index_fill_ over the state-item dim (dim=1) for every layer at once. + self._recurrent_state.index_fill_(1, idx, 0) + self._conv_q.index_fill_(1, idx, 0) + self._conv_k.index_fill_(1, idx, 0) + self._conv_v.index_fill_(1, idx, 0) + # Per-slot auxiliary rows (block_reps on branches that carry them) + # are zeroed too so the F4 zero-on-alloc covers every pool. + block_reps = getattr(self, "_block_reps", None) + if block_reps is not None: + block_reps.index_fill_(0, idx, 0) + + # ------------------------------------------------------------------ + # view export + # ------------------------------------------------------------------ + def get_layer_recurrent_view(self, logical_layer: int) -> torch.Tensor: + """Recurrent state view [num_state_items, HV, head_dim, head_dim].""" + self._ensure_initialized() + physical = self.resolve_physical_layer(logical_layer) + return self._recurrent_state[physical] + + def get_layer_conv_views( + self, logical_layer: int + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Conv views (q, k, v), each [num_state_items, conv_dim, width-1].""" + self._ensure_initialized() + physical = self.resolve_physical_layer(logical_layer) + return ( + self._conv_q[physical], + self._conv_k[physical], + self._conv_v[physical], + ) + + def get_recurrent_tensors(self) -> torch.Tensor: + self._ensure_initialized() + return self._recurrent_state + + def get_conv_tensors( + self, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + self._ensure_initialized() + return (self._conv_q, self._conv_k, self._conv_v) + + # ------------------------------------------------------------------ + # decode-step slot preparation (CUDA graph static buffer) + # ------------------------------------------------------------------ + def prepare_decode_step(self, sequence_ids: Sequence[int]) -> torch.Tensor: + """Fill the static slot-index buffer with each sequence's slot. + + Returns the (view onto the) prepared slot buffer for the batch. + """ + self._ensure_initialized() + slots = [ + self._get_sequence_state_item(int(seq_id)) + for seq_id in sequence_ids + ] + self._write_prepared_state_slots(slots) + return self._prepared_state_slots[: len(slots)] + + def get_prepared_state_slots(self) -> torch.Tensor: + self._ensure_initialized() + return self._prepared_state_slots[: self._prepared_state_slot_count] + + # ------------------------------------------------------------------ + # sequence -> slot lookup + # ------------------------------------------------------------------ + @property + def sequence_state_items(self) -> dict[int, int]: + return dict(self._sequence_state_items) + + def get_sequence_state_item(self, sequence_id: int) -> int: + self._ensure_initialized() + return self._get_sequence_state_item(int(sequence_id)) + + def has_sequence_state_item(self, sequence_id: int) -> bool: + self._ensure_initialized() + return int(sequence_id) in self._sequence_state_items + + def get_stats(self) -> KDAStateGPUStats: + self._ensure_initialized() + used = self.config.num_state_items - self._free_state_items.size + return KDAStateGPUStats( + num_total_state_items=self.config.num_state_items, + num_free_state_items=self._free_state_items.size, + num_used_state_items=used, + num_active_sequences=len(self._sequence_state_items), + ) + + # ------------------------------------------------------------------ + # layer mapping + # ------------------------------------------------------------------ + @property + def uses_logical_layer_mapping(self) -> bool: + return self._logical_to_physical_layer is not None + + def resolve_physical_layer(self, logical_layer_id: int) -> int: + logical_layer_id = int(logical_layer_id) + if logical_layer_id < 0: + raise IndexError("logical layer id must be >= 0") + if self._logical_to_physical_layer is None: + if logical_layer_id >= self.config.num_kda_layers: + raise IndexError( + f"layer_idx {logical_layer_id} out of range" + ) + return logical_layer_id + return resolve_from_layer_mapping( + "GPU KDA state", + "state", + self._logical_to_physical_layer, + logical_layer_id, + ) + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + def _reset_runtime_state(self) -> None: + self._is_initialized = False + self._recurrent_state: Optional[torch.Tensor] = None + self._conv_q: Optional[torch.Tensor] = None + self._conv_k: Optional[torch.Tensor] = None + self._conv_v: Optional[torch.Tensor] = None + self._free_state_items: Optional[_TensorStack] = None + self._sequence_state_items: dict[int, int] = {} + self._prepared_state_slots: Optional[torch.Tensor] = None + self._prepared_state_slot_count = 0 + + def _ensure_initialized(self) -> None: + if not self._is_initialized: + raise RuntimeError( + "KDAStateGPUManager.initialize must be called before use" + ) + + def _ensure_state_item(self, sequence_id: int) -> int: + state_item_id = self._sequence_state_items.get(sequence_id) + if state_item_id is not None: + return state_item_id + if self._free_state_items.size <= 0: + raise RuntimeError("Insufficient free KDA state items") + state_item = self._free_state_items.pop(1) + state_item_id = int(state_item[0].item()) + # F4 fix: zero the (possibly recycled) slot across every layer's + # conv + recurrent pool on FRESH alloc — a recycled slot must not + # leak the previous sequence's state, and layers > 0 seeing + # has_initial_state=True for a just-allocated sequence stays + # harmless (zero state == no state). Idempotent re-allocs return + # above and never re-zero live state. + self.reset_state_items([state_item_id]) + self._sequence_state_items[sequence_id] = state_item_id + return state_item_id + + def _get_sequence_state_item(self, sequence_id: int) -> int: + state_item_id = self._sequence_state_items.get(int(sequence_id)) + if state_item_id is None: + raise KeyError( + f"Sequence {sequence_id} has no KDA state item" + ) + return state_item_id + + def _ensure_prepared_state_slot_buffer(self) -> torch.Tensor: + if self._prepared_state_slots is None: + max_slots = self.config.cuda_graph_max_slots + if max_slots is None: + max_slots = 1024 + self._prepared_state_slots = torch.full( + (int(max_slots),), -1, dtype=torch.int32, device=self.device + ) + return self._prepared_state_slots + + def _write_prepared_state_slots(self, slots: Sequence[int]) -> None: + buffer = self._ensure_prepared_state_slot_buffer() + count = len(slots) + if count > int(buffer.numel()): + raise ValueError( + "prepare_decode_step batch exceeds prepared state-slot buffer" + ) + if count: + buffer[:count].copy_( + torch.as_tensor(slots, dtype=torch.int32, device=self.device) + ) + previous_count = self._prepared_state_slot_count + if previous_count > count: + buffer[count:previous_count].fill_(-1) + self._prepared_state_slot_count = count + + +__all__ = [ + "KDAStateGPUConfig", + "KDAStateGPUManager", + "KDAStateGPUStats", +] diff --git a/batchgen/kv_cache/kimi_linear_kv_coordinator.py b/batchgen/kv_cache/kimi_linear_kv_coordinator.py new file mode 100644 index 000000000..6faa8b6cb --- /dev/null +++ b/batchgen/kv_cache/kimi_linear_kv_coordinator.py @@ -0,0 +1,448 @@ +"""Hybrid KV/state coordinator for the ``kimi_linear`` model family. + +Kimi-Linear (and Kimi-K3) interleave two fundamentally different attention +mechanisms across their layers: + + * **MLA layers** (NoPE Multi-head Latent Attention) keep a *paged* compressed + KV cache. Storage grows with sequence length (one entry per token), so it is + served by :class:`GPUPagedKVCacheManager` with ``num_k_heads=1`` and + ``k_head_dim=compressed_kv_dim`` (=576: ``kv_lora_rank`` 512 + ``qk_rope`` 64) + and ``num_v_heads=0`` (MLA stores only the joint compressed KV, no separate V). + + * **KDA layers** (Kimi Delta Attention, a gated linear-attention variant) keep a + *fixed-size recurrent state* plus short-conv states per sequence — the storage + does NOT grow with sequence length. These are served by ``KDAStateGPUManager`` + (one state item per active sequence). + +Which mechanism a given global layer uses is decided by +``KimiLinearConfig.is_kda_layer(idx)`` (KDA layers are 1-indexed in +``linear_attn_config.kda_layers`` — layer ``idx`` is KDA iff ``idx + 1`` is in +that list). This coordinator composes the two sub-managers and builds two +**complementary** ``logical_to_physical_layer`` maps over the global layer index +space: + + global layer idx is KDA -> kda_map[idx] = , + mla_map[idx] = -1 + global layer idx is MLA -> mla_map[idx] = , + kda_map[idx] = -1 + +Because the "other" manager is given ``-1`` for every layer it does not own, +asking the wrong manager to resolve a layer raises ``KeyError`` loudly (via +``coordinator_utils.resolve_from_layer_mapping``) instead of silently reading the +wrong physical slot — this is the miswiring guardrail. + +The coordinator exposes a single lifecycle / allocation / release surface that +keeps the two sub-managers in lock-step: + + * :meth:`initialize` / :meth:`shutdown` fan out to both managers. + * :meth:`allocate` reserves MLA pages **and** a KDA state slot for a batch of + sequences atomically, rolling back everything if either side fails. + * :meth:`release_sequence` frees the MLA pages **and** the KDA state slot (and + the KDA manager's per-sequence bookkeeping / block_reps row) in one call. + * routing accessors send a global layer to the correct sub-manager, letting a + miswired layer surface as ``KeyError``. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch + +from batchgen.kv_cache.gpu_paged_kv_manager import ( + GPUPagedKVCacheManager, + GPUPagedKVConfig, +) + +logger = logging.getLogger(__name__) + + +def build_kimi_linear_layer_maps( + config: Any, +) -> Tuple[List[int], List[int], List[bool], int, int]: + """Builds complementary MLA/KDA ``logical_to_physical_layer`` maps. + + Returns ``(mla_map, kda_map, layer_is_kda, num_mla_layers, num_kda_layers)`` + where ``mla_map``/``kda_map`` are length ``num_hidden_layers`` and use ``-1`` + for layers the corresponding manager does not own. + """ + num_layers = int(config.num_hidden_layers) + mla_map: List[int] = [] + kda_map: List[int] = [] + layer_is_kda: List[bool] = [] + mla_slot = 0 + kda_slot = 0 + for idx in range(num_layers): + if config.is_kda_layer(idx): + kda_map.append(kda_slot) + mla_map.append(-1) + layer_is_kda.append(True) + kda_slot += 1 + else: + mla_map.append(mla_slot) + kda_map.append(-1) + layer_is_kda.append(False) + mla_slot += 1 + return mla_map, kda_map, layer_is_kda, mla_slot, kda_slot + + +class KimiLinearGPUKVCoordinator: + """Composes an MLA paged-KV manager and a KDA state manager. + + See the module docstring for the layer-split rationale. Sequences are tracked + identically by both managers: every active sequence owns MLA pages (for the + MLA layers) *and* one KDA state slot (for the KDA layers). + + Args: + mla_manager: paged-KV manager covering the MLA layers. Its + ``logical_to_physical_layer`` must map KDA layers to ``-1``. + kda_manager: KDA state manager covering the KDA layers. Its + ``logical_to_physical_layer`` must map MLA layers to ``-1``. + layer_is_kda: optional per-global-layer boolean classification. If not + given it is derived from ``config`` (if provided) or from the KDA + manager's layer mapping. + config: optional ``KimiLinearConfig`` used to derive ``layer_is_kda`` and + for diagnostics. + """ + + def __init__( + self, + *, + mla_manager: GPUPagedKVCacheManager, + kda_manager: Any, + layer_is_kda: Optional[Sequence[bool]] = None, + config: Any = None, + ) -> None: + self.mla_manager = mla_manager + self.kda_manager = kda_manager + self.config = config + + if layer_is_kda is not None: + self._layer_is_kda = [bool(v) for v in layer_is_kda] + elif config is not None: + self._layer_is_kda = [ + bool(config.is_kda_layer(idx)) + for idx in range(int(config.num_hidden_layers)) + ] + else: + self._layer_is_kda = self._derive_layer_is_kda(kda_manager) + + self._active_sequences: set[int] = set() + + # ------------------------------------------------------------------ # + # Factory + # ------------------------------------------------------------------ # + @classmethod + def from_config( + cls, + config: Any, + *, + device: Any, + num_pages: int, + page_size_tokens: int, + num_state_items: int, + kv_dtype: torch.dtype = torch.bfloat16, + state_dtype: torch.dtype = torch.float32, + conv_dtype: torch.dtype = torch.bfloat16, + cuda_graph_max_pages_per_sequence: Optional[int] = None, + cuda_graph_max_slots: Optional[int] = None, + ) -> "KimiLinearGPUKVCoordinator": + """Builds both sub-managers with complementary layer maps. + + Requires ``KDAStateGPUManager``/``KDAStateGPUConfig`` to be importable from + ``batchgen.kv_cache.kda_state_gpu_manager``. + """ + try: + from batchgen.kv_cache.kda_state_gpu_manager import ( + KDAStateGPUConfig, + KDAStateGPUManager, + ) + except ImportError as exc: # pragma: no cover - depends on peer module + raise ImportError( + "KimiLinearGPUKVCoordinator.from_config requires " + "batchgen.kv_cache.kda_state_gpu_manager (KDAStateGPUManager, " + "KDAStateGPUConfig). Construct the coordinator directly with " + "pre-built managers if that module is unavailable." + ) from exc + + mla_map, kda_map, layer_is_kda, num_mla, num_kda = ( + build_kimi_linear_layer_maps(config) + ) + + compressed_kv_dim = int( + getattr(config, "compressed_kv_dim", None) + or (int(config.kv_lora_rank) + int(config.qk_rope_head_dim)) + ) + + mla_cfg = GPUPagedKVConfig( + num_layers=num_mla, + num_pages=int(num_pages), + page_size_tokens=int(page_size_tokens), + num_k_heads=1, + k_head_dim=compressed_kv_dim, + num_v_heads=0, + v_head_dim=0, + kv_dtype=kv_dtype, + cuda_graph_max_pages_per_sequence=cuda_graph_max_pages_per_sequence, + cuda_graph_max_slots=cuda_graph_max_slots, + logical_to_physical_layer=mla_map, + ) + mla_manager = GPUPagedKVCacheManager(config=mla_cfg, device=device) + + conv_dim = int(config.kda_num_heads) * int(config.kda_head_dim) + kda_cfg = KDAStateGPUConfig( + num_kda_layers=num_kda, + num_state_items=int(num_state_items), + num_heads=int(config.kda_num_heads), + head_dim=int(config.kda_head_dim), + conv_dim=conv_dim, + conv_width=int(config.kda_conv_size), + logical_to_physical_layer=kda_map, + ) + # KDAStateGPUManager mirrors the compressed-state manager constructor + # signature (config= + device=). + kda_manager = KDAStateGPUManager(config=kda_cfg, device=device) + + return cls( + mla_manager=mla_manager, + kda_manager=kda_manager, + layer_is_kda=layer_is_kda, + config=config, + ) + + # ------------------------------------------------------------------ # + # Lifecycle + # ------------------------------------------------------------------ # + def initialize(self, device: Any = None) -> Dict[str, Any]: + """Initializes both sub-managers (device is fixed at construction).""" + results: Dict[str, Any] = {} + results["mla"] = self._call_first(self.mla_manager, ("initialize",)) + results["kda"] = self._call_first(self.kda_manager, ("initialize",)) + logger.info( + "KimiLinearGPUKVCoordinator initialized (mla_layers=%d, kda_layers=%d)", + self.num_mla_layers, + self.num_kda_layers, + ) + return results + + def shutdown(self, *, empty_cuda_cache: bool = False) -> Dict[str, Any]: + """Tears down both sub-managers.""" + results: Dict[str, Any] = {} + results["kda"] = self._call_first( + self.kda_manager, + ("shutdown", "destroy"), + empty_cuda_cache=empty_cuda_cache, + ) + results["mla"] = self._call_first( + self.mla_manager, + ("destroy", "shutdown"), + empty_cuda_cache=empty_cuda_cache, + ) + self._active_sequences.clear() + return results + + # alias for callers that mirror the paged-manager API + def destroy(self, *, empty_cuda_cache: bool = False) -> Dict[str, Any]: + return self.shutdown(empty_cuda_cache=empty_cuda_cache) + + @property + def is_initialized(self) -> bool: + mla_ok = bool(getattr(self.mla_manager, "is_initialized", False)) + kda_ok = getattr(self.kda_manager, "is_initialized", None) + if kda_ok is None: + kda_ok = getattr(self.kda_manager, "_is_initialized", True) + return bool(mla_ok and kda_ok) + + # ------------------------------------------------------------------ # + # Allocation (atomic across both managers) + # ------------------------------------------------------------------ # + def allocate( + self, + sequence_ids: Sequence[int], + num_tokens: Sequence[int], + ) -> Dict[int, List[int]]: + """Allocates MLA pages *and* a KDA state slot for a batch of sequences. + + Both sides are reserved atomically: if the KDA side (or the MLA side + mid-batch) fails, every allocation performed in this call is rolled back + before the exception propagates. + + Returns the MLA page allocation dict ``{seq_id: [page, ...]}``. + """ + seq_ids = [int(s) for s in sequence_ids] + toks = [int(t) for t in num_tokens] + if len(seq_ids) != len(toks): + raise ValueError( + "allocate: sequence_ids and num_tokens must be the same length" + ) + if not seq_ids: + return {} + + pre_mla = set(self.mla_manager._sequences.keys()) + kda_done: List[int] = [] + try: + pages = self.mla_manager.allocate_pages_for_sequences(seq_ids, toks) + for seq_id in seq_ids: + self.kda_manager.allocate_state_item(seq_id) + kda_done.append(seq_id) + except Exception: + self._rollback_allocation(pre_mla, kda_done) + raise + + self._active_sequences.update(seq_ids) + return pages + + def _rollback_allocation( + self, pre_mla: set[int], kda_done: Sequence[int] + ) -> None: + if kda_done: + try: + self.kda_manager.release_sequence_states(list(kda_done)) + except Exception: # pragma: no cover - best-effort cleanup + logger.exception("KDA rollback failed during allocate()") + new_mla = [ + seq_id + for seq_id in self.mla_manager._sequences.keys() + if seq_id not in pre_mla + ] + if new_mla: + try: + self.mla_manager.free_pages_for_sequences(new_mla) + except Exception: # pragma: no cover - best-effort cleanup + logger.exception("MLA rollback failed during allocate()") + + def release_sequence(self, sequence_ids: Sequence[int] | int) -> None: + """Frees MLA pages and the KDA state slot for the given sequence(s). + + Releases the KDA per-sequence recurrent/conv state item (and the + manager's block_reps row) together with the MLA pages, keeping both + managers in lock-step. + """ + if isinstance(sequence_ids, int): + seq_ids = [int(sequence_ids)] + else: + seq_ids = [int(s) for s in sequence_ids] + if not seq_ids: + return + + # KDA release is idempotent-friendly (silently ignores unknown ids in the + # peer manager); the paged manager raises on unknown ids, so filter. + self.kda_manager.release_sequence_states(seq_ids) + known_mla = [s for s in seq_ids if s in self.mla_manager._sequences] + if known_mla: + self.mla_manager.free_pages_for_sequences(known_mla) + for seq_id in seq_ids: + self._active_sequences.discard(seq_id) + + # ------------------------------------------------------------------ # + # Layer routing + # ------------------------------------------------------------------ # + def is_kda_layer(self, layer_idx: int) -> bool: + return bool(self._layer_is_kda[int(layer_idx)]) + + def is_mla_layer(self, layer_idx: int) -> bool: + return not self.is_kda_layer(layer_idx) + + def manager_for_layer(self, layer_idx: int) -> Tuple[str, Any]: + """Returns ``("kda", kda_manager)`` or ``("mla", mla_manager)``.""" + if self.is_kda_layer(layer_idx): + return "kda", self.kda_manager + return "mla", self.mla_manager + + def resolve_mla_physical_layer(self, layer_idx: int) -> int: + """Resolves an MLA physical slot; raises ``KeyError`` for a KDA layer.""" + return int(self.mla_manager.resolve_physical_layer(int(layer_idx))) + + def resolve_kda_physical_layer(self, layer_idx: int) -> int: + """Resolves a KDA physical slot; raises ``KeyError`` for an MLA layer.""" + return int(self.kda_manager.resolve_physical_layer(int(layer_idx))) + + def get_mla_layer_kv_with_page_table(self, layer_idx: int): + """Routes an MLA layer to the paged manager. + + The paged manager resolves the layer through its ``logical_to_physical`` + map, so a KDA layer (mapped to ``-1``) raises ``KeyError``. + """ + return self.mla_manager.get_layer_kv_with_page_table(int(layer_idx)) + + def get_kda_layer_recurrent_view(self, layer_idx: int): + """Routes a KDA layer to the state manager (KeyError if it is an MLA layer).""" + if self.is_mla_layer(layer_idx): + raise KeyError( + f"get_kda_layer_recurrent_view: layer {layer_idx} is an MLA layer, " + "not served by the KDA state manager" + ) + return self.kda_manager.get_layer_recurrent_view(int(layer_idx)) + + def get_kda_layer_conv_views(self, layer_idx: int): + """Routes a KDA layer to the state manager (KeyError if it is an MLA layer).""" + if self.is_mla_layer(layer_idx): + raise KeyError( + f"get_kda_layer_conv_views: layer {layer_idx} is an MLA layer, " + "not served by the KDA state manager" + ) + return self.kda_manager.get_layer_conv_views(int(layer_idx)) + + def prepare_decode_step( + self, + sequence_ids: Sequence[int], + raw_positions: Sequence[int] | torch.Tensor, + ) -> None: + """Prepares KDA decode-step state bookkeeping for the batch.""" + if hasattr(self.kda_manager, "prepare_decode_step"): + self.kda_manager.prepare_decode_step(sequence_ids, raw_positions) + + # ------------------------------------------------------------------ # + # Introspection + # ------------------------------------------------------------------ # + @property + def num_layers(self) -> int: + return len(self._layer_is_kda) + + @property + def num_kda_layers(self) -> int: + return sum(self._layer_is_kda) + + @property + def num_mla_layers(self) -> int: + return len(self._layer_is_kda) - self.num_kda_layers + + @property + def active_sequence_ids(self) -> List[int]: + return sorted(self._active_sequences) + + # ------------------------------------------------------------------ # + # Helpers + # ------------------------------------------------------------------ # + @staticmethod + def _derive_layer_is_kda(kda_manager: Any) -> List[bool]: + mapping = getattr(kda_manager, "_logical_to_physical_layer", None) + if mapping is None: + cfg = getattr(kda_manager, "config", None) + mapping = getattr(cfg, "logical_to_physical_layer", None) + if mapping is None: + raise ValueError( + "Cannot derive layer classification: KDA manager exposes no " + "logical_to_physical_layer; pass layer_is_kda or config explicitly" + ) + return [int(v) >= 0 for v in mapping] + + @staticmethod + def _call_first(manager: Any, method_names: Sequence[str], **kwargs) -> Any: + for name in method_names: + method = getattr(manager, name, None) + if callable(method): + try: + return method(**kwargs) + except TypeError: + # method does not accept the passed kwargs (e.g. no + # empty_cuda_cache); retry without them. + return method() + return None + + +__all__ = [ + "KimiLinearGPUKVCoordinator", + "build_kimi_linear_layer_maps", +] diff --git a/batchgen/server/batch_scheduler.py b/batchgen/server/batch_scheduler.py index f2400285b..efd939ea6 100644 --- a/batchgen/server/batch_scheduler.py +++ b/batchgen/server/batch_scheduler.py @@ -223,9 +223,29 @@ async def _process_batch(self, batch_id: str) -> None: ) return - prompts, per_request_max_tokens, sampling_params = self._convert_requests_to_worker_inputs( - requests, batch - ) + try: + prompts, per_request_max_tokens, sampling_params = self._convert_requests_to_worker_inputs( + requests, batch + ) + except Exception as exc: + # Prompt construction can reject a request: an unknown chat role, a + # conversation the tokenizer cannot render faithfully, a malformed + # tool call. Without this the exception propagates to _run's bare + # `except Exception: logger.exception(...)`, which never sets a + # terminal status -- and the batch was already marked IN_PROGRESS + # above, so every request in it is lost and the client polls + # forever. + # + # NOTE this still fails the whole batch on one bad request. + # Per-request rejection at admission is the proper fix and is left + # as a follow-up; this only makes the failure visible and terminal. + logger.exception("Batch %s: prompt construction failed", batch_id) + self.storage.update_batch_status( + batch_id, + BatchStatus.FAILED, + error=f"Prompt construction failed: {type(exc).__name__}: {exc}", + ) + return # Apply batch-level max_decoding_length as fallback for requests without explicit value default_max = batch.max_decoding_length if default_max is None: @@ -403,9 +423,18 @@ def _convert_requests_to_worker_inputs( template_kwargs["tools"] = body.tools if body.preserve_thinking is not None: template_kwargs["preserve_thinking"] = body.preserve_thinking - prompt = self._format_chat_messages( - messages, body.model, **template_kwargs - ) + try: + prompt = self._format_chat_messages( + messages, body.model, **template_kwargs + ) + except Exception as exc: + # Attach the custom_id. The caller fails the batch; without + # the id there is no way to tell which of N requests did it. + custom_id = request.custom_id or "" + raise ValueError( + f"request {custom_id!r}: the chat template rejected " + f"this conversation: {type(exc).__name__}: {exc}" + ) from exc # Priority: max_completion_tokens > max_tokens > None current_max_tokens = body.max_completion_tokens if body.max_completion_tokens is not None else body.max_tokens elif isinstance(body, CompletionRequest): diff --git a/batchgen/server/http_server.py b/batchgen/server/http_server.py index 93b969eb6..779a43041 100644 --- a/batchgen/server/http_server.py +++ b/batchgen/server/http_server.py @@ -38,9 +38,11 @@ ListFilesResponse, ListModelsResponse, ModelObject, - RawInferenceRequest, build_batch_object_from_create_request, - normalize_inference_results, +) +from batchgen.deprecation import ( + LEGACY_INFERENCE_ERROR_CODE, + LEGACY_INFERENCE_MESSAGE, ) from batchgen.server.health import ServerHealthState from batchgen.server.server_args import ServerArgs @@ -418,77 +420,29 @@ async def cancel_batch(request: Request, batch_id: str): return updated @app.post("/v1/inference") - async def run_inference(request: Request, body: RawInferenceRequest): - worker: WorkerManager = request.app.state.worker - server_args: ServerArgs = request.app.state.server_args - storage: StorageManager = request.app.state.storage - - max_input_len = body.max_input_len # None = dynamically determined - max_output_len = body.max_output_len or 128 # Default max output tokens - - start = time.perf_counter() - try: - results = await asyncio.to_thread( - worker.infer, - body.prompts, - max_input_len, - max_output_len, - body.ignore_eos, - body.temperature, # None = greedy decoding - body.top_p, # None = disabled - ) - except Exception as exc: - logger.exception("Inference failed") - raise HTTPException(status_code=500, detail=str(exc)) - - latency_ms = int((time.perf_counter() - start) * 1000) - # Worker returns dict {global_idx: str} — convert to ordered list - if isinstance(results, dict): - results = [results[k] for k in sorted(results.keys())] - normalized_results = normalize_inference_results(results) - - response_data = { - "status": "success", - "results": normalized_results, - "latency_ms": latency_ms, - } - - # Save results to file if save_result is enabled - if server_args.save_result: - output_file_id = f"file-{uuid.uuid4().hex}" - output_path = storage.output_dir / f"{output_file_id}.jsonl" - output_path.parent.mkdir(parents=True, exist_ok=True) - - with output_path.open("w", encoding="utf-8") as f: - for idx, result in enumerate(normalized_results): - record = { - "custom_id": f"inference-{idx}", - "prompt": body.prompts[idx] if idx < len(body.prompts) else "", - "response": result, - } - f.write(json.dumps(record, ensure_ascii=False) + "\n") - - # Save file metadata - file_meta = FileObject( - id=output_file_id, - bytes=output_path.stat().st_size, - created_at=int(time.time()), - filename=f"{output_file_id}.jsonl", - purpose=FilePurpose.BATCH_OUTPUT.value, - status=FileStatus.PROCESSED.value, - status_details=None, - checksum=None, - ) - storage.save_metadata(output_file_id, file_meta.dict()) - - # Also copy to files_dir for download via /v1/files/{id}/content - import shutil - shutil.copy(output_path, storage.files_dir / output_file_id) - - response_data["output_file_id"] = output_file_id - logger.info(f"Saved inference results to {output_path}") - - return response_data + async def run_inference(): + """DEPRECATED and disabled. All inference goes through the batch API. + + Rejected here, at the door, before any queue interaction. The old body's + first act was to put a payload on the shared worker request queue, and + in pool mode nothing downstream could undo that: the worker admission + loop recognises no legacy message, so the payload was dropped there + while the caller blocked on the shared response queue and took the next + batch's completion. + + The route is kept rather than deleted so callers get this explanation + instead of a 404 they would read as a typo. It takes no request body: + an unparseable legacy payload must still get the deprecation notice, + not a 422 about its schema. + """ + raise HTTPException( + status_code=410, + detail={ + "code": LEGACY_INFERENCE_ERROR_CODE, + "message": LEGACY_INFERENCE_MESSAGE, + "use_instead": "/v1/batches", + }, + ) @app.post("/v1/reload") async def reload_worker(request: Request): diff --git a/batchgen/server/worker_manager.py b/batchgen/server/worker_manager.py index 313c62f6e..a170d186a 100644 --- a/batchgen/server/worker_manager.py +++ b/batchgen/server/worker_manager.py @@ -803,6 +803,18 @@ def _load_model_locally( self.args.enable_hugetlbfs, enable_memfd=self.args.fast_init, ) + elif "kimi-linear" in self.args.model.lower() or "kimi-k3" in self.args.model.lower(): + from batchgen.models.moonshotai.kimi_linear.kimi_parameter_server import ( + KimiLinear_Parameter_Server, + ) + + parameter_server = KimiLinear_Parameter_Server( + self.args.model, + self.args.cache_dir, + converted_ckpt_dir, + self.args.enable_hugetlbfs, + enable_memfd=self.args.fast_init, + ) elif is_kimi_k25_backend_model(self.args.model): from batchgen.models.moonshotai.kimi_k25.kimi_parameter_server import ( KimiK25_Parameter_Server, diff --git a/core/HtoD_Engine/HtoD_Engine.cu b/core/HtoD_Engine/HtoD_Engine.cu index 327b9f790..a3ae364d4 100644 --- a/core/HtoD_Engine/HtoD_Engine.cu +++ b/core/HtoD_Engine/HtoD_Engine.cu @@ -439,16 +439,42 @@ void HtoD_Engine::HtoD_Worker() { for (auto& [tensor_name, host_tensor_storage] : src) { src_ptr = host_tensor_storage.data_ptr; src_byte_size = host_tensor_storage.byte_size; - if (dst[tensor_name].defined() && - dst[tensor_name].has_storage()) { - this->blocking_copy_(dst[tensor_name].data_ptr(), - src_ptr, src_byte_size); - } else { + // find(), not operator[]: dst is an unordered_map and + // operator[] DEFAULT-INSERTS an undefined torch::Tensor on + // every miss, permanently growing a buffer map that is + // reused across ring slots. + auto slot = dst.find(tensor_name); + if (slot == dst.end() || !slot->second.defined() || + !slot->second.has_storage()) { + // The host map carries a tensor module_shapes declares + // no slot for. Continuing drops it silently and the + // consumer reads whatever the slot last held. this->logger_->error( - "Tensor {} doesn't have valid storage", + "Module {}: host tensor {} has no GPU slot -- " + "module_shapes declares no such key", + module_name, tensor_name); + throw std::runtime_error( + "HtoD: host tensor has no GPU slot: " + tensor_name); - std::runtime_error("Tensor doesn't have valid storage"); } + int64_t dst_byte_size = slot->second.nbytes(); + if (src_byte_size != dst_byte_size) { + // blocking_copy_ writes src_byte_size bytes with no + // bound check: a short slot is overrun into its + // neighbour, a long one keeps a stale tail. Both are + // silent and both produce wrong weights. + this->logger_->error( + "Module {}: tensor {} size mismatch -- host {} B, " + "GPU slot {} B (module_shapes/dtype disagrees with " + "the checkpoint)", + module_name, tensor_name, src_byte_size, + dst_byte_size); + throw std::runtime_error( + "HtoD: host/GPU byte size mismatch for " + + tensor_name); + } + this->blocking_copy_(slot->second.data_ptr(), src_ptr, + src_byte_size); } this->logger_->debug("Copied module: {} to buffer: {}", module_name, buffer_idx); diff --git a/tests/unit/test_query_book_pool_grow_rebind.py b/tests/unit/test_query_book_pool_grow_rebind.py new file mode 100644 index 000000000..f484b0c44 --- /dev/null +++ b/tests/unit/test_query_book_pool_grow_rebind.py @@ -0,0 +1,272 @@ +"""CPU unit test for the QueryBook buffer-pool GROW + REBIND-with-live-sequences path. + +This exercises the branch that GPU runs have never hit. Wave admission serialises +batches, so a pool grow has never happened while a live, mid-decode sequence still +occupies a row -- every observed grow rebound 0 sequences. That copy-live-rows / +re-point-views branch is therefore untested code that could corrupt on a real +concurrent grow. Here we force exactly that and assert the live row survives the +grow byte-for-byte. + +The code under test is the REAL shipping source of + - ``QueryBookBufferPool`` (including ``.adopt`` -- the live-row copy) + - ``BatchGenWorker._ensure_buffer_pool`` (the grow orchestration) + - ``BatchGenWorker._rebind_buffer_pool_views``(re-point seq + query_book views) + - ``BatchGenWorker._retire_buffer_pool`` +extracted verbatim from ``batchgen/batchgen_worker.py`` by AST and exec'd against a +fake worker ``self``. We extract instead of importing because importing +``batchgen.batchgen_worker`` pulls in the JIT-compiled ``core_engine`` (see its +module-level ``from batchgen.models.engine_loader import core_engine``), which is +not available on a CPU box. Extraction keeps the test bound to the real source: +any edit to adopt/_ensure_buffer_pool/_rebind flows straight into these asserts. + +The node-shared input_ids segment is allocated through the real +``allocate_node_shared_int64`` (POSIX /dev/shm) with a no-op barrier, exactly as its +own docstring prescribes for tests. +""" + +import ast +import logging +import os +import textwrap +from types import SimpleNamespace +from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple + +import torch + +import batchgen +from batchgen.query_book import QueryBookEntry +from batchgen.sequence import SequenceEntry + +WORKER_PATH = os.path.join(os.path.dirname(batchgen.__file__), "batchgen_worker.py") + +# The single line inside QueryBookBufferPool.adopt that copies live input_ids rows +# from the superseded pool into the grown one. The mutation test neuters exactly +# this line to prove the positive assertions are load-bearing. +COPY_LINE = "self.input_ids_buffer[:rows, :cols] = old.input_ids_buffer[:rows, :cols]" + + +def _extract_segments(): + """Pull the exact source text of the symbols under test from the real file.""" + src = open(WORKER_PATH).read() + tree = ast.parse(src) + lines = src.splitlines(keepends=True) + + def grab(node): + # Slice full physical lines (they keep their leading tabs) then dedent, so a + # tab-indented method becomes a top-level function. + return textwrap.dedent("".join(lines[node.lineno - 1 : node.end_lineno])) + + wanted_top = {"allocate_node_shared_int64", "QueryBookPoolCapacityError", "QueryBookBufferPool"} + wanted_methods = { + "_node_shared_tag", + "_ensure_buffer_pool", + "_rebind_buffer_pool_views", + "_retire_buffer_pool", + } + seg = {} + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in wanted_top: + seg[node.name] = grab(node) + elif isinstance(node, ast.ClassDef) and node.name == "BatchGenWorker": + for m in node.body: + if isinstance(m, ast.FunctionDef) and m.name in wanted_methods: + seg[m.name] = grab(m) + missing = (wanted_top | wanted_methods) - set(seg) + assert not missing, f"failed to extract from real source: {missing}" + return seg + + +# Definitions must land in this order: exception -> shm helper -> pool -> methods +# (_retire_buffer_pool annotates a param with QueryBookBufferPool, so the class must +# already exist when its def executes). +_EXEC_ORDER = [ + "QueryBookPoolCapacityError", + "allocate_node_shared_int64", + "QueryBookBufferPool", + "_node_shared_tag", + "_ensure_buffer_pool", + "_rebind_buffer_pool_views", + "_retire_buffer_pool", +] + + +def _build_worker(mutate_adopt=False): + """Return (FakeWorker class, QueryBookBufferPool, QueryBookPoolCapacityError). + + All functions share one globals dict ``g`` so their cross-references resolve. + ``mutate_adopt`` neuters the live-row input_ids copy inside ``adopt``. + """ + seg = _extract_segments() + g = { + "torch": torch, + "os": os, + "logging": logging, + "dist": SimpleNamespace(barrier=lambda *a, **k: None), + "NUM_GPUS_PER_NODE": 8, + "Tuple": Tuple, + "Optional": Optional, + "Dict": Dict, + "List": List, + "Callable": Callable, + "Sequence": Sequence, + "Set": Set, + } + for name in _EXEC_ORDER: + code = seg[name] + if name == "QueryBookBufferPool" and mutate_adopt: + assert COPY_LINE in code, "adopt row-copy line not found -- source drifted" + code = code.replace(COPY_LINE, "pass # MUTATION: live-row input_ids copy skipped") + exec(compile(code, WORKER_PATH, "exec"), g) + + method_names = ( + "_node_shared_tag", + "_ensure_buffer_pool", + "_rebind_buffer_pool_views", + "_retire_buffer_pool", + ) + FakeWorker = type("FakeWorker", (), {m: g[m] for m in method_names}) + return FakeWorker, g["QueryBookBufferPool"], g["QueryBookPoolCapacityError"] + + +def _setup_live_pool(Pool): + """A small (2 x 8) pool with row 0 occupied by a live, mid-decode sequence.""" + pool = Pool(num_sequences=2, input_ids_width=8, max_decoding_length=4, pad_token_id=0) + slot = pool.allocate_slot() # -> 0 + seq = SequenceEntry("seq-live", global_idx=0, prompt_length=5, max_decode_length=3, text="live") + # kv_token_budget = 5 + 3 = 8 == input_ids_width + prompt = torch.tensor([11, 12, 13, 14, 15], dtype=torch.long) + iv = pool.get_input_ids_view(slot, seq.kv_token_budget) # (1, 8) + iv[0, :5] = prompt + dv = pool.get_decoded_tokens_view(slot) # (1, 4) + dv[0, :2] = torch.tensor([901, 902], dtype=torch.int64) + seq.decoded_length = 2 + seq._buffer_slot = slot + seq.input_ids = iv + seq.decoded_tokens = dv + entry = QueryBookEntry( + encoded={"input_ids": iv}, decoded_tokens=dv, kv_token_budget=seq.kv_token_budget + ) + return pool, seq, entry, prompt + + +def _make_worker(FakeWorker, pool, seq, entry): + w = FakeWorker() + w._buffer_pool = pool + w._buffer_pool_generation = 1 + w.rank = 0 + w.pad_token_id = 0 + w._retired_buffer_pools = [] + # short unique tag -> POSIX shm name stays under the macOS 31-char limit + w._shared_buffer_tag = os.urandom(2).hex() + w.global_batch = [seq] + w._uuid_to_local_map = {seq.uuid: 0} + w.query_book = {0: entry} + return w + + +def _cleanup(pool): + shm = getattr(pool, "input_ids_shm", None) + if shm is not None: + try: + shm.close() + except Exception: + pass + try: + shm.unlink() + except Exception: + pass + + +def test_extraction_covers_real_source(): + """Guard: the harness really pulled the grow/rebind code, not empty stubs.""" + seg = _extract_segments() + assert COPY_LINE in seg["QueryBookBufferPool"] + assert "def adopt(self" in seg["QueryBookBufferPool"] + assert "new_pool.adopt(old)" in seg["_ensure_buffer_pool"] + assert "self._rebind_buffer_pool_views()" in seg["_ensure_buffer_pool"] + assert "rebound += 1" in seg["_rebind_buffer_pool_views"] + + +def test_grow_copies_live_row_rebinds_and_admits_second(): + FakeWorker, Pool, _ = _build_worker() + pool, seq, entry, prompt = _setup_live_pool(Pool) + + old_buf_ptr = pool.input_ids_buffer.data_ptr() + prompt_snap = seq.input_ids[0, :5].clone() + dec_snap = seq.decoded_tokens[0, :2].clone() + + w = _make_worker(FakeWorker, pool, seq, entry) + try: + # Force a grow: width 8 -> 16 (also rows 2 -> 4, decode 4 -> 8). + w._ensure_buffer_pool( + required_rows=4, + required_input_width=16, + required_decode_width=8, + reason="unit test forced grow with a live sequence", + ) + grown = w._buffer_pool + + # A real grow occurred into a fresh, larger allocation. + assert grown is not pool + assert (grown.num_sequences, grown.input_ids_width, grown.max_decoding_length) == (4, 16, 8) + assert grown.input_ids_buffer.data_ptr() != old_buf_ptr + + # (1) live row copied byte-identically into the grown buffer + assert torch.equal(grown.input_ids_buffer[0, :5], prompt_snap) + assert grown.input_ids_buffer[0, 5:].sum() == 0 # remainder is padding + assert torch.equal(grown.decoded_tokens_buffer[0, :2], dec_snap) + + # (2) slot mapping intact; seq + query_book rebound onto the grown buffer + assert seq._buffer_slot == 0 + assert seq.input_ids.shape == (1, seq.kv_token_budget) # (1, 8) + assert torch.equal(seq.input_ids[0, :5], prompt_snap) + # the rebound view actually aliases the grown buffer (not a stale mapping) + grown.input_ids_buffer[0, 7] = 4242 + assert seq.input_ids[0, 7] == 4242 + grown.input_ids_buffer[0, 7] = 0 + # query_book entry rebound to the SAME grown view object as the sequence + assert entry.encoded["input_ids"].data_ptr() == seq.input_ids.data_ptr() + assert torch.equal(entry.decoded_tokens[0, :2], dec_snap) + + # (3) a second sequence admits into the grown pool without clobbering the first + slot2 = grown.allocate_slot() + assert slot2 == 1 # _next_slot carried over from the old pool (1 row used) + iv2 = grown.get_input_ids_view(slot2, 16) + iv2[0, :10] = torch.full((10,), 777, dtype=torch.long) + assert torch.equal(grown.input_ids_buffer[0, :5], prompt_snap) # row 0 untouched + assert torch.equal(seq.input_ids[0, :5], prompt_snap) + finally: + _cleanup(w._buffer_pool) + + +def test_mutation_row_copy_removed_is_detected(): + """Falsifiability: with the live-row copy skipped, the survival assertion fails. + + Proves test_grow_copies_live_row_rebinds_and_admits_second is not vacuous. + """ + FakeWorker, Pool, _ = _build_worker(mutate_adopt=True) + pool, seq, entry, prompt = _setup_live_pool(Pool) + prompt_snap = seq.input_ids[0, :5].clone() + + w = _make_worker(FakeWorker, pool, seq, entry) + try: + w._ensure_buffer_pool( + required_rows=4, + required_input_width=16, + required_decode_width=8, + reason="unit test forced grow (mutated adopt)", + ) + grown = w._buffer_pool + # The live row was NOT carried over -> grown row 0 is the zero-filled segment. + assert not torch.equal(grown.input_ids_buffer[0, :5], prompt_snap) + assert grown.input_ids_buffer[0, :5].sum() == 0 + # and the rebound live view now reads zeros -- the corruption the copy prevents. + assert seq.input_ids[0, :5].sum() == 0 + finally: + _cleanup(w._buffer_pool) + + +if __name__ == "__main__": + import pytest + + raise SystemExit(pytest.main([__file__, "-v", "-s"]))