feat(mocker): add explicit max_model_len support to vLLM mocker (#11069)

Signed-off-by: Yongming Ding <yongmingd@nvidia.com>

Author: dreamtalen

components/src/dynamo/mocker/args.py

@@ -18,6 +18,16 @@
 logger = logging.getLogger(__name__)
 
 
+def positive_int(value: str) -> int:
+    try:
+        parsed = int(value)
+    except ValueError as error:
+        raise argparse.ArgumentTypeError(str(error)) from error
+    if parsed <= 0:
+        raise argparse.ArgumentTypeError(f"must be positive, got {parsed}")
+    return parsed
+
+
 def non_negative_int(value: str) -> int:
     try:
         parsed = int(value)
@@ -211,6 +221,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
         help="Token block size for KV cache blocks. When unset, the default "
         "depends on engine: vLLM 64, SGLang 1, TRTLLM 32.",
     )
+    parser.add_argument(
+        "--max-model-len",
+        type=positive_int,
+        default=None,
+        help="Maximum vLLM sequence length, including prompt and generated tokens. "
+        "When omitted, no model-length limit is enforced.",
+    )
     parser.add_argument(
         "--max-num-seqs",
         type=int,

components/src/dynamo/mocker/config.py

@@ -243,6 +243,7 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs:
         aic_moe_ep_size = getattr(args, "aic_moe_ep_size", None)
         aic_attention_dp_size = getattr(args, "aic_attention_dp_size", None)
     engine_type = getattr(args, "engine_type", None) or "vllm"
+    max_model_len = getattr(args, "max_model_len", None)
     num_gpu_blocks = _resolve_num_gpu_blocks(
         explicit_num_gpu_blocks=getattr(args, "num_gpu_blocks", None),
         engine_type=engine_type,
@@ -267,6 +268,7 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs:
         engine_type=engine_type,
         num_gpu_blocks=num_gpu_blocks,
         block_size=getattr(args, "block_size", 0) or 0,
+        max_model_len=max_model_len,
         max_num_seqs=getattr(args, "max_num_seqs", _DEFAULT_MAX_NUM_SEQS),
         max_num_batched_tokens=getattr(
             args, "max_num_batched_tokens", _DEFAULT_MAX_NUM_BATCHED_TOKENS
@@ -349,8 +351,7 @@ def build_runtime_config(
     engine_args: MockEngineArgs,
 ) -> tuple[int, ModelRuntimeConfig]:
     rc = ModelRuntimeConfig()
-    # Mocker does not enforce a model context limit.
-    rc.context_length = 0
+    rc.context_length = engine_args.max_model_len or 0
     rc.total_kv_blocks = engine_args.num_gpu_blocks
     rc.max_num_seqs = engine_args.max_num_seqs
     if rc.max_num_seqs is None:

components/src/dynamo/mocker/tests/unit/test_config.py

@@ -35,6 +35,7 @@ def make_args(**overrides):
         "engine_type": "vllm",
         "num_gpu_blocks": None,
         "block_size": None,
+        "max_model_len": None,
         "max_num_seqs": 256,
         "max_num_batched_tokens": 8192,
         "enable_prefix_caching": True,
@@ -397,6 +398,73 @@ def test_mocker_cli_accepts_mtp_configuration():
     assert args.aic_mtp_seed == 99
 
 
+def test_mocker_cli_accepts_max_model_len():
+    args = parse_args(["--max-model-len", "32768"])
+
+    engine_args = CONFIG.build_mocker_engine_args(args)
+    _, runtime_config = CONFIG.build_runtime_config(engine_args)
+
+    assert engine_args.max_model_len == 32768
+    assert runtime_config.context_length == 32768
+
+
+@pytest.mark.parametrize("value", ["0", "-1"])
+def test_mocker_cli_rejects_non_positive_max_model_len(value):
+    with pytest.raises(SystemExit):
+        parse_args(["--max-model-len", value])
+
+
+def test_build_mocker_engine_args_keeps_max_model_len_explicit_only():
+    engine_args = CONFIG.build_mocker_engine_args(
+        make_args(model_path="/models/mock", num_gpu_blocks=4096)
+    )
+
+    assert engine_args.max_model_len is None
+
+
+def test_build_mocker_engine_args_preserves_explicit_max_model_len():
+    engine_args = CONFIG.build_mocker_engine_args(
+        make_args(
+            model_path="/models/mock",
+            max_model_len=32768,
+            num_gpu_blocks=4096,
+        )
+    )
+
+    assert engine_args.max_model_len == 32768
+
+
+def test_replay_engine_args_keeps_max_model_len_explicit_only():
+    import dynamo.replay.main as replay_main
+
+    engine_args = replay_main._load_engine_args(
+        json.dumps(
+            {
+                "num_gpu_blocks": 4096,
+                "aic_model_path": "/models/mock",
+            }
+        )
+    )
+
+    assert engine_args.max_model_len is None
+
+
+def test_replay_engine_args_preserves_explicit_max_model_len():
+    import dynamo.replay.main as replay_main
+
+    engine_args = replay_main._load_engine_args(
+        json.dumps(
+            {
+                "num_gpu_blocks": 4096,
+                "max_model_len": 32768,
+                "aic_model_path": "/models/mock",
+            }
+        )
+    )
+
+    assert engine_args.max_model_len == 32768
+
+
 def test_replay_engine_args_compute_kv_bytes_for_g3_before_validation(monkeypatch):
     import dynamo.replay.main as replay_main
 

components/src/dynamo/replay/main.py

@@ -295,7 +295,7 @@ def _engine_caps(args: MockEngineArgs) -> EngineCapabilities:
         num_gpu=1,
         max_num_batched_tokens=args.max_num_batched_tokens,
         max_num_seqs=args.max_num_seqs,
-        context_length=max_kv_tokens if max_kv_tokens > 0 else None,
+        context_length=args.max_model_len,
         max_kv_tokens=max_kv_tokens if max_kv_tokens > 0 else None,
         speculative_nextn=args.aic_nextn,
     )

lib/bindings/python/rust/llm/replay.rs

@@ -171,7 +171,7 @@ impl MockEngineArgs {
 #[pymethods]
 impl MockEngineArgs {
     #[new]
-    #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=None, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, aic_moe_tp_size=None, aic_moe_ep_size=None, aic_attention_dp_size=None, aic_nextn=None, aic_nextn_accept_rates=None, aic_mtp_seed=42, aic_gemm_dtype=None, aic_moe_dtype=None, aic_fmha_dtype=None, aic_kv_cache_dtype=None, aic_comm_dtype=None, gpu_memory_utilization=None, mem_fraction_static=None, free_gpu_memory_fraction=None, enable_local_indexer=false, bootstrap_port=None, handoff_session_timeout_ms=300000, kv_bytes_per_token=None, kv_transfer_bandwidth=None, kv_transfer_timing_mode="full_prompt", reasoning=None, response_replay_trace_path=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None, trtllm=None, num_g2_blocks=None, num_g3_blocks=None, offload_batch_size=None, bandwidth_g1_to_g2_gbps=None, bandwidth_g2_to_g1_gbps=None, bandwidth_g2_to_g3_gbps=None, bandwidth_g3_to_g2_gbps=None, enable_g4_storage=false, bandwidth_g2_to_g4_gbps=None, bandwidth_g4_to_g2_gbps=None))]
+    #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=None, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, aic_moe_tp_size=None, aic_moe_ep_size=None, aic_attention_dp_size=None, aic_nextn=None, aic_nextn_accept_rates=None, aic_mtp_seed=42, aic_gemm_dtype=None, aic_moe_dtype=None, aic_fmha_dtype=None, aic_kv_cache_dtype=None, aic_comm_dtype=None, gpu_memory_utilization=None, mem_fraction_static=None, free_gpu_memory_fraction=None, enable_local_indexer=false, bootstrap_port=None, handoff_session_timeout_ms=300000, kv_bytes_per_token=None, kv_transfer_bandwidth=None, kv_transfer_timing_mode="full_prompt", reasoning=None, response_replay_trace_path=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None, trtllm=None, num_g2_blocks=None, num_g3_blocks=None, offload_batch_size=None, bandwidth_g1_to_g2_gbps=None, bandwidth_g2_to_g1_gbps=None, bandwidth_g2_to_g3_gbps=None, bandwidth_g3_to_g2_gbps=None, enable_g4_storage=false, bandwidth_g2_to_g4_gbps=None, bandwidth_g4_to_g2_gbps=None, max_model_len=None))]
     #[allow(clippy::too_many_arguments)]
     fn new(
         engine_type: &str,
@@ -230,6 +230,7 @@ impl MockEngineArgs {
         enable_g4_storage: bool,
         bandwidth_g2_to_g4_gbps: Option<f64>,
         bandwidth_g4_to_g2_gbps: Option<f64>,
+        max_model_len: Option<usize>,
     ) -> PyResult<Self> {
         let engine_type = parse_mocker_engine_type(engine_type)?;
         let worker_type = parse_worker_type(worker_type)?;
@@ -248,6 +249,7 @@ impl MockEngineArgs {
         let mut builder = RsMockEngineArgs::builder()
             .engine_type(engine_type)
             .block_size(block_size)
+            .max_model_len(max_model_len)
             .max_num_seqs(max_num_seqs)
             .max_num_batched_tokens(max_num_batched_tokens)
             .enable_prefix_caching(enable_prefix_caching)
@@ -365,6 +367,11 @@ impl MockEngineArgs {
         self.inner.num_gpu_blocks
     }
 
+    #[getter]
+    fn max_model_len(&self) -> Option<usize> {
+        self.inner.max_model_len
+    }
+
     #[getter]
     fn max_num_seqs(&self) -> Option<usize> {
         self.inner.max_num_seqs

lib/bindings/python/src/dynamo/_core.pyi

@@ -1965,6 +1965,7 @@ class MockEngineArgs:
         enable_g4_storage: bool = False,
         bandwidth_g2_to_g4_gbps: Optional[float] = None,
         bandwidth_g4_to_g2_gbps: Optional[float] = None,
+        max_model_len: Optional[int] = None,
     ) -> None:
         ...
 
@@ -1983,6 +1984,9 @@ class MockEngineArgs:
     @num_gpu_blocks.setter
     def num_gpu_blocks(self, value: int) -> None: ...
 
+    @property
+    def max_model_len(self) -> Optional[int]: ...
+
     @property
     def max_num_seqs(self) -> Optional[int]: ...
 

lib/llm/src/mocker.rs

@@ -791,6 +791,12 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
         let max_output_tokens = planned_output_token_ids
             .as_ref()
             .map_or(requested_max_output_tokens, Vec::len);
+        let effective_max_output_tokens =
+            self.engine_args
+                .max_model_len
+                .map_or(max_output_tokens, |max_model_len| {
+                    max_output_tokens.min(max_model_len.saturating_sub(request.token_ids.len()))
+                });
         let native_timing = self
             .native_metrics
             .request_timing(&request.model, dp_rank, is_prefill, request_start)
@@ -1013,14 +1019,14 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
                             break;
                         };
 
-                        // A terminally rejected request never ran (its footprint
-                        // exceeds the KV pool): emit no token and do not complete the
-                        // bootstrap room — surface the rejection and end the stream
-                        // before any token/prefill bookkeeping.
+                        // A terminally rejected request never ran because it violated
+                        // a worker admission limit. Emit no token and do not complete
+                        // the bootstrap room; surface the rejection before any
+                        // token/prefill bookkeeping.
                         if signal.rejected {
                             handoff_cancel.cancel();
                             let _ = stream_tx.send(LLMEngineOutput::error(
-                                "request rejected: KV footprint exceeds pool capacity".to_string(),
+                                "request rejected: request exceeds worker admission limits".to_string(),
                             ));
                             break;
                         }
@@ -1043,7 +1049,7 @@ impl AsyncEngine<SingleIn<PreprocessedRequest>, ManyOut<LLMEngineOutput>, Error>
                             ..Default::default()
                         };
 
-                        if signal.completed && token_count < max_output_tokens {
+                        if signal.completed && token_count < effective_max_output_tokens {
                             let _ = stream_tx.send(LLMEngineOutput::error("Completion signal received before max tokens reached".to_string()));
                             break;
                         }
@@ -1239,6 +1245,22 @@ mod tests {
             .unwrap()
     }
 
+    fn decode_request(prompt_tokens: usize, max_tokens: u32) -> PreprocessedRequest {
+        PreprocessedRequest::builder()
+            .model("mock".to_string())
+            .token_ids(vec![1; prompt_tokens])
+            .stop_conditions(StopConditions {
+                max_tokens: Some(max_tokens),
+                ..Default::default()
+            })
+            .sampling_options(SamplingOptions::default())
+            .output_options(OutputOptions::default())
+            .eos_token_ids(vec![])
+            .annotations(vec![])
+            .build()
+            .unwrap()
+    }
+
     #[tokio::test(start_paused = true)]
     async fn no_bootstrap_prefill_delays_terminal_finish_once() {
         let args = MockEngineArgs::builder()
@@ -1284,6 +1306,43 @@ mod tests {
         assert!(stream.next().await.is_none());
     }
 
+    #[tokio::test]
+    async fn context_capped_completion_maps_to_length() {
+        let args = MockEngineArgs::builder()
+            .max_model_len(Some(4))
+            .build()
+            .unwrap();
+        let engine = MockEngine::new(args);
+        let (request_tx, mut request_rx) = tokio::sync::mpsc::unbounded_channel();
+        engine.request_senders.set(vec![request_tx]).unwrap();
+
+        let mut stream = engine
+            .generate(SingleIn::new(decode_request(3, 4)))
+            .await
+            .unwrap();
+        let request = request_rx.recv().await.unwrap();
+        assert_eq!(request.max_output_tokens, 4);
+        let request_id = request.uuid.unwrap();
+        engine
+            .active_requests
+            .get(&request_id)
+            .unwrap()
+            .send(OutputSignal {
+                uuid: request_id,
+                token_id: Some(42),
+                completed: true,
+                rejected: false,
+                handoff_delay_ms: None,
+            })
+            .unwrap();
+
+        let token = stream.next().await.unwrap();
+        assert_eq!(token.token_ids.len(), 1);
+        assert!(token.finish_reason.is_none());
+        assert_eq!(stream.next().await.unwrap(), LLMEngineOutput::length());
+        assert!(stream.next().await.is_none());
+    }
+
     #[test]
     fn unbounded_sequence_limit_uses_finite_multi_handoff_capacity() {
         let args = MockEngineArgs::builder()