fix: use decode-phase Mamba2 rows and gate the Mamba2 operator by database mode

crates/core/parity_tests/perfmodel/goldens/engine_step.json

@@ -1035,8 +1035,8 @@
       },
       "disagg": {
         "values": {
-          "disagg_request": 79.212,
-          "disagg_tpot": 4.901,
+          "disagg_request": 79.456,
+          "disagg_tpot": 5.145,
           "disagg_ttft": 74.311
         }
       },
@@ -1048,8 +1048,8 @@
       "static": {
         "values": {
           "static_ctx": 37.53,
-          "static_gen": 4.608,
-          "static_total": 42.138
+          "static_gen": 4.806,
+          "static_total": 42.336
         }
       }
     },
@@ -2483,15 +2483,15 @@
     },
     "nemotron-h-56b-b200-trtllm-isl1024-osl2::disagg": {
       "engine": "rust",
-      "git_head": "ed91582f6b9dff7c5cb7b5842fd1de5236ba4726"
+      "git_head": "19067832cc1478b1edf26fed3ac7cae3ac7c56cc"
     },
     "nemotron-h-56b-b200-trtllm-isl1024-osl2::mixed": {
       "engine": "rust",
       "git_head": "ed91582f6b9dff7c5cb7b5842fd1de5236ba4726"
     },
     "nemotron-h-56b-b200-trtllm-isl1024-osl2::static": {
       "engine": "rust",
-      "git_head": "ed91582f6b9dff7c5cb7b5842fd1de5236ba4726"
+      "git_head": "19067832cc1478b1edf26fed3ac7cae3ac7c56cc"
     },
     "nemotron-h-56b-b200-vllm-empirical::agg": {
       "engine": "rust",

crates/core/src/perfmodel/operators/mamba.rs

@@ -6,6 +6,7 @@
 //!
 //! Each wraps `db.state_space.query_*` with `scale_factor` + `clamp`.
 
+use crate::common::enums::DatabaseMode;
 use crate::common::error::AicError;
 use crate::operators::base::{PerformanceResult, Source};
 use crate::perf_database::PerfDatabase;
@@ -36,6 +37,19 @@ impl Mamba2Op {
         batch_size: u32,
         seq_len: u32,
     ) -> Result<PerformanceResult, AicError> {
+        // SOL (and the retired SOL_FULL alias) never reads the table, as in
+        // the GEMM, attention, MoE and mHC operators. Mamba2 has no util-space
+        // empirical layer, so EMPIRICAL answers with the same analytic SOL
+        // instead of reading measured rows.
+        if matches!(
+            db.database_mode,
+            DatabaseMode::Sol | DatabaseMode::SolFull | DatabaseMode::Empirical
+        ) {
+            let latency = self.sol_latency_ms(db, batch_size as f64, seq_len as f64);
+            return Ok(PerformanceResult::new(latency, Source::Sol)
+                .clamp_non_negative()
+                .scaled(self.scale_factor));
+        }
         // Mirrors Python `Mamba2Kernel.query`: silicon-first, SOL fallback
         // on perf-DB miss. The op's arg-style SOL is threaded into the table
         // query so the perf_interp engine can util-hold beyond-range queries
@@ -748,4 +762,156 @@ mod tests {
         let op = kda_op("mystery_kernel", "generation", 0);
         assert_eq!(op.sol_total_bytes(4.0, 0.0), 4.0 * 4096.0 * 2.0 * 2.0);
     }
+
+    /// A decode op reads its measured row instead of the analytic estimate.
+    /// Synthetic fixture: one `causal_conv1d_update` row at `batch=2`; the
+    /// query hits it exactly, so the expected latency is the stored value
+    /// times `scale_factor`. A decode kernel without rows keeps the analytic
+    /// estimate.
+    #[test]
+    fn mamba2_decode_op_reads_its_measured_row() {
+        use crate::perf_database::energy_test_fixtures::{
+            Col, write_energy_systems_root, write_parquet,
+        };
+        let root = tempfile::tempdir().expect("tmpdir");
+        write_energy_systems_root(root.path());
+        let data = root.path().join("data/test/1.0");
+        std::fs::create_dir_all(&data).expect("mkdir data");
+        write_parquet(
+            &data.join("mamba2_perf.parquet"),
+            &[
+                Col::Str("kernel_source", vec!["causal_conv1d_update"]),
+                Col::Str("phase", vec!["generation"]),
+                Col::I64("batch_size", vec![2]),
+                Col::I64("seq_len", vec![1]),
+                Col::I64("d_model", vec![1024]),
+                Col::I64("d_state", vec![64]),
+                Col::I64("d_conv", vec![4]),
+                Col::I64("nheads", vec![16]),
+                Col::I64("head_dim", vec![64]),
+                Col::I64("n_groups", vec![4]),
+                Col::I64("chunk_size", vec![128]),
+                Col::F64("latency", vec![0.25]),
+            ],
+        );
+        let db = PerfDatabase::load(root.path(), "testsys", "test", "1.0").expect("load");
+        let op = |kernel_source: &str| Mamba2Op {
+            name: "generation_mamba".into(),
+            scale_factor: 2.0,
+            kernel_source: kernel_source.into(),
+            phase: "generation".into(),
+            d_model: 1024,
+            d_state: 64,
+            d_conv: 4,
+            nheads: 16,
+            head_dim: 64,
+            n_groups: 4,
+            chunk_size: 128,
+        };
+        let measured = op("causal_conv1d_update").query(&db, 2, 1024).unwrap();
+        assert_eq!(measured.source, Source::Silicon);
+        assert_eq!(measured.latency_ms, 0.5);
+        let analytic = op("selective_state_update").query(&db, 2, 1024).unwrap();
+        assert_eq!(analytic.source, Source::Sol);
+    }
+
+    /// Mamba2 follows the database-mode semantics of the other table-backed
+    /// operators: SILICON and HYBRID read measured rows, while SOL, the
+    /// retired SOL_FULL alias and EMPIRICAL never read them. Mamba2 has no
+    /// util-space empirical layer, so EMPIRICAL answers with the analytic SOL.
+    ///
+    /// Synthetic fixture: one context row at `[batch=2][seq=512]` and one
+    /// decode row at `[batch=2]`; both queries hit them exactly, so a row
+    /// read returns the stored value times `scale_factor` (2.0). Hand-derived
+    /// SOL for this shape: conv_dim = 16*64 + 2*4*64 = 1536 channels and
+    /// bytes = x*conv_dim*(d_conv+1)*2 + x*conv_dim*2, with x = b*s = 1024 for
+    /// context (18_874_368 bytes) and x = b = 2 for decode (36_864 bytes),
+    /// over the fixture's 7.7e12 B/s memory bandwidth.
+    #[test]
+    fn mamba2_reads_measured_rows_only_in_silicon_and_hybrid() {
+        use crate::common::enums::{DatabaseMode, TransferPolicy};
+        use crate::perf_database::energy_test_fixtures::{
+            Col, write_energy_systems_root, write_parquet,
+        };
+        let root = tempfile::tempdir().expect("tmpdir");
+        write_energy_systems_root(root.path());
+        let data = root.path().join("data/test/1.0");
+        std::fs::create_dir_all(&data).expect("mkdir data");
+        write_parquet(
+            &data.join("mamba2_perf.parquet"),
+            &[
+                Col::Str(
+                    "kernel_source",
+                    vec!["causal_conv1d_fn", "causal_conv1d_update"],
+                ),
+                Col::Str("phase", vec!["context", "generation"]),
+                Col::I64("batch_size", vec![2, 2]),
+                Col::I64("seq_len", vec![512, 1]),
+                Col::I64("d_model", vec![1024, 1024]),
+                Col::I64("d_state", vec![64, 64]),
+                Col::I64("d_conv", vec![4, 4]),
+                Col::I64("nheads", vec![16, 16]),
+                Col::I64("head_dim", vec![64, 64]),
+                Col::I64("n_groups", vec![4, 4]),
+                Col::I64("chunk_size", vec![128, 128]),
+                Col::F64("latency", vec![0.5, 0.25]),
+            ],
+        );
+        let op = |kernel_source: &str, phase: &str| Mamba2Op {
+            name: format!("{phase}_mamba"),
+            scale_factor: 2.0,
+            kernel_source: kernel_source.into(),
+            phase: phase.into(),
+            d_model: 1024,
+            d_state: 64,
+            d_conv: 4,
+            nheads: 16,
+            head_dim: 64,
+            n_groups: 4,
+            chunk_size: 128,
+        };
+        let context = op("causal_conv1d_fn", "context");
+        let decode = op("causal_conv1d_update", "generation");
+        let sol_ms = |bytes: f64| bytes / 7.7e12 * 1000.0 * 2.0;
+        for mode in [
+            DatabaseMode::Silicon,
+            DatabaseMode::Hybrid,
+            DatabaseMode::Empirical,
+            DatabaseMode::Sol,
+            DatabaseMode::SolFull,
+        ] {
+            let db = PerfDatabase::load(root.path(), "testsys", "test", "1.0")
+                .expect("load")
+                .with_mode(mode, TransferPolicy::ALL);
+            let ctx = context.query(&db, 2, 512).unwrap();
+            let dec = decode.query(&db, 2, 1024).unwrap();
+            if matches!(mode, DatabaseMode::Silicon | DatabaseMode::Hybrid) {
+                assert_eq!(
+                    (ctx.source, ctx.latency_ms),
+                    (Source::Silicon, 1.0),
+                    "{mode:?}"
+                );
+                assert_eq!(
+                    (dec.source, dec.latency_ms),
+                    (Source::Silicon, 0.5),
+                    "{mode:?}"
+                );
+            } else {
+                assert_eq!(
+                    (ctx.source, dec.source),
+                    (Source::Sol, Source::Sol),
+                    "{mode:?}"
+                );
+                let (ctx_sol, dec_sol) = (sol_ms(18_874_368.0), sol_ms(36_864.0));
+                assert!(
+                    (ctx.latency_ms - ctx_sol).abs() <= 1e-12 * ctx_sol,
+                    "{mode:?}"
+                );
+                assert!(
+                    (dec.latency_ms - dec_sol).abs() <= 1e-12 * dec_sol,
+                    "{mode:?}"
+                );
+            }
+        }
+    }
 }

crates/core/src/perfmodel/perf_database/state_space.rs

@@ -282,30 +282,6 @@ impl StateSpaceTable {
         chunk_size: u32,
         sol: &dyn Fn(f64, f64) -> f64,
     ) -> Result<LeafValue, AicError> {
-        // Mirror Python v2's `load_mamba2_data` defaultdict bug (still
-        // present today, verified 2026-07-09): the row-population pattern
-        // `try { data[ks][ph][mk][bs] } except KeyError: ... = entry` never
-        // reaches the `except` branch for generation rows because the
-        // fourth level is a `defaultdict(dict)` that lazily materialises an
-        // empty dict on `[bs]` access (no KeyError). Generation leaves end
-        // up as empty `{}`; `_query_mamba2_table`'s generation branch then
-        // normalizes them to an empty curve, the perf_interp engine raises,
-        // and the op falls back to SOL (`db.query_mamba2(generation)`
-        // returns `source="sol"` for every query). The Rust parquet loader
-        // populates the rows correctly, so returning silicon here would
-        // give a numerically different (and arguably "more correct")
-        // answer — but for apple-to-apple parity we mirror Python by
-        // returning a PerfDatabase error so the operator-layer SOL branch
-        // fires. GDN's loader is fine (uses explicit `in` checks), so this
-        // workaround is Mamba2-generation-only.
-        if phase == "generation" {
-            return Err(AicError::PerfDatabase(format!(
-                "Mamba2 generation data intentionally not used (matches Python v2 \
-                 `load_mamba2_data` defaultdict bug in operations/mamba.py — generation \
-                 leaves load empty, so every generation query degrades to SOL); \
-                 operator must fall to SOL. ks={kernel_source}, d_model={d_model}"
-            )));
-        }
         let grids = self.load_mamba2()?;
         let key = Mamba2Key {
             kernel_source: kernel_source.to_string(),
@@ -325,9 +301,6 @@ impl StateSpaceTable {
         // shares d_model, surface as `PerfDatabase` so the operator layer's
         // SOL fallback applies.
         //
-        // Only the context phase reaches this point — generation queries are
-        // short-circuited above to match Python's degenerate behaviour.
-        //
         // The backend's own persisted physical label (TRT-LLM rc29 renamed
         // the context convolution `causal_conv1d_fn_token_major`) takes
         // precedence over the logical lane, which after the shared-layer
@@ -358,8 +331,8 @@ impl StateSpaceTable {
     }
 
     /// GDN latency for a layer instance. Same engine resolution as Mamba2
-    /// (and, unlike Mamba2, generation queries really resolve: 1-axis
-    /// `(batch,)` Grid RAW over the generation curve, per Python v2).
+    /// (generation queries resolve on a 1-axis `(batch,)` Grid RAW over the
+    /// generation curve, per Python v2).
     #[allow(clippy::too_many_arguments)]
     pub fn query_gdn(
         &self,
@@ -733,9 +706,7 @@ fn load_mamba2_parquet(sources: &[PerfSource]) -> Result<Mamba2Grids, AicError>
             };
             // First-wins parity with Python `load_mamba2_data`, extended across
             // shared-layer sources (earlier source wins). Generation rows are
-            // keyed by batch only (Python drops seq for generation). Note the
-            // stored generation leaves are never read: `query_mamba2` mirrors
-            // Python's empty-generation-leaves bug by erroring first.
+            // keyed by batch only (Python drops seq for generation).
             let node = by_keys.entry(key).or_insert_with(Node::branch);
             let batch = row.u32(batch_size_col)?;
             let latency = row.f64(latency_col)?;
@@ -2113,4 +2084,60 @@ mod tests {
             in_memory_mamba2_table("trtllm", "1.3.0rc29", &[("mamba_chunk_scan_combined", 3.0)]);
         assert!(query_mamba2_conv(&table).is_err());
     }
+
+    /// One Mamba2 shape written through the parquet loader: a context row at
+    /// `[batch=2][seq=512]` and a decode row at `[batch=2]`. Each query hits
+    /// its stored coordinate, so Grid RAW returns the stored fixture value
+    /// itself; `dummy_sol` returns 1.0, which would mark an analytic estimate.
+    #[test]
+    fn mamba2_decode_rows_resolve_like_context_rows() {
+        use crate::perf_database::energy_test_fixtures::{Col, write_parquet};
+        let tmp = tempfile::tempdir().expect("tmpdir");
+        write_parquet(
+            &tmp.path().join("mamba2_perf.parquet"),
+            &[
+                Col::Str(
+                    "kernel_source",
+                    vec!["causal_conv1d_fn", "causal_conv1d_update"],
+                ),
+                Col::Str("phase", vec!["context", "generation"]),
+                Col::I64("batch_size", vec![2, 2]),
+                Col::I64("seq_len", vec![512, 1]),
+                Col::I64("d_model", vec![1024, 1024]),
+                Col::I64("d_state", vec![64, 64]),
+                Col::I64("d_conv", vec![4, 4]),
+                Col::I64("nheads", vec![16, 16]),
+                Col::I64("head_dim", vec![64, 64]),
+                Col::I64("n_groups", vec![4, 4]),
+                Col::I64("chunk_size", vec![128, 128]),
+                Col::F64("latency", vec![0.5, 0.25]),
+            ],
+        );
+        let table = StateSpaceTable::new(tmp.path().to_path_buf(), "test", "1.0");
+        let query = |kernel_source: &str, phase: &str, seq_len: u32| {
+            table
+                .query_mamba2(
+                    kernel_source,
+                    phase,
+                    2,
+                    seq_len,
+                    1024,
+                    64,
+                    4,
+                    16,
+                    64,
+                    4,
+                    128,
+                    &dummy_sol,
+                )
+                .map(|v| v.latency)
+        };
+        // The context path keeps resolving on its 2-axis grid.
+        assert_eq!(query("causal_conv1d_fn", "context", 512).unwrap(), 0.5);
+        // Decode resolves on the 1-axis batch curve; its seq feeds SOL only.
+        assert_eq!(
+            query("causal_conv1d_update", "generation", 1024).unwrap(),
+            0.25
+        );
+    }
 }