perf: reduce KV admission over-reservation

[lgeln10]

Summary

Reduce KV admission over-reservation and add a fixed benchmark suite for validating the change with the same request set across unit tests, scheduler microbenchmarks, and OpenAI-compatible serving benchmarks.

This PR changes scheduler admission by:

• reserving only the additional decode KV blocks a running request may still need, instead of reserving all remaining generated tokens as fresh blocks
• tracking running decode reservations incrementally with num_reserved_decode_blocks, so admission no longer scans the full running queue on every candidate request
• applying the same exact extra-block formula to pending remote-KV requests
• exposing num_reserved_decode_blocks through cache stats for tests and benchmark validation

This PR also adds:

• a deterministic TinyLlama KV admission manifest with identical wave1/wave2 request shapes across validation layers
• manifest loading and validation helpers shared by all benchmark scripts
• a scheduler admission microbenchmark comparing legacy, exact-scan, and incremental variants
• a streaming serving benchmark that measures wave2 TTFT and overall output throughput
• a result summarizer that checks all runs use the same manifest before comparing variants
• unit tests for manifest validation, pressure prediction, fixed workload assumptions, and exact reservation accounting

What is measured

The benchmark separates the change into three variants so the over-reservation fix and the bookkeeping optimization can be evaluated independently:

Variant	Reservation formula	Running request accounting	What it shows
legacy	Old formula	Scan running queue	Baseline with KV over-reservation.
exact-scan	Exact extra-block formula	Scan running queue	Effect of fixing over-reservation only.
incremental	Exact extra-block formula	Incremental counter	Effect of avoiding the per-admission running-queue scan.

The over-reservation bug is that admission used the remaining output token budget as if all of it required new KV blocks. For requests that already own partially filled or sufficient KV blocks, this double-counts capacity. The fixed formula computes the final block requirement for prompt_len + max_tokens, subtracts the request's current block_table length, and reserves only the extra blocks that may still be needed.

End-to-end data

Archived serving runs on RTX 4090 D + TinyLlama-1.1B-Chat-v1.0 used the fixed 384-request workload: 256 first-wave requests, 128 second-wave requests, 128 max output tokens, input lengths cycling through 64/128/192/256/320/512/768/1024 tokens, block_size=256, num_blocks=864, prefix caching disabled, and EOS ignored. Each variant was run 3 times and the table reports medians.

Variant	Wave2 TTFT P50	Wave2 TTFT P95	Overall output throughput	Successful requests
legacy	7.68s	8.87s	2704.85 tok/s	384/384
exact-scan	1.54s	11.53s	2509.70 tok/s	384/384
incremental	1.49s	11.47s	2507.13 tok/s	384/384

Key observations:

• The old admission logic over-reserves decode KV by 96 blocks in this workload: 256 blocks reserved by legacy versus 160 blocks required by the exact formula.
• Fixing over-reservation increases predicted immediately admitted second-wave requests from 53 to 91.
• Fixing over-reservation alone reduces second-wave TTFT P50 by 79.9% (7.68s -> 1.54s).
• Adding incremental accounting reduces second-wave TTFT P50 by another 3.4% (1.54s -> 1.49s).
• Legacy to incremental reduces second-wave TTFT P50 by 80.6% overall (7.68s -> 1.49s).
• Incremental accounting reduces admission-check median time from 509.73 us to 23.29 us, a 21.9x speedup over exact scanning.
• Overall output throughput decreases by 7.3% in this mixed workload, and wave2 P95 regresses. This PR improves admission accuracy, scheduler CPU overhead, and P50 latency; it does not claim an end-to-end throughput win.

The archived results used the same request contents before compacting the JSON manifest in this PR, so the recorded file SHA differs from the committed compact JSON while the workload fields are unchanged.

Additional workload sanity check

Before the fixed TinyLlama manifest, the same idea was checked with a small head_dim=64 random Llama model on RTX 4090 D. That run is less representative for production model speed, but it is useful as an ablation because it includes a normal control workload where the exact formula should not release extra capacity.

Each row below is the median of 3 serving runs with 128 first-wave requests and 64 second-wave requests:

Workload	Input shape	Legacy admissions	Exact admissions	What it checks	Legacy wave2 TTFT P50	Incremental wave2 TTFT P50
aligned	Prompt lengths aligned to block boundaries	32	32	Control case with no over-reservation benefit.	603 ms	671 ms
mixed	Mixed prompt lengths	27	46	More normal mixed-length case.	760 ms	551 ms
boundary	Boundary prompt length	0	64	Stress case that maximizes the bug.	856 ms	621 ms

This sanity check shows the optimization is not expected to help every workload. When there is no false reservation (aligned), exact admission does not increase accepted requests. When prompt lengths create false reservations (mixed and boundary), the exact formula releases capacity and admits more second-wave requests.

Testing

Passed locally:

PATH=/tmp/infinilm-pr-pytest/bin:$PATH /tmp/infinilm-pr-pytest/bin/python scripts/format.py --check --path \
  python/infinilm/llm/scheduler.py \
  scripts/kv_admission_manifest.py \
  scripts/generate_kv_admission_manifest.py \
  scripts/benchmark_scheduler_admission.py \
  scripts/benchmark_serving_admission.py \
  scripts/summarize_kv_admission_results.py \
  test/scripts/test_kv_admission_manifest.py \
  test/scripts/test_benchmark_scheduler_admission.py \
  test/scripts/test_benchmark_serving_admission.py \
  test/scripts/test_fixed_kv_admission_workload.py

PYTHONPATH=python:. /tmp/infinilm-pr-pytest/bin/python -m py_compile \
  python/infinilm/llm/scheduler.py \
  scripts/kv_admission_manifest.py \
  scripts/generate_kv_admission_manifest.py \
  scripts/benchmark_scheduler_admission.py \
  scripts/benchmark_serving_admission.py \
  scripts/summarize_kv_admission_results.py \
  test/scripts/test_kv_admission_manifest.py \
  test/scripts/test_benchmark_scheduler_admission.py \
  test/scripts/test_benchmark_serving_admission.py \
  test/scripts/test_fixed_kv_admission_workload.py

PYTHONPATH=python:. /tmp/infinilm-pr-pytest/bin/python -m pytest \
  test/scripts/test_kv_admission_manifest.py \
  test/scripts/test_benchmark_serving_admission.py

Passed on RTX 4090 D server with the PR source uploaded from the local checkout and the existing _infinilm.so copied into the test tree:

export PYTHONPATH=/data/InfiniLM-pr590/python:/data/InfiniCore/python:/data/InfiniCore/build/linux/x86_64/release:$PYTHONPATH
export LD_LIBRARY_PATH=/data/InfiniCore/build/linux/x86_64/release:/data/InfiniCore/install/lib:$LD_LIBRARY_PATH
export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libfmt.so.9
/data/venv/bin/python -m py_compile \
  python/infinilm/llm/scheduler.py \
  scripts/kv_admission_manifest.py \
  scripts/generate_kv_admission_manifest.py \
  scripts/benchmark_scheduler_admission.py \
  scripts/benchmark_serving_admission.py \
  scripts/summarize_kv_admission_results.py \
  test/scripts/test_kv_admission_manifest.py \
  test/scripts/test_benchmark_scheduler_admission.py \
  test/scripts/test_benchmark_serving_admission.py \
  test/scripts/test_fixed_kv_admission_workload.py

/data/venv/bin/python -m pytest \
  test/scripts/test_kv_admission_manifest.py \
  test/scripts/test_benchmark_scheduler_admission.py \
  test/scripts/test_benchmark_serving_admission.py \
  test/scripts/test_fixed_kv_admission_workload.py

Result:

7 passed in 4.65s