Skip to content

perf: avoid scanning KV cache pages for capacity checks - #576

Open
T4t4KAU wants to merge 1 commit into
InfiniTensor:InfiniLM-v0.2.9cfrom
T4t4KAU:perf/kv-cache-capacity-index
Open

T4t4KAU wants to merge 1 commit into
InfiniTensor:InfiniLM-v0.2.9cfrom
T4t4KAU:perf/kv-cache-capacity-index

Conversation

@T4t4KAU

@T4t4KAU T4t4KAU commented Sep 17, 2026

Copy link
Copy Markdown

Summary

KV admission checks currently scan every allocated cache page to count reclaimable pages. Track the number of physical pages with live references so BlockManager.get_total_usable_blocks() becomes O(1).

  • Update the count when allocating a page, pinning a previously unreferenced prefix page, releasing its final reference, or rolling back speculative pages.
  • Preserve prefix lookup, reference counts, allocation and eviction order, and admission rules.
  • Add eight CPU-only cache/scheduler tests. Their module loader avoids the GPU package initializer and restores only the module entries it replaces.

Targets InfiniLM-v0.2.9c at 270feb3. One commit: 21c1444.

Motivation

Each candidate request calls the capacity query during admission. Repeated scans grow with the page pool and the number of candidate requests. Counting referenced physical pages directly removes this scan; shared prefix pages count once regardless of the number of owners. No linked issue.

Type of Change

  • perf — performance improvement with unchanged cache behavior.
  • Breaking change

Test Results of Involved Models on Supported Platforms

Environment / check Result
Linux x86_64, Intel Core i5-14600KF, Python 3.12.13 PASS: all eight CPU tests, including 10,000 randomized cache operations.
Cache lifecycle PASS: shared prefixes, duplicate hashes, prefix reuse followed by decode allocation, speculative rollback, rejected truncation, failed allocation/retry, and delayed release for remote transfer completion.
Scheduler comparison against the original 270feb3 source PASS: identical request order, prefix hits, physical pages, slot mappings, and available capacity for the exercised lifecycle.
Test-loader manual checks PASS: subsequent NumPy import, repeated loading, and restoration of pre-existing module entries.
Formatting PASS: repository scripts/format.py --check, Ruff lint/format, and git diff --check.
Full-model inference, service throughput, other platforms Not run. This change is confined to Python cache bookkeeping; validation covers CPU state transitions and scheduler behavior. No full-model speedup is claimed.

Tests ran noninteractively; results are recorded as text rather than screenshots.

python -m unittest discover -s test/llm -v
Ran 8 tests
OK

The named model/service tests were skipped: examples/test_infer.py, examples/bench.py, test/bench/test_benchmark.py, and inference_server.py + scripts/test_perf.py. The available validation is the CPU cache/scheduler suite and synthetic admission microbenchmark.

Benchmark / Performance Impact

CPU-only synthetic admission fixture on the same i5-14600KF / Python 3.12.13 machine. Baseline is the original cache-manager source at 270feb3; after is this commit. No tensor dtype or GPU kernels are involved.

  • Block size 16; all pages previously allocated, with references released for half the pages.
  • 32 running requests, each with max_tokens=128 and 64 generated tokens.
  • Candidate prompt length 128, max_tokens=128.
  • Each measurement repeats 32 can_accept_request() calls against fixed state. It does not enqueue/execute those candidates.
  • Five warmup groups; median of five rounds, each containing five groups of 32 checks. Timing uses perf_counter_ns; fixture setup is excluded. CPU affinity/frequency were not pinned.
Cache pages Before: 32 checks (ms) After: 32 checks (ms)
512 0.8834 0.5545
8,192 5.4589 0.5588
32,768 20.4412 0.5604

These measurements isolate admission checks, not full scheduling or model generation. The running-queue scan and actual eviction scan remain unchanged. The standalone timing harness was run locally and is not included in this PR.

Notes for Reviewers

The invariant is _num_referenced_blocks == count(block.ref_count > 0). Therefore usable capacity is num_blocks - _num_referenced_blocks.

  • Prefix pinning increments the count only on 0→1.
  • Releasing a shared page decrements it only on 1→0.
  • Speculative truncation validates all discarded pages before decrementing their counts.
  • Evicting an already unreferenced page must not decrement the count again.
  • Tests compare capacity with independently reconstructed ownership and the prior scanning query.

CI / ChatOps

Local checks passed. Remote CI status is pending. Hardware CI has not been manually triggered; no ChatOps command or reviewer notification has been posted.


Checklist

Title, Branch, and Commits

  • PR title and commit message follow Conventional Commits.
  • Branch follows <type>/xxx-yyyy-zzzz: perf/kv-cache-capacity-index.
  • One commit, with no fixup/squash/wip or stray merge commits.
  • N/A: Rebase onto main — this PR intentionally targets InfiniLM-v0.2.9c, directly based on its current 270feb3 head.
  • N/A: Legacy issue-format branch or commit.

Scope and Design

  • Changes are limited to cache capacity bookkeeping and its CPU regression tests.
  • No dead code, debug output, unowned TODOs, or unrelated formatting changes.
  • N/A: Public API changes — signatures and cache/admission behavior are unchanged.

General Code Hygiene

  • Comments explain the counter invariant and test import isolation.
  • Modified files end in one newline, with no trailing whitespace, tab/space mixing, or BOM.
  • Added comments and docstrings are English and follow Python conventions.
  • N/A: New diagnostic identifiers requiring Markdown escaping — no new production diagnostics.

C++ Specific

  • N/A: No C++ files changed.

Python Specific

  • Ruff lint and formatting pass.
  • Changed files pass scripts/format.py --path python/infinilm/llm/cache_manager.py test/llm/test_cache_manager.py --check.
  • No changes to python/infinilm/auto_config.py.

Testing

  • Eight CPU tests pass, including 10,000 randomized operations.
  • Scheduler lifecycle matches the original release source.
  • Other platforms and real-model/service validation — not run; scope and reasons are stated above. No hardware reviewer has been tagged.
  • Single request test — skipped; CPU cache/scheduler validation used.
  • Offline model performance test — skipped; admission microbenchmark only.
  • Model accuracy benchmark — skipped; no model computation changed.
  • Service performance test — skipped; no end-to-end throughput claim.

Build, CI, and Tooling

  • N/A: Fresh native build — this PR changes only Python code and CPU tests.
  • Hardware CI manually triggered or /retest requested — pending.

Documentation

  • N/A: README/CONTRIBUTING updates — no user behavior, build flags, or developer workflow changes.
  • N/A: Breaking-change documentation — no breaking change.

Security and Safety

  • No secrets, tokens, internal URLs, customer data, or personal hardware identifiers included.
  • N/A: Third-party code additions.
  • N/A: Pointer arithmetic or uninitialized native reads — Python-only change.

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.

1 participant