From 2e4a3d8814e1e580908c3cf251856191d7482284 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Wed, 21 Oct 2026 22:01:54 -0700 Subject: [PATCH 1/4] snapshot v4.1 index-k slots in prefill graph replay; stub it in the replay-context test --- python/sglang/srt/layers/attention/deepseek_v4_backend.py | 7 ++++++- .../sglang/srt/layers/attention/dsv4/v41_indexer/types.py | 7 ++++--- .../unit/layers/test_dsv41_late_layer_tail_readers.py | 5 ++++- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 9c71a22e90b6..7b2328bc5c4a 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -3132,6 +3132,11 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay( if capture_metadata.prefill_graph_tail else None ) + if ( + self.low_ratios + and metadata_batch.forward_mode.is_extend_without_speculative() + ): + self._snapshot_low_ratio_k_slots(metadata_batch) def _prefill_graph_tail_metadata( self, forward_batch: ForwardBatch, max_seq_len: int @@ -3834,7 +3839,7 @@ def _make_low_ratio_prefill_indexer_inputs( if k_slots is not None: k_slots = k_slots[layer.compress_ratio] else: - # Graph-captured metadata carries no snapshot. + # No snapshot, e.g. the prefill graph's capture pass. k_slots = prefill_k_slots( req_to_token=self.req_to_token, req_pool_indices=forward_batch.req_pool_indices, diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py index 36487d083bf8..848b66a58df1 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py @@ -48,11 +48,12 @@ class PrefillInputs(msgspec.Struct, frozen=True, kw_only=True): req_rows: Optional[torch.Tensor] req_pool_indices: torch.Tensor # [requests] req_to_token row of each request # [columns] int64 index-K pool slot of each flattened-K column, from - # scoring.prefill_k_slots; gathered with the metadata, never per layer. + # scoring.prefill_k_slots; gathered with the metadata, per layer only when the + # metadata has no snapshot (e.g. the prefill graph's capture pass). k_slots: torch.Tensor # ratio -> scoring.PrefillIndexPlan, filled by the first layer of each ratio and - # reused by the rest of the forward; None when the metadata outlives a forward - # (captured graphs), so every layer builds its own. + # reused by the rest of the forward; None when the metadata has no snapshot, so + # every layer builds its own. index_plans: Optional[Dict[int, object]] # [rows, pages] int32 at the KV page size; the index-K pool pages it with # expand_index_page_table. diff --git a/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py b/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py index b227a612f784..12c608ea8458 100644 --- a/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py +++ b/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py @@ -291,6 +291,7 @@ def build_tail(forward_batch, max_seq_len): backend._build_forward_metadata = build_full backend._prefill_graph_tail_metadata = build_tail backend._activate_request_window = lambda: None + backend._snapshot_low_ratio_k_slots = lambda forward_batch: None backend._prefill_graph_replays_any_context = lambda: any_context capture = object.__new__(DSV4Metadata) capture.prefill_graph_tail = True @@ -310,7 +311,9 @@ def build_tail(forward_batch, max_seq_len): backend, capture, SimpleNamespace( - seq_lens_cpu=torch.tensor([seq_len]), max_seq_len_override=1 << 20 + seq_lens_cpu=torch.tensor([seq_len]), + max_seq_len_override=1 << 20, + forward_mode=ForwardMode.EXTEND, ), ) return seen From bad5bc6f75a4095ac2a311acee80300c755f4979 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Wed, 21 Oct 2026 22:06:00 -0700 Subject: [PATCH 2/4] move the snapshot condition into the snapshot helper; drop the test batch's forward mode --- .../srt/layers/attention/deepseek_v4_backend.py | 17 +++++++---------- .../test_dsv41_late_layer_tail_readers.py | 4 +--- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 7b2328bc5c4a..d7a5079539e7 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -2724,17 +2724,18 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: and forward_batch.forward_mode.is_extend_without_speculative() else None ) - if ( - self.low_ratios - and forward_batch.forward_mode.is_extend_without_speculative() - ): - self._snapshot_low_ratio_k_slots(forward_batch) + self._snapshot_low_ratio_k_slots(forward_batch) self._activate_request_window() def _snapshot_low_ratio_k_slots(self, forward_batch: ForwardBatch) -> None: # Gathered once with the metadata instead of in every indexer layer, so no # indexer layer reads req_to_token. + if not ( + self.low_ratios + and forward_batch.forward_mode.is_extend_without_speculative() + ): + return seq_lens_cpu = _as_int_list(forward_batch.seq_lens_cpu) k_slots = { ratio: prefill_k_slots( @@ -3132,11 +3133,7 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay( if capture_metadata.prefill_graph_tail else None ) - if ( - self.low_ratios - and metadata_batch.forward_mode.is_extend_without_speculative() - ): - self._snapshot_low_ratio_k_slots(metadata_batch) + self._snapshot_low_ratio_k_slots(metadata_batch) def _prefill_graph_tail_metadata( self, forward_batch: ForwardBatch, max_seq_len: int diff --git a/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py b/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py index 12c608ea8458..da1f68b3357c 100644 --- a/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py +++ b/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py @@ -311,9 +311,7 @@ def build_tail(forward_batch, max_seq_len): backend, capture, SimpleNamespace( - seq_lens_cpu=torch.tensor([seq_len]), - max_seq_len_override=1 << 20, - forward_mode=ForwardMode.EXTEND, + seq_lens_cpu=torch.tensor([seq_len]), max_seq_len_override=1 << 20 ), ) return seen From d3f2acd6f0e3a08a6c056cb054824216767ec34b Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Wed, 21 Oct 2026 23:14:39 -0700 Subject: [PATCH 3/4] share the index-k snapshot with the tail that leaves the prefill graph --- .../srt/layers/attention/deepseek_v4_backend.py | 11 +++++++++-- python/sglang/srt/models/deepseek_v4_replay_tail.py | 1 + 2 files changed, 10 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 d7a5079539e7..30da455c2b39 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -2760,8 +2760,15 @@ def _snapshot_low_ratio_k_slots(self, forward_batch: ForwardBatch) -> None: torch.int64 ) if self.tail_forward_metadata is not None: - self.tail_forward_metadata.low_ratio_k_slots = k_slots - self.tail_forward_metadata.low_ratio_index_plans = {} + self._share_low_ratio_k_slots(self.tail_forward_metadata) + + def _share_low_ratio_k_slots(self, metadata: DSV4Metadata) -> None: + # The late-layer tail's metadata reads this forward's snapshot; its rows + # differ from the extend's, so it caches its own index plans. + metadata.low_ratio_k_slots = self.forward_metadata.low_ratio_k_slots + metadata.low_ratio_index_plans = ( + None if metadata.low_ratio_k_slots is None else {} + ) def _activate_request_window(self) -> None: if self.token_to_kv_pool.request_window is not None: diff --git a/python/sglang/srt/models/deepseek_v4_replay_tail.py b/python/sglang/srt/models/deepseek_v4_replay_tail.py index b5e552929818..c83e1bf17a40 100644 --- a/python/sglang/srt/models/deepseek_v4_replay_tail.py +++ b/python/sglang/srt/models/deepseek_v4_replay_tail.py @@ -630,6 +630,7 @@ def _enter_eager_tail_in_prefill_graph( backend.tail_forward_metadata = backend._build_late_layer_tail_metadata( forward_batch ) + backend._share_low_ratio_k_slots(backend.tail_forward_metadata) try: saved = backend.enter_late_layer_tail(forward_batch) finally: From 4c26ac45d65a2d7bc9d89762c9fe306ada7d5b71 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Wed, 21 Oct 2026 23:28:57 -0700 Subject: [PATCH 4/4] take the index-k snapshot in the metadata builders; the tail builder inherits it and replay refresh carries it; graph metadata publishes no source rows --- .../layers/attention/deepseek_v4_backend.py | 39 +++++++++---------- .../attention/dsv4/v41_indexer/types.py | 4 +- .../srt/models/deepseek_v4_replay_tail.py | 1 - .../test_dsv41_late_layer_tail_readers.py | 1 - 4 files changed, 21 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 30da455c2b39..ab7b80facd7f 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -1043,11 +1043,11 @@ 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 - # Eager extend: each low ratio's index-K slots (scoring.prefill_k_slots), - # shared with the late-layer tail metadata. + # Each low ratio's index-K slots (scoring.prefill_k_slots), gathered once by the + # extend's metadata builder; the late-layer tail's metadata reads the same dict. low_ratio_k_slots: Optional[Dict[int, torch.Tensor]] = None - # Eager extend: the indexer's batch-only plan per low ratio, filled by its first - # layer (scoring.prefill_index_plan); one per metadata, as the rows differ. + # The indexer's batch-only plan per low ratio, filled by its first layer + # (scoring.prefill_index_plan); one per metadata, as the rows differ. low_ratio_index_plans: Optional[Dict[int, object]] = None # Per-step scratch for TP-padded query heads, zeroed by the first user. @@ -1115,6 +1115,8 @@ def refresh_for_breakable_cuda_graph_replay_( maybe_copy_inplace( self.low_ratio_pos_i64, src=static_metadata.low_ratio_pos_i64 ) + self.low_ratio_k_slots = static_metadata.low_ratio_k_slots + self.low_ratio_index_plans = static_metadata.low_ratio_index_plans maybe_copy_inplace( self.c4_compress_metadata, src=static_metadata.c4_compress_metadata ) @@ -1935,6 +1937,11 @@ def _build_late_layer_tail_metadata( req_global=metadata.low_ratio_req_indices, pos_global=metadata.low_ratio_pos_i64, ) + # Same requests as the extend in use, so the same index-K slots; other rows, + # so its own plans. + if self.low_ratios and self.forward_metadata.low_ratio_k_slots is not None: + metadata.low_ratio_k_slots = self.forward_metadata.low_ratio_k_slots + metadata.low_ratio_index_plans = {} return metadata def _late_layer_tail_cp_layout( @@ -2724,11 +2731,12 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: and forward_batch.forward_mode.is_extend_without_speculative() else None ) - self._snapshot_low_ratio_k_slots(forward_batch) self._activate_request_window() - def _snapshot_low_ratio_k_slots(self, forward_batch: ForwardBatch) -> None: + def _snapshot_low_ratio_k_slots( + self, forward_batch: ForwardBatch, metadata: DSV4Metadata, *, publish_rows: bool + ) -> None: # Gathered once with the metadata instead of in every indexer layer, so no # indexer layer reads req_to_token. if not ( @@ -2746,10 +2754,9 @@ def _snapshot_low_ratio_k_slots(self, forward_batch: ForwardBatch) -> None: ) for ratio in self.low_ratios } - metadata = self.forward_metadata metadata.low_ratio_k_slots = k_slots metadata.low_ratio_index_plans = {} - if metadata.low_ratio_req_indices is None: + if publish_rows: # Every source layer would rebuild these rows; the tail and the # prefill graph publish theirs at metadata init the same way. num_tokens = sum(forward_batch.extend_seq_lens_cpu) @@ -2759,16 +2766,6 @@ def _snapshot_low_ratio_k_slots(self, forward_batch: ForwardBatch) -> None: metadata.low_ratio_pos_i64 = forward_batch.positions[:num_tokens].to( torch.int64 ) - if self.tail_forward_metadata is not None: - self._share_low_ratio_k_slots(self.tail_forward_metadata) - - def _share_low_ratio_k_slots(self, metadata: DSV4Metadata) -> None: - # The late-layer tail's metadata reads this forward's snapshot; its rows - # differ from the extend's, so it caches its own index plans. - metadata.low_ratio_k_slots = self.forward_metadata.low_ratio_k_slots - metadata.low_ratio_index_plans = ( - None if metadata.low_ratio_k_slots is None else {} - ) def _activate_request_window(self) -> None: if self.token_to_kv_pool.request_window is not None: @@ -3017,6 +3014,9 @@ def _build_forward_metadata( use_prefill_cuda_graph=use_prefill_cuda_graph, forward_batch=forward_batch, ) + self._snapshot_low_ratio_k_slots( + forward_batch, metadata, publish_rows=not use_prefill_cuda_graph + ) else: raise NotImplementedError(f"unsupported mode {forward_batch.forward_mode=}") @@ -3140,7 +3140,6 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay( if capture_metadata.prefill_graph_tail else None ) - self._snapshot_low_ratio_k_slots(metadata_batch) def _prefill_graph_tail_metadata( self, forward_batch: ForwardBatch, max_seq_len: int @@ -3843,7 +3842,7 @@ def _make_low_ratio_prefill_indexer_inputs( if k_slots is not None: k_slots = k_slots[layer.compress_ratio] else: - # No snapshot, e.g. the prefill graph's capture pass. + # Only a non-speculative extend's metadata carries a snapshot. k_slots = prefill_k_slots( req_to_token=self.req_to_token, req_pool_indices=forward_batch.req_pool_indices, diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py index 848b66a58df1..a36df1b43261 100644 --- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py +++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/types.py @@ -48,8 +48,8 @@ class PrefillInputs(msgspec.Struct, frozen=True, kw_only=True): req_rows: Optional[torch.Tensor] req_pool_indices: torch.Tensor # [requests] req_to_token row of each request # [columns] int64 index-K pool slot of each flattened-K column, from - # scoring.prefill_k_slots; gathered with the metadata, per layer only when the - # metadata has no snapshot (e.g. the prefill graph's capture pass). + # scoring.prefill_k_slots; gathered once with a non-speculative extend's + # metadata, per layer otherwise. k_slots: torch.Tensor # ratio -> scoring.PrefillIndexPlan, filled by the first layer of each ratio and # reused by the rest of the forward; None when the metadata has no snapshot, so diff --git a/python/sglang/srt/models/deepseek_v4_replay_tail.py b/python/sglang/srt/models/deepseek_v4_replay_tail.py index c83e1bf17a40..b5e552929818 100644 --- a/python/sglang/srt/models/deepseek_v4_replay_tail.py +++ b/python/sglang/srt/models/deepseek_v4_replay_tail.py @@ -630,7 +630,6 @@ def _enter_eager_tail_in_prefill_graph( backend.tail_forward_metadata = backend._build_late_layer_tail_metadata( forward_batch ) - backend._share_low_ratio_k_slots(backend.tail_forward_metadata) try: saved = backend.enter_late_layer_tail(forward_batch) finally: diff --git a/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py b/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py index da1f68b3357c..b227a612f784 100644 --- a/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py +++ b/test/registered/unit/layers/test_dsv41_late_layer_tail_readers.py @@ -291,7 +291,6 @@ def build_tail(forward_batch, max_seq_len): backend._build_forward_metadata = build_full backend._prefill_graph_tail_metadata = build_tail backend._activate_request_window = lambda: None - backend._snapshot_low_ratio_k_slots = lambda forward_batch: None backend._prefill_graph_replays_any_context = lambda: any_context capture = object.__new__(DSV4Metadata) capture.prefill_graph_tail = True