From fffb681f3c01a571a177b5e4b7d724748040afc2 Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 12:23:35 +0800 Subject: [PATCH 01/10] test(rust-core): add HYBRID/EMPIRICAL parity cases + open the rust routing gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Test-first guard for porting the util-space empirical layer (PR #1238, issue #1333 §4.6 option b) to the compiled engine: - EngineStepParityCase grows database_mode / transfer_policy / moe_quant_mode knobs, threaded through all four surfaces (static, mixed-step, agg, disagg). - 16 HYBRID_CASES probed with capture_provenance() so every transfer tier is pinned by real shipped data: xop+xshape (MiniMax-M3 on vllm and sglang), xquant / xprofile (forced MoE quants on b200/vllm), attention head_size xshape (MiMo-V2-Flash head_dim=192), own-data EMPIRICAL across 7 op families, HYBRID==SILICON invariance (Kimi-K2.5), policy gating (off/balanced -> symmetric EmpiricalNotImplementedError), and a genuine ladder miss (NVFP4 MoE on h200). - Hybrid cases assert at rtol=1e-4: EMPIRICAL is close to SILICON by design on collected shapes (and identical at exact grid hits), so the 1% smoke tolerance cannot distinguish a mode-ignoring engine. The GLM-5 EMPIRICAL case runs at isl=1536 for the same reason. - should_use_rust_engine_step now admits HYBRID/EMPIRICAL (SOL modes still delegate to the Python step), so these cases exercise the Rust path for real. Current census: 46 failed / 18 passed — the failures are the port's guard; the passes are the error-symmetry contracts. Existing SILICON smoke cases are unaffected. Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- .../parity_tests/test_engine_step_parity.py | 310 ++++++++++++++++-- src/aiconfigurator/sdk/rust_engine_step.py | 21 +- 2 files changed, 301 insertions(+), 30 deletions(-) diff --git a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py index 9be25d664..31b649324 100644 --- a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py +++ b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py @@ -19,7 +19,7 @@ import pytest from aiconfigurator.cli.api import cli_estimate -from aiconfigurator.sdk import config, perf_database, rust_engine_step +from aiconfigurator.sdk import common, config, perf_database, rust_engine_step from aiconfigurator.sdk.backends.factory import get_backend from aiconfigurator.sdk.models import get_model @@ -48,6 +48,12 @@ class EngineStepParityCase: disagg_prefill_num_workers: int = 1 disagg_decode_batch_size: int = 4 disagg_decode_num_workers: int = 1 + # HYBRID/EMPIRICAL parity knobs: database mode, transfer-policy preset + # (None = default ALL_TRANSFERS), and a forced MoE quant used to steer a + # query into a specific transfer tier (xquant/xprofile) on real data. + database_mode: str = "SILICON" + transfer_policy: str | None = None + moe_quant_mode: str | None = None SMOKE_CASES = [ @@ -603,6 +609,9 @@ def _static_metrics( "moe_ep_size": case.moe_ep_size, "stride": 1, "engine_step_backend": engine_step_backend, + "database_mode": case.database_mode, + "transfer_policy": case.transfer_policy, + "moe_quant_mode": case.moe_quant_mode, } context_ms = _safe_value( lambda: ( @@ -647,6 +656,9 @@ def call(): moe_tp_size=case.moe_tp_size, moe_ep_size=case.moe_ep_size, engine_step_backend=engine_step_backend, + database_mode=case.database_mode, + transfer_policy=case.transfer_policy, + moe_quant_mode=case.moe_quant_mode, ) # Errors propagate from a single call site — capture once, surface @@ -688,6 +700,9 @@ def call(): decode_batch_size=case.disagg_decode_batch_size, decode_num_workers=case.disagg_decode_num_workers, engine_step_backend=engine_step_backend, + database_mode=case.database_mode, + transfer_policy=case.transfer_policy, + moe_quant_mode=case.moe_quant_mode, ) err: _ErrorSentinel | None = None @@ -721,22 +736,45 @@ def _mix_step_shape(case: EngineStepParityCase) -> dict: } -def _python_mixed_step_ms(case: EngineStepParityCase) -> float: - """Python `_get_mix_step_latency` for the case's mix-step shape.""" - database = _quiet_call(perf_database.get_database, case.system_name, case.backend_name, case.backend_version) - if database is None: - raise RuntimeError( - f"failed to load perf database for {case.system_name}/{case.backend_name}/{case.backend_version}" - ) - model_config = config.ModelConfig( +def _case_database(case: EngineStepParityCase): + """Perf database for a case: the plain cached database for SILICON + defaults, or a mode/policy-configured query view for HYBRID/EMPIRICAL + cases (mirrors what `cli_estimate` builds internally).""" + if case.database_mode == "SILICON" and case.transfer_policy is None: + return _quiet_call(perf_database.get_database, case.system_name, case.backend_name, case.backend_version) + return _quiet_call( + perf_database.get_database_view, + case.system_name, + case.backend_name, + case.backend_version, + # Mirror `cli_estimate`: non-SILICON modes tolerate missing tables at + # load (the empirical layer covers the gaps at query time). + allow_missing_data=case.database_mode != "SILICON", + database_mode=case.database_mode, + transfer_policy=case.transfer_policy, + ) + + +def _case_model_config(case: EngineStepParityCase) -> config.ModelConfig: + return config.ModelConfig( tp_size=case.tp_size, pp_size=case.pp_size, attention_dp_size=case.attention_dp_size, moe_tp_size=case.moe_tp_size, moe_ep_size=case.moe_ep_size, cp_size=case.cp_size, + moe_quant_mode=(common.MoEQuantMode[case.moe_quant_mode] if case.moe_quant_mode else None), ) - model = _quiet_call(get_model, case.model_path, model_config, case.backend_name) + + +def _python_mixed_step_ms(case: EngineStepParityCase) -> float: + """Python `_get_mix_step_latency` for the case's mix-step shape.""" + database = _case_database(case) + if database is None: + raise RuntimeError( + f"failed to load perf database for {case.system_name}/{case.backend_name}/{case.backend_version}" + ) + model = _quiet_call(get_model, case.model_path, _case_model_config(case), case.backend_name) backend = get_backend(case.backend_name) runtime_config = config.RuntimeConfig( batch_size=case.batch_size, @@ -762,20 +800,12 @@ def _python_mixed_step_ms(case: EngineStepParityCase) -> float: def _rust_mixed_step_ms(case: EngineStepParityCase) -> float: """Rust mix-step latency for the case's mix-step shape (same FPM).""" - database = _quiet_call(perf_database.get_database, case.system_name, case.backend_name, case.backend_version) + database = _case_database(case) if database is None: raise RuntimeError( f"failed to load perf database for {case.system_name}/{case.backend_name}/{case.backend_version}" ) - model_config = config.ModelConfig( - tp_size=case.tp_size, - pp_size=case.pp_size, - attention_dp_size=case.attention_dp_size, - moe_tp_size=case.moe_tp_size, - moe_ep_size=case.moe_ep_size, - cp_size=case.cp_size, - ) - model = _quiet_call(get_model, case.model_path, model_config, case.backend_name) + model = _quiet_call(get_model, case.model_path, _case_model_config(case), case.backend_name) shape = _mix_step_shape(case) return rust_engine_step.estimate_mixed_step_latency_with_rust( model, @@ -790,6 +820,7 @@ def _rust_mixed_step_ms(case: EngineStepParityCase) -> float: def _parity_mismatch_reason( comparisons: dict[str, tuple[float | _ErrorSentinel, float | _ErrorSentinel]], + rtol: float = PARITY_RTOL, ) -> str | None: """Compare Python and Rust per-metric outputs with three valid pairings: @@ -818,14 +849,14 @@ def _parity_mismatch_reason( rows.append(f"{name:<{metric_width}} {py_repr:>10} {rs_repr:>10} {'-':>10} {'-':>10} {'-':>10} asym") continue # Both compute — apply numeric tolerance. - allowed = max(abs(python_value) * PARITY_RTOL, 1e-9) + allowed = max(abs(python_value) * rtol, 1e-9) delta = rust_value - python_value delta_pct = delta / abs(python_value) * 100 if python_value else float("inf") status = "drift" if abs(delta) > allowed else "ok" has_mismatch = has_mismatch or status == "drift" rows.append( f"{name:<{metric_width}} {python_value:>10.3f} {rust_value:>10.3f} " - f"{delta:>10.3f} {delta_pct:>9.2f}% {PARITY_RTOL * 100:>9.2f}% {status:>6}" + f"{delta:>10.3f} {delta_pct:>9.2f}% {rtol * 100:>9.2f}% {status:>6}" ) if not has_mismatch: return None @@ -1019,3 +1050,238 @@ def test_cp_parity( reason = _parity_mismatch_reason(_mixed_step_comparison_metrics(case)) assert reason is None, reason + + +# HYBRID / EMPIRICAL parity cases: the util-space empirical layer +# (`sdk/operations/util_empirical.py`, PR #1238) ported to the compiled engine +# (issue #1333 §4.6 option b). Every case below was probed on the Python side +# with `capture_provenance()` so each transfer tier of the ladder +# (own-data empirical -> xshape -> xquant -> xprofile -> xop) is pinned by at +# least one real-data case, plus the two contract cases (HYBRID==SILICON +# invariance on covered configs; symmetric EmpiricalNotImplementedError on +# genuine coverage misses). +# +# These cases assert at a much tighter tolerance than the SILICON smoke suite. +# util-space empirical is close to silicon BY DESIGN on collected configs +# (offline study: ~1.9% mean APE, many rows < 1%), so at the 1% smoke rtol a +# Rust engine that silently ignored the mode and computed pure SILICON would +# still pass several EMPIRICAL cases. Both engines run the same f64 math over +# the same rows, so the faithful port agrees to ~1e-9; 1e-4 keeps headroom +# while making a silicon fallback (0.3-5% off) unmissable. +HYBRID_PARITY_RTOL = 1e-4 +HYBRID_CASES = [ + # xop + xshape: MiniMax-M3 is all-MoE + MSA sparse attention with NO own + # silicon data anywhere — MSA borrows DSA's util (xop, `msa.py`), the MoE + # shape borrows the nearest collected sibling (xshape, `moe.py`). Probed: + # static ctx/gen = 69.588/5.426 ms, tags {xop, xshape}. + pytest.param( + EngineStepParityCase( + model_path="MiniMaxAI/MiniMax-M3", + database_mode="HYBRID", + ), + id="minimax-m3-b200-vllm-019-hybrid-xop", + ), + # Same rescue on the sglang tables (DSA util source = sglang 0.5.14 data). + # Probed: 39.490/2.994 ms, tags {xop, xshape}. + pytest.param( + EngineStepParityCase( + model_path="MiniMaxAI/MiniMax-M3", + backend_name="sglang", + backend_version="0.5.14", + database_mode="HYBRID", + ), + id="minimax-m3-b200-sglang-0514-hybrid-xop", + ), + # Transfer-policy gating: with xop disabled ("off" and "balanced" presets) + # MSA must raise EmpiricalNotImplementedError on BOTH engines + # (error-symmetry). Guards that the Rust port honours the policy at query + # time instead of always transferring. + pytest.param( + EngineStepParityCase( + model_path="MiniMaxAI/MiniMax-M3", + database_mode="HYBRID", + transfer_policy="off", + ), + id="minimax-m3-b200-vllm-019-hybrid-policy-off", + ), + pytest.param( + EngineStepParityCase( + model_path="MiniMaxAI/MiniMax-M3", + database_mode="HYBRID", + transfer_policy="balanced", + ), + id="minimax-m3-b200-vllm-019-hybrid-policy-balanced", + ), + # xquant: forced MoE quant w4a16_mxfp4_cutlass is uncollected on + # b200/vllm/0.19.0 but shares the (memory=0.5, compute=1) profile with + # collected int4_wo / w4a16_mxfp4 — the ladder lands on the xquant tier. + # Probed: 90.578/10.908 ms, tags {xquant}. + pytest.param( + EngineStepParityCase( + model_path="Qwen/Qwen3-235B-A22B", + database_mode="HYBRID", + moe_quant_mode="w4a16_mxfp4_cutlass", + ), + id="qwen3-235b-a22b-b200-vllm-019-hybrid-xquant", + ), + # xprofile: forced MoE quant w4afp8 (memory=0.5, compute=2) has NO + # collected same-profile sibling on b200/vllm/0.19.0 — the ladder falls + # through to the cross-profile tier with the util-level rescale. + # Probed: 47.755/8.450 ms, tags {xprofile}. + pytest.param( + EngineStepParityCase( + model_path="Qwen/Qwen3-235B-A22B", + database_mode="HYBRID", + moe_quant_mode="w4afp8", + ), + id="qwen3-235b-a22b-b200-vllm-019-hybrid-xprofile", + ), + # Attention cross-head_size xshape: MiMo-V2-Flash has head_dim=192 while + # b200/vllm/0.19.0 collected only {128, 256} — SILICON raises, HYBRID + # borrows the nearest collected head_size (`attention.py` ctx + gen + # reference grids). Probed: 33.253/3.499 ms, tags {xshape}. + pytest.param( + EngineStepParityCase( + model_path="XiaomiMiMo/MiMo-V2-Flash", + database_mode="HYBRID", + ), + id="mimo-v2-flash-b200-vllm-019-hybrid-attn-xshape", + ), + # HYBRID==SILICON invariance: Kimi-K2.5 on b200/vllm/0.19.0 is fully + # covered by silicon data (probed worst-provenance = silicon, no empirical + # tier fires). The hybrid layer must not perturb covered queries; this + # pins Rust-HYBRID == Python-HYBRID (== SILICON) on a collected config. + pytest.param( + EngineStepParityCase( + model_path="moonshotai/Kimi-K2.5", + database_mode="HYBRID", + ), + id="kimi-k25-b200-vllm-019-hybrid-invariant", + ), + # Ladder miss: NVFP4 MoE on Hopper has no own-shape, cross-shape, or + # sibling reference anywhere in the h200/vllm/0.19.0 tables — Python + # raises EmpiricalNotImplementedError; the Rust port must fail the same + # query point (error-symmetry), never fabricate a SOL/constant value. + pytest.param( + EngineStepParityCase( + model_path="nvidia/MiniMax-M2.5-NVFP4", + system_name="h200_sxm", + database_mode="HYBRID", + ), + id="minimax-m25-nvfp4-h200-vllm-019-hybrid-miss", + ), + # EMPIRICAL mode: every data-backed op answers SOL(query)/util from its + # own collected slice — the broadest guard of the ported util math (grid + # build, k=2 IDW in normalized log space, per-axis boundary clamp) across + # op families: dense GEMM+GQA, MoE, MLA, DSA (vllm + sglang), fp8_block + # MoE, and the state-space (Mamba2 SOL-degradation) path. + pytest.param( + EngineStepParityCase( + model_path="Qwen/Qwen3-32B", + database_mode="EMPIRICAL", + ), + id="qwen3-32b-b200-vllm-019-empirical", + ), + pytest.param( + EngineStepParityCase( + model_path="Qwen/Qwen3-235B-A22B", + database_mode="EMPIRICAL", + ), + id="qwen3-235b-a22b-b200-vllm-019-empirical", + ), + pytest.param( + EngineStepParityCase( + model_path="moonshotai/Kimi-K2.5", + database_mode="EMPIRICAL", + ), + id="kimi-k25-b200-vllm-019-empirical", + ), + pytest.param( + EngineStepParityCase( + model_path="deepseek-ai/DeepSeek-V3.2", + database_mode="EMPIRICAL", + ), + id="deepseek-v32-b200-vllm-019-empirical", + ), + # Off-grid shape on purpose: at the smoke shape (isl=1024, b=1) every GLM-5 + # op lands exactly on collected grid points, where util reconstruction + # returns SOL/util == measured — EMPIRICAL degenerates to SILICON and the + # case cannot distinguish a mode-ignoring engine. isl=1536 separates the + # two by ~3.1% (probed) while both still compute. + pytest.param( + EngineStepParityCase( + model_path="zai-org/GLM-5", + backend_name="sglang", + backend_version="0.5.14", + isl=1536, + database_mode="EMPIRICAL", + ), + id="glm5-b200-sglang-0514-empirical", + ), + pytest.param( + EngineStepParityCase( + model_path="MiniMaxAI/MiniMax-M2.5", + database_mode="EMPIRICAL", + ), + id="minimax-m25-b200-vllm-019-empirical", + ), + pytest.param( + EngineStepParityCase( + model_path="nvidia/Nemotron-H-56B-Base-8K", + database_mode="EMPIRICAL", + ), + id="nemotron-h-56b-b200-vllm-019-empirical", + ), +] + + +class TestRustEngineStepHybridStaticParity: + @pytest.mark.parametrize("case", HYBRID_CASES) + def test_hybrid_parity( + self, + case: EngineStepParityCase, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _prepare_rust_core(monkeypatch) + + reason = _parity_mismatch_reason(_static_comparison_metrics(case), rtol=HYBRID_PARITY_RTOL) + assert reason is None, reason + + +class TestRustEngineStepHybridMixedStepParity: + @pytest.mark.parametrize("case", HYBRID_CASES) + def test_hybrid_parity( + self, + case: EngineStepParityCase, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _prepare_rust_core(monkeypatch) + + reason = _parity_mismatch_reason(_mixed_step_comparison_metrics(case), rtol=HYBRID_PARITY_RTOL) + assert reason is None, reason + + +class TestRustEngineStepHybridAggParity: + @pytest.mark.parametrize("case", HYBRID_CASES) + def test_hybrid_parity( + self, + case: EngineStepParityCase, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _prepare_rust_core(monkeypatch) + + reason = _parity_mismatch_reason(_agg_comparison_metrics(case), rtol=HYBRID_PARITY_RTOL) + assert reason is None, reason + + +class TestRustEngineStepHybridDisaggParity: + @pytest.mark.parametrize("case", HYBRID_CASES) + def test_hybrid_parity( + self, + case: EngineStepParityCase, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _prepare_rust_core(monkeypatch) + + reason = _parity_mismatch_reason(_disagg_comparison_metrics(case), rtol=HYBRID_PARITY_RTOL) + assert reason is None, reason diff --git a/src/aiconfigurator/sdk/rust_engine_step.py b/src/aiconfigurator/sdk/rust_engine_step.py index 13fcd600c..23a6b02d6 100644 --- a/src/aiconfigurator/sdk/rust_engine_step.py +++ b/src/aiconfigurator/sdk/rust_engine_step.py @@ -221,25 +221,30 @@ def _normalize_tuning_iterations(iterations: dict[str, Any] | list[Any]) -> list return iterations +# Database modes the compiled engine answers itself. SILICON plus the +# util-space empirical layer (HYBRID / EMPIRICAL, mirroring +# `sdk/operations/util_empirical.py`); the SOL diagnostic modes stay on the +# Python step. +_RUST_SUPPORTED_DATABASE_MODES = {"SILICON", "HYBRID", "EMPIRICAL"} + + def should_use_rust_engine_step(runtime_config: RuntimeConfig, database: Any = None) -> bool: """Route to the compiled engine only when it can give the SAME answer. - The compiled engine implements the SILICON path only (no util_empirical - layer), so HYBRID/EMPIRICAL databases must stay on the Python step: - wherever silicon data misses, Python fills in empirically while the - compiled engine would fail the config -- delegating keeps the two - backends answer-identical instead of capability-divergent (parity by - delegation; the empirical-layer port is tracked in issue #1333). + The compiled engine implements the SILICON path and the util-space + empirical layer (HYBRID / EMPIRICAL). The SOL/SOL_FULL diagnostic modes + stay on the Python step -- delegating keeps the two backends + answer-identical instead of capability-divergent. """ backend = getattr(runtime_config, "engine_step_backend", None) or os.environ.get(ENGINE_STEP_BACKEND_ENV) if str(backend or "python").lower() != "rust": return False if database is not None: mode = getattr(database, "get_default_database_mode", lambda: None)() - if mode is not None and getattr(mode, "name", str(mode)) != "SILICON": + if mode is not None and getattr(mode, "name", str(mode)) not in _RUST_SUPPORTED_DATABASE_MODES: logger.debug( "engine-step backend 'rust' requested but database_mode=%s; " - "using the python step (compiled engine is SILICON-only).", + "using the python step (compiled engine implements SILICON/HYBRID/EMPIRICAL only).", getattr(mode, "name", mode), ) return False From 6b9f866b9cf7078ee3d6e18b617f3e7bc5aeb1ae Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 12:40:38 +0800 Subject: [PATCH 02/10] feat(rust-core): util-empirical core + database_mode/transfer_policy plumbing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First half of the HYBRID/EMPIRICAL port (issue #1333 §4.6 option b): - operators/util_empirical.rs: UtilGrid (k=2 inverse-distance weighting in per-axis normalised log space, exact-hit preservation, stable-order tie breaks, per-axis boundary clamp), build_samples, estimate (typed EmpiricalNotImplemented miss), nearest_candidate_index, and a keyed UtilGridCache (tables are immutable per database instance, so no id()-based invalidation is needed). Unit tests mirror the math anchors of tests/unit/sdk/test_util_empirical.py. - AicError::EmpiricalNotImplemented mirrors Python's EmpiricalNotImplementedError semantics (coverage gap, never fabricated). - common/enums.rs: TransferKind + TransferPolicy (wire = explicit kind tokens; Python resolves presets before serialising). MoeQuantMode gains w4a8_mxfp4_mxfp8_trtllm / w4a16_mxfp4_cutlass, which the DataType wire enum previously rejected outright (found by the xquant parity case). - EngineConfig carries database_mode + transfer_policy (serde defaults = SILICON/ALL for old specs); Engine::from_spec_bytes configures the loaded PerfDatabase, which now owns mode + policy + a util-grid cache. - Python side: _engine_config_dict emits both fields off the live database view, and the EngineHandle cache key includes them so HYBRID/EMPIRICAL views never reuse a SILICON handle. No op consumes the mode yet — per-op estimate dispatch lands next. The kimi HYBRID==SILICON invariance parity case stays green through the plumbing; the xquant case's Rust failure moved from spec-decode rejection to the MoE table miss where the empirical layer will plug in. Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- rust/aiconfigurator-core/src/common/enums.rs | 80 +++ rust/aiconfigurator-core/src/common/error.rs | 6 + rust/aiconfigurator-core/src/config.rs | 18 + .../aiconfigurator-core/src/engine/runtime.rs | 8 +- rust/aiconfigurator-core/src/engine/spec.rs | 2 + rust/aiconfigurator-core/src/fpm/tests.rs | 2 + rust/aiconfigurator-core/src/memory.rs | 2 + rust/aiconfigurator-core/src/operators/mod.rs | 1 + .../src/operators/util_empirical.rs | 455 ++++++++++++++++++ .../src/perf_database/mod.rs | 27 ++ rust/aiconfigurator-core/src/py.rs | 2 + .../tests/memory_round_trip.rs | 2 + src/aiconfigurator/sdk/engine.py | 34 ++ src/aiconfigurator/sdk/rust_engine_step.py | 25 +- 14 files changed, 662 insertions(+), 2 deletions(-) create mode 100644 rust/aiconfigurator-core/src/operators/util_empirical.rs diff --git a/rust/aiconfigurator-core/src/common/enums.rs b/rust/aiconfigurator-core/src/common/enums.rs index b53c4011f..c4b3a9d07 100644 --- a/rust/aiconfigurator-core/src/common/enums.rs +++ b/rust/aiconfigurator-core/src/common/enums.rs @@ -66,6 +66,78 @@ impl Default for DatabaseMode { } } +/// A way the empirical path may borrow utilisation when an op's own slice +/// has no data. Mirrors Python `common.TransferKind`; each kind is +/// independently enabled by the transfer policy carried on the engine +/// config. Ordered by decreasing confidence. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum TransferKind { + /// Cross-shape within the SAME quant (nearest collected config). + XShape, + /// Cross-quant within the SAME (memory, compute) profile. + XQuant, + /// Cross-quant across profiles (rescaled by a util-level ratio). + XProfile, + /// Cross-op (borrow a related op's util, e.g. MSA<-DSA, via util_scale). + XOp, +} + +/// Enabled transfer kinds. Python resolves presets +/// (`off`/`conservative`/`balanced`/`aggressive`) before serialising, so the +/// wire form is always an explicit list of kind tokens; `None` on the wire +/// means the default ALL-transfers policy. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub struct TransferPolicy { + pub xshape: bool, + pub xquant: bool, + pub xprofile: bool, + pub xop: bool, +} + +impl TransferPolicy { + pub const ALL: TransferPolicy = TransferPolicy { + xshape: true, + xquant: true, + xprofile: true, + xop: true, + }; + pub const OFF: TransferPolicy = TransferPolicy { + xshape: false, + xquant: false, + xprofile: false, + xop: false, + }; + + pub fn contains(&self, kind: TransferKind) -> bool { + match kind { + TransferKind::XShape => self.xshape, + TransferKind::XQuant => self.xquant, + TransferKind::XProfile => self.xprofile, + TransferKind::XOp => self.xop, + } + } + + /// Parse the wire form: `None` -> ALL (Python's default policy), a list + /// of kind tokens -> exactly those kinds (an empty list is `off`). + pub fn from_wire(kinds: Option<&[String]>) -> Result { + let Some(kinds) = kinds else { + return Ok(TransferPolicy::ALL); + }; + let mut policy = TransferPolicy::OFF; + for token in kinds { + match token.as_str() { + "xshape" => policy.xshape = true, + "xquant" => policy.xquant = true, + "xprofile" => policy.xprofile = true, + "xop" => policy.xop = true, + other => return Err(format!("unknown transfer kind {other:?}")), + } + } + Ok(policy) + } +} + /// Per-variant payload mirroring Python's /// `QuantMapping = namedtuple("QuantMapping", ["memory", "compute", "name"])`. /// @@ -126,6 +198,8 @@ pub enum MoeQuantMode { Nvfp4, W4a16Mxfp4, W4a8Mxfp4Mxfp8, + W4a8Mxfp4Mxfp8Trtllm, + W4a16Mxfp4Cutlass, } impl MoeQuantMode { @@ -141,6 +215,12 @@ impl MoeQuantMode { Self::W4a8Mxfp4Mxfp8 => { QuantMapping { memory: 0.5, compute: 2.0, name: "w4a8_mxfp4_mxfp8" } } + Self::W4a8Mxfp4Mxfp8Trtllm => { + QuantMapping { memory: 0.5, compute: 2.0, name: "w4a8_mxfp4_mxfp8_trtllm" } + } + Self::W4a16Mxfp4Cutlass => { + QuantMapping { memory: 0.5, compute: 1.0, name: "w4a16_mxfp4_cutlass" } + } } } diff --git a/rust/aiconfigurator-core/src/common/error.rs b/rust/aiconfigurator-core/src/common/error.rs index 6234e7006..e5d635f22 100644 --- a/rust/aiconfigurator-core/src/common/error.rs +++ b/rust/aiconfigurator-core/src/common/error.rs @@ -34,6 +34,12 @@ pub enum AicError { ModelConfig(String), #[error("perf database error: {0}")] PerfDatabase(String), + /// The HYBRID/EMPIRICAL path found no calibration data for the requested + /// slice — no own-shape, cross-shape, or sibling transfer reference. + /// Mirrors Python's `EmpiricalNotImplementedError`: a coverage gap, never + /// converted into a fabricated `SOL / constant` value. + #[error("empirical estimation not implemented: {0}")] + EmpiricalNotImplemented(String), #[error("I/O error at {path}: {source}")] Io { path: PathBuf, diff --git a/rust/aiconfigurator-core/src/config.rs b/rust/aiconfigurator-core/src/config.rs index fbfce9835..03f00c6f5 100644 --- a/rust/aiconfigurator-core/src/config.rs +++ b/rust/aiconfigurator-core/src/config.rs @@ -72,6 +72,20 @@ pub struct EngineConfig { #[serde(default)] pub perf_db_sources: PerfDbSources, + /// Perf-database lookup mode (Python's `database._default_database_mode`). + /// SILICON queries collected tables only; HYBRID falls back to the + /// util-space empirical layer on a typed silicon miss; EMPIRICAL always + /// answers `SOL/util`. Absent on old specs -> Silicon (back-compat). + #[serde(default)] + pub database_mode: crate::common::enums::DatabaseMode, + + /// Enabled empirical transfer kinds as explicit tokens (`xshape` / + /// `xquant` / `xprofile` / `xop`). Python resolves preset names before + /// serialising, so no preset vocabulary exists on the wire. `None` = + /// the default ALL-transfers policy (mirrors `common.ALL_TRANSFERS`). + #[serde(default)] + pub transfer_policy: Option>, + #[serde(default)] pub extra: BTreeMap, } @@ -193,6 +207,10 @@ pub enum DataType { W4a16Mxfp4, #[serde(rename = "w4a8_mxfp4_mxfp8")] W4a8Mxfp4Mxfp8, + #[serde(rename = "w4a8_mxfp4_mxfp8_trtllm")] + W4a8Mxfp4Mxfp8Trtllm, + #[serde(rename = "w4a16_mxfp4_cutlass")] + W4a16Mxfp4Cutlass, } #[cfg(test)] diff --git a/rust/aiconfigurator-core/src/engine/runtime.rs b/rust/aiconfigurator-core/src/engine/runtime.rs index a6d2a3b92..bf8d1c2da 100644 --- a/rust/aiconfigurator-core/src/engine/runtime.rs +++ b/rust/aiconfigurator-core/src/engine/runtime.rs @@ -18,6 +18,7 @@ use std::sync::Arc; +use crate::common::enums::TransferPolicy; use crate::common::error::AicError; use crate::engine::spec::EngineSpec; use crate::operators::Op; @@ -164,13 +165,16 @@ impl Engine { // The spec's own `systems_path` wins when present; otherwise fall back // to the `systems_root` argument. let systems_root = spec.engine.systems_path.as_deref().unwrap_or(systems_root); + let transfer_policy = TransferPolicy::from_wire(spec.engine.transfer_policy.as_deref()) + .map_err(AicError::InvalidEngineConfig)?; let db = PerfDatabase::load_with_sources( systems_root, &spec.engine.system_name, spec.engine.backend.as_str(), version, &spec.engine.perf_db_sources, - )?; + )? + .with_mode(spec.engine.database_mode, transfer_policy); Engine::build(spec, Arc::new(db)) } @@ -606,6 +610,8 @@ mod tests { nextn_accept_rates: None, }), perf_db_sources: Default::default(), + database_mode: Default::default(), + transfer_policy: None, extra: BTreeMap::new(), } } diff --git a/rust/aiconfigurator-core/src/engine/spec.rs b/rust/aiconfigurator-core/src/engine/spec.rs index 8a950e91e..181c7034b 100644 --- a/rust/aiconfigurator-core/src/engine/spec.rs +++ b/rust/aiconfigurator-core/src/engine/spec.rs @@ -618,6 +618,8 @@ mod tests { nextn_accept_rates: Some(vec![0.85, 0.3, 0.0, 0.0, 0.0]), }), perf_db_sources: Default::default(), + database_mode: Default::default(), + transfer_policy: None, extra: BTreeMap::new(), } } diff --git a/rust/aiconfigurator-core/src/fpm/tests.rs b/rust/aiconfigurator-core/src/fpm/tests.rs index ee3ac6a0f..39018bf39 100644 --- a/rust/aiconfigurator-core/src/fpm/tests.rs +++ b/rust/aiconfigurator-core/src/fpm/tests.rs @@ -106,6 +106,8 @@ fn fixture_engine_config() -> EngineConfig { }, speculative: None, perf_db_sources: Default::default(), + database_mode: Default::default(), + transfer_policy: None, extra: BTreeMap::new(), } } diff --git a/rust/aiconfigurator-core/src/memory.rs b/rust/aiconfigurator-core/src/memory.rs index d7a7e107e..1d903897e 100644 --- a/rust/aiconfigurator-core/src/memory.rs +++ b/rust/aiconfigurator-core/src/memory.rs @@ -404,6 +404,8 @@ fn dtype_str(dt: &crate::DataType) -> &'static str { W4afp8 => "w4afp8", W4a16Mxfp4 => "w4a16_mxfp4", W4a8Mxfp4Mxfp8 => "w4a8_mxfp4_mxfp8", + W4a8Mxfp4Mxfp8Trtllm => "w4a8_mxfp4_mxfp8_trtllm", + W4a16Mxfp4Cutlass => "w4a16_mxfp4_cutlass", } } diff --git a/rust/aiconfigurator-core/src/operators/mod.rs b/rust/aiconfigurator-core/src/operators/mod.rs index d7e3ad7a4..0bae8d64d 100644 --- a/rust/aiconfigurator-core/src/operators/mod.rs +++ b/rust/aiconfigurator-core/src/operators/mod.rs @@ -28,6 +28,7 @@ pub mod moe; pub mod moe_dispatch; pub mod op; pub mod overlap; +pub mod util_empirical; pub mod vision; pub mod wideep_mla; pub mod wideep_moe; diff --git a/rust/aiconfigurator-core/src/operators/util_empirical.rs b/rust/aiconfigurator-core/src/operators/util_empirical.rs new file mode 100644 index 000000000..4ec7bbb76 --- /dev/null +++ b/rust/aiconfigurator-core/src/operators/util_empirical.rs @@ -0,0 +1,455 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Data-calibrated empirical estimation via SOL-utilization. +//! +//! Mirrors `aiconfigurator.sdk.operations.util_empirical`: each op's +//! empirical estimate is `latency = SOL(query) / util` where +//! `util = SOL / measured > 0` is read best-effort from collected samples in +//! per-axis normalised log space. `util` is an effective calibration factor, +//! not a bounded physical efficiency (it may exceed 1); it is never clamped. +//! Every grid uses the same two-neighbour inverse-distance weighting +//! (`k=2`, `p=1`) without requiring a Cartesian product; queries outside the +//! measured range are clamped per axis before neighbour selection, so +//! extrapolation freezes boundary utilization. +//! +//! When *no* samples exist for the requested slice (no own-shape, no +//! cross-shape/sibling transfer reference), [`estimate`] returns +//! [`AicError::EmpiricalNotImplemented`] rather than a fabricated +//! `SOL / constant` — coverage gaps surface honestly, exactly like Python's +//! `EmpiricalNotImplementedError`. Genuinely table-less ops (mem / p2p / +//! element-wise) keep their analytic formulas and never call [`estimate`]. +//! +//! Divergences from the Python module, by design: +//! - No provenance contextvar: provenance capture feeds the Python-side +//! support matrix; the compiled engine only returns latencies. Reference +//! grids still carry their provenance tag for cache keying. +//! - Caching is the caller's concern: Python keys grids by `id(node)` because +//! database views share mutable table objects; Rust perf tables are +//! immutable after load, so per-op wiring caches grids in a +//! [`UtilGridCache`] keyed by the op's slice identity. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use crate::common::error::AicError; + +/// One collected calibration point: continuous-axis coordinates plus the +/// positive effective calibration factor `util = SOL / measured`. +#[derive(Clone, Debug, PartialEq)] +pub struct UtilSample { + pub coords: Vec, + pub util: f64, +} + +impl UtilSample { + pub fn new(coords: Vec, util: f64) -> Self { + Self { coords, util } + } +} + +/// Build util samples from `(coords, latency_ms)` points and an analytic SOL. +/// +/// Mirrors Python `build_samples`: a point is kept only when both the +/// measured latency and its SOL are strictly positive (NaN fails both +/// comparisons and is dropped, matching Python truthiness + `> 0`). +pub fn build_samples(points: I, sol_fn: F) -> Vec +where + I: IntoIterator, f64)>, + F: Fn(&[f64]) -> f64, +{ + let mut samples = Vec::new(); + for (coords, latency_ms) in points { + if latency_ms > 0.0 { + let sol = sol_fn(&coords); + if sol > 0.0 { + samples.push(UtilSample::new(coords, sol / latency_ms)); + } + } + } + samples +} + +/// Two-neighbour util lookup in per-axis normalised log space. +/// +/// The query is clamped independently on every axis, then the two nearest +/// samples are combined with inverse-distance weights (`k=2`, `p=1`). Exact +/// hits return the collected utilization unchanged. Works for ragged grids +/// without operation-specific Cartesian bracketing; callers remain +/// responsible for slicing categorical/kernel-regime axes. +#[derive(Debug, Clone)] +pub struct UtilGrid { + /// Normalised log-space coordinates, one row per sample. + norm: Vec>, + utils: Vec, + mins: Vec, + spans: Vec, + /// Transfer tag of the reference slice this grid was built from + /// (`xshape` / `xquant` / ...), when borrowed from a sibling. + pub reference_provenance: Option<&'static str>, +} + +fn log_floor(value: f64) -> f64 { + value.max(1e-9).ln() +} + +impl UtilGrid { + pub fn new(samples: Vec) -> Self { + if samples.is_empty() { + return Self { + norm: Vec::new(), + utils: Vec::new(), + mins: Vec::new(), + spans: Vec::new(), + reference_provenance: None, + }; + } + let dims = samples[0].coords.len(); + let logc: Vec> = samples + .iter() + .map(|s| s.coords.iter().map(|&c| log_floor(c)).collect()) + .collect(); + let mut mins = vec![f64::INFINITY; dims]; + let mut maxs = vec![f64::NEG_INFINITY; dims]; + for row in &logc { + for (a, &v) in row.iter().enumerate() { + mins[a] = mins[a].min(v); + maxs[a] = maxs[a].max(v); + } + } + let spans: Vec = mins + .iter() + .zip(&maxs) + .map(|(&lo, &hi)| if hi - lo > 0.0 { hi - lo } else { 1.0 }) + .collect(); + let norm: Vec> = logc + .iter() + .map(|row| { + row.iter() + .enumerate() + .map(|(a, &v)| (v - mins[a]) / spans[a]) + .collect() + }) + .collect(); + let utils = samples.iter().map(|s| s.util).collect(); + Self { + norm, + utils, + mins, + spans, + reference_provenance: None, + } + } + + pub fn is_empty(&self) -> bool { + self.utils.is_empty() + } + + /// Interpolated utilization at `query`, or `None` for an empty grid. + pub fn util(&self, query: &[f64]) -> Option { + if self.utils.is_empty() { + return None; + } + // Per-axis clamp to [0, 1] freezes boundary utilization for + // out-of-range queries (mirrors `np.clip`). + let q: Vec = query + .iter() + .enumerate() + .map(|(a, &v)| ((log_floor(v) - self.mins[a]) / self.spans[a]).clamp(0.0, 1.0)) + .collect(); + let distances: Vec = self + .norm + .iter() + .map(|row| { + row.iter() + .zip(&q) + .map(|(&x, &y)| (x - y) * (x - y)) + .sum::() + .sqrt() + }) + .collect(); + // Stable argsort: ties keep sample order, so duplicate / + // log-floor-collapsed coordinates deterministically prefer the first + // sample (mirrors `np.argsort(kind="stable")`). + let mut order: Vec = (0..distances.len()).collect(); + order.sort_by(|&i, &j| distances[i].partial_cmp(&distances[j]).expect("finite distances")); + + if distances[order[0]] == 0.0 { + return Some(self.utils[order[0]]); + } + + let nearest = &order[..order.len().min(2)]; + let mut weighted = 0.0; + let mut weight_sum = 0.0; + for &i in nearest { + let w = 1.0 / distances[i]; + weighted += self.utils[i] * w; + weight_sum += w; + } + Some(weighted / weight_sum) + } +} + +/// Return `(latency_ms, util)` from the util grid, or the typed coverage +/// error. +/// +/// Mirrors Python `estimate`: `None`, empty grids, and non-positive utils all +/// surface as [`AicError::EmpiricalNotImplemented`] — there is no own-shape, +/// cross-shape, or sibling data to calibrate from, so the gap surfaces +/// instead of inventing a `SOL / constant` placeholder. +/// +/// `util_scale` is the cross-op level-alignment hook (1.0 = no change). When +/// a CROSS-OP transfer borrows a *different* op's util grid, the caller +/// passes a manual scale `k` so `latency = SOL / (util * k)`. +pub fn estimate( + sol_query: f64, + query: &[f64], + grid: Option<&UtilGrid>, + util_scale: f64, +) -> Result<(f64, f64), AicError> { + if let Some(util) = grid.and_then(|g| g.util(query)) { + if util > 0.0 { + return Ok((sol_query / (util * util_scale), util)); + } + } + Err(AicError::EmpiricalNotImplemented(format!( + "No empirical utilisation data to estimate this op at query={query:?}: \ + no own-shape, cross-shape, or sibling transfer reference available." + ))) +} + +/// Nearest reference index by categorical shape features in per-dim +/// normalised log space (mirrors Python `_nearest_candidate`; ties keep the +/// first candidate, matching `np.argmin`). Returns `None` for an empty list. +pub fn nearest_candidate_index(query_features: &[f64], candidates: &[Vec]) -> Option { + if candidates.is_empty() { + return None; + } + let dims = query_features.len(); + let feats: Vec> = candidates + .iter() + .map(|c| c.iter().map(|&v| log_floor(v)).collect()) + .collect(); + let mut mins = vec![f64::INFINITY; dims]; + let mut maxs = vec![f64::NEG_INFINITY; dims]; + for row in &feats { + for (a, &v) in row.iter().enumerate() { + mins[a] = mins[a].min(v); + maxs[a] = maxs[a].max(v); + } + } + let spans: Vec = mins + .iter() + .zip(&maxs) + .map(|(&lo, &hi)| if hi - lo > 0.0 { hi - lo } else { 1.0 }) + .collect(); + // NOTE: the query is intentionally NOT clamped here (unlike UtilGrid) — + // Python normalises the query into the candidates' span without clipping. + let q: Vec = query_features + .iter() + .enumerate() + .map(|(a, &v)| (log_floor(v) - mins[a]) / spans[a]) + .collect(); + let mut best = 0; + let mut best_dist2 = f64::INFINITY; + for (i, row) in feats.iter().enumerate() { + let dist2: f64 = row + .iter() + .enumerate() + .map(|(a, &v)| { + let n = (v - mins[a]) / spans[a]; + (n - q[a]) * (n - q[a]) + }) + .sum(); + if dist2 < best_dist2 { + best_dist2 = dist2; + best = i; + } + } + Some(best) +} + +/// Process-lifetime memo of built util grids, keyed by the caller's slice +/// identity. Rust perf tables are immutable after load and each +/// `PerfDatabase` owns its cache, so a plain keyed map replaces Python's +/// `id(node)`-qualified module cache. +#[derive(Debug, Default)] +pub struct UtilGridCache { + grids: Mutex>>>, +} + +impl UtilGridCache { + pub fn new() -> Self { + Self::default() + } + + /// Fetch or build the grid for `key`. `builder` returns `None` when the + /// slice has no usable calibration data (typed coverage miss) — that + /// outcome is memoised too, mirroring Python's `grid_for` returning + /// `None` (the caller then raises via [`estimate`] with `grid=None`). + pub fn get_or_build(&self, key: &str, builder: F) -> Option> + where + F: FnOnce() -> Option, + { + let mut grids = self.grids.lock().expect("util grid cache poisoned"); + if let Some(cached) = grids.get(key) { + return cached.clone(); + } + let built = builder().map(Arc::new); + grids.insert(key.to_string(), built.clone()); + built + } +} + +#[cfg(test)] +mod tests { + //! Mirrors the math anchors of `tests/unit/sdk/test_util_empirical.py`. + //! The Python cache/`grid_for` contract tests are id()-specific and are + //! covered by `UtilGridCache`'s own semantics instead. + + use super::*; + + fn approx(a: f64, b: f64) { + assert!((a - b).abs() < 1e-12, "expected {b}, got {a}"); + } + + #[test] + fn exact_singleton_duplicate_and_empty_grid_contracts() { + let exact = UtilGrid::new(vec![ + UtilSample::new(vec![16.0], 0.8), + UtilSample::new(vec![8.0], 0.2), + UtilSample::new(vec![9.0], 0.4), + ]); + let duplicate = UtilGrid::new(vec![ + UtilSample::new(vec![4.0], 0.6), + UtilSample::new(vec![4.0], 0.7), + ]); + + approx(exact.util(&[9.0]).unwrap(), 0.4); + approx( + UtilGrid::new(vec![UtilSample::new(vec![0.0], 0.3)]) + .util(&[100.0]) + .unwrap(), + 0.3, + ); + // Duplicate coordinates: first sample wins (stable ordering). + approx(duplicate.util(&[4.0]).unwrap(), 0.6); + assert!(UtilGrid::new(vec![]).util(&[1.0, 2.0]).is_none()); + } + + #[test] + fn one_dim_k2_idw_uses_nearest_samples_in_normalized_log_space() { + let grid = UtilGrid::new(vec![ + UtilSample::new(vec![16.0], 0.8), + UtilSample::new(vec![8.0], 0.2), + UtilSample::new(vec![9.0], 0.4), + ]); + let distance_9 = 11.0_f64.ln() - 9.0_f64.ln(); + let distance_8 = 11.0_f64.ln() - 8.0_f64.ln(); + let expected = (0.4 / distance_9 + 0.2 / distance_8) / (1.0 / distance_9 + 1.0 / distance_8); + + approx(grid.util(&[11.0]).unwrap(), expected); + } + + #[test] + fn multidimensional_k2_idw_uses_nearest_samples() { + let grid = UtilGrid::new(vec![ + UtilSample::new(vec![1.0, 1.0], 0.2), + UtilSample::new(vec![100.0, 1.0], 0.6), + UtilSample::new(vec![100.0, 100.0], 1.0), + ]); + + // (10, 1) is equidistant from the first two normalized-log samples. + approx(grid.util(&[10.0, 1.0]).unwrap(), 0.4); + } + + #[test] + fn one_dim_extrapolation_clamps_to_measured_bounds() { + let grid = UtilGrid::new(vec![ + UtilSample::new(vec![8.0], 0.2), + UtilSample::new(vec![16.0], 0.8), + ]); + + approx(grid.util(&[1.0]).unwrap(), 0.2); + approx(grid.util(&[128.0]).unwrap(), 0.8); + } + + #[test] + fn multidimensional_extrapolation_clamps_each_axis() { + let grid = UtilGrid::new(vec![ + UtilSample::new(vec![1.0, 1.0], 0.2), + UtilSample::new(vec![1.0, 10.0], 0.4), + UtilSample::new(vec![10.0, 1.0], 0.8), + ]); + + // Clamping (0.1, 100) produces the exact measured boundary (1, 10). + approx(grid.util(&[0.1, 100.0]).unwrap(), 0.4); + } + + #[test] + fn build_samples_filters_non_positive_latency_and_sol() { + let samples = build_samples( + vec![ + (vec![2.0], 4.0), // kept: util = sol/lat = 2/4 + (vec![3.0], 0.0), // dropped: latency <= 0 + (vec![4.0], -1.0), // dropped: latency < 0 + (vec![0.5], 1.0), // dropped: sol_fn returns 0 below 1.0 + ], + |coords| if coords[0] >= 1.0 { coords[0] } else { 0.0 }, + ); + assert_eq!(samples.len(), 1); + approx(samples[0].util, 0.5); + } + + #[test] + fn estimate_returns_latency_and_util_or_typed_miss() { + let grid = UtilGrid::new(vec![UtilSample::new(vec![1.0], 0.5)]); + let (latency, util) = estimate(1.0, &[1.0], Some(&grid), 1.0).unwrap(); + approx(latency, 2.0); + approx(util, 0.5); + + // Cross-op level alignment: latency = SOL / (util * k). + let (latency, _) = estimate(1.0, &[1.0], Some(&grid), 2.0).unwrap(); + approx(latency, 1.0); + + let missing = estimate(1.0, &[1.0], None, 1.0); + assert!(matches!(missing, Err(AicError::EmpiricalNotImplemented(_)))); + let empty = UtilGrid::new(vec![]); + let empty_res = estimate(1.0, &[1.0], Some(&empty), 1.0); + assert!(matches!(empty_res, Err(AicError::EmpiricalNotImplemented(_)))); + } + + #[test] + fn nearest_candidate_matches_python_normalised_log_selection() { + // Query (90,) between features (1,) and (100,): log-nearer to 100. + let idx = nearest_candidate_index(&[90.0], &[vec![1.0], vec![100.0]]); + assert_eq!(idx, Some(1)); + // Single candidate always selected; empty list yields None. + assert_eq!(nearest_candidate_index(&[5.0], &[vec![1.0]]), Some(0)); + assert_eq!(nearest_candidate_index(&[5.0], &[]), None); + } + + #[test] + fn util_grid_cache_memoises_including_misses() { + let cache = UtilGridCache::new(); + let mut builds = 0; + let grid = cache.get_or_build("k", || { + builds += 1; + Some(UtilGrid::new(vec![UtilSample::new(vec![1.0], 0.5)])) + }); + assert!(grid.is_some()); + let again = cache.get_or_build("k", || { + builds += 1; + None + }); + assert!(again.is_some()); + assert_eq!(builds, 1); + + let miss = cache.get_or_build("missing", || None); + assert!(miss.is_none()); + let miss_again = cache.get_or_build("missing", || { + panic!("memoised miss must not rebuild") + }); + assert!(miss_again.is_none()); + } +} diff --git a/rust/aiconfigurator-core/src/perf_database/mod.rs b/rust/aiconfigurator-core/src/perf_database/mod.rs index 2ea868ba0..ac7746f95 100644 --- a/rust/aiconfigurator-core/src/perf_database/mod.rs +++ b/rust/aiconfigurator-core/src/perf_database/mod.rs @@ -12,9 +12,11 @@ use std::path::{Path, PathBuf}; +use crate::common::enums::{DatabaseMode, TransferPolicy}; use crate::common::error::AicError; use crate::common::system_spec::SystemSpec; use crate::config::{PerfDbSources, PerfSource}; +use crate::operators::util_empirical::UtilGridCache; /// Resolve the ordered source list for one op-file basename: the Python-supplied /// shared-layer sources when present, else a single primary `data_root/` @@ -102,6 +104,19 @@ pub struct PerfDatabase { pub wideep_mla: WideEpMlaTable, pub wideep_moe: WideEpMoeTable, pub state_space: StateSpaceTable, + /// Query mode (Python's `database._default_database_mode`). SILICON + /// queries collected tables only; HYBRID falls back to the util-space + /// empirical layer on a typed silicon miss; EMPIRICAL always answers + /// `SOL/util`. Defaults to SILICON; `Engine::from_spec_bytes` overwrites + /// it from the spec. + pub database_mode: DatabaseMode, + /// Enabled empirical transfer kinds (Python's `database.transfer_policy`). + pub transfer_policy: TransferPolicy, + /// Memo of built util-calibration grids, keyed by op/slice identity + /// (see `operators::util_empirical`). Tables are immutable after load and + /// mode/policy are fixed per database instance, so the cache never needs + /// invalidation. + pub util_grids: UtilGridCache, } impl PerfDatabase { @@ -193,8 +208,20 @@ impl PerfDatabase { ), system_spec: spec, data_root, + database_mode: DatabaseMode::default(), + transfer_policy: TransferPolicy::ALL, + util_grids: UtilGridCache::new(), }) } + + /// Configure the query mode + transfer policy (both immutable per + /// database instance afterwards, mirroring Python's configured query + /// views). Called by `Engine::from_spec_bytes` with the spec's values. + pub fn with_mode(mut self, database_mode: DatabaseMode, transfer_policy: TransferPolicy) -> Self { + self.database_mode = database_mode; + self.transfer_policy = transfer_policy; + self + } } #[cfg(test)] diff --git a/rust/aiconfigurator-core/src/py.rs b/rust/aiconfigurator-core/src/py.rs index 082ffba35..264a8d010 100644 --- a/rust/aiconfigurator-core/src/py.rs +++ b/rust/aiconfigurator-core/src/py.rs @@ -766,6 +766,8 @@ mod tests { }, speculative: None, perf_db_sources: Default::default(), + database_mode: Default::default(), + transfer_policy: None, extra: BTreeMap::new(), } } diff --git a/rust/aiconfigurator-core/tests/memory_round_trip.rs b/rust/aiconfigurator-core/tests/memory_round_trip.rs index 62f6b28dd..d7c8dd3d7 100644 --- a/rust/aiconfigurator-core/tests/memory_round_trip.rs +++ b/rust/aiconfigurator-core/tests/memory_round_trip.rs @@ -105,6 +105,8 @@ fn request(tolerance_fraction: Option) -> KvCacheEstimateRequest { }, speculative: None, perf_db_sources: BTreeMap::new(), + database_mode: Default::default(), + transfer_policy: None, extra: BTreeMap::new(), }, max_num_tokens: 8192, diff --git a/src/aiconfigurator/sdk/engine.py b/src/aiconfigurator/sdk/engine.py index e4e6bbe0c..56dc7668f 100644 --- a/src/aiconfigurator/sdk/engine.py +++ b/src/aiconfigurator/sdk/engine.py @@ -693,6 +693,13 @@ def _engine_config_dict( # in Python so the Rust core inherits the same rows. Empty/absent = Rust # uses its single-``data_root`` default (back-compat with old specs). "perf_db_sources": _compute_perf_db_sources(database), + # Perf-database query mode + enabled empirical transfer kinds, read off + # the live database view so the compiled engine answers HYBRID/EMPIRICAL + # queries the same way the Python step does. Presets are resolved here + # (single source of truth in ``common.TRANSFER_PRESETS``); the wire form + # is always explicit kind tokens, ``None`` = the default ALL policy. + "database_mode": _database_mode_name(database), + "transfer_policy": _transfer_policy_tokens(database), "extra": {}, } # SpeculativeConfig (flattened, Option<>): emit nextn/nextn_accept_rates at @@ -708,6 +715,33 @@ def _opt_int(value: Any) -> int | None: return None if value is None else int(value) +def _database_mode_name(database: Any) -> str: + """The database view's query mode as the wire token (default SILICON).""" + if database is None: + return "SILICON" + mode = getattr(database, "get_default_database_mode", lambda: None)() + return getattr(mode, "name", str(mode)) if mode is not None else "SILICON" + + +def _transfer_policy_tokens(database: Any) -> list[str] | None: + """The view's enabled transfer kinds as explicit wire tokens. + + ``None`` = the default ALL-transfers policy (backward-compatible absent + key). A non-default policy serialises as a sorted list of kind values so + the Rust side never needs the preset vocabulary. + """ + if database is None: + return None + policy = getattr(database, "transfer_policy", None) + if policy is None: + return None + from aiconfigurator.sdk.common import ALL_TRANSFERS + + if frozenset(policy) == ALL_TRANSFERS: + return None + return sorted(kind.value for kind in policy) + + # --------------------------------------------------------------------------- # # Public entry points. # --------------------------------------------------------------------------- # diff --git a/src/aiconfigurator/sdk/rust_engine_step.py b/src/aiconfigurator/sdk/rust_engine_step.py index 23a6b02d6..fbad61f29 100644 --- a/src/aiconfigurator/sdk/rust_engine_step.py +++ b/src/aiconfigurator/sdk/rust_engine_step.py @@ -432,11 +432,28 @@ def _engine_config_json(model: Any, database: Any) -> str: "kv_block_size": None, "nextn": int(nextn) if nextn is not None else None, "nextn_accept_rates": ([float(r) for r in nextn_accept_rates] if nextn_accept_rates is not None else None), + # Mode + transfer policy are part of the engine identity: a HYBRID or + # EMPIRICAL view of the same model/system must not reuse a SILICON + # handle (the compiled engine bakes the mode into its query dispatch). + "database_mode": _database_mode_key(database), + "transfer_policy": _transfer_policy_key(database), "extra": {}, } return json.dumps(config, sort_keys=True, separators=(",", ":")) +def _database_mode_key(database: Any) -> str: + mode = getattr(database, "get_default_database_mode", lambda: None)() + return getattr(mode, "name", str(mode)) if mode is not None else "SILICON" + + +def _transfer_policy_key(database: Any) -> list[str] | None: + policy = getattr(database, "transfer_policy", None) + if policy is None: + return None + return sorted(getattr(kind, "value", str(kind)) for kind in policy) + + def _backend_name(value: Any) -> str: if hasattr(value, "value"): return str(value.value) @@ -480,7 +497,13 @@ def _moe_quant_to_dtype(value: Any) -> str | None: value_name = getattr(getattr(value, "value", None), "name", None) if value_name: name = value_name.lower() - if name in {"w4afp8", "w4a16_mxfp4", "w4a8_mxfp4_mxfp8"}: + if name in { + "w4afp8", + "w4a16_mxfp4", + "w4a8_mxfp4_mxfp8", + "w4a8_mxfp4_mxfp8_trtllm", + "w4a16_mxfp4_cutlass", + }: return name return _quant_to_dtype(value) From c11b8162bdb1906126e9068ecec583f93094bf87 Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 12:49:10 +0800 Subject: [PATCH 03/10] feat(rust-core): GEMM-family HYBRID/EMPIRICAL dispatch (reference port) Mode-aware dispatch for the three GEMM-family tables, mirroring the Python `_query_*_table` classmethods (`operations/gemm.py`): - query_gemm_table / query_scale_matrix_table: SILICON queries the table; HYBRID converts a typed silicon miss (is_missing_perf_data: PerfDatabase | Io, the exact set Op::Fallback catches) into the util-space empirical estimate; EMPIRICAL always estimates. - compute_scale uses the ported ZeroAwareDeltaLookup (zeroes are measured deltas; nearest-point first, frozen-util rescale) with its own DeltaLookupCache on PerfDatabase. - UtilGridCache::get_or_try_build now mirrors grid_for's full contract: Ok(None) coverage misses are memoised, schema/load errors propagate un-memoised. - perf_database gains algorithm-free point accessors (gemm_points / compute_scale_points / scale_matrix_points via perf_interp node_points); the dispatch lives in the operator layer per the crate's layering policy. Anchored against Python oracles on b200_sxm/vllm/0.19.0 at 1e-9: off-grid m on a collected site, fully off-site query, exact collected hit (util reconstruction returns the measured value), small-shape corner, and the HYBRID terminal EmpiricalNotImplemented miss on a quant with no table (int4_wo). Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- rust/aiconfigurator-core/src/common/error.rs | 13 ++ .../aiconfigurator-core/src/operators/gemm.rs | 199 ++++++++++++++++- .../src/operators/util_empirical.rs | 209 ++++++++++++++++-- .../src/perf_database/gemm.rs | 62 +++++- .../src/perf_database/mod.rs | 6 +- .../src/perf_database/perf_interp.rs | 12 + 6 files changed, 472 insertions(+), 29 deletions(-) diff --git a/rust/aiconfigurator-core/src/common/error.rs b/rust/aiconfigurator-core/src/common/error.rs index e5d635f22..0d68c7bda 100644 --- a/rust/aiconfigurator-core/src/common/error.rs +++ b/rust/aiconfigurator-core/src/common/error.rs @@ -60,3 +60,16 @@ pub enum AicError { }, } +impl AicError { + /// The "missing silicon data" class that HYBRID converts into an + /// empirical estimate. Mirrors Python's + /// `_MISSING_SILICON_DATA_EXCEPTIONS` (`PerfDataNotAvailableError`, + /// `InterpolationDataNotAvailableError`) — the same set `Op::Fallback` + /// catches. `EmpiricalNotImplemented` is deliberately NOT in this set: + /// it is the terminal miss raised after the empirical path itself found + /// no calibration data. + pub fn is_missing_perf_data(&self) -> bool { + matches!(self, Self::PerfDatabase(_) | Self::Io { .. }) + } +} + diff --git a/rust/aiconfigurator-core/src/operators/gemm.rs b/rust/aiconfigurator-core/src/operators/gemm.rs index 0e0811527..5f70ba6c8 100644 --- a/rust/aiconfigurator-core/src/operators/gemm.rs +++ b/rust/aiconfigurator-core/src/operators/gemm.rs @@ -14,9 +14,11 @@ //! to mirror Python's `max(0.0, latency)` behavior. use serde::{Deserialize, Serialize}; -use crate::common::enums::GemmQuantMode; +use crate::common::enums::{DatabaseMode, GemmQuantMode}; use crate::common::error::AicError; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid, ZeroAwareDeltaLookup}; +use crate::perf_database::gemm::{gemm_sol_latency_ms, normalize_fp8_static_quant}; use crate::perf_database::PerfDatabase; /// GEMM operation: a dense matrix multiply of shape `M=x, N=n, K=k`. @@ -75,17 +77,17 @@ impl GemmOp { let m = m.div_ceil(self.seq_split.max(1)); let quant = quant_override.unwrap_or(self.quant_mode); - let mut latency = db.gemm.query(quant, m, self.n, self.k)?; - let mut source = Source::Silicon; + let (mut latency, mut source) = query_gemm_table(db, quant, m, self.n, self.k)?; if quant == GemmQuantMode::Fp8Static { - let cs_latency = db.gemm.query_compute_scale(quant, m, self.k)?; + let (cs_latency, cs_source) = query_compute_scale_table(db, quant, m, self.k)?; latency -= cs_latency; - source = source.combine(Source::Silicon); // both silicon -> still silicon + source = source.combine(cs_source); if self.low_precision_input { - let sm_latency = db.gemm.query_scale_matrix(quant, m, self.k)?; + let (sm_latency, sm_source) = query_scale_matrix_table(db, quant, m, self.k)?; latency -= sm_latency; + source = source.combine(sm_source); } } @@ -104,6 +106,152 @@ impl GemmOp { } } +// --------------------------------------------------------------------------- +// Database-mode dispatch, mirroring the Python `_query_*_table` classmethods +// (`operations/gemm.py`): SILICON queries the table; HYBRID converts a typed +// silicon miss into the util-space empirical estimate; EMPIRICAL always +// estimates. The SOL diagnostic modes never reach the compiled engine (the +// routing gate delegates them to the Python step). +// --------------------------------------------------------------------------- + +/// GEMM latency for `(m, n, k, quant)` under the database's query mode. +fn query_gemm_table( + db: &PerfDatabase, + quant: GemmQuantMode, + m: u32, + n: u32, + k: u32, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok((gemm_empirical(db, quant, m, n, k)?, Source::Empirical)), + DatabaseMode::Hybrid => match db.gemm.query(quant, m, n, k) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => { + Ok((gemm_empirical(db, quant, m, n, k)?, Source::Empirical)) + } + Err(err) => Err(err), + }, + _ => Ok((db.gemm.query(quant, m, n, k)?, Source::Silicon)), + } +} + +/// `SOL(query)/util` over the quant's own collected `(m, n, k)` grid. +/// Mirrors Python `_query_gemm_table::get_empirical` (grid depth 3; the SOL +/// uses the ORIGINAL quant, numerically identical to the fp8_static->fp8 +/// table quant since the profiles match). +fn gemm_empirical( + db: &PerfDatabase, + quant: GemmQuantMode, + m: u32, + n: u32, + k: u32, +) -> Result { + let spec = &db.system_spec; + let sol = |c: &[f64]| gemm_sol_latency_ms(spec, quant, c[0], c[1], c[2]); + let key = format!("gemm:{}", normalize_fp8_static_quant(quant).name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.gemm.gemm_points(quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [m as f64, n as f64, k as f64]; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + +/// compute_scale delta for `(m, k)` under the database's query mode. +fn query_compute_scale_table( + db: &PerfDatabase, + quant: GemmQuantMode, + m: u32, + k: u32, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok((compute_scale_empirical(db, quant, m, k)?, Source::Empirical)), + DatabaseMode::Hybrid => match db.gemm.query_compute_scale(quant, m, k) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => { + Ok((compute_scale_empirical(db, quant, m, k)?, Source::Empirical)) + } + Err(err) => Err(err), + }, + _ => Ok((db.gemm.query_compute_scale(quant, m, k)?, Source::Silicon)), + } +} + +/// Zero-aware nearest-point delta estimate over the compute_scale grid. +/// Mirrors Python `_query_compute_scale_table::get_empirical`: a typed miss +/// on the slice itself becomes the terminal `EmpiricalNotImplemented`. +fn compute_scale_empirical( + db: &PerfDatabase, + quant: GemmQuantMode, + m: u32, + k: u32, +) -> Result { + let key = format!("compute_scale:{}", normalize_fp8_static_quant(quant).name()); + let lookup = db + .delta_lookups + .get_or_try_build(&key, || match db.gemm.compute_scale_points(quant) { + Ok(points) => Ok(ZeroAwareDeltaLookup::new(points)), + Err(err) if err.is_missing_perf_data() => Err(AicError::EmpiricalNotImplemented( + format!("No empirical compute_scale data is available for m={m}, k={k}."), + )), + Err(err) => Err(err), + })?; + // sol_mem = 2 m k / bw * 1000 (read + write of the activation). + let spec = &db.system_spec; + lookup.estimate(&[m as f64, k as f64], |c| 2.0 * c[0] * c[1] / spec.gpu.mem_bw * 1000.0) +} + +/// scale_matrix latency for `(m, k)` under the database's query mode. +fn query_scale_matrix_table( + db: &PerfDatabase, + quant: GemmQuantMode, + m: u32, + k: u32, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok((scale_matrix_empirical(db, quant, m, k)?, Source::Empirical)), + DatabaseMode::Hybrid => match db.gemm.query_scale_matrix(quant, m, k) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => { + Ok((scale_matrix_empirical(db, quant, m, k)?, Source::Empirical)) + } + Err(err) => Err(err), + }, + _ => Ok((db.gemm.query_scale_matrix(quant, m, k)?, Source::Silicon)), + } +} + +/// `SOL(query)/util` over the scale_matrix `(m, k)` grid (a real memory +/// kernel, unlike the compute_scale delta). Mirrors Python +/// `_query_scale_matrix_table::get_empirical` (grid depth 2, +/// `sol_mem = 3 m k / bw * 1000`). +fn scale_matrix_empirical( + db: &PerfDatabase, + quant: GemmQuantMode, + m: u32, + k: u32, +) -> Result { + let spec = &db.system_spec; + let sol = |c: &[f64]| 3.0 * c[0] * c[1] / spec.gpu.mem_bw * 1000.0; + let key = format!("scale_matrix:{}", normalize_fp8_static_quant(quant).name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.gemm.scale_matrix_points(quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [m as f64, k as f64]; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + #[cfg(test)] mod tests { use super::*; @@ -191,6 +339,45 @@ mod tests { ); } + /// Oracle values generated from the Python reference on the same data: + /// `GEMM._query_gemm_table(db, m, n, k, quant, database_mode=EMPIRICAL)` + /// on b200_sxm/vllm/0.19.0. Regenerate if the shipped GEMM table or the + /// util-empirical math changes. + #[test] + fn gemm_empirical_matches_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + let cases = [ + // off-grid m on a collected (n, k) site + (3000u32, 65536u32, 16384u32, GemmQuantMode::Bfloat16, 3.7278025700902204), + // fully off-site query + (777, 4000, 5000, GemmQuantMode::Bfloat16, 0.023767651577298037), + // exact collected hit: util reconstruction returns the measured value + (32768, 65536, 16384, GemmQuantMode::Nvfp4, 20.538665771484375), + // small-shape corner + (1, 129, 130, GemmQuantMode::Fp8, 0.004668885990466126), + ]; + for (m, n, k, quant, expected) in cases { + let (latency, source) = query_gemm_table(&db, quant, m, n, k).expect("empirical query"); + assert!( + (latency - expected).abs() < 1e-9, + "({m}, {n}, {k}, {quant:?}): expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Empirical); + } + } + + /// HYBRID on a quant with NO collected table (int4_wo on b200/vllm) must + /// surface the terminal EmpiricalNotImplemented miss, never a fabricated + /// value (mirrors the Python contract). + #[test] + fn gemm_hybrid_missing_quant_raises_empirical_not_implemented() { + let mut db = b200_vllm_db(); + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let result = query_gemm_table(&db, GemmQuantMode::Int4Wo, 64, 64, 64); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } + #[test] fn gemm_op_weights_bytes_matches_python_formula() { let op = GemmOp::new("w", 1024, 4096, GemmQuantMode::Bfloat16); diff --git a/rust/aiconfigurator-core/src/operators/util_empirical.rs b/rust/aiconfigurator-core/src/operators/util_empirical.rs index 4ec7bbb76..3b59c441b 100644 --- a/rust/aiconfigurator-core/src/operators/util_empirical.rs +++ b/rust/aiconfigurator-core/src/operators/util_empirical.rs @@ -283,21 +283,159 @@ impl UtilGridCache { Self::default() } - /// Fetch or build the grid for `key`. `builder` returns `None` when the - /// slice has no usable calibration data (typed coverage miss) — that - /// outcome is memoised too, mirroring Python's `grid_for` returning - /// `None` (the caller then raises via [`estimate`] with `grid=None`). - pub fn get_or_build(&self, key: &str, builder: F) -> Option> + /// Fetch or build the grid for `key`. + /// + /// `builder` mirrors Python's `grid_for` contract: `Ok(None)` when the + /// slice has no usable calibration data (a typed coverage miss — memoised, + /// the caller then raises via [`estimate`] with `grid=None`), `Err` for + /// programming/schema errors (propagated, NOT memoised, never converted + /// into a fallback). + pub fn get_or_try_build(&self, key: &str, builder: F) -> Result>, AicError> where - F: FnOnce() -> Option, + F: FnOnce() -> Result, AicError>, { let mut grids = self.grids.lock().expect("util grid cache poisoned"); if let Some(cached) = grids.get(key) { - return cached.clone(); + return Ok(cached.clone()); } - let built = builder().map(Arc::new); + let built = builder()?.map(Arc::new); grids.insert(key.to_string(), built.clone()); - built + Ok(built) + } +} + +/// Nearest-point lookup over a non-negative latency *delta* table (the +/// `compute_scale` mechanism; Python `gemm._ZeroAwareDeltaLookup`). +/// +/// `compute_scale` stores `max(dynamic_quant - static_quant, 0)`: zero is a +/// measured, meaningful delta, not a missing latency. A normal util grid +/// cannot represent it (`SOL / 0`), and dropping zeroes can make one positive +/// noise sample the reference for the whole table. Select the nearest point +/// on the complete 2-D grid first: a selected zero stays zero; a positive +/// point uses frozen utilization so extrapolation scales with the query's +/// amount of work. +#[derive(Debug)] +pub struct ZeroAwareDeltaLookup { + coords: Vec>, + latencies: Vec, + mins: Vec, + spans: Vec, + norm: Vec>, +} + +impl ZeroAwareDeltaLookup { + /// Keep every point with `latency >= 0` (zero INCLUDED, unlike + /// [`build_samples`]). + pub fn new(points: Vec<(Vec, f64)>) -> Self { + let kept: Vec<(Vec, f64)> = points.into_iter().filter(|(_, lat)| *lat >= 0.0).collect(); + if kept.is_empty() { + return Self { + coords: Vec::new(), + latencies: Vec::new(), + mins: Vec::new(), + spans: Vec::new(), + norm: Vec::new(), + }; + } + let dims = kept[0].0.len(); + let logc: Vec> = kept + .iter() + .map(|(c, _)| c.iter().map(|&v| log_floor(v)).collect()) + .collect(); + let mut mins = vec![f64::INFINITY; dims]; + let mut maxs = vec![f64::NEG_INFINITY; dims]; + for row in &logc { + for (a, &v) in row.iter().enumerate() { + mins[a] = mins[a].min(v); + maxs[a] = maxs[a].max(v); + } + } + let spans: Vec = mins + .iter() + .zip(&maxs) + .map(|(&lo, &hi)| if hi - lo > 0.0 { hi - lo } else { 1.0 }) + .collect(); + let norm: Vec> = logc + .iter() + .map(|row| { + row.iter() + .enumerate() + .map(|(a, &v)| (v - mins[a]) / spans[a]) + .collect() + }) + .collect(); + Self { + latencies: kept.iter().map(|(_, lat)| *lat).collect(), + coords: kept.into_iter().map(|(c, _)| c).collect(), + mins, + spans, + norm, + } + } + + /// Nearest-point delta estimate (query is NOT clamped; the frozen-util + /// rescale `query_sol / (ref_sol / ref_latency)` carries extrapolation). + pub fn estimate(&self, query: &[f64], sol_fn: F) -> Result + where + F: Fn(&[f64]) -> f64, + { + if self.latencies.is_empty() { + return Err(AicError::EmpiricalNotImplemented(format!( + "No empirical compute_scale delta data is available at query={query:?}." + ))); + } + let q: Vec = query + .iter() + .enumerate() + .map(|(a, &v)| (log_floor(v) - self.mins[a]) / self.spans[a]) + .collect(); + let mut best = 0; + let mut best_dist2 = f64::INFINITY; + for (i, row) in self.norm.iter().enumerate() { + let dist2: f64 = row.iter().zip(&q).map(|(&x, &y)| (x - y) * (x - y)).sum(); + if dist2 < best_dist2 { + best_dist2 = dist2; + best = i; + } + } + let reference_latency = self.latencies[best]; + if reference_latency == 0.0 { + return Ok(0.0); + } + let reference_sol = sol_fn(&self.coords[best]); + let query_sol = sol_fn(query); + if reference_sol <= 0.0 || query_sol <= 0.0 { + return Err(AicError::EmpiricalNotImplemented(format!( + "No positive SOL reference is available for compute_scale at query={query:?}." + ))); + } + Ok(query_sol / (reference_sol / reference_latency)) + } +} + +/// Memo of built [`ZeroAwareDeltaLookup`]s (same keying/lifetime rationale as +/// [`UtilGridCache`]). +#[derive(Debug, Default)] +pub struct DeltaLookupCache { + lookups: Mutex>>, +} + +impl DeltaLookupCache { + pub fn new() -> Self { + Self::default() + } + + pub fn get_or_try_build(&self, key: &str, builder: F) -> Result, AicError> + where + F: FnOnce() -> Result, + { + let mut lookups = self.lookups.lock().expect("delta lookup cache poisoned"); + if let Some(cached) = lookups.get(key) { + return Ok(cached.clone()); + } + let built = Arc::new(builder()?); + lookups.insert(key.to_string(), built.clone()); + Ok(built) } } @@ -430,26 +568,55 @@ mod tests { } #[test] - fn util_grid_cache_memoises_including_misses() { + fn util_grid_cache_memoises_hits_and_misses_but_not_errors() { let cache = UtilGridCache::new(); let mut builds = 0; - let grid = cache.get_or_build("k", || { + let grid = cache.get_or_try_build("k", || { builds += 1; - Some(UtilGrid::new(vec![UtilSample::new(vec![1.0], 0.5)])) + Ok(Some(UtilGrid::new(vec![UtilSample::new(vec![1.0], 0.5)]))) }); - assert!(grid.is_some()); - let again = cache.get_or_build("k", || { + assert!(grid.unwrap().is_some()); + let again = cache.get_or_try_build("k", || { builds += 1; - None + Ok(None) }); - assert!(again.is_some()); + assert!(again.unwrap().is_some()); assert_eq!(builds, 1); - let miss = cache.get_or_build("missing", || None); - assert!(miss.is_none()); - let miss_again = cache.get_or_build("missing", || { - panic!("memoised miss must not rebuild") + // Typed coverage miss (Ok(None)) is memoised, like Python's grid_for. + let miss = cache.get_or_try_build("missing", || Ok(None)); + assert!(miss.unwrap().is_none()); + let miss_again = + cache.get_or_try_build("missing", || panic!("memoised miss must not rebuild")); + assert!(miss_again.unwrap().is_none()); + + // Programming/schema errors propagate and are NOT memoised. + let err = cache.get_or_try_build("broken", || { + Err(AicError::PerfDatabase("schema".to_string())) }); - assert!(miss_again.is_none()); + assert!(err.is_err()); + let recovered = cache.get_or_try_build("broken", || Ok(None)); + assert!(recovered.unwrap().is_none()); + } + + #[test] + fn zero_aware_delta_keeps_zeroes_and_freezes_utilization() { + let lookup = ZeroAwareDeltaLookup::new(vec![ + (vec![16.0, 16.0], 0.0), + (vec![1024.0, 1024.0], 2.0), + ]); + let sol = |c: &[f64]| c[0] * c[1]; + + // Nearest to the zero point: the measured zero delta stays zero. + assert_eq!(lookup.estimate(&[8.0, 8.0], sol).unwrap(), 0.0); + // Nearest to the positive point: frozen util scales with query SOL. + let expected = (2048.0 * 2048.0) / ((1024.0 * 1024.0) / 2.0); + assert!((lookup.estimate(&[2048.0, 2048.0], sol).unwrap() - expected).abs() < 1e-12); + // Empty table -> typed empirical miss. + let empty = ZeroAwareDeltaLookup::new(vec![]); + assert!(matches!( + empty.estimate(&[1.0, 1.0], sol), + Err(AicError::EmpiricalNotImplemented(_)) + )); } } diff --git a/rust/aiconfigurator-core/src/perf_database/gemm.rs b/rust/aiconfigurator-core/src/perf_database/gemm.rs index 13fcc385c..ad06fe59f 100644 --- a/rust/aiconfigurator-core/src/perf_database/gemm.rs +++ b/rust/aiconfigurator-core/src/perf_database/gemm.rs @@ -170,6 +170,66 @@ impl GemmTable { query_scale_table(&grids.by_quant, lookup.name(), m, k, &sol, true, &self.data_root) } + /// Collected `(m, n, k) -> latency` points for the quant's table, for the + /// operator-layer util-calibration grid (Python's + /// `require_data_slice(_gemm_data, tqm)` + `iter_grid(..., depth=3)`). + /// Missing quant / empty table is a typed `PerfDatabase` miss. + pub fn gemm_points(&self, quant: GemmQuantMode) -> Result, f64)>, AicError> { + let grids = self.load_gemm()?; + let quant_name = normalize_fp8_static_quant(quant).name(); + let (node, _) = grids.by_quant.get(quant_name).ok_or_else(|| { + AicError::PerfDatabase(format!( + "GEMM perf data missing for quant '{quant_name}' at {}", + self.data_root.display() + )) + })?; + let points = crate::perf_database::perf_interp::node_points(node); + if points.is_empty() { + return Err(AicError::PerfDatabase(format!( + "GEMM perf data empty for quant '{quant_name}' at {}", + self.data_root.display() + ))); + } + Ok(points) + } + + /// Collected `(m, k) -> delta` points of the compute_scale table (zeroes + /// included — they are measured deltas). Typed miss when absent. + pub fn compute_scale_points(&self, quant: GemmQuantMode) -> Result, f64)>, AicError> { + let grids = self.load_compute_scale()?; + Self::two_d_points(grids, normalize_fp8_static_quant(quant), "compute_scale", &self.data_root) + } + + /// Collected `(m, k) -> latency` points of the scale_matrix table. + /// Typed miss when absent. + pub fn scale_matrix_points(&self, quant: GemmQuantMode) -> Result, f64)>, AicError> { + let grids = self.load_scale_matrix()?; + Self::two_d_points(grids, normalize_fp8_static_quant(quant), "scale_matrix", &self.data_root) + } + + fn two_d_points( + grids: &TwoDGrids, + quant: GemmQuantMode, + table_name: &str, + data_root: &Path, + ) -> Result, f64)>, AicError> { + let quant_name = quant.name(); + let node = grids.by_quant.get(quant_name).ok_or_else(|| { + AicError::PerfDatabase(format!( + "{table_name} perf data missing for quant '{quant_name}' at {}", + data_root.display() + )) + })?; + let points = crate::perf_database::perf_interp::node_points(node); + if points.is_empty() { + return Err(AicError::PerfDatabase(format!( + "{table_name} perf data empty for quant '{quant_name}' at {}", + data_root.display() + ))); + } + Ok(points) + } + fn load_gemm(&self) -> Result<&GemmEngineGrids, AicError> { let cell = self.gemm.get_or_init(|| { let mut grids = load_gemm_parquet(&self.gemm_sources)?; @@ -303,7 +363,7 @@ fn gemm_quant_by_name(name: &str) -> Option { /// the fp8 perf tables — the perf-DB never stores rows under /// `fp8_static`. Applied uniformly to the GEMM, compute_scale, and /// scale_matrix table queries. -fn normalize_fp8_static_quant(quant: GemmQuantMode) -> GemmQuantMode { +pub(crate) fn normalize_fp8_static_quant(quant: GemmQuantMode) -> GemmQuantMode { if quant == GemmQuantMode::Fp8Static { GemmQuantMode::Fp8 } else { diff --git a/rust/aiconfigurator-core/src/perf_database/mod.rs b/rust/aiconfigurator-core/src/perf_database/mod.rs index ac7746f95..2d710e048 100644 --- a/rust/aiconfigurator-core/src/perf_database/mod.rs +++ b/rust/aiconfigurator-core/src/perf_database/mod.rs @@ -16,7 +16,7 @@ use crate::common::enums::{DatabaseMode, TransferPolicy}; use crate::common::error::AicError; use crate::common::system_spec::SystemSpec; use crate::config::{PerfDbSources, PerfSource}; -use crate::operators::util_empirical::UtilGridCache; +use crate::operators::util_empirical::{DeltaLookupCache, UtilGridCache}; /// Resolve the ordered source list for one op-file basename: the Python-supplied /// shared-layer sources when present, else a single primary `data_root/` @@ -117,6 +117,9 @@ pub struct PerfDatabase { /// mode/policy are fixed per database instance, so the cache never needs /// invalidation. pub util_grids: UtilGridCache, + /// Memo of zero-aware delta lookups (the `compute_scale` empirical + /// mechanism); same keying/lifetime rationale as `util_grids`. + pub delta_lookups: DeltaLookupCache, } impl PerfDatabase { @@ -211,6 +214,7 @@ impl PerfDatabase { database_mode: DatabaseMode::default(), transfer_policy: TransferPolicy::ALL, util_grids: UtilGridCache::new(), + delta_lookups: DeltaLookupCache::new(), }) } diff --git a/rust/aiconfigurator-core/src/perf_database/perf_interp.rs b/rust/aiconfigurator-core/src/perf_database/perf_interp.rs index b88d487ca..9e53330ee 100644 --- a/rust/aiconfigurator-core/src/perf_database/perf_interp.rs +++ b/rust/aiconfigurator-core/src/perf_database/perf_interp.rs @@ -680,6 +680,18 @@ impl SiteIndex { } } +/// Flatten a nested table into `(coords, latency_ms)` points with `f64` +/// coordinates — the input shape of the operator-layer util-calibration +/// grids (`operators::util_empirical::build_samples`, mirroring Python's +/// `iter_grid` over a nested dict slice). +pub(crate) fn node_points(node: &Node) -> Vec<(Vec, f64)> { + let mut out = Vec::new(); + walk_leaves(node, &mut Vec::new(), &mut out); + out.into_iter() + .map(|(coords, latency)| (coords.into_iter().map(|c| c as f64).collect(), latency)) + .collect() +} + fn walk_leaves(node: &Node, prefix: &mut Vec, out: &mut Vec<(Vec, f64)>) { match node { Node::Leaf(v) => out.push((prefix.clone(), *v)), From 28582e7f5bafa723d23725a995f6d4885d74d4c3 Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 13:12:08 +0800 Subject: [PATCH 04/10] feat(rust-core): communication-family HYBRID/EMPIRICAL dispatch CustomAllReduce + NCCL mode dispatch mirroring operations/communication.py: per-rank size math and tp<=1 short-circuits stay ahead of the branch; empirical mirrors the eff = min(tp|num_gpus, collected) slice selection, the rank-overflow node-boundary borrow (policy-exempt compatibility exception, TODO #1260), per-collective SOL (2x all_reduce), the NCCL-if-loaded-else-oneCCL source selection, and the terminal EmpiricalNotImplemented conditions (no comm data / missing bucket). P2P and mem-ops stay analytic in all modes (verified, untouched). Anchored against Python oracles on b200_sxm/vllm/0.19.0 at 1e-9 (12 value anchors + policy/miss contracts). Two PRE-EXISTING silicon divergences observed and documented, not changed: the GB200 (gpus_per_node=72, tp>4) allreduce->NCCL redirect is absent in the Rust silicon path, and Rust query_nccl falls back to oneCCL per-slice where Python picks one source per query. Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- .../src/operators/communication.rs | 474 ++++++++++++++++-- .../src/perf_database/communication.rs | 90 ++++ 2 files changed, 524 insertions(+), 40 deletions(-) diff --git a/rust/aiconfigurator-core/src/operators/communication.rs b/rust/aiconfigurator-core/src/operators/communication.rs index 14032415a..ad2ba47fb 100644 --- a/rust/aiconfigurator-core/src/operators/communication.rs +++ b/rust/aiconfigurator-core/src/operators/communication.rs @@ -4,8 +4,9 @@ //! Communication operators: custom allreduce, NCCL collectives, P2P. //! //! Mirrors `aiconfigurator.sdk.operations.communication.{CustomAllReduce, -//! NCCL, P2P}` SILICON paths. This is where the topology-aware scaling -//! lives: +//! NCCL, P2P}`, including the database-mode dispatch of +//! `_query_custom_allreduce_table` / `_query_nccl_table`. This is where the +//! topology-aware scaling lives: //! //! - `CustomAllReduceOp`: caps `tp_size` to `num_gpus_per_node` before //! the table lookup, then scales by `(tp-1)/tp * (per_node)/(per_node-1) @@ -13,13 +14,15 @@ //! - `NcclOp`: caps `num_gpus` to the table's max recorded fan-out, then //! scales by `(num_gpus-1)/num_gpus * max/(max-1) * max_bw/req_bw`. //! - `P2POp`: pure analytic formula — `(bytes / inter_node_bw + -//! p2p_latency) * 1000`. No CSV. +//! p2p_latency) * 1000`. No CSV, no mode dispatch (analytic in every +//! mode, like Python's `_query_p2p_table`). use serde::{Deserialize, Serialize}; -use crate::common::enums::CommQuantMode; +use crate::common::enums::{CommQuantMode, DatabaseMode}; use crate::common::error::AicError; use crate::common::system_spec::SystemSpec; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid}; use crate::perf_database::PerfDatabase; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -63,30 +66,131 @@ impl CustomAllReduceOp { num_tokens: u32, ) -> Result { if self.tp_size <= 1 { - return Ok(PerformanceResult::zero()); + // No-op short-circuit: tp_size=1 has no allreduce. Python tags + // it "empirical" so EMPIRICAL/SOL breakdowns don't report a + // spurious silicon leakage. + return Ok(PerformanceResult::new(0.0, Source::Empirical)); } let per_rank_tokens = num_tokens.div_ceil(self.seq_split.max(1)); // CP: busiest rank let message_size = (per_rank_tokens as u64) * (self.hidden_size as u64); - let spec = &db.system_spec; - let per_node = spec.node.num_gpus_per_node; - let effective_tp = self.tp_size.min(per_node); - let mut latency = - db.communication - .query_custom_allreduce(self.quant, effective_tp, message_size)?; - if self.tp_size > per_node { - let base_bw = p2p_bandwidth(spec, per_node); - let target_bw = p2p_bandwidth(spec, self.tp_size); - let f_tp = self.tp_size as f64; - let f_pn = per_node as f64; - let scale = (f_tp - 1.0) / f_tp * f_pn / (f_pn - 1.0).max(1.0) * base_bw / target_bw; - latency *= scale; - } - Ok(PerformanceResult::new(latency, Source::Silicon) + let (latency, source) = + query_custom_allreduce_table(db, self.quant, self.tp_size, message_size)?; + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } } +// --------------------------------------------------------------------------- +// Database-mode dispatch, mirroring the Python `_query_*_table` classmethods +// (`operations/communication.py`): SILICON queries the table (+ topology +// scaling); HYBRID converts a typed silicon miss into the util-space +// empirical estimate; EMPIRICAL always estimates. The SOL diagnostic modes +// never reach the compiled engine (the routing gate delegates them to the +// Python step). +// --------------------------------------------------------------------------- + +/// Custom-allreduce latency for `(quant, tp_size, size)` under the +/// database's query mode. `size` is an element count (see +/// [`CustomAllReduceOp::query`]). +fn query_custom_allreduce_table( + db: &PerfDatabase, + quant: CommQuantMode, + tp_size: u32, + size: u64, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + custom_allreduce_empirical(db, quant, tp_size, size)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match custom_allreduce_silicon(db, quant, tp_size, size) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + custom_allreduce_empirical(db, quant, tp_size, size)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((custom_allreduce_silicon(db, quant, tp_size, size)?, Source::Silicon)), + } +} + +/// Python `_query_custom_allreduce_table.get_silicon` minus the GB200 +/// (`num_gpus_per_node == 72 && tp > 4`) redirect to NCCL, which the Rust +/// silicon path has never carried (pre-existing divergence, reported, not +/// silently changed here). +fn custom_allreduce_silicon( + db: &PerfDatabase, + quant: CommQuantMode, + tp_size: u32, + size: u64, +) -> Result { + let spec = &db.system_spec; + let per_node = spec.node.num_gpus_per_node; + let effective_tp = tp_size.min(per_node); + let mut latency = db + .communication + .query_custom_allreduce(quant, effective_tp, size)?; + if tp_size > per_node { + let base_bw = p2p_bandwidth(spec, per_node); + let target_bw = p2p_bandwidth(spec, tp_size); + let f_tp = tp_size as f64; + let f_pn = per_node as f64; + let scale = (f_tp - 1.0) / f_tp * f_pn / (f_pn - 1.0).max(1.0) * base_bw / target_bw; + latency *= scale; + } + Ok(latency) +} + +/// Ring-allreduce SOL in ms. Mirrors Python +/// `_query_custom_allreduce_table.get_sol`: assume ring allreduce, ignore +/// constant latency, assume bfloat16 (`size` elements x 2 bytes). +fn custom_allreduce_sol_ms(spec: &SystemSpec, tp_size: u32, size: f64) -> f64 { + if tp_size == 1 { + return 0.0; + } + let p2p_bw = p2p_bandwidth(spec, tp_size); + let tp = tp_size as f64; + 2.0 * size * 2.0 / tp * (tp - 1.0) / p2p_bw * 1000.0 +} + +/// `SOL(query)/util` over the collected custom-allreduce size curve. +/// Mirrors Python `_query_custom_allreduce_table.get_empirical`: SOL uses +/// the real `tp_size`; the util grid is built from the effective +/// (node-capped) tp slice, so the SOL ratio carries any multi-node +/// bandwidth scaling. Rank-count overflow borrows the node-boundary util +/// slice regardless of the transfer policy — Python's documented +/// compatibility exception (`xshape` provenance, TODO #1260). +fn custom_allreduce_empirical( + db: &PerfDatabase, + quant: CommQuantMode, + tp_size: u32, + size: u64, +) -> Result { + let spec = &db.system_spec; + let sol_q = custom_allreduce_sol_ms(spec, tp_size, size as f64); + if tp_size <= 1 || sol_q <= 0.0 { + // No communication for a single rank -> 0/SOL, not a data gap. + return Ok(sol_q); + } + let eff = tp_size.min(spec.node.num_gpus_per_node); + let sol = |c: &[f64]| custom_allreduce_sol_ms(spec, eff, c[0]); + let key = format!("custom_allreduce:{}:{eff}", quant.name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.communication.custom_allreduce_points(quant, eff) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [size as f64]; + let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct NcclOp { pub name: String, @@ -131,35 +235,150 @@ impl NcclOp { num_tokens: u32, ) -> Result { if self.num_gpus <= 1 { - return Ok(PerformanceResult::zero()); + // No communication for a single rank. Python has no op-level + // short-circuit, but every mode branch returns 0 tagged + // "empirical" for num_gpus <= 1 (silicon's `num_gpus == 1` + // early return / empirical's `sol_q = 0` guard). + return Ok(PerformanceResult::new(0.0, Source::Empirical)); } let per_rank_tokens = num_tokens.div_ceil(self.seq_split.max(1)); // CP: busiest rank // Python: message_size = ceil(x/seq_split) * num_elements_per_token (float). let message_size = ((per_rank_tokens as f64) * self.hidden_size) as u64; - let max_recorded = db - .communication - .nccl_max_num_gpus(self.dtype, &self.operation)? - .unwrap_or(self.num_gpus); - let effective = self.num_gpus.min(max_recorded); - let mut latency = db - .communication - .query_nccl(self.dtype, &self.operation, effective, message_size)?; - if self.num_gpus > max_recorded { - let spec = &db.system_spec; - let max_bw = p2p_bandwidth(spec, max_recorded); - let req_bw = p2p_bandwidth(spec, self.num_gpus); - let f_n = self.num_gpus as f64; - let f_m = max_recorded as f64; - let scale = - (f_n - 1.0) / f_n * f_m / (f_m - 1.0).max(1.0) * max_bw / req_bw; - latency *= scale; - } - Ok(PerformanceResult::new(latency, Source::Silicon) + let (latency, source) = + query_nccl_table(db, self.dtype, self.num_gpus, &self.operation, message_size)?; + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } } +/// NCCL collective latency for `(dtype, num_gpus, operation, +/// message_size)` under the database's query mode. +fn query_nccl_table( + db: &PerfDatabase, + dtype: CommQuantMode, + num_gpus: u32, + operation: &str, + message_size: u64, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + nccl_empirical(db, dtype, num_gpus, operation, message_size)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match nccl_silicon(db, dtype, num_gpus, operation, message_size) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + nccl_empirical(db, dtype, num_gpus, operation, message_size)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((nccl_silicon(db, dtype, num_gpus, operation, message_size)?, Source::Silicon)), + } +} + +/// Python `_query_nccl_table.get_silicon`: cap `num_gpus` to the max +/// recorded fan-out, RAW lerp on the size curve, then bandwidth-scale for +/// out-of-range fan-outs. +fn nccl_silicon( + db: &PerfDatabase, + dtype: CommQuantMode, + num_gpus: u32, + operation: &str, + message_size: u64, +) -> Result { + let max_recorded = db + .communication + .nccl_max_num_gpus(dtype, operation)? + .unwrap_or(num_gpus); + let effective = num_gpus.min(max_recorded); + let mut latency = db + .communication + .query_nccl(dtype, operation, effective, message_size)?; + if num_gpus > max_recorded { + let spec = &db.system_spec; + let max_bw = p2p_bandwidth(spec, max_recorded); + let req_bw = p2p_bandwidth(spec, num_gpus); + let f_n = num_gpus as f64; + let f_m = max_recorded as f64; + let scale = (f_n - 1.0) / f_n * f_m / (f_m - 1.0).max(1.0) * max_bw / req_bw; + latency *= scale; + } + Ok(latency) +} + +/// Per-collective SOL in ms. Mirrors Python `_query_nccl_table.get_sol`: +/// one-directional ring traffic for gather/scatter-style collectives, 2x +/// for all_reduce, and 0 for unknown collectives (which the empirical +/// path then returns unchanged, not as a data gap). +fn nccl_sol_ms( + spec: &SystemSpec, + dtype: CommQuantMode, + num_gpus: u32, + operation: &str, + message_size: f64, +) -> f64 { + let p2p_bw = p2p_bandwidth(spec, num_gpus); + let n = num_gpus as f64; + let mem = dtype.mapping().memory; + match operation { + "all_gather" | "alltoall" | "reduce_scatter" => { + mem * message_size * (n - 1.0) / n / p2p_bw * 1000.0 + } + "all_reduce" => 2.0 * mem * message_size * (n - 1.0) / n / p2p_bw * 1000.0, + _ => 0.0, + } +} + +/// `SOL(query)/util` over the collected NCCL size curve for this +/// `(dtype, operation, num_gpus)` slice. Mirrors Python +/// `_query_nccl_table.get_empirical`: the grid is built from the available +/// (capped) num_gpus slice; SOL uses the real `num_gpus` so the SOL ratio +/// carries scaling beyond the largest collected fan-out (`xshape` +/// borrowing regardless of transfer policy, TODO #1260). A source with no +/// NCCL/OneCCL data loaded at all, or no `(dtype, operation)` bucket, is +/// the terminal empirical miss. +fn nccl_empirical( + db: &PerfDatabase, + dtype: CommQuantMode, + num_gpus: u32, + operation: &str, + message_size: u64, +) -> Result { + let spec = &db.system_spec; + let sol_q = nccl_sol_ms(spec, dtype, num_gpus, operation, message_size as f64); + if num_gpus <= 1 || sol_q <= 0.0 { + // No communication for a single rank -> 0, not a data gap. + return Ok(sol_q); + } + let max_collected = match db.communication.nccl_empirical_max_num_gpus(dtype, operation) { + Ok(max) => max, + Err(err) if err.is_missing_perf_data() => { + return Err(AicError::EmpiricalNotImplemented(format!( + "No NCCL data for operation '{operation}' ({}, num_gpus={num_gpus}).", + dtype.name() + ))); + } + Err(err) => return Err(err), + }; + let eff = num_gpus.min(max_collected); + let sol = |c: &[f64]| nccl_sol_ms(spec, dtype, eff, operation, c[0]); + let key = format!("nccl:{}:{operation}:{eff}", dtype.name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.communication.nccl_empirical_points(dtype, operation, eff) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + // Typed coverage miss (e.g. an uncollected intermediate gpu + // count) -> no grid; estimate() raises the empirical miss. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [message_size as f64]; + let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + /// Pure analytic P2P latency — no CSV. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct P2POp { @@ -202,3 +421,178 @@ impl P2POp { fn p2p_bandwidth(spec: &SystemSpec, num_gpus: u32) -> f64 { spec.get_p2p_bandwidth(num_gpus) } + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + const REPO_ROOT_HINT: &str = env!("CARGO_MANIFEST_DIR"); + + fn b200_vllm_db() -> PerfDatabase { + let systems_root = PathBuf::from(REPO_ROOT_HINT) + .join("../..") + .join("src/aiconfigurator/systems"); + PerfDatabase::load(&systems_root, "b200_sxm", "vllm", "0.19.0").expect("db must load") + } + + /// Oracle values generated from the Python reference on the same data: + /// `CustomAllReduce._query_custom_allreduce_table(db, half, tp, size, + /// database_mode=EMPIRICAL)` on b200_sxm/vllm/0.19.0 (collected tp + /// slices: {2, 4, 8}; sizes 128..512Mi elements). Regenerate if the + /// shipped table or the util-empirical math changes. + #[test] + fn custom_allreduce_empirical_matches_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = DatabaseMode::Empirical; + let cases = [ + // off-grid size on a collected tp slice + (8u32, 300_000u64, 0.0210877421663319), + // exact collected hit: util reconstruction returns the measured value + (8, 1024, 0.007716159820556641), + // rank overflow: tp=16 on an 8-GPU node borrows the tp=8 boundary + // util; SOL(tp=16, inter-node bw) carries the multi-node scaling + (16, 524_288, 0.1607595384120941), + (2, 777_777, 0.019098769751423886), + (4, 65_536, 0.008213120102882384), + ]; + for (tp, size, expected) in cases { + let (latency, source) = + query_custom_allreduce_table(&db, CommQuantMode::Half, tp, size) + .expect("empirical query"); + assert!( + (latency - expected).abs() < 1e-9, + "(tp={tp}, size={size}): expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Empirical); + } + } + + /// An uncollected tp slice (tp=5: eff=5 < 8, no data) is the terminal + /// EmpiricalNotImplemented miss in both EMPIRICAL and HYBRID modes, + /// never a fabricated value (mirrors the Python contract). + #[test] + fn custom_allreduce_missing_tp_raises_empirical_not_implemented() { + for mode in [DatabaseMode::Empirical, DatabaseMode::Hybrid] { + let mut db = b200_vllm_db(); + db.database_mode = mode; + let result = query_custom_allreduce_table(&db, CommQuantMode::Half, 5, 1024); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "{mode:?}: got {result:?}" + ); + } + } + + /// HYBRID prefers silicon whenever the table covers the slice. Oracles: + /// `CustomAllReduce._query_custom_allreduce_table(..., database_mode=HYBRID)` + /// (source reported as `silicon` by Python for both). + #[test] + fn custom_allreduce_hybrid_prefers_silicon() { + let mut db = b200_vllm_db(); + db.database_mode = DatabaseMode::Hybrid; + let cases = [ + (8u32, 300_000u64, 0.024691462737973777), + // rank overflow still resolves on silicon (tp=8 slice + scale) + (16, 524_288, 0.16075953841209412), + ]; + for (tp, size, expected) in cases { + let (latency, source) = + query_custom_allreduce_table(&db, CommQuantMode::Half, tp, size) + .expect("hybrid query"); + assert!( + (latency - expected).abs() < 1e-9, + "(tp={tp}, size={size}): expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Silicon); + } + } + + /// Oracle values generated from the Python reference: + /// `NCCL._query_nccl_table(db, half, num_gpus, op, msg, + /// database_mode=EMPIRICAL)` on b200_sxm/vllm/0.19.0 (nccl 2.27.3 + /// data; collected num_gpus {2, 4, 8}; sizes 256..256Mi). + #[test] + fn nccl_empirical_matches_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = DatabaseMode::Empirical; + let cases = [ + // off-grid message sizes on collected fan-outs + (8u32, "all_reduce", 300_000u64, 0.027067167868039803), + (4, "all_gather", 10_000_000, 0.05871496201240777), + // gpu-count overflow: 32 > max collected 8 borrows the boundary + // util; SOL(32, inter-node bw) carries the scaling + (32, "all_reduce", 1_048_576, 0.31676464285714284), + // below-min size: boundary util-hold + (8, "reduce_scatter", 999, 0.015899137756552793), + // exact collected hit + (2, "alltoall", 4096, 0.009470000000000001), + ]; + for (num_gpus, op, msg, expected) in cases { + let (latency, source) = query_nccl_table(&db, CommQuantMode::Half, num_gpus, op, msg) + .expect("empirical query"); + assert!( + (latency - expected).abs() < 1e-9, + "({num_gpus}, {op}, {msg}): expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Empirical); + } + } + + /// An uncollected intermediate gpu count (num_gpus=3: eff=3, no slice) + /// is the terminal EmpiricalNotImplemented miss in EMPIRICAL and + /// HYBRID modes. + #[test] + fn nccl_missing_gpu_count_raises_empirical_not_implemented() { + for mode in [DatabaseMode::Empirical, DatabaseMode::Hybrid] { + let mut db = b200_vllm_db(); + db.database_mode = mode; + let result = query_nccl_table(&db, CommQuantMode::Half, 3, "all_reduce", 1024); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "{mode:?}: got {result:?}" + ); + } + } + + /// Unknown collectives SOL to 0 and return it unchanged in EMPIRICAL + /// mode (Python `get_empirical`'s `sol_q <= 0` early return — a no-op, + /// not a data gap). Same for a single rank. + #[test] + fn nccl_empirical_zero_sol_returns_zero_not_a_gap() { + let mut db = b200_vllm_db(); + db.database_mode = DatabaseMode::Empirical; + let (latency, source) = + query_nccl_table(&db, CommQuantMode::Half, 8, "broadcast", 1024).expect("no-op query"); + assert_eq!(latency, 0.0); + assert_eq!(source, Source::Empirical); + let (latency, _) = + query_nccl_table(&db, CommQuantMode::Half, 1, "all_reduce", 4096).expect("single rank"); + assert_eq!(latency, 0.0); + // tp=1 allreduce likewise returns 0 in EMPIRICAL mode. + let (latency, _) = query_custom_allreduce_table(&db, CommQuantMode::Half, 1, 4096) + .expect("single rank allreduce"); + assert_eq!(latency, 0.0); + } + + /// HYBRID prefers silicon whenever the table covers the slice. Oracles: + /// `NCCL._query_nccl_table(..., database_mode=HYBRID)`. + #[test] + fn nccl_hybrid_prefers_silicon() { + let mut db = b200_vllm_db(); + db.database_mode = DatabaseMode::Hybrid; + let cases = [ + (8u32, "all_reduce", 300_000u64, 0.02807729736328125), + (32, "all_reduce", 1_048_576, 0.3167646428571429), + ]; + for (num_gpus, op, msg, expected) in cases { + let (latency, source) = query_nccl_table(&db, CommQuantMode::Half, num_gpus, op, msg) + .expect("hybrid query"); + assert!( + (latency - expected).abs() < 1e-9, + "({num_gpus}, {op}, {msg}): expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Silicon); + } + } +} diff --git a/rust/aiconfigurator-core/src/perf_database/communication.rs b/rust/aiconfigurator-core/src/perf_database/communication.rs index 33ee7e471..faf53050a 100644 --- a/rust/aiconfigurator-core/src/perf_database/communication.rs +++ b/rust/aiconfigurator-core/src/perf_database/communication.rs @@ -167,6 +167,96 @@ impl CommunicationTable { interp_message_size(by_size, message_size) } + /// Collected `(message_size,) -> latency_ms` points of the + /// custom-allreduce curve for `(quant, tp_size)` — the input of the + /// operator-layer util grid (mirrors Python's + /// `require_data_slice(dw, quant_mode, eff, "AUTO")`). Typed miss when + /// the slice is absent or empty. + pub fn custom_allreduce_points( + &self, + quant: CommQuantMode, + tp_size: u32, + ) -> Result, f64)>, AicError> { + let grids = self.load_custom_allreduce()?; + let key = (quant.name().to_string(), tp_size); + let by_size = grids.by_keys.get(&key).ok_or_else(|| { + AicError::PerfDatabase(format!( + "custom_allreduce data missing for {key:?} at {}", + self.data_root.display() + )) + })?; + if by_size.is_empty() { + return Err(AicError::PerfDatabase(format!( + "custom_allreduce data empty for {key:?} at {}", + self.data_root.display() + ))); + } + Ok(by_size.iter().map(|(&size, &lat)| (vec![size as f64], lat)).collect()) + } + + /// The single NCCL source the empirical path calibrates from, with + /// Python's selection order (`NCCL._query_nccl_table.get_empirical`): + /// the NCCL table when loaded, else the OneCCL fallback; a typed miss + /// when neither is loaded. Unlike [`Self::query_nccl`], there is NO + /// per-slice fallback across sources. + fn nccl_empirical_source(&self) -> Result<&NcclGrids, AicError> { + if let Ok(grids) = self.load_nccl() { + return Ok(grids); + } + self.load_oneccl() + } + + /// Maximum collected `num_gpus` for `(dtype, operation)` in the NCCL + /// empirical source (single source, Python parity — unlike + /// [`Self::nccl_max_num_gpus`], which unions NCCL and OneCCL for the + /// silicon cap). Typed miss when the source has no such bucket. + pub fn nccl_empirical_max_num_gpus( + &self, + dtype: CommQuantMode, + operation: &str, + ) -> Result { + let grids = self.nccl_empirical_source()?; + let dtype_name = dtype.name(); + grids + .by_keys + .keys() + .filter(|(d, op, _)| d.as_str() == dtype_name && op.as_str() == operation) + .map(|(_, _, n)| *n) + .max() + .ok_or_else(|| { + AicError::PerfDatabase(format!( + "NCCL data missing for dtype='{dtype_name}', operation='{operation}' at {}", + self.data_root.display() + )) + }) + } + + /// Collected `(message_size,) -> latency_ms` points for + /// `(dtype, operation, num_gpus)` in the NCCL empirical source. Typed + /// miss when the slice is absent or empty. + pub fn nccl_empirical_points( + &self, + dtype: CommQuantMode, + operation: &str, + num_gpus: u32, + ) -> Result, f64)>, AicError> { + let grids = self.nccl_empirical_source()?; + let key = (dtype.name().to_string(), operation.to_string(), num_gpus); + let by_size = grids.by_keys.get(&key).ok_or_else(|| { + AicError::PerfDatabase(format!( + "NCCL data missing for {key:?} at {}", + self.data_root.display() + )) + })?; + if by_size.is_empty() { + return Err(AicError::PerfDatabase(format!( + "NCCL data empty for {key:?} at {}", + self.data_root.display() + ))); + } + Ok(by_size.iter().map(|(&size, &lat)| (vec![size as f64], lat)).collect()) + } + /// Maximum recorded `num_gpus` for an NCCL (dtype, operation) tuple. /// Operator layer uses this to decide whether to apply a bandwidth /// scale factor for out-of-range fan-outs. From ccf64897527b0f7536172e9b0d0ac6233f22f82b Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 13:18:30 +0800 Subject: [PATCH 05/10] feat(rust-core): HYBRID/EMPIRICAL dispatch for attention/MoE/MLA/DSA/DSV4 + silicon-view FallbackOp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-family mode dispatch mirroring the Python _query_*_table classmethods, all following the GEMM reference pattern (EMPIRICAL -> estimate; HYBRID -> silicon, typed miss -> estimate; EmpiricalNotImplemented never caught), each anchored against Python-generated oracles at 1e-9: - attention: ctx/gen/encoder own-slice grids, the window->0 util-carrier fallback, and the XSHAPE cross-head_size ladder (ctx util_scale = prefill hs-ratio quotient, gen = 1.0), policy-gated. - moe: the full transfer ladder (xshape -> xquant same-profile -> xprofile with util-level rescale), kernel-table (ll/std) selection folded into cache keys, policy fingerprints in reference-grid keys. - mla: six tables (ctx/gen, bmm with the bfloat16 slice fallback, ctx/gen module, wideep ctx/gen with the round-ties-even sample head mapping). - dsa: ctx/gen module dispatch with the full calibration-variant ladder (interp-prefix / p0-anchor / boundary-freeze / legacy, exact-head splits) keyed by the Python cache tags. - dsv4 + mhc: ctx (interp-prefix vs p0-anchor) / gen (kv-derived SOL dtype) attention and the mhc pre/post/both composition. MegaMoE verified to have no Python empirical path (nothing to port). - mamba/GDN verified mode-independent in Python (SOL-degradation contract, already mirrored) — no change needed. PerfDatabase is now a mode-configured view over Arc-shared PerfTables (Deref keeps table access unchanged); silicon_view() gives Op::Fallback the Python semantics: under HYBRID the primary op is evaluated silicon-only so a missing module table falls to the granular fallback chain instead of being hybrid-estimated at module level. Pre-existing divergences observed and documented, not changed: GB200 allreduce->NCCL redirect absent in Rust silicon; NCCL per-slice oneCCL fallback vs Python single-source; dsv4 generation silicon SOL still uses the opspec fmha (Python derives from kv, ~2e-5 rel); wideep MLA opspec num_heads slot carries tp_size (engine.py) skewing tp>1 silicon. cargo test --lib: 242 passed, 5 failed (all pre-existing: stale parse_b200_sxm constant + 4 py::tests needing a linked interpreter). Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- .../src/operators/attention.rs | 652 +++++++++++++- rust/aiconfigurator-core/src/operators/dsa.rs | 813 +++++++++++++++++- .../aiconfigurator-core/src/operators/dsv4.rs | 440 +++++++++- rust/aiconfigurator-core/src/operators/mhc.rs | 144 +++- rust/aiconfigurator-core/src/operators/mla.rs | 646 +++++++++++++- rust/aiconfigurator-core/src/operators/moe.rs | 747 ++++++++++++++-- .../src/operators/moe_dispatch.rs | 4 +- rust/aiconfigurator-core/src/operators/op.rs | 17 +- .../src/operators/wideep_mla.rs | 333 ++++++- .../src/perf_database/attention.rs | 196 ++++- .../src/perf_database/dsa.rs | 55 +- .../src/perf_database/dsv4.rs | 92 +- .../src/perf_database/mhc.rs | 37 + .../src/perf_database/mla.rs | 175 +++- .../src/perf_database/mod.rs | 80 +- .../src/perf_database/moe.rs | 127 +++ .../src/perf_database/wideep_mla.rs | 122 ++- 17 files changed, 4459 insertions(+), 221 deletions(-) diff --git a/rust/aiconfigurator-core/src/operators/attention.rs b/rust/aiconfigurator-core/src/operators/attention.rs index 6840c0d76..4d22b0cdd 100644 --- a/rust/aiconfigurator-core/src/operators/attention.rs +++ b/rust/aiconfigurator-core/src/operators/attention.rs @@ -15,12 +15,20 @@ //! - `seq_imbalance_correction_scale` / `gen_seq_imbalance_correction_scale` //! multiplier for unbalanced sequence distributions //! - `scale_factor` scaling at the end +//! - database-mode dispatch (SILICON / HYBRID / EMPIRICAL) with the +//! util-space empirical layer: exact-window then window=0 util carriers, +//! plus the cross-head_size (XSHAPE) transfer ladder use serde::{Deserialize, Serialize}; -use crate::common::enums::{FmhaQuantMode, KvCacheQuantMode}; +use crate::common::enums::{DatabaseMode, FmhaQuantMode, KvCacheQuantMode, TransferKind}; use crate::common::error::AicError; use crate::common::system_spec::SystemSpec; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid}; +use crate::perf_database::attention::{ + context_attention_sol_ms, context_attention_sol_with_prefix_ms, encoder_attention_sol_ms, + generation_attention_sol_ms, +}; use crate::perf_database::PerfDatabase; /// Analytic memory-op latency (ms). Matches Python's @@ -42,6 +50,67 @@ fn prefix_correction(full_s: u32, prefix: u32) -> f64 { (f * f - p * p) / (f * f) } +// --------------------------------------------------------------------------- +// Cross-head_size (XSHAPE) transfer support, mirroring the module-level +// helpers of `operations/attention.py`. +// --------------------------------------------------------------------------- + +/// Prefill-attention util vs head_size, relative to head_size=128 (Python's +/// `_ATTN_PREFILL_HS_RATIO`). Used to rescale a borrowed util curve when the +/// exact head_size has no collected data. DECODE util is ~head_size- +/// independent (memory-bound KV read), so decode transfer uses no table. +const ATTN_PREFILL_HS_RATIO_TRTLLM: &[(u32, f64)] = + &[(64, 0.58), (128, 1.00), (192, 1.10), (256, 1.17), (512, 1.20)]; +const ATTN_PREFILL_HS_RATIO_SGLANG: &[(u32, f64)] = + &[(64, 0.60), (128, 1.00), (192, 1.18), (256, 1.32), (512, 1.38)]; +const ATTN_PREFILL_HS_RATIO_VLLM: &[(u32, f64)] = + &[(64, 0.60), (128, 1.00), (192, 1.27), (256, 1.51), (512, 1.60)]; + +/// Prefill-attention util ratio vs head_size=128, log2-interpolated between +/// table points and clamped at the ends. Unknown backend -> 1.0 (no +/// correction). Mirrors Python `_attn_prefill_hs_ratio`. +fn attn_prefill_hs_ratio(backend: &str, head_size: u32) -> f64 { + let table = match backend { + "trtllm" => ATTN_PREFILL_HS_RATIO_TRTLLM, + "sglang" => ATTN_PREFILL_HS_RATIO_SGLANG, + "vllm" => ATTN_PREFILL_HS_RATIO_VLLM, + _ => return 1.0, + }; + if let Some(&(_, ratio)) = table.iter().find(|&&(h, _)| h == head_size) { + return ratio; + } + let (first, last) = (table[0], table[table.len() - 1]); + if head_size <= first.0 { + return first.1; + } + if head_size >= last.0 { + return last.1; + } + // Bracketing keys exist by the checks above (table is sorted ascending). + let (lo, lo_ratio) = *table.iter().rev().find(|&&(h, _)| h < head_size).expect("lower bracket"); + let (hi, hi_ratio) = *table.iter().find(|&&(h, _)| h > head_size).expect("upper bracket"); + let t = ((head_size as f64).log2() - (lo as f64).log2()) / ((hi as f64).log2() - (lo as f64).log2()); + lo_ratio + t * (hi_ratio - lo_ratio) +} + +/// Pick the reference head_size to transfer util from. Prefer 128 — the +/// canonical, most densely collected head_size and the ratio table's +/// reference point; otherwise the nearest collected head_size in log space +/// (1-D normalised-log argmin == `|log2 h − log2 target|` argmin; ties keep +/// the first candidate). Mirrors Python `_ref_head_size`. +fn ref_head_size(available: &[u32], target: u32) -> Option { + let avail: Vec = available.iter().copied().filter(|&h| h != 0).collect(); + if avail.is_empty() { + return None; + } + if avail.contains(&128) { + return Some(128); + } + let features: Vec> = avail.iter().map(|&h| vec![h as f64]).collect(); + let idx = util_empirical::nearest_candidate_index(&[target as f64], &features)?; + Some(avail[idx]) +} + // --------------------------------------------------------------------------- // Context attention // --------------------------------------------------------------------------- @@ -97,31 +166,34 @@ impl ContextAttentionOp { prefix: u32, seq_imbalance_correction_scale: f64, ) -> Result { - // Mirror Python's `ContextAttention._ctx(s, pfx)`: query the table at - // the full sequence `s + pfx` and apply the prefix correction for - // `pfx`. The Rust table is keyed by the full sequence length. - let ctx = |s: u32, pfx: u32| -> Result { - let full_s = s + pfx; - let table = db.attention.query_context( + // Mirror Python's `ContextAttention._ctx(s, pfx)`: each chunk + // dispatches through the database mode — the silicon table at the + // full sequence `s + pfx` with the prefix correction, or the + // util-space empirical estimate (whose SOL already discounts prefix). + let ctx = |s: u32, pfx: u32| -> Result<(f64, Source), AicError> { + query_context_attention_table( + db, batch_size, - full_s, + s, + pfx, self.n, self.n_kv, self.head_size, self.window_size, self.kv_cache_dtype, self.fmha_quant_mode, - )?; - Ok(table * prefix_correction(full_s, pfx)) + ) }; // Context parallelism (SGLang AllGather / zigzag): model rank 0's two // balanced chunks. c = ceil(isl / 2cp); rank 0 owns chunk 0 (prefix // unchanged) and chunk 2cp-1 (attends almost the full sequence). Only // the FMHA table term is split; the fused extras below are added once. - let mut latency = if self.cp_size > 1 { + let (mut latency, source) = if self.cp_size > 1 { let c = isl.div_ceil(2 * self.cp_size).max(1); - ctx(c, prefix)? + ctx(c, prefix + isl - c)? + let (first, first_source) = ctx(c, prefix)?; + let (second, second_source) = ctx(c, prefix + isl - c)?; + (first + second, first_source.combine(second_source)) } else { ctx(isl, prefix)? }; @@ -149,7 +221,7 @@ impl ContextAttentionOp { latency *= seq_imbalance_correction_scale; } - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -196,7 +268,8 @@ impl GenerationAttentionOp { kv_seq_tokens: u32, gen_seq_imbalance_correction_scale: f64, ) -> Result { - let table_latency = db.attention.query_generation( + let (table_latency, source) = query_generation_attention_table( + db, batch_size, kv_seq_tokens, self.n, @@ -209,7 +282,7 @@ impl GenerationAttentionOp { if gen_seq_imbalance_correction_scale != 1.0 { latency *= gen_seq_imbalance_correction_scale; } - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -250,19 +323,419 @@ impl EncoderAttentionOp { batch_size: u32, s: u32, ) -> Result { - let latency = db.attention.query_encoder( + let (latency, source) = query_encoder_attention_table( + db, batch_size, s, self.n, self.head_size, self.fmha_quant_mode, )?; - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } } +// --------------------------------------------------------------------------- +// Database-mode dispatch, mirroring the Python `_query_*_attention_table` +// classmethods (`operations/attention.py`): SILICON queries the table; HYBRID +// converts a typed silicon miss into the util-space empirical estimate; +// EMPIRICAL always estimates. The SOL diagnostic modes never reach the +// compiled engine (the routing gate delegates them to the Python step). +// --------------------------------------------------------------------------- + +/// Context attention latency for `(b, s, prefix, shape)` under the database's +/// query mode. The silicon path queries the table at `full_s = s + prefix` +/// and applies the prefix correction; the empirical path bakes the prefix +/// into the query SOL instead (mirroring Python's `get_sol`). +#[allow(clippy::too_many_arguments)] +fn query_context_attention_table( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + n: u32, + n_kv: u32, + head_size: u32, + window_size: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, +) -> Result<(f64, Source), AicError> { + let silicon = || -> Result { + let full_s = s + prefix; + let latency = db.attention.query_context( + b, full_s, n, n_kv, head_size, window_size, kv_quant, fmha_quant, + )?; + Ok(latency * prefix_correction(full_s, prefix)) + }; + match db.database_mode { + DatabaseMode::Empirical => Ok(( + context_attention_empirical( + db, b, s, prefix, n, n_kv, head_size, window_size, kv_quant, fmha_quant, + )?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + context_attention_empirical( + db, b, s, prefix, n, n_kv, head_size, window_size, kv_quant, fmha_quant, + )?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((silicon()?, Source::Silicon)), + } +} + +/// `SOL(query)/util` for context (prefill) attention. Mirrors Python +/// `_query_context_attention_table::get_empirical`: the query SOL always uses +/// the real window/prefix; the UTIL carrier is borrowed by slice — exact +/// window first, then window=0 (full attention), each trying the own +/// head_size grid and then the cross-head_size (XSHAPE) ladder before moving +/// to the next window. No basis at all surfaces the typed empirical miss. +#[allow(clippy::too_many_arguments)] +fn context_attention_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + n: u32, + n_kv: u32, + head_size: u32, + window_size: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, +) -> Result { + let spec = &db.system_spec; + let sol_time = context_attention_sol_with_prefix_ms( + spec, + b as f64, + s as f64, + prefix as f64, + n as f64, + n_kv as f64, + head_size, + window_size, + kv_quant, + fmha_quant, + ); + let n_kv_lookup = if n == n_kv { 0 } else { n_kv }; + let query = [n as f64, (s + prefix) as f64, b as f64]; + + let windows: Vec = if window_size > 0 { vec![window_size, 0] } else { vec![window_size] }; + for &slice_window in &windows { + // Own-slice grid: samples are full attention (prefix=0), so the + // per-sample SOL is the prefix=0 specialization at the slice's own + // head_size/window (c = [n, full_s, b]). + let key = format!( + "ctx_attn:{}:{}:{}:{}:{}", + fmha_quant.name(), + kv_quant.name(), + n_kv_lookup, + head_size, + slice_window + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.attention.context_points(fmha_quant, kv_quant, n_kv_lookup, head_size, slice_window) { + Ok(points) => { + let sol = |c: &[f64]| { + context_attention_sol_ms( + spec, n_kv_lookup, head_size, slice_window, kv_quant, fmha_quant, c[0], + c[1], c[2], + ) + }; + Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))) + } + // Typed coverage miss -> no grid (fall through the ladder); + // schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + if grid.as_deref().is_some_and(|g| !g.is_empty()) { + let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + return Ok(latency); + } + // Cross-head_size transfer (XSHAPE): this head_size has no data, but + // num_heads is already an in-grid axis, so only head_size differs. + // Borrow the nearest collected head_size's util and rescale by the + // prefill head_size-util ratio (SOL still uses the query's own + // head_size). + if db.transfer_policy.contains(TransferKind::XShape) { + if let Some((ref_grid, ref_hs)) = + ctx_headsize_ref_grid(db, fmha_quant, kv_quant, n_kv_lookup, head_size, slice_window)? + { + let scale = attn_prefill_hs_ratio(&db.backend, head_size) + / attn_prefill_hs_ratio(&db.backend, ref_hs); + let (latency, _) = + util_empirical::estimate(sol_time, &query, Some(&ref_grid), scale)?; + return Ok(latency); + } + } + } + + // No own-window, full-attention, or cross-head basis -> typed miss. + util_empirical::estimate(sol_time, &query, None, 1.0).map(|(latency, _)| latency) +} + +/// Reference util grid for context attention borrowed from the nearest +/// collected head_size (same fmha/kv/n_kv/window). Mirrors Python +/// `_ctx_headsize_ref_grid`: built with the REFERENCE slice's own SOL +/// (reference head_size in the formula). `Ok(None)` when nothing usable is +/// collected. +fn ctx_headsize_ref_grid( + db: &PerfDatabase, + fmha_quant: FmhaQuantMode, + kv_quant: KvCacheQuantMode, + n_kv_lookup: u32, + target_hs: u32, + window_size: u32, +) -> Result, u32)>, AicError> { + let head_sizes = match db.attention.context_head_sizes(fmha_quant, kv_quant, n_kv_lookup) { + Ok(sizes) => sizes, + Err(err) if err.is_missing_perf_data() => return Ok(None), + Err(err) => return Err(err), + }; + let Some(ref_hs) = ref_head_size(&head_sizes, target_hs) else { + return Ok(None); + }; + let spec = &db.system_spec; + // Reference identity (ref_hs) + provenance in the key, so a policy that + // later reuses the same slice as own-shape cannot alias this grid. + let key = format!( + "ctx_attn_xhs:{}:{}:{}:{}:{}:xshape", + fmha_quant.name(), + kv_quant.name(), + n_kv_lookup, + ref_hs, + window_size + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.attention.context_points(fmha_quant, kv_quant, n_kv_lookup, ref_hs, window_size) { + Ok(points) => { + let sol = |c: &[f64]| { + context_attention_sol_ms( + spec, n_kv_lookup, ref_hs, window_size, kv_quant, fmha_quant, c[0], c[1], + c[2], + ) + }; + let mut grid = UtilGrid::new(util_empirical::build_samples(points, sol)); + grid.reference_provenance = Some("xshape"); + Ok(Some(grid)) + } + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + Ok(grid.filter(|g| !g.is_empty()).map(|g| (g, ref_hs))) +} + +/// Generation attention latency for `(b, s, shape)` under the database's +/// query mode. +#[allow(clippy::too_many_arguments)] +fn query_generation_attention_table( + db: &PerfDatabase, + b: u32, + s: u32, + n: u32, + n_kv: u32, + head_size: u32, + window_size: u32, + kv_quant: KvCacheQuantMode, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + generation_attention_empirical(db, b, s, n, n_kv, head_size, window_size, kv_quant)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => { + match db.attention.query_generation(b, s, n, n_kv, head_size, window_size, kv_quant) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + generation_attention_empirical( + db, b, s, n, n_kv, head_size, window_size, kv_quant, + )?, + Source::Empirical, + )), + Err(err) => Err(err), + } + } + _ => Ok(( + db.attention.query_generation(b, s, n, n_kv, head_size, window_size, kv_quant)?, + Source::Silicon, + )), + } +} + +/// `SOL(query)/util` for generation (decode) attention. Mirrors Python +/// `_query_generation_attention_table::get_empirical`: the query SOL uses the +/// real window (`kv_len` capped at `w`); the UTIL carrier is borrowed by +/// slice — exact window then window=0 — calibrated from the RAW (SOL-clamped) +/// generation table. Decode util is ~head_size-independent, so the XSHAPE +/// transfer keeps `util_scale = 1.0`. +#[allow(clippy::too_many_arguments)] +fn generation_attention_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + n: u32, + n_kv: u32, + head_size: u32, + window_size: u32, + kv_quant: KvCacheQuantMode, +) -> Result { + let spec = &db.system_spec; + let n_kv_lookup = if n_kv == n { 0 } else { n_kv }; + let sol_time = generation_attention_sol_ms( + spec, n_kv_lookup, head_size, window_size, kv_quant, n as f64, b as f64, s as f64, + ); + let query = [n as f64, b as f64, s as f64]; + + let windows: Vec = if window_size > 0 { vec![window_size, 0] } else { vec![window_size] }; + for &slice_window in &windows { + let key = format!( + "gen_attn:{}:{}:{}:{}", + kv_quant.name(), + n_kv_lookup, + head_size, + slice_window + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.attention.generation_points(kv_quant, n_kv_lookup, head_size, slice_window) { + Ok(points) => { + let sol = |c: &[f64]| { + generation_attention_sol_ms( + spec, n_kv_lookup, head_size, slice_window, kv_quant, c[0], c[1], c[2], + ) + }; + Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))) + } + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + if grid.as_deref().is_some_and(|g| !g.is_empty()) { + let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + return Ok(latency); + } + if db.transfer_policy.contains(TransferKind::XShape) { + if let Some((ref_grid, _ref_hs)) = + gen_headsize_ref_grid(db, kv_quant, n_kv_lookup, head_size, slice_window)? + { + let (latency, _) = + util_empirical::estimate(sol_time, &query, Some(&ref_grid), 1.0)?; + return Ok(latency); + } + } + } + + util_empirical::estimate(sol_time, &query, None, 1.0).map(|(latency, _)| latency) +} + +/// Reference util grid for generation attention borrowed from the nearest +/// collected head_size (same kv/n_kv/window). Mirrors Python +/// `_gen_headsize_ref_grid` (reference head_size in the sample SOL). +fn gen_headsize_ref_grid( + db: &PerfDatabase, + kv_quant: KvCacheQuantMode, + n_kv_lookup: u32, + target_hs: u32, + window_size: u32, +) -> Result, u32)>, AicError> { + let head_sizes = match db.attention.generation_head_sizes(kv_quant, n_kv_lookup) { + Ok(sizes) => sizes, + Err(err) if err.is_missing_perf_data() => return Ok(None), + Err(err) => return Err(err), + }; + let Some(ref_hs) = ref_head_size(&head_sizes, target_hs) else { + return Ok(None); + }; + let spec = &db.system_spec; + let key = format!( + "gen_attn_xhs:{}:{}:{}:{}:xshape", + kv_quant.name(), + n_kv_lookup, + ref_hs, + window_size + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.attention.generation_points(kv_quant, n_kv_lookup, ref_hs, window_size) { + Ok(points) => { + let sol = |c: &[f64]| { + generation_attention_sol_ms( + spec, n_kv_lookup, ref_hs, window_size, kv_quant, c[0], c[1], c[2], + ) + }; + let mut grid = UtilGrid::new(util_empirical::build_samples(points, sol)); + grid.reference_provenance = Some("xshape"); + Ok(Some(grid)) + } + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + Ok(grid.filter(|g| !g.is_empty()).map(|g| (g, ref_hs))) +} + +/// Encoder attention latency for `(b, s, shape)` under the database's query +/// mode. +fn query_encoder_attention_table( + db: &PerfDatabase, + b: u32, + s: u32, + n: u32, + head_size: u32, + fmha_quant: FmhaQuantMode, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + encoder_attention_empirical(db, b, s, n, head_size, fmha_quant)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match db.attention.query_encoder(b, s, n, head_size, fmha_quant) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + encoder_attention_empirical(db, b, s, n, head_size, fmha_quant)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok(( + db.attention.query_encoder(b, s, n, head_size, fmha_quant)?, + Source::Silicon, + )), + } +} + +/// `SOL(query)/util` over the encoder slice's own `(n, s, b)` grid. Mirrors +/// Python `_query_encoder_attention_table::get_empirical`: own-shape only, no +/// window ladder, no transfer. +fn encoder_attention_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + n: u32, + head_size: u32, + fmha_quant: FmhaQuantMode, +) -> Result { + let spec = &db.system_spec; + let sol = |c: &[f64]| encoder_attention_sol_ms(spec, head_size, fmha_quant, c[0], c[1], c[2]); + let query = [n as f64, s as f64, b as f64]; + let key = format!("encoder_attn:{}:{}", fmha_quant.name(), head_size); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.attention.encoder_points(fmha_quant, head_size) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + #[cfg(test)] mod tests { use super::*; @@ -343,10 +816,149 @@ mod tests { #[test] fn mem_op_latency_uses_empirical_formula() { let db = b200_vllm_db(); - // b200_sxm.yaml: mem_bw=8e12, scaling=0.8, constant=3e-6. - // For 1MB: latency = (1e6 / (8e12 * 0.8) + 3e-6) * 1000 = 0.156... + 0.003 = 0.159ms. + // b200_sxm.yaml: mem_bw=7.7e12 (per-datasheet fix, PR #1246), + // scaling=0.8, constant=3e-6. + // For 1MB: latency = (1e6 / (7.7e12 * 0.8) + 3e-6) * 1000 ≈ 0.165ms. let latency = mem_op_latency_ms(&db.system_spec, 1_000_000.0); - let expected = (1_000_000.0_f64 / (8e12 * 0.8) + 3e-6) * 1000.0; + let expected = (1_000_000.0_f64 / (7.7e12 * 0.8) + 3e-6) * 1000.0; assert!((latency - expected).abs() < 1e-12); } + + /// Oracle values generated from the Python reference on the same data: + /// `ContextAttention._query_context_attention_table(db, b, s, prefix, n, + /// n_kv, kv, fmha, database_mode=EMPIRICAL, window_size=w, head_size=hs)` + /// on b200_sxm/vllm/0.19.0. Regenerate if the shipped attention tables or + /// the util-empirical math changes. + #[test] + fn context_attention_empirical_matches_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + // (b, s, prefix, n, n_kv, hs, w, kv, expected) + let cases: &[(u32, u32, u32, u32, u32, u32, u32, KvCacheQuantMode, f64)] = &[ + // own-shape off-grid query on the collected hs=128 slice + (7, 3000, 0, 64, 1, 128, 0, KvCacheQuantMode::Fp8, 0.771381792089557), + // exact collected hit: util reconstruction returns the measured value + (8, 16384, 0, 64, 1, 128, 0, KvCacheQuantMode::Fp8, 19.820667266845703), + // prefix baked into the query SOL (util from the full-seq point) + (4, 8192, 8192, 64, 1, 128, 0, KvCacheQuantMode::Fp8, 7.964372158050536), + // head_size=192 XSHAPE transfer (collected head sizes are {128, 256}; + // ref=128, util_scale = ratio(vllm,192)/ratio(vllm,128) = 1.27) + (4, 4096, 0, 48, 8, 192, 0, KvCacheQuantMode::Fp8, 0.7588535312592514), + // collected windowed slice (bfloat16 kv, w=8192) as its own carrier + (2, 10000, 0, 32, 1, 128, 8192, KvCacheQuantMode::Bfloat16, 6.254832211751053), + // uncollected window (w=4096) -> window=0 slice as the util carrier + (2, 10000, 0, 32, 1, 128, 4096, KvCacheQuantMode::Bfloat16, 1.0547865593548398), + ]; + for &(b, s, prefix, n, n_kv, hs, w, kv, expected) in cases { + let (latency, source) = query_context_attention_table( + &db, b, s, prefix, n, n_kv, hs, w, kv, FmhaQuantMode::Bfloat16, + ) + .expect("empirical query"); + assert!( + (latency - expected).abs() < 1e-9, + "(b={b}, s={s}, p={prefix}, n={n}, n_kv={n_kv}, hs={hs}, w={w}): \ + expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Empirical); + } + } + + /// Oracle values generated from the Python reference: + /// `GenerationAttention._query_generation_attention_table(db, b, s, n, + /// n_kv, kv, database_mode=EMPIRICAL, window_size=w, head_size=hs)` on + /// b200_sxm/vllm/0.19.0. + #[test] + fn generation_attention_empirical_matches_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + // (b, s, n, n_kv, hs, w, kv, expected) + let cases: &[(u32, u32, u32, u32, u32, u32, KvCacheQuantMode, f64)] = &[ + // own-shape off-grid query on the collected hs=128 slice + (48, 7777, 64, 8, 128, 0, KvCacheQuantMode::Fp8, 0.1302149492334821), + // exact collected hit (isl=1 + step=1 -> stored s=2), calibrated + // from the RAW (SOL-clamped) table -- NOT the 5-sample silicon avg + (32, 2, 64, 4, 128, 0, KvCacheQuantMode::Fp8, 0.008661333471536636), + // head_size=192 XSHAPE transfer (decode util_scale stays 1.0) + (16, 4096, 48, 8, 192, 0, KvCacheQuantMode::Fp8, 0.03992800042033196), + // collected windowed slice (bfloat16 kv, w=8192) as its own carrier + (8, 12000, 32, 1, 128, 8192, KvCacheQuantMode::Bfloat16, 0.07096281754412269), + // uncollected window (w=2048) -> window=0 slice as the util carrier + (8, 12000, 32, 1, 128, 2048, KvCacheQuantMode::Bfloat16, 0.0023706380832401778), + ]; + for &(b, s, n, n_kv, hs, w, kv, expected) in cases { + let (latency, source) = + query_generation_attention_table(&db, b, s, n, n_kv, hs, w, kv) + .expect("empirical query"); + assert!( + (latency - expected).abs() < 1e-9, + "(b={b}, s={s}, n={n}, n_kv={n_kv}, hs={hs}, w={w}): \ + expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Empirical); + } + } + + /// Oracle values generated from the Python reference: + /// `EncoderAttention._query_encoder_attention_table(db, 3, 900, 16, 64, + /// bfloat16, database_mode=...)` on b200_sxm/vllm/0.19.0. EMPIRICAL + /// estimates from the util grid; HYBRID resolves on silicon (the slice is + /// collected) and must NOT detour through the empirical layer. + #[test] + fn encoder_attention_empirical_and_hybrid_match_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + let (latency, source) = + query_encoder_attention_table(&db, 3, 900, 16, 64, FmhaQuantMode::Bfloat16) + .expect("empirical query"); + assert!((latency - 0.03625488888618745).abs() < 1e-9, "got {latency}"); + assert_eq!(source, Source::Empirical); + + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let (latency, source) = + query_encoder_attention_table(&db, 3, 900, 16, 64, FmhaQuantMode::Bfloat16) + .expect("hybrid query"); + assert!((latency - 0.038151752523614205).abs() < 1e-9, "got {latency}"); + assert_eq!(source, Source::Silicon); + } + + /// HYBRID: an uncollected head_size (192) misses silicon and falls back + /// to the XSHAPE empirical estimate (same value as EMPIRICAL mode), while + /// a collected slice keeps resolving on silicon. Oracle from Python + /// `_query_context_attention_table(..., database_mode=HYBRID)`. + #[test] + fn context_attention_hybrid_dispatch_matches_python() { + let mut db = b200_vllm_db(); + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let (latency, source) = query_context_attention_table( + &db, 4, 4096, 0, 48, 8, 192, 0, KvCacheQuantMode::Fp8, FmhaQuantMode::Bfloat16, + ) + .expect("hybrid query"); + assert!((latency - 0.7588535312592514).abs() < 1e-9, "got {latency}"); + assert_eq!(source, Source::Empirical); + + // Collected slice: silicon exact hit, untouched by the fallback. + let (latency, source) = query_context_attention_table( + &db, 8, 16384, 0, 64, 1, 128, 0, KvCacheQuantMode::Fp8, FmhaQuantMode::Bfloat16, + ) + .expect("hybrid query"); + assert!((latency - 19.820667266845703).abs() < 1e-9, "got {latency}"); + assert_eq!(source, Source::Silicon); + } + + /// With XSHAPE disabled and no own-slice data (head_size=192), the + /// estimate must surface the terminal EmpiricalNotImplemented miss — + /// verified against Python `db.set_transfer_policy("off")` raising + /// `EmpiricalNotImplementedError` for both ctx and gen. + #[test] + fn attention_xshape_disabled_raises_empirical_not_implemented() { + let mut db = b200_vllm_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + db.transfer_policy = crate::common::enums::TransferPolicy::OFF; + let ctx = query_context_attention_table( + &db, 4, 4096, 0, 48, 8, 192, 0, KvCacheQuantMode::Fp8, FmhaQuantMode::Bfloat16, + ); + assert!(matches!(ctx, Err(AicError::EmpiricalNotImplemented(_))), "got {ctx:?}"); + let gen = query_generation_attention_table(&db, 16, 4096, 48, 8, 192, 0, KvCacheQuantMode::Fp8); + assert!(matches!(gen, Err(AicError::EmpiricalNotImplemented(_))), "got {gen:?}"); + } } diff --git a/rust/aiconfigurator-core/src/operators/dsa.rs b/rust/aiconfigurator-core/src/operators/dsa.rs index a50f47794..78220deb5 100644 --- a/rust/aiconfigurator-core/src/operators/dsa.rs +++ b/rust/aiconfigurator-core/src/operators/dsa.rs @@ -11,13 +11,17 @@ //! `index_topk` is the top-k boundary (per-architecture; 2048 for both //! DeepSeek-V3.2 and GLM-5). It is plumbed from the Python op-spec emitter. +use std::collections::{BTreeMap, BTreeSet}; + use serde::{Deserialize, Serialize}; -use crate::common::enums::{FmhaQuantMode, GemmQuantMode, KvCacheQuantMode}; +use crate::common::enums::{DatabaseMode, FmhaQuantMode, GemmQuantMode, KvCacheQuantMode}; use crate::common::error::AicError; use crate::operators::base::{PerformanceResult, Source}; use crate::operators::communication::NcclOp; +use crate::operators::util_empirical::{self, UtilGrid}; use crate::perf_database::dsa::{ - bs_slice, dsa_dims, dsa_sparse_file_prefix, lookup_2d, DsaSparseTables, + bs_slice, dsa_context_sol_ms, dsa_dims, dsa_generation_sol_ms, dsa_sparse_file_prefix, + lookup_2d, DsaHeadGrid, DsaKey, DsaSparseTables, }; use crate::perf_database::PerfDatabase; @@ -90,20 +94,8 @@ impl DsaModuleOp { // correction (it had no Python counterpart and under-counted context // latency ~75%). `dsa_backend="trtllm"` mirrors Python's non-CP // default (`_query_context_dsa_module_table(dsa_backend="trtllm")`). - let latency = db.dsa.query_context( - &db.system_spec, - batch_size, - isl, - self.num_heads, - self.kv_cache_dtype, - self.fmha_quant_mode, - self.gemm_quant_mode, - &self.architecture, - prefix, - self.index_topk, - "trtllm", - )?; - Ok(PerformanceResult::new(latency, Source::Silicon) + let (latency, source) = query_context_table(db, self, batch_size, isl, prefix, "trtllm")?; + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -128,21 +120,12 @@ impl DsaModuleOp { ) -> Result { let sparse = db.dsa.load_cp_sparse(&self.architecture, self.num_heads)?; let mut base = |per_card: u32| { - db.dsa.query_context( - &db.system_spec, - batch_size, - per_card, - self.num_heads, - self.kv_cache_dtype, - self.fmha_quant_mode, - self.gemm_quant_mode, - &self.architecture, - prefix, - self.index_topk, - // Python `_query_cp` queries the CP base on the flashmla_kv - // slice (the kernel used under CP). - "flashmla_kv", - ) + // Python `_query_cp` queries the CP base through the full + // `query_context_dsa_module` dispatch (no explicit database_mode + // => the database default), on the flashmla_kv slice (the kernel + // used under CP); `float(...)` drops the source. + query_context_table(db, self, batch_size, per_card, prefix, "flashmla_kv") + .map(|(latency, _)| latency) }; let mut ag = |elems: u64| { NcclOp::new( @@ -255,23 +238,582 @@ impl DsaModuleOp { ) -> Result { // `dsa_backend="trtllm"` mirrors Python's generation default // (`_query_generation_dsa_module_table(dsa_backend="trtllm")`). - let latency = db.dsa.query_generation( + let (latency, source) = query_generation_table(db, self, batch_size, s, "trtllm")?; + Ok(PerformanceResult::new(latency, source) + .clamp_non_negative() + .scaled(self.scale_factor)) + } +} + +// --------------------------------------------------------------------------- +// Database-mode dispatch, mirroring the Python `_query_*_dsa_module_table` +// classmethods (`operations/dsa.py`). Python does NOT use the shared +// `_query_silicon_or_hybrid` helper here: it wraps the silicon lookup in an +// explicit try/except over `(PerfDataNotAvailableError, +// InterpolationDataNotAvailableError)` and only falls to `get_empirical` when +// `database_mode == HYBRID` — in Rust that catch set is exactly +// `err.is_missing_perf_data()` (which deliberately excludes +// `EmpiricalNotImplemented`). EMPIRICAL always estimates; the SOL diagnostic +// modes never reach the compiled engine. +// --------------------------------------------------------------------------- + +/// Context DSA module latency for the op's slice under the database's mode. +fn query_context_table( + db: &PerfDatabase, + op: &DsaModuleOp, + b: u32, + isl: u32, + prefix: u32, + dsa_backend: &str, +) -> Result<(f64, Source), AicError> { + let silicon = || { + db.dsa.query_context( &db.system_spec, - batch_size, + b, + isl, + op.num_heads, + op.kv_cache_dtype, + op.fmha_quant_mode, + op.gemm_quant_mode, + &op.architecture, + prefix, + op.index_topk, + dsa_backend, + ) + }; + match db.database_mode { + DatabaseMode::Empirical => Ok(( + context_empirical(db, op, b, isl, prefix, dsa_backend)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + context_empirical(db, op, b, isl, prefix, dsa_backend)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((silicon()?, Source::Silicon)), + } +} + +/// Generation DSA module latency for the op's slice under the database's mode. +fn query_generation_table( + db: &PerfDatabase, + op: &DsaModuleOp, + b: u32, + s: u32, + dsa_backend: &str, +) -> Result<(f64, Source), AicError> { + let silicon = || { + db.dsa.query_generation( + &db.system_spec, + b, s, - self.num_heads, - self.kv_cache_dtype, - self.fmha_quant_mode, - self.gemm_quant_mode, - &self.architecture, - "trtllm", + op.num_heads, + op.kv_cache_dtype, + op.fmha_quant_mode, + op.gemm_quant_mode, + &op.architecture, + dsa_backend, + ) + }; + match db.database_mode { + DatabaseMode::Empirical => Ok(( + generation_empirical(db, op, b, s, dsa_backend)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + generation_empirical(db, op, b, s, dsa_backend)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((silicon()?, Source::Silicon)), + } +} + +// --------------------------------------------------------------------------- +// Util-space empirical estimation (`latency = SOL(query) / util`), mirroring +// the `get_empirical` closures of `operations/dsa.py` branch for branch. +// --------------------------------------------------------------------------- + +/// Sample-coordinate → SOL-argument mapping for the selected context +/// calibration variant (the Python per-branch `_sol(c)` closures). +#[derive(Clone, Copy)] +enum CtxSolShape { + /// Samples `(prefix, s, b)` on the exact head slice: + /// `get_sol(c[2], c[1], c[0], num_heads, ...)`. + PrefixSeqBatch, + /// Samples `(num_heads, prefix, s, b)` across heads: + /// `get_sol(c[3], c[2], c[1], c[0], ...)`. + HeadPrefixSeqBatch, + /// Prefix=0-anchored samples `(s, b)` on the exact head slice: + /// `get_sol(c[1], c[0], 0, num_heads, ...)`. + SeqBatchP0, + /// Prefix=0-anchored samples `(num_heads, s, b)` across heads: + /// `get_sol(c[2], c[1], 0, c[0], ...)`. + HeadSeqBatchP0, +} + +/// Calibration data feeding the selected context variant's util grid. +enum CtxCalibration<'a> { + /// `[prefix][s][b]` — one head's full prefix-carrying sub-grid. + HeadPrefix(&'a BTreeMap>>), + /// `[num_heads][prefix][s][b]` — the whole backend-selected slice. + AllHeads(&'a DsaHeadGrid), + /// `[s][b]` — one head's prefix=0 anchor sub-grid. + HeadP0(&'a BTreeMap>), + /// `[num_heads] -> [s][b]` at prefix=0, heads without a 0 anchor dropped + /// (the Python dict comprehension). + AllHeadsP0(&'a DsaHeadGrid), + /// Python `calibration_data is None`: `slice_fn` raises the typed + /// coverage miss, so the grid is `None` and `estimate` reports the gap. + Missing, +} + +/// `SOL(query)/util` for the context DSA module. Mirrors Python +/// `ContextDSAModule._query_context_dsa_module_table::get_empirical` +/// (`skip_indexer=False`): detects the prefix axis, keeps the query on its +/// exact measured head slice when present, then selects the calibration +/// variant keyed by (has_prefix, exact_head, interp_prefix). The prefix=0 +/// anchor and boundary-freeze variants compute util at a DIFFERENT anchor +/// coordinate than the query (full_s = s + prefix at prefix=0), while the +/// prefix effect stays in the true query SOL. +fn context_empirical( + db: &PerfDatabase, + op: &DsaModuleOp, + b: u32, + s: u32, + prefix: u32, + dsa_backend: &str, +) -> Result { + let spec = &db.system_spec; + let dims = dsa_dims(&op.architecture); + let topk = op.index_topk as i64; + let (kv, fmha, gemm) = (op.kv_cache_dtype, op.fmha_quant_mode, op.gemm_quant_mode); + // Python's inner `get_sol(b, s, prefix, num_heads, kv, fmha)` closing + // over gemm_quant_mode / index_topk / dims. + let sol = |b: f64, s: f64, prefix: f64, num_heads: f64| { + dsa_context_sol_ms( + spec, + dims, + topk, + kv, + fmha, + gemm, + b as i64, + s as i64, + prefix as i64, + num_heads as i64, + ) + }; + let sol_time = sol(b as f64, s as f64, prefix as f64, op.num_heads as f64); + + // Raw calibration slice; a typed miss mirrors Python's + // `except PerfDataNotAvailableError: slc = None`. + let key = DsaKey { + architecture: op.architecture.clone(), + fmha_quant: fmha.name().to_string(), + kv_quant: kv.name().to_string(), + gemm_quant: gemm.name().to_string(), + }; + let slice = match db.dsa.context_raw_slice(&key, dsa_backend) { + Ok(slice) => Some(slice), + Err(err) if err.is_missing_perf_data() => None, + Err(err) => return Err(err), + }; + // Prefix-axis detection (`_dsa_module_has_prefix_axis`): the Rust loader + // always materialises the step level (legacy CSVs load as step=0), so any + // loaded slice carries the prefix axis; with the data unavailable Python + // falls back to the prior per-arch heuristic. + let has_prefix = match slice { + Some(_) => true, + None => op.architecture == "GlmMoeDsaForCausalLM", + }; + + // num_heads identifies a TP/model shape: stay on the exact measured head + // slice whenever it exists (`exact_head = isinstance(head_data, dict) and + // bool(head_data)`), cross-head fallback only when absent. + let head_grid = slice + .and_then(|slc| slc.get(&op.num_heads)) + .filter(|head_grid| !head_grid.is_empty()); + let exact_head = head_grid.is_some(); + + // Collected prefix values: the exact head's keys, else the sorted union + // across heads (BTreeMap/BTreeSet iteration == Python `sorted`). + let prefix_keys: Vec = if let Some(head_grid) = head_grid { + head_grid.keys().copied().collect() + } else if let Some(slc) = slice { + let mut seen: BTreeSet = BTreeSet::new(); + for by_prefix in slc.values() { + seen.extend(by_prefix.keys().copied()); + } + seen.into_iter().collect() + } else { + Vec::new() + }; + // Genuine prefix interpolation needs >= 2 collected prefix points + // bracketing the query; otherwise anchor util at the prefix=0 slice at + // full_s = s + prefix (regime-matched) so the indexer on/off boundary and + // small-s overhead floor stay in the (true) SOL. + let interp_prefix = prefix_keys.len() >= 2 + && prefix_keys[0] <= prefix + && prefix <= *prefix_keys.last().expect("non-empty prefix keys"); + + let (heads_f, prefix_f, s_f, b_f) = ( + op.num_heads as f64, + prefix as f64, + s as f64, + b as f64, + ); + let (tag, depth, query, shape, calibration, head_scoped): ( + &'static str, + usize, + Vec, + CtxSolShape, + CtxCalibration, + bool, + ) = if has_prefix && interp_prefix { + // Genuine measured prefix axis. Samples are (prefix, s, b) on an + // exact head, otherwise (num_heads, prefix, s, b). + if exact_head { + ( + "ctx_dsa_exact_head", + 3, + vec![prefix_f, s_f, b_f], + CtxSolShape::PrefixSeqBatch, + CtxCalibration::HeadPrefix(head_grid.expect("exact head slice")), + true, + ) + } else { + ( + "ctx_dsa", + 4, + vec![heads_f, prefix_f, s_f, b_f], + CtxSolShape::HeadPrefixSeqBatch, + CtxCalibration::AllHeads(slice.expect("populated slice")), + false, + ) + } + } else if has_prefix && prefix_keys.contains(&0) { + // Degenerate/out-of-range prefix axis: anchor utilization at prefix=0 + // and full_s = s + prefix so indexer/top-k regime changes remain in + // the true query SOL. + if exact_head { + ( + "ctx_dsa_p0anchor_exact_head", + 2, + vec![s_f + prefix_f, b_f], + CtxSolShape::SeqBatchP0, + CtxCalibration::HeadP0(&head_grid.expect("exact head slice")[&0]), + true, + ) + } else { + ( + "ctx_dsa_p0anchor", + 3, + vec![heads_f, s_f + prefix_f, b_f], + CtxSolShape::HeadSeqBatchP0, + CtxCalibration::AllHeadsP0(slice.expect("populated slice")), + false, + ) + } + } else if has_prefix { + // No prefix=0 anchor exists. Preserve coverage by freezing the + // generic raw-grid utilization at the nearest measured prefix. + if exact_head { + ( + "ctx_dsa_prefix_boundary_exact_head", + 3, + vec![prefix_f, s_f, b_f], + CtxSolShape::PrefixSeqBatch, + CtxCalibration::HeadPrefix(head_grid.expect("exact head slice")), + true, + ) + } else { + ( + "ctx_dsa_prefix_boundary", + 4, + vec![heads_f, prefix_f, s_f, b_f], + CtxSolShape::HeadPrefixSeqBatch, + slice.map_or(CtxCalibration::Missing, CtxCalibration::AllHeads), + false, + ) + } + } else { + // Legacy [heads][s][b] tables. Unreachable with data present in Rust + // (the loader always materialises the step level), so this only fires + // for the missing-slice non-GLM heuristic (grid stays None); the tags + // and query coordinates still mirror Python for the estimate() error. + if exact_head { + ( + "ctx_dsa_legacy_exact_head", + 2, + vec![s_f + prefix_f, b_f], + CtxSolShape::SeqBatchP0, + CtxCalibration::Missing, + true, + ) + } else { + ( + "ctx_dsa_legacy", + 3, + vec![heads_f, s_f + prefix_f, b_f], + CtxSolShape::HeadSeqBatchP0, + CtxCalibration::Missing, + false, + ) + } + }; + + // Python's per-branch `_sol(c)` closure over the selected variant. + let sample_sol = |c: &[f64]| match shape { + CtxSolShape::PrefixSeqBatch => sol(c[2], c[1], c[0], heads_f), + CtxSolShape::HeadPrefixSeqBatch => sol(c[3], c[2], c[1], c[0]), + CtxSolShape::SeqBatchP0 => sol(c[1], c[0], 0.0, heads_f), + CtxSolShape::HeadSeqBatchP0 => sol(c[2], c[1], 0.0, c[0]), + }; + + // Python keys the grid on (tag, system, backend, version, quants, arch, + // dsa_backend, depth) + id(node); the cache here is per-database, and the + // exact-head node identity is restored by suffixing num_heads. + let cache_key = if head_scoped { + format!( + "{tag}:{}:{}:{}:{}:{dsa_backend}:{depth}:h{}", + fmha.name(), + kv.name(), + gemm.name(), + op.architecture, + op.num_heads + ) + } else { + format!( + "{tag}:{}:{}:{}:{}:{dsa_backend}:{depth}", + fmha.name(), + kv.name(), + gemm.name(), + op.architecture + ) + }; + let grid = db.util_grids.get_or_try_build(&cache_key, || { + let points = match calibration { + CtxCalibration::HeadPrefix(head_grid) => context_points_head(head_grid), + CtxCalibration::AllHeads(slc) => context_points_all(slc), + CtxCalibration::HeadP0(p0) => points_2d(p0), + CtxCalibration::AllHeadsP0(slc) => context_points_all_p0(slc), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss), mirroring Python's slice_fn raising + // PerfDataNotAvailableError inside grid_for. + CtxCalibration::Missing => return Ok(None), + }; + Ok(Some(UtilGrid::new(util_empirical::build_samples( + points, + sample_sol, + )))) + })?; + let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + +/// `SOL(query)/util` for the generation DSA module. Mirrors Python +/// `GenerationDSAModule._query_generation_dsa_module_table::get_empirical` +/// (`skip_indexer=False`): exact measured head slice when present +/// (`num_heads in data_slice`), cross-head nearest-neighbour otherwise, +/// typed empirical miss when the slice is absent. +fn generation_empirical( + db: &PerfDatabase, + op: &DsaModuleOp, + b: u32, + s: u32, + dsa_backend: &str, +) -> Result { + let spec = &db.system_spec; + let dims = dsa_dims(&op.architecture); + let (kv, gemm) = (op.kv_cache_dtype, op.gemm_quant_mode); + // Python's inner `get_sol(b, s, num_heads, kv_cache_dtype)` (the + // attention group is hardcoded bfloat16 inside). + let sol = |b: f64, s: f64, num_heads: f64| { + dsa_generation_sol_ms(spec, dims, kv, gemm, b as i64, s as i64, num_heads as i64) + }; + let heads_f = op.num_heads as f64; + let sol_time = sol(b as f64, s as f64, heads_f); + + // NOTE: Python's generation slice is keyed (kv, gemm, arch) only — no + // mla_dtype axis. The Rust table retains `mla_dtype` in `DsaKey` (see + // `query_generation`); collected generation files carry uniformly + // `bfloat16`, and the empirical slice follows the silicon keying so + // HYBRID's silicon and empirical stages always agree on the slice. + let key = DsaKey { + architecture: op.architecture.clone(), + fmha_quant: op.fmha_quant_mode.name().to_string(), + kv_quant: kv.name().to_string(), + gemm_quant: gemm.name().to_string(), + }; + let slice = match db.dsa.generation_raw_slice(&key, dsa_backend) { + Ok(slice) => Some(slice), + // Match `grid_for`'s best-effort contract: unavailable table data is + // reported by `estimate` as an empirical coverage miss. + Err(err) if err.is_missing_perf_data() => None, + Err(err) => return Err(err), + }; + + if let Some(head_grid) = slice.and_then(|slc| slc.get(&op.num_heads)) { + // `num_heads` is a TP/model-shape identity: stay on the exact head + // slice whenever it exists (Python `num_heads in data_slice`). + let cache_key = format!( + "gen_dsa_exact_heads:{}:{}:{}:{}:{dsa_backend}:h{}", + op.fmha_quant_mode.name(), + kv.name(), + gemm.name(), + op.architecture, + op.num_heads + ); + let grid = db.util_grids.get_or_try_build(&cache_key, || { + Ok(Some(UtilGrid::new(util_empirical::build_samples( + generation_points_head(head_grid), + // c = (b, s) + |c: &[f64]| sol(c[0], c[1], heads_f), + )))) + })?; + let (latency, _) = + util_empirical::estimate(sol_time, &[b as f64, s as f64], grid.as_deref(), 1.0)?; + Ok(latency) + } else if let Some(slc) = slice { + let cache_key = format!( + "gen_dsa:{}:{}:{}:{}:{dsa_backend}", + op.fmha_quant_mode.name(), + kv.name(), + gemm.name(), + op.architecture + ); + let grid = db.util_grids.get_or_try_build(&cache_key, || { + Ok(Some(UtilGrid::new(util_empirical::build_samples( + generation_points_all(slc), + // c = (num_heads, b, s) + |c: &[f64]| sol(c[1], c[2], c[0]), + )))) + })?; + let (latency, _) = util_empirical::estimate( + sol_time, + &[heads_f, b as f64, s as f64], + grid.as_deref(), + 1.0, )?; - Ok(PerformanceResult::new(latency, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) + Ok(latency) + } else { + let (latency, _) = util_empirical::estimate( + sol_time, + &[heads_f, b as f64, s as f64], + None, + 1.0, + )?; + Ok(latency) } } +// --------------------------------------------------------------------------- +// Raw-grid point extraction (Python `iter_grid` over the selected variant's +// nested dict). Rust iterates the BTreeMaps in ascending key order — Python +// walks insertion (file first-occurrence) order; the k=2 IDW estimate is +// order-independent except for exact distance ties. +// --------------------------------------------------------------------------- + +/// `(s, b)` points of one `[s][b]` sub-grid. +fn points_2d(grid: &BTreeMap>) -> Vec<(Vec, f64)> { + let mut points = Vec::new(); + for (&s, by_b) in grid { + for (&b, &latency) in by_b { + points.push((vec![s as f64, b as f64], latency)); + } + } + points +} + +/// `(prefix, s, b)` points of one head's `[prefix][s][b]` sub-grid. +fn context_points_head( + head_grid: &BTreeMap>>, +) -> Vec<(Vec, f64)> { + let mut points = Vec::new(); + for (&prefix, by_s) in head_grid { + for (&s, by_b) in by_s { + for (&b, &latency) in by_b { + points.push((vec![prefix as f64, s as f64, b as f64], latency)); + } + } + } + points +} + +/// `(num_heads, prefix, s, b)` points of a whole `[heads][prefix][s][b]` slice. +fn context_points_all(slice: &DsaHeadGrid) -> Vec<(Vec, f64)> { + let mut points = Vec::new(); + for (&heads, head_grid) in slice { + for (&prefix, by_s) in head_grid { + for (&s, by_b) in by_s { + for (&b, &latency) in by_b { + points.push(( + vec![heads as f64, prefix as f64, s as f64, b as f64], + latency, + )); + } + } + } + } + points +} + +/// `(num_heads, s, b)` points of the prefix=0 anchor across heads; heads +/// without a prefix=0 sub-grid are dropped (the Python dict comprehension). +fn context_points_all_p0(slice: &DsaHeadGrid) -> Vec<(Vec, f64)> { + let mut points = Vec::new(); + for (&heads, head_grid) in slice { + if let Some(by_s) = head_grid.get(&0) { + for (&s, by_b) in by_s { + for (&b, &latency) in by_b { + points.push((vec![heads as f64, s as f64, b as f64], latency)); + } + } + } + } + points +} + +/// `(b, s)` points of one head's generation sub-grid. The load-time +/// (isl, step) -> seq collapse stores `[0][seq][batch]`; Python's raw view is +/// `[b][s]`, so the coordinate order flips to (batch, seq). +fn generation_points_head( + head_grid: &BTreeMap>>, +) -> Vec<(Vec, f64)> { + let mut points = Vec::new(); + for by_seq in head_grid.values() { + for (&seq, by_b) in by_seq { + for (&b, &latency) in by_b { + points.push((vec![b as f64, seq as f64], latency)); + } + } + } + points +} + +/// `(num_heads, b, s)` points of a whole generation slice. +fn generation_points_all(slice: &DsaHeadGrid) -> Vec<(Vec, f64)> { + let mut points = Vec::new(); + for (&heads, head_grid) in slice { + for by_seq in head_grid.values() { + for (&seq, by_b) in by_seq { + for (&b, &latency) in by_b { + points.push((vec![heads as f64, b as f64, seq as f64], latency)); + } + } + } + } + points +} + #[cfg(test)] mod tests { use super::*; @@ -410,4 +952,191 @@ mod tests { let de: DsaModuleOp = serde_json::from_value(v).expect("deserialize"); assert_eq!(de.cp_size, 1); } + + // ------------------------------------------------------------------ + // Util-space empirical (EMPIRICAL / HYBRID) oracle parity. + // + // Oracle values generated from the Python reference on the same data + // (shared layer OFF so both sides read the same single parquet): + // + // ```text + // uv run --no-sync python - <<'PY' + // from aiconfigurator.sdk import perf_database, common + // db = perf_database.get_database("b200_sxm", , , + // shared_layer=False) + // r = db.query_context_dsa_module( # or query_generation_dsa_module + // b=..., s=..., prefix=..., num_heads=..., + // kvcache_quant_mode=..., fmha_quant_mode=..., gemm_quant_mode=..., + // database_mode=common.DatabaseMode.EMPIRICAL, architecture=...) + // print(float(r)) + // PY + // ``` + // + // The fired calibration variant was verified per anchor by instrumenting + // `util_empirical.grid_for` in the oracle script (not committed). + // Regenerate if the shipped DSA tables or the util-empirical math change. + // ------------------------------------------------------------------ + + use crate::common::enums::DatabaseMode; + use crate::operators::base::Source; + use crate::perf_database::PerfDatabase; + + const DSV32: &str = "DeepseekV32ForCausalLM"; + const GLM: &str = "GlmMoeDsaForCausalLM"; + + fn b200_db(backend: &str, version: &str, mode: DatabaseMode) -> PerfDatabase { + let systems_root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + let mut db = PerfDatabase::load(&systems_root, "b200_sxm", backend, version) + .expect("b200_sxm db must load"); + db.database_mode = mode; + db + } + + fn dsa_op(architecture: &str, num_heads: u32, kv: KvCacheQuantMode) -> DsaModuleOp { + DsaModuleOp::new( + "dsa_module", + num_heads, + kv, + FmhaQuantMode::Bfloat16, + GemmQuantMode::Bfloat16, + architecture, + 2048, + ) + } + + fn approx_rel_1e9(got: f64, want: f64) { + assert!( + ((got - want) / want).abs() < 1e-9, + "rust {got} vs python {want}" + ); + } + + /// Context EMPIRICAL parity on b200_sxm/vllm/0.19.0. Fired variants + /// (verified in Python): + /// - DSV32 h=128 prefix=0 off-grid s AND prefix=4096: the DSV32 slice + /// collects prefix=[0] only, so BOTH anchor util at the prefix=0 slice + /// with full_s = s + prefix -> `ctx_dsa_p0anchor_exact_head` (depth 2). + /// - GLM h=32 prefix=64: bracketed by the collected [0, 128] axis -> + /// `ctx_dsa_exact_head` (depth 3, genuine prefix interpolation). + /// - GLM h=32 prefix=1000: beyond 128 -> `ctx_dsa_p0anchor_exact_head`. + /// - DSV32 h=96 (uncollected head) -> cross-head `ctx_dsa_p0anchor` + /// (depth 3). + #[test] + fn context_empirical_matches_python_oracles() { + let db = b200_db("vllm", "0.19.0", DatabaseMode::Empirical); + let cases: [(&str, u32, u32, u32, u32, f64); 5] = [ + (DSV32, 128, 4, 3000, 0, 11.266464115749121), + (DSV32, 128, 2, 2048, 4096, 4.417586773123638), + (GLM, 32, 1, 512, 64, 2.62689536098129), + (GLM, 32, 1, 512, 1000, 0.514264970789275), + (DSV32, 96, 8, 1024, 0, 5.571143850099078), + ]; + for (arch, heads, b, s, prefix, expected) in cases { + let op = dsa_op(arch, heads, KvCacheQuantMode::Bfloat16); + let result = op.query_context(&db, b, s, prefix).expect("empirical query"); + approx_rel_1e9(result.latency_ms, expected); + assert_eq!(result.source, Source::Empirical, "({arch}, h={heads})"); + } + } + + /// Context EMPIRICAL parity on b200_sxm/sglang/0.5.14 (GLM-5, bf16 KV; + /// heads collected: 8/16/32/64 with a rich measured prefix axis): + /// - h=128 (uncollected) prefix=512 in range -> cross-head `ctx_dsa` + /// (depth 4). + /// - h=64 prefix=10000 at a collected (prefix, s, b) point -> + /// `ctx_dsa_exact_head` exact hit returns the measured latency. + #[test] + fn context_empirical_sglang_glm_matches_python_oracles() { + let db = b200_db("sglang", "0.5.14", DatabaseMode::Empirical); + let cross = dsa_op(GLM, 128, KvCacheQuantMode::Bfloat16); + let result = cross.query_context(&db, 2, 4096, 512).expect("cross-head query"); + approx_rel_1e9(result.latency_ms, 13.731357702671428); + assert_eq!(result.source, Source::Empirical); + + let exact = dsa_op(GLM, 64, KvCacheQuantMode::Bfloat16); + let result = exact.query_context(&db, 2, 4096, 10000).expect("exact-hit query"); + approx_rel_1e9(result.latency_ms, 8.2065); + assert_eq!(result.source, Source::Empirical); + } + + /// HYBRID keeps the silicon answer whenever the silicon lookup resolves + /// (exact collected hit AND an interpolated interior point) — the + /// empirical layer must not preempt it. + #[test] + fn context_hybrid_prefers_silicon() { + let db = b200_db("vllm", "0.19.0", DatabaseMode::Hybrid); + let op = dsa_op(DSV32, 128, KvCacheQuantMode::Bfloat16); + let exact = op.query_context(&db, 4, 2048, 0).expect("silicon exact hit"); + approx_rel_1e9(exact.latency_ms, 7.6471); + assert_eq!(exact.source, Source::Silicon); + let interior = op.query_context(&db, 4, 3000, 0).expect("silicon interp"); + approx_rel_1e9(interior.latency_ms, 11.388627343749999); + assert_eq!(interior.source, Source::Silicon); + } + + /// An uncollected kv-quant slice (int8 on b200/vllm) is a typed silicon + /// miss; HYBRID falls to the empirical layer, which also has no + /// calibration data (DSV32 missing-slice heuristic -> legacy variant, + /// grid None) and must surface the terminal EmpiricalNotImplemented — + /// never a fabricated value. Same terminal miss under EMPIRICAL. + #[test] + fn context_missing_slice_raises_empirical_not_implemented() { + for mode in [DatabaseMode::Hybrid, DatabaseMode::Empirical] { + let db = b200_db("vllm", "0.19.0", mode); + let op = dsa_op(DSV32, 128, KvCacheQuantMode::Int8); + let result = op.query_context(&db, 4, 3000, 0); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "{mode:?}: got {result:?}" + ); + } + } + + /// Generation EMPIRICAL parity on b200_sxm/vllm/0.19.0: + /// - h=128 off-grid (b=24, s=3000) and an exact collected hit + /// (b=16, s=4097 -> measured 0.2698) on the exact head slice + /// (`gen_dsa_exact_heads`, depth 2); + /// - h=96 (uncollected head) -> cross-head `gen_dsa` (depth 3). + #[test] + fn generation_empirical_matches_python_oracles() { + let db = b200_db("vllm", "0.19.0", DatabaseMode::Empirical); + let cases: [(u32, u32, u32, f64); 3] = [ + (128, 24, 3000, 0.259452522778543), + (128, 16, 4097, 0.2698), + (96, 8, 5000, 0.20801465850337103), + ]; + for (heads, b, s, expected) in cases { + let op = dsa_op(DSV32, heads, KvCacheQuantMode::Bfloat16); + let result = op.query_generation(&db, b, s).expect("empirical query"); + approx_rel_1e9(result.latency_ms, expected); + assert_eq!(result.source, Source::Empirical, "(h={heads}, b={b}, s={s})"); + } + } + + /// Generation EMPIRICAL parity on b200_sxm/sglang/0.5.14 (GLM-5, bf16 KV, + /// exact head slice off-grid). + #[test] + fn generation_empirical_sglang_glm_matches_python_oracle() { + let db = b200_db("sglang", "0.5.14", DatabaseMode::Empirical); + let op = dsa_op(GLM, 64, KvCacheQuantMode::Bfloat16); + let result = op.query_generation(&db, 48, 10000).expect("empirical query"); + approx_rel_1e9(result.latency_ms, 0.1927939670669063); + assert_eq!(result.source, Source::Empirical); + } + + /// Generation HYBRID on an uncollected kv-quant slice: typed silicon miss + /// -> empirical -> no calibration slice -> terminal + /// EmpiricalNotImplemented (mirrors the Python contract). + #[test] + fn generation_hybrid_missing_slice_raises_empirical_not_implemented() { + let db = b200_db("vllm", "0.19.0", DatabaseMode::Hybrid); + let op = dsa_op(DSV32, 128, KvCacheQuantMode::Int8); + let result = op.query_generation(&db, 24, 3000); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "got {result:?}" + ); + } } diff --git a/rust/aiconfigurator-core/src/operators/dsv4.rs b/rust/aiconfigurator-core/src/operators/dsv4.rs index 1ea690b81..24fe4c56a 100644 --- a/rust/aiconfigurator-core/src/operators/dsv4.rs +++ b/rust/aiconfigurator-core/src/operators/dsv4.rs @@ -17,11 +17,12 @@ //! `perf_database::dsv4` and Python `load_context_dsv4_kind_module_data`. use serde::{Deserialize, Serialize}; -use crate::common::enums::{FmhaQuantMode, GemmQuantMode, KvCacheQuantMode}; +use crate::common::enums::{DatabaseMode, FmhaQuantMode, GemmQuantMode, KvCacheQuantMode}; use crate::common::error::AicError; use crate::operators::base::{PerformanceResult, Source}; use crate::operators::communication::NcclOp; -use crate::perf_database::dsv4::{dsv4_dims, AttnKind, Dsv4SolDims}; +use crate::operators::util_empirical::{self, UtilGrid}; +use crate::perf_database::dsv4::{dsv4_attention_sol_ms, dsv4_dims, AttnKind, Dsv4SolDims}; use crate::perf_database::PerfDatabase; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -211,13 +212,41 @@ impl Dsv4ModuleOp { /// The per-(b, s, prefix) module base query shared by the plain and CP /// paths (Python `_module_base` -> the standard module query, which - /// includes the CSA topk DELTA correction). + /// includes the CSA topk DELTA correction and dispatches on the database + /// mode): SILICON queries the table; HYBRID converts a typed silicon miss + /// into the util-space empirical estimate; EMPIRICAL always estimates. + /// The SOL diagnostic modes never reach the compiled engine (the routing + /// gate delegates them to the Python step). fn module_base( &self, db: &PerfDatabase, batch_size: u32, s: u32, prefix: u32, + ) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + self.context_empirical(db, batch_size, s, prefix)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match self.context_silicon(db, batch_size, s, prefix) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + self.context_empirical(db, batch_size, s, prefix)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((self.context_silicon(db, batch_size, s, prefix)?, Source::Silicon)), + } + } + + fn context_silicon( + &self, + db: &PerfDatabase, + batch_size: u32, + s: u32, + prefix: u32, ) -> Result { db.dsv4.query_context( &db.system_spec, @@ -234,6 +263,102 @@ impl Dsv4ModuleOp { ) } + /// `SOL(query)/util` over the op's own context slice. Mirrors Python + /// `_query_context_attn_table::get_empirical`: + /// + /// - Genuine prefix interpolation needs >= 2 collected prefix points + /// bracketing the query: depth-3 grid over `(prefix, s, b)`, query + /// `(prefix, s, b)`. + /// - Otherwise (degenerate prefix axis or out-of-range query) anchor at + /// the prefix=0 slice at `full_s = s + prefix` (regime-matched), with + /// the prefix effect carried by `sol_q`: depth-2 grid over `(s, b)`, + /// query `(s + prefix, b)`. + fn context_empirical( + &self, + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + ) -> Result { + let spec = &db.system_spec; + let dims = self.sol_dims(); + let cr = self.attn_kind.compress_ratio(); + let heads = i64::from(self.num_heads); + let (kv, fmha, gemm) = (self.kv_cache_dtype, self.fmha_quant_mode, self.gemm_quant_mode); + let sol_at = move |b_: i64, s_: i64, p_: i64| { + dsv4_attention_sol_ms(spec, &dims, cr, true, kv, fmha, gemm, b_, s_, p_, heads) + }; + // True SOL(b, s, prefix) at the query (Python `sol_q = get_sol()[0]`). + let sol_q = sol_at(i64::from(b), i64::from(s), i64::from(prefix)); + + // Own-slice `(prefix, s, b)` points; a typed coverage miss means no + // prefix keys (Python: `prefix_keys = ()`), so the p0-anchor branch + // then finds no grid and estimate() raises the empirical miss. + let points = match db.dsv4.context_points( + self.attn_kind, + self.num_heads, + kv, + fmha, + gemm, + &self.architecture, + ) { + Ok(points) => Some(points), + Err(err) if err.is_missing_perf_data() => None, + Err(err) => return Err(err), + }; + let prefix_keys: std::collections::BTreeSet = points + .iter() + .flatten() + .map(|(coords, _)| coords[0] as u32) + .collect(); + let interp_prefix = prefix_keys.len() >= 2 + && *prefix_keys.first().expect("non-empty") <= prefix + && prefix <= *prefix_keys.last().expect("non-empty"); + + // Grid cache key mirrors Python's + // (key_tag, quants, num_heads, cr, depth); `architecture` stands in + // for Python's `id(node)` identity component (the Rust slice keys it). + let key_stem = format!( + "{}:{}:{}:{}:{}:{}", + self.architecture, + fmha.name(), + kv.name(), + gemm.name(), + self.num_heads, + cr + ); + if interp_prefix { + let sol3 = |c: &[f64]| sol_at(c[2] as i64, c[1] as i64, c[0] as i64); // c=(prefix, s, b) + let key = format!("dsv4_ctx_attn:{key_stem}:3"); + let grid = db.util_grids.get_or_try_build(&key, || { + Ok(points.map(|points| UtilGrid::new(util_empirical::build_samples(points, sol3)))) + })?; + let query = [f64::from(prefix), f64::from(s), f64::from(b)]; + let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + Ok(latency) + } else { + let sol2 = |c: &[f64]| sol_at(c[1] as i64, c[0] as i64, 0); // c=(s, b); anchor prefix=0 + let key = format!("dsv4_ctx_attn_p0anchor:{key_stem}:2"); + let grid = db.util_grids.get_or_try_build(&key, || { + // Python `require_data_slice(_slice(), 0)`: no prefix=0 rows + // is a typed coverage miss -> no grid. + let p0_points: Vec<(Vec, f64)> = points + .into_iter() + .flatten() + .filter(|(coords, _)| coords[0] == 0.0) + .map(|(coords, latency)| (vec![coords[1], coords[2]], latency)) + .collect(); + if p0_points.is_empty() { + return Ok(None); + } + Ok(Some(UtilGrid::new(util_empirical::build_samples(p0_points, sol2)))) + })?; + let query = [f64::from(s) + f64::from(prefix), f64::from(b)]; + let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + Ok(latency) + } + } + pub fn query_context( &self, db: &PerfDatabase, @@ -251,12 +376,12 @@ impl Dsv4ModuleOp { // (SILICON path): a 3-axis perf_interp v2 Grid query over // `(prefix, isl, batch)` with the step axis KEPT. The caller supplies // the new-token count as `isl` (Python's `s = effective_isl = - // isl - prefix`, computed in `run_context_ops`); the context CSVs - // collected to date carry a single `step=0` anchor, so `prefix=0` - // collapses that level exactly and `prefix>0` resolves via util-hold - // with the prefix-aware SOL carrying the effect — matching Python. - let raw = self.module_base(db, batch_size, isl, prefix)?; - Ok(PerformanceResult::new(raw, Source::Silicon) + // isl - prefix`, computed in `run_context_ops`); a prefix beyond the + // collected range resolves via util-hold with the prefix-aware SOL + // carrying the effect — matching Python. HYBRID/EMPIRICAL route + // through `module_base`'s mode dispatch. + let (raw, source) = self.module_base(db, batch_size, isl, prefix)?; + Ok(PerformanceResult::new(raw, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -286,7 +411,10 @@ impl Dsv4ModuleOp { batch_size, isl, prefix, - &mut |per_card| self.module_base(db, batch_size, per_card, prefix), + &mut |per_card| { + self.module_base(db, batch_size, per_card, prefix) + .map(|(latency, _)| latency) + }, &mut |chunk_isl, past| { db.dsv4.query_paged_mqa_logits( batch_size, @@ -405,13 +533,38 @@ impl Dsv4ModuleOp { Ok(PerformanceResult::new(latency, Source::Estimated).scaled(self.scale_factor)) } + /// Database-mode dispatch mirroring Python + /// `_query_generation_attn_table` (`operations/dsv4.py`): SILICON queries + /// the table; HYBRID converts a typed silicon miss into the util-space + /// empirical estimate; EMPIRICAL always estimates. pub fn query_generation( &self, db: &PerfDatabase, batch_size: u32, s: u32, ) -> Result { - let latency = db.dsv4.query_generation( + let (latency, source) = match db.database_mode { + DatabaseMode::Empirical => ( + self.generation_empirical(db, batch_size, s)?, + Source::Empirical, + ), + DatabaseMode::Hybrid => match self.generation_silicon(db, batch_size, s) { + Ok(latency) => (latency, Source::Silicon), + Err(err) if err.is_missing_perf_data() => ( + self.generation_empirical(db, batch_size, s)?, + Source::Empirical, + ), + Err(err) => return Err(err), + }, + _ => (self.generation_silicon(db, batch_size, s)?, Source::Silicon), + }; + Ok(PerformanceResult::new(latency, source) + .clamp_non_negative() + .scaled(self.scale_factor)) + } + + fn generation_silicon(&self, db: &PerfDatabase, batch_size: u32, s: u32) -> Result { + db.dsv4.query_generation( &db.system_spec, self.attn_kind, batch_size, @@ -422,10 +575,68 @@ impl Dsv4ModuleOp { self.gemm_quant_mode, &self.architecture, Some(self.sol_dims()), - )?; - Ok(PerformanceResult::new(latency, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) + ) + } + + /// `SOL(query)/util` over the op's own `(b, s_total)` generation slice. + /// Mirrors Python `_query_generation_attn_table::get_empirical` (grid + /// depth 2, `sol_fn = lambda c: get_sol(c[0], c[1])[0]`, query `(b, s)`). + fn generation_empirical(&self, db: &PerfDatabase, b: u32, s: u32) -> Result { + let spec = &db.system_spec; + let dims = self.sol_dims(); + let cr = self.attn_kind.compress_ratio(); + let heads = i64::from(self.num_heads); + let (kv, gemm) = (self.kv_cache_dtype, self.gemm_quant_mode); + // Python derives the decode SOL dtype from the kv-cache dtype at the + // top of `_query_generation_attn_table` (the fmha label is inert for + // generation; the table keys on kv dtype). + let fmha = if kv == KvCacheQuantMode::Fp8 { + FmhaQuantMode::Fp8 + } else { + FmhaQuantMode::Bfloat16 + }; + let sol = move |c: &[f64]| { + // c=(b, s_total); is_context=false, prefix=0. + dsv4_attention_sol_ms( + spec, + &dims, + cr, + false, + kv, + fmha, + gemm, + c[0] as i64, + c[1] as i64, + 0, + heads, + ) + }; + // Python keys the grid on (kv, gemm, num_heads, cr) — no fmha level + // (derived from kv above); `architecture` stands in for Python's + // `id(node)` identity component. + let key = format!( + "dsv4_gen_attn:{}:{}:{}:{}:{}", + self.architecture, + kv.name(), + gemm.name(), + self.num_heads, + cr + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db + .dsv4 + .generation_points(self.attn_kind, self.num_heads, kv, gemm, &self.architecture) + { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [f64::from(b), f64::from(s)]; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) } } @@ -655,6 +866,205 @@ mod tests { ); } + fn systems_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems") + } + + fn b200_sglang_root() -> PathBuf { + systems_root().join("data/b200_sxm/sglang/0.5.10") + } + + /// DeepSeek-V4-Pro op at tp=8: rank-local num_heads = 128/8 = 16, + /// o_groups: None derives local 2 from the pinned totals — matching the + /// Python oracle calls (num_heads=16, o_groups=2, window_size=128). + fn dsv4_pro_op(attn_kind: AttnKind) -> Dsv4ModuleOp { + Dsv4ModuleOp { + name: "dsv4_pro".into(), + scale_factor: 1.0, + attn_kind, + num_heads: 16, + native_heads: 128, + tp_size: 8, + kv_cache_dtype: KvCacheQuantMode::Fp8, + fmha_quant_mode: FmhaQuantMode::Bfloat16, + gemm_quant_mode: GemmQuantMode::Fp8Block, + architecture: "DeepseekV4ForCausalLM".into(), + cp_size: 1, + window_size: None, + hidden_size: default_hidden_size(), + q_lora_rank: default_q_lora_rank(), + o_lora_rank: default_o_lora_rank(), + head_dim: default_head_dim(), + rope_head_dim: default_rope_head_dim(), + index_n_heads: default_index_n_heads(), + index_head_dim: default_index_head_dim(), + index_topk: default_index_topk(), + o_groups: None, + } + } + + fn approx(got: f64, want: f64) { + assert!( + ((got - want) / want).abs() < 1e-9, + "rust {got} vs python {want}" + ); + } + + /// Oracle values generated from the Python reference on the same data: + /// + /// ```text + /// uv run --no-sync python3 -c " + /// from aiconfigurator.sdk import perf_database, common + /// from aiconfigurator.sdk.operations.dsv4 import ( + /// ContextDeepSeekV4AttentionModule as CTX, + /// GenerationDeepSeekV4AttentionModule as GEN) + /// db = perf_database.get_database('b200_sxm', 'sglang', '0.5.10') + /// pro = dict(num_heads=16, native_heads=128, tp_size=8, hidden_size=7168, + /// q_lora_rank=1536, o_lora_rank=1024, head_dim=512, rope_head_dim=64, + /// index_n_heads=64, index_head_dim=128, index_topk=1024, window_size=128, + /// o_groups=2, kvcache_quant_mode=common.KVCacheQuantMode.fp8, + /// fmha_quant_mode=common.FMHAQuantMode.bfloat16, + /// gemm_quant_mode=common.GEMMQuantMode.fp8_block) + /// EMP = common.DatabaseMode.EMPIRICAL + /// for b, s, p, cr in [(8,512,0,4), (8,700,0,4), (8,512,1024,4), (3,8192,4096,4), + /// (1,128,0,128), (8,512,1024,128)]: + /// print(float(CTX._query_context_attn_table(db, b=b, s=s, prefix=p, + /// compress_ratio=cr, database_mode=EMP, **pro))) + /// for b, s, cr in [(16,385,4), (16,200,4), (13,100000,4), (16,385,128)]: + /// print(float(GEN._query_generation_attn_table(db, b=b, s=s, + /// compress_ratio=cr, database_mode=EMP, **pro)))" + /// ``` + /// + /// The 0.5.10 context tables carry a single step=0 anchor, so every + /// context case exercises the p0-anchor branch (depth-2 grid, query + /// `(s + prefix, b)`, the prefix effect carried by sol_q). Covers per + /// phase: exact collected hit, off-grid interior IDW, prefix>0 anchored + /// at full_s, and the generation kv->fmha SOL derivation. Regenerate if + /// the shipped tables or the util-empirical math change. + #[test] + fn dsv4_empirical_matches_python_oracles() { + let root = b200_sglang_root(); + if !root.join("dsv4_csa_context_module_perf.parquet").exists() { + return; // git-lfs data not materialized + } + let mut db = PerfDatabase::load(&systems_root(), "b200_sxm", "sglang", "0.5.10") + .expect("b200_sxm/sglang/0.5.10 must load"); + db.database_mode = DatabaseMode::Empirical; + + let ctx_cases: &[(AttnKind, u32, u32, u32, f64)] = &[ + // exact collected hit (prefix=0) reconstructs the measured value + (AttnKind::Csa, 8, 512, 0, 0.9819), + // off-grid interior isl + (AttnKind::Csa, 8, 700, 0, 1.3345620964544214), + // prefix>0: p0-anchor at full_s = s + prefix, SOL carries prefix + (AttnKind::Csa, 8, 512, 1024, 1.0674237160346303), + (AttnKind::Csa, 3, 8192, 4096, 11.05984646800717), + (AttnKind::Hca, 1, 128, 0, 0.0802), + (AttnKind::Hca, 8, 512, 1024, 0.5717204610218155), + ]; + for &(kind, b, s, prefix, expected) in ctx_cases { + let result = dsv4_pro_op(kind) + .query_context(&db, b, s, prefix) + .expect("empirical context query"); + approx(result.latency_ms, expected); + assert_eq!(result.source, Source::Empirical); + } + + let gen_cases: &[(AttnKind, u32, u32, f64)] = &[ + (AttnKind::Csa, 16, 385, 0.11423651225442899), + (AttnKind::Csa, 16, 200, 0.11330535071492717), + // beyond-range s_total: boundary util frozen, decode SOL carries it + (AttnKind::Csa, 13, 100_000, 0.16844612496931166), + (AttnKind::Hca, 16, 385, 0.07241533398068029), + ]; + for &(kind, b, s, expected) in gen_cases { + let result = dsv4_pro_op(kind) + .query_generation(&db, b, s) + .expect("empirical generation query"); + approx(result.latency_ms, expected); + assert_eq!(result.source, Source::Empirical); + } + } + + /// The genuine prefix-interpolation branch (>= 2 collected prefix keys + /// bracketing the query -> depth-3 grid over `(prefix, s, b)`), on the + /// shipped b200_sxm/sglang/0.5.14 tables (28 prefix keys, 0..1048575). + /// Same Python oracle command as above with '0.5.14': + /// + /// ```text + /// for b, s, p, cr in [(8,512,1024,4), (8,700,1000,4), (8,512,2097150,4), (8,700,1000,128)]: ... + /// ``` + /// + /// Covers: exact 3-axis hit, off-grid (prefix, s) IDW, HCA, and a prefix + /// beyond the collected range falling back to the p0-anchor branch + /// (query `(s + prefix, b)` = (2097662, 8)). + #[test] + fn dsv4_empirical_prefix_interp_matches_python_oracles() { + let systems_root = systems_root(); + let data_root = systems_root.join("data/b200_sxm/sglang/0.5.14"); + if !data_root.join("dsv4_csa_context_module_perf.parquet").exists() { + return; // git-lfs data not materialized + } + let mut db = PerfDatabase::load(&systems_root, "b200_sxm", "sglang", "0.5.14") + .expect("b200_sxm/sglang/0.5.14 must load"); + db.database_mode = DatabaseMode::Empirical; + + let cases: &[(AttnKind, u32, u32, u32, f64)] = &[ + (AttnKind::Csa, 8, 512, 1024, 1.0183), + (AttnKind::Csa, 8, 700, 1000, 1.429314779309641), + (AttnKind::Csa, 8, 512, 2_097_150, 32.64516668240808), + (AttnKind::Hca, 8, 700, 1000, 0.836433276456113), + ]; + for &(kind, b, s, prefix, expected) in cases { + let result = dsv4_pro_op(kind) + .query_context(&db, b, s, prefix) + .expect("empirical context query"); + approx(result.latency_ms, expected); + assert_eq!(result.source, Source::Empirical); + } + } + + /// HYBRID with silicon data present must stay on the silicon path; a kv + /// dtype with NO collected slice (bfloat16 — only fp8 ships) must fall + /// through the empirical path and surface the terminal + /// EmpiricalNotImplemented miss, never a fabricated value (Python: the + /// silicon miss falls to `get_empirical`, whose own typed miss raises + /// `EmpiricalNotImplementedError`). + #[test] + fn dsv4_hybrid_silicon_passthrough_and_terminal_miss() { + let root = b200_sglang_root(); + if !root.join("dsv4_csa_context_module_perf.parquet").exists() { + return; // git-lfs data not materialized + } + let mut db = PerfDatabase::load(&systems_root(), "b200_sxm", "sglang", "0.5.10") + .expect("b200_sxm/sglang/0.5.10 must load"); + db.database_mode = DatabaseMode::Hybrid; + + // Data present: silicon passthrough (exact grid point, same value as + // the silicon parity test). + let result = dsv4_pro_op(AttnKind::Csa) + .query_context(&db, 8, 512, 0) + .expect("hybrid context query"); + approx(result.latency_ms, 0.9819); + assert_eq!(result.source, Source::Silicon); + + // kv=bfloat16 is not collected: typed silicon miss -> empirical miss. + let mut op = dsv4_pro_op(AttnKind::Csa); + op.kv_cache_dtype = KvCacheQuantMode::Bfloat16; + let ctx = op.query_context(&db, 8, 512, 0); + assert!( + matches!(ctx, Err(AicError::EmpiricalNotImplemented(_))), + "got {ctx:?}" + ); + let gen = op.query_generation(&db, 16, 385); + assert!( + matches!(gen, Err(AicError::EmpiricalNotImplemented(_))), + "got {gen:?}" + ); + } + /// `cp_size` / `window_size` are absent from every opspec the Python /// emitter produces today (`engine.py::_reject_cp` still guards CP) — /// they must default to 1 / None so existing specs keep the plain non-CP diff --git a/rust/aiconfigurator-core/src/operators/mhc.rs b/rust/aiconfigurator-core/src/operators/mhc.rs index df9279e65..5246a4ca6 100644 --- a/rust/aiconfigurator-core/src/operators/mhc.rs +++ b/rust/aiconfigurator-core/src/operators/mhc.rs @@ -11,9 +11,10 @@ //! the raw latency by `scale_factor`. use serde::{Deserialize, Serialize}; -use crate::common::enums::GemmQuantMode; +use crate::common::enums::{DatabaseMode, GemmQuantMode}; use crate::common::error::AicError; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid}; use crate::perf_database::gemm::tc_flops_for_compute; use crate::perf_database::PerfDatabase; @@ -111,15 +112,71 @@ impl MhcModuleOp { sol_math.max(sol_mem) } + /// Database-mode dispatch mirroring Python `_query_mhc_table` + /// (`operations/dsv4.py`): SILICON queries the table; HYBRID converts a + /// typed silicon miss into the util-space empirical estimate; EMPIRICAL + /// always estimates. The SOL diagnostic modes never reach the compiled + /// engine (the routing gate delegates them to the Python step). pub fn query(&self, db: &PerfDatabase, num_tokens: u32) -> Result { let sol = |op_name: &str, t: f64| self.sol_ms(db, op_name, t.round() as i64); - let latency = + let silicon = || { db.mhc - .query_module(&self.op, num_tokens, self.hc_mult, self.hidden_size, &sol)?; - Ok(PerformanceResult::new(latency, Source::Silicon) + .query_module(&self.op, num_tokens, self.hc_mult, self.hidden_size, &sol) + }; + let (latency, source) = match db.database_mode { + DatabaseMode::Empirical => (self.mhc_empirical(db, num_tokens)?, Source::Empirical), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => (latency, Source::Silicon), + Err(err) if err.is_missing_perf_data() => { + (self.mhc_empirical(db, num_tokens)?, Source::Empirical) + } + Err(err) => return Err(err), + }, + _ => (silicon()?, Source::Silicon), + }; + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } + + /// Mirrors Python `_query_mhc_table::get_empirical`: for `op == "both"` + /// the empirical estimate is the SUM of the two halves' own estimates + /// (`_emp_for_op("pre") + _emp_for_op("post")`), each half calibrated on + /// its own token curve with its own SOL. + fn mhc_empirical(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + if self.op == "both" { + return Ok(self.emp_for_op(db, "pre", num_tokens)? + + self.emp_for_op(db, "post", num_tokens)?); + } + self.emp_for_op(db, &self.op, num_tokens) + } + + /// `SOL(query)/util` over one op half's own `(num_tokens,)` curve. + /// Mirrors Python `_query_mhc_table::get_empirical::_emp_for_op` (grid + /// depth 1, `sol_fn = lambda c: get_sol(c[0], op_name)[0]`). + fn emp_for_op(&self, db: &PerfDatabase, op_name: &str, num_tokens: u32) -> Result { + let sol = |c: &[f64]| self.sol_ms(db, op_name, c[0].round() as i64); + // Python keys the grid on (op_name, hc_mult, hidden_size, quant) — + // NOT sinkhorn_iters, which is mirrored deliberately. + let key = format!( + "dsv4_mhc:{op_name}:{}:{}:{}", + self.hc_mult, + self.hidden_size, + self.quant_mode.name() + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.mhc.module_points(op_name, self.hc_mult, self.hidden_size) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [f64::from(num_tokens)]; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) + } } #[cfg(test)] @@ -202,6 +259,85 @@ mod tests { } } + /// Oracle values generated from the Python reference on the same data: + /// + /// ```text + /// uv run --no-sync python3 -c " + /// from aiconfigurator.sdk import perf_database, common + /// from aiconfigurator.sdk.operations.dsv4 import DeepSeekV4MHCModule as MHC + /// db = perf_database.get_database('b200_sxm', 'sglang', '0.5.10') + /// for nt, op in [(3000,'pre'), (3000,'post'), (3000,'both'), (8,'pre'), (8,'both'), (1048576,'pre')]: + /// r = MHC._query_mhc_table(db, num_tokens=nt, hidden_size=7168, hc_mult=4, + /// sinkhorn_iters=20, op=op, quant_mode=common.GEMMQuantMode.bfloat16, + /// database_mode=common.DatabaseMode.EMPIRICAL) + /// print(nt, op, repr(float(r)))" + /// ``` + /// + /// Covers: off-grid interior IDW (nt=3000), the exact collected hit + /// reconstructing the measured value (nt=8), the `both` = emp(pre) + + /// emp(post) composition at both points, and the beyond-range clamp + /// (nt=1048576: boundary util frozen, the mHC roofline ratio carries the + /// growth). Regenerate if the shipped mHC table or the util-empirical + /// math changes. + #[test] + fn mhc_empirical_matches_python_oracles() { + let mut db = b200_sglang_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + let cases: &[(&str, u32, f64)] = &[ + ("pre", 3000, 0.28677656188924283), + ("post", 3000, 0.12520766094146765), + // "both" empirical = emp("pre") + emp("post"), each half on its + // own curve with its own SOL (Python `_emp_for_op` composition). + ("both", 3000, 0.41198422283071046), + ("pre", 8, 0.0251), + ("both", 8, 0.0357), + ("pre", 1_048_576, 71.55398179178216), + ]; + for &(op, nt, expected) in cases { + let result = mhc_op(op).query(&db, nt).expect("empirical query"); + assert!( + ((result.latency_ms - expected) / expected).abs() < 1e-9, + "op={op}, nt={nt}: rust {} vs python {expected}", + result.latency_ms + ); + assert_eq!(result.source, Source::Empirical); + } + } + + /// HYBRID with silicon data present must stay on the silicon path + /// (in-range RAW lerp, Source::Silicon) — same oracle as the silicon test. + #[test] + fn mhc_hybrid_with_data_stays_silicon() { + let mut db = b200_sglang_db(); + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let result = mhc_op("pre").query(&db, 3).expect("hybrid query"); + let expected = 0.025050000000000003; + assert!( + ((result.latency_ms - expected) / expected).abs() < 1e-9, + "rust {} vs python {expected}", + result.latency_ms + ); + assert_eq!(result.source, Source::Silicon); + } + + /// HYBRID on a slice with NO collected curve (hidden_size=1234 is not in + /// the mHC table) must surface the terminal EmpiricalNotImplemented miss, + /// never a fabricated value (mirrors Python: the silicon miss falls to + /// `get_empirical`, whose own typed miss raises + /// `EmpiricalNotImplementedError`). + #[test] + fn mhc_hybrid_missing_slice_raises_empirical_not_implemented() { + let mut db = b200_sglang_db(); + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let mut op = mhc_op("pre"); + op.hidden_size = 1234; + let result = op.query(&db, 8); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "got {result:?}" + ); + } + /// `sinkhorn_iters` / `quant_mode` are new opspec fields; old specs lack /// them and must default to (20, bfloat16). #[test] diff --git a/rust/aiconfigurator-core/src/operators/mla.rs b/rust/aiconfigurator-core/src/operators/mla.rs index 7ac4a1ac4..fab240ff1 100644 --- a/rust/aiconfigurator-core/src/operators/mla.rs +++ b/rust/aiconfigurator-core/src/operators/mla.rs @@ -6,14 +6,22 @@ //! //! Mirrors `aiconfigurator.sdk.operations.mla.{ContextMLA, GenerationMLA, //! MLAModule, MLABmm}`. Op-level paths apply Python's prefix-correction -//! multiplier inside the operator; module-level paths do the same since the -//! perf-DB layer returns raw table values. MLA BMM has a quant-mode -//! fallback to bfloat16 inside the perf-DB query. +//! multiplier inside the mode dispatch (silicon branch only — the empirical +//! branch's SOL carries prefix natively, exactly like Python's +//! `get_empirical`); module-level paths do the same since the perf-DB layer +//! returns raw table values. MLA BMM has a quant-mode fallback to bfloat16 +//! as part of slice selection (silicon inside the perf-DB query, empirical +//! before grid construction). use serde::{Deserialize, Serialize}; -use crate::common::enums::{FmhaQuantMode, GemmQuantMode, KvCacheQuantMode}; +use crate::common::enums::{DatabaseMode, FmhaQuantMode, GemmQuantMode, KvCacheQuantMode}; use crate::common::error::AicError; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid}; +use crate::perf_database::mla::{ + context_mla_sol_ms, context_mla_sol_prefix_ms, generation_mla_module_sol_ms, + generation_mla_sol_ms, mla_bmm_sol_ms, +}; use crate::perf_database::PerfDatabase; fn prefix_correction(full_s: u32, prefix: u32) -> f64 { @@ -66,28 +74,31 @@ impl ContextMlaOp { prefix: u32, ) -> Result { // ctx(s, pfx): the un-sharded context-MLA query for a sequence chunk of - // length `s` at prefix `pfx`, with the prefix correction applied. - let ctx = |s: u32, pfx: u32| -> Result { - let full_s = s + pfx; - let raw = db.mla.query_context( + // length `s` at prefix `pfx` (mode dispatch handles prefix correction + // for silicon and the prefix-aware SOL for empirical). + let ctx = |s: u32, pfx: u32| -> Result<(f64, Source), AicError> { + query_context_mla_table( + db, batch_size, - full_s, + s, + pfx, self.num_heads, self.kv_cache_dtype, self.fmha_quant_mode, - )?; - Ok(raw * prefix_correction(full_s, pfx)) + ) }; // Context parallelism (SGLang AllGather / zigzag): model rank 0's two // balanced chunks, c = ceil(isl / 2cp). Mirrors Python // `ContextMLA.query` and `operators/attention.rs::ContextAttentionOp`. - let latency = if self.cp_size > 1 { + let (latency, source) = if self.cp_size > 1 { let c = isl.div_ceil(2 * self.cp_size).max(1); - ctx(c, prefix)? + ctx(c, prefix + isl - c)? + let (first, first_source) = ctx(c, prefix)?; + let (second, second_source) = ctx(c, prefix + isl - c)?; + (first + second, first_source.combine(second_source)) } else { ctx(isl, prefix)? }; - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -117,10 +128,9 @@ impl GenerationMlaOp { batch_size: u32, s: u32, ) -> Result { - let latency = db - .mla - .query_generation(batch_size, s, self.num_heads, self.kv_cache_dtype)?; - Ok(PerformanceResult::new(latency, Source::Silicon) + let (latency, source) = + query_generation_mla_table(db, batch_size, s, self.num_heads, self.kv_cache_dtype)?; + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -163,17 +173,17 @@ impl MlaModuleOp { isl: u32, prefix: u32, ) -> Result { - let full_s = isl + prefix; - let raw = db.mla.query_context_module( + let (latency, source) = query_context_mla_module_table( + db, batch_size, - full_s, + isl, + prefix, self.num_heads, self.kv_cache_dtype, self.fmha_quant_mode, self.gemm_quant_mode, )?; - let latency = raw * prefix_correction(full_s, prefix); - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -186,14 +196,15 @@ impl MlaModuleOp { ) -> Result { // No fmha arg: the generation module table has no fmha axis (decode // compute dtype follows the kv-cache dtype). - let latency = db.mla.query_generation_module( + let (latency, source) = query_generation_mla_module_table( + db, batch_size, s, self.num_heads, self.kv_cache_dtype, self.gemm_quant_mode, )?; - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -225,15 +236,366 @@ impl MlaBmmOp { } pub fn query(&self, db: &PerfDatabase, num_tokens: u32) -> Result { - let latency = db - .mla - .query_bmm(num_tokens, self.num_heads, self.quant_mode, self.is_pre)?; - Ok(PerformanceResult::new(latency, Source::Silicon) + let (latency, source) = + query_mla_bmm_table(db, num_tokens, self.num_heads, self.quant_mode, self.is_pre)?; + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } } +// --------------------------------------------------------------------------- +// Database-mode dispatch, mirroring the Python `_query_*_table` classmethods +// (`operations/mla.py`): SILICON queries the table; HYBRID converts a typed +// silicon miss into the util-space empirical estimate; EMPIRICAL always +// estimates. The SOL diagnostic modes never reach the compiled engine (the +// routing gate delegates them to the Python step). +// --------------------------------------------------------------------------- + +/// Op-level context MLA latency (prefix correction applied on the silicon +/// branch) under the database's query mode. +fn query_context_mla_table( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, +) -> Result<(f64, Source), AicError> { + let silicon = || -> Result { + let full_s = s + prefix; + let raw = db + .mla + .query_context(b, full_s, num_heads, kv_quant, fmha_quant)?; + Ok(raw * prefix_correction(full_s, prefix)) + }; + match db.database_mode { + DatabaseMode::Empirical => Ok(( + context_mla_empirical(db, b, s, prefix, num_heads, kv_quant, fmha_quant)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + context_mla_empirical(db, b, s, prefix, num_heads, kv_quant, fmha_quant)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((silicon()?, Source::Silicon)), + } +} + +/// `SOL(query)/util` over the (fmha, kv) slice's own `(num_heads, s, b)` +/// grid. Mirrors Python `ContextMLA._query_context_mla_table::get_empirical` +/// (depth 3; samples are prefix=0, the query SOL carries prefix natively). +fn context_mla_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, +) -> Result { + let spec = &db.system_spec; + // c = (num_heads, full_s, b), prefix = 0 for collected samples. + let sol = |c: &[f64]| context_mla_sol_ms(spec, kv_quant, fmha_quant, c[0], c[1], c[2]); + let key = format!("ctx_mla:{}:{}", fmha_quant.name(), kv_quant.name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.mla.context_points(kv_quant, fmha_quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let sol_query = context_mla_sol_prefix_ms( + spec, + kv_quant, + fmha_quant, + num_heads as f64, + s as f64, + prefix as f64, + b as f64, + ); + let query = [num_heads as f64, (s + prefix) as f64, b as f64]; + let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + +/// Op-level generation MLA latency under the database's query mode. +fn query_generation_mla_table( + db: &PerfDatabase, + b: u32, + s: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + generation_mla_empirical(db, b, s, num_heads, kv_quant)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match db.mla.query_generation(b, s, num_heads, kv_quant) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + generation_mla_empirical(db, b, s, num_heads, kv_quant)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((db.mla.query_generation(b, s, num_heads, kv_quant)?, Source::Silicon)), + } +} + +/// `SOL(query)/util` over the kv slice's own `(num_heads, b, s)` grid. +/// Mirrors `GenerationMLA._query_generation_mla_table::get_empirical`. +fn generation_mla_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, +) -> Result { + let spec = &db.system_spec; + // c = (num_heads, b, s). + let sol = |c: &[f64]| generation_mla_sol_ms(spec, kv_quant, c[0], c[1], c[2]); + let key = format!("gen_mla:{}", kv_quant.name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.mla.generation_points(kv_quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [num_heads as f64, b as f64, s as f64]; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + +/// MLA BMM (pre/post) latency under the database's query mode. +fn query_mla_bmm_table( + db: &PerfDatabase, + num_tokens: u32, + num_heads: u32, + quant: GemmQuantMode, + is_pre: bool, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + mla_bmm_empirical(db, num_tokens, num_heads, quant, is_pre)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match db.mla.query_bmm(num_tokens, num_heads, quant, is_pre) { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + mla_bmm_empirical(db, num_tokens, num_heads, quant, is_pre)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((db.mla.query_bmm(num_tokens, num_heads, quant, is_pre)?, Source::Silicon)), + } +} + +/// `SOL(query)/util` over the 1-D `num_tokens` curve of the selected +/// `(quant, op_name, num_heads)` slice. Mirrors +/// `MLABmm._query_mla_bmm_table::get_empirical`: slice selection falls back +/// to the bfloat16 quant when the requested quant has no BMM data at all +/// (BEFORE estimate()); the SOL keeps using the REQUESTED quant either way. +fn mla_bmm_empirical( + db: &PerfDatabase, + num_tokens: u32, + num_heads: u32, + quant: GemmQuantMode, + is_pre: bool, +) -> Result { + let spec = &db.system_spec; + // c = (num_tokens,); the SOL is bound to the REQUESTED quant. + let sol = |c: &[f64]| mla_bmm_sol_ms(spec, quant, num_heads as f64, c[0]); + let op_name = if is_pre { "mla_gen_pre" } else { "mla_gen_post" }; + // Slice selection first (Python: `qm = quant if quant in wrapper else + // bfloat16`); a typed miss here means the whole BMM table is absent. + let grid = match db.mla.bmm_selected_quant(quant) { + Ok(selected) => { + // The key carries the REQUESTED and the ACTUALLY-selected quant: + // on the bfloat16 fallback the samples still get the requested + // quant's SOL, so grids from the same slice under different + // requested quants must not alias (Python keys on the requested + // quant plus the concrete node identity). + let key = format!( + "mla_bmm:{}:{}:{}:{}", + quant.name(), + selected.name(), + op_name, + num_heads + ); + db.util_grids.get_or_try_build(&key, || { + match db.mla.bmm_points(selected, is_pre, num_heads) { + Ok(points) => { + Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))) + } + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })? + } + Err(err) if err.is_missing_perf_data() => None, + Err(err) => return Err(err), + }; + let query = [num_tokens as f64]; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + +/// Module-level context MLA latency (prefix correction on the silicon +/// branch) under the database's query mode. +#[allow(clippy::too_many_arguments)] +fn query_context_mla_module_table( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + gemm_quant: GemmQuantMode, +) -> Result<(f64, Source), AicError> { + let silicon = || -> Result { + let full_s = s + prefix; + let raw = db + .mla + .query_context_module(b, full_s, num_heads, kv_quant, fmha_quant, gemm_quant)?; + Ok(raw * prefix_correction(full_s, prefix)) + }; + match db.database_mode { + DatabaseMode::Empirical => Ok(( + context_mla_module_empirical(db, b, s, prefix, num_heads, kv_quant, fmha_quant, gemm_quant)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + context_mla_module_empirical( + db, b, s, prefix, num_heads, kv_quant, fmha_quant, gemm_quant, + )?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((silicon()?, Source::Silicon)), + } +} + +/// Mirrors `MLAModule._query_context_mla_module_table::get_empirical`: same +/// SOL as the op-level context MLA (the gemm quant only selects the slice), +/// over the module's own `(num_heads, s, b)` grid. +#[allow(clippy::too_many_arguments)] +fn context_mla_module_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + gemm_quant: GemmQuantMode, +) -> Result { + let spec = &db.system_spec; + // c = (num_heads, full_s, b), prefix = 0 for collected samples. + let sol = |c: &[f64]| context_mla_sol_ms(spec, kv_quant, fmha_quant, c[0], c[1], c[2]); + let key = format!( + "ctx_mla_mod:{}:{}:{}", + fmha_quant.name(), + kv_quant.name(), + gemm_quant.name() + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.mla.context_module_points(kv_quant, fmha_quant, gemm_quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let sol_query = context_mla_sol_prefix_ms( + spec, + kv_quant, + fmha_quant, + num_heads as f64, + s as f64, + prefix as f64, + b as f64, + ); + let query = [num_heads as f64, (s + prefix) as f64, b as f64]; + let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + +/// Module-level generation MLA latency under the database's query mode. +fn query_generation_mla_module_table( + db: &PerfDatabase, + b: u32, + s: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + gemm_quant: GemmQuantMode, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + generation_mla_module_empirical(db, b, s, num_heads, kv_quant, gemm_quant)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => { + match db + .mla + .query_generation_module(b, s, num_heads, kv_quant, gemm_quant) + { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + generation_mla_module_empirical(db, b, s, num_heads, kv_quant, gemm_quant)?, + Source::Empirical, + )), + Err(err) => Err(err), + } + } + _ => Ok(( + db.mla + .query_generation_module(b, s, num_heads, kv_quant, gemm_quant)?, + Source::Silicon, + )), + } +} + +/// Mirrors `MLAModule._query_generation_mla_module_table::get_empirical`: +/// generation MLA SOL + BMM pre/post terms (the module SOL closes over the +/// gemm quant), over the module's own `(num_heads, b, s)` grid. +fn generation_mla_module_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + gemm_quant: GemmQuantMode, +) -> Result { + let spec = &db.system_spec; + // c = (num_heads, b, s). + let sol = |c: &[f64]| generation_mla_module_sol_ms(spec, kv_quant, gemm_quant, c[0], c[1], c[2]); + let key = format!("gen_mla_mod:{}:{}", kv_quant.name(), gemm_quant.name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.mla.generation_module_points(kv_quant, gemm_quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [num_heads as f64, b as f64, s as f64]; + let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + #[cfg(test)] mod tests { use super::*; @@ -283,4 +645,230 @@ mod tests { other => panic!("unexpected error: {other:?}"), } } + + fn gb200_trtllm_db() -> PerfDatabase { + let systems_root = PathBuf::from(REPO_ROOT_HINT) + .join("../..") + .join("src/aiconfigurator/systems"); + PerfDatabase::load(&systems_root, "gb200", "trtllm", "1.3.0rc10").expect("db must load") + } + + fn assert_close(got: f64, expected: f64, what: &str) { + assert!( + (got - expected).abs() < 1e-9, + "{what}: expected {expected}, got {got}" + ); + } + + /// Oracle values generated from the Python reference on the same data: + /// `ContextMLA._query_context_mla_table(db, b, s, prefix, num_heads, + /// kv, fmha, database_mode=EMPIRICAL)` on gb200/trtllm/1.3.0rc10 + /// (`get_database(..., shared_layer=False)`, matching this single-primary + /// loader). Regenerate if the shipped table or the util math changes. + #[test] + fn context_mla_empirical_matches_python_oracles() { + let mut db = gb200_trtllm_db(); + db.database_mode = DatabaseMode::Empirical; + let cases: &[(u32, u32, u32, u32, f64)] = &[ + // off-grid seq + (4, 5000, 0, 128, 3.5591050930761825), + // prefix > 0: the query SOL carries prefix natively + (2, 3000, 1024, 16, 0.17932261401789712), + // exact collected hit: util reconstruction returns the measured value + (4, 4096, 0, 128, 2.4523092905680337), + ]; + for &(b, s, prefix, n, expected) in cases { + let (latency, source) = query_context_mla_table( + &db, + b, + s, + prefix, + n, + KvCacheQuantMode::Bfloat16, + FmhaQuantMode::Bfloat16, + ) + .expect("empirical query"); + assert_close(latency, expected, &format!("ctx_mla(b={b}, s={s}, pfx={prefix}, n={n})")); + assert_eq!(source, Source::Empirical); + } + } + + /// HYBRID on a slice with NO collected data (fmha=fp8 on gb200/trtllm, + /// whose context MLA table is bfloat16-only) must surface the terminal + /// EmpiricalNotImplemented miss (mirrors the Python contract). + #[test] + fn context_mla_hybrid_missing_slice_raises_empirical_not_implemented() { + let mut db = gb200_trtllm_db(); + db.database_mode = DatabaseMode::Hybrid; + let result = query_context_mla_table( + &db, + 4, + 4096, + 0, + 128, + KvCacheQuantMode::Bfloat16, + FmhaQuantMode::Fp8, + ); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } + + /// Python oracle: `GenerationMLA._query_generation_mla_table` in + /// EMPIRICAL mode on gb200/trtllm/1.3.0rc10 (shared_layer=False). + #[test] + fn generation_mla_empirical_matches_python_oracles() { + let mut db = gb200_trtllm_db(); + db.database_mode = DatabaseMode::Empirical; + let cases: &[(u32, u32, u32, f64)] = &[ + // off-grid (b, s) + (7, 9000, 128, 0.02734810076798607), + // exact collected hit + (1, 4096, 128, 0.02057066683967908), + ]; + for &(b, s, n, expected) in cases { + let (latency, source) = + query_generation_mla_table(&db, b, s, n, KvCacheQuantMode::Bfloat16) + .expect("empirical query"); + assert_close(latency, expected, &format!("gen_mla(b={b}, s={s}, n={n})")); + assert_eq!(source, Source::Empirical); + } + + // HYBRID with a kv dtype that has no table (int8) -> terminal miss. + db.database_mode = DatabaseMode::Hybrid; + let result = query_generation_mla_table(&db, 1, 4096, 128, KvCacheQuantMode::Int8); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } + + /// Python oracle: `MLABmm._query_mla_bmm_table` in EMPIRICAL mode on + /// gb200/trtllm/1.3.0rc10 (shared_layer=False). The fp8 cases exercise + /// the bfloat16 slice fallback (gb200's BMM table is bfloat16-only) with + /// the SOL still bound to the REQUESTED fp8 quant. + #[test] + fn mla_bmm_empirical_matches_python_oracles() { + let mut db = gb200_trtllm_db(); + db.database_mode = DatabaseMode::Empirical; + let cases: &[(u32, u32, GemmQuantMode, bool, f64)] = &[ + // off-grid tokens on the requested bf16 slice + (100, 128, GemmQuantMode::Bfloat16, true, 0.008883413307229573), + // exact collected hit + (256, 128, GemmQuantMode::Bfloat16, true, 0.010847999900579452), + // fp8 requested -> bfloat16 fallback slice, fp8 SOL + (20000, 128, GemmQuantMode::Fp8, true, 0.5326748099591996), + // fallback on the post BMM at another head count + (777, 64, GemmQuantMode::Fp8, false, 0.010838556565365292), + ]; + for &(t, n, quant, is_pre, expected) in cases { + let (latency, source) = + query_mla_bmm_table(&db, t, n, quant, is_pre).expect("empirical query"); + assert_close( + latency, + expected, + &format!("mla_bmm(t={t}, n={n}, {quant:?}, pre={is_pre})"), + ); + assert_eq!(source, Source::Empirical); + } + + // HYBRID at a head count absent from both the requested and the + // bfloat16 fallback slice -> terminal miss. + db.database_mode = DatabaseMode::Hybrid; + let result = query_mla_bmm_table(&db, 64, 7, GemmQuantMode::Bfloat16, true); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } + + /// Python oracle: `MLAModule._query_context_mla_module_table` in + /// EMPIRICAL mode on b200_sxm/vllm/0.19.0 (shared_layer=False). + #[test] + fn context_mla_module_empirical_matches_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = DatabaseMode::Empirical; + type Case = (u32, u32, u32, u32, FmhaQuantMode, KvCacheQuantMode, GemmQuantMode, f64); + let cases: &[Case] = &[ + // off-grid seq, bf16^3 slice + ( + 2, 5000, 0, 128, + FmhaQuantMode::Bfloat16, KvCacheQuantMode::Bfloat16, GemmQuantMode::Bfloat16, + 5.030970266382403, + ), + // prefix > 0 + ( + 1, 2000, 2048, 16, + FmhaQuantMode::Bfloat16, KvCacheQuantMode::Bfloat16, GemmQuantMode::Bfloat16, + 0.3328645307516498, + ), + // exact collected hit + ( + 1, 1, 0, 128, + FmhaQuantMode::Bfloat16, KvCacheQuantMode::Bfloat16, GemmQuantMode::Bfloat16, + 0.1351, + ), + // fp8 fmha/kv with fp8_block gemm slice + ( + 2, 5000, 0, 128, + FmhaQuantMode::Fp8, KvCacheQuantMode::Fp8, GemmQuantMode::Fp8Block, + 4.68971474347924, + ), + ]; + for &(b, s, prefix, n, fmha, kv, gemm, expected) in cases { + let (latency, source) = + query_context_mla_module_table(&db, b, s, prefix, n, kv, fmha, gemm) + .expect("empirical query"); + assert_close( + latency, + expected, + &format!("ctx_mla_mod(b={b}, s={s}, pfx={prefix}, n={n}, {fmha:?}, {kv:?}, {gemm:?})"), + ); + assert_eq!(source, Source::Empirical); + } + + // HYBRID with a gemm quant slice that has no data (fp8) -> miss. + db.database_mode = DatabaseMode::Hybrid; + let result = query_context_mla_module_table( + &db, + 2, + 5000, + 0, + 128, + KvCacheQuantMode::Bfloat16, + FmhaQuantMode::Bfloat16, + GemmQuantMode::Fp8, + ); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } + + /// Python oracle: `MLAModule._query_generation_mla_module_table` in + /// EMPIRICAL mode on b200_sxm/vllm/0.19.0 (shared_layer=False). The + /// fp8/fp8_block case exercises the module SOL's dependence on the gemm + /// quant (the BMM terms close over it). + #[test] + fn generation_mla_module_empirical_matches_python_oracles() { + let mut db = b200_vllm_db(); + db.database_mode = DatabaseMode::Empirical; + let cases: &[(u32, u32, u32, KvCacheQuantMode, GemmQuantMode, f64)] = &[ + (8, 3000, 128, KvCacheQuantMode::Bfloat16, GemmQuantMode::Bfloat16, 0.16175364801657854), + // exact collected hit (s = isl + step) + (1, 4097, 128, KvCacheQuantMode::Bfloat16, GemmQuantMode::Bfloat16, 0.1497), + (8, 3000, 16, KvCacheQuantMode::Fp8, GemmQuantMode::Fp8Block, 0.1146692546817411), + ]; + for &(b, s, n, kv, gemm, expected) in cases { + let (latency, source) = + query_generation_mla_module_table(&db, b, s, n, kv, gemm).expect("empirical query"); + assert_close( + latency, + expected, + &format!("gen_mla_mod(b={b}, s={s}, n={n}, {kv:?}, {gemm:?})"), + ); + assert_eq!(source, Source::Empirical); + } + + // HYBRID with a gemm quant slice that has no data (fp8) -> miss. + db.database_mode = DatabaseMode::Hybrid; + let result = query_generation_mla_module_table( + &db, + 8, + 3000, + 128, + KvCacheQuantMode::Bfloat16, + GemmQuantMode::Fp8, + ); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } } diff --git a/rust/aiconfigurator-core/src/operators/moe.rs b/rust/aiconfigurator-core/src/operators/moe.rs index b1c05674f..f9defa7dc 100644 --- a/rust/aiconfigurator-core/src/operators/moe.rs +++ b/rust/aiconfigurator-core/src/operators/moe.rs @@ -3,21 +3,140 @@ //! MoE operator. //! -//! Mirrors `aiconfigurator.sdk.operations.moe.MoE` SILICON path. The perf-DB -//! layer handles workload-distribution fallback to `"uniform"` and resolves -//! the token curve on the perf_interp v2 engine; this operator supplies the -//! MoE roofline SOL closure the engine's beyond-range util-hold anchors on -//! (Python v2 deleted the op-level overflow estimator — the engine's -//! `k_tail=1`, unclamped util-hold replaces it). +//! Mirrors `aiconfigurator.sdk.operations.moe.MoE._query_moe_table`. The +//! perf-DB layer handles workload-distribution fallback to `"uniform"` and +//! resolves the token curve on the perf_interp v2 engine; this operator +//! supplies the MoE roofline SOL closure the engine's beyond-range util-hold +//! anchors on (Python v2 deleted the op-level overflow estimator — the +//! engine's `k_tail=1`, unclamped util-hold replaces it). +//! +//! Database-mode dispatch follows the gemm.rs reference pattern: EMPIRICAL +//! always estimates `SOL(query)/util`; HYBRID queries silicon and converts a +//! typed missing-data error into the estimate; SILICON is unchanged. When +//! the op's own `(quant, shape)` slice has no collected data the empirical +//! path walks the transfer ladder (`operations/moe.py:446-570`): +//! +//! 1. **xshape** — nearest collected shape within the query quant; +//! 2. **xquant** — sibling quant with the SAME (memory, compute) profile, +//! util reconstructed with the QUERY quant's SOL; +//! 3. **xprofile** — nearest-profile collected quant, util reconstructed +//! with the REFERENCE quant's own SOL, rescaled by the per-quant +//! util-LEVEL ratio `e(query)/e(ref)`. +//! +//! Policy-disabled tiers are skipped (not an error); the terminal +//! [`AicError::EmpiricalNotImplemented`] only surfaces when every permitted +//! tier found nothing. +//! +//! Scope: the SGLang `moe_backend == "deepep_moe"` branch of Python's +//! `_moe_table` does not exist here — the Rust engine routes WideEP MoE +//! through `WideEpMoeOp` (`operators/wideep_moe.rs`), so this op only ever +//! addresses the regular / low-latency tables. //! //! Weights accounting (per-expert FFN weights + router) is in the model //! layer; the operator returns latency only. use serde::{Deserialize, Serialize}; -use crate::common::enums::MoeQuantMode; +use crate::common::enums::{DatabaseMode, MoeQuantMode, TransferKind, TransferPolicy}; use crate::common::error::AicError; +use crate::common::system_spec::SystemSpec; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid}; +use crate::perf_database::moe::{MoeKernel, MoeSiblingSlice}; use crate::perf_database::PerfDatabase; +use std::sync::Arc; + +/// Per-quant achieved-util LEVEL `e(q)` for MoE, keyed by the +/// `(memory, compute)` profile. Mirrors `_MOE_QUANT_UTIL_LEVEL` +/// (`operations/moe.py:87-97`); consumed ONLY by the cross-profile tier, +/// and only as the ratio `e(query)/e(ref)`. +const MOE_QUANT_UTIL_LEVEL: &[(f64, f64, f64)] = &[ + (2.0, 1.0, 0.53), // w16a16 / bfloat16 [data] + (1.0, 1.0, 0.45), // w8a16 [inferred] + (0.5, 1.0, 0.07), // w4a16 (int4_wo, mxfp4) [data] + (1.0, 2.0, 0.40), // w8a8 / fp8(_block) [data] + (0.5, 2.0, 0.15), // w4a8 (w4afp8, mxfp4_mxfp8) [data] + (1.0, 4.0, 0.30), // w8a4 [inferred] + (0.5, 4.0, 0.23), // w4a4 [data ≈ nvfp4] + (0.5625, 4.0, 0.23), // w4a4 / nvfp4 [data] +]; +/// Unlisted profile: mid-range relative level (Python `_MOE_QUANT_UTIL_DEFAULT`). +const MOE_QUANT_UTIL_DEFAULT: f64 = 0.30; + +/// Achieved-util level `e(q)` for a MoE quant, by `(memory, compute)` +/// profile (mirrors `_moe_quant_util_level`, `operations/moe.py:100-102`). +fn moe_quant_util_level(quant: MoeQuantMode) -> f64 { + let mapping = quant.mapping(); + MOE_QUANT_UTIL_LEVEL + .iter() + .find(|(memory, compute, _)| *memory == mapping.memory && *compute == mapping.compute) + .map(|(_, _, level)| *level) + .unwrap_or(MOE_QUANT_UTIL_DEFAULT) +} + +/// Every MoE quant variant, for parsing perf-table `moe_dtype` strings back +/// into the enum (Python's table is keyed by enum members directly). +const ALL_MOE_QUANTS: &[MoeQuantMode] = &[ + MoeQuantMode::Bfloat16, + MoeQuantMode::Fp8, + MoeQuantMode::Int4Wo, + MoeQuantMode::Fp8Block, + MoeQuantMode::W4afp8, + MoeQuantMode::Nvfp4, + MoeQuantMode::W4a16Mxfp4, + MoeQuantMode::W4a8Mxfp4Mxfp8, + MoeQuantMode::W4a8Mxfp4Mxfp8Trtllm, + MoeQuantMode::W4a16Mxfp4Cutlass, +]; + +fn moe_quant_from_name(name: &str) -> Option { + ALL_MOE_QUANTS.iter().copied().find(|q| q.name() == name) +} + +/// Collected quants with a DIFFERENT `(memory, compute)` profile than the +/// query, nearest-profile first (stable sort by `|Δmemory| + |Δcompute|`; +/// mirrors `_xprofile_moe_quants`, `operations/moe.py:105-117`). NOTE: +/// Python breaks distance ties by table insertion (file row) order; here +/// the input arrives in the accessor's sorted-name order. +fn xprofile_moe_quants(query: MoeQuantMode, table_quants: &[MoeQuantMode]) -> Vec { + let qp = query.mapping(); + let mut refs: Vec = table_quants + .iter() + .copied() + .filter(|q| { + let m = q.mapping(); + *q != query && !(m.memory == qp.memory && m.compute == qp.compute) + }) + .collect(); + let dist = |q: MoeQuantMode| { + let m = q.mapping(); + (m.memory - qp.memory).abs() + (m.compute - qp.compute).abs() + }; + refs.sort_by(|a, b| dist(*a).partial_cmp(&dist(*b)).expect("finite profile distances")); + refs +} + +/// Enabled-tier fingerprint folded into reference-grid cache keys so grids +/// selected under different policies cannot alias (Python's +/// `selection_key`/`identity_key` include the policy frozenset). +fn policy_fingerprint(policy: TransferPolicy) -> String { + format!( + "xshape={},xquant={},xprofile={},xop={}", + policy.xshape as u8, policy.xquant as u8, policy.xprofile as u8, policy.xop as u8 + ) +} + +/// A sibling slice the transfer ladder may borrow: the reference slice's +/// shape + token curve, the quant whose SOL reconstructs its util +/// (`sol_quant`: QUERY quant for same-profile tiers, REFERENCE quant for +/// cross-profile), and the transfer-tier provenance tag. Mirrors +/// `util_empirical.ReferenceCandidate` as built by `_collect` +/// (`operations/moe.py:454-486`). +struct MoeReferenceCandidate { + slice: MoeSiblingSlice, + ref_quant: MoeQuantMode, + sol_quant: MoeQuantMode, + provenance: &'static str, +} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct MoeOp { @@ -81,6 +200,30 @@ impl MoeOp { // `MoE.query` (`x = x * attention_dp_size`). Applied exactly once, // before the perf-DB resolution keys off the token count. let num_tokens = num_tokens.saturating_mul(self.attention_dp_size.max(1)); + + // Database-mode dispatch, mirroring the Python `_query_moe_table` + // tail (`database._query_silicon_or_hybrid`): EMPIRICAL always + // estimates; HYBRID converts a typed silicon miss into the estimate; + // SILICON is unchanged. The SOL diagnostic modes never reach the + // compiled engine. + let (latency, source) = match db.database_mode { + DatabaseMode::Empirical => (self.empirical_latency(db, num_tokens)?, Source::Empirical), + DatabaseMode::Hybrid => match self.silicon_latency(db, num_tokens) { + Ok(latency) => (latency, Source::Silicon), + Err(err) if err.is_missing_perf_data() => { + (self.empirical_latency(db, num_tokens)?, Source::Empirical) + } + Err(err) => return Err(err), + }, + _ => (self.silicon_latency(db, num_tokens)?, Source::Silicon), + }; + Ok(PerformanceResult::new(latency, source) + .clamp_non_negative() + .scaled(self.scale_factor)) + } + + /// SILICON table resolution (the pre-empirical behaviour, unchanged). + fn silicon_latency(&self, db: &PerfDatabase, num_tokens: u32) -> Result { // The roofline SOL the perf-DB engine anchors its beyond-range // util-hold on (Python `_resolve_tokens` passes the same closure). // Coordinates arriving from the engine are always integral (table @@ -96,7 +239,7 @@ impl MoeOp { // default grid on a shape miss. Other backends (vLLM, SGLang) never // have `kernel_source` populated, so `low_latency_available()` // returns false and this short-circuits. - let latency = if num_tokens <= 128 + if num_tokens <= 128 && self.quant_mode == MoeQuantMode::Nvfp4 && self.is_gated && db.moe.low_latency_available()? @@ -113,38 +256,319 @@ impl MoeOp { &self.workload_distribution, &sol, )? { - ll - } else { - db.moe.query( - num_tokens, - self.hidden_size, - self.inter_size, - self.topk, - self.num_experts, - self.moe_tp_size, - self.moe_ep_size, - self.quant_mode, - &self.workload_distribution, - &sol, - )? + return Ok(ll); + } + } + db.moe.query( + num_tokens, + self.hidden_size, + self.inter_size, + self.topk, + self.num_experts, + self.moe_tp_size, + self.moe_ep_size, + self.quant_mode, + &self.workload_distribution, + &sol, + ) + } + + /// `SOL(query)/util` with the full transfer ladder. Mirrors Python + /// `MoE._query_moe_table::get_empirical` (`operations/moe.py:327-572`): + /// own-shape grid → xshape → xquant → xprofile → typed empirical miss. + fn empirical_latency(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + let spec = &db.system_spec; + let quant = self.quant_mode; + let num_gemms: u64 = if self.is_gated { 3 } else { 2 }; + let sol_time = moe_sol_latency_ms( + spec, + quant, + num_gemms, + num_tokens, + self.hidden_size, + self.inter_size, + self.topk, + self.num_experts, + self.moe_tp_size, + self.moe_ep_size, + ); + + // Kernel-table selection mirrors get_silicon's (`_moe_table`, + // `operations/moe.py:355-386`): nvfp4 + small tokens + gated probes + // the low-latency table for the FULL slice and falls back to the + // default table on a shape miss. Building util from the wrong table + // would over-estimate by the ~3x kernel gap. The `kernel_tag` folds + // the choice into every grid cache key so a low-latency grid can't + // be served to a regular query (or vice versa) at the same shape. + let kernel = if num_tokens <= 128 + && quant == MoeQuantMode::Nvfp4 + && self.is_gated + && db.moe.low_latency_available()? + { + match self.slice_points(db, MoeKernel::LowLatency) { + Ok(_) => MoeKernel::LowLatency, + Err(err) if err.is_missing_perf_data() => MoeKernel::Standard, + Err(err) => return Err(err), } } else { - db.moe.query( - num_tokens, + MoeKernel::Standard + }; + let kernel_tag = match kernel { + MoeKernel::LowLatency => "ll", + MoeKernel::Standard => "std", + }; + + // Own-shape grid over this slice's num_tokens curve (depth 1). + let own_sol = |c: &[f64]| { + moe_sol_latency_ms( + spec, + quant, + num_gemms, + c[0].round() as u32, self.hidden_size, self.inter_size, self.topk, self.num_experts, self.moe_tp_size, self.moe_ep_size, - self.quant_mode, - &self.workload_distribution, - &sol, - )? + ) }; - Ok(PerformanceResult::new(latency, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) + let own_key = format!( + "moe:{}:{}:{}:{}:{}:{}:{}:{}:{}:{}", + quant.name(), + kernel_tag, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + self.workload_distribution, + num_gemms, + ); + let mut grid = db.util_grids.get_or_try_build(&own_key, || { + match self.slice_points(db, kernel) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples( + points.into_iter().map(|(t, lat)| (vec![t as f64], lat)), + own_sol, + )))), + // Typed coverage miss -> no grid (the ladder below, then + // estimate(), takes over); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + + let mut util_scale = 1.0; + if grid.as_deref().is_none_or(UtilGrid::is_empty) { + let policy = db.transfer_policy; + + // Tiers 1+2 flow through ONE reference selection (`_moe_candidates` + // + a single grid_from_reference, `operations/moe.py:490-532`): + // xshape candidates win outright when any exist; only an empty + // xshape set falls through to same-profile xquant siblings. + let mut candidates: Vec = Vec::new(); + if policy.contains(TransferKind::XShape) { + self.collect_candidates(db, kernel, quant, quant, "xshape", &mut candidates)?; + } + if candidates.is_empty() && policy.contains(TransferKind::XQuant) { + let qp = quant.mapping(); + for name in db.moe.available_quants(kernel)? { + let Some(sibling) = moe_quant_from_name(&name) else { + continue; + }; + let mapping = sibling.mapping(); + if sibling == quant + || mapping.memory != qp.memory + || mapping.compute != qp.compute + { + continue; + } + self.collect_candidates(db, kernel, sibling, quant, "xquant", &mut candidates)?; + } + } + if let Some(reference) = + self.reference_grid(db, kernel_tag, num_gemms, policy, &candidates)? + { + grid = Some(reference); + } + + // Tier 3: cross-PROFILE. No own- or same-profile data at all -> + // borrow the nearest collected quant's util curve, built with the + // REFERENCE quant's own SOL, and rescale by the per-quant + // util-LEVEL ratio e(query)/e(ref) (`operations/moe.py:541-570`). + if grid.as_deref().is_none_or(UtilGrid::is_empty) + && policy.contains(TransferKind::XProfile) + { + let table_quants: Vec = db + .moe + .available_quants(kernel)? + .iter() + .filter_map(|name| moe_quant_from_name(name)) + .collect(); + for ref_quant in xprofile_moe_quants(quant, &table_quants) { + let mut candidates: Vec = Vec::new(); + self.collect_candidates( + db, + kernel, + ref_quant, + ref_quant, + "xprofile", + &mut candidates, + )?; + if let Some(reference) = + self.reference_grid(db, kernel_tag, num_gemms, policy, &candidates)? + { + if !reference.is_empty() { + grid = Some(reference); + util_scale = + moe_quant_util_level(quant) / moe_quant_util_level(ref_quant); + break; + } + } + } + } + } + + let query = [num_tokens as f64]; + let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), util_scale)?; + Ok(latency) + } + + /// This op's own-slice token curve on the chosen kernel table. + fn slice_points(&self, db: &PerfDatabase, kernel: MoeKernel) -> Result, AicError> { + db.moe.slice_points( + kernel, + self.quant_mode.name(), + &self.workload_distribution, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + ) + } + + /// Enumerate `source_quant`'s collected sibling slices (same table, + /// same wl-after-fallback / moe_tp / moe_ep) as ladder candidates. + /// Mirrors `_collect` (`operations/moe.py:454-486`); a typed data miss + /// (table failed to load) yields no candidates, exactly like Python's + /// `grid_from_reference` catching the raise from `_collect`. + fn collect_candidates( + &self, + db: &PerfDatabase, + kernel: MoeKernel, + source_quant: MoeQuantMode, + sol_quant: MoeQuantMode, + provenance: &'static str, + out: &mut Vec, + ) -> Result<(), AicError> { + let slices = match db.moe.sibling_slices( + kernel, + source_quant.name(), + &self.workload_distribution, + self.moe_tp_size, + self.moe_ep_size, + ) { + Ok(slices) => slices, + Err(err) if err.is_missing_perf_data() => return Ok(()), + Err(err) => return Err(err), + }; + out.extend(slices.into_iter().map(|slice| MoeReferenceCandidate { + slice, + ref_quant: source_quant, + sol_quant, + provenance, + })); + Ok(()) + } + + /// Nearest-candidate selection + reference-grid build, mirroring + /// `util_empirical.grid_from_reference`: the candidate nearest to the + /// query's `(topk, num_experts, hidden, inter)` in normalised log space + /// wins (first-wins on ties), and its grid is built with the CANDIDATE's + /// own shape bound into the SOL. `None` when there are no candidates. + /// The cache key carries the op identity, the query slice, the + /// enabled-tier fingerprint, the selected reference identity, and the + /// provenance so differently-policied lookups cannot alias. + fn reference_grid( + &self, + db: &PerfDatabase, + kernel_tag: &str, + num_gemms: u64, + policy: TransferPolicy, + candidates: &[MoeReferenceCandidate], + ) -> Result>, AicError> { + if candidates.is_empty() { + return Ok(None); + } + let query_features = [ + self.topk as f64, + self.num_experts as f64, + self.hidden_size as f64, + self.inter_size as f64, + ]; + let feature_rows: Vec> = candidates + .iter() + .map(|c| { + vec![ + c.slice.topk as f64, + c.slice.num_experts as f64, + c.slice.hidden_size as f64, + c.slice.inter_size as f64, + ] + }) + .collect(); + let chosen = &candidates[util_empirical::nearest_candidate_index(&query_features, &feature_rows) + .expect("candidate list is non-empty")]; + + let spec = &db.system_spec; + let key = format!( + "moe_{}:{}:{}:{}:{}:{}:{}:{}:{}:{}:{}:policy={}:ref={}:{}x{}x{}x{}", + chosen.provenance, + self.quant_mode.name(), + kernel_tag, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + self.workload_distribution, + num_gemms, + policy_fingerprint(policy), + chosen.ref_quant.name(), + chosen.slice.topk, + chosen.slice.num_experts, + chosen.slice.hidden_size, + chosen.slice.inter_size, + ); + db.util_grids.get_or_try_build(&key, || { + // ReferenceCandidate contract: the SOL uses THE CANDIDATE's + // shape (not the query's) so util carries only the shared + // kernel efficiency. + let sol = |c: &[f64]| { + moe_sol_latency_ms( + spec, + chosen.sol_quant, + num_gemms, + c[0].round() as u32, + chosen.slice.hidden_size, + chosen.slice.inter_size, + chosen.slice.topk, + chosen.slice.num_experts, + self.moe_tp_size, + self.moe_ep_size, + ) + }; + let mut grid = UtilGrid::new(util_empirical::build_samples( + chosen.slice.points.iter().map(|&(t, lat)| (vec![t as f64], lat)), + sol, + )); + grid.reference_provenance = Some(chosen.provenance); + Ok(Some(grid)) + }) } /// SOL MoE latency (ms) mirroring Python `MoE._query_moe_table`'s @@ -156,32 +580,63 @@ impl MoeOp { // non-gated Relu² (up + down). Matches Python `num_gemms = 3 if // is_gated else 2` (`operations/moe.py:115, 239`). let num_gemms: u64 = if self.is_gated { 3 } else { 2 }; - let total_tokens = num_tokens as u64 * self.topk as u64; - let moe_ep = (self.moe_ep_size as u64).max(1); - let moe_tp = (self.moe_tp_size as u64).max(1); - let h = self.hidden_size as u64; - let inter = self.inter_size as u64; - let ne = self.num_experts as u64; - - let ops = total_tokens * h * inter * num_gemms * 2 / moe_ep / moe_tp; - let mem_bytes_int = total_tokens / moe_ep * h * 2 // input + output - + total_tokens / moe_ep * inter * num_gemms / moe_tp // intermediate - + h * inter * num_gemms / moe_tp - * std::cmp::min(ne / moe_ep, total_tokens / moe_ep); - let mem_bytes = (mem_bytes_int as f64) * self.quant_mode.mapping().memory; - - let spec = &db.system_spec; - // Python uses `system_spec["gpu"]["bfloat16_tc_flops"]` directly - // (KeyError if missing). Rust exposes it as Option; fall back to 1.0 - // to make the math identity (sol_math → ops, sol_mem dominates) - // rather than dividing by zero. Every shipped system populates it. - let tc_flops = spec.gpu.bfloat16_tc_flops.unwrap_or(1.0); - let sol_math = (ops as f64) / (tc_flops * self.quant_mode.mapping().compute) * 1000.0; - let sol_mem = mem_bytes / spec.gpu.mem_bw * 1000.0; - sol_math.max(sol_mem) + moe_sol_latency_ms( + &db.system_spec, + self.quant_mode, + num_gemms, + num_tokens, + self.hidden_size, + self.inter_size, + self.topk, + self.num_experts, + self.moe_tp_size, + self.moe_ep_size, + ) } } +/// MoE roofline SOL (ms) mirroring Python `MoE._query_moe_table.get_sol` +/// (`operations/moe.py:297-325`), parameterised over the slice's shape and +/// quant so the transfer ladder can bind it to a REFERENCE candidate's shape +/// (`num_experts` folds into the min() weight term; `workload_distribution` +/// never enters the math). +#[allow(clippy::too_many_arguments)] +fn moe_sol_latency_ms( + spec: &SystemSpec, + quant: MoeQuantMode, + num_gemms: u64, + num_tokens: u32, + hidden_size: u32, + inter_size: u32, + topk: u32, + num_experts: u32, + moe_tp_size: u32, + moe_ep_size: u32, +) -> f64 { + let total_tokens = num_tokens as u64 * topk as u64; + let moe_ep = (moe_ep_size as u64).max(1); + let moe_tp = (moe_tp_size as u64).max(1); + let h = hidden_size as u64; + let inter = inter_size as u64; + let ne = num_experts as u64; + + let ops = total_tokens * h * inter * num_gemms * 2 / moe_ep / moe_tp; + let mem_bytes_int = total_tokens / moe_ep * h * 2 // input + output + + total_tokens / moe_ep * inter * num_gemms / moe_tp // intermediate + + h * inter * num_gemms / moe_tp + * std::cmp::min(ne / moe_ep, total_tokens / moe_ep); + let mem_bytes = (mem_bytes_int as f64) * quant.mapping().memory; + + // Python uses `system_spec["gpu"]["bfloat16_tc_flops"]` directly + // (KeyError if missing). Rust exposes it as Option; fall back to 1.0 + // to make the math identity (sol_math → ops, sol_mem dominates) + // rather than dividing by zero. Every shipped system populates it. + let tc_flops = spec.gpu.bfloat16_tc_flops.unwrap_or(1.0); + let sol_math = (ops as f64) / (tc_flops * quant.mapping().compute) * 1000.0; + let sol_mem = mem_bytes / spec.gpu.mem_bw * 1000.0; + sol_math.max(sol_mem) +} + #[cfg(test)] mod tests { use super::*; @@ -209,6 +664,190 @@ mod tests { } } + fn b200_trtllm_db() -> PerfDatabase { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..").join("src/aiconfigurator/systems"); + PerfDatabase::load(&root, "b200_sxm", "trtllm", "1.2.0rc5").expect("db loads") + } + + /// Python `resolve_transfer_policy("conservative")`: xshape only. + const CONSERVATIVE: TransferPolicy = TransferPolicy { + xshape: true, + xquant: false, + xprofile: false, + xop: false, + }; + + /// Qwen3-235B-A22B expert shape on b200/vllm/0.19.0 (collected for + /// bfloat16 at tp=1/ep=8 under power_law_1.2). + fn qwen3_op(quant: MoeQuantMode) -> MoeOp { + MoeOp { + name: "moe".into(), + scale_factor: 1.0, + hidden_size: 4096, + inter_size: 1536, + topk: 8, + num_experts: 128, + moe_tp_size: 1, + moe_ep_size: 8, + quant_mode: quant, + workload_distribution: "power_law_1.2".into(), + attention_dp_size: 1, + is_gated: true, + } + } + + fn assert_oracle(result: &PerformanceResult, expected: f64, source: Source, label: &str) { + assert!( + (result.latency_ms - expected).abs() < 1e-9, + "{label}: expected {expected}, got {}", + result.latency_ms + ); + assert_eq!(result.source, source, "{label}: wrong source"); + } + + /// Oracle values generated from the Python reference on the same data + /// (shared layer pinned OFF so Python reads exactly the primary parquet + /// the Rust table loads): + /// + /// ```text + /// db = perf_database.get_database_view("b200_sxm", "vllm", "0.19.0", + /// allow_missing_data=True, database_mode=..., transfer_policy=..., + /// shared_layer=False) + /// float(MoE._query_moe_table(db, num_tokens=..., hidden_size=4096, + /// inter_size=1536, topk=8, num_experts=128, moe_tp_size=1, + /// moe_ep_size=8, quant_mode=..., workload_distribution="power_law_1.2", + /// database_mode=...)) + /// ``` + /// + /// Regenerate if the shipped MoE table or the util-empirical math changes. + #[test] + fn moe_empirical_own_shape_matches_python_oracles() { + let db = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Empirical, TransferPolicy::ALL); + let op = qwen3_op(MoeQuantMode::Bfloat16); + let r333 = op.query(&db, 333).expect("own-shape empirical t=333"); + assert_oracle(&r333, 0.19184494219320924, Source::Empirical, "own_emp_t333"); + let r96 = op.query(&db, 96).expect("own-shape empirical t=96"); + assert_oracle(&r96, 0.13852159976959227, Source::Empirical, "own_emp_t96"); + } + + /// HYBRID with data present must stay on silicon (exact hit and in-range + /// interpolation), and its interpolated value differs from the empirical + /// reconstruction at the same point (0.19178... vs 0.19184...). + #[test] + fn moe_hybrid_prefers_silicon_when_covered() { + let db = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, TransferPolicy::ALL); + let op = qwen3_op(MoeQuantMode::Bfloat16); + let hit = op.query(&db, 128).expect("collected token point"); + assert_oracle(&hit, 0.146451199054718, Source::Silicon, "hyb_silicon_t128"); + let interp = op.query(&db, 333).expect("in-range token interp"); + assert_oracle(&interp, 0.19178520552814007, Source::Silicon, "hyb_t333"); + } + + /// XQUANT tier: `w4a16_mxfp4_cutlass` is uncollected on b200/vllm/0.19.0 + /// but shares the (0.5, 1) profile with collected int4_wo / w4a16_mxfp4; + /// the borrowed util curve is reconstructed with the QUERY quant's SOL. + #[test] + fn moe_xquant_transfer_matches_python_oracle() { + let db = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, TransferPolicy::ALL); + let op = qwen3_op(MoeQuantMode::W4a16Mxfp4Cutlass); + let r = op.query(&db, 96).expect("xquant transfer"); + assert_oracle(&r, 0.329638409614563, Source::Empirical, "xquant_t96"); + } + + /// XPROFILE tier: `w4afp8` ((0.5, 2)) has no same-profile sibling in the + /// table; the nearest-profile quant (fp8, distance 0.5) supplies the util + /// curve built with ITS own SOL, rescaled by e(w4afp8)/e(fp8) = 0.15/0.40. + #[test] + fn moe_xprofile_transfer_matches_python_oracle() { + let db = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, TransferPolicy::ALL); + let op = qwen3_op(MoeQuantMode::W4afp8); + let r = op.query(&db, 96).expect("xprofile transfer"); + assert_oracle(&r, 0.13701972961425785, Source::Empirical, "xprofile_t96"); + } + + /// XSHAPE tier: same quant (bfloat16), uncollected inter_size 1600 → + /// nearest collected sibling (8, 128, 4096, 1536). Also reachable under + /// the conservative (xshape-only) policy with the identical value. + #[test] + fn moe_xshape_transfer_matches_python_oracle() { + let mut op = qwen3_op(MoeQuantMode::Bfloat16); + op.inter_size = 1600; + + let db = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, TransferPolicy::ALL); + let r = op.query(&db, 96).expect("xshape transfer"); + assert_oracle(&r, 0.14427836344943168, Source::Empirical, "xshape_t96"); + + let conservative = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, CONSERVATIVE); + let rc = op.query(&conservative, 96).expect("xshape under conservative policy"); + assert_oracle(&rc, 0.14427836344943168, Source::Empirical, "conservative_xshape_t96"); + } + + /// Policy gating: disabled tiers are SKIPPED, and the terminal + /// EmpiricalNotImplemented only fires when every permitted tier found + /// nothing — `off` blocks everything for an uncollected quant, and + /// `conservative` (xshape only) blocks the xquant-needing case. + #[test] + fn moe_transfer_policy_gates_tiers() { + let op = qwen3_op(MoeQuantMode::W4a16Mxfp4Cutlass); + + let off = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, TransferPolicy::OFF); + let blocked = op.query(&off, 96); + assert!( + matches!(blocked, Err(AicError::EmpiricalNotImplemented(_))), + "off policy must surface the typed empirical miss, got {blocked:?}" + ); + + let conservative = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, CONSERVATIVE); + let blocked = op.query(&conservative, 96); + assert!( + matches!(blocked, Err(AicError::EmpiricalNotImplemented(_))), + "conservative policy must not reach the xquant tier, got {blocked:?}" + ); + } + + /// Low-latency kernel-table selection inside the EMPIRICAL path + /// (b200/trtllm/1.2.0rc5 carries `moe_torch_flow_min_latency` rows): + /// nvfp4 gated at t<=128 with the slice present builds the util grid + /// from the LL table; t>128 uses the regular table at the same shape + /// (~3x apart); an uncollected shape at t<=128 fails the LL probe and + /// runs the xshape ladder over the REGULAR table. Oracles: + /// + /// ```text + /// db = perf_database.get_database_view("b200_sxm", "trtllm", "1.2.0rc5", + /// allow_missing_data=True, database_mode="EMPIRICAL", shared_layer=False) + /// float(MoE._query_moe_table(db, num_tokens=..., hidden_size=6144, + /// inter_size=..., topk=2, num_experts=8, moe_tp_size=32, moe_ep_size=1, + /// quant_mode=common.MoEQuantMode.nvfp4, workload_distribution="balanced", + /// database_mode=common.DatabaseMode.EMPIRICAL)) + /// ``` + #[test] + fn moe_empirical_low_latency_table_selection_matches_python_oracles() { + let db = b200_trtllm_db().with_mode(crate::common::enums::DatabaseMode::Empirical, TransferPolicy::ALL); + let op = MoeOp { + name: "moe-ll".into(), + scale_factor: 1.0, + hidden_size: 6144, + inter_size: 16384, + topk: 2, + num_experts: 8, + moe_tp_size: 32, + moe_ep_size: 1, + quant_mode: MoeQuantMode::Nvfp4, + workload_distribution: "balanced".into(), + attention_dp_size: 1, + is_gated: true, + }; + let ll = op.query(&db, 100).expect("ll-table empirical t=100"); + assert_oracle(&ll, 0.023113779703977197, Source::Empirical, "ll_own_t100"); + let std_table = op.query(&db, 200).expect("std-table empirical t=200"); + assert_oracle(&std_table, 0.058452753259364186, Source::Empirical, "std_own_t200"); + + let mut off_shape = op.clone(); + off_shape.inter_size = 17000; + let xshape = off_shape.query(&db, 100).expect("failed ll probe -> std xshape"); + assert_oracle(&xshape, 0.05842286435922407, Source::Empirical, "nvfp4_xshape_t100"); + } + /// With attention-dp, all dp ranks' tokens funnel into the shared expert /// pool: query(dp=4, t) must equal query(dp=1, 4t). Dropping the /// multiplier under-predicted MoE latency ~4.7x on dp=8 DeepSeek configs diff --git a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs index 9a4123a1f..a80aa1843 100644 --- a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs +++ b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs @@ -637,7 +637,7 @@ mod tests { &[(2, 16, 64, 100.0), (8, 16, 64, 900.0)], ); let mut db = b200_sglang_db(); // b200_sxm: num_gpus_per_node = 8 - db.wideep = crate::perf_database::wideep::WideEpTable::new(tmp.path().to_path_buf()); + db.tables_mut().wideep = crate::perf_database::wideep::WideEpTable::new(tmp.path().to_path_buf()); // num_gpus = moe_tp * moe_ep = 16 -> node_num = 16 / 8 = 2. let got = deepep_op(16, true, DispatchFlavor::DeepEpNormal) @@ -688,7 +688,7 @@ mod tests { &[(1, 64, 30.0, 20.0)], // full sum = 50us ); let mut db = b200_sglang_db(); - db.wideep = crate::perf_database::wideep::WideEpTable::new(tmp.path().to_path_buf()); + db.tables_mut().wideep = crate::perf_database::wideep::WideEpTable::new(tmp.path().to_path_buf()); // Context (DeepEP normal): pre == combine == full point sum. let pre = deepep_op(8, true, DispatchFlavor::DeepEpNormal) diff --git a/rust/aiconfigurator-core/src/operators/op.rs b/rust/aiconfigurator-core/src/operators/op.rs index 37a480d1d..1f694cb2f 100644 --- a/rust/aiconfigurator-core/src/operators/op.rs +++ b/rust/aiconfigurator-core/src/operators/op.rs @@ -326,7 +326,22 @@ impl Op { // skip subsequent retries — we don't bother here because the // hot-path penalty is one `OnceLock::get` per call on a // populated path and one retry on a missing one.) - match op.primary.query(db, ctx) { + // + // Under HYBRID the primary is evaluated against a SILICON + // view (Python swaps in `_get_configured_database_view(db, + // SILICON, transfer_policy)`): a missing module table must + // fall to the granular fallback chain, not be hybrid- + // estimated at module level. The fallback ops then run + // against the ORIGINAL (hybrid) database. + let silicon_db; + let primary_db: &PerfDatabase = + if db.database_mode == crate::common::enums::DatabaseMode::Hybrid { + silicon_db = db.silicon_view(); + &silicon_db + } else { + db + }; + match op.primary.query(primary_db, ctx) { Ok(r) => Ok(r), Err(AicError::PerfDatabase(_)) | Err(AicError::Io { .. }) => { let mut total = 0.0_f64; diff --git a/rust/aiconfigurator-core/src/operators/wideep_mla.rs b/rust/aiconfigurator-core/src/operators/wideep_mla.rs index bd49d1faa..5e16dbb9d 100644 --- a/rust/aiconfigurator-core/src/operators/wideep_mla.rs +++ b/rust/aiconfigurator-core/src/operators/wideep_mla.rs @@ -17,9 +17,13 @@ //! no prefix concept. use serde::{Deserialize, Serialize}; -use crate::common::enums::{FmhaQuantMode, KvCacheQuantMode}; +use crate::common::enums::{DatabaseMode, FmhaQuantMode, KvCacheQuantMode}; use crate::common::error::AicError; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid}; +use crate::perf_database::wideep_mla::{ + wideep_context_mla_sol_ms, wideep_generation_mla_sol_ms, +}; use crate::perf_database::PerfDatabase; fn prefix_correction(full_s: u32, prefix: u32) -> f64 { @@ -76,30 +80,33 @@ impl WideEpContextMlaOp { prefix: u32, ) -> Result { // ctx(s, pfx): the un-sharded wideep context-MLA query for a sequence - // chunk of length `s` at prefix `pfx`, with prefix correction applied. - let ctx = |s: u32, pfx: u32| -> Result { - let full_s = s + pfx; - let raw = db.wideep_mla.query_context( + // chunk of length `s` at prefix `pfx` (mode dispatch handles prefix + // correction for silicon and the prefix-aware SOL for empirical). + let ctx = |s: u32, pfx: u32| -> Result<(f64, Source), AicError> { + query_wideep_context_mla_table( + db, batch_size, - full_s, + s, + pfx, self.num_heads, self.kv_cache_dtype, self.fmha_quant_mode, &self.attn_backend, - )?; - Ok(raw * prefix_correction(full_s, pfx)) + ) }; // Context parallelism (SGLang AllGather / zigzag): model rank 0's two // balanced chunks, c = ceil(isl / 2cp). Mirrors Python // `WideEPContextMLA.query` (mla.py:1505-1510) and // `operators/mla.rs::ContextMlaOp`. - let latency = if self.cp_size > 1 { + let (latency, source) = if self.cp_size > 1 { let c = isl.div_ceil(2 * self.cp_size).max(1); - ctx(c, prefix)? + ctx(c, prefix + isl - c)? + let (first, first_source) = ctx(c, prefix)?; + let (second, second_source) = ctx(c, prefix + isl - c)?; + (first + second, first_source.combine(second_source)) } else { ctx(isl, prefix)? }; - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } @@ -182,6 +189,115 @@ mod tests { baseline.latency_ms ); } + + fn h200_sglang_db() -> PerfDatabase { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + PerfDatabase::load(&root, "h200_sxm", "sglang", "0.5.10").expect("db loads") + } + + fn assert_close(got: f64, expected: f64, what: &str) { + assert!( + (got - expected).abs() < 1e-9, + "{what}: expected {expected}, got {got}" + ); + } + + /// Oracle values generated from the Python reference on the same data: + /// `WideEPContextMLA._query_wideep_context_mla_table(db, b, s, prefix, + /// tp_size, kv=fp8, fmha=fp8_block, attention_backend, database_mode= + /// EMPIRICAL)` on h200_sxm/sglang/0.5.10 (`get_database(..., + /// shared_layer=False)`). `num_heads` here is Python's `128 // tp_size`. + /// Regenerate if the shipped table or the util math changes. + #[test] + fn wideep_context_mla_empirical_matches_python_oracles() { + let mut db = h200_sglang_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + let cases: &[(u32, u32, u32, u32, &str, f64)] = &[ + // tp=8 (n=16), prefix > 0: the query SOL carries prefix natively + (2, 3000, 1000, 16, "flashinfer", 0.7957964702937849), + // tp=1 (n=128), exact collected hit + (4, 4096, 0, 128, "flashinfer", 9.6274), + // fa3 kernel slice, exact collected hit + (4, 4096, 0, 128, "fa3", 9.7168), + ]; + for &(b, s, prefix, n, backend, expected) in cases { + let (latency, source) = query_wideep_context_mla_table( + &db, + b, + s, + prefix, + n, + KvCacheQuantMode::Fp8, + FmhaQuantMode::Fp8Block, + backend, + ) + .expect("empirical query"); + assert_close( + latency, + expected, + &format!("wideep_ctx_mla(b={b}, s={s}, pfx={prefix}, n={n}, {backend})"), + ); + assert_eq!(source, Source::Empirical); + } + + // HYBRID on a kv slice with no data (bfloat16; h200's wideep tables + // are fp8-KV only) -> terminal EmpiricalNotImplemented miss. + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let result = query_wideep_context_mla_table( + &db, + 4, + 4096, + 0, + 128, + KvCacheQuantMode::Bfloat16, + FmhaQuantMode::Fp8Block, + "flashinfer", + ); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } + + /// Python oracle: `WideEPGenerationMLA._query_wideep_generation_mla_table` + /// in EMPIRICAL mode on h200_sxm/sglang/0.5.10 (shared_layer=False). The + /// caller's fmha label is inert (derived from the fp8 KV cache), verified + /// on the Python side: fp8_block and bfloat16 labels give the same value. + #[test] + fn wideep_generation_mla_empirical_matches_python_oracles() { + let mut db = h200_sglang_db(); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + let cases: &[(u32, u32, u32, &str, f64)] = &[ + // tp=8 (n=16), off-grid (b, s) + (3, 5000, 16, "flashinfer", 0.0509930128230621), + // tp=1 (n=128), exact collected hit + (1, 4096, 128, "flashinfer", 0.1017), + // fa3 kernel slice, off-grid seq + (2, 9000, 128, "fa3", 0.12152241047815426), + ]; + for &(b, s, n, backend, expected) in cases { + let (latency, source) = query_wideep_generation_mla_table( + &db, + b, + s, + n, + KvCacheQuantMode::Fp8, + backend, + ) + .expect("empirical query"); + assert_close( + latency, + expected, + &format!("wideep_gen_mla(b={b}, s={s}, n={n}, {backend})"), + ); + assert_eq!(source, Source::Empirical); + } + + // HYBRID on a kv slice with no data (bfloat16) -> terminal miss. + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let result = + query_wideep_generation_mla_table(&db, 1, 4096, 128, KvCacheQuantMode::Bfloat16, "flashinfer"); + assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); + } } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -220,15 +336,206 @@ impl WideEpGenerationMlaOp { batch_size: u32, s: u32, ) -> Result { - let latency = db.wideep_mla.query_generation( + let (latency, source) = query_wideep_generation_mla_table( + db, batch_size, s, self.num_heads, self.kv_cache_dtype, &self.attn_backend, )?; - Ok(PerformanceResult::new(latency, Source::Silicon) + Ok(PerformanceResult::new(latency, source) .clamp_non_negative() .scaled(self.scale_factor)) } } + +// --------------------------------------------------------------------------- +// Database-mode dispatch, mirroring the Python `_query_wideep_*_table` +// classmethods (`operations/mla.py`): SILICON queries the table; HYBRID +// converts a typed silicon miss into the util-space empirical estimate; +// EMPIRICAL always estimates. The SOL diagnostic modes never reach the +// compiled engine. Python's classmethods take `tp_size` and use +// `num_head = 128 // tp_size` everywhere (query coordinate AND query SOL); +// these dispatches take that `num_heads` value directly, matching the op +// struct / table surface. +// --------------------------------------------------------------------------- + +/// Sample-coordinate head mapping for the util-empirical grids. Python's +/// `get_empirical` sol_fn maps a collected `num_heads` coordinate via +/// `tp = round(128 / c[0])` (banker's rounding) then `128 // tp` — unlike +/// the silicon sol_fn, which floors both divisions. +fn wideep_sample_num_head(n: f64) -> f64 { + let tp_size = (128.0 / n).round_ties_even(); + (128.0 / tp_size).floor() +} + +/// WideEP context MLA latency (prefix correction applied on the silicon +/// branch) under the database's query mode. +#[allow(clippy::too_many_arguments)] +fn query_wideep_context_mla_table( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + attn_backend: &str, +) -> Result<(f64, Source), AicError> { + let silicon = || -> Result { + let full_s = s + prefix; + let raw = db.wideep_mla.query_context( + b, + full_s, + num_heads, + kv_quant, + fmha_quant, + attn_backend, + )?; + Ok(raw * prefix_correction(full_s, prefix)) + }; + match db.database_mode { + DatabaseMode::Empirical => Ok(( + wideep_context_mla_empirical(db, b, s, prefix, num_heads, kv_quant, fmha_quant, attn_backend)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + wideep_context_mla_empirical( + db, b, s, prefix, num_heads, kv_quant, fmha_quant, attn_backend, + )?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((silicon()?, Source::Silicon)), + } +} + +/// `SOL(query)/util` over the `(attn_backend, fmha, kv)` slice's own +/// `(num_heads, s, b)` grid. Mirrors +/// `WideEPContextMLA._query_wideep_context_mla_table::get_empirical` +/// (depth 3; samples are prefix=0, the query SOL carries prefix natively; +/// the kv quant keys the slice only — the context SOL never reads it). +#[allow(clippy::too_many_arguments)] +fn wideep_context_mla_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + prefix: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + attn_backend: &str, +) -> Result { + let spec = &db.system_spec; + // c = (num_heads, full_s, b); tp = round(128 / c[0]) per Python. + let sol = |c: &[f64]| { + wideep_context_mla_sol_ms(spec, fmha_quant, wideep_sample_num_head(c[0]), c[1], 0.0, c[2]) + }; + let key = format!( + "wideep_ctx_mla:{attn_backend}:{}:{}", + fmha_quant.name(), + kv_quant.name() + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.wideep_mla.context_points(attn_backend, kv_quant, fmha_quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + // Query SOL uses the op's own head count (= Python's `128 // tp_size`). + let sol_query = wideep_context_mla_sol_ms( + spec, + fmha_quant, + num_heads as f64, + s as f64, + prefix as f64, + b as f64, + ); + let query = [num_heads as f64, (s + prefix) as f64, b as f64]; + let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + +/// WideEP generation MLA latency under the database's query mode. The fmha +/// mode is derived from the kv-cache dtype (fp8 KV -> fp8, else bfloat16), +/// exactly like the top of Python's +/// `_query_wideep_generation_mla_table` — the caller's fmha label is inert +/// for generation. +fn query_wideep_generation_mla_table( + db: &PerfDatabase, + b: u32, + s: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + attn_backend: &str, +) -> Result<(f64, Source), AicError> { + match db.database_mode { + DatabaseMode::Empirical => Ok(( + wideep_generation_mla_empirical(db, b, s, num_heads, kv_quant, attn_backend)?, + Source::Empirical, + )), + DatabaseMode::Hybrid => { + match db + .wideep_mla + .query_generation(b, s, num_heads, kv_quant, attn_backend) + { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + wideep_generation_mla_empirical(db, b, s, num_heads, kv_quant, attn_backend)?, + Source::Empirical, + )), + Err(err) => Err(err), + } + } + _ => Ok(( + db.wideep_mla + .query_generation(b, s, num_heads, kv_quant, attn_backend)?, + Source::Silicon, + )), + } +} + +/// `SOL(query)/util` over the `(attn_backend, kv)` slice's own +/// `(num_heads, b, s)` grid. Mirrors +/// `WideEPGenerationMLA._query_wideep_generation_mla_table::get_empirical`. +fn wideep_generation_mla_empirical( + db: &PerfDatabase, + b: u32, + s: u32, + num_heads: u32, + kv_quant: KvCacheQuantMode, + attn_backend: &str, +) -> Result { + // Decode compute dtype follows the kv-cache dtype (Python overrides the + // passed fmha label before get_sol is defined). + let fmha_quant = if kv_quant == KvCacheQuantMode::Fp8 { + FmhaQuantMode::Fp8 + } else { + FmhaQuantMode::Bfloat16 + }; + let spec = &db.system_spec; + // c = (num_heads, b, s); tp = round(128 / c[0]) per Python. + let sol = |c: &[f64]| { + wideep_generation_mla_sol_ms(spec, fmha_quant, wideep_sample_num_head(c[0]), c[1], c[2]) + }; + let key = format!("wideep_gen_mla:{attn_backend}:{}", kv_quant.name()); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.wideep_mla.generation_points(attn_backend, kv_quant) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples(points, sol)))), + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + // Query SOL uses the op's own head count (= Python's `128 // tp_size`). + let sol_query = wideep_generation_mla_sol_ms(spec, fmha_quant, num_heads as f64, b as f64, s as f64); + let query = [num_heads as f64, b as f64, s as f64]; + let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} diff --git a/rust/aiconfigurator-core/src/perf_database/attention.rs b/rust/aiconfigurator-core/src/perf_database/attention.rs index f7658aba2..a6b9e8af2 100644 --- a/rust/aiconfigurator-core/src/perf_database/attention.rs +++ b/rust/aiconfigurator-core/src/perf_database/attention.rs @@ -255,6 +255,153 @@ impl AttentionTable { perf_interp::query(&cfg, node, &[n as f64, s as f64, b as f64]) } + /// Collected `(num_heads, full_seq, batch) -> latency` points of one + /// context slice, for the operator-layer util-calibration grid (Python's + /// `require_data_slice(_context_attention_data, fmha, kv, n_kv, hs, w)` + + /// `iter_grid(..., depth=3)`). `n_kv_lookup` is the MHA-normalized + /// kv-head count (`0` == MHA). Missing slice / empty node is a typed + /// `PerfDatabase` miss. No estimation logic here — callers own the + /// SOL/util math. + pub fn context_points( + &self, + fmha_quant: FmhaQuantMode, + kv_quant: KvCacheQuantMode, + n_kv_lookup: u32, + head_size: u32, + window_size: u32, + ) -> Result, f64)>, AicError> { + let grids = self.load_context()?; + let key = ContextKey { + fmha_quant: fmha_quant.name().to_string(), + kv_quant: kv_quant.name().to_string(), + n_kv_lookup, + head_size, + window_size, + }; + let node = grids.by_keys.get(&key).ok_or_else(|| missing_key(&self.data_root, &key))?; + let points = perf_interp::node_points(node); + if points.is_empty() { + return Err(missing_key(&self.data_root, &key)); + } + Ok(points) + } + + /// Distinct collected `head_size` keys under `(fmha, kv, n_kv_lookup)`, + /// any window — the cross-head_size (XSHAPE) candidate list (Python's + /// `require_data_slice(wrapper, fmha, kv, n_kv).keys()`). Returned in + /// ascending order; Python yields CSV insertion order instead, which is + /// observable only on exact log-distance ties in the reference pick. + /// Typed `PerfDatabase` miss when nothing matches. + pub fn context_head_sizes( + &self, + fmha_quant: FmhaQuantMode, + kv_quant: KvCacheQuantMode, + n_kv_lookup: u32, + ) -> Result, AicError> { + let grids = self.load_context()?; + let fmha = fmha_quant.name(); + let kv = kv_quant.name(); + let mut sizes: Vec = Vec::new(); + for key in grids.by_keys.keys() { + if key.fmha_quant == fmha + && key.kv_quant == kv + && key.n_kv_lookup == n_kv_lookup + && !sizes.contains(&key.head_size) + { + sizes.push(key.head_size); + } + } + if sizes.is_empty() { + return Err(AicError::PerfDatabase(format!( + "context attention data missing for fmha={fmha}, kv={kv}, \ + n_kv={n_kv_lookup} at {}", + self.data_root.display() + ))); + } + Ok(sizes) + } + + /// Collected `(num_heads, batch, seq) -> latency` points of one + /// generation slice. Python calibrates from + /// `_raw_generation_attention_data`, which in v2 is an alias of the + /// SOL-clamped working table (`_correct_sol` runs before the alias is + /// taken) — exactly what [`AttentionTable::load_generation`] produces, so + /// this IS the RAW-table equivalent. Typed `PerfDatabase` miss when the + /// slice is absent/empty. + pub fn generation_points( + &self, + kv_quant: KvCacheQuantMode, + n_kv_lookup: u32, + head_size: u32, + window_size: u32, + ) -> Result, f64)>, AicError> { + let grids = self.load_generation()?; + let key = GenerationKey { + kv_quant: kv_quant.name().to_string(), + n_kv_lookup, + head_size, + window_size, + }; + let node = grids + .by_keys + .get(&key) + .ok_or_else(|| missing_gen_key(&self.data_root, &key))?; + let points = perf_interp::node_points(node); + if points.is_empty() { + return Err(missing_gen_key(&self.data_root, &key)); + } + Ok(points) + } + + /// Distinct collected `head_size` keys under `(kv, n_kv_lookup)`, any + /// window — the decode XSHAPE candidate list. Same ordering note as + /// [`AttentionTable::context_head_sizes`]. + pub fn generation_head_sizes( + &self, + kv_quant: KvCacheQuantMode, + n_kv_lookup: u32, + ) -> Result, AicError> { + let grids = self.load_generation()?; + let kv = kv_quant.name(); + let mut sizes: Vec = Vec::new(); + for key in grids.by_keys.keys() { + if key.kv_quant == kv && key.n_kv_lookup == n_kv_lookup && !sizes.contains(&key.head_size) { + sizes.push(key.head_size); + } + } + if sizes.is_empty() { + return Err(AicError::PerfDatabase(format!( + "generation attention data missing for kv={kv}, n_kv={n_kv_lookup} at {}", + self.data_root.display() + ))); + } + Ok(sizes) + } + + /// Collected `(num_heads, seq, batch) -> latency` points of one encoder + /// slice (own-shape only; encoder has no transfer ladder). Typed + /// `PerfDatabase` miss when the slice is absent/empty. + pub fn encoder_points( + &self, + fmha_quant: FmhaQuantMode, + head_size: u32, + ) -> Result, f64)>, AicError> { + let grids = self.load_encoder()?; + let key = EncoderKey { + fmha_quant: fmha_quant.name().to_string(), + head_size, + }; + let node = grids + .by_keys + .get(&key) + .ok_or_else(|| missing_encoder_key(&self.data_root, &key))?; + let points = perf_interp::node_points(node); + if points.is_empty() { + return Err(missing_encoder_key(&self.data_root, &key)); + } + Ok(points) + } + fn load_context(&self) -> Result<&ContextGrids, AicError> { let cell = self.context.get_or_init(|| { let raw = load_context_parquet(&self.context_sources)?; @@ -451,7 +598,7 @@ fn load_generation_parquet( /// `n_kv_lookup == 0` means MHA (n_kv tracks n); a positive `window_size` /// smaller than the seq cuts the O(s^2) causal work to O(s*w). #[allow(clippy::too_many_arguments)] -fn context_attention_sol_ms( +pub(crate) fn context_attention_sol_ms( spec: &SystemSpec, n_kv_lookup: u32, head_size: u32, @@ -483,12 +630,55 @@ fn context_attention_sol_ms( sol_math.max(sol_mem) } +/// Speed-of-light context-attention latency in ms for a QUERY with a prefix. +/// +/// Mirrors Python's `ContextAttention._query_context_attention_table::get_sol` +/// verbatim: `full_s = s + prefix`; the windowed branch fires when `w > 0 && +/// full_s > w`; the causal branch discounts already-computed prefix work +/// (`full_s² − prefix²`) and the Q/output traffic covers only the new tokens +/// (`full_s − prefix`) while the KV write spans the full sequence. +/// [`context_attention_sol_ms`] is the `prefix = 0` specialization used as +/// the per-sample sol_fn (table rows are full attention); this variant feeds +/// the empirical query SOL. `n_kv` is the REAL kv-head count (not the MHA +/// sentinel). +#[allow(clippy::too_many_arguments)] +pub(crate) fn context_attention_sol_with_prefix_ms( + spec: &SystemSpec, + b: f64, + s: f64, + prefix: f64, + n: f64, + n_kv: f64, + head_size: u32, + window_size: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, +) -> f64 { + let bf16_flops = spec.gpu.bfloat16_tc_flops.unwrap_or(0.0); + if bf16_flops <= 0.0 { + return 0.0; + } + let h = head_size as f64; + let w = window_size as f64; + let full_s = s + prefix; + let ops = if window_size > 0 && full_s > w { + 2.0 * b * (full_s - prefix) * w * n * h * 2.0 + } else { + 2.0 * b * (full_s * full_s - prefix * prefix) * n * h * 2.0 / 2.0 + }; + let mem_bytes = 2.0 * b * (n * (full_s - prefix) * h + n * (full_s - prefix) * h) + + kv_quant.mapping().memory * b * (2.0 * n_kv * full_s * h); + let sol_math = ops / bf16_flops * 1000.0 / fmha_quant.mapping().compute; + let sol_mem = mem_bytes / spec.gpu.mem_bw * 1000.0; + sol_math.max(sol_mem) +} + /// Speed-of-light encoder-attention latency in ms. /// /// Mirrors Python's `EncoderAttention._query_encoder_attention_table::get_sol`: /// non-causal full N^2 (no /2), no KV-cache read — Q/K/V read + output write /// in bf16 only. -fn encoder_attention_sol_ms( +pub(crate) fn encoder_attention_sol_ms( spec: &SystemSpec, head_size: u32, fmha_quant: FmhaQuantMode, @@ -514,7 +704,7 @@ fn encoder_attention_sol_ms( /// as wired into the perf_interp sol_fn: c = [n, b, s]. `n_kv_lookup == 0` /// means MHA (n_kv tracks n); `window_size > 0` clamps `kv_len` to /// `min(s-1, window_size)`. -fn generation_attention_sol_ms( +pub(crate) fn generation_attention_sol_ms( spec: &SystemSpec, n_kv_lookup: u32, head_size: u32, diff --git a/rust/aiconfigurator-core/src/perf_database/dsa.rs b/rust/aiconfigurator-core/src/perf_database/dsa.rs index 1499ce489..49c04aa79 100644 --- a/rust/aiconfigurator-core/src/perf_database/dsa.rs +++ b/rust/aiconfigurator-core/src/perf_database/dsa.rs @@ -414,6 +414,57 @@ impl DsaTable { ) } + /// RAW context slice `[num_heads][prefix][isl][batch]` for the exact + /// (arch, quants) key, with Python `_select_dsa_backend`'s fallback chain + /// applied. This is the calibration view the util-empirical path reads: + /// with load-time pre-expansion gone, Python's `_raw_context_dsa_module_data` + /// is a plain alias of the loaded table, and these grids ARE that table. + /// Typed `AicError::PerfDatabase` when the table/slice is missing — the + /// operator maps it to a typed coverage miss, mirroring Python + /// `_raw_slice` raising `PerfDataNotAvailableError`. + pub fn context_raw_slice( + &self, + key: &DsaKey, + dsa_backend: &str, + ) -> Result<&DsaHeadGrid, AicError> { + let grids = self.load_context()?; + grids + .by_keys + .get(key) + .and_then(|by_backend| select_dsa_backend(by_backend, dsa_backend)) + .ok_or_else(|| { + missing( + "raw context DSA module", + &self.data_root, + format!("{key:?} (dsa_backend={dsa_backend})"), + ) + }) + } + + /// RAW generation slice for the exact (arch, quants) key, backend-selected + /// like [`Self::context_raw_slice`]. The load-time (isl, step) -> seq + /// collapse stores `step = 0, isl = seq`, so the returned grid reads as + /// `[num_heads][0][seq][batch]` — the same measured rows Python's + /// `_raw_generation_dsa_module_data` alias exposes as `[num_heads][b][s]`. + pub fn generation_raw_slice( + &self, + key: &DsaKey, + dsa_backend: &str, + ) -> Result<&DsaHeadGrid, AicError> { + let grids = self.load_generation()?; + grids + .by_keys + .get(key) + .and_then(|by_backend| select_dsa_backend(by_backend, dsa_backend)) + .ok_or_else(|| { + missing( + "raw generation DSA module", + &self.data_root, + format!("{key:?} (dsa_backend={dsa_backend})"), + ) + }) + } + fn load_context(&self) -> Result<&DsaGrids, AicError> { let cell = self .context @@ -537,7 +588,7 @@ fn indexer_cache_entry_bytes(index_head_dim: i64) -> i64 { /// attention group (fmha_quant) whose exact KV pair count is /// `sum_{i=0..s-1} min(prefix+i+1, index_topk)`. #[allow(clippy::too_many_arguments)] -fn dsa_context_sol_ms( +pub(crate) fn dsa_context_sol_ms( spec: &SystemSpec, dims: &DsaDims, index_topk: i64, @@ -632,7 +683,7 @@ fn dsa_context_sol_ms( /// `GenerationDSAModule._query_generation_dsa_module_table::get_sol` /// (1 token per request; the attention group is hardcoded bfloat16 in /// Python — `fmha_mode = FMHAQuantMode.bfloat16` — so no fmha arg here). -fn dsa_generation_sol_ms( +pub(crate) fn dsa_generation_sol_ms( spec: &SystemSpec, dims: &DsaDims, kv_quant: KvCacheQuantMode, diff --git a/rust/aiconfigurator-core/src/perf_database/dsv4.rs b/rust/aiconfigurator-core/src/perf_database/dsv4.rs index 364a8648b..e7054ab69 100644 --- a/rust/aiconfigurator-core/src/perf_database/dsv4.rs +++ b/rust/aiconfigurator-core/src/perf_database/dsv4.rs @@ -508,6 +508,96 @@ impl Dsv4Table { } } + /// Collected context-module points `(prefix, s, b) -> latency` for the + /// operator-layer util-calibration grid (Python + /// `_query_context_attn_table::get_empirical`'s `_slice()`: + /// `require_data_slice(data, fmha, kv, gemm)` -> head resolution -> + /// `require_data_slice(quant_data, head, compress_ratio)`). Slices exactly + /// like the silicon query ([`select_resolved`], which additionally keys on + /// `architecture` — Python's loader merges architectures at load; shipped + /// files carry one). Missing slice / empty node is a typed miss. + pub fn context_points( + &self, + attn_kind: AttnKind, + local_heads: u32, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + gemm_quant: GemmQuantMode, + architecture: &str, + ) -> Result, f64)>, AicError> { + let grids = match attn_kind { + AttnKind::Csa => self.load_csa_context()?, + AttnKind::Hca => self.load_hca_context()?, + }; + let node = select_resolved(grids, architecture, fmha_quant, kv_quant, gemm_quant, local_heads)?; + let points = perf_interp::node_points(node); + if points.is_empty() { + return Err(AicError::PerfDatabase(format!( + "DSV4 context module data empty for local_heads={local_heads}, \ + attn_kind={attn_kind:?}, architecture='{architecture}'" + ))); + } + Ok(points) + } + + /// Collected generation-module points `(b, s_total) -> latency` for the + /// operator-layer util-calibration grid (Python + /// `_query_generation_attn_table::get_empirical`'s `_slice()`: + /// `require_data_slice(data, kv, gemm)` -> head resolution -> + /// `require_data_slice(quant_data, head, compress_ratio)`). + /// + /// Python's GENERATION loader keys `[kv][gemm][head]` with NO fmha level + /// (the CSV `mla_dtype` column is ignored on load), while the Rust table + /// keys it from the column: gather every matching `(arch, kv, gemm)` key + /// regardless of fmha. Shipped files carry one `mla_dtype` per + /// `(kv, gemm)`; if several existed Python would last-wins-merge shared + /// cells while this concatenation keeps both (the util grid then prefers + /// the first at an exact duplicate) — flagged here, not silently + /// divergent, because no shipped table hits it. + pub fn generation_points( + &self, + attn_kind: AttnKind, + local_heads: u32, + kv_quant: KvCacheQuantMode, + gemm_quant: GemmQuantMode, + architecture: &str, + ) -> Result, f64)>, AicError> { + let grids = match attn_kind { + AttnKind::Csa => self.load_csa_generation()?, + AttnKind::Hca => self.load_hca_generation()?, + }; + let mut by_head: BTreeMap> = BTreeMap::new(); + for (key, per_head) in &grids.by_keys { + if key.architecture == architecture + && key.kv_quant == kv_quant.name() + && key.gemm_quant == gemm_quant.name() + { + for (&head, node) in per_head { + by_head.entry(head).or_default().push(node); + } + } + } + let head = resolve_head_key(&by_head, local_heads).ok_or_else(|| { + AicError::PerfDatabase(format!( + "DSV4 generation module data missing for local_heads={local_heads}, \ + attn_kind={attn_kind:?}, kv='{}', gemm='{}', architecture='{architecture}'", + kv_quant.name(), + gemm_quant.name() + )) + })?; + let points: Vec<(Vec, f64)> = by_head[&head] + .iter() + .flat_map(|node| perf_interp::node_points(node)) + .collect(); + if points.is_empty() { + return Err(AicError::PerfDatabase(format!( + "DSV4 generation module data empty for local_heads={local_heads}, \ + attn_kind={attn_kind:?}, architecture='{architecture}'" + ))); + } + Ok(points) + } + /// Absolute top_last topk latency at `(isl, step)` for batch `b`. CP /// composition only — reads the RAW top_last rows the calib loader /// retains alongside the DELTA pairs (Python `_csa_topk_top_last` / @@ -974,7 +1064,7 @@ fn compressed_context_pairs(batch: i128, query_len: i128, prefix: i128, ratio: i /// `dims.local_o_groups` is the rank-local group count Python passes as /// `o_groups` (see [`Dsv4SolDims`]). #[allow(clippy::too_many_arguments)] -fn dsv4_attention_sol_ms( +pub(crate) fn dsv4_attention_sol_ms( spec: &SystemSpec, dims: &Dsv4SolDims, compress_ratio: i64, diff --git a/rust/aiconfigurator-core/src/perf_database/mhc.rs b/rust/aiconfigurator-core/src/perf_database/mhc.rs index 6e96e4b63..40fe6be31 100644 --- a/rust/aiconfigurator-core/src/perf_database/mhc.rs +++ b/rust/aiconfigurator-core/src/perf_database/mhc.rs @@ -129,6 +129,43 @@ impl MhcTable { query_token_curve(by_tokens, num_tokens as f64, &|t| sol(op, t)) } + /// Collected `(num_tokens,) -> latency` points for one RESOLVED op half + /// (`pre` / `post`), for the operator-layer util-calibration grid (Python + /// `_query_mhc_table::get_empirical`'s + /// `require_data_slice(mhc_data, op_name, hc_mult, hidden_size)` + + /// `iter_grid(..., depth=1)`). Missing key / empty curve is a typed + /// `PerfDatabase` miss; the `both` composition lives in the operator + /// (Python sums `_emp_for_op("pre") + _emp_for_op("post")`). + pub fn module_points( + &self, + op: &str, + hc_mult: u32, + hidden_size: u32, + ) -> Result, f64)>, AicError> { + let grids = self.load()?; + let key = MhcKey { + op_name: op.to_string(), + hc_mult, + hidden_size, + }; + let by_tokens = grids.by_keys.get(&key).ok_or_else(|| { + AicError::PerfDatabase(format!( + "MHC module data missing for {key:?} at {}", + self.data_root.display() + )) + })?; + if by_tokens.is_empty() { + return Err(AicError::PerfDatabase(format!( + "MHC module data empty for {key:?} at {}", + self.data_root.display() + ))); + } + Ok(by_tokens + .iter() + .map(|(&tokens, &latency)| (vec![f64::from(tokens)], latency)) + .collect()) + } + fn load(&self) -> Result<&MhcGrids, AicError> { let cell = self .module diff --git a/rust/aiconfigurator-core/src/perf_database/mla.rs b/rust/aiconfigurator-core/src/perf_database/mla.rs index f8c06c112..532725763 100644 --- a/rust/aiconfigurator-core/src/perf_database/mla.rs +++ b/rust/aiconfigurator-core/src/perf_database/mla.rs @@ -326,6 +326,136 @@ impl MlaTable { perf_interp::query(&cfg, node, &[num_heads as f64, b as f64, s as f64]) } + // ----------------------------------------------------------------------- + // Point accessors for the util-space empirical layer (algorithm-free: + // typed `AicError::PerfDatabase` miss on absent slice / empty node, no + // estimation logic). Coordinate order matches the Python `depth=3` + // iteration of each `require_data_slice` slice. + // ----------------------------------------------------------------------- + + /// Collected `(num_heads, seq, batch) -> latency` points of the op-level + /// context MLA `(fmha, kv)` slice. Typed miss when absent/empty. + pub fn context_points( + &self, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + ) -> Result, f64)>, AicError> { + let grids = self.load_context()?; + let key = ContextKey { + fmha_quant: fmha_quant.name().to_string(), + kv_quant: kv_quant.name().to_string(), + }; + let node = grids + .by_keys + .get(&key) + .ok_or_else(|| missing("context MLA", &self.data_root, format!("{key:?}")))?; + non_empty_points(node, "context MLA", &self.data_root) + } + + /// Collected `(num_heads, batch, seq) -> latency` points of the op-level + /// generation MLA kv slice. Typed miss when absent/empty. + pub fn generation_points( + &self, + kv_quant: KvCacheQuantMode, + ) -> Result, f64)>, AicError> { + let grids = self.load_generation()?; + let key = KvOnlyKey { + kv_quant: kv_quant.name().to_string(), + }; + let node = grids + .by_keys + .get(&key) + .ok_or_else(|| missing("generation MLA", &self.data_root, format!("{key:?}")))?; + non_empty_points(node, "generation MLA", &self.data_root) + } + + /// The BMM quant slice Python's `quant_mode in wrapper` membership check + /// selects: the requested quant when ANY BMM data exists for it (either + /// op_name), otherwise `bfloat16` — even when bfloat16 has no data + /// either (the follow-up [`Self::bmm_points`] then reports the typed + /// miss, matching Python's `require_data_slice`). Errs only when the + /// whole BMM table failed to load. + pub fn bmm_selected_quant(&self, quant: GemmQuantMode) -> Result { + let grids = self.load_bmm()?; + let has_quant = grids + .by_keys + .keys() + .any(|key| key.bmm_quant == quant.name()); + Ok(if has_quant { quant } else { GemmQuantMode::Bfloat16 }) + } + + /// Collected `(num_tokens,) -> latency` points of the + /// `(quant, op_name, num_heads)` BMM slice. No fallback here — the caller + /// resolves the quant via [`Self::bmm_selected_quant`] first (mirroring + /// Python, where the membership check is on the quant level only). + /// Typed miss when the slice is absent/empty. + pub fn bmm_points( + &self, + quant: GemmQuantMode, + is_pre: bool, + num_heads: u32, + ) -> Result, f64)>, AicError> { + let grids = self.load_bmm()?; + let pre_or_post = if is_pre { "mla_gen_pre" } else { "mla_gen_post" }; + let key = BmmKey { + bmm_quant: quant.name().to_string(), + pre_or_post: pre_or_post.to_string(), + }; + let node = grids + .by_keys + .get(&key) + .and_then(|by_heads| by_heads.get(&num_heads)) + .ok_or_else(|| { + missing( + "MLA BMM", + &self.data_root, + format!("quant={}, {pre_or_post}, num_heads={num_heads}", quant.name()), + ) + })?; + non_empty_points(node, "MLA BMM", &self.data_root) + } + + /// Collected `(num_heads, seq, batch) -> latency` points of the + /// module-level context MLA `(fmha, kv, gemm)` slice. Typed miss when + /// absent/empty. + pub fn context_module_points( + &self, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + gemm_quant: GemmQuantMode, + ) -> Result, f64)>, AicError> { + let grids = self.load_context_module()?; + let key = ModuleKey { + fmha_quant: fmha_quant.name().to_string(), + kv_quant: kv_quant.name().to_string(), + gemm_quant: gemm_quant.name().to_string(), + }; + let node = grids + .by_keys + .get(&key) + .ok_or_else(|| missing("context MLA module", &self.data_root, format!("{key:?}")))?; + non_empty_points(node, "context MLA module", &self.data_root) + } + + /// Collected `(num_heads, batch, seq) -> latency` points of the + /// module-level generation MLA `(kv, gemm)` slice. Typed miss when + /// absent/empty. + pub fn generation_module_points( + &self, + kv_quant: KvCacheQuantMode, + gemm_quant: GemmQuantMode, + ) -> Result, f64)>, AicError> { + let grids = self.load_generation_module()?; + let key = GenModuleKey { + kv_quant: kv_quant.name().to_string(), + gemm_quant: gemm_quant.name().to_string(), + }; + let node = grids.by_keys.get(&key).ok_or_else(|| { + missing("generation MLA module", &self.data_root, format!("{key:?}")) + })?; + non_empty_points(node, "generation MLA module", &self.data_root) + } + fn load_context(&self) -> Result<&ContextMlaGrids, AicError> { let cell = self .context @@ -378,22 +508,37 @@ fn bf16_tc_flops(spec: &SystemSpec) -> f64 { /// Context MLA SOL in ms, evaluated at prefix = 0 (the perf_interp `sol_fn` /// contract — samples are prefix=0 and the operator layer owns the prefix -/// correction). Mirrors `ContextMLA._query_context_mla_table::get_sol` and +/// correction). See [`context_mla_sol_prefix_ms`] for the formula. +pub(crate) fn context_mla_sol_ms( + spec: &SystemSpec, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + n: f64, + s: f64, + b: f64, +) -> f64 { + context_mla_sol_prefix_ms(spec, kv_quant, fmha_quant, n, s, 0.0, b) +} + +/// Prefix-aware context MLA SOL in ms (`s` is the chunk / isl length; the +/// util-empirical query SOL carries prefix natively, unlike the prefix=0 +/// sample SOL). Mirrors `ContextMLA._query_context_mla_table::get_sol` and /// the identical `MLAModule._query_context_mla_module_table::get_sol`: +/// - `full_s = s + prefix` /// - `ops = b * n * 2/2 * (192 + 128) * (full_s^2 - prefix^2)` /// - `mem = b * n * (kv.memory * full_s * (192+128) + 2 * s * (192+128))` /// - `sol_math = ops / bf16_tc_flops * 1000 / fmha.compute` /// - `sol_mem = mem / mem_bw * 1000` /// - `sol = max(sol_math, sol_mem)` -fn context_mla_sol_ms( +pub(crate) fn context_mla_sol_prefix_ms( spec: &SystemSpec, kv_quant: KvCacheQuantMode, fmha_quant: FmhaQuantMode, n: f64, s: f64, + prefix: f64, b: f64, ) -> f64 { - let prefix = 0.0_f64; let full_s = s + prefix; let ops = b * n * 2.0 / 2.0 * (192.0 + 128.0) * (full_s * full_s - prefix * prefix); let mem_bytes = @@ -411,7 +556,7 @@ fn context_mla_sol_ms( /// - `sol_math = ops / bf16_tc_flops * 1000 / quant_gen.compute` /// - `sol_mem = mem / mem_bw * 1000` /// - `sol = max(sol_math, sol_mem)` -fn generation_mla_sol_ms( +pub(crate) fn generation_mla_sol_ms( spec: &SystemSpec, kv_quant: KvCacheQuantMode, n: f64, @@ -437,7 +582,7 @@ fn generation_mla_sol_ms( /// - `sol_math = ops / (bf16_tc_flops * quant.compute) * 1000` /// - `sol_mem = mem / mem_bw * 1000` /// - `sol = max(sol_math, sol_mem)` -fn mla_bmm_sol_ms(spec: &SystemSpec, quant: GemmQuantMode, n: f64, t: f64) -> f64 { +pub(crate) fn mla_bmm_sol_ms(spec: &SystemSpec, quant: GemmQuantMode, n: f64, t: f64) -> f64 { let ops = 2.0 * t * n * 128.0 * 512.0; let mem_bytes = n * (t * 640.0 + 128.0 * 512.0) * quant.mapping().memory; let sol_math = ops / (bf16_tc_flops(spec) * quant.mapping().compute) * 1000.0; @@ -454,7 +599,7 @@ fn mla_bmm_sol_ms(spec: &SystemSpec, quant: GemmQuantMode, n: f64, t: f64) -> f6 /// - `sol_math = attn_ops/bf16/quant_gen.compute + bmm_ops/(bf16*gemm.compute)` /// - `sol_mem = (attn_mem + bmm_mem) / mem_bw` /// - `sol = max(sol_math, sol_mem)` -fn generation_mla_module_sol_ms( +pub(crate) fn generation_mla_module_sol_ms( spec: &SystemSpec, kv_quant: KvCacheQuantMode, gemm_quant: GemmQuantMode, @@ -810,6 +955,24 @@ fn missing(table: &str, data_root: &Path, descriptor: String) -> AicError { )) } +/// Flatten a slice node into `(coords, latency)` points, treating an empty +/// node as a typed coverage miss (mirrors `require_data_slice`'s empty-node +/// check). +fn non_empty_points( + node: &Node, + table: &str, + data_root: &Path, +) -> Result, f64)>, AicError> { + let points = perf_interp::node_points(node); + if points.is_empty() { + return Err(AicError::PerfDatabase(format!( + "{table} perf data empty for the requested slice at {}", + data_root.display() + ))); + } + Ok(points) +} + fn clone_err(err: &AicError) -> AicError { AicError::PerfDatabase(err.to_string()) } diff --git a/rust/aiconfigurator-core/src/perf_database/mod.rs b/rust/aiconfigurator-core/src/perf_database/mod.rs index 2d710e048..669e7b4e6 100644 --- a/rust/aiconfigurator-core/src/perf_database/mod.rs +++ b/rust/aiconfigurator-core/src/perf_database/mod.rs @@ -11,6 +11,7 @@ //! the WideEP/DeepEP all-to-all variants. use std::path::{Path, PathBuf}; +use std::sync::Arc; use crate::common::enums::{DatabaseMode, TransferPolicy}; use crate::common::error::AicError; @@ -80,13 +81,10 @@ pub use wideep::WideEpTable; pub use wideep_mla::WideEpMlaTable; pub use wideep_moe::WideEpMoeTable; -/// Modular performance database for a specific -/// `//` tuple. -/// -/// `load` does the cheap work: resolves the data directory from the system -/// YAML and constructs empty per-family tables. The first query on each -/// family triggers the CSV read. -pub struct PerfDatabase { +/// The loaded per-family perf tables for one `//` +/// tuple — the mode-independent, immutable-after-load data half of +/// [`PerfDatabase`]. Shared (via `Arc`) between mode views. +pub struct PerfTables { pub system: String, pub backend: String, pub version: String, @@ -104,6 +102,23 @@ pub struct PerfDatabase { pub wideep_mla: WideEpMlaTable, pub wideep_moe: WideEpMoeTable, pub state_space: StateSpaceTable, +} + +/// Modular performance database for a specific +/// `//` tuple. +/// +/// `load` does the cheap work: resolves the data directory from the system +/// YAML and constructs empty per-family tables. The first query on each +/// family triggers the CSV read. +/// +/// Structurally this is a mode-configured VIEW over shared [`PerfTables`] +/// (mirroring Python's configured query views): the tables deref through, so +/// `db.gemm` / `db.system_spec` read as before, while `database_mode` / +/// `transfer_policy` are per-view. [`PerfDatabase::silicon_view`] derives the +/// SILICON view `FallbackOp` (and MSA's cross-op DSA probe) evaluate against +/// under HYBRID. +pub struct PerfDatabase { + tables: Arc, /// Query mode (Python's `database._default_database_mode`). SILICON /// queries collected tables only; HYBRID falls back to the util-space /// empirical layer on a typed silicon miss; EMPIRICAL always answers @@ -114,12 +129,20 @@ pub struct PerfDatabase { pub transfer_policy: TransferPolicy, /// Memo of built util-calibration grids, keyed by op/slice identity /// (see `operators::util_empirical`). Tables are immutable after load and - /// mode/policy are fixed per database instance, so the cache never needs - /// invalidation. - pub util_grids: UtilGridCache, + /// grid keys name their slice, so the memo is shared across mode views + /// and never needs invalidation. + pub util_grids: Arc, /// Memo of zero-aware delta lookups (the `compute_scale` empirical /// mechanism); same keying/lifetime rationale as `util_grids`. - pub delta_lookups: DeltaLookupCache, + pub delta_lookups: Arc, +} + +impl std::ops::Deref for PerfDatabase { + type Target = PerfTables; + + fn deref(&self) -> &PerfTables { + &self.tables + } } impl PerfDatabase { @@ -178,7 +201,7 @@ impl PerfDatabase { .oneccl_version .as_ref() .map(|v| system_data_root.join("oneccl").join(v)); - Ok(Self { + let tables = PerfTables { system: system.to_string(), backend: backend.to_string(), version: version.to_string(), @@ -211,10 +234,13 @@ impl PerfDatabase { ), system_spec: spec, data_root, + }; + Ok(Self { + tables: Arc::new(tables), database_mode: DatabaseMode::default(), transfer_policy: TransferPolicy::ALL, - util_grids: UtilGridCache::new(), - delta_lookups: DeltaLookupCache::new(), + util_grids: Arc::new(UtilGridCache::new()), + delta_lookups: Arc::new(DeltaLookupCache::new()), }) } @@ -226,6 +252,32 @@ impl PerfDatabase { self.transfer_policy = transfer_policy; self } + + /// Test-only mutable access to the shared tables (panics if the view has + /// been cloned — synthetic-table injection must happen before any + /// `silicon_view`). Production code never mutates loaded tables. + #[cfg(test)] + pub(crate) fn tables_mut(&mut self) -> &mut PerfTables { + Arc::get_mut(&mut self.tables).expect("tables Arc must be unique for test mutation") + } + + /// The SILICON view over the same loaded tables (cheap: `Arc` clones). + /// + /// Mirrors Python `FallbackOp.query`'s + /// `_get_configured_database_view(database, SILICON, transfer_policy)`: + /// under HYBRID the primary op is evaluated silicon-only so the module + /// table's absence falls to the granular fallback chain instead of being + /// hybrid-estimated at module level. Also used by MSA's cross-op DSA + /// utilisation probe. + pub fn silicon_view(&self) -> PerfDatabase { + PerfDatabase { + tables: Arc::clone(&self.tables), + database_mode: DatabaseMode::Silicon, + transfer_policy: self.transfer_policy, + util_grids: Arc::clone(&self.util_grids), + delta_lookups: Arc::clone(&self.delta_lookups), + } + } } #[cfg(test)] diff --git a/rust/aiconfigurator-core/src/perf_database/moe.rs b/rust/aiconfigurator-core/src/perf_database/moe.rs index 9e68be6ba..4e05addae 100644 --- a/rust/aiconfigurator-core/src/perf_database/moe.rs +++ b/rust/aiconfigurator-core/src/perf_database/moe.rs @@ -79,6 +79,30 @@ pub struct MoeTable { moe: OnceLock>, } +/// Which kernel grid a MoE accessor addresses: the default table or the +/// TRT-LLM `moe_torch_flow_min_latency` low-latency split (Python's +/// `_moe_data` vs `_moe_low_latency_data`). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MoeKernel { + Standard, + LowLatency, +} + +/// One collected sibling slice of the MoE table for a fixed +/// `(quant, distribution, moe_tp, moe_ep)`: the categorical shape features +/// plus its `num_tokens -> latency_ms` curve. Consumed by the operator +/// layer's cross-shape/cross-quant transfer ladder (the algorithm lives in +/// `operators/moe.rs`; this is a data accessor payload only). +#[derive(Clone, Debug)] +pub struct MoeSiblingSlice { + pub topk: u32, + pub num_experts: u32, + pub hidden_size: u32, + pub inter_size: u32, + /// `(num_tokens, latency_ms)` in ascending token order. + pub points: Vec<(u32, f64)>, +} + /// Two parallel grids split by `kernel_source`. Mirrors Python's split in /// `aiconfigurator.sdk.operations.moe.MoE.load_data`, where rows tagged /// `kernel_source == "moe_torch_flow_min_latency"` route to a separate @@ -257,6 +281,109 @@ impl MoeTable { Ok(!loaded.low_latency.by_keys.is_empty()) } + /// Own-slice `num_tokens -> latency_ms` curve for a full MoE key, after + /// the per-quant `"uniform"` distribution fallback. A typed miss + /// (`AicError::PerfDatabase`) means the slice is absent or empty — + /// mirroring Python `util_empirical.require_data_slice` as used by the + /// empirical own-shape grid (`_slice`) and the low-latency table probe + /// (`_moe_table`) in `MoE._query_moe_table`. + #[allow(clippy::too_many_arguments)] + pub fn slice_points( + &self, + kernel: MoeKernel, + quant_name: &str, + workload_distribution: &str, + topk: u32, + num_experts: u32, + hidden_size: u32, + inter_size: u32, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { + let grids = self.grids_for(kernel)?; + let dist = self.resolve_distribution(grids, quant_name, workload_distribution); + let key = MoeKey { + quant: quant_name.to_string(), + distribution: dist, + topk, + num_experts, + hidden_size, + inter_size, + moe_tp_size, + moe_ep_size, + }; + let by_tokens = grids.by_keys.get(&key).filter(|curve| !curve.is_empty()).ok_or_else(|| { + AicError::PerfDatabase(format!( + "MoE data missing for {key:?} ({kernel:?}) at {}", + self.data_root.display() + )) + })?; + Ok(by_tokens.iter().map(|(&t, &lat)| (t, lat)).collect()) + } + + /// All collected sibling slices for `(quant, distribution-after-uniform- + /// fallback, moe_tp, moe_ep)`; empty curves skipped, an empty result is + /// data (not an error). Mirrors the enumeration in Python `_collect` + /// (`MoE._query_moe_table`), which walks the nested + /// `topk -> num_experts -> hidden -> inter` dicts. NOTE: Python yields + /// dict insertion (file row) order; the `BTreeMap` yields sorted + /// `(topk, num_experts, hidden, inter)` order instead — observable only + /// through exact ties in nearest-candidate selection. + pub fn sibling_slices( + &self, + kernel: MoeKernel, + quant_name: &str, + workload_distribution: &str, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { + let grids = self.grids_for(kernel)?; + let dist = self.resolve_distribution(grids, quant_name, workload_distribution); + let mut slices = Vec::new(); + for (key, curve) in &grids.by_keys { + if key.quant != quant_name + || key.distribution != dist + || key.moe_tp_size != moe_tp_size + || key.moe_ep_size != moe_ep_size + || curve.is_empty() + { + continue; + } + slices.push(MoeSiblingSlice { + topk: key.topk, + num_experts: key.num_experts, + hidden_size: key.hidden_size, + inter_size: key.inter_size, + points: curve.iter().map(|(&t, &lat)| (t, lat)).collect(), + }); + } + Ok(slices) + } + + /// Distinct quant names present in the kernel grid, in sorted (BTreeMap) + /// order. Python iterates the table dict in insertion (file row) order; + /// the difference is observable only through exact ties in the transfer + /// ladder's profile-distance sort / nearest-candidate selection. + pub fn available_quants(&self, kernel: MoeKernel) -> Result, AicError> { + let grids = self.grids_for(kernel)?; + let mut names: Vec = Vec::new(); + for key in grids.by_keys.keys() { + // `MoeKey` sorts by quant first, so consecutive dedup suffices. + if names.last().map(String::as_str) != Some(key.quant.as_str()) { + names.push(key.quant.clone()); + } + } + Ok(names) + } + + fn grids_for(&self, kernel: MoeKernel) -> Result<&MoeGrids, AicError> { + let loaded = self.load()?; + Ok(match kernel { + MoeKernel::Standard => &loaded.default, + MoeKernel::LowLatency => &loaded.low_latency, + }) + } + /// Mirrors Python's: /// `dist = workload if workload in moe_data[quant] else "uniform"` fn resolve_distribution( diff --git a/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs b/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs index 3b266e5e6..54a30c857 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs @@ -160,7 +160,11 @@ impl WideEpMlaTable { // reads it (memory scales by fmha.memory). let _ = kv_quant; let spec = &self.system_spec; - let sol = move |c: &[f64]| wideep_context_mla_sol_ms(spec, fmha_quant, c[0], c[1], c[2]); + // Silicon sol_fn: `tp = 128 // n` then `num_head = 128 // tp` + // (Python `get_silicon`'s lambda), prefix = 0 (samples are prefix=0). + let sol = move |c: &[f64]| { + wideep_context_mla_sol_ms(spec, fmha_quant, wideep_num_head(c[0]), c[1], 0.0, c[2]) + }; let cfg = OpInterpConfig::grid_sqrt_axis(CONTEXT_AXES, 1, &sol); perf_interp::query( &cfg, @@ -209,7 +213,10 @@ impl WideEpMlaTable { FmhaQuantMode::Bfloat16 }; let spec = &self.system_spec; - let sol = move |c: &[f64]| wideep_generation_mla_sol_ms(spec, fmha_quant, c[0], c[1], c[2]); + // Silicon sol_fn: `tp = 128 // n` then `num_head = 128 // tp`. + let sol = move |c: &[f64]| { + wideep_generation_mla_sol_ms(spec, fmha_quant, wideep_num_head(c[0]), c[1], c[2]) + }; let cfg = OpInterpConfig::grid(GENERATION_AXES, &sol); perf_interp::query( &cfg, @@ -218,6 +225,54 @@ impl WideEpMlaTable { ) } + // ----------------------------------------------------------------------- + // Point accessors for the util-space empirical layer (algorithm-free: + // typed `AicError::PerfDatabase` miss on absent slice / empty node, no + // estimation logic). Coordinate order matches the Python `depth=3` + // iteration of each `require_data_slice` slice. + // ----------------------------------------------------------------------- + + /// Collected `(num_heads, seq, batch) -> latency` points of the + /// `(kernel_source, fmha, kv)` context slice. Typed miss when + /// absent/empty. + pub fn context_points( + &self, + kernel_source: &str, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + ) -> Result, f64)>, AicError> { + let grids = self.load_context()?; + let key = ContextKey { + kernel_source: kernel_source.to_string(), + fmha_quant: fmha_quant.name().to_string(), + kv_quant: kv_quant.name().to_string(), + }; + let node = grids + .by_keys + .get(&key) + .ok_or_else(|| missing("WideEP context MLA", &self.data_root, format!("{key:?}")))?; + non_empty_points(node, "WideEP context MLA", &self.data_root) + } + + /// Collected `(num_heads, batch, seq) -> latency` points of the + /// `(kernel_source, kv)` generation slice. Typed miss when absent/empty. + pub fn generation_points( + &self, + kernel_source: &str, + kv_quant: KvCacheQuantMode, + ) -> Result, f64)>, AicError> { + let grids = self.load_generation()?; + let key = GenerationKey { + kernel_source: kernel_source.to_string(), + kv_quant: kv_quant.name().to_string(), + }; + let node = grids + .by_keys + .get(&key) + .ok_or_else(|| missing("WideEP generation MLA", &self.data_root, format!("{key:?}")))?; + non_empty_points(node, "WideEP generation MLA", &self.data_root) + } + fn load_context(&self) -> Result<&WideEpContextMlaGrids, AicError> { let cell = self .context @@ -246,30 +301,42 @@ fn bf16_tc_flops(spec: &SystemSpec) -> f64 { spec.gpu.bfloat16_tc_flops.unwrap_or(0.0) } -/// The tables key by `num_heads`; the Python SOLs take `tp_size` and derive -/// `num_head = 128 // tp_size`, with the sol_fn lambdas mapping -/// `tp = 128 // n`. Compose both floor divisions exactly. -fn wideep_num_head(n: f64) -> f64 { +/// The tables key by `num_heads`; the Python SILICON `sol_fn` lambdas take +/// `tp_size` derived as `tp = 128 // n` and the SOLs then use +/// `num_head = 128 // tp_size`. Compose both floor divisions exactly. (The +/// util-empirical sample mapping differs — Python rounds there: +/// `tp = round(128 / n)`; see `operators/wideep_mla.rs`.) +pub(crate) fn wideep_num_head(n: f64) -> f64 { let tp_size = (128.0 / n).floor(); (128.0 / tp_size).floor() } -/// WideEP context MLA SOL in ms at prefix = 0. Structure (per Python): +/// WideEP context MLA SOL in ms. `num_head` is the per-rank head count +/// (Python's `128 // tp_size`; the `n -> num_head` mapping lives at the +/// call sites because silicon and empirical map differently). `s` is the +/// chunk / isl length; silicon sol_fns pass `prefix = 0` (samples are +/// prefix=0), the util-empirical query SOL carries the real prefix. +/// Structure (per Python): /// - q_b / kv_b projections + attention output projection -> `ops` /// (divided by `bf16_tc_flops * fmha.compute`) -/// - attention flops `2 * nh * (nope*2 + rope) * b * full_s^2 // 2` added at -/// full bf16 throughput (no fmha compute factor) +/// - attention flops `2 * nh * (nope*2 + rope) * b * (full_s^2 - prefix^2) // 2` +/// added at full bf16 throughput (no fmha compute factor) /// - `mem = (q_b_mem + kv_b_mem + attn_mem * 2 + attn_out_mem) * fmha.memory` /// - `sol = max(sol_math, sol_mem)` -fn wideep_context_mla_sol_ms(spec: &SystemSpec, fmha_quant: FmhaQuantMode, n: f64, s: f64, b: f64) -> f64 { +pub(crate) fn wideep_context_mla_sol_ms( + spec: &SystemSpec, + fmha_quant: FmhaQuantMode, + num_head: f64, + s: f64, + prefix: f64, + b: f64, +) -> f64 { let hidden_size = 7168.0_f64; let q_lora_rank = 1536.0_f64; let kv_lora_rank = 512.0_f64; let qk_rope_head_dim = 64.0_f64; let qk_nope_head_dim = 128.0_f64; let v_head_dim = 128.0_f64; - let num_head = wideep_num_head(n); - let prefix = 0.0_f64; // q_b projection let q_b_flop = 2.0 * q_lora_rank * num_head * (qk_rope_head_dim + qk_nope_head_dim) * b * s; @@ -311,20 +378,27 @@ fn wideep_context_mla_sol_ms(spec: &SystemSpec, fmha_quant: FmhaQuantMode, n: f6 sol_math.max(sol_mem) } -/// WideEP generation MLA SOL in ms. Structure (per Python): q_b, q_w_kc, +/// WideEP generation MLA SOL in ms. `num_head` is the per-rank head count +/// (see [`wideep_context_mla_sol_ms`] for the mapping split between +/// silicon and empirical call sites). Structure (per Python): q_b, q_w_kc, /// s_w_vc and attention-output projections -> `ops` (divided by /// `bf16_tc_flops * fmha.compute`); the MQA attention flops /// `2 * b * s * nh * (rope + kv_lora*2)` added at full bf16 throughput; /// `mem = (q_b + q_w_kc + attn*2 + s_w_vc + attn_out) * fmha.memory`; /// `sol = max(sol_math, sol_mem)`. -fn wideep_generation_mla_sol_ms(spec: &SystemSpec, fmha_quant: FmhaQuantMode, n: f64, b: f64, s: f64) -> f64 { +pub(crate) fn wideep_generation_mla_sol_ms( + spec: &SystemSpec, + fmha_quant: FmhaQuantMode, + num_head: f64, + b: f64, + s: f64, +) -> f64 { let hidden_size = 7168.0_f64; let q_lora_rank = 1536.0_f64; let kv_lora_rank = 512.0_f64; let qk_rope_head_dim = 64.0_f64; let qk_nope_head_dim = 128.0_f64; let v_head_dim = 128.0_f64; - let num_head = wideep_num_head(n); // NOTE: qkv_a projection is modeled as a standalone GEMM op // (generation_qkv_a_proj_gemm) outside the MLA attention forward path, @@ -503,6 +577,24 @@ fn missing(table: &str, data_root: &Path, descriptor: String) -> AicError { )) } +/// Flatten a slice node into `(coords, latency)` points, treating an empty +/// node as a typed coverage miss (mirrors `require_data_slice`'s empty-node +/// check). +fn non_empty_points( + node: &Node, + table: &str, + data_root: &Path, +) -> Result, f64)>, AicError> { + let points = perf_interp::node_points(node); + if points.is_empty() { + return Err(AicError::PerfDatabase(format!( + "{table} perf data empty for the requested slice at {}", + data_root.display() + ))); + } + Ok(points) +} + fn clone_err(err: &AicError) -> AicError { AicError::PerfDatabase(err.to_string()) } From b72456f96810934b97a31fcdadfd6cc836088480 Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 13:53:09 +0800 Subject: [PATCH 06/10] =?UTF-8?q?feat(rust-core):=20MSA=20op=20(XOP=20tran?= =?UTF-8?q?sfer)=20+=20dispatch/wideep=20hybrid=20=E2=80=94=20hybrid=20par?= =?UTF-8?q?ity=20green?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the util-space empirical port (issue #1333 §4.6 option b): - operators/msa.rs: MiniMax Sparse Attention module (MiniMax-M3). No silicon data by design — SILICON raises the perf-data miss; HYBRID/EMPIRICAL run the cross-op (XOP) transfer borrowing DSA's measured utilisation via a silicon_view probe (SOL_dsa / silicon_dsa), policy-gated, with the manual dsa_scale_k level alignment. The verbatim _msa_attention_sol port (three-group GQA/indexer/sparse-attention SOL, i128 pair counts) anchors at 1e-9 against Python on both sglang 0.5.14 and vllm 0.19.0 GLM-architecture DSA rows. - Op::MsaContext/MsaGeneration variants + opspec conversion in engine.py (MiniMax-M3 now compiles). ENGINE_SPEC_SCHEMA_VERSION bumped to 3 on both sides (bincode enum indices after DsaGeneration shifted). - moe_dispatch/wideep_moe: deepep ll/normal EMPIRICAL gates (Python has no empirical for these — EMPIRICAL raises, HYBRID stays silicon), the full trtllm-alltoall table mirror (kernel/op_name/num_nodes layout, NotEnabled -> 0.0, kernel selection) with own-slice empirical, and the WideEpMoe compute dispatch with kernel/distribution fallback mirrors + the num_slots-aware roofline anchoring the silicon token curve. Python's latent HYBRID-alltoall-fallback TypeError (moe.py:2266 closure omits kernel_source) is documented; Rust implements the intended fallback. - test_rust_engine_step gate test updated to the new contract: SILICON / HYBRID / EMPIRICAL route to the compiled engine, SOL modes delegate. Full verification on this commit: - hybrid parity suite: 64/64 at rtol=1e-4 (was 46 failed at the red baseline) — every transfer tier pinned by real shipped data - existing SILICON parity: 231/231; compile-engine parity: 47/47; perf gate: 3/3; whole parity file: 295/295 - Python unit suite: 2071 passed - cargo test --lib: 254 passed, 5 pre-existing failures (stale parse_b200_sxm constant + 4 py::tests needing a linked interpreter) Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- rust/aiconfigurator-core/src/config.rs | 4 +- rust/aiconfigurator-core/src/engine/spec.rs | 25 + rust/aiconfigurator-core/src/operators/mod.rs | 2 + .../src/operators/moe_dispatch.rs | 504 +++++++++++++++++- rust/aiconfigurator-core/src/operators/msa.rs | 445 ++++++++++++++++ rust/aiconfigurator-core/src/operators/op.rs | 10 + .../src/operators/wideep_moe.rs | 311 ++++++++++- .../src/perf_database/wideep.rs | 214 +++++++- .../src/perf_database/wideep_moe.rs | 228 +++++--- src/aiconfigurator/sdk/engine.py | 36 +- tests/unit/sdk/test_rust_engine_step.py | 18 +- 11 files changed, 1673 insertions(+), 124 deletions(-) create mode 100644 rust/aiconfigurator-core/src/operators/msa.rs diff --git a/rust/aiconfigurator-core/src/config.rs b/rust/aiconfigurator-core/src/config.rs index 03f00c6f5..c48509773 100644 --- a/rust/aiconfigurator-core/src/config.rs +++ b/rust/aiconfigurator-core/src/config.rs @@ -21,7 +21,9 @@ pub const ENGINE_CONFIG_SCHEMA_VERSION: u32 = 1; // only distinguishable by this version — `EngineSpec::from_bincode` reads and // checks it before decoding the op lists. Bump whenever an `OpSpec` field // changes; keep in lockstep with `sdk/engine.py::ENGINE_SPEC_SCHEMA_VERSION`. -pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 2; +// Bumped to 3 when the `Msa{Context,Generation}` variants were inserted +// (bincode enum indices after `DsaGeneration` shifted). +pub const ENGINE_SPEC_SCHEMA_VERSION: u32 = 3; /// Static engine identity and setup information carried by an /// [`crate::engine::spec::EngineSpec`]. diff --git a/rust/aiconfigurator-core/src/engine/spec.rs b/rust/aiconfigurator-core/src/engine/spec.rs index 181c7034b..303a08424 100644 --- a/rust/aiconfigurator-core/src/engine/spec.rs +++ b/rust/aiconfigurator-core/src/engine/spec.rs @@ -382,6 +382,27 @@ mod tests { } } + fn msa_module() -> crate::operators::MsaModuleOp { + crate::operators::MsaModuleOp { + name: "context_attention".into(), + scale_factor: 62.0, + num_heads: 8, + num_kv_heads: 1, + hidden_size: 7168, + head_dim: 128, + v_head_dim: 128, + index_n_heads: 64, + index_head_dim: 128, + index_topk: 2048, + block_size: 64, + kv_cache_dtype: KvCacheQuantMode::Bfloat16, + fmha_quant_mode: FmhaQuantMode::Bfloat16, + gemm_quant_mode: GemmQuantMode::Fp8Block, + dsa_architecture: "GlmMoeDsaForCausalLM".into(), + dsa_scale_k: 1.0, + } + } + fn dsv4_module() -> Dsv4ModuleOp { Dsv4ModuleOp { name: "dsv4_module".into(), @@ -540,6 +561,8 @@ mod tests { OpSpec::Vision(vision()), OpSpec::DsaContext(dsa_module()), OpSpec::DsaGeneration(dsa_module()), + OpSpec::MsaContext(msa_module()), + OpSpec::MsaGeneration(msa_module()), OpSpec::Dsv4Context(dsv4_module()), OpSpec::Dsv4Generation(dsv4_module()), OpSpec::Mhc(mhc()), @@ -575,6 +598,8 @@ mod tests { | OpSpec::Vision(_) | OpSpec::DsaContext(_) | OpSpec::DsaGeneration(_) + | OpSpec::MsaContext(_) + | OpSpec::MsaGeneration(_) | OpSpec::Dsv4Context(_) | OpSpec::Dsv4Generation(_) | OpSpec::Mhc(_) diff --git a/rust/aiconfigurator-core/src/operators/mod.rs b/rust/aiconfigurator-core/src/operators/mod.rs index 0bae8d64d..4a7dddb81 100644 --- a/rust/aiconfigurator-core/src/operators/mod.rs +++ b/rust/aiconfigurator-core/src/operators/mod.rs @@ -26,6 +26,7 @@ pub mod mhc; pub mod mla; pub mod moe; pub mod moe_dispatch; +pub mod msa; pub mod op; pub mod overlap; pub mod util_empirical; @@ -46,6 +47,7 @@ pub use mhc::MhcModuleOp; pub use mla::{ContextMlaOp, GenerationMlaOp, MlaBmmOp, MlaModuleOp}; pub use moe::MoeOp; pub use moe_dispatch::{DispatchFlavor, MoEDispatchOp}; +pub use msa::MsaModuleOp; pub use op::{FallbackOp, Op, OverlapOp, RuntimeContext}; pub use overlap::overlap_composition; pub use vision::VisionEncoderOp; diff --git a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs index a80aa1843..f27d98f1c 100644 --- a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs +++ b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs @@ -21,11 +21,12 @@ use serde::{Deserialize, Serialize}; -use crate::common::enums::{BackendKind, CommQuantMode, MoeQuantMode}; +use crate::common::enums::{BackendKind, CommQuantMode, DatabaseMode, MoeQuantMode}; use crate::common::error::AicError; use crate::common::system_spec::SystemSpec; use crate::operators::base::{PerformanceResult, Source}; use crate::operators::communication::{CustomAllReduceOp, NcclOp}; +use crate::operators::util_empirical::{self, UtilGrid}; use crate::perf_database::PerfDatabase; /// MoE dispatch flavor. @@ -296,6 +297,19 @@ impl MoEDispatchOp { .scaled(self.scale_factor)) } DispatchFlavor::DeepEpNormal => { + // Python `_query_wideep_deepep_normal_table` has NO empirical + // path: EMPIRICAL mode raises (`NotImplementedError("WideEP + // deepep normal operation's empirical is not implemented + // yet")`), and HYBRID goes STRAIGHT to the silicon interp with + // no fallback (the method never routes through + // `_query_silicon_or_hybrid`) — so a silicon miss under HYBRID + // must propagate unchanged, not convert into an estimate. + if db.database_mode == DatabaseMode::Empirical { + return Err(AicError::EmpiricalNotImplemented( + "WideEP deepep normal operation's empirical is not implemented yet" + .to_string(), + )); + } let point = db.wideep.query_deepep_normal( self.deepep_node_num(spec)?, self.hidden_size, @@ -323,6 +337,13 @@ impl MoEDispatchOp { .scaled(self.scale_factor)) } DispatchFlavor::DeepEpLowLatency => { + // Same no-empirical rule as DeepEpNormal (Python + // `_query_wideep_deepep_ll_table`). + if db.database_mode == DatabaseMode::Empirical { + return Err(AicError::EmpiricalNotImplemented( + "WideEP deepep ll operation's empirical is not implemented yet".to_string(), + )); + } let point = db.wideep.query_deepep_ll( self.deepep_node_num(spec)?, self.hidden_size, @@ -367,18 +388,31 @@ impl MoEDispatchOp { let attention_tp = self.attention_tp_size(); if enable_alltoall { - let latency = db.wideep.query_trtllm_alltoall( + // Python queries op_name="alltoall_dispatch" (pre) / + // "alltoall_combine" (post) through + // `database.query_trtllm_alltoall` with `moe_backend = + // self._moe_backend` (None in all current callers) and + // `node_num=None` (computed from moe_ep_size inside the + // table method). `float(result)` + tail scaling == + // `.scaled(scale_factor)` here. + let op_name = if self.pre_dispatch { + "alltoall_dispatch" + } else { + "alltoall_combine" + }; + let result = query_alltoall_table( + db, + op_name, num_tokens, self.hidden_size, self.topk, self.num_experts, self.moe_ep_size, self.moe_quant, - "uniform", + None, + None, )?; - Ok(PerformanceResult::new(latency, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) + Ok(result.clamp_non_negative().scaled(self.scale_factor)) } else if self.attention_dp_size > 1 { // Python: query_nccl(half, num_gpus, "all_gather" or // "reduce_scatter", volume * attention_dp_size). @@ -415,6 +449,261 @@ impl MoEDispatchOp { } } +// --------------------------------------------------------------------------- +// TRT-LLM alltoall table (Python +// `TrtLLMWideEPMoEDispatch._query_alltoall_table`, also reached from +// `MoEDispatch.query`'s trtllm SM100 branch via +// `database.query_trtllm_alltoall`). Own-shape empirical only — no transfer +// ladder. +// --------------------------------------------------------------------------- + +/// Mirror of `TrtLLMWideEPMoEDispatch._normalize_quant_mode_for_table`: +/// `fp8_block` is a behavioral mode that reuses the `fp8` alltoall tables. +fn normalize_alltoall_quant_for_table(quant: MoeQuantMode) -> MoeQuantMode { + if quant == MoeQuantMode::Fp8Block { + MoeQuantMode::Fp8 + } else { + quant + } +} + +/// Mirror of `TrtLLMWideEPMoEDispatch._select_alltoall_kernel` (aligned with +/// TRT-LLM's per-backend `select_alltoall_method_type`). Pure logic: the +/// Python tail's data-availability check only LOGS a warning and returns the +/// preferred kernel anyway, so it has no behavioral counterpart here. +/// (Python's `quant_mode` parameter is unused by the selection — dropped.) +fn select_alltoall_kernel( + spec: &SystemSpec, + moe_ep_size: u32, + topk: u32, + moe_backend: Option<&str>, +) -> &'static str { + if let Some(backend) = moe_backend { + let upper = backend.to_uppercase(); + if upper == "DEEPGEMM" || upper == "CUTE_DSL" { + return "NotEnabled"; + } + } + let sm_version = spec.gpu.sm_version.unwrap_or(0); + let num_gpus_per_node = spec.node.num_gpus_per_node; + let is_inter_node = moe_ep_size > num_gpus_per_node; + let is_wideep = moe_backend.is_some_and(|b| b.to_uppercase() == "WIDEEP"); + // Python approximates supports_mnnvl() as SM >= 100. + let supports_mnnvl = sm_version >= 100; + + if is_wideep { + if supports_mnnvl { + "NVLinkTwoSided" + } else { + let deepep_feasible = moe_ep_size > 1 && topk <= 8; + if deepep_feasible && is_inter_node { + "DeepEP" + } else if deepep_feasible { + "DeepEPLowLatency" + } else { + "NotEnabled" + } + } + } else if supports_mnnvl { + "NVLinkOneSided" + } else { + "NotEnabled" + } +} + +/// Alltoall communication SOL (ms), mirroring `_query_alltoall_table.get_sol`: +/// - prepare: lightweight metadata exchange (`topk * 4` bytes per token); +/// - combine: bfloat16 results (2 B/elem), or fp4 (0.5 B/elem) for the +/// low-precision variant; +/// - dispatch: per-rank deduplication at the RAW quant's precision. +/// `remote_ranks = min(topk, num_experts, ep - 1)`; bandwidth is inter-node +/// when the group spans more than one node. Linear in `num_tokens` with zero +/// intercept — coordinates arrive as f64 from the util-grid engine, and the +/// math is float-exact against Python's (no floor division here). +#[allow(clippy::too_many_arguments)] +fn alltoall_sol_ms( + spec: &SystemSpec, + op_name: &str, + quant: MoeQuantMode, + node_num: u32, + num_tokens: f64, + hidden_size: u32, + topk: u32, + num_experts: u32, + moe_ep_size: u32, +) -> f64 { + let is_inter_node = node_num > 1; + let bw = if is_inter_node { + spec.node.inter_node_bw + } else { + spec.node.intra_node_bw + }; + let remote_ranks = topk.min(num_experts).min(moe_ep_size.saturating_sub(1)) as f64; + let data_bytes = if op_name == "alltoall_prepare" { + num_tokens * topk as f64 * 4.0 // token routing indices, ~4 bytes each + } else if op_name.contains("combine") { + let bytes_per_element = if op_name.contains("low_precision") { 0.5 } else { 2.0 }; + num_tokens * remote_ranks * hidden_size as f64 * bytes_per_element + } else { + // dispatch: per-rank deduplication, use quant_mode precision + num_tokens * remote_ranks * hidden_size as f64 * quant.mapping().memory + }; + data_bytes / bw * 1000.0 +} + +/// Verbatim port of `TrtLLMWideEPMoEDispatch._query_alltoall_table` (minus +/// the SOL diagnostic modes, which never reach the compiled engine): +/// normalize the table quant, default `node_num` from `moe_ep_size`, select +/// the kernel, early-return 0 for `NotEnabled`, then dispatch on the +/// database mode — EMPIRICAL always estimates `SOL(query)/util` from the +/// own-slice token grid; HYBRID converts a typed silicon miss into the +/// estimate; SILICON queries the table. +/// +/// KNOWN PYTHON DIVERGENCE: Python's HYBRID fallback closure calls +/// `get_empirical_from_sol` without its `kernel_source` argument and raises +/// `TypeError` (a latent bug — the branch has no working oracle). This port +/// implements the intended fallback (kernel threaded through); when the +/// slice has no data anywhere both languages still end in a typed failure +/// (`EmpiricalNotImplemented` here, `TypeError` there). +#[allow(clippy::too_many_arguments)] +fn query_alltoall_table( + db: &PerfDatabase, + op_name: &str, + num_tokens: u32, + hidden_size: u32, + topk: u32, + num_experts: u32, + moe_ep_size: u32, + quant: MoeQuantMode, + node_num: Option, + moe_backend: Option<&str>, +) -> Result { + let table_quant = normalize_alltoall_quant_for_table(quant); + + // Python: `node_num = 1 if moe_ep_size < 4 else moe_ep_size // 4` when + // not provided (no Rust caller provides one today, matching + // `MoEDispatch.query` which never passes node_num). + let node_num = node_num.unwrap_or(if moe_ep_size < 4 { 1 } else { moe_ep_size / 4 }); + + const VALID_OP_NAMES: [&str; 4] = [ + "alltoall_prepare", + "alltoall_dispatch", + "alltoall_combine", + "alltoall_combine_low_precision", + ]; + if !VALID_OP_NAMES.contains(&op_name) { + // Python raises ValueError — a programming error, deliberately NOT a + // missing-data signal (must not trigger HYBRID/fallback handling). + return Err(AicError::InvalidEngineConfig(format!( + "Invalid op_name '{op_name}'. Must be one of {VALID_OP_NAMES:?}" + ))); + } + + let kernel_source = select_alltoall_kernel(&db.system_spec, moe_ep_size, topk, moe_backend); + if kernel_source == "NotEnabled" { + // Python returns PerformanceResult(0.0, source="empirical") for every + // non-SOL mode. + return Ok(PerformanceResult::new(0.0, Source::Empirical)); + } + + let silicon = || { + db.wideep.query_trtllm_alltoall( + kernel_source, + op_name, + table_quant, + node_num, + hidden_size, + topk, + num_experts, + moe_ep_size, + num_tokens, + ) + }; + let empirical = || { + alltoall_empirical( + db, + kernel_source, + op_name, + quant, + table_quant, + node_num, + num_tokens, + hidden_size, + topk, + num_experts, + moe_ep_size, + ) + }; + match db.database_mode { + DatabaseMode::Empirical => Ok(PerformanceResult::new(empirical()?, Source::Empirical)), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok(PerformanceResult::new(latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => { + Ok(PerformanceResult::new(empirical()?, Source::Empirical)) + } + Err(err) => Err(err), + }, + _ => Ok(PerformanceResult::new(silicon()?, Source::Silicon)), + } +} + +/// `SOL(query)/util` over the own-slice token curve (depth 1, no transfer +/// ladder). Mirrors `_query_alltoall_table::get_empirical_from_sol`: the SOL +/// uses the RAW quant (only the dispatch op consults its memory width), the +/// table slice uses the NORMALIZED quant. +#[allow(clippy::too_many_arguments)] +fn alltoall_empirical( + db: &PerfDatabase, + kernel_source: &str, + op_name: &str, + quant: MoeQuantMode, + table_quant: MoeQuantMode, + node_num: u32, + num_tokens: u32, + hidden_size: u32, + topk: u32, + num_experts: u32, + moe_ep_size: u32, +) -> Result { + let spec = &db.system_spec; + let sol = |c: &[f64]| { + alltoall_sol_ms(spec, op_name, quant, node_num, c[0], hidden_size, topk, num_experts, moe_ep_size) + }; + let sol_time = sol(&[num_tokens as f64]); + + // Python keys on ("alltoall", system, backend, version, kernel, op_name, + // tqm.name, node_num, hidden, topk, experts, ep) + id(node); the cache + // here is per-database. + let key = format!( + "alltoall:{kernel_source}:{op_name}:{}:{node_num}:{hidden_size}:{topk}:{num_experts}:{moe_ep_size}", + table_quant.name(), + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.wideep.alltoall_slice_points( + kernel_source, + op_name, + table_quant, + node_num, + hidden_size, + topk, + num_experts, + moe_ep_size, + ) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples( + points.into_iter().map(|(t, lat)| (vec![t as f64], lat)), + sol, + )))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [num_tokens as f64]; + let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + Ok(latency) +} + #[cfg(test)] mod tests { use super::*; @@ -723,4 +1012,207 @@ mod tests { combine.latency_ms ); } + + // ----------------------------------------------------------------- + // HYBRID / EMPIRICAL parity (util-space empirical port). + // ----------------------------------------------------------------- + + use crate::common::enums::TransferPolicy; + + fn gb200_trtllm_db(mode: DatabaseMode) -> PerfDatabase { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + PerfDatabase::load(&root, "gb200", "trtllm", "1.3.0rc10") + .expect("db loads") + .with_mode(mode, TransferPolicy::ALL) + } + + fn h100_sglang_db(mode: DatabaseMode) -> PerfDatabase { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + PerfDatabase::load(&root, "h100_sxm", "sglang", "0.5.6.post2") + .expect("db loads") + .with_mode(mode, TransferPolicy::ALL) + } + + /// Shorthand for the collected gb200 alltoall shape (hidden=7168, + /// topk=8, experts=256). + fn a2a( + db: &PerfDatabase, + op_name: &str, + num_tokens: u32, + quant: MoeQuantMode, + moe_backend: Option<&str>, + ) -> Result { + query_alltoall_table(db, op_name, num_tokens, 7168, 8, 256, 8, quant, None, moe_backend) + } + + fn assert_oracle(result: &PerformanceResult, expected: f64, source: Source, label: &str) { + assert!( + (result.latency_ms - expected).abs() < 1e-9, + "{label}: expected {expected}, got {}", + result.latency_ms + ); + assert_eq!(result.source, source, "{label}: wrong source"); + } + + /// Oracle values generated from the Python reference on the same data + /// (shared layer pinned OFF so Python reads exactly the primary parquet): + /// + /// ```text + /// db = perf_database.get_database_view("gb200", "trtllm", "1.3.0rc10", + /// allow_missing_data=True, database_mode=..., shared_layer=False) + /// float(TrtLLMWideEPMoEDispatch._query_alltoall_table(db, op_name=..., + /// num_tokens=..., hidden_size=7168, topk=8, num_experts=256, + /// moe_ep_size=8, quant_mode=..., node_num=None, + /// database_mode=..., moe_backend=...)) + /// ``` + /// + /// `moe_backend=None` selects NVLinkOneSided (gb200 SM100, non-WideEP; + /// only nvfp4 collected there); `moe_backend="wideep"` selects + /// NVLinkTwoSided (bfloat16/fp8/nvfp4 + prepare/combine_low_precision). + /// ep=8 -> node_num = 8 // 4 = 2 (both computed and, in the loader, the + /// num-nodes-column default `max(1, ep // 4)`). nt=64 is a collected + /// point (exact hit); nt=333 is an off-grid interior point. + #[test] + fn alltoall_empirical_matches_python_oracles() { + let db = gb200_trtllm_db(DatabaseMode::Empirical); + let hit = a2a(&db, "alltoall_dispatch", 64, MoeQuantMode::Nvfp4, None).expect("exact hit"); + assert_oracle(&hit, 0.018886399269104005, Source::Empirical, "emp_dispatch_t64"); + let off = a2a(&db, "alltoall_dispatch", 333, MoeQuantMode::Nvfp4, None).expect("off-grid"); + assert_oracle(&off, 0.033548976838374114, Source::Empirical, "emp_dispatch_t333"); + let combine = a2a(&db, "alltoall_combine", 333, MoeQuantMode::Nvfp4, None).expect("combine"); + assert_oracle(&combine, 0.07118654040018774, Source::Empirical, "emp_combine_t333"); + + // WideEP kernel + the ops only that path uses. + let wideep_dispatch = + a2a(&db, "alltoall_dispatch", 333, MoeQuantMode::Fp8, Some("wideep")).expect("wideep"); + assert_oracle(&wideep_dispatch, 0.04266341407888801, Source::Empirical, "emp_wideep_fp8"); + let prepare = + a2a(&db, "alltoall_prepare", 333, MoeQuantMode::Fp8, Some("wideep")).expect("prepare"); + assert_oracle(&prepare, 0.015176159059580488, Source::Empirical, "emp_wideep_prepare"); + let combine_lp = + a2a(&db, "alltoall_combine_low_precision", 333, MoeQuantMode::Nvfp4, Some("wideep")) + .expect("combine_lp"); + assert_oracle(&combine_lp, 0.06946014693369004, Source::Empirical, "emp_wideep_combine_lp"); + } + + /// HYBRID with data present stays on silicon; the in-range interpolation + /// differs from the empirical reconstruction at the same points. + #[test] + fn alltoall_hybrid_prefers_silicon_when_covered() { + let db = gb200_trtllm_db(DatabaseMode::Hybrid); + let hit = a2a(&db, "alltoall_dispatch", 64, MoeQuantMode::Nvfp4, None).expect("exact hit"); + assert_oracle(&hit, 0.018886399269104005, Source::Silicon, "hyb_dispatch_t64"); + let off = a2a(&db, "alltoall_dispatch", 333, MoeQuantMode::Nvfp4, None).expect("off-grid"); + assert_oracle(&off, 0.033547499962151055, Source::Silicon, "hyb_dispatch_t333"); + let combine = a2a(&db, "alltoall_combine", 333, MoeQuantMode::Nvfp4, None).expect("combine"); + assert_oracle(&combine, 0.07116495203226805, Source::Silicon, "hyb_combine_t333"); + let wideep_dispatch = + a2a(&db, "alltoall_dispatch", 333, MoeQuantMode::Fp8, Some("wideep")).expect("wideep"); + assert_oracle(&wideep_dispatch, 0.042655749106779696, Source::Silicon, "hyb_wideep_fp8"); + let prepare = + a2a(&db, "alltoall_prepare", 333, MoeQuantMode::Fp8, Some("wideep")).expect("prepare"); + assert_oracle(&prepare, 0.015222300426103175, Source::Silicon, "hyb_wideep_prepare"); + let combine_lp = + a2a(&db, "alltoall_combine_low_precision", 333, MoeQuantMode::Nvfp4, Some("wideep")) + .expect("combine_lp"); + assert_oracle(&combine_lp, 0.069468751270324, Source::Silicon, "hyb_wideep_combine_lp"); + } + + /// fp8 is uncollected under NVLinkOneSided: EMPIRICAL surfaces the typed + /// empirical miss (own-shape only — no transfer ladder), and fp8_block + /// normalizes onto the same missing fp8 slice. Under HYBRID, Python's + /// fallback closure raises `TypeError` (it drops the `kernel_source` + /// argument — a latent bug, see `query_alltoall_table`); the Rust port + /// runs the intended fallback, which finds no data and surfaces the same + /// typed empirical miss. + #[test] + fn alltoall_missing_slice_is_typed_empirical_miss() { + let emp = gb200_trtllm_db(DatabaseMode::Empirical); + for quant in [MoeQuantMode::Fp8, MoeQuantMode::Fp8Block] { + let result = a2a(&emp, "alltoall_dispatch", 333, quant, None); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "EMPIRICAL {quant:?} must be a typed empirical miss, got {result:?}" + ); + } + let hyb = gb200_trtllm_db(DatabaseMode::Hybrid); + let result = a2a(&hyb, "alltoall_dispatch", 333, MoeQuantMode::Fp8, None); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "HYBRID fallback on a data-less slice must be a typed empirical miss, got {result:?}" + ); + } + + /// `_select_alltoall_kernel` mirror + the NotEnabled early return + /// (0.0, source "empirical", regardless of mode — before any table I/O). + #[test] + fn alltoall_kernel_selection_matches_python() { + let db = gb200_trtllm_db(DatabaseMode::Hybrid); + let spec = &db.system_spec; // SM100, 4 GPUs per node + assert_eq!(select_alltoall_kernel(spec, 8, 8, None), "NVLinkOneSided"); + assert_eq!(select_alltoall_kernel(spec, 8, 8, Some("wideep")), "NVLinkTwoSided"); + assert_eq!(select_alltoall_kernel(spec, 8, 8, Some("DEEPGEMM")), "NotEnabled"); + assert_eq!(select_alltoall_kernel(spec, 8, 8, Some("cute_dsl")), "NotEnabled"); + + // Pre-Blackwell WideEP: DeepEP when feasible, split on inter-node. + let mut hopper = db.system_spec.clone(); + hopper.gpu.sm_version = Some(90); + assert_eq!(select_alltoall_kernel(&hopper, 8, 8, Some("wideep")), "DeepEP"); + assert_eq!(select_alltoall_kernel(&hopper, 4, 8, Some("wideep")), "DeepEPLowLatency"); + assert_eq!(select_alltoall_kernel(&hopper, 8, 16, Some("wideep")), "NotEnabled"); + assert_eq!(select_alltoall_kernel(&hopper, 8, 8, None), "NotEnabled"); + + let zero = a2a(&db, "alltoall_dispatch", 333, MoeQuantMode::Nvfp4, Some("deepgemm")) + .expect("NotEnabled early return"); + assert_oracle(&zero, 0.0, Source::Empirical, "not_enabled_zero"); + } + + /// Python `_query_wideep_deepep_{ll,normal}_table` raise + /// `NotImplementedError` in EMPIRICAL mode — mirrored as the typed + /// `EmpiricalNotImplemented` (the gate fires before any table I/O, so + /// any sglang database works). + #[test] + fn deepep_empirical_mode_is_typed_not_implemented() { + let db = b200_sglang_db().with_mode(DatabaseMode::Empirical, TransferPolicy::ALL); + for flavor in [DispatchFlavor::DeepEpNormal, DispatchFlavor::DeepEpLowLatency] { + let result = deepep_op(8, true, flavor).query(&db, 64); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "{flavor:?} EMPIRICAL must be a typed empirical miss, got {result:?}" + ); + } + } + + /// Python's deepep tables have NO hybrid fallback (the `else:` branch + /// serves both SILICON and HYBRID and never routes through + /// `_query_silicon_or_hybrid`), so HYBRID answers the silicon interp + /// value unchanged. Oracles from the shipped h100 data: + /// + /// ```text + /// float(MoEDispatch._query_wideep_deepep_ll_table(db, node_num=1, + /// num_tokens=20, num_experts=256, topk=8, hidden_size=7168, + /// database_mode=common.DatabaseMode.HYBRID)) # -> 0.03853445 + /// float(MoEDispatch._query_wideep_deepep_normal_table(db, node_num=2, + /// num_tokens=64, num_experts=256, topk=8, hidden_size=7168, sms=20, + /// database_mode=common.DatabaseMode.HYBRID)) # -> 0.20963 + /// ``` + #[test] + fn deepep_hybrid_equals_silicon_interp() { + let db = h100_sglang_db(DatabaseMode::Hybrid); + // moe_ep=8 on h100 (8 GPUs/node) -> node_num = 1. + let ll = deepep_op(8, false, DispatchFlavor::DeepEpLowLatency) + .query(&db, 20) + .expect("ll hybrid query"); + assert_oracle(&ll, 0.03853445, Source::Silicon, "hyb_deepep_ll_t20"); + // moe_ep=16 -> node_num = 2; sms=20 with node_num != 1 resolves the + // 2-axis (sms, tokens) grid. + let mut normal = deepep_op(16, true, DispatchFlavor::DeepEpNormal); + normal.sms = 20; + let got = normal.query(&db, 64).expect("normal hybrid query"); + assert_oracle(&got, 0.20963, Source::Silicon, "hyb_deepep_normal_t64"); + } } diff --git a/rust/aiconfigurator-core/src/operators/msa.rs b/rust/aiconfigurator-core/src/operators/msa.rs new file mode 100644 index 000000000..8acf2aa91 --- /dev/null +++ b/rust/aiconfigurator-core/src/operators/msa.rs @@ -0,0 +1,445 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! MiniMax Sparse Attention (MSA) module ops for MiniMax-M3. +//! +//! Mirrors `aiconfigurator.sdk.operations.msa`. MSA is structurally a GQA +//! version of DSA: an indexer scores KV *blocks* (block_size tokens each), +//! the top-k blocks are selected, and full attention runs over only the +//! selected tokens. +//! +//! There is NO collected MSA silicon data. These ops therefore answer only +//! under HYBRID / EMPIRICAL: the SOL is analytic (same three-group split as +//! DSA/DSV4 — GEMM projections, FP8 indexer, sparse attention) and the +//! empirical value is a CROSS-OP (XOP) transfer from DSA's measured +//! utilisation at the same workload, scaled by the manual `dsa_scale_k` +//! level-alignment hook: `latency = SOL_msa / (util_dsa * k)`. SILICON mode +//! raises the perf-data miss; a policy with XOP disabled raises the terminal +//! empirical miss (there is nothing else to fall back on). + +use serde::{Deserialize, Serialize}; + +use crate::common::enums::{ + DatabaseMode, FmhaQuantMode, GemmQuantMode, KvCacheQuantMode, TransferKind, +}; +use crate::common::error::AicError; +use crate::common::system_spec::SystemSpec; +use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::dsa::DsaModuleOp; +use crate::perf_database::dsa::{dsa_context_sol_ms, dsa_dims, dsa_generation_sol_ms}; +use crate::perf_database::gemm::tc_flops_for_compute; +use crate::perf_database::PerfDatabase; + +/// One MSA module block (context or generation — the phase is chosen by the +/// `Op` variant). Field-for-field mirror of Python `_BaseMSAModule`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct MsaModuleOp { + pub name: String, + pub scale_factor: f64, + /// Local (per-rank) query heads. + pub num_heads: u32, + pub num_kv_heads: u32, + pub hidden_size: u32, + pub head_dim: u32, + pub v_head_dim: u32, + pub index_n_heads: u32, + pub index_head_dim: u32, + pub index_topk: u32, + pub block_size: u32, + pub kv_cache_dtype: KvCacheQuantMode, + pub fmha_quant_mode: FmhaQuantMode, + pub gemm_quant_mode: GemmQuantMode, + /// The DSA architecture whose measured utilisation is borrowed (XOP). + pub dsa_architecture: String, + /// Manual cross-op level-alignment scale `k` (`latency = SOL/(util*k)`). + pub dsa_scale_k: f64, +} + +impl MsaModuleOp { + /// Context (prefill) query. Mirrors `ContextMSAModule.query`. + pub fn query_context( + &self, + db: &PerfDatabase, + batch_size: u32, + s: u32, + prefix: u32, + ) -> Result { + let sol = self.sol_ms(db, batch_size, s, prefix, true); + match db.database_mode { + DatabaseMode::Sol | DatabaseMode::SolFull => { + Ok(PerformanceResult::new(sol * self.scale_factor, Source::Sol)) + } + DatabaseMode::Silicon => Err(AicError::PerfDatabase( + "MSA has no silicon data; use HYBRID or EMPIRICAL.".to_string(), + )), + DatabaseMode::Hybrid | DatabaseMode::Empirical => { + if !db.transfer_policy.contains(TransferKind::XOp) { + return Err(AicError::EmpiricalNotImplemented( + "MSA context: cross-op transfer (xop) is disabled by the transfer policy \ + and MSA has no own silicon data." + .to_string(), + )); + } + let util = self.dsa_context_util(db, batch_size, s, prefix); + match util { + Some(util) if util > 0.0 => { + let latency = sol / (util * self.dsa_scale_k); + Ok(PerformanceResult::new( + latency * self.scale_factor, + Source::Empirical, + )) + } + _ => Err(AicError::EmpiricalNotImplemented(format!( + "MSA context: no DSA util to transfer from (arch={}, b={batch_size}, \ + s={s}); collect DSA data or set msa_dsa_scale_k against an available \ + quant.", + self.dsa_architecture + ))), + } + } + } + } + + /// Generation (decode) query; `s` is the total KV length. Mirrors + /// `GenerationMSAModule.query`. + pub fn query_generation( + &self, + db: &PerfDatabase, + batch_size: u32, + s: u32, + ) -> Result { + let sol = self.sol_ms(db, batch_size, s, 0, false); + match db.database_mode { + DatabaseMode::Sol | DatabaseMode::SolFull => { + Ok(PerformanceResult::new(sol * self.scale_factor, Source::Sol)) + } + DatabaseMode::Silicon => Err(AicError::PerfDatabase( + "MSA has no silicon data; use HYBRID or EMPIRICAL.".to_string(), + )), + DatabaseMode::Hybrid | DatabaseMode::Empirical => { + if !db.transfer_policy.contains(TransferKind::XOp) { + return Err(AicError::EmpiricalNotImplemented( + "MSA generation: cross-op transfer (xop) is disabled by the transfer \ + policy and MSA has no own silicon data." + .to_string(), + )); + } + let util = self.dsa_generation_util(db, batch_size, s); + match util { + Some(util) if util > 0.0 => { + let latency = sol / (util * self.dsa_scale_k); + Ok(PerformanceResult::new( + latency * self.scale_factor, + Source::Empirical, + )) + } + _ => Err(AicError::EmpiricalNotImplemented(format!( + "MSA generation: no DSA util to transfer from (arch={}, b={batch_size}, \ + s={s}); collect DSA data or set msa_dsa_scale_k against an available \ + quant.", + self.dsa_architecture + ))), + } + } + } + } + + fn sol_ms(&self, db: &PerfDatabase, b: u32, s: u32, prefix: u32, is_context: bool) -> f64 { + msa_attention_sol_ms( + &db.system_spec, + is_context, + b as i128, + s as i128, + prefix as i128, + self.num_heads as i128, + self.num_kv_heads as i128, + self.hidden_size as i128, + self.head_dim as i128, + self.v_head_dim as i128, + self.index_n_heads as i128, + self.index_head_dim as i128, + self.index_topk as i128, + self.block_size as i128, + self.kv_cache_dtype, + self.fmha_quant_mode, + self.gemm_quant_mode, + ) + } + + /// DSA's measured utilisation (SOL / silicon) at the same context + /// workload, or `None`. Mirrors Python `_dsa_context_util`: the SOL comes + /// from the analytic DSA formula (`database_mode=SOL`), the silicon value + /// from a SILICON-view probe of the DSA module table; ANY failure means + /// no transfer source. + fn dsa_context_util(&self, db: &PerfDatabase, b: u32, s: u32, prefix: u32) -> Option { + let dims = dsa_dims(&self.dsa_architecture); + let sol = dsa_context_sol_ms( + &db.system_spec, + dims, + dims.index_topk, + self.kv_cache_dtype, + self.fmha_quant_mode, + self.gemm_quant_mode, + b as i64, + s as i64, + prefix as i64, + self.num_heads as i64, + ); + let probe = self.dsa_probe(dims.index_topk); + let silicon = probe + .query_context(&db.silicon_view(), b, s, prefix) + .ok()? + .latency_ms; + if sol > 0.0 && silicon > 0.0 { + Some(sol / silicon) + } else { + None + } + } + + /// DSA's measured utilisation at the same decode workload, or `None`. + /// Mirrors Python `_dsa_generation_util`. + fn dsa_generation_util(&self, db: &PerfDatabase, b: u32, s: u32) -> Option { + let dims = dsa_dims(&self.dsa_architecture); + let sol = dsa_generation_sol_ms( + &db.system_spec, + dims, + self.kv_cache_dtype, + self.gemm_quant_mode, + b as i64, + s as i64, + self.num_heads as i64, + ); + let probe = self.dsa_probe(dims.index_topk); + let silicon = probe + .query_generation(&db.silicon_view(), b, s) + .ok()? + .latency_ms; + if sol > 0.0 && silicon > 0.0 { + Some(sol / silicon) + } else { + None + } + } + + fn dsa_probe(&self, index_topk: i64) -> DsaModuleOp { + DsaModuleOp::new( + format!("{}_dsa_probe", self.name), + self.num_heads, + self.kv_cache_dtype, + self.fmha_quant_mode, + self.gemm_quant_mode, + self.dsa_architecture.clone(), + index_topk as u32, + ) + } +} + +/// SOL for one MSA block. Verbatim port of Python `_msa_attention_sol` +/// (`operations/msa.py`): GQA projections + per-block FP8 indexer + sparse +/// attention over the top-k selected tokens, integer pair counts in i128. +#[allow(clippy::too_many_arguments)] +fn msa_attention_sol_ms( + spec: &SystemSpec, + is_context: bool, + b: i128, + s: i128, + prefix: i128, + num_heads: i128, + num_kv_heads: i128, + hidden_size: i128, + head_dim: i128, + v_head_dim: i128, + index_n_heads: i128, + index_head_dim: i128, + index_topk: i128, + block_size: i128, + kv_quant: KvCacheQuantMode, + fmha_quant: FmhaQuantMode, + gemm_quant: GemmQuantMode, +) -> f64 { + let qk_head_dim = head_dim; + let tokens = if is_context { b * s } else { b }; + // context: full prefill of `s` new tokens on top of `prefix` cached. + // generation: 1 query token, kv_len = s - 1 cached. + let full_s = if is_context { prefix + s } else { s }; + let kv_len = if is_context { full_s } else { (s - 1).max(0) }; + + // ── GEMM group (Q / GQA-KV / O / indexer-Q projections) ────────────── + let gemm_ops = 2 * tokens * hidden_size * (num_heads * qk_head_dim) + + 2 * tokens * hidden_size * (2 * num_kv_heads * head_dim) + + 2 * tokens * (num_heads * v_head_dim) * hidden_size + + 2 * tokens * hidden_size * (index_n_heads * index_head_dim); + + // ── sparse attention: top-k saturated causal (query, kv) pair count ── + let (pairs, score_len) = if is_context { + let pairs = if full_s <= index_topk { + b * (full_s * (full_s + 1) - prefix * (prefix + 1)) / 2 + } else if prefix >= index_topk { + tokens * index_topk + } else { + let ramp = b * (index_topk * (index_topk + 1) - prefix * (prefix + 1)) / 2; + let sat = b * (full_s - index_topk) * index_topk; + ramp + sat + }; + (pairs, full_s) + } else { + (tokens * kv_len.min(index_topk), kv_len) + }; + let effective_kv = if is_context { + full_s.min(index_topk) + } else { + kv_len.min(index_topk) + }; + let attention_ops = 2 * num_heads * (qk_head_dim + v_head_dim) * pairs; // QK^T + AV + + // ── indexer: per-block scoring (block_size tokens per block), FP8 ──── + let num_blocks = if score_len > index_topk { + (score_len + block_size - 1) / block_size + } else { + 0 + }; + let indexer_ops = 2 * tokens * index_n_heads * index_head_dim * num_blocks; + + // ── memory ─────────────────────────────────────────────────────────── + let gemm_weight_elems = hidden_size * num_heads * qk_head_dim + + hidden_size * 2 * num_kv_heads * head_dim + + num_heads * v_head_dim * hidden_size + + hidden_size * index_n_heads * index_head_dim; + let gemm_weight_bytes = gemm_weight_elems as f64 * gemm_quant.mapping().memory; + let kv_cache_bytes = (b * num_kv_heads * effective_kv * (qk_head_dim + v_head_dim)) as f64 + * kv_quant.mapping().memory; + // FP8 index keys, per block (1 byte per element). + let indexer_cache_bytes = (b * num_blocks * index_n_heads * index_head_dim) as f64; + let q_io_bytes = (tokens * num_heads * qk_head_dim) as f64 * fmha_quant.mapping().memory * 2.0; + let total_mem = gemm_weight_bytes + kv_cache_bytes + indexer_cache_bytes + q_io_bytes; + + let gemm_flops = tc_flops_for_compute(spec, gemm_quant.mapping().compute); + // Python passes `common.FMHAQuantMode.fp8` (compute factor 2). + let fp8_flops = tc_flops_for_compute(spec, 2.0); + let attn_flops = tc_flops_for_compute(spec, fmha_quant.mapping().compute); + + let sol_math = (gemm_ops as f64 / gemm_flops + + indexer_ops as f64 / fp8_flops + + attention_ops as f64 / attn_flops) + * 1000.0; + let sol_mem = total_mem / spec.gpu.mem_bw * 1000.0; + sol_math.max(sol_mem) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::enums::TransferPolicy; + use std::path::PathBuf; + + const REPO_ROOT_HINT: &str = env!("CARGO_MANIFEST_DIR"); + + fn db(backend: &str, version: &str) -> PerfDatabase { + let systems_root = PathBuf::from(REPO_ROOT_HINT) + .join("../..") + .join("src/aiconfigurator/systems"); + let mut db = PerfDatabase::load(&systems_root, "b200_sxm", backend, version).expect("db must load"); + db.database_mode = DatabaseMode::Hybrid; + db + } + + /// MiniMax-M3 MSA block at tp=8 (local heads 8, GQA kv heads 1), all + /// bfloat16 — the shapes the M3 model builder emits. + fn msa_op() -> MsaModuleOp { + MsaModuleOp { + name: "msa".to_string(), + scale_factor: 1.0, + num_heads: 8, + num_kv_heads: 1, + hidden_size: 6144, + head_dim: 128, + v_head_dim: 128, + index_n_heads: 4, + index_head_dim: 128, + index_topk: 2048, + block_size: 128, + kv_cache_dtype: KvCacheQuantMode::Bfloat16, + fmha_quant_mode: FmhaQuantMode::Bfloat16, + gemm_quant_mode: GemmQuantMode::Bfloat16, + dsa_architecture: "GlmMoeDsaForCausalLM".to_string(), + dsa_scale_k: 1.0, + } + } + + fn approx(a: f64, b: f64) { + assert!((a - b).abs() < 1e-9 * b.abs().max(1.0), "expected {b}, got {a}"); + } + + /// Oracle values from the Python reference (`ContextMSAModule` / + /// `GenerationMSAModule` with `scale_factor=1.0` on a HYBRID + /// `shared_layer=False` view). Regenerate if the DSA tables or the + /// util math change. + #[test] + fn msa_xop_transfer_matches_python_oracles() { + for (backend, version, anchors) in [ + ( + "sglang", + "0.5.14", + [ + (1u32, 1024u32, 0u32, true, 0.15945479750992286), + (2, 3000, 512, true, 1.0573128907623435), + (8, 1025, 0, false, 0.03198051529058935), + (4, 7777, 0, false, 0.03568683502696848), + ], + ), + ( + "vllm", + "0.19.0", + [ + (1, 1024, 0, true, 0.44046553899838176), + (2, 3000, 512, true, 12.825449653876566), + (8, 1025, 0, false, 0.07392686165512992), + (4, 7777, 0, false, 0.0756313776883858), + ], + ), + ] { + let db = db(backend, version); + let op = msa_op(); + for (b, s, prefix, is_context, expected) in anchors { + let result = if is_context { + op.query_context(&db, b, s, prefix) + } else { + op.query_generation(&db, b, s) + } + .unwrap_or_else(|e| panic!("{backend} b={b} s={s}: {e}")); + approx(result.latency_ms, expected); + assert_eq!(result.source, Source::Empirical); + } + } + } + + /// XOP disabled ("balanced" preset) -> the terminal empirical miss; and + /// SILICON mode -> the perf-data miss ("MSA has no silicon data"). + #[test] + fn msa_policy_and_silicon_contracts() { + let mut hybrid = db("vllm", "0.19.0"); + hybrid.transfer_policy = TransferPolicy { + xshape: true, + xquant: true, + xprofile: false, + xop: false, + }; + let op = msa_op(); + assert!(matches!( + op.query_context(&hybrid, 1, 1024, 0), + Err(AicError::EmpiricalNotImplemented(_)) + )); + assert!(matches!( + op.query_generation(&hybrid, 8, 1025), + Err(AicError::EmpiricalNotImplemented(_)) + )); + + let mut silicon = db("vllm", "0.19.0"); + silicon.database_mode = DatabaseMode::Silicon; + assert!(matches!( + op.query_context(&silicon, 1, 1024, 0), + Err(AicError::PerfDatabase(_)) + )); + } +} diff --git a/rust/aiconfigurator-core/src/operators/op.rs b/rust/aiconfigurator-core/src/operators/op.rs index 1f694cb2f..f2d6a4191 100644 --- a/rust/aiconfigurator-core/src/operators/op.rs +++ b/rust/aiconfigurator-core/src/operators/op.rs @@ -21,6 +21,7 @@ use crate::operators::{ ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4ModuleOp, ElementwiseOp, EmbeddingOp, EncoderAttentionOp, GdnOp, GemmOp, GenerationAttentionOp, GenerationMlaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, MoEDispatchOp, MoeOp, + MsaModuleOp, NcclOp, P2POp, PerformanceResult, Source, VisionEncoderOp, WideEpContextMlaOp, WideEpGenerationMlaOp, WideEpMoeOp, }; @@ -109,6 +110,11 @@ pub enum Op { Vision(VisionEncoderOp), DsaContext(DsaModuleOp), DsaGeneration(DsaModuleOp), + /// MiniMax Sparse Attention (MSA) context module — no silicon data; + /// answers only under HYBRID/EMPIRICAL via cross-op DSA util transfer. + MsaContext(MsaModuleOp), + /// MSA generation module (`s` = total KV length). + MsaGeneration(MsaModuleOp), Dsv4Context(Dsv4ModuleOp), Dsv4Generation(Dsv4ModuleOp), Mhc(MhcModuleOp), @@ -200,6 +206,8 @@ impl Op { Op::Vision(o) => &o.name, Op::DsaContext(o) => &o.name, Op::DsaGeneration(o) => &o.name, + Op::MsaContext(o) => &o.name, + Op::MsaGeneration(o) => &o.name, Op::Dsv4Context(o) => &o.name, Op::Dsv4Generation(o) => &o.name, Op::Mhc(o) => &o.name, @@ -280,6 +288,8 @@ impl Op { Op::Vision(op) => op.query(db, ctx.num_image_tokens), Op::DsaContext(op) => op.query_context(db, ctx.batch_size, ctx.s, ctx.prefix), Op::DsaGeneration(op) => op.query_generation(db, ctx.batch_size, ctx.s), + Op::MsaContext(op) => op.query_context(db, ctx.batch_size, ctx.s, ctx.prefix), + Op::MsaGeneration(op) => op.query_generation(db, ctx.batch_size, ctx.s), Op::Dsv4Context(op) => op.query_context(db, ctx.batch_size, ctx.s, ctx.prefix), Op::Dsv4Generation(op) => op.query_generation(db, ctx.batch_size, ctx.s), Op::Mhc(op) => op.query(db, ctx.num_tokens), diff --git a/rust/aiconfigurator-core/src/operators/wideep_moe.rs b/rust/aiconfigurator-core/src/operators/wideep_moe.rs index 0dff0af80..2576aa5af 100644 --- a/rust/aiconfigurator-core/src/operators/wideep_moe.rs +++ b/rust/aiconfigurator-core/src/operators/wideep_moe.rs @@ -3,11 +3,11 @@ //! TensorRT-LLM WideEP MoE compute operator. //! -//! Apple-to-apple port of `aiconfigurator.sdk.operations.moe.TrtLLMWideEPMoE`. -//! Pure-compute kernel timing (no All2All). The dispatch / combine cost -//! belongs to `MoEDispatchOp` (with the TrtllmAlltoall or DeepEP flavor) -//! or the `wideep` table — depending on which path the model variant -//! exercises. +//! Apple-to-apple port of `aiconfigurator.sdk.operations.moe.TrtLLMWideEPMoE` +//! (`_query_compute_table`). Pure-compute kernel timing (no All2All). The +//! dispatch / combine cost belongs to `MoEDispatchOp` (with the +//! TrtllmAlltoall or DeepEP flavor) or the `wideep` table — depending on +//! which path the model variant exercises. //! //! EPLB modes: //! - EPLB off: `workload_distribution` without `_eplb` suffix, @@ -20,11 +20,22 @@ //! Mirrors Python: `query` multiplies `num_tokens` by `attention_dp_size` //! before the lookup (the perf table is collected per-rank but the //! op-level input is per-attention-DP-rank). +//! +//! Database-mode dispatch follows the gemm.rs reference pattern: EMPIRICAL +//! always estimates `SOL(query)/util` from the op's OWN slice (this table +//! has no transfer ladder — Python's `get_empirical_from_sol` builds a +//! depth-1 grid over the own token curve only); HYBRID queries silicon and +//! converts a typed missing-data error into the estimate; SILICON is +//! unchanged. Both paths resolve the kernel via the `_select_kernel` mirror +//! (architecture-preferred kernel, falling back to whatever the table +//! actually collected). use serde::{Deserialize, Serialize}; -use crate::common::enums::MoeQuantMode; +use crate::common::enums::{DatabaseMode, MoeQuantMode}; use crate::common::error::AicError; +use crate::common::system_spec::SystemSpec; use crate::operators::base::{PerformanceResult, Source}; +use crate::operators::util_empirical::{self, UtilGrid}; use crate::perf_database::PerfDatabase; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -42,11 +53,19 @@ pub struct WideEpMoeOp { pub workload_distribution: String, /// EPLB slots; defaults to `num_experts` (no EPLB redundancy). pub num_slots: u32, - /// WideEP MoE compute kernel: `"moe_torch_flow"` (Cutlass; SM<100 - /// default) or `"deepgemm"` (SM>=100 with fp8_block). + /// UNUSED — retained for opspec wire-format compatibility. Kernel + /// selection now mirrors Python `TrtLLMWideEPMoE._select_kernel` + /// (architecture + quant preferred, table-availability fallback); Python + /// never lets the caller pick the compute kernel. pub kernel_source: String, } +/// `num_gemms` for the WideEP compute SOL. Python `TrtLLMWideEPMoE.query` +/// never forwards its `is_gated` flag to `query_wideep_moe_compute`, so +/// `_query_compute_table` always runs with its `is_gated=True` default +/// (3 GEMMs: gate + up + down). +const NUM_GEMMS: u64 = 3; + impl WideEpMoeOp { #[allow(clippy::too_many_arguments)] pub fn new( @@ -85,8 +104,39 @@ impl WideEpMoeOp { ) -> Result { // Python: `x = num_tokens * self._attention_dp_size`. let scaled = num_tokens.saturating_mul(self.attention_dp_size.max(1)); - let latency = db.wideep_moe.query_compute( - scaled, + + // Database-mode dispatch, mirroring the Python `_query_compute_table` + // tail (`database._query_silicon_or_hybrid`). The SOL diagnostic + // modes never reach the compiled engine. + let (latency, source) = match db.database_mode { + DatabaseMode::Empirical => (self.empirical_latency(db, scaled)?, Source::Empirical), + DatabaseMode::Hybrid => match self.silicon_latency(db, scaled) { + Ok(latency) => (latency, Source::Silicon), + Err(err) if err.is_missing_perf_data() => { + (self.empirical_latency(db, scaled)?, Source::Empirical) + } + Err(err) => return Err(err), + }, + _ => (self.silicon_latency(db, scaled)?, Source::Silicon), + }; + Ok(PerformanceResult::new(latency, source) + .clamp_non_negative() + .scaled(self.scale_factor)) + } + + /// SILICON table resolution: resolved kernel + level-wise distribution + /// fallback in the table layer, token curve anchored on the WideEP MoE + /// roofline (Python `get_silicon` threads the same `get_sol` closure + /// through `OpInterpConfig`). + fn silicon_latency(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + let kernel = self.select_kernel(db)?; + let spec = &db.system_spec; + // Engine coordinates are always integral (table keys / the u32 + // query); rounding to u32 keeps integer floor-division parity with + // Python's `get_sol` (same convention as operators/moe.rs). + let sol = |t: f64| self.sol_latency_ms(spec, t.round() as u32); + db.wideep_moe.query_compute( + num_tokens, self.hidden_size, self.inter_size, self.topk, @@ -96,10 +146,241 @@ impl WideEpMoeOp { self.moe_ep_size, self.quant_mode, &self.workload_distribution, - &self.kernel_source, - )?; - Ok(PerformanceResult::new(latency, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) + &kernel, + &sol, + ) + } + + /// `SOL(query)/util` over the op's OWN token curve (depth 1, no transfer + /// ladder). Mirrors Python `_query_compute_table::get_empirical_from_sol`. + fn empirical_latency(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + let spec = &db.system_spec; + let kernel = self.select_kernel(db)?; + let sol = |c: &[f64]| self.sol_latency_ms(spec, c[0].round() as u32); + let sol_time = self.sol_latency_ms(spec, num_tokens); + + // Python keys on ("wideep_moe", system, backend, version, kernel, + // quant, topk, num_experts, hidden, inter, num_slots, moe_tp, + // moe_ep) + id(node); the cache here is per-database, and the + // requested workload_distribution stands in for the node identity + // (the resolved slice is a function of it). + let key = format!( + "wideep_moe:{}:{}:{}:{}:{}:{}:{}:{}:{}:{}", + kernel, + self.quant_mode.name(), + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.num_slots, + self.moe_tp_size, + self.moe_ep_size, + self.workload_distribution, + ); + let grid = db.util_grids.get_or_try_build(&key, || { + match db.wideep_moe.slice_points( + &kernel, + self.quant_mode, + &self.workload_distribution, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.num_slots, + self.moe_tp_size, + self.moe_ep_size, + ) { + Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples( + points.into_iter().map(|(t, lat)| (vec![t as f64], lat)), + sol, + )))), + // Typed coverage miss -> no grid (estimate() raises the + // empirical miss); schema/load errors propagate. + Err(err) if err.is_missing_perf_data() => Ok(None), + Err(err) => Err(err), + } + })?; + let query = [num_tokens as f64]; + let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + Ok(latency) + } + + /// Mirror of Python `TrtLLMWideEPMoE._select_kernel`: + /// 1. SM >= 100 (Blackwell) with fp8_block -> `deepgemm`; + /// 2. otherwise -> `moe_torch_flow` (Cutlass); + /// then, when the table actually loaded, keep the preferred kernel if + /// collected, else fall back to the first collected kernel. + fn select_kernel(&self, db: &PerfDatabase) -> Result { + let is_blackwell = db.system_spec.gpu.sm_version.unwrap_or(0) >= 100; + // Python: `"fp8_block" in quant_mode_str` (substring, not equality). + let is_fp8_block = self.quant_mode.name().contains("fp8_block"); + let preferred = if is_blackwell && is_fp8_block { + "deepgemm" + } else { + "moe_torch_flow" + }; + let available = db.wideep_moe.available_kernels()?; + if available.iter().any(|k| k == preferred) { + return Ok(preferred.to_string()); + } + if let Some(first) = available.into_iter().next() { + return Ok(first); + } + Ok(preferred.to_string()) + } + + /// WideEP MoE roofline SOL (ms) mirroring Python + /// `_query_compute_table.get_sol`: identical to the plain-MoE roofline + /// (`operators/moe.rs`) except the weight-read term uses `num_slots` + /// instead of `num_experts` (WideEP EPLB redundant mode may replicate + /// experts across slots). `num_experts` never enters the math. + fn sol_latency_ms(&self, spec: &SystemSpec, num_tokens: u32) -> f64 { + let total_tokens = num_tokens as u64 * self.topk as u64; + let moe_ep = (self.moe_ep_size as u64).max(1); + let moe_tp = (self.moe_tp_size as u64).max(1); + let h = self.hidden_size as u64; + let inter = self.inter_size as u64; + let slots = self.num_slots as u64; + + let ops = total_tokens * h * inter * NUM_GEMMS * 2 / moe_ep / moe_tp; + let mem_bytes_int = total_tokens / moe_ep * h * 2 // input + output + + total_tokens / moe_ep * inter * NUM_GEMMS / moe_tp // intermediate + + h * inter * NUM_GEMMS / moe_tp + * std::cmp::min(slots / moe_ep, total_tokens / moe_ep); // weights (num_slots) + let mem_bytes = (mem_bytes_int as f64) * self.quant_mode.mapping().memory; + + // Python indexes `bfloat16_tc_flops` directly (KeyError if absent); + // every shipped system populates it — same fallback convention as + // operators/moe.rs. + let tc_flops = spec.gpu.bfloat16_tc_flops.unwrap_or(1.0); + let sol_math = (ops as f64) / (tc_flops * self.quant_mode.mapping().compute) * 1000.0; + let sol_mem = mem_bytes / spec.gpu.mem_bw * 1000.0; + sol_math.max(sol_mem) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::enums::TransferPolicy; + use std::path::PathBuf; + + fn b200_trtllm_db(mode: DatabaseMode) -> PerfDatabase { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + PerfDatabase::load(&root, "b200_sxm", "trtllm", "1.3.0rc10") + .expect("db loads") + .with_mode(mode, TransferPolicy::ALL) + } + + /// Collected b200 WideEP shape: nvfp4, power_law_1.01, topk=8, + /// experts=256, hidden=6144, inter=2048, slots=256, tp=1, ep=2. + fn op() -> WideEpMoeOp { + WideEpMoeOp { + name: "wideep_moe".into(), + scale_factor: 1.0, + hidden_size: 6144, + inter_size: 2048, + topk: 8, + num_experts: 256, + moe_tp_size: 1, + moe_ep_size: 2, + attention_dp_size: 1, + quant_mode: MoeQuantMode::Nvfp4, + workload_distribution: "power_law_1.01".into(), + num_slots: 256, + kernel_source: "moe_torch_flow".into(), + } + } + + fn assert_oracle(result: &PerformanceResult, expected: f64, source: Source, label: &str) { + assert!( + (result.latency_ms - expected).abs() < 1e-9, + "{label}: expected {expected}, got {}", + result.latency_ms + ); + assert_eq!(result.source, source, "{label}: wrong source"); + } + + /// Oracle values generated from the Python reference on the same data + /// (shared layer pinned OFF so Python reads exactly the primary parquet): + /// + /// ```text + /// db = perf_database.get_database_view("b200_sxm", "trtllm", "1.3.0rc10", + /// allow_missing_data=True, database_mode=..., shared_layer=False) + /// float(TrtLLMWideEPMoE._query_compute_table(db, num_tokens=..., + /// hidden_size=6144, inter_size=2048, topk=8, num_experts=256, + /// num_slots=256, moe_tp_size=1, moe_ep_size=2, + /// quant_mode=common.MoEQuantMode.nvfp4, + /// workload_distribution=..., database_mode=...)) + /// ``` + /// + /// nt=1 is a collected point (exact hit, EMPIRICAL == measured); + /// nt=333 is an off-grid interior point; nt=100000 is beyond the + /// collected max (16384) — the util-hold anchors on the num_slots-aware + /// roofline, so it locks the SOL threading through both paths (a linear + /// token proxy fails this at 1e-9). The kernel resolution also runs the + /// `_select_kernel` fallback: preferred `moe_torch_flow` is uncollected, + /// the table only carries `wideep_compute_cutlass`. + #[test] + fn wideep_compute_empirical_matches_python_oracles() { + let db = b200_trtllm_db(DatabaseMode::Empirical); + let exact = op().query(&db, 1).expect("exact-hit empirical"); + assert_oracle(&exact, 0.08600959777832032, Source::Empirical, "emp_t1"); + let off = op().query(&db, 333).expect("off-grid empirical"); + assert_oracle(&off, 0.4523388956415573, Source::Empirical, "emp_t333"); + let hold = op().query(&db, 100000).expect("beyond-range empirical"); + assert_oracle(&hold, 15.171406557783484, Source::Empirical, "emp_t100000"); + } + + /// HYBRID with data present stays on silicon; the in-range interpolation + /// differs from the empirical reconstruction at the same point + /// (0.45251... vs 0.45233...), and the beyond-range hold coincides with + /// the empirical value (same boundary anchor + same roofline). + #[test] + fn wideep_compute_hybrid_prefers_silicon_when_covered() { + let db = b200_trtllm_db(DatabaseMode::Hybrid); + let hit = op().query(&db, 1).expect("collected token point"); + assert_oracle(&hit, 0.08600959777832032, Source::Silicon, "hyb_t1"); + let interp = op().query(&db, 333).expect("in-range token interp"); + assert_oracle(&interp, 0.45251359716057776, Source::Silicon, "hyb_t333"); + let hold = op().query(&db, 100000).expect("beyond-range hold"); + assert_oracle(&hold, 15.171406557783484, Source::Silicon, "hyb_t100000"); + } + + /// Distribution fallback: `"balanced"` is uncollected; both silicon and + /// empirical fall back to the first distribution under (kernel, quant) + /// (`power_law_1.01`) — nt=64 is a collected point there, so both modes + /// return the measured value. + #[test] + fn wideep_compute_distribution_fallback_matches_python_oracle() { + let mut fallback_op = op(); + fallback_op.workload_distribution = "balanced".into(); + let emp = fallback_op + .query(&b200_trtllm_db(DatabaseMode::Empirical), 64) + .expect("empirical dist fallback"); + assert_oracle(&emp, 0.4135615825653076, Source::Empirical, "emp_dist_fb_t64"); + let hyb = fallback_op + .query(&b200_trtllm_db(DatabaseMode::Hybrid), 64) + .expect("hybrid dist fallback"); + assert_oracle(&hyb, 0.4135615825653076, Source::Silicon, "hyb_dist_fb_t64"); + } + + /// HYBRID on a slice with no data anywhere (hidden=9999): the silicon + /// miss falls to the empirical path, whose own-slice grid is also empty + /// (this table has NO transfer ladder) -> the typed + /// EmpiricalNotImplemented miss surfaces, exactly like Python + /// (`EmpiricalNotImplementedError`). + #[test] + fn wideep_compute_hybrid_typed_miss_surfaces_empirical_not_implemented() { + let db = b200_trtllm_db(DatabaseMode::Hybrid); + let mut missing = op(); + missing.hidden_size = 9999; + let result = missing.query(&db, 64); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "expected the typed empirical miss, got {result:?}" + ); } } diff --git a/rust/aiconfigurator-core/src/perf_database/wideep.rs b/rust/aiconfigurator-core/src/perf_database/wideep.rs index 7284f6379..78d8430ae 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep.rs @@ -10,13 +10,19 @@ //! - `wideep_generation_moe_perf.txt` //! - `wideep_moe_perf.txt` (TRT-LLM WideEP MoE compute; extra columns //! handled by tolerant deserialization) -//! - `trtllm_alltoall_perf.txt` (TRT-LLM alltoall dispatch; subset of -//! MoE columns) //! //! 2. DeepEP dispatch layout (separate notify/transmit latencies): //! - `wideep_deepep_normal_perf.txt` //! - `wideep_deepep_ll_perf.txt` //! +//! 3. TRT-LLM alltoall layout (`trtllm_alltoall_perf.txt`), keyed exactly +//! like Python `load_trtllm_alltoall_data`: +//! `[kernel_source][op_name][quant][num_nodes][hidden_size][topk]` +//! `[num_experts][moe_ep_size][num_tokens] -> latency_ms`. +//! `kernel_source` defaults to `"NVLinkTwoSided"` when the column is +//! absent/null; `num_nodes` defaults to `max(1, moe_ep_size // 4)` when +//! the column is absent (the GB200-era collector convention). +//! //! All loaders are lazy. Token curves resolve on the shared perf_interp v2 //! engine (Grid, RAW lerp in range; beyond the collected range the //! boundary util is held with `k_tail=1` and a LINEAR num_tokens proxy SOL @@ -70,7 +76,7 @@ pub struct WideEpTable { context_moe: OnceLock>, generation_moe: OnceLock>, trtllm_wideep_moe: OnceLock>, - trtllm_alltoall: OnceLock>, + trtllm_alltoall: OnceLock>, deepep_normal: OnceLock>, deepep_ll: OnceLock>, } @@ -79,6 +85,24 @@ struct MoeGrids { by_keys: BTreeMap>, } +struct AlltoallGrids { + by_keys: BTreeMap>, +} + +/// Python `load_trtllm_alltoall_data` coordinate (minus the token level): +/// `[kernel_source][op_name][quant][num_nodes][hidden][topk][experts][ep]`. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct AlltoallKey { + kernel_source: String, + op_name: String, + quant: String, + num_nodes: u32, + hidden_size: u32, + topk: u32, + num_experts: u32, + moe_ep_size: u32, +} + struct DispatchGrids { by_keys: BTreeMap>, } @@ -267,35 +291,105 @@ impl WideEpTable { ) } - /// TRT-LLM alltoall dispatch latency. CSV uses `moe_ep_size` for fan-out - /// (`moe_tp_size`/`inter_size` are not present); the query API still - /// accepts them for shape symmetry but they're effectively ignored. + /// TRT-LLM alltoall latency (ms) for one `(kernel_source, op_name, + /// table_quant, node_num, hidden, topk, experts, ep)` slice, resolved on + /// the shared token-curve engine. Mirrors the silicon body of Python + /// `TrtLLMWideEPMoEDispatch._query_alltoall_table`: RAW lerp in range; + /// beyond-range holds ride the LINEAR num_tokens proxy, which is + /// ratio-equivalent to Python's alltoall SOL (`data_bytes / bw` with + /// `data_bytes = const * num_tokens` in every op branch, so + /// `sol(q)/sol(b) = q/b` exactly). + /// + /// `quant` must be the TABLE quant, i.e. already normalized by the + /// operator (`fp8_block -> fp8`); this layer does no algorithm work. #[allow(clippy::too_many_arguments)] pub fn query_trtllm_alltoall( &self, + kernel_source: &str, + op_name: &str, + quant: MoeQuantMode, + node_num: u32, + hidden_size: u32, + topk: u32, + num_experts: u32, + moe_ep_size: u32, num_tokens: u32, + ) -> Result { + let by_tokens = self.alltoall_slice( + kernel_source, + op_name, + quant, + node_num, + hidden_size, + topk, + num_experts, + moe_ep_size, + )?; + query_token_curve(by_tokens, num_tokens as f64, &|t| t) + } + + /// Own-slice `num_tokens -> latency_ms` points for the alltoall grid. + /// Typed `AicError::PerfDatabase` miss when the slice is absent or empty + /// (mirrors `util_empirical.require_data_slice` as used by the empirical + /// `_slice` closure in `_query_alltoall_table`). + #[allow(clippy::too_many_arguments)] + pub fn alltoall_slice_points( + &self, + kernel_source: &str, + op_name: &str, + quant: MoeQuantMode, + node_num: u32, hidden_size: u32, topk: u32, num_experts: u32, moe_ep_size: u32, + ) -> Result, AicError> { + let by_tokens = self.alltoall_slice( + kernel_source, + op_name, + quant, + node_num, + hidden_size, + topk, + num_experts, + moe_ep_size, + )?; + Ok(by_tokens.iter().map(|(&t, &lat)| (t, lat)).collect()) + } + + #[allow(clippy::too_many_arguments)] + fn alltoall_slice( + &self, + kernel_source: &str, + op_name: &str, quant: MoeQuantMode, - workload_distribution: &str, - ) -> Result { + node_num: u32, + hidden_size: u32, + topk: u32, + num_experts: u32, + moe_ep_size: u32, + ) -> Result<&BTreeMap, AicError> { let grids = self.load_trtllm_alltoall()?; - query_moe( - grids, - num_tokens, + let key = AlltoallKey { + kernel_source: kernel_source.to_string(), + op_name: op_name.to_string(), + quant: quant.name().to_string(), + num_nodes: node_num, hidden_size, - 0, topk, num_experts, - 1, moe_ep_size, - quant, - workload_distribution, - &self.data_root, - false, - ) + }; + grids + .by_keys + .get(&key) + .filter(|curve| !curve.is_empty()) + .ok_or_else(|| { + AicError::PerfDatabase(format!( + "TRT-LLM alltoall data missing for {key:?} at {}", + self.data_root.display() + )) + }) } /// DeepEP normal-mode dispatch point. @@ -394,10 +488,10 @@ impl WideEpTable { .get_or_init(|| load_moe_parquet(&self.moe_sources)); cell.as_ref().map_err(clone_err) } - fn load_trtllm_alltoall(&self) -> Result<&MoeGrids, AicError> { + fn load_trtllm_alltoall(&self) -> Result<&AlltoallGrids, AicError> { let cell = self .trtllm_alltoall - .get_or_init(|| load_moe_parquet(&self.alltoall_sources)); + .get_or_init(|| load_alltoall_parquet(&self.alltoall_sources)); cell.as_ref().map_err(clone_err) } fn load_deepep_normal(&self) -> Result<&NormalDispatchGrids, AicError> { @@ -633,7 +727,8 @@ fn dispatch_lookup( /// (`or_insert`), mirroring Python's `_read_filtered_rows` concatenation + /// `load_wideep_*_moe_data` skip-on-key-conflict. Missing files are skipped; an /// error is returned only when no source yields rows. Reused for the -/// context/generation/wideep-moe/alltoall parquets. +/// context/generation/wideep-moe parquets (alltoall has its own layout, see +/// [`load_alltoall_parquet`]). fn load_moe_parquet(sources: &[PerfSource]) -> Result { let mut by_keys: BTreeMap> = BTreeMap::new(); let mut any_source = false; @@ -647,9 +742,8 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result { let moe_dtype_col = reader.col("moe_dtype")?; let num_tokens_col = reader.col("num_tokens")?; let hidden_size_col = reader.col("hidden_size")?; - // `inter_size` / `moe_tp_size` are absent in `trtllm_alltoall_perf.parquet`; - // shared with the wideep_*_moe parquets which carry them. Optional lookup - // mirrors the prior `Option` deserialization plus `unwrap_or` default. + // Optional lookup mirrors the prior `Option` deserialization plus + // `unwrap_or` default (tolerant of older collector schemas). let inter_size_col = reader.col_optional("inter_size"); let topk_col = reader.col("topk")?; let num_experts_col = reader.col("num_experts")?; @@ -692,6 +786,78 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result { Ok(MoeGrids { by_keys }) } +/// Load the TRT-LLM alltoall table from an ordered, priority-sorted source +/// list. Missing files are skipped; an error is returned only when no source +/// yields rows. +/// +/// Keying mirrors Python `load_trtllm_alltoall_data` exactly: +/// - `kernel_source` is optional/nullable, defaulting to `"NVLinkTwoSided"`; +/// - `num_nodes` is optional, defaulting to `max(1, moe_ep_size // 4)` (the +/// GB200 NVL4 convention Python applies when the column is absent); +/// - duplicates resolve LAST-wins (Python uses plain dict assignment here, +/// unlike the skip-on-conflict deepep/moe-compute loaders), over +/// `_read_filtered_rows`' priority-ordered concatenation. +fn load_alltoall_parquet(sources: &[PerfSource]) -> Result { + let mut by_keys: BTreeMap> = BTreeMap::new(); + let mut any_source = false; + for source in sources { + let path = source.path(); + if !path.exists() { + continue; + } + any_source = true; + let reader = PerfReader::open(path)?; + let op_name_col = reader.col("op_name")?; + let moe_dtype_col = reader.col("moe_dtype")?; + let num_tokens_col = reader.col("num_tokens")?; + let hidden_size_col = reader.col("hidden_size")?; + let topk_col = reader.col("topk")?; + let num_experts_col = reader.col("num_experts")?; + let moe_ep_size_col = reader.col("moe_ep_size")?; + let latency_col = reader.col("latency")?; + let kernel_source_col = reader.col_optional("kernel_source"); + let num_nodes_col = reader.col_optional("num_nodes"); + for row in reader.rows()? { + let row = row?; + if !kernel_source_ok(source.kernel_sources(), kernel_source_col, &row)? { + continue; + } + let moe_ep_size = row.u32(moe_ep_size_col)?; + let kernel_source = row + .str_optional(kernel_source_col)? + .map(|s| s.to_string()) + .unwrap_or_else(|| "NVLinkTwoSided".to_string()); + let num_nodes = match row.u32_optional(num_nodes_col)? { + Some(n) => n, + None => (moe_ep_size / 4).max(1), + }; + let key = AlltoallKey { + kernel_source, + op_name: row.str_owned(op_name_col)?, + quant: row.str_owned(moe_dtype_col)?, + num_nodes, + hidden_size: row.u32(hidden_size_col)?, + topk: row.u32(topk_col)?, + num_experts: row.u32(num_experts_col)?, + moe_ep_size, + }; + // Last-wins on the full coordinate (Python plain assignment). + by_keys + .entry(key) + .or_default() + .insert(row.u32(num_tokens_col)?, row.f64(latency_col)?); + } + } + if !any_source || by_keys.is_empty() { + return Err(AicError::PerfDatabase(format!( + "no TRT-LLM alltoall rows loaded from {} source(s) (first: {})", + sources.len(), + sources.first().map(|s| s.path().display().to_string()).unwrap_or_default() + ))); + } + Ok(AlltoallGrids { by_keys }) +} + /// Load the DeepEP-normal dispatch table from an ordered source list. Missing /// files are skipped; an error is returned only when no source yields rows. /// diff --git a/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs b/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs index 3c240f32d..7c2d1acfc 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs @@ -95,18 +95,17 @@ impl WideEpMoeTable { /// `(kernel_source, quant, distribution, topk, num_experts, hidden, /// inter, num_slots, moe_tp_size, moe_ep_size)` key. The token curve /// rides the perf_interp v2 engine (1-axis Grid, RAW lerp in range, - /// boundary util-hold beyond it), mirroring Python's - /// `TrtllmWideEPMoE._query_compute_table`. If the exact `distribution` - /// isn't in the table, falls back to the first distribution available - /// under the matched quant — same as Python. + /// boundary util-hold beyond it), mirroring the silicon body of Python's + /// `TrtLLMWideEPMoE._query_compute_table`. `kernel_source` must be the + /// RESOLVED kernel (the operator mirrors `_select_kernel`, including its + /// data-availability fallback via [`Self::available_kernels`]); the + /// distribution falls back level-wise exactly like Python's + /// `get_silicon` (first distribution under the matched `(kernel, quant)` + /// when the exact string is absent — see [`Self::slice_points`]). /// - /// The util-hold SOL is a LINEAR num_tokens proxy: Python anchors on - /// the WideEP MoE roofline (`_query_compute_table.get_sol`, num_slots- - /// aware), which this table layer cannot compute (no SystemSpec). The - /// proxy only affects beyond-range holds; in-range lerp is SOL-free and - /// matches Python exactly. Thread the roofline through (like - /// `moe.rs::MoeTable::query`) if the operator layer ever needs - /// beyond-range parity here. + /// `sol` is the WideEP MoE roofline (num_slots-aware) the beyond-range + /// util-hold anchors on — Python threads the same closure through + /// `OpInterpConfig(sol_fn=...)`. In-range lerp never consults it. #[allow(clippy::too_many_arguments)] pub fn query_compute( &self, @@ -121,36 +120,46 @@ impl WideEpMoeTable { quant: MoeQuantMode, distribution: &str, kernel_source: &str, + sol: &dyn Fn(f64) -> f64, ) -> Result { - let grids = self.load_compute()?; - - // Mirror Python's `TrtllmWideEPMoE._select_kernel` fallback: if - // the requested kernel_source isn't in the loaded table (e.g. the - // caller asks for "moe_torch_flow" but the collected data is - // tagged "wideep_compute_cutlass"), fall back to any kernel - // present in the grid (Python takes `available_kernels[0]`). - let resolved_kernel = if grids - .by_keys - .keys() - .any(|k| k.kernel_source == kernel_source) - { - kernel_source.to_string() - } else { - grids - .by_keys - .keys() - .next() - .map(|k| k.kernel_source.clone()) - .unwrap_or_else(|| kernel_source.to_string()) - }; + let by_tokens = self.resolve_slice( + kernel_source, + quant.name(), + distribution, + topk, + num_experts, + hidden_size, + inter_size, + num_slots, + moe_tp_size, + moe_ep_size, + )?; + query_token_curve(by_tokens, num_tokens as f64, sol) + } - // Find a matching key. Python falls back to "first distribution - // under the same (kernel, quant)" when the exact distribution - // string isn't in the loaded data. - let exact_key = WideEpMoeKey { - kernel_source: resolved_kernel, - quant: quant.name().to_string(), - distribution: distribution.to_string(), + /// Own-slice `num_tokens -> latency_ms` points, after the level-wise + /// distribution fallback. Typed `AicError::PerfDatabase` miss when the + /// slice is absent or empty — mirrors the `_slice` closure of Python + /// `_query_compute_table.get_empirical_from_sol` + /// (`util_empirical.require_data_slice` semantics). + #[allow(clippy::too_many_arguments)] + pub fn slice_points( + &self, + kernel_source: &str, + quant: MoeQuantMode, + distribution: &str, + topk: u32, + num_experts: u32, + hidden_size: u32, + inter_size: u32, + num_slots: u32, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { + let by_tokens = self.resolve_slice( + kernel_source, + quant.name(), + distribution, topk, num_experts, hidden_size, @@ -158,42 +167,120 @@ impl WideEpMoeTable { num_slots, moe_tp_size, moe_ep_size, + )?; + Ok(by_tokens.iter().map(|(&t, &lat)| (t, lat)).collect()) + } + + /// Distinct kernel names present in the loaded table, in sorted + /// (BTreeMap) order; EMPTY when the table failed to load with a typed + /// data miss. Mirrors the `if database._wideep_moe_compute_data:` guard + /// of Python `_select_kernel` (an unloaded/empty table is falsy there — + /// the caller then keeps its architecture-preferred kernel). NOTE: + /// Python yields dict insertion (file row) order; sorted order is + /// observable only when several kernels coexist and the preferred one is + /// absent. + pub fn available_kernels(&self) -> Result, AicError> { + let grids = match self.load_compute() { + Ok(grids) => grids, + Err(err) if err.is_missing_perf_data() => return Ok(Vec::new()), + Err(err) => return Err(err), }; - let by_tokens = match grids.by_keys.get(&exact_key) { - Some(t) => t, - None => { - let fallback = grids - .by_keys - .iter() - .find(|(k, _)| { - k.kernel_source == exact_key.kernel_source - && k.quant == exact_key.quant - && k.topk == exact_key.topk - && k.num_experts == exact_key.num_experts - && k.hidden_size == exact_key.hidden_size - && k.inter_size == exact_key.inter_size - && k.num_slots == exact_key.num_slots - && k.moe_tp_size == exact_key.moe_tp_size - && k.moe_ep_size == exact_key.moe_ep_size - }) - .map(|(_, t)| t) - .ok_or_else(|| { - AicError::PerfDatabase(format!( - "WideEP MoE compute data missing for {exact_key:?} at {}", - self.data_root.display() - )) - })?; - fallback + let mut names: Vec = Vec::new(); + for key in grids.by_keys.keys() { + // `WideEpMoeKey` sorts by kernel_source first, so consecutive + // dedup suffices. + if names.last().map(String::as_str) != Some(key.kernel_source.as_str()) { + names.push(key.kernel_source.clone()); } - }; + } + Ok(names) + } - if by_tokens.is_empty() { + /// Level-wise slice resolution mirroring Python's nested-dict walk: + /// `require_data_slice(data, kernel)` -> `require_data_slice(kd, quant)` + /// -> distribution fallback (`workload if present else first available + /// under (kernel, quant)`) -> exact remaining coordinate. Each level + /// misses with a typed `AicError::PerfDatabase`. NOTE: the Python + /// fallback takes the FIRST distribution in dict-insertion (file row) + /// order and then requires the full shape under it (a shape present only + /// under a later distribution still misses) — mirrored here with sorted + /// (BTreeMap) order for the distribution list. + #[allow(clippy::too_many_arguments)] + fn resolve_slice( + &self, + kernel_source: &str, + quant_name: &str, + distribution: &str, + topk: u32, + num_experts: u32, + hidden_size: u32, + inter_size: u32, + num_slots: u32, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result<&BTreeMap, AicError> { + let grids = self.load_compute()?; + let mut kernel_seen = false; + let mut quant_seen = false; + let mut first_distribution: Option<&str> = None; + let mut requested_distribution_seen = false; + for key in grids.by_keys.keys() { + if key.kernel_source != kernel_source { + continue; + } + kernel_seen = true; + if key.quant != quant_name { + continue; + } + quant_seen = true; + if first_distribution.is_none() { + first_distribution = Some(&key.distribution); + } + if key.distribution == distribution { + requested_distribution_seen = true; + break; + } + } + if !kernel_seen { return Err(AicError::PerfDatabase(format!( - "WideEP MoE compute data has no token points for {exact_key:?} at {}", + "WideEP MoE compute data missing for kernel_source={kernel_source:?} at {}", self.data_root.display() ))); } - query_token_curve(by_tokens, num_tokens as f64, &|t| t) + if !quant_seen { + return Err(AicError::PerfDatabase(format!( + "WideEP MoE compute data missing for kernel_source={kernel_source:?} \ + quant={quant_name:?} at {}", + self.data_root.display() + ))); + } + let dist = if requested_distribution_seen { + distribution + } else { + first_distribution.expect("quant_seen implies at least one distribution") + }; + let key = WideEpMoeKey { + kernel_source: kernel_source.to_string(), + quant: quant_name.to_string(), + distribution: dist.to_string(), + topk, + num_experts, + hidden_size, + inter_size, + num_slots, + moe_tp_size, + moe_ep_size, + }; + grids + .by_keys + .get(&key) + .filter(|curve| !curve.is_empty()) + .ok_or_else(|| { + AicError::PerfDatabase(format!( + "WideEP MoE compute data missing for {key:?} at {}", + self.data_root.display() + )) + }) } fn load_compute(&self) -> Result<&WideEpMoeGrids, AicError> { @@ -310,11 +397,16 @@ mod tests { MoeQuantMode::Nvfp4, "power_law_1.01", "wideep_compute_cutlass", + &|t| t, ) .expect("WideEP MoE compute query must succeed"); assert!( (latency - 0.086_009_597_778_320_32).abs() < 1e-6, "expected recorded latency, got {latency}" ); + assert_eq!( + table.available_kernels().expect("kernels list"), + vec!["wideep_compute_cutlass".to_string()] + ); } } diff --git a/src/aiconfigurator/sdk/engine.py b/src/aiconfigurator/sdk/engine.py index 56dc7668f..f71713685 100644 --- a/src/aiconfigurator/sdk/engine.py +++ b/src/aiconfigurator/sdk/engine.py @@ -48,6 +48,7 @@ ContextDeepSeekV4AttentionModule, ContextDSAModule, ContextMLA, + ContextMSAModule, CustomAllReduce, DeepSeekV4MHCModule, ElementWise, @@ -59,6 +60,7 @@ GenerationDeepSeekV4AttentionModule, GenerationDSAModule, GenerationMLA, + GenerationMSAModule, Mamba2Kernel, MLABmm, MLAModule, @@ -86,10 +88,11 @@ # Schema versions must match the Rust crate constants # (`ENGINE_SPEC_SCHEMA_VERSION` / `ENGINE_CONFIG_SCHEMA_VERSION` in `lib.rs`). # ENGINE_SPEC bumped to 2 for the 0.10.0 op-payload layout change (CP + perf-DB -# refactor added serialized `OpSpec` fields); the Rust consumer gates on this -# version before decoding the positional op lists. Keep in lockstep with the -# Rust `ENGINE_SPEC_SCHEMA_VERSION`. -ENGINE_SPEC_SCHEMA_VERSION = 2 +# refactor added serialized `OpSpec` fields), and to 3 when the MSA op variants +# were inserted (bincode enum indices after `DsaGeneration` shifted); the Rust +# consumer gates on this version before decoding the positional op lists. Keep +# in lockstep with the Rust `ENGINE_SPEC_SCHEMA_VERSION`. +ENGINE_SPEC_SCHEMA_VERSION = 3 ENGINE_CONFIG_SCHEMA_VERSION = 1 @@ -334,6 +337,27 @@ def _p2p(op: P2P) -> dict: } +def _msa_module(op: ContextMSAModule | GenerationMSAModule) -> dict: + return { + "name": op._name, + "scale_factor": op._scale_factor, + "num_heads": op._num_heads, + "num_kv_heads": op._num_kv_heads, + "hidden_size": op._hidden_size, + "head_dim": op._head_dim, + "v_head_dim": op._v_head_dim, + "index_n_heads": op._index_n_heads, + "index_head_dim": op._index_head_dim, + "index_topk": op._index_topk, + "block_size": op._block_size, + "kv_cache_dtype": _quant_name(op._kvcache_quant_mode), + "fmha_quant_mode": _quant_name(op._fmha_quant_mode), + "gemm_quant_mode": _quant_name(op._gemm_quant_mode), + "dsa_architecture": op._dsa_architecture, + "dsa_scale_k": op._dsa_scale_k, + } + + def _dsa_module(op: ContextDSAModule | GenerationDSAModule, *, architecture: str) -> dict: # GenerationDSAModule stores `_kv_cache_dtype`; ContextDSAModule stores # `_kvcache_quant_mode` + `_fmha_quant_mode`. The Rust `DsaModuleOp` carries @@ -545,6 +569,10 @@ def recurse(child: Any) -> dict: return {"DsaContext": _dsa_module(op, architecture=architecture)} if isinstance(op, GenerationDSAModule): return {"DsaGeneration": _dsa_module(op, architecture=architecture)} + if isinstance(op, ContextMSAModule): + return {"MsaContext": _msa_module(op)} + if isinstance(op, GenerationMSAModule): + return {"MsaGeneration": _msa_module(op)} if isinstance(op, ContextDeepSeekV4AttentionModule): return {"Dsv4Context": _dsv4_module(op, architecture=architecture)} if isinstance(op, GenerationDeepSeekV4AttentionModule): diff --git a/tests/unit/sdk/test_rust_engine_step.py b/tests/unit/sdk/test_rust_engine_step.py index 01c7dd347..1953046bc 100644 --- a/tests/unit/sdk/test_rust_engine_step.py +++ b/tests/unit/sdk/test_rust_engine_step.py @@ -367,11 +367,11 @@ class _Dsv4Op: assert spec["window_size"] == 2048 -def test_non_silicon_database_mode_falls_back_to_python_step(): - """The compiled engine is SILICON-only (no util_empirical layer); HYBRID / - EMPIRICAL databases must stay on the Python step so both backends give - the SAME answer (parity by delegation) instead of the rust side failing - configs Python fills in empirically.""" +def test_sol_database_modes_fall_back_to_python_step(): + """The compiled engine implements SILICON plus the util-space empirical + layer (HYBRID / EMPIRICAL) — those modes route to Rust and are guarded by + the hybrid parity suite. The SOL/SOL_FULL diagnostic modes stay on the + Python step so both backends give the SAME answer.""" from enum import Enum from aiconfigurator.sdk.config import RuntimeConfig @@ -380,6 +380,9 @@ def test_non_silicon_database_mode_falls_back_to_python_step(): class _Mode(Enum): SILICON = "SILICON" HYBRID = "HYBRID" + EMPIRICAL = "EMPIRICAL" + SOL = "SOL" + SOL_FULL = "SOL_FULL" class _DB: def __init__(self, mode): @@ -390,5 +393,8 @@ def get_default_database_mode(self): rc = RuntimeConfig(engine_step_backend="rust") assert should_use_rust_engine_step(rc, _DB(_Mode.SILICON)) - assert not should_use_rust_engine_step(rc, _DB(_Mode.HYBRID)) + assert should_use_rust_engine_step(rc, _DB(_Mode.HYBRID)) + assert should_use_rust_engine_step(rc, _DB(_Mode.EMPIRICAL)) + assert not should_use_rust_engine_step(rc, _DB(_Mode.SOL)) + assert not should_use_rust_engine_step(rc, _DB(_Mode.SOL_FULL)) assert should_use_rust_engine_step(rc) # no database context -> unchanged From 35aa053043ee1ea7c8a3a476567be9685283e890 Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 19:32:50 +0800 Subject: [PATCH 07/10] fix(rust-core): close all SILICON-path Python/Rust parity gaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Full-surface parity pass against the Python SDK (audit + fixes; hybrid/ empirical port stays out of scope, tracked in #1333). Every fix lands with a pinning test; new end-to-end coverage is bit-identical or within the 1% gate. Numeric fixes (previously divergent or erroring on shipped configs): - WideEP MLA: emit per-rank heads (128 // tp) instead of raw tp_size on the opspec wire; mirror Python's {flashinfer, fa3} attn-backend whitelist. - wideep_moe + wideep context/generation MoE loaders: first-wins -> last-wins (Python direct-assigns; 270 real dup keys on rtx_pro_6000_server). - DSv4: port #1337 into dsv4.rs (generation table keys [kv][gemm] only, SOL fmha derived from kv dtype; drop the architecture key both phases); regenerate stale oracles. - trtllm alltoall: dedicated loader keyed [kernel_source][op_name][quant] [num_nodes][hidden][topk][ne][ep] (the shared MoE loader collapsed 1,556 of 2,096 gb200 rows and keyed a distribution axis Python ignores); port _select_alltoall_kernel + fp8_block->fp8 normalization; full trtllm MoEDispatch branch (sm==100 gate — gb300 is sm103, dp>1 all_gather/ reduce_scatter with quant-compressed volumes, NVL72 reroute). - NVL72 custom-AR -> NCCL reroute + node-clamp + overflow scaling sunk into DB-level query_custom_allreduce_scaled / query_nccl_scaled so every consumer inherits them (mirrors Python's _query_*_table funnels). - Mixed/decode step: rewritten as a LITERAL mirror of Python's three-pass _get_mix_step_latency (pass-1 floor prefix multiplier + raw gen tokens, pass-2 full-isl query / ceil(isl/ctx), pass-3 s = isl + osl//2 + 1) — bit-identical on chunked-prefill/prefix/MTP shapes. - Imbalance-correction scales threaded FFI -> session -> attention ops (were accepted on the wire and hardcoded 1.0). - DSA skip-indexer (GLM-5.2): skip tables + skip SOL terms + CP semantics + full_frac amortization; GLM-5.2 added to SMOKE_CASES. - sglang MoE routing: deepep_moe compute -> wideep tables (with the MoE roofline threaded for beyond-range holds), EPLB int(x*0.8) prefill correction, mxfp4 kernel->quant remaps + the two missing MoEQuantMode members, DeepEP normal/low-latency dispatch flavors + scale_num_tokens. - TrtLLMWideEPMoEDispatch: new WideEpMoeDispatch op + opspec branch (prepare+dispatch pre / combine[_low_precision] post) — trtllm WideEP DeepSeek previously failed opspec conversion outright. - P1: GEMM fp8_static residual floors at the GEMM SOL (not 0) and tags estimated; elementwise floor-divides tokens before the CP split; NCCL message sizes stay fractional (f64) through the interp; encoder partial-RoPE extra; wideep_moe beyond-range holds on the num_slots-aware roofline; generation token count scales by beam_width. Robustness / infra: - OpConversionError now falls back to the Python step (parity by delegation, failure memoized per engine identity) instead of crashing the sweep. - Engine-handle cache key carries un-collapsed quant names + the op-shaping ModelConfig fields (sq vs int8_wo etc. no longer alias one handle); the identity is memoized on the model object (the per-step json.dumps cost tripped the perf gate). - NCCL comm dtype passes through the wire; CommQuantMode grows int8/fp8. - Shared-layer source resolution failures warn loudly instead of silently degrading the Rust side to primary-only rows. - cargo test --lib added to the parity CI job (196 tests; two stale pins fixed — b200 mem_bw 7.7TB/s — and the py:: binding tests self-initialize the interpreter); ModelFamily/PerfDataFilename mirrors corrected; the wrong-composition caller-less overlap_composition helper deleted. Validation: engine-step parity 235/235, compile-engine parity 55/55 (incl. new WideEP sglang + trtllm, imbalance-scale, and chunked-prefill suites), perf gate 3/3, cargo test 196/196, sdk unit 1323/1323. Audit + status: rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md. Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- .github/workflows/build-test.yml | 10 + .../docs/parity-audit-2026-07-14.md | 177 ++++++++++ .../test_compile_engine_parity.py | 269 ++++++++++++++++ .../parity_tests/test_engine_step_parity.py | 12 + rust/aiconfigurator-core/src/common/enums.rs | 126 +++++--- .../src/common/system_spec.rs | 4 +- .../aiconfigurator-core/src/engine/runtime.rs | 237 +++++++++----- rust/aiconfigurator-core/src/engine/spec.rs | 28 +- rust/aiconfigurator-core/src/fpm/tests.rs | 6 + .../src/operators/attention.rs | 32 +- .../src/operators/communication.rs | 58 ++-- rust/aiconfigurator-core/src/operators/dsa.rs | 154 ++++++--- .../aiconfigurator-core/src/operators/dsv4.rs | 4 +- .../src/operators/elementwise.rs | 19 +- .../aiconfigurator-core/src/operators/gemm.rs | 23 +- rust/aiconfigurator-core/src/operators/mod.rs | 3 +- rust/aiconfigurator-core/src/operators/moe.rs | 69 ++++ .../src/operators/moe_dispatch.rs | 237 ++++++++++---- rust/aiconfigurator-core/src/operators/op.rs | 9 +- .../src/operators/overlap.rs | 77 +---- .../src/operators/wideep_mla.rs | 55 +++- .../src/operators/wideep_moe.rs | 33 ++ .../src/perf_database/communication.rs | 96 +++++- .../src/perf_database/dsa.rs | 97 +++++- .../src/perf_database/dsv4.rs | 67 ++-- .../src/perf_database/moe.rs | 24 +- .../src/perf_database/wideep.rs | 301 ++++++++++++++++-- .../src/perf_database/wideep_moe.rs | 56 +++- rust/aiconfigurator-core/src/py.rs | 57 +++- rust/aiconfigurator-core/src/session.rs | 85 ++++- .../sdk/backends/base_backend.py | 136 ++++---- src/aiconfigurator/sdk/engine.py | 148 +++++++-- src/aiconfigurator/sdk/rust_engine_step.py | 137 ++++++-- tests/unit/sdk/test_rust_engine_step.py | 116 ++++++- uv.lock | 27 +- 35 files changed, 2384 insertions(+), 605 deletions(-) create mode 100644 rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml index ff3a30775..8f3c33688 100644 --- a/.github/workflows/build-test.yml +++ b/.github/workflows/build-test.yml @@ -233,6 +233,16 @@ jobs: run: | VIRTUAL_ENV="${GITHUB_WORKSPACE}/.venv" \ "${GITHUB_WORKSPACE}/.venv/bin/maturin" develop --release + - name: Run Rust unit tests (cargo test) + # The crate's own oracle/unit tests were previously NOT in any gate, + # which let "match Python" pins go stale silently (e.g. the b200 + # mem_bw test survived a YAML correction unnoticed). The py:: binding + # tests self-initialize the interpreter (prepare_freethreaded_python), + # so plain `cargo test --lib` covers them; PYO3_PYTHON pins the venv + # interpreter for linking. + env: + PYO3_PYTHON: ${{ github.workspace }}/.venv/bin/python + run: cargo test --release --lib - name: Run Rust/Python engine-step parity suite run: | "${GITHUB_WORKSPACE}/.venv/bin/pytest" -q -rx \ diff --git a/rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md b/rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md new file mode 100644 index 000000000..73d2cf3da --- /dev/null +++ b/rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md @@ -0,0 +1,177 @@ +# Python ↔ Rust Parity Audit — 2026-07-14 + +Full-repo audit of the Rust core (`rust/aiconfigurator-core`) against the Python +SDK, in preparation for full migration. Baseline: `origin/main` @ `2feebac53`. + +> **STATUS UPDATE (same day, this branch): ALL items below are FIXED** except +> the intentionally-excluded HYBRID/EMPIRICAL port (#1333) and the by-design +> Python-only surfaces (AFD, vision-encoder phase, power/energy, `.txt` legacy +> loading, scheduling/summary orchestration). Every P0/P1 fix landed with a +> pinning test; new end-to-end parity coverage: GLM-5.2 skip-indexer smoke +> case, sglang WideEP/DeepEP suite (h200/0.5.6.post2), trtllm WideEP alltoall +> suite (gb200), chunked-prefill mixed-step cases, imbalance-scale threading +> cases — all bit-identical or within the 1% gate. The mixed/decode-step +> composition was rewritten as a literal mirror of Python's three-pass +> `_get_mix_step_latency` (bit-identical on chunked/prefix/MTP shapes). +> `cargo test --lib` (196 tests, incl. regenerated oracles) is now a CI gate. +> `OpConversionError` now falls back to the Python step instead of crashing; +> the engine-handle cache key carries the un-collapsed quant + ModelConfig +> identity. See the fix-order list at the bottom — items 1-12 all landed. + +**Scope exclusion:** HYBRID / EMPIRICAL / `util_empirical` differences are out of +scope (Rust is SILICON-only by design; delegation gate at +`sdk/rust_engine_step.py:224-246`; port tracked in issue #1333 with a separate +PR in flight). + +## Verdict + +**Not yet answer-identical.** The dense-model core (attention, MLA, GEMM +interpolation, standard MoE, communication, mamba/GDN, DSA full-layer, DSv4 +context) is a faithful, oracle-verified port, and the June 2026 full-matrix +parity gate (2,158 entries, 0 regression) still holds for that surface. But the +audit found **6 live high-severity numeric divergences** (concentrated in +WideEP / trtllm-alltoall / DSv4-generation / mixed-step composition), a set of +narrow-trigger divergences, and several whole features with no Rust +counterpart. Deleting the Python query layer today would regress WideEP +DeepSeek, GLM-5.2, DSv4 fp8-decode, GB200-NVL72 comm, and chunked-prefill agg +sweeps. + +Items 1, 3, 4, 6 below were independently found by two audit passes (operators ++ DB layers) and hand-verified in source. + +--- + +## P0 — live numeric divergences (fix before any default flip) + +### 1. WideEP MLA: bridge writes `tp_size` into the `num_heads` table axis ✅ hand-verified +- Python queries at `num_heads = 128 // tp_size` (`sdk/operations/mla.py:1195` gen, `:1464` ctx). +- The spec emitter writes raw tp: `"num_heads": op._tp_size` (`sdk/engine.py:463,475` — the comment claims "per-rank head split" but the value is tp). +- Rust passes it straight to the heads axis (`operators/wideep_mla.rs:85,226` → `perf_database/wideep_mla.rs:163-170,214-218`). tp=1 ⇒ Python reads the heads=128 slice, Rust queries heads=1 ⇒ out-of-grid SOL-rescale, order-of-magnitude error. +- Rust unit tests mask it by constructing ops with `num_heads=128` directly (`operators/wideep_mla.rs:121-131`). +- Fix: one line in `engine.py` (`128 // op._tp_size`) + a parity smoke case with tp>1 WideEP. + +### 2. trtllm `enable_alltoall` dispatch: loader collapses the table; query errors on shipped data +- Python keys `[kernel_source][op_name][quant][num_nodes][...]` and queries `alltoall_dispatch` + `alltoall_combine` separately (`sdk/operations/moe.py:1128-1172, 2921-2927, 1976-2042`). +- Rust loader drops `kernel_source`/`op_name`/`num_nodes` (`perf_database/wideep.rs:637-693`) — verified **1,556 / 2,096 rows in gb200 `trtllm_alltoall_perf.parquet` collide** with differing latencies (first-wins). +- Rust queries distribution `"uniform"` (`operators/moe_dispatch.rs:370-378`) but shipped data is `"balanced"` ⇒ hard error. +- Missing `fp8_block → fp8` normalization (`moe.py:1963-1973`). + +### 3. DSv4 generation: PR #1337 (c4a574a65) ported to `mla.rs` but not `dsv4.rs` +- Python derives generation SOL fmha from KV dtype (`sdk/operations/dsv4.py:1246-1254`); Rust uses the raw label (`operators/dsv4.rs:414-424`, `perf_database/dsv4.rs:368-381`). +- Python's generation table key ignores `mla_dtype` (`dsv4.py:1938-1966`); Rust requires exact fmha+arch match (`perf_database/dsv4.rs:855-880`) ⇒ hard error on fp8-labeled decode where Python resolves. +- The Rust oracle fixtures (`perf_database/dsv4.rs:1237-1284`) were generated pre-#1337 and now certify parity against stale Python. Regenerate after the fix. + +### 4. `wideep_moe_perf` dedup: first-wins vs Python's last-wins +- Python direct-assigns (last row wins, `moe.py:2819-2825`); Rust `or_insert` (first wins) with an incorrect "parity" comment (`perf_database/wideep_moe.rs:258-264`). +- Verified **270 duplicate keys with differing latencies** in rtx_pro_6000_server trtllm 1.3.0rc10. +- Same bug class was already fixed for MLA modules (`mla.rs:653-674`); wideep was missed. Fix: `or_insert` → `insert`. + +### 5. DSA skip-indexer (GLM-5.2) amortization missing +- Python prices layers as `w*full + (1-w)*skip` (`sdk/operations/dsa.py:774-781, 1412-1420`, CP `:754-759`); Rust drops skip rows at load (`perf_database/dsa.rs:774-779`) and has no `full_frac` field (`operators/dsa.rs:24-44`); the emitter neither sends nor rejects (`engine.py:337-360`). +- Triggered by `model_configs/nvidia--GLM-5.2-NVFP4_config.json` (`index_topk_freq: 4`) on sglang ⇒ **silent** overestimate. DeepSeek-V3.2 / GLM-5 (full_frac=1) unaffected. + +### 6. Sequence-imbalance correction scales silently ignored in Rust +- Python threads `seq_imbalance_correction_scale` / `gen_seq_imbalance_correction_scale` into every attention query (`backends/base_backend.py:331,372,950-965`). +- Rust accepts them on the wire but hardcodes `1.0` (`session.rs:45-46,74-75,146-147,178-179,203-204`; admitted at `engine/runtime.rs:33-35`); the facade passes them through implying support (`rust_engine_step.py:275-276`). +- Latent today (no shipped task sets ≠1.0) but ungated: either thread them through `RuntimeContext` or refuse rust routing when ≠1.0. + +### 7. Mixed-step composition (agg / chunked prefill) — three formula gaps +Largest first: +- **Pass-2 context attention under chunked prefill** (`ctx_tokens < isl`): Python queries at full per-request isl then divides by `ceil(isl/ctx_tokens)` (`base_backend.py:984-1003`); Rust queries at `ctx_tokens` with scale 1 (`session.rs:153-183`, self-documented limitation at `:158-164`). Attention is superlinear in s. +- **Pass-1 prefix/nextn accounting**: Python `prefix * floor(ctx_tokens/isl)`, raw combined tokens (`base_backend.py:954-964`); Rust `prefix * ceil(...)` (`runtime.rs:348-349`) and inflates pass-1 tokens by `(nextn+1)` (`runtime.rs:362`, `session.rs:129`). +- **±1 kv-token convention**: Python decode queries at `s = isl + osl//2 + 1` (`base_backend.py:371,1075-1085`); Rust at `isl + osl/2` (`runtime.rs:403`, deliberate legacy-FPM convention). Sub-tolerance; either align or codify. + +### 8. sglang MoE routing gaps in the bridge +- `deepep_moe` backend: Python routes MoE compute to wideep tables (`moe.py:685-691`); Rust `MoeOp` always reads `moe_perf` — `WideEpTable::query_{context,generation}_moe` exist but are caller-less (`wideep.rs:179-238`). +- EPLB prefill correction `int(num_tokens*0.8)` (`moe.py:684`) dropped by the emitter (`engine.py:241-255`). +- sglang WideEP dispatch flavor: emitter maps everything non-trtllm to CustomAllReduce (`engine.py:258-274`, documented limitation) — DeepEP normal/low-latency never emitted ⇒ **silently wrong** rather than rejected. +- MoE mxfp4 kernel→quant remap unported: Python reroutes `w4a8_mxfp4_mxfp8` / `w4a16_mxfp4` to `_trtllm` / `_cutlass` variants at load (`moe.py:2465-2474`); those quant names don't exist in Rust (`enums.rs:120-129`). Affects shipped sglang 0.5.14 data on 7 systems. + +--- + +## P1 — narrow-trigger / dormant divergences + +| # | Area | Python | Rust | Trigger | +|---|---|---|---|---| +| 9 | GEMM fp8_static residual floor | clamp to GEMM SOL (`gemm.py:787-800`) | clamp to 0 (`operators/gemm.rs:92-94`, stale comment) | fp8_static when subtraction dips below SOL | +| 10 | GB200/NVL72 comm reroute | custom-AR → NCCL when `num_gpus_per_node==72 && tp>4` (`communication.py:187-190`; same in `moe.py:1148-1155`) | always custom-AR (`operators/communication.rs:60-87`, `moe_dispatch.rs:396-406`); also missing `sm_version==100` gate (`moe.py:1099`) | GB200/GB300 NVL72 (Rust errors or mis-prices) | +| 11 | ElementWise `scale_num_tokens` | floor tokens then split (`elementwise.py:56-59`) | scale folded into bytes (`engine.py:153-154`, `elementwise.rs:47-48`) | exact only when scale divides tokens; live on DSv3.2 CP add_norm | +| 12 | NCCL fractional message size | float element counts (`models/gemma4.py:374`) | truncates to u64 (`communication.rs:138`) | gemma4 CP KV all-gather | +| 13 | Encoder partial-RoPE extra | `factor * 2 * mem_op(...) * 1.1` (`attention.py:1100-1107`) | field absent (`operators/attention.rs:234-296`) | Qwen3-VL-class (moot while encoder stays Python-side) | +| 14 | WideEP MoE beyond-range hold | num_slots-aware roofline (`moe.py:1782-1797`) | linear token proxy (`wideep_moe.rs:103-109,196`) | out-of-grid tokens only | +| 15 | MLA BMM quant fallback granularity | top-level (`mla.py:564`) | per-(quant,op) (`mla.rs:242-248`) — Rust answers where Python raises | partial-data dirs | +| 16 | Interp exact-distance tie-break | insertion order (`perf_interp/engine.py:278,386`) | sorted keys (`perf_interp.rs:424-432,573-580`) | exact ties only | +| 17 | beam_width | applied to gen batch (`base_backend.py:368-369`) | never applied (`runtime.rs:39-41`), ungated | beam>1 (unexercised) | + +Known-and-accepted (from the June gate, still present): GEMM large-n +extrapolation ~30% on the logits op; context-attention ragged-grid ~1.5%. + +--- + +## Unported features (no Rust counterpart) + +**Loud (raise `OpConversionError` at spec build — but note this CRASHES the +sweep instead of falling back; `base_backend.py:422` has no try/except):** +- MiniMax-M3: `Context/GenerationMSAModule` (`operations/msa.py:258,302`). In practice usually blocked earlier by the SILICON gate (MSA data is HYBRID). +- DeepSeek-V4 MegaMoE (sglang): `DeepSeekV4MegaMoEModule` (`dsv4.py:1417`); `query_dsv4_megamoe_module` (`perf_database.py:2761`). +- TRT-LLM WideEP DeepSeek / DSv3.2: `TrtLLMWideEPMoEDispatch` (`moe.py:1872`, used by `models/deepseek.py:772-1009`, `models/deepseek_v32.py:573-692`). The paired compute op converts, the dispatch does not — cheap to add, Rust dispatch machinery exists. + +**Python-only by design (must stay or be ported later):** +- AFD estimation end-to-end (`inference_session.py:937-1080`, `afd_transfer.py`, `afd_partition.py`, `AFDConfig`). +- Vision encoder phase (`base_backend.py:251-301`; spec excludes encoder ops `engine.py:807-818`; isl-folding compensates for the decoder). +- `query_gemm(database_mode=SOL/SOL_FULL)` — no `DatabaseMode` in the crate (root cause of #9). +- Power/energy everywhere (Rust is latency-only, `perf_interp.rs:23-24`; rust-routed runs return zero energy and a collapsed single-key breakdown — `e2e_power_avg` and per-op detail silently degrade). +- Support-matrix / feasibility / version discovery (`_LazySupportMatrix`, `get_supported_databases`, `get_latest_database_version` incl. `INCOMPLETE.txt` rejection, `get_database_view`) — Rust takes explicit tuples + Python-precomputed `perf_db_sources`. +- `.txt` legacy data loading: Rust is parquet-only (`parquet_loader.rs:42` vs `operations/base.py:62-70`) — txt-only system dirs (e.g. `local_systems/`, new a100 drops) answer in Python, error in Rust. +- Scheduling/summary math around the steps (`run_agg` composition, TTFT/TPOT/throughput, memory/OOM) — by design, "Python builds & orchestrates". + +--- + +## Infra / cross-cutting risks + +1. **Engine-handle cache collisions**: `_ENGINE_HANDLE_CACHE` keys on `_engine_config_json` (`rust_engine_step.py:399-432`) but the op payload comes from the actual model. Collapsed dtypes (`sq`→`int8_wo`, `fp8_ootb`→`fp8`, four 4-bit modes→`int4_wo`, two DSv4 MoE modes→`None`) and ~10 identity-omitted `ModelConfig` fields (`comm_quant_mode`, `cp_style`, `moe_backend`, `attention_backend`, `enable_wideep`, `enable_eplb`, `workload_distribution`, `overwrite_num_layers`, `sms`, `wideep_num_slots`) can alias two different models to one cached handle ⇒ silent wrong latencies. +2. **`comm_quant_mode` hardcoded `"half"`** in the emitter (`engine.py:294,310,321`) while Python NCCL queries use the real dtype (`communication.py:385-431`). Rust `CommQuantMode` only has `Half` (`enums.rs:203-205`, stale comment — Python has int8/fp8 at `common.py:1104-1106`). +3. **Silent shared-layer drop**: `_compute_perf_db_sources` swallows all exceptions → `{}` (`engine.py:635-636`); Rust then loads primary-only rows while Python uses the shared layer in the same run. No error surfaced. +4. **Enum drift**: Rust missing MoE `w4a8_mxfp4_mxfp8_trtllm` / `w4a16_mxfp4_cutlass`; `ModelFamily` has `GEMMA4MOE` vs Python `GEMMA4MIX` and lacks `MINIMAXM3`; `PerfDataFilename` still `.txt` (dead code, but its test claims Python parity). +5. **Stale Rust unit tests not in CI**: `parse_b200_sxm` asserts mem_bw 8.0 TB/s but PR #1246 corrected the YAML to 7.7 ⇒ the crate's "match Python" tests currently fail and nothing notices. Add `cargo test` to the gating CI. +6. **system_spec strictness**: Rust serde defaults (`mem_bw_empirical_scaling_factor`=1.0 etc., `system_spec.rs:36-39,77`) where Python KeyErrors — a malformed YAML parses silently with different physics. +7. **`overlap_composition`** (`operators/overlap.rs:20-40`) implements the opposite composition of Python's OverlapOp and has no caller — delete or fix before anyone wires it. (The production `op.rs:291-321` OverlapOp is correct.) +8. **Rust-only leniency**: missing `attn_backend` whitelist in WideEP MLA (`mla.py:1459-1460` raises, Rust succeeds); beam>1 accepted silently. + +--- + +## What is already solid + +- Parity CI gates: 40-case engine-step smoke + compile-engine suite + perf gate, every PR (`.github/workflows/build-test.yml:204-258`), 1% rtol, error-symmetry contract. +- Full-matrix scan (`tools/support_matrix/scan_rust_parity.py`, 2,158 entries, probe + pareto layers): gate CLOSED 2026-06-16 with 0 regression, max drift 0.41% ttft / 0.18% tpot (`docs/parity-scan-report.md`). +- Memory: single implementation — Rust `memory.rs` is a pure forwarder into Python (`fetch_python_estimate`); only shared constant is the 0.80 naive KV reservation. +- Interpolation core: step-for-step port, no scipy left; most table layers carry 1e-9 cross-language oracles. +- Static (`run_static`) composition: stride quadrature, `(nextn+1)` scaling, prefix handling formula-identical. + +Note the tension: the smoke matrix + June scan passed while P0 items 1–5 exist, +because the affected configs (WideEP tp>1, trtllm alltoall on gb200, GLM-5.2 +skip-indexer, post-#1337 DSv4 fp8 decode, deep chunked-prefill shapes) are +outside or under-sampled in the scan matrix, and Rust-side oracles were +constructed to match (heads=128, pre-#1337 fixtures). **Every P0 fix should add +its trigger config to SMOKE_CASES / the scan matrix.** + +## Migration gates (from `docs/python-dedup-plan.md`, status unchanged) + +- Gate 1 (default flip): blocked by #1333 (SILICON-only) **and now by P0 1–8 above**. +- Gate 2 (golden oracle capture): NOT STARTED — must land before any Python deletion, or the regression detector disappears. +- Gate 3 (delete Python latency code): blocked on 1+2. + +## Suggested fix order + +1. `engine.py:463,475` → `128 // op._tp_size` (one line) + tp>1 WideEP smoke case. +2. `wideep_moe.rs` `or_insert`→`insert` (one line) + dup-key regression test. +3. Port #1337 into `dsv4.rs` + regenerate DSV4 oracles. +4. trtllm alltoall loader keys (`op_name`/`kernel_source`/`num_nodes`) + distribution + fp8_block normalization. +5. Imbalance scales: thread through `session.rs` or gate in `should_use_rust_engine_step`. +6. Mixed-step pass-1/pass-2 accounting (or codify as accepted tolerance with a test pinning the delta). +7. DSA `full_frac` port or emit-time rejection for GLM-5.2. +8. sglang deepep/EPLB/mxfp4 + WideEP dispatch flavor: port or reject at emit (no silent mis-model). +9. `TrtLLMWideEPMoEDispatch` opspec conversion. +10. Wrap `build_engine_spec_json` call site (`base_backend.py:422`) so `OpConversionError` falls back to Python instead of crashing. +11. Widen `_engine_config_json` identity (add collapsed-dtype + missing ModelConfig fields) to kill cache aliasing. +12. Add `cargo test` to gating CI; fix `parse_b200_sxm`; update stale enum mirrors. diff --git a/rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py b/rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py index 823f0e982..52fd9aee8 100644 --- a/rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py +++ b/rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py @@ -325,6 +325,29 @@ def test_mixed_step(self, case: EngineStepParityCase) -> None: py_val = _python_mixed(case) _assert_within("mixed_step", py_val, new_val, backend=case.backend_name) + @pytest.mark.parametrize( + "ctx_tokens,gen_tokens,isl,osl,prefix", + [ + (512, 4, 4096, 128, 0), # chunked prefill: ctx_tokens < isl + (512, 4, 4096, 128, 256), # chunked + cached prefix + (300, 7, 1000, 64, 100), # ragged chunk + prefix + decode overlap + ], + ) + def test_mixed_step_chunked_prefill(self, ctx_tokens, gen_tokens, isl, osl, prefix) -> None: + """Chunked prefill (ctx_tokens < isl) was the largest pre-rewrite + composition gap: Python queries context attention at the FULL per-req + isl then divides by ceil(isl/ctx), the old Rust queried the chunk + directly. The rewritten three-pass mirror must match exactly.""" + case = _SUBSET_BY_ID["minimax-m25-b200-vllm-019-isl1024-osl2"].values[0] + model, backend, database = _build_python_model(case) + rc = config.RuntimeConfig(batch_size=1, beam_width=1, isl=isl, osl=osl, prefix=prefix) + py_val, _, _, _ = _quiet( + backend._get_mix_step_latency, model, database, rc, ctx_tokens, gen_tokens, isl, osl, prefix + ) + handle = _compile_handle(case) + new_val = handle.mixed_step_latency(ctx_tokens, gen_tokens, isl, osl, prefix) + _assert_within("mixed_step_chunked", float(py_val), new_val, backend=case.backend_name) + class TestCompileEngineDecodeStepParity: @pytest.mark.parametrize("case", _SUBSET_CASES) @@ -335,6 +358,252 @@ def test_decode_step(self, case: EngineStepParityCase) -> None: _assert_within("decode_step", py_val, new_val, backend=case.backend_name) +# --------------------------------------------------------------------------- # +# 2b. Imbalance-correction scale threading (session.rs used to hardcode 1.0). +# --------------------------------------------------------------------------- # + + +class TestImbalanceScaleParity: + """Non-1.0 seq/gen imbalance-correction scales must produce identical + Python and Rust numbers. Regression for the session.rs hardcode: the wire + accepted the scales but every RuntimeContext pinned them to 1.0, so any + task setting them diverged silently on the rust path.""" + + _CASE_ID = "minimax-m25-b200-vllm-019-isl1024-osl2" + _CTX_SCALE = 1.3 + _GEN_SCALE = 0.85 + + def _case(self) -> EngineStepParityCase: + return _SUBSET_BY_ID[self._CASE_ID].values[0] + + def test_static_scales_thread_through(self) -> None: + case = self._case() + model, backend, database = _build_python_model(case) + rc = config.RuntimeConfig( + batch_size=case.batch_size, + beam_width=1, + isl=case.isl, + osl=max(case.osl, 2), + prefix=case.prefix, + seq_imbalance_correction_scale=self._CTX_SCALE, + gen_seq_imbalance_correction_scale=self._GEN_SCALE, + ) + ctx_lat, _, gen_lat, _, _, _ = _quiet(backend._run_static_breakdown, model, database, rc, "static", 1) + py_ctx = float(sum(ctx_lat.values())) + py_gen = float(sum(gen_lat.values())) + + handle = _compile_handle(case) + new_ctx, new_gen, _ = handle.run_static( + batch_size=case.batch_size, + isl=case.isl, + osl=max(case.osl, 2), + prefix=case.prefix, + seq_imbalance_correction_scale=self._CTX_SCALE, + gen_seq_imbalance_correction_scale=self._GEN_SCALE, + stride=1, + ) + _assert_within("static_ctx@scale", py_ctx, new_ctx, backend=case.backend_name) + _assert_within("static_gen@scale", py_gen, new_gen, backend=case.backend_name) + + # The scales must actually bite: a scaled run differs from unscaled. + base_ctx, base_gen, _ = handle.run_static( + batch_size=case.batch_size, isl=case.isl, osl=max(case.osl, 2), prefix=case.prefix, stride=1 + ) + assert new_ctx != base_ctx, "ctx scale did not affect the rust static path" + assert new_gen != base_gen, "gen scale did not affect the rust static path" + + def test_mixed_and_decode_scales_thread_through(self) -> None: + case = self._case() + model, backend, database = _build_python_model(case) + rc = config.RuntimeConfig( + batch_size=case.batch_size, + beam_width=1, + isl=case.isl, + osl=max(case.osl, 2), + prefix=case.prefix, + seq_imbalance_correction_scale=self._CTX_SCALE, + gen_seq_imbalance_correction_scale=self._GEN_SCALE, + ) + py_mixed, _, _, _ = _quiet( + backend._get_mix_step_latency, + model, + database, + rc, + case.isl, + case.batch_size, + case.isl, + max(case.osl, 2), + case.prefix, + ) + py_decode, _, _, _ = _quiet( + backend._get_genonly_step_latency, + model, + database, + rc, + case.batch_size, + case.isl, + max(case.osl, 2), + ) + + handle = _compile_handle(case) + new_mixed = handle.mixed_step_latency( + case.isl, + case.batch_size, + case.isl, + max(case.osl, 2), + case.prefix, + seq_imbalance_correction_scale=self._CTX_SCALE, + gen_seq_imbalance_correction_scale=self._GEN_SCALE, + ) + new_decode = handle.decode_step_latency( + case.batch_size, + case.isl, + max(case.osl, 2), + gen_seq_imbalance_correction_scale=self._GEN_SCALE, + ) + _assert_within("mixed_step@scale", float(py_mixed), new_mixed, backend=case.backend_name) + _assert_within("decode_step@scale", float(py_decode), new_decode, backend=case.backend_name) + + +# --------------------------------------------------------------------------- # +# 2c. SGLang WideEP (deepep_moe) — MLA + MoE + DeepEP dispatch routing. +# --------------------------------------------------------------------------- # + + +class TestWideEpDeepEpParity: + """SGLang WideEP DeepSeek (moe_backend=deepep_moe) end-to-end parity. + + Covers three previously-divergent surfaces at once: the WideEP MLA + per-rank-heads table coordinate (tp=8 -> heads=16; the bridge used to emit + raw tp), the deepep MoE compute routing (Rust used to read `moe_perf` + where Python reads the wideep context/generation tables), and the DeepEP + dispatch flavor emission (the emitter used to map every sglang dispatch to + CustomAllReduce). Data lives on h200_sxm/sglang/0.5.6.post2 (the only + shipped version with the deepep dispatch parquets).""" + + _MODEL = "deepseek-ai/DeepSeek-V3" + _SYSTEM = "h200_sxm" + _VERSION = "0.5.6.post2" + + def _build(self): + from aiconfigurator.sdk import common + + database = _quiet(perf_database.get_database, self._SYSTEM, "sglang", self._VERSION) + if database is None: + pytest.skip(f"no perf database for {self._SYSTEM}/sglang/{self._VERSION}") + model_config = config.ModelConfig( + tp_size=8, + moe_tp_size=1, + moe_ep_size=8, + moe_backend="deepep_moe", + attention_backend="flashinfer", + gemm_quant_mode=common.GEMMQuantMode.fp8_block, + moe_quant_mode=common.MoEQuantMode.fp8_block, + kvcache_quant_mode=common.KVCacheQuantMode.fp8, + fmha_quant_mode=common.FMHAQuantMode.fp8_block, + ) + model = _quiet(get_model, self._MODEL, model_config, "sglang") + backend = get_backend("sglang") + spec_json = _quiet( + engine.build_engine_spec_json, + model, + model_path=self._MODEL, + system=self._SYSTEM, + backend="sglang", + backend_version=self._VERSION, + kv_block_size=None, + systems_path=None, + nextn=0, + nextn_accept_rates=None, + database=database, + ) + import aiconfigurator_core + + handle = engine.EngineHandle(bytes(aiconfigurator_core.engine_spec_bincode_from_json(spec_json))) + return model, backend, database, handle + + def test_wideep_static_parity(self) -> None: + model, backend, database, handle = self._build() + rc = config.RuntimeConfig(batch_size=1, beam_width=1, isl=1024, osl=4, prefix=0) + ctx_lat, _, gen_lat, _, _, _ = _quiet(backend._run_static_breakdown, model, database, rc, "static", 1) + py_ctx, py_gen = float(sum(ctx_lat.values())), float(sum(gen_lat.values())) + new_ctx, new_gen, _ = handle.run_static(batch_size=1, isl=1024, osl=4, prefix=0, stride=1) + _assert_within("wideep_static_ctx", py_ctx, new_ctx, backend="sglang") + _assert_within("wideep_static_gen", py_gen, new_gen, backend="sglang") + + def test_wideep_mixed_and_decode_parity(self) -> None: + model, backend, database, handle = self._build() + rc = config.RuntimeConfig(batch_size=1, beam_width=1, isl=1024, osl=4, prefix=0) + py_mixed, _, _, _ = _quiet(backend._get_mix_step_latency, model, database, rc, 1024, 2, 1024, 4, 0) + py_decode, _, _, _ = _quiet(backend._get_genonly_step_latency, model, database, rc, 2, 1024, 4) + new_mixed = handle.mixed_step_latency(1024, 2, 1024, 4, 0) + new_decode = handle.decode_step_latency(2, 1024, 4) + _assert_within("wideep_mixed", float(py_mixed), new_mixed, backend="sglang") + _assert_within("wideep_decode", float(py_decode), new_decode, backend="sglang") + + +# --------------------------------------------------------------------------- # +# 2d. TRT-LLM WideEP (NVLink Two-Sided alltoall) — gb200. +# --------------------------------------------------------------------------- # + + +class TestTrtllmWideEpParity: + """TRT-LLM WideEP DeepSeek (enable_wideep, attention_dp=8) on gb200. + + Covers the `TrtLLMWideEPMoEDispatch` port (prepare+dispatch pre / + combine post through the trtllm_alltoall table, kernel auto-selected as + NVLinkTwoSided via moe_backend="wideep") and the alltoall loader keying + (kernel_source/op_name/num_nodes — the pre-fix loader collapsed 1,556 of + 2,096 gb200 rows). This path used to fail opspec conversion entirely + (`TrtLLMWideEPMoEDispatch` had no `_to_opspec` branch).""" + + def _build(self): + from aiconfigurator.sdk import common + + database = _quiet(perf_database.get_database, "gb200", "trtllm", "1.3.0rc10") + if database is None: + pytest.skip("no perf database for gb200/trtllm/1.3.0rc10") + model_config = config.ModelConfig( + tp_size=1, + attention_dp_size=8, + moe_tp_size=1, + moe_ep_size=8, + enable_wideep=True, + gemm_quant_mode=common.GEMMQuantMode.nvfp4, + moe_quant_mode=common.MoEQuantMode.nvfp4, + kvcache_quant_mode=common.KVCacheQuantMode.fp8, + fmha_quant_mode=common.FMHAQuantMode.bfloat16, + ) + model = _quiet(get_model, "deepseek-ai/DeepSeek-V3", model_config, "trtllm") + backend = get_backend("trtllm") + spec_json = _quiet( + engine.build_engine_spec_json, + model, + model_path="deepseek-ai/DeepSeek-V3", + system="gb200", + backend="trtllm", + backend_version="1.3.0rc10", + kv_block_size=None, + systems_path=None, + nextn=0, + nextn_accept_rates=None, + database=database, + ) + import aiconfigurator_core + + handle = engine.EngineHandle(bytes(aiconfigurator_core.engine_spec_bincode_from_json(spec_json))) + return model, backend, database, handle + + def test_trtllm_wideep_static_parity(self) -> None: + model, backend, database, handle = self._build() + rc = config.RuntimeConfig(batch_size=1, beam_width=1, isl=1024, osl=4, prefix=0) + ctx_lat, _, gen_lat, _, _, _ = _quiet(backend._run_static_breakdown, model, database, rc, "static", 1) + py_ctx, py_gen = float(sum(ctx_lat.values())), float(sum(gen_lat.values())) + new_ctx, new_gen, _ = handle.run_static(batch_size=1, isl=1024, osl=4, prefix=0, stride=1) + _assert_within("trtllm_wideep_static_ctx", py_ctx, new_ctx, backend="trtllm") + _assert_within("trtllm_wideep_static_gen", py_gen, new_gen, backend="trtllm") + + # --------------------------------------------------------------------------- # # 3. Determinism across rayon thread counts. # --------------------------------------------------------------------------- # diff --git a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py index 9be25d664..f084602a1 100644 --- a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py +++ b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py @@ -214,6 +214,18 @@ class EngineStepParityCase: ), id="glm5-b200-sglang-0514-isl16384-osl2", ), + # GLM-5.2 shared-index amortization (full_frac = 21/78): per-layer DSA is + # w*full + (1-w)*skip using the skip_indexer rows collected in the same + # parquet. Tripwire for the Rust skip-table port — dropping the skip rows + # (the pre-port behavior) silently overestimates every GLM-5.2 sweep. + pytest.param( + EngineStepParityCase( + model_path="nvidia/GLM-5.2-NVFP4", + backend_name="sglang", + backend_version="0.5.14", + ), + id="glm52-b200-sglang-0514-isl1024-osl2", + ), # Phase 4 D7-F: backend coverage for newly-ported families. The # builders are backend-independent (the per-backend conditional # `_attn_dp` / `_tp_allreduce` branches are handled inside each diff --git a/rust/aiconfigurator-core/src/common/enums.rs b/rust/aiconfigurator-core/src/common/enums.rs index b53c4011f..dc42fd6c1 100644 --- a/rust/aiconfigurator-core/src/common/enums.rs +++ b/rust/aiconfigurator-core/src/common/enums.rs @@ -126,6 +126,15 @@ pub enum MoeQuantMode { Nvfp4, W4a16Mxfp4, W4a8Mxfp4Mxfp8, + /// Blackwell trtllm-gen MXFP4xMXFP8 kernel rows + /// (`sglang_mxfp4_flashinfer_trtllm_moe`); a distinct precision from the + /// flashinfer cutedsl kernel logged under the same `moe_dtype`. Loader + /// remaps rows into this mode (mirrors Python `load_moe_data`). + W4a8Mxfp4Mxfp8Trtllm, + /// Hopper flashinfer cutlass SM90 mixed-GEMM rows + /// (`sglang_flashinfer_cutlass_moe`); loader-remapped from + /// `w4a16_mxfp4` (mirrors Python `load_moe_data`). + W4a16Mxfp4Cutlass, } impl MoeQuantMode { @@ -141,6 +150,12 @@ impl MoeQuantMode { Self::W4a8Mxfp4Mxfp8 => { QuantMapping { memory: 0.5, compute: 2.0, name: "w4a8_mxfp4_mxfp8" } } + Self::W4a8Mxfp4Mxfp8Trtllm => { + QuantMapping { memory: 0.5, compute: 2.0, name: "w4a8_mxfp4_mxfp8_trtllm" } + } + Self::W4a16Mxfp4Cutlass => { + QuantMapping { memory: 0.5, compute: 1.0, name: "w4a16_mxfp4_cutlass" } + } } } @@ -196,18 +211,22 @@ impl KvCacheQuantMode { } } -/// Collective communication quantization mode. Mirrors `common.CommQuantMode`. -/// Only the `half` variant is present in Python today. +/// Collective communication quantization mode. Mirrors `common.CommQuantMode` +/// (half / int8 / fp8; byte widths 2 / 1 / 1). #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum CommQuantMode { Half, + Int8, + Fp8, } impl CommQuantMode { pub fn mapping(self) -> QuantMapping { match self { Self::Half => QuantMapping { memory: 2.0, compute: 0.0, name: "half" }, + Self::Int8 => QuantMapping { memory: 1.0, compute: 0.0, name: "int8" }, + Self::Fp8 => QuantMapping { memory: 1.0, compute: 0.0, name: "fp8" }, } } @@ -231,7 +250,8 @@ pub enum ModelFamily { NemotronH, HybridMoe, Qwen35, - Gemma4Moe, + Gemma4Mix, + MinimaxM3, Qwen3Vl, Qwen3VlMoe, } @@ -252,7 +272,8 @@ impl ModelFamily { Self::NemotronH => "NEMOTRONH", Self::HybridMoe => "HYBRIDMOE", Self::Qwen35 => "QWEN35", - Self::Gemma4Moe => "GEMMA4MOE", + Self::Gemma4Mix => "GEMMA4MIX", + Self::MinimaxM3 => "MINIMAXM3", Self::Qwen3Vl => "QWEN3VL", Self::Qwen3VlMoe => "QWEN3VL_MOE", } @@ -265,10 +286,12 @@ impl fmt::Display for ModelFamily { } } -/// Performance data CSV filenames. Mirrors `common.PerfDataFilename`. +/// Performance data parquet basenames. Mirrors `common.PerfDataFilename` +/// (which migrated from `.txt` to `.parquet`). /// -/// Used by `perf_database/` modules to locate per-table CSV files under -/// `systems/data////`. +/// NOTE: currently a reference mirror only — the `perf_database/` loaders +/// hardcode their basenames (via `resolve_op_sources`) rather than routing +/// through this enum. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub enum PerfDataFilename { Gemm, @@ -305,47 +328,52 @@ pub enum PerfDataFilename { Dsv4HcaGenerationModule, Dsv4PagedMqaLogitsModule, Dsv4HcaAttnModule, + Dsv4CsaAttnModule, + Dsv4CsaTopkCalib, + Dsv4MegamoeModule, } impl PerfDataFilename { - /// Stable CSV basename used by the perf DB loaders. Matches Python's - /// `PerfDataFilename` enum values. + /// Stable parquet basename. Matches Python's `PerfDataFilename` values. pub fn as_str(self) -> &'static str { match self { - Self::Gemm => "gemm_perf.txt", - Self::Nccl => "nccl_perf.txt", - Self::Oneccl => "oneccl_perf.txt", - Self::GenerationAttention => "generation_attention_perf.txt", - Self::ContextAttention => "context_attention_perf.txt", - Self::EncoderAttention => "encoder_attention_perf.txt", - Self::ContextMla => "context_mla_perf.txt", - Self::GenerationMla => "generation_mla_perf.txt", - Self::MlaBmm => "mla_bmm_perf.txt", - Self::Moe => "moe_perf.txt", - Self::CustomAllreduce => "custom_allreduce_perf.txt", - Self::WideepContextMla => "wideep_context_mla_perf.txt", - Self::WideepGenerationMla => "wideep_generation_mla_perf.txt", - Self::WideepContextMoe => "wideep_context_moe_perf.txt", - Self::WideepGenerationMoe => "wideep_generation_moe_perf.txt", - Self::WideepDeepepNormal => "wideep_deepep_normal_perf.txt", - Self::WideepDeepepLl => "wideep_deepep_ll_perf.txt", - Self::WideepMoeCompute => "wideep_moe_perf.txt", - Self::TrtllmAlltoall => "trtllm_alltoall_perf.txt", - Self::ComputeScale => "computescale_perf.txt", - Self::ScaleMatrix => "scale_matrix_perf.txt", - Self::Mamba2 => "mamba2_perf.txt", - Self::Gdn => "gdn_perf.txt", - Self::MlaContextModule => "mla_context_module_perf.txt", - Self::MlaGenerationModule => "mla_generation_module_perf.txt", - Self::DsaContextModule => "dsa_context_module_perf.txt", - Self::DsaGenerationModule => "dsa_generation_module_perf.txt", - Self::MhcModule => "mhc_module_perf.txt", - Self::Dsv4CsaContextModule => "dsv4_csa_context_module_perf.txt", - Self::Dsv4HcaContextModule => "dsv4_hca_context_module_perf.txt", - Self::Dsv4CsaGenerationModule => "dsv4_csa_generation_module_perf.txt", - Self::Dsv4HcaGenerationModule => "dsv4_hca_generation_module_perf.txt", - Self::Dsv4PagedMqaLogitsModule => "dsv4_paged_mqa_logits_module_perf.txt", - Self::Dsv4HcaAttnModule => "dsv4_hca_attn_module_perf.txt", + Self::Gemm => "gemm_perf.parquet", + Self::Nccl => "nccl_perf.parquet", + Self::Oneccl => "oneccl_perf.parquet", + Self::GenerationAttention => "generation_attention_perf.parquet", + Self::ContextAttention => "context_attention_perf.parquet", + Self::EncoderAttention => "encoder_attention_perf.parquet", + Self::ContextMla => "context_mla_perf.parquet", + Self::GenerationMla => "generation_mla_perf.parquet", + Self::MlaBmm => "mla_bmm_perf.parquet", + Self::Moe => "moe_perf.parquet", + Self::CustomAllreduce => "custom_allreduce_perf.parquet", + Self::WideepContextMla => "wideep_context_mla_perf.parquet", + Self::WideepGenerationMla => "wideep_generation_mla_perf.parquet", + Self::WideepContextMoe => "wideep_context_moe_perf.parquet", + Self::WideepGenerationMoe => "wideep_generation_moe_perf.parquet", + Self::WideepDeepepNormal => "wideep_deepep_normal_perf.parquet", + Self::WideepDeepepLl => "wideep_deepep_ll_perf.parquet", + Self::WideepMoeCompute => "wideep_moe_perf.parquet", + Self::TrtllmAlltoall => "trtllm_alltoall_perf.parquet", + Self::ComputeScale => "computescale_perf.parquet", + Self::ScaleMatrix => "scale_matrix_perf.parquet", + Self::Mamba2 => "mamba2_perf.parquet", + Self::Gdn => "gdn_perf.parquet", + Self::MlaContextModule => "mla_context_module_perf.parquet", + Self::MlaGenerationModule => "mla_generation_module_perf.parquet", + Self::DsaContextModule => "dsa_context_module_perf.parquet", + Self::DsaGenerationModule => "dsa_generation_module_perf.parquet", + Self::MhcModule => "mhc_module_perf.parquet", + Self::Dsv4CsaContextModule => "dsv4_csa_context_module_perf.parquet", + Self::Dsv4HcaContextModule => "dsv4_hca_context_module_perf.parquet", + Self::Dsv4CsaGenerationModule => "dsv4_csa_generation_module_perf.parquet", + Self::Dsv4HcaGenerationModule => "dsv4_hca_generation_module_perf.parquet", + Self::Dsv4PagedMqaLogitsModule => "dsv4_paged_mqa_logits_module_perf.parquet", + Self::Dsv4HcaAttnModule => "dsv4_hca_attn_module_perf.parquet", + Self::Dsv4CsaAttnModule => "dsv4_csa_attn_module_perf.parquet", + Self::Dsv4CsaTopkCalib => "dsv4_csa_topk_calib_perf.parquet", + Self::Dsv4MegamoeModule => "dsv4_megamoe_module_perf.parquet", } } } @@ -413,17 +441,17 @@ mod tests { #[test] fn perf_data_filenames_match_python_enum() { // Sampled across categories from common.PerfDataFilename. - assert_eq!(PerfDataFilename::Gemm.as_str(), "gemm_perf.txt"); + assert_eq!(PerfDataFilename::Gemm.as_str(), "gemm_perf.parquet"); assert_eq!( PerfDataFilename::ContextAttention.as_str(), - "context_attention_perf.txt" + "context_attention_perf.parquet" ); - assert_eq!(PerfDataFilename::CustomAllreduce.as_str(), "custom_allreduce_perf.txt"); - assert_eq!(PerfDataFilename::WideepDeepepLl.as_str(), "wideep_deepep_ll_perf.txt"); - assert_eq!(PerfDataFilename::TrtllmAlltoall.as_str(), "trtllm_alltoall_perf.txt"); + assert_eq!(PerfDataFilename::CustomAllreduce.as_str(), "custom_allreduce_perf.parquet"); + assert_eq!(PerfDataFilename::WideepDeepepLl.as_str(), "wideep_deepep_ll_perf.parquet"); + assert_eq!(PerfDataFilename::TrtllmAlltoall.as_str(), "trtllm_alltoall_perf.parquet"); assert_eq!( PerfDataFilename::Dsv4HcaGenerationModule.as_str(), - "dsv4_hca_generation_module_perf.txt" + "dsv4_hca_generation_module_perf.parquet" ); } } diff --git a/rust/aiconfigurator-core/src/common/system_spec.rs b/rust/aiconfigurator-core/src/common/system_spec.rs index 16609c426..3566a6bb9 100644 --- a/rust/aiconfigurator-core/src/common/system_spec.rs +++ b/rust/aiconfigurator-core/src/common/system_spec.rs @@ -161,7 +161,9 @@ mod tests { let spec = SystemSpec::load(&systems_root().join("b200_sxm.yaml")) .expect("b200_sxm.yaml parse must succeed"); assert_eq!(spec.data_dir, PathBuf::from("data/b200_sxm")); - assert_eq!(spec.gpu.mem_bw, 8_000_000_000_000.0); + // 7.7 TB/s since PR #1246 (was 8.0; this pin went stale unnoticed + // because cargo test was not in the gating CI). + assert_eq!(spec.gpu.mem_bw, 7_700_000_000_000.0); assert_eq!(spec.gpu.mem_bw_empirical_scaling_factor, 0.8); assert_eq!(spec.gpu.bfloat16_tc_flops, Some(2_250_000_000_000_000.0)); assert_eq!(spec.gpu.sm_version, Some(100)); diff --git a/rust/aiconfigurator-core/src/engine/runtime.rs b/rust/aiconfigurator-core/src/engine/runtime.rs index a6d2a3b92..f0a91c297 100644 --- a/rust/aiconfigurator-core/src/engine/runtime.rs +++ b/rust/aiconfigurator-core/src/engine/runtime.rs @@ -22,22 +22,23 @@ use crate::common::error::AicError; use crate::engine::spec::EngineSpec; use crate::operators::Op; use crate::perf_database::PerfDatabase; -use crate::session::{get_mix_step_ops, run_context_ops, run_generation_ops_step}; +use crate::session::{get_mix_step_ops, run_context_ops, run_generation_ops_step, ContextOpFilter}; use crate::{validate_forward_pass_metrics, ForwardPassMetrics}; /// Per-call runtime inputs. Field-for-field mirror of the Python /// `sdk/config.RuntimeConfig`. /// -/// Only the fields the static composition reads are consumed (`batch_size`, -/// `beam_width`, `isl`, `osl`, `prefix`). The two imbalance-correction scales -/// are carried for wire parity but are not yet threaded into the op -/// queries (the live FPM path hard-codes them to 1.0; see -/// `session::run_context_ops`). +/// The imbalance-correction scales thread into the per-op queries exactly +/// where Python applies them (`base_backend.py:331,372`): context-attention +/// ops multiply by `seq_imbalance_correction_scale`, generation-attention ops +/// by `gen_seq_imbalance_correction_scale`. (The FPM telemetry path has no +/// scale concept and keeps 1.0.) #[derive(Clone, Copy, Debug, PartialEq)] pub struct RuntimeConfig { pub batch_size: u32, - /// Beam width. Defaults to 1; the engine-step path is not exercised for - /// beam > 1 (matches `session.rs`), so it does not scale the gen batch. + /// Beam width. The generation phase queries token-major ops at + /// `x = batch_size * beam_width` (Python `_run_generation_phase`); + /// attention ops key on the raw decode batch. pub beam_width: u32, pub isl: u32, pub osl: u32, @@ -235,6 +236,8 @@ impl Engine { runtime.batch_size, effective_isl, runtime.prefix, + runtime.seq_imbalance_correction_scale, + ContextOpFilter::All, ) } @@ -264,7 +267,15 @@ impl Engine { // Python `s = isl + i + 1`. NOTE the `+1` — distinct from the FPM // bridge's `context_length = isl + i` packing convention. let s = runtime.isl + i + 1; - let step = run_generation_ops_step(&self.generation_ops, &self.db, bs, s)?; + let step = crate::session::run_generation_ops_step_beamed( + &self.generation_ops, + &self.db, + bs, + runtime.beam_width, + s, + runtime.gen_seq_imbalance_correction_scale, + false, + )?; let repeat_count = stride.min(upper - i); total += step * repeat_count as f64; i += stride; @@ -305,33 +316,34 @@ impl Engine { .total_ms) } - /// One mixed (chunked-prefill + decode) step latency. Mirrors Python - /// `_get_mix_step_latency` (`base_backend.py:706`) via the same FPM - /// parameter packing the live ctypes bridge - /// (`rust_engine_step.estimate_mixed_step_latency_with_rust`) uses, so the - /// numbers match Python to within the parity tolerance. - /// - /// The packing reproduces the bridge's FPM build + the session unpack in - /// `session::rank_latency_ms`, then calls the shared - /// [`get_mix_step_ops`] composition over this engine's op lists: + /// One mixed (chunked-prefill + decode) step latency. LITERAL mirror of + /// Python `_get_mix_step_latency` (`base_backend.py:925-1050`), which + /// composes three `run_static` calls and filters the per-op breakdown by + /// name: /// /// ```text - /// // prefill chunk (FPM build) - /// n_prefill = max(ceil(ctx_tokens / isl), 1) - /// cached_total= prefix * n_prefill - /// new_prefill = ctx_tokens - cached_total (>=1) if cached_total else ctx_tokens - /// // decode (FPM build, with the (nextn+1) MTP multiplier) - /// eff_gen = gen_tokens * (nextn + 1) - /// kv_total = eff_gen * (isl + osl/2) - /// // session unpack -> get_mix_step_ops args - /// new_tokens_per_req = new_prefill / n_prefill - /// prefix_per_req = cached_total / n_prefill - /// combined_prefix = cached_total - /// kv_per_decode = kv_total / eff_gen - /// decode_batch = eff_gen - /// ctx_tokens(arg) = new_prefill // == sum_prefill_tokens - /// gen_tokens(arg) = eff_gen // == num_decode_requests + /// // 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 the RAW `ctx + gen` + /// token count (no `(nextn+1)` inflation — the MTP multiplier applies only + /// to the pass-3 decode batch via `_run_generation_phase`), 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`). pub fn mixed_step_latency( &self, ctx_tokens: u32, @@ -339,69 +351,107 @@ impl Engine { isl: u32, osl: u32, prefix: u32, + seq_imbalance_correction_scale: f64, + gen_seq_imbalance_correction_scale: f64, ) -> Result { + if ctx_tokens == 0 && gen_tokens == 0 { + return Ok(0.0); + } + // Python divides by `isl` (`floor(ctx/isl)`, `ceil(ctx/isl)`) without + // a guard — callers always pass isl >= 1. Clamp to avoid a Rust + // div-by-zero panic on degenerate input Python would crash on. let isl = isl.max(1); - let osl = osl.max(1); - - // ---- Prefill chunk FPM build (mirrors the bridge) ---- - let (sum_prefill_tokens, sum_prefill_kv_tokens, n_prefill) = if ctx_tokens > 0 { - let n_prefill = ctx_tokens.div_ceil(isl).max(1); - let cached_total = prefix * n_prefill; - let new_prefill = if cached_total > 0 { - ctx_tokens.saturating_sub(cached_total).max(1) - } else { - ctx_tokens - }; - (new_prefill, cached_total, n_prefill) - } else { - (0, 0, 0) - }; + let mut total = 0.0_f64; - // ---- Decode FPM build with the (nextn + 1) MTP multiplier ---- - let (num_decode_requests, sum_decode_kv_tokens) = if gen_tokens > 0 { - let eff_gen = gen_tokens.saturating_mul(self.nextn.saturating_add(1)); - let kv_per_req = isl.saturating_add(osl / 2); - (eff_gen, eff_gen.saturating_mul(kv_per_req)) - } else { - (0, 0) - }; + // ---- Pass 1: combined non-attention work ---- + let combined = ctx_tokens + gen_tokens; + let prefix1 = prefix * (ctx_tokens / isl); // prefix * floor(ctx/isl) + if prefix1 >= combined { + return Err(AicError::InvalidEngineConfig(format!( + "isl must be greater than 0 after removing prefix, but got {}", + combined as i64 - prefix1 as i64 + ))); + } + total += run_context_ops( + &self.context_ops, + &self.db, + 1, + combined - prefix1, + prefix1, + seq_imbalance_correction_scale, + ContextOpFilter::SkipContextAttention, + )?; - if sum_prefill_tokens == 0 && num_decode_requests == 0 { - return Ok(0.0); + // ---- Pass 2: context attention at the prefill shape ---- + // Python: batch = ceil(ctx/isl), effective_isl = isl - prefix, then + // latency["context_attention"] / ceil(isl/ctx). With ctx_tokens == 0 + // Python's `np.ceil(isl/0)` is +inf and the division yields 0 — skip. + if ctx_tokens > 0 { + if prefix >= isl { + return Err(AicError::InvalidEngineConfig(format!( + "isl must be greater than 0 after removing prefix, but got {}", + isl as i64 - prefix as i64 + ))); + } + let batch2 = ctx_tokens.div_ceil(isl); + let scale2 = isl.div_ceil(ctx_tokens) as f64; + let attn = run_context_ops( + &self.context_ops, + &self.db, + batch2, + isl - prefix, + prefix, + seq_imbalance_correction_scale, + ContextOpFilter::OnlyContextAttention, + )?; + total += attn / scale2; } - // ---- Session unpack (mirrors session::rank_latency_ms) ---- - let n_prefill_safe = n_prefill.max(1); - let new_tokens_per_req = (sum_prefill_tokens / n_prefill_safe).max(1); - let prefix_per_req = sum_prefill_kv_tokens / n_prefill_safe; - let n_decode_safe = num_decode_requests.max(1); - let kv_per_decode = sum_decode_kv_tokens / n_decode_safe; + // ---- Pass 3: decode attention ---- + if gen_tokens > 0 { + let bs = gen_tokens.saturating_mul(self.nextn.saturating_add(1)); + // `_run_generation_phase` queries at s = isl_pass3 + i + 1 with + // isl_pass3 = isl + osl//2 and a single step (osl=2, i=0). + let s = isl + osl / 2 + 1; + total += run_generation_ops_step( + &self.generation_ops, + &self.db, + bs, + s, + gen_seq_imbalance_correction_scale, + true, + )?; + } - crate::session::get_mix_step_ops( - &self.context_ops, - &self.generation_ops, - &self.db, - sum_prefill_tokens, - num_decode_requests, - new_tokens_per_req, - prefix_per_req, - sum_prefill_kv_tokens, - kv_per_decode, - num_decode_requests, - ) + Ok(total) } - /// One generation-only step latency. Mirrors Python - /// `_get_genonly_step_latency` (`base_backend.py:834`) / - /// `rust_engine_step.estimate_decode_step_latency_with_rust`: one decode - /// step at `s = isl + osl/2` with the decode batch scaled by `(nextn + 1)`. - pub fn decode_step_latency(&self, gen_tokens: u32, isl: u32, osl: u32) -> Result { + /// 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)`. + pub fn decode_step_latency( + &self, + gen_tokens: u32, + isl: u32, + osl: u32, + gen_seq_imbalance_correction_scale: f64, + ) -> Result { if gen_tokens == 0 { return Ok(0.0); } let effective_batch = gen_tokens.saturating_mul(self.nextn.saturating_add(1)); - let context_length = isl.max(1).saturating_add(osl.max(1) / 2); - run_generation_ops_step(&self.generation_ops, &self.db, effective_batch, context_length) + let s = isl.max(1).saturating_add(osl.max(1) / 2).saturating_add(1); + run_generation_ops_step( + &self.generation_ops, + &self.db, + effective_batch, + s, + gen_seq_imbalance_correction_scale, + false, + ) } /// Compute one forward-pass latency from a list of per-rank FPM entries. @@ -490,13 +540,22 @@ impl Engine { n_prefill, new_tokens_per_req, prefix_per_req, + 1.0, + ContextOpFilter::All, )?; } if has_decode { let n_decode = sched.num_decode_requests.max(1); let kv_per_req = sched.sum_decode_kv_tokens / n_decode; - total += run_generation_ops_step(&self.generation_ops, &self.db, n_decode, kv_per_req)?; + total += run_generation_ops_step( + &self.generation_ops, + &self.db, + n_decode, + kv_per_req, + 1.0, + false, + )?; } Ok(total) @@ -531,6 +590,7 @@ mod tests { name: "rmsnorm".into(), scale_factor: 1.0, bytes_per_token: 8192.0, + scale_num_tokens: 1, seq_split: 1, }), Op::Gemm(GemmOp { @@ -564,6 +624,7 @@ mod tests { name: "rmsnorm".into(), scale_factor: 1.0, bytes_per_token: 8192.0, + scale_num_tokens: 1, seq_split: 1, }), Op::GenerationAttention(GenerationAttentionOp { @@ -686,6 +747,8 @@ mod tests { engine.database(), 1, // batch_size * (nextn+1), nextn=0 1024 + 0 + 1, + 1.0, + false, ) .unwrap(); assert!((coarse.generation_ms - one_step * 8.0).abs() < 1e-6); @@ -715,7 +778,7 @@ mod tests { #[test] fn mixed_step_empty_is_zero() { let engine = build_engine(None); - assert_eq!(engine.mixed_step_latency(0, 0, 1024, 8, 0).unwrap(), 0.0); + assert_eq!(engine.mixed_step_latency(0, 0, 1024, 8, 0, 1.0, 1.0).unwrap(), 0.0); } #[test] @@ -725,7 +788,7 @@ mod tests { // End-to-end parity is covered by the mixed-step parity cases; this is // the fast pure-Rust smoke that the composition actually computes. let engine = build_engine(None); - let ms = engine.mixed_step_latency(1024, 2, 1024, 8, 0).unwrap(); + let ms = engine.mixed_step_latency(1024, 2, 1024, 8, 0, 1.0, 1.0).unwrap(); assert!(ms > 0.0 && ms.is_finite(), "mixed-step latency must be > 0, got {ms}"); } @@ -752,6 +815,8 @@ mod tests { engine_nextn1.database(), 2, 1024 + 1, + 1.0, + false, ) .unwrap(); assert!( diff --git a/rust/aiconfigurator-core/src/engine/spec.rs b/rust/aiconfigurator-core/src/engine/spec.rs index 8a950e91e..de9ed77bf 100644 --- a/rust/aiconfigurator-core/src/engine/spec.rs +++ b/rust/aiconfigurator-core/src/engine/spec.rs @@ -161,7 +161,8 @@ mod tests { ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4ModuleOp, ElementwiseOp, EmbeddingOp, EncoderAttentionOp, GdnOp, GemmOp, GenerationAttentionOp, GenerationMlaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, MoEDispatchOp, MoeOp, - NcclOp, P2POp, VisionEncoderOp, WideEpContextMlaOp, WideEpGenerationMlaOp, WideEpMoeOp, + NcclOp, P2POp, TrtllmWideEpMoEDispatchOp, VisionEncoderOp, WideEpContextMlaOp, + WideEpGenerationMlaOp, WideEpMoeOp, }; use crate::operators::moe_dispatch::DispatchFlavor; use crate::perf_database::dsv4::AttnKind; @@ -201,6 +202,7 @@ mod tests { name: "rmsnorm".into(), scale_factor: 1.5, bytes_per_token: 8192.0, + scale_num_tokens: 1, seq_split: 1, } } @@ -239,6 +241,7 @@ mod tests { n: 16, head_size: 80, fmha_quant_mode: FmhaQuantMode::Fp8, + partial_rotary_factor: 0.0, } } @@ -297,6 +300,9 @@ mod tests { quant_mode: MoeQuantMode::Fp8Block, workload_distribution: "power_law_1.2".into(), is_gated: true, + moe_backend: None, + enable_eplb: false, + is_context: false, } } @@ -318,6 +324,7 @@ mod tests { attn_cp_size: 1, is_context: false, sms: 12, + scale_num_tokens: 1, } } @@ -379,6 +386,7 @@ mod tests { architecture: "DeepseekV32ForCausalLM".into(), index_topk: 2048, cp_size: 1, + full_frac: 1.0, } } @@ -493,6 +501,22 @@ mod tests { } } + fn wideep_moe_dispatch() -> TrtllmWideEpMoEDispatchOp { + TrtllmWideEpMoEDispatchOp { + name: "wideep_moe_dispatch".into(), + scale_factor: 61.0, + hidden_size: 7168, + topk: 8, + num_experts: 256, + moe_tp_size: 1, + moe_ep_size: 8, + attention_dp_size: 8, + pre_dispatch: true, + quant_mode: MoeQuantMode::Nvfp4, + use_low_precision_combine: false, + } + } + fn overlap() -> OverlapOp { // Recursive: nested children on both groups. OverlapOp { @@ -548,6 +572,7 @@ mod tests { OpSpec::WideEpContextMla(wideep_context_mla()), OpSpec::WideEpGenerationMla(wideep_generation_mla()), OpSpec::WideEpMoe(wideep_moe()), + OpSpec::WideEpMoeDispatch(wideep_moe_dispatch()), OpSpec::Overlap(overlap()), OpSpec::Fallback(fallback()), ]; @@ -583,6 +608,7 @@ mod tests { | OpSpec::WideEpContextMla(_) | OpSpec::WideEpGenerationMla(_) | OpSpec::WideEpMoe(_) + | OpSpec::WideEpMoeDispatch(_) | OpSpec::Overlap(_) | OpSpec::Fallback(_) => {} } diff --git a/rust/aiconfigurator-core/src/fpm/tests.rs b/rust/aiconfigurator-core/src/fpm/tests.rs index ee3ac6a0f..7ce75941e 100644 --- a/rust/aiconfigurator-core/src/fpm/tests.rs +++ b/rust/aiconfigurator-core/src/fpm/tests.rs @@ -34,6 +34,7 @@ fn context_ops() -> Vec { name: "rmsnorm".into(), scale_factor: 1.0, bytes_per_token: 8192.0, + scale_num_tokens: 1, seq_split: 1, }), Op::Gemm(GemmOp { @@ -67,6 +68,7 @@ fn generation_ops() -> Vec { name: "rmsnorm".into(), scale_factor: 1.0, bytes_per_token: 8192.0, + scale_num_tokens: 1, seq_split: 1, }), Op::GenerationAttention(GenerationAttentionOp { @@ -203,6 +205,8 @@ fn forward_pass_prefill_matches_run_context_ops() { 4, 1024, 0, + 1.0, + crate::session::ContextOpFilter::All, ) .unwrap(); assert_close(via_fpm, direct); @@ -228,6 +232,8 @@ fn forward_pass_decode_matches_run_generation_ops_step() { engine.database(), 8, 2048, + 1.0, + false, ) .unwrap(); assert_close(via_fpm, direct); diff --git a/rust/aiconfigurator-core/src/operators/attention.rs b/rust/aiconfigurator-core/src/operators/attention.rs index 6840c0d76..9dae15b79 100644 --- a/rust/aiconfigurator-core/src/operators/attention.rs +++ b/rust/aiconfigurator-core/src/operators/attention.rs @@ -226,6 +226,12 @@ pub struct EncoderAttentionOp { pub n: u32, pub head_size: u32, pub fmha_quant_mode: FmhaQuantMode, + /// Partial-RoPE fraction (Python `_partial_rotary_factor`): 1.0 = full + /// rotation, 0.5 = half head_dim rotated (Qwen3-VL), 0.0 = no RoPE. + /// Adds `factor * 2 * mem_op(Q+K bytes) * 1.1` on top of the table + /// latency. Defaults to 0.0 for pre-field opspecs. + #[serde(default)] + pub partial_rotary_factor: f64, } impl EncoderAttentionOp { @@ -241,6 +247,7 @@ impl EncoderAttentionOp { n, head_size, fmha_quant_mode, + partial_rotary_factor: 0.0, } } @@ -250,13 +257,23 @@ impl EncoderAttentionOp { batch_size: u32, s: u32, ) -> Result { - let latency = db.attention.query_encoder( + let mut latency = db.attention.query_encoder( batch_size, s, self.n, self.head_size, self.fmha_quant_mode, )?; + // Partial RoPE extra (Python `EncoderAttention.query`, + // operations/attention.py): Q + K bytes (bf16) over all tokens, + // rotated fractionally, with the 1.1 correction factor. + if self.partial_rotary_factor > 0.0 { + let qk_num = (self.n as u64) * (self.head_size as u64); // MHA: q == k + let qk_bytes = 2 * (qk_num * 2) * (batch_size as u64) * (s as u64); + let apply_rope = + self.partial_rotary_factor * 2.0 * mem_op_latency_ms(&db.system_spec, qk_bytes as f64); + latency += apply_rope * 1.1; + } Ok(PerformanceResult::new(latency, Source::Silicon) .clamp_non_negative() .scaled(self.scale_factor)) @@ -343,10 +360,15 @@ mod tests { #[test] fn mem_op_latency_uses_empirical_formula() { let db = b200_vllm_db(); - // b200_sxm.yaml: mem_bw=8e12, scaling=0.8, constant=3e-6. - // For 1MB: latency = (1e6 / (8e12 * 0.8) + 3e-6) * 1000 = 0.156... + 0.003 = 0.159ms. - let latency = mem_op_latency_ms(&db.system_spec, 1_000_000.0); - let expected = (1_000_000.0_f64 / (8e12 * 0.8) + 3e-6) * 1000.0; + // Formula check against the LIVE spec values (hardcoding mem_bw went + // stale when PR #1246 corrected b200 to 7.7 TB/s): + // latency = (bytes / (mem_bw * scaling) + constant) * 1000. + let spec = &db.system_spec; + let latency = mem_op_latency_ms(spec, 1_000_000.0); + let expected = (1_000_000.0_f64 + / (spec.gpu.mem_bw * spec.gpu.mem_bw_empirical_scaling_factor) + + spec.gpu.mem_empirical_constant_latency) + * 1000.0; assert!((latency - expected).abs() < 1e-12); } } diff --git a/rust/aiconfigurator-core/src/operators/communication.rs b/rust/aiconfigurator-core/src/operators/communication.rs index 14032415a..04dafb831 100644 --- a/rust/aiconfigurator-core/src/operators/communication.rs +++ b/rust/aiconfigurator-core/src/operators/communication.rs @@ -57,6 +57,9 @@ impl CustomAllReduceOp { /// count, not bytes) and passes it directly to /// `query_custom_allreduce`. Mirror that here — the underlying table /// is keyed by element count, with dtype implicit in the quant mode. + /// Node-fan-out capping, beyond-node bandwidth scaling and the GB200 + /// NVL72 -> NCCL reroute all live in the DB-level `_scaled` query + /// (mirroring Python's `_query_custom_allreduce_table` funnel). pub fn query( &self, db: &PerfDatabase, @@ -66,21 +69,13 @@ impl CustomAllReduceOp { return Ok(PerformanceResult::zero()); } let per_rank_tokens = num_tokens.div_ceil(self.seq_split.max(1)); // CP: busiest rank - let message_size = (per_rank_tokens as u64) * (self.hidden_size as u64); - let spec = &db.system_spec; - let per_node = spec.node.num_gpus_per_node; - let effective_tp = self.tp_size.min(per_node); - let mut latency = - db.communication - .query_custom_allreduce(self.quant, effective_tp, message_size)?; - if self.tp_size > per_node { - let base_bw = p2p_bandwidth(spec, per_node); - let target_bw = p2p_bandwidth(spec, self.tp_size); - let f_tp = self.tp_size as f64; - let f_pn = per_node as f64; - let scale = (f_tp - 1.0) / f_tp * f_pn / (f_pn - 1.0).max(1.0) * base_bw / target_bw; - latency *= scale; - } + let message_size = (per_rank_tokens as f64) * (self.hidden_size as f64); + let latency = db.communication.query_custom_allreduce_scaled( + &db.system_spec, + self.quant, + self.tp_size, + message_size, + )?; Ok(PerformanceResult::new(latency, Source::Silicon) .clamp_non_negative() .scaled(self.scale_factor)) @@ -134,26 +129,19 @@ impl NcclOp { return Ok(PerformanceResult::zero()); } let per_rank_tokens = num_tokens.div_ceil(self.seq_split.max(1)); // CP: busiest rank - // Python: message_size = ceil(x/seq_split) * num_elements_per_token (float). - let message_size = ((per_rank_tokens as f64) * self.hidden_size) as u64; - let max_recorded = db - .communication - .nccl_max_num_gpus(self.dtype, &self.operation)? - .unwrap_or(self.num_gpus); - let effective = self.num_gpus.min(max_recorded); - let mut latency = db - .communication - .query_nccl(self.dtype, &self.operation, effective, message_size)?; - if self.num_gpus > max_recorded { - let spec = &db.system_spec; - let max_bw = p2p_bandwidth(spec, max_recorded); - let req_bw = p2p_bandwidth(spec, self.num_gpus); - let f_n = self.num_gpus as f64; - let f_m = max_recorded as f64; - let scale = - (f_n - 1.0) / f_n * f_m / (f_m - 1.0).max(1.0) * max_bw / req_bw; - latency *= scale; - } + // Python: message_size = ceil(x/seq_split) * num_elements_per_token — + // kept as a FLOAT (fractional element counts are real: the gemma4 CP + // KV all-gather sizes per-token elements as kv_bytes / comm_bytes). + let message_size = (per_rank_tokens as f64) * self.hidden_size; + // Fan-out capping + beyond-range bandwidth correction live in the + // DB-level `_scaled` query (mirroring Python's `_query_nccl_table`). + let latency = db.communication.query_nccl_scaled( + &db.system_spec, + self.dtype, + &self.operation, + self.num_gpus, + message_size, + )?; Ok(PerformanceResult::new(latency, Source::Silicon) .clamp_non_negative() .scaled(self.scale_factor)) diff --git a/rust/aiconfigurator-core/src/operators/dsa.rs b/rust/aiconfigurator-core/src/operators/dsa.rs index a50f47794..29bef0423 100644 --- a/rust/aiconfigurator-core/src/operators/dsa.rs +++ b/rust/aiconfigurator-core/src/operators/dsa.rs @@ -41,12 +41,23 @@ pub struct DsaModuleOp { /// once BOTH the dsa and dsv4 Rust CP paths land. #[serde(default = "default_cp_size")] pub cp_size: u32, + /// GLM-5.2 shared-index amortization weight (Python `_full_frac`): the + /// exact fraction of indexer-computing layers. Per-layer cost is + /// `full_frac*full + (1-full_frac)*skip` using the directly-collected + /// skip-indexer table. 1.0 (DeepSeek-V3.2 / GLM-5, and pre-field opspecs) + /// keeps the pure-full path — the skip table is never touched. + #[serde(default = "default_full_frac")] + pub full_frac: f64, } fn default_cp_size() -> u32 { 1 } +fn default_full_frac() -> f64 { + 1.0 +} + impl DsaModuleOp { #[allow(clippy::too_many_arguments)] pub fn new( @@ -68,6 +79,7 @@ impl DsaModuleOp { architecture: architecture.into(), index_topk, cp_size: 1, + full_frac: 1.0, } } @@ -78,11 +90,22 @@ impl DsaModuleOp { isl: u32, prefix: u32, ) -> Result { + let w = self.full_frac; // CP (round-robin sequence split) prefill takes the sparse-delta // composition path (Python `ContextDSAModule.query` -> `_query_cp` - // when `_cp_size > 1`). + // when `_cp_size > 1`). GLM-5.2 amortizes full/skip on the CP path + // too (both carry the same scale_factor, so the weighted sum of the + // already-scaled results is exact — Python `_amortize`). if self.cp_size > 1 { - return self.query_context_cp(db, batch_size, isl, prefix); + let full = self.query_context_cp(db, batch_size, isl, prefix, false)?; + if w >= 1.0 { + return Ok(full); + } + let skip = self.query_context_cp(db, batch_size, isl, prefix, true)?; + return Ok(PerformanceResult::new( + w * full.latency_ms + (1.0 - w) * skip.latency_ms, + full.source, + )); } // Query at `isl` (new-token count) for the exact `prefix` slice — NOT // `isl + prefix`. The perf-DB layer resolves one 4-axis RAW grid via @@ -90,19 +113,29 @@ impl DsaModuleOp { // correction (it had no Python counterpart and under-counted context // latency ~75%). `dsa_backend="trtllm"` mirrors Python's non-CP // default (`_query_context_dsa_module_table(dsa_backend="trtllm")`). - let latency = db.dsa.query_context( - &db.system_spec, - batch_size, - isl, - self.num_heads, - self.kv_cache_dtype, - self.fmha_quant_mode, - self.gemm_quant_mode, - &self.architecture, - prefix, - self.index_topk, - "trtllm", - )?; + let q = |skip_indexer: bool| { + db.dsa.query_context( + &db.system_spec, + batch_size, + isl, + self.num_heads, + self.kv_cache_dtype, + self.fmha_quant_mode, + self.gemm_quant_mode, + &self.architecture, + prefix, + self.index_topk, + "trtllm", + skip_indexer, + ) + }; + let full = q(false)?; + let latency = if w >= 1.0 { + full + } else { + // GLM-5.2 shared-index amortization (Python ContextDSAModule.query). + w * full + (1.0 - w) * q(true)? + }; Ok(PerformanceResult::new(latency, Source::Silicon) .clamp_non_negative() .scaled(self.scale_factor)) @@ -125,6 +158,7 @@ impl DsaModuleOp { batch_size: u32, isl: u32, prefix: u32, + skip_indexer: bool, ) -> Result { let sparse = db.dsa.load_cp_sparse(&self.architecture, self.num_heads)?; let mut base = |per_card: u32| { @@ -142,6 +176,7 @@ impl DsaModuleOp { // Python `_query_cp` queries the CP base on the flashmla_kv // slice (the kernel used under CP). "flashmla_kv", + skip_indexer, ) }; let mut ag = |elems: u64| { @@ -155,7 +190,7 @@ impl DsaModuleOp { .query(db, 1) .map(|r| r.latency_ms) }; - self.query_cp_with(&sparse, batch_size, isl, prefix, &mut base, &mut ag) + self.query_cp_with(&sparse, batch_size, isl, prefix, skip_indexer, &mut base, &mut ag) } /// CP (round-robin split) per-layer DSA composition. Verbatim mirror of @@ -179,6 +214,7 @@ impl DsaModuleOp { b: u32, isl: u32, prefix: u32, + skip_indexer: bool, base: &mut dyn FnMut(u32) -> Result, ag: &mut dyn FnMut(u64) -> Result, ) -> Result { @@ -190,15 +226,21 @@ impl DsaModuleOp { // step, so an empty grid below means the table is absent entirely // (parquet not collected) — degrading silently to dsa_base would hide // that. Message shape mirrors Python's fail-loud contract. - let missing: Vec<&str> = [ - ("mqa", &sparse.mqa), - ("topk_last", &sparse.topk_last), - ("topk_flat", &sparse.topk_flat), - ] - .into_iter() - .filter(|(_, grid)| grid.is_empty()) - .map(|(name, _)| name) - .collect(); + // skip_indexer layers carry NO indexer -> no mqa/topk deltas needed, + // so don't require the sparse tables for them (Python dsa.py:835-837). + let missing: Vec<&str> = if skip_indexer { + Vec::new() + } else { + [ + ("mqa", &sparse.mqa), + ("topk_last", &sparse.topk_last), + ("topk_flat", &sparse.topk_flat), + ] + .into_iter() + .filter(|(_, grid)| grid.is_empty()) + .map(|(name, _)| name) + .collect() + }; if !missing.is_empty() { return Err(AicError::PerfDatabase(format!( "DSA CP modeling needs sparse tables ['{}'] for {} (num_heads={}); \ @@ -225,12 +267,17 @@ impl DsaModuleOp { let tl_full = lookup_2d(tl_tab, isl, prefix)?; let tf_perc = lookup_2d(tf_tab, per_card, prefix)?; let mut latency = dsa_base; - if let (Some(mqa_full), Some(mqa_perc), Some(tl_full), Some(tf_perc)) = - (mqa_full, mqa_perc, tl_full, tf_perc) - { - let delta_mqa = mqa_full / f64::from(cp) - mqa_perc; - let delta_topk = tl_full / f64::from(cp) - tf_perc; - latency += delta_mqa + delta_topk; + // skip layers reuse a sibling's topk index: no per-layer mqa/topk, so + // no full/cp deltas — just the per-card skip base + the attention + // all-gathers (Python dsa.py:871-876). + if !skip_indexer { + if let (Some(mqa_full), Some(mqa_perc), Some(tl_full), Some(tf_perc)) = + (mqa_full, mqa_perc, tl_full, tf_perc) + { + let delta_mqa = mqa_full / f64::from(cp) - mqa_perc; + let delta_topk = tl_full / f64::from(cp) - tf_perc; + latency += delta_mqa + delta_topk; + } } // CP attention all-gathers, per current-chunk tokens (isl, not // isl+prefix; prefix KV is already replicated), bf16 (see the Python @@ -241,7 +288,14 @@ impl DsaModuleOp { // MoE dispatch ops, not here.) let dims = dsa_dims(&self.architecture); let tokens = u64::from(b) * u64::from(isl); - let ag_kv = ag(tokens * dims.index_head_dim as u64)?; + // A skip-indexer (reuse) layer never runs the per-layer indexer, so it + // does not all-gather the DSA indexer key — only the MLA + // compressed-KV/LSE gather remains (Python dsa.py:895-902). + let ag_kv = if skip_indexer { + 0.0 + } else { + ag(tokens * dims.index_head_dim as u64)? + }; let ag_lse = ag(tokens * (dims.kv_lora_rank + dims.qk_rope_head_dim) as u64)?; latency += ag_kv + ag_lse; Ok(PerformanceResult::new(latency, Source::Estimated).scaled(self.scale_factor)) @@ -255,17 +309,28 @@ impl DsaModuleOp { ) -> Result { // `dsa_backend="trtllm"` mirrors Python's generation default // (`_query_generation_dsa_module_table(dsa_backend="trtllm")`). - let latency = db.dsa.query_generation( - &db.system_spec, - batch_size, - s, - self.num_heads, - self.kv_cache_dtype, - self.fmha_quant_mode, - self.gemm_quant_mode, - &self.architecture, - "trtllm", - )?; + let q = |skip_indexer: bool| { + db.dsa.query_generation( + &db.system_spec, + batch_size, + s, + self.num_heads, + self.kv_cache_dtype, + self.fmha_quant_mode, + self.gemm_quant_mode, + &self.architecture, + "trtllm", + skip_indexer, + ) + }; + let w = self.full_frac; + let full = q(false)?; + let latency = if w >= 1.0 { + full + } else { + // GLM-5.2 shared-index amortization (Python GenerationDSAModule.query). + w * full + (1.0 - w) * q(true)? + }; Ok(PerformanceResult::new(latency, Source::Silicon) .clamp_non_negative() .scaled(self.scale_factor)) @@ -289,6 +354,7 @@ mod tests { architecture: "GlmMoeDsaForCausalLM".into(), index_topk: 2048, cp_size, + full_frac: 1.0, } } @@ -329,6 +395,7 @@ mod tests { 1, isl, prefix, + false, &mut |per_card| { base_calls.push(per_card); Ok(4300.0) // per-card monolithic base @@ -366,6 +433,7 @@ mod tests { 1, 32768, // > grid max 16384 0, + false, &mut |_| Ok(4300.0), &mut |_| Ok(50.0), ) diff --git a/rust/aiconfigurator-core/src/operators/dsv4.rs b/rust/aiconfigurator-core/src/operators/dsv4.rs index 1ea690b81..e5fc4b8d8 100644 --- a/rust/aiconfigurator-core/src/operators/dsv4.rs +++ b/rust/aiconfigurator-core/src/operators/dsv4.rs @@ -411,6 +411,9 @@ impl Dsv4ModuleOp { batch_size: u32, s: u32, ) -> Result { + // No fmha argument: the generation table keys on kv dtype only and the + // SOL dtype is derived from kv inside `query_generation` (PR #1337). + // `self.fmha_quant_mode` stays on the op for the context path. let latency = db.dsv4.query_generation( &db.system_spec, self.attn_kind, @@ -418,7 +421,6 @@ impl Dsv4ModuleOp { s, self.num_heads, self.kv_cache_dtype, - self.fmha_quant_mode, self.gemm_quant_mode, &self.architecture, Some(self.sol_dims()), diff --git a/rust/aiconfigurator-core/src/operators/elementwise.rs b/rust/aiconfigurator-core/src/operators/elementwise.rs index 9c247152a..143e88994 100644 --- a/rust/aiconfigurator-core/src/operators/elementwise.rs +++ b/rust/aiconfigurator-core/src/operators/elementwise.rs @@ -22,13 +22,19 @@ use crate::perf_database::PerfDatabase; pub struct ElementwiseOp { pub name: String, pub scale_factor: f64, - /// Bytes touched per input token. The query multiplies this by the - /// runtime token count. + /// Bytes touched per input token (`dim_in * 2 + dim_out * 2`, bf16 + /// activations). The query multiplies this by the runtime token count. pub bytes_per_token: f64, + /// Token-count divisor (Python `_scale_num_tokens`), applied as an + /// INTEGER floor-division BEFORE the CP ceil-split — matching Python's + /// `x //= scale; x = ceil(x / seq_split)` order exactly. (Old specs + /// folded the divisor into `bytes_per_token` instead, which is only + /// exact when the divisor divides the token count; they deserialize + /// with the default 1 and keep their folded semantics.) + #[serde(default = "crate::operators::gemm::default_seq_split")] + pub scale_num_tokens: u32, /// CP sequence-shard factor (Python's `_seq_split`, = `cp_size`): the /// per-rank token count is `ceil(num_tokens / seq_split)`. Defaults to 1. - /// (Python's `scale_num_tokens` is folded into `bytes_per_token` by the - /// serializer; `seq_split` is applied here to the token count.) #[serde(default = "crate::operators::gemm::default_seq_split")] pub seq_split: u32, } @@ -39,11 +45,16 @@ impl ElementwiseOp { name: name.into(), scale_factor: 1.0, bytes_per_token, + scale_num_tokens: 1, seq_split: 1, } } pub fn query(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + // Python order: `x //= scale_num_tokens` (integer floor) THEN the CP + // ceil-split. Folding the divisor into bytes diverges whenever it + // does not divide the token count (live on DSv3.2 CP add_norm ops). + let num_tokens = num_tokens / self.scale_num_tokens.max(1); let num_tokens = num_tokens.div_ceil(self.seq_split.max(1)); // CP: busiest rank let bytes = self.bytes_per_token * (num_tokens as f64); let latency = mem_op_latency_ms(&db.system_spec, bytes); diff --git a/rust/aiconfigurator-core/src/operators/gemm.rs b/rust/aiconfigurator-core/src/operators/gemm.rs index 0e0811527..34a23c04a 100644 --- a/rust/aiconfigurator-core/src/operators/gemm.rs +++ b/rust/aiconfigurator-core/src/operators/gemm.rs @@ -10,8 +10,9 @@ //! //! For `fp8_static` quant mode, subtracts `compute_scale` overhead from the //! GEMM latency (and additionally subtracts `scale_matrix` when the input -//! is also low-precision). Latency is clamped to `>= 0` post-subtraction -//! to mirror Python's `max(0.0, latency)` behavior. +//! is also low-precision). Post-subtraction latency is floored at the GEMM's +//! own SOL roofline — mirroring Python's `max(latency_floor, latency)` where +//! `latency_floor = query_gemm(..., DatabaseMode.SOL)` — not at 0. use serde::{Deserialize, Serialize}; use crate::common::enums::GemmQuantMode; @@ -77,19 +78,33 @@ impl GemmOp { let mut latency = db.gemm.query(quant, m, self.n, self.k)?; let mut source = Source::Silicon; + let mut latency_floor = 0.0_f64; if quant == GemmQuantMode::Fp8Static { let cs_latency = db.gemm.query_compute_scale(quant, m, self.k)?; latency -= cs_latency; - source = source.combine(Source::Silicon); // both silicon -> still silicon if self.low_precision_input { let sm_latency = db.gemm.query_scale_matrix(quant, m, self.k)?; latency -= sm_latency; } + // Python (`operations/gemm.py`): the subtraction leaves a path + // that still contains the GEMM; independently interpolated + // component tables can cross, but that path cannot be faster than + // the GEMM's own roofline. Floor at the SOL (NOT at 0), and tag + // the result "estimated" (fp8_static is modeled from dynamic FP8 + // plus overhead tables, not measured directly). + latency_floor = crate::perf_database::gemm::gemm_sol_latency_ms( + &db.system_spec, + quant, + m as f64, + self.n as f64, + self.k as f64, + ); + source = Source::Estimated; } - Ok(PerformanceResult::new(latency, source) + Ok(PerformanceResult::new(latency.max(latency_floor), source) .clamp_non_negative() .scaled(self.scale_factor)) } diff --git a/rust/aiconfigurator-core/src/operators/mod.rs b/rust/aiconfigurator-core/src/operators/mod.rs index d7e3ad7a4..8a44b5e8a 100644 --- a/rust/aiconfigurator-core/src/operators/mod.rs +++ b/rust/aiconfigurator-core/src/operators/mod.rs @@ -44,9 +44,8 @@ pub use mamba::{GdnOp, Mamba2Op}; pub use mhc::MhcModuleOp; pub use mla::{ContextMlaOp, GenerationMlaOp, MlaBmmOp, MlaModuleOp}; pub use moe::MoeOp; -pub use moe_dispatch::{DispatchFlavor, MoEDispatchOp}; +pub use moe_dispatch::{DispatchFlavor, MoEDispatchOp, TrtllmWideEpMoEDispatchOp}; pub use op::{FallbackOp, Op, OverlapOp, RuntimeContext}; -pub use overlap::overlap_composition; pub use vision::VisionEncoderOp; pub use wideep_mla::{WideEpContextMlaOp, WideEpGenerationMlaOp}; pub use wideep_moe::WideEpMoeOp; diff --git a/rust/aiconfigurator-core/src/operators/moe.rs b/rust/aiconfigurator-core/src/operators/moe.rs index b1c05674f..dfc7cec7c 100644 --- a/rust/aiconfigurator-core/src/operators/moe.rs +++ b/rust/aiconfigurator-core/src/operators/moe.rs @@ -45,6 +45,22 @@ pub struct MoeOp { /// `moe_torch_flow_min_latency` kernel is only valid for gated nvfp4 /// MoE; non-gated paths (e.g. NemotronH) must skip it. pub is_gated: bool, + /// SGLang MoE backend (Python `MoE._moe_backend`). `Some("deepep_moe")` + /// routes the compute lookup to the wideep context/generation MoE tables + /// instead of `moe_perf` (operations/moe.py sglang branch). Absent in + /// pre-existing specs -> None -> the regular table. + #[serde(default)] + pub moe_backend: Option, + /// EPLB enabled (Python `MoE._enable_eplb`). On the sglang branch the + /// PREFILL token count is corrected to `int(num_tokens * 0.8)` + /// (operations/moe.py: expert-parallel load balancing evens the + /// per-expert token distribution). + #[serde(default)] + pub enable_eplb: bool, + /// Context (prefill) op — selects the wideep CONTEXT MoE table under + /// deepep and gates the EPLB prefill correction (Python `MoE._is_context`). + #[serde(default)] + pub is_context: bool, } impl MoeOp { @@ -72,6 +88,9 @@ impl MoeOp { quant_mode, workload_distribution: workload_distribution.into(), is_gated: true, + moe_backend: None, + enable_eplb: false, + is_context: false, } } @@ -81,6 +100,15 @@ impl MoeOp { // `MoE.query` (`x = x * attention_dp_size`). Applied exactly once, // before the perf-DB resolution keys off the token count. let num_tokens = num_tokens.saturating_mul(self.attention_dp_size.max(1)); + let is_sglang = db.backend == "sglang"; + // SGLang EPLB prefill correction (Python operations/moe.py: + // `num_tokens_corrected = int(num_tokens * 0.8) if enable_eplb and + // is_context else num_tokens`, sglang branch only). + let num_tokens = if is_sglang && self.enable_eplb && self.is_context { + (num_tokens as f64 * 0.8) as u32 + } else { + num_tokens + }; // The roofline SOL the perf-DB engine anchors its beyond-range // util-hold on (Python `_resolve_tokens` passes the same closure). // Coordinates arriving from the engine are always integral (table @@ -90,6 +118,44 @@ impl MoeOp { // `k_tail=1` util-hold handles beyond-range queries). let sol = |t: f64| self.sol_latency_ms(db, t.round() as u32); + // SGLang DeepEP (wideep) routes MoE compute to the wideep + // context/generation tables (Python operations/moe.py: + // `if moe_backend == "deepep_moe": moe_data = _wideep_*_moe_data`), + // resolved through the SAME `_resolve_tokens` semantics (singleton + // guard + MoE-roofline util-hold, threaded via `sol`). + if is_sglang && self.moe_backend.as_deref() == Some("deepep_moe") { + let latency = if self.is_context { + db.wideep.query_context_moe( + num_tokens, + self.hidden_size, + self.inter_size, + self.topk, + self.num_experts, + self.moe_tp_size, + self.moe_ep_size, + self.quant_mode, + &self.workload_distribution, + &sol, + )? + } else { + db.wideep.query_generation_moe( + num_tokens, + self.hidden_size, + self.inter_size, + self.topk, + self.num_experts, + self.moe_tp_size, + self.moe_ep_size, + self.quant_mode, + &self.workload_distribution, + &sol, + )? + }; + return Ok(PerformanceResult::new(latency, Source::Silicon) + .clamp_non_negative() + .scaled(self.scale_factor)); + } + // Mirrors Python's MoE._query_moe_table TRT-LLM gate: for nvfp4 // gated MoE at num_tokens <= 128, probe the // `moe_torch_flow_min_latency` grid first and fall back to the @@ -206,6 +272,9 @@ mod tests { workload_distribution: "power_law_1.2".into(), attention_dp_size, is_gated: true, + moe_backend: None, + enable_eplb: false, + is_context: false, } } diff --git a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs index 9a4123a1f..3de3140f9 100644 --- a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs +++ b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs @@ -71,12 +71,76 @@ pub struct MoEDispatchOp { /// default, so old opspecs keep the same query point. #[serde(default = "default_sms")] pub sms: u32, + /// Token-count divisor for the DeepEP branches (Python + /// `MoEDispatch._scale_num_tokens`, applied as + /// `num_tokens // scale_num_tokens` before the deepep table lookup). + /// Defaults to 1 (no scaling) for pre-field opspecs. + #[serde(default = "crate::operators::gemm::default_seq_split")] + pub scale_num_tokens: u32, } fn default_sms() -> u32 { 12 } +/// TensorRT-LLM WideEP MoE dispatch (NVLink Two-Sided All2All). Mirrors +/// Python `operations.moe.TrtLLMWideEPMoEDispatch`: +/// +/// - pre-dispatch op: `alltoall_prepare` + `alltoall_dispatch` (two queries, +/// summed); +/// - post-dispatch op: `alltoall_combine`, or `alltoall_combine_low_precision` +/// when `use_low_precision_combine` (models set it for nvfp4 generation); +/// - every query passes `moe_backend="wideep"`, so the kernel auto-selection +/// resolves `NVLinkTwoSided` on MNNVL (SM >= 100) and the DeepEP variants +/// on Hopper (see `perf_database::wideep::select_alltoall_kernel`). +/// +/// `node_num` is never set by the model builders (Python passes None), so the +/// perf-DB default `1 if ep < 4 else ep // 4` applies. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct TrtllmWideEpMoEDispatchOp { + pub name: String, + pub scale_factor: f64, + pub hidden_size: u32, + pub topk: u32, + pub num_experts: u32, + pub moe_tp_size: u32, + pub moe_ep_size: u32, + pub attention_dp_size: u32, + pub pre_dispatch: bool, + pub quant_mode: MoeQuantMode, + #[serde(default)] + pub use_low_precision_combine: bool, +} + +impl TrtllmWideEpMoEDispatchOp { + pub fn query(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + let spec: &SystemSpec = &db.system_spec; + let q = |op_name: &str| { + db.wideep.query_trtllm_alltoall( + spec, + op_name, + num_tokens, + self.hidden_size, + self.topk, + self.num_experts, + self.moe_ep_size, + self.quant_mode, + Some("wideep"), + ) + }; + let latency = if self.pre_dispatch { + q("alltoall_prepare")? + q("alltoall_dispatch")? + } else if self.use_low_precision_combine { + q("alltoall_combine_low_precision")? + } else { + q("alltoall_combine")? + }; + Ok(PerformanceResult::new(latency, Source::Silicon) + .clamp_non_negative() + .scaled(self.scale_factor)) + } +} + impl MoEDispatchOp { pub fn new( name: impl Into, @@ -107,6 +171,7 @@ impl MoEDispatchOp { attn_cp_size: 1, is_context: false, sms: default_sms(), + scale_num_tokens: 1, } } @@ -296,6 +361,9 @@ impl MoEDispatchOp { .scaled(self.scale_factor)) } DispatchFlavor::DeepEpNormal => { + // Python DeepEP branch: `num_tokens = num_tokens // + // self._scale_num_tokens` before the table lookup. + let num_tokens = num_tokens / self.scale_num_tokens.max(1); let point = db.wideep.query_deepep_normal( self.deepep_node_num(spec)?, self.hidden_size, @@ -323,6 +391,9 @@ impl MoEDispatchOp { .scaled(self.scale_factor)) } DispatchFlavor::DeepEpLowLatency => { + // Python DeepEP branch: `num_tokens = num_tokens // + // self._scale_num_tokens` before the table lookup. + let num_tokens = num_tokens / self.scale_num_tokens.max(1); let point = db.wideep.query_deepep_ll( self.deepep_node_num(spec)?, self.hidden_size, @@ -340,76 +411,112 @@ impl MoEDispatchOp { .scaled(self.scale_factor)) } DispatchFlavor::TrtllmAlltoall => { - // Port of Python `MoEDispatch.query` trtllm SM100 branch - // (operations/moe.py). The Python control flow: - // if backend_supports_alltoall && attention_dp > 1 - // && moe_tp == 1 && is_nvl72: -> trtllm_alltoall table - // elif attention_dp > 1: -> NCCL all_gather/reduce_scatter - // elif attention_tp > 1: -> custom_allreduce (when - // reduce_results) else 0 - // else: -> 0 + // Full port of Python `MoEDispatch.query`'s trtllm branch + // (operations/moe.py:1095-1221). Selecting the *flavor* up + // front (as the model builder does) cannot encode the gating — + // it depends on the system's SM version and NVLink topology — + // so `DispatchFlavor::TrtllmAlltoall` means "trtllm dispatch + // op; the *path* is picked here at query time with the spec + // in hand". Two regimes, split on `sm_version == 100` + // EXACTLY like Python (`.get("sm_version", -1) == 100` — + // gb300 is sm 103 and takes the else-branch): + // + // sm == 100: + // if alltoall (dp>1 && moe_tp==1 && NVL72): alltoall table + // (pre -> alltoall_dispatch, combine -> alltoall_combine) + // elif dp>1: NCCL all_gather((x+sf volumes)*dp) pre / + // reduce_scatter(volume*dp) combine — + // quant-compressed volumes on the pre side + // elif tp>1 && reduce_results: custom_allreduce(num_gpus) + // (the NVL72 tp>4 -> NCCL all_reduce reroute lives in + // the DB-level `query_custom_allreduce_scaled`) + // else 0 + // sm != 100 (checks tp FIRST, mirroring Python): + // if tp>1: custom_allreduce(num_gpus, volume) + // elif dp>1: all_gather(volume*dp) pre / + // reduce_scatter(volume*dp) combine + // else 0 // - // Selecting the *flavor* up front (as the model builder does) - // cannot encode this gating — the choice depends on the system's - // NVLink topology (`num_gpus_per_node`) and on tp/dp shapes that - // are only known with the system spec in hand. So - // `DispatchFlavor::TrtllmAlltoall` now means "trtllm SM100 - // dispatch op; the *table* is picked here at query time". - let num_gpus_per_node = spec.node.num_gpus_per_node; - let is_nvl72 = num_gpus_per_node >= 72; - // `moe_backend` is `None` in all current callers (no caller - // sets a non-default backend); treat as supporting alltoall. - let backend_supports_alltoall = true; - let enable_alltoall = backend_supports_alltoall - && self.attention_dp_size > 1 - && self.moe_tp_size == 1 - && is_nvl72; + // `reduce_results` defaults True in Python and no model + // overrides it; `moe_backend` is None in all current callers + // (backend_supports_alltoall = true). Python raises when the + // alltoall path is taken with quant_mode=None; the Rust op's + // `moe_quant` is non-optional so that guard is structural. + let num_gpus = (self.moe_tp_size * self.moe_ep_size).max(1); let attention_tp = self.attention_tp_size(); + let pre = self.pre_dispatch; + let sm_version = spec.gpu.sm_version.map(i64::from).unwrap_or(-1); - if enable_alltoall { - let latency = db.wideep.query_trtllm_alltoall( - num_tokens, - self.hidden_size, - self.topk, - self.num_experts, - self.moe_ep_size, - self.moe_quant, - "uniform", - )?; - Ok(PerformanceResult::new(latency, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) - } else if self.attention_dp_size > 1 { - // Python: query_nccl(half, num_gpus, "all_gather" or - // "reduce_scatter", volume * attention_dp_size). - // No smoke case exercises this branch today (all smoke - // configs use attention_dp_size == 1). The implementation - // is intentionally absent rather than added as untested - // code; a `PerfDatabase` error will surface and any - // future smoke case that hits this path will document - // the need. - Err(AicError::PerfDatabase(format!( - "trtllm MoEDispatch attention_dp_size={}>1 path not yet ported; \ - add a smoke case to fix.", - self.attention_dp_size - ))) + let comm_latency_ms = if sm_version == 100 { + let is_nvl72 = spec.node.num_gpus_per_node >= 72; + let enable_alltoall = + self.attention_dp_size > 1 && self.moe_tp_size == 1 && is_nvl72; + // Quantize-aware dispatch volume factors (elements per + // hidden element): nvfp4 -> x/4 + sf/32; fp8/fp8_block -> + // x/2; others -> full volume (moe.py:1109-1123). + let (x_factor, sf_factor) = match self.moe_quant { + MoeQuantMode::Nvfp4 => (0.25, 0.25 / 8.0), + MoeQuantMode::Fp8 | MoeQuantMode::Fp8Block => (0.5, 0.0), + _ => (1.0, 0.0), + }; + if enable_alltoall { + let op_name = if pre { "alltoall_dispatch" } else { "alltoall_combine" }; + db.wideep.query_trtllm_alltoall( + spec, + op_name, + num_tokens, + self.hidden_size, + self.topk, + self.num_experts, + self.moe_ep_size, + self.moe_quant, + None, + )? + } else if self.attention_dp_size > 1 { + if pre { + // all_gather((dispatch_x + dispatch_sf) * dp): + // fold the quant compression into the per-token + // element count so NcclOp's `tokens * hidden` + // reproduces Python's message size exactly. + let nccl = NcclOp::new( + &self.name, + 1.0, + self.hidden_size as f64 * (x_factor + sf_factor), + num_gpus, + "all_gather", + ); + nccl.query(db, num_tokens * self.attention_dp_size)?.latency_ms + } else { + let nccl = NcclOp::new( + &self.name, + 1.0, + self.hidden_size as f64, + num_gpus, + "reduce_scatter", + ); + nccl.query(db, num_tokens * self.attention_dp_size)?.latency_ms + } + } else if attention_tp > 1 { + let ar = CustomAllReduceOp::new(&self.name, 1.0, self.hidden_size, num_gpus); + ar.query(db, num_tokens)?.latency_ms + } else { + 0.0 + } } else if attention_tp > 1 { - // reduce_results path: Python defaults `_reduce_results` - // to True, and the smoke configs do not override it, so - // the branch we replicate is the custom_allreduce one. - let ar = CustomAllReduceOp::new( - &self.name, - self.scale_factor, - self.hidden_size, - attention_tp, - ); - ar.query(db, num_tokens) + let ar = CustomAllReduceOp::new(&self.name, 1.0, self.hidden_size, num_gpus); + ar.query(db, num_tokens)?.latency_ms + } else if self.attention_dp_size > 1 { + let op_name = if pre { "all_gather" } else { "reduce_scatter" }; + let nccl = + NcclOp::new(&self.name, 1.0, self.hidden_size as f64, num_gpus, op_name); + nccl.query(db, num_tokens * self.attention_dp_size)?.latency_ms } else { - // attn_tp == 1 and attn_dp == 1: no communication needed. - Ok(PerformanceResult::new(0.0, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) - } + 0.0 + }; + + Ok(PerformanceResult::new(comm_latency_ms, Source::Silicon) + .clamp_non_negative() + .scaled(self.scale_factor)) } } } diff --git a/rust/aiconfigurator-core/src/operators/op.rs b/rust/aiconfigurator-core/src/operators/op.rs index 37a480d1d..033551b14 100644 --- a/rust/aiconfigurator-core/src/operators/op.rs +++ b/rust/aiconfigurator-core/src/operators/op.rs @@ -21,6 +21,7 @@ use crate::operators::{ ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4ModuleOp, ElementwiseOp, EmbeddingOp, EncoderAttentionOp, GdnOp, GemmOp, GenerationAttentionOp, GenerationMlaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, MoEDispatchOp, MoeOp, + TrtllmWideEpMoEDispatchOp, NcclOp, P2POp, PerformanceResult, Source, VisionEncoderOp, WideEpContextMlaOp, WideEpGenerationMlaOp, WideEpMoeOp, }; @@ -122,8 +123,12 @@ pub enum Op { WideEpGenerationMla(WideEpGenerationMlaOp), /// TensorRT-LLM WideEP MoE compute. Used by the /// `TrtllmWideEPDeepSeekModel` variant; dispatch / combine cost is - /// modeled separately by `MoEDispatchOp` (TrtllmAlltoall flavor). + /// modeled separately by `WideEpMoeDispatch`. WideEpMoe(WideEpMoeOp), + /// TensorRT-LLM WideEP All2All dispatch (prepare+dispatch / combine). + /// Mirrors Python `TrtLLMWideEPMoEDispatch` — a direct `Operation` + /// subclass, NOT a `MoEDispatch` flavor. + WideEpMoeDispatch(TrtllmWideEpMoEDispatchOp), /// Two op groups that execute in parallel on different CUDA streams. /// Mirrors Python `aiconfigurator.sdk.operations.overlap.OverlapOp`: /// `latency = max(sum(group_a), sum(group_b))`. @@ -208,6 +213,7 @@ impl Op { Op::WideEpContextMla(o) => &o.name, Op::WideEpGenerationMla(o) => &o.name, Op::WideEpMoe(o) => &o.name, + Op::WideEpMoeDispatch(o) => &o.name, Op::Overlap(o) => &o.name, Op::Fallback(o) => &o.name, } @@ -288,6 +294,7 @@ impl Op { Op::WideEpContextMla(op) => op.query(db, ctx.batch_size, ctx.s, ctx.prefix), Op::WideEpGenerationMla(op) => op.query(db, ctx.batch_size, ctx.s), Op::WideEpMoe(op) => op.query(db, ctx.num_tokens), + Op::WideEpMoeDispatch(op) => op.query(db, ctx.num_tokens), Op::Overlap(op) => { // Mirrors Python `OverlapOp.query`: each group is summed // independently, then `max(group_a_total, group_b_total)` is diff --git a/rust/aiconfigurator-core/src/operators/overlap.rs b/rust/aiconfigurator-core/src/operators/overlap.rs index e6f63512d..dcc32c691 100644 --- a/rust/aiconfigurator-core/src/operators/overlap.rs +++ b/rust/aiconfigurator-core/src/operators/overlap.rs @@ -1,73 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Overlap composition: max-of-groups latency aggregation. -//! -//! Mirrors `aiconfigurator.sdk.operations.overlap.Overlap`. When two -//! kernels can execute in parallel (e.g. compute on one stream while a -//! collective runs on another), AIC composes their effective latency as -//! the max of the group rather than the sum. This helper takes a list of -//! `(latency_ms, group_id)` entries and returns `sum(group_max)` across -//! groups. - -use crate::operators::base::{PerformanceResult, Source}; - -/// Compose a list of latencies grouped by overlap-group identifier. -/// Within each group, the effective latency is the max; total is the -/// sum across groups. Sources combine the same way `Source::combine` -/// behaves — same tag across all entries keeps it; any disagreement -/// becomes `Source::Mixed`. -pub fn overlap_composition(entries: &[(PerformanceResult, u32)]) -> PerformanceResult { - use std::collections::BTreeMap; - - let mut by_group: BTreeMap = BTreeMap::new(); - for &(result, group_id) in entries { - let slot = by_group.entry(group_id).or_insert((0.0, result.source)); - if result.latency_ms > slot.0 { - slot.0 = result.latency_ms; - } - slot.1 = slot.1.combine(result.source); - } - let mut total = 0.0_f64; - let mut source = Source::Silicon; - let mut first = true; - for (_, (lat, src)) in by_group { - total += lat; - source = if first { src } else { source.combine(src) }; - first = false; - } - PerformanceResult::new(total, source) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn max_within_group_then_sum_across_groups() { - let entries = vec![ - (PerformanceResult::silicon(10.0), 0), - (PerformanceResult::silicon(8.0), 0), // grouped with above -> max = 10.0 - (PerformanceResult::silicon(5.0), 1), // own group - ]; - let result = overlap_composition(&entries); - assert_eq!(result.latency_ms, 15.0); // 10 + 5 - assert_eq!(result.source, Source::Silicon); - } - - #[test] - fn mixed_sources_yield_mixed_tag() { - let entries = vec![ - (PerformanceResult::silicon(5.0), 0), - (PerformanceResult::new(3.0, Source::Empirical), 1), - ]; - let result = overlap_composition(&entries); - assert_eq!(result.source, Source::Mixed); - } - - #[test] - fn empty_list_is_zero() { - let result = overlap_composition(&[]); - assert_eq!(result.latency_ms, 0.0); - } -} +//! (Retired.) The `overlap_composition` helper that lived here implemented +//! the OPPOSITE composition of Python\'s `OverlapOp` (max-within-group / +//! sum-across-groups instead of sum-within / max-across) and had no +//! production caller — the correct overlap semantics live on +//! `operators::op::OverlapOp` (`latency = max(sum(group_a), sum(group_b))`, +//! mirroring `aiconfigurator.sdk.operations.overlap.OverlapOp`). Deleted +//! rather than fixed so nobody wires the wrong composition by accident. diff --git a/rust/aiconfigurator-core/src/operators/wideep_mla.rs b/rust/aiconfigurator-core/src/operators/wideep_mla.rs index bd49d1faa..d4374e9e0 100644 --- a/rust/aiconfigurator-core/src/operators/wideep_mla.rs +++ b/rust/aiconfigurator-core/src/operators/wideep_mla.rs @@ -31,6 +31,20 @@ fn prefix_correction(full_s: u32, prefix: u32) -> f64 { (f * f - p * p) / (f * f) } +/// Python's `get_silicon` whitelists the attention backend before slicing +/// (`mla.py:1191-1192` generation, `:1459-1460` context: +/// `if attn_backend not in {"flashinfer", "fa3"}: raise ValueError`). +/// Mirror it so both engines error symmetrically on unsupported backends +/// instead of Rust answering from whatever `kernel_source` slice exists. +fn check_attn_backend(attn_backend: &str) -> Result<(), AicError> { + if attn_backend != "flashinfer" && attn_backend != "fa3" { + return Err(AicError::PerfDatabase(format!( + "Unsupported attention backend: {attn_backend}" + ))); + } + Ok(()) +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct WideEpContextMlaOp { pub name: String, @@ -75,6 +89,7 @@ impl WideEpContextMlaOp { isl: u32, prefix: u32, ) -> Result { + check_attn_backend(&self.attn_backend)?; // ctx(s, pfx): the un-sharded wideep context-MLA query for a sequence // chunk of length `s` at prefix `pfx`, with prefix correction applied. let ctx = |s: u32, pfx: u32| -> Result { @@ -110,11 +125,11 @@ mod tests { use super::*; use std::path::PathBuf; - fn b200_sglang_db() -> PerfDatabase { + fn h200_sglang_db() -> PerfDatabase { let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("../..") .join("src/aiconfigurator/systems"); - PerfDatabase::load(&root, "b200_sxm", "sglang", "0.5.10").expect("db loads") + PerfDatabase::load(&root, "h200_sxm", "sglang", "0.5.10").expect("db loads") } fn op(cp_size: u32) -> WideEpContextMlaOp { @@ -124,8 +139,9 @@ mod tests { KvCacheQuantMode::Fp8, FmhaQuantMode::Fp8Block, ); - // b200 sglang 0.5.10 wideep table carries kernel_source=trtllm_mla. - op.attn_backend = "trtllm_mla".to_string(); + // h200 sglang 0.5.10 wideep table carries kernel_source in + // {flashinfer, fa3} — the only backends the Python query whitelists. + op.attn_backend = "flashinfer".to_string(); op.cp_size = cp_size; op } @@ -136,7 +152,7 @@ mod tests { /// cp=1 stays the single full-length query. #[test] fn cp_zigzag_composition() { - let db = b200_sglang_db(); + let db = h200_sglang_db(); let (b, isl, prefix) = (4u32, 4096u32, 0u32); // cp=1 unchanged: exactly the raw single-chunk table query (prefix=0 @@ -150,7 +166,7 @@ mod tests { 128, KvCacheQuantMode::Fp8, FmhaQuantMode::Fp8Block, - "trtllm_mla", + "flashinfer", ) .expect("raw table query"); assert!( @@ -182,6 +198,32 @@ mod tests { baseline.latency_ms ); } + + /// Python whitelists attn_backend in {flashinfer, fa3} before slicing + /// (mla.py:1191-1192, 1459-1460) — an out-of-whitelist backend must error + /// here too, even when a matching kernel_source slice exists in the data + /// (b200's wideep tables carry kernel_source=trtllm_mla). + #[test] + fn attn_backend_whitelist_mirrors_python() { + let db = h200_sglang_db(); + let mut bad = op(1); + bad.attn_backend = "trtllm_mla".to_string(); + assert!(bad.query(&db, 1, 1024, 0).is_err()); + + let mut bad_gen = WideEpGenerationMlaOp::new( + "wideep_gen_mla", + 128, + KvCacheQuantMode::Fp8, + FmhaQuantMode::Fp8Block, + ); + bad_gen.attn_backend = "trtllm_mla".to_string(); + assert!(bad_gen.query(&db, 1, 1024).is_err()); + + // fa3 stays allowed (h200 carries an fa3 slice). + let mut fa3 = op(1); + fa3.attn_backend = "fa3".to_string(); + assert!(fa3.query(&db, 1, 1024, 0).is_ok()); + } } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -220,6 +262,7 @@ impl WideEpGenerationMlaOp { batch_size: u32, s: u32, ) -> Result { + check_attn_backend(&self.attn_backend)?; let latency = db.wideep_mla.query_generation( batch_size, s, diff --git a/rust/aiconfigurator-core/src/operators/wideep_moe.rs b/rust/aiconfigurator-core/src/operators/wideep_moe.rs index 0dff0af80..fbe631d4c 100644 --- a/rust/aiconfigurator-core/src/operators/wideep_moe.rs +++ b/rust/aiconfigurator-core/src/operators/wideep_moe.rs @@ -85,6 +85,11 @@ impl WideEpMoeOp { ) -> Result { // Python: `x = num_tokens * self._attention_dp_size`. let scaled = num_tokens.saturating_mul(self.attention_dp_size.max(1)); + // Beyond-range util-hold roofline (Python `_query_compute_table`'s + // get_sol): num_slots-aware weight-read term keeps small-token holds + // above the launch/weight floor. Coordinates from the engine are + // integral, so rounding keeps floor-division parity with Python. + let sol = |t: f64| self.sol_latency_ms(db, t.round() as u64); let latency = db.wideep_moe.query_compute( scaled, self.hidden_size, @@ -97,9 +102,37 @@ impl WideEpMoeOp { self.quant_mode, &self.workload_distribution, &self.kernel_source, + &sol, )?; Ok(PerformanceResult::new(latency, Source::Silicon) .clamp_non_negative() .scaled(self.scale_factor)) } + + /// WideEP MoE roofline. Verbatim port of Python + /// `TrtllmWideEPMoE._query_compute_table::get_sol` (operations/moe.py): + /// weight memory reads `min(num_slots // ep, total_tokens // ep)` experts + /// (EPLB redundant mode replicates experts across slots). The WideEP MoE + /// path is always gated (SwiGLU, 3 GEMMs). + fn sol_latency_ms(&self, db: &PerfDatabase, num_tokens: u64) -> f64 { + let num_gemms: u64 = 3; + let total_tokens = num_tokens * self.topk as u64; + let moe_ep = (self.moe_ep_size as u64).max(1); + let moe_tp = (self.moe_tp_size as u64).max(1); + let h = self.hidden_size as u64; + let inter = self.inter_size as u64; + let slots = self.num_slots as u64; + + let ops = total_tokens * h * inter * num_gemms * 2 / moe_ep / moe_tp; + let mem_bytes_int = total_tokens / moe_ep * h * 2 + + total_tokens / moe_ep * inter * num_gemms / moe_tp + + h * inter * num_gemms / moe_tp * std::cmp::min(slots / moe_ep, total_tokens / moe_ep); + let mem_bytes = (mem_bytes_int as f64) * self.quant_mode.mapping().memory; + + let spec = &db.system_spec; + let tc_flops = spec.gpu.bfloat16_tc_flops.unwrap_or(1.0); + let sol_math = (ops as f64) / (tc_flops * self.quant_mode.mapping().compute) * 1000.0; + let sol_mem = mem_bytes / spec.gpu.mem_bw * 1000.0; + sol_math.max(sol_mem) + } } diff --git a/rust/aiconfigurator-core/src/perf_database/communication.rs b/rust/aiconfigurator-core/src/perf_database/communication.rs index 33ee7e471..b258a4b61 100644 --- a/rust/aiconfigurator-core/src/perf_database/communication.rs +++ b/rust/aiconfigurator-core/src/perf_database/communication.rs @@ -8,9 +8,12 @@ //! P2P latency is computed analytically by the operator layer from //! `SystemSpec` fields, not from a CSV, so there's no `P2PTable` here. //! -//! Query APIs take *effective* tp_size / num_gpus values — the operator is -//! responsible for capping to the node fan-out and applying any -//! cross-rack bandwidth correction factor (those depend on `SystemSpec`). +//! The `*_scaled` query APIs take RAW tp_size / num_gpus values and own the +//! full Python DB-level semantics (node-fan-out capping, beyond-range +//! bandwidth correction, and the GB200-NVL72 custom-AR -> NCCL reroute) so +//! every consumer inherits them, exactly like Python's `_query_*_table` +//! funnels. The non-`_scaled` variants take *effective* values and only +//! interpolate the table. //! Rows with `_eager` kernel sources are filtered out at load time per //! Python's `CustomAllReduce.load_data` behavior; the production path uses //! CUDA-graph variants. @@ -25,6 +28,7 @@ use std::sync::OnceLock; use crate::common::enums::CommQuantMode; use crate::common::error::AicError; +use crate::common::system_spec::SystemSpec; use crate::config::{PerfDbSources, PerfSource}; use super::{kernel_source_ok, resolve_op_sources}; use super::perf_interp::{self, Node, OpInterpConfig}; @@ -114,7 +118,7 @@ impl CommunicationTable { &self, quant: CommQuantMode, tp_size_effective: u32, - message_size: u64, + message_size: f64, ) -> Result { if tp_size_effective <= 1 { return Ok(0.0); @@ -130,6 +134,70 @@ impl CommunicationTable { interp_message_size(by_size, message_size) } + /// Custom-allreduce latency at a RAW tp_size, mirroring the full Python + /// DB-level `_query_custom_allreduce_table.get_silicon` + /// (operations/communication.py) so every consumer inherits the same + /// semantics: + /// 1. `tp == 1` -> 0; + /// 2. GB200 NVL72 (`num_gpus_per_node == 72`) with `tp > 4` -> reroute + /// to NCCL all_reduce at the RAW tp (custom AR is only collected up + /// to tp4 there); + /// 3. clamp tp to the node size and interpolate the table; + /// 4. beyond-node overflow: scale by the p2p-bandwidth ratio. + pub fn query_custom_allreduce_scaled( + &self, + spec: &SystemSpec, + quant: CommQuantMode, + tp_size: u32, + message_size: f64, + ) -> Result { + if tp_size <= 1 { + return Ok(0.0); + } + let per_node = spec.node.num_gpus_per_node; + if per_node == 72 && tp_size > 4 { + return self.query_nccl_scaled(spec, quant, "all_reduce", tp_size, message_size); + } + let effective_tp = tp_size.min(per_node); + let mut latency = self.query_custom_allreduce(quant, effective_tp, message_size)?; + if tp_size > per_node { + let base_bw = spec.get_p2p_bandwidth(per_node); + let target_bw = spec.get_p2p_bandwidth(tp_size); + let f_tp = tp_size as f64; + let f_pn = per_node as f64; + latency *= (f_tp - 1.0) / f_tp * f_pn / (f_pn - 1.0).max(1.0) * base_bw / target_bw; + } + Ok(latency) + } + + /// NCCL collective latency at a RAW num_gpus, mirroring the Python + /// DB-level `_query_nccl_table.get_silicon`: fan-out capped to the max + /// recorded `num_gpus` for the (dtype, operation) slice, with the + /// p2p-bandwidth correction applied beyond it. + pub fn query_nccl_scaled( + &self, + spec: &SystemSpec, + dtype: CommQuantMode, + operation: &str, + num_gpus: u32, + message_size: f64, + ) -> Result { + if num_gpus <= 1 { + return Ok(0.0); + } + let max_recorded = self.nccl_max_num_gpus(dtype, operation)?.unwrap_or(num_gpus); + let effective = num_gpus.min(max_recorded); + let mut latency = self.query_nccl(dtype, operation, effective, message_size)?; + if num_gpus > max_recorded { + let max_bw = spec.get_p2p_bandwidth(max_recorded); + let req_bw = spec.get_p2p_bandwidth(num_gpus); + let f_n = num_gpus as f64; + let f_m = max_recorded as f64; + latency *= (f_n - 1.0) / f_n * f_m / (f_m - 1.0).max(1.0) * max_bw / req_bw; + } + Ok(latency) + } + /// Raw NCCL collective latency in ms. /// /// `operation` is one of `"all_reduce"`, `"all_gather"`, @@ -144,7 +212,7 @@ impl CommunicationTable { dtype: CommQuantMode, operation: &str, num_gpus_effective: u32, - message_size: u64, + message_size: f64, ) -> Result { if num_gpus_effective <= 1 { return Ok(0.0); @@ -239,7 +307,11 @@ impl CommunicationTable { /// The query coordinate is passed as `f64` without truncation (Python does /// none). Table keys clamp to `u32` only as a defensive bound; every shipped /// comm table tops out at 512 MiB message sizes, well under `u32::MAX`. -fn interp_message_size(by_size: &BTreeMap, message_size: u64) -> Result { +/// Interpolate the 1-D size curve at a possibly FRACTIONAL message size — +/// Python keeps float element counts (e.g. the gemma4 CP KV all-gather sizes +/// `kvcache_bytes_per_token / comm_bytes`), and the engine query coordinate +/// is float anyway. Truncating to integer first shifted the lerp point. +fn interp_message_size(by_size: &BTreeMap, message_size: f64) -> Result { if by_size.is_empty() { return Err(AicError::PerfDatabase( "comm data has no message_size points".to_string(), @@ -251,7 +323,7 @@ fn interp_message_size(by_size: &BTreeMap, message_size: u64) -> Resul } let sol = |c: &[f64]| c[0]; let cfg = OpInterpConfig::grid(&["message_bytes"], &sol); - perf_interp::query(&cfg, &node, &[message_size as f64]) + perf_interp::query(&cfg, &node, &[message_size]) } fn load_custom_allreduce_parquet(sources: &[PerfSource]) -> Result { @@ -377,7 +449,7 @@ mod tests { fn custom_allreduce_tp1_is_zero() { let table = CommunicationTable::new(b200_vllm_data_root(), None, None); let latency = table - .query_custom_allreduce(CommQuantMode::Half, 1, 1024) + .query_custom_allreduce(CommQuantMode::Half, 1, 1024.0) .expect("tp=1 is a no-op"); assert_eq!(latency, 0.0); } @@ -395,7 +467,7 @@ mod tests { let table = CommunicationTable::new(b200_sglang_data_root(), None, None); // SGLang b200 ships custom_allreduce data; pick a small message // and a TP that exists. - let result = table.query_custom_allreduce(CommQuantMode::Half, 2, 1024); + let result = table.query_custom_allreduce(CommQuantMode::Half, 2, 1024.0); match result { Ok(latency) => assert!(latency > 0.0, "expected positive latency"), Err(AicError::PerfDatabase(_)) => { @@ -409,7 +481,7 @@ mod tests { fn nccl_num_gpus_1_is_zero() { let table = CommunicationTable::new(b200_vllm_data_root(), None, None); let latency = table - .query_nccl(CommQuantMode::Half, "all_reduce", 1, 1024) + .query_nccl(CommQuantMode::Half, "all_reduce", 1, 1024.0) .expect("num_gpus=1 is a no-op"); assert_eq!(latency, 0.0); } @@ -455,7 +527,7 @@ mod tests { ]; for &(msg, expected) in cases { let got = table - .query_nccl(CommQuantMode::Half, "all_gather", 8, msg) + .query_nccl(CommQuantMode::Half, "all_gather", 8, msg as f64) .expect("query must succeed"); assert!( ((got - expected) / expected).abs() < 1e-9, @@ -471,7 +543,7 @@ mod tests { // rather than silently degrading. let table = CommunicationTable::new(b200_vllm_data_root(), None, None); let err = table - .query_nccl(CommQuantMode::Half, "all_reduce", 2, 1024) + .query_nccl(CommQuantMode::Half, "all_reduce", 2, 1024.0) .unwrap_err(); match err { AicError::PerfDatabase(msg) => { diff --git a/rust/aiconfigurator-core/src/perf_database/dsa.rs b/rust/aiconfigurator-core/src/perf_database/dsa.rs index 1499ce489..818827a73 100644 --- a/rust/aiconfigurator-core/src/perf_database/dsa.rs +++ b/rust/aiconfigurator-core/src/perf_database/dsa.rs @@ -54,6 +54,14 @@ pub struct DsaTable { generation_sources: Vec, context: OnceLock>, generation: OnceLock>, + /// GLM-5.2 skip-indexer (reuse-layer) tables — same files, rows tagged by + /// `op_name` `*_skip_indexer`. Loaded lazily; empty when the parquet + /// carries no skip rows (DeepSeek-V3.2 / GLM-5), in which case the skip + /// query fails loud exactly like Python's `None` slot. + context_skip: OnceLock>, + generation_skip: OnceLock>, + context_skip_nodes: OnceLock>, + generation_skip_nodes: OnceLock>, /// Engine-ready per-`DsaKey` context tables with the raw shape /// `[num_heads][step][isl][batch]`, built once from the loaded grids. context_nodes: OnceLock>, @@ -208,6 +216,10 @@ impl DsaTable { generation_sources, context: OnceLock::new(), generation: OnceLock::new(), + context_skip: OnceLock::new(), + generation_skip: OnceLock::new(), + context_skip_nodes: OnceLock::new(), + generation_skip_nodes: OnceLock::new(), context_nodes: OnceLock::new(), generation_nodes: OnceLock::new(), perf_db_sources: perf_db_sources.clone(), @@ -310,8 +322,16 @@ impl DsaTable { prefix: u32, index_topk: u32, dsa_backend: &str, + skip_indexer: bool, ) -> Result { - let nodes = self.load_context_nodes()?; + // `skip_indexer=true` reads the GLM-5.2 reuse-layer table (rows tagged + // `*_skip_indexer` in the same parquet) and zeroes the indexer terms + // in the SOL — mirroring Python `_query_context_dsa_module_table`. + let nodes = if skip_indexer { + self.load_context_skip_nodes()? + } else { + self.load_context_nodes()? + }; let key = DsaKey { architecture: architecture.to_string(), fmha_quant: fmha_quant.name().to_string(), @@ -342,6 +362,7 @@ impl DsaTable { c[2] as i64, // s c[1] as i64, // prefix c[0] as i64, // num_heads + skip_indexer, ) }; let cfg = OpInterpConfig::grid(&["num_heads", "prefix", "seq_len", "batch"], &sol); @@ -374,8 +395,16 @@ impl DsaTable { gemm_quant: GemmQuantMode, architecture: &str, dsa_backend: &str, + skip_indexer: bool, ) -> Result { - let nodes = self.load_generation_nodes()?; + // `skip_indexer=true` reads the GLM-5.2 reuse-layer generation table. + // The generation SOL is skip-independent (Python's generation get_sol + // has no skip branch) — only the table slice differs. + let nodes = if skip_indexer { + self.load_generation_skip_nodes()? + } else { + self.load_generation_nodes()? + }; // NOTE: Python's generation table is keyed (kv, gemm, arch) only — // no mla_dtype axis. The Rust key retains `fmha_quant` from the // parquet `mla_dtype` column (uniformly `bfloat16` in collected @@ -417,7 +446,14 @@ impl DsaTable { fn load_context(&self) -> Result<&DsaGrids, AicError> { let cell = self .context - .get_or_init(|| load_dsa_parquet(&self.context_sources, false)); + .get_or_init(|| load_dsa_parquet(&self.context_sources, false, false)); + cell.as_ref().map_err(clone_err) + } + + fn load_context_skip(&self) -> Result<&DsaGrids, AicError> { + let cell = self + .context_skip + .get_or_init(|| load_dsa_parquet(&self.context_sources, false, true)); cell.as_ref().map_err(clone_err) } @@ -427,7 +463,15 @@ impl DsaTable { // (last file row wins within a source; first source wins across). let cell = self .generation - .get_or_init(|| load_dsa_parquet(&self.generation_sources, true)); + .get_or_init(|| load_dsa_parquet(&self.generation_sources, true, false)); + cell.as_ref().map_err(clone_err) + } + + fn load_generation_skip(&self) -> Result<&DsaGrids, AicError> { + // Same load-time (isl, step) -> seq collapse as the full table. + let cell = self + .generation_skip + .get_or_init(|| load_dsa_parquet(&self.generation_sources, true, true)); cell.as_ref().map_err(clone_err) } @@ -439,6 +483,14 @@ impl DsaTable { cell.as_ref().map_err(clone_err) } + fn load_context_skip_nodes(&self) -> Result<&NodeCache, AicError> { + let cell = self.context_skip_nodes.get_or_init(|| { + let grids = self.load_context_skip()?; + Ok(build_context_nodes(grids)) + }); + cell.as_ref().map_err(clone_err) + } + fn load_generation_nodes(&self) -> Result<&NodeCache, AicError> { let cell = self.generation_nodes.get_or_init(|| { let grids = self.load_generation()?; @@ -446,6 +498,14 @@ impl DsaTable { }); cell.as_ref().map_err(clone_err) } + + fn load_generation_skip_nodes(&self) -> Result<&NodeCache, AicError> { + let cell = self.generation_skip_nodes.get_or_init(|| { + let grids = self.load_generation_skip()?; + Ok(build_generation_nodes(grids)) + }); + cell.as_ref().map_err(clone_err) + } } /// Materialise the per-`(DsaKey, dsa_backend)` engine table for @@ -537,6 +597,7 @@ fn indexer_cache_entry_bytes(index_head_dim: i64) -> i64 { /// attention group (fmha_quant) whose exact KV pair count is /// `sum_{i=0..s-1} min(prefix+i+1, index_topk)`. #[allow(clippy::too_many_arguments)] +#[allow(clippy::too_many_arguments)] fn dsa_context_sol_ms( spec: &SystemSpec, dims: &DsaDims, @@ -548,6 +609,7 @@ fn dsa_context_sol_ms( s: i64, prefix: i64, num_heads: i64, + skip_indexer: bool, ) -> f64 { let (hidden, q_lora, kv_lora) = (dims.hidden_size, dims.q_lora_rank, dims.kv_lora_rank); let (inh, ihd) = (dims.index_n_heads, dims.index_head_dim); @@ -575,8 +637,10 @@ fn dsa_context_sol_ms( + 2 * num_heads * tokens * kv_lora * v_dim; // Indexer logits group: always FP8; off when the full sequence fits the - // top-k window (regime split). - let indexer_logits_ops = if full_s <= topk { + // top-k window (regime split). A skip-indexer (reuse) layer never runs + // the per-layer indexer, so the group is zero regardless of full_s + // (Python operations/dsa.py get_sol). + let indexer_logits_ops = if skip_indexer || full_s <= topk { 0 } else { 2 * tokens * inh * ihd * full_s @@ -606,7 +670,8 @@ fn dsa_context_sol_ms( let kv_cache_bytes = (b * num_heads * effective_kv * attn_head_dim) as f64 * kv_quant.mapping().memory; - let indexer_cache_bytes = if full_s <= topk { + // Skip layers never store the index-K cache (the indexer never runs). + let indexer_cache_bytes = if skip_indexer || full_s <= topk { 0.0 } else { (b * full_s * indexer_cache_entry_bytes(dims.index_head_dim) as i128) as f64 @@ -738,6 +803,7 @@ fn dsa_generation_sol_ms( fn load_dsa_parquet( sources: &[PerfSource], collapse_isl_step_to_seq: bool, + want_skip_rows: bool, ) -> Result { let mut by_keys: BTreeMap> = BTreeMap::new(); let mut any_source = false; @@ -770,12 +836,12 @@ fn load_dsa_parquet( } // Full vs skip-indexer share one file, split by op_name (Python: // `if ("skip_indexer" in (row.get("op_name") or "")) != (op_kind - // == "skip"): continue` with op_kind="full"). - if row + // == "skip"): continue`). + let is_skip_row = row .str_optional(op_name_col)? .unwrap_or("") - .contains("skip_indexer") - { + .contains("skip_indexer"); + if is_skip_row != want_skip_rows { continue; } let key = DsaKey { @@ -1068,6 +1134,7 @@ mod tests { 0, INDEX_TOPK, "trtllm", + false, ) .expect("DSA context query must succeed"); assert!( @@ -1123,6 +1190,7 @@ mod tests { 0, INDEX_TOPK, "trtllm", + false, ) .expect("DSA context query must succeed"); approx_rel(latency, 7.756); @@ -1145,6 +1213,7 @@ mod tests { 0, INDEX_TOPK, "trtllm", + false, ) .unwrap_err(); assert!(matches!(err, AicError::PerfDatabase(_))); @@ -1176,6 +1245,7 @@ mod tests { prefix, INDEX_TOPK, "trtllm", + false, ) .unwrap() }; @@ -1401,6 +1471,7 @@ mod tests { GemmQuantMode::Bfloat16, "DeepseekV32ForCausalLM", "trtllm", + false, ) .unwrap() }; @@ -1539,6 +1610,7 @@ mod tests { 0, INDEX_TOPK, dsa_backend, + false, ) .expect("query must succeed") }; @@ -1581,6 +1653,7 @@ mod tests { 0, INDEX_TOPK, dsa_backend, + false, ) .expect("query must succeed") }; @@ -1699,6 +1772,7 @@ mod tests { GemmQuantMode::Bfloat16, "DeepseekV32ForCausalLM", "flashmla_kv", + false, ) .expect("query must succeed"); assert_eq!(got, 222.0); @@ -1729,6 +1803,7 @@ mod tests { 0, INDEX_TOPK, "flashmla_kv", // absent: falls back to the trtllm slice + false, ) .expect("query must succeed"); assert_eq!(got, 3.5); diff --git a/rust/aiconfigurator-core/src/perf_database/dsv4.rs b/rust/aiconfigurator-core/src/perf_database/dsv4.rs index 364a8648b..d8a2b20a2 100644 --- a/rust/aiconfigurator-core/src/perf_database/dsv4.rs +++ b/rust/aiconfigurator-core/src/perf_database/dsv4.rs @@ -121,9 +121,14 @@ struct ModuleNodes { by_keys: BTreeMap>, } +/// Table key mirroring the Python loaders (PR #1337 alignment): +/// - context modules key `[fmha][kv][gemm]` (`load_context_dsv4_kind_module_data`); +/// - generation modules key `[kv][gemm]` only — the `mla_dtype` column is +/// ignored at load (`load_generation_dsv4_kind_module_data`), so +/// `fmha_quant` is the empty sentinel for generation keys. +/// Neither keys on `architecture` (Python never reads that column). #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] struct ModuleKey { - architecture: String, fmha_quant: String, kv_quant: String, gemm_quant: String, @@ -295,7 +300,7 @@ impl Dsv4Table { AttnKind::Csa => self.load_csa_context()?, AttnKind::Hca => self.load_hca_context()?, }; - let node = select_resolved(grids, architecture, fmha_quant, kv_quant, gemm_quant, local_heads)?; + let node = select_resolved(grids, Some(fmha_quant), kv_quant, gemm_quant, local_heads)?; let dims = sol_dims .unwrap_or_else(|| Dsv4SolDims::from_pinned(dsv4_dims(architecture), local_heads as i64)); @@ -348,7 +353,6 @@ impl Dsv4Table { sequence_tokens: u32, local_heads: u32, kv_quant: KvCacheQuantMode, - fmha_quant: FmhaQuantMode, gemm_quant: GemmQuantMode, architecture: &str, sol_dims: Option, @@ -357,7 +361,18 @@ impl Dsv4Table { AttnKind::Csa => self.load_csa_generation()?, AttnKind::Hca => self.load_hca_generation()?, }; - let node = select_resolved(grids, architecture, fmha_quant, kv_quant, gemm_quant, local_heads)?; + // PR #1337: decode attention compute dtype follows the kv-cache dtype; + // the fmha label is inert for generation (the table keys on kv dtype). + // Derive the SOL dtype from kv so label changes cannot move decode SOL + // — mirrors `GenerationDeepSeekV4AttentionModule` (operations/dsv4.py). + // The table key carries no fmha level at all (see `ModuleKey`), so this + // query takes no fmha parameter. + let fmha_quant = if kv_quant == KvCacheQuantMode::Fp8 { + FmhaQuantMode::Fp8 + } else { + FmhaQuantMode::Bfloat16 + }; + let node = select_resolved(grids, None, kv_quant, gemm_quant, local_heads)?; let dims = sol_dims .unwrap_or_else(|| Dsv4SolDims::from_pinned(dsv4_dims(architecture), local_heads as i64)); @@ -394,25 +409,25 @@ impl Dsv4Table { fn load_csa_context(&self) -> Result<&ModuleNodes, AicError> { let cell = self.csa_context.get_or_init(|| { - load_module_parquet(&self.csa_context_sources).map(context_nodes) + load_module_parquet(&self.csa_context_sources, true).map(context_nodes) }); cell.as_ref().map_err(clone_err) } fn load_hca_context(&self) -> Result<&ModuleNodes, AicError> { let cell = self.hca_context.get_or_init(|| { - load_module_parquet(&self.hca_context_sources).map(context_nodes) + load_module_parquet(&self.hca_context_sources, true).map(context_nodes) }); cell.as_ref().map_err(clone_err) } fn load_csa_generation(&self) -> Result<&ModuleNodes, AicError> { let cell = self.csa_generation.get_or_init(|| { - load_module_parquet(&self.csa_generation_sources).map(generation_nodes) + load_module_parquet(&self.csa_generation_sources, false).map(generation_nodes) }); cell.as_ref().map_err(clone_err) } fn load_hca_generation(&self) -> Result<&ModuleNodes, AicError> { let cell = self.hca_generation.get_or_init(|| { - load_module_parquet(&self.hca_generation_sources).map(generation_nodes) + load_module_parquet(&self.hca_generation_sources, false).map(generation_nodes) }); cell.as_ref().map_err(clone_err) } @@ -849,20 +864,19 @@ fn topk_delta_ms(exact: &BTreeMap<(u32, u32, u32), f64>, prefix: u32, isl: u32, } } -/// Resolve the `(quant, architecture)` key, then resolve the model's rank-LOCAL +/// Resolve the quant key ([fmha][kv][gemm] for context, [kv][gemm] for +/// generation — pass `fmha = None`), then resolve the model's rank-LOCAL /// head count against the CSV head keys, returning the engine table for that /// head. fn select_resolved<'a>( grids: &'a ModuleNodes, - architecture: &str, - fmha: FmhaQuantMode, + fmha: Option, kv: KvCacheQuantMode, gemm: GemmQuantMode, local_heads: u32, ) -> Result<&'a Node, AicError> { let key = ModuleKey { - architecture: architecture.to_string(), - fmha_quant: fmha.name().to_string(), + fmha_quant: fmha.map(|f| f.name().to_string()).unwrap_or_default(), kv_quant: kv.name().to_string(), gemm_quant: gemm.name().to_string(), }; @@ -1109,7 +1123,13 @@ fn normalize_dsv4_dtype(name: &str) -> String { /// tp_size axis is collapsed with last-write-wins, so a later source overwrites /// an earlier one at a shared cell (mirroring Python's flat-dict overwrite). An /// error is returned only when no source yields rows. -fn load_module_parquet(sources: &[PerfSource]) -> Result { +/// +/// `key_on_fmha` selects the Python keying (PR #1337): context loaders key +/// `[fmha][kv][gemm]`, generation loaders ignore the `mla_dtype` column and +/// key `[kv][gemm]` (the fmha level of `ModuleKey` stays the empty sentinel, +/// and duplicate rows that differed only in `mla_dtype` collapse last-wins — +/// exactly what Python's direct-assign loader does). +fn load_module_parquet(sources: &[PerfSource], key_on_fmha: bool) -> Result { let mut by_keys: BTreeMap = BTreeMap::new(); let mut any_source = false; for source in sources { @@ -1119,8 +1139,7 @@ fn load_module_parquet(sources: &[PerfSource]) -> Result } any_source = true; let reader = PerfReader::open(path)?; - let arch_col = reader.col("architecture")?; - let mla_dtype_col = reader.col("mla_dtype")?; + let mla_dtype_col = if key_on_fmha { Some(reader.col("mla_dtype")?) } else { None }; let kv_cache_dtype_col = reader.col("kv_cache_dtype")?; let gemm_type_col = reader.col("gemm_type")?; let num_heads_col = reader.col("num_heads")?; @@ -1136,14 +1155,16 @@ fn load_module_parquet(sources: &[PerfSource]) -> Result continue; } let key = ModuleKey { - architecture: row.str_owned(arch_col)?, // CSV columns use sglang dtype naming; the query side builds keys // from the enum `.name()` (canonical short names). Normalize on // load to match Python `_dsv4_normalize_dtype`, which aliases // `fp8_e4m3` -> `fp8` for `mla_dtype` (fmha) and `kv_cache_dtype` // (kv). `gemm_type` is intentionally left untouched, matching // Python (e.g. `fp8_block` is a real value that must pass through). - fmha_quant: normalize_dsv4_dtype(&row.str_owned(mla_dtype_col)?), + fmha_quant: match mla_dtype_col { + Some(col) => normalize_dsv4_dtype(&row.str_owned(col)?), + None => String::new(), + }, kv_quant: normalize_dsv4_dtype(&row.str_owned(kv_cache_dtype_col)?), gemm_quant: row.str_owned(gemm_type_col)?, }; @@ -1252,7 +1273,7 @@ mod tests { let q_gen = |kind, b, s| { table .query_generation( - &spec, kind, b, s, 16, KvCacheQuantMode::Fp8, FmhaQuantMode::Bfloat16, + &spec, kind, b, s, 16, KvCacheQuantMode::Fp8, GemmQuantMode::Fp8Block, "DeepseekV4ForCausalLM", None, ) @@ -1275,9 +1296,11 @@ mod tests { approx(q_ctx(AttnKind::Hca, 1, 128, 0), 0.0802); approx(q_ctx(AttnKind::Hca, 8, 8192, 0), 9.0088); // Generation CSA: exact / interior s / s util-hold / ragged batch. + // Util-hold oracle regenerated post-#1337: the generation SOL now + // derives its fmha dtype from the kv dtype (fp8 here), not the label. approx(q_gen(AttnKind::Csa, 16, 385), 0.1142); approx(q_gen(AttnKind::Csa, 16, 200), 0.11328828125); - approx(q_gen(AttnKind::Csa, 16, 100000), 0.17550076017464525); + approx(q_gen(AttnKind::Csa, 16, 100000), 0.17550461244541488); approx(q_gen(AttnKind::Csa, 15, 385), 0.1129625); // Generation HCA: exact. approx(q_gen(AttnKind::Hca, 16, 385), 0.07239999999999999); @@ -1303,7 +1326,7 @@ mod tests { let q_gen = |kind, b, s| { table .query_generation( - &spec, kind, b, s, 16, KvCacheQuantMode::Fp8, FmhaQuantMode::Bfloat16, + &spec, kind, b, s, 16, KvCacheQuantMode::Fp8, GemmQuantMode::Fp8Block, "DeepseekV4ForCausalLM", None, ) @@ -1750,7 +1773,7 @@ mod tests { let q_gen = |table: &Dsv4Table, kind| { table .query_generation( - &spec, kind, 16, 385, 64, KvCacheQuantMode::Fp8, FmhaQuantMode::Bfloat16, + &spec, kind, 16, 385, 64, KvCacheQuantMode::Fp8, GemmQuantMode::Fp8Block, "DeepseekV4ForCausalLM", None, ) .unwrap() diff --git a/rust/aiconfigurator-core/src/perf_database/moe.rs b/rust/aiconfigurator-core/src/perf_database/moe.rs index 9e68be6ba..fa350705e 100644 --- a/rust/aiconfigurator-core/src/perf_database/moe.rs +++ b/rust/aiconfigurator-core/src/perf_database/moe.rs @@ -321,8 +321,29 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result if !kernel_source_ok(source.kernel_sources(), kernel_source_col, &row)? { continue; } + let kernel_source = row.str_optional(kernel_source_col)?.unwrap_or("").to_string(); + // Kernel-specific mxfp4 remaps (mirror Python `load_moe_data`): + // the collector logs two distinct kernels under one `moe_dtype`; + // route them to dedicated quant modes so DeepSeek-V4 modeling can + // select the right one per GPU generation. + // - Blackwell trtllm-gen MXFP4xMXFP8: + // w4a8_mxfp4_mxfp8 + sglang_mxfp4_flashinfer_trtllm_moe + // -> w4a8_mxfp4_mxfp8_trtllm + // - Hopper flashinfer cutlass SM90 mixed-GEMM: + // w4a16_mxfp4 + sglang_flashinfer_cutlass_moe + // -> w4a16_mxfp4_cutlass + let raw_quant = row.str_owned(moe_dtype_col)?; + let quant = match (raw_quant.as_str(), kernel_source.as_str()) { + ("w4a8_mxfp4_mxfp8", "sglang_mxfp4_flashinfer_trtllm_moe") => { + "w4a8_mxfp4_mxfp8_trtllm".to_string() + } + ("w4a16_mxfp4", "sglang_flashinfer_cutlass_moe") => { + "w4a16_mxfp4_cutlass".to_string() + } + _ => raw_quant, + }; let key = MoeKey { - quant: row.str_owned(moe_dtype_col)?, + quant, distribution: row.str_owned(distribution_col)?, topk: row.u32(topk_col)?, num_experts: row.u32(num_experts_col)?, @@ -331,7 +352,6 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result moe_tp_size: row.u32(moe_tp_size_col)?, moe_ep_size: row.u32(moe_ep_size_col)?, }; - let kernel_source = row.str_optional(kernel_source_col)?.unwrap_or(""); let target = if kernel_source == "moe_torch_flow_min_latency" { &mut low_latency_keys } else { diff --git a/rust/aiconfigurator-core/src/perf_database/wideep.rs b/rust/aiconfigurator-core/src/perf_database/wideep.rs index 7284f6379..c00f639f8 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep.rs @@ -50,6 +50,7 @@ use std::sync::OnceLock; use crate::common::enums::MoeQuantMode; use crate::common::error::AicError; +use crate::common::system_spec::SystemSpec; use crate::config::{PerfDbSources, PerfSource}; use super::perf_interp::{self, Node, OpInterpConfig}; use super::{kernel_source_ok, resolve_op_sources}; @@ -70,7 +71,7 @@ pub struct WideEpTable { context_moe: OnceLock>, generation_moe: OnceLock>, trtllm_wideep_moe: OnceLock>, - trtllm_alltoall: OnceLock>, + trtllm_alltoall: OnceLock>, deepep_normal: OnceLock>, deepep_ll: OnceLock>, } @@ -79,6 +80,27 @@ struct MoeGrids { by_keys: BTreeMap>, } +/// TRT-LLM alltoall grids. Keying mirrors Python `load_trtllm_alltoall_data` +/// exactly: `[kernel_source][op_name][quant][num_nodes][hidden_size][topk] +/// [num_experts][moe_ep_size][num_tokens]`. Note the table has NO +/// `distribution` axis (the parquet column is ignored, as in Python) and NO +/// `inter_size`/`moe_tp_size`. +struct AlltoallGrids { + by_keys: BTreeMap>, +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct AlltoallKey { + kernel_source: String, + op_name: String, + quant: String, + num_nodes: u32, + hidden_size: u32, + topk: u32, + num_experts: u32, + moe_ep_size: u32, +} + struct DispatchGrids { by_keys: BTreeMap>, } @@ -187,6 +209,7 @@ impl WideEpTable { moe_ep_size: u32, quant: MoeQuantMode, workload_distribution: &str, + sol: &dyn Fn(f64) -> f64, ) -> Result { let grids = self.load_context_moe()?; query_moe( @@ -202,6 +225,7 @@ impl WideEpTable { workload_distribution, &self.data_root, true, + sol, ) } @@ -219,6 +243,7 @@ impl WideEpTable { moe_ep_size: u32, quant: MoeQuantMode, workload_distribution: &str, + sol: &dyn Fn(f64) -> f64, ) -> Result { let grids = self.load_generation_moe()?; query_moe( @@ -234,6 +259,7 @@ impl WideEpTable { workload_distribution, &self.data_root, true, + sol, ) } @@ -251,6 +277,8 @@ impl WideEpTable { workload_distribution: &str, ) -> Result { let grids = self.load_trtllm_wideep_moe()?; + // Caller-less duplicate of `wideep_moe.rs::WideEpMoeTable` (the live + // TRT-LLM WideEP compute table); kept on the linear token proxy. query_moe( grids, num_tokens, @@ -264,38 +292,79 @@ impl WideEpTable { workload_distribution, &self.data_root, false, + &|t| t, ) } - /// TRT-LLM alltoall dispatch latency. CSV uses `moe_ep_size` for fan-out - /// (`moe_tp_size`/`inter_size` are not present); the query API still - /// accepts them for shape symmetry but they're effectively ignored. + /// TRT-LLM alltoall latency for one phase op. Mirrors Python + /// `TrtLLMWideEPMoEDispatch._query_alltoall_table` (SILICON path): + /// + /// 1. `node_num` defaults to `1 if ep < 4 else ep // 4` (the Python + /// default when the caller doesn't pass one — no Rust caller does); + /// 2. `op_name` must be one of the four collected phases; + /// 3. the kernel is auto-selected from the system architecture + MoE + /// backend ([`select_alltoall_kernel`]); `NotEnabled` short-circuits + /// to 0.0 (a dense-fallback config has no alltoall cost); + /// 4. `fp8_block` reuses the `fp8` tables (behavioral mode, + /// `_normalize_quant_mode_for_table`); + /// 5. the 1-D token curve resolves RAW-lerp in range; beyond it the + /// boundary util holds on the linear token proxy, ratio-identical to + /// Python's per-slice alltoall SOL (`const * num_tokens`). #[allow(clippy::too_many_arguments)] pub fn query_trtllm_alltoall( &self, + spec: &SystemSpec, + op_name: &str, num_tokens: u32, hidden_size: u32, topk: u32, num_experts: u32, moe_ep_size: u32, quant: MoeQuantMode, - workload_distribution: &str, + moe_backend: Option<&str>, ) -> Result { + const VALID_OP_NAMES: [&str; 4] = [ + "alltoall_prepare", + "alltoall_dispatch", + "alltoall_combine", + "alltoall_combine_low_precision", + ]; + if !VALID_OP_NAMES.contains(&op_name) { + return Err(AicError::PerfDatabase(format!( + "Invalid op_name '{op_name}'. Must be one of {VALID_OP_NAMES:?}" + ))); + } + let kernel_source = select_alltoall_kernel(spec, moe_ep_size, topk, moe_backend); + if kernel_source == "NotEnabled" { + return Ok(0.0); + } + let node_num = if moe_ep_size < 4 { 1 } else { moe_ep_size / 4 }; + // fp8_block reuses the fp8 alltoall tables (Python + // `_normalize_quant_mode_for_table`); the table key is the only + // consumer — the (linear-proxy) SOL is quant-independent. + let table_quant = if quant == MoeQuantMode::Fp8Block { + MoeQuantMode::Fp8 + } else { + quant + }; let grids = self.load_trtllm_alltoall()?; - query_moe( - grids, - num_tokens, + let key = AlltoallKey { + kernel_source: kernel_source.to_string(), + op_name: op_name.to_string(), + quant: table_quant.name().to_string(), + num_nodes: node_num, hidden_size, - 0, topk, num_experts, - 1, moe_ep_size, - quant, - workload_distribution, - &self.data_root, - false, - ) + }; + let by_tokens = grids.by_keys.get(&key).ok_or_else(|| { + AicError::PerfDatabase(format!( + "trtllm alltoall data missing for {key:?} at {}", + self.data_root.display() + )) + })?; + query_token_curve(by_tokens, num_tokens as f64, &|t| t) } /// DeepEP normal-mode dispatch point. @@ -394,10 +463,10 @@ impl WideEpTable { .get_or_init(|| load_moe_parquet(&self.moe_sources)); cell.as_ref().map_err(clone_err) } - fn load_trtllm_alltoall(&self) -> Result<&MoeGrids, AicError> { + fn load_trtllm_alltoall(&self) -> Result<&AlltoallGrids, AicError> { let cell = self .trtllm_alltoall - .get_or_init(|| load_moe_parquet(&self.alltoall_sources)); + .get_or_init(|| load_alltoall_parquet(&self.alltoall_sources)); cell.as_ref().map_err(clone_err) } fn load_deepep_normal(&self) -> Result<&NormalDispatchGrids, AicError> { @@ -428,6 +497,7 @@ fn query_moe( workload_distribution: &str, data_root: &Path, guard_singleton_underflow: bool, + sol: &dyn Fn(f64) -> f64, ) -> Result { let quant_name = quant.name(); let requested_exists = grids @@ -469,7 +539,7 @@ fn query_moe( ))); } } - query_token_curve(by_tokens, num_tokens as f64, &|t| t) + query_token_curve(by_tokens, num_tokens as f64, sol) } /// Resolve one `DispatchPoint` field's token curve on the engine, with the @@ -628,12 +698,14 @@ fn dispatch_lookup( }) } -/// Load a MoE-shape table from an ordered, priority-sorted source list. Sources -/// are read in order; the first source containing a `(key, num_tokens)` wins -/// (`or_insert`), mirroring Python's `_read_filtered_rows` concatenation + -/// `load_wideep_*_moe_data` skip-on-key-conflict. Missing files are skipped; an -/// error is returned only when no source yields rows. Reused for the -/// context/generation/wideep-moe/alltoall parquets. +/// Load a MoE-shape table from an ordered, priority-sorted source list. +/// Sources are read in order and duplicates resolve LAST-wins, mirroring +/// Python `load_wideep_context_moe_data` / `load_wideep_generation_moe_data` +/// / `load_wideep_moe_compute_data`, which direct-assign per coordinate with +/// no `try/except KeyError` guard — both within a file and across the +/// concatenated shared-layer rows. Missing files are skipped; an error is +/// returned only when no source yields rows. Reused for the +/// context/generation/wideep-moe parquets (alltoall has its own loader). fn load_moe_parquet(sources: &[PerfSource]) -> Result { let mut by_keys: BTreeMap> = BTreeMap::new(); let mut any_source = false; @@ -673,13 +745,12 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result { moe_tp_size: row.u32_optional(moe_tp_size_col)?.unwrap_or(1), moe_ep_size: row.u32(moe_ep_size_col)?, }; - // First-wins parity with Python `load_wideep_*_moe_data`, extended - // across shared-layer sources (earlier source wins). + // Last-wins parity with Python `load_wideep_*_moe_data` + // (direct-assign, no skip-on-conflict guard). by_keys .entry(key) .or_default() - .entry(row.u32(num_tokens_col)?) - .or_insert(row.f64(latency_col)?); + .insert(row.u32(num_tokens_col)?, row.f64(latency_col)?); } } if !any_source || by_keys.is_empty() { @@ -692,6 +763,120 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result { Ok(MoeGrids { by_keys }) } +/// Auto-select the TRT-LLM All2All kernel. Verbatim port of Python +/// `TrtLLMWideEPMoEDispatch._select_alltoall_kernel` (operations/moe.py), +/// aligned with TensorRT-LLM's per-backend `select_alltoall_method_type`: +/// +/// - `DEEPGEMM` / `CUTE_DSL` MoE backends never use alltoall; +/// - WideEP: MNNVL (SM >= 100) -> `NVLinkTwoSided`; else DeepEP when feasible +/// (`ep > 1 && topk <= 8`): inter-node (`ep > num_gpus_per_node`) -> +/// `DeepEP`, intra-node -> `DeepEPLowLatency`; else `NotEnabled`; +/// - non-WideEP (Cutlass/TRTLLM): MNNVL -> `NVLinkOneSided`, else `NotEnabled`. +/// +/// Python additionally warns when the preferred kernel is absent from the +/// loaded table but still returns it (the downstream slice miss surfaces the +/// error) — behavior-identical, so the warning is not replicated. +pub(crate) fn select_alltoall_kernel( + spec: &SystemSpec, + moe_ep_size: u32, + topk: u32, + moe_backend: Option<&str>, +) -> &'static str { + if let Some(backend) = moe_backend { + let upper = backend.to_uppercase(); + if upper == "DEEPGEMM" || upper == "CUTE_DSL" { + return "NotEnabled"; + } + } + let supports_mnnvl = spec.gpu.sm_version.unwrap_or(0) >= 100; + let is_wideep = moe_backend.map(|b| b.to_uppercase() == "WIDEEP").unwrap_or(false); + if is_wideep { + if supports_mnnvl { + return "NVLinkTwoSided"; + } + let deepep_feasible = moe_ep_size > 1 && topk <= 8; + let is_inter_node = moe_ep_size > spec.node.num_gpus_per_node; + if deepep_feasible && is_inter_node { + "DeepEP" + } else if deepep_feasible { + "DeepEPLowLatency" + } else { + "NotEnabled" + } + } else if supports_mnnvl { + "NVLinkOneSided" + } else { + "NotEnabled" + } +} + +/// Load the TRT-LLM alltoall table from an ordered source list. Mirrors +/// Python `load_trtllm_alltoall_data` exactly: +/// +/// - key `[kernel_source][op_name][quant][num_nodes][hidden][topk] +/// [num_experts][moe_ep_size]` -> `{num_tokens -> latency}`; +/// - `kernel_source` defaults to `"NVLinkTwoSided"` when the column is +/// absent; `num_nodes` defaults to `max(1, moe_ep_size // 4)` (GB200 NVL4); +/// - the `distribution` column is IGNORED (Python never reads it); +/// - duplicates resolve LAST-wins (Python direct-assigns per coordinate over +/// the concatenated shared-layer rows). +fn load_alltoall_parquet(sources: &[PerfSource]) -> Result { + let mut by_keys: BTreeMap> = BTreeMap::new(); + let mut any_source = false; + for source in sources { + let path = source.path(); + if !path.exists() { + continue; + } + any_source = true; + let reader = PerfReader::open(path)?; + let op_name_col = reader.col("op_name")?; + let moe_dtype_col = reader.col("moe_dtype")?; + let num_tokens_col = reader.col("num_tokens")?; + let hidden_size_col = reader.col("hidden_size")?; + let topk_col = reader.col("topk")?; + let num_experts_col = reader.col("num_experts")?; + let moe_ep_size_col = reader.col("moe_ep_size")?; + let latency_col = reader.col("latency")?; + let ks_col = reader.col_optional("kernel_source"); + let num_nodes_col = reader.col_optional("num_nodes"); + for row in reader.rows()? { + let row = row?; + if !kernel_source_ok(source.kernel_sources(), ks_col, &row)? { + continue; + } + let moe_ep_size = row.u32(moe_ep_size_col)?; + let key = AlltoallKey { + kernel_source: row + .str_optional(ks_col)? + .map(|s| s.to_string()) + .unwrap_or_else(|| "NVLinkTwoSided".to_string()), + op_name: row.str_owned(op_name_col)?, + quant: row.str_owned(moe_dtype_col)?, + num_nodes: row + .u32_optional(num_nodes_col)? + .unwrap_or_else(|| (moe_ep_size / 4).max(1)), + hidden_size: row.u32(hidden_size_col)?, + topk: row.u32(topk_col)?, + num_experts: row.u32(num_experts_col)?, + moe_ep_size, + }; + by_keys + .entry(key) + .or_default() + .insert(row.u32(num_tokens_col)?, row.f64(latency_col)?); + } + } + if !any_source || by_keys.is_empty() { + return Err(AicError::PerfDatabase(format!( + "no TRT-LLM alltoall rows loaded from {} source(s) (first: {})", + sources.len(), + sources.first().map(|s| s.path().display().to_string()).unwrap_or_default() + ))); + } + Ok(AlltoallGrids { by_keys }) +} + /// Load the DeepEP-normal dispatch table from an ordered source list. Missing /// files are skipped; an error is returned only when no source yields rows. /// @@ -1036,6 +1221,65 @@ mod tests { assert_eq!(point.dispatch_transmit_us, 100.0); } + fn gb200_spec() -> SystemSpec { + let yaml = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../src/aiconfigurator/systems/gb200.yaml"); + SystemSpec::load(&yaml).expect("gb200.yaml must parse") + } + + fn gb200_trtllm_table() -> WideEpTable { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../src/aiconfigurator/systems/data/gb200/trtllm/1.3.0rc10"); + WideEpTable::new(root) + } + + /// Kernel auto-selection mirrors Python `_select_alltoall_kernel`: + /// gb200 (sm 100): non-WideEP -> NVLinkOneSided, WideEP -> NVLinkTwoSided, + /// DEEPGEMM/CUTE_DSL -> NotEnabled (query returns 0.0). + #[test] + fn alltoall_kernel_selection_matches_python() { + let spec = gb200_spec(); + assert_eq!(select_alltoall_kernel(&spec, 4, 8, None), "NVLinkOneSided"); + assert_eq!(select_alltoall_kernel(&spec, 4, 8, Some("WIDEEP")), "NVLinkTwoSided"); + assert_eq!(select_alltoall_kernel(&spec, 4, 8, Some("DeepGemm")), "NotEnabled"); + assert_eq!(select_alltoall_kernel(&spec, 4, 8, Some("cute_dsl")), "NotEnabled"); + let table = gb200_trtllm_table(); + let zero = table + .query_trtllm_alltoall(&spec, "alltoall_dispatch", 1, 7168, 8, 256, 4, MoeQuantMode::Fp8, Some("DEEPGEMM")) + .expect("NotEnabled short-circuits"); + assert_eq!(zero, 0.0); + } + + /// Exact-hit anchors from gb200/trtllm/1.3.0rc10/trtllm_alltoall_perf.parquet. + /// The pre-fix loader collapsed kernel_source/op_name/num_nodes (1,556 of + /// 2,096 rows collided) and the query keyed distribution="uniform" (data is + /// "balanced") — these anchors fail on both bugs. + #[test] + fn alltoall_exact_hits_match_parquet_rows() { + let spec = gb200_spec(); + let table = gb200_trtllm_table(); + // WideEP -> NVLinkTwoSided; fp8 dispatch row (ep=4 -> node_num=1). + let dispatch = table + .query_trtllm_alltoall(&spec, "alltoall_dispatch", 1, 7168, 8, 256, 4, MoeQuantMode::Fp8, Some("WIDEEP")) + .expect("dispatch row"); + assert!((dispatch - 0.011_372_800_171_375_274).abs() < 1e-12, "got {dispatch}"); + // Same slice, combine phase: distinct value proves op_name keys the table. + let combine = table + .query_trtllm_alltoall(&spec, "alltoall_combine", 1, 7168, 8, 256, 4, MoeQuantMode::Fp8, Some("WIDEEP")) + .expect("combine row"); + assert!((combine - 0.012_921_600_043_773_651).abs() < 1e-12, "got {combine}"); + // fp8_block reuses the fp8 tables (Python `_normalize_quant_mode_for_table`). + let block = table + .query_trtllm_alltoall(&spec, "alltoall_dispatch", 1, 7168, 8, 256, 4, MoeQuantMode::Fp8Block, Some("WIDEEP")) + .expect("fp8_block reroutes to fp8"); + assert_eq!(block, dispatch); + // Non-WideEP -> NVLinkOneSided (nvfp4-only slice, ep=2 -> node_num=1). + let one_sided = table + .query_trtllm_alltoall(&spec, "alltoall_dispatch", 1, 7168, 8, 256, 2, MoeQuantMode::Nvfp4, None) + .expect("one-sided row"); + assert!((one_sided - 0.012_895_999_848_842_621).abs() < 1e-12, "got {one_sided}"); + } + #[test] fn wideep_loaders_smoke() { // None of the WideEP/DeepEP data exists on vLLM b200 (TRT-LLM/SGLang @@ -1054,6 +1298,7 @@ mod tests { 8, MoeQuantMode::Bfloat16, "uniform", + &|t| t, ) .unwrap_err(); match err { diff --git a/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs b/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs index 3c240f32d..fbc1fc8d7 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs @@ -121,6 +121,7 @@ impl WideEpMoeTable { quant: MoeQuantMode, distribution: &str, kernel_source: &str, + sol: &dyn Fn(f64) -> f64, ) -> Result { let grids = self.load_compute()?; @@ -193,7 +194,10 @@ impl WideEpMoeTable { self.data_root.display() ))); } - query_token_curve(by_tokens, num_tokens as f64, &|t| t) + // Beyond-range holds anchor on the caller's num_slots-aware roofline + // (Python `_query_compute_table`'s get_sol); in-range lerp never + // consults it. + query_token_curve(by_tokens, num_tokens as f64, sol) } fn load_compute(&self) -> Result<&WideEpMoeGrids, AicError> { @@ -255,13 +259,20 @@ fn load_compute_parquet(sources: &[PerfSource]) -> Result, @@ -236,10 +240,19 @@ impl AicEngine { isl: u32, osl: u32, prefix: u32, + seq_imbalance_correction_scale: f64, + gen_seq_imbalance_correction_scale: f64, ) -> PyResult { py.allow_threads(|| { - self.inner - .mixed_step_latency(ctx_tokens, gen_tokens, isl, osl, prefix) + self.inner.mixed_step_latency( + ctx_tokens, + gen_tokens, + isl, + osl, + prefix, + seq_imbalance_correction_scale, + gen_seq_imbalance_correction_scale, + ) }) .map_err(aic_to_py) } @@ -247,17 +260,20 @@ impl AicEngine { /// One generation-only engine-step latency in ms. Binds /// [`Engine::decode_step_latency`]; the Python agg orchestration /// (`base_backend._get_genonly_step_latency`) calls this per genonly step. - /// Mirrors the live FPM bridge `estimate_decode_step_latency_with_rust`. - #[pyo3(signature = (gen_tokens, isl, osl))] + #[pyo3(signature = (gen_tokens, isl, osl, gen_seq_imbalance_correction_scale=1.0))] fn decode_step_latency( &self, py: Python<'_>, gen_tokens: u32, isl: u32, osl: u32, + gen_seq_imbalance_correction_scale: f64, ) -> PyResult { - py.allow_threads(|| self.inner.decode_step_latency(gen_tokens, isl, osl)) - .map_err(aic_to_py) + py.allow_threads(|| { + self.inner + .decode_step_latency(gen_tokens, isl, osl, gen_seq_imbalance_correction_scale) + }) + .map_err(aic_to_py) } } @@ -678,6 +694,14 @@ mod tests { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../src/aiconfigurator/systems") } + /// `cargo test` runs without an embedding host, so the interpreter must + /// be initialized before any `Python::with_gil` (pyo3's + /// `auto-initialize` feature is intentionally off for the extension + /// build). Idempotent. + fn py_init() { + pyo3::prepare_freethreaded_python(); + } + const TEST_MODEL: &str = "MiniMaxAI/MiniMax-M2.5"; /// Hand-built context op list against the b200_sxm/vllm/0.19.0 perf tables. @@ -692,6 +716,7 @@ mod tests { name: "rmsnorm".into(), scale_factor: 1.0, bytes_per_token: 8192.0, + scale_num_tokens: 1, seq_split: 1, }), Op::Gemm(GemmOp { @@ -727,6 +752,7 @@ mod tests { name: "rmsnorm".into(), scale_factor: 1.0, bytes_per_token: 8192.0, + scale_num_tokens: 1, seq_split: 1, }), Op::GenerationAttention(GenerationAttentionOp { @@ -783,6 +809,7 @@ mod tests { /// `Engine` built from the same bytes via `from_spec_bytes`. #[test] fn aic_engine_matches_raw_engine() { + py_init(); let bytes = fixture_spec_bytes(); let root = systems_root(); @@ -859,6 +886,7 @@ mod tests { /// unknown mode must raise (not silently default). #[test] fn mode_strings_map_correctly() { + py_init(); let bytes = fixture_spec_bytes(); let root = systems_root(); let aic = AicEngine::from_spec(&bytes, root.to_str()).unwrap(); @@ -885,17 +913,19 @@ mod tests { /// the raw `Engine` numbers unchanged. #[test] fn per_step_bindings_match_raw_engine() { + py_init(); let bytes = fixture_spec_bytes(); let root = systems_root(); let raw = Engine::from_spec_bytes(&bytes, &root).unwrap(); let aic = AicEngine::from_spec(&bytes, root.to_str()).unwrap(); - let raw_mixed = raw.mixed_step_latency(1024, 2, 1024, 8, 0).unwrap(); - let mixed = Python::with_gil(|py| aic.mixed_step_latency(py, 1024, 2, 1024, 8, 0)).unwrap(); + let raw_mixed = raw.mixed_step_latency(1024, 2, 1024, 8, 0, 1.0, 1.0).unwrap(); + let mixed = + Python::with_gil(|py| aic.mixed_step_latency(py, 1024, 2, 1024, 8, 0, 1.0, 1.0)).unwrap(); assert!((mixed - raw_mixed).abs() < 1e-12); - let raw_decode = raw.decode_step_latency(4, 1024, 8).unwrap(); - let decode = Python::with_gil(|py| aic.decode_step_latency(py, 4, 1024, 8)).unwrap(); + let raw_decode = raw.decode_step_latency(4, 1024, 8, 1.0).unwrap(); + let decode = Python::with_gil(|py| aic.decode_step_latency(py, 4, 1024, 8, 1.0)).unwrap(); assert!((decode - raw_decode).abs() < 1e-12); } @@ -917,6 +947,7 @@ mod tests { /// exercises end-to-end. #[test] fn inherent_predict_matches_raw_engine() { + py_init(); let bytes = fixture_spec_bytes(); let root = systems_root(); let raw = Engine::from_spec_bytes(&bytes, &root).unwrap(); diff --git a/rust/aiconfigurator-core/src/session.rs b/rust/aiconfigurator-core/src/session.rs index 999fe323e..81dad1951 100644 --- a/rust/aiconfigurator-core/src/session.rs +++ b/rust/aiconfigurator-core/src/session.rs @@ -14,23 +14,46 @@ use crate::common::error::AicError; use crate::operators::{Op, RuntimeContext}; use crate::perf_database::PerfDatabase; -/// Python `_run_context_phase` (`base_backend.py:144`) — one full pass over -/// the context op list. A free function so the compiled +/// Context-op selection for [`run_context_ops`]. Mirrors the name-based +/// filtering Python's `_get_mix_step_latency` applies to `run_static`'s +/// per-op breakdown: pass 1 keeps every op EXCEPT `"context_attention"`, +/// pass 2 keeps ONLY `"context_attention"`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ContextOpFilter { + All, + SkipContextAttention, + OnlyContextAttention, +} + +/// Python `_run_context_phase` (`base_backend.py:144`) — one pass over the +/// context op list. A free function so the compiled /// [`crate::engine::Engine`] iterates one canonical body over its op list. /// /// `effective_isl = isl - prefix` is the caller's responsibility (Python /// computes it in `_run_context_phase` and `_run_static_breakdown`); this /// fn takes the already-effective ISL and does NOT validate it (the Engine /// caller performs Python's `effective_isl > 0` check before calling). +/// +/// `seq_imbalance_correction_scale` mirrors Python's per-op kwarg +/// (`base_backend.py:331`) — the context-attention ops multiply their result +/// by it (`operations/attention.py:550-552`); every other op ignores it. pub(crate) fn run_context_ops( ops: &[Op], db: &PerfDatabase, batch_size: u32, effective_isl: u32, prefix: u32, + seq_imbalance_correction_scale: f64, + filter: ContextOpFilter, ) -> Result { let mut total = 0.0_f64; for op in ops { + match filter { + ContextOpFilter::All => {} + ContextOpFilter::SkipContextAttention if op.is_context_attention() => continue, + ContextOpFilter::OnlyContextAttention if !op.is_context_attention() => continue, + _ => {} + } let x = if op.is_logits_gemm() { batch_size } else { @@ -42,7 +65,7 @@ pub(crate) fn run_context_ops( s: effective_isl, prefix, num_tokens: x, - seq_imbalance_correction_scale: 1.0, + seq_imbalance_correction_scale, gen_seq_imbalance_correction_scale: 1.0, num_image_tokens: 0, }; @@ -57,22 +80,59 @@ pub(crate) fn run_context_ops( /// `kv_seq_tokens` (= Python `s = isl + i + 1`); the stride quadrature /// (`for i in range(0, osl-1, stride)` × `repeat_count`) and the `(nextn + 1)` /// decode-batch multiplier are applied by the caller (the Engine). +/// +/// `gen_seq_imbalance_correction_scale` mirrors Python's per-op kwarg +/// (`base_backend.py:372`) — generation-attention ops multiply their result +/// by it (`operations/attention.py:899-906`). `only_generation_attention` +/// mirrors the mixed-step pass-3 name filter (only +/// `latency_dict["generation_attention"]` is read). pub(crate) fn run_generation_ops_step( ops: &[Op], db: &PerfDatabase, batch_size: u32, kv_seq_tokens: u32, + gen_seq_imbalance_correction_scale: f64, + only_generation_attention: bool, +) -> Result { + run_generation_ops_step_beamed( + ops, + db, + batch_size, + 1, + kv_seq_tokens, + gen_seq_imbalance_correction_scale, + only_generation_attention, + ) +} + +/// [`run_generation_ops_step`] with an explicit beam width. Python's +/// `_run_generation_phase` queries with `x = batch_size * beam_width` while +/// `batch_size` itself stays unscaled (`base_backend.py:368-372`) — token-major +/// ops (GEMM, elementwise, comm) see the beam-scaled token count, attention +/// ops key on the raw decode batch. +#[allow(clippy::too_many_arguments)] +pub(crate) fn run_generation_ops_step_beamed( + ops: &[Op], + db: &PerfDatabase, + batch_size: u32, + beam_width: u32, + kv_seq_tokens: u32, + gen_seq_imbalance_correction_scale: f64, + only_generation_attention: bool, ) -> Result { let mut total = 0.0_f64; for op in ops { + if only_generation_attention && !op.is_generation_attention() { + continue; + } let ctx = RuntimeContext { batch_size, - beam_width: 1, + beam_width: beam_width.max(1), s: kv_seq_tokens, prefix: 0, - num_tokens: batch_size, + num_tokens: batch_size.saturating_mul(beam_width.max(1)), seq_imbalance_correction_scale: 1.0, - gen_seq_imbalance_correction_scale: 1.0, + gen_seq_imbalance_correction_scale, num_image_tokens: 0, }; total += op.query(db, &ctx)?.latency_ms; @@ -80,13 +140,12 @@ pub(crate) fn run_generation_ops_step( Ok(total) } -/// Python `_get_mix_step_latency` (Rust-shaped) — the three-pass mix-step -/// composition for the agg path's chunked-prefill + decode step. A free -/// function driven by the compiled [`crate::engine::Engine`]. The passes filter -/// differently from the full-pass loops — pass 1 skips `context_attention`, -/// pass 2 keeps only `context_attention`, pass 3 keeps only -/// `generation_attention` — so it does NOT route through [`run_context_ops`] / -/// [`run_generation_ops_step`]. +/// FPM-telemetry mix-step composition (the Dynamo mocker contract). Consumed +/// ONLY by `Engine::rank_latency_ms` (observed `ForwardPassMetrics` +/// dispatch); the live SDK engine-step path (`Engine::mixed_step_latency`) +/// mirrors Python's `_get_mix_step_latency` three-pass composition directly +/// via [`run_context_ops`] / [`run_generation_ops_step`] filters and does NOT +/// route through here. /// /// Algorithm (mirrors Python): /// 1. Combined non-attention pass: iterate `context_ops`, **skip** diff --git a/src/aiconfigurator/sdk/backends/base_backend.py b/src/aiconfigurator/sdk/backends/base_backend.py index ead86b89c..f5711fac9 100644 --- a/src/aiconfigurator/sdk/backends/base_backend.py +++ b/src/aiconfigurator/sdk/backends/base_backend.py @@ -16,6 +16,7 @@ from aiconfigurator.sdk.models import BaseModel from aiconfigurator.sdk.perf_database import PerfDatabase from aiconfigurator.sdk.rust_engine_step import ( + RustEngineUnsupportedError, estimate_decode_step_latency_with_rust, estimate_mixed_step_latency_with_rust, estimate_static_latency_breakdown_with_rust, @@ -420,33 +421,43 @@ def _run_static_breakdown( generation_latency_dict, generation_energy_wms_dict, generation_source_dict = {}, {}, {} if should_use_rust_engine_step(runtime_config, database): - rust_runtime_config = runtime_config - if img_ctx_tokens: - rust_runtime_config = copy.copy(runtime_config) - rust_runtime_config.isl = isl_eff - ( - context_latency_dict, - generation_latency_dict, - context_source_dict, - generation_source_dict, - ) = estimate_static_latency_breakdown_with_rust( - model, - database, - rust_runtime_config, - mode, - stride, - latency_correction_scale, - ) - context_energy_wms_dict = dict.fromkeys(context_latency_dict, 0.0) - generation_energy_wms_dict = dict.fromkeys(generation_latency_dict, 0.0) - return ( - context_latency_dict, - context_energy_wms_dict, - generation_latency_dict, - generation_energy_wms_dict, - context_source_dict, - generation_source_dict, - ) + try: + rust_runtime_config = runtime_config + if img_ctx_tokens: + rust_runtime_config = copy.copy(runtime_config) + rust_runtime_config.isl = isl_eff + ( + context_latency_dict, + generation_latency_dict, + context_source_dict, + generation_source_dict, + ) = estimate_static_latency_breakdown_with_rust( + model, + database, + rust_runtime_config, + mode, + stride, + latency_correction_scale, + ) + context_energy_wms_dict = dict.fromkeys(context_latency_dict, 0.0) + generation_energy_wms_dict = dict.fromkeys(generation_latency_dict, 0.0) + return ( + context_latency_dict, + context_energy_wms_dict, + generation_latency_dict, + generation_energy_wms_dict, + context_source_dict, + generation_source_dict, + ) + except RustEngineUnsupportedError as exc: + # Op graph not expressible as a compiled EngineSpec: fall back + # to the Python step (parity by delegation — Python computes + # what Rust cannot yet express). Perf-data misses are NOT + # caught here; they must stay error-symmetric. + logger.warning( + "engine-step backend 'rust' cannot compile this model; using the python step: %s", + exc, + ) if mode == "static_ctx": context_latency_dict, context_energy_wms_dict, context_source_dict = self._run_context_phase( @@ -931,21 +942,29 @@ def _get_mix_step_latency( per-ops summary. """ if should_use_rust_engine_step(runtime_config, database): - latency_ms = estimate_mixed_step_latency_with_rust( - model, - database, - ctx_tokens=ctx_tokens, - gen_tokens=gen_tokens, - isl=isl, - osl=osl, - prefix=prefix, - ) - return ( - latency_ms, - 0.0, - {"rust_engine_step_mixed": latency_ms}, - {"rust_engine_step_mixed": "rust"}, - ) + try: + latency_ms = estimate_mixed_step_latency_with_rust( + model, + database, + ctx_tokens=ctx_tokens, + gen_tokens=gen_tokens, + isl=isl, + osl=osl, + prefix=prefix, + seq_imbalance_correction_scale=runtime_config.seq_imbalance_correction_scale, + gen_seq_imbalance_correction_scale=runtime_config.gen_seq_imbalance_correction_scale, + ) + return ( + latency_ms, + 0.0, + {"rust_engine_step_mixed": latency_ms}, + {"rust_engine_step_mixed": "rust"}, + ) + except RustEngineUnsupportedError as exc: + logger.warning( + "engine-step backend 'rust' cannot compile this model; using the python step: %s", + exc, + ) ctx_scale = runtime_config.seq_imbalance_correction_scale gen_scale = runtime_config.gen_seq_imbalance_correction_scale @@ -1057,19 +1076,26 @@ def _get_genonly_step_latency( if gen_tokens <= 0: return 0.0, 0.0, {}, {} if should_use_rust_engine_step(runtime_config, database): - latency_ms = estimate_decode_step_latency_with_rust( - model, - database, - gen_tokens=gen_tokens, - isl=isl, - osl=osl, - ) - return ( - latency_ms, - 0.0, - {"rust_engine_step_generation": latency_ms}, - {"rust_engine_step_generation": "rust"}, - ) + try: + latency_ms = estimate_decode_step_latency_with_rust( + model, + database, + gen_tokens=gen_tokens, + isl=isl, + osl=osl, + gen_seq_imbalance_correction_scale=runtime_config.gen_seq_imbalance_correction_scale, + ) + return ( + latency_ms, + 0.0, + {"rust_engine_step_generation": latency_ms}, + {"rust_engine_step_generation": "rust"}, + ) + except RustEngineUnsupportedError as exc: + logger.warning( + "engine-step backend 'rust' cannot compile this model; using the python step: %s", + exc, + ) gen_scale = runtime_config.gen_seq_imbalance_correction_scale summary = self.run_static( diff --git a/src/aiconfigurator/sdk/engine.py b/src/aiconfigurator/sdk/engine.py index e4e6bbe0c..8d25d4924 100644 --- a/src/aiconfigurator/sdk/engine.py +++ b/src/aiconfigurator/sdk/engine.py @@ -34,6 +34,7 @@ from __future__ import annotations import json +import logging import os from typing import Any @@ -66,6 +67,7 @@ MoEDispatch, OverlapOp, TrtLLMWideEPMoE, + TrtLLMWideEPMoEDispatch, WideEPContextMLA, WideEPGenerationMLA, ) @@ -92,6 +94,8 @@ ENGINE_SPEC_SCHEMA_VERSION = 2 ENGINE_CONFIG_SCHEMA_VERSION = 1 +logger = logging.getLogger(__name__) + class OpConversionError(RuntimeError): """Raised when an ``Operation`` cannot be converted to an ``OpSpec``.""" @@ -147,15 +151,15 @@ def _embedding(op: Embedding) -> dict: def _elementwise(op: ElementWise) -> dict: - # Rust `ElementwiseOp.bytes_per_token * num_tokens` must equal Python's - # `(x * dim_in * 2 + x * dim_out * 2)` after `x //= scale_num_tokens`. - # Per-token bytes = (dim_in * 2 + dim_out * 2) / scale_num_tokens. - scale = op._scale_num_tokens if op._scale_num_tokens else 1 - bytes_per_token = (op._dim_in * 2 + op._dim_out * 2) / scale + # `scale_num_tokens` rides the wire as its own field so the Rust op can + # reproduce Python's integer order exactly (`x //= scale` THEN the CP + # ceil-split). Folding it into bytes_per_token (the old encoding) is only + # exact when the divisor divides the token count. return { "name": op._name, "scale_factor": op._scale_factor, - "bytes_per_token": float(bytes_per_token), + "bytes_per_token": float(op._dim_in * 2 + op._dim_out * 2), + "scale_num_tokens": op._scale_num_tokens if op._scale_num_tokens else 1, "seq_split": op._seq_split, } @@ -194,6 +198,9 @@ def _encoder_attention(op: EncoderAttention) -> dict: "n": op._n, "head_size": op._head_size, "fmha_quant_mode": _quant_name(op._fmha_quant_mode), + # Partial-RoPE extra (Qwen3-VL uses 0.5); Rust adds + # `factor * 2 * mem_op(Q+K bytes) * 1.1` on top of the table latency. + "partial_rotary_factor": float(getattr(op, "_partial_rotary_factor", 0.0) or 0.0), } @@ -252,26 +259,33 @@ def _moe(op: MoE) -> dict: "quant_mode": _quant_name(op._quant_mode), "workload_distribution": op._workload_distribution, "is_gated": op._is_gated, + # SGLang MoE routing (Python `MoE.query` sglang branch): deepep_moe + # reads the wideep context/generation tables; EPLB corrects the + # prefill token count to int(x * 0.8). + "moe_backend": op._moe_backend, + "enable_eplb": bool(op._enable_eplb), + "is_context": bool(op._is_context), } -def _dispatch_flavor(backend: str) -> str: - """Mirror Rust `models/moe.rs::dispatch_flavor`: trtllm -> TrtllmAlltoall, - everything else -> CustomAllReduce. The `DispatchFlavor` enum has no serde - rename, so emit the exact Rust variant name. The fine SGLang-DeepEP / - NVL72 gating stays inside the Rust `MoEDispatchOp::query`. - - KNOWN LIMITATION (WideEP / DeepEP, out of scope): this binary - `trtllm -> TrtllmAlltoall, else -> CustomAllReduce` mapping is correct for - the standard (non-WideEP) MoE path that the gated parity harness - (`test_engine_step_parity.py`) and `cli_estimate` exercise. It does NOT emit - the SGLang-WideEP `DeepEpNormal` / `DeepEpLowLatency` flavors — that DeepEP - path is out-of-scope scan territory, not gated by the smoke harness. A future - WideEP stage must thread the WideEP/DeepEP flavor selection through here - (and likely needs the resolved dispatch kind off the Python op rather than a - pure backend-string mapping). +def _dispatch_flavor(backend: str, op: MoEDispatch) -> str: + """Resolve the Rust `DispatchFlavor` variant for a dispatch op. The enum + has no serde rename, so emit the exact variant name: + + - trtllm -> `TrtllmAlltoall` (the fine SM/NVL72 gating stays inside the + Rust `MoEDispatchOp::query`); + - sglang with `moe_backend == "deepep_moe"` -> the WideEP DeepEP tables: + context ops use DeepEP-normal (high-throughput), decode ops use + DeepEP-low-latency — mirroring the Python branch split + (`operations/moe.py`: `if self._is_context: query_wideep_deepep_normal + else query_wideep_deepep_ll`); + - everything else -> `CustomAllReduce`. """ - return "TrtllmAlltoall" if backend == "trtllm" else "CustomAllReduce" + if backend == "trtllm": + return "TrtllmAlltoall" + if backend == "sglang" and getattr(op, "_moe_backend", None) == "deepep_moe": + return "DeepEpNormal" if op._is_context else "DeepEpLowLatency" + return "CustomAllReduce" def _moe_dispatch(op: MoEDispatch, *, backend: str) -> dict: @@ -290,7 +304,10 @@ def _moe_dispatch(op: MoEDispatch, *, backend: str) -> dict: "attention_dp_size": op._attention_dp_size, "pre_dispatch": op._pre_dispatch, "backend": backend, - "flavor": _dispatch_flavor(backend), + "flavor": _dispatch_flavor(backend, op), + # DeepEP branches divide the dispatch token count by this (Python + # `num_tokens // self._scale_num_tokens`, moe.py sglang DeepEP path). + "scale_num_tokens": op._scale_num_tokens, "comm_quant": "half", "moe_quant": _quant_name(quant) if quant is not None else "bfloat16", "attn_cp_size": op._attn_cp_size, @@ -318,7 +335,12 @@ def _nccl(op: NCCL) -> dict: "scale_factor": op._scale_factor, "hidden_size": op._num_elements_per_token, "num_gpus": op._num_gpus, - "dtype": "half", + # Pass the op's real comm dtype through (every current model builder + # passes half, but the NCCL op supports int8/fp8 and the Rust enum + # carries them). CustomAllReduce / MoEDispatch stay "half" by + # construction — Python hardcodes CommQuantMode.half at their query + # sites, so "half" IS the parity value there. + "dtype": op._comm_quant_mode.name, "operation": op._nccl_op, "seq_split": op._seq_split, } @@ -357,6 +379,10 @@ def _dsa_module(op: ContextDSAModule | GenerationDSAModule, *, architecture: str "gemm_quant_mode": _quant_name(op._gemm_quant_mode), "architecture": arch, "index_topk": int(dims["index_topk"]), + # GLM-5.2 shared-index amortization: per-layer cost is + # `full_frac*full + (1-full_frac)*skip` (see `ContextDSAModule.query`). + # 1.0 (DeepSeek-V3.2 / GLM-5) keeps the pure-full path. + "full_frac": float(getattr(op, "_full_frac", 1.0)), } @@ -460,7 +486,10 @@ def _wideep_context_mla(op: WideEPContextMLA) -> dict: return { "name": op._name, "scale_factor": op._scale_factor, - "num_heads": op._tp_size, # Rust num_heads slot carries the per-rank head split + # The op stores tp_size; the Rust table axis is per-rank heads. Mirror + # the Python query's conversion (mla.py: `num_heads = 128 // tp_size`, + # DeepSeek's 128 total heads). + "num_heads": 128 // op._tp_size, "kv_cache_dtype": _quant_name(op._kvcache_quant_mode), "fmha_quant_mode": _quant_name(op._fmha_quant_mode), "attn_backend": op._attn_backend, @@ -472,7 +501,7 @@ def _wideep_generation_mla(op: WideEPGenerationMLA) -> dict: return { "name": op._name, "scale_factor": op._scale_factor, - "num_heads": op._tp_size, + "num_heads": 128 // op._tp_size, "kv_cache_dtype": _quant_name(op._kvcache_quant_mode), "fmha_quant_mode": _quant_name(op._fmha_quant_mode), "attn_backend": op._attn_backend, @@ -506,6 +535,26 @@ def _wideep_moe(op: TrtLLMWideEPMoE, *, database: Any) -> dict: } +def _wideep_moe_dispatch(op: TrtLLMWideEPMoEDispatch) -> dict: + """TRT-LLM WideEP All2All dispatch (Python `TrtLLMWideEPMoEDispatch`). + Field names match the Rust `TrtllmWideEpMoEDispatchOp`. `node_num` is + intentionally NOT on the wire: the model builders never set it and the + Rust query applies Python's default (`1 if ep < 4 else ep // 4`).""" + return { + "name": op._name, + "scale_factor": op._scale_factor, + "hidden_size": op._hidden_size, + "topk": op._topk, + "num_experts": op._num_experts, + "moe_tp_size": op._moe_tp_size, + "moe_ep_size": op._moe_ep_size, + "attention_dp_size": op._attention_dp_size, + "pre_dispatch": op._pre_dispatch, + "quant_mode": _quant_name(op._quant_mode), + "use_low_precision_combine": bool(op._use_low_precision_combine), + } + + def _to_opspec(op: Any, *, backend: str, architecture: str, database: Any) -> dict: """Convert one Python ``Operation`` to its externally-tagged ``OpSpec`` dict. @@ -557,6 +606,10 @@ def recurse(child: Any) -> dict: return {"WideEpGenerationMla": _wideep_generation_mla(op)} if isinstance(op, TrtLLMWideEPMoE): return {"WideEpMoe": _wideep_moe(op, database=database)} + # MUST precede the MoEDispatch check: TrtLLMWideEPMoEDispatch is a direct + # `Operation` subclass (not a MoEDispatch), but keep the guard explicit. + if isinstance(op, TrtLLMWideEPMoEDispatch): + return {"WideEpMoeDispatch": _wideep_moe_dispatch(op)} # Plain (single-variant) ops. if isinstance(op, GEMM): @@ -633,6 +686,16 @@ def _compute_perf_db_sources(database: Any) -> dict: ] return out except Exception: + logger.warning( + "Failed to resolve shared-layer perf-DB sources for %s/%s/%s; the " + "compiled engine will load PRIMARY-ONLY rows while Python uses the " + "shared layer in the same run — cross-engine drift is possible. " + "Investigate rather than ignore.", + getattr(database, "system", "?"), + getattr(database, "backend", "?"), + getattr(database, "version", "?"), + exc_info=True, + ) return {} @@ -904,8 +967,33 @@ def predict_prefill_latency(self, bs: int, isl: int, prefix: int = 0) -> float: def predict_decode_latency(self, bs: int, isl: int, osl: int = 2) -> float: return self._engine.predict_decode_latency(int(bs), int(isl), int(osl)) - def mixed_step_latency(self, ctx_tokens: int, gen_tokens: int, isl: int, osl: int, prefix: int = 0) -> float: - return self._engine.mixed_step_latency(int(ctx_tokens), int(gen_tokens), int(isl), int(osl), int(prefix)) + def mixed_step_latency( + self, + ctx_tokens: int, + gen_tokens: int, + isl: int, + osl: int, + prefix: int = 0, + seq_imbalance_correction_scale: float = 1.0, + gen_seq_imbalance_correction_scale: float = 1.0, + ) -> float: + return self._engine.mixed_step_latency( + int(ctx_tokens), + int(gen_tokens), + int(isl), + int(osl), + int(prefix), + float(seq_imbalance_correction_scale), + float(gen_seq_imbalance_correction_scale), + ) - def decode_step_latency(self, gen_tokens: int, isl: int, osl: int) -> float: - return self._engine.decode_step_latency(int(gen_tokens), int(isl), int(osl)) + def decode_step_latency( + self, + gen_tokens: int, + isl: int, + osl: int, + gen_seq_imbalance_correction_scale: float = 1.0, + ) -> float: + return self._engine.decode_step_latency( + int(gen_tokens), int(isl), int(osl), float(gen_seq_imbalance_correction_scale) + ) diff --git a/src/aiconfigurator/sdk/rust_engine_step.py b/src/aiconfigurator/sdk/rust_engine_step.py index 13fcd600c..d1a85d535 100644 --- a/src/aiconfigurator/sdk/rust_engine_step.py +++ b/src/aiconfigurator/sdk/rust_engine_step.py @@ -26,6 +26,14 @@ ENGINE_STEP_BACKEND_ENV = "AICONFIGURATOR_ENGINE_STEP_BACKEND" +class RustEngineUnsupportedError(RuntimeError): + """The model's op graph cannot be expressed as a compiled ``EngineSpec`` + (``engine.OpConversionError``). Python CAN compute these configs, so the + ``base_backend`` gates catch this and fall back to the Python step + (parity by delegation) instead of crashing the sweep. Distinct from + perf-data misses, which must stay error-symmetric on both engines.""" + + class RustForwardPassPerfModel: """Facade over the compiled Rust forward-pass perf model (PR #1152). @@ -298,16 +306,17 @@ def estimate_mixed_step_latency_with_rust( isl: int, osl: int, prefix: int, + seq_imbalance_correction_scale: float = 1.0, + gen_seq_imbalance_correction_scale: float = 1.0, ) -> float: """Estimate one mixed prefill/decode engine step through the compiled engine. Delegates to ``EngineHandle.mixed_step_latency``. The Rust - ``Engine::mixed_step_latency`` (``engine/runtime.rs:280``) reproduces the - full FPM packing the old ctypes bridge did inline — the - ``ceil(ctx_tokens / isl)`` prefill-request count, the cached-prefix - subtraction, the ``(nextn + 1)`` decode multiplier, and the kv-token - packing — so the raw step args pass straight through with no Python-side - pre-math. + ``Engine::mixed_step_latency`` is a literal mirror of Python's + ``_get_mix_step_latency`` three-pass composition (combined non-attention, + context attention / ceil(isl/ctx), decode attention with the ``(nextn+1)`` + batch), so the raw step args plus the runtime imbalance scales pass + straight through with no Python-side pre-math. """ handle = _cached_engine_handle(model, database) return handle.mixed_step_latency( @@ -316,6 +325,8 @@ def estimate_mixed_step_latency_with_rust( int(isl), int(osl), int(prefix or 0), + seq_imbalance_correction_scale=float(seq_imbalance_correction_scale or 1.0), + gen_seq_imbalance_correction_scale=float(gen_seq_imbalance_correction_scale or 1.0), ) @@ -326,16 +337,23 @@ def estimate_decode_step_latency_with_rust( gen_tokens: int, isl: int, osl: int, + gen_seq_imbalance_correction_scale: float = 1.0, ) -> float: """Estimate one decode-only engine step through the compiled engine. Delegates to ``EngineHandle.decode_step_latency``. The Rust - ``Engine::decode_step_latency`` (``engine/runtime.rs:342``) applies the - ``(nextn + 1)`` decode-batch scaling and the ``s = isl + osl/2`` sequence - length internally, so the raw args pass straight through. + ``Engine::decode_step_latency`` mirrors Python's + ``_get_genonly_step_latency``: one step over the full generation op list + at ``s = isl + osl//2 + 1`` with the ``(nextn + 1)`` decode-batch scaling + applied internally, so the raw args pass straight through. """ handle = _cached_engine_handle(model, database) - return handle.decode_step_latency(int(gen_tokens), int(isl), int(osl)) + return handle.decode_step_latency( + int(gen_tokens), + int(isl), + int(osl), + gen_seq_imbalance_correction_scale=float(gen_seq_imbalance_correction_scale or 1.0), + ) # Memo of compiled ``EngineHandle`` objects, keyed by the engine identity @@ -364,8 +382,26 @@ def _cached_engine_handle(model: Any, database: Any) -> Any: ``AICONFIGURATOR_SYSTEMS_PATH`` so it resolves to the same systems tree the Python ``database`` came from. """ - key = _engine_config_json(model, database) + # The identity JSON is a hot-path cost: the engine-step helpers call this + # per step and `_engine_config_json` runs ~2-3us of getattr + json.dumps + # (which the perf regression gate measures against a ~20us step). The + # identity is immutable for a given (model, database) pair, so memoize the + # computed key on the model object and only recompute when the database + # object changes. + memo = getattr(model, "_aic_engine_identity_memo", None) + if memo is not None and memo[0] is database: + key = memo[1] + else: + key = _engine_config_json(model, database) + try: + model._aic_engine_identity_memo = (database, key) + except (AttributeError, TypeError): + pass # slotted/frozen model objects: recompute per call handle = _ENGINE_HANDLE_CACHE.get(key) + if isinstance(handle, RustEngineUnsupportedError): + # Compilation already failed for this engine identity; re-raise the + # cached error instead of re-walking the op graph every step. + raise handle if handle is not None: return handle @@ -374,22 +410,27 @@ def _cached_engine_handle(model: Any, database: Any) -> Any: # (``_quant_to_dtype`` / ``_moe_quant_to_dtype``), so a top-level import # here would be a circular import. import aiconfigurator_core - from aiconfigurator.sdk.engine import EngineHandle, build_engine_spec_json + from aiconfigurator.sdk.engine import EngineHandle, OpConversionError, build_engine_spec_json systems_path = os.environ.get("AICONFIGURATOR_SYSTEMS_PATH") nextn = getattr(model, "_nextn", None) - spec_json = build_engine_spec_json( - model, - model_path=getattr(model, "model_path", getattr(model, "model_name", "")), - system=database.system, - backend=_backend_name(database.backend), - backend_version=getattr(database, "version", None), - kv_block_size=None, - systems_path=systems_path, - nextn=int(nextn) if nextn is not None else 0, - nextn_accept_rates=getattr(model, "_nextn_accept_rates", None), - database=database, - ) + try: + spec_json = build_engine_spec_json( + model, + model_path=getattr(model, "model_path", getattr(model, "model_name", "")), + system=database.system, + backend=_backend_name(database.backend), + backend_version=getattr(database, "version", None), + kv_block_size=None, + systems_path=systems_path, + nextn=int(nextn) if nextn is not None else 0, + nextn_accept_rates=getattr(model, "_nextn_accept_rates", None), + database=database, + ) + except OpConversionError as exc: + unsupported = RustEngineUnsupportedError(str(exc)) + _ENGINE_HANDLE_CACHE[key] = unsupported + raise unsupported from exc spec_bytes = bytes(aiconfigurator_core.engine_spec_bincode_from_json(spec_json)) handle = EngineHandle(spec_bytes, systems_path=systems_path) _ENGINE_HANDLE_CACHE[key] = handle @@ -427,11 +468,57 @@ def _engine_config_json(model: Any, database: Any) -> str: "kv_block_size": None, "nextn": int(nextn) if nextn is not None else None, "nextn_accept_rates": ([float(r) for r in nextn_accept_rates] if nextn_accept_rates is not None else None), - "extra": {}, + # Cache-identity widening. The dtype fields above collapse distinct + # quant modes onto one wire string (`sq`/`int8_wo` -> "int8", + # `fp8_ootb` -> "fp8", four 4-bit modes -> "int4", the DSv4 MoE modes + # -> None), and several ModelConfig fields shape the compiled op list + # without appearing in the identity at all. Two models differing only + # in those would otherwise share one cached handle and silently return + # each other's latencies. `extra` participates in the JSON key, so + # carrying the RAW enum names + the op-shaping fields here + # disambiguates the memo without touching the wire schema. + # Rust `EngineConfig.extra` is `BTreeMap`, so the + # identity payload is one JSON-encoded STRING value (deserializable + # if this dict ever crosses the wire, and a stable cache key today). + "extra": { + "identity": json.dumps( + { + "raw_quant_modes": { + "gemm": _raw_quant_name(getattr(model_config, "gemm_quant_mode", None)), + "moe": _raw_quant_name(getattr(model_config, "moe_quant_mode", None)), + "fmha": _raw_quant_name(getattr(model_config, "fmha_quant_mode", None)), + "kvcache": _raw_quant_name(getattr(model_config, "kvcache_quant_mode", None)), + "comm": _raw_quant_name(getattr(model_config, "comm_quant_mode", None)), + }, + "model_config": { + "cp_style": getattr(model_config, "cp_style", None), + "workload_distribution": getattr(model_config, "workload_distribution", None), + "overwrite_num_layers": getattr(model_config, "overwrite_num_layers", None), + "sms": getattr(model_config, "sms", None), + "moe_backend": getattr(model_config, "moe_backend", None), + "attention_backend": getattr(model_config, "attention_backend", None), + "enable_wideep": bool(getattr(model_config, "enable_wideep", False)), + "enable_eplb": bool(getattr(model_config, "enable_eplb", False)), + "wideep_num_slots": getattr(model_config, "wideep_num_slots", None), + }, + }, + sort_keys=True, + separators=(",", ":"), + ), + }, } return json.dumps(config, sort_keys=True, separators=(",", ":")) +def _raw_quant_name(value: Any) -> str | None: + """Un-collapsed quant identity for the handle-cache key: the Python enum + member name (e.g. ``sq``, ``fp8_ootb``, ``w4afp8``) rather than the lossy + wire ``DataType`` string.""" + if value is None: + return None + return getattr(value, "name", str(value)) + + def _backend_name(value: Any) -> str: if hasattr(value, "value"): return str(value.value) diff --git a/tests/unit/sdk/test_rust_engine_step.py b/tests/unit/sdk/test_rust_engine_step.py index 01c7dd347..880bd5dfa 100644 --- a/tests/unit/sdk/test_rust_engine_step.py +++ b/tests/unit/sdk/test_rust_engine_step.py @@ -94,12 +94,12 @@ def test_mixed_and_decode_helpers_pass_raw_step_args(monkeypatch) -> None: decode_calls = [] class _FakeHandle: - def mixed_step_latency(self, *args): - mixed_calls.append(args) + def mixed_step_latency(self, *args, **kwargs): + mixed_calls.append((args, kwargs)) return 8.5 - def decode_step_latency(self, *args): - decode_calls.append(args) + def decode_step_latency(self, *args, **kwargs): + decode_calls.append((args, kwargs)) return 9.5 monkeypatch.setattr(rust_engine_step, "_cached_engine_handle", lambda model, database: _FakeHandle()) @@ -126,8 +126,15 @@ def decode_step_latency(self, *args): assert mixed_ms == 8.5 assert decode_ms == 9.5 - assert mixed_calls == [(384, 7, 256, 256, 128)] - assert decode_calls == [(7, 256, 256)] + # Raw step args pass through positionally; the runtime imbalance scales + # ride as kwargs (default 1.0 when the caller doesn't set them). + assert mixed_calls == [ + ( + (384, 7, 256, 256, 128), + {"seq_imbalance_correction_scale": 1.0, "gen_seq_imbalance_correction_scale": 1.0}, + ) + ] + assert decode_calls == [((7, 256, 256), {"gen_seq_imbalance_correction_scale": 1.0})] def test_engine_config_json_preserves_moe_specific_quant_mode() -> None: @@ -367,6 +374,103 @@ class _Dsv4Op: assert spec["window_size"] == 2048 +def test_engine_config_json_identity_disambiguates_collapsed_quant_modes(): + """Two models differing only in a wire-collapsed dtype (sq vs int8_wo both + -> "int8") or an identity-omitted ModelConfig field (moe_backend) must get + DISTINCT handle-cache keys — sharing one cached handle silently returns + the other model's latencies.""" + from aiconfigurator.sdk import common + + def _model(gemm_mode, moe_backend=None): + cfg = SimpleNamespace( + tp_size=8, + pp_size=1, + moe_tp_size=1, + moe_ep_size=8, + attention_dp_size=1, + cp_size=None, + gemm_quant_mode=gemm_mode, + moe_quant_mode=None, + fmha_quant_mode=None, + kvcache_quant_mode=None, + comm_quant_mode=None, + moe_backend=moe_backend, + attention_backend=None, + enable_wideep=False, + enable_eplb=False, + wideep_num_slots=None, + cp_style=None, + workload_distribution=None, + overwrite_num_layers=None, + sms=None, + ) + return SimpleNamespace(model_path="test/model", architecture=None, config=cfg, _nextn=None) + + database = SimpleNamespace(system="test_sxm", backend="vllm", version="1.0.0") + key_sq = rust_engine_step._engine_config_json(_model(common.GEMMQuantMode.sq), database) + key_int8 = rust_engine_step._engine_config_json(_model(common.GEMMQuantMode.int8_wo), database) + assert key_sq != key_int8, "sq and int8_wo must not alias one cached handle" + + key_deepep = rust_engine_step._engine_config_json( + _model(common.GEMMQuantMode.sq, moe_backend="deepep_moe"), database + ) + assert key_sq != key_deepep, "moe_backend must participate in the cache identity" + + +def test_op_conversion_error_falls_back_to_python_step(monkeypatch): + """An OpConversionError (op graph not expressible in Rust) must be + surfaced as RustEngineUnsupportedError, cached per engine identity, and + caught by the base_backend gates (fallback to the Python step) — NOT + crash the sweep.""" + import pytest + + from aiconfigurator.sdk.engine import OpConversionError + from aiconfigurator.sdk.rust_engine_step import RustEngineUnsupportedError + + calls = {"n": 0} + + def _raise_conversion(*args, **kwargs): + calls["n"] += 1 + raise OpConversionError("unsupported op: ContextMSAModule") + + monkeypatch.setattr("aiconfigurator.sdk.engine.build_engine_spec_json", _raise_conversion) + rust_engine_step._engine_handle_cache_clear() + + model = _dense_model() + database = SimpleNamespace(system="test_sxm", backend="vllm", version="1.0.0") + + with pytest.raises(RustEngineUnsupportedError): + rust_engine_step._cached_engine_handle(model, database) + # Second call re-raises from the cache without recompiling. + with pytest.raises(RustEngineUnsupportedError): + rust_engine_step._cached_engine_handle(model, database) + assert calls["n"] == 1, "compile failure must be memoized per engine identity" + rust_engine_step._engine_handle_cache_clear() + + +def test_wideep_mla_spec_emits_per_rank_heads_not_tp(): + """The WideEP MLA table axis is per-rank heads; the Python query converts + ``num_heads = 128 // tp_size`` (mla.py). The spec emitter must apply the + same conversion — emitting raw tp_size makes Rust query the wrong table + slice (tp=8 would read the heads=8 extrapolation instead of heads=16).""" + from aiconfigurator.sdk import common + from aiconfigurator.sdk.engine import _wideep_context_mla, _wideep_generation_mla + + class _WideEpOp: + _name = "context_attention" + _scale_factor = 1.0 + _tp_size = 8 + _cp_size = 1 + _kvcache_quant_mode = common.KVCacheQuantMode.fp8 + _fmha_quant_mode = common.FMHAQuantMode.fp8_block + _attn_backend = "flashinfer" + + ctx_spec = _wideep_context_mla(_WideEpOp()) + gen_spec = _wideep_generation_mla(_WideEpOp()) + assert ctx_spec["num_heads"] == 16 # 128 // 8, NOT tp_size=8 + assert gen_spec["num_heads"] == 16 + + def test_non_silicon_database_mode_falls_back_to_python_step(): """The compiled engine is SILICON-only (no util_empirical layer); HYBRID / EMPIRICAL databases must stay on the Python step so both backends give diff --git a/uv.lock b/uv.lock index efba61bd8..cb18cdb4a 100644 --- a/uv.lock +++ b/uv.lock @@ -46,8 +46,6 @@ dependencies = [ { name = "pyarrow" }, { name = "pydantic" }, { name = "pyyaml" }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "tqdm" }, ] @@ -122,7 +120,6 @@ requires-dist = [ { name = "pytest-xdist", marker = "extra == 'dev'", specifier = "==3.8.0" }, { name = "pyyaml", specifier = ">=6.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = "==0.14.1" }, - { name = "scipy", specifier = ">=1.13.1" }, { name = "tensorflow-probability", extras = ["jax"], marker = "python_full_version < '3.13' and extra == 'spica'", specifier = "==0.25.0" }, { name = "tqdm", specifier = ">=4.0.0" }, { name = "uv", marker = "extra == 'dev'" }, @@ -1239,14 +1236,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f5/a1/1f7f0c555f5858fd2906fe9f7b0a3554fddb85cb70df7a6aaec41dc292c2/greenlet-3.5.3-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:c180d22d325fb613956b443c3c6f4406eb70e6defc70d3974da2a7b59e06f48c", size = 285838, upload-time = "2026-06-26T18:21:05.167Z" }, { url = "https://files.pythonhosted.org/packages/0a/29/be9f43ed61677a5759b38c8a9389248133c8c731bbfc0574ecdff66c99fc/greenlet-3.5.3-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:483d08c11181c83a6ce1a7a61df0f624a208ec40817a3bb2302714592eee4f04", size = 602342, upload-time = "2026-06-26T19:07:06.908Z" }, { url = "https://files.pythonhosted.org/packages/b9/42/ba41c97ec36aa4b3ec25e5aa691d79561254805fad7f2f826dd6770587e2/greenlet-3.5.3-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1dae6e0091eae084317e411f047f0b7cb241c6db570f7c45fd6b900a274914ce", size = 615541, upload-time = "2026-06-26T19:10:04.909Z" }, + { url = "https://files.pythonhosted.org/packages/2e/8c/231ca675b0df779816950ca66b40b1fa14dbff4a0ed9814a9a29ec399140/greenlet-3.5.3-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0f6ff50ff8dbd51fae9b37f4101648b04ea0df19b3f50ab2beb5061e7716a5c8", size = 622473, upload-time = "2026-06-26T19:24:12.786Z" }, { url = "https://files.pythonhosted.org/packages/f5/c7/28747042e1df8a9cd120a1ebe15529fc4be3b486e13e8d551ff307a82412/greenlet-3.5.3-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9bcd2d72ccd70a1ec68ba6ef93e7fbb4420ef9997dabc7010d893bd4015e0bec", size = 615675, upload-time = "2026-06-26T18:32:14.444Z" }, + { url = "https://files.pythonhosted.org/packages/81/fe/dd97c483a3ff82849196ccd07851600edd3ac9de74669ca8a6022ada9ea1/greenlet-3.5.3-cp310-cp310-manylinux_2_39_riscv64.whl", hash = "sha256:37bf9c538f5ae6e63d643f88dec37c0c83bdf0e2ebc62961dedcf458822f7b71", size = 418421, upload-time = "2026-06-26T19:25:34.503Z" }, { url = "https://files.pythonhosted.org/packages/cc/a8/b85525a6c8fba9f009a5f7c8df1545de8fb0f0bf3e0179194ef4e500317f/greenlet-3.5.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:73f152c895e09907e0dbe24f6c2db37beb085cd63db91c3825a0fcd0064124a8", size = 1575057, upload-time = "2026-06-26T19:09:00.264Z" }, { url = "https://files.pythonhosted.org/packages/03/79/fb76edb218fe6735ab0edeba176c7ab80df9618f7c02ce4208979f3ae7db/greenlet-3.5.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8bdb43e1a1d1873721acab2be99c5befd4d2044ddfd52e4d610801019880a702", size = 1641692, upload-time = "2026-06-26T18:31:41.454Z" }, { url = "https://files.pythonhosted.org/packages/6b/79/86fe3ee50ed55d9b3907eecd3208b5c3fe8a79515519aae98b4753c3fa1d/greenlet-3.5.3-cp310-cp310-win_amd64.whl", hash = "sha256:0909f9355a9f24845d3299f3112e266a06afb68302041989fd26bd68894933db", size = 238742, upload-time = "2026-06-26T18:20:40.758Z" }, { url = "https://files.pythonhosted.org/packages/51/58/5404031044f55afad7aad1aff8be3f22b1bed03e237cfeabbc7e5c8cfde0/greenlet-3.5.3-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:aca9b4ce85b152b5524ef7d88170efdff80dc0032aa8b75f9aaf7f3479ea95b4", size = 287424, upload-time = "2026-06-26T18:20:31.469Z" }, { url = "https://files.pythonhosted.org/packages/b4/bf/1c65e9b94a54d547068fa5b5a8a06f221f3316b48908e08668d29c77cb50/greenlet-3.5.3-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0f71be4920368fe1fabeeaa53d1e3548337e2b223d9565f8ad5e392a75ba23fc", size = 606523, upload-time = "2026-06-26T19:07:08.859Z" }, { url = "https://files.pythonhosted.org/packages/b8/c7/b66baacc95775ad511287acb0137b95574a9ce5491902372b7564799d790/greenlet-3.5.3-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4d77e67f65f98449e3fb83f795b5d0a8437aead2f874ca89c96576caf4be3af6", size = 618315, upload-time = "2026-06-26T19:10:06.055Z" }, + { url = "https://files.pythonhosted.org/packages/b0/a0/68afd1ebad40db87dac0a28ffa120726b98bf9c7c40c481b0f63c105d298/greenlet-3.5.3-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e18619ba655ac05d78d80fc83cac4ba892bd6927b99e3b8237aee861aaacc8bb", size = 626155, upload-time = "2026-06-26T19:24:14.44Z" }, { url = "https://files.pythonhosted.org/packages/78/2b/28ed29463522fdbe4c15b1f63922041626a7478316b34ab4adda3f0a4aba/greenlet-3.5.3-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8540f1e6205bd13ca0ce685581037219ca54a1b41a0a15d228c6c9b8ad5903d7", size = 617381, upload-time = "2026-06-26T18:32:16.077Z" }, + { url = "https://files.pythonhosted.org/packages/07/7f/e327d912239ec4b3b49999e3967389bcf1ee8722b9ee9194d2752ecd558a/greenlet-3.5.3-cp311-cp311-manylinux_2_39_riscv64.whl", hash = "sha256:d27c0c653a60d9535f690226474a5cc1036a8b0d7b57504d1c4f89c44a07a80c", size = 421083, upload-time = "2026-06-26T19:25:35.804Z" }, { url = "https://files.pythonhosted.org/packages/2a/7b/ad04e9d1337fc04965dc9fc616b6a72cb65a24b800a014c011ec812f5489/greenlet-3.5.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:7ef56fe650f50575bf843acde967b9c567687f3c22340941a899b7bc56e956a8", size = 1577771, upload-time = "2026-06-26T19:09:01.537Z" }, { url = "https://files.pythonhosted.org/packages/d8/33/6c87ab7ba663f70ca21f3022aad1ffe56d3f3e0521e836c2415e13abcc3c/greenlet-3.5.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:5121af01cf911e70056c00d4b46d5e9b5d1415550038573d744138bacb59e6b8", size = 1644048, upload-time = "2026-06-26T18:31:42.996Z" }, { url = "https://files.pythonhosted.org/packages/1c/35/f0d8ee998b422cf8693b270f098e55d8d4ec8006b061b333f54f177d28d9/greenlet-3.5.3-cp311-cp311-win_amd64.whl", hash = "sha256:0f41e4a05a3c0cb31b17023eff28dd111e1d16bf7d7d00406cd7df23f31398a7", size = 239137, upload-time = "2026-06-26T18:23:21.664Z" }, @@ -1254,7 +1255,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/6e/4c37d51a2b7f82d2ff11bb6b5f7d766d9a011726624af255e843727627a3/greenlet-3.5.3-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:719757059f5a53fd0dde23f78cffeafcdd97b21c850ddb7ca684a3c1a1f122e2", size = 288685, upload-time = "2026-06-26T18:22:08.977Z" }, { url = "https://files.pythonhosted.org/packages/7a/73/815dd90131c1b71ebdf53dbc7c276cafec2a1173b97559f97aba72724a87/greenlet-3.5.3-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:efa9f765dd09f9d0cdac651ffdf631ee59ec5dc6ee7a73e0c012ba9c52fbdf5b", size = 604761, upload-time = "2026-06-26T19:07:10.114Z" }, { url = "https://files.pythonhosted.org/packages/9f/57/079cfe76bcef36b153b25607ee91c6fcb58f17f8b23c86bbbeabe0c88d72/greenlet-3.5.3-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7faba15ac005376e02a0384504e0243be3370ce010296a44a820feb342b505ab", size = 617044, upload-time = "2026-06-26T19:10:07.25Z" }, + { url = "https://files.pythonhosted.org/packages/fb/fb/d97dc261209c80744b7c8132693a30d70ec6e7315e632cb0a10b3fec94dd/greenlet-3.5.3-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5795cd1101371140551c645f2d408b8d3c01a5a29cf8a9bce6e759c983682d23", size = 622351, upload-time = "2026-06-26T19:24:16.32Z" }, { url = "https://files.pythonhosted.org/packages/37/87/b4d095775a3fb1bcafbb483fc206b27ebb785724c83051447737085dc54e/greenlet-3.5.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:87142215824be6ac05e2e8e2786eec307ccbc27c36723c3881959df654af6861", size = 614244, upload-time = "2026-06-26T18:32:17.594Z" }, + { url = "https://files.pythonhosted.org/packages/8e/ac/e5fee13cbbd0e8de312d9a146584b8a51891c68847330ef9dc8b5109d23f/greenlet-3.5.3-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:af4923b3096e26a36d7e9cf24ab88083a20f97d191e3b97f253731ce9b41b28c", size = 425395, upload-time = "2026-06-26T19:25:37.144Z" }, { url = "https://files.pythonhosted.org/packages/8a/70/7559b609683650fa2b95b8ab84b4ab0b26556a635d19675e12aa832d826d/greenlet-3.5.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:215275b1b49320987352e6c1b054acca0064f965a2c66992bed9a6f7d913f149", size = 1574210, upload-time = "2026-06-26T19:09:03.077Z" }, { url = "https://files.pythonhosted.org/packages/ae/73/be55392074c60fc37655ca40fa6022457bfbf6718e9e342a7b0b41f96dd2/greenlet-3.5.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6b1b0eed82364b0e32c4ea0f221452d33e6bb17ae094d9f72aed9851812747ea", size = 1638627, upload-time = "2026-06-26T18:31:44.748Z" }, { url = "https://files.pythonhosted.org/packages/14/40/c57489acf8e37d74e2913d4eff63aa0dba17acccc4bdeef874dde2dbbec9/greenlet-3.5.3-cp312-cp312-win_amd64.whl", hash = "sha256:cde8adafa2365676f74a979744629589999093bc86e2484214f58e61df08902c", size = 239882, upload-time = "2026-06-26T18:23:27.518Z" }, @@ -1262,7 +1265,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9b/ff/a620267401db30a50cc8450ee90730e2d4a85658c055c0e760d4ed47fb13/greenlet-3.5.3-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:c8d87c2134d871df96ecdea9cec7cbaab286dadab0f56476e57aaf9e8ac11550", size = 287609, upload-time = "2026-06-26T18:21:14.724Z" }, { url = "https://files.pythonhosted.org/packages/d6/fa/5401ac78021c826a25b6dde0c705e0a8f29b617509f9185a31dac15fbe1b/greenlet-3.5.3-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a2d185dd1621757e70c3861cceffd5317ab4e7ed7eb09c82994828468527ade5", size = 607435, upload-time = "2026-06-26T19:07:11.412Z" }, { url = "https://files.pythonhosted.org/packages/e9/76/1dc144a2e56e65d36405078ed774224375ea520a1870a6e46e08bb4ac7bf/greenlet-3.5.3-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1c514a468149bf8fbbab874188a3535cd8a48a3e353eb53a3d424296f8dbacd3", size = 619787, upload-time = "2026-06-26T19:10:08.396Z" }, + { url = "https://files.pythonhosted.org/packages/57/61/2f5b1adf256d039f5dab8005de8d3d7ad2b0070a3219c0e036b3fbfeb440/greenlet-3.5.3-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9ad04dd75458c6300b047c61b8639092433d205a25a14e310d6582a480efcca1", size = 625580, upload-time = "2026-06-26T19:24:18.344Z" }, { url = "https://files.pythonhosted.org/packages/bf/87/c298cee62df1de4ad7fec32abda73526cff347fd143a6ed4ac369246668a/greenlet-3.5.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:915f887cf2682b66419b879423a2e072634aa7b7dce6f3ada4957cfced3f1e9a", size = 616786, upload-time = "2026-06-26T18:32:19.128Z" }, + { url = "https://files.pythonhosted.org/packages/3e/d9/ab7fc9e543e44d6879b0a6ef9a4b2188940fd180cc65d6f646883ddf7201/greenlet-3.5.3-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:afaabdd554cd7ae9bbb3ca070b0d7fdfd207dbf1d16865f7233837709d354bda", size = 427933, upload-time = "2026-06-26T19:25:38.219Z" }, { url = "https://files.pythonhosted.org/packages/9e/2e/e6f009885ed0705ccf33fe0583c117cfd03cde77e31a596dd5785a30762b/greenlet-3.5.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:766cfd421c13e450feb340cd472a3ed9957d438727b7b4593ad7c76c5d2b0deb", size = 1574316, upload-time = "2026-06-26T19:09:04.273Z" }, { url = "https://files.pythonhosted.org/packages/ef/fe/43fd110b01e40da0adb7c90ac7ea744bef2d43dca00de5095fd2351c2a68/greenlet-3.5.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:2ecda9ec22edf38fa389369eaed8c3d37c05f3c54e69f69438dbb2cc1de1458b", size = 1638614, upload-time = "2026-06-26T18:31:46.297Z" }, { url = "https://files.pythonhosted.org/packages/0f/7c/062447147a61f8b4337b156fe70d32a165fcf2f89d7ca6255e572806705c/greenlet-3.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:c82304750f057167ff60d188df1d0cc1764ce9567eadf03e6a7443bcedd0b30b", size = 239850, upload-time = "2026-06-26T18:21:54.613Z" }, @@ -1270,7 +1275,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c3/93/43e116ee114b28737ba7e12952a0d4e2f55944d0f84e42bc91ba7192a3c9/greenlet-3.5.3-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:fd2e02fa07485778536a036222d616ab957b1d533f36b3ed98ce725d9c9d3117", size = 288202, upload-time = "2026-06-26T18:23:49.604Z" }, { url = "https://files.pythonhosted.org/packages/82/2f/146d218299046a43d1f029fd544b3d110d0f175a09c715c7e8da4a4a345d/greenlet-3.5.3-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:df0a0628d1597eb0897b62f55d1343f772405fd25f3b2a796c76874b0c2e22e8", size = 654096, upload-time = "2026-06-26T19:07:12.71Z" }, { url = "https://files.pythonhosted.org/packages/a0/cc/04738cafb3f45fa991ea44f9de94c47dcec964f5a972300988a6751f49d9/greenlet-3.5.3-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ebd933a6adabc298bab47731a130fe6bfb888bd934eee37810f151159544540d", size = 666304, upload-time = "2026-06-26T19:10:09.503Z" }, + { url = "https://files.pythonhosted.org/packages/86/a9/73fa62893d5b84b4205544e6b673c654cc43aa5b9899bac00f04d64af73d/greenlet-3.5.3-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8d19fe6c39ebff9259f07bcc685d3290f8fa4ea2278e51dd0008e4d6b0f2d814", size = 670657, upload-time = "2026-06-26T19:24:19.967Z" }, { url = "https://files.pythonhosted.org/packages/ce/aa/4e0dad5e605c270c784ab911c43da6adb136ccd4d81180f763ca429a723d/greenlet-3.5.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4b9d501b40e80b70e32323c799dd9b420a5577a9601469d362ae1ffb690f3a7c", size = 663635, upload-time = "2026-06-26T18:32:20.802Z" }, + { url = "https://files.pythonhosted.org/packages/29/7e/2ffce64929fb3cab7b65d5a0b20aaf9764e227681d731b041077fc9a525a/greenlet-3.5.3-cp314-cp314-manylinux_2_39_riscv64.whl", hash = "sha256:962c5df2db8cb446da51edf1ca5296c389d93b99c9d8aa2ee4c7d0d8f1218260", size = 473497, upload-time = "2026-06-26T19:25:39.421Z" }, { url = "https://files.pythonhosted.org/packages/d1/50/13efdbea246fe3d3b735e191fec08fb50809f53cd2383ebe123d0809e44b/greenlet-3.5.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a1fad1d11e7d6aab184107baa8e4ece11ccba3ec9599cd7efa5ff4d70d43256a", size = 1621252, upload-time = "2026-06-26T19:09:05.647Z" }, { url = "https://files.pythonhosted.org/packages/f7/22/c0a336ae4a1410fd5f5121098e5bfbf1865f64c5ef80b4b5412886c4a332/greenlet-3.5.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:fad5aec764399f1b5cc347ad250a59660f20c8f8888ea6bae1f93b769cce1154", size = 1684824, upload-time = "2026-06-26T18:31:47.738Z" }, { url = "https://files.pythonhosted.org/packages/7a/94/91aec0030bea75c4b3244251d0de60a1f3432d1ecb53ab6c437fb5c3ba61/greenlet-3.5.3-cp314-cp314-win_amd64.whl", hash = "sha256:7669aa24cf2a1041d6f7899575b494a3ab4cf68bfcc8609b1dc0be7272db835e", size = 240754, upload-time = "2026-06-26T18:22:15.669Z" }, @@ -1278,14 +1285,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/95/3e161213d7f1d378d15aa9e792093e9bfe01844680d04b7fd6e0107c9098/greenlet-3.5.3-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:271a8ea7c1024e8a0d7dd2be66dd66dda8a07193f41a17b9e924f7600f5b62be", size = 296389, upload-time = "2026-06-26T18:22:20.657Z" }, { url = "https://files.pythonhosted.org/packages/00/92/715c44721abe2b4d1ae9abde4179411868a5bff312479f54e105d372f131/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:19131729ae0ddc3c2e1ef85e650169b5e37ee32e400f215f78b94d7b0d567310", size = 653382, upload-time = "2026-06-26T19:07:14.209Z" }, { url = "https://files.pythonhosted.org/packages/a0/83/37a10372a1090a6624cca8e74c12df1a36c2dc36429ed0255b7fb1aeee23/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1540dd8e5fc2a5aec40fbb98ef8e149fa47c89a4b4a1cf2575a14d3d1869d7a8", size = 659401, upload-time = "2026-06-26T19:10:10.876Z" }, + { url = "https://files.pythonhosted.org/packages/cb/73/8faec206b851c22b1733545fda900829a1f3f5b1c78ae7e0fb3dba57d9f4/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b897d97759425953f69a9c0fac67f8fe333ec0ce7377ef186fb2b0c3ad5e354d", size = 659582, upload-time = "2026-06-26T19:24:21.357Z" }, { url = "https://files.pythonhosted.org/packages/db/e2/d1509cad4207da559cc42986ecdd8fc67ad0d1bba2bf03023c467fd5e0f3/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e81fa194a1d20967877bdf9c7794db2bc99063e5be36aee710c08f04c5bb087f", size = 656969, upload-time = "2026-06-26T18:32:22.272Z" }, + { url = "https://files.pythonhosted.org/packages/b4/55/50c19e49f8045834ada71ef12f8ad048eba8517c6aa41161bed676328fae/greenlet-3.5.3-cp314-cp314t-manylinux_2_39_riscv64.whl", hash = "sha256:3236754d423955ea08e9bb5f6c04a7895f9e22c290b66aa7653fcb922d839eb0", size = 491037, upload-time = "2026-06-26T19:25:40.672Z" }, { url = "https://files.pythonhosted.org/packages/86/7d/eaf70de20aadca3a5884aec58362861c64ce45e7b277f47ed026926a3b89/greenlet-3.5.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:55cf4d777485d43110e47133cbba6d74a8885a87ec1227ef0267f9ee80c5aa21", size = 1617822, upload-time = "2026-06-26T19:09:06.893Z" }, { url = "https://files.pythonhosted.org/packages/8a/f9/414d38fc400ae4350d4185eaad1827676f7cf5287b9136e0ed1cbbe20a7f/greenlet-3.5.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:12a248ba75f6a9a236375f52296c498c89ff1d8badf32deb9eca7abd5853f7da", size = 1677983, upload-time = "2026-06-26T18:31:49.396Z" }, { url = "https://files.pythonhosted.org/packages/e4/15/7edb977e08f9bff702fe42d6c902702786ff6b9694058b4e6a2a6ac90e57/greenlet-3.5.3-cp314-cp314t-win_amd64.whl", hash = "sha256:efc6bd60ea02e085862c74a3ef64b147ffc6f1a5ea7d9f26e7a939943f68c1e3", size = 243626, upload-time = "2026-06-26T18:24:41.485Z" }, { url = "https://files.pythonhosted.org/packages/2c/8a/93928dce91e6b3598b5e779e8d1fd6576a504640c58e78627077f6a7a91a/greenlet-3.5.3-cp315-cp315-macosx_11_0_universal2.whl", hash = "sha256:ea03f2f04367845d6b58eeed276e1e56e51f0b97d8ad5a88a7d20a91dc9056cc", size = 288860, upload-time = "2026-06-26T18:22:48.07Z" }, { url = "https://files.pythonhosted.org/packages/4f/ca/69db42d447a1378043e2c8f19c09cbbd1263371505053c496b49066d3d16/greenlet-3.5.3-cp315-cp315-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:78dbef602fda6d97d957eb7937f70c9ce9e9527330347f8f6b6f9e554a9e7a47", size = 659747, upload-time = "2026-06-26T19:07:15.565Z" }, { url = "https://files.pythonhosted.org/packages/a8/0b/af7ac2ef8dd41e3da1a40dda6305c23b9a03e13ba975ec916357b50f8575/greenlet-3.5.3-cp315-cp315-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6f73857adb8fee13fa56c172bd11262f888c0c648f9fea113e777bb2c7904a81", size = 670419, upload-time = "2026-06-26T19:10:12.293Z" }, + { url = "https://files.pythonhosted.org/packages/25/aa/952cf28c2ff949a8c971134fb43854dd7eaa737218723aaef758f8c9aead/greenlet-3.5.3-cp315-cp315-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:cefa9cef4b371f9844c6053db71f1138bc6807bab1578b0dae5149c1f1141357", size = 674261, upload-time = "2026-06-26T19:24:22.79Z" }, { url = "https://files.pythonhosted.org/packages/51/1e/1d51640cacbfc455dbe9f9a9f594c49e4e244f63b9971a2f4764e46cc53d/greenlet-3.5.3-cp315-cp315-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:232fec92e823addaf02d9472cf7381e24a1d046a6ced1103c5caa4c21b9dfc1d", size = 668787, upload-time = "2026-06-26T18:32:24.298Z" }, + { url = "https://files.pythonhosted.org/packages/dc/f2/b00d6f5e63e531a93562b2ec1a4c320fbee91f580fc42e6417af69d706e5/greenlet-3.5.3-cp315-cp315-manylinux_2_39_riscv64.whl", hash = "sha256:6219b6d04dbf6ba6084d77dc609e8473060dc55f759cbf626d512122781fa128", size = 480322, upload-time = "2026-06-26T19:25:41.852Z" }, { url = "https://files.pythonhosted.org/packages/21/66/4030d5b0b5894500023f003bb054d9bb354dfbd1e186c3a296759172f5f5/greenlet-3.5.3-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:2421c3564da9429d5586d46ca31ebb26516b5498a802cf65c041a8e8a8980d34", size = 1626305, upload-time = "2026-06-26T19:09:08.281Z" }, { url = "https://files.pythonhosted.org/packages/0e/50/5221371c7550108dfa3c378debc41d032aa9c78e89abb01d8011cfc93289/greenlet-3.5.3-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:e0f0d160f0b2e558e6c75f7930967183255dc9735e5f5b8cae58ee09c9576d8b", size = 1688631, upload-time = "2026-06-26T18:31:51.278Z" }, { url = "https://files.pythonhosted.org/packages/68/5d/00d469daae3c65d2bf620b10eee82eb022127d483c6bc8c69fae6f3fbf17/greenlet-3.5.3-cp315-cp315-win_amd64.whl", hash = "sha256:dd99329bbc15ca78dcc583dba05d0b1b0bae01ab6c2174989f5aaee3e41ac930", size = 241027, upload-time = "2026-06-26T18:22:38.203Z" }, @@ -1293,7 +1304,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1c/da/4f4a8450962fad137c1c8981a3f1b8919d06c829993d4d476f9c525d5173/greenlet-3.5.3-cp315-cp315t-macosx_11_0_universal2.whl", hash = "sha256:176bc16a721fa5fc294d70b87b4dfa5fbdd251b3da5d5372735ecef9bd7d6d0c", size = 297221, upload-time = "2026-06-26T18:23:27.176Z" }, { url = "https://files.pythonhosted.org/packages/57/66/b3bfae3e220a9b63ea539a0eea681800c69ab1aada757eae8789f183e7ce/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:629b614d2b786e89c50440e246f33eea78f58a962d0bdbbcc809e6d13605903f", size = 657221, upload-time = "2026-06-26T19:07:16.973Z" }, { url = "https://files.pythonhosted.org/packages/7b/81/b6d4d73a709684fc77e7fa034d7c2fe82cffa9fc920fadcaa659c2626213/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2b2e857ae16f5f72142edf75f9f176fe7526ba19a2841df1420516f83831c9f2", size = 663226, upload-time = "2026-06-26T19:10:13.723Z" }, + { url = "https://files.pythonhosted.org/packages/e9/39/0e0938a75115b939d42733a2a12e1d349653c9531fe6fe563e8a681f04e6/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:16d192579ed281051396dddd7f7754dac6259e6b1fb26378c87b66622f8e3f91", size = 663706, upload-time = "2026-06-26T19:24:24.312Z" }, { url = "https://files.pythonhosted.org/packages/f5/07/e210b02b589f16e74ff48b730690e4a34ffe984219fce4f3c1a0e7ec8545/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e515757e2e36bcbf1fad09a46e1557e8b1ae1797d4b44d09da7deed88ad28608", size = 660802, upload-time = "2026-06-26T18:32:26.081Z" }, + { url = "https://files.pythonhosted.org/packages/5b/41/35d1c678cdb3c3b9e6bee691728e563cfb294202b23c7a4c3c2ccc343589/greenlet-3.5.3-cp315-cp315t-manylinux_2_39_riscv64.whl", hash = "sha256:4399eb8d041f20b68d943918bc55502a93d6fdc0a37c14da7881c04139acee9d", size = 498803, upload-time = "2026-06-26T19:25:43.063Z" }, { url = "https://files.pythonhosted.org/packages/eb/2e/5303eb3fa06bca089060f479707182a93e360683bc252acf846c3090d34e/greenlet-3.5.3-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:b363d46ed1ea431825fdb01471bb024fc08399bad1572a616e853c7684415adb", size = 1622157, upload-time = "2026-06-26T19:09:09.527Z" }, { url = "https://files.pythonhosted.org/packages/54/70/50de47a488f14df260b50ae34fb5d56016e308b098eab02c878b5223c26a/greenlet-3.5.3-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:e44da2f5bbdaabaf7d80b73dbb430c7035771e9f244e3c8b769715c9d8fa0a16", size = 1681159, upload-time = "2026-06-26T18:31:52.986Z" }, { url = "https://files.pythonhosted.org/packages/a7/13/1055e1dda7882073eda533e2b96c62e55bbd2db7fda6d5ece992febc7071/greenlet-3.5.3-cp315-cp315t-win_amd64.whl", hash = "sha256:8ff8bed3e3baa20a3ea261ce00526f1898ad4801d4886fd2220580ee0ad8fadf", size = 244007, upload-time = "2026-06-26T18:22:04.353Z" }, @@ -3811,12 +3824,6 @@ name = "scipy" version = "1.17.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14' and sys_platform == 'win32'", - "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.12.*' and sys_platform == 'emscripten'", @@ -3825,7 +3832,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.11'" }, + { name = "numpy", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ From 53f2a1aab9d02915cd64c27e0a673b2dd622dde5 Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 22:11:16 +0800 Subject: [PATCH 08/10] fix(rust-core): close capability + selection-rule audit findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Post-merge whole-surface audit (capability asymmetries + table/slice selection consistency) findings, all fixed with Python-oracle anchors: Selection rules (2 live on shipped data, 4 reachable): - WideEP MoE distribution fallback now follows file order via a load-time first-seen record (gb200/h100/h200 file order starts power_law_1.01_eplb; lexicographic-first diverged live). - MoE xquant/xprofile candidate enumeration follows quant file order (fp8 vs fp8_block xprofile tie on b200/vllm/0.24.0 was a measured ~13% divergence). - MoeOp EMPIRICAL now routes sglang deepep_moe calibration (own grid + full ladder) to the wideep ctx/gen tables like Python's _moe_table. - EPLB 0.8 token correction scoped to the silicon path only (Python applies it inside get_silicon; empirical uses raw tokens). - WideEP MLA attn_backend whitelist moved into the silicon arm (EMPIRICAL may calibrate from a non-whitelisted slice, per Python); whitelist violations are InvalidEngineConfig so HYBRID propagates. - GB200/NVL72 custom-allreduce reroute is mode-aware in the operator (HYBRID miss falls to NCCL-empirical like Python, not AR-empirical). Capability gaps (freeze blockers from the audit): - Typed errors across the FFI: AicError::PerfDatabase/Io -> sdk PerfDataNotAvailableError, EmpiricalNotImplemented -> sdk EmpiricalNotImplementedError (lazy, cycle-free class resolution; ValueError fallback without the sdk). has_perf_data_not_available_cause now classifies rust misses. - Empirical provenance across the FFI: ProvenanceTier max-rank cell on PerfDatabase (shared across silicon views), note sites mirroring Python one-for-one in every family (incl. the MoE ladder tiers via reference-grid tags and the comm rank-overflow xshape), per-call reset + last_provenance() on AicEngine, and a bridge hook feeding util_empirical.note_provenance — support-matrix HYBRID_PASS tier labelling now works on the rust path. - DSV4 MegaMoE ported (new table loader with Python's row invariants + duplicate-leaf load error, Op::Dsv4MegaMoe appended tail-position, no schema bump; SILICON/HYBRID passthrough contract, no empirical) — DeepSeek-V4 megamoe models no longer fall back to the Python step. - Python fix: the trtllm-alltoall HYBRID fallback closure omitted the mandatory kernel_source arg of get_empirical_from_sol (latent TypeError on every HYBRID silicon miss); fixed + regression test. Verification: parity_tests/ 361 passed (engine-step 303 incl. new typed-error + provenance capture tests; compile-engine; perf gate); cargo test --lib 276 passed; Python unit suite 2077 passed. Known remaining (deliberate): energy/power over the FFI (rust returns latency only; power_w reports 0.0 on rust-routed runs) is a follow-up PR — it threads the energy data column through the interpolation core, every loader, and every op. Rust remains parquet-only by decision (.txt-only data drops must ship parquet). Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- .../parity_tests/test_engine_step_parity.py | 111 +++++- .../aiconfigurator-core/src/engine/runtime.rs | 21 + rust/aiconfigurator-core/src/engine/spec.rs | 30 +- .../src/operators/attention.rs | 10 + .../src/operators/communication.rs | 163 +++++++- rust/aiconfigurator-core/src/operators/dsa.rs | 6 + .../aiconfigurator-core/src/operators/dsv4.rs | 278 +++++++++++++- .../aiconfigurator-core/src/operators/gemm.rs | 11 +- rust/aiconfigurator-core/src/operators/mhc.rs | 2 + rust/aiconfigurator-core/src/operators/mla.rs | 10 + rust/aiconfigurator-core/src/operators/mod.rs | 2 +- rust/aiconfigurator-core/src/operators/moe.rs | 360 +++++++++++++++--- .../src/operators/moe_dispatch.rs | 4 +- rust/aiconfigurator-core/src/operators/msa.rs | 11 + rust/aiconfigurator-core/src/operators/op.rs | 18 +- .../src/operators/util_empirical.rs | 92 +++++ .../src/operators/wideep_mla.rs | 178 +++++++-- .../src/operators/wideep_moe.rs | 55 +++ .../src/perf_database/dsv4_megamoe.rs | 334 ++++++++++++++++ .../src/perf_database/mod.rs | 67 +++- .../src/perf_database/moe.rs | 47 ++- .../src/perf_database/parquet_loader.rs | 14 + .../src/perf_database/wideep.rs | 132 ++++++- .../src/perf_database/wideep_mla.rs | 13 + .../src/perf_database/wideep_moe.rs | 33 +- rust/aiconfigurator-core/src/py.rs | 167 +++++++- src/aiconfigurator/sdk/engine.py | 42 ++ src/aiconfigurator/sdk/operations/moe.py | 1 + src/aiconfigurator/sdk/rust_engine_step.py | 31 +- tests/unit/sdk/database/test_moe_mla.py | 40 ++ tests/unit/sdk/test_rust_engine_step.py | 60 +++ 31 files changed, 2179 insertions(+), 164 deletions(-) create mode 100644 rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs diff --git a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py index ce5efa83b..493089994 100644 --- a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py +++ b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py @@ -19,9 +19,10 @@ import pytest from aiconfigurator.cli.api import cli_estimate -from aiconfigurator.sdk import common, config, perf_database, rust_engine_step +from aiconfigurator.sdk import common, config, errors, perf_database, rust_engine_step from aiconfigurator.sdk.backends.factory import get_backend from aiconfigurator.sdk.models import get_model +from aiconfigurator.sdk.operations import util_empirical pytestmark = pytest.mark.integration @@ -563,8 +564,11 @@ class EngineStepParityCase: # Error-symmetry contract: when Python raises one of these, Rust is # expected to raise (Python's `PerfDataNotAvailableError` and friends # travel through `cli_estimate` as `ValueError` / `RuntimeError` -# subclasses; the Rust FFI surfaces perf-DB misses as -# `RustCoreError`). Tests count any exception in either as a sentinel +# subclasses; the Rust FFI maps `AicError::PerfDatabase`/`Io` to the same +# `PerfDataNotAvailableError` and `AicError::EmpiricalNotImplemented` to +# `EmpiricalNotImplementedError` — see `TestRustTypedErrorsAcrossFfi` — +# with everything else as `ValueError`). Tests count any exception in +# either as a sentinel # value `_ERROR` and assert that *both* engines either compute or # error consistently. Numeric tolerance only applies when both # compute. @@ -1297,3 +1301,104 @@ def test_hybrid_parity( reason = _parity_mismatch_reason(_disagg_comparison_metrics(case), rtol=HYBRID_PARITY_RTOL) assert reason is None, reason + + +def _rust_static_breakdown(case: EngineStepParityCase): + """Drive the rust engine-step bridge directly (no cli_estimate error + wrapping) so the exception object crossing the FFI is what the test sees.""" + database = _case_database(case) + model = _quiet_call(get_model, case.model_path, _case_model_config(case), case.backend_name) + runtime_config = config.RuntimeConfig( + batch_size=case.batch_size, + beam_width=1, + isl=case.isl, + osl=case.osl, + prefix=case.prefix, + ) + return rust_engine_step.estimate_static_latency_breakdown_with_rust( + model, database, runtime_config, "static", 1, 1.0 + ) + + +class TestRustTypedErrorsAcrossFfi: + """FFI typed-error contract (audit finding #8): `aic_to_py` used to map + every `AicError` to `ValueError`, so Python-side classifiers + (`perf_database.has_perf_data_not_available_cause`, the support-matrix + HYBRID-miss triage on `EmpiricalNotImplementedError`) could not recognize + rust-path misses. The boundary now raises the canonical + `aiconfigurator.sdk.errors` classes for the typed variants.""" + + def test_silicon_data_gap_raises_typed_perf_data_miss( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # NVFP4 GEMM tables are not collected on h200/vllm/0.19.0: under + # SILICON, Rust hits `AicError::PerfDatabase` ("GEMM perf data missing + # for quant 'nvfp4'") at the query point — which must cross the FFI as + # the SAME sdk class Python raises, recognized by the cause-chain + # walker (the miss-classification the sweep/support-matrix rely on). + _prepare_rust_core(monkeypatch) + case = EngineStepParityCase(model_path="nvidia/MiniMax-M2.5-NVFP4", system_name="h200_sxm") + with pytest.raises(errors.PerfDataNotAvailableError) as excinfo: + _rust_static_breakdown(case) + assert perf_database.has_perf_data_not_available_cause(excinfo.value) + # The AicError display prefix pins the raise to the rust side of the + # FFI (a Python-side miss would carry the sdk's own wording). + message = str(excinfo.value) + assert "perf database error" in message or "I/O error" in message, message + + def test_hybrid_ladder_miss_raises_typed_empirical_miss( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # NVFP4 MoE on Hopper: no own-shape/cross-shape/sibling reference + # anywhere (the HYBRID_CASES ladder-miss config). Rust raises + # `AicError::EmpiricalNotImplemented`, which must surface as the sdk's + # EmpiricalNotImplementedError — the typed hybrid-miss — and NOT be + # classified as a plain perf-data miss. + _prepare_rust_core(monkeypatch) + case = EngineStepParityCase( + model_path="nvidia/MiniMax-M2.5-NVFP4", + system_name="h200_sxm", + database_mode="HYBRID", + ) + with pytest.raises(errors.EmpiricalNotImplementedError) as excinfo: + _rust_static_breakdown(case) + message = str(excinfo.value) + assert "empirical estimation not implemented" in message, message + assert not perf_database.has_perf_data_not_available_cause(excinfo.value) + + +class TestRustProvenanceCapture: + """FFI provenance contract (audit finding #6): the compiled engine records + the empirical tier that fired (max-rank, mirroring Python's + `PROVENANCE_ORDER`) and the bridge forwards it into + `util_empirical.capture_provenance`, so support-matrix HYBRID_PASS tier + labelling works identically for rust-routed runs.""" + + def test_hybrid_xop_run_records_tier( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # MiniMax-M3 HYBRID: MSA borrows DSA's util (xop) — the run's worst + # tier. Python probes record {xop, xshape}; the rust path must land on + # the same worst_provenance. + _prepare_rust_core(monkeypatch) + case = EngineStepParityCase(model_path="MiniMaxAI/MiniMax-M3", database_mode="HYBRID") + with util_empirical.capture_provenance() as tags: + metrics = _static_metrics(case, engine_step_backend="rust") + assert not isinstance(metrics["total_ms"], _ErrorSentinel), repr(metrics) + assert util_empirical.worst_provenance(tags) == "xop", tags + + def test_pure_silicon_run_records_nothing( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # Fully-collected config on SILICON: no empirical path fires, so the + # capture stays empty and worst_provenance defaults to "silicon". + _prepare_rust_core(monkeypatch) + case = EngineStepParityCase(model_path="moonshotai/Kimi-K2.5") + with util_empirical.capture_provenance() as tags: + metrics = _static_metrics(case, engine_step_backend="rust") + assert not isinstance(metrics["total_ms"], _ErrorSentinel), repr(metrics) + assert util_empirical.worst_provenance(tags) == "silicon", tags diff --git a/rust/aiconfigurator-core/src/engine/runtime.rs b/rust/aiconfigurator-core/src/engine/runtime.rs index 2571f9283..3f48c5d21 100644 --- a/rust/aiconfigurator-core/src/engine/runtime.rs +++ b/rust/aiconfigurator-core/src/engine/runtime.rs @@ -184,6 +184,27 @@ impl Engine { &self.db } + /// 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. + pub fn reset_provenance(&self) { + self.db.reset_provenance(); + } + + /// 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). + pub fn last_provenance(&self) -> Option<&'static str> { + match self.db.worst_provenance() { + crate::operators::util_empirical::ProvenanceTier::Silicon => None, + tier => Some(tier.as_str()), + } + } + /// Test-only accessor for the context op list (the field is private, but /// `fpm`'s `#[cfg(test)]` parity tests compare `forward_pass_time_ms` /// against the shared session free fns over these exact ops). diff --git a/rust/aiconfigurator-core/src/engine/spec.rs b/rust/aiconfigurator-core/src/engine/spec.rs index 66411722c..b2a615e55 100644 --- a/rust/aiconfigurator-core/src/engine/spec.rs +++ b/rust/aiconfigurator-core/src/engine/spec.rs @@ -158,7 +158,8 @@ mod tests { }; use crate::operators::op::{FallbackOp, OverlapOp}; use crate::operators::{ - ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4ModuleOp, + ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4MegaMoeOp, + Dsv4ModuleOp, ElementwiseOp, EmbeddingOp, EncoderAttentionOp, GdnOp, GemmOp, GenerationAttentionOp, GenerationMlaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, MoEDispatchOp, MoeOp, NcclOp, P2POp, TrtllmWideEpMoEDispatchOp, VisionEncoderOp, WideEpContextMlaOp, @@ -437,6 +438,27 @@ mod tests { } } + fn dsv4_megamoe() -> Dsv4MegaMoeOp { + Dsv4MegaMoeOp { + name: "context_megamoe".into(), + scale_factor: 61.0, + hidden_size: 7168, + inter_size: 3072, + topk: 6, + num_experts: 384, + moe_tp_size: 1, + moe_ep_size: 8, + quant_mode: MoeQuantMode::W4a8Mxfp4Mxfp8, + workload_distribution: "balanced".into(), + is_context: true, + source_policy: "random".into(), + pre_dispatch: "sglang_jit".into(), + num_fused_shared_experts: 0, + kernel_source: "deepgemm_megamoe".into(), + kernel_dtype: "fp8_fp4".into(), + } + } + fn mhc() -> MhcModuleOp { MhcModuleOp { name: "mhc_module".into(), @@ -598,6 +620,9 @@ mod tests { OpSpec::WideEpMoeDispatch(wideep_moe_dispatch()), OpSpec::Overlap(overlap()), OpSpec::Fallback(fallback()), + // Appended AFTER Fallback (bincode enum indices are positional; + // appending shifts nothing, so no ENGINE_SPEC_SCHEMA_VERSION bump). + OpSpec::Dsv4MegaMoe(dsv4_megamoe()), ]; // Exhaustiveness guard: if a variant is added to `Op`, this match @@ -635,7 +660,8 @@ mod tests { | OpSpec::WideEpMoe(_) | OpSpec::WideEpMoeDispatch(_) | OpSpec::Overlap(_) - | OpSpec::Fallback(_) => {} + | OpSpec::Fallback(_) + | OpSpec::Dsv4MegaMoe(_) => {} } } ops diff --git a/rust/aiconfigurator-core/src/operators/attention.rs b/rust/aiconfigurator-core/src/operators/attention.rs index 018f9656b..72ba4c233 100644 --- a/rust/aiconfigurator-core/src/operators/attention.rs +++ b/rust/aiconfigurator-core/src/operators/attention.rs @@ -476,6 +476,8 @@ fn context_attention_empirical( })?; if grid.as_deref().is_some_and(|g| !g.is_empty()) { let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python attention.py:406, default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); return Ok(latency); } // Cross-head_size transfer (XSHAPE): this head_size has no data, but @@ -491,6 +493,8 @@ fn context_attention_empirical( / attn_prefill_hs_ratio(&db.backend, ref_hs); let (latency, _) = util_empirical::estimate(sol_time, &query, Some(&ref_grid), scale)?; + // Cross-head_size borrow (Python attention.py:424 "xshape"). + db.note_provenance(util_empirical::ProvenanceTier::XShape); return Ok(latency); } } @@ -638,6 +642,8 @@ fn generation_attention_empirical( })?; if grid.as_deref().is_some_and(|g| !g.is_empty()) { let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python attention.py:791, default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); return Ok(latency); } if db.transfer_policy.contains(TransferKind::XShape) { @@ -646,6 +652,8 @@ fn generation_attention_empirical( { let (latency, _) = util_empirical::estimate(sol_time, &query, Some(&ref_grid), 1.0)?; + // Cross-head_size borrow (Python attention.py:802 "xshape"). + db.note_provenance(util_empirical::ProvenanceTier::XShape); return Ok(latency); } } @@ -752,6 +760,8 @@ fn encoder_attention_empirical( } })?; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python attention.py:1044, default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } diff --git a/rust/aiconfigurator-core/src/operators/communication.rs b/rust/aiconfigurator-core/src/operators/communication.rs index 3044c3aa2..adde4fd68 100644 --- a/rust/aiconfigurator-core/src/operators/communication.rs +++ b/rust/aiconfigurator-core/src/operators/communication.rs @@ -102,9 +102,26 @@ fn query_custom_allreduce_table( tp_size: u32, size: f64, ) -> Result<(f64, Source), AicError> { + // GB200 NVL72 reroute — inside get_silicon only (Python + // `communication.py:188-190`): custom AR is only collected up to tp4 + // there, so (SILICON and HYBRID) queries reroute to the MODE-AWARE + // `database.query_nccl` dispatch. Under HYBRID an NCCL silicon miss + // therefore falls to the NCCL empirical (never the custom-AR + // empirical), and a terminal NCCL EmpiricalNotImplemented propagates + // (Python's outer catch only handles the missing-silicon-data class). + // EMPIRICAL mode never reaches get_silicon, so it never reroutes. + // Handled here rather than via the DB-internal reroute in + // `query_custom_allreduce_scaled`, which is silicon-only; in SILICON + // mode both routes evaluate the identical `query_nccl_scaled` call. + if db.system_spec.node.num_gpus_per_node == 72 + && tp_size > 4 + && db.database_mode != DatabaseMode::Empirical + { + return query_nccl_table(db, quant, tp_size, "all_reduce", size); + } // The silicon path is the DB-level `_scaled` query (fan-out capping + - // beyond-range bandwidth correction inside, incl. the GB200 NVL72 - // reroute), mirroring Python `_query_custom_allreduce_table.get_silicon`. + // beyond-range bandwidth correction inside), mirroring Python + // `_query_custom_allreduce_table.get_silicon`. let silicon = |db: &PerfDatabase| db.communication.query_custom_allreduce_scaled(&db.system_spec, quant, tp_size, size); match db.database_mode { @@ -169,6 +186,13 @@ fn custom_allreduce_empirical( })?; let query = [size]; let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + // Rank-count overflow borrowed the node-boundary tp slice -> "xshape"; + // otherwise own-slice "empirical" (Python communication.py:167-168). + db.note_provenance(if eff != tp_size { + util_empirical::ProvenanceTier::XShape + } else { + util_empirical::ProvenanceTier::Empirical + }); Ok(latency) } @@ -336,6 +360,14 @@ fn nccl_empirical( })?; let query = [message_size]; let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + // Fan-out beyond the largest collected num_gpus borrowed the boundary + // slice -> "xshape"; otherwise own-slice "empirical" (Python + // communication.py:439-440). + db.note_provenance(if eff != num_gpus { + util_empirical::ProvenanceTier::XShape + } else { + util_empirical::ProvenanceTier::Empirical + }); Ok(latency) } @@ -385,6 +417,7 @@ fn p2p_bandwidth(spec: &SystemSpec, num_gpus: u32) -> f64 { #[cfg(test)] mod tests { use super::*; + use crate::perf_database::CommunicationTable; use std::path::PathBuf; const REPO_ROOT_HINT: &str = env!("CARGO_MANIFEST_DIR"); @@ -405,18 +438,21 @@ mod tests { fn custom_allreduce_empirical_matches_python_oracles() { let mut db = b200_vllm_db(); db.database_mode = DatabaseMode::Empirical; + use util_empirical::ProvenanceTier; let cases = [ // off-grid size on a collected tp slice - (8u32, 300_000u64, 0.0210877421663319), + (8u32, 300_000u64, 0.0210877421663319, ProvenanceTier::Empirical), // exact collected hit: util reconstruction returns the measured value - (8, 1024, 0.007716159820556641), + (8, 1024, 0.007716159820556641, ProvenanceTier::Empirical), // rank overflow: tp=16 on an 8-GPU node borrows the tp=8 boundary - // util; SOL(tp=16, inter-node bw) carries the multi-node scaling - (16, 524_288, 0.1607595384120941), - (2, 777_777, 0.019098769751423886), - (4, 65_536, 0.008213120102882384), + // util; SOL(tp=16, inter-node bw) carries the multi-node scaling. + // Python capture: {"xshape"} (communication.py:167). + (16, 524_288, 0.1607595384120941, ProvenanceTier::XShape), + (2, 777_777, 0.019098769751423886, ProvenanceTier::Empirical), + (4, 65_536, 0.008213120102882384, ProvenanceTier::Empirical), ]; - for (tp, size, expected) in cases { + for (tp, size, expected, tier) in cases { + db.reset_provenance(); let (latency, source) = query_custom_allreduce_table(&db, CommQuantMode::Half, tp, size as f64) .expect("empirical query"); @@ -425,6 +461,7 @@ mod tests { "(tp={tp}, size={size}): expected {expected}, got {latency}" ); assert_eq!(source, Source::Empirical); + assert_eq!(db.worst_provenance(), tier, "(tp={tp}, size={size}): wrong tier"); } } @@ -476,19 +513,22 @@ mod tests { fn nccl_empirical_matches_python_oracles() { let mut db = b200_vllm_db(); db.database_mode = DatabaseMode::Empirical; + use util_empirical::ProvenanceTier; let cases = [ // off-grid message sizes on collected fan-outs - (8u32, "all_reduce", 300_000u64, 0.027067167868039803), - (4, "all_gather", 10_000_000, 0.05871496201240777), + (8u32, "all_reduce", 300_000u64, 0.027067167868039803, ProvenanceTier::Empirical), + (4, "all_gather", 10_000_000, 0.05871496201240777, ProvenanceTier::Empirical), // gpu-count overflow: 32 > max collected 8 borrows the boundary - // util; SOL(32, inter-node bw) carries the scaling - (32, "all_reduce", 1_048_576, 0.31676464285714284), + // util; SOL(32, inter-node bw) carries the scaling. + // Python capture: {"xshape"} (communication.py:439). + (32, "all_reduce", 1_048_576, 0.31676464285714284, ProvenanceTier::XShape), // below-min size: boundary util-hold - (8, "reduce_scatter", 999, 0.015899137756552793), + (8, "reduce_scatter", 999, 0.015899137756552793, ProvenanceTier::Empirical), // exact collected hit - (2, "alltoall", 4096, 0.009470000000000001), + (2, "alltoall", 4096, 0.009470000000000001, ProvenanceTier::Empirical), ]; - for (num_gpus, op, msg, expected) in cases { + for (num_gpus, op, msg, expected, tier) in cases { + db.reset_provenance(); let (latency, source) = query_nccl_table(&db, CommQuantMode::Half, num_gpus, op, msg as f64) .expect("empirical query"); assert!( @@ -496,6 +536,7 @@ mod tests { "({num_gpus}, {op}, {msg}): expected {expected}, got {latency}" ); assert_eq!(source, Source::Empirical); + assert_eq!(db.worst_provenance(), tier, "({num_gpus}, {op}, {msg}): wrong tier"); } } @@ -535,6 +576,96 @@ mod tests { assert_eq!(latency, 0.0); } + /// GB200 NVL72 reroute (num_gpus_per_node == 72 && tp > 4): Python's + /// get_silicon returns the MODE-AWARE `database.query_nccl(...)` + /// (communication.py:188-190), so under HYBRID an NCCL miss falls to the + /// NCCL empirical — never the custom-AR empirical — and EMPIRICAL mode + /// never reroutes at all. No shipped spec has 72 GPUs per node, so the + /// spec is patched like the Python oracle run: + /// + /// ```text + /// db = perf_database.get_database_view("b200_sxm", "vllm", "0.19.0", + /// allow_missing_data=True, database_mode=..., shared_layer=False) + /// db.system_spec = copy.deepcopy(db.system_spec) + /// db.system_spec["node"]["num_gpus_per_node"] = 72 + /// CustomAllReduce._query_custom_allreduce_table(db, half, tp, size, mode) + /// ``` + #[test] + fn custom_allreduce_nvl72_reroute_is_mode_aware() { + let nvl72_db = |mode: DatabaseMode| { + let mut db = b200_vllm_db(); + db.tables_mut().system_spec.node.num_gpus_per_node = 72; + db.database_mode = mode; + db + }; + + // HYBRID with NCCL data present: the reroute answers from the NCCL + // silicon table (tp=16 exercises the beyond-max fan-out scaling). + let db = nvl72_db(DatabaseMode::Hybrid); + for (tp, size, expected) in [ + (8u32, 300_000u64, 0.02807729736328125), + (16, 524_288, 0.031017857142857142), + ] { + let (latency, source) = + query_custom_allreduce_table(&db, CommQuantMode::Half, tp, size as f64) + .expect("hybrid reroute"); + assert!( + (latency - expected).abs() < 1e-9, + "(tp={tp}, size={size}): expected {expected}, got {latency}" + ); + assert_eq!(source, Source::Silicon); + } + + // SILICON mode is unchanged: the operator-level reroute evaluates the + // identical `query_nccl_scaled` the DB-internal reroute did. + let db = nvl72_db(DatabaseMode::Silicon); + let (latency, source) = + query_custom_allreduce_table(&db, CommQuantMode::Half, 8, 300_000.0) + .expect("silicon reroute"); + let scaled = db + .communication + .query_custom_allreduce_scaled(&db.system_spec, CommQuantMode::Half, 8, 300_000.0) + .expect("db-level reroute"); + assert!( + (latency - scaled).abs() < 1e-12, + "operator ({latency}) and DB-level ({scaled}) reroutes must agree" + ); + assert_eq!(source, Source::Silicon); + + // EMPIRICAL never reaches get_silicon, hence never reroutes: the + // value stays the custom-AR empirical estimate (same as the + // unpatched-spec oracle at tp=8/300k). + let db = nvl72_db(DatabaseMode::Empirical); + let (latency, source) = + query_custom_allreduce_table(&db, CommQuantMode::Half, 8, 300_000.0) + .expect("empirical never reroutes"); + assert!( + (latency - 0.0210877421663319).abs() < 1e-9, + "expected the custom-AR empirical value, got {latency}" + ); + assert_eq!(source, Source::Empirical); + + // HYBRID with NO NCCL data at all: the nested NCCL empirical raises + // the terminal EmpiricalNotImplemented, which PROPAGATES (Python: + // "No NCCL data to estimate all_reduce (half, num_gpus=8)." — not in + // `_MISSING_SILICON_DATA_EXCEPTIONS`, so the outer catch never runs + // the custom-AR empirical). The pre-fix DB-internal reroute wrongly + // fell back to the custom-AR empirical here. + let mut db = b200_vllm_db(); + { + let tables = db.tables_mut(); + tables.system_spec.node.num_gpus_per_node = 72; + tables.communication = + CommunicationTable::new(tables.data_root.clone(), None, None); + } + db.database_mode = DatabaseMode::Hybrid; + let result = query_custom_allreduce_table(&db, CommQuantMode::Half, 8, 300_000.0); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "expected the NCCL empirical miss to propagate, got {result:?}" + ); + } + /// HYBRID prefers silicon whenever the table covers the slice. Oracles: /// `NCCL._query_nccl_table(..., database_mode=HYBRID)`. #[test] diff --git a/rust/aiconfigurator-core/src/operators/dsa.rs b/rust/aiconfigurator-core/src/operators/dsa.rs index 4df1c09d0..77e156cd8 100644 --- a/rust/aiconfigurator-core/src/operators/dsa.rs +++ b/rust/aiconfigurator-core/src/operators/dsa.rs @@ -688,6 +688,8 @@ fn context_empirical( )))) })?; let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python dsa.py:620, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -755,6 +757,8 @@ fn generation_empirical( })?; let (latency, _) = util_empirical::estimate(sol_time, &[b as f64, s as f64], grid.as_deref(), 1.0)?; + // Own-shape util fired (Python dsa.py:1296, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } else if let Some(slc) = slice { let skip_tag = if skip_indexer { ":skip" } else { "" }; @@ -778,6 +782,8 @@ fn generation_empirical( grid.as_deref(), 1.0, )?; + // Own-shape util fired (Python dsa.py:1296, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } else { let (latency, _) = util_empirical::estimate( diff --git a/rust/aiconfigurator-core/src/operators/dsv4.rs b/rust/aiconfigurator-core/src/operators/dsv4.rs index b5c32bb03..6b5e3e760 100644 --- a/rust/aiconfigurator-core/src/operators/dsv4.rs +++ b/rust/aiconfigurator-core/src/operators/dsv4.rs @@ -17,7 +17,7 @@ //! `perf_database::dsv4` and Python `load_context_dsv4_kind_module_data`. use serde::{Deserialize, Serialize}; -use crate::common::enums::{DatabaseMode, FmhaQuantMode, GemmQuantMode, KvCacheQuantMode}; +use crate::common::enums::{DatabaseMode, FmhaQuantMode, GemmQuantMode, KvCacheQuantMode, MoeQuantMode}; use crate::common::error::AicError; use crate::operators::base::{PerformanceResult, Source}; use crate::operators::communication::NcclOp; @@ -334,6 +334,8 @@ impl Dsv4ModuleOp { })?; let query = [f64::from(prefix), f64::from(s), f64::from(b)]; let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python dsv4.py:889, default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } else { let sol2 = |c: &[f64]| sol_at(c[1] as i64, c[0] as i64, 0); // c=(s, b); anchor prefix=0 @@ -354,6 +356,8 @@ impl Dsv4ModuleOp { })?; let query = [f64::from(s) + f64::from(prefix), f64::from(b)]; let (latency, _) = util_empirical::estimate(sol_q, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python dsv4.py:889, default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } } @@ -638,10 +642,110 @@ impl Dsv4ModuleOp { })?; let query = [f64::from(b), f64::from(s)]; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python dsv4.py:1313, default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } } +/// SGLang DeepSeek-V4 MegaMoE routed module. +/// +/// Mirrors Python `operations/dsv4.py::DeepSeekV4MegaMoEModule`: the measured +/// routed MegaMoE module boundary (prepared hidden states + top-k tensors -> +/// SGLang pre-dispatch -> `deep_gemm.fp8_fp4_mega_moe` -> routed output +/// scaling). Gate/top-k and shared experts are modeled outside this op. One +/// class serves both phases via `is_context` (same as Python), so a single +/// `Op::Dsv4MegaMoe` variant dispatches on `ctx.num_tokens`. +/// +/// Mode contract (Python `_query_megamoe_table`): SILICON/HYBRID only. The op +/// has NO empirical path — EMPIRICAL (and the SOL diagnostics) raise a typed +/// perf-data error, and under HYBRID a silicon miss propagates RAW (Python +/// never routes this op through `_query_silicon_or_hybrid`). +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Dsv4MegaMoeOp { + pub name: String, + pub scale_factor: f64, + pub hidden_size: u32, + /// FULL (un-partitioned) MoE inter size — Python passes + /// `self._moe_inter_size` (models/deepseek_v4.py), not the tp-local one. + pub inter_size: u32, + pub topk: u32, + pub num_experts: u32, + pub moe_tp_size: u32, + pub moe_ep_size: u32, + pub quant_mode: MoeQuantMode, + /// Normalized by the Python op ctor (`uniform` -> `balanced`); the query + /// re-applies the normalization defensively (see `query`). + pub workload_distribution: String, + pub is_context: bool, + pub source_policy: String, + pub pre_dispatch: String, + pub num_fused_shared_experts: u32, + pub kernel_source: String, + pub kernel_dtype: String, +} + +impl Dsv4MegaMoeOp { + /// Query measured MegaMoE routed-module latency at the rank-LOCAL token + /// count `num_tokens` (Python `query`'s `x`; the perf rows are indexed by + /// local-rank tokens — do NOT pre-multiply by attention_dp_size). + pub fn query(&self, db: &PerfDatabase, num_tokens: u32) -> Result { + // Python `DeepSeekV4MegaMoEModule.query`: Blackwell-only guard + // (`ValueError`, not a perf-data miss). + let sm_version = db + .system_spec + .gpu + .sm_version + .map_or(-1_i64, i64::from); + if sm_version < 100 { + return Err(AicError::ModelConfig(format!( + "DeepSeek-V4 MegaMoE is only supported on Blackwell-class GPUs (SM >= 100); \ + got sm_version={sm_version}." + ))); + } + // Python `_query_megamoe_table`: `database_mode not in (SILICON, + // HYBRID)` -> PerfDataNotAvailableError. HYBRID takes the SAME + // silicon table path as SILICON, and a miss propagates raw. + match db.database_mode { + DatabaseMode::Silicon | DatabaseMode::Hybrid => {} + mode => { + return Err(AicError::PerfDatabase(format!( + "DSv4 MegaMoE module only supports measured SILICON data, got \ + database_mode={mode:?}." + ))) + } + } + // Python `_normalize_distribution` (op ctor): uniform -> balanced. + // The Python emitter sends the already-normalized value; re-apply for + // hand-written specs. + let distribution = if self.workload_distribution == "uniform" { + "balanced" + } else { + self.workload_distribution.as_str() + }; + let latency = db.dsv4_megamoe.query_module( + num_tokens, + self.hidden_size, + self.inter_size, + self.topk, + self.num_experts, + self.moe_tp_size, + self.moe_ep_size, + self.quant_mode, + distribution, + self.is_context, + &self.source_policy, + &self.pre_dispatch, + self.num_fused_shared_experts, + &self.kernel_source, + &self.kernel_dtype, + )?; + // Python: `PerformanceResult(float(result) * scale, source="silicon")` + // — no clamp (nothing is subtracted on this path). + Ok(PerformanceResult::new(latency, Source::Silicon).scaled(self.scale_factor)) + } +} + #[cfg(test)] mod tests { use super::*; @@ -1067,6 +1171,178 @@ mod tests { ); } + /// DeepSeek-V4-Pro MegaMoE routed-module op mirroring the model builder + /// (`models/deepseek_v4.py::_moe_ops` with `use_megamoe`): full + /// moe_inter_size, defaults for source_policy / pre_dispatch / fused + /// shared experts / kernel identity. + fn megamoe_op(is_context: bool, distribution: &str) -> Dsv4MegaMoeOp { + Dsv4MegaMoeOp { + name: if is_context { "context_megamoe" } else { "generation_megamoe" }.into(), + scale_factor: 1.0, + hidden_size: 7168, + inter_size: 3072, + topk: 6, + num_experts: 384, + moe_tp_size: 1, + moe_ep_size: 8, + quant_mode: MoeQuantMode::W4a8Mxfp4Mxfp8, + workload_distribution: distribution.into(), + is_context, + source_policy: "random".into(), + pre_dispatch: "sglang_jit".into(), + num_fused_shared_experts: 0, + kernel_source: "deepgemm_megamoe".into(), + kernel_dtype: "fp8_fp4".into(), + } + } + + fn gb200_db() -> Option { + let systems_root = systems_root(); + let data_root = systems_root.join("data/gb200/sglang/0.5.10"); + if !data_root.join("dsv4_megamoe_module_perf.parquet").exists() { + return None; // git-lfs data not materialized + } + Some( + PerfDatabase::load(&systems_root, "gb200", "sglang", "0.5.10") + .expect("gb200/sglang/0.5.10 must load"), + ) + } + + /// SILICON parity anchors on the shipped gb200/sglang/0.5.10 MegaMoE + /// table. Python oracle: + /// + /// ```text + /// uv run --no-sync python -c " + /// from aiconfigurator.sdk import common + /// from aiconfigurator.sdk.perf_database import get_database + /// from aiconfigurator.sdk.operations.dsv4 import DeepSeekV4MegaMoEModule as M + /// from aiconfigurator.sdk.common import DatabaseMode + /// db = get_database('gb200', 'sglang', '0.5.10') + /// base = dict(hidden_size=7168, inter_size=3072, topk=6, num_experts=384, + /// moe_tp_size=1, moe_ep_size=8, + /// quant_mode=common.MoEQuantMode.w4a8_mxfp4_mxfp8) + /// for tok, dist, ctx in [(1024,'balanced',True), (3000,'power_law_1.2',True), + /// (100,'balanced',True), (64,'power_law_1.01',False), + /// (2000,'balanced',False)]: + /// print(float(M._query_megamoe_table(db, num_tokens=tok, + /// workload_distribution=dist, is_context=ctx, + /// database_mode=DatabaseMode.SILICON, **base)))" + /// ``` + /// + /// Covers: exact collected hit, off-grid in-range lerp, below-range and + /// beyond-range boundary util-hold (linear token-proxy SOL), both phases. + #[test] + fn megamoe_silicon_matches_python_oracles() { + let Some(db) = gb200_db() else { return }; + assert_eq!(db.database_mode, DatabaseMode::Silicon); + + let cases: &[(bool, &str, u32, f64)] = &[ + // exact collected hit (context, tokens=1024, balanced, ep=8) + (true, "balanced", 1024, 0.508394), + // off-grid in-range lerp (context, tokens=3000 between 2048/4096) + (true, "power_law_1.2", 3000, 2.01460530859375), + // below the collected range (context, tokens=100 < 1024) + (true, "balanced", 100, 0.0496478515625), + // exact collected hit (generation, tokens=64) + (false, "power_law_1.01", 64, 0.312203), + // beyond the collected range (generation, tokens=2000 > 512) + (false, "balanced", 2000, 1.5932929687500001), + ]; + for &(is_context, dist, tokens, expected) in cases { + let result = megamoe_op(is_context, dist) + .query(&db, tokens) + .expect("silicon megamoe query"); + approx(result.latency_ms, expected); + assert_eq!(result.source, Source::Silicon); + } + } + + /// Mode + miss contract, mirroring Python `_query_megamoe_table`: + /// - EMPIRICAL -> typed PerfDataNotAvailableError (the op has NO + /// empirical path); + /// - an absent shape is a typed miss under SILICON; + /// - HYBRID == SILICON: hit passes through, a miss propagates RAW + /// (never converted into an empirical estimate). + /// - `uniform` normalizes to `balanced` (Python ctor normalization). + #[test] + fn megamoe_mode_contract_and_miss() { + let Some(mut db) = gb200_db() else { return }; + + // EMPIRICAL -> error, Python message mirrored. + db.database_mode = DatabaseMode::Empirical; + let err = megamoe_op(true, "balanced").query(&db, 1024).unwrap_err(); + assert!(err.is_missing_perf_data(), "got {err:?}"); + assert!( + err.to_string() + .contains("DSv4 MegaMoE module only supports measured SILICON data"), + "unexpected message: {err}" + ); + + // SILICON miss: absent shape is a typed miss naming the key. + db.database_mode = DatabaseMode::Silicon; + let mut op = megamoe_op(true, "balanced"); + op.num_experts = 999; + let err = op.query(&db, 1024).unwrap_err(); + assert!(err.is_missing_perf_data(), "got {err:?}"); + assert!( + err.to_string().contains("No DSv4 MegaMoE context module data") + && err.to_string().contains("num_experts=999"), + "unexpected message: {err}" + ); + + // HYBRID: hit == silicon value; miss propagates raw (same typed + // error, no empirical conversion). + db.database_mode = DatabaseMode::Hybrid; + let result = megamoe_op(true, "balanced") + .query(&db, 1024) + .expect("hybrid megamoe hit"); + approx(result.latency_ms, 0.508394); + assert_eq!(result.source, Source::Silicon); + let err = op.query(&db, 1024).unwrap_err(); + assert!(err.is_missing_perf_data(), "got {err:?}"); + + // `uniform` -> `balanced` (Python `_normalize_distribution`). + db.database_mode = DatabaseMode::Silicon; + let result = megamoe_op(true, "uniform") + .query(&db, 1024) + .expect("uniform must normalize to balanced"); + approx(result.latency_ms, 0.508394); + } + + /// The exact externally-tagged JSON the Python emitter + /// (`engine.py::_dsv4_megamoe`) produces must decode into the op — + /// pins the wire field names. + #[test] + fn megamoe_python_wire_json_decodes() { + let json = serde_json::json!({ + "Dsv4MegaMoe": { + "name": "context_megamoe", + "scale_factor": 61.0, + "hidden_size": 7168, + "inter_size": 3072, + "topk": 6, + "num_experts": 384, + "moe_tp_size": 1, + "moe_ep_size": 8, + "quant_mode": "w4a8_mxfp4_mxfp8", + "workload_distribution": "balanced", + "is_context": true, + "source_policy": "random", + "pre_dispatch": "sglang_jit", + "num_fused_shared_experts": 0, + "kernel_source": "deepgemm_megamoe", + "kernel_dtype": "fp8_fp4", + } + }); + let op: crate::operators::op::Op = serde_json::from_value(json).expect("decode"); + let crate::operators::op::Op::Dsv4MegaMoe(op) = op else { + panic!("expected Dsv4MegaMoe variant"); + }; + assert_eq!(op.quant_mode, MoeQuantMode::W4a8Mxfp4Mxfp8); + assert_eq!(op.scale_factor, 61.0); + assert!(op.is_context); + } + /// `cp_size` / `window_size` are absent from every opspec the Python /// emitter produces today (`engine.py::_reject_cp` still guards CP) — /// they must default to 1 / None so existing specs keep the plain non-CP diff --git a/rust/aiconfigurator-core/src/operators/gemm.rs b/rust/aiconfigurator-core/src/operators/gemm.rs index cd26c5c4c..be74fb9b7 100644 --- a/rust/aiconfigurator-core/src/operators/gemm.rs +++ b/rust/aiconfigurator-core/src/operators/gemm.rs @@ -176,6 +176,8 @@ fn gemm_empirical( })?; let query = [m as f64, n as f64, k as f64]; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python estimate()'s default provenance). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -220,7 +222,12 @@ fn compute_scale_empirical( })?; // sol_mem = 2 m k / bw * 1000 (read + write of the activation). let spec = &db.system_spec; - lookup.estimate(&[m as f64, k as f64], |c| 2.0 * c[0] * c[1] / spec.gpu.mem_bw * 1000.0) + let latency = + lookup.estimate(&[m as f64, k as f64], |c| 2.0 * c[0] * c[1] / spec.gpu.mem_bw * 1000.0)?; + // The delta lookup fired (Python `_ZeroAwareDeltaLookup.estimate` notes + // "empirical"; zero deltas count — they are measured values). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); + Ok(latency) } /// scale_matrix latency for `(m, k)` under the database's query mode. @@ -265,6 +272,8 @@ fn scale_matrix_empirical( })?; let query = [m as f64, k as f64]; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python estimate()'s default provenance). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } diff --git a/rust/aiconfigurator-core/src/operators/mhc.rs b/rust/aiconfigurator-core/src/operators/mhc.rs index 5246a4ca6..04e4c3a66 100644 --- a/rust/aiconfigurator-core/src/operators/mhc.rs +++ b/rust/aiconfigurator-core/src/operators/mhc.rs @@ -175,6 +175,8 @@ impl MhcModuleOp { })?; let query = [f64::from(num_tokens)]; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python mhc.py, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } } diff --git a/rust/aiconfigurator-core/src/operators/mla.rs b/rust/aiconfigurator-core/src/operators/mla.rs index fab240ff1..91be8b1e4 100644 --- a/rust/aiconfigurator-core/src/operators/mla.rs +++ b/rust/aiconfigurator-core/src/operators/mla.rs @@ -323,6 +323,8 @@ fn context_mla_empirical( ); let query = [num_heads as f64, (s + prefix) as f64, b as f64]; let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python mla.py, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -373,6 +375,8 @@ fn generation_mla_empirical( })?; let query = [num_heads as f64, b as f64, s as f64]; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python mla.py, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -448,6 +452,8 @@ fn mla_bmm_empirical( }; let query = [num_tokens as f64]; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python mla.py, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -531,6 +537,8 @@ fn context_mla_module_empirical( ); let query = [num_heads as f64, (s + prefix) as f64, b as f64]; let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python mla.py, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -593,6 +601,8 @@ fn generation_mla_module_empirical( })?; let query = [num_heads as f64, b as f64, s as f64]; let (latency, _) = util_empirical::estimate(sol(&query), &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python mla.py, estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } diff --git a/rust/aiconfigurator-core/src/operators/mod.rs b/rust/aiconfigurator-core/src/operators/mod.rs index a89ca25f0..af29e2aa1 100644 --- a/rust/aiconfigurator-core/src/operators/mod.rs +++ b/rust/aiconfigurator-core/src/operators/mod.rs @@ -38,7 +38,7 @@ pub use attention::{ContextAttentionOp, EncoderAttentionOp, GenerationAttentionO pub use base::{PerformanceResult, Source}; pub use communication::{CustomAllReduceOp, NcclOp, P2POp}; pub use dsa::DsaModuleOp; -pub use dsv4::Dsv4ModuleOp; +pub use dsv4::{Dsv4MegaMoeOp, Dsv4ModuleOp}; pub use elementwise::ElementwiseOp; pub use embedding::EmbeddingOp; pub use gemm::GemmOp; diff --git a/rust/aiconfigurator-core/src/operators/moe.rs b/rust/aiconfigurator-core/src/operators/moe.rs index 438fbad86..406aefbb4 100644 --- a/rust/aiconfigurator-core/src/operators/moe.rs +++ b/rust/aiconfigurator-core/src/operators/moe.rs @@ -27,10 +27,11 @@ //! [`AicError::EmpiricalNotImplemented`] only surfaces when every permitted //! tier found nothing. //! -//! Scope: the SGLang `moe_backend == "deepep_moe"` branch of Python's -//! `_moe_table` does not exist here — the Rust engine routes WideEP MoE -//! through `WideEpMoeOp` (`operators/wideep_moe.rs`), so this op only ever -//! addresses the regular / low-latency tables. +//! The SGLang `moe_backend == "deepep_moe"` branch of Python's `_moe_table` +//! routes BOTH the silicon lookup and the empirical calibration (own-shape +//! grid + transfer ladder) to the wideep context/generation MoE tables — +//! mirrored here via [`MoeTableSel`]. (The TRT-LLM WideEP compute table is a +//! different op: `WideEpMoeOp`, `operators/wideep_moe.rs`.) //! //! Weights accounting (per-expert FFN weights + router) is in the model //! layer; the operator returns latency only. @@ -94,9 +95,9 @@ fn moe_quant_from_name(name: &str) -> Option { /// Collected quants with a DIFFERENT `(memory, compute)` profile than the /// query, nearest-profile first (stable sort by `|Δmemory| + |Δcompute|`; -/// mirrors `_xprofile_moe_quants`, `operations/moe.py:105-117`). NOTE: -/// Python breaks distance ties by table insertion (file row) order; here -/// the input arrives in the accessor's sorted-name order. +/// mirrors `_xprofile_moe_quants`, `operations/moe.py:105-117`). Python +/// breaks distance ties by table insertion (file row) order — the accessor +/// feeds that load order in, and the stable sort preserves it. fn xprofile_moe_quants(query: MoeQuantMode, table_quants: &[MoeQuantMode]) -> Vec { let qp = query.mapping(); let mut refs: Vec = table_quants @@ -125,6 +126,42 @@ fn policy_fingerprint(policy: TransferPolicy) -> String { ) } +/// Which perf table calibrates the EMPIRICAL path. Mirrors Python's +/// `_moe_table()` selection (`operations/moe.py:364-397`): SGLang +/// `moe_backend == "deepep_moe"` routes to the wideep context/generation +/// MoE tables; nvfp4 small-token gated probes the TRT-LLM low-latency +/// split; everything else uses the default table. +#[derive(Clone, Copy, PartialEq)] +enum MoeTableSel { + Standard, + LowLatency, + Wideep { is_context: bool }, +} + +impl MoeTableSel { + /// Grid cache-key tag. Python folds `kernel_tag` ("std" / "ll" / + /// "wideep") plus `id(node)` into the key; the ctx/gen split here plays + /// the node-identity role, so the two wideep tables cannot alias. + fn tag(self) -> &'static str { + match self { + Self::Standard => "std", + Self::LowLatency => "ll", + Self::Wideep { is_context: true } => "wideep_ctx", + Self::Wideep { is_context: false } => "wideep_gen", + } + } +} + +/// The perf-DB kernel grid behind a non-wideep selector. Callers match the +/// wideep variants off to `WideEpTable` accessors before reaching this. +fn moe_kernel(table: MoeTableSel) -> MoeKernel { + match table { + MoeTableSel::Standard => MoeKernel::Standard, + MoeTableSel::LowLatency => MoeKernel::LowLatency, + MoeTableSel::Wideep { .. } => unreachable!("wideep selectors dispatch to WideEpTable"), + } +} + /// A sibling slice the transfer ladder may borrow: the reference slice's /// shape + token curve, the quant whose SOL reconstructs its util /// (`sol_quant`: QUERY quant for same-profile tiers, REFERENCE quant for @@ -219,17 +256,6 @@ impl MoeOp { // `MoE.query` (`x = x * attention_dp_size`). Applied exactly once, // before the perf-DB resolution keys off the token count. let num_tokens = num_tokens.saturating_mul(self.attention_dp_size.max(1)); - let is_sglang = db.backend == "sglang"; - // SGLang EPLB prefill correction (Python operations/moe.py: - // `num_tokens_corrected = int(num_tokens * 0.8) if enable_eplb and - // is_context else num_tokens`, sglang branch only). Applied before - // BOTH the silicon resolution and the empirical estimate (Python - // corrects the token count ahead of the table query). - let num_tokens = if is_sglang && self.enable_eplb && self.is_context { - (num_tokens as f64 * 0.8) as u32 - } else { - num_tokens - }; // Database-mode dispatch, mirroring the Python `_query_moe_table` // tail (`database._query_silicon_or_hybrid`): EMPIRICAL always @@ -261,6 +287,16 @@ impl MoeOp { /// grid, scale/clamp applied per branch — the audit-PR body, unchanged). fn silicon_pr(&self, db: &PerfDatabase, num_tokens: u32) -> Result { let is_sglang = db.backend == "sglang"; + // SGLang EPLB prefill correction — INSIDE get_silicon only (Python + // operations/moe.py:684, sglang branch: `num_tokens_corrected = + // int(num_tokens * 0.8) if enable_eplb and is_context else + // num_tokens`). The EMPIRICAL closures receive RAW tokens + // (moe.py:637-647, 803-813), so the correction must not leak there. + let num_tokens = if is_sglang && self.enable_eplb && self.is_context { + (num_tokens as f64 * 0.8) as u32 + } else { + num_tokens + }; // The roofline SOL the perf-DB engine anchors its beyond-range // util-hold on (Python `_resolve_tokens` passes the same closure). // Coordinates arriving from the engine are always integral (table @@ -373,30 +409,34 @@ impl MoeOp { self.moe_ep_size, ); - // Kernel-table selection mirrors get_silicon's (`_moe_table`, - // `operations/moe.py:355-386`): nvfp4 + small tokens + gated probes - // the low-latency table for the FULL slice and falls back to the - // default table on a shape miss. Building util from the wrong table - // would over-estimate by the ~3x kernel gap. The `kernel_tag` folds - // the choice into every grid cache key so a low-latency grid can't - // be served to a regular query (or vice versa) at the same shape. - let kernel = if num_tokens <= 128 + // Table selection mirrors get_silicon's (`_moe_table`, + // `operations/moe.py:364-397`): the SGLang deepep branch comes FIRST + // and routes the whole calibration (own-shape grid + ladder) to the + // wideep context/generation tables; otherwise nvfp4 + small tokens + + // gated probes the low-latency table for the FULL slice and falls + // back to the default table on a shape miss. Building util from the + // wrong table would over-estimate by the ~3x kernel gap. The tag + // folds the choice into every grid cache key so one table's grid + // can't be served to another's query at the same shape. + let table = if db.backend == "sglang" && self.moe_backend.as_deref() == Some("deepep_moe") + { + MoeTableSel::Wideep { + is_context: self.is_context, + } + } else if num_tokens <= 128 && quant == MoeQuantMode::Nvfp4 && self.is_gated && db.moe.low_latency_available()? { - match self.slice_points(db, MoeKernel::LowLatency) { - Ok(_) => MoeKernel::LowLatency, - Err(err) if err.is_missing_perf_data() => MoeKernel::Standard, + match self.slice_points(db, MoeTableSel::LowLatency) { + Ok(_) => MoeTableSel::LowLatency, + Err(err) if err.is_missing_perf_data() => MoeTableSel::Standard, Err(err) => return Err(err), } } else { - MoeKernel::Standard - }; - let kernel_tag = match kernel { - MoeKernel::LowLatency => "ll", - MoeKernel::Standard => "std", + MoeTableSel::Standard }; + let kernel_tag = table.tag(); // Own-shape grid over this slice's num_tokens curve (depth 1). let own_sol = |c: &[f64]| { @@ -427,7 +467,7 @@ impl MoeOp { num_gemms, ); let mut grid = db.util_grids.get_or_try_build(&own_key, || { - match self.slice_points(db, kernel) { + match self.slice_points(db, table) { Ok(points) => Ok(Some(UtilGrid::new(util_empirical::build_samples( points.into_iter().map(|(t, lat)| (vec![t as f64], lat)), own_sol, @@ -449,11 +489,11 @@ impl MoeOp { // xshape set falls through to same-profile xquant siblings. let mut candidates: Vec = Vec::new(); if policy.contains(TransferKind::XShape) { - self.collect_candidates(db, kernel, quant, quant, "xshape", &mut candidates)?; + self.collect_candidates(db, table, quant, quant, "xshape", &mut candidates)?; } if candidates.is_empty() && policy.contains(TransferKind::XQuant) { let qp = quant.mapping(); - for name in db.moe.available_quants(kernel)? { + for name in self.table_available_quants(db, table)? { let Some(sibling) = moe_quant_from_name(&name) else { continue; }; @@ -464,7 +504,7 @@ impl MoeOp { { continue; } - self.collect_candidates(db, kernel, sibling, quant, "xquant", &mut candidates)?; + self.collect_candidates(db, table, sibling, quant, "xquant", &mut candidates)?; } } if let Some(reference) = @@ -480,9 +520,8 @@ impl MoeOp { if grid.as_deref().is_none_or(UtilGrid::is_empty) && policy.contains(TransferKind::XProfile) { - let table_quants: Vec = db - .moe - .available_quants(kernel)? + let table_quants: Vec = self + .table_available_quants(db, table)? .iter() .filter_map(|name| moe_quant_from_name(name)) .collect(); @@ -490,7 +529,7 @@ impl MoeOp { let mut candidates: Vec = Vec::new(); self.collect_candidates( db, - kernel, + table, ref_quant, ref_quant, "xprofile", @@ -512,22 +551,59 @@ impl MoeOp { let query = [num_tokens as f64]; let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), util_scale)?; + // Tier fired = the reference grid's transfer kind (xshape / xquant / + // xprofile), or own-shape "empirical" when no borrow happened — + // Python `prov` at moe.py:447/534/569 passed to estimate() at :571. + db.note_provenance( + grid.as_deref() + .and_then(|g| g.reference_provenance) + .and_then(util_empirical::ProvenanceTier::from_tag) + .unwrap_or(util_empirical::ProvenanceTier::Empirical), + ); Ok(latency) } - /// This op's own-slice token curve on the chosen kernel table. - fn slice_points(&self, db: &PerfDatabase, kernel: MoeKernel) -> Result, AicError> { - db.moe.slice_points( - kernel, - self.quant_mode.name(), - &self.workload_distribution, - self.topk, - self.num_experts, - self.hidden_size, - self.inter_size, - self.moe_tp_size, - self.moe_ep_size, - ) + /// This op's own-slice token curve on the selected table. + fn slice_points(&self, db: &PerfDatabase, table: MoeTableSel) -> Result, AicError> { + match table { + MoeTableSel::Standard | MoeTableSel::LowLatency => db.moe.slice_points( + moe_kernel(table), + self.quant_mode.name(), + &self.workload_distribution, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + ), + MoeTableSel::Wideep { is_context } => db.wideep.moe_slice_points( + is_context, + self.quant_mode.name(), + &self.workload_distribution, + self.topk, + self.num_experts, + self.hidden_size, + self.inter_size, + self.moe_tp_size, + self.moe_ep_size, + ), + } + } + + /// Distinct quant names of the selected table, in first-seen (file row) + /// order — Python's `for q in moe_table` dict-insertion iteration. + fn table_available_quants( + &self, + db: &PerfDatabase, + table: MoeTableSel, + ) -> Result, AicError> { + match table { + MoeTableSel::Standard | MoeTableSel::LowLatency => { + db.moe.available_quants(moe_kernel(table)) + } + MoeTableSel::Wideep { is_context } => db.wideep.moe_available_quants(is_context), + } } /// Enumerate `source_quant`'s collected sibling slices (same table, @@ -538,19 +614,29 @@ impl MoeOp { fn collect_candidates( &self, db: &PerfDatabase, - kernel: MoeKernel, + table: MoeTableSel, source_quant: MoeQuantMode, sol_quant: MoeQuantMode, provenance: &'static str, out: &mut Vec, ) -> Result<(), AicError> { - let slices = match db.moe.sibling_slices( - kernel, - source_quant.name(), - &self.workload_distribution, - self.moe_tp_size, - self.moe_ep_size, - ) { + let slices = match table { + MoeTableSel::Standard | MoeTableSel::LowLatency => db.moe.sibling_slices( + moe_kernel(table), + source_quant.name(), + &self.workload_distribution, + self.moe_tp_size, + self.moe_ep_size, + ), + MoeTableSel::Wideep { is_context } => db.wideep.moe_sibling_slices( + is_context, + source_quant.name(), + &self.workload_distribution, + self.moe_tp_size, + self.moe_ep_size, + ), + }; + let slices = match slices { Ok(slices) => slices, Err(err) if err.is_missing_perf_data() => return Ok(()), Err(err) => return Err(err), @@ -814,6 +900,8 @@ mod tests { assert_oracle(&r333, 0.19184494219320924, Source::Empirical, "own_emp_t333"); let r96 = op.query(&db, 96).expect("own-shape empirical t=96"); assert_oracle(&r96, 0.13852159976959227, Source::Empirical, "own_emp_t96"); + // Python capture: {"empirical"} (own-shape grid, no borrow). + assert_eq!(db.worst_provenance(), util_empirical::ProvenanceTier::Empirical); } /// HYBRID with data present must stay on silicon (exact hit and in-range @@ -838,6 +926,8 @@ mod tests { let op = qwen3_op(MoeQuantMode::W4a16Mxfp4Cutlass); let r = op.query(&db, 96).expect("xquant transfer"); assert_oracle(&r, 0.329638409614563, Source::Empirical, "xquant_t96"); + // Python capture: {"xquant"} (reference grid's tier, moe.py:534). + assert_eq!(db.worst_provenance(), util_empirical::ProvenanceTier::XQuant); } /// XPROFILE tier: `w4afp8` ((0.5, 2)) has no same-profile sibling in the @@ -849,6 +939,8 @@ mod tests { let op = qwen3_op(MoeQuantMode::W4afp8); let r = op.query(&db, 96).expect("xprofile transfer"); assert_oracle(&r, 0.13701972961425785, Source::Empirical, "xprofile_t96"); + // Python capture: {"xprofile"} (tier-3 borrow, moe.py:569). + assert_eq!(db.worst_provenance(), util_empirical::ProvenanceTier::XProfile); } /// XSHAPE tier: same quant (bfloat16), uncollected inter_size 1600 → @@ -862,6 +954,8 @@ mod tests { let db = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, TransferPolicy::ALL); let r = op.query(&db, 96).expect("xshape transfer"); assert_oracle(&r, 0.14427836344943168, Source::Empirical, "xshape_t96"); + // Python capture: {"xshape"} (tier-1 borrow, moe.py:534). + assert_eq!(db.worst_provenance(), util_empirical::ProvenanceTier::XShape); let conservative = b200_vllm_db().with_mode(crate::common::enums::DatabaseMode::Hybrid, CONSERVATIVE); let rc = op.query(&conservative, 96).expect("xshape under conservative policy"); @@ -937,6 +1031,146 @@ mod tests { assert_oracle(&xshape, 0.05842286435922407, Source::Empirical, "nvfp4_xshape_t100"); } + /// XPROFILE tie-break follows FILE-ROW quant order, not sorted order: + /// b200/vllm/0.24.0 lists `fp8_block` before `fp8` (both profile (1,2), + /// distance 0.5 from w4afp8), so Python's stable sort borrows the + /// `fp8_block` util curve. Oracles: + /// + /// ```text + /// db = perf_database.get_database_view("b200_sxm", "vllm", "0.24.0", + /// allow_missing_data=True, database_mode="EMPIRICAL", shared_layer=False) + /// float(MoE._query_moe_table(db, num_tokens=..., hidden_size=5120, + /// inter_size=8192, topk=1, num_experts=16, moe_tp_size=1, + /// moe_ep_size=1, quant_mode=common.MoEQuantMode.w4afp8, + /// workload_distribution="power_law_1.01", + /// database_mode=common.DatabaseMode.EMPIRICAL)) + /// ``` + /// + /// The fp8-referenced value (the old sorted-name order) is + /// 0.42024958928426115 at t=96 — a live ~13% divergence this pins. + #[test] + fn moe_xprofile_tie_break_follows_file_order() { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..").join("src/aiconfigurator/systems"); + let db = PerfDatabase::load(&root, "b200_sxm", "vllm", "0.24.0") + .expect("db loads") + .with_mode(crate::common::enums::DatabaseMode::Empirical, TransferPolicy::ALL); + let op = MoeOp { + name: "moe".into(), + scale_factor: 1.0, + hidden_size: 5120, + inter_size: 8192, + topk: 1, + num_experts: 16, + moe_tp_size: 1, + moe_ep_size: 1, + quant_mode: MoeQuantMode::W4afp8, + workload_distribution: "power_law_1.01".into(), + attention_dp_size: 1, + is_gated: true, + moe_backend: None, + enable_eplb: false, + is_context: false, + }; + let r96 = op.query(&db, 96).expect("xprofile tie t=96"); + assert_oracle(&r96, 0.47411200205485027, Source::Empirical, "xprofile_tie_t96"); + let r512 = op.query(&db, 512).expect("xprofile tie t=512"); + assert_oracle(&r512, 0.8249173482259117, Source::Empirical, "xprofile_tie_t512"); + } + + /// SGLang deepep op used by the routing tests below: the h200 sglang + /// 0.5.10 wideep context/generation MoE tables cover the DSv3 expert + /// shape (7168, 2048, topk 8, experts 256) at tp=1/ep=8 under + /// power_law_0.8 — while the REGULAR h200 moe table also carries + /// fp8_block, so mis-routing to it yields a value, not an error. + fn h200_deepep_op(is_context: bool) -> MoeOp { + MoeOp { + name: "moe".into(), + scale_factor: 1.0, + hidden_size: 7168, + inter_size: 2048, + topk: 8, + num_experts: 256, + moe_tp_size: 1, + moe_ep_size: 8, + quant_mode: MoeQuantMode::Fp8Block, + workload_distribution: "power_law_0.8".into(), + attention_dp_size: 1, + is_gated: true, + moe_backend: Some("deepep_moe".into()), + enable_eplb: false, + is_context, + } + } + + fn h200_sglang_db(mode: crate::common::enums::DatabaseMode) -> PerfDatabase { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..").join("src/aiconfigurator/systems"); + PerfDatabase::load(&root, "h200_sxm", "sglang", "0.5.10") + .expect("db loads") + .with_mode(mode, TransferPolicy::ALL) + } + + /// EMPIRICAL under SGLang deepep calibrates from the WIDEEP + /// context/generation MoE tables, not the regular one — Python + /// `_moe_table()` (`operations/moe.py:364-397`) routes the util grid by + /// `moe_backend == "deepep_moe"` + `is_context`. Oracles: + /// + /// ```text + /// db = perf_database.get_database_view("h200_sxm", "sglang", "0.5.10", + /// allow_missing_data=True, database_mode="EMPIRICAL", shared_layer=False) + /// float(MoE._query_moe_table(db, num_tokens=..., hidden_size=7168, + /// inter_size=2048, topk=8, num_experts=256, moe_tp_size=1, + /// moe_ep_size=8, quant_mode=common.MoEQuantMode.fp8_block, + /// workload_distribution="power_law_0.8", is_context=..., + /// moe_backend="deepep_moe", database_mode=common.DatabaseMode.EMPIRICAL)) + /// ``` + /// + /// ctx t=200000 sits beyond the collected range (max 131072) so the + /// util-hold anchors on the MoE roofline through the wideep grid. + #[test] + fn moe_empirical_deepep_routes_to_wideep_tables() { + let db = h200_sglang_db(crate::common::enums::DatabaseMode::Empirical); + let ctx = h200_deepep_op(true); + let r = ctx.query(&db, 300).expect("deepep ctx t=300"); + assert_oracle(&r, 0.6444098182832491, Source::Empirical, "deepep_ctx_t300"); + let r = ctx.query(&db, 200000).expect("deepep ctx t=200000"); + assert_oracle(&r, 21.3889914448373, Source::Empirical, "deepep_ctx_t200000"); + let gen = h200_deepep_op(false); + let r = gen.query(&db, 100).expect("deepep gen t=100"); + assert_oracle(&r, 0.34094198365735795, Source::Empirical, "deepep_gen_t100"); + let r = gen.query(&db, 3000).expect("deepep gen t=3000"); + assert_oracle(&r, 0.4024570594575049, Source::Empirical, "deepep_gen_t3000"); + // Python capture: {"empirical"} (own-slice wideep calibration). + assert_eq!(db.worst_provenance(), util_empirical::ProvenanceTier::Empirical); + } + + /// The EPLB 0.8 prefill token correction applies INSIDE the silicon + /// path only (Python moe.py:684); the empirical estimate uses RAW + /// tokens (moe.py:637-647, 803-813). Python oracles (same call shape as + /// `moe_empirical_deepep_routes_to_wideep_tables`, `enable_eplb=True`): + /// SILICON eplb-on t=160 = SILICON eplb-off t=128 = 0.6220973747117179 + /// (int(160*0.8) = 128, a collected point); EMPIRICAL eplb-on t=300 = + /// eplb-off t=300 = 0.6444098182832491. + #[test] + fn moe_eplb_correction_scoped_to_silicon_only() { + let mut eplb_op = h200_deepep_op(true); + eplb_op.enable_eplb = true; + + let silicon = h200_sglang_db(crate::common::enums::DatabaseMode::Hybrid); + let corrected = eplb_op.query(&silicon, 160).expect("eplb-on silicon t=160"); + assert_oracle(&corrected, 0.6220973747117179, Source::Silicon, "eplb_sil_t160"); + let baseline = h200_deepep_op(true).query(&silicon, 128).expect("eplb-off silicon t=128"); + assert!( + (corrected.latency_ms - baseline.latency_ms).abs() < 1e-12, + "silicon eplb-on(160) ({}) must equal eplb-off(128) ({})", + corrected.latency_ms, + baseline.latency_ms + ); + + let empirical = h200_sglang_db(crate::common::enums::DatabaseMode::Empirical); + let raw = eplb_op.query(&empirical, 300).expect("eplb-on empirical t=300"); + assert_oracle(&raw, 0.6444098182832491, Source::Empirical, "eplb_emp_t300"); + } + /// With attention-dp, all dp ranks' tokens funnel into the shared expert /// pool: query(dp=4, t) must equal query(dp=1, 4t). Dropping the /// multiplier under-predicted MoE latency ~4.7x on dp=8 DeepSeek configs diff --git a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs index 655e8f23e..a205800a6 100644 --- a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs +++ b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs @@ -450,7 +450,7 @@ impl MoEDispatchOp { // quant-compressed volumes on the pre side // elif tp>1 && reduce_results: custom_allreduce(num_gpus) // (the NVL72 tp>4 -> NCCL all_reduce reroute lives in - // the DB-level `query_custom_allreduce_scaled`) + // the operator-level `query_custom_allreduce_table`) // else 0 // sm != 100 (checks tp FIRST, mirroring Python): // if tp>1: custom_allreduce(num_gpus, volume) @@ -804,6 +804,8 @@ fn alltoall_empirical( })?; let query = [num_tokens as f64]; let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + // Own-shape util fired (Python estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } diff --git a/rust/aiconfigurator-core/src/operators/msa.rs b/rust/aiconfigurator-core/src/operators/msa.rs index bc7f80822..abea60c79 100644 --- a/rust/aiconfigurator-core/src/operators/msa.rs +++ b/rust/aiconfigurator-core/src/operators/msa.rs @@ -84,6 +84,8 @@ impl MsaModuleOp { match util { Some(util) if util > 0.0 => { let latency = sol / (util * self.dsa_scale_k); + // Cross-op transfer from DSA (Python msa.py:297 "xop"). + db.note_provenance(crate::operators::util_empirical::ProvenanceTier::XOp); Ok(PerformanceResult::new( latency * self.scale_factor, Source::Empirical, @@ -128,6 +130,8 @@ impl MsaModuleOp { match util { Some(util) if util > 0.0 => { let latency = sol / (util * self.dsa_scale_k); + // Cross-op transfer from DSA (Python msa.py:335 "xop"). + db.note_provenance(crate::operators::util_empirical::ProvenanceTier::XOp); Ok(PerformanceResult::new( latency * self.scale_factor, Source::Empirical, @@ -405,6 +409,7 @@ mod tests { let db = db(backend, version); let op = msa_op(); for (b, s, prefix, is_context, expected) in anchors { + db.reset_provenance(); let result = if is_context { op.query_context(&db, b, s, prefix) } else { @@ -413,6 +418,12 @@ mod tests { .unwrap_or_else(|e| panic!("{backend} b={b} s={s}: {e}")); approx(result.latency_ms, expected); assert_eq!(result.source, Source::Empirical); + // The xop tier must be recorded on the db accumulator + // (Python msa.py:297/335 note_provenance("xop")). + assert_eq!( + db.worst_provenance(), + crate::operators::util_empirical::ProvenanceTier::XOp + ); } } } diff --git a/rust/aiconfigurator-core/src/operators/op.rs b/rust/aiconfigurator-core/src/operators/op.rs index 913da8e1e..1d37b03cd 100644 --- a/rust/aiconfigurator-core/src/operators/op.rs +++ b/rust/aiconfigurator-core/src/operators/op.rs @@ -18,7 +18,8 @@ use serde::{Deserialize, Serialize}; use crate::common::error::AicError; use crate::operators::{ - ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4ModuleOp, + ContextAttentionOp, ContextMlaOp, CustomAllReduceOp, DsaModuleOp, Dsv4MegaMoeOp, + Dsv4ModuleOp, ElementwiseOp, EmbeddingOp, EncoderAttentionOp, GdnOp, GemmOp, GenerationAttentionOp, GenerationMlaOp, Mamba2Op, MhcModuleOp, MlaBmmOp, MlaModuleOp, MoEDispatchOp, MoeOp, MsaModuleOp, TrtllmWideEpMoEDispatchOp, @@ -144,6 +145,16 @@ pub enum Op { /// transitional state where some systems have module-level profiling /// data and others still ship per-kernel granular data. Fallback(FallbackOp), + /// SGLang DeepSeek-V4 MegaMoE routed module (Python + /// `DeepSeekV4MegaMoEModule`): one variant for both phases — the op's + /// `is_context` field selects the phase inside the unified table. + /// Measured-SILICON-only; see `operators/dsv4.rs::Dsv4MegaMoeOp`. + /// + /// APPENDED after `Fallback` on purpose: bincode enum indices are + /// positional, so appending does not shift existing variants and + /// `ENGINE_SPEC_SCHEMA_VERSION` stays unchanged. Do NOT insert new + /// variants mid-enum. + Dsv4MegaMoe(Dsv4MegaMoeOp), } /// Inline-defined here (rather than a sibling module under `operators/`) @@ -223,6 +234,7 @@ impl Op { Op::WideEpMoeDispatch(o) => &o.name, Op::Overlap(o) => &o.name, Op::Fallback(o) => &o.name, + Op::Dsv4MegaMoe(o) => &o.name, } } @@ -376,6 +388,10 @@ impl Op { Err(other) => Err(other), } } + // Rank-LOCAL token count, like Moe/MoeDispatch (Python passes the + // same `x`); the megamoe table is indexed by local-rank tokens + // and the op must NOT re-multiply by attention_dp_size. + Op::Dsv4MegaMoe(op) => op.query(db, ctx.num_tokens), } } } diff --git a/rust/aiconfigurator-core/src/operators/util_empirical.rs b/rust/aiconfigurator-core/src/operators/util_empirical.rs index 3b59c441b..32ba32df0 100644 --- a/rust/aiconfigurator-core/src/operators/util_empirical.rs +++ b/rust/aiconfigurator-core/src/operators/util_empirical.rs @@ -34,6 +34,75 @@ use std::sync::{Arc, Mutex}; use crate::common::error::AicError; +/// Empirical provenance tiers, mirroring Python's `PROVENANCE_ORDER` +/// (`sdk/operations/util_empirical.py`): ordered by DECREASING confidence, +/// so the max rank fired during a run is the run's effective data source +/// (Python `worst_provenance`). `Silicon` is the default when nothing fired; +/// operators never note it. +/// +/// The accumulation cell lives on `PerfDatabase` (shared across mode views); +/// operators call `PerfDatabase::note_provenance` at the same sites Python +/// calls `note_provenance` — after a successful `estimate` with the tier the +/// call site knows (Python passes it as `estimate`'s `provenance` param). +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +#[repr(u8)] +pub enum ProvenanceTier { + /// Pure silicon table data (default; never recorded by operators). + Silicon = 0, + /// Own-shape util (no transfer). + Empirical = 1, + /// Cross-shape, same quant. + XShape = 2, + /// Cross-quant, same profile. + XQuant = 3, + /// Cross-quant, cross profile. + XProfile = 4, + /// Cross-op (borrowed a different op's util). + XOp = 5, +} + +impl ProvenanceTier { + /// The Python tag string (`PROVENANCE_ORDER` spelling) for this tier. + pub fn as_str(self) -> &'static str { + match self { + ProvenanceTier::Silicon => "silicon", + ProvenanceTier::Empirical => "empirical", + ProvenanceTier::XShape => "xshape", + ProvenanceTier::XQuant => "xquant", + ProvenanceTier::XProfile => "xprofile", + ProvenanceTier::XOp => "xop", + } + } + + /// Inverse of [`Self::as_str`] for call sites that carry the Python tag + /// string (e.g. a reference grid's `reference_provenance`). Unknown tags + /// yield `None` so callers choose their own default tier. + pub fn from_tag(tag: &str) -> Option { + match tag { + "silicon" => Some(ProvenanceTier::Silicon), + "empirical" => Some(ProvenanceTier::Empirical), + "xshape" => Some(ProvenanceTier::XShape), + "xquant" => Some(ProvenanceTier::XQuant), + "xprofile" => Some(ProvenanceTier::XProfile), + "xop" => Some(ProvenanceTier::XOp), + _ => None, + } + } + + /// Inverse of `tier as u8` for reading the accumulation cell back. + /// Out-of-range ranks clamp to the least-confident tier. + pub fn from_rank(rank: u8) -> ProvenanceTier { + match rank { + 0 => ProvenanceTier::Silicon, + 1 => ProvenanceTier::Empirical, + 2 => ProvenanceTier::XShape, + 3 => ProvenanceTier::XQuant, + 4 => ProvenanceTier::XProfile, + _ => ProvenanceTier::XOp, + } + } +} + /// One collected calibration point: continuous-axis coordinates plus the /// positive effective calibration factor `util = SOL / measured`. #[derive(Clone, Debug, PartialEq)] @@ -599,6 +668,29 @@ mod tests { assert!(recovered.unwrap().is_none()); } + #[test] + fn provenance_tier_mirrors_python_provenance_order() { + // Rank order == PROVENANCE_ORDER index; tag strings match Python. + let tiers = [ + (ProvenanceTier::Silicon, 0u8, "silicon"), + (ProvenanceTier::Empirical, 1, "empirical"), + (ProvenanceTier::XShape, 2, "xshape"), + (ProvenanceTier::XQuant, 3, "xquant"), + (ProvenanceTier::XProfile, 4, "xprofile"), + (ProvenanceTier::XOp, 5, "xop"), + ]; + for (tier, rank, tag) in tiers { + assert_eq!(tier as u8, rank); + assert_eq!(tier.as_str(), tag); + assert_eq!(ProvenanceTier::from_rank(rank), tier); + assert_eq!(ProvenanceTier::from_tag(tag), Some(tier)); + } + assert_eq!(ProvenanceTier::from_tag("unknown"), None); + // max-rank == worst_provenance semantics; overflow clamps to XOp. + assert!(ProvenanceTier::XOp > ProvenanceTier::Empirical); + assert_eq!(ProvenanceTier::from_rank(200), ProvenanceTier::XOp); + } + #[test] fn zero_aware_delta_keeps_zeroes_and_freezes_utilization() { let lookup = ZeroAwareDeltaLookup::new(vec![ diff --git a/rust/aiconfigurator-core/src/operators/wideep_mla.rs b/rust/aiconfigurator-core/src/operators/wideep_mla.rs index fa1a9bd2b..d631d466e 100644 --- a/rust/aiconfigurator-core/src/operators/wideep_mla.rs +++ b/rust/aiconfigurator-core/src/operators/wideep_mla.rs @@ -35,14 +35,18 @@ fn prefix_correction(full_s: u32, prefix: u32) -> f64 { (f * f - p * p) / (f * f) } -/// Python's `get_silicon` whitelists the attention backend before slicing +/// Python whitelists the attention backend INSIDE `get_silicon` only /// (`mla.py:1191-1192` generation, `:1459-1460` context: -/// `if attn_backend not in {"flashinfer", "fa3"}: raise ValueError`). -/// Mirror it so both engines error symmetrically on unsupported backends -/// instead of Rust answering from whatever `kernel_source` slice exists. +/// `if attn_backend not in {"flashinfer", "fa3"}: raise ValueError`), AFTER +/// the table's `raise_if_not_loaded()`. The EMPIRICAL path slices whatever +/// backend is requested (`mla.py:1147, 1410`) — e.g. b200's `trtllm_mla` +/// wideep slices calibrate fine there — so the check must not run for +/// EMPIRICAL or for the HYBRID fallback. The error is a Python `ValueError`, +/// deliberately NOT a missing-data signal: under HYBRID it propagates +/// instead of triggering the empirical fallback. fn check_attn_backend(attn_backend: &str) -> Result<(), AicError> { if attn_backend != "flashinfer" && attn_backend != "fa3" { - return Err(AicError::PerfDatabase(format!( + return Err(AicError::InvalidEngineConfig(format!( "Unsupported attention backend: {attn_backend}" ))); } @@ -93,7 +97,6 @@ impl WideEpContextMlaOp { isl: u32, prefix: u32, ) -> Result { - check_attn_backend(&self.attn_backend)?; // ctx(s, pfx): the un-sharded wideep context-MLA query for a sequence // chunk of length `s` at prefix `pfx` (mode dispatch handles prefix // correction for silicon and the prefix-aware SOL for empirical). @@ -251,6 +254,11 @@ mod tests { ); assert_eq!(source, Source::Empirical); } + // Python capture: {"empirical"} (own-slice util, mla.py:1431). + assert_eq!( + db.worst_provenance(), + util_empirical::ProvenanceTier::Empirical + ); // HYBRID on a kv slice with no data (bfloat16; h200's wideep tables // are fp8-KV only) -> terminal EmpiricalNotImplemented miss. @@ -301,6 +309,11 @@ mod tests { ); assert_eq!(source, Source::Empirical); } + // Python capture: {"empirical"} (own-slice util, mla.py:1162). + assert_eq!( + db.worst_provenance(), + util_empirical::ProvenanceTier::Empirical + ); // HYBRID on a kv slice with no data (bfloat16) -> terminal miss. db.database_mode = crate::common::enums::DatabaseMode::Hybrid; @@ -309,10 +322,12 @@ mod tests { assert!(matches!(result, Err(AicError::EmpiricalNotImplemented(_))), "got {result:?}"); } - /// Python whitelists attn_backend in {flashinfer, fa3} before slicing - /// (mla.py:1191-1192, 1459-1460) — an out-of-whitelist backend must error - /// here too, even when a matching kernel_source slice exists in the data - /// (b200's wideep tables carry kernel_source=trtllm_mla). + /// Python whitelists attn_backend in {flashinfer, fa3} inside + /// `get_silicon` only (mla.py:1191-1192, 1459-1460) — SILICON (the + /// db default here) errors on an out-of-whitelist backend even when a + /// matching kernel_source slice exists in the data, and under HYBRID + /// with a loaded table the ValueError propagates (a config error, not + /// a missing-data fallback trigger). #[test] fn attn_backend_whitelist_mirrors_python() { let db = h200_sglang_db(); @@ -334,6 +349,101 @@ mod tests { fa3.attn_backend = "fa3".to_string(); assert!(fa3.query(&db, 1, 1024, 0).is_ok()); + // HYBRID + loaded table + bad backend: the whitelist error is NOT a + // missing-data signal, so it propagates instead of falling back to + // the empirical estimate (Python: plain ValueError is not in + // `_MISSING_SILICON_DATA_EXCEPTIONS`). + let mut hybrid = h200_sglang_db(); + hybrid.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let mut bad = op(1); + bad.attn_backend = "trtllm_mla".to_string(); + let result = bad.query(&hybrid, 1, 1024, 0); + assert!( + matches!(result, Err(AicError::InvalidEngineConfig(_))), + "HYBRID must propagate the whitelist error, got {result:?}" + ); + } + + /// EMPIRICAL mode never runs the whitelist: it slices whatever backend + /// is requested (mla.py:1147, 1410) — b200's wideep MLA tables are + /// `trtllm_mla`-only and calibrate fine. Oracles: + /// + /// ```text + /// db = perf_database.get_database_view("b200_sxm", "sglang", "0.5.10", + /// allow_missing_data=True, database_mode="EMPIRICAL", shared_layer=False) + /// float(WideEPContextMLA._query_wideep_context_mla_table(db, b, s, + /// prefix, tp_size, KVCacheQuantMode.fp8, FMHAQuantMode.fp8_block, + /// "trtllm_mla", DatabaseMode.EMPIRICAL)) # + generation variant + /// ``` + #[test] + fn empirical_estimates_from_non_whitelisted_backend_slice() { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + let mut db = PerfDatabase::load(&root, "b200_sxm", "sglang", "0.5.10").expect("db loads"); + db.database_mode = crate::common::enums::DatabaseMode::Empirical; + + // (b, s, prefix, num_heads = 128 // tp_size, expected) + let ctx_cases: &[(u32, u32, u32, u32, f64)] = &[ + (4, 4096, 0, 128, 5.1478), + (2, 3000, 1000, 16, 0.5668068670651026), + ]; + for &(b, s, prefix, n, expected) in ctx_cases { + let (latency, source) = query_wideep_context_mla_table( + &db, + b, + s, + prefix, + n, + KvCacheQuantMode::Fp8, + FmhaQuantMode::Fp8Block, + "trtllm_mla", + ) + .expect("trtllm_mla slice estimates"); + assert_close(latency, expected, &format!("trtllm_mla ctx(b={b}, s={s}, pfx={prefix})")); + assert_eq!(source, Source::Empirical); + } + + let (latency, source) = query_wideep_generation_mla_table( + &db, + 3, + 5000, + 16, + KvCacheQuantMode::Fp8, + "trtllm_mla", + ) + .expect("trtllm_mla gen slice estimates"); + assert_close(latency, 0.056915046282426024, "trtllm_mla gen(b=3, s=5000)"); + assert_eq!(source, Source::Empirical); + } + + /// HYBRID with the table NOT loaded (vLLM DBs ship no wideep MLA files): + /// the load probe is the typed miss that precedes the whitelist — the + /// fallback runs the empirical path with the requested backend and + /// surfaces the terminal EmpiricalNotImplemented (never the whitelist + /// config error). Mirrors Python get_silicon's `raise_if_not_loaded()` + /// -> whitelist ordering (mla.py:1449-1461). + #[test] + fn hybrid_unloaded_table_falls_to_empirical_before_whitelist() { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + let mut db = PerfDatabase::load(&root, "b200_sxm", "vllm", "0.19.0").expect("db loads"); + db.database_mode = crate::common::enums::DatabaseMode::Hybrid; + let result = query_wideep_context_mla_table( + &db, + 1, + 1024, + 0, + 128, + KvCacheQuantMode::Fp8, + FmhaQuantMode::Fp8Block, + "trtllm_mla", + ); + assert!( + matches!(result, Err(AicError::EmpiricalNotImplemented(_))), + "expected the typed empirical miss, got {result:?}" + ); } } @@ -373,9 +483,6 @@ impl WideEpGenerationMlaOp { batch_size: u32, s: u32, ) -> Result { - // Python whitelists the attention backend BEFORE any table/mode - // dispatch (mla.py:1459-1460) — a bad backend errors in every mode. - check_attn_backend(&self.attn_backend)?; let (latency, source) = query_wideep_generation_mla_table( db, batch_size, @@ -424,6 +531,11 @@ fn query_wideep_context_mla_table( attn_backend: &str, ) -> Result<(f64, Source), AicError> { let silicon = || -> Result { + // Python get_silicon order (mla.py:1449-1461): raise_if_not_loaded + // (typed miss -> HYBRID may still estimate) BEFORE the attn-backend + // whitelist ValueError (propagates, never falls back). + db.wideep_mla.ensure_context_loaded()?; + check_attn_backend(attn_backend)?; let full_s = s + prefix; let raw = db.wideep_mla.query_context( b, @@ -500,6 +612,8 @@ fn wideep_context_mla_empirical( ); let query = [num_heads as f64, (s + prefix) as f64, b as f64]; let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + // Own-slice util fired (Python mla.py:1431 estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -516,29 +630,29 @@ fn query_wideep_generation_mla_table( kv_quant: KvCacheQuantMode, attn_backend: &str, ) -> Result<(f64, Source), AicError> { + let silicon = || -> Result { + // Python get_silicon order (mla.py:1188-1192): raise_if_not_loaded + // (typed miss -> HYBRID may still estimate) BEFORE the attn-backend + // whitelist ValueError (propagates, never falls back). + db.wideep_mla.ensure_generation_loaded()?; + check_attn_backend(attn_backend)?; + db.wideep_mla + .query_generation(b, s, num_heads, kv_quant, attn_backend) + }; match db.database_mode { DatabaseMode::Empirical => Ok(( wideep_generation_mla_empirical(db, b, s, num_heads, kv_quant, attn_backend)?, Source::Empirical, )), - DatabaseMode::Hybrid => { - match db - .wideep_mla - .query_generation(b, s, num_heads, kv_quant, attn_backend) - { - Ok(latency) => Ok((latency, Source::Silicon)), - Err(err) if err.is_missing_perf_data() => Ok(( - wideep_generation_mla_empirical(db, b, s, num_heads, kv_quant, attn_backend)?, - Source::Empirical, - )), - Err(err) => Err(err), - } - } - _ => Ok(( - db.wideep_mla - .query_generation(b, s, num_heads, kv_quant, attn_backend)?, - Source::Silicon, - )), + DatabaseMode::Hybrid => match silicon() { + Ok(latency) => Ok((latency, Source::Silicon)), + Err(err) if err.is_missing_perf_data() => Ok(( + wideep_generation_mla_empirical(db, b, s, num_heads, kv_quant, attn_backend)?, + Source::Empirical, + )), + Err(err) => Err(err), + }, + _ => Ok((silicon()?, Source::Silicon)), } } @@ -577,5 +691,7 @@ fn wideep_generation_mla_empirical( let sol_query = wideep_generation_mla_sol_ms(spec, fmha_quant, num_heads as f64, b as f64, s as f64); let query = [num_heads as f64, b as f64, s as f64]; let (latency, _) = util_empirical::estimate(sol_query, &query, grid.as_deref(), 1.0)?; + // Own-slice util fired (Python mla.py:1162 estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } diff --git a/rust/aiconfigurator-core/src/operators/wideep_moe.rs b/rust/aiconfigurator-core/src/operators/wideep_moe.rs index d76a1d938..56b0c0adc 100644 --- a/rust/aiconfigurator-core/src/operators/wideep_moe.rs +++ b/rust/aiconfigurator-core/src/operators/wideep_moe.rs @@ -230,6 +230,8 @@ impl WideEpMoeOp { })?; let query = [num_tokens as f64]; let (latency, _) = util_empirical::estimate(sol_time, &query, grid.as_deref(), 1.0)?; + // Own-slice util fired (Python moe.py:1697 estimate()'s default tier). + db.note_provenance(util_empirical::ProvenanceTier::Empirical); Ok(latency) } @@ -339,6 +341,11 @@ mod tests { assert_oracle(&off, 0.4523388956415573, Source::Empirical, "emp_t333"); let hold = op().query(&db, 100000).expect("beyond-range empirical"); assert_oracle(&hold, 15.171406557783484, Source::Empirical, "emp_t100000"); + // Python capture: {"empirical"} (own-slice util, no transfer ladder). + assert_eq!( + db.worst_provenance(), + util_empirical::ProvenanceTier::Empirical + ); } /// HYBRID with data present stays on silicon; the in-range interpolation @@ -392,4 +399,52 @@ mod tests { } + /// Distribution fallback is FILE-ROW order, not sorted order: gb200's + /// `wideep_moe_perf.parquet` lists `power_law_1.01_eplb` before + /// `power_law_1.01` (b200's happens to be sorted, which masked this). + /// Python's `dists[0]` (dict insertion order — moe.py:1650-1653 + /// empirical, :1755-1765 silicon) therefore answers an uncollected + /// distribution from the EPLB slice. Oracles: + /// + /// ```text + /// db = perf_database.get_database_view("gb200", "trtllm", "1.3.0rc10", + /// allow_missing_data=True, database_mode=..., shared_layer=False) + /// float(TrtLLMWideEPMoE._query_compute_table(db, num_tokens=..., + /// hidden_size=6144, inter_size=2048, topk=8, num_experts=256, + /// num_slots=256, moe_tp_size=1, moe_ep_size=2, + /// quant_mode=common.MoEQuantMode.nvfp4, + /// workload_distribution="balanced", database_mode=...)) + /// ``` + /// + /// nt=64 is a collected point of the EPLB slice (0.34764...; the sorted + /// -first `power_law_1.01` slice holds 0.34145... there, so this anchor + /// fails on lexicographic fallback); nt=333 is an interior lerp, + /// covering both the silicon and the empirical fallback sites. + #[test] + fn wideep_compute_distribution_fallback_uses_file_order() { + let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../..") + .join("src/aiconfigurator/systems"); + let gb200 = |mode: DatabaseMode| { + PerfDatabase::load(&root, "gb200", "trtllm", "1.3.0rc10") + .expect("db loads") + .with_mode(mode, TransferPolicy::ALL) + }; + let mut fallback_op = op(); + fallback_op.workload_distribution = "balanced".into(); + + let hit = fallback_op + .query(&gb200(DatabaseMode::Hybrid), 64) + .expect("silicon dist fallback"); + assert_oracle(&hit, 0.34764800071716306, Source::Silicon, "gb200_dist_fb_t64"); + let interp = fallback_op + .query(&gb200(DatabaseMode::Hybrid), 333) + .expect("silicon dist fallback interp"); + assert_oracle(&interp, 0.42581739127635954, Source::Silicon, "gb200_dist_fb_t333"); + let emp = fallback_op + .query(&gb200(DatabaseMode::Empirical), 333) + .expect("empirical dist fallback"); + assert_oracle(&emp, 0.4259303480537969, Source::Empirical, "gb200_emp_dist_fb_t333"); + } + } diff --git a/rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs b/rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs new file mode 100644 index 000000000..27ca7ff79 --- /dev/null +++ b/rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs @@ -0,0 +1,334 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! DeepSeek-V4 MegaMoE routed-module perf table. +//! +//! `dsv4_megamoe_module_perf.parquet` — the measured SGLang/DeepGEMM MegaMoE +//! routed path (prepared hidden states + top-k tensors -> SGLang pre-dispatch +//! -> `deep_gemm.fp8_fp4_mega_moe` -> routed output scaling). Gate/top-k and +//! shared experts are modeled outside this table. +//! +//! Port of Python `operations/dsv4.py::load_dsv4_megamoe_module_data` + +//! `DeepSeekV4MegaMoEModule._query_megamoe_table`. +//! +//! ## Loading (mirrors the Python loader exactly) +//! +//! - SINGLE primary file only: Python's `load_data` passes one unified path +//! (`_read_filtered_rows()`), NOT the shared-layer source list — the +//! loader even `raise`s `TypeError` on a list. No sibling/cross-version +//! inheritance and no `kernel_source` row filter here. +//! - Row invariants (any violation fails the WHOLE load, mirroring the +//! Python `ValueError`s): `used_cuda_graph` must be true, +//! `includes_gate_topk` must be false, `includes_routed_scale` must be +//! true; `phase` must be `context` | `generation`. +//! - Keying: `[phase][kernel_source][kernel_dtype][moe_dtype][pre_dispatch]` +//! `[source_policy][distribution][topk][num_experts]` +//! `[num_fused_shared_experts][hidden_size][inter_size][moe_tp_size]` +//! `[moe_ep_size][num_tokens]`. Duplicate leaves are a load ERROR +//! (Python `_put_nested`), not last-wins. +//! - `moe_dtype` must name a valid `MoEQuantMode` member (Python +//! `common.MoEQuantMode[row["moe_dtype"]]` KeyErrors otherwise). +//! +//! ## Query (mirrors `_query_megamoe_table`) +//! +//! Strict measured-only table: exact key walk (typed miss on any absent +//! level), then a 1-D `num_tokens` Grid resolution with a LINEAR token proxy +//! SOL (`sol_fn = lambda t: float(t)` — routed-expert work scales ~linearly +//! with tokens at fixed topk/experts/hidden, and util-hold only needs the +//! SOL RATIO). The SILICON/HYBRID-only mode contract lives on the operator +//! (`operators/dsv4.rs::Dsv4MegaMoeOp`), matching Python's split. Latency +//! only — the Rust engine does not track energy. + +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::sync::OnceLock; + +use crate::common::enums::MoeQuantMode; +use crate::common::error::AicError; +use super::moe::query_token_curve; +use crate::perf_database::parquet_loader::PerfReader; + +pub struct Dsv4MegaMoeTable { + /// The single unified perf file (see the module note: the Python loader + /// reads ONE primary path, never the shared-layer source list). + primary_path: PathBuf, + module: OnceLock>, +} + +struct Dsv4MegaMoeGrids { + by_keys: BTreeMap>, +} + +/// Full table key (every level of the Python nested dict except the trailing +/// `num_tokens` axis). `quant` holds the `MoeQuantMode` name — Python keys +/// the enum member, whose name is exactly this string. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct Dsv4MegaMoeKey { + phase: String, + kernel_source: String, + kernel_dtype: String, + quant: String, + pre_dispatch: String, + source_policy: String, + distribution: String, + topk: u32, + num_experts: u32, + num_fused_shared_experts: u32, + hidden_size: u32, + inter_size: u32, + moe_tp_size: u32, + moe_ep_size: u32, +} + +/// Python `common.MoEQuantMode[name]` (member-name lookup). The Rust serde +/// names are the same snake_case strings, so a serde round-trip is an exact +/// mirror of the Python KeyError contract. +fn moe_dtype_from_name(name: &str) -> Option { + serde_json::from_value(serde_json::Value::String(name.to_string())).ok() +} + +impl Dsv4MegaMoeTable { + /// Construct for the given data directory. No I/O. + pub fn new(data_root: PathBuf) -> Self { + Self { + primary_path: data_root.join("dsv4_megamoe_module_perf.parquet"), + module: OnceLock::new(), + } + } + + /// Query the measured MegaMoE routed-module latency (ms) at `num_tokens` + /// (rank-LOCAL token count — callers must NOT pre-multiply by + /// attention_dp_size). Mirrors the table body of Python + /// `_query_megamoe_table`: exact key walk, then the 1-axis Grid engine + /// with the linear token-proxy SOL. + #[allow(clippy::too_many_arguments)] + pub fn query_module( + &self, + num_tokens: u32, + hidden_size: u32, + inter_size: u32, + topk: u32, + num_experts: u32, + moe_tp_size: u32, + moe_ep_size: u32, + quant: MoeQuantMode, + workload_distribution: &str, + is_context: bool, + source_policy: &str, + pre_dispatch: &str, + num_fused_shared_experts: u32, + kernel_source: &str, + kernel_dtype: &str, + ) -> Result { + let grids = self.load_module()?; + let phase = if is_context { "context" } else { "generation" }; + let key = Dsv4MegaMoeKey { + phase: phase.to_string(), + kernel_source: kernel_source.to_string(), + kernel_dtype: kernel_dtype.to_string(), + quant: quant.name().to_string(), + pre_dispatch: pre_dispatch.to_string(), + source_policy: source_policy.to_string(), + distribution: workload_distribution.to_string(), + topk, + num_experts, + num_fused_shared_experts, + hidden_size, + inter_size, + moe_tp_size, + moe_ep_size, + }; + let curve = grids.by_keys.get(&key).ok_or_else(|| { + // Python's KeyError -> PerfDataNotAvailableError message. + AicError::PerfDatabase(format!( + "No DSv4 MegaMoE {phase} module data for kernel_source={kernel_source:?}, \ + kernel_dtype={kernel_dtype:?}, quant_mode={}, pre_dispatch={pre_dispatch:?}, \ + source_policy={source_policy:?}, workload_distribution={workload_distribution:?}, \ + topk={topk}, num_experts={num_experts}, \ + num_fused_shared_experts={num_fused_shared_experts}, hidden_size={hidden_size}, \ + inter_size={inter_size}, moe_tp_size={moe_tp_size}, moe_ep_size={moe_ep_size}.", + quant.name() + )) + })?; + // Python: OpInterpConfig(axes=("num_tokens",), resolver=Grid(), + // sol_fn=lambda t: float(t)) — in-range RAW lerp, boundary util-hold + // beyond the collected range with the linear token proxy. + query_token_curve(curve, f64::from(num_tokens), &|t| t) + } + + fn load_module(&self) -> Result<&Dsv4MegaMoeGrids, AicError> { + let cell = self + .module + .get_or_init(|| load_module_parquet(&self.primary_path)); + cell.as_ref().map_err(clone_err) + } +} + +/// Load the unified MegaMoE module parquet. Mirrors Python +/// `load_dsv4_megamoe_module_data` (see the module doc): required columns, +/// bool invariants, phase validation, duplicate-leaf ERROR. A missing file is +/// a typed miss (Python: `LoadedOpData.raise_if_not_loaded` -> +/// `PerfDataNotAvailableError`). +fn load_module_parquet(path: &PathBuf) -> Result { + if !path.exists() { + return Err(AicError::PerfDatabase(format!( + "DSv4 MegaMoE module data not loaded: perf file not found at {}. This combination \ + of model, system, backend, and backend version is not supported by AIC in SILICON \ + mode.", + path.display() + ))); + } + let reader = PerfReader::open(path)?; + // Required columns (a missing column fails the load, matching Python's + // KeyError / phase ValueError contract). + let phase_col = reader.col("phase")?; + let kernel_dtype_col = reader.col("kernel_dtype")?; + let moe_dtype_col = reader.col("moe_dtype")?; + let pre_dispatch_col = reader.col("pre_dispatch")?; + let source_policy_col = reader.col("source_policy")?; + let distribution_col = reader.col("distribution")?; + let topk_col = reader.col("topk")?; + let num_experts_col = reader.col("num_experts")?; + let hidden_size_col = reader.col("hidden_size")?; + let inter_size_col = reader.col("inter_size")?; + let moe_ep_size_col = reader.col("moe_ep_size")?; + let num_tokens_col = reader.col("num_tokens")?; + let latency_col = reader.col("latency")?; + // Python reads `row["routed_scaling_factor"]` unconditionally into the + // leaf metadata (KeyError if absent), even though the query never + // consumes it — require the column for load parity. + let routed_scaling_col = reader.col("routed_scaling_factor")?; + // Bool invariants: Python's defaults for absent columns are chosen so + // that a missing column ALSO fails the invariant (used_cuda_graph + // default None -> false != true; includes_gate_topk default "true" != + // false; includes_routed_scale default None -> false != true). Required + // columns mirror that fail-on-absence exactly. + let used_cuda_graph_col = reader.col("used_cuda_graph")?; + let includes_gate_topk_col = reader.col("includes_gate_topk")?; + let includes_routed_scale_col = reader.col("includes_routed_scale")?; + // Optional columns with Python defaults. + let kernel_source_col = reader.col_optional("kernel_source"); + let num_fused_shared_col = reader.col_optional("num_fused_shared_experts"); + let moe_tp_size_col = reader.col_optional("moe_tp_size"); + + let mut by_keys: BTreeMap> = BTreeMap::new(); + for row in reader.rows()? { + let row = row?; + for (col, expected, error) in [ + (used_cuda_graph_col, true, "DSv4 MegaMoE perf row was not collected with CUDA Graph"), + ( + includes_gate_topk_col, + false, + "DSv4 MegaMoE perf row includes gate/top-k outside the supported boundary", + ), + ( + includes_routed_scale_col, + true, + "DSv4 MegaMoE perf row does not include SGLang routed output scaling", + ), + ] { + if row.bool(col)? != expected { + return Err(AicError::PerfDatabase(format!("{error}: {}", path.display()))); + } + } + let phase = row.str_owned(phase_col)?; + if phase != "context" && phase != "generation" { + return Err(AicError::PerfDatabase(format!( + "DSv4 MegaMoE perf row has unsupported phase={phase:?}: {}", + path.display() + ))); + } + let moe_dtype = row.str_owned(moe_dtype_col)?; + let Some(quant) = moe_dtype_from_name(&moe_dtype) else { + return Err(AicError::PerfDatabase(format!( + "DSv4 MegaMoE perf row has unknown moe_dtype={moe_dtype:?} at {}", + path.display() + ))); + }; + // routed_scaling_factor: read for the required-column/parse contract + // only (the query never consumes it; energy is not modeled in Rust). + let _ = row.f64(routed_scaling_col)?; + let key = Dsv4MegaMoeKey { + phase, + kernel_source: row + .str_optional(kernel_source_col)? + .map(str::to_string) + .unwrap_or_else(|| "deepgemm_megamoe".to_string()), + kernel_dtype: row.str_owned(kernel_dtype_col)?, + quant: quant.name().to_string(), + pre_dispatch: row.str_owned(pre_dispatch_col)?, + source_policy: row.str_owned(source_policy_col)?, + distribution: row.str_owned(distribution_col)?, + topk: row.u32(topk_col)?, + num_experts: row.u32(num_experts_col)?, + num_fused_shared_experts: row.u32_optional(num_fused_shared_col)?.unwrap_or(0), + hidden_size: row.u32(hidden_size_col)?, + inter_size: row.u32(inter_size_col)?, + moe_tp_size: row.u32_optional(moe_tp_size_col)?.unwrap_or(1), + moe_ep_size: row.u32(moe_ep_size_col)?, + }; + let num_tokens = row.u32(num_tokens_col)?; + let latency = row.f64(latency_col)?; + // Python `_put_nested`: a duplicate leaf is a load ERROR, not + // last-wins. + if by_keys + .entry(key.clone()) + .or_default() + .insert(num_tokens, latency) + .is_some() + { + return Err(AicError::PerfDatabase(format!( + "duplicate DSv4 MegaMoE data row for {} {key:?} num_tokens={num_tokens}", + path.display() + ))); + } + } + Ok(Dsv4MegaMoeGrids { by_keys }) +} + +fn clone_err(err: &AicError) -> AicError { + AicError::PerfDatabase(err.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn moe_dtype_names_mirror_python_member_lookup() { + assert_eq!( + moe_dtype_from_name("w4a8_mxfp4_mxfp8"), + Some(MoeQuantMode::W4a8Mxfp4Mxfp8) + ); + assert_eq!(moe_dtype_from_name("fp8_block"), Some(MoeQuantMode::Fp8Block)); + assert_eq!(moe_dtype_from_name("not_a_dtype"), None); + } + + /// Missing perf file is a typed miss (Python + /// `LoadedOpData.raise_if_not_loaded` -> `PerfDataNotAvailableError`). + #[test] + fn missing_file_is_typed_miss() { + let table = Dsv4MegaMoeTable::new(PathBuf::from("/nonexistent/dir")); + let err = table + .query_module( + 1024, + 7168, + 3072, + 6, + 384, + 1, + 8, + MoeQuantMode::W4a8Mxfp4Mxfp8, + "balanced", + true, + "random", + "sglang_jit", + 0, + "deepgemm_megamoe", + "fp8_fp4", + ) + .unwrap_err(); + assert!(err.is_missing_perf_data(), "got {err:?}"); + assert!(err.to_string().contains("DSv4 MegaMoE module data not loaded")); + } +} diff --git a/rust/aiconfigurator-core/src/perf_database/mod.rs b/rust/aiconfigurator-core/src/perf_database/mod.rs index 669e7b4e6..e0083afc5 100644 --- a/rust/aiconfigurator-core/src/perf_database/mod.rs +++ b/rust/aiconfigurator-core/src/perf_database/mod.rs @@ -11,13 +11,14 @@ //! the WideEP/DeepEP all-to-all variants. use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; use crate::common::enums::{DatabaseMode, TransferPolicy}; use crate::common::error::AicError; use crate::common::system_spec::SystemSpec; use crate::config::{PerfDbSources, PerfSource}; -use crate::operators::util_empirical::{DeltaLookupCache, UtilGridCache}; +use crate::operators::util_empirical::{DeltaLookupCache, ProvenanceTier, UtilGridCache}; /// Resolve the ordered source list for one op-file basename: the Python-supplied /// shared-layer sources when present, else a single primary `data_root/` @@ -56,6 +57,7 @@ pub mod attention; pub mod communication; pub mod dsa; pub mod dsv4; +pub mod dsv4_megamoe; pub mod gemm; mod interpolation; pub mod mhc; @@ -72,6 +74,7 @@ pub use attention::AttentionTable; pub use communication::CommunicationTable; pub use dsa::DsaTable; pub use dsv4::{AttnKind, Dsv4Table}; +pub use dsv4_megamoe::Dsv4MegaMoeTable; pub use gemm::GemmTable; pub use mhc::MhcTable; pub use mla::MlaTable; @@ -97,6 +100,7 @@ pub struct PerfTables { pub communication: CommunicationTable, pub dsa: DsaTable, pub dsv4: Dsv4Table, + pub dsv4_megamoe: Dsv4MegaMoeTable, pub mhc: MhcTable, pub wideep: WideEpTable, pub wideep_mla: WideEpMlaTable, @@ -135,6 +139,14 @@ pub struct PerfDatabase { /// Memo of zero-aware delta lookups (the `compute_scale` empirical /// mechanism); same keying/lifetime rationale as `util_grids`. pub delta_lookups: Arc, + /// Max-rank empirical provenance tier fired since the last + /// [`PerfDatabase::reset_provenance`] (Rust mirror of Python's + /// `capture_provenance` contextvar in `sdk/operations/util_empirical.py`). + /// Shared across mode views (`silicon_view` clones) like `util_grids`, so + /// notes made through a derived view land on the run's accumulator. The + /// engine FFI resets it per `run_static`/step call and reads the worst + /// tier back for the Python bridge / support-matrix HYBRID labelling. + provenance: Arc, } impl std::ops::Deref for PerfDatabase { @@ -222,6 +234,10 @@ impl PerfDatabase { ), dsa: DsaTable::with_sources(data_root.clone(), perf_db_sources), dsv4: Dsv4Table::with_sources(data_root.clone(), perf_db_sources), + // Single-primary by design: the Python MegaMoE loader reads one + // unified path and never the shared-layer source list (see + // `dsv4_megamoe.rs`). + dsv4_megamoe: Dsv4MegaMoeTable::new(data_root.clone()), mhc: MhcTable::with_sources(data_root.clone(), perf_db_sources), wideep: WideEpTable::with_sources(data_root.clone(), perf_db_sources), wideep_mla: WideEpMlaTable::with_sources(data_root.clone(), spec.clone(), perf_db_sources), @@ -241,9 +257,35 @@ impl PerfDatabase { transfer_policy: TransferPolicy::ALL, util_grids: Arc::new(UtilGridCache::new()), delta_lookups: Arc::new(DeltaLookupCache::new()), + provenance: Arc::new(AtomicU8::new(ProvenanceTier::Silicon as u8)), }) } + /// Record that an empirical path of tier `tier` produced a value. + /// Max-rank accumulation, so the cell always holds the run's + /// least-confident (worst) tier — Python's `worst_provenance` semantics. + /// + /// Call sites mirror Python `note_provenance` exactly: after each + /// successful `util_empirical::estimate` with the tier the site knows + /// (own-data "empirical", attention head_size ref grid "xshape", the MoE + /// ladder's `reference_provenance`, communication rank-overflow "xshape", + /// MSA cross-op "xop", ...). + pub fn note_provenance(&self, tier: ProvenanceTier) { + self.provenance.fetch_max(tier as u8, Ordering::Relaxed); + } + + /// Clear the accumulator back to `Silicon` (start of a run). + pub fn reset_provenance(&self) { + self.provenance + .store(ProvenanceTier::Silicon as u8, Ordering::Relaxed); + } + + /// The least-confident tier fired since the last reset; `Silicon` when no + /// empirical path fired (Python `worst_provenance` of an empty capture). + pub fn worst_provenance(&self) -> ProvenanceTier { + ProvenanceTier::from_rank(self.provenance.load(Ordering::Relaxed)) + } + /// Configure the query mode + transfer policy (both immutable per /// database instance afterwards, mirroring Python's configured query /// views). Called by `Engine::from_spec_bytes` with the spec's values. @@ -276,6 +318,7 @@ impl PerfDatabase { transfer_policy: self.transfer_policy, util_grids: Arc::clone(&self.util_grids), delta_lookups: Arc::clone(&self.delta_lookups), + provenance: Arc::clone(&self.provenance), } } } @@ -303,6 +346,28 @@ mod tests { assert!(db.data_root.join("gemm_perf.parquet").is_file()); } + #[test] + fn provenance_cell_accumulates_worst_tier_and_is_shared_with_views() { + let db = PerfDatabase::load(&systems_root(), "b200_sxm", "vllm", "0.19.0") + .expect("b200_sxm/vllm/0.19.0 must load"); + assert_eq!(db.worst_provenance(), ProvenanceTier::Silicon); + + // Max-rank accumulation: a lower tier never overwrites a higher one. + db.note_provenance(ProvenanceTier::XShape); + db.note_provenance(ProvenanceTier::Empirical); + assert_eq!(db.worst_provenance(), ProvenanceTier::XShape); + + // Notes through a derived silicon view land on the same accumulator + // (the MSA xop probe evaluates against `silicon_view`). + let view = db.silicon_view(); + view.note_provenance(ProvenanceTier::XOp); + assert_eq!(db.worst_provenance(), ProvenanceTier::XOp); + + db.reset_provenance(); + assert_eq!(db.worst_provenance(), ProvenanceTier::Silicon); + assert_eq!(view.worst_provenance(), ProvenanceTier::Silicon); + } + #[test] fn load_unknown_version_errors() { match PerfDatabase::load(&systems_root(), "b200_sxm", "vllm", "99.99.99") { diff --git a/rust/aiconfigurator-core/src/perf_database/moe.rs b/rust/aiconfigurator-core/src/perf_database/moe.rs index e2ee883e5..78dea036e 100644 --- a/rust/aiconfigurator-core/src/perf_database/moe.rs +++ b/rust/aiconfigurator-core/src/perf_database/moe.rs @@ -115,6 +115,12 @@ struct LoadedMoeGrids { struct MoeGrids { by_keys: BTreeMap>, + /// Distinct quant names in first-seen (file row) order. Python's + /// transfer ladder iterates the table dict in INSERTION order + /// (`for q in moe_table`), which breaks profile-distance ties by file + /// order — live on shards whose file order differs from sorted order + /// (e.g. b200/vllm/0.24.0 lists `fp8_block` before `fp8`). + quants_in_load_order: Vec, } #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] @@ -360,20 +366,12 @@ impl MoeTable { Ok(slices) } - /// Distinct quant names present in the kernel grid, in sorted (BTreeMap) - /// order. Python iterates the table dict in insertion (file row) order; - /// the difference is observable only through exact ties in the transfer - /// ladder's profile-distance sort / nearest-candidate selection. + /// Distinct quant names present in the kernel grid, in first-seen + /// (file row) order — Python iterates the table dict in insertion + /// order (`for q in moe_table`), and the transfer ladder's stable + /// profile-distance sort breaks ties by that order. pub fn available_quants(&self, kernel: MoeKernel) -> Result, AicError> { - let grids = self.grids_for(kernel)?; - let mut names: Vec = Vec::new(); - for key in grids.by_keys.keys() { - // `MoeKey` sorts by quant first, so consecutive dedup suffices. - if names.last().map(String::as_str) != Some(key.quant.as_str()) { - names.push(key.quant.clone()); - } - } - Ok(names) + Ok(self.grids_for(kernel)?.quants_in_load_order.clone()) } fn grids_for(&self, kernel: MoeKernel) -> Result<&MoeGrids, AicError> { @@ -421,6 +419,8 @@ impl MoeTable { fn load_moe_parquet(sources: &[PerfSource]) -> Result { let mut default_keys: BTreeMap> = BTreeMap::new(); let mut low_latency_keys: BTreeMap> = BTreeMap::new(); + let mut default_quants: Vec = Vec::new(); + let mut low_latency_quants: Vec = Vec::new(); let mut any_source = false; for source in sources { let path = source.path(); @@ -479,11 +479,16 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result moe_tp_size: row.u32(moe_tp_size_col)?, moe_ep_size: row.u32(moe_ep_size_col)?, }; - let target = if kernel_source == "moe_torch_flow_min_latency" { - &mut low_latency_keys + let (target, target_quants) = if kernel_source == "moe_torch_flow_min_latency" { + (&mut low_latency_keys, &mut low_latency_quants) } else { - &mut default_keys + (&mut default_keys, &mut default_quants) }; + // First-seen (file row) quant order — Python's dict insertion + // order, consumed by `available_quants`. + if !target_quants.iter().any(|q| q == &key.quant) { + target_quants.push(key.quant.clone()); + } // Python's `load_moe_data` wraps the leaf insert in a try/except KeyError // and skips on conflict, i.e. it keeps the FIRST occurrence of each // (shape, num_tokens) tuple. Some perf files contain duplicate rows @@ -504,8 +509,14 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result ))); } Ok(LoadedMoeGrids { - default: MoeGrids { by_keys: default_keys }, - low_latency: MoeGrids { by_keys: low_latency_keys }, + default: MoeGrids { + by_keys: default_keys, + quants_in_load_order: default_quants, + }, + low_latency: MoeGrids { + by_keys: low_latency_keys, + quants_in_load_order: low_latency_quants, + }, }) } diff --git a/rust/aiconfigurator-core/src/perf_database/parquet_loader.rs b/rust/aiconfigurator-core/src/perf_database/parquet_loader.rs index 7f9ee298b..00d97717d 100644 --- a/rust/aiconfigurator-core/src/perf_database/parquet_loader.rs +++ b/rust/aiconfigurator-core/src/perf_database/parquet_loader.rs @@ -211,6 +211,20 @@ impl PerfRow { }) } + /// BOOLEAN column, with fallbacks mirroring Python's `_to_bool` string + /// coercion (`str(value).strip().lower() in {"1","true","yes","y"}`) for + /// files that store flags as INT64 or strings. + pub fn bool(&self, col: usize) -> Result { + if let Ok(v) = self.row.get_bool(col) { + return Ok(v); + } + if let Ok(v) = self.row.get_long(col) { + return Ok(v == 1); + } + let s = self.str(col)?; + Ok(matches!(s.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "y")) + } + /// Optional double. Returns None when the column lookup is None OR the /// cell is null. Used by dispatch tables whose split-latency columns /// can legitimately be absent for one mode of the schema (DeepEP normal diff --git a/rust/aiconfigurator-core/src/perf_database/wideep.rs b/rust/aiconfigurator-core/src/perf_database/wideep.rs index e86a06093..2e7617976 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep.rs @@ -54,7 +54,7 @@ use crate::common::system_spec::SystemSpec; use crate::config::{PerfDbSources, PerfSource}; use super::perf_interp::{self, Node, OpInterpConfig}; use super::{kernel_source_ok, resolve_op_sources}; -use super::moe::{query_token_curve, singleton_underflow}; +use super::moe::{query_token_curve, singleton_underflow, MoeSiblingSlice}; use crate::perf_database::parquet_loader::PerfReader; pub struct WideEpTable { @@ -78,6 +78,10 @@ pub struct WideEpTable { struct MoeGrids { by_keys: BTreeMap>, + /// Distinct quant names in first-seen (file row) order — Python's dict + /// insertion order, consumed by the operator-layer transfer ladder + /// (same contract as `moe.rs::MoeTable::available_quants`). + quants_in_load_order: Vec, } /// TRT-LLM alltoall grids. Keying mirrors Python `load_trtllm_alltoall_data` @@ -263,6 +267,97 @@ impl WideEpTable { ) } + /// Own-slice `num_tokens -> latency_ms` curve of the context/generation + /// WideEP MoE table, after the per-quant `"uniform"` distribution + /// fallback. Typed miss when the slice is absent or empty — the + /// EMPIRICAL calibration input of the SGLang deepep branch of + /// `MoE._query_moe_table` (`_moe_table` routes the util grid to these + /// tables, `operations/moe.py:364-397`). + #[allow(clippy::too_many_arguments)] + pub fn moe_slice_points( + &self, + is_context: bool, + quant_name: &str, + workload_distribution: &str, + topk: u32, + num_experts: u32, + hidden_size: u32, + inter_size: u32, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { + let grids = self.load_ctx_or_gen_moe(is_context)?; + let key = MoeKey { + quant: quant_name.to_string(), + distribution: resolve_moe_distribution(grids, quant_name, workload_distribution), + topk, + num_experts, + hidden_size, + inter_size, + moe_tp_size, + moe_ep_size, + }; + let by_tokens = grids.by_keys.get(&key).filter(|curve| !curve.is_empty()).ok_or_else(|| { + AicError::PerfDatabase(format!( + "WideEP MoE data missing for {key:?} at {}", + self.data_root.display() + )) + })?; + Ok(by_tokens.iter().map(|(&t, &lat)| (t, lat)).collect()) + } + + /// All collected sibling slices of the context/generation WideEP MoE + /// table for `(quant, distribution-after-uniform-fallback, moe_tp, + /// moe_ep)`; empty curves skipped, an empty result is data (not an + /// error). Same contract (incl. the tie-order NOTE) as + /// `moe.rs::MoeTable::sibling_slices` — the transfer ladder walks THIS + /// table when Python's `_moe_table` routes an SGLang deepep query here. + pub fn moe_sibling_slices( + &self, + is_context: bool, + quant_name: &str, + workload_distribution: &str, + moe_tp_size: u32, + moe_ep_size: u32, + ) -> Result, AicError> { + let grids = self.load_ctx_or_gen_moe(is_context)?; + let dist = resolve_moe_distribution(grids, quant_name, workload_distribution); + let mut slices = Vec::new(); + for (key, curve) in &grids.by_keys { + if key.quant != quant_name + || key.distribution != dist + || key.moe_tp_size != moe_tp_size + || key.moe_ep_size != moe_ep_size + || curve.is_empty() + { + continue; + } + slices.push(MoeSiblingSlice { + topk: key.topk, + num_experts: key.num_experts, + hidden_size: key.hidden_size, + inter_size: key.inter_size, + points: curve.iter().map(|(&t, &lat)| (t, lat)).collect(), + }); + } + Ok(slices) + } + + /// Distinct quant names of the context/generation WideEP MoE table, in + /// first-seen (file row) order (Python dict insertion order — same + /// contract as `moe.rs::MoeTable::available_quants`). + pub fn moe_available_quants(&self, is_context: bool) -> Result, AicError> { + Ok(self.load_ctx_or_gen_moe(is_context)?.quants_in_load_order.clone()) + } + + fn load_ctx_or_gen_moe(&self, is_context: bool) -> Result<&MoeGrids, AicError> { + if is_context { + self.load_context_moe() + } else { + self.load_generation_moe() + } + } + #[allow(clippy::too_many_arguments)] pub fn query_trtllm_wideep_moe( &self, @@ -527,6 +622,24 @@ impl WideEpTable { } } +/// Mirrors Python's `wl = workload if workload in moe_data[quant] else +/// "uniform"` on a wideep MoE grid. +fn resolve_moe_distribution( + grids: &MoeGrids, + quant_name: &str, + workload_distribution: &str, +) -> String { + let requested_exists = grids + .by_keys + .keys() + .any(|k| k.quant == quant_name && k.distribution == workload_distribution); + if requested_exists { + workload_distribution.to_string() + } else { + "uniform".to_string() + } +} + #[allow(clippy::too_many_arguments)] fn query_moe( grids: &MoeGrids, @@ -544,15 +657,7 @@ fn query_moe( sol: &dyn Fn(f64) -> f64, ) -> Result { let quant_name = quant.name(); - let requested_exists = grids - .by_keys - .keys() - .any(|k| k.quant == quant_name && k.distribution == workload_distribution); - let distribution = if requested_exists { - workload_distribution.to_string() - } else { - "uniform".to_string() - }; + let distribution = resolve_moe_distribution(grids, quant_name, workload_distribution); let key = MoeKey { quant: quant_name.to_string(), distribution, @@ -752,6 +857,7 @@ fn dispatch_lookup( /// context/generation/wideep-moe parquets (alltoall has its own loader). fn load_moe_parquet(sources: &[PerfSource]) -> Result { let mut by_keys: BTreeMap> = BTreeMap::new(); + let mut quants_in_load_order: Vec = Vec::new(); let mut any_source = false; for source in sources { let path = source.path(); @@ -789,6 +895,10 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result { moe_tp_size: row.u32_optional(moe_tp_size_col)?.unwrap_or(1), moe_ep_size: row.u32(moe_ep_size_col)?, }; + // First-seen (file row) quant order (Python dict insertion order). + if !quants_in_load_order.iter().any(|q| q == &key.quant) { + quants_in_load_order.push(key.quant.clone()); + } // Last-wins parity with Python `load_wideep_*_moe_data` // (direct-assign, no skip-on-conflict guard). by_keys @@ -804,7 +914,7 @@ fn load_moe_parquet(sources: &[PerfSource]) -> Result { sources.first().map(|s| s.path().display().to_string()).unwrap_or_default() ))); } - Ok(MoeGrids { by_keys }) + Ok(MoeGrids { by_keys, quants_in_load_order }) } /// Auto-select the TRT-LLM All2All kernel. Verbatim port of Python diff --git a/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs b/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs index 54a30c857..46312e183 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep_mla.rs @@ -273,6 +273,19 @@ impl WideEpMlaTable { non_empty_points(node, "WideEP generation MLA", &self.data_root) } + /// Probe the context table load. Typed missing-data error when the + /// perf file is absent — Python `get_silicon`'s `raise_if_not_loaded()` + /// step, which PRECEDES the attn-backend whitelist (`mla.py:1449-1461`). + pub fn ensure_context_loaded(&self) -> Result<(), AicError> { + self.load_context().map(|_| ()) + } + + /// Generation-table counterpart of [`Self::ensure_context_loaded`] + /// (`mla.py:1188-1192`). + pub fn ensure_generation_loaded(&self) -> Result<(), AicError> { + self.load_generation().map(|_| ()) + } + fn load_context(&self) -> Result<&WideEpContextMlaGrids, AicError> { let cell = self .context diff --git a/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs b/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs index a18608c80..4991d277f 100644 --- a/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs +++ b/rust/aiconfigurator-core/src/perf_database/wideep_moe.rs @@ -50,6 +50,12 @@ pub struct WideEpMoeTable { /// num_slots, moe_tp, moe_ep)` -> `num_tokens -> latency`. pub struct WideEpMoeGrids { pub by_keys: BTreeMap>, + /// First-seen (file row order) distribution per `(kernel_source, quant)`. + /// Python's fallback takes `list(quant_data.keys())[0]` — dict INSERTION + /// order, i.e. the distribution of the first row loaded for that + /// `(kernel, quant)` — which differs from sorted order on real shards + /// (e.g. gb200/h100 wideep files start with `power_law_1.01_eplb`). + pub first_distribution: BTreeMap<(String, String), String>, } #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] @@ -200,11 +206,11 @@ impl WideEpMoeTable { /// `require_data_slice(data, kernel)` -> `require_data_slice(kd, quant)` /// -> distribution fallback (`workload if present else first available /// under (kernel, quant)`) -> exact remaining coordinate. Each level - /// misses with a typed `AicError::PerfDatabase`. NOTE: the Python - /// fallback takes the FIRST distribution in dict-insertion (file row) - /// order and then requires the full shape under it (a shape present only - /// under a later distribution still misses) — mirrored here with sorted - /// (BTreeMap) order for the distribution list. + /// misses with a typed `AicError::PerfDatabase`. The Python fallback + /// takes the FIRST distribution in dict-insertion (file row) order — + /// served here from the load-time `first_distribution` map — and then + /// requires the full shape under it (a shape present only under a later + /// distribution still misses). #[allow(clippy::too_many_arguments)] fn resolve_slice( &self, @@ -222,7 +228,6 @@ impl WideEpMoeTable { let grids = self.load_compute()?; let mut kernel_seen = false; let mut quant_seen = false; - let mut first_distribution: Option<&str> = None; let mut requested_distribution_seen = false; for key in grids.by_keys.keys() { if key.kernel_source != kernel_source { @@ -233,9 +238,6 @@ impl WideEpMoeTable { continue; } quant_seen = true; - if first_distribution.is_none() { - first_distribution = Some(&key.distribution); - } if key.distribution == distribution { requested_distribution_seen = true; break; @@ -257,7 +259,10 @@ impl WideEpMoeTable { let dist = if requested_distribution_seen { distribution } else { - first_distribution.expect("quant_seen implies at least one distribution") + grids + .first_distribution + .get(&(kernel_source.to_string(), quant_name.to_string())) + .expect("quant_seen implies a recorded first distribution") }; let key = WideEpMoeKey { kernel_source: kernel_source.to_string(), @@ -298,6 +303,7 @@ impl WideEpMoeTable { /// yields rows. fn load_compute_parquet(sources: &[PerfSource]) -> Result { let mut by_keys: BTreeMap> = BTreeMap::new(); + let mut first_distribution: BTreeMap<(String, String), String> = BTreeMap::new(); let mut any_source = false; for source in sources { let path = source.path(); @@ -342,6 +348,11 @@ fn load_compute_parquet(sources: &[PerfSource]) -> Result Result AicError { diff --git a/rust/aiconfigurator-core/src/py.rs b/rust/aiconfigurator-core/src/py.rs index 2e7cce586..6912f67d2 100644 --- a/rust/aiconfigurator-core/src/py.rs +++ b/rust/aiconfigurator-core/src/py.rs @@ -30,6 +30,8 @@ use std::sync::Arc; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::sync::GILOnceCell; +use pyo3::types::PyType; use crate::common::error::AicError; use crate::engine::runtime::{ @@ -44,10 +46,70 @@ fn _build_smoke() -> u32 { ENGINE_CONFIG_SCHEMA_VERSION } +/// Cached handles to the canonical SDK exception classes +/// (`aiconfigurator.sdk.errors`). Filled lazily on first use so importing the +/// extension never imports the sdk (the sdk imports aiconfigurator_core — an +/// eager import here would be a cycle), and left empty in pure-Rust contexts +/// where the sdk is not installed (fallback to `PyValueError`). +static PERF_DATA_NOT_AVAILABLE_ERROR: GILOnceCell> = GILOnceCell::new(); +static EMPIRICAL_NOT_IMPLEMENTED_ERROR: GILOnceCell> = GILOnceCell::new(); + +/// Resolve (and memoize) one sdk error class; `None` when the sdk is not +/// importable, so the caller falls back to `PyValueError`. +fn sdk_error_type( + py: Python<'_>, + cell: &'static GILOnceCell>, + name: &str, +) -> Option> { + cell.get_or_try_init(py, || -> PyResult> { + Ok(py + .import("aiconfigurator.sdk.errors")? + .getattr(name)? + .downcast_into::()? + .unbind()) + }) + .ok() + .map(|ty| ty.clone_ref(py)) +} + /// Convert a crate error into a Python exception at the `#[pymethods]` boundary. /// Inline (not a `From` impl in `error.rs`) so the error module stays pyo3-free. +/// +/// Typed mapping so Python-side classifiers keep working across the FFI: +/// * missing-perf-data errors (`AicError::PerfDatabase` / `Io` — the +/// `is_missing_perf_data` set) raise the canonical +/// `aiconfigurator.sdk.errors.PerfDataNotAvailableError`, so +/// `perf_database.has_perf_data_not_available_cause` recognizes rust-path +/// data misses; +/// * `AicError::EmpiricalNotImplemented` raises +/// `aiconfigurator.sdk.errors.EmpiricalNotImplementedError` (the typed +/// HYBRID/EMPIRICAL coverage miss); +/// * everything else stays `PyValueError`. +/// +/// The sdk import is lazy and failure-tolerant: in pure-Rust test contexts +/// (cargo test without the sdk on `sys.path`) the conversion degrades to +/// `PyValueError` with the same message. fn aic_to_py(e: AicError) -> PyErr { - PyValueError::new_err(e.to_string()) + let sdk_class: Option<(&'static GILOnceCell>, &str)> = if e.is_missing_perf_data() { + Some((&PERF_DATA_NOT_AVAILABLE_ERROR, "PerfDataNotAvailableError")) + } else if matches!(e, AicError::EmpiricalNotImplemented(_)) { + Some((&EMPIRICAL_NOT_IMPLEMENTED_ERROR, "EmpiricalNotImplementedError")) + } else { + None + }; + let message = e.to_string(); + if let Some((cell, name)) = sdk_class { + // `with_gil` is re-entrant, so this is safe whether the caller holds + // the GIL (from_spec) or just re-acquired it (post-allow_threads). + let typed = Python::with_gil(|py| { + sdk_error_type(py, cell, name) + .map(|ty| PyErr::from_type(ty.into_bound(py), message.clone())) + }); + if let Some(err) = typed { + return err; + } + } + PyValueError::new_err(message) } /// Map the `mode` string (Python's `_run_static_breakdown` convention) to the @@ -192,6 +254,8 @@ impl AicEngine { gen_seq_imbalance_correction_scale, }; let mode = parse_mode(mode)?; + // Per-call provenance scope (see `last_provenance`). + self.inner.reset_provenance(); // Rust compute runs with the GIL released. let result: StaticResult = py .allow_threads(|| self.inner.run_static(&rt, mode, stride)) @@ -199,6 +263,21 @@ impl AicEngine { Ok((result.context_ms, result.generation_ms, result.total_ms)) } + /// Empirical provenance tier fired during the most recent compute call on + /// this handle (max-rank across ops, Python `worst_provenance` semantics), + /// or `None` when the call was answered purely from silicon tables. + /// + /// Every compute method (`run_static`, `predict_*_latency`, + /// `mixed_step_latency`, `decode_step_latency`) resets the accumulator on + /// entry, so this is per-call state — the Python bridge reads it right + /// after each call and forwards non-silicon tiers into + /// `util_empirical.note_provenance` (support-matrix HYBRID labelling). + /// Not meaningful under concurrent calls on one handle (the sweep drives + /// each handle sequentially). + fn last_provenance(&self) -> Option<&'static str> { + self.inner.last_provenance() + } + /// Mocker H1: prefill-step latency in ms. Thin shim over `run_static` with /// `mode=Context` (osl is irrelevant for the context phase, so it is fixed /// at 1). Returns the total ms (== context_ms in this mode). @@ -210,6 +289,7 @@ impl AicEngine { isl: u32, prefix: u32, ) -> PyResult { + self.inner.reset_provenance(); py.allow_threads(|| self.inner.predict_prefill_latency(bs, isl, prefix)) .map_err(aic_to_py) } @@ -219,6 +299,7 @@ impl AicEngine { /// `s = isl + 1`). Returns the total ms (== generation_ms in this mode). #[pyo3(signature = (bs, isl, osl=2))] fn predict_decode_latency(&self, py: Python<'_>, bs: u32, isl: u32, osl: u32) -> PyResult { + self.inner.reset_provenance(); py.allow_threads(|| self.inner.predict_decode_latency(bs, isl, osl)) .map_err(aic_to_py) } @@ -243,6 +324,7 @@ impl AicEngine { seq_imbalance_correction_scale: f64, gen_seq_imbalance_correction_scale: f64, ) -> PyResult { + self.inner.reset_provenance(); py.allow_threads(|| { self.inner.mixed_step_latency( ctx_tokens, @@ -269,6 +351,7 @@ impl AicEngine { osl: u32, gen_seq_imbalance_correction_scale: f64, ) -> PyResult { + self.inner.reset_provenance(); py.allow_threads(|| { self.inner .decode_step_latency(gen_tokens, isl, osl, gen_seq_imbalance_correction_scale) @@ -965,4 +1048,86 @@ mod tests { assert!((aic_decode - raw_decode).abs() < 1e-12); assert!(aic_decode > 0.0 && aic_decode.is_finite()); } + + /// `aic_to_py` must raise the canonical sdk exception classes for the + /// typed variants when the sdk is importable, and degrade to `ValueError` + /// when it is not (pure-Rust test contexts). Everything else stays + /// `ValueError` unconditionally. + #[test] + fn aic_to_py_maps_typed_errors_to_sdk_classes() { + py_init(); + Python::with_gil(|py| { + let sdk_available = py.import("aiconfigurator.sdk.errors").is_ok(); + + let check = |err: AicError, sdk_name: &str| { + let pyerr = aic_to_py(err); + let type_name = pyerr.get_type(py).name().unwrap().to_string(); + if sdk_available { + assert_eq!(type_name, sdk_name); + } else { + assert_eq!(type_name, "ValueError"); + } + }; + check( + AicError::PerfDatabase("missing table".to_string()), + "PerfDataNotAvailableError", + ); + check( + AicError::Io { + path: PathBuf::from("/nope"), + source: std::io::Error::new(std::io::ErrorKind::NotFound, "nope"), + }, + "PerfDataNotAvailableError", + ); + check( + AicError::EmpiricalNotImplemented("no basis".to_string()), + "EmpiricalNotImplementedError", + ); + + // Non-typed variants stay ValueError regardless of the sdk. + let other = aic_to_py(AicError::InvalidEngineConfig("bad".to_string())); + assert_eq!(other.get_type(py).name().unwrap().to_string(), "ValueError"); + + // The message survives the mapping. + let msg_err = aic_to_py(AicError::PerfDatabase("missing table xyz".to_string())); + assert!(msg_err.value(py).to_string().contains("missing table xyz")); + }); + } + + /// Compute bindings reset the provenance accumulator on entry, and + /// `last_provenance` reads it back (None == pure silicon). The silicon + /// fixture fires no empirical path, so a pre-seeded tier must be cleared + /// by the next compute call. + #[test] + fn compute_bindings_reset_provenance_per_call() { + use crate::operators::util_empirical::ProvenanceTier; + + py_init(); + let bytes = fixture_spec_bytes(); + let root = systems_root(); + let aic = AicEngine::from_spec(&bytes, root.to_str()).unwrap(); + + assert_eq!(aic.last_provenance(), None); + aic.inner.database().note_provenance(ProvenanceTier::XOp); + assert_eq!(aic.last_provenance(), Some("xop")); + + // Positional order: (bs, beam, isl, osl, prefix, seq_corr, gen_seq_corr, mode, stride). + Python::with_gil(|py| { + aic.run_static(py, 1, 1, 1024, 8, 0, 1.0, 1.0, "static", 32) + .unwrap(); + }); + assert_eq!(aic.last_provenance(), None); + + aic.inner.database().note_provenance(ProvenanceTier::Empirical); + Python::with_gil(|py| { + aic.mixed_step_latency(py, 1024, 2, 1024, 8, 0, 1.0, 1.0).unwrap(); + }); + assert_eq!(aic.last_provenance(), None); + + aic.inner.database().note_provenance(ProvenanceTier::XShape); + Python::with_gil(|py| { + aic.decode_step_latency(py, 4, 1024, 8, 1.0).unwrap(); + }); + assert_eq!(aic.last_provenance(), None); + } } diff --git a/src/aiconfigurator/sdk/engine.py b/src/aiconfigurator/sdk/engine.py index 7c64e117a..12f97bb83 100644 --- a/src/aiconfigurator/sdk/engine.py +++ b/src/aiconfigurator/sdk/engine.py @@ -51,6 +51,7 @@ ContextMLA, ContextMSAModule, CustomAllReduce, + DeepSeekV4MegaMoEModule, DeepSeekV4MHCModule, ElementWise, Embedding, @@ -457,6 +458,32 @@ def _dsv4_module( } +def _dsv4_megamoe(op: DeepSeekV4MegaMoEModule) -> dict: + """SGLang DeepSeek-V4 MegaMoE routed module (Python + ``DeepSeekV4MegaMoEModule``). One class serves both phases via + ``is_context``, so a single Rust variant carries the flag. Field names + match the Rust ``Dsv4MegaMoeOp``; ``workload_distribution`` is already + normalized by the ctor (``uniform`` -> ``balanced``).""" + return { + "name": op._name, + "scale_factor": op._scale_factor, + "hidden_size": op._hidden_size, + "inter_size": op._inter_size, + "topk": op._topk, + "num_experts": op._num_experts, + "moe_tp_size": op._moe_tp_size, + "moe_ep_size": op._moe_ep_size, + "quant_mode": _quant_name(op._quant_mode), + "workload_distribution": op._workload_distribution, + "is_context": op._is_context, + "source_policy": op._source_policy, + "pre_dispatch": op._pre_dispatch, + "num_fused_shared_experts": op._num_fused_shared_experts, + "kernel_source": op._kernel_source, + "kernel_dtype": op._kernel_dtype, + } + + def _mhc_module(op: DeepSeekV4MHCModule, *, architecture: str) -> dict: return { "name": op._name, @@ -626,6 +653,11 @@ def recurse(child: Any) -> dict: return {"Dsv4Context": _dsv4_module(op, architecture=architecture)} if isinstance(op, GenerationDeepSeekV4AttentionModule): return {"Dsv4Generation": _dsv4_module(op, architecture=architecture)} + # Rust `Op::Dsv4MegaMoe` is APPENDED after `Fallback` (bincode enum + # indices are positional; appending shifts nothing), so no + # ENGINE_SPEC_SCHEMA_VERSION bump. + if isinstance(op, DeepSeekV4MegaMoEModule): + return {"Dsv4MegaMoe": _dsv4_megamoe(op)} # WideEP variants (must precede their non-WideEP base classes if any). if isinstance(op, WideEPContextMLA): @@ -1059,3 +1091,13 @@ def decode_step_latency( return self._engine.decode_step_latency( int(gen_tokens), int(isl), int(osl), float(gen_seq_imbalance_correction_scale) ) + + def last_provenance(self) -> str | None: + """Empirical provenance tier fired during the most recent engine call + on this handle (worst tier across ops, per Python's + ``util_empirical.PROVENANCE_ORDER``), or ``None`` for a pure-silicon + answer. Per-call state: every compute method resets the accumulator on + entry. The rust engine-step bridge forwards non-silicon tiers into + ``util_empirical.note_provenance`` so ``capture_provenance()`` / + support-matrix HYBRID labelling behave identically on both engines.""" + return self._engine.last_provenance() diff --git a/src/aiconfigurator/sdk/operations/moe.py b/src/aiconfigurator/sdk/operations/moe.py index 9e849e9a7..afe8f0293 100644 --- a/src/aiconfigurator/sdk/operations/moe.py +++ b/src/aiconfigurator/sdk/operations/moe.py @@ -2272,6 +2272,7 @@ def get_empirical() -> float: moe_ep_size, quant_mode, node_num, + kernel_source, ) return database._query_silicon_or_hybrid( diff --git a/src/aiconfigurator/sdk/rust_engine_step.py b/src/aiconfigurator/sdk/rust_engine_step.py index e699aa864..981ab0834 100644 --- a/src/aiconfigurator/sdk/rust_engine_step.py +++ b/src/aiconfigurator/sdk/rust_engine_step.py @@ -259,6 +259,28 @@ def should_use_rust_engine_step(runtime_config: RuntimeConfig, database: Any = N return True +def _note_rust_provenance(handle: Any) -> None: + """Forward the compiled engine's per-call empirical provenance tier into + Python's capture (``util_empirical.note_provenance``). + + ``EngineHandle.last_provenance()`` returns the worst tier fired during the + engine call just made (``None`` for a pure-silicon answer). Forwarding it + keeps ``capture_provenance()`` consumers — the support matrix's + HYBRID_PASS tier labelling — working unchanged when the engine step is + rust-routed. ``note_provenance`` is a no-op outside an active capture, so + this costs one getattr-free call per step. Only the worst tier crosses the + FFI (not the full tag set); ``worst_provenance`` over the captured tags is + unaffected because max(worst) == worst(all). + """ + tier = handle.last_provenance() + if tier is not None and tier != "silicon": + # Deferred import: keep module import light and cycle-free + # (sdk.engine imports this module at top level). + from aiconfigurator.sdk.operations import util_empirical + + util_empirical.note_provenance(tier) + + def estimate_static_latency_breakdown_with_rust( model: Any, database: Any, @@ -290,6 +312,7 @@ def estimate_static_latency_breakdown_with_rust( mode=engine_mode, stride=int(stride), ) + _note_rust_provenance(handle) if latency_correction_scale != 1.0: context_latency_ms *= latency_correction_scale @@ -324,7 +347,7 @@ def estimate_mixed_step_latency_with_rust( straight through with no Python-side pre-math. """ handle = _cached_engine_handle(model, database) - return handle.mixed_step_latency( + latency_ms = handle.mixed_step_latency( int(ctx_tokens), int(gen_tokens), int(isl), @@ -333,6 +356,8 @@ def estimate_mixed_step_latency_with_rust( seq_imbalance_correction_scale=float(seq_imbalance_correction_scale or 1.0), gen_seq_imbalance_correction_scale=float(gen_seq_imbalance_correction_scale or 1.0), ) + _note_rust_provenance(handle) + return latency_ms def estimate_decode_step_latency_with_rust( @@ -353,12 +378,14 @@ def estimate_decode_step_latency_with_rust( applied internally, so the raw args pass straight through. """ handle = _cached_engine_handle(model, database) - return handle.decode_step_latency( + latency_ms = handle.decode_step_latency( int(gen_tokens), int(isl), int(osl), gen_seq_imbalance_correction_scale=float(gen_seq_imbalance_correction_scale or 1.0), ) + _note_rust_provenance(handle) + return latency_ms # Memo of compiled ``EngineHandle`` objects, keyed by the engine identity diff --git a/tests/unit/sdk/database/test_moe_mla.py b/tests/unit/sdk/database/test_moe_mla.py index 850c32e1c..b2cc1f10e 100644 --- a/tests/unit/sdk/database/test_moe_mla.py +++ b/tests/unit/sdk/database/test_moe_mla.py @@ -912,3 +912,43 @@ def test_transfer_policy_gates_and_tags_cross_profile(self, comprehensive_perf_d assert util_empirical.worst_provenance(tags) == "xprofile" finally: comprehensive_perf_db.set_transfer_policy(None) + + +class TestAlltoallHybridFallbackClosure: + """Regression: the HYBRID fallback closure of + `TrtLLMWideEPMoEDispatch._query_alltoall_table` omitted the mandatory + `kernel_source` argument of `get_empirical_from_sol`, so a silicon miss + under HYBRID raised `TypeError` instead of running the empirical + estimate. The fallback must execute and surface the TYPED empirical + outcome (a value, or `EmpiricalNotImplementedError` when the slice has no + calibration data) — never a `TypeError`.""" + + def test_hybrid_silicon_miss_runs_empirical_closure(self): + from aiconfigurator.sdk import perf_database + from aiconfigurator.sdk.operations.moe import TrtLLMWideEPMoEDispatch + + db = perf_database.get_database_view( + "gb200", + "trtllm", + "1.3.0rc10", + allow_missing_data=True, + database_mode="HYBRID", + shared_layer=False, + ) + if db is None: + pytest.skip("gb200/trtllm/1.3.0rc10 data unavailable") + + # Off-shape hidden_size forces a silicon miss; the HYBRID fallback + # closure must run (typed empirical miss here — the slice has no + # own-shape calibration data), not crash with TypeError. + with pytest.raises(EmpiricalNotImplementedError): + TrtLLMWideEPMoEDispatch._query_alltoall_table( + db, + op_name="alltoall_dispatch", + num_tokens=64, + hidden_size=7000, + topk=8, + num_experts=256, + moe_ep_size=8, + quant_mode=common.MoEQuantMode.nvfp4, + ) diff --git a/tests/unit/sdk/test_rust_engine_step.py b/tests/unit/sdk/test_rust_engine_step.py index 217e5a981..c027bbb11 100644 --- a/tests/unit/sdk/test_rust_engine_step.py +++ b/tests/unit/sdk/test_rust_engine_step.py @@ -56,6 +56,9 @@ def run_static(self, **kwargs): # (context_ms, generation_ms, total_ms) return (10.0, 6.0, 16.0) + def last_provenance(self): + return None # pure-silicon answer + monkeypatch.setattr(rust_engine_step, "_cached_engine_handle", lambda model, database: _FakeHandle()) model = _dense_model() @@ -102,6 +105,9 @@ def decode_step_latency(self, *args, **kwargs): decode_calls.append((args, kwargs)) return 9.5 + def last_provenance(self): + return None # pure-silicon answer + monkeypatch.setattr(rust_engine_step, "_cached_engine_handle", lambda model, database: _FakeHandle()) model = _dense_model() @@ -137,6 +143,60 @@ def decode_step_latency(self, *args, **kwargs): assert decode_calls == [((7, 256, 256), {"gen_seq_imbalance_correction_scale": 1.0})] +def test_rust_provenance_tier_forwarded_into_python_capture(monkeypatch) -> None: + """The engine-step helpers forward the compiled engine's per-call + empirical provenance tier into Python's ``capture_provenance`` (used by + the support matrix to label HYBRID_PASS rows). Silicon answers + (``last_provenance() is None`` or ``"silicon"``) record nothing.""" + from aiconfigurator.sdk.operations import util_empirical + + class _FakeHandle: + def __init__(self, tier): + self._tier = tier + + def run_static(self, **kwargs): + return (10.0, 6.0, 16.0) + + def mixed_step_latency(self, *args, **kwargs): + return 8.5 + + def decode_step_latency(self, *args, **kwargs): + return 9.5 + + def last_provenance(self): + return self._tier + + model = _dense_model() + database = SimpleNamespace(system="test_sxm", backend="vllm", version="1.0.0") + runtime_config = RuntimeConfig(batch_size=1, beam_width=1, isl=8, osl=4) + + def run_all_helpers(tier): + monkeypatch.setattr(rust_engine_step, "_cached_engine_handle", lambda model, database: _FakeHandle(tier)) + rust_engine_step.estimate_static_latency_breakdown_with_rust( + model, database, runtime_config, mode="static", stride=2, latency_correction_scale=1.0 + ) + rust_engine_step.estimate_mixed_step_latency_with_rust( + model, database, ctx_tokens=8, gen_tokens=1, isl=8, osl=4, prefix=0 + ) + rust_engine_step.estimate_decode_step_latency_with_rust(model, database, gen_tokens=1, isl=8, osl=4) + + # A rust-routed HYBRID rescue records its tier (all three step surfaces). + with util_empirical.capture_provenance() as tags: + run_all_helpers("xop") + assert tags == {"xop"} + assert util_empirical.worst_provenance(tags) == "xop" + + # Pure-silicon runs record nothing -> worst_provenance stays "silicon". + for silicon_tier in (None, "silicon"): + with util_empirical.capture_provenance() as tags: + run_all_helpers(silicon_tier) + assert tags == set() + assert util_empirical.worst_provenance(tags) == "silicon" + + # Outside a capture, forwarding is a no-op (note_provenance no-ops). + run_all_helpers("xshape") + + def test_engine_config_json_preserves_moe_specific_quant_mode() -> None: model = SimpleNamespace( model_path="Test/Moe", From 6a9e260398119943543109d0b1c80331be836985 Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Tue, 14 Jul 2026 22:58:08 +0800 Subject: [PATCH 09/10] =?UTF-8?q?chore:=20parity-freeze=20defenses=20?= =?UTF-8?q?=E2=80=94=20opspec=20coverage=20tripwire,=20rules=20doc,=20rust?= =?UTF-8?q?=20CODEOWNERS?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-freeze drift defenses for community PRs: - tests/unit/sdk/test_opspec_coverage.py: every Operation subclass must have a _to_opspec branch or a justified EXEMPT entry (current exemptions: the four AFD session-level ops pending the op-list FFI, and the dead Mamba2 class). Complements the Rust-side spec.rs exhaustiveness guard. - .claude/rules/rust-core/parity.md: the freeze discipline (mirror + oracle + parity case for any op/query-math change; tail-append rule for Op variants; selection rules incl. dict-order fallbacks are parity surface; the intentional splits that must not be "fixed" ad hoc). - CODEOWNERS: /rust/ now owned by the same teams as /src/ so parity-coupled changes route to the same reviewers (was falling through to the catch-all). Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- .claude/rules/rust-core/parity.md | 59 ++++++++++++++++++++ .github/CODEOWNERS | 5 ++ tests/unit/sdk/test_opspec_coverage.py | 77 ++++++++++++++++++++++++++ 3 files changed, 141 insertions(+) create mode 100644 .claude/rules/rust-core/parity.md create mode 100644 tests/unit/sdk/test_opspec_coverage.py diff --git a/.claude/rules/rust-core/parity.md b/.claude/rules/rust-core/parity.md new file mode 100644 index 000000000..b697a8abf --- /dev/null +++ b/.claude/rules/rust-core/parity.md @@ -0,0 +1,59 @@ +# Rust-Core Parity Discipline (Python-path freeze) + +The compiled engine (`rust/aiconfigurator-core`) is at full numeric parity +with the Python engine step for the SILICON, HYBRID, and EMPIRICAL database +modes, guarded by `rust/aiconfigurator-core/parity_tests/` (engine-step, +compile-engine, and perf gates). The Python op/query math is **frozen**: it +is the reference the parity suite compares against, and it is scheduled for +retirement (see the Python-path freeze tracking issue). + +## The rule + +Any PR that changes latency-affecting behavior in: + +- `src/aiconfigurator/sdk/operations/**` (op query math, SOL formulas, + slice/kernel selection, transfer ladder), +- the query/loader layer of `src/aiconfigurator/sdk/perf_database.py`, +- `src/aiconfigurator/sdk/perf_interp/**`, + +MUST in the same PR: + +1. **Mirror the change** in the corresponding `rust/aiconfigurator-core/src/` + operator/table (the layering matches: dispatch + estimators in + `operators/`, algorithm-free loaders/accessors in `perf_database/`). +2. **Anchor it**: a Python-generated oracle in the Rust `#[cfg(test)]` module + (1e-9, `uv run python` against a `shared_layer=False` view — copy the + existing oracle-test pattern in `operators/gemm.rs`), and/or a case in + `parity_tests/test_engine_step_parity.py` when a new config class becomes + reachable. +3. Keep the `rust-engine-step-parity` CI job green — it is the enforcement + mechanism, not this document. + +## Adding a new Operation + +A new `Operation` subclass must get a `_to_opspec` branch in +`sdk/engine.py`, an `Op` variant in `operators/op.rs` (**append at the tail** +— bincode variant indices are positional; mid-enum insertion requires an +`ENGINE_SPEC_SCHEMA_VERSION` bump on BOTH sides), the `engine/spec.rs` +round-trip fixture, and a parity case. `tests/unit/sdk/test_opspec_coverage.py` +fails until the op converts or carries a justified `EXEMPT` entry. + +## Selection rules are parity surface too + +Table/slice/kernel selection must match rule-for-rule, including fallback +order and tie-breaks. Python dicts iterate in file/insertion order; Rust +`BTreeMap` iterates sorted — any "first available" fallback needs a +load-order record on the Rust side (see `quants_in_load_order` / +`first_distribution` in `perf_database/{moe,wideep_moe}.rs` for the pattern). + +## Known intentional splits (do not "fix" without the tracking issue) + +- SOL / SOL_FULL modes delegate to the Python step (routing gate in + `sdk/rust_engine_step.py`). +- AFD and the VL encoder phase are Python-side orchestration; their per-op + values move to Rust via the planned op-list evaluation FFI, not by porting + the orchestration. +- Rust reads parquet only (no `.txt` legacy loading) — new data drops must + ship parquet. +- Energy/power does not yet cross the FFI (rust-routed reports show 0.0W) + until the energy follow-up PR lands. diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index df2c28671..a3507397b 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -16,6 +16,11 @@ # Runtime, SDK, backends, systems, model configs, and UI /src/ @ai-dynamo/aiconfigurator-runtime @ai-dynamo/maintainers-aiconfigurator +# The compiled engine mirrors the SDK op/query math (parity-coupled: changes +# to sdk/operations or the perf-DB query layer require a matching change +# here — see .claude/rules/rust-core/parity.md). Same owners as /src/ so the +# same reviewers see both sides. +/rust/ @ai-dynamo/aiconfigurator-runtime @ai-dynamo/maintainers-aiconfigurator /src/aiconfigurator/sdk/ @ai-dynamo/aiconfigurator-runtime @ai-dynamo/maintainers-aiconfigurator /src/aiconfigurator/sdk/backends/sglang_backend.py @ai-dynamo/aiconfigurator-runtime @ai-dynamo/maintainers-aiconfigurator /src/aiconfigurator/sdk/backends/trtllm_backend.py @ai-dynamo/aiconfigurator-runtime @ai-dynamo/maintainers-aiconfigurator diff --git a/tests/unit/sdk/test_opspec_coverage.py b/tests/unit/sdk/test_opspec_coverage.py new file mode 100644 index 000000000..2ce05e001 --- /dev/null +++ b/tests/unit/sdk/test_opspec_coverage.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Structural tripwire: every ``Operation`` must reach the compiled engine. + +The Rust engine executes whatever ``engine.py::_to_opspec`` emits. An op class +added on the Python side WITHOUT a conversion branch silently forces the whole +model back onto the (frozen) Python step — exactly the drift the Python-path +freeze forbids. This test makes that state un-mergeable: a new ``Operation`` +subclass must either get a ``_to_opspec`` branch (plus its Rust mirror and a +parity case — see ``.claude/rules/rust-core/parity.md``) or an explicit, +justified entry in ``EXEMPT`` below. + +The Rust side has the symmetric guard: ``engine/spec.rs::all_op_variants`` +fails to compile when an ``Op`` variant is added without wiring. +""" + +from __future__ import annotations + +import inspect +import re + +import pytest + +pytestmark = pytest.mark.unit + +# Op classes deliberately NOT convertible to the compiled engine. Every entry +# needs a reason; removing the op or porting it must also remove the entry +# (the staleness assertion below enforces that). +EXEMPT: dict[str, str] = { + # AFD (attention-FFN disagg) is session-level Python orchestration by + # design (`inference_session.py` builds and sums these directly; the + # engine-step path is never involved). Retirement prerequisite: the thin + # op-list evaluation FFI (see the Python-path freeze tracking issue). + "AFDTransfer": "AFD orchestration is Python-side; op-list FFI planned", + "AFDCombine": "AFD orchestration is Python-side; op-list FFI planned", + "AFDFAllGather": "AFD orchestration is Python-side; op-list FFI planned", + "AFDFReduceScatter": "AFD orchestration is Python-side; op-list FFI planned", + # Dead class: no model instantiates it (Mamba2Kernel is the live op and + # converts). Remove the class or this entry together. + "Mamba2": "dead code — never instantiated; Mamba2Kernel is the live op", +} + + +def _operation_subclasses() -> set[str]: + # Importing the package registers every op subclass. + import aiconfigurator.sdk.operations # noqa: F401 + from aiconfigurator.sdk.operations.base import Operation + + seen: set[type] = set() + stack: list[type] = [Operation] + while stack: + for sub in stack.pop().__subclasses__(): + if sub not in seen: + seen.add(sub) + stack.append(sub) + # Private bases (e.g. _BaseMSAModule) are implementation details; their + # public leaves are what models instantiate. + return {cls.__name__ for cls in seen if not cls.__name__.startswith("_")} + + +def test_every_operation_converts_to_an_opspec_or_is_exempt(): + from aiconfigurator.sdk import engine + + handled = set(re.findall(r"isinstance\(op, (\w+)\)", inspect.getsource(engine._to_opspec))) + all_ops = _operation_subclasses() + + missing = sorted(all_ops - handled - set(EXEMPT)) + assert not missing, ( + f"Operation classes without a _to_opspec branch: {missing}. " + "Port them to the compiled engine (opspec branch + Rust mirror + " + "parity case, see .claude/rules/rust-core/parity.md) or add an " + "explicit EXEMPT entry with a reason." + ) + + stale = sorted(set(EXEMPT) & handled) + assert not stale, f"EXEMPT entries now have conversion branches; remove them: {stale}" From bef449250983f7de1b709c7ee37c2b0f8a1e528e Mon Sep 17 00:00:00 2001 From: Tianhao Xu Date: Wed, 15 Jul 2026 00:31:39 +0800 Subject: [PATCH 10/10] =?UTF-8?q?fix(rust-core):=20address=20PR=20#1355=20?= =?UTF-8?q?review=20=E2=80=94=20mode-aware=20wideep=20dispatch=20+=20guard?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review fixes (coderabbit on PR #1355): - moe_dispatch: TrtllmWideEpMoEDispatchOp::query now routes through the mode-aware query_alltoall_table (Python TrtLLMWideEPMoEDispatch.query passes database_mode=None -> the view's mode): EMPIRICAL estimates SOL/util and HYBRID converts a typed silicon miss into the estimate instead of returning silicon values / hard-erroring. Phase sources combine like Python's PerformanceResult.__add__. Anchored by standalone_wideep_dispatch_is_mode_aware on the existing gb200 Python oracles (SILICON values unchanged). - enums: explicit TransferPolicy Default = ALL (the derived all-false Default silently meant OFF, contradicting the unset-means-ALL wire semantics). - rust_engine_step: forward imbalance-correction scales with a None-only default (_scale_or_one) — `or 1.0` clobbered an explicit 0.0 that the Python path would apply. - parity comparator: both-error pairs must agree on exception kind; a panic paired with a typed miss no longer counts as error symmetry. - wideep_mla: whitelist test pins AicError::InvalidEngineConfig (a plain is_err() would also pass on a lookup miss). - opspec coverage: EXEMPT is a closed set — entries must name live Operation classes and carry non-empty justifications. - parity.md: add description/paths frontmatter per repo-guide rule. - dsa: refresh stale query_context_cp doc (skip_indexer now passes through for the GLM-5.2 CP blend). Validation: cargo test --lib 277 passed; engine-step parity suite 303 passed; touched pytest files green; ruff clean. Deferred to the freeze follow-up (#1357): error-classification fidelity through the clone_err caches / is_missing_perf_data Io narrowing, util-grid cache lock granularity, and the suggested loader/beam-width test coverage additions. Co-Authored-By: Claude Fable 5 Signed-off-by: Tianhao Xu --- .claude/rules/rust-core/parity.md | 15 ++++ .../parity_tests/test_engine_step_parity.py | 19 ++++- rust/aiconfigurator-core/src/common/enums.rs | 10 ++- rust/aiconfigurator-core/src/operators/dsa.rs | 6 +- .../src/operators/moe_dispatch.rs | 82 +++++++++++++++++-- .../src/operators/wideep_mla.rs | 15 +++- src/aiconfigurator/sdk/rust_engine_step.py | 20 +++-- tests/unit/sdk/test_opspec_coverage.py | 8 ++ 8 files changed, 151 insertions(+), 24 deletions(-) diff --git a/.claude/rules/rust-core/parity.md b/.claude/rules/rust-core/parity.md index b697a8abf..4ab9e44bd 100644 --- a/.claude/rules/rust-core/parity.md +++ b/.claude/rules/rust-core/parity.md @@ -1,3 +1,18 @@ +--- +description: > + Parity discipline for the compiled engine: the frozen Python op/query math + is the spec; latency-affecting changes must land on both sides with an + oracle anchor in the same PR. +paths: + - "rust/aiconfigurator-core/**" + - "src/aiconfigurator/sdk/operations/**" + - "src/aiconfigurator/sdk/perf_database.py" + - "src/aiconfigurator/sdk/perf_interp/**" + - "src/aiconfigurator/sdk/engine.py" + - "src/aiconfigurator/sdk/rust_engine_step.py" + - "tests/unit/sdk/test_opspec_coverage.py" +--- + # Rust-Core Parity Discipline (Python-path freeze) The compiled engine (`rust/aiconfigurator-core`) is at full numeric parity diff --git a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py index 493089994..4d3603231 100644 --- a/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py +++ b/rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py @@ -840,9 +840,12 @@ def _parity_mismatch_reason( ) -> str | None: """Compare Python and Rust per-metric outputs with three valid pairings: - - both compute -> numeric tolerance check - - both error -> pass (error-symmetry contract) - - only one errors -> fail with the asymmetric reason + - both compute -> numeric tolerance check + - both error with the SAME kind -> pass (error-symmetry contract; the + typed-FFI mapping raises the canonical sdk classes on the rust side, + so `type(exc).__name__` must agree — a panic/TypeError paired with a + typed miss is a real divergence, not symmetry) + - anything else -> fail with the asymmetric reason Returns ``None`` when every metric in `comparisons` matches under one of those rules; otherwise returns a formatted multi-row diff. @@ -854,7 +857,15 @@ def _parity_mismatch_reason( py_err = isinstance(python_value, _ErrorSentinel) rs_err = isinstance(rust_value, _ErrorSentinel) if py_err and rs_err: - # Both errored — symmetric. Pass. + if python_value.kind != rust_value.kind: + # Both errored, but with different exception classes. + has_mismatch = True + rows.append( + f"{name:<{metric_width}} {python_value!r:>10} {rust_value!r:>10} " + f"{'-':>10} {'-':>10} {'-':>10} kind" + ) + continue + # Both errored with the same kind — symmetric. Pass. rows.append(f"{name:<{metric_width}} {'ERROR':>10} {'ERROR':>10} {'-':>10} {'-':>10} {'-':>10} sym") continue if py_err != rs_err: diff --git a/rust/aiconfigurator-core/src/common/enums.rs b/rust/aiconfigurator-core/src/common/enums.rs index cb249728d..e223fef7a 100644 --- a/rust/aiconfigurator-core/src/common/enums.rs +++ b/rust/aiconfigurator-core/src/common/enums.rs @@ -87,7 +87,7 @@ pub enum TransferKind { /// (`off`/`conservative`/`balanced`/`aggressive`) before serialising, so the /// wire form is always an explicit list of kind tokens; `None` on the wire /// means the default ALL-transfers policy. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct TransferPolicy { pub xshape: bool, pub xquant: bool, @@ -95,6 +95,14 @@ pub struct TransferPolicy { pub xop: bool, } +/// "Unset" means ALL everywhere (`from_wire(None)`, `PerfDatabase::load`); +/// a derived all-`false` Default would silently mean OFF. +impl Default for TransferPolicy { + fn default() -> Self { + Self::ALL + } +} + impl TransferPolicy { pub const ALL: TransferPolicy = TransferPolicy { xshape: true, diff --git a/rust/aiconfigurator-core/src/operators/dsa.rs b/rust/aiconfigurator-core/src/operators/dsa.rs index 77e156cd8..18b18c73d 100644 --- a/rust/aiconfigurator-core/src/operators/dsa.rs +++ b/rust/aiconfigurator-core/src/operators/dsa.rs @@ -136,9 +136,9 @@ impl DsaModuleOp { /// Context-Parallel (CP) prefill — GLM-5/DSA sparse composition. /// /// Wires the real data dependencies and delegates to [`Self::query_cp_with`] - /// (the verbatim mirror of Python `ContextDSAModule._query_cp`, - /// `skip_indexer=False` — the Rust table loads Python's full slice, so - /// this matches the full-layer-only opspec path): + /// (the verbatim mirror of Python `ContextDSAModule._query_cp`); the + /// caller passes `skip_indexer` through, so GLM-5.2's CP + shared-index + /// amortization can query both the full and skip-indexer slices: /// - base = the existing 4-axis engine query at `(b, per_card, prefix)` /// with `dsa_backend="flashmla_kv"`, exactly like Python's `_query_cp` /// base query; diff --git a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs index a205800a6..0715185c3 100644 --- a/rust/aiconfigurator-core/src/operators/moe_dispatch.rs +++ b/rust/aiconfigurator-core/src/operators/moe_dispatch.rs @@ -14,7 +14,8 @@ //! - **SGLang DeepEP**: dispatch + combine latencies come from the //! `wideep_deepep_normal` / `wideep_deepep_ll` tables (see //! `db.wideep.query_deepep_normal/ll`). -//! - **TRT-LLM WideEP**: uses `db.wideep.query_trtllm_alltoall`. +//! - **TRT-LLM WideEP**: uses the mode-aware [`query_alltoall_table`] +//! (silicon arm = `db.wideep.query_trtllm_alltoall`). //! //! All paths route through the corresponding tables; the higher-level //! model is responsible for choosing the dispatch flavor. @@ -114,11 +115,19 @@ pub struct TrtllmWideEpMoEDispatchOp { } impl TrtllmWideEpMoEDispatchOp { + /// Mode-aware, like every phase-op query in Python's + /// `TrtLLMWideEPMoEDispatch.query` (it calls + /// `database.query_trtllm_alltoall` with `database_mode=None`, which + /// resolves to the view's mode inside `_query_alltoall_table`): + /// EMPIRICAL estimates `SOL/util`, HYBRID converts a typed silicon miss + /// into the estimate, SILICON queries the table. `node_num` stays the + /// perf-DB default (Python's model builders always pass `None`). The + /// pre-dispatch sum combines sources exactly like Python's + /// `PerformanceResult.__add__` (same -> same, mismatch -> mixed). pub fn query(&self, db: &PerfDatabase, num_tokens: u32) -> Result { - let spec: &SystemSpec = &db.system_spec; let q = |op_name: &str| { - db.wideep.query_trtllm_alltoall( - spec, + query_alltoall_table( + db, op_name, num_tokens, self.hidden_size, @@ -126,19 +135,23 @@ impl TrtllmWideEpMoEDispatchOp { self.num_experts, self.moe_ep_size, self.quant_mode, + None, Some("wideep"), ) }; - let latency = if self.pre_dispatch { - q("alltoall_prepare")? + q("alltoall_dispatch")? + let result = if self.pre_dispatch { + let prepare = q("alltoall_prepare")?; + let dispatch = q("alltoall_dispatch")?; + PerformanceResult::new( + prepare.latency_ms + dispatch.latency_ms, + prepare.source.combine(dispatch.source), + ) } else if self.use_low_precision_combine { q("alltoall_combine_low_precision")? } else { q("alltoall_combine")? }; - Ok(PerformanceResult::new(latency, Source::Silicon) - .clamp_non_negative() - .scaled(self.scale_factor)) + Ok(result.clamp_non_negative().scaled(self.scale_factor)) } } @@ -1227,6 +1240,57 @@ mod tests { assert_oracle(&combine_lp, 0.069468751270324, Source::Silicon, "hyb_wideep_combine_lp"); } + /// The standalone WideEP dispatch op must route through the mode-aware + /// [`query_alltoall_table`] (Python `TrtLLMWideEPMoEDispatch.query` + /// passes `database_mode=None`, i.e. the view's mode) — a direct + /// silicon-table call would return silicon values under EMPIRICAL and + /// hard-error instead of falling back under HYBRID. Expected values are + /// the per-phase-op Python oracles already pinned in + /// `alltoall_empirical_matches_python_oracles` / + /// `alltoall_hybrid_prefers_silicon_when_covered`, composed per phase + /// (pre-dispatch = prepare + dispatch). + #[test] + fn standalone_wideep_dispatch_is_mode_aware() { + let wideep_op = |pre_dispatch: bool, quant: MoeQuantMode, low_precision: bool| { + TrtllmWideEpMoEDispatchOp { + name: "trtllm_wideep_dispatch".to_string(), + scale_factor: 1.0, + hidden_size: 7168, + topk: 8, + num_experts: 256, + moe_tp_size: 1, + moe_ep_size: 8, + attention_dp_size: 8, + pre_dispatch, + quant_mode: quant, + use_low_precision_combine: low_precision, + } + }; + + let emp = gb200_trtllm_db(DatabaseMode::Empirical); + let pre = wideep_op(true, MoeQuantMode::Fp8, false).query(&emp, 333).expect("emp pre"); + assert_oracle( + &pre, + 0.015176159059580488 + 0.04266341407888801, // prepare + dispatch + Source::Empirical, + "standalone_emp_pre", + ); + let combine_lp = + wideep_op(false, MoeQuantMode::Nvfp4, true).query(&emp, 333).expect("emp combine_lp"); + assert_oracle(&combine_lp, 0.06946014693369004, Source::Empirical, "standalone_emp_combine_lp"); + + // HYBRID with data present stays on silicon (same values as the + // covered-slice oracles above). + let hyb = gb200_trtllm_db(DatabaseMode::Hybrid); + let pre = wideep_op(true, MoeQuantMode::Fp8, false).query(&hyb, 333).expect("hyb pre"); + assert_oracle( + &pre, + 0.015222300426103175 + 0.042655749106779696, + Source::Silicon, + "standalone_hyb_pre", + ); + } + /// fp8 is uncollected under NVLinkOneSided: EMPIRICAL surfaces the typed /// empirical miss (own-shape only — no transfer ladder), and fp8_block /// normalizes onto the same missing fp8 slice. Under HYBRID, Python's diff --git a/rust/aiconfigurator-core/src/operators/wideep_mla.rs b/rust/aiconfigurator-core/src/operators/wideep_mla.rs index d631d466e..64bdca46c 100644 --- a/rust/aiconfigurator-core/src/operators/wideep_mla.rs +++ b/rust/aiconfigurator-core/src/operators/wideep_mla.rs @@ -333,7 +333,14 @@ mod tests { let db = h200_sglang_db(); let mut bad = op(1); bad.attn_backend = "trtllm_mla".to_string(); - assert!(bad.query(&db, 1, 1024, 0).is_err()); + // Pin the WHITELIST error specifically: h200 has no trtllm_mla slice, + // so a plain `.is_err()` would also pass on a mere lookup miss + // (`PerfDatabase`) if the whitelist were removed. + let ctx = bad.query(&db, 1, 1024, 0); + assert!( + matches!(ctx, Err(AicError::InvalidEngineConfig(_))), + "context whitelist must reject trtllm_mla, got {ctx:?}" + ); let mut bad_gen = WideEpGenerationMlaOp::new( "wideep_gen_mla", @@ -342,7 +349,11 @@ mod tests { FmhaQuantMode::Fp8Block, ); bad_gen.attn_backend = "trtllm_mla".to_string(); - assert!(bad_gen.query(&db, 1, 1024).is_err()); + let gen = bad_gen.query(&db, 1, 1024); + assert!( + matches!(gen, Err(AicError::InvalidEngineConfig(_))), + "generation whitelist must reject trtllm_mla, got {gen:?}" + ); // fa3 stays allowed (h200 carries an fa3 slice). let mut fa3 = op(1); diff --git a/src/aiconfigurator/sdk/rust_engine_step.py b/src/aiconfigurator/sdk/rust_engine_step.py index 981ab0834..d9744b1dc 100644 --- a/src/aiconfigurator/sdk/rust_engine_step.py +++ b/src/aiconfigurator/sdk/rust_engine_step.py @@ -281,6 +281,16 @@ def _note_rust_provenance(handle: Any) -> None: util_empirical.note_provenance(tier) +def _scale_or_one(value: Any) -> float: + """Imbalance-scale forwarding: default to ``1.0`` only for ``None``. + + The Python engine path multiplies by the raw scale, so an explicit + ``0.0`` must pass through unchanged — a truthiness fallback (``or 1.0``) + would silently clobber it into ``1.0``. + """ + return 1.0 if value is None else float(value) + + def estimate_static_latency_breakdown_with_rust( model: Any, database: Any, @@ -307,8 +317,8 @@ def estimate_static_latency_breakdown_with_rust( osl=int(runtime_config.osl), prefix=int(runtime_config.prefix or 0), beam_width=int(runtime_config.beam_width or 1), - seq_imbalance_correction_scale=float(runtime_config.seq_imbalance_correction_scale or 1.0), - gen_seq_imbalance_correction_scale=float(runtime_config.gen_seq_imbalance_correction_scale or 1.0), + seq_imbalance_correction_scale=_scale_or_one(runtime_config.seq_imbalance_correction_scale), + gen_seq_imbalance_correction_scale=_scale_or_one(runtime_config.gen_seq_imbalance_correction_scale), mode=engine_mode, stride=int(stride), ) @@ -353,8 +363,8 @@ def estimate_mixed_step_latency_with_rust( int(isl), int(osl), int(prefix or 0), - seq_imbalance_correction_scale=float(seq_imbalance_correction_scale or 1.0), - gen_seq_imbalance_correction_scale=float(gen_seq_imbalance_correction_scale or 1.0), + seq_imbalance_correction_scale=_scale_or_one(seq_imbalance_correction_scale), + gen_seq_imbalance_correction_scale=_scale_or_one(gen_seq_imbalance_correction_scale), ) _note_rust_provenance(handle) return latency_ms @@ -382,7 +392,7 @@ def estimate_decode_step_latency_with_rust( int(gen_tokens), int(isl), int(osl), - gen_seq_imbalance_correction_scale=float(gen_seq_imbalance_correction_scale or 1.0), + gen_seq_imbalance_correction_scale=_scale_or_one(gen_seq_imbalance_correction_scale), ) _note_rust_provenance(handle) return latency_ms diff --git a/tests/unit/sdk/test_opspec_coverage.py b/tests/unit/sdk/test_opspec_coverage.py index 2ce05e001..2598bf401 100644 --- a/tests/unit/sdk/test_opspec_coverage.py +++ b/tests/unit/sdk/test_opspec_coverage.py @@ -75,3 +75,11 @@ def test_every_operation_converts_to_an_opspec_or_is_exempt(): stale = sorted(set(EXEMPT) & handled) assert not stale, f"EXEMPT entries now have conversion branches; remove them: {stale}" + + # EXEMPT is a closed set: every entry must name a live Operation class + # (deleted/renamed ops must drop their entry) and carry a real reason. + unknown = sorted(set(EXEMPT) - all_ops) + assert not unknown, f"EXEMPT entries are not discovered Operation classes; remove them: {unknown}" + + blank = sorted(name for name, reason in EXEMPT.items() if not reason.strip()) + assert not blank, f"EXEMPT entries need a non-empty justification: {blank}"