fix(rust-core): close all Python/Rust parity gaps — SILICON audit + HYBRID/EMPIRICAL port#1355
fix(rust-core): close all Python/Rust parity gaps — SILICON audit + HYBRID/EMPIRICAL port#1355tianhaox wants to merge 11 commits into
Conversation
…uting gate Test-first guard for porting the util-space empirical layer (PR ai-dynamo#1238, issue ai-dynamo#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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
…plumbing First half of the HYBRID/EMPIRICAL port (issue ai-dynamo#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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
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 ai-dynamo#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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
…DSV4 + silicon-view FallbackOp 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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
…rid parity green Completes the util-space empirical port (issue ai-dynamo#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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
Full-surface parity pass against the Python SDK (audit + fixes; hybrid/ empirical port stays out of scope, tracked in ai-dynamo#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 ai-dynamo#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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe change expands Rust/Python parity coverage across operator schemas, empirical performance modeling, database-mode dispatch, runtime scaling, engine bindings, cache identity, fallback handling, and CI validation. ChangesParity alignment
Estimated code review effort: 5 (Critical) | ~120 minutes Poem
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/aiconfigurator/sdk/rust_engine_step.py (2)
333-356: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winSame
x or 1.0falsy-zero issue here forgen_seq_imbalance_correction_scale.See the sibling comment on
estimate_mixed_step_latency_with_rust(Line 309-330) — same root cause, same fix.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiconfigurator/sdk/rust_engine_step.py` around lines 333 - 356, Update estimate_decode_step_latency_with_rust so gen_seq_imbalance_correction_scale preserves an explicitly provided zero instead of replacing it with 1.0; only use 1.0 when the value is None, matching the fix in estimate_mixed_step_latency_with_rust.
309-330: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
x or 1.0silently discards an explicit0.0scale.
seq_imbalance_correction_scale/gen_seq_imbalance_correction_scaleare typedfloat = 1.0(never documented asOptional), so theor 1.0guard isn't protecting against a realNonecase here — it only has the side effect of clobbering a legitimately-passed0.0into1.0, since0.0is falsy in Python. If a caller ever intentionally sets the correction scale to0.0(e.g. to zero out an attention term for analysis/ablation), this silently changes the result instead of applying it.Same pattern repeats in
estimate_decode_step_latency_with_rustbelow (Line 355). Contrast withint(prefix or 0)on Line 327, which is safe only because the default equals the falsy value.🐛 Proposed fix (explicit None-check instead of truthiness)
return handle.mixed_step_latency( int(ctx_tokens), int(gen_tokens), 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=float( + seq_imbalance_correction_scale if seq_imbalance_correction_scale is not None else 1.0 + ), + gen_seq_imbalance_correction_scale=float( + gen_seq_imbalance_correction_scale if gen_seq_imbalance_correction_scale is not None else 1.0 + ), )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aiconfigurator/sdk/rust_engine_step.py` around lines 309 - 330, Replace the truthiness fallback for seq_imbalance_correction_scale and gen_seq_imbalance_correction_scale in the mixed-step latency call with an explicit None check, preserving explicitly passed 0.0 values while retaining 1.0 only for None. Apply the same change in estimate_decode_step_latency_with_rust, and leave the safe int(prefix or 0) handling unchanged.rust/aiconfigurator-core/src/operators/dsa.rs (1)
211-243: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the missing Clippy allow on
query_cp_with.
This signature is over Clippy’s defaulttoo_many_argumentslimit, and the neighboring DSA helpers already use the same allow.🔧 Suggested change
+ #[allow(clippy::too_many_arguments)] fn query_cp_with(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/operators/dsa.rs` around lines 211 - 243, Add a #[allow(clippy::too_many_arguments)] attribute directly above query_cp_with, matching the neighboring DSA helper methods’ existing Clippy allowances.
🧹 Nitpick comments (5)
.github/workflows/build-test.yml (1)
236-245: 🚀 Performance & Scalability | 🔵 TrivialNew gating step is a solid fix for the stale-oracle gap. Wires
cargo testinto CI as the audit doc calls for, and pinningPYO3_PYTHONto the venv matches the pattern used by the preceding maturin build step.One thing worth keeping an eye on:
cargo test --release --librecompiles the crate a second time (test harness, different from thematurin develop --releaseartifact), adding to the 30-minute job budget. Consider caching~/.cargo/targetacross steps if this job starts running long — not a blocker, just future-proofing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In @.github/workflows/build-test.yml around lines 236 - 245, Optionally optimize the new Rust test step by caching the existing Cargo registry and target artifacts across workflow steps, using the established build configuration around “Run Rust unit tests (cargo test)” and the preceding maturin build. Preserve the PYO3_PYTHON setting and cargo test command while avoiding unnecessary recompilation if the CI job approaches its time budget.tests/unit/sdk/test_rust_engine_step.py (1)
377-417: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winIdentity test covers only
moe_backend+ quant collapsing — the other 7 new op-shaping fields are untested.
_engine_config_json(rust_engine_step.py:493-503) now also foldscp_style,workload_distribution,overwrite_num_layers,sms,attention_backend,enable_wideep,enable_eplb,wideep_num_slotsinto the identity, but this test only exercisesgemm_quant_modecollapsing andmoe_backend. Per path instructions fortests/**, changed behavior should be covered beyond the happy path — a regression that silently drops one of these fields from the identity (or introduces a typo, see therust_engine_step.pycomment) would go undetected.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" + + key_wideep = rust_engine_step._engine_config_json( + _model(common.GEMMQuantMode.sq, moe_backend=None), database + ) + # repeat the same pattern for enable_wideep / enable_eplb / cp_style / etc.Based on path instructions for
tests/**: "Check that tests cover the changed behavior rather than only the happy path."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/sdk/test_rust_engine_step.py` around lines 377 - 417, Expand test_engine_config_json_identity_disambiguates_collapsed_quant_modes to independently vary each newly identity-participating field—cp_style, workload_distribution, overwrite_num_layers, sms, attention_backend, enable_wideep, enable_eplb, and wideep_num_slots—and assert each variation produces a distinct _engine_config_json key from the baseline. Preserve the existing quantization and moe_backend assertions.Source: Path instructions
rust/aiconfigurator-core/src/perf_database/communication.rs (1)
147-199: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated bandwidth-scaling formula between
query_custom_allreduce_scaledandquery_nccl_scaled.Both functions compute
(x-1)/x * y/(y-1).max(1.0) * bw_a/bw_bfor their respective beyond-range correction. Worth extracting into a shared helper so future formula tweaks can't drift between the two paths.♻️ Proposed refactor: extract shared bandwidth-ratio helper
+/// Shared beyond-range bandwidth-ratio correction: `(req-1)/req * cap/(cap-1).max(1) * bw_cap/bw_req`. +fn beyond_range_bw_scale(spec: &SystemSpec, requested: u32, cap: u32) -> f64 { + let bw_cap = spec.get_p2p_bandwidth(cap); + let bw_req = spec.get_p2p_bandwidth(requested); + let f_req = requested as f64; + let f_cap = cap as f64; + (f_req - 1.0) / f_req * f_cap / (f_cap - 1.0).max(1.0) * bw_cap / bw_req +}Then both call sites reduce to
latency *= beyond_range_bw_scale(spec, tp_size, per_node);andlatency *= beyond_range_bw_scale(spec, num_gpus, max_recorded);respectively.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/perf_database/communication.rs` around lines 147 - 199, Extract the duplicated beyond-range bandwidth correction from query_custom_allreduce_scaled and query_nccl_scaled into a shared helper, such as beyond_range_bw_scale, accepting the requested and reference GPU counts plus SystemSpec. Preserve the existing formula and apply the helper at both call sites, keeping all surrounding query and boundary behavior unchanged.rust/aiconfigurator-core/src/operators/communication.rs (1)
63-82: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
p2p_bandwidthhelper. The DB-level_scaledqueries no longer call it, so it can be dropped.Suggested change
- fn p2p_bandwidth(spec: &SystemSpec, num_gpus: u32) -> f64 { - spec.get_p2p_bandwidth(num_gpus) - }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/operators/communication.rs` around lines 63 - 82, Remove the unused p2p_bandwidth helper from the communication operator implementation, leaving query and the DB-level scaled query flow unchanged.rust/aiconfigurator-core/src/perf_database/moe.rs (1)
294-346: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winLGTM! Reusing the single
kernel_sourceread for both the MXFP4 remap and the low-latency/default grid split is a clean simplification, and the remap mirrors the documented Pythonload_moe_databehavior.One suggestion: since this remap is exact-string-match driven and feeds GPU-generation-specific model selection, a small parquet-fixture test asserting the two remap arms (and the fallthrough case) would guard against a silent regression if either kernel_source string ever changes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/perf_database/moe.rs` around lines 294 - 346, The load_moe_parquet quantization remap lacks regression coverage for its exact kernel_source matches. Add a small parquet-fixture test targeting load_moe_parquet that verifies both MXFP4 remap arms produce their dedicated quantization keys and that a nonmatching quant/kernel_source pair preserves the original quantization value.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@rust/aiconfigurator-core/src/operators/dsa.rs`:
- Around line 211-243: Add a #[allow(clippy::too_many_arguments)] attribute
directly above query_cp_with, matching the neighboring DSA helper methods’
existing Clippy allowances.
In `@src/aiconfigurator/sdk/rust_engine_step.py`:
- Around line 333-356: Update estimate_decode_step_latency_with_rust so
gen_seq_imbalance_correction_scale preserves an explicitly provided zero instead
of replacing it with 1.0; only use 1.0 when the value is None, matching the fix
in estimate_mixed_step_latency_with_rust.
- Around line 309-330: Replace the truthiness fallback for
seq_imbalance_correction_scale and gen_seq_imbalance_correction_scale in the
mixed-step latency call with an explicit None check, preserving explicitly
passed 0.0 values while retaining 1.0 only for None. Apply the same change in
estimate_decode_step_latency_with_rust, and leave the safe int(prefix or 0)
handling unchanged.
---
Nitpick comments:
In @.github/workflows/build-test.yml:
- Around line 236-245: Optionally optimize the new Rust test step by caching the
existing Cargo registry and target artifacts across workflow steps, using the
established build configuration around “Run Rust unit tests (cargo test)” and
the preceding maturin build. Preserve the PYO3_PYTHON setting and cargo test
command while avoiding unnecessary recompilation if the CI job approaches its
time budget.
In `@rust/aiconfigurator-core/src/operators/communication.rs`:
- Around line 63-82: Remove the unused p2p_bandwidth helper from the
communication operator implementation, leaving query and the DB-level scaled
query flow unchanged.
In `@rust/aiconfigurator-core/src/perf_database/communication.rs`:
- Around line 147-199: Extract the duplicated beyond-range bandwidth correction
from query_custom_allreduce_scaled and query_nccl_scaled into a shared helper,
such as beyond_range_bw_scale, accepting the requested and reference GPU counts
plus SystemSpec. Preserve the existing formula and apply the helper at both call
sites, keeping all surrounding query and boundary behavior unchanged.
In `@rust/aiconfigurator-core/src/perf_database/moe.rs`:
- Around line 294-346: The load_moe_parquet quantization remap lacks regression
coverage for its exact kernel_source matches. Add a small parquet-fixture test
targeting load_moe_parquet that verifies both MXFP4 remap arms produce their
dedicated quantization keys and that a nonmatching quant/kernel_source pair
preserves the original quantization value.
In `@tests/unit/sdk/test_rust_engine_step.py`:
- Around line 377-417: Expand
test_engine_config_json_identity_disambiguates_collapsed_quant_modes to
independently vary each newly identity-participating field—cp_style,
workload_distribution, overwrite_num_layers, sms, attention_backend,
enable_wideep, enable_eplb, and wideep_num_slots—and assert each variation
produces a distinct _engine_config_json key from the baseline. Preserve the
existing quantization and moe_backend assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c711105b-9d1f-44fa-ab00-17f76ae05b94
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lockand included by none
📒 Files selected for processing (34)
.github/workflows/build-test.ymlrust/aiconfigurator-core/docs/parity-audit-2026-07-14.mdrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pyrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/common/system_spec.rsrust/aiconfigurator-core/src/engine/runtime.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/elementwise.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/mod.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/src/operators/overlap.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/perf_database/dsa.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/perf_database/wideep.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/src/session.rssrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/engine.pysrc/aiconfigurator/sdk/rust_engine_step.pytests/unit/sdk/test_rust_engine_step.py
📜 Review details
⏰ Context from checks skipped due to timeout. (6)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Cargo Deny
- GitHub Check: Rust/Python engine-step parity
- GitHub Check: Build and Test (unit)
- GitHub Check: Build and Test (e2e)
🧰 Additional context used
📓 Path-based instructions (9)
rust/aiconfigurator-core/**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
Treat
rust/aiconfigurator-core/as the Rust port of modeling operators.
Files:
rust/aiconfigurator-core/src/common/system_spec.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/overlap.rsrust/aiconfigurator-core/src/operators/mod.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/docs/parity-audit-2026-07-14.mdrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/session.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pyrust/aiconfigurator-core/src/operators/elementwise.rsrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/wideep.rsrust/aiconfigurator-core/src/perf_database/dsa.rsrust/aiconfigurator-core/src/engine/runtime.rs
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
rust/aiconfigurator-core/src/common/system_spec.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/overlap.rsrust/aiconfigurator-core/src/operators/mod.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/operators/communication.rssrc/aiconfigurator/sdk/backends/base_backend.pyrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/fpm/tests.rstests/unit/sdk/test_rust_engine_step.pyrust/aiconfigurator-core/docs/parity-audit-2026-07-14.mdrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/session.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/operators/dsa.rssrc/aiconfigurator/sdk/rust_engine_step.pyrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pyrust/aiconfigurator-core/src/operators/elementwise.rsrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/common/enums.rssrc/aiconfigurator/sdk/engine.pyrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/wideep.rsrust/aiconfigurator-core/src/perf_database/dsa.rsrust/aiconfigurator-core/src/engine/runtime.rs
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.pysrc/aiconfigurator/sdk/backends/base_backend.pytests/unit/sdk/test_rust_engine_step.pysrc/aiconfigurator/sdk/rust_engine_step.pyrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pysrc/aiconfigurator/sdk/engine.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/rust_engine_step.pysrc/aiconfigurator/sdk/engine.py
.github/workflows/**
⚙️ CodeRabbit configuration file
.github/workflows/**: - Review workflow edits for least-privilege permissions, safe trigger scopes, secret exposure, fork behavior, and release/cherry-pick side effects.
- Verify matrix jobs, artifact uploads, and scheduled support-matrix workflows still produce debuggable evidence.
Files:
.github/workflows/build-test.yml
tests/**/*.{py,yaml,txt,sh}
📄 CodeRabbit inference engine (.claude/rules/generator/testing.md)
Use integration tests for the full input-to-artifacts pipeline, comparing output against golden snapshots without external dependencies.
Files:
tests/unit/sdk/test_rust_engine_step.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Run unit tests with
pytest -m unit; run the PR build subset withpytest -m "unit or build"when applicable.
Files:
tests/unit/sdk/test_rust_engine_step.py
tests/unit/sdk/test_rust_engine_step.py
📄 CodeRabbit inference engine (AGENTS.md)
This test requires
cargowith network access to crates.io and may fail when that domain is blocked.
Files:
tests/unit/sdk/test_rust_engine_step.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/sdk/test_rust_engine_step.py
🧠 Learnings (3)
📚 Learning: 2026-06-22T17:22:01.021Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 1235
File: rust/aiconfigurator-core/parity_tests/drift_map.py:18-25
Timestamp: 2026-06-22T17:22:01.021Z
Learning: When referencing the Nemotron-3-Nano-30B model in parity scan code (e.g., when building queries/keys into `scan.sqlite`’s `entries` support matrix), use the exact canonical model identifier string: `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` (including the `NVIDIA-` prefix). Do not omit or alter the `NVIDIA-` prefix—otherwise the lookup/mapping will fail (this was the cause of the earlier typo in `parity-scan-report.md` §6, while the parity test scripts such as `drift_map.py` should keep using the correct identifier as-is).
Applied to files:
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py
📚 Learning: 2026-03-17T07:41:41.843Z
Learnt from: YijiaZhao
Repo: ai-dynamo/aiconfigurator PR: 600
File: src/aiconfigurator/sdk/backends/vllm_backend.py:42-43
Timestamp: 2026-03-17T07:41:41.843Z
Learning: In src/aiconfigurator/sdk/backends/vllm_backend.py, do not include seq_imbalance_correction_scale or gen_seq_imbalance_correction_scale in the AGG cache key isl/osl/b/ctx_tokens. These correction scales are per-instance and constant for a given usage of VLLMBackend, so they do not affect cache hits. Update the cache key computation to exclude these fields and adjust any tests that validate the cache key content. This guidance targets this file; apply similar logic to other backends only if they share the same per-instance scale semantics.
Applied to files:
src/aiconfigurator/sdk/backends/base_backend.py
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/rust_engine_step.pysrc/aiconfigurator/sdk/engine.py
🪛 ast-grep (0.44.1)
src/aiconfigurator/sdk/rust_engine_step.py
[info] 483-506: use jsonify instead of json.dumps for JSON output
Context: 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=(",", ":"),
)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 509-509: use jsonify instead of json.dumps for JSON output
Context: json.dumps(config, sort_keys=True, separators=(",", ":"))
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🪛 LanguageTool
rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md
[uncategorized] ~145-~145: The official name of this software platform is spelled with a capital “H”.
Context: ...ile-engine suite + perf gate, every PR (.github/workflows/build-test.yml:204-258), 1% ...
(GITHUB)
[style] ~162-~162: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...r the regression detector disappears. - Gate 3 (delete Python latency code): blocked...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
🔇 Additional comments (47)
rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md (1)
1-178: LGTM!rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py (1)
328-350: LGTM!Also applies to: 361-466
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py (1)
217-228: LGTM!rust/aiconfigurator-core/src/operators/dsv4.rs (1)
408-431: LGTM!rust/aiconfigurator-core/src/common/system_spec.rs (1)
164-166: 🎯 Functional CorrectnessNo change needed:
spec.gpu.mem_bwalready matchesb200_sxm.yaml.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/engine/runtime.rs (5)
25-41: LGTM!Also applies to: 222-284
319-427:mixed_step_latencythree-pass rewrite correctly mirrors_get_mix_step_latency.Verified each pass against
base_backend._get_mix_step_latency(base_backend.py:926-1060): pass-1prefix1 = prefix * (ctx_tokens / isl)usesu32floor division matchingnp.floor(ctx_tokens/isl), pass-2batch2/scale2matchceil(ctx_tokens/isl)/ceil(isl/ctx_tokens), pass-3s = isl + osl/2 + 1matches the_run_generation_phase+1convention. Thectx_tokens == 0skip in pass 2 correctly mirrors Python'snp.ceil(isl/0) == inf→latency/inf == 0behavior. Nicely documented.
429-455: LGTM!
543-558: LGTM! Hardcoding1.0/ContextOpFilter::All/falsehere is correct since the FPM telemetry bridge (rank_latency_ms) has no imbalance-scale concept, per the doc comment at Line 34-35.
593-593: LGTM! Test fixtures and call-site updates correctly track the newscale_num_tokensfield and expandedrun_generation_ops_step/mixed_step_latencysignatures.Also applies to: 627-627, 750-751, 781-781, 791-791, 818-819
rust/aiconfigurator-core/src/fpm/tests.rs (1)
37-37: LGTM! Fixture and direct-call updates correctly track the newElementwiseOp::scale_num_tokensfield and the expandedrun_context_ops/run_generation_ops_stepsignatures (hardcoded1.0/falsecorrectly matches the FPM bridge's no-scale convention).Also applies to: 71-71, 208-209, 235-236
rust/aiconfigurator-core/src/session.rs (2)
17-75: LGTM!ContextOpFilter's guarded match arms (SkipContextAttention if is_context_attention() => continue,OnlyContextAttention if !is_context_attention() => continue, else fall through to_ => {}) correctly implement include/exclude-only-context_attention semantics, andseq_imbalance_correction_scaleis threaded intoRuntimeContextas intended.
83-141: LGTM!run_generation_ops_step→run_generation_ops_step_beameddelegation withbeam_width=1correctly matches Python's_get_mix_step_latency/_get_genonly_step_latency, which always query withbeam_width=1.rust/aiconfigurator-core/src/py.rs (2)
226-278: LGTM!mixed_step_latency/decode_step_latencyPyO3 signatures (arg order + defaults) matchEngine::mixed_step_latency/Engine::decode_step_latencyand are consistent withEngineHandle's positional forwarding inengine.py.
697-704: LGTM!py_init()correctly initializes the freethreaded interpreter forcargo test(no embedding host), and the updated test call sites correctly thread the new scale arguments through.Also applies to: 719-719, 755-755, 812-812, 889-889, 916-928, 950-950
src/aiconfigurator/sdk/backends/base_backend.py (1)
19-19: LGTM! Thetry/except RustEngineUnsupportedErrorwrapping correctly falls through to the existing Python step logic on catch (noelseblocking the fallthrough, and no partial mutation of the output dicts before the exception can propagate), matching the intended "parity by delegation" fallback design.Also applies to: 424-460, 945-967, 1079-1098
tests/unit/sdk/test_rust_engine_step.py (2)
97-137: LGTM!_FakeHandleaccepting**kwargsand the updated assertions correctly verify the imbalance-scale kwargs pass through with their1.0defaults.
420-473: LGTM! TheOpConversionError→RustEngineUnsupportedErrormemoization test and the WideEP MLA per-rank-heads test both correctly validate the documented contracts.src/aiconfigurator/sdk/rust_engine_step.py (1)
471-521: 📐 Maintainability & Code QualityNo issue: these
ModelConfigfields already exist with the same names, so this lookup list matches the current schema.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/common/enums.rs (2)
129-137: LGTM!Also applies to: 153-158, 214-229, 289-376, 444-454
253-254: 🎯 Functional CorrectnessNo code change needed
ModelFamily::Gemma4MixandMINIMAXM3are already wired through the enum andas_str()mapping; the remainingGEMMA4MOEmention is in a docs-only parity audit note.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/operators/attention.rs (1)
223-281: LGTM!Also applies to: 363-373
rust/aiconfigurator-core/src/operators/moe.rs (1)
48-63: LGTM!Also applies to: 91-93, 103-111, 121-158, 275-277
src/aiconfigurator/sdk/engine.py (2)
37-37: LGTM!Also applies to: 70-70, 97-98, 154-162, 201-203, 262-288, 332-346, 382-386, 485-497, 500-508, 538-557, 558-652, 689-699, 914-999
307-318: 🎯 Functional Correctness
_scale_num_tokensis always initialized onMoEDispatch, so this read is safe.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/operators/overlap.rs (1)
4-10: LGTM!rust/aiconfigurator-core/src/perf_database/communication.rs (1)
11-16: LGTM!Also applies to: 31-31, 121-136, 201-236, 310-327, 448-557
rust/aiconfigurator-core/src/perf_database/wideep_moe.rs (1)
111-201: LGTM!Also applies to: 216-286, 304-365
rust/aiconfigurator-core/src/operators/elementwise.rs (1)
25-37: LGTM!Also applies to: 48-57
rust/aiconfigurator-core/src/operators/op.rs (1)
24-24: LGTM!Also applies to: 126-131, 216-216, 297-297
rust/aiconfigurator-core/src/operators/mod.rs (1)
47-47: LGTM!rust/aiconfigurator-core/src/engine/spec.rs (1)
164-165: LGTM!Also applies to: 205-205, 244-244, 303-305, 327-327, 389-389, 504-519, 575-575, 611-611
rust/aiconfigurator-core/src/operators/wideep_mla.rs (1)
34-47: LGTM!Also applies to: 92-92, 128-133, 142-144, 155-170, 201-226, 265-265
rust/aiconfigurator-core/src/operators/moe_dispatch.rs (2)
100-142: 🎯 Functional CorrectnessNo token scaling change needed
TrtLLMWideEPMoEDispatchOp::queryalready matches the Python WideEP path: it passesnum_tokensdirectly toquery_trtllm_alltoall, and that model sizes the alltoall message fromnum_tokens,topk,num_experts, andmoe_ep_sizerather thanattention_dp_size/moe_tp_size.> Likely an incorrect or invalid review comment.
74-79: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winGuard
scale_num_tokensin the MoE dispatch payload
src/aiconfigurator/sdk/engine.pyshould normalize this like_elementwise;#[serde(default = ...)]only helps when the key is missing, not whenscale_num_tokensis emitted asnull.Suggested Change
- "scale_num_tokens": op._scale_num_tokens, + "scale_num_tokens": op._scale_num_tokens if op._scale_num_tokens else 1,> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/operators/dsa.rs (4)
44-142: LGTM!full_fracamortization logic inquery_context(CP and non-CP branches) correctly mirrors the documentedw*full + (1-w)*skipblend, and the early-return atw >= 1.0avoids an unnecessary skip-table query for the default (pure-full) architectures.
144-161: 📐 Maintainability & Code Quality | ⚡ Quick winStale doc:
query_context_cpno longer always usesskip_indexer=False.The doc comment says this delegates to
query_cp_withwithskip_indexer=Falseverbatim, butquery_contextnow calls this withskip_indexer=truetoo (the CP +full_frac<1.0blend branch at lines 104-108). Worth a one-line update so future readers don't assume the CP path never exercises the skip table.📝 Suggested doc tweak
- /// 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): + /// Wires the real data dependencies and delegates to [`Self::query_cp_with`] + /// (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:
270-337: LGTM! Delta/AG-KV skip gating (lines 270-298) and thequery_generationamortization (lines 304-337) correctly mirror the context-side logic and the documented Python skip-indexer semantics.
357-436: LGTM! Test updates correctly thread the newskip_indexer/full_fracparams through without changing existing expectations.rust/aiconfigurator-core/src/operators/wideep_moe.rs (2)
81-104: LGTM! Rounding the interpolated token count before feeding it to the integralsol_latency_ms(instead of the previous linear-proxy|t| t) is a reasonable way to keep floor-division parity with Python's integer-typed inputs.
105-137: 🎯 Functional CorrectnessNo truncation mismatch here Python’s
get_soluses//foropsand everymem_bytesterm, so the Rustu64arithmetic matches the oracle. The missing extrapolation test is a separate gap, not a bug in this function.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/operators/gemm.rs (1)
68-110: LGTM! Flooring at the GEMM's own SOL roofline (instead of 0) for thefp8_staticpath is a sound fix, and the non-fp8_staticpath is untouched (latency_floorstays at 0, so behavior there is unchanged).rust/aiconfigurator-core/src/perf_database/dsa.rs (1)
57-64: LGTM! The skip-indexer table plumbing (newOnceLockcaches,load_dsa_parquet'swant_skip_rowsrow filter, and the analytic-SOL zeroing indsa_context_sol_ms) is internally consistent and correctly mirrors the documented Python skip-indexer semantics, with tests updated to match the new signatures.Also applies to: 219-222, 312-374, 387-444, 446-508, 600-694, 803-844, 1116-1810
rust/aiconfigurator-core/src/perf_database/dsv4.rs (2)
124-135: LGTM! DroppingarchitecturefromModuleKeyand gatingmla_dtypereads viakey_on_fmhacorrectly mirrors the documented Python keying split (context keys on[fmha][kv][gemm], generation on[kv][gemm]only), andselect_resolved'sfmha: Option<FmhaQuantMode>cleanly represents the sentinel-empty generation case.Also applies to: 867-894, 1126-1197
348-375: 🎯 Functional CorrectnessNo change needed in
rust/aiconfigurator-core/src/operators/dsv4.rsquery_generationalready uses the shortened call signature here.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/perf_database/wideep.rs (2)
83-103: LGTM! The newAlltoallGrids/AlltoallKeymodel,select_alltoall_kernel's backend/topology auto-selection, andload_alltoall_parquet's defaulting (num_nodesdefault algebraically matches the query-sidenode_numformula) are well documented and backed by GB200 exact-hit + kernel-selection tests.Also applies to: 295-368, 766-878, 1224-1282
201-263: 🎯 Functional CorrectnessNo action here The WideEP MoE callers already pass an analytic
solclosure (MoeOp::queryandWideEpMoeOp::query); the identity proxy only shows up in the smoke test.> Likely an incorrect or invalid review comment.
…D port = full parity Integrates the SILICON-path parity audit into the HYBRID/EMPIRICAL util-empirical port so the compiled engine reaches full Python parity in one branch. Conflict-resolution policy: the audit branch's silicon semantics are authoritative (fp8_static SOL floor, _scaled comm queries with fan-out/GB200 handling inside the DB layer, DSV4 [kv][gemm] generation keying, GLM-5.2 skip-indexer tables + full_frac blend, the full trtllm dispatch branch, EPLB/deepep MoE routing); the hybrid branch's mode dispatch and empirical estimators are re-grafted on top, with skip_indexer threaded through the DSA dispatch/empirical layer and grid cache keys. Two integration decisions of note: - perf_database/wideep_moe keeps the hybrid branch's level-wise distribution resolution (kernel -> quant -> distribution, then exact shape) over the audit branch's any-shape fallback, per the Python oracle verification that found the latter divergent. - WideEpMoeOp re-resolves the compute kernel at QUERY time (Python `_select_kernel(database, quant)` sees loaded tables); the opspec kernel_source stays as a wire hint only — the compile-time resolution is load-order dependent (returns the preferred kernel on a fresh database) and broke the TrtllmWideEp compile-parity case. Verification on the merged tree: - engine-step parity: 299/299 (64 HYBRID/EMPIRICAL surfaces at rtol=1e-4 + the full SILICON suite incl. the audit's new cases) - compile-engine parity + perf gate: 58/58 - cargo test --lib: 260 passed, 0 failed - Python unit suite: 2075 passed Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
…D port = full parity Integrates the SILICON-path parity audit into the HYBRID/EMPIRICAL util-empirical port so the compiled engine reaches full Python parity in one branch. Conflict-resolution policy: the audit branch's silicon semantics are authoritative (fp8_static SOL floor, _scaled comm queries with fan-out/GB200 handling inside the DB layer, DSV4 [kv][gemm] generation keying, GLM-5.2 skip-indexer tables + full_frac blend, the full trtllm dispatch branch, EPLB/deepep MoE routing); the hybrid branch's mode dispatch and empirical estimators are re-grafted on top, with skip_indexer threaded through the DSA dispatch/empirical layer and grid cache keys. Two integration decisions of note: - perf_database/wideep_moe keeps the hybrid branch's level-wise distribution resolution (kernel -> quant -> distribution, then exact shape) over the audit branch's any-shape fallback, per the Python oracle verification that found the latter divergent. - WideEpMoeOp re-resolves the compute kernel at QUERY time (Python `_select_kernel(database, quant)` sees loaded tables); the opspec kernel_source stays as a wire hint only — the compile-time resolution is load-order dependent (returns the preferred kernel on a fresh database) and broke the TrtllmWideEp compile-parity case. Verification on the merged tree: - engine-step parity: 299/299 (64 HYBRID/EMPIRICAL surfaces at rtol=1e-4 + the full SILICON suite incl. the audit's new cases) - compile-engine parity + perf gate: 58/58 - cargo test --lib: 260 passed, 0 failed - Python unit suite: 2075 passed Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
tianhaox
changed the title
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
rust/aiconfigurator-core/src/operators/util_empirical.rs (1)
272-305: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache builders run while holding the
Mutex, serializing unrelated keys.
UtilGridCache::get_or_try_buildandDeltaLookupCache::get_or_try_buildboth keep the lock held across thebuilder()call, which triggers parquet loads + grid construction. Since these caches areArc-shared across all mode-views of onePerfDatabase(seemod.rs'ssilicon_view), any thread building gridAblocks every other thread's lookup of gridB,C, ... — even ones already cached. If the engine ever evaluates configurations concurrently against a sharedPerfDatabase, this becomes a global choke point during warm-up.♻️ Proposed two-phase lock/build/insert
pub fn get_or_try_build<F>(&self, key: &str, builder: F) -> Result<Option<Arc<UtilGrid>>, AicError> where F: FnOnce() -> Result<Option<UtilGrid>, AicError>, { - let mut grids = self.grids.lock().expect("util grid cache poisoned"); - if let Some(cached) = grids.get(key) { - return Ok(cached.clone()); - } - let built = builder()?.map(Arc::new); - grids.insert(key.to_string(), built.clone()); - Ok(built) + if let Some(cached) = self.grids.lock().expect("util grid cache poisoned").get(key) { + return Ok(cached.clone()); + } + // Build outside the lock so a slow load doesn't block unrelated keys. + let built = builder()?.map(Arc::new); + let mut grids = self.grids.lock().expect("util grid cache poisoned"); + Ok(grids.entry(key.to_string()).or_insert(built).clone()) }The same restructuring applies to
DeltaLookupCache::get_or_try_build.Also applies to: 416-441
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/operators/util_empirical.rs` around lines 272 - 305, Update UtilGridCache::get_or_try_build and DeltaLookupCache::get_or_try_build to release their Mutex before invoking the builder, then reacquire it to insert the built result. Recheck the cache after rebuilding so concurrent builders do not overwrite an existing value, while preserving memoization of successful None results and avoiding memoization of errors.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/aiconfigurator-core/src/common/enums.rs`:
- Around line 90-96: Update the Default implementation for TransferPolicy so its
default value matches the documented ALL semantics rather than the derived
all-false OFF state. Remove the derived Default behavior and implement or
delegate Default to TransferPolicy::ALL, preserving the existing Clone, Copy,
Debug, PartialEq, and Eq derives.
In `@rust/aiconfigurator-core/src/common/error.rs`:
- Around line 63-75: Narrow AicError::is_missing_perf_data() so it matches
PerfDatabase errors and only Io errors representing a missing file, excluding
permission, read, and corruption failures. Update the Io pattern to inspect its
underlying error kind while preserving the existing fallback behavior for
genuine file-not-found cases.
In `@rust/aiconfigurator-core/src/operators/msa.rs`:
- Around line 60-101: Update MSA::query_context and query_generation to reject
or safely handle non-positive dsa_scale_k before dividing, preventing negative
or infinite latency; preserve the existing empirical error path for invalid
transfer data. Apply the sibling operators’ result-shaping convention by
chaining clamp_non_negative() (and the established scaling helper where
applicable) before constructing each PerformanceResult, including the Sol and
empirical return paths.
---
Nitpick comments:
In `@rust/aiconfigurator-core/src/operators/util_empirical.rs`:
- Around line 272-305: Update UtilGridCache::get_or_try_build and
DeltaLookupCache::get_or_try_build to release their Mutex before invoking the
builder, then reacquire it to insert the built result. Recheck the cache after
rebuilding so concurrent builders do not overwrite an existing value, while
preserving memoization of successful None results and avoiding memoization of
errors.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 51011e29-7465-4e62-8fba-91f3499c2326
📒 Files selected for processing (41)
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/common/error.rsrust/aiconfigurator-core/src/config.rsrust/aiconfigurator-core/src/engine/runtime.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/memory.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/mod.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/attention.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/perf_database/dsa.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/gemm.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/perf_database/perf_interp.rsrust/aiconfigurator-core/src/perf_database/wideep.rsrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/tests/memory_round_trip.rssrc/aiconfigurator/sdk/engine.pysrc/aiconfigurator/sdk/rust_engine_step.pytests/unit/sdk/test_rust_engine_step.py
🚧 Files skipped from review as they are similar to previous changes (11)
- rust/aiconfigurator-core/src/operators/mod.rs
- tests/unit/sdk/test_rust_engine_step.py
- rust/aiconfigurator-core/src/perf_database/wideep_moe.rs
- rust/aiconfigurator-core/src/operators/wideep_mla.rs
- rust/aiconfigurator-core/src/py.rs
- rust/aiconfigurator-core/src/perf_database/communication.rs
- rust/aiconfigurator-core/src/operators/moe.rs
- src/aiconfigurator/sdk/rust_engine_step.py
- rust/aiconfigurator-core/src/operators/moe_dispatch.rs
- rust/aiconfigurator-core/src/perf_database/wideep.rs
- rust/aiconfigurator-core/src/engine/runtime.rs
👮 Files not reviewed due to content moderation or server errors (7)
- rust/aiconfigurator-core/src/perf_database/attention.rs
- rust/aiconfigurator-core/src/operators/communication.rs
- rust/aiconfigurator-core/src/operators/dsa.rs
- rust/aiconfigurator-core/src/operators/dsv4.rs
- rust/aiconfigurator-core/src/operators/attention.rs
- rust/aiconfigurator-core/src/perf_database/dsv4.rs
- rust/aiconfigurator-core/src/perf_database/dsa.rs
📜 Review details
⏰ Context from checks skipped due to timeout. (4)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Build and Test (e2e)
- GitHub Check: Rust/Python engine-step parity
🧰 Additional context used
📓 Path-based instructions (4)
rust/aiconfigurator-core/**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
Treat
rust/aiconfigurator-core/as the Rust port of modeling operators.
Files:
rust/aiconfigurator-core/src/memory.rsrust/aiconfigurator-core/tests/memory_round_trip.rsrust/aiconfigurator-core/src/common/error.rsrust/aiconfigurator-core/src/perf_database/perf_interp.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/src/config.rsrust/aiconfigurator-core/src/perf_database/gemm.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/perf_database/attention.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/dsa.rs
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
rust/aiconfigurator-core/src/memory.rsrust/aiconfigurator-core/tests/memory_round_trip.rsrust/aiconfigurator-core/src/common/error.rsrust/aiconfigurator-core/src/perf_database/perf_interp.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/src/config.rsrust/aiconfigurator-core/src/perf_database/gemm.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/perf_database/attention.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/communication.rssrc/aiconfigurator/sdk/engine.pyrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/dsa.rs
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.pysrc/aiconfigurator/sdk/engine.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/engine.py
🧠 Learnings (3)
📚 Learning: 2026-07-09T16:27:36.054Z
Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 1337
File: rust/aiconfigurator-core/src/perf_database/mla.rs:691-751
Timestamp: 2026-07-09T16:27:36.054Z
Learning: For the MLA module loaders in `perf_database` (Rust: `load_context_module_parquet`/`load_generation_module_parquet`; Python: `load_context_mla_module_data`/`load_generation_mla_module_data`), reviews should assume and preserve the intended `last-wins-across-sources` behavior: when the same key appears in multiple sources within the priority-ordered source list, values from later (higher-priority) sources overwrite earlier ones. Do not flag or “fix” these loaders to use the DSA loader’s different semantics (e.g., two-phase first-source-wins / last-row-wins-within-file). If any refactor changes the merge/overwrite semantics, require updating the in-code documentation/comments and adjusting/adding tests to cover the overwrite/duplicate-row behavior (notably the incident that motivated last-wins).
Applied to files:
rust/aiconfigurator-core/src/perf_database/mla.rs
📚 Learning: 2026-06-22T17:22:01.021Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 1235
File: rust/aiconfigurator-core/parity_tests/drift_map.py:18-25
Timestamp: 2026-06-22T17:22:01.021Z
Learning: When referencing the Nemotron-3-Nano-30B model in parity scan code (e.g., when building queries/keys into `scan.sqlite`’s `entries` support matrix), use the exact canonical model identifier string: `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` (including the `NVIDIA-` prefix). Do not omit or alter the `NVIDIA-` prefix—otherwise the lookup/mapping will fail (this was the cause of the earlier typo in `parity-scan-report.md` §6, while the parity test scripts such as `drift_map.py` should keep using the correct identifier as-is).
Applied to files:
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/engine.py
🪛 OpenGrep (1.25.0)
rust/aiconfigurator-core/src/operators/wideep_moe.rs
[ERROR] 339-339: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 370-370: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 374-374: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
rust/aiconfigurator-core/src/operators/mla.rs
[ERROR] 674-674: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 678-678: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 755-755: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🔇 Additional comments (34)
rust/aiconfigurator-core/tests/memory_round_trip.rs (1)
108-109: LGTM!rust/aiconfigurator-core/src/fpm/tests.rs (1)
111-113: LGTM!rust/aiconfigurator-core/src/config.rs (1)
24-26: LGTM!Also applies to: 77-90, 212-215
rust/aiconfigurator-core/src/perf_database/moe.rs (2)
82-105: LGTM!Also applies to: 284-386, 664-686
363-386: 🗄️ Data Integrity & IntegrationNo action needed for the new quant names.
w4a8_mxfp4_mxfp8_trtllmandw4a16_mxfp4_cutlassare already inALL_MOE_QUANTS, somoe_quant_from_name()round-trips them correctly.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py (2)
22-22: LGTM!Also applies to: 51-56, 602-627, 651-690, 693-732, 751-830, 833-882, 1065-1299
770-780: 🎯 Functional Correctness
common.MoEQuantMode[...]matches the enum names here.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/engine/spec.rs (1)
160-172: LGTM!Also applies to: 200-208, 237-246, 289-307, 309-329, 378-412, 507-540, 564-642, 644-676
rust/aiconfigurator-core/src/common/enums.rs (2)
69-84: LGTM!Also applies to: 201-230, 286-301, 361-450, 514-528
334-352: 🗄️ Data Integrity & IntegrationNo action needed: both Rust and Python already use
GEMMA4MIXandMINIMAXM3.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/common/error.rs (1)
37-42: LGTM!rust/aiconfigurator-core/src/perf_database/perf_interp.rs (1)
683-694: LGTM!rust/aiconfigurator-core/src/perf_database/mhc.rs (1)
132-168: LGTM!rust/aiconfigurator-core/src/perf_database/gemm.rs (1)
173-231: LGTM!Also applies to: 366-366
rust/aiconfigurator-core/src/operators/util_empirical.rs (2)
37-270: LGTM!
307-414: LGTM!Also applies to: 442-622
rust/aiconfigurator-core/src/operators/gemm.rs (3)
182-269: LGTM!
358-396: LGTM!
100-106: 🎯 Functional CorrectnessNo action needed
Fp8StaticandFp8share the same SOL mapping here, so passing the raw quant intogemm_sol_latency_msis fine as-is.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/perf_database/mod.rs (1)
84-146: LGTM!Also applies to: 204-280
rust/aiconfigurator-core/src/perf_database/wideep_mla.rs (1)
137-226: LGTM!Also applies to: 228-289, 304-401, 580-597
rust/aiconfigurator-core/src/operators/msa.rs (2)
147-331: LGTM!
333-448: LGTM!rust/aiconfigurator-core/src/operators/mhc.rs (1)
115-180: LGTM!Also applies to: 262-340
rust/aiconfigurator-core/src/operators/wideep_moe.rs (1)
56-69: LGTM!Also applies to: 101-263, 266-395
rust/aiconfigurator-core/src/perf_database/mla.rs (1)
329-458: LGTM!Also applies to: 511-628, 958-975
rust/aiconfigurator-core/src/operators/op.rs (1)
24-24: LGTM!Also applies to: 113-117, 131-136, 213-214, 223-223, 296-297, 306-306, 338-380
rust/aiconfigurator-core/src/operators/mla.rs (2)
125-245: LGTM!Also applies to: 247-598, 600-874
93-100: 🩺 Stability & Availability
isl == 0is already excluded upstream.
engine/runtime.rsrejectsprefix >= islbefore constructingRuntimeContext, so this branch can’t receive a zero-length context chunk through the engine path.> Likely an incorrect or invalid review comment.src/aiconfigurator/sdk/engine.py (4)
52-52: LGTM!Also applies to: 64-64, 362-382, 637-640, 1032-1061
93-97: 🗄️ Data Integrity & IntegrationNo change needed
787-793: 🗄️ Data Integrity & IntegrationNo issue:
PerfDatabasealready exposes both accessors.get_default_database_mode()and thetransfer_policyproperty are defined, so the engine helpers read the configured mode and policy correctly.> Likely an incorrect or invalid review comment.
617-624: 🎯 Functional CorrectnessMSA branches are fine here.
ContextMSAModuleandGenerationMSAModuleinherit from_BaseMSAModule, so the DSAisinstancechecks won’t shadow them.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/memory.rs (1)
393-409: LGTM!
| 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 { .. }) | ||
| } | ||
| } | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find every construction site of AicError::Io to see whether it's scoped to "file not found"
# or wraps arbitrary io::Error kinds (permission, corruption, etc.).
rg -n --type=rust 'AicError::Io' rust/aiconfigurator-core/src -B3 -A1
rg -n --type=rust 'fn load_gemm_parquet|fn load_two_d_parquet|fn resolve_op_sources' rust/aiconfigurator-core/src -A15Repository: ai-dynamo/aiconfigurator
Length of output: 11000
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== error.rs =="
sed -n '1,180p' rust/aiconfigurator-core/src/common/error.rs
echo
echo "== perf_database/parquet_loader.rs =="
sed -n '1,140p' rust/aiconfigurator-core/src/perf_database/parquet_loader.rs
echo
echo "== perf_database/mod.rs =="
sed -n '1,220p' rust/aiconfigurator-core/src/perf_database/mod.rs
echo
echo "== operators/op.rs fallback =="
sed -n '330,390p' rust/aiconfigurator-core/src/operators/op.rsRepository: ai-dynamo/aiconfigurator
Length of output: 20050
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== AicError usages in loaders =="
rg -n --type=rust 'AicError::Io \{|AicError::PerfDatabase\(' rust/aiconfigurator-core/src/perf_database rust/aiconfigurator-core/src/operators rust/aiconfigurator-core/src/common -A2 -B2Repository: ai-dynamo/aiconfigurator
Length of output: 50380
Narrow is_missing_perf_data() to file-not-found I/O.
PerfReader::open() maps both open and read failures into AicError::Io, and op.rs treats every Io as a safe fallback trigger. That can turn permission/read/corruption failures into empirical estimates instead of surfacing a hard error.
Suggested Change
pub fn is_missing_perf_data(&self) -> bool {
- matches!(self, Self::PerfDatabase(_) | Self::Io { .. })
+ matches!(self, Self::PerfDatabase(_))
+ || matches!(self, Self::Io { source, .. } if source.kind() == std::io::ErrorKind::NotFound)
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 { .. }) | |
| } | |
| } | |
| 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(_)) | |
| || matches!(self, Self::Io { source, .. } if source.kind() == std::io::ErrorKind::NotFound) | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/aiconfigurator-core/src/common/error.rs` around lines 63 - 75, Narrow
AicError::is_missing_perf_data() so it matches PerfDatabase errors and only Io
errors representing a missing file, excluding permission, read, and corruption
failures. Update the Io pattern to inspect its underlying error kind while
preserving the existing fallback behavior for genuine file-not-found cases.
There was a problem hiding this comment.
Agreed directionally — Python only treats PerfDataNotAvailableError/InterpolationDataNotAvailableError as fallback-able (_MISSING_SILICON_DATA_EXCEPTIONS, perf_database.py:29); OSErrors surface as hard errors. But the one-line narrowing is ineffective on its own: every table loader memoizes load failures behind clone_err, which collapses ALL variants (incl. Parquet/Io) to PerfDatabase(msg) before is_missing_perf_data ever sees them (13 identical clone_err sites), and several loaders match Io | PerfDatabase directly (gemm.rs:769, mhc.rs:304, mla.rs:1020, dsv4.rs:1301, wideep.rs:1459, op.rs:374). The proper fix is an error-classification pass that preserves the class through the caches, then narrows Io to NotFound — error-path only, no latency effect on valid data. Deferred to the freeze follow-up: #1357
There was a problem hiding this comment.
@tianhaox, thanks for the thorough investigation — that's a much more complete picture than my one-line suggestion. Makes sense that narrowing Io alone is moot while clone_err collapses all cached load failures (including Parquet/Io) down to PerfDatabase(msg) before is_missing_perf_data() ever runs, and that the direct Io | PerfDatabase match sites need the same classification-preserving treatment. Deferring the full error-classification pass to #1357 sounds like the right call for this freeze — error-path only, no impact on valid-data latency in the meantime.
✏️ Learnings added
Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 1355
File: rust/aiconfigurator-core/src/common/error.rs:63-75
Timestamp: 2026-07-14T16:43:06.967Z
Learning: In rust/aiconfigurator-core, table loaders (e.g., gemm.rs, mhc.rs, mla.rs, dsv4.rs, wideep.rs) memoize load failures via a shared `clone_err` helper that collapses all `AicError` variants (including `Parquet` and `Io`) into `AicError::PerfDatabase(msg)` before `AicError::is_missing_perf_data()` (in rust/aiconfigurator-core/src/common/error.rs) ever inspects them. Additionally, several call sites (gemm.rs:769, mhc.rs:304, mla.rs:1020, dsv4.rs:1301, wideep.rs:1459, op.rs:374) match `Io | PerfDatabase` directly. Narrowing `is_missing_perf_data()`'s `Io` arm to `ErrorKind::NotFound` alone is therefore ineffective; a full error-classification pass that preserves the error class through these caches (then narrows `Io` to `NotFound`) is needed. This is tracked as a follow-up in https://github.com/ai-dynamo/aiconfigurator/issues/1357.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/aiconfigurator-core/src/operators/moe_dispatch.rs (1)
117-141: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRoute the standalone WideEP dispatch op through mode-aware querying.
This production path calls the silicon-only DB API directly, bypassing the new EMPIRICAL/HYBRID handling in
query_alltoall_table. Consequently, EMPIRICAL returns silicon results and HYBRID misses cannot fall back.Proposed fix
pub fn query(&self, db: &PerfDatabase, num_tokens: u32) -> Result<PerformanceResult, AicError> { - let spec: &SystemSpec = &db.system_spec; - let q = |op_name: &str| { - db.wideep.query_trtllm_alltoall( - spec, + let q = |op_name: &str| { + query_alltoall_table( + db, op_name, num_tokens, self.hidden_size, self.topk, 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)) }As per path instructions, this fix is concrete and limited to the commented hunk.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/operators/moe_dispatch.rs` around lines 117 - 141, Update the query closure in MoeDispatch::query to use the mode-aware all-toall querying path, query_alltoall_table, instead of db.wideep.query_trtllm_alltoall. Preserve the existing operation names, parameters, and latency-selection logic so EMPIRICAL and HYBRID modes receive their intended results and fallback behavior.Source: Path instructions
🧹 Nitpick comments (2)
rust/aiconfigurator-core/src/common/enums.rs (1)
90-98: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMake
TransferPolicy::default()returnALL.The derived default disables every transfer, while omitted wire values and newly loaded databases default to ALL. Rust-side default construction can therefore silently change estimation behavior.
Proposed change
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct TransferPolicy { pub xshape: bool, pub xquant: bool, pub xprofile: bool, pub xop: bool, } +impl Default for TransferPolicy { + fn default() -> Self { + Self::ALL + } +} + impl TransferPolicy {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/common/enums.rs` around lines 90 - 98, Implement an explicit Default for TransferPolicy so TransferPolicy::default() returns the existing ALL policy with every transfer flag enabled. Remove the derived Default implementation while preserving the current Clone, Copy, Debug, PartialEq, and Eq derives.rust/aiconfigurator-core/src/operators/wideep_mla.rs (1)
318-330: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the whitelist-specific failure.
H200 has no
trtllm_mlaslice, sois_err()still passes ifcheck_attn_backendis removed and lookup merely misses. Assert the validation error message, or use data containing that backend. A one-click suggestion is unsafe because the helper’s exact error text is not included here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/operators/wideep_mla.rs` around lines 318 - 330, Strengthen the assertions in the H200 backend validation tests around op and WideEpGenerationMlaOp::query so they verify the whitelist-specific validation error, not merely any failure. Inspect check_attn_backend for the exact error text and assert it in both bad.query calls, or use fixture data containing trtllm_mla so lookup cannot fail for an unrelated missing slice.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py`:
- Around line 1107-1125: Update the shared parity error comparator and the
affected EngineStepParityCase declarations for the transfer-policy
error-symmetry cases. Require both engines’ _ErrorSentinel.kind values to
normalize to and match EmpiricalNotImplemented, rather than treating any pair of
exceptions as equivalent; preserve existing comparison behavior for cases
without an expected error kind.
In `@rust/aiconfigurator-core/src/common/error.rs`:
- Around line 71-72: Update Error::is_missing_perf_data to classify only
file-not-found I/O errors, while retaining PerfDatabase errors as missing
performance data. Match the specific NotFound kind within the Io variant and
ensure permission or transient I/O failures return false so they propagate
normally.
In `@rust/aiconfigurator-core/src/perf_database/attention.rs`:
- Around line 289-321: Update context_head_sizes and the corresponding accessor
at the referenced second API to accept window_size and restrict candidate
aggregation to entries matching that window before reference head-size
selection. Propagate the new parameter through all API callers and update tests
to cover mixed coverage where one head size exists only at window 0 and another
exists at the requested window, ensuring the same-window candidate is selected.
In `@rust/aiconfigurator-core/src/perf_database/wideep_moe.rs`:
- Around line 174-261: Preserve perf-file insertion order by retaining row-order
metadata during loading, then update
rust/aiconfigurator-core/src/perf_database/wideep_moe.rs:174-261 so
available_kernels and resolve_slice use first-seen kernel/distribution order for
fallbacks instead of BTreeMap order; update
rust/aiconfigurator-core/src/perf_database/perf_interp.rs:687-692 to flatten
samples in source order for empirical tie-breaking; and update
rust/aiconfigurator-core/src/perf_database/moe.rs:324-376 to enumerate sibling
slices and quant modes in source order.
---
Outside diff comments:
In `@rust/aiconfigurator-core/src/operators/moe_dispatch.rs`:
- Around line 117-141: Update the query closure in MoeDispatch::query to use the
mode-aware all-toall querying path, query_alltoall_table, instead of
db.wideep.query_trtllm_alltoall. Preserve the existing operation names,
parameters, and latency-selection logic so EMPIRICAL and HYBRID modes receive
their intended results and fallback behavior.
---
Nitpick comments:
In `@rust/aiconfigurator-core/src/common/enums.rs`:
- Around line 90-98: Implement an explicit Default for TransferPolicy so
TransferPolicy::default() returns the existing ALL policy with every transfer
flag enabled. Remove the derived Default implementation while preserving the
current Clone, Copy, Debug, PartialEq, and Eq derives.
In `@rust/aiconfigurator-core/src/operators/wideep_mla.rs`:
- Around line 318-330: Strengthen the assertions in the H200 backend validation
tests around op and WideEpGenerationMlaOp::query so they verify the
whitelist-specific validation error, not merely any failure. Inspect
check_attn_backend for the exact error text and assert it in both bad.query
calls, or use fixture data containing trtllm_mla so lookup cannot fail for an
unrelated missing slice.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 51011e29-7465-4e62-8fba-91f3499c2326
📒 Files selected for processing (41)
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/common/error.rsrust/aiconfigurator-core/src/config.rsrust/aiconfigurator-core/src/engine/runtime.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/memory.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/mod.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/attention.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/perf_database/dsa.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/gemm.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/perf_database/perf_interp.rsrust/aiconfigurator-core/src/perf_database/wideep.rsrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/tests/memory_round_trip.rssrc/aiconfigurator/sdk/engine.pysrc/aiconfigurator/sdk/rust_engine_step.pytests/unit/sdk/test_rust_engine_step.py
🚧 Files skipped from review as they are similar to previous changes (8)
- rust/aiconfigurator-core/src/operators/mod.rs
- tests/unit/sdk/test_rust_engine_step.py
- rust/aiconfigurator-core/src/py.rs
- src/aiconfigurator/sdk/rust_engine_step.py
- rust/aiconfigurator-core/src/perf_database/wideep.rs
- rust/aiconfigurator-core/src/engine/runtime.rs
- rust/aiconfigurator-core/src/operators/dsa.rs
- rust/aiconfigurator-core/src/operators/moe.rs
📜 Review details
🧰 Additional context used
📓 Path-based instructions (4)
rust/aiconfigurator-core/**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
Treat
rust/aiconfigurator-core/as the Rust port of modeling operators.
Files:
rust/aiconfigurator-core/src/common/error.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/tests/memory_round_trip.rsrust/aiconfigurator-core/src/memory.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/perf_database/gemm.rsrust/aiconfigurator-core/src/perf_database/perf_interp.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/config.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/perf_database/attention.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/perf_database/dsa.rs
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
rust/aiconfigurator-core/src/common/error.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/tests/memory_round_trip.rsrust/aiconfigurator-core/src/memory.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/perf_database/gemm.rsrust/aiconfigurator-core/src/perf_database/perf_interp.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/config.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/perf_database/attention.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/operators/mla.rssrc/aiconfigurator/sdk/engine.pyrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/perf_database/dsa.rs
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.pysrc/aiconfigurator/sdk/engine.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/engine.py
🧠 Learnings (3)
📚 Learning: 2026-07-09T16:27:36.054Z
Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 1337
File: rust/aiconfigurator-core/src/perf_database/mla.rs:691-751
Timestamp: 2026-07-09T16:27:36.054Z
Learning: For the MLA module loaders in `perf_database` (Rust: `load_context_module_parquet`/`load_generation_module_parquet`; Python: `load_context_mla_module_data`/`load_generation_mla_module_data`), reviews should assume and preserve the intended `last-wins-across-sources` behavior: when the same key appears in multiple sources within the priority-ordered source list, values from later (higher-priority) sources overwrite earlier ones. Do not flag or “fix” these loaders to use the DSA loader’s different semantics (e.g., two-phase first-source-wins / last-row-wins-within-file). If any refactor changes the merge/overwrite semantics, require updating the in-code documentation/comments and adjusting/adding tests to cover the overwrite/duplicate-row behavior (notably the incident that motivated last-wins).
Applied to files:
rust/aiconfigurator-core/src/perf_database/mla.rs
📚 Learning: 2026-06-22T17:22:01.021Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 1235
File: rust/aiconfigurator-core/parity_tests/drift_map.py:18-25
Timestamp: 2026-06-22T17:22:01.021Z
Learning: When referencing the Nemotron-3-Nano-30B model in parity scan code (e.g., when building queries/keys into `scan.sqlite`’s `entries` support matrix), use the exact canonical model identifier string: `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` (including the `NVIDIA-` prefix). Do not omit or alter the `NVIDIA-` prefix—otherwise the lookup/mapping will fail (this was the cause of the earlier typo in `parity-scan-report.md` §6, while the parity test scripts such as `drift_map.py` should keep using the correct identifier as-is).
Applied to files:
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/engine.py
🪛 OpenGrep (1.25.0)
rust/aiconfigurator-core/src/operators/wideep_moe.rs
[ERROR] 339-339: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 370-370: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 374-374: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
rust/aiconfigurator-core/src/operators/mla.rs
[ERROR] 674-674: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 678-678: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 755-755: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🔇 Additional comments (33)
rust/aiconfigurator-core/src/common/error.rs (1)
37-42: LGTM!rust/aiconfigurator-core/tests/memory_round_trip.rs (1)
108-109: LGTM!rust/aiconfigurator-core/src/perf_database/gemm.rs (1)
173-232: LGTM!Also applies to: 366-372
rust/aiconfigurator-core/src/config.rs (1)
24-26: LGTM!Also applies to: 77-90, 212-215
rust/aiconfigurator-core/src/common/enums.rs (1)
69-89: LGTM!Also applies to: 99-139, 201-230, 286-301, 325-348, 361-448, 516-526
rust/aiconfigurator-core/src/operators/communication.rs (1)
7-174: LGTM!Also applies to: 219-341, 384-558
rust/aiconfigurator-core/src/operators/util_empirical.rs (1)
1-622: LGTM!rust/aiconfigurator-core/src/perf_database/mhc.rs (1)
132-168: LGTM!rust/aiconfigurator-core/src/memory.rs (1)
407-408: LGTM!rust/aiconfigurator-core/src/fpm/tests.rs (1)
37-37: LGTM!Also applies to: 71-71, 111-112, 210-211, 237-238
rust/aiconfigurator-core/src/operators/wideep_mla.rs (1)
20-26: LGTM!Also applies to: 38-51, 96-124, 135-151, 162-162, 176-176, 208-310, 312-317, 331-337, 376-581
rust/aiconfigurator-core/src/perf_database/wideep_mla.rs (1)
160-167: LGTM!Also applies to: 216-219, 228-275, 304-339, 381-401, 580-597
src/aiconfigurator/sdk/engine.py (1)
37-37: LGTM!Also applies to: 52-72, 93-101, 157-165, 204-206, 265-291, 310-313, 341-346, 362-382, 406-409, 513-516, 528-528, 562-581, 621-624, 637-640, 717-726, 809-835, 1032-1061
rust/aiconfigurator-core/src/operators/msa.rs (1)
1-448: LGTM!rust/aiconfigurator-core/src/engine/spec.rs (1)
164-165: LGTM!Also applies to: 205-205, 244-244, 303-305, 327-327, 389-410, 525-540, 588-589, 598-598, 626-627, 636-636, 672-673
rust/aiconfigurator-core/src/operators/op.rs (2)
24-24: LGTM!Also applies to: 213-214, 223-223, 296-297, 306-306, 345-360
113-136: 🗄️ Data Integrity & IntegrationDrop this schema-version concern
ENGINE_SPEC_SCHEMA_VERSIONis already bumped to3alongside this op enum change, and the version-gate test covers rejection before op decoding.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py (1)
22-22: LGTM!Also applies to: 51-56, 223-234, 624-626, 671-673, 715-717, 751-820, 835-835, 864-871, 1067-1106, 1127-1172, 1185-1247
rust/aiconfigurator-core/src/operators/wideep_moe.rs (1)
6-10: LGTM!Also applies to: 23-38, 56-69, 108-263, 266-394
rust/aiconfigurator-core/src/operators/moe_dispatch.rs (1)
24-29: LGTM!Also applies to: 75-116, 145-809, 1032-1032, 1083-1083, 1119-1320
rust/aiconfigurator-core/src/perf_database/mla.rs (1)
329-458: LGTM!Also applies to: 511-541, 559-559, 585-585, 602-602, 958-975
rust/aiconfigurator-core/src/perf_database/moe.rs (1)
82-105: LGTM!Also applies to: 284-323, 379-386, 451-473, 482-482
rust/aiconfigurator-core/src/perf_database/dsv4.rs (1)
124-131: LGTM!Also applies to: 303-303, 355-375, 412-430, 526-583, 925-937, 1049-1049, 1184-1225, 1334-1334, 1357-1361, 1387-1387, 1834-1834
rust/aiconfigurator-core/src/perf_database/communication.rs (1)
11-16: LGTM!Also applies to: 31-31, 121-121, 137-200, 215-215, 238-327, 400-404, 416-416, 542-542, 560-560, 574-574, 620-620, 636-636
rust/aiconfigurator-core/src/perf_database/wideep_moe.rs (1)
98-172: LGTM!Also applies to: 262-283, 345-358, 407-451
rust/aiconfigurator-core/src/perf_database/mod.rs (1)
14-20: LGTM!Also applies to: 84-147, 204-280
rust/aiconfigurator-core/src/operators/mhc.rs (1)
14-17: LGTM!Also applies to: 115-179, 262-340
rust/aiconfigurator-core/src/operators/gemm.rs (1)
18-22: LGTM!Also applies to: 81-110, 125-270, 358-396
rust/aiconfigurator-core/src/perf_database/attention.rs (1)
258-287: LGTM!Also applies to: 324-354, 381-404, 601-601, 633-681, 707-707
rust/aiconfigurator-core/src/operators/attention.rs (1)
18-31: LGTM!Also applies to: 53-113, 169-196, 224-224, 271-285, 302-307, 323-323, 333-357, 359-757, 838-986
rust/aiconfigurator-core/src/operators/mla.rs (1)
9-24: LGTM!Also applies to: 77-101, 131-133, 176-186, 199-207, 239-244, 247-598, 648-873
rust/aiconfigurator-core/src/operators/dsv4.rs (1)
20-25: LGTM!Also applies to: 215-249, 266-360, 378-383, 413-416, 535-641, 871-1069
rust/aiconfigurator-core/src/perf_database/dsa.rs (1)
57-64: LGTM!Also applies to: 219-222, 312-365, 387-407, 446-520, 524-539, 550-572, 664-676, 704-707, 737-738, 764-764, 867-910, 1180-1874
| pub fn is_missing_perf_data(&self) -> bool { | ||
| matches!(self, Self::PerfDatabase(_) | Self::Io { .. }) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Only treat file-not-found I/O as missing performance data.
Io also covers permission and transient filesystem failures. HYBRID currently hides those failures behind empirical results instead of surfacing broken data access.
Proposed change
- matches!(self, Self::PerfDatabase(_) | Self::Io { .. })
+ matches!(self, Self::PerfDatabase(_))
+ || matches!(
+ self,
+ Self::Io { source, .. }
+ if source.kind() == std::io::ErrorKind::NotFound
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| pub fn is_missing_perf_data(&self) -> bool { | |
| matches!(self, Self::PerfDatabase(_) | Self::Io { .. }) | |
| pub fn is_missing_perf_data(&self) -> bool { | |
| matches!(self, Self::PerfDatabase(_)) | |
| || matches!( | |
| self, | |
| Self::Io { source, .. } | |
| if source.kind() == std::io::ErrorKind::NotFound | |
| ) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/aiconfigurator-core/src/common/error.rs` around lines 71 - 72, Update
Error::is_missing_perf_data to classify only file-not-found I/O errors, while
retaining PerfDatabase errors as missing performance data. Match the specific
NotFound kind within the Io variant and ensure permission or transient I/O
failures return false so they propagate normally.
There was a problem hiding this comment.
Same finding as the sibling comment on lines 63-75 — deferred to #1357: the narrowing alone is ineffective because the loaders' clone_err memoization collapses all variants to PerfDatabase before classification reaches is_missing_perf_data; it needs the error-classification pass described there (preserve the class through the caches, then narrow Io to NotFound).
There was a problem hiding this comment.
@tianhaox Understood, thanks for the context. Deferring to #1357 makes sense here since a narrow fix to is_missing_perf_data alone wouldn't help while clone_err memoization already collapses the error class upstream. I'll leave this as tracked follow-up work in that issue rather than pushing for a partial fix here.
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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rust/aiconfigurator-core/src/operators/moe.rs (1)
247-249: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winBump
ENGINE_SPEC_SCHEMA_VERSIONfor the newMoeOpfields (rust/aiconfigurator-core/src/operators/moe.rs:247-249)
MoeOpis part of the bincode-backedOpSpecpayload inEngineSpec; addingmoe_backend,enable_eplb, andis_contextchanges the wire layout and will break older blobs unless this is version-gated.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/operators/moe.rs` around lines 247 - 249, Increment ENGINE_SPEC_SCHEMA_VERSION to reflect the newly added MoeOp fields moe_backend, enable_eplb, and is_context in the bincode-backed EngineSpec payload. Update the version wherever the schema version is defined or validated, while preserving existing version-gating behavior for older blobs.
🧹 Nitpick comments (2)
rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs (1)
172-287: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a synthetic-data unit test for the loader's validation paths.
Only
moe_dtype_from_nameand the missing-file case are unit-tested; the invariant checks (used_cuda_graph/includes_gate_topk/includes_routed_scale) and duplicate-leaf detection inload_module_parquetaren't directly covered here (likely only indirectly via the parity suite's real data).Also applies to: 293-334
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs` around lines 172 - 287, Add a synthetic-data unit test covering load_module_parquet validation paths, including failures for each boolean invariant (used_cuda_graph, includes_gate_topk, and includes_routed_scale) and duplicate (key, num_tokens) leaves. Reuse the existing test utilities and assert that each invalid dataset returns an AicError, while preserving the current missing-file and moe_dtype_from_name tests.rust/aiconfigurator-core/src/engine/runtime.rs (1)
269-309: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a
beam_width > 1case tonextn_scales_decode_batch. The current assertion still only exercises the defaultbeam_width = 1path, so it doesn’t cover the new beamed generation branch.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/engine/runtime.rs` around lines 269 - 309, Update the tests for nextn_scales_decode_batch to add a case with beam_width greater than 1, exercising the beamed generation path in run_generation_phase and validating its expected result alongside the existing beam_width = 1 coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs`:
- Around line 159-165: Update load_module and its cached-result error handling
so re-emitting a failed load preserves the original AicError variant, including
Parquet and schema-corruption errors, instead of converting everything through
clone_err to AicError::PerfDatabase. Use an owned cached error representation if
required, while retaining the existing successful module-loading behavior.
---
Outside diff comments:
In `@rust/aiconfigurator-core/src/operators/moe.rs`:
- Around line 247-249: Increment ENGINE_SPEC_SCHEMA_VERSION to reflect the newly
added MoeOp fields moe_backend, enable_eplb, and is_context in the
bincode-backed EngineSpec payload. Update the version wherever the schema
version is defined or validated, while preserving existing version-gating
behavior for older blobs.
---
Nitpick comments:
In `@rust/aiconfigurator-core/src/engine/runtime.rs`:
- Around line 269-309: Update the tests for nextn_scales_decode_batch to add a
case with beam_width greater than 1, exercising the beamed generation path in
run_generation_phase and validating its expected result alongside the existing
beam_width = 1 coverage.
In `@rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs`:
- Around line 172-287: Add a synthetic-data unit test covering
load_module_parquet validation paths, including failures for each boolean
invariant (used_cuda_graph, includes_gate_topk, and includes_routed_scale) and
duplicate (key, num_tokens) leaves. Reuse the existing test utilities and assert
that each invalid dataset returns an AicError, while preserving the current
missing-file and moe_dtype_from_name tests.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e1b3f72f-e2d1-4570-9f70-ac9d798dc362
📒 Files selected for processing (31)
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/engine/runtime.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/mod.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/perf_database/parquet_loader.rsrust/aiconfigurator-core/src/perf_database/wideep.rsrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/src/py.rssrc/aiconfigurator/sdk/engine.pysrc/aiconfigurator/sdk/operations/moe.pysrc/aiconfigurator/sdk/rust_engine_step.pytests/unit/sdk/database/test_moe_mla.pytests/unit/sdk/test_rust_engine_step.py
🚧 Files skipped from review as they are similar to previous changes (21)
- rust/aiconfigurator-core/src/operators/mhc.rs
- rust/aiconfigurator-core/src/operators/mod.rs
- rust/aiconfigurator-core/src/operators/gemm.rs
- rust/aiconfigurator-core/src/perf_database/wideep_moe.rs
- rust/aiconfigurator-core/src/perf_database/wideep_mla.rs
- rust/aiconfigurator-core/src/perf_database/moe.rs
- rust/aiconfigurator-core/src/operators/op.rs
- src/aiconfigurator/sdk/rust_engine_step.py
- rust/aiconfigurator-core/src/operators/util_empirical.rs
- rust/aiconfigurator-core/src/operators/msa.rs
- rust/aiconfigurator-core/src/operators/wideep_mla.rs
- rust/aiconfigurator-core/src/engine/spec.rs
- rust/aiconfigurator-core/src/operators/attention.rs
- rust/aiconfigurator-core/src/operators/communication.rs
- rust/aiconfigurator-core/src/operators/dsv4.rs
- rust/aiconfigurator-core/src/operators/dsa.rs
- rust/aiconfigurator-core/src/operators/mla.rs
- rust/aiconfigurator-core/src/perf_database/mod.rs
- rust/aiconfigurator-core/src/operators/moe_dispatch.rs
- src/aiconfigurator/sdk/engine.py
- rust/aiconfigurator-core/src/perf_database/wideep.rs
📜 Review details
⏰ Context from checks skipped due to timeout. (4)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Rust/Python engine-step parity
- GitHub Check: Build and Test (e2e)
🧰 Additional context used
📓 Path-based instructions (8)
rust/aiconfigurator-core/**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
Treat
rust/aiconfigurator-core/as the Rust port of modeling operators.
Files:
rust/aiconfigurator-core/src/perf_database/parquet_loader.rsrust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rsrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/engine/runtime.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/wideep_moe.rs
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
rust/aiconfigurator-core/src/perf_database/parquet_loader.rssrc/aiconfigurator/sdk/operations/moe.pytests/unit/sdk/database/test_moe_mla.pyrust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rstests/unit/sdk/test_rust_engine_step.pyrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/engine/runtime.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/wideep_moe.rs
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
src/aiconfigurator/sdk/operations/moe.pytests/unit/sdk/database/test_moe_mla.pytests/unit/sdk/test_rust_engine_step.pyrust/aiconfigurator-core/parity_tests/test_engine_step_parity.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/operations/moe.py
tests/**/*.{py,yaml,txt,sh}
📄 CodeRabbit inference engine (.claude/rules/generator/testing.md)
Use integration tests for the full input-to-artifacts pipeline, comparing output against golden snapshots without external dependencies.
Files:
tests/unit/sdk/database/test_moe_mla.pytests/unit/sdk/test_rust_engine_step.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Run unit tests with
pytest -m unit; run the PR build subset withpytest -m "unit or build"when applicable.
Files:
tests/unit/sdk/database/test_moe_mla.pytests/unit/sdk/test_rust_engine_step.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/sdk/database/test_moe_mla.pytests/unit/sdk/test_rust_engine_step.py
tests/unit/sdk/test_rust_engine_step.py
📄 CodeRabbit inference engine (AGENTS.md)
This test requires
cargowith network access to crates.io and may fail when that domain is blocked.
Files:
tests/unit/sdk/test_rust_engine_step.py
🧠 Learnings (2)
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/operations/moe.py
📚 Learning: 2026-06-22T17:22:01.021Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 1235
File: rust/aiconfigurator-core/parity_tests/drift_map.py:18-25
Timestamp: 2026-06-22T17:22:01.021Z
Learning: When referencing the Nemotron-3-Nano-30B model in parity scan code (e.g., when building queries/keys into `scan.sqlite`’s `entries` support matrix), use the exact canonical model identifier string: `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` (including the `NVIDIA-` prefix). Do not omit or alter the `NVIDIA-` prefix—otherwise the lookup/mapping will fail (this was the cause of the earlier typo in `parity-scan-report.md` §6, while the parity test scripts such as `drift_map.py` should keep using the correct identifier as-is).
Applied to files:
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py
🪛 OpenGrep (1.25.0)
rust/aiconfigurator-core/src/operators/moe.rs
[ERROR] 1141-1141: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
rust/aiconfigurator-core/src/operators/wideep_moe.rs
[ERROR] 447-447: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🔇 Additional comments (9)
rust/aiconfigurator-core/src/operators/moe.rs (1)
6-176: LGTM!Also applies to: 259-353, 370-737, 830-1173
rust/aiconfigurator-core/src/operators/wideep_moe.rs (1)
6-10: LGTM!Also applies to: 23-38, 56-69, 108-142, 152-180, 182-236, 238-265, 266-450
rust/aiconfigurator-core/src/perf_database/parquet_loader.rs (1)
214-227: LGTM!src/aiconfigurator/sdk/operations/moe.py (1)
2266-2276: LGTM!tests/unit/sdk/test_rust_engine_step.py (1)
47-90: LGTM!Also applies to: 93-143, 146-199
rust/aiconfigurator-core/src/py.rs (1)
49-113: LGTM!Also applies to: 257-360, 1051-1133
rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py (1)
1306-1404: LGTM!rust/aiconfigurator-core/src/engine/runtime.rs (1)
21-42: LGTM!Also applies to: 169-178, 187-207, 247-267, 344-452, 460-480, 568-583, 618-618, 652-652, 695-696, 777-778, 808-808, 818-818
tests/unit/sdk/database/test_moe_mla.py (1)
917-954: 🎯 Functional CorrectnessNo change needed:
EmpiricalNotImplementedErroris already imported here. The regression test is targeted and correctly covers the HYBRID fallback path.> Likely an incorrect or invalid review comment.
…rust CODEOWNERS 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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.claude/rules/rust-core/parity.md:
- Around line 1-2: Add YAML frontmatter to the Rust-Core Parity Discipline rule,
including a paths list covering the relevant Python sources, Rust sources, and
parity tests. Keep the existing rule content unchanged and scope it to those
paths rather than leaving it always-on; do not use an always-on rule without
explicit human approval.
In `@tests/unit/sdk/test_opspec_coverage.py`:
- Around line 27-30: Strengthen the EXEMPT validation and stale-check logic in
the coverage test: require each exemption reason to be non-empty, and assert
every EXEMPT operation name exists in all_ops so deleted or renamed entries fail
the gate. Update the stale calculation/assertion to detect exemptions missing
from all_ops while preserving the requirement that unhandled operations are
either convertible or explicitly justified.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 158d7369-b3e2-4edf-9453-2b7c92895a39
📒 Files selected for processing (3)
.claude/rules/rust-core/parity.md.github/CODEOWNERStests/unit/sdk/test_opspec_coverage.py
📜 Review details
🧰 Additional context used
📓 Path-based instructions (7)
tests/**/*.{py,yaml,txt,sh}
📄 CodeRabbit inference engine (.claude/rules/generator/testing.md)
Use integration tests for the full input-to-artifacts pipeline, comparing output against golden snapshots without external dependencies.
Files:
tests/unit/sdk/test_opspec_coverage.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
tests/unit/sdk/test_opspec_coverage.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Run unit tests with
pytest -m unit; run the PR build subset withpytest -m "unit or build"when applicable.
Files:
tests/unit/sdk/test_opspec_coverage.py
tests/unit/sdk/test_opspec_coverage.py
📄 CodeRabbit inference engine (.claude/rules/rust-core/parity.md)
Every new operation must be convertible to an opspec or carry a justified
EXEMPTentry; the opspec coverage test must pass.
Files:
tests/unit/sdk/test_opspec_coverage.py
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
tests/unit/sdk/test_opspec_coverage.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/sdk/test_opspec_coverage.py
**/.claude/rules/**/*.md
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
New rule files must include
paths:frontmatter; do not add new always-on rules without human approval.
Files:
.claude/rules/rust-core/parity.md
🧠 Learnings (1)
📚 Learning: 2026-04-28T16:18:19.130Z
Learnt from: liyuanzhe1991
Repo: ai-dynamo/aiconfigurator PR: 903
File: src/aiconfigurator/sdk/perf_database.py:3529-3568
Timestamp: 2026-04-28T16:18:19.130Z
Learning: In src/aiconfigurator/sdk/perf_database.py, ensure DSA context interpolation uses the pre-extrapolation (raw) context data for the piecewise topk-regime-safe path. Specifically: build the raw_dsa_dict from self._raw_context_dsa_module_data (deep-copied before _extrapolate_data_grid) and call _interp_dsa_context_topk_piecewise_from_raw(...). Only if this call returns None should you fall back to _interp_3d using the extrapolated self._context_dsa_module_data. Reviewers should flag any interpolation logic that blends across regimes around index_topk, which can underestimate latency just above the boundary.
Applied to files:
.github/CODEOWNERS
🔇 Additional comments (3)
.github/CODEOWNERS (1)
19-23: LGTM!tests/unit/sdk/test_opspec_coverage.py (2)
1-25: LGTM!Also applies to: 45-60
65-65: 🎯 Functional CorrectnessAvoid coupling the coverage gate to one source spelling.
Line 65 recognizes only the exact text
isinstance(op, Name). Valid branches using qualified names, tuple types, or different formatting will be reported as missing, while unrelated matching expressions could be counted as coverage. Please verify the current_to_opspecbranch shapes and consider AST-based or runtime-based discovery; a one-click suggestion is not safe until that contract is confirmed.
ee6557e to
6a9e260
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
rust/aiconfigurator-core/src/operators/mla.rs (1)
257-288: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the repeated database-mode dispatch skeleton into a shared helper.
The same
match db.database_mode { Empirical => estimate, Hybrid => try silicon then fall back on a typed miss, _ => silicon }skeleton is copy-pasted 7 times across these files (and, per each file's own doc comment, mirrors an 8th copy ingemm.rs). This is precisely the pattern of duplicated logic that this PR's own audit doc found silently drifting elsewhere (e.g.or_insertvsinsertinwideep_moe.rs's dedup). A small generic helper collapses all sites to one definition:fn dispatch_database_mode<S, E>( mode: DatabaseMode, silicon: S, empirical: E, ) -> Result<(f64, Source), AicError> where S: FnOnce() -> Result<f64, AicError>, E: FnOnce() -> Result<f64, AicError>, { match mode { DatabaseMode::Empirical => Ok((empirical()?, Source::Empirical)), DatabaseMode::Hybrid => match silicon() { Ok(latency) => Ok((latency, Source::Silicon)), Err(err) if err.is_missing_perf_data() => Ok((empirical()?, Source::Empirical)), Err(err) => Err(err), }, _ => Ok((silicon()?, Source::Silicon)), } }
rust/aiconfigurator-core/src/operators/mla.rs#L257-L288: replacequery_context_mla_table's match body withdispatch_database_mode(db.database_mode, silicon, || context_mla_empirical(db, b, s, prefix, num_heads, kv_quant, fmha_quant)).rust/aiconfigurator-core/src/operators/mla.rs#L332-L354: same collapse inquery_generation_mla_tableusinggeneration_mla_empirical.rust/aiconfigurator-core/src/operators/mla.rs#L384-L406: same collapse inquery_mla_bmm_tableusingmla_bmm_empirical.rust/aiconfigurator-core/src/operators/mla.rs#L463-L497: same collapse inquery_context_mla_module_tableusingcontext_mla_module_empirical.rust/aiconfigurator-core/src/operators/mla.rs#L546-L578: same collapse inquery_generation_mla_module_tableusinggeneration_mla_module_empirical.rust/aiconfigurator-core/src/operators/moe.rs#L259-L284: same collapse inMoeOp::queryusingself.silicon_pr(db, num_tokens)andself.empirical_latency(db, num_tokens).rust/aiconfigurator-core/src/operators/wideep_moe.rs#L108-L122: same collapse inWideEpMoeOp::queryusingself.silicon_latency(db, scaled)andself.empirical_latency(db, scaled).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/aiconfigurator-core/src/operators/mla.rs` around lines 257 - 288, Repeated database-mode dispatch should be centralized in a shared generic helper. Add dispatch_database_mode with the existing Empirical, Hybrid missing-data fallback, and Silicon behaviors, then replace the matches in rust/aiconfigurator-core/src/operators/mla.rs:257-288, 332-354, 384-406, 463-497, and 546-578 using the corresponding empirical functions; update rust/aiconfigurator-core/src/operators/moe.rs:259-284 to use silicon_pr and empirical_latency, and rust/aiconfigurator-core/src/operators/wideep_moe.rs:108-122 to use silicon_latency and empirical_latency. Preserve each closure’s existing arguments and error behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/sdk/test_rust_engine_step.py`:
- Around line 534-539: Remove the obsolete zero-assertion
test_non_silicon_database_mode_falls_back_to_python_step entirely; retain
test_sol_database_modes_fall_back_to_python_step as the authoritative coverage
for HYBRID/EMPIRICAL routing.
---
Nitpick comments:
In `@rust/aiconfigurator-core/src/operators/mla.rs`:
- Around line 257-288: Repeated database-mode dispatch should be centralized in
a shared generic helper. Add dispatch_database_mode with the existing Empirical,
Hybrid missing-data fallback, and Silicon behaviors, then replace the matches in
rust/aiconfigurator-core/src/operators/mla.rs:257-288, 332-354, 384-406,
463-497, and 546-578 using the corresponding empirical functions; update
rust/aiconfigurator-core/src/operators/moe.rs:259-284 to use silicon_pr and
empirical_latency, and
rust/aiconfigurator-core/src/operators/wideep_moe.rs:108-122 to use
silicon_latency and empirical_latency. Preserve each closure’s existing
arguments and error behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3eec47cf-3f8e-4f93-be6a-7aa6d7cb7caa
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lockand included by none
📒 Files selected for processing (56)
.claude/rules/rust-core/parity.md.github/CODEOWNERS.github/workflows/build-test.ymlrust/aiconfigurator-core/docs/parity-audit-2026-07-14.mdrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pyrust/aiconfigurator-core/parity_tests/test_engine_step_parity.pyrust/aiconfigurator-core/src/common/enums.rsrust/aiconfigurator-core/src/common/error.rsrust/aiconfigurator-core/src/common/system_spec.rsrust/aiconfigurator-core/src/config.rsrust/aiconfigurator-core/src/engine/runtime.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/fpm/tests.rsrust/aiconfigurator-core/src/memory.rsrust/aiconfigurator-core/src/operators/attention.rsrust/aiconfigurator-core/src/operators/communication.rsrust/aiconfigurator-core/src/operators/dsa.rsrust/aiconfigurator-core/src/operators/dsv4.rsrust/aiconfigurator-core/src/operators/elementwise.rsrust/aiconfigurator-core/src/operators/gemm.rsrust/aiconfigurator-core/src/operators/mhc.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/mod.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/moe_dispatch.rsrust/aiconfigurator-core/src/operators/msa.rsrust/aiconfigurator-core/src/operators/op.rsrust/aiconfigurator-core/src/operators/overlap.rsrust/aiconfigurator-core/src/operators/util_empirical.rsrust/aiconfigurator-core/src/operators/wideep_mla.rsrust/aiconfigurator-core/src/operators/wideep_moe.rsrust/aiconfigurator-core/src/perf_database/attention.rsrust/aiconfigurator-core/src/perf_database/communication.rsrust/aiconfigurator-core/src/perf_database/dsa.rsrust/aiconfigurator-core/src/perf_database/dsv4.rsrust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rsrust/aiconfigurator-core/src/perf_database/gemm.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/perf_database/mod.rsrust/aiconfigurator-core/src/perf_database/moe.rsrust/aiconfigurator-core/src/perf_database/parquet_loader.rsrust/aiconfigurator-core/src/perf_database/perf_interp.rsrust/aiconfigurator-core/src/perf_database/wideep.rsrust/aiconfigurator-core/src/perf_database/wideep_mla.rsrust/aiconfigurator-core/src/perf_database/wideep_moe.rsrust/aiconfigurator-core/src/py.rsrust/aiconfigurator-core/src/session.rsrust/aiconfigurator-core/tests/memory_round_trip.rssrc/aiconfigurator/sdk/backends/base_backend.pysrc/aiconfigurator/sdk/engine.pysrc/aiconfigurator/sdk/operations/moe.pysrc/aiconfigurator/sdk/rust_engine_step.pytests/unit/sdk/database/test_moe_mla.pytests/unit/sdk/test_opspec_coverage.pytests/unit/sdk/test_rust_engine_step.py
🚧 Files skipped from review as they are similar to previous changes (45)
- .github/workflows/build-test.yml
- tests/unit/sdk/test_opspec_coverage.py
- rust/aiconfigurator-core/src/operators/overlap.rs
- .github/CODEOWNERS
- rust/aiconfigurator-core/src/memory.rs
- rust/aiconfigurator-core/src/fpm/tests.rs
- tests/unit/sdk/database/test_moe_mla.py
- src/aiconfigurator/sdk/backends/base_backend.py
- rust/aiconfigurator-core/tests/memory_round_trip.rs
- rust/aiconfigurator-core/src/perf_database/dsv4_megamoe.rs
- rust/aiconfigurator-core/src/operators/mhc.rs
- rust/aiconfigurator-core/src/perf_database/parquet_loader.rs
- src/aiconfigurator/sdk/operations/moe.py
- rust/aiconfigurator-core/src/operators/msa.rs
- rust/aiconfigurator-core/src/perf_database/perf_interp.rs
- .claude/rules/rust-core/parity.md
- rust/aiconfigurator-core/src/operators/util_empirical.rs
- rust/aiconfigurator-core/src/operators/elementwise.rs
- rust/aiconfigurator-core/src/operators/mod.rs
- rust/aiconfigurator-core/src/perf_database/moe.rs
- rust/aiconfigurator-core/src/common/error.rs
- rust/aiconfigurator-core/src/perf_database/gemm.rs
- rust/aiconfigurator-core/src/operators/op.rs
- rust/aiconfigurator-core/src/perf_database/wideep_moe.rs
- rust/aiconfigurator-core/src/config.rs
- rust/aiconfigurator-core/src/perf_database/attention.rs
- rust/aiconfigurator-core/src/operators/dsv4.rs
- rust/aiconfigurator-core/src/operators/gemm.rs
- rust/aiconfigurator-core/src/perf_database/wideep_mla.rs
- rust/aiconfigurator-core/src/operators/wideep_mla.rs
- rust/aiconfigurator-core/src/common/enums.rs
- rust/aiconfigurator-core/src/operators/attention.rs
- rust/aiconfigurator-core/src/session.rs
- rust/aiconfigurator-core/src/perf_database/mod.rs
- rust/aiconfigurator-core/src/engine/runtime.rs
- rust/aiconfigurator-core/parity_tests/test_engine_step_parity.py
- rust/aiconfigurator-core/src/perf_database/dsv4.rs
- rust/aiconfigurator-core/src/perf_database/communication.rs
- src/aiconfigurator/sdk/engine.py
- rust/aiconfigurator-core/src/operators/communication.rs
- rust/aiconfigurator-core/src/operators/dsa.rs
- rust/aiconfigurator-core/src/operators/moe_dispatch.rs
- rust/aiconfigurator-core/src/perf_database/dsa.rs
- rust/aiconfigurator-core/src/py.rs
- rust/aiconfigurator-core/src/perf_database/wideep.rs
📜 Review details
⏰ Context from checks skipped due to timeout. (6)
- GitHub Check: Collect snapshot (old)
- GitHub Check: Collect snapshot (new)
- GitHub Check: Cargo Deny
- GitHub Check: Build and Test (e2e)
- GitHub Check: Build and Test (unit)
- GitHub Check: Rust/Python engine-step parity
🧰 Additional context used
📓 Path-based instructions (11)
rust/aiconfigurator-core/**/*
📄 CodeRabbit inference engine (.claude/rules/repo-guide.md)
Treat
rust/aiconfigurator-core/as the Rust port of modeling operators.
Files:
rust/aiconfigurator-core/src/common/system_spec.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pyrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/docs/parity-audit-2026-07-14.mdrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/wideep_moe.rs
rust/aiconfigurator-core/src/**/*.rs
📄 CodeRabbit inference engine (.claude/rules/rust-core/parity.md)
Anchor Rust parity-sensitive changes with a Python-generated oracle in a Rust
#[cfg(test)]module, using 1e-9 tolerance and auv run pythoninvocation against ashared_layer=Falseview; alternatively or additionally add a case toparity_tests/test_engine_step_parity.pywhen a new configuration class becomes reachable.
Files:
rust/aiconfigurator-core/src/common/system_spec.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/engine/spec.rsrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/wideep_moe.rs
**/*
⚙️ CodeRabbit configuration file
**/*: - Prefer applicable inline comments. When the correct fix is clear, small, and limited to the commented diff hunk, include it as a GitHub Suggested Change so the author can apply it with one click.
- Do not use a suggested change when the fix requires broader design choices, multiple files, generated artifacts, unavailable context, or validation that cannot be inferred from the diff.
- If a comment is not directly applicable, state the smallest concrete next step and why a one-click suggestion is not safe.
Files:
rust/aiconfigurator-core/src/common/system_spec.rsrust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pyrust/aiconfigurator-core/src/perf_database/mla.rsrust/aiconfigurator-core/src/engine/spec.rstests/unit/sdk/test_rust_engine_step.pyrust/aiconfigurator-core/docs/parity-audit-2026-07-14.mdrust/aiconfigurator-core/src/operators/mla.rsrust/aiconfigurator-core/src/operators/moe.rsrust/aiconfigurator-core/src/operators/wideep_moe.rssrc/aiconfigurator/sdk/rust_engine_step.py
rust/aiconfigurator-core/src/**/perf_database/*.rs
📄 CodeRabbit inference engine (.claude/rules/rust-core/parity.md)
Preserve Python selection semantics exactly, including fallback order and tie-breaks. Because Python dictionaries preserve insertion order while Rust
BTreeMapsorts keys, retain an explicit Rust load-order record for first-available fallbacks.
Files:
rust/aiconfigurator-core/src/perf_database/mhc.rsrust/aiconfigurator-core/src/perf_database/mla.rs
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use the project’s
uv-managed environment and run linting withruff check .andruff format --check ..
Files:
rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.pytests/unit/sdk/test_rust_engine_step.pysrc/aiconfigurator/sdk/rust_engine_step.py
rust/aiconfigurator-core/src/engine/spec.rs
📄 CodeRabbit inference engine (.claude/rules/rust-core/parity.md)
Add an engine/spec round-trip fixture for every new operation.
Files:
rust/aiconfigurator-core/src/engine/spec.rs
tests/**/*.{py,yaml,txt,sh}
📄 CodeRabbit inference engine (.claude/rules/generator/testing.md)
Use integration tests for the full input-to-artifacts pipeline, comparing output against golden snapshots without external dependencies.
Files:
tests/unit/sdk/test_rust_engine_step.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Run unit tests with
pytest -m unit; run the PR build subset withpytest -m "unit or build"when applicable.
Files:
tests/unit/sdk/test_rust_engine_step.py
tests/unit/sdk/test_rust_engine_step.py
📄 CodeRabbit inference engine (AGENTS.md)
This test requires
cargowith network access to crates.io and may fail when that domain is blocked.
Files:
tests/unit/sdk/test_rust_engine_step.py
tests/**
⚙️ CodeRabbit configuration file
tests/**: - Check that tests cover the changed behavior rather than only the happy path.
- Watch for fixtures or golden outputs that mask backend drift, support-matrix ordering changes, or CLI output regressions.
Files:
tests/unit/sdk/test_rust_engine_step.py
src/aiconfigurator/sdk/**
⚙️ CodeRabbit configuration file
src/aiconfigurator/sdk/**: - Verify SDK API changes remain compatible with generator inputs, profiler data flow, and documented examples.
- Flag silent schema or field-name drift between SDK models and generator/module bridge code.
Files:
src/aiconfigurator/sdk/rust_engine_step.py
🧠 Learnings (4)
📓 Common learnings
Learnt from: CR
Repo: ai-dynamo/aiconfigurator
Timestamp: 2026-07-14T16:08:43.081Z
Learning: Keep Python and Rust engine behavior in numeric parity; treat the Python op/query math as the frozen reference and keep the Rust engine-step parity CI job green.
Learnt from: CR
Repo: ai-dynamo/aiconfigurator
Timestamp: 2026-07-14T16:08:43.081Z
Learning: Do not port or otherwise alter the intentional Python/Rust splits for SOL and SOL_FULL routing, AFD and VL encoder orchestration, legacy `.txt` loading, or energy/power FFI behavior without the relevant tracking issue or follow-up work.
📚 Learning: 2026-06-22T17:22:01.021Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 1235
File: rust/aiconfigurator-core/parity_tests/drift_map.py:18-25
Timestamp: 2026-06-22T17:22:01.021Z
Learning: When referencing the Nemotron-3-Nano-30B model in parity scan code (e.g., when building queries/keys into `scan.sqlite`’s `entries` support matrix), use the exact canonical model identifier string: `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` (including the `NVIDIA-` prefix). Do not omit or alter the `NVIDIA-` prefix—otherwise the lookup/mapping will fail (this was the cause of the earlier typo in `parity-scan-report.md` §6, while the parity test scripts such as `drift_map.py` should keep using the correct identifier as-is).
Applied to files:
rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py
📚 Learning: 2026-07-09T16:27:36.054Z
Learnt from: tianhaox
Repo: ai-dynamo/aiconfigurator PR: 1337
File: rust/aiconfigurator-core/src/perf_database/mla.rs:691-751
Timestamp: 2026-07-09T16:27:36.054Z
Learning: For the MLA module loaders in `perf_database` (Rust: `load_context_module_parquet`/`load_generation_module_parquet`; Python: `load_context_mla_module_data`/`load_generation_mla_module_data`), reviews should assume and preserve the intended `last-wins-across-sources` behavior: when the same key appears in multiple sources within the priority-ordered source list, values from later (higher-priority) sources overwrite earlier ones. Do not flag or “fix” these loaders to use the DSA loader’s different semantics (e.g., two-phase first-source-wins / last-row-wins-within-file). If any refactor changes the merge/overwrite semantics, require updating the in-code documentation/comments and adjusting/adding tests to cover the overwrite/duplicate-row behavior (notably the incident that motivated last-wins).
Applied to files:
rust/aiconfigurator-core/src/perf_database/mla.rs
📚 Learning: 2026-05-01T00:39:37.334Z
Learnt from: simone-chen
Repo: ai-dynamo/aiconfigurator PR: 956
File: src/aiconfigurator/sdk/perf_database.py:4144-4151
Timestamp: 2026-05-01T00:39:37.334Z
Learning: In src/aiconfigurator/sdk/**/*.py, preserve upstream metadata for PerformanceResult.source: since PerformanceResult defaults source to "silicon", callers should not force-set result.source (e.g., PerfDatabase._query_silicon_or_hybrid should rely on the default and keep any existing source information rather than overwriting it). Only set source explicitly when you truly intend to change it.
Applied to files:
src/aiconfigurator/sdk/rust_engine_step.py
🪛 ast-grep (0.44.1)
src/aiconfigurator/sdk/rust_engine_step.py
[info] 520-543: use jsonify instead of json.dumps for JSON output
Context: 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=(",", ":"),
)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 546-546: use jsonify instead of json.dumps for JSON output
Context: json.dumps(config, sort_keys=True, separators=(",", ":"))
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🪛 LanguageTool
rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md
[uncategorized] ~145-~145: The official name of this software platform is spelled with a capital “H”.
Context: ...ile-engine suite + perf gate, every PR (.github/workflows/build-test.yml:204-258), 1% ...
(GITHUB)
[style] ~162-~162: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...r the regression detector disappears. - Gate 3 (delete Python latency code): blocked...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
🪛 OpenGrep (1.25.0)
rust/aiconfigurator-core/src/operators/mla.rs
[ERROR] 684-684: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 688-688: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 765-765: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
rust/aiconfigurator-core/src/operators/moe.rs
[ERROR] 1141-1141: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
rust/aiconfigurator-core/src/operators/wideep_moe.rs
[ERROR] 341-341: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 377-377: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 381-381: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 447-447: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🔇 Additional comments (13)
rust/aiconfigurator-core/parity_tests/test_compile_engine_parity.py (1)
328-350: LGTM!Also applies to: 361-466, 473-543, 550-605
rust/aiconfigurator-core/src/common/system_spec.rs (1)
159-174: LGTM!rust/aiconfigurator-core/src/engine/spec.rs (2)
161-166: LGTM!Also applies to: 206-206, 245-245, 304-306, 328-328, 390-411, 441-461, 547-562, 651-664, 698-699
610-625: 🗄️ Data Integrity & IntegrationNo schema bump is needed here. This block only extends
all_op_variants()for the round-trip test; it doesn’t changeOp/OpSpecdeclaration order or wire discriminants.> Likely an incorrect or invalid review comment.rust/aiconfigurator-core/src/perf_database/mhc.rs (1)
132-168: LGTM!rust/aiconfigurator-core/src/perf_database/mla.rs (1)
329-458: LGTM!Also applies to: 511-541, 559-559, 585-585, 602-602, 958-975
tests/unit/sdk/test_rust_engine_step.py (1)
59-61: LGTM!Also applies to: 100-110, 135-144, 146-197, 437-478, 480-509, 511-532, 554-556, 567-570
rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md (1)
1-178: LGTM! The audit trail is thorough and consistent with the actual fixes present inmla.rs,moe.rs,wideep_moe.rs, andrust_engine_step.py. The static-analysis PII and grammar hints on this file are false positives (floating-point oracle literals and file-path code spans respectively).rust/aiconfigurator-core/src/operators/mla.rs (1)
9-24: LGTM! Call sites correctly route through the new mode-dispatch helpers with matching signatures, and the empirical helpers (own-shape grid keyed on the right quant tuple, prefix-aware SOL, BMM bfloat16 fallback bound to the requested quant's SOL) match the documented Python parity contract. Oracle tests are thorough and tolerance-tight (1e-9).Also applies to: 77-104, 125-136, 169-211, 238-245, 290-329, 356-381, 408-458, 499-543, 580-607, 658-883
rust/aiconfigurator-core/src/operators/moe.rs (1)
6-176: LGTM! EPLB token correction is correctly scoped to the silicon-only path (confirmed againstmoe_eplb_correction_scoped_to_silicon_only), the deepep/wideep routing and low-latency/xshape/xquant/xprofile transfer ladder match the documented Python parity contract, and the file-row tie-break fix forxprofile_moe_quantsis verified correct (stable sort preserving insertion order).Also applies to: 204-219, 247-249, 286-390, 392-564, 566-738, 744-805, 830-1192
rust/aiconfigurator-core/src/operators/wideep_moe.rs (1)
6-69: LGTM! Query-time kernel re-resolution, thenum_slots-aware roofline divergence from the plain-MoE SOL, and the file-row-order distribution-fallback regression test all check out against the documented Python parity contract.Also applies to: 128-266, 268-450
src/aiconfigurator/sdk/rust_engine_step.py (2)
29-36: LGTM! Database-mode gating, the mixed/decode-step kwarg forwarding matches the Rustmixed_step_latency/decode_step_latencysignatures, and the widened cache identity (database_mode, transfer_policy, raw quant names, previously-omitted ModelConfig fields) correctly addresses the audit doc's own cache-aliasing findings.Also applies to: 232-283, 315-315, 328-360, 363-388, 437-469, 472-548, 550-570, 607-622
417-436: 🗄️ Data Integrity & IntegrationDrop this stale-key concern The repo builds fresh model instances per config point, and there’s no in-repo mutation of
model.configafterget_model(), so this memoization doesn’t introduce a stale-handle path here.> Likely an incorrect or invalid review comment.
…dispatch + guards Review fixes (coderabbit on PR ai-dynamo#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 (ai-dynamo#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 <noreply@anthropic.com> Signed-off-by: Tianhao Xu <tianhaox@nvidia.com>
|
Responses to the review items that have no inline thread (outside-diff and review-body comments). Fixes are in bef4492; deferred items are tracked on #1357. Fixed (bef4492)
Rejected (stale or contrary to the Python spec)
Deferred to #1357
|
| **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). |
There was a problem hiding this comment.
If the PR aims for HYBRID/EMPIRICAL port, should we include the parity scan in this PR?
|
|
||
|
|
||
| class TestRustProvenanceCapture: | ||
| """FFI provenance contract (audit finding #6): the compiled engine records |
There was a problem hiding this comment.
Should we avoid having "audit finding" in the codebase?
| # 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`. |
There was a problem hiding this comment.
Shall we refactor the paragraph into bullet points and map the schema version to a specifc version+commit?
3 participants
Summary
Full-surface parity between the Python SDK and
rust/aiconfigurator-core, in two halves merged into one branch:rust/aiconfigurator-core/docs/parity-audit-2026-07-14.md).util_empirical) port (issue Follow-ups from the perf_interp v2 redesign (PR #1303): rust CP model, rust shared-layer, DSV4 CP calib data, minor gaps #1333 §4.6 option b): the compiled engine now implements the full util-space empirical layer —latency = SOL(query)/utilwith the xshape/xquant/xprofile/xop transfer ladder, transfer-policy gating, andEmpiricalNotImplementedhonest-miss semantics — so--database-mode HYBRID/EMPIRICALruns on the Rust step instead of delegating to Python. Only the SOL/SOL_FULL diagnostic modes still delegate.Out of scope remain the by-design Python-only surfaces (AFD, vision-encoder phase, power/energy,
.txtlegacy loading, scheduling/summary orchestration).Every fix lands with a pinning test, and each previously-uncovered trigger config is now in the parity matrices.
P0 fixes (live numeric divergences)
tp_sizeinto the per-rank-heads table axis (tp=8 queried heads=8 instead of 16 → off-grid extrapolation)128 // tp_size; mirror the{flashinfer, fa3}whitelistmla.rsbut notdsv4.rs; oracles certified stale behavior[kv][gemm]only, SOL fmha derived from kv; arch key dropped; oracles regenerateddistributionaxis Python ignores → hard error on shipped data_select_alltoall_kernelport; fp8_block→fp8; full trtllm dispatch branch (sm==100 gate — gb300 is sm103, dp>1 volumes, NVL72 reroute)_get_mix_step_latency— bit-identical on chunked/prefix/MTP shapesRuntimeContextfull_fracamortization → silent overestimatefull_fracblendmoe_perf; EPLB ×0.8 dropped; mxfp4 remaps missing; DeepEP dispatch flavors never emittedscale_num_tokenson the wireWideEpMoeDispatchop (prepare+dispatch / combine)Plus the P1 batch: GEMM fp8_static SOL floor, elementwise integer ordering, fractional NCCL message sizes, encoder partial-RoPE, wideep num_slots roofline holds, beam-width token scaling, NVL72 reroute sunk to the DB layer.
Robustness / infra
OpConversionErrorfalls back to the Python step (parity by delegation, memoized) instead of crashing the sweep.sqvsint8_wo, …) or identity-omitted ModelConfig fields; the key is memoized on the model object (the per-stepjson.dumpscost tripped the perf gate mid-branch and was fixed).cargo test --libjoins the parity CI job (196 tests): two stale pins fixed (b200 mem_bw 7.7 TB/s survived PR fix(systems): correct B200 SXM HBM bandwidth to 7.7 TB/s #1246 unnoticed), py:: binding tests self-initialize the interpreter.overlap_compositiondeleted; shared-layer resolution failures now warn loudly.New parity coverage
Validation
Follow-ups (not in this PR)
scan_rust_parity.pyrun before merge (hours; not run locally).🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Chores
HYBRID/EMPIRICAL port (merged from the hybrid-rust-parity branch)
Test-first: 16 new
HYBRID_CASESintest_engine_step_parity.py(×4 surfaces = 64 checks) were landed as a red baseline (46 failing) before the port, each tier pinned by acapture_provenance()-verified real-data config:moe_quant=w4a16_mxfp4_cutlass(same (0.5,1) profile as collected int4_wo)w4afp8(no same-profile sibling; util-level rescale)EmpiricalNotImplementedError, ladder miss (NVFP4 on Hopper)Hybrid cases assert at rtol=1e-4 (not the 1% smoke tolerance): EMPIRICAL≈SILICON by design on collected shapes — and identical at exact grid hits — so 1% cannot distinguish a mode-ignoring engine.
Mechanics:
operators/util_empirical.rs:UtilGrid(k=2 IDW in per-axis normalised log space, exact-hit preservation, boundary clamp),ZeroAwareDeltaLookup(compute_scale deltas), keyed grid caches. Unit tests mirrortests/unit/sdk/test_util_empirical.py.EmpiricalNotImplementednever caught) for GEMM/attention/MoE-ladder/MLA/DSA/DSV4/MHC/communication/alltoall/wideep-MoE, each anchored against Python-generated oracles at 1e-9 in#[cfg(test)].PerfDatabaseis now a mode-configured view overArc<PerfTables>(Deref keeps table access unchanged);silicon_view()givesOp::FallbackPython's semantics (under HYBRID the primary is evaluated silicon-only) and powers MSA's cross-op DSA probe.Op::MsaContext/MsaGeneration+ opspec conversion — MiniMax-M3 compiles and matches at 1e-4 end-to-end.ENGINE_SPEC_SCHEMA_VERSION→ 3 (variant insertion shifted bincode indices).EngineConfigcarriesdatabase_mode+transfer_policy(serde defaults keep old specs valid); the handle cache key includes both.Integration notes (merge of the two branches):
perf_database/wideep_moekeeps the hybrid branch's level-wise distribution resolution (kernel → quant → distribution, then exact shape) — the any-shape fallback was verified divergent against Python.WideEpMoeOpre-resolves the compute kernel at query time (Python_select_kernel(database, quant)needs loaded tables; the compile-time opspec value is load-order dependent and broke the TrtllmWideEp compile-parity case on a fresh database). The opspeckernel_sourcestays as a wire hint.Post-merge audit (capability + selection-rule) — findings closed
A whole-surface audit after the merge found and fixed:
PerfDataNotAvailableError/EmpiricalNotImplementedError(classifier tools work on the rust path).capture_provenance, so support-matrixHYBRID_PASSlabelling works on rust-routed runs.get_empirical_from_sol(latentTypeErroron every HYBRID silicon miss) + regression test.Deliberately deferred: energy/power over the FFI (rust-routed reports show
power=0.0W; follow-up PR — threads the energy column through the interp core, every loader, every op). Rust stays parquet-only.Verification (final tree)
rust/aiconfigurator-core/parity_tests/: 361 passed (engine-step 303 — 64 hybrid surfaces at rtol=1e-4, full SILICON suite, new typed-error + provenance-capture tests; compile-engine parity; perf gate)cargo test --lib: 276 passed, 0 failed🤖 Generated with Claude Code