Skip to content

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

Open
big-hip wants to merge 8 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:feat/cache-chunk-graphs
Open

big-hip wants to merge 8 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:feat/cache-chunk-graphs

Conversation

@big-hip

@big-hip big-hip commented Sep 16, 2026

Copy link
Copy Markdown

Summary

Integrate paged-prefix LRU/optional SLRU with bounded Prefill scheduling, TP2 Decode graphs and PP2 eager execution. Intermediate chunks publish only completed KV pages and emit no token; completed requests and cancellation release shared page references consistently.

  • Add O(1) ordered reclaim queues with tail-first final release, bounded SLRU protection and successful-admission-only promotion. The chunk path explicitly records admitted hits.
  • Reuse the original Decode scheduler and common output construction across Prefill policies; share configuration-test module loading.
  • Rotate Decode, Prefill continuation and new admission. Support dense paged TP/PP=1/1, 2/1 and 1/2; chunking defaults to disabled and LRU remains the default policy.
  • Skip unused intermediate LM-head/sampling/output work while executing every Transformer layer fully. Propagate the outputless contract to both PP stages.
  • Combine eager chunked Prefill with the existing Decode device graphs on TP1/TP2. This PR adds no Prefill graph compiler or Prefill-graph configuration switch. Single-token final tails may use the existing Decode graph.
  • Add focused CPU/native lifecycle regressions, actual graph-launch counting and a concise usage guide. Historical measurements and experiment scripts are linked below, outside the final source diff.

Motivation

Long uninterrupted Prefill delays other requests, while cache reclamation can discard reusable prefixes. These mechanisms share admission/publication/release state, so the PR includes their integration checks. It follows the input slicing/continuation idea in #371 and adds explicit phase rotation and current per-request KV ownership; #571's admission priority is a separate feature. Neither proposal is merged here.

Type of Change

  • feat — optional cache policy and chunked execution combinations
  • fix — integrated cache admission and outputless execution contracts
  • Breaking change

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

Current retained suite at cb243d9a (2026-09-17): 83 CPU tests passed. Consolidate the cache/chunk configuration suites and duplicate output checks, removing 230 test lines while retaining cache ownership, admission failure, cancellation, deferred remote-KV release, capacity and phase-progress coverage. Production implementation is unchanged. The earlier 101-test counts below describe the historical suite before consolidation, not the current test count. Saved current output: immutable evidence. The historical GPU/performance matrix was not repeated for this test-only update.

Current usage guide and regression-test commands. Archived machine-readable results and historical reproduction commands: integrated validation. The screenshot below renders archived audited results, not a new GPU run or an upstream CI pass. It also includes the separately deferred Prefill-graph experiment; those entries are outside this PR. The scope table below identifies the relevant historical coverage.

Audited model and performance results

Platform/model Relevant archived validation
Two NVIDIA A6000s, Qwen2.5-1.5B FP16 13 TP1/TP2/PP2 eager/Decode-graph configurations, LRU/SLRU, prefix on/off, chunk300 across page256 boundaries. Completed shared/repeated outputs matched one baseline.
MetaX C500 50% compute slice, Qwen3-0.6B BF16 Two archived lifecycle configurations with eager Prefill + Decode graphs: LRU and SLRU, chunk512. The four configurations using Prefill graphs are deferred and do not establish coverage for this PR head. Qwen3-4B Decode performance evidence is reported separately below.
CPU 101 runtime/configuration/input-boundary regressions passed, including admission rejection, cancellation, shared/deferred KV ownership, SLRU promotion, capacity accounting, phase progress and PP worker config forwarding.
Ascend/Moore/other devices No connected test hardware; no new platform support claim.

First/last local KV layers were poisoned and checked for expected writes and unchanged values outside scheduled slots. PP stage 1 was checked separately. All cases ended with zero references and all 16 pages usable. Device-graph launches were counted directly, including both TP ranks. A final output-suppression check and an asynchronous 160-output-token mixed-request comparison also passed.

Earlier validation after deferring Prefill graphs (before test consolidation): all 101 CPU regressions passed. Five affected C++ translation units were rebuilt and linked with 122 unchanged cached objects; the resulting library contains no experimental Prefill-capture/replay symbols. A two-A6000 Qwen2.5-1.5B FP16 check (TP2/PP1, SLRU, page256/pool16, chunk300) observed zero graph launches for intermediate Prefill and 28 Decode launches across both ranks. Shared/repeated output IDs matched the archived same-configuration reference; cancellation returned all 16 pages with zero references. This uses the new native binary and current Python source. It is a scoped correctness rerun, not a clean full build or a new performance/C500 matrix.

Template test entrypoints and disposition:

Requested check Evidence / reason for substitution
examples/test_infer.py single request Integrated direct scheduler/model runners exercise the changed paged/chunk path, output suppression, repeated outputs and cancellation. The named example was not rerun for this integration.
examples/bench.py offline performance This direct-native example bypasses the scheduler and rejects chunking; use the archived mixed-request experiment linked below for the scheduling comparison.
test/bench/test_benchmark.py sanity Broad quality matrix not run; scoped greedy token controls and the retained FP16 mismatch diagnosis are documented.
HTTP server + scripts/test_perf.py service Integrated HTTP load test not run; AsyncLLMEngine mixed-request runners cover scheduler behavior, not HTTP overhead or sustained load.

These substitutions and untested entries are explicit; they do not claim that the named scripts or the broader quality/service matrices passed.

Benchmark / Performance Impact

Measured value: retain reusable prefixes under pressure, bound the pause caused by long Prefill, and combine chunking with device graphs. These are distinct workload-dependent benefits. The table reports archived component experiments before final integration, not a full performance rerun of the PR head. Final integration correctness is reported separately above.

Mechanism / baseline Conditions and samples Before -> after Meaning and cost
LRU vs former set-based reclamation A6000, Qwen2.5-1.5B FP16, TP1 eager, pool64/page256, concurrency1; two independent repeats, 112 measured requests each Prefill 56,832 -> 32,256 tokens (-43.24%); mean TTFT median paired change -28.24% Hotspot trace; throughput changes only +0.87% / +0.32%, no stable throughput gain
SLRU vs LRU Same GPU/model, pool8, protected ratio0.5; four warmups + ten measured requests, one run per mode Hits 1,024 -> 2,048; Prefill 4,416 -> 3,392 tokens (-23.19%) Established hot prefixes survive a one-use scan; not a universal policy improvement
Chunk512 vs unchunked eager Two A6000s, Qwen2.5-1.5B FP16, TP2/PP1, prefix off; 8192-token long prompt + two 128-token prompts; three runs per mode Active maximum ITL 1,022.94 -> 142.74 ms (-86.05%); late short TTFT 1,001.54 -> 68.16 ms (-93.19%) Long TTFT +41.85%, active p95 ITL 24.88 -> 92.68 ms, window output rate 75.27 -> 62.83 token/s (-16.53%)
Omit intermediate LM-head/sampling/output Same TP2 chunk512 workload; three alternating runs per mode Long-Prefill dispatch total 1,202.99 -> 1,175.69 ms (-2.27%) Output rate 64.19 -> 62.87 token/s; no stable serving-throughput gain; all Transformer layers still execute
Add Decode graphs to chunk eager C500 50% slice, Qwen3-4B BF16, TP1, Flash Attention, chunk512; three single requests and two mixed windows per mode Single-request ITL 12.96 -> 9.37 ms (-27.64%); mixed-window rate 69.80 -> 88.06 token/s (+26.16%) Reuses the existing Decode compiler; measures compatibility with chunking, not a new Decode-graph implementation

How measured: NVIDIA uses CUDA12.4/driver580.105.08; paged KV, page256, greedy fixed output budgets. The TP2 mixed test emits 160 tokens total; the long request arrives after the active request's eighth output and the late short request 20 ms later. TP2 table values are medians of three runs. LRU uses 1056-input/32-output requests with 16 excluded warmups; SLRU uses 544-input/8-output requests. Cache counts are actual admitted hits and remaining prompt work.

C500 is a 50% compute / 32,000 MiB slice, six CPU cores, MACA3.5.3.20, driver3.8.30, torch2.8.0+metax3.5.3.9 and flash-attn2.6.3+metax3.5.3.9torch2.8. The archived modes use the same binary, BF16, TP1, page256/pool32 and prefix reuse off. Single requests: 2048 input/16 output, three measurements after warmup; mean TTFT and mean of per-request median ITL. Mixed: 128-input/48-output active request, 2048-input/16-output competitor after the eighth token, two windows; rate = 56 remaining output tokens / window duration. This is not sustained HTTP serving throughput. The archived output controls passed and actual graph launches were counted. The archive also contains Prefill-graph modes; their measurements are outside the submitted scope.

Tradeoffs and negative controls: LRU under a 128-page cyclic working set in a 64-block pool (7B, one run) reduced hits 33,792 -> 0 and output rate 34.96 -> 33.16 token/s (-5.14%). SLRU has no work reduction on the tested no-reuse, overcapacity-cycle and hotspot-shift metadata traces. Chunking can worsen p95 ITL despite reducing the maximum pause. Decode graphs require initialization and retained workspace; full graph memory overhead was not measured. FP16 prefix-shape output differences remain documented.

CPU-only reclaim microbenchmark: pool65,536, 90% pinned, Python3.11.15, five method-call samples: one-block median 2,058.758 us -> 18.217 us; retained Python metadata +5,767,328 bytes. Pool construction and hashing are excluded; this is not a GPU/model speedup.

Final integrated TP2 check: 2048-token long prompt, two 128-token prompts, 160 outputs, chunk512, prefix off, one window per mode: 161.71 token/s eager vs 161.02 token/s Decode graph; all output IDs match. No throughput gain is established by this pair. PP2 has lifecycle/correctness evidence, not a measured performance gain.

Archived C500 tables, setup, aggregation, historical version provenance, reproduction entrypoints and limits (including the separately deferred Prefill-graph experiment): performance report. Audited per-run values and source artifact hashes: measurement JSON. Benefits from different experiments must not be multiplied or added together.

Notes for Reviewers

Submission scope: the final diff retains implementation, regression tests and concise usage documentation. Two experiment-specific benchmarks, their 12 harness tests, result JSON/PNG files and repetitive reports are retained at archive revision a7345fa2, not proposed for merge. The prefix-cache experiment and mixed-request experiment remain available there. The archived evidence links and screenshot use immutable revisions; the full historical hardware/performance matrix has not been rerun; current scoped validation is reported above.

Review the pin/promote/release invariants, publication after all PP stages complete, and eager-Prefill/Decode-graph dispatch. Experimental fixed-size Prefill graph capture/replay has been removed from this PR and retained on a local experiment branch. PP graphs, TP2/PP2, remote-KV/chunk integration, MoE and multimodal chunking remain excluded. Quantized-model chunking has not been validated; the dense-model check does not reject quantization metadata. Full graph memory overhead and sustained service throughput are not measured. A historical near-tied FP16 token mismatch between prefix execution shapes is preserved in the documentation; this PR does not promise universal bitwise equivalence.

The C500 checks use the companion InfiniCore varlen ABI fix. Device graphs require an InfiniCore build with --graph=y; the LM flag alone can otherwise replay host operators. Native builds succeeded using unchanged cached objects where available; a clean full build remains CI validation.

CI / ChatOps

At this head, the upstream CI run reports action_required; repository-side approval is required before jobs can start. Fork formatting success is reported separately below.

Latest cb243d9a fork push checks passed: CI formatting and Ruff. This does not replace the pending full accelerator build.

Local repository formatting and diff checks passed. The formatting repair at 32d00aae passed fork CI format and Ruff. The scheduler reuse commit a7345fa2 passed fork CI format and Ruff, full-PR local formatting, and the targeted CPU/GPU checks above. Upstream PR workflows still need approval; the full hardware matrix also requires manual dispatch and upstream runner/secret access. Uncompleted checklist items remain unchecked, including a clean affected-platform build, remaining detailed style review and untested-platform reviewer coverage.


Checklist

Every contributor must verify every item below before requesting
review. Tick each box only after the check has actually been performed —
do not tick speculatively. If an item truly does not apply, replace the
checkbox with N/A and briefly explain why in an inline comment.

Title, Branch, and Commits

  • PR title follows Conventional Commits (e.g. feat(nvidia): …, fix(cuda/gemm): …).
  • Branch name follows <type>/xxx-yyyy-zzzz where <type> matches the PR title's Conventional Commits type and words are joined with hyphens (see CONTRIBUTING.md §Branches).
  • Each commit message follows Conventional Commits.
  • Small PR is a single squashable commit; or, for a large PR, every commit is meaningful, well-formed, and independently reviewable (see CONTRIBUTING.md §Pull Requests).
  • The branch is based on the requested InfiniLM-v0.2.9c base without development merge history; implementation, formatting and evidence commits are separately reviewable.
  • No fixup! / squash! / wip commits remain.
  • N/A: No legacy issue-format branch or commit.

Scope and Design

  • Changes are minimal — nothing unrelated to the stated motivation was added (CONTRIBUTING.md §Code/General).
  • No dead code, commented-out blocks, debug prints, printf/std::cout/print(...) left behind, or TODO without an owner and issue link.
  • No unrelated formatting churn that would obscure the diff.
  • Public API changes (if any) are intentional, documented, and reflected in affected callers/tests.

General Code Hygiene (applies to all languages)

  • The code is self-explanatory; comments were added only where the why is non-obvious (CONTRIBUTING.md §Code/General).
  • Every modified or added file ends with a single trailing newline (CONTRIBUTING.md §Code/General).
  • No trailing whitespace, tab/space mixing, or stray BOMs.
  • Identifiers in comments and error messages are wrapped in backticks (e.g. the `seqlens_k` tensor) (CONTRIBUTING.md §Code/General).
  • All comments and error messages are in English (CONTRIBUTING.md §Code/General).
  • Comments and error messages are complete sentences — capitalized first letter, terminal punctuation — unless the language/framework convention says otherwise (CONTRIBUTING.md §Code/General; §Python).

C++ Specific (if C++ files changed)

  • Code follows the Google C++ Style Guide strictly.
  • Error and warning message wording follows the LLVM Coding Standards (CONTRIBUTING.md §C++).
  • Constructor initializer list order matches member declaration order (CONTRIBUTING.md §C++).
  • No raw new/delete; RAII / smart pointers / existing allocators are used.
  • Changed files are formatted by scripts/format.py.
  • No changes/reference to csrc/models/llama_legacy/.

Python Specific (if Python files changed)

  • Code is PEP 8 compliant.
  • Comments are complete English sentences, starting with a capital letter and ending with punctuation; Markdown backticks are used for code references (CONTRIBUTING.md §Python).
  • Docstrings (if any) follow PEP 257 (CONTRIBUTING.md §Python).
  • Changed files are formatted by scripts/format.py.
  • No changes/reference to python/infinilm/auto_config.py.

Testing

  • For any platform that could not be tested, an explicit reason is given in the table and a reviewer with access has been tagged.
  • Passed single request test (examples/test_infer.py), or specify the reason for skipping.
  • Passed offline performance test (examples/bench.py), or specify the reason for skipping.
  • Passed sanity test (test/bench/test_benchmark.py), or specify the reason for skipping.
  • Passed service test (python/infinilm/server/inference_server.py + scripts/test_perf.py), or specify the reason for skipping.

Build, CI, and Tooling

  • The project builds cleanly from a fresh directory on at least one affected platform.
  • CI has been triggered manually (Actions → CI on this branch), or /retest was requested.

Documentation

  • README.md, CONTRIBUTING.md, or inline docs updated when behavior, build flags, or developer workflow changed.
  • N/A: Chunking and SLRU are optional. The unmerged experimental Prefill-graph environment switch is deferred with that implementation; existing upstream public APIs are preserved.

Security and Safety

  • No secrets, access tokens, internal URLs, customer data, or personal hardware identifiers have been committed.
  • Third-party code is license-compatible and attributed.
  • No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.

CI status after submission: the upstream run reports action_required with zero jobs started. Repository-side action is still required; this is not a completed build/test result.

@xiaoba17

Copy link
Copy Markdown

Thanks for calling out #571. I agree that the two changes are orthogonal: #573 controls phase rotation and chunked-prefill lifecycle, while #571 only orders initial admission from the waiting queue and keeps Decode non-preemptive.
If both changes land, priority and aging should be applied inside _admit_chunk_request() when selecting a new waiting request. Requests already admitted into chunking_queue should continue to follow #573’s FIFO/round-robin continuation policy, so priority does not let a long prompt monopolize Prefill.
I’ll rebase #571 and add focused chunked-prefill integration tests if #573 lands first.

Exercise shared CLI, service and example forwarding once with both policies enabled. Retain capability rejection, PP worker forwarding, cache lifetimes and native output checks while removing duplicated setup.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants