[Observability] Break down prefill #cached-token into device / host / storage tiers

[hxieustc]

Motivation

The prefill log reports a single aggregate #cached-token for prefix-cache hits. With HiCache (L1 GPU / L2 host / L3 storage), this hides where hits come from — so you can't tell from the logs whether load-back from L2 or storage reuse is happening at all.

Modifications

Split cached-token accounting in PrefillAdder into three tiers and surface them in the scheduler log:

• New counters log_hit_tokens_device / _host / _storage (aggregate log_hit_tokens kept).
• In add_one_req, the hit is split by source before L2 load-back mutates req.prefix_indices: device = current prefix_indices; host/storage from req.host_hit_length / req.storage_hit_length. Kept mutually exclusive (device + host + storage == total), warning on mismatch; per-request cached_tokens_{device,host,storage} recorded.
• _update_prefill_budget / _add_dllm_req take optional keyword splits; default (device_hit_tokens=None means all device) keeps existing call sites correct.
• PrefillStats carries the new fields; the log line emits #cached-token-device/-host/-storage (storage shown only when enabled or non-zero).

How to enable

No flag required — the device/host tier split is emitted in the existing prefill log line automatically. The values are only non-trivial when the corresponding tiers exist.

• The #cached-token-device / #cached-token-host fields always appear in the prefill log (visible at the default INFO log level).
• #cached-token-storage is appended only when storage is enabled (enable_hicache_storage) or a non-zero storage hit is observed.

Benefits

• Direct visibility into HiCache hit distribution across tiers, right in the existing prefill log line — no extra tooling.

Expected output

Prefill log line — without L3 storage

For example, HiCache L1/L2, 2048-token prefix served partly from host after load-back:

Before:

Prefill batch, #new-seq: 1, #new-token: 512, #cached-token: 2048, token usage: 0.06, #running-req: 1, #queue-req: 0, #pending-token: 0

After:

Prefill batch, #new-seq: 1, #new-token: 512, #cached-token: 2048, #cached-token-device: 1536, #cached-token-host: 512, token usage: 0.06, #running-req: 1, #queue-req: 0, #pending-token: 0

Prefill log line — with L3 storage enabled (#cached-token-storage appended):

Before:

Prefill batch, #new-seq: 1, #new-token: 256, #cached-token: 4096, token usage: 0.11, #running-req: 1, #queue-req: 0, #pending-token: 0

After:

Prefill batch, #new-seq: 1, #new-token: 256, #cached-token: 4096, #cached-token-device: 1024, #cached-token-host: 2048, #cached-token-storage: 1024, token usage: 0.11, #running-req: 1, #queue-req: 0, #pending-token: 0

Non-HiCache (backward compatibility)

host/storage are 0, device equals the old aggregate:

Prefill batch, #new-seq: 1, #new-token: 512, #cached-token: 1536, #cached-token-device: 1536, #cached-token-host: 0, token usage: 0.05, #running-req: 1, #queue-req: 0, #pending-token: 0

Correctness

• Backward compatible by construction: device_hit_tokens defaults to full prefix length, so without HiCache device == old aggregate, host/storage == 0.
• getattr(..., default) fallbacks on req and adder keep older/partial paths valid.
• device + host + storage == total invariant logs a warning instead of failing.

Caveats

• Logging/accounting only — no scheduling, memory, or output change.
• Relies on req.host_hit_length / req.storage_hit_length being populated; backends that don't set them report all hits as device (same as today).
• storage/host exclusivity uses min(storage, host) (common case: storage materialized into host first); exotic tiering may need refinement.

Accuracy Tests

It does not affect model outputs.

Speed Tests and Profiling

Negligible overhead.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #28265459028
Latest PR Test (Extra): ❌ Run #28265458955

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ishandhanani]

Tested this branch with HiCache enabled (--enable-hierarchical-cache --hicache-write-policy write_through, page-size 64, a small KV pool so prefixes get evicted to host and then loaded back).

The device/host split itself is correct. On a host-tier hit (1408-token prefix served entirely from host):

Prefill batch, #new-token: 64, #cached-token: 1408, #cached-token-device: 0, #cached-token-host: 1408, ...

and after the block is promoted back to device, a re-send flips to #cached-token-device: 1408, #cached-token-host: 0. No-hit traffic reports all zeros, so it's backward compatible.

Two things noted inline: a WARN that fires on every host hit, and a breakdown that's already computed elsewhere. Also one stray unrelated file. Thanks for adding this — the tiering is genuinely useful for HiCache tuning.

[ishandhanani]

This invariant fires on essentially every request with a host-tier hit. At this point init_load_back hasn't run yet, so prefix_len = len(req.prefix_indices) is the device-only count, while host_hit_tokens comes from req.host_hit_length. When the prefix is served entirely from host, prefix_len == 0 but host_hit_tokens > 0, so device + host + storage != prefix_len.

Observed (1408-token prefix fully evicted to host):

WARN schedule_policy.add_one_req: Cached-token split mismatch: total=0, device=0, host=1408, storage=0, rid=...

The split values themselves are right (the prefill log line correctly shows device=0/host=1408) — it's the invariant baseline that's off, and the total=0 in the message is misleading. Suggest comparing against prefix_len + host_hit_length (the post-load-back total), or running the check after init_load_back. As is, it emits a WARN on every load-back.

[hxieustc]

Good catch. The warning baseline was wrong because this block runs before init_load_back, so len(req.prefix_indices) is still device-only while host_hit_length already includes host-tier hits. I removed this pre-load-back split/check and now rely on the post-load-back request-level accounting instead.

[ishandhanani]

Heads up: the cached-token accounting in schedule_batch.py already derives the same per-tier split after load-back (device_portion = len(prefix_indices) - host_hit_length) and writes req.cached_tokens_device/host/storage. This block recomputes it pre-load-back and those fields get overwritten there later. Might be cleaner to compute it once and reuse, to avoid two sources of truth that can drift.

[hxieustc]

Agreed. The main branch had evolved and schedule_batch.py already computes req.cached_tokens_device/host/storage after load-back, which is the right point in the flow. I removed the recomputation in schedule_policy.py and changed PrefillStats.from_adder() to aggregate from those request fields after prepare_for_extend().

[ishandhanani]

Minor: inserting the defaulted log_hit_tokens_* fields forces defaults onto the fields below, and num_running_reqs: QueueCount = None is a None default on a non-Optional type. It only works because from_adder always passes it — a little fragile if another constructor path is added.

[hxieustc]

Thanks, fixed. I reordered PrefillStats so the required fields stay non-default (new_token_ratio, num_running_reqs, num_new_seqs) and the new tier fields remain optional defaults afterward. This avoids the QueueCount = None workaround.

[ishandhanani]

Re-tested the updated branch (77394e6cf, Qwen3-0.6B on 1x L40S, HiCache + write_through + page_size 64). Everything flagged in the previous review checks out now:

• No more Cached-token split mismatch warnings — re-ran the evict-to-host scenario that previously fired it on every host hit: zero warnings.
• Host-tier attribution correct: #cached-token: 1536, #cached-token-device: 0, #cached-token-host: 1536 on load-back, flipping to device: 1536, host: 0 after promotion. device + host == total in all cases.
• Split now reads the canonical req.cached_tokens_* breakdown, PrefillStats field defaults are clean, and the docs file is out of the diff.

One remaining bug: with chunked prefill, the tier split is re-counted on every continuation chunk — chunks 2..N log #cached-token: 0 but repeat the full #cached-token-device value, breaking the log_hit_tokens == device + host + storage invariant documented on PrefillStats. Any prompt longer than chunked_prefill_size (default 4096) hits this.

Fix + before/after evidence: hxieustc#1. With it applied, the aggregate over a full load-test log holds exactly: sum #cached-token = 133952 == sum device 121536 + sum host 12416.

[ishandhanani]

Re-tested the latest revision (ca5c53bac, Qwen3-0.6B on 1x L40S, HiCache + write_through + page_size 64). The chunked-prefill re-count from the previous round is fixed, and the cleaner adder-side accumulation (_accumulate_hit_token_split guarded on prefix_len > 0, with compute_cache_hit_split summing to prefix_len by construction) holds up empirically:

• Chunked prefill (--chunked-prefill-size 512, 1344-token cached prefix): chunk 1 logs #cached-token: 1344, #cached-token-device: 1344; continuation chunks log 0 / device 0 — no more per-chunk re-count.
• Tier attribution: host load-back #cached-token: 1536, device: 0, host: 1536, flipping to device: 1536, host: 0 after promotion. Using loaded_host_hit_length (actual restored tokens) is the right call.
• Invariant under load (aiperf c=16, 256 reqs, 1024-token prompts vs 512 chunk size, live host hits): aggregated over the full server log, sum #cached-token = 133952 == device 122560 + host 11392. Holds exactly.
• No Cached-token split mismatch warnings anywhere.

LGTM. One non-blocking suggestion: compute_cache_hit_split and the adder accounting are pure arithmetic — a small CPU unit test asserting log_hit_tokens == device + host + storage (including the chunked continuation case) would lock in the invariant cheaply, since there's currently no test coverage for the split.

[nvpohanh]

/tag-and-rerun-ci

[hxieustc]

/tag-and-rerun-ci

[hxieustc]

/tag-and-rerun-ci

[nvpohanh]

/tag-and-rerun-ci