Replace the radix-tree KV allocator with vLLM's block pool - #60
Merged
Conversation
SchedulerStats has no gpu_cache_usage field; it is kv_cache_usage (vllm/v1/metrics/stats.py:181 in v0.19.0). The getattr default silently returned 0.0, so the kv_cache_pct column in every timeseries.csv was empty and sim-vs-vLLM comparisons had no way to see KV pressure. Read the attribute directly so a future rename fails loudly instead of zeroing the column again.
Port of vLLM v0.19.0's block_pool.py, the block/queue primitives in kv_cache_utils.py, and kv_cache_manager.py, extended to a NPU/CPU/CXL hierarchy. Nothing imports these modules yet, so this commit cannot change any result. Each pool is the single authority for its tier -- free list, prefix-cache index and refcounts in one object -- so num_free_blocks is exact and an allocation either succeeds or reports failure in the same call. The radix tree could not do that: evictable_size_ counted every unlocked token while evict() could only drop unlocked leaves, and memory was charged later via tree events rather than in the call that can fail, so the scheduler could believe in space that did not exist. The tiers share one key space. Block hashes are chained once at the NPU block size (hash(parent_hash, block_tokens)), and a lower tier whose blocks are N times larger keys on every Nth hash -- the last fine hash of each coarse block, which because the chain is cumulative identifies the whole prefix up to that point. That is offloading/scheduler.py::_get_block_hashes, and it keeps LMCache's 256-token chunk granularity without a second hash function. It also fixes a latent collision: radix_tree.py hashed hash(tuple(page_tokens)), so two different prefixes ending in the same 16 tokens collided. Traffic: recall (lower tier -> NPU) is charged; the write-through is reported for energy only, because vLLM's OffloadingConnector defers it to the next engine step on a dedicated stream specifically so it cannot delay token generation. Eviction from the NPU costs nothing -- the data is either a finished request's cache or already has a copy below. Both modules carry a self-test runnable without Docker: python3 -m serving.core.block_pool python3 -m serving.core.kv_cache_manager
Wires block_pool.py / kv_cache_manager.py in, deletes radix_tree.py, and rewrites Scheduler to vLLM V1's schedule(). These have to land together: the scheduler called the radix-era API (avail_size, lock_prefix, cache_unfinished_req, ...), so removing it without the rewrite leaves the tree unrunnable. Scheduler - one schedule() for prefix caching on and off, in two phases. Phase A serves the persistent self.running set, preempting only from its tail and retrying. Phase B admits from self.waiting while budget and sequence slots remain, breaking on the first allocation failure -- admission never preempts -- and is skipped entirely on any step that preempted. That anti-thrash rule is load-bearing. - no prefill phase and no decode phase: a request catches up to num_tokens_reached, so num_new = num_tokens_reached - num_computed_tokens, which is 1 in steady-state decode and the whole sequence for a resumed request. The trace classifies by scheduled token count instead (>1 = prefill chunk, ==1 = decode), which is what the attention profile axes expect and the only classification that survives a resumed request. Request.is_prefill() is deleted. - num_computed_tokens advances when the batch is formed, as in vLLM's _update_after_schedule, with Batch.scheduled_tokens as the snapshot add_done works from. Advancing at completion let pp_size > 1 schedule the same tokens twice. - preemption is vLLM verbatim, num_computed_tokens = 0 included. That is not re-prefill: free_blocks keeps the blocks' hashes, so on re-admission the still-resident prefix is found, a lower tier returns what was written down, and only the remainder is recomputed. Recovery comes from the tier hierarchy, not from a "preserve the decode state" special case. - schedule_base and schedule_with_prefix collapse into one, taking with them a duplicate shadowed definition of schedule_with_prefix (228 dead lines), _rollback_locks, _get_reload_size, the STEP 2/3 shrink search, the commit-time reclaim/shrink/preempt loop, and get_first_arrival_time (which read an attribute never assigned and had no callers). 1307 -> ~510 lines. Memory model - KV capacity is npu_mem * gpu_memory_utilization - weight, divided into blocks, mirroring vLLM's requested_memory - non_kv_cache_memory. New --gpu-memory-utilization flag (default 0.9) with a per-instance cluster-config override. vLLM also subtracts the activation peak and CUDA context, which are not modelled, so this capacity is an upper bound at the same value. - npu_used / cpu_used become properties over the pools, so there is one ledger per tier. The previous pair of ledgers (npu_used alongside RadixCache.capacity/total_memory_usage) is what let PR #59's mismatch through. - the three prefix-cache modes now map onto three real vLLM configurations: --no-enable-prefix-caching behaves like vLLM with prefix caching off, where a resumed request recomputes its whole sequence; --enable-prefix-caching is default vLLM; adding --prefix-storage CPU/CXL is vLLM with LMCache or OffloadingConnector attached. The previous middle case billed a KV transfer against a tier that held nothing. - a host offload tier now uses 256-token chunks (LMCache's default) uniformly. The non-shared second tier and the shared CXL pool used page size 1, which matched at token granularity and over-reported hits relative to any real offload tier. - 885 -> ~545 lines. Removed --prioritize-prefill and the per-instance prioritize_prefill key, along with _merge_by_arrival_id (its only caller). vLLM v0.19.0 has no equivalent: vllm/core/ -- the V0 scheduler whose _schedule_default served prefills first -- no longer exists, and SchedulerPolicy is fcfs or priority, which is request priority rather than prefill-vs-decode. The only committed config that set the key paired it with enable_chunked_prefill, where the branch could not fire anyway. P/D disaggregation: the KV transfer is charged in convert_prefill's per-layer SEND/RECV, but it was sized from the *v_proj layer's output_size, i.e. the whole QKV activation, so it shipped Q as well -- a factor of (q_dim + 2*kv_dim) / (2*kv_dim), 3x for Llama-3.1-8B -- and it ignored kv_cache_dtype, making --kv-cache-dtype fp8 6x high. The frontend now puts the per-layer, per-rank K+V bytes in the trace's comm_size column via Batch.pd_kv_send_tokens, which also counts a request's prefix-cache hit on its first step: the decode side needs that KV even though the prefill side read it from cache. The converter half of this lives in the chakra submodule and is not in this commit, so the new value is inert until that is bumped. Also fixed here: - a P/D handoff left num_tokens_reached un-advanced, so the decode instance received a request with nothing to schedule and the run never terminated. The prefill instance ran lm_head and the sampler, so the first token exists. A guard now raises instead of spinning if a running request has nothing to schedule and no batch in flight. - Total input tokens was reported as total_prompt - total_recompute, and those are not complements: a request preempted again mid-recompute is charged its full remaining work each time it is re-admitted. Summed from the requests' own original_input now. - the per-tick Running column counted the in-flight batch, which is only the subset that fit in the step's token budget. len(running) is the exact analogue of vLLM's num_running_reqs, which is what bench compares it against. Measured at 24 GiB / max_num_seqs 128 on the RTX 4090 ShareGPT replay (gpu_memory_utilization 1.0, to stay comparable with the old baselines): preemptions 20,540 -> 208, makespan 245.28 s -> 151.32 s, median TTFT 121,809 -> 49,078 ms, median TPOT 54.95 -> 45.32 ms, prefix hit ratio 1.91% -> 3.96%. That is issue #40; on a 24 GB card with an 8B bf16 model it fires at any load, not only under overload. Regressions: single_node_single_instance + example_trace stays bit-identical at 1,665,077,255 ns; rtx4090 128 GiB msq 256 moves -0.027% (106,668,561,596 -> 106,639,710,621), explained by the prefix hit ratio rising 1.91% -> 7.59% from the chained hashes and from generated-token blocks becoming recoverable; single_node_pd_instance, single_node_pd_per_instance_config and single_node_moe_pd_instance are all bit-identical. rtx4090 24 GiB --no-enable-prefix-caching is 6.7% slower on purpose (163.33 -> 174.23 s), since CPU swap is replaced by vLLM's full recompute, with 518,327 recomputed tokens accounting for it.
AGENTS.md: repository-structure table, and the "Scheduler and memory model" section rewritten around the two-phase schedule, the absence of a prefill/decode phase, schedule-time advancement of num_computed_tokens, and the three prefix-cache modes as three real vLLM configurations. Adds the trace-format note that comm_size on qkv_proj carries the P/D KV amount, plus four pitfalls that exist because each has already cost a wrong turn: do not reintroduce is_prefill(), do not special-case "preserve the decode state" on preemption, do not derive sequence length from num_computed_tokens, and note that kv_load/kv_evict fire only with --prefix-storage. docs site: prefix-caching.md largely rewritten (chained hashes, tiers sharing one key space, lookup/insert/evict flow, what preemption costs); continuous-batching.md's "two scheduling paths" replaced by the two phases and the uniform tokens-to-catch-up rule; plus the smaller references in architecture.mdx, request-lifecycle.md, reading-output.md, codebase-tour.md, validating-changes.md, jsonl-format.md, examples/memory-tiers/prefix-caching.mdx, and the two reference tables. No RadixCache / RadixAttention / --prioritize-prefill references remain. CHANGELOG: Added / Changed / Removed / Fixed entries for the redesign.
Re-ran all three examples on the block-pool scheduler and regenerated summary.txt plus the throughput / latency / requests plots. The plots move because the per-tick Running column now reports len(running) rather than the in-flight batch size. Accuracy against the recorded vLLM runs is preserved, and slightly better in most places. Total clocks move by -0.03% (Llama-3.1-8B), -0.004% (Qwen3-32B) and -0.14% (Qwen3-30B-A3B-Instruct-2507). Of the 45 vLLM-vs-sim Diff% figures across the three summaries, none moves by more than 0.1pp except on the MoE example, where TTFT P90 improves +3.7% -> +2.1% and P99 +4.7% -> +3.9%. These configs have 96 GiB of NPU memory, so the eviction path is essentially inert -- which is why the same defect that cost 20,540 preemptions on a 24 GB card was invisible here.
…ll-ISL gate Naming: --gpu-memory-utilization / gpu_memory_utilization becomes --npu-memory-utilization / npu_mem.mem_util. It was the only surface in the repo that said GPU -- the cluster config says npu_mem 40 times and num_npus 18 -- and the simulator models accelerators that are not GPUs. The flag help names vLLM's --gpu-memory-utilization so the correspondence is not lost. Placement: the per-instance override now lives inside npu_mem as mem_util, beside the mem_size it multiplies and matching its mem_* siblings. Validated at load: a non-numeric value or one outside (0, 1] is refused with a message saying it is a fraction, since mistyping 0.9 as 90 would silently size the cache 100x. Admission gate: --reserve-full-isl (on by default, per-instance reserve_full_isl) admits a request only if its whole sequence fits, not merely its first chunk. Port of vLLM's scheduler_reserve_full_isl, True there too, documented as preventing "over-admission and KV cache thrashing with chunked prefill". allocate_slots' block arithmetic is extracted into _num_blocks_to_allocate so the gate and the allocation cannot drift apart -- two estimates disagreeing was the radix tree's central defect. Startup output: a KV Cache Initialization section between the input config and the run loop, listing each instance's derived capacity: • Instance [0] : 54400 tokens / 3400 blocks (6.64 GiB/rank at util 0.90) The utilization fraction alone does not tell you where memory pressure will land; the token count does. It is per instance and only known once the schedulers exist, so it cannot be a row in the input-config block. Dropped that block's global utilization row, which would have lied as soon as one instance overrode it. New logger helper print_heading() centres a section title without drawing a second rule. "▶ Starting simulation..." moves to just before the loop. It used to print before config_builder, the prefix pools, the schedulers, the controller and the router were built. Docs: AGENTS.md gains the gate and npu_mem.mem_util; the KV-cache and memory page is rewritten around the block pool, the capacity formula and the gate; CONTRIBUTORS.md picks up #51, #53, #56, #57 and credits the two reports that drove this work (#40's scheduler analysis, and #58/#59, whose 24 GB profiling run was the first time the simulator saw real KV pressure); the CHANGELOG says explicitly what replaced the radix tree and that prefix caching's user-visible flags are unchanged. bench/examples/configs/*.json state mem_util: 0.9 explicitly. At 96 GiB those configs have 458k-617k tokens of KV, far past what the workload wants, so the value is inert there and all three examples reproduce their previous clocks exactly.
meta.json carried ten hand-picked engine kwargs, which is what we asked
vLLM for rather than what it ran with. Three things this investigation
needed were therefore unavailable: how many KV blocks vLLM actually
allocated, whether scheduler_reserve_full_isl was on, and which GPU the
run was on. The KV block count had to be inferred from the ratio of
running requests to token capacity.
Adds three keys, all optional and all read with meta.get() rather than
off schema_version, since they are absent from older runs and can be {}
if collection failed:
- kv_cache: num_gpu_blocks, block_size, num_kv_tokens,
gpu_memory_utilization. num_gpu_blocks is the number a simulator has to
match, and the only place the activation peak and CUDA context that
vLLM subtracts from its budget become visible -- every other term in
that budget is known up front.
- hardware: accelerator name, total memory, compute capability, CUDA and
torch versions. Enough to match a run against a profiler/perf/<hw>/
bundle without a separate host_metadata file.
- resolved_config: the whole VllmConfig, one key per sub-config, with
defaults filled in and inferred values settled. Built by walking the
config's own field list rather than a hand-kept list of interesting
knobs, so a vLLM upgrade that adds one shows up without a change here.
26 sections, ~296 leaf fields on v0.19.0.
Values that JSON cannot hold are replaced by a short type tag, so one HF
config object cannot bury the file -- meta.json goes from ~700 bytes to
~12 KB rather than megabytes. A property that raises while being read is
recorded as unreadable instead of aborting. The whole collection is
wrapped so that gathering metadata can never lose an otherwise finished
run.
META_SCHEMA_VERSION stays at 1: these are additions, so an existing
reader keeps working, and presence is better tested directly than
inferred from a version number. recorder.py documents all three, since
its module docstring already claims to be the single source of the
schema.
CHANGELOG and bench/README.md describe the three new meta.json keys and why num_gpu_blocks matters: every other term in vLLM's KV budget is known up front, so it is the only place the activation peak and CUDA context it subtracts become visible.
Completes the half-landed change from 2cf3f12. The frontend already writes the per-layer, per-rank K+V bytes into the trace's comm_size column on qkv_proj; the Chakra converter that reads it lives in a nested submodule, so until now the SEND/RECV pair still used the qkv_proj output size and shipped Q along with K and V -- 3x too much for Llama-3.1-8B, and blind to kv_cache_dtype. Verified on single_node_pd_instance: the trace carries comm_size=40960 against output_size=122880 (exactly 3.0x), and the generated graphs hold 32 COMM_SEND_NODE of 40960 B on the prefill NPU against 32 COMM_RECV_NODE of the same size and tags on the decode NPU, with no COMM_COLL_NODE introduced -- collectives are gated on comm_type, which stays NONE. All five reference scenarios stay bit-identical (single_node_pd_instance, single_node_pd_per_instance_config, single_node_moe_pd_instance, single_node_heterogeneous, single_node_single_instance): example_trace's 10-token prompts put the KV send 3 orders of magnitude below the layer compute, so the clock cannot see it either way. A realistic-prompt workload is what will show the difference.
_axis_bracket blended two profiled samples on a log scale. The profiler sweeps every axis geometrically, which makes log space look like the matching choice, but those are separate decisions: the grid spacing decides where the kernel is sampled, the blend decides how two samples are combined, and the kernel is linear in each axis. Profiled decode attention depends on the total KV read, not on how it is spread across requests. Fitting time_us = a + b * (n_decode * kv_decode) on the RTX 4090 Llama-3.1-8B grid gives R^2 = 1.0000 at an implied 953 GB/s -- 95% of the card's 1008 GB/s spec, i.e. a pure KV-bandwidth stream. At fixed n_decode * kv_decode, sweeping n_decode from 8 to 256 moves the measured time by only ~7%. Blending a per-axis-linear function in log space is convex-biased upward: up to +6.0% per axis on a doubling grid (worst at query/x0 = 1/ln2 = 1.443), compounding across axes. Leave-one-out over every profiled attention row -- predict a grid point from its two neighbours, compare against what the GPU actually reported, no fitted model anywhere -- puts log space at +11.6% to +14.4% mean error across all seven bundles in the repo against +2.3% to +3.7% for linear, with linear ahead on all four axes: axis log linear prefill_chunk +12.59% +3.81% kv_prefill +6.67% +0.23% n_decode +18.44% +4.31% kv_decode +11.83% +1.05% At real batch shapes the runtime bracket spans 2x rather than the LOO test's 4x, giving log/linear = 1.091 on attention, i.e. +3.0% of a decode iteration. End-to-end against the three real vLLM runs in bench/examples, mean |Diff%| over 15 metrics each: TPOT improves on all 15 of 15 metrics (1.54 -> 0.84, 1.82 -> 1.38, 1.78 -> 0.76) and Latency on 13 of 15 (0.74 -> 0.68, 2.02 -> 1.42, 0.88 -> 0.20). TTFT splits 8 improved / 6 worse: the change makes every iteration cheaper, so a run that was already predicted too fast gets faster. Qwen3-32B, the one example with no compensating error, improves on all 15 metrics. Reference scenarios move by ~0.02%: single_node_single_instance 1,665,077,255 -> 1,664,699,047, single_node_pd_instance and single_node_pd_per_instance_config 1,662,668,674 -> 1,662,366,658, single_node_moe_pd_instance 2,082,386,488 -> 2,081,832,520, single_node_heterogeneous 3,003,118,275 -> 3,002,576,771. example_trace's 10-token prompts sit at the bottom of the grid where the two bases nearly coincide.
_ATTN_SKEW_ALPHA_FALLBACK was 0.093, a constant no bundle in the repo reproduces. It only applies when a profile carries no skew_fit block at all -- today that is RTX4090, profiled with SKIP_SKEW=1; bundles with a real fit resolve alpha per bucket and are unaffected. Guessing that constant is worse than omitting the correction: - The measured pooled value for this same model on RTXPRO6000 is 0.0543, and resolving a saturated RTX 4090 run's own 3,853 decode batches against that bucket table gives alpha p50 0.059 (min -0.004, max 0.068). 0.093 is ~1.6x too high. - The parameter is badly conditioned. The endpoint gap (t_max - t_mean) * num_layers is ~12.6 ms on a ~29 ms iteration, so each 0.1 of alpha is ~4.3% of iteration time and alpha would have to be known to +/-0.023 to keep attention within 1%. - alpha is not bounded to [0, 1] either. Across the 12,984 raw shots in RTXPRO6000/meta-llama/Llama-3.1-8B/bf16/tp1/skew.csv, where t_mean, t_max and t_skew are each measured directly on the GPU, alpha spans -68.6 to +9.28 with 16.6% negative and only 80.8% inside [0, 1]. alpha=0 returns t_mean, which is also the physically correct anchor: attention cost tracks the total KV read Sigma_k, and a uniform batch at the arithmetic mean has n * mean(k) = Sigma_k exactly. The median would understate it -- the runtime kv distribution is right-skewed, measured kv_max/kv_mean p50 = 2.61 on the ShareGPT replay. The profiler uses the same definition (skew.py: kv_mean = total_kv // n). Measured t_skew/t_mean at 4090-like shapes is 1.029, so the correction being dropped is worth ~3%. Also records why kv_decode_mean is the lookup coordinate, so the next reader does not reach for a median. On the RTX 4090 ShareGPT replay at the default mem_util 0.9 this moves clocks 168,187,564,812 -> 163,114,307,308 ns and TPOT p50 36.85 -> 35.66 ms. The five reference scenarios are unchanged: they all run on RTXPRO6000, which has a real skew fit, so the fallback never fires.
Regenerated with ./bench/examples/run.sh followed by
./bench/examples/validate.sh, so paths and flags match the committed
layout. Clocks move by about a percent, all in the same direction --
the linear basis makes every attention lookup cheaper:
Llama-3.1-8B 65,705,603,364 -> 65,049,871,040 (-1.00%)
Qwen3-32B 163,422,751,517 -> 162,573,714,653 (-0.52%)
Qwen3-30B-A3B 62,942,560,948 -> 62,060,744,672 (-1.40%)
Accuracy against the recorded vLLM runs, mean |Diff%| over each
example's 15 metrics:
before after
Llama-3.1-8B 1.18% 1.47%
Qwen3-32B 2.03% 1.49%
Qwen3-30B-A3B 3.06% 2.56%
TPOT improves on all 15 of 15 metrics across the three examples and
Latency on 13 of 15; TTFT splits 8 improved / 6 worse. A cheaper
iteration helps wherever the simulator was predicting too slow and hurts
where it was already too fast, so the sign of each example's existing
bias decides the outcome. Qwen3-32B, whose 15 metrics were uniformly
+1.4% to +2.8% before, improves on every one of them.
Two figures unchanged by this and worth chasing separately: the MoE
example's TTFT Median (-19.9%, identical before and after) and its
makespan (-9.0%). They account for most of that example's remaining
error -- Llama and Qwen3-32B now peak at 7.0% and 2.3%.
Updates AGENTS.md, the docs site, profiler/README.md and the changelog for the two model changes, and corrects three descriptions that were already wrong before them: - The attention lookup was documented as "nearest-neighbour on (prefill_chunk, n_decode), bilinear on (kv_prefill, kv_decode)" in both AGENTS.md and the trace-generation page. It has always bracketed and interpolated all four axes; only the scale was ever in question. - The trace file was called tab-separated. utils.py::_FMT emits fixed-width space-padded columns, and the converter splits on any whitespace. Parsing on '\t' silently yields one field per line. - The SKIP_SKEW=1 fallback alpha was documented as "roughly 0.3 across observed hardware", a figure no bundle in the repo reproduces. The code constant was 0.093 and is now 0. Also adds the gitignore rules AGENTS.md already claimed existed: outputs/* with !outputs/example_*.csv, plus astra-sim/inputs/runs/. Without them every run's scratch accumulated in git status, and --no-cleanup-inputs left 1.3 GB of ASTRA-Sim inputs behind. The AGENTS.md entry for _axis_bracket carries the evidence inline, so the next reader does not "fix" the linear blend back to log space on the grounds that the sweep grid is geometric. Those are separate choices. Site builds clean (cd docs && pnpm build).
The model-architecture page did not match the YAML schema. `LayerEntry`
sets `extra="forbid"`, so the fields it documented are rejected at load
time:
cls: -> vllm:
category: dense -> the catalog block the layer sits in
tp_collective: -> not a field; the simulator attaches TP ALLREDUCE
after o_proj / down_proj from the cluster config
ep_collective: -> likewise for EP ALLTOALL around moe
(missing) -> within:
`within` filters on an ancestor class, which is what lets one `RMSNorm`
entry serve the input norm, the post-attention norm and the final norm:
`within: LlamaDecoderLayer` for the two block-level ones,
`within: LlamaForCausalLM` for the last. The same filter separates
`o_proj` from `down_proj`, both `RowParallelLinear`. The page now also
records that `(vllm, within)` has to be globally unique.
Further corrections there: `attention` is listed explicitly in
`sequence.pre_attn`; `lm_head` maps to `LogitsProcessor`; Llama 3 uses
`Llama3RotaryEmbedding`; the MoE catalog names the sparse block
(`Qwen3MoeSparseMoeBlock`); and both dense and MoE models declare all six
sequence groups with the unused MLP group empty. The example is checked
against `Architecture` / `Catalog` / `Sequence`, and every layer its
sequence references exists in its catalog.
Same class of drift elsewhere:
- `_lookup_attention_with_skew` was described as always doing two 4D
lookups (AGENTS.md, serving/README.md, its docstring,
trace-generation.md, skew-alpha-fit.md). The second lookup is
conditional: `t_mean` is returned after one lookup for `n_decode <= 1`,
for a batch whose decode kv lengths are all equal, or for `alpha == 0`,
which is the default without a skew profile.
- `output-bundle.md` described the attention lookup as nearest-neighbour
plus bilinear, and `_attn_slice_lookup`'s docstring said "log on each
axis".
- `codebase-tour.md` pointed at `astra-sim/astra-sim/system/Workload.cc`;
the file is under `workload/`.
- Two example pages pointed at `configs/pim/DDR4_8GB_3200_pim/` as a
directory. `configs/pim/*.ini` are files, one per device.
Drops the reasoning that landed on trace-generation.md alongside the
linear interpolation change: why the blend is linear, the range of the
fitted alpha, and the arithmetic-mean lookup coordinate. That belongs
with the code, in `_axis_bracket`'s docstring and AGENTS.md.
Adds @bui-thanh-lam to CONTRIBUTORS for the report, and brings the
report-only entries to one line each to match the rest of the list.
Checked mechanically and clean: canonical layer names against every
`profiler/models/*.yaml` catalog, cluster-config keys against
`config_builder.py`, the 11-column trace header against `utils.py::_FMT`,
model-config fields against what the simulator reads, and documented CLI
defaults against argparse.
Site builds clean (cd docs && pnpm build).
The committed outputs/example_*.csv were last regenerated in v1.1.0 (3012eb1) and had gone stale against `main`, not just against this branch: re-running `main` today produces different numbers, and the committed files disagree with the dataset they claim to replay. `workloads/example_trace.jsonl` gives request 3 `output_toks: 12` and request 0 `output_toks: 70`; the committed CSVs recorded 16 and 80. Cause: `serving/run.sh` has ten of its eleven examples commented out, so only `single` was ever refreshed. The others drifted unnoticed. That is the intended shape of the file -- a menu you uncomment one line of -- so it stays as it is; the eleven commands were checked flag by flag against what was actually run here, and they match. Its shebang was commented out too (`# #!/bin/bash`) while the file carried the execute bit, so invoking it directly fell through to `sh`. All eleven scenarios were run on this branch and on `origin/main` for comparison. Every one completes with no traceback, and the branch differs from `main` only by the size of the linear-interpolation change: scenario origin/main branch diff cxl 12,113,387,362 12,113,387,362 +0.0000% single 1,665,077,255 1,664,699,047 -0.0227% power 1,665,077,255 1,664,699,047 -0.0227% multi 1,665,607,238 1,665,338,118 -0.0162% prefix_cpu_mem_pool 1,665,607,238 1,665,338,118 -0.0162% dual_prefix_cpu_mem_pool 1,669,378,488 1,669,153,368 -0.0135% pd 1,662,668,674 1,662,366,658 -0.0182% pim 1,643,260,840 1,643,260,424 -0.0000% pim_sub_batch 2,634,146,962 2,634,146,546 -0.0000% moe 1,946,195,383 1,945,725,751 -0.0241% moe_dp_ep 36,636,490,356 36,584,504,388 -0.1419% `cxl` is unchanged because it is CXL-bandwidth bound, and the two PIM cases because PIM attention replaces the NPU attention lookup the change touches. `moe_dp_ep` moves furthest simply by running the most iterations: 52 ms over a 36.6 s simulation. Also removes four tracked outputs that nothing regenerates: - `example_weight_offload_run.csv` and `example_prefix_cpu_mem_run.csv` have no command in `run.sh` at all, so they could not be reproduced or checked. - `example_ns3_run.csv` and `example_prefix_run.csv` belong to the commands under run.sh's "Deprecated examples" heading -- NS-3 integration is still a work in progress, and the prefix-caching flag became the default, which is why that example was disabled. What is left is eleven tracked outputs for the eleven scenarios that were actually verified.
This was referenced Aug 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replaces the radix-tree KV allocator with a port of vLLM v0.19.0's block pool,
and aligns the scheduler with vLLM V1's two-phase
schedule(). Fixes #40 and#52. Also corrects the attention latency lookup, which was interpolating on the
wrong scale.
Changes
KV cache & scheduler
serving/core/block_pool.py,kv_cache_manager.py: per-tier free-blockqueues, chained block hashes (
hash(parent_hash, block_tokens)), eviction asa side effect of allocation.
serving/core/radix_tree.pydeleted.schedule()for prefix caching on and off.scheduler.py1300 → ~510lines,
memory_model.py885 → ~545.--npu-memory-utilization(default0.9, per-instancenpu_mem.mem_util),--reserve-full-isl(on, mirrors vLLM'sscheduler_reserve_full_isl).--prioritize-prefill(no equivalent in vLLM).Attention latency model
_axis_bracketnow blends linearly, not in log space. The profiler sweepsgeometrically, but the kernel is linear per axis — profiled decode attention
fits
a + b·(n_decode·kv_decode)at R² = 1.0000. Leave-one-out on measuredgrid points: log-space +11.6…+14.4% mean error vs +2.3…+3.7% for linear,
across all 7 bundles in
profiler/perf/._ATTN_SKEW_ALPHA_FALLBACK0.093 → 0. It applies only to bundles with noskew_fit, and 0.093 is a constant no bundle reproduces.P/D
trace's
comm_sizecolumn instead of the qkv_proj output size, which shippedQ as well — 3× too large for Llama-3.1-8B, and blind to
kv_cache_dtype.bench
meta.jsonrecords vLLM's resolved config, includingkv_cache.num_gpu_blocks.kv_cache_pctalways logging 0.0 intimeseries.csv.Docs
adding-model-architecturedocumented four fields that do not exist(
cls,category,tp_collective,ep_collective) and omittedwithin;its example failed pydantic validation. Rewritten and now checked against the
schema. (Outdated model YAML format #52)
the
Workload.ccpath, the PIM config paths.outputs/*andastra-sim/inputs/runs/are gitignored, whichAGENTS.mdalready claimed.
Validation
All 11 example scenarios run clean (single / multi / P/D / CXL / prefix CPU pool
×2 / power / PIM / PIM sub-batch / MoE / MoE DP+EP agentic), no tracebacks.
Branch vs
main, worst case −0.142% (moe_dp_ep), rest within −0.025%.Against real vLLM runs, mean |Diff%| over 15 metrics:
bench/examples/Llama-3.1-8Bbench/examples/Qwen3-32Bbench/examples/Qwen3-30B-A3BThe bench examples run on a 96 GB card at 17–36% KV occupancy, so they were
already accurate and barely move. The case under real KV pressure is the RTX 4090
24 GB run from #58/#59 — that is #40's regime, and where this PR is aimed. #59
reported TTFT mean −87.5% / TPOT mean +173.5% there. Now:
npu_mem.mem_utilThe gap between those two rows is the known limitation below, not a tuning knob
worth shipping: 0.87 is standing in for ~0.72 GiB the capacity model does not
account for.
outputs/example_*.csvare regenerated. Four with no generator are removed(
weight_offload,prefix_cpu_mem,ns3,prefix); they had been stale sincev1.1.0 and disagreed with the dataset they claim to replay.
Known limitation
KV capacity is
mem_size · mem_util − weight. vLLM also subtracts itsactivation peak and CUDA context, which are not modelled, so capacity is an
upper bound — on the 4090 that is ~0.72 GiB, and it is what the residual 7.75%
is made of.
meta.json::kv_cache.num_gpu_blocksnow records the real number, soone bench run on any card turns that estimate into a measurement. Follow-up, not
a blocker.