Guard unsafe prompt cache reuse

Author: apstenku123

cppmega_mlx/inference/generation.py

@@ -34,6 +34,7 @@
 _SEQUENCE_ALIGNED_MODEL_KWARGS = (
     _ZERO_GENERATED_MODEL_KWARGS | _REPEAT_GENERATED_MODEL_KWARGS
 )
+_PROMPT_CACHE_UNSAFE_ROUTE_ROLES = frozenset({"mamba3", "m2rnn", "engram"})
 
 
 @dataclass(frozen=True)
@@ -474,6 +475,7 @@ def build_prompt_cache(
     max_seq_length = _model_max_seq_length(model)
     if max_seq_length is not None and prompt_ids.shape[1] > max_seq_length:
         raise ValueError("prompt_ids already exceed model.config.max_seq_length")
+    _require_prompt_cache_safe_model(model)
 
     kv_cache = _resolve_kv_cache(
         cache=cache,
@@ -538,6 +540,7 @@ def generate_tokens_with_prompt_cache(
     max_seq_length = _model_max_seq_length(model)
     if max_seq_length is not None and prompt_ids.shape[1] > max_seq_length:
         raise ValueError("prompt_ids already exceed model.config.max_seq_length")
+    _require_prompt_cache_safe_model(model)
 
     tokens = prompt_ids
     kv_cache = clone_contiguous_kv_cache(prompt_cache.cache)
@@ -910,6 +913,20 @@ def _validate_prompt_cache_prefix(
         raise ValueError("prompt_ids must start with prompt_cache.prompt_ids")
 
 
+def _require_prompt_cache_safe_model(model: Any) -> None:
+    route_roles = getattr(model, "route_roles", None)
+    if route_roles is None:
+        return
+    roles = tuple(str(role) for role in route_roles)
+    unsafe = tuple(role for role in roles if role in _PROMPT_CACHE_UNSAFE_ROUTE_ROLES)
+    if unsafe:
+        joined = ", ".join(unsafe)
+        raise ValueError(
+            "prompt cache is only validated for attention/stateless routes; "
+            f"route roles requiring their own state cache are present: {joined}"
+        )
+
+
 def _propose_speculative_draft_window(
     draft_model: Any,
     tokens: mx.array,

docs/mlx_port_master_plan.md

@@ -43,7 +43,7 @@ This document is a planning synthesis from parallel research lanes (local-code a
 - **Sharded checkpoints** (model.safetensors.index.json).
 - **Sequence packing** (cumulative-doc-id attention mask).
 - **Tokenizer**: M0.1 is closed. The vendored GB10 BPE artifact and special-token contract exist for vocab=65536 with id7=<CODE_START>, id45=<FIM_INSTRUCTION>, id46=<SPACE>, id47=<NL>. Encode normalizes whitespace runs ([\r\n]+-><NL>, [ \t]+-><SPACE>); decode is plain token concat with sentinel substitution. Same wrapper deployed on Mac and gb10/CUDA side.
-- **Inference / generation**: temperature/top-k/top-p sampling, eager no-KV full-prefix generation, contiguous KV-cache generation/streaming, paged KV scheduler metadata with fail-closed model-integration guard, contiguous prompt-cache reuse, fail-closed standard MTP/draft-output guard, prompt-only FIM/iFIM infilling construction, eager vanilla speculative decoding, and eager MTP self-speculative decoding exist locally. Still open: model-integrated paged attention/API serving, sliding-window/SSM prompt-cache safety, integrated FIM decode/postprocess loop, EAGLE/token-recycling speculative paths, q4 path, and KV-q4 path.
+- **Inference / generation**: temperature/top-k/top-p sampling, eager no-KV full-prefix generation, contiguous KV-cache generation/streaming, paged KV scheduler metadata with fail-closed model-integration guard, contiguous prompt-cache reuse, fail-closed prompt-cache safety guards for SSM/stateful routes, fail-closed standard MTP/draft-output guard, prompt-only FIM/iFIM infilling construction, eager vanilla speculative decoding, and eager MTP self-speculative decoding exist locally. Still open: model-integrated paged attention/API serving, integrated FIM decode/postprocess loop, EAGLE/token-recycling speculative paths, q4 path, and KV-q4 path.
 - **Quantization**: bf16 only; no mx.quantize integration, no q4 inference path.
 - **Structural parity anchors, not CUDA tensor parity**: existing tests/test_cppmega_parity_anchors.py checks NAM56R route/layer constants, DSA/MLA layer derivation, vocab/MoE anchors, fail-closed parity wording, and optional sibling cppmega source-anchor presence when that checkout exists. The source-file existence check is only one guard; the test is broader than file-presence coverage, but it still does not compare CUDA golden tensors or prove numerical agreement.
 
@@ -887,7 +887,7 @@ Excluded as Hopper-only dead-end on Metal: cppmega/megatron/cute_dsl_mimo/ (sm_9
 169. **SCOPED DONE**: add `generate_tokens_mtp_self_speculative(...)` in `cppmega_mlx/inference/generation.py`, exported from `cppmega_mlx.inference`. The eager batch=1 loop uses the model-owned `mtp_head` to draft a bounded K-token window from the last verified hidden state, verifies the candidate prefix with one normal target forward over K+1 logits, and reuses the Leviathan-style acceptance-rejection helper. It supports the real `MinimalMTPHead` fallback plus explicit `mtp_head.draft(...)` adapters, and fails closed when no trained MTP head is attached or the requested draft window exceeds `mtp_head.config.depth`. This is not KV-cache/paged self-speculative serving, EAGLE-2, token recycling, throughput/quality benchmarking, or GB10/CUDA parity.
 170. **SCOPED DONE (local streaming slice)**: add `GenerationChunk` and `stream_generate_tokens(...)` in `cppmega_mlx/inference/generation.py`, exported from `cppmega_mlx.inference`. Covered eager full-prefix and contiguous-KV token streaming, batched rows, all-rows EOS stopping, optional per-token text decode, zero-new-token no-op, package export boundary, quantized-KV fail-closed behavior, and paged-attention non-integration guard. This is not full mlx-lm registry export, model-integrated paged attention, prompt cache, q4/KV-q4, speculative/self-spec decoding, OpenAI serving, throughput/quality/long-context benching, or GB10/CUDA parity.
 171. **SCOPED DONE (contiguous prompt-cache reuse)**: add PromptCacheEntry, clone_contiguous_kv_cache(...), build_prompt_cache(...), and generate_tokens_with_prompt_cache(...) for repeated-prefix reuse on the Mac-local contiguous KV path. Covered build/reuse, suffix decode, prefix mismatch fail-closed behavior, zero-new-token no-op, package exports, independent cache cloning, and real HybridTinyLM prompt-cache-vs-full-KV parity. This is not paged attention, OpenAI serving, sliding-window/SSM safety, q4/KV-q4, speculative decode, or CUDA/H200 parity.
-172. Validate prompt-cache safety with sliding-window/SSM against the installed mlx-lm prompt-cache APIs; do not assume a generic KV cache is safe for every route.
+172. **SCOPED DONE**: validate the local contiguous prompt-cache path as attention/stateless-route only. `build_prompt_cache(...)` and `generate_tokens_with_prompt_cache(...)` now fail closed for stateful route roles such as Mamba3/M2RNN/engram, because the contiguous KV entry does not carry SSM/recurrent/ngram state. Attention-only prompt-cache parity remains covered against full-prefix KV generation. This is not model-integrated sliding-window cache support, SSM-state prompt caching, paged prompt caching, or mlx-lm API parity for every upstream route.
 173. Add OpenAI-compatible serving endpoint (vllm-mlx pattern); cppmega_mlx/inference/serve.py.
 174. Add throughput benchmark: prefill / decode tok/s on Qwen3-4B-class and NAM56R-class models.
 175. Add quality benchmark: ARC / MMLU / HumanEval on q4-quantized model.

tests/test_inference_generation.py

@@ -293,6 +293,25 @@ def _tiny_attention_lm() -> HybridTinyLM:
     )
 
 
+def _tiny_mamba_lm() -> HybridTinyLM:
+    return HybridTinyLM(
+        HybridTinyConfig(
+            vocab_size=17,
+            hidden_size=8,
+            pattern="M",
+            depth=1,
+            dsa_a_layer_ranks=(),
+            num_attention_heads=1,
+            max_seq_length=8,
+            mamba_expand=1,
+            mamba_head_dim=4,
+            mamba_state_dim=4,
+            mamba_groups=1,
+            mamba_chunk_size=4,
+        )
+    )
+
+
 def test_generate_tokens_greedy_appends_full_prefix_steps() -> None:
     model = _ScriptedLogitsModel([[4], [5], [6]])
     prompt = mx.array([[1, 2, 3]], dtype=mx.int32)
@@ -1195,6 +1214,30 @@ def test_generate_tokens_with_prompt_cache_returns_prompt_for_zero_new_tokens()
     assert model.calls == calls_after_build
 
 
+def test_build_prompt_cache_rejects_mamba_route_without_ssm_state_cache() -> None:
+    model = _tiny_mamba_lm()
+    prompt = mx.array([[1, 2, 3]], dtype=mx.int32)
+
+    with pytest.raises(ValueError, match="prompt cache is only validated"):
+        build_prompt_cache(model, prompt)
+
+
+def test_generate_tokens_with_prompt_cache_rejects_mamba_route_reuse() -> None:
+    safe_model = _UpdatingKVScriptedLogitsModel([[4]])
+    prompt = mx.array([[1, 2, 3]], dtype=mx.int32)
+    prompt_cache = build_prompt_cache(safe_model, prompt, cache=_make_cache())
+    unsafe_model = _tiny_mamba_lm()
+
+    with pytest.raises(ValueError, match="prompt cache is only validated"):
+        generate_tokens_with_prompt_cache(
+            unsafe_model,
+            prompt,
+            prompt_cache=prompt_cache,
+            max_new_tokens=1,
+            temperature=0.0,
+        )
+
+
 def test_real_hybrid_tiny_lm_greedy_kv_cache_matches_full_prefix_generation() -> None:
     model = _tiny_attention_lm()
     prompt = mx.array([[1, 2, 3]], dtype=mx.int32)