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

[refactor] init_forward_metadata 3-method ABC + side-channel removal + ForwardMetadata type rename (#26735)

Browse Source 
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-06-02 10:33:33 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 6c69756fa8
commit 99da43b900
42 changed files with 1999 additions and 2081 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/attention/test_trtllm_mla_backend.py
+11 -11
View File
@@ -1,5 +1,6 @@
import math
import unittest
from types import SimpleNamespace
import numpy as np
import torch
@@ -1030,15 +1031,15 @@ class TestTRTLLMMLA(CustomTestCase):
)
req_pool_indices = torch.arange(batch_size, device=config["device"])
backend.init_forward_metadata_capture_cuda_graph(
bs=batch_size,
num_tokens=batch_size,
capture_fb = SimpleNamespace(
batch_size=batch_size,
forward_mode=ForwardMode.DECODE,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
encoder_lens=None,
forward_mode=ForwardMode.DECODE,
positions=torch.arange(batch_size, device=config["device"]),
spec_info=None,
)
backend.init_forward_metadata_out_graph(capture_fb, in_capture=True)
# Verify capture metadata
self.assertIn(batch_size, backend.decode_cuda_graph_metadata)
@@ -1054,16 +1055,15 @@ class TestTRTLLMMLA(CustomTestCase):
device=config["device"],
)
backend.init_forward_metadata_replay_cuda_graph(
bs=batch_size,
replay_fb = SimpleNamespace(
batch_size=batch_size,
forward_mode=ForwardMode.DECODE,
req_pool_indices=req_pool_indices,
seq_lens=new_seq_lens,
seq_lens_sum=new_seq_lens.sum().item(),
encoder_lens=None,
forward_mode=ForwardMode.DECODE,
spec_info=None,
seq_lens_cpu=new_seq_lens.cpu(),
spec_info=None,
)
backend.init_forward_metadata_out_graph(replay_fb)
# Verify replay updated the metadata
replay_metadata = backend.forward_decode_metadata
test/registered/attention/unittests/dense/test_extend_init_contract.py
+147
View File
@@ -0,0 +1,147 @@
"""Unit tests for the init contract that piecewise + breakable cuda graph
capture rely on with ``forward_mode=EXTEND``.
The piecewise + breakable runners (unlike the full ``cuda_graph_runner``)
capture prefill chunks, which means their capture path passes
``forward_mode=ForwardMode.EXTEND`` to ``attn_backend.init_forward_metadata*``.
Backends like FlashInfer and FA3 implement two separate init bodies:
- ``init_forward_metadata(fb)`` — the eager entry. Handles all modes
including plain ``EXTEND`` (full prefill / chunked prefill).
- ``init_forward_metadata_out_graph(fb, in_capture=True)`` — the
bucket-keyed wrapper prep used by the full cuda graph runner. Only
handles modes the full runner captures: ``DECODE`` / ``IDLE`` /
``TARGET_VERIFY`` / ``DRAFT_EXTEND`` / ``DLLM_EXTEND``. Plain
``EXTEND`` is not in scope here and the body raises on it.
These tests pin both halves of that contract so a future refactor that
incorrectly routes piecewise/breakable capture through ``_out_graph(in_capture=True)``
fails at unit-test time instead of at GPU CI / e2e time. This is the
specific regression #26735 introduced and then fixed
(see ``piecewise_cuda_graph_runner.py`` /
``breakable_cuda_graph_runner.py`` capture sites).
"""
import sys
import unittest
from pathlib import Path
import torch
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import get_device_sm
from sglang.test.test_utils import CustomTestCase
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
DenseAttentionCase,
build_dense_attention_fixture,
)
register_cuda_ci(est_time=10, stage="base-a", runner_config="1-gpu-small")
_EXTEND_CASE = DenseAttentionCase(
name="extend_no_prefix_smoke",
backend="fa3", # overridden per backend below
forward_mode=ForwardMode.EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(0,),
extend_lens=(16,),
)
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
class TestExtendInitContract(CustomTestCase):
def _make_case(self, backend: str) -> DenseAttentionCase:
return DenseAttentionCase(
name=f"extend_no_prefix_{backend}",
backend=backend,
forward_mode=_EXTEND_CASE.forward_mode,
num_heads=_EXTEND_CASE.num_heads,
num_kv_heads=_EXTEND_CASE.num_kv_heads,
page_size=_EXTEND_CASE.page_size,
prefix_lens=_EXTEND_CASE.prefix_lens,
extend_lens=_EXTEND_CASE.extend_lens,
)
def _build_fixture(self, backend: str, *, head_dim: int = 16):
case = self._make_case(backend)
try:
return build_dense_attention_fixture(self, case, head_dim=head_dim)
except (AssertionError, ImportError, ModuleNotFoundError) as exc:
self.skipTest(f"backend {backend} unavailable: {exc}")
def _assert_extend_eager_init_well_formed(
self, backend: str, *, head_dim: int = 16
):
fixture = self._build_fixture(backend, head_dim=head_dim)
fixture.backend.init_forward_metadata(fixture.forward_batch)
meta = fixture.backend.forward_metadata
self.assertIsNotNone(
meta,
f"{backend}: init_forward_metadata(EXTEND) left forward_metadata as None — "
"piecewise/breakable capture would crash on the first forward_extend call.",
)
page_table = getattr(meta, "page_table", None)
if page_table is not None:
self.assertIsInstance(page_table, torch.Tensor)
@unittest.skipIf(
get_device_sm() >= 100 or get_device_sm() < 80,
"FA3 backend requires SM 80-90",
)
def test_fa3_extend_eager_init(self):
self._assert_extend_eager_init_well_formed("fa3")
def test_flashinfer_extend_eager_init(self):
# FlashInfer's JIT prefill kernel needs head_dim ≥ 64.
self._assert_extend_eager_init_well_formed("flashinfer", head_dim=128)
def test_triton_extend_eager_init(self):
self._assert_extend_eager_init_well_formed("triton")
def _assert_out_graph_in_capture_rejects_extend(
self, backend: str, *, head_dim: int = 16
):
"""``init_forward_metadata_out_graph(fb, in_capture=True)`` is the
bucket-prep path. Plain EXTEND mode isn't in its supported set —
piecewise/breakable capture used to route through here and crash.
This pins the constraint: if the path starts handling EXTEND
(raises become passes), revisit the piecewise/breakable capture
wiring and consider unifying the API.
"""
fixture = self._build_fixture(backend, head_dim=head_dim)
try:
fixture.backend.init_forward_metadata_out_graph(
fixture.forward_batch, in_capture=True
)
except (ValueError, AttributeError, KeyError, AssertionError):
return
meta = fixture.backend.forward_metadata
self.assertIsNotNone(
meta,
f"{backend}: _out_graph(in_capture=True) accepted EXTEND but left "
"forward_metadata as None — this is the specific bug pattern that "
"broke piecewise capture (FA3._apply_cuda_graph_metadata falls "
"through for unsupported modes and sets self.forward_metadata = None).",
)
@unittest.skipIf(
get_device_sm() >= 100 or get_device_sm() < 80,
"FA3 backend requires SM 80-90",
)
def test_fa3_out_graph_capture_rejects_extend(self):
self._assert_out_graph_in_capture_rejects_extend("fa3")
def test_flashinfer_out_graph_capture_rejects_extend(self):
self._assert_out_graph_in_capture_rejects_extend("flashinfer", head_dim=128)
if __name__ == "__main__":
unittest.main()
test/registered/attention/unittests/dense/test_tbo.py
+96 -14
View File
@@ -94,14 +94,15 @@ class TestTboAttnDenseAttentionBackendCorrectness(CustomTestCase):
"FA3 backend requires SM 80-90",
)
def test_tbo_target_verify_cuda_graph_capture_delegates_to_primary_capture(self):
"""TBO capture must invoke ``primary.init_forward_metadata_capture_cuda_graph``,
not ``primary.init_forward_metadata_replay_cuda_graph``.
"""TBO capture must dispatch primary's
``init_forward_metadata_out_graph(fb, in_capture=True)`` (the capture
path), not the replay path.
Backends like FlashAttention store per-bs metadata in dicts populated
only by their capture path (via ``_bind_metadata_buffers``). If TBO
short-circuits its capture to its own replay (which delegates to
``primary.replay``), those dicts are empty and replay raises
``KeyError: bs``. Reproduces the deepep-4-gpu-h100 failure where
only by the in_capture=True branch (via ``_bind_metadata_buffers``).
If TBO short-circuits its capture to its own replay path, those dicts
are empty and replay raises ``KeyError: bs``. Reproduces the
deepep-4-gpu-h100 failure where
``flashattention_backend.target_verify_metadata[bs]`` lookup blew up
during ``init_device_graphs``.
@@ -126,19 +127,100 @@ class TestTboAttnDenseAttentionBackendCorrectness(CustomTestCase):
capture_bs = case.batch_size
num_tokens = sum(case.extend_lens)
wrapper.init_cuda_graph_state(max_bs=capture_bs, max_num_tokens=num_tokens)
# This is the failing call before the fix: TBO.capture delegating to
# primary.replay (instead of primary.capture) reads an unpopulated
# ``target_verify_metadata[bs]`` dict and raises KeyError.
wrapper.init_forward_metadata_capture_cuda_graph(
bs=capture_bs,
num_tokens=num_tokens,
wrapper.init_forward_metadata_out_graph(batch, in_capture=True)
@unittest.skipIf(
get_device_sm() >= 100 or get_device_sm() < 80,
"FA3 backend requires SM 80-90",
)
def test_tbo_target_verify_cuda_graph_replay_splits_children(self):
"""TBO replay must split the padded capture-time buffers into per-child
views before dispatching to ``init_forward_metadata_out_graph(fb_view,
in_capture=False)``.
``cuda_graph_runner.replay_prepare`` constructs a ``SimpleNamespace``
fb_view via ``build_replay_fb_view`` — it has no ``tbo_children``
attribute because ``tbo_plugin.replay_prepare`` does not call
``prepare_raw``. Without the in-line split, TBO either crashes
(AttributeError on ``fb_view.tbo_children``) or silently leaves child
metadata stale.
Asserts both children's ``init_forward_metadata_out_graph`` is invoked
with sliced ``req_pool_indices`` / ``seq_lens`` / ``seq_lens_cpu``
whose lengths match the split derived from
``compute_split_indices_for_cuda_graph_replay``.
"""
from types import SimpleNamespace
from unittest.mock import MagicMock
from sglang.srt.batch_overlap.two_batch_overlap import (
compute_split_indices_for_cuda_graph_replay,
)
case = self.TARGET_VERIFY_CAPTURE_CASE
fixture = self._build_and_wrap(case)
wrapper = fixture.backend
batch = fixture.forward_batch
_prepare_spec_verify_batch(
case,
batch,
topk=1,
spec_kind="eagle",
device=str(batch.seq_lens.device),
)
capture_bs = case.batch_size
num_tokens_per_bs = sum(case.extend_lens) // capture_bs
num_tokens = capture_bs * num_tokens_per_bs
split_seq_index, split_token_index = (
compute_split_indices_for_cuda_graph_replay(
forward_mode=batch.forward_mode,
cuda_graph_num_tokens=num_tokens,
spec_info=batch.spec_info,
)
)
self.assertGreater(split_seq_index, 0)
self.assertLess(split_seq_index, capture_bs)
# fb_view shaped like build_replay_fb_view's output: a SimpleNamespace
# with no tbo_children attribute.
fb_view = SimpleNamespace(
batch_size=capture_bs,
forward_mode=batch.forward_mode,
actual_forward_mode=batch.forward_mode,
input_ids=batch.input_ids,
req_pool_indices=batch.req_pool_indices,
seq_lens=batch.seq_lens,
encoder_lens=batch.encoder_lens,
forward_mode=batch.forward_mode,
seq_lens_sum=int(batch.seq_lens_cpu.sum()),
seq_lens_cpu=batch.seq_lens_cpu,
encoder_lens=None,
out_cache_loc=batch.out_cache_loc,
spec_info=batch.spec_info,
)
# Pure mocks (no `wraps=...`) so the dispatcher's slicing/contract is
# observed without invoking real backend bodies.
primary_mock = MagicMock()
child_mocks = [MagicMock(), MagicMock()]
wrapper.primary = primary_mock
wrapper.children = child_mocks
wrapper.init_forward_metadata_out_graph(fb_view, in_capture=False)
primary_mock.init_forward_metadata_out_graph.assert_called_once()
for child_mock in child_mocks:
child_mock.init_forward_metadata_out_graph.assert_called_once()
child_fbs = [
m.init_forward_metadata_out_graph.call_args.kwargs["forward_batch"]
for m in child_mocks
]
self.assertEqual(child_fbs[0].batch_size, split_seq_index)
self.assertEqual(child_fbs[1].batch_size, capture_bs - split_seq_index)
self.assertEqual(child_fbs[0].req_pool_indices.shape[0], split_seq_index)
self.assertEqual(
child_fbs[1].req_pool_indices.shape[0], capture_bs - split_seq_index
)
if __name__ == "__main__":
unittest.main()
test/registered/attention/unittests/gdn/test_triton.py
+22 -40
View File
@@ -399,38 +399,35 @@ class TestTritonGDNBackendCorrectness(CustomTestCase):
linear_attn_backend.init_forward_metadata, sentinel_forward_batch
)
def _make_sentinel_fb(self):
return SimpleNamespace(
batch_size=3,
forward_mode=ForwardMode.DECODE,
req_pool_indices=object(),
seq_lens=object(),
seq_lens_cpu=object(),
seq_lens_sum=42,
spec_info=object(),
encoder_lens=None,
positions=object(),
input_ids=object(),
out_cache_loc=None,
)
def test_hybrid_dispatch_replay_init_forward_metadata_fan_out(self):
backend, full_attn_backend, linear_attn_backend = (
self._make_dispatch_spy_backend()
)
sentinel_req_pool = object()
sentinel_seq_lens = object()
sentinel_seq_lens_cpu = object()
sentinel_spec_info = object()
backend.init_forward_metadata_replay_cuda_graph(
bs=3,
req_pool_indices=sentinel_req_pool,
seq_lens=sentinel_seq_lens,
seq_lens_sum=42,
encoder_lens=None,
forward_mode=ForwardMode.DECODE,
spec_info=sentinel_spec_info,
seq_lens_cpu=sentinel_seq_lens_cpu,
)
fb = self._make_sentinel_fb()
backend.init_forward_metadata_out_graph(fb)
# We assert sentinel identity rather than exact (args, kwargs) shape
# so a positional↔keyword refactor inside `HybridLinearAttnBackend`
# doesn't trip the test as long as the values still flow through.
# doesn't trip the test as long as the fb still flows through.
for sub_backend in (full_attn_backend, linear_attn_backend):
self._assert_fanout_forwarded(
sub_backend.init_forward_metadata_replay_cuda_graph,
sentinel_req_pool,
sentinel_seq_lens,
sentinel_seq_lens_cpu,
sentinel_spec_info,
ForwardMode.DECODE,
sub_backend.init_forward_metadata_out_graph, fb
)
def test_hybrid_dispatch_capture_init_forward_metadata_fan_out(self):
@@ -439,27 +436,12 @@ class TestTritonGDNBackendCorrectness(CustomTestCase):
backend, full_attn_backend, linear_attn_backend = (
self._make_dispatch_spy_backend()
)
sentinel_req_pool = object()
sentinel_seq_lens = object()
sentinel_spec_info = object()
backend.init_forward_metadata_capture_cuda_graph(
bs=3,
num_tokens=3,
req_pool_indices=sentinel_req_pool,
seq_lens=sentinel_seq_lens,
encoder_lens=None,
forward_mode=ForwardMode.DECODE,
spec_info=sentinel_spec_info,
)
fb = self._make_sentinel_fb()
backend.init_forward_metadata_out_graph(fb, in_capture=True)
for sub_backend in (full_attn_backend, linear_attn_backend):
self._assert_fanout_forwarded(
sub_backend.init_forward_metadata_capture_cuda_graph,
sentinel_req_pool,
sentinel_seq_lens,
sentinel_spec_info,
ForwardMode.DECODE,
sub_backend.init_forward_metadata_out_graph, fb
)
test/registered/attention/unittests/mamba/test_mamba2.py
+29 -46
View File
@@ -216,7 +216,7 @@ class TestTritonMamba2BackendCorrectness(CustomTestCase):
# padding the `forward_batch.input_ids` / `out_cache_loc`.
def test_mamba2_replay_metadata_padding_indices(self):
# Drive `init_forward_metadata_replay_cuda_graph` directly with
# Drive `init_forward_metadata_out_graph` (replay path) directly with
# `seq_lens_cpu=[5, 1, 1]` (two trailing rows at the cuda-graph
# fill value 1) so the padding-row count is observable in
# `state_indices_list[bs - 1]`.
@@ -245,16 +245,17 @@ class TestTritonMamba2BackendCorrectness(CustomTestCase):
req_pool_indices
] = torch.tensor([7, 0, 0], dtype=torch.int32, device=device)
backend.init_forward_metadata_replay_cuda_graph(
bs=bs,
fb = SimpleNamespace(
batch_size=bs,
forward_mode=ForwardMode.DECODE,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_sum=int(seq_lens_cpu.sum().item()),
encoder_lens=None,
forward_mode=ForwardMode.DECODE,
spec_info=None,
seq_lens_cpu=seq_lens_cpu,
seq_lens_sum=int(seq_lens_cpu.sum().item()),
spec_info=None,
encoder_lens=None,
)
backend.init_forward_metadata_out_graph(fb)
state_indices = backend.state_indices_list[bs - 1].cpu().tolist()
self.assertEqual(
@@ -326,62 +327,44 @@ class TestTritonMamba2BackendCorrectness(CustomTestCase):
linear_attn_backend.init_forward_metadata, sentinel_forward_batch
)
def _make_sentinel_fb(self):
return SimpleNamespace(
batch_size=3,
forward_mode=ForwardMode.DECODE,
req_pool_indices=object(),
seq_lens=object(),
seq_lens_cpu=object(),
seq_lens_sum=42,
spec_info=object(),
encoder_lens=None,
positions=object(),
input_ids=object(),
out_cache_loc=None,
)
def test_hybrid_dispatch_replay_init_forward_metadata_fan_out(self):
backend, full_attn_backend, linear_attn_backend = (
self._make_dispatch_spy_backend()
)
sentinel_req_pool = object()
sentinel_seq_lens = object()
sentinel_seq_lens_cpu = object()
sentinel_spec_info = object()
backend.init_forward_metadata_replay_cuda_graph(
bs=3,
req_pool_indices=sentinel_req_pool,
seq_lens=sentinel_seq_lens,
seq_lens_sum=42,
encoder_lens=None,
forward_mode=ForwardMode.DECODE,
spec_info=sentinel_spec_info,
seq_lens_cpu=sentinel_seq_lens_cpu,
)
fb = self._make_sentinel_fb()
backend.init_forward_metadata_out_graph(fb)
for sub_backend in (full_attn_backend, linear_attn_backend):
self._assert_fanout_forwarded(
sub_backend.init_forward_metadata_replay_cuda_graph,
sentinel_req_pool,
sentinel_seq_lens,
sentinel_seq_lens_cpu,
sentinel_spec_info,
ForwardMode.DECODE,
sub_backend.init_forward_metadata_out_graph, fb
)
def test_hybrid_dispatch_capture_init_forward_metadata_fan_out(self):
backend, full_attn_backend, linear_attn_backend = (
self._make_dispatch_spy_backend()
)
sentinel_req_pool = object()
sentinel_seq_lens = object()
sentinel_spec_info = object()
backend.init_forward_metadata_capture_cuda_graph(
bs=3,
num_tokens=3,
req_pool_indices=sentinel_req_pool,
seq_lens=sentinel_seq_lens,
encoder_lens=None,
forward_mode=ForwardMode.DECODE,
spec_info=sentinel_spec_info,
)
fb = self._make_sentinel_fb()
backend.init_forward_metadata_out_graph(fb, in_capture=True)
for sub_backend in (full_attn_backend, linear_attn_backend):
self._assert_fanout_forwarded(
sub_backend.init_forward_metadata_capture_cuda_graph,
sentinel_req_pool,
sentinel_seq_lens,
sentinel_spec_info,
ForwardMode.DECODE,
sub_backend.init_forward_metadata_out_graph, fb
)