[DSv4.1] Run encoder SWA bounded replay inside the breakable prefill CUDA graph

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -280,9 +280,12 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
                     or (get_platform().is_hip and is_gfx95_supported())
                 ),
             ),
+            # Breakable is admitted: the fold's floor, replay-row skip and window
+            # layout are bucket-sized metadata the replay refreshes in place.
             (
-                "prefill CUDA graphs",
-                cfg.cuda_graph_config.prefill.backend != Backend.DISABLED,
+                "prefill CUDA graphs other than breakable",
+                cfg.cuda_graph_config.prefill.backend
+                not in (Backend.DISABLED, Backend.BREAKABLE),
             ),
             ("context parallelism", cfg.attn_cp_size > 1),
             ("external cache linker", cfg.enable_unified_cache_external_linker),

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

@@ -705,6 +705,9 @@ def refresh_for_breakable_cuda_graph_replay_(self, other: DSV4AttnMetadata) -> N
         # the field is produced in-graph before the attention graph break reads it.
         for field_name in reference_assign_fields:
             setattr(self, field_name, getattr(other, field_name))
+        if self.request_window_layout is not None:
+            # Captured window copy/commit kernels read the layout by address.
+            self.request_window_layout.copy_(other.request_window_layout)
 
     def init_compression_metadata(
         self, num_tokens: Optional[int] = None, low_ratio_buffers=None
@@ -976,6 +979,9 @@ class DSV4Metadata:
     # Shared by all low-ratio source layers; graph replay refreshes them live.
     low_ratio_req_indices: Optional[torch.Tensor] = None
     low_ratio_pos_i64: Optional[torch.Tensor] = None
+    # Folded encoder-replay rows the ratio-1/2 compressors skip under the prefill
+    # graph (their groups are cached already), [num_q_tokens] bool.
+    low_ratio_replay_rows: Optional[torch.Tensor] = None
 
     # Per-step scratch for TP-padded query heads, zeroed by the first user.
     # Later layers overwrite real heads and preserve the zero padding.
@@ -1029,6 +1035,9 @@ def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata
         maybe_copy_inplace(
             self.low_ratio_pos_i64, src=static_metadata.low_ratio_pos_i64
         )
+        maybe_copy_inplace(
+            self.low_ratio_replay_rows, src=static_metadata.low_ratio_replay_rows
+        )
         maybe_copy_inplace(
             self.c4_compress_metadata, src=static_metadata.c4_compress_metadata
         )
@@ -1446,7 +1455,13 @@ def init_forward_metadata_prefill(
             dspark_swa_buffers=dspark_swa_buffers,
             num_tokens=num_write_tokens if cp_active else None,
             swa_replay_start=swa_replay_start,
-            num_groups=len(extend_seq_lens_cpu),
+            num_groups=(
+                # One captured window layout serves every batch size of a bucket.
+                self.req_to_token_pool.size
+                if use_prefill_cuda_graph
+                and self.token_to_kv_pool.request_window is not None
+                else len(extend_seq_lens_cpu)
+            ),
         )
         if cp_active:
             core_attn_metadata.apply_cp_reindex(
@@ -1524,8 +1539,26 @@ def create(compress_ratio: Literal[4, 128]):
             metadata.low_ratio_pos_i64 = core_attn_metadata.positions_casual.to(
                 torch.int64
             )
+            if self.token_to_kv_pool.request_window is not None:
+                metadata.low_ratio_replay_rows = self._encoder_replay_rows(
+                    forward_batch,
+                    num_tokens=num_tokens,
+                    num_rows=out_cache_loc.shape[0],
+                )
         return metadata
 
+    def _encoder_replay_rows(
+        self, forward_batch: Optional[ForwardBatch], *, num_tokens: int, num_rows: int
+    ) -> torch.Tensor:
+        rows = torch.zeros(num_rows, dtype=torch.bool, device=self.device)
+        keep = (
+            None if forward_batch is None else forward_batch.encoder_swa_compress_rows
+        )
+        if keep is not None:
+            rows[:num_tokens] = True
+            rows[keep] = False
+        return rows
+
     def _low_ratio_prefill_indexer_metadata(
         self, core: DSV4AttnMetadata, compress_ratio: int
     ) -> PagedIndexerMetadata:
@@ -2822,11 +2855,14 @@ def init_forward_metadata_for_breakable_cuda_graph_capture(
         self, forward_batch: ForwardBatch
     ):
         max_seq_len = forward_batch.max_seq_len_override or self.MAX_SEQ_LEN_FOR_CAPTURE
+        self.encoder_replay = forward_batch.encoder_swa_replay
+        self.encoder_row_floor = forward_batch.encoder_swa_row_floor
         self.forward_metadata = self._build_forward_metadata(
             forward_batch,
             max_seq_len_override=max_seq_len,
             use_prefill_cuda_graph=True,
         )
+        self._activate_request_window(self.forward_metadata, reserve=True)
         with_tail = (
             self.enable_decoder_swa_bounded_replay
             and forward_batch.forward_mode.is_extend_without_speculative()
@@ -2886,6 +2922,8 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay(
         max_seq_len = (
             metadata_batch.max_seq_len_override or self.MAX_SEQ_LEN_FOR_CAPTURE
         )
+        self.encoder_replay = metadata_batch.encoder_swa_replay
+        self.encoder_row_floor = metadata_batch.encoder_swa_row_floor
         static_metadata = self._build_forward_metadata(
             metadata_batch,
             max_seq_len_override=max_seq_len,
@@ -2894,12 +2932,26 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay(
         assert isinstance(capture_metadata, DSV4Metadata)
         capture_metadata.refresh_for_breakable_cuda_graph_replay_(static_metadata)
         self.forward_metadata = capture_metadata
+        self._activate_request_window(capture_metadata)
         self.tail_forward_metadata = (
             self._prefill_graph_tail_metadata(metadata_batch, max_seq_len)
             if capture_metadata.prefill_graph_tail
             else None
         )
 
+    def _activate_request_window(
+        self, metadata: DSV4Metadata, *, reserve: bool = False
+    ) -> None:
+        window = self.token_to_kv_pool.request_window
+        if window is None:
+            return
+        layout = metadata.core_attn_metadata.request_window_layout
+        if reserve:
+            # Prefill graphs capture first and largest bucket first, so no graph
+            # holds the workspace address yet.
+            window.reserve_workspace(layout.size)
+        window.activate(layout)
+
     def _prefill_graph_tail_metadata(
         self, forward_batch: ForwardBatch, max_seq_len: int
     ) -> DSV4Metadata:
@@ -3053,7 +3105,12 @@ def forward_low_ratio_sources(
             _bcg_low_ratio_source_projections(layer, x, q_lora, pos, bufs)
             if run_compressor and layer.compressor is not None:
                 self._low_ratio_compress_torch(
-                    layer, x, req, pos, projected=(bufs["kv"], bufs.get("score"))
+                    layer,
+                    x,
+                    req,
+                    pos,
+                    projected=(bufs["kv"], bufs.get("score")),
+                    skip_rows=meta.low_ratio_replay_rows,
                 )
             if run_indexer and layer.indexer is not None:
                 self._low_ratio_index_topk_captured(layer, bufs["q"], bufs["w"])
@@ -3314,18 +3371,33 @@ def _low_ratio_compress_fused(
             )
 
     def _low_ratio_compress_torch(
-        self, layer, x, req, pos, projected=None, *, fuse_index_store=False
+        self,
+        layer,
+        x,
+        req,
+        pos,
+        projected=None,
+        *,
+        fuse_index_store=False,
+        skip_rows: Optional[torch.Tensor] = None,
     ) -> None:
         core = self.forward_metadata.core_metadata
         num_tokens = pos.shape[0]
         kv, score = projected if projected is not None else layer.compressor.project(x)
         if not num_tokens:
             return
+        # Skipped rows act as padding: no ring state, writes land on the dummy slot.
+        pad = core.raw_out_loc[:num_tokens] == 0
+        if skip_rows is not None:
+            pad = pad | skip_rows[:num_tokens]
         if layer.compress_ratio == 1:
+            slots = core.c1_out_loc[:num_tokens]
+            if skip_rows is not None:
+                slots = torch.where(pad, torch.zeros_like(slots), slots)
             self._low_ratio_write_group(
                 layer,
                 kv,
-                core.c1_out_loc[:num_tokens],
+                slots,
                 pos,
                 fuse_index_store=fuse_index_store,
             )
@@ -3337,15 +3409,16 @@ def _low_ratio_compress_torch(
             score=score,
             req=req,
             pos=pos,
-            pad=core.raw_out_loc[:num_tokens] == 0,
+            pad=pad,
         )
         pooled = layer.compressor.pool_pairs(
             torch.stack([partner_kv, kv], dim=1),
             torch.stack([partner_score, score], dim=1),
         )
         group_pos = torch.where(pos % 2 == 1, pos - 1, pos)
         out_loc = core.c2_out_loc[:num_tokens]
-        slots = torch.where(out_loc >= 0, out_loc, torch.zeros_like(out_loc))
+        writable = out_loc >= 0 if skip_rows is None else (out_loc >= 0) & ~pad
+        slots = torch.where(writable, out_loc, torch.zeros_like(out_loc))
         self._low_ratio_write_group(
             layer, pooled, slots, group_pos, fuse_index_store=fuse_index_store
         )
@@ -4308,12 +4381,14 @@ def make_core_attn_metadata(
                 swa_replay_start = _pad_rows(
                     self.encoder_row_floor, num_rows=raw_positions.shape[0]
                 )
+            # Padding rows hold the dummy slot and take the spare group.
             request_layout = window_layout(
                 req_pool_indices_repeated,
                 raw_positions,
                 capacity=self.token_to_kv_pool.request_window.capacity,
                 floor=swa_replay_start,
                 num_groups=num_groups,
+                pad=out_loc == 0,
             )
             swa_page_indices = _pad_last_dim(request_layout.indices)
             swa_topk_lengths = request_layout.lengths

python/sglang/srt/managers/scheduler_components/dp_attn.py

@@ -389,6 +389,7 @@ def _local_prefill_cuda_graph_vote(
     breakable_prefill: bool,
     spec_algorithm: SpeculativeAlgorithm,
     model_config,
+    folds_encoder_swa_replay: bool = False,
 ) -> bool:
     """This rank's vote for the prefill graph (min-reduced across dp
     ranks). Extend and mixed batches share the runner's rank-local replay
@@ -402,6 +403,9 @@ def _local_prefill_cuda_graph_vote(
     mode = local_batch.forward_mode
     if mode in (ForwardMode.EXTEND, ForwardMode.MIXED):
         num_tokens = local_batch.extend_num_tokens
+        if folds_encoder_swa_replay:
+            # The forward replays the folded rows, so the bucket vote counts them.
+            num_tokens += folded_replay_rows(local_batch)
         input_embeds = local_batch.input_embeds
         replace_embeds = local_batch.replace_embeds
         prefix_lens = local_batch.prefix_lens
@@ -577,6 +581,12 @@ def prepare_mlp_sync_batch_raw(
         breakable_prefill=breakable_prefill,
         spec_algorithm=model_runner.spec_algorithm,
         model_config=model_runner.model_config,
+        folds_encoder_swa_replay=(
+            local_batch is not None
+            and local_batch.forward_mode.is_extend()
+            and local_batch.encoder_swa_reset is not None
+            and model_runner.attn_backend.folds_encoder_swa_replay
+        ),
     )
 
     is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False

python/sglang/srt/mem_cache/dsv41_request_window.py

@@ -65,16 +65,20 @@ def window_layout(
     capacity: int = 256,
     floor: Optional[torch.Tensor] = None,
     num_groups: Optional[int] = None,
+    pad: Optional[torch.Tensor] = None,
 ):
     n = pos.numel()
     if n == 0:
         raise ValueError("request-window layout needs at least one query")
     req = req.to(torch.int64)
     pos = pos.to(torch.int64)
     device = pos.device
-    groups = n if num_groups is None else int(num_groups)
+    # Padding rows (DP or graph-bucket fill) take a spare group behind the live
+    # requests, so they never shorten a request's commit range or commit themselves.
+    groups = n if num_groups is None else int(num_groups) + (pad is not None)
     offset = torch.arange(n, device=device)
-    group, group_first, group_last = _first_row_offsets(req)
+    group_req = req if pad is None else torch.where(pad, -1, req)
+    group, group_first, group_last = _first_row_offsets(group_req)
     first_pos = pos - (offset - group_first)
     history_rows = groups * window
 
@@ -111,6 +115,8 @@ def window_layout(
     history_loc = torch.arange(history_rows, device=device)
 
     commit_mask = (group_last - offset) < capacity
+    if pad is not None:
+        commit_mask &= ~pad
     return WindowLayout(
         req,
         pos,
@@ -176,6 +182,10 @@ def _ensure_workspace(self, rows: int) -> None:
         size = ((rows + self.page_size - 1) // self.page_size) * self.page_size
         self.workspace = self.pool_factory(size, 1)
 
+    def reserve_workspace(self, rows: int) -> None:
+        """Grow the workspace for a graph layout; only before any graph holds it."""
+        self._ensure_workspace(rows)
+
     def reset(self, slots):
         loc = slots.to(torch.int64)[:, None] * self.capacity + torch.arange(
             self.capacity, device=slots.device

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -1994,6 +1994,15 @@ def _slot(name):
         # The n-gram hasher runs outside the graph and reads this at replay.
         static_forward_batch.ngram_embedding_info = forward_batch.ngram_embedding_info
         static_forward_batch.engram_history = forward_batch.engram_history
+        # Encoder-SWA fold fields: eager breaks read them off this batch view.
+        static_forward_batch.encoder_swa_replay = forward_batch.encoder_swa_replay
+        static_forward_batch.encoder_swa_row_floor = forward_batch.encoder_swa_row_floor
+        static_forward_batch.encoder_swa_compress_skip = (
+            forward_batch.encoder_swa_compress_skip
+        )
+        static_forward_batch.encoder_swa_compress_rows = (
+            forward_batch.encoder_swa_compress_rows
+        )
         static_forward_batch = self.model_runner.prepare_dummy_forward_batch(
             static_forward_batch
         )

test/registered/unit/managers/scheduler_components/test_dp_attn.py

@@ -497,6 +497,40 @@ def test_extend_batch_votes_for_prefill_graph(self):
         runner.can_replay_locally.assert_called_once()
         self.assertFalse(runner.can_replay_locally.call_args.kwargs["is_mixed"])
 
+    def test_vote_counts_folded_encoder_replay_rows(self):
+        """A backend that folds the encoder-SWA replay forwards the replay rows
+        with the extend, so the bucket vote must see them: voting on the scheduled
+        rows alone let a rank pick the graph for a fold the runner cannot replay."""
+        runner = Mock(spec=dp_attn.PrefillCudaGraphRunner)
+        runner.enable_lora = False
+        runner.max_context_size = runner.max_seq_len = None
+        runner._qwen_bcg_hc_sidechannel = False
+        runner.can_replay_locally.return_value = True
+        batch = SimpleNamespace(
+            forward_mode=ForwardMode.EXTEND,
+            extend_num_tokens=10,
+            encoder_swa_reset=[True, False],
+            prefix_lens=[512, 64],
+            multimodal_inputs=None,
+            input_embeds=None,
+            replace_embeds=None,
+            return_logprob=False,
+            batch_size=lambda: 2,
+        )
+        for folds, expected in ((True, 10 + 128), (False, 10)):
+            dp_attn._local_prefill_cuda_graph_vote(
+                local_batch=batch,
+                prefill_graph_runner=runner,
+                coordinated_prefill=True,
+                breakable_prefill=True,
+                spec_algorithm=SpeculativeAlgorithm.NONE,
+                model_config=object(),
+                folds_encoder_swa_replay=folds,
+            )
+            self.assertEqual(
+                runner.can_replay_locally.call_args.kwargs["num_tokens"], expected
+            )
+
     def test_mixed_batch_delegates_to_runner_policy(self):
         vote, runner = self._vote(ForwardMode.MIXED)