minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Perf] Hoist DSv4 draft-extend SWA write locs; unify SWA graph buffer naming (#34890)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-18 00:13:28 -07:00
committed by GitHub
parent a779a2a2a5
commit 0111b29031
4 changed files with 68 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+4 -4
View File
@@ -607,7 +607,7 @@ class AscendAttnBackend(AttentionBackend):
) )
if self.use_sliding_window_kv_pool: if self.use_sliding_window_kv_pool:
# refilled in place at replay; the captured graph reads this storage # 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, max_num_tokens,
dtype=torch.int64, dtype=torch.int64,
device=self.device, device=self.device,
@@ -638,7 +638,7 @@ class AscendAttnBackend(AttentionBackend):
metadata.swa_mask = self.graph_metadata["swa_mask"][:bs, :, :] metadata.swa_mask = self.graph_metadata["swa_mask"][:bs, :, :]
if self.use_sliding_window_kv_pool and out_cache_loc is not None: if self.use_sliding_window_kv_pool and out_cache_loc is not None:
num_tokens = out_cache_loc.shape[0] 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_cpu_list = seq_lens.cpu().int().tolist()
metadata.seq_lens = seq_lens metadata.seq_lens = seq_lens
if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): 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 # 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: if self.use_sliding_window_kv_pool and out_cache_loc is not None:
n = out_cache_loc.shape[0] n = out_cache_loc.shape[0]
self.swa_out_cache_loc_buf[n:].zero_() self.cuda_graph_swa_out_cache_loc[n:].zero_()
self.swa_out_cache_loc_buf[:n].copy_( self.cuda_graph_swa_out_cache_loc[:n].copy_(
self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc) self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc)
) )
max_len = seq_lens_cpu[:bs].max().item() max_len = seq_lens_cpu[:bs].max().item()
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+45 -9
View File
@@ -588,6 +588,7 @@ class DeepseekV4AttnBackend(
self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark() self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark()
self.is_draft_runner = model_runner.is_draft_worker self.is_draft_runner = model_runner.is_draft_worker
self._verify_mask = None 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: def _move_to_device(self, x: List[int]) -> torch.Tensor:
pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True)
@@ -1000,6 +1001,13 @@ class DeepseekV4AttnBackend(
) -> DSV4Metadata: ) -> DSV4Metadata:
batch_size = len(seq_lens) batch_size = len(seq_lens)
num_tokens = num_tokens_per_req * batch_size 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: if out_cache_loc is None:
out_cache_loc = seq_lens.new_zeros(num_tokens) out_cache_loc = seq_lens.new_zeros(num_tokens)
@@ -1022,11 +1030,36 @@ class DeepseekV4AttnBackend(
need_compress=False, need_compress=False,
is_prefill=True, 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( return DSV4Metadata(
core_attn_metadata=core_attn_metadata, core_attn_metadata=core_attn_metadata,
indexer_metadata=None, 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: def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None:
# Upgrade Raw->Full so the c4/c128 compress + core_attn + indexer # Upgrade Raw->Full so the c4/c128 compress + core_attn + indexer
# materialization is recorded inside the cuda graph; a no-op (Full # materialization is recorded inside the cuda graph; a no-op (Full
@@ -1446,6 +1479,12 @@ class DeepseekV4AttnBackend(
self.draft_extend_num_tokens_per_req = ( self.draft_extend_num_tokens_per_req = (
max_num_tokens // max_bs if max_bs > 0 else 1 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. # Verify metadata never extracts the mask. No skip_prefill notion here.
self._verify_mask = maybe_create_verify_mask( self._verify_mask = maybe_create_verify_mask(
is_draft_runner=self.is_draft_runner, 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: def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
"""Resolve the SWA KV-store write target for the current forward. """Resolve the SWA KV-store write target for the current forward.
Fast path: the per-forward value cached by init_forward_metadata_in_graph Prefer the value cached by the metadata init: in-graph for
(recorded inside cuda graphs, so replay re-reads live buffers). Fallback: decode/verify, the hoisted cuda_graph_swa_out_cache_loc buffer for
translate at store time, matching the pre-cache behavior, for paths that draft-extend. Translate at store time when nothing matching is cached
never run the in-graph init — eager idle (forward_idle skips attn init), (paths that skip the init, or a batch re-padded after init). Idle
runners that only run the out-graph prep (e.g. always falls back: its metadata may be stale, and
EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after translating the zero-padded out_cache_loc writes to the dummy slot.
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.
""" """
out_cache_loc = forward_batch.out_cache_loc out_cache_loc = forward_batch.out_cache_loc
core = getattr(self.forward_metadata, "core_attn_metadata", None) core = getattr(self.forward_metadata, "core_attn_metadata", None)
python/sglang/srt/layers/attention/flashattention_backend.py
+16 -8
View File
@@ -2184,7 +2184,7 @@ class FlashAttentionBackend(AttentionBackend):
) )
# SWA write-target buffer; metadata binds a [:num_tokens] view, # SWA write-target buffer; metadata binds a [:num_tokens] view,
# refilled from the live out_cache_loc before each replay. # 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, max_num_tokens,
dtype=torch.int64, dtype=torch.int64,
device=self.device, device=self.device,
@@ -2463,7 +2463,7 @@ class FlashAttentionBackend(AttentionBackend):
metadata.swa_page_table = self.decode_cuda_graph_metadata[ metadata.swa_page_table = self.decode_cuda_graph_metadata[
"swa_page_table" "swa_page_table"
][:bs, :] ][: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 :num_tokens
] ]
self.decode_cuda_graph_metadata[bs] = metadata self.decode_cuda_graph_metadata[bs] = metadata
@@ -2525,7 +2525,9 @@ class FlashAttentionBackend(AttentionBackend):
metadata.swa_page_table = self.decode_cuda_graph_metadata[ metadata.swa_page_table = self.decode_cuda_graph_metadata[
"swa_page_table" "swa_page_table"
][:bs, :] ][: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 self.decode_cuda_graph_metadata[bs] = metadata
elif forward_mode.is_target_verify(): elif forward_mode.is_target_verify():
@@ -2545,7 +2547,9 @@ class FlashAttentionBackend(AttentionBackend):
metadata.swa_page_table = self.target_verify_metadata[ metadata.swa_page_table = self.target_verify_metadata[
"swa_page_table" "swa_page_table"
][:bs, :] ][: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 self.target_verify_metadata[bs] = metadata
else: else:
# Target Verify topk>1: two (or three with SWA) metadata objects # 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 # topk>1 target-verify early-returns before _apply; bind the
# view here (buffer refilled at replay). # view here (buffer refilled at replay).
if self.use_sliding_window_kv_pool: 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: if self.has_swa:
metadata_swa = FlashAttentionMetadata() metadata_swa = FlashAttentionMetadata()
@@ -2621,7 +2627,9 @@ class FlashAttentionBackend(AttentionBackend):
metadata.swa_page_table = self.draft_extend_metadata["swa_page_table"][ metadata.swa_page_table = self.draft_extend_metadata["swa_page_table"][
:bs, : :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 self.draft_extend_metadata[bs] = metadata
if encoder_lens is not None: if encoder_lens is not None:
@@ -2684,8 +2692,8 @@ class FlashAttentionBackend(AttentionBackend):
# _bind_metadata_buffers) from the live out_cache_loc before replay. # _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: if self.use_sliding_window_kv_pool and out_cache_loc is not None:
n = out_cache_loc.shape[0] n = out_cache_loc.shape[0]
self.swa_out_cache_loc_buf[n:].zero_() self.cuda_graph_swa_out_cache_loc[n:].zero_()
self.swa_out_cache_loc_buf[:n].copy_( self.cuda_graph_swa_out_cache_loc[:n].copy_(
self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc) self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc)
) )
python/sglang/srt/layers/attention/xpu_backend.py
+3 -9
View File
@@ -78,12 +78,6 @@ class XPUAttentionBackend(AttentionBackend):
isinstance(model_runner.token_to_kv_pool, SWAKVPool) isinstance(model_runner.token_to_kv_pool, SWAKVPool)
and model_runner.token_to_kv_pool.swa_layer_nums > 0 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.topk = model_runner.server_args.speculative_eagle_topk or 0
self.speculative_num_steps = speculative_num_steps self.speculative_num_steps = speculative_num_steps
self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens 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 cu_seqlens_q = metadata.cu_seqlens_q
cache_seqlens_int32 = metadata.cache_seqlens_int32 cache_seqlens_int32 = metadata.cache_seqlens_int32
if self.is_hybrid_swa: if self.is_hybrid_swa:
page_table = self.full_to_swa_index_mapping[metadata.page_table].to( page_table = self.token_to_kv_pool.full_to_swa_index_mapping[
torch.int32 metadata.page_table
) ].to(torch.int32)
else: else:
page_table = metadata.page_table page_table = metadata.page_table
if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None: if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None: