diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 2a02dc21465d..379fcfba02a5 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/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), diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 5eb17b471c8f..668b01018d81 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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,7 +3409,7 @@ 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), @@ -3345,7 +3417,8 @@ def _low_ratio_compress_torch( ) 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 diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index aeb21bbb09d8..81cd4d305963 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/dsv41_request_window.py b/python/sglang/srt/mem_cache/dsv41_request_window.py index 5180aec572c1..a0274e37bb74 100644 --- a/python/sglang/srt/mem_cache/dsv41_request_window.py +++ b/python/sglang/srt/mem_cache/dsv41_request_window.py @@ -65,6 +65,7 @@ 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: @@ -72,9 +73,12 @@ def window_layout( 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 diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 119997e5bf88..029a4173ae93 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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 ) diff --git a/test/registered/unit/managers/scheduler_components/test_dp_attn.py b/test/registered/unit/managers/scheduler_components/test_dp_attn.py index 11ba81d228f5..d0c8e70a5bce 100644 --- a/test/registered/unit/managers/scheduler_components/test_dp_attn.py +++ b/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) diff --git a/test/registered/unit/model_executor/test_encoder_swa_replay.py b/test/registered/unit/model_executor/test_encoder_swa_replay.py index 0d953d241ec6..ae6d15b22d92 100644 --- a/test/registered/unit/model_executor/test_encoder_swa_replay.py +++ b/test/registered/unit/model_executor/test_encoder_swa_replay.py @@ -219,6 +219,7 @@ def test_tail_keeps_its_own_floor_under_fold(self): capacity=2 * SWA_WINDOW, floor=floor, num_groups=3, + pad=torch.zeros(req.shape[0], dtype=torch.bool), ) layout = metadata.core_attn_metadata.request_window_layout torch.testing.assert_close(layout.lengths, expected.lengths) @@ -272,6 +273,133 @@ def test_tail_hidden_rows_map_to_scheduled_rows(self): ) +class TestWindowLayoutPadding(CustomTestCase): + def test_padding_rows_group_apart_and_never_commit(self): + """Padding rows (DP fill, prefill graph bucket) carry the last request's + slot at position 0. Grouped with it, they shortened its commit range by + their count and committed garbage into its ring slot 0.""" + live, fill = 300, 100 + req = torch.full((live + fill,), 7, dtype=torch.int64) + pos = torch.cat([torch.arange(live), torch.zeros(fill, dtype=torch.int64)]) + pad = torch.cat( + [torch.zeros(live, dtype=torch.bool), torch.ones(fill, dtype=torch.bool)] + ) + layout = window_layout(req, pos, capacity=256, num_groups=1, pad=pad) + alone = window_layout(req[:live], pos[:live], capacity=256, num_groups=2) + + self.assertFalse(layout.commit_mask[live:].any()) + torch.testing.assert_close(layout.commit_mask[:live], alone.commit_mask) + self.assertEqual(int(layout.commit_mask.sum()), 256) + torch.testing.assert_close(layout.indices[:live], alone.indices) + torch.testing.assert_close(layout.lengths[:live], alone.lengths) + # The spare group holds no history, so nothing reads the pad rows' slot. + self.assertFalse(layout.history_valid[SWA_WINDOW:].any()) + + +class TestPrefillGraphReplay(CustomTestCase): + """--enable-encoder-swa-bounded-replay under the breakable prefill CUDA graph.""" + + def _layout(self, seed): + g = torch.Generator().manual_seed(seed) + req = torch.randint(1, 5, (8,), generator=g).sort().values + pos = torch.randint(0, 600, (8,), generator=g) + return window_layout(req, pos, capacity=256, num_groups=4) + + def _core(self, layout): + n = layout.req.shape[0] + metadata = DSV4AttnMetadata( + page_size=256, + page_table=torch.zeros((n, 2), dtype=torch.int32), + raw_out_loc=torch.zeros(n, dtype=torch.int32), + cuda_int32_kwargs={"dtype": torch.int32}, + seq_lens_casual=torch.zeros(n, dtype=torch.int32), + positions_casual=torch.zeros(n, dtype=torch.int32), + swa_page_indices=torch.zeros((n, 2), dtype=torch.int32), + swa_topk_lengths=torch.zeros(n, dtype=torch.int32), + index_topk=128, + present_ratios=(0,), + request_window_layout=layout, + ) + for name in ( + "c4_sparse_topk_lengths", + "c4_sparse_page_indices", + "c0_flashmla_metadata", + "c4_flashmla_metadata", + "c128_flashmla_metadata", + ): + setattr(metadata, name, None) + return metadata + + def test_replay_refreshes_the_captured_window_layout_in_place(self): + """The captured window copy/commit kernels read the layout tensors by + address, so a replay must rewrite the captured object's contents.""" + captured, live = self._layout(1), self._layout(2) + metadata = self._core(captured) + metadata.refresh_for_breakable_cuda_graph_replay_(self._core(live)) + self.assertIs(metadata.request_window_layout, captured) + for name in ("req", "pos", "indices", "lengths", "commit_mask", "history_pos"): + torch.testing.assert_close(getattr(captured, name), getattr(live, name)) + + def _backend(self, c1_out_loc, c2_out_loc, raw_out_loc): + backend = object.__new__(DeepseekV4AttnBackend) + backend.device = torch.device("cpu") + backend.forward_metadata = SimpleNamespace( + core_metadata=SimpleNamespace( + c1_out_loc=c1_out_loc, c2_out_loc=c2_out_loc, raw_out_loc=raw_out_loc + ) + ) + backend._low_ratio_write_group = MagicMock() + return backend + + def test_in_graph_compressor_skips_folded_replay_rows(self): + """Inside the graph the ratio-1/2 compressors see every bucket row, replay + rows included; those groups are cached already, so their writes must land + on the dummy slot and they must not feed the ring state.""" + fx = _Fixture() + folded = fx.folded + n, rows = 320, folded.num_rows # bucket rows, folded live rows + fb = SimpleNamespace(encoder_swa_compress_rows=folded.keep_rows) + backend = self._backend( + c1_out_loc=torch.arange(1, n + 1, dtype=torch.int32), + c2_out_loc=torch.arange(1, n + 1, dtype=torch.int32), + raw_out_loc=torch.cat([torch.ones(rows), torch.zeros(n - rows)]).int(), + ) + skip = backend._encoder_replay_rows(fb, num_tokens=rows, num_rows=n) + expected = torch.zeros(n, dtype=torch.bool) + expected[40:168] = True # slot 2 replays 128 rows before its 30 new ones + expected[198:262] = True # slot 3 replays 64 rows before its 20 new ones + torch.testing.assert_close(skip, expected) + self.assertEqual(int(skip.sum()), REPLAY_ROWS) + + kv, score, pos = torch.zeros(n, 4), torch.zeros(n, 4), torch.arange(n) + req = torch.zeros(n, dtype=torch.int64) + pad = (torch.arange(n) >= rows) | skip + layer = SimpleNamespace(compress_ratio=1, compressor=MagicMock(), layer_id=0) + backend._low_ratio_compress_torch( + layer, None, req, pos, projected=(kv, score), skip_rows=skip + ) + slots = backend._low_ratio_write_group.call_args.args[2] + torch.testing.assert_close( + slots[pad], torch.zeros(int(pad.sum()), dtype=torch.int32) + ) + torch.testing.assert_close( + slots[~pad], torch.arange(1, n + 1, dtype=torch.int32)[~pad] + ) + + layer.compress_ratio = 2 + layer.compressor.pool_pairs.return_value = kv + backend._low_ratio_pair_partners = MagicMock(return_value=(kv, score)) + backend._low_ratio_compress_torch( + layer, None, req, pos, projected=(kv, score), skip_rows=skip + ) + torch.testing.assert_close( + backend._low_ratio_pair_partners.call_args.kwargs["pad"], pad + ) + slots = backend._low_ratio_write_group.call_args.args[2] + self.assertFalse(slots[pad].any()) + self.assertTrue((slots[~pad] > 0).all()) + + class TestFoldMlpSync(CustomTestCase): """Pure TP that still syncs the MLP (e.g. TP4 with --moe-dense-tp-size 1)."""