diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 4cc25b584..356281ab7 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -607,7 +607,7 @@ class AscendAttnBackend(AttentionBackend): ) if self.use_sliding_window_kv_pool: # refilled in place at replay; the captured graph reads this storage - self.swa_out_cache_loc_buf = torch.zeros( + self.cuda_graph_swa_out_cache_loc = torch.zeros( max_num_tokens, dtype=torch.int64, device=self.device, @@ -638,7 +638,7 @@ class AscendAttnBackend(AttentionBackend): metadata.swa_mask = self.graph_metadata["swa_mask"][:bs, :, :] if self.use_sliding_window_kv_pool and out_cache_loc is not None: num_tokens = out_cache_loc.shape[0] - metadata.swa_out_cache_loc = self.swa_out_cache_loc_buf[:num_tokens] + metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[:num_tokens] metadata.seq_lens_cpu_list = seq_lens.cpu().int().tolist() metadata.seq_lens = seq_lens if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): @@ -712,8 +712,8 @@ class AscendAttnBackend(AttentionBackend): # refill the captured SWA write-target buffer in place from the live loc if self.use_sliding_window_kv_pool and out_cache_loc is not None: n = out_cache_loc.shape[0] - self.swa_out_cache_loc_buf[n:].zero_() - self.swa_out_cache_loc_buf[:n].copy_( + self.cuda_graph_swa_out_cache_loc[n:].zero_() + self.cuda_graph_swa_out_cache_loc[:n].copy_( self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc) ) max_len = seq_lens_cpu[:bs].max().item() diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 7dc9a7fe4..37081e06e 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -588,6 +588,7 @@ class DeepseekV4AttnBackend( self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark() self.is_draft_runner = model_runner.is_draft_worker self._verify_mask = None + self.cuda_graph_swa_out_cache_loc: Optional[torch.Tensor] = None def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) @@ -1000,6 +1001,13 @@ class DeepseekV4AttnBackend( ) -> DSV4Metadata: batch_size = len(seq_lens) num_tokens = num_tokens_per_req * batch_size + swa_out_cache_loc = self._fill_cuda_graph_swa_out_cache_loc(out_cache_loc) + if swa_out_cache_loc is None and out_cache_loc is not None: + # Eager-only miss (no graph state / oversized batch): translate once + # per step instead of per layer at store time. + swa_out_cache_loc = self.token_to_kv_pool.translate_loc_from_full_to_swa( + out_cache_loc + ).to(torch.int32) if out_cache_loc is None: out_cache_loc = seq_lens.new_zeros(num_tokens) @@ -1022,11 +1030,36 @@ class DeepseekV4AttnBackend( need_compress=False, is_prefill=True, ) + if swa_out_cache_loc is not None: + # Captures store_cache's cached path instead of a per-layer + # in-graph mapping translate. + core_attn_metadata.swa_out_cache_loc = swa_out_cache_loc return DSV4Metadata( core_attn_metadata=core_attn_metadata, indexer_metadata=None, ) + def _fill_cuda_graph_swa_out_cache_loc( + self, out_cache_loc: Optional[torch.Tensor] + ) -> Optional[torch.Tensor]: + # None (buffer absent / too small) is an eager-only miss: capture and + # replay always fit the pre-sized buffer. + buf = self.cuda_graph_swa_out_cache_loc + if ( + buf is None + or out_cache_loc is None + or out_cache_loc.shape[0] > buf.shape[0] + ): + return None + n = out_cache_loc.shape[0] + buf[n:].zero_() + buf[:n].copy_( + self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( + torch.int32 + ) + ) + return buf[:n] + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: # Upgrade Raw->Full so the c4/c128 compress + core_attn + indexer # materialization is recorded inside the cuda graph; a no-op (Full @@ -1446,6 +1479,12 @@ class DeepseekV4AttnBackend( self.draft_extend_num_tokens_per_req = ( max_num_tokens // max_bs if max_bs > 0 else 1 ) + if self.is_draft_runner: + # Draft-extend SWA write-target buffer; bound as a [:num_tokens] + # view and refilled outside the graph each step. + self.cuda_graph_swa_out_cache_loc = torch.zeros( + max_num_tokens, dtype=torch.int32, device=self.device + ) # Verify metadata never extracts the mask. No skip_prefill notion here. self._verify_mask = maybe_create_verify_mask( is_draft_runner=self.is_draft_runner, @@ -1502,15 +1541,12 @@ class DeepseekV4AttnBackend( def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: """Resolve the SWA KV-store write target for the current forward. - Fast path: the per-forward value cached by init_forward_metadata_in_graph - (recorded inside cuda graphs, so replay re-reads live buffers). Fallback: - translate at store time, matching the pre-cache behavior, for paths that - never run the in-graph init — eager idle (forward_idle skips attn init), - runners that only run the out-graph prep (e.g. - EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after - init (shape mismatch). Idle always falls back: its metadata is absent or - left over from a previous forward, and translating the zero-padded - out_cache_loc writes to the dummy slot. + Prefer the value cached by the metadata init: in-graph for + decode/verify, the hoisted cuda_graph_swa_out_cache_loc buffer for + draft-extend. Translate at store time when nothing matching is cached + (paths that skip the init, or a batch re-padded after init). Idle + always falls back: its metadata may be stale, and + translating the zero-padded out_cache_loc writes to the dummy slot. """ out_cache_loc = forward_batch.out_cache_loc core = getattr(self.forward_metadata, "core_attn_metadata", None) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 82189ff7d..60a46b146 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -2184,7 +2184,7 @@ class FlashAttentionBackend(AttentionBackend): ) # SWA write-target buffer; metadata binds a [:num_tokens] view, # refilled from the live out_cache_loc before each replay. - self.swa_out_cache_loc_buf = torch.zeros( + self.cuda_graph_swa_out_cache_loc = torch.zeros( max_num_tokens, dtype=torch.int64, device=self.device, @@ -2463,7 +2463,7 @@ class FlashAttentionBackend(AttentionBackend): metadata.swa_page_table = self.decode_cuda_graph_metadata[ "swa_page_table" ][:bs, :] - metadata.swa_out_cache_loc = self.swa_out_cache_loc_buf[ + metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[ :num_tokens ] self.decode_cuda_graph_metadata[bs] = metadata @@ -2525,7 +2525,9 @@ class FlashAttentionBackend(AttentionBackend): metadata.swa_page_table = self.decode_cuda_graph_metadata[ "swa_page_table" ][:bs, :] - metadata.swa_out_cache_loc = self.swa_out_cache_loc_buf[:num_tokens] + metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[ + :num_tokens + ] self.decode_cuda_graph_metadata[bs] = metadata elif forward_mode.is_target_verify(): @@ -2545,7 +2547,9 @@ class FlashAttentionBackend(AttentionBackend): metadata.swa_page_table = self.target_verify_metadata[ "swa_page_table" ][:bs, :] - metadata.swa_out_cache_loc = self.swa_out_cache_loc_buf[:num_tokens] + metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[ + :num_tokens + ] self.target_verify_metadata[bs] = metadata else: # Target Verify topk>1: two (or three with SWA) metadata objects @@ -2584,7 +2588,9 @@ class FlashAttentionBackend(AttentionBackend): # topk>1 target-verify early-returns before _apply; bind the # view here (buffer refilled at replay). if self.use_sliding_window_kv_pool: - metadata.swa_out_cache_loc = self.swa_out_cache_loc_buf[:num_tokens] + metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[ + :num_tokens + ] if self.has_swa: metadata_swa = FlashAttentionMetadata() @@ -2621,7 +2627,9 @@ class FlashAttentionBackend(AttentionBackend): metadata.swa_page_table = self.draft_extend_metadata["swa_page_table"][ :bs, : ] - metadata.swa_out_cache_loc = self.swa_out_cache_loc_buf[:num_tokens] + metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[ + :num_tokens + ] self.draft_extend_metadata[bs] = metadata if encoder_lens is not None: @@ -2684,8 +2692,8 @@ class FlashAttentionBackend(AttentionBackend): # _bind_metadata_buffers) from the live out_cache_loc before replay. if self.use_sliding_window_kv_pool and out_cache_loc is not None: n = out_cache_loc.shape[0] - self.swa_out_cache_loc_buf[n:].zero_() - self.swa_out_cache_loc_buf[:n].copy_( + self.cuda_graph_swa_out_cache_loc[n:].zero_() + self.cuda_graph_swa_out_cache_loc[:n].copy_( self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc) ) diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 9a4561921..4020c067f 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/python/sglang/srt/layers/attention/xpu_backend.py @@ -78,12 +78,6 @@ class XPUAttentionBackend(AttentionBackend): isinstance(model_runner.token_to_kv_pool, SWAKVPool) and model_runner.token_to_kv_pool.swa_layer_nums > 0 ) - if self.use_sliding_window_kv_pool: - self.token_to_kv_pool = model_runner.token_to_kv_pool - if self.is_hybrid_swa: - self.full_to_swa_index_mapping = ( - model_runner.token_to_kv_pool.full_to_swa_index_mapping - ) self.topk = model_runner.server_args.speculative_eagle_topk or 0 self.speculative_num_steps = speculative_num_steps self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens @@ -1215,9 +1209,9 @@ class XPUAttentionBackend(AttentionBackend): cu_seqlens_q = metadata.cu_seqlens_q cache_seqlens_int32 = metadata.cache_seqlens_int32 if self.is_hybrid_swa: - page_table = self.full_to_swa_index_mapping[metadata.page_table].to( - torch.int32 - ) + page_table = self.token_to_kv_pool.full_to_swa_index_mapping[ + metadata.page_table + ].to(torch.int32) else: page_table = metadata.page_table if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None: