fix: allow full-block state alignment with speculation

crates/core/src/engine/config.rs

@@ -804,8 +804,10 @@ impl EngineConfig {
                 "state_cache does not support g3_offload in the G1-only implementation"
             );
             ensure!(
-                self.prefix_match_unit.is_none() || self.aic_nextn.is_none(),
-                "prefix_match_unit does not support speculative decoding"
+                self.prefix_match_unit.is_none()
+                    || self.prefix_match_unit == Some(self.block_size)
+                    || self.aic_nextn.is_none(),
+                "sub-block prefix_match_unit does not support speculative decoding"
             );
             let bytes_per_token = self
                 .kv_cache_bytes_per_token
@@ -1129,11 +1131,16 @@ mod tests {
         );
         input["aic_nextn"] = serde_json::json!(1);
         assert!(
-            serde_json::from_value::<EngineConfig>(input)
+            serde_json::from_value::<EngineConfig>(input.clone())
                 .unwrap_err()
                 .to_string()
                 .contains("speculative")
         );
+        input["prefix_match_unit"] = serde_json::json!(64);
+        serde_json::from_value::<EngineConfig>(input)
+            .unwrap()
+            .validate()
+            .unwrap();
     }
 
     #[test]

crates/core/src/engine/scheduler/vllm/core.rs

@@ -4163,6 +4163,34 @@ mod state_cache_tests {
         }
     }
 
+    #[test]
+    fn full_block_alignment_retains_state_with_mtp() {
+        for rates in ["0,0,0", "1,1,1", "1,0.57,0.5"] {
+            let mut config = args(64, true, Some(3));
+            config.prefix_match_unit = Some(4);
+            config.aic_nextn_accept_rates = Some(rates.to_owned());
+            let mut core = VllmCore::new(config);
+            submit(&mut core, 93_032, 17, 16);
+            for _ in 0..32 {
+                if core.state.requests.is_empty() {
+                    break;
+                }
+                step(&mut core);
+            }
+            assert!(core.state.requests.is_empty());
+            let second = submit(&mut core, 93_033, 17, 1);
+            let pass = step(&mut core);
+            assert_eq!(
+                pass.admissions
+                    .iter()
+                    .find(|admission| admission.uuid == second)
+                    .unwrap()
+                    .reused_input_tokens,
+                16
+            );
+        }
+    }
+
     #[test]
     fn empty_native_prompts_are_rejected_before_state_or_speculative_generation() {
         for nextn in [None, Some(3)] {

docs/replay/engine/kv-cache.md

@@ -306,7 +306,10 @@ transfers the state for the whole prompt and does not model the recompute.
 
 Supported: aggregated and P/D vLLM, fixed `capacity`, `predict --stack engine`;
 inferred geometry for Kimi K3 and DeepSeek V4 at PP=1.
-Rejected: `recommend`, G2/G3 offload, and speculative decoding with
+Speculative decoding supports `prefix_match_unit` equal to `block_size`.
+Sub-block matching with speculation remains unsupported.
+
+Rejected: `recommend`, G2/G3 offload, and speculative decoding with sub-block
 `prefix_match_unit`.
 
 ## Other cache modes

docs/replay/engine/speculation.md

@@ -105,7 +105,7 @@ follow the [cost-model contract](../../perf-model/configuration.md#speculative-d
 an explicit FPM request with `fallback_policy: deny` fails, while an allowed
 fallback can select OP without discarding the scheme.
 
-Grouped caches, G2/G3 offload, state-cache `prefix_match_unit`, AFD/EPD,
+Grouped caches, G2/G3 offload, sub-block state-cache `prefix_match_unit`, AFD/EPD,
 TensorRT-LLM, and optional Dynamo-stack execution are outside this generic
 replay path. Functional replay does not establish AgentX prediction accuracy,
 model quality, or measured hardware speedup; see [feature support](../features.md).
@@ -204,7 +204,8 @@ HBM-only agentic MTP/EAGLE replay is supported by the offline engine stack,
 including aggregated and P/D configurations supported by the selected backend.
 This is functional replay, not a claim of AgentX hardware accuracy.
 Rejected with MTP: grouped caches, G2/G3 offload, and state-cache prefix matching
-with `prefix_match_unit`. Combinations accepted by the
+with `prefix_match_unit` smaller than `block_size`. Full-block alignment is supported.
+Combinations accepted by the
 timing model are not necessarily supported by Replay; see
 [feature support](../features.md).
 

docs/replay/features.md

@@ -15,8 +15,8 @@ is a separate [performance-model question](../perf-model/support-matrix.md).
 | --- | --- | --- | --- |
 | Native token scheduling | vLLM, SGLang, TensorRT-LLM | Aggregated or P/D; synthetic and compatible traces | Engine-specific admission/cache behavior; no real network or host-runtime overhead |
 | Attention-DP | All three native backends | Aggregated and independently sized P/D pools | Rank-local scheduling and G1; one grouped pass completes at the maximum rank time |
-| Native MTP | Supported engine/model path | Token or HBM-only agentic replay with explicit acceptance | Separate from ngram; no grouped-cache/G2/G3/fine state-cache combination; agentic accuracy is functional-only |
-| [Registered speculative schemes](engine/speculation.md#registered-schemes) | Offline vLLM/SGLang, subject to scheme/model constraints | Prediction, aggregated or P/D; token or HBM-only agentic replay | Existing linear 1–5-token executor; fixed KV capacity; default timing with supported OP cost; no trees, grouped-cache/G2/G3/fine state-cache combination or generic-scheme recommendation |
+| Native MTP | Supported engine/model path | Token or HBM-only agentic replay with explicit acceptance | Separate from ngram; no grouped-cache/G2/G3/sub-block state-cache combination; agentic accuracy is functional-only |
+| [Registered speculative schemes](engine/speculation.md#registered-schemes) | Offline vLLM/SGLang, subject to scheme/model constraints | Prediction, aggregated or P/D; token or HBM-only agentic replay | Existing linear 1–5-token executor; fixed KV capacity; default timing with supported OP cost; no trees, grouped-cache/G2/G3/sub-block state-cache combination or generic-scheme recommendation |
 | Ngram prompt lookup | Offline vLLM | Aggregated or P/D, op-level or explicit timing; HBM-only agentic replay | Fixed draft/acceptance assumption; no lookup CPU cost, adapters or deployment generation |
 | Sessions | Native token replay | Synthetic or compatible trace sessions | Concurrency counts sessions through think time; see trace-format restrictions |
 | Prefix reuse | Native backend cache | G1 with prefix caching enabled | Actual reuse differs from Router overlap; complete block and backend admission rules apply |

python/aisimulate/src/aisimulate/config/engine.py

@@ -1053,8 +1053,14 @@ def _validate_prediction_state_cache(engine: EnginePredictionConfig) -> None:
         if worker is None or worker.kv_cache.state_cache is None:
             continue
         roles.append(role)
-        if worker.kv_cache.prefix_match_unit is not None and (engine.speculation is not None or engine.nextn > 0):
-            raise ValueError("prefix_match_unit does not support speculative decoding (speculation or nextn > 0)")
+        if (
+            worker.kv_cache.prefix_match_unit is not None
+            and worker.kv_cache.prefix_match_unit != worker.kv_cache.block_size
+            and (engine.speculation is not None or engine.nextn > 0)
+        ):
+            raise ValueError(
+                "sub-block prefix_match_unit does not support speculative decoding (speculation or nextn > 0)"
+            )
         supported = role == "aggregated" if engine.mode == "aggregated" else engine.mode == "disaggregated"
         if engine.backend != "vllm" or not supported:
             raise ValueError("state_cache requires backend=vllm and mode=aggregated or disaggregated (G1 only)")

python/aisimulate/src/aisimulate/runner.py

@@ -1936,8 +1936,12 @@ def _manual_state_cache(rank: Mapping[str, JSONValue], backend: str, role: str)
     if raw is None:
         return None
     state_cache = StateCacheConfig.model_validate(raw)
-    if rank.get("prefix_match_unit") is not None and rank.get("aic_nextn") is not None:
-        raise ValueError("prefix_match_unit does not support speculative decoding")
+    if (
+        rank.get("prefix_match_unit") is not None
+        and rank["prefix_match_unit"] != rank.get("block_size")
+        and rank.get("aic_nextn") is not None
+    ):
+        raise ValueError("sub-block prefix_match_unit does not support speculative decoding")
     if backend != "vllm" or role not in ("aggregated", "prefill", "decode"):
         raise ValueError("state_cache requires backend=vllm and an aggregated, prefill or decode G1 worker")
     if rank.get("native_host_offload") is not None:

tests/test_state_cache_config.py

@@ -420,6 +420,28 @@ def test_aligned_state_config_allows_zero_nextn_through_engine_handoff(forbid_es
     assert rank["prefix_match_unit"] == 16
 
 
+@pytest.mark.parametrize("canonical", [False, True])
+def test_full_block_aligned_state_allows_mtp_through_engine_handoff(forbid_estimators, canonical):
+    raw = _public(timing="default", prefix_match_unit=64)
+    if canonical:
+        raw["engine"]["speculation"] = {
+            "kind": "mtp",
+            "params": {"depth": 3},
+            "expected_accepted_tokens": 2.12,
+            "seed": 42,
+        }
+    else:
+        raw["engine"].update(nextn=3, nextn_accepted=2.12)
+    config = CorePredictionConfig.model_validate(raw)
+    config = CorePredictionConfig.model_validate_json(config.model_dump_json())
+    spec = prediction_to_replay_spec(config)
+    runtime = RecordingRuntime()
+    EngineReplayRunnerFactory(runtime=runtime).create(0).run(spec)
+    rank = runtime.execution_spec["spec"]["engine"]["rank"]
+    assert rank["aic_nextn"] == 3
+    assert rank["prefix_match_unit"] == rank["block_size"] == 64
+
+
 @pytest.mark.parametrize(
     "shared_prefix_tokens, expected_reused_tokens",
     [