Skip to main content

Module memory

Module memory 

Source
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:

  1. crosses into Python once (with_gil → import → call estimate_kv_cache → extract dict), forwarding tolerance_fraction through;
  2. rebuilds a KvCacheEstimate from that dict (including tolerance_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§

KvCacheEstimate
KV-cache memory estimate.
KvCacheEstimateAdjusted
KvCacheEstimateOptions
KvCacheEstimateRequest
KV-cache memory request. engine reuses the modularised EngineConfig; the remaining fields describe the runtime sizing budget.
MemoryBreakdown
Non-KV memory components, in bytes. Maps AIC’s _get_memory_usage dict: weights → weights, activations → activations, others → runtime_overhead, nccl → comm_overhead.

Enums§

EstimateSource
KvCacheEstimateError
KvCacheMemoryFraction
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_cache validates it against engine.backend and returns an error (mapped to KvCacheEstimateError::IncompatibleMemoryFraction) if mismatched.