From 52d048081f98714419cb014ff136fa7575539301 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Sun, 11 Oct 2026 17:33:24 -0700 Subject: [PATCH 1/3] [DeepSeek-V4.1] Run decoder SWA bounded replay inside the breakable prefill CUDA graph Bounded replay refused every prefill graph backend, so V4.1 prefill ran eager and host-launch-bound on GB300. Keep the late layers on their per-request tail inside the captured graph: a per-shape tail block (min(max_running_requests x 128, shape tokens)) with persistent block-sized extend/scatter indices and SWA write locations, one tail metadata per block size built the prefill-graph way and refreshed in place on replay, enter/exit eager breaks around the late layers, and the logits row indices carried through the runner's output trim for the DSpark injector. The decode/target-verify graph metadata init clears the tail metadata the prefill graph leaves behind, and the DSV4.1 feature check admits the breakable backend with decoder bounded replay. --- .../sglang/srt/arg_groups/deepseek_v4_hook.py | 5 +- .../layers/attention/deepseek_v4_backend.py | 246 ++++++++++++++++-- .../runner/prefill_cuda_graph_runner.py | 3 + python/sglang/srt/models/deepseek_v4.py | 63 ++++- 4 files changed, 295 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 14b9c29553a5..cc32a1d95696 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -320,10 +320,13 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None: if cfg.enable_decoder_swa_bounded_replay: # Late layers see a per-request tail slice, not the captured prefill shape. + # The breakable backend is admitted: the tail switch runs as eager breaks and + # the tail is a fixed-capacity row block inside the captured segments. incompatible = ( ( "the prefill CUDA graph", - cfg.cuda_graph_config.prefill.backend != Backend.DISABLED, + cfg.cuda_graph_config.prefill.backend + not in (Backend.DISABLED, Backend.BREAKABLE), ), # input_ids_global is a DP-wide gather, so the tail slice cannot apply. ("DP attention", attn_dp_enabled_of(cfg)), diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 4420fdecd383..c8027592aed9 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -915,6 +915,11 @@ class LateLayerTail(msgspec.Struct, frozen=True): local_lens_cpu: Optional[List[int]] = None req_global: Optional[torch.Tensor] = None pos_global: Optional[torch.Tensor] = None + # Prefill CUDA graph: token_indices/positions are persistent buffers of a fixed + # capacity (real rows first, padding after); scatter_indices maps every row back + # into the extend, padding rows to a dummy row past the end. + scatter_indices: Optional[torch.Tensor] = None + real_rows_count: Optional[int] = None def rows(self, t: torch.Tensor) -> torch.Tensor: rows = self.real_rows(t) @@ -1533,10 +1538,69 @@ def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool: seq_lens_cpu = forward_batch.seq_lens_cpu if max_seq_len is None or seq_lens_cpu is None or seq_lens_cpu.numel() == 0: return True - return int(seq_lens_cpu.max().item()) <= max_seq_len + if int(seq_lens_cpu.max().item()) > max_seq_len: + return False + if self.enable_decoder_swa_bounded_replay: + extend_lens_cpu = forward_batch.extend_seq_lens_cpu + if extend_lens_cpu is not None: + tail_rows = sum(min(SWA_WINDOW, int(n)) for n in extend_lens_cpu) + if tail_rows > self._prefill_graph_tail_capacity(): + return False + return True + + def _prefill_graph_tail_capacity(self) -> int: + """Rows the late layers run over inside the prefill CUDA graph: every + request's SWA_WINDOW tail at the request cap, so one capture fits any batch.""" + cap = getattr(self, "_bcg_tail_capacity", None) + if cap is None: + max_reqs = getattr( + getattr(self.model_runner, "server_args", None), + "max_running_requests", + None, + ) + if not max_reqs: + max_reqs = int(self.model_runner.req_to_token_pool.size) + cap = int(max_reqs) * SWA_WINDOW + self._bcg_tail_capacity = cap + return cap + + def _prefill_graph_tail_block_rows(self, forward_batch: ForwardBatch) -> int: + """Tail block rows for the captured shape this batch is padded to: the + shape's token count, capped at the request-cap capacity.""" + return min( + self._prefill_graph_tail_capacity(), + int(forward_batch.out_cache_loc.shape[0]), + ) + + def _prefill_graph_tail_buffers( + self, device: torch.device, pos_dtype: torch.dtype, loc_dtype: torch.dtype + ) -> tuple: + # (gather indices, positions, scatter indices, SWA write locations, full + # write locations); the positions buffer keeps the batch's dtype so the + # captured graph reads it directly, no copies. The late layers' KV stores + # and low-ratio compressors run inside the captured segments over every + # row of the block, so the write locations must be persistent buffers of + # the same capacity: padding rows hold slot 0, the dummy slot that also + # marks them as padding for the compressors (raw_out_loc == 0). + bufs = getattr(self, "_bcg_tail_buffers", None) + if bufs is None: + cap = self._prefill_graph_tail_capacity() + bufs = ( + torch.zeros((cap,), dtype=torch.int64, device=device), + torch.zeros((cap,), dtype=pos_dtype, device=device), + torch.zeros((cap,), dtype=torch.int64, device=device), + torch.zeros((cap,), dtype=torch.int32, device=device), + torch.zeros((cap,), dtype=loc_dtype, device=device), + ) + self._bcg_tail_buffers = bufs + return bufs def _build_late_layer_tail_metadata( - self, forward_batch: ForwardBatch + self, + forward_batch: ForwardBatch, + *, + for_prefill_graph: bool = False, + max_seq_len_override: Optional[int] = None, ) -> DSV4Metadata: # Each request contributes only its last SWA_WINDOW extend tokens, with the # window floored at the tail start: window KV before it is never written here. @@ -1553,11 +1617,60 @@ def _build_late_layer_tail_metadata( contiguous_start = ( extend_lens_cpu[0] - tail_lens_cpu[0] if len(extend_lens_cpu) == 1 else None ) - out_cache_loc = _tail_rows( - forward_batch.out_cache_loc, - token_indices=token_indices, - contiguous_start=contiguous_start, - ) + graph_tail = None + if for_prefill_graph: + # Captured segments gather/scatter through persistent buffers of a + # fixed capacity: real rows first, padding rows gather row 0 and + # scatter to a dummy row past the extend. + assert not is_cp_active(forward_batch), ( + "decoder SWA bounded replay inside the prefill CUDA graph does not " + "support prefill CP" + ) + real_rows = int(sum(tail_lens_cpu)) + # The block is sized per captured shape: a shape of N tokens can + # never hold more than N tail rows, so small shapes get small + # blocks and only the shapes past the request cap use the full one. + cap = self._prefill_graph_tail_block_rows(forward_batch) + assert real_rows <= cap, (real_rows, cap) + idx_buf, pos_buf, scatter_buf, swa_buf, loc_buf = ( + buf[:cap] + for buf in self._prefill_graph_tail_buffers( + device, + forward_batch.positions.dtype, + forward_batch.out_cache_loc.dtype, + ) + ) + real_idx = ( + torch.arange(contiguous_start, extend_lens_cpu[0], device=device) + if contiguous_start is not None + else token_indices + ) + idx_buf[:real_rows].copy_(real_idx) + idx_buf[real_rows:].zero_() + scatter_buf[:real_rows].copy_(real_idx) + scatter_buf[real_rows:].fill_(int(forward_batch.out_cache_loc.shape[0])) + pos_buf[:real_rows].copy_(forward_batch.positions[real_idx]) + pos_buf[real_rows:].zero_() + loc_buf[:real_rows].copy_(forward_batch.out_cache_loc[real_idx]) + loc_buf[real_rows:].zero_() + # The window floor is combined with the block-sized causal positions. + swa_replay_start = torch.nn.functional.pad( + swa_replay_start, (0, cap - real_rows) + ) + graph_tail = (real_rows, idx_buf, pos_buf, scatter_buf, swa_buf) + contiguous_start = None + if graph_tail is not None: + # The block's metadata is built the way the full prefill graph builds + # its own: live rows against a padded write-location buffer, so the + # causal expansion, compressor plans and low-ratio indexer metadata + # are bucket-sized and graph-safe. + out_cache_loc = loc_buf + else: + out_cache_loc = _tail_rows( + forward_batch.out_cache_loc, + token_indices=token_indices, + contiguous_start=contiguous_start, + ) tail_lens = torch.tensor(tail_lens_cpu, dtype=torch.int32, device=device) cp_tail = ( self._late_layer_tail_cp_layout(forward_batch, token_indices, tail_lens) @@ -1566,7 +1679,11 @@ def _build_late_layer_tail_metadata( ) metadata = self.init_forward_metadata_prefill( - max_seq_len=int(seq_lens_cpu.max().item()), + max_seq_len=( + max_seq_len_override + if max_seq_len_override is not None + else int(seq_lens_cpu.max().item()) + ), req_pool_indices=forward_batch.req_pool_indices, seq_lens=forward_batch.seq_lens.to(torch.int32), seq_lens_cpu=seq_lens_cpu.tolist(), @@ -1576,6 +1693,7 @@ def _build_late_layer_tail_metadata( extend_seq_lens_cpu=tail_lens_cpu, extend_start_loc=torch.cumsum(tail_lens, dim=0) - tail_lens, swa_replay_start=swa_replay_start, + use_prefill_cuda_graph=graph_tail is not None, forward_batch=forward_batch if cp_tail is not None else None, cp_metadata=cp_tail["cp_metadata"] if cp_tail is not None else None, ) @@ -1587,16 +1705,46 @@ def _build_late_layer_tail_metadata( ) ) metadata.core_attn_metadata.swa_out_cache_loc = swa_out_cache_loc - metadata.low_ratio_req_indices = torch.repeat_interleave( - forward_batch.req_pool_indices.to(torch.int64), tail_lens.to(torch.int64) - ) - positions = _tail_rows( - forward_batch.positions, - token_indices=token_indices, - contiguous_start=contiguous_start, - ) - metadata.low_ratio_pos_i64 = positions.to(torch.int64) - if cp_tail is None: + if graph_tail is None: + # The graph build keeps the bucket-sized request/position rows the + # prefill-graph metadata init published; eager uses the exact rows. + metadata.low_ratio_req_indices = torch.repeat_interleave( + forward_batch.req_pool_indices.to(torch.int64), + tail_lens.to(torch.int64), + ) + positions = _tail_rows( + forward_batch.positions, + token_indices=token_indices, + contiguous_start=contiguous_start, + ) + metadata.low_ratio_pos_i64 = positions.to(torch.int64) + if graph_tail is not None: + real_rows, idx_buf, pos_buf, scatter_buf, swa_buf = graph_tail + # The captured KV stores take the whole block: real rows write their + # SWA slots, padding rows the dummy slot (as the full graph's padding does). + swa_buf.copy_(swa_out_cache_loc.to(torch.int32)) + if metadata.low_ratio_req_indices is None: + # Low-ratio sources outside the graph path still take block-sized + # rows here; padding rows repeat the last request at position 0. + cap = idx_buf.shape[0] + req = torch.repeat_interleave( + forward_batch.req_pool_indices.to(torch.int64), + tail_lens.to(torch.int64), + ) + metadata.low_ratio_req_indices = torch.cat( + [req, req[-1:].expand(cap - real_rows)] + ) + metadata.low_ratio_pos_i64 = pos_buf.to(torch.int64) + metadata.late_layer_tail = LateLayerTail( + token_indices=idx_buf, + positions=pos_buf, + extend_seq_lens=tail_lens, + extend_seq_lens_cpu=tail_lens_cpu, + swa_out_cache_loc=swa_buf, + scatter_indices=scatter_buf, + real_rows_count=real_rows, + ) + elif cp_tail is None: # Without CP, tail rows index the full extend on this rank. metadata.late_layer_tail = LateLayerTail( token_indices=token_indices, @@ -1698,6 +1846,13 @@ def enter_late_layer_tail(self, forward_batch: ForwardBatch) -> tuple: get_local_dp_buffer_len(), ) tail = tail_metadata.late_layer_tail + # Inside the prefill CUDA graph the tail block is padded to a fixed capacity; + # the eager breaks (attention, low-ratio sources) clip their inputs to the + # batch's non-padded count, so point it at the real tail rows while the + # late layers run. + saved = saved + (forward_batch.global_num_token_non_padded_cpu,) + if tail.real_rows_count is not None: + forward_batch.global_num_token_non_padded_cpu = tail.real_rows_count # The layers before the switch published top-k into the full metadata's # buffers; carry the tail rows into the tail metadata's (padding stays -1). full_core = saved[0].core_attn_metadata @@ -1720,6 +1875,8 @@ def enter_late_layer_tail(self, forward_batch: ForwardBatch) -> tuple: if full_buf is None or tail_buf is None: continue rows = tail.real_rows(full_buf) + if tail.real_rows_count is not None: + rows = rows[: tail.real_rows_count] tail_buf[: rows.shape[0]].copy_(rows) self.forward_metadata = tail_metadata if self.token_to_kv_pool.request_window is not None: @@ -1736,8 +1893,20 @@ def exit_late_layer_tail(self, saved: tuple, forward_batch: ForwardBatch) -> Non self.forward_metadata, forward_batch.attn_cp_metadata, local_dp_buffer_len, + non_padded, ) = saved set_local_dp_buffer_len(local_dp_buffer_len) + forward_batch.global_num_token_non_padded_cpu = non_padded + + # The prefill CUDA graph runs the tail switch as eager breaks: these two run + # at capture and again at every replay (see deepseek_v4.bcg_enter/exit). + def prefill_graph_enter_late_layer_tail(self, forward_batch: ForwardBatch) -> None: + self.prefill_graph_late_layer_saved = self.enter_late_layer_tail(forward_batch) + + def prefill_graph_exit_late_layer_tail(self, forward_batch: ForwardBatch) -> None: + saved = self.prefill_graph_late_layer_saved + self.prefill_graph_late_layer_saved = None + self.exit_late_layer_tail(saved, forward_batch) def init_forward_metadata_target_verify( self, @@ -2058,10 +2227,14 @@ def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: raw_metadata=self.forward_metadata, online_c128_state_slot_offset=self.online_c128_mtp.state_slot_offset(), ) + # Decode/verify graphs have no late-layer tail; drop the one the + # prefill graph capture or replay before this batch left behind. + self.tail_forward_metadata = None elif isinstance(self.forward_metadata, DSV4RawDecodeMetadata): self.forward_metadata = self.make_forward_metadata_from_raw_decode( raw_metadata=self.forward_metadata, ) + self.tail_forward_metadata = None metadata = self.forward_metadata if isinstance(metadata, DSV4Metadata): @@ -2586,6 +2759,12 @@ def init_forward_metadata_for_breakable_cuda_graph_capture( max_seq_len_override=max_seq_len, use_prefill_cuda_graph=True, ) + self.tail_forward_metadata = ( + self._prefill_graph_tail_metadata(forward_batch, max_seq_len) + if self.enable_decoder_swa_bounded_replay + and forward_batch.forward_mode.is_extend_without_speculative() + else None + ) if self.low_ratio_prefill_graph and forward_batch.forward_mode.is_extend(): for ratio in self.low_ratios: self._source_projection_buffers( @@ -2643,6 +2822,37 @@ 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.tail_forward_metadata = ( + self._prefill_graph_tail_metadata(metadata_batch, max_seq_len) + if self.enable_decoder_swa_bounded_replay + and metadata_batch.forward_mode.is_extend_without_speculative() + else None + ) + + def _prefill_graph_tail_metadata( + self, forward_batch: ForwardBatch, max_seq_len: int + ) -> DSV4Metadata: + """The late layers' metadata for the prefill CUDA graph: one object per + tail block size (shared by every captured shape of that size), refreshed + in place like the full metadata, since the captured low-ratio + compressors and indexers read its tensors by address.""" + fresh = self._build_late_layer_tail_metadata( + forward_batch, for_prefill_graph=True, max_seq_len_override=max_seq_len + ) + held_by_rows = getattr(self, "_bcg_tail_metadata", None) + if held_by_rows is None: + held_by_rows = self._bcg_tail_metadata = {} + rows = self._prefill_graph_tail_block_rows(forward_batch) + held = held_by_rows.get(rows) + if held is None: + held_by_rows[rows] = fresh + return fresh + held.refresh_for_breakable_cuda_graph_replay_(fresh) + held.core_attn_metadata.swa_out_cache_loc = ( + fresh.core_attn_metadata.swa_out_cache_loc + ) + held.late_layer_tail = fresh.late_layer_tail + return held def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: self.cuda_graph_metadata_of_bucket_and_bs: Dict[ 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 7a11965ca46d..b5b91a05f039 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 @@ -2203,6 +2203,9 @@ def _trim_logits_output( input_logprobs_copy_done=output.input_logprobs_copy_done, customized_info=output.customized_info, mm_input_embeds=mm_input_embeds, + # Bounded SWA replay returns the late layers' tail rows plus the + # extend indices they came from; the DSpark injector needs both. + hidden_states_token_indices=output.hidden_states_token_indices, ) def _finalize_execute_output( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index cd8db6534d2f..5abd7c3e2e2e 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -793,6 +793,29 @@ def deepseek_v4_engram_hash_ids(hasher, input_ids: torch.Tensor) -> torch.Tensor bcg_deepseek_v4_engram_hash_ids = eager_on_graph(True)(deepseek_v4_engram_hash_ids) +def deepseek_v4_enter_late_layer_tail(sentinel: torch.Tensor) -> torch.Tensor: + # Decoder SWA bounded replay under the prefill CUDA graph: swap the attention + # backend onto the tail metadata of the *live* batch (the captured args hold + # the capture-time batch). The sentinel keeps the break in the replay sequence. + forward_batch = get_tc_piecewise_forward_context().forward_batch + get_attn_backend().prefill_graph_enter_late_layer_tail(forward_batch) + return sentinel + + +def deepseek_v4_exit_late_layer_tail(sentinel: torch.Tensor) -> torch.Tensor: + forward_batch = get_tc_piecewise_forward_context().forward_batch + get_attn_backend().prefill_graph_exit_late_layer_tail(forward_batch) + return sentinel + + +bcg_deepseek_v4_enter_late_layer_tail = eager_on_graph(True)( + deepseek_v4_enter_late_layer_tail +) +bcg_deepseek_v4_exit_late_layer_tail = eager_on_graph(True)( + deepseek_v4_exit_late_layer_tail +) + + class MqaAttentionBase(nn.Module): # Class-level default for subclasses that read it without running __init__. wo_a_fp8: bool = False @@ -3813,6 +3836,12 @@ def _scatter_tail_rows( tail: LateLayerTail, rows: torch.Tensor, num_tokens: int ) -> torch.Tensor: # Rows outside the tail are never read (see _check_late_layer_tail_readers). + if tail.scatter_indices is not None: + # Prefill CUDA graph: fixed-capacity tail block; padding rows land on the + # dummy row past the extend and are dropped by the slice. + full = rows.new_zeros((num_tokens + 1, rows.shape[1])) + full.index_copy_(0, tail.scatter_indices, rows) + return full[:num_tokens] full = rows.new_empty((num_tokens, rows.shape[1])) if tail.contiguous_start is not None: full[tail.contiguous_start :].copy_(rows) @@ -4059,6 +4088,12 @@ def _forward_layers_hc_pre_from_prev( self._check_late_layer_tail_readers(forward_batch) attn_backend = get_attn_backend() tail = attn_backend.tail_forward_metadata.late_layer_tail + tail_in_graph = ( + tail is not None + and tail.scatter_indices is not None + and is_in_breakable_cuda_graph() + ) + tail_sentinel = positions.new_zeros((1,)) if tail_in_graph else None saved_full = None # A pending post never meets a residual reader: the HIP boundary's defer_post # gate excludes Engram/DSpark-capture layers and the model end. @@ -4068,7 +4103,14 @@ def _forward_layers_hc_pre_from_prev( for i in range(self.start_layer, self.end_layer): if tail is not None and i == self.late_layer_start: # Decode reaches back at most SWA_WINDOW positions. - saved_full = attn_backend.enter_late_layer_tail(forward_batch) + if tail_in_graph: + # Eager break: re-run at replay so the backend swaps onto the + # live batch's tail metadata; the gathers below are captured + # against the persistent index buffers. + tail_sentinel = bcg_deepseek_v4_enter_late_layer_tail(tail_sentinel) + saved_full = attn_backend.prefill_graph_late_layer_saved + else: + saved_full = attn_backend.enter_late_layer_tail(forward_batch) state = state.take_rows(tail.rows) input_ids, input_ids_global = ( tail.rows(input_ids), @@ -4134,7 +4176,10 @@ def _forward_layers_hc_pre_from_prev( ) state = state.materialized(self.layers[self.end_layer - 1].hc_cfg) if saved_full is not None: - attn_backend.exit_late_layer_tail(saved_full, forward_batch) + if tail_in_graph: + tail_sentinel = bcg_deepseek_v4_exit_late_layer_tail(tail_sentinel) + else: + attn_backend.exit_late_layer_tail(saved_full, forward_batch) return state.residual, state.pre, tail return state.residual, state.pre, None @@ -4732,6 +4777,14 @@ def forward( ): tail = get_attn_backend().tail_forward_metadata.late_layer_tail input_ids = tail.rows(input_ids) + if tail.real_rows_count is not None: + # Fixed-capacity tail block from the prefill CUDA graph: drop padding. + n = tail.real_rows_count + input_ids = input_ids[:n] + hidden_states = hidden_states[:n] + pre_hc_head = pre_hc_head[:n] + if aux_hidden_states is not None: + aux_hidden_states = [a[:n] for a in aux_hidden_states] logits_metadata = LogitsMetadata.from_forward_batch(forward_batch) logits_metadata.extend_seq_lens = tail.extend_seq_lens logits_metadata.extend_seq_lens_cpu = tail.extend_seq_lens_cpu @@ -4748,7 +4801,11 @@ def forward( ), ) if tail is not None: - output.hidden_states_token_indices = tail.token_indices + output.hidden_states_token_indices = ( + tail.token_indices[: tail.real_rows_count] + if tail.real_rows_count is not None + else tail.token_indices + ) return output def _setup_fp8_wo_a_scales(self, is_nextn: bool) -> None: From 6dfa52314671072afc23d15d02a90e23d3ae83a2 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Mon, 12 Oct 2026 20:28:20 -0700 Subject: [PATCH 2/3] [DeepSeek-V4.1] bounded replay in prefill graph: tolerate backends built without the feature flag test_backend_replay_keeps_captured_metadata_active builds the backend with object.__new__, so the graph capture/replay paths must not assume enable_decoder_swa_bounded_replay exists. --- python/sglang/srt/layers/attention/deepseek_v4_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 0e59f7c9ce62..003693ed0a9b 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -2785,7 +2785,7 @@ def init_forward_metadata_for_breakable_cuda_graph_capture( ) self.tail_forward_metadata = ( self._prefill_graph_tail_metadata(forward_batch, max_seq_len) - if self.enable_decoder_swa_bounded_replay + if getattr(self, "enable_decoder_swa_bounded_replay", False) and forward_batch.forward_mode.is_extend_without_speculative() else None ) @@ -2848,7 +2848,7 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay( self.forward_metadata = capture_metadata self.tail_forward_metadata = ( self._prefill_graph_tail_metadata(metadata_batch, max_seq_len) - if self.enable_decoder_swa_bounded_replay + if getattr(self, "enable_decoder_swa_bounded_replay", False) and metadata_batch.forward_mode.is_extend_without_speculative() else None ) From a4310242ac9c1b56e58ffccdeb1c78b9ca58dc39 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Mon, 12 Oct 2026 21:23:35 -0700 Subject: [PATCH 3/3] [DeepSeek-V4.1] bounded replay in prefill graph: no lazy attributes Initialise the tail block state (capacity, persistent buffers, held metadata per block size) in __init__, and record on the captured DSV4Metadata whether the late layers were captured on a tail block (prefill_graph_tail); replay rebuilds the tail metadata from that flag instead of consulting the backend's feature flag, so a backend built without __init__ (test_backend_replay_keeps_captured_metadata_active) replays without a tail. --- .../layers/attention/deepseek_v4_backend.py | 56 +++++++++++-------- 1 file changed, 32 insertions(+), 24 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 003693ed0a9b..59f7df8207fa 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -983,6 +983,12 @@ class DSV4Metadata: # Set only on the metadata built for the late layers under bounded SWA replay. late_layer_tail: Optional[LateLayerTail] = None + # Captured with the late layers on a tail block (decoder SWA bounded replay + # inside the prefill CUDA graph); replay rebuilds the tail metadata iff set. + prefill_graph_tail: bool = False + # Captured with the late layers on a tail block (decoder SWA bounded replay + # inside the prefill CUDA graph); replay rebuilds the tail metadata iff set. + prefill_graph_tail: bool = False @property def core_metadata(self) -> DSV4AttnMetadata: @@ -1181,6 +1187,16 @@ def __init__( ) # The model switches onto this metadata in enter_late_layer_tail. self.tail_forward_metadata: Optional[DSV4Metadata] = None + # Prefill CUDA graph tail block (bounded replay inside the graph): row + # capacity, persistent index/location buffers, held metadata per block size. + self._bcg_tail_capacity: Optional[int] = None + self._bcg_tail_buffers: Optional[tuple] = None + self._bcg_tail_metadata: Dict[int, DSV4Metadata] = {} + # Prefill CUDA graph tail block (bounded replay inside the graph): row + # capacity, persistent index/location buffers, held metadata per block size. + self._bcg_tail_capacity: Optional[int] = None + self._bcg_tail_buffers: Optional[tuple] = None + self._bcg_tail_metadata: Dict[int, DSV4Metadata] = {} self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner) self.dsv4_prefill_backend = getattr(kernel, "dsv4_prefill_backend", "auto") if use_dsv4_q8kv8_sparse_prefill(self.dsv4_prefill_backend): @@ -1561,18 +1577,12 @@ def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool: def _prefill_graph_tail_capacity(self) -> int: """Rows the late layers run over inside the prefill CUDA graph: every request's SWA_WINDOW tail at the request cap, so one capture fits any batch.""" - cap = getattr(self, "_bcg_tail_capacity", None) - if cap is None: - max_reqs = getattr( - getattr(self.model_runner, "server_args", None), - "max_running_requests", - None, - ) + if self._bcg_tail_capacity is None: + max_reqs = self.model_runner.server_args.max_running_requests if not max_reqs: max_reqs = int(self.model_runner.req_to_token_pool.size) - cap = int(max_reqs) * SWA_WINDOW - self._bcg_tail_capacity = cap - return cap + self._bcg_tail_capacity = int(max_reqs) * SWA_WINDOW + return self._bcg_tail_capacity def _prefill_graph_tail_block_rows(self, forward_batch: ForwardBatch) -> int: """Tail block rows for the captured shape this batch is padded to: the @@ -1592,18 +1602,16 @@ def _prefill_graph_tail_buffers( # row of the block, so the write locations must be persistent buffers of # the same capacity: padding rows hold slot 0, the dummy slot that also # marks them as padding for the compressors (raw_out_loc == 0). - bufs = getattr(self, "_bcg_tail_buffers", None) - if bufs is None: + if self._bcg_tail_buffers is None: cap = self._prefill_graph_tail_capacity() - bufs = ( + self._bcg_tail_buffers = ( torch.zeros((cap,), dtype=torch.int64, device=device), torch.zeros((cap,), dtype=pos_dtype, device=device), torch.zeros((cap,), dtype=torch.int64, device=device), torch.zeros((cap,), dtype=torch.int32, device=device), torch.zeros((cap,), dtype=loc_dtype, device=device), ) - self._bcg_tail_buffers = bufs - return bufs + return self._bcg_tail_buffers def _build_late_layer_tail_metadata( self, @@ -2783,10 +2791,14 @@ def init_forward_metadata_for_breakable_cuda_graph_capture( max_seq_len_override=max_seq_len, use_prefill_cuda_graph=True, ) + with_tail = ( + self.enable_decoder_swa_bounded_replay + and forward_batch.forward_mode.is_extend_without_speculative() + ) + self.forward_metadata.prefill_graph_tail = with_tail self.tail_forward_metadata = ( self._prefill_graph_tail_metadata(forward_batch, max_seq_len) - if getattr(self, "enable_decoder_swa_bounded_replay", False) - and forward_batch.forward_mode.is_extend_without_speculative() + if with_tail else None ) if self.low_ratio_prefill_graph and forward_batch.forward_mode.is_extend(): @@ -2848,8 +2860,7 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay( self.forward_metadata = capture_metadata self.tail_forward_metadata = ( self._prefill_graph_tail_metadata(metadata_batch, max_seq_len) - if getattr(self, "enable_decoder_swa_bounded_replay", False) - and metadata_batch.forward_mode.is_extend_without_speculative() + if capture_metadata.prefill_graph_tail else None ) @@ -2863,13 +2874,10 @@ def _prefill_graph_tail_metadata( fresh = self._build_late_layer_tail_metadata( forward_batch, for_prefill_graph=True, max_seq_len_override=max_seq_len ) - held_by_rows = getattr(self, "_bcg_tail_metadata", None) - if held_by_rows is None: - held_by_rows = self._bcg_tail_metadata = {} rows = self._prefill_graph_tail_block_rows(forward_batch) - held = held_by_rows.get(rows) + held = self._bcg_tail_metadata.get(rows) if held is None: - held_by_rows[rows] = fresh + self._bcg_tail_metadata[rows] = fresh return fresh held.refresh_for_breakable_cuda_graph_replay_(fresh) held.core_attn_metadata.swa_out_cache_loc = (