diff --git a/batchgen/attention/forward_metadata.py b/batchgen/attention/forward_metadata.py new file mode 100644 index 000000000..4d7da288b --- /dev/null +++ b/batchgen/attention/forward_metadata.py @@ -0,0 +1,152 @@ +"""First-class forward metadata for attention execution. + +These dataclasses describe the logical forward batch without depending on +legacy wrapper class variables. They intentionally do not mutate runtime state; +builders are responsible for constructing them from already-validated static +planning inputs. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, Optional + +import torch + +ForwardPhase = Literal["prefill", "decode"] + + +@dataclass(frozen=True) +class PrefillAttentionMetadata: + """Attention metadata for prefill or suffix-only prefill. + + Prefix reuse is represented by q/kv length divergence: + ``kv_seq_lens[i] - q_seq_lens[i]`` is the cached prefix length for sequence + ``i``. The legacy wrapper context mirrors these derived values into + ``AttnWrapperBase.prepack_prefix_*`` for model wrappers. + """ + + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + max_seqlen_q: int + max_seqlen_k: int + q_seq_lens: list[int] + kv_seq_lens: list[int] + position_ids: torch.Tensor + append_seq_lens: Optional[list[int]] = None + + @property + def batch_size(self) -> int: + return len(self.q_seq_lens) + + +@dataclass(frozen=True) +class DecodeAttentionMetadata: + """Attention metadata for decode forward batches.""" + + cache_seqlens: torch.Tensor + max_seqlen: int + page_table: Optional[torch.Tensor] = None + slot_indices: Optional[torch.Tensor] = None + batch_slice: Optional[slice] = None + + @property + def batch_size(self) -> int: + return int(self.cache_seqlens.numel()) + + +@dataclass(frozen=True) +class KVCacheMetadata: + """KV cache handles associated with a forward batch.""" + + gpu_paged_kv_manager: Optional[object] = None + host_worker_view: Optional[object] = None + aux_gpu_paged_kv_manager: Optional[object] = None + aux_host_worker_view: Optional[object] = None + prefill_prefix_materialization: Optional[object] = None + + +@dataclass(frozen=True) +class ForwardBatchMetadata: + """Top-level metadata object for one model forward batch.""" + + phase: ForwardPhase + global_sequence_ids: list[int] + prefill: Optional[PrefillAttentionMetadata] = None + decode: Optional[DecodeAttentionMetadata] = None + kv_cache: Optional[KVCacheMetadata] = None + + def require_prefill(self) -> PrefillAttentionMetadata: + if self.phase != "prefill" or self.prefill is None: + raise RuntimeError( + "Prefix cache prepack metadata requires prefill metadata" + ) + return self.prefill + + @property + def cu_seqlens(self) -> torch.Tensor: + return self.require_prefill().cu_seqlens_q + + @property + def cu_seqlens_cpu(self) -> list[int]: + return _build_cu_seqlens_values(self.seq_lengths) + + @property + def max_seqlen(self) -> int: + return int(self.require_prefill().max_seqlen_q) + + @property + def num_sequences(self) -> int: + return int(self.require_prefill().batch_size) + + @property + def seq_lengths(self) -> list[int]: + return [int(length) for length in self.require_prefill().q_seq_lens] + + @property + def append_seq_lengths(self) -> list[int]: + prefill = self.require_prefill() + if prefill.append_seq_lens is None: + return [int(length) for length in prefill.q_seq_lens] + return [int(length) for length in prefill.append_seq_lens] + + @property + def prefix_shared_tokens(self) -> Optional[list[int]]: + tokens = [ + int(kv_len) - int(append_len) + for kv_len, append_len in zip( + self.require_prefill().kv_seq_lens, + self.append_seq_lengths, + ) + ] + if any(token < 0 for token in tokens): + raise RuntimeError( + "Prefix cache metadata requires kv lengths >= append lengths" + ) + return tokens if any(token > 0 for token in tokens) else None + + @property + def prefix_reuse_mode(self) -> bool: + tokens = self.prefix_shared_tokens + return tokens is not None and any(token > 0 for token in tokens) + + @property + def full_seq_lengths(self) -> Optional[list[int]]: + if not self.prefix_reuse_mode: + return None + return [int(length) for length in self.require_prefill().kv_seq_lens] + + def cu_seqlens_list(self) -> list[int]: + return list(self.cu_seqlens_cpu) + + def append_seq_lengths_list(self) -> list[int]: + return list(self.append_seq_lengths) + + +def _build_cu_seqlens_values(seq_lengths: list[int]) -> list[int]: + values = [0] + running = 0 + for length in seq_lengths: + running += int(length) + values.append(running) + return values diff --git a/batchgen/attention/forward_metadata_context.py b/batchgen/attention/forward_metadata_context.py new file mode 100644 index 000000000..09f02a6f1 --- /dev/null +++ b/batchgen/attention/forward_metadata_context.py @@ -0,0 +1,177 @@ +"""Context binding for first-class attention forward metadata. + +This module is the compatibility bridge between explicit +``ForwardBatchMetadata`` and the legacy ``AttnWrapperBase`` class variables. +The metadata object remains the source of truth; legacy fields are only +populated for the dynamic extent of a single forward call. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Iterator, Optional + +from batchgen.attention.forward_metadata import ( + DecodeAttentionMetadata, + ForwardBatchMetadata, + KVCacheMetadata, + PrefillAttentionMetadata, +) + + +_CURRENT_FORWARD_BATCH_METADATA: ContextVar[Optional[ForwardBatchMetadata]] = ( + ContextVar("current_forward_batch_metadata", default=None) +) + +_LEGACY_ATTENTION_FIELDS = ( + "phase", + "cur_batch", + "position_ids", + "prepack_mode", + "prepack_cu_seqlens", + "prepack_max_seqlen", + "prepack_num_sequences", + "prepack_seq_lengths", + "prepack_append_seq_lengths", + "prepack_prefix_reuse_mode", + "prepack_prefix_shared_tokens", + "prepack_full_seq_lengths", + "cache_seqlens", + "max_seqlen", + "gpu_paged_kv_manager", + "host_paged_kv_worker_view", + "prefill_prefix_materialization", + "gpu_paged_kv_manager_aux", + "host_paged_kv_worker_view_aux", +) + + +def get_current_forward_batch_metadata( + required: bool = False, +) -> Optional[ForwardBatchMetadata]: + """Return the metadata bound to the current execution context.""" + + metadata = _CURRENT_FORWARD_BATCH_METADATA.get() + if metadata is None and required: + raise RuntimeError("ForwardBatchMetadata is required but is not bound") + return metadata + + +@contextmanager +def bind_forward_batch_metadata( + metadata: ForwardBatchMetadata, +) -> Iterator[ForwardBatchMetadata]: + """Bind metadata for one forward and mirror it into legacy wrapper fields.""" + + if not isinstance(metadata, ForwardBatchMetadata): + raise TypeError("metadata must be a ForwardBatchMetadata instance") + + # Import lazily so metadata users can be unit-tested without importing model + # wrappers unless the compatibility bridge is actually used. + from batchgen.models.wrappers.attention import AttnWrapperBase + + previous_values = { + field: getattr(AttnWrapperBase, field, None) + for field in _LEGACY_ATTENTION_FIELDS + } + token = _CURRENT_FORWARD_BATCH_METADATA.set(metadata) + try: + _sync_legacy_attention_wrapper(AttnWrapperBase, metadata) + yield metadata + finally: + _CURRENT_FORWARD_BATCH_METADATA.reset(token) + for field, value in previous_values.items(): + setattr(AttnWrapperBase, field, value) + + +def _sync_legacy_attention_wrapper( + wrapper_cls: type, + metadata: ForwardBatchMetadata, +) -> None: + wrapper_cls.phase = metadata.phase + wrapper_cls.cur_batch = list(metadata.global_sequence_ids) + + if metadata.phase == "prefill": + assert metadata.prefill is not None + _sync_prefill_fields(wrapper_cls, metadata.prefill) + else: + assert metadata.decode is not None + _sync_decode_fields(wrapper_cls, metadata.decode) + + if metadata.kv_cache is not None: + _sync_kv_cache_fields(wrapper_cls, metadata.kv_cache) + + +def _sync_prefill_fields( + wrapper_cls: type, + prefill: PrefillAttentionMetadata, +) -> None: + wrapper_cls.position_ids = prefill.position_ids + wrapper_cls.prepack_mode = True + wrapper_cls.prepack_cu_seqlens = prefill.cu_seqlens_q + wrapper_cls.prepack_max_seqlen = int(prefill.max_seqlen_q) + wrapper_cls.prepack_num_sequences = prefill.batch_size + wrapper_cls.prepack_seq_lengths = list(prefill.q_seq_lens) + wrapper_cls.prepack_append_seq_lengths = _append_seq_lens(prefill) + wrapper_cls.cache_seqlens = None + wrapper_cls.max_seqlen = None + + _sync_prefix_reuse_fields(wrapper_cls, prefill) + + +def _sync_prefix_reuse_fields( + wrapper_cls: type, + prefill: PrefillAttentionMetadata, +) -> None: + prefix_lens = [ + int(kv_len) - int(q_len) + for q_len, kv_len in zip(prefill.q_seq_lens, prefill.kv_seq_lens) + ] + if any(length < 0 for length in prefix_lens): + raise ValueError( + "prefill kv sequence lengths must be >= query sequence lengths" + ) + if not any(length > 0 for length in prefix_lens): + wrapper_cls.prepack_prefix_reuse_mode = False + wrapper_cls.prepack_prefix_shared_tokens = None + wrapper_cls.prepack_full_seq_lengths = None + return + + full_seq_lens = [int(length) for length in prefill.kv_seq_lens] + wrapper_cls.prepack_prefix_reuse_mode = True + wrapper_cls.prepack_prefix_shared_tokens = prefix_lens + wrapper_cls.prepack_full_seq_lengths = full_seq_lens + + +def _sync_decode_fields( + wrapper_cls: type, decode: DecodeAttentionMetadata +) -> None: + wrapper_cls.position_ids = None + wrapper_cls.prepack_mode = False + wrapper_cls.prepack_cu_seqlens = None + wrapper_cls.prepack_max_seqlen = None + wrapper_cls.prepack_num_sequences = None + wrapper_cls.prepack_seq_lengths = None + wrapper_cls.prepack_append_seq_lengths = None + wrapper_cls.prepack_prefix_reuse_mode = False + wrapper_cls.prepack_prefix_shared_tokens = None + wrapper_cls.prepack_full_seq_lengths = None + wrapper_cls.cache_seqlens = decode.cache_seqlens + wrapper_cls.max_seqlen = int(decode.max_seqlen) + + +def _sync_kv_cache_fields(wrapper_cls: type, kv_cache: KVCacheMetadata) -> None: + wrapper_cls.gpu_paged_kv_manager = kv_cache.gpu_paged_kv_manager + wrapper_cls.host_paged_kv_worker_view = kv_cache.host_worker_view + wrapper_cls.prefill_prefix_materialization = ( + kv_cache.prefill_prefix_materialization + ) + wrapper_cls.gpu_paged_kv_manager_aux = kv_cache.aux_gpu_paged_kv_manager + wrapper_cls.host_paged_kv_worker_view_aux = kv_cache.aux_host_worker_view + + +def _append_seq_lens(prefill: PrefillAttentionMetadata) -> list[int]: + if prefill.append_seq_lens is None: + return list(prefill.q_seq_lens) + return [int(length) for length in prefill.append_seq_lens] diff --git a/batchgen/attention/gqa/__init__.py b/batchgen/attention/gqa/__init__.py index cd74449b5..229f2d079 100644 --- a/batchgen/attention/gqa/__init__.py +++ b/batchgen/attention/gqa/__init__.py @@ -6,6 +6,7 @@ Key components: - gqa_prefill_fa: Prefill using flash_attn_varlen_func (unpadded sequences) - gqa_decode_fa: Decode using flash_attn_with_kvcache (paged KV cache) +- gqa_extend_fa: Extend prefill using flash_attn_with_kvcache (paged KV cache) - apply_sink_correction: Post-correction for attention sinks - attention_ref: Reference implementation for testing @@ -16,6 +17,7 @@ from .fa_prefill import gqa_prefill_fa from .fa_decode import gqa_decode_fa, gqa_decode_fa_contiguous +from .fa_extend import gqa_extend_fa from .sink_correction import apply_sink_correction from .reference import attention_ref, attention_ref_no_sinks from .gqa_mode3 import gqa_decoding_mode_3_bf16 @@ -25,6 +27,7 @@ 'gqa_prefill_fa', 'gqa_decode_fa', 'gqa_decode_fa_contiguous', + 'gqa_extend_fa', 'apply_sink_correction', 'attention_ref', 'attention_ref_no_sinks', diff --git a/batchgen/attention/gqa/fa_extend.py b/batchgen/attention/gqa/fa_extend.py new file mode 100644 index 000000000..caa761d39 --- /dev/null +++ b/batchgen/attention/gqa/fa_extend.py @@ -0,0 +1,95 @@ +"""GQA extend prefill using FlashAttention paged KV cache.""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +_USE_FA3 = False +_flash_with_kvcache = None + +try: + from flash_attn_interface import ( + flash_attn_with_kvcache as _fa3_with_kvcache, + ) + + _USE_FA3 = True + _flash_with_kvcache = _fa3_with_kvcache +except ImportError: + pass + +if _flash_with_kvcache is None: + try: + from flash_attn import flash_attn_with_kvcache as _fa2_with_kvcache + + _flash_with_kvcache = _fa2_with_kvcache + except ImportError: + pass + + +def gqa_extend_fa( + *, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + page_table: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + sinks: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + sliding_window: Optional[int] = None, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run batched suffix-prefill attention over paged prefix+suffix KV. + + This is the paged-KV extend counterpart of varlen prefill attention. The + caller is responsible for writing the freshly computed suffix K/V into the + paged cache before calling this function. ``cache_seqlens`` and + ``cu_seqlens_k`` therefore describe the full logical KV lengths, while + ``cu_seqlens_q`` describes only the suffix query lengths. + """ + + if _flash_with_kvcache is None: + raise ImportError( + "Neither flash_attn_interface (FA3) nor flash_attn (FA2) is available" + ) + + if sliding_window is not None and sliding_window > 0: + window_size = (sliding_window - 1, 0) + else: + window_size = (-1, -1) + + if softmax_scale is None: + softmax_scale = q.shape[-1] ** -0.5 + + page_table_kwarg = "page_table" if _USE_FA3 else "block_table" + extra_kwargs = {page_table_kwarg: page_table} + result = _flash_with_kvcache( + q, + k_cache, + v_cache, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max_seqlen_q, + softmax_scale=softmax_scale, + causal=True, + window_size=window_size, + return_softmax_lse=sinks is not None, + **extra_kwargs, + ) + + if isinstance(result, tuple): + output = result[0] + lse = result[1] if len(result) > 1 else None + else: + output = result + lse = None + + if sinks is not None and lse is not None: + from .sink_correction import apply_sink_correction + + output = apply_sink_correction(output, lse, sinks) + + return output, lse diff --git a/batchgen/attention/mla/fa3_backend.py b/batchgen/attention/mla/fa3_backend.py index f1ec03420..835c9e6df 100644 --- a/batchgen/attention/mla/fa3_backend.py +++ b/batchgen/attention/mla/fa3_backend.py @@ -11,7 +11,9 @@ import deep_gemm # from deep_gemm import get_col_major_tma_aligned_tensor import logging -from typing import Tuple +import os +from dataclasses import dataclass +from typing import Callable, Optional, Tuple import torch.distributed as dist from ...moe.fused_dequant_gemm import fused_fp8_bf16_gemm @@ -658,6 +660,167 @@ def w8a16_gemm_dequant( return out +@dataclass(frozen=True) +class MlaPrepackProjection: + """Q and compressed-KV tensors shared by MLA prefill variants.""" + + q_nope: torch.Tensor + q_pe: torch.Tensor + normed_kv: Optional[torch.Tensor] = None + k_pe: Optional[torch.Tensor] = None + offload_kv: Optional[torch.Tensor] = None + + +def select_w8a16_gemm() -> Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor], + torch.Tensor, +]: + """Return the default W8A16 GEMM implementation used by MLA prefill.""" + use_dequant_path = os.environ.get("BATCHGEN_W8A16_DEQUANT", "0") == "1" + return w8a16_gemm_dequant if use_dequant_path else w8a16_gemm + + +def _apply_prepacked_mla_rope( + self, + q_pe: torch.Tensor, + k_pe: Optional[torch.Tensor], + position_ids: torch.Tensor, + rotary_seq_len: int, + *, + interleaved: bool, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + if interleaved: + from batchgen.attention.mla.rotary_embedding import ( + rotary_pos_emb_interleaved_native, + ) + rope_fn = rotary_pos_emb_interleaved_native + else: + rope_fn = rotary_pos_emb + cos, sin = self.rotary_emb(q_pe.unsqueeze(0), seq_len=rotary_seq_len) + q_pe = rope_fn( + q_pe.unsqueeze(0), + cos, + sin, + position_ids.unsqueeze(0), + 2, + ).squeeze(0) + if k_pe is None: + return q_pe, None + k_pe = rope_fn( + k_pe.unsqueeze(0), + cos, + sin, + position_ids.unsqueeze(0), + 2, + ).squeeze(0) + return q_pe, k_pe + + +def project_bf16_mla_q_and_compressed_kv_prepacked( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + rotary_seq_len: int, +) -> MlaPrepackProjection: + """Project prepacked MLA Q and compressed KV using BF16 module linears.""" + total_tokens = hidden_states.shape[0] + query_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + query_states = query_states.view(total_tokens, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split( + query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + normed_kv = self.kv_a_layernorm(compressed_kv) + k_pe = k_pe.view(total_tokens, 1, self.qk_rope_head_dim) + + q_pe, k_pe = _apply_prepacked_mla_rope( + self, + q_pe, + k_pe, + position_ids, + rotary_seq_len, + interleaved=False, + ) + if k_pe is None: + raise RuntimeError("BF16 MLA prepack projection failed to build k_pe") + k_pe_flat = k_pe.view(total_tokens, self.qk_rope_head_dim) + offload_kv = torch.cat([normed_kv, k_pe_flat], dim=-1) + del compressed_kv, k_pe_flat + return MlaPrepackProjection( + q_nope=q_nope, + q_pe=q_pe, + normed_kv=normed_kv, + k_pe=k_pe, + offload_kv=offload_kv, + ) + + +def project_w8a16_mla_q_and_compressed_kv_prepacked( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + rotary_seq_len: int, + weight_scale: dict, + gemm: Optional[ + Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] + ] = None, +) -> MlaPrepackProjection: + """Project prepacked MLA Q and compressed KV using the default W8A16 path.""" + gemm = select_w8a16_gemm() if gemm is None else gemm + total_tokens = hidden_states.shape[0] + query_states = gemm( + self.q_a_proj.weight.data, + weight_scale["q_a_proj.weight_scale_inv"], + hidden_states, + ) + query_states = self.q_a_layernorm(query_states) + query_states = gemm( + self.q_b_proj.weight.data, + weight_scale["q_b_proj.weight_scale_inv"], + query_states, + ) + query_states = query_states.view(total_tokens, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split( + query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = gemm( + self.kv_a_proj_with_mqa.weight.data, + weight_scale["kv_a_proj_with_mqa.weight_scale_inv"], + hidden_states, + ) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + normed_kv = self.kv_a_layernorm(compressed_kv) + k_pe = k_pe.view(total_tokens, 1, self.qk_rope_head_dim) + + q_pe, k_pe = _apply_prepacked_mla_rope( + self, + q_pe, + k_pe, + position_ids, + rotary_seq_len, + interleaved=True, + ) + if k_pe is None: + raise RuntimeError("W8A16 MLA prepack projection failed to build k_pe") + k_pe_flat = k_pe.view(total_tokens, self.qk_rope_head_dim) + offload_kv = torch.cat([normed_kv, k_pe_flat], dim=-1) + del compressed_kv, k_pe_flat + return MlaPrepackProjection( + q_nope=q_nope, + q_pe=q_pe, + normed_kv=normed_kv, + k_pe=k_pe, + offload_kv=offload_kv, + ) + + @torch.inference_mode() def mla_prefill_flashattention3_w8a16_deepgemm( self, @@ -1038,6 +1201,7 @@ def mla_prefill_flashattention3_prepacked( cu_seqlens: torch.Tensor, max_seqlen: int, num_sequences: int, + prefix_context=None, ) -> tuple[torch.Tensor, torch.Tensor]: """ MLA prefill on Hopper device for PREPACKED sequences. @@ -1059,31 +1223,27 @@ def mla_prefill_flashattention3_prepacked( offload_kv: [total_tokens, kv_lora_rank + qk_rope_head_dim] for KV cache """ total_tokens = hidden_states.shape[0] - - # Project Q - query_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - query_states = query_states.view(total_tokens, self.num_heads, self.q_head_dim) - q_nope, q_pe = torch.split( - query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + rotary_seq_len = max_seqlen + if prefix_context is not None: + rotary_seq_len = prefix_context.rotary_seq_len(position_ids, max_seqlen) + + projection = project_bf16_mla_q_and_compressed_kv_prepacked( + self, + hidden_states, + position_ids, + rotary_seq_len, ) - - # Project KV - compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - compressed_kv, k_pe = torch.split( - compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) - normed_kv = self.kv_a_layernorm(compressed_kv) - k_pe = k_pe.view(total_tokens, 1, self.qk_rope_head_dim) - - # Apply rotary embeddings - cos, sin = self.rotary_emb(q_pe.unsqueeze(0), seq_len=max_seqlen) - # For prepacked, position_ids is 1D [total_tokens] - q_pe = rotary_pos_emb(q_pe.unsqueeze(0), cos, sin, position_ids.unsqueeze(0), 2).squeeze(0) - k_pe = rotary_pos_emb(k_pe.unsqueeze(0), cos, sin, position_ids.unsqueeze(0), 2).squeeze(0) - - k_pe_flat = k_pe.view(total_tokens, self.qk_rope_head_dim) - offload_kv = torch.cat([normed_kv, k_pe_flat], dim=-1) - del compressed_kv, k_pe_flat + q_nope = projection.q_nope + q_pe = projection.q_pe + normed_kv = projection.normed_kv + k_pe = projection.k_pe + offload_kv = projection.offload_kv + if normed_kv is None or k_pe is None or offload_kv is None: + raise RuntimeError("BF16 MLA prepack projection returned incomplete KV") + if prefix_context is not None: + if not prefix_context.prefix_reuse_mode: + raise RuntimeError("MLA prefix context has no enabled reuse mode") + return prefix_context.run_suffix_prefill(projection) # Expand KV kv = self.kv_b_proj(normed_kv) @@ -1140,6 +1300,7 @@ def mla_prefill_flashattention3_w8a16_deepgemm_prepacked( max_seqlen: int, num_sequences: int, weight_scale: dict, + prefix_context=None, ) -> tuple[torch.Tensor, torch.Tensor]: """ MLA prefill with W8A16 quantization for PREPACKED sequences. @@ -1162,50 +1323,29 @@ def mla_prefill_flashattention3_w8a16_deepgemm_prepacked( # Default: FP8 act_quant + DeepGEMM fp8_gemm_nt (matches SGLang/DeepGEMM # blockwise FP8 semantics and the decode path's w8a8_deepgemm). Opt into # the dequant-to-BF16 path via BATCHGEN_W8A16_DEQUANT=1. - import os as _os_gemm - _w8a16_dequant_path = _os_gemm.environ.get("BATCHGEN_W8A16_DEQUANT", "0") == "1" - _gemm = w8a16_gemm_dequant if _w8a16_dequant_path else w8a16_gemm - - # Project Q - query_states = _gemm( - self.q_a_proj.weight.data, - weight_scale["q_a_proj.weight_scale_inv"], - hidden_states - ) - query_states = self.q_a_layernorm(query_states) - query_states = _gemm( - self.q_b_proj.weight.data, - weight_scale["q_b_proj.weight_scale_inv"], - query_states - ) - - query_states = query_states.view(total_tokens, self.num_heads, self.q_head_dim) - q_nope, q_pe = torch.split( - query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) - - # Project KV - compressed_kv = _gemm( - self.kv_a_proj_with_mqa.weight.data, - weight_scale["kv_a_proj_with_mqa.weight_scale_inv"], - hidden_states - ) - compressed_kv, k_pe = torch.split( - compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + _gemm = select_w8a16_gemm() + rotary_seq_len = max_seqlen + if prefix_context is not None: + rotary_seq_len = prefix_context.rotary_seq_len(position_ids, max_seqlen) + projection = project_w8a16_mla_q_and_compressed_kv_prepacked( + self, + hidden_states, + position_ids, + rotary_seq_len, + weight_scale, + gemm=_gemm, ) - normed_kv = self.kv_a_layernorm(compressed_kv) - k_pe = k_pe.view(total_tokens, 1, self.qk_rope_head_dim) - - # Native interleaved RoPE (matches HF / SGLang / vLLM is_neox_style=False - # when rope_interleave=true). - from batchgen.attention.mla.rotary_embedding import rotary_pos_emb_interleaved_native - cos, sin = self.rotary_emb(q_pe.unsqueeze(0), seq_len=max_seqlen) - q_pe = rotary_pos_emb_interleaved_native(q_pe.unsqueeze(0), cos, sin, position_ids.unsqueeze(0), 2).squeeze(0) - k_pe = rotary_pos_emb_interleaved_native(k_pe.unsqueeze(0), cos, sin, position_ids.unsqueeze(0), 2).squeeze(0) - - k_pe_flat = k_pe.view(total_tokens, self.qk_rope_head_dim) - offload_kv = torch.cat([normed_kv, k_pe_flat], dim=-1) - del compressed_kv, k_pe_flat + q_nope = projection.q_nope + q_pe = projection.q_pe + normed_kv = projection.normed_kv + k_pe = projection.k_pe + offload_kv = projection.offload_kv + if normed_kv is None or k_pe is None or offload_kv is None: + raise RuntimeError("W8A16 MLA prepack projection returned incomplete KV") + if prefix_context is not None: + if not prefix_context.prefix_reuse_mode: + raise RuntimeError("MLA prefix context has no enabled reuse mode") + return prefix_context.run_suffix_prefill(projection) # Expand KV kv = _gemm( diff --git a/batchgen/attention/mla/flashinfer_extend.py b/batchgen/attention/mla/flashinfer_extend.py new file mode 100644 index 000000000..9c5833b15 --- /dev/null +++ b/batchgen/attention/mla/flashinfer_extend.py @@ -0,0 +1,279 @@ +"""FlashInfer MLA paged-KV extend helpers.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Optional + +import torch +from flashinfer import BatchMLAPagedAttentionWrapper + +_DEFAULT_WORKSPACE_BYTES = 384 * 1024 * 1024 +_WORKSPACE_CACHE: dict[tuple[str, Optional[int]], torch.Tensor] = {} +_WRAPPER_CACHE: dict[tuple[str, Optional[int], str], object] = {} +_PLAN_CACHE_KEY = "flashinfer_mla_extend_prefill" + + +@dataclass +class _FlashInferMlaExtendPlanState: + signature: tuple[object, ...] + wrapper: object + qo_indptr: torch.Tensor + kv_indptr: torch.Tensor + kv_indices: torch.Tensor + kv_len_arr: torch.Tensor + + +def run_flashinfer_mla_extend_prefill( + *, + query_states: torch.Tensor, + compressed_kv_cache: torch.Tensor, + page_table: torch.Tensor, + slot_indices: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + kv_lora_rank: int, + num_heads: int, + softmax_scale: float, + plan_cache: dict[str, object] | None = None, +) -> torch.Tensor: + """Run prefix-hit MLA extend prefill through FlashInfer paged attention. + + ``compressed_kv_cache`` is BatchGen's materialized GPU paged MLA cache with + shape ``[num_pages, page_size, 1, kv_lora_rank + rope_dim]``. The returned + tensor is packed as ``[1, tokens, heads, rank]`` so existing MLA + output-projection glue can stay unchanged. Exact full hits are represented + as one query token per sequence. + """ + + packed_query = _packed_query_view(query_states) + q_nope = packed_query[..., :kv_lora_rank].contiguous() + q_pe = packed_query[..., kv_lora_rank:].contiguous() + ckv_cache, kpe_cache = _split_compressed_mla_cache( + compressed_kv_cache, + kv_lora_rank=kv_lora_rank, + ) + device = packed_query.device + page_size = int(compressed_kv_cache.shape[1]) + kv_len_arr = cache_seqlens.to(device=device, dtype=torch.int32) + kv_indptr, kv_indices = _build_flashinfer_page_metadata( + page_table=page_table, + slot_indices=slot_indices, + cache_seqlens=kv_len_arr, + page_size=page_size, + ) + qo_indptr = cu_seqlens_q.to(device=device, dtype=torch.int32) + + wrapper = _get_or_plan_flashinfer_mla_wrapper( + device=device, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_len_arr=kv_len_arr, + num_heads=int(num_heads), + kv_lora_rank=int(kv_lora_rank), + rope_head_dim=int(q_pe.shape[-1]), + page_size=page_size, + softmax_scale=float(softmax_scale), + q_dtype=q_nope.dtype, + kv_dtype=ckv_cache.dtype, + plan_cache=plan_cache, + ) + output = wrapper.run(q_nope, q_pe, ckv_cache, kpe_cache) + return output.unsqueeze(0).contiguous() + + +def _packed_query_view(query_states: torch.Tensor) -> torch.Tensor: + if query_states.dim() == 4 and query_states.shape[0] == 1: + return query_states.squeeze(0) + if query_states.dim() == 4 and query_states.shape[1] == 1: + return query_states.reshape( + query_states.shape[0], + query_states.shape[2], + query_states.shape[3], + ) + if query_states.dim() == 3: + return query_states + raise RuntimeError( + "FlashInfer MLA extend prefill expects packed query states shaped " + "[1, tokens, heads, dim], [batch, 1, heads, dim], or " + "[tokens, heads, dim]" + ) + + +def _split_compressed_mla_cache( + compressed_kv_cache: torch.Tensor, + *, + kv_lora_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if compressed_kv_cache.dim() != 4 or compressed_kv_cache.shape[2] != 1: + raise RuntimeError( + "FlashInfer MLA extend prefill expects K-only compressed MLA cache " + "shaped [pages, page_size, 1, dim]" + ) + cache = compressed_kv_cache.squeeze(2) + return cache[..., :kv_lora_rank], cache[..., kv_lora_rank:] + + +def _build_flashinfer_page_metadata( + *, + page_table: torch.Tensor, + slot_indices: torch.Tensor, + cache_seqlens: torch.Tensor, + page_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + device = cache_seqlens.device + slot_indices = slot_indices.to(device=page_table.device, dtype=torch.long) + selected_table = page_table.index_select(0, slot_indices).to( + dtype=torch.int32 + ) + pages_per_sequence = torch.div( + cache_seqlens + (int(page_size) - 1), + int(page_size), + rounding_mode="floor", + ) + kv_indptr = torch.empty( + pages_per_sequence.numel() + 1, + dtype=torch.int32, + device=device, + ) + kv_indptr[0] = 0 + kv_indptr[1:] = torch.cumsum(pages_per_sequence, dim=0, dtype=torch.int32) + + page_offsets = torch.arange( + selected_table.shape[1], + dtype=torch.int32, + device=selected_table.device, + ) + valid_pages = page_offsets.unsqueeze(0) < pages_per_sequence.to( + device=selected_table.device + ).unsqueeze(1) + kv_indices = selected_table[valid_pages].to( + device=device, dtype=torch.int32 + ) + return kv_indptr, kv_indices.contiguous() + + +def _get_flashinfer_mla_wrapper(device: torch.device) -> object: + backend = os.getenv("BATCHGEN_FLASHINFER_MLA_BACKEND", "auto") + key = _cache_key(device) + (backend,) + wrapper = _WRAPPER_CACHE.get(key) + if wrapper is not None: + return wrapper + + workspace = _get_workspace(device) + wrapper = BatchMLAPagedAttentionWrapper(workspace, backend=backend) + _WRAPPER_CACHE[key] = wrapper + return wrapper + + +def _get_or_plan_flashinfer_mla_wrapper( + *, + device: torch.device, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + kv_len_arr: torch.Tensor, + num_heads: int, + kv_lora_rank: int, + rope_head_dim: int, + page_size: int, + softmax_scale: float, + q_dtype: torch.dtype, + kv_dtype: torch.dtype, + plan_cache: dict[str, object] | None, +) -> object: + signature = ( + _cache_key(device), + os.getenv("BATCHGEN_FLASHINFER_MLA_BACKEND", "auto"), + tuple(qo_indptr.shape), + tuple(kv_indptr.shape), + tuple(kv_indices.shape), + tuple(kv_len_arr.shape), + int(num_heads), + int(kv_lora_rank), + int(rope_head_dim), + int(page_size), + float(softmax_scale), + str(q_dtype), + str(kv_dtype), + ) + if plan_cache is not None: + cached = plan_cache.get(_PLAN_CACHE_KEY) + if ( + isinstance(cached, _FlashInferMlaExtendPlanState) + and cached.signature == signature + ): + return cached.wrapper + + if plan_cache is None: + wrapper = _get_flashinfer_mla_wrapper(device) + else: + wrapper = _new_flashinfer_mla_wrapper(device) + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + int(num_heads), + int(kv_lora_rank), + int(rope_head_dim), + int(page_size), + True, + float(softmax_scale), + q_dtype, + kv_dtype, + ) + if plan_cache is not None: + plan_cache[_PLAN_CACHE_KEY] = _FlashInferMlaExtendPlanState( + signature=signature, + wrapper=wrapper, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_len_arr=kv_len_arr, + ) + return wrapper + + +def _new_flashinfer_mla_wrapper(device: torch.device) -> object: + backend = os.getenv("BATCHGEN_FLASHINFER_MLA_BACKEND", "auto") + return BatchMLAPagedAttentionWrapper(_get_workspace(device), backend=backend) + + +def _get_workspace(device: torch.device) -> torch.Tensor: + key = _cache_key(device) + workspace = _WORKSPACE_CACHE.get(key) + if workspace is None: + workspace_bytes = int( + os.getenv("BATCHGEN_FLASHINFER_WORKSPACE_BYTES", "0") + or "0" + ) + if workspace_bytes <= 0: + workspace_mb = int( + os.getenv("BATCHGEN_FLASHINFER_WORKSPACE_MB", "0") + or "0" + ) + workspace_bytes = ( + workspace_mb * 1024 * 1024 + if workspace_mb > 0 + else _DEFAULT_WORKSPACE_BYTES + ) + workspace = torch.empty( + workspace_bytes, + dtype=torch.uint8, + device=device, + ) + _WORKSPACE_CACHE[key] = workspace + return workspace + + +def _cache_key(device: torch.device) -> tuple[str, Optional[int]]: + normalized = torch.device(device) + return normalized.type, normalized.index + + +def _reset_flashinfer_mla_extend_prefill_cache_for_tests() -> None: + _WORKSPACE_CACHE.clear() + _WRAPPER_CACHE.clear() diff --git a/batchgen/attention/mla/prefix_absorb.py b/batchgen/attention/mla/prefix_absorb.py new file mode 100644 index 000000000..07961f28f --- /dev/null +++ b/batchgen/attention/mla/prefix_absorb.py @@ -0,0 +1,99 @@ +"""MLA absorb helpers used by prefix-cache prefill paths.""" + +from __future__ import annotations + +from typing import Callable + +import torch + +W8A16GemmFn = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] + + +def prefix_rotary_seq_len(full_length: int, position_ids: torch.Tensor) -> int: + """Return the RoPE seq-len needed for a prefix-aware prefill batch.""" + + return max(int(full_length), int(position_ids.max().item()) + 1) + + +def build_absorbed_mla_query_states( + *, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + q_absorb: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + """Build FlashMLA query states from projected MLA q_nope/q_pe tensors.""" + + total_tokens = q_nope.shape[0] + num_heads = q_nope.shape[1] + kv_lora_rank = q_absorb.shape[2] + query_states = torch.empty( + 1, + total_tokens, + num_heads, + kv_lora_rank + q_pe.shape[-1], + dtype=dtype, + device=q_pe.device, + ) + query_states[0, :, :, :kv_lora_rank] = torch.einsum( + "thd,hdc->thc", + q_nope, + q_absorb, + ) + query_states[0, :, :, kv_lora_rank:] = q_pe + return query_states.contiguous() + + +def absorb_mla_attention_output( + *, + attn_out: torch.Tensor, + out_absorb: torch.Tensor, + v_head_dim: int, +) -> torch.Tensor: + """Apply MLA out-absorb and flatten heads for the final output projection.""" + + attn_output = torch.einsum("bqhc,hdc->bqhd", attn_out, out_absorb) + return attn_output.reshape( + attn_out.shape[0] * attn_out.shape[1], + attn_out.shape[2] * int(v_head_dim), + ) + + +def project_absorbed_mla_output( + *, + attn_out: torch.Tensor, + out_absorb: torch.Tensor, + v_head_dim: int, + output_projection: Callable[[torch.Tensor], torch.Tensor], +) -> torch.Tensor: + """Apply out-absorb followed by a BF16/regular output projection.""" + + return output_projection( + absorb_mla_attention_output( + attn_out=attn_out, + out_absorb=out_absorb, + v_head_dim=v_head_dim, + ) + ) + + +def project_absorbed_mla_output_w8a16( + *, + attn_out: torch.Tensor, + out_absorb: torch.Tensor, + v_head_dim: int, + o_proj_weight: torch.Tensor, + o_proj_scale: torch.Tensor, + gemm: W8A16GemmFn, +) -> torch.Tensor: + """Apply out-absorb followed by the selected W8A16 output GEMM.""" + + return gemm( + o_proj_weight, + o_proj_scale, + absorb_mla_attention_output( + attn_out=attn_out, + out_absorb=out_absorb, + v_head_dim=v_head_dim, + ), + ) diff --git a/batchgen/attention/prefix_aware_backend.py b/batchgen/attention/prefix_aware_backend.py new file mode 100644 index 000000000..86ec506b7 --- /dev/null +++ b/batchgen/attention/prefix_aware_backend.py @@ -0,0 +1,174 @@ +"""Prefix-aware attention backend adapters. + +The adapters in this module provide a small explicit interface for prefill +attention where query length and KV length can differ because cached prefix KV +is prepended to freshly computed suffix KV. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import logging +import os +from typing import Callable, Optional + +import torch + +from batchgen.prefix_reuse.materialization import ( + get_prefix_materialization_for_group, +) + + +@dataclass(frozen=True) +class GqaPrefixAwareAttentionBackend: + """GQA backend adapter for varlen prefill and paged extend prefill.""" + + layer_idx: int + num_kv_heads: int + head_dim: int + sinks: Optional[torch.Tensor] = None + softmax_scale: Optional[float] = None + sliding_window: Optional[int] = None + attention_fn: Optional[Callable[..., tuple[torch.Tensor, object]]] = None + + def forward_prefill( + self, + *, + query: torch.Tensor, + key: torch.Tensor, + value: Optional[torch.Tensor], + metadata, + kv_cache_metadata=None, + ) -> torch.Tensor: + if value is None: + raise RuntimeError("GQA prefix-aware prefill requires value tensor") + + from batchgen.models.wrappers.prefix_cache import ( + ensure_prefix_cache_forward_metadata, + ) + + metadata = ensure_prefix_cache_forward_metadata(metadata) + + cu_q = metadata.cu_seqlens.to(query.device) + materialization = ( + getattr(kv_cache_metadata, "prefill_prefix_materialization", None) + if kv_cache_metadata is not None + else None + ) + materialization = get_prefix_materialization_for_group( + materialization, + group_id=0, + consumer="GQA prefix-aware prefill", + ) + if metadata.prefix_reuse_mode and materialization is None: + raise RuntimeError( + "GQA partial-hit prefix reuse requires GPU paged materialization" + ) + + if metadata.prefix_reuse_mode: + return self._forward_paged_extend_prefill( + query=query, + key=key, + value=value, + metadata=metadata, + materialization=materialization, + ) + + key_for_attn = key + value_for_attn = value + cu_k = cu_q + max_seqlen_k = metadata.max_seqlen + + attention_fn = self.attention_fn + if attention_fn is None: + from batchgen.attention.gqa import gqa_prefill_fa + + attention_fn = gqa_prefill_fa + attn_output, _ = attention_fn( + q=query, + k=key_for_attn, + v=value_for_attn, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=metadata.max_seqlen, + max_seqlen_k=max_seqlen_k, + sinks=self.sinks, + softmax_scale=self.softmax_scale, + sliding_window=self.sliding_window, + ) + return attn_output + + def _forward_paged_extend_prefill( + self, + *, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + metadata, + materialization, + ) -> torch.Tensor: + """Run prefix-hit suffix prefill over materialized GPU paged KV.""" + + from batchgen.attention.gqa import gqa_extend_fa + + layer_idx = int(self.layer_idx) + debug_sync = os.environ.get("BATCHGEN_PREFIX_DEBUG_SYNC", "0") == "1" + if debug_sync: + logging.info( + "[PREFIX_DEBUG] layer=%s begin q_shape=%s k_shape=%s v_shape=%s " + "cache_seqlens_minmax=(%s,%s) page_table_shape=%s", + layer_idx, + tuple(query.shape), + tuple(key.shape), + tuple(value.shape), + int(materialization.append_plan.cache_seqlens.min().item()), + int(materialization.append_plan.cache_seqlens.max().item()), + tuple(materialization.append_plan.page_table.shape), + ) + materialization.wait_for_layer(layer_idx) + if debug_sync: + torch.cuda.synchronize(query.device) + logging.info("[PREFIX_DEBUG] layer=%s prefix_load_ready", layer_idx) + materialization.manager.append_layer_prefill_suffix_tokens( + k_tensor=key, + v_tensor=value, + append_plan=materialization.append_plan, + layer_idx=layer_idx, + ) + if debug_sync: + torch.cuda.synchronize(query.device) + logging.info("[PREFIX_DEBUG] layer=%s suffix_append_done", layer_idx) + k_cache, v_cache, page_table = ( + materialization.manager.get_layer_kv_with_page_table(layer_idx) + ) + if v_cache is None: + raise RuntimeError("GQA paged prefix prefill requires V cache") + + cu_k = torch.nn.functional.pad( + torch.cumsum( + materialization.append_plan.cache_seqlens, + dim=0, + dtype=torch.int32, + ), + (1, 0), + ) + attn_output, _ = gqa_extend_fa( + q=query, + k_cache=k_cache, + v_cache=v_cache, + cache_seqlens=materialization.append_plan.cache_seqlens, + page_table=page_table, + cu_seqlens_q=metadata.cu_seqlens.to( + device=query.device, + dtype=torch.int32, + ), + cu_seqlens_k=cu_k, + max_seqlen_q=int(metadata.max_seqlen), + sinks=self.sinks, + softmax_scale=self.softmax_scale, + sliding_window=self.sliding_window, + ) + if debug_sync: + torch.cuda.synchronize(query.device) + logging.info("[PREFIX_DEBUG] layer=%s extend_attention_done", layer_idx) + return attn_output diff --git a/batchgen/batchgen_server.py b/batchgen/batchgen_server.py index b1f0835c1..678435bae 100644 --- a/batchgen/batchgen_server.py +++ b/batchgen/batchgen_server.py @@ -149,8 +149,22 @@ def allocate_host_kv_cache(self, host_kv_cache_size_gb: int): indexer, splitting the budget proportionally between primary and aux. """ from batchgen.kv_cache.dual_host_kv_coordinator import DualHostKVCoordinator + from batchgen.kv_cache.glm5_kv_coordinator import GLM5HostKVCoordinator - # DSA models: split budget into primary + auxiliary + # GLM-5 uses a model-specific group coordinator so prefix cache can + # manage primary/indexer pages independently. + glm5 = GLM5HostKVCoordinator.create_managers( + model_name=self.args.model, + host_kv_cache_size=int(host_kv_cache_size_gb * (1024**3)), + ) + if glm5 is not None: + primary_mgr, indexer_mgr = glm5 + logging.info( + "Allocated GLM-5 host KV cache: primary + indexer" + ) + return primary_mgr, indexer_mgr + + # Other DSA models keep the existing dual coordinator path. dual = DualHostKVCoordinator.create_managers( model_name=self.args.model, host_kv_cache_size=int(host_kv_cache_size_gb * (1024**3)), @@ -238,6 +252,8 @@ def spawn_workers(self): adaptive_chunk_max=getattr(self.args, 'adaptive_chunk_max', 65536), adaptive_chunk_ema_alpha=getattr(self.args, 'adaptive_chunk_ema_alpha', 0.1), adaptive_chunk_multiplier=getattr(self.args, 'adaptive_chunk_multiplier', 1.5), + enable_prefix_cache=getattr(self.args, 'enable_prefix_cache', False), + prefix_cache_debug_stats=getattr(self.args, 'prefix_cache_debug_stats', False), # Place holder local_rank=-1, @@ -257,6 +273,33 @@ def spawn_workers(self): daemon=True ) + def _initialize_prefix_cache_owner(self): + self.prefix_cache_runtime_config = None + self.prefix_cache_coordinator_owner = None + if not getattr(self.args, "enable_prefix_cache", False): + return + if self.args.host_kv_cache_size is None: + raise RuntimeError( + "--enable-prefix-cache requires --host-kv-cache-size" + ) + from batchgen.prefix_reuse.config import ( + build_prefix_cache_runtime_config, + create_host_prefix_cache_coordinator, + ) + + runtime_config = build_prefix_cache_runtime_config( + model_name=self.args.model, + kv_dtype=self.args.kv_dtype, + host_kv_cache_size_bytes=int(self.args.host_kv_cache_size * (1024**3)), + debug_stats=getattr(self.args, "prefix_cache_debug_stats", False), + ) + self.prefix_cache_runtime_config = runtime_config + self.prefix_cache_coordinator_owner = create_host_prefix_cache_coordinator( + core_engine_module=bg_lib, + runtime_config=runtime_config, + create_region=True, + ) + def start(self): """Start the TCP Server loop""" try: @@ -269,6 +312,7 @@ def start(self): # 1. Allocate KV & Load Model & Spawn Workers self.allocate_host_kv_cache(self.args.host_kv_cache_size) self.load_model_resources() + self._initialize_prefix_cache_owner() self.spawn_workers() # 2. Start TCP Listener @@ -664,6 +708,8 @@ def parse_args(): parser.add_argument("--nnodes", type=int, default=1) parser.add_argument("--node-rank", type=int, default=0) parser.add_argument("--world-size", type=int, default=1) + parser.add_argument("--enable-prefix-cache", action="store_true", default=False) + parser.add_argument("--prefix-cache-debug-stats", action="store_true", default=False) parser.add_argument( "--allow-model-download", action="store_true", @@ -697,4 +743,4 @@ def parse_args(): mp.set_start_method("spawn", force=True) args = parse_args() server = BatchGenServer(args) - server.start() \ No newline at end of file + server.start() diff --git a/batchgen/batchgen_worker.py b/batchgen/batchgen_worker.py index 35427a22b..f90abede9 100644 --- a/batchgen/batchgen_worker.py +++ b/batchgen/batchgen_worker.py @@ -92,6 +92,26 @@ def _check_repeating_pattern(token_ids: torch.Tensor, decoded_length: int, ) from batchgen.utils import config_torch_module_initializer from batchgen.config.model_name_utils import is_kimi_k25_backend_model +from batchgen.prefix_reuse.prefill import ( + PrefixCachePrefillLookup, + build_prefix_cache_prefill_inputs, + effective_prefix_shared_tokens, + estimate_prefix_cache_for_prefill, + lookup_prefix_cache_for_prefill, +) +from batchgen.prefill.prefix_reuse import split_prefix_reuse_plan_for_micro_batch +from batchgen.prefix_reuse.materialization import ( + PrefixMaterializationBundle, + materialize_single_group_lookup_results, +) +from batchgen.prefix_reuse.eviction import ( + commit_prefix_pages_with_capacity_retry, + evict_prefix_pages_for_host_allocation, +) +from batchgen.prefix_reuse.worker_commit import ( + build_sequence_prefix_commit_request, + retain_newly_committed_prefix_pages, +) from batchgen.models.glm.glm5.cuda_graph_policy import ( glm5_any_cuda_graph_requested_for_model, glm5_dsa_cuda_graph_requested_for_model, @@ -100,7 +120,10 @@ def _check_repeating_pattern(token_ids: torch.Tensor, decoded_length: int, glm5_segmented_cuda_graph_requested_for_model, glm5_whole_model_cuda_graph_requested_for_model, ) -from batchgen.kv_cache.gpu_paged_kv_manager import GPUPagedKVCacheManager +from batchgen.kv_cache.gpu_paged_kv_manager import ( + GPUPagedKVCacheManager, + GPUPagedKVConfig, +) from batchgen.models.engine_loader import core_engine from batchgen.worker.indexing import IndexLookupRequest, IndexManager from batchgen.worker.completion import CompletionContext, CompletionHandler @@ -114,6 +137,7 @@ def _check_repeating_pattern(token_ids: torch.Tensor, decoded_length: int, PrefillCandidate, PrefillScheduler, PrefillSelectionRequest, + PrefillWaveGateRequest, ) from batchgen.worker.host_rebalancer import HostKVRebalancer from batchgen.worker.boundary import ( @@ -146,6 +170,12 @@ def _check_repeating_pattern(token_ids: torch.Tensor, decoded_length: int, ) from batchgen.kv_cache.dual_kv_cache_coordinator import DualKVCacheCoordinator from batchgen.kv_cache.dual_host_kv_coordinator import DualAsyncKVTask, DualHostKVCoordinator +from batchgen.kv_cache.glm5_kv_coordinator import ( + GLM5AsyncKVTask, + GLM5GPUKVCoordinator, + GLM5HostKVCoordinator, + is_glm5_dual_kv_model, +) from batchgen.sequence import SequenceBatch, SequenceEntry, SequenceStatus, INITIAL_GPU_PAGE_BUFFER, EXTENSION_GPU_PAGE_BUFFER, DECISION_FREQUENCY_PAGES, configure_page_buffers from batchgen.prefill.prepack import ( prepack_sequences, @@ -274,6 +304,11 @@ class _DualKVLoadPointers: aux_page_counts: torch.Tensor +GroupedGPUKVCoordinator = (DualKVCacheCoordinator, GLM5GPUKVCoordinator) +GroupedHostKVCoordinator = (DualHostKVCoordinator, GLM5HostKVCoordinator) +GroupedAsyncKVTask = (DualAsyncKVTask, GLM5AsyncKVTask) + + class QueryBookBufferPool: """Pre-allocated contiguous buffers for query book tensors. @@ -470,6 +505,9 @@ class BatchGenWorkerArgs: weights_memfd_fd: int = -1 # Request pool: max QueryBook capacity (pre-allocated, metadata only) max_pool_size: int = 10240 # Default enables pool mode. 0 = legacy batch-FIFO. + # Host-side prefix cache. Detailed runtime config is derived inside the worker. + enable_prefix_cache: bool = False + prefix_cache_debug_stats: bool = False class BatchGenWorker: @@ -654,11 +692,13 @@ def __init__(self, args: BatchGenWorkerArgs): # 5. Initialize Host KV Cache Manager View (cudaHostRegister for Host KV) self.host_kv_cache_size = args.host_kv_cache_size self.global_host_kv_cache_size_gb = args.global_host_kv_cache_size_gb + self.host_paged_kv_worker_view_aux = None - # DSA models: create DualHostKVCoordinator with proportional budget split. - # Non-DSA models get a single-view worker below. + # GLM-5 uses a model-specific primary/indexer coordinator so prefix + # cache can manage each logical KV group independently. Other DSA models + # keep the existing DualHostKVCoordinator path. host_budget_bytes = int(args.global_host_kv_cache_size_gb * (1024**3)) - dual_host = DualHostKVCoordinator.from_budget( + glm5_host = GLM5HostKVCoordinator.from_budget( model_name=args.model_name, host_kv_cache_size=host_budget_bytes, core_engine_module=core_engine, @@ -667,11 +707,34 @@ def __init__(self, args: BatchGenWorkerArgs): memfd_fd=args.kv_memfd_fd if args.fast_init else -1, aux_memfd_fd=args.kv_aux_memfd_fd if args.fast_init else -1, ) - if dual_host is not None: - self.host_paged_kv_worker_view = dual_host - logging.info(f"Rank {self.rank}: Initializing DualHostKVCoordinator with parallel cudaHostRegister (local_rank={self.local_rank})") - dual_host.initialize(device_index=self.local_rank, create_region=False) - logging.info(f"Rank {self.rank}: DualHostKVCoordinator cudaHostRegister completed (local_rank={self.local_rank})") + dual_host = None + if glm5_host is None: + dual_host = DualHostKVCoordinator.from_budget( + model_name=args.model_name, + host_kv_cache_size=host_budget_bytes, + core_engine_module=core_engine, + enable_memfd=args.fast_init, + memfd_creator_pid=args.kv_memfd_pid if args.fast_init else -1, + memfd_fd=args.kv_memfd_fd if args.fast_init else -1, + aux_memfd_fd=args.kv_aux_memfd_fd if args.fast_init else -1, + ) + grouped_host = glm5_host if glm5_host is not None else dual_host + if grouped_host is not None: + self.host_paged_kv_worker_view = grouped_host + logging.info( + "Rank %s: Initializing %s with parallel cudaHostRegister " + "(local_rank=%s)", + self.rank, + type(grouped_host).__name__, + self.local_rank, + ) + grouped_host.initialize(device_index=self.local_rank, create_region=False) + logging.info( + "Rank %s: %s cudaHostRegister completed (local_rank=%s)", + self.rank, + type(grouped_host).__name__, + self.local_rank, + ) else: worker_kv_config = build_host_kv_config( model_name=args.model_name, @@ -760,9 +823,764 @@ def __init__(self, args: BatchGenWorkerArgs): self._response_queue = None # mp.Queue, set via set_response_queue() self._shutdown_requested = False self._max_pool_size = args.max_pool_size # 0 = legacy mode + self._final_response_completed_outputs: Dict[int, str] = {} + self.enable_prefix_cache = bool(args.enable_prefix_cache) + self.prefix_cache_debug_stats = bool(args.prefix_cache_debug_stats) + self.prefix_cache_runtime_config = None + self.prefix_cache_coordinator = None + self._prefix_prefill_lookup_by_local_idx = {} + self._prefix_cache_attachment_by_global_idx = {} + self._initialize_prefix_cache_worker(args) logging.info(f"Rank {self.rank}: BatchGenWorker __init__ completed.") + def _initialize_prefix_cache_worker( + self, args: BatchGenWorkerArgs + ) -> None: + if not self.enable_prefix_cache: + return + if args.host_kv_cache_size is None: + raise RuntimeError( + "Prefix cache worker requires resolved Host KV cache budget" + ) + + from batchgen.prefix_reuse.config import ( + build_prefix_cache_runtime_config, + create_host_prefix_cache_coordinator, + ) + + runtime_config = build_prefix_cache_runtime_config( + model_name=args.model_name, + kv_dtype=args.kv_dtype, + host_kv_cache_size_bytes=int(args.host_kv_cache_size * (1024**3)), + debug_stats=bool(args.prefix_cache_debug_stats), + ) + self.prefix_cache_runtime_config = runtime_config + self.prefix_cache_coordinator = create_host_prefix_cache_coordinator( + core_engine_module=core_engine, + runtime_config=runtime_config, + create_region=False, + ) + logging.info( + "Rank %s attached Host prefix cache: shm=%s groups=%d", + self.rank, + runtime_config.shm_name, + len(runtime_config.group_specs), + ) + + def _lookup_prefix_cache_for_prefill( + self, + *, + local_indices: Sequence[int], + input_ids_list: Sequence[torch.Tensor], + prompt_lengths: Sequence[int], + ): + if not self.enable_prefix_cache: + return None + if self.prefix_cache_coordinator is None: + raise RuntimeError( + "Prefix cache is enabled but coordinator is not attached" + ) + if self.prefix_cache_runtime_config is None: + raise RuntimeError( + "Prefix cache is enabled but runtime config is missing" + ) + + prompt_token_ids = [] + for input_ids, prompt_length in zip(input_ids_list, prompt_lengths): + prompt_token_ids.append( + [ + int(token_id) + for token_id in input_ids.reshape(-1)[: int(prompt_length)].tolist() + ] + ) + lookup = lookup_prefix_cache_for_prefill( + coordinator=self.prefix_cache_coordinator, + namespace_digest=self.prefix_cache_runtime_config.namespace_digest, + prompt_token_ids=prompt_token_ids, + ) + for local_idx, cached_tokens, result in zip( + local_indices, lookup.prefix_shared_tokens, lookup.lookup_results + ): + uuid = self._local_to_uuid_map.get(int(local_idx)) + if uuid is None: + raise RuntimeError( + f"Missing UUID for prefix-cache prefill local_idx={local_idx}" + ) + seq = self.global_batch.get_sequence(uuid) + if seq is None: + raise RuntimeError( + f"Missing sequence for prefix-cache prefill uuid={uuid[:8]}" + ) + raw_cached_tokens = int(result.common_cached_tokens) + cached_tokens = int(cached_tokens) + if raw_cached_tokens < 0 or raw_cached_tokens > int(seq.prompt_length): + raise RuntimeError( + f"Prefix cache returned invalid hit for {uuid[:8]}: " + f"cached={raw_cached_tokens}, prompt={seq.prompt_length}" + ) + if raw_cached_tokens % int(seq.PAGE_SIZE) != 0: + raise RuntimeError( + f"Prefix cache returned non-page-aligned hit for " + f"{uuid[:8]}: cached={raw_cached_tokens}, page_size={seq.PAGE_SIZE}" + ) + if cached_tokens < 0 or cached_tokens >= int(seq.prompt_length): + raise RuntimeError( + f"Prefix cache normalized invalid effective hit for " + f"{uuid[:8]}: cached={cached_tokens}, " + f"prompt={seq.prompt_length}" + ) + seq.prefix_shared_tokens = cached_tokens + seq.prefix_committed_tokens = cached_tokens + return lookup + + def _prefill_inputs_for_local_indices( + self, + local_indices: Sequence[int], + ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[int]]: + input_ids_list = [] + attention_mask_list = [] + seq_lengths = [] + + for query_idx in local_indices: + uuid = self._local_to_uuid_map[int(query_idx)] + seq = self.global_batch.get_sequence(uuid) + query_entry = self.query_book[int(query_idx)] + encoded = query_entry.encoded["input_ids"] + if encoded.data_ptr() != seq.input_ids.data_ptr(): + raise RuntimeError( + f"Rank {self.rank}: stale query_book input_ids binding for " + f"local_idx={query_idx} uuid={uuid[:8]} " + f"(query_book_ptr={encoded.data_ptr():#x}, " + f"seq_ptr={seq.input_ids.data_ptr():#x})" + ) + if query_entry.decoded_tokens.data_ptr() != seq.decoded_tokens.data_ptr(): + raise RuntimeError( + f"Rank {self.rank}: stale query_book decoded_tokens binding for " + f"local_idx={query_idx} uuid={uuid[:8]} " + f"(query_book_ptr={query_entry.decoded_tokens.data_ptr():#x}, " + f"seq_ptr={seq.decoded_tokens.data_ptr():#x})" + ) + + prompt_length = int(seq.prompt_length) + if encoded.size(-1) < prompt_length: + raise RuntimeError( + f"encoded prompt length {encoded.size(-1)} < " + f"seq.prompt_length {prompt_length} for " + f"query_idx={query_idx} uuid={uuid[:8]}" + ) + input_ids = encoded[:, :prompt_length] + attention_mask = torch.zeros_like(input_ids, dtype=torch.int64) + attention_mask[0, :prompt_length] = 1 + + input_ids_list.append(input_ids) + attention_mask_list.append(attention_mask) + seq_lengths.append(prompt_length) + + return input_ids_list, attention_mask_list, seq_lengths + + def _prefix_cache_lookup_for_prefill_batch( + self, + local_indices: Sequence[int], + ) -> PrefixCachePrefillLookup | None: + if not self.enable_prefix_cache: + return None + lookup_results = [] + prefix_shared_tokens = [] + for local_idx in local_indices: + result = self._prefix_prefill_lookup_by_local_idx.get(int(local_idx)) + if result is None: + raise RuntimeError( + f"Prefix cache enabled but missing prefill lookup for " + f"local_idx={local_idx}" + ) + lookup_results.append(result) + uuid = self._local_to_uuid_map[int(local_idx)] + seq = self.global_batch.get_sequence(uuid) + if seq is None: + raise RuntimeError( + f"Missing sequence for prefix-cache prefill uuid={uuid[:8]}" + ) + prefix_shared_tokens.append( + effective_prefix_shared_tokens( + raw_cached_tokens=int(result.common_cached_tokens), + prompt_length=int(seq.prompt_length), + ) + ) + return PrefixCachePrefillLookup( + lookup_results=tuple(lookup_results), + prefix_shared_tokens=tuple(prefix_shared_tokens), + ) + + def _host_page_ids_from_prefix_lookup_group( + self, + result: object, + *, + group_id: int, + ) -> List[int]: + if int(result.common_cached_tokens) <= 0: + return [] + spans = result.materialization_spans + if spans is None: + raise RuntimeError("Prefix lookup result has no materialization spans") + for span in spans: + if int(span.group_id) != int(group_id): + continue + return [ + int(page.page_id) + for page in span.pages + ] + raise RuntimeError( + f"Prefix lookup hit has no materialization span for group {group_id}" + ) + + def _prefix_cache_worker_views_by_group(self) -> Dict[int, object]: + host_view = self.core_engine.host_paged_kv_worker_view + if isinstance(self.host_paged_kv_worker_view, GroupedHostKVCoordinator): + return self.host_paged_kv_worker_view.views_by_group() + if hasattr(host_view, "views_by_group"): + return host_view.views_by_group() + views = {0: host_view} + aux_view = self.host_paged_kv_worker_view_aux + if aux_view is not None: + views[1] = aux_view + return views + + def _prefix_cache_raw_page_tokens_by_group(self) -> Dict[int, int]: + runtime_config = self.prefix_cache_runtime_config + if runtime_config is None: + return {} + return { + int(spec.group_id): int(spec.raw_page_tokens) + for spec in runtime_config.group_specs + } + + def _prefix_cache_required_group_ids(self) -> Set[int]: + runtime_config = self.prefix_cache_runtime_config + if runtime_config is None: + return set() + return { + int(spec.group_id) + for spec in runtime_config.group_specs + if spec.required_for_reuse + } + + def _prefix_cache_private_page_requirements_by_group( + self, + sequence_tokens: Sequence[int], + ) -> Dict[int, int]: + runtime_config = self.prefix_cache_runtime_config + if runtime_config is None: + return {} + required_pages_by_group: Dict[int, int] = {} + for spec in runtime_config.group_specs: + if not spec.required_for_reuse: + continue + group_id = int(spec.group_id) + raw_page_tokens = int(spec.raw_page_tokens) + compression_ratio = max(1, int(spec.compression_ratio)) + if raw_page_tokens <= 0: + raise RuntimeError( + f"Invalid prefix cache raw_page_tokens for group {group_id}: " + f"{raw_page_tokens}" + ) + storage_page_tokens = max(1, raw_page_tokens // compression_ratio) + pages = 0 + for raw_tokens in sequence_tokens: + raw_tokens = int(raw_tokens) + if raw_tokens <= 0: + continue + if compression_ratio == 1: + storage_tokens = raw_tokens + else: + storage_tokens = max(1, raw_tokens // compression_ratio) + pages += math.ceil(storage_tokens / storage_page_tokens) + required_pages_by_group[group_id] = pages + return required_pages_by_group + + def _ensure_prefix_cache_host_pages_for_allocation( + self, + *, + sequence_tokens: Sequence[int], + reason: str, + ) -> None: + if not self.enable_prefix_cache: + return + if self.prefix_cache_coordinator is None: + raise RuntimeError( + "Prefix cache is enabled but coordinator is not attached" + ) + if self.prefix_cache_runtime_config is None: + raise RuntimeError( + "Prefix cache is enabled but runtime config is missing" + ) + + worker_views_by_group = self._prefix_cache_worker_views_by_group() + local_required_pages_by_group = ( + self._prefix_cache_private_page_requirements_by_group(sequence_tokens) + ) + group_ids = sorted(worker_views_by_group) + required_pages_by_group = {group_id: 0 for group_id in group_ids} + if self.world_size > 1 and dist.is_initialized(): + node_id = self.rank // NUM_GPUS_PER_NODE + payload = torch.tensor( + [ + node_id, + *[ + int(local_required_pages_by_group.get(group_id, 0)) + for group_id in group_ids + ], + ], + dtype=torch.int64, + device=self.torch_device, + ) + gathered = [torch.zeros_like(payload) for _ in range(self.world_size)] + dist.all_gather(gathered, payload) + for item in gathered: + if int(item[0].item()) != node_id: + continue + for index, group_id in enumerate(group_ids, start=1): + required_pages_by_group[group_id] += int(item[index].item()) + else: + required_pages_by_group.update(local_required_pages_by_group) + + page_deficit_by_group: Dict[int, int] = {} + eviction_error = "" + if self.local_rank == 0: + try: + for group_id, required_pages in required_pages_by_group.items(): + worker_view = worker_views_by_group.get(group_id) + if worker_view is None: + raise RuntimeError( + "Missing Host KV worker view for prefix cache " + f"group {group_id}" + ) + free_pages = int(worker_view.get_stats().num_free_pages) + deficit = int(required_pages) - free_pages + if deficit > 0: + page_deficit_by_group[group_id] = deficit + + if page_deficit_by_group: + eviction = evict_prefix_pages_for_host_allocation( + core_engine_module=core_engine, + coordinator=self.prefix_cache_coordinator, + worker_views_by_group=worker_views_by_group, + page_deficit_by_group=page_deficit_by_group, + ) + if self.prefix_cache_debug_stats and self.rank == 0: + evicted_nodes = ( + 0 + if eviction.eviction_result is None + else int(eviction.eviction_result.evicted_nodes) + ) + protected_nodes = ( + 0 + if eviction.eviction_result is None + else int(eviction.eviction_result.protected_nodes) + ) + logging.info( + "Prefix cache allocation eviction: reason=%s " + "deficits=%s released=%s evicted_nodes=%s " + "protected_nodes=%s", + reason, + page_deficit_by_group, + eviction.released_pages_by_group, + evicted_nodes, + protected_nodes, + ) + + remaining_deficits: Dict[int, int] = {} + for group_id, required_pages in required_pages_by_group.items(): + worker_view = worker_views_by_group[group_id] + free_pages = int(worker_view.get_stats().num_free_pages) + deficit = int(required_pages) - free_pages + if deficit > 0: + remaining_deficits[group_id] = deficit + if remaining_deficits: + raise RuntimeError( + "Prefix cache eviction did not free enough Host KV " + f"pages for {reason}: remaining={remaining_deficits}" + ) + except Exception as exc: + eviction_error = str(exc) + logging.exception("Prefix cache allocation eviction failed") + + if self.world_size > 1 and dist.is_initialized(): + error_flag = torch.tensor( + [1 if eviction_error else 0], + dtype=torch.int64, + device=self.torch_device, + ) + dist.all_reduce(error_flag, op=dist.ReduceOp.MAX) + if int(error_flag.item()) != 0: + if eviction_error: + raise RuntimeError(eviction_error) + raise RuntimeError( + "Prefix cache allocation eviction failed on another rank" + ) + elif eviction_error: + raise RuntimeError(eviction_error) + + def _prefix_cache_gpu_managers_by_group( + self, + manager: object, + ) -> Dict[int, object]: + if isinstance(manager, GroupedGPUKVCoordinator): + return manager.managers_by_group() + if hasattr(manager, "managers_by_group"): + return manager.managers_by_group() + return {0: manager} + + def _attach_prefix_cache_lookup_pages( + self, + *, + local_indices: Sequence[int], + lookup: PrefixCachePrefillLookup, + ) -> None: + views_by_group = self._prefix_cache_worker_views_by_group() + for local_idx, result in zip(local_indices, lookup.lookup_results): + uuid = self._local_to_uuid_map[int(local_idx)] + seq = self.global_batch.get_sequence(uuid) + global_idx = int(seq.global_idx) + for group_id, worker_view in views_by_group.items(): + page_ids = self._host_page_ids_from_prefix_lookup_group( + result, + group_id=group_id, + ) + if page_ids: + worker_view.attach_shared_prefix_pages(global_idx, page_ids) + + attachment_handle = int(result.attachment_handle) + if attachment_handle: + self._prefix_cache_attachment_by_global_idx[global_idx] = ( + attachment_handle + ) + + def _materialize_prefix_cache_prefill( + self, + *, + lookup: PrefixCachePrefillLookup, + prefix_plan, + ) -> PrefixMaterializationBundle | None: + if lookup is None or not lookup.has_hit: + return None + + if self.gpu_paged_kv_cache_manager is not None: + self._destroy_gpu_paged_kv_cache(empty_cuda_cache=True) + + sequence_ids = [ + int(item.sequence_id) for item in prefix_plan.sequences + ] + prompt_lengths = [ + int(item.full_logical_context_length) + for item in prefix_plan.sequences + ] + prefix_shared_tokens = [ + int(item.prefix_shared_tokens) for item in prefix_plan.sequences + ] + if self.rank == 0 or any(tokens > 0 for tokens in prefix_shared_tokens): + page_size = int(SequenceEntry.PAGE_SIZE) + planned_pages = [ + (math.ceil(max(1, int(tokens)) / page_size)) + for tokens in prompt_lengths + ] + try: + free_mem_bytes, total_mem_bytes = torch.cuda.mem_get_info( + self.local_rank + ) + hbm_msg = ( + f"hbm_free_gb={free_mem_bytes / (1024**3):.2f} " + f"hbm_total_gb={total_mem_bytes / (1024**3):.2f}" + ) + except Exception: + hbm_msg = "hbm_free_gb=" + logging.info( + "Rank %s prefix materialization sizing: seq_ids=%s " + "prompt_lengths=%s shared_tokens=%s planned_pages=%s %s", + self.rank, + sequence_ids, + prompt_lengths, + prefix_shared_tokens, + planned_pages, + hbm_msg, + ) + manager, rolling_layers_by_group = ( + self._ensure_prefix_prefill_gpu_manager(prompt_lengths) + ) + host_views_by_group = self._prefix_cache_worker_views_by_group() + gpu_managers_by_group = self._prefix_cache_gpu_managers_by_group(manager) + raw_page_tokens_by_group = self._prefix_cache_raw_page_tokens_by_group() + required_group_ids = self._prefix_cache_required_group_ids() + missing_host_groups = required_group_ids - set(host_views_by_group) + if missing_host_groups: + raise RuntimeError( + "Missing Host KV worker views for required prefix cache " + f"groups: {sorted(missing_host_groups)}" + ) + missing_gpu_groups = required_group_ids - set(gpu_managers_by_group) + if missing_gpu_groups: + raise RuntimeError( + "Missing GPU KV managers for required prefix cache groups: " + f"{sorted(missing_gpu_groups)}" + ) + + by_group = {} + for group_id in sorted(required_group_ids): + gpu_group_manager = gpu_managers_by_group.get(group_id) + if gpu_group_manager is None: + raise RuntimeError( + f"Missing GPU KV manager for prefix cache group {group_id}" + ) + by_group[group_id] = materialize_single_group_lookup_results( + gpu_manager=gpu_group_manager, + host_worker_view=host_views_by_group[group_id], + lookup_results=lookup.lookup_results, + sequence_ids=sequence_ids, + prompt_lengths=prompt_lengths, + group_id=group_id, + prefix_shared_tokens=prefix_shared_tokens, + raw_page_tokens=raw_page_tokens_by_group.get(group_id), + prefix_cache_coordinator=self.prefix_cache_coordinator, + rolling_logical_layer_count=( + rolling_layers_by_group.get(group_id) + ), + ) + + return PrefixMaterializationBundle(by_group_id=by_group) + + def _release_prefix_cache_attachments_for_global_ids( + self, + global_sequence_ids: Sequence[int], + ) -> None: + if not self.enable_prefix_cache or self.prefix_cache_coordinator is None: + return + handles = [] + for global_idx in global_sequence_ids: + handle = self._prefix_cache_attachment_by_global_idx.pop( + int(global_idx), + 0, + ) + if handle: + handles.append(int(handle)) + for handle in dict.fromkeys(handles): + self.prefix_cache_coordinator.release_attachment(handle) + + def _commit_prefix_cache_for_sequences( + self, + uuids: Sequence[str], + *, + include_new_decode_tokens: bool, + reason: str, + ) -> None: + if not self.enable_prefix_cache: + return + if self.prefix_cache_coordinator is None: + raise RuntimeError( + "Prefix cache is enabled but coordinator is not attached" + ) + if self.prefix_cache_runtime_config is None: + raise RuntimeError( + "Prefix cache is enabled but runtime config is missing" + ) + + worker_views_by_group = self._prefix_cache_worker_views_by_group() + + total_start = time.perf_counter() + build_seconds = 0.0 + commit_seconds = 0.0 + retain_seconds = 0.0 + planned_count = 0 + committed_count = 0 + inserted_nodes = 0 + existing_nodes = 0 + evicted_nodes = 0 + + for uuid in uuids: + if uuid not in self._uuid_to_local_map: + continue + seq = self.global_batch.get_sequence(uuid) + if seq is None: + continue + build_start = time.perf_counter() + request_pair = build_sequence_prefix_commit_request( + core_engine_module=core_engine, + runtime_config=self.prefix_cache_runtime_config, + worker_views_by_group=worker_views_by_group, + seq=seq, + include_new_decode_tokens=include_new_decode_tokens, + ) + build_seconds += time.perf_counter() - build_start + if request_pair is None: + continue + planned_count += 1 + request, commit_tokens = request_pair + commit_start = time.perf_counter() + retry_result = commit_prefix_pages_with_capacity_retry( + request=request, + coordinator=self.prefix_cache_coordinator, + worker_views_by_group=worker_views_by_group, + ) + commit_seconds += time.perf_counter() - commit_start + result = retry_result.commit_result + committed_count += 1 + inserted_nodes += int(result.inserted_nodes) + existing_nodes += int(result.existing_nodes) + if retry_result.eviction_result is not None: + evicted_nodes += int(retry_result.eviction_result.evicted_nodes) + existing_tokens = ( + int(result.existing_nodes) + * int(request.publish_boundary_tokens) + ) + retain_start_tokens = max( + int(seq.prefix_committed_tokens), + int(seq.prefix_shared_tokens), + existing_tokens, + ) + if ( + int(result.inserted_nodes) > 0 + and int(commit_tokens) > retain_start_tokens + ): + retain_start = time.perf_counter() + retain_newly_committed_prefix_pages( + runtime_config=self.prefix_cache_runtime_config, + worker_views_by_group=worker_views_by_group, + sequence_id=int(seq.global_idx), + previous_committed_tokens=retain_start_tokens, + commit_tokens=int(commit_tokens), + page_ids_by_group=request.page_ids_by_group, + ) + retain_seconds += time.perf_counter() - retain_start + seq.prefix_committed_tokens = int(commit_tokens) + if self.prefix_cache_debug_stats and self.rank == 0: + logging.info( + "Prefix cache %s commit: seq=%s gid=%s tokens=%s " + "inserted=%s existing=%s evicted=%s", + reason, + uuid[:8], + seq.global_idx, + result.committed_tokens, + result.inserted_nodes, + result.existing_nodes, + 0 + if retry_result.eviction_result is None + else retry_result.eviction_result.evicted_nodes, + ) + + total_seconds = time.perf_counter() - total_start + if planned_count > 0 and ( + self.prefix_cache_debug_stats or total_seconds >= 1.0 + ): + logging.info( + "Prefix cache %s commit timings: rank=%s uuids=%d " + "planned=%d committed=%d inserted=%d existing=%d " + "evicted=%d total_s=%.3f build_s=%.3f " + "coordinator_s=%.3f retain_s=%.3f", + reason, + self.rank, + len(uuids), + planned_count, + committed_count, + inserted_nodes, + existing_nodes, + evicted_nodes, + total_seconds, + build_seconds, + commit_seconds, + retain_seconds, + ) + + def _commit_prefix_cache_prompt_pages( + self, + uuids: Sequence[str], + ) -> None: + self._commit_prefix_cache_for_sequences( + uuids, + include_new_decode_tokens=False, + reason="prompt", + ) + + def _commit_prefix_cache_completed_pages( + self, + uuids: Sequence[str], + ) -> None: + # Only prompt/prefill pages are published to prefix cache. + # Generated decode KV remains private runtime state. + return + + def _estimate_prefix_cache_for_prefill( + self, + *, + input_ids_list: Sequence[torch.Tensor], + prompt_lengths: Sequence[int], + ): + if not self.enable_prefix_cache: + return None + if self.prefix_cache_coordinator is None: + raise RuntimeError( + "Prefix cache is enabled but coordinator is not attached" + ) + if self.prefix_cache_runtime_config is None: + raise RuntimeError( + "Prefix cache is enabled but runtime config is missing" + ) + + prompt_token_ids = [] + for input_ids, prompt_length in zip(input_ids_list, prompt_lengths): + prompt_token_ids.append( + [ + int(token_id) + for token_id in input_ids.reshape(-1)[: int(prompt_length)].tolist() + ] + ) + estimate = estimate_prefix_cache_for_prefill( + coordinator=self.prefix_cache_coordinator, + namespace_digest=self.prefix_cache_runtime_config.namespace_digest, + prompt_token_ids=prompt_token_ids, + ) + if self.rank == 0 and self.prefix_cache_debug_stats: + hit_count = sum( + 1 for tokens in estimate.prefix_shared_tokens if tokens > 0 + ) + logging.info( + "Prefix cache estimate: %d/%d requests have reusable prefix " + "(forced miss until Host KV alias/copy is implemented)", + hit_count, + len(estimate.prefix_shared_tokens), + ) + return estimate + + def _build_prefix_reuse_prepack_inputs( + self, + *, + local_indices: Sequence[int], + input_ids_list: Sequence[torch.Tensor], + prompt_lengths: Sequence[int], + lookup, + ): + if lookup is None: + return None + + sequence_ids = [] + for local_idx in local_indices: + uuid = self._local_to_uuid_map.get(int(local_idx)) + if uuid is None: + raise RuntimeError( + f"Missing UUID for prefix-cache prepack local_idx={local_idx}" + ) + seq = self.global_batch.get_sequence(uuid) + if seq is None: + raise RuntimeError( + f"Missing sequence for prefix-cache prepack uuid={uuid[:8]}" + ) + sequence_ids.append(int(seq.global_idx)) + return build_prefix_cache_prefill_inputs( + local_indices=local_indices, + sequence_ids=sequence_ids, + input_ids=input_ids_list, + prompt_lengths=prompt_lengths, + lookup=lookup, + ) + def Init(self, max_input_length, max_decoding_length, num_queries, max_context_length=None): """ Initialize/reconfigure for a new batch. @@ -876,7 +1694,9 @@ def _initialize_gpu_kv_manager_fixed_size(self) -> GPUPagedKVCacheManager: Called once at the start of decoding. For DSA models, splits the memory budget between primary (MLA) and - auxiliary (indexer) caches, wrapping both in a DualKVCacheCoordinator. + auxiliary/indexer caches. GLM-5 uses GLM5GPUKVCoordinator so prefix + cache can address each logical KV group independently; other DSA + models keep DualKVCacheCoordinator. """ from batchgen.kv_cache.host_kv_mananger_config import ( build_gpu_kv_config_fixed_size, @@ -932,14 +1752,17 @@ def _initialize_gpu_kv_manager_fixed_size(self) -> GPUPagedKVCacheManager: primary.initialize() auxiliary = GPUPagedKVCacheManager(config=aux_config, device=self.local_rank) auxiliary.initialize() - manager = DualKVCacheCoordinator(primary, auxiliary) + if is_glm5_dual_kv_model(self.huggingface_ckpt_name): + manager = GLM5GPUKVCoordinator(primary, auxiliary) + else: + manager = DualKVCacheCoordinator(primary, auxiliary) self._bind_gpu_paged_kv_manager(manager) if self.rank == 0: primary_gb = (primary_bytes_per_page * num_pages) / (1024 ** 3) aux_gb = (aux_bytes_per_page * num_pages) / (1024 ** 3) logging.info( - f"[GPU-KV] DualKVCacheCoordinator initialized: " + f"[GPU-KV] {type(manager).__name__} initialized: " f"{num_pages} pages, primary={primary_gb:.2f} GB (dim={primary_profile.k_head_dim}), " f"auxiliary={aux_gb:.2f} GB (dim={aux_profile.k_head_dim})" ) @@ -1465,16 +2288,22 @@ def _build_local_query_book_for_admitted(self, uuids: List[str]) -> None: continue self._bind_local_sequence_to_query_book(uuid) - def _report_completion(self, uuid: str, gathered_text: str = None) -> None: + def _report_completion( + self, + uuid: str, + gathered_text: str = None, + cached_tokens: Optional[int] = None, + ) -> None: """Report a single sequence completion to the response queue. Also frees the QueryBook buffer slot so it can be reused by new admissions. Args: uuid: Sequence UUID. - gathered_text: Pre-gathered decoded text from _gather_completed_tokens. + gathered_text: Pre-gathered decoded text from _gather_completed_outputs. If provided, uses this instead of reading from local decoded_tokens (which may be empty on rank 0 for sequences owned by other ranks). + cached_tokens: Prefix-cache hit tokens gathered from the owner rank. """ seq = self.global_batch.get_sequence(uuid) if seq is None: @@ -1514,6 +2343,11 @@ def _report_completion(self, uuid: str, gathered_text: str = None) -> None: text = self.tokenizer.decode(token_ids) except Exception: text = "" + reported_cached_tokens = ( + int(cached_tokens) + if cached_tokens is not None + else int(getattr(seq, "prefix_shared_tokens", 0)) + ) self._response_queue.put({ "type": "completion", "request_id": uuid, @@ -1522,47 +2356,81 @@ def _report_completion(self, uuid: str, gathered_text: str = None) -> None: "text": text, "prompt_length": seq.prompt_length, "decoded_length": seq.decoded_length, + "cached_tokens": reported_cached_tokens, "finish_reason": self._get_finish_reason(seq), }) - def _gather_completed_tokens(self, completed_uuids: List[str]) -> dict: - """Gather decoded tokens from owning ranks for completed sequences. + def _gather_completed_outputs(self, completed_uuids: List[str]) -> dict: + """Gather completion outputs from owning ranks for completed sequences. - Each rank writes decoded tokens only for sequences it owns. This method - uses all_gather_object to collect tokens from all ranks so rank 0 can - report them correctly. + Each rank writes decoded tokens and prefix-cache metadata only for + sequences it owns. This method uses all_gather_object to collect that + owner-rank state so rank 0 reports correct text and usage. Returns: - Dict mapping uuid -> decoded text string. + Dict mapping uuid -> {"text": str, "cached_tokens": int}. """ if not completed_uuids: return {} - # Each rank provides tokens for its locally-owned completed sequences - my_tokens = {} + # Each rank provides outputs for its locally-owned completed sequences. + my_outputs = {} for uuid in completed_uuids: if uuid in self._uuid_to_local_map: local_idx = self._uuid_to_local_map[uuid] seq = self.global_batch.get_sequence(uuid) if seq is not None and local_idx in self.query_book: - token_ids = self.query_book[local_idx].decoded_tokens[0, :seq.decoded_length].tolist() + token_ids = self.query_book[local_idx].decoded_tokens[ + 0, :seq.decoded_length + ].tolist() try: text = self.tokenizer.decode(token_ids) except Exception: text = "" - my_tokens[uuid] = text + my_outputs[uuid] = { + "text": text, + "cached_tokens": int( + getattr(seq, "prefix_shared_tokens", 0) + ), + } # All ranks participate in gather - all_tokens = [None] * self.world_size - dist.all_gather_object(all_tokens, my_tokens) + all_outputs = [None] * self.world_size + dist.all_gather_object(all_outputs, my_outputs) # Merge: each uuid is owned by exactly one rank merged = {} - for rank_tokens in all_tokens: - if rank_tokens: - merged.update(rank_tokens) + for rank_outputs in all_outputs: + if rank_outputs: + merged.update(rank_outputs) return merged + def _record_completed_outputs_for_final_response( + self, + completed_uuids: Sequence[str], + gathered_outputs: dict, + ) -> None: + """Keep owner-rank decoded text for legacy synchronous responses. + + Completed sequences release their local query slots immediately so the + final result gather can no longer read their decoded_tokens from + query_book. Rank 0 stores the already-gathered text for this batch only. + """ + if self.rank != 0 or not gathered_outputs: + return + + for uuid in completed_uuids: + seq = self.global_batch.get_sequence(uuid) + if seq is None: + continue + output = gathered_outputs.get(uuid) + if output is None: + continue + text = output.get("text", "") + self._final_response_completed_outputs[seq.global_idx] = ( + text if isinstance(text, str) else str(text) + ) + # ============ End Request Pool Methods ============ def _build_sampling_tensors(self, batch_sequences: list) -> tuple: @@ -1698,6 +2566,156 @@ def _log_decode_timing(self): except ImportError: pass # Not GPT-OSS or module not available + def _reduce_runtime_metrics( + self, + values: Sequence[float], + op: "dist.ReduceOp", + ) -> List[float]: + if not dist.is_initialized(): + return [float(value) for value in values] + tensor = torch.tensor( + [float(value) for value in values], + dtype=torch.float64, + device=self.torch_device, + ) + dist.all_reduce(tensor, op=op) + return [float(value) for value in tensor.cpu().tolist()] + + def _log_prefill_phase_metrics( + self, + *, + prefill_uuids: Sequence[str], + local_prefill_indices: Sequence[int], + config_s: float, + prefill_s: float, + total_s: float, + ) -> None: + local_sequence_count = 0 + local_prompt_tokens = 0 + local_cached_tokens = 0 + local_requests_with_cache = 0 + for local_idx in local_prefill_indices: + uuid = self._local_to_uuid_map.get(int(local_idx)) + if uuid is None: + continue + seq = self.global_batch.get_sequence(uuid) + if seq is None: + continue + prompt_tokens = int(seq.prompt_length) + cached_tokens = int(getattr(seq, "prefix_shared_tokens", 0)) + local_sequence_count += 1 + local_prompt_tokens += prompt_tokens + local_cached_tokens += cached_tokens + if cached_tokens > 0: + local_requests_with_cache += 1 + + ( + global_sequence_count, + global_prompt_tokens, + global_cached_tokens, + global_requests_with_cache, + ) = self._reduce_runtime_metrics( + [ + local_sequence_count, + local_prompt_tokens, + local_cached_tokens, + local_requests_with_cache, + ], + dist.ReduceOp.SUM, + ) + ( + max_config_s, + max_prefill_s, + max_total_s, + ) = self._reduce_runtime_metrics( + [config_s, prefill_s, total_s], + dist.ReduceOp.MAX, + ) + + if self.rank != 0: + return + token_hit_rate = ( + global_cached_tokens / global_prompt_tokens + if global_prompt_tokens > 0 + else 0.0 + ) + request_hit_rate = ( + global_requests_with_cache / global_sequence_count + if global_sequence_count > 0 + else 0.0 + ) + prefill_tps = ( + (global_prompt_tokens - global_cached_tokens) / max_prefill_s + if max_prefill_s > 0 + else 0.0 + ) + logging.info( + "[PREFILL_METRICS] completed sequences=%d selected=%d " + "prompt_tokens=%d cached_tokens=%d token_hit_rate=%.2f%% " + "request_hit_rate=%.2f%% config_s=%.3f prefill_s=%.3f " + "total_s=%.3f effective_prefill_tps=%.1f", + int(global_sequence_count), + len(prefill_uuids), + int(global_prompt_tokens), + int(global_cached_tokens), + token_hit_rate * 100.0, + request_hit_rate * 100.0, + max_config_s, + max_prefill_s, + max_total_s, + prefill_tps, + ) + + def _log_decode_phase_metrics( + self, + *, + active_start: int, + local_generated_tokens: int, + elapsed_s: float, + iteration_delta: int, + boundary_delta: int, + forward_ms_delta: float, + boundary_ms_delta: float, + ) -> None: + (global_generated_tokens,) = self._reduce_runtime_metrics( + [local_generated_tokens], + dist.ReduceOp.SUM, + ) + (max_elapsed_s,) = self._reduce_runtime_metrics( + [elapsed_s], + dist.ReduceOp.MAX, + ) + if self.rank != 0: + return + decode_tps = ( + global_generated_tokens / max_elapsed_s + if max_elapsed_s > 0 + else 0.0 + ) + avg_forward_ms = ( + forward_ms_delta / iteration_delta + if iteration_delta > 0 + else 0.0 + ) + avg_boundary_ms = ( + boundary_ms_delta / boundary_delta + if boundary_delta > 0 + else 0.0 + ) + logging.info( + "[DECODE_METRICS] completed active_start=%d generated_tokens=%d " + "elapsed_s=%.3f decode_tps=%.1f iterations=%d boundaries=%d " + "avg_forward_ms=%.3f avg_boundary_ms=%.3f", + active_start, + int(global_generated_tokens), + max_elapsed_s, + decode_tps, + iteration_delta, + boundary_delta, + avg_forward_ms, + avg_boundary_ms, + ) + def set_watchdog(self, watchdog) -> None: """ Set the watchdog for stuck detection during inference. @@ -2380,7 +3398,7 @@ def _append_decode_kv_to_host_aux_async( Mirrors _append_decode_kv_to_host_async but uses the auxiliary host worker view. Shares the same pending task list for unified flushing. """ - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view = self.host_paged_kv_worker_view_aux if aux_view is None or not batch: return @@ -2515,7 +3533,7 @@ def _initialize_core_components(self, num_queries: int) -> None: [m.value for m in self._cuda_graph_adapter.advertised_modes()], ) - if isinstance(self.host_paged_kv_worker_view, DualHostKVCoordinator): + if isinstance(self.host_paged_kv_worker_view, GroupedHostKVCoordinator): self.core_engine.host_paged_kv_worker_view = self.host_paged_kv_worker_view.primary self.host_paged_kv_worker_view_aux = self.host_paged_kv_worker_view.auxiliary else: @@ -2651,12 +3669,12 @@ def _compute_host_kv_sequence_tokens(self, sequence_ids: List[int]) -> List[int] def _bind_gpu_paged_kv_manager(self, manager) -> None: """Bind GPU KV manager to both worker and core_engine. - If manager is a DualKVCacheCoordinator, the primary manager is bound - to existing gpu_paged_kv_manager slots and the auxiliary (indexer) is - bound to gpu_paged_kv_manager_aux slots. + If manager owns multiple logical KV groups, the primary manager is + bound to existing gpu_paged_kv_manager slots and the indexer/auxiliary + manager is bound to gpu_paged_kv_manager_aux slots. """ self.gpu_paged_kv_cache_manager = manager - if isinstance(manager, DualKVCacheCoordinator): + if isinstance(manager, GroupedGPUKVCoordinator): if hasattr(self.core_engine, "gpu_paged_kv_manager"): self.core_engine.gpu_paged_kv_manager = manager.primary if hasattr(self.core_engine, "gpu_paged_kv_manager_aux"): @@ -2665,10 +3683,17 @@ def _bind_gpu_paged_kv_manager(self, manager) -> None: if hasattr(self.core_engine, "gpu_paged_kv_manager"): self.core_engine.gpu_paged_kv_manager = manager + def _unbind_gpu_paged_kv_manager(self) -> None: + """Clear stale GPU KV manager references after destroying the manager.""" + self.gpu_paged_kv_cache_manager = None + Attn_Wrapper.gpu_paged_kv_manager = None + AttnWrapperBase.gpu_paged_kv_manager = None + AttnWrapperBase.gpu_paged_kv_manager_aux = None + def _get_cuda_graph_gpu_manager(self): """Return the GPU KV manager object to use for CUDA graph setup.""" manager = self.gpu_paged_kv_cache_manager - if isinstance(manager, DualKVCacheCoordinator): + if isinstance(manager, GroupedGPUKVCoordinator): return manager if manager is not None: return manager @@ -2743,15 +3768,18 @@ def _make_gpu_kv_manager_request( ) -> GpuKvManagerRequest: """Snapshot the worker state `plan_gpu_kv_manager` consumes.""" manager = self.gpu_paged_kv_cache_manager + manager_initialized = ( + manager is not None and bool(getattr(manager, "is_initialized", False)) + ) current_pages = ( getattr(getattr(manager, "config", None), "num_pages", 0) - if manager is not None + if manager_initialized else 0 ) return GpuKvManagerRequest( model_name=self.huggingface_ckpt_name, sequence_tokens=tuple(int(t) for t in sequence_tokens), - has_manager=manager is not None, + has_manager=manager_initialized, current_num_pages=int(current_pages), capacity=self._make_page_table_capacity_request(sequence_tokens), ) @@ -2795,14 +3823,17 @@ def _apply_gpu_kv_manager_plan( config=plan.aux_config, device=self.local_rank, ) - manager = DualKVCacheCoordinator(primary, auxiliary) + if is_glm5_dual_kv_model(self.huggingface_ckpt_name): + manager = GLM5GPUKVCoordinator(primary, auxiliary) + else: + manager = DualKVCacheCoordinator(primary, auxiliary) manager.initialize() self._bind_gpu_paged_kv_manager(manager) logging.info( - "Rank %s initialized DualKVCacheCoordinator on %s: " + "Rank %s initialized %s on %s: " "primary=%d pages (dim=%d), auxiliary=%d pages (dim=%d)", - self.rank, self.local_rank, + self.rank, type(manager).__name__, self.local_rank, plan.primary_config.num_pages, plan.primary_config.k_head_dim, plan.aux_config.num_pages, plan.aux_config.k_head_dim, ) @@ -2832,6 +3863,96 @@ def _ensure_gpu_paged_kv_manager(self, sequence_tokens: Sequence[int]) -> GPUPag ) return self._apply_gpu_kv_manager_plan(plan) + def _rolling_prefix_prefill_config( + self, + config: GPUPagedKVConfig, + sequence_tokens: Sequence[int], + ) -> tuple[GPUPagedKVConfig, Optional[int]]: + """Return a two-slot layer-mapped config for prefix-hit prefill. + + The temporary prefix materialization only needs the current attention + layer and the next prefetched layer resident on GPU. This applies to + GQA/MHA K+V caches and MLA K-only compressed caches alike. + """ + page_size_tokens = int(config.page_size_tokens) + fa_page_size_tokens = self._fa_paged_kv_page_size_tokens(page_size_tokens) + num_pages = sum( + math.ceil(max(1, int(tokens)) / fa_page_size_tokens) + for tokens in sequence_tokens + ) + logical_layer_count = int(config.num_layers) + physical_layer_count = min(2, logical_layer_count) + layer_mapping = tuple( + layer_idx % physical_layer_count + for layer_idx in range(logical_layer_count) + ) + return ( + replace( + config, + num_pages=max(1, int(num_pages)), + page_size_tokens=fa_page_size_tokens, + num_layers=physical_layer_count, + logical_to_physical_layer=layer_mapping, + ), + logical_layer_count, + ) + + def _fa_paged_kv_page_size_tokens(self, page_size_tokens: int) -> int: + """Return a GPU page size accepted by FlashAttention paged KV.""" + + page_size_tokens = int(page_size_tokens) + fa_block = 256 + if page_size_tokens >= fa_block and page_size_tokens % fa_block == 0: + return page_size_tokens + return math.ceil(max(page_size_tokens, fa_block) / fa_block) * fa_block + + def _ensure_prefix_prefill_gpu_manager( + self, + sequence_tokens: Sequence[int], + ) -> tuple[object, Dict[int, int]]: + """Create the temporary GPU KV manager used by prefix-hit prefill.""" + + plan = KVCacheManager.plan_gpu_kv_manager( + self._make_gpu_kv_manager_request(sequence_tokens) + ) + if plan.aux_config is not None: + return self._apply_gpu_kv_manager_plan(plan), {} + + config, logical_layer_count = self._rolling_prefix_prefill_config( + plan.primary_config, + sequence_tokens, + ) + if logical_layer_count is None: + return self._apply_gpu_kv_manager_plan(plan), {} + + manager = self.gpu_paged_kv_cache_manager + current_pages = ( + getattr(getattr(manager, "config", None), "num_pages", 0) + if manager is not None + else 0 + ) + if manager is not None: + manager.destroy() + + logging.info( + "Rank %s creating rolling prefix prefill GPUPagedKVCacheManager " + "on %s: current pages=%d, required pages=%d, " + "logical_layers=%d, physical_layers=%d", + self.rank, + self.local_rank, + current_pages, + config.num_pages, + logical_layer_count, + config.num_layers, + ) + manager = GPUPagedKVCacheManager( + config=config, + device=self.local_rank, + ) + manager.initialize() + self._bind_gpu_paged_kv_manager(manager) + return manager, {0: logical_layer_count} + def _prepare_gpu_paged_kv_cache(self, local_sequence_ids: List[int]) -> None: """Allocate GPU KV pages and load host-resident KV for the batch.""" if not local_sequence_ids: @@ -2865,10 +3986,10 @@ def _launch_aux_host_kv_load(self, sequence_tensor: torch.Tensor): effects. The returned task must be .wait()'d before the first decode step that consumes the aux cache. """ - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view = self.host_paged_kv_worker_view_aux if aux_view is None: return None - if not isinstance(self.gpu_paged_kv_cache_manager, DualKVCacheCoordinator): + if not isinstance(self.gpu_paged_kv_cache_manager, GroupedGPUKVCoordinator): return None aux_mgr = self.gpu_paged_kv_cache_manager.auxiliary k_ptrs_aux, v_ptrs_aux = aux_mgr.get_padded_3d_page_pointers() @@ -2882,12 +4003,12 @@ def _launch_aux_host_kv_load(self, sequence_tensor: torch.Tensor): def _prepare_dual_kv_load_pointers( self, - gpu_manager: DualKVCacheCoordinator, + gpu_manager, new_global_ids: List[int], existing_global_ids: Optional[List[int]] = None, ) -> _DualKVLoadPointers: - if not isinstance(gpu_manager, DualKVCacheCoordinator): - raise RuntimeError("DSA dual KV load requires DualKVCacheCoordinator") + if not isinstance(gpu_manager, GroupedGPUKVCoordinator): + raise RuntimeError("Grouped KV load requires a grouped GPU KV coordinator") if not new_global_ids: raise ValueError("_prepare_dual_kv_load_pointers requires non-empty sequence ids") @@ -2928,10 +4049,10 @@ def _prepare_dual_kv_load_pointers( else: gpu_manager.clear_page_table() - def _launch_dual_host_kv_load(self, pointers: _DualKVLoadPointers) -> DualAsyncKVTask: + def _launch_dual_host_kv_load(self, pointers: _DualKVLoadPointers): host_view = self.host_paged_kv_worker_view - if not isinstance(host_view, DualHostKVCoordinator): - raise RuntimeError("DSA dual KV load requires DualHostKVCoordinator") + if not isinstance(host_view, GroupedHostKVCoordinator): + raise RuntimeError("Grouped KV load requires a grouped Host KV coordinator") return host_view.async_load_layer_paged_kv_to_device_dual( sequence_ids=pointers.sequence_tensor, primary_active_page_counts=pointers.primary_page_counts, @@ -2980,7 +4101,7 @@ def _load_host_kv_to_gpu( f"{len(global_sequence_ids)} sequences..." ) - if isinstance(manager, DualKVCacheCoordinator): + if isinstance(manager, GroupedGPUKVCoordinator): pointers = self._prepare_dual_kv_load_pointers(manager, global_sequence_ids) load_task = self._launch_dual_host_kv_load(pointers) else: @@ -3062,8 +4183,11 @@ def _destroy_gpu_paged_kv_cache(self, *, empty_cuda_cache: bool = False) -> None f"First 5: {seqs_with_gpu_alloc[:5]}" ) + if empty_cuda_cache: + torch.cuda.synchronize(self.torch_device) manager.destroy(empty_cuda_cache=empty_cuda_cache) - + self._unbind_gpu_paged_kv_manager() + # FIX Bug 2: Clear tracking set when GPU KV is destroyed self._sequences_with_gpu_kv.clear() @@ -3537,7 +4661,7 @@ def _execute_single_kv_migration(self, uuid: str, from_rank: int, to_rank: int) # mirrored explicitly below — the coordinator does not implement # read/write_sequence_kv_to_cpu, so go direct on primary and aux. worker_view = self.core_engine.host_paged_kv_worker_view - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view = self.host_paged_kv_worker_view_aux if self.rank == from_rank: # ===== SOURCE RANK: Read host KV directly to CPU, send via Gloo ===== @@ -3917,6 +5041,7 @@ def process_new_batch( logging.info( f"Rank {self.rank}: Processing global batch of {len(global_prompts)} sequences" ) + self._final_response_completed_outputs = {} # Step 1: Initialize global batch self.global_batch = SequenceBatch() @@ -4202,6 +5327,105 @@ def _sync_decode_uuids_tensor( self._make_sync_context(), decode_uuids ) + def _handle_completed_decode_uuids( + self, + completed_uuids: Sequence[str], + ) -> None: + if not completed_uuids: + return + + completed_list = sorted( + dict.fromkeys(completed_uuids), + key=lambda uuid: ( + self.global_batch.get_sequence(uuid).global_idx + if self.global_batch.get_sequence(uuid) is not None + else 2**63 - 1 + ), + ) + self._submit_completed_to_incremental_writer(completed_list) + gathered_outputs = self._gather_completed_outputs(completed_list) + self._record_completed_outputs_for_final_response( + completed_list, + gathered_outputs, + ) + + if self.enable_prefix_cache: + self._wait_pending_kv_append_tasks(sync_distributed_errors=True) + + my_completed = [ + uuid for uuid in completed_list if uuid in self._uuid_to_local_map + ] + host_kv_stats_before = None + if my_completed and self.host_paged_kv_worker_view is not None: + host_kv_stats_before = self.host_paged_kv_worker_view.get_stats() + if my_completed: + if self.rank == 0 or self.prefix_cache_debug_stats: + global_ids = [ + self.global_batch.get_sequence(uuid).global_idx + for uuid in my_completed + if self.global_batch.get_sequence(uuid) is not None + ] + before_msg = "" + if host_kv_stats_before is not None: + before_msg = ( + f" before_used={host_kv_stats_before.num_used_pages}" + f" before_free={host_kv_stats_before.num_free_pages}" + f" total={host_kv_stats_before.num_total_pages}" + ) + logging.info( + "[DECODE_RELEASE] Rank %s: releasing completed host KV " + "local=%d global_ids_sample=%s%s", + self.rank, + len(my_completed), + global_ids[:8], + before_msg, + ) + self._commit_prefix_cache_completed_pages(my_completed) + gpu_allocated = [ + uuid for uuid in my_completed if uuid 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) + if self.rank == 0 or self.prefix_cache_debug_stats: + stats_after = self.host_paged_kv_worker_view.get_stats() + logging.info( + "[DECODE_RELEASE] Rank %s: completed host KV release " + "local=%d after_used=%d after_free=%d total=%d", + self.rank, + len(my_completed), + stats_after.num_used_pages, + stats_after.num_free_pages, + stats_after.num_total_pages, + ) + + for uuid in completed_list: + 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) + + for uuid in completed_list: + seq = self.global_batch.get_sequence(uuid) + if seq is not None and seq.status == SequenceStatus.COMPLETED: + output = gathered_outputs.get(uuid, {}) + self._report_completion( + uuid, + gathered_text=output.get("text"), + cached_tokens=output.get("cached_tokens"), + ) + elif seq is not None: + logging.warning( + f"Rank {self.rank}: Skipping _report_completion for " + f"{uuid[:8]} (status={seq.status.name}, expected " + "COMPLETED). Likely stale eos_reached from pre-eviction " + "cycle." + ) + # ============ Tokenization and Assignment ============ def _tokenize_global_batch(self) -> None: @@ -4533,6 +5757,23 @@ def _get_effective_chunk_size(self) -> int: chunk = math.ceil(chunk / SequenceEntry.PAGE_SIZE) * SequenceEntry.PAGE_SIZE return chunk + def _get_prefill_initial_capacity_tokens( + self, + seq: SequenceEntry, + chunk_size: int, + ) -> int: + post_prefill_length = seq.prompt_length + 1 + gpu_initial_pages = ( + math.ceil(post_prefill_length / seq.PAGE_SIZE) + + INITIAL_GPU_PAGE_BUFFER + ) + gpu_initial_tokens = gpu_initial_pages * seq.PAGE_SIZE + initial_capacity = max( + seq.prompt_length + chunk_size, + gpu_initial_tokens, + ) + return min(initial_capacity, seq.kv_token_budget) + def _prepare_prefill_batch(self) -> List[str]: """ Select sequences for prefill based on HOST KV cache capacity. @@ -4566,69 +5807,119 @@ def _prepare_prefill_batch(self) -> List[str]: if not all_candidates: return [] - gpus_per_node = NUM_GPUS_PER_NODE num_nodes = self._get_num_nodes() chunk_size = self._get_effective_chunk_size() - # Step 1: Get this node's host KV free pages - local_host_free = self._get_host_kv_free_pages() - - # Step 2: Gather host KV free pages from first rank on each node - # Only rank 0, 8, 16, ... (first on each node) reports actual value - if self.local_rank == 0: - report_node = self.rank // gpus_per_node - report_free = local_host_free - else: - report_node = -1 - report_free = 0 # Non-first ranks report 0 - - free_tensor = torch.tensor([report_node, report_free], dtype=torch.int64, device=self.torch_device) - gathered = [torch.zeros_like(free_tensor) for _ in range(self.world_size)] - dist.all_gather(gathered, free_tensor) - - # Extract per-node host KV free pages - reports_by_node = {} - for item in gathered: - node_id = int(item[0].item()) - if node_id >= 0: - reports_by_node[node_id] = int(item[1].item()) - per_node_host_free = [] - for node in range(num_nodes): - per_node_host_free.append(reports_by_node.get(node, 0)) + # Step 1: Gather host KV stats from the first rank on each node. + node_host_stats = self._gather_host_kv_stats_by_node( + self.host_paged_kv_worker_view + ) + per_node_host_free = [ + int(stats["num_free_pages"]) for stats in node_host_stats + ] + per_node_host_total = [ + int(stats["num_total_pages"]) for stats in node_host_stats + ] if self.rank == 0: logging.info(f"Per-node host KV free pages: {per_node_host_free} (chunk_size={chunk_size})") # Step 3: Select sequences considering per-node host KV capacity. - # The NCCL gather above and the logging below stay here; the greedy - # per-node admission is delegated to PrefillScheduler. + # The greedy per-node admission is delegated to PrefillScheduler. + # With prefix cache enabled and no live sequences, used Host KV pages + # are reclaimable prefix-resident pages. Keep admission prefix-agnostic: + # select up to physical capacity, then let allocation evict cached + # prefix pages if the current free stack is insufficient. + has_active_work = ( + self.global_batch.has_prefilled() + or self.global_batch.has_in_decode() + or self.global_batch.has_on_hold() + ) + if self.enable_prefix_cache and not has_active_work: + per_node_effective_free = list(per_node_host_total) + charge_shared_prefix_pages = True + if self.rank == 0 and per_node_effective_free != per_node_host_free: + logging.info( + "[PREFILL] Using Host KV total capacity for selection " + "because no live sequences are holding private Host KV " + "pages; charging unique shared prefix pages selected " + "for this wave: " + f"total_pages={per_node_effective_free}, " + f"free_pages={per_node_host_free}" + ) + else: + per_node_effective_free = list(per_node_host_free) + charge_shared_prefix_pages = False + prefill_batch = PrefillScheduler.select_prefill_batch( self._make_prefill_selection_request( - all_candidates, per_node_host_free, num_nodes, chunk_size + all_candidates, + per_node_effective_free, + num_nodes, + chunk_size, + charge_shared_prefix_pages=charge_shared_prefix_pages, + ) + ) + + if self.rank == 0: + n_evicted = sum( + 1 for u in prefill_batch + if self.global_batch.get_sequence(u).status == SequenceStatus.EVICTED ) + logging.info( + f"[PREFILL] Selected {len(prefill_batch)} sequences " + f"({n_evicted} recompute from eviction)" + ) + + return prefill_batch + + def _has_active_prefill_decode_work(self) -> bool: + """Return whether current live work can make a small prefill wave costly.""" + return ( + self.global_batch.has_prefilled() + or self.global_batch.has_in_decode() + or self.global_batch.has_on_hold() + ) + + def _should_run_selected_prefill_wave( + self, + prefill_uuids: List[str], + *, + reason: str, + ) -> bool: + req = PrefillWaveGateRequest( + selected_count=len(prefill_uuids), + prefix_cache_enabled=bool(self.enable_prefix_cache), + has_active_work=self._has_active_prefill_decode_work(), + world_size=int(self.world_size), ) - - if self.rank == 0: - n_evicted = sum( - 1 for u in prefill_batch - if self.global_batch.get_sequence(u).status == SequenceStatus.EVICTED + should_run = PrefillScheduler.should_run_prefill_wave(req) + if not should_run and self.rank == 0: + min_sequences = PrefillScheduler.min_prefix_cache_wave_sequences( + self.world_size ) logging.info( - f"[PREFILL] Selected {len(prefill_batch)} sequences " - f"({n_evicted} recompute from eviction)" + "[PREFILL] Deferring small prefix-cache prefill wave: " + f"selected={len(prefill_uuids)} min_sequences={min_sequences} " + f"reason={reason}" ) - - return prefill_batch + return should_run def _make_prefill_selection_request( self, all_candidates: List[str], per_node_host_free: List[int], - num_nodes: int, chunk_size: int, + num_nodes: int, chunk_size: int, *, charge_shared_prefix_pages: bool, ) -> PrefillSelectionRequest: """Snapshot the candidate metadata `select_prefill_batch` consumes.""" from batchgen.sequence import INITIAL_GPU_PAGE_BUFFER + prefix_estimates = self._estimate_prefix_cache_for_admission( + all_candidates + ) candidates = [] for uuid in all_candidates: seq = self.global_batch.get_sequence(uuid) + estimated_cached_tokens, estimated_page_ids = ( + prefix_estimates.get(uuid, (0, ())) + ) candidates.append(PrefillCandidate( uuid=uuid, assigned_rank=seq.assigned_rank, @@ -4640,6 +5931,8 @@ def _make_prefill_selection_request( prompt_length=seq.prompt_length, kv_token_budget=seq.kv_token_budget, page_size=seq.PAGE_SIZE, + estimated_shared_prefix_tokens=estimated_cached_tokens, + estimated_shared_prefix_page_ids=estimated_page_ids, )) return PrefillSelectionRequest( candidates=tuple(candidates), @@ -4648,8 +5941,118 @@ def _make_prefill_selection_request( num_nodes=num_nodes, gpus_per_node=NUM_GPUS_PER_NODE, initial_gpu_page_buffer=INITIAL_GPU_PAGE_BUFFER, + charge_shared_prefix_pages=charge_shared_prefix_pages, ) + def _estimate_prefix_cache_for_admission( + self, + all_candidates: Sequence[str], + ) -> Dict[str, Tuple[int, Tuple[Tuple[int, int], ...]]]: + """Return per-candidate prefix-hit estimates for prefill admission. + + The scheduler must be deterministic on every rank, but only the owning + rank is guaranteed to have the prompt tensor needed for lookup. Each rank + estimates its owned candidates and the small ``uuid -> tokens`` maps are + gathered before building the pure scheduler request. + """ + if not self.enable_prefix_cache: + return {} + if self.prefix_cache_coordinator is None: + raise RuntimeError( + "Prefix cache is enabled but coordinator is not attached" + ) + if self.prefix_cache_runtime_config is None: + raise RuntimeError( + "Prefix cache is enabled but runtime config is missing" + ) + + estimate_start = time.perf_counter() + local_estimates = {} + + for uuid in all_candidates: + seq = self.global_batch.get_sequence(uuid) + if seq.assigned_rank != self.rank: + continue + + prompt_length = int(seq.prompt_length) + if prompt_length <= 0: + local_estimates[uuid] = (0, ()) + continue + + token_tensor = ( + seq.evicted_token_ids + if ( + seq.status == SequenceStatus.EVICTED + and seq.evicted_token_ids is not None + ) + else seq.input_ids + ) + prompt_token_ids = [ + int(token_id) + for token_id in token_tensor.reshape(-1)[:prompt_length].tolist() + ] + result = self.prefix_cache_coordinator.estimate_lookup( + list(self.prefix_cache_runtime_config.namespace_digest), + prompt_token_ids, + ) + cached_tokens = effective_prefix_shared_tokens( + raw_cached_tokens=int(result.common_cached_tokens), + prompt_length=prompt_length, + ) + shared_page_ids = self._prefix_admission_page_ids(result) + local_estimates[uuid] = (int(cached_tokens), shared_page_ids) + + if dist.is_available() and dist.is_initialized() and self.world_size > 1: + gathered = [None] * int(self.world_size) + dist.all_gather_object(gathered, local_estimates) + prefix_estimates = {} + for item in gathered: + if item: + prefix_estimates.update(item) + else: + prefix_estimates = dict(local_estimates) + + if self.rank == 0 and prefix_estimates: + hit_count = sum( + 1 for tokens, _ in prefix_estimates.values() if tokens > 0 + ) + cached_tokens = sum( + int(tokens) for tokens, _ in prefix_estimates.values() + ) + shared_pages = len({ + page_key + for _, page_ids in prefix_estimates.values() + for page_key in page_ids + }) + logging.info( + "[PREFIX_ADMISSION] estimated candidates=%d hit_seqs=%d " + "cached_tokens=%d unique_shared_pages=%d elapsed_ms=%.1f", + len(prefix_estimates), + hit_count, + cached_tokens, + shared_pages, + (time.perf_counter() - estimate_start) * 1000, + ) + + return prefix_estimates + + def _prefix_admission_page_ids( + self, + lookup_result, + ) -> Tuple[Tuple[int, int], ...]: + """Return unique ``(group_id, page_id)`` keys from an estimate result.""" + page_ids = [] + seen = set() + for span in lookup_result.materialization_spans: + group_id = int(span.group_id) + for page in span.pages: + page_key = (group_id, int(page.page_id)) + if page_key in seen: + continue + seen.add(page_key) + page_ids.append(page_key) + return tuple(page_ids) + def _put_sequences_on_hold(self, uuids: List[str]) -> None: """Move IN_DECODE sequences to ON_HOLD, freeing GPU KV but keeping host KV.""" if not uuids: @@ -4874,7 +6277,14 @@ def _submit_completed_to_incremental_writer( if seq is not None and local_idx in self.query_book: finish_reason = self._get_finish_reason(seq) my_completed_tokens.append( - (seq.global_idx, self.query_book[local_idx].decoded_tokens[:, :seq.decoded_length].clone(), finish_reason) + ( + seq.global_idx, + self.query_book[local_idx].decoded_tokens[ + :, : seq.decoded_length + ].clone(), + finish_reason, + int(getattr(seq, "prefix_shared_tokens", 0)), + ) ) # All ranks participate in gather (NCCL collective requirement) @@ -4886,8 +6296,13 @@ def _submit_completed_to_incremental_writer( if writer is not None: for rank_tokens in all_completed_tokens: if rank_tokens: - for global_idx, tokens, finish_reason in rank_tokens: - writer.submit(global_idx, tokens, finish_reason=finish_reason) + for global_idx, tokens, finish_reason, cached_tokens in rank_tokens: + writer.submit( + global_idx, + tokens, + finish_reason=finish_reason, + cached_tokens=cached_tokens, + ) def _try_load_new_sequences( self, @@ -5509,10 +6924,17 @@ def generate(self): ) prefill_uuids = self._prepare_prefill_batch() + prefill_ran = False + if prefill_uuids and not self._should_run_selected_prefill_wave( + prefill_uuids, + reason="active_work", + ): + prefill_uuids = [] if prefill_uuids: if self.rank == 0: logging.info(f"[PREFILL] Starting for {len(prefill_uuids)} sequences") + prefill_phase_start = time.perf_counter() for uuid in prefill_uuids: seq = self.global_batch.get_sequence(uuid) is_reentry = seq.evicted_token_ids is not None @@ -5523,12 +6945,14 @@ def generate(self): # A. Config Prefill (this adds new sequences to _uuid_to_local_map) config_start = time.perf_counter() self._config_prefill_for_batch(prefill_uuids) - config_prefill_time += time.perf_counter() - config_start + config_elapsed = time.perf_counter() - config_start + config_prefill_time += config_elapsed # Get local indices AFTER config (new sequences now in map) local_prefill_indices = self._get_local_indices_for_uuids(prefill_uuids) # B. Execute Prefill + prefill_elapsed = 0.0 if local_prefill_indices: if torch.cuda.is_available(): free_mem, total_mem = torch.cuda.mem_get_info(self.local_rank) @@ -5543,7 +6967,8 @@ def generate(self): self.prefill_prepacked(local_prefill_indices) else: self.prefill(local_prefill_indices) - prefill_time += time.perf_counter() - prefill_start + prefill_elapsed = time.perf_counter() - prefill_start + prefill_time += prefill_elapsed # CRITICAL: Wait for all async KV offloads to complete before decode. # async_offload_layer_kv_to_host returns a future backed by a @@ -5561,6 +6986,7 @@ def generate(self): logging.info( f"[PREFILL_SYNC] waited on {num_retired} async KV offload tasks" ) + self._commit_prefix_cache_prompt_pages(prefill_uuids) # Cleanup & Status Update self._unregister_fp8_weights() @@ -5569,6 +6995,14 @@ def generate(self): seq.log_event(SeqEvent.PREFILL_DONE, self.rank, f"decoded_len={seq.decoded_length}") self._update_batch_status(prefill_uuids, SequenceStatus.PREFILLED) + self._log_prefill_phase_metrics( + prefill_uuids=prefill_uuids, + local_prefill_indices=local_prefill_indices, + config_s=config_elapsed, + prefill_s=prefill_elapsed, + total_s=time.perf_counter() - prefill_phase_start, + ) + prefill_ran = True dist.barrier() # After prefill completes, poll for newly arrived sequences. @@ -5576,9 +7010,15 @@ def generate(self): # loop back to prefill instead of entering decode. if self._admission_queue is not None: self._poll_admissions() - if self.global_batch.has_queueing(): + if prefill_ran and self.global_batch.has_queueing(): next_prefill = self._prepare_prefill_batch() - if next_prefill: + if ( + next_prefill + and self._should_run_selected_prefill_wave( + next_prefill, + reason="back_to_back", + ) + ): if self.rank == 0: logging.info( f"[PREFILL] Back-to-back prefill: {len(next_prefill)} new sequences ready" @@ -5651,51 +7091,7 @@ def generate(self): # Incremental write: submit sequences completed between decode rounds if 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) - gathered_texts = self._gather_completed_tokens(list(global_completed)) - # ORDERING FIX: release resources BEFORE _report_completion - # pops local_map entries. See matching fix in _page_boundary_fast - # Phase 4.A and in the legacy decode path. - completed_list = list(global_completed) - my_completed = [u for u in completed_list if u in self._uuid_to_local_map] - if my_completed: - # Only release GPU pages for seqs that were actually GPU-allocated. - # prefill_prepacked writes KV directly to host (never registers - # with the GPU paged manager), so zero-tok-EOS prefill completions - # are in _uuid_to_local_map but never in manager._sequences. - # _sequences_with_gpu_kv is the source-of-truth tracking set - # (added at :1619/:4904/:6191, discarded on release/eviction). - 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) - # All-ranks scalar cleanup - for uuid in completed_list: - 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) - # Report completions (pops local_map; runs LAST). - # Guard: only report if status actually reached COMPLETED. - # _sync_completion_status_tensor may detect eos_reached=True - # for a PREFILLED sequence (stale from pre-eviction), but - # PREFILLED→COMPLETED is an invalid transition. Without this - # guard, _report_completion pops local_map for a sequence - # whose status never changed, creating an orphan. - for uuid in completed_list: - seq = self.global_batch.get_sequence(uuid) - if seq is not None and seq.status == SequenceStatus.COMPLETED: - self._report_completion(uuid, gathered_text=gathered_texts.get(uuid)) - elif seq is not None: - logging.warning( - f"Rank {self.rank}: Skipping _report_completion for {uuid[:8]} " - f"(status={seq.status.name}, expected COMPLETED). " - f"Likely stale eos_reached from pre-eviction cycle." - ) + self._handle_completed_decode_uuids(global_completed) if not decode_uuids: break @@ -5778,6 +7174,9 @@ def generate(self): new_tokens = torch.empty((0, 1), dtype=torch.int64, device=self.torch_device) self.decoding_continuous(new_tokens, decode_uuids, local_decode_indices) + global_completed, decode_uuids = self._sync_completion_status_tensor(decode_uuids) + if global_completed: + self._handle_completed_decode_uuids(global_completed) decoding_time += time.perf_counter() - decode_start # D. Cleanup @@ -5873,6 +7272,8 @@ def generate(self): # With 12K sequences × 1MB tensors = 12GB, all_gather_object OOMs. # Gathering strings (~KB each) instead reduces memory by ~100x. local_results = [] + if self.rank == 0: + local_results.extend(self._final_response_completed_outputs.items()) for local_idx, uuid in self._local_to_uuid_map.items(): seq = self.global_batch.get_sequence(uuid) if seq is None: @@ -5949,6 +7350,31 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: in_decode = self.global_batch.get_sequences_by_status(SequenceStatus.IN_DECODE) on_hold = self.global_batch.get_sequences_by_status(SequenceStatus.ON_HOLD) prefilling = self.global_batch.get_sequences_by_status(SequenceStatus.PREFILLED) + if self.rank == 0: + evicted_prefill = [ + uuid for uuid in prefill_uuids + if self.global_batch.get_sequence(uuid).status == SequenceStatus.EVICTED + ] + prompt_lengths = [ + int(self.global_batch.get_sequence(uuid).prompt_length) + for uuid in prefill_uuids + ] + decoded_before = [ + int(self.global_batch.get_sequence(uuid).total_decoded_before_eviction) + for uuid in evicted_prefill + ] + prompt_min = min(prompt_lengths) if prompt_lengths else 0 + prompt_max = max(prompt_lengths) if prompt_lengths else 0 + decoded_max = max(decoded_before) if decoded_before else 0 + logging.info( + "[PREFILL_REENTRY] selected_batch: " + f"total={len(prefill_uuids)} evicted={len(evicted_prefill)} " + f"queueing={len(prefill_uuids) - len(evicted_prefill)} " + f"in_decode={len(in_decode)} on_hold={len(on_hold)} " + f"prefilled={len(prefilling)} prompt_tokens={sum(prompt_lengths)} " + f"prompt_len_range=[{prompt_min},{prompt_max}] " + f"max_decoded_before_eviction={decoded_max}" + ) if (in_decode or on_hold) and BATCHGEN_CB_DEBUG: logging.debug( f"Rank {self.rank}: _config_prefill_for_batch called while " @@ -5986,7 +7412,7 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: # CRITICAL: Destroy GPU KV cache BEFORE configure_prefill (Bug Fix 7.2) # The GPU KV cache holds ~20-30GB that must be freed before loading prefill model # Previously this was called AFTER configure_prefill() which caused OOM - self._destroy_gpu_paged_kv_cache() + self._destroy_gpu_paged_kv_cache(empty_cuda_cache=True) if torch.cuda.is_available(): free_mem, total_mem = torch.cuda.mem_get_info(self.local_rank) @@ -6034,6 +7460,10 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: # until the next _sync_sequence_metadata call. # (a) All-ranks scalar metadata update for re-entering sequences. + reentry_scalar_start = time.perf_counter() + reentry_scalar_count = 0 + reentry_scalar_prompt_tokens = 0 + reentry_scalar_decoded_max = 0 for uuid in prefill_uuids: seq = self.global_batch.get_sequence(uuid) # total_decoded_before_eviction > 0 identifies sequences that have @@ -6077,7 +7507,40 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: if hasattr(seq, '_rep_detected'): seq._rep_detected = False + reentry_scalar_count += 1 + reentry_scalar_prompt_tokens += int(seq.prompt_length) + reentry_scalar_decoded_max = max( + reentry_scalar_decoded_max, + int(seq.total_decoded_before_eviction), + ) + + if reentry_scalar_count: + logging.info( + "[PREFILL_REENTRY] " + f"Rank {self.rank}: scalar update end: " + f"seqs={reentry_scalar_count} " + f"prompt_tokens={reentry_scalar_prompt_tokens} " + f"max_decoded_before_eviction={reentry_scalar_decoded_max} " + f"elapsed_ms={(time.perf_counter() - reentry_scalar_start) * 1000:.1f}" + ) + # (b) Owner-only tensor buffer setup. Also clears seq.evicted_token_ids. + owner_reentry_start = time.perf_counter() + owner_reentry_count = sum( + 1 + for uuid in prefill_uuids + if self.global_batch.get_sequence(uuid).evicted_token_ids is not None + ) + if owner_reentry_count: + logging.info( + "[PREFILL_REENTRY] " + f"Rank {self.rank}: owner tensor setup begin: " + f"owner_seqs={owner_reentry_count}" + ) + owner_prompt_tokens = 0 + owner_prompt_min = None + owner_prompt_max = 0 + owner_prev_decoded_max = 0 for uuid in prefill_uuids: seq = self.global_batch.get_sequence(uuid) # Gate on evicted_token_ids (owner-only tensor); non-owners fall @@ -6088,6 +7551,14 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: evicted_ids = seq.evicted_token_ids # 1D tensor new_prompt_len = len(evicted_ids) prev_decoded = seq.total_decoded_before_eviction + owner_prompt_tokens += int(new_prompt_len) + owner_prompt_min = ( + int(new_prompt_len) + if owner_prompt_min is None + else min(owner_prompt_min, int(new_prompt_len)) + ) + owner_prompt_max = max(owner_prompt_max, int(new_prompt_len)) + owner_prev_decoded_max = max(owner_prev_decoded_max, int(prev_decoded)) seq.log_event(SeqEvent.REENTRY_START, self.rank, f"new_prompt_len={new_prompt_len}, prev_decoded={prev_decoded}") @@ -6138,10 +7609,21 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: local_idx = self._uuid_to_local_map[uuid] self.query_book[local_idx] = make_query_book_entry(seq) + logging.info( + f"Rank {self.rank}: Prepared EVICTED seq {uuid[:8]} for re-entry: " + f"new_prompt={new_prompt_len}, prev_decoded={prev_decoded}, " + f"remaining_decode={seq.max_decode_length}, kv_budget={seq.kv_token_budget}" + ) + + if owner_reentry_count: logging.info( - f"Rank {self.rank}: Prepared EVICTED seq {uuid[:8]} for re-entry: " - f"new_prompt={new_prompt_len}, prev_decoded={prev_decoded}, " - f"remaining_decode={seq.max_decode_length}, kv_budget={seq.kv_token_budget}" + "[PREFILL_REENTRY] " + f"Rank {self.rank}: owner tensor setup end: " + f"owner_seqs={owner_reentry_count} " + f"prompt_tokens={owner_prompt_tokens} " + f"prompt_len_range=[{owner_prompt_min},{owner_prompt_max}] " + f"max_prev_decoded={owner_prev_decoded_max} " + f"elapsed_ms={(time.perf_counter() - owner_reentry_start) * 1000:.1f}" ) # STEP 4: Allocate host KV pages for sequences (only THIS RANK's sequences) @@ -6159,70 +7641,191 @@ def _config_prefill_for_batch(self, prefill_uuids: List[str]) -> None: f"(local_idx={new_local_idx})" ) + if self.enable_prefix_cache: + self._prefix_prefill_lookup_by_local_idx.clear() + global_sequence_ids = [] + sequence_tokens = [] + prefill_local_indices = [] + prefix_lookup = None + lookup_results_by_uuid = {} + chunk_size = self._get_effective_chunk_size() + total_private_pages = 0 + total_shared_pages = 0 + total_append_tokens = 0 + if my_prefill_uuids: - global_sequence_ids = [] - sequence_tokens = [] - chunk_size = self._get_effective_chunk_size() + prefill_local_indices = [ + self._uuid_to_local_map[uuid] for uuid in my_prefill_uuids + ] + if self.enable_prefix_cache: + input_ids_for_lookup, _, prompt_lengths_for_lookup = ( + self._prefill_inputs_for_local_indices(prefill_local_indices) + ) + lookup_start = time.perf_counter() + logging.info( + "[PREFIX_LOOKUP] " + f"Rank {self.rank}: prefill lookup begin: " + f"local_seqs={len(prefill_local_indices)} " + f"prompt_tokens={sum(prompt_lengths_for_lookup)} " + f"prompt_len_range=[{min(prompt_lengths_for_lookup)}," + f"{max(prompt_lengths_for_lookup)}]" + ) + prefix_lookup = self._lookup_prefix_cache_for_prefill( + local_indices=prefill_local_indices, + input_ids_list=input_ids_for_lookup, + prompt_lengths=prompt_lengths_for_lookup, + ) + lookup_hits = sum( + 1 for tokens in prefix_lookup.prefix_shared_tokens + if int(tokens) > 0 + ) + raw_cached_tokens = sum( + int(result.common_cached_tokens) + for result in prefix_lookup.lookup_results + ) + logging.info( + "[PREFIX_LOOKUP] " + f"Rank {self.rank}: prefill lookup end: " + f"local_seqs={len(prefill_local_indices)} " + f"hit_seqs={lookup_hits} " + f"effective_cached_tokens=" + f"{sum(int(t) for t in prefix_lookup.prefix_shared_tokens)} " + f"raw_cached_tokens={raw_cached_tokens} " + f"elapsed_ms={(time.perf_counter() - lookup_start) * 1000:.1f}" + ) + for local_idx, result in zip( + prefill_local_indices, + prefix_lookup.lookup_results, + ): + self._prefix_prefill_lookup_by_local_idx[int(local_idx)] = result + lookup_results_by_uuid = dict( + zip(my_prefill_uuids, prefix_lookup.lookup_results) + ) for uuid in my_prefill_uuids: seq = self.global_batch.get_sequence(uuid) global_sequence_ids.append(seq.global_idx) + shared_prefix_tokens = ( + int(seq.prefix_shared_tokens) + if self.enable_prefix_cache and prefix_lookup is not None + else 0 + ) + if not self.enable_prefix_cache: + seq.prefix_shared_tokens = 0 + seq.prefix_committed_tokens = 0 + if shared_prefix_tokens >= int(seq.prompt_length): + raise RuntimeError( + f"Rank {self.rank}: prefix cache hit exceeds prompt " + f"for gid={seq.global_idx}: hit={shared_prefix_tokens}, " + f"prompt={seq.prompt_length}" + ) + lookup_result = lookup_results_by_uuid.get(uuid) + shared_pages = 0 + if lookup_result is not None: + shared_pages = len( + self._host_page_ids_from_prefix_lookup_group( + lookup_result, + group_id=0, + ) + ) + shared_page_tokens = shared_pages * seq.PAGE_SIZE + if shared_page_tokens < shared_prefix_tokens: + raise RuntimeError( + f"Rank {self.rank}: prefix cache shared pages do not " + f"cover effective hit for gid={seq.global_idx}: " + f"pages={shared_pages}, page_size={seq.PAGE_SIZE}, " + f"hit={shared_prefix_tokens}" + ) # Dynamic reservation: allocate prompt + chunk_size, not full budget. # Must also cover the GPU initial load which needs # ceil((prompt+1)/PAGE_SIZE) + INITIAL_GPU_PAGE_BUFFER pages. # The +1 accounts for the first decoded token produced during prefill # (current_context_length = prompt_length + 1 after prefill). - from batchgen.sequence import INITIAL_GPU_PAGE_BUFFER post_prefill_length = seq.prompt_length + 1 # prefill produces 1 decode token gpu_initial_pages = math.ceil(post_prefill_length / seq.PAGE_SIZE) + INITIAL_GPU_PAGE_BUFFER gpu_initial_tokens = gpu_initial_pages * seq.PAGE_SIZE initial_capacity = max(seq.prompt_length + chunk_size, gpu_initial_tokens) initial_capacity = min(initial_capacity, seq.kv_token_budget) - seq.host_pages_allocated = math.ceil(initial_capacity / seq.PAGE_SIZE) + append_tokens = int(seq.prompt_length) - shared_prefix_tokens + private_capacity = max( + initial_capacity - shared_page_tokens, + append_tokens, + ) + private_pages = math.ceil(private_capacity / seq.PAGE_SIZE) + seq.host_pages_allocated = shared_pages + private_pages seq.host_token_capacity = seq.host_pages_allocated * seq.PAGE_SIZE - sequence_tokens.append(seq.host_token_capacity) + sequence_tokens.append(private_pages * seq.PAGE_SIZE) + total_private_pages += private_pages + total_shared_pages += shared_pages + total_append_tokens += append_tokens - # Safety assertion: log if selection over-admitted. This should not - # happen after the EVICTED-length fix in _prepare_prefill_batch — - # if it fires, there's another selection bug to investigate. - kv_stats = self.core_engine.host_paged_kv_worker_view.get_stats() - total_pages_needed = sum(math.ceil(t / seq.PAGE_SIZE) for t in sequence_tokens) - if total_pages_needed > kv_stats.num_free_pages: - # Log per-sequence breakdown to help diagnose the selection bug. - seq_details = [] - for gid, tokens in list(zip(global_sequence_ids, sequence_tokens))[:10]: - s = self.global_batch.get_sequence( - next(u for u in my_prefill_uuids if self.global_batch.get_sequence(u).global_idx == gid) - ) - seq_details.append( - f"gid={gid} prompt_len={s.prompt_length} " - f"was_evicted={s.total_decoded_before_eviction > 0} " - f"tokens={tokens}" - ) - logging.error( - f"Rank {self.rank}: Host KV OVER-ADMISSION: need {total_pages_needed} pages, " - f"have {kv_stats.num_free_pages}. Selection should have prevented this. " - f"First 10 seqs: {seq_details}" + if self.enable_prefix_cache: + if sequence_tokens: + logging.info( + "[PREFILL_ALLOC] " + f"Rank {self.rank}: ensure private host pages begin: " + f"local_seqs={len(sequence_tokens)} " + f"private_pages={total_private_pages} " + f"shared_pages={total_shared_pages} " + f"append_tokens={total_append_tokens}" + ) + ensure_start = time.perf_counter() + self._ensure_prefix_cache_host_pages_for_allocation( + sequence_tokens=sequence_tokens, + reason="prefill_private_allocation", + ) + if sequence_tokens: + logging.info( + "[PREFILL_ALLOC] " + f"Rank {self.rank}: ensure private host pages end: " + f"elapsed_ms={(time.perf_counter() - ensure_start) * 1000:.1f}" ) - logging.debug( - f"Rank {self.rank}: Registering {len(global_sequence_ids)} sequences for host KV " - f"(chunk_size={chunk_size})" + if my_prefill_uuids: + alloc_start = time.perf_counter() + logging.info( + "[PREFILL_ALLOC] " + f"Rank {self.rank}: host KV allocation begin: " + f"local_seqs={len(global_sequence_ids)} chunk_size={chunk_size} " + f"private_pages={total_private_pages} shared_pages={total_shared_pages}" ) self.core_engine.host_paged_kv_worker_view.register_sequences(global_sequence_ids) + aux_view = self.host_paged_kv_worker_view_aux + if aux_view is not None: + aux_view.register_sequences(global_sequence_ids) + if prefix_lookup is not None: + attach_start = time.perf_counter() + logging.info( + "[PREFILL_ALLOC] " + f"Rank {self.rank}: prefix attachment begin: " + f"local_seqs={len(prefill_local_indices)}" + ) + self._attach_prefix_cache_lookup_pages( + local_indices=prefill_local_indices, + lookup=prefix_lookup, + ) + logging.info( + "[PREFILL_ALLOC] " + f"Rank {self.rank}: prefix attachment end: " + f"elapsed_ms={(time.perf_counter() - attach_start) * 1000:.1f}" + ) self.core_engine.host_paged_kv_worker_view.allocate_pages_for_sequences( list(zip(global_sequence_ids, sequence_tokens)) ) # DSA: mirror registration on auxiliary host KV - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) if aux_view is not None: - aux_view.register_sequences(global_sequence_ids) aux_view.allocate_pages_for_sequences( list(zip(global_sequence_ids, sequence_tokens)) ) kv_stats = self.core_engine.host_paged_kv_worker_view.get_stats() + logging.info( + "[PREFILL_ALLOC] " + f"Rank {self.rank}: host KV allocation end: " + f"used={kv_stats.num_used_pages}/{kv_stats.num_total_pages} " + f"elapsed_ms={(time.perf_counter() - alloc_start) * 1000:.1f}" + ) if self.rank == 0: logging.info(f"[PREFILL] Host KV allocated: {kv_stats.num_used_pages}/{kv_stats.num_total_pages} pages") @@ -6625,9 +8228,12 @@ def _release_host_kv_pages_for_batch(self, uuids: List[str]) -> None: # so we don't need to call unregister_sequences separately worker_view.release_sequence_pages(global_sequence_ids) # DSA: release auxiliary host KV pages too - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view = self.host_paged_kv_worker_view_aux if aux_view is not None: aux_view.release_sequence_pages(global_sequence_ids) + self._release_prefix_cache_attachments_for_global_ids( + global_sequence_ids + ) # Rebuild GPU page table with remaining active sequences manager = self.gpu_paged_kv_cache_manager @@ -6661,7 +8267,7 @@ def prefill(self, batch: list[int]): # ensures `_offload_prepacked_indexer_kv` actually pushes indexer K to # the host aux cache instead of early-returning on a None view. AttnWrapperBase.host_paged_kv_worker_view = getattr(self.core_engine, "host_paged_kv_worker_view", None) - AttnWrapperBase.host_paged_kv_worker_view_aux = getattr(self, "host_paged_kv_worker_view_aux", None) + AttnWrapperBase.host_paged_kv_worker_view_aux = self.host_paged_kv_worker_view_aux if "deepseek" in self.model_config.model_type: self.model.model._use_flash_attention_2 = False @@ -6772,9 +8378,51 @@ def prefill(self, batch: list[int]): # MODIFIED: Check for EOS respecting ignore_eos flag if self._should_stop_at_eos(new_tokens_cpu[i].item()): seq.eos_reached = True + if seq.decoded_length >= seq.max_decode_length: + seq.eos_reached = True return new_tokens + def _reset_prefill_prepack_runtime_state(self) -> None: + # Reset prepack mode + Attn_Wrapper.prepack_mode = False + Attn_Wrapper.prepack_cu_seqlens = None + Attn_Wrapper.prepack_max_seqlen = None + Attn_Wrapper.prepack_num_sequences = None + Attn_Wrapper.prepack_seq_lengths = None + Attn_Wrapper.prepack_append_seq_lengths = None + Attn_Wrapper.prepack_prefix_reuse_mode = False + Attn_Wrapper.prepack_prefix_shared_tokens = None + Attn_Wrapper.prepack_full_seq_lengths = None + + # Also reset AttnWrapperBase for models using new wrapper system (GPT-OSS) + AttnWrapperBase.prepack_mode = False + AttnWrapperBase.prepack_cu_seqlens = None + AttnWrapperBase.prepack_max_seqlen = None + AttnWrapperBase.prepack_num_sequences = None + AttnWrapperBase.prepack_seq_lengths = None + AttnWrapperBase.prepack_append_seq_lengths = None + AttnWrapperBase.prepack_prefix_reuse_mode = False + AttnWrapperBase.prepack_prefix_shared_tokens = None + AttnWrapperBase.prepack_full_seq_lengths = None + AttnWrapperBase.prefill_prefix_materialization = None + + @contextmanager + def _prefill_prepack_runtime_scope(self, prefix_materialization): + try: + yield + finally: + AttnWrapperBase.retire_pending_prefill_offloads( + device=self.torch_device, + reason="end of prepack microbatch", + ) + self._reset_prefill_prepack_runtime_state() + if prefix_materialization is not None: + try: + prefix_materialization.close(empty_cuda_cache=False) + finally: + self._destroy_gpu_paged_kv_cache(empty_cuda_cache=True) + def prefill_prepacked(self, batch: list[int]): """ Handle prefill for a batch using prepack optimization. @@ -6796,55 +8444,28 @@ def prefill_prepacked(self, batch: list[int]): # ensures `_offload_prepacked_indexer_kv` actually pushes indexer K to # the host aux cache instead of early-returning on a None view. AttnWrapperBase.host_paged_kv_worker_view = getattr(self.core_engine, "host_paged_kv_worker_view", None) - AttnWrapperBase.host_paged_kv_worker_view_aux = getattr(self, "host_paged_kv_worker_view_aux", None) + AttnWrapperBase.host_paged_kv_worker_view_aux = self.host_paged_kv_worker_view_aux if "deepseek" in self.model_config.model_type: self.model.model._use_flash_attention_2 = False - # Collect input_ids and attention_masks as lists for prepacking - input_ids_list = [] - attention_mask_list = [] - seq_lengths = [] - - for query_idx in batch: - uuid = self._local_to_uuid_map[query_idx] - seq = self.global_batch.get_sequence(uuid) - query_entry = self.query_book[query_idx] - encoded = query_entry.encoded["input_ids"] - if encoded.data_ptr() != seq.input_ids.data_ptr(): - raise RuntimeError( - f"Rank {self.rank}: stale query_book input_ids binding for " - f"local_idx={query_idx} uuid={uuid[:8]} " - f"(query_book_ptr={encoded.data_ptr():#x}, seq_ptr={seq.input_ids.data_ptr():#x})" - ) - if query_entry.decoded_tokens.data_ptr() != seq.decoded_tokens.data_ptr(): - raise RuntimeError( - f"Rank {self.rank}: stale query_book decoded_tokens binding for " - f"local_idx={query_idx} uuid={uuid[:8]} " - f"(query_book_ptr={query_entry.decoded_tokens.data_ptr():#x}, " - f"seq_ptr={seq.decoded_tokens.data_ptr():#x})" - ) - # NO truncation: every prompt is tokenized to its OWN length. - # An earlier `[:, :self.max_input_length]` slice silently dropped - # the tail of long LongBench prompts when max_input_length was - # carried over from a smaller earlier admit batch, causing the - # model to "continue" mid-sentence instead of answering. Bind - # everything to seq.prompt_length directly. - L = seq.prompt_length - assert encoded.size(-1) >= L, ( - f"encoded prompt length {encoded.size(-1)} < seq.prompt_length {L} " - f"for query_idx={query_idx} uuid={uuid[:8]}" - ) - input_ids = encoded[:, :L] - seq_lengths.append(L) - - # Per-seq mask marks the L valid positions for the prepacker. - # Causal attention is enforced by FA varlen + cu_seqlens. - attention_mask = torch.zeros_like(input_ids, dtype=torch.int64) - attention_mask[0, :L] = 1 - - input_ids_list.append(input_ids) - attention_mask_list.append(attention_mask) + full_input_ids_list, attention_mask_list, seq_lengths = ( + self._prefill_inputs_for_local_indices(batch) + ) + input_ids_list = full_input_ids_list + prefix_plan = None + prefix_lookup = self._prefix_cache_lookup_for_prefill_batch(batch) + if prefix_lookup is not None: + prefix_inputs = self._build_prefix_reuse_prepack_inputs( + local_indices=batch, + input_ids_list=full_input_ids_list, + prompt_lengths=seq_lengths, + lookup=prefix_lookup, + ) + prefix_plan = prefix_inputs.plan + input_ids_list = prefix_inputs.input_ids_list + attention_mask_list = prefix_inputs.attention_mask_list + seq_lengths = [item.suffix_length for item in prefix_plan.sequences] # Prepack sequences # Row capacity is set by planner in config (None = no limit, use max sequence length) @@ -6881,8 +8502,14 @@ def prefill_prepacked(self, batch: list[int]): seq_input_ids = prepack_meta.packed_input_ids[row_idx, start_pos:start_pos + seq_len] packed_input_ids_flat.append(seq_input_ids) - # Position IDs are 0, 1, 2, ... for each sequence - packed_position_ids_flat.append(torch.arange(seq_len, device=self.torch_device)) + if prefix_plan is None: + packed_position_ids_flat.append( + torch.arange(seq_len, device=self.torch_device) + ) + else: + packed_position_ids_flat.append( + prefix_plan.suffix_position_ids[seq_idx].to(self.torch_device) + ) packed_input_ids_flat = torch.cat(packed_input_ids_flat, dim=0) # [total_tokens] packed_position_ids_flat = torch.cat(packed_position_ids_flat, dim=0) # [total_tokens] @@ -6891,8 +8518,26 @@ def prefill_prepacked(self, batch: list[int]): # This prevents OOM when sequences have varying lengths # Token cap is set by planner in config, worker reads from config (no hardcoded values) MAX_TOKENS_PER_MICRO_BATCH = self.engine_config.Module_Batching_Config.prefill_micro_batch_token_cap + if prefix_plan is not None and prefix_plan.saved_prefill_tokens > 0: + prefix_reuse_cap = int( + os.environ.get( + "BATCHGEN_PREFIX_REUSE_PREFILL_MICRO_BATCH_TOKEN_CAP", + "131072", + ) + ) + if prefix_reuse_cap > 0: + MAX_TOKENS_PER_MICRO_BATCH = min( + MAX_TOKENS_PER_MICRO_BATCH, + prefix_reuse_cap, + ) num_sequences = prepack_meta.num_original_sequences seq_lengths_list = prepack_meta.original_seq_lengths + micro_batch_admission_lengths = seq_lengths_list + if prefix_plan is not None and prefix_plan.saved_prefill_tokens > 0: + micro_batch_admission_lengths = [ + max(1, int(item.full_logical_context_length)) + for item in prefix_plan.sequences + ] # Create micro-batches bounded by token count, optionally also by sum(L^2) # so the per-microbatch attention work (which is O(L^2)) doesn't pile up @@ -6900,16 +8545,22 @@ def prefill_prepacked(self, batch: list[int]): import os as _os_mb _USE_L2_MB = _os_mb.environ.get("BATCHGEN_L2_BALANCE", "1") == "1" micro_batches, l2_cap = build_prefill_micro_batches( - seq_lengths_list, + micro_batch_admission_lengths, MAX_TOKENS_PER_MICRO_BATCH, l2_balance=_USE_L2_MB, ) total_tokens_all = sum(seq_lengths_list) + total_admission_tokens = sum(micro_batch_admission_lengths) if self.rank == 0: logging.info( f"Prepacked prefill: {len(micro_batches)} micro batches, " f"{total_tokens_all:,} total tokens, max {MAX_TOKENS_PER_MICRO_BATCH:,} tokens/batch" + + ( + f", admission_tokens={total_admission_tokens:,}" + if total_admission_tokens != total_tokens_all + else "" + ) + (f", l2_cap={l2_cap:,}" if l2_cap > 0 else "") ) @@ -6973,91 +8624,156 @@ def prefill_prepacked(self, batch: list[int]): device=self.torch_device, ) batch_max_seqlen = max(batch_seq_lengths) + if prefix_plan is None: + batch_append_seq_lengths = list(batch_seq_lengths) + batch_prefix_shared_tokens = None + batch_full_seq_lengths = None + batch_prefix_reuse_mode = False + else: + batch_plan_items = prefix_plan.sequences[seq_start:seq_end] + batch_append_seq_lengths = [ + int(item.suffix_length) for item in batch_plan_items + ] + batch_prefix_shared_tokens = [ + int(item.prefix_shared_tokens) + for item in batch_plan_items + ] + batch_full_seq_lengths = [ + int(item.full_logical_context_length) + for item in batch_plan_items + ] + batch_prefix_reuse_mode = any( + tokens > 0 for tokens in batch_prefix_shared_tokens + ) - # Set up Attn_Wrapper for this micro-batch - Attn_Wrapper.prepack_mode = True - Attn_Wrapper.prepack_cu_seqlens = batch_cu_seqlens - Attn_Wrapper.prepack_max_seqlen = batch_max_seqlen - Attn_Wrapper.prepack_num_sequences = batch_num_seqs - Attn_Wrapper.prepack_seq_lengths = batch_seq_lengths - Attn_Wrapper.position_ids = batch_position_ids_flat - Attn_Wrapper.cur_batch = prefill_sequence_spans_to_global_seq_ids(batch_spans) - - # CRITICAL: Also bind to AttnWrapperBase for models using new wrapper system (GPT-OSS) - # Without this, GPT-OSS uses _forward_prefill instead of _forward_prefill_prepacked, - # which does NOT offload KV to host, causing decode to read garbage. - AttnWrapperBase.prepack_mode = True - AttnWrapperBase.prepack_cu_seqlens = batch_cu_seqlens - AttnWrapperBase.prepack_max_seqlen = batch_max_seqlen - AttnWrapperBase.prepack_num_sequences = batch_num_seqs - AttnWrapperBase.prepack_seq_lengths = batch_seq_lengths - AttnWrapperBase.position_ids = batch_position_ids_flat - AttnWrapperBase.cur_batch = Attn_Wrapper.cur_batch - - # Embed tokens - inputs_embeds = self.model.model.embed_tokens(batch_input_ids_flat.to(self.torch_device)) - - # Reshape to 3D: [1, batch_total_tokens, hidden_dim] - hidden_states = inputs_embeds.unsqueeze(0) + batch_prefix_materialization = None + if ( + prefix_lookup is not None + and prefix_plan is not None + and batch_prefix_reuse_mode + ): + batch_lookup = PrefixCachePrefillLookup( + lookup_results=tuple( + prefix_lookup.lookup_results[seq_start:seq_end] + ), + prefix_shared_tokens=tuple( + prefix_lookup.prefix_shared_tokens[seq_start:seq_end] + ), + ) + batch_prefix_materialization = ( + self._materialize_prefix_cache_prefill( + lookup=batch_lookup, + prefix_plan=split_prefix_reuse_plan_for_micro_batch( + prefix_plan, + seq_start, + seq_end, + ), + ) + ) - 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, + with self._prefill_prepack_runtime_scope(batch_prefix_materialization): + # Set up Attn_Wrapper for this micro-batch + Attn_Wrapper.prepack_mode = True + Attn_Wrapper.prepack_cu_seqlens = batch_cu_seqlens + Attn_Wrapper.prepack_max_seqlen = batch_max_seqlen + Attn_Wrapper.prepack_num_sequences = batch_num_seqs + Attn_Wrapper.prepack_seq_lengths = batch_seq_lengths + Attn_Wrapper.prepack_append_seq_lengths = batch_append_seq_lengths + Attn_Wrapper.prepack_prefix_reuse_mode = batch_prefix_reuse_mode + Attn_Wrapper.prepack_prefix_shared_tokens = ( + batch_prefix_shared_tokens if batch_prefix_reuse_mode else None ) - hidden_states = layer_outputs[0] - - # Final norm - hidden_states = self.model.model.norm(hidden_states) - - # Extract last token hidden states for each sequence - last_token_indices = batch_cu_seqlens[1:] - 1 - last_token_hidden = hidden_states[0, last_token_indices, :] - - # lm_head matmul: BF16 by default (matches HF / SGLang / vLLM). - # Opt into FP32-cast via BATCHGEN_GLM5_LMHEAD_FP32=1 for debugging. - if os.environ.get("BATCHGEN_GLM5_LMHEAD_FP32", "0") == "1": - logits = torch.nn.functional.linear( - last_token_hidden.float(), - self.model.lm_head.weight.float(), - self.model.lm_head.bias.float() if hasattr(self.model.lm_head, 'bias') and self.model.lm_head.bias is not None else None + Attn_Wrapper.prepack_full_seq_lengths = ( + batch_full_seq_lengths if batch_prefix_reuse_mode else None ) - else: - logits = torch.nn.functional.linear( - last_token_hidden, - self.model.lm_head.weight, - self.model.lm_head.bias if hasattr(self.model.lm_head, 'bias') and self.model.lm_head.bias is not None else None - ).float() + Attn_Wrapper.position_ids = batch_position_ids_flat + Attn_Wrapper.cur_batch = prefill_sequence_spans_to_global_seq_ids(batch_spans) - batch_sequences = [ - self.global_batch.get_sequence(self._local_to_uuid_map[local_idx]) - for local_idx in batch_local_indices - ] - batch_new_tokens = self._select_tokens(logits, batch_sequences) - if batch_new_tokens.shape[0] != batch_num_seqs: - raise RuntimeError( - f"Rank {self.rank}: prefill token selection shape mismatch, " - f"got {batch_new_tokens.shape[0]} rows for {batch_num_seqs} sequences" + # CRITICAL: Also bind to AttnWrapperBase for models using new wrapper system (GPT-OSS) + # Without this, GPT-OSS uses _forward_prefill instead of _forward_prefill_prepacked, + # which does NOT offload KV to host, causing decode to read garbage. + AttnWrapperBase.prepack_mode = True + AttnWrapperBase.prepack_cu_seqlens = batch_cu_seqlens + AttnWrapperBase.prepack_max_seqlen = batch_max_seqlen + AttnWrapperBase.prepack_num_sequences = batch_num_seqs + AttnWrapperBase.prepack_seq_lengths = batch_seq_lengths + AttnWrapperBase.prepack_append_seq_lengths = batch_append_seq_lengths + AttnWrapperBase.prepack_prefix_reuse_mode = batch_prefix_reuse_mode + AttnWrapperBase.prepack_prefix_shared_tokens = ( + batch_prefix_shared_tokens if batch_prefix_reuse_mode else None + ) + AttnWrapperBase.prepack_full_seq_lengths = ( + batch_full_seq_lengths if batch_prefix_reuse_mode else None + ) + AttnWrapperBase.position_ids = batch_position_ids_flat + AttnWrapperBase.cur_batch = Attn_Wrapper.cur_batch + AttnWrapperBase.prefill_prefix_materialization = ( + batch_prefix_materialization ) - output_tokens.append(batch_new_tokens) - # Reset prepack mode - Attn_Wrapper.prepack_mode = False - Attn_Wrapper.prepack_cu_seqlens = None - Attn_Wrapper.prepack_max_seqlen = None - Attn_Wrapper.prepack_num_sequences = None - Attn_Wrapper.prepack_seq_lengths = None + # Embed tokens + inputs_embeds = self.model.model.embed_tokens(batch_input_ids_flat.to(self.torch_device)) - # Also reset AttnWrapperBase for models using new wrapper system (GPT-OSS) - AttnWrapperBase.prepack_mode = False - AttnWrapperBase.prepack_cu_seqlens = None - AttnWrapperBase.prepack_max_seqlen = None - AttnWrapperBase.prepack_num_sequences = None - AttnWrapperBase.prepack_seq_lengths = None + # Reshape to 3D: [1, batch_total_tokens, hidden_dim] + hidden_states = inputs_embeds.unsqueeze(0) + + layer_outputs = None + for layer_idx, decoder_layer in enumerate(self.model.model.layers): + AttnWrapperBase.retire_pending_prefill_offloads_before_layer( + layer_idx, + device=self.torch_device, + ) + 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] + + # Final norm + hidden_states = self.model.model.norm(hidden_states) + + # Extract last token hidden states for each sequence + last_token_indices = batch_cu_seqlens[1:] - 1 + last_token_hidden = hidden_states[0, last_token_indices, :] + + # lm_head matmul: BF16 by default (matches HF / SGLang / vLLM). + # Opt into FP32-cast via BATCHGEN_GLM5_LMHEAD_FP32=1 for debugging. + if os.environ.get("BATCHGEN_GLM5_LMHEAD_FP32", "0") == "1": + logits = torch.nn.functional.linear( + last_token_hidden.float(), + self.model.lm_head.weight.float(), + self.model.lm_head.bias.float() if hasattr(self.model.lm_head, 'bias') and self.model.lm_head.bias is not None else None + ) + else: + logits = torch.nn.functional.linear( + last_token_hidden, + self.model.lm_head.weight, + self.model.lm_head.bias if hasattr(self.model.lm_head, 'bias') and self.model.lm_head.bias is not None else None + ).float() + + batch_sequences = [ + self.global_batch.get_sequence(self._local_to_uuid_map[local_idx]) + for local_idx in batch_local_indices + ] + batch_new_tokens = self._select_tokens(logits, batch_sequences) + if batch_new_tokens.shape[0] != batch_num_seqs: + raise RuntimeError( + f"Rank {self.rank}: prefill token selection shape mismatch, " + f"got {batch_new_tokens.shape[0]} rows for {batch_num_seqs} sequences" + ) + output_tokens.append(batch_new_tokens.cpu()) + del ( + inputs_embeds, + hidden_states, + layer_outputs, + last_token_hidden, + logits, + batch_new_tokens, + ) # Log timing summary for GPT-OSS if timing was enabled self._log_prefill_timing() @@ -7083,6 +8799,8 @@ def prefill_prepacked(self, batch: list[int]): # Check for EOS respecting ignore_eos flag if self._should_stop_at_eos(new_tokens_cpu[i].item()): seq.eos_reached = True + if seq.decoded_length >= seq.max_decode_length: + seq.eos_reached = True return new_tokens @@ -7475,40 +9193,7 @@ def _page_boundary_fast( completed_uuids = decisions.completed_uuids if completed_uuids: self._update_batch_status(completed_uuids, SequenceStatus.COMPLETED) - # Incremental write: gather completed tokens to rank 0 - self._submit_completed_to_incremental_writer(completed_uuids) - # Gather decoded tokens from owning ranks before reporting - gathered_texts = self._gather_completed_tokens(completed_uuids) - - # Release resources on owners BEFORE popping local_map entries via - # _report_completion (see ordering fix note above). - my_completed = [u for u in completed_uuids if u in self._uuid_to_local_map] - if my_completed: - # Only release GPU pages for seqs that were actually GPU-allocated. - # See note at the matching site (~line 5435) — zero-tok-EOS - # prefill completions are in _uuid_to_local_map but 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) - - # All-ranks: zero scalar counters so downstream reads (e.g. - # migration planning iterating all sequences) never see a stale - # non-zero page count for completed sequences. - 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) - - # Report completions (this is what pops local_map on the owner). - # Must run LAST so the _release_*_pages calls above see the - # correct local_map state. - for uuid in completed_uuids: - self._report_completion(uuid, gathered_text=gathered_texts.get(uuid)) + self._handle_completed_decode_uuids(completed_uuids) # Report completions to adaptive chunk sizer if self.adaptive_chunk_sizer is not None: for uuid in completed_uuids: @@ -7596,10 +9281,13 @@ def _page_boundary_fast( worker_view.release_sequence_pages(evicted_global_ids) worker_view.unregister_sequences(evicted_global_ids) # DSA: mirror release + unregister on auxiliary host KV - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view = self.host_paged_kv_worker_view_aux if aux_view is not None: aux_view.release_sequence_pages(evicted_global_ids) aux_view.unregister_sequences(evicted_global_ids) + self._release_prefix_cache_attachments_for_global_ids( + evicted_global_ids + ) # All-ranks: update scalar metadata deterministically. Compute # new_reentry_len from already-synced prompt_length, decoded_length, @@ -7662,7 +9350,7 @@ def _page_boundary_fast( if host_grow_requests and worker_view is not None: worker_view.grow_pages_for_sequences(host_grow_requests) # DSA: mirror growth on auxiliary host KV - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view = self.host_paged_kv_worker_view_aux if aux_view is not None: aux_view.grow_pages_for_sequences(host_grow_requests) if self.rank == 0: @@ -7812,7 +9500,7 @@ def _page_boundary_fast( t_launch = time.perf_counter() if worker_view is not None: existing_global_ids = self._local_indices_to_global_seq_ids(batch) - if isinstance(gpu_manager, DualKVCacheCoordinator): + if isinstance(gpu_manager, GroupedGPUKVCoordinator): pointers = self._prepare_dual_kv_load_pointers( gpu_manager, new_load_global, existing_global_ids ) @@ -7967,10 +9655,10 @@ def _finalize_async_load_minimal( Attn_Wrapper.async_kv_load_active = False Attn_Wrapper.async_kv_load_task = None - if pending_local_indices and isinstance(gpu_manager, DualKVCacheCoordinator): - if not isinstance(async_task, DualAsyncKVTask): + if pending_local_indices and isinstance(gpu_manager, GroupedGPUKVCoordinator): + if not isinstance(async_task, GroupedAsyncKVTask): raise RuntimeError( - "DSA async load finalize requires a completed DualAsyncKVTask" + "Grouped KV async load finalize requires a completed grouped task" ) pending_local_uuid_set = { @@ -9461,14 +11149,14 @@ def decoding_continuous( Attn_Wrapper.cur_batch = self._local_indices_to_global_seq_ids(batch) if batch else [] # Also bind to AttnWrapperBase for models using new wrapper system (e.g., GPT-OSS) - if isinstance(gpu_manager, DualKVCacheCoordinator): + if isinstance(gpu_manager, GroupedGPUKVCoordinator): AttnWrapperBase.gpu_paged_kv_manager = gpu_manager.primary AttnWrapperBase.gpu_paged_kv_manager_aux = gpu_manager.auxiliary else: AttnWrapperBase.gpu_paged_kv_manager = gpu_manager AttnWrapperBase.gpu_paged_kv_manager_aux = None AttnWrapperBase.host_paged_kv_worker_view = worker_view - AttnWrapperBase.host_paged_kv_worker_view_aux = getattr(self, "host_paged_kv_worker_view_aux", None) + AttnWrapperBase.host_paged_kv_worker_view_aux = self.host_paged_kv_worker_view_aux AttnWrapperBase.cur_batch = Attn_Wrapper.cur_batch # CRITICAL FIX: Ensure page table matches cur_batch at entry @@ -9520,9 +11208,27 @@ def decoding_continuous( self._cumulative_forward_ms = 0.0 # Local iteration counter (for boundary interval tracking within this decode round) + decode_round_start = time.perf_counter() + round_start_iterations = self._cumulative_decode_iterations + round_start_boundaries = self._cumulative_decode_boundaries + round_start_forward_ms = self._cumulative_forward_ms + round_start_boundary_ms = self._cumulative_boundary_ms + active_start = len(decode_uuids) + local_generated_tokens = 0 local_iteration = 0 last_boundary = 0 global_batch_size = len(self.global_batch) + decode_terminal_sync_iteration = None + if decode_uuids: + max_remaining_decode = 0 + for uuid in decode_uuids: + seq = self.global_batch.get_sequence(uuid) + if seq is None: + continue + remaining = int(seq.max_decode_length) - int(seq.decoded_length) + max_remaining_decode = max(max_remaining_decode, remaining) + if max_remaining_decode > 0: + decode_terminal_sync_iteration = max_remaining_decode # ========== INITIAL MOE BUFFER SYNC ========== # Sync buffer size BEFORE first forward pass to prevent overflow. @@ -9903,7 +11609,7 @@ def decoding_continuous( self._deferred_kv_entries = [] self._deferred_kv_entries_aux = [] self._deferred_kv_worker_view = _kv_worker_view - self._deferred_kv_worker_view_aux = getattr(self, "host_paged_kv_worker_view_aux", None) + self._deferred_kv_worker_view_aux = self.host_paged_kv_worker_view_aux if BATCHGEN_SYNC_KV and _kv_worker_view is not None: # SYNC MODE: Immediately write each layer's KV to host (no deferral) @@ -9937,7 +11643,7 @@ def kv_append_callback(layer_idx: int, k_tensor: torch.Tensor, v_tensor: torch.T # In deferred mode (BATCHGEN_SYNC_KV=0, the default) layers push # to _deferred_kv_entries_aux; a single event.synchronize in # _flush_deferred_kv_to_host covers both primary and aux caches. - aux_view = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view = self.host_paged_kv_worker_view_aux if aux_view is not None: if BATCHGEN_SYNC_KV: def kv_append_callback_aux(layer_idx: int, k_tensor: torch.Tensor, v_tensor: torch.Tensor = None): @@ -10409,6 +12115,7 @@ def kv_append_callback_aux(layer_idx: int, k_tensor: torch.Tensor, v_tensor: tor new_tokens_cpu = _new_tokens_pinned[:bs] # Update sequences (reuse batch_sequences from forward pass setup) + step_generated_tokens = 0 for i, (local_idx, seq) in enumerate(zip(batch, batch_sequences)): if self._is_sequence_completed(seq): continue @@ -10427,6 +12134,7 @@ def kv_append_callback_aux(layer_idx: int, k_tensor: torch.Tensor, v_tensor: tor seq.decoded_length += 1 seq.current_context_length += 1 + step_generated_tokens += 1 # Use CPU tensor to avoid GPU sync token_id = new_tokens_cpu[i].item() @@ -10472,6 +12180,34 @@ def kv_append_callback_aux(layer_idx: int, k_tensor: torch.Tensor, v_tensor: tor f"Rank {self.rank}: REPETITION (ngram) {seq.uuid[:8]} " f"gid={seq.global_idx} at decoded_len={_dl}" ) + local_generated_tokens += step_generated_tokens + + if ( + decode_terminal_sync_iteration is not None + and local_iteration >= decode_terminal_sync_iteration + ): + global_completed, decode_uuids = ( + self._sync_completion_status_tensor(decode_uuids) + ) + if global_completed: + self._handle_completed_decode_uuids(global_completed) + batch = self._get_local_indices_for_uuids(decode_uuids) + if not decode_uuids: + break + if gpu_manager is not None and gpu_manager.is_initialized: + self._rebuild_page_table_for_batch(batch, gpu_manager) + max_remaining_decode = 0 + for uuid in decode_uuids: + seq = self.global_batch.get_sequence(uuid) + if seq is None: + continue + remaining = int(seq.max_decode_length) - int(seq.decoded_length) + max_remaining_decode = max(max_remaining_decode, remaining) + decode_terminal_sync_iteration = ( + local_iteration + max_remaining_decode + if max_remaining_decode > 0 + else None + ) self._cumulative_forward_ms += (time.perf_counter() - forward_start) * 1000 @@ -10536,6 +12272,15 @@ def kv_append_callback_aux(layer_idx: int, k_tensor: torch.Tensor, v_tensor: tor f"{'='*50}" ) + self._log_decode_phase_metrics( + active_start=active_start, + local_generated_tokens=local_generated_tokens, + elapsed_s=time.perf_counter() - decode_round_start, + iteration_delta=self._cumulative_decode_iterations - round_start_iterations, + boundary_delta=self._cumulative_decode_boundaries - round_start_boundaries, + forward_ms_delta=self._cumulative_forward_ms - round_start_forward_ms, + boundary_ms_delta=self._cumulative_boundary_ms - round_start_boundary_ms, + ) self.disable_decode_watchdog() return decode_uuids, batch @@ -10838,7 +12583,7 @@ def _launch_async_load_new_sequences( gpu_manager.rebuild_page_table(existing_global_ids) return None, new_uuids, new_local_indices, new_global_ids - if isinstance(gpu_manager, DualKVCacheCoordinator): + if isinstance(gpu_manager, GroupedGPUKVCoordinator): pointers = self._prepare_dual_kv_load_pointers( gpu_manager, new_global_ids, existing_global_ids ) @@ -10993,7 +12738,7 @@ def _launch_async_load_new_sequences_timed( # Capture existing batch for later restoration existing_global_ids = self._local_indices_to_global_seq_ids(current_batch) - if isinstance(gpu_manager, DualKVCacheCoordinator): + if isinstance(gpu_manager, GroupedGPUKVCoordinator): pointers = self._prepare_dual_kv_load_pointers( gpu_manager, new_global_ids, existing_global_ids ) @@ -11019,7 +12764,7 @@ def _launch_async_load_new_sequences_timed( timing['launch_ms'] = (time.perf_counter() - t0) * 1000 return None, new_uuids, new_local_indices, new_global_ids, timing - if isinstance(gpu_manager, DualKVCacheCoordinator): + if isinstance(gpu_manager, GroupedGPUKVCoordinator): async_task = self._launch_dual_host_kv_load(pointers) else: async_task = worker_view.async_load_layer_paged_kv_to_device( @@ -11037,7 +12782,7 @@ def _launch_async_load_new_sequences_timed( timing['launch_ms'] = (time.perf_counter() - t0) * 1000 # Store tensor references to prevent GC during async operation - self._async_load_tensors = pointers if isinstance(gpu_manager, DualKVCacheCoordinator) else { + self._async_load_tensors = pointers if isinstance(gpu_manager, GroupedGPUKVCoordinator) else { 'k_ptrs': k_ptrs, 'v_ptrs': v_ptrs, 'sequence_tensor': sequence_tensor, @@ -11285,8 +13030,8 @@ def _decoding_legacy_modes( self._update_batch_status(completed_uuids, SequenceStatus.COMPLETED) # Incremental write: gather completed tokens to rank 0 self._submit_completed_to_incremental_writer(completed_uuids) - # Gather decoded tokens from owning ranks before reporting - gathered_texts = self._gather_completed_tokens(completed_uuids) + # Gather decoded tokens and usage metadata from owning ranks before reporting + gathered_outputs = self._gather_completed_outputs(completed_uuids) # ORDERING FIX: release GPU/host KV BEFORE _report_completion # pops local_map entries. Previously the filter below # captured an empty list because _report_completion ran @@ -11301,7 +13046,12 @@ def _decoding_legacy_modes( self._release_host_kv_pages_for_batch(completed_uuids) # Report completions (this pops local_map; must run LAST). for uuid in completed_uuids: - self._report_completion(uuid, gathered_text=gathered_texts.get(uuid)) + output = gathered_outputs.get(uuid, {}) + self._report_completion( + uuid, + gathered_text=output.get("text"), + cached_tokens=output.get("cached_tokens"), + ) if decode_uuids: decode_uuids, batch = self._try_load_new_sequences(decode_uuids, batch) @@ -12061,12 +13811,13 @@ def _reset_for_new_batch(self) -> None: ) # Try to release each sequence individually to handle already-released ones released_count = 0 - aux_view_shutdown = getattr(self, "host_paged_kv_worker_view_aux", None) + aux_view_shutdown = self.host_paged_kv_worker_view_aux for seq_id in global_ids_to_release: try: worker_view.release_sequence_pages([seq_id]) if aux_view_shutdown is not None: aux_view_shutdown.release_sequence_pages([seq_id]) + self._release_prefix_cache_attachments_for_global_ids([seq_id]) released_count += 1 except Exception: # Sequence was already released during decode - this is normal @@ -12077,6 +13828,7 @@ def _reset_for_new_batch(self) -> None: # 2. Reset batch completion flag self._batch_completed = False + self._final_response_completed_outputs = {} # 3. Destroy GPU KV cache (but keep the manager reference for reuse) self._destroy_gpu_paged_kv_cache(empty_cuda_cache=True) diff --git a/batchgen/kv_cache/__init__.py b/batchgen/kv_cache/__init__.py index 16d18fe8b..fd0a38fab 100644 --- a/batchgen/kv_cache/__init__.py +++ b/batchgen/kv_cache/__init__.py @@ -22,6 +22,12 @@ DeepSeekV4GPUKVCoordinator, DeepSeekV4HostKVCoordinator, ) +from batchgen.kv_cache.glm5_kv_coordinator import ( + GLM5_INDEXER_GROUP_ID, + GLM5_PRIMARY_GROUP_ID, + GLM5GPUKVCoordinator, + GLM5HostKVCoordinator, +) from batchgen.kv_cache.swa_gpu_paged_kv_manager import ( SWAGPUPagedKVCacheManager, ) @@ -40,5 +46,9 @@ "CompressedStateGPUStats", "DeepSeekV4GPUKVCoordinator", "DeepSeekV4HostKVCoordinator", + "GLM5_INDEXER_GROUP_ID", + "GLM5_PRIMARY_GROUP_ID", + "GLM5GPUKVCoordinator", + "GLM5HostKVCoordinator", "SWAGPUPagedKVCacheManager", ] diff --git a/batchgen/kv_cache/deepseek_v4_kv_coordinator.py b/batchgen/kv_cache/deepseek_v4_kv_coordinator.py index 8dd00b2ec..b5a32b49e 100644 --- a/batchgen/kv_cache/deepseek_v4_kv_coordinator.py +++ b/batchgen/kv_cache/deepseek_v4_kv_coordinator.py @@ -53,6 +53,18 @@ def __init__( self.compressor_c128_state = compressor_c128_state self.indexer_c4_state = indexer_c4_state + def views_by_group(self) -> dict[int, Any]: + return { + group_id: manager + for group_id, manager in ( + (0, self.swa), + (1, self.compressor_c4), + (2, self.compressor_c128), + (3, self.indexer_c4), + ) + if manager is not None + } + def initialize( self, device_index: int, create_region: bool = False ) -> dict[str, Any]: @@ -105,6 +117,18 @@ def __init__( self.compressor_c128_state = compressor_c128_state self.indexer_c4_state = indexer_c4_state + def managers_by_group(self) -> dict[int, Any]: + return { + group_id: manager + for group_id, manager in ( + (0, self.swa), + (1, self.compressor_c4), + (2, self.compressor_c128), + (3, self.indexer_c4), + ) + if manager is not None + } + def initialize(self) -> dict[str, Any]: results: dict[str, Any] = {} for component_name in _COMPONENT_NAMES: diff --git a/batchgen/kv_cache/dual_host_kv_coordinator.py b/batchgen/kv_cache/dual_host_kv_coordinator.py index 1da74dcde..3da95974b 100644 --- a/batchgen/kv_cache/dual_host_kv_coordinator.py +++ b/batchgen/kv_cache/dual_host_kv_coordinator.py @@ -144,6 +144,9 @@ def __init__(self, primary, auxiliary) -> None: self.primary = primary self.auxiliary = auxiliary + def views_by_group(self) -> dict[int, Any]: + return {0: self.primary, 1: self.auxiliary} + @classmethod def from_budget( cls, diff --git a/batchgen/kv_cache/dual_kv_cache_coordinator.py b/batchgen/kv_cache/dual_kv_cache_coordinator.py index 84dbb0acf..b888fee95 100644 --- a/batchgen/kv_cache/dual_kv_cache_coordinator.py +++ b/batchgen/kv_cache/dual_kv_cache_coordinator.py @@ -52,6 +52,9 @@ def __init__( f"aux={auxiliary.config.page_size_tokens}" ) + def managers_by_group(self) -> dict[int, GPUPagedKVCacheManager]: + return {0: self.primary, 1: self.auxiliary} + # -- Lifecycle -- def initialize(self) -> None: diff --git a/batchgen/kv_cache/glm5_kv_coordinator.py b/batchgen/kv_cache/glm5_kv_coordinator.py new file mode 100644 index 000000000..82cae454d --- /dev/null +++ b/batchgen/kv_cache/glm5_kv_coordinator.py @@ -0,0 +1,633 @@ +"""GLM-5 KV coordinators. + +GLM-5 uses two logical KV groups: + +- group 0: primary MLA compressed KV +- group 1: DSA/indexer KV + +Unlike the legacy dual coordinators, these classes do not require primary and +indexer managers to allocate identical physical page ids. The shared invariant +is the logical sequence set, token/page counts, and active slot order. Prefix +cache metadata keeps the per-group physical page handles separate. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any, Optional, Sequence + +import torch + +from batchgen.config.model_name_utils import is_glm5_backend_model +from batchgen.kv_cache.gpu_paged_kv_manager import ( + GPUPagedKVCacheManager, + GPUPagedKVConfig, + GPUPagedKVStats, +) +from batchgen.kv_cache.host_kv_mananger_config import ( + HOST_KV_AUX_SHM_NAME, + HOST_KV_SHM_NAME, + HostKVGroupProfile, + _dtype_size_bytes, + resolve_host_kv_group_profiles, +) +from batchgen.models.engine_loader import core_engine as bg_lib + +logger = logging.getLogger(__name__) + +GLM5_PRIMARY_GROUP_ID = 0 +GLM5_INDEXER_GROUP_ID = 1 + + +@dataclass +class GLM5AsyncKVTask: + """Composite async task for primary + indexer Host->GPU KV loads.""" + + primary_task: Any + indexer_task: Any + tensors: Any = None + + def wait(self) -> None: + errors = [] + for name, task in ( + ("primary", self.primary_task), + ("indexer", self.indexer_task), + ): + try: + task.wait() + except Exception as exc: + errors.append((name, exc)) + if not errors: + return + if len(errors) == 1: + raise errors[0][1] + names = ", ".join(name for name, _ in errors) + raise RuntimeError( + f"GLM5AsyncKVTask wait failed for KV loads: {names}" + ) from errors[0][1] + + +def is_glm5_dual_kv_model(model_name: str | None) -> bool: + """Return whether this model should use the GLM-5 KV coordinator.""" + + return is_glm5_backend_model(model_name) + + +def _try_set_logger_name(config: Any, name: str) -> None: + try: + config.logger_name = name + except AttributeError: + return + + +def _glm5_group_profiles( + model_name: str, +) -> tuple[HostKVGroupProfile, HostKVGroupProfile]: + if not is_glm5_dual_kv_model(model_name): + raise ValueError(f"Model '{model_name}' is not a GLM-5 KV model") + profiles = { + int(profile.group_id): profile + for profile in resolve_host_kv_group_profiles(model_name) + } + try: + primary = profiles[GLM5_PRIMARY_GROUP_ID] + indexer = profiles[GLM5_INDEXER_GROUP_ID] + except KeyError as exc: + raise ValueError( + f"GLM-5 KV profiles must include groups " + f"{GLM5_PRIMARY_GROUP_ID} and {GLM5_INDEXER_GROUP_ID}" + ) from exc + if primary.raw_page_tokens != indexer.raw_page_tokens: + raise ValueError( + "GLM-5 primary/indexer raw page mismatch: " + f"primary={primary.raw_page_tokens}, indexer={indexer.raw_page_tokens}" + ) + return primary, indexer + + +def _compute_glm5_page_count( + model_name: str, + host_kv_cache_size: int, +) -> tuple[HostKVGroupProfile, HostKVGroupProfile, int]: + primary, indexer = _glm5_group_profiles(model_name) + combined_bytes_per_page = ( + primary.bytes_per_page() * primary.num_layers + + indexer.bytes_per_page() * indexer.num_layers + ) + num_pages = int(host_kv_cache_size) // combined_bytes_per_page + if num_pages <= 0: + raise ValueError( + f"host_kv_cache_size ({host_kv_cache_size}) too small for " + f"GLM-5 KV cache (combined bytes per page = " + f"{combined_bytes_per_page})" + ) + return primary, indexer, num_pages + + +def _build_host_config_from_group( + profile: HostKVGroupProfile, + *, + shm_name: str, + num_pages: int, +) -> Any: + config = bg_lib.HostPagedKVConfig() + config.shm_name = shm_name + config.num_layers = profile.num_layers + config.num_pages = int(num_pages) + config.page_size_tokens = profile.storage_page_tokens + config.num_k_heads = profile.num_k_heads + config.k_head_dim = profile.k_head_dim + config.num_v_heads = profile.num_v_heads + config.v_head_dim = profile.v_head_dim + config.k_element_size_bytes = _dtype_size_bytes(profile.kv_dtype) + config.v_element_size_bytes = ( + 0 if profile.num_v_heads == 0 else config.k_element_size_bytes + ) + config.sequence_table_capacity = ( + profile.sequence_table_capacity or config.num_pages + ) + config.alignment_bytes = profile.alignment_bytes + return config + + +class GLM5HostKVCoordinator: + """Host-side GLM-5 KV facade with independent primary/indexer pages.""" + + def __init__(self, primary: Any, indexer: Any) -> None: + self.primary = primary + self.indexer = indexer + self.auxiliary = indexer + + def views_by_group(self) -> dict[int, Any]: + return { + GLM5_PRIMARY_GROUP_ID: self.primary, + GLM5_INDEXER_GROUP_ID: self.indexer, + } + + @classmethod + def from_budget( + cls, + *, + model_name: str, + host_kv_cache_size: int, + core_engine_module: Any, + enable_memfd: bool = False, + memfd_creator_pid: int = -1, + memfd_fd: int = -1, + aux_memfd_fd: int = -1, + ) -> Optional["GLM5HostKVCoordinator"]: + if not is_glm5_dual_kv_model(model_name): + return None + primary_profile, indexer_profile, num_pages = _compute_glm5_page_count( + model_name, host_kv_cache_size + ) + primary_config = _build_host_config_from_group( + primary_profile, + shm_name=HOST_KV_SHM_NAME, + num_pages=num_pages, + ) + indexer_config = _build_host_config_from_group( + indexer_profile, + shm_name=HOST_KV_AUX_SHM_NAME, + num_pages=num_pages, + ) + _try_set_logger_name(primary_config, "GLM5HostPagedKVWorkerView") + _try_set_logger_name(indexer_config, "GLM5IndexerHostPagedKVWorkerView") + + if enable_memfd: + primary_config.enable_memfd = True + primary_config.memfd_creator_pid = memfd_creator_pid + primary_config.memfd_fd = memfd_fd + indexer_config.enable_memfd = True + indexer_config.memfd_creator_pid = memfd_creator_pid + indexer_config.memfd_fd = aux_memfd_fd + + primary_view = core_engine_module.MLAHostPagedKVWorkerView( + primary_config + ) + indexer_view = core_engine_module.MLAHostPagedKVWorkerView( + indexer_config + ) + logger.info( + "GLM5HostKVCoordinator created: %d pages, primary dim=%d, " + "indexer dim=%d", + num_pages, + primary_profile.k_head_dim, + indexer_profile.k_head_dim, + ) + return cls(primary_view, indexer_view) + + @classmethod + def create_managers( + cls, + *, + model_name: str, + host_kv_cache_size: int, + enable_memfd: bool = False, + ) -> Optional[tuple[Any, Any]]: + if not is_glm5_dual_kv_model(model_name): + return None + primary_profile, indexer_profile, num_pages = _compute_glm5_page_count( + model_name, host_kv_cache_size + ) + primary_config = _build_host_config_from_group( + primary_profile, + shm_name=HOST_KV_SHM_NAME, + num_pages=num_pages, + ) + indexer_config = _build_host_config_from_group( + indexer_profile, + shm_name=HOST_KV_AUX_SHM_NAME, + num_pages=num_pages, + ) + _try_set_logger_name(primary_config, "GLM5HostPagedKVManager") + _try_set_logger_name(indexer_config, "GLM5IndexerHostPagedKVManager") + + if enable_memfd: + primary_config.enable_memfd = True + indexer_config.enable_memfd = True + + primary_manager = bg_lib.MLAHostPagedKVManager(primary_config) + primary_manager.initialize(True) + indexer_manager = bg_lib.MLAHostPagedKVManager(indexer_config) + indexer_manager.initialize(True) + logger.info( + "GLM5HostKVCoordinator managers created: %d pages, primary dim=%d, " + "indexer dim=%d", + num_pages, + primary_profile.k_head_dim, + indexer_profile.k_head_dim, + ) + return primary_manager, indexer_manager + + def initialize(self, **kwargs: Any) -> None: + self.primary.initialize(**kwargs) + self.indexer.initialize(**kwargs) + + def register_sequences(self, sequence_ids: Sequence[int]) -> None: + self.primary.register_sequences(sequence_ids) + try: + self.indexer.register_sequences(sequence_ids) + except Exception: + self.primary.unregister_sequences(sequence_ids) + raise + + def allocate_pages_for_sequences(self, seq_token_pairs: Sequence[Any]) -> None: + pairs = list(seq_token_pairs) + sequence_ids = [int(seq_id) for seq_id, _ in pairs] + self.primary.allocate_pages_for_sequences(pairs) + try: + self.indexer.allocate_pages_for_sequences(pairs) + except Exception: + self.primary.release_sequence_pages(sequence_ids) + raise + + def grow_pages_for_sequences(self, seq_page_pairs: Sequence[Any]) -> None: + pairs = list(seq_page_pairs) + needed = sum(int(pages) for _, pages in pairs) + primary_free = int(self.primary.get_stats().num_free_pages) + indexer_free = int(self.indexer.get_stats().num_free_pages) + if needed > primary_free or needed > indexer_free: + raise RuntimeError( + "GLM-5 grow_pages_for_sequences: insufficient Host KV free " + f"pages: need={needed}, primary_free={primary_free}, " + f"indexer_free={indexer_free}" + ) + self.primary.grow_pages_for_sequences(pairs) + self.indexer.grow_pages_for_sequences(pairs) + + def release_sequence_pages(self, sequence_ids: Sequence[int]) -> None: + self.primary.release_sequence_pages(sequence_ids) + self.indexer.release_sequence_pages(sequence_ids) + + def unregister_sequences(self, sequence_ids: Sequence[int]) -> None: + self.primary.unregister_sequences(sequence_ids) + self.indexer.unregister_sequences(sequence_ids) + + def get_stats(self) -> Any: + primary_stats = self.primary.get_stats() + indexer_stats = self.indexer.get_stats() + if indexer_stats.num_free_pages < primary_stats.num_free_pages: + return indexer_stats + return primary_stats + + def async_load_layer_paged_kv_to_device(self, **kwargs: Any) -> None: + raise RuntimeError( + "GLM-5 KV load must use " + "async_load_layer_paged_kv_to_device_dual()" + ) + + def async_load_layer_paged_kv_to_device_dual( + self, + *, + sequence_ids: torch.Tensor, + primary_active_page_counts: torch.Tensor, + primary_k_device_ptrs: torch.Tensor, + primary_v_device_ptrs: Optional[torch.Tensor], + aux_active_page_counts: torch.Tensor, + aux_k_device_ptrs: torch.Tensor, + aux_v_device_ptrs: Optional[torch.Tensor], + tensors: Any = None, + ) -> GLM5AsyncKVTask: + if primary_active_page_counts.tolist() != aux_active_page_counts.tolist(): + raise RuntimeError( + "GLM-5 primary/indexer load page-count mismatch: " + f"primary={primary_active_page_counts.tolist()}, " + f"indexer={aux_active_page_counts.tolist()}" + ) + primary_task = self.primary.async_load_layer_paged_kv_to_device( + sequence_ids=sequence_ids, + active_page_counts=primary_active_page_counts, + k_device_ptrs=primary_k_device_ptrs, + v_device_ptrs=primary_v_device_ptrs, + ) + try: + indexer_task = self.indexer.async_load_layer_paged_kv_to_device( + sequence_ids=sequence_ids, + active_page_counts=aux_active_page_counts, + k_device_ptrs=aux_k_device_ptrs, + v_device_ptrs=aux_v_device_ptrs, + ) + except Exception: + primary_task.wait() + raise + return GLM5AsyncKVTask( + primary_task=primary_task, + indexer_task=indexer_task, + tensors=tensors, + ) + + def async_offload_layer_kv_to_host(self, **kwargs: Any) -> None: + raise RuntimeError( + "GLM-5 KV offload must explicitly offload primary and indexer KV" + ) + + +class GLM5GPUKVCoordinator: + """GPU-side GLM-5 KV facade with per-group physical page ownership.""" + + def __init__( + self, + primary: GPUPagedKVCacheManager, + indexer: GPUPagedKVCacheManager, + ) -> None: + self.primary = primary + self.indexer = indexer + self.auxiliary = indexer + if primary.config.page_size_tokens != indexer.config.page_size_tokens: + raise ValueError( + "GLM-5 primary/indexer GPU KV page size mismatch: " + f"primary={primary.config.page_size_tokens}, " + f"indexer={indexer.config.page_size_tokens}" + ) + + def managers_by_group(self) -> dict[int, GPUPagedKVCacheManager]: + return { + GLM5_PRIMARY_GROUP_ID: self.primary, + GLM5_INDEXER_GROUP_ID: self.indexer, + } + + def initialize(self) -> None: + self.primary.initialize() + self.indexer.initialize() + logger.info( + "GLM5GPUKVCoordinator initialized: primary=%s, indexer=%s", + self.primary.get_stats(), + self.indexer.get_stats(), + ) + + def destroy(self, *, empty_cuda_cache: bool = False) -> None: + self.primary.destroy(empty_cuda_cache=empty_cuda_cache) + self.indexer.destroy(empty_cuda_cache=empty_cuda_cache) + + @property + def is_initialized(self) -> bool: + return self.primary.is_initialized and self.indexer.is_initialized + + def allocate_pages_for_sequences( + self, + sequence_ids: Sequence[int], + num_tokens: Sequence[int], + ) -> Any: + result = self.primary.allocate_pages_for_sequences( + sequence_ids, num_tokens + ) + try: + self.indexer.allocate_pages_for_sequences(sequence_ids, num_tokens) + except Exception: + self._rollback_primary_allocations(result) + raise + self.assert_aligned_state("allocate_pages_for_sequences", sequence_ids) + return result + + def grow_pages_for_sequences( + self, + sequence_ids: Sequence[int], + additional_tokens: Sequence[int], + ) -> Any: + needed = sum(int(tokens) for tokens in additional_tokens) + primary_free = int(self.primary.get_stats().num_free_pages) + indexer_free = int(self.indexer.get_stats().num_free_pages) + if needed > primary_free or needed > indexer_free: + raise RuntimeError( + "GLM-5 grow_pages_for_sequences: insufficient GPU KV free " + f"pages: need={needed}, primary_free={primary_free}, " + f"indexer_free={indexer_free}" + ) + result = self.primary.grow_pages_for_sequences( + sequence_ids, additional_tokens + ) + self.indexer.grow_pages_for_sequences(sequence_ids, additional_tokens) + self.assert_aligned_state("grow_pages_for_sequences", sequence_ids) + return result + + def extend_pages_for_sequence( + self, + sequence_id: int, + new_total_tokens: int, + ) -> int: + primary_state = self.primary._sequences.get(sequence_id) + indexer_state = self.indexer._sequences.get(sequence_id) + if primary_state is None or indexer_state is None: + raise KeyError( + f"extend_pages_for_sequence: GLM-5 sequence {sequence_id} " + "is not allocated in both primary and indexer managers" + ) + primary_required = int( + self.primary._geometry.required_pages(new_total_tokens) + ) + indexer_required = int( + self.indexer._geometry.required_pages(new_total_tokens) + ) + primary_missing = max(0, primary_required - int(primary_state.pages.numel())) + indexer_missing = max(0, indexer_required - int(indexer_state.pages.numel())) + if primary_missing != indexer_missing: + raise RuntimeError( + "GLM-5 primary/indexer page growth mismatch: " + f"primary_missing={primary_missing}, " + f"indexer_missing={indexer_missing}" + ) + if primary_missing <= 0: + return 0 + if primary_missing > self.primary.get_stats().num_free_pages: + raise RuntimeError( + "GLM-5 primary GPU KV has insufficient free pages: " + f"need={primary_missing}" + ) + if indexer_missing > self.indexer.get_stats().num_free_pages: + raise RuntimeError( + "GLM-5 indexer GPU KV has insufficient free pages: " + f"need={indexer_missing}" + ) + added = self.primary.extend_pages_for_sequence( + sequence_id, new_total_tokens + ) + self.indexer.extend_pages_for_sequence(sequence_id, new_total_tokens) + self.assert_aligned_state("extend_pages_for_sequence", [sequence_id]) + return added + + def rebuild_page_table(self, sequence_ids: Sequence[int]) -> torch.Tensor: + table = self.primary.rebuild_page_table(sequence_ids) + self.indexer.rebuild_page_table(sequence_ids) + self._assert_slot_order("rebuild_page_table") + return table + + def clear_page_table(self) -> None: + self.primary.clear_page_table() + self.indexer.clear_page_table() + + def free_pages_for_sequences(self, sequence_ids: Sequence[int]) -> None: + self.primary.free_pages_for_sequences(sequence_ids) + self.indexer.free_pages_for_sequences(sequence_ids) + + def get_stats(self) -> GPUPagedKVStats: + primary_stats = self.primary.get_stats() + indexer_stats = self.indexer.get_stats() + if indexer_stats.num_free_pages < primary_stats.num_free_pages: + return indexer_stats + return primary_stats + + def get_page_table_version(self) -> int: + return self.primary.get_page_table_version() + + @property + def config(self) -> GPUPagedKVConfig: + return self.primary.config + + @property + def device(self) -> torch.device: + return self.primary.device + + @property + def _gpu_page_table_manager(self) -> Any: + return self.primary._gpu_page_table_manager + + @property + def _sequences(self) -> Any: + return self.primary._sequences + + def copy_kv_to_tensor(self, sequence_id: int) -> torch.Tensor: + return self.primary.copy_kv_to_tensor(sequence_id) + + def copy_tensor_to_kv(self, sequence_id: int, k_tensor: torch.Tensor) -> None: + self.primary.copy_tensor_to_kv(sequence_id, k_tensor) + + def get_context_kv_page_ptrs(self, *args: Any, **kwargs: Any) -> Any: + return self.primary.get_context_kv_page_ptrs(*args, **kwargs) + + def get_sequence_layer_page_pointers(self, *args: Any, **kwargs: Any) -> Any: + return self.primary.get_sequence_layer_page_pointers(*args, **kwargs) + + def export_layer_page_pointer_table(self, *args: Any, **kwargs: Any) -> Any: + return self.primary.export_layer_page_pointer_table(*args, **kwargs) + + def get_kv_tensors(self) -> Any: + raise RuntimeError( + "GLM5GPUKVCoordinator.get_kv_tensors() is primary-only; " + "use .primary or .indexer explicitly" + ) + + def get_layer_kv_with_page_table(self, layer_idx: int) -> Any: + raise RuntimeError( + "GLM5GPUKVCoordinator.get_layer_kv_with_page_table() is " + "primary-only; use .primary or .indexer explicitly" + ) + + def export_active_sequence_page_counts(self) -> torch.Tensor: + raise RuntimeError( + "GLM5GPUKVCoordinator.export_active_sequence_page_counts() is " + "primary-only; use .primary or .indexer explicitly" + ) + + def get_padded_3d_page_pointers(self) -> Any: + raise RuntimeError( + "GLM5GPUKVCoordinator.get_padded_3d_page_pointers() is " + "primary-only; use .primary or .indexer explicitly" + ) + + def assert_aligned_state( + self, + op_name: str, + sequence_ids: Optional[Sequence[int]] = None, + ) -> None: + primary_ids = set(self.primary._sequences.keys()) + indexer_ids = set(self.indexer._sequences.keys()) + if primary_ids != indexer_ids: + raise RuntimeError( + f"{op_name}: GLM-5 primary/indexer sequence set mismatch: " + f"primary_only={sorted(primary_ids - indexer_ids)[:10]}, " + f"indexer_only={sorted(indexer_ids - primary_ids)[:10]}" + ) + check_ids = ( + list(sequence_ids) if sequence_ids is not None else sorted(primary_ids) + ) + for seq_id in check_ids: + primary_state = self.primary._sequences.get(seq_id) + indexer_state = self.indexer._sequences.get(seq_id) + if primary_state is None or indexer_state is None: + continue + primary_pages = int(primary_state.pages.numel()) + indexer_pages = int(indexer_state.pages.numel()) + if primary_pages != indexer_pages: + raise RuntimeError( + f"{op_name}: GLM-5 primary/indexer page-count mismatch " + f"for seq {seq_id}: primary={primary_pages}, " + f"indexer={indexer_pages}" + ) + self._assert_slot_order(op_name) + + def _assert_slot_order(self, op_name: str) -> None: + primary_slots = list(self.primary._gpu_page_table_manager.slot_to_seq_id) + indexer_slots = list(self.indexer._gpu_page_table_manager.slot_to_seq_id) + if primary_slots != indexer_slots: + raise RuntimeError( + f"{op_name}: GLM-5 primary/indexer slot order mismatch: " + f"primary={primary_slots[:10]}, indexer={indexer_slots[:10]}" + ) + + def _rollback_primary_allocations(self, allocations: Any) -> None: + if not allocations: + return + reclaimed = [] + for seq_id, pages in allocations.items(): + if not pages: + continue + state = self.primary._sequences.get(seq_id) + if state is None: + continue + count = len(pages) + tail = state.pages[-count:].tolist() + if tail != pages: + raise RuntimeError( + f"Cannot rollback GLM-5 primary KV allocation for seq " + f"{seq_id}: tail={tail}, allocated={pages}" + ) + reclaimed.append(state.pages[-count:].clone()) + if state.pages.numel() == count: + del self.primary._sequences[seq_id] + else: + state.pages = state.pages[:-count].clone() + if reclaimed: + self.primary._free_pages.push(torch.cat(reclaimed, dim=0)) + self.primary._clear_active_page_pointer_tables() diff --git a/batchgen/kv_cache/gpu_paged_kv_manager.py b/batchgen/kv_cache/gpu_paged_kv_manager.py index 079974377..34ae158bb 100644 --- a/batchgen/kv_cache/gpu_paged_kv_manager.py +++ b/batchgen/kv_cache/gpu_paged_kv_manager.py @@ -113,6 +113,27 @@ class CUDAGraphPageTableState: rebuild_version: int +@dataclass(frozen=True) +class GPUPagedKVSuffixAppendPlan: + """Destination metadata for multi-token suffix writes into GPU paged KV.""" + + sequence_ids: List[int] + prefix_values: Tuple[int, ...] + suffix_values: Tuple[int, ...] + slot_values: Tuple[int, ...] + total_suffix_tokens: int + prefix_lens: torch.Tensor + suffix_lens: torch.Tensor + cache_seqlens: torch.Tensor + token_starts: torch.Tensor + slot_indices: torch.Tensor + page_table: torch.Tensor + + @property + def batch_size(self) -> int: + return len(self.sequence_ids) + + @dataclass(frozen=True) class GPUPagedKVConfig: num_layers: int @@ -759,6 +780,8 @@ def destroy(self, *, empty_cuda_cache: bool = False) -> None: "GPUPagedKVCacheManager.destroy called while uninitialized; " "no-op (state was already reset by a prior destroy call)" ) + if empty_cuda_cache: + self._release_cached_cuda_memory() return self._reset_runtime_state() @@ -890,6 +913,170 @@ def grow_pages_for_sequences( self._clear_active_page_pointer_tables() return allocations + def prepare_prefill_suffix_append( + self, + *, + sequence_ids: Sequence[int], + prefix_lens: Sequence[int] | torch.Tensor, + suffix_lens: Sequence[int] | torch.Tensor, + rebuild_page_table: bool = True, + ) -> GPUPagedKVSuffixAppendPlan: + """Prepare page-table metadata for multi-token prefill suffix writes. + + The returned plan maps each suffix segment to destination token positions + ``[prefix_len, prefix_len + suffix_len)`` for the matching sequence. + Sequences with reused prefixes must already have full-context GPU pages + allocated by the caller; miss sequences without an allocation are + allocated normally. + """ + + self._ensure_initialized() + sequence_ids = [int(seq_id) for seq_id in sequence_ids] + prefix_values = self._normalize_cpu_int_vector( + prefix_lens, + expected_len=len(sequence_ids), + name="prefix_lens", + allow_zero=True, + ) + suffix_values = self._normalize_cpu_int_vector( + suffix_lens, + expected_len=len(sequence_ids), + name="suffix_lens", + allow_zero=True, + ) + if not sequence_ids: + raise ValueError("prepare_prefill_suffix_append: sequence_ids must be non-empty") + + full_lengths = [ + int(prefix_len) + int(suffix_len) + for prefix_len, suffix_len in zip(prefix_values, suffix_values) + ] + for seq_id, prefix_len, suffix_len, full_len in zip( + sequence_ids, + prefix_values, + suffix_values, + full_lengths, + ): + if full_len <= 0: + raise ValueError( + "prepare_prefill_suffix_append: full sequence length must be " + f"positive for seq {seq_id}, got prefix={prefix_len}, suffix={suffix_len}" + ) + state = self._sequences.get(seq_id) + if state is None: + if prefix_len > 0: + raise KeyError( + "prepare_prefill_suffix_append: prefix-reused sequence " + f"{seq_id} is not allocated on GPU" + ) + self.allocate_pages(seq_id, full_len) + continue + required_pages = int(self._geometry.required_pages(full_len)) + missing_pages = max(0, required_pages - int(state.pages.numel())) + if missing_pages > 0: + self.grow_sequence_pages(seq_id, missing_pages) + + if rebuild_page_table: + page_table = self.rebuild_page_table(sequence_ids) + else: + page_table = self._gpu_page_table_manager.gpu_table + if page_table is None: + raise RuntimeError( + "prepare_prefill_suffix_append: GPU page table is not initialized" + ) + + slot_indices = [] + for seq_id in sequence_ids: + slot = self._gpu_page_table_manager.seq_id_to_slot.get(seq_id) + if slot is None: + raise RuntimeError( + "prepare_prefill_suffix_append: missing page-table slot for " + f"sequence {seq_id}" + ) + slot_indices.append(int(slot)) + + return GPUPagedKVSuffixAppendPlan( + sequence_ids=sequence_ids, + prefix_values=tuple(prefix_values), + suffix_values=tuple(suffix_values), + slot_values=tuple(slot_indices), + total_suffix_tokens=sum(suffix_values), + prefix_lens=torch.tensor(prefix_values, dtype=torch.int32, device=self.device), + suffix_lens=torch.tensor(suffix_values, dtype=torch.int32, device=self.device), + cache_seqlens=torch.tensor(full_lengths, dtype=torch.int32, device=self.device), + token_starts=torch.tensor(prefix_values, dtype=torch.int32, device=self.device), + slot_indices=torch.tensor(slot_indices, dtype=torch.int32, device=self.device), + page_table=page_table, + ) + + def append_layer_prefill_suffix_tokens( + self, + *, + k_tensor: torch.Tensor, + v_tensor: Optional[torch.Tensor], + append_plan: GPUPagedKVSuffixAppendPlan, + layer_idx: int, + ) -> None: + """Write flattened multi-token suffix K/V into GPU paged KV.""" + + op_name = "append_layer_prefill_suffix_tokens" + self._ensure_initialized() + layer_idx = self.resolve_physical_layer(layer_idx) + self._geometry.ensure_layer_bounds(layer_idx, op_name) + k_tensor = self._prepare_flat_suffix_tensor( + k_tensor, + expected_heads=self.config.num_k_heads, + expected_dim=self.config.k_head_dim, + expected_tokens=append_plan.total_suffix_tokens, + name="k_tensor", + op_name=op_name, + ) + if v_tensor is not None: + if not self.config.has_v_cache: + raise ValueError(f"{op_name}: V tensor provided but V cache disabled") + v_tensor = self._prepare_flat_suffix_tensor( + v_tensor, + expected_heads=int(self.config.num_v_heads), + expected_dim=int(self.config.v_head_dim), + expected_tokens=append_plan.total_suffix_tokens, + name="v_tensor", + op_name=op_name, + ) + elif self.config.has_v_cache: + logging.debug("%s: V cache enabled but v_tensor is None", op_name) + + k_layer = self._k_cache[layer_idx] + v_layer = self._v_cache[layer_idx] if self._v_cache is not None else None + source_offset = 0 + for seq_id, prefix_len, suffix_len, slot_idx in zip( + append_plan.sequence_ids, + append_plan.prefix_values, + append_plan.suffix_values, + append_plan.slot_values, + ): + end_offset = source_offset + int(suffix_len) + if suffix_len > 0: + self._write_token_range_to_cache_by_page_table( + cache_layer=k_layer, + page_table=append_plan.page_table, + slot_index=int(slot_idx), + sequence_id=seq_id, + token_start=int(prefix_len), + values=k_tensor[source_offset:end_offset], + context=op_name, + ) + if v_layer is not None and v_tensor is not None: + self._write_token_range_to_cache_by_page_table( + cache_layer=v_layer, + page_table=append_plan.page_table, + slot_index=int(slot_idx), + sequence_id=seq_id, + token_start=int(prefix_len), + values=v_tensor[source_offset:end_offset], + context=op_name, + ) + source_offset = end_offset + def clear_page_table(self) -> None: """Clear the GPU page table to empty state (0 sequences). @@ -1432,6 +1619,116 @@ def _get_sequence_state(self, sequence_id: int) -> _SequenceState: raise KeyError(f"Sequence {sequence_id} not registered on GPU") return state + def _normalize_cpu_int_vector( + self, + values: Sequence[int] | torch.Tensor, + *, + expected_len: int, + name: str, + allow_zero: bool, + ) -> List[int]: + tensor = torch.as_tensor(values, dtype=torch.long, device="cpu") + if tensor.dim() != 1: + raise ValueError( + f"{name} must be 1-D, got shape={tuple(tensor.shape)}" + ) + if tensor.numel() != expected_len: + raise ValueError( + f"{name} length must match sequence_ids: " + f"{tensor.numel()} != {expected_len}" + ) + limit_ok = tensor >= 0 if allow_zero else tensor > 0 + if not bool(torch.all(limit_ok).item()): + requirement = "non-negative" if allow_zero else "positive" + raise ValueError(f"{name} values must be {requirement}") + return [int(value) for value in tensor.tolist()] + + def _prepare_flat_suffix_tensor( + self, + tensor: torch.Tensor, + *, + expected_heads: int, + expected_dim: int, + expected_tokens: int, + name: str, + op_name: str, + ) -> torch.Tensor: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{op_name}: {name} must be a torch.Tensor") + if tensor.dim() == 2 and expected_heads == 1: + tensor = tensor.unsqueeze(1) + if tensor.dim() != 3: + raise ValueError( + f"{op_name}: {name} must have shape [tokens, heads, dim], " + f"got {tuple(tensor.shape)}" + ) + if tensor.device != self.device: + raise ValueError(f"{op_name}: {name} must be on device {self.device}") + if int(tensor.shape[0]) != int(expected_tokens): + raise ValueError( + f"{op_name}: {name} token count mismatch: " + f"{tensor.shape[0]} != {expected_tokens}" + ) + if int(tensor.shape[1]) != int(expected_heads) or int(tensor.shape[2]) != int(expected_dim): + raise ValueError( + f"{op_name}: {name} head shape mismatch, got " + f"{tuple(tensor.shape[1:])}, expected " + f"({int(expected_heads)}, {int(expected_dim)})" + ) + return tensor.contiguous() + + def _write_token_range_to_cache_by_page_table( + self, + *, + cache_layer: torch.Tensor, + page_table: torch.Tensor, + slot_index: int, + sequence_id: int, + token_start: int, + values: torch.Tensor, + context: str, + ) -> None: + if values.numel() == 0: + return + if page_table is None: + raise RuntimeError(f"{context}: append plan has no page_table") + if page_table.ndim != 2: + raise ValueError( + f"{context}: page_table must be 2-D, got {tuple(page_table.shape)}" + ) + if slot_index < 0 or slot_index >= page_table.shape[0]: + raise ValueError( + f"{context}: slot index {slot_index} is outside page_table rows " + f"{page_table.shape[0]}" + ) + if token_start < 0: + raise ValueError(f"{context}: token_start must be non-negative") + page_size = self.config.page_size_tokens + remaining = int(values.shape[0]) + source_offset = 0 + token_index = int(token_start) + while remaining > 0: + page_slot = token_index // page_size + if page_slot >= page_table.shape[1]: + raise RuntimeError( + f"{context}: sequence {sequence_id} token range exceeds " + f"page_table width {page_table.shape[1]}" + ) + gpu_page = int(page_table[slot_index, page_slot].item()) + if gpu_page < 0: + raise RuntimeError( + f"{context}: sequence {sequence_id} slot {slot_index} " + f"has no GPU page for logical page {page_slot}" + ) + page_offset = token_index % page_size + take = min(remaining, page_size - page_offset) + cache_layer[gpu_page, page_offset : page_offset + take].copy_( + values[source_offset : source_offset + take] + ) + remaining -= take + source_offset += take + token_index += take + def _validate_token_inputs( self, k_tensor: torch.Tensor, @@ -1867,6 +2164,43 @@ def get_padded_3d_page_pointers( return k_tensor, v_tensor + def get_page_pointer_matrix( + self, + gpu_pages: Sequence[int] | torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Return layer-major device pointer matrices for explicit GPU pages. + + The returned tensors are CPU ``int64`` matrices shaped + ``[num_layers, num_pages]``. They are suitable for C++ page-level H2D + copy APIs that take explicit destination pages instead of active + sequence layouts. + """ + + self._ensure_initialized() + pages = torch.as_tensor(gpu_pages, dtype=torch.long, device="cpu") + if pages.dim() != 1: + raise ValueError( + "get_page_pointer_matrix: gpu_pages must be 1-D, " + f"got shape={tuple(pages.shape)}" + ) + if pages.numel() == 0: + empty = self._k_page_ptr_table.new_empty( + (self.config.num_layers, 0) + ) + return empty, None if self._v_page_ptr_table is None else empty.clone() + if torch.any(pages < 0) or torch.any(pages >= self.config.num_pages): + raise ValueError( + "get_page_pointer_matrix: gpu_pages contains out-of-range page IDs" + ) + + k_ptrs = self._select_active_page_columns(self._k_page_ptr_table, pages) + v_ptrs = None + if self.config.has_v_cache: + v_ptrs = self._select_active_page_columns( + self._v_page_ptr_table, pages + ) + return k_ptrs.contiguous(), None if v_ptrs is None else v_ptrs.contiguous() + # In gpu_paged_kv_manager.py def extend_pages_for_sequence(self, sequence_id: int, new_total_tokens: int) -> int: self._ensure_initialized() diff --git a/batchgen/kv_cache/host_kv_mananger_config.py b/batchgen/kv_cache/host_kv_mananger_config.py index a8ec16aef..701505b15 100644 --- a/batchgen/kv_cache/host_kv_mananger_config.py +++ b/batchgen/kv_cache/host_kv_mananger_config.py @@ -9,14 +9,25 @@ from batchgen.kv_cache.gpu_paged_kv_manager import GPUPagedKVConfig from batchgen.models.engine_loader import core_engine as bg_lib -HOST_KV_SHM_NAME = "batchgen_host_kv_cache" +HOST_KV_SHM_NAME = os.environ.get( + "BATCHGEN_HOST_KV_SHM_NAME", "batchgen_host_kv_cache" +) __all__ = [ "build_host_kv_config", + "build_host_kv_config_from_group_profile", "build_gpu_kv_config", + "build_gpu_kv_config_from_group_profile", + "HostKVGroupProfile", + "resolve_host_kv_group_profiles", "HOST_KV_SHM_NAME", ] +HOST_KV_SEMANTIC_FULL_KV = "full_kv" +HOST_KV_SEMANTIC_MLA_COMPRESSED_KV = "mla_compressed_kv" +HOST_KV_SEMANTIC_SWA_KV = "swa_kv" +HOST_KV_SEMANTIC_COMPRESSED_RATIO_KV = "compressed_ratio_kv" + def _dtype_size_bytes(dtype: str) -> int: """Returns the storage size in bytes for the provided dtype string.""" @@ -68,6 +79,41 @@ def bytes_per_page(self) -> int: return k_bytes + v_bytes +@dataclass(frozen=True) +class HostKVGroupProfile: + group_id: int + group_name: str + semantic: str + required_for_reuse: bool + num_layers: int + num_k_heads: int + k_head_dim: int + storage_page_tokens: int + raw_page_tokens: int + compression_ratio: int = 1 + num_v_heads: int = 0 + v_head_dim: int = 0 + kv_dtype: str = "bfloat16" + sequence_table_capacity: int | None = None + alignment_bytes: int = 64 + + def bytes_per_page(self) -> int: + element_bytes = _dtype_size_bytes(self.kv_dtype) + k_bytes = ( + self.storage_page_tokens + * self.num_k_heads + * self.k_head_dim + * element_bytes + ) + v_bytes = ( + self.storage_page_tokens + * self.num_v_heads + * self.v_head_dim + * element_bytes + ) + return k_bytes + v_bytes + + _DEEPSEEK_MLA_PROFILE = _HostKVModelProfile( num_layers=61, num_k_heads=1, @@ -156,6 +202,115 @@ def bytes_per_page(self) -> int: "glm5_indexer": _GLM5_INDEXER_PROFILE, } + +def _legacy_group_profile( + *, + group_id: int, + group_name: str, + profile: _HostKVModelProfile, + required_for_reuse: bool, +) -> HostKVGroupProfile: + semantic = ( + HOST_KV_SEMANTIC_MLA_COMPRESSED_KV + if int(profile.num_v_heads) == 0 + else HOST_KV_SEMANTIC_FULL_KV + ) + return HostKVGroupProfile( + group_id=group_id, + group_name=group_name, + semantic=semantic, + required_for_reuse=required_for_reuse, + num_layers=profile.num_layers, + num_k_heads=profile.num_k_heads, + k_head_dim=profile.k_head_dim, + storage_page_tokens=profile.page_size, + raw_page_tokens=profile.page_size, + compression_ratio=1, + num_v_heads=profile.num_v_heads, + v_head_dim=profile.v_head_dim, + kv_dtype=profile.kv_dtype, + sequence_table_capacity=profile.sequence_table_capacity, + alignment_bytes=profile.alignment_bytes, + ) + + +_DEEPSEEK_V4_FLASH_GROUP_PROFILES: tuple[HostKVGroupProfile, ...] = ( + HostKVGroupProfile( + group_id=0, + group_name="swa", + semantic=HOST_KV_SEMANTIC_SWA_KV, + required_for_reuse=True, + num_layers=43, + num_k_heads=1, + k_head_dim=512, + storage_page_tokens=64, + raw_page_tokens=64, + compression_ratio=1, + ), + HostKVGroupProfile( + group_id=1, + group_name="compressor_c4", + semantic=HOST_KV_SEMANTIC_COMPRESSED_RATIO_KV, + required_for_reuse=True, + num_layers=43, + num_k_heads=1, + k_head_dim=512, + storage_page_tokens=64, + raw_page_tokens=256, + compression_ratio=4, + ), + HostKVGroupProfile( + group_id=2, + group_name="compressor_c128", + semantic=HOST_KV_SEMANTIC_COMPRESSED_RATIO_KV, + required_for_reuse=True, + num_layers=43, + num_k_heads=1, + k_head_dim=512, + storage_page_tokens=2, + raw_page_tokens=256, + compression_ratio=128, + ), + HostKVGroupProfile( + group_id=3, + group_name="indexer_c4", + semantic=HOST_KV_SEMANTIC_COMPRESSED_RATIO_KV, + required_for_reuse=True, + num_layers=43, + num_k_heads=1, + k_head_dim=128, + storage_page_tokens=64, + raw_page_tokens=256, + compression_ratio=4, + ), +) + +_DEEPSEEK_V4_PRO_GROUP_PROFILES: tuple[HostKVGroupProfile, ...] = tuple( + HostKVGroupProfile( + group_id=profile.group_id, + group_name=profile.group_name, + semantic=profile.semantic, + required_for_reuse=profile.required_for_reuse, + num_layers=61, + num_k_heads=profile.num_k_heads, + k_head_dim=profile.k_head_dim, + storage_page_tokens=profile.storage_page_tokens, + raw_page_tokens=profile.raw_page_tokens, + compression_ratio=profile.compression_ratio, + num_v_heads=profile.num_v_heads, + v_head_dim=profile.v_head_dim, + kv_dtype=profile.kv_dtype, + sequence_table_capacity=profile.sequence_table_capacity, + alignment_bytes=profile.alignment_bytes, + ) + for profile in _DEEPSEEK_V4_FLASH_GROUP_PROFILES +) + +_GROUP_PROFILE_REGISTRY: Dict[str, tuple[HostKVGroupProfile, ...]] = { + "deepseek_v4_flash": _DEEPSEEK_V4_FLASH_GROUP_PROFILES, + "deepseek_v4_pro": _DEEPSEEK_V4_PRO_GROUP_PROFILES, +} + _PROFILE_ALIASES: Dict[str, str] = {} for canonical, aliases in { "deepseek_mla": ( @@ -249,12 +404,54 @@ def _resolve_indexer_profile(model_name: str) -> _HostKVModelProfile | None: def _resolve_profile(model_name: str) -> _HostKVModelProfile: """Maps a user supplied model name to a cached profile.""" + return _PROFILE_REGISTRY[_resolve_profile_key(model_name)] + + +def _resolve_profile_key(model_name: str) -> str: + """Maps a user supplied model name to its canonical profile key.""" + if not isinstance(model_name, str): raise ValueError("model_name must be a string") alias = model_name.strip().lower() if alias not in _PROFILE_ALIASES: raise ValueError(f"Unsupported model '{model_name}' for host KV cache") - return _PROFILE_REGISTRY[_PROFILE_ALIASES[alias]] + return _PROFILE_ALIASES[alias] + + +def resolve_host_kv_group_profiles( + model_name: str, +) -> tuple[HostKVGroupProfile, ...]: + """Return logical Host KV groups required to reuse a model prefix. + + Most existing models have one primary Host KV group plus an optional DSA + indexer group. Multi-rate models, such as DeepSeek-V4, override this with a + model-specific group profile that records each reusable KV component's raw + token boundary and physical storage page shape. + """ + + profile_key = _resolve_profile_key(model_name) + if profile_key in _GROUP_PROFILE_REGISTRY: + return _GROUP_PROFILE_REGISTRY[profile_key] + + group_profiles = [ + _legacy_group_profile( + group_id=0, + group_name="primary", + profile=_PROFILE_REGISTRY[profile_key], + required_for_reuse=True, + ) + ] + aux_profile = _resolve_indexer_profile(model_name) + if aux_profile is not None: + group_profiles.append( + _legacy_group_profile( + group_id=1, + group_name="aux", + profile=aux_profile, + required_for_reuse=True, + ) + ) + return tuple(group_profiles) def build_host_kv_config(model_name: str, host_kv_cache_size: int) -> Any: @@ -304,6 +501,50 @@ def build_host_kv_config(model_name: str, host_kv_cache_size: int) -> Any: return config +def build_host_kv_config_from_group_profile( + profile: HostKVGroupProfile, + host_kv_cache_size: int, + *, + shm_name: str | None = None, +) -> Any: + """Builds a HostPagedKVConfig for one logical KV group profile.""" + + if host_kv_cache_size is None: + raise ValueError("host_kv_cache_size must be a positive integer") + host_budget = int(host_kv_cache_size) + if host_budget <= 0: + raise ValueError("host_kv_cache_size must be a positive integer") + + bytes_per_page = profile.bytes_per_page() + denom = profile.num_layers * bytes_per_page + if denom <= 0: + raise ValueError(f"Invalid KV group profile '{profile.group_name}'") + if host_budget < denom: + raise ValueError( + "host_kv_cache_size is too small to allocate even one page per layer" + ) + + num_pages_per_layer = host_budget // denom + config = bg_lib.HostPagedKVConfig() + config.shm_name = shm_name or f"{HOST_KV_SHM_NAME}_{profile.group_name}" + config.num_layers = profile.num_layers + config.num_pages = num_pages_per_layer + config.page_size_tokens = profile.storage_page_tokens + config.num_k_heads = profile.num_k_heads + config.k_head_dim = profile.k_head_dim + config.num_v_heads = profile.num_v_heads + config.v_head_dim = profile.v_head_dim + config.k_element_size_bytes = _dtype_size_bytes(profile.kv_dtype) + config.v_element_size_bytes = ( + 0 if profile.num_v_heads == 0 else config.k_element_size_bytes + ) + config.sequence_table_capacity = ( + profile.sequence_table_capacity or config.num_pages + ) + config.alignment_bytes = profile.alignment_bytes + return config + + def _normalize_sequence_tokens(sequence_tokens: Sequence[int]) -> list[int]: if not sequence_tokens: raise ValueError("sequence_tokens must contain at least one element") @@ -338,6 +579,30 @@ def _compute_gpu_page_capacity( return total_pages +def _compute_gpu_page_capacity_for_group( + sequence_tokens: Sequence[int], profile: HostKVGroupProfile +) -> int: + normalized = _normalize_sequence_tokens(sequence_tokens) + total_pages = 0 + for raw_token_count in normalized: + storage_tokens = _raw_tokens_to_storage_tokens( + raw_token_count, profile.compression_ratio + ) + total_pages += (storage_tokens // profile.storage_page_tokens) + 1 + if total_pages <= 0: + raise ValueError("Computed GPU page capacity must be positive") + return total_pages + + +def _raw_tokens_to_storage_tokens( + raw_token_count: int, compression_ratio: int +) -> int: + ratio = int(compression_ratio) + if ratio <= 1: + return int(raw_token_count) + return max(1, int(raw_token_count) // ratio) + + def build_gpu_kv_config( model_name: str, sequence_tokens: Sequence[int] ) -> GPUPagedKVConfig: @@ -357,7 +622,27 @@ def build_gpu_kv_config( ) -HOST_KV_AUX_SHM_NAME = "batchgen_host_kv_cache_aux" +def build_gpu_kv_config_from_group_profile( + profile: HostKVGroupProfile, sequence_tokens: Sequence[int] +) -> GPUPagedKVConfig: + """Builds a GPUPagedKVConfig for one logical KV group profile.""" + + num_pages = _compute_gpu_page_capacity_for_group(sequence_tokens, profile) + return GPUPagedKVConfig( + num_layers=profile.num_layers, + num_pages=num_pages, + page_size_tokens=profile.storage_page_tokens, + num_k_heads=profile.num_k_heads, + k_head_dim=profile.k_head_dim, + num_v_heads=profile.num_v_heads, + v_head_dim=profile.v_head_dim, + kv_dtype=_torch_dtype_from_string(profile.kv_dtype), + ) + + +HOST_KV_AUX_SHM_NAME = os.environ.get( + "BATCHGEN_HOST_KV_AUX_SHM_NAME", "batchgen_host_kv_cache_aux" +) def is_dsa_model(model_name: str) -> bool: diff --git a/batchgen/kv_cache/prefill_offload.py b/batchgen/kv_cache/prefill_offload.py new file mode 100644 index 000000000..bfbe71266 --- /dev/null +++ b/batchgen/kv_cache/prefill_offload.py @@ -0,0 +1,194 @@ +"""Prefill Host KV offload helpers.""" + +from __future__ import annotations + +from typing import Callable, List, Optional + +import torch + +from batchgen.attention.forward_metadata import ForwardBatchMetadata +from batchgen.models.wrappers.prefix_cache import ( + ensure_prefix_cache_forward_metadata, +) + + +class PrefillHostKVOffloader: + """Offload prepacked KV with optional destination offsets.""" + + def __init__( + self, + *, + worker_view: object, + layer_idx: int, + metadata: ForwardBatchMetadata, + track_task: Optional[Callable[[object, int], None]] = None, + pin_tensor: Optional[Callable[[torch.Tensor, int], None]] = None, + ): + if worker_view is None: + raise RuntimeError("Prefill offload requires host KV view") + self.worker_view = worker_view + self.layer_idx = int(layer_idx) + self.metadata = ensure_prefix_cache_forward_metadata(metadata) + self.track_task = track_task + self.pin_tensor = pin_tensor + + def _track(self, task: object) -> None: + if task is not None and self.track_task is not None: + self.track_task(task, self.layer_idx) + + def _pin(self, tensor: torch.Tensor) -> None: + if self.pin_tensor is not None: + self.pin_tensor(tensor, self.layer_idx) + + def _pin_parent_tensors(self, *tensors: torch.Tensor) -> None: + should_sync = False + for tensor in tensors: + self._pin(tensor) + should_sync = should_sync or bool(getattr(tensor, "is_cuda", False)) + if should_sync: + event = torch.cuda.Event() + event.record(torch.cuda.current_stream()) + event.synchronize() + + def _destination_starts(self) -> Optional[List[int]]: + if not self.metadata.prefix_reuse_mode: + return None + if not hasattr( + self.worker_view, "async_offload_layer_kv_range_to_host" + ): + raise RuntimeError( + "Prefill offset offload requires " + "async_offload_layer_kv_range_to_host" + ) + return [int(tokens) for tokens in self.metadata.prefix_shared_tokens] + + def _append_lengths(self) -> List[int]: + return self.metadata.append_seq_lengths_list() + + def _offload_one( + self, + *, + sequence_id: int, + k_tensor: torch.Tensor, + v_tensor: Optional[torch.Tensor], + sequence_length: int, + destination_start: Optional[int], + ) -> None: + if destination_start is None: + task = self.worker_view.async_offload_layer_kv_to_host( + layer_idx=self.layer_idx, + sequence_ids=[int(sequence_id)], + k_tensor=k_tensor, + v_tensor=v_tensor, + sequence_lengths=[int(sequence_length)], + ) + else: + task = self.worker_view.async_offload_layer_kv_range_to_host( + layer_idx=self.layer_idx, + sequence_ids=[int(sequence_id)], + k_tensor=k_tensor, + v_tensor=v_tensor, + raw_start_positions=[int(destination_start)], + token_counts=[int(sequence_length)], + ) + self._track(task) + + def offload_gqa( + self, + *, + key: torch.Tensor, + value: torch.Tensor, + sequence_callback: Optional[ + Callable[[int, int, int, torch.Tensor, torch.Tensor], None] + ] = None, + ) -> None: + self._pin_parent_tensors(key, value) + cu = self.metadata.cu_seqlens_list() + destination_starts = self._destination_starts() + append_lengths = self._append_lengths() + for seq_idx, sequence_id in enumerate( + self.metadata.global_sequence_ids + ): + start_idx = int(cu[seq_idx]) + end_idx = int(cu[seq_idx + 1]) + query_len = end_idx - start_idx + seq_len = int(append_lengths[seq_idx]) + if seq_len < 0 or seq_len > query_len: + raise RuntimeError( + "Prefill offload append length must be within query length: " + f"sequence={sequence_id}, append={seq_len}, query={query_len}" + ) + if seq_len == 0: + continue + append_start = end_idx - seq_len + seq_key = key[append_start:end_idx].unsqueeze(0) + seq_value = value[append_start:end_idx].unsqueeze(0) + self._pin(seq_key) + self._pin(seq_value) + if sequence_callback is not None: + sequence_callback( + seq_idx, sequence_id, seq_len, seq_key, seq_value + ) + self._offload_one( + sequence_id=sequence_id, + k_tensor=seq_key, + v_tensor=seq_value, + sequence_length=seq_len, + destination_start=( + None + if destination_starts is None + else destination_starts[seq_idx] + ), + ) + + def offload_mla( + self, + *, + key: torch.Tensor, + sequence_callback: Optional[ + Callable[[int, int, int, torch.Tensor], None] + ] = None, + ) -> None: + self._pin_parent_tensors(key) + cu = self.metadata.cu_seqlens_list() + destination_starts = self._destination_starts() + append_lengths = self._append_lengths() + for seq_idx, sequence_id in enumerate( + self.metadata.global_sequence_ids + ): + start_idx = int(cu[seq_idx]) + end_idx = int(cu[seq_idx + 1]) + query_len = end_idx - start_idx + seq_len = int(append_lengths[seq_idx]) + if seq_len < 0 or seq_len > query_len: + raise RuntimeError( + "Prefill offload append length must be within query length: " + f"sequence={sequence_id}, append={seq_len}, query={query_len}" + ) + if seq_len == 0: + continue + append_start = end_idx - seq_len + seq_key = key[append_start:end_idx] + if seq_key.dim() == 2: + seq_key = seq_key.unsqueeze(0).unsqueeze(2) + elif seq_key.dim() == 3: + seq_key = seq_key.unsqueeze(0) + else: + raise RuntimeError( + "MLA prefill offload expects 2D or 3D KV, " + f"got {seq_key.dim()}D" + ) + self._pin(seq_key) + if sequence_callback is not None: + sequence_callback(seq_idx, sequence_id, seq_len, seq_key) + self._offload_one( + sequence_id=sequence_id, + k_tensor=seq_key, + v_tensor=None, + sequence_length=seq_len, + destination_start=( + None + if destination_starts is None + else destination_starts[seq_idx] + ), + ) diff --git a/batchgen/models/deepseek/deepseek_parameter_server.py b/batchgen/models/deepseek/deepseek_parameter_server.py index 09399e429..57d73911f 100644 --- a/batchgen/models/deepseek/deepseek_parameter_server.py +++ b/batchgen/models/deepseek/deepseek_parameter_server.py @@ -138,7 +138,7 @@ def Init(self): self.shm_name, self.tensor_meta_shm_name, byte_size, - self.converted_ckpt_dir, + str(self.converted_ckpt_dir), self.state_dict_name_map, ) return self.shm_name, self.tensor_meta_shm_name diff --git a/batchgen/models/deepseek/deepseekv3/modeling_deepseek_v3.py b/batchgen/models/deepseek/deepseekv3/modeling_deepseek_v3.py index 76b38dbd9..360963bf4 100755 --- a/batchgen/models/deepseek/deepseekv3/modeling_deepseek_v3.py +++ b/batchgen/models/deepseek/deepseekv3/modeling_deepseek_v3.py @@ -903,9 +903,8 @@ def moe_gate_forward_hybrid(self, hidden_states): self.e_score_correction_bias, self.n_group, self.topk_group, - self.n_routed_experts, self.top_k, - self.routed_scaling_factor + routed_scaling_factor=self.routed_scaling_factor, ) return topk_idx, topk_weight @@ -4964,4 +4963,4 @@ def forward( past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, - ) \ No newline at end of file + ) diff --git a/batchgen/models/deepseek/deepseekv3/wrappers.py b/batchgen/models/deepseek/deepseekv3/wrappers.py index 9571cd4fe..049c68eb6 100644 --- a/batchgen/models/deepseek/deepseekv3/wrappers.py +++ b/batchgen/models/deepseek/deepseekv3/wrappers.py @@ -29,6 +29,9 @@ import torch import torch.nn as nn +from batchgen.models.wrappers.prefix_mla_model_adapters import ( + build_deepseek_prefix_backend_context, +) from batchgen.models.wrappers import ExpertWrapperBase, AttnWrapperBase from batchgen.quantization.fp8e4m3 import deepseek_v3_dequantization @@ -268,18 +271,28 @@ def _forward_prefill(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: # Prepack mode: hidden_states is [1, total_tokens, hidden_dim] # Prepacked attention expects [total_tokens, hidden_dim] hidden_states_2d = hidden_states.squeeze(0) + metadata = self.prefix_cache_metadata() + position_ids = self.position_ids.to(hidden_states_2d.device) + prefix_context = None + if metadata.prefix_reuse_mode: + prefix_context = build_deepseek_prefix_backend_context( + wrapper=self, + metadata=metadata, + ) attn_output, offload_kv = self.module.prefill_attn_w8a16_prepacked( hidden_states_2d, - self.position_ids.to(hidden_states_2d.device), - self.prepack_cu_seqlens.to(hidden_states_2d.device), - self.prepack_max_seqlen, - self.prepack_num_sequences, - self.weight_dequant_scale + position_ids, + metadata.cu_seqlens.to(hidden_states_2d.device), + metadata.max_seqlen, + metadata.num_sequences, + self.weight_dequant_scale, + prefix_context=prefix_context, ) - # Offload KV cache per-sequence to host # offload_kv is [total_tokens, kv_lora_rank + qk_rope_head_dim] + if offload_kv is None: + raise RuntimeError("DeepSeek prepacked prefill returned no KV") self._offload_prepacked_kv(offload_kv) # Reshape back to [1, total_tokens, hidden_dim] for decoder_layer @@ -311,31 +324,7 @@ def _offload_prepacked_kv(self, offload_kv: torch.Tensor): Args: offload_kv: [total_tokens, kv_lora_rank + qk_rope_head_dim] """ - cu_seqlens = self.prepack_cu_seqlens - num_sequences = self.prepack_num_sequences - global_sequence_ids = self.cur_batch - - for seq_idx in range(num_sequences): - start_idx = cu_seqlens[seq_idx].item() - end_idx = cu_seqlens[seq_idx + 1].item() - seq_len = end_idx - start_idx - - # Extract KV for this sequence - seq_kv = offload_kv[start_idx:end_idx] # [seq_len, kv_dim] - - # Reshape to [1, seq_len, 1, kv_dim] for KV cache API - seq_kv = seq_kv.unsqueeze(0).unsqueeze(2) - - seq_global_id = [global_sequence_ids[seq_idx]] - - # MLA has no V (K contains compressed KV + k_pe) - self.core_engine.host_paged_kv_worker_view.async_offload_layer_kv_to_host( - layer_idx=self.layer_idx, - sequence_ids=seq_global_id, - k_tensor=seq_kv, - v_tensor=None, - sequence_lengths=[seq_len], - ) + self.offload_prepacked_mla_kv(offload_kv) def _forward_decode(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: """Decode forward using FlashMLA backend. diff --git a/batchgen/models/glm/glm5/wrappers.py b/batchgen/models/glm/glm5/wrappers.py index a9efafc51..b5852cf8f 100644 --- a/batchgen/models/glm/glm5/wrappers.py +++ b/batchgen/models/glm/glm5/wrappers.py @@ -26,6 +26,10 @@ import torch.nn.functional as F +from batchgen.models.wrappers.prefix_mla_model_adapters import ( + build_glm5_prefix_backend_context, + offload_glm5_prepacked_mla_kv, +) from batchgen.models.wrappers import ExpertWrapperBase, AttnWrapperBase from batchgen.timing import init_decode_timer @@ -682,15 +686,23 @@ def _forward_prefill(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: ) if self.prepack_mode: hidden_states_2d = hidden_states.squeeze(0) + metadata = self.prefix_cache_metadata() + position_ids = self.position_ids.to(hidden_states_2d.device) + prefix_context = None + if metadata.prefix_reuse_mode: + prefix_context = build_glm5_prefix_backend_context( + wrapper=self, + metadata=metadata, + ) attn_output, offload_kv = self.module.prefill_attn_w8a16_prepacked( hidden_states_2d, - self.position_ids.to(hidden_states_2d.device), - self.prepack_cu_seqlens.to(hidden_states_2d.device), - self.prepack_max_seqlen, - self.prepack_num_sequences, - self.weight_dequant_scale + position_ids, + metadata.cu_seqlens.to(hidden_states_2d.device), + metadata.max_seqlen, + metadata.num_sequences, + self.weight_dequant_scale, + prefix_context=prefix_context, ) - # DSA: compute indexer K and offload to auxiliary host cache. # This path MUST run for every prompt token during prefill — otherwise # aux cache is unpopulated and any later decode past 2048 tokens reads @@ -702,7 +714,7 @@ def _forward_prefill(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: ) indexer_kv = self.module.indexer.compute_indexer_kv( hidden_states_2d.unsqueeze(0), - positions=self.position_ids.to(hidden_states_2d.device), + positions=position_ids, ) if indexer_kv is None: raise RuntimeError( @@ -710,6 +722,8 @@ def _forward_prefill(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: ) self._offload_prepacked_indexer_kv(indexer_kv.squeeze(0)) + if offload_kv is None: + raise RuntimeError("GLM-5 prepacked prefill returned no KV") self._offload_prepacked_kv(offload_kv) attn_output = attn_output.unsqueeze(0) return (attn_output, None, None) @@ -727,75 +741,21 @@ def _offload_prepacked_indexer_kv(self, offload_kv: torch.Tensor): "GLM-5 DSA auxiliary host KV worker view is required for " "indexer KV offload" ) - cu_seqlens = self.prepack_cu_seqlens - num_sequences = self.prepack_num_sequences - global_sequence_ids = self.cur_batch - - # Lifespan management mirrored from decode-side `_pending_kv_append_*` - # (worker.py:1898-1925). Drain compute stream via a CUDA event so the - # FA3 prefill kernel that wrote `offload_kv` has fully retired before - # the C++ async lambda's d2h memcpy reads the source memory; pin the - # source tensor (and the parent `offload_kv`) in the class-level list - # so PyTorch's caching allocator cannot re-hand the same physical - # pages to a later layer's K/V tensor while the d2h is in flight. - AttnWrapperBase.pin_prefill_offload_tensor(offload_kv, self.layer_idx) - evt = torch.cuda.Event() - evt.record(torch.cuda.current_stream()) - evt.synchronize() - - # Single D2H sync for all seq boundaries instead of 2N per-seq .item() calls. - cu = cu_seqlens.tolist() - for seq_idx in range(num_sequences): - start_idx = cu[seq_idx] - end_idx = cu[seq_idx + 1] - seq_len = end_idx - start_idx - # indexer_kv is already [T, H=1, D=128] after caller's .squeeze(0), - # so only .unsqueeze(0) is needed to add the B dim; don't also - # .unsqueeze(2) (that would make 5D — the primary-MLA path copy-paste - # of this code was for a 2D [T, kv_lora+rope] input). - seq_kv = offload_kv[start_idx:end_idx].unsqueeze(0) - seq_global_id = [global_sequence_ids[seq_idx]] - task = AttnWrapperBase.host_paged_kv_worker_view_aux.async_offload_layer_kv_to_host( - layer_idx=self.layer_idx, - sequence_ids=seq_global_id, - k_tensor=seq_kv, - v_tensor=None, - sequence_lengths=[seq_len], - ) - # Pin both the per-seq view AND the parent offload_kv (already - # pinned outside the loop) so neither's storage is reclaimed. - AttnWrapperBase.pin_prefill_offload_tensor(seq_kv, self.layer_idx) - AttnWrapperBase.track_prefill_offload_task(task, self.layer_idx) + offload_glm5_prepacked_mla_kv( + key=offload_kv, + worker_view=AttnWrapperBase.host_paged_kv_worker_view_aux, + layer_idx=self.layer_idx, + metadata=self.prefix_cache_metadata(), + ) def _offload_prepacked_kv(self, offload_kv: torch.Tensor): """Offload KV cache per-sequence to host memory.""" - cu_seqlens = self.prepack_cu_seqlens - num_sequences = self.prepack_num_sequences - global_sequence_ids = self.cur_batch - - # See _offload_prepacked_indexer_kv for rationale. - AttnWrapperBase.pin_prefill_offload_tensor(offload_kv, self.layer_idx) - evt = torch.cuda.Event() - evt.record(torch.cuda.current_stream()) - evt.synchronize() - - # Single D2H sync for all seq boundaries instead of 2N per-seq .item() calls. - cu = cu_seqlens.tolist() - for seq_idx in range(num_sequences): - start_idx = cu[seq_idx] - end_idx = cu[seq_idx + 1] - seq_len = end_idx - start_idx - seq_kv = offload_kv[start_idx:end_idx].unsqueeze(0).unsqueeze(2) - seq_global_id = [global_sequence_ids[seq_idx]] - task = self.core_engine.host_paged_kv_worker_view.async_offload_layer_kv_to_host( - layer_idx=self.layer_idx, - sequence_ids=seq_global_id, - k_tensor=seq_kv, - v_tensor=None, - sequence_lengths=[seq_len], - ) - AttnWrapperBase.pin_prefill_offload_tensor(seq_kv, self.layer_idx) - AttnWrapperBase.track_prefill_offload_task(task, self.layer_idx) + offload_glm5_prepacked_mla_kv( + key=offload_kv, + worker_view=self.core_engine.host_paged_kv_worker_view, + layer_idx=self.layer_idx, + metadata=self.prefix_cache_metadata(), + ) def _forward_decode(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: """Decode forward with DSA sparse attention. diff --git a/batchgen/models/minimax/minimax_m25/model.py b/batchgen/models/minimax/minimax_m25/model.py index 2a702a8d0..329f101cc 100644 --- a/batchgen/models/minimax/minimax_m25/model.py +++ b/batchgen/models/minimax/minimax_m25/model.py @@ -663,15 +663,45 @@ def __init__( ) def resize_if_needed(self, global_bsz: int): - """Resize communication/routing buffers if global_bsz exceeds capacity.""" - if global_bsz <= self.max_global_bsz: + """Resize buffers for the current global decode batch. + + ``dispatch_scatter_3d`` writes up to one row per original token into a + single local expert slot. In the worst case all tokens route to the same + local expert, so the per-expert stride must cover ``global_bsz``. + """ + grew_comm = global_bsz > self.max_global_bsz + grew_mtp = global_bsz > self.max_tokens_padded + if not grew_comm and not grew_mtp: return - logging.info(f"[MoEBufferManager] Resizing: {self.max_global_bsz} -> {global_bsz}") - self.max_global_bsz = global_bsz - NK = global_bsz * self.topk - self.all_tokens = torch.zeros(global_bsz, self.H, dtype=torch.bfloat16, device=self.device) - self.topk_pos = torch.full((NK,), -1, dtype=torch.int32, device=self.device) - self.result_buffer = torch.empty(global_bsz, self.H, dtype=torch.bfloat16, device=self.device) + + if grew_comm: + logging.info( + f"[MoEBufferManager] Resizing comm buffers: {self.max_global_bsz} -> {global_bsz}" + ) + self.max_global_bsz = global_bsz + NK = global_bsz * self.topk + self.all_tokens = torch.zeros( + global_bsz, self.H, dtype=torch.bfloat16, device=self.device, + ) + self.topk_pos = torch.full((NK,), -1, dtype=torch.int32, device=self.device) + self.result_buffer = torch.empty( + global_bsz, self.H, dtype=torch.bfloat16, device=self.device, + ) + + if grew_mtp: + new_mtp = ((global_bsz + _DEFAULT_MTP - 1) // _DEFAULT_MTP) * _DEFAULT_MTP + logging.info( + f"[MoEBufferManager] Resizing 3D buffers: " + f"mtp {self.max_tokens_padded} -> {new_mtp}" + ) + self.max_tokens_padded = new_mtp + buf_rows = self.E_local * new_mtp + self.dispatched_x = torch.zeros( + buf_rows, self.H, dtype=torch.bfloat16, device=self.device, + ) + self.expert_out = torch.zeros( + buf_rows, self.H, dtype=torch.bfloat16, device=self.device, + ) def _total_bytes(self): total = 0 diff --git a/batchgen/models/minimax/minimax_m25/wrappers.py b/batchgen/models/minimax/minimax_m25/wrappers.py index 568036c29..aa1f904c5 100644 --- a/batchgen/models/minimax/minimax_m25/wrappers.py +++ b/batchgen/models/minimax/minimax_m25/wrappers.py @@ -33,6 +33,10 @@ import torch.nn.functional as F from batchgen.models.wrappers import ExpertWrapperBase, AttnWrapperBase +from batchgen.models.wrappers.prefix_gqa_extend import ( + GqaExtendSpec, + run_prefix_gqa_prefill_attention, +) from batchgen.quantization.fp8e4m3 import deepseek_v3_dequantization from .model import rotate_half @@ -422,7 +426,6 @@ def dequantize_weights(self, weights_dict): def _forward_prefill(self, hidden_states, **kwargs): """Prefill forward: FP8 Q/K/V projection + QK norm + partial RoPE + FA varlen.""" - from batchgen.attention.gqa import gqa_prefill_fa from batchgen.attention.fused_kernels import cuda_rmsnorm fp8_q, q_scale, fp8_k, k_scale, fp8_v, v_scale, fp8_o, o_scale = self._get_attn_weights() @@ -435,9 +438,14 @@ def _forward_prefill(self, hidden_states, **kwargs): hidden_states_2d = hidden_states total_tokens = hidden_states_2d.shape[0] - cu_seqlens = self.prepack_cu_seqlens.to(hidden_states_2d.device) - max_seqlen = self.prepack_max_seqlen + metadata = self.prefix_cache_metadata() + max_seqlen = metadata.max_seqlen position_ids = self.position_ids.to(hidden_states_2d.device) + full_seq_lengths = metadata.full_seq_lengths + if metadata.prefix_reuse_mode and full_seq_lengths: + rotary_seq_len = max(max(int(length) for length in full_seq_lengths), int(max_seqlen)) + else: + rotary_seq_len = int(max_seqlen) # Q/K/V projection — packed FP8 GEMM (2.94× faster) or fallback if _HAS_PACKED_QKV and hasattr(self, 'packed_qkv_w') and self.packed_qkv_w is not None: @@ -464,7 +472,7 @@ def _forward_prefill(self, hidden_states, **kwargs): value = value.view(total_tokens, num_kv_heads, head_dim) # Partial RoPE (rotate first rotary_dim=64 dims, passthrough rest) - cos, sin = self.module.rotary_emb(value, seq_len=max_seqlen) + cos, sin = self.module.rotary_emb(value, seq_len=rotary_seq_len) cos = cos[position_ids] # [total_tokens, rotary_dim] sin = sin[position_ids] @@ -486,15 +494,16 @@ def _forward_prefill(self, hidden_states, **kwargs): k_pass, ], dim=-1) - # FlashAttention varlen GQA - attn_output, lse = gqa_prefill_fa( - q=query, - k=key, - v=value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, + attn_output = run_prefix_gqa_prefill_attention( + wrapper=self, + query=query, + key=key, + value=value, + metadata=metadata, + spec=GqaExtendSpec( + num_kv_heads=num_kv_heads, + head_dim=head_dim, + ), ) # Output projection via FP8 GEMM @@ -516,26 +525,7 @@ def _forward_prefill(self, hidden_states, **kwargs): def _offload_prepacked_kv_gqa(self, k_cache, v_cache): """Offload GQA KV cache per-sequence to host memory.""" - cu_seqlens = self.prepack_cu_seqlens - num_sequences = self.prepack_num_sequences - global_sequence_ids = self.cur_batch - - for seq_idx in range(num_sequences): - start_idx = cu_seqlens[seq_idx].item() - end_idx = cu_seqlens[seq_idx + 1].item() - seq_len = end_idx - start_idx - - seq_k = k_cache[start_idx:end_idx].unsqueeze(0) - seq_v = v_cache[start_idx:end_idx].unsqueeze(0) - seq_global_id = [global_sequence_ids[seq_idx]] - - self.core_engine.host_paged_kv_worker_view.async_offload_layer_kv_to_host( - layer_idx=self.layer_idx, - sequence_ids=seq_global_id, - k_tensor=seq_k, - v_tensor=seq_v, - sequence_lengths=[seq_len], - ) + self.offload_prepacked_gqa_kv(k_cache, v_cache) def _forward_decode(self, hidden_states, **kwargs): """Decode forward: FP8 Q/K/V + QK norm + partial RoPE + paged KV attention. diff --git a/batchgen/models/moonshotai/kimi_k25/Parallel_Strategy_Manager.py b/batchgen/models/moonshotai/kimi_k25/Parallel_Strategy_Manager.py index d807fd455..63da12b12 100644 --- a/batchgen/models/moonshotai/kimi_k25/Parallel_Strategy_Manager.py +++ b/batchgen/models/moonshotai/kimi_k25/Parallel_Strategy_Manager.py @@ -28,7 +28,12 @@ - Loads INT4 packed/scale tensors for persistent experts """ -from .model import KimiK25ForCausalLM, KimiK25MoE, KimiK25MoEBufferManager +from .model import ( + KimiK25ForCausalLM, + KimiK25MoE, + KimiK25MoEBufferManager, + round_moe_buffer_tokens, +) from .wrappers import KimiK25ExpertWrapper, KimiK25AttnWrapper import logging import types @@ -524,6 +529,7 @@ def _log_hbm(step_name): # Allocate shared MoE buffer manager (one instance for all 60 MoE layers) max_global_bsz = self.world_size * effective_padding_bsz + max_tokens_padded = round_moe_buffer_tokens(max_global_bsz) KimiK25MoE._buf = KimiK25MoEBufferManager( E_local=NUM_LOCAL_EXPERT_PER_LAYER, max_global_bsz=max_global_bsz, @@ -532,6 +538,7 @@ def _log_hbm(step_name): topk=self.loaded_model_config.num_experts_per_tok, num_tokens_per_rank=effective_padding_bsz, device=device, + max_tokens_padded=max_tokens_padded, ) _log_hbm("MoEBufferManager") diff --git a/batchgen/models/moonshotai/kimi_k25/assets/modeling_deepseek.py b/batchgen/models/moonshotai/kimi_k25/assets/modeling_deepseek.py index fd3f1b443..fdf1526d2 100644 --- a/batchgen/models/moonshotai/kimi_k25/assets/modeling_deepseek.py +++ b/batchgen/models/moonshotai/kimi_k25/assets/modeling_deepseek.py @@ -889,9 +889,8 @@ def moe_gate_forward_hybrid(self, hidden_states): self.e_score_correction_bias, self.n_group, self.topk_group, - self.n_routed_experts, self.top_k, - self.routed_scaling_factor + routed_scaling_factor=self.routed_scaling_factor, ) return topk_idx, topk_weight @@ -4785,4 +4784,4 @@ def forward( # K2.5-specific aliases for external code -KimiK25ForCausalLM = DeepseekV3ForCausalLM \ No newline at end of file +KimiK25ForCausalLM = DeepseekV3ForCausalLM diff --git a/batchgen/models/moonshotai/kimi_k25/model.py b/batchgen/models/moonshotai/kimi_k25/model.py index 0609fdf57..6740aa276 100644 --- a/batchgen/models/moonshotai/kimi_k25/model.py +++ b/batchgen/models/moonshotai/kimi_k25/model.py @@ -233,6 +233,13 @@ def _get_k25_timer(num_layers: int = 61) -> Optional[K25DecodeTimer]: _DEFAULT_MTP = 4096 # Default max_tokens_padded (stride per expert in 3D buffer) +def round_moe_buffer_tokens(num_tokens: int) -> int: + """Round MoE 3D-stride capacity to the WGMMA/TMA tile requirement.""" + if num_tokens <= 0: + return _BLOCK_M + return max(_BLOCK_M, ((num_tokens + _BLOCK_M - 1) // _BLOCK_M) * _BLOCK_M) + + class KimiK25MoEBufferManager: """Pre-allocated buffers for K2.5 MoE decode pipeline (3D strided layout). @@ -272,12 +279,11 @@ def __init__( # -> no runtime resize -> no corner-case OOM. (Previously a fixed _DEFAULT_MTP=4096, # which over-reserved ~6 GiB and OOM'd single-node no-offload init.) if max_tokens_padded is None: - max_tokens_padded = max( - ((max_global_bsz + _BLOCK_M - 1) // _BLOCK_M) * _BLOCK_M, _BLOCK_M) - self.max_tokens_padded = max_tokens_padded + max_tokens_padded = max_global_bsz + self.max_tokens_padded = round_moe_buffer_tokens(max_tokens_padded) NK = max_global_bsz * topk - buf_rows = E_local * max_tokens_padded # 3D strided: E * mtp + buf_rows = E_local * self.max_tokens_padded # 3D strided: E * mtp # Communication buffers self.all_tokens = torch.zeros(max_global_bsz, H, dtype=torch.bfloat16, device=device) @@ -305,7 +311,7 @@ def __init__( self._init_tma_descriptors() logging.debug( - f"[MoEBufferManager] 3D strided layout: E_local={E_local}, mtp={max_tokens_padded}, " + f"[MoEBufferManager] 3D strided layout: E_local={E_local}, mtp={self.max_tokens_padded}, " f"buf_rows={buf_rows}, H={H}, N_inter={N_inter}, " f"total={self._total_bytes() / (1024**3):.2f} GiB" ) @@ -336,7 +342,7 @@ def resize_if_needed(self, global_bsz: int): # Resize 3D buffers only if needed if global_bsz > self.max_tokens_padded: - new_mtp = ((global_bsz + _BLOCK_M - 1) // _BLOCK_M) * _BLOCK_M + new_mtp = round_moe_buffer_tokens(global_bsz) logging.info(f"[MoEBufferManager] Resizing 3D buffers: mtp {self.max_tokens_padded} → {new_mtp}") self.max_tokens_padded = new_mtp buf_rows = self.E_local * new_mtp diff --git a/batchgen/models/moonshotai/kimi_k25/wrappers.py b/batchgen/models/moonshotai/kimi_k25/wrappers.py index a499c2af6..ed3381dda 100644 --- a/batchgen/models/moonshotai/kimi_k25/wrappers.py +++ b/batchgen/models/moonshotai/kimi_k25/wrappers.py @@ -37,6 +37,9 @@ import torch.nn as nn import torch.nn.functional as F +from batchgen.models.wrappers.prefix_mla_model_adapters import ( + build_kimi_prefix_backend_context, +) from batchgen.models.wrappers import ExpertWrapperBase, AttnWrapperBase @@ -434,16 +437,26 @@ def _forward_prefill(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: if self.prepack_mode: # Prepacked mode: varlen flash attention hidden_states_2d = hidden_states.squeeze(0) + metadata = self.prefix_cache_metadata() + position_ids = self.position_ids.to(hidden_states_2d.device) + prefix_context = None + if metadata.prefix_reuse_mode: + prefix_context = build_kimi_prefix_backend_context( + wrapper=self, + metadata=metadata, + ) attn_output, offload_kv = self.module.prefill_attn_bf16_prepacked( hidden_states_2d, - self.position_ids.to(hidden_states_2d.device), - self.prepack_cu_seqlens.to(hidden_states_2d.device), - self.prepack_max_seqlen, - self.prepack_num_sequences, + position_ids, + metadata.cu_seqlens.to(hidden_states_2d.device), + metadata.max_seqlen, + metadata.num_sequences, + prefix_context=prefix_context, ) - # Offload KV cache per-sequence to host + if offload_kv is None: + raise RuntimeError("Kimi prepacked prefill returned no KV") self._offload_prepacked_kv(offload_kv) attn_output = attn_output.unsqueeze(0) @@ -468,26 +481,7 @@ def _offload_prepacked_kv(self, offload_kv: torch.Tensor): Args: offload_kv: [total_tokens, kv_lora_rank + qk_rope_head_dim] """ - cu_seqlens = self.prepack_cu_seqlens - num_sequences = self.prepack_num_sequences - global_sequence_ids = self.cur_batch - - for seq_idx in range(num_sequences): - start_idx = cu_seqlens[seq_idx].item() - end_idx = cu_seqlens[seq_idx + 1].item() - seq_len = end_idx - start_idx - - seq_kv = offload_kv[start_idx:end_idx].unsqueeze(0).unsqueeze(2) - seq_global_id = [global_sequence_ids[seq_idx]] - - # MLA: K contains compressed KV + k_pe, no separate V - self.core_engine.host_paged_kv_worker_view.async_offload_layer_kv_to_host( - layer_idx=self.layer_idx, - sequence_ids=seq_global_id, - k_tensor=seq_kv, - v_tensor=None, - sequence_lengths=[seq_len], - ) + self.offload_prepacked_mla_kv(offload_kv) def _forward_decode(self, hidden_states: torch.Tensor, **kwargs) -> Tuple: """Decode forward using BF16 MLA attention. diff --git a/batchgen/models/openai/gpt_oss_120b/model.py b/batchgen/models/openai/gpt_oss_120b/model.py index 10e4f9da3..ea931cc10 100644 --- a/batchgen/models/openai/gpt_oss_120b/model.py +++ b/batchgen/models/openai/gpt_oss_120b/model.py @@ -844,10 +844,11 @@ def _forward_local(self, hidden_flat: torch.Tensor) -> torch.Tensor: # Phase 1: Grouped kernel for persistent experts num_persistent = len(self.persistent_expert_indices) if num_persistent > 0: - output = self._grouped_forward( + self._grouped_forward( hidden_flat, topk_indices, topk_weights, expert_start=self.expert_start, num_local_experts=num_persistent, + output=output, ) # Phase 2: Single-expert kernel for non-persistent experts @@ -902,21 +903,23 @@ def _forward_ep(self, x: torch.Tensor) -> torch.Tensor: global_results = self.global_results_buffer global_results.zero_() num_global_tokens = all_tokens.shape[0] + global_results_view = global_results[:num_global_tokens] # Phase 1: Grouped kernel for persistent experts num_persistent = len(self.persistent_expert_indices) if num_persistent > 0: - global_results[:num_global_tokens] = self._grouped_forward( + self._grouped_forward( all_tokens, topk_indices, topk_weights, expert_start=self.expert_start, num_local_experts=num_persistent, + output=global_results_view, ) # Phase 2: Single-expert kernel for non-persistent experts if self.non_persistent_expert_indices: self._single_expert_forward( all_tokens, topk_indices, topk_weights, - global_results[:num_global_tokens], + global_results_view, ) # 4) AllReduce @@ -938,6 +941,7 @@ def _grouped_forward( topk_weights: torch.Tensor, expert_start: int, num_local_experts: int, + output: torch.Tensor | None = None, ) -> torch.Tensor: """Grouped kernel for persistent experts.""" if self.weight_format == "mxfp4": @@ -959,6 +963,7 @@ def _grouped_forward( gate_bias_ptrs=self.gate_bias_ptrs, up_bias_ptrs=self.up_bias_ptrs, down_bias_ptrs=self.down_bias_ptrs, + output=output, ) elif self.weight_format == "bf16": # Placeholder: grouped BF16 kernel to be ported from @@ -1110,6 +1115,23 @@ def _get_bf16_buffer(self, shape: Tuple[int, int], device: torch.device) -> torc self._buffer_shape = shape return self._bf16_buffer + @staticmethod + def _debug_sync_mlp(label: str, tensor: torch.Tensor) -> None: + if os.environ.get("BATCHGEN_GPT_OSS_MLP_DEBUG_SYNC", "0") != "1": + return + logging.getLogger(__name__).info( + "[GPT_OSS_MLP_DEBUG] before_sync label=%s shape=%s dtype=%s device=%s", + label, + tuple(tensor.shape), + tensor.dtype, + tensor.device, + ) + torch.cuda.synchronize(tensor.device) + logging.getLogger(__name__).info( + "[GPT_OSS_MLP_DEBUG] after_sync label=%s", + label, + ) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """Forward pass: grouped WGMMA for persistent experts, per-expert loop for the rest.""" import os @@ -1121,8 +1143,17 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, seq_len, hidden_dim = hidden_states.shape hidden_flat = hidden_states.view(-1, hidden_dim) # [total_tokens, hidden_size] + disable_fused_gate = ( + os.environ.get("BATCHGEN_GPT_OSS_PREFILL_DISABLE_FUSED_GATE", "0") + == "1" + ) + # Compute routing: fused gate (WGMMA GEMM + bias + TopK + Softmax) or fallback - if self._fused_gate_ctx is None and _HAS_CUDA_ROUTING: + if ( + not disable_fused_gate + and self._fused_gate_ctx is None + and _HAS_CUDA_ROUTING + ): w = self.router.weight # [E, K_dim] BF16 if w.dtype == torch.bfloat16: from batchgen.moe.routing import FusedGateContext @@ -1130,22 +1161,26 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: _bias_bf16 = _bias.to(torch.bfloat16) if _bias is not None else None self._fused_gate_ctx = FusedGateContext(w, _bias_bf16, topk=self.num_experts_per_tok) - if self._fused_gate_ctx is not None: + if self._fused_gate_ctx is not None and not disable_fused_gate: topk_indices, topk_weights = self._fused_gate_ctx.forward(hidden_flat) + self._debug_sync_mlp("routing_fused_gate", hidden_flat) elif _HAS_CUDA_ROUTING: router_logits = self.router(hidden_flat) # [total_tokens, num_experts] topk_indices, topk_weights = gate_topk_softmax_cuda( router_logits, k=self.num_experts_per_tok ) + self._debug_sync_mlp("routing_cuda_topk", hidden_flat) else: router_logits = self.router(hidden_flat) topk_weights, topk_indices = torch.topk( router_logits, k=self.num_experts_per_tok, dim=-1 ) topk_weights = F.softmax(topk_weights, dim=-1) + self._debug_sync_mlp("routing_torch_topk", hidden_flat) # Initialize output output = torch.zeros_like(hidden_flat) + self._debug_sync_mlp("zeros_like_output", output) # Phase 1: Grouped WGMMA for persistent experts num_persistent = len(self.persistent_expert_indices) @@ -1163,6 +1198,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: up_bias_ptrs=self.up_bias_ptrs, down_bias_ptrs=self.down_bias_ptrs, ) + self._debug_sync_mlp("grouped_moe", output) # If all experts are persistent, we're done if not self.non_persistent_expert_indices: return output.view(batch_size, seq_len, hidden_dim) diff --git a/batchgen/models/openai/gpt_oss_120b/wrappers.py b/batchgen/models/openai/gpt_oss_120b/wrappers.py index 9803b4aeb..5a25b3024 100644 --- a/batchgen/models/openai/gpt_oss_120b/wrappers.py +++ b/batchgen/models/openai/gpt_oss_120b/wrappers.py @@ -33,7 +33,7 @@ import math import os import time -from typing import Dict, Optional, Tuple +from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn @@ -1677,7 +1677,9 @@ def _forward_prefill_prepacked( Tuple of (output, None, None) - KV cache offloaded to host """ # Import here to avoid circular imports - from batchgen.attention.gqa import gqa_prefill_fa + from batchgen.attention.prefix_aware_backend import ( + GqaPrefixAwareAttentionBackend, + ) # Handle both 2D and 3D input if hidden_states.dim() == 3: @@ -1691,10 +1693,18 @@ def _forward_prefill_prepacked( total_tokens = hidden_states_2d.shape[0] # Get prepack metadata from class variables - cu_seqlens = AttnWrapperBase.prepack_cu_seqlens - max_seqlen = AttnWrapperBase.prepack_max_seqlen - num_sequences = AttnWrapperBase.prepack_num_sequences - seq_lengths = AttnWrapperBase.prepack_seq_lengths + metadata = self.prefix_cache_metadata() + cu_seqlens = metadata.cu_seqlens + max_seqlen = metadata.max_seqlen + num_sequences = metadata.num_sequences + seq_lengths = metadata.seq_lengths + prefix_reuse_mode = metadata.prefix_reuse_mode + global_sequence_ids = metadata.global_sequence_ids + full_seq_lengths = metadata.full_seq_lengths + if prefix_reuse_mode and full_seq_lengths: + rotary_seq_len = max(max(int(length) for length in full_seq_lengths), int(max_seqlen)) + else: + rotary_seq_len = int(max_seqlen) # DEBUG: Check input hidden_states before projection if self.layer_idx == 0 and os.environ.get("BATCHGEN_DEBUG_PREFILL_KV", "0") == "1": @@ -1766,8 +1776,8 @@ def _forward_prefill_prepacked( # hidden_states_2d: [total_tokens, hidden_size] if self._use_wgmma and position_ids is not None: from batchgen.attention.fused_kernels import cuda_qkv_wgmma - cos_table = self.module.rotary_emb.cos_cached[:max_seqlen].to(hidden_states_2d.dtype) - sin_table = self.module.rotary_emb.sin_cached[:max_seqlen].to(hidden_states_2d.dtype) + cos_table = self.module.rotary_emb.cos_cached[:rotary_seq_len].to(hidden_states_2d.dtype) + sin_table = self.module.rotary_emb.sin_cached[:rotary_seq_len].to(hidden_states_2d.dtype) rope_cos = cos_table[position_ids] # [total_tokens, head_dim] rope_sin = sin_table[position_ids] query, key, value = cuda_qkv_wgmma( @@ -1787,7 +1797,7 @@ def _forward_prefill_prepacked( # Apply RoPE per sequence using position_ids if position_ids is not None: - cos, sin = self.module.rotary_emb(value, seq_len=max_seqlen) + cos, sin = self.module.rotary_emb(value, seq_len=rotary_seq_len) cos = cos[position_ids] # [total_tokens, head_dim] sin = sin[position_ids] # [total_tokens, head_dim] @@ -1808,20 +1818,41 @@ def _forward_prefill_prepacked( k2 * cos_half + k1 * sin_half ], dim=-1) - # Use gqa_prefill_fa for varlen attention with sink correction - # q, k, v: [total_tokens, num_heads, head_dim] - attn_output, lse = gqa_prefill_fa( - q=query, - k=key, - v=value, - cu_seqlens_q=cu_seqlens.to(hidden_states_2d.device), - cu_seqlens_k=cu_seqlens.to(hidden_states_2d.device), - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, + backend = GqaPrefixAwareAttentionBackend( + layer_idx=self.layer_idx, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, sinks=self.sinks, softmax_scale=self.scale, sliding_window=self.sliding_window, ) + from batchgen.attention.forward_metadata_context import ( + get_current_forward_batch_metadata, + ) + + forward_metadata = get_current_forward_batch_metadata() + kv_cache_metadata = ( + None if forward_metadata is None else forward_metadata.kv_cache + ) + if ( + kv_cache_metadata is None + and AttnWrapperBase.prefill_prefix_materialization is not None + ): + from types import SimpleNamespace + + kv_cache_metadata = SimpleNamespace( + prefill_prefix_materialization=( + AttnWrapperBase.prefill_prefix_materialization + ) + ) + # q, k, v: [total_tokens, num_heads, head_dim] + attn_output = backend.forward_prefill( + query=query, + key=key, + value=value, + metadata=metadata, + kv_cache_metadata=kv_cache_metadata, + ) # attn_output: [total_tokens, num_heads, head_dim] # Reshape for output projection @@ -1830,9 +1861,6 @@ def _forward_prefill_prepacked( # Output projection attn_output = self.module.o_proj(attn_output) # [total_tokens, hidden_size] - # Offload KV cache per sequence to host - global_sequence_ids = AttnWrapperBase.cur_batch - torch.cuda.current_stream().synchronize() # Make sure KV is ready # DEBUG: Check if K values differ across sequences before offload @@ -1865,35 +1893,28 @@ def _forward_prefill_prepacked( else: print(f"[PREFILL L0] OK: seq0 and seq1 have DIFFERENT K at position 0") - # For GQA, we store both K and V (unlike MLA which only stores K) - # Split by cu_seqlens and offload each sequence - for seq_idx in range(num_sequences): - start_idx = cu_seqlens[seq_idx].item() - end_idx = cu_seqlens[seq_idx + 1].item() - seq_len = end_idx - start_idx - - # Extract KV for this sequence - seq_key = key[start_idx:end_idx] # [seq_len, num_kv_heads, head_dim] - seq_value = value[start_idx:end_idx] # [seq_len, num_kv_heads, head_dim] - - # Reshape to [1, seq_len, num_kv_heads, head_dim] for KV cache API - seq_key = seq_key.unsqueeze(0) - seq_value = seq_value.unsqueeze(0) - - seq_global_id = [global_sequence_ids[seq_idx]] - - # DEBUG: Print what's being offloaded per sequence - if self.layer_idx == 0 and os.environ.get("BATCHGEN_DEBUG_PREFILL_KV", "0") == "1" and seq_idx < 3: - k_sample = seq_key[0, 0, 0, :4].cpu().tolist() # [1, seq_len, heads, dim] -> position 0, head 0 - print(f"[PREFILL L0 OFFLOAD] seq{seq_idx}: global_id={seq_global_id[0]}, seq_len={seq_len}, K[0,0,:4]={k_sample}") - - self.core_engine.host_paged_kv_worker_view.async_offload_layer_kv_to_host( - layer_idx=self.layer_idx, - sequence_ids=seq_global_id, - k_tensor=seq_key, - v_tensor=seq_value, - sequence_lengths=[seq_len], - ) + def _debug_offload_sequence(seq_idx, sequence_id, seq_len, seq_key, seq_value): + del seq_value + if ( + self.layer_idx == 0 + and os.environ.get("BATCHGEN_DEBUG_PREFILL_KV", "0") == "1" + and seq_idx < 3 + ): + # [1, seq_len, heads, dim] -> position 0, head 0. + k_sample = seq_key[0, 0, 0, :4].cpu().tolist() + print( + f"[PREFILL L0 OFFLOAD] seq{seq_idx}: " + f"global_id={sequence_id}, seq_len={seq_len}, " + f"K[0,0,:4]={k_sample}" + ) + + self.offload_prepacked_gqa_kv( + key, + value, + metadata=metadata, + track_tasks=(metadata.full_seq_lengths is not None), + sequence_callback=_debug_offload_sequence, + ) logging.debug( f"[Layer {self.layer_idx}] GPT-OSS prepacked prefill complete. " diff --git a/batchgen/models/wrappers/__init__.py b/batchgen/models/wrappers/__init__.py index 9a8ad470c..451cef31a 100644 --- a/batchgen/models/wrappers/__init__.py +++ b/batchgen/models/wrappers/__init__.py @@ -37,9 +37,9 @@ ) """ +from .attention import AttnWrapperBase from .base import BaseModuleWrapper from .expert import ExpertWrapperBase -from .attention import AttnWrapperBase __all__ = [ "BaseModuleWrapper", diff --git a/batchgen/models/wrappers/attention.py b/batchgen/models/wrappers/attention.py index fe956b824..0bbecaa78 100644 --- a/batchgen/models/wrappers/attention.py +++ b/batchgen/models/wrappers/attention.py @@ -93,6 +93,17 @@ class AttnWrapperBase(BaseModuleWrapper): pending_prefill_offload_tensors: ClassVar[list] = [] pending_prefill_offload_layer_idx: ClassVar[Optional[int]] = None + @classmethod + def _finish_pending_prefix_materialization_layer( + cls, + layer_idx: Optional[int], + ) -> None: + del cls + materialization = AttnWrapperBase.prefill_prefix_materialization + if layer_idx is None or materialization is None: + return + materialization.finish_layer(int(layer_idx)) + @classmethod def record_glm5_dispatch( cls, @@ -155,21 +166,23 @@ def retire_pending_prefill_offloads( device: Optional[torch.device] = None, reason: str = "", ) -> int: - pending = cls.pending_prefill_offload_tasks - pinned = cls.pending_prefill_offload_tensors + del cls + pending = AttnWrapperBase.pending_prefill_offload_tasks + pinned = AttnWrapperBase.pending_prefill_offload_tensors if not pending and not pinned: - cls.pending_prefill_offload_layer_idx = None + AttnWrapperBase.pending_prefill_offload_layer_idx = None return 0 num_tasks = len(pending) for task in pending: task.wait() pending.clear() - cls._prefill_offload_sync_device(device) + AttnWrapperBase._prefill_offload_sync_device(device) pinned.clear() - layer_idx = cls.pending_prefill_offload_layer_idx - cls.pending_prefill_offload_layer_idx = None + layer_idx = AttnWrapperBase.pending_prefill_offload_layer_idx + AttnWrapperBase._finish_pending_prefix_materialization_layer(layer_idx) + AttnWrapperBase.pending_prefill_offload_layer_idx = None if num_tasks: suffix = f" ({reason})" if reason else "" logging.debug( @@ -185,24 +198,103 @@ def retire_pending_prefill_offloads_before_layer( *, device: Optional[torch.device] = None, ) -> int: - pending_layer = cls.pending_prefill_offload_layer_idx + del cls + pending_layer = AttnWrapperBase.pending_prefill_offload_layer_idx if pending_layer is None or pending_layer == layer_idx: return 0 - return cls.retire_pending_prefill_offloads( + return AttnWrapperBase.retire_pending_prefill_offloads( device=device, reason=f"before layer {layer_idx}", ) @classmethod def pin_prefill_offload_tensor(cls, tensor: torch.Tensor, layer_idx: int) -> None: - cls.pending_prefill_offload_layer_idx = layer_idx - cls.pending_prefill_offload_tensors.append(tensor) + del cls + AttnWrapperBase.pending_prefill_offload_layer_idx = layer_idx + AttnWrapperBase.pending_prefill_offload_tensors.append(tensor) @classmethod def track_prefill_offload_task(cls, task: object, layer_idx: int) -> None: - cls.pending_prefill_offload_layer_idx = layer_idx + del cls + AttnWrapperBase.pending_prefill_offload_layer_idx = layer_idx if task is not None: - cls.pending_prefill_offload_tasks.append(task) + AttnWrapperBase.pending_prefill_offload_tasks.append(task) + + def prefix_cache_metadata(self): + """Return validated metadata derived from AttnWrapperBase fields.""" + from .prefix_cache import current_or_legacy_prefix_cache_metadata + + return current_or_legacy_prefix_cache_metadata(AttnWrapperBase) + + def offload_prepacked_gqa_kv( + self, + key: torch.Tensor, + value: torch.Tensor, + *, + metadata=None, + track_tasks: bool = False, + sequence_callback=None, + ) -> None: + """Offload prepacked GQA KV with optional prefix-cache offsets.""" + from batchgen.kv_cache.prefill_offload import PrefillHostKVOffloader + + metadata = metadata or self.prefix_cache_metadata() + materialization = AttnWrapperBase.prefill_prefix_materialization + prefix_materialization_active = ( + metadata.prefix_reuse_mode and materialization is not None + ) + should_track = track_tasks or prefix_materialization_active + tracker = self.track_prefill_offload_task if should_track else None + tensor_pinner = self.pin_prefill_offload_tensor if should_track else None + offloader = PrefillHostKVOffloader( + worker_view=getattr(self.core_engine, "host_paged_kv_worker_view", None), + layer_idx=self.layer_idx, + metadata=metadata, + track_task=tracker, + pin_tensor=tensor_pinner, + ) + offloader.offload_gqa( + key=key, + value=value, + sequence_callback=sequence_callback, + ) + if ( + prefix_materialization_active + and AttnWrapperBase.pending_prefill_offload_layer_idx != self.layer_idx + ): + materialization.finish_layer(self.layer_idx) + + def offload_prepacked_mla_kv( + self, + key: torch.Tensor, + *, + metadata=None, + track_tasks: bool = False, + ) -> None: + """Offload prepacked MLA primary KV with optional prefix-cache offsets.""" + from batchgen.kv_cache.prefill_offload import PrefillHostKVOffloader + + metadata = metadata or self.prefix_cache_metadata() + materialization = AttnWrapperBase.prefill_prefix_materialization + prefix_materialization_active = ( + metadata.prefix_reuse_mode and materialization is not None + ) + should_track = track_tasks or prefix_materialization_active + tracker = self.track_prefill_offload_task if should_track else None + tensor_pinner = self.pin_prefill_offload_tensor if should_track else None + offloader = PrefillHostKVOffloader( + worker_view=getattr(self.core_engine, "host_paged_kv_worker_view", None), + layer_idx=self.layer_idx, + metadata=metadata, + track_task=tracker, + pin_tensor=tensor_pinner, + ) + offloader.offload_mla(key=key) + if ( + prefix_materialization_active + and AttnWrapperBase.pending_prefill_offload_layer_idx != self.layer_idx + ): + materialization.finish_layer(self.layer_idx) # Prepack mode state prepack_mode: ClassVar[bool] = False @@ -210,6 +302,10 @@ def track_prefill_offload_task(cls, task: object, layer_idx: int) -> None: prepack_max_seqlen: ClassVar[Optional[int]] = None prepack_num_sequences: ClassVar[Optional[int]] = None prepack_seq_lengths: ClassVar[Optional[List[int]]] = None + prepack_append_seq_lengths: ClassVar[Optional[List[int]]] = None + prepack_prefix_reuse_mode: ClassVar[bool] = False + prepack_prefix_shared_tokens: ClassVar[Optional[List[int]]] = None + prepack_full_seq_lengths: ClassVar[Optional[List[int]]] = None # KV cache state past_key_states: ClassVar[Optional[List[torch.Tensor]]] = None @@ -223,6 +319,7 @@ def track_prefill_offload_task(cls, task: object, layer_idx: int) -> None: # per audit §A finding #8. gpu_paged_kv_manager: ClassVar[Optional[object]] = None host_paged_kv_worker_view: ClassVar[Optional[object]] = None + prefill_prefix_materialization: ClassVar[Optional[object]] = None # DSA auxiliary caches (indexer KV for DeepSeek Sparse Attention) gpu_paged_kv_manager_aux: ClassVar[Optional[object]] = None host_paged_kv_worker_view_aux: ClassVar[Optional[object]] = None diff --git a/batchgen/models/wrappers/prefix_cache.py b/batchgen/models/wrappers/prefix_cache.py new file mode 100644 index 000000000..41f9cf8e3 --- /dev/null +++ b/batchgen/models/wrappers/prefix_cache.py @@ -0,0 +1,238 @@ +"""Prefix-cache metadata compatibility helpers. + +The source of truth for prefill execution metadata is +``ForwardBatchMetadata``. This module only provides the legacy conversion path +for wrappers that still receive state through ``AttnWrapperBase`` class fields. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from batchgen.attention.forward_metadata import ( + ForwardBatchMetadata, + PrefillAttentionMetadata, +) +from batchgen.attention.forward_metadata_context import ( + get_current_forward_batch_metadata, +) + + +def ensure_prefix_cache_forward_metadata(metadata) -> ForwardBatchMetadata: + """Return validated prefill ``ForwardBatchMetadata``. + + Prefix-cache-aware compute/offload paths should consume + ``ForwardBatchMetadata`` directly. Passing ``PrefillAttentionMetadata`` is + intentionally rejected because it lacks global sequence ids. + """ + + if isinstance(metadata, ForwardBatchMetadata): + _require_prefill(metadata) + _validate_forward_metadata(metadata) + return metadata + if isinstance(metadata, PrefillAttentionMetadata): + raise RuntimeError( + "PrefillAttentionMetadata does not carry global sequence ids; " + "pass ForwardBatchMetadata or use AttnWrapperBase-bound fields" + ) + raise TypeError("metadata must be ForwardBatchMetadata") + + +def current_or_legacy_prefix_cache_metadata( + wrapper_cls: type, +) -> ForwardBatchMetadata: + """Prefer bound metadata, otherwise build it from legacy wrapper fields.""" + + metadata = get_current_forward_batch_metadata() + if metadata is not None: + return ensure_prefix_cache_forward_metadata(metadata) + return build_prefix_cache_forward_metadata_from_wrapper_cls(wrapper_cls) + + +def build_prefix_cache_forward_metadata_from_wrapper_cls( + wrapper_cls: type, +) -> ForwardBatchMetadata: + """Build ``ForwardBatchMetadata`` from legacy prepack class variables.""" + + if getattr(wrapper_cls, "phase", None) not in (None, "prefill"): + raise RuntimeError("Prefix cache prepack metadata requires prefill metadata") + + cu_seqlens = _require_attr(wrapper_cls, "prepack_cu_seqlens") + max_seqlen = int(_require_attr(wrapper_cls, "prepack_max_seqlen")) + num_sequences = int(_require_attr(wrapper_cls, "prepack_num_sequences")) + seq_lengths = _int_list(_require_attr(wrapper_cls, "prepack_seq_lengths")) + global_sequence_ids = _int_list(_require_attr(wrapper_cls, "cur_batch")) + append_seq_lengths = _optional_int_list( + getattr(wrapper_cls, "prepack_append_seq_lengths", None) + ) + if append_seq_lengths is None: + append_seq_lengths = list(seq_lengths) + + if len(seq_lengths) != num_sequences: + raise RuntimeError( + "Prefix cache seq_lengths length does not match num_sequences: " + f"{len(seq_lengths)} != {num_sequences}" + ) + if len(global_sequence_ids) != num_sequences: + raise RuntimeError( + "Prefix cache cur_batch length does not match num_sequences: " + f"{len(global_sequence_ids)} != {num_sequences}" + ) + if len(append_seq_lengths) != num_sequences: + raise RuntimeError( + "Prefix cache append_seq_lengths length does not match " + f"num_sequences: {len(append_seq_lengths)} != {num_sequences}" + ) + _validate_append_lengths( + append_seq_lengths=append_seq_lengths, + query_seq_lengths=seq_lengths, + ) + + prefix_reuse_mode = bool( + getattr(wrapper_cls, "prepack_prefix_reuse_mode", False) + ) + if prefix_reuse_mode: + prefix_shared_tokens = _int_list( + _require_attr(wrapper_cls, "prepack_prefix_shared_tokens") + ) + full_seq_lengths = _int_list( + _require_attr(wrapper_cls, "prepack_full_seq_lengths") + ) + if len(prefix_shared_tokens) != num_sequences: + raise RuntimeError( + "Prefix shared token count length does not match batch: " + f"{len(prefix_shared_tokens)} != {num_sequences}" + ) + if len(full_seq_lengths) != num_sequences: + raise RuntimeError( + "Full sequence length metadata length does not match batch: " + f"{len(full_seq_lengths)} != {num_sequences}" + ) + for idx, (append_len, prefix_tokens, full_length) in enumerate( + zip(append_seq_lengths, prefix_shared_tokens, full_seq_lengths) + ): + expected = int(append_len) + int(prefix_tokens) + if expected != int(full_length): + raise RuntimeError( + "Prefix cache full length mismatch at sequence " + f"{idx}: append={append_len}, prefix={prefix_tokens}, " + f"full={full_length}" + ) + kv_seq_lengths = full_seq_lengths + else: + kv_seq_lengths = list(seq_lengths) + + position_ids = getattr(wrapper_cls, "position_ids", None) + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=global_sequence_ids, + prefill=PrefillAttentionMetadata( + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=_build_cu_seqlens_like( + kv_seq_lengths, + reference=cu_seqlens, + ), + max_seqlen_q=max_seqlen, + max_seqlen_k=max(kv_seq_lengths, default=0), + q_seq_lens=seq_lengths, + kv_seq_lens=kv_seq_lengths, + position_ids=position_ids, + append_seq_lens=append_seq_lengths, + ), + ) + + +def _require_prefill(metadata: ForwardBatchMetadata) -> None: + if metadata.phase != "prefill" or metadata.prefill is None: + raise RuntimeError("Prefix cache prepack metadata requires prefill metadata") + + +def _validate_forward_metadata(metadata: ForwardBatchMetadata) -> None: + prefill = metadata.require_prefill() + num_sequences = int(prefill.batch_size) + if len(metadata.global_sequence_ids) != num_sequences: + raise RuntimeError( + "Prefix cache global sequence id count does not match batch: " + f"{len(metadata.global_sequence_ids)} != {num_sequences}" + ) + if len(prefill.kv_seq_lens) != num_sequences: + raise RuntimeError( + "Prefix cache metadata KV length count does not match batch: " + f"{len(prefill.kv_seq_lens)} != {num_sequences}" + ) + if len(metadata.append_seq_lengths) != num_sequences: + raise RuntimeError( + "Prefix cache append length count does not match batch: " + f"{len(metadata.append_seq_lengths)} != {num_sequences}" + ) + if len(prefill.cu_seqlens_q) != num_sequences + 1: + raise RuntimeError( + "Prefix cache cu_seqlens length does not match batch: " + f"{len(prefill.cu_seqlens_q)} != {num_sequences + 1}" + ) + _validate_append_lengths( + append_seq_lengths=metadata.append_seq_lengths, + query_seq_lengths=prefill.q_seq_lens, + ) + prefix_tokens = [ + int(kv_len) - int(append_len) + for kv_len, append_len in zip( + prefill.kv_seq_lens, + metadata.append_seq_lengths, + ) + ] + if any(tokens < 0 for tokens in prefix_tokens): + raise RuntimeError("Prefix cache metadata requires kv lengths >= append lengths") + + +def _require_attr(wrapper_cls: type, name: str): + value = getattr(wrapper_cls, name, None) + if value is None: + raise RuntimeError(f"Prefix cache prepack metadata requires {name}") + return value + + +def _int_list(values: Sequence[int]) -> list[int]: + return [int(value) for value in values] + + +def _optional_int_list(values: Sequence[int] | None) -> list[int] | None: + if values is None: + return None + return _int_list(values) + + +def _validate_append_lengths( + *, + append_seq_lengths: Sequence[int], + query_seq_lengths: Sequence[int], +) -> None: + for idx, (append_len, query_len) in enumerate( + zip(append_seq_lengths, query_seq_lengths) + ): + if int(append_len) < 0 or int(append_len) > int(query_len): + raise RuntimeError( + "Prefix cache append length must be within query length at " + f"sequence {idx}: append={append_len}, query={query_len}" + ) + + +def _build_cu_seqlens_like( + seq_lengths: Sequence[int], + *, + reference, +): + values = [0] + running = 0 + for length in seq_lengths: + running += int(length) + values.append(running) + + if hasattr(reference, "new_tensor"): + return reference.new_tensor(values) + try: + return torch.tensor(values, dtype=torch.int32) + except AttributeError: + return values diff --git a/batchgen/models/wrappers/prefix_gqa_extend.py b/batchgen/models/wrappers/prefix_gqa_extend.py new file mode 100644 index 000000000..03de2d19b --- /dev/null +++ b/batchgen/models/wrappers/prefix_gqa_extend.py @@ -0,0 +1,57 @@ +"""Common GQA prefix-cache extend-prefill helpers for attention wrappers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + + +@dataclass(frozen=True) +class GqaExtendSpec: + """Static GQA dimensions and optional attention modifiers.""" + + num_kv_heads: int + head_dim: int + sinks: Optional[torch.Tensor] = None + softmax_scale: Optional[float] = None + sliding_window: Optional[int] = None + + +def run_prefix_gqa_prefill_attention( + *, + wrapper: object, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + metadata: object, + spec: GqaExtendSpec, +) -> torch.Tensor: + """Run GQA prefill attention with optional cached prefix K/V.""" + from batchgen.attention.forward_metadata_context import ( + get_current_forward_batch_metadata, + ) + from batchgen.attention.prefix_aware_backend import ( + GqaPrefixAwareAttentionBackend, + ) + + forward_metadata = get_current_forward_batch_metadata() + kv_cache_metadata = ( + None if forward_metadata is None else forward_metadata.kv_cache + ) + backend = GqaPrefixAwareAttentionBackend( + layer_idx=int(wrapper.layer_idx), + num_kv_heads=spec.num_kv_heads, + head_dim=spec.head_dim, + sinks=spec.sinks, + softmax_scale=spec.softmax_scale, + sliding_window=spec.sliding_window, + ) + return backend.forward_prefill( + query=query, + key=key, + value=value, + metadata=metadata, + kv_cache_metadata=kv_cache_metadata, + ) diff --git a/batchgen/models/wrappers/prefix_mla_extend.py b/batchgen/models/wrappers/prefix_mla_extend.py new file mode 100644 index 000000000..748a82428 --- /dev/null +++ b/batchgen/models/wrappers/prefix_mla_extend.py @@ -0,0 +1,144 @@ +"""Common MLA prefix-cache extend-prefill helpers for attention wrappers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +import torch + +from batchgen.models.wrappers.prefix_cache import ( + ensure_prefix_cache_forward_metadata, +) +from batchgen.prefix_reuse.materialization import ( + get_prefix_materialization_for_group, +) + + +@dataclass(frozen=True) +class MlaExtendSpec: + """Static MLA dimensions needed by the prefix extend-prefill path.""" + + num_heads: int + kv_lora_rank: int + softmax_scale: float + + +OutputProjectMlaFn = Callable[[torch.Tensor], torch.Tensor] + + +def run_prefix_mla_suffix_prefill_with_projected( + *, + wrapper: object, + query_states: torch.Tensor, + offload_kv: torch.Tensor, + metadata: object, + spec: MlaExtendSpec, + output_projection: OutputProjectMlaFn, + prefill_prefix_materialization: object | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run suffix-only MLA prefill from already projected suffix Q/KV.""" + if prefill_prefix_materialization is None: + raise RuntimeError( + "MLA prefix-cache suffix prefill requires GPU paged materialization" + ) + prefill_prefix_materialization = get_prefix_materialization_for_group( + prefill_prefix_materialization, + group_id=0, + consumer="MLA prefix-cache suffix prefill", + ) + attn_out = run_projected_mla_prefix_attention_from_gpu_pages( + layer_idx=int(wrapper.layer_idx), + query_states=query_states, + offload_kv=offload_kv, + metadata=metadata, + spec=spec, + materialization=prefill_prefix_materialization, + ) + return output_projection(attn_out), offload_kv + + +def run_projected_mla_prefix_attention_from_gpu_pages( + *, + layer_idx: int, + query_states: torch.Tensor, + offload_kv: torch.Tensor | None, + metadata: object, + spec: MlaExtendSpec, + materialization: object, +) -> torch.Tensor: + """Run MLA prefix attention from materialized GPU compressed KV.""" + + metadata = ensure_prefix_cache_forward_metadata(metadata) + manager = materialization.manager + if manager.config.has_v_cache: + raise RuntimeError( + "MLA GPU prefix materialization requires K-only compressed KV pages" + ) + + layer_idx = int(layer_idx) + materialization.wait_for_layer(layer_idx) + + if not metadata.prefix_reuse_mode: + raise RuntimeError( + "MLA GPU prefix materialization requires prefix reuse" + ) + if offload_kv is None: + raise RuntimeError("MLA GPU prefix extend requires suffix KV") + manager.append_layer_prefill_suffix_tokens( + k_tensor=offload_kv, + v_tensor=None, + append_plan=materialization.append_plan, + layer_idx=layer_idx, + ) + + blocked_k, blocked_v, block_table = manager.get_layer_kv_with_page_table( + layer_idx + ) + if blocked_v is not None: + raise RuntimeError( + "MLA GPU prefix materialization unexpectedly has V cache" + ) + if block_table is None: + raise RuntimeError("MLA GPU prefix materialization requires page table") + + return _run_flashinfer_mla_prefix_attention( + query_states=query_states.contiguous(), + blocked_k=blocked_k, + block_table=block_table, + cache_seqlens=materialization.append_plan.cache_seqlens, + slot_indices=materialization.append_plan.slot_indices, + metadata=metadata, + spec=spec, + plan_cache=getattr(materialization, "backend_state", None), + ) + + +def _run_flashinfer_mla_prefix_attention( + *, + query_states: torch.Tensor, + blocked_k: torch.Tensor, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + slot_indices: torch.Tensor, + metadata: object, + spec: MlaExtendSpec, + plan_cache: dict[str, object] | None = None, +) -> torch.Tensor: + """Run FlashInfer MLA paged attention against materialized prefix pages.""" + from batchgen.attention.mla.flashinfer_extend import ( + run_flashinfer_mla_extend_prefill, + ) + + return run_flashinfer_mla_extend_prefill( + query_states=query_states, + compressed_kv_cache=blocked_k, + page_table=block_table, + slot_indices=slot_indices, + cache_seqlens=cache_seqlens, + cu_seqlens_q=metadata.cu_seqlens, + kv_lora_rank=int(spec.kv_lora_rank), + num_heads=int(spec.num_heads), + softmax_scale=float(spec.softmax_scale), + plan_cache=plan_cache, + ) diff --git a/batchgen/models/wrappers/prefix_mla_model_adapters.py b/batchgen/models/wrappers/prefix_mla_model_adapters.py new file mode 100644 index 000000000..b4c92132c --- /dev/null +++ b/batchgen/models/wrappers/prefix_mla_model_adapters.py @@ -0,0 +1,318 @@ +"""Model-specific MLA prefix-cache adapters. + +The page lookup, GPU page materialization, and paged MLA extend prefill live in +generic prefix-cache helpers. This module keeps the remaining model glue in one +place: how each MLA model builds prefix extend contexts and projects attention +output. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +import torch + +from batchgen.attention.mla.prefix_absorb import ( + build_absorbed_mla_query_states, + prefix_rotary_seq_len, + project_absorbed_mla_output, + project_absorbed_mla_output_w8a16, +) +from batchgen.kv_cache.prefill_offload import PrefillHostKVOffloader + +from .attention import AttnWrapperBase +from .prefix_cache import ensure_prefix_cache_forward_metadata +from .prefix_mla_extend import ( + MlaExtendSpec, + run_prefix_mla_suffix_prefill_with_projected, +) + +OutputProjector = Callable[[torch.Tensor], torch.Tensor] +ProjectedQueryBuilder = Callable[[object], torch.Tensor] + + +@dataclass(frozen=True) +class MlaPrefixBackendContext: + """Prefix extend callbacks consumed by the existing MLA prepack backend.""" + + wrapper: object + metadata: object + spec: MlaExtendSpec + suffix_query_builder: ProjectedQueryBuilder + output_projection: OutputProjector + prefill_prefix_materialization: object | None = None + + @property + def prefix_reuse_mode(self) -> bool: + return self.metadata.prefix_reuse_mode + + def rotary_seq_len( + self, + position_ids: torch.Tensor, + fallback_seq_len: int, + ) -> int: + if self.metadata.full_seq_lengths: + return prefix_rotary_seq_len( + max(self.metadata.full_seq_lengths), + position_ids, + ) + return prefix_rotary_seq_len(fallback_seq_len, position_ids) + + def run_suffix_prefill( + self, + projection: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + offload_kv = getattr(projection, "offload_kv", None) + if offload_kv is None: + raise RuntimeError("MLA prefix backend context requires suffix KV") + return run_prefix_mla_suffix_prefill_with_projected( + wrapper=self.wrapper, + query_states=self.suffix_query_builder(projection), + offload_kv=offload_kv, + metadata=self.metadata, + spec=self.spec, + output_projection=self.output_projection, + prefill_prefix_materialization=self.prefill_prefix_materialization, + ) + + +def build_deepseek_prefix_backend_context( + *, + wrapper: object, + metadata: object, +) -> MlaPrefixBackendContext: + metadata = ensure_prefix_cache_forward_metadata(metadata) + return _build_w8a16_prefix_backend_context( + wrapper=wrapper, + metadata=metadata, + model_label="DeepSeek prefix extend", + use_cached_absorb=False, + ) + + +def build_glm5_prefix_backend_context( + *, + wrapper: object, + metadata: object, +) -> MlaPrefixBackendContext: + metadata = ensure_prefix_cache_forward_metadata(metadata) + return _build_w8a16_prefix_backend_context( + wrapper=wrapper, + metadata=metadata, + model_label="GLM-5 prefix prefill", + use_cached_absorb=True, + ) + + +def build_kimi_prefix_backend_context( + *, + wrapper: object, + metadata: object, +) -> MlaPrefixBackendContext: + metadata = ensure_prefix_cache_forward_metadata(metadata) + return MlaPrefixBackendContext( + wrapper=wrapper, + metadata=metadata, + spec=_mla_extend_spec(wrapper), + prefill_prefix_materialization=_prefill_prefix_materialization(wrapper), + suffix_query_builder=lambda projection: build_absorbed_mla_query_states( + q_nope=projection.q_nope, + q_pe=projection.q_pe, + dtype=projection.offload_kv.dtype, + q_absorb=_kimi_q_absorb_weights(wrapper), + ), + output_projection=lambda attn_out: project_absorbed_mla_output( + attn_out=attn_out, + out_absorb=_kimi_out_absorb_weights(wrapper), + v_head_dim=wrapper.module.v_head_dim, + output_projection=wrapper.module.o_proj, + ), + ) + + +def offload_glm5_prepacked_mla_kv( + *, + key: torch.Tensor, + worker_view: object, + layer_idx: int, + metadata: object, +) -> None: + """Offload prepacked GLM-5 k-only MLA/indexer KV with prefix offsets.""" + offloader = PrefillHostKVOffloader( + worker_view=worker_view, + layer_idx=layer_idx, + metadata=metadata, + track_task=AttnWrapperBase.track_prefill_offload_task, + pin_tensor=AttnWrapperBase.pin_prefill_offload_tensor, + ) + offloader.offload_mla(key=key) + + +def _mla_extend_spec(wrapper: object) -> MlaExtendSpec: + attn = wrapper.module + return MlaExtendSpec( + num_heads=attn.num_heads, + kv_lora_rank=attn.kv_lora_rank, + softmax_scale=attn.softmax_scale, + ) + + +def _prefill_prefix_materialization(wrapper: object) -> object | None: + return getattr(wrapper, "prefill_prefix_materialization", None) + + +def _build_w8a16_prefix_backend_context( + *, + wrapper: object, + metadata: object, + model_label: str, + use_cached_absorb: bool, +) -> MlaPrefixBackendContext: + return MlaPrefixBackendContext( + wrapper=wrapper, + metadata=metadata, + spec=_mla_extend_spec(wrapper), + prefill_prefix_materialization=_prefill_prefix_materialization(wrapper), + suffix_query_builder=lambda projection: build_absorbed_mla_query_states( + q_nope=projection.q_nope, + q_pe=projection.q_pe, + dtype=projection.offload_kv.dtype, + q_absorb=_w8a16_q_absorb_weights( + wrapper, + model_label=model_label, + use_cached_absorb=use_cached_absorb, + ), + ), + output_projection=lambda attn_out: _project_w8a16_absorbed_output( + wrapper=wrapper, + attn_out=attn_out, + out_absorb=_w8a16_out_absorb_weights( + wrapper, + model_label=model_label, + use_cached_absorb=use_cached_absorb, + ), + model_label=model_label, + ), + ) + + +def _project_w8a16_absorbed_output( + *, + wrapper: object, + attn_out: torch.Tensor, + out_absorb: torch.Tensor, + model_label: str, +) -> torch.Tensor: + attn = wrapper.module + from batchgen.attention.mla.fa3_backend import select_w8a16_gemm + + return project_absorbed_mla_output_w8a16( + attn_out=attn_out, + out_absorb=out_absorb, + v_head_dim=attn.v_head_dim, + o_proj_weight=attn.o_proj.weight.data, + o_proj_scale=_weight_scale( + wrapper, model_label, ("o_proj.weight_scale_inv",) + )["o_proj.weight_scale_inv"], + gemm=select_w8a16_gemm(), + ) + + +def _w8a16_q_absorb_weights( + wrapper: object, + *, + model_label: str, + use_cached_absorb: bool, +) -> torch.Tensor: + if ( + use_cached_absorb + and getattr(wrapper, "_cached_q_absorb", None) is not None + ): + return wrapper._cached_q_absorb + attn = wrapper.module + if use_cached_absorb and getattr(attn, "q_absorb", None) is not None: + return attn.q_absorb + kv_b_proj = _dequantized_kv_b_proj(wrapper, model_label) + return kv_b_proj[:, : attn.qk_nope_head_dim, :] + + +def _w8a16_out_absorb_weights( + wrapper: object, + *, + model_label: str, + use_cached_absorb: bool, +) -> torch.Tensor: + if ( + use_cached_absorb + and getattr(wrapper, "_cached_out_absorb", None) is not None + ): + return wrapper._cached_out_absorb + attn = wrapper.module + if use_cached_absorb and getattr(attn, "out_absorb", None) is not None: + return attn.out_absorb + kv_b_proj = _dequantized_kv_b_proj(wrapper, model_label) + return kv_b_proj[:, attn.qk_nope_head_dim :, :] + + +def _dequantized_kv_b_proj(wrapper: object, model_label: str) -> torch.Tensor: + attn = wrapper.module + weight_scale = _weight_scale( + wrapper, + model_label, + ("kv_b_proj.weight_scale_inv",), + ) + + from batchgen.attention.mla.flashmla_backend import ( + deepseek_v3_dequantization, + ) + + return deepseek_v3_dequantization( + attn.kv_b_proj.weight.data, + weight_scale["kv_b_proj.weight_scale_inv"], + ).view( + attn.num_heads, + -1, + attn.kv_lora_rank, + ) + + +def _kimi_q_absorb_weights(wrapper: object) -> torch.Tensor: + attn = wrapper.module + if getattr(attn, "q_absorb", None) is not None: + return attn.q_absorb + return _kimi_kv_b_proj(wrapper)[:, : attn.qk_nope_head_dim, :] + + +def _kimi_out_absorb_weights(wrapper: object) -> torch.Tensor: + attn = wrapper.module + if getattr(attn, "out_absorb", None) is not None: + return attn.out_absorb + return _kimi_kv_b_proj(wrapper)[:, attn.qk_nope_head_dim :, :] + + +def _kimi_kv_b_proj(wrapper: object) -> torch.Tensor: + attn = wrapper.module + return attn.kv_b_proj.weight.data.view( + attn.num_heads, + -1, + attn.kv_lora_rank, + ) + + +def _weight_scale( + wrapper: object, + model_label: str, + required_keys: tuple[str, ...], +) -> dict: + weight_scale = getattr(wrapper, "weight_dequant_scale", None) + missing = [ + key + for key in required_keys + if weight_scale is None or key not in weight_scale + ] + if missing: + raise RuntimeError( + f"{model_label} requires weight scales: {', '.join(missing)}" + ) + return weight_scale diff --git a/batchgen/moe/fused_wgmma_grouped.py b/batchgen/moe/fused_wgmma_grouped.py index ece624ba0..489bff4db 100644 --- a/batchgen/moe/fused_wgmma_grouped.py +++ b/batchgen/moe/fused_wgmma_grouped.py @@ -324,6 +324,7 @@ def fused_mxfp4_grouped_moe_forward_cuda_routing( gate_bias_ptrs: torch.Tensor = None, up_bias_ptrs: torch.Tensor = None, down_bias_ptrs: torch.Tensor = None, + output: torch.Tensor = None, ) -> torch.Tensor: """End-to-end grouped MXFP4 MoE forward using WGMMA + CUDA routing. @@ -346,6 +347,7 @@ def fused_mxfp4_grouped_moe_forward_cuda_routing( gate_bias_ptrs: Gate bias pointer array [num_experts] int64, or None up_bias_ptrs: Up bias pointer array [num_experts] int64, or None down_bias_ptrs: Down bias pointer array [num_experts] int64, or None + output: Optional pre-allocated output buffer [batch*seq, hidden] BF16 Returns: Output [batch*seq, hidden] BF16 @@ -410,6 +412,7 @@ def fused_mxfp4_grouped_moe_forward_cuda_routing( output = reduce_weighted_scatter_cuda( sorted_output, topk_pos, topk_weights, num_tokens, hidden_size, K_topk, + output=output, ) return output diff --git a/batchgen/prefill/__init__.py b/batchgen/prefill/__init__.py index fd0ef881d..d50bb5511 100644 --- a/batchgen/prefill/__init__.py +++ b/batchgen/prefill/__init__.py @@ -1,13 +1,20 @@ """Prefill utilities for efficient batch processing.""" +from .attention_metadata_builder import build_prefill_forward_metadata +from .prefix_reuse import ( + PrefixReusePrefillPlan, + PrefixReuseSequencePlan, + build_prefix_reuse_prefill_plan, + split_prefix_reuse_plan_for_micro_batch, +) from .prepack import ( PrepackMetadata, bin_pack_first_fit_decreasing, - prepack_sequences, - unpack_outputs, - unpack_last_token_logits, create_block_diagonal_attention_mask, get_prepack_stats, + prepack_sequences, + unpack_last_token_logits, + unpack_outputs, ) __all__ = [ @@ -18,4 +25,9 @@ "unpack_last_token_logits", "create_block_diagonal_attention_mask", "get_prepack_stats", + "PrefixReusePrefillPlan", + "PrefixReuseSequencePlan", + "build_prefix_reuse_prefill_plan", + "split_prefix_reuse_plan_for_micro_batch", + "build_prefill_forward_metadata", ] diff --git a/batchgen/prefill/attention_metadata_builder.py b/batchgen/prefill/attention_metadata_builder.py new file mode 100644 index 000000000..d01df6d6e --- /dev/null +++ b/batchgen/prefill/attention_metadata_builder.py @@ -0,0 +1,171 @@ +"""Builders for prefill attention forward metadata.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +from batchgen.attention.forward_metadata import ( + ForwardBatchMetadata, + KVCacheMetadata, + PrefillAttentionMetadata, +) +from batchgen.batch_order import PrefillSequenceSpan +from batchgen.prefill.prepack import PrepackMetadata +from batchgen.prefill.prefix_reuse import PrefixReusePrefillPlan + + +def build_prefill_forward_metadata( + *, + prepack_metadata: PrepackMetadata, + batch_spans: Sequence[PrefillSequenceSpan], + seq_start: int, + seq_end: int, + position_ids: torch.Tensor, + device: torch.device, + prefix_reuse_plan: Optional[PrefixReusePrefillPlan] = None, + kv_cache_metadata: Optional[KVCacheMetadata] = None, +) -> ForwardBatchMetadata: + """Build first-class metadata for one prepacked prefill micro-batch.""" + + if seq_start < 0 or seq_end < seq_start: + raise ValueError(f"Invalid sequence range [{seq_start}, {seq_end})") + q_seq_lens = [ + int(length) + for length in prepack_metadata.original_seq_lengths[seq_start:seq_end] + ] + if len(q_seq_lens) != len(batch_spans): + raise ValueError( + f"batch_spans length must match micro-batch sequence count: " + f"{len(batch_spans)} != {len(q_seq_lens)}" + ) + span_seq_lens = [int(span.seq_len) for span in batch_spans] + if span_seq_lens != q_seq_lens: + raise ValueError( + f"batch span sequence lengths do not match prepack lengths: " + f"{span_seq_lens} != {q_seq_lens}" + ) + + global_sequence_ids = [int(span.global_seq_id) for span in batch_spans] + total_query_tokens = sum(q_seq_lens) + if position_ids.numel() != total_query_tokens: + raise ValueError( + f"position_ids length must match micro-batch query tokens: " + f"{position_ids.numel()} != {total_query_tokens}" + ) + position_ids = position_ids.to(device=device) + cu_seqlens_q = _build_cu_seqlens(q_seq_lens, device=device) + + if prefix_reuse_plan is None: + kv_seq_lens = list(q_seq_lens) + append_seq_lens = list(q_seq_lens) + else: + kv_seq_lens = _build_prefix_reuse_kv_seq_lens( + plan=prefix_reuse_plan, + seq_start=seq_start, + seq_end=seq_end, + q_seq_lens=q_seq_lens, + global_sequence_ids=global_sequence_ids, + ) + append_seq_lens = _build_prefix_reuse_append_seq_lens( + plan=prefix_reuse_plan, + seq_start=seq_start, + seq_end=seq_end, + q_seq_lens=q_seq_lens, + ) + cu_seqlens_k = _build_cu_seqlens(kv_seq_lens, device=device) + + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=global_sequence_ids, + prefill=PrefillAttentionMetadata( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max(q_seq_lens, default=0), + max_seqlen_k=max(kv_seq_lens, default=0), + q_seq_lens=q_seq_lens, + kv_seq_lens=kv_seq_lens, + position_ids=position_ids, + append_seq_lens=append_seq_lens, + ), + kv_cache=kv_cache_metadata, + ) + + +def _build_prefix_reuse_kv_seq_lens( + *, + plan: PrefixReusePrefillPlan, + seq_start: int, + seq_end: int, + q_seq_lens: Sequence[int], + global_sequence_ids: Sequence[int], +) -> list[int]: + sequence_plans = plan.sequences[seq_start:seq_end] + if len(sequence_plans) != len(q_seq_lens): + raise ValueError( + f"prefix reuse plan slice length mismatch: " + f"{len(sequence_plans)} != {len(q_seq_lens)}" + ) + + suffix_lens: list[int] = [] + kv_seq_lens: list[int] = [] + plan_sequence_ids: list[int] = [] + for item in sequence_plans: + suffix_lens.append(int(item.suffix_length)) + kv_seq_lens.append(int(item.full_logical_context_length)) + plan_sequence_ids.append(int(item.sequence_id)) + + if len(suffix_lens) != len(q_seq_lens): + raise ValueError( + f"prefix reuse suffix length count does not match query lengths: " + f"{len(suffix_lens)} != {len(q_seq_lens)}" + ) + for idx, (suffix_len, query_len) in enumerate(zip(suffix_lens, q_seq_lens)): + if suffix_len < 0 or suffix_len > int(query_len): + raise ValueError( + f"prefix reuse append length must be within query length at " + f"sequence {idx}: append={suffix_len}, query={query_len}" + ) + if plan_sequence_ids != [int(seq_id) for seq_id in global_sequence_ids]: + raise ValueError( + f"prefix reuse sequence ids do not match batch spans: " + f"{plan_sequence_ids} != {list(global_sequence_ids)}" + ) + return kv_seq_lens + + +def _build_prefix_reuse_append_seq_lens( + *, + plan: PrefixReusePrefillPlan, + seq_start: int, + seq_end: int, + q_seq_lens: Sequence[int], +) -> list[int]: + sequence_plans = plan.sequences[seq_start:seq_end] + append_lens = [int(item.suffix_length) for item in sequence_plans] + if len(append_lens) != len(q_seq_lens): + raise ValueError( + f"prefix reuse append length mismatch: " + f"{len(append_lens)} != {len(q_seq_lens)}" + ) + for idx, (append_len, query_len) in enumerate(zip(append_lens, q_seq_lens)): + if append_len < 0 or append_len > int(query_len): + raise ValueError( + f"prefix reuse append length must be within query length at " + f"sequence {idx}: append={append_len}, query={query_len}" + ) + return append_lens + + +def _build_cu_seqlens( + seq_lens: Sequence[int], + *, + device: torch.device, +) -> torch.Tensor: + values = [0] + running = 0 + for length in seq_lens: + running += int(length) + values.append(running) + return torch.tensor(values, dtype=torch.int32, device=device) diff --git a/batchgen/prefill/prefix_reuse.py b/batchgen/prefill/prefix_reuse.py new file mode 100644 index 000000000..373073c87 --- /dev/null +++ b/batchgen/prefill/prefix_reuse.py @@ -0,0 +1,177 @@ +"""Side-effect-free planning helpers for prefix-reuse prefill.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + + +@dataclass(frozen=True) +class PrefixReuseSequencePlan: + local_idx: int + sequence_id: int + prompt_length: int + prefix_shared_tokens: int + suffix_start_pos: int + suffix_length: int + full_logical_context_length: int + fallback_reason: Optional[str] = None + + +@dataclass(frozen=True) +class PrefixReusePrefillPlan: + sequences: list[PrefixReuseSequencePlan] + suffix_input_ids: list[torch.Tensor] + suffix_position_ids: list[torch.Tensor] + cache_seqlens: torch.Tensor + total_prompt_tokens: int + total_suffix_tokens: int + saved_prefill_tokens: int + + +def _normalize_input_ids( + input_ids: torch.Tensor, prompt_length: int +) -> torch.Tensor: + if input_ids.dim() == 2: + if input_ids.size(0) != 1: + raise ValueError( + f"2D input_ids must have batch size 1, got shape={tuple(input_ids.shape)}" + ) + input_ids = input_ids[0] + elif input_ids.dim() != 1: + raise ValueError( + f"input_ids must be 1D or [1, S], got shape={tuple(input_ids.shape)}" + ) + if prompt_length < 0: + raise ValueError( + f"prompt_length must be non-negative, got {prompt_length}" + ) + if input_ids.numel() < prompt_length: + raise ValueError( + f"input_ids length {input_ids.numel()} is shorter than prompt_length {prompt_length}" + ) + return input_ids[:prompt_length] + + +def build_prefix_reuse_prefill_plan( + *, + local_indices: Sequence[int], + sequence_ids: Sequence[int], + input_ids: Sequence[torch.Tensor], + prompt_lengths: Sequence[int], + prefix_shared_tokens: Sequence[int], + device: Optional[torch.device] = None, +) -> PrefixReusePrefillPlan: + """Build suffix-only prefill metadata without mutating runtime state. + + ``prefix_shared_tokens`` must already use the canonical compute semantic: + it is the prefix length actually reused by this prefill. A raw full hit is + normalized by the lookup layer to ``prompt_length - 1`` so the final prompt + token is represented as a regular one-token extend prefill. + """ + + count = len(local_indices) + if not ( + len(sequence_ids) == count + and len(input_ids) == count + and len(prompt_lengths) == count + and len(prefix_shared_tokens) == count + ): + raise ValueError("All input sequences must have the same length") + + plans: list[PrefixReuseSequencePlan] = [] + suffix_input_ids: list[torch.Tensor] = [] + suffix_position_ids: list[torch.Tensor] = [] + cache_seqlens: list[int] = [] + total_prompt_tokens = 0 + total_suffix_tokens = 0 + + for idx in range(count): + prompt_length = int(prompt_lengths[idx]) + shared_tokens = int(prefix_shared_tokens[idx]) + prompt_ids = _normalize_input_ids(input_ids[idx], prompt_length) + if prompt_length <= 0: + raise ValueError( + f"prompt_length must be positive for prefix reuse, got {prompt_length}" + ) + if shared_tokens < 0: + raise ValueError( + f"prefix_shared_tokens must be non-negative, got {shared_tokens}" + ) + if shared_tokens >= prompt_length: + raise ValueError( + "prefix_shared_tokens must be smaller than prompt_length; " + f"got prefix_shared_tokens={shared_tokens}, " + f"prompt_length={prompt_length}. Raw full hits must be " + "normalized to prompt_length - 1 before planning." + ) + + suffix_start = shared_tokens + suffix_length = prompt_length - suffix_start + target_device = device if device is not None else prompt_ids.device + suffix_ids = prompt_ids[suffix_start:prompt_length].to(target_device) + position_ids = torch.arange( + suffix_start, + prompt_length, + dtype=torch.long, + device=target_device, + ) + + plans.append( + PrefixReuseSequencePlan( + local_idx=int(local_indices[idx]), + sequence_id=int(sequence_ids[idx]), + prompt_length=prompt_length, + prefix_shared_tokens=suffix_start, + suffix_start_pos=suffix_start, + suffix_length=suffix_length, + full_logical_context_length=prompt_length, + ) + ) + suffix_input_ids.append(suffix_ids) + suffix_position_ids.append(position_ids) + cache_seqlens.append(suffix_start) + total_prompt_tokens += prompt_length + total_suffix_tokens += suffix_length + + cache_device = device if device is not None else torch.device("cpu") + return PrefixReusePrefillPlan( + sequences=plans, + suffix_input_ids=suffix_input_ids, + suffix_position_ids=suffix_position_ids, + cache_seqlens=torch.tensor( + cache_seqlens, dtype=torch.int32, device=cache_device + ), + total_prompt_tokens=total_prompt_tokens, + total_suffix_tokens=total_suffix_tokens, + saved_prefill_tokens=total_prompt_tokens - total_suffix_tokens, + ) + + +def split_prefix_reuse_plan_for_micro_batch( + plan: PrefixReusePrefillPlan, + seq_start: int, + seq_end: int, +) -> PrefixReusePrefillPlan: + if seq_start < 0 or seq_end < seq_start or seq_end > len(plan.sequences): + raise ValueError( + f"Invalid micro-batch range [{seq_start}, {seq_end}) for " + f"{len(plan.sequences)} sequences" + ) + sequences = plan.sequences[seq_start:seq_end] + suffix_input_ids = plan.suffix_input_ids[seq_start:seq_end] + suffix_position_ids = plan.suffix_position_ids[seq_start:seq_end] + cache_seqlens = plan.cache_seqlens[seq_start:seq_end].clone() + total_prompt_tokens = sum(item.prompt_length for item in sequences) + total_suffix_tokens = sum(item.suffix_length for item in sequences) + return PrefixReusePrefillPlan( + sequences=list(sequences), + suffix_input_ids=list(suffix_input_ids), + suffix_position_ids=list(suffix_position_ids), + cache_seqlens=cache_seqlens, + total_prompt_tokens=total_prompt_tokens, + total_suffix_tokens=total_suffix_tokens, + saved_prefill_tokens=total_prompt_tokens - total_suffix_tokens, + ) diff --git a/batchgen/prefix_reuse/__init__.py b/batchgen/prefix_reuse/__init__.py new file mode 100644 index 000000000..c8e5694ba --- /dev/null +++ b/batchgen/prefix_reuse/__init__.py @@ -0,0 +1,87 @@ +"""Prefix KV reuse helpers.""" + +from .config import ( + PrefixCacheRuntimeConfig, + PrefixKVGroupSemantic, + PrefixKVGroupSpec, + build_prefix_cache_namespace_digest, + build_prefix_cache_runtime_config, + build_prefix_cache_runtime_config_from_specs, + create_host_prefix_cache_coordinator, + derive_prefix_cache_shm_name, +) +from .commit import ( + PrefixCommitRequest, + aligned_prefix_tokens, + build_committable_prefix_token_ids, + build_prefix_commit_request, + collect_required_group_pages_for_commit, +) +from .eviction import ( + PrefixAllocationEvictionResult, + PrefixCommitRetryResult, + commit_prefix_pages_with_capacity_retry, + evict_prefix_pages_for_host_allocation, + release_evicted_prefix_pages, +) +from .materialization import ( + PrefixMaterializationBundle, + PrefixMaterializationSequence, + RollingSingleGroupPrefixMaterialization, + SingleGroupPrefixMaterialization, + get_prefix_materialization_for_group, + materialize_single_group_lookup_results, + materialize_single_group_prefix_pages, +) +from .prefill import ( + PrefixCachePrefillInputs, + PrefixCachePrefillEstimate, + PrefixCachePrefillLookup, + build_prefix_cache_prefill_inputs, + effective_prefix_shared_tokens, + estimate_prefix_cache_for_prefill, + lookup_prefix_cache_for_prefill, +) +from .worker_commit import ( + build_sequence_prefix_commit_request, + retain_newly_committed_prefix_pages, + sequence_token_ids_for_prefix_commit, +) + +__all__ = [ + "PrefixCacheRuntimeConfig", + "PrefixKVGroupSemantic", + "PrefixKVGroupSpec", + "build_prefix_cache_namespace_digest", + "build_prefix_cache_runtime_config", + "build_prefix_cache_runtime_config_from_specs", + "create_host_prefix_cache_coordinator", + "derive_prefix_cache_shm_name", + "PrefixCommitRequest", + "aligned_prefix_tokens", + "build_committable_prefix_token_ids", + "build_prefix_commit_request", + "collect_required_group_pages_for_commit", + "PrefixAllocationEvictionResult", + "PrefixCommitRetryResult", + "commit_prefix_pages_with_capacity_retry", + "evict_prefix_pages_for_host_allocation", + "release_evicted_prefix_pages", + "PrefixMaterializationBundle", + "PrefixMaterializationSequence", + "RollingSingleGroupPrefixMaterialization", + "SingleGroupPrefixMaterialization", + "get_prefix_materialization_for_group", + "materialize_single_group_lookup_results", + "materialize_single_group_prefix_pages", + "PrefixCachePrefillInputs", + "PrefixCachePrefillEstimate", + "PrefixCachePrefillLookup", + "build_prefix_cache_prefill_inputs", + "effective_prefix_shared_tokens", + "estimate_prefix_cache_for_prefill", + "lookup_prefix_cache_for_prefill", + "build_sequence_prefix_commit_request", + "retain_newly_committed_prefix_pages", + "sequence_token_ids_for_prefix_commit", +] diff --git a/batchgen/prefix_reuse/admin.py b/batchgen/prefix_reuse/admin.py new file mode 100644 index 000000000..118c4e8f5 --- /dev/null +++ b/batchgen/prefix_reuse/admin.py @@ -0,0 +1,145 @@ +"""Admin helpers for managing the Host prefix cache.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from batchgen.prefix_reuse.eviction import release_evicted_prefix_pages + +_STATS_FIELDS = ( + "resident_nodes", + "active_attachments", + "pending_load_entries", + "pending_load_refs", + "used_group_entries", + "used_page_handles", + "lookup_hits", + "lookup_misses", + "evicted_nodes", + "eviction_protected_skips", +) + +_EVICTION_FIELDS = ( + "evicted_nodes", + "protected_nodes", + "freed_group_entries", + "freed_page_handles", +) + + +def clear_host_prefix_cache( + *, + coordinator: Any, + host_kv_views_by_group: Mapping[int, Any], +) -> dict[str, Any]: + """Clear unprotected prefix-cache entries and release their Host KV pages.""" + + stats_before = _object_int_fields(coordinator.get_stats(), _STATS_FIELDS) + eviction_result = coordinator.clear_unprotected() + released_pages_by_group = release_evicted_prefix_pages( + eviction_result=eviction_result, + worker_views_by_group=host_kv_views_by_group, + ) + stats_after = _object_int_fields(coordinator.get_stats(), _STATS_FIELDS) + + return { + "status": "success", + "cleared_all": stats_after["resident_nodes"] == 0, + "stats_before": stats_before, + "stats_after": stats_after, + "eviction": { + **_object_int_fields(eviction_result, _EVICTION_FIELDS), + "evicted_pages_by_group": _evicted_page_counts_by_group( + eviction_result + ), + "released_pages_by_group": { + int(group_id): int(count) + for group_id, count in sorted(released_pages_by_group.items()) + }, + }, + } + + +def pin_host_prefix_cache( + *, + coordinator: Any, + namespace_digest: tuple[int, int, int, int], + token_id_batches: list[list[int]], +) -> dict[str, Any]: + """Pin existing prefix-cache entries by holding lookup attachments.""" + + handles: list[int] = [] + cached_tokens: list[int] = [] + missed_count = 0 + for token_ids in token_id_batches: + result = coordinator.lookup_and_attach( + list(namespace_digest), + [int(token_id) for token_id in token_ids], + ) + handle = int(result.attachment_handle) + if handle == 0: + missed_count += 1 + continue + handles.append(handle) + cached_tokens.append(int(result.common_cached_tokens)) + + return { + "status": "success", + "requested": len(token_id_batches), + "pinned": len(handles), + "missed": missed_count, + "cached_tokens": sum(cached_tokens), + "cached_tokens_by_request": cached_tokens, + "attachment_handles": handles, + } + + +def unpin_host_prefix_cache( + *, + coordinator: Any, + attachment_handles: list[int], +) -> dict[str, Any]: + """Release previously pinned prefix-cache lookup attachments.""" + + released = 0 + for handle in attachment_handles: + coordinator.release_attachment(int(handle)) + released += 1 + return { + "status": "success", + "released": released, + } + + +def host_kv_views_by_prefix_group( + *, + primary_host_kv: Any, + auxiliary_host_kv: Any | None, +) -> dict[int, Any]: + """Build the prefix-cache group -> Host KV owner map used for page release.""" + + views_by_group = getattr(primary_host_kv, "views_by_group", None) + if views_by_group is not None: + return { + int(group_id): view for group_id, view in views_by_group().items() + } + + result = {0: primary_host_kv} + if auxiliary_host_kv is not None: + result[1] = auxiliary_host_kv + return result + + +def _object_int_fields(obj: Any, fields: tuple[str, ...]) -> dict[str, int]: + return {field: int(getattr(obj, field)) for field in fields} + + +def _evicted_page_counts_by_group(eviction_result: Any) -> dict[int, int]: + page_counts: dict[int, int] = {} + for group_pages in eviction_result.evicted_group_pages: + group_id = int(group_pages.group_id) + page_counts[group_id] = page_counts.get(group_id, 0) + len( + group_pages.pages + ) + return dict(sorted(page_counts.items())) diff --git a/batchgen/prefix_reuse/commit.py b/batchgen/prefix_reuse/commit.py new file mode 100644 index 000000000..c6def8522 --- /dev/null +++ b/batchgen/prefix_reuse/commit.py @@ -0,0 +1,194 @@ +"""Helpers for publishing completed Host KV pages to the prefix cache.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Sequence + +from batchgen.prefix_reuse.config import PrefixKVGroupSpec + + +@dataclass(frozen=True) +class PrefixCommitRequest: + namespace_digest: tuple[int, int, int, int] + token_ids: list[int] + commit_tokens: int + publish_boundary_tokens: int + group_pages: list[object] + page_ids_by_group: dict[int, list[int]] + raw_page_tokens_by_group: dict[int, int] + + def commit(self, coordinator: object): + if hasattr(coordinator, "commit_prefix_page_ids"): + return coordinator.commit_prefix_page_ids( + list(self.namespace_digest), + self.token_ids, + int(self.commit_tokens), + [ + (int(group_id), list(page_ids)) + for group_id, page_ids in sorted( + self.page_ids_by_group.items() + ) + ], + ) + return coordinator.commit_prefix_pages( + list(self.namespace_digest), + self.token_ids, + int(self.commit_tokens), + self.group_pages, + ) + + def capacity_requirements(self) -> tuple[int, int, int]: + """Return worst-case metadata slots needed for this commit. + + The coordinator skips entries that already exist, so this intentionally + overestimates on the capacity-failure path. It avoids inspecting + shared-memory internals from Python while still evicting enough metadata + before retrying once. + """ + + boundary = int(self.publish_boundary_tokens) + commit_tokens = int(self.commit_tokens) + node_count = commit_tokens // boundary + group_entry_count = 0 + page_handle_count = 0 + raw_start_token = 0 + for raw_end_token in range(boundary, commit_tokens + 1, boundary): + for group_pages in self.group_pages: + group_id = int(group_pages.group_id) + raw_page_tokens = int(self.raw_page_tokens_by_group[group_id]) + if ( + raw_start_token % raw_page_tokens != 0 + or raw_end_token % raw_page_tokens != 0 + ): + continue + first_page = raw_start_token // raw_page_tokens + page_count = ( + raw_end_token - raw_start_token + ) // raw_page_tokens + if len(group_pages.pages) < first_page + page_count: + continue + group_entry_count += 1 + page_handle_count += page_count + raw_start_token = raw_end_token + return node_count, group_entry_count, page_handle_count + + +def aligned_prefix_tokens(total_tokens: int, publish_boundary_tokens: int) -> int: + """Return the longest prefix length that can be safely published.""" + + boundary = int(publish_boundary_tokens) + if boundary <= 0: + raise ValueError("publish_boundary_tokens must be positive") + token_count = max(0, int(total_tokens)) + return (token_count // boundary) * boundary + + +def build_committable_prefix_token_ids( + *, + prompt_token_ids: Sequence[int], + decoded_token_ids: Sequence[int] = (), + decoded_start: int = 0, + max_tokens: int | None = None, +) -> list[int]: + """Build the logical token prefix represented by a sequence Host KV table.""" + + tokens = [int(token_id) for token_id in prompt_token_ids] + start = max(0, int(decoded_start)) + if start < len(decoded_token_ids): + tokens.extend(int(token_id) for token_id in decoded_token_ids[start:]) + if max_tokens is not None: + return tokens[: max(0, int(max_tokens))] + return tokens + + +def build_prefix_commit_request( + *, + core_engine_module: object, + namespace_digest: Sequence[int], + token_ids: Sequence[int], + publish_boundary_tokens: int, + pages_by_group: Mapping[int, Sequence[int | object]], + raw_page_tokens_by_group: Mapping[int, int], +) -> PrefixCommitRequest | None: + """Build a page-aligned prefix cache commit request. + + The coordinator indexes existing Host KV pages. Page allocation, page + ownership, and eviction-side page release stay with the Host KV managers. + """ + + commit_tokens = aligned_prefix_tokens( + len(token_ids), publish_boundary_tokens + ) + if commit_tokens == 0: + return None + + group_pages = [] + page_ids_by_group = {} + for group_id, page_handles in sorted(pages_by_group.items()): + page_ids = [_host_page_id(page) for page in page_handles] + page_ids_by_group[int(group_id)] = page_ids + group = core_engine_module.GroupCommitPages() + group.group_id = int(group_id) + group.pages = [ + _to_host_page_handle(core_engine_module, page) + for page in page_handles + ] + group_pages.append(group) + + return PrefixCommitRequest( + namespace_digest=tuple(int(value) for value in namespace_digest), + token_ids=[int(token_id) for token_id in token_ids], + commit_tokens=commit_tokens, + publish_boundary_tokens=int(publish_boundary_tokens), + group_pages=group_pages, + page_ids_by_group=page_ids_by_group, + raw_page_tokens_by_group={ + int(group_id): int(raw_page_tokens) + for group_id, raw_page_tokens in raw_page_tokens_by_group.items() + }, + ) + + +def collect_required_group_pages_for_commit( + *, + worker_views_by_group: Mapping[int, object], + sequence_id: int, + commit_tokens: int, + group_specs: Sequence[PrefixKVGroupSpec], +) -> dict[int, list[int]]: + """Collect existing Host KV page ids for a page-aligned commit.""" + + result: dict[int, list[int]] = {} + for spec in group_specs: + if not spec.required_for_reuse: + continue + worker_view = worker_views_by_group.get(int(spec.group_id)) + if worker_view is None: + raise RuntimeError( + f"missing Host KV worker view for prefix group {spec.group_id}" + ) + page_count = int(commit_tokens) // int(spec.raw_page_tokens) + page_table = worker_view.build_page_table([int(sequence_id)]) + pages = list(page_table[0])[:page_count] + if len(pages) != page_count: + raise RuntimeError( + f"prefix group {spec.group_id} has {len(pages)} pages for " + f"sequence {sequence_id}, expected {page_count}" + ) + result[int(spec.group_id)] = [int(page_id) for page_id in pages] + return result + + +def _to_host_page_handle(core_engine_module: object, page: int | object): + if hasattr(page, "page_id"): + return page + handle = core_engine_module.HostPageHandle() + handle.page_id = int(page) + return handle + + +def _host_page_id(page: int | object) -> int: + if hasattr(page, "page_id"): + return int(page.page_id) + return int(page) diff --git a/batchgen/prefix_reuse/config.py b/batchgen/prefix_reuse/config.py new file mode 100644 index 000000000..cc5fedf43 --- /dev/null +++ b/batchgen/prefix_reuse/config.py @@ -0,0 +1,275 @@ +"""Runtime configuration helpers for Host-side prefix reuse. + +This module deliberately keeps the Python-side configuration lightweight: +user-facing CLI only enables/disables prefix reuse, while shared-memory names, +group semantics, hash granularity, and table capacities are derived from the +model and Host KV profile. +""" + +from __future__ import annotations + +import hashlib +import math +import re +from dataclasses import dataclass +from enum import Enum +from typing import Iterable, Sequence + + +class PrefixKVGroupSemantic(str, Enum): + FULL_KV = "full_kv" + MLA_COMPRESSED_KV = "mla_compressed_kv" + SWA_KV = "swa_kv" + COMPRESSED_RATIO_KV = "compressed_ratio_kv" + + +@dataclass(frozen=True) +class PrefixKVGroupSpec: + group_id: int + semantic: PrefixKVGroupSemantic + required_for_reuse: bool + raw_page_tokens: int + compression_ratio: int = 1 + + +@dataclass(frozen=True) +class PrefixCacheRuntimeConfig: + shm_name: str + namespace_digest: tuple[int, int, int, int] + group_specs: tuple[PrefixKVGroupSpec, ...] + hash_block_tokens: int + publish_boundary_tokens: int + max_nodes: int + max_group_entries: int + max_page_handles: int + max_attachments: int + debug_stats: bool = False + + def to_core_config(self, core_engine_module): + """Build a core_engine.HostPrefixCacheConfig instance.""" + + core_config = core_engine_module.HostPrefixCacheConfig() + core_config.shm_name = self.shm_name + core_config.hash_block_tokens = int(self.hash_block_tokens) + core_config.max_nodes = int(self.max_nodes) + core_config.max_group_entries = int(self.max_group_entries) + core_config.max_page_handles = int(self.max_page_handles) + core_config.max_attachments = int(self.max_attachments) + core_config.group_specs = [ + _to_core_group_spec(core_engine_module, spec) + for spec in self.group_specs + ] + return core_config + + +def build_prefix_cache_runtime_config( + *, + model_name: str, + kv_dtype: str, + host_kv_cache_size_bytes: int, + debug_stats: bool = False, +) -> PrefixCacheRuntimeConfig: + """Derive a Host prefix-cache config from existing Host KV profiles.""" + + group_specs, required_pages = _derive_group_specs_and_page_count( + model_name=model_name, + host_kv_cache_size_bytes=host_kv_cache_size_bytes, + ) + return build_prefix_cache_runtime_config_from_specs( + model_name=model_name, + kv_dtype=kv_dtype, + host_kv_pages_per_required_group=required_pages, + group_specs=group_specs, + debug_stats=debug_stats, + ) + + +def build_prefix_cache_runtime_config_from_specs( + *, + model_name: str, + kv_dtype: str, + host_kv_pages_per_required_group: int, + group_specs: Sequence[PrefixKVGroupSpec], + debug_stats: bool = False, +) -> PrefixCacheRuntimeConfig: + """Build a runtime config from already-derived logical KV groups.""" + + specs = tuple(group_specs) + if not specs: + raise ValueError("prefix cache requires at least one KV group") + required_specs = tuple(spec for spec in specs if spec.required_for_reuse) + if not required_specs: + raise ValueError("prefix cache requires at least one required KV group") + + hash_block_tokens = _gcd(spec.raw_page_tokens for spec in required_specs) + publish_boundary_tokens = _lcm( + spec.raw_page_tokens for spec in required_specs + ) + if hash_block_tokens <= 0 or publish_boundary_tokens <= 0: + raise ValueError("prefix cache token boundaries must be positive") + + pages_per_group = int(host_kv_pages_per_required_group) + if pages_per_group <= 0: + raise ValueError("host_kv_pages_per_required_group must be positive") + + max_nodes = max(1024, pages_per_group + 1) + max_group_entries = max_nodes * len(specs) + max_page_handles = _derive_page_handle_capacity( + max_nodes=max_nodes, + pages_per_group=pages_per_group, + group_count=len(specs), + ) + max_attachments = max(1024, max_nodes // 4) + + return PrefixCacheRuntimeConfig( + shm_name=derive_prefix_cache_shm_name(model_name), + namespace_digest=build_prefix_cache_namespace_digest( + model_name=model_name, + kv_dtype=kv_dtype, + group_specs=specs, + ), + group_specs=specs, + hash_block_tokens=hash_block_tokens, + publish_boundary_tokens=publish_boundary_tokens, + max_nodes=max_nodes, + max_group_entries=max_group_entries, + max_page_handles=max_page_handles, + max_attachments=max_attachments, + debug_stats=debug_stats, + ) + + +def create_host_prefix_cache_coordinator( + *, + core_engine_module, + runtime_config: PrefixCacheRuntimeConfig, + create_region: bool, +): + coordinator = core_engine_module.HostPrefixCacheCoordinator( + runtime_config.to_core_config(core_engine_module) + ) + coordinator.initialize(bool(create_region)) + return coordinator + + +def derive_prefix_cache_shm_name(model_name: str) -> str: + normalized = re.sub(r"[^a-zA-Z0-9]+", "_", model_name).strip("_").lower() + normalized = normalized[:64] or "model" + digest = hashlib.blake2b(model_name.encode("utf-8"), digest_size=4) + suffix = int.from_bytes(digest.digest(), "little") + return f"batchgen_prefix_cache_{normalized}_{suffix:08x}" + + +def build_prefix_cache_namespace_digest( + *, + model_name: str, + kv_dtype: str, + group_specs: Sequence[PrefixKVGroupSpec], +) -> tuple[int, int, int, int]: + hasher = hashlib.blake2b(digest_size=32) + hasher.update(model_name.strip().lower().encode("utf-8")) + hasher.update(b"\0") + hasher.update(kv_dtype.strip().lower().encode("utf-8")) + for spec in sorted(group_specs, key=lambda item: int(item.group_id)): + hasher.update(b"\0") + hasher.update(int(spec.group_id).to_bytes(4, "little")) + hasher.update(spec.semantic.value.encode("ascii")) + hasher.update(b"\0") + hasher.update(int(spec.required_for_reuse).to_bytes(1, "little")) + hasher.update(int(spec.raw_page_tokens).to_bytes(4, "little")) + hasher.update(int(spec.compression_ratio).to_bytes(4, "little")) + digest = hasher.digest() + return tuple( + int.from_bytes(digest[offset : offset + 8], "little") + for offset in range(0, 32, 8) + ) + + +def _derive_group_specs_and_page_count( + *, model_name: str, host_kv_cache_size_bytes: int +) -> tuple[tuple[PrefixKVGroupSpec, ...], int]: + from batchgen.kv_cache.host_kv_mananger_config import ( + resolve_host_kv_group_profiles, + ) + + group_profiles = resolve_host_kv_group_profiles(model_name) + specs = tuple( + PrefixKVGroupSpec( + group_id=profile.group_id, + semantic=_semantic_from_group_profile(profile), + required_for_reuse=profile.required_for_reuse, + raw_page_tokens=profile.raw_page_tokens, + compression_ratio=profile.compression_ratio, + ) + for profile in group_profiles + ) + required_profiles = tuple( + profile for profile in group_profiles if profile.required_for_reuse + ) + if not required_profiles: + raise ValueError("prefix cache requires at least one required KV group") + + publish_boundary_tokens = _lcm( + profile.raw_page_tokens for profile in required_profiles + ) + bytes_per_publish_boundary = sum( + profile.bytes_per_page() + * profile.num_layers + * (publish_boundary_tokens // profile.raw_page_tokens) + for profile in required_profiles + ) + publish_units = int(host_kv_cache_size_bytes) // bytes_per_publish_boundary + if publish_units <= 0: + raise ValueError("host KV cache is too small for prefix cache") + + return specs, publish_units + + +def _semantic_from_group_profile(profile) -> PrefixKVGroupSemantic: + try: + return PrefixKVGroupSemantic(profile.semantic) + except ValueError as exc: + raise ValueError( + f"unsupported prefix KV group semantic {profile.semantic!r}" + ) from exc + + +def _to_core_group_spec(core_engine_module, spec: PrefixKVGroupSpec): + core_spec = core_engine_module.HostKVGroupSpec() + core_spec.group_id = int(spec.group_id) + core_spec.semantic = _to_core_semantic(core_engine_module, spec.semantic) + core_spec.required_for_reuse = bool(spec.required_for_reuse) + core_spec.raw_page_tokens = int(spec.raw_page_tokens) + core_spec.compression_ratio = int(spec.compression_ratio) + return core_spec + + +def _to_core_semantic(core_engine_module, semantic: PrefixKVGroupSemantic): + enum_cls = core_engine_module.HostKVGroupSemantic + return getattr(enum_cls, semantic.name) + + +def _gcd(values: Iterable[int]) -> int: + result = 0 + for value in values: + result = int(value) if result == 0 else math.gcd(result, int(value)) + return result + + +def _lcm(values: Iterable[int]) -> int: + result = 1 + for value in values: + result = math.lcm(result, int(value)) + return result + + +def _derive_page_handle_capacity( + *, max_nodes: int, pages_per_group: int, group_count: int +) -> int: + # The current C++ entry stores enough page handles to materialize a node, + # so long prompts need more than one handle per node. Use a derived, + # bounded estimate instead of a user-tunable knob. + average_pages_per_node = max( + 16, min(512, int(math.sqrt(max(1, pages_per_group)))) + ) + return max_nodes * max(1, group_count) * average_pages_per_node diff --git a/batchgen/prefix_reuse/eviction.py b/batchgen/prefix_reuse/eviction.py new file mode 100644 index 000000000..6dc52c4d6 --- /dev/null +++ b/batchgen/prefix_reuse/eviction.py @@ -0,0 +1,184 @@ +"""Eviction helpers for Host prefix-cache integration. + +The coordinator owns prefix metadata and chooses eviction victims. Host KV +worker views own physical Host pages, so released page handles must be routed +back to the worker view for the matching prefix group. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Iterable, Iterator + +from batchgen.prefix_reuse.commit import PrefixCommitRequest + + +_CAPACITY_ERROR_MARKERS = ( + "Host prefix cache node table is full", + "Host prefix cache group entry table is full", + "Host prefix cache page handle arena is full", +) + + +@dataclass(frozen=True) +class PrefixCommitRetryResult: + commit_result: object + eviction_result: object | None = None + released_pages_by_group: dict[int, int] | None = None + + +@dataclass(frozen=True) +class PrefixAllocationEvictionResult: + eviction_result: object | None + released_pages_by_group: dict[int, int] + + +def commit_prefix_pages_with_capacity_retry( + *, + request: PrefixCommitRequest, + coordinator: object, + worker_views_by_group: Mapping[int, object], + max_scan_nodes: int = 0, +) -> PrefixCommitRetryResult: + """Commit prefix pages, evicting and retrying once on metadata pressure.""" + + try: + return PrefixCommitRetryResult(commit_result=request.commit(coordinator)) + except RuntimeError as exc: + if not _is_capacity_error(exc): + raise + eviction_result = _evict_for_request_capacity( + request=request, + coordinator=coordinator, + max_scan_nodes=max_scan_nodes, + ) + released = release_evicted_prefix_pages( + eviction_result=eviction_result, + worker_views_by_group=worker_views_by_group, + ) + return PrefixCommitRetryResult( + commit_result=request.commit(coordinator), + eviction_result=eviction_result, + released_pages_by_group=released, + ) + + +def evict_prefix_pages_for_host_allocation( + *, + core_engine_module: object, + coordinator: object, + worker_views_by_group: Mapping[int, object], + page_deficit_by_group: Mapping[int, int], + max_scan_nodes: int = 0, +) -> PrefixAllocationEvictionResult: + """Evict common prefix nodes until enough physical Host pages are released. + + The input is per-group pressure, but the coordinator still evicts whole + prefix nodes. The returned pages are filtered by the coordinator so only + pages no longer referenced by any resident prefix node are released. + """ + + requirements = [] + for group_id, deficit in sorted(page_deficit_by_group.items()): + deficit = int(deficit) + if deficit <= 0: + continue + requirement = core_engine_module.GroupPageRequirement() + requirement.group_id = int(group_id) + requirement.min_pages = deficit + requirements.append(requirement) + if not requirements: + return PrefixAllocationEvictionResult( + eviction_result=None, + released_pages_by_group={}, + ) + + eviction_result = coordinator.evict_until_releasable_pages( + requirements, + int(max_scan_nodes), + ) + released = release_evicted_prefix_pages( + eviction_result=eviction_result, + worker_views_by_group=worker_views_by_group, + ) + + missing = { + int(requirement.group_id): int(requirement.min_pages) + - int(released.get(int(requirement.group_id), 0)) + for requirement in requirements + if int(released.get(int(requirement.group_id), 0)) + < int(requirement.min_pages) + } + if missing: + raise RuntimeError( + "prefix cache eviction could not release enough Host KV pages " + f"for allocation: missing={missing}, released={released}, " + f"evicted_nodes={int(eviction_result.evicted_nodes)}, " + f"protected_nodes={int(eviction_result.protected_nodes)}" + ) + + return PrefixAllocationEvictionResult( + eviction_result=eviction_result, + released_pages_by_group=released, + ) + + +def release_evicted_prefix_pages( + *, + eviction_result: object, + worker_views_by_group: Mapping[int, object], +) -> dict[int, int]: + """Release coordinator-evicted physical Host pages by KV group.""" + + pages_by_group: dict[int, list[int]] = {} + seen_by_group: dict[int, set[int]] = {} + for group_pages in eviction_result.evicted_group_pages: + group_id = int(group_pages.group_id) + seen = seen_by_group.setdefault(group_id, set()) + pages = pages_by_group.setdefault(group_id, []) + for page_id in _page_ids(group_pages.pages): + if page_id in seen: + continue + seen.add(page_id) + pages.append(page_id) + + released: dict[int, int] = {} + for group_id, page_ids in pages_by_group.items(): + worker_view = worker_views_by_group.get(group_id) + if worker_view is None: + raise RuntimeError( + f"missing Host KV worker view for evicted prefix group {group_id}" + ) + if not page_ids: + continue + worker_view.release_resident_pages(page_ids) + released[group_id] = len(page_ids) + return released + + +def _evict_for_request_capacity( + *, + request: PrefixCommitRequest, + coordinator: object, + max_scan_nodes: int, +) -> object: + node_count, group_entry_count, page_handle_count = ( + request.capacity_requirements() + ) + return coordinator.evict_until_free( + int(node_count), + int(group_entry_count), + int(page_handle_count), + int(max_scan_nodes), + ) + + +def _is_capacity_error(exc: RuntimeError) -> bool: + message = str(exc) + return any(marker in message for marker in _CAPACITY_ERROR_MARKERS) + + +def _page_ids(pages: Iterable[object]) -> Iterator[int]: + for page in pages: + yield int(getattr(page, "page_id", page)) diff --git a/batchgen/prefix_reuse/materialization.py b/batchgen/prefix_reuse/materialization.py new file mode 100644 index 000000000..c300a7d9f --- /dev/null +++ b/batchgen/prefix_reuse/materialization.py @@ -0,0 +1,732 @@ +"""GPU paged materialization helpers for prefix-reuse prefill.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import Optional, Protocol, Sequence + +import torch + +from batchgen.prefix_reuse.prefill import effective_prefix_shared_tokens + + +class _AsyncTask(Protocol): + def wait_for_layer(self, layer_idx: int) -> None: ... + + def wait(self) -> None: ... + + +class _PrefixCacheCoordinator(Protocol): + def begin_attachment_load(self, attachment_handle: int) -> None: ... + + def end_attachment_load(self, attachment_handle: int) -> None: ... + + +@dataclass(frozen=True) +class PrefixMaterializationSequence: + """Host prefix pages needed by one target GPU sequence.""" + + sequence_id: int + prefix_tokens: int + suffix_tokens: int + host_pages: Sequence[int | object] + attachment_handle: int = 0 + + +@dataclass +class SingleGroupPrefixMaterialization: + """Single KV-group materialization view consumed by current adapters.""" + + manager: object | None + append_plan: object | None + load_task: Optional[_AsyncTask] = None + backend_state: dict[str, object] = field(default_factory=dict) + _loaded: bool = False + _closed: bool = False + + def wait_for_layer(self, layer_idx: int) -> None: + if self._closed: + raise RuntimeError("prefix materialization is already closed") + if self._loaded or self.load_task is None: + return + self.load_task.wait_for_layer(int(layer_idx)) + + def wait(self) -> None: + if self._loaded: + return + if self.load_task is not None: + self.load_task.wait() + self._loaded = True + + def close(self, *, empty_cuda_cache: bool = False) -> None: + """Wait for outstanding loads and release GPU materialization buffers.""" + + if self._closed: + return + manager = self.manager + try: + self.wait() + finally: + self.manager = None + self.append_plan = None + self.load_task = None + self.backend_state.clear() + self._closed = True + if manager is not None: + manager.destroy(empty_cuda_cache=empty_cuda_cache) + + def finish_layer(self, layer_idx: int) -> None: + """Notify materialization that a logical layer no longer needs GPU KV.""" + + del layer_idx + + +@dataclass +class RollingSingleGroupPrefixMaterialization(SingleGroupPrefixMaterialization): + """Two-slot logical-layer materialization for prefix-hit prefill. + + The manager owns a small physical layer window and maps logical layers onto + those slots. Prefix pages are loaded layer-by-layer from Host KV, so prefill + does not retain full-model GPU KV for every layer. + """ + + host_worker_view: object | None = None + host_page_ids: torch.Tensor | None = None + active_page_counts: torch.Tensor | None = None + host_page_tokens: int | None = None + logical_layer_count: int = 0 + prefix_cache_coordinator: Optional[_PrefixCacheCoordinator] = None + attachment_handles: Sequence[int] = () + _scheduled_tasks: dict[int, _AsyncTask] = field(default_factory=dict) + _begun_handles: list[int] = field(default_factory=list) + _attachments_released: bool = False + + def start(self) -> None: + """Begin attachment protection and prefetch the first two layers.""" + + if self.host_page_ids is None or self.active_page_counts is None: + return + self._begin_attachment_loads() + try: + self._schedule_layer(0) + self._schedule_layer(1) + except Exception: + self.wait() + raise + + def wait_for_layer(self, layer_idx: int) -> None: + if self._closed: + raise RuntimeError("prefix materialization is already closed") + layer_idx = int(layer_idx) + task = self._scheduled_tasks.get(layer_idx) + if task is None: + raise RuntimeError( + "rolling prefix materialization layer was not scheduled; " + f"layer={layer_idx}. The previous prefill offload may not have " + "been retired before reusing the physical GPU KV slot." + ) + task.wait_for_layer(layer_idx) + + def finish_layer(self, layer_idx: int) -> None: + if self._closed: + return + # Reuse the just-consumed physical slot for the next non-resident + # logical layer. Callers invoke this after the layer's attention output + # and suffix offload have consumed the temporary GPU KV. + self._schedule_layer(int(layer_idx) + 2) + + def wait(self) -> None: + if self._loaded: + return + try: + for task in self._scheduled_tasks.values(): + task.wait() + finally: + self._release_attachment_loads() + self._loaded = True + + def close(self, *, empty_cuda_cache: bool = False) -> None: + if self._closed: + return + manager = self.manager + try: + self.wait() + finally: + self.manager = None + self.append_plan = None + self.load_task = None + self.host_worker_view = None + self.host_page_ids = None + self.active_page_counts = None + self._scheduled_tasks.clear() + self.backend_state.clear() + self._closed = True + if manager is not None: + manager.destroy(empty_cuda_cache=empty_cuda_cache) + + def _begin_attachment_loads(self) -> None: + if self._begun_handles or self.prefix_cache_coordinator is None: + return + for handle in self.attachment_handles: + self.prefix_cache_coordinator.begin_attachment_load(int(handle)) + self._begun_handles.append(int(handle)) + + def _release_attachment_loads(self) -> None: + if self._attachments_released: + return + coordinator = self.prefix_cache_coordinator + if coordinator is not None: + for handle in reversed(self._begun_handles): + coordinator.end_attachment_load(int(handle)) + self._begun_handles.clear() + self._attachments_released = True + + def _schedule_layer(self, layer_idx: int) -> None: + if ( + layer_idx < 0 + or layer_idx >= int(self.logical_layer_count) + or layer_idx in self._scheduled_tasks + or self.host_page_ids is None + or self.active_page_counts is None + ): + return + if self.manager is None or self.host_worker_view is None: + raise RuntimeError("rolling prefix materialization is not active") + + physical_layer = int(self.manager.resolve_physical_layer(layer_idx)) + selected_rows = torch.tensor([physical_layer], dtype=torch.int64) + k_ptrs, v_ptrs = self.manager.get_padded_3d_page_pointers() + selected_k_ptrs = k_ptrs.index_select(0, selected_rows).contiguous() + selected_v_ptrs = ( + None + if v_ptrs is None + else v_ptrs.index_select(0, selected_rows).contiguous() + ) + selected_k_ptrs, selected_v_ptrs = _expand_device_ptrs_for_host_pages( + gpu_manager=self.manager, + k_device_ptrs=selected_k_ptrs, + v_device_ptrs=selected_v_ptrs, + active_page_counts=self.active_page_counts, + host_page_tokens=self.host_page_tokens, + ) + logical_layers = torch.tensor([layer_idx], dtype=torch.int64) + task = self.host_worker_view.async_load_prefix_layers_to_device( + host_page_ids=self.host_page_ids, + active_page_counts=self.active_page_counts, + logical_layer_ids=logical_layers, + k_device_ptrs=selected_k_ptrs, + v_device_ptrs=selected_v_ptrs, + ) + self._scheduled_tasks[layer_idx] = task + + +@dataclass +class PrefixMaterializationBundle: + """Materialized prefix pages keyed by logical prefix-cache group id.""" + + by_group_id: dict[int, SingleGroupPrefixMaterialization] + + def get(self, group_id: int) -> Optional[SingleGroupPrefixMaterialization]: + return self.by_group_id.get(int(group_id)) + + def require( + self, group_id: int, *, consumer: str + ) -> SingleGroupPrefixMaterialization: + materialization = self.get(group_id) + if materialization is None: + raise RuntimeError( + f"{consumer} requires prefix materialization group {group_id}" + ) + return materialization + + def wait_for_layer(self, layer_idx: int) -> None: + for materialization in self.by_group_id.values(): + materialization.wait_for_layer(layer_idx) + + def wait(self) -> None: + for materialization in self.by_group_id.values(): + materialization.wait() + + def finish_layer(self, layer_idx: int) -> None: + for materialization in self.by_group_id.values(): + materialization.finish_layer(layer_idx) + + def close(self, *, empty_cuda_cache: bool = False) -> None: + for materialization in self.by_group_id.values(): + materialization.close(empty_cuda_cache=empty_cuda_cache) + + +def get_prefix_materialization_for_group( + materialization: object | None, + *, + group_id: int, + consumer: str, +) -> SingleGroupPrefixMaterialization | None: + """Return the materialization consumed by one attention backend.""" + + if materialization is None: + return None + if isinstance(materialization, PrefixMaterializationBundle): + return materialization.require(group_id, consumer=consumer) + raise RuntimeError( + f"{consumer} requires PrefixMaterializationBundle, " + f"got {type(materialization).__name__}" + ) + + +class _AttachmentLoadTask: + def __init__( + self, + *, + load_task: _AsyncTask, + coordinator: _PrefixCacheCoordinator, + attachment_handles: Sequence[int], + ) -> None: + self._load_task = load_task + self._coordinator = coordinator + self._attachment_handles = tuple( + int(handle) for handle in attachment_handles + ) + self._done = False + + def wait(self) -> None: + if self._done: + return + try: + self._load_task.wait() + finally: + for handle in reversed(self._attachment_handles): + self._coordinator.end_attachment_load(handle) + self._done = True + + def wait_for_layer(self, layer_idx: int) -> None: + if self._done: + return + self._load_task.wait_for_layer(int(layer_idx)) + + +def materialize_single_group_prefix_pages( + *, + gpu_manager: object, + host_worker_view: object, + sequences: Sequence[PrefixMaterializationSequence], + raw_page_tokens: int | None = None, + prefix_cache_coordinator: Optional[_PrefixCacheCoordinator] = None, + rolling_logical_layer_count: int | None = None, +) -> SingleGroupPrefixMaterialization: + """Materialize Host prefix pages into target GPU paged KV slots. + + This helper is intentionally below the Host prefix-cache coordinator. The + caller provides already attached/pinned Host page handles and target + sequence ids; this function only allocates GPU pages, starts the page-id + based Host->GPU copy, and prepares suffix append metadata. + """ + + if not sequences: + raise ValueError( + "prefix materialization requires at least one sequence" + ) + + sequence_ids = [int(item.sequence_id) for item in sequences] + prefix_lens = [int(item.prefix_tokens) for item in sequences] + suffix_lens = [int(item.suffix_tokens) for item in sequences] + full_lens = [ + prefix + suffix for prefix, suffix in zip(prefix_lens, suffix_lens) + ] + for seq_id, prefix_len, suffix_len, full_len in zip( + sequence_ids, prefix_lens, suffix_lens, full_lens + ): + if prefix_len < 0 or suffix_len < 0: + raise ValueError( + "prefix/suffix lengths must be non-negative for sequence " + f"{seq_id}: prefix={prefix_len}, suffix={suffix_len}" + ) + if full_len <= 0: + raise ValueError( + f"full sequence length must be positive for sequence {seq_id}" + ) + + host_page_tokens = int( + raw_page_tokens + if raw_page_tokens is not None + else gpu_manager.config.page_size_tokens + ) + if host_page_tokens <= 0: + raise ValueError("raw_page_tokens must be positive") + prefix_page_counts = [ + int(math.ceil(prefix_len / host_page_tokens)) if prefix_len > 0 else 0 + for prefix_len in prefix_lens + ] + has_prefix_pages = any(count > 0 for count in prefix_page_counts) + host_page_ids = None + active_page_counts = None + if has_prefix_pages: + host_page_ids = _build_host_page_id_tensor( + sequences, + prefix_page_counts=prefix_page_counts, + ) + active_page_counts = torch.tensor(prefix_page_counts, dtype=torch.int64) + + gpu_manager.allocate_pages_for_sequences(sequence_ids, full_lens) + gpu_manager.rebuild_page_table(sequence_ids) + k_ptrs, v_ptrs = gpu_manager.get_padded_3d_page_pointers() + copy_k_ptrs, copy_v_ptrs = _expand_device_ptrs_for_host_pages( + gpu_manager=gpu_manager, + k_device_ptrs=k_ptrs, + v_device_ptrs=v_ptrs, + active_page_counts=active_page_counts, + host_page_tokens=host_page_tokens, + ) + append_plan = gpu_manager.prepare_prefill_suffix_append( + sequence_ids=sequence_ids, + prefix_lens=prefix_lens, + suffix_lens=suffix_lens, + rebuild_page_table=False, + ) + + if rolling_logical_layer_count is not None: + attachment_handles = _attachment_handles_for_load( + sequences, + prefix_page_counts, + ) + if attachment_handles and prefix_cache_coordinator is None: + raise ValueError( + "prefix materialization sequences with attachment handles " + "require prefix_cache_coordinator" + ) + materialization = RollingSingleGroupPrefixMaterialization( + manager=gpu_manager, + append_plan=append_plan, + host_worker_view=host_worker_view, + host_page_ids=host_page_ids, + active_page_counts=active_page_counts, + host_page_tokens=host_page_tokens, + logical_layer_count=int(rolling_logical_layer_count), + prefix_cache_coordinator=prefix_cache_coordinator, + attachment_handles=tuple(attachment_handles), + ) + try: + materialization.start() + except Exception: + materialization.close() + raise + return materialization + + load_task = None + if has_prefix_pages: + attachment_handles = _attachment_handles_for_load( + sequences, + prefix_page_counts, + ) + if attachment_handles and prefix_cache_coordinator is None: + raise ValueError( + "prefix materialization sequences with attachment handles " + "require prefix_cache_coordinator" + ) + + begun_handles: list[int] = [] + try: + if prefix_cache_coordinator is not None: + for handle in attachment_handles: + prefix_cache_coordinator.begin_attachment_load(handle) + begun_handles.append(handle) + + load_task = host_worker_view.async_load_prefix_pages_to_device( + host_page_ids=host_page_ids, + active_page_counts=active_page_counts, + k_device_ptrs=copy_k_ptrs, + v_device_ptrs=copy_v_ptrs, + ) + except Exception: + if prefix_cache_coordinator is not None: + for handle in reversed(begun_handles): + prefix_cache_coordinator.end_attachment_load(handle) + raise + + if prefix_cache_coordinator is not None and begun_handles: + load_task = _AttachmentLoadTask( + load_task=load_task, + coordinator=prefix_cache_coordinator, + attachment_handles=begun_handles, + ) + + return SingleGroupPrefixMaterialization( + manager=gpu_manager, + append_plan=append_plan, + load_task=load_task, + ) + + +def materialize_single_group_lookup_results( + *, + gpu_manager: object, + host_worker_view: object, + lookup_results: Sequence[object], + sequence_ids: Sequence[int], + prompt_lengths: Sequence[int], + group_id: int, + prefix_shared_tokens: Sequence[int] | None = None, + raw_page_tokens: int | None = None, + prefix_cache_coordinator: Optional[_PrefixCacheCoordinator] = None, + rolling_logical_layer_count: int | None = None, +) -> SingleGroupPrefixMaterialization: + """Materialize a batch of C++ HostPrefixCache lookup results. + + The Host prefix-cache coordinator owns lookup, attachment lifetime, and + eviction. This function is only the compute-path producer: it converts + attached lookup results for one KV group into GPU paged KV materialization. + """ + + count = len(lookup_results) + if len(sequence_ids) != count or len(prompt_lengths) != count: + raise ValueError( + "lookup_results, sequence_ids, and prompt_lengths differ" + ) + if prefix_shared_tokens is not None and len(prefix_shared_tokens) != count: + raise ValueError( + "prefix_shared_tokens length differs from lookup_results" + ) + + sequences: list[PrefixMaterializationSequence] = [] + for idx, (result, sequence_id, prompt_length) in enumerate(zip( + lookup_results, + sequence_ids, + prompt_lengths, + )): + prompt_len = int(prompt_length) + raw_cached_tokens = int(result.common_cached_tokens) + if prompt_len <= 0: + raise ValueError( + f"prompt length must be positive for sequence {sequence_id}" + ) + if raw_cached_tokens < 0 or raw_cached_tokens > prompt_len: + raise ValueError( + "lookup cached token count must be within prompt length for " + f"sequence {sequence_id}: cached={raw_cached_tokens}, " + f"prompt={prompt_len}" + ) + if prefix_shared_tokens is None: + cached_tokens = effective_prefix_shared_tokens( + raw_cached_tokens=raw_cached_tokens, + prompt_length=prompt_len, + ) + else: + cached_tokens = int(prefix_shared_tokens[idx]) + if cached_tokens < 0 or cached_tokens >= prompt_len: + raise ValueError( + "effective cached token count must be within compute bounds " + f"for sequence {sequence_id}: cached={cached_tokens}, " + f"prompt={prompt_len}" + ) + span_pages = [] + attachment_handle = int(result.attachment_handle) + if cached_tokens > 0: + if attachment_handle == 0: + raise ValueError( + "lookup result with cached prefix must have non-zero " + f"attachment_handle for sequence {sequence_id}" + ) + span = _find_group_span(result, group_id=int(group_id)) + span_raw_end = int(span.raw_end_token) + if span_raw_end < cached_tokens: + raise ValueError( + "single-group prefix materialization requires lookup span " + "to cover the effective cached token boundary for sequence " + f"{sequence_id}: span={span_raw_end}, " + f"effective_cached={cached_tokens}" + ) + span_pages = list(span.pages) + + sequences.append( + PrefixMaterializationSequence( + sequence_id=int(sequence_id), + prefix_tokens=cached_tokens, + suffix_tokens=prompt_len - cached_tokens, + host_pages=span_pages, + attachment_handle=attachment_handle, + ) + ) + + return materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=host_worker_view, + sequences=sequences, + raw_page_tokens=raw_page_tokens, + prefix_cache_coordinator=prefix_cache_coordinator, + rolling_logical_layer_count=rolling_logical_layer_count, + ) + + +def _build_host_page_id_tensor( + sequences: Sequence[PrefixMaterializationSequence], + *, + prefix_page_counts: Sequence[int], +) -> torch.Tensor: + max_pages = max(int(count) for count in prefix_page_counts) + rows: list[list[int]] = [] + for item, page_count in zip(sequences, prefix_page_counts): + pages = [_host_page_id(handle) for handle in item.host_pages] + if len(pages) < int(page_count): + raise ValueError( + "host prefix page list is shorter than required for sequence " + f"{item.sequence_id}: need {page_count}, got {len(pages)}" + ) + row = pages[: int(page_count)] + row.extend([0] * (max_pages - len(row))) + rows.append(row) + return torch.tensor(rows, dtype=torch.int64) + + +def _expand_device_ptrs_for_host_pages( + *, + gpu_manager: object, + k_device_ptrs: torch.Tensor, + v_device_ptrs: torch.Tensor | None, + active_page_counts: torch.Tensor | None, + host_page_tokens: int | None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Map host-page copy slots onto potentially larger GPU pages. + + Host prefix pages are indexed by the Host KV group's raw page size. Some + GPU kernels impose a larger paged-cache block size; for example FA3 paged + ``flash_attn_with_kvcache`` requires a 256-token GPU page. In that case + each Host page is copied into a subrange of the larger GPU page by adding a + byte offset to the destination page pointer. The C++ copy path remains + asynchronous and still copies one Host page per entry. + """ + + gpu_page_tokens = int(gpu_manager.config.page_size_tokens) + host_tokens = int( + host_page_tokens if host_page_tokens is not None else gpu_page_tokens + ) + if ( + host_tokens == gpu_page_tokens + or host_tokens > gpu_page_tokens + or active_page_counts is None + ): + return k_device_ptrs, v_device_ptrs + if host_tokens <= 0 or gpu_page_tokens <= 0: + raise ValueError("host and GPU page sizes must be positive") + if gpu_page_tokens % host_tokens != 0: + raise ValueError( + "GPU page size must be a multiple of Host page size for prefix " + f"materialization, got gpu={gpu_page_tokens}, host={host_tokens}" + ) + + k_page_bytes = _host_page_bytes( + gpu_manager=gpu_manager, + host_page_tokens=host_tokens, + is_value=False, + ) + expanded_k = _expand_pointer_tensor_for_host_pages( + k_device_ptrs, + active_page_counts=active_page_counts, + host_page_bytes=k_page_bytes, + host_pages_per_gpu_page=gpu_page_tokens // host_tokens, + ) + + expanded_v = None + if v_device_ptrs is not None: + v_page_bytes = _host_page_bytes( + gpu_manager=gpu_manager, + host_page_tokens=host_tokens, + is_value=True, + ) + expanded_v = _expand_pointer_tensor_for_host_pages( + v_device_ptrs, + active_page_counts=active_page_counts, + host_page_bytes=v_page_bytes, + host_pages_per_gpu_page=gpu_page_tokens // host_tokens, + ) + + return expanded_k, expanded_v + + +def _host_page_bytes( + *, + gpu_manager: object, + host_page_tokens: int, + is_value: bool, +) -> int: + config = gpu_manager.config + if is_value: + heads = int(config.num_v_heads) + head_dim = int(config.v_head_dim) + else: + heads = int(config.num_k_heads) + head_dim = int(config.k_head_dim) + element_size = torch.empty((), dtype=config.kv_dtype).element_size() + return int(host_page_tokens) * heads * head_dim * int(element_size) + + +def _expand_pointer_tensor_for_host_pages( + pointer_tensor: torch.Tensor, + *, + active_page_counts: torch.Tensor, + host_page_bytes: int, + host_pages_per_gpu_page: int, +) -> torch.Tensor: + if int(active_page_counts.numel()) == 0: + return pointer_tensor[:, :, :0].contiguous() + max_host_pages = int(active_page_counts.max().item()) + if max_host_pages == 0: + return pointer_tensor[:, :, :0].contiguous() + + host_slots = torch.arange(max_host_pages, dtype=torch.long) + gpu_slots = torch.div( + host_slots, + int(host_pages_per_gpu_page), + rounding_mode="floor", + ) + if int(gpu_slots[-1].item()) >= int(pointer_tensor.shape[2]): + raise ValueError( + "GPU page pointer tensor is too small for Host prefix pages: " + f"max_host_pages={max_host_pages}, " + f"host_pages_per_gpu_page={host_pages_per_gpu_page}, " + f"gpu_pointer_pages={pointer_tensor.shape[2]}" + ) + offsets = ( + torch.remainder(host_slots, int(host_pages_per_gpu_page)).to( + dtype=torch.int64 + ) + * int(host_page_bytes) + ) + expanded = pointer_tensor.index_select(2, gpu_slots).contiguous() + return expanded + offsets.view(1, 1, -1) + + +def _find_group_span(result: object, *, group_id: int) -> object: + spans = result.materialization_spans + if spans is None: + raise TypeError("lookup result must expose materialization_spans") + for span in spans: + if int(span.group_id) == int(group_id): + return span + raise ValueError( + f"lookup result has no materialization span for group {group_id}" + ) + + +def _host_page_id(handle: int | object) -> int: + if isinstance(handle, int): + return int(handle) + page_id = getattr(handle, "page_id", None) + if page_id is None: + raise TypeError("host page handle must be an int or expose page_id") + return int(page_id) + + +def _attachment_handles_for_load( + sequences: Sequence[PrefixMaterializationSequence], + prefix_page_counts: Sequence[int], +) -> list[int]: + handles: list[int] = [] + seen: set[int] = set() + for item, page_count in zip(sequences, prefix_page_counts): + handle = int(item.attachment_handle) + if int(page_count) <= 0 or handle == 0 or handle in seen: + continue + seen.add(handle) + handles.append(handle) + return handles diff --git a/batchgen/prefix_reuse/prefill.py b/batchgen/prefix_reuse/prefill.py new file mode 100644 index 000000000..9db941bd2 --- /dev/null +++ b/batchgen/prefix_reuse/prefill.py @@ -0,0 +1,152 @@ +"""Host prefix-cache lookup helpers for prepacked prefill.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + +import torch + +from batchgen.prefill.prefix_reuse import ( + PrefixReusePrefillPlan, + build_prefix_reuse_prefill_plan, +) + + +@dataclass(frozen=True) +class PrefixCachePrefillLookup: + lookup_results: tuple[object, ...] + prefix_shared_tokens: tuple[int, ...] + + @property + def has_hit(self) -> bool: + return any(tokens > 0 for tokens in self.prefix_shared_tokens) + + +@dataclass(frozen=True) +class PrefixCachePrefillEstimate: + prefix_shared_tokens: tuple[int, ...] + + @property + def has_hit(self) -> bool: + return any(tokens > 0 for tokens in self.prefix_shared_tokens) + + +@dataclass(frozen=True) +class PrefixCachePrefillInputs: + plan: PrefixReusePrefillPlan + input_ids_list: list[torch.Tensor] + attention_mask_list: list[torch.Tensor] + + +def effective_prefix_shared_tokens( + *, raw_cached_tokens: int, prompt_length: int +) -> int: + """Normalize coordinator lookup tokens to the compute-path semantic. + + The coordinator reports raw page-cache hits. The prefill compute path always + runs at least one query token, so an exact full hit becomes a one-token + extend with ``prompt_length - 1`` cached tokens. After this boundary, + callers should propagate only the normalized value. + """ + + prompt_len = int(prompt_length) + cached = int(raw_cached_tokens) + if prompt_len <= 0: + raise ValueError( + f"prompt_length must be positive for prefix lookup, got {prompt_len}" + ) + if cached < 0 or cached > prompt_len: + raise ValueError( + "raw_cached_tokens must be within prompt length: " + f"cached={cached}, prompt_length={prompt_len}" + ) + if cached == prompt_len: + return max(prompt_len - 1, 0) + return cached + + +def lookup_prefix_cache_for_prefill( + *, + coordinator: object, + namespace_digest: Sequence[int], + prompt_token_ids: Sequence[Sequence[int]], +) -> PrefixCachePrefillLookup: + """Lookup reusable prompt prefixes for a local prefill batch.""" + + lookup_results = [] + prefix_shared_tokens = [] + for token_ids in prompt_token_ids: + result = coordinator.lookup_and_attach( + list(namespace_digest), + [int(token_id) for token_id in token_ids], + ) + lookup_results.append(result) + prefix_shared_tokens.append( + effective_prefix_shared_tokens( + raw_cached_tokens=int(result.common_cached_tokens), + prompt_length=len(token_ids), + ) + ) + + return PrefixCachePrefillLookup( + lookup_results=tuple(lookup_results), + prefix_shared_tokens=tuple(prefix_shared_tokens), + ) + + +def estimate_prefix_cache_for_prefill( + *, + coordinator: object, + namespace_digest: Sequence[int], + prompt_token_ids: Sequence[Sequence[int]], +) -> PrefixCachePrefillEstimate: + """Estimate reusable prefixes without attaching or pinning cache entries.""" + + prefix_shared_tokens = [] + for token_ids in prompt_token_ids: + result = coordinator.estimate_lookup( + list(namespace_digest), + [int(token_id) for token_id in token_ids], + ) + prefix_shared_tokens.append( + effective_prefix_shared_tokens( + raw_cached_tokens=int(result.common_cached_tokens), + prompt_length=len(token_ids), + ) + ) + + return PrefixCachePrefillEstimate( + prefix_shared_tokens=tuple(prefix_shared_tokens), + ) + + +def build_prefix_cache_prefill_inputs( + *, + local_indices: Sequence[int], + sequence_ids: Sequence[int], + input_ids: Sequence[torch.Tensor], + prompt_lengths: Sequence[int], + lookup: PrefixCachePrefillLookup, +) -> PrefixCachePrefillInputs: + """Build suffix-only prepack inputs from prefix lookup results.""" + + plan = build_prefix_reuse_prefill_plan( + local_indices=local_indices, + sequence_ids=sequence_ids, + input_ids=input_ids, + prompt_lengths=prompt_lengths, + prefix_shared_tokens=lookup.prefix_shared_tokens, + ) + suffix_inputs = [] + suffix_masks = [] + for suffix_ids in plan.suffix_input_ids: + suffix = suffix_ids.view(1, -1) + suffix_inputs.append(suffix) + suffix_masks.append(torch.ones_like(suffix, dtype=torch.int64)) + + return PrefixCachePrefillInputs( + plan=plan, + input_ids_list=suffix_inputs, + attention_mask_list=suffix_masks, + ) diff --git a/batchgen/prefix_reuse/worker_commit.py b/batchgen/prefix_reuse/worker_commit.py new file mode 100644 index 000000000..0285e14d6 --- /dev/null +++ b/batchgen/prefix_reuse/worker_commit.py @@ -0,0 +1,170 @@ +"""BatchGenWorker-facing helpers for publishing Host KV pages.""" + +from __future__ import annotations + +from typing import Mapping + +from batchgen.prefix_reuse.commit import ( + aligned_prefix_tokens, + build_committable_prefix_token_ids, + build_prefix_commit_request, + collect_required_group_pages_for_commit, +) +from batchgen.prefix_reuse.config import PrefixCacheRuntimeConfig +from batchgen.sequence import SequenceEntry + + +def sequence_token_ids_for_prefix_commit( + seq: SequenceEntry, + *, + include_new_decode_tokens: bool, + max_tokens: int, +) -> list[int]: + """Return token ids matching the logical Host KV prefix for a sequence.""" + + prompt_token_count = int(seq.prompt_length) + prompt_tensor = seq.input_ids.reshape(-1) + prompt_data_ptr = int(prompt_tensor.data_ptr()) + prompt_version = int(prompt_tensor._version) + if ( + seq.prefix_prompt_token_ids is not None + and int(seq.prefix_prompt_cache_data_ptr) == prompt_data_ptr + and int(seq.prefix_prompt_cache_length) == prompt_token_count + and int(seq.prefix_prompt_cache_version) == prompt_version + ): + prompt_token_ids = seq.prefix_prompt_token_ids + else: + prompt_token_ids = [ + int(token_id) + for token_id in prompt_tensor[:prompt_token_count].tolist() + ] + seq.prefix_prompt_token_ids = prompt_token_ids + seq.prefix_prompt_cache_data_ptr = prompt_data_ptr + seq.prefix_prompt_cache_length = prompt_token_count + seq.prefix_prompt_cache_version = prompt_version + + decoded_token_ids: list[int] = [] + decoded_start = 0 + if include_new_decode_tokens: + decoded_start = int(seq.reentry_decoded_baseline) + decoded_length = int(seq.decoded_length) + decoded_tensor = seq.decoded_tokens + if decoded_tensor is not None and decoded_length > 0: + decoded_token_ids = [ + int(token_id) + for token_id in decoded_tensor.reshape(-1)[ + :decoded_length + ].tolist() + ] + + return build_committable_prefix_token_ids( + prompt_token_ids=prompt_token_ids, + decoded_token_ids=decoded_token_ids, + decoded_start=decoded_start, + max_tokens=max_tokens, + ) + + +def build_sequence_prefix_commit_request( + *, + core_engine_module: object, + runtime_config: PrefixCacheRuntimeConfig, + worker_views_by_group: Mapping[int, object], + seq: SequenceEntry, + include_new_decode_tokens: bool, +) -> tuple[object, int] | None: + """Build a prefix-cache commit request for one worker-owned sequence.""" + + decoded_start = int(seq.reentry_decoded_baseline) + decoded_length = int(seq.decoded_length) + new_decode_tokens = ( + max(0, decoded_length - decoded_start) + if include_new_decode_tokens + else 0 + ) + total_tokens = int(seq.prompt_length) + new_decode_tokens + commit_tokens = aligned_prefix_tokens( + total_tokens, + int(runtime_config.publish_boundary_tokens), + ) + if commit_tokens <= 0: + return None + + already_committed_tokens = max( + int(seq.prefix_shared_tokens), + int(seq.prefix_committed_tokens), + ) + if commit_tokens <= already_committed_tokens: + return None + + token_ids = sequence_token_ids_for_prefix_commit( + seq, + include_new_decode_tokens=include_new_decode_tokens, + max_tokens=commit_tokens, + ) + if len(token_ids) < commit_tokens: + raise RuntimeError( + "prefix cache commit has fewer token ids than committed tokens: " + f"got {len(token_ids)}, expected {commit_tokens}" + ) + + pages_by_group = collect_required_group_pages_for_commit( + worker_views_by_group=worker_views_by_group, + sequence_id=int(seq.global_idx), + commit_tokens=commit_tokens, + group_specs=runtime_config.group_specs, + ) + request = build_prefix_commit_request( + core_engine_module=core_engine_module, + namespace_digest=runtime_config.namespace_digest, + token_ids=token_ids, + publish_boundary_tokens=int(runtime_config.publish_boundary_tokens), + pages_by_group=pages_by_group, + raw_page_tokens_by_group={ + int(spec.group_id): int(spec.raw_page_tokens) + for spec in runtime_config.group_specs + }, + ) + if request is None: + return None + return request, commit_tokens + + +def retain_newly_committed_prefix_pages( + *, + runtime_config: PrefixCacheRuntimeConfig, + worker_views_by_group: Mapping[int, object], + sequence_id: int, + previous_committed_tokens: int, + commit_tokens: int, + page_ids_by_group: Mapping[int, list[int]] | None = None, +) -> int: + """Move newly published sequence-owned pages into resident ownership.""" + + previous = max(0, int(previous_committed_tokens)) + target = int(commit_tokens) + if target <= previous: + return previous + + for spec in runtime_config.group_specs: + if not spec.required_for_reuse: + continue + raw_page_tokens = int(spec.raw_page_tokens) + previous_pages = previous // raw_page_tokens + target_pages = target // raw_page_tokens + new_pages = target_pages - previous_pages + if new_pages <= 0: + continue + worker_view = worker_views_by_group[int(spec.group_id)] + if page_ids_by_group is None: + logical_pages = worker_view.build_page_table([int(sequence_id)])[0] + retained_pages = logical_pages[previous_pages:target_pages] + else: + retained_pages = page_ids_by_group[int(spec.group_id)][ + previous_pages:target_pages + ] + worker_view.retain_sequence_pages( + int(sequence_id), + [int(page_id) for page_id in retained_pages], + ) + return target diff --git a/batchgen/sequence.py b/batchgen/sequence.py index db5ba6bc8..36d06d83c 100644 --- a/batchgen/sequence.py +++ b/batchgen/sequence.py @@ -76,6 +76,12 @@ class SequenceEntry: # Dynamic host KV reservation tracking 'host_token_capacity', # Current host KV capacity in tokens (grows by chunk) 'host_pages_allocated', # Current host page count + 'prefix_shared_tokens', # Effective tokens reused by this prefill + 'prefix_committed_tokens', # Tokens already owned by prefix cache metadata + 'prefix_prompt_token_ids', # Cached prompt token ids for prefix commit + 'prefix_prompt_cache_data_ptr', # input_ids pointer for cached tokens + 'prefix_prompt_cache_length', # prompt length for cached tokens + 'prefix_prompt_cache_version', # input_ids tensor version for cache # Eviction support 'evicted_token_ids', # Saved (prompt + decoded) tokens for recompute after eviction 'original_prompt_length', # Original prompt length before eviction (for tracking) @@ -152,6 +158,12 @@ def __init__( # Dynamic host KV reservation: starts at 0, set by worker at prefill time self.host_token_capacity: int = 0 self.host_pages_allocated: int = 0 + self.prefix_shared_tokens: int = 0 + self.prefix_committed_tokens: int = 0 + self.prefix_prompt_token_ids: Optional[List[int]] = None + self.prefix_prompt_cache_data_ptr: int = 0 + self.prefix_prompt_cache_length: int = 0 + self.prefix_prompt_cache_version: int = -1 # Eviction support self.evicted_token_ids: Optional[torch.Tensor] = None @@ -543,6 +555,24 @@ def is_resumable(self) -> bool: def remaining_decode_tokens(self) -> int: return self.max_decode_length - self.decoded_length + def clamp_reentry_decoded_length(self, decoded_length: int) -> int: + """Clamp reconstructed decoded progress using this request's limit.""" + return min(max(0, int(decoded_length)), int(self.original_max_decode_length)) + + def compute_reentry_decoded_length(self, reconstructed_prompt_length: int) -> int: + """Return cumulative decoded progress after host-KV eviction re-entry. + + Pool-mode workers can process batches with different request-level + max_tokens. Re-entry accounting must therefore use the sequence's own + original_max_decode_length, not the worker's current batch default. + """ + prompt_delta = max( + 0, + int(reconstructed_prompt_length) - int(self.original_prompt_length), + ) + cumulative_decoded = max(prompt_delta, int(self.total_decoded_before_eviction)) + return self.clamp_reentry_decoded_length(cumulative_decoded) + def should_check_completion(self) -> bool: """Check if we're at a page boundary (every PAGE_SIZE tokens in decoding).""" return self.decoded_length > 0 and self.decoded_length % self.PAGE_SIZE == 0 diff --git a/batchgen/server/batch_scheduler.py b/batchgen/server/batch_scheduler.py index f2400285b..3a409cbdc 100644 --- a/batchgen/server/batch_scheduler.py +++ b/batchgen/server/batch_scheduler.py @@ -7,8 +7,11 @@ import logging import time import uuid +from dataclasses import dataclass +from pathlib import Path from typing import Any, Dict, List, Optional, Tuple +from batchgen.server.intake_pool import IntakeEntry, IntakePool, Priority from batchgen.server.io_struct import ( BatchEndpoint, BatchError, @@ -31,15 +34,90 @@ ToolCallFunction, Usage, ) -from batchgen.server.intake_pool import IntakeEntry, IntakePool, Priority from batchgen.server.scheduling_pool import SchedulingPool from batchgen.server.server_args import ServerArgs from batchgen.server.storage import StorageManager +from batchgen.server.usage import build_usage as make_usage +from batchgen.server.usage import build_usage_dict from batchgen.server.worker_manager import WorkerManager logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class BatchOutputMetrics: + rows: int = 0 + errors: int = 0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + cached_tokens: int = 0 + requests_with_cache: int = 0 + + @property + def cache_hit_rate(self) -> float: + if self.prompt_tokens <= 0: + return 0.0 + return self.cached_tokens / self.prompt_tokens + + +def _summarize_batch_output_file(path: Path) -> BatchOutputMetrics: + metrics = BatchOutputMetrics() + if not path.exists() or path.stat().st_size == 0: + return metrics + + rows = 0 + errors = 0 + prompt_tokens = 0 + completion_tokens = 0 + total_tokens = 0 + cached_tokens = 0 + requests_with_cache = 0 + + with path.open("r", encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if not line: + continue + rows += 1 + try: + item = json.loads(line) + except json.JSONDecodeError: + errors += 1 + continue + + response = item.get("response") or {} + status_code = int(response.get("status_code") or 0) + body = response.get("body") or {} + if item.get("error") is not None or status_code >= 400 or not body: + errors += 1 + continue + + usage = body.get("usage") or {} + prompt = int(usage.get("prompt_tokens") or 0) + completion = int(usage.get("completion_tokens") or 0) + total = int(usage.get("total_tokens") or (prompt + completion)) + details = usage.get("prompt_tokens_details") or {} + cached = int(details.get("cached_tokens") or 0) + + prompt_tokens += prompt + completion_tokens += completion + total_tokens += total + cached_tokens += cached + if cached > 0: + requests_with_cache += 1 + + return BatchOutputMetrics( + rows=rows, + errors=errors, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + cached_tokens=cached_tokens, + requests_with_cache=requests_with_cache, + ) + + def completion_prompt_to_text(prompt: str | List[str]) -> str: if isinstance(prompt, list): return "\n".join(prompt) @@ -315,8 +393,8 @@ async def _process_batch(self, batch_id: str) -> None: # If incremental save was active, use the incremental JSONL as the output incremental_path = None if incremental_output_dir: - from pathlib import Path import shutil + from pathlib import Path incremental_path = Path(incremental_output_dir) / f"{batch_id}.jsonl" if incremental_path and incremental_path.exists() and incremental_path.stat().st_size > 0: @@ -349,6 +427,12 @@ async def _process_batch(self, batch_id: str) -> None: self.storage.save_metadata(output_file_id, output_meta.dict()) completed_at = int(time.time()) + self._log_completed_batch_metrics( + batch_id=batch_id, + mode="legacy", + request_count=len(requests), + output_path=output_path, + ) self.storage.update_batch_status( batch_id, BatchStatus.COMPLETED, @@ -486,6 +570,49 @@ def _format_chat_messages(self, messages: List[dict], model: str, **kwargs) -> s messages, tokenize=False, add_generation_prompt=True, **kwargs ) + def _log_completed_batch_metrics( + self, + *, + batch_id: str, + mode: str, + request_count: int, + output_path: Path, + ) -> None: + metrics = _summarize_batch_output_file(output_path) + batch = self.storage.load_batch(batch_id) + started_at = getattr(batch, "started_at", None) if batch else None + elapsed_s = ( + max(0.0, time.time() - float(started_at)) + if started_at is not None + else None + ) + elapsed_text = ( + f"{elapsed_s:.3f}" if elapsed_s is not None else "unknown" + ) + avg_cached = ( + metrics.cached_tokens / metrics.rows if metrics.rows else 0.0 + ) + logger.info( + "[BATCH_METRICS] batch=%s mode=%s requests=%d rows=%d errors=%d " + "elapsed_s=%s prompt_tokens=%d completion_tokens=%d " + "total_tokens=%d cached_tokens=%d cache_hit_rate=%.2f%% " + "requests_with_cache=%d avg_cached_tokens=%.1f output_bytes=%d", + batch_id, + mode, + request_count, + metrics.rows, + metrics.errors, + elapsed_text, + metrics.prompt_tokens, + metrics.completion_tokens, + metrics.total_tokens, + metrics.cached_tokens, + metrics.cache_hit_rate * 100.0, + metrics.requests_with_cache, + avg_cached, + output_path.stat().st_size if output_path.exists() else 0, + ) + def _build_output_items( self, requests: List[BatchRequestItem], @@ -718,11 +845,7 @@ def _build_usage_from_text( return None prompt_tokens = self._count_tokens(tokenizer, prompt_text) completion_tokens = self._count_tokens(tokenizer, completion_text) - return Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) + return make_usage(prompt_tokens, completion_tokens) def _build_usage( self, model: str, prompt_text: str, token_ids: List[int] @@ -732,11 +855,7 @@ def _build_usage( return None prompt_tokens = self._count_tokens(tokenizer, prompt_text) completion_tokens = len(token_ids) - return Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) + return make_usage(prompt_tokens, completion_tokens) def _count_tokens(self, tokenizer: Any, text: str) -> int: if not text: @@ -1152,11 +1271,13 @@ def _write_pool_completion( decoded_text = result.get("text", "") prompt_length = result.get("prompt_length", 0) decoded_length = result.get("decoded_length", 0) + cached_tokens = result.get("cached_tokens", 0) finish_reason = result.get("finish_reason", "stop") model = meta["model"] custom_id = meta["custom_id"] url = meta["url"] created_at = int(time.time()) + usage = build_usage_dict(prompt_length, decoded_length, cached_tokens) # Build response body based on endpoint type if url == "/v1/chat/completions": @@ -1174,11 +1295,7 @@ def _write_pool_completion( "logprobs": None, "finish_reason": finish_reason, }], - "usage": { - "prompt_tokens": prompt_length, - "completion_tokens": decoded_length, - "total_tokens": prompt_length + decoded_length, - }, + "usage": usage, } else: body = { @@ -1192,11 +1309,7 @@ def _write_pool_completion( "logprobs": None, "finish_reason": finish_reason, }], - "usage": { - "prompt_tokens": prompt_length, - "completion_tokens": decoded_length, - "total_tokens": prompt_length + decoded_length, - }, + "usage": usage, } result_item = { @@ -1241,8 +1354,8 @@ def _finalize_batch_output( else None ) if incremental_output_dir: - from pathlib import Path import shutil + from pathlib import Path incremental_path = Path(incremental_output_dir) / f"{batch_id}.jsonl" if incremental_path and incremental_path.exists() and incremental_path.stat().st_size > 0: @@ -1273,6 +1386,12 @@ def _finalize_batch_output( self.storage.save_metadata(output_file_id, output_meta.dict()) completed_at = int(time.time()) + self._log_completed_batch_metrics( + batch_id=batch_id, + mode="pool", + request_count=len(requests), + output_path=output_path, + ) self.storage.update_batch_status( batch_id, BatchStatus.COMPLETED, diff --git a/batchgen/server/http_server.py b/batchgen/server/http_server.py index 93b969eb6..f8e1b89f1 100644 --- a/batchgen/server/http_server.py +++ b/batchgen/server/http_server.py @@ -89,6 +89,55 @@ async def dispatch(self, request: Request, call_next): asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) +def _count_active_batches(scheduler: BatchScheduler) -> int: + spool = scheduler._scheduling_pool + active_batches = 0 + with spool._lock: + for tracker in spool._batch_trackers.values(): + if not tracker.is_complete and not tracker.is_failed: + active_batches += 1 + return active_batches + + +def _ensure_no_active_batches(scheduler: BatchScheduler, action: str) -> None: + active_batches = _count_active_batches(scheduler) + if active_batches: + raise HTTPException( + status_code=409, + detail=( + f"Cannot {action} prefix cache while batches are active: " + f"active_batches={active_batches}" + ), + ) + + +def _token_id_batches_from_payload(payload: object) -> list[list[int]]: + if not isinstance(payload, dict): + raise HTTPException(status_code=400, detail="Expected JSON object") + token_ids = payload.get("token_ids") + if not isinstance(token_ids, list): + raise HTTPException( + status_code=400, + detail="Expected token_ids to be a list of token-id lists", + ) + + batches: list[list[int]] = [] + for index, row in enumerate(token_ids): + if not isinstance(row, list): + raise HTTPException( + status_code=400, + detail=f"Expected token_ids[{index}] to be a list", + ) + try: + batches.append([int(token_id) for token_id in row]) + except (TypeError, ValueError) as exc: + raise HTTPException( + status_code=400, + detail=f"token_ids[{index}] contains a non-integer token id", + ) from exc + return batches + + def create_app( server_args: ServerArgs, worker_exit_state: Optional[WorkerExitState] = None, @@ -168,11 +217,7 @@ async def pool_status(request: Request): scheduler: BatchScheduler = request.app.state.scheduler spool = scheduler._scheduling_pool ipool = scheduler._intake_pool - active_batches = 0 - with spool._lock: - for t in spool._batch_trackers.values(): - if not t.is_complete and not t.is_failed: - active_batches += 1 + active_batches = _count_active_batches(scheduler) return { "intake_pool_size": ipool.size(), "intake_pool_capacity": ipool.max_capacity, @@ -183,6 +228,70 @@ async def pool_status(request: Request): "pool_mode": scheduler._pool_mode, } + @app.post("/v1/prefix-cache/clear") + @app.post("/v1/prefix_cache/clear", include_in_schema=False) + async def clear_prefix_cache(request: Request): + scheduler: BatchScheduler = request.app.state.scheduler + _ensure_no_active_batches(scheduler, "clear") + + worker: WorkerManager = request.app.state.worker + try: + return await asyncio.to_thread(worker.clear_prefix_cache) + except RuntimeError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except Exception as exc: + logger.exception("Failed to clear prefix cache") + raise HTTPException(status_code=500, detail=str(exc)) + + @app.post("/v1/prefix-cache/pin") + @app.post("/v1/prefix_cache/pin", include_in_schema=False) + async def pin_prefix_cache(request: Request): + scheduler: BatchScheduler = request.app.state.scheduler + _ensure_no_active_batches(scheduler, "pin") + payload = await request.json() + token_id_batches = _token_id_batches_from_payload(payload) + replace_existing = bool(payload.get("replace_existing", False)) + + worker: WorkerManager = request.app.state.worker + try: + return await asyncio.to_thread( + worker.pin_prefix_cache, + token_id_batches, + replace_existing=replace_existing, + ) + except RuntimeError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except Exception as exc: + logger.exception("Failed to pin prefix cache") + raise HTTPException(status_code=500, detail=str(exc)) + + @app.post("/v1/prefix-cache/unpin") + @app.post("/v1/prefix_cache/unpin", include_in_schema=False) + async def unpin_prefix_cache(request: Request): + scheduler: BatchScheduler = request.app.state.scheduler + _ensure_no_active_batches(scheduler, "unpin") + + worker: WorkerManager = request.app.state.worker + try: + return await asyncio.to_thread(worker.unpin_prefix_cache) + except RuntimeError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except Exception as exc: + logger.exception("Failed to unpin prefix cache") + raise HTTPException(status_code=500, detail=str(exc)) + + @app.get("/v1/prefix-cache/pins") + @app.get("/v1/prefix_cache/pins", include_in_schema=False) + async def prefix_cache_pins(request: Request): + worker: WorkerManager = request.app.state.worker + try: + return await asyncio.to_thread(worker.prefix_cache_pin_status) + except RuntimeError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except Exception as exc: + logger.exception("Failed to inspect prefix cache pins") + raise HTTPException(status_code=500, detail=str(exc)) + # ==================== File Endpoints ==================== @app.post("/v1/files", response_model=FileObject) diff --git a/batchgen/server/incremental_writer.py b/batchgen/server/incremental_writer.py index 97e76bd89..fa4cbc07e 100644 --- a/batchgen/server/incremental_writer.py +++ b/batchgen/server/incremental_writer.py @@ -31,8 +31,8 @@ CompletionResponse, ToolCall, ToolCallFunction, - Usage, ) +from batchgen.server.usage import build_usage as make_usage logger = logging.getLogger(__name__) @@ -90,13 +90,28 @@ def __init__( f"output={self._output_path}, sequences={len(custom_id_map)}" ) - def submit(self, global_idx: int, decoded_tokens: torch.Tensor, finish_reason: str = "stop") -> None: + def submit( + self, + global_idx: int, + decoded_tokens: torch.Tensor, + finish_reason: str = "stop", + cached_tokens: int = 0, + ) -> None: """Enqueue a completed sequence for async writing. Thread-safe.""" if self._closed: logger.warning("IncrementalWriter.submit() called after close()") return - tokens_cpu = decoded_tokens.cpu() if decoded_tokens.is_cuda else decoded_tokens.clone() - self._queue.put((global_idx, tokens_cpu, finish_reason)) + tokens_cpu = ( + decoded_tokens.cpu() + if decoded_tokens.is_cuda + else decoded_tokens.clone() + ) + self._queue.put(( + global_idx, + tokens_cpu, + finish_reason, + cached_tokens, + )) def submit_error(self, global_idx: int, error_code: str, error_message: str) -> None: """Enqueue an error result for a rejected sequence. Thread-safe.""" @@ -141,13 +156,30 @@ def _background_loop(self) -> None: try: # Error items: ("error", global_idx, error_code, error_message) - if isinstance(item, tuple) and len(item) == 4 and item[0] == "error": + if ( + isinstance(item, tuple) + and len(item) == 4 + and item[0] == "error" + ): _, global_idx, error_code, error_message = item - line = self._build_error_line(global_idx, error_code, error_message) + line = self._build_error_line( + global_idx, + error_code, + error_message, + ) else: - # Normal items: (global_idx, tokens, finish_reason) - global_idx, tokens, finish_reason = item - line = self._build_result_line(global_idx, tokens, finish_reason=finish_reason) + # Normal items: (global_idx, tokens, finish_reason[, cached_tokens]) + if len(item) == 4: + global_idx, tokens, finish_reason, cached_tokens = item + else: + global_idx, tokens, finish_reason = item + cached_tokens = 0 + line = self._build_result_line( + global_idx, + tokens, + finish_reason=finish_reason, + cached_tokens=cached_tokens, + ) fh.write(line) fh.write("\n") fh.flush() @@ -162,10 +194,19 @@ def _background_loop(self) -> None: # -------------------- Result building -------------------- - def _build_result_line(self, global_idx: int, tokens: torch.Tensor, finish_reason: str = "stop") -> str: + def _build_result_line( + self, + global_idx: int, + tokens: torch.Tensor, + finish_reason: str = "stop", + cached_tokens: int = 0, + ) -> str: """Build a BatchResultItem-compatible JSON line.""" custom_id = self._custom_id_map.get(global_idx, f"unknown_{global_idx}") - endpoint_url = self._request_urls.get(global_idx, BatchEndpoint.CHAT_COMPLETIONS.value) + endpoint_url = self._request_urls.get( + global_idx, + BatchEndpoint.CHAT_COMPLETIONS.value, + ) prompt_text = self._prompt_texts.get(global_idx, "") # Detokenize @@ -180,11 +221,7 @@ def _build_result_line(self, global_idx: int, tokens: torch.Tensor, finish_reaso content, reasoning_content, tool_calls = self._parse_output(decoded_text) created_at = int(time.time()) - usage = Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) + usage = make_usage(prompt_tokens, completion_tokens, cached_tokens) if endpoint_url == BatchEndpoint.CHAT_COMPLETIONS.value: body = ChatCompletionResponse( diff --git a/batchgen/server/io_struct.py b/batchgen/server/io_struct.py index e534600a9..792375dfe 100644 --- a/batchgen/server/io_struct.py +++ b/batchgen/server/io_struct.py @@ -248,10 +248,17 @@ class BatchError(BaseModel): message: str +class PromptTokensDetails(BaseModel): + cached_tokens: int = 0 + + class Usage(BaseModel): prompt_tokens: int completion_tokens: int total_tokens: int + prompt_tokens_details: PromptTokensDetails = Field( + default_factory=PromptTokensDetails + ) class ToolCallFunction(BaseModel): diff --git a/batchgen/server/server_args.py b/batchgen/server/server_args.py index 8a32d8f83..d82a591d5 100644 --- a/batchgen/server/server_args.py +++ b/batchgen/server/server_args.py @@ -137,6 +137,10 @@ class ServerArgs: # IntakePool capacity: max total requests that can be queued. # Prevents OOM under high-load. Default 1M. Set 0 for unlimited. max_intake_capacity: int = 1_000_000 + # Host-side prefix cache. Internal sizing and namespace settings are derived + # from model and Host KV config. + enable_prefix_cache: bool = False + prefix_cache_debug_stats: bool = False def __post_init__(self): if self.storage_path is None: @@ -281,6 +285,18 @@ def _build_parser() -> argparse.ArgumentParser: help="Max total requests in the intake pool. Prevents OOM under high load. " "Default: 1000000. Set to 0 for unlimited (not recommended).", ) + parser.add_argument( + "--enable-prefix-cache", + action="store_true", + default=False, + help="Enable Host-side prefix cache reuse. Internal cache sizing is derived from Host KV settings.", + ) + parser.add_argument( + "--prefix-cache-debug-stats", + action="store_true", + default=False, + help="Emit additional Host prefix cache lookup/commit statistics.", + ) parser.add_argument( "--enable-prepack", action="store_true", @@ -592,6 +608,8 @@ def prepare_server_args(argv: Optional[list[str]] = None) -> ServerArgs: startup_timeout=parsed.startup_timeout, max_pool_size=parsed.max_pool_size, max_intake_capacity=parsed.max_intake_capacity, + enable_prefix_cache=parsed.enable_prefix_cache, + prefix_cache_debug_stats=parsed.prefix_cache_debug_stats, ) server_args.resolve_paths() validate_server_args(server_args) diff --git a/batchgen/server/usage.py b/batchgen/server/usage.py new file mode 100644 index 000000000..cc8b387d4 --- /dev/null +++ b/batchgen/server/usage.py @@ -0,0 +1,46 @@ +"""Helpers for OpenAI-compatible token usage reporting.""" + +from __future__ import annotations + +from typing import Any, Dict + +from batchgen.server.io_struct import PromptTokensDetails, Usage + + +def build_usage( + prompt_tokens: Any, + completion_tokens: Any, + cached_tokens: Any = 0, +) -> Usage: + """Build a usage model with normalized cached prompt token count.""" + prompt_count = _non_negative_int(prompt_tokens) + completion_count = _non_negative_int(completion_tokens) + cached_count = min(_non_negative_int(cached_tokens), prompt_count) + return Usage( + prompt_tokens=prompt_count, + completion_tokens=completion_count, + total_tokens=prompt_count + completion_count, + prompt_tokens_details=PromptTokensDetails(cached_tokens=cached_count), + ) + + +def build_usage_dict( + prompt_tokens: Any, + completion_tokens: Any, + cached_tokens: Any = 0, +) -> Dict[str, Any]: + usage = build_usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cached_tokens=cached_tokens, + ) + if hasattr(usage, "model_dump"): + return usage.model_dump() + return usage.dict() + + +def _non_negative_int(value: Any) -> int: + try: + return max(int(value), 0) + except (TypeError, ValueError): + return 0 diff --git a/batchgen/server/worker_manager.py b/batchgen/server/worker_manager.py index 313c62f6e..4f45a9f5c 100644 --- a/batchgen/server/worker_manager.py +++ b/batchgen/server/worker_manager.py @@ -22,8 +22,15 @@ from batchgen.kv_cache.host_kv_mananger_config import build_host_kv_config from batchgen.models.engine_loader import core_engine as bg_lib from batchgen.parameter_server_client import ParameterServerClient +from batchgen.prefix_reuse.admin import ( + clear_host_prefix_cache, + host_kv_views_by_prefix_group, + pin_host_prefix_cache, + unpin_host_prefix_cache, +) from batchgen.server.gpu_arch import detect_gpu_arch # noqa: F401 (re-export) from batchgen.server.process_utils import ( + cleanup_shm_files, cleanup_resources, get_hugepage_size, get_model_byte_size, @@ -121,6 +128,7 @@ def __init__( self._stopping = False self._monitor_interval_s = 1.0 self._ready_event = self._mp_ctx.Event() + self._prefix_cache_pin_handles: list[int] = [] # Register cleanup for skeleton state dict temp file atexit.register(self._cleanup_skeleton_state_dict_file) @@ -193,6 +201,7 @@ def _diag(msg): _diag("<<< _load_model_resources") logger.info("[startup] Model resources loaded in %.2fs", _time.monotonic() - model_start) + self._initialize_prefix_cache_owner() spawn_start = _time.monotonic() _diag(">>> _spawn_workers") @@ -279,6 +288,9 @@ def stop(self) -> None: clean_hugepages=self._hugepages_enabled, kill_workers=False, # Already handled above ) + prefix_config = getattr(self, "prefix_cache_runtime_config", None) + if prefix_config is not None: + cleanup_shm_files(prefix_config.shm_name) self.started = False logger.info("WorkerManager stopped") @@ -299,6 +311,99 @@ def _get_worker_pids(self) -> List[int]: def get_worker_exit_state(self) -> WorkerExitState: return self._worker_exit_state + def clear_prefix_cache(self) -> dict[str, Any]: + """Clear unprotected Host prefix-cache entries and release Host pages.""" + + if not self.args.enable_prefix_cache: + raise RuntimeError("Prefix cache is not enabled") + coordinator = self.prefix_cache_coordinator_owner + if coordinator is None: + raise RuntimeError("Prefix cache coordinator is not initialized") + host_kv_manager = self.host_kv_manager + if host_kv_manager is None: + raise RuntimeError("Host KV manager is not initialized") + + host_kv_views = host_kv_views_by_prefix_group( + primary_host_kv=host_kv_manager, + auxiliary_host_kv=self.host_kv_aux_manager, + ) + with self._lock: + unpin_result = self._unpin_prefix_cache_locked(coordinator) + clear_result = clear_host_prefix_cache( + coordinator=coordinator, + host_kv_views_by_group=host_kv_views, + ) + clear_result["unpin"] = unpin_result + return clear_result + + def pin_prefix_cache( + self, + token_id_batches: list[list[int]], + *, + replace_existing: bool = False, + ) -> dict[str, Any]: + """Pin cache entries matching token-id batches against eviction.""" + + if not self.args.enable_prefix_cache: + raise RuntimeError("Prefix cache is not enabled") + coordinator = self.prefix_cache_coordinator_owner + if coordinator is None: + raise RuntimeError("Prefix cache coordinator is not initialized") + runtime_config = self.prefix_cache_runtime_config + if runtime_config is None: + raise RuntimeError("Prefix cache runtime config is not initialized") + + with self._lock: + unpin_result = None + if replace_existing: + unpin_result = self._unpin_prefix_cache_locked(coordinator) + pin_result = pin_host_prefix_cache( + coordinator=coordinator, + namespace_digest=runtime_config.namespace_digest, + token_id_batches=token_id_batches, + ) + handles = [ + int(handle) + for handle in pin_result.pop("attachment_handles") + ] + self._prefix_cache_pin_handles.extend(handles) + pin_result["total_pinned"] = len(self._prefix_cache_pin_handles) + if unpin_result is not None: + pin_result["unpin"] = unpin_result + return pin_result + + def unpin_prefix_cache(self) -> dict[str, Any]: + """Release all server-owned prefix-cache pins.""" + + if not self.args.enable_prefix_cache: + raise RuntimeError("Prefix cache is not enabled") + coordinator = self.prefix_cache_coordinator_owner + if coordinator is None: + raise RuntimeError("Prefix cache coordinator is not initialized") + + with self._lock: + return self._unpin_prefix_cache_locked(coordinator) + + def _unpin_prefix_cache_locked(self, coordinator: Any) -> dict[str, Any]: + handles = list(self._prefix_cache_pin_handles) + result = unpin_host_prefix_cache( + coordinator=coordinator, + attachment_handles=handles, + ) + self._prefix_cache_pin_handles.clear() + return result + + def prefix_cache_pin_status(self) -> dict[str, Any]: + """Return server-owned prefix-cache pin state.""" + + if not self.args.enable_prefix_cache: + raise RuntimeError("Prefix cache is not enabled") + with self._lock: + return { + "status": "success", + "pinned": len(self._prefix_cache_pin_handles), + } + def infer( self, prompts: List[str], @@ -612,6 +717,8 @@ def _spawn_workers(self) -> None: adaptive_chunk_multiplier=self.args.adaptive_chunk_multiplier, fast_init=self.args.fast_init, max_pool_size=self.args.max_pool_size, + enable_prefix_cache=self.args.enable_prefix_cache, + prefix_cache_debug_stats=self.args.prefix_cache_debug_stats, kv_memfd_pid=self._get_kv_memfd_pid(), kv_memfd_fd=self._get_kv_memfd_fd(), kv_aux_memfd_fd=self._get_kv_aux_memfd_fd(), @@ -632,6 +739,46 @@ def _spawn_workers(self) -> None: daemon=True, ) + def _initialize_prefix_cache_owner(self) -> None: + self.prefix_cache_runtime_config = None + self.prefix_cache_coordinator_owner = None + if not self.args.enable_prefix_cache: + return + if getattr(self, "host_kv_manager", None) is None: + raise RuntimeError( + "--enable-prefix-cache requires Host KV cache allocation" + ) + + from batchgen.prefix_reuse.config import ( + build_prefix_cache_runtime_config, + create_host_prefix_cache_coordinator, + ) + + host_budget_gb = self.args_dict.get("host_kv_cache_size_per_rank") + if host_budget_gb is None: + raise RuntimeError( + "Prefix cache requires resolved Host KV cache budget" + ) + runtime_config = build_prefix_cache_runtime_config( + model_name=self.args.model, + kv_dtype=self.args.kv_dtype, + host_kv_cache_size_bytes=int(host_budget_gb * (1024**3)), + debug_stats=self.args.prefix_cache_debug_stats, + ) + self.prefix_cache_runtime_config = runtime_config + self.prefix_cache_coordinator_owner = create_host_prefix_cache_coordinator( + core_engine_module=bg_lib, + runtime_config=runtime_config, + create_region=True, + ) + logger.info( + "Host prefix cache initialized: shm=%s groups=%d hash_block=%d publish_boundary=%d", + runtime_config.shm_name, + len(runtime_config.group_specs), + runtime_config.hash_block_tokens, + runtime_config.publish_boundary_tokens, + ) + def _get_kv_memfd_pid(self) -> int: if self.args.fast_init and getattr(self, 'host_kv_manager', None) is not None: return os.getpid() @@ -993,8 +1140,21 @@ def allocate_host_kv_cache( enable_memfd: bool = False, ) -> Any: from batchgen.kv_cache.dual_host_kv_coordinator import DualHostKVCoordinator + from batchgen.kv_cache.glm5_kv_coordinator import GLM5HostKVCoordinator + + # GLM-5 uses model-specific logical KV groups so prefix cache can + # manage primary/indexer pages independently. + glm5 = GLM5HostKVCoordinator.create_managers( + model_name=model_name, + host_kv_cache_size=int(host_kv_cache_size_gb * (1024**3)), + enable_memfd=enable_memfd, + ) + if glm5 is not None: + primary_mgr, indexer_mgr = glm5 + logger.info("Allocated GLM-5 host KV cache: primary + indexer") + return primary_mgr, indexer_mgr - # DSA models: split budget into primary + auxiliary + # Other DSA models keep the existing dual coordinator path. dual = DualHostKVCoordinator.create_managers( model_name=model_name, host_kv_cache_size=int(host_kv_cache_size_gb * (1024**3)), diff --git a/batchgen/worker/prefill.py b/batchgen/worker/prefill.py index b8965514a..3c631b2e7 100644 --- a/batchgen/worker/prefill.py +++ b/batchgen/worker/prefill.py @@ -26,7 +26,7 @@ import math from dataclasses import dataclass -from typing import List, Sequence, Tuple +from typing import List, Optional, Sequence, Tuple @dataclass(frozen=True) @@ -46,6 +46,8 @@ class PrefillCandidate: prompt_length: int kv_token_budget: int page_size: int + estimated_shared_prefix_tokens: int = 0 + estimated_shared_prefix_page_ids: Tuple[Tuple[int, int], ...] = () @dataclass(frozen=True) @@ -63,11 +65,50 @@ class PrefillSelectionRequest: num_nodes: int gpus_per_node: int initial_gpu_page_buffer: int + charge_shared_prefix_pages: bool = False + + +@dataclass(frozen=True) +class PrefillWaveGateRequest: + """Inputs for deciding whether to start a selected prefill wave. + + Prefix-cache hits can make the selected wave's real append work much + smaller than the prompt-length-based admission estimate. When active decode + work already exists, running a tiny extra prefill wave pays the full + decode->prefill transition cost for little compute. The gate only applies + to that case; first waves and non-prefix-cache runs keep the legacy path. + """ + + selected_count: int + prefix_cache_enabled: bool + has_active_work: bool + world_size: int + min_sequences: Optional[int] = None class PrefillScheduler: """Prefill admission decision — pure, deterministic across ranks.""" + @staticmethod + def min_prefix_cache_wave_sequences(world_size: int) -> int: + """Minimum selected sequence count for prefill while decode is active.""" + return max(128, max(1, int(world_size)) * 16) + + @staticmethod + def should_run_prefill_wave(req: PrefillWaveGateRequest) -> bool: + """Return whether the selected wave should be launched immediately.""" + if req.selected_count <= 0: + return False + if not req.prefix_cache_enabled or not req.has_active_work: + return True + + min_sequences = ( + req.min_sequences + if req.min_sequences is not None + else PrefillScheduler.min_prefix_cache_wave_sequences(req.world_size) + ) + return req.selected_count >= int(min_sequences) + @staticmethod def select_prefill_batch(req: PrefillSelectionRequest) -> List[str]: """Select which candidate sequences to prefill, bounded by host KV. @@ -81,8 +122,10 @@ def select_prefill_batch(req: PrefillSelectionRequest) -> List[str]: ``max(prompt_length + chunk_size, gpu_initial_tokens)`` capped at ``kv_token_budget``, rounded up to whole pages — where ``gpu_initial_tokens`` covers ``prompt_length + 1`` plus the GPU - page buffer. No safety margin: selection and allocation use the - same formula by design. + page buffer. With prefix-cache estimates, page-aligned shared prefix + pages are charged as already resident and only the private append + capacity is admitted. No safety margin: selection and allocation use + the same formula by design. Pure: reads only the candidate snapshots + per-node free pages. The NCCL gather and the ``global_batch`` enumeration stay on the @@ -99,6 +142,7 @@ def select_prefill_batch(req: PrefillSelectionRequest) -> List[str]: per_node_effective_free = list(req.per_node_host_free) node_pages_used = [0] * req.num_nodes + protected_shared_pages = [set() for _ in range(req.num_nodes)] prefill_batch: List[str] = [] for c in all_candidates: @@ -111,10 +155,29 @@ def select_prefill_batch(req: PrefillSelectionRequest) -> List[str]: gpu_initial_tokens = gpu_initial_pages * c.page_size initial_capacity = max(c.prompt_length + req.chunk_size, gpu_initial_tokens) initial_capacity = min(initial_capacity, c.kv_token_budget) - req_pages = math.ceil(initial_capacity / c.page_size) + shared_tokens = max(0, int(c.estimated_shared_prefix_tokens)) + shared_tokens = min(shared_tokens, int(c.prompt_length)) + shared_page_tokens = (shared_tokens // c.page_size) * c.page_size + append_tokens = max(0, int(c.prompt_length) - shared_tokens) + private_capacity = max( + initial_capacity - shared_page_tokens, + append_tokens, + ) + req_pages = math.ceil(private_capacity / c.page_size) + if req.charge_shared_prefix_pages: + shared_pages = protected_shared_pages[seq_node] + req_pages += sum( + 1 + for page_key in c.estimated_shared_prefix_page_ids + if page_key not in shared_pages + ) if node_pages_used[seq_node] + req_pages <= per_node_effective_free[seq_node]: prefill_batch.append(c.uuid) node_pages_used[seq_node] += req_pages + if req.charge_shared_prefix_pages: + protected_shared_pages[seq_node].update( + c.estimated_shared_prefix_page_ids + ) return prefill_batch diff --git a/batchgen/worker/sync.py b/batchgen/worker/sync.py index 946cd376d..6486c107c 100644 --- a/batchgen/worker/sync.py +++ b/batchgen/worker/sync.py @@ -153,7 +153,11 @@ def sync_completion_status_tensor( if uuid in ctx.uuid_to_local: seq = ctx.global_batch.get_sequence(uuid) if seq is not None and uuid in uuid_to_idx: - is_completed = (seq.status == SequenceStatus.COMPLETED or seq.eos_reached) + is_completed = ( + seq.status == SequenceStatus.COMPLETED + or seq.eos_reached + or seq.decoded_length >= seq.max_decode_length + ) if is_completed: completion_tensor[uuid_to_idx[uuid]] = 1 diff --git a/core/KV_Storage/host_kv_page_table.cpp b/core/KV_Storage/host_kv_page_table.cpp index ba8fbf602..9c7cbb346 100644 --- a/core/KV_Storage/host_kv_page_table.cpp +++ b/core/KV_Storage/host_kv_page_table.cpp @@ -32,6 +32,15 @@ void HostKVPageTable::AppendPages( additional_pages.end()); } +void HostKVPageTable::PrependPages( + std::int64_t sequence_id, + const std::vector& prefix_pages) { + std::unique_lock lock(mutex_); + SequenceRecord& record = RequireRecordLocked(sequence_id, lock); + record.pages.insert(record.pages.begin(), prefix_pages.begin(), + prefix_pages.end()); +} + std::vector HostKVPageTable::PopPrefixPages( std::int64_t sequence_id, std::size_t num_pages) { if (num_pages == 0) { diff --git a/core/KV_Storage/host_kv_page_table.h b/core/KV_Storage/host_kv_page_table.h index 3bb6433e8..cb324388f 100644 --- a/core/KV_Storage/host_kv_page_table.h +++ b/core/KV_Storage/host_kv_page_table.h @@ -32,6 +32,9 @@ class HostKVPageTable { void AppendPages(std::int64_t sequence_id, const std::vector& additional_pages); + void PrependPages(std::int64_t sequence_id, + const std::vector& prefix_pages); + std::vector PopPrefixPages(std::int64_t sequence_id, std::size_t num_pages); diff --git a/core/KV_Storage/host_paged_kv_backend.cpp b/core/KV_Storage/host_paged_kv_backend.cpp index 6312a08d4..36737deee 100644 --- a/core/KV_Storage/host_paged_kv_backend.cpp +++ b/core/KV_Storage/host_paged_kv_backend.cpp @@ -1,5 +1,7 @@ #include "host_paged_kv_backend.h" +#include "shared_memory_utils.h" + #include #include #include @@ -22,6 +24,7 @@ #include #include #include +#include #include namespace batchgen::kv { @@ -34,63 +37,7 @@ constexpr std::int32_t kInvalidPageIndex = -1; constexpr std::int64_t kEmptySequenceId = std::numeric_limits::min(); constexpr std::int64_t kTombstoneSequenceId = kEmptySequenceId + 1; - -enum class InitState : std::uint32_t { - kUninitialized = 0, - kInitializing = 1, - kReady = 2, -}; - -std::size_t AlignUp(std::size_t value, std::size_t alignment) { - if (alignment == 0) { - return value; - } - const std::size_t remainder = value % alignment; - if (remainder == 0) { - return value; - } - return value + (alignment - remainder); -} - -std::size_t GetSystemPageSize() { - const long page_size = sysconf(_SC_PAGESIZE); - if (page_size <= 0) { - const int err = errno; - throw std::system_error(err, std::generic_category(), - "sysconf(_SC_PAGESIZE) failed"); - } - return static_cast(page_size); -} - -class ScopedMutexLock { - public: - explicit ScopedMutexLock(pthread_mutex_t* mu) : mu_(mu) { - int rc = pthread_mutex_lock(mu_); - if (rc == EOWNERDEAD) { - const int consistent_rc = pthread_mutex_consistent(mu_); - if (consistent_rc != 0) { - throw std::system_error(consistent_rc, std::generic_category(), - "pthread_mutex_consistent failed"); - } - } else if (rc != 0) { - throw std::system_error(rc, std::generic_category(), - "pthread_mutex_lock failed"); - } - } - - ScopedMutexLock(const ScopedMutexLock&) = delete; - ScopedMutexLock& operator=(const ScopedMutexLock&) = delete; - - ~ScopedMutexLock() { - const int rc = pthread_mutex_unlock(mu_); - if (rc != 0) { - std::terminate(); // Unlock failure is irrecoverable here. - } - } - - private: - pthread_mutex_t* mu_; -}; +constexpr std::int64_t kPrefixResidentSequenceId = kEmptySequenceId + 2; struct SequenceEntry { std::int64_t sequence_id = kEmptySequenceId; @@ -101,7 +48,7 @@ struct SequenceEntry { struct SharedHeader { std::atomic init_state{ - static_cast(InitState::kUninitialized)}; + static_cast(SharedMemoryInitState::kUninitialized)}; std::uint64_t magic = kSharedMemoryMagic; std::uint64_t layout_fingerprint = 0; std::uint64_t config_hash = 0; @@ -203,8 +150,20 @@ struct HostPagedKVBackend::SharedState { void ReleaseSequence(std::int64_t sequence_id); std::vector ReleasePrefixPages(std::int64_t sequence_id, std::size_t num_pages); + std::vector RetainPrefixPages(std::int64_t sequence_id, + std::size_t num_pages); + std::vector RetainPageRange(std::int64_t sequence_id, + std::size_t start_page, + std::size_t num_pages); + std::vector RetainPages( + std::int64_t sequence_id, + const std::vector& page_ids); + void ReleaseResidentPages(const std::vector& page_ids); std::vector SequencePages( std::int64_t sequence_id, std::optional max_pages) const; + std::vector SequencePageRange( + std::int64_t sequence_id, std::size_t start_page, + std::size_t page_count) const; HostPagedKVStats CollectStats() const; std::byte* DataBase() { return data_base; } @@ -330,52 +289,31 @@ void HostPagedKVBackend::SharedState::ConstructSharedState() { sequence_table[i] = SequenceEntry(); } - pthread_mutexattr_t attr; - if (const int rc = pthread_mutexattr_init(&attr); rc != 0) { - throw std::system_error(rc, std::generic_category(), - "pthread_mutexattr_init failed"); - } - if (const int rc = - pthread_mutexattr_setpshared(&attr, PTHREAD_PROCESS_SHARED); - rc != 0) { - pthread_mutexattr_destroy(&attr); - throw std::system_error(rc, std::generic_category(), - "pthread_mutexattr_setpshared failed"); - } - if (const int rc = pthread_mutexattr_setrobust(&attr, PTHREAD_MUTEX_ROBUST); - rc != 0) { - pthread_mutexattr_destroy(&attr); - throw std::system_error(rc, std::generic_category(), - "pthread_mutexattr_setrobust failed"); - } - - if (const int rc = pthread_mutex_init(&header->allocation_mutex, &attr); - rc != 0) { - pthread_mutexattr_destroy(&attr); - throw std::system_error(rc, std::generic_category(), - "pthread_mutex_init allocation_mutex failed"); - } - if (const int rc = pthread_mutex_init(&header->sequence_mutex, &attr); - rc != 0) { + InitProcessSharedRobustMutex( + &header->allocation_mutex, + "pthread_mutex_init allocation_mutex failed"); + try { + InitProcessSharedRobustMutex( + &header->sequence_mutex, + "pthread_mutex_init sequence_mutex failed"); + } catch (...) { pthread_mutex_destroy(&header->allocation_mutex); - pthread_mutexattr_destroy(&attr); - throw std::system_error(rc, std::generic_category(), - "pthread_mutex_init sequence_mutex failed"); + throw; } - pthread_mutexattr_destroy(&attr); - header->init_state.store(static_cast(InitState::kReady), - std::memory_order_release); + header->init_state.store( + static_cast(SharedMemoryInitState::kReady), + std::memory_order_release); } void HostPagedKVBackend::SharedState::WaitForInitialization() const { while (true) { - const auto state = static_cast( + const auto state = static_cast( header->init_state.load(std::memory_order_acquire)); - if (state == InitState::kReady) { + if (state == SharedMemoryInitState::kReady) { return; } - if (state == InitState::kUninitialized) { + if (state == SharedMemoryInitState::kUninitialized) { std::this_thread::sleep_for(std::chrono::milliseconds(1)); continue; } @@ -487,7 +425,7 @@ SequenceEntry* HostPagedKVBackend::SharedState::FindOrInsertSequenceEntryLocked( } void HostPagedKVBackend::SharedState::Initialize(bool create_region) { - const std::size_t page_size = GetSystemPageSize(); + const std::size_t page_size = SystemPageSize(); total_bytes = AlignUp(total_bytes_unaligned, page_size); constexpr std::size_t kHugePageSize = 2 * 1024 * 1024; const std::size_t alignment = std::max(kHugePageSize, page_size); @@ -531,7 +469,8 @@ void HostPagedKVBackend::SharedState::Initialize(bool create_region) { mapping = static_cast(mapped); MapPointers(); header->init_state.store( - static_cast(InitState::kInitializing), + static_cast( + SharedMemoryInitState::kInitializing), std::memory_order_relaxed); ConstructSharedState(); } else { @@ -666,7 +605,7 @@ void HostPagedKVBackend::SharedState::Initialize(bool create_region) { if (created_region) { header->init_state.store( - static_cast(InitState::kInitializing), + static_cast(SharedMemoryInitState::kInitializing), std::memory_order_relaxed); ConstructSharedState(); } else { @@ -682,7 +621,7 @@ std::vector HostPagedKVBackend::SharedState::AcquirePages( } std::vector pages(num_pages); { - ScopedMutexLock lock(&header->allocation_mutex); + ScopedPthreadMutexLock lock(&header->allocation_mutex); const std::uint32_t top = header->free_stack_top.load(std::memory_order_relaxed); if (top < num_pages) { @@ -700,7 +639,7 @@ std::vector HostPagedKVBackend::SharedState::AcquirePages( } { - ScopedMutexLock lock(&header->sequence_mutex); + ScopedPthreadMutexLock lock(&header->sequence_mutex); bool is_new = false; SequenceEntry* entry = FindOrInsertSequenceEntryLocked(sequence_id, &is_new); @@ -728,7 +667,7 @@ void HostPagedKVBackend::SharedState::ReleaseSequence( std::int64_t sequence_id) { std::vector pages; { - ScopedMutexLock lock(&header->sequence_mutex); + ScopedPthreadMutexLock lock(&header->sequence_mutex); SequenceEntry* entry = FindSequenceEntryLocked(sequence_id); if (entry == nullptr) { throw std::out_of_range("Sequence ID " + @@ -752,7 +691,7 @@ void HostPagedKVBackend::SharedState::ReleaseSequence( } if (!pages.empty()) { - ScopedMutexLock lock(&header->allocation_mutex); + ScopedPthreadMutexLock lock(&header->allocation_mutex); std::uint32_t top = header->free_stack_top.load(std::memory_order_relaxed); for (std::int32_t page : pages) { @@ -769,7 +708,7 @@ std::vector HostPagedKVBackend::SharedState::ReleasePrefixPages( } std::vector pages; { - ScopedMutexLock lock(&header->sequence_mutex); + ScopedPthreadMutexLock lock(&header->sequence_mutex); SequenceEntry* entry = FindSequenceEntryLocked(sequence_id); if (entry == nullptr) { throw std::out_of_range("Sequence ID " + @@ -806,7 +745,7 @@ std::vector HostPagedKVBackend::SharedState::ReleasePrefixPages( } if (!pages.empty()) { - ScopedMutexLock lock(&header->allocation_mutex); + ScopedPthreadMutexLock lock(&header->allocation_mutex); std::uint32_t top = header->free_stack_top.load(std::memory_order_relaxed); for (std::int32_t page : pages) { @@ -817,9 +756,193 @@ std::vector HostPagedKVBackend::SharedState::ReleasePrefixPages( return pages; } +std::vector HostPagedKVBackend::SharedState::RetainPrefixPages( + std::int64_t sequence_id, std::size_t num_pages) { + return RetainPageRange(sequence_id, 0, num_pages); +} + +std::vector HostPagedKVBackend::SharedState::RetainPageRange( + std::int64_t sequence_id, std::size_t start_page, std::size_t num_pages) { + if (num_pages == 0) { + return {}; + } + std::vector pages; + { + ScopedPthreadMutexLock lock(&header->sequence_mutex); + SequenceEntry* entry = FindSequenceEntryLocked(sequence_id); + if (entry == nullptr) { + throw std::out_of_range("Sequence ID " + + std::to_string(sequence_id) + + " not found during prefix retain"); + } + if (start_page > entry->num_pages || + num_pages > entry->num_pages - start_page) { + throw std::out_of_range( + "Requested retain of " + std::to_string(num_pages) + + " pages from offset " + std::to_string(start_page) + + " but sequence " + std::to_string(sequence_id) + + " only owns " + std::to_string(entry->num_pages) + + " pages"); + } + + pages.reserve(num_pages); + std::int32_t page = entry->head_page; + std::int32_t previous_page = kInvalidPageIndex; + for (std::size_t i = 0; i < start_page; ++i) { + if (page == kInvalidPageIndex) { + throw std::logic_error( + "Corrupt page chain before retained range for sequence " + + std::to_string(sequence_id)); + } + previous_page = page; + page = page_links[page]; + } + + for (std::size_t i = 0; i < num_pages; ++i) { + if (page == kInvalidPageIndex) { + throw std::logic_error( + "Corrupt page chain during range retain for sequence " + + std::to_string(sequence_id)); + } + pages.push_back(page); + const std::int32_t next = page_links[page]; + page_links[page] = kInvalidPageIndex; + page_owners[page] = kPrefixResidentSequenceId; + page = next; + } + + if (previous_page == kInvalidPageIndex) { + entry->head_page = page; + } else { + page_links[previous_page] = page; + } + entry->num_pages -= static_cast(num_pages); + if (entry->num_pages == 0) { + entry->tail_page = kInvalidPageIndex; + } else if (page == kInvalidPageIndex) { + entry->tail_page = previous_page; + } + } + return pages; +} + +std::vector HostPagedKVBackend::SharedState::RetainPages( + std::int64_t sequence_id, const std::vector& page_ids) { + if (page_ids.empty()) { + return {}; + } + std::vector pages; + pages.reserve(page_ids.size()); + { + ScopedPthreadMutexLock lock(&header->sequence_mutex); + SequenceEntry* entry = FindSequenceEntryLocked(sequence_id); + if (entry == nullptr) { + throw std::out_of_range("Sequence ID " + + std::to_string(sequence_id) + + " not found during page retain"); + } + + std::unordered_set requested_pages; + requested_pages.reserve(page_ids.size()); + for (const std::int32_t page : page_ids) { + if (page < 0 || + static_cast(page) >= config.num_pages) { + throw std::out_of_range("Retained page id out of range: " + + std::to_string(page)); + } + if (!requested_pages.insert(page).second) { + throw std::runtime_error( + "Duplicate retained page id: " + std::to_string(page)); + } + if (page_owners[page] != sequence_id) { + throw std::runtime_error( + "Cannot retain page " + std::to_string(page) + + " for sequence " + std::to_string(sequence_id) + + " because it is not sequence-owned"); + } + } + + std::size_t found = 0; + std::int32_t page = entry->head_page; + while (page != kInvalidPageIndex) { + if (requested_pages.find(page) != requested_pages.end()) { + ++found; + } + page = page_links[page]; + } + if (found != requested_pages.size()) { + throw std::logic_error( + "Sequence-owned retained pages are not present in the " + "sequence page chain for sequence " + + std::to_string(sequence_id)); + } + + std::int32_t previous_page = kInvalidPageIndex; + page = entry->head_page; + while (page != kInvalidPageIndex) { + const std::int32_t next = page_links[page]; + if (requested_pages.find(page) != requested_pages.end()) { + if (previous_page == kInvalidPageIndex) { + entry->head_page = next; + } else { + page_links[previous_page] = next; + } + if (entry->tail_page == page) { + entry->tail_page = previous_page; + } + page_links[page] = kInvalidPageIndex; + page_owners[page] = kPrefixResidentSequenceId; + pages.push_back(page); + --entry->num_pages; + } else { + previous_page = page; + } + page = next; + } + if (entry->num_pages == 0) { + entry->head_page = kInvalidPageIndex; + entry->tail_page = kInvalidPageIndex; + } + } + return pages; +} + +void HostPagedKVBackend::SharedState::ReleaseResidentPages( + const std::vector& page_ids) { + if (page_ids.empty()) { + return; + } + { + ScopedPthreadMutexLock lock(&header->sequence_mutex); + for (const std::int32_t page : page_ids) { + if (page < 0 || + static_cast(page) >= config.num_pages) { + throw std::out_of_range( + "Resident page id out of range: " + + std::to_string(page)); + } + if (page_owners[page] != kPrefixResidentSequenceId) { + throw std::runtime_error( + "Cannot release page " + std::to_string(page) + + " because it is not prefix-resident"); + } + page_owners[page] = kEmptySequenceId; + page_links[page] = kInvalidPageIndex; + } + } + + ScopedPthreadMutexLock lock(&header->allocation_mutex); + std::uint32_t top = + header->free_stack_top.load(std::memory_order_relaxed); + for (const std::int32_t page : page_ids) { + free_stack[top++] = page; + } + header->free_stack_top.store(top, std::memory_order_relaxed); +} + std::vector HostPagedKVBackend::SharedState::SequencePages( std::int64_t sequence_id, std::optional max_pages) const { - ScopedMutexLock lock(&header->sequence_mutex); + ScopedPthreadMutexLock lock(&header->sequence_mutex); SequenceEntry* entry = FindSequenceEntryLocked(sequence_id); if (entry == nullptr) { throw std::out_of_range("Sequence ID " + std::to_string(sequence_id) + @@ -848,6 +971,53 @@ std::vector HostPagedKVBackend::SharedState::SequencePages( return pages; } +std::vector HostPagedKVBackend::SharedState::SequencePageRange( + std::int64_t sequence_id, std::size_t start_page, + std::size_t page_count) const { + ScopedPthreadMutexLock lock(&header->sequence_mutex); + SequenceEntry* entry = FindSequenceEntryLocked(sequence_id); + if (entry == nullptr) { + throw std::out_of_range("Sequence ID " + std::to_string(sequence_id) + + " not found when fetching page range"); + } + const std::size_t available_pages = entry->num_pages; + if (start_page > available_pages) { + throw std::out_of_range( + "Requested page range starting at " + std::to_string(start_page) + + " but sequence " + std::to_string(sequence_id) + " only owns " + + std::to_string(available_pages) + " pages"); + } + if (page_count > available_pages - start_page) { + throw std::out_of_range( + "Requested " + std::to_string(page_count) + + " pages from offset " + std::to_string(start_page) + + " but sequence " + std::to_string(sequence_id) + " only has " + + std::to_string(available_pages) + " pages"); + } + + std::vector pages; + pages.reserve(page_count); + std::int32_t page = entry->head_page; + for (std::size_t skipped = 0; skipped < start_page; ++skipped) { + if (page == kInvalidPageIndex) { + throw std::logic_error( + "Corrupt page chain while skipping to page range for sequence " + + std::to_string(sequence_id)); + } + page = page_links[page]; + } + for (std::size_t count = 0; count < page_count; ++count) { + if (page == kInvalidPageIndex) { + throw std::logic_error( + "Corrupt page chain while reading page range for sequence " + + std::to_string(sequence_id)); + } + pages.push_back(page); + page = page_links[page]; + } + return pages; +} + HostPagedKVStats HostPagedKVBackend::SharedState::CollectStats() const { HostPagedKVStats stats; stats.num_total_pages = config.num_pages; @@ -979,11 +1149,39 @@ std::vector HostPagedKVBackend::ReleaseSequencePrefixPages( return state_->ReleasePrefixPages(sequence_id, num_pages); } +std::vector HostPagedKVBackend::RetainSequencePrefixPages( + std::int64_t sequence_id, std::size_t num_pages) { + return state_->RetainPrefixPages(sequence_id, num_pages); +} + +std::vector HostPagedKVBackend::RetainSequencePageRange( + std::int64_t sequence_id, std::size_t start_page, + std::size_t num_pages) { + return state_->RetainPageRange(sequence_id, start_page, num_pages); +} + +std::vector HostPagedKVBackend::RetainSequencePages( + std::int64_t sequence_id, + const std::vector& page_ids) { + return state_->RetainPages(sequence_id, page_ids); +} + +void HostPagedKVBackend::ReleaseResidentPages( + const std::vector& page_ids) { + state_->ReleaseResidentPages(page_ids); +} + std::vector HostPagedKVBackend::SequencePages( std::int64_t sequence_id, std::optional max_pages) const { return state_->SequencePages(sequence_id, max_pages); } +std::vector HostPagedKVBackend::SequencePageRange( + std::int64_t sequence_id, std::size_t start_page, + std::size_t page_count) const { + return state_->SequencePageRange(sequence_id, start_page, page_count); +} + HostPagedKVStats HostPagedKVBackend::CollectStats() const { return state_->CollectStats(); } diff --git a/core/KV_Storage/host_paged_kv_backend.h b/core/KV_Storage/host_paged_kv_backend.h index 90ba1e562..7444021b2 100644 --- a/core/KV_Storage/host_paged_kv_backend.h +++ b/core/KV_Storage/host_paged_kv_backend.h @@ -234,9 +234,26 @@ class HostPagedKVBackend { std::vector ReleaseSequencePrefixPages( std::int64_t sequence_id, std::size_t num_pages); + std::vector RetainSequencePrefixPages( + std::int64_t sequence_id, std::size_t num_pages); + + std::vector RetainSequencePageRange( + std::int64_t sequence_id, std::size_t start_page, + std::size_t num_pages); + + std::vector RetainSequencePages( + std::int64_t sequence_id, + const std::vector& page_ids); + + void ReleaseResidentPages(const std::vector& page_ids); + std::vector SequencePages( std::int64_t sequence_id, std::optional max_pages) const; + std::vector SequencePageRange( + std::int64_t sequence_id, std::size_t start_page, + std::size_t page_count) const; + HostPagedKVStats CollectStats() const; std::byte* DataBase(); diff --git a/core/KV_Storage/host_paged_kv_manager.h b/core/KV_Storage/host_paged_kv_manager.h index 7ac08a78c..108cb1cea 100644 --- a/core/KV_Storage/host_paged_kv_manager.h +++ b/core/KV_Storage/host_paged_kv_manager.h @@ -164,6 +164,10 @@ class HostPagedKVManager { backend_.ReleaseSequences(sequence_ids); } + void ReleaseResidentPages(const std::vector& page_ids) { + backend_.ReleaseResidentPages(page_ids); + } + std::pair, std::optional>> GetSequenceLayerPagePointers( std::int64_t sequence_id, std::size_t layer_idx, @@ -200,6 +204,37 @@ class HostPagedKVManager { return {std::move(k_ptrs), std::move(v_ptrs)}; } + std::pair, std::optional>> + GetSequenceLayerPageRangePointers(std::int64_t sequence_id, + std::size_t layer_idx, + std::size_t start_page, + std::size_t page_count) const { + geometry_.EnsureLayerBounds( + layer_idx, + "HostPagedKVManager::GetSequenceLayerPageRangePointers"); + auto page_indices = + backend_.SequencePageRange(sequence_id, start_page, page_count); + std::vector k_ptrs; + k_ptrs.reserve(page_indices.size()); + std::optional> v_ptrs; + if constexpr (Layout::kHasVCache) { + v_ptrs.emplace(); + v_ptrs->reserve(page_indices.size()); + } + std::byte* base = const_cast(backend_.DataBase()); + for (std::int32_t page : page_indices) { + void* k_ptr = + static_cast(layout_.KPageAddress(base, layer_idx, page)); + k_ptrs.emplace_back(k_ptr); + if constexpr (Layout::kHasVCache) { + void* v_ptr = static_cast( + layout_.VPageAddress(base, layer_idx, page)); + v_ptrs->emplace_back(v_ptr); + } + } + return {std::move(k_ptrs), std::move(v_ptrs)}; + } + std::vector> BuildPageTable( const std::vector& sequence_ids) const { std::vector> table; @@ -281,4 +316,4 @@ using MLAHostPagedKVManager = HostPagedKVManager; } // namespace batchgen::kv -#endif // HOST_PAGED_KV_MANAGER_H_ \ No newline at end of file +#endif // HOST_PAGED_KV_MANAGER_H_ diff --git a/core/KV_Storage/host_paged_kv_worker_view.h b/core/KV_Storage/host_paged_kv_worker_view.h index a25e7b6bd..10efd1f37 100644 --- a/core/KV_Storage/host_paged_kv_worker_view.h +++ b/core/KV_Storage/host_paged_kv_worker_view.h @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -166,6 +167,151 @@ struct KVAsyncTask { std::shared_future future_; }; +struct LayeredLoadState { + LayeredLoadState(std::uint64_t task_id, int target_device, + std::size_t layer_count, + std::function layer_resolver, + std::shared_ptr task_logger) + : id(task_id), + device_index(target_device), + layer_events(layer_count, nullptr), + resolve_layer(std::move(layer_resolver)), + logger(std::move(task_logger)) {} + + LayeredLoadState(const LayeredLoadState&) = delete; + LayeredLoadState& operator=(const LayeredLoadState&) = delete; + + ~LayeredLoadState() noexcept { + if (device_index >= 0) { + c10::cuda::OptionalCUDAGuard guard(device_index); + if (final_event != nullptr && final_event_recorded) { + const auto status = cudaEventSynchronize(final_event); + if (status != cudaSuccess && logger != nullptr) { + logger->error( + "Failed to synchronize layered load final event: {}", + cudaGetErrorString(status)); + } + } else if (h2d_stream != nullptr && has_enqueued_work) { + const auto status = cudaStreamSynchronize(h2d_stream); + if (status != cudaSuccess && logger != nullptr) { + logger->error( + "Failed to synchronize layered load stream: {}", + cudaGetErrorString(status)); + } + } + k_device_src_ptrs.Reset(); + k_device_dst_ptrs.Reset(); + v_device_src_ptrs.Reset(); + v_device_dst_ptrs.Reset(); + + for (cudaEvent_t event : layer_events) { + DestroyEvent(event); + } + DestroyEvent(final_event); + return; + } + + for (cudaEvent_t event : layer_events) { + DestroyEvent(event); + } + DestroyEvent(final_event); + } + + static void DestroyEvent(cudaEvent_t event) noexcept { + if (event == nullptr) { + return; + } + const auto status = cudaEventDestroy(event); + (void)status; + } + + std::uint64_t id = 0; + int device_index = -1; + cudaStream_t h2d_stream = nullptr; + std::vector layer_events; + cudaEvent_t final_event = nullptr; + bool final_event_recorded = false; + bool has_enqueued_work = false; + worker_detail::DeviceBuffer k_device_src_ptrs; + worker_detail::DeviceBuffer k_device_dst_ptrs; + worker_detail::DeviceBuffer v_device_src_ptrs; + worker_detail::DeviceBuffer v_device_dst_ptrs; + std::function resolve_layer; + std::shared_ptr logger; +}; + +class KVLayeredAsyncTask { + public: + KVLayeredAsyncTask() = default; + explicit KVLayeredAsyncTask(std::shared_ptr state) + : state_(std::move(state)) {} + + [[nodiscard]] std::uint64_t id() const { + return state_ != nullptr ? state_->id : 0; + } + + [[nodiscard]] std::size_t num_layers() const { + return state_ != nullptr ? state_->layer_events.size() : 0; + } + + [[nodiscard]] bool done() const { + if (state_ == nullptr || state_->final_event == nullptr || + !state_->final_event_recorded) { + return true; + } + c10::cuda::OptionalCUDAGuard guard(state_->device_index); + const auto status = cudaEventQuery(state_->final_event); + if (status == cudaSuccess) { + return true; + } + if (status == cudaErrorNotReady) { + return false; + } + CUDA_CHECK(status); + return false; + } + + void wait() const { + if (state_ == nullptr || state_->final_event == nullptr || + !state_->final_event_recorded) { + return; + } + c10::cuda::OptionalCUDAGuard guard(state_->device_index); + CUDA_CHECK(cudaEventSynchronize(state_->final_event)); + } + + void result() const { wait(); } + + void wait_for_layer(std::size_t layer_idx) const { + if (state_ == nullptr || state_->layer_events.empty()) { + return; + } + const std::size_t event_idx = + state_->resolve_layer != nullptr ? state_->resolve_layer(layer_idx) + : layer_idx; + if (event_idx >= state_->layer_events.size()) { + std::ostringstream oss; + oss << "KVLayeredAsyncTask::wait_for_layer: layer " << layer_idx + << " resolved to event " << event_idx + << " but task has " << state_->layer_events.size() + << " events"; + throw std::out_of_range(oss.str()); + } + cudaEvent_t event = state_->layer_events[event_idx]; + if (event == nullptr) { + throw std::runtime_error( + "KVLayeredAsyncTask::wait_for_layer: missing layer event"); + } + c10::cuda::OptionalCUDAGuard guard(state_->device_index); + const auto compute_stream = + at::cuda::getCurrentCUDAStream(state_->device_index).stream(); + CUDA_CHECK(cudaStreamWaitEvent(compute_stream, event, 0)); + } + + private: + std::shared_ptr state_; +}; + using SequenceLengthMap = std::unordered_map; using SequenceLengthVector = std::vector; using SequenceLengths = std::variant; @@ -359,6 +505,33 @@ class HostPagedKVWorkerView : private LayerMapper { return new_pages; } + void AttachSharedPrefixPages( + std::int64_t sequence_id, + const std::vector& page_ids) { + if (page_ids.empty()) { + return; + } + EnsureSequenceRegistered(sequence_id); + page_table_.PrependPages(sequence_id, page_ids); + } + + void AttachSharedPrefixPagesForSequences( + const std::vector& sequence_ids, + const std::vector>& page_ids_by_sequence) { + if (sequence_ids.size() != page_ids_by_sequence.size()) { + throw std::invalid_argument( + "sequence_ids and page_ids_by_sequence must have the same " + "length"); + } + EnsureSequencesRegistered(sequence_ids); + for (std::size_t i = 0; i < sequence_ids.size(); ++i) { + if (!page_ids_by_sequence[i].empty()) { + page_table_.PrependPages(sequence_ids[i], + page_ids_by_sequence[i]); + } + } + } + std::vector> GrowPagesForSequences( const std::vector& sequence_ids, const std::vector& num_pages) { @@ -612,20 +785,8 @@ class HostPagedKVWorkerView : private LayerMapper { validated_counts, "active_page_counts", kOpName); auto page_table = BuildPageTable(sequence_vector); - const auto max_sequence_pages = - static_cast(validated_k_ptrs.size(2)); - std::vector sequence_offsets(batch_size, 0); - std::size_t total_pages = 0; for (std::size_t seq_idx = 0; seq_idx < batch_size; ++seq_idx) { const std::size_t requested = page_counts[seq_idx]; - if (requested > max_sequence_pages) { - std::ostringstream oss; - oss << kOpName << ": requested pages " << requested - << " exceed provided pointer tensor capacity " - << max_sequence_pages << " for sequence index " - << seq_idx; - throw std::out_of_range(oss.str()); - } const auto available = page_table[seq_idx].size(); if (requested > available) { std::ostringstream oss; @@ -634,180 +795,137 @@ class HostPagedKVWorkerView : private LayerMapper { << " for sequence " << sequence_vector[seq_idx]; throw std::out_of_range(oss.str()); } - sequence_offsets[seq_idx] = total_pages; page_table[seq_idx].resize(requested); - total_pages += requested; } - if (total_pages == 0) { - return LaunchAsyncTask([] {}); - } + return LaunchHostPageTableLoadToDevice( + std::move(page_table), page_counts, std::move(validated_k_ptrs), + std::move(validated_v_ptrs), kOpName, prep_start); + } - auto flattened_k_ptrs = FlattenActivePointerTensor( - validated_k_ptrs, sequence_offsets, page_counts, total_pages, - "k_device_ptrs", kOpName); + KVLayeredAsyncTask AsyncLoadPrefixPagesToDevice( + torch::Tensor host_page_ids, torch::Tensor active_page_counts, + torch::Tensor k_device_ptrs, + std::optional v_device_ptrs = std::nullopt) { + EnsureDeviceReady(); + constexpr std::string_view kOpName = + "AsyncLoadPrefixPagesToDevice"; - std::optional flattened_v_ptrs; - if (validated_v_ptrs.has_value()) { - flattened_v_ptrs = FlattenActivePointerTensor( - *validated_v_ptrs, sequence_offsets, page_counts, - total_pages, "v_device_ptrs", kOpName); - } + auto validated_host_pages = ValidatePageIdTensor2D( + std::move(host_page_ids), "host_page_ids", kOpName); + const auto batch_size = + static_cast(validated_host_pages.size(0)); - const std::size_t num_layers = config_.num_layers; - const std::size_t copy_entries = num_layers * total_pages; - if (copy_entries == 0) { - return LaunchAsyncTask([] {}); - } - const auto kernel_limit = - static_cast(std::numeric_limits::max()); - if (copy_entries > kernel_limit) { - std::ostringstream oss; - oss << kOpName << ": num_pages=" << copy_entries - << " exceeds kernel limit=" << kernel_limit; - throw std::invalid_argument(oss.str()); - } + auto validated_counts = ValidatePageCountTensor( + std::move(active_page_counts), batch_size, kOpName); - const auto prep_end = std::chrono::high_resolution_clock::now(); - const double prep_ms = - std::chrono::duration_cast< - std::chrono::duration>(prep_end - - prep_start) - .count(); - logger_->debug( - "Prepared AsyncLoadLayerPagedKVToDevice (num_layers={}, total_pages={}, max_sequence_pages={}, prep_time_ms={:.3f})", - num_layers, total_pages, max_sequence_pages, prep_ms); + auto validated_k_ptrs = ValidatePointerTensor3D( + std::move(k_device_ptrs), "k_device_ptrs", batch_size, + kOpName); - return LaunchAsyncTask([ - this, - page_table = std::move(page_table), - sequence_offsets = std::move(sequence_offsets), - k_tensor = std::move(flattened_k_ptrs), - v_tensor = std::move(flattened_v_ptrs), - total_pages, - num_layers, - copy_entries, - kOpName - ]() mutable { - const auto start = std::chrono::high_resolution_clock::now(); - c10::cuda::OptionalCUDAGuard device_guard(device_index_); - const auto cuda_stream = CopyStream(CopyDirection::kHostToDevice); - const std::size_t k_page_bytes = layout_.KPageBytes(); - if (k_page_bytes == 0) { - return; + std::optional validated_v_ptrs; + if (v_device_ptrs.has_value()) { + if constexpr (!kHasVCache) { + throw std::invalid_argument(std::string(kOpName) + + ": V cache is disabled"); } - - auto* k_dest_ptr = k_tensor.template data_ptr(); - const std::int64_t* v_dest_ptr = - v_tensor.has_value() - ? v_tensor->data_ptr() - : nullptr; - const std::size_t row_stride = total_pages; - auto build_plan = [&](const std::int64_t* dest_ptrs, - auto&& host_ptr_provider) { - if (dest_ptrs == nullptr) { - throw std::invalid_argument(std::string(kOpName) + - ": null device pointers"); - } - return this->BuildPageCopyPlan( - page_table, sequence_offsets, num_layers, row_stride, - copy_entries, dest_ptrs, - std::forward( - host_ptr_provider), - kOpName); - }; - - const auto k_plan = build_plan( - k_dest_ptr, - [this](std::size_t layer_idx, std::int32_t page_idx) -> void* { - return this->KPhysicalPagePtr(layer_idx, page_idx); - }); - const auto plan_end = std::chrono::high_resolution_clock::now(); - const double plan_ms = - std::chrono::duration_cast< - std::chrono::duration>(plan_end - - start) - .count(); - logger_->debug( - "Built paged K copy plan (num_layers={}, total_pages={}, plan_time_ms={:.3f})", - num_layers, total_pages, plan_ms); - - std::optional v_plan; - if constexpr (kHasVCache) { - if (v_dest_ptr != nullptr) { - v_plan = build_plan( - v_dest_ptr, [this](std::size_t layer_idx, - std::int32_t page_idx) -> void* { - return this->template VPhysicalPagePtr<>( - layer_idx, page_idx); - }); - } + auto tensor = ValidatePointerTensor3D( + std::move(*v_device_ptrs), "v_device_ptrs", batch_size, + kOpName); + if (tensor.sizes() != validated_k_ptrs.sizes()) { + std::ostringstream oss; + oss << kOpName + << ": v_device_ptrs must match k_device_ptrs shape"; + throw std::invalid_argument(oss.str()); } + validated_v_ptrs = std::move(tensor); + } - worker_detail::DeviceBuffer k_device_src_ptrs( - copy_entries); - worker_detail::DeviceBuffer k_device_dst_ptrs( - copy_entries); - worker_detail::DeviceBuffer v_device_src_ptrs( - v_plan.has_value() ? copy_entries : 0); - worker_detail::DeviceBuffer v_device_dst_ptrs( - v_plan.has_value() ? copy_entries : 0); + if (batch_size == 0) { + return KVLayeredAsyncTask{}; + } - auto enqueue_plan = - [&](const PageCopyPlan& plan, - worker_detail::DeviceBuffer& dev_src_ptrs, - worker_detail::DeviceBuffer& dev_dst_ptrs, - std::size_t page_bytes) { - if (plan.host_sources.empty() || page_bytes == 0) { - return; - } - const std::size_t ptr_bytes = - plan.host_sources.size() * sizeof(uint8_t*); - EnqueueCopy( - reinterpret_cast( - plan.host_sources.data()), - reinterpret_cast(dev_src_ptrs.get()), - ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); - EnqueueCopy( - reinterpret_cast( - plan.device_dests.data()), - reinterpret_cast(dev_dst_ptrs.get()), - ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); - worker_detail::LaunchUvaPageCopyKernel( - dev_src_ptrs.get(), dev_dst_ptrs.get(), page_bytes, - static_cast(plan.host_sources.size()), - cuda_stream); - }; + const auto prep_start = std::chrono::high_resolution_clock::now(); + auto page_counts = TensorToSizeVector( + validated_counts, "active_page_counts", kOpName); + auto page_table = TensorToPageTable(validated_host_pages, page_counts, + kOpName); - enqueue_plan(k_plan, k_device_src_ptrs, k_device_dst_ptrs, - k_page_bytes); + return LaunchHostPageTableLayeredLoadToDevice( + std::move(page_table), page_counts, std::move(validated_k_ptrs), + std::move(validated_v_ptrs), kOpName, prep_start); + } - if constexpr (kHasVCache) { - if (v_plan.has_value()) { - const std::size_t v_page_bytes = layout_.VPageBytes(); - enqueue_plan(*v_plan, v_device_src_ptrs, - v_device_dst_ptrs, v_page_bytes); - } - } + KVLayeredAsyncTask AsyncLoadPrefixLayersToDevice( + torch::Tensor host_page_ids, torch::Tensor active_page_counts, + torch::Tensor logical_layer_ids, torch::Tensor k_device_ptrs, + std::optional v_device_ptrs = std::nullopt) { + EnsureDeviceReady(); + constexpr std::string_view kOpName = + "AsyncLoadPrefixLayersToDevice"; - logger_->debug( - "AsyncLoadLayerPagedKVToDevice completed (num_layers={}, total_pages={}, k_page_bytes={})", - num_layers, total_pages, k_page_bytes); - this->SynchronizeWithEvent(cuda_stream); - }); - } + auto validated_host_pages = ValidatePageIdTensor2D( + std::move(host_page_ids), "host_page_ids", kOpName); + const auto batch_size = + static_cast(validated_host_pages.size(0)); - std::byte* DataBase() { return backend_.DataBase(); } - const std::byte* DataBase() const { return backend_.DataBase(); } + auto validated_counts = ValidatePageCountTensor( + std::move(active_page_counts), batch_size, kOpName); + auto validated_layers = ValidateCpuTensor1D( + std::move(logical_layer_ids), torch::kInt64, + "logical_layer_ids", kOpName); + const auto layer_ids = TensorToInt64Vector(validated_layers); + const auto layer_count = layer_ids.size(); - void* KPagePtr(std::size_t layer_idx, std::int32_t page_idx) { - const std::size_t physical_layer_idx = - ResolvePhysicalLayer(layer_idx, kClassTag); - return KPhysicalPagePtr(physical_layer_idx, page_idx); - } + auto validated_k_ptrs = ValidatePointerTensor3DWithLayerCount( + std::move(k_device_ptrs), "k_device_ptrs", batch_size, + layer_count, kOpName); - const void* KPagePtr(std::size_t layer_idx, std::int32_t page_idx) const { - const std::size_t physical_layer_idx = + std::optional validated_v_ptrs; + if (v_device_ptrs.has_value()) { + if constexpr (!kHasVCache) { + throw std::invalid_argument(std::string(kOpName) + + ": V cache is disabled"); + } + auto tensor = ValidatePointerTensor3DWithLayerCount( + std::move(*v_device_ptrs), "v_device_ptrs", batch_size, + layer_count, kOpName); + if (tensor.sizes() != validated_k_ptrs.sizes()) { + std::ostringstream oss; + oss << kOpName + << ": v_device_ptrs must match k_device_ptrs shape"; + throw std::invalid_argument(oss.str()); + } + validated_v_ptrs = std::move(tensor); + } + + if (batch_size == 0 || layer_count == 0) { + return KVLayeredAsyncTask{}; + } + + const auto prep_start = std::chrono::high_resolution_clock::now(); + auto page_counts = TensorToSizeVector( + validated_counts, "active_page_counts", kOpName); + auto page_table = TensorToPageTable(validated_host_pages, page_counts, + kOpName); + + return LaunchHostPageTableSelectedLayeredLoadToDevice( + std::move(page_table), page_counts, std::move(layer_ids), + std::move(validated_k_ptrs), std::move(validated_v_ptrs), + kOpName, prep_start); + } + + std::byte* DataBase() { return backend_.DataBase(); } + const std::byte* DataBase() const { return backend_.DataBase(); } + + void* KPagePtr(std::size_t layer_idx, std::int32_t page_idx) { + const std::size_t physical_layer_idx = + ResolvePhysicalLayer(layer_idx, kClassTag); + return KPhysicalPagePtr(physical_layer_idx, page_idx); + } + + const void* KPagePtr(std::size_t layer_idx, std::int32_t page_idx) const { + const std::size_t physical_layer_idx = ResolvePhysicalLayer(layer_idx, kClassTag); return KPhysicalPagePtr(physical_layer_idx, page_idx); } @@ -911,6 +1029,37 @@ class HostPagedKVWorkerView : private LayerMapper { return {std::move(k_ptrs), std::move(v_ptrs)}; } + std::pair, std::optional>> + GetSequenceLayerPageRangePointers(std::int64_t sequence_id, + std::size_t layer_idx, + std::size_t start_page, + std::size_t page_count) const { + const std::size_t physical_layer_idx = ResolvePhysicalLayer( + layer_idx, + "HostPagedKVWorkerView::GetSequenceLayerPageRangePointers"); + auto page_indices = + backend_.SequencePageRange(sequence_id, start_page, page_count); + std::vector k_ptrs; + k_ptrs.reserve(page_indices.size()); + std::optional> v_ptrs; + if constexpr (Layout::kHasVCache) { + v_ptrs.emplace(); + v_ptrs->reserve(page_indices.size()); + } + std::byte* base = const_cast(backend_.DataBase()); + for (std::int32_t page : page_indices) { + void* k_ptr = static_cast( + layout_.KPageAddress(base, physical_layer_idx, page)); + k_ptrs.emplace_back(k_ptr); + if constexpr (Layout::kHasVCache) { + void* v_ptr = static_cast( + layout_.VPageAddress(base, physical_layer_idx, page)); + v_ptrs->emplace_back(v_ptr); + } + } + return {std::move(k_ptrs), std::move(v_ptrs)}; + } + void RegisterSequences(const std::vector& sequence_ids) { if (sequence_ids.empty()) { return; @@ -991,126 +1140,67 @@ class HostPagedKVWorkerView : private LayerMapper { return released; } + std::vector RetainSequencePrefixPages( + std::int64_t sequence_id, std::size_t num_pages) { + return RetainSequencePageRange(sequence_id, 0, num_pages); + } + + std::vector RetainSequencePageRange( + std::int64_t sequence_id, std::size_t start_page, + std::size_t num_pages) { + if (num_pages == 0) { + return {}; + } + EnsureSequenceRegistered(sequence_id); + const auto current_pages = page_table_.Pages(sequence_id); + if (start_page > current_pages.size() || + num_pages > current_pages.size() - start_page) { + std::ostringstream oss; + oss << "RetainSequencePageRange: cannot retain " << num_pages + << " pages from offset " << start_page << " for sequence " + << sequence_id << " with only " << current_pages.size() + << " pages in the worker page table"; + throw std::out_of_range(oss.str()); + } + return backend_.RetainSequencePageRange(sequence_id, start_page, + num_pages); + } + + std::vector RetainSequencePages( + std::int64_t sequence_id, + const std::vector& page_ids) { + if (page_ids.empty()) { + return {}; + } + EnsureSequenceRegistered(sequence_id); + return backend_.RetainSequencePages(sequence_id, page_ids); + } + + void ReleaseResidentPages(const std::vector& page_ids) { + backend_.ReleaseResidentPages(page_ids); + } + KVAsyncTask AsyncOffloadLayerKVToHost( std::size_t layer_idx, std::vector sequence_ids, torch::Tensor k_tensor, std::optional v_tensor, // [B, S, H, D] SequenceLengths sequence_lengths) { - const std::size_t physical_layer_idx = - ResolvePhysicalLayer(layer_idx, "AsyncOffloadLayerKVToHost"); - EnsureDeviceReady(); const std::size_t batch = sequence_ids.size(); - if (batch == 0) { - return LaunchAsyncTask([] {}); - } - const std::size_t tokens_per_sequence = - ValidateKTensorShape(k_tensor, batch); - ValidateSequenceLengthsInput(sequence_lengths, batch, - "AsyncOffloadLayerKVToHost"); - torch::Tensor prepared_k = k_tensor; - std::optional prepared_v; - if (v_tensor.has_value()) { - if constexpr (kHasVCache) { - ValidateVTensorShape(*v_tensor, batch, tokens_per_sequence); - prepared_v = *v_tensor; - } else { - throw std::invalid_argument( - "V tensor provided but V cache is disabled"); - } - } - c10::cuda::OptionalCUDAGuard producer_guard(device_index_); - const auto producer_cuda_stream = - at::cuda::getCurrentCUDAStream(device_index_).stream(); - - return LaunchAsyncTask([this, physical_layer_idx, - sequence_ids = std::move(sequence_ids), - sequence_lengths = std::move(sequence_lengths), - prepared_k, prepared_v, tokens_per_sequence, - producer_cuda_stream]() { - c10::cuda::OptionalCUDAGuard device_guard(device_index_); - const auto cuda_stream = CopyStream(CopyDirection::kDeviceToHost); - this->WaitForProducerStream(cuda_stream, producer_cuda_stream); - - const auto* k_base = - static_cast(prepared_k.data_ptr()); - const std::size_t k_token_bytes = geometry_.KTokenBytes(); - const std::size_t k_seq_stride = - tokens_per_sequence * k_token_bytes; - - const std::byte* v_base = nullptr; - std::size_t v_token_bytes = 0; - std::size_t v_seq_stride = 0; - if (prepared_v.has_value()) { - if constexpr (kHasVCache) { - v_base = - static_cast(prepared_v->data_ptr()); - v_token_bytes = - geometry_.template VTokenBytes(); - v_seq_stride = tokens_per_sequence * v_token_bytes; - } - } - - std::byte* host_base = backend_.DataBase(); - - for (std::size_t batch_idx = 0; batch_idx < sequence_ids.size(); - ++batch_idx) { - const std::int64_t sequence_id = sequence_ids[batch_idx]; - const auto pages = page_table_.Pages(sequence_id); - const std::size_t tokens_to_copy = ResolveSequenceLength( - sequence_lengths, batch_idx, sequence_id, - tokens_per_sequence, "AsyncOffloadLayerKVToHost"); - if (tokens_to_copy == 0) { - continue; - } - geometry_.ValidatePageCapacity(pages, tokens_to_copy, - "AsyncOffloadLayerKVToHost"); - - const auto* seq_k_src = k_base + batch_idx * k_seq_stride; - - ForEachPageChunk( - pages, 0, tokens_to_copy, - [&](std::int32_t page_idx, std::size_t page_offset_tokens, - std::size_t chunk_tokens, - std::size_t relative_token_offset) { - std::byte* dst = layout_.KPageAddress( - host_base, physical_layer_idx, - page_idx) + - page_offset_tokens * k_token_bytes; - const std::byte* src = - seq_k_src + relative_token_offset * k_token_bytes; - EnqueueCopy(src, dst, chunk_tokens * k_token_bytes, - CopyDirection::kDeviceToHost, cuda_stream); - }); - if constexpr (kHasVCache) { - if (v_base != nullptr) { - const auto* seq_v_src = - v_base + batch_idx * v_seq_stride; - ForEachPageChunk( - pages, 0, tokens_to_copy, - [&](std::int32_t page_idx, - std::size_t page_offset_tokens, - std::size_t chunk_tokens, - std::size_t relative_token_offset) { - std::byte* dst = - layout_.template VPageAddress<>( - host_base, physical_layer_idx, - page_idx) + - page_offset_tokens * v_token_bytes; - const std::byte* src = - seq_v_src + - relative_token_offset * v_token_bytes; - EnqueueCopy( - src, dst, chunk_tokens * v_token_bytes, - CopyDirection::kDeviceToHost, cuda_stream); - }); - } - } - } + SequenceLengthVector raw_start_positions(batch, 0); + return AsyncOffloadLayerKVRangeToHostImpl( + layer_idx, std::move(sequence_ids), std::move(k_tensor), + std::move(v_tensor), std::move(raw_start_positions), + std::move(sequence_lengths), "AsyncOffloadLayerKVToHost"); + } - this->SynchronizeWithEvent(cuda_stream); - // LogFirstTokenPerPage(layer_idx, sequence_ids, sequence_lengths, - // tokens_per_sequence, host_base); - }); + KVAsyncTask AsyncOffloadLayerKVRangeToHost( + std::size_t layer_idx, std::vector sequence_ids, + torch::Tensor k_tensor, std::optional v_tensor, + SequenceLengths raw_start_positions, SequenceLengths token_counts) { + return AsyncOffloadLayerKVRangeToHostImpl( + layer_idx, std::move(sequence_ids), std::move(k_tensor), + std::move(v_tensor), std::move(raw_start_positions), + std::move(token_counts), "AsyncOffloadLayerKVRangeToHost"); } KVAsyncTask AsyncAppendDecodeKVToHost( @@ -1556,6 +1646,144 @@ class HostPagedKVWorkerView : private LayerMapper { static inline constexpr std::string_view kClassTag = "HostPagedKVWorkerView"; + KVAsyncTask AsyncOffloadLayerKVRangeToHostImpl( + std::size_t layer_idx, std::vector sequence_ids, + torch::Tensor k_tensor, std::optional v_tensor, + SequenceLengths raw_start_positions, SequenceLengths token_counts, + std::string_view op_name) { + const std::size_t physical_layer_idx = + ResolvePhysicalLayer(layer_idx, op_name); + EnsureDeviceReady(); + const std::size_t batch = sequence_ids.size(); + if (batch == 0) { + return LaunchAsyncTask([] {}); + } + const std::size_t tokens_per_sequence = + ValidateKTensorShape(k_tensor, batch); + ValidateSequenceLengthsInput(raw_start_positions, batch, op_name); + ValidateSequenceLengthsInput(token_counts, batch, op_name); + torch::Tensor prepared_k = std::move(k_tensor); + std::optional prepared_v; + if (v_tensor.has_value()) { + if constexpr (kHasVCache) { + ValidateVTensorShape(*v_tensor, batch, tokens_per_sequence); + prepared_v = std::move(*v_tensor); + } else { + throw std::invalid_argument( + "V tensor provided but V cache is disabled"); + } + } + c10::cuda::OptionalCUDAGuard producer_guard(device_index_); + const auto producer_cuda_stream = + at::cuda::getCurrentCUDAStream(device_index_).stream(); + std::string op_name_string(op_name); + + return LaunchAsyncTask( + [this, physical_layer_idx, sequence_ids = std::move(sequence_ids), + raw_start_positions = std::move(raw_start_positions), + token_counts = std::move(token_counts), prepared_k, prepared_v, + tokens_per_sequence, producer_cuda_stream, + op_name = std::move(op_name_string)]() { + c10::cuda::OptionalCUDAGuard device_guard(device_index_); + const auto cuda_stream = + CopyStream(CopyDirection::kDeviceToHost); + this->WaitForProducerStream(cuda_stream, producer_cuda_stream); + + const auto* k_base = + static_cast(prepared_k.data_ptr()); + const std::size_t k_token_bytes = geometry_.KTokenBytes(); + const std::size_t k_seq_stride = + tokens_per_sequence * k_token_bytes; + + const std::byte* v_base = nullptr; + std::size_t v_token_bytes = 0; + std::size_t v_seq_stride = 0; + if (prepared_v.has_value()) { + if constexpr (kHasVCache) { + v_base = static_cast( + prepared_v->data_ptr()); + v_token_bytes = + geometry_.template VTokenBytes(); + v_seq_stride = tokens_per_sequence * v_token_bytes; + } + } + + std::byte* host_base = backend_.DataBase(); + for (std::size_t batch_idx = 0; + batch_idx < sequence_ids.size(); ++batch_idx) { + const std::int64_t sequence_id = sequence_ids[batch_idx]; + const auto pages = page_table_.Pages(sequence_id); + const std::size_t raw_start = ResolveSequenceLength( + raw_start_positions, batch_idx, sequence_id, + std::nullopt, op_name); + const std::size_t tokens_to_copy = ResolveSequenceLength( + token_counts, batch_idx, sequence_id, + tokens_per_sequence, op_name); + if (tokens_to_copy == 0) { + continue; + } + if (raw_start > + std::numeric_limits::max() - + tokens_to_copy) { + std::ostringstream oss; + oss << op_name << ": raw range overflows size_t"; + throw std::out_of_range(oss.str()); + } + geometry_.ValidatePageCapacity( + pages, raw_start + tokens_to_copy, op_name); + + const auto* seq_k_src = k_base + batch_idx * k_seq_stride; + ForEachPageChunk( + pages, raw_start, tokens_to_copy, + [&](std::int32_t page_idx, + std::size_t page_offset_tokens, + std::size_t chunk_tokens, + std::size_t relative_token_offset) { + std::byte* dst = + layout_.KPageAddress(host_base, + physical_layer_idx, + page_idx) + + page_offset_tokens * k_token_bytes; + const std::byte* src = + seq_k_src + + relative_token_offset * k_token_bytes; + EnqueueCopy(src, dst, + chunk_tokens * k_token_bytes, + CopyDirection::kDeviceToHost, + cuda_stream); + }); + if constexpr (kHasVCache) { + if (v_base != nullptr) { + const auto* seq_v_src = + v_base + batch_idx * v_seq_stride; + ForEachPageChunk( + pages, raw_start, tokens_to_copy, + [&](std::int32_t page_idx, + std::size_t page_offset_tokens, + std::size_t chunk_tokens, + std::size_t relative_token_offset) { + std::byte* dst = + layout_.template VPageAddress<>( + host_base, physical_layer_idx, + page_idx) + + page_offset_tokens * v_token_bytes; + const std::byte* src = + seq_v_src + + relative_token_offset * v_token_bytes; + EnqueueCopy( + src, dst, + chunk_tokens * v_token_bytes, + CopyDirection::kDeviceToHost, + cuda_stream); + }); + } + } + } + + this->SynchronizeWithEvent(cuda_stream); + }); + } + void* KPhysicalPagePtr(std::size_t physical_layer_idx, std::int32_t page_idx) { geometry_.EnsureLayerBounds(physical_layer_idx, kClassTag); @@ -2063,36 +2291,684 @@ class HostPagedKVWorkerView : private LayerMapper { return plan; } - torch::Tensor ValidateCpuTensor1D(torch::Tensor tensor, - torch::ScalarType dtype, - std::string_view tensor_name, - std::string_view op_name) const { - if (tensor.device().type() != torch::kCPU) { - std::ostringstream oss; - oss << op_name << ": " << tensor_name - << " must be on CPU (got device " << tensor.device().str() - << ")"; - throw std::invalid_argument(oss.str()); + KVAsyncTask LaunchHostPageTableLoadToDevice( + std::vector> page_table, + const std::vector& page_counts, + torch::Tensor k_device_ptrs, + std::optional v_device_ptrs, + std::string_view op_name, + std::chrono::high_resolution_clock::time_point prep_start) { + const std::string op_name_text(op_name); + const auto batch_size = page_table.size(); + if (batch_size == 0) { + return LaunchAsyncTask([] {}); } - if (tensor.dim() != 1) { - std::ostringstream oss; - oss << op_name << ": " << tensor_name << " must be 1-D"; - throw std::invalid_argument(oss.str()); + if (page_counts.size() != batch_size) { + throw std::logic_error(op_name_text + ": page_counts size mismatch"); } - if (tensor.scalar_type() != dtype) { - std::ostringstream oss; - oss << op_name << ": " << tensor_name << " must have dtype " - << c10::toString(dtype) << " (got " - << c10::toString(tensor.scalar_type()) << ")"; - throw std::invalid_argument(oss.str()); + + const auto max_sequence_pages = + static_cast(k_device_ptrs.size(2)); + std::vector sequence_offsets(batch_size, 0); + std::size_t total_pages = 0; + for (std::size_t seq_idx = 0; seq_idx < batch_size; ++seq_idx) { + const std::size_t requested = page_counts[seq_idx]; + if (requested > max_sequence_pages) { + std::ostringstream oss; + oss << op_name_text << ": requested pages " << requested + << " exceed provided pointer tensor capacity " + << max_sequence_pages << " for sequence index " + << seq_idx; + throw std::out_of_range(oss.str()); + } + if (requested > page_table[seq_idx].size()) { + std::ostringstream oss; + oss << op_name_text << ": requested pages " << requested + << " exceed resolved host pages " + << page_table[seq_idx].size() + << " for sequence index " << seq_idx; + throw std::out_of_range(oss.str()); + } + page_table[seq_idx].resize(requested); + sequence_offsets[seq_idx] = total_pages; + total_pages += requested; } - if (!tensor.is_contiguous()) { + + if (total_pages == 0) { + return LaunchAsyncTask([] {}); + } + + auto flattened_k_ptrs = FlattenActivePointerTensor( + k_device_ptrs, sequence_offsets, page_counts, total_pages, + "k_device_ptrs", op_name_text); + + std::optional flattened_v_ptrs; + if (v_device_ptrs.has_value()) { + flattened_v_ptrs = FlattenActivePointerTensor( + *v_device_ptrs, sequence_offsets, page_counts, total_pages, + "v_device_ptrs", op_name_text); + } + + const std::size_t num_layers = config_.num_layers; + const std::size_t copy_entries = num_layers * total_pages; + if (copy_entries == 0) { + return LaunchAsyncTask([] {}); + } + const auto kernel_limit = + static_cast(std::numeric_limits::max()); + if (copy_entries > kernel_limit) { std::ostringstream oss; - oss << op_name << ": " << tensor_name << " must be contiguous"; + oss << op_name_text << ": num_pages=" << copy_entries + << " exceeds kernel limit=" << kernel_limit; throw std::invalid_argument(oss.str()); } - return tensor; - } + + const auto prep_end = std::chrono::high_resolution_clock::now(); + const double prep_ms = + std::chrono::duration_cast< + std::chrono::duration>(prep_end - + prep_start) + .count(); + logger_->debug( + "Prepared {} (num_layers={}, total_pages={}, max_sequence_pages={}, prep_time_ms={:.3f})", + op_name_text, num_layers, total_pages, max_sequence_pages, prep_ms); + + return LaunchAsyncTask([ + this, + page_table = std::move(page_table), + sequence_offsets = std::move(sequence_offsets), + k_tensor = std::move(flattened_k_ptrs), + v_tensor = std::move(flattened_v_ptrs), + total_pages, + num_layers, + copy_entries, + op_name = op_name_text + ]() mutable { + const auto start = std::chrono::high_resolution_clock::now(); + c10::cuda::OptionalCUDAGuard device_guard(device_index_); + const auto cuda_stream = CopyStream(CopyDirection::kHostToDevice); + const std::size_t k_page_bytes = layout_.KPageBytes(); + if (k_page_bytes == 0) { + return; + } + + auto* k_dest_ptr = k_tensor.template data_ptr(); + const std::int64_t* v_dest_ptr = + v_tensor.has_value() + ? v_tensor->data_ptr() + : nullptr; + const std::size_t row_stride = total_pages; + auto build_plan = [&](const std::int64_t* dest_ptrs, + auto&& host_ptr_provider) { + if (dest_ptrs == nullptr) { + throw std::invalid_argument(op_name + + ": null device pointers"); + } + return this->BuildPageCopyPlan( + page_table, sequence_offsets, num_layers, row_stride, + copy_entries, dest_ptrs, + std::forward( + host_ptr_provider), + op_name); + }; + + const auto k_plan = build_plan( + k_dest_ptr, + [this](std::size_t layer_idx, std::int32_t page_idx) -> void* { + return this->KPhysicalPagePtr(layer_idx, page_idx); + }); + const auto plan_end = std::chrono::high_resolution_clock::now(); + const double plan_ms = + std::chrono::duration_cast< + std::chrono::duration>(plan_end - + start) + .count(); + logger_->debug( + "Built {} K copy plan (num_layers={}, total_pages={}, plan_time_ms={:.3f})", + op_name, num_layers, total_pages, plan_ms); + + std::optional v_plan; + if constexpr (kHasVCache) { + if (v_dest_ptr != nullptr) { + v_plan = build_plan( + v_dest_ptr, [this](std::size_t layer_idx, + std::int32_t page_idx) -> void* { + return this->template VPhysicalPagePtr<>( + layer_idx, page_idx); + }); + } + } + + worker_detail::DeviceBuffer k_device_src_ptrs( + copy_entries); + worker_detail::DeviceBuffer k_device_dst_ptrs( + copy_entries); + worker_detail::DeviceBuffer v_device_src_ptrs( + v_plan.has_value() ? copy_entries : 0); + worker_detail::DeviceBuffer v_device_dst_ptrs( + v_plan.has_value() ? copy_entries : 0); + + auto enqueue_plan = + [&](const PageCopyPlan& plan, + worker_detail::DeviceBuffer& dev_src_ptrs, + worker_detail::DeviceBuffer& dev_dst_ptrs, + std::size_t page_bytes) { + if (plan.host_sources.empty() || page_bytes == 0) { + return; + } + const std::size_t ptr_bytes = + plan.host_sources.size() * sizeof(uint8_t*); + EnqueueCopy( + reinterpret_cast( + plan.host_sources.data()), + reinterpret_cast(dev_src_ptrs.get()), + ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); + EnqueueCopy( + reinterpret_cast( + plan.device_dests.data()), + reinterpret_cast(dev_dst_ptrs.get()), + ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); + worker_detail::LaunchUvaPageCopyKernel( + dev_src_ptrs.get(), dev_dst_ptrs.get(), page_bytes, + static_cast(plan.host_sources.size()), + cuda_stream); + }; + + enqueue_plan(k_plan, k_device_src_ptrs, k_device_dst_ptrs, + k_page_bytes); + + if constexpr (kHasVCache) { + if (v_plan.has_value()) { + const std::size_t v_page_bytes = layout_.VPageBytes(); + enqueue_plan(*v_plan, v_device_src_ptrs, + v_device_dst_ptrs, v_page_bytes); + } + } + + logger_->debug( + "{} completed (num_layers={}, total_pages={}, k_page_bytes={})", + op_name, num_layers, total_pages, k_page_bytes); + this->SynchronizeWithEvent(cuda_stream); + }); + } + + KVLayeredAsyncTask LaunchHostPageTableLayeredLoadToDevice( + std::vector> page_table, + const std::vector& page_counts, + torch::Tensor k_device_ptrs, + std::optional v_device_ptrs, + std::string_view op_name, + std::chrono::high_resolution_clock::time_point prep_start) { + const std::string op_name_text(op_name); + const auto batch_size = page_table.size(); + if (batch_size == 0) { + return KVLayeredAsyncTask{}; + } + if (page_counts.size() != batch_size) { + throw std::logic_error(op_name_text + ": page_counts size mismatch"); + } + + const auto max_sequence_pages = + static_cast(k_device_ptrs.size(2)); + std::vector sequence_offsets(batch_size, 0); + std::size_t total_pages = 0; + for (std::size_t seq_idx = 0; seq_idx < batch_size; ++seq_idx) { + const std::size_t requested = page_counts[seq_idx]; + if (requested > max_sequence_pages) { + std::ostringstream oss; + oss << op_name_text << ": requested pages " << requested + << " exceed provided pointer tensor capacity " + << max_sequence_pages << " for sequence index " + << seq_idx; + throw std::out_of_range(oss.str()); + } + if (requested > page_table[seq_idx].size()) { + std::ostringstream oss; + oss << op_name_text << ": requested pages " << requested + << " exceed resolved host pages " + << page_table[seq_idx].size() + << " for sequence index " << seq_idx; + throw std::out_of_range(oss.str()); + } + page_table[seq_idx].resize(requested); + sequence_offsets[seq_idx] = total_pages; + total_pages += requested; + } + + if (total_pages == 0) { + return KVLayeredAsyncTask{}; + } + + auto flattened_k_ptrs = FlattenActivePointerTensor( + k_device_ptrs, sequence_offsets, page_counts, total_pages, + "k_device_ptrs", op_name_text); + + std::optional flattened_v_ptrs; + if (v_device_ptrs.has_value()) { + flattened_v_ptrs = FlattenActivePointerTensor( + *v_device_ptrs, sequence_offsets, page_counts, total_pages, + "v_device_ptrs", op_name_text); + } + + const std::size_t num_layers = config_.num_layers; + const std::size_t copy_entries = num_layers * total_pages; + if (copy_entries == 0) { + return KVLayeredAsyncTask{}; + } + const auto kernel_limit = + static_cast(std::numeric_limits::max()); + if (total_pages > kernel_limit) { + std::ostringstream oss; + oss << op_name_text << ": total_pages=" << total_pages + << " exceeds kernel limit=" << kernel_limit; + throw std::invalid_argument(oss.str()); + } + + const auto prep_end = std::chrono::high_resolution_clock::now(); + const double prep_ms = + std::chrono::duration_cast< + std::chrono::duration>(prep_end - + prep_start) + .count(); + logger_->debug( + "Prepared layered {} (num_layers={}, total_pages={}, max_sequence_pages={}, prep_time_ms={:.3f})", + op_name_text, num_layers, total_pages, max_sequence_pages, prep_ms); + + c10::cuda::OptionalCUDAGuard device_guard(device_index_); + const auto cuda_stream = CopyStream(CopyDirection::kHostToDevice); + auto state = std::make_shared( + task_id_counter_.fetch_add(1, std::memory_order_relaxed) + 1, + device_index_, num_layers, + [this](std::size_t logical_layer_idx) -> std::size_t { + return this->ResolvePhysicalLayer( + logical_layer_idx, + "KVLayeredAsyncTask::wait_for_layer"); + }, + logger_); + state->h2d_stream = cuda_stream; + + const std::size_t k_page_bytes = layout_.KPageBytes(); + if (k_page_bytes == 0) { + return KVLayeredAsyncTask{std::move(state)}; + } + + for (auto& event : state->layer_events) { + CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); + } + CUDA_CHECK(cudaEventCreateWithFlags(&state->final_event, + cudaEventDisableTiming)); + + auto* k_dest_ptr = flattened_k_ptrs.template data_ptr(); + const std::int64_t* v_dest_ptr = + flattened_v_ptrs.has_value() + ? flattened_v_ptrs->data_ptr() + : nullptr; + const std::size_t row_stride = total_pages; + auto build_plan = [&](const std::int64_t* dest_ptrs, + auto&& host_ptr_provider) { + if (dest_ptrs == nullptr) { + throw std::invalid_argument(op_name_text + + ": null device pointers"); + } + return this->BuildPageCopyPlan( + page_table, sequence_offsets, num_layers, row_stride, + copy_entries, dest_ptrs, + std::forward(host_ptr_provider), + op_name_text); + }; + + const auto k_plan = build_plan( + k_dest_ptr, + [this](std::size_t layer_idx, std::int32_t page_idx) -> void* { + return this->KPhysicalPagePtr(layer_idx, page_idx); + }); + + std::optional v_plan; + if constexpr (kHasVCache) { + if (v_dest_ptr != nullptr) { + v_plan = build_plan( + v_dest_ptr, + [this](std::size_t layer_idx, + std::int32_t page_idx) -> void* { + return this->template VPhysicalPagePtr<>(layer_idx, + page_idx); + }); + } + } + + state->k_device_src_ptrs.Allocate(copy_entries); + state->k_device_dst_ptrs.Allocate(copy_entries); + if (v_plan.has_value()) { + state->v_device_src_ptrs.Allocate(copy_entries); + state->v_device_dst_ptrs.Allocate(copy_entries); + } + + auto enqueue_layer_plan = + [&](const PageCopyPlan& plan, + worker_detail::DeviceBuffer& dev_src_ptrs, + worker_detail::DeviceBuffer& dev_dst_ptrs, + std::size_t layer_idx, std::size_t page_bytes) { + if (page_bytes == 0) { + return; + } + const std::size_t layer_offset = layer_idx * total_pages; + const std::size_t ptr_bytes = + total_pages * sizeof(uint8_t*); + EnqueueCopy( + reinterpret_cast( + plan.host_sources.data() + layer_offset), + reinterpret_cast( + dev_src_ptrs.get() + layer_offset), + ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); + EnqueueCopy( + reinterpret_cast( + plan.device_dests.data() + layer_offset), + reinterpret_cast( + dev_dst_ptrs.get() + layer_offset), + ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); + worker_detail::LaunchUvaPageCopyKernel( + dev_src_ptrs.get() + layer_offset, + dev_dst_ptrs.get() + layer_offset, page_bytes, + static_cast(total_pages), cuda_stream); + state->has_enqueued_work = true; + }; + + for (std::size_t layer_idx = 0; layer_idx < num_layers; ++layer_idx) { + enqueue_layer_plan(k_plan, state->k_device_src_ptrs, + state->k_device_dst_ptrs, layer_idx, + k_page_bytes); + + if constexpr (kHasVCache) { + if (v_plan.has_value()) { + enqueue_layer_plan(*v_plan, state->v_device_src_ptrs, + state->v_device_dst_ptrs, layer_idx, + layout_.VPageBytes()); + } + } + + CUDA_CHECK(cudaEventRecord(state->layer_events[layer_idx], + cuda_stream)); + } + + CUDA_CHECK(cudaEventRecord(state->final_event, cuda_stream)); + state->final_event_recorded = true; + logger_->debug( + "{} layered enqueue complete (num_layers={}, total_pages={}, k_page_bytes={})", + op_name_text, num_layers, total_pages, k_page_bytes); + + return KVLayeredAsyncTask{std::move(state)}; + } + + KVLayeredAsyncTask LaunchHostPageTableSelectedLayeredLoadToDevice( + std::vector> page_table, + const std::vector& page_counts, + std::vector logical_layer_ids, + torch::Tensor k_device_ptrs, + std::optional v_device_ptrs, + std::string_view op_name, + std::chrono::high_resolution_clock::time_point prep_start) { + const std::string op_name_text(op_name); + const auto batch_size = page_table.size(); + const auto selected_layers = logical_layer_ids.size(); + if (batch_size == 0 || selected_layers == 0) { + return KVLayeredAsyncTask{}; + } + if (page_counts.size() != batch_size) { + throw std::logic_error(op_name_text + ": page_counts size mismatch"); + } + + const auto max_sequence_pages = + static_cast(k_device_ptrs.size(2)); + std::vector sequence_offsets(batch_size, 0); + std::size_t total_pages = 0; + for (std::size_t seq_idx = 0; seq_idx < batch_size; ++seq_idx) { + const std::size_t requested = page_counts[seq_idx]; + if (requested > max_sequence_pages) { + std::ostringstream oss; + oss << op_name_text << ": requested pages " << requested + << " exceed provided pointer tensor capacity " + << max_sequence_pages << " for sequence index " + << seq_idx; + throw std::out_of_range(oss.str()); + } + if (requested > page_table[seq_idx].size()) { + std::ostringstream oss; + oss << op_name_text << ": requested pages " << requested + << " exceed resolved host pages " + << page_table[seq_idx].size() + << " for sequence index " << seq_idx; + throw std::out_of_range(oss.str()); + } + page_table[seq_idx].resize(requested); + sequence_offsets[seq_idx] = total_pages; + total_pages += requested; + } + + if (total_pages == 0) { + return KVLayeredAsyncTask{}; + } + + std::vector host_physical_layers; + host_physical_layers.reserve(selected_layers); + for (std::int64_t logical_layer_id : logical_layer_ids) { + if (logical_layer_id < 0) { + std::ostringstream oss; + oss << op_name_text << ": logical_layer_ids must be " + << "non-negative, got " << logical_layer_id; + throw std::out_of_range(oss.str()); + } + host_physical_layers.push_back(ResolvePhysicalLayer( + static_cast(logical_layer_id), op_name_text)); + } + + auto flattened_k_ptrs = FlattenActivePointerTensor( + k_device_ptrs, sequence_offsets, page_counts, total_pages, + "k_device_ptrs", op_name_text); + + std::optional flattened_v_ptrs; + if (v_device_ptrs.has_value()) { + flattened_v_ptrs = FlattenActivePointerTensor( + *v_device_ptrs, sequence_offsets, page_counts, total_pages, + "v_device_ptrs", op_name_text); + } + + const std::size_t copy_entries = selected_layers * total_pages; + if (copy_entries == 0) { + return KVLayeredAsyncTask{}; + } + const auto kernel_limit = + static_cast(std::numeric_limits::max()); + if (total_pages > kernel_limit) { + std::ostringstream oss; + oss << op_name_text << ": total_pages=" << total_pages + << " exceeds kernel limit=" << kernel_limit; + throw std::invalid_argument(oss.str()); + } + + const auto prep_end = std::chrono::high_resolution_clock::now(); + const double prep_ms = + std::chrono::duration_cast< + std::chrono::duration>(prep_end - + prep_start) + .count(); + logger_->debug( + "Prepared selected-layer {} (selected_layers={}, total_pages={}, max_sequence_pages={}, prep_time_ms={:.3f})", + op_name_text, selected_layers, total_pages, max_sequence_pages, + prep_ms); + + c10::cuda::OptionalCUDAGuard device_guard(device_index_); + const auto cuda_stream = CopyStream(CopyDirection::kHostToDevice); + auto state = std::make_shared( + task_id_counter_.fetch_add(1, std::memory_order_relaxed) + 1, + device_index_, selected_layers, + [layers = logical_layer_ids]( + std::size_t logical_layer_idx) -> std::size_t { + const auto iter = std::find( + layers.begin(), layers.end(), + static_cast(logical_layer_idx)); + if (iter == layers.end()) { + std::ostringstream oss; + oss << "KVLayeredAsyncTask::wait_for_layer: layer " + << logical_layer_idx + << " is not part of this selected-layer load"; + throw std::out_of_range(oss.str()); + } + return static_cast( + std::distance(layers.begin(), iter)); + }, + logger_); + state->h2d_stream = cuda_stream; + + const std::size_t k_page_bytes = layout_.KPageBytes(); + if (k_page_bytes == 0) { + return KVLayeredAsyncTask{std::move(state)}; + } + + for (auto& event : state->layer_events) { + CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); + } + CUDA_CHECK(cudaEventCreateWithFlags(&state->final_event, + cudaEventDisableTiming)); + + auto* k_dest_ptr = flattened_k_ptrs.template data_ptr(); + const std::int64_t* v_dest_ptr = + flattened_v_ptrs.has_value() + ? flattened_v_ptrs->data_ptr() + : nullptr; + const std::size_t row_stride = total_pages; + auto build_plan = [&](const std::int64_t* dest_ptrs, + auto&& host_ptr_provider) { + if (dest_ptrs == nullptr) { + throw std::invalid_argument(op_name_text + + ": null device pointers"); + } + return this->BuildPageCopyPlan( + page_table, sequence_offsets, selected_layers, row_stride, + copy_entries, dest_ptrs, + std::forward(host_ptr_provider), + op_name_text); + }; + + const auto k_plan = build_plan( + k_dest_ptr, + [&host_physical_layers, this]( + std::size_t selected_layer_idx, + std::int32_t page_idx) -> void* { + return this->KPhysicalPagePtr( + host_physical_layers[selected_layer_idx], page_idx); + }); + + std::optional v_plan; + if constexpr (kHasVCache) { + if (v_dest_ptr != nullptr) { + v_plan = build_plan( + v_dest_ptr, + [&host_physical_layers, this]( + std::size_t selected_layer_idx, + std::int32_t page_idx) -> void* { + return this->template VPhysicalPagePtr<>( + host_physical_layers[selected_layer_idx], + page_idx); + }); + } + } + + state->k_device_src_ptrs.Allocate(copy_entries); + state->k_device_dst_ptrs.Allocate(copy_entries); + if (v_plan.has_value()) { + state->v_device_src_ptrs.Allocate(copy_entries); + state->v_device_dst_ptrs.Allocate(copy_entries); + } + + auto enqueue_layer_plan = + [&](const PageCopyPlan& plan, + worker_detail::DeviceBuffer& dev_src_ptrs, + worker_detail::DeviceBuffer& dev_dst_ptrs, + std::size_t selected_layer_idx, std::size_t page_bytes) { + if (page_bytes == 0) { + return; + } + const std::size_t layer_offset = + selected_layer_idx * total_pages; + const std::size_t ptr_bytes = + total_pages * sizeof(uint8_t*); + EnqueueCopy( + reinterpret_cast( + plan.host_sources.data() + layer_offset), + reinterpret_cast( + dev_src_ptrs.get() + layer_offset), + ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); + EnqueueCopy( + reinterpret_cast( + plan.device_dests.data() + layer_offset), + reinterpret_cast( + dev_dst_ptrs.get() + layer_offset), + ptr_bytes, CopyDirection::kHostToDevice, cuda_stream); + worker_detail::LaunchUvaPageCopyKernel( + dev_src_ptrs.get() + layer_offset, + dev_dst_ptrs.get() + layer_offset, page_bytes, + static_cast(total_pages), cuda_stream); + state->has_enqueued_work = true; + }; + + for (std::size_t selected_layer_idx = 0; + selected_layer_idx < selected_layers; ++selected_layer_idx) { + enqueue_layer_plan(k_plan, state->k_device_src_ptrs, + state->k_device_dst_ptrs, selected_layer_idx, + k_page_bytes); + + if constexpr (kHasVCache) { + if (v_plan.has_value()) { + enqueue_layer_plan(*v_plan, state->v_device_src_ptrs, + state->v_device_dst_ptrs, + selected_layer_idx, + layout_.VPageBytes()); + } + } + + CUDA_CHECK(cudaEventRecord(state->layer_events[selected_layer_idx], + cuda_stream)); + } + + CUDA_CHECK(cudaEventRecord(state->final_event, cuda_stream)); + state->final_event_recorded = true; + logger_->debug( + "{} selected-layer enqueue complete (selected_layers={}, total_pages={}, k_page_bytes={})", + op_name_text, selected_layers, total_pages, k_page_bytes); + + return KVLayeredAsyncTask{std::move(state)}; + } + + torch::Tensor ValidateCpuTensor1D(torch::Tensor tensor, + torch::ScalarType dtype, + std::string_view tensor_name, + std::string_view op_name) const { + if (tensor.device().type() != torch::kCPU) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must be on CPU (got device " << tensor.device().str() + << ")"; + throw std::invalid_argument(oss.str()); + } + if (tensor.dim() != 1) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name << " must be 1-D"; + throw std::invalid_argument(oss.str()); + } + if (tensor.scalar_type() != dtype) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name << " must have dtype " + << c10::toString(dtype) << " (got " + << c10::toString(tensor.scalar_type()) << ")"; + throw std::invalid_argument(oss.str()); + } + if (!tensor.is_contiguous()) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name << " must be contiguous"; + throw std::invalid_argument(oss.str()); + } + return tensor; + } torch::Tensor ValidatePointerMatrix(torch::Tensor tensor, std::string_view tensor_name, @@ -2177,6 +3053,87 @@ class HostPagedKVWorkerView : private LayerMapper { return tensor; } + torch::Tensor ValidatePointerTensor3DWithLayerCount( + torch::Tensor tensor, std::string_view tensor_name, + std::size_t expected_sequences, std::size_t expected_layers, + std::string_view op_name) const { + if (tensor.device().type() != torch::kCPU) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must reside on CPU (got " << tensor.device().str() + << ')'; + throw std::invalid_argument(oss.str()); + } + if (tensor.scalar_type() != torch::kInt64) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must have dtype int64 (got " + << c10::toString(tensor.scalar_type()) << ')'; + throw std::invalid_argument(oss.str()); + } + if (!tensor.is_contiguous()) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must be contiguous"; + throw std::invalid_argument(oss.str()); + } + if (tensor.dim() != 3) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must be 3-D (got dim=" << tensor.dim() << ')'; + throw std::invalid_argument(oss.str()); + } + if (tensor.size(0) != static_cast(expected_layers)) { + std::ostringstream oss; + oss << op_name << ": first dimension of " << tensor_name + << " must equal selected layer count (expected " + << expected_layers << ", got " << tensor.size(0) << ')'; + throw std::out_of_range(oss.str()); + } + if (tensor.size(1) != static_cast(expected_sequences)) { + std::ostringstream oss; + oss << op_name << ": second dimension of " << tensor_name + << " must equal sequence count (expected " + << expected_sequences << ", got " << tensor.size(1) + << ')'; + throw std::out_of_range(oss.str()); + } + return tensor; + } + + torch::Tensor ValidatePageIdTensor2D( + torch::Tensor tensor, std::string_view tensor_name, + std::string_view op_name) const { + if (tensor.device().type() != torch::kCPU) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must reside on CPU (got " << tensor.device().str() + << ')'; + throw std::invalid_argument(oss.str()); + } + if (tensor.scalar_type() != torch::kInt64 && + tensor.scalar_type() != torch::kInt32) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must have dtype int32 or int64 (got " + << c10::toString(tensor.scalar_type()) << ')'; + throw std::invalid_argument(oss.str()); + } + if (!tensor.is_contiguous()) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must be contiguous"; + throw std::invalid_argument(oss.str()); + } + if (tensor.dim() != 2) { + std::ostringstream oss; + oss << op_name << ": " << tensor_name + << " must be 2-D (got dim=" << tensor.dim() << ')'; + throw std::invalid_argument(oss.str()); + } + return tensor; + } + torch::Tensor ValidatePageCountTensor( torch::Tensor tensor, std::size_t expected_length, std::string_view op_name) const { @@ -2250,6 +3207,54 @@ class HostPagedKVWorkerView : private LayerMapper { return values; } + std::vector> TensorToPageTable( + const torch::Tensor& tensor, + const std::vector& page_counts, + std::string_view op_name) const { + const auto batch_size = static_cast(tensor.size(0)); + const auto max_pages = static_cast(tensor.size(1)); + if (batch_size != page_counts.size()) { + throw std::logic_error(std::string(op_name) + + ": page_counts size mismatch"); + } + std::vector> page_table(batch_size); + auto read_value = [&](std::size_t index) -> std::int64_t { + if (tensor.scalar_type() == torch::kInt64) { + return tensor.data_ptr()[index]; + } + return tensor.data_ptr()[index]; + }; + for (std::size_t seq_idx = 0; seq_idx < batch_size; ++seq_idx) { + const std::size_t count = page_counts[seq_idx]; + if (count > max_pages) { + std::ostringstream oss; + oss << op_name << ": host_page_ids lacks capacity for " + << "sequence index " << seq_idx; + throw std::out_of_range(oss.str()); + } + auto& pages = page_table[seq_idx]; + pages.reserve(count); + const std::size_t row_offset = seq_idx * max_pages; + for (std::size_t slot = 0; slot < count; ++slot) { + const std::int64_t value = read_value(row_offset + slot); + if (value < 0 || + value > + static_cast( + std::numeric_limits::max())) { + std::ostringstream oss; + oss << op_name << ": invalid host page id " << value + << " at sequence index " << seq_idx << " slot " + << slot; + throw std::out_of_range(oss.str()); + } + const auto page_idx = static_cast(value); + geometry_.EnsurePageBounds(page_idx, op_name); + pages.push_back(page_idx); + } + } + return page_table; + } + torch::Tensor FlattenActivePointerTensor( const torch::Tensor& tensor, const std::vector& sequence_offsets, diff --git a/core/KV_Storage/host_prefix_cache_coordinator.cpp b/core/KV_Storage/host_prefix_cache_coordinator.cpp new file mode 100644 index 000000000..e6a10f2e3 --- /dev/null +++ b/core/KV_Storage/host_prefix_cache_coordinator.cpp @@ -0,0 +1,2139 @@ +#include "host_prefix_cache_coordinator.h" + +#include "shared_memory_utils.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace batchgen::kv { + +namespace { + +constexpr std::uint64_t kPrefixCacheMagic = 0x484f535450434348ULL; +constexpr std::uint32_t kPrefixCacheAbiVersion = 4; +constexpr std::uint32_t kNodeIndexRebuildMinTombstones = 1024; +constexpr std::uint32_t kPageRefRebuildMinTombstones = 1024; +constexpr std::uint32_t kArenaCompactMinDeadGroupEntries = 1024; +constexpr std::uint32_t kArenaCompactMinDeadPageHandles = 4096; + +enum class EntryState : std::uint32_t { + kEmpty = 0, + kResident = 1, + kTombstone = 2, +}; + +enum class IndexSlotState : std::uint32_t { + kEmpty = 0, + kResident = 1, + kTombstone = 2, +}; + +struct SharedHeader { + std::atomic init_state{ + static_cast(SharedMemoryInitState::kUninitialized)}; + std::uint64_t magic = kPrefixCacheMagic; + std::uint32_t abi_version = kPrefixCacheAbiVersion; + std::uint64_t create_time_ns = 0; + + std::uint32_t group_count = 0; + std::uint32_t hash_block_tokens = 0; + std::uint32_t commit_boundary_tokens = 0; + std::uint32_t max_nodes = 0; + std::uint32_t max_node_index_slots = 0; + std::uint32_t max_page_ref_slots = 0; + std::uint32_t max_group_entries = 0; + std::uint32_t max_page_handles = 0; + + std::atomic next_group_entry{0}; + std::atomic next_page_handle{0}; + std::atomic global_epoch{0}; + std::atomic lookup_hits{0}; + std::atomic lookup_misses{0}; + std::atomic evicted_nodes{0}; + std::atomic eviction_protected_skips{0}; + pthread_mutex_t mutex{}; +}; + +struct SharedGroupSpec { + std::uint32_t group_id = 0; + std::uint32_t semantic = 0; + std::uint32_t required_for_reuse = 0; + std::uint32_t raw_page_tokens = 0; + std::uint32_t compression_ratio = 1; +}; + +struct SharedPrefixNode { + std::uint32_t state = static_cast(EntryState::kEmpty); + PrefixDigest namespace_digest{}; + PrefixDigest digest{}; + std::uint32_t raw_start_token = 0; + std::uint32_t raw_end_token = 0; + std::uint32_t first_group_entry = 0; + std::uint32_t group_entry_count = 0; + std::uint64_t last_access_epoch = 0; +}; + +struct SharedNodeIndexSlot { + std::uint32_t state = static_cast(IndexSlotState::kEmpty); + std::uint32_t node_index = 0; + PrefixDigest digest{}; +}; + +struct SharedGroupEntry { + std::uint32_t state = static_cast(EntryState::kEmpty); + std::uint32_t group_id = 0; + std::uint32_t raw_start_token = 0; + std::uint32_t raw_end_token = 0; + std::uint32_t first_page_handle = 0; + std::uint32_t page_handle_count = 0; + std::atomic active_ref_count{0}; + std::atomic pending_load_count{0}; +}; + +struct SharedPageHandle { + std::uint32_t page_id = 0; +}; + +struct SharedPageRefSlot { + std::uint32_t state = static_cast(IndexSlotState::kEmpty); + std::uint32_t group_id = 0; + std::uint32_t page_id = 0; + std::uint32_t ref_count = 0; +}; + +std::uint64_t NowNs() { + const auto now = std::chrono::steady_clock::now().time_since_epoch(); + return static_cast( + std::chrono::duration_cast(now).count()); +} + +std::uint64_t SplitMix64(std::uint64_t value) { + value += 0x9e3779b97f4a7c15ULL; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); +} + +std::uint32_t NextPowerOfTwo(std::uint32_t value) { + if (value <= 1) { + return 1; + } + --value; + value |= value >> 1; + value |= value >> 2; + value |= value >> 4; + value |= value >> 8; + value |= value >> 16; + return value + 1; +} + +std::uint64_t DigestHash(const PrefixDigest& digest) { + std::uint64_t value = 0x9e3779b97f4a7c15ULL; + for (std::size_t lane = 0; lane < digest.size(); ++lane) { + value ^= SplitMix64(digest[lane] + lane); + } + return SplitMix64(value); +} + +std::uint64_t PageRefHash(std::uint32_t group_id, std::uint32_t page_id) { + return SplitMix64((static_cast(group_id) << 32) | + static_cast(page_id)); +} + +PrefixDigest HashPrefixBlock(PrefixDigest namespace_digest, + PrefixDigest parent_digest, + const std::int64_t* tokens, + std::uint32_t token_count) { + PrefixDigest state{ + 0x243f6a8885a308d3ULL, + 0x13198a2e03707344ULL, + 0xa4093822299f31d0ULL, + 0x082efa98ec4e6c89ULL, + }; + for (std::size_t lane = 0; lane < state.size(); ++lane) { + state[lane] ^= SplitMix64(namespace_digest[lane]); + state[lane] ^= SplitMix64(parent_digest[lane] + lane); + } + for (std::uint32_t idx = 0; idx < token_count; ++idx) { + const auto token = static_cast(tokens[idx]); + const std::size_t lane = idx % state.size(); + state[lane] = SplitMix64(state[lane] ^ token ^ + (static_cast(idx) << 32)); + state[(lane + 1) % state.size()] ^= state[lane]; + } + for (std::size_t lane = 0; lane < state.size(); ++lane) { + state[lane] = SplitMix64(state[lane] ^ token_count ^ lane); + } + return state; +} + +bool DigestEquals(const PrefixDigest& lhs, const PrefixDigest& rhs) { + return lhs == rhs; +} + +void ResetGroupEntry(SharedGroupEntry& entry) { + entry.state = static_cast(EntryState::kEmpty); + entry.group_id = 0; + entry.raw_start_token = 0; + entry.raw_end_token = 0; + entry.first_page_handle = 0; + entry.page_handle_count = 0; + entry.active_ref_count.store(0, std::memory_order_relaxed); + entry.pending_load_count.store(0, std::memory_order_relaxed); +} + +std::uint32_t Gcd(std::uint32_t lhs, std::uint32_t rhs) { + return static_cast(std::gcd(lhs, rhs)); +} + +std::uint32_t Lcm(std::uint32_t lhs, std::uint32_t rhs) { + if (lhs == 0 || rhs == 0) { + return 0; + } + return static_cast(std::lcm(lhs, rhs)); +} + +std::map NormalizePageRequirements( + const std::vector& requirements) { + std::map result; + for (const GroupPageRequirement& requirement : requirements) { + if (requirement.min_pages == 0) { + continue; + } + result[requirement.group_id] += requirement.min_pages; + } + return result; +} + +bool HasEnoughReleasablePages( + const PrefixEvictionResult& result, + const std::map& requirements) { + if (requirements.empty()) { + return true; + } + std::map released_by_group; + for (const GroupCommitPages& group_pages : result.evicted_group_pages) { + released_by_group[group_pages.group_id] += + static_cast(group_pages.pages.size()); + } + for (const auto& [group_id, min_pages] : requirements) { + const auto iter = released_by_group.find(group_id); + const std::uint32_t released = + iter == released_by_group.end() ? 0 : iter->second; + if (released < min_pages) { + return false; + } + } + return true; +} + +void ValidateGroupSpec(const HostKVGroupSpec& spec) { + if (spec.raw_page_tokens == 0) { + throw std::invalid_argument( + "HostKVGroupSpec.raw_page_tokens must be > 0"); + } + if (spec.compression_ratio == 0) { + throw std::invalid_argument( + "HostKVGroupSpec.compression_ratio must be > 0"); + } +} + +std::uint32_t ComputeHashBlockTokens( + const std::vector& group_specs) { + std::uint32_t result = 0; + for (const auto& spec : group_specs) { + if (!spec.required_for_reuse) { + continue; + } + result = result == 0 ? spec.raw_page_tokens + : Gcd(result, spec.raw_page_tokens); + } + return result; +} + +std::uint32_t ComputeCommitBoundaryTokens( + const std::vector& group_specs) { + std::uint32_t result = 1; + bool has_required_group = false; + for (const auto& spec : group_specs) { + if (!spec.required_for_reuse) { + continue; + } + has_required_group = true; + result = Lcm(result, spec.raw_page_tokens); + } + return has_required_group ? result : 0; +} + +} // namespace + +std::string ToString(const HostKVGroupSpec& spec) { + std::ostringstream oss; + oss << "HostKVGroupSpec(group_id=" << spec.group_id + << ", semantic=" << static_cast(spec.semantic) + << ", required_for_reuse=" << spec.required_for_reuse + << ", raw_page_tokens=" << spec.raw_page_tokens + << ", compression_ratio=" << spec.compression_ratio << ")"; + return oss.str(); +} + +std::string ToString(const HostPrefixCacheStats& stats) { + std::ostringstream oss; + oss << "HostPrefixCacheStats(resident_nodes=" << stats.resident_nodes + << ", active_attachments=" << stats.active_attachments + << ", pending_load_entries=" << stats.pending_load_entries + << ", pending_load_refs=" << stats.pending_load_refs + << ", used_group_entries=" << stats.used_group_entries + << ", used_page_handles=" << stats.used_page_handles + << ", lookup_hits=" << stats.lookup_hits + << ", lookup_misses=" << stats.lookup_misses + << ", evicted_nodes=" << stats.evicted_nodes + << ", eviction_protected_skips=" << stats.eviction_protected_skips + << ")"; + return oss.str(); +} + +std::vector> BuildPrefixHashChain( + PrefixDigest namespace_digest, const std::vector& token_ids, + std::uint32_t block_tokens) { + if (block_tokens == 0) { + throw std::invalid_argument("block_tokens must be > 0"); + } + std::vector> chain; + const std::uint32_t full_tokens = + static_cast(token_ids.size()) - + (static_cast(token_ids.size()) % block_tokens); + chain.reserve(full_tokens / block_tokens); + PrefixDigest parent_digest{}; + for (std::uint32_t start = 0; start < full_tokens; start += block_tokens) { + parent_digest = HashPrefixBlock(namespace_digest, parent_digest, + token_ids.data() + start, block_tokens); + chain.emplace_back(start + block_tokens, parent_digest); + } + return chain; +} + +struct HostPrefixCacheCoordinator::SharedState { + struct LocalAttachment { + std::vector node_indices; + std::uint32_t pending_load_count = 0; + bool release_requested = false; + }; + struct ArenaUsage { + std::uint32_t group_entries = 0; + std::uint32_t page_handles = 0; + }; + + explicit SharedState(HostPrefixCacheConfig cfg, + std::uint32_t hash_block_tokens, + std::uint32_t commit_boundary_tokens) + : config(std::move(cfg)), + hash_block_tokens(hash_block_tokens), + commit_boundary_tokens(commit_boundary_tokens) { + ComputeOffsets(); + } + + void Initialize(bool create_region); + PrefixCommitResult CommitPrefixPages( + PrefixDigest namespace_digest, + const std::vector& token_ids, std::uint32_t commit_tokens, + const std::vector& group_pages); + PrefixLookupResult LookupAndAttach( + PrefixDigest namespace_digest, + const std::vector& token_ids); + PrefixLookupResult EstimateLookup( + PrefixDigest namespace_digest, + const std::vector& token_ids); + void ReleaseAttachment(std::uint64_t attachment_handle); + void BeginAttachmentLoad(std::uint64_t attachment_handle); + void EndAttachmentLoad(std::uint64_t attachment_handle); + PrefixEvictionResult EvictUntilFree(std::uint32_t min_free_nodes, + std::uint32_t min_free_group_entries, + std::uint32_t min_free_page_handles, + std::uint32_t max_scan_nodes); + PrefixEvictionResult EvictUntilReleasablePages( + const std::vector& requirements, + std::uint32_t max_scan_nodes); + PrefixEvictionResult ClearUnprotected(); + PrefixEvictionResult ClearNamespace(PrefixDigest namespace_digest); + HostPrefixCacheStats GetStats() const; + + HostPrefixCacheConfig config; + std::uint32_t hash_block_tokens = 0; + std::uint32_t commit_boundary_tokens = 0; + + int shm_fd = -1; + std::size_t total_bytes = 0; + std::byte* mapping = nullptr; + SharedHeader* header = nullptr; + SharedGroupSpec* group_specs = nullptr; + SharedPrefixNode* nodes = nullptr; + SharedNodeIndexSlot* node_index_slots = nullptr; + SharedPageRefSlot* page_ref_slots = nullptr; + SharedGroupEntry* group_entries = nullptr; + SharedPageHandle* page_handles = nullptr; + + std::size_t header_offset = 0; + std::size_t group_spec_offset = 0; + std::size_t node_offset = 0; + std::size_t node_index_offset = 0; + std::size_t page_ref_offset = 0; + std::size_t group_entry_offset = 0; + std::size_t page_handle_offset = 0; + std::size_t total_bytes_unaligned = 0; + + private: + mutable std::mutex local_attachment_mutex; + std::unordered_map local_attachments; + std::uint64_t next_local_attachment_handle = 1; + + void ComputeOffsets(); + void MapPointers(); + void ConstructSharedState(); + void WaitForInitialization() const; + void ValidateSharedState() const; + std::optional FindNodeLocked( + const PrefixDigest& digest) const; + std::uint32_t AllocateNodeLocked(); + bool NodeHasRequiredGroupsLocked(const SharedPrefixNode& node) const; + std::vector BuildMaterializationSpansLocked( + const std::vector& node_indices) const; + std::uint64_t AttachNodesLocked( + const std::vector& node_indices); + std::uint32_t CountFreeNodeSlotsLocked() const; + bool NodeIsProtectedLocked(const SharedPrefixNode& node) const; + void IncrementActiveRefsLocked( + const std::vector& node_indices); + void DecrementActiveRefsLocked( + const std::vector& node_indices); + void UpdateLoadRefsLocked(const std::vector& node_indices, + int delta); + void EvictNodeLocked(SharedPrefixNode* node, PrefixEvictionResult* result); + void AppendEvictedPagesLocked(const SharedPrefixNode& node, + PrefixEvictionResult* result); + void CompactArenasLocked(); + std::uint32_t MaxNodeIndexSlots() const; + std::optional FindNodeIndexSlotLocked( + const PrefixDigest& digest) const; + void InsertNodeIndexLocked(const PrefixDigest& digest, + std::uint32_t node_index); + void RemoveNodeIndexLocked(const PrefixDigest& digest); + void RebuildNodeIndexLocked(); + std::uint32_t CountNodeIndexTombstonesLocked() const; + std::uint32_t MaxPageRefSlots() const; + std::optional FindPageRefSlotLocked( + std::uint32_t group_id, std::uint32_t page_id) const; + SharedPageRefSlot& InsertOrGetPageRefSlotLocked(std::uint32_t group_id, + std::uint32_t page_id); + void IncrementPageRefLocked(std::uint32_t group_id, std::uint32_t page_id); + bool DecrementPageRefLocked(std::uint32_t group_id, std::uint32_t page_id); + void RebuildPageRefIndexLocked(); + std::uint32_t CountPageRefTombstonesLocked() const; + void RebuildPageRefIndexIfNeededLocked(); + ArenaUsage ResidentArenaUsageLocked() const; + bool TailArenaCapacityEnoughLocked(std::uint32_t min_group_entries, + std::uint32_t min_page_handles) const; + bool CompactedArenaCapacityEnoughLocked( + std::uint32_t min_group_entries, + std::uint32_t min_page_handles) const; + void CompactArenasForCapacityLocked(std::uint32_t min_group_entries, + std::uint32_t min_page_handles); + void CompactArenasAfterEvictionIfUsefulLocked( + std::uint32_t min_group_entries, std::uint32_t min_page_handles); + void RebuildNodeIndexIfNeededLocked(); +}; + +void HostPrefixCacheCoordinator::SharedState::ComputeOffsets() { + std::size_t offset = 0; + offset = AlignUp(offset, alignof(SharedHeader)); + header_offset = offset; + offset += sizeof(SharedHeader); + + offset = AlignUp(offset, alignof(SharedGroupSpec)); + group_spec_offset = offset; + offset += sizeof(SharedGroupSpec) * config.group_specs.size(); + + offset = AlignUp(offset, alignof(SharedPrefixNode)); + node_offset = offset; + offset += sizeof(SharedPrefixNode) * config.max_nodes; + + offset = AlignUp(offset, alignof(SharedNodeIndexSlot)); + node_index_offset = offset; + offset += sizeof(SharedNodeIndexSlot) * MaxNodeIndexSlots(); + + offset = AlignUp(offset, alignof(SharedPageRefSlot)); + page_ref_offset = offset; + offset += sizeof(SharedPageRefSlot) * MaxPageRefSlots(); + + offset = AlignUp(offset, alignof(SharedGroupEntry)); + group_entry_offset = offset; + offset += sizeof(SharedGroupEntry) * config.max_group_entries; + + offset = AlignUp(offset, alignof(SharedPageHandle)); + page_handle_offset = offset; + offset += sizeof(SharedPageHandle) * config.max_page_handles; + + total_bytes_unaligned = offset; +} + +void HostPrefixCacheCoordinator::SharedState::MapPointers() { + header = reinterpret_cast(mapping + header_offset); + group_specs = + reinterpret_cast(mapping + group_spec_offset); + nodes = reinterpret_cast(mapping + node_offset); + node_index_slots = reinterpret_cast( + mapping + node_index_offset); + page_ref_slots = + reinterpret_cast(mapping + page_ref_offset); + group_entries = + reinterpret_cast(mapping + group_entry_offset); + page_handles = + reinterpret_cast(mapping + page_handle_offset); +} + +void HostPrefixCacheCoordinator::SharedState::ConstructSharedState() { + std::memset(mapping, 0, total_bytes); + MapPointers(); + header->magic = kPrefixCacheMagic; + header->abi_version = kPrefixCacheAbiVersion; + header->create_time_ns = NowNs(); + header->group_count = static_cast(config.group_specs.size()); + header->hash_block_tokens = hash_block_tokens; + header->commit_boundary_tokens = commit_boundary_tokens; + header->max_nodes = config.max_nodes; + header->max_node_index_slots = MaxNodeIndexSlots(); + header->max_page_ref_slots = MaxPageRefSlots(); + header->max_group_entries = config.max_group_entries; + header->max_page_handles = config.max_page_handles; + header->next_group_entry.store(0, std::memory_order_relaxed); + header->next_page_handle.store(0, std::memory_order_relaxed); + header->global_epoch.store(0, std::memory_order_relaxed); + header->lookup_hits.store(0, std::memory_order_relaxed); + header->lookup_misses.store(0, std::memory_order_relaxed); + header->evicted_nodes.store(0, std::memory_order_relaxed); + header->eviction_protected_skips.store(0, std::memory_order_relaxed); + + for (std::size_t i = 0; i < config.group_specs.size(); ++i) { + const HostKVGroupSpec& spec = config.group_specs[i]; + group_specs[i].group_id = spec.group_id; + group_specs[i].semantic = static_cast(spec.semantic); + group_specs[i].required_for_reuse = spec.required_for_reuse ? 1 : 0; + group_specs[i].raw_page_tokens = spec.raw_page_tokens; + group_specs[i].compression_ratio = spec.compression_ratio; + } + + InitProcessSharedRobustMutex(&header->mutex, "pthread_mutex_init failed"); + header->init_state.store( + static_cast(SharedMemoryInitState::kReady), + std::memory_order_release); +} + +void HostPrefixCacheCoordinator::SharedState::WaitForInitialization() const { + while (true) { + const auto state = static_cast( + header->init_state.load(std::memory_order_acquire)); + if (state == SharedMemoryInitState::kReady) { + return; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } +} + +void HostPrefixCacheCoordinator::SharedState::ValidateSharedState() const { + if (header->magic != kPrefixCacheMagic) { + throw std::runtime_error( + "Host prefix cache shared memory magic mismatch"); + } + if (header->abi_version != kPrefixCacheAbiVersion) { + throw std::runtime_error("Host prefix cache ABI version mismatch"); + } + if (header->group_count != config.group_specs.size()) { + throw std::runtime_error("Host prefix cache group count mismatch"); + } + if (header->hash_block_tokens != hash_block_tokens || + header->commit_boundary_tokens != commit_boundary_tokens || + header->max_nodes != config.max_nodes || + header->max_node_index_slots != MaxNodeIndexSlots() || + header->max_page_ref_slots != MaxPageRefSlots() || + header->max_group_entries != config.max_group_entries || + header->max_page_handles != config.max_page_handles) { + throw std::runtime_error("Host prefix cache config mismatch"); + } + for (std::size_t i = 0; i < config.group_specs.size(); ++i) { + const HostKVGroupSpec& expected = config.group_specs[i]; + const SharedGroupSpec& actual = group_specs[i]; + if (actual.group_id != expected.group_id || + actual.semantic != static_cast(expected.semantic) || + actual.required_for_reuse != + (expected.required_for_reuse ? 1U : 0U) || + actual.raw_page_tokens != expected.raw_page_tokens || + actual.compression_ratio != expected.compression_ratio) { + throw std::runtime_error("Host prefix cache group spec mismatch"); + } + } +} + +void HostPrefixCacheCoordinator::SharedState::Initialize(bool create_region) { + const std::size_t page_size = SystemPageSize(); + total_bytes = AlignUp(total_bytes_unaligned, page_size); + int flags = O_RDWR; + if (create_region) { + flags |= O_CREAT; + } + shm_fd = shm_open(config.shm_name.c_str(), flags, 0660); + if (shm_fd == -1) { + throw std::system_error(errno, std::generic_category(), + "host prefix cache shm_open failed"); + } + if (create_region) { + if (ftruncate(shm_fd, static_cast(total_bytes)) == -1) { + const int err = errno; + close(shm_fd); + shm_fd = -1; + throw std::system_error(err, std::generic_category(), + "host prefix cache ftruncate failed"); + } + } else { + struct stat stat_buffer{}; + if (fstat(shm_fd, &stat_buffer) == -1) { + const int err = errno; + close(shm_fd); + shm_fd = -1; + throw std::system_error(err, std::generic_category(), + "host prefix cache fstat failed"); + } + if (static_cast(stat_buffer.st_size) < total_bytes) { + close(shm_fd); + shm_fd = -1; + throw std::runtime_error( + "host prefix cache shared memory segment is too small"); + } + } + + void* mapped = mmap(nullptr, total_bytes, PROT_READ | PROT_WRITE, + MAP_SHARED, shm_fd, 0); + if (mapped == MAP_FAILED) { + const int err = errno; + close(shm_fd); + shm_fd = -1; + throw std::system_error(err, std::generic_category(), + "host prefix cache mmap failed"); + } + mapping = static_cast(mapped); + MapPointers(); + + if (create_region) { + header->init_state.store( + static_cast(SharedMemoryInitState::kInitializing), + std::memory_order_relaxed); + ConstructSharedState(); + } else { + WaitForInitialization(); + ValidateSharedState(); + } +} + +std::uint32_t +HostPrefixCacheCoordinator::SharedState::MaxNodeIndexSlots() const { + if (config.max_nodes == 0) { + return 1; + } + return NextPowerOfTwo(config.max_nodes * 2); +} + +std::optional +HostPrefixCacheCoordinator::SharedState::FindNodeIndexSlotLocked( + const PrefixDigest& digest) const { + const std::uint32_t slot_count = MaxNodeIndexSlots(); + const std::uint32_t start = + static_cast(DigestHash(digest) & (slot_count - 1)); + for (std::uint32_t probe = 0; probe < slot_count; ++probe) { + const std::uint32_t slot_index = (start + probe) & (slot_count - 1); + const SharedNodeIndexSlot& slot = node_index_slots[slot_index]; + const auto state = static_cast(slot.state); + if (state == IndexSlotState::kEmpty) { + return std::nullopt; + } + if (state == IndexSlotState::kResident && + DigestEquals(slot.digest, digest)) { + return slot_index; + } + } + return std::nullopt; +} + +void HostPrefixCacheCoordinator::SharedState::InsertNodeIndexLocked( + const PrefixDigest& digest, std::uint32_t node_index) { + const std::uint32_t slot_count = MaxNodeIndexSlots(); + const std::uint32_t start = + static_cast(DigestHash(digest) & (slot_count - 1)); + std::optional first_tombstone; + for (std::uint32_t probe = 0; probe < slot_count; ++probe) { + const std::uint32_t slot_index = (start + probe) & (slot_count - 1); + SharedNodeIndexSlot& slot = node_index_slots[slot_index]; + const auto state = static_cast(slot.state); + if (state == IndexSlotState::kResident) { + if (DigestEquals(slot.digest, digest)) { + slot.node_index = node_index; + return; + } + continue; + } + if (state == IndexSlotState::kTombstone) { + if (!first_tombstone.has_value()) { + first_tombstone = slot_index; + } + continue; + } + const std::uint32_t target = + first_tombstone.has_value() ? first_tombstone.value() : slot_index; + SharedNodeIndexSlot& target_slot = node_index_slots[target]; + target_slot.state = + static_cast(IndexSlotState::kResident); + target_slot.node_index = node_index; + target_slot.digest = digest; + return; + } + + if (first_tombstone.has_value()) { + SharedNodeIndexSlot& slot = node_index_slots[first_tombstone.value()]; + slot.state = static_cast(IndexSlotState::kResident); + slot.node_index = node_index; + slot.digest = digest; + return; + } + throw std::runtime_error("Host prefix cache node index table is full"); +} + +void HostPrefixCacheCoordinator::SharedState::RemoveNodeIndexLocked( + const PrefixDigest& digest) { + const auto slot_index = FindNodeIndexSlotLocked(digest); + if (!slot_index.has_value()) { + return; + } + SharedNodeIndexSlot& slot = node_index_slots[slot_index.value()]; + slot.state = static_cast(IndexSlotState::kTombstone); + slot.node_index = 0; + slot.digest = PrefixDigest{}; +} + +void HostPrefixCacheCoordinator::SharedState::RebuildNodeIndexLocked() { + std::fill(node_index_slots, node_index_slots + MaxNodeIndexSlots(), + SharedNodeIndexSlot{}); + for (std::uint32_t node_index = 0; node_index < config.max_nodes; + ++node_index) { + const SharedPrefixNode& node = nodes[node_index]; + if (node.state == static_cast(EntryState::kResident)) { + InsertNodeIndexLocked(node.digest, node_index); + } + } +} + +std::uint32_t +HostPrefixCacheCoordinator::SharedState::CountNodeIndexTombstonesLocked() + const { + std::uint32_t tombstones = 0; + for (std::uint32_t slot_index = 0; slot_index < MaxNodeIndexSlots(); + ++slot_index) { + if (node_index_slots[slot_index].state == + static_cast(IndexSlotState::kTombstone)) { + ++tombstones; + } + } + return tombstones; +} + +std::uint32_t +HostPrefixCacheCoordinator::SharedState::MaxPageRefSlots() const { + if (config.max_page_handles == 0) { + return 1; + } + return NextPowerOfTwo(config.max_page_handles * 2); +} + +std::optional +HostPrefixCacheCoordinator::SharedState::FindPageRefSlotLocked( + std::uint32_t group_id, std::uint32_t page_id) const { + const std::uint32_t slot_count = MaxPageRefSlots(); + const std::uint32_t start = + static_cast(PageRefHash(group_id, page_id) & + (slot_count - 1)); + for (std::uint32_t probe = 0; probe < slot_count; ++probe) { + const std::uint32_t slot_index = (start + probe) & (slot_count - 1); + const SharedPageRefSlot& slot = page_ref_slots[slot_index]; + const auto state = static_cast(slot.state); + if (state == IndexSlotState::kEmpty) { + return std::nullopt; + } + if (state == IndexSlotState::kResident && + slot.group_id == group_id && slot.page_id == page_id) { + return slot_index; + } + } + return std::nullopt; +} + +SharedPageRefSlot& +HostPrefixCacheCoordinator::SharedState::InsertOrGetPageRefSlotLocked( + std::uint32_t group_id, std::uint32_t page_id) { + const std::uint32_t slot_count = MaxPageRefSlots(); + const std::uint32_t start = + static_cast(PageRefHash(group_id, page_id) & + (slot_count - 1)); + std::optional first_tombstone; + for (std::uint32_t probe = 0; probe < slot_count; ++probe) { + const std::uint32_t slot_index = (start + probe) & (slot_count - 1); + SharedPageRefSlot& slot = page_ref_slots[slot_index]; + const auto state = static_cast(slot.state); + if (state == IndexSlotState::kResident) { + if (slot.group_id == group_id && slot.page_id == page_id) { + return slot; + } + continue; + } + if (state == IndexSlotState::kTombstone) { + if (!first_tombstone.has_value()) { + first_tombstone = slot_index; + } + continue; + } + const std::uint32_t target = + first_tombstone.has_value() ? first_tombstone.value() : slot_index; + SharedPageRefSlot& target_slot = page_ref_slots[target]; + target_slot.state = + static_cast(IndexSlotState::kResident); + target_slot.group_id = group_id; + target_slot.page_id = page_id; + target_slot.ref_count = 0; + return target_slot; + } + + if (first_tombstone.has_value()) { + SharedPageRefSlot& slot = page_ref_slots[first_tombstone.value()]; + slot.state = static_cast(IndexSlotState::kResident); + slot.group_id = group_id; + slot.page_id = page_id; + slot.ref_count = 0; + return slot; + } + throw std::runtime_error("Host prefix cache page ref table is full"); +} + +void HostPrefixCacheCoordinator::SharedState::IncrementPageRefLocked( + std::uint32_t group_id, std::uint32_t page_id) { + SharedPageRefSlot& slot = InsertOrGetPageRefSlotLocked(group_id, page_id); + if (slot.ref_count == std::numeric_limits::max()) { + throw std::runtime_error("Host prefix cache page ref count overflow"); + } + ++slot.ref_count; +} + +bool HostPrefixCacheCoordinator::SharedState::DecrementPageRefLocked( + std::uint32_t group_id, std::uint32_t page_id) { + const auto slot_index = FindPageRefSlotLocked(group_id, page_id); + if (!slot_index.has_value()) { + throw std::runtime_error("Host prefix cache page ref is missing"); + } + SharedPageRefSlot& slot = page_ref_slots[slot_index.value()]; + if (slot.ref_count == 0) { + throw std::runtime_error("Host prefix cache page ref underflow"); + } + --slot.ref_count; + if (slot.ref_count != 0) { + return false; + } + slot.state = static_cast(IndexSlotState::kTombstone); + slot.group_id = 0; + slot.page_id = 0; + return true; +} + +void HostPrefixCacheCoordinator::SharedState::RebuildPageRefIndexLocked() { + std::vector refs; + refs.reserve(MaxPageRefSlots()); + for (std::uint32_t slot_index = 0; slot_index < MaxPageRefSlots(); + ++slot_index) { + const SharedPageRefSlot& slot = page_ref_slots[slot_index]; + if (slot.state == static_cast(IndexSlotState::kResident)) { + refs.push_back(slot); + } + } + std::fill(page_ref_slots, page_ref_slots + MaxPageRefSlots(), + SharedPageRefSlot{}); + for (const SharedPageRefSlot& ref : refs) { + SharedPageRefSlot& slot = + InsertOrGetPageRefSlotLocked(ref.group_id, ref.page_id); + slot.ref_count = ref.ref_count; + } +} + +std::uint32_t +HostPrefixCacheCoordinator::SharedState::CountPageRefTombstonesLocked() const { + std::uint32_t tombstones = 0; + for (std::uint32_t slot_index = 0; slot_index < MaxPageRefSlots(); + ++slot_index) { + if (page_ref_slots[slot_index].state == + static_cast(IndexSlotState::kTombstone)) { + ++tombstones; + } + } + return tombstones; +} + +void HostPrefixCacheCoordinator::SharedState:: + RebuildPageRefIndexIfNeededLocked() { + const std::uint32_t tombstones = CountPageRefTombstonesLocked(); + const std::uint32_t threshold = + std::max(kPageRefRebuildMinTombstones, MaxPageRefSlots() / 4); + if (tombstones >= threshold) { + RebuildPageRefIndexLocked(); + } +} + +HostPrefixCacheCoordinator::SharedState::ArenaUsage +HostPrefixCacheCoordinator::SharedState::ResidentArenaUsageLocked() const { + ArenaUsage usage; + for (std::uint32_t node_index = 0; node_index < config.max_nodes; + ++node_index) { + const SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + continue; + } + usage.group_entries += node.group_entry_count; + for (std::uint32_t offset = 0; offset < node.group_entry_count; + ++offset) { + const SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + if (entry.state == + static_cast(EntryState::kResident)) { + usage.page_handles += entry.page_handle_count; + } + } + } + return usage; +} + +bool HostPrefixCacheCoordinator::SharedState::TailArenaCapacityEnoughLocked( + std::uint32_t min_group_entries, + std::uint32_t min_page_handles) const { + const std::uint32_t next_group_entry = + header->next_group_entry.load(std::memory_order_relaxed); + const std::uint32_t next_page_handle = + header->next_page_handle.load(std::memory_order_relaxed); + return config.max_group_entries - next_group_entry >= min_group_entries && + config.max_page_handles - next_page_handle >= min_page_handles; +} + +bool HostPrefixCacheCoordinator::SharedState:: + CompactedArenaCapacityEnoughLocked(std::uint32_t min_group_entries, + std::uint32_t min_page_handles) const { + const ArenaUsage usage = ResidentArenaUsageLocked(); + return config.max_group_entries - usage.group_entries >= + min_group_entries && + config.max_page_handles - usage.page_handles >= min_page_handles; +} + +void HostPrefixCacheCoordinator::SharedState::CompactArenasForCapacityLocked( + std::uint32_t min_group_entries, std::uint32_t min_page_handles) { + if (TailArenaCapacityEnoughLocked(min_group_entries, min_page_handles)) { + return; + } + if (CompactedArenaCapacityEnoughLocked(min_group_entries, + min_page_handles)) { + CompactArenasLocked(); + } +} + +void HostPrefixCacheCoordinator::SharedState:: + RebuildNodeIndexIfNeededLocked() { + const std::uint32_t tombstones = CountNodeIndexTombstonesLocked(); + const std::uint32_t threshold = std::max( + kNodeIndexRebuildMinTombstones, MaxNodeIndexSlots() / 4); + if (tombstones >= threshold) { + RebuildNodeIndexLocked(); + } +} + +void HostPrefixCacheCoordinator::SharedState:: + CompactArenasAfterEvictionIfUsefulLocked( + std::uint32_t min_group_entries, std::uint32_t min_page_handles) { + if (!TailArenaCapacityEnoughLocked(min_group_entries, min_page_handles) && + CompactedArenaCapacityEnoughLocked(min_group_entries, + min_page_handles)) { + CompactArenasLocked(); + return; + } + + const ArenaUsage usage = ResidentArenaUsageLocked(); + const std::uint32_t next_group_entry = + header->next_group_entry.load(std::memory_order_relaxed); + const std::uint32_t next_page_handle = + header->next_page_handle.load(std::memory_order_relaxed); + const std::uint32_t dead_group_entries = + next_group_entry >= usage.group_entries + ? next_group_entry - usage.group_entries + : 0; + const std::uint32_t dead_page_handles = + next_page_handle >= usage.page_handles + ? next_page_handle - usage.page_handles + : 0; + const std::uint32_t group_threshold = std::max( + kArenaCompactMinDeadGroupEntries, config.max_group_entries / 4); + const std::uint32_t page_threshold = std::max( + kArenaCompactMinDeadPageHandles, config.max_page_handles / 4); + if (dead_group_entries >= group_threshold || + dead_page_handles >= page_threshold) { + CompactArenasLocked(); + return; + } + + RebuildNodeIndexIfNeededLocked(); + RebuildPageRefIndexIfNeededLocked(); +} + +std::optional +HostPrefixCacheCoordinator::SharedState::FindNodeLocked( + const PrefixDigest& digest) const { + const auto slot_index = FindNodeIndexSlotLocked(digest); + if (!slot_index.has_value()) { + return std::nullopt; + } + const SharedNodeIndexSlot& slot = node_index_slots[slot_index.value()]; + if (slot.node_index >= config.max_nodes) { + return std::nullopt; + } + const SharedPrefixNode& node = nodes[slot.node_index]; + if (node.state != static_cast(EntryState::kResident) || + !DigestEquals(node.digest, digest)) { + return std::nullopt; + } + return slot.node_index; +} + +std::uint32_t HostPrefixCacheCoordinator::SharedState::AllocateNodeLocked() { + for (std::uint32_t index = 0; index < config.max_nodes; ++index) { + SharedPrefixNode& node = nodes[index]; + if (node.state == static_cast(EntryState::kEmpty) || + node.state == static_cast(EntryState::kTombstone)) { + node = SharedPrefixNode(); + return index; + } + } + throw std::runtime_error("Host prefix cache node table is full"); +} + +bool HostPrefixCacheCoordinator::SharedState::NodeHasRequiredGroupsLocked( + const SharedPrefixNode& node) const { + for (std::size_t spec_idx = 0; spec_idx < config.group_specs.size(); + ++spec_idx) { + const SharedGroupSpec& spec = group_specs[spec_idx]; + if (spec.required_for_reuse == 0) { + continue; + } + bool found = false; + for (std::uint32_t offset = 0; offset < node.group_entry_count; + ++offset) { + const SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + if (entry.state == + static_cast(EntryState::kResident) && + entry.group_id == spec.group_id && + entry.raw_start_token <= node.raw_start_token && + entry.raw_end_token >= node.raw_end_token) { + found = true; + break; + } + } + if (!found) { + return false; + } + } + return true; +} + +std::vector +HostPrefixCacheCoordinator::SharedState::BuildMaterializationSpansLocked( + const std::vector& node_indices) const { + std::map spans_by_group; + for (std::uint32_t node_index : node_indices) { + if (node_index >= config.max_nodes) { + throw std::out_of_range("prefix cache node index out of range"); + } + const SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + throw std::runtime_error( + "prefix cache materialization refers to non-resident node"); + } + for (std::uint32_t offset = 0; offset < node.group_entry_count; + ++offset) { + const SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + if (entry.state != + static_cast(EntryState::kResident)) { + continue; + } + GroupMaterializationSpan& span = spans_by_group[entry.group_id]; + span.group_id = entry.group_id; + span.raw_end_token = std::max(span.raw_end_token, + entry.raw_end_token); + span.pages.reserve(span.pages.size() + entry.page_handle_count); + for (std::uint32_t page_idx = 0; + page_idx < entry.page_handle_count; ++page_idx) { + const SharedPageHandle& page = + page_handles[entry.first_page_handle + page_idx]; + span.pages.push_back({page.page_id}); + } + } + } + + std::vector spans; + spans.reserve(spans_by_group.size()); + for (std::size_t spec_idx = 0; spec_idx < config.group_specs.size(); + ++spec_idx) { + const std::uint32_t group_id = group_specs[spec_idx].group_id; + auto iter = spans_by_group.find(group_id); + if (iter == spans_by_group.end()) { + continue; + } + spans.emplace_back(std::move(iter->second)); + spans_by_group.erase(iter); + } + for (auto& [_, span] : spans_by_group) { + spans.emplace_back(std::move(span)); + } + return spans; +} + +std::uint64_t HostPrefixCacheCoordinator::SharedState::AttachNodesLocked( + const std::vector& node_indices) { + if (node_indices.empty()) { + throw std::invalid_argument( + "prefix cache attach needs at least one node"); + } + LocalAttachment attachment; + attachment.node_indices = node_indices; + const std::uint64_t epoch = + header->global_epoch.fetch_add(1, std::memory_order_relaxed) + 1; + for (std::uint32_t node_index : attachment.node_indices) { + if (node_index >= config.max_nodes) { + throw std::out_of_range("prefix cache node index out of range"); + } + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + throw std::runtime_error( + "prefix cache attach refers to non-resident node"); + } + node.last_access_epoch = epoch; + } + IncrementActiveRefsLocked(attachment.node_indices); + try { + std::lock_guard attachment_lock(local_attachment_mutex); + const std::uint64_t handle = next_local_attachment_handle++; + local_attachments.emplace(handle, std::move(attachment)); + return handle; + } catch (...) { + DecrementActiveRefsLocked(node_indices); + throw; + } +} + +std::uint32_t +HostPrefixCacheCoordinator::SharedState::CountFreeNodeSlotsLocked() const { + std::uint32_t free_slots = 0; + for (std::uint32_t index = 0; index < config.max_nodes; ++index) { + const SharedPrefixNode& node = nodes[index]; + if (node.state == static_cast(EntryState::kEmpty) || + node.state == static_cast(EntryState::kTombstone)) { + ++free_slots; + } + } + return free_slots; +} + +bool HostPrefixCacheCoordinator::SharedState::NodeIsProtectedLocked( + const SharedPrefixNode& node) const { + for (std::uint32_t offset = 0; offset < node.group_entry_count; ++offset) { + const SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + if (entry.state != static_cast(EntryState::kResident)) { + continue; + } + if (entry.active_ref_count.load(std::memory_order_relaxed) != 0 || + entry.pending_load_count.load(std::memory_order_relaxed) != 0) { + return true; + } + } + return false; +} + +void HostPrefixCacheCoordinator::SharedState::IncrementActiveRefsLocked( + const std::vector& node_indices) { + for (std::uint32_t node_index : node_indices) { + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + throw std::runtime_error( + "host prefix cache attachment refers to non-resident node"); + } + for (std::uint32_t offset = 0; offset < node.group_entry_count; + ++offset) { + SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + entry.active_ref_count.fetch_add(1, std::memory_order_relaxed); + } + } +} + +void HostPrefixCacheCoordinator::SharedState::DecrementActiveRefsLocked( + const std::vector& node_indices) { + for (std::uint32_t node_index : node_indices) { + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + continue; + } + for (std::uint32_t offset = 0; offset < node.group_entry_count; + ++offset) { + SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + const std::uint32_t refs = + entry.active_ref_count.load(std::memory_order_relaxed); + if (refs == 0) { + throw std::runtime_error( + "host prefix cache active attachment ref underflow"); + } + entry.active_ref_count.store(refs - 1, std::memory_order_relaxed); + } + } +} + +void HostPrefixCacheCoordinator::SharedState::UpdateLoadRefsLocked( + const std::vector& node_indices, int delta) { + for (std::uint32_t node_index : node_indices) { + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + throw std::runtime_error( + "host prefix cache attachment refers to non-resident node"); + } + for (std::uint32_t offset = 0; offset < node.group_entry_count; + ++offset) { + SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + const std::uint32_t pending = + entry.pending_load_count.load(std::memory_order_relaxed); + if (delta > 0) { + entry.pending_load_count.store(pending + 1, + std::memory_order_relaxed); + } else { + if (pending == 0) { + throw std::runtime_error( + "host prefix cache pending load ref underflow"); + } + entry.pending_load_count.store(pending - 1, + std::memory_order_relaxed); + } + } + } +} + +void HostPrefixCacheCoordinator::SharedState::EvictNodeLocked( + SharedPrefixNode* node, PrefixEvictionResult* result) { + AppendEvictedPagesLocked(*node, result); + RemoveNodeIndexLocked(node->digest); + result->freed_group_entries += node->group_entry_count; + for (std::uint32_t offset = 0; offset < node->group_entry_count; ++offset) { + const SharedGroupEntry& entry = + group_entries[node->first_group_entry + offset]; + if (entry.state == static_cast(EntryState::kResident)) { + result->freed_page_handles += entry.page_handle_count; + } + } + *node = SharedPrefixNode(); + node->state = static_cast(EntryState::kTombstone); + ++result->evicted_nodes; +} + +void HostPrefixCacheCoordinator::SharedState::AppendEvictedPagesLocked( + const SharedPrefixNode& node, PrefixEvictionResult* result) { + for (std::uint32_t offset = 0; offset < node.group_entry_count; ++offset) { + const SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + if (entry.state != static_cast(EntryState::kResident)) { + continue; + } + for (std::uint32_t page_idx = 0; page_idx < entry.page_handle_count; + ++page_idx) { + const SharedPageHandle& handle = + page_handles[entry.first_page_handle + page_idx]; + if (DecrementPageRefLocked(entry.group_id, handle.page_id)) { + auto iter = std::find_if( + result->evicted_group_pages.begin(), + result->evicted_group_pages.end(), + [&entry](const GroupCommitPages& group_pages) { + return group_pages.group_id == entry.group_id; + }); + if (iter == result->evicted_group_pages.end()) { + result->evicted_group_pages.push_back( + GroupCommitPages{entry.group_id, {}}); + iter = result->evicted_group_pages.end() - 1; + } + iter->pages.push_back({handle.page_id}); + } + } + } +} + +void HostPrefixCacheCoordinator::SharedState::CompactArenasLocked() { + struct GroupEntrySnapshot { + std::uint32_t group_id = 0; + std::uint32_t raw_start_token = 0; + std::uint32_t raw_end_token = 0; + std::uint32_t active_ref_count = 0; + std::uint32_t pending_load_count = 0; + std::vector pages; + }; + struct NodeSnapshot { + std::uint32_t node_index = 0; + PrefixDigest namespace_digest{}; + PrefixDigest digest{}; + std::uint32_t raw_start_token = 0; + std::uint32_t raw_end_token = 0; + std::uint64_t last_access_epoch = 0; + std::vector groups; + }; + + std::vector snapshots; + snapshots.reserve(config.max_nodes); + for (std::uint32_t node_index = 0; node_index < config.max_nodes; + ++node_index) { + const SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + continue; + } + NodeSnapshot snapshot; + snapshot.node_index = node_index; + snapshot.namespace_digest = node.namespace_digest; + snapshot.digest = node.digest; + snapshot.raw_start_token = node.raw_start_token; + snapshot.raw_end_token = node.raw_end_token; + snapshot.last_access_epoch = node.last_access_epoch; + snapshot.groups.reserve(node.group_entry_count); + for (std::uint32_t offset = 0; offset < node.group_entry_count; + ++offset) { + const SharedGroupEntry& entry = + group_entries[node.first_group_entry + offset]; + if (entry.state != + static_cast(EntryState::kResident)) { + continue; + } + GroupEntrySnapshot group; + group.group_id = entry.group_id; + group.raw_start_token = entry.raw_start_token; + group.raw_end_token = entry.raw_end_token; + group.active_ref_count = + entry.active_ref_count.load(std::memory_order_relaxed); + group.pending_load_count = + entry.pending_load_count.load(std::memory_order_relaxed); + group.pages.reserve(entry.page_handle_count); + for (std::uint32_t page_idx = 0; page_idx < entry.page_handle_count; + ++page_idx) { + group.pages.push_back( + page_handles[entry.first_page_handle + page_idx]); + } + snapshot.groups.emplace_back(std::move(group)); + } + snapshots.emplace_back(std::move(snapshot)); + } + + for (std::uint32_t index = 0; index < config.max_group_entries; ++index) { + ResetGroupEntry(group_entries[index]); + } + std::fill(page_handles, page_handles + config.max_page_handles, + SharedPageHandle{}); + + std::uint32_t next_group_entry = 0; + std::uint32_t next_page_handle = 0; + for (const NodeSnapshot& snapshot : snapshots) { + SharedPrefixNode& node = nodes[snapshot.node_index]; + node.state = static_cast(EntryState::kResident); + node.namespace_digest = snapshot.namespace_digest; + node.digest = snapshot.digest; + node.raw_start_token = snapshot.raw_start_token; + node.raw_end_token = snapshot.raw_end_token; + node.first_group_entry = next_group_entry; + node.group_entry_count = + static_cast(snapshot.groups.size()); + node.last_access_epoch = snapshot.last_access_epoch; + + for (const GroupEntrySnapshot& group : snapshot.groups) { + SharedGroupEntry& entry = group_entries[next_group_entry++]; + ResetGroupEntry(entry); + entry.state = static_cast(EntryState::kResident); + entry.group_id = group.group_id; + entry.raw_start_token = group.raw_start_token; + entry.raw_end_token = group.raw_end_token; + entry.first_page_handle = next_page_handle; + entry.page_handle_count = + static_cast(group.pages.size()); + entry.active_ref_count.store(group.active_ref_count, + std::memory_order_relaxed); + entry.pending_load_count.store(group.pending_load_count, + std::memory_order_relaxed); + for (const SharedPageHandle& page : group.pages) { + page_handles[next_page_handle++] = page; + } + } + } + header->next_group_entry.store(next_group_entry, std::memory_order_relaxed); + header->next_page_handle.store(next_page_handle, std::memory_order_relaxed); + RebuildNodeIndexLocked(); +} + +PrefixCommitResult HostPrefixCacheCoordinator::SharedState::CommitPrefixPages( + PrefixDigest namespace_digest, const std::vector& token_ids, + std::uint32_t commit_tokens, + const std::vector& group_pages) { + const std::uint32_t token_count = + static_cast(token_ids.size()); + commit_tokens = std::min(commit_tokens, token_count); + commit_tokens -= commit_tokens % commit_boundary_tokens; + if (commit_tokens == 0) { + return {}; + } + + std::unordered_map*> + pages_by_group; + pages_by_group.reserve(group_pages.size()); + for (const auto& pages : group_pages) { + pages_by_group[pages.group_id] = &pages.pages; + } + for (const auto& spec : config.group_specs) { + if (!spec.required_for_reuse) { + continue; + } + const auto iter = pages_by_group.find(spec.group_id); + if (iter == pages_by_group.end()) { + throw std::invalid_argument("missing pages for required group " + + std::to_string(spec.group_id)); + } + const std::uint32_t required_pages = + commit_tokens / spec.raw_page_tokens; + if (iter->second->size() < required_pages) { + throw std::invalid_argument( + "insufficient pages for required group " + + std::to_string(spec.group_id)); + } + } + + const auto chain = + BuildPrefixHashChain(namespace_digest, token_ids, hash_block_tokens); + PrefixCommitResult result; + result.committed_tokens = commit_tokens; + + ScopedPthreadMutexLock lock(&header->mutex); + std::uint32_t new_nodes_needed = 0; + std::uint32_t group_entries_needed = 0; + std::uint32_t page_handles_needed = 0; + std::uint32_t raw_start_token = 0; + for (const auto& [raw_end_token, digest] : chain) { + if (raw_end_token > commit_tokens || + raw_end_token % commit_boundary_tokens != 0) { + continue; + } + if (FindNodeLocked(digest).has_value()) { + raw_start_token = raw_end_token; + continue; + } + ++new_nodes_needed; + for (const auto& spec : config.group_specs) { + const auto iter = pages_by_group.find(spec.group_id); + if (iter == pages_by_group.end()) { + continue; + } + if (raw_start_token % spec.raw_page_tokens != 0 || + raw_end_token % spec.raw_page_tokens != 0) { + continue; + } + const std::uint32_t first_page = + raw_start_token / spec.raw_page_tokens; + const std::uint32_t pages_needed = + (raw_end_token - raw_start_token) / spec.raw_page_tokens; + if (iter->second->size() < first_page + pages_needed) { + continue; + } + ++group_entries_needed; + page_handles_needed += pages_needed; + } + raw_start_token = raw_end_token; + } + + std::uint32_t free_node_slots = 0; + for (std::uint32_t index = 0; index < config.max_nodes; ++index) { + const SharedPrefixNode& node = nodes[index]; + if (node.state == static_cast(EntryState::kEmpty) || + node.state == static_cast(EntryState::kTombstone)) { + ++free_node_slots; + } + } + if (free_node_slots < new_nodes_needed) { + throw std::runtime_error("Host prefix cache node table is full"); + } + CompactArenasForCapacityLocked(group_entries_needed, page_handles_needed); + const std::uint32_t first_group_entry = + header->next_group_entry.load(std::memory_order_relaxed); + const std::uint32_t first_page_handle = + header->next_page_handle.load(std::memory_order_relaxed); + if (first_group_entry + group_entries_needed > config.max_group_entries) { + throw std::runtime_error("Host prefix cache group entry table is full"); + } + if (first_page_handle + page_handles_needed > config.max_page_handles) { + throw std::runtime_error("Host prefix cache page handle arena is full"); + } + + raw_start_token = 0; + for (const auto& [raw_end_token, digest] : chain) { + if (raw_end_token > commit_tokens || + raw_end_token % commit_boundary_tokens != 0) { + continue; + } + if (FindNodeLocked(digest).has_value()) { + ++result.existing_nodes; + raw_start_token = raw_end_token; + continue; + } + + std::uint32_t group_entry_count = 0; + std::uint32_t page_handle_count = 0; + for (const auto& spec : config.group_specs) { + const auto iter = pages_by_group.find(spec.group_id); + if (iter == pages_by_group.end()) { + continue; + } + if (raw_start_token % spec.raw_page_tokens != 0 || + raw_end_token % spec.raw_page_tokens != 0) { + if (spec.required_for_reuse) { + throw std::runtime_error( + "required group is not aligned to raw page tokens"); + } + continue; + } + const std::uint32_t first_page = + raw_start_token / spec.raw_page_tokens; + const std::uint32_t pages_needed = + (raw_end_token - raw_start_token) / spec.raw_page_tokens; + if (iter->second->size() < first_page + pages_needed) { + if (spec.required_for_reuse) { + throw std::runtime_error( + "required group page list became too short"); + } + continue; + } + ++group_entry_count; + page_handle_count += pages_needed; + } + + const std::uint32_t first_group_entry = + header->next_group_entry.load(std::memory_order_relaxed); + const std::uint32_t first_page_handle = + header->next_page_handle.load(std::memory_order_relaxed); + if (first_group_entry + group_entry_count > config.max_group_entries) { + throw std::runtime_error( + "Host prefix cache group entry table is full"); + } + if (first_page_handle + page_handle_count > config.max_page_handles) { + throw std::runtime_error( + "Host prefix cache page handle arena is full"); + } + + std::uint32_t next_group_entry = first_group_entry; + std::uint32_t next_page_handle = first_page_handle; + for (const auto& spec : config.group_specs) { + const auto iter = pages_by_group.find(spec.group_id); + if (iter == pages_by_group.end() || + raw_start_token % spec.raw_page_tokens != 0 || + raw_end_token % spec.raw_page_tokens != 0) { + continue; + } + const std::uint32_t first_page = + raw_start_token / spec.raw_page_tokens; + const std::uint32_t pages_needed = + (raw_end_token - raw_start_token) / spec.raw_page_tokens; + if (iter->second->size() < first_page + pages_needed) { + continue; + } + + SharedGroupEntry& entry = group_entries[next_group_entry++]; + ResetGroupEntry(entry); + entry.state = static_cast(EntryState::kResident); + entry.group_id = spec.group_id; + entry.raw_start_token = raw_start_token; + entry.raw_end_token = raw_end_token; + entry.first_page_handle = next_page_handle; + entry.page_handle_count = pages_needed; + for (std::uint32_t page_idx = 0; page_idx < pages_needed; + ++page_idx) { + const HostPageHandle& handle = + (*iter->second)[first_page + page_idx]; + page_handles[next_page_handle++] = + SharedPageHandle{handle.page_id}; + IncrementPageRefLocked(spec.group_id, handle.page_id); + } + } + + const std::uint32_t node_index = AllocateNodeLocked(); + SharedPrefixNode& node = nodes[node_index]; + node.state = static_cast(EntryState::kResident); + node.namespace_digest = namespace_digest; + node.digest = digest; + node.raw_start_token = raw_start_token; + node.raw_end_token = raw_end_token; + node.first_group_entry = first_group_entry; + node.group_entry_count = group_entry_count; + node.last_access_epoch = + header->global_epoch.fetch_add(1, std::memory_order_relaxed) + 1; + header->next_group_entry.store(next_group_entry, + std::memory_order_relaxed); + header->next_page_handle.store(next_page_handle, + std::memory_order_relaxed); + InsertNodeIndexLocked(digest, node_index); + ++result.inserted_nodes; + raw_start_token = raw_end_token; + } + return result; +} + +PrefixLookupResult HostPrefixCacheCoordinator::SharedState::LookupAndAttach( + PrefixDigest namespace_digest, const std::vector& token_ids) { + const auto chain = + BuildPrefixHashChain(namespace_digest, token_ids, hash_block_tokens); + PrefixLookupResult result; + ScopedPthreadMutexLock lock(&header->mutex); + std::vector hit_node_indices; + for (const auto& [raw_end_token, digest] : chain) { + if (raw_end_token % commit_boundary_tokens != 0) { + continue; + } + const auto node_index = FindNodeLocked(digest); + if (!node_index.has_value()) { + break; + } + SharedPrefixNode& node = nodes[node_index.value()]; + if (!NodeHasRequiredGroupsLocked(node)) { + break; + } + hit_node_indices.push_back(node_index.value()); + result.common_cached_tokens = node.raw_end_token; + } + if (!hit_node_indices.empty()) { + result.attachment_handle = AttachNodesLocked(hit_node_indices); + result.materialization_spans = + BuildMaterializationSpansLocked(hit_node_indices); + header->lookup_hits.fetch_add(1, std::memory_order_relaxed); + return result; + } + result.miss_reason_mask = 1; + header->lookup_misses.fetch_add(1, std::memory_order_relaxed); + return result; +} + +PrefixLookupResult HostPrefixCacheCoordinator::SharedState::EstimateLookup( + PrefixDigest namespace_digest, const std::vector& token_ids) { + const auto chain = + BuildPrefixHashChain(namespace_digest, token_ids, hash_block_tokens); + PrefixLookupResult result; + ScopedPthreadMutexLock lock(&header->mutex); + std::vector hit_node_indices; + for (const auto& [raw_end_token, digest] : chain) { + if (raw_end_token % commit_boundary_tokens != 0) { + continue; + } + const auto node_index = FindNodeLocked(digest); + if (!node_index.has_value()) { + break; + } + const SharedPrefixNode& node = nodes[node_index.value()]; + if (!NodeHasRequiredGroupsLocked(node)) { + break; + } + hit_node_indices.push_back(node_index.value()); + result.common_cached_tokens = node.raw_end_token; + } + if (!hit_node_indices.empty()) { + result.materialization_spans = + BuildMaterializationSpansLocked(hit_node_indices); + return result; + } + result.miss_reason_mask = 1; + return result; +} + +void HostPrefixCacheCoordinator::SharedState::ReleaseAttachment( + std::uint64_t attachment_handle) { + if (attachment_handle == 0) { + return; + } + LocalAttachment attachment; + { + std::lock_guard attachment_lock(local_attachment_mutex); + auto iter = local_attachments.find(attachment_handle); + if (iter == local_attachments.end()) { + throw std::out_of_range( + "unknown host prefix cache attachment handle"); + } + if (iter->second.release_requested) { + throw std::runtime_error( + "host prefix cache attachment release was already requested"); + } + if (iter->second.pending_load_count != 0) { + iter->second.release_requested = true; + return; + } + attachment = std::move(iter->second); + local_attachments.erase(iter); + } + ScopedPthreadMutexLock lock(&header->mutex); + DecrementActiveRefsLocked(attachment.node_indices); +} + +void HostPrefixCacheCoordinator::SharedState::BeginAttachmentLoad( + std::uint64_t attachment_handle) { + if (attachment_handle == 0) { + throw std::invalid_argument( + "host prefix cache load attachment handle must be non-zero"); + } + std::vector node_indices; + { + std::lock_guard attachment_lock(local_attachment_mutex); + auto iter = local_attachments.find(attachment_handle); + if (iter == local_attachments.end()) { + throw std::out_of_range( + "unknown host prefix cache attachment handle"); + } + if (iter->second.release_requested) { + throw std::runtime_error( + "cannot begin load for a released host prefix cache attachment"); + } + ++iter->second.pending_load_count; + node_indices = iter->second.node_indices; + } + try { + ScopedPthreadMutexLock lock(&header->mutex); + UpdateLoadRefsLocked(node_indices, 1); + } catch (...) { + std::lock_guard attachment_lock(local_attachment_mutex); + auto iter = local_attachments.find(attachment_handle); + if (iter != local_attachments.end() && + iter->second.pending_load_count != 0) { + --iter->second.pending_load_count; + } + throw; + } +} + +void HostPrefixCacheCoordinator::SharedState::EndAttachmentLoad( + std::uint64_t attachment_handle) { + if (attachment_handle == 0) { + throw std::invalid_argument( + "host prefix cache load attachment handle must be non-zero"); + } + std::vector node_indices; + bool finalize_release = false; + { + std::lock_guard attachment_lock(local_attachment_mutex); + auto iter = local_attachments.find(attachment_handle); + if (iter == local_attachments.end()) { + throw std::out_of_range( + "unknown host prefix cache attachment handle"); + } + if (iter->second.pending_load_count == 0) { + throw std::runtime_error( + "host prefix cache attachment pending load underflow"); + } + --iter->second.pending_load_count; + node_indices = iter->second.node_indices; + finalize_release = iter->second.release_requested && + iter->second.pending_load_count == 0; + if (finalize_release) { + local_attachments.erase(iter); + } + } + ScopedPthreadMutexLock lock(&header->mutex); + UpdateLoadRefsLocked(node_indices, -1); + if (finalize_release) { + DecrementActiveRefsLocked(node_indices); + } +} + +PrefixEvictionResult HostPrefixCacheCoordinator::SharedState::EvictUntilFree( + std::uint32_t min_free_nodes, std::uint32_t min_free_group_entries, + std::uint32_t min_free_page_handles, std::uint32_t max_scan_nodes) { + PrefixEvictionResult result; + ScopedPthreadMutexLock lock(&header->mutex); + + const auto has_enough_free_capacity = [this, min_free_nodes, + min_free_group_entries, + min_free_page_handles]() { + const std::uint32_t free_nodes = CountFreeNodeSlotsLocked(); + return free_nodes >= min_free_nodes && TailArenaCapacityEnoughLocked( + min_free_group_entries, + min_free_page_handles); + }; + const auto can_satisfy_after_compact = [this, min_free_nodes, + min_free_group_entries, + min_free_page_handles]() { + const std::uint32_t free_nodes = CountFreeNodeSlotsLocked(); + return free_nodes >= min_free_nodes && + CompactedArenaCapacityEnoughLocked(min_free_group_entries, + min_free_page_handles); + }; + + if (has_enough_free_capacity()) { + return result; + } + CompactArenasForCapacityLocked(min_free_group_entries, + min_free_page_handles); + if (has_enough_free_capacity()) { + return result; + } + + std::vector candidates; + candidates.reserve(config.max_nodes); + for (std::uint32_t index = 0; index < config.max_nodes; ++index) { + if (nodes[index].state == + static_cast(EntryState::kResident)) { + candidates.push_back(index); + } + } + std::sort(candidates.begin(), candidates.end(), + [this](std::uint32_t lhs, std::uint32_t rhs) { + return nodes[lhs].last_access_epoch < + nodes[rhs].last_access_epoch; + }); + + std::uint32_t scanned = 0; + for (std::uint32_t node_index : candidates) { + if (max_scan_nodes != 0 && scanned >= max_scan_nodes) { + break; + } + ++scanned; + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + continue; + } + if (NodeIsProtectedLocked(node)) { + ++result.protected_nodes; + continue; + } + + EvictNodeLocked(&node, &result); + + if (has_enough_free_capacity() || can_satisfy_after_compact()) { + break; + } + } + + if (result.evicted_nodes != 0) { + CompactArenasAfterEvictionIfUsefulLocked(min_free_group_entries, + min_free_page_handles); + } + header->evicted_nodes.fetch_add(result.evicted_nodes, + std::memory_order_relaxed); + header->eviction_protected_skips.fetch_add(result.protected_nodes, + std::memory_order_relaxed); + return result; +} + +PrefixEvictionResult +HostPrefixCacheCoordinator::SharedState::EvictUntilReleasablePages( + const std::vector& requirements, + std::uint32_t max_scan_nodes) { + const auto required_pages = NormalizePageRequirements(requirements); + PrefixEvictionResult result; + if (required_pages.empty()) { + return result; + } + + ScopedPthreadMutexLock lock(&header->mutex); + + std::vector candidates; + candidates.reserve(config.max_nodes); + for (std::uint32_t index = 0; index < config.max_nodes; ++index) { + if (nodes[index].state == + static_cast(EntryState::kResident)) { + candidates.push_back(index); + } + } + std::sort(candidates.begin(), candidates.end(), + [this](std::uint32_t lhs, std::uint32_t rhs) { + return nodes[lhs].last_access_epoch < + nodes[rhs].last_access_epoch; + }); + + std::uint32_t scanned = 0; + for (std::uint32_t node_index : candidates) { + if (max_scan_nodes != 0 && scanned >= max_scan_nodes) { + break; + } + ++scanned; + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + continue; + } + if (NodeIsProtectedLocked(node)) { + ++result.protected_nodes; + continue; + } + + EvictNodeLocked(&node, &result); + + if (HasEnoughReleasablePages(result, required_pages)) { + break; + } + } + + if (result.evicted_nodes != 0) { + CompactArenasAfterEvictionIfUsefulLocked(0, 0); + } + header->evicted_nodes.fetch_add(result.evicted_nodes, + std::memory_order_relaxed); + header->eviction_protected_skips.fetch_add(result.protected_nodes, + std::memory_order_relaxed); + return result; +} + +PrefixEvictionResult +HostPrefixCacheCoordinator::SharedState::ClearUnprotected() { + PrefixEvictionResult result; + ScopedPthreadMutexLock lock(&header->mutex); + + for (std::uint32_t node_index = 0; node_index < config.max_nodes; + ++node_index) { + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + continue; + } + if (NodeIsProtectedLocked(node)) { + ++result.protected_nodes; + continue; + } + EvictNodeLocked(&node, &result); + } + + if (result.evicted_nodes != 0) { + CompactArenasAfterEvictionIfUsefulLocked(0, 0); + } + header->evicted_nodes.fetch_add(result.evicted_nodes, + std::memory_order_relaxed); + header->eviction_protected_skips.fetch_add(result.protected_nodes, + std::memory_order_relaxed); + return result; +} + +PrefixEvictionResult HostPrefixCacheCoordinator::SharedState::ClearNamespace( + PrefixDigest namespace_digest) { + PrefixEvictionResult result; + ScopedPthreadMutexLock lock(&header->mutex); + + for (std::uint32_t node_index = 0; node_index < config.max_nodes; + ++node_index) { + SharedPrefixNode& node = nodes[node_index]; + if (node.state != static_cast(EntryState::kResident)) { + continue; + } + if (!DigestEquals(node.namespace_digest, namespace_digest)) { + continue; + } + if (NodeIsProtectedLocked(node)) { + ++result.protected_nodes; + continue; + } + EvictNodeLocked(&node, &result); + } + + if (result.evicted_nodes != 0) { + CompactArenasAfterEvictionIfUsefulLocked(0, 0); + } + header->evicted_nodes.fetch_add(result.evicted_nodes, + std::memory_order_relaxed); + header->eviction_protected_skips.fetch_add(result.protected_nodes, + std::memory_order_relaxed); + return result; +} + +HostPrefixCacheStats HostPrefixCacheCoordinator::SharedState::GetStats() const { + HostPrefixCacheStats stats; + { + ScopedPthreadMutexLock lock(&header->mutex); + for (std::uint32_t index = 0; index < config.max_nodes; ++index) { + if (nodes[index].state == + static_cast(EntryState::kResident)) { + ++stats.resident_nodes; + } + } + for (std::uint32_t index = 0; + index < header->next_group_entry.load(std::memory_order_relaxed); + ++index) { + const SharedGroupEntry& entry = group_entries[index]; + if (entry.state != + static_cast(EntryState::kResident)) { + continue; + } + const std::uint32_t pending = + entry.pending_load_count.load(std::memory_order_relaxed); + if (pending != 0) { + ++stats.pending_load_entries; + stats.pending_load_refs += pending; + } + } + stats.used_group_entries = + header->next_group_entry.load(std::memory_order_relaxed); + stats.used_page_handles = + header->next_page_handle.load(std::memory_order_relaxed); + stats.lookup_hits = + header->lookup_hits.load(std::memory_order_relaxed); + stats.lookup_misses = + header->lookup_misses.load(std::memory_order_relaxed); + stats.evicted_nodes = + header->evicted_nodes.load(std::memory_order_relaxed); + stats.eviction_protected_skips = + header->eviction_protected_skips.load(std::memory_order_relaxed); + } + { + std::lock_guard attachment_lock(local_attachment_mutex); + stats.active_attachments = + static_cast(local_attachments.size()); + } + return stats; +} + +HostPrefixCacheCoordinator::HostPrefixCacheCoordinator( + HostPrefixCacheConfig config) + : config_(std::move(config)) { + if (config_.shm_name.empty()) { + throw std::invalid_argument("HostPrefixCacheConfig.shm_name is empty"); + } + if (config_.group_specs.empty()) { + throw std::invalid_argument( + "HostPrefixCacheConfig.group_specs is empty"); + } + if (config_.max_nodes == 0 || config_.max_group_entries == 0 || + config_.max_page_handles == 0) { + throw std::invalid_argument( + "HostPrefixCacheConfig capacities must be positive"); + } + bool has_required_group = false; + for (const auto& spec : config_.group_specs) { + ValidateGroupSpec(spec); + has_required_group = has_required_group || spec.required_for_reuse; + } + if (!has_required_group) { + throw std::invalid_argument( + "HostPrefixCacheConfig needs at least one required group"); + } + hash_block_tokens_ = config_.hash_block_tokens == 0 + ? ComputeHashBlockTokens(config_.group_specs) + : config_.hash_block_tokens; + commit_boundary_tokens_ = ComputeCommitBoundaryTokens(config_.group_specs); + if (hash_block_tokens_ == 0 || commit_boundary_tokens_ == 0) { + throw std::invalid_argument( + "HostPrefixCacheConfig computed zero token boundary"); + } + state_ = + new SharedState(config_, hash_block_tokens_, commit_boundary_tokens_); +} + +HostPrefixCacheCoordinator::~HostPrefixCacheCoordinator() { + if (state_ != nullptr) { + if (state_->mapping != nullptr && state_->total_bytes != 0) { + munmap(state_->mapping, state_->total_bytes); + } + if (state_->shm_fd >= 0) { + close(state_->shm_fd); + } + delete state_; + } +} + +void HostPrefixCacheCoordinator::Initialize(bool create_region) { + state_->Initialize(create_region); +} + +PrefixCommitResult HostPrefixCacheCoordinator::CommitPrefixPages( + PrefixDigest namespace_digest, const std::vector& token_ids, + std::uint32_t commit_tokens, + const std::vector& group_pages) { + return state_->CommitPrefixPages(namespace_digest, token_ids, commit_tokens, + group_pages); +} + +PrefixLookupResult HostPrefixCacheCoordinator::LookupAndAttach( + PrefixDigest namespace_digest, const std::vector& token_ids) { + return state_->LookupAndAttach(namespace_digest, token_ids); +} + +PrefixLookupResult HostPrefixCacheCoordinator::EstimateLookup( + PrefixDigest namespace_digest, const std::vector& token_ids) { + return state_->EstimateLookup(namespace_digest, token_ids); +} + +void HostPrefixCacheCoordinator::ReleaseAttachment( + std::uint64_t attachment_handle) { + state_->ReleaseAttachment(attachment_handle); +} + +void HostPrefixCacheCoordinator::BeginAttachmentLoad( + std::uint64_t attachment_handle) { + state_->BeginAttachmentLoad(attachment_handle); +} + +void HostPrefixCacheCoordinator::EndAttachmentLoad( + std::uint64_t attachment_handle) { + state_->EndAttachmentLoad(attachment_handle); +} + +PrefixEvictionResult HostPrefixCacheCoordinator::EvictUntilFree( + std::uint32_t min_free_nodes, std::uint32_t min_free_group_entries, + std::uint32_t min_free_page_handles, std::uint32_t max_scan_nodes) { + return state_->EvictUntilFree(min_free_nodes, min_free_group_entries, + min_free_page_handles, max_scan_nodes); +} + +PrefixEvictionResult HostPrefixCacheCoordinator::EvictUntilReleasablePages( + const std::vector& requirements, + std::uint32_t max_scan_nodes) { + return state_->EvictUntilReleasablePages(requirements, max_scan_nodes); +} + +PrefixEvictionResult HostPrefixCacheCoordinator::ClearUnprotected() { + return state_->ClearUnprotected(); +} + +PrefixEvictionResult HostPrefixCacheCoordinator::ClearNamespace( + PrefixDigest namespace_digest) { + return state_->ClearNamespace(namespace_digest); +} + +HostPrefixCacheStats HostPrefixCacheCoordinator::GetStats() const { + return state_->GetStats(); +} + +} // namespace batchgen::kv diff --git a/core/KV_Storage/host_prefix_cache_coordinator.h b/core/KV_Storage/host_prefix_cache_coordinator.h new file mode 100644 index 000000000..c54da3918 --- /dev/null +++ b/core/KV_Storage/host_prefix_cache_coordinator.h @@ -0,0 +1,160 @@ +#ifndef HOST_PREFIX_CACHE_COORDINATOR_H_ +#define HOST_PREFIX_CACHE_COORDINATOR_H_ + +#include +#include +#include +#include +#include + +namespace batchgen::kv { + +using PrefixDigest = std::array; + +enum class HostKVGroupSemantic : std::uint32_t { + kFullKV = 0, + kMlaCompressedKV = 1, + kSwaKV = 2, + kCompressedRatioKV = 3, +}; + +struct HostKVGroupSpec { + std::uint32_t group_id = 0; + HostKVGroupSemantic semantic = HostKVGroupSemantic::kFullKV; + bool required_for_reuse = true; + std::uint32_t raw_page_tokens = 0; + std::uint32_t compression_ratio = 1; +}; + +struct HostPageHandle { + std::uint32_t page_id = 0; +}; + +struct GroupCommitPages { + std::uint32_t group_id = 0; + std::vector pages; +}; + +struct GroupPageRequirement { + std::uint32_t group_id = 0; + std::uint32_t min_pages = 0; +}; + +struct GroupMaterializationSpan { + std::uint32_t group_id = 0; + std::uint32_t raw_end_token = 0; + std::vector pages; +}; + +struct PrefixLookupResult { + std::uint64_t attachment_handle = 0; + std::uint32_t common_cached_tokens = 0; + std::vector materialization_spans; + std::uint64_t miss_reason_mask = 0; +}; + +struct PrefixCommitResult { + std::uint32_t committed_tokens = 0; + std::uint32_t inserted_nodes = 0; + std::uint32_t existing_nodes = 0; +}; + +struct PrefixEvictionResult { + std::uint32_t evicted_nodes = 0; + std::uint32_t protected_nodes = 0; + std::uint32_t freed_group_entries = 0; + std::uint32_t freed_page_handles = 0; + std::vector evicted_group_pages; +}; + +struct HostPrefixCacheStats { + std::uint32_t resident_nodes = 0; + std::uint32_t active_attachments = 0; + std::uint32_t pending_load_entries = 0; + std::uint32_t pending_load_refs = 0; + std::uint32_t used_group_entries = 0; + std::uint32_t used_page_handles = 0; + std::uint64_t lookup_hits = 0; + std::uint64_t lookup_misses = 0; + std::uint64_t evicted_nodes = 0; + std::uint64_t eviction_protected_skips = 0; +}; + +struct HostPrefixCacheConfig { + std::string shm_name; + std::vector group_specs; + std::uint32_t hash_block_tokens = 0; + std::uint32_t max_nodes = 0; + std::uint32_t max_group_entries = 0; + std::uint32_t max_page_handles = 0; + std::uint32_t max_attachments = 0; +}; + +std::string ToString(const HostKVGroupSpec& spec); +std::string ToString(const HostPrefixCacheStats& stats); +std::vector> BuildPrefixHashChain( + PrefixDigest namespace_digest, const std::vector& token_ids, + std::uint32_t block_tokens); + +class HostPrefixCacheCoordinator { + public: + explicit HostPrefixCacheCoordinator(HostPrefixCacheConfig config); + HostPrefixCacheCoordinator(const HostPrefixCacheCoordinator&) = delete; + HostPrefixCacheCoordinator& operator=(const HostPrefixCacheCoordinator&) = + delete; + HostPrefixCacheCoordinator(HostPrefixCacheCoordinator&&) = delete; + HostPrefixCacheCoordinator& operator=(HostPrefixCacheCoordinator&&) = + delete; + ~HostPrefixCacheCoordinator(); + + void Initialize(bool create_region); + + PrefixCommitResult CommitPrefixPages( + PrefixDigest namespace_digest, + const std::vector& token_ids, std::uint32_t commit_tokens, + const std::vector& group_pages); + + PrefixLookupResult LookupAndAttach( + PrefixDigest namespace_digest, + const std::vector& token_ids); + + PrefixLookupResult EstimateLookup( + PrefixDigest namespace_digest, + const std::vector& token_ids); + + void ReleaseAttachment(std::uint64_t attachment_handle); + + void BeginAttachmentLoad(std::uint64_t attachment_handle); + void EndAttachmentLoad(std::uint64_t attachment_handle); + + PrefixEvictionResult EvictUntilFree(std::uint32_t min_free_nodes, + std::uint32_t min_free_group_entries, + std::uint32_t min_free_page_handles, + std::uint32_t max_scan_nodes); + PrefixEvictionResult EvictUntilReleasablePages( + const std::vector& requirements, + std::uint32_t max_scan_nodes); + + PrefixEvictionResult ClearUnprotected(); + PrefixEvictionResult ClearNamespace(PrefixDigest namespace_digest); + + HostPrefixCacheStats GetStats() const; + + std::uint32_t hash_block_tokens() const { return hash_block_tokens_; } + std::uint32_t commit_boundary_tokens() const { + return commit_boundary_tokens_; + } + const HostPrefixCacheConfig& config() const { return config_; } + + private: + struct SharedState; + + HostPrefixCacheConfig config_; + std::uint32_t hash_block_tokens_ = 0; + std::uint32_t commit_boundary_tokens_ = 0; + SharedState* state_ = nullptr; +}; + +} // namespace batchgen::kv + +#endif // HOST_PREFIX_CACHE_COORDINATOR_H_ diff --git a/core/KV_Storage/shared_memory_utils.h b/core/KV_Storage/shared_memory_utils.h new file mode 100644 index 000000000..7d03ee886 --- /dev/null +++ b/core/KV_Storage/shared_memory_utils.h @@ -0,0 +1,100 @@ +#ifndef SHARED_MEMORY_UTILS_H_ +#define SHARED_MEMORY_UTILS_H_ + +#include +#include + +#include +#include +#include +#include +#include + +namespace batchgen::kv { + +enum class SharedMemoryInitState : std::uint32_t { + kUninitialized = 0, + kInitializing = 1, + kReady = 2, +}; + +inline std::size_t AlignUp(std::size_t value, std::size_t alignment) { + if (alignment == 0) { + return value; + } + const std::size_t remainder = value % alignment; + if (remainder == 0) { + return value; + } + return value + (alignment - remainder); +} + +inline std::size_t SystemPageSize() { + const long page_size = sysconf(_SC_PAGESIZE); + if (page_size <= 0) { + throw std::system_error(errno, std::generic_category(), + "sysconf(_SC_PAGESIZE) failed"); + } + return static_cast(page_size); +} + +class ScopedPthreadMutexLock { + public: + explicit ScopedPthreadMutexLock(pthread_mutex_t* mutex) : mutex_(mutex) { + const int rc = pthread_mutex_lock(mutex_); + if (rc == EOWNERDEAD) { + const int consistent_rc = pthread_mutex_consistent(mutex_); + if (consistent_rc != 0) { + throw std::system_error(consistent_rc, std::generic_category(), + "pthread_mutex_consistent failed"); + } + } else if (rc != 0) { + throw std::system_error(rc, std::generic_category(), + "pthread_mutex_lock failed"); + } + } + + ScopedPthreadMutexLock(const ScopedPthreadMutexLock&) = delete; + ScopedPthreadMutexLock& operator=(const ScopedPthreadMutexLock&) = delete; + + ~ScopedPthreadMutexLock() { + const int rc = pthread_mutex_unlock(mutex_); + if (rc != 0) { + std::terminate(); + } + } + + private: + pthread_mutex_t* mutex_; +}; + +inline void InitProcessSharedRobustMutex(pthread_mutex_t* mutex, + const char* name) { + pthread_mutexattr_t attr; + if (const int rc = pthread_mutexattr_init(&attr); rc != 0) { + throw std::system_error(rc, std::generic_category(), + "pthread_mutexattr_init failed"); + } + if (const int rc = + pthread_mutexattr_setpshared(&attr, PTHREAD_PROCESS_SHARED); + rc != 0) { + pthread_mutexattr_destroy(&attr); + throw std::system_error(rc, std::generic_category(), + "pthread_mutexattr_setpshared failed"); + } + if (const int rc = pthread_mutexattr_setrobust(&attr, PTHREAD_MUTEX_ROBUST); + rc != 0) { + pthread_mutexattr_destroy(&attr); + throw std::system_error(rc, std::generic_category(), + "pthread_mutexattr_setrobust failed"); + } + if (const int rc = pthread_mutex_init(mutex, &attr); rc != 0) { + pthread_mutexattr_destroy(&attr); + throw std::system_error(rc, std::generic_category(), name); + } + pthread_mutexattr_destroy(&attr); +} + +} // namespace batchgen::kv + +#endif // SHARED_MEMORY_UTILS_H_ diff --git a/core/KV_Storage/swa_host_paged_kv_worker_view.h b/core/KV_Storage/swa_host_paged_kv_worker_view.h index 3c8603b79..e755bc736 100644 --- a/core/KV_Storage/swa_host_paged_kv_worker_view.h +++ b/core/KV_Storage/swa_host_paged_kv_worker_view.h @@ -1,30 +1,31 @@ #ifndef SWA_HOST_PAGED_KV_WORKER_VIEW_H_ #define SWA_HOST_PAGED_KV_WORKER_VIEW_H_ -#include -#include #include #include -#include -#include -#include #include #include #include #include -#include -#include #include -#include #include #include "host_paged_kv_worker_view.h" -#include "transformed_host_paged_kv_utils.h" namespace batchgen::kv { +struct SWAHostPageRange { + std::int64_t sequence_id = 0; + std::size_t raw_context_len = 0; + std::size_t window_start_token = 0; + std::size_t first_page = 0; + std::size_t page_count = 0; + std::size_t local_kv_len = 0; + std::size_t mask_start = 0; +}; + template -class SWAHostPagedKVWorkerView { +class SWAHostPagedKVWorkerView : public BaseView { public: using BatchedKVEntry = typename BaseView::BatchedKVEntry; static constexpr bool kHasVCache = BaseView::kHasVCache; @@ -34,17 +35,14 @@ class SWAHostPagedKVWorkerView { SWAHostPagedKVWorkerView(const EngineConfig& engine_config, const ModelConfig& model_config, std::size_t window_size_tokens) - : base_view_(engine_config, model_config), - page_size_tokens_(base_view_.config().page_size_tokens), + : BaseView(engine_config, model_config), window_size_tokens_(window_size_tokens) { ValidateWindowConfig(); } explicit SWAHostPagedKVWorkerView(const HostPagedKVConfig& config, std::size_t window_size_tokens) - : base_view_(config), - page_size_tokens_(base_view_.config().page_size_tokens), - window_size_tokens_(window_size_tokens) { + : BaseView(config), window_size_tokens_(window_size_tokens) { ValidateWindowConfig(); } @@ -54,280 +52,89 @@ class SWAHostPagedKVWorkerView { SWAHostPagedKVWorkerView(SWAHostPagedKVWorkerView&&) = delete; SWAHostPagedKVWorkerView& operator=(SWAHostPagedKVWorkerView&&) = delete; - void Initialize(int device_index, bool create_region = false) { - base_view_.Initialize(device_index, create_region); - } - - void Shutdown() { - { - std::lock_guard lock(mutex_); - pending_host_writes_.Drain(); - sequence_states_.clear(); - } - base_view_.Shutdown(); - } - - std::byte* DataBase() { return base_view_.DataBase(); } - const std::byte* DataBase() const { return base_view_.DataBase(); } - - void* KPagePtr(std::size_t layer_idx, std::int32_t page_idx) { - return base_view_.KPagePtr(layer_idx, page_idx); - } - - const void* KPagePtr(std::size_t layer_idx, - std::int32_t page_idx) const { - return base_view_.KPagePtr(layer_idx, page_idx); + std::size_t page_size_tokens() const { + return this->config().page_size_tokens; } - template > - void* VPagePtr(std::size_t layer_idx, std::int32_t page_idx) { - return base_view_.VPagePtr(layer_idx, page_idx); - } - - template > - const void* VPagePtr(std::size_t layer_idx, - std::int32_t page_idx) const { - return base_view_.VPagePtr(layer_idx, page_idx); - } + std::size_t window_size_tokens() const { return window_size_tokens_; } - [[nodiscard]] std::size_t ResolvePhysicalLayer( - std::size_t logical_layer_idx, std::string_view context) const { - return base_view_.ResolvePhysicalLayer(logical_layer_idx, context); + std::size_t window_pages() const { + return CeilDiv(window_size_tokens_, page_size_tokens()); } - const HostPagedKVConfig& config() const { return base_view_.config(); } - const auto& layout() const { return base_view_.layout(); } - HostPagedKVStats GetStats() const { return base_view_.GetStats(); } - int device_index() const { return base_view_.device_index(); } - std::size_t page_size_tokens() const { return page_size_tokens_; } - std::size_t window_size_tokens() const { return window_size_tokens_; } - std::size_t window_pages() const { return window_pages_; } - std::string DebugString() const { std::ostringstream oss; - oss << "SWAHostPagedKVWorkerView(window_size_tokens=" - << window_size_tokens_ << ", page_size_tokens=" - << page_size_tokens_ << ", window_pages=" << window_pages_ - << ", base=" << base_view_.DebugString() << ")"; + oss << "SWAHostPagedKVWorkerView(full_history=true, " + << "window_size_tokens=" << window_size_tokens_ + << ", page_size_tokens=" << page_size_tokens() + << ", window_pages=" << window_pages() + << ", base=" << BaseView::DebugString() << ")"; return oss.str(); } - std::vector> BuildPageTable( - const std::vector& sequence_ids) const { - return base_view_.BuildPageTable(sequence_ids); - } - - std::pair, std::optional>> - GetSequenceLayerPagePointers( - std::int64_t sequence_id, std::size_t layer_idx, - std::optional max_tokens = std::nullopt) const { - return base_view_.GetSequenceLayerPagePointers(sequence_id, layer_idx, - max_tokens); - } - - void RegisterSequences(const std::vector& sequence_ids) { - base_view_.RegisterSequences(sequence_ids); - std::lock_guard lock(mutex_); - for (std::int64_t sequence_id : sequence_ids) { - sequence_states_.try_emplace(sequence_id); - } - } - - void UnregisterSequence(std::int64_t sequence_id) { - base_view_.UnregisterSequence(sequence_id); - std::lock_guard lock(mutex_); - sequence_states_.erase(sequence_id); - } - - void UnregisterSequences(const std::vector& sequence_ids) { - base_view_.UnregisterSequences(sequence_ids); - std::lock_guard lock(mutex_); - for (std::int64_t sequence_id : sequence_ids) { - sequence_states_.erase(sequence_id); - } - } - - std::vector> AllocatePagesForSequences( + SWAHostPageRange ComputeSWAHostPageRange( + std::int64_t sequence_id, std::size_t raw_context_len) const { + const std::size_t page_size = page_size_tokens(); + const std::size_t window_start_token = + raw_context_len > window_size_tokens_ + ? raw_context_len - window_size_tokens_ + : 0; + const std::size_t first_page = window_start_token / page_size; + const std::size_t last_page_exclusive = + CeilDiv(raw_context_len, page_size); + const std::size_t first_page_token = first_page * page_size; + const std::size_t page_count = + last_page_exclusive > first_page + ? last_page_exclusive - first_page + : 0; + return SWAHostPageRange{ + sequence_id, + raw_context_len, + window_start_token, + first_page, + page_count, + raw_context_len - first_page_token, + window_start_token - first_page_token, + }; + } + + std::vector ComputeSWAHostPageRanges( const std::vector& sequence_ids, - const std::vector& raw_num_tokens) { - if (sequence_ids.size() != raw_num_tokens.size()) { + const std::vector& raw_context_lens) const { + if (sequence_ids.size() != raw_context_lens.size()) { throw std::invalid_argument( - "sequence_ids and raw_num_tokens must have the same length"); + "ComputeSWAHostPageRanges: sequence_ids and " + "raw_context_lens must have the same length"); } - std::vector active_tokens; - active_tokens.reserve(raw_num_tokens.size()); - std::vector windows; - windows.reserve(raw_num_tokens.size()); - for (std::size_t raw_tokens : raw_num_tokens) { - const auto window = ComputeWindowForRawEnd(raw_tokens); - if (window.active_tokens == 0) { - throw std::invalid_argument( - "AllocatePagesForSequences: raw_num_tokens entries must " - "be greater than zero"); - } - active_tokens.push_back(window.active_tokens); - windows.push_back(window); - } - - auto allocations = - base_view_.AllocatePagesForSequences(sequence_ids, active_tokens); - std::lock_guard lock(mutex_); + std::vector ranges; + ranges.reserve(sequence_ids.size()); for (std::size_t i = 0; i < sequence_ids.size(); ++i) { - auto& state = sequence_states_[sequence_ids[i]]; - state.window_start_page = windows[i].window_start_page; - state.active_pages = windows[i].required_pages; - state.max_seen_raw_pos = raw_num_tokens[i] - 1; - state.has_tokens = true; - } - return allocations; - } - - void ReleaseSequencePages(const std::vector& sequence_ids) { - std::lock_guard lock(mutex_); - pending_host_writes_.Drain(); - base_view_.ReleaseSequencePages(sequence_ids); - for (std::int64_t sequence_id : sequence_ids) { - sequence_states_.erase(sequence_id); - } - } - - KVAsyncTask AsyncLoadLayerKVToDevice( - torch::Tensor sequence_ids, torch::Tensor k_device_ptrs, - std::optional v_device_ptrs = std::nullopt) { - return base_view_.AsyncLoadLayerKVToDevice( - std::move(sequence_ids), std::move(k_device_ptrs), - std::move(v_device_ptrs)); - } - - KVAsyncTask AsyncLoadLayerPagedKVToDevice( - torch::Tensor sequence_ids, torch::Tensor active_page_counts, - torch::Tensor k_device_ptrs, - std::optional v_device_ptrs = std::nullopt) { - return base_view_.AsyncLoadLayerPagedKVToDevice( - std::move(sequence_ids), std::move(active_page_counts), - std::move(k_device_ptrs), std::move(v_device_ptrs)); - } - - KVAsyncTask AsyncOffloadLayerKVToHost( - std::size_t layer_idx, std::vector sequence_ids, - torch::Tensor k_tensor, std::optional v_tensor, - SequenceLengths raw_sequence_lengths) { - if (sequence_ids.empty()) { - return transformed_detail::MakeAsyncTask([] {}); - } - std::vector tasks; - { - std::lock_guard lock(mutex_); - const std::size_t batch = sequence_ids.size(); - for (std::size_t batch_idx = 0; batch_idx < batch; ++batch_idx) { - const std::int64_t sequence_id = sequence_ids[batch_idx]; - const std::size_t raw_tokens = - transformed_detail::ResolveLength( - raw_sequence_lengths, batch_idx, sequence_id, - "SWAHostPagedKVWorkerView::" - "AsyncOffloadLayerKVToHost"); - const auto active_tokens = - UpdateWindowForRawEndLocked(sequence_id, raw_tokens); - if (active_tokens == 0) { - continue; - } - const auto source_start = - static_cast(raw_tokens - active_tokens); - auto k_slice = k_tensor - .narrow(0, static_cast( - batch_idx), - 1) - .narrow(1, source_start, - static_cast( - active_tokens)) - .contiguous(); - std::optional v_slice; - if (v_tensor.has_value()) { - v_slice = v_tensor->narrow( - 0, static_cast( - batch_idx), - 1) - .narrow(1, source_start, - static_cast( - active_tokens)) - .contiguous(); - } - auto task = base_view_.AsyncOffloadLayerKVToHost( - layer_idx, {sequence_id}, std::move(k_slice), - std::move(v_slice), SequenceLengthVector{active_tokens}); - pending_host_writes_.Track(task); - tasks.emplace_back(std::move(task)); - } + ranges.push_back( + ComputeSWAHostPageRange(sequence_ids[i], raw_context_lens[i])); } - return transformed_detail::MakeCombinedTask(std::move(tasks)); + return ranges; } - KVAsyncTask AsyncAppendDecodeKVToHost( - std::size_t layer_idx, std::vector sequence_ids, - torch::Tensor k_tensor, std::optional v_tensor, - SequenceLengths raw_positions) { - if (sequence_ids.empty()) { - return transformed_detail::MakeAsyncTask([] {}); - } - KVAsyncTask task; - { - std::lock_guard lock(mutex_); - auto storage_positions = - PrepareStoragePositionsLocked(sequence_ids, raw_positions); - task = base_view_.AsyncAppendDecodeKVToHost( - layer_idx, std::move(sequence_ids), std::move(k_tensor), - std::move(v_tensor), std::move(storage_positions)); - pending_host_writes_.Track(task); - } - return task; + std::pair, std::optional>> + GetSequenceLayerSWAWindowPagePointers( + std::int64_t sequence_id, std::size_t layer_idx, + std::size_t raw_context_len) const { + const SWAHostPageRange range = + ComputeSWAHostPageRange(sequence_id, raw_context_len); + return this->GetSequenceLayerPageRangePointers( + sequence_id, layer_idx, range.first_page, range.page_count); } - KVAsyncTask AsyncAppendDecodeKVToHostBatchedKernel( - std::vector entries, - std::vector sequence_ids, SequenceLengths raw_positions) { - if (entries.empty() || sequence_ids.empty()) { - return transformed_detail::MakeAsyncTask([] {}); - } - KVAsyncTask task; - { - std::lock_guard lock(mutex_); - auto storage_positions = - PrepareStoragePositionsLocked(sequence_ids, raw_positions); - task = base_view_.AsyncAppendDecodeKVToHostBatchedKernel( - std::move(entries), std::move(sequence_ids), - std::move(storage_positions)); - pending_host_writes_.Track(task); + private: + static std::size_t CeilDiv(std::size_t value, std::size_t divisor) { + if (divisor == 0) { + throw std::invalid_argument("CeilDiv divisor must be non-zero"); } - return task; - } - - std::pair ReadSequenceKVToCPU( - std::int64_t sequence_id) const { - return base_view_.ReadSequenceKVToCPU(sequence_id); - } - - void WriteSequenceKVFromCPU( - std::int64_t sequence_id, const torch::Tensor& k_tensor, - const std::optional& v_tensor = std::nullopt) { - base_view_.WriteSequenceKVFromCPU(sequence_id, k_tensor, v_tensor); + return (value + divisor - 1) / divisor; } - private: - struct SWASequenceState { - std::size_t window_start_page = 0; - std::size_t active_pages = 0; - std::size_t max_seen_raw_pos = 0; - bool has_tokens = false; - }; - - struct WindowForRawEnd { - std::size_t window_start_page = 0; - std::size_t active_tokens = 0; - std::size_t required_pages = 0; - }; - - void ValidateWindowConfig() { - if (page_size_tokens_ == 0) { + void ValidateWindowConfig() const { + if (page_size_tokens() == 0) { throw std::invalid_argument( "SWAHostPagedKVWorkerView requires page_size_tokens > 0"); } @@ -335,130 +142,9 @@ class SWAHostPagedKVWorkerView { throw std::invalid_argument( "SWAHostPagedKVWorkerView requires window_size_tokens > 0"); } - if (window_size_tokens_ % page_size_tokens_ != 0) { - throw std::invalid_argument( - "SWAHostPagedKVWorkerView requires window_size_tokens to be " - "divisible by page_size_tokens"); - } - window_pages_ = window_size_tokens_ / page_size_tokens_; - } - - WindowForRawEnd ComputeWindowForRawEnd(std::size_t raw_end_tokens) const { - if (raw_end_tokens == 0) { - return {}; - } - const std::size_t first_needed_token = - raw_end_tokens > window_size_tokens_ - ? raw_end_tokens - window_size_tokens_ - : 0; - const std::size_t window_start_page = - first_needed_token / page_size_tokens_; - const std::size_t window_start_token = - window_start_page * page_size_tokens_; - const std::size_t active_tokens = - raw_end_tokens - window_start_token; - const std::size_t required_pages = - (active_tokens + page_size_tokens_ - 1) / page_size_tokens_; - if (required_pages > window_pages_ + 1) { - std::ostringstream oss; - oss << "SWAHostPagedKVWorkerView: active pages " - << required_pages << " exceed window_pages + 1 (" - << (window_pages_ + 1) << ")"; - throw std::logic_error(oss.str()); - } - return {window_start_page, active_tokens, required_pages}; - } - - std::size_t UpdateWindowForRawEndLocked(std::int64_t sequence_id, - std::size_t raw_end_tokens) { - const auto window = ComputeWindowForRawEnd(raw_end_tokens); - auto& state = sequence_states_[sequence_id]; - if (state.has_tokens && - window.window_start_page < state.window_start_page) { - throw std::out_of_range( - "SWAHostPagedKVWorkerView does not support writing a raw " - "token range that is older than the current SWA window"); - } - if (state.has_tokens && - window.window_start_page > state.window_start_page) { - const std::size_t pages_to_release = - window.window_start_page - state.window_start_page; - if (pages_to_release > state.active_pages) { - std::ostringstream oss; - oss << "SWAHostPagedKVWorkerView: sequence " << sequence_id - << " cannot release " << pages_to_release - << " pages with only " << state.active_pages - << " active pages"; - throw std::out_of_range(oss.str()); - } - pending_host_writes_.Drain(); - base_view_.ReleaseSequencePrefixPages(sequence_id, - pages_to_release); - state.active_pages -= pages_to_release; - } - state.window_start_page = window.window_start_page; - EnsureCapacityForActivePagesLocked(sequence_id, state, - window.required_pages); - if (raw_end_tokens > 0) { - state.max_seen_raw_pos = - std::max(state.max_seen_raw_pos, raw_end_tokens - 1); - state.has_tokens = true; - } - return window.active_tokens; - } - - void EnsureCapacityForActivePagesLocked(std::int64_t sequence_id, - SWASequenceState& state, - std::size_t required_pages) { - if (required_pages == 0) { - return; - } - if (required_pages > window_pages_ + 1) { - std::ostringstream oss; - oss << "SWAHostPagedKVWorkerView: sequence " << sequence_id - << " requires " << required_pages - << " active pages, exceeding window_pages + 1 (" - << (window_pages_ + 1) << ")"; - throw std::out_of_range(oss.str()); - } - if (state.active_pages < required_pages) { - const std::size_t missing_pages = - required_pages - state.active_pages; - base_view_.GrowSequencePages(sequence_id, - missing_pages); - state.active_pages += missing_pages; - } - } - - SequenceLengthVector PrepareStoragePositionsLocked( - const std::vector& sequence_ids, - const SequenceLengths& raw_positions) { - SequenceLengthVector storage_positions; - storage_positions.reserve(sequence_ids.size()); - for (std::size_t batch_idx = 0; batch_idx < sequence_ids.size(); - ++batch_idx) { - const std::int64_t sequence_id = sequence_ids[batch_idx]; - const std::size_t raw_pos = transformed_detail::ResolveLength( - raw_positions, batch_idx, sequence_id, - "SWAHostPagedKVWorkerView::AsyncAppendDecodeKVToHost"); - if (raw_pos == std::numeric_limits::max()) { - throw std::out_of_range( - "SWAHostPagedKVWorkerView: raw position overflow"); - } - const std::size_t active_tokens_after = - UpdateWindowForRawEndLocked(sequence_id, raw_pos + 1); - storage_positions.push_back(active_tokens_after - 1); - } - return storage_positions; } - BaseView base_view_; - std::size_t page_size_tokens_ = 0; std::size_t window_size_tokens_ = 0; - std::size_t window_pages_ = 0; - mutable std::mutex mutex_; - std::unordered_map sequence_states_; - transformed_detail::PendingHostWriteTasks pending_host_writes_; }; using SWADefaultHostPagedKVWorkerView = diff --git a/core/batchgen_Binding.cpp b/core/batchgen_Binding.cpp index af0679bcf..c184fe35e 100644 --- a/core/batchgen_Binding.cpp +++ b/core/batchgen_Binding.cpp @@ -20,6 +20,7 @@ #include "KV_Storage/host_paged_kv_manager.h" #include "KV_Storage/host_paged_kv_worker_view.h" +#include "KV_Storage/host_prefix_cache_coordinator.h" #include "KV_Storage/compressed_state_host_manager.h" #include "KV_Storage/compressed_ratio_host_paged_kv_worker_view.h" #include "KV_Storage/swa_host_paged_kv_worker_view.h" @@ -71,6 +72,25 @@ struct HasGrowPagesForSequences< std::declval&>()))>> : std::true_type {}; +inline py::tuple PagePointersToPyTuple( + std::pair, std::optional>> result) { + py::list k_ptrs; + for (void* ptr : result.first) { + k_ptrs.append( + py::int_(reinterpret_cast(ptr))); + } + py::object v_ptrs = py::none(); + if (result.second.has_value()) { + py::list v_list; + for (void* ptr : result.second.value()) { + v_list.append( + py::int_(reinterpret_cast(ptr))); + } + v_ptrs = std::move(v_list); + } + return py::make_tuple(std::move(k_ptrs), v_ptrs); +} + template void BindHostPagedManager(py::module& m, const char* name) { py::class_(m, name) @@ -83,40 +103,40 @@ void BindHostPagedManager(py::module& m, const char* name) { return self.AllocatePages(sequence_id, num_tokens); }, py::arg("sequence_id"), py::arg("num_tokens")) - .def("free_sequence", &Manager::FreeSequence, - py::arg("sequence_id")) - .def("free_sequences", &Manager::FreeSequences, - py::arg("sequence_ids")) + .def("free_sequence", &Manager::FreeSequence, + py::arg("sequence_id")) + .def("free_sequences", &Manager::FreeSequences, + py::arg("sequence_ids")) + .def("release_resident_pages", &Manager::ReleaseResidentPages, + py::arg("page_ids"), + "Release prefix-cache resident pages returned by coordinator " + "eviction.") .def("build_page_table", &Manager::BuildPageTable, py::arg("sequence_ids")) .def("get_stats", &Manager::GetStats) .def("memfd_fd", &Manager::memfd_fd) - .def("__repr__", - [](const Manager& self) { return self.DebugString(); }) + .def("__repr__", + [](const Manager& self) { return self.DebugString(); }) .def("get_sequence_layer_page_pointers", [](Manager& self, std::int64_t sequence_id, std::size_t layer_idx, std::optional max_tokens) { - auto result = self.GetSequenceLayerPagePointers( - sequence_id, layer_idx, max_tokens); - py::list k_ptrs; - for (void* ptr : result.first) { - k_ptrs.append(py::int_( - reinterpret_cast(ptr))); - } - py::object v_ptrs = py::none(); - if (result.second.has_value()) { - py::list v_list; - for (void* ptr : result.second.value()) { - v_list.append(py::int_( - reinterpret_cast(ptr))); - } - v_ptrs = std::move(v_list); - } - return py::make_tuple(std::move(k_ptrs), v_ptrs); + return PagePointersToPyTuple( + self.GetSequenceLayerPagePointers( + sequence_id, layer_idx, max_tokens)); }, py::arg("sequence_id"), py::arg("layer_idx"), - py::arg("max_tokens") = py::none()); + py::arg("max_tokens") = py::none()) + .def("get_sequence_layer_page_range_pointers", + [](Manager& self, std::int64_t sequence_id, + std::size_t layer_idx, std::size_t start_page, + std::size_t page_count) { + return PagePointersToPyTuple( + self.GetSequenceLayerPageRangePointers( + sequence_id, layer_idx, start_page, page_count)); + }, + py::arg("sequence_id"), py::arg("layer_idx"), + py::arg("start_page"), py::arg("page_count")); } template @@ -209,12 +229,42 @@ void BindCommonHostPagedWorkerViewMethods(py::class_& cls) { py::arg("sequence_ids")) .def("register_sequences", &WorkerView::RegisterSequences, py::arg("sequence_ids")) + .def("attach_shared_prefix_pages", + &WorkerView::AttachSharedPrefixPages, + py::arg("sequence_id"), py::arg("page_ids"), + "Prepend shared prefix Host page ids to a registered sequence's " + "logical page table. Ownership remains with the prefix cache.") + .def("attach_shared_prefix_pages_for_sequences", + &WorkerView::AttachSharedPrefixPagesForSequences, + py::arg("sequence_ids"), py::arg("page_ids_by_sequence"), + "Prepend shared prefix Host pages for multiple registered " + "sequences.") .def("unregister_sequence", &WorkerView::UnregisterSequence, py::arg("sequence_id")) .def("unregister_sequences", &WorkerView::UnregisterSequences, py::arg("sequence_ids")) .def("release_sequence_pages", &WorkerView::ReleaseSequencePages, py::arg("sequence_ids")) + .def("retain_sequence_prefix_pages", + &WorkerView::RetainSequencePrefixPages, + py::arg("sequence_id"), py::arg("num_pages"), + "Move sequence-owned prefix pages into prefix-cache resident " + "ownership without changing the worker logical page table.") + .def("retain_sequence_page_range", + &WorkerView::RetainSequencePageRange, + py::arg("sequence_id"), py::arg("start_page"), + py::arg("num_pages"), + "Move a sequence-owned page range into prefix-cache resident " + "ownership without changing the worker logical page table.") + .def("retain_sequence_pages", + &WorkerView::RetainSequencePages, + py::arg("sequence_id"), py::arg("page_ids"), + "Move exact sequence-owned pages into prefix-cache resident " + "ownership without changing the worker logical page table.") + .def("release_resident_pages", &WorkerView::ReleaseResidentPages, + py::arg("page_ids"), + "Release prefix-cache resident pages returned by coordinator " + "eviction.") .def("read_sequence_kv_to_cpu", &WorkerView::ReadSequenceKVToCPU, py::arg("sequence_id"), "Read all KV pages for a sequence directly to CPU tensors (no GPU). " @@ -236,6 +286,14 @@ void BindCommonHostPagedWorkerViewMethods(py::class_& cls) { "Offload one layer of prefill KV into host pages. For mapped " "worker views, layer_idx is a logical layer id and is resolved to " "a physical layer id before writing.") + .def("async_offload_layer_kv_range_to_host", + &WorkerView::AsyncOffloadLayerKVRangeToHost, + py::arg("layer_idx"), py::arg("sequence_ids"), + py::arg("k_tensor"), py::arg("v_tensor") = py::none(), + py::arg("raw_start_positions"), py::arg("token_counts"), + "Offload one layer of KV into a raw token range in host pages. " + "The source tensor starts at offset 0 while raw_start_positions " + "select each sequence's destination offset.") .def("async_append_decode_kv_to_host", &WorkerView::AsyncAppendDecodeKVToHost, py::arg("layer_idx"), py::arg("sequence_ids"), @@ -294,6 +352,41 @@ void BindCommonHostPagedWorkerViewMethods(py::class_& cls) { "tables. This loads all physical layers; destination pointer " "tensors are indexed by physical layer id even for mapped worker " "views.") + .def( + "async_load_prefix_pages_to_device", + [](WorkerView& self, torch::Tensor host_page_ids, + torch::Tensor active_page_counts, + torch::Tensor k_device_ptrs, + std::optional v_device_ptrs) { + return self.AsyncLoadPrefixPagesToDevice( + std::move(host_page_ids), std::move(active_page_counts), + std::move(k_device_ptrs), std::move(v_device_ptrs)); + }, + py::arg("host_page_ids"), py::arg("active_page_counts"), + py::arg("k_device_ptrs"), + py::arg("v_device_ptrs") = py::none(), + "Load prefix-cache Host page ids into pre-allocated GPU pages. " + "Unlike async_load_layer_paged_kv_to_device, this reads directly " + "from the provided physical Host page ids instead of resolving " + "pages through sequence ids.") + .def( + "async_load_prefix_layers_to_device", + [](WorkerView& self, torch::Tensor host_page_ids, + torch::Tensor active_page_counts, + torch::Tensor logical_layer_ids, + torch::Tensor k_device_ptrs, + std::optional v_device_ptrs) { + return self.AsyncLoadPrefixLayersToDevice( + std::move(host_page_ids), std::move(active_page_counts), + std::move(logical_layer_ids), std::move(k_device_ptrs), + std::move(v_device_ptrs)); + }, + py::arg("host_page_ids"), py::arg("active_page_counts"), + py::arg("logical_layer_ids"), py::arg("k_device_ptrs"), + py::arg("v_device_ptrs") = py::none(), + "Load selected logical prefix-cache Host layers into provided " + "GPU destination pointer rows. The destination tensor first " + "dimension must match logical_layer_ids length.") .def("__repr__", [](const WorkerView& self) { return self.DebugString(); }) .def( @@ -333,30 +426,29 @@ void BindCommonHostPagedWorkerViewMethods(py::class_& cls) { [](WorkerView& self, std::int64_t sequence_id, std::size_t layer_idx, std::optional max_tokens) { - auto result = self.GetSequenceLayerPagePointers( - sequence_id, layer_idx, max_tokens); - py::list k_ptrs; - for (void* ptr : result.first) { - k_ptrs.append(py::int_( - reinterpret_cast(ptr))); - } - py::object v_ptrs = py::none(); - if (result.second.has_value()) { - py::list v_list; - for (void* ptr : result.second.value()) { - v_list.append(py::int_( - reinterpret_cast(ptr))); - } - v_ptrs = std::move(v_list); - } - return py::make_tuple(std::move(k_ptrs), v_ptrs); + return PagePointersToPyTuple( + self.GetSequenceLayerPagePointers( + sequence_id, layer_idx, max_tokens)); }, py::arg("sequence_id"), py::arg("layer_idx"), py::arg("max_tokens") = py::none(), "Return per-page K/V host pointers for one sequence and one " "layer. For mapped worker views, layer_idx is a logical layer id " "and is resolved to a physical layer id before address " - "calculation."); + "calculation.") + .def("get_sequence_layer_page_range_pointers", + [](WorkerView& self, std::int64_t sequence_id, + std::size_t layer_idx, std::size_t start_page, + std::size_t page_count) { + return PagePointersToPyTuple( + self.GetSequenceLayerPageRangePointers( + sequence_id, layer_idx, start_page, page_count)); + }, + py::arg("sequence_id"), py::arg("layer_idx"), + py::arg("start_page"), py::arg("page_count"), + "Return K/V host pointers for a raw page range of one sequence " + "and one layer. For mapped worker views, layer_idx is a logical " + "layer id and is resolved before address calculation."); if constexpr (HasGrowSequencePages::value) { cls.def("grow_sequence_pages", @@ -411,7 +503,24 @@ void BindSWAHostPagedWorkerView(py::module& m, const char* name) { &WorkerView::page_size_tokens) .def_property_readonly("window_size_tokens", &WorkerView::window_size_tokens) - .def_property_readonly("window_pages", &WorkerView::window_pages); + .def_property_readonly("window_pages", &WorkerView::window_pages) + .def("compute_swa_host_page_range", + &WorkerView::ComputeSWAHostPageRange, + py::arg("sequence_id"), py::arg("raw_context_len")) + .def("compute_swa_host_page_ranges", + &WorkerView::ComputeSWAHostPageRanges, + py::arg("sequence_ids"), py::arg("raw_context_lens")) + .def("get_sequence_layer_swa_window_page_pointers", + [](WorkerView& self, std::int64_t sequence_id, + std::size_t layer_idx, std::size_t raw_context_len) { + return PagePointersToPyTuple( + self.GetSequenceLayerSWAWindowPagePointers( + sequence_id, layer_idx, raw_context_len)); + }, + py::arg("sequence_id"), py::arg("layer_idx"), + py::arg("raw_context_len"), + "Return K/V host pointers for the raw pages covering the current " + "SWA window. The underlying Host KV remains full-history."); } template @@ -606,6 +715,215 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { return kv::ToString(self); }); + py::enum_(m, "HostKVGroupSemantic") + .value("FULL_KV", kv::HostKVGroupSemantic::kFullKV) + .value("MLA_COMPRESSED_KV", + kv::HostKVGroupSemantic::kMlaCompressedKV) + .value("SWA_KV", kv::HostKVGroupSemantic::kSwaKV) + .value("COMPRESSED_RATIO_KV", + kv::HostKVGroupSemantic::kCompressedRatioKV); + + py::class_(m, "HostKVGroupSpec") + .def(py::init<>()) + .def_readwrite("group_id", &kv::HostKVGroupSpec::group_id) + .def_readwrite("semantic", &kv::HostKVGroupSpec::semantic) + .def_readwrite("required_for_reuse", + &kv::HostKVGroupSpec::required_for_reuse) + .def_readwrite("raw_page_tokens", + &kv::HostKVGroupSpec::raw_page_tokens) + .def_readwrite("compression_ratio", + &kv::HostKVGroupSpec::compression_ratio) + .def("__repr__", [](const kv::HostKVGroupSpec& self) { + return kv::ToString(self); + }); + + py::class_(m, "HostPageHandle") + .def(py::init<>()) + .def_readwrite("page_id", &kv::HostPageHandle::page_id); + + py::class_(m, "GroupCommitPages") + .def(py::init<>()) + .def_readwrite("group_id", &kv::GroupCommitPages::group_id) + .def_readwrite("pages", &kv::GroupCommitPages::pages); + + py::class_(m, "GroupPageRequirement") + .def(py::init<>()) + .def_readwrite("group_id", &kv::GroupPageRequirement::group_id) + .def_readwrite("min_pages", &kv::GroupPageRequirement::min_pages); + + py::class_( + m, "GroupMaterializationSpan") + .def_readonly("group_id", &kv::GroupMaterializationSpan::group_id) + .def_readonly("raw_end_token", + &kv::GroupMaterializationSpan::raw_end_token) + .def_readonly("pages", &kv::GroupMaterializationSpan::pages); + + py::class_(m, "PrefixLookupResult") + .def_readonly("attachment_handle", + &kv::PrefixLookupResult::attachment_handle) + .def_readonly("common_cached_tokens", + &kv::PrefixLookupResult::common_cached_tokens) + .def_readonly("materialization_spans", + &kv::PrefixLookupResult::materialization_spans) + .def_readonly("miss_reason_mask", + &kv::PrefixLookupResult::miss_reason_mask); + + py::class_(m, "PrefixCommitResult") + .def_readonly("committed_tokens", + &kv::PrefixCommitResult::committed_tokens) + .def_readonly("inserted_nodes", + &kv::PrefixCommitResult::inserted_nodes) + .def_readonly("existing_nodes", + &kv::PrefixCommitResult::existing_nodes); + + py::class_(m, "PrefixEvictionResult") + .def_readonly("evicted_nodes", + &kv::PrefixEvictionResult::evicted_nodes) + .def_readonly("protected_nodes", + &kv::PrefixEvictionResult::protected_nodes) + .def_readonly("freed_group_entries", + &kv::PrefixEvictionResult::freed_group_entries) + .def_readonly("freed_page_handles", + &kv::PrefixEvictionResult::freed_page_handles) + .def_readonly("evicted_group_pages", + &kv::PrefixEvictionResult::evicted_group_pages); + + py::class_(m, "HostPrefixCacheStats") + .def(py::init<>()) + .def_readwrite("resident_nodes", + &kv::HostPrefixCacheStats::resident_nodes) + .def_readwrite("active_attachments", + &kv::HostPrefixCacheStats::active_attachments) + .def_readwrite("pending_load_entries", + &kv::HostPrefixCacheStats::pending_load_entries) + .def_readwrite("pending_load_refs", + &kv::HostPrefixCacheStats::pending_load_refs) + .def_readwrite("used_group_entries", + &kv::HostPrefixCacheStats::used_group_entries) + .def_readwrite("used_page_handles", + &kv::HostPrefixCacheStats::used_page_handles) + .def_readwrite("lookup_hits", &kv::HostPrefixCacheStats::lookup_hits) + .def_readwrite("lookup_misses", + &kv::HostPrefixCacheStats::lookup_misses) + .def_readwrite("evicted_nodes", + &kv::HostPrefixCacheStats::evicted_nodes) + .def_readwrite("eviction_protected_skips", + &kv::HostPrefixCacheStats::eviction_protected_skips) + .def("__repr__", [](const kv::HostPrefixCacheStats& self) { + return kv::ToString(self); + }); + + py::class_(m, "HostPrefixCacheConfig") + .def(py::init<>()) + .def_readwrite("shm_name", &kv::HostPrefixCacheConfig::shm_name) + .def_readwrite("group_specs", + &kv::HostPrefixCacheConfig::group_specs) + .def_readwrite("hash_block_tokens", + &kv::HostPrefixCacheConfig::hash_block_tokens) + .def_readwrite("max_nodes", &kv::HostPrefixCacheConfig::max_nodes) + .def_readwrite("max_group_entries", + &kv::HostPrefixCacheConfig::max_group_entries) + .def_readwrite("max_page_handles", + &kv::HostPrefixCacheConfig::max_page_handles) + .def_readwrite("max_attachments", + &kv::HostPrefixCacheConfig::max_attachments); + + py::class_( + m, "HostPrefixCacheCoordinator") + .def(py::init(), py::arg("config")) + .def("initialize", &kv::HostPrefixCacheCoordinator::Initialize, + py::arg("create_region")) + .def("commit_prefix_pages", + &kv::HostPrefixCacheCoordinator::CommitPrefixPages, + py::arg("namespace_digest"), py::arg("token_ids"), + py::arg("commit_tokens"), py::arg("group_pages")) + .def( + "commit_prefix_page_ids", + [](kv::HostPrefixCacheCoordinator& self, + kv::PrefixDigest namespace_digest, + const std::vector& token_ids, + std::uint32_t commit_tokens, + const std::vector< + std::pair>>& + group_page_ids) { + std::vector group_pages; + group_pages.reserve(group_page_ids.size()); + for (const auto& [group_id, page_ids] : group_page_ids) { + kv::GroupCommitPages group; + group.group_id = group_id; + group.pages.reserve(page_ids.size()); + for (std::uint32_t page_id : page_ids) { + group.pages.push_back(kv::HostPageHandle{page_id}); + } + group_pages.emplace_back(std::move(group)); + } + return self.CommitPrefixPages(namespace_digest, token_ids, + commit_tokens, group_pages); + }, + py::arg("namespace_digest"), py::arg("token_ids"), + py::arg("commit_tokens"), py::arg("group_page_ids")) + .def("lookup_and_attach", + &kv::HostPrefixCacheCoordinator::LookupAndAttach, + py::arg("namespace_digest"), py::arg("token_ids")) + .def("estimate_lookup", + &kv::HostPrefixCacheCoordinator::EstimateLookup, + py::arg("namespace_digest"), py::arg("token_ids")) + .def("release_attachment", + &kv::HostPrefixCacheCoordinator::ReleaseAttachment, + py::arg("attachment_handle")) + .def("begin_attachment_load", + &kv::HostPrefixCacheCoordinator::BeginAttachmentLoad, + py::arg("attachment_handle")) + .def("end_attachment_load", + &kv::HostPrefixCacheCoordinator::EndAttachmentLoad, + py::arg("attachment_handle")) + .def("evict_until_free", + &kv::HostPrefixCacheCoordinator::EvictUntilFree, + py::arg("min_free_nodes"), + py::arg("min_free_group_entries"), + py::arg("min_free_page_handles"), py::arg("max_scan_nodes")) + .def("evict_until_releasable_pages", + &kv::HostPrefixCacheCoordinator::EvictUntilReleasablePages, + py::arg("requirements"), py::arg("max_scan_nodes")) + .def("clear_unprotected", + &kv::HostPrefixCacheCoordinator::ClearUnprotected) + .def("clear_namespace", + &kv::HostPrefixCacheCoordinator::ClearNamespace, + py::arg("namespace_digest")) + .def("get_stats", &kv::HostPrefixCacheCoordinator::GetStats) + .def_property_readonly( + "hash_block_tokens", + &kv::HostPrefixCacheCoordinator::hash_block_tokens) + .def_property_readonly( + "commit_boundary_tokens", + &kv::HostPrefixCacheCoordinator::commit_boundary_tokens); + + m.def("build_prefix_hash_chain", &kv::BuildPrefixHashChain, + py::arg("namespace_digest"), py::arg("token_ids"), + py::arg("block_tokens")); + + py::class_(m, "SWAHostPageRange") + .def_readonly("sequence_id", &kv::SWAHostPageRange::sequence_id) + .def_readonly("raw_context_len", + &kv::SWAHostPageRange::raw_context_len) + .def_readonly("window_start_token", + &kv::SWAHostPageRange::window_start_token) + .def_readonly("first_page", &kv::SWAHostPageRange::first_page) + .def_readonly("page_count", &kv::SWAHostPageRange::page_count) + .def_readonly("local_kv_len", &kv::SWAHostPageRange::local_kv_len) + .def_readonly("mask_start", &kv::SWAHostPageRange::mask_start) + .def("__repr__", [](const kv::SWAHostPageRange& range) { + std::ostringstream oss; + oss << "SWAHostPageRange(sequence_id=" << range.sequence_id + << ", raw_context_len=" << range.raw_context_len + << ", window_start_token=" << range.window_start_token + << ", first_page=" << range.first_page + << ", page_count=" << range.page_count + << ", local_kv_len=" << range.local_kv_len + << ", mask_start=" << range.mask_start << ")"; + return oss.str(); + }); + py::class_(m, "CompressedStateHostStats") .def(py::init<>()) @@ -663,6 +981,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("done", &kv::KVAsyncTask::done) .def("result", &kv::KVAsyncTask::result); + py::class_(m, "KVLayeredAsyncTask") + .def_property_readonly("id", &kv::KVLayeredAsyncTask::id) + .def_property_readonly("num_layers", + &kv::KVLayeredAsyncTask::num_layers) + .def("wait", &kv::KVLayeredAsyncTask::wait) + .def("wait_for_layer", &kv::KVLayeredAsyncTask::wait_for_layer) + .def("done", &kv::KVLayeredAsyncTask::done) + .def("result", &kv::KVLayeredAsyncTask::result); + BindHostPagedManager( m, "DefaultHostPagedKVManager"); BindHostPagedManager( diff --git a/docs/prefix-cache-worker-integration-plan.md b/docs/prefix-cache-worker-integration-plan.md new file mode 100644 index 000000000..7b842b749 --- /dev/null +++ b/docs/prefix-cache-worker-integration-plan.md @@ -0,0 +1,964 @@ +# Prefix Cache Worker Integration Plan + +## Scope + +This document reviews the current `feature/add-staged-page-level-prefix-reuse` +code path and defines the remaining work needed to make Host-side prefix cache +reuse active in `BatchGenWorker`. + +The immediate target is GPT-OSS / GQA because its model path already has a +prefix-aware extend-prefill backend. The worker integration should still be +written through generic Host prefix-cache and forward-metadata abstractions so +MLA/SWA/auxiliary groups can reuse the same lifecycle later. + +Prefix cache remains a Host-side shared-memory index over existing Host KV +pages. It does not allocate Host KV pages. KV managers allocate, write, load, +and release physical pages; the coordinator indexes resident page handles, +attaches them during lookup, protects them during Host-to-GPU loads, and returns +evicted handles to the caller. + +## Current Code Review + +### Already Wired + +- Server flags exist: + - `batchgen/server/server_args.py` + - legacy `batchgen/batchgen_server.py` + - flags: `--enable-prefix-cache`, `--prefix-cache-debug-stats` +- Server creates the C++ Host coordinator with `create_region=True`: + - `batchgen/server/worker_manager.py::_initialize_prefix_cache_owner` + - legacy `batchgen/batchgen_server.py::_initialize_prefix_cache_owner` + - The new `server/worker_manager.py` path derives prefix-cache capacity from + `host_kv_cache_size_per_rank`. + - The legacy `batchgen_server.py` path currently derives owner capacity from + the user-facing total Host KV size. This can mismatch the worker attach + config and should be fixed or the legacy path should be disabled for prefix + cache validation. +- Workers attach to the coordinator with `create_region=False`: + - `BatchGenWorker._initialize_prefix_cache_worker` +- Runtime config is derived from Host KV profiles: + - `batchgen/prefix_reuse/config.py` + - GPT-OSS resolves to one required `FULL_KV` group. + - MLA models resolve to `MLA_COMPRESSED_KV`. + - DSA/indexer models can add an auxiliary required group. +- Lookup helper exists: + - `batchgen/prefix_reuse/prefill.py::lookup_prefix_cache_for_prefill` + - `BatchGenWorker._lookup_prefix_cache_for_prefill` +- Estimate-only helper exists: + - `batchgen/prefix_reuse/prefill.py::estimate_prefix_cache_for_prefill` + - `BatchGenWorker._estimate_prefix_cache_for_prefill` +- Prefix-reuse prefill planning exists: + - `batchgen/prefill/prefix_reuse.py` + - current full-hit planning recomputes the final prompt token. This remains + acceptable for the first shared-page implementation as an idempotent + overwrite of already-cached KV. +- First-class prefill metadata can express prefix reuse: + - `batchgen/attention/forward_metadata.py` + - `batchgen/prefill/attention_metadata_builder.py` + - `q_seq_lens`, `kv_seq_lens`, and `append_seq_lens` are separate. +- Legacy wrapper compatibility exists: + - `batchgen/attention/forward_metadata_context.py` + - It mirrors metadata into `AttnWrapperBase` fields for current wrappers. +- GQA compute path exists: + - `batchgen/attention/prefix_aware_backend.py::GqaPrefixAwareAttentionBackend` + - `batchgen/attention/gqa/fa_extend.py::gqa_extend_fa` + - GPT-OSS wrapper calls the prefix-aware backend in prepacked prefill. +- Host KV offload can append only newly computed suffix tokens: + - `batchgen/kv_cache/prefill_offload.py` + - It uses `async_offload_layer_kv_range_to_host` when prefix reuse is active. +- GPU materialization helper exists: + - `batchgen/prefix_reuse/materialization.py` + - It can convert attached Host prefix pages into GPU paged KV pages. +- Commit helper exists: + - `batchgen/prefix_reuse/commit.py` + - It can build aligned `commit_prefix_pages` requests from existing Host KV + page tables. +- API usage has a `cached_tokens` field: + - `batchgen/server/usage.py` + - `SequenceEntry.prefix_shared_tokens` already exists. + +### Current Blocker + +`BatchGenWorker.prefill_prepacked()` still does not run the actual reuse path. + +Current production flow: + +```text +collect full prompt tensors + -> _estimate_prefix_cache_for_prefill(...) + -> prepack full prompt + -> manually set Attn_Wrapper / AttnWrapperBase fields + -> run model + -> offload full prompt KV to Host + -> select first decode token +``` + +The comment in `batchgen_worker.py` explicitly keeps prefix cache in +estimate-only mode until Host sequence KV completeness is solved. That is the +right safety guard: suffix-only prefill is incorrect unless the sequence's Host +KV state also contains the reused prefix pages before the sequence enters +decode or later Host-to-GPU reload. + +### Missing End-to-End Pieces + +- Replace estimate-only lookup with real lookup/attach in the prefix-enabled + prefill admission path. +- Build suffix-only prepack inputs from lookup results. +- Build `ForwardBatchMetadata` for each micro-batch instead of manually setting + parallel wrapper class variables. +- Materialize attached Host prefix pages into GPU paged KV for the current + micro-batch. +- Make the sequence Host KV table logically complete: + - shared prefix pages from the coordinator + - private suffix/decode pages from the sequence allocation +- Ensure decode load/reload sees the complete logical KV, not just private + suffix pages. +- Move lookup early enough to affect Host KV allocation, or explicitly accept a + correctness-only first version that allocates full private Host capacity. +- Commit completed aligned prompt pages into the coordinator after prefill + offload completes. +- Commit aligned prompt+decode pages at request completion before Host pages + are released or recycled. +- Keep lookup attachments alive while any sequence Host page table references + shared prefix pages; release them only when the sequence detaches those pages. +- Feed evicted prefix page handles back to the owning Host KV manager before + relying on those pages as free capacity. +- Populate `cached_tokens` from the page-aligned tokens actually attached and + reused by the worker, not from the raw lookup result. + +## Required Invariants + +### Prefix Lookup + +For every sequence, the coordinator must return an attachable page-boundary +hit. Page alignment is a lookup/admission invariant, not a planner +responsibility. The worker should validate the invariant before attaching pages +and fail loudly if it is violated; it should not silently floor the hit length. + +```text +raw_cached_tokens = coordinator.common_cached_tokens +assert raw_cached_tokens % page_size == 0 +assert raw_cached_tokens <= prompt_length +shared_prefix_tokens = raw_cached_tokens +``` + +For normal partial hits: + +```text +shared_prefix_tokens < prompt_length +query_tokens = prompt[shared_prefix_tokens : prompt_length] +position_ids = range(shared_prefix_tokens, prompt_length) +logical_kv_len = prompt_length +append_tokens = prompt_length - shared_prefix_tokens +usage.cached_tokens = shared_prefix_tokens +``` + +For raw full hits, attach the full prompt but still run the existing one-token +continuation step. In practice this only applies when the full prompt has been +published at the prefix-cache boundary; otherwise the lookup returns the +largest published page-aligned prefix and the request is a partial hit. + +```text +raw_cached_tokens = prompt_length +shared_prefix_tokens = prompt_length +compute_cached_tokens = prompt_length - 1 +query_tokens = [prompt[-1]] +position_ids = [prompt_length - 1] +logical_kv_len = prompt_length +append_tokens = 1 +usage.cached_tokens = shared_prefix_tokens +``` + +This writes the final token KV back to the same logical page that already +contains it. That is intentionally treated as an idempotent overwrite: the +request has the same prompt tokens and the same prefix context, so the produced +KV is semantically the same cached KV. Do not introduce page rollback, overlay +pages, or a separate query-only full-hit path for the first implementation. + +### Host KV Completeness + +Before a sequence transitions from prefill to decode, the Host KV representation +for that sequence must cover the full logical prompt: + +```text +[shared prefix pages] + [private suffix pages] +``` + +GPU materialization alone is not sufficient because GPU pages are transient. +Decode ON_HOLD reload, migration, host eviction/re-entry, and completion commit +all depend on Host KV being the source of truth. + +### Host Allocation Timing + +Current Host KV pages for prefill are allocated before `prefill_prepacked()`. +The allocation code reserves capacity from `seq.prompt_length + chunk_size` +before the worker currently performs the estimate-only prefix lookup. + +That means real prefix reuse cannot simply be inserted inside the existing +`prefill_prepacked()` body if the goal is Host page sharing: + +```text +current order: + allocate private Host pages for full prompt + -> run prefill_prepacked() + -> estimate prefix cache +``` + +For correctness-only validation, this is acceptable if reused prefix pages are +copied into the already allocated private sequence pages. It does not save Host +memory, but it lets compute reuse be tested. + +For the target shared-page design, lookup must move earlier: + +```text +target order: + collect prompt token ids + -> prefix lookup + -> reserve/attach shared prefix pages + -> allocate private Host pages for the existing initial Host KV reserve, + minus attached shared prefix tokens + -> run suffix prefill +``` + +This also affects `SequenceEntry` metadata. Today `host_pages_allocated` and +`host_token_capacity` mean private sequence-owned capacity. With shared prefix +attachment, do not reinterpret those fields as logical capacity. Keep them as +private capacity and add explicit shared-prefix metadata: + +Validation should compare logical capacity as: + +```text +logical_host_tokens = + shared_prefix_tokens + private_host_token_capacity +``` + +Do not silently reinterpret `host_pages_allocated`; it is already used for host +KV pressure planning and release ordering. + +### Sequence Page Layout + +The target design must treat a sequence's Host KV page table as a flat logical +address map. The page table should not own pages and should not need to know +whether a page is shared or private. + +Prefix reuse is page-granular. A shared prefix attachment is valid only when it +is page-aligned: + +```text +shared_prefix_tokens % page_size == 0 +shared_prefix_pages == shared_prefix_tokens / page_size +``` + +If lookup returns a token hit that cannot be represented as full pages, that is +a coordinator/configuration bug. Do not clamp it in the planner, and do not +attach partial pages. + +For a prefix-hit sequence: + +```text +logical Host KV page table + + token range: [0 ................................ prompt_length) + [cached prefix pages] [private suffix pages] + + ownership: prefix coordinator Host KV manager sequence allocation + lifecycle: resident/attached released with the sequence +``` + +In this design, "prepare pages for a sequence" means two different operations: + +- attach existing shared prefix pages returned by the coordinator lookup +- allocate new private pages for suffix and future decode growth + +Shared prefix pages are not allocated again. They are inserted into the +sequence's flat logical Host page table and protected by coordinator +attachments while the sequence uses them. Private pages are allocated by the +Host KV manager and remain owned by that sequence. + +Keep ownership out of `HostKVPageTable`: + +```text +HostKVPageTable: + sequence_id -> [shared prefix page handles..., private page handles...] + no ownership, no shared/private flags + +HostPrefixCacheCoordinator: + owns shared-resident page references, attachment refs, eviction state + +HostPagedKVBackend / allocator: + owns sequence-private pages only +``` + +The worker should not implement this by manually concatenating Python lists of +page ids. Add a per-Host-KV-worker-view API that creates or updates that +manager's sequence logical page table in C++: + +```python +host_worker_view.prepare_sequence_with_shared_prefix( + sequence_id: int, + shared_prefix_pages: Sequence[HostPageHandle], + shared_prefix_tokens: int, + private_token_capacity: int, +) +``` + +For multi-KV-manager models, the worker integration calls the same API on each +required worker view with that group's own pages. Cross-group hit consistency +is enforced by the coordinator/lookup result, not by overloading one page-table +API with group maps. + +Equivalent split APIs are acceptable only if they are used as one transaction: + +```python +host_worker_view.attach_shared_prefix_pages(...) +host_worker_view.allocate_private_pages_for_sequence(...) +``` + +The resulting Host worker view must expose logical page tables for later load +and commit paths: + +```text +build_page_table(sequence_id) + -> [shared page 0, shared page 1, ..., private page 0, private page 1, ...] +``` + +`build_page_table(...)` should not need a new public shape. It can continue to +return the flat logical page vector. The important requirement is that all +logical KV read/write paths use this table instead of asking the backend for +sequence-private pages only. + +Suffix offload with `raw_start_position` must use the original logical token +position. Because shared prefix attachment is page-aligned, normal page-table +indexing is sufficient: + +```text +page_index = raw_start_position / page_size +page_offset = raw_start_position % page_size +target_page = logical_pages[page_index] +``` + +For example, if `shared_prefix_tokens == 128` and `page_size == 64`, suffix +offload at `raw_start_position=128` resolves to `logical_pages[2]`, the first +private page. No special shared/private check is needed in the page table. + +Raw full-hit is the exception to the "suffix writes private pages" intuition: +the one-token continuation has `raw_start_position=prompt_length - 1`, so it +resolves to the last shared prompt page and idempotently overwrites that KV. Do +not allocate a private overlay page for this case. + +Completion and eviction release rules: + +- sequence completion asks the backend to release sequence-private pages, + releases coordinator attachments for shared prefix pages, and removes the + flat logical page-table record +- shared resident prefix pages are only returned to the Host KV manager after + prefix coordinator eviction +- decode commit collects the logical page table, so it can publish chains that + contain both shared prefix pages and private decode pages + +GPU materialization is separate from this Host logical layout. For prefill +compute, `materialize_single_group_lookup_results(...)` allocates temporary GPU +pages for the full logical KV, loads shared Host prefix pages into those GPU +pages, and lets the attention backend append suffix KV. Those GPU pages are +runtime scratch for attention and do not replace the Host sequence page table. + +### Page Ownership + +- KV managers allocate physical Host pages. +- Prefix coordinator stores resident references to already-written pages. +- A page can be resident in prefix cache with zero active lookup/load + references. +- A resident page with zero active references is evictable, not free. +- Only explicit prefix eviction returns page handles to the owning Host KV + manager. +- Sequence cleanup must release private pages but must not free shared resident + prefix pages unless the coordinator evicts them. + +## Recommended Worker Architecture + +### New Worker-Side Integration Helper + +Add a small module instead of expanding `batchgen_worker.py` further: + +```text +batchgen/prefix_reuse/worker_integration.py +``` + +Responsibilities: + +- Convert worker-local batch data into prompt token lists. +- Run coordinator lookup. +- Update `SequenceEntry.prefix_shared_tokens`. +- Keep sequence-level prefix attachment handles until Host shared pages are + detached from the sequence logical page table. +- Build `PrefixReusePrefillPlan`. +- Build suffix-only prepack input lists. +- Slice prefix plans for micro-batches. +- Build `ForwardBatchMetadata` and `KVCacheMetadata`. +- Materialize required compute groups for a micro-batch. +- Release prefix attachments during sequence cleanup, not immediately after the + first GPU materialization. +- Build prompt/decode commit requests. + +Keep side effects explicit. The helper may mutate sequence usage fields and +call coordinator/KV manager APIs, but it should not own the model forward loop. + +### Worker State To Add + +Keep these fields inside `BatchGenWorker`: + +```python +self.prefix_cache_runtime_config +self.prefix_cache_coordinator +self._active_prefix_sequence_attachments +``` + +Do not add user-configurable prefix metadata to worker args. Derived config +stays runtime-only. + +### Main Worker Flow + +Prefix lookup is part of prefill admission/configuration, not part of the +model-forward body. The target order is: + +```text +_prepare_prefill_batch() + -> collect prompt token ids for admitted requests + -> lookup_and_attach prefix cache entries + -> derive per-sequence private Host KV capacity + -> register sequence Host KV tables + -> attach shared prefix pages + -> allocate private suffix/decode pages + -> prefill_prepacked() runs suffix/continuation forward +``` + +This order is required because the private Host page allocation depends on the +attached shared token length returned by lookup. If lookup happens after +`allocate_pages_for_sequences`, the worker has already allocated pages for the +full prompt and cannot realize Host memory sharing. + +After admission/configuration, `prefill_prepacked()` should branch once: + +```text +if not enable_prefix_cache: + run current full-prompt path unchanged +else: + run prefix-aware prepacked path +``` + +Do not mix manual `Attn_Wrapper` assignment with metadata context in the same +prefix path. The prefix path should use `bind_forward_batch_metadata(...)`; the +disabled path can keep the current behavior until it is separately cleaned up. + +## Detailed Prefill Plan + +### Prefill Step 1: Admission Lookup And Plan + +Input: + +- admitted prefill request ids from `_prepare_prefill_batch()` +- collected full `input_ids_list` +- `prompt_lengths` + +Steps: + +1. Call `_lookup_prefix_cache_for_prefill(...)`. +2. Validate the raw hit is page-aligned and within the prompt length. +3. Store it as `SequenceEntry.prefix_shared_tokens`; this value is the + validated page-aligned shared prefix length, not a planner-derived clamp. +4. Attach shared prefix pages and allocate only private suffix/decode pages. +5. Build `PrefixCachePrefillInputs` using + `_build_prefix_reuse_prepack_inputs(...)`. +6. Use `plan.suffix_input_ids` and `plan.suffix_position_ids` as the query + input source. +7. If every `prefix_shared_tokens == 0`, the path may either: + - fall back to the current full-prompt path; or + - keep the unified path with full suffix inputs. + +Recommended first implementation: keep the unified path even on miss. It tests +one path and should produce identical metadata when `prefix_reuse_mode` is +false. + +Placement: + +- For copy fallback, lookup can initially live inside `prefill_prepacked()` + because full private pages are still allocated before the copy. +- For shared-page attachment, this must be lifted into the prefill admission / + Host allocation stage so allocation can reserve private suffix capacity only. +- The target implementation must not perform the first real lookup inside + `prefill_prepacked()`. + +### Prefill Step 2: Suffix Prepack + +Build prepack metadata from suffix inputs: + +```text +prepack_sequences(prefix_inputs.input_ids_list, prefix_inputs.attention_mask_list) +``` + +Use the plan's position ids, not `torch.arange(seq_len)`. Partial-hit suffix +positions start at `prefix_shared_tokens`; raw full-hit continuation starts +at `prompt_length - 1` even though `prefix_shared_tokens == prompt_length`. + +For each sequence in packed order: + +```text +query_len = plan.suffix_length +position_ids = plan.suffix_position_ids +global_sequence_id = plan.sequence_id +``` + +### Prefill Step 3: Micro-Batch Metadata + +For each micro-batch: + +1. Slice `PrefixReusePrefillPlan`. +2. Build spans with global sequence ids in the same order as suffix prepack. +3. Build `ForwardBatchMetadata`: + +```python +ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[...], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=..., + cu_seqlens_k=..., + q_seq_lens=suffix_query_lens, + kv_seq_lens=prompt_lengths, + append_seq_lens=append_lengths, + position_ids=suffix_position_ids, + ), + kv_cache=KVCacheMetadata( + gpu_paged_kv_manager=..., + host_worker_view=..., + aux_gpu_paged_kv_manager=..., + aux_host_worker_view=..., + prefill_prefix_materialization=..., + ), +) +``` + +The prefix path should not manually set: + +- `Attn_Wrapper.prepack_*` +- `AttnWrapperBase.prepack_*` +- `AttnWrapperBase.prefill_prefix_materialization` + +Those should be derived by `bind_forward_batch_metadata(...)`. + +### Prefill Step 4: GPU Prefix Materialization + +For GPT-OSS / GQA first: + +1. Create a temporary `GPUPagedKVCacheManager` or reuse the existing worker GPU + manager if it can safely isolate prefill materialization from active decode + pages. +2. Call `materialize_single_group_lookup_results(...)` for group `0`. +3. Wrap it in `PrefixMaterializationBundle.from_single(0, materialization)`. +4. Pass the bundle in `KVCacheMetadata.prefill_prefix_materialization`. +5. The attention backend will: + - wait for the layer load through `wait_for_layer(layer_idx)` + - append suffix KV into GPU paged KV + - call `gqa_extend_fa(...)`. + +Important implementation choice: + +- If the same `gpu_paged_kv_cache_manager` is used for decode and prefix + prefill, release materialization pages immediately after the micro-batch + forward and rebuild the decode page table if needed. +- If a separate temporary prefill manager is used, keep ownership simpler: + destroy/free it after the micro-batch. This is safer for the first worker + integration. + +### Prefill Step 5: Host KV Table Completeness + +This must be implemented before enabling real suffix-only prefill. + +Preferred API: + +```python +host_worker_view.attach_shared_prefix_pages( + sequence_id: int, + pages: Sequence[HostPageHandle], + prefix_tokens: int, +) +``` + +Semantics: + +- The sequence Host KV table is updated to start with shared resident pages. +- `prefix_tokens` must be page-aligned. +- Partial-hit suffix offload writes into private pages through ordinary logical + `raw_start_position` indexing. +- Raw full-hit one-token continuation may idempotently overwrite the final KV + in a shared page. +- `HostKVPageTable` remains a flat logical page list and does not store + shared/private flags. +- Releasing the sequence releases only backend-owned private pages, drops the + coordinator attachment to shared pages, then removes the flat page-table + entry. +- Shared prefix pages remain resident until prefix eviction. +- Sequence metadata tracks shared-prefix length separately from private Host + capacity, but page-table entries do not need ownership metadata. + +If this API is too invasive, use a temporary correctness fallback: + +```text +copy shared prefix pages into the sequence's private Host KV allocation +``` + +The fallback is slower and loses Host memory sharing, but it proves the compute +path before flat logical page-table attachment lands. + +Fallback requirements: + +- Prefix pages must be copied before suffix offload or before the sequence + enters decode. +- `host_pages_allocated` may remain the full private allocation. +- Commit can read one private Host page table. +- This fallback should be marked temporary because it does not exercise the + resident shared-page lifecycle. + +Do not ship the suffix-only path without either shared-page attachment or copy. +Otherwise decode reload will observe incomplete Host KV. + +### Prefill Step 6: Forward Execution + +Inside the micro-batch loop: + +```python +with bind_forward_batch_metadata(forward_metadata): + inputs_embeds = model.model.embed_tokens(batch_input_ids_flat) + hidden_states = inputs_embeds.unsqueeze(0) + for layer in model.model.layers: + hidden_states = layer(...)[0] +``` + +After the forward: + +- select logits from `batch_cu_seqlens[1:] - 1` +- write the first generated token exactly as the existing path does +- keep `seq.current_context_length = seq.original_prompt_length + seq.decoded_length` + +### Prefill Step 7: Attachment Lifetime + +Lookup attachments protect resident prefix nodes from eviction. With shared +Host page-table attachment, they must cover the whole period where the +sequence's logical Host page table references shared pages, not just the first +Host-to-GPU materialization. + +Load order: + +1. Lookup attaches node. +2. Materialization calls `begin_attachment_load(handle)`. +3. Host-to-GPU load task completes. +4. Materialization calls `end_attachment_load(handle)`. + +The attachment itself remains active after step 4. Release it only when the +sequence detaches shared prefix pages: + +```text +sequence complete / cancelled / migrated away + -> stop using the logical page table entry + -> release private Host pages + -> release prefix-cache attachment handle + -> remove sequence page-table entry +``` + +If the implementation uses the temporary copy fallback instead of shared page +attachment, the attachment may be released after the copied prefix pages and all +dependent GPU loads are complete, because the sequence no longer references +shared resident Host pages. + +Use `try/finally` around prefill admission and forward errors. On failure before +the sequence page table owns the shared pages, release the lookup attachment +immediately. On failure after attachment, run the normal sequence cleanup path. + +## Commit Plan + +### Prompt Commit + +Commit after prefill Host offload tasks are complete. + +Steps: + +1. Retire pending prefill offload tasks. +2. For each owner-local sequence, compute: + +```text +commit_tokens = + floor(prompt_length / publish_boundary_tokens) * publish_boundary_tokens +``` + +3. Collect pages for all required groups: + - primary worker view + - aux worker view if the runtime config has group `1` +4. Call `build_prefix_commit_request(...)`. +5. Call `request.commit(prefix_cache_coordinator)`. +6. On metadata capacity failure, evict unprotected prefix nodes and retry once. + +Partial-hit commit must publish a semantically complete prefix chain. If +shared-attachment mode is used, the page list can include already-shared prefix +pages plus private suffix pages. If copy mode is used, the page list is simply +the sequence's private table. + +If a sequence has no newly computed aligned prompt pages beyond +`prefix_shared_tokens` (for example a raw full hit), prompt commit should be a +no-op for that sequence. Recommitting an already resident chain is unnecessary. + +### Decode Commit + +Decode-generated tokens should enter prefix cache, but only after they are no +longer being mutated. + +First implementation: + +1. At completion, before `_release_host_kv_pages_for_batch(...)`, wait for + pending decode Host KV append tasks. +2. Compute: + +```text +total_tokens = prompt_length + decoded_length +commit_tokens = + floor(total_tokens / publish_boundary_tokens) * publish_boundary_tokens +``` + +3. Commit only full aligned pages. +4. Skip the final partial page. +5. Then run sequence cleanup, which releases private pages and drops shared + prefix attachments. + +Do not commit decode pages at every step initially. Completion-time commit is +simpler and avoids publishing pages that are still being appended. + +## Eviction Integration + +Coordinator eviction returns page handles; it does not release physical Host +pages by itself. + +Worker/server integration must add: + +```text +evicted = coordinator.evict_until_free(...) +for group in evicted.evicted_group_pages: + owning_host_worker_view.release_prefix_resident_pages(group.pages) +``` + +Required semantics: + +- Evict whole prefix nodes, not individual groups. +- Do not evict active lookup/load attachments. +- Do not put resident pages into the normal Host KV free list until the + coordinator has removed the node. +- Host KV allocation pressure and prefix metadata pressure should both be able + to trigger eviction. + +If the current Host KV manager lacks an API to free page handles that are not +attached to a live sequence, add one in C++ rather than faking it in Python. + +## Distributed Behavior + +Initial scope: + +- Prefix cache is per node. +- Each node owns one coordinator shared-memory region. +- Workers on the same node attach to the same region. +- No cross-node prefix sharing. + +For tensor/expert parallel: + +- Every rank that writes a Host KV shard must commit its own pages. +- Lookup token decisions must be deterministic across ranks. +- Page handles are rank/node local. Do not broadcast raw page handles across + nodes. +- Usage accounting can be owner-rank only, but compute materialization must + happen on ranks that execute attention for that sequence. + +## Concrete Implementation Order + +### Step 1: Worker Prefix Path Skeleton + +- Add `batchgen/prefix_reuse/worker_integration.py`. +- Move prompt-token extraction and prefix lookup helpers out of + `batchgen_worker.py`. +- Add a prefix-enabled branch in `prefill_prepacked()`. +- Keep behavior unchanged when disabled. +- Add unit tests for helper-only code. +- Fix legacy server owner/worker config symmetry or explicitly block + `--enable-prefix-cache` on the legacy path until it is fixed. + +### Step 2: Metadata Binding + +- Replace manual wrapper field writes in the prefix branch with + `ForwardBatchMetadata` and `bind_forward_batch_metadata(...)`. +- Use `build_prefill_forward_metadata(...)`. +- Ensure no-prefix metadata still produces the same `AttnWrapperBase` fields in + unit tests. + +### Step 3: Suffix Prepack + +- Prepack `PrefixCachePrefillInputs.input_ids_list`. +- Use `plan.suffix_position_ids` for flattened position ids. +- Cover miss, partial hit, full hit, mixed hit/miss in tests. + +### Step 4: GPU Materialization For GQA + +- Materialize group `0` lookup results for each micro-batch. +- Attach the resulting bundle to `KVCacheMetadata`. +- Use a temporary prefill GPU KV manager first unless reusing the decode manager + can be proven safe. +- Free materialization GPU pages after the micro-batch. +- Add tests with fake materialization managers. + +### Step 5: Host KV Completeness + +- Implement either: + - shared-prefix logical page-table attachment in Host KV worker view; or + - a correctness-only copy fallback. +- Ensure `_load_host_kv_to_gpu(...)` can reload a prefixed sequence and see the + full logical prompt. +- Add integration tests at C++/binding level for shared prefix + private suffix + page tables. +- If logical attachment is implemented, update `SequenceEntry` metadata and + validation so private Host capacity and shared logical prefix length are not + conflated. Do not add shared/private ownership state to `HostKVPageTable`. + +### Step 6: Enable Real Lookup Before Private Host Allocation + +- Required for the target shared-page attachment design. +- Prefix lookup must run during prefill admission/configuration, before + `register_sequences(...)` and `allocate_pages_for_sequences(...)`. +- Allocation should reuse the existing non-prefix Host KV reserve formula. Do + not introduce a new prefix-specific runway knob: + +```text +post_prefill_length = prompt_length + 1 +gpu_initial_pages = ceil(post_prefill_length / page_size) + INITIAL_GPU_PAGE_BUFFER +gpu_initial_tokens = gpu_initial_pages * page_size + +logical_initial_capacity = + min( + max(prompt_length + chunk_size, gpu_initial_tokens), + kv_token_budget, + ) + +private_initial_capacity = + max(logical_initial_capacity - prefix_shared_tokens, append_tokens) + +private_pages = ceil(private_initial_capacity / page_size) +``` + +This preserves the current `chunk_size` and `INITIAL_GPU_PAGE_BUFFER` behavior +while avoiding private allocation for attached shared prefix pages. + +- The sequence Host KV view then attaches shared prefix pages and allocates + private suffix pages. +- Replace `_estimate_prefix_cache_for_prefill(...)` with real + `lookup_and_attach(...)` in the prefix admission branch. +- Keep estimate-only logs as an optional debug mode if useful. +- Store attachment handles on the sequence or worker state until sequence + cleanup. +- Set `SequenceEntry.prefix_shared_tokens` from the validated page-aligned + lookup result. +- For the copy fallback only, moving lookup before private allocation can be + deferred because the fallback still allocates full private Host pages. + +### Step 7: Prompt Commit + +- Wait for prefill Host KV offload tasks. +- Collect aligned prompt pages. +- Commit required groups together. +- Evict/retry on coordinator metadata pressure. +- Add tests for aligned/unaligned prompt lengths and multi-group page lists. + +### Step 8: Decode Commit At Completion + +- Before completed sequence Host page release: + - wait for pending decode append tasks + - commit aligned prompt+decode pages + - then run sequence cleanup, which releases private pages and drops shared + prefix attachments +- Ensure completion reporting includes `cached_tokens`. +- Add tests for completion-time commit ordering. + +### Step 9: Eviction Hook + +- Add a Host KV manager API for releasing evicted resident page handles if + missing. +- Wire coordinator eviction into Host KV allocation pressure and commit retry. +- Add tests that protected attachments are not evicted. + +### Step 10: Remote Validation + +Run in increasing scale: + +1. GPT-OSS prefix disabled: sanity baseline. +2. GPT-OSS prefix enabled, empty cache: miss path. +3. GPT-OSS repeated page-aligned prompts: hit path with nonzero + `cached_tokens`. +4. 20 MMLU Pro requests, short decode, output sanity. +5. 1000 MMLU Pro requests, larger decode. +6. Compare accuracy and output length with main/baseline. +7. Inspect coordinator stats, Host KV stats, `/dev/shm`, and GPU processes + after shutdown. + +## Tests To Add Or Update + +- `tests/unit/test_prefix_worker_integration.py` + - prompt extraction + - lookup result to sequence usage + - suffix prepack source selection + - micro-batch slicing + - attachment release on exception + - attachment handle remains active after materialization when shared pages + are part of the sequence logical page table + - copy fallback may release lookup attachment after copy/load completion +- `tests/unit/test_prefill_attention_metadata_builder.py` + - keep existing mixed hit/full hit cases + - add assertion that raw full hit has one query token, full prompt KV length, + and append length one +- `tests/unit/test_prefix_materialization.py` + - temporary manager release behavior + - bundle group lookup behavior +- `tests/unit/test_prefix_commit_helpers.py` + - prompt commit alignment + - decode completion commit alignment + - multi-group required pages +- C++/binding tests: + - `HostPrefixCacheCoordinator` lookup/commit/evict with GPT-OSS full-KV group + - flat logical Host KV page table containing shared prefix pages followed by + private suffix pages + - eviction returns page handles and does not free active attachments + +## Risks And Open Questions + +- Host logical page-table attachment is the main correctness blocker. Without + either attachment or a copy fallback, suffix-only prefill cannot safely enter + decode. +- Reusing the global decode GPU KV manager for prefill materialization may + disturb active decode page tables. Prefer a temporary prefill manager first. +- Prompt commit must wait for asynchronous prefill offload completion; + otherwise the coordinator may publish pages before all layers are written. +- Completion-time decode commit must run before Host page release. +- Multi-group models require all required groups to hit together. Partial group + hit must be a miss for reuse correctness. +- Auxiliary/indexer groups may be required for later decode correctness even if + the current attention kernel only materializes group `0`. +- Raw full-hit semantics intentionally allow idempotent overwrite of the final + prompt token KV in the shared page. If a future backend cannot tolerate this + write pattern, handle it in that backend; do not make the generic worker path + more complex upfront. + +## Definition Of Done + +Prefix cache is considered worker-integrated for GPT-OSS when all are true: + +- `--enable-prefix-cache` triggers real lookup, not estimate-only. +- Repeated page-aligned prompts report nonzero `cached_tokens`. +- Prefix-hit prefill uses suffix/continuation inputs. +- Host KV remains complete for decode reload and ON_HOLD reload. +- Prompt pages are committed after prefill. +- Decode pages are committed at request completion. +- Shared resident pages are not freed until coordinator eviction. +- Prefix-disabled behavior is unchanged. +- Remote GPT-OSS sanity and MMLU Pro runs complete without output corruption. diff --git a/op_builder/core_engine.py b/op_builder/core_engine.py index 1a03f1c2f..893d15bf4 100644 --- a/op_builder/core_engine.py +++ b/op_builder/core_engine.py @@ -32,6 +32,7 @@ def sources(self): f"{BATCHGEN_CORE_ROOT}/GPU_KV_Buffer/GPU_KV_Buffer.cpp", f"{BATCHGEN_CORE_ROOT}/KV_Storage/host_paged_kv_manager.cpp", f"{BATCHGEN_CORE_ROOT}/KV_Storage/host_paged_kv_backend.cpp", + f"{BATCHGEN_CORE_ROOT}/KV_Storage/host_prefix_cache_coordinator.cpp", f"{BATCHGEN_CORE_ROOT}/KV_Storage/host_paged_kv_worker_view.cpp", f"{BATCHGEN_CORE_ROOT}/KV_Storage/host_kv_page_table.cpp", f"{BATCHGEN_CORE_ROOT}/KV_Storage/uva_copy_kernel.cu", @@ -106,4 +107,4 @@ def extra_ldflags(self): return flags def is_compatible(self, verbose=True): - return super().is_compatible(verbose) \ No newline at end of file + return super().is_compatible(verbose) diff --git a/requirements.txt b/requirements.txt index d951d7156..adb07fbec 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,15 +4,22 @@ aiosignal==1.4.0 annotated-doc==0.0.4 annotated-types==0.7.0 anyio==4.11.0 +apache-tvm-ffi==0.1.11 attrs==25.4.0 certifi==2025.11.12 charset-normalizer==3.4.4 click==8.3.0 +cuda-bindings==13.2.0 +cuda-pathfinder==1.5.4 +cuda-python==13.2.0 +cuda-tile==1.3.0 datasets==2.16.1 dill==0.3.7 einops==0.8.1 fastapi==0.121.1 filelock==3.20.0 +flashinfer-cubin==0.6.11.post1 +flashinfer-python==0.6.11.post1 frozenlist==1.8.0 fsspec==2023.10.0 h11==0.16.0 @@ -34,12 +41,16 @@ nvidia-cuda-cupti-cu12==12.8.90 nvidia-cuda-nvrtc-cu12==12.8.93 nvidia-cuda-runtime-cu12==12.8.90 nvidia-cudnn-cu12==9.10.2.21 +nvidia-cudnn-frontend==1.23.0 nvidia-cufft-cu12==11.3.3.83 nvidia-cufile-cu12==1.13.1.3 nvidia-curand-cu12==10.3.9.90 nvidia-cusolver-cu12==11.7.3.90 nvidia-cusparse-cu12==12.5.8.93 nvidia-cusparselt-cu12==0.7.1 +nvidia-cutlass-dsl-libs-base==4.5.0 +nvidia-cutlass-dsl==4.5.0 +nvidia-ml-py==13.595.45 nvidia-nccl-cu12==2.27.5 nvidia-nvjitlink-cu12==12.8.93 nvidia-nvshmem-cu12==3.3.20 @@ -72,6 +83,7 @@ six==1.17.0 sniffio==1.3.1 starlette==0.49.3 sympy==1.14.0 +tabulate==0.10.0 test-kernel-0ab602a9==0.1.5 tiktoken==0.12.0 tokenizers==0.22.2 diff --git a/tests/integration/paged_kv/test_host_paged_kv_manager.py b/tests/integration/paged_kv/test_host_paged_kv_manager.py index 30de854f3..3fa995ef2 100644 --- a/tests/integration/paged_kv/test_host_paged_kv_manager.py +++ b/tests/integration/paged_kv/test_host_paged_kv_manager.py @@ -148,6 +148,168 @@ def test_parallel_worker_allocate_sequences(): _shm_unlink(shm_name) +def test_worker_view_attaches_shared_prefix_pages_without_owning_them(): + shm_name = _random_shm_name() + cfg = _make_deepseek_r1_config(shm_name) + cfg.num_pages = 32 + worker = bg.MLAHostPagedKVWorkerView(cfg) + + try: + worker.initialize(0, True) + source_seq = 101 + target_seq = 202 + worker.register_sequences([source_seq, target_seq]) + + shared_pages = worker.allocate_pages_for_sequences( + [(source_seq, cfg.page_size_tokens * 2)] + )[0] + worker.attach_shared_prefix_pages(target_seq, shared_pages) + private_pages = worker.allocate_pages_for_sequences( + [(target_seq, cfg.page_size_tokens)] + )[0] + + assert worker.build_page_table([target_seq]) == [ + shared_pages + private_pages + ] + + before_release = worker.get_stats() + worker.release_sequence_pages([target_seq]) + after_release = worker.get_stats() + + assert ( + after_release.num_used_pages + == before_release.num_used_pages - len(private_pages) + ) + assert worker.build_page_table([source_seq]) == [shared_pages] + finally: + for sequence_id in (202, 101): + try: + worker.release_sequence_pages([sequence_id]) + except Exception: + pass + try: + worker.shutdown() + except Exception: + pass + del worker + _shm_unlink(shm_name) + + +def test_worker_view_retains_prefix_resident_pages_until_eviction_release(): + shm_name = _random_shm_name() + cfg = _make_deepseek_r1_config(shm_name) + cfg.num_pages = 16 + worker = bg.MLAHostPagedKVWorkerView(cfg) + + try: + worker.initialize(0, True) + sequence_id = 303 + worker.register_sequences([sequence_id]) + + pages = worker.allocate_pages_for_sequences( + [(sequence_id, cfg.page_size_tokens * 3)] + )[0] + retained = worker.retain_sequence_prefix_pages(sequence_id, 2) + + assert retained == pages[:2] + assert worker.build_page_table([sequence_id]) == [pages] + + before_release = worker.get_stats() + worker.release_sequence_pages([sequence_id]) + after_sequence_release = worker.get_stats() + + assert ( + after_sequence_release.num_used_pages + == before_release.num_used_pages - 1 + ) + assert after_sequence_release.num_active_sequences == 0 + + worker.release_resident_pages(retained) + after_eviction_release = worker.get_stats() + + assert after_eviction_release.num_used_pages == 0 + + grow_sequence_id = 404 + worker.register_sequences([grow_sequence_id]) + prefix_pages = worker.allocate_pages_for_sequences( + [(grow_sequence_id, cfg.page_size_tokens * 2)] + )[0] + retained_prefix = worker.retain_sequence_prefix_pages( + grow_sequence_id, 2 + ) + grown_pages = worker.grow_sequence_pages(grow_sequence_id, 1) + + assert retained_prefix == prefix_pages + assert worker.build_page_table([grow_sequence_id]) == [ + prefix_pages + grown_pages + ] + + before_grow_release = worker.get_stats() + worker.release_sequence_pages([grow_sequence_id]) + after_grow_sequence_release = worker.get_stats() + + assert ( + after_grow_sequence_release.num_used_pages + == before_grow_release.num_used_pages - len(grown_pages) + ) + worker.release_resident_pages(retained_prefix) + assert worker.get_stats().num_used_pages == 0 + + range_sequence_id = 505 + worker.register_sequences([range_sequence_id]) + range_pages = worker.allocate_pages_for_sequences( + [(range_sequence_id, cfg.page_size_tokens * 4)] + )[0] + retained_range = worker.retain_sequence_page_range( + range_sequence_id, 2, 2 + ) + + assert retained_range == range_pages[2:4] + assert worker.build_page_table([range_sequence_id]) == [range_pages] + + before_range_release = worker.get_stats() + worker.release_sequence_pages([range_sequence_id]) + after_range_sequence_release = worker.get_stats() + + assert ( + after_range_sequence_release.num_used_pages + == before_range_release.num_used_pages - 2 + ) + worker.release_resident_pages(retained_range) + assert worker.get_stats().num_used_pages == 0 + + exact_sequence_id = 606 + worker.register_sequences([exact_sequence_id]) + exact_pages = worker.allocate_pages_for_sequences( + [(exact_sequence_id, cfg.page_size_tokens * 4)] + )[0] + retained_exact = worker.retain_sequence_pages( + exact_sequence_id, + [exact_pages[1], exact_pages[3]], + ) + + assert retained_exact == [exact_pages[1], exact_pages[3]] + assert worker.build_page_table([exact_sequence_id]) == [exact_pages] + + before_exact_release = worker.get_stats() + worker.release_sequence_pages([exact_sequence_id]) + after_exact_sequence_release = worker.get_stats() + + assert ( + after_exact_sequence_release.num_used_pages + == before_exact_release.num_used_pages - 2 + ) + worker.release_resident_pages(retained_exact) + assert worker.get_stats().num_used_pages == 0 + finally: + try: + worker.shutdown() + except Exception: + pass + del worker + _shm_unlink(shm_name) + + def _worker_proc_copy_prefill(shm_name, device_index, requests): # 每个进程里重新构造 cfg，shm_name 必须一致 cfg = _make_deepseek_r1_config(shm_name) diff --git a/tests/integration/paged_kv/test_host_prefix_cache_coordinator.py b/tests/integration/paged_kv/test_host_prefix_cache_coordinator.py new file mode 100644 index 000000000..41a631335 --- /dev/null +++ b/tests/integration/paged_kv/test_host_prefix_cache_coordinator.py @@ -0,0 +1,532 @@ +import ctypes +import errno +import random +import string + +from batchgen.models.engine_loader import core_engine as bg + + +_LIBC = ctypes.CDLL("libc.so.6", use_errno=True) + + +def _random_shm_name() -> str: + suffix = "".join( + random.choices(string.ascii_lowercase + string.digits, k=10) + ) + return f"/batchgen_prefix_cache_{suffix}" + + +def _shm_unlink(name: str) -> None: + result = _LIBC.shm_unlink(name.encode("utf-8")) + if result != 0: + err = ctypes.get_errno() + if err != errno.ENOENT: + raise OSError(err, f"shm_unlink({name}) failed") + + +def _group_spec(group_id: int, raw_page_tokens: int): + spec = bg.HostKVGroupSpec() + spec.group_id = group_id + spec.semantic = bg.HostKVGroupSemantic.FULL_KV + spec.required_for_reuse = True + spec.raw_page_tokens = raw_page_tokens + spec.compression_ratio = 1 + return spec + + +def _page(page_id: int): + handle = bg.HostPageHandle() + handle.page_id = page_id + return handle + + +def _group_pages(group_id: int, pages): + group = bg.GroupCommitPages() + group.group_id = group_id + group.pages = list(pages) + return group + + +def _requirement(group_id: int, min_pages: int): + requirement = bg.GroupPageRequirement() + requirement.group_id = group_id + requirement.min_pages = min_pages + return requirement + + +def _config(shm_name: str): + config = bg.HostPrefixCacheConfig() + config.shm_name = shm_name + config.group_specs = [_group_spec(0, 4), _group_spec(1, 8)] + config.max_nodes = 16 + config.max_group_entries = 32 + config.max_page_handles = 128 + config.max_attachments = 16 + return config + + +def _small_config(shm_name: str): + config = _config(shm_name) + config.max_nodes = 2 + config.max_group_entries = 8 + config.max_page_handles = 32 + config.max_attachments = 4 + return config + + +def _single_node_config(shm_name: str): + config = _config(shm_name) + config.max_nodes = 1 + config.max_group_entries = 4 + config.max_page_handles = 16 + config.max_attachments = 2 + return config + + +def _compact_pressure_config(shm_name: str): + config = _single_node_config(shm_name) + config.max_group_entries = 2 + config.max_page_handles = 3 + return config + + +def test_host_prefix_cache_lookup_attach_release(): + shm_name = _random_shm_name() + namespace = [11, 22, 33, 44] + token_ids = list(range(16)) + try: + coordinator = bg.HostPrefixCacheCoordinator(_config(shm_name)) + coordinator.initialize(True) + + assert coordinator.hash_block_tokens == 4 + assert coordinator.commit_boundary_tokens == 8 + + commit = coordinator.commit_prefix_pages( + namespace, + token_ids, + 16, + [ + _group_pages(0, [_page(idx) for idx in range(4)]), + _group_pages(1, [_page(idx) for idx in range(2)]), + ], + ) + assert commit.committed_tokens == 16 + assert commit.inserted_nodes == 2 + assert commit.existing_nodes == 0 + + estimated = coordinator.estimate_lookup(namespace, token_ids[:12]) + assert estimated.common_cached_tokens == 8 + assert estimated.attachment_handle == 0 + assert [span.group_id for span in estimated.materialization_spans] == [ + 0, + 1, + ] + assert coordinator.get_stats().active_attachments == 0 + + attached = coordinator.lookup_and_attach(namespace, token_ids[:12]) + assert attached.common_cached_tokens == 8 + assert attached.attachment_handle != 0 + assert [span.group_id for span in attached.materialization_spans] == [ + 0, + 1, + ] + assert [len(span.pages) for span in attached.materialization_spans] == [ + 2, + 1, + ] + + stats = coordinator.get_stats() + assert stats.resident_nodes == 2 + assert stats.active_attachments == 1 + assert stats.lookup_hits == 1 + assert stats.lookup_misses == 0 + + coordinator.release_attachment(attached.attachment_handle) + assert coordinator.get_stats().active_attachments == 0 + + full = coordinator.lookup_and_attach(namespace, token_ids) + assert full.common_cached_tokens == 16 + assert [ + [page.page_id for page in span.pages] + for span in full.materialization_spans + ] == [ + [0, 1, 2, 3], + [0, 1], + ] + coordinator.release_attachment(full.attachment_handle) + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_evicts_lru_and_preserves_active_attachment(): + shm_name = _random_shm_name() + namespace = [101, 202, 303, 404] + token_ids = list(range(16)) + try: + coordinator = bg.HostPrefixCacheCoordinator(_small_config(shm_name)) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace, + token_ids, + 16, + [ + _group_pages(0, [_page(idx) for idx in range(4)]), + _group_pages(1, [_page(idx) for idx in range(2)]), + ], + ) + + active = coordinator.lookup_and_attach(namespace, token_ids) + assert active.common_cached_tokens == 16 + + evicted = coordinator.evict_until_free(2, 0, 0, 2) + assert evicted.evicted_nodes == 0 + assert evicted.protected_nodes == 2 + assert evicted.freed_group_entries == 0 + assert evicted.freed_page_handles == 0 + assert len(evicted.evicted_group_pages) == 0 + + miss = coordinator.estimate_lookup(namespace, token_ids[:8]) + hit = coordinator.estimate_lookup(namespace, token_ids) + assert miss.common_cached_tokens == 8 + assert hit.common_cached_tokens == 16 + stats = coordinator.get_stats() + assert stats.resident_nodes == 2 + assert stats.used_group_entries == 4 + assert stats.used_page_handles == 6 + + coordinator.release_attachment(active.attachment_handle) + evicted = coordinator.evict_until_free(2, 0, 0, 2) + assert evicted.evicted_nodes == 2 + assert [pages.group_id for pages in evicted.evicted_group_pages] == [ + 0, + 1, + ] + assert [len(pages.pages) for pages in evicted.evicted_group_pages] == [ + 4, + 2, + ] + assert coordinator.get_stats().resident_nodes == 0 + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_evicts_common_nodes_until_pages_releasable(): + shm_name = _random_shm_name() + namespace = [301, 302, 303, 304] + token_ids = list(range(16)) + try: + coordinator = bg.HostPrefixCacheCoordinator(_small_config(shm_name)) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace, + token_ids, + 16, + [ + _group_pages(0, [_page(idx) for idx in range(4)]), + _group_pages(1, [_page(idx) for idx in range(2)]), + ], + ) + + evicted = coordinator.evict_until_releasable_pages( + [_requirement(0, 1)], + 0, + ) + + # Nodes store only their own block interval, so the first LRU node can + # release physical pages immediately. + assert evicted.evicted_nodes == 1 + assert evicted.protected_nodes == 0 + assert [pages.group_id for pages in evicted.evicted_group_pages] == [ + 0, + 1, + ] + assert [ + [page.page_id for page in pages.pages] + for pages in evicted.evicted_group_pages + ] == [ + [0, 1], + [0], + ] + assert coordinator.get_stats().resident_nodes == 1 + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_releases_shared_physical_page_after_last_ref(): + shm_name = _random_shm_name() + namespace_a = [401, 402, 403, 404] + namespace_b = [501, 502, 503, 504] + token_ids = list(range(8)) + try: + coordinator = bg.HostPrefixCacheCoordinator(_config(shm_name)) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace_a, + token_ids, + 8, + [ + _group_pages(0, [_page(42), _page(43)]), + _group_pages(1, [_page(7)]), + ], + ) + coordinator.commit_prefix_pages( + namespace_b, + token_ids, + 8, + [ + _group_pages(0, [_page(42), _page(44)]), + _group_pages(1, [_page(7)]), + ], + ) + + first = coordinator.clear_namespace(namespace_a) + assert first.evicted_nodes == 1 + assert [ + [page.page_id for page in pages.pages] + for pages in first.evicted_group_pages + ] == [[43]] + assert coordinator.estimate_lookup( + namespace_b, + token_ids, + ).common_cached_tokens == 8 + + second = coordinator.clear_namespace(namespace_b) + assert second.evicted_nodes == 1 + assert [ + [page.page_id for page in pages.pages] + for pages in second.evicted_group_pages + ] == [ + [42, 44], + [7], + ] + assert coordinator.get_stats().resident_nodes == 0 + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_clear_skips_active_entries(): + shm_name = _random_shm_name() + namespace = [505, 606, 707, 808] + token_ids = list(range(16)) + try: + coordinator = bg.HostPrefixCacheCoordinator(_small_config(shm_name)) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace, + token_ids, + 16, + [ + _group_pages(0, [_page(idx) for idx in range(4)]), + _group_pages(1, [_page(idx) for idx in range(2)]), + ], + ) + + active = coordinator.lookup_and_attach(namespace, token_ids) + clear = coordinator.clear_unprotected() + assert clear.evicted_nodes == 0 + assert clear.protected_nodes == 2 + assert coordinator.get_stats().resident_nodes == 2 + miss = coordinator.estimate_lookup(namespace, token_ids[:8]) + hit = coordinator.estimate_lookup(namespace, token_ids) + assert miss.common_cached_tokens == 8 + assert hit.common_cached_tokens == 16 + + coordinator.release_attachment(active.attachment_handle) + clear = coordinator.clear_unprotected() + assert clear.evicted_nodes == 2 + assert clear.protected_nodes == 0 + assert coordinator.get_stats().resident_nodes == 0 + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_pending_load_protects_after_release(): + shm_name = _random_shm_name() + namespace = [909, 808, 707, 606] + token_ids = list(range(16)) + try: + coordinator = bg.HostPrefixCacheCoordinator(_small_config(shm_name)) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace, + token_ids, + 16, + [ + _group_pages(0, [_page(idx) for idx in range(4)]), + _group_pages(1, [_page(idx) for idx in range(2)]), + ], + ) + + active = coordinator.lookup_and_attach(namespace, token_ids) + coordinator.begin_attachment_load(active.attachment_handle) + coordinator.release_attachment(active.attachment_handle) + stats = coordinator.get_stats() + assert stats.active_attachments == 1 + assert stats.pending_load_entries == 4 + assert stats.pending_load_refs == 4 + + evicted = coordinator.evict_until_free(2, 0, 0, 2) + assert evicted.evicted_nodes == 0 + assert evicted.protected_nodes == 2 + assert coordinator.get_stats().eviction_protected_skips == 2 + + coordinator.end_attachment_load(active.attachment_handle) + stats = coordinator.get_stats() + assert stats.active_attachments == 0 + assert stats.pending_load_entries == 0 + assert stats.pending_load_refs == 0 + evicted = coordinator.evict_until_free(2, 0, 0, 2) + assert evicted.evicted_nodes == 2 + assert coordinator.get_stats().evicted_nodes == 2 + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_clear_namespace_only_removes_matching_domain(): + shm_name = _random_shm_name() + namespace_a = [1, 3, 5, 7] + namespace_b = [2, 4, 6, 8] + token_ids = list(range(8)) + try: + coordinator = bg.HostPrefixCacheCoordinator(_config(shm_name)) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace_a, + token_ids, + 8, + [ + _group_pages(0, [_page(0), _page(1)]), + _group_pages(1, [_page(0)]), + ], + ) + coordinator.commit_prefix_pages( + namespace_b, + token_ids, + 8, + [ + _group_pages(0, [_page(10), _page(11)]), + _group_pages(1, [_page(10)]), + ], + ) + + cleared = coordinator.clear_namespace(namespace_a) + assert cleared.evicted_nodes == 1 + miss = coordinator.estimate_lookup(namespace_a, token_ids) + hit = coordinator.estimate_lookup(namespace_b, token_ids) + assert miss.miss_reason_mask + assert hit.common_cached_tokens == 8 + assert coordinator.get_stats().resident_nodes == 1 + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_is_shared_across_process_attachments(): + shm_name = _random_shm_name() + namespace = [7, 8, 9, 10] + token_ids = list(range(8)) + try: + owner = bg.HostPrefixCacheCoordinator(_config(shm_name)) + owner.initialize(True) + owner.commit_prefix_pages( + namespace, + token_ids, + 8, + [ + _group_pages(0, [_page(0), _page(1)]), + _group_pages(1, [_page(0)]), + ], + ) + + worker = bg.HostPrefixCacheCoordinator(_config(shm_name)) + worker.initialize(False) + attached = worker.lookup_and_attach(namespace, token_ids) + + assert attached.common_cached_tokens == 8 + assert owner.get_stats().active_attachments == 0 + assert worker.get_stats().active_attachments == 1 + evicted = owner.evict_until_free(16, 0, 0, 1) + assert evicted.evicted_nodes == 0 + assert evicted.protected_nodes == 1 + + worker.release_attachment(attached.attachment_handle) + assert worker.get_stats().active_attachments == 0 + evicted = owner.evict_until_free(16, 0, 0, 1) + assert evicted.evicted_nodes == 1 + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_index_drops_evicted_nodes(): + shm_name = _random_shm_name() + namespace = [17, 18, 19, 20] + token_ids = list(range(8)) + try: + coordinator = bg.HostPrefixCacheCoordinator( + _single_node_config(shm_name) + ) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace, + token_ids, + 8, + [ + _group_pages(0, [_page(0), _page(1)]), + _group_pages(1, [_page(0)]), + ], + ) + + evicted = coordinator.evict_until_free(1, 0, 0, 1) + assert evicted.evicted_nodes == 1 + + miss = coordinator.estimate_lookup(namespace, token_ids) + assert miss.common_cached_tokens == 0 + assert miss.miss_reason_mask + finally: + _shm_unlink(shm_name) + + +def test_host_prefix_cache_compacts_lazily_when_arena_tail_is_full(): + shm_name = _random_shm_name() + namespace = [21, 22, 23, 24] + first_tokens = list(range(8)) + second_tokens = list(range(10, 18)) + try: + coordinator = bg.HostPrefixCacheCoordinator( + _compact_pressure_config(shm_name) + ) + coordinator.initialize(True) + coordinator.commit_prefix_pages( + namespace, + first_tokens, + 8, + [ + _group_pages(0, [_page(0), _page(1)]), + _group_pages(1, [_page(0)]), + ], + ) + + evicted = coordinator.evict_until_free(1, 0, 0, 1) + assert evicted.evicted_nodes == 1 + # The eviction itself does not compact small dead arenas. + assert coordinator.get_stats().used_group_entries == 2 + assert coordinator.get_stats().used_page_handles == 3 + + committed = coordinator.commit_prefix_pages( + namespace, + second_tokens, + 8, + [ + _group_pages(0, [_page(2), _page(3)]), + _group_pages(1, [_page(1)]), + ], + ) + + assert committed.inserted_nodes == 1 + assert coordinator.estimate_lookup( + namespace, + second_tokens, + ).common_cached_tokens == 8 + assert coordinator.get_stats().used_group_entries == 2 + assert coordinator.get_stats().used_page_handles == 3 + finally: + _shm_unlink(shm_name) diff --git a/tests/integration/paged_kv/test_mapped_host_paged_kv_worker_view.py b/tests/integration/paged_kv/test_mapped_host_paged_kv_worker_view.py index f21c1acfc..76f500949 100644 --- a/tests/integration/paged_kv/test_mapped_host_paged_kv_worker_view.py +++ b/tests/integration/paged_kv/test_mapped_host_paged_kv_worker_view.py @@ -72,6 +72,14 @@ def _make_mapped_mla_config(bg, shm_name: str): return cfg +def _make_mapped_default_config(bg, shm_name: str): + cfg = _make_mapped_mla_config(bg, shm_name) + cfg.num_v_heads = NUM_K_HEADS + cfg.v_head_dim = K_HEAD_DIM + cfg.v_element_size_bytes = K_ELEMENT_SIZE_BYTES + return cfg + + def _k_page_bytes() -> int: return PAGE_TOKENS * NUM_K_HEADS * K_HEAD_DIM * K_ELEMENT_SIZE_BYTES @@ -211,6 +219,37 @@ def test_mapped_mla_view_routes_logical_layers_after_cpu_write(bg): _expected_token(expected_value), ) + range_sequence_id = 404 + range_logical_layer = 4 + all_k_ptrs, all_v_ptrs = view.get_sequence_layer_page_pointers( + range_sequence_id, range_logical_layer, None + ) + range_k_ptrs, range_v_ptrs = ( + view.get_sequence_layer_page_range_pointers( + range_sequence_id, range_logical_layer, 1, 2 + ) + ) + assert all_v_ptrs is None + assert range_v_ptrs is None + assert range_k_ptrs == all_k_ptrs[1:3] + + empty_k_ptrs, empty_v_ptrs = ( + view.get_sequence_layer_page_range_pointers( + range_sequence_id, range_logical_layer, len(all_k_ptrs), 0 + ) + ) + assert empty_k_ptrs == [] + assert empty_v_ptrs is None + + with pytest.raises(IndexError): + view.get_sequence_layer_page_range_pointers( + range_sequence_id, 0, 1, 1 + ) + with pytest.raises(IndexError): + view.get_sequence_layer_page_range_pointers( + range_sequence_id, range_logical_layer, len(all_k_ptrs), 1 + ) + with pytest.raises(IndexError): view.resolve_physical_layer(len(LOGICAL_TO_PHYSICAL)) @@ -319,6 +358,143 @@ def test_mapped_mla_view_routes_prefill_and_batched_decode_writes(bg): _shm_unlink(shm_name) +def test_mapped_mla_view_offloads_prefill_range_to_raw_offset(bg): + shm_name = _random_shm_name() + cfg = _make_mapped_mla_config(bg, shm_name) + view = None + sequence_ids = [501, 502] + raw_starts = [PAGE_TOKENS + 3, PAGE_TOKENS * 2 - 2] + token_counts = [5, 7] + capacity_tokens = PAGE_TOKENS * 4 + + try: + torch.cuda.set_device(0) + view = bg.MappedMLAHostPagedKVWorkerView(cfg) + view.initialize(0, True) + view.register_sequences(sequence_ids) + allocations = view.allocate_pages_for_sequences( + [(sequence_id, capacity_tokens) for sequence_id in sequence_ids] + ) + assert all(len(pages) == 4 for pages in allocations) + + device = torch.device("cuda:0") + max_token_count = max(token_counts) + prefill = torch.zeros( + ( + len(sequence_ids), + max_token_count, + NUM_K_HEADS, + K_HEAD_DIM, + ), + dtype=torch.bfloat16, + device=device, + ) + for batch_idx, count in enumerate(token_counts): + prefill[batch_idx, :count].fill_(float(80 + batch_idx)) + + # logical layer 5 routes to physical layer 1. + task = view.async_offload_layer_kv_range_to_host( + 5, + sequence_ids, + prefill, + None, + raw_starts, + token_counts, + ) + task.wait() + + for batch_idx, sequence_id in enumerate(sequence_ids): + k_cpu, v_cpu = view.read_sequence_kv_to_cpu(sequence_id) + assert v_cpu.numel() == 0 + value = float(80 + batch_idx) + for token_idx in range( + raw_starts[batch_idx], + raw_starts[batch_idx] + token_counts[batch_idx], + ): + page_ordinal = token_idx // PAGE_TOKENS + page_offset = token_idx % PAGE_TOKENS + actual = k_cpu[1, page_ordinal, page_offset].flatten() + assert torch.equal(actual, _expected_token(value)) + + before_start = raw_starts[batch_idx] - 1 + actual_before = k_cpu[ + 1, + before_start // PAGE_TOKENS, + before_start % PAGE_TOKENS, + ].flatten() + assert torch.equal(actual_before, torch.zeros_like(actual_before)) + + view.release_sequence_pages(sequence_ids) + view.shutdown() + view = None + finally: + _close_view(view, sequence_ids) + _shm_unlink(shm_name) + + +def test_mapped_default_view_offloads_prefill_range_to_raw_offset(bg): + shm_name = _random_shm_name() + cfg = _make_mapped_default_config(bg, shm_name) + view = None + sequence_ids = [601] + raw_start = PAGE_TOKENS - 2 + token_count = 4 + capacity_tokens = PAGE_TOKENS * 2 + + try: + torch.cuda.set_device(0) + view = bg.MappedDefaultHostPagedKVWorkerView(cfg) + view.initialize(0, True) + assert view.has_v_cache is True + view.register_sequences(sequence_ids) + allocations = view.allocate_pages_for_sequences( + [(sequence_ids[0], capacity_tokens)] + ) + assert len(allocations[0]) == 2 + + device = torch.device("cuda:0") + k_prefill = torch.full( + (1, token_count, NUM_K_HEADS, K_HEAD_DIM), + 91.0, + dtype=torch.bfloat16, + device=device, + ) + v_prefill = torch.full( + (1, token_count, NUM_K_HEADS, K_HEAD_DIM), + 101.0, + dtype=torch.bfloat16, + device=device, + ) + + # logical layer 1 routes to physical layer 3. + task = view.async_offload_layer_kv_range_to_host( + 1, + sequence_ids, + k_prefill, + v_prefill, + [raw_start], + [token_count], + ) + task.wait() + + k_cpu, v_cpu = view.read_sequence_kv_to_cpu(sequence_ids[0]) + assert v_cpu.numel() != 0 + for token_idx in range(raw_start, raw_start + token_count): + page_ordinal = token_idx // PAGE_TOKENS + page_offset = token_idx % PAGE_TOKENS + actual_k = k_cpu[3, page_ordinal, page_offset].flatten() + actual_v = v_cpu[3, page_ordinal, page_offset].flatten() + assert torch.equal(actual_k, _expected_token(91.0)) + assert torch.equal(actual_v, _expected_token(101.0)) + + view.release_sequence_pages(sequence_ids) + view.shutdown() + view = None + finally: + _close_view(view, sequence_ids) + _shm_unlink(shm_name) + + def test_mapped_mla_view_rejects_all_absent_mapping(bg): cfg = _make_mapped_mla_config(bg, _random_shm_name()) cfg.logical_to_physical_layer = [-1, -1, -1] diff --git a/tests/integration/paged_kv/test_prefix_page_materialization.py b/tests/integration/paged_kv/test_prefix_page_materialization.py new file mode 100644 index 000000000..f628370d7 --- /dev/null +++ b/tests/integration/paged_kv/test_prefix_page_materialization.py @@ -0,0 +1,178 @@ +import ctypes +import errno +import math +import random +import string + +import pytest +import torch + +from batchgen.kv_cache.gpu_paged_kv_manager import ( + GPUPagedKVCacheManager, + GPUPagedKVConfig, +) +from batchgen.models.engine_loader import core_engine as bg + + +_LIBC = ctypes.CDLL("libc.so.6", use_errno=True) + + +def _random_shm_name() -> str: + suffix = "".join( + random.choices(string.ascii_lowercase + string.digits, k=10) + ) + return f"/batchgen_prefix_pages_{suffix}" + + +def _shm_unlink(name: str) -> None: + result = _LIBC.shm_unlink(name.encode("utf-8")) + if result != 0: + err = ctypes.get_errno() + if err != errno.ENOENT: + raise OSError(err, f"shm_unlink({name}) failed") + + +def _host_config(shm_name: str) -> bg.HostPagedKVConfig: + cfg = bg.HostPagedKVConfig() + cfg.shm_name = shm_name + cfg.num_layers = 2 + cfg.num_pages = 16 + cfg.page_size_tokens = 4 + cfg.num_k_heads = 1 + cfg.k_head_dim = 2 + cfg.num_v_heads = 1 + cfg.v_head_dim = 2 + cfg.k_element_size_bytes = 2 + cfg.v_element_size_bytes = 2 + cfg.sequence_table_capacity = 16 + cfg.alignment_bytes = 64 + return cfg + + +def _gpu_config() -> GPUPagedKVConfig: + return GPUPagedKVConfig( + num_layers=2, + num_pages=16, + page_size_tokens=4, + num_k_heads=1, + k_head_dim=2, + num_v_heads=1, + v_head_dim=2, + kv_dtype=torch.bfloat16, + ) + + +def _read_sequence_tokens( + manager: GPUPagedKVCacheManager, + *, + sequence_id: int, + layer_idx: int, + length: int, + value_cache: bool, +) -> torch.Tensor: + cache = manager._v_cache if value_cache else manager._k_cache + pages = manager._sequences[sequence_id].pages.tolist() + chunks = [] + remaining = int(length) + for page in pages: + if remaining <= 0: + break + take = min(remaining, manager.config.page_size_tokens) + chunks.append(cache[layer_idx, page, :take].detach().cpu()) + remaining -= take + return torch.cat(chunks, dim=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_async_load_prefix_pages_to_device_uses_host_page_ids(): + shm_name = _random_shm_name() + source_seq = 101 + target_seq = 202 + prefix_tokens = 5 + full_tokens = 7 + page_size = 4 + prefix_pages = math.ceil(prefix_tokens / page_size) + device = torch.device("cuda:0") + torch.cuda.set_device(device) + + host_manager = bg.DefaultHostPagedKVManager(_host_config(shm_name)) + host_manager.initialize(True) + worker = bg.DefaultHostPagedKVWorkerView(_host_config(shm_name)) + worker.initialize(0, False) + + try: + worker.register_sequences([source_seq]) + host_pages = worker.allocate_pages_for_sequences( + [(source_seq, prefix_pages * page_size)] + )[0] + + expected_k = {} + expected_v = {} + for layer_idx in range(2): + base = float(10 * (layer_idx + 1)) + k_tensor = ( + torch.arange( + prefix_tokens * 2, dtype=torch.float32, device=device + ) + .reshape(1, prefix_tokens, 1, 2) + .add(base) + .to(torch.bfloat16) + ) + v_tensor = (k_tensor + 100).contiguous() + expected_k[layer_idx] = k_tensor.detach().cpu().squeeze(0) + expected_v[layer_idx] = v_tensor.detach().cpu().squeeze(0) + task = worker.async_offload_layer_kv_to_host( + layer_idx=layer_idx, + sequence_ids=[source_seq], + k_tensor=k_tensor.contiguous(), + v_tensor=v_tensor, + sequence_lengths=[prefix_tokens], + ) + task.wait() + + gpu_manager = GPUPagedKVCacheManager( + config=_gpu_config(), + device=device, + ) + gpu_manager.initialize() + gpu_manager.allocate_pages_for_sequences([target_seq], [full_tokens]) + gpu_manager.rebuild_page_table([target_seq]) + k_ptrs, v_ptrs = gpu_manager.get_padded_3d_page_pointers() + active_page_counts = torch.tensor([prefix_pages], dtype=torch.int64) + host_page_ids = torch.tensor( + [host_pages[:prefix_pages]], + dtype=torch.int64, + ) + + load_task = worker.async_load_prefix_pages_to_device( + host_page_ids=host_page_ids, + active_page_counts=active_page_counts, + k_device_ptrs=k_ptrs, + v_device_ptrs=v_ptrs, + ) + load_task.wait() + torch.cuda.synchronize(device) + + for layer_idx in range(2): + actual_k = _read_sequence_tokens( + gpu_manager, + sequence_id=target_seq, + layer_idx=layer_idx, + length=prefix_tokens, + value_cache=False, + ) + actual_v = _read_sequence_tokens( + gpu_manager, + sequence_id=target_seq, + layer_idx=layer_idx, + length=prefix_tokens, + value_cache=True, + ) + torch.testing.assert_close(actual_k, expected_k[layer_idx]) + torch.testing.assert_close(actual_v, expected_v[layer_idx]) + finally: + try: + host_manager.free_sequence(source_seq) + except Exception: + pass + _shm_unlink(shm_name) diff --git a/tests/integration/paged_kv/test_transformed_host_paged_kv_worker_view.py b/tests/integration/paged_kv/test_transformed_host_paged_kv_worker_view.py index d083ce49e..016df633d 100644 --- a/tests/integration/paged_kv/test_transformed_host_paged_kv_worker_view.py +++ b/tests/integration/paged_kv/test_transformed_host_paged_kv_worker_view.py @@ -137,7 +137,7 @@ def test_compressed_ratio_host_batched_append_skips_pending_rows(bg): _shm_unlink(shm_name) -def test_swa_host_view_keeps_page_aligned_tail(bg): +def test_swa_host_view_keeps_full_history_and_exposes_window_range(bg): shm_name = _random_shm_name("swa_host") view = None sequence_ids = [201] @@ -150,8 +150,8 @@ def test_swa_host_view_keeps_page_aligned_tail(bg): ) view.initialize(0, True) view.register_sequences(sequence_ids) - allocations = view.allocate_pages_for_sequences([(201, 9)]) - assert len(allocations[0]) == 3 + allocations = view.allocate_pages_for_sequences([(201, 13)]) + assert len(allocations[0]) == 4 first_page = allocations[0][0] token = torch.full( @@ -169,13 +169,41 @@ def test_swa_host_view_keeps_page_aligned_tail(bg): ).wait() page_table = view.build_page_table(sequence_ids) - assert len(page_table[0]) == 3 - assert page_table[0][0] != first_page + assert len(page_table[0]) == 4 + assert page_table[0][0] == first_page + + window_range = view.compute_swa_host_page_range(201, 13) + assert window_range.sequence_id == 201 + assert window_range.raw_context_len == 13 + assert window_range.window_start_token == 5 + assert window_range.first_page == 1 + assert window_range.page_count == 3 + assert window_range.local_kv_len == 9 + assert window_range.mask_start == 1 + + ranges = view.compute_swa_host_page_ranges([201], [13]) + assert len(ranges) == 1 + assert ranges[0].first_page == window_range.first_page + assert ranges[0].page_count == window_range.page_count + + all_k_ptrs, all_v_ptrs = view.get_sequence_layer_page_pointers( + 201, + 1, + None, + ) + window_k_ptrs, window_v_ptrs = ( + view.get_sequence_layer_swa_window_page_pointers( + 201, + 1, + 13, + ) + ) + assert all_v_ptrs is None + assert window_v_ptrs is None + assert window_k_ptrs == all_k_ptrs[1:4] - # raw position 12 maps to page-local active position 8 after the first - # page is released: page ordinal 2, offset 0. k_cpu, _ = view.read_sequence_kv_to_cpu(201) - assert torch.equal(k_cpu[0, 2, 0], _expected_token(42.0)) + assert torch.equal(k_cpu[0, 3, 0], _expected_token(42.0)) finally: _close_view(view, sequence_ids) _shm_unlink(shm_name) diff --git a/tests/integration/test_dynamic_host_kv.py b/tests/integration/test_dynamic_host_kv.py index a315168d7..6eb962790 100644 --- a/tests/integration/test_dynamic_host_kv.py +++ b/tests/integration/test_dynamic_host_kv.py @@ -531,15 +531,12 @@ def test_eviction_reentry_lifecycle(self): new_prompt_len = len(evicted_ids) prev_decoded = seq.total_decoded_before_eviction - # Rebuild input_ids (2D) and attention_mask + # Rebuild input_ids (2D) seq_extended_size = seq.kv_token_budget input_ids_extended = torch.zeros((1, seq_extended_size), dtype=torch.long) - attention_mask_extended = torch.zeros((1, seq_extended_size), dtype=torch.int64) input_ids_extended[0, :new_prompt_len] = evicted_ids - attention_mask_extended[0, :new_prompt_len] = 1 seq.input_ids = input_ids_extended - seq.attention_mask = attention_mask_extended seq.prompt_length = new_prompt_len seq.current_context_length = new_prompt_len @@ -550,19 +547,22 @@ def test_eviction_reentry_lifecycle(self): n_old = min(len(old_decoded), max_decoding_length) seq.decoded_tokens[0, :n_old] = old_decoded[:n_old] seq.decoded_length = n_old + seq.reentry_decoded_baseline = n_old - remaining_decode = seq.original_max_decode_length - prev_decoded - seq.max_decode_length = remaining_decode + assert prev_decoded == n_old + seq.max_decode_length = seq.original_max_decode_length # kv_token_budget stays unchanged seq.evicted_token_ids = None batch.update_status("s1", SequenceStatus.IN_PREFILL) assert seq.prompt_length == 5512 - assert seq.max_decode_length == 32768 - 5000 + assert seq.max_decode_length == 32768 assert seq.decoded_length == 5000 # Pre-filled with old tokens + assert seq.reentry_decoded_baseline == 5000 assert seq.kv_token_budget == 512 + 32768 # Unchanged assert seq.original_prompt_length == 512 # Original preserved assert seq.status == SequenceStatus.IN_PREFILL + seq.validate_metadata("test_eviction_reentry_lifecycle") # Verify pre-filled tokens match original decoded tokens assert torch.equal(seq.decoded_tokens[0, :n_old], decoded[:n_old]) @@ -600,6 +600,26 @@ def test_eviction_reentry_token_write_offset(self): # Old tokens still intact assert torch.equal(seq.decoded_tokens[0, :200], old_decoded) + def test_reentry_uses_sequence_decode_limit_not_stale_worker_limit(self): + """Re-entry accounting must not clamp to a previous pool batch max_tokens.""" + seq = make_seq( + uuid="s1", + prompt_length=1000, + max_decode_length=512, + decoded_length=128, + status=SequenceStatus.EVICTED, + ) + seq.original_prompt_length = 1000 + seq.original_max_decode_length = 512 + seq.total_decoded_before_eviction = 128 + reconstructed_prompt_len = seq.original_prompt_length + seq.total_decoded_before_eviction + + stale_worker_max_decoding_length = 1 + + assert min(seq.total_decoded_before_eviction, stale_worker_max_decoding_length) == 1 + assert seq.compute_reentry_decoded_length(reconstructed_prompt_len) == 128 + assert seq.clamp_reentry_decoded_length(128) == 128 + def test_adaptive_chunk_reduces_waste(self): """Demonstrate that adaptive sizing reduces over-reservation.""" sizer = AdaptiveChunkSizer( diff --git a/tests/test_decode_transition_metadata.py b/tests/test_decode_transition_metadata.py index 9887a5cc2..ffc131f8b 100644 --- a/tests/test_decode_transition_metadata.py +++ b/tests/test_decode_transition_metadata.py @@ -49,19 +49,19 @@ def test_initial_host_kv_capacity_is_page_rounded_before_metadata_validation(): seq.validate_metadata("unit") -def test_synchronous_host_to_gpu_load_uses_dual_dsa_path(): +def test_synchronous_host_to_gpu_load_uses_grouped_kv_path(): source = WORKER.read_text() start = source.index("\tdef _load_host_kv_to_gpu(") end = source.index("\n\tdef _release_gpu_kv_pages", start) body = source[start:end] - dual_branch = body.index("if isinstance(manager, DualKVCacheCoordinator):") + grouped_branch = body.index("if isinstance(manager, GroupedGPUKVCoordinator):") dual_prepare = body.index( "pointers = self._prepare_dual_kv_load_pointers(manager, global_sequence_ids)" ) dual_launch = body.index("load_task = self._launch_dual_host_kv_load(pointers)") primary_only_call = body.index("k_ptrs, v_ptrs = manager.get_padded_3d_page_pointers()") - assert dual_branch < dual_prepare < dual_launch < primary_only_call + assert grouped_branch < dual_prepare < dual_launch < primary_only_call assert "async_load_layer_paged_kv_to_device_dual" not in body assert "host_paged_kv_worker_view_aux" not in body diff --git a/tests/test_flashinfer_mla_extend_prefill.py b/tests/test_flashinfer_mla_extend_prefill.py new file mode 100644 index 000000000..9a54dd372 --- /dev/null +++ b/tests/test_flashinfer_mla_extend_prefill.py @@ -0,0 +1,154 @@ +import sys +import types + +import torch + +_FLASHINFER_STUB = types.ModuleType("flashinfer") +_FLASHINFER_STUB.BatchMLAPagedAttentionWrapper = object +sys.modules.setdefault("flashinfer", _FLASHINFER_STUB) + +from batchgen.attention.mla import flashinfer_extend + + +def test_flashinfer_mla_extend_prefill_builds_wrapper_inputs(monkeypatch): + calls = {} + + class FakeWrapper: + def __init__(self, workspace, backend): + calls["workspace"] = workspace + calls["backend"] = backend + + def plan( + self, + qo_indptr, + kv_indptr, + kv_indices, + kv_len_arr, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + causal, + sm_scale, + q_data_type, + kv_data_type, + ): + calls["plan"] = { + "qo_indptr": qo_indptr.clone(), + "kv_indptr": kv_indptr.clone(), + "kv_indices": kv_indices.clone(), + "kv_len_arr": kv_len_arr.clone(), + "num_heads": num_heads, + "head_dim_ckv": head_dim_ckv, + "head_dim_kpe": head_dim_kpe, + "page_size": page_size, + "causal": causal, + "sm_scale": sm_scale, + "q_data_type": q_data_type, + "kv_data_type": kv_data_type, + } + + def run(self, q_nope, q_pe, ckv_cache, kpe_cache): + calls["run"] = { + "q_nope": q_nope, + "q_pe": q_pe, + "ckv_cache": ckv_cache, + "kpe_cache": kpe_cache, + } + return torch.ones_like(q_nope) + + flashinfer_extend._reset_flashinfer_mla_extend_prefill_cache_for_tests() + monkeypatch.setattr( + flashinfer_extend, + "BatchMLAPagedAttentionWrapper", + FakeWrapper, + ) + + query_states = torch.zeros(1, 3, 2, 6) + compressed_kv_cache = torch.zeros(5, 16, 1, 6) + page_table = torch.tensor( + [ + [3, 4, 1], + [8, 7, 6], + [2, 0, 9], + ], + dtype=torch.int32, + ) + slot_indices = torch.tensor([2, 0], dtype=torch.int32) + cache_seqlens = torch.tensor([17, 33], dtype=torch.int32) + cu_seqlens_q = torch.tensor([0, 1, 3], dtype=torch.int32) + + output = flashinfer_extend.run_flashinfer_mla_extend_prefill( + query_states=query_states, + compressed_kv_cache=compressed_kv_cache, + page_table=page_table, + slot_indices=slot_indices, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + kv_lora_rank=4, + num_heads=2, + softmax_scale=0.25, + ) + + plan = calls["plan"] + assert output.shape == (1, 3, 2, 4) + assert calls["backend"] == "auto" + assert torch.equal(plan["qo_indptr"], cu_seqlens_q) + assert torch.equal( + plan["kv_indptr"], torch.tensor([0, 2, 5], dtype=torch.int32) + ) + assert torch.equal( + plan["kv_indices"], + torch.tensor([2, 0, 3, 4, 1], dtype=torch.int32), + ) + assert torch.equal(plan["kv_len_arr"], cache_seqlens) + assert plan["num_heads"] == 2 + assert plan["head_dim_ckv"] == 4 + assert plan["head_dim_kpe"] == 2 + assert plan["page_size"] == 16 + assert plan["causal"] is True + assert plan["sm_scale"] == 0.25 + assert calls["run"]["q_nope"].shape == (3, 2, 4) + assert calls["run"]["q_pe"].shape == (3, 2, 2) + assert calls["run"]["ckv_cache"].shape == (5, 16, 4) + assert calls["run"]["kpe_cache"].shape == (5, 16, 2) + + +def test_flashinfer_mla_extend_prefill_accepts_full_hit_query_layout( + monkeypatch, +): + calls = {} + + class FakeWrapper: + def __init__(self, workspace, backend): + del workspace, backend + + def plan(self, *args): + calls["plan"] = args + + def run(self, q_nope, q_pe, ckv_cache, kpe_cache): + del q_pe, ckv_cache, kpe_cache + calls["q_nope_shape"] = q_nope.shape + return torch.ones_like(q_nope) + + flashinfer_extend._reset_flashinfer_mla_extend_prefill_cache_for_tests() + monkeypatch.setattr( + flashinfer_extend, + "BatchMLAPagedAttentionWrapper", + FakeWrapper, + ) + + output = flashinfer_extend.run_flashinfer_mla_extend_prefill( + query_states=torch.zeros(2, 1, 3, 6), + compressed_kv_cache=torch.zeros(5, 16, 1, 6), + page_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + slot_indices=torch.tensor([0, 1], dtype=torch.int32), + cache_seqlens=torch.tensor([17, 18], dtype=torch.int32), + cu_seqlens_q=torch.tensor([0, 1, 2], dtype=torch.int32), + kv_lora_rank=4, + num_heads=3, + softmax_scale=0.25, + ) + + assert output.shape == (1, 2, 3, 4) + assert calls["q_nope_shape"] == (2, 3, 4) diff --git a/tests/test_gqa_extend_fa.py b/tests/test_gqa_extend_fa.py new file mode 100644 index 000000000..a7514fd06 --- /dev/null +++ b/tests/test_gqa_extend_fa.py @@ -0,0 +1,49 @@ +import torch + +from batchgen.attention.gqa import fa_extend + + +def test_gqa_extend_fa_passes_paged_extend_metadata(monkeypatch): + calls = {} + + def fake_flash_with_kvcache(*args, **kwargs): + calls["args"] = args + calls["kwargs"] = kwargs + return torch.ones_like(args[0]) + + monkeypatch.setattr(fa_extend, "_USE_FA3", True) + monkeypatch.setattr( + fa_extend, "_flash_with_kvcache", fake_flash_with_kvcache + ) + + q = torch.zeros(5, 4, 8) + k_cache = torch.zeros(3, 64, 1, 8) + v_cache = torch.zeros(3, 64, 1, 8) + cache_seqlens = torch.tensor([66, 67], dtype=torch.int32) + page_table = torch.tensor([[0, 1], [2, -1]], dtype=torch.int32) + cu_q = torch.tensor([0, 2, 5], dtype=torch.int32) + cu_k = torch.tensor([0, 66, 133], dtype=torch.int32) + + output, lse = fa_extend.gqa_extend_fa( + q=q, + k_cache=k_cache, + v_cache=v_cache, + cache_seqlens=cache_seqlens, + page_table=page_table, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=3, + sliding_window=128, + ) + + assert lse is None + assert torch.equal(output, torch.ones_like(q)) + assert calls["args"] == (q, k_cache, v_cache) + assert calls["kwargs"]["page_table"] is page_table + assert calls["kwargs"]["cache_seqlens"] is cache_seqlens + assert calls["kwargs"]["cu_seqlens_q"] is cu_q + assert "cu_seqlens_k_new" not in calls["kwargs"] + assert calls["kwargs"]["max_seqlen_q"] == 3 + assert calls["kwargs"]["causal"] is True + assert calls["kwargs"]["window_size"] == (127, 0) + assert calls["kwargs"]["return_softmax_lse"] is False diff --git a/tests/test_prepack_micro_batches.py b/tests/test_prepack_micro_batches.py index a810200d9..498f56525 100644 --- a/tests/test_prepack_micro_batches.py +++ b/tests/test_prepack_micro_batches.py @@ -25,6 +25,28 @@ def test_build_prefill_micro_batches_can_force_single_sequence_batches(): assert l2_cap == 0 +def test_build_prefill_micro_batches_uses_supplied_admission_lengths(): + micro_batches, l2_cap = build_prefill_micro_batches( + [400, 400, 400, 400], + token_cap=1000, + l2_balance=False, + ) + + assert micro_batches == [(0, 2), (2, 4)] + assert l2_cap == 0 + + +def test_build_prefill_micro_batches_keeps_oversized_sequence_whole(): + micro_batches, l2_cap = build_prefill_micro_batches( + [1200, 100], + token_cap=1000, + l2_balance=False, + ) + + assert micro_batches == [(0, 1), (1, 2)] + assert l2_cap == 0 + + def test_build_prefill_micro_batches_requires_positive_token_cap(): with pytest.raises(ValueError, match="token_cap must be positive"): build_prefill_micro_batches([16, 32], token_cap=0) diff --git a/tests/unit/test_forward_metadata_context.py b/tests/unit/test_forward_metadata_context.py new file mode 100644 index 000000000..f393caeb7 --- /dev/null +++ b/tests/unit/test_forward_metadata_context.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +import pytest +import torch + +from batchgen.attention.forward_metadata import ( + DecodeAttentionMetadata, + ForwardBatchMetadata, + KVCacheMetadata, + PrefillAttentionMetadata, +) +from batchgen.attention.forward_metadata_context import ( + _LEGACY_ATTENTION_FIELDS, + bind_forward_batch_metadata, + get_current_forward_batch_metadata, +) +from batchgen.models.wrappers.attention import AttnWrapperBase + + +@pytest.fixture(autouse=True) +def restore_legacy_attention_fields(): + previous_values = { + field: getattr(AttnWrapperBase, field, None) + for field in _LEGACY_ATTENTION_FIELDS + } + yield + for field, value in previous_values.items(): + setattr(AttnWrapperBase, field, value) + + +def _prefill_metadata(prefix_reuse: bool = True) -> ForwardBatchMetadata: + q_seq_lens = [2, 1, 1] + kv_seq_lens = [5, 1, 4] if prefix_reuse else list(q_seq_lens) + cu_seqlens_k = ( + torch.tensor([0, 5, 6, 10], dtype=torch.int32) + if prefix_reuse + else torch.tensor([0, 2, 3, 4], dtype=torch.int32) + ) + + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[11, 12, 13], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=torch.tensor([0, 2, 3, 4], dtype=torch.int32), + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=2, + max_seqlen_k=max(kv_seq_lens), + q_seq_lens=q_seq_lens, + kv_seq_lens=kv_seq_lens, + position_ids=torch.tensor([3, 4, 0, 3], dtype=torch.int64), + ), + kv_cache=KVCacheMetadata( + gpu_paged_kv_manager=object(), + host_worker_view=object(), + aux_gpu_paged_kv_manager=object(), + aux_host_worker_view=object(), + ), + ) + + +def _decode_metadata() -> ForwardBatchMetadata: + return ForwardBatchMetadata( + phase="decode", + global_sequence_ids=[21, 22], + decode=DecodeAttentionMetadata( + cache_seqlens=torch.tensor([5, 7], dtype=torch.int32), + max_seqlen=7, + page_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + slot_indices=torch.tensor([4, 6], dtype=torch.int64), + ), + ) + + +def _partial_reuse_prefill_metadata() -> ForwardBatchMetadata: + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[31, 32], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=torch.tensor([0, 2, 3], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 5, 6], dtype=torch.int32), + max_seqlen_q=2, + max_seqlen_k=5, + q_seq_lens=[2, 1], + kv_seq_lens=[5, 1], + position_ids=torch.tensor([3, 4, 0], dtype=torch.int64), + ), + ) + + +def test_get_required_raises_when_unbound(): + assert get_current_forward_batch_metadata() is None + with pytest.raises(RuntimeError, match="required"): + get_current_forward_batch_metadata(required=True) + + +def test_bind_forward_batch_metadata_sets_and_restores_legacy_fields(): + AttnWrapperBase.phase = "decode" + AttnWrapperBase.cur_batch = [99] + AttnWrapperBase.prepack_mode = False + AttnWrapperBase.prepack_cu_seqlens = None + AttnWrapperBase.prepack_max_seqlen = None + AttnWrapperBase.prepack_num_sequences = None + AttnWrapperBase.prepack_seq_lengths = None + AttnWrapperBase.prepack_prefix_reuse_mode = False + AttnWrapperBase.prepack_prefix_shared_tokens = None + AttnWrapperBase.prepack_full_seq_lengths = None + + metadata = _prefill_metadata() + with bind_forward_batch_metadata(metadata) as bound: + assert bound is metadata + assert get_current_forward_batch_metadata(required=True) is metadata + assert AttnWrapperBase.phase == "prefill" + assert AttnWrapperBase.cur_batch == [11, 12, 13] + assert AttnWrapperBase.prepack_mode is True + assert ( + AttnWrapperBase.prepack_cu_seqlens is metadata.prefill.cu_seqlens_q + ) + assert AttnWrapperBase.prepack_max_seqlen == 2 + assert AttnWrapperBase.prepack_num_sequences == 3 + assert AttnWrapperBase.prepack_seq_lengths == [2, 1, 1] + assert AttnWrapperBase.prepack_prefix_reuse_mode is True + assert AttnWrapperBase.prepack_prefix_shared_tokens == [3, 0, 3] + assert AttnWrapperBase.prepack_full_seq_lengths == [5, 1, 4] + assert AttnWrapperBase.position_ids is metadata.prefill.position_ids + assert AttnWrapperBase.cache_seqlens is None + assert AttnWrapperBase.max_seqlen is None + assert ( + AttnWrapperBase.gpu_paged_kv_manager + is metadata.kv_cache.gpu_paged_kv_manager + ) + assert ( + AttnWrapperBase.host_paged_kv_worker_view + is metadata.kv_cache.host_worker_view + ) + assert ( + AttnWrapperBase.gpu_paged_kv_manager_aux + is metadata.kv_cache.aux_gpu_paged_kv_manager + ) + assert ( + AttnWrapperBase.host_paged_kv_worker_view_aux + is metadata.kv_cache.aux_host_worker_view + ) + + assert get_current_forward_batch_metadata() is None + assert AttnWrapperBase.phase == "decode" + assert AttnWrapperBase.cur_batch == [99] + assert AttnWrapperBase.prepack_mode is False + assert AttnWrapperBase.prepack_cu_seqlens is None + assert AttnWrapperBase.prepack_prefix_shared_tokens is None + + +def test_bind_forward_batch_metadata_restores_on_exception(): + AttnWrapperBase.phase = "decode" + metadata = _prefill_metadata() + + with pytest.raises(ValueError, match="boom"): + with bind_forward_batch_metadata(metadata): + assert AttnWrapperBase.phase == "prefill" + raise ValueError("boom") + + assert get_current_forward_batch_metadata() is None + assert AttnWrapperBase.phase == "decode" + + +def test_bind_forward_batch_metadata_supports_nested_contexts(): + outer = _prefill_metadata(prefix_reuse=False) + inner = _decode_metadata() + + with bind_forward_batch_metadata(outer): + assert get_current_forward_batch_metadata(required=True) is outer + assert AttnWrapperBase.phase == "prefill" + assert AttnWrapperBase.prepack_prefix_reuse_mode is False + with bind_forward_batch_metadata(inner): + assert get_current_forward_batch_metadata(required=True) is inner + assert AttnWrapperBase.phase == "decode" + assert AttnWrapperBase.prepack_mode is False + assert AttnWrapperBase.cache_seqlens is inner.decode.cache_seqlens + assert AttnWrapperBase.max_seqlen == 7 + + assert get_current_forward_batch_metadata(required=True) is outer + assert AttnWrapperBase.phase == "prefill" + assert AttnWrapperBase.prepack_mode is True + assert AttnWrapperBase.cache_seqlens is None + + +def test_legacy_fields_do_not_leak_across_batches(): + prefix_batch = _prefill_metadata(prefix_reuse=True) + plain_batch = _prefill_metadata(prefix_reuse=False) + + with bind_forward_batch_metadata(prefix_batch): + assert AttnWrapperBase.prepack_prefix_reuse_mode is True + assert AttnWrapperBase.prepack_prefix_shared_tokens == [3, 0, 3] + + with bind_forward_batch_metadata(plain_batch): + assert AttnWrapperBase.prepack_prefix_reuse_mode is False + assert AttnWrapperBase.prepack_prefix_shared_tokens is None + assert AttnWrapperBase.prepack_full_seq_lengths is None + + +def test_prefix_cache_metadata_prefers_bound_forward_metadata(): + class WrapperWithBadLegacyState(AttnWrapperBase): + prepack_cu_seqlens = None + prepack_max_seqlen = None + prepack_num_sequences = None + prepack_seq_lengths = None + cur_batch = None + + wrapper = object.__new__(WrapperWithBadLegacyState) + metadata = _prefill_metadata() + + with bind_forward_batch_metadata(metadata): + prefix_metadata = wrapper.prefix_cache_metadata() + + assert prefix_metadata.global_sequence_ids == [11, 12, 13] + assert prefix_metadata.seq_lengths == [2, 1, 1] + assert prefix_metadata.append_seq_lengths == [2, 1, 1] + assert prefix_metadata.prefix_shared_tokens == [3, 0, 3] + assert prefix_metadata.full_seq_lengths == [5, 1, 4] + assert prefix_metadata.prefix_reuse_mode is True + + +def test_prefix_cache_metadata_rejects_bound_decode_metadata(): + wrapper = object.__new__(AttnWrapperBase) + + with bind_forward_batch_metadata(_decode_metadata()): + with pytest.raises(RuntimeError, match="prefill metadata"): + wrapper.prefix_cache_metadata() + + +def test_prefix_cache_metadata_explicit_matches_legacy_fields(): + from batchgen.models.wrappers.prefix_cache import ( + ensure_prefix_cache_forward_metadata, + ) + + metadata = _partial_reuse_prefill_metadata() + wrapper = object.__new__(AttnWrapperBase) + + with bind_forward_batch_metadata(metadata): + explicit_metadata = wrapper.prefix_cache_metadata() + + assert explicit_metadata.cu_seqlens_list() == [0, 2, 3] + assert explicit_metadata.max_seqlen == 2 + assert explicit_metadata.num_sequences == 2 + assert explicit_metadata.seq_lengths == [2, 1] + expected_append_lens = ( + metadata.prefill.append_seq_lens + if metadata.prefill.append_seq_lens is not None + else metadata.prefill.q_seq_lens + ) + assert explicit_metadata.append_seq_lengths == expected_append_lens + assert explicit_metadata.global_sequence_ids == metadata.global_sequence_ids + assert explicit_metadata.prefix_reuse_mode is True + assert explicit_metadata.prefix_shared_tokens == [3, 0] + assert explicit_metadata.full_seq_lengths == [5, 1] + assert ( + ensure_prefix_cache_forward_metadata(metadata).global_sequence_ids + == metadata.global_sequence_ids + ) + with pytest.raises(RuntimeError, match="global sequence ids"): + ensure_prefix_cache_forward_metadata(metadata.prefill) diff --git a/tests/unit/test_gpt_oss_prefix_reuse_attention.py b/tests/unit/test_gpt_oss_prefix_reuse_attention.py new file mode 100644 index 000000000..05a34c502 --- /dev/null +++ b/tests/unit/test_gpt_oss_prefix_reuse_attention.py @@ -0,0 +1,74 @@ +import pytest +import torch + +from batchgen.models.openai.gpt_oss_120b.wrappers import GptOssAttnWrapper +from batchgen.models.wrappers import AttnWrapperBase + + +@pytest.fixture(autouse=True) +def _reset_prefix_reuse_metadata(): + old_cu = AttnWrapperBase.prepack_cu_seqlens + old_max = AttnWrapperBase.prepack_max_seqlen + old_num = AttnWrapperBase.prepack_num_sequences + old_seq_lengths = AttnWrapperBase.prepack_seq_lengths + old_batch = AttnWrapperBase.cur_batch + old_mode = AttnWrapperBase.prepack_prefix_reuse_mode + old_tokens = AttnWrapperBase.prepack_prefix_shared_tokens + old_lengths = AttnWrapperBase.prepack_full_seq_lengths + yield + AttnWrapperBase.prepack_cu_seqlens = old_cu + AttnWrapperBase.prepack_max_seqlen = old_max + AttnWrapperBase.prepack_num_sequences = old_num + AttnWrapperBase.prepack_seq_lengths = old_seq_lengths + AttnWrapperBase.cur_batch = old_batch + AttnWrapperBase.prepack_prefix_reuse_mode = old_mode + AttnWrapperBase.prepack_prefix_shared_tokens = old_tokens + AttnWrapperBase.prepack_full_seq_lengths = old_lengths + + +def _make_wrapper() -> GptOssAttnWrapper: + wrapper = GptOssAttnWrapper.__new__(GptOssAttnWrapper) + wrapper.layer_idx = 0 + return wrapper + + +def test_gpt_oss_wrapper_no_longer_exposes_host_prefix_kv_reader(): + wrapper = _make_wrapper() + + assert not hasattr(wrapper, "host_prefix_reader") + assert not hasattr(wrapper, "prefix_attention_kv_builder") + + +def test_prefix_cache_metadata_rejects_inconsistent_lengths(): + wrapper = _make_wrapper() + + AttnWrapperBase.prepack_cu_seqlens = torch.tensor([0, 2], dtype=torch.int32) + AttnWrapperBase.prepack_max_seqlen = 2 + AttnWrapperBase.prepack_num_sequences = 1 + AttnWrapperBase.prepack_seq_lengths = [2] + AttnWrapperBase.cur_batch = [101] + AttnWrapperBase.prepack_prefix_reuse_mode = True + AttnWrapperBase.prepack_prefix_shared_tokens = [4] + AttnWrapperBase.prepack_full_seq_lengths = [7] + + with pytest.raises(RuntimeError, match="full length mismatch"): + wrapper.prefix_cache_metadata() + + +def test_clamped_full_hit_metadata_is_normal_prefix_reuse(): + wrapper = _make_wrapper() + + AttnWrapperBase.prepack_cu_seqlens = torch.tensor([0, 1], dtype=torch.int32) + AttnWrapperBase.prepack_max_seqlen = 1 + AttnWrapperBase.prepack_num_sequences = 1 + AttnWrapperBase.prepack_seq_lengths = [1] + AttnWrapperBase.cur_batch = [101] + AttnWrapperBase.prepack_prefix_reuse_mode = True + AttnWrapperBase.prepack_prefix_shared_tokens = [4] + AttnWrapperBase.prepack_full_seq_lengths = [5] + + metadata = wrapper.prefix_cache_metadata() + + assert metadata.prefix_reuse_mode is True + assert metadata.prefix_shared_tokens == [4] + assert metadata.full_seq_lengths == [5] diff --git a/tests/unit/test_gpu_paged_kv_manager_lifecycle.py b/tests/unit/test_gpu_paged_kv_manager_lifecycle.py new file mode 100644 index 000000000..c9dbc0895 --- /dev/null +++ b/tests/unit/test_gpu_paged_kv_manager_lifecycle.py @@ -0,0 +1,77 @@ +import torch + +from batchgen.kv_cache.gpu_paged_kv_manager import ( + GPUPagedKVCacheManager, + GPUPagedKVConfig, +) + + +def test_destroy_releases_cuda_cache_even_after_runtime_state_reset(): + manager = object.__new__(GPUPagedKVCacheManager) + manager._is_initialized = False + calls = [] + + def release_cached_cuda_memory(): + calls.append("released") + + manager._release_cached_cuda_memory = release_cached_cuda_memory + + manager.destroy(empty_cuda_cache=True) + + assert calls == ["released"] + + +def test_destroy_skips_cuda_cache_release_for_uninitialized_noop(): + manager = object.__new__(GPUPagedKVCacheManager) + manager._is_initialized = False + calls = [] + + def release_cached_cuda_memory(): + calls.append("released") + + manager._release_cached_cuda_memory = release_cached_cuda_memory + + manager.destroy(empty_cuda_cache=False) + + assert calls == [] + + +def test_append_prefill_suffix_resolves_logical_layer_mapping(): + config = GPUPagedKVConfig( + num_layers=2, + num_pages=4, + page_size_tokens=4, + num_k_heads=1, + k_head_dim=2, + num_v_heads=1, + v_head_dim=2, + kv_dtype=torch.float32, + logical_to_physical_layer=(0, 1, 0), + ) + manager = GPUPagedKVCacheManager(config=config, device="cpu") + manager.initialize() + manager.allocate_pages_for_sequences([101], [6]) + manager.rebuild_page_table([101]) + plan = manager.prepare_prefill_suffix_append( + sequence_ids=[101], + prefix_lens=[4], + suffix_lens=[2], + rebuild_page_table=False, + ) + + k_tensor = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]) + v_tensor = torch.tensor([[[5.0, 6.0]], [[7.0, 8.0]]]) + manager.append_layer_prefill_suffix_tokens( + k_tensor=k_tensor, + v_tensor=v_tensor, + append_plan=plan, + layer_idx=2, + ) + + k_cache, v_cache = manager.get_kv_tensors() + page_for_token_four = int(manager._sequences[101].pages[1].item()) + assert k_cache[0, page_for_token_four, 0, 0].tolist() == [1.0, 2.0] + assert k_cache[0, page_for_token_four, 1, 0].tolist() == [3.0, 4.0] + assert v_cache[0, page_for_token_four, 0, 0].tolist() == [5.0, 6.0] + assert v_cache[0, page_for_token_four, 1, 0].tolist() == [7.0, 8.0] + assert torch.count_nonzero(k_cache[1]).item() == 0 diff --git a/tests/unit/test_gpu_prefill_suffix_append.py b/tests/unit/test_gpu_prefill_suffix_append.py new file mode 100644 index 000000000..f4ec6c49a --- /dev/null +++ b/tests/unit/test_gpu_prefill_suffix_append.py @@ -0,0 +1,260 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch + + +def _load_gpu_manager_module(): + repo_root = Path(__file__).resolve().parents[2] + config_path = repo_root / "batchgen" / "config" / "config.py" + config_spec = importlib.util.spec_from_file_location( + "_batchgen_config_config_for_gpu_suffix_append_test", + config_path, + ) + config_module = importlib.util.module_from_spec(config_spec) + sys.modules[config_spec.name] = config_module + config_spec.loader.exec_module(config_module) + + previous_config_pkg = sys.modules.get("batchgen.config") + previous_config_module = sys.modules.get("batchgen.config.config") + previous_gpu_kv_kernels = sys.modules.get( + "batchgen.kv_cache.gpu_kv_kernels" + ) + config_pkg = types.ModuleType("batchgen.config") + config_pkg.__path__ = [str(repo_root / "batchgen" / "config")] + config_pkg.config = config_module + sys.modules["batchgen.config"] = config_pkg + sys.modules["batchgen.config.config"] = config_module + + gpu_kv_kernels = types.ModuleType("batchgen.kv_cache.gpu_kv_kernels") + + def _unused_gpu_kernel(*args, **kwargs): + raise RuntimeError( + "GPU KV kernels are not used by this suffix append test" + ) + + gpu_kv_kernels.run_paged_kv_token_update = _unused_gpu_kernel + gpu_kv_kernels.run_paged_kv_token_update_fused = _unused_gpu_kernel + sys.modules["batchgen.kv_cache.gpu_kv_kernels"] = gpu_kv_kernels + + try: + manager_path = ( + repo_root / "batchgen" / "kv_cache" / "gpu_paged_kv_manager.py" + ) + manager_spec = importlib.util.spec_from_file_location( + "_batchgen_gpu_paged_kv_manager_for_suffix_append_test", + manager_path, + ) + manager_module = importlib.util.module_from_spec(manager_spec) + sys.modules[manager_spec.name] = manager_module + manager_spec.loader.exec_module(manager_module) + return manager_module + finally: + if previous_config_pkg is None: + sys.modules.pop("batchgen.config", None) + else: + sys.modules["batchgen.config"] = previous_config_pkg + if previous_config_module is None: + sys.modules.pop("batchgen.config.config", None) + else: + sys.modules["batchgen.config.config"] = previous_config_module + if previous_gpu_kv_kernels is None: + sys.modules.pop("batchgen.kv_cache.gpu_kv_kernels", None) + else: + sys.modules["batchgen.kv_cache.gpu_kv_kernels"] = ( + previous_gpu_kv_kernels + ) + + +_gpu_manager_module = _load_gpu_manager_module() +GPUPagedKVCacheManager = _gpu_manager_module.GPUPagedKVCacheManager +GPUPagedKVConfig = _gpu_manager_module.GPUPagedKVConfig + + +def _make_config( + *, + has_v: bool = True, +) -> GPUPagedKVConfig: + return GPUPagedKVConfig( + num_layers=2, + num_pages=16, + page_size_tokens=4, + num_k_heads=1, + k_head_dim=2, + num_v_heads=1 if has_v else 0, + v_head_dim=2 if has_v else 0, + kv_dtype=torch.float32, + ) + + +def _make_manager(*, has_v: bool = True) -> GPUPagedKVCacheManager: + manager = GPUPagedKVCacheManager( + config=_make_config(has_v=has_v), device="cpu" + ) + manager.initialize() + return manager + + +def _read_sequence_k( + manager: GPUPagedKVCacheManager, sequence_id: int, length: int +): + k_cache, _ = manager.get_kv_tensors() + pages = manager._sequences[sequence_id].pages.tolist() + chunks = [] + remaining = length + for page in pages: + if remaining <= 0: + break + take = min(remaining, manager.config.page_size_tokens) + chunks.append(k_cache[0, page, :take].clone()) + remaining -= take + return torch.cat(chunks, dim=0) + + +def _read_sequence_v( + manager: GPUPagedKVCacheManager, sequence_id: int, length: int +): + _, v_cache = manager.get_kv_tensors() + pages = manager._sequences[sequence_id].pages.tolist() + chunks = [] + remaining = length + for page in pages: + if remaining <= 0: + break + take = min(remaining, manager.config.page_size_tokens) + chunks.append(v_cache[0, page, :take].clone()) + remaining -= take + return torch.cat(chunks, dim=0) + + +def test_prepare_prefill_suffix_append_auto_allocates_miss_sequence(): + manager = _make_manager() + + plan = manager.prepare_prefill_suffix_append( + sequence_ids=[101], + prefix_lens=[0], + suffix_lens=[3], + ) + + assert plan.sequence_ids == [101] + assert plan.prefix_lens.tolist() == [0] + assert plan.suffix_lens.tolist() == [3] + assert plan.cache_seqlens.tolist() == [3] + assert plan.token_starts.tolist() == [0] + assert plan.slot_indices.tolist() == [0] + assert plan.page_table.shape[0] == 1 + assert 101 in manager._sequences + + +def test_prepare_prefill_suffix_append_requires_allocated_reused_prefix(): + manager = _make_manager() + + with pytest.raises(KeyError, match="prefix-reused sequence"): + manager.prepare_prefill_suffix_append( + sequence_ids=[101], + prefix_lens=[2], + suffix_lens=[3], + ) + + +def test_append_layer_prefill_suffix_tokens_writes_across_page_boundary(): + manager = _make_manager() + manager.allocate_pages_for_sequences([101], [7]) + plan = manager.prepare_prefill_suffix_append( + sequence_ids=[101], + prefix_lens=[3], + suffix_lens=[4], + ) + suffix_k = torch.tensor( + [[[1.0, 1.5]], [[2.0, 2.5]], [[3.0, 3.5]], [[4.0, 4.5]]], + dtype=torch.float32, + ) + suffix_v = suffix_k + 10 + + manager.append_layer_prefill_suffix_tokens( + k_tensor=suffix_k, + v_tensor=suffix_v, + append_plan=plan, + layer_idx=0, + ) + + full_k = _read_sequence_k(manager, 101, 7) + full_v = _read_sequence_v(manager, 101, 7) + torch.testing.assert_close(full_k[:3], torch.zeros_like(full_k[:3])) + torch.testing.assert_close(full_v[:3], torch.zeros_like(full_v[:3])) + torch.testing.assert_close(full_k[3:7], suffix_k) + torch.testing.assert_close(full_v[3:7], suffix_v) + + +def test_append_layer_prefill_suffix_tokens_handles_mixed_batch(): + manager = _make_manager() + manager.allocate_pages_for_sequences([101], [5]) + manager.allocate_pages_for_sequences([103], [4]) + plan = manager.prepare_prefill_suffix_append( + sequence_ids=[101, 102, 103], + prefix_lens=[3, 0, 4], + suffix_lens=[2, 3, 0], + ) + suffix_k = torch.arange(10, dtype=torch.float32).view(5, 1, 2) + suffix_v = suffix_k + 100 + + manager.append_layer_prefill_suffix_tokens( + k_tensor=suffix_k, + v_tensor=suffix_v, + append_plan=plan, + layer_idx=0, + ) + + torch.testing.assert_close( + _read_sequence_k(manager, 101, 5)[3:5], suffix_k[:2] + ) + torch.testing.assert_close( + _read_sequence_v(manager, 101, 5)[3:5], suffix_v[:2] + ) + torch.testing.assert_close(_read_sequence_k(manager, 102, 3), suffix_k[2:5]) + torch.testing.assert_close(_read_sequence_v(manager, 102, 3), suffix_v[2:5]) + torch.testing.assert_close( + _read_sequence_k(manager, 103, 4), torch.zeros(4, 1, 2) + ) + + +def test_append_layer_prefill_suffix_tokens_accepts_mla_2d_k_tensor(): + manager = _make_manager(has_v=False) + plan = manager.prepare_prefill_suffix_append( + sequence_ids=[101], + prefix_lens=[0], + suffix_lens=[2], + ) + suffix_k = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) + + manager.append_layer_prefill_suffix_tokens( + k_tensor=suffix_k, + v_tensor=None, + append_plan=plan, + layer_idx=0, + ) + + torch.testing.assert_close( + _read_sequence_k(manager, 101, 2), + suffix_k.unsqueeze(1), + ) + + +def test_append_layer_prefill_suffix_tokens_rejects_bad_token_count(): + manager = _make_manager() + plan = manager.prepare_prefill_suffix_append( + sequence_ids=[101], + prefix_lens=[0], + suffix_lens=[2], + ) + + with pytest.raises(ValueError, match="token count mismatch"): + manager.append_layer_prefill_suffix_tokens( + k_tensor=torch.zeros(3, 1, 2), + v_tensor=None, + append_plan=plan, + layer_idx=0, + ) diff --git a/tests/unit/test_host_kv_group_profiles.py b/tests/unit/test_host_kv_group_profiles.py new file mode 100644 index 000000000..3424dc73f --- /dev/null +++ b/tests/unit/test_host_kv_group_profiles.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from batchgen.kv_cache.host_kv_mananger_config import ( + build_gpu_kv_config_from_group_profile, + resolve_host_kv_group_profiles, +) + + +def test_deepseek_v4_group_profiles_capture_storage_and_raw_rates(): + profiles = resolve_host_kv_group_profiles("deepseek-v4-flash") + + assert [ + ( + profile.group_id, + profile.group_name, + profile.storage_page_tokens, + profile.raw_page_tokens, + profile.compression_ratio, + ) + for profile in profiles + ] == [ + (0, "swa", 64, 64, 1), + (1, "compressor_c4", 64, 256, 4), + (2, "compressor_c128", 2, 256, 128), + (3, "indexer_c4", 64, 256, 4), + ] + + +def test_compressed_group_gpu_config_uses_storage_page_capacity(): + profiles = { + profile.group_name: profile + for profile in resolve_host_kv_group_profiles("deepseek-v4-flash") + } + + swa_config = build_gpu_kv_config_from_group_profile(profiles["swa"], [1024]) + c4_config = build_gpu_kv_config_from_group_profile( + profiles["compressor_c4"], [1024] + ) + c128_config = build_gpu_kv_config_from_group_profile( + profiles["compressor_c128"], [1024] + ) + + assert swa_config.page_size_tokens == 64 + assert swa_config.num_pages == 17 + assert c4_config.page_size_tokens == 64 + assert c4_config.num_pages == 5 + assert c128_config.page_size_tokens == 2 + assert c128_config.num_pages == 5 diff --git a/tests/unit/test_moe_fused_gate_calls.py b/tests/unit/test_moe_fused_gate_calls.py new file mode 100644 index 000000000..151979543 --- /dev/null +++ b/tests/unit/test_moe_fused_gate_calls.py @@ -0,0 +1,39 @@ +import ast +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _moe_fused_gate_calls(relative_path: str) -> list[ast.Call]: + tree = ast.parse((REPO_ROOT / relative_path).read_text(encoding="utf-8")) + calls: list[ast.Call] = [] + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + if node.func.id == "moe_fused_gate": + calls.append(node) + return calls + + +def _assert_moe_fused_gate_wrapper_signature(relative_path: str) -> None: + calls = _moe_fused_gate_calls(relative_path) + assert calls, f"expected at least one moe_fused_gate call in {relative_path}" + for call in calls: + positional = [ast.unparse(arg) for arg in call.args] + keywords = {keyword.arg for keyword in call.keywords} + + assert len(positional) == 5 + assert not any("n_routed_experts" in arg for arg in positional) + assert "routed_scaling_factor" in keywords + + +def test_deepseek_moe_fused_gate_uses_python_wrapper_signature(): + _assert_moe_fused_gate_wrapper_signature( + "batchgen/models/deepseek/deepseekv3/modeling_deepseek_v3.py" + ) + + +def test_kimi_asset_moe_fused_gate_uses_python_wrapper_signature(): + _assert_moe_fused_gate_wrapper_signature( + "batchgen/models/moonshotai/kimi_k25/assets/modeling_deepseek.py" + ) diff --git a/tests/unit/test_prefill_attention_metadata_builder.py b/tests/unit/test_prefill_attention_metadata_builder.py new file mode 100644 index 000000000..f8aede4b9 --- /dev/null +++ b/tests/unit/test_prefill_attention_metadata_builder.py @@ -0,0 +1,264 @@ +from __future__ import annotations + +import pytest +import torch + +from batchgen.batch_order import PrefillSequenceSpan +from batchgen.prefill.attention_metadata_builder import ( + build_prefill_forward_metadata, +) +from batchgen.prefill.prepack import PrepackMetadata, prepack_sequences +from batchgen.prefill.prefix_reuse import ( + PrefixReusePrefillPlan, + PrefixReuseSequencePlan, +) + + +def _span( + row_index: int, global_seq_id: int, seq_len: int +) -> PrefillSequenceSpan: + return PrefillSequenceSpan( + row_index=row_index, + local_idx=10 + row_index, + uuid=f"uuid-{row_index}", + global_seq_id=global_seq_id, + seq_len=seq_len, + start=0, + end=seq_len, + ) + + +def _spans( + global_ids: list[int], seq_lens: list[int] +) -> list[PrefillSequenceSpan]: + cursor = 0 + spans = [] + for row_index, (global_seq_id, seq_len) in enumerate( + zip(global_ids, seq_lens) + ): + spans.append( + PrefillSequenceSpan( + row_index=row_index, + local_idx=10 + row_index, + uuid=f"uuid-{row_index}", + global_seq_id=global_seq_id, + seq_len=seq_len, + start=cursor, + end=cursor + seq_len, + ) + ) + cursor += seq_len + return spans + + +def _prepack_metadata(seq_lens: list[int]) -> PrepackMetadata: + return PrepackMetadata( + packed_input_ids=torch.empty((0,), dtype=torch.long), + packed_attention_mask=torch.empty((0,), dtype=torch.long), + packed_position_ids=torch.empty((0,), dtype=torch.long), + sequence_ids=torch.empty((0,), dtype=torch.long), + cu_seqlens_per_row=[], + max_seqlen_per_row=[], + original_seq_lengths=seq_lens, + num_original_sequences=len(seq_lens), + num_packed_rows=0, + row_length=max(seq_lens, default=0), + pack_assignment=[], + ) + + +def _prefix_plan( + global_ids: list[int], + prefix_lens: list[int], + suffix_lens: list[int], +) -> PrefixReusePrefillPlan: + sequences = [] + suffix_input_ids = [] + suffix_position_ids = [] + for local_idx, (global_id, prefix_len, suffix_len) in enumerate( + zip(global_ids, prefix_lens, suffix_lens) + ): + prompt_length = prefix_len + suffix_len + sequences.append( + PrefixReuseSequencePlan( + local_idx=local_idx, + sequence_id=global_id, + prompt_length=prompt_length, + prefix_shared_tokens=prefix_len, + suffix_start_pos=prefix_len, + suffix_length=suffix_len, + full_logical_context_length=prompt_length, + ) + ) + suffix_input_ids.append(torch.arange(suffix_len, dtype=torch.long)) + suffix_position_ids.append( + torch.arange(prefix_len, prompt_length, dtype=torch.long) + ) + + return PrefixReusePrefillPlan( + sequences=sequences, + suffix_input_ids=suffix_input_ids, + suffix_position_ids=suffix_position_ids, + cache_seqlens=torch.tensor(prefix_lens, dtype=torch.int32), + total_prompt_tokens=sum( + prefix + suffix for prefix, suffix in zip(prefix_lens, suffix_lens) + ), + total_suffix_tokens=sum(suffix_lens), + saved_prefill_tokens=sum(prefix_lens), + ) + + +def _prefix_lens(metadata) -> list[int]: + return [ + int(kv_len) - int(append_len) + for append_len, kv_len in zip( + metadata.prefill.append_seq_lens, + metadata.prefill.kv_seq_lens, + ) + ] + + +def test_build_prefill_forward_metadata_without_prefix_reuse(): + prepack = prepack_sequences( + [ + torch.tensor([[1, 2, 3]], dtype=torch.long), + torch.tensor([[4, 5]], dtype=torch.long), + ], + [ + torch.tensor([[1, 1, 1]], dtype=torch.long), + torch.tensor([[1, 1]], dtype=torch.long), + ], + device=torch.device("cpu"), + ) + + metadata = build_prefill_forward_metadata( + prepack_metadata=prepack, + batch_spans=_spans([100, 101], [3, 2]), + seq_start=0, + seq_end=2, + position_ids=torch.tensor([0, 1, 2, 0, 1], dtype=torch.long), + device=torch.device("cpu"), + ) + + assert metadata.phase == "prefill" + assert metadata.global_sequence_ids == [100, 101] + assert metadata.prefill.q_seq_lens == [3, 2] + assert metadata.prefill.append_seq_lens == [3, 2] + assert metadata.prefill.kv_seq_lens == [3, 2] + assert metadata.prefill.cu_seqlens_q.tolist() == [0, 3, 5] + assert metadata.prefill.cu_seqlens_k.tolist() == [0, 3, 5] + assert _prefix_lens(metadata) == [0, 0] + + +def test_build_prefill_forward_metadata_with_prefix_reuse_slice(): + prepack = _prepack_metadata([99, 2, 1, 88]) + plan = _prefix_plan( + global_ids=[90, 100, 101, 91], + prefix_lens=[0, 3, 0, 0], + suffix_lens=[99, 2, 1, 88], + ) + + metadata = build_prefill_forward_metadata( + prepack_metadata=prepack, + batch_spans=_spans([100, 101], [2, 1]), + seq_start=1, + seq_end=3, + position_ids=torch.tensor([3, 4, 0], dtype=torch.long), + device=torch.device("cpu"), + prefix_reuse_plan=plan, + ) + + assert metadata.prefill.q_seq_lens == [2, 1] + assert metadata.prefill.append_seq_lens == [2, 1] + assert metadata.prefill.kv_seq_lens == [5, 1] + assert metadata.prefill.cu_seqlens_q.tolist() == [0, 2, 3] + assert metadata.prefill.cu_seqlens_k.tolist() == [0, 5, 6] + assert _prefix_lens(metadata) == [3, 0] + + +def test_build_prefill_forward_metadata_with_mixed_hit_miss_and_full_hit(): + prepack = _prepack_metadata([2, 1, 1]) + plan = _prefix_plan( + global_ids=[100, 101, 102], + prefix_lens=[3, 0, 3], + suffix_lens=[2, 1, 1], + ) + + metadata = build_prefill_forward_metadata( + prepack_metadata=prepack, + batch_spans=_spans([100, 101, 102], [2, 1, 1]), + seq_start=0, + seq_end=3, + position_ids=torch.tensor([3, 4, 0, 3], dtype=torch.long), + device=torch.device("cpu"), + prefix_reuse_plan=plan, + ) + + assert metadata.prefill.q_seq_lens == [2, 1, 1] + assert metadata.prefill.append_seq_lens == [2, 1, 1] + assert metadata.prefill.kv_seq_lens == [5, 1, 4] + assert metadata.prefill.cu_seqlens_q.tolist() == [0, 2, 3, 4] + assert metadata.prefill.cu_seqlens_k.tolist() == [0, 5, 6, 10] + assert _prefix_lens(metadata) == [3, 0, 3] + + +def test_build_prefill_forward_metadata_one_token_full_hit_is_plain_query(): + prepack = _prepack_metadata([1]) + plan = _prefix_plan( + global_ids=[100], + prefix_lens=[0], + suffix_lens=[1], + ) + + metadata = build_prefill_forward_metadata( + prepack_metadata=prepack, + batch_spans=_spans([100], [1]), + seq_start=0, + seq_end=1, + position_ids=torch.tensor([0], dtype=torch.long), + device=torch.device("cpu"), + prefix_reuse_plan=plan, + ) + + assert metadata.prefill.q_seq_lens == [1] + assert metadata.prefill.append_seq_lens == [1] + assert metadata.prefill.kv_seq_lens == [1] + assert metadata.prefill.cu_seqlens_q.tolist() == [0, 1] + assert metadata.prefill.cu_seqlens_k.tolist() == [0, 1] + assert _prefix_lens(metadata) == [0] + + +def test_build_prefill_forward_metadata_allows_shorter_append_length(): + prepack = _prepack_metadata([3]) + plan = _prefix_plan(global_ids=[100], prefix_lens=[2], suffix_lens=[1]) + + metadata = build_prefill_forward_metadata( + prepack_metadata=prepack, + batch_spans=[_span(0, 100, 3)], + seq_start=0, + seq_end=1, + position_ids=torch.tensor([2, 3, 4], dtype=torch.long), + device=torch.device("cpu"), + prefix_reuse_plan=plan, + ) + + assert metadata.prefill.q_seq_lens == [3] + assert metadata.prefill.append_seq_lens == [1] + assert metadata.prefill.kv_seq_lens == [3] + assert _prefix_lens(metadata) == [2] + + +def test_build_prefill_forward_metadata_rejects_sequence_id_mismatch(): + prepack = _prepack_metadata([1]) + plan = _prefix_plan(global_ids=[200], prefix_lens=[0], suffix_lens=[1]) + + with pytest.raises(ValueError, match="sequence ids"): + build_prefill_forward_metadata( + prepack_metadata=prepack, + batch_spans=[_span(0, 100, 1)], + seq_start=0, + seq_end=1, + position_ids=torch.tensor([0], dtype=torch.long), + device=torch.device("cpu"), + prefix_reuse_plan=plan, + ) diff --git a/tests/unit/test_prefill_offload_retire.py b/tests/unit/test_prefill_offload_retire.py new file mode 100644 index 000000000..6226ae63d --- /dev/null +++ b/tests/unit/test_prefill_offload_retire.py @@ -0,0 +1,166 @@ +from types import SimpleNamespace + +import torch + +from batchgen.attention.forward_metadata import ( + ForwardBatchMetadata, + PrefillAttentionMetadata, +) +from batchgen.models.wrappers.attention import AttnWrapperBase + + +class _FakeTask: + def __init__(self): + self.wait_calls = 0 + + def wait(self): + self.wait_calls += 1 + + +class _FakeHostWorkerView: + def __init__(self): + self.task = _FakeTask() + self.range_calls = [] + + def async_offload_layer_kv_range_to_host(self, **kwargs): + self.range_calls.append(kwargs) + return self.task + + +class _FakePrefixMaterialization: + def __init__(self): + self.finished_layers = [] + + def finish_layer(self, layer_idx): + self.finished_layers.append(int(layer_idx)) + + +def _metadata(*, append_len: int) -> ForwardBatchMetadata: + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[101], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=torch.tensor([0, 2], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 6], dtype=torch.int32), + max_seqlen_q=2, + max_seqlen_k=6, + q_seq_lens=[2], + kv_seq_lens=[6], + position_ids=torch.tensor([4, 5], dtype=torch.int64), + append_seq_lens=[append_len], + ), + ) + + +def _reset_pending_state() -> None: + AttnWrapperBase.pending_prefill_offload_tasks = [] + AttnWrapperBase.pending_prefill_offload_tensors = [] + AttnWrapperBase.pending_prefill_offload_layer_idx = None + AttnWrapperBase.prefill_prefix_materialization = None + + +def test_prefix_reuse_finish_layer_waits_for_tracked_prefill_offload(): + _reset_pending_state() + host_view = _FakeHostWorkerView() + materialization = _FakePrefixMaterialization() + wrapper = object.__new__(AttnWrapperBase) + wrapper.layer_idx = 7 + wrapper.core_engine = SimpleNamespace(host_paged_kv_worker_view=host_view) + AttnWrapperBase.prefill_prefix_materialization = materialization + + key = torch.ones(2, 1, 4) + value = torch.ones(2, 1, 4) + wrapper.offload_prepacked_gqa_kv( + key, + value, + metadata=_metadata(append_len=2), + track_tasks=False, + ) + + assert materialization.finished_layers == [] + assert AttnWrapperBase.pending_prefill_offload_layer_idx == 7 + assert len(AttnWrapperBase.pending_prefill_offload_tasks) == 1 + assert len(AttnWrapperBase.pending_prefill_offload_tensors) >= 2 + assert host_view.range_calls[0]["raw_start_positions"] == [4] + assert host_view.range_calls[0]["token_counts"] == [2] + + AttnWrapperBase.retire_pending_prefill_offloads(device=None) + + assert host_view.task.wait_calls == 1 + assert materialization.finished_layers == [7] + assert AttnWrapperBase.pending_prefill_offload_layer_idx is None + assert AttnWrapperBase.pending_prefill_offload_tasks == [] + assert AttnWrapperBase.pending_prefill_offload_tensors == [] + + _reset_pending_state() + + +def test_prefix_reuse_zero_append_finishes_layer_on_retire(): + _reset_pending_state() + host_view = _FakeHostWorkerView() + materialization = _FakePrefixMaterialization() + wrapper = object.__new__(AttnWrapperBase) + wrapper.layer_idx = 3 + wrapper.core_engine = SimpleNamespace(host_paged_kv_worker_view=host_view) + AttnWrapperBase.prefill_prefix_materialization = materialization + + key = torch.ones(2, 1, 4) + value = torch.ones(2, 1, 4) + wrapper.offload_prepacked_gqa_kv( + key, + value, + metadata=_metadata(append_len=0), + track_tasks=False, + ) + + assert host_view.range_calls == [] + assert materialization.finished_layers == [] + assert AttnWrapperBase.pending_prefill_offload_layer_idx == 3 + assert AttnWrapperBase.pending_prefill_offload_tasks == [] + assert len(AttnWrapperBase.pending_prefill_offload_tensors) == 2 + + AttnWrapperBase.retire_pending_prefill_offloads(device=None) + + assert materialization.finished_layers == [3] + assert AttnWrapperBase.pending_prefill_offload_layer_idx is None + + _reset_pending_state() + + +def test_subclass_prefill_offload_state_retires_through_base_wrapper(): + class _ModelWrapper(AttnWrapperBase): + pass + + _reset_pending_state() + host_view = _FakeHostWorkerView() + materialization = _FakePrefixMaterialization() + wrapper = object.__new__(_ModelWrapper) + wrapper.layer_idx = 5 + wrapper.core_engine = SimpleNamespace(host_paged_kv_worker_view=host_view) + AttnWrapperBase.prefill_prefix_materialization = materialization + + key = torch.ones(2, 1, 4) + wrapper.offload_prepacked_mla_kv( + key, + metadata=_metadata(append_len=2), + track_tasks=False, + ) + + assert "pending_prefill_offload_layer_idx" not in _ModelWrapper.__dict__ + assert AttnWrapperBase.pending_prefill_offload_layer_idx == 5 + assert len(AttnWrapperBase.pending_prefill_offload_tasks) == 1 + assert len(AttnWrapperBase.pending_prefill_offload_tensors) >= 2 + + retired = AttnWrapperBase.retire_pending_prefill_offloads_before_layer( + 6, + device=None, + ) + + assert retired == 1 + assert host_view.task.wait_calls == 1 + assert materialization.finished_layers == [5] + assert AttnWrapperBase.pending_prefill_offload_layer_idx is None + assert AttnWrapperBase.pending_prefill_offload_tasks == [] + assert AttnWrapperBase.pending_prefill_offload_tensors == [] + + _reset_pending_state() diff --git a/tests/unit/test_prefix_aware_backend.py b/tests/unit/test_prefix_aware_backend.py new file mode 100644 index 000000000..20b0b57d2 --- /dev/null +++ b/tests/unit/test_prefix_aware_backend.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from batchgen.attention.forward_metadata import ( + ForwardBatchMetadata, + KVCacheMetadata, + PrefillAttentionMetadata, +) +from batchgen.attention.forward_metadata_context import ( + bind_forward_batch_metadata, +) +from batchgen.attention.prefix_aware_backend import ( + GqaPrefixAwareAttentionBackend, +) +from batchgen.models.wrappers.prefix_gqa_extend import ( + GqaExtendSpec, + run_prefix_gqa_prefill_attention, +) +from batchgen.prefix_reuse.materialization import PrefixMaterializationBundle + +_LAYER_IDX = 2 + + +def _metadata( + *, + prefix_reuse: bool = False, +) -> ForwardBatchMetadata: + cu_seqlens = torch.tensor([0, 2], dtype=torch.int32) + max_seqlen = 2 + seq_lengths = [2] + kv_seq_lengths = [5] if prefix_reuse else list(seq_lengths) + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[100], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=torch.tensor( + [0, kv_seq_lengths[0]], + dtype=torch.int32, + ), + max_seqlen_q=max_seqlen, + max_seqlen_k=max(kv_seq_lengths), + q_seq_lens=seq_lengths, + kv_seq_lens=kv_seq_lengths, + position_ids=torch.tensor([0, 1], dtype=torch.int64), + append_seq_lens=seq_lengths, + ), + ) + + +def _clamped_full_hit_metadata() -> ForwardBatchMetadata: + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[100], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 5], dtype=torch.int32), + max_seqlen_q=1, + max_seqlen_k=5, + q_seq_lens=[1], + kv_seq_lens=[5], + position_ids=torch.tensor([4], dtype=torch.int64), + append_seq_lens=[1], + ), + ) + + +def test_gqa_backend_no_prefix_uses_query_cu_seqlens_for_kv(): + recorded = {} + + def attention_fn(**kwargs): + recorded.update(kwargs) + return kwargs["q"] + 1, None + + backend = GqaPrefixAwareAttentionBackend( + layer_idx=_LAYER_IDX, + num_kv_heads=1, + head_dim=2, + attention_fn=attention_fn, + ) + query = torch.zeros((2, 2, 2)) + key = torch.ones((2, 1, 2)) + value = key + 10 + + output = backend.forward_prefill( + query=query, + key=key, + value=value, + metadata=_metadata(), + ) + + torch.testing.assert_close(output, query + 1) + assert recorded["k"] is key + assert recorded["v"] is value + assert recorded["cu_seqlens_q"].tolist() == [0, 2] + assert recorded["cu_seqlens_k"].tolist() == [0, 2] + assert recorded["max_seqlen_q"] == 2 + assert recorded["max_seqlen_k"] == 2 + + +def test_gqa_backend_prefix_reuse_requires_gpu_materialization(): + backend = GqaPrefixAwareAttentionBackend( + layer_idx=_LAYER_IDX, + num_kv_heads=1, + head_dim=2, + ) + query = torch.zeros((2, 2, 2)) + key = torch.ones((2, 1, 2)) + value = key + 10 + + with pytest.raises(RuntimeError, match="GPU paged materialization"): + backend.forward_prefill( + query=query, + key=key, + value=value, + metadata=_metadata(prefix_reuse=True), + ) + + +class _FakeGqaMaterializedManager: + def __init__(self): + self.k_cache = torch.zeros((4, 4, 1, 2)) + self.v_cache = torch.ones((4, 4, 1, 2)) + self.append_calls = [] + self.page_table = torch.tensor( + [ + [0, 1], + [2, 3], + ], + dtype=torch.int32, + ) + + def get_layer_kv_with_page_table(self, layer_idx): + assert layer_idx == _LAYER_IDX + return self.k_cache, self.v_cache, self.page_table + + def append_layer_prefill_suffix_tokens(self, **kwargs): + self.append_calls.append(kwargs) + + +class _FakeGqaMaterialization: + def __init__(self): + self.manager = _FakeGqaMaterializedManager() + self.append_plan = SimpleNamespace( + slot_values=torch.tensor([0], dtype=torch.int32), + cache_seqlens=torch.tensor([5], dtype=torch.int32), + ) + self.waited_layers = [] + + def wait_for_layer(self, layer_idx): + self.waited_layers.append(int(layer_idx)) + + +def test_gqa_backend_clamped_full_hit_uses_extend_prefill(monkeypatch): + recorded = {} + + import batchgen.attention.gqa as gqa + + def fake_extend(**kwargs): + recorded.update(kwargs) + return kwargs["q"] + 10, None + + monkeypatch.setattr(gqa, "gqa_extend_fa", fake_extend) + + materialization = _FakeGqaMaterialization() + backend = GqaPrefixAwareAttentionBackend( + layer_idx=_LAYER_IDX, + num_kv_heads=1, + head_dim=2, + ) + query = torch.arange(4, dtype=torch.float32).reshape(1, 2, 2) + key = torch.ones((1, 1, 2)) + value = key + 1 + + output = backend.forward_prefill( + query=query, + key=key, + value=value, + metadata=_clamped_full_hit_metadata(), + kv_cache_metadata=SimpleNamespace( + prefill_prefix_materialization=PrefixMaterializationBundle( + by_group_id={0: materialization} + ) + ), + ) + + torch.testing.assert_close(output, query + 10) + assert materialization.waited_layers == [_LAYER_IDX] + assert materialization.manager.append_calls[0]["k_tensor"] is key + assert materialization.manager.append_calls[0]["v_tensor"] is value + assert recorded["q"] is query + assert recorded["k_cache"] is materialization.manager.k_cache + assert recorded["v_cache"] is materialization.manager.v_cache + assert recorded["cache_seqlens"].tolist() == [5] + + +class _FakeGqaExtendWrapper: + layer_idx = _LAYER_IDX + + +def test_gqa_extend_passes_bound_kv_cache_metadata(monkeypatch): + recorded = {} + + def fake_forward_prefill(self, **kwargs): + recorded.update(kwargs) + return kwargs["query"] + + monkeypatch.setattr( + GqaPrefixAwareAttentionBackend, + "forward_prefill", + fake_forward_prefill, + ) + kv_cache = KVCacheMetadata( + prefill_prefix_materialization=object(), + ) + forward_metadata = ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[100], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=torch.tensor([0, 2], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 5], dtype=torch.int32), + max_seqlen_q=2, + max_seqlen_k=5, + q_seq_lens=[2], + kv_seq_lens=[5], + position_ids=torch.tensor([3, 4], dtype=torch.int64), + ), + kv_cache=kv_cache, + ) + query = torch.zeros((2, 2, 2)) + key = torch.ones((2, 1, 2)) + value = key + 10 + + with bind_forward_batch_metadata(forward_metadata): + output = run_prefix_gqa_prefill_attention( + wrapper=_FakeGqaExtendWrapper(), + query=query, + key=key, + value=value, + metadata=_metadata(prefix_reuse=True), + spec=GqaExtendSpec(num_kv_heads=1, head_dim=2), + ) + + assert output is query + assert recorded["kv_cache_metadata"] is kv_cache + + +def test_gqa_backend_missing_value_raises(): + backend = GqaPrefixAwareAttentionBackend( + layer_idx=_LAYER_IDX, + num_kv_heads=1, + head_dim=2, + attention_fn=lambda **kwargs: (kwargs["q"], None), + ) + + with pytest.raises(RuntimeError, match="value tensor"): + backend.forward_prefill( + query=torch.zeros((1, 1, 2)), + key=torch.zeros((1, 1, 2)), + value=None, + metadata=_metadata(), + ) + + +def test_gqa_backend_missing_metadata_raises(): + backend = GqaPrefixAwareAttentionBackend( + layer_idx=_LAYER_IDX, + num_kv_heads=1, + head_dim=2, + attention_fn=lambda **kwargs: (kwargs["q"], None), + ) + + with pytest.raises(TypeError, match="metadata"): + backend.forward_prefill( + query=torch.zeros((1, 1, 2)), + key=torch.zeros((1, 1, 2)), + value=torch.zeros((1, 1, 2)), + metadata=object(), + ) diff --git a/tests/unit/test_prefix_cache_admin.py b/tests/unit/test_prefix_cache_admin.py new file mode 100644 index 000000000..73a4a362a --- /dev/null +++ b/tests/unit/test_prefix_cache_admin.py @@ -0,0 +1,176 @@ +from types import SimpleNamespace + +from batchgen.prefix_reuse.admin import ( + clear_host_prefix_cache, + host_kv_views_by_prefix_group, +) + + +class _Coordinator: + def __init__(self): + self._stats = [ + _stats(resident_nodes=3, used_group_entries=4), + _stats(resident_nodes=0, evicted_nodes=3), + ] + self.clear_calls = 0 + + def get_stats(self): + return self._stats.pop(0) + + def clear_unprotected(self): + self.clear_calls += 1 + return SimpleNamespace( + evicted_nodes=3, + protected_nodes=0, + freed_group_entries=4, + freed_page_handles=5, + evicted_group_pages=[ + _group_pages(0, [10, 11, 11]), + _group_pages(1, [20, 21]), + ], + ) + + +class _PinCoordinator: + def __init__(self): + self.lookup_calls = [] + self.released_handles = [] + self._next_handle = 100 + + def lookup_and_attach(self, namespace_digest, token_ids): + self.lookup_calls.append((list(namespace_digest), list(token_ids))) + if not token_ids or token_ids[0] < 0: + return SimpleNamespace( + attachment_handle=0, + common_cached_tokens=0, + ) + self._next_handle += 1 + return SimpleNamespace( + attachment_handle=self._next_handle, + common_cached_tokens=len(token_ids), + ) + + def release_attachment(self, handle): + self.released_handles.append(int(handle)) + + +class _HostKV: + def __init__(self): + self.released_pages = [] + + def release_resident_pages(self, page_ids): + self.released_pages.append(list(page_ids)) + + +class _GroupedHostKV: + def __init__(self, views): + self._views = views + + def views_by_group(self): + return self._views + + +def _stats(**overrides): + values = { + "resident_nodes": 0, + "active_attachments": 0, + "pending_load_entries": 0, + "pending_load_refs": 0, + "used_group_entries": 0, + "used_page_handles": 0, + "lookup_hits": 0, + "lookup_misses": 0, + "evicted_nodes": 0, + "eviction_protected_skips": 0, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _group_pages(group_id, pages): + return SimpleNamespace( + group_id=group_id, + pages=[SimpleNamespace(page_id=page_id) for page_id in pages], + ) + + +def test_clear_host_prefix_cache_releases_evicted_pages_by_group(): + primary = _HostKV() + auxiliary = _HostKV() + coordinator = _Coordinator() + + result = clear_host_prefix_cache( + coordinator=coordinator, + host_kv_views_by_group={0: primary, 1: auxiliary}, + ) + + assert coordinator.clear_calls == 1 + assert primary.released_pages == [[10, 11]] + assert auxiliary.released_pages == [[20, 21]] + assert result["cleared_all"] is True + assert result["stats_before"]["resident_nodes"] == 3 + assert result["stats_after"]["resident_nodes"] == 0 + assert result["eviction"]["evicted_nodes"] == 3 + assert result["eviction"]["evicted_pages_by_group"] == {0: 3, 1: 2} + assert result["eviction"]["released_pages_by_group"] == {0: 2, 1: 2} + + +def test_host_kv_views_by_prefix_group_uses_grouped_coordinator_first(): + views = {0: _HostKV(), 3: _HostKV()} + grouped = _GroupedHostKV(views) + + assert ( + host_kv_views_by_prefix_group( + primary_host_kv=grouped, + auxiliary_host_kv=_HostKV(), + ) + == views + ) + + +def test_host_kv_views_by_prefix_group_maps_primary_and_auxiliary(): + primary = _HostKV() + auxiliary = _HostKV() + + assert host_kv_views_by_prefix_group( + primary_host_kv=primary, + auxiliary_host_kv=auxiliary, + ) == {0: primary, 1: auxiliary} + + +def test_pin_host_prefix_cache_holds_only_lookup_hits(): + from batchgen.prefix_reuse.admin import pin_host_prefix_cache + + coordinator = _PinCoordinator() + + result = pin_host_prefix_cache( + coordinator=coordinator, + namespace_digest=(1, 2, 3, 4), + token_id_batches=[[10, 11, 12], [-1, 2], [20]], + ) + + assert result["requested"] == 3 + assert result["pinned"] == 2 + assert result["missed"] == 1 + assert result["cached_tokens"] == 4 + assert result["cached_tokens_by_request"] == [3, 1] + assert result["attachment_handles"] == [101, 102] + assert coordinator.lookup_calls == [ + ([1, 2, 3, 4], [10, 11, 12]), + ([1, 2, 3, 4], [-1, 2]), + ([1, 2, 3, 4], [20]), + ] + + +def test_unpin_host_prefix_cache_releases_handles(): + from batchgen.prefix_reuse.admin import unpin_host_prefix_cache + + coordinator = _PinCoordinator() + + result = unpin_host_prefix_cache( + coordinator=coordinator, + attachment_handles=[101, 102], + ) + + assert result == {"status": "success", "released": 2} + assert coordinator.released_handles == [101, 102] diff --git a/tests/unit/test_prefix_cache_config.py b/tests/unit/test_prefix_cache_config.py new file mode 100644 index 000000000..2f0bd4b8c --- /dev/null +++ b/tests/unit/test_prefix_cache_config.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from batchgen.prefix_reuse.config import ( + PrefixKVGroupSemantic, + PrefixKVGroupSpec, + build_prefix_cache_namespace_digest, + build_prefix_cache_runtime_config, + build_prefix_cache_runtime_config_from_specs, + create_host_prefix_cache_coordinator, + derive_prefix_cache_shm_name, +) +from batchgen.server.server_args import _build_parser + + +def test_prefix_cache_runtime_config_derives_boundaries_and_capacities(): + config = build_prefix_cache_runtime_config_from_specs( + model_name="test/model", + kv_dtype="bfloat16", + host_kv_pages_per_required_group=128, + group_specs=[ + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.MLA_COMPRESSED_KV, + required_for_reuse=True, + raw_page_tokens=64, + ), + PrefixKVGroupSpec( + group_id=1, + semantic=PrefixKVGroupSemantic.SWA_KV, + required_for_reuse=True, + raw_page_tokens=128, + ), + ], + ) + + assert config.hash_block_tokens == 64 + assert config.publish_boundary_tokens == 128 + assert config.max_nodes >= 1024 + assert config.max_group_entries == config.max_nodes * 2 + assert config.max_page_handles >= config.max_group_entries + assert config.max_attachments >= 1024 + + +def test_prefix_cache_runtime_config_uses_multi_rate_kv_groups(): + config = build_prefix_cache_runtime_config( + model_name="deepseek-v4-flash", + kv_dtype="bfloat16", + host_kv_cache_size_bytes=1 << 30, + ) + + assert config.hash_block_tokens == 64 + assert config.publish_boundary_tokens == 256 + assert [ + ( + spec.group_id, + spec.semantic, + spec.raw_page_tokens, + spec.compression_ratio, + ) + for spec in config.group_specs + ] == [ + (0, PrefixKVGroupSemantic.SWA_KV, 64, 1), + (1, PrefixKVGroupSemantic.COMPRESSED_RATIO_KV, 256, 4), + (2, PrefixKVGroupSemantic.COMPRESSED_RATIO_KV, 256, 128), + (3, PrefixKVGroupSemantic.COMPRESSED_RATIO_KV, 256, 4), + ] + + +def test_prefix_cache_namespace_digest_is_stable_and_group_sensitive(): + group = PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=True, + raw_page_tokens=64, + ) + same = build_prefix_cache_namespace_digest( + model_name="OpenAI/GPT-OSS-120B", + kv_dtype="bfloat16", + group_specs=[group], + ) + reordered_case = build_prefix_cache_namespace_digest( + model_name="openai/gpt-oss-120b", + kv_dtype="BFLOAT16", + group_specs=[group], + ) + changed = build_prefix_cache_namespace_digest( + model_name="openai/gpt-oss-120b", + kv_dtype="bfloat16", + group_specs=[ + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=True, + raw_page_tokens=128, + ) + ], + ) + + assert same == reordered_case + assert same != changed + assert len(same) == 4 + + +def test_prefix_cache_core_config_conversion_uses_bound_classes(): + class _CoreGroupSpec(SimpleNamespace): + pass + + class _CoreConfig(SimpleNamespace): + pass + + class _Core: + HostKVGroupSpec = _CoreGroupSpec + HostPrefixCacheConfig = _CoreConfig + HostKVGroupSemantic = SimpleNamespace( + FULL_KV="full", + MLA_COMPRESSED_KV="mla", + SWA_KV="swa", + COMPRESSED_RATIO_KV="compressed", + ) + + config = build_prefix_cache_runtime_config_from_specs( + model_name="test/model", + kv_dtype="bfloat16", + host_kv_pages_per_required_group=8, + group_specs=[ + PrefixKVGroupSpec( + group_id=3, + semantic=PrefixKVGroupSemantic.COMPRESSED_RATIO_KV, + required_for_reuse=False, + raw_page_tokens=256, + compression_ratio=4, + ), + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=True, + raw_page_tokens=64, + ), + ], + ) + + core_config = config.to_core_config(_Core) + + assert core_config.shm_name == config.shm_name + assert core_config.hash_block_tokens == 64 + assert len(core_config.group_specs) == 2 + assert core_config.group_specs[0].group_id == 3 + assert core_config.group_specs[0].semantic == "compressed" + assert core_config.group_specs[0].compression_ratio == 4 + + +def test_create_host_prefix_cache_coordinator_initializes_requested_region(): + class _CoreGroupSpec(SimpleNamespace): + pass + + class _CoreConfig(SimpleNamespace): + pass + + class _Coordinator: + instances = [] + + def __init__(self, config): + self.config = config + self.initialize_calls = [] + self.instances.append(self) + + def initialize(self, create_region): + self.initialize_calls.append(bool(create_region)) + + class _Core: + HostKVGroupSpec = _CoreGroupSpec + HostPrefixCacheConfig = _CoreConfig + HostPrefixCacheCoordinator = _Coordinator + HostKVGroupSemantic = SimpleNamespace( + FULL_KV="full", + MLA_COMPRESSED_KV="mla", + SWA_KV="swa", + COMPRESSED_RATIO_KV="compressed", + ) + + runtime_config = build_prefix_cache_runtime_config_from_specs( + model_name="test/model", + kv_dtype="bfloat16", + host_kv_pages_per_required_group=8, + group_specs=[ + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=True, + raw_page_tokens=64, + ) + ], + ) + + coordinator = create_host_prefix_cache_coordinator( + core_engine_module=_Core, + runtime_config=runtime_config, + create_region=True, + ) + + assert coordinator.initialize_calls == [True] + assert coordinator.config.shm_name == runtime_config.shm_name + + +def test_prefix_cache_runtime_config_rejects_no_required_group(): + with pytest.raises(ValueError, match="required KV group"): + build_prefix_cache_runtime_config_from_specs( + model_name="test/model", + kv_dtype="bfloat16", + host_kv_pages_per_required_group=8, + group_specs=[ + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=False, + raw_page_tokens=64, + ) + ], + ) + + +def test_prefix_cache_shm_name_is_sanitized_and_node_agnostic(): + shm_name = derive_prefix_cache_shm_name("Org/Model-Name") + + assert shm_name.startswith("batchgen_prefix_cache_org_model_name_") + assert "_node" not in shm_name + + +def test_server_parser_exposes_only_prefix_cache_user_flags(): + parsed = _build_parser().parse_args( + [ + "--model", + "openai/gpt-oss-120b", + "--enable-prefix-cache", + "--prefix-cache-debug-stats", + ] + ) + + assert parsed.enable_prefix_cache is True + assert parsed.prefix_cache_debug_stats is True + assert not hasattr(parsed, "prefix_cache_size_gb") + assert not hasattr(parsed, "prefix_cache_hash_block_tokens") diff --git a/tests/unit/test_prefix_cache_wrapper_helpers.py b/tests/unit/test_prefix_cache_wrapper_helpers.py new file mode 100644 index 000000000..01de20f87 --- /dev/null +++ b/tests/unit/test_prefix_cache_wrapper_helpers.py @@ -0,0 +1,163 @@ +import importlib +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +class _FakeCuSeqlens: + def __init__(self, values): + self._values = list(values) + + def __len__(self): + return len(self._values) + + def detach(self): + return self + + def cpu(self): + return self + + def tolist(self): + return list(self._values) + + +class _FakeSeqTensor: + def __init__(self, name, dim=3): + self.name = name + self._dim = dim + + def unsqueeze(self, dim): + return _FakeSeqTensor(f"{self.name}.unsqueeze({dim})", self._dim + 1) + + def dim(self): + return self._dim + + +class _FakeFlatTensor: + def __init__(self, name, dim=3): + self.name = name + self._dim = dim + + def __getitem__(self, key): + return _FakeSeqTensor(f"{self.name}[{key.start}:{key.stop}]", self._dim) + + +class _FakeWorkerView: + def __init__(self): + self.calls = [] + + def async_offload_layer_kv_to_host(self, **kwargs): + self.calls.append(("normal", kwargs)) + return SimpleNamespace(done=lambda: True, wait=lambda: None) + + def async_offload_layer_kv_range_to_host(self, **kwargs): + self.calls.append(("offset", kwargs)) + return SimpleNamespace(done=lambda: True, wait=lambda: None) + + +class _NoOffsetWorkerView: + def async_offload_layer_kv_to_host(self, **kwargs): + del kwargs + return None + + +def _install_torch_stub(monkeypatch): + torch_stub = types.ModuleType("torch") + torch_stub.Tensor = object + torch_stub.bfloat16 = "bfloat16" + torch_stub.float16 = "float16" + torch_stub.int32 = "int32" + monkeypatch.setitem(sys.modules, "torch", torch_stub) + batchgen_stub = types.ModuleType("batchgen") + batchgen_stub.__path__ = [str(REPO_ROOT / "batchgen")] + monkeypatch.setitem(sys.modules, "batchgen", batchgen_stub) + models_stub = types.ModuleType("batchgen.models") + models_stub.__path__ = [str(REPO_ROOT / "batchgen" / "models")] + monkeypatch.setitem(sys.modules, "batchgen.models", models_stub) + wrappers_stub = types.ModuleType("batchgen.models.wrappers") + wrappers_stub.__path__ = [ + str(REPO_ROOT / "batchgen" / "models" / "wrappers") + ] + monkeypatch.setitem(sys.modules, "batchgen.models.wrappers", wrappers_stub) + kv_cache_stub = types.ModuleType("batchgen.kv_cache") + kv_cache_stub.__path__ = [str(REPO_ROOT / "batchgen" / "kv_cache")] + monkeypatch.setitem(sys.modules, "batchgen.kv_cache", kv_cache_stub) + + +def _prefix_cache_module(monkeypatch): + _install_torch_stub(monkeypatch) + return importlib.import_module("batchgen.models.wrappers.prefix_cache") + + +def _prefill_offload_module(monkeypatch): + _install_torch_stub(monkeypatch) + return importlib.import_module("batchgen.kv_cache.prefill_offload") + + +class _Wrapper: + prepack_cu_seqlens = _FakeCuSeqlens([0, 2, 5]) + prepack_max_seqlen = 3 + prepack_num_sequences = 2 + prepack_seq_lengths = [2, 3] + cur_batch = [10, 20] + prepack_prefix_reuse_mode = True + prepack_prefix_shared_tokens = [7, 11] + prepack_full_seq_lengths = [9, 14] + + +def test_prefix_cache_metadata_validates_prefix_lengths(monkeypatch): + mod = _prefix_cache_module(monkeypatch) + + metadata = mod.build_prefix_cache_forward_metadata_from_wrapper_cls(_Wrapper) + + assert metadata.global_sequence_ids == [10, 20] + assert metadata.prefix_shared_tokens == [7, 11] + + +def test_prefix_offloader_uses_destination_offsets(monkeypatch): + prefix_mod = _prefix_cache_module(monkeypatch) + offload_mod = _prefill_offload_module(monkeypatch) + metadata = prefix_mod.build_prefix_cache_forward_metadata_from_wrapper_cls( + _Wrapper + ) + worker_view = _FakeWorkerView() + tracked = [] + offloader = offload_mod.PrefillHostKVOffloader( + worker_view=worker_view, + layer_idx=3, + metadata=metadata, + track_task=lambda task, layer_idx: tracked.append((task, layer_idx)), + ) + + offloader.offload_gqa( + key=_FakeFlatTensor("k"), + value=_FakeFlatTensor("v"), + ) + + assert [kind for kind, _ in worker_view.calls] == ["offset", "offset"] + assert worker_view.calls[0][1]["raw_start_positions"] == [7] + assert worker_view.calls[0][1]["token_counts"] == [2] + assert worker_view.calls[1][1]["raw_start_positions"] == [11] + assert worker_view.calls[1][1]["token_counts"] == [3] + assert [layer_idx for _, layer_idx in tracked] == [3, 3] + + +def test_prefix_offloader_rejects_missing_offset_api(monkeypatch): + prefix_mod = _prefix_cache_module(monkeypatch) + offload_mod = _prefill_offload_module(monkeypatch) + metadata = prefix_mod.build_prefix_cache_forward_metadata_from_wrapper_cls( + _Wrapper + ) + offloader = offload_mod.PrefillHostKVOffloader( + worker_view=_NoOffsetWorkerView(), + layer_idx=0, + metadata=metadata, + ) + + with pytest.raises(RuntimeError, match="range_to_host"): + offloader.offload_mla(key=_FakeFlatTensor("kv", dim=2)) diff --git a/tests/unit/test_prefix_commit_helpers.py b/tests/unit/test_prefix_commit_helpers.py new file mode 100644 index 000000000..2d765cc6b --- /dev/null +++ b/tests/unit/test_prefix_commit_helpers.py @@ -0,0 +1,716 @@ +from __future__ import annotations + +import torch + +from batchgen.prefix_reuse.commit import ( + aligned_prefix_tokens, + build_committable_prefix_token_ids, + build_prefix_commit_request, + collect_required_group_pages_for_commit, +) +from batchgen.prefix_reuse.config import ( + PrefixCacheRuntimeConfig, + PrefixKVGroupSemantic, + PrefixKVGroupSpec, +) +from batchgen.prefix_reuse.eviction import ( + commit_prefix_pages_with_capacity_retry, + evict_prefix_pages_for_host_allocation, + release_evicted_prefix_pages, +) +from batchgen.prefix_reuse.worker_commit import ( + build_sequence_prefix_commit_request, + retain_newly_committed_prefix_pages, + sequence_token_ids_for_prefix_commit, +) + + +class _HostPageHandle: + def __init__(self): + self.page_id = 0 + + +class _GroupCommitPages: + def __init__(self): + self.group_id = 0 + self.pages = [] + + +class _GroupPageRequirement: + def __init__(self): + self.group_id = 0 + self.min_pages = 0 + + +class _Core: + HostPageHandle = _HostPageHandle + GroupCommitPages = _GroupCommitPages + GroupPageRequirement = _GroupPageRequirement + + +class _Coordinator: + def __init__( + self, + *, + fail_once_with: RuntimeError | None = None, + eviction_result=None, + ): + self.calls = [] + self.evict_calls = [] + self.fail_once_with = fail_once_with + self.eviction_result = eviction_result + + def commit_prefix_pages( + self, namespace_digest, token_ids, commit_tokens, group_pages + ): + self.calls.append( + (namespace_digest, token_ids, commit_tokens, group_pages) + ) + if self.fail_once_with is not None: + exc = self.fail_once_with + self.fail_once_with = None + raise exc + return "committed" + + def evict_until_free( + self, + min_free_nodes, + min_free_group_entries, + min_free_page_handles, + max_scan_nodes, + ): + self.evict_calls.append( + ( + min_free_nodes, + min_free_group_entries, + min_free_page_handles, + max_scan_nodes, + ) + ) + return self.eviction_result + + def evict_until_releasable_pages(self, requirements, max_scan_nodes): + self.evict_calls.append( + ( + [ + (int(requirement.group_id), int(requirement.min_pages)) + for requirement in requirements + ], + max_scan_nodes, + ) + ) + return self.eviction_result + + +class _WorkerView: + def __init__(self, pages): + self.pages = list(pages) + self.calls = [] + self.retained = [] + self.released = [] + + def build_page_table(self, sequence_ids): + self.calls.append(list(sequence_ids)) + return [list(self.pages) for _ in sequence_ids] + + def retain_sequence_prefix_pages(self, sequence_id, num_pages): + self.retained.append((int(sequence_id), int(num_pages))) + return self.pages[: int(num_pages)] + + def retain_sequence_page_range(self, sequence_id, start_page, num_pages): + self.retained.append( + (int(sequence_id), int(start_page), int(num_pages)) + ) + start = int(start_page) + end = start + int(num_pages) + return self.pages[start:end] + + def retain_sequence_pages(self, sequence_id, page_ids): + self.retained.append((int(sequence_id), list(page_ids))) + return list(page_ids) + + def release_resident_pages(self, page_ids): + self.released.append(list(page_ids)) + + +class _FastCoordinator(_Coordinator): + def commit_prefix_page_ids( + self, namespace_digest, token_ids, commit_tokens, group_page_ids + ): + self.calls.append( + (namespace_digest, token_ids, commit_tokens, group_page_ids) + ) + return "committed-fast" + + +class _EvictionResult: + def __init__(self, evicted_group_pages): + self.evicted_nodes = len(evicted_group_pages) + self.protected_nodes = 0 + self.evicted_group_pages = evicted_group_pages + + +def _page(page_id: int) -> _HostPageHandle: + handle = _HostPageHandle() + handle.page_id = int(page_id) + return handle + + +def _group_pages(group_id: int, pages) -> _GroupCommitPages: + group = _GroupCommitPages() + group.group_id = int(group_id) + group.pages = list(pages) + return group + + +class _Seq: + def __init__( + self, + *, + global_idx=7, + prompt=None, + decoded=None, + decoded_length=0, + reentry_decoded_baseline=0, + prefix_shared_tokens=0, + prefix_committed_tokens=0, + ): + prompt = [1, 2, 3, 4] if prompt is None else list(prompt) + decoded = [] if decoded is None else list(decoded) + self.global_idx = global_idx + self.prompt_length = len(prompt) + self.input_ids = torch.tensor([prompt], dtype=torch.long) + decoded_capacity = max(len(decoded), decoded_length, 1) + self.decoded_tokens = torch.zeros( + (1, decoded_capacity), dtype=torch.long + ) + if decoded: + self.decoded_tokens[0, : len(decoded)] = torch.tensor( + decoded, dtype=torch.long + ) + self.decoded_length = decoded_length + self.reentry_decoded_baseline = reentry_decoded_baseline + self.prefix_shared_tokens = prefix_shared_tokens + self.prefix_committed_tokens = prefix_committed_tokens + self.prefix_prompt_token_ids = None + self.prefix_prompt_cache_data_ptr = 0 + self.prefix_prompt_cache_length = 0 + self.prefix_prompt_cache_version = -1 + + +def _runtime_config() -> PrefixCacheRuntimeConfig: + return PrefixCacheRuntimeConfig( + shm_name="test", + namespace_digest=(1, 2, 3, 4), + group_specs=( + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=True, + raw_page_tokens=4, + ), + ), + hash_block_tokens=4, + publish_boundary_tokens=4, + max_nodes=16, + max_group_entries=16, + max_page_handles=32, + max_attachments=16, + ) + + +def test_aligned_prefix_tokens_floor_to_publish_boundary(): + assert aligned_prefix_tokens(0, 64) == 0 + assert aligned_prefix_tokens(63, 64) == 0 + assert aligned_prefix_tokens(64, 64) == 64 + assert aligned_prefix_tokens(191, 64) == 128 + + +def test_build_committable_prefix_token_ids_appends_only_new_decode_tokens(): + token_ids = build_committable_prefix_token_ids( + prompt_token_ids=[1, 2, 3, 4], + decoded_token_ids=[10, 11, 12], + decoded_start=2, + max_tokens=5, + ) + + assert token_ids == [1, 2, 3, 4, 12] + + +def test_build_committable_prefix_token_ids_clamps_negative_inputs(): + token_ids = build_committable_prefix_token_ids( + prompt_token_ids=[1, 2], + decoded_token_ids=[3, 4], + decoded_start=-8, + max_tokens=-1, + ) + + assert token_ids == [] + + +def test_sequence_token_ids_for_prefix_commit_skips_reentry_baseline(): + seq = _Seq( + prompt=[1, 2, 3, 10], + decoded=[10, 11, 12], + decoded_length=3, + reentry_decoded_baseline=1, + ) + + token_ids = sequence_token_ids_for_prefix_commit( + seq, + include_new_decode_tokens=True, + max_tokens=6, + ) + + assert token_ids == [1, 2, 3, 10, 11, 12] + + +def test_build_sequence_prefix_commit_request_collects_logical_pages(): + seq = _Seq( + global_idx=42, + prompt=[1, 2, 3, 4], + decoded=[5, 6, 7, 8], + decoded_length=4, + ) + + request_pair = build_sequence_prefix_commit_request( + core_engine_module=_Core, + runtime_config=_runtime_config(), + worker_views_by_group={0: _WorkerView([100, 101])}, + seq=seq, + include_new_decode_tokens=True, + ) + + assert request_pair is not None + request, commit_tokens = request_pair + assert commit_tokens == 8 + assert request.token_ids == [1, 2, 3, 4, 5, 6, 7, 8] + assert request.page_ids_by_group == {0: [100, 101]} + assert [page.page_id for page in request.group_pages[0].pages] == [ + 100, + 101, + ] + + +def test_build_sequence_prefix_commit_request_skips_already_shared_prefix(): + seq = _Seq( + prompt=[1, 2, 3, 4], + decoded=[5], + decoded_length=1, + prefix_shared_tokens=4, + ) + + request_pair = build_sequence_prefix_commit_request( + core_engine_module=_Core, + runtime_config=_runtime_config(), + worker_views_by_group={0: _WorkerView([100])}, + seq=seq, + include_new_decode_tokens=False, + ) + + assert request_pair is None + + +def test_build_sequence_prefix_commit_request_skips_already_committed_prefix(): + seq = _Seq( + prompt=[1, 2, 3, 4], + decoded=[5], + decoded_length=1, + prefix_committed_tokens=4, + ) + + request_pair = build_sequence_prefix_commit_request( + core_engine_module=_Core, + runtime_config=_runtime_config(), + worker_views_by_group={0: _WorkerView([100])}, + seq=seq, + include_new_decode_tokens=False, + ) + + assert request_pair is None + + +def test_build_prefix_commit_request_skips_unaligned_short_prefix(): + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=[10, 11, 12], + publish_boundary_tokens=4, + pages_by_group={0: [7]}, + raw_page_tokens_by_group={0: 4}, + ) + + assert request is None + + +def test_build_prefix_commit_request_uses_existing_group_pages(): + existing = _HostPageHandle() + existing.page_id = 9 + + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=[10, 11, 12, 13, 14], + publish_boundary_tokens=4, + pages_by_group={1: [existing], 0: [5, 6]}, + raw_page_tokens_by_group={0: 4, 1: 4}, + ) + + assert request is not None + assert request.namespace_digest == (1, 2, 3, 4) + assert request.token_ids == [10, 11, 12, 13, 14] + assert request.commit_tokens == 4 + assert [group.group_id for group in request.group_pages] == [0, 1] + assert [page.page_id for page in request.group_pages[0].pages] == [5, 6] + assert request.group_pages[1].pages == [existing] + + +def test_prefix_commit_request_invokes_coordinator(): + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=[10, 11, 12, 13], + publish_boundary_tokens=4, + pages_by_group={0: [5]}, + raw_page_tokens_by_group={0: 4}, + ) + coordinator = _Coordinator() + + result = request.commit(coordinator) + + assert result == "committed" + assert len(coordinator.calls) == 1 + namespace_digest, token_ids, commit_tokens, group_pages = ( + coordinator.calls[0] + ) + assert namespace_digest == [1, 2, 3, 4] + assert token_ids == [10, 11, 12, 13] + assert commit_tokens == 4 + assert [page.page_id for page in group_pages[0].pages] == [5] + + +def test_prefix_commit_request_uses_page_id_fast_path_when_available(): + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=[10, 11, 12, 13], + publish_boundary_tokens=4, + pages_by_group={0: [5]}, + raw_page_tokens_by_group={0: 4}, + ) + coordinator = _FastCoordinator() + + result = request.commit(coordinator) + + assert result == "committed-fast" + assert coordinator.calls == [ + ([1, 2, 3, 4], [10, 11, 12, 13], 4, [(0, [5])]) + ] + + +def test_sequence_token_ids_for_prefix_commit_reuses_prompt_cache(): + seq = _Seq(prompt=[1, 2, 3, 4], decoded=[5], decoded_length=1) + + first = sequence_token_ids_for_prefix_commit( + seq, + include_new_decode_tokens=True, + max_tokens=5, + ) + cached_prompt_ids = seq.prefix_prompt_token_ids + second = sequence_token_ids_for_prefix_commit( + seq, + include_new_decode_tokens=True, + max_tokens=5, + ) + + assert first == [1, 2, 3, 4, 5] + assert second == [1, 2, 3, 4, 5] + assert seq.prefix_prompt_token_ids is cached_prompt_ids + + +def test_sequence_token_ids_for_prefix_commit_invalidates_mutated_prompt_cache(): + seq = _Seq(prompt=[1, 2, 3, 4], decoded=[5], decoded_length=1) + + first = sequence_token_ids_for_prefix_commit( + seq, + include_new_decode_tokens=True, + max_tokens=5, + ) + seq.input_ids[0, 0] = 99 + second = sequence_token_ids_for_prefix_commit( + seq, + include_new_decode_tokens=True, + max_tokens=5, + ) + + assert first == [1, 2, 3, 4, 5] + assert second == [99, 2, 3, 4, 5] + + +def test_prefix_commit_request_capacity_requirements_use_raw_page_rates(): + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=list(range(16)), + publish_boundary_tokens=8, + pages_by_group={0: [0, 1, 2, 3], 1: [10, 11]}, + raw_page_tokens_by_group={0: 4, 1: 8}, + ) + + assert request is not None + assert request.capacity_requirements() == (2, 4, 6) + + +def test_retain_newly_committed_prefix_pages_reuses_collected_page_ids(): + worker_view = _WorkerView([100, 101, 102, 103]) + + retained = retain_newly_committed_prefix_pages( + runtime_config=_runtime_config(), + worker_views_by_group={0: worker_view}, + sequence_id=42, + previous_committed_tokens=4, + commit_tokens=12, + page_ids_by_group={0: [100, 101, 102]}, + ) + + assert retained == 12 + assert worker_view.calls == [] + assert worker_view.retained == [(42, [101, 102])] + + +def test_prefix_commit_request_capacity_requirements_cover_c128_groups(): + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=list(range(512)), + publish_boundary_tokens=256, + pages_by_group={ + 0: list(range(8)), + 1: [100, 101], + 2: [200, 201], + }, + raw_page_tokens_by_group={0: 64, 1: 256, 2: 256}, + ) + + assert request is not None + assert request.capacity_requirements() == (2, 6, 12) + + +def test_release_evicted_prefix_pages_requires_matching_worker_view(): + evicted = _EvictionResult([_group_pages(9, [_page(100)])]) + + try: + release_evicted_prefix_pages( + eviction_result=evicted, + worker_views_by_group={}, + ) + except RuntimeError as exc: + assert "evicted prefix group 9" in str(exc) + else: # pragma: no cover - failure path assertion + raise AssertionError("missing evicted group worker view should fail") + + +def test_commit_prefix_pages_retries_after_capacity_eviction(): + evicted = _EvictionResult( + [ + _group_pages(0, [_page(100), _page(100), _page(101)]), + _group_pages(1, [_page(200)]), + ] + ) + coordinator = _Coordinator( + fail_once_with=RuntimeError("Host prefix cache node table is full"), + eviction_result=evicted, + ) + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=list(range(16)), + publish_boundary_tokens=8, + pages_by_group={0: [0, 1, 2, 3], 1: [10, 11]}, + raw_page_tokens_by_group={0: 4, 1: 8}, + ) + primary = _WorkerView([]) + compressed = _WorkerView([]) + + result = commit_prefix_pages_with_capacity_retry( + request=request, + coordinator=coordinator, + worker_views_by_group={0: primary, 1: compressed}, + max_scan_nodes=7, + ) + + assert result.commit_result == "committed" + assert result.eviction_result is evicted + assert result.released_pages_by_group == {0: 2, 1: 1} + assert coordinator.evict_calls == [(2, 4, 6, 7)] + assert len(coordinator.calls) == 2 + assert primary.released == [[100, 101]] + assert compressed.released == [[200]] + + +def test_evict_prefix_pages_for_host_allocation_uses_page_requirements(): + evicted = _EvictionResult( + [ + _group_pages(0, [_page(100), _page(101)]), + _group_pages(1, [_page(200), _page(201), _page(201)]), + ] + ) + coordinator = _Coordinator(eviction_result=evicted) + primary = _WorkerView([]) + compressed = _WorkerView([]) + + result = evict_prefix_pages_for_host_allocation( + core_engine_module=_Core, + coordinator=coordinator, + worker_views_by_group={0: primary, 1: compressed}, + page_deficit_by_group={0: 2, 1: 1, 2: 0}, + max_scan_nodes=9, + ) + + assert result.eviction_result is evicted + assert result.released_pages_by_group == {0: 2, 1: 2} + assert coordinator.evict_calls == [([(0, 2), (1, 1)], 9)] + assert primary.released == [[100, 101]] + assert compressed.released == [[200, 201]] + + +def test_evict_prefix_pages_for_host_allocation_raises_when_short(): + evicted = _EvictionResult([_group_pages(0, [_page(100)])]) + coordinator = _Coordinator(eviction_result=evicted) + + try: + evict_prefix_pages_for_host_allocation( + core_engine_module=_Core, + coordinator=coordinator, + worker_views_by_group={0: _WorkerView([])}, + page_deficit_by_group={0: 2}, + ) + except RuntimeError as exc: + assert "could not release enough Host KV pages" in str(exc) + assert "missing={0: 1}" in str(exc) + else: # pragma: no cover - failure path assertion + raise AssertionError("short eviction result should fail") + + +def test_commit_prefix_pages_does_not_retry_non_capacity_errors(): + coordinator = _Coordinator(fail_once_with=RuntimeError("other failure")) + request = build_prefix_commit_request( + core_engine_module=_Core, + namespace_digest=(1, 2, 3, 4), + token_ids=list(range(8)), + publish_boundary_tokens=4, + pages_by_group={0: [0, 1]}, + raw_page_tokens_by_group={0: 4}, + ) + + try: + commit_prefix_pages_with_capacity_retry( + request=request, + coordinator=coordinator, + worker_views_by_group={0: _WorkerView([])}, + ) + except RuntimeError as exc: + assert str(exc) == "other failure" + else: # pragma: no cover - failure path assertion + raise AssertionError("non-capacity RuntimeError should be re-raised") + assert coordinator.evict_calls == [] + + +def test_retain_newly_committed_prefix_pages_uses_group_raw_page_rates(): + config = PrefixCacheRuntimeConfig( + shm_name="test", + namespace_digest=(1, 2, 3, 4), + group_specs=( + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=True, + raw_page_tokens=4, + ), + PrefixKVGroupSpec( + group_id=1, + semantic=PrefixKVGroupSemantic.COMPRESSED_RATIO_KV, + required_for_reuse=True, + raw_page_tokens=8, + compression_ratio=2, + ), + PrefixKVGroupSpec( + group_id=2, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=False, + raw_page_tokens=4, + ), + ), + hash_block_tokens=4, + publish_boundary_tokens=8, + max_nodes=16, + max_group_entries=16, + max_page_handles=32, + max_attachments=16, + ) + primary = _WorkerView([0, 1, 2, 3]) + compressed = _WorkerView([10, 11]) + + committed = retain_newly_committed_prefix_pages( + runtime_config=config, + worker_views_by_group={0: primary, 1: compressed}, + sequence_id=123, + previous_committed_tokens=8, + commit_tokens=16, + ) + + assert committed == 16 + assert primary.retained == [(123, [2, 3])] + assert compressed.retained == [(123, [11])] + + +def test_retain_newly_committed_prefix_pages_skips_already_committed_tokens(): + config = _runtime_config() + primary = _WorkerView([0, 1]) + + committed = retain_newly_committed_prefix_pages( + runtime_config=config, + worker_views_by_group={0: primary}, + sequence_id=123, + previous_committed_tokens=8, + commit_tokens=8, + ) + + assert committed == 8 + assert primary.retained == [] + + +def test_collect_required_group_pages_for_commit_reads_worker_page_tables(): + specs = [ + PrefixKVGroupSpec( + group_id=0, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=True, + raw_page_tokens=4, + ), + PrefixKVGroupSpec( + group_id=1, + semantic=PrefixKVGroupSemantic.MLA_COMPRESSED_KV, + required_for_reuse=True, + raw_page_tokens=8, + ), + PrefixKVGroupSpec( + group_id=2, + semantic=PrefixKVGroupSemantic.FULL_KV, + required_for_reuse=False, + raw_page_tokens=4, + ), + ] + primary = _WorkerView([10, 11, 12, 13]) + mla = _WorkerView([20, 21]) + + pages = collect_required_group_pages_for_commit( + worker_views_by_group={0: primary, 1: mla}, + sequence_id=100, + commit_tokens=8, + group_specs=specs, + ) + + assert pages == {0: [10, 11], 1: [20]} + assert primary.calls == [[100]] + assert mla.calls == [[100]] diff --git a/tests/unit/test_prefix_materialization.py b/tests/unit/test_prefix_materialization.py new file mode 100644 index 000000000..ee7ca59cd --- /dev/null +++ b/tests/unit/test_prefix_materialization.py @@ -0,0 +1,619 @@ +from types import SimpleNamespace + +import pytest +import torch + +from batchgen.prefix_reuse.materialization import ( + PrefixMaterializationBundle, + PrefixMaterializationSequence, + SingleGroupPrefixMaterialization, + get_prefix_materialization_for_group, + materialize_single_group_lookup_results, + materialize_single_group_prefix_pages, +) + + +class _FakeTask: + def __init__(self): + self.wait_count = 0 + self.waited_layers = [] + + def wait(self): + self.wait_count += 1 + + def wait_for_layer(self, layer_idx): + self.waited_layers.append(int(layer_idx)) + + +class _FakeHostWorkerView: + def __init__(self): + self.task = _FakeTask() + self.calls = [] + self.layer_calls = [] + self.layer_tasks = [] + + def async_load_prefix_pages_to_device(self, **kwargs): + self.calls.append(kwargs) + return self.task + + def async_load_prefix_layers_to_device(self, **kwargs): + task = _FakeTask() + self.layer_calls.append(kwargs) + self.layer_tasks.append(task) + return task + + +class _FailingHostWorkerView(_FakeHostWorkerView): + def async_load_prefix_pages_to_device(self, **kwargs): + super().async_load_prefix_pages_to_device(**kwargs) + raise RuntimeError("load failed") + + +class _FakePrefixCoordinator: + def __init__(self): + self.begin_calls = [] + self.end_calls = [] + + def begin_attachment_load(self, attachment_handle): + self.begin_calls.append(int(attachment_handle)) + + def end_attachment_load(self, attachment_handle): + self.end_calls.append(int(attachment_handle)) + + +class _FakeGpuManager: + def __init__(self): + self.config = SimpleNamespace( + page_size_tokens=4, + num_k_heads=1, + k_head_dim=1, + num_v_heads=1, + v_head_dim=1, + kv_dtype=torch.bfloat16, + ) + self.allocations = [] + self.rebuilt = [] + self.prepared = [] + self.destroy_calls = [] + self.k_ptrs = torch.tensor( + [ + [[1000, 2000, 3000], [4000, 5000, 6000]], + [[7000, 8000, 9000], [10000, 11000, 12000]], + ], + dtype=torch.int64, + ) + self.v_ptrs = self.k_ptrs + 100000 + self.append_plan = SimpleNamespace( + cache_seqlens=torch.tensor([7, 3], dtype=torch.int32), + slot_indices=torch.tensor([0, 1], dtype=torch.int32), + slot_values=(0, 1), + ) + + def allocate_pages_for_sequences(self, sequence_ids, num_tokens): + self.allocations.append((list(sequence_ids), list(num_tokens))) + + def rebuild_page_table(self, sequence_ids): + self.rebuilt.append(list(sequence_ids)) + + def get_padded_3d_page_pointers(self): + return self.k_ptrs, self.v_ptrs + + def prepare_prefill_suffix_append( + self, + *, + sequence_ids, + prefix_lens, + suffix_lens, + rebuild_page_table, + ): + self.prepared.append( + ( + list(sequence_ids), + list(prefix_lens), + list(suffix_lens), + rebuild_page_table, + ) + ) + return self.append_plan + + def destroy(self, *, empty_cuda_cache=False): + self.destroy_calls.append(bool(empty_cuda_cache)) + + def resolve_physical_layer(self, layer_idx): + return int(layer_idx) % int(self.k_ptrs.shape[0]) + + +class _FailingAppendPlanGpuManager(_FakeGpuManager): + def prepare_prefill_suffix_append(self, **kwargs): + super().prepare_prefill_suffix_append(**kwargs) + raise RuntimeError("append plan failed") + + +def test_prefix_materialization_bundle_returns_group_materialization(): + primary = SingleGroupPrefixMaterialization( + manager=object(), + append_plan=object(), + ) + aux = SingleGroupPrefixMaterialization( + manager=object(), + append_plan=object(), + ) + bundle = PrefixMaterializationBundle(by_group_id={0: primary, 1: aux}) + + assert bundle.get(0) is primary + assert bundle.require(1, consumer="test") is aux + assert ( + get_prefix_materialization_for_group( + bundle, group_id=0, consumer="test" + ) + is primary + ) + + +def test_prefix_materialization_bundle_rejects_missing_group(): + bundle = PrefixMaterializationBundle(by_group_id={}) + + with pytest.raises(RuntimeError, match="group 2"): + bundle.require(2, consumer="test") + + +def test_get_prefix_materialization_rejects_legacy_single_group(): + materialization = SingleGroupPrefixMaterialization( + manager=object(), + append_plan=object(), + ) + + with pytest.raises(RuntimeError, match="PrefixMaterializationBundle"): + get_prefix_materialization_for_group( + materialization, + group_id=0, + consumer="test", + ) + + +def test_materialize_single_group_prefix_pages_starts_page_id_load(): + gpu_manager = _FakeGpuManager() + host_view = _FakeHostWorkerView() + + materialization = materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=host_view, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=5, + suffix_tokens=2, + host_pages=[11, 12], + ), + PrefixMaterializationSequence( + sequence_id=102, + prefix_tokens=0, + suffix_tokens=3, + host_pages=[], + ), + ], + ) + + assert materialization.manager is gpu_manager + assert materialization.append_plan is gpu_manager.append_plan + assert gpu_manager.allocations == [([101, 102], [7, 3])] + assert gpu_manager.rebuilt == [[101, 102]] + assert gpu_manager.prepared == [([101, 102], [5, 0], [2, 3], False)] + assert len(host_view.calls) == 1 + call = host_view.calls[0] + assert call["host_page_ids"].tolist() == [[11, 12], [0, 0]] + assert call["active_page_counts"].tolist() == [2, 0] + assert call["k_device_ptrs"] is gpu_manager.k_ptrs + assert call["v_device_ptrs"] is gpu_manager.v_ptrs + + materialization.wait_for_layer(0) + materialization.wait_for_layer(1) + assert host_view.task.waited_layers == [0, 1] + assert host_view.task.wait_count == 0 + + +def test_materialize_prefix_pages_uses_raw_page_tokens_for_compressed_groups(): + gpu_manager = _FakeGpuManager() + gpu_manager.config.page_size_tokens = 2 + host_view = _FakeHostWorkerView() + + materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=host_view, + raw_page_tokens=256, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=256, + suffix_tokens=8, + host_pages=[11], + ), + PrefixMaterializationSequence( + sequence_id=102, + prefix_tokens=512, + suffix_tokens=8, + host_pages=[21, 22], + ), + ], + ) + + assert len(host_view.calls) == 1 + call = host_view.calls[0] + assert call["host_page_ids"].tolist() == [[11, 0], [21, 22]] + assert call["active_page_counts"].tolist() == [1, 2] + + +def test_materialize_prefix_pages_offsets_host_pages_inside_larger_gpu_pages(): + gpu_manager = _FakeGpuManager() + gpu_manager.config.page_size_tokens = 4 + host_view = _FakeHostWorkerView() + + materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=host_view, + raw_page_tokens=2, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=6, + suffix_tokens=2, + host_pages=[11, 12, 13], + ), + PrefixMaterializationSequence( + sequence_id=102, + prefix_tokens=2, + suffix_tokens=2, + host_pages=[21], + ), + ], + ) + + call = host_view.calls[0] + assert call["host_page_ids"].tolist() == [[11, 12, 13], [21, 0, 0]] + assert call["active_page_counts"].tolist() == [3, 1] + # raw_page_tokens=2, BF16, 1 head, dim=1 -> 4 bytes per Host page. + assert call["k_device_ptrs"].tolist() == [ + [[1000, 1004, 2000], [4000, 4004, 5000]], + [[7000, 7004, 8000], [10000, 10004, 11000]], + ] + assert call["v_device_ptrs"].tolist() == [ + [[101000, 101004, 102000], [104000, 104004, 105000]], + [[107000, 107004, 108000], [110000, 110004, 111000]], + ] + + +def test_rolling_materialization_prefetches_two_layers_and_advances(): + gpu_manager = _FakeGpuManager() + host_view = _FakeHostWorkerView() + coordinator = _FakePrefixCoordinator() + + materialization = materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=host_view, + prefix_cache_coordinator=coordinator, + rolling_logical_layer_count=4, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=4, + suffix_tokens=2, + host_pages=[11], + attachment_handle=91, + ), + ], + ) + + assert coordinator.begin_calls == [91] + assert coordinator.end_calls == [] + assert len(host_view.layer_calls) == 2 + assert host_view.layer_calls[0]["logical_layer_ids"].tolist() == [0] + assert host_view.layer_calls[1]["logical_layer_ids"].tolist() == [1] + assert host_view.layer_calls[0]["k_device_ptrs"].tolist() == [ + gpu_manager.k_ptrs[0].tolist() + ] + assert host_view.layer_calls[1]["k_device_ptrs"].tolist() == [ + gpu_manager.k_ptrs[1].tolist() + ] + + materialization.wait_for_layer(0) + assert host_view.layer_tasks[0].waited_layers == [0] + + materialization.backend_state["flashinfer"] = object() + materialization.finish_layer(0) + assert len(host_view.layer_calls) == 3 + assert host_view.layer_calls[2]["logical_layer_ids"].tolist() == [2] + assert host_view.layer_calls[2]["k_device_ptrs"].tolist() == [ + gpu_manager.k_ptrs[0].tolist() + ] + + materialization.close(empty_cuda_cache=True) + assert coordinator.end_calls == [91] + assert gpu_manager.destroy_calls == [True] + assert materialization.backend_state == {} + + +def test_materialize_single_group_prefix_pages_skips_load_for_all_miss(): + gpu_manager = _FakeGpuManager() + host_view = _FakeHostWorkerView() + + materialization = materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=host_view, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=0, + suffix_tokens=3, + host_pages=[], + ), + ], + ) + + assert host_view.calls == [] + materialization.wait_for_layer(0) + assert gpu_manager.allocations == [([101], [3])] + + +def test_materialize_single_group_prefix_pages_guards_attachment_load(): + gpu_manager = _FakeGpuManager() + host_view = _FakeHostWorkerView() + coordinator = _FakePrefixCoordinator() + + materialization = materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=host_view, + prefix_cache_coordinator=coordinator, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=4, + suffix_tokens=1, + host_pages=[11], + attachment_handle=91, + ), + PrefixMaterializationSequence( + sequence_id=102, + prefix_tokens=4, + suffix_tokens=1, + host_pages=[12], + attachment_handle=91, + ), + ], + ) + + assert coordinator.begin_calls == [91] + assert coordinator.end_calls == [] + materialization.wait_for_layer(0) + materialization.wait_for_layer(1) + assert host_view.task.waited_layers == [0, 1] + assert host_view.task.wait_count == 0 + assert coordinator.end_calls == [] + materialization.wait() + assert host_view.task.wait_count == 1 + assert coordinator.end_calls == [91] + + +def test_bundle_full_wait_waits_all_groups(): + primary_manager = _FakeGpuManager() + aux_manager = _FakeGpuManager() + primary = SingleGroupPrefixMaterialization( + manager=primary_manager, + append_plan=object(), + load_task=_FakeTask(), + ) + aux = SingleGroupPrefixMaterialization( + manager=aux_manager, + append_plan=object(), + load_task=_FakeTask(), + ) + bundle = PrefixMaterializationBundle(by_group_id={0: primary, 1: aux}) + + bundle.wait_for_layer(3) + assert primary.load_task.waited_layers == [3] + assert aux.load_task.waited_layers == [3] + + bundle.wait() + assert primary.load_task.wait_count == 1 + assert aux.load_task.wait_count == 1 + + bundle.close(empty_cuda_cache=True) + assert primary_manager.destroy_calls == [True] + assert aux_manager.destroy_calls == [True] + assert primary.manager is None + assert primary.append_plan is None + assert primary.load_task is None + + +def test_materialize_single_group_prefix_pages_unwinds_attachment_on_load_error(): + gpu_manager = _FakeGpuManager() + coordinator = _FakePrefixCoordinator() + + with pytest.raises(RuntimeError, match="load failed"): + materialize_single_group_prefix_pages( + gpu_manager=gpu_manager, + host_worker_view=_FailingHostWorkerView(), + prefix_cache_coordinator=coordinator, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=4, + suffix_tokens=1, + host_pages=[11], + attachment_handle=91, + ), + ], + ) + + assert coordinator.begin_calls == [91] + assert coordinator.end_calls == [91] + + +def test_materialize_single_group_prefix_pages_does_not_load_before_append_plan(): + host_view = _FakeHostWorkerView() + coordinator = _FakePrefixCoordinator() + + with pytest.raises(RuntimeError, match="append plan failed"): + materialize_single_group_prefix_pages( + gpu_manager=_FailingAppendPlanGpuManager(), + host_worker_view=host_view, + prefix_cache_coordinator=coordinator, + sequences=[ + PrefixMaterializationSequence( + sequence_id=101, + prefix_tokens=4, + suffix_tokens=1, + host_pages=[11], + attachment_handle=91, + ), + ], + ) + + assert host_view.calls == [] + assert coordinator.begin_calls == [] + assert coordinator.end_calls == [] + + +def test_materialize_single_group_lookup_results_builds_sequences(): + gpu_manager = _FakeGpuManager() + host_view = _FakeHostWorkerView() + coordinator = _FakePrefixCoordinator() + lookup_result = SimpleNamespace( + attachment_handle=91, + common_cached_tokens=5, + materialization_spans=[ + SimpleNamespace( + group_id=7, + raw_end_token=5, + pages=[ + SimpleNamespace(page_id=11), + SimpleNamespace(page_id=12), + ], + ) + ], + ) + + materialization = materialize_single_group_lookup_results( + gpu_manager=gpu_manager, + host_worker_view=host_view, + prefix_cache_coordinator=coordinator, + lookup_results=[lookup_result], + sequence_ids=[101], + prompt_lengths=[7], + group_id=7, + prefix_shared_tokens=[5], + ) + + assert materialization.append_plan is gpu_manager.append_plan + assert gpu_manager.allocations == [([101], [7])] + assert gpu_manager.prepared == [([101], [5], [2], False)] + assert host_view.calls[0]["host_page_ids"].tolist() == [[11, 12]] + assert host_view.calls[0]["active_page_counts"].tolist() == [2] + assert coordinator.begin_calls == [91] + materialization.wait() + assert coordinator.end_calls == [91] + + +def test_materialize_single_group_lookup_results_clamps_full_hit_to_extend_one(): + gpu_manager = _FakeGpuManager() + host_view = _FakeHostWorkerView() + coordinator = _FakePrefixCoordinator() + lookup_result = SimpleNamespace( + attachment_handle=91, + common_cached_tokens=7, + materialization_spans=[ + SimpleNamespace( + group_id=7, + raw_end_token=7, + pages=[ + SimpleNamespace(page_id=11), + SimpleNamespace(page_id=12), + ], + ) + ], + ) + + materialize_single_group_lookup_results( + gpu_manager=gpu_manager, + host_worker_view=host_view, + prefix_cache_coordinator=coordinator, + lookup_results=[lookup_result], + sequence_ids=[101], + prompt_lengths=[7], + group_id=7, + ) + + assert gpu_manager.allocations == [([101], [7])] + assert gpu_manager.prepared == [([101], [6], [1], False)] + assert host_view.calls[0]["host_page_ids"].tolist() == [[11, 12]] + assert host_view.calls[0]["active_page_counts"].tolist() == [2] + assert coordinator.begin_calls == [91] + + +def test_materialize_single_group_lookup_results_skips_load_for_one_token_full_hit(): + gpu_manager = _FakeGpuManager() + host_view = _FakeHostWorkerView() + lookup_result = SimpleNamespace( + attachment_handle=91, + common_cached_tokens=1, + materialization_spans=[ + SimpleNamespace( + group_id=7, + raw_end_token=1, + pages=[SimpleNamespace(page_id=11)], + ) + ], + ) + + materialize_single_group_lookup_results( + gpu_manager=gpu_manager, + host_worker_view=host_view, + lookup_results=[lookup_result], + sequence_ids=[101], + prompt_lengths=[1], + group_id=7, + prefix_shared_tokens=[0], + ) + + assert gpu_manager.allocations == [([101], [1])] + assert gpu_manager.prepared == [([101], [0], [1], False)] + assert host_view.calls == [] + + +def test_materialize_single_group_lookup_results_rejects_mismatched_span(): + lookup_result = SimpleNamespace( + attachment_handle=91, + common_cached_tokens=5, + materialization_spans=[ + SimpleNamespace(group_id=7, raw_end_token=4, pages=[11]) + ], + ) + + with pytest.raises(ValueError, match="effective cached token boundary"): + materialize_single_group_lookup_results( + gpu_manager=_FakeGpuManager(), + host_worker_view=_FakeHostWorkerView(), + lookup_results=[lookup_result], + sequence_ids=[101], + prompt_lengths=[7], + group_id=7, + ) + + +def test_materialize_single_group_lookup_results_requires_attachment_for_hit(): + lookup_result = SimpleNamespace( + attachment_handle=0, + common_cached_tokens=5, + materialization_spans=[ + SimpleNamespace(group_id=7, raw_end_token=5, pages=[11]) + ], + ) + + with pytest.raises(ValueError, match="attachment_handle"): + materialize_single_group_lookup_results( + gpu_manager=_FakeGpuManager(), + host_worker_view=_FakeHostWorkerView(), + lookup_results=[lookup_result], + sequence_ids=[101], + prompt_lengths=[7], + group_id=7, + ) diff --git a/tests/unit/test_prefix_mla_absorb.py b/tests/unit/test_prefix_mla_absorb.py new file mode 100644 index 000000000..455be1be4 --- /dev/null +++ b/tests/unit/test_prefix_mla_absorb.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import torch + +from batchgen.attention.mla.prefix_absorb import ( + absorb_mla_attention_output, + build_absorbed_mla_query_states, + prefix_rotary_seq_len, + project_absorbed_mla_output, + project_absorbed_mla_output_w8a16, +) + + +def test_build_absorbed_mla_query_states_matches_manual_einsum(): + q_nope = torch.arange(12, dtype=torch.float32).view(2, 2, 3) + q_pe = torch.arange(8, dtype=torch.float32).view(2, 2, 2) + q_absorb = torch.arange(24, dtype=torch.float32).view(2, 3, 4) + + actual = build_absorbed_mla_query_states( + q_nope=q_nope, + q_pe=q_pe, + q_absorb=q_absorb, + dtype=torch.float32, + ) + + expected = torch.empty(1, 2, 2, 6) + expected[0, :, :, :4] = torch.einsum("thd,hdc->thc", q_nope, q_absorb) + expected[0, :, :, 4:] = q_pe + assert torch.equal(actual, expected.contiguous()) + + +def test_project_absorbed_mla_output_uses_common_absorb_layout(): + attn_out = torch.arange(16, dtype=torch.float32).view(1, 2, 2, 4) + out_absorb = torch.arange(24, dtype=torch.float32).view(2, 3, 4) + projection = torch.nn.Linear(6, 5, bias=False) + + absorbed = absorb_mla_attention_output( + attn_out=attn_out, + out_absorb=out_absorb, + v_head_dim=3, + ) + actual = project_absorbed_mla_output( + attn_out=attn_out, + out_absorb=out_absorb, + v_head_dim=3, + output_projection=projection, + ) + + expected_absorbed = torch.einsum( + "bqhc,hdc->bqhd", + attn_out, + out_absorb, + ).reshape(2, 6) + assert torch.equal(absorbed, expected_absorbed) + assert torch.equal(actual, projection(absorbed)) + + +def test_project_absorbed_mla_output_w8a16_delegates_to_gemm(): + attn_out = torch.arange(16, dtype=torch.float32).view(1, 2, 2, 4) + out_absorb = torch.arange(24, dtype=torch.float32).view(2, 3, 4) + weight = torch.randn(5, 6) + scale = torch.ones(5) + calls = {} + expected_result = torch.randn(2, 5) + + def fake_gemm(w, s, x): + calls["weight"] = w + calls["scale"] = s + calls["input"] = x + return expected_result + + actual = project_absorbed_mla_output_w8a16( + attn_out=attn_out, + out_absorb=out_absorb, + v_head_dim=3, + o_proj_weight=weight, + o_proj_scale=scale, + gemm=fake_gemm, + ) + + assert actual is expected_result + assert calls["weight"] is weight + assert calls["scale"] is scale + assert torch.equal( + calls["input"], + absorb_mla_attention_output( + attn_out=attn_out, + out_absorb=out_absorb, + v_head_dim=3, + ), + ) + + +def test_prefix_rotary_seq_len_covers_prefix_and_position_ids(): + position_ids = torch.tensor([3, 7, 8], dtype=torch.long) + + assert prefix_rotary_seq_len(5, position_ids) == 9 + assert prefix_rotary_seq_len(16, position_ids) == 16 diff --git a/tests/unit/test_prefix_mla_extend_path.py b/tests/unit/test_prefix_mla_extend_path.py new file mode 100644 index 000000000..721b2f36b --- /dev/null +++ b/tests/unit/test_prefix_mla_extend_path.py @@ -0,0 +1,204 @@ +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + +_FLASHINFER_STUB = types.ModuleType("flashinfer") +_FLASHINFER_STUB.BatchMLAPagedAttentionWrapper = object +sys.modules.setdefault("flashinfer", _FLASHINFER_STUB) + +from batchgen.attention.forward_metadata import ( # noqa: E402 + ForwardBatchMetadata, + PrefillAttentionMetadata, +) +from batchgen.attention.mla import flashinfer_extend # noqa: E402 +from batchgen.models.wrappers.prefix_mla_extend import ( # noqa: E402 + MlaExtendSpec, + run_projected_mla_prefix_attention_from_gpu_pages, +) + + +class _FakeMlaGpuManager: + def __init__(self, *, has_v_cache: bool = False): + self.config = SimpleNamespace(has_v_cache=has_v_cache) + self.append_calls = [] + self.blocked_k = torch.arange( + 4 * 8 * 1 * 6, + dtype=torch.float32, + ).reshape(4, 8, 1, 6) + self.block_table = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32) + + def append_layer_prefill_suffix_tokens( + self, + *, + k_tensor, + v_tensor, + append_plan, + layer_idx, + ): + self.append_calls.append( + { + "k_tensor": k_tensor, + "v_tensor": v_tensor, + "append_plan": append_plan, + "layer_idx": int(layer_idx), + } + ) + + def get_layer_kv_with_page_table(self, layer_idx): + return self.blocked_k, None, self.block_table + + +class _FakeMaterialization: + def __init__(self, *, has_v_cache: bool = False): + self.manager = _FakeMlaGpuManager(has_v_cache=has_v_cache) + self.backend_state = {} + self.append_plan = SimpleNamespace( + cache_seqlens=torch.tensor([9, 10], dtype=torch.int32), + slot_indices=torch.tensor([1, 0], dtype=torch.int32), + ) + self.waited_layers = [] + + def wait_for_layer(self, layer_idx): + self.waited_layers.append(int(layer_idx)) + + +def _metadata() -> ForwardBatchMetadata: + return ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[101, 102], + prefill=PrefillAttentionMetadata( + cu_seqlens_q=torch.tensor([0, 1, 3], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 9, 19], dtype=torch.int32), + max_seqlen_q=2, + max_seqlen_k=10, + q_seq_lens=[1, 2], + kv_seq_lens=[9, 10], + position_ids=torch.tensor([8, 8, 9], dtype=torch.int64), + append_seq_lens=[1, 2], + ), + ) + + +def test_projected_mla_prefix_attention_appends_suffix_and_runs_flashinfer( + monkeypatch, +): + materialization = _FakeMaterialization() + query_states = torch.zeros((1, 3, 2, 6), dtype=torch.float32) + offload_kv = torch.ones((3, 1, 6), dtype=torch.float32) + expected_output = torch.full((1, 3, 2, 4), 7.0, dtype=torch.float32) + call = {} + + def fake_flashinfer_extend(**kwargs): + call.update(kwargs) + return expected_output + + monkeypatch.setattr( + flashinfer_extend, + "run_flashinfer_mla_extend_prefill", + fake_flashinfer_extend, + ) + + output = run_projected_mla_prefix_attention_from_gpu_pages( + layer_idx=5, + query_states=query_states, + offload_kv=offload_kv, + metadata=_metadata(), + spec=MlaExtendSpec( + num_heads=2, + kv_lora_rank=4, + softmax_scale=0.25, + ), + materialization=materialization, + ) + + assert output is expected_output + assert materialization.waited_layers == [5] + assert len(materialization.manager.append_calls) == 1 + append_call = materialization.manager.append_calls[0] + assert append_call["k_tensor"] is offload_kv + assert append_call["v_tensor"] is None + assert append_call["append_plan"] is materialization.append_plan + assert append_call["layer_idx"] == 5 + assert call["query_states"].shape == query_states.shape + assert call["compressed_kv_cache"] is materialization.manager.blocked_k + assert call["page_table"] is materialization.manager.block_table + assert call["slot_indices"] is materialization.append_plan.slot_indices + assert call["cache_seqlens"] is materialization.append_plan.cache_seqlens + assert call["cu_seqlens_q"].tolist() == [0, 1, 3] + assert call["kv_lora_rank"] == 4 + assert call["num_heads"] == 2 + assert call["softmax_scale"] == 0.25 + assert call["plan_cache"] is materialization.backend_state + + +def test_flashinfer_mla_extend_prefill_reuses_materialization_plan( + monkeypatch, +): + flashinfer_extend._reset_flashinfer_mla_extend_prefill_cache_for_tests() + created_wrappers = [] + + class FakeWrapper: + def __init__(self, workspace, backend="auto"): + self.workspace = workspace + self.backend = backend + self.plan_calls = 0 + self.run_calls = 0 + created_wrappers.append(self) + + def plan(self, *args): + self.plan_calls += 1 + + def run(self, q_nope, q_pe, ckv_cache, kpe_cache): + self.run_calls += 1 + return torch.zeros_like(q_nope) + + monkeypatch.setattr( + flashinfer_extend, + "BatchMLAPagedAttentionWrapper", + FakeWrapper, + ) + + plan_cache = {} + kwargs = dict( + query_states=torch.zeros((1, 3, 2, 6), dtype=torch.float32), + compressed_kv_cache=torch.zeros((4, 8, 1, 6), dtype=torch.float32), + page_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + slot_indices=torch.tensor([1, 0], dtype=torch.int32), + cache_seqlens=torch.tensor([9, 10], dtype=torch.int32), + cu_seqlens_q=torch.tensor([0, 1, 3], dtype=torch.int32), + kv_lora_rank=4, + num_heads=2, + softmax_scale=0.25, + plan_cache=plan_cache, + ) + + flashinfer_extend.run_flashinfer_mla_extend_prefill(**kwargs) + flashinfer_extend.run_flashinfer_mla_extend_prefill(**kwargs) + + assert len(created_wrappers) == 1 + assert created_wrappers[0].plan_calls == 1 + assert created_wrappers[0].run_calls == 2 + + +def test_projected_mla_prefix_attention_rejects_v_cache_before_append(): + materialization = _FakeMaterialization(has_v_cache=True) + + with pytest.raises(RuntimeError, match="K-only compressed KV"): + run_projected_mla_prefix_attention_from_gpu_pages( + layer_idx=5, + query_states=torch.zeros((1, 1, 2, 6), dtype=torch.float32), + offload_kv=torch.ones((1, 1, 6), dtype=torch.float32), + metadata=_metadata(), + spec=MlaExtendSpec( + num_heads=2, + kv_lora_rank=4, + softmax_scale=0.25, + ), + materialization=materialization, + ) + + assert materialization.waited_layers == [] + assert materialization.manager.append_calls == [] diff --git a/tests/unit/test_prefix_mla_model_adapters.py b/tests/unit/test_prefix_mla_model_adapters.py new file mode 100644 index 000000000..e38ecf961 --- /dev/null +++ b/tests/unit/test_prefix_mla_model_adapters.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import importlib +import sys +import types +from types import SimpleNamespace + +import torch + +from batchgen.attention.forward_metadata import ( + ForwardBatchMetadata, + PrefillAttentionMetadata, +) + + +def _prefill_metadata() -> PrefillAttentionMetadata: + return PrefillAttentionMetadata( + cu_seqlens_q=torch.tensor([0, 2], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 5], dtype=torch.int32), + max_seqlen_q=2, + max_seqlen_k=5, + q_seq_lens=[2], + kv_seq_lens=[5], + position_ids=torch.tensor([3, 4], dtype=torch.int64), + ) + + +def _wrapper(): + module = SimpleNamespace( + kv_lora_rank=4, + qk_rope_head_dim=2, + num_heads=2, + softmax_scale=0.5, + ) + return SimpleNamespace(module=module) + + +def _prefix_mla_adapters(monkeypatch): + kv_cache_stub = types.ModuleType("batchgen.kv_cache") + kv_cache_stub.__path__ = [] + monkeypatch.setitem(sys.modules, "batchgen.kv_cache", kv_cache_stub) + prefill_offload_stub = types.ModuleType("batchgen.kv_cache.prefill_offload") + prefill_offload_stub.PrefillHostKVOffloader = object + monkeypatch.setitem( + sys.modules, + "batchgen.kv_cache.prefill_offload", + prefill_offload_stub, + ) + return importlib.import_module( + "batchgen.models.wrappers.prefix_mla_model_adapters" + ) + + +def test_mla_model_adapters_accept_explicit_prefill_metadata(monkeypatch): + adapters = _prefix_mla_adapters(monkeypatch) + prefill = _prefill_metadata() + metadata = ForwardBatchMetadata( + phase="prefill", + global_sequence_ids=[100], + prefill=prefill, + ) + wrapper = _wrapper() + + contexts = [ + adapters.build_deepseek_prefix_backend_context( + wrapper=wrapper, metadata=metadata + ), + adapters.build_glm5_prefix_backend_context( + wrapper=wrapper, metadata=metadata + ), + adapters.build_kimi_prefix_backend_context( + wrapper=wrapper, metadata=metadata + ), + ] + + for context in contexts: + assert context.prefix_reuse_mode is True + assert context.metadata.global_sequence_ids == [100] + assert context.metadata.prefix_shared_tokens == [3] + assert context.metadata.full_seq_lengths == [5] + assert ( + context.rotary_seq_len(prefill.position_ids, fallback_seq_len=2) + == 5 + ) diff --git a/tests/unit/test_prefix_prefill_lookup.py b/tests/unit/test_prefix_prefill_lookup.py new file mode 100644 index 000000000..8736c6596 --- /dev/null +++ b/tests/unit/test_prefix_prefill_lookup.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from batchgen.prefix_reuse.prefill import ( + build_prefix_cache_prefill_inputs, + estimate_prefix_cache_for_prefill, + lookup_prefix_cache_for_prefill, +) + + +class _Coordinator: + def __init__(self, cached_tokens: list[int], handles: list[int]): + self.cached_tokens = list(cached_tokens) + self.handles = list(handles) + self.lookup_calls = [] + self.estimate_calls = [] + + def lookup_and_attach(self, namespace_digest, token_ids): + index = len(self.lookup_calls) + self.lookup_calls.append((list(namespace_digest), list(token_ids))) + return SimpleNamespace( + common_cached_tokens=self.cached_tokens[index], + attachment_handle=self.handles[index], + ) + + def estimate_lookup(self, namespace_digest, token_ids): + index = len(self.estimate_calls) + self.estimate_calls.append((list(namespace_digest), list(token_ids))) + return SimpleNamespace( + common_cached_tokens=self.cached_tokens[index], + attachment_handle=0, + ) + + +def test_lookup_prefix_cache_for_prefill_preserves_request_order(): + coordinator = _Coordinator(cached_tokens=[4, 0, 8], handles=[11, 0, 12]) + + lookup = lookup_prefix_cache_for_prefill( + coordinator=coordinator, + namespace_digest=(1, 2, 3, 4), + prompt_token_ids=[ + [10, 11, 12, 13, 14], + [20, 21], + [30, 31, 32, 33, 34, 35, 36, 37], + ], + ) + + assert lookup.prefix_shared_tokens == (4, 0, 7) + assert lookup.has_hit is True + assert coordinator.lookup_calls == [ + ([1, 2, 3, 4], [10, 11, 12, 13, 14]), + ([1, 2, 3, 4], [20, 21]), + ([1, 2, 3, 4], [30, 31, 32, 33, 34, 35, 36, 37]), + ] + + +def test_estimate_prefix_cache_for_prefill_does_not_attach(): + coordinator = _Coordinator(cached_tokens=[4, 0], handles=[11, 0]) + + estimate = estimate_prefix_cache_for_prefill( + coordinator=coordinator, + namespace_digest=(1, 2, 3, 4), + prompt_token_ids=[ + [10, 11, 12, 13, 14], + [20, 21], + ], + ) + + assert estimate.prefix_shared_tokens == (4, 0) + assert estimate.has_hit is True + assert coordinator.estimate_calls == [ + ([1, 2, 3, 4], [10, 11, 12, 13, 14]), + ([1, 2, 3, 4], [20, 21]), + ] + assert coordinator.lookup_calls == [] + + +def test_lookup_prefix_cache_for_prefill_normalizes_full_hit(): + coordinator = _Coordinator(cached_tokens=[5], handles=[11]) + + lookup = lookup_prefix_cache_for_prefill( + coordinator=coordinator, + namespace_digest=(1, 2, 3, 4), + prompt_token_ids=[[10, 11, 12, 13, 14]], + ) + + assert lookup.prefix_shared_tokens == (4,) + + +def test_build_prefix_cache_prefill_inputs_uses_suffix_only_tokens(): + coordinator = _Coordinator(cached_tokens=[3, 0, 5], handles=[11, 0, 12]) + lookup = lookup_prefix_cache_for_prefill( + coordinator=coordinator, + namespace_digest=(1, 2, 3, 4), + prompt_token_ids=[ + [10, 11, 12, 13, 14], + [20, 21], + [30, 31, 32, 33, 34], + ], + ) + + inputs = build_prefix_cache_prefill_inputs( + local_indices=[7, 8, 9], + sequence_ids=[100, 101, 102], + input_ids=[ + torch.tensor([[10, 11, 12, 13, 14]]), + torch.tensor([[20, 21]]), + torch.tensor([[30, 31, 32, 33, 34]]), + ], + prompt_lengths=[5, 2, 5], + lookup=lookup, + ) + + assert lookup.prefix_shared_tokens == (3, 0, 4) + assert [item.tolist() for item in inputs.plan.suffix_input_ids] == [ + [13, 14], + [20, 21], + [34], + ] + assert [item.tolist() for item in inputs.plan.suffix_position_ids] == [ + [3, 4], + [0, 1], + [4], + ] + assert [item.tolist() for item in inputs.input_ids_list] == [ + [[13, 14]], + [[20, 21]], + [[34]], + ] + assert [item.tolist() for item in inputs.attention_mask_list] == [ + [[1, 1]], + [[1, 1]], + [[1]], + ] diff --git a/tests/unit/test_prefix_reuse_prefill_plan.py b/tests/unit/test_prefix_reuse_prefill_plan.py new file mode 100644 index 000000000..4d6960522 --- /dev/null +++ b/tests/unit/test_prefix_reuse_prefill_plan.py @@ -0,0 +1,111 @@ +import pytest +import torch + +from batchgen.prefill.prefix_reuse import ( + build_prefix_reuse_prefill_plan, + split_prefix_reuse_plan_for_micro_batch, +) + + +def test_build_prefix_reuse_prefill_plan_mixed_hit_and_miss(): + input_ids = [ + torch.tensor([[10, 11, 12, 13, 14, 15]]), + torch.tensor([[20, 21, 22, 23]]), + torch.tensor([[30, 31, 32, 33, 34]]), + ] + plan = build_prefix_reuse_prefill_plan( + local_indices=[0, 1, 2], + sequence_ids=[100, 101, 102], + input_ids=input_ids, + prompt_lengths=[6, 4, 5], + prefix_shared_tokens=[4, 0, 4], + ) + + assert [item.suffix_length for item in plan.sequences] == [2, 4, 1] + assert [item.suffix_start_pos for item in plan.sequences] == [4, 0, 4] + assert [tensor.tolist() for tensor in plan.suffix_input_ids] == [ + [14, 15], + [20, 21, 22, 23], + [34], + ] + assert [tensor.tolist() for tensor in plan.suffix_position_ids] == [ + [4, 5], + [0, 1, 2, 3], + [4], + ] + assert plan.cache_seqlens.tolist() == [4, 0, 4] + assert plan.total_prompt_tokens == 15 + assert plan.total_suffix_tokens == 7 + assert plan.saved_prefill_tokens == 8 + + +def test_split_prefix_reuse_prefill_plan_recomputes_stats(): + plan = build_prefix_reuse_prefill_plan( + local_indices=[0, 1, 2], + sequence_ids=[100, 101, 102], + input_ids=[ + torch.arange(0, 6), + torch.arange(10, 14), + torch.arange(20, 25), + ], + prompt_lengths=[6, 4, 5], + prefix_shared_tokens=[4, 0, 2], + ) + + micro = split_prefix_reuse_plan_for_micro_batch(plan, 1, 3) + + assert [item.sequence_id for item in micro.sequences] == [101, 102] + assert [tensor.tolist() for tensor in micro.suffix_input_ids] == [ + [10, 11, 12, 13], + [22, 23, 24], + ] + assert micro.cache_seqlens.tolist() == [0, 2] + assert micro.total_prompt_tokens == 9 + assert micro.total_suffix_tokens == 7 + assert micro.saved_prefill_tokens == 2 + + +def test_build_prefix_reuse_prefill_plan_uses_effective_full_hit_tokens(): + plan = build_prefix_reuse_prefill_plan( + local_indices=[0], + sequence_ids=[100], + input_ids=[torch.arange(0, 4)], + prompt_lengths=[4], + prefix_shared_tokens=[3], + ) + + assert plan.sequences[0].prefix_shared_tokens == 3 + assert plan.sequences[0].suffix_start_pos == 3 + assert plan.sequences[0].suffix_length == 1 + assert plan.suffix_input_ids[0].tolist() == [3] + + +def test_build_prefix_reuse_prefill_plan_one_token_full_hit_has_no_saved_tokens(): + plan = build_prefix_reuse_prefill_plan( + local_indices=[0], + sequence_ids=[100], + input_ids=[torch.tensor([42])], + prompt_lengths=[1], + prefix_shared_tokens=[0], + ) + + assert plan.sequences[0].prefix_shared_tokens == 0 + assert plan.sequences[0].suffix_start_pos == 0 + assert plan.sequences[0].suffix_length == 1 + assert plan.cache_seqlens.tolist() == [0] + assert plan.total_prompt_tokens == 1 + assert plan.total_suffix_tokens == 1 + assert plan.saved_prefill_tokens == 0 + assert plan.suffix_input_ids[0].tolist() == [42] + assert plan.suffix_position_ids[0].tolist() == [0] + + +def test_build_prefix_reuse_prefill_plan_validates_lengths(): + with pytest.raises(ValueError, match="must be smaller than prompt_length"): + build_prefix_reuse_prefill_plan( + local_indices=[0], + sequence_ids=[100], + input_ids=[torch.arange(0, 4)], + prompt_lengths=[4], + prefix_shared_tokens=[5], + ) diff --git a/tests/unit/test_prefix_worker_cleanup_invariants.py b/tests/unit/test_prefix_worker_cleanup_invariants.py new file mode 100644 index 000000000..db727eb69 --- /dev/null +++ b/tests/unit/test_prefix_worker_cleanup_invariants.py @@ -0,0 +1,62 @@ +from pathlib import Path + + +_WORKER_SOURCE = ( + Path(__file__).resolve().parents[2] / "batchgen" / "batchgen_worker.py" +) + + +def _source() -> str: + return _WORKER_SOURCE.read_text() + + +def _method_body(source: str, name: str, next_name: str) -> str: + start = source.index(f"\tdef {name}(") + end = source.index(f"\n\tdef {next_name}(", start) + return source[start:end] + + +def test_prefill_prepack_scope_cleans_global_state_in_finally(): + source = _source() + scope = _method_body( + source, + "_prefill_prepack_runtime_scope", + "prefill_prepacked", + ) + + assert "\n\t\tfinally:\n" in scope + assert "self._reset_prefill_prepack_runtime_state()" in scope + assert "prefix_materialization.close(empty_cuda_cache=False)" in scope + assert "self._destroy_gpu_paged_kv_cache(empty_cuda_cache=True)" in scope + + +def test_prefill_prepacked_uses_cleanup_scope_around_inference_loop(): + source = _source() + body = _method_body( + source, + "prefill_prepacked", + "_compute_boundary_decisions", + ) + scope_call = ( + "self._prefill_prepack_runtime_scope(batch_prefix_materialization)" + ) + inference_call = "torch.inference_mode()" + + assert scope_call in body + assert inference_call in body + assert body.index("batch_prefix_materialization = None") < body.index( + scope_call + ) + assert body.index(scope_call) > body.index("Prepacked Prefill") + + +def test_prefix_reuse_prefill_preserves_default_microbatch_cap(): + source = _source() + body = _method_body( + source, + "prefill_prepacked", + "_compute_boundary_decisions", + ) + + assert "BATCHGEN_PREFIX_REUSE_PREFILL_MICRO_BATCH_TOKEN_CAP" in body + assert '"131072"' in body diff --git a/tests/unit/test_result_gathering.py b/tests/unit/test_result_gathering.py index baefda7fa..4234f249d 100644 --- a/tests/unit/test_result_gathering.py +++ b/tests/unit/test_result_gathering.py @@ -8,6 +8,7 @@ import sys import torch +import pytest from dataclasses import dataclass from typing import List, Set, Optional @@ -199,6 +200,45 @@ def test_old_vs_new_equivalence(): print(" PASS: test_old_vs_new_equivalence") +def test_completed_outputs_cached_for_final_response(): + """Completed sequences remain available after local query slots are released.""" + try: + from batchgen.batchgen_worker import BatchGenWorker + except ImportError as exc: + pytest.skip(f"BatchGenWorker import requires runtime extensions: {exc}") + + @dataclass + class Sequence: + global_idx: int + + class Batch: + def __init__(self): + self._sequences = { + "seq_a": Sequence(global_idx=3), + "seq_b": Sequence(global_idx=1), + } + + def get_sequence(self, uuid: str): + return self._sequences.get(uuid) + + worker = object.__new__(BatchGenWorker) + worker.rank = 0 + worker.global_batch = Batch() + worker._final_response_completed_outputs = {} + + worker._record_completed_outputs_for_final_response( + ["seq_a", "seq_b", "missing"], + { + "seq_a": {"text": "alpha"}, + "seq_b": {"text": "beta"}, + "missing": {"text": "ignored"}, + }, + ) + + assert worker._final_response_completed_outputs == {3: "alpha", 1: "beta"} + print(" PASS: test_completed_outputs_cached_for_final_response") + + if __name__ == "__main__": print("Running result gathering tests...\n") @@ -214,6 +254,7 @@ def test_old_vs_new_equivalence(): test_gather_sorting, test_empty_rank, test_old_vs_new_equivalence, + test_completed_outputs_cached_for_final_response, ] passed = 0 diff --git a/tests/unit/test_usage_cached_tokens.py b/tests/unit/test_usage_cached_tokens.py new file mode 100644 index 000000000..b8790c054 --- /dev/null +++ b/tests/unit/test_usage_cached_tokens.py @@ -0,0 +1,191 @@ +import importlib.util +import json +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _load_module(monkeypatch, module_name: str, path: Path): + spec = importlib.util.spec_from_file_location(module_name, path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def _load_lightweight_usage_modules(monkeypatch): + batchgen_pkg = types.ModuleType("batchgen") + batchgen_pkg.__path__ = [str(REPO_ROOT / "batchgen")] + server_pkg = types.ModuleType("batchgen.server") + server_pkg.__path__ = [str(REPO_ROOT / "batchgen" / "server")] + monkeypatch.setitem(sys.modules, "batchgen", batchgen_pkg) + monkeypatch.setitem(sys.modules, "batchgen.server", server_pkg) + + io_struct = _load_module( + monkeypatch, + "batchgen.server.io_struct", + REPO_ROOT / "batchgen" / "server" / "io_struct.py", + ) + usage = _load_module( + monkeypatch, + "batchgen.server.usage", + REPO_ROOT / "batchgen" / "server" / "usage.py", + ) + return io_struct, usage + + +def _stub_batch_scheduler_deps(monkeypatch): + dependencies = { + "batchgen.server.intake_pool": { + "IntakeEntry": object, + "IntakePool": object, + "Priority": object, + }, + "batchgen.server.scheduling_pool": {"SchedulingPool": object}, + "batchgen.server.server_args": {"ServerArgs": object}, + "batchgen.server.storage": {"StorageManager": object}, + "batchgen.server.worker_manager": {"WorkerManager": object}, + } + for module_name, attrs in dependencies.items(): + module = types.ModuleType(module_name) + for attr_name, attr_value in attrs.items(): + setattr(module, attr_name, attr_value) + monkeypatch.setitem(sys.modules, module_name, module) + + +def _load_batch_scheduler(monkeypatch): + _load_lightweight_usage_modules(monkeypatch) + _stub_batch_scheduler_deps(monkeypatch) + return _load_module( + monkeypatch, + "batchgen.server.batch_scheduler", + REPO_ROOT / "batchgen" / "server" / "batch_scheduler.py", + ) + + +def _model_dict(model): + if hasattr(model, "model_dump"): + return model.model_dump() + return model.dict() + + +def test_usage_serializes_prompt_cached_tokens(monkeypatch): + _, usage_module = _load_lightweight_usage_modules(monkeypatch) + + usage = usage_module.build_usage( + prompt_tokens=128, + completion_tokens=16, + cached_tokens=64, + ) + + usage_dict = _model_dict(usage) + assert usage_dict["prompt_tokens_details"] == {"cached_tokens": 64} + assert usage_dict["total_tokens"] == 144 + + +def test_usage_clamps_cached_tokens_to_prompt_tokens(monkeypatch): + _, usage_module = _load_lightweight_usage_modules(monkeypatch) + + usage = usage_module.build_usage( + prompt_tokens=32, + completion_tokens=4, + cached_tokens=128, + ) + + assert usage.prompt_tokens_details.cached_tokens == 32 + + +def test_pool_completion_writes_cached_tokens(tmp_path, monkeypatch): + batch_scheduler = _load_batch_scheduler(monkeypatch) + scheduler = batch_scheduler.BatchScheduler.__new__( + batch_scheduler.BatchScheduler + ) + scheduler.server_args = SimpleNamespace( + incremental_output_dir=str(tmp_path) + ) + scheduler._pool_request_meta = { + "batch_1": { + "req_1": { + "custom_id": "custom-1", + "model": "openai/gpt-oss-120b", + "url": "/v1/chat/completions", + } + } + } + + scheduler._write_pool_completion( + "batch_1", + "req_1", + { + "text": "answer", + "prompt_length": 128, + "decoded_length": 8, + "cached_tokens": 64, + "finish_reason": "stop", + }, + ) + + output_path = tmp_path / "batch_1.jsonl" + line = json.loads(output_path.read_text().strip()) + usage = line["response"]["body"]["usage"] + + assert usage["prompt_tokens"] == 128 + assert usage["completion_tokens"] == 8 + assert usage["total_tokens"] == 136 + assert usage["prompt_tokens_details"] == {"cached_tokens": 64} + + +def test_batch_output_metrics_summarize_cached_tokens(tmp_path, monkeypatch): + batch_scheduler = _load_batch_scheduler(monkeypatch) + output_path = tmp_path / "batch.jsonl" + rows = [ + { + "custom_id": "req-1", + "response": { + "status_code": 200, + "body": { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "prompt_tokens_details": {"cached_tokens": 40}, + } + }, + }, + "error": None, + }, + { + "custom_id": "req-2", + "response": { + "status_code": 200, + "body": { + "usage": { + "prompt_tokens": 50, + "completion_tokens": 5, + "total_tokens": 55, + "prompt_tokens_details": {"cached_tokens": 0}, + } + }, + }, + "error": None, + }, + {"custom_id": "req-3", "response": None, "error": {"message": "bad"}}, + ] + output_path.write_text( + "\n".join(json.dumps(row) for row in rows) + "\n", + encoding="utf-8", + ) + + metrics = batch_scheduler._summarize_batch_output_file(output_path) + + assert metrics.rows == 3 + assert metrics.errors == 1 + assert metrics.prompt_tokens == 150 + assert metrics.completion_tokens == 15 + assert metrics.total_tokens == 165 + assert metrics.cached_tokens == 40 + assert metrics.requests_with_cache == 1 + assert metrics.cache_hit_rate == 40 / 150 diff --git a/tests/worker/test_prefill.py b/tests/worker/test_prefill.py index 5714eab41..896296fc9 100644 --- a/tests/worker/test_prefill.py +++ b/tests/worker/test_prefill.py @@ -19,6 +19,7 @@ PrefillCandidate, PrefillScheduler, PrefillSelectionRequest, + PrefillWaveGateRequest, ) _PAGE = 64 @@ -39,7 +40,38 @@ def _cand(uuid, *, rank=0, evicted=False, gidx=0, decoded=0, prompt=100, budget= ) -def _req(candidates, per_node_free, *, chunk=128, gpus_per_node=_GPN): +def _prefix_cand( + uuid, + *, + rank=0, + gidx=0, + prompt=4096, + cached=0, + page_ids=(), + budget=100000, +): + return PrefillCandidate( + uuid=uuid, + assigned_rank=rank, + is_evicted=False, + global_idx=gidx, + total_decoded_before_eviction=0, + prompt_length=prompt, + kv_token_budget=budget, + page_size=_PAGE, + estimated_shared_prefix_tokens=cached, + estimated_shared_prefix_page_ids=tuple(page_ids), + ) + + +def _req( + candidates, + per_node_free, + *, + chunk=128, + gpus_per_node=_GPN, + charge_shared_prefix_pages=False, +): return PrefillSelectionRequest( candidates=tuple(candidates), per_node_host_free=tuple(per_node_free), @@ -47,6 +79,7 @@ def _req(candidates, per_node_free, *, chunk=128, gpus_per_node=_GPN): num_nodes=len(per_node_free), gpus_per_node=gpus_per_node, initial_gpu_page_buffer=_BUF, + charge_shared_prefix_pages=charge_shared_prefix_pages, ) @@ -150,6 +183,56 @@ def test_no_eviction_candidates_pure_queueing_order(): assert plan == ["q2", "q1", "q0"] # uuids q2(gidx0), q1(gidx1), q0(gidx2) +def test_prefix_estimate_reduces_admission_pages(): + # Without prefix estimate, prompt 4096 needs 97 pages: + # max(prompt + chunk = 4224, gpu_tokens = (65 + 32) * 64). + # With a 3072-token page-aligned hit, the private charge drops to + # ceil((6208 - 3072) / 64) = 49 pages, so two candidates fit in 98 pages. + c0 = _prefix_cand("c0", gidx=0, prompt=4096, cached=3072) + c1 = _prefix_cand("c1", gidx=1, prompt=4096, cached=3072) + plan = PrefillScheduler.select_prefill_batch(_req([c0, c1], [98])) + assert plan == ["c0", "c1"] + plan = PrefillScheduler.select_prefill_batch(_req([c0, c1], [97])) + assert plan == ["c0"] + + +def test_non_page_aligned_prefix_estimate_is_conservative(): + # A full-hit compute path may normalize to prompt_length - 1. Admission + # must only credit fully page-aligned shared pages. + c = _prefix_cand("c", prompt=4096, cached=4095) + plan = PrefillScheduler.select_prefill_batch(_req([c], [33])) + assert plan == [] + plan = PrefillScheduler.select_prefill_batch(_req([c], [34])) + assert plan == ["c"] + + +def test_prefix_admission_charges_unique_shared_pages_when_requested(): + shared_pages = tuple((0, page_id) for page_id in range(48)) + c0 = _prefix_cand("c0", gidx=0, prompt=4096, cached=3072, page_ids=shared_pages) + c1 = _prefix_cand("c1", gidx=1, prompt=4096, cached=3072, page_ids=shared_pages) + + plan = PrefillScheduler.select_prefill_batch( + _req([c0, c1], [145], charge_shared_prefix_pages=True) + ) + assert plan == ["c0"] + + plan = PrefillScheduler.select_prefill_batch( + _req([c0, c1], [146], charge_shared_prefix_pages=True) + ) + assert plan == ["c0", "c1"] + + +def test_prefix_admission_does_not_charge_shared_pages_against_free_capacity(): + shared_pages = tuple((0, page_id) for page_id in range(48)) + c0 = _prefix_cand("c0", gidx=0, prompt=4096, cached=3072, page_ids=shared_pages) + c1 = _prefix_cand("c1", gidx=1, prompt=4096, cached=3072, page_ids=shared_pages) + + plan = PrefillScheduler.select_prefill_batch( + _req([c0, c1], [98], charge_shared_prefix_pages=False) + ) + assert plan == ["c0", "c1"] + + def test_request_and_candidate_are_frozen(): req = _req([_cand("a")], [34]) with pytest.raises((AttributeError, Exception)): @@ -157,3 +240,48 @@ def test_request_and_candidate_are_frozen(): c = _cand("a") with pytest.raises((AttributeError, Exception)): c.uuid = "b" # type: ignore[misc] + + +def test_prefix_cache_wave_gate_allows_first_wave(): + req = PrefillWaveGateRequest( + selected_count=1, + prefix_cache_enabled=True, + has_active_work=False, + world_size=8, + ) + assert PrefillScheduler.should_run_prefill_wave(req) + + +def test_prefix_cache_wave_gate_defers_small_wave_with_active_work(): + req = PrefillWaveGateRequest( + selected_count=35, + prefix_cache_enabled=True, + has_active_work=True, + world_size=8, + ) + assert not PrefillScheduler.should_run_prefill_wave(req) + + +def test_prefix_cache_wave_gate_allows_large_wave_with_active_work(): + req = PrefillWaveGateRequest( + selected_count=128, + prefix_cache_enabled=True, + has_active_work=True, + world_size=8, + ) + assert PrefillScheduler.should_run_prefill_wave(req) + + +def test_prefix_cache_wave_gate_does_not_change_non_prefix_cache_path(): + req = PrefillWaveGateRequest( + selected_count=1, + prefix_cache_enabled=False, + has_active_work=True, + world_size=8, + ) + assert PrefillScheduler.should_run_prefill_wave(req) + + +def test_prefix_cache_wave_gate_uses_world_size_threshold(): + assert PrefillScheduler.min_prefix_cache_wave_sequences(world_size=1) == 128 + assert PrefillScheduler.min_prefix_cache_wave_sequences(world_size=32) == 512 diff --git a/tests/worker/test_sync.py b/tests/worker/test_sync.py index 0131d0c66..d202b83bf 100644 --- a/tests/worker/test_sync.py +++ b/tests/worker/test_sync.py @@ -143,6 +143,22 @@ def test_completion_status_one_eos(coordinator, ctx): assert ctx.global_batch.get_sequence("bravo").status == SequenceStatus.COMPLETED +def test_completion_status_one_length_completed(coordinator, ctx): + """Length-complete sequences must not wait for a decode boundary.""" + seq = ctx.global_batch.get_sequence("bravo") + seq.decoded_length = seq.max_decode_length + seq.eos_reached = False + + completed, active = coordinator.sync_completion_status_tensor( + ctx, ["alpha", "bravo", "charlie"] + ) + + assert completed == {"bravo"} + assert active == ["alpha", "charlie"] + assert ctx.global_batch.get_sequence("bravo").status == SequenceStatus.COMPLETED + assert ctx.global_batch.get_sequence("bravo").eos_reached is True + + def test_completion_status_idempotent_mutation(coordinator, ctx): """Running twice in succession yields the same result; status guard prevents double-transition.""" ctx.global_batch.get_sequence("bravo").eos_reached = True