From cf288d765e4faa9ab69b91c02d384cbee5714091 Mon Sep 17 00:00:00 2001
From: yidingcheng0206 <300775907+yidingcheng0206@users.noreply.github.com>
Date: Tue, 20 Oct 2026 20:24:00 +0800
Subject: [PATCH 1/4] fix(pytorch): enforce prefill checkpoint save alignment

---
 lmdeploy/cli/utils.py                         |   4 +-
 lmdeploy/messages.py                          |  10 +-
 .../paging/block_trie/checkpoint_lifecycle.py |  14 +-
 .../test_checkpoint_lifecycle.py              |   3 +-
 .../paging/test_prefill_save_alignment.py     | 210 ++++++++++++++++++
 tests/pytorch/paging/test_scheduler_ssm.py    |   2 +
 6 files changed, 234 insertions(+), 9 deletions(-)
 create mode 100644 tests/pytorch/paging/test_prefill_save_alignment.py

diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py
index c026811191..04bc4b4f6f 100644
--- a/lmdeploy/cli/utils.py
+++ b/lmdeploy/cli/utils.py
@@ -646,7 +646,9 @@ def prefill_save_alignment(parser):
                                    default=8192,
                                    help='Token alignment for linear-attention prefill state checkpoints, shared by '
                                    'local prefix caching and Mooncake Store. Must be a positive multiple of cache '
-                                   'block size; intermediate aligned positions may be skipped. '
+                                   'block size. Both internal and prompt-end checkpoints require an aligned absolute '
+                                   'prefix length and must pass model safety checks. Also guides prefill chunking; '
+                                   'aligned positions may be skipped. Decode checkpoint intervals are unchanged. '
                                    'Ignored by full-attention-only models. '
                                    'Only used by the PyTorch engine.')
 
diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py
index 3f22316ed7..70d37cd9d3 100644
--- a/lmdeploy/messages.py
+++ b/lmdeploy/messages.py
@@ -503,9 +503,13 @@ class PytorchEngineConfig:
             must be multiples of the cache block size.
         prefill_save_alignment: Token alignment for prefill state checkpoints,
             shared by local prefix caching and Mooncake Store. Defaults to
-            8192 and must be a positive multiple of block_size. Only applies
-            to models with linear-attention state caches; full-attention-only
-            models ignore it. Intermediate aligned positions may be skipped.
+            8192 and must be a positive multiple of block_size. Both internal
+            and prompt-end checkpoints must satisfy this alignment in addition
+            to model safety checks. Positions are absolute prefix lengths, not
+            chunk lengths. This also guides prefill chunking; aligned positions
+            may be skipped. Decode checkpoints use their separate interval.
+            Only applies to models with linear-attention state caches;
+            full-attention-only models ignore it.
         mooncake_state_save_slots: Number of temporary hybrid state snapshots
             reserved for asynchronous Mooncake saves. Defaults to 8. A full
             snapshot pool skips the save without blocking inference.
diff --git a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
index 124ccdbe34..9b042cfc49 100644
--- a/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
+++ b/lmdeploy/pytorch/paging/block_trie/checkpoint_lifecycle.py
@@ -103,13 +103,19 @@ def reserve_prefill_save_batch(
         seqs: list[SchedulerSequence],
         steps: tuple[int, ...] | None = None,
     ) -> CheckpointCopyPlan:
-        """Reserve prefill checkpoints and expose their logical copy pairs."""
+        """Reserve aligned prefill checkpoints and expose their copy pairs."""
         if steps is not None and len(steps) != len(seqs):
             raise ValueError('steps must have one entry per prefill sequence.')
         if steps is None:
-            state_offsets = [self.reserve_save(seq) for seq in seqs]
-        else:
-            state_offsets = [self.reserve_save(seq, step=step) for seq, step in zip(seqs, steps)]
+            steps = tuple(seq.num_valid_ids for seq in seqs)
+        state_offsets = []
+        for seq, step in zip(seqs, steps):
+            alignment = seq._seq_meta.prefill_save_alignment
+            if alignment and step % alignment:
+                self.discard_save(seq)
+                state_offsets.append(-1)
+                continue
+            state_offsets.append(self.reserve_save(seq, step=step))
 
         state_pairs = []
         kv_block_pairs = []
diff --git a/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py b/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
index 96feb64fb7..1a28bbd69b 100644
--- a/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
+++ b/tests/pytorch/paging/test_block_trie/test_checkpoint_lifecycle.py
@@ -519,7 +519,8 @@ def test_prepare_restore_batch_owns_partial_checkpoint_copy_pairs(self, ssm_sche
         assert not seq.prefix_cache.restore.is_selected
         assert checkpoint.pin_count == 0
 
-    def test_reserve_prefill_save_batch_owns_partial_checkpoint_copy_pairs(self, ssm_scheduler):
+    def test_reserve_prefill_save_batch_owns_partial_copy_pairs_without_alignment(self, ssm_scheduler):
+        ssm_scheduler.seq_meta.prefill_save_alignment = 0
         block_mgr = ssm_scheduler.block_manager
         block_trie = ssm_scheduler.block_trie
         state_manager = ssm_scheduler.state_manager
diff --git a/tests/pytorch/paging/test_prefill_save_alignment.py b/tests/pytorch/paging/test_prefill_save_alignment.py
new file mode 100644
index 0000000000..8e7826ae44
--- /dev/null
+++ b/tests/pytorch/paging/test_prefill_save_alignment.py
@@ -0,0 +1,210 @@
+import pytest
+import torch
+
+from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
+from lmdeploy.pytorch.long_context import get_long_context_chunk_limit, plan_long_context_chunk
+from lmdeploy.pytorch.messages import SequenceMeta
+from lmdeploy.pytorch.paging import Scheduler
+from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
+
+
+def _make_scheduler(alignment=512, budget=2048, lookahead=0):
+    cache_config = CacheConfig(max_batches=4,
+                               block_size=64,
+                               num_cpu_blocks=0,
+                               num_gpu_blocks=1024,
+                               max_prefill_token_num=budget,
+                               enable_prefix_caching=True,
+                               prefix_cache_block_aligned=True,
+                               prefill_save_alignment=alignment,
+                               prefix_cache_token_lookahead=lookahead,
+                               num_state_caches=21,
+                               prefix_cache_state_budget=16,
+                               states_shapes=[((1, ), torch.float32)])
+    scheduler_config = SchedulerConfig(max_batches=4,
+                                       max_session_len=65536,
+                                       max_request_output_len=128,
+                                       eviction_type='recompute')
+    return Scheduler(scheduler_config, cache_config, SequenceMeta(64, strategy=ARSequenceStrategy()))
+
+
+def _allocate_sequence(scheduler, length, session_id=0):
+    seq = scheduler.add_session(session_id).add_sequence([session_id + 1] * length)
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    scheduler.state_manager.allocate(seq)
+    return seq
+
+
+@pytest.mark.parametrize(('length', 'alignment', 'budget', 'expected_chunks', 'expected_saves'), [
+    (65, 64, 2048, [64, 1], [64]),
+    (65, 512, 2048, [65], []),
+    (128, 512, 2048, [128], []),
+    (129, 64, 2048, [128, 1], [128]),
+    (129, 512, 2048, [129], []),
+    (512, 512, 2048, [512], [512]),
+    (513, 512, 2048, [512, 1], [512]),
+    (576, 512, 2048, [512, 64], [512]),
+    (1024, 512, 2048, [512, 512], [512, 1024]),
+    (1024, 1024, 2048, [1024], [1024]),
+    (1088, 512, 2048, [1024, 64], [1024]),
+    (3000, 64, 2048, [2048, 896, 56], [2048, 2944]),
+    (3000, 512, 2048, [2048, 512, 440], [2048, 2560]),
+    (3000, 8192, 2048, [2048, 952], []),
+    (3000, 8192, 8192, [3000], []),
+    (4096, 512, 2048, [2048, 1536, 512], [2048, 3584, 4096]),
+    (16384, 8192, 16384, [8192, 8192], [8192, 16384]),
+    (16385, 8192, 32768, [16384, 1], [16384]),
+    (16384, 64, 16384, [16320, 64], [16320, 16384]),
+])
+def test_prefill_saves_only_at_configured_alignment(length, alignment, budget, expected_chunks, expected_saves):
+    scheduler = _make_scheduler(alignment, budget)
+    seq = scheduler.add_session(0).add_sequence(list(range(length)))
+    assert scheduler.schedule(is_prefill=True).running == [seq]
+    chunks = []
+    saved_steps = []
+    while seq.num_history_ids < length:
+        plan = plan_long_context_chunk(seq, get_long_context_chunk_limit(seq, budget))
+        assert 0 < plan.chunk_size <= budget
+        if chunks:
+            assert scheduler.reserve_long_context_chunk(seq, plan.chunk_size, is_last_chunk=plan.is_last_chunk)
+        else:
+            assert (seq.kv_token_limit or seq.num_valid_ids) == plan.chunk_end
+        copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch(
+            [seq], None if plan.is_last_chunk else (plan.chunk_end, ))
+        pending = seq.prefix_cache.pending_save
+        if plan.chunk_end % alignment == 0:
+            assert pending.is_pending
+            assert pending.step == plan.chunk_end
+            assert copy_plan.state_pairs == ((seq.logical_state, pending.slot), )
+            saved_steps.append(pending.step)
+            assert scheduler.state_checkpoints.publish_save(seq)
+        else:
+            assert not pending.is_pending
+            assert not copy_plan.state_pairs
+            assert not copy_plan.kv_block_pairs
+        chunks.append(plan.chunk_size)
+        seq.set_step(plan.chunk_end)
+    assert chunks == expected_chunks
+    assert saved_steps == expected_saves
+    published = sorted(node.state_checkpoint.step for node in scheduler.block_trie._checkpoint_index.unique_nodes()
+                       if node.state_checkpoint.published)
+    assert published == expected_saves
+
+
+@pytest.mark.parametrize('explicit_steps', [False, True])
+def test_prefill_alignment_filters_mixed_batch(explicit_steps):
+    scheduler = _make_scheduler()
+    seqs = [_allocate_sequence(scheduler, length, session_id)
+            for session_id, length in enumerate((128, 512, 576))]
+    steps = tuple(seq.num_valid_ids for seq in seqs) if explicit_steps else None
+    copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch(seqs, steps)
+    assert copy_plan.state_pairs == ((seqs[1].logical_state, seqs[1].prefix_cache.pending_save.slot), )
+    assert not copy_plan.kv_block_pairs
+    assert not seqs[0].prefix_cache.pending_save.is_pending
+    assert seqs[1].prefix_cache.pending_save.step == 512
+    assert not seqs[2].prefix_cache.pending_save.is_pending
+
+
+def test_prefill_alignment_rejection_does_not_allocate_or_evict(monkeypatch):
+    scheduler = _make_scheduler()
+    producer = _allocate_sequence(scheduler, 512)
+    scheduler.state_checkpoints.reserve_prefill_save_batch([producer])
+    node = producer.prefix_cache.pending_save.node
+    assert scheduler.state_checkpoints.publish_save(producer)
+    seq = _allocate_sequence(scheduler, 128, 1)
+    free_states = scheduler.state_manager.get_num_free_checkpoint()
+    free_blocks = scheduler.block_manager.get_num_free_gpu_blocks()
+
+    def unexpected_reservation(*args):
+        pytest.fail('An unaligned prefill must not try to allocate or evict a checkpoint.')
+
+    monkeypatch.setattr(scheduler.state_checkpoints, '_reserve_checkpoint', unexpected_reservation)
+    copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch([seq])
+    assert not copy_plan.state_pairs
+    assert not copy_plan.kv_block_pairs
+    assert node.state_checkpoint.published
+    assert scheduler.state_manager.get_num_free_checkpoint() == free_states
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == free_blocks
+
+
+def test_prefill_alignment_rejection_discards_old_unpublished_reservation():
+    scheduler = _make_scheduler()
+    seq = _allocate_sequence(scheduler, 1024)
+    free_states = scheduler.state_manager.get_num_free_checkpoint()
+    scheduler.state_checkpoints.reserve_prefill_save_batch([seq], (512, ))
+    node = seq.prefix_cache.pending_save.node
+    assert seq.prefix_cache.pending_save.is_pending
+    assert scheduler.state_manager.get_num_free_checkpoint() == free_states - 1
+
+    copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch([seq], (768, ))
+
+    assert not copy_plan.state_pairs
+    assert not copy_plan.kv_block_pairs
+    assert not seq.prefix_cache.pending_save.is_pending
+    assert node.state_checkpoint is None
+    assert scheduler.state_manager.get_num_free_checkpoint() == free_states
+
+
+def test_prefill_alignment_rejection_preserves_published_checkpoint():
+    scheduler = _make_scheduler()
+    seq = _allocate_sequence(scheduler, 1024)
+    scheduler.state_checkpoints.reserve_prefill_save_batch([seq], (512, ))
+    node = seq.prefix_cache.pending_save.node
+    scheduler.state_checkpoints.finish_forward_dispatch([seq], has_save_plan=True)
+    scheduler.state_checkpoints.unpin_saves([seq])
+    free_states = scheduler.state_manager.get_num_free_checkpoint()
+
+    copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch([seq], (768, ))
+
+    assert not copy_plan.state_pairs
+    assert not seq.prefix_cache.pending_save.is_pending
+    assert node.state_checkpoint.published
+    assert node.state_checkpoint.step == 512
+    assert scheduler.state_manager.get_num_free_checkpoint() == free_states
+
+
+@pytest.mark.parametrize(('length', 'history', 'saved'), [(1024, 768, True), (768, 512, False)])
+def test_prefill_alignment_uses_absolute_prefix_position(length, history, saved):
+    scheduler = _make_scheduler()
+    seq = _allocate_sequence(scheduler, length)
+    seq.set_step(history)
+    assert seq.num_token_ids == 256
+    copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch([seq])
+    assert bool(copy_plan.state_pairs) == saved
+    assert seq.prefix_cache.pending_save.is_pending == saved
+    if saved:
+        assert seq.prefix_cache.pending_save.step == length
+
+
+@pytest.mark.parametrize(('lookahead', 'length', 'saved'), [(0, 512, True), (1, 512, False), (1, 513, True)])
+def test_prefill_alignment_preserves_lookahead_guard(lookahead, length, saved):
+    scheduler = _make_scheduler(lookahead=lookahead)
+    seq = _allocate_sequence(scheduler, length)
+    copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch([seq], (512, ))
+    assert bool(copy_plan.state_pairs) == saved
+    assert seq.prefix_cache.pending_save.is_pending == saved
+
+
+def test_prefill_alignment_does_not_filter_decode_checkpoints():
+    scheduler = _make_scheduler()
+    seq = _allocate_sequence(scheduler, 768)
+    copy_plan = scheduler.state_checkpoints.reserve_decode_save_batch([seq], interval=64)
+    assert copy_plan.state_pairs == ((seq.logical_state, seq.prefix_cache.pending_save.slot), )
+    assert seq.prefix_cache.pending_save.step == 768
+    assert seq.prefix_cache.pending_save.is_decode
+
+
+@pytest.mark.parametrize('steps', [(), (512, 1024)])
+def test_prefill_alignment_preserves_steps_length_validation(steps):
+    scheduler = _make_scheduler()
+    seq = _allocate_sequence(scheduler, 512)
+    with pytest.raises(ValueError, match='one entry per prefill sequence'):
+        scheduler.state_checkpoints.reserve_prefill_save_batch([seq], steps)
+
+
+def test_prefill_alignment_preserves_empty_batch():
+    scheduler = _make_scheduler()
+    copy_plan = scheduler.state_checkpoints.reserve_prefill_save_batch([])
+    assert not copy_plan.state_pairs
+    assert not copy_plan.kv_block_pairs
diff --git a/tests/pytorch/paging/test_scheduler_ssm.py b/tests/pytorch/paging/test_scheduler_ssm.py
index 796565e682..f1e0d674f4 100644
--- a/tests/pytorch/paging/test_scheduler_ssm.py
+++ b/tests/pytorch/paging/test_scheduler_ssm.py
@@ -110,6 +110,8 @@ def test_ssm_runtime_state_waits_when_only_checkpoint_slot_is_pinned():
 def test_ssm_same_batch_duplicate_checkpoint_save_has_unique_dst_offsets():
     scheduler = _make_ssm_scheduler(max_batch_size=2, prefix_cache_state_budget=2)
     block_size = scheduler.seq_meta.block_size
+    scheduler.cache_config.prefill_save_alignment = block_size * 2
+    scheduler.seq_meta.prefill_save_alignment = block_size * 2
     token_ids = [1] * block_size * 2
 
     seq_a = scheduler.add_session(100).add_sequence(token_ids)

From ebba3bdc91884b39d34e2cdd30c9ef02e975644e Mon Sep 17 00:00:00 2001
From: yidingcheng0206 <300775907+yidingcheng0206@users.noreply.github.com>
Date: Tue, 20 Oct 2026 20:26:34 +0800
Subject: [PATCH 2/4] feat(pytorch): add optional native FlashKDA prefill
 backend

---
 docs/en/advance/internal_checkpoint.md        | 20 +++++
 lmdeploy/pytorch/backends/cuda/flash_kda.py   | 38 ++++++++++
 lmdeploy/pytorch/backends/cuda/kda.py         | 10 +++
 lmdeploy/pytorch/envs.py                      |  2 +
 tests/pytorch/kernel/test_native_flash_kda.py | 76 +++++++++++++++++++
 5 files changed, 146 insertions(+)
 create mode 100644 docs/en/advance/internal_checkpoint.md
 create mode 100644 lmdeploy/pytorch/backends/cuda/flash_kda.py
 create mode 100644 tests/pytorch/kernel/test_native_flash_kda.py

diff --git a/docs/en/advance/internal_checkpoint.md b/docs/en/advance/internal_checkpoint.md
new file mode 100644
index 0000000000..3ff9481ffb
--- /dev/null
+++ b/docs/en/advance/internal_checkpoint.md
@@ -0,0 +1,20 @@
+# GLM-5.3 native KDA prefill backend
+
+The default `LMDEPLOY_KDA_PREFILL_BACKEND=fla` keeps the existing FLA entry
+point. FLA may dispatch to its own FlashKDA backend when its conditions are met.
+Set `LMDEPLOY_KDA_PREFILL_BACKEND=flash_kda` to explicitly select LMDeploy's
+native FlashKDA adapter instead. This does not change the decode backend.
+
+The native adapter requires BF16 inputs, head dimension 128, FP32 recurrent
+state, and a gate lower bound. The installed FlashKDA operator must expose
+`checkpoint_state` and `checkpoint_offsets`; incompatible installations fail
+the capability check instead of silently switching backends.
+
+The validated FlashKDA source revision is
+`ee0be888cd0e972f9409bf53756f8c38c6652173`, with Torch 2.10.0+cu128 and
+flash-linear-attention 0.5.2 on H200. No FlashKDA, FLA, convolution, KPool, or
+TileLang decode kernel source is modified by this integration.
+
+For controlled comparisons, select the same prefill backend in both arms.
+Use `FLA_FLASH_KDA=0` when an explicit FLA Triton reference is required.
+Selecting the native adapter does not by itself enable internal checkpoints.
diff --git a/lmdeploy/pytorch/backends/cuda/flash_kda.py b/lmdeploy/pytorch/backends/cuda/flash_kda.py
new file mode 100644
index 0000000000..bfded46035
--- /dev/null
+++ b/lmdeploy/pytorch/backends/cuda/flash_kda.py
@@ -0,0 +1,38 @@
+"""Optional native FlashKDA adapter with explicit intermediate-state outputs."""
+
+import torch
+
+
+def check_flash_kda():
+    """Fail before inference if the installed native operator lacks IC outputs."""
+    import flash_kda
+
+    arguments = {argument.name for argument in torch.ops.flash_kda.fwd.default._schema.arguments}
+    if not {'checkpoint_state', 'checkpoint_offsets'} <= arguments:
+        raise RuntimeError('FlashKDA must provide native checkpoint_state and checkpoint_offsets outputs.')
+    return flash_kda
+
+
+def flash_kda_prefill(query, key, value, gate, beta, a_log, dt_bias,
+                      initial_state, cu_seqlens, lower_bound, checkpoint=None):
+    """Use the same native prefill path with or without checkpoint export."""
+    import flash_kda
+
+    if query.dtype != torch.bfloat16 or query.size(-1) != 128 or initial_state.dtype != torch.float32:
+        raise ValueError('Native FlashKDA prefill requires BF16 inputs, D=128 and FP32 recurrent state.')
+    if lower_bound is None:
+        raise ValueError('Native FlashKDA requires a finite gate lower bound.')
+    output = torch.empty_like(query, memory_format=torch.contiguous_format)
+    final_state = torch.empty_like(initial_state, memory_format=torch.contiguous_format)
+    workspace = torch.empty(
+        flash_kda.get_workspace_size(query.size(1), query.size(2), initial_state.size(0)),
+        dtype=torch.uint8, device=query.device)
+    torch.ops.flash_kda.fwd(
+        query.contiguous(), key.contiguous(), value.contiguous(), gate.contiguous(), beta.contiguous(),
+        query.size(-1) ** -0.5, output, workspace, a_log,
+        dt_bias.reshape(query.size(2), query.size(3)), lower_bound,
+        initial_state=initial_state.contiguous(), final_state=final_state, cu_seqlens=cu_seqlens,
+        checkpoint_state=None if checkpoint is None else checkpoint.recurrent_state_out,
+        checkpoint_offsets=None if checkpoint is None else checkpoint.offsets,
+    )
+    return output, final_state
diff --git a/lmdeploy/pytorch/backends/cuda/kda.py b/lmdeploy/pytorch/backends/cuda/kda.py
index c74c76bf1b..b43e5d32c5 100644
--- a/lmdeploy/pytorch/backends/cuda/kda.py
+++ b/lmdeploy/pytorch/backends/cuda/kda.py
@@ -12,6 +12,7 @@
 
 import torch
 
+from lmdeploy.pytorch import envs
 from lmdeploy.pytorch.backends.kda import KdaImpl
 
 from .gated_delta_rule import GatedDeltaStepMetaUpdater, _state_scatter, _state_select
@@ -52,6 +53,11 @@ def __init__(self):
         self.causal_conv1d_fwd = causal_conv1d_fwd
         self.causal_conv1d_update = causal_conv1d_update
         self.chunk_kda = chunk_kda
+        self.flash_prefill = None
+        if envs.kda_prefill_backend == 'flash_kda':
+            from .flash_kda import check_flash_kda, flash_kda_prefill
+            check_flash_kda()
+            self.flash_prefill = flash_kda_prefill
         from lmdeploy.pytorch.kernels.cuda.causal_conv1d import causal_conv1d_update as shared_conv_update
         from lmdeploy.pytorch.kernels.cuda.gated_delta_rule import fused_recurrent_gated_delta_rule
         self.shared_conv_update = shared_conv_update
@@ -292,6 +298,10 @@ def decode_view(x: torch.Tensor) -> torch.Tensor:
                 state_v_first=True,
             )
             output = output.flatten(0, 1).unsqueeze(0)
+        elif self.flash_prefill is not None:
+            output, final_state = self.flash_prefill(
+                q, k, v, raw_gate, raw_beta, a_log, dt_bias, selected_state,
+                metadata.cu_seqlens, lower_bound)
         else:
             output, final_state = self.chunk_kda(
                 q=q.contiguous(),
diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py
index e8cc6eb11d..54424b172c 100644
--- a/lmdeploy/pytorch/envs.py
+++ b/lmdeploy/pytorch/envs.py
@@ -231,6 +231,8 @@ def _patched_get_env(
     # qwen3.5 recurrent_state dtype
     fp32_mamba_ssm_dtype = env_to_bool('LMDEPLOY_FP32_MAMBA_SSM_DTYPE', False)
 
+    kda_prefill_backend = env_to_choice('LMDEPLOY_KDA_PREFILL_BACKEND', 'fla', {'fla', 'flash_kda'})
+
     # DSA indexer fusion
     disable_dsa_indexer_fusion = env_to_bool('LMDEPLOY_DISABLE_DSA_INDEXER_FUSION', False)
 
diff --git a/tests/pytorch/kernel/test_native_flash_kda.py b/tests/pytorch/kernel/test_native_flash_kda.py
new file mode 100644
index 0000000000..1d4818d841
--- /dev/null
+++ b/tests/pytorch/kernel/test_native_flash_kda.py
@@ -0,0 +1,76 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from lmdeploy.pytorch.backends.cuda.flash_kda import check_flash_kda, flash_kda_prefill
+
+
+@pytest.mark.parametrize('arguments', [
+    ['checkpoint_state', 'checkpoint_offsets'],
+    ['checkpoint_state'],
+    ['checkpoint_offsets'],
+])
+def test_native_flash_kda_capability_check(monkeypatch, arguments):
+    module = SimpleNamespace()
+    monkeypatch.setitem(sys.modules, 'flash_kda', module)
+    schema = SimpleNamespace(arguments=[SimpleNamespace(name=name) for name in arguments])
+    monkeypatch.setattr(torch.ops, 'flash_kda',
+                        SimpleNamespace(fwd=SimpleNamespace(default=SimpleNamespace(_schema=schema))))
+    if len(arguments) == 2:
+        assert check_flash_kda() is module
+    else:
+        with pytest.raises(RuntimeError, match='checkpoint_state and checkpoint_offsets'):
+            check_flash_kda()
+
+
+def make_inputs():
+    query = torch.zeros((1, 64, 2, 128), dtype=torch.bfloat16)
+    return dict(query=query, key=query.clone(), value=query.clone(), gate=query.clone(),
+                beta=torch.zeros((1, 64, 2), dtype=torch.bfloat16),
+                a_log=torch.zeros(2), dt_bias=torch.zeros(256),
+                initial_state=torch.zeros((1, 2, 128, 128)),
+                cu_seqlens=torch.tensor([0, 64], dtype=torch.int64), lower_bound=-5.0)
+
+
+@pytest.mark.parametrize('capture', [False, True])
+def test_native_flash_kda_forwards_optional_checkpoint(monkeypatch, capture):
+    recorded = {}
+
+    def forward(*args, **kwargs):
+        recorded.update(kwargs)
+        recorded['bias_shape'] = args[9].shape
+        args[6].zero_()
+        kwargs['final_state'].fill_(2)
+
+    monkeypatch.setitem(sys.modules, 'flash_kda', SimpleNamespace(get_workspace_size=lambda *args: 64))
+    monkeypatch.setattr(torch.ops, 'flash_kda', SimpleNamespace(fwd=forward))
+    inputs = make_inputs()
+    checkpoint = (SimpleNamespace(recurrent_state_out=torch.empty_like(inputs['initial_state']),
+                                  offsets=torch.tensor([32], dtype=torch.int64)) if capture else None)
+    output, final_state = flash_kda_prefill(**inputs, checkpoint=checkpoint)
+    assert output.shape == inputs['query'].shape
+    assert torch.count_nonzero(output) == 0
+    assert torch.all(final_state == 2)
+    assert recorded['bias_shape'] == (2, 128)
+    assert recorded['checkpoint_state'] is (checkpoint.recurrent_state_out if capture else None)
+    assert recorded['checkpoint_offsets'] is (checkpoint.offsets if capture else None)
+    assert recorded['cu_seqlens'] is inputs['cu_seqlens']
+
+
+@pytest.mark.parametrize('field', ['query', 'initial_state'])
+def test_native_flash_kda_rejects_incompatible_dtype(monkeypatch, field):
+    monkeypatch.setitem(sys.modules, 'flash_kda', SimpleNamespace())
+    inputs = make_inputs()
+    inputs[field] = inputs[field].to(torch.float16)
+    with pytest.raises(ValueError, match='BF16 inputs'):
+        flash_kda_prefill(**inputs)
+
+
+def test_native_flash_kda_requires_gate_bound(monkeypatch):
+    monkeypatch.setitem(sys.modules, 'flash_kda', SimpleNamespace())
+    inputs = make_inputs()
+    inputs['lower_bound'] = None
+    with pytest.raises(ValueError, match='gate lower bound'):
+        flash_kda_prefill(**inputs)

From a639759b5b9b1c75bd8073b15491ca988d2dee2a Mon Sep 17 00:00:00 2001
From: yidingcheng0206 <300775907+yidingcheng0206@users.noreply.github.com>
Date: Tue, 20 Oct 2026 20:30:05 +0800
Subject: [PATCH 3/4] feat(pytorch): support GLM internal prefill checkpoints

---
 docs/en/advance/internal_checkpoint.md        |  36 ++++
 lmdeploy/pytorch/backends/cuda/kda.py         |  20 +-
 lmdeploy/pytorch/backends/kda.py              |  11 +
 lmdeploy/pytorch/config.py                    |   1 +
 lmdeploy/pytorch/configurations/glm5_next.py  |  11 +
 lmdeploy/pytorch/engine/cache_inputs.py       |  23 ++
 lmdeploy/pytorch/engine/engine_loop.py        |   3 +-
 lmdeploy/pytorch/engine/inputs_maker.py       |  72 ++++++-
 lmdeploy/pytorch/engine/model_agent/agent.py  |   2 +
 lmdeploy/pytorch/envs.py                      |   1 +
 lmdeploy/pytorch/model_inputs.py              |   2 +
 lmdeploy/pytorch/models/glm5_next.py          |  43 +++-
 lmdeploy/pytorch/nn/kda.py                    |   4 +-
 tests/pytorch/engine/test_inputs_maker.py     |   8 +-
 .../test_internal_prefill_checkpoint.py       | 199 ++++++++++++++++++
 .../kernel/test_kda_internal_checkpoint.py    |  66 ++++++
 16 files changed, 482 insertions(+), 20 deletions(-)
 create mode 100644 tests/pytorch/engine/test_internal_prefill_checkpoint.py
 create mode 100644 tests/pytorch/kernel/test_kda_internal_checkpoint.py

diff --git a/docs/en/advance/internal_checkpoint.md b/docs/en/advance/internal_checkpoint.md
index 3ff9481ffb..9951702ae6 100644
--- a/docs/en/advance/internal_checkpoint.md
+++ b/docs/en/advance/internal_checkpoint.md
@@ -18,3 +18,39 @@ TileLang decode kernel source is modified by this integration.
 For controlled comparisons, select the same prefill backend in both arms.
 Use `FLA_FLASH_KDA=0` when an explicit FLA Triton reference is required.
 Selecting the native adapter does not by itself enable internal checkpoints.
+
+## Internal prefill checkpoints
+
+Set `LMDEPLOY_INTERNAL_PREFILL_CHECKPOINT=1` together with the native backend
+and local prefix caching to enable IC. IC is disabled by default.
+
+The scheduler first reserves the same safe checkpoint H as the original
+split-prefill path. If the remaining prompt fits the existing prefill budget,
+it can expand that forward to the prompt end E while retaining its reservation
+at H. FlashKDA exports the recurrent state directly into that checkpoint slot;
+the backend copies the preceding convolution-input window into its conv state.
+The normal runtime state still advances to E. An endpoint copy must not
+overwrite the internal checkpoint.
+
+The capture description contains the input start/end, checkpoint position, and
+source/destination slots. Public engine components transport this description
+and retain existing reservation, pinning, publication, and cleanup mechanisms.
+GLM-specific code maps slots to per-layer KDA output views. KPool runtime tails
+remain excluded from checkpoint copies, as defined by PR #5037.
+
+Initial support is non-MTP, text-only, exclusive single-sequence prefill with
+one internal save. A non-MTP prompt end that also requires an aligned checkpoint
+keeps the original split rather than dropping either save. Connector/Mooncake,
+multimodal, scoring, mixed local decode/prefill, insufficient capacity, and
+unsupported capture layouts do not enter the IC path. Invalid enabled-backend
+configurations fail validation; per-request ineligibility preserves the
+ordinary split path.
+
+Merged intervals within `piecewise_cudagraph_max_tokens` retain the original
+path. Current IC runs outside prefill PCG; it does not disable decode graphs.
+`prefill_save_alignment` constrains absolute checkpoint positions, including
+prompt-end saves, but does not guarantee a save at every aligned position.
+
+Comparisons must hold the backend, prefill budget, PCG cap, alignment, and
+request set fixed. Export/copy work and a larger single forward can consume
+time and memory; IC is not a zero-cost operation or a universal speedup.
diff --git a/lmdeploy/pytorch/backends/cuda/kda.py b/lmdeploy/pytorch/backends/cuda/kda.py
index b43e5d32c5..f0acfa06a1 100644
--- a/lmdeploy/pytorch/backends/cuda/kda.py
+++ b/lmdeploy/pytorch/backends/cuda/kda.py
@@ -13,7 +13,7 @@
 import torch
 
 from lmdeploy.pytorch import envs
-from lmdeploy.pytorch.backends.kda import KdaImpl
+from lmdeploy.pytorch.backends.kda import KdaImpl, KdaPrefillCheckpoint
 
 from .gated_delta_rule import GatedDeltaStepMetaUpdater, _state_scatter, _state_select
 from .step_metadata import register_step_metadata_impl
@@ -251,13 +251,29 @@ def forward(
         num_heads: int,
         head_dim: int,
         lower_bound: float,
+        checkpoint: KdaPrefillCheckpoint | None = None,
     ) -> torch.Tensor:
         if (getattr(metadata, 'num_spec_tokens', 0) or 0) > 0:
+            if checkpoint is not None:
+                raise ValueError('MTP internal checkpoints are not supported.')
             return self._forward_spec(
                 mixed_qkv, raw_gate, raw_beta, conv_state, recurrent_state,
                 metadata, conv_weight=conv_weight, conv_bias=conv_bias,
                 a_log=a_log, dt_bias=dt_bias, num_heads=num_heads,
                 head_dim=head_dim, lower_bound=lower_bound)
+        if checkpoint is not None:
+            if (self.flash_prefill is None or metadata.is_decoding
+                    or metadata.state_ids.numel() != 1):
+                raise ValueError('KDA IC requires native FlashKDA and a single prefill sequence.')
+            width = conv_weight.size(-1)
+            if not width <= checkpoint.offset < mixed_qkv.size(1) or checkpoint.offset % 16:
+                raise ValueError('KDA IC offset must contain the conv window and align to 16 tokens.')
+            if (checkpoint.offsets.shape != (1,) or checkpoint.offsets.dtype != metadata.cu_seqlens.dtype
+                    or checkpoint.recurrent_state_out.dtype != torch.float32
+                    or not checkpoint.recurrent_state_out.is_contiguous()):
+                raise ValueError('Invalid native KDA checkpoint output layout.')
+            window = mixed_qkv[:, checkpoint.offset - width:checkpoint.offset].transpose(1, 2)
+            checkpoint.conv_state_out.copy_(window)
         batch_size = metadata.state_ids.numel()
         if metadata.is_decoding:
             query_length = mixed_qkv.size(1) // batch_size
@@ -301,7 +317,7 @@ def decode_view(x: torch.Tensor) -> torch.Tensor:
         elif self.flash_prefill is not None:
             output, final_state = self.flash_prefill(
                 q, k, v, raw_gate, raw_beta, a_log, dt_bias, selected_state,
-                metadata.cu_seqlens, lower_bound)
+                metadata.cu_seqlens, lower_bound, checkpoint)
         else:
             output, final_state = self.chunk_kda(
                 q=q.contiguous(),
diff --git a/lmdeploy/pytorch/backends/kda.py b/lmdeploy/pytorch/backends/kda.py
index 783d8b8b0a..337e02c978 100644
--- a/lmdeploy/pytorch/backends/kda.py
+++ b/lmdeploy/pytorch/backends/kda.py
@@ -8,6 +8,16 @@
 from .base import BuildSpec
 
 
+@dataclass(frozen=True)
+class KdaPrefillCheckpoint:
+    """Single-sequence prefill outputs, separate from the live state."""
+
+    offset: int
+    offsets: torch.Tensor
+    conv_state_out: torch.Tensor
+    recurrent_state_out: torch.Tensor
+
+
 class KdaImpl(ABC):
     """Backend interface for Kimi Delta Attention inference."""
 
@@ -27,6 +37,7 @@ def forward(
         num_heads: int,
         head_dim: int,
         lower_bound: float,
+        checkpoint: KdaPrefillCheckpoint | None = None,
     ) -> torch.Tensor:
         """Run short convolution followed by the KDA recurrence."""
         raise NotImplementedError
diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py
index 75b2540b12..65326b6aa2 100644
--- a/lmdeploy/pytorch/config.py
+++ b/lmdeploy/pytorch/config.py
@@ -172,6 +172,7 @@ class CacheConfig:
     prefix_cache_state_budget: int = 0
     prefix_cache_decode_state_interval: int = 0
     prefill_save_alignment: int = 8192
+    enable_internal_prefill_checkpoint: bool = False
     mooncake_state_save_slots: int = 8
     # Model-owned checkpoint validity constraint, independent of prefill chunking.
     prefix_cache_block_aligned: bool = False
diff --git a/lmdeploy/pytorch/configurations/glm5_next.py b/lmdeploy/pytorch/configurations/glm5_next.py
index e85c81761c..a479a021a4 100644
--- a/lmdeploy/pytorch/configurations/glm5_next.py
+++ b/lmdeploy/pytorch/configurations/glm5_next.py
@@ -5,6 +5,7 @@
 
 import torch
 
+from lmdeploy.pytorch import envs
 from lmdeploy.pytorch.config import StateCacheSpec
 from lmdeploy.pytorch.consts import (
     GLM5_KDA_CONV_STATE,
@@ -43,6 +44,16 @@ def update_cache_config(cache_config, token_lookahead=0):
         raise ValueError('GLM-5.3 MTP prefix caching only publishes prefill checkpoints.')
     if cache_config.prefix_cache_decode_state_interval % 64:
         raise ValueError('GLM-5.3 prefix_cache_decode_state_interval must be a multiple of 64.')
+    if envs.internal_prefill_checkpoint:
+        if (token_lookahead != 0 or cache_config.kv_transfer_config is not None
+                or not cache_config.enable_prefix_caching or cache_config.device_type != 'cuda'
+                or cache_config.prefill_save_alignment % 64
+                or envs.kda_prefill_backend != 'flash_kda'):
+            raise ValueError('GLM IC requires CUDA, local prefix caching, non-MTP AR without a connector, '
+                             '64-aligned prefill_save_alignment and the native FlashKDA backend.')
+        from lmdeploy.pytorch.backends.cuda.flash_kda import check_flash_kda
+        check_flash_kda()
+        cache_config.enable_internal_prefill_checkpoint = True
 
 
 def is_glm5_kda_layer(text_config, layer_idx: int) -> bool:
diff --git a/lmdeploy/pytorch/engine/cache_inputs.py b/lmdeploy/pytorch/engine/cache_inputs.py
index 7cf94b01d8..45038a362d 100644
--- a/lmdeploy/pytorch/engine/cache_inputs.py
+++ b/lmdeploy/pytorch/engine/cache_inputs.py
@@ -7,6 +7,23 @@
 StateCacheCopyPlan = tuple[tuple[int, ...], tuple[int, ...]]
 
 
+@dataclass(frozen=True)
+class PrefillStateCapture:
+    """One internal checkpoint produced during an exclusive prefill."""
+
+    input_start: int
+    input_end: int
+    checkpoint_step: int
+    source_slot: int
+    target_slot: int
+
+    def __post_init__(self):
+        if not 0 <= self.input_start < self.checkpoint_step < self.input_end:
+            raise ValueError('Internal checkpoint must lie inside the input interval.')
+        if min(self.source_slot, self.target_slot) < 0 or self.source_slot == self.target_slot:
+            raise ValueError('Internal checkpoint requires distinct allocated state slots.')
+
+
 @dataclass(frozen=True)
 class CacheCheckpointInputs:
     """Cache-copy plans consumed by exactly one model forward.
@@ -25,6 +42,11 @@ class CacheCheckpointInputs:
     kv_save_plan: torch.Tensor | None = None
     state_restore_plan: StateCacheCopyPlan | None = None
     state_save_plan: StateCacheCopyPlan | None = None
+    state_capture: PrefillStateCapture | None = None
+
+    def __post_init__(self):
+        if self.state_capture is not None and (self.state_save_plan is not None or self.kv_save_plan is not None):
+            raise ValueError('Internal capture cannot share a forward with an endpoint checkpoint save.')
 
     @torch.inference_mode()
     def to_device(self, device: str, non_blocking: bool = False):
@@ -41,6 +63,7 @@ def to_device(self, device: str, non_blocking: bool = False):
             kv_save_plan=kv_save_plan,
             state_restore_plan=self.state_restore_plan,
             state_save_plan=self.state_save_plan,
+            state_capture=self.state_capture,
         )
 
     def record_stream(self, stream: torch.cuda.Stream) -> None:
diff --git a/lmdeploy/pytorch/engine/engine_loop.py b/lmdeploy/pytorch/engine/engine_loop.py
index dcd5027093..0e516e4ac0 100644
--- a/lmdeploy/pytorch/engine/engine_loop.py
+++ b/lmdeploy/pytorch/engine/engine_loop.py
@@ -496,7 +496,8 @@ async def _main_loop_get_outputs(
         if has_model_work:
             self.inputs_maker.update_running_seqs(running, model_inputs)
         has_state_checkpoint_save = (cache_inputs is not None
-                                     and cache_inputs.state_save_plan is not None)
+                                     and (cache_inputs.state_save_plan is not None
+                                          or cache_inputs.state_capture is not None))
 
         # ModelAgent executes queued forwards in send order.  Once the current
         # input is queued, matched checkpoints can be published before waiting
diff --git a/lmdeploy/pytorch/engine/inputs_maker.py b/lmdeploy/pytorch/engine/inputs_maker.py
index 47df448506..9dee4c91ce 100644
--- a/lmdeploy/pytorch/engine/inputs_maker.py
+++ b/lmdeploy/pytorch/engine/inputs_maker.py
@@ -30,7 +30,7 @@
 from lmdeploy.pytorch.model_inputs import ModelInputs, ModelInputsDelta, VisionModelInputs
 from lmdeploy.utils import get_logger
 
-from .cache_inputs import CacheCheckpointInputs, StateCacheCopyPlan
+from .cache_inputs import CacheCheckpointInputs, PrefillStateCapture, StateCacheCopyPlan
 
 if TYPE_CHECKING:
     from lmdeploy.pytorch.adapter.adapter import AdapterManager
@@ -135,6 +135,8 @@ class InputsMakerConfig:
     enable_chunked_prefill: bool = False
     use_mrope: bool = False
     prefill_interval: int = 16
+    enable_internal_prefill_checkpoint: bool = False
+    prefill_graph_max_tokens: int = 0
 
     @staticmethod
     def from_engine(engine: 'Engine'):
@@ -142,6 +144,8 @@ def from_engine(engine: 'Engine'):
         model_config = engine.model_config
         prefill_interval = engine.engine_config.prefill_interval
         kwargs = dict()
+        if _envs.internal_prefill_checkpoint and not cache_config.enable_internal_prefill_checkpoint:
+            raise ValueError('This model does not support internal prefill checkpoints.')
         if prefill_interval is not None:
             if not isinstance(prefill_interval, int) or prefill_interval <= 0:
                 raise ValueError('engine.engine_config.prefill_interval must be a positive int '
@@ -161,6 +165,8 @@ def from_engine(engine: 'Engine'):
             dp=engine.dist_config.dp,
             enable_chunked_prefill=engine.misc_config.enable_chunked_prefill,
             use_mrope=model_config.use_mrope,
+            enable_internal_prefill_checkpoint=cache_config.enable_internal_prefill_checkpoint,
+            prefill_graph_max_tokens=(engine.engine_config.piecewise_cudagraph_max_tokens or 0),
             **kwargs,
         )
 
@@ -610,6 +616,7 @@ def _schedule_prefill(self,
                  result.cache_inputs) = self._build_chunk_inputs(running, chunk_size, multimodals)
                 result.inputs.is_first_chunk = True
                 result.inputs.is_chunk_multimodal = maker.long_context_chunker.has_multimodal
+                self._try_internal_prefill(result)
                 maker._short_prefill_turns_since_long_chunk = 0
         elif len(running) > 0:
             (result.inputs, result.delta, result.extra_inputs,
@@ -637,11 +644,59 @@ def _build_active_chunk(self):
         inputs.is_first_chunk = False
         inputs.is_chunk_multimodal = is_chunk_multimodal
         maker._short_prefill_turns_since_long_chunk = 0
-        return _ForwardInputsResult(running=running,
-                                    inputs=inputs,
-                                    delta=delta,
-                                    cache_inputs=cache_inputs,
-                                    extra_inputs=extra_inputs)
+        result = _ForwardInputsResult(running=running,
+                                      inputs=inputs,
+                                      delta=delta,
+                                      cache_inputs=cache_inputs,
+                                      extra_inputs=extra_inputs)
+        if not is_last_chunk:
+            self._try_internal_prefill(result)
+        return result
+
+    def _try_internal_prefill(self, result: _ForwardInputsResult) -> bool:
+        """Expand an admitted final split without changing its pending save."""
+        maker = self.maker
+        if not getattr(maker.config, 'enable_internal_prefill_checkpoint', False):
+            return False
+        if (any(seq.status == MessageStatus.RUNNING for seq in maker.running_seqs)
+                or self.scheduler.has_kv_connector()
+                or len(result.running) != 1 or result.delta is not None):
+            return False
+        seq = result.running[0]
+        if (seq.logprob_start_pos >= 0 or seq.return_logits or seq.return_ce_loss or seq.return_routed_experts
+                or seq.history_embeddings or seq.history_multimodals.multimodals):
+            return False
+        start, end = seq.num_history_ids, seq.input_end_pos
+        step = start + result.inputs.max_q_seqlen
+        alignment = seq._seq_meta.prefill_save_alignment
+        lookahead = seq._seq_meta.prefix_cache_token_lookahead
+        if (maker.spec_decoding or lookahead
+                or not alignment or end % alignment == 0 or step % alignment
+                or step + lookahead > end
+                or not start < step < end or step - start < 64 or (step - start) % 16
+                or end - start != seq.num_token_ids or end - start > maker.config.max_prefill_token_num
+                or end - start <= maker.config.prefill_graph_max_tokens):
+            return False
+        cache_inputs = result.cache_inputs
+        if (cache_inputs is None or cache_inputs.kv_save_plan is not None
+                or cache_inputs.state_save_plan is None or len(cache_inputs.state_save_plan[0]) != 1):
+            return False
+        source, target = cache_inputs.state_save_plan
+        capture = PrefillStateCapture(start, end, step, source[0], target[0])
+        if not self._reserve_chunk(seq, end - start, is_last_chunk=True):
+            return False
+        cache_inputs = replace(cache_inputs, state_save_plan=None, state_capture=capture)
+        was_first = result.inputs.is_first_chunk
+        result.inputs, result.delta, result.extra_inputs, result.cache_inputs = self._build_prefill_inputs(
+            result.running, cache_inputs=cache_inputs)
+        result.inputs.is_chunk = True
+        result.inputs.is_first_chunk = was_first
+        result.inputs.is_last_chunk = True
+        maker.long_context_chunker.clear()
+        logger.debug('Internal prefill checkpoint: session_id=%s seq_id=%s start=%s end=%s step=%s '
+                     'source_slot=%s target_slot=%s',
+                     seq.session_id, seq.seq_id, start, end, step, source[0], target[0])
+        return True
 
     def _reserve_chunk(self,
                        seq: 'SchedulerSequence',
@@ -659,12 +714,13 @@ def _reserve_chunk(self,
                                                          prealloc_size=prealloc_size,
                                                          is_last_chunk=is_last_chunk)
 
-    def _build_prefill_inputs(self, seqs: 'SeqList'):
+    def _build_prefill_inputs(self, seqs: 'SeqList', cache_inputs: CacheCheckpointInputs | None = None):
         assert all(seq.kv_token_limit is None for seq in seqs), (
             'Non-final chunks must use exclusive chunk inputs, not full prefill')
         maker = self.maker
         inputs = maker.create_model_inputs(seqs, True)
-        cache_inputs = maker._prepare_prefill_cache_inputs(seqs)
+        if cache_inputs is None:
+            cache_inputs = maker._prepare_prefill_cache_inputs(seqs)
         delta, valid_seqs, _ = maker.create_model_inputs_delta_valid_only()
         maker.running_seqs = valid_seqs
         extra_inputs = maker.model_agent_strategy.make_extra_inputs(seqs, inputs)
diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py
index 28f31e01d6..a782a96683 100644
--- a/lmdeploy/pytorch/engine/model_agent/agent.py
+++ b/lmdeploy/pytorch/engine/model_agent/agent.py
@@ -247,6 +247,8 @@ def model_forward(
         # Attach operator-owned block caches and configured state caches.
         context.block_caches = cache_engine.block_caches
         context.named_state_caches = state_cache_engine.named_state_caches
+        if cache_inputs is not None and not inputs.is_dummy:
+            context.state_capture = cache_inputs.state_capture
 
         with ctx_mgr.context(context):
             # Some backends synchronize while building context. Queue restore
diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py
index 54424b172c..3c89b2c8e2 100644
--- a/lmdeploy/pytorch/envs.py
+++ b/lmdeploy/pytorch/envs.py
@@ -232,6 +232,7 @@ def _patched_get_env(
     fp32_mamba_ssm_dtype = env_to_bool('LMDEPLOY_FP32_MAMBA_SSM_DTYPE', False)
 
     kda_prefill_backend = env_to_choice('LMDEPLOY_KDA_PREFILL_BACKEND', 'fla', {'fla', 'flash_kda'})
+    internal_prefill_checkpoint = env_to_bool('LMDEPLOY_INTERNAL_PREFILL_CHECKPOINT', False)
 
     # DSA indexer fusion
     disable_dsa_indexer_fusion = env_to_bool('LMDEPLOY_DISABLE_DSA_INDEXER_FUSION', False)
diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py
index ceb4b5c5b4..ae66432992 100644
--- a/lmdeploy/pytorch/model_inputs.py
+++ b/lmdeploy/pytorch/model_inputs.py
@@ -17,6 +17,7 @@
 from lmdeploy.pytorch.utils import CtxMgrBase, singleton
 
 if TYPE_CHECKING:
+    from lmdeploy.pytorch.engine.cache_inputs import PrefillStateCapture
     from lmdeploy.pytorch.strategies.base import StrategyFactoryBase
 
 
@@ -355,6 +356,7 @@ class StepContext:
     # named views for operator-owned block caches and configured state caches
     block_caches: Mapping[str, torch.Tensor] | None = None
     named_state_caches: Mapping[str, torch.Tensor] | None = None
+    state_capture: 'PrefillStateCapture | None' = None
 
     # mrope
     mrope_position_ids: torch.Tensor | None = None
diff --git a/lmdeploy/pytorch/models/glm5_next.py b/lmdeploy/pytorch/models/glm5_next.py
index 62588a1690..2810f8b633 100644
--- a/lmdeploy/pytorch/models/glm5_next.py
+++ b/lmdeploy/pytorch/models/glm5_next.py
@@ -31,6 +31,7 @@
     kpool_select_prefill_cuda,
     kpool_write_token_cache_cuda,
 )
+from lmdeploy.pytorch.backends.kda import KdaPrefillCheckpoint
 from lmdeploy.pytorch.configurations.glm5_next import is_glm5_kda_layer
 from lmdeploy.pytorch.consts import (
     GLM5_KDA_CONV_STATE,
@@ -40,6 +41,7 @@
 )
 from lmdeploy.pytorch.distributed import get_dist_group, get_dist_manager, get_tp_world_rank
 from lmdeploy.pytorch.engine.cache_engine.schema import BlockCacheRequest
+from lmdeploy.pytorch.engine.cache_inputs import PrefillStateCapture
 from lmdeploy.pytorch.kernels.cuda.activation import silu_and_mul_masked_post_quant_fwd
 from lmdeploy.pytorch.model_inputs import StepContext, StepContextManager, get_step_ctx_manager
 from lmdeploy.pytorch.models.patch import get_build_model_context
@@ -723,7 +725,8 @@ def __init__(self,
 
     def forward(self, hidden_states: torch.Tensor,
                 past_key_value: Sequence[torch.Tensor],
-                kda_metadata: GatedDeltaMeta) -> torch.Tensor:
+                kda_metadata: GatedDeltaMeta,
+                checkpoint: KdaPrefillCheckpoint | None = None) -> torch.Tensor:
         mixed_qkv = self.qkv_proj(hidden_states)
         raw_beta = self.b_proj(hidden_states)
         raw_gate, norm_gate = self.fg_b_proj(self.fg_a_proj(hidden_states))
@@ -742,6 +745,7 @@ def forward(self, hidden_states: torch.Tensor,
             num_heads=self.local_num_heads,
             head_dim=self.head_dim,
             lower_bound=self.lower_bound,
+            checkpoint=checkpoint,
         )
         norm_gate = norm_gate.unflatten(-1,
                                         (self.local_num_heads, self.head_dim))
@@ -1400,7 +1404,8 @@ def forward(self, hidden_states: torch.Tensor,
                 kpool_metadata: Any = None,
                 hc_input_fp32: torch.Tensor | None = None,
                 prepare_hc_fp32: bool = False,
-                all_routed_experts: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor | None]:
+                all_routed_experts: torch.Tensor | None = None,
+                kda_checkpoint: KdaPrefillCheckpoint | None = None) -> tuple[torch.Tensor, torch.Tensor | None]:
         residual = hidden_states
         hidden_states, post, comb = self._hc_pre(
             hidden_states,
@@ -1413,7 +1418,8 @@ def forward(self, hidden_states: torch.Tensor,
         if self.is_linear_attention:
             hidden_states = self.self_attn(hidden_states,
                                            past_key_value=past_key_value,
-                                           kda_metadata=kda_metadata)
+                                           kda_metadata=kda_metadata,
+                                           checkpoint=kda_checkpoint)
         else:
             hidden_states = self.self_attn(hidden_states,
                                            past_key_value=past_key_value,
@@ -1475,6 +1481,7 @@ def forward(
         state_ids: torch.Tensor,
         inputs_embeds: torch.Tensor | None = None,
         all_routed_experts: torch.Tensor | None = None,
+        state_capture: PrefillStateCapture | None = None,
     ) -> torch.Tensor:
         del position_ids  # GLM-5.3 text attention has qk_rope_head_dim == 0.
         if state_ids is None:
@@ -1487,6 +1494,14 @@ def forward(
         kda_metadata = self.gated_delta_meta_builder(hidden_states.size(1),
                                       self.config.linear_conv_kernel_dim,
                                       state_ids, attn_metadata)
+        capture_offsets = None
+        if state_capture is not None:
+            if (attn_metadata.is_decoding or state_ids.numel() != 1
+                    or hidden_states.size(1) != state_capture.input_end - state_capture.input_start):
+                raise ValueError('GLM IC requires an exclusive prefill matching the capture interval.')
+            capture_offsets = torch.tensor(
+                [state_capture.checkpoint_step - state_capture.input_start],
+                dtype=kda_metadata.cu_seqlens.dtype, device=hidden_states.device)
         if len(past_key_values) != len(self.layers):
             raise RuntimeError(
                 f'GLM-5.3 expects {len(self.layers)} layer caches, got {len(past_key_values)}.'
@@ -1500,7 +1515,18 @@ def forward(
                 get_tp_world_rank('attn')[1] == 0, page_size=page_size)
         hc_input_fp32 = None
         prepare_hc_fp32 = hidden_states.is_cuda and hidden_states.size(0) * hidden_states.size(1) <= 32
+        captured_layers = 0
         for layer, past_key_value in zip(self.layers, past_key_values):
+            kda_checkpoint = None
+            if state_capture is not None and layer.is_linear_attention:
+                target = state_capture.target_slot
+                if max(target, state_capture.source_slot) >= past_key_value[0].size(0):
+                    raise ValueError('GLM IC slot is outside the allocated state cache.')
+                kda_checkpoint = KdaPrefillCheckpoint(
+                    offset=state_capture.checkpoint_step - state_capture.input_start,
+                    offsets=capture_offsets,
+                    conv_state_out=past_key_value[0][target:target + 1],
+                    recurrent_state_out=past_key_value[1][target:target + 1])
             kpool_tail_state = None
             if not layer.is_linear_attention:
                 past_key_value, kpool_tail_state = past_key_value[:-2], past_key_value[-2:]
@@ -1513,7 +1539,11 @@ def forward(
                                   kpool_metadata=kpool_metadata,
                                   hc_input_fp32=hc_input_fp32,
                                   prepare_hc_fp32=prepare_hc_fp32,
-                                  all_routed_experts=all_routed_experts)
+                                  all_routed_experts=all_routed_experts,
+                                  kda_checkpoint=kda_checkpoint)
+            captured_layers += kda_checkpoint is not None
+        if state_capture is not None and captured_layers != len(self.config.linear_layer_ids):
+            raise RuntimeError('GLM IC did not submit all KDA checkpoint writes.')
         hidden_states = hidden_states.mean(dim=2)
         return self.norm(hidden_states)
 
@@ -1574,6 +1604,7 @@ def forward(
         vision_prompt_order: list[int] | None = None,
         multimodal_mask: torch.Tensor | None = None,
         return_input_embeds: bool = False,
+        state_capture: PrefillStateCapture | None = None,
         **kwargs,
     ) -> torch.Tensor:
         if inputs_embeds is None and (vision_groups or
@@ -1627,7 +1658,8 @@ def forward(
                           attn_metadata=attn_metadata,
                           inputs_embeds=inputs_embeds,
                           state_ids=state_ids,
-                          all_routed_experts=all_routed_experts)
+                          all_routed_experts=all_routed_experts,
+                          state_capture=state_capture)
         if all_routed_experts is None and not return_input_embeds:
             return hidden_states
         outputs = dict(hidden_states=hidden_states)
@@ -1837,6 +1869,7 @@ def prepare_inputs_for_generation(
                     attn_metadata=context.attn_metadata,
                     inputs_embeds=inputs_embeds,
                     state_ids=context.state_offsets,
+                    state_capture=context.state_capture,
                     pixel_values=pixel_values,
                     grid_thw=grid_thw,
                     vision_groups=vision_groups,
diff --git a/lmdeploy/pytorch/nn/kda.py b/lmdeploy/pytorch/nn/kda.py
index 17072d109e..16727a6a2d 100644
--- a/lmdeploy/pytorch/nn/kda.py
+++ b/lmdeploy/pytorch/nn/kda.py
@@ -5,7 +5,7 @@
 from torch import nn
 
 from lmdeploy.pytorch.backends import get_backend
-from lmdeploy.pytorch.backends.kda import KdaBuildSpec
+from lmdeploy.pytorch.backends.kda import KdaBuildSpec, KdaPrefillCheckpoint
 from lmdeploy.pytorch.models.patch import get_build_model_context
 
 
@@ -34,6 +34,7 @@ def forward(
         num_heads: int,
         head_dim: int,
         lower_bound: float,
+        checkpoint: KdaPrefillCheckpoint | None = None,
     ) -> torch.Tensor:
         return self.impl.forward(
             mixed_qkv=mixed_qkv,
@@ -49,4 +50,5 @@ def forward(
             num_heads=num_heads,
             head_dim=head_dim,
             lower_bound=lower_bound,
+            checkpoint=checkpoint,
         )
diff --git a/tests/pytorch/engine/test_inputs_maker.py b/tests/pytorch/engine/test_inputs_maker.py
index 855c52448b..e5dea52575 100644
--- a/tests/pytorch/engine/test_inputs_maker.py
+++ b/tests/pytorch/engine/test_inputs_maker.py
@@ -10,7 +10,7 @@
 import lmdeploy.pytorch.engine.engine_loop as engine_loop_module
 import lmdeploy.pytorch.engine.inputs_maker as inputs_maker_module
 from lmdeploy.pytorch.disagg.config import EngineRole
-from lmdeploy.pytorch.engine.cache_inputs import CacheCheckpointInputs
+from lmdeploy.pytorch.engine.cache_inputs import CacheCheckpointInputs, PrefillStateCapture
 from lmdeploy.pytorch.engine.engine_loop import EngineLoop, RunableEventAsync
 from lmdeploy.pytorch.engine.inputs_maker import (
     InputsMakerAsync,
@@ -361,7 +361,8 @@ def _fake_model_inputs(is_chunk: bool = False):
                            seq_length=torch.tensor([1]))
 
 
-def test_engine_loop_keeps_state_save_pinned_until_output_boundary():
+@pytest.mark.parametrize('internal', [False, True])
+def test_engine_loop_keeps_state_save_pinned_until_output_boundary(internal):
     events = []
 
     class _StateCheckpoints:
@@ -405,7 +406,8 @@ async def get_output_async(self):
     loop.executor = _Executor(state_checkpoints)
     loop._sleep_requested = False
     model_inputs = SimpleNamespace()
-    cache_inputs = CacheCheckpointInputs(state_save_plan=((0, ), (1, )))
+    cache_inputs = (CacheCheckpointInputs(state_capture=PrefillStateCapture(0, 8256, 8192, 0, 1))
+                    if internal else CacheCheckpointInputs(state_save_plan=((0, ), (1, ))))
     forward_inputs = dict(inputs=model_inputs, delta=None, cache_inputs=cache_inputs)
 
     forward_inputs, next_running = asyncio.run(loop._main_loop_get_outputs([object()], forward_inputs))
diff --git a/tests/pytorch/engine/test_internal_prefill_checkpoint.py b/tests/pytorch/engine/test_internal_prefill_checkpoint.py
new file mode 100644
index 0000000000..ce7b7ef7ca
--- /dev/null
+++ b/tests/pytorch/engine/test_internal_prefill_checkpoint.py
@@ -0,0 +1,199 @@
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from lmdeploy.pytorch.disagg.config import EngineRole
+from lmdeploy.pytorch.engine.cache_inputs import CacheCheckpointInputs, PrefillStateCapture
+from lmdeploy.pytorch.engine.inputs_maker import _ForwardInputsResult, _ForwardInputsTask
+from lmdeploy.pytorch.messages import MessageStatus
+
+
+def make_task(start=0, end=8256, step=8192, budget=16384, reserve=True, mtp=False):
+    seq = SimpleNamespace(
+        num_history_ids=start, input_end_pos=end, num_token_ids=end - start,
+        logprob_start_pos=-1, return_logits=False, return_ce_loss=False, return_routed_experts=False,
+        history_embeddings=[], history_multimodals=SimpleNamespace(multimodals={}),
+        _seq_meta=SimpleNamespace(prefill_save_alignment=8192, prefix_cache_token_lookahead=int(mtp)),
+        kv_token_limit=step,
+        session_id=1, seq_id=2)
+
+    def reserve_chunk(seq, chunk_size, prealloc_size, is_last_chunk):
+        assert chunk_size == end - start and is_last_chunk and prealloc_size == 1
+        if reserve:
+            seq.kv_token_limit = None
+        return reserve
+
+    maker = SimpleNamespace(
+        config=SimpleNamespace(enable_internal_prefill_checkpoint=True, max_prefill_token_num=budget,
+                               prefill_graph_max_tokens=2048, role=EngineRole.Hybrid),
+        spec_decoding=mtp, running_seqs=[],
+        scheduler=SimpleNamespace(has_kv_connector=lambda: False,
+                                  reserve_long_context_chunk=Mock(side_effect=reserve_chunk)),
+        engine_strategy=SimpleNamespace(get_prealloc_size=lambda prefill: 1),
+        long_context_chunker=SimpleNamespace(clear=Mock()),
+        model_agent_strategy=SimpleNamespace(make_extra_inputs=lambda seqs, inputs: None),
+        create_model_inputs=Mock(return_value=SimpleNamespace()),
+        create_model_inputs_delta_valid_only=lambda: (None, [], []),
+        _prepare_prefill_cache_inputs=Mock(side_effect=AssertionError('Do not reserve H a second time')))
+    task = _ForwardInputsTask(maker, prefill=True)
+    result = _ForwardInputsResult(
+        running=[seq], inputs=SimpleNamespace(max_q_seqlen=step - start, is_first_chunk=start == 0),
+        cache_inputs=CacheCheckpointInputs(state_restore_plan=((2,), (1,)), state_save_plan=((1,), (3,))))
+    return task, result
+
+
+@pytest.mark.parametrize('mtp', [False])
+@pytest.mark.parametrize(('start', 'end', 'step'), [(0, 8193, 8192), (0, 8256, 8192), (0, 10000, 8192),
+                                                   (8192, 20000, 16384)])
+def test_internal_capture_preserves_one_reservation_and_restore(start, end, step, mtp):
+    task, result = make_task(start, end, step, mtp=mtp)
+    assert task._try_internal_prefill(result)
+    assert result.cache_inputs.state_capture == PrefillStateCapture(start, end, step, 1, 3)
+    assert result.cache_inputs.state_save_plan is None
+    assert result.cache_inputs.state_restore_plan == ((2,), (1,))
+    assert result.inputs.is_chunk and result.inputs.is_last_chunk
+    assert result.inputs.is_first_chunk == (start == 0)
+    task.maker._prepare_prefill_cache_inputs.assert_not_called()
+    task.maker.long_context_chunker.clear.assert_called_once()
+    assert result.cache_inputs.to_device('cpu').state_capture == result.cache_inputs.state_capture
+
+
+@pytest.mark.parametrize('condition', ['disabled', 'unsupported_spec', 'mixed', 'connector', 'logprob', 'logits',
+                                      'routes', 'embeddings',
+                                      'multimodal', 'end_checkpoint', 'budget', 'pcg', 'no_slot', 'kv_copy',
+                                      'unaligned_start', 'no_expansion_capacity'])
+def test_internal_capture_falls_back_without_losing_original_plan(condition):
+    task, result = make_task(reserve=condition != 'no_expansion_capacity')
+    maker, seq = task.maker, result.running[0]
+    if condition == 'disabled':
+        maker.config.enable_internal_prefill_checkpoint = False
+    elif condition == 'unsupported_spec':
+        maker.spec_decoding = True
+    elif condition == 'mixed':
+        maker.running_seqs = [SimpleNamespace(status=MessageStatus.RUNNING)]
+    elif condition == 'connector':
+        maker.scheduler.has_kv_connector = lambda: True
+    elif condition == 'logprob':
+        seq.logprob_start_pos = 0
+    elif condition == 'logits':
+        seq.return_logits = True
+    elif condition == 'routes':
+        seq.return_routed_experts = True
+    elif condition == 'embeddings':
+        seq.history_embeddings = [object()]
+    elif condition == 'multimodal':
+        seq.history_multimodals.multimodals = {'image': [object()]}
+    elif condition == 'end_checkpoint':
+        seq.input_end_pos = seq.num_token_ids = 16384
+    elif condition == 'budget':
+        maker.config.max_prefill_token_num = 8192
+    elif condition == 'pcg':
+        maker.config.prefill_graph_max_tokens = 16384
+    elif condition == 'no_slot':
+        result.cache_inputs = replace(result.cache_inputs, state_save_plan=None)
+    elif condition == 'kv_copy':
+        result.cache_inputs = replace(result.cache_inputs, kv_save_plan=torch.tensor([[1], [2]]))
+    elif condition == 'unaligned_start':
+        seq.num_history_ids = 1
+        seq.num_token_ids -= 1
+        result.inputs.max_q_seqlen -= 1
+    original_cache, original_inputs = result.cache_inputs, result.inputs
+    assert not task._try_internal_prefill(result)
+    assert result.cache_inputs is original_cache and result.inputs is original_inputs
+    assert seq.kv_token_limit == 8192
+    maker._prepare_prefill_cache_inputs.assert_not_called()
+    maker.long_context_chunker.clear.assert_not_called()
+    maker.create_model_inputs.assert_not_called()
+    assert maker.scheduler.reserve_long_context_chunk.call_count == (condition == 'no_expansion_capacity')
+
+
+@pytest.mark.parametrize('arguments', [(0, 64, 64, 1, 2), (64, 128, 64, 1, 2), (0, 128, 64, 1, 1),
+                                      (0, 128, 64, -1, 2)])
+def test_invalid_internal_capture_is_rejected(arguments):
+    with pytest.raises(ValueError):
+        PrefillStateCapture(*arguments)
+
+
+def test_capture_cannot_be_overwritten_by_final_copy():
+    with pytest.raises(ValueError):
+        CacheCheckpointInputs(state_capture=PrefillStateCapture(0, 8256, 8192, 1, 3),
+                              state_save_plan=((1,), (3,)))
+
+
+def test_completed_tracked_sequence_does_not_block_capture_or_drop_cleanup_delta():
+    task, result = make_task(start=8192, end=20000, step=16384)
+    finished = SimpleNamespace(status=MessageStatus.STOPPED)
+    task.maker.running_seqs = [finished]
+    cleanup_delta = SimpleNamespace(indice_cpu=[])
+    task.maker.create_model_inputs_delta_valid_only = lambda: (cleanup_delta, [], [finished])
+    assert task._try_internal_prefill(result)
+    assert result.delta is cleanup_delta
+    assert task.maker.running_seqs == []
+    assert result.cache_inputs.state_capture == PrefillStateCapture(8192, 20000, 16384, 1, 3)
+
+
+def test_non_mtp_ic_config_rejects_mtp(monkeypatch):
+    from lmdeploy.pytorch import envs
+    from lmdeploy.pytorch.backends.cuda import flash_kda
+    from lmdeploy.pytorch.config import CacheConfig
+    from lmdeploy.pytorch.configurations.glm5_next import update_cache_config
+
+    monkeypatch.setattr(envs, 'internal_prefill_checkpoint', True)
+    monkeypatch.setattr(envs, 'kda_prefill_backend', 'flash_kda')
+    monkeypatch.setattr(flash_kda, 'check_flash_kda', Mock())
+    config = CacheConfig(max_batches=32, block_size=64, num_cpu_blocks=0, num_gpu_blocks=64,
+                         enable_prefix_caching=True, prefill_save_alignment=8192)
+    update_cache_config(config)
+    assert config.enable_internal_prefill_checkpoint
+    with pytest.raises(ValueError, match='non-MTP AR'):
+        update_cache_config(config, token_lookahead=1)
+
+
+@pytest.mark.parametrize(('internal', 'routes'), [(False, False), (False, True), (True, False)])
+@pytest.mark.parametrize('return_embeds', [False, True])
+def test_wrapper_preserves_capture_routes_and_mtp_embeddings(internal, routes, return_embeds):
+    from lmdeploy.pytorch.models.glm5_next import Glm5NextForConditionalGeneration
+
+    class Backbone(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.embed_tokens = torch.nn.Embedding(16, 4)
+            self.seen_capture = None
+            self.seen_routes = None
+
+        def get_input_embeddings(self):
+            return self.embed_tokens
+
+        def forward(self, input_ids, inputs_embeds, all_routed_experts, state_capture, **kwargs):
+            self.seen_capture = state_capture
+            self.seen_routes = all_routed_experts
+            if all_routed_experts is not None:
+                all_routed_experts[:, 1].fill_(7)
+            return self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
+
+    model = Glm5NextForConditionalGeneration.__new__(Glm5NextForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.enable_return_routed_experts = routes
+    model.config = SimpleNamespace(num_hidden_layers=3, num_experts_per_tok=2)
+    model.model = Backbone()
+    input_ids = torch.arange(65).remainder(16).unsqueeze(0)
+    capture = PrefillStateCapture(0, 65, 64, 1, 2) if internal else None
+    output = model(input_ids, input_ids, [], state_capture=capture, return_input_embeds=return_embeds)
+    expected = model.get_input_embeddings()(input_ids)
+    assert model.model.seen_capture is capture
+    if not routes and not return_embeds:
+        torch.testing.assert_close(output, expected, rtol=0, atol=0)
+        assert model.model.seen_routes is None
+        return
+    torch.testing.assert_close(output['hidden_states'], expected, rtol=0, atol=0)
+    assert ('target_inputs_embeds' in output) == return_embeds
+    assert ('all_routed_experts' in output) == routes
+    if return_embeds:
+        torch.testing.assert_close(output['target_inputs_embeds'], expected, rtol=0, atol=0)
+    if routes:
+        assert output['all_routed_experts'] is model.model.seen_routes
+        assert output['all_routed_experts'].shape == (65, 3, 2)
+        assert torch.all(output['all_routed_experts'][:, 1].int() == 7)
diff --git a/tests/pytorch/kernel/test_kda_internal_checkpoint.py b/tests/pytorch/kernel/test_kda_internal_checkpoint.py
new file mode 100644
index 0000000000..232009d1cf
--- /dev/null
+++ b/tests/pytorch/kernel/test_kda_internal_checkpoint.py
@@ -0,0 +1,66 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA')
+@pytest.mark.parametrize(('heads', 'length', 'offset', 'history'), [(4, 65, 64, 0), (64, 8256, 8192, 0),
+                                                                  (4, 11808, 8192, 8192)])
+def test_native_backend_internal_checkpoint_and_restore(monkeypatch, heads, length, offset, history):
+    pytest.importorskip('flash_kda')
+    from lmdeploy.pytorch import envs
+    from lmdeploy.pytorch.backends.cuda.kda import CudaKdaImpl
+    from lmdeploy.pytorch.backends.kda import KdaPrefillCheckpoint
+
+    monkeypatch.setattr(envs, 'kda_prefill_backend', 'flash_kda')
+    impl = CudaKdaImpl()
+    torch.manual_seed(42)
+    width, dim, slots = 4, 128, 5
+    conv_pool = torch.full((slots, 2, heads * dim * 3, width), float('nan'), dtype=torch.bfloat16, device='cuda')
+    state_pool = torch.full((slots, 2, heads, dim, dim), float('nan'), dtype=torch.float32, device='cuda')
+    conv, state = conv_pool[:, 1], state_pool[:, 1]
+    conv[1].zero_()
+    state[1].zero_()
+    if history:
+        conv[1].normal_(std=0.1)
+        state[1].normal_(std=0.1)
+    conv[2].copy_(conv[1])
+    state[2].copy_(state[1])
+    values = dict(mixed_qkv=torch.randn(1, length + 1, 3 * heads * dim, dtype=torch.bfloat16, device='cuda') / 4,
+                  raw_gate=torch.randn(1, length + 1, heads * dim, dtype=torch.bfloat16, device='cuda'),
+                  raw_beta=torch.randn(1, length + 1, heads, dtype=torch.bfloat16, device='cuda'))
+    kwargs = dict(conv_state=conv, recurrent_state=state,
+                  conv_weight=torch.randn(3 * heads * dim, width, dtype=torch.bfloat16, device='cuda') / 8,
+                  conv_bias=None, a_log=torch.zeros(heads, device='cuda'),
+                  dt_bias=torch.zeros(heads * dim, device='cuda'),
+                  num_heads=heads, head_dim=dim, lower_bound=-5.0)
+
+    def forward(start, end, slot, checkpoint=None, decode=False):
+        meta = SimpleNamespace(is_decoding=decode, num_spec_tokens=0, is_init=None,
+                               valid_state=torch.tensor([True], device='cuda'),
+                               state_ids=torch.tensor([slot], device='cuda'),
+                               cu_seqlens=torch.tensor([0, end - start], dtype=torch.int64, device='cuda'))
+        return impl.forward(**{name: value[:, start:end] for name, value in values.items()},
+                            metadata=meta, checkpoint=checkpoint, **kwargs)
+
+    checkpoint = KdaPrefillCheckpoint(offset, torch.tensor([offset], dtype=torch.int64, device='cuda'),
+                                     conv[3:4], state[3:4])
+    full = forward(0, length, 1, checkpoint)
+    prefix = forward(0, offset, 2)
+    torch.testing.assert_close(conv[3], conv[2], atol=0, rtol=0)
+    torch.testing.assert_close(state[3], state[2], atol=0, rtol=0)
+    conv[4].copy_(conv[3])
+    state[4].copy_(state[3])
+    saved_state = state[3].clone()
+    suffix = forward(offset, length, 2)
+    restored_suffix = forward(offset, length, 4)
+    torch.testing.assert_close(full, torch.cat((prefix, suffix), dim=1), atol=0, rtol=0)
+    torch.testing.assert_close(suffix, restored_suffix, atol=0, rtol=0)
+    torch.testing.assert_close(state[1], state[2], atol=0, rtol=0)
+    torch.testing.assert_close(conv[1], conv[2], atol=0, rtol=0)
+    torch.testing.assert_close(forward(length, length + 1, 1, decode=True),
+                               forward(length, length + 1, 2, decode=True), atol=0, rtol=0)
+    torch.testing.assert_close(state[3], saved_state, atol=0, rtol=0)
+    assert torch.isnan(conv_pool[:, 0]).all() and torch.isnan(state_pool[:, 0]).all()
+    assert torch.isnan(conv[0]).all() and torch.isnan(state[0]).all()

From a447f3768c5a7818563f85a40b9e4d5e55a0ec78 Mon Sep 17 00:00:00 2001
From: yidingcheng0206 <300775907+yidingcheng0206@users.noreply.github.com>
Date: Tue, 20 Oct 2026 20:33:29 +0800
Subject: [PATCH 4/4] feat(pytorch): support MTP internal prefill checkpoints

---
 docs/en/advance/internal_checkpoint.md        | 114 +++++++++---------
 lmdeploy/pytorch/backends/cuda/kda.py         |  31 ++++-
 lmdeploy/pytorch/backends/kda.py              |   1 +
 lmdeploy/pytorch/configurations/glm5_next.py  |   4 +-
 lmdeploy/pytorch/engine/inputs_maker.py       |   4 +-
 lmdeploy/pytorch/models/glm5_next.py          |   3 +-
 .../test_internal_prefill_checkpoint.py       |  25 +++-
 .../kernel/test_kda_internal_checkpoint.py    |  90 ++++++++++++++
 8 files changed, 201 insertions(+), 71 deletions(-)

diff --git a/docs/en/advance/internal_checkpoint.md b/docs/en/advance/internal_checkpoint.md
index 9951702ae6..29aceb9f98 100644
--- a/docs/en/advance/internal_checkpoint.md
+++ b/docs/en/advance/internal_checkpoint.md
@@ -1,56 +1,58 @@
-# GLM-5.3 native KDA prefill backend
-
-The default `LMDEPLOY_KDA_PREFILL_BACKEND=fla` keeps the existing FLA entry
-point. FLA may dispatch to its own FlashKDA backend when its conditions are met.
-Set `LMDEPLOY_KDA_PREFILL_BACKEND=flash_kda` to explicitly select LMDeploy's
-native FlashKDA adapter instead. This does not change the decode backend.
-
-The native adapter requires BF16 inputs, head dimension 128, FP32 recurrent
-state, and a gate lower bound. The installed FlashKDA operator must expose
-`checkpoint_state` and `checkpoint_offsets`; incompatible installations fail
-the capability check instead of silently switching backends.
-
-The validated FlashKDA source revision is
-`ee0be888cd0e972f9409bf53756f8c38c6652173`, with Torch 2.10.0+cu128 and
-flash-linear-attention 0.5.2 on H200. No FlashKDA, FLA, convolution, KPool, or
-TileLang decode kernel source is modified by this integration.
-
-For controlled comparisons, select the same prefill backend in both arms.
-Use `FLA_FLASH_KDA=0` when an explicit FLA Triton reference is required.
-Selecting the native adapter does not by itself enable internal checkpoints.
-
-## Internal prefill checkpoints
-
-Set `LMDEPLOY_INTERNAL_PREFILL_CHECKPOINT=1` together with the native backend
-and local prefix caching to enable IC. IC is disabled by default.
-
-The scheduler first reserves the same safe checkpoint H as the original
-split-prefill path. If the remaining prompt fits the existing prefill budget,
-it can expand that forward to the prompt end E while retaining its reservation
-at H. FlashKDA exports the recurrent state directly into that checkpoint slot;
-the backend copies the preceding convolution-input window into its conv state.
-The normal runtime state still advances to E. An endpoint copy must not
-overwrite the internal checkpoint.
-
-The capture description contains the input start/end, checkpoint position, and
-source/destination slots. Public engine components transport this description
-and retain existing reservation, pinning, publication, and cleanup mechanisms.
-GLM-specific code maps slots to per-layer KDA output views. KPool runtime tails
-remain excluded from checkpoint copies, as defined by PR #5037.
-
-Initial support is non-MTP, text-only, exclusive single-sequence prefill with
-one internal save. A non-MTP prompt end that also requires an aligned checkpoint
-keeps the original split rather than dropping either save. Connector/Mooncake,
-multimodal, scoring, mixed local decode/prefill, insufficient capacity, and
-unsupported capture layouts do not enter the IC path. Invalid enabled-backend
-configurations fail validation; per-request ineligibility preserves the
-ordinary split path.
-
-Merged intervals within `piecewise_cudagraph_max_tokens` retain the original
-path. Current IC runs outside prefill PCG; it does not disable decode graphs.
-`prefill_save_alignment` constrains absolute checkpoint positions, including
-prompt-end saves, but does not guarantee a save at every aligned position.
-
-Comparisons must hold the backend, prefill budget, PCG cap, alignment, and
-request set fixed. Export/copy work and a larger single forward can consume
-time and memory; IC is not a zero-cost operation or a universal speedup.
+# GLM-5.3 internal prefill checkpoints
+
+IC avoids a model forward used only to reach a checkpoint boundary. It is
+disabled by default and supports non-MTP AR and `deepseek_mtp`.
+
+## Enablement
+
+Use local prefix caching (`--enable-prefix-caching`) and set:
+
+```bash
+export LMDEPLOY_KDA_PREFILL_BACKEND=flash_kda
+export LMDEPLOY_INTERNAL_PREFILL_CHECKPOINT=1
+```
+
+The default backend remains the FLA entry, which may dispatch to FlashKDA.
+Native selection alone does not enable IC or change decode. The native adapter
+requires BF16 inputs, head dimension 128, FP32 recurrent state, a gate lower
+bound, and operator outputs `checkpoint_state` / `checkpoint_offsets`.
+Incompatible enabled configurations fail validation rather than silently fall back.
+
+Validated FlashKDA revision: `ee0be888cd0e972f9409bf53756f8c38c6652173`,
+Torch 2.10.0+cu128 and FLA 0.5.2 on H200. No kernel sources are modified.
+
+## Execution and safety
+
+The scheduler retains the original reservation at H and, when capacity permits,
+expands execution to prompt end E within the existing prefill budget:
+
+```text
+OFF: forward [start,H) -> save S_H -> forward [H,E)
+ ON: forward [start,E), exporting S_H internally; runtime advances to S_E
+```
+
+FlashKDA writes recurrent state into the checkpoint slot. Conv state comes from
+the preceding projected-input window. Existing pinning, publication and cleanup
+are reused; endpoint copying must not overwrite the internal state. KPool raw
+tails remain runtime-only.
+
+For MTP K, operator offset is `H - start`, recurrent destination is `H % (K+1)`,
+and conv window `[H-width,H)` is placed in the `width+K` ring, wrapping if needed.
+Draft shifting, known-prompt lookahead checks and accepted-history recovery are
+unchanged. Predicting a next token does not make a prompt-end MTP checkpoint
+publishable; MTP decode checkpoint saving remains disabled.
+
+## Limits and validation
+
+- Only exclusive single-sequence text prefill with one internal save is supported.
+- Required non-MTP endpoint saves retain the original split; no checkpoint is dropped.
+- Connector/Mooncake, multimodal, scoring/routed-expert output, mixed local
+  decode/prefill, insufficient capacity and unsupported layouts do not enter IC.
+- Intervals at or below `piecewise_cudagraph_max_tokens` retain the original path.
+  IC currently runs outside prefill PCG; decode graphs remain enabled.
+- `prefill_save_alignment` filters absolute internal and endpoint save positions,
+  including with IC OFF; aligned positions may be skipped. Decode intervals are unchanged.
+- Comparisons must fix backend, budget, PCG cap, alignment and requests. Export
+  work and larger forwards cost time/memory; IC is not a universal speedup.
+- Tests cover alignment, fallback, restore and K=1/3/5 rings/partial acceptance;
+  model checks cover K=3 at TP1/DP8/EP8, not full IC accuracy or production stress.
diff --git a/lmdeploy/pytorch/backends/cuda/kda.py b/lmdeploy/pytorch/backends/cuda/kda.py
index f0acfa06a1..a1175aec36 100644
--- a/lmdeploy/pytorch/backends/cuda/kda.py
+++ b/lmdeploy/pytorch/backends/cuda/kda.py
@@ -7,6 +7,7 @@
 """
 
 from copy import copy
+from dataclasses import replace
 from functools import partial
 from typing import Any
 
@@ -106,17 +107,27 @@ def piecewise_forward(mixed_qkv, raw_gate, raw_beta, *args, **kwargs):
         self.forward = piecewise_forward
 
     def _forward_spec(self, mixed_qkv, raw_gate, raw_beta, conv_state,
-                      recurrent_state, metadata, **kwargs):
+                      recurrent_state, metadata, checkpoint=None, **kwargs):
         """Checkpoint every verified token at its accepted-history ring slot.
 
         State is addressed by accepted history length, not the last proposed length. The ring therefore also handles
         zero/partial acceptance and request reordering without a scheduler-side rollback hook.
         """
         if metadata.is_decoding:
+            if checkpoint is not None:
+                raise ValueError('Internal checkpoints are only supported during prefill.')
             return self._forward_spec_decode(mixed_qkv, raw_gate, raw_beta,
                                              conv_state, recurrent_state, metadata, **kwargs)
         history = metadata.cache_seqlens.long()
         ring_size = metadata.num_spec_tokens + 1
+        if checkpoint is not None:
+            if (metadata.state_ids.numel() != 1 or checkpoint.step is None
+                    or checkpoint.step < checkpoint.offset
+                    or checkpoint.recurrent_state_out.shape != (1, *recurrent_state.shape[1:])
+                    or checkpoint.conv_state_out.shape != (1, *conv_state.shape[1:])):
+                raise ValueError('MTP IC requires one checkpoint ring and an absolute save step.')
+            checkpoint = replace(checkpoint,
+                                 recurrent_state_out=checkpoint.recurrent_state_out[:, checkpoint.step % ring_size])
         ids = torch.where(metadata.valid_state, metadata.state_ids, -1).long()
         read_slot = history.remainder(ring_size)
         # FLA prefill consumes a chronological window; the persistent cache
@@ -139,7 +150,7 @@ def store(state, values, lengths):
 
         output = self.forward(mixed_qkv, raw_gate, raw_beta,
                               conv_state=conv, recurrent_state=recurrent,
-                              metadata=local, **kwargs)
+                              metadata=local, checkpoint=checkpoint, **kwargs)
         lengths = history + metadata.cu_seqlens.diff()
         write_offsets = (lengths[:, None] + offsets).remainder(conv_state.size(-1))
         conv_cache.scatter_(2, write_offsets[:, None].expand(-1, conv_cache.size(1), -1), conv)
@@ -254,11 +265,9 @@ def forward(
         checkpoint: KdaPrefillCheckpoint | None = None,
     ) -> torch.Tensor:
         if (getattr(metadata, 'num_spec_tokens', 0) or 0) > 0:
-            if checkpoint is not None:
-                raise ValueError('MTP internal checkpoints are not supported.')
             return self._forward_spec(
                 mixed_qkv, raw_gate, raw_beta, conv_state, recurrent_state,
-                metadata, conv_weight=conv_weight, conv_bias=conv_bias,
+                metadata, checkpoint=checkpoint, conv_weight=conv_weight, conv_bias=conv_bias,
                 a_log=a_log, dt_bias=dt_bias, num_heads=num_heads,
                 head_dim=head_dim, lower_bound=lower_bound)
         if checkpoint is not None:
@@ -273,7 +282,17 @@ def forward(
                     or not checkpoint.recurrent_state_out.is_contiguous()):
                 raise ValueError('Invalid native KDA checkpoint output layout.')
             window = mixed_qkv[:, checkpoint.offset - width:checkpoint.offset].transpose(1, 2)
-            checkpoint.conv_state_out.copy_(window)
+            destination = checkpoint.conv_state_out
+            if destination.size(-1) == width:
+                destination.copy_(window)
+            else:
+                if destination.size(-1) < width or checkpoint.step is None:
+                    raise ValueError('KDA IC convolution ring requires an absolute save step.')
+                start = (checkpoint.step - width) % destination.size(-1)
+                first = min(width, destination.size(-1) - start)
+                destination[..., start:start + first].copy_(window[..., :first])
+                if first < width:
+                    destination[..., :width - first].copy_(window[..., first:])
         batch_size = metadata.state_ids.numel()
         if metadata.is_decoding:
             query_length = mixed_qkv.size(1) // batch_size
diff --git a/lmdeploy/pytorch/backends/kda.py b/lmdeploy/pytorch/backends/kda.py
index 337e02c978..abe988b268 100644
--- a/lmdeploy/pytorch/backends/kda.py
+++ b/lmdeploy/pytorch/backends/kda.py
@@ -16,6 +16,7 @@ class KdaPrefillCheckpoint:
     offsets: torch.Tensor
     conv_state_out: torch.Tensor
     recurrent_state_out: torch.Tensor
+    step: int | None = None
 
 
 class KdaImpl(ABC):
diff --git a/lmdeploy/pytorch/configurations/glm5_next.py b/lmdeploy/pytorch/configurations/glm5_next.py
index a479a021a4..6488e7cf63 100644
--- a/lmdeploy/pytorch/configurations/glm5_next.py
+++ b/lmdeploy/pytorch/configurations/glm5_next.py
@@ -45,11 +45,11 @@ def update_cache_config(cache_config, token_lookahead=0):
     if cache_config.prefix_cache_decode_state_interval % 64:
         raise ValueError('GLM-5.3 prefix_cache_decode_state_interval must be a multiple of 64.')
     if envs.internal_prefill_checkpoint:
-        if (token_lookahead != 0 or cache_config.kv_transfer_config is not None
+        if (token_lookahead not in (0, 1) or cache_config.kv_transfer_config is not None
                 or not cache_config.enable_prefix_caching or cache_config.device_type != 'cuda'
                 or cache_config.prefill_save_alignment % 64
                 or envs.kda_prefill_backend != 'flash_kda'):
-            raise ValueError('GLM IC requires CUDA, local prefix caching, non-MTP AR without a connector, '
+            raise ValueError('GLM IC requires CUDA, local prefix caching, AR/deepseek_mtp without a connector, '
                              '64-aligned prefill_save_alignment and the native FlashKDA backend.')
         from lmdeploy.pytorch.backends.cuda.flash_kda import check_flash_kda
         check_flash_kda()
diff --git a/lmdeploy/pytorch/engine/inputs_maker.py b/lmdeploy/pytorch/engine/inputs_maker.py
index 9dee4c91ce..5de4414a4f 100644
--- a/lmdeploy/pytorch/engine/inputs_maker.py
+++ b/lmdeploy/pytorch/engine/inputs_maker.py
@@ -670,8 +670,8 @@ def _try_internal_prefill(self, result: _ForwardInputsResult) -> bool:
         step = start + result.inputs.max_q_seqlen
         alignment = seq._seq_meta.prefill_save_alignment
         lookahead = seq._seq_meta.prefix_cache_token_lookahead
-        if (maker.spec_decoding or lookahead
-                or not alignment or end % alignment == 0 or step % alignment
+        if ((maker.spec_decoding and lookahead != 1)
+                or not alignment or (not lookahead and end % alignment == 0) or step % alignment
                 or step + lookahead > end
                 or not start < step < end or step - start < 64 or (step - start) % 16
                 or end - start != seq.num_token_ids or end - start > maker.config.max_prefill_token_num
diff --git a/lmdeploy/pytorch/models/glm5_next.py b/lmdeploy/pytorch/models/glm5_next.py
index 2810f8b633..ed84407941 100644
--- a/lmdeploy/pytorch/models/glm5_next.py
+++ b/lmdeploy/pytorch/models/glm5_next.py
@@ -1526,7 +1526,8 @@ def forward(
                     offset=state_capture.checkpoint_step - state_capture.input_start,
                     offsets=capture_offsets,
                     conv_state_out=past_key_value[0][target:target + 1],
-                    recurrent_state_out=past_key_value[1][target:target + 1])
+                    recurrent_state_out=past_key_value[1][target:target + 1],
+                    step=state_capture.checkpoint_step)
             kpool_tail_state = None
             if not layer.is_linear_attention:
                 past_key_value, kpool_tail_state = past_key_value[:-2], past_key_value[-2:]
diff --git a/tests/pytorch/engine/test_internal_prefill_checkpoint.py b/tests/pytorch/engine/test_internal_prefill_checkpoint.py
index ce7b7ef7ca..f309ec1f32 100644
--- a/tests/pytorch/engine/test_internal_prefill_checkpoint.py
+++ b/tests/pytorch/engine/test_internal_prefill_checkpoint.py
@@ -45,7 +45,7 @@ def reserve_chunk(seq, chunk_size, prealloc_size, is_last_chunk):
     return task, result
 
 
-@pytest.mark.parametrize('mtp', [False])
+@pytest.mark.parametrize('mtp', [False, True])
 @pytest.mark.parametrize(('start', 'end', 'step'), [(0, 8193, 8192), (0, 8256, 8192), (0, 10000, 8192),
                                                    (8192, 20000, 16384)])
 def test_internal_capture_preserves_one_reservation_and_restore(start, end, step, mtp):
@@ -135,7 +135,22 @@ def test_completed_tracked_sequence_does_not_block_capture_or_drop_cleanup_delta
     assert result.cache_inputs.state_capture == PrefillStateCapture(8192, 20000, 16384, 1, 3)
 
 
-def test_non_mtp_ic_config_rejects_mtp(monkeypatch):
+def test_mtp_can_merge_aligned_end_without_losing_an_endpoint_checkpoint():
+    task, result = make_task(end=16384, mtp=True)
+    assert task._try_internal_prefill(result)
+    assert result.cache_inputs.state_capture == PrefillStateCapture(0, 16384, 8192, 1, 3)
+    assert result.cache_inputs.state_save_plan is None
+
+
+def test_mtp_capacity_failure_preserves_pending_checkpoint():
+    task, result = make_task(end=16384, mtp=True, reserve=False)
+    original = result.cache_inputs
+    assert not task._try_internal_prefill(result)
+    assert result.cache_inputs is original
+    task.maker.long_context_chunker.clear.assert_not_called()
+
+
+def test_mtp_ic_config_accepts_lookahead_without_enabling_decode_saves(monkeypatch):
     from lmdeploy.pytorch import envs
     from lmdeploy.pytorch.backends.cuda import flash_kda
     from lmdeploy.pytorch.config import CacheConfig
@@ -146,9 +161,11 @@ def test_non_mtp_ic_config_rejects_mtp(monkeypatch):
     monkeypatch.setattr(flash_kda, 'check_flash_kda', Mock())
     config = CacheConfig(max_batches=32, block_size=64, num_cpu_blocks=0, num_gpu_blocks=64,
                          enable_prefix_caching=True, prefill_save_alignment=8192)
-    update_cache_config(config)
+    update_cache_config(config, token_lookahead=1)
     assert config.enable_internal_prefill_checkpoint
-    with pytest.raises(ValueError, match='non-MTP AR'):
+    assert config.prefix_cache_token_lookahead == 1
+    config.prefix_cache_decode_state_interval = 64
+    with pytest.raises(ValueError, match='only publishes prefill'):
         update_cache_config(config, token_lookahead=1)
 
 
diff --git a/tests/pytorch/kernel/test_kda_internal_checkpoint.py b/tests/pytorch/kernel/test_kda_internal_checkpoint.py
index 232009d1cf..2cef5a1df1 100644
--- a/tests/pytorch/kernel/test_kda_internal_checkpoint.py
+++ b/tests/pytorch/kernel/test_kda_internal_checkpoint.py
@@ -64,3 +64,93 @@ def forward(start, end, slot, checkpoint=None, decode=False):
     torch.testing.assert_close(state[3], saved_state, atol=0, rtol=0)
     assert torch.isnan(conv_pool[:, 0]).all() and torch.isnan(state_pool[:, 0]).all()
     assert torch.isnan(conv[0]).all() and torch.isnan(state[0]).all()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA')
+@pytest.mark.parametrize(('draft', 'heads', 'length', 'offset', 'history'),
+                         [(1, 4, 65, 64, 0), (3, 4, 8256, 8192, 0), (5, 4, 8256, 8192, 0),
+                          (3, 4, 11808, 8192, 8192), (5, 64, 8256, 8192, 8192)])
+def test_mtp_internal_checkpoint_ring_restore_and_rejection(monkeypatch, draft, heads, length, offset, history):
+    pytest.importorskip('flash_kda')
+    from lmdeploy.pytorch import envs
+    from lmdeploy.pytorch.backends.cuda.kda import CudaKdaImpl
+    from lmdeploy.pytorch.backends.kda import KdaPrefillCheckpoint
+
+    monkeypatch.setattr(envs, 'kda_prefill_backend', 'flash_kda')
+    impl = CudaKdaImpl()
+    torch.manual_seed(42)
+    width, dim, slots, ring = 4, 128, 7, draft + 1
+    channels = 3 * heads * dim
+    conv_pool = torch.full((slots, 2, channels, width + draft), float('nan'), dtype=torch.bfloat16, device='cuda')
+    state_pool = torch.full((slots, 2, ring, heads, dim, dim), float('nan'), dtype=torch.float32, device='cuda')
+    conv, state = conv_pool[:, 1], state_pool[:, 1]
+    conv[1].zero_()
+    state[1, history % ring].zero_()
+    if history:
+        conv[1].normal_(std=0.1)
+        state[1, history % ring].normal_(std=0.1)
+    conv[2].copy_(conv[1])
+    state[2].copy_(state[1])
+    values = dict(mixed_qkv=torch.randn(1, length, channels, dtype=torch.bfloat16, device='cuda') / 4,
+                  raw_gate=torch.randn(1, length, heads * dim, dtype=torch.bfloat16, device='cuda'),
+                  raw_beta=torch.randn(1, length, heads, dtype=torch.bfloat16, device='cuda'))
+    kwargs = dict(conv_state=conv, recurrent_state=state,
+                  conv_weight=torch.randn(channels, width, dtype=torch.bfloat16, device='cuda') / 8,
+                  conv_bias=None, a_log=torch.zeros(heads, device='cuda'),
+                  dt_bias=torch.zeros(heads * dim, device='cuda'),
+                  num_heads=heads, head_dim=dim, lower_bound=-5.0)
+
+    def forward(inputs, slot, accepted_history, checkpoint=None, decode=False):
+        count = inputs['mixed_qkv'].size(1)
+        metadata = SimpleNamespace(is_decoding=decode, num_spec_tokens=draft, is_init=None,
+                                   state_ids=torch.tensor([slot], device='cuda'),
+                                   valid_state=torch.tensor([True], device='cuda'),
+                                   cache_seqlens=torch.tensor([accepted_history], dtype=torch.int32, device='cuda'),
+                                   cu_seqlens=torch.tensor([0, count], dtype=torch.int64, device='cuda'))
+        return impl.forward(**inputs, metadata=metadata, checkpoint=checkpoint, **kwargs)
+
+    step = history + offset
+    checkpoint = KdaPrefillCheckpoint(offset, torch.tensor([offset], dtype=torch.int64, device='cuda'),
+                                     conv[3:4], state[3:4], step=step)
+    full = forward(values, 1, history, checkpoint)
+    prefix = forward({name: value[:, :offset] for name, value in values.items()}, 2, history)
+    window = (step + torch.arange(-width, 0, device='cuda')) % conv.size(-1)
+    torch.testing.assert_close(conv[3, :, window], conv[2, :, window], atol=0, rtol=0)
+    torch.testing.assert_close(state[3, step % ring], state[2, step % ring], atol=0, rtol=0)
+    untouched = [index for index in range(ring) if index != step % ring]
+    assert torch.isnan(state[3, untouched]).all()
+    conv[4].copy_(conv[3])
+    state[4].copy_(state[3])
+    saved_conv, saved_state = conv[3].clone(), state[3].clone()
+    tail = {name: value[:, offset:] for name, value in values.items()}
+    suffix = forward(tail, 2, step)
+    restored = forward(tail, 4, step)
+    torch.testing.assert_close(full, torch.cat((prefix, suffix), dim=1), atol=0, rtol=0)
+    torch.testing.assert_close(suffix, restored, atol=0, rtol=0)
+    end = history + length
+    torch.testing.assert_close(state[1, end % ring], state[2, end % ring], atol=0, rtol=0)
+    torch.testing.assert_close(state[4, end % ring], state[2, end % ring], atol=0, rtol=0)
+    window = (end + torch.arange(-width, 0, device='cuda')) % conv.size(-1)
+    torch.testing.assert_close(conv[1, :, window], conv[2, :, window], atol=0, rtol=0)
+
+    for slot in (5, 6):
+        conv[slot].copy_(conv[3])
+        state[slot].copy_(state[3])
+    accepted_history = step
+    for accepted in (1, draft, ring, 1):
+        verify = dict(mixed_qkv=torch.randn(1, ring, channels, dtype=torch.bfloat16, device='cuda') / 4,
+                      raw_gate=torch.randn(1, ring, heads * dim, dtype=torch.bfloat16, device='cuda'),
+                      raw_beta=torch.randn(1, ring, heads, dtype=torch.bfloat16, device='cuda'))
+        actual = forward(verify, 5, accepted_history, decode=True)
+        expected = [forward({name: value[:, index:index + 1] for name, value in verify.items()},
+                            6, accepted_history + index, decode=True) for index in range(accepted)]
+        torch.testing.assert_close(actual[:, :accepted], torch.cat(expected, dim=1), rtol=1e-3, atol=1e-3)
+        accepted_history += accepted
+        torch.testing.assert_close(state[5, accepted_history % ring], state[6, accepted_history % ring],
+                                   rtol=1e-4, atol=1e-4)
+        window = (accepted_history + torch.arange(-width, 0, device='cuda')) % conv.size(-1)
+        torch.testing.assert_close(conv[5, :, window], conv[6, :, window], atol=0, rtol=0)
+    torch.testing.assert_close(conv[3], saved_conv, atol=0, rtol=0, equal_nan=True)
+    torch.testing.assert_close(state[3], saved_state, atol=0, rtol=0, equal_nan=True)
+    assert torch.isnan(conv_pool[:, 0]).all() and torch.isnan(state_pool[:, 0]).all()
+    assert torch.isnan(conv[0]).all() and torch.isnan(state[0]).all()