feat(sweeper): preserve engine and request controls on canonical API [AIC-1771]

crates/core/src/python.rs

@@ -54,6 +54,8 @@ struct RuntimeTraffic {
     #[serde(default)]
     osl: Option<usize>,
     #[serde(default)]
+    cached_prefix_tokens: Option<usize>,
+    #[serde(default)]
     request_count: Option<usize>,
     #[serde(default)]
     turns_per_session: Option<usize>,
@@ -120,6 +122,16 @@ struct AicTimingConfig {
     free_gpu_memory_fraction: Option<f64>,
     #[serde(default)]
     systems_path: Option<String>,
+    #[serde(default)]
+    moe_backend: Option<String>,
+    #[serde(default)]
+    attention_backend: Option<String>,
+    #[serde(default)]
+    enable_wideep: bool,
+    #[serde(default)]
+    enable_eplb: bool,
+    #[serde(default)]
+    wideep_num_slots: Option<u32>,
 }
 
 const fn one() -> u32 {
@@ -231,6 +243,11 @@ impl AicTimingModel {
             kwargs.set_item("nextn", config.nextn)?;
             kwargs.set_item("kv_block_size", config.kv_block_size)?;
             kwargs.set_item("systems_path", config.systems_path.as_deref())?;
+            kwargs.set_item("moe_backend", config.moe_backend.as_deref())?;
+            kwargs.set_item("attention_backend", config.attention_backend.as_deref())?;
+            kwargs.set_item("enable_wideep", config.enable_wideep)?;
+            kwargs.set_item("enable_eplb", config.enable_eplb)?;
+            kwargs.set_item("wideep_num_slots", config.wideep_num_slots)?;
             let spec = sdk.getattr("compile_engine")?.call(
                 (
                     config.model.as_str(),
@@ -660,8 +677,16 @@ fn build_runtime_input(
     } else {
         1
     };
+    let cached_prefix_tokens = traffic.cached_prefix_tokens.unwrap_or(0);
     let trace = Trace::synthetic(SyntheticTraceSpec {
-        block_size: engine_block_size,
+        // A one-token trace block preserves prefixes that are not aligned to
+        // the engine's scheduler block size. The driver still hashes them at
+        // `engine_block_size` when it constructs replay requests.
+        block_size: if cached_prefix_tokens == 0 {
+            engine_block_size
+        } else {
+            1
+        },
         num_sessions: sessions,
         turns_per_session: turns,
         input_tokens: LengthSpec {
@@ -672,6 +697,7 @@ fn build_runtime_input(
             mean: traffic.osl.context("synthetic traffic requires osl")?,
             stddev: 0.0,
         },
+        cached_prefix_tokens,
         shared_prefix_ratio: traffic.shared_prefix_ratio.unwrap_or(0.0),
         num_prefix_groups: traffic.num_prefix_groups.unwrap_or(0),
         first_turn_arrivals: synthetic_arrivals(&traffic)?,
@@ -877,6 +903,11 @@ mod tests {
             mem_fraction_static: None,
             free_gpu_memory_fraction: None,
             systems_path: None,
+            moe_backend: None,
+            attention_backend: None,
+            enable_wideep: false,
+            enable_eplb: false,
+            wideep_num_slots: None,
         }
     }
 
@@ -926,6 +957,28 @@ mod tests {
         );
     }
 
+    #[test]
+    fn aic_timing_config_deserializes_engine_controls() {
+        let config: AicTimingConfig = serde_json::from_value(serde_json::json!({
+            "model": "test-model",
+            "backend": "sglang",
+            "system": "test-system",
+            "tp": 1,
+            "moe_backend": "deepep_moe",
+            "attention_backend": "fa3",
+            "enable_wideep": true,
+            "enable_eplb": true,
+            "wideep_num_slots": 64
+        }))
+        .unwrap();
+
+        assert_eq!(config.moe_backend.as_deref(), Some("deepep_moe"));
+        assert_eq!(config.attention_backend.as_deref(), Some("fa3"));
+        assert!(config.enable_wideep);
+        assert!(config.enable_eplb);
+        assert_eq!(config.wideep_num_slots, Some(64));
+    }
+
     #[test]
     fn json_binding_retains_authored_per_request_correlation() {
         let spec = ReplaySpec {

crates/core/src/replay/loadgen/tests.rs

@@ -718,6 +718,7 @@ fn test_partition_by_session_round_robin_keeps_sessions_intact() {
             mean: 2,
             stddev: 0.0,
         },
+        cached_prefix_tokens: 0,
         shared_prefix_ratio: 0.5,
         num_prefix_groups: 2,
         first_turn_arrivals: ArrivalSpec::Burst,
@@ -754,6 +755,7 @@ fn test_synthetic_prefix_groups_share_prefixes_within_group() {
             mean: 2,
             stddev: 0.0,
         },
+        cached_prefix_tokens: 0,
         shared_prefix_ratio: 0.5,
         num_prefix_groups: 2,
         first_turn_arrivals: ArrivalSpec::Burst,
@@ -778,6 +780,46 @@ fn test_synthetic_prefix_groups_share_prefixes_within_group() {
     assert_eq!(suffixes.len(), trace.sessions.len());
 }
 
+#[test]
+fn test_synthetic_exact_cached_prefix_shares_only_requested_tokens() {
+    let trace = Trace::synthetic(SyntheticTraceSpec {
+        block_size: 1,
+        num_sessions: 3,
+        turns_per_session: 1,
+        input_tokens: LengthSpec {
+            mean: 8,
+            stddev: 0.0,
+        },
+        output_tokens: LengthSpec {
+            mean: 2,
+            stddev: 0.0,
+        },
+        cached_prefix_tokens: 3,
+        shared_prefix_ratio: 0.0,
+        num_prefix_groups: 0,
+        first_turn_arrivals: ArrivalSpec::Burst,
+        inter_turn_delays: DelaySpec::None,
+        seed: 42,
+        arrival_seed: 42,
+    })
+    .unwrap();
+
+    let prompts = trace
+        .sessions
+        .iter()
+        .map(|session| session.turns[0].synthesize_tokens(1).unwrap())
+        .collect::<Vec<_>>();
+    assert!(prompts.windows(2).all(|pair| pair[0][..3] == pair[1][..3]));
+    assert_eq!(
+        prompts
+            .iter()
+            .map(|tokens| tokens[3..].to_vec())
+            .collect::<HashSet<_>>()
+            .len(),
+        prompts.len()
+    );
+}
+
 #[test]
 fn test_synthetic_arrival_mode_changes_timestamps_only() {
     let build = |first_turn_arrivals, arrival_seed| {
@@ -793,6 +835,7 @@ fn test_synthetic_arrival_mode_changes_timestamps_only() {
                 mean: 4,
                 stddev: 1.0,
             },
+            cached_prefix_tokens: 0,
             shared_prefix_ratio: 0.5,
             num_prefix_groups: 5,
             first_turn_arrivals,

crates/core/src/replay/loadgen/trace.rs

@@ -784,6 +784,13 @@ impl Trace {
                 spec.shared_prefix_ratio
             );
         }
+        if spec.cached_prefix_tokens % spec.block_size != 0 {
+            bail!(
+                "cached_prefix_tokens {} must align to synthetic trace block_size {}",
+                spec.cached_prefix_tokens,
+                spec.block_size
+            );
+        }
 
         let mut rng = StdRng::seed_from_u64(spec.seed);
         let mut sessions = Vec::with_capacity(spec.num_sessions);
@@ -803,13 +810,24 @@ impl Trace {
             for turn_idx in 0..spec.turns_per_session {
                 let input_length = sample_length(&spec.input_tokens, 1, &mut rng);
                 let max_output_tokens = sample_length(&spec.output_tokens, 1, &mut rng);
+                if spec.cached_prefix_tokens > input_length {
+                    bail!(
+                        "cached_prefix_tokens {} exceeds sampled synthetic input length {}",
+                        spec.cached_prefix_tokens,
+                        input_length
+                    );
+                }
                 let num_blocks = input_length.div_ceil(spec.block_size);
-                let prefix_blocks =
+                let grouped_prefix_blocks =
                     ((num_blocks as f64) * spec.shared_prefix_ratio).round() as usize;
-                let prefix_blocks = prefix_blocks.min(num_blocks);
+                let grouped_prefix_blocks = grouped_prefix_blocks.min(num_blocks);
+                let cached_prefix_blocks = spec.cached_prefix_tokens / spec.block_size;
                 let mut hash_ids = Vec::with_capacity(num_blocks);
 
-                for block_idx in 0..prefix_blocks {
+                for block_idx in 0..cached_prefix_blocks {
+                    hash_ids.push(0xCA_C0_0000_0000_0000 | block_idx as u64);
+                }
+                for block_idx in cached_prefix_blocks..grouped_prefix_blocks {
                     if let Some(group_id) = group_id {
                         hash_ids.push(0xD00D_0000_0000_0000 | (group_id << 32) | block_idx as u64);
                     }

crates/core/src/replay/loadgen/types.rs

@@ -335,6 +335,7 @@ pub struct SyntheticTraceSpec {
     pub turns_per_session: usize,
     pub input_tokens: LengthSpec,
     pub output_tokens: LengthSpec,
+    pub cached_prefix_tokens: usize,
     pub shared_prefix_ratio: f64,
     pub num_prefix_groups: usize,
     pub first_turn_arrivals: ArrivalSpec,

docs/cli/migrate-from-aiconfigurator.md

@@ -27,6 +27,114 @@ by the unified command surface. A legacy command becomes an AISimulate recommend
 | `--tpot` | `evaluation.sla.itl_ms` | Per-request goodput uses average ITL; strict mode compares aggregate mean TPOT |
 | `--strict-sla` | `optimization.strict_sla: true` | Reject before scalar ranking or Pareto dominance |
 
+## Engine and request controls
+
+The additional controls below remain flat fields under `engine`. They are concrete, pinned inputs;
+adding them does not create new recommendation dimensions or enlarge the search space.
+
+| AIConfigurator input | AISimulate recommendation | Notes |
+|---|---|---|
+| `--prefix` | `traffic.source.cached_prefix_tokens` | Exact cached tokens shared by synthetic requests; must not exceed the input length |
+| `--max-seq-len` | `engine.context_length` | Search capacity limit; AISimulate also accepts `max`, and the engine runner enforces it directly for vLLM |
+| `--nextn` | `engine.nextn` | Concrete MTP draft depth from 0 through 5 |
+| `--nextn-accepted` | `engine.nextn_accepted` | Required when `nextn > 0`; measured average accepted draft tokens in `[0, nextn]` |
+| `--enable-chunked-prefill` | `engine.enable_chunked_prefill: true` | Enables chunking on aggregated and prefill workers; decode retains its backend default |
+| `--free-gpu-memory-fraction` | `engine.free_gpu_memory_fraction` | Shared override for every active worker role |
+| `--enable-wideep` / `Task.enable_wideep` | `engine.enable_wideep: true` | Requires an MoE model; the legacy CLI flag is deprecated even though the Task field remains available |
+| `Task.enable_eplb` | `engine.enable_eplb: true` | Requires an MoE model |
+| `Task.wideep_num_slots` | `engine.wideep_num_slots` | Positive EPLB slot count |
+| `--moe-backend` / `Task.moe_backend` | `engine.moe_backend` | `deepep_moe` or `megamoe`; SGLang only |
+| `Task.attention_backend` | `engine.attention_backend` | `flashinfer` or `fa3`; SGLang MLA models only |
+| `--gemm-quant-mode` | `engine.gemm_quant_mode` | Same AIConfigurator enum name |
+| `--moe-quant-mode` | `engine.moe_quant_mode` | Same AIConfigurator enum name; requires an MoE model |
+| `--kvcache-quant-mode` | `engine.kvcache_quant_mode` | Same AIConfigurator enum name |
+| `--fmha-quant-mode` | `engine.fmha_quant_mode` | Same AIConfigurator enum name |
+| `--comm-quant-mode` | `engine.comm_quant_mode` | Same AIConfigurator enum name |
+
+`--nextn auto` does not transfer as the string `auto`. Resolve it with AIConfigurator first, then
+copy the resulting integer depth and the measured `nextn_accepted` value. AISimulate does not infer
+an acceptance assumption.
+
+### Example with the new controls
+
+Legacy AIConfigurator command:
+
+```bash
+aiconfigurator cli default \
+  --model-path deepseek-ai/DeepSeek-R1 \
+  --system h200_sxm \
+  --backend sglang \
+  --total-gpus 8 \
+  --isl 4096 \
+  --osl 1024 \
+  --prefix 1024 \
+  --max-seq-len 8192 \
+  --nextn 3 \
+  --nextn-accepted 1.5 \
+  --enable-chunked-prefill \
+  --free-gpu-memory-fraction 0.85 \
+  --moe-backend deepep_moe \
+  --gemm-quant-mode fp8 \
+  --moe-quant-mode fp8 \
+  --kvcache-quant-mode fp8 \
+  --fmha-quant-mode fp8 \
+  --comm-quant-mode fp8
+```
+
+The equivalent pinned controls in `recommendation.yaml` are:
+
+```yaml
+traffic:
+  source:
+    type: synthetic
+    input_tokens: 4096
+    output_tokens: 1024
+    cached_prefix_tokens: 1024
+  load:
+    type: constant_rate
+    requests_per_second: 4
+  stop:
+    requests_per_load_unit: 10
+
+engine:
+  mode: aggregated
+  model: deepseek-ai/DeepSeek-R1
+  hardware: h200_sxm
+  backend: sglang
+  context_length: 8192
+  nextn: 3
+  nextn_accepted: 1.5
+  enable_chunked_prefill: true
+  free_gpu_memory_fraction: 0.85
+  moe_backend: deepep_moe
+  gemm_quant_mode: fp8
+  moe_quant_mode: fp8
+  kvcache_quant_mode: fp8
+  fmha_quant_mode: fp8
+  comm_quant_mode: fp8
+  workers:
+    aggregated: {}
+
+evaluation: {}
+
+optimization:
+  target: throughput
+  constraints:
+    max_candidate_gpus: 8
+```
+
+Run it with:
+
+```bash
+aisimulate recommend --config recommendation.yaml
+```
+
+The traffic rate and stop condition are explicit AISimulate choices; the legacy capacity sweep did
+not define an equivalent offered load. For disaggregated configurations, the shared engine controls
+apply to both roles, except decode retains its backend chunking default. Keep using
+`workers.prefill.kv_cache.capacity.memory_fraction` and
+`workers.decode.kv_cache.capacity.memory_fraction` when the two roles need different memory limits.
+
 ## Illustrative strict SLA translation
 
 The examples below demonstrate how the strict-SLA fields map, but they are not behaviorally

docs/sweeper/configuration.md

@@ -61,13 +61,24 @@ configuration for each candidate.
 | `backend` | `[vllm]` | engine backends to search |
 | `gpu_budget` | `32` | maximum GPUs per candidate |
 | `min_gpu_budget` | `None` | optional lower bound during enumeration |
-| `context_length` | `None` | optional KV-feasibility sequence length |
+| `context_length` | `None` | compatibility alias for `max_seq_len`; both must match if set |
+| `max_seq_len` | model maximum | sequence capacity used by KV feasibility and every engine's `max_model_len` |
 | `parallel_configs` | `[]` | optional pinned parallel configurations |
 | `startup_time` | `None` | optional simulated worker startup time |
 | `aic_nextn` | `None` | optional speculative-decoding depth |
+| `nextn_accepted` | `None` | required explicit expected accepted draft tokens when `aic_nextn` is set |
+| `enable_chunked_prefill` | `false` | enable chunking on aggregated/prefill roles; selected `max_num_batched_tokens` remains the exact context-token budget |
+| `enable_wideep`, `enable_eplb` | `false` | shared MoE WideEP/EPLB controls |
+| `wideep_num_slots` | `None` | positive EPLB slot count |
+| `moe_backend`, `attention_backend` | `None` | explicit supported MoE/MLA kernel backends |
+| `gemm_quant_mode`, `moe_quant_mode`, `kvcache_quant_mode`, `fmha_quant_mode`, `comm_quant_mode` | `None` | shared quantization overrides used by KV feasibility and AIC timing |
+| `free_gpu_memory_fraction` | role default | shared memory fraction, mapped to backend-native total/free-memory semantics |
 
 Each engine role also has lists for `max_num_batched_tokens` and `max_num_seqs`, plus pinned block
 size, GPU-memory-utilization, and prefix-caching fields. A one-item list pins a searched field.
+Shared controls apply to both active roles in a disaggregated candidate. Unsupported model/backend
+combinations are rejected before adapter preparation or replay, and the resolved values are retained
+in `Candidate.config.engine_request` and `ReplaySpec.backend_deployment.engine_request`.
 
 ## Pinned Parallel Configurations
 

docs/sweeper/traffic.md

@@ -58,6 +58,7 @@ Every `Workload` field:
 | `num_request_ratio` | `float \| None` | `None` | Synthetic request count relative to the load: `num_requests = round(num_request_ratio * load)`. Required for synthetic. See below. |
 | `random_range_ratio` | `float` | `1.0` | Uniformly sample each synthetic ISL and OSL from `[int(ratio * configured_length), configured_length]`. Must be in `(0.0, 1.0]`; `1.0` preserves fixed lengths. Single-turn only. |
 | `random_seed` | `int` | `0` | Unsigned 64-bit seed for deterministic synthetic ISL/OSL sampling. |
+| `cached_prefix_tokens` | `int` | `0` | Exact common prompt-prefix length materialized as token IDs for KV-aware replay. It must not exceed the shortest randomized ISL and is independent of ratio-based locality controls. |
 | `shared_prefix_ratio` | `float` | `0.0` | Fraction of shared prefix across requests (cache-locality / prefix sharing). |
 | `num_prefix_groups` | `int` | `0` | Number of distinct shared-prefix groups. |
 | `turns_per_session` | `int` | `1` | Turns per multi-turn session. |
@@ -68,7 +69,7 @@ Every `Workload` field:
 | `replay_concurrency` | `int \| None` | `None` | Closed-loop in-flight cap **for a trace** (shape 1c); when set, trace timestamps are ignored. For synthetic closed-loop use `concurrency` instead. |
 
 The synthetic fields are `isl`, `osl`, `request_rate`, `concurrency`, `kv_load_ratio`,
-`num_request_ratio`, `random_range_ratio`, and `random_seed`;
+`num_request_ratio`, `random_range_ratio`, `random_seed`, and `cached_prefix_tokens`;
 `shared_prefix_ratio`, `num_prefix_groups`, `turns_per_session`, `inter_turn_delay_ms` are
 shared synthetic knobs carried by `ReplaySpec.workload`.