Expand description
KV-cache memory estimation.
[estimate_kv_cache] is a top-level crate function (NOT a method on
AicEngine): estimation runs once at startup, uses overlapping but not
identical inputs to AicEngineBuilder, and is a separate concern from
latency prediction. The Dynamo Mocker is the primary external consumer; it
calls this once and derives num_gpu_blocks_per_rank from
total_kv_size_tokens.
§Rust is a pure forwarder; the estimate is computed in Python
This mirrors how AicEngineBuilder forwards to Python’s compile_engine:
ALL of the work – fraction + tolerance validation, HF-config parsing, the
AIC backend memory model, the OfFree/OfTotal budget math, the naive heuristic
fallback, AND the tolerance margin – lives in
aiconfigurator.sdk.memory.estimate_kv_cache. The Rust side:
- crosses into Python once (
with_gil → import → call estimate_kv_cache → extract dict), forwardingtolerance_fractionthrough; - rebuilds a
KvCacheEstimatefrom that dict (includingtolerance_adjusted), with no math of its own.
The two budget formulas (TRT-LLM free-fraction vs vLLM/SGLang total-fraction),
the naive fallback, and the tolerance margin all live on the Python side; see
the docstring of aiconfigurator.sdk.memory.estimate_kv_cache. The
KvCacheMemoryFraction enum still encodes the backend↔fraction XOR so the
request shape is unambiguous; the variant is validated against
engine.backend in Python.
Structs§
- KvCache
Estimate - KV-cache memory estimate.
- KvCache
Estimate Adjusted - KvCache
Estimate Options - KvCache
Estimate Request - KV-cache memory request.
enginereuses the modularisedEngineConfig; the remaining fields describe the runtime sizing budget. - Memory
Breakdown - Non-KV memory components, in bytes. Maps AIC’s
_get_memory_usagedict:weights → weights,activations → activations,others → runtime_overhead,nccl → comm_overhead.
Enums§
- Estimate
Source - KvCache
Estimate Error - KvCache
Memory Fraction - Backend-tagged memory fraction. The variant encodes the XOR between
TRT-LLM’s free-fraction and vLLM/SGLang’s total-fraction semantics; the
Python
estimate_kv_cachevalidates it againstengine.backendand returns an error (mapped toKvCacheEstimateError::IncompatibleMemoryFraction) if mismatched.