Skip to content
Merged
Show file tree
Hide file tree
Changes from 15 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 82 additions & 2 deletions crates/core/src/perfmodel/engine/readiness.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,10 @@
use std::collections::HashMap;

use crate::AicError;
use crate::common::enums::{DatabaseMode, GemmQuantMode, TransferKind};
use crate::common::enums::{DatabaseMode, GemmQuantMode, MoeQuantMode, TransferKind};
use crate::config::PerfSource;
use crate::operators::Op;
use crate::perf_database::moe::MoeKernel;
use crate::perf_database::{PerfDatabase, kernel_source_ok, parquet_loader::PerfReader};

pub(super) fn validate<'a>(
Expand All @@ -36,6 +37,15 @@ impl Availability<'_> {
if let Some(found) = self.tables.get(name) {
return Ok(*found);
}
if name == "moe_perf.parquet" {
let found = match self.db.moe.available_quants(MoeKernel::Standard) {
Ok(quants) => !quants.is_empty(),
Err(error) if error.is_missing_perf_data() => false,
Err(error) => return Err(error),
};
self.tables.insert(name.to_owned(), found);
return Ok(found);
}
let sources = match name {
"nccl_perf.parquet" => self
.db
Expand Down Expand Up @@ -184,7 +194,30 @@ impl Availability<'_> {
MlaModuleContext(_) => self.any(&["mla_context_module_perf.parquet"]),
MlaModuleGeneration(_) => self.any(&["mla_generation_module_perf.parquet"]),
MlaBmm(_) => self.any(&["mla_bmm_perf.parquet"]),
Moe(_) => self.any(&["moe_perf.parquet"]),
Moe(moe) => {
let found = match moe.moe_kernel_source.as_deref() {
Some(source) => !self
.db
.moe
.available_quants_for_kernel_source(source)?
.is_empty(),
None => {
self.has("moe_perf.parquet")?
|| (moe.is_gated
&& moe.quant_mode == MoeQuantMode::Nvfp4
&& self.db.moe.low_latency_available()?)
}
};
if found {
Ok(())
} else {
Err(AicError::PerfDatabase(format!(
"required MoE data unavailable for kernel_source={:?} at {}",
moe.moe_kernel_source,
self.db.data_root.display()
)))
}
}
// State-space kernels explicitly fall back to their analytic SOL
// on missing tables in every database mode.
Mamba2(_) | Gdn(_) | Kda(_) => Ok(()),
Expand Down Expand Up @@ -571,4 +604,51 @@ mod tests {
}
}
}

#[test]
fn moe_readiness_requires_eligible_or_exact_source_rows() {
use crate::common::enums::MoeQuantMode;
use crate::operators::MoeOp;

let root = systems();
let data = root
.path()
.join("data/b200_sxm/vllm/0.24.0/moe_perf.parquet");
for eligibility in [None, Some(true), Some(false)] {
let mut columns = vec![
Col::Str("moe_dtype", vec!["fp8_block"]),
Col::I64("num_tokens", vec![32]),
Col::I64("hidden_size", vec![8192]),
Col::I64("inter_size", vec![2048]),
Col::I64("topk", vec![8]),
Col::I64("num_experts", vec![256]),
Col::I64("moe_tp_size", vec![1]),
Col::I64("moe_ep_size", vec![1]),
Col::Str("distribution", vec!["uniform"]),
Col::Str("kernel_source", vec!["exact"]),
Col::F64("latency", vec![0.25]),
];
if let Some(eligible) = eligibility {
columns.push(Col::Bool("default_eligible", vec![eligible]));
}
write_parquet(&data, &columns);
for source in [None, Some("exact")] {
let db = PerfDatabase::load(root.path(), "b200_sxm", "vllm", "0.24.0").unwrap();
let mut op = MoeOp::new(
"moe",
8192,
2048,
8,
256,
1,
1,
MoeQuantMode::Fp8Block,
"uniform",
);
op.moe_kernel_source = source.map(str::to_owned);
let expected = source.is_some() || eligibility != Some(false);
assert_eq!(engine_readiness(db, Op::Moe(op)).is_ok(), expected);
}
}
}
}
72 changes: 58 additions & 14 deletions crates/core/src/perfmodel/perf_database/moe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ use super::{SourceResolver, kernel_source_ok};
use crate::common::enums::MoeQuantMode;
use crate::common::error::AicError;
use crate::config::{PerfDbSources, PerfSource};
use crate::perf_database::parquet_loader::PerfReader;
use crate::perf_database::parquet_loader::{PerfReader, PerfRow};

pub struct MoeTable {
data_root: PathBuf,
Expand Down Expand Up @@ -75,8 +75,9 @@ pub struct MoeSiblingSlice {
}

/// Parallel grids split by `kernel_source`. `default` and `low_latency`
/// preserve the legacy automatic lookup behavior, while `by_kernel_source`
/// retains every named lane for an exact user-selected lookup.
/// contain default-eligible rows (all legacy rows when the metadata column is
/// absent), while `by_kernel_source` retains every named lane for an exact
/// user-selected lookup.
struct LoadedMoeGrids {
default: MoeGrids,
low_latency: MoeGrids,
Expand Down Expand Up @@ -611,6 +612,37 @@ pub(crate) fn moe_kernel_quant_rewrite(raw_quant: String, kernel_source: &str) -
}
}

/// Optional selection metadata shared by the query loader and coverage view.
/// An absent column preserves legacy automatic selection; malformed metadata
/// is not a coverage miss and must never trigger an estimator fallback.
pub(crate) fn moe_default_eligible(
reader: &PerfReader,
row: &PerfRow,
column: Option<usize>,
kernel_source_column: Option<usize>,
) -> Result<bool, AicError> {
let Some(column) = column else {
return Ok(true);
};
let eligible = row.bool_strict(column).map_err(|_| {
AicError::InvalidPerfData(format!(
"default_eligible must be a non-null Boolean at {}",
reader.path().display()
))
})?;
if !eligible
&& row
.str_optional(kernel_source_column)?
.is_none_or(|source| source.trim().is_empty())
{
return Err(AicError::InvalidPerfData(format!(
"default_eligible=false requires a nonblank string kernel_source at {}",
reader.path().display()
)));
}
Ok(eligible)
}

fn load_moe_parquet(sources: &[PerfSource]) -> Result<LoadedMoeGrids, AicError> {
let mut default_index: MoeIndex<MoeShapeKey, BTreeMap<u32, LeafValue>> = MoeIndex::default();
let mut low_latency_index: MoeIndex<MoeShapeKey, BTreeMap<u32, LeafValue>> =
Expand Down Expand Up @@ -645,11 +677,14 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result<LoadedMoeGrids, AicError>
// the `default` grid (matching the pre-split behavior). The same column
// gates the per-source shared-layer `kernel_source` allowlist.
let kernel_source_col = reader.col_optional("kernel_source");
let default_eligible_col = reader.col_optional("default_eligible");
for row in reader.rows()? {
let row = row?;
if !kernel_source_ok(source.kernel_sources(), kernel_source_col, &row)? {
continue;
}
let default_eligible =
moe_default_eligible(&reader, &row, default_eligible_col, kernel_source_col)?;
let kernel_source = row
.str_optional(kernel_source_col)?
.unwrap_or("")
Expand Down Expand Up @@ -683,15 +718,17 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result<LoadedMoeGrids, AicError>
let power = row.f64_optional(power_col)?.unwrap_or(0.0);
let value = LeafValue::with_power(latency, power);
let num_tokens = row.u32(num_tokens_col)?;
insert_moe_row(
target,
target_quants,
&quant,
&distribution,
shape,
num_tokens,
value,
);
if default_eligible {
insert_moe_row(
target,
target_quants,
&quant,
&distribution,
shape,
num_tokens,
value,
);
}
if !kernel_source.is_empty() {
let source_index = kernel_source_indices
.entry(kernel_source.clone())
Expand All @@ -709,7 +746,11 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result<LoadedMoeGrids, AicError>
}
}
}
if !any_source || (default_index.is_empty() && low_latency_index.is_empty()) {
if !any_source
|| (default_index.is_empty()
&& low_latency_index.is_empty()
&& kernel_source_indices.is_empty())
{
return Err(AicError::PerfDatabase(format!(
"no rows loaded from {} source(s) (first: {})",
sources.len(),
Expand Down Expand Up @@ -749,7 +790,10 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result<LoadedMoeGrids, AicError>
}

fn clone_err(err: &AicError) -> AicError {
AicError::PerfDatabase(err.to_string())
match err {
AicError::InvalidPerfData(message) => AicError::InvalidPerfData(message.clone()),
_ => AicError::PerfDatabase(err.to_string()),
}
}

#[cfg(test)]
Expand Down
12 changes: 11 additions & 1 deletion crates/core/src/perfmodel/perf_database/table_view.rs
Original file line number Diff line number Diff line change
Expand Up @@ -912,12 +912,22 @@ pub fn view_kda(sources: &[PerfSource]) -> Result<Option<ViewNode>, AicError> {
/// `operations/moe.py::load_moe_data` — 9-level
/// `[moe_dtype][distribution][topk][num_experts][hidden][inter][moe_tp][moe_ep][num_tokens]`,
/// with `kernel_source == "moe_torch_flow_min_latency"` rows routed to the
/// low-latency twin table. Returns (default, low_latency).
/// low-latency twin table. Only default-eligible rows enter these automatic
/// coverage views; raw parquet enumeration retains every measured row.
/// Returns (default, low_latency).
pub fn view_moe(sources: &[PerfSource]) -> Result<Option<(ViewNode, ViewNode)>, AicError> {
let mut default = ViewNode::branch();
let mut low_latency = ViewNode::branch();
let found = fold_sources(sources, |ctx| {
let r = ctx.reader;
if !crate::perf_database::moe::moe_default_eligible(
r,
ctx.row,
r.col_optional("default_eligible"),
ctx.ks_col,
)? {
return Ok(());
}
let num_tokens = ctx.row.u32(r.col("num_tokens")?)?;
let hidden = ctx.row.u32(r.col("hidden_size")?)?;
let inter = ctx.row.u32(r.col("inter_size")?)?;
Expand Down
2 changes: 2 additions & 0 deletions docs/core-api.md
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,8 @@ The canonical configuration also carries quantization overrides and `attention_b

`moe_kernel_source` selects an exact, nonblank collected `kernel_source` label for fused MoE compute. It is distinct from the existing `moe_backend` graph/backend control; source labels are not backend aliases and are preserved without trimming. `None` keeps the existing default source-selection policy, including eligible low-latency NVFP4 selection. `SILICON` reads only the requested source's table; `EMPIRICAL` derives its estimate from that same source; `HYBRID` may fall back to empirical estimation within that source, but does not substitute a different source. Missing source data remains an error. An explicit `moe_torch_flow_min_latency` requires gated NVFP4 and at most 128 tokens after attention-DP gathering. Pure-roofline `SOL` remains table-independent and does not claim measured support for the requested source.

`moe_perf.parquet` may include Boolean selection metadata `default_eligible`. An absent column preserves legacy automatic selection; when present, every value must be a non-null Boolean. A `false` row requires a nonblank string `kernel_source`, preserved exactly, and is available only through that named source. It cannot enter automatic standard or low-latency grids, including empirical cross-shape and cross-quant reference selection. Among eligible rows, existing source priority and first-row precedence remain unchanged. Default table views, coverage, and readiness use those same eligible grids; raw Parquet enumeration retains all measured rows. Malformed eligibility metadata is an invalid-data error, never a missing-data fallback. The flag is not a measurement identity dimension. Collector finalization preserves it during merges, including existing annotations when a legacy recollection omits the column; genuinely new legacy keys remain eligible.

An explicit source is rejected for dense graphs, MegaMoE modules, large-EP expert-compute graphs, and any constructed timing phase with no compatible fused MoE operator. It is also incompatible with whole-forward FPM, including the legacy Task `forward_model='fpm'` rewrite. Invalid graph/source combinations fail as invalid configuration rather than triggering estimator fallback. An untrained `fpm_regression` model remains not-ready; retaining a source in its configuration is not evidence of source-specific prediction support.

AFD regular companions currently reject exact-source requests through their legacy estimator and fixed timing paths. The external-FPM companion forwards the source to canonical validation, which rejects the incompatible FPM request. These controls describe standalone AISimulate behavior, not downstream Dynamo planner integration.
Expand Down
2 changes: 2 additions & 0 deletions python/aisimulate/collector/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -587,6 +587,8 @@ Large-EP MoE uses stock `moe_perf` for modeled local expert compute and
`moe_a2a_perf` records latency in microseconds; `load_moe_a2a_data` converts
leaves to milliseconds. Stock `moe_perf` retains its existing timing contract.

Stock `moe_perf` also supports optional Boolean `default_eligible` selection metadata. The [Core API contract](../../../docs/core-api.md#choosing-a-forward-pass-api) defines automatic and exact-source selection. Finalization validates non-null Boolean flags and the named source required by `false`, excludes the flag from measurement identity, and preserves existing annotations when merging a legacy collection that omits the column. Collectors do not infer eligibility from kernel labels or row order.

**Legacy-overwrite caveats.** A new-schema row replaces a legacy-adapted
leaf only at the *same* key, and the legacy adapters derive their node/EP
geometry rather than reading it:
Expand Down
47 changes: 45 additions & 2 deletions python/aisimulate/collector/helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -962,7 +962,7 @@ def convert_perf_csv_to_parquet(
When ``merge_existing`` is True and a parquet already exists at the target,
the new rows are merged into it instead of overwriting: the existing and new
rows are concatenated (new last) and deduplicated on their identity key —
every column except the measured metrics (``PERF_METRIC_COLUMNS``) — keeping
every column except measured metrics and MoE selection metadata — keeping
the newest row per key. This makes finalization idempotent and accumulative:
a resumed / ``--resume-retry-failed`` / batched collection extends the
parquet instead of clobbering it with only the current run's subset.
Expand Down Expand Up @@ -1885,6 +1885,7 @@ def _prepare_perf_file(
raise RuntimeError(f"Collector staging file changed after preflight: {csv_path}")
snapshot_file.seek(0)
table = pc_csv.read_csv(snapshot_file)
_validate_moe_default_eligibility(table, parquet_path, pa=pa)
new_rows = table.num_rows
merged_existing = False
merge_target = None
Expand Down Expand Up @@ -2523,6 +2524,23 @@ def cleanup_unjournaled_perf_preparations(
return True


def _validate_moe_default_eligibility(table, parquet_path: Path, *, pa):
if parquet_path.name != "moe_perf.parquet" or "default_eligible" not in table.column_names:
return
flags = table.column("default_eligible")
if not pa.types.is_boolean(flags.type) or flags.null_count:
raise ValueError(f"{parquet_path}: default_eligible must contain only non-null Boolean values")
if "kernel_source" in table.column_names:
sources = table.column("kernel_source").to_pylist()
else:
sources = [None] * table.num_rows
if any(
not eligible and (not isinstance(source, str) or not source.strip())
for eligible, source in zip(flags.to_pylist(), sources, strict=True)
):
raise ValueError(f"{parquet_path}: default_eligible=false requires a nonblank string kernel_source")


def _merge_perf_rows(new_table, old_table, parquet_path: Path, *, pa):
"""Merge freshly-collected rows into an existing perf parquet, keeping the
newest row per identity key. Returns ``(table, merged_existing,
Expand All @@ -2531,6 +2549,25 @@ def _merge_perf_rows(new_table, old_table, parquet_path: Path, *, pa):
import pandas as pd

log = logging.getLogger(__name__)
_validate_moe_default_eligibility(old_table, parquet_path, pa=pa)

# Eligibility annotates a measurement; it does not create a new identity.
# Legacy rows are eligible, but a legacy recollection must not erase an
# existing annotation on the same physical key.
metadata_columns = {"default_eligible"} if parquet_path.name == "moe_perf.parquet" else set()
inherit_eligibility = False
if metadata_columns:
old_has_flag = "default_eligible" in old_table.column_names
new_has_flag = "default_eligible" in new_table.column_names
inherit_eligibility = old_has_flag and not new_has_flag
if old_has_flag and not new_has_flag:
new_table = new_table.append_column(
"default_eligible", pa.array([True] * new_table.num_rows, type=pa.bool_())
)
elif new_has_flag and not old_has_flag:
old_table = old_table.append_column(
"default_eligible", pa.array([True] * old_table.num_rows, type=pa.bool_())
)

# Compare (name, type) pairs for IDENTITY columns only: matching names
# with drifted types would otherwise round-trip through pandas and
Expand Down Expand Up @@ -2586,7 +2623,13 @@ def fields(schema):
old_df = old_table.to_pandas()

new_df = new_df[old_df.columns.tolist()] # align column order
identity = [c for c in old_df.columns if c not in PERF_METRIC_COLUMNS]
identity = [c for c in old_df.columns if c not in PERF_METRIC_COLUMNS and c not in metadata_columns]
if inherit_eligibility:
annotations = old_df.drop_duplicates(subset=identity, keep="last")[identity + ["default_eligible"]]
new_df = new_df.drop(columns="default_eligible").merge(annotations, on=identity, how="left", sort=False)
flags = new_df["default_eligible"]
new_df["default_eligible"] = flags.where(flags.notna(), True).astype(bool)
new_df = new_df[old_df.columns.tolist()]
current_event_rows = len(new_df.drop_duplicates(subset=identity, keep="last"))
combined = pd.concat([old_df, new_df], ignore_index=True)
deduped = combined.drop_duplicates(subset=identity, keep="last").reset_index(drop=True)
Expand Down
4 changes: 4 additions & 0 deletions python/aisimulate/src/aisimulate/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -1154,6 +1154,10 @@ def _configured_in_flight_cap(spec: ReplaySpec) -> int | None:
def _pop_aic_timing_overrides(rank: dict[str, JSONValue], role: str) -> dict[str, JSONValue]:
aic_timing_overrides: dict[str, JSONValue] = {}
for target, aliases in _AIC_TIMING_FIELD_ALIASES.items():
if target == "moe_kernel_source":
for alias in aliases:
if rank.get(alias) is None:
rank.pop(alias, None)
configured = [alias for alias in aliases if alias in rank]
if len(configured) > 1:
names = ", ".join(configured)
Expand Down
Loading
Loading