From 00198a3d8a3d9692e4bf118fc5223b6d31814d69 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 22:53:20 +0800 Subject: [PATCH 01/19] feat(core): core-side wiring for kimi-linear serving Worker carry-loop, KDA state GPU manager and kimi-linear KV coordinator under kv_cache/, host-KV profile entry, model/tokenizer registry rows, and the worker_manager seam. Core half of the kimi-linear bring-up; the model-side stack (models/moonshotai/kimi_linear/**) lands in the companion model PR. --- batchgen/batchgen_worker.py | 45 +- batchgen/config/model_registry.py | 33 +- batchgen/config/tokenizer_registry.py | 9 + batchgen/kv_cache/host_kv_mananger_config.py | 21 + batchgen/kv_cache/kda_state_gpu_manager.py | 344 ++++++++++++++ .../kv_cache/kimi_linear_kv_coordinator.py | 448 ++++++++++++++++++ batchgen/server/worker_manager.py | 12 + 7 files changed, 898 insertions(+), 14 deletions(-) create mode 100644 batchgen/kv_cache/kda_state_gpu_manager.py create mode 100644 batchgen/kv_cache/kimi_linear_kv_coordinator.py diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 92cf18bb5..2d913b94f 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -2006,11 +2006,19 @@ def _put_sequences_onhold(self, uuids: List[str]) -> None: local_indices = self._get_local_indices_for_uuids(my_uuids) global_ids = self._local_indices_to_global_seq_ids(local_indices) - manager = self.gpu_paged_kv_cache_manager - if manager is not None: - manager.free_pages_for_sequences(global_ids) - - for uuid in my_uuids: + manager = self.gpu_paged_kv_cache_manager + if manager is not None: + manager.free_pages_for_sequences(global_ids) + + # Kimi-Linear: release KDA state slots alongside GPU KV pages. + try: + from batchgen.models.moonshotai.kimi_linear.wrappers import KimiLinearKDAWrapper + if KimiLinearKDAWrapper.slot_manager is not None: + KimiLinearKDAWrapper.free_sequences(global_ids) + except ImportError: + pass + + for uuid in my_uuids: seq = self.global_batch.get_sequence(uuid) seq.gpu_pages_allocated = 0 self._sequences_with_gpu_kv.discard(uuid) @@ -3026,6 +3034,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) @@ -5998,6 +6015,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") @@ -6974,7 +6995,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 @@ -11816,6 +11840,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/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..dc643360d 100644 --- a/batchgen/config/tokenizer_registry.py +++ b/batchgen/config/tokenizer_registry.py @@ -85,6 +85,10 @@ "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": "kimi_linear", + "kimi-k3": "kimi_linear", } for model_id in KIMI_K25_BACKEND_MODEL_IDS: @@ -189,6 +193,11 @@ def _import_tokenizers(): except ImportError: pass + try: + from batchgen.models.moonshotai.kimi_linear import tokenizer as _ # noqa: F401 + except ImportError: + pass + try: from batchgen.models.glm.glm5 import tokenizer as _ # noqa: F401 except ImportError: diff --git a/batchgen/kv_cache/host_kv_mananger_config.py b/batchgen/kv_cache/host_kv_mananger_config.py index a8ec16aef..ddbf91647 100644 --- a/batchgen/kv_cache/host_kv_mananger_config.py +++ b/batchgen/kv_cache/host_kv_mananger_config.py @@ -145,6 +145,17 @@ def bytes_per_page(self) -> int: 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 +165,7 @@ 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, } _PROFILE_ALIASES: Dict[str, str] = {} @@ -197,6 +209,15 @@ 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", + "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..d7678a44e --- /dev/null +++ b/batchgen/kv_cache/kda_state_gpu_manager.py @@ -0,0 +1,344 @@ +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. + """ + + 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 + 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-recycle. + """ + + 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 <= 0: + raise ValueError("conv_width must be > 0") + 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, + ) + conv_shape = ( + cfg.num_kda_layers, + cfg.num_state_items, + conv_dim, + cfg.conv_width, + ) + 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 state for recycled 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) + + # ------------------------------------------------------------------ + # 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-state views (q, k, v), each [num_state_items, conv_dim, width].""" + 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 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()) + 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/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, From a52e95cc56d04f7a208086a1a0e75fce7bd389f6 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 22:53:49 +0800 Subject: [PATCH 02/19] fix(worker): M1 core-side fixes for kimi-linear (F5, decode heartbeat) Worker half of the kimi-linear M1 fix batch: F5 worker-side sequencing and the decode heartbeat. The planner-override half (F6) rides in the companion model PR. --- batchgen/batchgen_worker.py | 36 +++++++++++++++++++++++------------- 1 file changed, 23 insertions(+), 13 deletions(-) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 2d913b94f..bd565881d 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -2006,19 +2006,11 @@ def _put_sequences_onhold(self, uuids: List[str]) -> None: local_indices = self._get_local_indices_for_uuids(my_uuids) global_ids = self._local_indices_to_global_seq_ids(local_indices) - manager = self.gpu_paged_kv_cache_manager - if manager is not None: - manager.free_pages_for_sequences(global_ids) - - # Kimi-Linear: release KDA state slots alongside GPU KV pages. - try: - from batchgen.models.moonshotai.kimi_linear.wrappers import KimiLinearKDAWrapper - if KimiLinearKDAWrapper.slot_manager is not None: - KimiLinearKDAWrapper.free_sequences(global_ids) - except ImportError: - pass - - for uuid in my_uuids: + manager = self.gpu_paged_kv_cache_manager + if manager is not None: + manager.free_pages_for_sequences(global_ids) + + for uuid in my_uuids: seq = self.global_batch.get_sequence(uuid) seq.gpu_pages_allocated = 0 self._sequences_with_gpu_kv.discard(uuid) @@ -9566,6 +9558,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: @@ -9576,6 +9572,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 From f524e91c863a039b62d0b7aca9338b8e04bfe9ff Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 30 Jul 2026 16:36:52 +0100 Subject: [PATCH 03/19] fix(worker): stale completion metadata and slot reuse on decode drain _report_completion read prompt/decoded_length from rank-0's local replica, stale for non-rank-0-owned sequences on the outer-loop drain path (only the completion bit is all-reduced there) -> short usage.completion_tokens now, routinely wrong at 128-seq L2. Sync metadata before reporting. Also reset seq._buffer_slot after free_slot so a re-entered report cannot free a slot already owned by another sequence. --- batchgen/batchgen_worker.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index bd565881d..a393fba2e 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -1484,6 +1484,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 @@ -5660,6 +5663,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) From 72957f1f302655c81649b976216dda8714969c52 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 22:54:25 +0800 Subject: [PATCH 04/19] feat(kv_cache): manager-authoritative graph-ready KDA state pools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KDAStateGPUManager owns allocation, slot accounting, and F4 zero-on-alloc; conv pools migrate to the causal_conv1d.cu contract ((L, slots, dim, W-1), contiguous per-layer 3-D views), and each pool is a single fixed-address allocation — the CUDA-graph capture requirement. The wrapper-side views and the graph-readiness tests ride in the companion model PR. --- batchgen/kv_cache/kda_state_gpu_manager.py | 49 ++++++++++++++++++---- 1 file changed, 42 insertions(+), 7 deletions(-) diff --git a/batchgen/kv_cache/kda_state_gpu_manager.py b/batchgen/kv_cache/kda_state_gpu_manager.py index d7678a44e..712a3f7e6 100644 --- a/batchgen/kv_cache/kda_state_gpu_manager.py +++ b/batchgen/kv_cache/kda_state_gpu_manager.py @@ -22,6 +22,10 @@ class KDAStateGPUConfig: 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 @@ -29,7 +33,7 @@ class KDAStateGPUConfig: 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 + 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 @@ -55,7 +59,15 @@ class KDAStateGPUManager: 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-recycle. + 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" @@ -76,8 +88,10 @@ def __init__( 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 <= 0: - raise ValueError("conv_width 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( @@ -107,11 +121,16 @@ def initialize(self) -> None: 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, + cfg.conv_width - 1, ) self._conv_q = torch.zeros( conv_shape, dtype=cfg.conv_dtype, device=self.device @@ -160,7 +179,7 @@ def release_sequence_states(self, sequence_ids: Sequence[int]) -> None: self._free_state_items.push(reclaimed) def reset_state_items(self, state_item_ids: Sequence[int]) -> None: - """Zero recurrent + conv state for recycled slots.""" + """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: @@ -174,6 +193,11 @@ def reset_state_items(self, state_item_ids: Sequence[int]) -> None: 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 @@ -187,7 +211,7 @@ def get_layer_recurrent_view(self, logical_layer: int) -> torch.Tensor: def get_layer_conv_views( self, logical_layer: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Conv-state views (q, k, v), each [num_state_items, conv_dim, width].""" + """Conv views (q, k, v), each [num_state_items, conv_dim, width-1].""" self._ensure_initialized() physical = self.resolve_physical_layer(logical_layer) return ( @@ -237,6 +261,10 @@ 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 @@ -299,6 +327,13 @@ def _ensure_state_item(self, sequence_id: int) -> int: 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 From 87aea662e6934a112230b773c2c206a1998fab96 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Sat, 1 Aug 2026 13:31:21 +0100 Subject: [PATCH 05/19] fix(worker): log the full sequence id in REPETITION warnings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit uuid[:8] is unique for random-hex ids but /v1/batches reuses structured custom_ids (mmlu-{run}-{idx}) whose first 8 chars are identical across a whole run — the warning named 'mmlu-030' for every sequence, hiding which one actually looped and costing real debugging time. The remaining 58 uuid[:8] log sites are queued for the pre-PR sweep. --- batchgen/batchgen_worker.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index a393fba2e..407305098 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -4821,7 +4821,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}" ) @@ -10506,7 +10506,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: @@ -10520,7 +10520,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}" ) From 96436cbef37766b52eec2021b43059755ab6f24d Mon Sep 17 00:00:00 2001 From: tairanxu Date: Tue, 4 Aug 2026 22:58:18 +0100 Subject: [PATCH 06/19] fix(core): route kimi-k3 tokenizer; fail prompt construction loudly Registry: kimi-k3 / moonshotai/Kimi-K3 now resolve to KimiK3Tokenizer, and _import_tokenizers warns per-module on import failure instead of one module's error aborting the loop. kimi_linear no longer owns the kimi_k3 name. Scheduler: prompt construction (chat rendering, tool formatting) can raise; previously the exception reached _run's bare logger.exception with the batch already marked IN_PROGRESS, so no terminal status was ever set and the client polled forever. The batch now fails LOUDLY with the exception in the error field. Deliberately still batch-granular -- per-request rejection at admission is a named follow-up; this change only makes the failure visible and terminal, per the no-silent-fallback rule. Core PR per PR_MERGE_POLICY (batchgen/config/ and batchgen/server/ are outside MODEL_ALLOW_RE). --- batchgen/config/tokenizer_registry.py | 99 +++++++++++++-------------- batchgen/server/batch_scheduler.py | 41 +++++++++-- 2 files changed, 82 insertions(+), 58 deletions(-) diff --git a/batchgen/config/tokenizer_registry.py b/batchgen/config/tokenizer_registry.py index dc643360d..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 @@ -87,8 +88,15 @@ "MiniMaxAI/MiniMax-M2.5": "minimax_m25", "Kimi-Linear": "kimi_linear", "kimi-linear": "kimi_linear", - "Kimi-K3": "kimi_linear", - "kimi-k3": "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: @@ -142,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}. " @@ -162,56 +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.moonshotai.kimi_linear 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/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): From 6e3cf462eae18c76b07b5837669a78b96d3493b7 Mon Sep 17 00:00:00 2001 From: tairanxu Date: Tue, 4 Aug 2026 23:22:57 +0100 Subject: [PATCH 07/19] fix(htod): throw on missing slot and size mismatch in the copy loop Three guards on the host-to-device weight copy, replacing an error path that constructed a std::runtime_error and discarded it before falling through to weights_copy_complete: - dst.find() instead of operator[], which default-inserts an undefined torch::Tensor on every miss into a buffer map reused across ring slots. - A host tensor with no GPU slot now throws. Continuing dropped it silently and the consumer read whatever the slot last held -- the GLM-5 Q/K RMSNorm incident class (glm5_initializer.py:141-146). - src/dst byte-size equality before the copy. 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 silent. Latent today (48B name maps agree, no occurrences in current logs), but K3 lines up 497,220 checkpoint tensors against a freshly written module_shapes -- exactly the condition these guards exist for. Per the 2026-08-04 ledger this lands as its own fix PR before the M3 force-stream rehearsal, whose negative tests (byte-size mismatch, omitted tensor) must fail loudly against THIS code and hang/corrupt without it. Compile verification is staged for the GPU machine; this translation unit builds only there. --- core/HtoD_Engine/HtoD_Engine.cu | 40 +++++++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 7 deletions(-) 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); From d6f42a9069a2cd486c7d54224d382d7560802b3c Mon Sep 17 00:00:00 2001 From: tairanxu Date: Tue, 4 Aug 2026 23:23:13 +0100 Subject: [PATCH 08/19] fix: ckpt_converter refuses MXFP4 tensors in the INT4 marlin repack The repack reinterprets the uint8 buffer as packed uint4b8 and converts the scale tensor VALUE to bfloat16 -- an E8M0 exponent byte of 121 would become 121.0 instead of 2**(121-127). Both corruptions are silent, and convert_checkpoint.py:144 passes marlin=True by default, so converting the K3 checkpoint today would quietly produce garbage marlin weights. Scoped to tensors the repack actually matches (uint8 scale = E8M0 = MXFP4): an unscoped check would reject checkpoints the repack never touches. The MXFP4-aware repack is the in-flight kernel work unit; until it lands, K3 conversion must run --no-marlin. --- batchgen/ckpt_converter/ckpt_converter.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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] From cd8f1c4f661e4442d81fad4e38eb414c516dfc6a Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 22:55:40 +0800 Subject: [PATCH 09/19] feat(worker): carry K3 block residuals through the serving prefill loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Worker half of the block-residual carry (the model half, kimi_linear model.py's _apply_output_attn_res, rides in the companion model PR). K3 replaces the classic residual body with Block Attention Residuals: a depth-mix across block boundaries whose state is carried layer to layer. The prepacked prefill loop did `hidden_states = layer_outputs[0]`, which never supplied that state and discarded the second return value. The model would have loaded, streamed and produced logits while every layer saw a zero-width residual — wrong text, no error, nothing to grep for. That is the exact silent-wrong-output class this project bans, and it is why the first real 8K run had not been attempted. The loop now mirrors KimiLinearModel.forward, the eager reference: initialise the (N, 0, hidden) residual, thread it through layers that ask for it, and apply the output depth-mix BEFORE the final norm — mix-then-norm is the reference order and swapping it changes every output. Models without attn_res_block_size keep the exact previous code path, so kimi-linear-48B, GLM-5 and K2.5 are untouched. The worker calls the model's _apply_output_attn_res method instead of importing a model-private function, so core stays model-agnostic and there is exactly one implementation of the mix — the eager and serving paths cannot drift. --- batchgen/batchgen_worker.py | 51 ++++++++++++++++++++++++++++++------- 1 file changed, 42 insertions(+), 9 deletions(-) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 407305098..02af524da 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -7024,16 +7024,49 @@ def prefill_prepacked(self, batch: list[int]): # Reshape to 3D: [1, batch_total_tokens, hidden_dim] hidden_states = inputs_embeds.unsqueeze(0) + # 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: + block_residual = hidden_states.new_zeros( + hidden_states.shape[0] * hidden_states.shape[1], 0, + hidden_states.shape[2]) + 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) From 61e43dc802f0815eacd49eff6b229f9cf9075143 Mon Sep 17 00:00:00 2001 From: tairanxu Date: Thu, 6 Aug 2026 20:23:01 +0100 Subject: [PATCH 10/19] fix(kv_cache): give kimi-k3 its own host-KV profile K3 was aliased onto _KIMI_LINEAR_MLA_PROFILE, which declares num_layers=27 because that is the 48B's depth. K3 has 93 engine layers with 24 MLA layers at engine indices 3,7,...,87,91,92, and wrappers.py::_offload_prepacked_kv indexes the pool by the ENGINE layer index rather than by a dense MLA counter. Every MLA layer at index >= 27 therefore wrote past the end of the pool -- silent host-memory corruption starting at layer 28 of the very first prefill, with no error and nothing to grep for. The geometry is otherwise identical (compressed_kv_dim 576, MQA single head, bf16), so this is purely the layer count. K3 now resolves to its own kimi_k3_mla key; the kimi-linear aliases are unchanged, so the 48B keeps the 27-layer profile it wants. Found by reading the load path while diagnosing an unrelated OOM, before it could execute -- the first real K3 prefill had not reached layer 28. --- batchgen/kv_cache/host_kv_mananger_config.py | 21 ++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/batchgen/kv_cache/host_kv_mananger_config.py b/batchgen/kv_cache/host_kv_mananger_config.py index ddbf91647..83a5a1a91 100644 --- a/batchgen/kv_cache/host_kv_mananger_config.py +++ b/batchgen/kv_cache/host_kv_mananger_config.py @@ -145,6 +145,22 @@ 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( @@ -166,6 +182,7 @@ def bytes_per_page(self) -> int: "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] = {} @@ -214,6 +231,10 @@ def bytes_per_page(self) -> int: "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", From af34da963a9ea9ed762c1ec6abdd86a3488b6924 Mon Sep 17 00:00:00 2001 From: tairanxu Date: Fri, 7 Aug 2026 00:08:15 +0100 Subject: [PATCH 11/19] fix(worker): log the first token prefill actually sampled Prefill computes the first generated token (`_select_tokens` on the last-token logits) and appends it to output_tokens. For a prefill-only model that token is then thrown away: a max_tokens=1 request still enters decode (PREFILL_PLAN C4), decode raises, and the client receives the decode error string instead of the token. Full 93-layer Kimi-K3 prefills have now completed successfully (all 82,432 experts streamed) without anyone being able to see what the model produced. Rank 0, ids only -- the worker has no tokenizer, so decode them client-side. The proper fix is to return the prefill token when max_tokens=1 instead of entering decode; that is C4 and a larger change. --- batchgen/batchgen_worker.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 02af524da..c66c250b4 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -7102,6 +7102,17 @@ def prefill_prepacked(self, batch: list[int]): ) output_tokens.append(batch_new_tokens) + # The FIRST generated token, straight out of prefill. Logged + # because a max_tokens=1 request currently still enters decode + # (PREFILL_PLAN C4) and returns decode's error instead of this, + # so for a prefill-only model the token is computed and then + # thrown away with no way to see it. 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]) + # Reset prepack mode Attn_Wrapper.prepack_mode = False Attn_Wrapper.prepack_cu_seqlens = None From 1842a65e2a90dcfb2de15c66745f303e6eb67bac Mon Sep 17 00:00:00 2001 From: TairanXU Date: Sat, 8 Aug 2026 20:41:02 +0100 Subject: [PATCH 12/19] perf(worker): seed preallocated block_residual; free inputs_embeds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two prefill-memory fixes from batchgen_design/model_support/kimi_k3/PREFILL_MEMORY_AUDIT.md section 7. Depends on the model-side prealloc commit in the companion model PR, which adds KimiLinearModel._new_block_residual. fix 3 — the prepack-prefill loop seeded block_residual itself with new_zeros(S, 0, H), so every block boundary reallocated it by cat, and the old and new buffers were co-live at K3's last boundary. It now asks the model for a zero-column view of a buffer preallocated for all 8 boundaries. The existing `block_residual = None` immediately above is load-bearing: it drops the previous micro-batch's view before the next buffer is allocated. fix 4 — `hidden_states = inputs_embeds.unsqueeze(0)` is a VIEW, so hidden_states alone already keeps the embedding storage alive for exactly as long as layer 0 needs it. Keeping inputs_embeds bound as well pinned that storage for the whole 93-layer forward, dead weight across layers 1-92. One `del`. --- batchgen/batchgen_worker.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index c66c250b4..f9413b874 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -7023,6 +7023,14 @@ 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 @@ -7033,9 +7041,14 @@ def prefill_prepacked(self, batch: list[int]): use_attn_res = getattr(self.model.model, "use_attn_residuals", False) block_residual = None if use_attn_res: - block_residual = hidden_states.new_zeros( - hidden_states.shape[0] * hidden_states.shape[1], 0, - hidden_states.shape[2]) + # 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): if use_attn_res: From f2ee9be3e61b57970055408cac3005bcc12b231b Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 08:25:04 +0800 Subject: [PATCH 13/19] feat(worker): return the prefill-sampled token for max_tokens=1 (C4) Contract -------- Prefill already samples the first token and appends it through the normal decode write path: the writeback loop at the end of prefill()/ prefill_prepacked() writes it into query_book[local_idx].decoded_tokens at seq.decoded_length and bumps decoded_length. But _update_batch_status( prefill_uuids, PREFILLED) then handed EVERY prefilled sequence to the decode phase unconditionally, so a max_tokens=1 request had to make a full decode round-trip to be told it was already finished. _finish_prefill_completed_sequences() now runs right after that status update and completes the sequences whose decode budget is already satisfied (decoded_length >= max_decode_length). They are reported through the same path decode uses -- incremental writer, _gather_completed_tokens, KV release, PREFILLED -> COMPLETED, _report_completion -- so the HTTP response carries the decoded token text with the existing length-capped finish_reason. No parallel token-recording path is introduced and the handoff of output_tokens to decode for the surviving sequences is untouched, so nothing is double-appended. Only the length budget is tested. That is decode's own first completion test (CompletionHandler.is_sequence_completed / _check_and_handle_completions) and the test that wins in get_finish_reason, so the reported finish_reason is unchanged. Sequences that stop for any other reason (EOS, context limit) stay PREFILLED and reach decode exactly as before. Skipping decode --------------- Sequences still needing tokens remain PREFILLED, so the decode while-loop condition (has_prefilled() or has_in_decode() or has_on_hold()) is false only when NOTHING remains -- no configure_decoding, no decode-model load. The skip therefore falls out of replicated batch state rather than a separate branch. Rank alignment: decoded_length is advanced only on the owning rank, so the completed set is derived AFTER _sync_sequence_metadata() replicates it to every rank. Every rank computes the identical set from identical batch-global state; a rank-divergent skip would deadlock the next collective. That sync is needed anyway, because _report_completion reads decoded_length on rank 0 for sequences owned elsewhere (same reason the decode path syncs before reporting). Behaviour change ---------------- max_tokens > 1: unchanged on every model, including the 48B path -- still enters decode (and under Kimi-K3 stream_all_modules still raises the M-PR-6 prefill-only error, which is correct until decode exists). max_tokens == 1: now returns the prefill token WITHOUT the decode round-trip. This is a deliberate, documented improvement, not just a K3 fix -- the 48B path used to load the decode model and configure decoding purely to discover the sequence was already complete. Same token, same finish_reason ("length"), one less model load. For a prefill-only model it is the difference between the answer and decode's NotImplementedError string reaching the client. Verified on CPU: the branch logic was extracted from this source with ast and driven with a stubbed 3-rank mixed batch (max_tokens 1,1,5,3) -- identical completed set on every rank, correct COMPLETED/PREFILLED split, token recorded exactly once, identical collective sequence on every rank, decode skipped only when all sequences finish. py_compile + tests/test_kimi_k3_model.py (51 passed). --- batchgen/batchgen_worker.py | 116 +++++++++++++++++++++++++++++++++--- 1 file changed, 109 insertions(+), 7 deletions(-) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index f9413b874..638be3f99 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -4901,8 +4901,101 @@ 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) + # 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]]: @@ -5583,6 +5676,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. @@ -7115,12 +7216,13 @@ def prefill_prepacked(self, batch: list[int]): ) output_tokens.append(batch_new_tokens) - # The FIRST generated token, straight out of prefill. Logged - # because a max_tokens=1 request currently still enters decode - # (PREFILL_PLAN C4) and returns decode's error instead of this, - # so for a prefill-only model the token is computed and then - # thrown away with no way to see it. Rank 0 only; ids only - # (the worker has no tokenizer -- decode them client-side). + # 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", From 30a88e06e775a1c949d6ff66dbee9b6d120d6825 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 08:39:11 +0800 Subject: [PATCH 14/19] perf(worker): share the input_ids pool per node and size it by need Nothing model-specific is introduced. The problem ----------- QueryBookBufferPool allocated torch.zeros((num_sequences, model_context_length)) per worker. With a large --max-pool-size and K3's 1,048,576-token context that is 8 B x rows x context of int64 zeros for input_ids, per worker, times every rank on the node -- and generate_persistent() passed model_context_length as max_decoding_length too, so decoded_tokens was a second buffer of the same size. That OOM-killed the node twice and made --max-pool-size 1 load-bearing. get_input_ids_view(slot, seq_extended_size) always slices to the real length; the full width was never used. Fix 1: share it. The tokenized global batch is IDENTICAL on every rank (_tokenize_global_batch all-gathers the results to all of them), and every rank fills a row for every sequence, so the ranks were holding world_size byte-identical copies. input_ids_buffer is now ONE multiprocessing shared_memory segment per node: the node's rank%NUM_GPUS_PER_NODE==0 creates it, dist.barrier, everyone else attaches (allocate_node_shared_int64). Not the parent-allocated allocate_host_kv_cache pattern, because the parent cannot size this buffer -- the widths come from the tokenized requests, which only exist inside the worker loop after dist is up. decoded_tokens_buffer stays PRIVATE per rank: only the owning rank writes a sequence's decoded tokens, so those copies legitimately differ. It is, however, now sized by need as well -- the pool-mode path set it to model_context_length, the second full-width buffer, not what the spec's "already width-bounded" assumed. Fix 2: size by need. Width = the widest seq_extended_size the batch will actually ask for (prompt + THAT request's decode budget), capped at the model context length. Never the context length, never --max-pool-size, which keeps its row-count meaning: allocate_slot() still hard-fails past it. In pool mode the allocation is deferred out of generate_persistent() into the first admission, because that is the first moment either width is known. Overflow is never silent. A later admission needing more grows the pool with a WARNING naming both old and new sizes, copies the live rows across (adopt()) and rebinds every seq.input_ids / query_book view; the superseded segment stays mapped (untracked views must not be unmapped under) with its name unlinked. get_input_ids_view() raises QueryBookPoolCapacityError rather than silently returning a SHORT view if anything slips past that, and allocate_slot()'s exhaustion error is now the same named type. The slot trap ------------- A shared buffer is only safe if slot -> row is globally consistent. It is, almost: _tokenize_global_batch Phase 3 and _tokenize_admitted_sequences allocate a slot for EVERY sequence on EVERY rank in the same order, and _report_completion frees on every rank. The one exception was host-KV migration -- the source rank called free_slot() and set _buffer_slot = -1 while the destination reused that same slot index and every other rank kept it. So the source could hand row S to a new admission while everyone else still read S as the migrated sequence, and the -1 left behind made an eviction re-entry write into row -1, i.e. the LAST row of the pool. Chosen fix: derive nothing new -- just stop the divergence. The source-side free is removed (the slot IS still that sequence's, globally), and the two places that could reintroduce a rank-local slot now hard-fail with QueryBookPoolCapacityError instead of quietly allocating: migration-receive with _buffer_slot < 0, and re-entry with _buffer_slot < 0. Per-rank disjoint row ranges were rejected: every rank needs a row for every sequence, so disjoint ranges would save nothing. Measured (CPU, 4 processes, real code lifted from this file with ast) --------------------------------------------------------------------- - one segment: each process sees every other process's disjoint row writes; a private-pool control run sees none of them - per-process private growth for a 64 MiB buffer, 4 procs: shared 0.2/0.2/0.2/0.1 MiB (total 0.7) vs private 64.2 x4 (total 256.8) - identical buffer digest in every process - QueryBookPoolCapacityError fires on both width overflow and row overflow, naming both sizes - grow: adopt() preserves every live row, carries slot bookkeeping, the wider view then works py_compile + tests/test_kimi_k3_model.py (51 passed). --- batchgen/batchgen_worker.py | 375 ++++++++++++++++++++++++++++++++---- 1 file changed, 342 insertions(+), 33 deletions(-) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 638be3f99..f112cdf17 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -274,24 +274,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 +402,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 +413,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: @@ -761,6 +872,17 @@ def __init__(self, args: BatchGenWorkerArgs): 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): @@ -1261,7 +1383,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 +1468,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. @@ -3590,11 +3743,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: @@ -3823,9 +3981,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) @@ -4214,6 +4382,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: @@ -4384,17 +4674,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): @@ -5219,19 +5522,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 @@ -6228,6 +6529,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) From 211ae2aed485d0d1782e6cdee5f83b8b9a088fe9 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 09:54:05 +0800 Subject: [PATCH 15/19] fix(worker): keep C4-completed seqs visible to legacy result gather Fixes a regression introduced by the preceding C4 commit. Symptom ------- Every /v1/inference request whose sequences complete at prefill (max_tokens=1, the C4 path) came back with "Results unexpectedly empty after inference" (server_worker_main_loop.py:552). The batch endpoint was unaffected. Root cause ---------- generate() gathers legacy results at the very end by iterating self._local_to_uuid_map. _finish_prefill_completed_sequences() completes C4 sequences during PREFILL and calls _report_completion(), which calls release_local_query_slot() (query_book.py:51) -- and that pops uuid_to_local_map, local_to_uuid_map AND the query_book entry. By the time the gather runs, a C4-completed sequence is gone from every structure the gather reads, so local_results is empty on every rank and rank 0 returns {}. The batch path never saw this because C4 already feeds it through _submit_completed_to_incremental_writer() before the pop. Fix --- Capture the decoded text while the slot still exists -- immediately before the _report_completion loop -- into self._prefill_completed_results (global_idx -> str), and have the legacy gather seed local_results from it before walking the still-live map. Ordering is the whole fix: capture must precede the pop. The capture is gated on `self._response_queue is None`, i.e. legacy mode only. Pool/batch mode is fed by _report_completion and the incremental writer, so accumulating there as well would both be redundant and grow without bound in a persistent server that never drains the store. Nothing is reported twice on either path: the store feeds only generate()'s return value, which pool mode discards. The store is cleared in process_new_batch() alongside the new SequenceBatch, so it never carries across batches. Verified on CPU (real method lifted from this source with ast, stubbed so that _report_completion actually pops -- a capture-after-pop ordering mistake fails the test): - legacy (no response queue): gather is NON-empty, contains both C4 texts plus the survivor gathered from the live map - pool (response queue present): store stays empty, each sequence reported exactly once, incremental writer called exactly once - no C4 sequences: store untouched, gather unchanged py_compile + tests/test_kimi_k3_model.py. --- batchgen/batchgen_worker.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index f112cdf17..698c1335b 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -869,6 +869,10 @@ 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 @@ -4100,6 +4104,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): @@ -5290,6 +5295,28 @@ def _finish_prefill_completed_sequences(self, prefill_uuids: List[str]) -> List[ 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: + 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] + self._prefill_completed_results[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: @@ -6292,6 +6319,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(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: From 5f92b08306fdeb34f8cdac77285a33777469564a Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 09:56:27 +0800 Subject: [PATCH 16/19] feat(worker): emit a structured [METRICS] prefill record Implements the prefill-metrics proposal, with one deliberate departure documented below. Why --- prefill_s -- the most-quoted BatchGen performance number -- is scraped out of a tqdm progress bar. That string is presentation, not an API: it breaks on any desc=/bar_format/width change, it is emitted with \r so one physical log line can carry several bars, it reports a rate rather than an elapsed time, and tqdm prints the final bar twice with the copies disagreeing in the last digit. prefill_prepacked now emits one JSON line after the forward loop: [METRICS] {"phase":"prefill","prefill_s":...,"sequences":..., "tokens_total":...,"seq_len_min":...,"seq_len_max":..., "micro_batches":...,"max_tokens_per_micro_batch":..., "world_size":...,"rank":...,"first_sampled_token_ids":[...]} The report tooling already prefers this over the tqdm scrape and drops the "came from a progress bar" warning when it is present. Additive-only: the presence of "phase" is the whole handshake, so older readers are unaffected. The timing window starts immediately before the `with torch.inference_mode():` that wraps the loop, so it is the pure forward pass. configure_prefill is NOT folded in -- it is already reported separately as `Config completed`. Departure from the proposal: emitted by EVERY prefilling rank, not rank 0 ------------------------------------------------------------------------- The proposal specifies "rank 0 only" and fixes "rank": 0. That is wrong, and a 131,069-token run proved it: the batch was owned by rank 2 and the run was left unmeasurable, reporting no prefill record at all. Two independent reasons rank-0 gating cannot work: 1. prefill_prepacked is called only under `if local_prefill_indices:`, so a rank that owns none of the batch never reaches the emit point at all -- the line is simply absent, not zero. The same is true of the tqdm bar, which is additionally disable=(self.rank != 0), so nothing is emitted anywhere. 2. The proposal's snippet reads output_tokens[0] unguarded, which is an IndexError on a rank that ran zero micro-batches. So the line carries the emitting rank's real id and its LOCAL sequence/token counts. The wall time for the batch is the MAX of prefill_s over the emitting ranks; aggregating that is the parser's job. PARSER DEPENDENCY: the report tooling currently takes prefill_s from the first matching phase=="prefill" line. Until it aggregates over ranks, a multi-rank prefill will record an arbitrary rank's forward time rather than the max. This is still strictly better than today, where such a run yields no number at all. Verified on CPU (dict literal lifted from this source with ast and evaluated against stubs, so the contract is checked against shipping code): all 11 contract fields present and correctly typed, no undeclared fields, rank reflects the emitting rank, compact single-line JSON that re-parses, and both crashes the proposal's snippet would have hit (empty output_tokens, empty seq_lengths_list) return safe values instead. py_compile + tests/test_kimi_k3_model.py (51 passed). --- batchgen/batchgen_worker.py | 38 +++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 698c1335b..87976252c 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 @@ -7376,6 +7377,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), @@ -7567,6 +7571,40 @@ def prefill_prepacked(self, batch: list[int]): "[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 From 8fc02f29cef9b58ee1fb2f3dcb5cbb507fd6d421 Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 09:58:37 +0800 Subject: [PATCH 17/19] fix(worker): make the C4 result store survive a hot reload Follow-up to the C4 gather fix, which introduced self._prefill_completed_results in __init__. Hot reload (/v1/reload) rebinds methods on a LIVE worker instance and never re-runs __init__, so a reloaded process would get the new code without the new state and raise AttributeError on the first C4 completion. This is not hypothetical: _validate_reload (server_worker_main_loop.py:68) diffs the two __init__ sources, lists attributes the running worker lacks, and logs "These will cause AttributeError if accessed." It reports the hazard and deliberately does not repair it. Both read sites now tolerate the attribute being absent and create it on first use. __init__ still declares it, so a cold start is unchanged. Verified: the C4 result-store checks still pass on both paths (legacy gather non-empty with C4 texts; pool mode accumulates nothing and reports once). py_compile. --- batchgen/batchgen_worker.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 87976252c..af02730a2 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -5308,15 +5308,21 @@ def _finish_prefill_completed_sequences(self, prefill_uuids: List[str]) -> List[ # 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] - self._prefill_completed_results[seq.global_idx] = ( - self._decode_tokens_to_string(_decoded) - ) + 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. @@ -6322,7 +6328,7 @@ def generate(self): 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(self._prefill_completed_results.items()) + 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: From 7cf61f461bcf367e86d78e582f443a9251f4fcff Mon Sep 17 00:00:00 2001 From: TairanXU Date: Wed, 12 Aug 2026 21:40:50 +0800 Subject: [PATCH 18/19] feat(server): reject /v1/inference at the door instead of parking it Implements the ruling that /v1/inference is DEPRECATED and all inference goes through the batch API. The hazard ---------- The legacy path carried no request-id routing. On a pool-mode server worker_manager.infer() put its payload on the shared worker request queue and then blocked on response_queue.get(). The worker admission loop recognises only None and {"type": "admit"}, so the legacy payload matched no branch and was SILENTLY DROPPED -- while the caller sat on the shared response queue and took the next completion belonging to somebody else's batch. That is not theory: it corrupted a real production batch, which was left parked in in_progress with its completion consumed by a /v1/inference request. Rejected at the earliest seam ----------------------------- The HTTP handler now raises 410 immediately, before any queue interaction, with a JSON detail naming the deprecation and pointing at /v1/batches. The route is kept rather than deleted so a caller gets the explanation instead of a 404 it would read as a typo, and it takes NO request body so an unparseable legacy payload still gets the deprecation notice rather than a 422 about its schema. The old body is deleted, not left unreachable below the raise. BatchGenHttpClient.submit_inference raises the same named error without touching the network. The symbol is kept deliberately: deleting it yields an AttributeError, which tells a caller only that something vanished. Defense in depth, and what it deliberately does NOT catch -------------------------------------------------------- Both worker admission polls -- _poll_admissions and the idle-wait block inside generate() -- now raise LegacyInferenceDeprecated on a dict carrying "prompts", the shape worker_manager.infer builds. Reaching either line means some producer other than the now-closed HTTP route is putting legacy payloads on the queue, so failing loudly is right; parking is exactly the bug. This is deliberately NOT a catch-all on unrecognised messages. {"command": "reload"} also reaches _poll_admissions and is dropped there today, so a catch-all would convert every mid-batch /v1/reload into a dead server. The guard keys on the legacy shape only, and a test asserts the reload message still passes through untouched. Relationship to the C4 gather fix --------------------------------- The C4 legacy-gather fix STANDS and is not reverted. Verifying this change showed its premise was too narrow: the legacy result gather at the end of generate() is not exclusive to /v1/inference. worker.infer() has a second caller, batch_scheduler.py:313, which is the /v1/batches path taken whenever max_pool_size == 0. So the gather remains live for legacy-mode batches -- a batch-API mode, squarely inside the ruling -- while the deprecated HTTP entry is rejected at the door. Reverting it would have re-broken max_tokens=1 batches on a --max-pool-size 0 server. (It would also have conflicted: the hot-reload commit rewrote both hunks the gather fix added.) The shared seam --------------- batchgen/deprecation.py holds the message and exception. The three consumers cannot import one another: batchgen_client is imported by batchgen/__init__ and must carry no heavy deps, http_server pulls FastAPI/pydantic, batchgen_worker pulls torch. A module that imports nothing is the only seam all three can share, and a test asserts it stays that way. Verified on CPU (py3.11): - tests/test_kimi_k3_model.py: 51 passed - py_compile on all four files - 38 checks over the SHIPPING source: the handler is checked structurally with ast (http_server will not import on the dev box -- JIT core_engine needs ninja) and asserts the body is one raise, 410, the shared constants, no body param, and no reference to worker/infer/request_queue/response_queue/put/get; _poll_admissions is LIFTED from source with ast and executed against stubs, confirming a legacy payload raises and admits nothing, while admit, reload, None and an empty queue all keep their prior behaviour. Mutation-checked: stubbing the guard out makes the legacy-payload assertion fail, so the test is not vacuous. NOT verified live. The handler is not hot-reloadable: _reload_worker_module reloads batchgen_worker plus four dependency modules, and http_server is not among them; the routes are closures registered in create_app() in the parent process. The curl check rides the next scheduled restart rather than forcing one. --- batchgen/batchgen_client.py | 47 ++++------------ batchgen/batchgen_worker.py | 18 ++++++ batchgen/deprecation.py | 26 +++++++++ batchgen/server/http_server.py | 100 +++++++++------------------------ 4 files changed, 81 insertions(+), 110 deletions(-) create mode 100644 batchgen/deprecation.py 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 af02730a2..cef5bd2f2 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -17,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 @@ -1294,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 @@ -5851,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) 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/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): From 8d961b7576f6cd1e02ab3a3d2284b3cf16b69b6f Mon Sep 17 00:00:00 2001 From: TairanXU Date: Thu, 13 Aug 2026 05:40:58 +0800 Subject: [PATCH 19/19] test(k3): exercise the buffer-pool grow rebind-with-live-sequences branch The copy-live-rows / re-point-views path in _ensure_buffer_pool has never run with a live sequence: wave admission serialises batches so every observed grow rebound 0. This CPU test forces a grow while a mid-decode sequence occupies a row and asserts the row survives byte-for-byte, the slot mapping and query_book entry rebind onto the grown buffer, and a second admission does not clobber the first. It exercises the real batchgen_worker.py source (AST-extracted, since the module import pulls the JIT core_engine) and carries a mutation check that neuters the row-copy to prove the assertions are load-bearing. --- .../unit/test_query_book_pool_grow_rebind.py | 272 ++++++++++++++++++ 1 file changed, 272 insertions(+) create mode 100644 tests/unit/test_query_book_pool_grow_rebind.py 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"]))