aisimulate-core 0.12.0

Engine-neutral inference simulation, deterministic replay, and performance modeling
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Operator primitives: GEMM, attention, MLA, MoE, communication, elementwise,
//! embedding, overlap. Submodules each define a focused op type that holds
//! its config-time parameters and exposes a `query` method against
//! `PerfDatabase`.
//!
//! Each operator wraps the raw `perf_database/<family>.rs` table with
//! op-specific extras: SOL/EMPIRICAL/HYBRID database-mode dispatch, prefix
//! correction, fused-op accounting (rope/kv-write/qk-norm for attention),
//! and bandwidth scaling for collectives. The perf-DB layer stays
//! algorithm-free; this layer is where Python's `_query_*_table` static
//! methods live.

pub mod attention;
pub mod base;
pub mod communication;
pub mod dsa;
pub mod dsv4;
pub mod elementwise;
pub mod embedding;
pub mod fpm_forward;
pub(crate) mod fpm_sol;
pub mod gemm;
pub mod mamba;
pub mod mhc;
pub mod mla;
pub mod moe;
pub mod moe_a2a;
pub mod moe_dispatch;
pub mod moe_expert_compute;
pub mod msa;
pub mod op;
pub mod overlap;
pub mod util_empirical;
pub mod vision;
pub mod wideep_mla;

pub use attention::{ContextAttentionOp, EncoderAttentionOp, GenerationAttentionOp};
pub use base::{PerformanceResult, Source};
pub use communication::{CustomAllReduceOp, NcclOp, P2POp};
pub use dsa::DsaModuleOp;
pub use dsv4::{Dsv4MegaMoeOp, Dsv4ModuleOp};
pub use elementwise::ElementwiseOp;
pub use embedding::EmbeddingOp;
pub use fpm_forward::{FpmForwardOp, FpmPhase};
pub use gemm::GemmOp;
pub use mamba::{GdnOp, KdaOp, Mamba2Op};
pub use mhc::MhcModuleOp;
pub use mla::{ContextMlaOp, GenerationMlaOp, MlaBmmOp, MlaModuleOp};
pub use moe::MoeOp;
pub use moe_a2a::MoeAllToAllOp;
#[allow(unused_imports)]
pub use moe_dispatch::{DispatchFlavor, MoEDispatchOp};
pub use moe_expert_compute::MoeExpertComputeOp;
pub use msa::MsaModuleOp;
#[allow(unused_imports)]
pub use op::{FallbackOp, Op, OverlapOp, RuntimeContext};
pub use vision::VisionEncoderOp;
pub use wideep_mla::{WideEpContextMlaOp, WideEpGenerationMlaOp};