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

Publish gated DSV4 DFLASH-family target-prefill read completion (#35947)

Browse Source 
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-08-26 21:33:08 -07:00
committed by GitHub
co-authored by weireweire
parent 7324021e6c
commit 4d5d506486
8 changed files with 249 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/base_attn_backend.py
+10
View File
@@ -152,6 +152,16 @@ class AttentionBackend(ABC):
return SharedReadEnds.IN_REPLAY
return SharedReadEnds.UNKNOWN
def prepare_prefill_shared_read_snapshot(
self, forward_batch: ForwardBatch, *, num_qo_tokens: int
) -> None:
"""Snapshot late prefill reads before a PRE_REPLAY event is published.
Runners call this only after the actual eager/replay query geometry is
known. Backends that retain scheduler-shared reads into the model
forward keep the default no-op and must not declare PRE_REPLAY.
"""
# Chunked-prefix FullCG capture has a second model topology and stable
# prefix buffers. Backends must opt in explicitly so the runner does not
# assume that generic ForwardBatch metadata is sufficient for every
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+76 -47
View File
@@ -408,10 +408,10 @@ class DSV4Metadata:
c4_compress_metadata: Optional[FusedCompressMetadata] = None
c128_compress_metadata: Optional[FusedCompressMetadata] = None
# Lazily populated on the first call to ``_forward_prefill_sparse`` and
# reused across every layer in the chunk. Reset to ``None`` when graph
# metadata is refreshed so replay rebuilds it from the live batch.
# Built at the runner's prefill WAR boundary when the fast path is on,
# otherwise lazily by ``_forward_prefill_sparse``.
sparse_prefill_cache: Optional[SparsePrefillChunkCache] = None
prefill_shared_reads_snapshotted: bool = False
@property
def core_metadata(self) -> DSV4AttnMetadata:
@@ -425,6 +425,7 @@ class DSV4Metadata:
self.c128_compress_metadata, src=other.c128_compress_metadata
)
self.sparse_prefill_cache = None
self.prefill_shared_reads_snapshotted = False
def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata):
self.core_attn_metadata.refresh_for_breakable_cuda_graph_replay_(
@@ -444,6 +445,7 @@ class DSV4Metadata:
src=static_metadata.c128_compress_metadata,
)
self.sparse_prefill_cache = None
self.prefill_shared_reads_snapshotted = False
@dataclass
@@ -519,6 +521,13 @@ class DeepseekV4AttnBackend(
if self.model_runner.spec_algorithm.is_dspark():
return SharedReadEnds.IN_REPLAY
return SharedReadEnds.POST_REPLAY
metadata = self.forward_metadata
if (
fm == ForwardMode.EXTEND
and isinstance(metadata, DSV4Metadata)
and metadata.prefill_shared_reads_snapshotted
):
return SharedReadEnds.PRE_REPLAY
return super().shared_read_ends(fm)
def __init__(
@@ -1352,6 +1361,66 @@ class DeepseekV4AttnBackend(
self.forward_metadata = self._build_forward_metadata(forward_batch)
self.init_forward_metadata_in_graph(forward_batch)
def prepare_prefill_shared_read_snapshot(
self, forward_batch: ForwardBatch, *, num_qo_tokens: int
) -> None:
# Sparse prefill otherwise reads req_to_token/full_to_swa lazily in its
# first layer. DFLASH/DSPARK have no later prefill draft-extend reader;
# CP-v2 shards the query layout that this global snapshot assumes.
metadata = self.forward_metadata
if isinstance(metadata, DSV4Metadata):
metadata.prefill_shared_reads_snapshotted = False
snapshot_shared_prefill_reads = (
envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.get()
and forward_batch.forward_mode == ForwardMode.EXTEND
and self.model_runner.spec_algorithm.is_dflash_family()
and not is_cp_v2_active(forward_batch)
)
if not snapshot_shared_prefill_reads:
return
assert isinstance(metadata, DSV4Metadata)
use_sparse_prefill = not _is_sm120 and (
num_qo_tokens > _LARGE_INDEXER_QUERY_THRESHOLD
or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get()
)
if use_sparse_prefill:
metadata.sparse_prefill_cache = self._build_sparse_prefill_chunk_cache(
forward_batch, num_qo_tokens=num_qo_tokens
)
# Marked for dense prefill too: that path reads only core_attn_metadata,
# which init_forward_metadata already snapshotted.
metadata.prefill_shared_reads_snapshotted = True
def _build_sparse_prefill_chunk_cache(
self, forward_batch: ForwardBatch, *, num_qo_tokens: int
) -> SparsePrefillChunkCache:
seq_lens_cpu = forward_batch.seq_lens_cpu
assert seq_lens_cpu is not None
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
assert extend_seq_lens_cpu is not None
seq_lens_cpu_list = seq_lens_cpu.tolist()
total_swa = sum(
min(int(seq_len), int(extend_len) + SWA_WINDOW - 1)
for seq_len, extend_len in zip(
seq_lens_cpu_list, extend_seq_lens_cpu, strict=True
)
)
# ``swa_window_size`` on the pool is its storage page size, not the
# model's SWA window, so pass both explicitly.
return SparsePrefillChunkCache.build(
seq_lens=forward_batch.seq_lens.to(torch.int32),
extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32),
req_pool_indices=forward_batch.req_pool_indices.to(torch.int32),
req_to_token=self.req_to_token,
full_to_swa=self.token_to_kv_pool.full_to_swa_index_mapping,
swa_window_size=SWA_WINDOW,
swa_page_size=self.token_to_kv_pool.swa_window_size,
num_qo_tokens=num_qo_tokens,
max_seq_len=max(seq_lens_cpu_list),
total_swa=total_swa,
)
def _build_forward_metadata(
self,
forward_batch: ForwardBatch,
@@ -1791,29 +1860,8 @@ class DeepseekV4AttnBackend(
cache = self.forward_metadata.sparse_prefill_cache
if cache is None:
seq_lens_cpu = forward_batch.seq_lens_cpu
assert seq_lens_cpu is not None
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
assert extend_seq_lens_cpu is not None
total_swa = sum(
min(int(seq_len), int(extend_len) + SWA_WINDOW - 1)
for seq_len, extend_len in zip(
seq_lens_cpu.tolist(), extend_seq_lens_cpu, strict=True
)
)
# ``swa_window_size`` on the pool is its storage page size, not
# the model's SWA window — pass both explicitly.
cache = SparsePrefillChunkCache.build(
seq_lens=forward_batch.seq_lens.to(torch.int32),
extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32),
req_pool_indices=forward_batch.req_pool_indices.to(torch.int32),
req_to_token=self.req_to_token,
full_to_swa=token_to_kv_pool.full_to_swa_index_mapping,
swa_window_size=SWA_WINDOW,
swa_page_size=token_to_kv_pool.swa_window_size,
num_qo_tokens=q_flat.shape[0],
max_seq_len=int(seq_lens_cpu.max().item()),
total_swa=total_swa,
cache = self._build_sparse_prefill_chunk_cache(
forward_batch, num_qo_tokens=q_flat.shape[0]
)
self.forward_metadata.sparse_prefill_cache = cache
@@ -1987,27 +2035,8 @@ class DeepseekV4AttnBackend(
cache = self.forward_metadata.sparse_prefill_cache
if cache is None:
seq_lens_cpu = forward_batch.seq_lens_cpu
assert seq_lens_cpu is not None
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
assert extend_seq_lens_cpu is not None
total_swa = sum(
min(int(seq_len), int(extend_len) + SWA_WINDOW - 1)
for seq_len, extend_len in zip(
seq_lens_cpu.tolist(), extend_seq_lens_cpu, strict=True
)
)
cache = SparsePrefillChunkCache.build(
seq_lens=forward_batch.seq_lens.to(torch.int32),
extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32),
req_pool_indices=forward_batch.req_pool_indices.to(torch.int32),
req_to_token=self.req_to_token,
full_to_swa=token_to_kv_pool.full_to_swa_index_mapping,
swa_window_size=SWA_WINDOW,
swa_page_size=token_to_kv_pool.swa_window_size,
num_qo_tokens=q_flat.shape[0],
max_seq_len=int(seq_lens_cpu.max().item()),
total_swa=total_swa,
cache = self._build_sparse_prefill_chunk_cache(
forward_batch, num_qo_tokens=q_flat.shape[0]
)
self.forward_metadata.sparse_prefill_cache = cache
python/sglang/srt/model_executor/runner/eager_runner.py
+12
View File
@@ -49,6 +49,9 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
enable_tc_piecewise_cuda_graph,
set_tc_piecewise_forward_context,
)
from sglang.srt.model_executor.runner_utils import (
maybe_publish_prefill_shared_read_done,
)
from sglang.srt.runtime_context import (
get_parallel,
get_spec,
@@ -305,6 +308,15 @@ class EagerRunner(BaseRunner):
# e.g. Moss-VL's prefill cross-attention custom mask.
model_runner.model.prepare_forward_batch(forward_batch)
model_runner.attn_backend.init_forward_metadata(forward_batch)
model_runner.attn_backend.prepare_prefill_shared_read_snapshot(
forward_batch,
num_qo_tokens=len(forward_batch.input_ids),
)
maybe_publish_prefill_shared_read_done(
model_runner,
forward_batch,
torch.get_device_module(model_runner.device),
)
if not cp_v2_active:
forward_batch.attn_cp_metadata = None
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+3
View File
@@ -1044,6 +1044,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
return
if not self.use_captured_attn_metadata:
attn_backend.init_forward_metadata(forward_batch)
attn_backend.prepare_prefill_shared_read_snapshot(
forward_batch, num_qo_tokens=num_tokens
)
return
assert self.attn_metadata_buffers is not None
metadata = self.attn_metadata_buffers[num_tokens]
python/sglang/srt/model_executor/runner_utils/shared_read_event.py
+8 -4
View File
@@ -1,4 +1,4 @@
"""Shared-read-done event utilities for CUDA graph runners."""
"""Shared-read-done event utilities for graph and eager runners."""
import logging
from typing import Optional
@@ -31,9 +31,13 @@ def maybe_publish_prefill_shared_read_done(
return
if forward_batch.forward_mode != ForwardMode.EXTEND:
return
# TODO(Jialin): Relax this gate for speculative decoding after its prefill
# WAR boundaries are validated.
if not model_runner.spec_algorithm.is_none():
# TODO(Jialin): Relax for EAGLE/MTP after validating the later
# draft-extend reader's WAR boundary.
if (
not model_runner.spec_algorithm.is_none()
and not model_runner.spec_algorithm.is_dflash_family()
):
# Other speculative algorithms may have a later draft-extend reader.
return
# The record lands right after replay prep, so PRE_REPLAY only.
declared = model_runner.attn_backend.shared_read_ends(forward_batch.forward_mode)