Size CP HiCache host budget from compact index layers
The CP shared target+draft HiCache budget was still estimating NSA target bytes per token as if every target layer had an index cache. After index skip compacts L1 and L2 index buffers to active layers, that stale estimate kept host/L2 token capacity artificially low even though the physical host index allocation was already smaller. Use the pool's index_active_layer_ids when estimating NSA HiCache bytes per token, with the existing full-layer fallback for pools that do not expose compact metadata. This keeps the shared target+draft capacity computation consistent with the actual L2 host layout. Constraint: Index skip compacts NSA index cache by active logical layers, but MLA KV remains full-layer. Rejected: Keep full-layer HiCache estimate | preserves correctness but wastes the L2 capacity recovered by compact index allocation. Confidence: high Scope-risk: narrow Directive: Keep HiCache byte estimates aligned with NSATokenToKVPool allocation metadata; otherwise target+draft token budgets will drift from real host memory use. Tested: Remote container targeted pytest for active-index HiCache estimate and shared-budget tests -> 3 passed. Tested: Remote container P3-P6 regression suite plus new capacity tests -> 203 passed, 2 subtests passed. Not-tested: Full ETE replay/GSM8K after this budget correction. Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
@@ -75,10 +75,17 @@ def _estimate_hicache_size_per_token(kv_cache) -> int:
|
||||
kv_cache.index_head_dim
|
||||
+ kv_cache.index_head_dim // kv_cache.quant_block_size * 4
|
||||
)
|
||||
index_layer_count = len(
|
||||
getattr(
|
||||
kv_cache,
|
||||
"index_active_layer_ids",
|
||||
range(kv_cache.layer_num),
|
||||
)
|
||||
)
|
||||
return (
|
||||
kv_cache.kv_cache_dim * dtype_size * kv_cache.layer_num
|
||||
+ indexer_size_per_token
|
||||
* kv_cache.layer_num
|
||||
* index_layer_count
|
||||
* NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize
|
||||
)
|
||||
if isinstance(kv_cache, MLATokenToKVPool):
|
||||
|
||||
Reference in New Issue
Block a user