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

[NPU] Gate DFlash replay metadata refresh behind spec_algorithm check (#39879)

Browse Source
This commit is contained in:
iridiumine
2026-09-18 17:07:07 +08:00
committed by GitHub
parent bbfcda48ce
commit 7e6d5cbfac
2 changed files with 139 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+7 -2
View File
@@ -831,8 +831,13 @@ class AscendAttnBackend(AttentionBackend):
if forward_mode.is_target_verify(): if forward_mode.is_target_verify():
seq_lens = seq_lens + self.speculative_num_draft_tokens seq_lens = seq_lens + self.speculative_num_draft_tokens
# For DFlash, seq_lens_cpu (= prefix + block_size) is the true KV # For DFlash, seq_lens_cpu (= prefix + block_size) is the true KV
# length; other spec algorithms already added the draft tokens above. # length; refresh the graph-bound list from that CPU tensor.
if _is_dflash_verify(spec_info) and seq_lens_cpu is not None: # Other spec algorithms keep the capture-time list (graph_runner
# recomputes it): gathering seq_lens[:bs] here would force a D2H
# sync for models (e.g. DSA+MTP) whose replay path is otherwise
# sync-free.
if _is_dflash_verify(spec_info):
if seq_lens_cpu is not None:
kv_lens = seq_lens_cpu[:bs] kv_lens = seq_lens_cpu[:bs]
else: else:
kv_lens = seq_lens[:bs] kv_lens = seq_lens[:bs]
python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
+47 -11
View File
@@ -243,10 +243,29 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
if forward_batch.needs_forward_metadata_init(): if forward_batch.needs_forward_metadata_init():
self.load_batch(forward_batch, pp_proxy_tensors) self.load_batch(forward_batch, pp_proxy_tensors)
else: else:
if not self.model_runner.spec_algorithm.is_dflash():
# In speculative decoding, these two fields are still needed. # In speculative decoding, these two fields are still needed.
# Non-DFlash keeps the historical pre-planned replay path: no
# attention-metadata refresh, no seq_lens_cpu work, and no
# device sync (DSA/DSV4 rely on this staying sync-free).
self.buffers.input_ids[: self.raw_num_token].copy_(
forward_batch.input_ids
)
self.buffers.positions[: self.raw_num_token].copy_(
forward_batch.positions
)
if (
envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get()
and forward_batch.mrope_positions is not None
):
self.buffers.mrope_positions[:, : self.raw_num_token].copy_(
forward_batch.mrope_positions
)
else:
# NPU skips the DFLASH verify pre-planning, so load_batch may # NPU skips the DFLASH verify pre-planning, so load_batch may
# never have recorded the padded batch shapes; recompute them on # never have recorded the padded batch shapes; recompute them
# every batch (the verify batch size varies with concurrency). # on every batch (the verify batch size varies with
# concurrency).
raw_bs = forward_batch.batch_size raw_bs = forward_batch.batch_size
if self.require_mlp_tp_gather: if self.require_mlp_tp_gather:
bs = self._pad_to_bucket( bs = self._pad_to_bucket(
@@ -262,12 +281,15 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
# switched it. # switched it.
self.deepep_adapter.replay() self.deepep_adapter.replay()
# Refresh the static DP token buffers bound by the captured # Refresh the static DP token buffers bound by the captured
# graph (stale values misalign dp-gather segments across ranks); # graph (stale values misalign dp-gather segments across
# mirror the capture-side uniform [padded_num_tokens] * dp_size. # ranks); mirror the capture-side uniform
# [padded_num_tokens] * dp_size.
if self.require_mlp_tp_gather: if self.require_mlp_tp_gather:
_padded_num_tokens = bs * self.captured_req_width _padded_num_tokens = bs * self.captured_req_width
self.buffers.global_num_tokens_gpu.fill_(_padded_num_tokens) self.buffers.global_num_tokens_gpu.fill_(_padded_num_tokens)
self.buffers.global_num_tokens_for_logprob_gpu.fill_(_padded_num_tokens) self.buffers.global_num_tokens_for_logprob_gpu.fill_(
_padded_num_tokens
)
if ( if (
enable_num_token_non_padded() enable_num_token_non_padded()
and self.require_gathered_buffer and self.require_gathered_buffer
@@ -284,11 +306,14 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
), ),
) )
) )
self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.input_ids[: self.raw_num_token].copy_(
self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) forward_batch.input_ids
)
self.buffers.positions[: self.raw_num_token].copy_(
forward_batch.positions
)
if ( if (
self.model_runner.spec_algorithm.is_dflash() self.model_runner.is_draft_worker
and self.model_runner.is_draft_worker
and forward_batch.input_embeds is not None and forward_batch.input_embeds is not None
): ):
self.buffers.input_embeds[: self.raw_num_token].copy_( self.buffers.input_embeds[: self.raw_num_token].copy_(
@@ -307,7 +332,9 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
self.buffers.seq_lens[: self.raw_bs].copy_( self.buffers.seq_lens[: self.raw_bs].copy_(
forward_batch.seq_lens_cpu[: self.raw_bs] forward_batch.seq_lens_cpu[: self.raw_bs]
) )
self.buffers.seq_lens[self.raw_bs : self.bs].fill_(self.seq_len_fill_value) self.buffers.seq_lens[self.raw_bs : self.bs].fill_(
self.seq_len_fill_value
)
self.buffers.seq_lens_cpu[: self.raw_bs].copy_( self.buffers.seq_lens_cpu[: self.raw_bs].copy_(
forward_batch.seq_lens_cpu[: self.raw_bs] forward_batch.seq_lens_cpu[: self.raw_bs]
) )
@@ -324,7 +351,9 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
if forward_batch.out_cache_loc is not None: if forward_batch.out_cache_loc is not None:
_padded_num_token = self.bs * self.captured_req_width _padded_num_token = self.bs * self.captured_req_width
_n = min(self.raw_num_token, forward_batch.out_cache_loc.shape[0]) _n = min(self.raw_num_token, forward_batch.out_cache_loc.shape[0])
self.buffers.out_cache_loc[:_n].copy_(forward_batch.out_cache_loc[:_n]) self.buffers.out_cache_loc[:_n].copy_(
forward_batch.out_cache_loc[:_n]
)
self.buffers.out_cache_loc[_n:_padded_num_token].zero_() self.buffers.out_cache_loc[_n:_padded_num_token].zero_()
fb_view = build_replay_fb_view( fb_view = build_replay_fb_view(
forward_batch=forward_batch, forward_batch=forward_batch,
@@ -345,9 +374,16 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
or is_deepseek_v4(self.model_runner.model_config.hf_config) or is_deepseek_v4(self.model_runner.model_config.hf_config)
): ):
if forward_batch.forward_mode.is_target_verify(): if forward_batch.forward_mode.is_target_verify():
# Only DFlash refreshes forward_metadata.seq_lens_cpu_list at
# replay; other algorithms keep the capture-time list, so
# recompute from the live batch as before — reading the
# backend list there would replay stale capture-time lengths.
if self.model_runner.spec_algorithm.is_dflash():
_attn = self._replay_attn_backend() _attn = self._replay_attn_backend()
_meta = getattr(_attn, "forward_metadata", None) _meta = getattr(_attn, "forward_metadata", None)
_meta_list = getattr(_meta, "seq_lens_cpu_list", None) _meta_list = getattr(_meta, "seq_lens_cpu_list", None)
else:
_meta_list = None
if _meta_list is not None: if _meta_list is not None:
# graph.update must carry the exact KV length already # graph.update must carry the exact KV length already
# computed in forward_metadata.seq_lens_cpu_list (it # computed in forward_metadata.seq_lens_cpu_list (it