fix(mocker): reject prompts above max model length

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

Author: dreamtalen

components/src/dynamo/mocker/args.py

@@ -211,6 +211,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=int,
+        default=None,
+        help="Reject vLLM prompts longer than this many tokens. When omitted, "
+        "derive the limit from --model-path when available; otherwise disabled.",
+    )
     parser.add_argument(
         "--max-num-seqs",
         type=int,

components/src/dynamo/mocker/config.py

@@ -14,6 +14,7 @@
 from dynamo.common.utils.topology import apply_topology_config
 from dynamo.llm import ModelRuntimeConfig
 from dynamo.mocker import MockEngineArgs, ReasoningConfig, SglangArgs, TrtllmArgs
+from dynamo.mocker.utils.kv_cache import get_model_max_context_length
 
 _DEFAULT_NUM_GPU_BLOCKS = 16384
 _DEFAULT_MAX_NUM_SEQS = 256
@@ -183,14 +184,20 @@ def _resolve_raw_engine_args(
     *,
     fallback_model_path: str | None = None,
 ) -> dict:
+    engine_type = raw.get("engine_type") or "vllm"
+    model_path = raw.get("aic_model_path") or fallback_model_path
+    if raw.get("max_model_len") is None and engine_type == "vllm" and model_path:
+        max_model_len = get_model_max_context_length(model_path)
+        if max_model_len is not None:
+            raw["max_model_len"] = max_model_len
+
     if raw.get("num_gpu_blocks") is not None:
         return raw
 
     aic_backend = raw.get("aic_backend")
     if aic_backend is None:
         return raw
 
-    engine_type = raw.get("engine_type") or "vllm"
     sglang = raw.get("sglang")
     sglang_page_size = sglang.get("page_size") if isinstance(sglang, dict) else None
     raw["num_gpu_blocks"] = _estimate_aic_num_gpu_blocks(
@@ -201,7 +208,7 @@ def _resolve_raw_engine_args(
         aic_system=raw.get("aic_system"),
         aic_backend_version=raw.get("aic_backend_version"),
         aic_tp_size=raw.get("aic_tp_size"),
-        aic_model_path=raw.get("aic_model_path") or fallback_model_path,
+        aic_model_path=model_path,
         aic_moe_tp_size=raw.get("aic_moe_tp_size"),
         aic_moe_ep_size=raw.get("aic_moe_ep_size"),
         aic_attention_dp_size=raw.get("aic_attention_dp_size"),
@@ -243,6 +250,10 @@ 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)
+    model_path = getattr(args, "model_path", None)
+    if max_model_len is None and engine_type == "vllm" and model_path:
+        max_model_len = get_model_max_context_length(model_path)
     num_gpu_blocks = _resolve_num_gpu_blocks(
         explicit_num_gpu_blocks=getattr(args, "num_gpu_blocks", None),
         engine_type=engine_type,
@@ -267,6 +278,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 +361,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,127 @@ 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
+
+
+def test_model_max_context_length_uses_shared_model_info(monkeypatch):
+    import dynamo.mocker.utils.kv_cache as kv_cache
+
+    monkeypatch.setattr(
+        kv_cache,
+        "get_model_info",
+        lambda _model_path: SimpleNamespace(max_context_length=204800),
+    )
+
+    assert kv_cache.get_model_max_context_length("/models/mock") == 204800
+
+
+def test_model_max_context_length_returns_none_on_failure(monkeypatch):
+    import dynamo.mocker.utils.kv_cache as kv_cache
+
+    def fail(_model_path):
+        raise ValueError("bad model config")
+
+    monkeypatch.setattr(kv_cache, "get_model_info", fail)
+
+    assert kv_cache.get_model_max_context_length("/models/mock") is None
+
+
+def test_build_mocker_engine_args_derives_max_model_len(monkeypatch):
+    monkeypatch.setattr(
+        CONFIG, "get_model_max_context_length", lambda _model_path: 204800
+    )
+
+    engine_args = CONFIG.build_mocker_engine_args(
+        make_args(model_path="/models/mock", num_gpu_blocks=4096)
+    )
+
+    assert engine_args.max_model_len == 204800
+
+
+def test_build_mocker_engine_args_preserves_explicit_max_model_len(monkeypatch):
+    def fail(_model_path):
+        raise AssertionError("model context lookup should not be called")
+
+    monkeypatch.setattr(CONFIG, "get_model_max_context_length", fail)
+
+    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_derives_max_model_len(monkeypatch):
+    import dynamo.replay.main as replay_main
+
+    monkeypatch.setattr(
+        replay_main, "get_model_max_context_length", lambda _model_path: 204800
+    )
+
+    engine_args = replay_main._load_engine_args(
+        json.dumps(
+            {
+                "num_gpu_blocks": 4096,
+                "aic_model_path": "/models/mock",
+            }
+        )
+    )
+
+    assert engine_args.max_model_len == 204800
+
+
+def test_replay_engine_args_keeps_max_model_len_unset_on_lookup_failure(monkeypatch):
+    import dynamo.replay.main as replay_main
+
+    monkeypatch.setattr(
+        replay_main, "get_model_max_context_length", lambda _model_path: None
+    )
+
+    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(monkeypatch):
+    import dynamo.replay.main as replay_main
+
+    def fail(_model_path):
+        raise AssertionError("model context lookup should not be called")
+
+    monkeypatch.setattr(replay_main, "get_model_max_context_length", fail)
+
+    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/mocker/utils/kv_cache.py

@@ -35,6 +35,18 @@
 DEFAULT_KV_TRANSFER_BANDWIDTH_GBPS = 64.0
 
 
+def get_model_max_context_length(model_path: str) -> int | None:
+    """Best-effort model context length from the shared model metadata loader."""
+    try:
+        max_context_length = get_model_info(model_path).max_context_length
+        if max_context_length is None or max_context_length <= 0:
+            return None
+        return max_context_length
+    except Exception as e:
+        logger.warning("Could not determine max context length from model config: %s", e)
+        return None
+
+
 def _normalize_dtype_str(dtype) -> str:
     """Normalize a dtype to a plain string like 'float16'.
 

components/src/dynamo/replay/main.py

@@ -29,7 +29,10 @@
 )
 from dynamo.llm import AicPerfConfig, KvRouterConfig
 from dynamo.mocker import MockEngineArgs
-from dynamo.mocker.utils.kv_cache import compute_kv_bytes_per_token
+from dynamo.mocker.utils.kv_cache import (
+    compute_kv_bytes_per_token,
+    get_model_max_context_length,
+)
 from dynamo.replay import run_synthetic_trace_replay, run_trace_replay
 from dynamo.replay.reporting import format_report_table, write_report_json
 
@@ -197,6 +200,21 @@ def _resolve_kv_bytes_per_token(raw: dict) -> None:
         raw["kv_bytes_per_token"] = kv_bytes_per_token
 
 
+def _resolve_max_model_len(raw: dict) -> None:
+    if raw.get("max_model_len") is not None:
+        return
+    if (raw.get("engine_type") or "vllm") != "vllm":
+        return
+
+    model_path = raw.get("aic_model_path")
+    if not model_path:
+        return
+
+    max_model_len = get_model_max_context_length(model_path)
+    if max_model_len is not None:
+        raw["max_model_len"] = max_model_len
+
+
 def _load_engine_args(raw_args: str | None):
     if raw_args is None:
         return None
@@ -229,6 +247,7 @@ def _load_engine_args(raw_args: str | None):
                 raw["planner_profile_data"] = str(profile_data_result.npz_path)
             else:
                 del raw["planner_profile_data"]
+    _resolve_max_model_len(raw)
     _resolve_aic_num_gpu_blocks(raw)
     _resolve_kv_bytes_per_token(raw)
     return MockEngineArgs.from_json(json.dumps(raw))
@@ -295,7 +314,8 @@ 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
+        or (max_kv_tokens if max_kv_tokens > 0 else None),
         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

@@ -1958,6 +1958,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:
         ...
 
@@ -1976,6 +1977,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]: ...