feat(cache): add exact hybrid checkpoints for agent continuations

lightllm/common/basemodel/attention/base_att.py

@@ -66,8 +66,12 @@ def uses_dynamic_spec_verify_layout(self) -> bool:
         draft_step = self.model.mtp_manager.get_decode_draft_step(self.model.is_mtp_draft_model)
         is_main_model = not self.model.is_mtp_draft_model
         has_decode_draft_step = draft_step > 0
-        dynamic_verify_enabled = args.mtp_dynamic_verify
-        return is_main_model and has_decode_draft_step and dynamic_verify_enabled
+        # Exact HEAD replay starts without proposals, so fixed planning can
+        # also mix one-row and full-width requests. Keep the attention layout
+        # service-wide to make CUDA Graph capture and replay use the same shape.
+        exact_head_enabled = getattr(args, "enable_exact_prefix_cache", False) and args.run_mode == "normal"
+        variable_layout_enabled = args.mtp_dynamic_verify or exact_head_enabled
+        return is_main_model and has_decode_draft_step and variable_layout_enabled
 
     def uses_causal_attention(self) -> bool:
         args = get_env_start_args()

lightllm/common/basemodel/basemodel.py

@@ -310,6 +310,55 @@ def _init_custom(self):
     def _init_hidden_collector(self):
         self.hidden_collector_prototype = self.mtp_manager.create_hidden_collector(model=self)
 
+    def supports_exact_output_seed(self) -> bool:
+        """Whether the target's output head has the implemented replay format.
+
+        This only describes the target head. Speculative modes additionally
+        require the proposer's auxiliary-resume capability before a full hit
+        can be admitted. A matching seed alone is not sufficient for MTP.
+        """
+        from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer
+
+        return not self.is_mtp_draft_model and type(self.post_infer) is LlamaPostLayerInfer
+
+    def _capture_output_seed(self, hidden: torch.Tensor, infer_state: InferStateInfo):
+        if not getattr(self.args, "enable_exact_prefix_cache", False) or not self.supports_exact_output_seed():
+            return None
+        if infer_state.is_prefill:
+            last_rows = torch.cumsum(infer_state.b_seq_len - infer_state.b_ready_cache_len, dim=0).long() - 1
+            return hidden.index_select(0, last_rows)
+        return hidden[-infer_state.batch_size :].clone()
+
+    @torch.no_grad()
+    def forward_output_seed(self, output_seed: torch.Tensor, microbatch_index: int = 0) -> ModelOutput:
+        """Run only the normal final norm/head/gather for a batch of exact hits.
+
+        The scheduler must call this on its usual compute stream, in the same
+        collective order on all ranks of this TP group. This method neither
+        touches recurrent/KV state nor samples, updates request lengths, or
+        initializes a drafter. Those remain normal backend responsibilities.
+        """
+        if not self.supports_exact_output_seed():
+            raise NotImplementedError("this model does not implement exact output-seed replay")
+        if output_seed.ndim != 2 or output_seed.shape[1] != self.config["hidden_size"]:
+            raise ValueError("output seed must contain raw final hidden rows for this model")
+        if not output_seed.is_cuda or output_seed.dtype != self.data_type:
+            raise ValueError("output seed must use the model CUDA device and hidden dtype")
+        infer_state = self.infer_state_class()
+        infer_state.dist_group = dist_group_manager.get_group(microbatch_index)
+        batch_size = output_seed.shape[0]
+        if batch_size == 0:
+            vocab_size = self.pre_post_weight.lm_head_weight_.vocab_size
+            return ModelOutput(logits=torch.empty((0, vocab_size), dtype=torch.float32, device=output_seed.device))
+        g_cache_manager.cache_env_in()
+        try:
+            logits = self.post_infer._lm_head_and_gather(
+                output_seed, batch_size, self.pre_post_weight, infer_state
+            ).clone()
+        finally:
+            g_cache_manager.cache_env_out()
+        return ModelOutput(logits=logits)
+
     @torch.no_grad()
     def forward(self, model_input: ModelInput):
         model_input.to_cuda()
@@ -472,14 +521,18 @@ def _create_padded_prefill_model_input(self, model_input: ModelInput, new_handle
 
     def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int):
         padded_batch_size = model_output.logits.shape[0]
-        if padded_batch_size == origin_batch_size:
+        if padded_batch_size == origin_batch_size and model_output.output_seed is None:
             return model_output
         new_model_output = copy.copy(model_output)
         new_model_output.logits = new_model_output.logits[0:origin_batch_size]
         new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode(
             padded_batch_size=padded_batch_size,
             origin_batch_size=origin_batch_size,
         )
+        if model_output.output_seed is not None:
+            # Graph output addresses are overwritten by the next replay even
+            # without padding. Detach them before returning to the backend.
+            new_model_output.output_seed = model_output.output_seed[:origin_batch_size].clone()
         return new_model_output
 
     def _create_unpad_prefill_model_output(
@@ -488,6 +541,8 @@ def _create_unpad_prefill_model_output(
         new_model_output = copy.copy(padded_model_output)
         # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。
         new_model_output.logits = new_model_output.logits[0:origin_batch_size]
+        if padded_model_output.output_seed is not None:
+            new_model_output.output_seed = padded_model_output.output_seed[:origin_batch_size].clone()
         new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill(
             origin_handle_token_num=origin_handle_token_num
         )
@@ -585,9 +640,15 @@ def _decode(
         # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的
         # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。
         infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len)
-        use_cuda_graph = self.graph is not None and self.graph.can_run(
-            batch_size=infer_batch_size,
-            max_len_in_batch=infer_max_kv_seq_len,
+        # Auxiliary single-batch calls cannot replay a graph captured with two
+        # microbatches; normal DP decode uses _microbatch_overlap_decode_cuda.
+        use_cuda_graph = (
+            self.graph is not None
+            and not self.graph.enable_decode_microbatch_overlap
+            and self.graph.can_run(
+                batch_size=infer_batch_size,
+                max_len_in_batch=infer_max_kv_seq_len,
+            )
         )
         need_capture = False
         if use_cuda_graph:
@@ -620,7 +681,6 @@ def _decode(
 
     @final
     def _context_forward(self, infer_state: InferStateInfo):
-
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"
@@ -674,13 +734,15 @@ def prefill_func(input_tensors, _infer_state):
         if infer_state.need_dp_prefill_balance:
             last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs)
 
+        output_seed = self._capture_output_seed(last_input_embs, infer_state)
         predict_logits = self.post_infer.token_forward(last_input_embs, infer_state, self.pre_post_weight)
         hidden_collector = infer_state.hidden_collector
         hidden_collector.add_final_hidden(last_input_embs)
         model_output = ModelOutput(
             logits=predict_logits.contiguous(),
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
             prompt_logics=infer_state.prompt_logics,
+            output_seed=output_seed,
         )
 
         # 在开启使用deepep的时候，需要调用clear_deepep_buffer做资源清理，没有启用的时候
@@ -701,6 +763,7 @@ def _token_forward(self, infer_state: InferStateInfo):
             hidden_collector.add(layer_index=i, hidden=input_embs)
 
         last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
+        output_seed = self._capture_output_seed(last_input_embs, infer_state)
         predict_logits: torch.Tensor = self.post_infer.token_forward(
             last_input_embs, infer_state=infer_state, layer_weight=self.pre_post_weight
         )
@@ -709,6 +772,7 @@ def _token_forward(self, infer_state: InferStateInfo):
         model_output = ModelOutput(
             logits=predict_logits.contiguous(),
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
+            output_seed=output_seed,
         )
 
         # 在 cuda graph 模式下，输出需要转为 no ref tensor, 加强mem pool 的复用，降低显存的使用。
@@ -953,6 +1017,8 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state
             last_input_embs = infer_state._all_to_all_unbalance_get(data=last_input_embs)
             last_input_embs1 = infer_state1._all_to_all_unbalance_get(data=last_input_embs1)
 
+        output_seed = self._capture_output_seed(last_input_embs, infer_state)
+        output_seed1 = self._capture_output_seed(last_input_embs1, infer_state1)
         predict_logits, predict_logits1 = self.post_infer.overlap_tpsp_token_forward(
             last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight
         )
@@ -964,11 +1030,13 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state
             logits=predict_logits.contiguous(),
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
             prompt_logics=infer_state.prompt_logics,
+            output_seed=output_seed,
         )
         model_output1 = ModelOutput(
             logits=predict_logits1.contiguous(),
             mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1),
             prompt_logics=infer_state1.prompt_logics,
+            output_seed=output_seed1,
         )
 
         return model_output, model_output1
@@ -1003,6 +1071,8 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1:
         last_input_embs = self.post_infer._tpsp_allgather(input=input_embs, infer_state=infer_state)
         last_input_embs1 = self.post_infer._tpsp_allgather(input=input_embs1, infer_state=infer_state1)
 
+        output_seed = self._capture_output_seed(last_input_embs, infer_state)
+        output_seed1 = self._capture_output_seed(last_input_embs1, infer_state1)
         predict_logits, predict_logits1 = self.post_infer.overlap_tpsp_token_forward(
             last_input_embs, last_input_embs1, infer_state, infer_state1, self.pre_post_weight
         )
@@ -1012,10 +1082,12 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1:
         model_output = ModelOutput(
             logits=predict_logits.contiguous(),
             mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state),
+            output_seed=output_seed,
         )
         model_output1 = ModelOutput(
             logits=predict_logits1.contiguous(),
             mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1),
+            output_seed=output_seed1,
         )
 
         if infer_state.is_cuda_graph:

lightllm/common/basemodel/batch_objs.py

@@ -200,10 +200,19 @@ class ModelOutput:
     # 需要返回 prompt logprobs 信息时才会非空。
     prompt_logics: Optional[torch.Tensor] = None
 
+    # Exact-prefix replay seed: one raw final hidden per logits row, before
+    # final norm. Public model.forward() returns independently owned storage;
+    # graph-internal outputs are cloned when leaving the graph replay wrapper.
+    # This is distinct from spec_hidden, which may contain intermediate layers
+    # and may be normalized in-place by an MTP draft model.
+    output_seed: Optional[torch.Tensor] = None
+
     def __post_init__(self) -> None:
         if self.mtp_collector is None:
             self.mtp_collector = ModelMtpOutputCollector()
 
     def to_no_ref_tensor(self):
         self.logits = tensor_to_no_ref_tensor(self.logits)
         self.mtp_collector.to_no_ref_tensor()
+        if self.output_seed is not None:
+            self.output_seed = tensor_to_no_ref_tensor(self.output_seed)

lightllm/common/basemodel/triton_kernel/linear_att/capture_state.py

@@ -0,0 +1,164 @@
+"""Freeze a bounded number of recurrent states without a GPU-to-CPU decision.
+
+The output buffers are owned by the caller and must not be reused while a
+checkpoint transfer still reads them. Candidate selection is stable across TP
+ranks: logical input order, rather than the order of GPU atomics, chooses slots.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit(do_not_specialize=["N"])
+def _select_capture_candidates(
+    req_indices,
+    state_rows,
+    exact_lengths,
+    capture_mask,
+    selected_reqs,
+    selected_rows,
+    selected_lengths,
+    source_rows,
+    N,
+    CAPACITY: tl.constexpr,
+    MAX_REQS: tl.constexpr,
+    MTP_SIZE: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    index = tl.arange(0, BLOCK)
+    tl.store(selected_reqs + index, -1, index < CAPACITY)
+    tl.store(selected_rows + index, 0, index < CAPACITY)
+    tl.store(selected_lengths + index, 0, index < CAPACITY)
+    tl.store(source_rows + index, -1, index < CAPACITY)
+    req = tl.load(req_indices + index, index < N, other=-1)
+    row = tl.load(state_rows + index, index < N, other=-1)
+    length = tl.load(exact_lengths + index, index < N, other=0)
+    enabled = tl.load(capture_mask + index, index < N, other=0)
+    valid = enabled & (req >= 0) & (req < MAX_REQS) & (row >= 0) & (row < MTP_SIZE) & (length > 0)
+    slot = tl.cumsum(valid.to(tl.int32)) - 1
+    keep = valid & (slot < CAPACITY)
+    tl.debug_barrier()
+    tl.store(selected_reqs + slot, req, keep)
+    tl.store(selected_rows + slot, row, keep)
+    tl.store(selected_lengths + slot, length, keep)
+    tl.store(source_rows + slot, index, keep)
+
+
+@triton.jit
+def _freeze_linear_states(
+    conv,
+    ssm,
+    selected_reqs,
+    selected_rows,
+    out_conv,
+    out_ssm,
+    CONV_LAYER_STRIDE: tl.constexpr,
+    CONV_REQ_STRIDE: tl.constexpr,
+    CONV_DIM_STRIDE: tl.constexpr,
+    CONV_WIDTH_STRIDE: tl.constexpr,
+    SSM_LAYER_STRIDE: tl.constexpr,
+    SSM_REQ_STRIDE: tl.constexpr,
+    LAYERS: tl.constexpr,
+    CONV_WIDTH: tl.constexpr,
+    CONV_ELEMENTS: tl.constexpr,
+    SSM_ELEMENTS: tl.constexpr,
+    MTP_SIZE: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    slot, layer, block = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    req = tl.load(selected_reqs + slot)
+    if req < 0:
+        return
+    row = tl.load(selected_rows + slot)
+    index = block * BLOCK + tl.arange(0, BLOCK)
+    conv_source = (
+        layer * CONV_LAYER_STRIDE
+        + req * CONV_REQ_STRIDE
+        + (index // CONV_WIDTH) * CONV_DIM_STRIDE
+        + (row + index % CONV_WIDTH) * CONV_WIDTH_STRIDE
+    )
+    conv_value = tl.load(conv + conv_source, index < CONV_ELEMENTS, other=0)
+    tl.store(out_conv + (slot * LAYERS + layer) * CONV_ELEMENTS + index, conv_value, index < CONV_ELEMENTS)
+    ssm_source = layer * SSM_LAYER_STRIDE + (req * MTP_SIZE + row) * SSM_REQ_STRIDE + index
+    ssm_value = tl.load(ssm + ssm_source, index < SSM_ELEMENTS, other=0)
+    tl.store(out_ssm + (slot * LAYERS + layer) * SSM_ELEMENTS + index, ssm_value, index < SSM_ELEMENTS)
+
+
+def freeze_linear_states(
+    conv: torch.Tensor,
+    ssm: torch.Tensor,
+    req_indices: torch.Tensor,
+    state_rows: torch.Tensor,
+    exact_lengths: torch.Tensor,
+    capture_mask: torch.Tensor,
+    selected_reqs: torch.Tensor,
+    selected_rows: torch.Tensor,
+    selected_lengths: torch.Tensor,
+    source_rows: torch.Tensor,
+    out_conv: torch.Tensor,
+    out_ssm: torch.Tensor,
+    mtp_size: int,
+) -> None:
+    """Gather at most ``out_conv.shape[0]`` masked candidates into owned slots.
+
+    ``state_rows`` contains request-local MTP row numbers, not flattened model
+    output rows. A row r is the state *after processing input row r*, before the
+    token sampled from that row. The caller supplies that exact prefix length.
+    No candidate metadata is read back to the CPU by this function.
+    """
+    count = req_indices.numel()
+    capacity, layers, conv_dim, conv_width = out_conv.shape
+    assert capacity > 0 and mtp_size > 0
+    for tensor in (req_indices, state_rows, exact_lengths, capture_mask):
+        assert tensor.is_cuda and tensor.ndim == 1 and tensor.numel() == count and tensor.is_contiguous()
+        assert tensor.device == conv.device
+    for tensor in (selected_reqs, selected_rows, selected_lengths, source_rows):
+        assert tensor.is_cuda and tensor.shape == (capacity,) and tensor.is_contiguous()
+        assert tensor.device == conv.device and tensor.dtype in (torch.int32, torch.int64)
+    assert req_indices.dtype in (torch.int32, torch.int64)
+    assert state_rows.dtype in (torch.int32, torch.int64)
+    assert exact_lengths.dtype in (torch.int32, torch.int64)
+    assert capture_mask.dtype == torch.bool
+    assert conv.is_cuda and ssm.is_cuda and conv.device == ssm.device
+    assert out_conv.device == conv.device and out_ssm.device == conv.device
+    assert conv.ndim == 4 and conv.shape[0] == layers and conv.shape[2] == conv_dim
+    assert conv.shape[-1] == conv_width + mtp_size - 1
+    assert ssm.shape[0] == layers and ssm.shape[1] == conv.shape[1] * mtp_size
+    assert ssm.is_contiguous() and out_conv.is_contiguous() and out_ssm.is_contiguous()
+    assert out_ssm.shape == (capacity, layers, *ssm.shape[2:])
+    assert conv.dtype == out_conv.dtype and ssm.dtype == out_ssm.dtype
+    _select_capture_candidates[(1,)](
+        req_indices,
+        state_rows,
+        exact_lengths,
+        capture_mask,
+        selected_reqs,
+        selected_rows,
+        selected_lengths,
+        source_rows,
+        N=count,
+        CAPACITY=capacity,
+        MAX_REQS=conv.shape[1] - 1,  # The final request slot is graph padding.
+        MTP_SIZE=mtp_size,
+        BLOCK=triton.next_power_of_2(max(count, capacity)),
+    )
+    ssm_elements = ssm[0, 0].numel()
+    conv_elements = conv_dim * conv_width
+    _freeze_linear_states[(capacity, layers, triton.cdiv(max(conv_elements, ssm_elements), 256))](
+        conv,
+        ssm,
+        selected_reqs,
+        selected_rows,
+        out_conv,
+        out_ssm,
+        *conv.stride(),
+        ssm.stride(0),
+        ssm.stride(1),
+        LAYERS=layers,
+        CONV_WIDTH=conv_width,
+        CONV_ELEMENTS=conv_elements,
+        SSM_ELEMENTS=ssm_elements,
+        MTP_SIZE=mtp_size,
+        BLOCK=256,
+    )