minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
+11 -6
View File
@@ -831,12 +831,17 @@ 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
kv_lens = seq_lens_cpu[:bs] # recomputes it): gathering seq_lens[:bs] here would force a D2H
else: # sync for models (e.g. DSA+MTP) whose replay path is otherwise
kv_lens = seq_lens[:bs] # sync-free.
metadata.seq_lens_cpu_list = kv_lens.cpu().int().tolist() if _is_dflash_verify(spec_info):
if seq_lens_cpu is not None:
kv_lens = seq_lens_cpu[:bs]
else:
kv_lens = seq_lens[:bs]
metadata.seq_lens_cpu_list = kv_lens.cpu().int().tolist()
elif forward_mode.is_decode_or_idle() and spec_info is not None: elif forward_mode.is_decode_or_idle() and spec_info is not None:
seq_lens = seq_lens + self.speculative_step_offset_npu seq_lens = seq_lens + self.speculative_step_offset_npu
metadata.seq_lens[:bs].copy_(seq_lens[:bs]) metadata.seq_lens[:bs].copy_(seq_lens[:bs])
python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
+128 -92
View File
@@ -243,100 +243,129 @@ 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:
# In speculative decoding, these two fields are still needed. if not self.model_runner.spec_algorithm.is_dflash():
# NPU skips the DFLASH verify pre-planning, so load_batch may # In speculative decoding, these two fields are still needed.
# never have recorded the padded batch shapes; recompute them on # Non-DFlash keeps the historical pre-planned replay path: no
# every batch (the verify batch size varies with concurrency). # attention-metadata refresh, no seq_lens_cpu work, and no
raw_bs = forward_batch.batch_size # device sync (DSA/DSV4 rely on this staying sync-free).
if self.require_mlp_tp_gather: self.buffers.input_ids[: self.raw_num_token].copy_(
bs = self._pad_to_bucket( forward_batch.input_ids
self._max_dp_batch_size(forward_batch), self.capture_bs
) )
else: self.buffers.positions[: self.raw_num_token].copy_(
bs = self._pad_to_bucket(raw_bs, self.capture_bs) forward_batch.positions
self.raw_bs = raw_bs )
self.raw_num_token = raw_bs * self.captured_req_width if (
self.bs = bs envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get()
# Restore the DeepEP dispatch mode recorded at capture time and forward_batch.mrope_positions is not None
# (mirrors load_batch); an interleaved eager extend may have ):
# switched it. self.buffers.mrope_positions[:, : self.raw_num_token].copy_(
self.deepep_adapter.replay() forward_batch.mrope_positions
# Refresh the static DP token buffers bound by the captured
# graph (stale values misalign dp-gather segments across ranks);
# mirror the capture-side uniform [padded_num_tokens] * dp_size.
if self.require_mlp_tp_gather:
_padded_num_tokens = bs * self.captured_req_width
self.buffers.global_num_tokens_gpu.fill_(_padded_num_tokens)
self.buffers.global_num_tokens_for_logprob_gpu.fill_(_padded_num_tokens)
if (
enable_num_token_non_padded()
and self.require_gathered_buffer
and not self.enable_prefill_cp
):
self.buffers.num_token_non_padded.fill_(
compute_local_num_token_non_padded_cpu(
global_num_token_non_padded=(
forward_batch.global_num_token_non_padded_cpu
),
num_tokens_per_dp=bs * self.captured_req_width,
sharded=self.model_runner.attn_tp_sequence_sharded(
bs * self.captured_req_width
),
) )
else:
# NPU skips the DFLASH verify pre-planning, so load_batch may
# never have recorded the padded batch shapes; recompute them
# on every batch (the verify batch size varies with
# concurrency).
raw_bs = forward_batch.batch_size
if self.require_mlp_tp_gather:
bs = self._pad_to_bucket(
self._max_dp_batch_size(forward_batch), self.capture_bs
)
else:
bs = self._pad_to_bucket(raw_bs, self.capture_bs)
self.raw_bs = raw_bs
self.raw_num_token = raw_bs * self.captured_req_width
self.bs = bs
# Restore the DeepEP dispatch mode recorded at capture time
# (mirrors load_batch); an interleaved eager extend may have
# switched it.
self.deepep_adapter.replay()
# Refresh the static DP token buffers bound by the captured
# graph (stale values misalign dp-gather segments across
# ranks); mirror the capture-side uniform
# [padded_num_tokens] * dp_size.
if self.require_mlp_tp_gather:
_padded_num_tokens = bs * self.captured_req_width
self.buffers.global_num_tokens_gpu.fill_(_padded_num_tokens)
self.buffers.global_num_tokens_for_logprob_gpu.fill_(
_padded_num_tokens
)
if (
enable_num_token_non_padded()
and self.require_gathered_buffer
and not self.enable_prefill_cp
):
self.buffers.num_token_non_padded.fill_(
compute_local_num_token_non_padded_cpu(
global_num_token_non_padded=(
forward_batch.global_num_token_non_padded_cpu
),
num_tokens_per_dp=bs * self.captured_req_width,
sharded=self.model_runner.attn_tp_sequence_sharded(
bs * self.captured_req_width
),
)
)
self.buffers.input_ids[: self.raw_num_token].copy_(
forward_batch.input_ids
) )
self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(
self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) forward_batch.positions
if (
self.model_runner.spec_algorithm.is_dflash()
and self.model_runner.is_draft_worker
and forward_batch.input_embeds is not None
):
self.buffers.input_embeds[: self.raw_num_token].copy_(
forward_batch.input_embeds
)
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
) )
if (
self.model_runner.is_draft_worker
and forward_batch.input_embeds is not None
):
self.buffers.input_embeds[: self.raw_num_token].copy_(
forward_batch.input_embeds
)
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
)
# The pre-planned path skipped init_forward_metadata_out_graph; # The pre-planned path skipped init_forward_metadata_out_graph;
# refresh attention metadata so replay reads correct KV pages. # refresh attention metadata so replay reads correct KV pages.
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.buffers.seq_lens_cpu[: self.raw_bs].copy_( self.seq_len_fill_value
forward_batch.seq_lens_cpu[: self.raw_bs] )
) self.buffers.seq_lens_cpu[: self.raw_bs].copy_(
self.buffers.seq_lens_cpu[self.raw_bs : self.bs].fill_( forward_batch.seq_lens_cpu[: self.raw_bs]
self.seq_len_fill_value )
) self.buffers.seq_lens_cpu[self.raw_bs : self.bs].fill_(
self.buffers.req_pool_indices[: self.raw_bs].copy_( self.seq_len_fill_value
forward_batch.req_pool_indices[: self.raw_bs] )
) self.buffers.req_pool_indices[: self.raw_bs].copy_(
self.buffers.req_pool_indices[self.raw_bs : self.bs].fill_(0) forward_batch.req_pool_indices[: self.raw_bs]
# Refresh the static out_cache_loc bound by the captured graph )
# for full-pool KV writes in save_kv_cache (replay would self.buffers.req_pool_indices[self.raw_bs : self.bs].fill_(0)
# otherwise write verify KV to stale capture-time slots). # Refresh the static out_cache_loc bound by the captured graph
if forward_batch.out_cache_loc is not None: # for full-pool KV writes in save_kv_cache (replay would
_padded_num_token = self.bs * self.captured_req_width # otherwise write verify KV to stale capture-time slots).
_n = min(self.raw_num_token, forward_batch.out_cache_loc.shape[0]) if forward_batch.out_cache_loc is not None:
self.buffers.out_cache_loc[:_n].copy_(forward_batch.out_cache_loc[:_n]) _padded_num_token = self.bs * self.captured_req_width
self.buffers.out_cache_loc[_n:_padded_num_token].zero_() _n = min(self.raw_num_token, forward_batch.out_cache_loc.shape[0])
fb_view = build_replay_fb_view( self.buffers.out_cache_loc[:_n].copy_(
forward_batch=forward_batch, forward_batch.out_cache_loc[:_n]
buffers=self.buffers, )
bs=self.bs, self.buffers.out_cache_loc[_n:_padded_num_token].zero_()
raw_bs=self.raw_bs, fb_view = build_replay_fb_view(
num_tokens=self.bs * self.captured_req_width, forward_batch=forward_batch,
seq_len_fill_value=self.seq_len_fill_value, buffers=self.buffers,
capture_forward_mode=self.capture_forward_mode, bs=self.bs,
is_encoder_decoder=self.is_encoder_decoder, raw_bs=self.raw_bs,
) num_tokens=self.bs * self.captured_req_width,
self._replay_attn_backend().init_forward_metadata_out_graph(fb_view) seq_len_fill_value=self.seq_len_fill_value,
capture_forward_mode=self.capture_forward_mode,
is_encoder_decoder=self.is_encoder_decoder,
)
self._replay_attn_backend().init_forward_metadata_out_graph(fb_view)
graph_key = self._make_graph_key(self.bs) graph_key = self._make_graph_key(self.bs)
@@ -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():
_attn = self._replay_attn_backend() # Only DFlash refreshes forward_metadata.seq_lens_cpu_list at
_meta = getattr(_attn, "forward_metadata", None) # replay; other algorithms keep the capture-time list, so
_meta_list = getattr(_meta, "seq_lens_cpu_list", None) # 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()
_meta = getattr(_attn, "forward_metadata", 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