diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index dee3123c9..9700642df 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -180,6 +180,9 @@ void PagedCompiler::compile() { } PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &input) { + if (input.prefill_only) { + return {nullptr, nullptr}; + } if (model_->get_cache_config() != nullptr && dynamic_cast(model_->get_cache_config())) { size_t batch_size = input.block_tables.value()->size(0); size_t block_per_req = input.block_tables.value()->size(1); diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 422b5df73..947be6e2e 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -240,7 +240,8 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { image_req_ids, visual_token_ranges, to_device(target_hidden_states), - sample_all_positions}; + sample_all_positions, + prefill_only}; if (serialize_host_copy) { infinicore::context::syncStream(); @@ -269,6 +270,13 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } InferEngine::Output InferEngine::forward(const InferEngine::Input &input) { + // Every PP stage receives prefill_only and skips the sampled-ID exchange. + if (input.prefill_only && input.sample_all_positions) { + throw std::invalid_argument("prefill_only requires sample_all_positions=false"); + } + if (input.prefill_only && (!input.input_offsets.has_value() || input.input_offsets.value()->numel() < 2)) { + throw std::invalid_argument("prefill_only requires request input_offsets"); + } // Trigger each worker to run inference for (auto &worker : workers_) { worker->run(input); diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 0fa0a84cf..2b23b8b18 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -435,6 +435,24 @@ void RankWorker::thread_loop() { hidden_states = model_output.hidden_states; } + if (local_args.prefill_only) { + if (rank_info_.tp_rank == 0) { + // Preserve the old sampler's RNG advancement, but + // avoid LM output sampling and the token D2H copy. + const auto n_req = local_args.input_offsets.value()->numel() - 1; + for (size_t i = 0; i < n_req; ++i) { + (void)std::uniform_real_distribution(0, 1)(rng_); + } + // Keep the ordinary forward completion contract. + // Publication/cancellation follows this return. + infinicore::context::syncStream(); + } + output_ = Output{}; + job_done_ = true; + cv_.notify_all(); + continue; + } + if (rank_info_.pp_size > 1 && rank_info_.pp_stage + 1 != rank_info_.pp_size) { infinicore::Tensor output_ids; if (rank_info_.pp_stage == 0 && rank_info_.tp_rank == 0) { diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f1..10028cf20 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -79,6 +79,9 @@ class RankWorker { float top_p{1}; + /// Compute KV for an intermediate prefill without returning model output. + bool prefill_only{false}; + infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const; }; diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp index e39f27d9e..8dd307421 100644 --- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp +++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp @@ -52,7 +52,7 @@ class TextCausalLM : public InfinilmModel { */ Output forward(const Input &input) const override { auto hidden_states = model_->forward(input); - if (!is_last_pp_stage()) { + if (!is_last_pp_stage() || input.prefill_only) { return {infinicore::Tensor(), hidden_states}; } diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index 02677318d..c729cf4f7 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -57,6 +57,8 @@ class InfinilmModel : public infinicore::nn::Module { std::optional target_hidden_states; /// Preserve logits for every packed position for speculative/MTP callers. bool sample_all_positions{false}; + /// Intermediate prefill does not need the language-model head. + bool prefill_only{false}; }; struct Output { diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index c5e85577c..025abc1f5 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -122,18 +122,14 @@ inline void bind_infer_engine(py::module &m) { return state_dict_tp_all; }) .def("process_weights_after_loading", &InferEngine::process_weights_after_loading, "Process the weights after loading on all workers (e.g., for quantization)") - .def( - "forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { + .def("forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { // IMPORTANT: Release the GIL before calling forward() to allow other Python threads // to run concurrently during inference (which may block for a long time). // Do NOT remove this — without it, the GIL is held throughout inference and will // deadlock or stall any other Python thread (e.g., request handling, scheduling). py::gil_scoped_release release; - return self.forward(input); - }, - "Run inference on all ranks with arbitrary arguments") - .def( - "reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) + return self.forward(input); }, "Run inference on all ranks with arbitrary arguments") + .def("reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) .def("get_kv_cache", &InferEngine::get_kv_cache, "Get per-rank kv cache list") .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr { auto cfg = self.get_cache_config(); @@ -190,6 +186,7 @@ inline void bind_infer_engine(py::module &m) { // Allowed keyword arguments static const std::unordered_set allowed_kwargs = { + "prefill_only", "temperature", "top_p", "top_k", @@ -203,7 +200,9 @@ inline void bind_infer_engine(py::module &m) { "InferEngine.Input got an unexpected keyword argument '" + key + "'"); } - if (key == "temperature") { + if (key == "prefill_only") { + input.prefill_only = py::cast(item.second); + } else if (key == "temperature") { input.temperature = py::cast(item.second); } else if (key == "top_p") { input.top_p = py::cast(item.second); @@ -249,6 +248,7 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("image_req_ids", &InferEngine::Input::image_req_ids) .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) + .def_readwrite("prefill_only", &InferEngine::Input::prefill_only) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) diff --git a/docs/cache-and-chunking.md b/docs/cache-and-chunking.md new file mode 100644 index 000000000..8bfdd8e11 --- /dev/null +++ b/docs/cache-and-chunking.md @@ -0,0 +1,97 @@ +# Paged prefix caching and chunked Prefill + +Paged prefix caching uses LRU by default. Chunking is disabled by default. +For a dense text model, optional SLRU and bounded Prefill can be configured +through `EngineConfig`, `LLM`, or `AsyncLLMEngine`: + +```python +from infinilm.llm import AsyncLLMEngine + +engine = AsyncLLMEngine( + model_path="/path/to/model", + cache_type="paged", + prefix_cache_policy="slru", + prefix_cache_protected_ratio=0.8, + prefill_chunk_size=512, + tensor_parallel_size=1, + enable_graph=False, +) +``` + +The inference server and `examples/test_infer.py` expose the same settings: + +```sh +python python/infinilm/server/inference_server.py --model /path/to/model \ + --enable-paged-attn --prefix-cache-policy slru \ + --prefix-cache-protected-ratio 0.8 --prefill-chunk-size 512 +``` + +## Cache policies + +Only zero-reference pages can be reclaimed. Final release processes a request's +pages tail first, favoring its earlier prefix. Shared or remote-transfer owners +keep their pages pinned until the last reference is released. + +SLRU reclaims probationary pages before protected pages. A resident prefix is +promoted only after successful admission; lookup, rejected admission and newly +computed KV do not promote it. Protected membership includes pinned pages and +is capped at `floor(num_blocks * prefix_cache_protected_ratio)`. The oldest +protected member is demoted when this cap is exceeded. The ratio must be +strictly between zero and one; this cap does not reserve GPU memory. + +SLRU requires paged cache. Disabling prefix caching disables reuse and promotion. +It retains no history of evicted hashes. It can preserve established hotspots +through one-use traffic, but stale protected pages can delay adaptation to new +hotspots, and overcapacity cyclic traffic can still miss on every request. + +## Scheduling and parallel execution + +A Prefill dispatch computes at most +`min(prefill_chunk_size, max_num_batched_tokens)` tokens from one request. +Successful dispatches rotate among Decode, Prefill continuation and admission; +empty or blocked phases are skipped. Decode reuses the ordinary scheduler and +its batch limit. This bounds dispatch opportunities, not elapsed latency. + +Admission still reserves all prompt pages and future Decode headroom. Prefix +lookup pins only published full pages; failed admission releases temporary pins. +After successful execution, only completed full pages are published. Cancellation +releases request ownership while preserving reusable completed pages. +Intermediate chunks execute every Transformer layer but omit the LM head, +sampling and output tokens. The final chunk enters normal generation handling. + +Supported chunk configurations are TP/PP = 1/1, 2/1 and 1/2. TP2 requires working +CUDA collectives; PP2 uses eager execution and publishes only after both stages +complete. Static cache, MLA, draft models, remote KV connectors, Mamba, MoE and +multimodal inputs are excluded from chunking. The direct-native `bench.py`, +`llama.py` and `bench_videonsa.py` examples reject chunking because they bypass +the scheduler; `bench.py` also disables prefix caching. + +## Device graphs + +With `device="cuda"` and `enable_graph=True`, TP1/TP2 with PP1 can combine eager +Prefill and Decode graphs. InfiniCore must be built with `--graph=y`. Tested +attention backends are `paged-attn` on NVIDIA A6000 and `flash-attn` on MetaX +C500; `cuda` maps to MACA on the MetaX build. Other devices have not been +validated for these combinations. + +Intermediate Prefill chunks use eager execution. A single-token final tail +may use the existing Decode graph because it has the same one-query attention +semantics. This change adds no Prefill graph capture or configuration switch. +PP2 chunking remains eager. + +## Validation and tradeoffs + +See [test instructions](../test/llm/README.md) for CPU regressions and opt-in +native lifecycle checks. Hardware validation covered A6000 Qwen2.5-1.5B FP16 +and C500 Qwen3-0.6B/4B BF16. This does not establish support for every dense +architecture or backend. Quantized-model chunking has not been validated; +the dense-model check does not reject quantization metadata. + +Chunking can reduce long output pauses and short-request waiting while reducing +throughput and increasing long-request TTFT. Graphs add initialization time and +retained workspace; full memory overhead and sustained HTTP throughput were +not measured. A historical FP16 near-tied-token mismatch across prefix execution +shapes remains documented; universal bitwise equivalence is not promised. + +[PR #573](https://github.com/InfiniTensor/InfiniLM/pull/573) contains the measured +benefits, costs, test conditions and links to archived reports and experiments. diff --git a/examples/bench.py b/examples/bench.py index 17bfe1a6d..6d95db09c 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -753,6 +753,10 @@ def run( if __name__ == "__main__": cfg = BaseConfig() + if cfg.prefill_chunk_size: + raise ValueError( + "Chunked prefill requires the LLM engine; use test_infer.py or the server." + ) logging.basicConfig( level=getattr(logging, cfg.log_level.upper(), logging.INFO), format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", @@ -903,9 +907,11 @@ def run( ) cfg.max_cache_len = max( max_benchmark_cache_len, - next(iter(cases_dict.values()))["input_len"] + _WARMUP_DECODE_LEN - if cfg.warmup - else 0, + ( + next(iter(cases_dict.values()))["input_len"] + _WARMUP_DECODE_LEN + if cfg.warmup + else 0 + ), ) cfg.attn = attn_backend if enable_paged_attn: @@ -945,9 +951,11 @@ def run( block_size=cfg.block_size, max_cache_len=max( max_benchmark_cache_len, - next(iter(cases_dict.values()))["input_len"] + _WARMUP_DECODE_LEN - if cfg.warmup - else 0, + ( + next(iter(cases_dict.values()))["input_len"] + _WARMUP_DECODE_LEN + if cfg.warmup + else 0 + ), ), temperature=cfg.temperature, top_p=cfg.top_p, diff --git a/examples/bench_videonsa.py b/examples/bench_videonsa.py index 8c0338a83..6a2770e6e 100644 --- a/examples/bench_videonsa.py +++ b/examples/bench_videonsa.py @@ -179,6 +179,10 @@ def run_case(model, tokenizer, cfg, video_payload, batch_size, input_len, output def main(): cfg = BaseConfig() + if cfg.prefill_chunk_size: + raise ValueError( + "Chunked prefill requires the LLM engine; use test_infer.py or the server." + ) cfg.enable_prefix_caching = False normalize_bench_defaults(cfg) diff --git a/examples/llama.py b/examples/llama.py index a3f0f11f8..57ff7781e 100644 --- a/examples/llama.py +++ b/examples/llama.py @@ -1,13 +1,13 @@ -import infinicore -from transformers import AutoTokenizer -from tokenizers import decoders as _dec -from infinilm.modeling_utils import get_model_state_dict -import infinilm -import argparse +import os import sys import time -import os + +import infinicore +import infinilm from infinilm.base_config import BaseConfig +from infinilm.modeling_utils import get_model_state_dict +from tokenizers import decoders as _dec +from transformers import AutoTokenizer sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../python")) @@ -103,6 +103,10 @@ def test( if __name__ == "__main__": cfg = BaseConfig() + if cfg.prefill_chunk_size: + raise ValueError( + "Chunked prefill requires the LLM engine; use test_infer.py or the server." + ) device_str = cfg.get_device_str(cfg.device) diff --git a/examples/test_infer.py b/examples/test_infer.py index f17a8baa0..da6e2956e 100644 --- a/examples/test_infer.py +++ b/examples/test_infer.py @@ -41,6 +41,9 @@ def test( use_legacy_moe=False, enable_prefix_caching=True, pre_transpose=False, + prefix_cache_policy="lru", + prefix_cache_protected_ratio=0.8, + prefill_chunk_size=0, ): model_path = os.path.expanduser(model_path) # ---------------------------------------------------------------------------- # @@ -77,6 +80,9 @@ def test( use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, pre_transpose=pre_transpose, + prefix_cache_policy=prefix_cache_policy, + prefix_cache_protected_ratio=prefix_cache_protected_ratio, + prefill_chunk_size=prefill_chunk_size, ) conversations = [ @@ -185,4 +191,7 @@ def test( use_legacy_moe=cfg.use_legacy_moe, enable_prefix_caching=cfg.enable_prefix_caching, pre_transpose=cfg.pre_transpose, + prefix_cache_policy=cfg.prefix_cache_policy, + prefix_cache_protected_ratio=cfg.prefix_cache_protected_ratio, + prefill_chunk_size=cfg.prefill_chunk_size, ) diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c0..2f46c3927 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -8,6 +8,17 @@ from infinilm.moe_config import MOE_EP_BACKEND_HELP +def parse_nonnegative_int(value: str) -> int: + """Parse an integer option that uses zero to disable its feature.""" + try: + result = int(value) + except ValueError: + raise argparse.ArgumentTypeError("value must be a nonnegative integer") + if result < 0: + raise argparse.ArgumentTypeError("value must be a nonnegative integer") + return result + + def parse_list(value: str): """Parse parse_list argument: can be a single int or a list of ints. @@ -76,6 +87,18 @@ def __init__(self): self.enable_graph = self.args.enable_graph self.enable_paged_attn = self.args.enable_paged_attn self.enable_prefix_caching = self.args.enable_prefix_caching + self.prefix_cache_policy = self.args.prefix_cache_policy + self.prefix_cache_protected_ratio = self.args.prefix_cache_protected_ratio + if not 0 < self.prefix_cache_protected_ratio < 1: + self.parser.error("--prefix-cache-protected-ratio must be between 0 and 1") + self.prefill_chunk_size = self.args.prefill_chunk_size + # Worker entrypoints branch before constructing the host LLM engine. + if self.prefill_chunk_size and ( + (self.tp, self.pp) not in {(1, 1), (2, 1), (1, 2)} + ): + self.parser.error("--prefill-chunk-size requires TP/PP=1/1, 2/1 or 1/2") + if self.prefill_chunk_size and self.draft_model: + self.parser.error("--prefill-chunk-size does not support --draft-model") self.use_mla = self.args.use_mla self.pre_transpose = self.args.pre_transpose self.num_blocks = self.args.num_blocks @@ -274,6 +297,24 @@ def _add_common_args(self): default=True, help="disable KV prefix cache reuse", ) + self.parser.add_argument( + "--prefix-cache-policy", + choices=["lru", "slru"], + default="lru", + help="paged prefix-cache eviction policy", + ) + self.parser.add_argument( + "--prefix-cache-protected-ratio", + type=float, + default=0.8, + help="fraction of paged blocks protected by SLRU (strictly between 0 and 1)", + ) + self.parser.add_argument( + "--prefill-chunk-size", + type=parse_nonnegative_int, + default=0, + help="maximum prompt tokens per prefill step (0 disables; paged TP/PP=1/1, 2/1 or 1/2; graphs require PP=1)", + ) self.parser.add_argument( "--num-blocks", type=int, default=512, help="number of KV cache blocks" ) diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index bccb7e758..cdab542c6 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -28,6 +28,9 @@ class EngineConfig: block_size: Size of each KV cache block (only for paged cache). max_cache_len: Maximum sequence length (only for static cache). enable_prefix_caching: Whether to reuse KV cache across requests. + prefix_cache_policy: Paged prefix-cache eviction policy ('lru' or 'slru'). + prefix_cache_protected_ratio: Fraction of paged blocks protected by SLRU. + prefill_chunk_size: Maximum prompt tokens per prefill step; 0 disables chunking. temperature: Default sampling temperature. top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. @@ -69,8 +72,49 @@ class EngineConfig: use_legacy_moe: bool = False kv_transfer_config: Optional[KVTransferConfig] = None enable_prefix_caching: bool = True + prefix_cache_policy: str = "lru" + prefix_cache_protected_ratio: float = 0.8 + + prefill_chunk_size: int = 0 def __post_init__(self) -> None: + if self.prefix_cache_policy not in {"lru", "slru"}: + raise ValueError("prefix_cache_policy must be either 'lru' or 'slru'") + if not 0 < self.prefix_cache_protected_ratio < 1: + raise ValueError("prefix_cache_protected_ratio must be between 0 and 1") + if self.prefix_cache_policy == "slru" and self.cache_type != "paged": + raise ValueError("prefix_cache_policy='slru' requires cache_type='paged'") + + if ( + isinstance(self.prefill_chunk_size, bool) + or not isinstance(self.prefill_chunk_size, int) + or self.prefill_chunk_size < 0 + ): + raise ValueError("prefill_chunk_size must be a nonnegative integer") + if self.prefill_chunk_size: + if self.cache_type != "paged": + raise ValueError("prefill_chunk_size requires cache_type='paged'") + if (self.tensor_parallel_size, self.pipeline_parallel_size) not in { + (1, 1), + (2, 1), + (1, 2), + }: + raise ValueError("prefill_chunk_size requires TP/PP=1/1, 2/1 or 1/2") + if self.enable_graph and ( + self.pipeline_parallel_size != 1 + or self.device != "cuda" + or self.attn_backend not in {"default", "paged-attn", "flash-attn"} + ): + raise ValueError( + "prefill_chunk_size with enable_graph requires PP=1, " + "device='cuda' and a supported paged attention backend" + ) + if self.use_mla: + raise ValueError("prefill_chunk_size does not support MLA") + if self.draft_model_path: + raise ValueError("prefill_chunk_size does not support draft models") + if self.kv_transfer_config and self.kv_transfer_config.kv_connector: + raise ValueError("prefill_chunk_size does not support KV transfer") if self.num_draft_tokens < 1: raise ValueError("num_draft_tokens must be >= 1") if self.pipeline_parallel_size < 1: diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 117d82f9a..e012bcb91 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -275,6 +275,7 @@ def _build_input( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=False, + prefill_only=False, temperature=None, top_k=None, top_p=None, @@ -333,6 +334,7 @@ def convert_tensor_list(tensor_list_): visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + prefill_only=prefill_only, temperature=temperature, top_k=top_k, top_p=top_p, @@ -358,85 +360,40 @@ def forward( image_req_ids=None, visual_token_ranges=None, target_hidden_states=None, + prefill_only=False, temperature=None, top_k=None, top_p=None, ): try: - # TODO: Remove `_underlying` and simplify the corresponding code. - input_ids = input_ids._underlying if input_ids is not None else None - position_ids = ( - position_ids._underlying if position_ids is not None else None - ) - past_kv_lengths = ( - past_kv_lengths._underlying if past_kv_lengths is not None else None - ) - total_kv_lengths = ( - total_kv_lengths._underlying if total_kv_lengths is not None else None - ) - input_offsets = ( - input_offsets._underlying if input_offsets is not None else None - ) - block_tables = ( - block_tables._underlying if block_tables is not None else None - ) - cu_seqlens = cu_seqlens._underlying if cu_seqlens is not None else None - slot_mapping = ( - slot_mapping._underlying if slot_mapping is not None else None - ) - mamba_init_state_indices = ( - mamba_init_state_indices._underlying - if mamba_init_state_indices is not None - else None - ) - mamba_final_state_indices = ( - mamba_final_state_indices._underlying - if mamba_final_state_indices is not None - else None - ) - - def convert_tensor_list(tensor_list_): - if tensor_list_ is None: - return None - if not isinstance(tensor_list_, list): - tensor_list_ = [tensor_list_] - if len(tensor_list_) == 0: - return None - return [tensor._underlying for tensor in tensor_list_] - - pixel_values = convert_tensor_list(pixel_values) - image_bound = convert_tensor_list(image_bound) - tgt_sizes = convert_tensor_list(tgt_sizes) - image_grid_thw = convert_tensor_list(image_grid_thw) - - return infinicore.Tensor( - super() - .forward( - self._build_input( - input_ids, - position_ids=position_ids, - past_kv_lengths=past_kv_lengths, - total_kv_lengths=total_kv_lengths, - input_offsets=input_offsets, - cu_seqlens=cu_seqlens, - block_tables=block_tables, - slot_mapping=slot_mapping, - mamba_init_state_indices=mamba_init_state_indices, - mamba_final_state_indices=mamba_final_state_indices, - pixel_values=pixel_values, - image_bound=image_bound, - tgt_sizes=tgt_sizes, - image_grid_thw=image_grid_thw, - image_req_ids=image_req_ids, - visual_token_ranges=visual_token_ranges, - target_hidden_states=target_hidden_states, - temperature=temperature, - top_k=top_k, - top_p=top_p, - ) + output = super().forward( + self._build_input( + input_ids, + position_ids=position_ids, + past_kv_lengths=past_kv_lengths, + total_kv_lengths=total_kv_lengths, + input_offsets=input_offsets, + cu_seqlens=cu_seqlens, + block_tables=block_tables, + slot_mapping=slot_mapping, + mamba_init_state_indices=mamba_init_state_indices, + mamba_final_state_indices=mamba_final_state_indices, + pixel_values=pixel_values, + image_bound=image_bound, + tgt_sizes=tgt_sizes, + image_grid_thw=image_grid_thw, + image_req_ids=image_req_ids, + visual_token_ranges=visual_token_ranges, + target_hidden_states=target_hidden_states, + prefill_only=prefill_only, + temperature=temperature, + top_k=top_k, + top_p=top_p, ) - .output_ids ) + if prefill_only: + return None + return infinicore.Tensor(output.output_ids) except BaseException as e: handle_oom_and_exit(e) raise diff --git a/python/infinilm/llm/cache_manager.py b/python/infinilm/llm/cache_manager.py index 6c045e4ee..f58a33afb 100644 --- a/python/infinilm/llm/cache_manager.py +++ b/python/infinilm/llm/cache_manager.py @@ -1,6 +1,6 @@ """Paged KV cache allocation and source-agnostic prefix lookup.""" -from collections import deque +from collections import OrderedDict, deque from collections.abc import Sequence from typing import Dict, List, Set @@ -71,16 +71,32 @@ def get_num_free_blocks(self) -> int: class BlockManager: """Manage physical paged-cache blocks and published prefix hashes.""" - def __init__(self, num_blocks: int, block_size: int): + def __init__( + self, + num_blocks: int, + block_size: int, + prefix_cache_policy: str = "lru", + prefix_cache_protected_ratio: float = 0.8, + ): if num_blocks <= 0 or block_size <= 0: raise ValueError("num_blocks and block_size must be positive") + if prefix_cache_policy not in {"lru", "slru"}: + raise ValueError("`prefix_cache_policy` must be 'lru' or 'slru'.") + if not 0 < prefix_cache_protected_ratio < 1: + raise ValueError("`prefix_cache_protected_ratio` must be between 0 and 1.") self.num_blocks = num_blocks self.block_size = block_size + self.prefix_cache_policy = prefix_cache_policy + self._protected_capacity = int(num_blocks * prefix_cache_protected_ratio) self.blocks: List[Block] = [Block(i) for i in range(num_blocks)] self.hash_to_block_ids: Dict[BlockHash, Set[int]] = {} self.free_block_ids: deque[int] = deque(range(num_blocks)) self.used_block_ids: Set[int] = set() + self._evictable_blocks: OrderedDict[int, None] = OrderedDict() + self._protected_evictable_blocks: OrderedDict[int, None] = OrderedDict() + # Membership survives pinning so active requests cannot bypass the cap. + self._protected_blocks: OrderedDict[int, None] = OrderedDict() def __repr__(self) -> str: return ( @@ -115,6 +131,9 @@ def _deallocate_block(self, block_id: int) -> None: f"Block {block_id} ref_count not zero, cannot deallocate" ) self._remove_block_hash(block) + self._evictable_blocks.pop(block_id, None) + self._protected_evictable_blocks.pop(block_id, None) + self._protected_blocks.pop(block_id, None) block.free() self.used_block_ids.remove(block_id) self.free_block_ids.append(block_id) @@ -126,12 +145,11 @@ def get_num_free_blocks(self) -> int: return len(self.free_block_ids) def get_total_usable_blocks(self) -> int: - freeable_used_blocks = sum( - 1 - for block_id in self.used_block_ids - if self.blocks[block_id].ref_count == 0 + return ( + len(self.free_block_ids) + + len(self._evictable_blocks) + + len(self._protected_evictable_blocks) ) - return len(self.free_block_ids) + freeable_used_blocks def get_computed_blocks( self, @@ -151,10 +169,30 @@ def get_computed_blocks( block_id = next(iter(block_ids)) block = self.blocks[block_id] assert block.hash == block_hash and block_id in self.used_block_ids + if block.ref_count == 0: + if block_id in self._protected_blocks: + del self._protected_evictable_blocks[block_id] + else: + del self._evictable_blocks[block_id] block.ref_count += 1 cached_block_table.append(block_id) return cached_block_table, len(cached_block_table) * self.block_size + def record_cache_hit(self, block_table: Sequence[int]) -> None: + """Promote locally matched blocks only after request admission succeeds.""" + if self.prefix_cache_policy != "slru": + return + for block_id in reversed(block_table): + block = self.blocks[block_id] + assert block.ref_count > 0 and block.hash != EMPTY_BLOCK_HASH + self._protected_blocks[block_id] = None + self._protected_blocks.move_to_end(block_id) + while len(self._protected_blocks) > self._protected_capacity: + block_id, _ = self._protected_blocks.popitem(last=False) + if self.blocks[block_id].ref_count == 0: + del self._protected_evictable_blocks[block_id] + self._evictable_blocks[block_id] = None + def allocate_slots( self, num_new_tokens: int, @@ -328,23 +366,28 @@ def append_slot( return block_table, last_block_id * self.block_size + offset def free_blocks(self, block_table: Sequence[int]) -> None: - """Release request references while retaining computed blocks for reuse.""" + """Release references and retain only reusable computed blocks.""" for block_id in reversed(block_table): block = self.blocks[block_id] - assert block.ref_count > 0, "block ref_count must be greater than 0" + assert block.ref_count > 0, "Block reference count must be positive." block.ref_count -= 1 + if block.ref_count == 0: + if block.hash == EMPTY_BLOCK_HASH: + self._deallocate_block(block_id) + elif block_id in self._protected_blocks: + self._protected_blocks.move_to_end(block_id) + self._protected_evictable_blocks[block_id] = None + else: + self._evictable_blocks[block_id] = None def try_free_blocks(self, num_required: int) -> bool: - """Evict unreferenced blocks until the requested capacity is available.""" - to_free = [ - block_id - for block_id in self.used_block_ids - if self.blocks[block_id].ref_count == 0 - ] - for block_id in to_free: + """Reclaim probationary blocks before protected blocks, oldest first.""" + while not self.can_allocate(num_required): + candidates = self._evictable_blocks or self._protected_evictable_blocks + if not candidates: + break + block_id = next(iter(candidates)) self._deallocate_block(block_id) - if self.can_allocate(num_required): - return True return self.can_allocate(num_required) def update_blocks_slot( diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..946c6aa28 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -42,6 +42,18 @@ def __init__(self, config: EngineConfig): self.config = config hf_config = read_hf_config(config.model_path) has_mamba_cache = model_uses_mamba_cache(hf_config) + if config.prefill_chunk_size: + text_config = hf_config.get("text_config", hf_config) + if ( + has_mamba_cache + or any( + text_config.get(key) + for key in ("num_experts", "num_local_experts", "n_routed_experts") + ) + or "vision_config" in hf_config + or "audio_config" in hf_config + ): + raise ValueError("Chunked prefill supports dense text models only.") if has_mamba_cache and config.enable_prefix_caching: model_type = hf_config["model_type"] raise RuntimeError( @@ -110,6 +122,9 @@ def __init__(self, config: EngineConfig): has_mamba_cache=has_mamba_cache, num_mamba_cache_blocks=num_mamba_cache_blocks, enable_prefix_caching=config.enable_prefix_caching, + prefix_cache_policy=config.prefix_cache_policy, + prefix_cache_protected_ratio=config.prefix_cache_protected_ratio, + prefill_chunk_size=config.prefill_chunk_size, ) logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") if has_mamba_cache: @@ -157,6 +172,19 @@ def step(self) -> tuple[bool, list[tuple]]: runner_output = self.model_runner.execute_model(scheduler_output) sampled_token_ids = runner_output.sampled_token_ids self.scheduler.update_from_output(runner_output) + end = getattr(scheduler_output, "prefill_end", None) + if end is not None: + req = scheduler_output.scheduled_requests[0] + req.num_computed_tokens = end + if end < req.get_prompt_length(): + self.scheduler.commit_computed_tokens(req, end) + if req.is_aborted() or req.is_finished(): + if not req.is_finished(): + req.mark_canceled() + self.scheduler.complete_requests([req]) + else: + self.scheduler.requeue_prefill(req) + return True, [] pending = self._update_requests( scheduler_output.scheduled_requests, sampled_token_ids, @@ -366,6 +394,9 @@ def __init__( skip_load: bool = False, use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + prefix_cache_policy: str = "lru", + prefix_cache_protected_ratio: float = 0.8, + prefill_chunk_size: int = 0, ): """Initialize LLM. @@ -379,7 +410,10 @@ def __init__( max_tokens: Default maximum tokens to generate. num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). + prefill_chunk_size: Prompt tokens per prefill step; 0 disables chunking. max_cache_len: Maximum sequence length (only for static cache). + prefix_cache_policy: Paged prefix-cache eviction policy ('lru' or 'slru'). + prefix_cache_protected_ratio: Fraction of paged blocks protected by SLRU. temperature: Default sampling temperature. top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. @@ -418,6 +452,9 @@ def __init__( skip_load=skip_load, use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, + prefix_cache_policy=prefix_cache_policy, + prefix_cache_protected_ratio=prefix_cache_protected_ratio, + prefill_chunk_size=prefill_chunk_size, ) self.engine = LLMEngine(config) self.config = config @@ -594,6 +631,9 @@ def __init__( weight_load_mode: str = "async", use_legacy_moe: bool = False, enable_prefix_caching: bool = True, + prefix_cache_policy: str = "lru", + prefix_cache_protected_ratio: float = 0.8, + prefill_chunk_size: int = 0, ): """Initialize AsyncLLMEngine. @@ -607,7 +647,10 @@ def __init__( max_tokens: Default maximum tokens to generate. num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). + prefill_chunk_size: Prompt tokens per prefill step; 0 disables chunking. max_cache_len: Maximum sequence length (only for static cache). + prefix_cache_policy: Paged prefix-cache eviction policy ('lru' or 'slru'). + prefix_cache_protected_ratio: Fraction of paged blocks protected by SLRU. temperature: Default sampling temperature. top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. @@ -651,6 +694,9 @@ def __init__( weight_load_mode=weight_load_mode, use_legacy_moe=use_legacy_moe, enable_prefix_caching=enable_prefix_caching, + prefix_cache_policy=prefix_cache_policy, + prefix_cache_protected_ratio=prefix_cache_protected_ratio, + prefill_chunk_size=prefill_chunk_size, ) self.engine = LLMEngine(config) self.config = config diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index a1696f848..499b108c0 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -221,6 +221,14 @@ def _model_forward(self, scheduler_output): if self.speculative_runner is not None: return self._model_forward_with_speculative(scheduler_output, model_input) + end = getattr(scheduler_output, "prefill_end", None) + prefill_only = ( + end is not None + and end < scheduler_output.scheduled_requests[0].get_prompt_length() + ) + if prefill_only: + model_input["prefill_only"] = True + # Wake every stage before stage 0 enters forward. Each worker receives # the same metadata and then blocks in its model on the activation from # the preceding stage. Stage 0 waits for all acknowledgements afterward. @@ -237,6 +245,8 @@ def _model_forward(self, scheduler_output): raise if self.pipeline_control is not None: self.pipeline_control.wait_forward() + if prefill_only: + return [] sampled_tokens_list = sampled_tokens.to_numpy().tolist() return sampled_tokens_list diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index c10b55f2f..3c1c91a52 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -4,6 +4,7 @@ import logging import queue +from collections import deque from typing import List, Optional import janus @@ -44,12 +45,14 @@ def __init__( scheduled_requests: List[InferenceRequest], is_prefill: bool = False, speculative_cache_ops: Optional[SpeculativeCacheOps] = None, + prefill_end: int | None = None, ): self.scheduled_requests = scheduled_requests self.num_requests = len(scheduled_requests) self.is_prefill = is_prefill self.speculative_cache_ops = speculative_cache_ops self.kv_connector_metadata = None + self.prefill_end = prefill_end class Scheduler: @@ -71,7 +74,16 @@ def __init__( has_mamba_cache: bool = False, num_mamba_cache_blocks: int | None = None, enable_prefix_caching: bool = True, + prefix_cache_policy: str = "lru", + prefix_cache_protected_ratio: float = 0.8, + prefill_chunk_size: int = 0, ): + if type(prefill_chunk_size) is not int or prefill_chunk_size < 0: + raise ValueError("`prefill_chunk_size` must be a nonnegative integer.") + if prefill_chunk_size and (connector is not None or has_mamba_cache): + raise ValueError("Chunked prefill does not support remote KV or Mamba.") + if prefill_chunk_size and max_num_batched_tokens <= 0: + raise ValueError("Chunked prefill requires a positive token budget.") self.waiting_queue = janus.Queue() self.running_queue = janus.Queue() self.max_batch_size = max_batch_size @@ -82,7 +94,12 @@ def __init__( self.pending_kv_decode_blocks: int = 0 self.remote_kv_requests: dict[str, InferenceRequest] = {} - self.cache_manager = BlockManager(num_blocks=num_blocks, block_size=block_size) + self.cache_manager = BlockManager( + num_blocks=num_blocks, + block_size=block_size, + prefix_cache_policy=prefix_cache_policy, + prefix_cache_protected_ratio=prefix_cache_protected_ratio, + ) self.has_mamba_cache = has_mamba_cache self.mamba_cache_manager = ( MambaCacheManager(num_mamba_cache_blocks or max(2, num_blocks // 4)) @@ -94,9 +111,16 @@ def __init__( self.max_num_batched_tokens = max_num_batched_tokens self.connector = connector self.enable_prefix_caching = enable_prefix_caching + self.prefill_chunk_size = prefill_chunk_size + self.chunking_queue: deque[InferenceRequest] = deque() + self._next_chunk_phase = 0 def add_request(self, request: InferenceRequest): if request is not None: + if self.prefill_chunk_size and request.has_multimodal_inputs: + raise ValueError( + "Chunked prefill does not support multimodal requests." + ) # TODO: Remove the multimodal exclusion once media-aware prefix # hashing and model-side cache-boundary handling are supported. request.initialize_block_hashes( @@ -128,9 +152,10 @@ def _exceeds_token_budget( def schedule(self) -> Optional[SchedulerOutput]: """Schedule and return batch of requests to execute.""" + if self.prefill_chunk_size: + return self._schedule_chunked() deferred_requests = [] scheduled_requests = [] - is_prefill = False current_num_batched_tokens = 0 current_prefill_extra_blocks = 0 @@ -258,6 +283,7 @@ def schedule(self) -> Optional[SchedulerOutput]: num_external_computed_tokens, self.block_size, ) + self.cache_manager.record_cache_hit(cached_block_table) else: load_kv_async = False num_tokens_this_step = ( @@ -294,17 +320,13 @@ def schedule(self) -> Optional[SchedulerOutput]: # Return prefill batch if any waiting requests were scheduled if scheduled_requests: - is_prefill = True - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output + return self._make_output(scheduled_requests, is_prefill=True) + return self._schedule_decode() + + def _schedule_decode(self) -> Optional[SchedulerOutput]: + """Schedule Decode and remote-KV progress for either Prefill policy.""" + scheduled_requests = [] # Process Running queue (decode phase) while len(scheduled_requests) < self.max_batch_size: try: @@ -351,31 +373,111 @@ def schedule(self) -> Optional[SchedulerOutput]: else: break # Defer promotion to next schedule() if batch is full - # Return decode batch if any running requests were scheduled - if scheduled_requests: - is_prefill = False - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output + if scheduled_requests or self.connector is not None: + return self._make_output(scheduled_requests) + return None + def _make_output( + self, + requests: List[InferenceRequest], + is_prefill: bool = False, + prefill_end: int | None = None, + ) -> SchedulerOutput: + output = SchedulerOutput( + requests, + is_prefill=is_prefill, + speculative_cache_ops=self.speculative_cache_ops, + prefill_end=prefill_end, + ) if self.connector is not None: - scheduler_output = SchedulerOutput( - scheduled_requests=[], - speculative_cache_ops=self.speculative_cache_ops, - ) - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output + output.kv_connector_metadata = self.connector.build_connector_meta() + return output + + def _schedule_chunked(self) -> Optional[SchedulerOutput]: + """Rotate dispatch opportunities across decode, continuation, and admission.""" + phases = ( + self._schedule_decode, + self._schedule_continuation, + self._admit_chunk_request, + ) + for offset in range(len(phases)): + phase = (self._next_chunk_phase + offset) % len(phases) + output = phases[phase]() + if output is not None: + self._next_chunk_phase = (phase + 1) % len(phases) + return output + return None + def _schedule_continuation(self) -> Optional[SchedulerOutput]: + while self.chunking_queue: + req = self.chunking_queue.popleft() + if req.is_finished(): + self.complete_requests([req]) + continue + return self._prefill_chunk(req) return None + def _admit_chunk_request(self) -> Optional[SchedulerOutput]: + while True: + try: + req = self.waiting_queue.sync_q.get_nowait() + except queue.Empty: + return None + if req.is_finished(): + self.complete_requests([req]) + continue + cached_table, cached_tokens = ( + self.cache_manager.get_computed_blocks( + req.block_hashes, req.get_prompt_length() - 1 + ) + if self.enable_prefix_caching + else ([], 0) + ) + # Partial prompts own their prompt pages but still need decode headroom. + chunk_headroom = sum( + self._get_prefill_extra_blocks(other) + for other in self.chunking_queue + if not other.is_finished() + ) + allocation = None + if self.can_accept_request(req, cached_tokens, chunk_headroom): + allocation = self.cache_manager.allocate_slots( + req.get_prompt_length() - cached_tokens, + num_computed_tokens=cached_tokens, + cached_block_table=cached_table, + ) + if allocation is None: + self.cache_manager.free_blocks(cached_table) + self.waiting_queue.sync_q.put(req) + return None + self.cache_manager.record_cache_hit(cached_table) + req.block_table, _ = allocation + req.num_blocks = len(req.block_table) + req.num_cache_indexed_blocks = len(cached_table) + req.num_computed_tokens = cached_tokens + req.status = RequestStatus.RUNNING + return self._prefill_chunk(req) + + def _prefill_chunk(self, req: InferenceRequest) -> SchedulerOutput: + start = req.num_computed_tokens + end = min( + req.get_prompt_length(), + start + min(self.prefill_chunk_size, self.max_num_batched_tokens), + ) + req.num_local_cached_tokens = start + req.slot_mapping = self.cache_manager.update_blocks_slot( + req.block_table, start, end + ) + return self._make_output( + [req], + is_prefill=True, + prefill_end=end, + ) + + def requeue_prefill(self, request: InferenceRequest) -> None: + """Retain ownership while waiting for the next prefill segment.""" + self.chunking_queue.append(request) + def update_waiting_for_remote_kv(self, request: InferenceRequest): self.remote_kv_requests.pop(request.request_id, None) self.pending_kv_decode_blocks -= ( diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..0076ab34e 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -207,11 +207,12 @@ def _build_model_input_from_batch_scheduler_output( if scheduler_output.is_prefill: # Prefill phase req_tokens = req.get_input_tokens() - tokens_to_compute = req_tokens[num_cached:] + prefill_end = scheduler_output.prefill_end + seq_len = len(req_tokens) if prefill_end is None else prefill_end + tokens_to_compute = req_tokens[num_cached:seq_len] tokens.extend(tokens_to_compute) compute_len = len(tokens_to_compute) - seq_len = len(req_tokens) seq_lens.append(seq_len) current_offset += compute_len diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..9651d1881 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -125,6 +125,9 @@ def __init__( kv_transfer_config: Optional[KVTransferConfig] = None, enable_prefix_caching: bool = True, pre_transpose: bool = False, + prefix_cache_policy: str = "lru", + prefix_cache_protected_ratio: float = 0.8, + prefill_chunk_size: int = 0, ): """Initialize inference server. @@ -142,6 +145,8 @@ def __init__( num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). max_cache_len: Maximum sequence length (only for static cache). + prefix_cache_policy: Paged prefix-cache eviction policy ('lru' or 'slru'). + prefix_cache_protected_ratio: Fraction of paged blocks protected by SLRU. temperature: Default sampling temperature. top_p: Default top-p sampling parameter. top_k: Default top-k sampling parameter. @@ -154,6 +159,7 @@ def __init__( weight_load_mode: Weight loading mode across tensor-parallel workers. ignore_eos: Whether to ignore EOS tokens during generation. kv_transfer_config: Optional configuration for the KV transfer mechanism. + prefill_chunk_size: Maximum prompt tokens per prefill step; 0 disables chunking. """ self.model_path = model_path # vLLM-like served model id: directory name of model_path @@ -187,7 +193,10 @@ def __init__( self.ignore_eos = ignore_eos self.kv_transfer_config = kv_transfer_config self.enable_prefix_caching = enable_prefix_caching + self.prefix_cache_policy = prefix_cache_policy + self.prefix_cache_protected_ratio = prefix_cache_protected_ratio self.pre_transpose = pre_transpose + self.prefill_chunk_size = prefill_chunk_size self.engine: AsyncLLMEngine = None @@ -231,7 +240,10 @@ async def lifespan(app: FastAPI): weight_load_mode=self.weight_load_mode, kv_transfer_config=self.kv_transfer_config, enable_prefix_caching=self.enable_prefix_caching, + prefix_cache_policy=self.prefix_cache_policy, + prefix_cache_protected_ratio=self.prefix_cache_protected_ratio, pre_transpose=self.pre_transpose, + prefill_chunk_size=self.prefill_chunk_size, ) self.engine.start() logger.info(f"Engine initialized with model at {self.model_path}") @@ -666,7 +678,10 @@ def main(): ignore_eos=cfg.ignore_eos, kv_transfer_config=kv_transfer_config, enable_prefix_caching=cfg.enable_prefix_caching, + prefix_cache_policy=cfg.prefix_cache_policy, + prefix_cache_protected_ratio=cfg.prefix_cache_protected_ratio, pre_transpose=cfg.pre_transpose, + prefill_chunk_size=cfg.prefill_chunk_size, ) server.start() diff --git a/python/infinilm/server/pipeline_worker.py b/python/infinilm/server/pipeline_worker.py index 9e4bb8e6b..d38c05cde 100644 --- a/python/infinilm/server/pipeline_worker.py +++ b/python/infinilm/server/pipeline_worker.py @@ -38,6 +38,10 @@ def run_worker(cfg: BaseConfig) -> None: weight_load_mode=cfg.weight_load_mode, skip_load=cfg.skip_load, use_legacy_moe=cfg.use_legacy_moe, + prefill_chunk_size=cfg.prefill_chunk_size, + enable_prefix_caching=cfg.enable_prefix_caching, + prefix_cache_policy=cfg.prefix_cache_policy, + prefix_cache_protected_ratio=cfg.prefix_cache_protected_ratio, ) runner = ModelRunner(config, initialize_processor=False) diff --git a/test/llm/README.md b/test/llm/README.md new file mode 100644 index 000000000..a8bb05364 --- /dev/null +++ b/test/llm/README.md @@ -0,0 +1,34 @@ +# Cache and chunking regression tests + +Run the CPU regressions with the project Python dependencies installed: + +```sh +python -m unittest discover -s test/llm -p 'test_*.py' +``` + +These tests load isolated Python modules without constructing a native model. +They cover cache ownership/capacity, LRU/SLRU eviction, admission rollback, +remote-KV delayed release, configuration forwarding, chunk boundaries, +phase progress, cancellation and final-only output. + +With matching built InfiniLM/InfiniCore extensions and a dense FP16 model, +run the short native regression (use `--tp 1` for a single GPU): + +```sh +CUDA_VISIBLE_DEVICES=0,1 python test/llm/check_chunk_output.py \ + --model /path/to/model --tp 2 --chunk-size 17 +``` + +The check compares chunked and ordinary greedy output, verifies native +intermediate-output suppression and invalid-input rejection, and exercises +prefix reuse, cancellation and complete page-reference reclamation. Add `--graph` +to run ordinary Decode with graphs; compile InfiniCore with `--graph=y` first. +The model must accept token IDs 1–67 and meet the scheduler’s minimum +`max_position_embeddings` of 1024. + +Longer TP/PP experiments, KV poisoning and graph-launch interception are archived +in the [validation tools](https://github.com/big-hip/InfiniCore/tree/b635f35f359d2f536b9ba5ca82686b6b2b988cb7/docs/validation/cache-chunk-tools-20260921). + +Configuration and limits: [cache and chunking](../../docs/cache-and-chunking.md). +Historical benchmark scripts and measurements are linked from +[PR #573](https://github.com/InfiniTensor/InfiniLM/pull/573). diff --git a/test/llm/cache_test_support.py b/test/llm/cache_test_support.py new file mode 100644 index 000000000..96777dc8c --- /dev/null +++ b/test/llm/cache_test_support.py @@ -0,0 +1,84 @@ +import importlib.util +import sys +from pathlib import Path +from unittest.mock import patch + +# Keep native NumPy modules alive when restoring the isolated package imports. +import numpy as np # noqa: F401 + + +def load_modules(): + source = Path(__file__).resolve().parents[2] / "python/infinilm/llm" + modules = {} + with patch.dict(sys.modules): + for name in ( + "prefix_cache", + "sampling_params", + "request", + "cache_manager", + "scheduler", + ): + fullname = f"infinilm.llm.{name}" + spec = importlib.util.spec_from_file_location( + fullname, source / f"{name}.py" + ) + module = importlib.util.module_from_spec(spec) + sys.modules[fullname] = module + spec.loader.exec_module(module) + modules[name] = module + return modules + + +MODULES = load_modules() +BlockManager = MODULES["cache_manager"].BlockManager + + +def chain_hashes(tokens, block_size=16): + hashes = [] + parent = MODULES["prefix_cache"].EMPTY_BLOCK_HASH + for start in range(0, len(tokens) - block_size + 1, block_size): + parent = MODULES["prefix_cache"].hash_block_tokens( + tokens[start : start + block_size], parent + ) + hashes.append(parent) + return hashes + + +def publish(manager, tokens): + table, _ = manager.allocate_slots(len(tokens)) + hashes = chain_hashes(tokens, manager.block_size) + manager.publish_computed_blocks(table, hashes, 0, len(tokens)) + manager.free_blocks(table) + return table, hashes + + +def assert_state(case, manager): + free = list(manager.free_block_ids) + used = set(manager.used_block_ids) + case.assertEqual(len(free), len(set(free))) + case.assertFalse(set(free) & used) + case.assertEqual(set(free) | used, set(range(manager.num_blocks))) + evictable = set() + index = {} + for block in manager.blocks: + case.assertGreaterEqual(block.ref_count, 0) + if block.block_id in free: + case.assertEqual((block.ref_count, block.hash), (0, b"")) + elif block.ref_count == 0: + case.assertNotEqual(block.hash, b"") + evictable.add(block.block_id) + if block.hash: + case.assertIn(block.block_id, used) + index.setdefault(block.hash, set()).add(block.block_id) + probationary = set(manager._evictable_blocks) + protected = set(manager._protected_evictable_blocks) + case.assertFalse(probationary & protected) + case.assertEqual(probationary | protected, evictable) + case.assertTrue(protected <= set(manager._protected_blocks) <= used) + case.assertFalse(probationary & set(manager._protected_blocks)) + case.assertLessEqual(len(manager._protected_blocks), manager._protected_capacity) + for block_id in manager._protected_blocks: + case.assertTrue(manager.blocks[block_id].hash) + case.assertEqual(manager.blocks[block_id].ref_count == 0, block_id in protected) + case.assertEqual(manager.hash_to_block_ids, index) + case.assertEqual(manager.get_total_usable_blocks(), len(free) + len(evictable)) diff --git a/test/llm/check_chunk_output.py b/test/llm/check_chunk_output.py new file mode 100644 index 000000000..28139a1d3 --- /dev/null +++ b/test/llm/check_chunk_output.py @@ -0,0 +1,127 @@ +"""Opt-in native chunk/output/cancellation regression for a dense FP16 model.""" + +import argparse +from unittest.mock import patch + + +def check(model, tp=1, chunk_size=17, graph=False): + import infinicore + from infinilm.config.engine_config import EngineConfig + from infinilm.lib import _infinilm + from infinilm.llm.llm import LLMEngine + from infinilm.llm.request import InferenceRequest + from infinilm.llm.sampling_params import SamplingParams + + outputs = [] + for chunk in (0, chunk_size): + engine = LLMEngine( + EngineConfig( + model, + device="cuda", + dtype="float16", + tensor_parallel_size=tp, + enable_graph=graph, + attn_backend="paged-attn", + num_blocks=16, + block_size=64, + max_batch_size=1, + prefill_chunk_size=chunk, + enable_prefix_caching=True, + prefix_cache_policy="slru", + ) + ) + raw = engine.model_runner.model_engine + native = _infinilm.InferEngine.forward + calls = [] + + def forward(instance, inputs): + result = native(instance, inputs) + calls.append(inputs.prefill_only) + if inputs.prefill_only: + assert ( + not result.output_ids + and not result.logits + and not result.hidden_states + ) + else: + assert result.output_ids and result.logits + return result + + def generate(name, tokens, cancel=False, reused=False): + request = InferenceRequest( + name, + prompt_token_ids=tokens, + sampling_params=SamplingParams(max_tokens=4, ignore_eos=True, top_k=1), + ) + engine.add_request(request) + for step in range(100): + if request.is_finished(): + break + assert engine.step()[0] + if step == 0 and reused: + assert request.num_local_cached_tokens == 64 + if cancel and step == 0: + assert not request.generated_token_ids + request.abort() + assert request.is_finished() + cache = engine.scheduler.cache_manager + assert all(block.ref_count == 0 for block in cache.blocks) + assert cache.get_total_usable_blocks() == cache.num_blocks + if cancel: + assert ( + request.status.name == "CANCELED" + and not request.generated_token_ids + ) + else: + assert len(request.generated_token_ids) == 4 + return list(request.generated_token_ids) + + try: + caches = _infinilm.InferEngine.get_kv_cache(raw) + assert len(caches) == tp + for rank, tensors in enumerate(caches): + assert {infinicore.Tensor(t).device.index for t in tensors if t} == { + rank + } + # Invalid output suppression must fail before dispatching worker jobs. + for arguments, message in ( + ( + {"prefill_only": True, "sample_all_positions": True}, + "sample_all_positions=false", + ), + ({"prefill_only": True}, "input_offsets"), + ): + try: + native(raw, _infinilm.InferEngine.Input(**arguments)) + except ValueError as error: + assert message in str(error) + else: + raise AssertionError("Invalid outputless forward was accepted.") + with patch.object(_infinilm.InferEngine, "forward", forward): + tokens = list(range(1, 68)) + result = generate("first", tokens) + assert generate("prefix-reuse", tokens, reused=True) == result + if chunk: + assert any(calls) and not all(calls) + generate("cancel", [7] * len(tokens), cancel=True) + else: + assert not any(calls) + outputs.append(result) + finally: + engine.close() + assert outputs[0] == outputs[1], "Chunked and ordinary greedy tokens differ." + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--chunk-size", type=int, default=17) + parser.add_argument("--graph", action="store_true") + args = parser.parse_args() + if not 0 < args.chunk_size < 67: + parser.error( + "--chunk-size must be between 1 and 66 to exercise intermediate chunks" + ) + check(args.model, args.tp, args.chunk_size, args.graph) + print("Native chunk output, prefix reuse, cancellation and reclamation passed.") diff --git a/test/llm/config_test_support.py b/test/llm/config_test_support.py new file mode 100644 index 000000000..3696e29c5 --- /dev/null +++ b/test/llm/config_test_support.py @@ -0,0 +1,58 @@ +"""Load configuration entrypoints without constructing native model extensions.""" + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import patch + +from cache_test_support import MODULES + +SOURCE = Path(__file__).resolve().parents[2] / "python/infinilm" + + +def load_module(name, relative_path): + spec = importlib.util.spec_from_file_location(name, SOURCE / relative_path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def load_config_modules(): + with patch.dict(sys.modules): + # Keep package __init__ files from importing the native model extension. + for name in ("infinilm", "infinilm.config", "infinilm.llm"): + sys.modules[name] = ModuleType(name) + for name, module in MODULES.items(): + sys.modules[f"infinilm.llm.{name}"] = module + kv = load_module("infinilm.config.kv_transfer", "config/kv_transfer.py") + sys.modules["infinilm.config"].KVTransferConfig = kv.KVTransferConfig + engine_config = load_module( + "infinilm.config.engine_config", "config/engine_config.py" + ) + load_module("infinilm.moe_config", "moe_config.py") + base = load_module("infinilm.base_config", "base_config.py") + load_module("infinilm.llm.static_scheduler", "llm/static_scheduler.py") + sys.modules["infinilm.infer_engine"] = SimpleNamespace( + read_hf_config=lambda path: {}, model_uses_mamba_cache=lambda config: False + ) + sys.modules["infinilm.kv_connector"] = SimpleNamespace( + KVConnectorFactory=object, KVConnectorRole=object + ) + sys.modules["infinilm.llm.model_runner.model_runner"] = SimpleNamespace( + ModelRunner=object + ) + sys.modules["infinilm.multimodal.multimodal"] = SimpleNamespace( + resolve_multimodal_inputs=object + ) + llm = load_module("infinilm.llm.llm", "llm/llm.py") + for name in ("AsyncLLMEngine", "FinishReason", "SamplingParams"): + setattr(sys.modules["infinilm.llm"], name, getattr(llm, name)) + server = load_module( + "infinilm.server.inference_server", "server/inference_server.py" + ) + return engine_config.EngineConfig, base.BaseConfig, llm, server + + +EngineConfig, BaseConfig, LLM_MODULE, SERVER_MODULE = load_config_modules() diff --git a/test/llm/test_cache_manager.py b/test/llm/test_cache_manager.py new file mode 100644 index 000000000..bb497c7cf --- /dev/null +++ b/test/llm/test_cache_manager.py @@ -0,0 +1,255 @@ +import random +import unittest + +from cache_test_support import ( + BlockManager, + assert_state, + chain_hashes, + publish, +) + + +class CachePolicyTests(unittest.TestCase): + def test_recently_reused_prefix_survives_pressure(self): + manager = BlockManager(3, 16) + _, a = publish(manager, [11] * 16) + _, b = publish(manager, [22] * 16) + _, c = publish(manager, [33] * 16) + touched, hit = manager.get_computed_blocks(a, 16) + self.assertEqual(hit, 16) + manager.free_blocks(touched) + self.assertIsNotNone(manager.allocate_slots(16)) + self.assertIn(a[0], manager.hash_to_block_ids) + self.assertNotIn(b[0], manager.hash_to_block_ids) + self.assertIn(c[0], manager.hash_to_block_ids) + assert_state(self, manager) + + def test_request_tail_is_evicted_before_its_prefix(self): + manager = BlockManager(2, 16) + _, hashes = publish(manager, [11] * 16 + [22] * 16) + self.assertIsNotNone(manager.allocate_slots(16)) + pinned, hit = manager.get_computed_blocks(hashes, 32) + self.assertEqual(hit, 16) + self.assertNotIn(hashes[1], manager.hash_to_block_ids) + manager.free_blocks(pinned) + assert_state(self, manager) + + +class CacheStateTests(unittest.TestCase): + def test_shared_pin_survives_until_last_release(self): + manager = BlockManager(2, 16) + _, hashes = publish(manager, [11] * 16) + first, _ = manager.get_computed_blocks(hashes, 16) + second, _ = manager.get_computed_blocks(hashes, 16) + pressure, _ = manager.allocate_slots(16) + + manager.free_blocks(first) + self.assertFalse(manager.try_free_blocks(1)) + self.assertEqual(manager.blocks[second[0]].ref_count, 1) + + manager.free_blocks(second) + self.assertTrue(manager.try_free_blocks(1)) + self.assertNotIn(hashes[0], manager.hash_to_block_ids) + manager.free_blocks(pressure) + assert_state(self, manager) + + def test_duplicate_hash_keeps_other_block_indexed(self): + manager = BlockManager(2, 16) + first, _ = manager.allocate_slots(16) + second, _ = manager.allocate_slots(16) + block_hash = chain_hashes([7] * 16) + manager.publish_computed_blocks(first, block_hash, 0, 16) + manager.publish_computed_blocks(second, block_hash, 0, 16) + + manager.free_blocks(first) + self.assertTrue(manager.try_free_blocks(1)) + self.assertEqual(manager.hash_to_block_ids[block_hash[0]], {second[0]}) + manager.free_blocks(second) + assert_state(self, manager) + + def test_unpublished_release_returns_free_capacity(self): + manager = BlockManager(2, 16) + table, _ = manager.allocate_slots(17) + manager.free_blocks(table) + + self.assertEqual(manager.get_num_free_blocks(), 2) + self.assertEqual(manager.used_block_ids, set()) + self.assertEqual(manager.hash_to_block_ids, {}) + self.assertEqual(manager._evictable_blocks, {}) + assert_state(self, manager) + + def test_partial_publish_retains_only_full_block(self): + manager = BlockManager(2, 16) + table, _ = manager.allocate_slots(17) + hashes = chain_hashes([11] * 16 + [22] * 16) + manager.publish_computed_blocks(table, hashes, 0, 17) + manager.free_blocks(table) + + self.assertEqual(manager.get_num_free_blocks(), 1) + missed, missed_tokens = manager.get_computed_blocks(hashes, 15) + self.assertEqual((missed, missed_tokens), ([], 0)) + pinned, hit_tokens = manager.get_computed_blocks(hashes, 17) + self.assertEqual(hit_tokens, 16) + manager.free_blocks(pinned) + assert_state(self, manager) + + def test_insufficient_capacity_preserves_pins(self): + manager = BlockManager(2, 16) + _, hashes = publish(manager, [11] * 16) + private, _ = manager.allocate_slots(16) + private_id = private[0] + + self.assertFalse(manager.try_free_blocks(2)) + self.assertNotIn(hashes[0], manager.hash_to_block_ids) + self.assertEqual(manager.blocks[private_id].ref_count, 1) + self.assertEqual(manager.blocks[private_id].block_id, private_id) + self.assertEqual(manager.get_num_free_blocks(), 1) + manager.free_blocks(private) + assert_state(self, manager) + + def test_sufficient_capacity_does_not_evict(self): + manager = BlockManager(2, 16) + _, hashes = publish(manager, [11] * 16) + + self.assertTrue(manager.try_free_blocks(1)) + self.assertTrue(manager.try_free_blocks(0)) + self.assertIn(hashes[0], manager.hash_to_block_ids) + assert_state(self, manager) + + def test_speculative_truncate_releases_private_tail(self): + manager = BlockManager(3, 16) + table, _ = manager.allocate_slots(17) + table, _ = manager.append_slots(table, 18, 16) + self.assertEqual(len(table), 3) + + retained = manager.truncate_blocks(table, 17) + self.assertEqual(len(retained), 2) + self.assertEqual(manager.get_num_free_blocks(), 1) + self.assertEqual(manager._evictable_blocks, {}) + manager.free_blocks(retained) + assert_state(self, manager) + + def test_invalid_truncate_is_atomic(self): + for invalid_kind in ( + "discarded_published", + "discarded_shared", + "retained_published", + "retained_shared", + ): + with self.subTest(invalid_kind=invalid_kind): + manager = BlockManager(4, 16) + table, _ = manager.allocate_slots(48) + hashes = chain_hashes([11] * 16 + [22] * 16 + [33] * 16) + second_owner = [] + if invalid_kind == "discarded_published": + manager.publish_computed_blocks(table, hashes, 0, 48) + keep_tokens = 16 + elif invalid_kind == "discarded_shared": + second_owner = [table[-1]] + manager.blocks[second_owner[0]].ref_count += 1 + keep_tokens = 16 + elif invalid_kind == "retained_published": + manager.publish_computed_blocks(table, hashes, 0, 32) + keep_tokens = 17 + else: + second_owner = [table[1]] + manager.blocks[second_owner[0]].ref_count += 1 + keep_tokens = 17 + before = self._snapshot(manager, table) + + with self.assertRaises(RuntimeError): + manager.truncate_blocks(table, keep_tokens) + + self.assertEqual(self._snapshot(manager, table), before) + manager.free_blocks(table) + if second_owner: + manager.free_blocks(second_owner) + assert_state(self, manager) + + def test_append_slot_uses_lru_at_boundary(self): + manager = BlockManager(2, 16) + old_table, old_hashes = publish(manager, [11] * 16) + table, _ = manager.allocate_slots(16) + + table, slot = manager.append_slot(table, 17) + self.assertEqual(slot, old_table[0] * 16) + self.assertNotIn(old_hashes[0], manager.hash_to_block_ids) + self.assertEqual(len(table), 2) + manager.free_blocks(table) + assert_state(self, manager) + + def test_bounded_legal_operation_sequences_preserve_invariants(self): + for seed in range(20): + rng = random.Random(seed) + manager = BlockManager(8, 16) + owners = [] + published_sequences = [] + next_token = 1 + history = [] + for step in range(500): + operation = rng.choice( + ("allocate", "publish", "pin", "release", "free") + ) + try: + if operation == "allocate": + blocks = rng.randint(1, 3) + allocation = manager.allocate_slots(blocks * 16) + if allocation is not None: + owners.append(allocation[0]) + elif operation == "publish": + candidates = [ + table + for table in owners + if all( + manager.blocks[block_id].ref_count == 1 + and not manager.blocks[block_id].hash + for block_id in table + ) + ] + if candidates: + table = rng.choice(candidates) + tokens = list( + range(next_token, next_token + len(table) * 16) + ) + next_token += len(tokens) + hashes = chain_hashes(tokens) + manager.publish_computed_blocks( + table, hashes, 0, len(tokens) + ) + published_sequences.append(hashes) + elif operation == "pin" and published_sequences: + hashes = rng.choice(published_sequences) + table, _ = manager.get_computed_blocks(hashes, len(hashes) * 16) + if table: + owners.append(table) + elif operation == "release" and owners: + owner = owners.pop(rng.randrange(len(owners))) + manager.free_blocks(owner) + elif operation == "free": + manager.try_free_blocks(rng.randint(0, 9)) + history.append(operation) + assert_state(self, manager) + except Exception as error: + self.fail( + f"seed={seed} step={step} operation={operation} " + f"history={history[-30:]} error={error!r}" + ) + for owner in owners: + manager.free_blocks(owner) + self.assertEqual(manager.get_total_usable_blocks(), 8, f"seed={seed}") + assert_state(self, manager) + + @staticmethod + def _snapshot(manager, table): + return ( + list(table), + list(manager.free_block_ids), + set(manager.used_block_ids), + list(manager._evictable_blocks), + {key: set(value) for key, value in manager.hash_to_block_ids.items()}, + [(block.ref_count, block.hash) for block in manager.blocks], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/llm/test_cache_manager_slru.py b/test/llm/test_cache_manager_slru.py new file mode 100644 index 000000000..c419b2b77 --- /dev/null +++ b/test/llm/test_cache_manager_slru.py @@ -0,0 +1,149 @@ +import random +import unittest + +from cache_test_support import BlockManager, assert_state, publish + + +class SegmentedCacheTests(unittest.TestCase): + def make_manager(self, blocks=4, ratio=0.5): + return BlockManager( + blocks, + 16, + prefix_cache_policy="slru", + prefix_cache_protected_ratio=ratio, + ) + + def reuse(self, manager, hashes): + table, _ = manager.get_computed_blocks(hashes, len(hashes) * 16) + manager.record_cache_hit(table) + manager.free_blocks(table) + + def test_admitted_hot_prefix_survives_one_off_scan(self): + for policy, retained in (("lru", False), ("slru", True)): + with self.subTest(policy=policy): + manager = BlockManager(4, 16, prefix_cache_policy=policy) + _, hot = publish(manager, [11] * 16) + self.reuse(manager, hot) + for token in range(20, 30): + publish(manager, [token] * 16) + assert_state(self, manager) + self.assertEqual(hot[0] in manager.hash_to_block_ids, retained) + + def test_probe_does_not_promote(self): + manager = self.make_manager() + table, hashes = publish(manager, [11] * 16) + probe, _ = manager.get_computed_blocks(hashes, 16) + manager.free_blocks(probe) + self.assertNotIn(table[0], manager._protected_blocks) + for token in range(20, 25): + publish(manager, [token] * 16) + self.assertNotIn(hashes[0], manager.hash_to_block_ids) + + def test_protected_capacity_demotes_old_hotspot(self): + manager = self.make_manager(ratio=0.25) + old_table, old = publish(manager, [11] * 16) + self.reuse(manager, old) + new_table, new = publish(manager, [22] * 16) + self.reuse(manager, new) + self.assertEqual(list(manager._protected_blocks), new_table) + self.assertIn(old_table[0], manager._evictable_blocks) + for token in range(30, 35): + publish(manager, [token] * 16) + self.assertNotIn(old[0], manager.hash_to_block_ids) + self.assertIn(new[0], manager.hash_to_block_ids) + assert_state(self, manager) + + def test_pinned_demotion_never_releases_shared_owner(self): + manager = self.make_manager(ratio=0.25) + table, hashes = publish(manager, [11] * 16) + first, _ = manager.get_computed_blocks(hashes, 16) + second, _ = manager.get_computed_blocks(hashes, 16) + manager.record_cache_hit(first) + _, other = publish(manager, [22] * 16) + self.reuse(manager, other) + self.assertNotIn(table[0], manager._protected_blocks) + self.assertEqual(manager.blocks[table[0]].ref_count, 2) + self.assertFalse(manager.try_free_blocks(4)) + manager.free_blocks(first) + self.assertFalse(manager.try_free_blocks(4)) + manager.free_blocks(second) + self.assertTrue(manager.try_free_blocks(4)) + assert_state(self, manager) + + def test_tail_loses_protection_before_prefix_when_cap_is_small(self): + manager = self.make_manager(ratio=0.25) + table, hashes = publish(manager, [11] * 16 + [22] * 16) + self.reuse(manager, hashes) + self.assertEqual(list(manager._protected_blocks), table[:1]) + self.assertTrue(manager.try_free_blocks(3)) + self.assertIn(hashes[0], manager.hash_to_block_ids) + self.assertNotIn(hashes[1], manager.hash_to_block_ids) + + def test_protected_tail_evicted_before_prefix_when_no_probation_remains(self): + manager = self.make_manager(ratio=0.75) + _, hashes = publish(manager, [11] * 16 + [22] * 16) + self.reuse(manager, hashes) + self.assertTrue(manager.try_free_blocks(3)) + self.assertIn(hashes[0], manager.hash_to_block_ids) + self.assertNotIn(hashes[1], manager.hash_to_block_ids) + + def test_eviction_clears_protection_before_block_id_reuse(self): + manager = self.make_manager() + _, hot = publish(manager, [11] * 16) + self.reuse(manager, hot) + self.assertTrue(manager.try_free_blocks(4)) + self.assertFalse(manager._protected_blocks) + self.assertFalse(manager._protected_evictable_blocks) + for token in range(20, 24): + publish(manager, [token] * 16) + self.assertFalse(manager._protected_blocks) + assert_state(self, manager) + + def test_one_block_pool_has_no_permanent_protection(self): + manager = self.make_manager(blocks=1) + _, hashes = publish(manager, [11] * 16) + self.reuse(manager, hashes) + self.assertFalse(manager._protected_blocks) + self.assertTrue(manager.try_free_blocks(1)) + assert_state(self, manager) + + def test_invalid_policy_configuration(self): + for policy, ratio in ( + ("unknown", 0.5), + ("slru", 0), + ("slru", 1), + ("slru", float("nan")), + ): + with self.subTest(policy=policy, ratio=ratio): + with self.assertRaises(ValueError): + BlockManager( + 4, + 16, + prefix_cache_policy=policy, + prefix_cache_protected_ratio=ratio, + ) + + def test_randomized_shared_lifetimes_preserve_capacity(self): + rng = random.Random(71) + manager = self.make_manager(blocks=12) + held = [] + for _ in range(400): + action = rng.randrange(4) + if action == 0 and manager.hash_to_block_ids: + key = rng.choice(list(manager.hash_to_block_ids)) + table, _ = manager.get_computed_blocks([key], 16) + if rng.choice((True, False)): + manager.record_cache_hit(table) + held.append(table) + elif action == 1 and held: + manager.free_blocks(held.pop(rng.randrange(len(held)))) + elif action == 2 and manager.get_total_usable_blocks(): + publish(manager, [rng.randrange(10, 1000)] * 16) + else: + manager.try_free_blocks(rng.randrange(1, 13)) + assert_state(self, manager) + self.assertLessEqual(len(manager._protected_blocks), 6) + for table in held: + manager.free_blocks(table) + self.assertTrue(manager.try_free_blocks(12)) + assert_state(self, manager) diff --git a/test/llm/test_chunk_execution.py b/test/llm/test_chunk_execution.py new file mode 100644 index 000000000..4b36b19b2 --- /dev/null +++ b/test/llm/test_chunk_execution.py @@ -0,0 +1,316 @@ +import sys +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from cache_test_support import MODULES +from config_test_support import LLM_MODULE, EngineConfig, load_module + + +def load_processor(): + with patch.dict(sys.modules): + for name, module in MODULES.items(): + sys.modules[f"infinilm.llm.{name}"] = module + sys.modules["transformers"] = SimpleNamespace(AutoTokenizer=object) + static = load_module("infinilm.llm.static_scheduler", "llm/static_scheduler.py") + load_module("infinilm.processors.processor", "processors/processor.py") + return ( + load_module( + "infinilm.processors.basic_llm_processor", + "processors/basic_llm_processor.py", + ).BasicLLMProcessor, + static, + ) + + +Processor, STATIC_MODULE = load_processor() + + +class ChunkExecutionTests(unittest.TestCase): + def test_legacy_static_output_does_not_require_chunk_metadata(self): + engine, req = self.setup_engine(length=1) + engine.scheduler = STATIC_MODULE.StaticScheduler() + self.addCleanup(engine.scheduler.waiting_queue.close) + engine.scheduler.add_request(req) + engine.step() + self.assertEqual(list(req.generated_token_ids), [77]) + + def test_static_processor_still_uses_its_prefix_metadata(self): + _, req = self.setup_engine(length=17) + output = STATIC_MODULE.StaticSchedulerOutput( + [req], is_prefill=True, prefix_hit_len=16 + ) + processor = Processor.__new__(Processor) + backend = SimpleNamespace( + from_list=lambda values, **kwargs: values, int64="int64", int32="int32" + ) + with patch.dict(sys.modules, {"infinicore": backend}): + inputs = processor.build_model_inputs(output) + self.assertEqual([list(row) for row in inputs["input_ids"]], [[26]]) + self.assertEqual(inputs["position_ids"], [[16]]) + self.assertEqual(inputs["past_kv_lengths"], [16]) + self.assertEqual(inputs["total_kv_lengths"], [17]) + self.assertIsNone(inputs["slot_mapping"]) + + def test_unsupported_models_rejected_before_native_initialization(self): + from config_test_support import EngineConfig + + for hf in ({"num_experts": 8}, {"vision_config": {}}, {"audio_config": {}}): + with ( + self.subTest(hf=hf), + patch.object(LLM_MODULE, "read_hf_config", return_value=hf), + ): + with self.assertRaisesRegex(ValueError, "dense text"): + LLM_MODULE.LLMEngine(EngineConfig("unused", prefill_chunk_size=16)) + with patch.object(LLM_MODULE, "model_uses_mamba_cache", return_value=True): + with self.assertRaisesRegex(ValueError, "dense text"): + LLM_MODULE.LLMEngine(EngineConfig("unused", prefill_chunk_size=16)) + + def setup_engine(self, length=35, prefix=True): + scheduler = MODULES["scheduler"].Scheduler( + num_blocks=16, + block_size=16, + prefill_chunk_size=16, + enable_prefix_caching=prefix, + ) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + req = MODULES["request"].InferenceRequest( + "long", + prompt_token_ids=list(range(10, 10 + length)), + sampling_params=MODULES["sampling_params"].SamplingParams( + max_tokens=1, ignore_eos=True + ), + ) + scheduler.add_request(req) + engine = LLM_MODULE.LLMEngine.__new__(LLM_MODULE.LLMEngine) + engine.scheduler = scheduler + engine.tokenizer = SimpleNamespace(decode=lambda tokens: "output") + engine.model_runner = SimpleNamespace( + execute_model=lambda output: SimpleNamespace( + sampled_token_ids=[77] * len(output.scheduled_requests), + kv_connector_output=None, + ) + ) + return engine, req + + def test_intermediate_steps_publish_only_computed_pages_and_emit_no_tokens(self): + engine, req = self.setup_engine() + for end, pages in ((16, 1), (32, 2)): + worked, pending = engine.step() + self.assertTrue(worked) + self.assertEqual(pending, []) + self.assertEqual(req.get_num_generated_tokens(), 0) + self.assertEqual(req.num_computed_tokens, end) + self.assertEqual(req.num_cache_indexed_blocks, pages) + self.assertEqual( + len(engine.scheduler.cache_manager.hash_to_block_ids), pages + ) + engine.step() + self.assertEqual(list(req.generated_token_ids), [77]) + self.assertEqual(req.status, MODULES["request"].RequestStatus.FINISHED) + self.assertEqual(engine.scheduler.cache_manager.get_total_usable_blocks(), 16) + + def test_abort_during_intermediate_forward_releases_ownership(self): + engine, req = self.setup_engine() + + def execute(output): + req._aborted = True + return SimpleNamespace(sampled_token_ids=[77], kv_connector_output=None) + + engine.model_runner.execute_model = execute + engine.step() + self.assertEqual(req.status, MODULES["request"].RequestStatus.CANCELED) + self.assertFalse(engine.scheduler.chunking_queue) + self.assertEqual(engine.scheduler.running_queue.sync_q.qsize(), 0) + self.assertEqual(req.get_num_generated_tokens(), 0) + self.assertEqual(req.num_cache_indexed_blocks, 1) + self.assertEqual(engine.scheduler.cache_manager.get_total_usable_blocks(), 16) + + def test_other_request_completion_never_publishes_future_long_pages(self): + engine, req = self.setup_engine(length=64) + engine.step() + short = MODULES["request"].InferenceRequest( + "short", + prompt_token_ids=[8], + sampling_params=MODULES["sampling_params"].SamplingParams( + max_tokens=1, ignore_eos=True + ), + ) + engine.scheduler.add_request(short) + engine.step() + engine.step() + self.assertEqual(short.get_num_generated_tokens(), 1) + self.assertEqual(req.num_computed_tokens, 32) + self.assertNotIn( + req.block_hashes[2], engine.scheduler.cache_manager.hash_to_block_ids + ) + self.assertNotIn( + req.block_hashes[3], engine.scheduler.cache_manager.hash_to_block_ids + ) + + def test_disabled_prefix_cache_still_completes_in_three_steps(self): + engine, req = self.setup_engine(prefix=False) + for _ in range(3): + engine.step() + self.assertEqual(list(req.generated_token_ids), [77]) + self.assertFalse(engine.scheduler.cache_manager.hash_to_block_ids) + + def test_abort_during_final_chunk_emits_nothing_and_releases(self): + engine, req = self.setup_engine(length=17) + engine.step() + + def execute(output): + req.abort() + return SimpleNamespace(sampled_token_ids=[77], kv_connector_output=None) + + engine.model_runner.execute_model = execute + self.assertEqual(engine.step(), (True, [])) + self.assertEqual(req.get_num_generated_tokens(), 0) + self.assertEqual(req.status, MODULES["request"].RequestStatus.CANCELED) + self.assertTrue( + all(b.ref_count == 0 for b in engine.scheduler.cache_manager.blocks) + ) + + def test_intermediate_eos_is_ignored_and_final_eos_finishes(self): + engine, req = self.setup_engine(length=17) + req.sampling_params.ignore_eos = False + req.sampling_params.max_tokens = 8 + engine.eos_token_ids = [77] + engine.step() + self.assertFalse(req.is_finished()) + self.assertEqual(req.get_num_generated_tokens(), 0) + engine.step() + self.assertTrue(req.is_finished()) + self.assertEqual(list(req.generated_token_ids), [77]) + self.assertEqual(req.finish_reason, MODULES["request"].FinishReason.EOS_TOKEN) + + def test_decode_extends_hashes_and_only_publishes_computed_tokens(self): + engine, req = self.setup_engine(length=31) + req.sampling_params.max_tokens = 3 + engine.step() + engine.step() + self.assertEqual(len(req.block_hashes), 2) + self.assertEqual(req.num_cache_indexed_blocks, 1) + engine.step() + self.assertEqual(req.num_cache_indexed_blocks, 2) + engine.step() + self.assertTrue(req.is_finished()) + self.assertTrue( + all(b.ref_count == 0 for b in engine.scheduler.cache_manager.blocks) + ) + + def test_processor_slices_partial_prefill_positions_and_lengths(self): + engine, req = self.setup_engine(length=35) + engine.step() + step = engine.scheduler.schedule() + processor = Processor.__new__(Processor) + backend = SimpleNamespace( + from_list=lambda values, **kwargs: values, int64="int64", int32="int32" + ) + with patch.dict(sys.modules, {"infinicore": backend}): + inputs = processor.build_model_inputs(step) + self.assertEqual(inputs["input_ids"], [list(range(26, 42))]) + self.assertEqual(inputs["position_ids"], list(range(16, 32))) + self.assertEqual(inputs["past_kv_lengths"], [16]) + self.assertEqual(inputs["total_kv_lengths"], [32]) + self.assertEqual(inputs["input_offsets"], [0, 16]) + self.assertEqual(len(inputs["slot_mapping"]), 16) + + +def load_runner(): + replacements = { + "infinicore": SimpleNamespace(), + "infinilm.cache.cache": SimpleNamespace( + PagedKVCacheConfig=object, StaticKVCacheConfig=object + ), + "infinilm.config.engine_config": SimpleNamespace(EngineConfig=EngineConfig), + "infinilm.distributed": SimpleNamespace(DistConfig=object), + "infinilm.distributed.pipeline_transport": SimpleNamespace( + PipelineControlServer=object + ), + "infinilm.infer_engine": SimpleNamespace(InferEngine=object), + "infinilm.kv_connector": SimpleNamespace( + KVConnectorFactory=object, KVConnectorRole=object + ), + "infinilm.llm.model_runner.speculative_runner": SimpleNamespace( + SpeculativeRunner=object + ), + "infinilm.modeling_utils": SimpleNamespace( + load_model_state_dict_by_file=object + ), + "infinilm.processors": SimpleNamespace(AutoInfinilmProcessor=object), + } + with patch.dict(sys.modules, replacements): + return load_module( + "infinilm.llm.model_runner.model_runner", "llm/model_runner/model_runner.py" + ).ModelRunner + + +Runner = load_runner() + + +class ChunkOutputTests(unittest.TestCase): + def setup_runner(self): + engine, req = ChunkExecutionTests.setup_engine(self) + runner = Runner.__new__(Runner) + runner.config = EngineConfig("unused", prefill_chunk_size=16) + runner.processor = SimpleNamespace(build_model_inputs=lambda *a: {}) + runner.speculative_runner = None + runner.pipeline_control = None + runner.kv_connector = None + calls = [] + + def forward(**kwargs): + calls.append(kwargs.get("prefill_only", False)) + if kwargs.get("prefill_only"): + return None + return SimpleNamespace( + to_numpy=lambda: SimpleNamespace(tolist=lambda: [77]) + ) + + runner.model_engine = SimpleNamespace(forward=forward) + engine.model_runner = runner + return engine, req, calls + + def test_only_intermediate_chunks_skip_native_output(self): + engine, req, calls = self.setup_runner() + req.sampling_params.max_tokens = 2 + for _ in range(2): + self.assertEqual(engine.step(), (True, [])) + self.assertEqual(list(req.generated_token_ids), []) + self.assertEqual(calls, [True, True]) + engine.step() + self.assertEqual(calls, [True, True, False]) + self.assertEqual(list(req.generated_token_ids), [77]) + engine.step() + self.assertEqual(calls, [True, True, False, False]) + self.assertEqual(list(req.generated_token_ids), [77, 77]) + self.assertTrue( + all(b.ref_count == 0 for b in engine.scheduler.cache_manager.blocks) + ) + + def test_legacy_output_without_chunk_metadata_keeps_sampling(self): + engine, req, calls = self.setup_runner() + output = SimpleNamespace( + scheduled_requests=[req], num_requests=1, is_prefill=True + ) + result = engine.model_runner.execute_model(output) + self.assertEqual(calls, [False]) + self.assertEqual(result.sampled_token_ids, [77]) + + def test_empty_native_output_still_finishes_cancelled_chunk(self): + engine, req, calls = self.setup_runner() + forward = engine.model_runner.model_engine.forward + + def abort(**kwargs): + req.abort() + return forward(**kwargs) + + engine.model_runner.model_engine.forward = abort + self.assertEqual(engine.step(), (True, [])) + self.assertEqual(calls, [True]) + self.assertEqual(list(req.generated_token_ids), []) + self.assertTrue( + all(b.ref_count == 0 for b in engine.scheduler.cache_manager.blocks) + ) diff --git a/test/llm/test_chunk_scheduler.py b/test/llm/test_chunk_scheduler.py new file mode 100644 index 000000000..50280cafe --- /dev/null +++ b/test/llm/test_chunk_scheduler.py @@ -0,0 +1,261 @@ +import unittest + +from cache_test_support import MODULES + +Scheduler = MODULES["scheduler"].Scheduler +Request = MODULES["request"].InferenceRequest +Status = MODULES["request"].RequestStatus +Sampling = MODULES["sampling_params"].SamplingParams + + +class ChunkSchedulerTests(unittest.TestCase): + def scheduler(self, **kwargs): + config = dict(num_blocks=64, block_size=16, prefill_chunk_size=16) + config.update(kwargs) + scheduler = Scheduler(**config) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + return scheduler + + @staticmethod + def request(name, length, output=4): + return Request( + name, + prompt_token_ids=[11] * length, + sampling_params=Sampling(max_tokens=output), + ) + + def finish_step(self, scheduler, output): + for req in output.scheduled_requests: + if output.prefill_end is not None: + req.num_computed_tokens = output.prefill_end + scheduler.commit_computed_tokens(req, output.prefill_end) + if output.prefill_end < req.prompt_length: + scheduler.requeue_prefill(req) + continue + req.append_generated_token_id(12) + scheduler.complete_requests([req]) + + def test_slru_chunk_admission_promotes_hits_only_after_allocation(self): + from cache_test_support import assert_state, publish + + for rejected in (False, True): + with self.subTest(rejected=rejected): + scheduler = self.scheduler(num_blocks=4, prefix_cache_policy="slru") + manager = scheduler.cache_manager + table, _ = publish(manager, [11] * 16) + req = self.request("hit", 1000 if rejected else 33, output=1) + scheduler.add_request(req) + step = scheduler.schedule() + if rejected: + self.assertIsNone(step) + self.assertFalse(manager._protected_blocks) + self.assertEqual(manager.blocks[table[0]].ref_count, 0) + else: + self.assertEqual(step.scheduled_requests, [req]) + self.assertIn(table[0], manager._protected_blocks) + req.status = Status.CANCELED + scheduler.complete_requests([req]) + assert_state(self, manager) + + def test_chunk_boundaries_and_last_partial_segment(self): + scheduler = self.scheduler() + req = self.request("long", 35) + scheduler.add_request(req) + for start, end in ((0, 16), (16, 32), (32, 35)): + step = scheduler.schedule() + self.assertTrue(step.is_prefill) + self.assertEqual(step.prefill_end, end) + self.assertEqual(req.num_local_cached_tokens, start) + self.assertEqual(len(req.slot_mapping), end - start) + self.assertEqual(req.get_num_generated_tokens(), 0) + self.finish_step(scheduler, step) + self.assertEqual(req.get_num_generated_tokens(), 1) + self.assertFalse(scheduler.schedule().is_prefill) + + def test_chunk_respects_smaller_token_budget(self): + scheduler = self.scheduler(max_num_batched_tokens=7) + req = self.request("limited", 20) + scheduler.add_request(req) + self.assertEqual(scheduler.schedule().prefill_end, 7) + + def test_disabled_mode_retains_whole_prompt_dispatch(self): + scheduler = self.scheduler(prefill_chunk_size=0, max_num_batched_tokens=16) + req = self.request("legacy", 35) + scheduler.add_request(req) + step = scheduler.schedule() + self.assertTrue(step.is_prefill) + self.assertIsNone(step.prefill_end) + self.assertEqual(len(req.slot_mapping), 35) + self.assertFalse(scheduler.chunking_queue) + + def test_decode_page_boundary_and_cancellation_match_without_chunking(self): + from cache_test_support import assert_state + + slots = [] + for chunk_size in (0, 16): + with self.subTest(chunk_size=chunk_size): + scheduler = self.scheduler( + num_blocks=8, max_batch_size=2, prefill_chunk_size=chunk_size + ) + requests = [self.request(name, 16) for name in ("cancel", "a", "b")] + for req in requests: + req.block_table, _ = scheduler.cache_manager.allocate_slots(16) + req.num_blocks = 1 + req.num_computed_tokens = 16 + req.status = Status.RUNNING + req.append_generated_token_id(12) + scheduler.running_queue.sync_q.put(req) + requests[0].mark_canceled() + + step = scheduler.schedule() + + self.assertFalse(step.is_prefill) + self.assertIsNone(step.prefill_end) + self.assertEqual(step.scheduled_requests, requests[1:]) + for req in step.scheduled_requests: + self.assertEqual(req.num_local_cached_tokens, 16) + self.assertEqual(req.num_blocks, 2) + self.assertEqual(req.slot_mapping, [req.block_table[1] * 16]) + slots.append([req.slot_mapping for req in step.scheduled_requests]) + self.assertEqual(scheduler.cache_manager.get_total_usable_blocks(), 4) + for req in step.scheduled_requests: + req.mark_canceled() + scheduler.complete_requests(step.scheduled_requests) + self.assertEqual(scheduler.cache_manager.get_total_usable_blocks(), 8) + assert_state(self, scheduler.cache_manager) + self.assertEqual(slots[0], slots[1]) + + def test_nonaligned_chunks_reconstruct_physical_slots(self): + scheduler = self.scheduler(prefill_chunk_size=11) + req = self.request("unaligned", 35) + scheduler.add_request(req) + for start, end in ((0, 11), (11, 22), (22, 33), (33, 35)): + step = scheduler.schedule() + expected = [ + req.block_table[i // 16] * 16 + i % 16 for i in range(start, end) + ] + self.assertEqual(req.slot_mapping, expected) + self.finish_step(scheduler, step) + + def test_published_partial_prefix_starts_at_hit_boundary(self): + scheduler = self.scheduler() + first = self.request("first", 49) + scheduler.add_request(first) + self.finish_step(scheduler, scheduler.schedule()) + first.status = Status.CANCELED + second = self.request("second", 49) + scheduler.add_request(second) + seen = False + for _ in range(4): + step = scheduler.schedule() + if step.scheduled_requests == [second]: + self.assertEqual(second.num_local_cached_tokens, 16) + self.assertEqual(step.prefill_end, 32) + seen = True + break + self.finish_step(scheduler, step) + self.assertTrue(seen) + + def test_decode_continuation_and_admission_each_get_dispatch_opportunities(self): + scheduler = self.scheduler(max_batch_size=1) + active = self.request("active", 1) + scheduler.add_request(active) + self.finish_step(scheduler, scheduler.schedule()) + long = self.request("long", 200) + scheduler.add_request(long) + kinds = [] + seen = {"active"} + for i in range(12): + scheduler.add_request(self.request(f"new-{i}", 100)) + step = scheduler.schedule() + req = step.scheduled_requests[0] + kinds.append( + "decode" + if not step.is_prefill + else "continue" + if req.request_id in seen + else "admit" + ) + seen.add(req.request_id) + self.finish_step(scheduler, step) + for start in (3, 6, 9): + self.assertEqual( + set(kinds[start : start + 3]), {"decode", "continue", "admit"} + ) + + def test_cancelled_middle_chunk_is_released_without_requeue(self): + scheduler = self.scheduler() + req = self.request("canceled", 64) + scheduler.add_request(req) + self.finish_step(scheduler, scheduler.schedule()) + req.status = Status.CANCELED + self.assertIsNone(scheduler.schedule()) + self.assertFalse(scheduler.chunking_queue) + self.assertEqual(scheduler.cache_manager.get_total_usable_blocks(), 64) + + def test_prefix_disabled_still_advances(self): + scheduler = self.scheduler(enable_prefix_caching=False) + req = self.request("no-cache", 33) + scheduler.add_request(req) + for end in (16, 32, 33): + step = scheduler.schedule() + self.assertEqual(step.prefill_end, end) + self.finish_step(scheduler, step) + self.assertFalse(scheduler.cache_manager.hash_to_block_ids) + + def test_partial_requests_reserve_future_decode_capacity(self): + scheduler = self.scheduler(num_blocks=6) + long = self.request("long", 48, output=32) + scheduler.add_request(long) + self.finish_step(scheduler, scheduler.schedule()) + other = self.request("other", 32, output=16) + scheduler.add_request(other) + for _ in range(2): + step = scheduler.schedule() + self.assertNotIn(other, step.scheduled_requests) + self.finish_step(scheduler, step) + self.assertEqual(other.status, Status.WAITING) + + def test_rejected_admission_returns_temporary_prefix_reference(self): + scheduler = self.scheduler(num_blocks=5) + req = self.request("long", 48, output=16) + scheduler.add_request(req) + self.finish_step(scheduler, scheduler.schedule()) + other = self.request("too-large", 48, output=128) + scheduler.add_request(other) + before = [b.ref_count for b in scheduler.cache_manager.blocks] + step = scheduler.schedule() + self.assertNotIn(other, step.scheduled_requests) + self.assertEqual(before, [b.ref_count for b in scheduler.cache_manager.blocks]) + + def test_shared_prefix_stays_pinned_until_both_owners_release(self): + scheduler = self.scheduler() + first = self.request("first", 64) + second = self.request("second", 64) + scheduler.add_request(first) + self.finish_step(scheduler, scheduler.schedule()) + scheduler.add_request(second) + for _ in range(3): + step = scheduler.schedule() + self.finish_step(scheduler, step) + if second.status == Status.RUNNING: + break + shared = first.block_table[0] + self.assertEqual(second.block_table[0], shared) + self.assertEqual(scheduler.cache_manager.blocks[shared].ref_count, 2) + first.mark_canceled() + scheduler.schedule() + self.assertEqual(scheduler.cache_manager.blocks[shared].ref_count, 1) + # The previous dispatch may have removed the second request for execution. + second.mark_canceled() + scheduler.complete_requests([second]) + self.assertTrue(all(b.ref_count == 0 for b in scheduler.cache_manager.blocks)) + + def test_multimodal_requests_fail_before_ownership_is_acquired(self): + scheduler = self.scheduler() + req = self.request("image", 16) + req.has_multimodal_inputs = True + with self.assertRaisesRegex(ValueError, "multimodal"): + scheduler.add_request(req) + self.assertEqual(scheduler.waiting_queue.sync_q.qsize(), 0) diff --git a/test/llm/test_engine_config.py b/test/llm/test_engine_config.py new file mode 100644 index 000000000..3da00be82 --- /dev/null +++ b/test/llm/test_engine_config.py @@ -0,0 +1,323 @@ +"""Validate cache/chunk configuration once across public entrypoints.""" + +import asyncio +import io +import os +import runpy +import sys +import unittest +from contextlib import contextmanager, redirect_stderr, redirect_stdout +from types import SimpleNamespace +from unittest.mock import patch + +from config_test_support import ( + LLM_MODULE, + SERVER_MODULE, + SOURCE, + BaseConfig, + EngineConfig, + load_module, +) + +OPTIONS = dict( + prefix_cache_policy="slru", + prefix_cache_protected_ratio=0.6, + prefill_chunk_size=128, + tensor_parallel_size=2, +) +CLI = [ + "--enable-paged-attn", + "--prefix-cache-policy", + "slru", + "--prefix-cache-protected-ratio", + "0.6", + "--prefill-chunk-size", + "128", + "--tp", + "2", +] + + +class EngineConfigTests(unittest.TestCase): + @contextmanager + def cli(self, *args): + with ( + patch.object( + sys, "argv", ["server", "--model", "unused", "--device", "cpu", *args] + ), + patch.dict(os.environ), + ): + yield + + def parse_cli(self, *args): + with self.cli(*args): + return BaseConfig() + + def assert_options(self, config, expected=OPTIONS): + for key, value in expected.items(): + self.assertEqual(getattr(config, key), value) + + def test_defaults_preserve_existing_behavior(self): + for config in ( + EngineConfig("unused", cache_type="paged"), + EngineConfig("unused", cache_type="static"), + self.parse_cli(), + ): + self.assert_options( + config, + dict( + prefix_cache_policy="lru", + prefix_cache_protected_ratio=0.8, + prefill_chunk_size=0, + ), + ) + + def test_invalid_policy_ratio_and_chunk_values(self): + invalid = [("prefix_cache_policy", "fifo")] + invalid += [ + ("prefix_cache_protected_ratio", r) + for r in (-0.1, 0, 1, 1.1, float("nan"), float("inf"), -float("inf")) + ] + invalid += [ + ("prefill_chunk_size", n) for n in (-1, 1.5, "128", None, True, False) + ] + for key, value in invalid: + with ( + self.subTest(key=key, value=value), + self.assertRaisesRegex(ValueError, key), + ): + EngineConfig("unused", **{key: value}) + for enabled in (True, False): + with self.assertRaisesRegex(ValueError, "paged"): + EngineConfig( + "unused", + cache_type="static", + prefix_cache_policy="slru", + enable_prefix_caching=enabled, + ) + + def test_chunk_execution_capabilities(self): + inactive = SERVER_MODULE.KVTransferConfig() + EngineConfig("unused", prefill_chunk_size=128, kv_transfer_config=inactive) + for tp in (1, 2): + for backend in ("default", "paged-attn", "flash-attn"): + EngineConfig( + "unused", + prefill_chunk_size=128, + tensor_parallel_size=tp, + enable_graph=True, + attn_backend=backend, + ) + for stage in (0, 1): + EngineConfig( + "unused", + prefill_chunk_size=300, + pipeline_parallel_size=2, + pipeline_parallel_stage=stage, + prefix_cache_policy="slru", + ) + active = SERVER_MODULE.KVTransferConfig( + kv_connector="MooncakeConnector", kv_role="kv_producer" + ) + for overrides in ( + {"cache_type": "static"}, + {"tensor_parallel_size": 4}, + {"pipeline_parallel_size": 3}, + {"use_mla": True}, + {"draft_model_path": "draft"}, + {"kv_transfer_config": active}, + {"tensor_parallel_size": 2, "pipeline_parallel_size": 2}, + ): + with self.subTest(overrides=overrides), self.assertRaises(ValueError): + EngineConfig("unused", prefill_chunk_size=128, **overrides) + EngineConfig("unused", prefill_chunk_size=0, **overrides) + for overrides in ( + {"pipeline_parallel_size": 2}, + {"device": "cpu"}, + {"attn_backend": "unsupported"}, + ): + with ( + self.subTest(overrides=overrides), + self.assertRaisesRegex(ValueError, "requires PP=1"), + ): + EngineConfig( + "unused", prefill_chunk_size=128, enable_graph=True, **overrides + ) + + def test_cli_validation_before_dispatch(self): + config = self.parse_cli(*CLI) + self.assertEqual( + ( + config.prefix_cache_policy, + config.prefix_cache_protected_ratio, + config.prefill_chunk_size, + config.tp, + ), + ("slru", 0.6, 128, 2), + ) + invalid = [("--prefix-cache-policy", "fifo")] + invalid += [ + ("--prefix-cache-protected-ratio", v) + for v in ("0", "1", "nan", "inf", "-inf") + ] + invalid += [("--prefill-chunk-size", v) for v in ("-1", "1.5", "True")] + invalid += [ + ("--prefill-chunk-size", "128", *extra) + for extra in ( + ("--tp", "4"), + ("--pp", "3"), + ("--pp", "3", "--node-rank", "1"), + ("--draft-model", "draft"), + ) + ] + for args in invalid: + with ( + self.subTest(args=args), + redirect_stderr(io.StringIO()), + self.assertRaises(SystemExit) as error, + ): + self.parse_cli(*args) + self.assertEqual(error.exception.code, 2) + for value in ("0", "1", "128"): + self.assertEqual( + self.parse_cli("--prefill-chunk-size", value).prefill_chunk_size, + int(value), + ) + + def test_convenience_apis_forward_and_validate_options(self): + with patch.object( + LLM_MODULE, "LLMEngine", lambda config: SimpleNamespace(config=config) + ): + for constructor in (LLM_MODULE.LLM, LLM_MODULE.AsyncLLMEngine): + self.assert_options(constructor("unused", **OPTIONS).engine.config) + default = constructor("unused").engine.config + self.assertEqual( + (default.prefix_cache_policy, default.prefill_chunk_size), + ("lru", 0), + ) + with self.assertRaisesRegex(ValueError, "paged"): + constructor( + "unused", cache_type="static", prefix_cache_policy="slru" + ) + with self.assertRaisesRegex(ValueError, "prefill_chunk_size"): + constructor("unused", cache_type="static", prefill_chunk_size=128) + + def test_engine_forwards_options_to_scheduler(self): + runner = SimpleNamespace( + device="cpu", + dtype="float16", + eos_token_id=[], + processor=SimpleNamespace(get_tokenizer=lambda: None), + model_engine=SimpleNamespace(hf_config={"max_position_embeddings": 4096}), + ) + with ( + patch.object(LLM_MODULE, "ModelRunner", lambda config: runner), + patch.object(LLM_MODULE, "Scheduler") as scheduler, + ): + LLM_MODULE.LLMEngine(EngineConfig("unused", **OPTIONS)) + for key in ( + "prefix_cache_policy", + "prefix_cache_protected_ratio", + "prefill_chunk_size", + ): + self.assertEqual(scheduler.call_args.kwargs[key], OPTIONS[key]) + + def test_pipeline_worker_forwards_options(self): + config = self.parse_cli(*CLI[:-2], "--pp", "2", "--node-rank", "1") + captured, closed = [], [] + + def make_runner(config, initialize_processor): + self.assertFalse(initialize_processor) + captured.append(config) + return SimpleNamespace(close=lambda: closed.append(True)) + + replacements = { + "infinilm.base_config": SimpleNamespace(BaseConfig=BaseConfig), + "infinilm.config.engine_config": SimpleNamespace(EngineConfig=EngineConfig), + "infinilm.distributed.pipeline_transport": SimpleNamespace( + PipelineWorkerClient=lambda *a, **kw: SimpleNamespace( + serve_forever=lambda: None + ) + ), + "infinilm.llm.model_runner.model_runner": SimpleNamespace( + ModelRunner=make_runner + ), + } + with patch.dict(sys.modules, replacements): + worker = load_module( + "infinilm.server.pipeline_worker", "server/pipeline_worker.py" + ) + worker.run_worker(config) + self.assertEqual(len(captured), 1) + self.assertEqual( + ( + captured[0].prefill_chunk_size, + captured[0].prefix_cache_policy, + captured[0].prefix_cache_protected_ratio, + captured[0].pipeline_parallel_stage, + ), + (128, "slru", 0.6, 1), + ) + self.assertEqual(closed, [True]) + + def check_entrypoint(self, main): + configs = [] + + def make_engine(config): + configs.append(config) + return SimpleNamespace(config=config, close=lambda: None) + + with ( + self.cli(*CLI), + patch.object(LLM_MODULE, "LLMEngine", make_engine), + redirect_stdout(io.StringIO()), + ): + main() + self.assertEqual(len(configs), 1) + self.assert_options(configs[0]) + + def test_cli_reaches_server_lifespan(self): + async def start_without_listener(server): + app = server._create_app() + async with app.router.lifespan_context(app): + self.assert_options(server.engine.config) + + with ( + patch.object(LLM_MODULE.AsyncLLMEngine, "start"), + patch.object(LLM_MODULE.AsyncLLMEngine, "stop"), + patch.object( + SERVER_MODULE.InferenceServer, + "start", + lambda server: asyncio.run(start_without_listener(server)), + ), + patch.object(SERVER_MODULE, "setup_logging"), + ): + self.check_entrypoint(SERVER_MODULE.main) + + def test_offline_cli_reaches_llm(self): + modules = { + "infinilm.base_config": SimpleNamespace(BaseConfig=BaseConfig), + "infinilm.llm.llm": LLM_MODULE, + "infinilm.moe_config": SimpleNamespace( + configure_moe_ep_backend=SERVER_MODULE.configure_moe_ep_backend + ), + "infinilm.processors.videonsa_processor": SimpleNamespace( + decode_video_frames=object + ), + } + + with ( + patch.dict(sys.modules, modules), + patch.object(LLM_MODULE.LLM, "chat", return_value=[]), + patch.object(SERVER_MODULE.logging, "basicConfig"), + ): + self.check_entrypoint( + lambda: runpy.run_path( + str(SOURCE.parents[1] / "examples/test_infer.py"), + run_name="__main__", + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/llm/test_infer_engine_input.py b/test/llm/test_infer_engine_input.py new file mode 100644 index 000000000..252dd627d --- /dev/null +++ b/test/llm/test_infer_engine_input.py @@ -0,0 +1,113 @@ +"""Check the Python/native input boundary without constructing a model.""" + +import sys +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from config_test_support import load_module + + +class InputConversionTests(unittest.TestCase): + def setUp(self): + self.calls = [] + calls = self.calls + + class NativeEngine: + @staticmethod + def Input(input_ids, **kwargs): + return SimpleNamespace(input_ids=input_ids, **kwargs) + + def forward(self, inputs): + calls.append(inputs) + return SimpleNamespace(output_ids=77, logits=88, hidden_states=99) + + replacements = { + "infinicore": SimpleNamespace(Tensor=lambda value: value), + "infinilm.cache": SimpleNamespace(PagedKVCacheConfig=object), + "infinilm.distributed": SimpleNamespace(DistConfig=object), + "infinilm.lib": SimpleNamespace( + _infinilm=SimpleNamespace(InferEngine=NativeEngine) + ), + "infinilm.exception_utils": SimpleNamespace( + handle_oom_and_exit=lambda e: None + ), + "infinilm.modeling_utils": SimpleNamespace(parse_dtype=object), + } + with patch.dict(sys.modules, replacements): + cls = load_module("infinilm.infer_engine", "infer_engine.py").InferEngine + self.engine = cls.__new__(cls) + + def test_forward_preserves_tensor_metadata_and_output_modes(self): + tensor = SimpleNamespace(_underlying=object()) + for prefill_only in (False, True): + with self.subTest(prefill_only=prefill_only): + output = self.engine.forward( + tensor, + position_ids=tensor, + past_kv_lengths=tensor, + total_kv_lengths=tensor, + input_offsets=tensor, + cu_seqlens=tensor, + block_tables=tensor, + slot_mapping=tensor, + mamba_init_state_indices=tensor, + mamba_final_state_indices=tensor, + target_hidden_states=tensor, + pixel_values=[tensor], + image_bound=tensor, + tgt_sizes=[], + image_grid_thw=None, + image_req_ids=[0], + visual_token_ranges=[(0, 1)], + temperature=0.0, + top_k=3, + top_p=0.9, + prefill_only=prefill_only, + ) + call = self.calls[-1] + for name in ( + "input_ids", + "position_ids", + "past_sequence_lengths", + "total_sequence_lengths", + "input_offsets", + "cu_seqlens", + "block_tables", + "slot_mapping", + "mamba_init_state_indices", + "mamba_final_state_indices", + "target_hidden_states", + ): + self.assertIs(getattr(call, name), tensor._underlying) + self.assertEqual(call.pixel_values, [tensor._underlying]) + self.assertEqual(call.image_bound, [tensor._underlying]) + self.assertIsNone(call.tgt_sizes) + self.assertIsNone(call.image_grid_thw) + self.assertEqual( + (call.image_req_ids, call.visual_token_ranges), ([0], [(0, 1)]) + ) + self.assertEqual( + (call.temperature, call.top_k, call.top_p), (0.0, 3, 0.9) + ) + self.assertEqual(call.prefill_only, prefill_only) + self.assertFalse(call.sample_all_positions) + self.assertEqual(output, None if prefill_only else 77) + + def test_forward_and_raw_keep_default_sampling_and_output_contracts(self): + tensor = SimpleNamespace(_underlying=object()) + self.assertEqual(self.engine.forward(tensor), 77) + self.assertEqual( + self.engine.forward_raw(tensor), + {"output_ids": 77, "logits": 88, "hidden_states": 99}, + ) + for call in self.calls: + self.assertIs(call.input_ids, tensor._underlying) + self.assertEqual((call.temperature, call.top_k, call.top_p), (1.0, 1, 1.0)) + self.assertFalse(call.prefill_only) + self.assertFalse(self.calls[0].sample_all_positions) + self.assertTrue(self.calls[1].sample_all_positions) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/llm/test_scheduler_cache_lifecycle.py b/test/llm/test_scheduler_cache_lifecycle.py new file mode 100644 index 000000000..6d18efbbf --- /dev/null +++ b/test/llm/test_scheduler_cache_lifecycle.py @@ -0,0 +1,383 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from cache_test_support import MODULES, assert_state, chain_hashes, publish + +Scheduler = MODULES["scheduler"].Scheduler +InferenceRequest = MODULES["request"].InferenceRequest +RequestStatus = MODULES["request"].RequestStatus +SamplingParams = MODULES["sampling_params"].SamplingParams +MambaCacheManager = MODULES["cache_manager"].MambaCacheManager + + +class DelayedConnector: + def request_finished(self, request, block_table, block_size): + return True, None + + +class SchedulerCacheTests(unittest.TestCase): + def make_scheduler(self, **kwargs): + scheduler = Scheduler(**kwargs) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + return scheduler + + @staticmethod + def output(**connector_output): + return SimpleNamespace(kv_connector_output=SimpleNamespace(**connector_output)) + + def test_connector_metadata_survives_prefill_decode_and_idle_dispatch(self): + metadata = object() + connector = SimpleNamespace( + build_connector_meta=lambda: metadata, + get_num_new_matched_tokens=lambda req, cached: (0, False), + update_state_after_alloc=lambda *args: None, + request_finished=lambda *args: (False, None), + ) + scheduler = self.make_scheduler( + num_blocks=4, block_size=16, connector=connector + ) + req = InferenceRequest( + "metadata", + prompt_token_ids=[11], + sampling_params=SamplingParams(max_tokens=2), + ) + scheduler.add_request(req) + prefill = scheduler.schedule() + self.assertTrue(prefill.is_prefill) + self.assertIs(prefill.kv_connector_metadata, metadata) + req.append_generated_token_id(12) + scheduler.complete_requests([req]) + decode = scheduler.schedule() + self.assertFalse(decode.is_prefill) + self.assertEqual(decode.scheduled_requests, [req]) + self.assertIs(decode.kv_connector_metadata, metadata) + req.mark_canceled() + scheduler.complete_requests([req]) + idle = scheduler.schedule() + self.assertEqual(idle.scheduled_requests, []) + self.assertIs(idle.kv_connector_metadata, metadata) + assert_state(self, scheduler.cache_manager) + + def test_send_completion_releases_exactly_once(self): + scheduler = self.make_scheduler( + num_blocks=1, block_size=16, connector=DelayedConnector() + ) + manager = scheduler.cache_manager + request = InferenceRequest("delayed", prompt_token_ids=[11] * 16) + request.block_table, _ = manager.allocate_slots(16) + manager.publish_computed_blocks( + request.block_table, chain_hashes([11] * 16), 0, 16 + ) + block_id = request.block_table[0] + request.status = RequestStatus.CANCELED + + scheduler.complete_requests([request]) + + self.assertFalse(manager.try_free_blocks(1)) + self.assertEqual(manager.blocks[block_id].ref_count, 1) + event = self.output(finished_sending={"delayed"}) + scheduler.update_from_output(event) + scheduler.update_from_output(event) + self.assertEqual(scheduler.pending_free_blocks, {}) + self.assertEqual(manager.get_total_usable_blocks(), 1) + assert_state(self, manager) + + def test_terminal_status_releases_tail_and_preserves_shared_prefix(self): + for status in ( + RequestStatus.FINISHED, + RequestStatus.CANCELED, + RequestStatus.FAILED, + RequestStatus.TIMEOUT, + ): + with self.subTest(status=status): + scheduler = self.make_scheduler(num_blocks=2, block_size=16) + manager = scheduler.cache_manager + request = InferenceRequest(status.value, prompt_token_ids=[11] * 17) + request.block_table, _ = manager.allocate_slots(17) + hashes = chain_hashes([11] * 16) + manager.publish_computed_blocks(request.block_table, hashes, 0, 16) + shared, hit = manager.get_computed_blocks(hashes, 16) + prefix_id, tail_id = request.block_table + request.status = status + + scheduler.complete_requests([request]) + + self.assertEqual(hit, 16) + self.assertEqual(manager.blocks[prefix_id].ref_count, 1) + self.assertIn(prefix_id, manager.used_block_ids) + self.assertIn(tail_id, manager.free_block_ids) + self.assertNotIn(prefix_id, manager._evictable_blocks) + manager.free_blocks(shared) + self.assertIn(prefix_id, manager._evictable_blocks) + assert_state(self, manager) + + def test_admission_failure_returns_temporary_prefix_pin(self): + scheduler = self.make_scheduler(num_blocks=2, block_size=16) + manager = scheduler.cache_manager + table, hashes = publish(manager, [11] * 16) + request = InferenceRequest( + "too-large", + prompt_token_ids=[11] * 16 + [12], + sampling_params=SamplingParams(max_tokens=64), + ) + scheduler.add_request(request) + + self.assertIsNone(scheduler.schedule()) + + self.assertEqual(request.status, RequestStatus.WAITING) + self.assertEqual(scheduler.waiting_queue.sync_q.qsize(), 1) + self.assertEqual(manager.blocks[table[0]].ref_count, 0) + self.assertEqual(list(manager._evictable_blocks), table) + self.assertEqual(manager.get_total_usable_blocks(), 2) + assert_state(self, manager) + + def test_failed_admission_refreshes_prefix_release_recency(self): + scheduler = self.make_scheduler(num_blocks=3, block_size=16) + manager = scheduler.cache_manager + a_table, _ = publish(manager, [11] * 16) + b_table, b_hashes = publish(manager, [22] * 16) + c_table, _ = publish(manager, [33] * 16) + request = InferenceRequest( + "touch-a", + prompt_token_ids=[11] * 16 + [12], + sampling_params=SamplingParams(max_tokens=64), + ) + scheduler.add_request(request) + + self.assertIsNone(scheduler.schedule()) + allocation = manager.allocate_slots(16) + + self.assertIsNotNone(allocation) + self.assertEqual(list(manager._evictable_blocks), [c_table[0], a_table[0]]) + self.assertNotIn(b_hashes[0], manager.hash_to_block_ids) + self.assertEqual(allocation[0], b_table) + manager.free_blocks(allocation[0]) + assert_state(self, manager) + + def test_token_budget_defer_returns_second_requests_temporary_pin(self): + scheduler = self.make_scheduler( + max_batch_size=2, + max_num_batched_tokens=16, + num_blocks=8, + block_size=16, + ) + manager = scheduler.cache_manager + prefix_table, _ = publish(manager, [11] * 16) + first = InferenceRequest( + "first", + prompt_token_ids=[11] * 16 + [21], + sampling_params=SamplingParams(max_tokens=1), + ) + second = InferenceRequest( + "second", + prompt_token_ids=[11] * 16 + [31] * 32, + sampling_params=SamplingParams(max_tokens=1), + ) + scheduler.add_request(first) + scheduler.add_request(second) + + batch = scheduler.schedule() + + self.assertEqual(batch.scheduled_requests, [first]) + self.assertEqual(first.status, RequestStatus.RUNNING) + self.assertEqual(second.status, RequestStatus.WAITING) + self.assertEqual(scheduler.waiting_queue.sync_q.qsize(), 1) + self.assertEqual(manager.blocks[prefix_table[0]].ref_count, 1) + self.assertNotIn(prefix_table[0], manager._evictable_blocks) + first.status = RequestStatus.CANCELED + scheduler.complete_requests([first]) + assert_state(self, manager) + + def test_prefix_disabled_releases_all_pages_as_free(self): + scheduler = self.make_scheduler( + num_blocks=3, block_size=16, enable_prefix_caching=False + ) + request = InferenceRequest( + "no-prefix", + prompt_token_ids=[11] * 17, + sampling_params=SamplingParams(max_tokens=1), + ) + scheduler.add_request(request) + batch = scheduler.schedule() + request.status = RequestStatus.CANCELED + + scheduler.complete_requests([request]) + + self.assertEqual(batch.scheduled_requests, [request]) + self.assertEqual(list(request.block_hashes), []) + self.assertEqual(scheduler.cache_manager.get_num_free_blocks(), 3) + self.assertEqual(scheduler.cache_manager.hash_to_block_ids, {}) + assert_state(self, scheduler.cache_manager) + + def test_receive_completion_releases_canceled_remote_request_once(self): + scheduler = self.make_scheduler( + num_blocks=1, block_size=16, connector=DelayedConnector() + ) + manager = scheduler.cache_manager + request = InferenceRequest( + "remote", + prompt_token_ids=[11] * 16, + sampling_params=SamplingParams(max_tokens=1), + ) + request.block_table, _ = manager.allocate_slots(16) + manager.publish_computed_blocks( + request.block_table, chain_hashes([11] * 16), 0, 16 + ) + block_id = request.block_table[0] + scheduler.remote_kv_requests[request.request_id] = request + scheduler.pending_kv_decode_blocks = 1 + request.status = RequestStatus.CANCELED + + scheduler.complete_requests([request]) + + self.assertEqual(scheduler.pending_kv_decode_blocks, 0) + self.assertFalse(manager.try_free_blocks(1)) + self.assertEqual(manager.blocks[block_id].ref_count, 1) + event = self.output(finished_recving={"remote", "unknown"}) + scheduler.update_from_output(event) + scheduler.update_from_output(event) + self.assertEqual(scheduler.pending_free_blocks, {}) + self.assertEqual(manager.get_total_usable_blocks(), 1) + assert_state(self, manager) + + def test_mamba_manager_and_scheduler_release_owned_rows_and_pages(self): + mamba = MambaCacheManager(3) + first = mamba.allocate() + second = mamba.allocate() + self.assertEqual((first, second), (1, 2)) + self.assertIsNone(mamba.allocate()) + mamba.free(first) + self.assertEqual(mamba.allocate(), 1) + + scheduler = self.make_scheduler( + num_blocks=3, + block_size=16, + has_mamba_cache=True, + num_mamba_cache_blocks=2, + ) + request = InferenceRequest( + "mamba", + prompt_token_ids=[11] * 17, + sampling_params=SamplingParams(max_tokens=1), + ) + scheduler.add_request(request) + batch = scheduler.schedule() + owned_row = request.mamba_cache_index + request.status = RequestStatus.CANCELED + + scheduler.complete_requests([request]) + + self.assertEqual(batch.scheduled_requests, [request]) + self.assertEqual(list(request.block_hashes), []) + self.assertIsNotNone(owned_row) + self.assertIsNone(request.mamba_cache_index) + self.assertEqual(scheduler.mamba_cache_manager.get_num_free_blocks(), 1) + self.assertEqual(scheduler.cache_manager.get_num_free_blocks(), 3) + assert_state(self, scheduler.cache_manager) + + +class SegmentedSchedulerCacheTests(SchedulerCacheTests): + def make_scheduler(self, **kwargs): + kwargs["prefix_cache_policy"] = "slru" + kwargs["prefix_cache_protected_ratio"] = 0.5 + return super().make_scheduler(**kwargs) + + def test_only_successful_admission_promotes_local_prefix(self): + scheduler = self.make_scheduler(num_blocks=4, block_size=16) + manager = scheduler.cache_manager + table, hashes = publish(manager, [11] * 16) + rejected = InferenceRequest( + "rejected", + prompt_token_ids=[11] * 16 + [12], + sampling_params=SamplingParams(max_tokens=128), + ) + scheduler.add_request(rejected) + self.assertIsNone(scheduler.schedule()) + self.assertFalse(manager._protected_blocks) + rejected.status = RequestStatus.CANCELED + admitted = InferenceRequest( + "admitted", + prompt_token_ids=[11] * 16 + [13], + sampling_params=SamplingParams(max_tokens=1), + ) + scheduler.add_request(admitted) + self.assertEqual(scheduler.schedule().scheduled_requests, [admitted]) + self.assertEqual(list(manager._protected_blocks), table) + self.assertFalse(manager._protected_evictable_blocks) + admitted.status = RequestStatus.CANCELED + scheduler.complete_requests([admitted]) + self.assertEqual(list(manager._protected_evictable_blocks), table) + for token in range(20, 28): + publish(manager, [token] * 16) + self.assertIn(hashes[0], manager.hash_to_block_ids) + assert_state(self, manager) + + def test_token_budget_probe_does_not_promote_unrelated_prefix(self): + scheduler = self.make_scheduler( + num_blocks=8, block_size=16, max_batch_size=2, max_num_batched_tokens=16 + ) + manager = scheduler.cache_manager + table, _ = publish(manager, [22] * 16) + first = InferenceRequest( + "first", + prompt_token_ids=[11] * 8, + sampling_params=SamplingParams(max_tokens=1), + ) + second = InferenceRequest( + "deferred", + prompt_token_ids=[22] * 16 + [33] * 16, + sampling_params=SamplingParams(max_tokens=1), + ) + scheduler.add_request(first) + scheduler.add_request(second) + self.assertEqual(scheduler.schedule().scheduled_requests, [first]) + self.assertEqual(second.status, RequestStatus.WAITING) + self.assertNotIn(table[0], manager._protected_blocks) + first.status = RequestStatus.CANCELED + scheduler.complete_requests([first]) + assert_state(self, manager) + + def test_allocation_failure_does_not_promote(self): + scheduler = self.make_scheduler(num_blocks=4, block_size=16) + manager = scheduler.cache_manager + table, _ = publish(manager, [11] * 16) + request = InferenceRequest( + "allocation-failed", + prompt_token_ids=[11] * 16 + [12], + sampling_params=SamplingParams(max_tokens=1), + ) + scheduler.add_request(request) + with patch.object(manager, "allocate_slots", return_value=None): + self.assertIsNone(scheduler.schedule()) + self.assertEqual(request.status, RequestStatus.WAITING) + self.assertFalse(manager._protected_blocks) + self.assertEqual(list(manager._evictable_blocks), table) + assert_state(self, manager) + + def test_protected_remote_owner_is_not_evictable_until_transfer_finishes(self): + scheduler = self.make_scheduler( + num_blocks=4, block_size=16, connector=DelayedConnector() + ) + manager = scheduler.cache_manager + _, hashes = publish(manager, [11] * 16) + table, _ = manager.get_computed_blocks(hashes, 16) + manager.record_cache_hit(table) + request = InferenceRequest("protected-send", prompt_token_ids=[11] * 16) + request.block_table = table + request.status = RequestStatus.CANCELED + scheduler.complete_requests([request]) + self.assertEqual(list(manager._protected_blocks), table) + self.assertFalse(manager.try_free_blocks(4)) + self.assertEqual(manager.blocks[table[0]].ref_count, 1) + event = self.output(finished_sending={request.request_id}) + scheduler.update_from_output(event) + scheduler.update_from_output(event) + self.assertEqual(list(manager._protected_evictable_blocks), table) + self.assertTrue(manager.try_free_blocks(4)) + assert_state(self, manager) + + +if __name__ == "__main__": + unittest.main()