feat(engine): integrate prefix eviction and chunked parallel execution

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<const cache::PagedKVCacheConfig *>(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);

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);

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<float>(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) {

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;
     };
 

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};
         }
 

csrc/models/infinilm_model.hpp

@@ -57,6 +57,8 @@ class InfinilmModel : public infinicore::nn::Module {
         std::optional<infinicore::Tensor> 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 {

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<cache::CacheConfig> 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<cache::CacheConfig> 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<cache::CacheConfig> {
             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<std::string> 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<bool>(item.second);
+                    } else if (key == "temperature") {
                         input.temperature = py::cast<float>(item.second);
                     } else if (key == "top_p") {
                         input.top_p = py::cast<float>(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)

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.

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,

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)
 

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)
 

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,
     )