[DSv4.1] Replay the trimmed late layers from CUDA graphs under decoder SWA bounded replay

python/sglang/srt/environ.py

@@ -1583,6 +1583,14 @@ class Envs:
     # Quantize the SWA fp8 KV cache from bf16-rounded values (matches
     # trainer-side QAT and the DSA-CP path) instead of fp32 registers.
     SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE = EnvBool(False)
+    # Decoder SWA bounded replay: eager prefill steps run the trimmed late layers
+    # from CUDA graphs captured per tail-row bucket (multiples of 128) up to this
+    # many rows; 0 keeps them eager. 2048 rows (16 buckets) took 0.11 GiB per GPU
+    # on GB300 TP4.
+    SGLANG_DSV4_DECODER_REPLAY_GRAPH_MAX_ROWS = EnvInt(2048)
+    # Debug only: replay decoder replay graphs one segment at a time with a sync,
+    # naming the segment or break that faults.
+    SGLANG_DSV4_DECODER_REPLAY_GRAPH_DEBUG = EnvBool(False)
     # Paged KV layout of the DeepSeek-V4 family pools: "v4" (584 B/token, every
     # GPU), "v41" (the SM100 FlashMLA V4.1 formats: 528 B fp8 SWA cache, fp8 or
     # fp4 compressed caches) or "auto" (v41 on SM100, v4 elsewhere).

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -2896,8 +2896,16 @@ def _low_ratio_in_prefill_graph(self) -> bool:
         from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
             is_in_breakable_cuda_graph,
         )
+        from sglang.srt.models.deepseek_v4_replay_graphs import (
+            in_decoder_replay_graph,
+        )
 
-        return self.low_ratio_prefill_graph and is_in_breakable_cuda_graph()
+        # Decoder replay graphs run the low-ratio sources as an eager break.
+        return (
+            self.low_ratio_prefill_graph
+            and is_in_breakable_cuda_graph()
+            and not in_decoder_replay_graph()
+        )
 
     def _low_ratio_compress_decode(self, layer, x, req, pos) -> None:
         # Projection layout and fused-write support are fixed together at load time.

python/sglang/srt/models/deepseek_v4.py

@@ -178,6 +178,11 @@
     _is_npu,
     _is_xpu,
 )
+from sglang.srt.models.deepseek_v4_replay_graphs import (
+    DecoderReplayGraphs,
+    bcg_late_kv_store,
+    in_decoder_replay_graph,
+)
 from sglang.srt.models.deepseek_v41_vit import Aligner, ViT
 from sglang.srt.models.utils import WeightsMapper
 from sglang.srt.multimodal.deepseek_v41_image_processing import (
@@ -2092,6 +2097,10 @@ def _forward_prepare(
                     positions=global_positions,
                 )
                 kv = None
+            elif in_decoder_replay_graph():
+                assert not fuse_q_rope
+                bcg_late_kv_store(self, x_linear, positions, qkv_a)
+                kv = None
             else:
                 self._compute_kv_to_cache(
                     x_linear,
@@ -2111,7 +2120,10 @@ def _forward_prepare(
             if (
                 forward_batch.forward_mode.is_extend()
                 and is_in_breakable_cuda_graph()
-                and not getattr(attn_backend, "low_ratio_prefill_graph", False)
+                and (
+                    not getattr(attn_backend, "low_ratio_prefill_graph", False)
+                    or in_decoder_replay_graph()
+                )
             ):
                 bcg_deepseek_v4_low_ratio_sources(self, x, q_lora, positions)
             else:
@@ -4068,6 +4080,12 @@ def __init__(
                 0,
                 1,
             }, f"late layers must not compress on their own, got ratios {late_ratios}"
+        self.decoder_replay_graphs: Optional[DecoderReplayGraphs] = None
+        replay_graph_rows = envs.SGLANG_DSV4_DECODER_REPLAY_GRAPH_MAX_ROWS.get()
+        if self.late_layer_start is not None and replay_graph_rows > 0 and not _is_hip:
+            self.decoder_replay_graphs = DecoderReplayGraphs(
+                model=self, run_layers=self._run_late_layers, max_rows=replay_graph_rows
+            )
 
     def get_input_embeddings(self) -> nn.Module:
         return self.embed_tokens
@@ -4108,6 +4126,50 @@ def hc_head(
             hc_eps=self.hc_eps,
         )
 
+    def _replays_late_layers(
+        self, tail: LateLayerTail, capture_dspark: bool, forward_batch: ForwardBatch
+    ) -> bool:
+        # Prefill graph steps keep the late layers inside their own graph.
+        graphs = self.decoder_replay_graphs
+        return (
+            graphs is not None
+            and graphs.ready(forward_batch.input_ids.shape[0])
+            and not capture_dspark
+            and tail.cp_metadata is None
+            and not is_in_breakable_cuda_graph()
+            and not get_is_capture_mode()
+            and graphs.bucket_rows(tail.positions.shape[0]) is not None
+        )
+
+    def _run_late_layers(
+        self,
+        state: mhc.HcState,
+        *,
+        positions: torch.Tensor,
+        input_ids: torch.Tensor,
+        input_ids_global: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """The late layers on the tail rows: the decoder replay graphs' body."""
+        for i in range(self.late_layer_start, self.end_layer):
+            assert self.layers[i].engram is None
+            ctx = (
+                nullcontext()
+                if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
+                else get_global_expert_distribution_recorder().with_current_layer(i)
+            )
+            with ctx:
+                state = self.layers[i].forward_hc_pre_from_prev(
+                    positions=positions,
+                    state=state,
+                    input_ids=input_ids,
+                    forward_batch=forward_batch,
+                    input_ids_global=input_ids_global,
+                    seam_open=False,
+                )
+        state = state.materialized(self.layers[self.end_layer - 1].hc_cfg)
+        return state.residual, state.pre
+
     def _check_late_layer_tail_readers(self, forward_batch: ForwardBatch) -> None:
         # Rows outside the tail are never computed past the last kv_source layer.
         if (
@@ -4209,6 +4271,16 @@ def _forward_layers_hc_pre_from_prev(
                 positions = tail.positions
                 if hash_ids is not None:
                     hash_ids = tail.rows(hash_ids)
+                if self._replays_late_layers(tail, capture_dspark, forward_batch):
+                    residual, pre = self.decoder_replay_graphs.run(
+                        state=state,
+                        positions=positions,
+                        input_ids=input_ids,
+                        input_ids_global=input_ids_global,
+                        forward_batch=forward_batch,
+                    )
+                    attn_backend.exit_late_layer_tail(saved_full, forward_batch)
+                    return residual, pre, tail
             # late_dp_counts is only not None on hip
             if late_dp_counts is not None and i == self.late_layer_start:
                 saved_dp = _enter_late_layer_dp(forward_batch, late_dp_counts)