fix: count only uncached tokens against the legacy agg prefill budget

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

@@ -16,7 +16,11 @@
   "post_freeze_pins": {
     "chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step": {
       "engine": "rust",
-      "git_head": "0a51476b6ab90bc0e475bd41d4b1c7abbef07b95"
+      "git_head": "3eae51edd74cd072546a1502bcc029ff67220e53"
+    },
+    "chunked_prefill::ctx4096_gen4_isl4096_osl128_prefix256::mixed_step": {
+      "engine": "rust",
+      "git_head": "3eae51edd74cd072546a1502bcc029ff67220e53"
     },
     "chunked_prefill::ctx512_gen4_isl4096_osl128_prefix0::mixed_step": {
       "engine": "rust",
@@ -228,7 +232,8 @@
     }
   },
   "references": {
-    "chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step": 16.081171352808852,
+    "chunked_prefill::ctx300_gen7_isl1000_osl64_prefix100::mixed_step": 16.414675748089216,
+    "chunked_prefill::ctx4096_gen4_isl4096_osl128_prefix256::mixed_step": 57.89342286133688,
     "chunked_prefill::ctx512_gen4_isl4096_osl128_prefix0::mixed_step": 17.664446149246828,
     "chunked_prefill::ctx512_gen4_isl4096_osl128_prefix256::mixed_step": 17.662495086777085,
     "gpt-oss-20b-b200-trtllm-isl1024-osl2::decode_step": 2.213361158220323,

crates/core/parity_tests/perfmodel/pin_goldens.py

@@ -177,7 +177,11 @@ def _compile_case_references(case) -> dict[str, float]:
     return {
         "static_ctx": float(ctx),
         "static_gen": float(gen),
-        "mixed_step": float(handle.mixed_step_latency(case.isl, case.batch_size, case.isl, osl, case.prefix)),
+        "mixed_step": float(
+            handle.mixed_step_latency(
+                compile_parity._default_ctx_tokens(case), case.batch_size, case.isl, osl, case.prefix
+            )
+        ),
         "decode_step": float(handle.decode_step_latency(case.batch_size, case.isl, osl)),
     }
 

crates/core/parity_tests/perfmodel/test_compile_engine_parity.py

@@ -306,9 +306,16 @@ def _assert_within(name: str, python_value: float, new_value: float, *, backend:
     (512, 4, 4096, 128, 0),  # chunked prefill: ctx_tokens < isl
     (512, 4, 4096, 128, 256),  # chunked + cached prefix
     (300, 7, 1000, 64, 100),  # ragged chunk + prefix + decode overlap
+    (4096, 4, 4096, 128, 256),  # isl-sized budget + cached prefix: one request plus a partial one
 ]
 
 
+def _default_ctx_tokens(case) -> int:
+    """The agg default budget: one request's uncached prefill. Shared by the
+    mixed-step test and the golden pin path (pin_goldens.py)."""
+    return max(case.isl - case.prefix, 1)
+
+
 def _chunked_prefill_key(ctx_tokens: int, gen_tokens: int, isl: int, osl: int, prefix: int) -> str:
     return f"chunked_prefill::ctx{ctx_tokens}_gen{gen_tokens}_isl{isl}_osl{osl}_prefix{prefix}::mixed_step"
 
@@ -333,7 +340,11 @@ class TestCompileEngineMixedStepParity:
     @pytest.mark.parametrize("case", _SUBSET_CASES)
     def test_mixed_step(self, case: EngineStepParityCase) -> None:
         handle = _compile_handle(case)
-        new_val = handle.mixed_step_latency(case.isl, case.batch_size, case.isl, max(case.osl, 2), case.prefix)
+        # One request's uncached prefill as the budget (the agg default), as in
+        # test_engine_step_parity._mix_step_shape.
+        new_val = handle.mixed_step_latency(
+            _default_ctx_tokens(case), case.batch_size, case.isl, max(case.osl, 2), case.prefix
+        )
         py_val = _case_reference(case, "mixed_step")
         _assert_within("mixed_step", py_val, new_val, backend=case.backend_name)
 

crates/core/parity_tests/perfmodel/test_engine_step_parity.py

@@ -758,7 +758,8 @@ def call():
             backend_name=case.backend_name,
             backend_version=case.backend_version,
             batch_size=case.agg_batch_size,
-            ctx_tokens=case.agg_ctx_tokens or case.isl,
+            # The CLI default budget: one request's uncached prefill.
+            ctx_tokens=case.agg_ctx_tokens or max(case.isl - case.prefix, 1),
             isl=case.isl,
             osl=case.osl,
             prefix=case.prefix,
@@ -896,12 +897,13 @@ def call():
 def _mix_step_shape(case: EngineStepParityCase) -> dict:
     """Mix-step (chunked-prefill + decode) shape for a smoke case.
 
-    Treats the case's single prefill request as one chunk with `case.isl`
-    isl-equivalent tokens (matching Python's agg orchestration). Decode
-    batch is `case.batch_size` (matches the FPM constructor below).
+    Treats the case's single prefill request as one chunk: the budget is its
+    `case.isl - case.prefix` uncached tokens (the agg default), over
+    `case.prefix` cached ones. Decode batch is `case.batch_size` (matches the
+    FPM constructor below).
     """
     return {
-        "ctx_tokens": case.isl,
+        "ctx_tokens": max(case.isl - case.prefix, 1),
         "gen_tokens": case.batch_size,
         "isl": case.isl,
         "osl": max(case.osl, 2),