pub struct Engine { /* private fields */ }Expand description
Compiled engine: precompiled op lists + the matching perf database.
Built from an EngineSpec (Python’s compile_engine output) plus a
loaded [PerfDatabase]. Holds only the scalars the static composition
reads: the two op lists and nextn (the MTP decode-batch multiplier).
Parallelism / quant scalars do not enter the latency sum — they drive
throughput and memory, which StaticResult omits — so they are not stored.
Implementations§
Source§impl Engine
impl Engine
Sourcepub fn build(
spec: EngineSpec,
db: Arc<PerfDatabase>,
) -> Result<Engine, AicError>
pub fn build( spec: EngineSpec, db: Arc<PerfDatabase>, ) -> Result<Engine, AicError>
Build an Engine from a spec and a pre-loaded database.
Extracts the op lists and the nextn scalar from spec.engine. The
caller (AicEngineBuilder / from_spec_bytes) is responsible for
having loaded the matching PerfDatabase from spec.engine’s identity.
Sourcepub fn from_spec_bytes(
bytes: &[u8],
systems_root: &Path,
) -> Result<Engine, AicError>
pub fn from_spec_bytes( bytes: &[u8], systems_root: &Path, ) -> Result<Engine, AicError>
Convenience constructor: deserialize a bincode EngineSpec and load the
matching PerfDatabase from its identity, then Engine::build.
Runs the Engine::from_spec_bytes(bytes) + PerfDatabase::load
flow. systems_root points at python/aisimulate/src/aiconfigurator_core/systems and is used
only as a fallback: when the decoded spec.engine.systems_path is
Some, that path is authoritative and overrides the systems_root
argument.
Sourcepub fn reset_provenance(&self)
pub fn reset_provenance(&self)
Clear the empirical-provenance accumulator (start of a run). The PyO3
boundary calls this at the top of every compute method so
Self::last_provenance carries per-call semantics, mirroring
Python’s capture_provenance() scope. Deliberately NOT called inside
run_static itself: mixed_step_latency composes multiple internal
passes whose tiers must accumulate into one answer.
Sourcepub fn last_provenance(&self) -> Option<&'static str>
pub fn last_provenance(&self) -> Option<&'static str>
The least-confident empirical tier fired since the last
Self::reset_provenance, as the Python tag string; None when the
run was answered purely from silicon tables (nothing to note — Python’s
note_provenance is skipped for silicon too).
Sourcepub fn run_static(
&self,
runtime: &RuntimeConfig,
mode: StaticMode,
stride: u32,
) -> Result<StaticResult, AicError>
pub fn run_static( &self, runtime: &RuntimeConfig, mode: StaticMode, stride: u32, ) -> Result<StaticResult, AicError>
Python run_static / run_static_latency_only (base_backend.py:347,
:322) restricted to the latency breakdown. Dispatches on mode the
way _run_static_breakdown does and sums context + generation.
Sourcepub fn predict_prefill_latency(
&self,
bs: u32,
isl: u32,
prefix: u32,
) -> Result<f64, AicError>
pub fn predict_prefill_latency( &self, bs: u32, isl: u32, prefix: u32, ) -> Result<f64, AicError>
Mocker H1: prefill-step latency in ms. Pure-Rust inherent method (no
PyO3 py token), so the Mocker hot path runs without acquiring the GIL.
Thin shim over Self::run_static with mode=Context (osl is
irrelevant for the context phase, so it is fixed at 1).
Sourcepub fn predict_decode_latency(
&self,
bs: u32,
isl: u32,
osl: u32,
) -> Result<f64, AicError>
pub fn predict_decode_latency( &self, bs: u32, isl: u32, osl: u32, ) -> Result<f64, AicError>
Mocker H2: decode-step latency in ms. Pure-Rust inherent method (no
PyO3 py token). Thin shim over Self::run_static with
mode=Generation. Mocker passes osl=2 (one decode step at
s = isl + 1).
Sourcepub fn predict_decode_latency_total(
&self,
batch_size: u32,
total_past_kv_tokens: u32,
) -> Result<f64, AicError>
pub fn predict_decode_latency_total( &self, batch_size: u32, total_past_kv_tokens: u32, ) -> Result<f64, AicError>
Predict one decode step from exact FPM iteration totals.
total_past_kv_tokens excludes the one current token processed by each
decode request, matching the collector’s total_kv_read_tokens axis.
Sourcepub fn fpm_decode_kv_ceiling(&self) -> Result<Option<u32>, AicError>
pub fn fpm_decode_kv_ceiling(&self) -> Result<Option<u32>, AicError>
Highest decode KV-read total covered by a compiled FPM engine.
Op-level engines return None.
Sourcepub fn mixed_step_latency(
&self,
ctx_tokens: u32,
gen_tokens: u32,
isl: u32,
osl: u32,
prefix: u32,
seq_imbalance_correction_scale: f64,
gen_seq_imbalance_correction_scale: f64,
) -> Result<f64, AicError>
pub fn mixed_step_latency( &self, ctx_tokens: u32, gen_tokens: u32, isl: u32, osl: u32, prefix: u32, seq_imbalance_correction_scale: f64, gen_seq_imbalance_correction_scale: f64, ) -> Result<f64, AicError>
One mixed (chunked-prefill + decode) step latency. LITERAL mirror of
Python _get_mix_step_latency / run_mixed, which composes three
filtered phase passes (_run_context_phase / _run_generation_phase
with op_filter) that query ONLY the ops each pass consumes — the
same name-keyed sets the ContextOpFilter /
only_generation_attention walks below visit (issue #1498
follow-through: Python used to run the full lists and discard, so a
raise in a discarded query was a one-sided error surface):
// Pass 1 — combined non-attention work:
// run_static(batch=1, isl=ctx+gen, osl=1,
// prefix=prefix*floor(ctx/isl), mode=static_ctx)
// sum every op EXCEPT "context_attention"
// Pass 2 — context attention at the prefill shape:
// run_static(batch=ceil(ctx/isl), isl=isl, osl=1, prefix=prefix)
// take ONLY "context_attention", divide by ceil(isl/ctx)
// Pass 3 — decode attention (only when gen_tokens > 0):
// run_static(batch=gen, isl=isl+osl//2, osl=2, mode=static_gen)
// -> one step at s = isl + osl//2 + 1 with the (nextn+1) batch
// take ONLY "generation_attention"Note the Python conventions this deliberately preserves (they differed
from the pre-rewrite FPM packing): pass 1 uses
ctx + gen * (nextn + 1) tokens (the speculative-progress model —
every decode request verifies one target plus all drafts in the
combined pass, mirroring Python run_mixed’s decode_query_tokens),
the cached prefix multiplier is floor(ctx/isl) (not ceil), and the
pass-3 kv position carries _run_generation_phase’s +1.
The imbalance-correction scales mirror the RuntimeConfig fields
Python threads into each pass (base_backend.py:950-1043).
Sourcepub fn mixed_step_breakdown(
&self,
ctx_tokens: u32,
gen_tokens: u32,
isl: u32,
osl: u32,
prefix: u32,
seq_imbalance_correction_scale: f64,
gen_seq_imbalance_correction_scale: f64,
) -> Result<[f64; 4], AicError>
pub fn mixed_step_breakdown( &self, ctx_tokens: u32, gen_tokens: u32, isl: u32, osl: u32, prefix: u32, seq_imbalance_correction_scale: f64, gen_seq_imbalance_correction_scale: f64, ) -> Result<[f64; 4], AicError>
Return [total, shared_non_attention, context_attention, decode_attention] for one mixed engine iteration — the three passes
of the _get_mix_step_latency composition reported separately: pass 1
is the shared non-attention work, pass 2 the context-attention slice
(already divided by ceil(isl/ctx)), pass 3 the decode-attention
slice. Engine::mixed_step_latency is their sum; the agg
speculative scheduler consumes the components.
Sourcepub fn decode_step_latency(
&self,
gen_tokens: u32,
isl: u32,
osl: u32,
gen_seq_imbalance_correction_scale: f64,
) -> Result<f64, AicError>
pub fn decode_step_latency( &self, gen_tokens: u32, isl: u32, osl: u32, gen_seq_imbalance_correction_scale: f64, ) -> Result<f64, AicError>
One generation-only step latency. LITERAL mirror of Python
_get_genonly_step_latency (base_backend.py:1040-1100):
run_static(batch=gen_tokens, isl=isl+osl//2, osl=2, mode=static_gen)
summed over the FULL generation op list — one step at
s = isl + osl//2 + 1 (note _run_generation_phase’s +1) with the
decode batch scaled by (nextn + 1).
Sourcepub fn run_static_per_op(
&self,
runtime: &RuntimeConfig,
mode: StaticMode,
stride: u32,
) -> Result<(Vec<PerOpValue>, Vec<PerOpValue>), AicError>
pub fn run_static_per_op( &self, runtime: &RuntimeConfig, mode: StaticMode, stride: u32, ) -> Result<(Vec<PerOpValue>, Vec<PerOpValue>), AicError>
Self::run_static with the per-op values kept instead of summed:
(context, generation) lists of (name, latency_ms, energy_wms, source), NAME-FOLDED (see PerOpValue): each name crosses once,
pre-accumulated with Python’s phase-dict semantics. Generation values
are per-step-folded, then weighted by the stride repeat_count.
Sourcepub fn mixed_step_breakdown_per_op(
&self,
ctx_tokens: u32,
gen_tokens: u32,
isl: u32,
osl: u32,
prefix: u32,
seq_imbalance_correction_scale: f64,
gen_seq_imbalance_correction_scale: f64,
) -> Result<(Vec<PerOpValue>, Vec<PerOpValue>, Vec<PerOpValue>), AicError>
pub fn mixed_step_breakdown_per_op( &self, ctx_tokens: u32, gen_tokens: u32, isl: u32, osl: u32, prefix: u32, seq_imbalance_correction_scale: f64, gen_seq_imbalance_correction_scale: f64, ) -> Result<(Vec<PerOpValue>, Vec<PerOpValue>, Vec<PerOpValue>), AicError>
Self::mixed_step_breakdown with the per-op values kept:
(shared_non_attention, context_attention, decode_attention) lists of
(name, latency_ms, energy_wms, source). Context-attention entries
arrive already divided by the ceil(isl/ctx) scale.
Sourcepub fn decode_step_per_op(
&self,
gen_tokens: u32,
isl: u32,
osl: u32,
gen_seq_imbalance_correction_scale: f64,
) -> Result<Vec<PerOpValue>, AicError>
pub fn decode_step_per_op( &self, gen_tokens: u32, isl: u32, osl: u32, gen_seq_imbalance_correction_scale: f64, ) -> Result<Vec<PerOpValue>, AicError>
Self::decode_step_latency with the per-op values kept.
Sourcepub fn evaluate_context_ops(
&self,
indices: &[usize],
batch_size: u32,
s: u32,
prefix: u32,
seq_imbalance_correction_scale: f64,
x_override: Option<u32>,
) -> Result<Vec<PerOpValue>, AicError>
pub fn evaluate_context_ops( &self, indices: &[usize], batch_size: u32, s: u32, prefix: u32, seq_imbalance_correction_scale: f64, x_override: Option<u32>, ) -> Result<Vec<PerOpValue>, AicError>
Evaluate an index-addressed sublist of the compiled CONTEXT op list at
the context-phase shape (the thin op-list evaluation FFI — Python-side
orchestration like AFD partitions the compiled list and sources per-op
values here instead of walking Operation.query()).
Sourcepub fn evaluate_generation_ops(
&self,
indices: &[usize],
batch_size: u32,
s: u32,
gen_seq_imbalance_correction_scale: f64,
prefix: u32,
x_override: Option<u32>,
) -> Result<Vec<PerOpValue>, AicError>
pub fn evaluate_generation_ops( &self, indices: &[usize], batch_size: u32, s: u32, gen_seq_imbalance_correction_scale: f64, prefix: u32, x_override: Option<u32>, ) -> Result<Vec<PerOpValue>, AicError>
Evaluate an index-addressed sublist of the compiled GENERATION op list
at the decode-step shape (see Self::evaluate_context_ops).
Sourcepub fn evaluate_ops_json(
&self,
ops_json: &str,
is_context: bool,
batch_size: u32,
s: u32,
prefix: u32,
imbalance_correction_scale: f64,
x_override: Option<u32>,
) -> Result<Vec<PerOpValue>, AicError>
pub fn evaluate_ops_json( &self, ops_json: &str, is_context: bool, batch_size: u32, s: u32, prefix: u32, imbalance_correction_scale: f64, x_override: Option<u32>, ) -> Result<Vec<PerOpValue>, AicError>
Evaluate an ad-hoc op list (a JSON array of OpSpec objects, the same
externally-tagged encoding EngineSpec uses) against this engine’s
database. Serves op lists that are deliberately NOT in the compiled
spec — the VL encoder phase — while the shape math stays Python-side.
Sourcepub fn evaluate_ops_sol_json(
&self,
ops_json: &str,
is_context: bool,
batch_size: u32,
s: u32,
prefix: u32,
imbalance_correction_scale: f64,
x_override: Option<u32>,
) -> Result<Vec<PerOpSolValue>, AicError>
pub fn evaluate_ops_sol_json( &self, ops_json: &str, is_context: bool, batch_size: u32, s: u32, prefix: u32, imbalance_correction_scale: f64, x_override: Option<u32>, ) -> Result<Vec<PerOpSolValue>, AicError>
Self::evaluate_ops_json under the SOL_FULL view: evaluate an
ad-hoc op list (JSON array of OpSpec objects) with every operator
forced onto its analytic SOL branch, and keep the roofline
decomposition. Returns (name, sol_time_ms, sol_math_ms, sol_mem_ms)
per op (see PerOpSolValue) — the compiled-engine replacement for
Python’s per-call query_*(..., database_mode=SOL_FULL) triples.
Errors when an op’s family does not export its decomposition yet.
Sourcepub fn forward_pass_time_ms(
&self,
metrics_by_rank: &[ForwardPassMetrics],
) -> Result<f64, AicError>
pub fn forward_pass_time_ms( &self, metrics_by_rank: &[ForwardPassMetrics], ) -> Result<f64, AicError>
Compute one forward-pass latency from a list of per-rank FPM entries.
Re-platformed from the (deleted) SessionEstimator::forward_pass_time_ms
(commit 520dcfff session.rs:289): validate every rank, dispatch each
rank on its scheduled workload via Self::rank_latency_ms, and take the
max across ranks (attention-DP ranks run in lockstep, so the slowest rank
gates the iteration).
Unlike Self::mixed_step_latency / Self::decode_step_latency, this
consumes ALREADY-PACKED telemetry: the FPM fields are the observed
per-iteration counts, so the (nextn + 1) MTP multiplier is NOT applied
here (it is already baked into the scheduled-decode counts the engine
emitted). The dispatch reuses the shared [run_context_ops] /
[run_generation_ops_step] / [get_mix_step_ops] free fns so this path
and the live engine-step path stay numerically identical.