aisimulate_core/perfmodel/operators/op.rs
1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! `Op` enum: unified typed dispatch for every operator family.
5//!
6//! Mirrors Python's `model.context_ops` / `model.generation_ops` /
7//! `model.encoder_ops` lists — each entry is one typed op carrying its
8//! config-time parameters. Session code iterates the list and calls
9//! `op.query(db, runtime)` exactly the way Python's `_run_context_phase` /
10//! `_run_generation_phase` iterate `for op in model.context_ops:
11//! op.query(database, **runtime_kwargs)`.
12//!
13//! Module-level ops with separate context/generation queries (MLA module,
14//! DSA, DSV4) get one variant per phase so a single `query` method handles
15//! dispatch.
16
17use serde::{Deserialize, Serialize};
18
19use crate::common::error::AicError;
20use crate::operators::{
21 ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4MegaMoeOp, Dsv4ModuleOp,
22 ElementwiseOp, EmbeddingOp, EncoderAttentionOp, FpmForwardOp, GdnOp, GemmOp,
23 GenerationAttentionOp, GenerationMlaOp, KdaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp,
24 MoEDispatchOp, MoeAllToAllOp, MoeExpertComputeOp, MoeOp, MsaModuleOp, NcclOp, P2POp,
25 PerformanceResult, Source, VisionEncoderOp, WideEpContextMlaOp, WideEpGenerationMlaOp,
26};
27use crate::perf_database::PerfDatabase;
28
29/// Runtime context passed to every op's `query`.
30///
31/// Mirrors Python's `**kwargs` payload to `op.query(database, ...)`. Each
32/// op variant extracts the fields it needs; non-applicable fields are
33/// safely ignored.
34#[derive(Clone, Copy, Debug)]
35pub struct RuntimeContext {
36 /// Per-rank batch size for attention queries (context: prefill batch;
37 /// generation: decode batch).
38 pub batch_size: u32,
39 /// Beam width (1 for static / agg / disagg; >1 for beam-search modes,
40 /// which are not currently exercised by the engine-step path).
41 pub beam_width: u32,
42 /// Sequence length passed to attention queries. For context phase this
43 /// is `effective_isl = isl - prefix`. For generation phase this is the
44 /// current `isl + step + 1` decode position.
45 pub s: u32,
46 /// Prefix length already in KV cache (context phase only; 0 otherwise).
47 pub prefix: u32,
48 /// Total per-rank token count for compute-bound ops (GEMM, Embedding,
49 /// Elementwise, MoE, MoE dispatch, comm). Python passes this as `x` to
50 /// `op.query`. For context: `batch_size * effective_isl`. For
51 /// generation: `batch_size * beam_width`.
52 pub num_tokens: u32,
53 /// Sequence-imbalance correction multiplier for context attention.
54 pub seq_imbalance_correction_scale: f64,
55 /// Sequence-imbalance correction multiplier for generation attention.
56 pub gen_seq_imbalance_correction_scale: f64,
57 /// Number of vision-encoder tokens per image (encoder phase only).
58 pub num_image_tokens: u32,
59}
60impl Default for RuntimeContext {
61 fn default() -> Self {
62 Self {
63 batch_size: 1,
64 beam_width: 1,
65 s: 1,
66 prefix: 0,
67 num_tokens: 1,
68 seq_imbalance_correction_scale: 1.0,
69 gen_seq_imbalance_correction_scale: 1.0,
70 num_image_tokens: 0,
71 }
72 }
73}
74
75/// Typed operator. One variant per Python `operations` family.
76///
77/// Module-level ops with separate context/generation queries become
78/// distinct variants so dispatch is unambiguous.
79///
80/// Serializes as the wire-format op for [`crate::perfmodel::engine::spec::EngineSpec`]
81/// (re-exported there as `OpSpec`). All config-time fields are plain
82/// serializable data, so the enum and its recursive `Overlap`/`Fallback`
83/// children round-trip through bincode.
84///
85/// `Op::Vision` is part of the shared session path and derives serde with
86/// the rest, but it is **never emitted into a compiled `EngineSpec`**:
87/// `compile_engine` decomposes the vision encoder into its child
88/// `Gemm`/`EncoderAttention`/`Elementwise` ops instead.
89/// Production specs therefore never contain a `Vision` variant.
90#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
91pub enum Op {
92 Gemm(GemmOp),
93 Embedding(EmbeddingOp),
94 Elementwise(ElementwiseOp),
95 ContextAttention(ContextAttentionOp),
96 GenerationAttention(GenerationAttentionOp),
97 EncoderAttention(EncoderAttentionOp),
98 ContextMla(ContextMlaOp),
99 GenerationMla(GenerationMlaOp),
100 MlaModuleContext(MlaModuleOp),
101 MlaModuleGeneration(MlaModuleOp),
102 MlaBmm(MlaBmmOp),
103 Moe(MoeOp),
104 MoeDispatch(MoEDispatchOp),
105 CustomAllReduce(CustomAllReduceOp),
106 Nccl(NcclOp),
107 P2P(P2POp),
108 Vision(VisionEncoderOp),
109 DsaContext(DsaModuleOp),
110 DsaGeneration(DsaModuleOp),
111 /// MiniMax Sparse Attention (MSA) context module — no silicon data;
112 /// answers only under HYBRID/EMPIRICAL via cross-op DSA util transfer.
113 MsaContext(MsaModuleOp),
114 /// MSA generation module (`s` = total KV length).
115 MsaGeneration(MsaModuleOp),
116 Dsv4Context(Dsv4ModuleOp),
117 Dsv4Generation(Dsv4ModuleOp),
118 Mhc(MhcModuleOp),
119 Mamba2(Mamba2Op),
120 Gdn(GdnOp),
121 /// SGLang WideEP context MLA — replaces `ContextMlaOp` in the
122 /// `WideEPDeepSeekModel` variant. SGLang-only perf data.
123 WideEpContextMla(WideEpContextMlaOp),
124 /// SGLang WideEP generation MLA — replaces `GenerationMlaOp` in the
125 /// `WideEPDeepSeekModel` variant.
126 WideEpGenerationMla(WideEpGenerationMlaOp),
127 /// Two op groups that execute in parallel on different CUDA streams.
128 /// Mirrors Python `aiconfigurator.sdk.operations.overlap.OverlapOp`:
129 /// `latency = max(sum(group_a), sum(group_b))`.
130 Overlap(OverlapOp),
131 /// Try a primary op; on perf-DB miss, fall back to summing a list of
132 /// granular ops. Mirrors Python
133 /// `aiconfigurator.sdk.operations.overlap.FallbackOp`: supports the
134 /// transitional state where some systems have module-level profiling
135 /// data and others still ship per-kernel granular data.
136 Fallback(FallbackOp),
137 /// SGLang DeepSeek-V4 MegaMoE routed module (Python
138 /// `DeepSeekV4MegaMoEModule`): one variant for both phases — the op's
139 /// `is_context` field selects the phase inside the unified table.
140 /// Measured-SILICON-only; see `operators/dsv4.rs::Dsv4MegaMoeOp`.
141 ///
142 /// APPENDED after `Fallback` on purpose: bincode enum indices are
143 /// positional, so appending does not shift existing variants and
144 /// `ENGINE_SPEC_SCHEMA_VERSION` stays unchanged. Do NOT insert new
145 /// variants mid-enum.
146 Dsv4MegaMoe(Dsv4MegaMoeOp),
147 /// Kimi Delta Attention (KDA) kernel for Kimi-K3 linear_attention
148 /// layers — Python `KDAKernel` (a `GDNKernel` subclass with a distinct
149 /// `kda_perf` table, an fp32-state SOL byte model, a "verify" phase and
150 /// a `draft_tokens` field). APPENDED at the end (see the bincode note on
151 /// `Dsv4MegaMoe`); the new serialized variant bumped
152 /// `ENGINE_SPEC_SCHEMA_VERSION` to 5 (renumbered to 6 at its merge).
153 Kda(KdaOp),
154 /// Whole-model forward pass (Python `forward_model="fpm"`): with the FPM
155 /// rewrite each phase op list is exactly one of these, answering from the
156 /// collected `fpm_forward_perf` cells instead of a granular composition.
157 /// NOT related to the `crate::fpm` (ForwardPassPerfModel) module.
158 /// APPENDED at the end (see the bincode note on `Dsv4MegaMoe`); claimed
159 /// `ENGINE_SPEC_SCHEMA_VERSION` 5 concurrently with #1460/#1435 and was
160 /// renumbered to 9 across the intervening wire-format landings.
161 FpmForward(FpmForwardOp),
162 /// Unified large-EP MoE all-to-all comm phase (Python
163 /// `operations.moe_comm.MoEAllToAll`) — one variant serves every backend
164 /// and every phase; the op's `phase` / `comm_backend` fields select the
165 /// slice. Measured-SILICON-only; see `operators/moe_a2a.rs`.
166 ///
167 /// APPENDED after `FpmForward` — same positional-index rule as above.
168 MoeAllToAll(MoeAllToAllOp),
169 /// Unified large-EP MoE expert compute (Python
170 /// `operations.moe_comm.MoEExpertCompute`) — one variant for both inference phases;
171 /// the op's `inference_phase` field selects the slice.
172 /// Measured-SILICON-only; see `operators/moe_expert_compute.rs`.
173 MoeExpertCompute(MoeExpertComputeOp),
174}
175
176/// Inline-defined here (rather than a sibling module under `operators/`)
177/// because the variants of an overlap group are themselves `Op` values —
178/// the definition is cyclic with `Op` and the implementation is small.
179#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
180pub struct OverlapOp {
181 pub name: String,
182 pub group_a: Vec<Op>,
183 pub group_b: Vec<Op>,
184}
185
186impl OverlapOp {
187 pub fn new(name: impl Into<String>, group_a: Vec<Op>, group_b: Vec<Op>) -> Self {
188 Self {
189 name: name.into(),
190 group_a,
191 group_b,
192 }
193 }
194}
195
196#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
197pub struct FallbackOp {
198 pub name: String,
199 /// Try this first. On `AicError::PerfDatabase`-class failure (missing
200 /// file / missing data point), the fallback chain is used instead.
201 pub primary: Box<Op>,
202 pub fallback: Vec<Op>,
203}
204
205impl FallbackOp {
206 pub fn new(name: impl Into<String>, primary: Op, fallback: Vec<Op>) -> Self {
207 Self {
208 name: name.into(),
209 primary: Box::new(primary),
210 fallback,
211 }
212 }
213}
214
215impl Op {
216 /// Constant per-op weight bytes (PR-6): the engine-side replacement for
217 /// Python's `Operation.get_weights` math. Structural, not data-driven —
218 /// computed from op fields alone, never from perf tables. Ops with no
219 /// resident weights (attention/MLA kernels — their weights live on the
220 /// adjacent GEMMs — comm ops, dispatch, elementwise, MSA, the mamba
221 /// KERNEL ops) are 0.0, exactly like their Python `_weights = 0.0`.
222 /// `FpmForward` carries its snapshot verbatim (Python returns
223 /// `_weight_bytes` WITHOUT the scale_factor multiply); every non-zero
224 /// family multiplies its own scale_factor inside its `weight_bytes`.
225 pub fn weight_bytes(&self) -> f64 {
226 match self {
227 Op::Gemm(o) => o.weights_bytes(),
228 Op::Embedding(o) => o.weights_bytes(),
229 Op::Moe(o) => o.weight_bytes(),
230 Op::MoeExpertCompute(o) => o.weight_bytes(),
231 Op::Dsv4MegaMoe(o) => o.weight_bytes(),
232 Op::Mhc(o) => o.weight_bytes(),
233 Op::DsaContext(o) | Op::DsaGeneration(o) => o.weight_bytes(),
234 Op::Dsv4Context(o) | Op::Dsv4Generation(o) => o.weight_bytes(),
235 Op::FpmForward(o) => o.weight_bytes,
236 // Python FallbackOp.get_weights: primary wins when positive,
237 // else the granular fallback chain sums.
238 Op::Fallback(o) => {
239 let primary = o.primary.weight_bytes();
240 if primary > 0.0 {
241 primary
242 } else {
243 o.fallback.iter().map(Op::weight_bytes).sum()
244 }
245 }
246 // Python OverlapOp.get_weights: both groups sum (unlike latency's max).
247 Op::Overlap(o) => {
248 o.group_a.iter().map(Op::weight_bytes).sum::<f64>()
249 + o.group_b.iter().map(Op::weight_bytes).sum::<f64>()
250 }
251 Op::Elementwise(_)
252 | Op::ContextAttention(_)
253 | Op::GenerationAttention(_)
254 | Op::EncoderAttention(_)
255 | Op::ContextMla(_)
256 | Op::GenerationMla(_)
257 | Op::MlaModuleContext(_)
258 | Op::MlaModuleGeneration(_)
259 | Op::MlaBmm(_)
260 | Op::MoeDispatch(_)
261 | Op::CustomAllReduce(_)
262 | Op::Nccl(_)
263 | Op::P2P(_)
264 | Op::Vision(_)
265 | Op::MsaContext(_)
266 | Op::MsaGeneration(_)
267 | Op::Mamba2(_)
268 | Op::Gdn(_)
269 | Op::Kda(_)
270 | Op::WideEpContextMla(_)
271 | Op::WideEpGenerationMla(_)
272 | Op::MoeAllToAll(_) => 0.0,
273 }
274 }
275
276 /// Stable op name (Python `op._name`). Used by session code to filter
277 /// (e.g. context-attention exclusion in mix-step composition) and for
278 /// debugging.
279 pub fn name(&self) -> &str {
280 match self {
281 Op::Gemm(o) => &o.name,
282 Op::Embedding(o) => &o.name,
283 Op::Elementwise(o) => &o.name,
284 Op::ContextAttention(o) => &o.name,
285 Op::GenerationAttention(o) => &o.name,
286 Op::EncoderAttention(o) => &o.name,
287 Op::ContextMla(o) => &o.name,
288 Op::GenerationMla(o) => &o.name,
289 Op::MlaModuleContext(o) => &o.name,
290 Op::MlaModuleGeneration(o) => &o.name,
291 Op::MlaBmm(o) => &o.name,
292 Op::Moe(o) => &o.name,
293 Op::MoeDispatch(o) => &o.name,
294 Op::CustomAllReduce(o) => &o.name,
295 Op::Nccl(o) => &o.name,
296 Op::P2P(o) => &o.name,
297 Op::Vision(o) => &o.name,
298 Op::DsaContext(o) => &o.name,
299 Op::DsaGeneration(o) => &o.name,
300 Op::MsaContext(o) => &o.name,
301 Op::MsaGeneration(o) => &o.name,
302 Op::Dsv4Context(o) => &o.name,
303 Op::Dsv4Generation(o) => &o.name,
304 Op::Mhc(o) => &o.name,
305 Op::Mamba2(o) => &o.name,
306 Op::Gdn(o) => &o.name,
307 Op::WideEpContextMla(o) => &o.name,
308 Op::WideEpGenerationMla(o) => &o.name,
309 Op::FpmForward(o) => &o.name,
310 Op::Overlap(o) => &o.name,
311 Op::Fallback(o) => &o.name,
312 Op::Dsv4MegaMoe(o) => &o.name,
313 Op::Kda(o) => &o.name,
314 Op::MoeAllToAll(o) => &o.name,
315 Op::MoeExpertCompute(o) => &o.name,
316 }
317 }
318
319 /// Rename the op (Python's post-construction `op._name = ...` rewiring:
320 /// hybrid layer-type prefixes rename block ops after the shared builder
321 /// returns them). Every variant carries `name`.
322 pub fn set_name(&mut self, name: String) {
323 match self {
324 Op::Gemm(o) => o.name = name,
325 Op::Embedding(o) => o.name = name,
326 Op::Elementwise(o) => o.name = name,
327 Op::ContextAttention(o) => o.name = name,
328 Op::GenerationAttention(o) => o.name = name,
329 Op::EncoderAttention(o) => o.name = name,
330 Op::ContextMla(o) => o.name = name,
331 Op::GenerationMla(o) => o.name = name,
332 Op::MlaModuleContext(o) => o.name = name,
333 Op::MlaModuleGeneration(o) => o.name = name,
334 Op::MlaBmm(o) => o.name = name,
335 Op::Moe(o) => o.name = name,
336 Op::MoeDispatch(o) => o.name = name,
337 Op::CustomAllReduce(o) => o.name = name,
338 Op::Nccl(o) => o.name = name,
339 Op::P2P(o) => o.name = name,
340 Op::Vision(o) => o.name = name,
341 Op::DsaContext(o) => o.name = name,
342 Op::DsaGeneration(o) => o.name = name,
343 Op::MsaContext(o) => o.name = name,
344 Op::MsaGeneration(o) => o.name = name,
345 Op::Dsv4Context(o) => o.name = name,
346 Op::Dsv4Generation(o) => o.name = name,
347 Op::Mhc(o) => o.name = name,
348 Op::Mamba2(o) => o.name = name,
349 Op::Gdn(o) => o.name = name,
350 Op::WideEpContextMla(o) => o.name = name,
351 Op::WideEpGenerationMla(o) => o.name = name,
352 Op::FpmForward(o) => o.name = name,
353 Op::Overlap(o) => o.name = name,
354 Op::Fallback(o) => o.name = name,
355 Op::Dsv4MegaMoe(o) => o.name = name,
356 Op::Kda(o) => o.name = name,
357 Op::MoeAllToAll(o) => o.name = name,
358 Op::MoeExpertCompute(o) => o.name = name,
359 }
360 }
361
362 /// CP sequence-shard factor for the token-major families that carry one;
363 /// 1 for every other variant (their constructors' CP audit gate refuses
364 /// `seq_split > 1`, so 1 is exact, not a guess). Backs the Python-side
365 /// `Operation._seq_split` default read.
366 pub fn seq_split(&self) -> u32 {
367 match self {
368 Op::Gemm(o) => o.seq_split,
369 Op::Embedding(o) => o.seq_split,
370 Op::Elementwise(o) => o.seq_split,
371 Op::CustomAllReduce(o) => o.seq_split,
372 Op::Nccl(o) => o.seq_split,
373 Op::P2P(o) => o.seq_split,
374 Op::Mhc(o) => o.seq_split,
375 _ => 1,
376 }
377 }
378
379 /// True if this op's name matches Python's mix-step filter for the
380 /// context-attention bucket. Python uses literal string equality on
381 /// `"context_attention"` — that's the LLAMA / MOE attention op name.
382 /// Models with module-level attention (e.g. Kimi's
383 /// `context_mla_module`) have names that don't match this filter, so
384 /// they're treated as non-attention in the mix-step composition
385 /// (matching Python's intent: the module already represents the full
386 /// fused attention+projection work and shouldn't be re-decomposed).
387 pub fn is_context_attention(&self) -> bool {
388 self.name() == "context_attention"
389 }
390
391 /// True if this op's name matches Python's mix-step filter for the
392 /// generation-attention bucket (`"generation_attention"`).
393 pub fn is_generation_attention(&self) -> bool {
394 self.name() == "generation_attention"
395 }
396
397 /// Identifies the logits projection GEMM by name. Python special-cases
398 /// `logits_gemm` in `_run_context_phase` to use `x=batch_size` instead
399 /// of `x=batch_size * effective_isl`.
400 pub fn is_logits_gemm(&self) -> bool {
401 matches!(self, Op::Gemm(_)) && self.name().contains("logits_gemm")
402 }
403
404 /// Query this op with the given runtime. Returns the scaled latency
405 /// from the underlying op's `query` method.
406 pub fn query(
407 &self,
408 db: &PerfDatabase,
409 ctx: &RuntimeContext,
410 ) -> Result<PerformanceResult, AicError> {
411 match self {
412 Op::Gemm(op) => op.query(db, ctx.num_tokens, None),
413 Op::Embedding(op) => op.query(db, ctx.num_tokens),
414 Op::Elementwise(op) => op.query(db, ctx.num_tokens),
415 Op::ContextAttention(op) => op.query(
416 db,
417 ctx.batch_size,
418 ctx.s,
419 ctx.prefix,
420 ctx.seq_imbalance_correction_scale,
421 ),
422 Op::GenerationAttention(op) => op.query(
423 db,
424 ctx.batch_size,
425 ctx.s,
426 ctx.gen_seq_imbalance_correction_scale,
427 ),
428 Op::EncoderAttention(op) => op.query(db, ctx.batch_size, ctx.s),
429 Op::ContextMla(op) => op.query(db, ctx.batch_size, ctx.s, ctx.prefix),
430 Op::GenerationMla(op) => op.query(db, ctx.batch_size, ctx.s),
431 Op::MlaModuleContext(op) => op.query_context(db, ctx.batch_size, ctx.s, ctx.prefix),
432 Op::MlaModuleGeneration(op) => op.query_generation(db, ctx.batch_size, ctx.s),
433 // Python's `MLABmm.query` uses `batch_size` as the BMM table's
434 // tokens-axis index (the table's `num_tokens` column equals the
435 // op's per-request count, which is `batch_size`). Pass
436 // `ctx.batch_size`, not `ctx.num_tokens`.
437 Op::MlaBmm(op) => op.query(db, ctx.batch_size),
438 Op::Moe(op) => op.query(db, ctx.num_tokens),
439 Op::MoeDispatch(op) => op.query(db, ctx.num_tokens),
440 Op::CustomAllReduce(op) => op.query(db, ctx.num_tokens),
441 Op::Nccl(op) => op.query(db, ctx.num_tokens),
442 Op::P2P(op) => op.query(db, ctx.num_tokens),
443 Op::Vision(op) => op.query(db, ctx.num_image_tokens),
444 Op::DsaContext(op) => op.query_context(db, ctx.batch_size, ctx.s, ctx.prefix),
445 Op::DsaGeneration(op) => op.query_generation(db, ctx.batch_size, ctx.s),
446 Op::MsaContext(op) => op.query_context(db, ctx.batch_size, ctx.s, ctx.prefix),
447 Op::MsaGeneration(op) => op.query_generation(db, ctx.batch_size, ctx.s),
448 Op::Dsv4Context(op) => op.query_context(db, ctx.batch_size, ctx.s, ctx.prefix),
449 Op::Dsv4Generation(op) => op.query_generation(db, ctx.batch_size, ctx.s),
450 Op::Mhc(op) => op.query(db, ctx.num_tokens),
451 Op::Mamba2(op) => op.query(db, ctx.batch_size, ctx.s),
452 Op::Gdn(op) => op.query(db, ctx.batch_size, ctx.s),
453 Op::WideEpContextMla(op) => op.query(db, ctx.batch_size, ctx.s, ctx.prefix),
454 Op::WideEpGenerationMla(op) => op.query(db, ctx.batch_size, ctx.s),
455 // Whole-model op: consumes batch_size/s/prefix/beam_width from the
456 // context (num_tokens is ignored, mirroring Python's kwargs use).
457 Op::FpmForward(op) => op.query(db, ctx),
458 Op::Overlap(op) => {
459 // Mirrors Python `OverlapOp.query`: each group accumulates
460 // through `PerformanceResult` addition from a zero/empirical
461 // seed (`total_a = PerformanceResult(0.0, energy=0.0,
462 // source="empirical"); total_a += ...`), then latency =
463 // max(group_a_total, group_b_total) while ENERGY = group_a +
464 // group_b (both groups consume power even though they overlap
465 // in time). The source is `(total_a + total_b).source` — the
466 // seeds and `plus`'s zero-identity rule make zero-valued
467 // members (empty groups, nested empty composites, zero-cost
468 // legs) source-NEUTRAL instead of poisoning the tag to Mixed.
469 let mut total_a = PerformanceResult::new(0.0, Source::Empirical);
470 for inner in &op.group_a {
471 total_a = total_a.plus(inner.query(db, ctx)?);
472 }
473 let mut total_b = PerformanceResult::new(0.0, Source::Empirical);
474 for inner in &op.group_b {
475 total_b = total_b.plus(inner.query(db, ctx)?);
476 }
477 let latency_ms = total_a.latency_ms.max(total_b.latency_ms);
478 let energy_wms = total_a.energy_wms + total_b.energy_wms;
479 let merged = total_a.plus(total_b);
480 Ok(
481 PerformanceResult::with_energy(latency_ms, energy_wms, merged.source)
482 .with_moe_comm_fallbacks(merged.moe_comm_fallbacks)
483 .clamp_non_negative(),
484 )
485 }
486 Op::Fallback(op) => {
487 // Mirrors Python `FallbackOp.query`: try the primary; on
488 // perf-DB-class failure, sum the fallback chain instead.
489 // (Python additionally caches `primary_unavailable=True` to
490 // skip subsequent retries — we don't bother here because the
491 // hot-path penalty is one `OnceLock::get` per call on a
492 // populated path and one retry on a missing one.)
493 //
494 // Under HYBRID the primary is evaluated against a SILICON
495 // view (Python swaps in `_get_configured_database_view(db,
496 // SILICON, transfer_policy)`): a missing module table must
497 // fall to the granular fallback chain, not be hybrid-
498 // estimated at module level. The fallback ops then run
499 // against the ORIGINAL (hybrid) database.
500 let silicon_db;
501 let primary_db: &PerfDatabase =
502 if db.database_mode == crate::common::enums::DatabaseMode::Hybrid {
503 silicon_db = db.silicon_view();
504 &silicon_db
505 } else {
506 db
507 };
508 match op.primary.query(primary_db, ctx) {
509 // Primary result passes through verbatim — its energy
510 // rides along (Python returns `self._primary.query(...)`).
511 Ok(r) => Ok(r),
512 Err(AicError::PerfDatabase(_)) | Err(AicError::Io { .. }) => {
513 // Fallback chain: Python sums PerformanceResults
514 // from a zero/empirical seed (`total =
515 // PerformanceResult(0.0, energy=0.0,
516 // source="empirical"); total += op.query(...)`), so
517 // latency AND energy both accumulate and an empty (or
518 // all-zero) chain keeps the empirical seed tag via
519 // `plus`'s zero-identity rule.
520 let mut total = PerformanceResult::new(0.0, Source::Empirical);
521 for inner in &op.fallback {
522 total = total.plus(inner.query(db, ctx)?);
523 }
524 // `with_energy` (sol: None) keeps the pre-existing
525 // composed-result behavior: only the SOURCE semantics
526 // change in this fix, not the SOL decomposition.
527 Ok(PerformanceResult::with_energy(
528 total.latency_ms,
529 total.energy_wms,
530 total.source,
531 )
532 .with_moe_comm_fallbacks(total.moe_comm_fallbacks)
533 .clamp_non_negative())
534 }
535 Err(other) => Err(other),
536 }
537 }
538 // Rank-LOCAL token count, like Moe/MoeDispatch (Python passes the
539 // same `x`); the megamoe table is indexed by local-rank tokens
540 // and the op must NOT re-multiply by attention_dp_size.
541 Op::Dsv4MegaMoe(op) => op.query(db, ctx.num_tokens),
542 // Like Gdn: the op derives its phase coordinates internally
543 // (verify divides the (nextn+1)-scaled batch by draft_tokens).
544 Op::Kda(op) => op.query(db, ctx.batch_size, ctx.s),
545 // Both large-EP ops take Python's `x` (moe_comm.py:657, :1291) —
546 // the same per-rank token count every other compute/comm op gets.
547 // The per-op token rescaling (`// attention_tp_size` for the comm
548 // side, `* attention_dp_size` for the compute side) happens INSIDE
549 // each `query`, exactly where Python does it.
550 Op::MoeAllToAll(op) => op.query(db, ctx.num_tokens),
551 Op::MoeExpertCompute(op) => op.query(db, ctx.num_tokens),
552 }
553 }
554}
555
556#[cfg(test)]
557mod tests {
558 use super::*;
559 use crate::common::enums::GemmQuantMode;
560 use crate::operators::gemm::GemmOp;
561 use crate::perf_database::PerfDatabase;
562 use crate::perf_database::energy_test_fixtures::{Col, write_parquet};
563
564 /// Minimal systems root with ONE bf16 GEMM row: an exact-hit silicon
565 /// leaf at `num_tokens=128`, and a guaranteed typed data miss for any
566 /// fp8 query (no fp8 table exists).
567 fn one_row_gemm_db() -> (tempfile::TempDir, PerfDatabase) {
568 let tmp = tempfile::tempdir().expect("tmpdir");
569 let data =
570 crate::perf_database::energy_test_fixtures::write_energy_systems_root(tmp.path());
571 write_parquet(
572 &data.join("gemm_perf.parquet"),
573 &[
574 Col::Str("gemm_dtype", vec!["bfloat16"]),
575 Col::I64("m", vec![128]),
576 Col::I64("n", vec![1024]),
577 Col::I64("k", vec![1024]),
578 Col::F64("latency", vec![1.0]),
579 Col::F64("power", vec![100.0]),
580 ],
581 );
582 let db = PerfDatabase::load(tmp.path(), "testsys", "vllm", "1.0").expect("db must load");
583 (tmp, db)
584 }
585
586 fn silicon_leaf() -> Op {
587 Op::Gemm(GemmOp::new("hit", 1024, 1024, GemmQuantMode::Bfloat16))
588 }
589
590 fn missing_leaf() -> Op {
591 Op::Gemm(GemmOp::new("miss", 1024, 1024, GemmQuantMode::Fp8))
592 }
593
594 fn ctx() -> RuntimeContext {
595 RuntimeContext {
596 num_tokens: 128,
597 ..RuntimeContext::default()
598 }
599 }
600
601 fn empty_overlap(name: &str) -> Op {
602 Op::Overlap(OverlapOp::new(name, vec![], vec![]))
603 }
604
605 // Zero-valued composite provenance oracle (review #1552 round 4): the
606 // legacy Python accumulators start from `PerformanceResult(0.0,
607 // energy=0.0, source="empirical")` and `__add__` treats a (0.0, 0.0)
608 // operand as a source-neutral identity, so zero-valued members must
609 // never poison a composite's tag to Mixed and empty composition must
610 // report `empirical`.
611
612 #[test]
613 fn empty_overlap_source_is_empirical() {
614 let (_tmp, db) = one_row_gemm_db();
615 let r = empty_overlap("e").query(&db, &ctx()).expect("query");
616 assert_eq!(r.latency_ms, 0.0);
617 assert_eq!(r.energy_wms, 0.0);
618 assert_eq!(r.source, Source::Empirical);
619 }
620
621 #[test]
622 fn half_empty_overlap_keeps_leaf_source() {
623 let (_tmp, db) = one_row_gemm_db();
624 let op = Op::Overlap(OverlapOp::new("half", vec![silicon_leaf()], vec![]));
625 let r = op.query(&db, &ctx()).expect("query");
626 assert!((r.latency_ms - 1.0).abs() < 1e-12);
627 assert_eq!(r.source, Source::Silicon);
628 }
629
630 #[test]
631 fn nested_zero_overlap_is_source_neutral_same_group() {
632 let (_tmp, db) = one_row_gemm_db();
633 let op = Op::Overlap(OverlapOp::new(
634 "same",
635 vec![silicon_leaf(), empty_overlap("nested")],
636 vec![],
637 ));
638 let r = op.query(&db, &ctx()).expect("query");
639 assert!((r.latency_ms - 1.0).abs() < 1e-12);
640 assert_eq!(
641 r.source,
642 Source::Silicon,
643 "zero-valued nested composite must be source-neutral"
644 );
645 }
646
647 #[test]
648 fn nested_zero_overlap_is_source_neutral_opposite_group() {
649 let (_tmp, db) = one_row_gemm_db();
650 let op = Op::Overlap(OverlapOp::new(
651 "opp",
652 vec![silicon_leaf()],
653 vec![empty_overlap("nested")],
654 ));
655 let r = op.query(&db, &ctx()).expect("query");
656 assert!((r.latency_ms - 1.0).abs() < 1e-12);
657 assert_eq!(
658 r.source,
659 Source::Silicon,
660 "zero-valued opposite group must be source-neutral"
661 );
662 }
663
664 #[test]
665 fn failed_primary_empty_fallback_is_empirical() {
666 let (_tmp, db) = one_row_gemm_db();
667 let op = Op::Fallback(FallbackOp::new("fb", missing_leaf(), vec![]));
668 let r = op.query(&db, &ctx()).expect("query");
669 assert_eq!(r.latency_ms, 0.0);
670 assert_eq!(r.energy_wms, 0.0);
671 assert_eq!(r.source, Source::Empirical);
672 }
673
674 #[test]
675 fn failed_primary_fallback_chain_keeps_leaf_source() {
676 let (_tmp, db) = one_row_gemm_db();
677 let op = Op::Fallback(FallbackOp::new("fb", missing_leaf(), vec![silicon_leaf()]));
678 let r = op.query(&db, &ctx()).expect("query");
679 assert!((r.latency_ms - 1.0).abs() < 1e-12);
680 assert_eq!(r.source, Source::Silicon);
681 }
682}