feat: decouple FPM collection limits and validate AgentX coverage

AGENTS.md

@@ -8,16 +8,17 @@ When asked to onboard a model for FPM simulation on designated hardware, follow
 [Onboard with an agent](docs/fpm-self-service.md#onboard-with-an-agent) in the FPM
 self-service guide. Use the checkout's `aisimulate onboard` CLI and current help.
 Follow its six stages: inspect the model and target; choose the worker and
-workload; derive, review and save the profile; plan collection; collect and verify
-data; run prediction/recommendation and report. Start by asking only for a missing
+review runtime/collection limits; derive, review and save the profile; plan
+collection; collect and verify data; validate replay and run ordinary
+prediction/recommendation. Start by asking only for a missing
 Hugging Face model ID (`organization/model-name`) and target GPU platform. Accept
 an already supplied local config, profile or checkpoint path instead of requiring
 a Hub ID; do not ask for both an ID and a config upfront. For a Hub ID, retrieve
 the config using available Hub access as described in stage 1. Reuse supplied
 facts and inspect the config/profile before asking for derivable metadata; defer
 other questions to their stage and help the user choose the worker topology.
 At stage 2, inspect `aisimulate onboard init --model-config PATH --suggest-parallel`
-with the actual checkpoint/revision, runtime, GPU, interconnect and workload
+with the actual checkpoint/revision, runtime, GPU, interconnect and collection
 options. This read-only JSON preview reports model/hardware-aware TP choices and
 MoE TP/DEP/TEP alternatives, exact flags, resource sources and missing inputs.
 Ask for unresolved shared precision/layout facts rather than guessing. Only a
@@ -34,6 +35,29 @@ interconnect characteristics; verify actual collection resources before executio
 Generated predict/recommend configs validate one worker. Deployment replicas and
 optimization GPU budgets belong to ordinary predict/recommend configurations.
 
+Review context, scheduler and prefill capture limits independently of validation
+traffic. Fresh config/profile defaults use the smaller of the declared context
+and 256,000 tokens, profile scheduler bounds or 8,192 tokens/256 sequences, and a
+2,048-token prefill CUDA graph capture limit. Explain and allow edits to these
+initial policies; they establish neither capacity nor timing coverage. Do not
+require fixed input/output lengths, concurrency, TTFT or TPOT during intake.
+Those flags customize optional synthetic examples only. Maximum sequences does
+not reserve maximum context for every sequence. Dynamo self-benchmark owns the
+grid generated from capture sizes and runtime bounds; do not invent a separate
+AgentX collection grid.
+
+After verifying the formal data pair, use `aisimulate onboard validate-fpm`
+with a local Weka trace and a separate validation output directory. Follow the
+guide's pinned AgentX reference and current cold aggregated, one-lane, HBM-only,
+non-speculative scope. Preserve target model projection and strict direct FPM
+with denied fallback. Read `validation.json` together with native query coverage
+and request completion evidence. Coverage counts native lookup resolutions,
+including interpolation and failures; cached timing reuse is not another query.
+Missing timing stops replay and preserves partial evidence, which cannot certify
+the remainder of the trace. Report coverage and silicon accuracy separately.
+Reuse verified v2 collection plans for validation-only changes; legacy v1 plans
+need a new directory as described in the guide.
+
 Do not reject a checkpoint solely because it is multimodal. Read its unambiguous
 `text_config`, or its flat text-decoder fields, and explain that FPM models only
 the text decoder. Multimodal encoders, projectors, preprocessing and other

crates/core/src/lib.rs

@@ -49,7 +49,8 @@ pub use perfmodel::{
 };
 pub use perfmodel::{
     CorrectionFactorBounds, CorrectionFeatureSpace, FpmInterpolationConfig, FpmInterpolationMethod,
-    OpLevelConfig, RegressionFitConfig, RegressionFitKind,
+    FpmQueryCoverage, FpmQueryCoverageCounts, FpmQueryGap, FpmQueryPurpose, OpLevelConfig,
+    RegressionFitConfig, RegressionFitKind,
 };
 
 #[cfg(feature = "python")]

crates/core/src/perfmodel/engine/runtime.rs

@@ -20,8 +20,10 @@
 
 use std::sync::Arc;
 
+use crate::FpmQueryCoverage;
 use crate::common::enums::{DatabaseMode, TransferPolicy};
 use crate::common::error::AicError;
+use crate::operators::RuntimeContext;
 use crate::operators::base::PerformanceResult;
 use crate::operators::{FpmForwardOp, FpmPhase, Op};
 use crate::perf_database::PerfDatabase;
@@ -585,6 +587,14 @@ impl Engine {
     /// Python `_run_context_phase` (`base_backend.py:144`): `effective_isl =
     /// isl - prefix`, validate `> 0`, then one full pass over `context_ops`.
     fn run_context_phase(&self, runtime: &RuntimeConfig) -> Result<f64, AicError> {
+        self.run_context_phase_with_coverage(runtime, None)
+    }
+
+    fn run_context_phase_with_coverage(
+        &self,
+        runtime: &RuntimeConfig,
+        coverage: Option<&mut FpmQueryCoverage>,
+    ) -> Result<f64, AicError> {
         // Python raises `ValueError` when `effective_isl <= 0`; mirror that.
         if runtime.prefix >= runtime.isl {
             return Err(AicError::InvalidEngineConfig(format!(
@@ -593,6 +603,33 @@ impl Engine {
             )));
         }
         let effective_isl = runtime.isl - runtime.prefix;
+        if let Some(coverage) = coverage {
+            let Some((prefill, _, ctx_tail, gen_tail)) = self.fpm_split() else {
+                return Err(AicError::InvalidEngineConfig(
+                    "query coverage requires a direct FPM engine".into(),
+                ));
+            };
+            if !ctx_tail.is_empty() || !gen_tail.is_empty() || self.nextn > 0 {
+                return Err(AicError::InvalidEngineConfig(
+                    "query coverage does not support speculative FPM".into(),
+                ));
+            }
+            return prefill
+                .query_with_coverage(
+                    &self.db,
+                    &RuntimeContext {
+                        batch_size: runtime.batch_size,
+                        beam_width: 1,
+                        s: effective_isl,
+                        prefix: runtime.prefix,
+                        num_tokens: runtime.batch_size * effective_isl,
+                        seq_imbalance_correction_scale: runtime.seq_imbalance_correction_scale,
+                        ..Default::default()
+                    },
+                    Some(coverage),
+                )
+                .map(|result| result.latency_ms);
+        }
         run_context_ops(
             &self.context_ops,
             &self.db,
@@ -688,16 +725,24 @@ impl Engine {
     /// Thin shim over [`Self::run_static`] with `mode=Context` (osl is
     /// irrelevant for the context phase, so it is fixed at 1).
     pub fn predict_prefill_latency(&self, bs: u32, isl: u32, prefix: u32) -> Result<f64, AicError> {
+        self.predict_prefill_latency_with_coverage(bs, isl, prefix, None)
+    }
+
+    pub(crate) fn predict_prefill_latency_with_coverage(
+        &self,
+        bs: u32,
+        isl: u32,
+        prefix: u32,
+        coverage: Option<&mut FpmQueryCoverage>,
+    ) -> Result<f64, AicError> {
         let rt = RuntimeConfig {
             batch_size: bs,
             isl,
             osl: 1,
             prefix,
             ..Default::default()
         };
-        Ok(self
-            .run_static(&rt, StaticMode::Context, DEFAULT_STATIC_STRIDE)?
-            .total_ms)
+        self.run_context_phase_with_coverage(&rt, coverage)
     }
 
     /// Mocker H2: decode-step latency in ms. Pure-Rust inherent method (no
@@ -725,14 +770,26 @@ impl Engine {
         batch_size: u32,
         total_past_kv_tokens: u32,
     ) -> Result<f64, AicError> {
-        self.forward_pass_time_ms(&[ForwardPassMetrics {
-            scheduled_requests: crate::ScheduledRequestMetrics {
-                num_decode_requests: batch_size,
-                sum_decode_kv_tokens: total_past_kv_tokens,
+        self.predict_decode_latency_total_with_coverage(batch_size, total_past_kv_tokens, None)
+    }
+
+    pub(crate) fn predict_decode_latency_total_with_coverage(
+        &self,
+        batch_size: u32,
+        total_past_kv_tokens: u32,
+        coverage: Option<&mut FpmQueryCoverage>,
+    ) -> Result<f64, AicError> {
+        self.forward_pass_time_ms_with_coverage(
+            &[ForwardPassMetrics {
+                scheduled_requests: crate::ScheduledRequestMetrics {
+                    num_decode_requests: batch_size,
+                    sum_decode_kv_tokens: total_past_kv_tokens,
+                    ..Default::default()
+                },
                 ..Default::default()
-            },
-            ..Default::default()
-        }])
+            }],
+            coverage,
+        )
     }
 
     /// Highest decode KV-read total covered by a compiled FPM engine.
@@ -1971,6 +2028,14 @@ impl Engine {
     pub fn forward_pass_time_ms(
         &self,
         metrics_by_rank: &[ForwardPassMetrics],
+    ) -> Result<f64, AicError> {
+        self.forward_pass_time_ms_with_coverage(metrics_by_rank, None)
+    }
+
+    pub(crate) fn forward_pass_time_ms_with_coverage(
+        &self,
+        metrics_by_rank: &[ForwardPassMetrics],
+        mut coverage: Option<&mut FpmQueryCoverage>,
     ) -> Result<f64, AicError> {
         if metrics_by_rank.is_empty() {
             return Err(AicError::InvalidForwardPassMetrics(
@@ -1982,7 +2047,7 @@ impl Engine {
         }
         let mut max_latency = 0.0_f64;
         for metrics in metrics_by_rank {
-            let rank_latency = self.rank_latency_ms(metrics)?;
+            let rank_latency = self.rank_latency_ms(metrics, coverage.as_deref_mut())?;
             if rank_latency > max_latency {
                 max_latency = rank_latency;
             }
@@ -1996,7 +2061,11 @@ impl Engine {
     /// [`run_context_ops`]; decode-only -> [`run_generation_ops_step`]. The FPM
     /// counts pass through unscaled (no `nextn` multiplier — see
     /// [`Self::forward_pass_time_ms`]).
-    fn rank_latency_ms(&self, metrics: &ForwardPassMetrics) -> Result<f64, AicError> {
+    fn rank_latency_ms(
+        &self,
+        metrics: &ForwardPassMetrics,
+        mut coverage: Option<&mut FpmQueryCoverage>,
+    ) -> Result<f64, AicError> {
         let sched = &metrics.scheduled_requests;
         // Token-based dispatch, aligned with `IterationFeatures` (fpm/model.rs):
         // a fully prefix-cached payload can retain prefill request/KV metadata
@@ -2033,35 +2102,38 @@ impl Engine {
             let mut total = 0.0_f64;
             if has_prefill {
                 total += prefill_op
-                    .query_totals(
+                    .query_totals_with_coverage(
                         &self.db,
                         &[
                             sched.num_prefill_requests as f64,
                             sched.sum_prefill_tokens as f64,
                             sched.sum_prefill_kv_tokens as f64,
                         ],
+                        coverage.as_deref_mut(),
                     )?
                     .latency_ms;
             }
             if has_decode {
                 let decode_ms = decode_op
-                    .query_totals(
+                    .query_totals_with_coverage(
                         &self.db,
                         &[
                             sched.num_decode_requests as f64,
                             sched.sum_decode_kv_tokens as f64,
                         ],
+                        coverage.as_deref_mut(),
                     )?
                     .latency_ms;
                 if has_prefill {
                     // Mixed rank: marginal-decode composition, mirroring
                     // `_get_fpm_mix_step_latency` (counts already packed, no
                     // `(nextn + 1)` — speculative FPM was rejected above).
                     let baseline_ms = decode_op
-                        .query_pass_baseline(
+                        .query_pass_baseline_with_coverage(
                             &self.db,
                             sched.num_decode_requests,
                             sched.sum_decode_kv_tokens as f64,
+                            coverage.as_deref_mut(),
                         )?
                         .latency_ms;
                     total += (decode_ms - baseline_ms).max(0.0);

crates/core/src/perfmodel/fpm/config.rs

@@ -274,6 +274,16 @@ impl ForwardPassPerfModelConfig {
     }
 
     pub(crate) fn validate(&self) -> Result<(), AicError> {
+        if self.estimator_config.fpm_interpolation.collect_coverage
+            && (self.estimation_mode != EstimationMode::FpmInterpolation
+                || self.fallback_policy != ForwardPassFallbackPolicy::Deny
+                || self.estimator_config.fpm_interpolation.method
+                    != super::FpmInterpolationMethod::Direct)
+        {
+            return Err(invalid_config(
+                "collect_coverage requires estimation_mode='fpm_interpolation', method='direct', and fallback_policy='deny'",
+            ));
+        }
         if self.model.trim().is_empty() {
             return Err(invalid_config("model cannot be empty"));
         }
@@ -436,6 +446,66 @@ mod tests {
         );
     }
 
+    #[test]
+    fn coverage_controls_preserve_default_serialization_and_require_strict_direct() {
+        let default = serde_json::to_value(EstimatorConfig::default()).unwrap();
+        assert_eq!(
+            default["fpm_interpolation"],
+            serde_json::json!({"method": "auto"})
+        );
+        let cfg = config(serde_json::json!({
+            "estimation_mode": "fpm_interpolation",
+            "estimator_config": {"fpm_interpolation": {"method": "direct", "collect_coverage": true}}
+        }));
+        cfg.validate().unwrap();
+        let reloaded: ForwardPassPerfModelConfig =
+            serde_json::from_value(serde_json::to_value(&cfg).unwrap()).unwrap();
+        assert_eq!(cfg, reloaded);
+        for mode in [
+            EstimationMode::Auto,
+            EstimationMode::OpLevel,
+            EstimationMode::FpmRegression,
+        ] {
+            let mut invalid = cfg.clone();
+            invalid.estimation_mode = mode;
+            assert!(
+                invalid
+                    .validate()
+                    .unwrap_err()
+                    .to_string()
+                    .contains("collect_coverage")
+            );
+        }
+        for method in [
+            super::super::FpmInterpolationMethod::Auto,
+            super::super::FpmInterpolationMethod::Sol,
+        ] {
+            let mut invalid = cfg.clone();
+            invalid.estimator_config.fpm_interpolation.method = method;
+            assert!(
+                invalid
+                    .validate()
+                    .unwrap_err()
+                    .to_string()
+                    .contains("collect_coverage")
+            );
+        }
+        for fallback in [
+            ForwardPassFallbackPolicy::Allow,
+            ForwardPassFallbackPolicy::LegacyRegression,
+        ] {
+            let mut invalid = cfg.clone();
+            invalid.fallback_policy = fallback;
+            assert!(
+                invalid
+                    .validate()
+                    .unwrap_err()
+                    .to_string()
+                    .contains("collect_coverage")
+            );
+        }
+    }
+
     #[test]
     fn engine_controls_round_trip_with_identity_and_policy() {
         let cfg = config(serde_json::json!({