[Feature] GLM-5.3 internal prefill checkpoints

docs/en/advance/internal_checkpoint.md

@@ -0,0 +1,58 @@
+# 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.

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.')
 

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.

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

lmdeploy/pytorch/backends/cuda/kda.py

@@ -7,12 +7,14 @@
 """
 
 from copy import copy
+from dataclasses import replace
 from functools import partial
 from typing import Any
 
 import torch
 
-from lmdeploy.pytorch.backends.kda import KdaImpl
+from lmdeploy.pytorch import envs
+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
@@ -52,6 +54,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
@@ -100,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
@@ -133,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)
@@ -245,13 +262,37 @@ 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:
             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:
+            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)
+            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
@@ -292,6 +333,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, checkpoint)
         else:
             output, final_state = self.chunk_kda(
                 q=q.contiguous(),

lmdeploy/pytorch/backends/kda.py

@@ -8,6 +8,17 @@
 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
+    step: int | None = None
+
+
 class KdaImpl(ABC):
     """Backend interface for Kimi Delta Attention inference."""
 
@@ -27,6 +38,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

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

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 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, 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()
+        cache_config.enable_internal_prefill_checkpoint = True
 
 
 def is_glm5_kda_layer(text_config, layer_idx: int) -> bool:

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:

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