Skip to main content

Engine

Struct Engine 

Source
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

Source

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.

Source

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.

Source

pub fn database(&self) -> &Arc<PerfDatabase>

Shared perf database handle.

Source

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.

Source

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).

Source

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.

Source

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).

Source

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).

Source

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.

Source

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.

Source

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).

Source

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.

Source

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).

Source

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.

Source

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.

Source

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.

Source

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()).

Source

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).

Source

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.

Source

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.

Source

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.

Trait Implementations§

Source§

impl Debug for Engine

Source§

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter. Read more

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
where ST: ?Sized, DT: ?Sized,

Source§

impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
where ST: ?Sized, DT: ?Sized,

Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T> Instrument for T

Source§

fn instrument(self, span: Span) -> Instrumented<Self>

Instruments this type with the provided Span, returning an Instrumented wrapper. Read more
Source§

fn in_current_span(self) -> Instrumented<Self>

Instruments this type with the current Span, returning an Instrumented wrapper. Read more
Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> Read<Exclusive, BecauseExclusive> for T
where T: ?Sized,

Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = !

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, !>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.
Source§

impl<V, T> VZip<V> for T
where V: MultiLane<T>,

Source§

fn vzip(self) -> V

Source§

impl<T> WithSubscriber for T

Source§

fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self>
where S: Into<Dispatch>,

Attaches the provided Subscriber to this type, returning a WithDispatch wrapper. Read more
Source§

fn with_current_subscriber(self) -> WithDispatch<Self>

Attaches the current default Subscriber to this type, returning a WithDispatch wrapper. Read more