diff --git a/experiment/verify_cached_decode_alignment.py b/experiment/verify_cached_decode_alignment.py
new file mode 100644
index 0000000..bb315d4
--- /dev/null
+++ b/experiment/verify_cached_decode_alignment.py
@@ -0,0 +1,394 @@
+"""
+ShadowPEFT Cached Decode 对齐验证 PoC.
+
+验证: 对纯 full-attention 模型 (Llama/Qwen2), ShadowPEFT 的 cached 单 token decode
+与 full-sequence 重算的逐 token logits 是否等价。
+
+核心结论 (已验证): ShadowPEFT 的 injection/update 是逐 token 纯函数, 无跨 token 递归,
+因此只需 base KV-cache + shadow backbone KV-cache 即可实现正确的增量解码。
+
+用法:
+    .venv/bin/python experiment/verify_cached_decode_alignment.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+from transformers import DynamicCache, PreTrainedModel
+
+from shadow_peft import ShadowConfig, ShadowPeftModel, get_shadow_model
+from shadow_peft.model_utils import _get_backbone
+
+# ---------------------------------------------------------------------------
+# Tiny model factories (mirrors tests/conftest.py patterns)
+# ---------------------------------------------------------------------------
+
+
+def _tiny_llama(vocab_size: int = 128, num_layers: int = 4, hidden_size: int = 32) -> PreTrainedModel:
+    from transformers import LlamaConfig, LlamaForCausalLM
+
+    cfg = LlamaConfig(
+        vocab_size=vocab_size,
+        hidden_size=hidden_size,
+        intermediate_size=hidden_size * 2,
+        num_hidden_layers=num_layers,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        max_position_embeddings=64,
+    )
+    return LlamaForCausalLM(cfg)
+
+
+def _tiny_qwen2(vocab_size: int = 128, num_layers: int = 4, hidden_size: int = 32) -> PreTrainedModel:
+    from transformers import Qwen2Config, Qwen2ForCausalLM
+
+    cfg = Qwen2Config(
+        vocab_size=vocab_size,
+        hidden_size=hidden_size,
+        intermediate_size=hidden_size * 2,
+        num_hidden_layers=num_layers,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        max_position_embeddings=64,
+    )
+    return Qwen2ForCausalLM(cfg)
+
+
+# ---------------------------------------------------------------------------
+# Shadow cached decode engine
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ShadowCache:
+    """复合 cache: 同时保存 base backbone 和 shadow backbone 的 KV-cache."""
+
+    base_cache: DynamicCache = field(default_factory=DynamicCache)
+    shadow_cache: DynamicCache = field(default_factory=DynamicCache)
+
+
+class ShadowCachedDecoder:
+    """
+    手动驱动 ShadowPEFT 的 cached decode, 绕过 ShadowPeftModel.forward (强制 use_cache=False)。
+
+    流程:
+      prefill: shadow_backbone(seq) -> shadow_hidden -> base_backbone(seq, shadow)
+      step:    shadow_backbone(token) -> shadow_token -> base_backbone(token, shadow)
+    """
+
+    def __init__(self, peft: ShadowPeftModel) -> None:
+        self.peft = peft
+        self.peft.eval()
+        self.base_backbone = _get_backbone(peft.base_model)
+        self.shadow_backbone = _get_backbone(peft.shadow_model)
+        self.base_embed = peft.base_model.get_input_embeddings()
+        self.lm_head = peft.base_model.get_output_embeddings()
+        self.shadow_projection = peft.shadow_hidden_projection
+
+        # Shadow embedding policy (mirrors ShadowPeftModel._compute_initial_shadow_hidden):
+        # - implicit shadow model / explicit w/ removed embed_tokens: share base embeddings
+        # - explicit shadow model w/ own embed_tokens: use shadow's own embeddings
+        self._shadow_share_base = (not peft._explicit_shadow_model) or peft._explicit_share_base_embeddings
+        self._shadow_embed = None
+        if not self._shadow_share_base:
+            shadow_get = getattr(peft.shadow_model, "get_input_embeddings", None)
+            if callable(shadow_get):
+                self._shadow_embed = shadow_get()
+
+    def _embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        emb = self.base_embed(input_ids)
+        if emb is None:
+            raise RuntimeError("Base model has no input embeddings.")
+        return emb
+
+    def _embed_shadow(self, input_ids: torch.Tensor) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        """
+        Return (inputs_embeds, input_ids) for the shadow backbone.
+
+        If sharing base embeddings, return (embeds, None).
+        If shadow has its own embeddings, return (None, input_ids).
+        """
+        if self._shadow_share_base:
+            return self._embed(input_ids), None
+        return None, input_ids
+
+    def _project_shadow(self, shadow_hidden: torch.Tensor) -> torch.Tensor:
+        if not isinstance(self.shadow_projection, torch.nn.Identity):
+            return self.shadow_projection(shadow_hidden)
+        return shadow_hidden
+
+    def _run_shadow(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        cache: DynamicCache,
+    ) -> torch.Tensor:
+        """Run the shadow backbone with the correct embedding source."""
+        inputs_embeds, shadow_input_ids = self._embed_shadow(input_ids)
+        out = self.shadow_backbone(
+            input_ids=shadow_input_ids,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            use_cache=True,
+            past_key_values=cache,
+        )
+        return self._project_shadow(out.last_hidden_state)
+
+    def _run_base(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        cache: DynamicCache,
+        shadow_hidden: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run base backbone with shadow injection/update active."""
+        self.peft._shadow_hidden_states = shadow_hidden
+        try:
+            out = self.base_backbone(
+                inputs_embeds=inputs_embeds,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                use_cache=True,
+                past_key_values=cache,
+            )
+        finally:
+            self.peft._shadow_hidden_states = None
+        return out.last_hidden_state
+
+    def prefill(self, input_ids: torch.Tensor, cache: ShadowCache) -> torch.Tensor:
+        """Prefill the full prompt; returns logits [batch, seq, vocab]."""
+        batch_size, seq_len = input_ids.shape
+        inputs_embeds = self._embed(input_ids)
+        position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
+        attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long, device=input_ids.device)
+
+        shadow_hidden = self._run_shadow(input_ids, position_ids, cache.shadow_cache)
+        hidden = self._run_base(inputs_embeds, position_ids, attention_mask, cache.base_cache, shadow_hidden)
+        return self.lm_head(hidden)
+
+    def step(self, input_ids: torch.Tensor, cache: ShadowCache, step_idx: int) -> torch.Tensor:
+        """Decode a single token; returns logits [batch, 1, vocab]."""
+        batch_size = input_ids.shape[0]
+        inputs_embeds = self._embed(input_ids)  # [batch, 1, hidden]
+        position_ids = torch.full((batch_size, 1), step_idx, dtype=torch.long, device=input_ids.device)
+        # Attention mask must cover all seen tokens (prefix + this token).
+        total_len = step_idx + 1
+        attention_mask = torch.ones(batch_size, total_len, dtype=torch.long, device=input_ids.device)
+
+        shadow_token = self._run_shadow(input_ids, position_ids, cache.shadow_cache)
+        hidden = self._run_base(inputs_embeds, position_ids, attention_mask, cache.base_cache, shadow_token)
+        return self.lm_head(hidden)
+
+
+# ---------------------------------------------------------------------------
+# Alignment verification
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class StepResult:
+    step: int
+    token_cached: int
+    token_full: int
+    maxdiff: float
+    logits_match: bool
+
+
+def greedy_next(logits: torch.Tensor) -> torch.Tensor:
+    """Return argmax token id from logits [batch, seq, vocab] or [batch, vocab]."""
+    return logits[:, -1, :].argmax(dim=-1)
+
+
+def verify_alignment(
+    peft: ShadowPeftModel,
+    input_ids: torch.Tensor,
+    max_new_tokens: int = 8,
+    *,
+    threshold: float = 1e-3,
+    label: str = "",
+) -> list[StepResult]:
+    """
+    Compare cached decode vs full-sequence recompute, step by step.
+
+    Path A (cached): prefill once, then single-token decode steps.
+    Path B (full):   recompute the entire sequence at each step (use_cache=False).
+    """
+    decoder = ShadowCachedDecoder(peft)
+    results: list[StepResult] = []
+
+    prefix_len = input_ids.shape[1]
+
+    if prefix_len + max_new_tokens > peft.base_model.config.max_position_embeddings:
+        max_new_tokens = peft.base_model.config.max_position_embeddings - prefix_len
+        print(f"  [warn] clamped max_new_tokens to {max_new_tokens} (max_position_embeddings)")
+
+    # === Path A: cached ===
+    cache = ShadowCache()
+    with torch.no_grad():
+        logits_prefill_cached = decoder.prefill(input_ids, cache)
+    # First generated token comes from prefill's last position.
+    cached_tokens = [greedy_next(logits_prefill_cached)]
+
+    # === Path B: full-seq (ground truth) ===
+    with torch.no_grad():
+        logits_full_prefix = peft(input_ids=input_ids).logits[:, -1, :]
+    full_tokens = [logits_full_prefix.argmax(dim=-1)]
+
+    # Check prefill alignment.
+    prefill_maxdiff = (logits_prefill_cached[:, -1, :] - logits_full_prefix).abs().max().item()
+    prefill_match = greedy_next(logits_prefill_cached) == logits_full_prefix.argmax(dim=-1)
+    results.append(
+        StepResult(
+            step=0,
+            token_cached=cached_tokens[0].item(),
+            token_full=full_tokens[0].item(),
+            maxdiff=prefill_maxdiff,
+            logits_match=bool(prefill_match.all().item()),
+        )
+    )
+
+    # Greedy decode loop.
+    for step in range(1, max_new_tokens):
+        pos = prefix_len + step - 1  # position of the token we're decoding
+        token_cached = cached_tokens[-1]
+        token_full = full_tokens[-1]
+
+        # Path A: cached single-token step.
+        with torch.no_grad():
+            logits_cached = decoder.step(token_cached.unsqueeze(-1), cache, step_idx=pos)
+        new_cached = greedy_next(logits_cached)
+        cached_tokens.append(new_cached)
+
+        # Path B: full-sequence recompute.
+        full_ids = torch.cat([input_ids, torch.stack(full_tokens, dim=1)], dim=1)
+        with torch.no_grad():
+            logits_full = peft(input_ids=full_ids).logits[:, -1, :]
+        new_full = logits_full.argmax(dim=-1)
+        full_tokens.append(new_full)
+
+        maxdiff = (logits_cached[:, -1, :] - logits_full).abs().max().item()
+        results.append(
+            StepResult(
+                step=step,
+                token_cached=new_cached.item(),
+                token_full=new_full.item(),
+                maxdiff=maxdiff,
+                logits_match=bool((new_cached == new_full).all().item()),
+            )
+        )
+
+    # Print results table.
+    header = f"{'step':>4}  {'cached':>7}  {'full':>7}  {'maxdiff':>12}  {'match':>5}"
+    print(f"\n  [{label}] Alignment results (threshold={threshold:.0e}):")
+    print(f"  {header}")
+    print(f"  {'-' * len(header)}")
+    all_match = True
+    for r in results:
+        ok = r.maxdiff < threshold and r.logits_match
+        all_match = all_match and ok
+        status = "OK" if ok else "FAIL"
+        print(f"  {r.step:>4}  {r.token_cached:>7}  {r.token_full:>7}  {r.maxdiff:>12.2e}  {status:>5}")
+
+    cached_seq = [input_ids[0].tolist()] + [t.item() for t in cached_tokens]
+    full_seq = [input_ids[0].tolist()] + [t.item() for t in full_tokens]
+    seq_match = cached_seq == full_seq
+    print(f"\n  Sequence match: {seq_match}")
+    if not seq_match:
+        print(f"    cached: {cached_seq}")
+        print(f"    full:   {full_seq}")
+
+    max_overall = max(r.maxdiff for r in results)
+    verdict = "PASS" if all_match and seq_match else "FAIL"
+    print(f"  Overall max diff: {max_overall:.2e}  ->  {verdict}")
+
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Test matrix
+# ---------------------------------------------------------------------------
+
+
+def build_llama_implicit() -> tuple[ShadowPeftModel, str]:
+    torch.manual_seed(42)
+    base = _tiny_llama(num_layers=4)
+    cfg = ShadowConfig(num_shadow_layers=1, injection_hidden_size=8, gate_hidden_size=10, alpha=0.1, dropout=0.0)
+    return get_shadow_model(base, cfg), "Llama-4L implicit shadow (hidden match)"
+
+
+def build_qwen2_implicit() -> tuple[ShadowPeftModel, str]:
+    torch.manual_seed(42)
+    base = _tiny_qwen2(num_layers=4)
+    cfg = ShadowConfig(num_shadow_layers=1, injection_hidden_size=8, gate_hidden_size=10, alpha=0.1, dropout=0.0)
+    return get_shadow_model(base, cfg), "Qwen2-4L implicit shadow (hidden match)"
+
+
+def build_llama_explicit_projection() -> tuple[ShadowPeftModel, str]:
+    """Explicit shadow model with DIFFERENT hidden size -> requires projection."""
+    torch.manual_seed(42)
+    base = _tiny_llama(num_layers=4, hidden_size=32)
+    # Shadow model with smaller hidden size (16) -> projection 16->32 kicks in.
+    # Keep shadow's own embed_tokens (16-dim) since base embeddings (32-dim) won't fit.
+    shadow_base = _tiny_llama(num_layers=1, hidden_size=16)
+    cfg = ShadowConfig(num_shadow_layers=1, injection_hidden_size=8, gate_hidden_size=10, alpha=0.1, dropout=0.0)
+    from shadow_peft import prepare_shadow_model
+
+    shadow_backbone = prepare_shadow_model(shadow_base, remove_embed_tokens=False)
+    peft = get_shadow_model(base, cfg, shadow_model=shadow_backbone)
+    return peft, "Llama-4L explicit shadow (hidden mismatch, projection 16->32)"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Verify ShadowPEFT cached decode alignment")
+    parser.add_argument(
+        "--models",
+        nargs="*",
+        default=["llama", "qwen2", "projection"],
+        choices=["llama", "qwen2", "projection"],
+        help="Which model configurations to test",
+    )
+    parser.add_argument("--max-new-tokens", type=int, default=8, help="Number of decode steps")
+    parser.add_argument("--threshold", type=float, default=1e-3, help="Max logits diff threshold")
+    parser.add_argument("--prompt-len", type=int, default=4, help="Prompt length")
+    args = parser.parse_args()
+
+    builders: dict[str, Any] = {
+        "llama": build_llama_implicit,
+        "qwen2": build_qwen2_implicit,
+        "projection": build_llama_explicit_projection,
+    }
+
+    print("=" * 70)
+    print("ShadowPEFT Cached Decode Alignment Verification")
+    print("=" * 70)
+    print(f"Config: prompt_len={args.prompt_len}, max_new_tokens={args.max_new_tokens}, threshold={args.threshold:.0e}")
+
+    all_pass = True
+    for key in args.models:
+        peft, label = builders[key]()
+        torch.manual_seed(123)  # deterministic input
+        vocab = peft.base_model.config.vocab_size
+        input_ids = torch.randint(0, vocab, (1, args.prompt_len))
+        print(f"\n{'─' * 70}")
+        print(f"Model: {label}")
+        print(f"Input: {input_ids[0].tolist()}")
+        results = verify_alignment(peft, input_ids, args.max_new_tokens, threshold=args.threshold, label=label)
+        max_diff = max(r.maxdiff for r in results)
+        passed = max_diff < args.threshold and all(r.logits_match for r in results)
+        all_pass = all_pass and passed
+
+    print(f"\n{'=' * 70}")
+    verdict = "ALL PASS" if all_pass else "SOME FAILED"
+    print(f"Final verdict: {verdict}")
+    print("=" * 70)
+    return 0 if all_pass else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/src/shadow_peft/peft_model.py b/src/shadow_peft/peft_model.py
index 200304d..0e2bae2 100644
--- a/src/shadow_peft/peft_model.py
+++ b/src/shadow_peft/peft_model.py
@@ -13,7 +13,7 @@
 from safetensors.torch import load_file as safetensors_load_file
 from safetensors.torch import save_file as safetensors_save_file
 from torch.nn.modules.module import _IncompatibleKeys
-from transformers import PreTrainedModel
+from transformers import DynamicCache, PreTrainedModel
 
 from .config import ShadowConfig
 from .model_utils import (
@@ -338,6 +338,12 @@ def __init__(
         # Internal mutable state during forward.
         self._shadow_hidden_states: torch.Tensor | None = None
 
+        # Shadow backbone KV-cache for incremental decode.
+        # Created lazily on the first cached prefill; reused for subsequent decode steps.
+        # Lifecycle is tied to the base model's past_key_values: when the caller starts a
+        # new generation (past_key_values=None), the shadow cache is rebuilt from scratch.
+        self._shadow_cache: DynamicCache | None = None
+
         # Wrap base layers in-place.
         wrapped = nn.ModuleList([])
         for i, layer in enumerate(base_layers):
@@ -378,8 +384,16 @@ def _compute_initial_shadow_hidden(
         attention_mask: torch.Tensor | None,
         position_ids: torch.Tensor | None,
         inputs_embeds: torch.Tensor | None,
+        shadow_cache: DynamicCache | None = None,
         **kwargs: Any,
     ) -> torch.Tensor:
+        """
+        Run the shadow backbone to produce per-token shadow hidden states.
+
+        When ``shadow_cache`` is provided, the shadow backbone runs in cached mode:
+        prefill (cache empty) or incremental decode (cache non-empty). When ``None``,
+        caching is disabled and the full sequence is recomputed (backward-compatible default).
+        """
         # We call the shadow *backbone* to get last_hidden_state (no logits needed).
         shadow_backbone = _get_backbone(self.shadow_model)
 
@@ -396,13 +410,19 @@ def _compute_initial_shadow_hidden(
                 raise AttributeError("Base model does not expose input embeddings.")
             inputs_embeds = base_embed(input_ids)
 
-        # Make a best-effort to disable caching, since Shadow updates require full sequences.
         kwargs = dict(kwargs)
-        kwargs["use_cache"] = False
-        kwargs["past_key_values"] = None
         kwargs["output_hidden_states"] = False
         kwargs["return_dict"] = True
 
+        if shadow_cache is not None:
+            # Cached mode: let the shadow backbone use/update its KV-cache.
+            kwargs["use_cache"] = True
+            kwargs["past_key_values"] = shadow_cache
+        else:
+            # Full-sequence mode (backward-compatible): disable caching.
+            kwargs["use_cache"] = False
+            kwargs["past_key_values"] = None
+
         # If we have `inputs_embeds`, use it. Otherwise, fall back to `input_ids` so an explicit
         # shadow model can use its own embedding table.
         if self._shadow_supports_inputs_embeds and inputs_embeds is not None:
@@ -483,23 +503,53 @@ def _configure_shadow_embedding_sharing(self) -> bool:
                 self._remove_embed_tokens(shadow_backbone)
         return supports_inputs_embeds
 
+    def _resolve_shadow_cache(self, use_cache: bool | None, past_key_values: Any) -> DynamicCache | None:
+        """
+        Decide whether to use a shadow backbone cache for this forward call.
+
+        Returns the shadow cache to use (creating one on the first cached prefill), or ``None``
+        for full-sequence uncached mode.
+        """
+        if not use_cache:
+            # No caching requested — reset internal shadow cache for a clean state next time.
+            self._shadow_cache = None
+            return None
+
+        # Cached mode. If the base has no cache yet, this is a fresh prefill → create shadow cache.
+        # If the base already has a cache, reuse the existing shadow cache for incremental decode.
+        if past_key_values is None or (hasattr(past_key_values, "get_seq_length") and past_key_values.get_seq_length() == 0):
+            shadow_backbone = _get_backbone(self.shadow_model)
+            shadow_cfg = getattr(shadow_backbone, "config", None)
+            try:
+                self._shadow_cache = DynamicCache(config=shadow_cfg)
+            except Exception:
+                # Fallback: some non-standard configs may not support config-based cache init.
+                self._shadow_cache = DynamicCache()
+        return self._shadow_cache
+
     def forward(self, *args, **kwargs):
         # Initialize per-forward shadow state.
         input_ids = kwargs.get("input_ids")
         attention_mask = kwargs.get("attention_mask")
         position_ids = kwargs.get("position_ids")
         inputs_embeds = kwargs.get("inputs_embeds")
+        use_cache = kwargs.get("use_cache")
+        past_key_values = kwargs.get("past_key_values")
+
+        shadow_cache = self._resolve_shadow_cache(use_cache, past_key_values)
 
         self._shadow_hidden_states = self._compute_initial_shadow_hidden(
             input_ids=input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,
             inputs_embeds=inputs_embeds,
+            shadow_cache=shadow_cache,
         )
 
-        # Force no-cache (same reason as raw implementation).
-        kwargs["use_cache"] = False
-        kwargs["past_key_values"] = None
+        if shadow_cache is None:
+            # Backward-compatible full-sequence mode.
+            kwargs["use_cache"] = False
+            kwargs["past_key_values"] = None
 
         try:
             return self.base_model(*args, **kwargs)
@@ -520,18 +570,24 @@ def forward_with_shadow(self, *args, **kwargs) -> tuple[object, torch.Tensor]:
         attention_mask = kwargs.get("attention_mask")
         position_ids = kwargs.get("position_ids")
         inputs_embeds = kwargs.get("inputs_embeds")
+        use_cache = kwargs.get("use_cache")
+        past_key_values = kwargs.get("past_key_values")
+
+        shadow_cache = self._resolve_shadow_cache(use_cache, past_key_values)
 
         initial_shadow_hidden = self._compute_initial_shadow_hidden(
             input_ids=input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,
             inputs_embeds=inputs_embeds,
+            shadow_cache=shadow_cache,
         )
         self._shadow_hidden_states = initial_shadow_hidden
 
-        # Force no-cache (same reason as raw implementation).
-        kwargs["use_cache"] = False
-        kwargs["past_key_values"] = None
+        if shadow_cache is None:
+            # Backward-compatible full-sequence mode.
+            kwargs["use_cache"] = False
+            kwargs["past_key_values"] = None
 
         try:
             outputs = self.base_model(*args, **kwargs)
diff --git a/src/shadow_peft/task_models.py b/src/shadow_peft/task_models.py
index 1a1c50e..6fd0e1f 100644
--- a/src/shadow_peft/task_models.py
+++ b/src/shadow_peft/task_models.py
@@ -143,14 +143,10 @@ def __init__(
 
         # Expose config/generation_config like a HF model.
         self.config = getattr(self.peft_model.base_model, "config", None)
-        if hasattr(self.config, "use_cache"):
-            self.config.use_cache = False
         base_gen_cfg = getattr(self.peft_model.base_model, "generation_config", None)
         self.generation_config = (
             base_gen_cfg if base_gen_cfg is not None else GenerationConfig.from_model_config(self.config)
         )
-        if hasattr(self.generation_config, "use_cache"):
-            self.generation_config.use_cache = False
 
         # Heads: use base model embeddings/head if available.
         self.lm_head = self.peft_model.base_model.get_output_embeddings()
@@ -311,15 +307,28 @@ def gradient_checkpointing_disable(self) -> None:
         if callable(fn):
             fn()
 
+    def get_experts_implementation(self) -> dict[str, str | None]:
+        """Delegate to base model (required by HF generate() in transformers >= 5.x)."""
+        fn = getattr(self.peft_model.base_model, "get_experts_implementation", None)
+        if callable(fn):
+            return fn()
+        return {}
+
+    def set_experts_implementation(self, experts_implementation: str | dict) -> None:
+        """Delegate to base model (required by HF generate() in transformers >= 5.x)."""
+        fn = getattr(self.peft_model.base_model, "set_experts_implementation", None)
+        if callable(fn):
+            fn(experts_implementation)
+
     def forward(
         self,
         *args,
         labels: torch.Tensor | None = None,
         **kwargs: Any,
     ) -> ShadowCausalLMOutputWithPast:
-        # Force disable caching (Shadow requires full-seq processing).
-        kwargs["use_cache"] = False
-        kwargs["past_key_values"] = None
+        # Default to no-cache for backward compatibility (training / full-seq inference).
+        kwargs.setdefault("use_cache", False)
+        kwargs.setdefault("past_key_values", None)
 
         # Normalize common HF kwargs so we never pass them twice (e.g. generation sets return_dict=True).
         if "labels" in kwargs:
@@ -374,37 +383,47 @@ def forward(
             loss=loss,
             logits=base_logits,
             shadow_logits=shadow_logits,
-            past_key_values=None,
+            past_key_values=getattr(outputs, "past_key_values", None),
             hidden_states=getattr(outputs, "hidden_states", None),
             attentions=getattr(outputs, "attentions", None),
         )
 
-    # GenerationMixin hooks (ensure full-sequence forward, no cache).
+    # GenerationMixin hooks.
     def prepare_inputs_for_generation(
         self,
         input_ids: torch.LongTensor,
         attention_mask: torch.Tensor | None = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
+        past_key_values = kwargs.get("past_key_values")
+        use_cache = kwargs.get("use_cache", False)
+
+        # When caching, slice input_ids to only the new tokens (prefill: full, decode: last token).
+        if use_cache and past_key_values is not None:
+            past_len = past_key_values.get_seq_length()
+            input_ids = input_ids[:, past_len:]
+
         position_ids = kwargs.get("position_ids")
         if attention_mask is not None and position_ids is None:
             position_ids = attention_mask.long().cumsum(-1) - 1
             position_ids.masked_fill_(attention_mask == 0, 1)
+            # In cached decode mode, only pass position_ids for the new tokens.
+            if use_cache and past_key_values is not None:
+                position_ids = position_ids[:, -(input_ids.shape[1]):]
 
         return {
             "input_ids": input_ids,
             "attention_mask": attention_mask,
             "position_ids": position_ids,
-            "use_cache": False,
-            "past_key_values": None,
+            "use_cache": use_cache,
+            "past_key_values": past_key_values,
         }
 
     def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder=False):
+        # Let HF's default implementation propagate past_key_values from outputs.
         model_kwargs = super()._update_model_kwargs_for_generation(
             outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder
         )
-        model_kwargs["past_key_values"] = None
-        model_kwargs["use_cache"] = False
         return model_kwargs
 
 
diff --git a/tests/test_cached_decode.py b/tests/test_cached_decode.py
new file mode 100644
index 0000000..6c7adfa
--- /dev/null
+++ b/tests/test_cached_decode.py
@@ -0,0 +1,233 @@
+"""
+Tests for ShadowPEFT KV-cache incremental decode.
+
+Verifies that cached generation (use_cache=True) produces identical results
+to full-sequence recompute (use_cache=False), proving the cache integration
+is mathematically correct.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+transformers = pytest.importorskip("transformers")
+
+from transformers import LlamaConfig, LlamaForCausalLM  # noqa: E402
+
+from shadow_peft import (  # noqa: E402
+    ShadowConfig,
+    ShadowForCausalLM,
+    get_shadow_model,
+)
+
+
+def _tiny_llama(vocab_size: int = 128, num_layers: int = 4) -> LlamaForCausalLM:
+    cfg = LlamaConfig(
+        vocab_size=vocab_size,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=num_layers,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        max_position_embeddings=64,
+    )
+    return LlamaForCausalLM(cfg)
+
+
+def _tiny_qwen2(num_layers: int = 4):
+    Qwen2Config = getattr(transformers, "Qwen2Config", None)
+    Qwen2ForCausalLM = getattr(transformers, "Qwen2ForCausalLM", None)
+    if Qwen2Config is None or Qwen2ForCausalLM is None:
+        pytest.skip("Qwen2 is not available in this transformers version")
+    cfg = Qwen2Config(
+        vocab_size=128,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=num_layers,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        max_position_embeddings=64,
+    )
+    return Qwen2ForCausalLM(cfg)
+
+
+def _build_causallm(base):
+    cfg = ShadowConfig(
+        num_shadow_layers=1,
+        injection_hidden_size=8,
+        gate_hidden_size=10,
+        alpha=0.1,
+        dropout=0.0,
+    )
+    peft = get_shadow_model(base, cfg)
+    return ShadowForCausalLM(peft, inference_mode="base_shadow")
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+def test_generate_cached_matches_nocache():
+    """generate(use_cache=True) and generate(use_cache=False) must produce identical tokens."""
+    torch.manual_seed(42)
+    base = _tiny_llama(num_layers=4)
+    m = _build_causallm(base)
+    m.eval()
+
+    input_ids = torch.randint(0, base.config.vocab_size, (1, 4))
+    with torch.no_grad():
+        gen_cached = m.generate(input_ids=input_ids, max_new_tokens=8, use_cache=True, do_sample=False)
+        m.peft_model._shadow_cache = None  # reset for the no-cache run
+        gen_nocache = m.generate(input_ids=input_ids, max_new_tokens=8, use_cache=False, do_sample=False)
+
+    assert gen_cached.tolist() == gen_nocache.tolist(), (
+        f"Cached {gen_cached.tolist()} != nocache {gen_nocache.tolist()}"
+    )
+
+
+def test_generate_cached_matches_nocache_qwen2():
+    """Same test for Qwen2 architecture."""
+    torch.manual_seed(42)
+    base = _tiny_qwen2(num_layers=4)
+    m = _build_causallm(base)
+    m.eval()
+
+    input_ids = torch.randint(0, base.config.vocab_size, (1, 4))
+    with torch.no_grad():
+        gen_cached = m.generate(input_ids=input_ids, max_new_tokens=8, use_cache=True, do_sample=False)
+        m.peft_model._shadow_cache = None
+        gen_nocache = m.generate(input_ids=input_ids, max_new_tokens=8, use_cache=False, do_sample=False)
+
+    assert gen_cached.tolist() == gen_nocache.tolist()
+
+
+def test_forward_prefill_then_decode():
+    """
+    Manually prefill with use_cache=True, then do a single decode step.
+    Compare with full-sequence forward on the extended input.
+    """
+    torch.manual_seed(42)
+    base = _tiny_llama(num_layers=4)
+    m = _build_causallm(base)
+    m.eval()
+
+    prefix = torch.randint(0, base.config.vocab_size, (1, 4))
+    next_token = torch.tensor([[77]])
+    full_ids = torch.cat([prefix, next_token], dim=1)
+
+    with torch.no_grad():
+        # Prefill
+        out_prefill = m(input_ids=prefix, use_cache=True)
+        base_cache = out_prefill.past_key_values
+        assert base_cache is not None, "Prefill should return a DynamicCache"
+        assert base_cache.get_seq_length() == 4
+
+        # Decode step: pass only the new token + cache
+        out_decode = m(input_ids=next_token, past_key_values=base_cache, use_cache=True)
+        logits_decode = out_decode.logits  # [1, 1, vocab]
+
+        # Ground truth: full-sequence forward
+        m.peft_model._shadow_cache = None  # reset shadow cache
+        out_full = m(input_ids=full_ids, use_cache=False)
+        logits_full = out_full.logits[:, -1:, :]  # [1, 1, vocab]
+
+    maxdiff = (logits_decode - logits_full).abs().max().item()
+    assert maxdiff < 1e-4, f"Cached decode logits maxdiff {maxdiff:.2e} exceeds threshold 1e-4"
+
+
+def test_step_by_step_logits_alignment():
+    """
+    Greedy decode token-by-token with cache, comparing logits at each step
+    against full-sequence recompute.
+
+    Uses two separate model instances (with identical seeds) so the cached path's
+    _shadow_cache is not corrupted by the full-seq ground-truth runs.
+    """
+    torch.manual_seed(42)
+    base = _tiny_llama(num_layers=4)
+    m_cached = _build_causallm(base)
+    m_cached.eval()
+
+    # Second instance for ground-truth full-seq runs (same seed → same weights).
+    torch.manual_seed(42)
+    base2 = _tiny_llama(num_layers=4)
+    m_full = _build_causallm(base2)
+    m_full.eval()
+
+    prefix = torch.randint(0, base.config.vocab_size, (1, 3))
+    max_new_tokens = 5
+
+    with torch.no_grad():
+        # Ground truth: full-sequence forward on prefix gives the logits for position len-1
+        out_full_prefix = m_full(input_ids=prefix, use_cache=False)
+        full_logits = out_full_prefix.logits[:, -1, :]
+
+        # Cached prefill
+        out = m_cached(input_ids=prefix, use_cache=True)
+        cache = out.past_key_values
+        cached_logits = out.logits[:, -1, :]
+
+        maxdiff = (cached_logits - full_logits).abs().max().item()
+        assert maxdiff < 1e-4, f"Prefill: cached vs full logits maxdiff {maxdiff:.2e}"
+
+        generated = []
+        for step in range(max_new_tokens):
+            # Pick next token from cached logits
+            next_tok = cached_logits.argmax(dim=-1, keepdim=True)
+            generated.append(next_tok.item())
+
+            # Full-seq ground truth for the extended sequence
+            full_ids = torch.cat([prefix] + [torch.tensor([[t]]) for t in generated], dim=1)
+            out_full = m_full(input_ids=full_ids, use_cache=False)
+            full_logits = out_full.logits[:, -1, :]
+
+            # Cached decode step
+            if step < max_new_tokens - 1:
+                out = m_cached(input_ids=next_tok, past_key_values=cache, use_cache=True)
+                cached_logits = out.logits[:, -1, :]
+                maxdiff = (cached_logits - full_logits).abs().max().item()
+                assert maxdiff < 1e-4, f"Step {step}: cached vs full logits maxdiff {maxdiff:.2e}"
+
+    # Verify the cached greedy sequence matches a cached generate() call.
+    # generate() may stop early on EOS; compare up to the shorter sequence.
+    cached_gen = m_cached.generate(input_ids=prefix, max_new_tokens=max_new_tokens, use_cache=True, do_sample=False)
+    gen_tokens = cached_gen[0, prefix.shape[1]:].tolist()
+    assert generated[:len(gen_tokens)] == gen_tokens, (
+        f"Manual decode {generated[:len(gen_tokens)]} != generate() {gen_tokens}"
+    )
+
+
+def test_backward_compat_nocache_default():
+    """
+    Default forward (use_cache not specified) must behave exactly as before:
+    no cache returned, same logits as use_cache=False.
+    """
+    torch.manual_seed(42)
+    base = _tiny_llama(num_layers=4)
+    m = _build_causallm(base)
+    m.eval()
+
+    input_ids = torch.randint(0, base.config.vocab_size, (1, 6))
+
+    with torch.no_grad():
+        out_default = m(input_ids=input_ids)
+        out_explicit = m(input_ids=input_ids, use_cache=False)
+
+    assert out_default.past_key_values is None, "Default forward should not return a cache"
+    assert torch.allclose(out_default.logits, out_explicit.logits, atol=0.0, rtol=0.0)
+
+
+def test_cached_forward_does_not_break_shadow_logits():
+    """Shadow logits must still be computed correctly in cached mode."""
+    torch.manual_seed(42)
+    base = _tiny_llama(num_layers=4)
+    m = _build_causallm(base)
+    m.eval()
+
+    input_ids = torch.randint(0, base.config.vocab_size, (1, 5))
+    with torch.no_grad():
+        out = m(input_ids=input_ids, use_cache=True)
+    assert out.shadow_logits is not None
+    assert out.shadow_logits.shape == out.logits.shape