From 1aea16883f7545e628d68d84123d9ef28b476f09 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 19 Oct 2026 20:22:55 -0700 Subject: [PATCH 1/3] allocate the padded query inside decoder tail graphs --- python/sglang/srt/models/deepseek_v4.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 362ea60bf0cf..3bb4396f0f61 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -2547,6 +2547,10 @@ def forward( # zero. The buffer is reused per layer; no consumer may keep it. want = (x.shape[0], kernel_num_heads, self.head_dim) meta = getattr(attn_backend, "forward_metadata", None) + if in_replay_tail_graph(): + # Every bucket a step captures shares its metadata, so a + # kept buffer would be freed under the earlier buckets' graphs. + meta = None q_padded = getattr(meta, "q_pad_buffer", None) if ( q_padded is None From 0ea54451e437d2c82eecefdc8184923431bdc606 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 19 Oct 2026 20:29:53 -0700 Subject: [PATCH 2/3] move the padded query allocation into a method --- python/sglang/srt/models/deepseek_v4.py | 45 ++++++++++++++----------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 3bb4396f0f61..0194871dbb5f 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1104,6 +1104,29 @@ def _kernel_num_heads(self, num_tokens: int) -> int: # Other FlashMLA implementations retain their existing padded shape. return 64 if self.n_local_heads <= 64 else self.n_heads + def _zero_padded_q(self, x, kernel_num_heads, attn_backend) -> torch.Tensor: + want = (x.shape[0], kernel_num_heads, self.head_dim) + if in_replay_tail_graph(): + # A tail graph step captures every bucket on one metadata, so a buffer + # kept there would be freed under the earlier buckets' graphs; each + # capture allocates its own from the graph pool. + return x.new_zeros(*want) + # Eager steps reuse one buffer per step and shape; no consumer may keep it. + meta = getattr(attn_backend, "forward_metadata", None) + q_padded = getattr(meta, "q_pad_buffer", None) + if ( + q_padded is None + or tuple(q_padded.shape) != want + or q_padded.dtype != x.dtype + ): + q_padded = x.new_zeros(*want) + if meta is not None: + try: + meta.q_pad_buffer = q_padded + except (AttributeError, TypeError): + pass + return q_padded + def _wo_a_quant_freqs( self, forward_batch: ForwardBatch, q_rope: Optional[torch.Tensor] ) -> Optional[torch.Tensor]: @@ -2543,26 +2566,8 @@ def forward( # Backends without an exact-head specialization retain the existing # padded shape. attn_sink is sliced to this rank and padded to match. if self.is_dsv41: - # V4.1 kernels read all padded heads, so the padding must be - # zero. The buffer is reused per layer; no consumer may keep it. - want = (x.shape[0], kernel_num_heads, self.head_dim) - meta = getattr(attn_backend, "forward_metadata", None) - if in_replay_tail_graph(): - # Every bucket a step captures shares its metadata, so a - # kept buffer would be freed under the earlier buckets' graphs. - meta = None - q_padded = getattr(meta, "q_pad_buffer", None) - if ( - q_padded is None - or tuple(q_padded.shape) != want - or q_padded.dtype != x.dtype - ): - q_padded = x.new_zeros(*want) - if meta is not None: - try: - meta.q_pad_buffer = q_padded - except (AttributeError, TypeError): - pass + # V4.1 kernels read all padded heads, so the padding must be zero. + q_padded = self._zero_padded_q(x, kernel_num_heads, attn_backend) elif _is_gfx942_supported: # Uninitialized padded TP heads inject NaN into attention on gfx942 # (fnuz), so zero-init there; other archs tolerate new_empty and skip From 117247ea2f9ba77cc8d47005553d3eed18e3a7bc Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Mon, 19 Oct 2026 21:08:42 -0700 Subject: [PATCH 3/3] keep one growing padded query buffer per metadata; assert it does not grow after a capture --- .../layers/attention/deepseek_v4_backend.py | 2 ++ .../deepseek_v4_backend_hip_radix.py | 7 +++- python/sglang/srt/models/deepseek_v4.py | 33 +++++++++++-------- 3 files changed, 27 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 4ebb31c22136..9b6dcd634c86 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -1046,6 +1046,8 @@ class DSV4Metadata: # Per-step scratch for TP-padded query heads, zeroed by the first user. # Later layers overwrite real heads and preserve the zero padding. q_pad_buffer: Optional[torch.Tensor] = None + # Set once a CUDA graph captures q_pad_buffer; it may not grow after that. + q_pad_captured: bool = False # What the candidate-source layer published for the index-source layers after # it, in the implementation's own type; never copied from the host. diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 74597972ae67..a71891e51257 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -691,6 +691,8 @@ class DSV4Metadata: low_ratio_pos_i64: Optional[torch.Tensor] = None # Per-step scratch for the TP-padded query heads (models/deepseek_v4.py). q_pad_buffer: Optional[torch.Tensor] = None + # Set once a CUDA graph captures q_pad_buffer; it may not grow after that. + q_pad_captured: bool = False # low-ratio FlyDSL indexer workspaces by ratio; address-pinned like the c4 ones, rebuilt in place fp4_low_ratio_decode_workspaces: Dict[ int, Union[FP4DecodeWorkspace, FP4RowgroupDecodeWorkspace] @@ -746,7 +748,10 @@ def copy_(self, other: DSV4Metadata): ) # allocated inside a captured segment (graph-pool storage), so a refresh must keep it - BREAKABLE_CUDA_GRAPH_KEEP_FIELDS: ClassVar[Tuple[str, ...]] = ("q_pad_buffer",) + BREAKABLE_CUDA_GRAPH_KEEP_FIELDS: ClassVar[Tuple[str, ...]] = ( + "q_pad_buffer", + "q_pad_captured", + ) def refresh_for_breakable_cuda_graph_replay_(self, other: DSV4Metadata) -> None: self.core_attn_metadata.refresh_for_breakable_cuda_graph_replay_( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 0194871dbb5f..2cbd84fd21fa 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1105,27 +1105,32 @@ def _kernel_num_heads(self, num_tokens: int) -> int: return 64 if self.n_local_heads <= 64 else self.n_heads def _zero_padded_q(self, x, kernel_num_heads, attn_backend) -> torch.Tensor: - want = (x.shape[0], kernel_num_heads, self.head_dim) - if in_replay_tail_graph(): - # A tail graph step captures every bucket on one metadata, so a buffer - # kept there would be freed under the earlier buckets' graphs; each - # capture allocates its own from the graph pool. - return x.new_zeros(*want) - # Eager steps reuse one buffer per step and shape; no consumer may keep it. + # One zeroed buffer per metadata, grown to the most rows it has served; each + # call takes a row slice. A graph that keeps the metadata keeps the buffer, + # so the buffer must not grow once a graph has captured it. meta = getattr(attn_backend, "forward_metadata", None) - q_padded = getattr(meta, "q_pad_buffer", None) + buf = getattr(meta, "q_pad_buffer", None) if ( - q_padded is None - or tuple(q_padded.shape) != want - or q_padded.dtype != x.dtype + buf is None + or buf.shape[0] < x.shape[0] + or buf.shape[1:] != (kernel_num_heads, self.head_dim) + or buf.dtype != x.dtype ): - q_padded = x.new_zeros(*want) + assert not getattr(meta, "q_pad_captured", False), ( + "the padded query buffer grew after a CUDA graph captured it" + ) + buf = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim) if meta is not None: try: - meta.q_pad_buffer = q_padded + meta.q_pad_buffer = buf except (AttributeError, TypeError): pass - return q_padded + if meta is not None and torch.cuda.is_current_stream_capturing(): + try: + meta.q_pad_captured = True + except (AttributeError, TypeError): + pass + return buf[: x.shape[0]] def _wo_a_quant_freqs( self, forward_batch: ForwardBatch, q_rope: Optional[torch.Tensor]