test: strengthen CG-replay coverage with prod-fill padding, metadata invariants, and pad-ratio sweep (#26658)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

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(

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)
+        )

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],

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,
     )