From 2dfbc3d78189e2aae1eb6fa2b7e737aef89f9b4d Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 28 May 2026 22:43:29 -0700 Subject: [PATCH] test: strengthen CG-replay coverage with prod-fill padding, metadata invariants, and pad-ratio sweep (#26658) Co-authored-by: Claude Sonnet 4.6 --- .../runner_modes/cuda_graph_decode_runner.py | 5 + .../runner_modes/metadata_invariants.py | 103 ++++++++++++++ .../speculative_cuda_graph_runner.py | 133 +++++++++++++++++- .../speculative_draft_extend_runner.py | 11 ++ .../attention/unittests/dense/test_fa3.py | 27 +++- .../attention/unittests/dense/test_fa4.py | 27 +++- .../attention/unittests/dense/test_triton.py | 22 ++- .../attention/unittests/mla/test_triton.py | 20 ++- 8 files changed, 323 insertions(+), 25 deletions(-) create mode 100644 python/sglang/test/kits/attention_unittest/runner_modes/metadata_invariants.py diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py index 7e63246c5..8cd74bda3 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py @@ -239,6 +239,7 @@ from ..attention_methods.mla_attention import ( run_mla_fixture_eager, run_mla_forward, ) +from .metadata_invariants import assert_cg_metadata_well_formed DENSE_CUDA_GRAPH_CAPTURE_BATCH_SIZE = 4 MLA_CUDA_GRAPH_CAPTURE_BATCH_SIZE = 4 @@ -320,6 +321,10 @@ def _init_cuda_graph_replay_metadata(backend, capture_batch_size: int, batch): ) finally: backend._replay_forward_batch = None + # Best-effort metadata-shape sanity check — catches negative kv_lens and + # non-monotonic indptr that would otherwise leave real-row output correct + # but corrupt padded-row scratch state. See `metadata_invariants.py`. + assert_cg_metadata_well_formed(backend, bs=capture_batch_size) def _run_cuda_graph_decode_case( diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/metadata_invariants.py b/python/sglang/test/kits/attention_unittest/runner_modes/metadata_invariants.py new file mode 100644 index 000000000..a45660830 --- /dev/null +++ b/python/sglang/test/kits/attention_unittest/runner_modes/metadata_invariants.py @@ -0,0 +1,103 @@ +"""Backend-agnostic assertions on `forward_metadata` after CG replay init. + +Catches corruption that the output-equality assertion misses — e.g., negative +per-request lengths or non-monotonic indptr that happen to leave real-row +output correct while corrupting padded-row scratch state. + +Usage from a CG runner kit: + + from .metadata_invariants import assert_cg_metadata_well_formed + _init_cuda_graph_replay_metadata(backend, capture_batch_size, replay_batch) + assert_cg_metadata_well_formed(backend, bs=capture_batch_size) +""" + +from __future__ import annotations + +from typing import Any + +import torch + +# Field names that, when present on `forward_metadata`, denote a CSR-style +# `indptr` array. Convention: a length-(bs+1) int tensor whose first element is +# 0 and that is non-decreasing. +_INDPTR_FIELDS = ( + "kv_indptr", + "qo_indptr", + "mask_indptr", + "window_kv_indptr", + "cu_seqlens_q", + "cu_seqlens_k", + "encoder_cu_seqlens_k", +) + +# Field names that, when present, denote a length-bs per-request length tensor. +# Each element must be >= 0. +_PER_REQ_LEN_FIELDS = ( + "cache_seqlens_int32", + "encoder_lens_int32", + "local_seqused_k", + "max_seq_len_k", # may be scalar; checked defensively +) + + +def _slice_bs_plus_one(t: torch.Tensor, bs: int) -> torch.Tensor: + # An indptr conventionally has length bs+1; some backends may pre-size to + # max_bs+1 and slice on demand. Take the first bs+1 either way. + return t[: bs + 1] if t.numel() >= bs + 1 else t + + +def _slice_bs(t: torch.Tensor, bs: int) -> torch.Tensor: + return t[:bs] if t.numel() >= bs else t + + +def assert_cg_metadata_well_formed(backend: Any, bs: int) -> None: + """Inspect ``backend.forward_metadata`` and flag obviously-corrupt buffers. + + Best-effort: a field is checked only when it exists on the metadata object + and is a non-None tensor. Backends without a ``forward_metadata`` attribute + (or with one set to None) are skipped silently — the assertion only fires + on tensors that are clearly malformed. + """ + meta = getattr(backend, "forward_metadata", None) + if meta is None: + return + + errors: list[str] = [] + + for field in _INDPTR_FIELDS: + t = getattr(meta, field, None) + if not isinstance(t, torch.Tensor): + continue + sliced = _slice_bs_plus_one(t, bs) + if sliced.numel() < 2: + continue + diff = sliced[1:].to(torch.int64) - sliced[:-1].to(torch.int64) + # Allow zero-length requests (diff == 0); reject only negative diffs. + if (diff < 0).any().item(): + min_diff = diff.min().item() + errors.append( + f"{field} is not monotonic non-decreasing at bs={bs} " + f"(min adjacent diff={min_diff}); slice={sliced[: min(bs + 1, 16)].tolist()}" + ) + if sliced[0].item() != 0: + errors.append(f"{field}[0] != 0 (got {sliced[0].item()}) at bs={bs}") + + for field in _PER_REQ_LEN_FIELDS: + t = getattr(meta, field, None) + if not isinstance(t, torch.Tensor): + continue + sliced = _slice_bs(t, bs) + if sliced.numel() == 0: + continue + if (sliced < 0).any().item(): + min_v = sliced.min().item() + errors.append( + f"{field} has negative values at bs={bs} " + f"(min={min_v}); slice={sliced[: min(bs, 16)].tolist()}" + ) + + if errors: + raise AssertionError( + "CG forward_metadata invariants violated after replay init:\n - " + + "\n - ".join(errors) + ) diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_cuda_graph_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_cuda_graph_runner.py index a2ec3d968..faf1b59a5 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_cuda_graph_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_cuda_graph_runner.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import Any, Callable +from typing import Any, Callable, Literal, Optional import torch @@ -10,6 +10,24 @@ from .cuda_graph_decode_runner import ( _init_cuda_graph_replay_metadata, ) +# How padded rows of a CG replay batch are filled. +# +# "small_real" (default, legacy): padded rows look like miniature real +# requests — `seq_lens[padded] = capture_prefix_len + num_tokens_per_req`, +# `extend_seq_lens[padded] = num_tokens_per_req`. Easy for the reference +# module to compute an expected attention output for, but doesn't match +# production's padding shape. +# +# "prod_fill": mirrors `eagle_draft_extend_cuda_graph_runner.py:466-474` +# (and similar in `multi_layer_eagle_draft_extend_cuda_graph_runner.py`): +# padded rows are pure scratch — `seq_lens[padded] = seq_len_fill_value`, +# `extend_seq_lens[padded] = num_tokens_per_bs`, `req_pool_indices[padded] = 0`, +# `out_cache_loc[padded] = 0`, `positions[padded] = 0`. seq_lens and +# extend_seq_lens are intentionally inconsistent for padded rows (their +# subtraction goes negative), so backends must defend against that — the +# attention output for padded rows is never used in production. +PadStyle = Literal["small_real", "prod_fill"] + @dataclass(frozen=True) class SpeculativeCudaGraphAdapter: @@ -32,6 +50,72 @@ class SpeculativeCudaGraphAdapter: compare_replay_to_graph_eager: bool = True atol: float = 0.0 rtol: float = 0.0 + # Padding behavior for replay batches when raw_bs < capture_batch_size. + # See PadStyle docstring above for semantics. + pad_style: PadStyle = "small_real" + # Required when pad_style == "prod_fill": draft tokens per request, + # used to fill the padded slots of extend_seq_lens / spec_info. + pad_num_tokens_per_bs: Optional[int] = None + + +def _apply_prod_fill_padding( + batch, + *, + real_bs: int, + capture_bs: int, + seq_len_fill_value: int, + num_tokens_per_bs: int, +) -> None: + """Overwrite padded slots of `batch` to match the production CG runner. + + Production sets padded rows to: seq_lens = fill, extend_seq_lens = N, + req_pool_indices = 0, out_cache_loc = 0, positions = 0. The subtraction + `seq_lens - extend_seq_lens` then goes negative for padded rows — a + well-behaved backend must clamp or otherwise handle this. + """ + if real_bs >= capture_bs: + return + + pad_lo, pad_hi = real_bs, capture_bs + + # Per-request length tensors. + batch.seq_lens[pad_lo:pad_hi] = seq_len_fill_value + if batch.seq_lens_cpu is not None: + batch.seq_lens_cpu[pad_lo:pad_hi] = seq_len_fill_value + batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) + + if getattr(batch, "extend_seq_lens", None) is not None: + batch.extend_seq_lens[pad_lo:pad_hi] = num_tokens_per_bs + if getattr(batch, "extend_seq_lens_cpu", None) is not None: + ext = list(batch.extend_seq_lens_cpu) + for i in range(pad_lo, min(pad_hi, len(ext))): + ext[i] = num_tokens_per_bs + batch.extend_seq_lens_cpu = ext + + # Per-request slot tensors. + batch.req_pool_indices[pad_lo:pad_hi] = 0 + + # Per-token tensors: padded rows occupy slots + # [real_bs * num_tokens_per_bs, capture_bs * num_tokens_per_bs). + tok_lo = pad_lo * num_tokens_per_bs + tok_hi = pad_hi * num_tokens_per_bs + for field in ("out_cache_loc", "positions", "input_ids"): + t = getattr(batch, field, None) + if t is not None and t.numel() >= tok_hi: + t[tok_lo:tok_hi] = 0 + + # Mirror spec_info's per-request length tensor if present (set by the + # V2 production runner at replay: `spec_info.extend_seq_lens_tensor = + # buffers.extend_seq_lens[:bs]`). + spec_info = getattr(batch, "spec_info", None) + if spec_info is not None: + eslt = getattr(spec_info, "extend_seq_lens_tensor", None) + if isinstance(eslt, torch.Tensor) and eslt.numel() >= pad_hi: + eslt[pad_lo:pad_hi] = num_tokens_per_bs + eslc = getattr(spec_info, "extend_seq_lens_cpu", None) + if isinstance(eslc, list): + for i in range(pad_lo, min(pad_hi, len(eslc))): + eslc[i] = num_tokens_per_bs def _check_speculative_cuda_graph_case( @@ -201,6 +285,30 @@ def run_speculative_cuda_graph_case( graph_initial_state, ) + # Optionally overwrite padded rows to match the production CG runner's + # fill pattern. Done after prepare_batch + expected_output so that the + # reference is computed against the "small_real" layout (where padded + # rows have an interpretable attention output); the assertion below + # then compares only the real-row slice. + real_bs = case.batch_size + if ( + adapter.pad_style == "prod_fill" + and adapter.allow_padding + and real_bs < capture_batch_size + ): + if adapter.pad_num_tokens_per_bs is None: + raise ValueError( + "SpeculativeCudaGraphAdapter.pad_num_tokens_per_bs must be set " + "when pad_style='prod_fill'." + ) + _apply_prod_fill_padding( + replay_batch, + real_bs=real_bs, + capture_bs=capture_batch_size, + seq_len_fill_value=capture_prefix_len, + num_tokens_per_bs=adapter.pad_num_tokens_per_bs, + ) + with torch.no_grad(), forward_context(ForwardContext(attn_backend=backend)): _init_cuda_graph_replay_metadata(backend, capture_batch_size, replay_batch) replay_actual = adapter.run_forward( @@ -209,12 +317,23 @@ def run_speculative_cuda_graph_case( replay_inputs, ) - torch.testing.assert_close( - replay_actual, - replay_expected, - atol=adapter.atol, - rtol=adapter.rtol, - ) + if adapter.pad_style == "prod_fill": + # Padded rows have undefined output (their state is scratch in + # production; the runner discards their result). Assert only on the + # real-row slice. + torch.testing.assert_close( + replay_actual[: case.num_input_tokens], + replay_expected[: case.num_input_tokens], + atol=adapter.atol, + rtol=adapter.rtol, + ) + else: + torch.testing.assert_close( + replay_actual, + replay_expected, + atol=adapter.atol, + rtol=adapter.rtol, + ) if adapter.compare_replay_to_graph_eager: torch.testing.assert_close( replay_actual[: case.num_input_tokens], diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index dc52c9ab7..aaace8799 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -289,6 +289,8 @@ def _run_draft_extend_cuda_graph_case( max_num_tokens=None, run_graph_eager: bool = True, compare_replay_to_graph_eager: bool = True, + pad_style: str = "small_real", + pad_num_tokens_per_bs: int | None = None, ): adapter = SpeculativeCudaGraphAdapter( build_fixture=build_fixture, @@ -312,6 +314,8 @@ def _run_draft_extend_cuda_graph_case( compare_replay_to_graph_eager=compare_replay_to_graph_eager, atol=atol, rtol=rtol, + pad_style=pad_style, + pad_num_tokens_per_bs=pad_num_tokens_per_bs, ) run_speculative_cuda_graph_case( testcase, @@ -439,6 +443,7 @@ def run_dense_draft_extend_v2_cuda_graph_case( dtype: torch.dtype = DENSE_DEFAULT_DTYPE, device: str = DENSE_DEFAULT_DEVICE, cuda_graph_capture_batch_size: int = 4, + pad_style: str = "small_real", ): if not case.forward_mode.is_draft_extend_v2(): raise ValueError("Draft-extend-v2 CUDA graph coverage expects DRAFT_EXTEND_V2.") @@ -495,6 +500,8 @@ def run_dense_draft_extend_v2_cuda_graph_case( rtol=DENSE_RTOL, run_graph_eager=False, compare_replay_to_graph_eager=False, + pad_style=pad_style, + pad_num_tokens_per_bs=num_tokens_per_req, ) @@ -509,6 +516,7 @@ def run_mla_draft_extend_v2_cuda_graph_case( dtype: torch.dtype = MLA_DEFAULT_DTYPE, device: str = MLA_DEFAULT_DEVICE, cuda_graph_capture_batch_size: int = 4, + pad_style: str = "small_real", ): if not case.forward_mode.is_draft_extend_v2(): raise ValueError("Draft-extend-v2 CUDA graph coverage expects DRAFT_EXTEND_V2.") @@ -516,6 +524,7 @@ def run_mla_draft_extend_v2_cuda_graph_case( raise ValueError( "Draft-extend-v2 CUDA graph coverage uses a fixed token count per request." ) + num_tokens_per_req = case.input_lens[0] _run_draft_extend_cuda_graph_case( testcase, @@ -559,6 +568,8 @@ def run_mla_draft_extend_v2_cuda_graph_case( rtol=MLA_RTOL, run_graph_eager=False, compare_replay_to_graph_eager=False, + pad_style=pad_style, + pad_num_tokens_per_bs=num_tokens_per_req, ) diff --git a/test/registered/attention/unittests/dense/test_fa3.py b/test/registered/attention/unittests/dense/test_fa3.py index 79bb7d3dd..bdbd14f7d 100644 --- a/test/registered/attention/unittests/dense/test_fa3.py +++ b/test/registered/attention/unittests/dense/test_fa3.py @@ -435,13 +435,26 @@ class TestFA3DenseAttentionBackendCorrectness(CustomTestCase): "and produces ~82 % wrong attention values. Tracked in KNOWN_FAILURES.md §C.3." ) for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dense_draft_extend_v2_cuda_graph_case( - self, - case, - head_dim=self.HEAD_DIM, - hidden_size=self.HIDDEN_SIZE, - ) + for pad_style in ("small_real", "prod_fill"): + for capture_bs in ( + case.batch_size, + case.batch_size * 2, + case.batch_size * 4, + ): + with self.subTest( + case=case.name, + backend=case.backend, + pad_style=pad_style, + capture_bs=capture_bs, + ): + run_dense_draft_extend_v2_cuda_graph_case( + self, + case, + head_dim=self.HEAD_DIM, + hidden_size=self.HIDDEN_SIZE, + cuda_graph_capture_batch_size=capture_bs, + pad_style=pad_style, + ) def test_runner_mode_eagle_draft_extend_cases(self): for case, spec_kind in self.DRAFT_EXTEND_CASES: diff --git a/test/registered/attention/unittests/dense/test_fa4.py b/test/registered/attention/unittests/dense/test_fa4.py index dabacf530..13e3c3ec3 100644 --- a/test/registered/attention/unittests/dense/test_fa4.py +++ b/test/registered/attention/unittests/dense/test_fa4.py @@ -424,13 +424,26 @@ class TestFA4DenseAttentionBackendCorrectness(CustomTestCase): "and produces ~82 % wrong attention values. Tracked in KNOWN_FAILURES.md §C.3." ) for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dense_draft_extend_v2_cuda_graph_case( - self, - case, - head_dim=self.HEAD_DIM, - hidden_size=self.HIDDEN_SIZE, - ) + for pad_style in ("small_real", "prod_fill"): + for capture_bs in ( + case.batch_size, + case.batch_size * 2, + case.batch_size * 4, + ): + with self.subTest( + case=case.name, + backend=case.backend, + pad_style=pad_style, + capture_bs=capture_bs, + ): + run_dense_draft_extend_v2_cuda_graph_case( + self, + case, + head_dim=self.HEAD_DIM, + hidden_size=self.HIDDEN_SIZE, + cuda_graph_capture_batch_size=capture_bs, + pad_style=pad_style, + ) def test_runner_mode_eagle_draft_extend_cases(self): for case, spec_kind in self.DRAFT_EXTEND_CASES: diff --git a/test/registered/attention/unittests/dense/test_triton.py b/test/registered/attention/unittests/dense/test_triton.py index 1aed78dca..1d438c399 100644 --- a/test/registered/attention/unittests/dense/test_triton.py +++ b/test/registered/attention/unittests/dense/test_triton.py @@ -413,9 +413,27 @@ class TestTritonDenseAttentionBackendCorrectness(CustomTestCase): ) def test_runner_mode_eagle_draft_extend_v2_cuda_graph_cases(self): + # pad_ratio is expressed as the captured batch size relative to the + # case's real batch size: 1.0x = no padding, 2.0x = 50% padded, etc. for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_dense_draft_extend_v2_cuda_graph_case(self, case) + for pad_style in ("small_real", "prod_fill"): + for capture_bs in ( + case.batch_size, + case.batch_size * 2, + case.batch_size * 4, + ): + with self.subTest( + case=case.name, + backend=case.backend, + pad_style=pad_style, + capture_bs=capture_bs, + ): + run_dense_draft_extend_v2_cuda_graph_case( + self, + case, + cuda_graph_capture_batch_size=capture_bs, + pad_style=pad_style, + ) def test_runner_mode_eagle_draft_extend_v2_cuda_graph_runner_cases(self): for case in self.EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES: diff --git a/test/registered/attention/unittests/mla/test_triton.py b/test/registered/attention/unittests/mla/test_triton.py index e4204312e..8c6c4281d 100644 --- a/test/registered/attention/unittests/mla/test_triton.py +++ b/test/registered/attention/unittests/mla/test_triton.py @@ -336,8 +336,24 @@ class TestTritonMLAAttentionBackendCorrectness(CustomTestCase): def test_runner_mode_eagle_draft_extend_v2_cuda_graph_cases(self): for case in self.DRAFT_EXTEND_V2_CUDA_GRAPH_CASES: - with self.subTest(case=case.name, backend=case.backend): - run_mla_draft_extend_v2_cuda_graph_case(self, case) + for pad_style in ("small_real", "prod_fill"): + for capture_bs in ( + case.batch_size, + case.batch_size * 2, + case.batch_size * 4, + ): + with self.subTest( + case=case.name, + backend=case.backend, + pad_style=pad_style, + capture_bs=capture_bs, + ): + run_mla_draft_extend_v2_cuda_graph_case( + self, + case, + cuda_graph_capture_batch_size=capture_bs, + pad_style=pad_style, + ) def test_runner_mode_eagle_draft_extend_v2_cuda_graph_runner_cases(self): for case in self.EAGLE_DRAFT_EXTEND_V2_RUNNER_CASES: