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)
test/registered/attention/unittests/dsv4/test_deepseek_v4.py
+109
View File
@@ -391,6 +391,113 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
DeepseekV4AttnBackend.use_captured_forward_metadata_for_breakable_cuda_graph
)
def test_prefill_snapshot_declares_pre_replay_boundary(self):
from sglang.srt.layers.attention.base_attn_backend import SharedReadEnds
from sglang.srt.layers.attention.deepseek_v4_backend import (
DeepseekV4AttnBackend,
DSV4Metadata,
)
backend = object.__new__(DeepseekV4AttnBackend)
backend.forward_metadata = DSV4Metadata(
self._make_core_metadata(0), indexer_metadata=None
)
self.assertIs(
backend.shared_read_ends(ForwardMode.EXTEND),
SharedReadEnds.UNKNOWN,
)
backend.forward_metadata.prefill_shared_reads_snapshotted = True
self.assertIs(
backend.shared_read_ends(ForwardMode.EXTEND),
SharedReadEnds.PRE_REPLAY,
)
def test_snapshot_builds_cache_only_for_sparse_prefill(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.deepseek_v4_backend import (
_LARGE_INDEXER_QUERY_THRESHOLD,
DeepseekV4AttnBackend,
DSV4Metadata,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND)
cache = object()
for num_qo_tokens, builds in (
(_LARGE_INDEXER_QUERY_THRESHOLD, False),
(_LARGE_INDEXER_QUERY_THRESHOLD + 1, True),
):
with self.subTest(num_qo_tokens=num_qo_tokens):
backend = object.__new__(DeepseekV4AttnBackend)
backend.model_runner = SimpleNamespace(
spec_algorithm=SpeculativeAlgorithm.DFLASH
)
backend.forward_metadata = DSV4Metadata(
self._make_core_metadata(0), indexer_metadata=None
)
backend._build_sparse_prefill_chunk_cache = mock.Mock(
return_value=cache
)
with (
envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True),
envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.override(False),
mock.patch(
"sglang.srt.layers.attention.deepseek_v4_backend._is_sm120",
False,
),
):
backend.prepare_prefill_shared_read_snapshot(
batch, num_qo_tokens=num_qo_tokens
)
metadata = backend.forward_metadata
if builds:
backend._build_sparse_prefill_chunk_cache.assert_called_once_with(
batch, num_qo_tokens=num_qo_tokens
)
self.assertIs(metadata.sparse_prefill_cache, cache)
else:
backend._build_sparse_prefill_chunk_cache.assert_not_called()
self.assertIsNone(metadata.sparse_prefill_cache)
# Dense declares the boundary too; it reads only the metadata
# that init_forward_metadata already snapshotted.
self.assertTrue(metadata.prefill_shared_reads_snapshotted)
def test_sparse_prefill_snapshot_marks_success_only_after_build(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.deepseek_v4_backend import (
DeepseekV4AttnBackend,
DSV4Metadata,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
backend = object.__new__(DeepseekV4AttnBackend)
backend.model_runner = SimpleNamespace(
spec_algorithm=SpeculativeAlgorithm.DFLASH
)
backend.forward_metadata = DSV4Metadata(
self._make_core_metadata(0), indexer_metadata=None
)
backend.forward_metadata.prefill_shared_reads_snapshotted = True
backend._build_sparse_prefill_chunk_cache = mock.Mock(
side_effect=RuntimeError("snapshot failed")
)
batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND)
with (
envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True),
envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.override(True),
mock.patch(
"sglang.srt.layers.attention.deepseek_v4_backend._is_sm120", False
),
self.assertRaisesRegex(RuntimeError, "snapshot failed"),
):
backend.prepare_prefill_shared_read_snapshot(batch, num_qo_tokens=12288)
self.assertFalse(backend.forward_metadata.prefill_shared_reads_snapshotted)
def test_refresh_replay_metadata_preserves_captured_tensor_storage(self):
capture_metadata = self._make_core_metadata(0)
replay_metadata = self._make_core_metadata(1000)
@@ -459,6 +566,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
self._make_core_metadata(0), indexer_metadata=None
)
capture_metadata.sparse_prefill_cache = object()
capture_metadata.prefill_shared_reads_snapshotted = True
replay_metadata = DSV4Metadata(
self._make_core_metadata(1000), indexer_metadata=None
)
@@ -487,6 +595,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
self.assertTrue(calls[0][2])
self.assertIs(backend.forward_metadata, capture_metadata)
self.assertIsNone(capture_metadata.sparse_prefill_cache)
self.assertFalse(capture_metadata.prefill_shared_reads_snapshotted)
self.assertTrue(
torch.equal(
capture_metadata.core_attn_metadata.seq_lens_casual,
test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py
+20
View File
@@ -1,5 +1,6 @@
import unittest
from types import SimpleNamespace
from unittest import mock
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.forward_batch_info import (
@@ -52,6 +53,25 @@ class TestPrefillCudaGraphPadding(CustomTestCase):
self.assertTrue(runner.can_run_graph(self._make_forward_batch(8)))
def test_replay_snapshot_uses_padded_token_count(self):
runner = self._make_runner()
runner.use_captured_attn_metadata = False
attn_backend = mock.Mock()
runner.model_runner = SimpleNamespace(attn_backend=attn_backend)
forward_batch = self._make_forward_batch(8)
static_forward_batch = self._make_forward_batch(16)
runner._prepare_forward_metadata_for_replay(
forward_batch,
static_forward_batch,
num_tokens=16,
)
attn_backend.init_forward_metadata.assert_called_once_with(forward_batch)
attn_backend.prepare_prefill_shared_read_snapshot.assert_called_once_with(
forward_batch, num_qo_tokens=16
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py
+11
View File
@@ -60,6 +60,17 @@ def test_disabled_when_flag_is_false():
assert runner.shared_read_done_event is None
@pytest.mark.parametrize(
"algorithm", (SpeculativeAlgorithm.DFLASH, SpeculativeAlgorithm.DSPARK)
)
def test_dflash_family_target_prefill_publishes(algorithm):
runner = _model_runner(spec_algorithm=algorithm)
with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True):
maybe_publish_prefill_shared_read_done(runner, _batch(), _DEVICE_MODULE)
published = runner.shared_read_done_event
assert isinstance(published, _Event) and published.recorded
def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backend():
with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True):
for runner, batch in (