chore: import upstream snapshot with attribution

python/sglang/test/kits/abort_timeout_kit.py

@@ -0,0 +1,144 @@
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import requests
+
+# Safety timeout for all HTTP requests to prevent CI from hanging forever.
+_REQUEST_TIMEOUT = 60
+
+
+class AbortAllMixin:
+    """Test /abort_request with abort_all=True.
+
+    Server needs sufficient --max-running-requests.
+    """
+
+    abort_all_num_requests: int = 32
+    abort_all_max_new_tokens: int = 16000
+    abort_all_sleep: float = 2
+
+    def test_abort_all(self):
+        num_requests = self.abort_all_num_requests
+
+        def run_decode():
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": "The capital of France is",
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": self.abort_all_max_new_tokens,
+                        "ignore_eos": True,
+                    },
+                },
+                timeout=_REQUEST_TIMEOUT,
+            )
+            return response.json()
+
+        with ThreadPoolExecutor(num_requests) as executor:
+            futures = [executor.submit(run_decode) for _ in range(num_requests)]
+
+            time.sleep(self.abort_all_sleep)
+
+            requests.post(
+                self.base_url + "/abort_request",
+                json={"abort_all": True},
+                timeout=10,
+            ).raise_for_status()
+
+            for future in as_completed(futures):
+                result = future.result()
+                self.assertEqual(result["meta_info"]["finish_reason"]["type"], "abort")
+
+            self.assertIsNone(self.process.poll())
+
+
+class WaitingTimeoutMixin:
+    """Test waiting queue timeout.
+
+    Server needs SGLANG_REQ_WAITING_TIMEOUT and --max-running-requests=1.
+    """
+
+    waiting_timeout_num_requests: int = 2
+    waiting_timeout_max_new_tokens: int = 512
+
+    def test_waiting_timeout(self):
+        num_requests = self.waiting_timeout_num_requests
+
+        def run_decode():
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": "Today is ",
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": self.waiting_timeout_max_new_tokens,
+                        "ignore_eos": True,
+                    },
+                },
+                timeout=_REQUEST_TIMEOUT,
+            )
+            return response.json()
+
+        with ThreadPoolExecutor(num_requests) as executor:
+            futures = [executor.submit(run_decode) for _ in range(num_requests)]
+
+            error_count = 0
+            for future in as_completed(futures):
+                result = future.result()
+                if result.get("object") == "error":
+                    error_count += 1
+                    self.assertEqual(result["code"], 503)
+
+            self.assertEqual(error_count, 1)
+            self.assertIsNone(self.process.poll())
+
+
+class RunningTimeoutTwoWaveMixin:
+    """Test running timeout with a two-wave pattern.
+
+    Sends two waves with different forward_entry_time so that timeouts are
+    triggered in separate batches. Regression test for
+    https://github.com/sgl-project/sglang/pull/18760
+
+    Server needs SGLANG_REQ_RUNNING_TIMEOUT and sufficient --max-running-requests
+    to hold both waves.
+    """
+
+    running_timeout_num_wave1: int = 8
+    running_timeout_num_wave2: int = 8
+    running_timeout_sleep: float = 3
+    running_timeout_max_new_tokens: int = 1024
+
+    def test_running_timeout_no_crash(self):
+        num_wave1 = self.running_timeout_num_wave1
+        num_wave2 = self.running_timeout_num_wave2
+
+        def run_decode():
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": "Write a long story about a magical kingdom.",
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": self.running_timeout_max_new_tokens,
+                        "ignore_eos": True,
+                    },
+                },
+                timeout=_REQUEST_TIMEOUT,
+            )
+            return response.json()
+
+        with ThreadPoolExecutor(num_wave1 + num_wave2) as executor:
+            futures1 = [executor.submit(run_decode) for _ in range(num_wave1)]
+
+            time.sleep(self.running_timeout_sleep)
+
+            futures2 = [executor.submit(run_decode) for _ in range(num_wave2)]
+
+            for future in as_completed(futures1 + futures2):
+                result = future.result()
+                if result.get("object") == "error":
+                    self.assertEqual(result["code"], 503)
+
+        self.assertIsNone(self.process.poll())

python/sglang/test/kits/attention_unittest/__init__.py

@@ -0,0 +1 @@
+"""Shared fixtures for manual attention backend unit tests."""

python/sglang/test/kits/attention_unittest/attention_methods/__init__.py

@@ -0,0 +1 @@
+"""Attention-method fixtures for attention backend unit tests."""

python/sglang/test/kits/attention_unittest/runner_modes/__init__.py

@@ -0,0 +1 @@
+"""Runner orchestration helpers for attention backend unit tests."""

python/sglang/test/kits/attention_unittest/runner_modes/cuda_graph_decode_runner.py

@@ -0,0 +1,947 @@
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import torch
+
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+
+from ..attention_methods.dense_attention import DEFAULT_DEVICE as DENSE_DEFAULT_DEVICE
+from ..attention_methods.dense_attention import DEFAULT_DTYPE as DENSE_DEFAULT_DTYPE
+from ..attention_methods.dense_attention import (
+    DEFAULT_HEAD_DIM,
+    DEFAULT_HIDDEN_SIZE,
+)
+from ..attention_methods.dense_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as DENSE_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.dense_attention import (
+    DENSE_ATOL,
+    DENSE_RTOL,
+    DenseAttentionCase,
+)
+from ..attention_methods.dense_attention import (
+    _make_forward_batch as _make_dense_forward_batch,
+)
+from ..attention_methods.dense_attention import (
+    build_dense_attention_fixture,
+    dense_fixture_inputs,
+    expected_dense_output_from_inputs,
+    make_dense_case_with_prefix_lens,
+    make_dense_padded_replay_inputs,
+    make_dense_random_inputs,
+    prepare_dense_runner_inputs,
+    run_dense_fixture_eager,
+    run_dense_forward,
+)
+from ..attention_methods.dsa_attention import (
+    DSA_PAGE_SIZE,
+    DSA_SPARSE_ATOL,
+    DSA_SPARSE_RTOL,
+    DSAAttentionCase,
+    _clone_dsa_sparse_cache,
+)
+from ..attention_methods.dsa_attention import (
+    _make_forward_batch as _make_dsa_forward_batch,
+)
+from ..attention_methods.dsa_attention import (
+    _restore_dsa_sparse_cache,
+    build_dsa_sparse_attention_fixture,
+    dsa_sparse_fixture_inputs,
+    expected_dsa_sparse_output_from_inputs,
+    make_dsa_sparse_case_with_prefix_lens,
+    make_dsa_sparse_random_inputs,
+    make_dsa_sparse_replay_inputs,
+    prepare_dsa_sparse_runner_inputs,
+    run_dsa_sparse_forward,
+)
+from ..attention_methods.dsv4_attention import (
+    DSV4_ATOL,
+    DSV4_RTOL,
+    DSV4AttentionCase,
+)
+from ..attention_methods.dsv4_attention import (
+    _make_forward_batch as _make_dsv4_forward_batch,
+)
+from ..attention_methods.dsv4_attention import (
+    build_dsv4_attention_fixture,
+    dsv4_fixture_inputs,
+    expected_dsv4_output_from_inputs,
+    make_dsv4_case_with_prefix_lens,
+    make_dsv4_padded_replay_inputs,
+    make_dsv4_random_inputs,
+    prepare_dsv4_runner_inputs,
+    run_dsv4_fixture_eager,
+    run_dsv4_forward,
+)
+from ..attention_methods.dual_chunk_attention import (
+    DualChunkAttentionCase,
+    _clone_dual_chunk_cache,
+    _restore_dual_chunk_cache,
+    build_dual_chunk_attention_fixture,
+    dual_chunk_fixture_inputs,
+    expected_dual_chunk_output_from_inputs,
+    make_dual_chunk_case_with_prefix_lens,
+    make_dual_chunk_random_inputs,
+    make_dual_chunk_replay_inputs,
+    prepare_dual_chunk_runner_inputs,
+    run_dual_chunk_fixture_eager,
+    run_dual_chunk_forward,
+)
+from ..attention_methods.gdn_attention import DEFAULT_DEVICE as GDN_DEFAULT_DEVICE
+from ..attention_methods.gdn_attention import DEFAULT_DTYPE as GDN_DEFAULT_DTYPE
+from ..attention_methods.gdn_attention import (
+    DEFAULT_HEAD_K_DIM,
+    DEFAULT_HEAD_V_DIM,
+)
+from ..attention_methods.gdn_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as GDN_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.gdn_attention import (
+    GDN_ATOL,
+    GDN_RTOL,
+    GDNAttentionCase,
+    _clone_gdn_cache,
+)
+from ..attention_methods.gdn_attention import (
+    _make_forward_batch as _make_gdn_forward_batch,
+)
+from ..attention_methods.gdn_attention import (
+    _restore_gdn_cache,
+    build_gdn_attention_fixture,
+    expected_gdn_output_from_inputs,
+    gdn_fixture_inputs,
+    make_gdn_case_with_prefix_lens,
+    make_gdn_random_inputs,
+    make_gdn_replay_inputs,
+    prepare_gdn_runner_inputs,
+    run_gdn_fixture_eager,
+    run_gdn_forward,
+)
+from ..attention_methods.kda_attention import DEFAULT_DEVICE as KDA_DEFAULT_DEVICE
+from ..attention_methods.kda_attention import DEFAULT_DTYPE as KDA_DEFAULT_DTYPE
+from ..attention_methods.kda_attention import (
+    DEFAULT_HEAD_K_DIM as KDA_DEFAULT_HEAD_K_DIM,
+)
+from ..attention_methods.kda_attention import (
+    DEFAULT_HEAD_V_DIM as KDA_DEFAULT_HEAD_V_DIM,
+)
+from ..attention_methods.kda_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as KDA_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.kda_attention import (
+    KDA_GRAPH_ATOL,
+    KDA_GRAPH_RTOL,
+    KDAAttentionCase,
+    _clone_kda_cache,
+)
+from ..attention_methods.kda_attention import (
+    _make_forward_batch as _make_kda_forward_batch,
+)
+from ..attention_methods.kda_attention import (
+    _restore_kda_cache,
+    build_kda_attention_fixture,
+    expected_kda_output_from_inputs,
+    kda_fixture_inputs,
+    make_kda_case_with_prefix_lens,
+    make_kda_random_inputs,
+    make_kda_replay_inputs,
+    prepare_kda_runner_inputs,
+    run_kda_fixture_eager,
+    run_kda_forward,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_DEVICE as LIGHTNING_DEFAULT_DEVICE,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_DTYPE as LIGHTNING_DEFAULT_DTYPE,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_HEAD_DIM as LIGHTNING_DEFAULT_HEAD_DIM,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as LIGHTNING_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.lightning_attention import (
+    LIGHTNING_GRAPH_ATOL,
+    LIGHTNING_GRAPH_RTOL,
+    LightningAttentionCase,
+    _clone_lightning_cache,
+)
+from ..attention_methods.lightning_attention import (
+    _make_forward_batch as _make_lightning_forward_batch,
+)
+from ..attention_methods.lightning_attention import (
+    _restore_lightning_cache,
+    build_lightning_attention_fixture,
+    expected_lightning_output_from_inputs,
+    lightning_fixture_inputs,
+    make_lightning_case_with_prefix_lens,
+    make_lightning_random_inputs,
+    make_lightning_replay_inputs,
+    prepare_lightning_runner_inputs,
+    run_lightning_fixture_eager,
+    run_lightning_forward,
+)
+from ..attention_methods.mamba2_attention import DEFAULT_DEVICE as MAMBA2_DEFAULT_DEVICE
+from ..attention_methods.mamba2_attention import DEFAULT_DTYPE as MAMBA2_DEFAULT_DTYPE
+from ..attention_methods.mamba2_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as MAMBA2_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.mamba2_attention import (
+    MAMBA2_GRAPH_ATOL,
+    MAMBA2_GRAPH_RTOL,
+    Mamba2AttentionCase,
+    _clone_mamba2_cache,
+)
+from ..attention_methods.mamba2_attention import (
+    _make_forward_batch as _make_mamba2_forward_batch,
+)
+from ..attention_methods.mamba2_attention import (
+    _restore_mamba2_cache,
+    build_mamba2_attention_fixture,
+    expected_mamba2_output_from_inputs,
+    make_mamba2_case_with_prefix_lens,
+    make_mamba2_random_inputs,
+    make_mamba2_replay_inputs,
+    mamba2_fixture_inputs,
+    prepare_mamba2_runner_inputs,
+    run_mamba2_fixture_eager,
+    run_mamba2_forward,
+)
+from ..attention_methods.mla_attention import DEFAULT_DEVICE as MLA_DEFAULT_DEVICE
+from ..attention_methods.mla_attention import DEFAULT_DTYPE as MLA_DEFAULT_DTYPE
+from ..attention_methods.mla_attention import (
+    DEFAULT_HIDDEN_SIZE as MLA_DEFAULT_HIDDEN_SIZE,
+)
+from ..attention_methods.mla_attention import (
+    DEFAULT_KV_LORA_RANK,
+)
+from ..attention_methods.mla_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as MLA_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.mla_attention import (
+    DEFAULT_QK_ROPE_HEAD_DIM,
+    MLA_ATOL,
+    MLA_RTOL,
+    MLAAttentionCase,
+)
+from ..attention_methods.mla_attention import (
+    _make_forward_batch as _make_mla_forward_batch,
+)
+from ..attention_methods.mla_attention import (
+    build_mla_attention_fixture,
+    expected_mla_output_from_inputs,
+    make_mla_case_with_prefix_lens,
+    make_mla_padded_replay_inputs,
+    make_mla_random_inputs,
+    mla_fixture_inputs,
+    prepare_mla_runner_inputs,
+    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
+GDN_CUDA_GRAPH_CAPTURE_BATCH_SIZE = 3
+DSV4_CUDA_GRAPH_CAPTURE_BATCH_SIZE = 2
+KDA_CUDA_GRAPH_CAPTURE_BATCH_SIZE = 3
+LIGHTNING_CUDA_GRAPH_CAPTURE_BATCH_SIZE = 3
+MAMBA2_CUDA_GRAPH_CAPTURE_BATCH_SIZE = 3
+
+
+@dataclass(frozen=True)
+class CudaGraphDecodeAdapter:
+    build_fixture: Callable[..., Any]
+    make_case: Callable[[Any, str, tuple[int, ...]], Any]
+    make_forward_batch: Callable[..., Any]
+    fixture_inputs: Callable[[Any], dict[str, Any]]
+    make_capture_inputs: Callable[..., dict[str, Any]]
+    make_replay_inputs: Callable[..., dict[str, Any]]
+    prepare_inputs: Callable[..., None]
+    run_eager: Callable[[Any], torch.Tensor]
+    run_forward: Callable[[Any, Any, dict[str, Any]], torch.Tensor]
+    expected_output: Callable[[Any, Any, dict[str, Any], Any], torch.Tensor]
+    clone_state: Callable[[Any], Any] = lambda _: None
+    restore_state: Callable[[Any, Any], None] = lambda _fixture, _state: None
+    allow_padding: bool = True
+    atol: float = 0.0
+    rtol: float = 0.0
+
+
+def _check_decode_cuda_graph_case(case, capture_batch_size: int, *, allow_padding=True):
+    if not case.forward_mode.is_decode():
+        raise ValueError(
+            "CUDA graph runner integration currently expects decode cases."
+        )
+    if allow_padding:
+        if case.batch_size > capture_batch_size:
+            raise ValueError(
+                "CUDA graph capture batch size must be at least the replay batch size."
+            )
+    elif case.batch_size != capture_batch_size:
+        raise ValueError(
+            "This CUDA graph coverage uses an unpadded replay batch; choose a case "
+            "whose batch size matches the capture batch size."
+        )
+
+
+def _init_cuda_graph_capture_metadata(backend, capture_batch_size: int, batch):
+    backend.init_cuda_graph_state(
+        max_bs=capture_batch_size,
+        max_num_tokens=batch.input_ids.numel(),
+    )
+    backend.init_forward_metadata_out_graph(batch, in_capture=True)
+    backend.init_forward_metadata_in_graph(batch)
+
+
+def _init_cuda_graph_replay_metadata(backend, capture_batch_size: int, batch):
+    from types import SimpleNamespace
+
+    fb_view = SimpleNamespace(
+        batch_size=capture_batch_size,
+        forward_mode=batch.forward_mode,
+        actual_forward_mode=batch.forward_mode,
+        input_ids=batch.input_ids,
+        positions=getattr(batch, "positions", None),
+        req_pool_indices=batch.req_pool_indices,
+        seq_lens=batch.seq_lens,
+        seq_lens_sum=batch.seq_lens_sum,
+        seq_lens_cpu=batch.seq_lens_cpu,
+        encoder_lens=batch.encoder_lens,
+        out_cache_loc=getattr(batch, "out_cache_loc", None),
+        spec_info=batch.spec_info,
+    )
+    backend.init_forward_metadata_out_graph(fb_view)
+    # No real cuda graph here, so run the in-graph step explicitly to produce
+    # the Full metadata the forward path expects (no-op for non-DSV4).
+    backend.init_forward_metadata_in_graph(fb_view)
+    # 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(
+    testcase,
+    case,
+    *,
+    adapter: CudaGraphDecodeAdapter,
+    build_kwargs: dict,
+    capture_batch_size: int,
+    max_context_len: int,
+    dtype: torch.dtype,
+    device: str,
+):
+    _check_decode_cuda_graph_case(
+        case,
+        capture_batch_size,
+        allow_padding=adapter.allow_padding,
+    )
+    # NOTE: `capture_prefix_len`-vs-replay assertion happens below once the
+    # graph fixture is built (we need `backend.get_cuda_graph_seq_len_fill_value`).
+
+    eager_fixture = adapter.build_fixture(testcase, case, **build_kwargs)
+    eager_inputs = adapter.fixture_inputs(eager_fixture)
+    eager_initial_state = adapter.clone_state(eager_fixture)
+    eager_actual = adapter.run_eager(eager_fixture)
+    eager_expected = adapter.expected_output(
+        eager_fixture,
+        case,
+        eager_inputs,
+        eager_initial_state,
+    )
+    torch.testing.assert_close(
+        eager_actual,
+        eager_expected,
+        atol=adapter.atol,
+        rtol=adapter.rtol,
+    )
+
+    graph_fixture = adapter.build_fixture(
+        testcase,
+        case,
+        **build_kwargs,
+        disable_cuda_graph=False,
+        runner_batch_size=capture_batch_size,
+    )
+    backend = graph_fixture.backend
+    graph_replay_inputs = adapter.fixture_inputs(graph_fixture)
+    graph_initial_state = adapter.clone_state(graph_fixture)
+    capture_prefix_len = max(0, backend.get_cuda_graph_seq_len_fill_value() - 1)
+    if any(p < capture_prefix_len for p in case.prefix_lens):
+        raise AssertionError(
+            f"replay prefix_lens must each be >= capture_prefix_len="
+            f"{capture_prefix_len} so capture-time random KV does not leak "
+            f"into replay; got prefix_lens={case.prefix_lens}"
+        )
+
+    capture_case = adapter.make_case(
+        case,
+        f"{case.name}_cuda_graph_capture",
+        (capture_prefix_len,) * capture_batch_size,
+    )
+    capture_inputs = adapter.make_capture_inputs(
+        capture_case,
+        graph_fixture,
+        dtype=dtype,
+        device=device,
+    )
+    capture_batch = adapter.make_forward_batch(
+        capture_case,
+        graph_fixture.runner,
+        max_context_len=max_context_len,
+        device=device,
+    )
+    adapter.prepare_inputs(
+        graph_fixture,
+        capture_case,
+        capture_batch,
+        capture_inputs,
+        max_context_len=max_context_len,
+    )
+
+    with torch.no_grad(), forward_context(ForwardContext(attn_backend=backend)):
+        _init_cuda_graph_capture_metadata(backend, capture_batch_size, capture_batch)
+        # Capture forward is a JIT warmup that mirrors production: the
+        # captured CUDA graph records kernel launches against buffers
+        # that *will* be populated by `init_forward_metadata_replay_cuda_graph`
+        # at replay. The capture-time output itself is discarded in
+        # production — and we discard it here too. Backends like FA3/FA4
+        # legitimately assign-but-don't-populate metadata buffers at
+        # capture, which makes the capture-time output undefined; only
+        # the replay output is contractually required to match the
+        # reference.
+        adapter.run_forward(graph_fixture, capture_batch, capture_inputs)
+        backend.on_after_cuda_graph_warmup()
+
+        adapter.restore_state(graph_fixture, graph_initial_state)
+        replay_pad_prefix_lens = (capture_prefix_len,) * (
+            capture_batch_size - case.batch_size
+        )
+        replay_case = adapter.make_case(
+            case,
+            f"{case.name}_cuda_graph_replay",
+            case.prefix_lens + replay_pad_prefix_lens,
+        )
+        replay_inputs = adapter.make_replay_inputs(
+            replay_case,
+            graph_fixture,
+            replay_pad_prefix_lens,
+            graph_replay_inputs,
+            dtype=dtype,
+            device=device,
+        )
+        replay_batch = adapter.make_forward_batch(
+            replay_case,
+            graph_fixture.runner,
+            max_context_len=max_context_len,
+            device=device,
+        )
+        adapter.prepare_inputs(
+            graph_fixture,
+            replay_case,
+            replay_batch,
+            replay_inputs,
+            max_context_len=max_context_len,
+        )
+        _init_cuda_graph_replay_metadata(backend, capture_batch_size, replay_batch)
+        replay_actual = adapter.run_forward(
+            graph_fixture,
+            replay_batch,
+            replay_inputs,
+        )
+
+    replay_expected = adapter.expected_output(
+        graph_fixture,
+        replay_case,
+        replay_inputs,
+        graph_initial_state,
+    )
+    torch.testing.assert_close(
+        replay_actual,
+        replay_expected,
+        atol=adapter.atol,
+        rtol=adapter.rtol,
+    )
+    torch.testing.assert_close(
+        replay_actual[: case.num_input_tokens],
+        eager_actual,
+        atol=adapter.atol,
+        rtol=adapter.rtol,
+    )
+
+
+def run_dense_cuda_graph_decode_case(
+    testcase,
+    case: DenseAttentionCase,
+    *,
+    head_dim: int = DEFAULT_HEAD_DIM,
+    hidden_size: int = DEFAULT_HIDDEN_SIZE,
+    max_context_len: int = DENSE_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = DENSE_DEFAULT_DTYPE,
+    device: str = DENSE_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int = DENSE_CUDA_GRAPH_CAPTURE_BATCH_SIZE,
+):
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_dense_attention_fixture,
+        make_case=make_dense_case_with_prefix_lens,
+        make_forward_batch=_make_dense_forward_batch,
+        fixture_inputs=dense_fixture_inputs,
+        make_capture_inputs=make_dense_random_inputs,
+        make_replay_inputs=make_dense_padded_replay_inputs,
+        prepare_inputs=prepare_dense_runner_inputs,
+        run_eager=run_dense_fixture_eager,
+        run_forward=run_dense_forward,
+        expected_output=expected_dense_output_from_inputs,
+        atol=DENSE_ATOL,
+        rtol=DENSE_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_dim=head_dim,
+            hidden_size=hidden_size,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=cuda_graph_capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def run_mla_cuda_graph_decode_case(
+    testcase,
+    case: MLAAttentionCase,
+    *,
+    kv_lora_rank: int = DEFAULT_KV_LORA_RANK,
+    qk_rope_head_dim: int = DEFAULT_QK_ROPE_HEAD_DIM,
+    hidden_size: int = MLA_DEFAULT_HIDDEN_SIZE,
+    max_context_len: int = MLA_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = MLA_DEFAULT_DTYPE,
+    device: str = MLA_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int = MLA_CUDA_GRAPH_CAPTURE_BATCH_SIZE,
+):
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_mla_attention_fixture,
+        make_case=make_mla_case_with_prefix_lens,
+        make_forward_batch=_make_mla_forward_batch,
+        fixture_inputs=mla_fixture_inputs,
+        make_capture_inputs=make_mla_random_inputs,
+        make_replay_inputs=make_mla_padded_replay_inputs,
+        prepare_inputs=prepare_mla_runner_inputs,
+        run_eager=run_mla_fixture_eager,
+        run_forward=run_mla_forward,
+        expected_output=expected_mla_output_from_inputs,
+        atol=MLA_ATOL,
+        rtol=MLA_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            kv_lora_rank=kv_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            hidden_size=hidden_size,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=cuda_graph_capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def run_dsv4_cuda_graph_decode_case(
+    testcase,
+    case: DSV4AttentionCase,
+    *,
+    swa_size: int = 1024,
+    max_context_len: int = 256,
+    dtype: torch.dtype = torch.bfloat16,
+    device: str = "cuda",
+    cuda_graph_capture_batch_size: int = DSV4_CUDA_GRAPH_CAPTURE_BATCH_SIZE,
+):
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_dsv4_attention_fixture,
+        make_case=make_dsv4_case_with_prefix_lens,
+        make_forward_batch=_make_dsv4_forward_batch,
+        fixture_inputs=dsv4_fixture_inputs,
+        make_capture_inputs=make_dsv4_random_inputs,
+        make_replay_inputs=make_dsv4_padded_replay_inputs,
+        prepare_inputs=prepare_dsv4_runner_inputs,
+        run_eager=run_dsv4_fixture_eager,
+        run_forward=run_dsv4_forward,
+        expected_output=expected_dsv4_output_from_inputs,
+        atol=DSV4_ATOL,
+        rtol=DSV4_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            swa_size=swa_size,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=cuda_graph_capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def run_gdn_cuda_graph_decode_case(
+    testcase,
+    case: GDNAttentionCase,
+    *,
+    head_k_dim: int = DEFAULT_HEAD_K_DIM,
+    head_v_dim: int = DEFAULT_HEAD_V_DIM,
+    max_context_len: int = GDN_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = GDN_DEFAULT_DTYPE,
+    device: str = GDN_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int = GDN_CUDA_GRAPH_CAPTURE_BATCH_SIZE,
+):
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_gdn_attention_fixture,
+        make_case=make_gdn_case_with_prefix_lens,
+        make_forward_batch=_make_gdn_forward_batch,
+        fixture_inputs=gdn_fixture_inputs,
+        make_capture_inputs=make_gdn_random_inputs,
+        make_replay_inputs=make_gdn_replay_inputs,
+        prepare_inputs=prepare_gdn_runner_inputs,
+        run_eager=run_gdn_fixture_eager,
+        run_forward=run_gdn_forward,
+        expected_output=expected_gdn_output_from_inputs,
+        clone_state=_clone_gdn_cache,
+        restore_state=_restore_gdn_cache,
+        allow_padding=False,
+        atol=GDN_ATOL,
+        rtol=GDN_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_k_dim=head_k_dim,
+            head_v_dim=head_v_dim,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=cuda_graph_capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def run_kda_cuda_graph_decode_case(
+    testcase,
+    case: KDAAttentionCase,
+    *,
+    head_k_dim: int = KDA_DEFAULT_HEAD_K_DIM,
+    head_v_dim: int = KDA_DEFAULT_HEAD_V_DIM,
+    max_context_len: int = KDA_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = KDA_DEFAULT_DTYPE,
+    device: str = KDA_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int = KDA_CUDA_GRAPH_CAPTURE_BATCH_SIZE,
+):
+    """KDA CUDA-graph decode replay. Mirrors `run_gdn_cuda_graph_decode_case`:
+    KDA inherits the same `MambaAttnBackendBase` capture/replay path through
+    `HybridLinearAttnBackend`, so the adapter wiring is identical to GDN.
+    Only DECODE / TARGET_VERIFY are reachable here (the underlying
+    `_replay_metadata` rejects other modes — see kda/README.md).
+    """
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_kda_attention_fixture,
+        make_case=make_kda_case_with_prefix_lens,
+        make_forward_batch=_make_kda_forward_batch,
+        fixture_inputs=kda_fixture_inputs,
+        make_capture_inputs=make_kda_random_inputs,
+        make_replay_inputs=make_kda_replay_inputs,
+        prepare_inputs=prepare_kda_runner_inputs,
+        run_eager=run_kda_fixture_eager,
+        run_forward=run_kda_forward,
+        expected_output=expected_kda_output_from_inputs,
+        clone_state=_clone_kda_cache,
+        restore_state=_restore_kda_cache,
+        allow_padding=False,
+        atol=KDA_GRAPH_ATOL,
+        rtol=KDA_GRAPH_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_k_dim=head_k_dim,
+            head_v_dim=head_v_dim,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=cuda_graph_capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def run_lightning_cuda_graph_decode_case(
+    testcase,
+    case: LightningAttentionCase,
+    *,
+    head_dim: int = LIGHTNING_DEFAULT_HEAD_DIM,
+    max_context_len: int = LIGHTNING_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = LIGHTNING_DEFAULT_DTYPE,
+    device: str = LIGHTNING_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int = LIGHTNING_CUDA_GRAPH_CAPTURE_BATCH_SIZE,
+):
+    """Lightning (Bailing seg_la) CUDA-graph decode replay. Mirrors GDN/KDA;
+    Lightning uses `LightningAttentionBackend` (installed directly via
+    ForwardContext rather than through `HybridLinearAttnBackend`), but the
+    capture/replay contract is the same shape because the backend also
+    inherits from `MambaAttnBackendBase`. Loose tolerance to absorb seg_la
+    Triton kernel CG-replay drift; eager tolerance preserved for non-graph
+    cases."""
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_lightning_attention_fixture,
+        make_case=make_lightning_case_with_prefix_lens,
+        make_forward_batch=_make_lightning_forward_batch,
+        fixture_inputs=lightning_fixture_inputs,
+        make_capture_inputs=make_lightning_random_inputs,
+        make_replay_inputs=make_lightning_replay_inputs,
+        prepare_inputs=prepare_lightning_runner_inputs,
+        run_eager=run_lightning_fixture_eager,
+        run_forward=run_lightning_forward,
+        expected_output=expected_lightning_output_from_inputs,
+        clone_state=_clone_lightning_cache,
+        restore_state=_restore_lightning_cache,
+        allow_padding=False,
+        atol=LIGHTNING_GRAPH_ATOL,
+        rtol=LIGHTNING_GRAPH_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_dim=head_dim,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=cuda_graph_capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def run_mamba2_cuda_graph_decode_case(
+    testcase,
+    case: Mamba2AttentionCase,
+    *,
+    max_context_len: int = MAMBA2_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = MAMBA2_DEFAULT_DTYPE,
+    device: str = MAMBA2_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int = MAMBA2_CUDA_GRAPH_CAPTURE_BATCH_SIZE,
+):
+    """Mamba2 CUDA-graph decode replay. The fixture's
+    `initialize_mamba_selective_state_update_backend` call makes
+    `MambaMixer2.forward_decode` reachable; this adapter then drives the
+    capture/replay lifecycle the same way as GDN/KDA/Lightning, snapshotting
+    both SSM and conv state between capture and replay so the recurrent
+    backend output is reproducible.
+
+    Loose `MAMBA2_GRAPH_ATOL=1e-1` absorbs CG-replay drift; eager
+    `MAMBA2_ATOL=5e-2` is kept for non-graph cases.
+    """
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_mamba2_attention_fixture,
+        make_case=make_mamba2_case_with_prefix_lens,
+        make_forward_batch=_make_mamba2_forward_batch,
+        fixture_inputs=mamba2_fixture_inputs,
+        make_capture_inputs=make_mamba2_random_inputs,
+        make_replay_inputs=make_mamba2_replay_inputs,
+        prepare_inputs=prepare_mamba2_runner_inputs,
+        run_eager=run_mamba2_fixture_eager,
+        run_forward=run_mamba2_forward,
+        expected_output=expected_mamba2_output_from_inputs,
+        clone_state=_clone_mamba2_cache,
+        restore_state=_restore_mamba2_cache,
+        allow_padding=False,
+        atol=MAMBA2_GRAPH_ATOL,
+        rtol=MAMBA2_GRAPH_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=cuda_graph_capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def _run_dsa_sparse_eager_for_cg(fixture):
+    """Eager wrapper for the DSA sparse CG decode adapter — wraps a
+    `forward_context` around `run_dsa_sparse_forward` so `module.attn`
+    sees the active backend (the existing
+    `run_dsa_sparse_fixture_eager` has its own context but takes an
+    extra `testcase` arg for `skipTest`, which doesn't fit the
+    adapter's `run_eager(fixture)` signature)."""
+    with torch.no_grad(), forward_context(ForwardContext(attn_backend=fixture.backend)):
+        fixture.backend.init_forward_metadata(fixture.forward_batch)
+        return run_dsa_sparse_forward(
+            fixture, fixture.forward_batch, dsa_sparse_fixture_inputs(fixture)
+        )
+
+
+def run_dsa_sparse_cuda_graph_decode_case(
+    testcase,
+    case: DSAAttentionCase,
+    *,
+    hidden_size: int = DEFAULT_HIDDEN_SIZE,
+    max_context_len: int | None = None,
+    dtype: torch.dtype = torch.bfloat16,
+    device: str = DENSE_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int | None = None,
+    dsa_decode_backend: str = "flashmla_kv",
+    fp8_kv_cache: bool = False,
+):
+    """DSA sparse-topk CUDA-graph decode replay (`flashmla_kv` path).
+    Sparse decode uses cached MLA latent KV (written by
+    `_populate_dsa_sparse_prefix_kv` at fixture build), so the
+    capture/replay K-cache boundary is compatible with piecewise CG —
+    unlike the dense-fallback MHA_ONE_SHOT path which passes prefix+
+    extend K inline."""
+    if not case.forward_mode.is_decode():
+        raise ValueError(
+            "run_dsa_sparse_cuda_graph_decode_case expects a DECODE case "
+            "(the sparse `flashmla_kv` path is the natural CG decode target)."
+        )
+    capture_batch_size = cuda_graph_capture_batch_size or case.batch_size
+    if max_context_len is None:
+        max_context_len = max(case.seq_lens) if case.seq_lens else DSA_PAGE_SIZE
+        # Round up to page_size multiple.
+        if max_context_len % case.page_size:
+            max_context_len = (
+                (max_context_len + case.page_size - 1) // case.page_size
+            ) * case.page_size
+    from ..attention_methods.dsa_attention import (
+        DSA_SPARSE_FP8_ATOL,
+        DSA_SPARSE_FP8_RTOL,
+    )
+
+    if fp8_kv_cache:
+        atol, rtol = DSA_SPARSE_FP8_ATOL, DSA_SPARSE_FP8_RTOL
+    else:
+        atol, rtol = DSA_SPARSE_ATOL, DSA_SPARSE_RTOL
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_dsa_sparse_attention_fixture,
+        make_case=make_dsa_sparse_case_with_prefix_lens,
+        make_forward_batch=_make_dsa_forward_batch,
+        fixture_inputs=dsa_sparse_fixture_inputs,
+        make_capture_inputs=make_dsa_sparse_random_inputs,
+        make_replay_inputs=make_dsa_sparse_replay_inputs,
+        prepare_inputs=prepare_dsa_sparse_runner_inputs,
+        run_eager=_run_dsa_sparse_eager_for_cg,
+        run_forward=run_dsa_sparse_forward,
+        expected_output=expected_dsa_sparse_output_from_inputs,
+        clone_state=_clone_dsa_sparse_cache,
+        restore_state=_restore_dsa_sparse_cache,
+        allow_padding=False,
+        atol=atol,
+        rtol=rtol,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            hidden_size=hidden_size,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+            dsa_decode_backend=dsa_decode_backend,
+            fp8_kv_cache=fp8_kv_cache,
+        ),
+        capture_batch_size=capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )
+
+
+def run_dual_chunk_cuda_graph_decode_case(
+    testcase,
+    case: DualChunkAttentionCase,
+    *,
+    head_dim: int = DEFAULT_HEAD_DIM,
+    hidden_size: int = DEFAULT_HIDDEN_SIZE,
+    max_context_len: int = DENSE_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = DENSE_DEFAULT_DTYPE,
+    device: str = DENSE_DEFAULT_DEVICE,
+    cuda_graph_capture_batch_size: int | None = None,
+):
+    """Dual-chunk CUDA-graph decode replay. Decode reads cached K/V (set
+    by `set_kv_buffer` inside `forward_decode`) so the capture/replay
+    contract is the same shape as dense attention. The
+    `_clone_dual_chunk_cache` / `_restore_dual_chunk_cache` hooks snapshot
+    both K and V buffers so the capture forward's writes don't bleed into
+    replay state."""
+    if not case.forward_mode.is_decode():
+        raise ValueError("run_dual_chunk_cuda_graph_decode_case expects a DECODE case.")
+    capture_batch_size = cuda_graph_capture_batch_size or case.batch_size
+    adapter = CudaGraphDecodeAdapter(
+        build_fixture=build_dual_chunk_attention_fixture,
+        make_case=make_dual_chunk_case_with_prefix_lens,
+        make_forward_batch=_make_dense_forward_batch,
+        fixture_inputs=dual_chunk_fixture_inputs,
+        make_capture_inputs=make_dual_chunk_random_inputs,
+        make_replay_inputs=make_dual_chunk_replay_inputs,
+        prepare_inputs=prepare_dual_chunk_runner_inputs,
+        run_eager=run_dual_chunk_fixture_eager,
+        run_forward=run_dual_chunk_forward,
+        expected_output=expected_dual_chunk_output_from_inputs,
+        clone_state=_clone_dual_chunk_cache,
+        restore_state=_restore_dual_chunk_cache,
+        allow_padding=True,
+        atol=DENSE_ATOL,
+        rtol=DENSE_RTOL,
+    )
+    _run_cuda_graph_decode_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_dim=head_dim,
+            hidden_size=hidden_size,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        capture_batch_size=capture_batch_size,
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+    )

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

@@ -0,0 +1,343 @@
+from dataclasses import dataclass
+from typing import Any, Callable, Literal, Optional
+
+import torch
+
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+
+from .cuda_graph_decode_runner import (
+    _init_cuda_graph_capture_metadata,
+    _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:
+    build_fixture: Callable[..., Any]
+    make_capture_case: Callable[[Any, str, int, int], Any]
+    make_replay_case: Callable[[Any, str, tuple[int, ...]], Any]
+    make_forward_batch: Callable[..., Any]
+    fixture_inputs: Callable[[Any], dict[str, Any]]
+    make_capture_inputs: Callable[..., dict[str, Any]]
+    make_replay_inputs: Callable[..., dict[str, Any]]
+    prepare_batch: Callable[[Any, Any], None]
+    prepare_inputs: Callable[..., None]
+    run_forward: Callable[[Any, Any, dict[str, Any]], torch.Tensor]
+    expected_output: Callable[[Any, Any, dict[str, Any], Any], torch.Tensor]
+    max_num_tokens: Callable[[Any, int], int] | None = None
+    clone_state: Callable[[Any], Any] = lambda _: None
+    restore_state: Callable[[Any, Any], None] = lambda _fixture, _state: None
+    allow_padding: bool = True
+    run_graph_eager: bool = True
+    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(
+    case,
+    capture_batch_size: int,
+    *,
+    allow_padding: bool,
+) -> None:
+    if allow_padding:
+        if case.batch_size > capture_batch_size:
+            raise ValueError("CUDA graph capture must cover replay batch size.")
+    elif case.batch_size != capture_batch_size:
+        raise ValueError(
+            "This CUDA graph coverage uses an unpadded replay batch; choose a case "
+            "whose batch size matches the capture batch size."
+        )
+
+
+def run_speculative_cuda_graph_case(
+    testcase,
+    case,
+    *,
+    adapter: SpeculativeCudaGraphAdapter,
+    build_kwargs: dict,
+    capture_batch_size: int,
+    max_context_len: int,
+    dtype: torch.dtype,
+    device: str,
+):
+    _check_speculative_cuda_graph_case(
+        case,
+        capture_batch_size,
+        allow_padding=adapter.allow_padding,
+    )
+
+    graph_fixture = adapter.build_fixture(
+        testcase,
+        case,
+        **build_kwargs,
+        disable_cuda_graph=False,
+        runner_batch_size=capture_batch_size,
+    )
+    backend = graph_fixture.backend
+    graph_inputs = adapter.fixture_inputs(graph_fixture)
+    graph_initial_state = adapter.clone_state(graph_fixture)
+    graph_eager_actual = None
+
+    if adapter.run_graph_eager:
+        if adapter.max_num_tokens is not None:
+            backend.init_cuda_graph_state(
+                max_bs=capture_batch_size,
+                max_num_tokens=adapter.max_num_tokens(case, capture_batch_size),
+            )
+        graph_batch = graph_fixture.forward_batch
+        adapter.prepare_batch(case, graph_batch)
+        # Run prepare_inputs in the eager leg too so backends whose reference
+        # depends on cache state / per-fixture stashes (e.g. DSV4 reads BF16
+        # K from `fixture._swa_bf16_k_per_req`, populated by
+        # `prepare_dsv4_runner_inputs`) work the same way as the
+        # capture/replay legs. Backends whose reference is self-contained
+        # (dense / MLA — they re-project from `inputs`) are unaffected;
+        # `prepare_inputs` just re-writes the SWA cache.
+        adapter.prepare_inputs(
+            graph_fixture,
+            case,
+            graph_batch,
+            graph_inputs,
+            max_context_len=max_context_len,
+        )
+        graph_expected = adapter.expected_output(
+            graph_fixture,
+            case,
+            graph_inputs,
+            graph_initial_state,
+        )
+
+        with torch.no_grad(), forward_context(ForwardContext(attn_backend=backend)):
+            backend.init_forward_metadata(graph_batch)
+            graph_eager_actual = adapter.run_forward(
+                graph_fixture,
+                graph_batch,
+                graph_inputs,
+            )
+
+        torch.testing.assert_close(
+            graph_eager_actual,
+            graph_expected,
+            atol=adapter.atol,
+            rtol=adapter.rtol,
+        )
+
+    capture_prefix_len = backend.get_cuda_graph_seq_len_fill_value()
+    capture_case = adapter.make_capture_case(
+        case,
+        f"{case.name}_cuda_graph_capture",
+        capture_prefix_len,
+        capture_batch_size,
+    )
+    capture_inputs = adapter.make_capture_inputs(
+        capture_case,
+        graph_fixture,
+        dtype=dtype,
+        device=device,
+    )
+    capture_batch = adapter.make_forward_batch(
+        capture_case,
+        graph_fixture.runner,
+        max_context_len=max_context_len,
+        device=device,
+    )
+    adapter.prepare_batch(capture_case, capture_batch)
+    adapter.prepare_inputs(
+        graph_fixture,
+        capture_case,
+        capture_batch,
+        capture_inputs,
+        max_context_len=max_context_len,
+    )
+    with torch.no_grad(), forward_context(ForwardContext(attn_backend=backend)):
+        _init_cuda_graph_capture_metadata(backend, capture_batch_size, capture_batch)
+        # Capture forward is a JIT warmup that mirrors production: the
+        # captured CUDA graph records kernel launches against buffers
+        # that *will* be populated by replay-init at replay. The
+        # capture-time output itself is discarded in production — and
+        # we discard it here too. Only the replay output is
+        # contractually required to match the reference.
+        adapter.run_forward(graph_fixture, capture_batch, capture_inputs)
+        backend.on_after_cuda_graph_warmup()
+
+    adapter.restore_state(graph_fixture, graph_initial_state)
+    replay_pad_prefix_lens = (
+        (capture_prefix_len,) * (capture_batch_size - case.batch_size)
+        if adapter.allow_padding
+        else ()
+    )
+    replay_case = adapter.make_replay_case(
+        case,
+        f"{case.name}_cuda_graph_replay",
+        replay_pad_prefix_lens,
+    )
+    replay_inputs = adapter.make_replay_inputs(
+        replay_case,
+        graph_fixture,
+        replay_pad_prefix_lens,
+        graph_inputs,
+        dtype=dtype,
+        device=device,
+    )
+    replay_batch = adapter.make_forward_batch(
+        replay_case,
+        graph_fixture.runner,
+        max_context_len=max_context_len,
+        device=device,
+    )
+    adapter.prepare_batch(replay_case, replay_batch)
+    adapter.prepare_inputs(
+        graph_fixture,
+        replay_case,
+        replay_batch,
+        replay_inputs,
+        max_context_len=max_context_len,
+    )
+    replay_expected = adapter.expected_output(
+        graph_fixture,
+        replay_case,
+        replay_inputs,
+        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(
+            graph_fixture,
+            replay_batch,
+            replay_inputs,
+        )
+
+    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],
+            graph_eager_actual,
+            atol=adapter.atol,
+            rtol=adapter.rtol,
+        )

python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py

@@ -0,0 +1,600 @@
+from dataclasses import dataclass, replace
+from typing import Any, Callable
+
+import torch
+
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    enable_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import (
+    enable_tc_piecewise_cuda_graph as enable_piecewise_cuda_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import (
+    set_tc_piecewise_forward_context as piecewise_forward_context,
+)
+
+from ..attention_methods.dense_attention import DEFAULT_DEVICE as DENSE_DEFAULT_DEVICE
+from ..attention_methods.dense_attention import DEFAULT_DTYPE as DENSE_DEFAULT_DTYPE
+from ..attention_methods.dense_attention import (
+    DEFAULT_HEAD_DIM,
+    DEFAULT_HIDDEN_SIZE,
+)
+from ..attention_methods.dense_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as DENSE_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.dense_attention import (
+    DENSE_ATOL,
+    DENSE_RTOL,
+    DenseAttentionCase,
+    build_dense_attention_fixture,
+    dense_attention_layers,
+    dense_fixture_inputs,
+    expected_dense_output_from_inputs,
+    make_dense_token_padded_inputs,
+    prepare_dense_runner_inputs,
+    run_dense_fixture_eager,
+    run_dense_forward,
+)
+from ..attention_methods.gdn_attention import DEFAULT_DEVICE as GDN_DEFAULT_DEVICE
+from ..attention_methods.gdn_attention import DEFAULT_DTYPE as GDN_DEFAULT_DTYPE
+from ..attention_methods.gdn_attention import (
+    DEFAULT_HEAD_K_DIM,
+    DEFAULT_HEAD_V_DIM,
+)
+from ..attention_methods.gdn_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as GDN_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.gdn_attention import (
+    GDN_ATOL,
+    GDN_RTOL,
+    GDNAttentionCase,
+    _clone_gdn_cache,
+    _restore_gdn_cache,
+    build_gdn_attention_fixture,
+    expected_gdn_output_from_inputs,
+    gdn_attention_layers,
+    gdn_fixture_inputs,
+    make_gdn_token_padded_inputs,
+    prepare_gdn_runner_inputs,
+    run_gdn_fixture_eager,
+    run_gdn_forward,
+)
+from ..attention_methods.kda_attention import DEFAULT_DEVICE as KDA_DEFAULT_DEVICE
+from ..attention_methods.kda_attention import DEFAULT_DTYPE as KDA_DEFAULT_DTYPE
+from ..attention_methods.kda_attention import (
+    DEFAULT_HEAD_K_DIM as KDA_DEFAULT_HEAD_K_DIM,
+)
+from ..attention_methods.kda_attention import (
+    DEFAULT_HEAD_V_DIM as KDA_DEFAULT_HEAD_V_DIM,
+)
+from ..attention_methods.kda_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as KDA_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.kda_attention import (
+    KDA_ATOL,
+    KDA_RTOL,
+    KDAAttentionCase,
+    _clone_kda_cache,
+    _restore_kda_cache,
+    build_kda_attention_fixture,
+    expected_kda_output_from_inputs,
+    kda_attention_layers,
+    kda_fixture_inputs,
+    make_kda_token_padded_inputs,
+    prepare_kda_runner_inputs,
+    run_kda_fixture_eager,
+    run_kda_forward,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_DEVICE as LIGHTNING_DEFAULT_DEVICE,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_DTYPE as LIGHTNING_DEFAULT_DTYPE,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_HEAD_DIM as LIGHTNING_DEFAULT_HEAD_DIM,
+)
+from ..attention_methods.lightning_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as LIGHTNING_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.lightning_attention import (
+    LIGHTNING_ATOL,
+    LIGHTNING_RTOL,
+    LightningAttentionCase,
+    _clone_lightning_cache,
+    _restore_lightning_cache,
+    build_lightning_attention_fixture,
+    expected_lightning_split_op_output_from_inputs,
+    lightning_attention_layers,
+    lightning_fixture_inputs,
+    make_lightning_token_padded_inputs,
+    prepare_lightning_runner_inputs,
+    run_lightning_fixture_eager,
+    run_lightning_forward,
+)
+from ..attention_methods.mamba2_attention import DEFAULT_DEVICE as MAMBA2_DEFAULT_DEVICE
+from ..attention_methods.mamba2_attention import DEFAULT_DTYPE as MAMBA2_DEFAULT_DTYPE
+from ..attention_methods.mamba2_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as MAMBA2_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.mamba2_attention import (
+    MAMBA2_ATOL,
+    MAMBA2_RTOL,
+    Mamba2AttentionCase,
+    _clone_mamba2_cache,
+    _restore_mamba2_cache,
+    build_mamba2_attention_fixture,
+    expected_mamba2_output_from_inputs,
+    make_mamba2_token_padded_inputs,
+    mamba2_attention_layers,
+    mamba2_fixture_inputs,
+    prepare_mamba2_runner_inputs,
+    run_mamba2_fixture_eager,
+    run_mamba2_forward,
+)
+from ..attention_methods.mla_attention import DEFAULT_DEVICE as MLA_DEFAULT_DEVICE
+from ..attention_methods.mla_attention import DEFAULT_DTYPE as MLA_DEFAULT_DTYPE
+from ..attention_methods.mla_attention import (
+    DEFAULT_HIDDEN_SIZE as MLA_DEFAULT_HIDDEN_SIZE,
+)
+from ..attention_methods.mla_attention import (
+    DEFAULT_KV_LORA_RANK,
+)
+from ..attention_methods.mla_attention import (
+    DEFAULT_MAX_CONTEXT_LEN as MLA_DEFAULT_MAX_CONTEXT_LEN,
+)
+from ..attention_methods.mla_attention import (
+    DEFAULT_QK_ROPE_HEAD_DIM,
+    MLA_ATOL,
+    MLA_RTOL,
+    MLAAttentionCase,
+    build_mla_attention_fixture,
+    expected_mla_output_from_inputs,
+    make_mla_token_padded_inputs,
+    mla_attention_layers,
+    mla_fixture_inputs,
+    prepare_mla_runner_inputs,
+    run_mla_fixture_eager,
+    run_mla_forward,
+)
+
+
+@dataclass(frozen=True)
+class SplitOpAdapter:
+    build_fixture: Callable[..., Any]
+    fixture_inputs: Callable[[Any], dict[str, Any]]
+    make_token_padded_inputs: Callable[..., dict[str, Any]]
+    prepare_inputs: Callable[..., None]
+    run_eager: Callable[[Any], torch.Tensor]
+    run_forward: Callable[[Any, Any, dict[str, Any]], torch.Tensor]
+    expected_output: Callable[[Any, Any, dict[str, Any], Any], torch.Tensor]
+    attention_layers: Callable[[Any], list[Any]]
+    clone_state: Callable[[Any], Any] = lambda _: None
+    restore_state: Callable[[Any, Any], None] = lambda _fixture, _state: None
+    atol: float = 0.0
+    rtol: float = 0.0
+
+
+def _check_extend_split_op_case(case) -> None:
+    if not case.forward_mode.is_extend_without_speculative():
+        raise ValueError("PCG/BCG split-op coverage expects non-spec extend cases.")
+
+
+def _split_op_context(*, breakable: bool):
+    if breakable:
+        return enable_breakable_cuda_graph()
+    return enable_piecewise_cuda_graph()
+
+
+def _make_static_forward_batch(raw_batch, static_num_tokens: int, device: str):
+    raw_num_tokens = raw_batch.input_ids.numel()
+    if static_num_tokens < raw_num_tokens:
+        raise ValueError("static_num_tokens must cover the live input token count.")
+    if static_num_tokens == raw_num_tokens:
+        input_ids = raw_batch.input_ids
+        positions = raw_batch.positions
+        out_cache_loc = raw_batch.out_cache_loc
+    else:
+        pad_tokens = static_num_tokens - raw_num_tokens
+        input_ids = torch.cat(
+            [
+                raw_batch.input_ids,
+                torch.zeros(pad_tokens, dtype=raw_batch.input_ids.dtype, device=device),
+            ],
+            dim=0,
+        )
+        positions = torch.cat(
+            [
+                raw_batch.positions,
+                torch.zeros(pad_tokens, dtype=raw_batch.positions.dtype, device=device),
+            ],
+            dim=0,
+        )
+        out_cache_loc = torch.cat(
+            [
+                raw_batch.out_cache_loc,
+                torch.zeros(
+                    pad_tokens,
+                    dtype=raw_batch.out_cache_loc.dtype,
+                    device=device,
+                ),
+            ],
+            dim=0,
+        )
+
+    raw_batch.num_token_non_padded_cpu = raw_num_tokens
+    return replace(
+        raw_batch,
+        input_ids=input_ids,
+        positions=positions,
+        out_cache_loc=out_cache_loc,
+        padded_static_len=static_num_tokens,
+        num_token_non_padded_cpu=raw_num_tokens,
+    )
+
+
+def _slice_live_tokens(output: torch.Tensor, num_tokens: int) -> torch.Tensor:
+    if output.dim() >= 2 and output.shape[0] == 1:
+        return output[:, :num_tokens]
+    return output[:num_tokens]
+
+
+def _run_split_op_extend_case(
+    testcase,
+    case,
+    *,
+    adapter: SplitOpAdapter,
+    build_kwargs: dict[str, Any],
+    max_context_len: int,
+    dtype: torch.dtype,
+    device: str,
+    breakable: bool,
+    static_num_tokens: int | None,
+):
+    _check_extend_split_op_case(case)
+
+    eager_fixture = adapter.build_fixture(testcase, case, **build_kwargs)
+    eager_inputs = adapter.fixture_inputs(eager_fixture)
+    eager_initial_state = adapter.clone_state(eager_fixture)
+    eager_actual = adapter.run_eager(eager_fixture)
+    eager_expected = adapter.expected_output(
+        eager_fixture,
+        case,
+        eager_inputs,
+        eager_initial_state,
+    )
+    torch.testing.assert_close(
+        eager_actual,
+        eager_expected,
+        atol=adapter.atol,
+        rtol=adapter.rtol,
+    )
+
+    split_fixture = adapter.build_fixture(
+        testcase,
+        case,
+        **build_kwargs,
+        disable_piecewise_cuda_graph=False,
+    )
+    split_inputs = adapter.fixture_inputs(split_fixture)
+    split_initial_state = adapter.clone_state(split_fixture)
+    expected = adapter.expected_output(
+        split_fixture,
+        case,
+        split_inputs,
+        split_initial_state,
+    )
+    raw_batch = split_fixture.forward_batch
+    raw_num_tokens = case.num_input_tokens
+    static_num_tokens = static_num_tokens or raw_num_tokens
+    static_batch = _make_static_forward_batch(raw_batch, static_num_tokens, device)
+    static_inputs = adapter.make_token_padded_inputs(
+        case,
+        split_fixture,
+        static_num_tokens,
+        split_inputs,
+        dtype=dtype,
+        device=device,
+    )
+    adapter.prepare_inputs(
+        split_fixture,
+        case,
+        raw_batch,
+        split_inputs,
+        max_context_len=max_context_len,
+    )
+
+    with (
+        torch.no_grad(),
+        _split_op_context(breakable=breakable),
+        forward_context(ForwardContext(attn_backend=split_fixture.backend)),
+        piecewise_forward_context(
+            static_batch,
+            adapter.attention_layers(split_fixture),
+            None,
+            [],
+            [],
+        ),
+    ):
+        split_fixture.backend.init_forward_metadata(raw_batch)
+        actual = adapter.run_forward(split_fixture, static_batch, static_inputs)
+
+    actual = _slice_live_tokens(actual, raw_num_tokens)
+    torch.testing.assert_close(actual, expected, atol=adapter.atol, rtol=adapter.rtol)
+    torch.testing.assert_close(
+        actual,
+        eager_actual,
+        atol=adapter.atol,
+        rtol=adapter.rtol,
+    )
+    adapter.restore_state(split_fixture, split_initial_state)
+
+
+def run_dense_split_op_extend_case(
+    testcase,
+    case: DenseAttentionCase,
+    *,
+    breakable: bool,
+    static_num_tokens: int | None = None,
+    head_dim: int = DEFAULT_HEAD_DIM,
+    hidden_size: int = DEFAULT_HIDDEN_SIZE,
+    max_context_len: int = DENSE_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = DENSE_DEFAULT_DTYPE,
+    device: str = DENSE_DEFAULT_DEVICE,
+):
+    adapter = SplitOpAdapter(
+        build_fixture=build_dense_attention_fixture,
+        fixture_inputs=dense_fixture_inputs,
+        make_token_padded_inputs=make_dense_token_padded_inputs,
+        prepare_inputs=prepare_dense_runner_inputs,
+        run_eager=run_dense_fixture_eager,
+        run_forward=run_dense_forward,
+        expected_output=expected_dense_output_from_inputs,
+        attention_layers=dense_attention_layers,
+        atol=DENSE_ATOL,
+        rtol=DENSE_RTOL,
+    )
+    _run_split_op_extend_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_dim=head_dim,
+            hidden_size=hidden_size,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+        breakable=breakable,
+        static_num_tokens=static_num_tokens,
+    )
+
+
+def run_mla_split_op_extend_case(
+    testcase,
+    case: MLAAttentionCase,
+    *,
+    breakable: bool,
+    static_num_tokens: int | None = None,
+    kv_lora_rank: int = DEFAULT_KV_LORA_RANK,
+    qk_rope_head_dim: int = DEFAULT_QK_ROPE_HEAD_DIM,
+    hidden_size: int = MLA_DEFAULT_HIDDEN_SIZE,
+    max_context_len: int = MLA_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = MLA_DEFAULT_DTYPE,
+    device: str = MLA_DEFAULT_DEVICE,
+):
+    adapter = SplitOpAdapter(
+        build_fixture=build_mla_attention_fixture,
+        fixture_inputs=mla_fixture_inputs,
+        make_token_padded_inputs=make_mla_token_padded_inputs,
+        prepare_inputs=prepare_mla_runner_inputs,
+        run_eager=run_mla_fixture_eager,
+        run_forward=run_mla_forward,
+        expected_output=expected_mla_output_from_inputs,
+        attention_layers=mla_attention_layers,
+        atol=MLA_ATOL,
+        rtol=MLA_RTOL,
+    )
+    _run_split_op_extend_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            kv_lora_rank=kv_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            hidden_size=hidden_size,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+        breakable=breakable,
+        static_num_tokens=static_num_tokens,
+    )
+
+
+def run_gdn_split_op_extend_case(
+    testcase,
+    case: GDNAttentionCase,
+    *,
+    breakable: bool,
+    static_num_tokens: int | None = None,
+    head_k_dim: int = DEFAULT_HEAD_K_DIM,
+    head_v_dim: int = DEFAULT_HEAD_V_DIM,
+    max_context_len: int = GDN_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = GDN_DEFAULT_DTYPE,
+    device: str = GDN_DEFAULT_DEVICE,
+):
+    adapter = SplitOpAdapter(
+        build_fixture=build_gdn_attention_fixture,
+        fixture_inputs=gdn_fixture_inputs,
+        make_token_padded_inputs=make_gdn_token_padded_inputs,
+        prepare_inputs=prepare_gdn_runner_inputs,
+        run_eager=run_gdn_fixture_eager,
+        run_forward=run_gdn_forward,
+        expected_output=expected_gdn_output_from_inputs,
+        attention_layers=gdn_attention_layers,
+        clone_state=_clone_gdn_cache,
+        restore_state=_restore_gdn_cache,
+        atol=GDN_ATOL,
+        rtol=GDN_RTOL,
+    )
+    _run_split_op_extend_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_k_dim=head_k_dim,
+            head_v_dim=head_v_dim,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+        breakable=breakable,
+        static_num_tokens=static_num_tokens,
+    )
+
+
+def run_kda_split_op_extend_case(
+    testcase,
+    case: KDAAttentionCase,
+    *,
+    breakable: bool,
+    static_num_tokens: int | None = None,
+    head_k_dim: int = KDA_DEFAULT_HEAD_K_DIM,
+    head_v_dim: int = KDA_DEFAULT_HEAD_V_DIM,
+    max_context_len: int = KDA_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = KDA_DEFAULT_DTYPE,
+    device: str = KDA_DEFAULT_DEVICE,
+):
+    """KDA PCG/BCG split-op extend. Verifies the live-token slicing contract
+    with a larger static token buffer, mirroring GDN's split_op coverage."""
+    adapter = SplitOpAdapter(
+        build_fixture=build_kda_attention_fixture,
+        fixture_inputs=kda_fixture_inputs,
+        make_token_padded_inputs=make_kda_token_padded_inputs,
+        prepare_inputs=prepare_kda_runner_inputs,
+        run_eager=run_kda_fixture_eager,
+        run_forward=run_kda_forward,
+        expected_output=expected_kda_output_from_inputs,
+        attention_layers=kda_attention_layers,
+        clone_state=_clone_kda_cache,
+        restore_state=_restore_kda_cache,
+        atol=KDA_ATOL,
+        rtol=KDA_RTOL,
+    )
+    _run_split_op_extend_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_k_dim=head_k_dim,
+            head_v_dim=head_v_dim,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+        breakable=breakable,
+        static_num_tokens=static_num_tokens,
+    )
+
+
+def run_lightning_split_op_extend_case(
+    testcase,
+    case: LightningAttentionCase,
+    *,
+    breakable: bool,
+    static_num_tokens: int | None = None,
+    head_dim: int = LIGHTNING_DEFAULT_HEAD_DIM,
+    max_context_len: int = LIGHTNING_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = LIGHTNING_DEFAULT_DTYPE,
+    device: str = LIGHTNING_DEFAULT_DEVICE,
+):
+    """Lightning PCG/BCG split-op extend. Same pattern as KDA/GDN."""
+    adapter = SplitOpAdapter(
+        build_fixture=build_lightning_attention_fixture,
+        fixture_inputs=lightning_fixture_inputs,
+        make_token_padded_inputs=make_lightning_token_padded_inputs,
+        prepare_inputs=prepare_lightning_runner_inputs,
+        run_eager=run_lightning_fixture_eager,
+        run_forward=run_lightning_forward,
+        expected_output=expected_lightning_split_op_output_from_inputs,
+        attention_layers=lightning_attention_layers,
+        clone_state=_clone_lightning_cache,
+        restore_state=_restore_lightning_cache,
+        atol=LIGHTNING_ATOL,
+        rtol=LIGHTNING_RTOL,
+    )
+    _run_split_op_extend_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            head_dim=head_dim,
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+        breakable=breakable,
+        static_num_tokens=static_num_tokens,
+    )
+
+
+def run_mamba2_split_op_extend_case(
+    testcase,
+    case: Mamba2AttentionCase,
+    *,
+    breakable: bool,
+    static_num_tokens: int | None = None,
+    max_context_len: int = MAMBA2_DEFAULT_MAX_CONTEXT_LEN,
+    dtype: torch.dtype = MAMBA2_DEFAULT_DTYPE,
+    device: str = MAMBA2_DEFAULT_DEVICE,
+):
+    """Mamba2 PCG/BCG split-op extend. Same pattern as KDA. Mamba2's
+    forward writes through an `empty_like(hidden_states)` buffer that
+    short-circuits the RadixAttention dispatch path, so the per-head-vs-flat
+    shape mismatch that blocks Lightning split-op doesn't apply."""
+    adapter = SplitOpAdapter(
+        build_fixture=build_mamba2_attention_fixture,
+        fixture_inputs=mamba2_fixture_inputs,
+        make_token_padded_inputs=make_mamba2_token_padded_inputs,
+        prepare_inputs=prepare_mamba2_runner_inputs,
+        run_eager=run_mamba2_fixture_eager,
+        run_forward=run_mamba2_forward,
+        expected_output=expected_mamba2_output_from_inputs,
+        attention_layers=mamba2_attention_layers,
+        clone_state=_clone_mamba2_cache,
+        restore_state=_restore_mamba2_cache,
+        atol=MAMBA2_ATOL,
+        rtol=MAMBA2_RTOL,
+    )
+    _run_split_op_extend_case(
+        testcase,
+        case,
+        adapter=adapter,
+        build_kwargs=dict(
+            max_context_len=max_context_len,
+            dtype=dtype,
+            device=device,
+        ),
+        max_context_len=max_context_len,
+        dtype=dtype,
+        device=device,
+        breakable=breakable,
+        static_num_tokens=static_num_tokens,
+    )

python/sglang/test/kits/basic_api_contract_kit.py

@@ -0,0 +1,87 @@
+"""Basic HTTP / SSE API contract sanity kit.
+
+Probes that catch the server failing at the protocol layer: endpoints
+missing, 5xx returned, response schema broken, or OpenAI-compatible
+routes drifting from the spec.
+
+Mix into any ``CustomTestCase`` subclass that exposes ``self.base_url``
+and ``self.process``. Override ``served_model_name`` if the OpenAI
+probes should pin a specific model id."""
+
+import requests
+
+_REQUEST_TIMEOUT = 60
+
+
+class BasicAPIContractMixin:
+    """Health endpoints + OpenAI /v1 surface probes."""
+
+    served_model_name: str = "default"
+
+    def test_health(self):
+        # Cheapest possible alive check; FastAPI route alone.
+        resp = requests.get(self.base_url + "/health", timeout=10)
+        self.assertEqual(resp.status_code, 200)
+
+    def test_health_generate(self):
+        # sglang's built-in minimal-forward sanity. 200 only if the
+        # scheduler can complete one prefill+decode end to end.
+        resp = requests.get(self.base_url + "/health_generate", timeout=60)
+        self.assertEqual(resp.status_code, 200)
+
+    def test_get_server_info(self):
+        resp = requests.get(self.base_url + "/get_server_info", timeout=10)
+        self.assertEqual(resp.status_code, 200)
+        info = resp.json()
+        # Must expose at least some scheduler/server-args bundle.
+        self.assertIsInstance(info, dict)
+        self.assertGreater(len(info), 0)
+
+    def test_get_model_info(self):
+        resp = requests.get(self.base_url + "/get_model_info", timeout=10)
+        self.assertEqual(resp.status_code, 200)
+        info = resp.json()
+        self.assertIn("model_path", info)
+        self.assertTrue(info["model_path"])
+
+    def test_openai_chat_completion(self):
+        resp = requests.post(
+            self.base_url + "/v1/chat/completions",
+            json={
+                "model": self.served_model_name,
+                "messages": [
+                    {"role": "user", "content": "Say hi in one word."},
+                ],
+                "temperature": 0.0,
+                "max_tokens": 16,
+            },
+            timeout=_REQUEST_TIMEOUT,
+        )
+        self.assertEqual(resp.status_code, 200, resp.text)
+        body = resp.json()
+        self.assertIn("choices", body)
+        self.assertGreater(len(body["choices"]), 0)
+        content = body["choices"][0]["message"]["content"]
+        self.assertIsInstance(content, str)
+        self.assertGreater(len(content), 0)
+        self.assertIn("usage", body)
+
+    def test_openai_completion(self):
+        resp = requests.post(
+            self.base_url + "/v1/completions",
+            json={
+                "model": self.served_model_name,
+                "prompt": "The capital of France is",
+                "temperature": 0.0,
+                "max_tokens": 16,
+            },
+            timeout=_REQUEST_TIMEOUT,
+        )
+        self.assertEqual(resp.status_code, 200, resp.text)
+        body = resp.json()
+        self.assertIn("choices", body)
+        self.assertGreater(len(body["choices"]), 0)
+        text = body["choices"][0]["text"]
+        self.assertIsInstance(text, str)
+        self.assertGreater(len(text), 0)
+        self.assertIn("usage", body)

python/sglang/test/kits/basic_decode_correctness_kit.py

@@ -0,0 +1,114 @@
+"""Basic decode correctness sanity kit.
+
+Probes that catch the model producing wrong output: weight load
+failure, sampling path bugs, KV / attention corruption, and cuda graph
+edge cases. Single-prompt coverage only -- dataset-driven accuracy gates
+belong to the consuming test class, not this kit.
+
+Mix into any ``CustomTestCase`` subclass that exposes ``self.base_url``
+and ``self.process``. Probes complete in well under a minute after
+warmup."""
+
+import requests
+
+_REQUEST_TIMEOUT = 120
+
+
+class BasicDecodeCorrectnessMixin:
+    """Cheap output-quality probes."""
+
+    sanity_max_new_tokens_short: int = 64
+    sanity_max_new_tokens_long: int = 128
+
+    def _decode_generate(self, prompt: str, max_new_tokens: int, stop=None) -> str:
+        sampling_params = {"temperature": 0.0, "max_new_tokens": max_new_tokens}
+        if stop is not None:
+            sampling_params["stop"] = stop
+        resp = requests.post(
+            self.base_url + "/generate",
+            json={"text": prompt, "sampling_params": sampling_params},
+            timeout=_REQUEST_TIMEOUT,
+        )
+        self.assertEqual(resp.status_code, 200)
+        return resp.json()["text"]
+
+    def test_capital_france(self):
+        out = self._decode_generate(
+            "Q: What is the capital of France?\nA:",
+            self.sanity_max_new_tokens_short,
+        )
+        self.assertIn("paris", out.lower())
+
+    def test_basic_math(self):
+        out = self._decode_generate(
+            "Q: What is 17 multiplied by 23? Reply with just the number.\nA:",
+            self.sanity_max_new_tokens_short,
+        )
+        self.assertIn("391", out)
+
+    def test_color_completion(self):
+        out = self._decode_generate(
+            "Q: The three primary colors are red, blue, and ___. "
+            "Fill in the blank.\nA:",
+            self.sanity_max_new_tokens_short,
+        )
+        self.assertIn("yellow", out.lower())
+
+    def test_ascii_ratio(self):
+        # Language-agnostic gibberish detector. Healthy English output is
+        # >90% printable ASCII; multilingual token salad / Unicode noise
+        # from broken weight load drops well below 50%.
+        out = self._decode_generate(
+            "Write a single sentence about a sunny day in the park.",
+            self.sanity_max_new_tokens_long,
+        )
+        printable = sum(1 for c in out if 32 <= ord(c) < 127 or c in "\n\t")
+        ratio = printable / max(len(out), 1)
+        self.assertGreater(
+            ratio,
+            0.85,
+            f"output looks like gibberish (printable ASCII ratio={ratio:.2f}): {out!r}",
+        )
+
+    def test_no_repetition_blowup(self):
+        # KV-cache / attn corruption often manifests as the model getting
+        # stuck looping the same n-gram.
+        out = self._decode_generate(
+            "Briefly explain what gravity is.",
+            self.sanity_max_new_tokens_long,
+        )
+        if len(out) >= 50:
+            windows = [out[i : i + 5] for i in range(len(out) - 5)]
+            most_common_count = max((windows.count(w) for w in set(windows)), default=0)
+            ratio = most_common_count / len(windows)
+            self.assertLess(
+                ratio,
+                0.25,
+                f"output appears to repeat heavily (top 5-gram ratio={ratio:.2f}): {out!r}",
+            )
+
+    def test_determinism_temp_zero(self):
+        # temp=0 must be byte-identical across runs. Stop on "\n" so we
+        # only compare the answer word; long continuations drift on
+        # near-tie tokens (EP MoE / EAGLE spec) and aren't the point.
+        prompt = "Q: What is the capital of France? Reply in one word.\nA:"
+        out1 = self._decode_generate(
+            prompt, self.sanity_max_new_tokens_short, stop=["\n"]
+        )
+        out2 = self._decode_generate(
+            prompt, self.sanity_max_new_tokens_short, stop=["\n"]
+        )
+        self.assertEqual(
+            out1.strip(),
+            out2.strip(),
+            f"temp=0 outputs diverged:\n  out1={out1!r}\n  out2={out2!r}",
+        )
+
+    def test_max_token_one(self):
+        # Degenerate spec step. cuda-graph capture path bugs that only
+        # fire on minimal-output requests.
+        out = self._decode_generate(
+            "Q: What is the capital of France? Just one word.\nA:",
+            max_new_tokens=1,
+        )
+        self.assertGreater(len(out), 0)

python/sglang/test/kits/basic_scheduler_stress_kit.py

@@ -0,0 +1,135 @@
+"""Basic scheduler / cache / streaming stress sanity kit.
+
+Probes that catch bugs which only fire under multi-request or large-
+prompt conditions: scheduler hangs, radix prefix-cache cross-
+contamination, chunked-prefill multi-chunk kernel crashes, and SSE
+streaming corruption.
+
+Mix into any ``CustomTestCase`` subclass that exposes ``self.base_url``
+and ``self.process``."""
+
+import json
+import threading
+
+import requests
+
+_REQUEST_TIMEOUT = 120
+
+# Shared prefix forces all concurrent requests through the same radix
+# match path; per-request suffix branches the tail so the model still
+# has to predict different tokens (otherwise outputs would be identical
+# and we'd be testing 1 request 8 times instead of 8 independent reqs).
+_CONCURRENT_PREFIX = "You are a helpful assistant. Answer with a single word.\n"
+_CONCURRENT_QA = [
+    ("Q: What is the capital of France?\nA:", "paris"),
+    ("Q: What is the capital of Germany?\nA:", "berlin"),
+    ("Q: What is the capital of Italy?\nA:", "rome"),
+    ("Q: What is the capital of Japan?\nA:", "tokyo"),
+    ("Q: What is the capital of Spain?\nA:", "madrid"),
+    ("Q: What is the capital of Egypt?\nA:", "cairo"),
+    ("Q: What is the capital of Russia?\nA:", "moscow"),
+    ("Q: What is the capital of Australia?\nA:", "canberra"),
+]
+
+
+class BasicSchedulerStressMixin:
+    """Streaming + concurrent + long-prompt path probes."""
+
+    sanity_max_new_tokens_short: int = 64
+
+    def _stress_generate(self, prompt: str, max_new_tokens: int) -> str:
+        resp = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompt,
+                "sampling_params": {
+                    "temperature": 0.0,
+                    "max_new_tokens": max_new_tokens,
+                },
+            },
+            timeout=_REQUEST_TIMEOUT,
+        )
+        self.assertEqual(resp.status_code, 200)
+        return resp.json()["text"]
+
+    def test_streaming_response(self):
+        # SSE streaming exercises a different return path than non-stream
+        # /generate. Catches token-by-token streaming corruption and SSE
+        # framing bugs without changing the model.
+        with requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": "Q: What is the capital of France?\nA:",
+                "sampling_params": {
+                    "temperature": 0.0,
+                    "max_new_tokens": self.sanity_max_new_tokens_short,
+                },
+                "stream": True,
+            },
+            stream=True,
+            timeout=_REQUEST_TIMEOUT,
+        ) as resp:
+            self.assertEqual(resp.status_code, 200)
+            chunks_seen = 0
+            last_text = ""
+            for raw in resp.iter_lines(decode_unicode=True):
+                if not raw or not raw.startswith("data:"):
+                    continue
+                payload = raw[len("data:") :].strip()
+                if payload == "[DONE]":
+                    break
+                obj = json.loads(payload)
+                last_text = obj.get("text", last_text)
+                chunks_seen += 1
+            self.assertGreater(chunks_seen, 0)
+            self.assertIn("paris", last_text.lower())
+
+    def test_concurrent_requests(self):
+        # 8 parallel reqs share a system prefix but each has a distinct
+        # question suffix. Shared prefix exercises radix prefix caching
+        # across concurrent reqs; per-request suffix forces independent
+        # decode tails (different canonical answers). Catches concurrent
+        # scheduler hangs and prefix-cache cross-contamination.
+        results = [None] * len(_CONCURRENT_QA)
+
+        def worker(idx, suffix, expected):
+            try:
+                out = self._stress_generate(
+                    _CONCURRENT_PREFIX + suffix,
+                    self.sanity_max_new_tokens_short,
+                )
+                results[idx] = expected in out.lower()
+            except Exception:
+                results[idx] = False
+
+        threads = [
+            threading.Thread(target=worker, args=(i, suffix, expected))
+            for i, (suffix, expected) in enumerate(_CONCURRENT_QA)
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join(timeout=_REQUEST_TIMEOUT)
+
+        passed = sum(1 for r in results if r)
+        # Tolerate one stochastic miss; gibberish would fail all 8.
+        self.assertGreaterEqual(
+            passed,
+            len(_CONCURRENT_QA) - 1,
+            f"concurrent answers correct: {passed}/{len(_CONCURRENT_QA)}; results={results}",
+        )
+
+    def test_long_prompt(self):
+        # ~8k-token filler drives the chunked-prefill path through
+        # multiple chunks. Catches DeepEP / large-prompt kernel crashes
+        # that only fire on multi-chunk prefill.
+        filler = "the quick brown fox jumps over the lazy dog. " * 800
+        out = self._stress_generate(
+            f"Read the following text and then answer.\n{filler}\n\n"
+            "Q: What is the capital of France?\nA:",
+            self.sanity_max_new_tokens_short,
+        )
+        # Long-prompt substring match is best-effort (model may get
+        # distracted); primary assertion is the 200 + non-empty inside
+        # _stress_generate.
+        self.assertGreater(len(out), 0)

python/sglang/test/kits/cache_hit_kit.py

@@ -0,0 +1,403 @@
+import asyncio
+import json
+import time
+
+import aiohttp
+import requests
+
+from sglang.benchmark.datasets.random import sample_random_requests
+from sglang.benchmark.serving import RequestFuncOutput
+from sglang.benchmark.utils import get_tokenizer, remove_prefix
+
+AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60)
+
+
+async def async_request_sglang_generate(
+    payload,
+    url,
+    pbar=None,
+):
+    """Send a streaming request to the server and collect cache metrics.
+
+    Returns a RequestFuncOutput with additional cached_tokens and output_ids attributes.
+    """
+    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
+        headers = {}
+        generated_text = ""
+        all_output_ids = []
+        ttft = 0.0
+        st = time.perf_counter()
+        most_recent_timestamp = st
+        output = RequestFuncOutput()
+
+        try:
+            async with session.post(url=url, json=payload, headers=headers) as response:
+                if response.status == 200:
+                    prompt_tokens = 0
+                    cached_tokens = 0
+
+                    async for chunk_bytes in response.content:
+                        chunk_bytes = chunk_bytes.strip()
+                        if not chunk_bytes:
+                            continue
+
+                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
+                        latency = time.perf_counter() - st
+
+                        if chunk == "[DONE]":
+                            pass
+                        else:
+                            data = json.loads(chunk)
+
+                            # output_ids and text are always returned together
+                            if data.get("output_ids"):
+                                all_output_ids = data["output_ids"]
+                                generated_text = data.get("text", "")
+                                timestamp = time.perf_counter()
+
+                                if ttft == 0.0:
+                                    ttft = time.perf_counter() - st
+                                    output.ttft = ttft
+                                    prompt_tokens = (data.get("meta_info") or {}).get(
+                                        "prompt_tokens", 0
+                                    )
+                                    cached_tokens = (data.get("meta_info") or {}).get(
+                                        "cached_tokens", 0
+                                    )
+                                else:
+                                    output.itl.append(timestamp - most_recent_timestamp)
+
+                                most_recent_timestamp = timestamp
+
+                    output.generated_text = generated_text
+                    output.output_ids = all_output_ids
+                    output.success = True
+                    output.latency = latency
+                    output.prompt_len = prompt_tokens
+                    output.cached_tokens = cached_tokens
+                    output.generated_len = len(output.itl) + 1
+                else:
+                    output.error = response.reason or ""
+                    output.success = False
+        except Exception as e:
+            output.success = False
+            output.error = str(e)
+            print(f"Request failed: {e}")
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+async def async_request_openai_chat_completions(
+    payload,
+    url,
+    pbar=None,
+):
+    """Send a streaming request to an OpenAI-compatible /v1/chat/completions endpoint.
+
+    Returns a RequestFuncOutput with the same dynamic attributes as
+    async_request_sglang_generate (except output_ids, which is unavailable).
+    """
+    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
+        generated_text = ""
+        ttft = 0.0
+        latency = 0.0
+        st = time.perf_counter()
+        most_recent_timestamp = st
+        output = RequestFuncOutput()
+
+        try:
+            async with session.post(url=url, json=payload) as response:
+                if response.status == 200:
+                    prompt_tokens = 0
+                    cached_tokens = 0
+                    completion_tokens = 0
+
+                    async for chunk_bytes in response.content:
+                        chunk_bytes = chunk_bytes.strip()
+                        if not chunk_bytes:
+                            continue
+
+                        chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
+                        latency = time.perf_counter() - st
+
+                        if chunk == "[DONE]":
+                            pass
+                        else:
+                            data = json.loads(chunk)
+
+                            # Streaming token chunks
+                            if data.get("choices"):
+                                raw_delta = data["choices"][0].get("delta")
+                                text = raw_delta.get("content", "") if raw_delta else ""
+                                if text:
+                                    generated_text += text
+                                    timestamp = time.perf_counter()
+
+                                    if ttft == 0.0:
+                                        ttft = time.perf_counter() - st
+                                        output.ttft = ttft
+                                    else:
+                                        output.itl.append(
+                                            timestamp - most_recent_timestamp
+                                        )
+
+                                    most_recent_timestamp = timestamp
+
+                            # Final chunk with usage stats
+                            usage = data.get("usage")
+                            if usage:
+                                prompt_tokens = usage.get("prompt_tokens", 0)
+                                completion_tokens = usage.get("completion_tokens", 0)
+                                details = usage.get("prompt_tokens_details", {}) or {}
+                                cached_tokens = details.get("cached_tokens", 0)
+
+                    output.generated_text = generated_text
+                    output.output_ids = []  # Not available from OpenAI endpoint
+                    output.success = True
+                    output.latency = latency
+                    output.prompt_len = prompt_tokens
+                    output.cached_tokens = cached_tokens
+                    output.generated_len = (
+                        completion_tokens if completion_tokens else len(output.itl) + 1
+                    )
+                else:
+                    output.error = response.reason or ""
+                    output.success = False
+        except Exception as e:
+            output.success = False
+            output.error = str(e)
+            print(f"Request failed: {e}")
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+def gen_payload_openai(messages, output_len, model):
+    return {
+        "model": model,
+        "messages": messages,
+        "max_tokens": output_len,
+        "temperature": 0.0,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+    }
+
+
+def gen_payload(input_ids, output_len, lora_path=""):
+    return {
+        "input_ids": input_ids,
+        "sampling_params": {
+            "temperature": 0.0,
+            "max_new_tokens": output_len,
+            "ignore_eos": True,
+        },
+        "stream": True,
+        "stream_options": {"include_usage": True},
+        "lora_path": lora_path,
+        "return_logprob": False,
+        "logprob_start_len": -1,
+    }
+
+
+async def _send_round(
+    payloads,
+    url,
+    max_parallel,
+):
+    """Send a batch of payloads concurrently with concurrency limit."""
+    semaphore = asyncio.Semaphore(max_parallel)
+
+    async def _send_one(payload):
+        async with semaphore:
+            return await async_request_sglang_generate(payload, url)
+
+    tasks = [asyncio.create_task(_send_one(p)) for p in payloads]
+    return await asyncio.gather(*tasks)
+
+
+def _get_page_size(base_url: str) -> int:
+    """Query server for page_size used by radix cache."""
+    try:
+        resp = requests.get(f"{base_url}/server_info", timeout=10)
+        resp.raise_for_status()
+        info = resp.json()
+        return info.get("page_size", 1)
+    except Exception:
+        return 1
+
+
+def run_multiturn_cache_hit_test(
+    base_url: str,
+    model_path: str,
+    num_clients: int = 8,
+    num_rounds: int = 3,
+    request_length: int = 256,
+    output_length: int = 32,
+    miss_tolerance: int = 1,
+    sub_question_input_length: int = 0,
+    lora_path: str = "",
+    dataset_path: str = "",
+    max_parallel: int = 64,
+    seed: int = 1,
+) -> dict:
+    """Run a multi-turn workload and verify cache hit rate.
+
+    Sends requests in round-barrier mode: all clients complete round i
+    before round i+1 starts, ensuring deterministic cache state.
+
+    The expected cache hit rate is self-computed from the workload structure:
+    - Round 0: expected cached_tokens = 0 (cold start after flush)
+    - Round r (r >= 1): each client's prefix from round r-1 should be cached,
+      minus up to previous round's (prompt_len + decoding output - miss_tolerance) // page * page.
+
+    Returns metrics dict with per-round and overall cache_hit_rate.
+    """
+    import random
+
+    import numpy as np
+
+    random.seed(seed)
+    np.random.seed(seed)
+
+    generate_url = f"{base_url}/generate"
+    page_size = _get_page_size(base_url)
+
+    # Flush cache for clean state
+    requests.post(f"{base_url}/flush_cache")
+    time.sleep(1)
+
+    # Resolve sub-question length (0 means same as request_length)
+    effective_sub_len = (
+        sub_question_input_length if sub_question_input_length != 0 else request_length
+    )
+
+    # Sample initial prompts and sub-question prompts as token ids
+    tokenizer = get_tokenizer(model_path)
+
+    initial_inputs = sample_random_requests(
+        input_len=request_length,
+        output_len=output_length,
+        num_prompts=num_clients,
+        range_ratio=1.0,
+        tokenizer=tokenizer,
+        dataset_path=dataset_path,
+        return_text=False,
+    )
+    # r.prompt is now List[int] when return_text=False
+    initial_token_ids = [list(r.prompt) for r in initial_inputs]
+
+    sub_question_inputs = sample_random_requests(
+        input_len=effective_sub_len,
+        output_len=output_length,
+        num_prompts=num_clients * max(num_rounds - 1, 1),
+        range_ratio=1.0,
+        tokenizer=tokenizer,
+        dataset_path=dataset_path,
+        return_text=False,
+    )
+    sub_question_token_ids = [list(r.prompt) for r in sub_question_inputs]
+
+    # Per-round metrics and per-client tracking for expected cache computation
+    round_metrics = {
+        i: {"prompt_len": [], "cached_tokens": [], "ttft": []}
+        for i in range(num_rounds)
+    }
+    # Track the previous round's prompt_len per client to compute expected cache
+    prev_prompt_lens = [0] * num_clients
+    # histories now stores List[int] (token ids) for each client
+    histories = [list(ids) for ids in initial_token_ids]
+    sub_idx = 0
+
+    for round_num in range(num_rounds):
+        payloads = [gen_payload(h, output_length, lora_path) for h in histories]
+        responses = asyncio.run(_send_round(payloads, generate_url, max_parallel))
+
+        for i, resp in enumerate(responses):
+            assert resp.success, f"Round {round_num}, client {i} failed: {resp.error}"
+
+            round_metrics[round_num]["prompt_len"].append(resp.prompt_len)
+            round_metrics[round_num]["cached_tokens"].append(resp.cached_tokens)
+            round_metrics[round_num]["ttft"].append(resp.ttft)
+
+            # Verify cache hit against expected value
+            if round_num == 0:
+                # Cold start: no cache expected
+                expected_cached = 0
+            else:
+                # Previous round's prompt + output are in cache.
+                # Radix cache aligns to page_size, so the last partial page
+                # may not be cached.
+                cacheable = prev_prompt_lens[i] + output_length - miss_tolerance
+                expected_cached = (cacheable // page_size) * page_size
+
+            msg = (
+                f"Round {round_num}, client {i}: "
+                f"cached_tokens={resp.cached_tokens}, "
+                f"expected>={expected_cached} "
+                f"(prev_prompt={prev_prompt_lens[i]}, "
+                f"output={output_length}, page_size={page_size})"
+            )
+
+            print(msg)
+
+            assert resp.cached_tokens >= expected_cached
+            # Upper bound: cached tokens are a subset of the prompt, so they can
+            # never exceed prompt_len. In PD disaggregation with decode radix
+            # cache, the shared prefix was previously counted on both the prefill
+            # and the decode node, making cached_tokens exceed prompt_len.
+            assert resp.cached_tokens <= resp.prompt_len, (
+                f"Round {round_num}, client {i}: cached_tokens="
+                f"{resp.cached_tokens} exceeds prompt_len={resp.prompt_len} "
+                f"(double-counted prefix across prefill/decode)"
+            )
+
+            # Record this round's prompt_len for next round's expected calc
+            prev_prompt_lens[i] = resp.prompt_len
+
+            # Accumulate history for next round using output_ids (token ids)
+            histories[i].extend(resp.output_ids)
+            if round_num < num_rounds - 1:
+                histories[i].extend(sub_question_token_ids[sub_idx])
+                sub_idx += 1
+
+    # Compute per-round and overall cache hit rate
+    total_prompt = 0
+    total_cached = 0
+    result = {"rounds": {}, "overall": {}}
+
+    for r in range(num_rounds):
+        rm = round_metrics[r]
+        r_prompt = sum(rm["prompt_len"])
+        r_cached = sum(rm["cached_tokens"])
+        r_hit_rate = r_cached / r_prompt if r_prompt > 0 else 0.0
+        r_avg_ttft = sum(rm["ttft"]) / len(rm["ttft"]) if rm["ttft"] else 0.0
+
+        result["rounds"][f"round_{r}"] = {
+            "cache_hit_rate": r_hit_rate,
+            "average_ttft": r_avg_ttft,
+            "total_prompt_tokens": r_prompt,
+            "total_cached_tokens": r_cached,
+            "request_count": len(rm["ttft"]),
+        }
+
+        total_prompt += r_prompt
+        total_cached += r_cached
+
+        print(
+            f"  Round {r}: cache_hit_rate={r_hit_rate:.4f}, "
+            f"avg_ttft={r_avg_ttft:.4f}s, "
+            f"cached={r_cached}/{r_prompt} tokens"
+        )
+
+    overall_hit_rate = total_cached / total_prompt if total_prompt > 0 else 0.0
+    result["overall"] = {
+        "cache_hit_rate": overall_hit_rate,
+        "total_prompt_tokens": total_prompt,
+        "total_cached_tokens": total_cached,
+    }
+    print(f"  Overall cache_hit_rate={overall_hit_rate:.4f}")
+
+    return result

python/sglang/test/kits/ebnf_constrained_kit.py

@@ -0,0 +1,225 @@
+import json
+
+import requests
+
+
+class EBNFConstrainedMixin:
+
+    ebnf_grammar = 'root ::= "test"'  # Default grammar
+
+    def _run_decode_ebnf(
+        self,
+        ebnf,
+        expected_patterns,
+        prompt,
+        return_logprob=False,
+        top_logprobs_num=0,
+        n=1,
+    ):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompt,
+                "sampling_params": {
+                    "temperature": 0 if n == 1 else 0.5,
+                    "max_new_tokens": 128,
+                    "n": n,
+                    "ebnf": ebnf,
+                },
+                "stream": False,
+                "return_logprob": return_logprob,
+                "top_logprobs_num": top_logprobs_num,
+                "logprob_start_len": 0,
+            },
+        )
+
+        ret = response.json()
+        print(json.dumps(ret, indent=2))
+        print("=" * 100)
+
+        if not isinstance(ret, list):
+            self.fail(f"Expected response to be a list, but got {type(ret)}")
+
+        for item in ret:
+            text = item.get("text", "").strip()
+            if not text:
+                self.fail("Generated text is empty.")
+
+            match = False
+            for pattern in expected_patterns:
+                if self.regex_match(text, pattern):
+                    match = True
+                    break
+            if not match:
+                self.fail(f"Text '{text}' does not match any of the allowed patterns.")
+
+    def regex_match(self, text, pattern):
+        import re
+
+        return re.match(pattern, text) is not None
+
+    def test_ebnf_generate_email(self):
+        self.__class__.ebnf_grammar = 'root ::= "user@example.com"'
+        allowed_patterns = [r"^user@example\.com$"]
+        prompt = "Generate an email address:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_greeting(self):
+        self.__class__.ebnf_grammar = 'root ::= "Hello" | "Hi" | "Hey"'
+        allowed_patterns = [r"^(Hello|Hi|Hey)$"]
+        prompt = "Generate a greeting:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_number(self):
+        self.__class__.ebnf_grammar = """
+        root ::= digit digit digit
+        digit ::= [0-9]
+        """
+        allowed_patterns = [r"^\d{3}$"]
+        prompt = "Generate a three-digit number:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_phone(self):
+        self.__class__.ebnf_grammar = """
+        root ::= "(" area ")" " " prefix "-" line
+        area ::= [0-9] [0-9] [0-9]
+        prefix ::= [0-9] [0-9] [0-9]
+        line ::= [0-9] [0-9] [0-9] [0-9]
+        """
+        allowed_patterns = [r"^\(\d{3}\) \d{3}-\d{4}$"]
+        prompt = "Generate a phone number:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_date(self):
+        self.__class__.ebnf_grammar = """
+        root ::= year "-" month "-" day
+        year ::= "2024"
+        month ::= "01" | "02" | "03" | "04" | "05" | "06" | "07" | "08" | "09" | "10" | "11" | "12"
+        day ::= "01" | "02" | "03" | "04" | "05" | "06" | "07" | "08" | "09" | "10" |
+               "11" | "12" | "13" | "14" | "15" | "16" | "17" | "18" | "19" | "20" |
+               "21" | "22" | "23" | "24" | "25" | "26" | "27" | "28" | "29" | "30" | "31"
+        """
+        allowed_patterns = [r"^2024-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$"]
+        prompt = "Generate a date in YYYY-MM-DD format:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_hex_color(self):
+        self.__class__.ebnf_grammar = """
+        root ::= "#" hex hex hex hex hex hex
+        hex ::= [0-9] | [A-F]
+        """
+        allowed_patterns = [r"^#[0-9A-F]{6}$"]
+        prompt = "Generate a hex color code:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_complex_json(self):
+        self.__class__.ebnf_grammar = """
+        root ::= object
+        object ::= "{" ws pair (ws "," ws pair)* ws "}"
+        pair ::= "\\"name\\"" ws ":" ws value |
+                 "\\"age\\"" ws ":" ws number |
+                 "\\"city\\"" ws ":" ws string
+        value ::= string | number
+        string ::= "\\"" [a-zA-Z0-9 ]+ "\\""
+        number ::= [1-9] [0-9]*
+        ws ::= [ ]*
+        """
+        allowed_patterns = [
+            r'^{\s*"name"\s*:\s*"[a-zA-Z0-9 ]+"\s*,\s*"age"\s*:\s*[1-9][0-9]*\s*,\s*"city"\s*:\s*"[a-zA-Z0-9 ]+"\s*}$',
+        ]
+        prompt = "Generate a simple JSON with name, age, and city:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_custom_log_format(self):
+        self.__class__.ebnf_grammar = """
+        root ::= logentry
+        logentry ::= "[" datetime "] " level ": System.process - " message
+        datetime ::= "2024-01-01T12:00:00Z"
+        level ::= "INFO"
+        message ::= "Operation " [a-z]+ " successfully"
+        """
+        allowed_patterns = [
+            r"^\[2024-01-01T12:00:00Z\] INFO: System\.process - Operation [a-z]+ successfully$"
+        ]
+        prompt = "Generate a log entry:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_ebnf_generate_all_optional_function_params(self):
+        """Test function call with all optional parameters - verifies flexible ordering."""
+        self.__class__.ebnf_grammar = """
+        root ::= function_call
+        function_call ::= call_config_service
+        call_config_service ::= "{" "\\"name\\"" ":" "\\"config_service\\"" ", " "\\"arguments\\"" ":" arguments_config_service "}"
+        arguments_config_service ::= "{" ( "\\"theme\\"" ":" ("\\"light\\"" | "\\"dark\\"") ( "," "\\"language\\"" ":" ("\\"en\\"" | "\\"es\\"" | "\\"fr\\"") )? ( "," "\\"notifications\\"" ":" ("true" | "false") )? | "\\"language\\"" ":" ("\\"en\\"" | "\\"es\\"" | "\\"fr\\"") ( "," "\\"notifications\\"" ":" ("true" | "false") )? | "\\"notifications\\"" ":" ("true" | "false") )? "}"
+        """
+        # Test patterns that should match - flexible ordering of optional parameters
+        allowed_patterns = [
+            # Empty arguments
+            r'^\{"name":"config_service",\s*"arguments":\{\}\}$',
+            # Single optional parameters (any can appear first)
+            r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)"\}\}$',
+            r'^\{"name":"config_service",\s*"arguments":\{"language":"(en|es|fr)"\}\}$',
+            r'^\{"name":"config_service",\s*"arguments":\{"notifications":(true|false)\}\}$',
+            # Two optional parameters (in any order)
+            r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)",\s*"language":"(en|es|fr)"\}\}$',
+            r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)",\s*"notifications":(true|false)\}\}$',
+            r'^\{"name":"config_service",\s*"arguments":\{"language":"(en|es|fr)",\s*"notifications":(true|false)\}\}$',
+            # All three optional parameters
+            r'^\{"name":"config_service",\s*"arguments":\{"theme":"(light|dark)",\s*"language":"(en|es|fr)",\s*"notifications":(true|false)\}\}$',
+        ]
+        prompt = "Configure the service with optional settings:"
+
+        self._run_decode_ebnf(
+            ebnf=self.__class__.ebnf_grammar,
+            expected_patterns=allowed_patterns,
+            prompt=prompt,
+            n=5,
+        )

python/sglang/test/kits/eval_accuracy_kit.py

@@ -0,0 +1,419 @@
+from types import SimpleNamespace
+from typing import Optional
+
+import requests
+
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import is_in_amd_ci, is_in_ci, write_github_step_summary
+
+_THRESHOLD_NOT_SET = float("nan")
+
+
+def _check_accept_length(test_case, base_url, threshold=None):
+    """Print speculative accept length; optionally assert it exceeds threshold."""
+    try:
+        server_info = requests.get(base_url + "/server_info").json()
+        val = server_info["internal_states"][0]["avg_spec_accept_length"]
+    except (KeyError, IndexError, requests.RequestException):
+        return
+    print(f"avg_spec_accept_length={val:.4f}")
+    if threshold is not None:
+        test_case.assertGreater(val, threshold)
+
+
+def _finalize_eval(
+    test_case,
+    *,
+    eval_name: str,
+    score: float,
+    score_threshold: float,
+    accept_length_thres: Optional[float] = None,
+    summary_label: Optional[str] = None,
+):
+    """Shared driver tail: CI step summary, accept-length check, threshold assert."""
+    if is_in_ci():
+        label = summary_label or f"test_{eval_name}"
+        write_github_step_summary(f"### {label}\n{eval_name}_score={score:.4f}\n")
+    _check_accept_length(test_case, test_case.base_url, accept_length_thres)
+    test_case.assertGreaterEqual(score, score_threshold)
+
+
+def _run_accuracy_eval(
+    test_case,
+    *,
+    eval_name: str,
+    score_threshold: float,
+    num_examples: Optional[int],
+    num_threads: int,
+    accept_length_thres: Optional[float] = None,
+    summary_label: Optional[str] = None,
+    **eval_overrides,
+):
+    """Shared driver for the accuracy mixins below.
+
+    Runs ``run_eval`` for ``eval_name`` against the test class's server
+    (``base_url`` / ``model``), records a CI step summary, asserts the score
+    meets ``score_threshold``, and checks the speculative accept length.
+
+    ``eval_overrides`` (e.g. ``api``, ``max_tokens``, ``temperature``,
+    ``top_p``, ``num_shots``) are forwarded to ``run_eval`` only when not
+    ``None``, so the common case stays identical to ``run_eval``'s defaults.
+    Returns the metrics dict.
+    """
+    assert (
+        score_threshold == score_threshold
+    ), f"{type(test_case).__name__} must set the {eval_name} score threshold"
+
+    kwargs = dict(
+        base_url=test_case.base_url,
+        model=getattr(test_case, "model", None),
+        eval_name=eval_name,
+        num_examples=num_examples,
+        num_threads=num_threads,
+    )
+    kwargs.update({k: v for k, v in eval_overrides.items() if v is not None})
+
+    metrics = run_eval(SimpleNamespace(**kwargs))
+    print(f"{eval_name} {metrics=}")
+    _finalize_eval(
+        test_case,
+        eval_name=eval_name,
+        score=metrics["score"],
+        score_threshold=score_threshold,
+        accept_length_thres=accept_length_thres,
+        summary_label=summary_label,
+    )
+    return metrics
+
+
+def _run_sgl_eval(
+    test_case,
+    *,
+    eval_name: str,
+    score_threshold: float,
+    metric: str = "score",
+    n_repeats: int = 1,
+    num_examples: Optional[int] = None,
+    num_threads: int = 512,
+    thinking: bool = True,
+    reasoning_effort: Optional[str] = None,
+    max_tokens: Optional[int] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    accept_length_thres: Optional[float] = None,
+    summary_label: Optional[str] = None,
+):
+    """Shared sgl-eval driver for the reasoning mixins and the ``sgl_eval`` backend.
+
+    Runs ``eval_name`` via the sgl-eval Python API (``registry.get`` ->
+    ``EvalSpec.run``) against the test class's server, records a CI step summary,
+    asserts the score meets ``score_threshold``, and checks the speculative accept
+    length. ``thinking=True`` sends per-request ``chat_template_kwargs={"thinking":
+    True}`` so the server separates reasoning from the final answer. Skips the test
+    if sgl-eval (git-only) is not installed. Returns the RunResult.
+    """
+    assert (
+        score_threshold == score_threshold
+    ), f"{type(test_case).__name__} must set the {eval_name} score threshold"
+
+    try:
+        from sgl_eval.registry import get as get_eval_spec
+        from sgl_eval.sampler import ChatCompletionSampler
+        from sgl_eval.types import GenConfig
+    except ImportError:
+        test_case.skipTest(
+            "sgl-eval not installed; pip install "
+            "'sgl-eval @ git+https://github.com/sgl-project/sgl-eval'"
+        )
+
+    base_url = test_case.base_url.rstrip("/")
+    if not base_url.endswith("/v1"):
+        base_url += "/v1"
+    sampler = ChatCompletionSampler(
+        base_url=base_url, model=getattr(test_case, "model", None), api_key="EMPTY"
+    )
+
+    gen_kwargs = dict(
+        max_tokens=max_tokens,
+        reasoning_effort=reasoning_effort,
+        chat_template_kwargs={"thinking": True} if thinking else None,
+    )
+    if temperature is not None:
+        gen_kwargs["temperature"] = temperature
+    if top_p is not None:
+        gen_kwargs["top_p"] = top_p
+
+    result = get_eval_spec(eval_name).run(
+        sampler=sampler,
+        gen=GenConfig(**gen_kwargs),
+        n_repeats=n_repeats,
+        num_examples=num_examples,
+        num_threads=num_threads,
+        predictions_writer=None,
+        load_examples=None,
+    )
+    score = result.aggregate[metric]
+    print(f"{eval_name} sgl-eval {metric}={score:.4f}")
+    _finalize_eval(
+        test_case,
+        eval_name=eval_name,
+        score=score,
+        score_threshold=score_threshold,
+        accept_length_thres=accept_length_thres,
+        summary_label=summary_label,
+    )
+    return result
+
+
+class GSM8KMixin:
+    """Mixin for GSM8K evaluation.
+
+    Backend is selectable via ``gsm8k_backend`` (default ``"run_eval"``: OpenAI
+    completion API, 5-shot; or ``"sgl_eval"``: sgl-eval chat + boxed/sympy grader,
+    skipped if sgl-eval is not installed). The canonical threshold/count knobs are
+    ``gsm8k_score_threshold`` / ``gsm8k_num_examples``; the legacy
+    ``gsm8k_accuracy_thres`` / ``gsm8k_num_questions`` are still honored.
+
+    Required attributes on the test class:
+        base_url: str
+        gsm8k_score_threshold: float
+
+    Optional attributes:
+        model: str (if not set, auto-detected from server)
+    """
+
+    gsm8k_score_threshold: float = _THRESHOLD_NOT_SET
+    gsm8k_accuracy_thres: float = _THRESHOLD_NOT_SET  # legacy alias
+    gsm8k_num_examples: Optional[int] = None
+    gsm8k_num_questions: int = 200  # legacy alias
+    gsm8k_accept_length_thres: Optional[float] = None
+    gsm8k_num_threads: int = 128
+    gsm8k_num_shots: int = 5  # run_eval backend only
+    gsm8k_backend: str = "run_eval"  # "run_eval" | "sgl_eval"
+    gsm8k_thinking: bool = False  # sgl_eval backend
+    gsm8k_n_repeats: int = 1  # sgl_eval backend
+
+    def test_gsm8k(self):
+        requests.get(self.base_url + "/flush_cache")
+        threshold = self.gsm8k_score_threshold
+        if threshold != threshold:  # canonical unset (NaN) -> legacy alias
+            threshold = self.gsm8k_accuracy_thres
+        num_examples = (
+            self.gsm8k_num_examples
+            if self.gsm8k_num_examples is not None
+            else self.gsm8k_num_questions
+        )
+        if self.gsm8k_backend == "sgl_eval":
+            _run_sgl_eval(
+                self,
+                eval_name="gsm8k",
+                score_threshold=threshold,
+                n_repeats=self.gsm8k_n_repeats,
+                num_examples=num_examples,
+                num_threads=self.gsm8k_num_threads,
+                thinking=self.gsm8k_thinking,
+                accept_length_thres=self.gsm8k_accept_length_thres,
+            )
+        else:
+            _run_accuracy_eval(
+                self,
+                eval_name="gsm8k",
+                score_threshold=threshold,
+                num_examples=num_examples,
+                num_threads=self.gsm8k_num_threads,
+                accept_length_thres=self.gsm8k_accept_length_thres,
+                api="completion",
+                max_tokens=512,
+                num_shots=self.gsm8k_num_shots,
+            )
+
+
+class MMLUMixin:
+    """Mixin for MMLU evaluation.
+
+    Backend is selectable via ``mmlu_backend`` (default ``"run_eval"``; or
+    ``"sgl_eval"``: sgl-eval multichoice grader, skipped if sgl-eval is not
+    installed).
+
+    Required attributes on the test class:
+        base_url: str
+        model: str
+        mmlu_score_threshold: float
+    """
+
+    mmlu_score_threshold: float = _THRESHOLD_NOT_SET
+    mmlu_accept_length_thres: Optional[float] = None
+    mmlu_num_examples: int = 5000
+    mmlu_num_threads: int = 1024
+    mmlu_backend: str = "run_eval"  # "run_eval" | "sgl_eval"
+    mmlu_thinking: bool = False  # sgl_eval backend
+    mmlu_n_repeats: int = 1  # sgl_eval backend
+
+    def test_mmlu(self):
+        if self.mmlu_backend == "sgl_eval":
+            _run_sgl_eval(
+                self,
+                eval_name="mmlu",
+                score_threshold=self.mmlu_score_threshold,
+                n_repeats=self.mmlu_n_repeats,
+                num_examples=self.mmlu_num_examples,
+                num_threads=self.mmlu_num_threads,
+                thinking=self.mmlu_thinking,
+                accept_length_thres=self.mmlu_accept_length_thres,
+            )
+        else:
+            _run_accuracy_eval(
+                self,
+                eval_name="mmlu",
+                score_threshold=self.mmlu_score_threshold,
+                num_examples=self.mmlu_num_examples,
+                num_threads=self.mmlu_num_threads,
+                accept_length_thres=self.mmlu_accept_length_thres,
+            )
+
+
+class GPQAMixin:
+    """Mixin for GPQA-Diamond evaluation (graduate-level multiple choice).
+
+    Runs via the sgl-eval Python API (the test is skipped if sgl-eval is not
+    installed). ``gpqa_thinking`` defaults to True, which
+    enables per-request thinking so the server separates reasoning from the final
+    answer.
+
+    Required attributes on the test class:
+        base_url: str
+        model: str
+        gpqa_score_threshold: float
+
+    Optional sampling knobs (default to sgl-eval's defaults when unset). Set these
+    for reasoning models -- e.g. DeepSeek-V4 Think-Max wants
+    gpqa_reasoning_effort="max", gpqa_max_tokens=200000, gpqa_temperature=1.0,
+    gpqa_top_p=1.0. GPQA-Diamond is 198 questions; raise gpqa_n_repeats (e.g. 16)
+    for a stable number.
+    """
+
+    gpqa_score_threshold: float = _THRESHOLD_NOT_SET
+    gpqa_accept_length_thres: Optional[float] = None
+    gpqa_num_examples: Optional[int] = None
+    gpqa_num_threads: int = 1024
+    gpqa_n_repeats: int = 1
+    gpqa_thinking: bool = True
+    gpqa_reasoning_effort: Optional[str] = None
+    gpqa_max_tokens: Optional[int] = None
+    gpqa_temperature: Optional[float] = None
+    gpqa_top_p: Optional[float] = None
+
+    def test_gpqa(self):
+        _run_sgl_eval(
+            self,
+            eval_name="gpqa",
+            score_threshold=self.gpqa_score_threshold,
+            n_repeats=self.gpqa_n_repeats,
+            num_examples=self.gpqa_num_examples,
+            num_threads=self.gpqa_num_threads,
+            thinking=self.gpqa_thinking,
+            reasoning_effort=self.gpqa_reasoning_effort,
+            max_tokens=self.gpqa_max_tokens,
+            temperature=self.gpqa_temperature,
+            top_p=self.gpqa_top_p,
+            accept_length_thres=self.gpqa_accept_length_thres,
+        )
+
+
+class AIME25Mixin:
+    """Mixin for AIME 2025 evaluation (competition math, integer answers).
+
+    Runs via the sgl-eval Python API (the test is skipped if sgl-eval is not
+    installed). ``aime25_thinking`` defaults to True, which
+    enables per-request thinking so the server separates reasoning from the final
+    answer.
+
+    Required attributes on the test class:
+        base_url: str
+        model: str
+        aime25_score_threshold: float
+
+    Optional sampling knobs (default to sgl-eval's defaults when unset). Set these
+    for reasoning models -- e.g. DeepSeek-V4 Think-Max wants
+    aime25_reasoning_effort="max", aime25_max_tokens=200000, aime25_temperature=1.0,
+    aime25_top_p=1.0. AIME25 has only 30 problems, so it is high variance; raise
+    aime25_n_repeats (e.g. 16) for a stable number.
+    """
+
+    aime25_score_threshold: float = _THRESHOLD_NOT_SET
+    aime25_accept_length_thres: Optional[float] = None
+    aime25_num_examples: Optional[int] = None
+    aime25_num_threads: int = 1024
+    aime25_n_repeats: int = 1
+    aime25_thinking: bool = True
+    aime25_reasoning_effort: Optional[str] = None
+    aime25_max_tokens: Optional[int] = None
+    aime25_temperature: Optional[float] = None
+    aime25_top_p: Optional[float] = None
+
+    def test_aime25(self):
+        _run_sgl_eval(
+            self,
+            eval_name="aime25",
+            score_threshold=self.aime25_score_threshold,
+            n_repeats=self.aime25_n_repeats,
+            num_examples=self.aime25_num_examples,
+            num_threads=self.aime25_num_threads,
+            thinking=self.aime25_thinking,
+            reasoning_effort=self.aime25_reasoning_effort,
+            max_tokens=self.aime25_max_tokens,
+            temperature=self.aime25_temperature,
+            top_p=self.aime25_top_p,
+            accept_length_thres=self.aime25_accept_length_thres,
+        )
+
+
+class HumanEvalMixin:
+    """Mixin for HumanEval evaluation.
+
+    Required attributes on the test class:
+        base_url: str
+        model: str
+        humaneval_score_threshold: float
+    """
+
+    humaneval_score_threshold: float = _THRESHOLD_NOT_SET
+    humaneval_score_threshold_amd: Optional[float] = None
+    humaneval_num_threads: int = 1024
+
+    def test_human_eval(self):
+        threshold = self.humaneval_score_threshold
+        if is_in_amd_ci() and self.humaneval_score_threshold_amd is not None:
+            threshold = self.humaneval_score_threshold_amd
+
+        _run_accuracy_eval(
+            self,
+            eval_name="humaneval",
+            score_threshold=threshold,
+            num_examples=None,
+            num_threads=self.humaneval_num_threads,
+            summary_label="test_human_eval",
+        )
+
+
+class MGSMEnMixin:
+    """Mixin for MGSM English evaluation.
+
+    Required attributes on the test class:
+        base_url: str
+        model: str
+        mgsm_en_score_threshold: float
+    """
+
+    mgsm_en_score_threshold: float = _THRESHOLD_NOT_SET
+    mgsm_en_num_examples: Optional[int] = None
+    mgsm_en_num_threads: int = 1024
+
+    def test_mgsm_en(self):
+        _run_accuracy_eval(
+            self,
+            eval_name="mgsm_en",
+            score_threshold=self.mgsm_en_score_threshold,
+            num_examples=self.mgsm_en_num_examples,
+            num_threads=self.mgsm_en_num_threads,
+        )

python/sglang/test/kits/fwd_occupancy_kit.py

@@ -0,0 +1,261 @@
+"""Single-batch decode GPU occupancy sanity kit.
+
+Probes ``sglang:fwd_occupancy`` (a 0-100 percentage averaged over the
+last ``decode_log_interval`` batches; resets to NaN at window
+boundaries) under one long single-batch ``/generate`` request, and
+asserts median above a threshold. Single-batch is where CPU overhead
+dominates -- overlap scheduler / cuda graph regressions surface here
+before batched throughput moves.
+
+Prerequisites on the consuming server:
+    env:          SGLANG_ENABLE_METRICS_DEVICE_TIMER=1
+    server flag:  --enable-metrics
+
+Mix into a ``CustomTestCase`` subclass exposing ``self.base_url``.
+"""
+
+import re
+import statistics
+import threading
+import time
+
+import requests
+import tabulate
+
+_FWD_OCCUPANCY_RE = re.compile(
+    r"^sglang:fwd_occupancy(?:\{[^}]*\})?\s+(\S+)", re.MULTILINE
+)
+_GENERATE_REQUEST_TIMEOUT = 600
+_METRICS_REQUEST_TIMEOUT = 10
+
+
+class FwdOccupancyMixin:
+    """Assert single-batch ``sglang:fwd_occupancy`` median > threshold."""
+
+    fwd_occupancy_threshold: float = 95.0
+    fwd_occupancy_min_samples: int = 5
+    fwd_occupancy_scrape_interval: float = 0.5
+
+    # Spec-decoding accept-length floor. Only enforced when the server
+    # is running with a spec algorithm (avg_spec_accept_length present
+    # in /server_info); silently skipped otherwise. EAGLE3 3/1/4 on
+    # 5090 + Llama-3.1-8B measured ~2.0 in CI; 1.8 leaves a small
+    # buffer while still catching silent fallback to vanilla (~1.0).
+    fwd_occupancy_acc_length_threshold: float = 1.8
+
+    # Warmup: one short request to fill cuda graphs + get the
+    # device-timer past its first NaN window.
+    fwd_occupancy_warmup_max_new_tokens: int = 64
+    fwd_occupancy_warmup_settle_seconds: float = 1.0
+
+    # Measurement: one long single-batch request -- max_new_tokens must
+    # span several decode_log_interval windows for enough samples.
+    fwd_occupancy_max_new_tokens: int = 2048
+    fwd_occupancy_prompt: str = (
+        "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:"
+    )
+
+    def _scrape_fwd_occupancy(self):
+        """Max non-NaN gauge value across exposed labels (e.g. dp ranks);
+        None on transient scrape failure."""
+        try:
+            resp = requests.get(
+                self.base_url + "/metrics", timeout=_METRICS_REQUEST_TIMEOUT
+            )
+        except requests.RequestException:
+            return None
+        if resp.status_code != 200:
+            return None
+        vals = []
+        for raw in _FWD_OCCUPANCY_RE.findall(resp.text):
+            try:
+                v = float(raw)
+            except ValueError:
+                continue
+            if v == v:  # NaN filter (gauge resets to NaN on window boundary)
+                vals.append(v)
+        return max(vals) if vals else None
+
+    def _assert_metrics_device_timer_enabled(self):
+        """Fail loudly on missing flag/env -- otherwise a NaN-only gauge
+        looks like a real occupancy regression."""
+        resp = requests.get(
+            self.base_url + "/metrics", timeout=_METRICS_REQUEST_TIMEOUT
+        )
+        if resp.status_code != 200:
+            raise AssertionError(
+                f"/metrics returned {resp.status_code}; the test class's "
+                "server must be launched with --enable-metrics"
+            )
+        if "sglang:fwd_occupancy" not in resp.text:
+            raise AssertionError(
+                "sglang:fwd_occupancy gauge not exposed; set "
+                "SGLANG_ENABLE_METRICS_DEVICE_TIMER=1 in the server's env "
+                "and pass --enable-metrics"
+            )
+
+    def _fwd_occupancy_fire(self, prompt: str, max_new_tokens: int):
+        """Fire one /generate, return (meta_info, wall_time).
+        Must not be called concurrently -- that would break the
+        single-batch invariant."""
+        t0 = time.perf_counter()
+        try:
+            resp = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": prompt,
+                    "sampling_params": {
+                        "temperature": 0.0,
+                        "max_new_tokens": max_new_tokens,
+                        "ignore_eos": True,
+                    },
+                },
+                timeout=_GENERATE_REQUEST_TIMEOUT,
+            )
+        except requests.RequestException:
+            # Final stats-vs-threshold is the signal; individual fire
+            # failure isn't.
+            return {}, 0.0
+        elapsed = time.perf_counter() - t0
+        try:
+            return resp.json().get("meta_info", {}), elapsed
+        except ValueError:  # non-JSON body
+            return {}, elapsed
+
+    def _fwd_occupancy_warmup(self):
+        """Fill cuda graphs + step the device-timer past its first NaN
+        window before measurement starts."""
+        self._fwd_occupancy_fire(
+            "warmup " + self.fwd_occupancy_prompt,
+            self.fwd_occupancy_warmup_max_new_tokens,
+        )
+        time.sleep(self.fwd_occupancy_warmup_settle_seconds)
+
+    def _fwd_occupancy_measure(self):
+        """Background-fire one long single-batch request, scrape
+        /metrics on the foreground; return (non-NaN samples, perf)."""
+        samples = []
+        request_done = threading.Event()
+        result = {"meta_info": {}, "elapsed": 0.0}
+
+        def fire_one():
+            try:
+                result["meta_info"], result["elapsed"] = self._fwd_occupancy_fire(
+                    self.fwd_occupancy_prompt,
+                    self.fwd_occupancy_max_new_tokens,
+                )
+            finally:
+                request_done.set()
+
+        firer = threading.Thread(target=fire_one, daemon=True)
+        firer.start()
+
+        while not request_done.is_set():
+            v = self._scrape_fwd_occupancy()
+            if v is not None:
+                samples.append(v)
+            time.sleep(self.fwd_occupancy_scrape_interval)
+
+        firer.join(timeout=_GENERATE_REQUEST_TIMEOUT)
+        return samples, self._fwd_occupancy_perf(result)
+
+    def _fwd_occupancy_perf(self, result):
+        """Aggregate per-request perf metrics from the fire result
+        (input/output tokens, decode tps, mean inter-token latency,
+        wall-clock tps)."""
+        meta = result["meta_info"] or {}
+        elapsed = result["elapsed"]
+        out = meta.get("completion_tokens", 0) or 0
+        decode_tps = meta.get("decode_throughput", 0.0) or 0.0
+        return {
+            "input_tokens": meta.get("prompt_tokens", 0) or 0,
+            "output_tokens": out,
+            "decode_tps": decode_tps,
+            "mean_itl_ms": (1000.0 / decode_tps) if decode_tps > 0 else 0.0,
+            "wall_tps": (out / elapsed) if elapsed > 0 else 0.0,
+        }
+
+    def test_fwd_occupancy(self):
+        self._assert_metrics_device_timer_enabled()
+        self._fwd_occupancy_warmup()
+        samples, perf = self._fwd_occupancy_measure()
+
+        # The 2048-token decode above populates the spec running average
+        # if a spec algorithm is enabled; absent otherwise (vanilla
+        # decode skips this check).
+        try:
+            info = requests.get(
+                self.base_url + "/server_info", timeout=_METRICS_REQUEST_TIMEOUT
+            ).json()
+            avg_accept = info["internal_states"][0].get("avg_spec_accept_length")
+        except (requests.RequestException, KeyError, IndexError):
+            avg_accept = None
+
+        # Median is the steady-state signal; peak / p10 included in the
+        # assertion message for triage. Both tables print before any
+        # assertion so the numbers surface even on assertion failure.
+        samples_sorted = sorted(samples)
+        if samples_sorted:
+            median = statistics.median(samples_sorted)
+            peak = samples_sorted[-1]
+            p10_idx = min(len(samples_sorted) - 1, max(0, len(samples_sorted) // 10))
+            p10 = samples_sorted[p10_idx]
+        else:
+            median = peak = p10 = float("nan")
+
+        perf_rows = [
+            ["input tokens", perf["input_tokens"]],
+            ["output tokens", perf["output_tokens"]],
+            ["decode tps", f"{perf['decode_tps']:.2f}"],
+            ["mean itl (ms)", f"{perf['mean_itl_ms']:.2f}"],
+            ["wall tps", f"{perf['wall_tps']:.2f}"],
+        ]
+        if avg_accept is not None:
+            perf_rows.append(["avg spec accept", f"{avg_accept:.3f}"])
+        # Lead each table with a text title line so the two tables stay
+        # visually separated even when CI prefixes every line with a timestamp
+        # (which turns blank separator lines into non-empty lines).
+        print(
+            "\n[perf metrics]\n"
+            + tabulate.tabulate(
+                perf_rows, headers=["perf metric", "value"], tablefmt="github"
+            )
+        )
+
+        print(
+            "\n[fwd_occupancy stats]\n"
+            + tabulate.tabulate(
+                [
+                    ["samples (n)", len(samples)],
+                    ["median", f"{median:.2f}"],
+                    ["peak", f"{peak:.2f}"],
+                    ["p10", f"{p10:.2f}"],
+                    ["threshold", f"{self.fwd_occupancy_threshold:.2f}"],
+                ],
+                headers=["fwd_occupancy", "value"],
+                tablefmt="github",
+            )
+        )
+
+        self.assertGreaterEqual(
+            len(samples),
+            self.fwd_occupancy_min_samples,
+            f"only {len(samples)} non-NaN occupancy samples collected "
+            f"(need >= {self.fwd_occupancy_min_samples}); the measurement "
+            "window may be too short or the gauge stuck at NaN",
+        )
+        self.assertGreater(
+            median,
+            self.fwd_occupancy_threshold,
+            f"sglang:fwd_occupancy median={median:.2f} did not exceed "
+            f"threshold {self.fwd_occupancy_threshold} "
+            f"(peak={peak:.2f}, p10={p10:.2f}, n={len(samples)})",
+        )
+        if avg_accept is not None:
+            self.assertGreater(
+                avg_accept,
+                self.fwd_occupancy_acc_length_threshold,
+                f"avg_spec_accept_length={avg_accept:.3f} did not exceed "
+                f"threshold {self.fwd_occupancy_acc_length_threshold} -- spec "
+                "barely accepted, possibly degraded to vanilla decode",
+            )

python/sglang/test/kits/hellaswag_kit.py

@@ -0,0 +1,29 @@
+"""Hellaswag sanity kit.
+
+Runs hellaswag via the sgl frontend DSL bound to ``self.base_url`` and
+asserts accuracy above a threshold. Catches systematic regressions that
+pass every cheap single-prompt probe but tank multi-choice reasoning.
+
+Mix into a ``CustomTestCase`` subclass exposing ``self.base_url``.
+"""
+
+
+class HellaswagMixin:
+    """Assert hellaswag accuracy > threshold."""
+
+    hellaswag_accuracy_threshold: float = 0.60
+
+    def test_accuracy_floor(self):
+        import sglang as sgl
+        from sglang.test.test_programs import test_hellaswag_select
+
+        sgl.set_default_backend(sgl.RuntimeEndpoint(self.base_url))
+        try:
+            accuracy, _ = test_hellaswag_select()
+        finally:
+            sgl.set_default_backend(None)
+        self.assertGreater(
+            accuracy,
+            self.hellaswag_accuracy_threshold,
+            f"hellaswag accuracy floor breached: {accuracy:.3f}",
+        )

python/sglang/test/kits/json_constrained_kit.py

@@ -0,0 +1,102 @@
+import json
+from concurrent.futures import ThreadPoolExecutor
+
+import openai
+import requests
+
+
+class JSONConstrainedMixin:
+
+    json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string", "pattern": "^[\\w]+$"},
+                "population": {"type": "integer"},
+            },
+            "required": ["name", "population"],
+            "additionalProperties": False,
+        }
+    )
+
+    def _run_decode_json(
+        self, json_schema, return_logprob=False, top_logprobs_num=0, n=1
+    ):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": (
+                    "Introduce the capital of France. Return in a JSON format. The JSON Schema is: "
+                    + json.dumps(json_schema)
+                ),
+                "sampling_params": {
+                    "temperature": 0 if n == 1 else 0.5,
+                    "max_new_tokens": 128,
+                    "n": n,
+                    "stop_token_ids": [119690],
+                    "json_schema": json_schema,
+                },
+                "stream": False,
+                "return_logprob": return_logprob,
+                "top_logprobs_num": top_logprobs_num,
+                "logprob_start_len": 0,
+            },
+        )
+        ret = response.json()
+        print(json.dumps(ret))
+        print("=" * 100)
+
+        if not json_schema or json_schema == "INVALID":
+            return
+
+        # Make sure the json output is valid
+        try:
+            js_obj = json.loads(ret["text"])
+        except (TypeError, json.decoder.JSONDecodeError):
+            raise
+
+        self.assertIsInstance(js_obj["name"], str)
+        self.assertIsInstance(js_obj["population"], int)
+
+    def test_json_generate(self):
+        self._run_decode_json(json_schema=self.json_schema)
+
+    def test_json_invalid(self):
+        self._run_decode_json(json_schema="INVALID")
+
+    def test_json_openai(self):
+        client = openai.Client(api_key="EMPTY", base_url=f"{self.base_url}/v1")
+
+        response = client.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "system", "content": "You are a helpful AI assistant"},
+                {
+                    "role": "user",
+                    "content": "Introduce the capital of France. Return in a JSON format. "
+                    "The JSON Schema is: " + json.dumps(self.json_schema),
+                },
+            ],
+            temperature=0,
+            max_tokens=128,
+            response_format={
+                "type": "json_schema",
+                "json_schema": {"name": "foo", "schema": json.loads(self.json_schema)},
+            },
+        )
+        text = response.choices[0].message.content
+
+        try:
+            js_obj = json.loads(text)
+        except (TypeError, json.decoder.JSONDecodeError):
+            print("JSONDecodeError", text)
+            raise
+
+        self.assertIsInstance(js_obj["name"], str)
+        self.assertIsInstance(js_obj["population"], int)
+
+    def test_mix_json_and_other(self):
+        json_schemas = [None, None, self.json_schema, self.json_schema] * 10
+
+        with ThreadPoolExecutor(len(json_schemas)) as executor:
+            list(executor.map(self._run_decode_json, json_schemas))

python/sglang/test/kits/kl_divergence_kit.py

@@ -0,0 +1,42 @@
+from sglang.test.kl_test_utils import (
+    test_input_output_logprobs_match_decode_cache_hit_helper,
+    test_input_output_logprobs_match_prefill_cache_hit_helper,
+)
+
+
+class KLDivergenceMixin:
+    kl_div_thres: float
+    kl_div_thres_decode: float | None = None
+    kl_div_thres_prefill: float | None = None
+    kl_div_max_samples: int = 32
+    kl_div_prefill_max_new_tokens: int = 512
+    kl_div_decode_max_new_tokens: int = 512
+
+    @classmethod
+    def _build_acc_thresholds(cls, threshold):
+        """Build an ACC_THRESHOLDS dict compatible with kl_test_utils."""
+        return {cls.model: {"kl_div": threshold}}
+
+    @classmethod
+    def test_input_output_logprobs_match_prefill_cache_hit(cls):
+        test_input_output_logprobs_match_prefill_cache_hit_helper(
+            base_url=cls.base_url,
+            ACC_THRESHOLDS=cls._build_acc_thresholds(
+                cls.kl_div_thres_prefill or cls.kl_div_thres
+            ),
+            model_name=cls.model,
+            max_samples=cls.kl_div_max_samples,
+            max_new_tokens=cls.kl_div_prefill_max_new_tokens,
+        )
+
+    @classmethod
+    def test_input_output_logprobs_match_decode_cache_hit(cls):
+        test_input_output_logprobs_match_decode_cache_hit_helper(
+            base_url=cls.base_url,
+            ACC_THRESHOLDS=cls._build_acc_thresholds(
+                cls.kl_div_thres_decode or cls.kl_div_thres
+            ),
+            model_name=cls.model,
+            max_samples=cls.kl_div_max_samples,
+            max_new_tokens=cls.kl_div_decode_max_new_tokens,
+        )

python/sglang/test/kits/lm_eval_kit.py

@@ -0,0 +1,120 @@
+"""
+This module provides a mixin class for running lm-eval harness evaluations
+against SGLang servers
+"""
+
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import requests
+import yaml
+
+from sglang.test.test_utils import dump_metric
+
+
+@contextmanager
+def scoped_env_vars(new_env: dict[str, str] | None):
+    """Context manager to temporarily set environment variables."""
+    if not new_env:
+        yield
+        return
+
+    old_values = {}
+    new_keys = []
+
+    try:
+        for key, value in new_env.items():
+            if key in os.environ:
+                old_values[key] = os.environ[key]
+            else:
+                new_keys.append(key)
+            os.environ[key] = str(value)
+        yield
+    finally:
+        for key, value in old_values.items():
+            os.environ[key] = value
+        for key in new_keys:
+            os.environ.pop(key, None)
+
+
+class LMEvalMixin:
+    """
+    Mixin class for running lm-eval harness evaluations.
+    """
+
+    other_args: list[str] = []
+    model_config_name: str = ""
+    default_rtol: float = 0.08
+
+    def test_lm_eval(self):
+        """Run lm-eval evaluation and validate results."""
+        # Flush cache before evaluation
+        requests.get(self.base_url + "/flush_cache")
+
+        eval_config = yaml.safe_load(
+            Path(self.model_config_name).read_text(encoding="utf-8")
+        )
+        results = self.launch_lm_eval(eval_config)
+
+        rtol = eval_config.get("rtol", self.default_rtol)
+
+        success = True
+        for task in eval_config["tasks"]:
+            for metric in task["metrics"]:
+                ground_truth = metric["value"]
+                measured_value = results["results"][task["name"]][metric["name"]]
+                print(
+                    f"{task['name']} | {metric['name']}: "
+                    f"ground_truth={ground_truth:.3f} | "
+                    f"measured={measured_value:.3f} | rtol={rtol}"
+                )
+                dump_metric(
+                    f"{task['name']}_{metric['name']}",
+                    measured_value,
+                    labels={
+                        "model": eval_config.get("model_name", ""),
+                        "eval": "lm-eval",
+                        "task": task["name"],
+                    },
+                )
+                success = success and np.isclose(
+                    ground_truth, measured_value, rtol=rtol
+                )
+
+        self.assertTrue(success, f"lm-eval validation failed")
+
+    def launch_lm_eval(self, eval_config: dict[str, Any]) -> dict:
+        """
+        Args:
+            eval_config: Configuration dictionary with model and task settings
+        """
+        import lm_eval
+
+        batch_size = eval_config.get("batch_size", "auto")
+        backend = eval_config.get("backend", "local-completions")
+        num_concurrent = eval_config.get("num_concurrent", 1)
+
+        model_args = {
+            "model": eval_config["model_name"],
+            "base_url": self.base_url + "/v1/completions",
+            "num_concurrent": num_concurrent,
+        }
+
+        env_vars = eval_config.get("env_vars", None)
+        with scoped_env_vars(env_vars):
+            results = lm_eval.simple_evaluate(
+                model=backend,
+                model_args=model_args,
+                tasks=[task["name"] for task in eval_config["tasks"]],
+                num_fewshot=eval_config.get("num_fewshot", 0),
+                limit=eval_config.get("limit", None),
+                apply_chat_template=eval_config.get("apply_chat_template", False),
+                fewshot_as_multiturn=eval_config.get("fewshot_as_multiturn", False),
+                gen_kwargs=eval_config.get("gen_kwargs"),
+                batch_size=batch_size,
+            )
+
+        return results

python/sglang/test/kits/matched_stop_kit.py

@@ -0,0 +1,157 @@
+import json
+
+import requests
+
+MANY_NEW_TOKENS_PROMPT = """
+Please write an extremely detailed and vivid fantasy story, set in a world full of intricate magic systems, political intrigue, and complex characters.
+Ensure that you thoroughly describe every scene, character's motivations, and the environment. Include long, engaging dialogues and elaborate on the inner thoughts of the characters.
+Each section should be as comprehensive as possible to create a rich and immersive experience for the reader.
+The story should span multiple events, challenges, and character developments over time. Aim to make the story at least 3,000 words long.
+"""
+
+
+class MatchedStopMixin:
+    def _run_completions_generation(
+        self,
+        prompt=MANY_NEW_TOKENS_PROMPT,
+        max_tokens=1,
+        stop=None,
+        stop_regex=None,
+        finish_reason=None,
+        matched_stop=None,
+    ):
+        payload = {
+            "prompt": prompt,
+            "model": self.model,
+            "temperature": 0,
+            "top_p": 1,
+            "max_tokens": max_tokens,
+        }
+
+        if stop is not None:
+            payload["stop"] = stop
+
+        if stop_regex is not None:
+            payload["stop_regex"] = stop_regex
+
+        response_completions = requests.post(
+            self.base_url + "/v1/completions",
+            json=payload,
+        )
+        res = response_completions.json()
+        print(json.dumps(res))
+        print("=" * 100)
+
+        if not isinstance(matched_stop, list):
+            matched_stop = [matched_stop]
+
+        assert (
+            res["choices"][0]["finish_reason"] == finish_reason
+        ), f"Expected finish_reason: {finish_reason}, but got: {res['choices'][0]['finish_reason']}"
+        assert (
+            res["choices"][0]["matched_stop"] in matched_stop
+        ), f"Expected matched_stop: {matched_stop}, but got: {res['choices'][0]['matched_stop']}"
+
+    def _run_chat_completions_generation(
+        self,
+        prompt=MANY_NEW_TOKENS_PROMPT,
+        max_tokens=1,
+        stop=None,
+        stop_regex=None,
+        finish_reason=None,
+        matched_stop=None,
+    ):
+        chat_payload = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": "You are a helpful AI assistant"},
+                {"role": "user", "content": prompt},
+            ],
+            "temperature": 0,
+            "top_p": 1,
+            "max_tokens": max_tokens,
+        }
+
+        if stop is not None:
+            chat_payload["stop"] = stop
+
+        if stop_regex is not None:
+            chat_payload["stop_regex"] = stop_regex
+
+        response_chat = requests.post(
+            self.base_url + "/v1/chat/completions",
+            json=chat_payload,
+        )
+        res = response_chat.json()
+        print(json.dumps(res))
+        print("=" * 100)
+
+        if not isinstance(matched_stop, list):
+            matched_stop = [matched_stop]
+
+        assert (
+            res["choices"][0]["finish_reason"] == finish_reason
+        ), f"Expected finish_reason: {finish_reason}, but got: {res['choices'][0]['finish_reason']}"
+        assert (
+            res["choices"][0]["matched_stop"] in matched_stop
+        ), f"Expected matched_stop: {matched_stop}, but got: {res['choices'][0]['matched_stop']}"
+
+    def test_finish_stop_str(self):
+        self._run_completions_generation(
+            max_tokens=1000, stop="\n", finish_reason="stop", matched_stop="\n"
+        )
+        self._run_chat_completions_generation(
+            max_tokens=1000, stop="\n", finish_reason="stop", matched_stop="\n"
+        )
+
+    def test_finish_stop_regex_str(self):
+        STOP_REGEX_STR = r"and|or"
+        self._run_completions_generation(
+            max_tokens=1000,
+            stop_regex=STOP_REGEX_STR,
+            finish_reason="stop",
+            matched_stop=STOP_REGEX_STR,
+        )
+        self._run_chat_completions_generation(
+            max_tokens=1000,
+            stop_regex=STOP_REGEX_STR,
+            finish_reason="stop",
+            matched_stop=STOP_REGEX_STR,
+        )
+
+        # Match a complete sentence
+        STOP_REGEX_STR_SENTENCE = r"[.!?]\s*$"
+        self._run_chat_completions_generation(
+            max_tokens=1000,
+            stop_regex=STOP_REGEX_STR_SENTENCE,
+            finish_reason="stop",
+            matched_stop=STOP_REGEX_STR_SENTENCE,
+        )
+
+    def test_finish_stop_eos(self):
+        llama_format_prompt = """\
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>
+What is 2 + 2?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+        """
+        eos_token_ids = [128000, 128009, 2]
+        self._run_completions_generation(
+            prompt=llama_format_prompt,
+            max_tokens=1000,
+            finish_reason="stop",
+            matched_stop=eos_token_ids,
+        )
+        self._run_chat_completions_generation(
+            prompt="What is 2 + 2?",
+            max_tokens=1000,
+            finish_reason="stop",
+            matched_stop=eos_token_ids,
+        )
+
+    def test_finish_length(self):
+        self._run_completions_generation(
+            max_tokens=5, finish_reason="length", matched_stop=None
+        )
+        self._run_chat_completions_generation(
+            max_tokens=5, finish_reason="length", matched_stop=None
+        )

python/sglang/test/kits/mmmu_vlm_kit.py

@@ -0,0 +1,470 @@
+import glob
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.common import temp_set_env
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    dump_metric,
+    popen_launch_server,
+)
+
+# Set default mem_fraction_static to 0.8
+DEFAULT_MEM_FRACTION_STATIC = 0.8
+
+
+def _is_mmmu_parquet_corruption(error_output: str) -> bool:
+    """Check if error is due to MMMU parquet file corruption."""
+    return (
+        "ArrowInvalid" in error_output
+        and "Parquet magic bytes not found" in error_output
+        and ("MMMU" in error_output or "lmms-lab--MMMU" in error_output)
+    )
+
+
+def _cleanup_mmmu_dataset_cache():
+    """Clean up corrupted MMMU dataset cache to allow fresh download."""
+    # Priority 1: Check CI convention path /hf_home first (used in Docker containers)
+    ci_hf_home = Path("/hf_home/hub/datasets--lmms-lab--MMMU")
+    if ci_hf_home.exists():
+        mmmu_cache_path = ci_hf_home
+    else:
+        # Priority 2: Use HF_HOME env var or default user cache
+        hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+        mmmu_cache_path = Path(hf_home) / "hub" / "datasets--lmms-lab--MMMU"
+
+    if mmmu_cache_path.exists():
+        print(f"Detected corrupted MMMU parquet cache. Cleaning up: {mmmu_cache_path}")
+        try:
+            shutil.rmtree(mmmu_cache_path)
+            print(f"Successfully removed corrupted cache: {mmmu_cache_path}")
+            return True
+        except OSError as e:
+            print(f"Warning: Failed to remove cache {mmmu_cache_path}: {e}")
+            return False
+    else:
+        print(f"MMMU cache not found at {mmmu_cache_path}, skipping cleanup")
+        return False
+
+
+def _run_lmms_eval_with_retry(cmd: list[str], timeout: int = 3600) -> None:
+    """Run lmms_eval command with automatic retry on MMMU parquet corruption."""
+    try:
+        result = subprocess.run(
+            cmd,
+            check=True,
+            timeout=timeout,
+            capture_output=True,
+            text=True,
+        )
+        # Check for errors in output even if exit code is 0
+        # lmms_eval sometimes returns 0 even when errors occur
+        combined_output = result.stdout + result.stderr
+        if _is_mmmu_parquet_corruption(combined_output):
+            print(
+                "Detected MMMU parquet corruption error in output. Attempting recovery..."
+            )
+            if _cleanup_mmmu_dataset_cache():
+                print("Retrying lmms_eval with fresh download...")
+                with temp_set_env(
+                    HF_HUB_OFFLINE="0",
+                    HF_DATASETS_DOWNLOAD_MODE="force_redownload",
+                ):
+                    retry_result = subprocess.run(
+                        cmd, check=True, timeout=timeout, capture_output=True, text=True
+                    )
+                    # Print retry output
+                    if retry_result.stdout:
+                        print(retry_result.stdout, end="")
+                    if retry_result.stderr:
+                        print(retry_result.stderr, end="")
+            else:
+                print(
+                    f"Failed to cleanup corrupted MMMU cache. Output from lmms_eval:\nStdout:\n{result.stdout}\nStderr:\n{result.stderr}"
+                )
+                raise RuntimeError("Failed to cleanup corrupted MMMU cache")
+        else:
+            # Print captured output to maintain visibility of successful runs
+            if result.stdout:
+                print(result.stdout, end="")
+            if result.stderr:
+                print(result.stderr, end="")
+    except subprocess.CalledProcessError as e:
+        error_output = e.stderr + e.stdout
+        if _is_mmmu_parquet_corruption(error_output):
+            print("Detected MMMU parquet corruption error. Attempting recovery...")
+            if _cleanup_mmmu_dataset_cache():
+                print("Retrying lmms_eval with fresh download...")
+                with temp_set_env(
+                    HF_HUB_OFFLINE="0",
+                    HF_DATASETS_DOWNLOAD_MODE="force_redownload",
+                ):
+                    retry_result = subprocess.run(
+                        cmd, check=True, timeout=timeout, capture_output=True, text=True
+                    )
+                    # Print retry output
+                    if retry_result.stdout:
+                        print(retry_result.stdout, end="")
+                    if retry_result.stderr:
+                        print(retry_result.stderr, end="")
+            else:
+                print(
+                    f"Failed to cleanup corrupted MMMU cache. Error from lmms_eval:\nStdout:\n{e.stdout}\nStderr:\n{e.stderr}"
+                )
+                raise
+        else:
+            print(
+                f"lmms_eval failed with an unhandled error.\nStdout:\n{e.stdout}\nStderr:\n{e.stderr}"
+            )
+            raise
+
+
+class MMMUMixin:
+    """Mixin for MMMU evaluation.
+
+    Use with MMMUServerBase for single-model tests:
+        class TestMyModel(MMMUMixin, MMMUServerBase):
+            model = "my/model"
+            accuracy = 0.4
+    """
+
+    accuracy: float
+    mmmu_args: list[str] = []
+
+    # For OpenAI API settings
+    api_key = "sk-123456"
+
+    def run_mmmu_eval(
+        self: CustomTestCase,
+        model_version: str,
+        output_path: str,
+    ):
+        """
+        Evaluate a VLM on the MMMU validation set with lmms-eval.
+        Only `model_version` (checkpoint) and `chat_template` vary;
+        We are focusing only on the validation set due to resource constraints.
+        """
+        # -------- fixed settings --------
+        model = "openai_compatible"
+        tp = 1
+        tasks = "mmmu_val"
+        batch_size = 64
+        log_suffix = "openai_compatible"
+        os.makedirs(output_path, exist_ok=True)
+
+        # -------- compose --model_args --------
+        model_args = f'model_version="{model_version}",' f"tp={tp}"
+
+        # -------- build command list --------
+        cmd = [
+            "python3",
+            "-m",
+            "lmms_eval",
+            "--model",
+            model,
+            "--model_args",
+            model_args,
+            "--tasks",
+            tasks,
+            "--batch_size",
+            str(batch_size),
+            "--log_samples",
+            "--log_samples_suffix",
+            log_suffix,
+            "--output_path",
+            str(output_path),
+            *self.mmmu_args,
+        ]
+
+        # Set OpenAI API key and base URL environment variables.
+        # Needed for lmms-eval to work.
+        with temp_set_env(
+            OPENAI_API_KEY=self.api_key,
+            OPENAI_API_BASE=f"{self.base_url}/v1",
+        ):
+            _run_lmms_eval_with_retry(cmd)
+
+    def test_mmmu(self: CustomTestCase):
+        """Run MMMU evaluation test."""
+        with tempfile.TemporaryDirectory() as output_path:
+            # Run evaluation
+            self.run_mmmu_eval(self.model, output_path)
+
+            result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
+
+            if not result_files:
+                raise FileNotFoundError(f"No JSON result files found in {output_path}")
+
+            result_file_path = result_files[0]
+
+            with open(result_file_path, "r") as f:
+                result = json.load(f)
+                print(f"Result: {result}")
+
+            # Process the result
+            mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
+            print(f"Model {self.model} achieved accuracy: {mmmu_accuracy:.4f}")
+
+            dump_metric(
+                "mmmu_score",
+                mmmu_accuracy,
+                labels={"model": self.model, "eval": "mmmu", "api": "lmms-eval"},
+            )
+
+            # Assert performance meets expected threshold
+            self.assertGreaterEqual(
+                mmmu_accuracy,
+                self.accuracy,
+                f"Model {self.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({self.accuracy:.4f})",
+            )
+
+
+class MMMUMultiModelTestBase(CustomTestCase):
+    """Base class for multi-model MMMU tests.
+
+    This class is for tests that need to evaluate multiple models,
+    starting and stopping a server for each model within the test method.
+    For single-model tests, use MMMUMixin with MMMUServerBase instead.
+    """
+
+    parsed_args = None  # Class variable to store args
+    other_args = []
+    mmmu_args = []
+
+    @classmethod
+    def setUpClass(cls):
+        # Removed argument parsing from here
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+
+        if cls.parsed_args is None:
+            cls.parsed_args = SimpleNamespace(
+                mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
+            )
+
+        # Save original environment variables for restoration in tearDownClass
+        cls._original_openai_api_key = os.environ.get("OPENAI_API_KEY")
+        cls._original_openai_api_base = os.environ.get("OPENAI_API_BASE")
+
+        # Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
+        os.environ["OPENAI_API_KEY"] = cls.api_key
+        os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
+
+    @classmethod
+    def tearDownClass(cls):
+        # Restore original environment variables
+        if cls._original_openai_api_key is not None:
+            os.environ["OPENAI_API_KEY"] = cls._original_openai_api_key
+        elif "OPENAI_API_KEY" in os.environ:
+            del os.environ["OPENAI_API_KEY"]
+
+        if cls._original_openai_api_base is not None:
+            os.environ["OPENAI_API_BASE"] = cls._original_openai_api_base
+        elif "OPENAI_API_BASE" in os.environ:
+            del os.environ["OPENAI_API_BASE"]
+
+    def run_mmmu_eval(
+        self,
+        model_version: str,
+        output_path: str,
+        *,
+        env: dict | None = None,
+    ):
+        """
+        Evaluate a VLM on the MMMU validation set with lmms-eval.
+        Only `model_version` (checkpoint) and `chat_template` vary;
+        We are focusing only on the validation set due to resource constraints.
+        """
+        # -------- fixed settings --------
+        model = "openai_compatible"
+        tp = 1
+        tasks = "mmmu_val"
+        batch_size = 64
+        log_suffix = "openai_compatible"
+        os.makedirs(output_path, exist_ok=True)
+
+        # -------- compose --model_args --------
+        model_args = f'model_version="{model_version}",' f"tp={tp}"
+
+        # -------- build command list --------
+        cmd = [
+            "python3",
+            "-m",
+            "lmms_eval",
+            "--model",
+            model,
+            "--model_args",
+            model_args,
+            "--tasks",
+            tasks,
+            "--batch_size",
+            str(batch_size),
+            "--log_samples",
+            "--log_samples_suffix",
+            log_suffix,
+            "--output_path",
+            str(output_path),
+            *self.mmmu_args,
+        ]
+
+        _run_lmms_eval_with_retry(cmd)
+
+    def _run_vlm_mmmu_test(
+        self,
+        model,
+        output_path,
+        test_name="",
+        custom_env=None,
+        log_level="info",
+        capture_output=False,
+    ):
+        """
+        Common method to run VLM MMMU benchmark test.
+
+        Args:
+            model: Model to test
+            output_path: Path for output logs
+            test_name: Optional test name for logging
+            custom_env: Optional custom environment variables
+            log_level: Log level for server (default: "info")
+            capture_output: Whether to capture server stdout/stderr
+        """
+        print(f"\nTesting model: {model.model}{test_name}")
+
+        process = None
+        mmmu_accuracy = 0  # Initialize to handle potential exceptions
+        server_output = ""
+
+        try:
+            # Prepare environment variables
+            process_env = os.environ.copy()
+            if custom_env:
+                process_env.update(custom_env)
+            # if test vlm with cuda_ipc feature, open this env_var
+            process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
+
+            # Prepare stdout/stderr redirection if needed
+            stdout_file = None
+            stderr_file = None
+            if capture_output:
+                stdout_file = open("/tmp/server_stdout.log", "w")
+                stderr_file = open("/tmp/server_stderr.log", "w")
+
+            # Launch server for testing
+            process = popen_launch_server(
+                model.model,
+                base_url=self.base_url,
+                timeout=self.time_out,
+                api_key=self.api_key,
+                other_args=[
+                    "--trust-remote-code",
+                    "--cuda-graph-max-bs-decode",
+                    "64",
+                    "--enable-multimodal",
+                    "--mem-fraction-static",
+                    str(self.parsed_args.mem_fraction_static),  # Use class variable
+                    "--log-level",
+                    log_level,
+                    *self.other_args,
+                ],
+                env=process_env,
+                return_stdout_stderr=(
+                    (stdout_file, stderr_file) if capture_output else None
+                ),
+            )
+
+            # Run evaluation
+            self.run_mmmu_eval(model.model, output_path)
+
+            result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
+
+            if not result_files:
+                raise FileNotFoundError(f"No JSON result files found in {output_path}")
+
+            result_file_path = result_files[0]
+
+            with open(result_file_path, "r") as f:
+                result = json.load(f)
+                print(f"Result{test_name}\n: {result}")
+
+            # Process the result
+            mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
+            print(
+                f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
+            )
+
+            dump_metric(
+                "mmmu_score",
+                mmmu_accuracy,
+                labels={"model": model.model, "eval": "mmmu", "api": "lmms-eval"},
+            )
+
+            # Capture server output if requested
+            if capture_output and process:
+                server_output = self._read_output_from_files()
+
+            # Assert performance meets expected threshold
+            self.assertGreaterEqual(
+                mmmu_accuracy,
+                model.mmmu_accuracy,
+                f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
+            )
+
+            return server_output
+
+        except Exception as e:
+            print(f"Error testing {model.model}{test_name}: {e}")
+            self.fail(f"Test failed for {model.model}{test_name}: {e}")
+
+        finally:
+            # Ensure process cleanup happens regardless of success/failure
+            if process is not None and process.poll() is None:
+                print(f"Cleaning up process {process.pid}")
+                try:
+                    kill_process_tree(process.pid)
+                except Exception as e:
+                    print(f"Error killing process: {e}")
+
+            # clean up temporary files
+            if capture_output:
+                if stdout_file:
+                    stdout_file.close()
+                if stderr_file:
+                    stderr_file.close()
+                for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
+                    try:
+                        if os.path.exists(filename):
+                            os.remove(filename)
+                    except Exception as e:
+                        print(f"Error removing {filename}: {e}")
+
+    def _read_output_from_files(self):
+        output_lines = []
+
+        log_files = [
+            ("/tmp/server_stdout.log", "[STDOUT]"),
+            ("/tmp/server_stderr.log", "[STDERR]"),
+        ]
+        for filename, tag in log_files:
+            try:
+                if os.path.exists(filename):
+                    with open(filename, "r") as f:
+                        for line in f:
+                            output_lines.append(f"{tag} {line.rstrip()}")
+            except Exception as e:
+                print(f"Error reading {tag.lower()} file: {e}")
+
+        return "\n".join(output_lines)
+
+
+# Backward compatibility alias
+MMMUVLMTestBase = MMMUMultiModelTestBase

python/sglang/test/kits/pause_generation_kit.py

@@ -0,0 +1,111 @@
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import requests
+
+_REQUEST_TIMEOUT = 180
+
+
+class PauseResumeInPlaceMixin:
+    """Test pause/resume with in_place mode.
+
+    Sends concurrent requests, pauses mid-generation, verifies no progress
+    during the pause window, then resumes and verifies all requests complete.
+
+    Subclass must set:
+      - pause_generate_url: URL to send /generate requests (or falls back to self.base_url)
+      - pause_target_urls: list of URLs to send /pause_generation and /continue_generation
+    """
+
+    pause_num_requests: int = 32
+    pause_max_new_tokens: int = 512
+    pause_duration: float = 5
+    pause_generate_url: str = ""
+    pause_target_urls: list = []
+
+    def test_pause_resume_in_place(self):
+        generate_url = self.pause_generate_url or self.base_url
+        target_urls = self.pause_target_urls or [self.base_url]
+        num_requests = self.pause_num_requests
+
+        def _generate(prompt_id):
+            return requests.post(
+                generate_url + "/generate",
+                json={
+                    "text": f"Question {prompt_id}: Write a short essay about the number {prompt_id}.",
+                    "sampling_params": {
+                        "temperature": 0.8,
+                        "max_new_tokens": self.pause_max_new_tokens,
+                    },
+                },
+                timeout=_REQUEST_TIMEOUT,
+            )
+
+        with ThreadPoolExecutor(max_workers=num_requests) as executor:
+            futures = {executor.submit(_generate, i): i for i in range(num_requests)}
+
+            time.sleep(1)
+
+            # Pause all targets
+            for url in target_urls:
+                requests.post(
+                    url + "/pause_generation",
+                    json={"mode": "in_place"},
+                    timeout=30,
+                ).raise_for_status()
+
+            time.sleep(0.5)
+            done_before = sum(1 for f in futures if f.done())
+
+            time.sleep(self.pause_duration)
+            done_after = sum(1 for f in futures if f.done())
+
+            self.assertLess(
+                done_before,
+                num_requests,
+                "All requests completed before pause took effect -- "
+                "increase pause_max_new_tokens to make the test meaningful.",
+            )
+
+            self.assertEqual(
+                done_after - done_before,
+                0,
+                f"{done_after - done_before} requests completed during pause "
+                f"({done_before} before, {done_after} after) -- "
+                f"pause_generation was not respected by the scheduler.",
+            )
+
+            # Resume all targets (reverse order to unblock downstream first)
+            for url in reversed(target_urls):
+                requests.post(
+                    url + "/continue_generation",
+                    json={},
+                    timeout=30,
+                ).raise_for_status()
+
+            completed = 0
+            errors = []
+            for future in as_completed(futures, timeout=_REQUEST_TIMEOUT):
+                prompt_id = futures[future]
+                try:
+                    resp = future.result()
+                    if resp.status_code == 200:
+                        body = resp.json()
+                        self.assertIn("text", body)
+                        self.assertGreater(len(body["text"]), 0)
+                        completed += 1
+                    else:
+                        errors.append(f"Request {prompt_id}: status={resp.status_code}")
+                except Exception as e:
+                    errors.append(f"Request {prompt_id}: exception={e}")
+
+        self.assertEqual(
+            completed + len(errors),
+            num_requests,
+            "Some requests did not resolve within the timeout -- likely hung during pause.",
+        )
+        self.assertEqual(
+            completed,
+            num_requests,
+            f"Some requests failed: {completed}/{num_requests} succeeded. Errors: {errors}",
+        )

python/sglang/test/kits/prefix_cache_branching_kit.py

@@ -0,0 +1,50 @@
+import requests
+
+
+class PrefixCacheBranchingMixin:
+    cache_chunk_size: int
+
+    @classmethod
+    def send_request_helper(cls, text: str):
+        response = requests.post(
+            cls.base_url + "/generate",
+            json={
+                "text": text,
+                "sampling_params": {
+                    "max_new_tokens": 1,
+                },
+            },
+        )
+        return response.json()
+
+    @classmethod
+    def test_prefix_cache_branching(cls):
+        cls.flush_cache()
+        branching_pos = 257
+        text_prefix = "hi" * branching_pos
+        suffix_list = [
+            "this" * cls.cache_chunk_size * 4,
+            "here" * cls.cache_chunk_size * 4,
+            "that" * cls.cache_chunk_size * 4,
+        ]
+        cache_hit_list = [False, False, True]
+
+        # First request only prefill the entire sequence
+        # Second request won't have cache hit, but will cache the branching point
+        # Third request will have cache hit on the branching point
+        for i, (suffix, cache_hit) in enumerate(
+            zip(suffix_list, cache_hit_list, strict=True)
+        ):
+            result = cls.send_request_helper(text_prefix + suffix)
+            cached_tokens = result["meta_info"]["cached_tokens"]
+            if cache_hit:
+                expected_cached_tokens = (
+                    branching_pos // cls.cache_chunk_size * cls.cache_chunk_size
+                )
+                assert (
+                    cached_tokens == expected_cached_tokens
+                ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}"
+            else:
+                assert (
+                    cached_tokens == 0
+                ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0"

python/sglang/test/kits/radix_cache_server_kit.py

@@ -0,0 +1,54 @@
+import random
+
+import requests
+
+
+def gen_radix_tree(num_nodes=400, chunk_len=256):
+    num0 = num_nodes // 2
+    num1 = num_nodes - num0
+    nodes = [{"input_ids": [37] * 117, "decode_len": 217}]
+    for _ in range(num0):
+        parent = random.choice(nodes)
+        unique_len = random.randint(0, chunk_len)
+        decode_len = random.randint(0, chunk_len)
+        token_id = random.randint(
+            0, 31999
+        )  # randint is inclusive; vocab_size-1 = 31999
+        child = {
+            "input_ids": parent["input_ids"] + [token_id] * unique_len,
+            "decode_len": decode_len,
+        }
+        nodes.append(child)
+
+    while num1 > 0:
+        num_branch = random.randint(1, min(num1, 10))
+        parent = random.choice(nodes)
+        for _ in range(num_branch):
+            unique_len = random.randint(0, chunk_len)
+            decode_len = random.randint(0, chunk_len)
+            token_id = random.randint(
+                0, 31999
+            )  # randint is inclusive; vocab_size-1 = 31999
+            child = {
+                "input_ids": parent["input_ids"] + [token_id] * unique_len,
+                "decode_len": decode_len,
+            }
+            nodes.append(child)
+
+        num1 -= num_branch
+
+    random.shuffle(nodes)
+    return nodes
+
+
+def run_radix_attention_test(base_url: str):
+    nodes = gen_radix_tree()
+    data = {
+        "input_ids": [node["input_ids"] for node in nodes],
+        "sampling_params": [
+            {"max_new_tokens": node["decode_len"], "temperature": 0} for node in nodes
+        ],
+    }
+
+    res = requests.post(base_url + "/generate", json=data)
+    assert res.status_code == 200

python/sglang/test/kits/reasoning_kit.py

@@ -0,0 +1,197 @@
+import json
+
+import openai
+import requests
+
+from sglang.srt.parser.reasoning_parser import ReasoningParser
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+
+
+class ReasoningTokenUsageMixin:
+    """Mixin for reasoning_tokens usage tests.
+
+    Required attributes on the test class:
+        model: str
+        base_url: str
+        reasoning_parser_name: str
+
+    Optional attributes:
+        api_key: str (if not set, no auth)
+
+    Call cls.init_reasoning_token_verifier() in setUpClass.
+    """
+
+    reasoning_parser_name = None
+
+    @classmethod
+    def init_reasoning_token_verifier(cls):
+        assert cls.reasoning_parser_name, "reasoning_parser_name must be set"
+        cls.tokenizer = get_tokenizer(cls.model)
+        parser = ReasoningParser(cls.reasoning_parser_name, tokenizer=cls.tokenizer)
+        cls.think_end_token_id = cls.tokenizer.convert_tokens_to_ids(
+            parser.detector.think_end_token
+        )
+        assert cls.think_end_token_id is not None
+
+    def _reasoning_chat_request(self, enable_thinking, stream=False):
+        api_key = getattr(self, "api_key", None)
+        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+        payload = {
+            "model": self.model,
+            "messages": [{"role": "user", "content": "What is 1+3?"}],
+            "max_tokens": 1024,
+            "chat_template_kwargs": {"enable_thinking": enable_thinking},
+        }
+        if stream:
+            payload["stream"] = True
+            payload["stream_options"] = {"include_usage": True}
+        return requests.post(
+            f"{self.base_url}/v1/chat/completions",
+            headers=headers,
+            json=payload,
+            stream=stream,
+        )
+
+    def _extract_streaming_usage(self, response):
+        usage = None
+        for line in response.iter_lines():
+            if not line:
+                continue
+            decoded = line.decode("utf-8")
+            if not decoded.startswith("data:") or decoded.startswith("data: [DONE]"):
+                continue
+            data = json.loads(decoded[len("data:") :].strip())
+            if data.get("usage"):
+                usage = data["usage"]
+        return usage
+
+    def test_reasoning_tokens_thinking(self):
+        resp = self._reasoning_chat_request(enable_thinking=True)
+        self.assertEqual(resp.status_code, 200, resp.text)
+        usage = resp.json()["usage"]
+        self.assertGreater(usage["reasoning_tokens"], 0)
+        self.assertLess(usage["reasoning_tokens"], usage["completion_tokens"])
+
+    def test_reasoning_tokens_non_thinking(self):
+        resp = self._reasoning_chat_request(enable_thinking=False)
+        self.assertEqual(resp.status_code, 200, resp.text)
+        self.assertEqual(resp.json()["usage"]["reasoning_tokens"], 0)
+
+    def test_reasoning_tokens_thinking_stream(self):
+        with self._reasoning_chat_request(enable_thinking=True, stream=True) as resp:
+            self.assertEqual(resp.status_code, 200, resp.text)
+            usage = self._extract_streaming_usage(resp)
+            self.assertIsNotNone(usage, "No usage in stream")
+            self.assertGreater(usage["reasoning_tokens"], 0)
+            self.assertLess(usage["reasoning_tokens"], usage["completion_tokens"])
+
+    def test_reasoning_tokens_non_thinking_stream(self):
+        with self._reasoning_chat_request(enable_thinking=False, stream=True) as resp:
+            self.assertEqual(resp.status_code, 200, resp.text)
+            usage = self._extract_streaming_usage(resp)
+            self.assertIsNotNone(usage, "No usage in stream")
+            self.assertEqual(usage["reasoning_tokens"], 0)
+
+    def test_reasoning_tokens_generate_exact_count(self):
+        api_key = getattr(self, "api_key", None)
+        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+        messages = [{"role": "user", "content": "What is 1+3?"}]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True, tokenize=False
+        )
+        resp = requests.post(
+            f"{self.base_url}/generate",
+            headers=headers,
+            json={
+                "text": prompt,
+                "sampling_params": {"max_new_tokens": 1024},
+                "require_reasoning": True,
+            },
+        )
+        self.assertEqual(resp.status_code, 200, resp.text)
+        data = resp.json()
+        reported = data["meta_info"]["reasoning_tokens"]
+        actual = data["output_ids"].index(self.think_end_token_id) + 1
+        self.assertEqual(reported, actual)
+
+
+class SeparateReasoningMixin:
+    """Mixin for separate_reasoning tests.
+
+    Required attributes on the test class:
+        model: str
+        base_url: str (without /v1)
+        api_key: str
+    """
+
+    def _openai_client(self):
+        return openai.Client(api_key=self.api_key, base_url=f"{self.base_url}/v1")
+
+    def _chat(self, stream=False, extra_body=None):
+        return self._openai_client().chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": "What is 1+3?"}],
+            max_tokens=1024,
+            stream=stream,
+            extra_body=extra_body,
+        )
+
+    def _collect_stream(self, response):
+        reasoning_content = ""
+        content = ""
+        for chunk in response:
+            if chunk.choices[0].delta.content:
+                content += chunk.choices[0].delta.content
+            elif chunk.choices[0].delta.reasoning_content:
+                reasoning_content += chunk.choices[0].delta.reasoning_content
+        return reasoning_content, content
+
+    def test_streaming_separate_reasoning_false(self):
+        response = self._chat(stream=True, extra_body={"separate_reasoning": False})
+        reasoning_content, content = self._collect_stream(response)
+        self.assertEqual(len(reasoning_content), 0)
+        self.assertGreater(len(content), 0)
+
+    def test_streaming_separate_reasoning_true(self):
+        response = self._chat(stream=True, extra_body={"separate_reasoning": True})
+        reasoning_content, content = self._collect_stream(response)
+        self.assertGreater(len(reasoning_content), 0)
+        self.assertGreater(len(content), 0)
+
+    def test_streaming_separate_reasoning_true_stream_reasoning_false(self):
+        response = self._chat(
+            stream=True,
+            extra_body={"separate_reasoning": True, "stream_reasoning": False},
+        )
+        reasoning_content = ""
+        content = ""
+        first_chunk = False
+        for chunk in response:
+            if chunk.choices[0].delta.reasoning_content:
+                reasoning_content = chunk.choices[0].delta.reasoning_content
+                first_chunk = True
+            if chunk.choices[0].delta.content:
+                content += chunk.choices[0].delta.content
+                if not first_chunk:
+                    reasoning_content = chunk.choices[0].delta.reasoning_content
+                first_chunk = True
+            if not first_chunk:
+                assert (
+                    not chunk.choices[0].delta.reasoning_content
+                    or len(chunk.choices[0].delta.reasoning_content) == 0
+                )
+        self.assertGreater(len(reasoning_content), 0)
+        self.assertGreater(len(content), 0)
+
+    def test_nonstreaming_separate_reasoning_false(self):
+        response = self._chat(extra_body={"separate_reasoning": False})
+        assert (
+            not response.choices[0].message.reasoning_content
+            or len(response.choices[0].message.reasoning_content) == 0
+        )
+        self.assertGreater(len(response.choices[0].message.content), 0)
+
+    def test_nonstreaming_separate_reasoning_true(self):
+        response = self._chat(extra_body={"separate_reasoning": True})
+        self.assertGreater(len(response.choices[0].message.reasoning_content), 0)
+        self.assertGreater(len(response.choices[0].message.content), 0)

python/sglang/test/kits/regex_constrained_kit.py

@@ -0,0 +1,131 @@
+import json
+
+import requests
+
+
+class RegexConstrainedMixin:
+
+    def _run_decode_regex(
+        self,
+        regex,
+        prompt,
+        return_logprob=False,
+        top_logprobs_num=0,
+        n=1,
+    ):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompt,
+                "sampling_params": {
+                    "temperature": 0 if n == 1 else 0.5,
+                    "max_new_tokens": 128,
+                    "n": n,
+                    "regex": regex,
+                },
+                "stream": False,
+                "return_logprob": return_logprob,
+                "top_logprobs_num": top_logprobs_num,
+                "logprob_start_len": 0,
+            },
+        )
+
+        ret = response.json()
+        print(json.dumps(ret, indent=2))
+        print("=" * 100)
+
+        if not isinstance(ret, list):
+            self.fail(f"Expected response to be a list, but got {type(ret)}")
+
+        for item in ret:
+            text = item.get("text", "").strip()
+            if not text:
+                self.fail("Generated text is empty.")
+
+            if not self.regex_match(text, regex):
+                self.fail(f"Text '{text}' does not match regex pattern.")
+
+    def regex_match(self, text, pattern):
+        import re
+
+        return re.match(pattern, text) is not None
+
+    def test_regex_generate_email(self):
+        pattern = r"^user@example\.com$"
+        prompt = "Generate an email address:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_regex_generate_greeting(self):
+        pattern = r"^(Hello|Hi|Hey)$"
+        prompt = "Generate a greeting:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_regex_generate_number(self):
+        pattern = r"^\d{3}$"
+        prompt = "Generate a three-digit number:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_regex_generate_phone(self):
+        pattern = r"^\(\d{3}\) \d{3}-\d{4}$"
+        prompt = "Generate a phone number:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_regex_generate_date(self):
+        pattern = r"^2024-(0[1-9]|1[0-2])-(0[1-9]|[12]\d|3[01])$"
+        prompt = "Generate a date in YYYY-MM-DD format:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_regex_generate_hex_color(self):
+        pattern = r"^#[0-9A-F]{6}$"
+        prompt = "Generate a hex color code:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_regex_generate_complex_json(self):
+        pattern = r'^\{\s*"name"\s*:\s*"[a-zA-Z0-9 ]+"\s*,\s*"age"\s*:\s*[1-9][0-9]*\s*,\s*"city"\s*:\s*"[a-zA-Z0-9 ]+"\s*\}$'
+        prompt = "Generate a simple JSON with name, age, and city:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )
+
+    def test_regex_generate_custom_log_format(self):
+        pattern = r"^\[2024-01-01T12:00:00Z\] INFO: System\.process - Operation [a-z]+ successfully$"
+        prompt = "Generate a log entry:"
+
+        self._run_decode_regex(
+            regex=pattern,
+            prompt=prompt,
+            n=3,
+        )

python/sglang/test/kits/spec_decoding_kit.py

@@ -0,0 +1,32 @@
+import requests
+
+from sglang.test.send_one import BenchArgs, send_one_prompt
+from sglang.test.test_utils import is_in_ci, write_github_step_summary
+
+
+class SpecDecodingMixin:
+    bs_1_speed_thres: float
+    accept_length_thres: float
+    bs_1_speed_attempts: int = 3
+
+    def test_bs_1_speed(self):
+        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        acc_length, speed = 0.0, 0.0
+        for attempt in range(1, self.bs_1_speed_attempts + 1):
+            requests.get(self.base_url + "/flush_cache")
+            acc_length, speed = send_one_prompt(
+                args, label=f"attempt {attempt}", print_output=False
+            )
+            if acc_length > self.accept_length_thres and speed > self.bs_1_speed_thres:
+                break
+        requests.get(self.base_url + "/flush_cache")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_bs_1_speed ({self.model})\n"
+                f"{acc_length=:.2f}\n"
+                f"{speed=:.2f} token/s\n"
+            )
+
+        self.assertGreater(acc_length, self.accept_length_thres)
+        self.assertGreater(speed, self.bs_1_speed_thres)

python/sglang/test/kits/spec_server_kits.py

@@ -0,0 +1,693 @@
+"""Reusable test-method mixins (kits) for EAGLE/EAGLE3 spec-decoding servers.
+
+Pair these with ``SpecEagleServerBase`` (sglang.test.server_fixtures.spec_eagle_fixture).
+Each kit is a cohesive group of ``test_*`` methods with no launch logic; concrete
+test classes mix in the fixture (which owns launch knobs) + whichever kits apply.
+
+Thresholds are read off ``self`` so a config can tune them as class attributes.
+"""
+
+import concurrent.futures
+import json
+import random
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from types import SimpleNamespace
+
+import numpy as np
+import requests
+
+from sglang.srt.utils.common import kill_process_tree
+from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    popen_launch_server,
+    run_logprob_check,
+)
+
+
+class SpecCorrectnessKit:
+    """Acceptance-quality + EOS checks (single server, cheap)."""
+
+    # Tunable thresholds (override per config class).
+    acc_length_thres = 3.1
+    batch_accept_len_thres = 1.75
+
+    def test_acc_length(self):
+        prompt = [
+            "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:",
+        ] * 5
+        sampling_params = {"temperature": 0, "max_new_tokens": 512}
+        output = requests.post(
+            self.base_url + "/generate",
+            json={"text": prompt, "sampling_params": sampling_params},
+        ).json()[0]
+
+        meta = output["meta_info"]
+        if "spec_verify_ct" in meta and meta["spec_verify_ct"] > 0:
+            acc_length = meta["completion_tokens"] / meta["spec_verify_ct"]
+        else:
+            acc_length = 1.0
+        print(f"{acc_length=:.4f}")
+        self.assertGreater(acc_length, self.acc_length_thres)
+
+    def test_batch_generation(self):
+        prompts = [
+            "Hello, my name is",
+            "The president of the United States is",
+            "The capital of France is",
+            "The future of AI is",
+        ]
+        results = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompts,
+                "sampling_params": {"temperature": 0, "max_new_tokens": 50},
+            },
+        ).json()
+        # Accept length from per-request meta_info (self-contained). The
+        # internal_states `avg_spec_accept_length` isn't populated on the v1 /
+        # disable-overlap path after a small batch, so don't read server_info.
+        total_completion, total_verify = 0, 0
+        for r in results:
+            self.assertIn("text", r, f"Server error: {r}")
+            meta = r["meta_info"]
+            total_completion += meta["completion_tokens"]
+            total_verify += meta.get("spec_verify_ct", 0)
+        if total_verify > 0:
+            acc_length = total_completion / total_verify
+            print(f"batch {acc_length=:.4f}")
+            self.assertGreater(acc_length, self.batch_accept_len_thres)
+
+    def test_eos_token(self):
+        prompt = "[INST] <<SYS>>\nYou are a helpful assistant.\n<</SYS>>\nToday is a sunny day and I like [/INST]"
+        res = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompt,
+                "sampling_params": {
+                    "temperature": 0.1,
+                    "max_new_tokens": 1024,
+                    "skip_special_tokens": False,
+                },
+            },
+        ).json()
+        output = res["text"]
+        tokens = self.tokenizer.encode(output, truncation=False)
+        self.assertNotIn(self.tokenizer.eos_token_id, tokens)
+
+    def test_first_token_finish(self):
+        # Very short max_new_tokens (1-3): exercise the immediate-finish path,
+        # where a request stops within the first draft window. Just must not crash.
+        prompts = [
+            f"There are {i} apples on the table. How to divide them equally?"
+            for i in range(8)
+        ]
+        sampling_params = [
+            {"temperature": 0, "max_new_tokens": random.randint(1, 3)} for _ in range(8)
+        ]
+        results = requests.post(
+            self.base_url + "/generate",
+            json={"text": prompts, "sampling_params": sampling_params},
+        ).json()
+        for r in results:
+            self.assertIn("text", r, f"Server error: {r}")
+
+
+def _greedy(url, text, max_new_tokens=48):
+    return requests.post(
+        url + "/generate",
+        json={
+            "text": text,
+            "sampling_params": {"temperature": 0, "max_new_tokens": max_new_tokens},
+        },
+    ).json()["text"]
+
+
+class SpecParityKit:
+    """Lossless output parity vs a non-spec reference.
+
+    Sequential (NOT concurrent): launch a non-spec reference server on the
+    standard port, capture greedy outputs, tear it down, THEN let the fixture
+    launch the spec server. Only one model is resident at a time -- two 8B
+    servers don't fit on one GPU. Mix this kit FIRST in the bases so its
+    setUpClass runs before the fixture's:  ``class T(SpecParityKit, Eagle3Base)``.
+    """
+
+    parity_prompts = [
+        "The capital of France is",
+        "Once upon a time, there was a",
+        "The three primary colors are",
+        "def fibonacci(n):",
+    ]
+
+    @classmethod
+    def setUpClass(cls):
+        ref_url = DEFAULT_URL_FOR_TEST
+        ref_proc = popen_launch_server(
+            cls.model,
+            ref_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--mem-fraction-static",
+                "0.8",  # ref alone -> full GPU available
+                "--attention-backend",
+                cls.attention_backend,
+                "--page-size",
+                "1",
+                "--dtype",
+                cls.dtype,
+                *(["--trust-remote-code"] if cls.trust_remote_code else []),
+            ],
+        )
+        try:
+            cls.parity_ref_outputs = {
+                p: _greedy(ref_url, p) for p in cls.parity_prompts
+            }
+        finally:
+            kill_process_tree(ref_proc.pid, wait_timeout=60)
+        # Now the spec server (same port; ref is gone).
+        super().setUpClass()
+
+    def test_parity_vs_reference(self):
+        """Spec decode greedy output must equal the non-spec reference."""
+        for prompt in self.parity_prompts:
+            spec_out = _greedy(self.base_url, prompt)
+            self.assertEqual(
+                spec_out,
+                self.parity_ref_outputs[prompt],
+                f"spec != ref for prompt {prompt!r}",
+            )
+
+
+class SpecAccuracyKit:
+    """gsm8k accuracy + acceptance length, and throughput at max_tokens=1."""
+
+    gsm8k_num_examples = 200
+    gsm8k_score_thres = 0.20
+    gsm8k_check_accept_len = True
+    # If set, use this; else fall back to topk-based default (2.5 / 3.47).
+    gsm8k_accept_len_thres = None
+
+    def test_gsm8k(self):
+        requests.get(self.base_url + "/flush_cache")
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=self.gsm8k_num_examples,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], self.gsm8k_score_thres)
+
+        if self.gsm8k_check_accept_len:
+            server_info = requests.get(self.base_url + "/server_info").json()
+            avg_spec_accept_length = server_info["internal_states"][0].get(
+                "avg_spec_accept_length"
+            )
+            print(f"{avg_spec_accept_length=}")
+            # The metric isn't always populated (e.g. v1 / disable-overlap).
+            # Only enforce the threshold when it's reported.
+            if avg_spec_accept_length is not None:
+                topk = server_info["speculative_eagle_topk"]
+                thres = self.gsm8k_accept_len_thres
+                if thres is None:
+                    thres = 2.5 if topk == 1 else 3.47
+                self.assertGreater(avg_spec_accept_length, thres)
+
+
+class SpecPerfKit:
+    """Throughput perf check (GPU-specific -> run on the reference/Hopper runner)."""
+
+    perf_output_throughput_thres = 50
+
+    def test_max_token_one(self):
+        requests.get(self.base_url + "/flush_cache")
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=1,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        self.assertGreater(
+            metrics["output_throughput"], self.perf_output_throughput_thres
+        )
+
+
+class SpecLogprobKit:
+    """Logprob correctness: start_len, prefill-rescore match, mixed sweep,
+    spec-v2 decode-vs-prefill match, and ragged token_ids_logprob."""
+
+    # Max |decode-path - prefill-rescore| logprob gap. The two paths run
+    # different kernels / batch shapes, so the gap is accumulated rounding
+    # noise of the fixture dtype: ~0.25 observed for bf16 (up to 0.36 on
+    # some CI runners), ~8x smaller for fp16 (3 extra mantissa bits).
+    logprob_match_delta = 0.5
+
+    def test_logprob_start_len(self):
+        logprob_start_len = 4
+        new_tokens = 4
+        prompts = [
+            "I have a very good idea on",
+            "Today is a sunndy day and",
+        ]
+
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompts,
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": new_tokens,
+                },
+                "return_logprob": True,
+                "top_logprobs_num": 5,
+                "logprob_start_len": logprob_start_len,
+            },
+        )
+        response_json = response.json()
+        for res in response_json:
+            self.assertEqual(
+                res["meta_info"]["prompt_tokens"],
+                logprob_start_len + len(res["meta_info"]["input_token_logprobs"]),
+            )
+            self.assertEqual(res["meta_info"]["completion_tokens"], new_tokens)
+            self.assertEqual(len(res["meta_info"]["output_token_logprobs"]), new_tokens)
+
+    def test_logprob_match(self):
+        """Output logprobs should match a fresh prefill of the same sequence."""
+
+        def run_generate(
+            prompt,
+            return_logprob=False,
+            max_new_tokens=512,
+            logprob_start_len=-1,
+            temperature=1.0,
+        ):
+            if isinstance(prompt, str):
+                prompt_kwargs = {"text": prompt}
+            else:
+                prompt_kwargs = {"input_ids": prompt}
+
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    **prompt_kwargs,
+                    "sampling_params": {
+                        "temperature": temperature,
+                        "max_new_tokens": max_new_tokens,
+                        "ignore_eos": True,
+                    },
+                    "return_logprob": return_logprob,
+                    "return_text_in_logprobs": True,
+                    "logprob_start_len": logprob_start_len,
+                },
+            )
+            return response.json()
+
+        prompt = "I have a very good idea on how to"
+
+        for temperature in [1.0]:
+            gen = run_generate(
+                prompt,
+                return_logprob=True,
+                logprob_start_len=0,
+                temperature=temperature,
+            )
+            output_logprobs = np.array(
+                [x[0] for x in gen["meta_info"]["output_token_logprobs"]]
+            )
+            num_prompts_tokens = gen["meta_info"]["prompt_tokens"]
+
+            input_tokens = [x[1] for x in gen["meta_info"]["input_token_logprobs"]]
+            output_tokens = [x[1] for x in gen["meta_info"]["output_token_logprobs"]]
+
+            new_prompt = input_tokens + output_tokens
+            score = run_generate(
+                new_prompt,
+                return_logprob=True,
+                logprob_start_len=0,
+                max_new_tokens=0,
+                temperature=temperature,
+            )
+            output_logprobs_score = np.array(
+                [
+                    x[0]
+                    for x in score["meta_info"]["input_token_logprobs"][
+                        num_prompts_tokens:
+                    ]
+                ]
+            )
+
+            diff = np.abs(output_logprobs - output_logprobs_score)
+            max_diff = np.max(diff)
+            self.assertLess(max_diff, self.logprob_match_delta)
+
+    def test_logprob_mixed(self):
+        args = []
+        temperature = 0
+        # input_len, output_len, temperature, logprob_start_len, return_logprob, top_logprobs_num
+        for input_len in [200, 500, 1000, 2000]:
+            for output_len in [4, 8]:
+                for logprob_start_len in [0, 100, 300, 800, 1998]:
+                    for return_logprob in [True, False]:
+                        for top_logprobs_num in [0, 5]:
+                            if logprob_start_len >= input_len:
+                                continue
+                            args.append(
+                                (
+                                    input_len,
+                                    output_len,
+                                    temperature,
+                                    logprob_start_len,
+                                    return_logprob,
+                                    top_logprobs_num,
+                                )
+                            )
+
+        random.shuffle(args)
+        func = partial(run_logprob_check, self)
+        with ThreadPoolExecutor(8) as executor:
+            list(executor.map(func, args))
+
+    def test_logprob_spec_v2_match(self):
+        """Verify spec v2 decode logprobs match prefill scoring logprobs."""
+        top_k = 5
+        probe_token_ids = [1, 2, 10, 100, 1000]
+        prompts = [
+            "The capital of France is",
+            "Explain quantum computing in simple terms:",
+        ]
+
+        for round_idx, prompt in enumerate(prompts):
+            with self.subTest(round=round_idx, prompt=prompt):
+                gen_res = requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "text": prompt,
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 32,
+                            "ignore_eos": True,
+                        },
+                        "return_logprob": True,
+                        "top_logprobs_num": top_k,
+                        "token_ids_logprob": probe_token_ids,
+                        "logprob_start_len": 0,
+                    },
+                ).json()
+
+                decode_logprobs = gen_res["meta_info"]["output_token_logprobs"]
+                decode_top_logprobs = gen_res["meta_info"]["output_top_logprobs"]
+                decode_tid_logprobs = gen_res["meta_info"]["output_token_ids_logprobs"]
+                input_token_ids = [
+                    t[1] for t in gen_res["meta_info"]["input_token_logprobs"]
+                ]
+                output_token_ids = [t[1] for t in decode_logprobs]
+                num_prompt_tokens = gen_res["meta_info"]["prompt_tokens"]
+
+                score_res = requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "input_ids": input_token_ids + output_token_ids,
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 0,
+                        },
+                        "return_logprob": True,
+                        "top_logprobs_num": top_k,
+                        "token_ids_logprob": probe_token_ids,
+                        "logprob_start_len": 0,
+                    },
+                ).json()
+
+                score_logprobs = score_res["meta_info"]["input_token_logprobs"][
+                    num_prompt_tokens:
+                ]
+                score_top_logprobs = score_res["meta_info"]["input_top_logprobs"][
+                    num_prompt_tokens:
+                ]
+                score_tid_logprobs = score_res["meta_info"]["input_token_ids_logprobs"][
+                    num_prompt_tokens:
+                ]
+
+                self.assertEqual(len(decode_logprobs), len(score_logprobs))
+
+                decode_vals = np.array([t[0] for t in decode_logprobs])
+                score_vals = np.array([t[0] for t in score_logprobs])
+                max_diff = np.max(np.abs(decode_vals - score_vals))
+                print(f"[round {round_idx}] logprob max_diff={max_diff:.6f}")
+                self.assertLess(max_diff, self.logprob_match_delta)
+
+                for pos in range(len(decode_logprobs)):
+                    dec_top = {t[1]: t[0] for t in decode_top_logprobs[pos]}
+                    scr_top = {t[1]: t[0] for t in score_top_logprobs[pos]}
+                    common_ids = set(dec_top.keys()) & set(scr_top.keys())
+                    self.assertGreater(len(common_ids), 0)
+                    for tid in common_ids:
+                        self.assertAlmostEqual(
+                            dec_top[tid], scr_top[tid], delta=self.logprob_match_delta
+                        )
+
+                self.assertEqual(len(decode_tid_logprobs), len(score_tid_logprobs))
+                for pos in range(len(decode_tid_logprobs)):
+                    dec_tid = {t[1]: t[0] for t in decode_tid_logprobs[pos]}
+                    scr_tid = {t[1]: t[0] for t in score_tid_logprobs[pos]}
+                    self.assertEqual(set(dec_tid.keys()), set(scr_tid.keys()))
+                    for tid in dec_tid:
+                        self.assertAlmostEqual(
+                            dec_tid[tid], scr_tid[tid], delta=self.logprob_match_delta
+                        )
+
+    def test_token_ids_logprob_ragged(self):
+        """Regression: ragged token_ids_logprob lists in one batch must not crash."""
+
+        def send(probe_ids):
+            return requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": "Hello world",
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+                    "return_logprob": True,
+                    "top_logprobs_num": 3,
+                    "token_ids_logprob": probe_ids,
+                },
+            ).json()
+
+        ragged_probes = [
+            [1, 2],
+            [3, 4, 5],
+            [6],
+            [10, 20, 30, 40],
+            [1, 2],
+            [3, 4, 5],
+            [6],
+            [10, 20, 30, 40],
+        ]
+        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
+            futs = [pool.submit(send, ids) for ids in ragged_probes]
+            for f in concurrent.futures.as_completed(futs):
+                res = f.result()
+                self.assertIn("text", res, f"Server error: {res}")
+
+
+class SpecPenaltyKit:
+    """Penalty parameters under concurrency must not crash / corrupt output."""
+
+    def test_penalty_mixed(self):
+        args = [
+            {},
+            {},
+            {},
+            {"frequency_penalty": 2},
+            {"presence_penalty": 1},
+            {"min_new_tokens": 16},
+            {"frequency_penalty": 0.2},
+            {"presence_penalty": 0.4},
+            {"min_new_tokens": 8},
+            {"frequency_penalty": 0.4, "presence_penalty": 0.8},
+            {"frequency_penalty": 0.4, "min_new_tokens": 12},
+            {"presence_penalty": 0.8, "min_new_tokens": 12},
+            {"presence_penalty": -0.3, "frequency_penalty": 1.3, "min_new_tokens": 32},
+            {"presence_penalty": 0.3, "frequency_penalty": -1.3, "min_new_tokens": 32},
+        ]
+        random.shuffle(args * 5)
+        with ThreadPoolExecutor(8) as executor:
+            list(executor.map(self.run_decode, args))
+
+
+class SpecFeatureKit:
+    """Radix attention, constrained decoding, concurrent abort."""
+
+    def test_radix_attention(self):
+        run_radix_attention_test(self.base_url)
+        self.assertIsNone(self.process.poll())
+
+    def test_request_abort(self):
+        concurrency = 4
+        threads = [
+            threading.Thread(target=self.send_request) for _ in range(concurrency)
+        ] + [
+            threading.Thread(target=self.send_requests_abort)
+            for _ in range(concurrency)
+        ]
+        for worker in threads:
+            worker.start()
+        for p in threads:
+            p.join()
+
+    def test_constrained_decoding(self):
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Give me a json"},
+        ]
+        response = requests.post(
+            self.base_url + "/v1/chat/completions",
+            json={
+                "model": self.model,
+                "messages": messages,
+                "temperature": 0,
+                "response_format": {"type": "json_object"},
+            },
+        )
+        self.assertEqual(response.status_code, 200)
+        res = response.json()
+        self.assertIn("choices", res)
+        self.assertEqual(len(res["choices"]), 1)
+        self.assertIn("message", res["choices"][0])
+        self.assertIn("content", res["choices"][0]["message"])
+
+        content_json = res["choices"][0]["message"]["content"]
+        try:
+            content = json.loads(content_json)
+            self.assertIsInstance(content, dict)
+        except Exception:
+            self.fail(f"parse JSON failed: {content_json}")
+
+
+class SpecHiddenStatesKit:
+    """return_hidden_states under spec V2 (regression for issue #26163).
+
+    Requires the server launched with --enable-return-hidden-states
+    (set ``enable_return_hidden_states = True`` on the fixture class).
+    """
+
+    def test_return_hidden_states(self):
+        # Two prompts of different lengths to exercise the per-req stride
+        # window: under spec V2 hidden_states is [bs * num_draft_tokens, dim],
+        # so a wrong index aliases a neighbor request's accepted rows.
+        prompts = [
+            "Repeat: the quick brown fox the quick brown fox the quick brown fox",
+            "Count down from ten: ten nine eight",
+        ]
+        res = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompts,
+                "sampling_params": {"temperature": 0, "max_new_tokens": 32},
+                "return_hidden_states": True,
+            },
+        )
+        self.assertEqual(res.status_code, 200)
+        outputs = res.json()
+
+        for out in outputs:
+            meta = out["meta_info"]
+            hs = meta["hidden_states"]
+            ct = meta["completion_tokens"]
+            # One hidden-state entry per completion token: hs[0] is the prefill
+            # block (List[List[float]]), hs[1:] are per-decode-token rows.
+            self.assertEqual(
+                len(hs),
+                ct,
+                f"len(hidden_states)={len(hs)} but completion_tokens={ct}",
+            )
+            decode_rows = hs[1:]
+            self.assertGreater(len(decode_rows), 0)
+            hidden_dim = len(decode_rows[0])
+            self.assertGreater(hidden_dim, 0)
+            for row in decode_rows:
+                self.assertIsInstance(row, list)
+                self.assertEqual(len(row), hidden_dim)
+
+
+class SpecGrammarKit:
+    """Grammar-constrained structured output under spec decoding.
+
+    Regression for spec verify accepting tokens past grammar termination: the
+    output must be valid JSON with nothing emitted after completion, and the
+    logprob count must match the (truncated) completion-token count.
+    """
+
+    # Override per config if a different schema is desired.
+    grammar_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string", "pattern": "^[\\w]+$"},
+                "population": {"type": "integer"},
+                "country": {"type": "string", "pattern": "^[\\w ]+$"},
+                "capital": {"type": "string", "pattern": "^[\\w ]+$"},
+            },
+            "required": ["name", "population", "country", "capital"],
+        }
+    )
+
+    def _generate_grammar(self, return_logprob: bool):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": "Here is the information of the capital of France in the JSON format.\n",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 256,
+                    "json_schema": self.grammar_json_schema,
+                },
+                "return_logprob": return_logprob,
+                "logprob_start_len": 0,
+            },
+        )
+        self.assertEqual(response.status_code, 200, response.text)
+        out = response.json()
+        self.assertGreater(
+            out["meta_info"]["spec_verify_ct"],
+            0,
+            "expected spec decoding to run (spec_verify_ct > 0)",
+        )
+        return out
+
+    def test_grammar_structured_output_no_trailing_tokens(self):
+        """Output is valid JSON with nothing emitted past grammar completion."""
+        out = self._generate_grammar(return_logprob=False)
+        text = out["text"]
+        parsed = json.loads(text)
+        for key in ("name", "population", "country", "capital"):
+            self.assertIn(key, parsed)
+        self.assertTrue(
+            text.strip().endswith("}"), f"unexpected trailing tokens: {text!r}"
+        )
+
+    def test_grammar_logprob_count_matches_completion_tokens(self):
+        """Trimmed spec tokens keep logprob count == completion token count."""
+        out = self._generate_grammar(return_logprob=True)
+        meta = out["meta_info"]
+        completion_tokens = meta["completion_tokens"]
+        output_logprobs = meta["output_token_logprobs"]
+        self.assertEqual(
+            len(output_logprobs),
+            completion_tokens,
+            "output logprobs must align with retained (trimmed) tokens: "
+            f"got {len(output_logprobs)} logprobs vs {completion_tokens} completion tokens",
+        )
+        json.loads(out["text"])

python/sglang/test/kits/streaming_session_kit.py

@@ -0,0 +1,442 @@
+"""Streaming-session test method mixins.
+
+Pair these with `StreamingSessionServerBase` (from sglang.test.server_fixtures.streaming_session_fixture)
+to assemble a concrete test class. Per the sglang fixture/kit split:
+the fixture only launches the server; the kit owns the `test_*` methods.
+
+- `StreamingSessionKitMixin`: KV-inheritance + chunked-prefill + abort-recovery
+  + concurrent-logprob/stress test methods.
+- `AbortLeakReproKitMixin`: single test method for abort-heavy chunked-prefill leak repro.
+"""
+
+import asyncio
+import time
+
+import requests
+
+from sglang.test.server_fixtures.streaming_session_fixture import (
+    _abort_repro_run_all,
+    _concurrent_logprob_run,
+    _stress_run_all,
+)
+
+
+class StreamingSessionKitMixin:
+    """Streaming-session KV-inheritance + retract/abort-recovery suite."""
+
+    # Allowed inherited-cache offsets vs the previous turn's total. Non-overlap
+    # spec decode can be off by 1: the bonus token's KV is only computed by the
+    # next forward, which sync skips at finish (overlap drains it, so it's 0).
+    kv_inherit_offsets = (0,)
+
+    def test_kv_cache_inheritance(self, gen_len=12):
+        """Each turn's cached_tokens must equal previous turn's prompt+completion
+        (modulo kv_inherit_offsets)."""
+        chunks = [
+            "Let me tell you something about France.",
+            "The capital of France is",
+            "The population of the city is",
+        ]
+        chunks_ids = [self.tokenizer.encode(x) for x in chunks]
+        for i in range(1, len(chunks_ids)):
+            if chunks_ids[i][0] == self.tokenizer.bos_token_id:
+                chunks_ids[i] = chunks_ids[i][1:]
+
+        # === Part 1: streaming session — check KV inheritance ===
+        requests.post(self.base_url + "/flush_cache")
+        session_id = requests.post(
+            self.base_url + "/open_session",
+            json={"capacity_of_str_len": 1000, "streaming": True},
+        ).json()
+        rid = None
+
+        prev_kv_len = 0
+        for turn_idx, chunk_ids in enumerate(chunks_ids):
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": chunk_ids,
+                    "session_params": {"id": session_id, "rid": rid},
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": gen_len,
+                        "no_stop_trim": True,
+                        "skip_special_tokens": False,
+                    },
+                },
+            ).json()
+            rid = response["meta_info"]["id"]
+            cached = response["meta_info"]["cached_tokens"]
+            prompt_tokens = response["meta_info"]["prompt_tokens"]
+            completion_tokens = response["meta_info"]["completion_tokens"]
+
+            if turn_idx == 0:
+                # Turn 1: cache flushed, no hit.
+                self.assertEqual(cached, 0, "Turn 1: clean start, no cache hit")
+            else:
+                # Turns 2+: cached_tokens reflects KV inherited from previous turn
+                # (via inherit_kv_states, not radix tree matching).
+                allowed = {prev_kv_len + off for off in self.kv_inherit_offsets}
+                self.assertIn(
+                    cached,
+                    allowed,
+                    f"Turn {turn_idx + 1}: inherited {cached} not in {sorted(allowed)}",
+                )
+            prev_kv_len = prompt_tokens + completion_tokens
+
+        # Close the session.
+        ret = requests.post(
+            self.base_url + "/close_session",
+            json={"session_id": session_id},
+        )
+        self.assertEqual(ret.status_code, 200)
+
+    def test_leak_logprob_concurrent(self) -> None:
+        """Concurrent multi-session × 3 logprob modes (output / input / none),
+        watch for KV leak."""
+        requests.post(self.base_url + "/flush_cache")
+        # Output logprob
+        asyncio.run(
+            _concurrent_logprob_run(self.base_url, self.tokenizer, return_logprob=True)
+        )
+        # Input logprob (logprob_start_len=0)
+        asyncio.run(
+            _concurrent_logprob_run(
+                self.base_url,
+                self.tokenizer,
+                return_logprob=True,
+                logprob_start_len=0,
+            )
+        )
+        # No logprob
+        asyncio.run(_concurrent_logprob_run(self.base_url, self.tokenizer))
+        time.sleep(3)
+        assert (
+            requests.get(self.base_url + "/health").status_code == 200
+        ), "Server unhealthy after concurrent logprob sessions."
+
+    def test_stress_concurrent_sessions(self) -> None:
+        """High concurrency streaming + non-streaming with retract pressure;
+        scheduler must roll back streaming KV without leaking."""
+        requests.post(self.base_url + "/flush_cache")
+        asyncio.run(_stress_run_all(self.base_url, self.tokenizer))
+
+        for i in range(3):
+            ids = self.tokenizer.encode(f"Post-stress cleanup {i}.")
+            requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": ids,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 4},
+                },
+            )
+
+        time.sleep(5)
+        health = requests.get(self.base_url + "/health")
+        self.assertEqual(
+            health.status_code,
+            200,
+            "Server unhealthy after concurrent stress test — "
+            "likely a token leak from retract/mixed-chunk + streaming session.",
+        )
+
+    def test_nth_mid_abort_recovery(self) -> None:
+        """Abort an Nth-turn request mid-decode; session rolls back to last
+        successful turn."""
+        requests.post(self.base_url + "/flush_cache")
+
+        resp = requests.post(
+            self.base_url + "/open_session",
+            json={"capacity_of_str_len": 50000, "streaming": True},
+        )
+        self.assertEqual(resp.status_code, 200)
+        session_id = resp.json()
+
+        try:
+            # Turn 1: normal generate to create slot.
+            ids_1 = self.tokenizer.encode("Tell me a very long story about a wizard.")
+            resp_1 = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": ids_1,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 16},
+                    "session_params": {"id": session_id, "rid": None},
+                },
+                timeout=30,
+            )
+            self.assertEqual(resp_1.status_code, 200, resp_1.text)
+            data_1 = resp_1.json()
+            turn_1_total = (
+                data_1["meta_info"]["prompt_tokens"]
+                + data_1["meta_info"]["completion_tokens"]
+            )
+
+            # Turn 2: long generate, then abort mid-decode.
+            ids_2 = self.tokenizer.encode(" Continue the story in great detail.")
+
+            import threading
+
+            result = [None]
+
+            def do_generate():
+                r = requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "input_ids": ids_2,
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 100000,
+                        },
+                        "session_params": {"id": session_id, "rid": None},
+                    },
+                    timeout=60,
+                )
+                result[0] = r
+
+            t = threading.Thread(target=do_generate)
+            t.start()
+            time.sleep(0.5)
+            abort_resp = requests.post(
+                self.base_url + "/abort_request",
+                json={"rid": "", "abort_all": True},
+                timeout=10,
+            )
+            self.assertEqual(abort_resp.status_code, 200, abort_resp.text)
+            t.join(timeout=30)
+
+            self.assertIsNotNone(result[0], "Turn 2 should have returned")
+            data_2 = result[0].json()
+            self.assertEqual(
+                data_2["meta_info"]["finish_reason"]["type"],
+                "abort",
+                "Turn 2 should be aborted, not finished normally",
+            )
+
+            # Turn 3: recovery. Rolls back to turn 1.
+            ids_3 = self.tokenizer.encode(" What happens next?")
+            for attempt in range(20):
+                resp_3 = requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "input_ids": ids_3,
+                        "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+                        "session_params": {"id": session_id, "rid": None},
+                    },
+                    timeout=30,
+                )
+                if resp_3.status_code == 200:
+                    break
+                time.sleep(0.5)
+            self.assertEqual(resp_3.status_code, 200, resp_3.text)
+            data_3 = resp_3.json()
+            # prompt_tokens = turn_1_total + append (BOS stripped).
+            bos = 1 if ids_3[0] == self.tokenizer.bos_token_id else 0
+            expected_prompt_3 = turn_1_total + len(ids_3) - bos
+            self.assertEqual(
+                data_3["meta_info"]["prompt_tokens"],
+                expected_prompt_3,
+                "prompt_tokens must equal turn_1_total + append (no stale abort context)",
+            )
+        finally:
+            requests.post(
+                self.base_url + "/close_session",
+                json={"session_id": session_id},
+            )
+
+        health = requests.get(self.base_url + "/health", timeout=10)
+        self.assertEqual(health.status_code, 200)
+
+    def test_first_mid_abort_recovery(self) -> None:
+        """Abort the very first request mid-decode (no slot yet; ephemeral
+        slot is created and nuked). Session must still be usable."""
+        requests.post(self.base_url + "/flush_cache")
+
+        resp = requests.post(
+            self.base_url + "/open_session",
+            json={"capacity_of_str_len": 50000, "streaming": True},
+        )
+        self.assertEqual(resp.status_code, 200)
+        session_id = resp.json()
+
+        try:
+            ids_1 = self.tokenizer.encode("Tell me a very long story about a wizard.")
+
+            import threading
+
+            result = [None]
+
+            def do_generate():
+                r = requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "input_ids": ids_1,
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 100000,
+                        },
+                        "session_params": {"id": session_id, "rid": None},
+                    },
+                    timeout=60,
+                )
+                result[0] = r
+
+            t = threading.Thread(target=do_generate)
+            t.start()
+            time.sleep(0.5)
+            abort_resp = requests.post(
+                self.base_url + "/abort_request",
+                json={"rid": "", "abort_all": True},
+                timeout=10,
+            )
+            self.assertEqual(abort_resp.status_code, 200, abort_resp.text)
+            t.join(timeout=30)
+
+            self.assertIsNotNone(result[0], "Turn 1 should have returned")
+            data_1 = result[0].json()
+            self.assertEqual(
+                data_1["meta_info"]["finish_reason"]["type"],
+                "abort",
+                "Turn 1 should be aborted, not finished normally",
+            )
+
+            # Turn 2: recovery. No inherited context (req_nodes empty).
+            ids_2 = self.tokenizer.encode("Tell me a short joke.")
+            for attempt in range(20):
+                resp_2 = requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "input_ids": ids_2,
+                        "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+                        "session_params": {"id": session_id, "rid": None},
+                    },
+                    timeout=30,
+                )
+                if resp_2.status_code == 200:
+                    break
+                time.sleep(0.5)
+            self.assertEqual(resp_2.status_code, 200, resp_2.text)
+            data_2 = resp_2.json()
+            self.assertEqual(
+                data_2["meta_info"]["prompt_tokens"],
+                len(ids_2),
+                "prompt_tokens must equal turn 2 input only (no inherited context)",
+            )
+        finally:
+            requests.post(
+                self.base_url + "/close_session",
+                json={"session_id": session_id},
+            )
+
+        health = requests.get(self.base_url + "/health", timeout=10)
+        self.assertEqual(health.status_code, 200)
+
+    def test_preabort_recovery(self) -> None:
+        """Pre-abort (rejected by create_req) preserves the slot; next turn
+        inherits correctly."""
+        requests.post(self.base_url + "/flush_cache")
+
+        resp = requests.post(
+            self.base_url + "/open_session",
+            json={"capacity_of_str_len": 50000, "streaming": True},
+        )
+        self.assertEqual(resp.status_code, 200)
+        session_id = resp.json()
+
+        try:
+            # Turn 1: normal generate to create slot.
+            ids_1 = self.tokenizer.encode("Tell me a very long story about a wizard.")
+            resp_1 = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": ids_1,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 16},
+                    "session_params": {"id": session_id, "rid": None},
+                },
+                timeout=30,
+            )
+            self.assertEqual(resp_1.status_code, 200, resp_1.text)
+            data_1 = resp_1.json()
+            turn_1_total = (
+                data_1["meta_info"]["prompt_tokens"]
+                + data_1["meta_info"]["completion_tokens"]
+            )
+
+            # Turn 2: pre-aborted via unsupported offset parameter.
+            ids_2 = self.tokenizer.encode(" This should be rejected.")
+            resp_2 = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": ids_2,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+                    "session_params": {
+                        "id": session_id,
+                        "rid": None,
+                        "offset": 1,
+                    },
+                },
+                timeout=30,
+            )
+            self.assertIn(resp_2.status_code, (200, 400), resp_2.text)
+
+            # Turn 3: normal append. Slot should be intact from turn 1.
+            ids_3 = self.tokenizer.encode(" What happens next?")
+            resp_3 = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": ids_3,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+                    "session_params": {"id": session_id, "rid": None},
+                },
+                timeout=30,
+            )
+            self.assertEqual(resp_3.status_code, 200, resp_3.text)
+            data_3 = resp_3.json()
+            bos = 1 if ids_3[0] == self.tokenizer.bos_token_id else 0
+            expected_prompt_3 = turn_1_total + len(ids_3) - bos
+            self.assertEqual(
+                data_3["meta_info"]["prompt_tokens"],
+                expected_prompt_3,
+                "prompt_tokens must equal turn_1_total + append (slot preserved)",
+            )
+        finally:
+            requests.post(
+                self.base_url + "/close_session",
+                json={"session_id": session_id},
+            )
+
+        health = requests.get(self.base_url + "/health", timeout=10)
+        self.assertEqual(health.status_code, 200)
+
+
+class AbortLeakReproKitMixin:
+    """Abort-heavy chunked-prefill leak repro."""
+
+    def test_abort_heavy_chunked_prefill_does_not_leak(self) -> None:
+        requests.post(self.base_url + "/flush_cache")
+
+        asyncio.run(_abort_repro_run_all(self.base_url, self.tokenizer))
+
+        for i in range(3):
+            ids = self.tokenizer.encode(f"Post-session cleanup request {i}.")
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": ids,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 4},
+                },
+                timeout=30,
+            )
+            self.assertEqual(response.status_code, 200, response.text)
+
+        time.sleep(5)
+        self.assertIsNone(
+            self.process.poll(),
+            "Server crashed during abort-heavy streaming session repro.",
+        )
+
+        health = requests.get(self.base_url + "/health", timeout=10)
+        self.assertEqual(
+            health.status_code,
+            200,
+            "Server unhealthy after abort-heavy streaming session cleanup.",
+        )

python/sglang/test/kits/unified_radix_cache_kit.py

@@ -0,0 +1,147 @@
+import random
+from types import SimpleNamespace
+from urllib.parse import urlparse
+
+from sglang.test.kl_multiturn_utils import (
+    test_input_output_logprobs_match_decode_cache_hit_helper,
+    test_input_output_logprobs_match_helper,
+    test_input_output_logprobs_match_prefill_cache_hit_helper,
+)
+
+
+def _random_suffixes(n, length, seed):
+    """Generate n random token-id lists of the given length."""
+    rng = random.Random(seed)
+    return [[rng.randint(1, 30000) for _ in range(length)] for _ in range(n)]
+
+
+class UnifiedRadixTreeTestMixin:
+    """Mixin: gsm8k, mmlu and multi-turn KL tests with multi-branch interleaving."""
+
+    kl_threshold: float = 0.003
+    max_new_tokens: int = 512
+    num_groups: int = 3
+    branches_per_group: int = 3
+    prefix_len: int = 512
+    prefill_cache_assert = None
+    decode_cache_assert = None
+    sampling_temperature: float = 1
+    decode_hit_request_batch_size: int | None = None
+    decode_hit_inter_batch_delay_s: float = 0
+
+    gsm8k_threshold: float = 0.93
+    mmlu_threshold: float = 0.8
+    num_gsm8k_questions: int = 200
+
+    def test_gsm8k(self):
+        """Few-shot GSM8K math reasoning accuracy."""
+        from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
+
+        url = urlparse(self.base_url)
+        args = SimpleNamespace(
+            num_shots=10,
+            data_path=None,
+            num_questions=self.num_gsm8k_questions,
+            max_new_tokens=16000,
+            parallel=128,
+            host=f"http://{url.hostname}",
+            port=int(url.port),
+        )
+        metrics = run_few_shot_gsm8k(args)
+        print(
+            f"[{self.__class__.__name__}] GSM8K accuracy: {metrics['accuracy']:.3f} "
+            f"(threshold: {self.gsm8k_threshold})"
+        )
+        self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
+
+    def test_mmlu(self):
+        """Simple-evals MMLU multi-task accuracy."""
+        from sglang.test.run_eval import run_eval as run_simple_eval
+
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmlu",
+            num_examples=64,
+            num_threads=32,
+        )
+        metrics = run_simple_eval(args)
+        print(
+            f"[{self.__class__.__name__}] MMLU score: {metrics['score']:.3f} "
+            f"(threshold: {self.mmlu_threshold})"
+        )
+        self.assertGreaterEqual(metrics["score"], self.mmlu_threshold)
+
+    def test_multiturn_logprobs_match(self):
+        """Helper 1: 3-turn, no explicit cache seeding."""
+        ids = self.input_ids[:4]
+        n = len(ids)
+        t2 = _random_suffixes(n, 512, seed=100)
+        t3 = _random_suffixes(n, 256, seed=200)
+        test_input_output_logprobs_match_helper(
+            self.base_url,
+            self.model,
+            self.kl_threshold,
+            ids,
+            turn_suffixes=[t2, t3],
+            assert_decode_cached_tokens=self.decode_cache_assert,
+            max_new_tokens=self.max_new_tokens,
+            sampling_temperature=self.sampling_temperature,
+        )
+
+    def test_multiturn_prefill_cache_hit_branching(self):
+        """Helper 2: prefill hit + 2 decode-hit turns, multi-branch interleaved."""
+        num_groups = self.num_groups
+        branches = self.branches_per_group
+        n = num_groups * branches
+        rng = random.Random(456)
+        prefix_ids, full_ids = [], []
+        for g in range(num_groups):
+            prefix = self.input_ids[g][: self.prefix_len]
+            for b in range(branches):
+                suffix = [rng.randint(1, 30000) for _ in range(256 + b * 64)]
+                prefix_ids.append(list(prefix))
+                full_ids.append(prefix + suffix)
+
+        t2 = _random_suffixes(n, 512, seed=789)
+        t3 = _random_suffixes(n, 256, seed=890)
+        test_input_output_logprobs_match_prefill_cache_hit_helper(
+            self.base_url,
+            self.model,
+            self.kl_threshold,
+            prefix_input_ids=prefix_ids,
+            full_input_ids=full_ids,
+            turn_suffixes=[t2, t3],
+            assert_prefill_cached_tokens=self.prefill_cache_assert,
+            assert_decode_cached_tokens=self.decode_cache_assert,
+            branches_per_group=branches,
+            max_new_tokens=self.max_new_tokens,
+            sampling_temperature=self.sampling_temperature,
+        )
+
+    def test_multiturn_decode_cache_hit_branching(self):
+        """Helper 3: 3-turn decode hit, multi-branch interleaved."""
+        num_groups = self.num_groups
+        branches = self.branches_per_group
+        n = num_groups * branches
+        first_turn = []
+        for g in range(num_groups):
+            base = self.input_ids[g][: self.prefix_len]
+            for _ in range(branches):
+                first_turn.append(list(base))
+
+        t2 = _random_suffixes(n, 512, seed=300)
+        t3 = _random_suffixes(n, 256, seed=400)
+        test_input_output_logprobs_match_decode_cache_hit_helper(
+            self.base_url,
+            self.model,
+            self.kl_threshold,
+            first_turn,
+            turn_suffixes=[t2, t3],
+            assert_decode_cached_tokens=self.decode_cache_assert,
+            branches_per_group=branches,
+            max_new_tokens=self.max_new_tokens,
+            sampling_temperature=self.sampling_temperature,
+            request_batch_size=self.decode_hit_request_batch_size,
+            inter_batch_delay_s=self.decode_hit_inter_batch_delay_s,
+        )