Merge branch 'main' into dsv41-pd

test/README.md

@@ -3,11 +3,11 @@
 This page covers principles and essentials: folder layout, how to run tests, registration, and suite selection. For complete references, see the skill guides:
 
 - **Writing tests** — templates, fixtures, model selection, complete suite tables, checklist: [`.claude/skills/write-sglang-test/SKILL.md`](../.claude/skills/write-sglang-test/SKILL.md)
-- **CI pipeline internals** — stage flow diagrams, fast-fail layers, gating, partitioning, execution modes, debugging failures: [`.claude/skills/ci-workflow-guide/SKILL.md`](../.claude/skills/ci-workflow-guide/SKILL.md)
+- **CI pipeline internals** — stage flow diagrams, fail-fast layers, gating, partitioning, execution modes, debugging failures: [`.claude/skills/ci-workflow-guide/SKILL.md`](../.claude/skills/ci-workflow-guide/SKILL.md)
 
 ## CI Pipeline Overview
 
-The CI pipeline runs in three sequential stages: **A** (pre-flight, ~3 min) → **B** (basic, ~30 min) → **C** (advanced, ~30 min). Kernel and multimodal-gen tests run in parallel with stage B. For details on stage gating, fast-fail mechanisms, execution modes (PR vs scheduled vs manual dispatch), and debugging CI failures, see the [CI workflow guide](../.claude/skills/ci-workflow-guide/SKILL.md).
+The CI pipeline runs in three sequential stages: **A** (pre-flight, ~3 min) → **B** (basic, ~30 min) → **C** (advanced, ~30 min). Kernel and multimodal-gen tests run in parallel with stage B. For details on stage gating, fail-fast mechanisms, execution modes (PR vs scheduled vs manual dispatch), and debugging CI failures, see the [CI workflow guide](../.claude/skills/ci-workflow-guide/SKILL.md).
 
 ## Folder Organization
 

test/manual/ep/test_flashinfer_dispatcher.py

@@ -14,7 +14,7 @@ from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatch
 from sglang.srt.layers.moe.utils import initialize_moe_config
 from sglang.srt.runtime_context import get_context, publish
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 
 class TestFlashinferDispatcher(CustomTestCase):
@@ -44,9 +44,8 @@ class TestFlashinferDispatcher(CustomTestCase):
         publish(server_args, role="scheduler")
         initialize_moe_config()
 
-        initialize_model_parallel(
-            tensor_model_parallel_size=world_size, expert_model_parallel_size=world_size
-        )
+        publish_build_topology(tp_size=world_size, ep_size=world_size, world_rank=rank)
+        initialize_model_parallel()
 
     @classmethod
     def tearDownClass(cls):

test/manual/kv_canary/test_xpu_kv_canary_pp.py

@@ -0,0 +1,36 @@
+"""KV-canary end-to-end on Intel XPU with pipeline parallelism.
+
+``--pp 2`` routes the run through ``Qwen3ForCausalLM.set_embed_and_head`` (mha mode
+is Qwen/Qwen3-0.6B), the embedding/head handoff that syncs and releases the device
+cache. Needs two XPU cards, so it is manual until a 2-card lane is confirmed.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.test.kv_canary.e2e_base import CanaryE2EBase
+
+
+class TestXPUCanaryPipelineParallel(CanaryE2EBase):
+    """Clean canary run across a pipeline-parallel XPU pair."""
+
+    model_mode = "mha"
+    kv_canary_mode = CanaryMode.LOG
+    # --disable-cuda-graph is mandatory, not tuning: install_canary refuses a captured decode
+    # on a device that routes to the torch reference (host work and D2H, so replay checks nothing).
+    extra_server_args = ("--device", "xpu", "--disable-cuda-graph", "--pp", "2")
+    # The torch reference folds the chain slot-by-slot on the host, so the workload is much
+    # smaller than the CUDA-tuned defaults on the shared base.
+    default_parallel_n = 2
+    default_max_new_tokens = 32
+    default_request_timeout = 120.0
+
+    def test_no_violation(self) -> None:
+        self.send_parallel_requests()
+        self.assert_no_violation(wait_seconds=2.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/kv_transfer/test_mooncake_transfer_engine_init.py

@@ -45,7 +45,6 @@ def test_mooncake_te_condition(server_args: ServerArgs) -> bool:
     """
     Test the condition logic for using MooncakeTransferEngine.
     """
-    from sglang.srt.model_executor.model_runner import ModelRunner
 
     dummy_runner = SimpleNamespace(server_args=server_args, gpu_id=0)
     init_called = False
@@ -69,7 +68,11 @@ def test_mooncake_te_condition(server_args: ServerArgs) -> bool:
             return_value="127.0.0.1",
         ),
     ):
-        ModelRunner.init_shared_mooncake_transfer_engine(dummy_runner)
+        from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+            maybe_init_shared_mooncake_transfer_engine,
+        )
+
+        maybe_init_shared_mooncake_transfer_engine(gpu_id=dummy_runner.gpu_id)
 
     return init_called
 

test/manual/test_custom_allreduce.py

@@ -18,7 +18,7 @@ from sglang.srt.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 
 def get_open_port() -> int:
@@ -98,7 +98,8 @@ class TestCustomAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         # Set global server args to avoid "Global server args is not set yet!" error
@@ -161,7 +162,8 @@ class TestCustomAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         # Set global server args to avoid "Global server args is not set yet!" error

test/manual/test_forward_split_prefill.py

@@ -15,7 +15,7 @@ import torch
 
 from sglang.benchmark.one_batch import TreeCacheNamespace
 from sglang.srt.configs.model_config import ModelConfig
-from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+from sglang.srt.distributed import bootstrap
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.forward_context import (
@@ -23,7 +23,7 @@ from sglang.srt.model_executor.forward_context import (
     set_forward_context,
 )
 from sglang.srt.model_executor.model_runner import ModelRunner
-from sglang.srt.runtime_context import publish
+from sglang.srt.runtime_context import SpawnRanks, publish
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -57,15 +57,24 @@ class TestForwardSplitPrefill(CustomTestCase):
 
         cls.port_args = PortArgs.init_new(cls.server_args)
 
-        publish(cls.server_args, role="scheduler")
+        publish(
+            cls.server_args,
+            role="scheduler",
+            ranks=SpawnRanks(world_rank=0, gpu_id=0),
+        )
 
         # Load model and tokenizer
         cls.model_config = ModelConfig.from_server_args(cls.server_args)
+        bootstrap.init_parallel_runtime(
+            server_args=cls.server_args,
+            model_config=cls.model_config,
+            device=cls.device,
+            dist_port=cls.port_args.nccl_port,
+        )
         cls.model_runner = ModelRunner(
             model_config=cls.model_config,
             mem_fraction_static=cls.server_args.mem_fraction_static,
             gpu_id=0,
-            ps=ParallelState.trivial(tp_size=cls.tp_size),
             nccl_port=cls.port_args.nccl_port,
             server_args=cls.server_args,
         )

test/manual/test_quick_allreduce.py

@@ -23,7 +23,7 @@ from sglang.srt.distributed.parallel_state import (
     graph_capture,
     initialize_model_parallel,
 )
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 torch.manual_seed(42)
 random.seed(44)  # keep the deterministic seed
@@ -117,7 +117,8 @@ class TestQuickAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         # A small all_reduce for warmup.
@@ -186,7 +187,8 @@ class TestQuickAllReduce(CustomTestCase):
             distributed_init_method=distributed_init_method,
             local_rank=rank,
         )
-        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        publish_build_topology(tp_size=world_size, world_rank=rank)
+        initialize_model_parallel()
         group = get_tensor_model_parallel_group().device_group
 
         for sz in self.TEST_SIZES:

test/manual/test_vlm_accuracy.py

@@ -9,7 +9,7 @@ import torch.nn.functional as F
 from transformers import AutoModel, AutoProcessor, AutoTokenizer
 
 from sglang.srt.configs.model_config import ModelConfig
-from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+from sglang.srt.distributed import bootstrap
 from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
 from sglang.srt.managers.mm_utils import embed_mm_inputs, init_mm_embedding_cache
 from sglang.srt.managers.schedule_batch import (
@@ -20,7 +20,7 @@ from sglang.srt.managers.schedule_batch import (
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
 from sglang.srt.parser.conversation import generate_chat_conv
-from sglang.srt.runtime_context import publish
+from sglang.srt.runtime_context import SpawnRanks, get_device, publish
 from sglang.srt.server_args import ServerArgs
 from sglang.test.test_utils import download_image_with_retry
 
@@ -146,12 +146,22 @@ class VisionLLMLogitsBase(unittest.IsolatedAsyncioTestCase):
             model_path=self.model_path,
             disable_cuda_graph=True,
         )
-        publish(server_args, role="scheduler")
+        publish(
+            server_args,
+            role="scheduler",
+            ranks=SpawnRanks(world_rank=0, gpu_id=0),
+        )
+        model_config = ModelConfig(self.model_path, model_override_args="{}")
+        bootstrap.init_parallel_runtime(
+            server_args=server_args,
+            model_config=model_config,
+            device=get_device().device,
+            dist_port=12435,
+        )
         self.model_runner = ModelRunner(
-            model_config=ModelConfig(self.model_path, model_override_args="{}"),
+            model_config=model_config,
             mem_fraction_static=0.8,
             gpu_id=0,
-            ps=ParallelState.trivial(),
             nccl_port=12435,
             server_args=server_args,
         )

test/manual/test_weight_loader_v2_equiv.py

@@ -24,6 +24,7 @@ import unittest
 import torch
 
 from sglang.srt.environ import envs
+from sglang.test.test_utils import publish_build_topology
 
 MODEL = "Qwen/Qwen2-0.5B"
 
@@ -43,7 +44,8 @@ def _init_model_parallel() -> None:
             local_rank=0,
             distributed_init_method="tcp://127.0.0.1:29634",
         )
-        initialize_model_parallel(tensor_model_parallel_size=1)
+        publish_build_topology(tp_size=1)
+        initialize_model_parallel()
         monkey_patch_vllm_parallel_state()
     except AssertionError:
         pass

test/registered/amd/test_qwen35_moe_softmax_topk.py

@@ -0,0 +1,127 @@
+"""ROCm coverage for the decode-sized Qwen3.5 MoE softmax router."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.layers.moe import topk as topk_module
+from sglang.srt.utils import is_gfx95_supported
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and torch.version.hip and is_gfx95_supported(),
+    "requires AMD gfx95",
+)
+class TestQwen35MoeSoftmaxTopK(CustomTestCase):
+    def test_triton_dispatch_matches_aiter(self):
+        from aiter.fused_moe import fused_topk as aiter_fused_topk
+
+        for num_tokens in (4, 12, 128):
+            with self.subTest(num_tokens=num_tokens):
+                torch.manual_seed(num_tokens)
+                hidden_states = torch.randn(
+                    num_tokens, 4096, device="cuda", dtype=torch.bfloat16
+                )
+                router_logits = torch.randn(
+                    num_tokens, 512, device="cuda", dtype=torch.bfloat16
+                )
+                ref_weights = torch.empty(
+                    num_tokens, 10, device="cuda", dtype=torch.float32
+                )
+                ref_ids = torch.empty(num_tokens, 10, device="cuda", dtype=torch.int32)
+                ref_weights, ref_ids = aiter_fused_topk(
+                    hidden_states,
+                    router_logits,
+                    10,
+                    True,
+                    topk_ids=ref_ids,
+                    topk_weights=ref_weights,
+                )
+
+                with (
+                    patch.object(topk_module, "_use_aiter", True),
+                    patch.object(topk_module, "_is_gfx95", True),
+                    patch.object(
+                        topk_module,
+                        "aiter_fused_topk",
+                        side_effect=AssertionError("AITER top-k should be bypassed"),
+                        create=True,
+                    ),
+                ):
+                    weights, ids = topk_module.fused_topk(
+                        hidden_states,
+                        router_logits,
+                        topk=10,
+                        renormalize=True,
+                    )
+
+                torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0)
+                torch.testing.assert_close(weights, ref_weights, rtol=1e-5, atol=1e-6)
+
+    def test_dispatch_envelope_is_narrow(self):
+        hidden_states = torch.empty(128, 4096, device="cuda", dtype=torch.bfloat16)
+        logits = torch.empty(128, 512, device="cuda", dtype=torch.bfloat16)
+        packed = torch.empty(1, device="cuda")
+        with (
+            patch.object(topk_module, "_use_aiter", True),
+            patch.object(topk_module, "_is_gfx95", True),
+        ):
+            self.assertTrue(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 10, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    torch.empty(129, 4096, device="cuda", dtype=torch.bfloat16),
+                    torch.empty(129, 512, device="cuda", dtype=torch.bfloat16),
+                    10,
+                    None,
+                    0,
+                    None,
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits.float(), 10, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states[:, :2048], logits, 10, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 8, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states,
+                    logits,
+                    10,
+                    torch.empty(512, device="cuda"),
+                    0,
+                    None,
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 10, None, 1, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 10, None, 0, packed
+                )
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py

@@ -137,6 +137,42 @@ class TestGdnReplayssmSpecFold(CustomTestCase):
             )
             self.assertTrue(torch.equal(out_plain, out_ring), f"{dtype=}")
 
+    def test_ring_write_accepts_strided_qkv_views(self):
+        inputs = _make_window(12)
+        packed_qkv = torch.cat(
+            [inputs[name].reshape(B * T, -1) for name in ("q", "k", "v")],
+            dim=-1,
+        )
+        q, k, v = packed_qkv.split([H * K, H * K, HV * V], dim=-1)
+        strided_qkv = {
+            "q": q.view(1, B * T, H, K),
+            "k": k.view(1, B * T, H, K),
+            "v": v.view(1, B * T, HV, V),
+        }
+        self.assertTrue(
+            all(not tensor.is_contiguous() for tensor in strided_qkv.values())
+        )
+
+        def run(qkv):
+            state = self._state(torch.float32).unsqueeze(0).contiguous()
+            rings = _make_rings()
+            output = _run_verify(
+                {**inputs, **qkv},
+                self.gating,
+                state[0],
+                self.slots,
+                rings=rings,
+            )
+            _fold(state, rings, self.slots, self.accept_lens)
+            return {"output": output, "state": state, **rings}
+
+        contiguous = run(
+            {name: tensor.contiguous() for name, tensor in strided_qkv.items()}
+        )
+        strided = run(strided_qkv)
+        for name in contiguous:
+            self.assertTrue(torch.equal(contiguous[name], strided[name]), name)
+
     def test_fold_matches_snapshot_baseline(self):
         for dtype in (torch.float32, torch.bfloat16):
             state = self._state(dtype)

test/registered/attention/unittests/gdn/test_triton.py

@@ -1,6 +1,6 @@
 import unittest
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import torch
 
@@ -8,6 +8,7 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
     HybridLinearAttnBackend,
     MambaAttnBackendBase,
 )
+from sglang.srt.layers.attention.linear import gdn_backend
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.utils import is_hip
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
@@ -316,6 +317,24 @@ class TestTritonGDNBackendCorrectness(CustomTestCase):
             ):
                 run_gdn_eagle_verify_case(self, case, topk=topk, spec_kind=spec_kind)
 
+    def test_triton_target_verify_skips_prefill_qkv_materialization(self):
+        case, topk, spec_kind = self.EAGLE_VERIFY_CASES[0]
+        with patch.object(
+            gdn_backend,
+            "fused_qkv_split_gdn_prefill",
+            side_effect=AssertionError,
+        ):
+            run_gdn_eagle_verify_case(self, case, topk=topk, spec_kind=spec_kind)
+
+    def test_triton_prefill_keeps_contiguous_qkv_materialization(self):
+        with patch.object(
+            gdn_backend,
+            "fused_qkv_split_gdn_prefill",
+            wraps=gdn_backend.fused_qkv_split_gdn_prefill,
+        ) as split_spy:
+            run_gdn_attention_case(self, self.CASES[0])
+        split_spy.assert_called()
+
     def test_runner_mode_eagle_verify_cuda_graph_cases(self):
         for case, topk, spec_kind in self.EAGLE_VERIFY_CUDA_GRAPH_CASES:
             with self.subTest(

test/registered/basic_perf/test_eagle3_latency.py

@@ -0,0 +1,60 @@
+"""The only test in the tree that bounds speculative decoding LATENCY; every
+other one bounds accept length. CUDA only -- AMD bounds are unmeasured.
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.perf_bench_kit import at_least, at_most, check_perf
+from sglang.test.test_utils import (
+    DEFAULT_DRAFT_MODEL_EAGLE3,
+    DEFAULT_TARGET_MODEL_EAGLE3,
+    CustomTestCase,
+    run_bench_serving,
+)
+
+register_cuda_ci(est_time=145, stage="extra-a", runner_config="1-gpu-large")
+
+
+class TestEagle3Latency(CustomTestCase):
+    def test_online_latency_eagle3(self):
+        res = run_bench_serving(
+            model=DEFAULT_TARGET_MODEL_EAGLE3,
+            num_prompts=300,
+            request_rate=8,
+            sharegpt_context_len=3072,
+            disable_ignore_eos=True,
+            dataset_name="sharegpt",
+            other_server_args=[
+                "--speculative-algorithm",
+                "EAGLE3",
+                "--speculative-draft-model-path",
+                DEFAULT_DRAFT_MODEL_EAGLE3,
+                "--speculative-num-steps",
+                "5",
+                "--speculative-eagle-topk",
+                "4",
+                "--speculative-num-draft-tokens",
+                "16",
+                "--mem-fraction-static",
+                "0.7",
+                # The draft checkpoint ships fp16 and the target bf16; the CUDA
+                # rmsnorm path rejects a weight and activation pair that disagree.
+                "--dtype",
+                "float16",
+            ],
+            need_warmup=True,
+            seed=42,
+        )
+
+        check_perf(
+            self,
+            at_most(
+                "median_e2e_latency_ms", res["median_e2e_latency_ms"], 1150, unit="ms"
+            ),
+            at_least("accept_length", res["accept_length"], 2.3),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_embeddings_api.py

@@ -0,0 +1,61 @@
+"""Latency and throughput of the /v1/embeddings endpoint."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import (
+    at_least,
+    at_most,
+    check_batch_scaling,
+    check_perf,
+)
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    run_embeddings_benchmark,
+    run_embeddings_benchmark_multi,
+)
+
+register_cuda_ci(est_time=245, stage="extra-a", runner_config="1-gpu-large")
+register_amd_ci(est_time=240, suite="stage-b-test-1-gpu-large-amd")
+
+
+class TestEmbeddingsAPI(CustomTestCase):
+    def test_embeddings_api_latency_throughput(self):
+        res = run_embeddings_benchmark(
+            model=DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
+            num_requests=1000,
+            batch_size=1,
+            input_tokens=500,
+            other_server_args=[],
+            need_warmup=True,
+        )
+
+        self.assertEqual(res["successful_requests"], res["total_requests"])
+        check_perf(
+            self,
+            at_most("avg_latency_ms", res["avg_latency_ms"], 23, amd=35, unit="ms"),
+            at_most("p95_latency_ms", res["p95_latency_ms"], 34, amd=40, unit="ms"),
+            at_least("throughput", res["throughput"], 45, amd=30, unit="req/s"),
+        )
+
+    def test_embeddings_api_batch_scaling(self):
+        check_batch_scaling(
+            self,
+            lambda batch_sizes: run_embeddings_benchmark_multi(
+                DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
+                batch_sizes,
+                num_requests=500,
+                input_tokens=500,
+            ),
+            # batch size, avg ms, p95 ms, then the same two relaxed for mi300x
+            [
+                (10, 44, 52, 80, 90),
+                (25, 72, 101, 140, 150),
+                (50, 126, 200, 230, 240),
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_lora_latency.py

@@ -0,0 +1,134 @@
+"""Latency of the LoRA serving path, with and without adapter churn."""
+
+import asyncio
+import itertools
+import unittest
+
+import requests
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import at_most, check_perf
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    run_bench_serving,
+)
+
+register_cuda_ci(est_time=490, stage="extra-a", runner_config="1-gpu-large")
+register_amd_ci(est_time=430, suite="stage-b-test-1-gpu-large-amd")
+
+
+class TestLoRALatency(CustomTestCase):
+    def test_online_lora_latency(self):
+        res = self._run_lora_latency_test(enable_background_task=False)
+
+        check_perf(
+            self,
+            at_most(
+                "median_e2e_latency_ms",
+                res["median_e2e_latency_ms"],
+                2270,
+                amd=3320,
+                unit="ms",
+            ),
+            # mi300x is about twice as slow as mi325 on LoRA TTFT.
+            at_most("median_ttft_ms", res["median_ttft_ms"], 52, amd=100, unit="ms"),
+        )
+
+    def test_online_lora_latency_with_concurrent_adapter_updates(self):
+        res = self._run_lora_latency_test(enable_background_task=True)
+
+        check_perf(
+            self,
+            at_most(
+                "median_e2e_latency_ms",
+                res["median_e2e_latency_ms"],
+                3420,
+                amd=6000,
+                unit="ms",
+            ),
+            at_most("median_ttft_ms", res["median_ttft_ms"], 55, amd=130, unit="ms"),
+        )
+
+    def _run_lora_latency_test(self, enable_background_task: bool):
+        async def lora_loader_unloader_task(
+            base_url: str,
+            start_event: asyncio.Event,
+            stop_event: asyncio.Event,
+        ):
+            """
+            A background task that repeatedly loads and unloads a LoRA adapter.
+            """
+            await start_event.wait()
+
+            path_cycler = itertools.cycle(
+                [
+                    "pbevan11/llama-3.1-8b-ocr-correction",
+                    "faridlazuarda/valadapt-llama-3.1-8B-it-chinese",
+                    "philschmid/code-llama-3-1-8b-text-to-sql-lora",
+                ]
+            )
+            load_url = f"{base_url}/load_lora_adapter"
+            unload_url = f"{base_url}/unload_lora_adapter"
+            num_updates = 0
+
+            while not stop_event.is_set():
+                lora_path = next(path_cycler)
+                response = await asyncio.to_thread(
+                    requests.post,
+                    load_url,
+                    json={"lora_name": lora_path, "lora_path": lora_path},
+                )
+                self.assertTrue(
+                    response.ok, f"Failed to load LoRA adapter: {response.text}"
+                )
+                num_updates += 1
+
+                if stop_event.is_set():
+                    break
+
+                await asyncio.sleep(1)
+
+                response = await asyncio.to_thread(
+                    requests.post,
+                    unload_url,
+                    json={"lora_name": lora_path},
+                )
+                self.assertTrue(
+                    response.ok, f"Failed to unload LoRA adapter: {response.text}"
+                )
+                num_updates += 1
+
+                await asyncio.sleep(1)
+
+        background_task = lora_loader_unloader_task if enable_background_task else None
+        res = run_bench_serving(
+            model=DEFAULT_MODEL_NAME_FOR_TEST,
+            num_prompts=400,
+            request_rate=8,
+            other_server_args=[
+                "--enable-lora",
+                "--max-loras-per-batch",
+                "1",
+                "--disable-radix-cache",
+                "--random-seed",
+                "42",
+                "--mem-fraction-static",
+                "0.8",
+                "--lora-paths",
+                "nvidia/llama-3.1-nemoguard-8b-topic-control",
+                "--max-lora-rank",
+                "256",
+            ],
+            dataset_name="random",
+            random_input_len=256,
+            random_output_len=256,
+            lora_name=["nvidia/llama-3.1-nemoguard-8b-topic-control"],
+            background_task=background_task,
+        )
+
+        return res
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_moe_throughput.py

@@ -0,0 +1,53 @@
+"""Throughput of the MoE model on two GPUs, batched and at batch size one."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import at_least, check_perf
+from sglang.test.test_utils import (
+    DEFAULT_MOE_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    run_bench_offline_throughput,
+    run_bench_serving,
+)
+
+register_cuda_ci(est_time=290, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=770, suite="stage-b-test-2-gpu-large-amd")
+
+
+class TestMoEThroughput(CustomTestCase):
+    def test_moe_offline_throughput_default(self):
+        res = run_bench_serving(
+            model=DEFAULT_MOE_MODEL_NAME_FOR_TEST,
+            num_prompts=300,
+            request_rate=float("inf"),
+            other_server_args=["--tp", "2"],
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput",
+                res["output_throughput"],
+                2660,
+                amd=2100,
+                unit="token/s",
+            ),
+        )
+
+    def test_moe_tp2_bs1(self):
+        output_throughput = run_bench_offline_throughput(
+            DEFAULT_MOE_MODEL_NAME_FOR_TEST,
+            ["--tp", "2", "--cuda-graph-max-bs-decode", "2"],
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput", output_throughput, 139, amd=85, unit="token/s"
+            ),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_pp_throughput.py

@@ -0,0 +1,70 @@
+"""Throughput of pipeline parallelism on two GPUs, decode and long prefill."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import at_least, check_perf
+from sglang.test.test_utils import (
+    DEFAULT_MOE_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    is_in_amd_ci,
+    run_bench_serving,
+)
+
+register_cuda_ci(est_time=490, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=1030, suite="stage-b-test-2-gpu-large-amd")
+
+
+class TestPPThroughput(CustomTestCase):
+    def test_pp_offline_throughput_default_decode(self):
+        res = run_bench_serving(
+            model=DEFAULT_MOE_MODEL_NAME_FOR_TEST,
+            num_prompts=1000,
+            request_rate=float("inf"),
+            random_input_len=1,
+            random_output_len=1024,
+            other_server_args=["--pp-size", "2"],
+            need_warmup=True,
+            seed=42,
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput", res["output_throughput"], 6250, unit="token/s"
+            ),
+        )
+
+    def test_pp_long_context_prefill(self):
+        res = run_bench_serving(
+            model="meta-llama/Llama-3.3-70B-Instruct",
+            num_prompts=4,
+            request_rate=float("inf"),
+            random_input_len=128000,
+            random_output_len=1,
+            dataset_name="random",
+            other_server_args=[
+                "--quantization",
+                "fp8",
+                "--pp-size",
+                "2",
+            ]
+            + (["--mem-fraction-static", "0.7"] if is_in_amd_ci() else []),
+            need_warmup=False,
+            seed=42,
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "input_throughput",
+                res["input_throughput"],
+                4380,
+                amd=2190,
+                unit="token/s",
+            ),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_score_api.py

@@ -0,0 +1,55 @@
+"""Latency and throughput of the /v1/score endpoint."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import (
+    at_least,
+    at_most,
+    check_batch_scaling,
+    check_perf,
+)
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST_SCORE,
+    CustomTestCase,
+    run_score_benchmark,
+    run_score_benchmark_multi,
+)
+
+register_cuda_ci(est_time=215, stage="extra-a", runner_config="1-gpu-large")
+register_amd_ci(est_time=210, suite="stage-b-test-1-gpu-large-amd")
+
+
+class TestScoreAPI(CustomTestCase):
+    def test_score_api_latency_throughput(self):
+        res = run_score_benchmark(
+            model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST_SCORE,
+            num_requests=1000,
+            batch_size=10,
+            other_server_args=[],
+            need_warmup=True,
+        )
+
+        self.assertEqual(res["successful_requests"], res["total_requests"])
+        check_perf(
+            self,
+            at_most("avg_latency_ms", res["avg_latency_ms"], 31, amd=60, unit="ms"),
+            at_most("p95_latency_ms", res["p95_latency_ms"], 37, amd=65, unit="ms"),
+            at_least("throughput", res["throughput"], 32, amd=16, unit="req/s"),
+        )
+
+    def test_score_api_batch_scaling(self):
+        check_batch_scaling(
+            self,
+            lambda batch_sizes: run_score_benchmark_multi(
+                DEFAULT_SMALL_MODEL_NAME_FOR_TEST_SCORE,
+                batch_sizes,
+                num_requests=500,
+            ),
+            # batch size, avg ms, p95 ms, then the same two relaxed for mi300x
+            [(10, 32, 40, 60, 65), (25, 37, 42, 70, 80), (50, 54, 64, 80, 90)],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_serving_latency.py

@@ -0,0 +1,40 @@
+"""Latency of the default serving path on one large GPU."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import at_most, check_perf
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    run_bench_serving,
+)
+
+register_cuda_ci(est_time=190, stage="extra-a", runner_config="1-gpu-large")
+register_amd_ci(est_time=165, suite="stage-b-test-1-gpu-large-amd")
+
+
+class TestServingLatency(CustomTestCase):
+    def test_online_latency_default(self):
+        res = run_bench_serving(
+            model=DEFAULT_MODEL_NAME_FOR_TEST,
+            num_prompts=100,
+            request_rate=1,
+            other_server_args=[],
+        )
+
+        check_perf(
+            self,
+            at_most(
+                "median_e2e_latency_ms",
+                res["median_e2e_latency_ms"],
+                9140,
+                unit="ms",
+            ),
+            at_most("median_ttft_ms", res["median_ttft_ms"], 84, amd=115, unit="ms"),
+            at_most("median_itl_ms", res["median_itl_ms"], 9, unit="ms"),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_serving_throughput.py

@@ -0,0 +1,107 @@
+"""Offline throughput of the default serving path on one large GPU."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import at_least, check_perf
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_MODEL_NAME_FOR_TEST_FP8,
+    CustomTestCase,
+    run_bench_serving,
+)
+
+register_cuda_ci(est_time=710, stage="extra-a", runner_config="1-gpu-large")
+register_amd_ci(est_time=810, suite="stage-b-test-1-gpu-large-amd")
+
+
+class TestServingThroughput(CustomTestCase):
+    def test_offline_throughput_default(self):
+        res = run_bench_serving(
+            model=DEFAULT_MODEL_NAME_FOR_TEST,
+            num_prompts=500,
+            request_rate=float("inf"),
+            other_server_args=[],
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput",
+                res["output_throughput"],
+                4000,
+                amd=3050,
+                unit="token/s",
+            ),
+        )
+
+    def test_offline_throughput_non_stream_small_batch_size(self):
+        res = run_bench_serving(
+            model=DEFAULT_MODEL_NAME_FOR_TEST,
+            num_prompts=200,
+            request_rate=float("inf"),
+            other_server_args=["--max-running-requests", "10"],
+            dataset_name="sharegpt",
+            random_input_len=None,
+            random_output_len=None,
+            disable_stream=True,
+            need_warmup=True,
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput",
+                res["output_throughput"],
+                1100,
+                amd=1000,
+                unit="token/s",
+            ),
+        )
+
+    def test_offline_throughput_with_triton_attention_backend(self):
+        res = run_bench_serving(
+            model=DEFAULT_MODEL_NAME_FOR_TEST,
+            num_prompts=500,
+            request_rate=float("inf"),
+            other_server_args=[
+                "--attention-backend",
+                "triton",
+                "--context-length",
+                "8192",
+            ],
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput",
+                res["output_throughput"],
+                3730,
+                amd=2700,
+                unit="token/s",
+            ),
+        )
+
+    def test_offline_throughput_default_fp8(self):
+        res = run_bench_serving(
+            model=DEFAULT_MODEL_NAME_FOR_TEST_FP8,
+            num_prompts=500,
+            request_rate=float("inf"),
+            other_server_args=[],
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput",
+                res["output_throughput"],
+                4860,
+                amd=3500,
+                unit="token/s",
+            ),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_torch_compile_throughput.py

@@ -0,0 +1,33 @@
+"""Throughput of torch.compile at batch size one across two GPUs."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kits.perf_bench_kit import at_least, check_perf
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    run_bench_offline_throughput,
+)
+
+register_cuda_ci(est_time=75, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=280, suite="stage-b-test-2-gpu-large-amd")
+
+
+class TestTorchCompileThroughput(CustomTestCase):
+    def test_torch_compile_tp2_bs1(self):
+        output_throughput = run_bench_offline_throughput(
+            DEFAULT_MODEL_NAME_FOR_TEST,
+            ["--tp", "2", "--enable-torch-compile", "--cuda-graph-max-bs-decode", "2"],
+        )
+
+        check_perf(
+            self,
+            at_least(
+                "output_throughput", output_throughput, 255, amd=200, unit="token/s"
+            ),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_vlm_serving_aiter.py

@@ -0,0 +1,25 @@
+"""VLM serving perf on the aiter attention backend."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.kits.vlm_perf_kit import check_vlm_serving_perf
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=300, suite="stage-b-test-1-gpu-small-amd")
+
+
+class TestVLMServingAiter(CustomTestCase):
+    def test_vlm_serving_aiter(self):
+        check_vlm_serving_perf(
+            self,
+            "aiter",
+            output_throughput=2000,
+            e2e_ms=16500,
+            ttft_ms=150,
+            itl_ms=8,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_vlm_serving_fa3.py

@@ -0,0 +1,26 @@
+"""VLM serving perf on the fa3 attention backend."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.vlm_perf_kit import check_vlm_serving_perf
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=150, stage="extra-a", runner_config="1-gpu-large")
+
+
+class TestVLMServingFa3(CustomTestCase):
+    def test_vlm_serving_fa3(self):
+        check_vlm_serving_perf(
+            self,
+            "fa3",
+            # No offline bound: never measured on this lane.
+            output_throughput=15640,
+            e2e_ms=11000,
+            ttft_ms=100,
+            itl_ms=5.2,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/basic_perf/test_vlm_serving_flashinfer.py

@@ -0,0 +1,25 @@
+"""VLM serving perf on the flashinfer attention backend."""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.vlm_perf_kit import check_vlm_serving_perf
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=195, stage="extra-a", runner_config="1-gpu-small")
+
+
+class TestVLMServingFlashinfer(CustomTestCase):
+    def test_vlm_serving_flashinfer(self):
+        check_vlm_serving_perf(
+            self,
+            "flashinfer",
+            output_throughput=5940,
+            e2e_ms=17480,
+            ttft_ms=83,
+            itl_ms=8.4,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_rank_consensus_checker.py

@@ -23,7 +23,11 @@ from sglang.srt.utils.rank_consensus_checker import (
     shutdown,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase, find_available_port
+from sglang.test.test_utils import (
+    CustomTestCase,
+    find_available_port,
+    publish_build_topology,
+)
 
 register_cpu_ci(est_time=193, suite="stage-a-test-cpu-intel")
 
@@ -80,11 +84,8 @@ def run_distributed_test(
                 backend="gloo",
             )
 
-            initialize_model_parallel(
-                tensor_model_parallel_size=tp_size,
-                pipeline_model_parallel_size=pp_size,
-                backend="gloo",
-            )
+            publish_build_topology(tp_size=tp_size, pp_size=pp_size, world_rank=rank)
+            initialize_model_parallel(backend="gloo")
 
             fn()
         except Exception as e:

test/registered/dcp/test_dcp_layout_unit.py

@@ -134,7 +134,9 @@ class TestFilterDcpLocalChunkKvIndices(CustomTestCase):
 
     def test_identity_without_dcp(self):
         kv = torch.arange(37)
-        with rc.get_parallel().override(dcp_enabled=False, dcp_size=1, dcp_rank=0):
+        with rc.get_parallel().override(
+            dcp_enabled=False, dcp_size=1, dcp_rank=0, attn_dcp_rank=0
+        ):
             self.assertIs(
                 filter_dcp_local_chunk_kv_indices(
                     kv, torch.tensor([0]), torch.tensor([37])

test/registered/dcp/test_reduce_scatter_along_dim.py

@@ -25,6 +25,7 @@ import torch
 import torch.distributed as dist
 
 import sglang.srt.distributed.parallel_state as ps
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(
@@ -130,6 +131,7 @@ def init_distributed():
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=coord)
 
     cpu_group = coord.cpu_group
     nccl_group = coord.device_group

test/registered/disaggregation/test_disaggregation_basic.py

@@ -459,29 +459,11 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
             self._run_pause_on_decode_running_batch("retract", weight_update=True)
         )
 
-    async def _get_decode_num_running_reqs(self, session):
-        """Query current decode running_batch size from /v1/loads."""
-        async with session.get(
-            self.decode_url + "/v1/loads?include=core",
-            timeout=aiohttp.ClientTimeout(total=5),
-        ) as resp:
-            resp.raise_for_status()
-            body = await resp.json()
-            return sum(load["num_running_reqs"] for load in body["loads"])
-
-    async def _wait_for_decode_running_batch(self, session, timeout):
-        deadline = asyncio.get_running_loop().time() + timeout
-        while asyncio.get_running_loop().time() < deadline:
-            if await self._get_decode_num_running_reqs(session) > 0:
-                return
-            await asyncio.sleep(0.2)
-
-        self.fail("Timed out waiting for decode running_batch to become non-empty")
-
     async def _run_pause_on_decode_running_batch(self, mode, weight_update=False):
         num_requests = 2
         max_new_tokens = 512
         prompt = "Write a detailed numbered explanation of distributed inference. " * 12
+        decode_started = [asyncio.Event() for _ in range(num_requests)]
 
         async def _post(session, url, json_data, timeout=30):
             async with session.post(
@@ -493,20 +475,37 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
                 return await resp.json()
 
         async def _generate(session, request_id):
-            return await _post(
-                session,
+            async with session.post(
                 self.lb_url + "/generate",
-                {
+                json={
                     "text": f"Request {request_id}: {prompt}",
                     "background": True,
+                    "stream": True,
                     "sampling_params": {
                         "temperature": 0,
                         "ignore_eos": True,
                         "max_new_tokens": max_new_tokens,
                     },
                 },
-                timeout=180,
-            )
+                timeout=aiohttp.ClientTimeout(total=180),
+            ) as resp:
+                resp.raise_for_status()
+                response = None
+                async for line in resp.content:
+                    line = line.strip()
+                    if not line.startswith(b"data: "):
+                        continue
+                    data = line[len(b"data: ") :]
+                    if data == b"[DONE]":
+                        break
+                    response = json.loads(data)
+                    self.assertNotIn("error", response)
+                    # Prefill produces the first token. A later token proves this
+                    # request has reached running_batch on the decode worker.
+                    if response["meta_info"]["completion_tokens"] > 1:
+                        decode_started[request_id].set()
+                self.assertIsNotNone(response, "Generation stream returned no output")
+                return response
 
         async with aiohttp.ClientSession() as session:
             tasks = [
@@ -515,12 +514,17 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
             decode_paused = False
 
             try:
-                await self._wait_for_decode_running_batch(session, timeout=30)
-                await asyncio.sleep(0.1)
+                # /v1/loads can still report a previous batch. Wait for every
+                # current request to decode so none can arrive in the prealloc
+                # queue after the pause and prevent the weight-update flush.
+                await asyncio.wait_for(
+                    asyncio.gather(*(event.wait() for event in decode_started)),
+                    timeout=30,
+                )
 
                 self.assertTrue(
-                    any(not task.done() for task in tasks),
-                    "All requests finished before decode retract pause was issued.",
+                    all(not task.done() for task in tasks),
+                    "A request finished before decode retract pause was issued.",
                 )
 
                 await _post(
@@ -580,6 +584,9 @@ class TestDisaggregationPauseResumeDecodeRetract(PDDisaggregationServerBase):
             for response in responses:
                 self.assertIn("text", response)
                 self.assertGreater(len(response["text"]), 0)
+                self.assertEqual(
+                    response["meta_info"]["completion_tokens"], max_new_tokens
+                )
 
             self.assertGreater(
                 sum(

test/registered/disaggregation/test_disaggregation_optimistic_prefill.py

@@ -27,9 +27,10 @@ from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
-from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool
 from sglang.srt.mem_cache.radix_cache import RadixKey
+from sglang.srt.mem_cache.unified_cache.components.base import ComponentType
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils import get_device
@@ -354,16 +355,8 @@ class TestOptimisticPrefillMambaAdmission(CustomTestCase):
 
 
 class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
-    """Optimistic retry cleanup must not treat the donated Mamba checkpoint
-    as a second donation.
-
-    Bug mechanism: the retry path first inserts the unfinished prefix, which
-    donates the tracked checkpoint and clears ``mamba_last_track_seqlen``.
-    Releasing with ``is_insert=True`` afterwards re-enters the donation path
-    with the cleared marker, inserting a zero-length radix entry that pins a
-    clone of the request's live state. Releasing with ``is_insert=False``
-    retains the donated prefix/checkpoint and frees only the uncached tail
-    and the request-owned Mamba buffers.
+    """An optimistic-prefill retry leaves the donated prefix and exactly one
+    Mamba checkpoint in the tree -- not zero, and not a second pinned clone.
     """
 
     SIZE = 128
@@ -373,7 +366,7 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
 
     def _setup_mamba_tree(self):
         server_args = ServerArgs(model_path="dummy", page_size=1)
-        # MambaRadixCache reads mamba_cache_chunk_size, whose property
+        # The mamba component reads mamba_cache_chunk_size, whose property
         # otherwise loads the HF config for the dummy model.
         server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE
         set_global_server_args_for_scheduler(server_args)
@@ -427,13 +420,14 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
             kvcache=pool,
             need_sort=False,
         )
-        tree = MambaRadixCache(
+        tree = UnifiedRadixCache(
             params=CacheInitParams(
                 disable=False,
                 req_to_token_pool=req_to_token_pool,
                 token_to_kv_pool_allocator=allocator,
                 page_size=1,
                 enable_mamba_extra_buffer=True,
+                tree_components=(ComponentType.FULL, ComponentType.MAMBA),
             )
         )
         return tree, allocator, req_to_token_pool
@@ -459,7 +453,7 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
         req.kv.kv_committed_len = len(self.PROMPT)
         req.kv.kv_allocated_len = len(self.PROMPT)
         req.kv.mamba_last_track_seqlen = self.TRACK_SEQLEN
-        req.last_node = tree.root_node
+        req.last_node = tree.root_node_handle()
 
         scheduler = SimpleNamespace(
             tree_cache=tree,
@@ -478,13 +472,15 @@ class TestOptimisticPrefillMambaRetryRelease(CustomTestCase):
                 scheduler, req
             )
 
-        # The donated prefix and exactly one checkpoint stay in the tree; the
-        # old double-donation path either asserts or pins a second state.
         self.assertEqual(tree.total_size(), (self.TRACK_SEQLEN, 1))
         match = tree.match_prefix(
             MatchPrefixParams(key=RadixKey(array("q", self.PROMPT)))
         )
-        self.assertIsNotNone(match.last_device_node.mamba_value)
+        self.assertIsNotNone(
+            tree.tree_core.get_component_device_value(
+                match.last_device_node, ComponentType.MAMBA
+            )
+        )
 
         # Only the uncached tail and the request-owned Mamba buffers are
         # freed; the tree keeps the donated checkpoint slot.

test/registered/disaggregation/test_disaggregation_xpu.py

@@ -18,17 +18,16 @@ import unittest
 import requests
 import torch
 
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_xpu_ci
 from sglang.test.server_fixtures.disaggregation_fixture import (
     PDDisaggregationServerBase,
 )
 from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN
 
-register_cuda_ci(
+register_xpu_ci(
     est_time=300,
-    stage="base-a",
-    runner_config="1-gpu-small",
-    disabled="Intel XPU only — not available in standard CUDA CI",
+    suite="stage-b-test-1-gpu-xpu",
+    disabled="XPU CI image does not include an XPU-compatible NIXL/UCX build",
 )
 
 _XPU_AVAILABLE = torch.xpu.is_available()

test/registered/disaggregation/test_epd_disaggregation.py

@@ -5,6 +5,7 @@ import subprocess
 import threading
 import time
 import unittest
+from concurrent.futures import ThreadPoolExecutor
 
 import grpc
 import openai
@@ -1522,23 +1523,28 @@ class TestEPDDisaggregationGrpcEncoderOnly(PDDisaggregationServerBase):
         image_path = os.path.abspath("examples/assets/example_image.png")
 
         try:
-            stub.SchedulerReceiveUrl(
-                sglang_encoder_pb2.SchedulerReceiveUrlRequest(
-                    req_id=req_id,
-                    receive_url=f"{self.base_host}:{recv_port}",
-                    receive_count=1,
-                ),
-                timeout=60,
-            )
-            stub.Encode(
-                sglang_encoder_pb2.EncodeRequest(
-                    mm_items=[image_path],
-                    req_id=req_id,
-                    num_parts=1,
-                    part_idx=0,
-                ),
-                timeout=300,
-            )
+            # A scheduler registers concurrently with Encode, never before it:
+            # the request state only exists once Encode dispatches.
+            with ThreadPoolExecutor(max_workers=1) as pool:
+                registration = pool.submit(
+                    stub.SchedulerReceiveUrl,
+                    sglang_encoder_pb2.SchedulerReceiveUrlRequest(
+                        req_id=req_id,
+                        receive_url=f"{self.base_host}:{recv_port}",
+                        receive_count=1,
+                    ),
+                    timeout=60,
+                )
+                stub.Encode(
+                    sglang_encoder_pb2.EncodeRequest(
+                        mm_items=[image_path],
+                        req_id=req_id,
+                        num_parts=1,
+                        part_idx=0,
+                    ),
+                    timeout=300,
+                )
+                registration.result(timeout=60)
 
             poller = zmq.Poller()
             poller.register(recv_socket, zmq.POLLIN)

test/registered/disaggregation/test_kimi_linear_pd_dcp4.py

@@ -464,6 +464,29 @@ class TestKimiLinearPDDCP4(GSM8KMixin, PDDisaggregationServerBase):
                         f"niah prompt_tokens={LONG_CONTEXT_TOKENS} depth={needle_depth}"
                     ),
                 )
+                # Prefill now holds the long prefix. Decode must receive it
+                # again, even though prefill computes almost no new tokens.
+                # OSL > 1 checks that decode actually reads the transferred KV.
+                self._flush_cache(self.decode_url)
+                cached_actual = self._generate(
+                    self.base_url,
+                    prompt,
+                    max_new_tokens=16,
+                    ignore_eos=False,
+                )
+                self.assertGreater(
+                    cached_actual["meta_info"]["cached_tokens"],
+                    CHUNKED_PREFILL_SIZE,
+                )
+                self.assertIn(NIAH_KEY, cached_actual["text"])
+                self._assert_output_parity(
+                    reference,
+                    cached_actual,
+                    label=(
+                        f"cached niah prompt_tokens={LONG_CONTEXT_TOKENS} "
+                        f"depth={needle_depth}"
+                    ),
+                )
 
     def _assert_batch_completes(self, batch_size: int):
         response = requests.post(

test/registered/dp_attn/test_dp_attention.py

@@ -50,6 +50,8 @@ class TestDPAttentionDP2TP2(
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=[
                 "--trust-remote-code",
+                "--constrained-json-max-whitespace-cnt",
+                "4",
                 "--tp",
                 "2",
                 "--enable-dp-attention",

test/registered/e2e/hicache/test_hicache_unified_memory.py

@@ -0,0 +1,240 @@
+"""Compare unified-memory HiCache reloads against a resident-cache reference.
+
+Evict a target prefix with distinct filler requests, require a host hit on
+reload, and compare generated text and output logprobs. Both servers use the
+same unified-memory configuration to keep attention reduction order comparable.
+Covers GDN, SWA, tri-pool, and MLA layouts.
+"""
+
+import os
+import time
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=600, stage="extra-a", runner_config="2-gpu-large")
+
+_COMMON_ARGS = [
+    "--trust-remote-code",
+    "--enable-unified-memory",
+    "--enable-cache-report",
+    "--max-running-requests",
+    "1",
+    "--context-length",
+    "4096",
+]
+
+# Distinct filler prefixes must evict the target from device memory.
+_SMALL_POOL = ["--max-total-tokens", "8192"]
+
+_PREFIX = (
+    "The following is a detailed technical description of a distributed inference "
+    "system with paged attention, radix prefix caching and hierarchical offload. "
+) * 90
+_TARGET = _PREFIX + " Question one:"
+_CONTINUATION = _TARGET + " Explain how it works."
+
+
+def _generate(base_url, text, max_new_tokens=32, logprobs=True):
+    payload = {
+        "text": text,
+        "sampling_params": {"temperature": 0.0, "max_new_tokens": max_new_tokens},
+    }
+    if logprobs:
+        payload["return_logprob"] = True
+        # Output logprobs suffice; asking for prompt logprobs from zero
+        # caps the reusable prefix at zero and bypasses HiCache entirely.
+        payload["logprob_start_len"] = -1
+    resp = requests.post(f"{base_url}/generate", json=payload, timeout=600)
+    assert resp.status_code == 200, resp.text
+    data = resp.json()
+    lp = (
+        [t[0] for t in data["meta_info"]["output_token_logprobs"]] if logprobs else None
+    )
+    return data["text"], lp, data["meta_info"]
+
+
+class UnifiedMemoryHiCacheBase(CustomTestCase):
+    """Compare identical unified-memory configurations with and without HiCache."""
+
+    model: str = ""
+    extra_args: list = []
+    server_env: dict = {}
+
+    @classmethod
+    def setUpClass(cls):
+        if cls is UnifiedMemoryHiCacheBase:
+            raise unittest.SkipTest("base class")
+        base_args = _COMMON_ARGS + cls.extra_args
+        cls.hicache_url = "http://127.0.0.1:8157"
+        cls.reference_url = "http://127.0.0.1:8158"
+        env = {**os.environ, **cls.server_env} if cls.server_env else None
+        hicache_args = ["--enable-hierarchical-cache"]
+        if "--hicache-size" not in base_args:
+            hicache_args += ["--hicache-ratio", "4"]
+        cls.process_hicache = popen_launch_server(
+            cls.model,
+            cls.hicache_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=base_args + hicache_args,
+            env=env,
+        )
+        cls.addClassCleanup(kill_process_tree, cls.process_hicache.pid)
+        cls.process_reference = popen_launch_server(
+            cls.model,
+            cls.reference_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=base_args + ["--base-gpu-id", "1"],
+            env=env,
+        )
+        cls.addClassCleanup(kill_process_tree, cls.process_reference.pid)
+
+    def _force_host_round_trip(self):
+        """Evict the target off the device so the next hit must come from L2."""
+        for i in range(8):
+            _generate(
+                self.hicache_url,
+                f"Document {i}. "
+                + (f"Unique filler {i} about an unrelated subject. " * 300),
+                max_new_tokens=8,
+                logprobs=False,
+            )
+
+    def _flush_both(self):
+        """Reset cache state to match prefill boundaries and reduction order."""
+        for url in (self.hicache_url, self.reference_url):
+            requests.post(f"{url}/flush_cache", timeout=180)
+        time.sleep(3)
+
+    def test_load_back_matches_no_hicache(self):
+        """Host reloads preserve generated text and logprobs within tolerance."""
+        self._flush_both()
+        cold_text, cold_lp, _ = _generate(self.hicache_url, _TARGET)
+        ref_cold_text, ref_cold_lp, _ = _generate(self.reference_url, _TARGET)
+        self._force_host_round_trip()
+        # Extend the prefix so both servers compute new KV rows. Repeating it
+        # would let only the resident reference reuse its original final-token KV.
+        warm_text, warm_lp, warm_meta = _generate(self.hicache_url, _CONTINUATION)
+        ref_text, ref_lp, _ = _generate(self.reference_url, _CONTINUATION)
+
+        self.assertGreater(
+            (warm_meta.get("cached_tokens_details") or {}).get("host", 0),
+            0,
+            msg=f"Target did not reload from host: {warm_meta}",
+        )
+        self.assertEqual(cold_text, ref_cold_text)
+        self.assertEqual(warm_text, ref_text)
+        # Match the reference's prefill boundary in each comparison: cold
+        # against cold, and an L2 prefix hit against a resident prefix hit.
+        for label, lp, reference in (
+            ("cold", cold_lp, ref_cold_lp),
+            ("after-L2-reload", warm_lp, ref_lp),
+        ):
+            self.assertEqual(len(lp), len(reference))
+            delta = max(abs(a - b) for a, b in zip(lp, reference))
+            self.assertAlmostEqual(
+                delta,
+                0.0,
+                places=5,
+                msg=f"{label} diverged from the no-HiCache reference by {delta}",
+            )
+
+    def test_server_survives_the_round_trip(self):
+        """Cache churn must leave both schedulers healthy."""
+        self._force_host_round_trip()
+        for url in (self.hicache_url, self.reference_url):
+            resp = requests.get(f"{url}/health", timeout=30)
+            self.assertEqual(resp.status_code, 200)
+
+
+class TestUnifiedMemoryHiCacheGDN(UnifiedMemoryHiCacheBase):
+    """MHA full attention with envelope-strided gated-delta-net state."""
+
+    model = "yujiepan/qwen3.5-tiny-random"
+    extra_args = _SMALL_POOL + [
+        "--linear-attn-backend",
+        "triton",
+        "--mamba-backend",
+        "triton",
+        "--max-mamba-cache-size",
+        "8",
+        "--mem-fraction-static",
+        "0.6",
+    ]
+
+
+class TestUnifiedMemoryHiCacheSWA(UnifiedMemoryHiCacheBase):
+    """Hybrid SWA reloads bind pages to the full-attention pool's virtual IDs."""
+
+    model = "yujiepan/gemma-4e-tiny-random"
+    extra_args = _SMALL_POOL + [
+        "--attention-backend",
+        "triton",
+        "--mem-fraction-static",
+        "0.7",
+    ]
+
+
+class TestUnifiedMemoryHiCacheTriPool(UnifiedMemoryHiCacheBase):
+    """Full attention, sliding-window attention, and ShortConv state together."""
+
+    # The test revision is the reduced checkpoint used by Inkling CI.
+    model = "thinkingmachines/Inkling"
+    server_env = {"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}
+    extra_args = _SMALL_POOL + [
+        "--revision",
+        "test",
+        "--attention-backend",
+        "triton",
+        "--page-size",
+        "128",
+        "--mamba-radix-cache-strategy",
+        "extra_buffer",
+        "--swa-full-tokens-ratio",
+        "0.8",
+        "--max-mamba-cache-size",
+        "8",
+        "--mamba-full-memory-ratio",
+        "0.1",
+        "--mem-fraction-static",
+        "0.5",
+        "--cuda-graph-backend-prefill",
+        "disabled",
+        # Bound total host memory across all three component pools.
+        "--hicache-size",
+        "8",
+    ]
+
+
+class TestUnifiedMemoryHiCacheMLA(UnifiedMemoryHiCacheBase):
+    """MLA full attention with KDA state and MLA-specific transfer pointers."""
+
+    model = "yujiepan/kimi-linear-tiny-random"
+    extra_args = _SMALL_POOL + [
+        "--max-mamba-cache-size",
+        "8",
+        "--mem-fraction-static",
+        "0.5",
+        "--linear-attn-backend",
+        "triton",
+        "--mamba-backend",
+        "triton",
+        "--attention-backend",
+        "triton",
+        "--cuda-graph-backend-decode",
+        "disabled",
+        "--cuda-graph-backend-prefill",
+        "disabled",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/e2e/models/test_deepseek_v3_mtp.py

@@ -56,11 +56,7 @@ class TestDeepseekV3MTP(GSM8KMixin, DefaultServerBase):
                 f"{acc_length=:.2f}\n"
                 f"{speed=:.2f} token/s\n"
             )
-            self.assertGreater(acc_length, 2.8)
-            if is_in_amd_ci():
-                self.assertGreater(speed, 15)
-            else:
-                self.assertGreater(speed, 130)
+            self.assertGreater(speed, 130)
 
 
 if __name__ == "__main__":

test/registered/e2e/models/test_dsa_glm52_dp_mtp.py

@@ -10,7 +10,7 @@ from sglang.test.server_fixtures.dsa_mtp_fixture import (
 
 register_cuda_ci(
     est_time=400,
-    stage="base-c",
+    stage="nightly",
     runner_config="8-gpu-h200",
 )
 

test/registered/e2e/models/test_dsa_glm52_tp_mtp.py

@@ -10,7 +10,7 @@ from sglang.test.server_fixtures.dsa_mtp_fixture import (
 
 register_cuda_ci(
     est_time=400,
-    stage="base-c",
+    stage="nightly",
     runner_config="8-gpu-h200",
 )
 

test/registered/e2e/models/test_gpt_oss_4gpu_mxfp4.py

@@ -3,7 +3,6 @@ import unittest
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.gpt_oss_common import BaseTestGptOss
 
-register_cuda_ci(est_time=128, stage="base-c", runner_config="4-gpu-h100")
 register_cuda_ci(est_time=119, stage="base-c", runner_config="4-gpu-b200")
 
 

test/registered/e2e/models/test_kimi_k3_b300_low_latency.py

@@ -17,7 +17,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=1472, stage="base-c", runner_config="8-gpu-b300")
+register_cuda_ci(est_time=1472, stage="extra-b", runner_config="8-gpu-b300")
 
 MODEL_PATH = "moonshotai/Kimi-K3"
 DSPARK_DRAFT_MODEL = "RadixArk/Kimi-K3-DSpark"

test/registered/e2e/models/test_mimo_v2.py

@@ -5,7 +5,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
 from sglang.test.server_fixtures.mmmu_fixture import MMMUServerBase
 
-register_cuda_ci(est_time=317, stage="base-c", runner_config="8-gpu-h200")
+register_cuda_ci(est_time=317, stage="extra-b", runner_config="8-gpu-h200")
 
 MIMO_V2_MODEL = "XiaomiMiMo/MiMo-V2.5"
 MIMO_V2_OTHER_ARGS = [

test/registered/e2e/moe/test_glm53_flash_quark_moe_mi35x.py

@@ -0,0 +1,397 @@
+"""Isolated gfx950 numerical tests for GLM-5.3-Flash Quark MoE."""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+import torch.nn.functional as F
+from aiter.ops.flydsl.moe_common import GateMode
+from aiter.ops.shuffle import shuffle_weight
+from aiter.ops.triton.quant import dynamic_mxfp4_quant
+from aiter.utility.fp4_utils import e8m0_shuffle
+
+from sglang.srt.layers.moe.moe_runner.aiter import (
+    AiterMoeQuantInfo,
+    AiterQuantType,
+    AiterRunnerCore,
+    AiterRunnerInput,
+)
+from sglang.srt.layers.quantization.fp8_utils import dequant_mxfp4
+from sglang.srt.utils import is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=180, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
+    "requires one gfx950 GPU",
+)
+class TestGLM53FlashQuarkMoE(CustomTestCase):
+    hidden_size = 4096
+    intermediate_size = 2048
+    num_experts = 9
+    swiglu_limit = 10.0
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        torch.manual_seed(7)
+        cls.weights = cls._make_mxfp4_bank()
+        cls.weights["w13_deq"] = cls._dequant(
+            cls.weights["w13_raw"], cls.weights["s13_raw"]
+        )
+        cls.weights["w2_deq"] = cls._dequant(
+            cls.weights["w2_raw"], cls.weights["s2_raw"]
+        )
+        cls.runner = AiterRunnerCore(
+            SimpleNamespace(
+                no_combine=False,
+                activation="silu",
+                gemm1_alpha=None,
+                gemm1_clamp_limit=None,
+            )
+        )
+
+    @classmethod
+    def _make_mxfp4_bank(cls):
+        gate_weights = []
+        up_weights = []
+        down_weights = []
+        gate_scales = []
+        up_scales = []
+        down_scales = []
+        for expert in range(cls.num_experts):
+            generator = torch.Generator(device="cuda")
+            generator.manual_seed(100 + expert)
+            gate = (
+                torch.randn(
+                    cls.intermediate_size,
+                    cls.hidden_size,
+                    generator=generator,
+                    device="cuda",
+                    dtype=torch.bfloat16,
+                )
+                * 0.05
+            )
+            up = (
+                torch.randn(
+                    cls.intermediate_size,
+                    cls.hidden_size,
+                    generator=generator,
+                    device="cuda",
+                    dtype=torch.bfloat16,
+                )
+                * 0.05
+            )
+            down = (
+                torch.randn(
+                    cls.hidden_size,
+                    cls.intermediate_size,
+                    generator=generator,
+                    device="cuda",
+                    dtype=torch.bfloat16,
+                )
+                * 0.01
+            )
+            gate_q, gate_s = dynamic_mxfp4_quant(gate)
+            up_q, up_s = dynamic_mxfp4_quant(up)
+            down_q, down_s = dynamic_mxfp4_quant(down)
+            gate_weights.append(gate_q)
+            up_weights.append(up_q)
+            down_weights.append(down_q)
+            gate_scales.append(gate_s)
+            up_scales.append(up_s)
+            down_scales.append(down_s)
+
+        w13 = torch.cat([torch.stack(gate_weights), torch.stack(up_weights)], dim=1)
+        w2 = torch.stack(down_weights)
+        s13 = torch.cat([torch.stack(gate_scales), torch.stack(up_scales)], dim=1)
+        s2 = torch.stack(down_scales)
+        return {
+            "w13_raw": w13,
+            "w2_raw": w2,
+            "s13_raw": s13,
+            "s2_raw": s2,
+            "w13": shuffle_weight(w13.contiguous(), (16, 16)),
+            "w2": shuffle_weight(w2.contiguous(), (16, 16)),
+            "s13": e8m0_shuffle(s13.view(-1, s13.shape[-1])).view_as(s13),
+            "s2": e8m0_shuffle(s2.view(-1, s2.shape[-1])).view_as(s2),
+        }
+
+    @staticmethod
+    def _quantize_fp8_weight(weight):
+        rows, width = weight.shape
+        blocks = (
+            weight.float().view(rows // 128, 128, width // 128, 128).permute(0, 2, 1, 3)
+        )
+        scale = blocks.abs().amax(dim=(2, 3)).clamp(min=1e-12) / 448.0
+        quantized = (blocks / scale[:, :, None, None]).to(torch.float8_e4m3fn)
+        return (
+            quantized.permute(0, 2, 1, 3).reshape(rows, width),
+            scale,
+        )
+
+    @staticmethod
+    def _dequantize_fp8_weight(weight, scale):
+        return weight.float() * scale.repeat_interleave(128, dim=0).repeat_interleave(
+            128, dim=1
+        )
+
+    @staticmethod
+    def _quant_dequant_fp8_activation(activation):
+        tokens, width = activation.shape
+        groups = activation.float().view(tokens, width // 128, 128)
+        scale = groups.abs().amax(dim=-1).clamp(min=1e-12) / 448.0
+        quantized = (groups / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
+        return (quantized.float() * scale.unsqueeze(-1)).reshape(tokens, width)
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "weights"):
+            del cls.weights
+        if hasattr(cls, "runner"):
+            del cls.runner
+        torch.cuda.empty_cache()
+        super().tearDownClass()
+
+    @classmethod
+    def _dequant(cls, weight, scale):
+        experts, rows, packed = weight.shape
+        blocks = packed // 16
+        return dequant_mxfp4(
+            weight.view(experts, rows, blocks, 16),
+            scale,
+            torch.bfloat16,
+        )
+
+    @classmethod
+    def _quant_dequant_activation(cls, activation):
+        quantized, scale = dynamic_mxfp4_quant(activation)
+        tokens, packed = quantized.shape
+        blocks = packed // 16
+        return dequant_mxfp4(
+            quantized.view(1, tokens, blocks, 16),
+            scale.view(1, tokens, blocks),
+            torch.bfloat16,
+        ).squeeze(0)
+
+    @classmethod
+    def _torch_oracle(cls, hidden_states, topk_ids, topk_weights):
+        w13 = cls.weights["w13_deq"]
+        w2 = cls.weights["w2_deq"]
+        output = torch.zeros_like(hidden_states)
+        hidden_qdq = cls._quant_dequant_activation(hidden_states)
+        for token in range(hidden_states.shape[0]):
+            for route in range(topk_ids.shape[1]):
+                expert = int(topk_ids[token, route])
+                gate = F.linear(
+                    hidden_qdq[token].float(),
+                    w13[expert, : cls.intermediate_size].float(),
+                )
+                up = F.linear(
+                    hidden_qdq[token].float(),
+                    w13[expert, cls.intermediate_size :].float(),
+                )
+                gate = gate.clamp(max=cls.swiglu_limit)
+                up = up.clamp(min=-cls.swiglu_limit, max=cls.swiglu_limit)
+                activated = F.silu(gate) * up
+                activated = cls._quant_dequant_activation(
+                    activated.unsqueeze(0).bfloat16()
+                ).squeeze(0)
+                expert_output = F.linear(activated.float(), w2[expert].float())
+                output[token] += (expert_output * topk_weights[token, route]).to(
+                    output.dtype
+                )
+        return output
+
+    @classmethod
+    def _aiter(cls, hidden_states, topk_ids, topk_weights):
+        w13 = cls.weights["w13"].view(torch.float4_e2m1fn_x2)
+        w2 = cls.weights["w2"].view(torch.float4_e2m1fn_x2)
+        w13.is_shuffled = True
+        w2.is_shuffled = True
+        quant_info = AiterMoeQuantInfo(
+            w13_weight=w13,
+            w2_weight=w2,
+            quant_type=AiterQuantType.PER_1X32,
+            w13_scale=cls.weights["s13"],
+            w2_scale=cls.weights["s2"],
+            swiglu_limit=cls.swiglu_limit,
+            fused_moe_kwargs={"gate_mode": GateMode.SEPARATED.value},
+        )
+        runner_input = AiterRunnerInput(
+            hidden_states=hidden_states,
+            topk_ids=topk_ids.to(torch.int32),
+            topk_weights=topk_weights.to(torch.float32),
+            quant_type=AiterQuantType.PER_1X32,
+        )
+        return cls.runner.run(runner_input, quant_info, {}).hidden_states
+
+    def _assert_numerics(self, actual, expected, max_abs=None):
+        self.assertTrue(torch.isfinite(actual).all())
+        actual_float = actual.float()
+        expected_float = expected.float()
+        cosine = F.cosine_similarity(
+            actual_float.flatten().unsqueeze(0),
+            expected_float.flatten().unsqueeze(0),
+        ).item()
+        self.assertGreater(cosine, 0.98)
+        relative_l2 = (
+            torch.linalg.vector_norm(actual_float - expected_float)
+            / torch.linalg.vector_norm(expected_float).clamp(min=1e-12)
+        ).item()
+        self.assertLess(relative_l2, 0.20)
+        if max_abs is not None:
+            self.assertLess(
+                (actual_float - expected_float).abs().max().item(),
+                max_abs,
+            )
+
+    def test_top1_and_top8_match_dequantized_oracle(self):
+        for tokens in (1, 8, 17, 32, 64, 128):
+            generator = torch.Generator(device="cuda")
+            generator.manual_seed(tokens)
+            hidden = (
+                torch.randn(
+                    tokens,
+                    self.hidden_size,
+                    generator=generator,
+                    device="cuda",
+                    dtype=torch.bfloat16,
+                )
+                * 0.5
+            )
+            # topk=9 models eight routed experts plus one fused shared slot.
+            for topk in (1, 8, 9):
+                with self.subTest(tokens=tokens, topk=topk):
+                    ids = torch.arange(topk, device="cuda", dtype=torch.int64).repeat(
+                        tokens, 1
+                    )
+                    weights = torch.rand(
+                        tokens,
+                        topk,
+                        generator=generator,
+                        device="cuda",
+                        dtype=torch.float32,
+                    )
+                    if topk > 1:
+                        weights /= weights.sum(dim=-1, keepdim=True)
+                    expected = self._torch_oracle(hidden, ids, weights)
+                    actual = self._aiter(hidden, ids, weights)
+                    repeated = self._aiter(hidden, ids, weights)
+                    self._assert_numerics(actual, expected, max_abs=0.75)
+                    if topk == 1:
+                        torch.testing.assert_close(actual, repeated, atol=0, rtol=0)
+                    else:
+                        # Stage-2 combines top-k routes with atomics; reduction
+                        # order may differ while remaining BF16-equivalent.
+                        torch.testing.assert_close(
+                            actual, repeated, atol=2e-2, rtol=1e-2
+                        )
+
+    def test_clamp_boundary(self):
+        hidden = torch.full(
+            (1, self.hidden_size),
+            4.0,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        ids = torch.tensor([[0]], device="cuda", dtype=torch.int64)
+        weights = torch.ones((1, 1), device="cuda", dtype=torch.float32)
+        expected = self._torch_oracle(hidden, ids, weights)
+        actual = self._aiter(hidden, ids, weights)
+        self._assert_numerics(actual, expected)
+
+    def test_plain_block_fp8_matches_separated_oracle(self):
+        generator = torch.Generator(device="cuda")
+        generator.manual_seed(1234)
+        gate = (
+            torch.randn(
+                self.intermediate_size,
+                self.hidden_size,
+                generator=generator,
+                device="cuda",
+                dtype=torch.bfloat16,
+            )
+            * 0.05
+        )
+        up = (
+            torch.randn(
+                self.intermediate_size,
+                self.hidden_size,
+                generator=generator,
+                device="cuda",
+                dtype=torch.bfloat16,
+            )
+            * 0.05
+        )
+        down = (
+            torch.randn(
+                self.hidden_size,
+                self.intermediate_size,
+                generator=generator,
+                device="cuda",
+                dtype=torch.bfloat16,
+            )
+            * 0.01
+        )
+        gate_q, gate_s = self._quantize_fp8_weight(gate)
+        up_q, up_s = self._quantize_fp8_weight(up)
+        down_q, down_s = self._quantize_fp8_weight(down)
+        w13_raw = torch.cat([gate_q, up_q], dim=0).unsqueeze(0)
+        w13_scale = torch.cat([gate_s, up_s], dim=0).unsqueeze(0)
+        w2_raw = down_q.unsqueeze(0)
+        w2_scale = down_s.unsqueeze(0)
+        w13 = shuffle_weight(w13_raw.contiguous(), (16, 16))
+        w2 = shuffle_weight(w2_raw.contiguous(), (16, 16))
+        quant_info = AiterMoeQuantInfo(
+            w13_weight=w13,
+            w2_weight=w2,
+            quant_type=AiterQuantType.PER_128X128,
+            w13_scale=w13_scale,
+            w2_scale=w2_scale,
+            swiglu_limit=self.swiglu_limit,
+            fused_moe_kwargs={"gate_mode": GateMode.SEPARATED.value},
+        )
+
+        gate_deq = self._dequantize_fp8_weight(gate_q, gate_s)
+        up_deq = self._dequantize_fp8_weight(up_q, up_s)
+        down_deq = self._dequantize_fp8_weight(down_q, down_s)
+        for tokens in (1, 8, 32):
+            with self.subTest(tokens=tokens):
+                hidden = (
+                    torch.randn(
+                        tokens,
+                        self.hidden_size,
+                        generator=generator,
+                        device="cuda",
+                        dtype=torch.bfloat16,
+                    )
+                    * 0.5
+                )
+                hidden_qdq = self._quant_dequant_fp8_activation(hidden)
+                gate_out = F.linear(hidden_qdq, gate_deq).clamp(max=self.swiglu_limit)
+                up_out = F.linear(hidden_qdq, up_deq).clamp(
+                    -self.swiglu_limit, self.swiglu_limit
+                )
+                activated = self._quant_dequant_fp8_activation(
+                    (F.silu(gate_out) * up_out).bfloat16()
+                )
+                expected = F.linear(activated, down_deq).bfloat16()
+
+                runner_input = AiterRunnerInput(
+                    hidden_states=hidden,
+                    topk_ids=torch.zeros((tokens, 1), device="cuda", dtype=torch.int32),
+                    topk_weights=torch.ones(
+                        (tokens, 1), device="cuda", dtype=torch.float32
+                    ),
+                    quant_type=AiterQuantType.PER_128X128,
+                )
+                actual = self.runner.run(runner_input, quant_info, {}).hidden_states
+                self._assert_numerics(actual, expected, max_abs=0.75)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/e2e/pp/test_pp_spec_embed_scan.py

@@ -10,7 +10,11 @@ import torch
 from transformers import MistralConfig, PretrainedConfig
 
 from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    publish_build_topology,
+)
 
 register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small")
 
@@ -77,7 +81,8 @@ class TestDraftEmbedScan(CustomTestCase):
         init_distributed_environment(
             world_size=1, rank=0, local_rank=0, distributed_init_method="env://"
         )
-        initialize_model_parallel(tensor_model_parallel_size=1)
+        publish_build_topology(tp_size=1)
+        initialize_model_parallel()
         torch.set_default_dtype(torch.bfloat16)
         torch.cuda.set_device(0)
 

test/registered/eplb/test_lplb_distributed.py

@@ -38,6 +38,7 @@ from sglang.srt.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import publish_build_topology
 
 register_cuda_ci(est_time=18, stage="base-b", runner_config="2-gpu-large")
 
@@ -238,10 +239,10 @@ def _worker_main(local_rank: int, world_size: int):
     init_distributed_environment(
         world_size=world_size, rank=local_rank, local_rank=local_rank
     )
-    initialize_model_parallel(
-        tensor_model_parallel_size=world_size,
-        expert_model_parallel_size=world_size,
+    publish_build_topology(
+        tp_size=world_size, ep_size=world_size, world_rank=local_rank
     )
+    initialize_model_parallel()
 
     from sglang.srt.eplb.lplb_solver import clear_global_lplb_solvers
 

test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py

@@ -14,6 +14,7 @@ from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import get_benchmark_range, multigpu_bench_main
 from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl
 from sglang.kernels.ops.communication.mp import register_comm_cleanup
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(
@@ -61,6 +62,7 @@ def _init_cpu_group() -> dist.ProcessGroup:
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=coord)
     atexit.register(dist.destroy_process_group)
     torch.cuda.set_stream(torch.cuda.Stream())
     return coord.cpu_group

test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py

@@ -33,6 +33,7 @@ from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
     all_gather_inner,
     create_state,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(
@@ -77,6 +78,7 @@ def _init_cpu_group() -> dist.ProcessGroup:
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=ps._WORLD)
     atexit.register(dist.destroy_process_group)
     logging.disable(logging.INFO)
     torch.cuda.set_stream(torch.cuda.Stream())

test/registered/kernels/benchmark/communication/bench_tp_qknorm.py

@@ -36,6 +36,7 @@ from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(
@@ -108,6 +109,7 @@ def _init_cpu_group() -> dist.ProcessGroup:
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=coord)
     atexit.register(dist.destroy_process_group)
     logging.disable(logging.INFO)
     torch.cuda.set_stream(torch.cuda.Stream())

test/registered/kernels/benchmark/kvcache/bench_hicache.py

@@ -24,8 +24,11 @@ from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import get_benchmark_range
 from sglang.kernels.ops.kvcache.hicache import (
-    transfer_hicache_all_layer,
-    transfer_hicache_one_layer,
+    DEFAULT_BLOCK_QUOTA,
+    TMA_BLOCK_QUOTA,
+    _default_unroll,
+    _jit_hicache_module,
+    _jit_hicache_tma_module,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
@@ -108,15 +111,19 @@ def sglang_jit_transfer_one(
     indices_src: torch.Tensor,
     element_dim: int,
 ) -> None:
-    """SGL JIT Kernel for single layer transfer."""
-    transfer_hicache_one_layer(
-        k_cache_dst,
-        v_cache_dst,
+    """SGL JIT register kernel for single layer transfer (bypasses TMA routing)."""
+    element_size = element_dim * k_cache_dst.element_size()
+    _jit_hicache_module(
+        element_size=element_size,
+        unroll=_default_unroll(element_size),
+        block_quota=DEFAULT_BLOCK_QUOTA,
+    ).launch_one(
+        k_cache_dst.view(-1, element_dim),
+        v_cache_dst.view(-1, element_dim),
         indices_dst,
-        k_cache_src,
-        v_cache_src,
+        k_cache_src.view(-1, element_dim),
+        v_cache_src.view(-1, element_dim),
         indices_src,
-        element_dim=element_dim,
     )
 
 
@@ -153,17 +160,58 @@ def sglang_jit_transfer_all(
     stride_bytes: int,
     element_size: int,
 ) -> None:
-    """SGL JIT Kernel for all layer transfer."""
-    transfer_hicache_all_layer(
+    """SGL JIT register kernel for all layer transfer (bypasses TMA routing)."""
+    _jit_hicache_module(
+        element_size=element_size,
+        unroll=_default_unroll(element_size),
+        block_quota=DEFAULT_BLOCK_QUOTA,
+    ).launch_all(
         k_ptrs_dst,
         v_ptrs_dst,
         indices_dst,
         k_ptrs_src,
         v_ptrs_src,
         indices_src,
-        kv_cache_src_stride_bytes=stride_bytes,
-        kv_cache_dst_stride_bytes=stride_bytes,
-        element_size=element_size,
+        stride_bytes,
+        stride_bytes,
+    )
+
+
+def sglang_tma_transfer_one(
+    k_cache_dst: torch.Tensor,
+    v_cache_dst: torch.Tensor,
+    indices_dst: torch.Tensor,
+    k_cache_src: torch.Tensor,
+    v_cache_src: torch.Tensor,
+    indices_src: torch.Tensor,
+) -> None:
+    """SGL TMA staging kernel for single layer transfer."""
+    _jit_hicache_tma_module(block_quota=TMA_BLOCK_QUOTA).launch_one(
+        k_cache_dst, v_cache_dst, indices_dst, k_cache_src, v_cache_src, indices_src
+    )
+
+
+def sglang_tma_transfer_all(
+    k_ptrs_dst: torch.Tensor,
+    v_ptrs_dst: torch.Tensor,
+    indices_dst: torch.Tensor,
+    k_ptrs_src: torch.Tensor,
+    v_ptrs_src: torch.Tensor,
+    indices_src: torch.Tensor,
+    stride_bytes: int,
+    element_size: int,
+) -> None:
+    """SGL TMA staging kernel for all layer transfer."""
+    _jit_hicache_tma_module(block_quota=TMA_BLOCK_QUOTA).launch_all(
+        k_ptrs_dst,
+        v_ptrs_dst,
+        indices_dst,
+        k_ptrs_src,
+        v_ptrs_src,
+        indices_src,
+        stride_bytes,
+        stride_bytes,
+        element_size,
     )
 
 
@@ -191,6 +239,13 @@ ELEMENT_SIZE_RANGE = get_benchmark_range(
 LINE_VALS = ["aot", "jit", "torch"]
 if DISABLE_TORCH:
     LINE_VALS.remove("torch")
+# The TMA staging kernel needs sm_90+ (cp.async.bulk); skip the line elsewhere.
+if (
+    torch.cuda.is_available()
+    and torch.version.hip is None
+    and torch.cuda.get_device_capability()[0] >= 9
+):
+    LINE_VALS.insert(2, "tma")
 
 
 # =============================================================================
@@ -246,6 +301,17 @@ def benchmark_one_layer_h2d(element_size: int, batch_size: int, provider: str):
             )
             for i in range(NUM_LAYERS)
         ],
+        "tma": lambda: [
+            sglang_tma_transfer_one(
+                k_cache_dst[i],
+                v_cache_dst[i],
+                indices_dst_gpu,
+                k_cache_src[i],
+                v_cache_src[i],
+                indices_src_gpu,
+            )
+            for i in range(NUM_LAYERS)
+        ],
         "torch": lambda: [
             pytorch_transfer(
                 k_cache_dst[i],
@@ -329,6 +395,16 @@ def benchmark_all_layer_d2h(element_size: int, batch_size: int, provider: str):
             element_bytes,
             element_bytes,
         ),
+        "tma": lambda: sglang_tma_transfer_all(
+            k_ptrs_dst,
+            v_ptrs_dst,
+            indices_dst_gpu,
+            k_ptrs_src,
+            v_ptrs_src,
+            indices_src_gpu,
+            element_bytes,
+            element_bytes,
+        ),
         "torch": lambda: [
             pytorch_transfer(
                 k_caches_dst[i],

test/registered/kernels/ops/attention/test_dense_prefill_indexer.py

@@ -0,0 +1,316 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4.fp4_indexer import quantize_fp4_indexer_tensor
+from sglang.srt.layers.attention.dsv4 import dense_prefill_indexer
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+def make_inputs(request_lengths, ratio=1, zero_queries=False, seed=17):
+    torch.manual_seed(seed)
+    rows = sum(q for q, _ in request_lengths)
+    q = torch.randn((rows, 32, 128), dtype=torch.bfloat16, device="cuda")
+    if zero_queries:
+        q.zero_()
+    packed, scales = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True)
+    kv = quantize_fp4_indexer_tensor(
+        torch.randn(
+            (sum(n for _, n in request_lengths), 128),
+            dtype=torch.bfloat16,
+            device="cuda",
+        ),
+        rne=True,
+    )
+    starts, lengths = [], []
+    start = 0
+    for queries, context in request_lengths:
+        starts.extend([start] * queries)
+        lengths.extend(
+            (position + 1) // ratio
+            for position in range(context * ratio - queries, context * ratio)
+        )
+        start += context
+    return dict(
+        q=(packed.view(rows, 32, 64), scales.view(rows, 32)),
+        kv=kv,
+        weights=torch.rand((rows, 32), dtype=torch.float32, device="cuda"),
+        starts=torch.tensor(starts, dtype=torch.int32, device="cuda"),
+        lengths=torch.tensor(lengths, dtype=torch.int32, device="cuda"),
+        request_lengths=request_lengths,
+        topk=512,
+        candidate_topk_blocks=2,
+        candidate_block_size=8,
+    )
+
+
+def dense_scores(inputs):
+    from deep_gemm import fp8_fp4_mqa_logits
+
+    width = (max(n for _, n in inputs["request_lengths"]) + 3) // 4 * 4
+    scores = fp8_fp4_mqa_logits(
+        inputs["q"],
+        inputs["kv"],
+        inputs["weights"],
+        inputs["starts"],
+        inputs["starts"] + inputs["lengths"],
+        False,
+        width,
+    )
+    return scores.masked_fill_(
+        torch.arange(width, device="cuda")[None, :] >= inputs["lengths"][:, None],
+        -torch.inf,
+    )
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10,
+    "requires SM100",
+)
+class TestDensePrefillIndexer(CustomTestCase):
+    def assert_topk(self, inputs, selected, scores):
+        columns = (selected - inputs["starts"][:, None]).long()
+        valid = selected >= 0
+        expected_count = torch.isfinite(scores).sum(-1).clamp_max(inputs["topk"])
+        torch.testing.assert_close(valid.sum(-1), expected_count)
+        self.assertTrue(
+            (~valid | ((columns >= 0) & (columns < inputs["lengths"][:, None]))).all()
+        )
+        actual = scores.gather(1, columns.clamp(0, scores.shape[1] - 1)).masked_fill(
+            ~valid, -torch.inf
+        )
+        expected = scores.topk(min(inputs["topk"], scores.shape[1]), dim=-1).values
+        expected = torch.nn.functional.pad(
+            expected, (0, inputs["topk"] - expected.shape[1]), value=-torch.inf
+        )
+        torch.testing.assert_close(
+            actual.sort(descending=True).values, expected, rtol=1e-5, atol=1e-5
+        )
+        ordered = (
+            columns.masked_fill(~valid, torch.iinfo(torch.int64).max).sort().values
+        )
+        self.assertTrue(
+            (
+                (ordered[:, 1:] != ordered[:, :-1])
+                | (ordered[:, 1:] == torch.iinfo(torch.int64).max)
+            ).all()
+        )
+
+    def test_ragged_source_consumer_and_replay(self):
+        for ratio in (1, 2):
+            for zero_queries in (False, True):
+                with self.subTest(ratio=ratio, zero_queries=zero_queries):
+                    inputs = make_inputs(
+                        [(0, 0), (1, 1), (33, 511), (257, 4097)],
+                        ratio=ratio,
+                        zero_queries=zero_queries,
+                    )
+                    inputs["candidate_topk_blocks"] = 128
+                    scores = dense_scores(inputs)
+                    consumer_inputs = make_inputs(
+                        inputs["request_lengths"],
+                        ratio=ratio,
+                        zero_queries=zero_queries,
+                        seed=29,
+                    )
+                    consumer_inputs["kv"] = inputs["kv"]
+                    consumer_inputs["candidate_topk_blocks"] = 128
+                    consumer_scores = dense_scores(consumer_inputs)
+                    with patch.object(
+                        dense_prefill_indexer, "_SCORE_BUDGET_BYTES", 128 << 10
+                    ):
+                        selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                            **inputs, publish_candidates=True, candidates=None
+                        )
+                        self.assert_topk(inputs, selected, scores)
+                        row = 0
+                        for (queries, context), blocks in zip(
+                            inputs["request_lengths"], candidates.request_blocks
+                        ):
+                            local = scores[row : row + queries, :context]
+                            if queries and context:
+                                padded = torch.nn.functional.pad(
+                                    local, (0, -context % 8), value=-torch.inf
+                                )
+                                block_scores = padded.unflatten(-1, (-1, 8)).amax(-1)
+                                last = (inputs["lengths"][row : row + queries] - 1) // 8
+                                block_scores.masked_fill_(
+                                    torch.arange(block_scores.shape[1], device="cuda")[
+                                        None, :
+                                    ]
+                                    == last[:, None],
+                                    torch.inf,
+                                )
+                                chosen_scores = block_scores.gather(
+                                    1, blocks.long().clamp_min(0)
+                                ).masked_fill(blocks < 0, -torch.inf)
+                                torch.testing.assert_close(
+                                    chosen_scores.sort(descending=True).values,
+                                    block_scores.topk(blocks.shape[1]).values,
+                                    rtol=1e-5,
+                                    atol=1e-5,
+                                )
+                                columns = torch.arange(context, device="cuda")
+                                member = (
+                                    columns[None, :, None] // 8 == blocks[:, None, :]
+                                ).any(-1)
+                                consumer_scores[
+                                    row : row + queries, :context
+                                ].masked_fill_(
+                                    ~member,
+                                    -torch.inf,
+                                )
+                            row += queries
+                        self.assertGreater(
+                            torch.isfinite(consumer_scores[-1]).sum().item(),
+                            inputs["topk"],
+                        )
+                        self.assertLess(
+                            torch.isfinite(consumer_scores[-1]).sum().item(),
+                            inputs["lengths"][-1].item(),
+                        )
+                        selected, published = dense_prefill_indexer.dense_prefill_topk(
+                            **consumer_inputs,
+                            publish_candidates=False,
+                            candidates=candidates,
+                        )
+                        self.assertIsNone(published)
+                        self.assert_topk(consumer_inputs, selected, consumer_scores)
+                        tail_lengths = [0, 0, 7, 31]
+                        rows, row = [], 0
+                        for (queries, _), tail in zip(
+                            inputs["request_lengths"], tail_lengths
+                        ):
+                            rows.extend(range(row + queries - tail, row + queries))
+                            row += queries
+                        rows = torch.tensor(rows, dtype=torch.int64, device="cuda")
+                        tail_inputs = dict(
+                            consumer_inputs,
+                            q=tuple(t[rows] for t in consumer_inputs["q"]),
+                            weights=consumer_inputs["weights"][rows],
+                            starts=inputs["starts"][rows],
+                            lengths=inputs["lengths"][rows],
+                            request_lengths=list(
+                                zip(
+                                    tail_lengths,
+                                    [n for _, n in inputs["request_lengths"]],
+                                )
+                            ),
+                        )
+                        selected, _ = dense_prefill_indexer.dense_prefill_topk(
+                            **tail_inputs,
+                            publish_candidates=False,
+                            candidates=candidates.tail(tail_lengths),
+                        )
+                        self.assert_topk(tail_inputs, selected, consumer_scores[rows])
+
+    def test_unfiltered_and_zero_length_requests(self):
+        for request_lengths in ([(257, 8192)], [(1, 0), (1, 1), (0, 7)]):
+            for publish in (False, True):
+                with self.subTest(request_lengths=request_lengths, publish=publish):
+                    inputs = make_inputs(request_lengths)
+                    selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                        **inputs, publish_candidates=publish, candidates=None
+                    )
+                    self.assert_topk(inputs, selected, dense_scores(inputs))
+                    if publish:
+                        self.assertEqual(
+                            [b.shape[0] for b in candidates.request_blocks],
+                            [q for q, _ in request_lengths],
+                        )
+
+    def test_empty_queries_or_context(self):
+        for request_lengths, shape, block_shapes in (
+            ([], (0, 512), []),
+            ([(0, 0)], (0, 512), [(0, 0)]),
+            ([(0, 0), (0, 17)], (0, 512), [(0, 0), (0, 2)]),
+            ([(1, 0)], (1, 512), [(1, 0)]),
+        ):
+            with self.subTest(request_lengths=request_lengths):
+                inputs = make_inputs(request_lengths)
+                selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                    **inputs, publish_candidates=True, candidates=None
+                )
+                torch.testing.assert_close(
+                    selected, torch.full(shape, -1, dtype=torch.int32, device="cuda")
+                )
+                self.assertEqual(
+                    [tuple(b.shape) for b in candidates.request_blocks], block_shapes
+                )
+
+    def test_score_budget_includes_allocation_padding(self):
+        from deep_gemm import fp8_fp4_mqa_logits
+
+        inputs = make_inputs([(13, 257)])
+        expected = dense_scores(inputs)
+        for budget in (32 << 10, 28 << 10, 8 << 10):
+            with self.subTest(budget=budget):
+                allocations = []
+
+                def checked_logits(*args, **kwargs):
+                    before = torch.cuda.memory_allocated()
+                    logits = fp8_fp4_mqa_logits(*args, **kwargs)
+                    allocations.append(torch.cuda.memory_allocated() - before)
+                    return logits
+
+                with (
+                    patch.object(dense_prefill_indexer, "_SCORE_BUDGET_BYTES", budget),
+                    patch("deep_gemm.fp8_fp4_mqa_logits", new=checked_logits),
+                ):
+                    selected, _ = dense_prefill_indexer.dense_prefill_topk(
+                        **inputs, publish_candidates=True, candidates=None
+                    )
+                self.assertTrue(allocations)
+                self.assertLessEqual(max(allocations), budget)
+                self.assert_topk(inputs, selected, expected)
+
+    def test_score_memory_is_bounded(self):
+        for context, limit_gib in ((65536, 3), (65535, 5)):
+            with self.subTest(context=context):
+                inputs = make_inputs([(16384, context)])
+                inputs["candidate_topk_blocks"] = 2048
+                torch.cuda.synchronize()
+                torch.cuda.reset_peak_memory_stats()
+                baseline = torch.cuda.memory_allocated()
+                selected, candidates = dense_prefill_indexer.dense_prefill_topk(
+                    **inputs, publish_candidates=True, candidates=None
+                )
+                torch.cuda.synchronize()
+                self.assertLess(
+                    torch.cuda.max_memory_allocated() - baseline, limit_gib << 30
+                )
+                self.assertEqual(tuple(selected.shape), (16384, 512))
+                self.assertEqual(
+                    tuple(candidates.request_blocks[0].shape), (16384, 2048)
+                )
+                del selected
+                selected, published = dense_prefill_indexer.dense_prefill_topk(
+                    **inputs, publish_candidates=False, candidates=candidates
+                )
+                self.assertIsNone(published)
+                del selected
+                torch.cuda.synchronize()
+                baseline = torch.cuda.memory_allocated()
+                for _ in range(3):
+                    torch.cuda.reset_peak_memory_stats()
+                    selected, published = dense_prefill_indexer.dense_prefill_topk(
+                        **inputs, publish_candidates=False, candidates=candidates
+                    )
+                    torch.cuda.synchronize()
+                    self.assertIsNone(published)
+                    self.assertLess(
+                        torch.cuda.max_memory_allocated() - baseline, 4 << 30
+                    )
+                    self.assertEqual(tuple(selected.shape), (16384, 512))
+                    del selected
+                    torch.cuda.synchronize()
+                    self.assertEqual(torch.cuda.memory_allocated(), baseline)
+                del inputs, candidates
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_dsa_metadata.py

@@ -10,7 +10,7 @@ from sglang.kernels.ops.attention.dsa_metadata import (
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_amd_ci(est_time=15, stage="stage-b", runner_config="1-gpu-large-amd")
 
 

test/registered/kernels/ops/attention/test_dsa_multi_ctas_counter.py

@@ -0,0 +1,139 @@
+"""Lifetime of the DSA multi-CTAs KV counter across CUDA graph capture.
+
+The decode graphs record this buffer's address, so nothing after capture may
+reallocate it. _forward_trtllm, the production caller, needs a live FlashInfer
+kernel and is not covered here.
+"""
+
+import unittest
+import weakref
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
+from sglang.srt.layers.attention.trtllm_mla_backend import (
+    TRTLLM_MLA_MAX_BATCH_SIZE,
+    grow_multi_ctas_kv_counter_buffer_if_needed,
+    make_persistent_multi_ctas_kv_counter_buffer,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+_NUM_Q_HEADS = 128
+_MAX_CTX_LEN = 64
+# 1024 captured requests x 9 draft tokens: a supported speculative capture whose
+# query-row count exceeds TRTLLM_MLA_MAX_BATCH_SIZE.
+_CAPTURED_BS = 1024
+_NUM_DRAFT_TOKENS = 9
+_NUM_CAPTURED_ROWS = _CAPTURED_BS * _NUM_DRAFT_TOKENS
+
+
+def _make_backend(*, allocate_counter: bool = True):
+    backend = object.__new__(DeepseekSparseAttnBackend)
+    backend.device = "cuda"
+    backend.num_q_heads = _NUM_Q_HEADS
+    backend.real_page_size = 64
+    backend.hisparse_coordinator = None
+    backend.speculative_num_draft_tokens = _NUM_DRAFT_TOKENS
+    backend.dsa_index_kpool = 1
+    backend.use_fused_topk = False
+    backend.dsa_topk_backend = SimpleNamespace(should_use_topk_v2=lambda: False)
+    backend.dsa_index_topk = 2048
+    backend.dsa_decode_impl = "trtllm"
+    backend.req_to_token = torch.zeros(
+        8, _MAX_CTX_LEN, dtype=torch.int32, device="cuda"
+    )
+    backend._multi_ctas_kv_counter_buffer = (
+        make_persistent_multi_ctas_kv_counter_buffer(
+            device=torch.device("cuda"),
+            num_q_heads=_NUM_Q_HEADS,
+            max_batch_size=48,
+        )
+        if allocate_counter
+        else None
+    )
+    return backend
+
+
+def _would_grow(backend, rows: int) -> bool:
+    return (
+        grow_multi_ctas_kv_counter_buffer_if_needed(
+            buffer=backend._multi_ctas_kv_counter_buffer,
+            device=torch.device("cuda"),
+            num_q_heads=backend.num_q_heads,
+            batch_size=rows,
+        )
+        is not backend._multi_ctas_kv_counter_buffer
+    )
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "needs a CUDA device")
+class TestMultiCtasKvCounterLifetime(CustomTestCase):
+    def test_request_sized_counter_would_grow_at_capture(self):
+        """The premise: sizing by requests undercounts captured query rows."""
+        self.assertGreater(_NUM_CAPTURED_ROWS, TRTLLM_MLA_MAX_BATCH_SIZE)
+        self.assertTrue(_would_grow(_make_backend(), _NUM_CAPTURED_ROWS))
+
+    def test_init_cuda_graph_state_sizes_for_query_rows(self):
+        backend = _make_backend()
+        backend.init_cuda_graph_state(
+            max_bs=_CAPTURED_BS, max_num_tokens=_NUM_CAPTURED_ROWS
+        )
+        self.assertFalse(_would_grow(backend, _NUM_CAPTURED_ROWS))
+
+    def test_init_cuda_graph_state_is_grow_only(self):
+        """A later, smaller graph must not discard an earlier allocation."""
+        backend = _make_backend()
+        backend.init_cuda_graph_state(
+            max_bs=_CAPTURED_BS, max_num_tokens=_NUM_CAPTURED_ROWS
+        )
+        sized = backend._multi_ctas_kv_counter_buffer
+        backend.init_cuda_graph_state(max_bs=8, max_num_tokens=64)
+        self.assertIs(backend._multi_ctas_kv_counter_buffer, sized)
+
+    def test_init_cuda_graph_state_tolerates_backends_without_a_counter(self):
+        """Non-TRT-LLM branches leave the counter None; sizing must not read it."""
+        backend = _make_backend(allocate_counter=False)
+        backend.init_cuda_graph_state(
+            max_bs=_CAPTURED_BS, max_num_tokens=_NUM_CAPTURED_ROWS
+        )
+        self.assertIsNone(backend._multi_ctas_kv_counter_buffer)
+
+    def test_counter_field_defaults_to_none_on_the_class(self):
+        """The sizing hook is unconditional, so every branch must leave it readable."""
+        self.assertIsNone(DeepseekSparseAttnBackend._multi_ctas_kv_counter_buffer)
+
+    def test_oversized_eager_call_keeps_the_captured_allocation(self):
+        """Holds only a weakref and an address, so a rebinding implementation
+        drops the last strong reference and the assertions see it."""
+        backend = _make_backend()
+        backend.init_cuda_graph_state(
+            max_bs=_CAPTURED_BS, max_num_tokens=_NUM_CAPTURED_ROWS
+        )
+        captured_ref = weakref.ref(backend._multi_ctas_kv_counter_buffer)
+        captured_ptr = backend._multi_ctas_kv_counter_buffer.data_ptr()
+
+        counter = backend._multi_ctas_kv_counter_for(_NUM_CAPTURED_ROWS * 2)
+        self.assertIsNot(counter, backend._multi_ctas_kv_counter_buffer)
+        del counter
+
+        self.assertIsNotNone(captured_ref())
+        self.assertIs(backend._multi_ctas_kv_counter_buffer, captured_ref())
+        self.assertEqual(backend._multi_ctas_kv_counter_buffer.data_ptr(), captured_ptr)
+
+    def test_within_capacity_eager_call_reuses_the_captured_allocation(self):
+        backend = _make_backend()
+        backend.init_cuda_graph_state(
+            max_bs=_CAPTURED_BS, max_num_tokens=_NUM_CAPTURED_ROWS
+        )
+        self.assertIs(
+            backend._multi_ctas_kv_counter_for(_NUM_CAPTURED_ROWS),
+            backend._multi_ctas_kv_counter_buffer,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/attention/test_dsa_transform_index.py

@@ -7,13 +7,19 @@ import sglang.kernels.ops.attention.dsa.transform_index as transform_index_modul
 from sglang.kernels.ops.attention.dsa.transform_index import (
     transform_index_page_table_decode_fast,
     transform_index_page_table_prefill_fast,
+    transform_index_page_table_prefill_ref,
 )
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd-mi35x")
 
 TOPK = 2048
+# k-pool appends up to index_kpool - 1 open-tail tokens to index_topk, so the
+# width the indexer hands over is not a power of two: 2048 + 4 - 1 for
+# GLM-5.3-Flash. See get_dsa_mtp_topk_width() in srt/configs/model_config.py.
+KPOOL_TOPK = 2051
 
 
 @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
@@ -33,9 +39,11 @@ class TestDSATransformIndex(CustomTestCase):
         )
         return columns.unsqueeze(0) + row_bias
 
-    def _make_topk(self, rows: int, context_length: int) -> torch.Tensor:
+    def _make_topk(
+        self, rows: int, context_length: int, topk_width: int = TOPK
+    ) -> torch.Tensor:
         topk = (
-            torch.arange(TOPK, dtype=torch.int64, device=self.device)
+            torch.arange(topk_width, dtype=torch.int64, device=self.device)
             .remainder(context_length)
             .repeat(rows, 1)
         )
@@ -52,10 +60,11 @@ class TestDSATransformIndex(CustomTestCase):
         extend_lens_cpu: list[int],
         output_num_tokens: int,
         page_table_is_expanded: bool,
+        topk_width: int = TOPK,
     ) -> torch.Tensor:
         real_num_tokens = sum(extend_lens_cpu)
         expected = torch.full(
-            (output_num_tokens, TOPK),
+            (output_num_tokens, topk_width),
             -1,
             dtype=torch.int32,
             device=self.device,
@@ -91,14 +100,15 @@ class TestDSATransformIndex(CustomTestCase):
         *,
         zero_row_stride: bool = False,
         provide_result: bool = False,
+        topk_width: int = TOPK,
     ) -> None:
         if zero_row_stride:
             page_table = self._make_page_table(1, context_length).expand(batch_size, -1)
         else:
             page_table = self._make_page_table(batch_size, context_length)
-        topk_indices = self._make_topk(batch_size, context_length)
+        topk_indices = self._make_topk(batch_size, context_length, topk_width)
         expected = torch.empty(
-            (batch_size, TOPK), dtype=torch.int32, device=self.device
+            (batch_size, topk_width), dtype=torch.int32, device=self.device
         )
         torch.gather(
             page_table,
@@ -127,6 +137,7 @@ class TestDSATransformIndex(CustomTestCase):
         page_table_is_expanded: bool,
         topk_padding: int = 0,
         output_padding: int = 0,
+        topk_width: int = TOPK,
     ) -> None:
         real_num_tokens = sum(extend_lens_cpu)
         page_table_rows = (
@@ -135,13 +146,14 @@ class TestDSATransformIndex(CustomTestCase):
         topk_num_tokens = real_num_tokens + topk_padding
         output_num_tokens = topk_num_tokens + output_padding
         page_table = self._make_page_table(page_table_rows, context_length)
-        topk_indices = self._make_topk(topk_num_tokens, context_length)
+        topk_indices = self._make_topk(topk_num_tokens, context_length, topk_width)
         expected = self._expected(
             page_table,
             topk_indices,
             extend_lens_cpu,
             output_num_tokens,
             page_table_is_expanded,
+            topk_width,
         )
 
         actual = transform_index_page_table_prefill_fast(
@@ -235,6 +247,43 @@ class TestDSATransformIndex(CustomTestCase):
             output_padding=8,
         )
 
+    def test_prefill_kpool_tail_width(self):
+        extend_lens_cpu = [0, 3, 1, 0, 4]
+        real_num_tokens = sum(extend_lens_cpu)
+        topk_num_tokens = real_num_tokens + 5
+        output_num_tokens = topk_num_tokens + 7
+        context_length = 8192
+
+        # KPool can append three tail positions after the 2048 history indices.
+        tail_indices = torch.tensor(
+            [context_length - 3, context_length - 2, context_length - 1],
+            dtype=torch.int64,
+            device=self.device,
+        ).repeat(topk_num_tokens, 1)
+        topk_indices = torch.cat(
+            [self._make_topk(topk_num_tokens, context_length), tail_indices], dim=1
+        )
+        topk_indices[0, -1] = -1
+        self.assertEqual(topk_indices.shape[1], 2051)
+
+        for page_table_is_expanded in (False, True):
+            with self.subTest(page_table_is_expanded=page_table_is_expanded):
+                page_table_rows = (
+                    real_num_tokens if page_table_is_expanded else len(extend_lens_cpu)
+                )
+                page_table = self._make_page_table(page_table_rows, context_length)
+                kwargs = dict(
+                    page_table=page_table,
+                    topk_indices=topk_indices,
+                    extend_lens_cpu=extend_lens_cpu,
+                    output_num_tokens=output_num_tokens,
+                    page_table_is_expanded=page_table_is_expanded,
+                )
+                expected = transform_index_page_table_prefill_ref(**kwargs)
+                actual = transform_index_page_table_prefill_fast(**kwargs)
+                torch.cuda.synchronize()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
     def test_large_batch_size(self):
         self._check_case(
             [1] * 8192,
@@ -265,6 +314,26 @@ class TestDSATransformIndex(CustomTestCase):
         self._check_decode_case(8192, 4096)
         self._check_decode_case(2, 1_000_000)
 
+    def test_decode_fast_kpool_widths(self):
+        # 2051 is the GLM-5.3-Flash k-pool width; the others cover a partial
+        # trailing tile and a width below one tile.
+        for topk_width in (KPOOL_TOPK, 515, 257):
+            with self.subTest(topk_width=topk_width):
+                self._check_decode_case(17, 8192, topk_width=topk_width)
+                self._check_decode_case(
+                    17, 8192, topk_width=topk_width, provide_result=True
+                )
+
+    def test_prefill_kpool_widths(self):
+        for topk_width in (KPOOL_TOPK, 515, 257):
+            with self.subTest(topk_width=topk_width):
+                self._check_case(
+                    [2, 1],
+                    4096,
+                    page_table_is_expanded=False,
+                    topk_width=topk_width,
+                )
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/kernels/ops/attention/test_flash_attention_4_sm120.py

@@ -23,7 +23,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(
     est_time=240,
-    stage="base-b",
+    stage="base-b-kernel-unit",
     runner_config="1-gpu-small",
 )
 

test/registered/kernels/ops/attention/test_fused_kda_conv_recurrent_verify.py

@@ -14,7 +14,7 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 
-register_cuda_ci(est_time=90, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 _DEVICE = "cuda"
 
@@ -204,7 +204,7 @@ def _compare_case(case, num_warps, use_ring=False):
         rings_fus = {name: buf.clone() for name, buf in template.items()}
     else:
         rings_ref = rings_fus = None
-    o_ref, conv_ref, win_ref, ic_ref = _run_reference(
+    o_ref, _, win_ref, ic_ref = _run_reference(
         inp, B, T, H, HV, K, V, lower_bound, rings=rings_ref
     )
     o_fus, conv_fus, win_fus, ic_fus = _run_fused(
@@ -213,13 +213,13 @@ def _compare_case(case, num_warps, use_ring=False):
 
     idx_vals = inp["idx_vals"]
     valid_rows = [i for i, slot in enumerate(idx_vals) if slot >= 0]
-    touched_slots = [slot for slot in idx_vals if slot >= 0]
 
     o_ref_v = o_ref.reshape(B, T, HV, V)[valid_rows]
     o_fus_v = o_fus.reshape(B, T, HV, V)[valid_rows]
     # One bf16 ulp: the fused and reference tiles reduce K in different orders.
     torch.testing.assert_close(o_fus_v, o_ref_v, rtol=2**-7, atol=1e-7)
-    assert torch.equal(conv_ref[touched_slots], conv_fus[touched_slots])
+    # conv_state is read-only in verify; the commit scatter advances it.
+    assert torch.equal(inp["conv_pool"], conv_fus)
     assert torch.equal(win_ref[valid_rows], win_fus[valid_rows])
     if use_ring:
         # Full-tensor bitwise: ring values are elementwise (conv FMA chain,
@@ -239,6 +239,28 @@ def test_matches_unfused_reference(case):
     _compare_case(case, num_warps=4)
 
 
+def test_output_does_not_depend_on_cta_scheduling():
+    """The verify output must not change with how the CTAs happen to be
+    scheduled. H=1 with HV=16 shares one Q/K history across 16 V tiles."""
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("green contexts need SM90 or newer")
+    from flashinfer.green_ctx import split_device_green_ctx_by_sm_count
+
+    case = (1, 6, 1, 16, 128, 128, 4, False, None, False, 1)
+    B, T, H, HV, K, V, W, has_bias, lower_bound, neg_slot, seed = case
+    inp = _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed)
+    full = _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps=4)[0]
+
+    streams, _ = split_device_green_ctx_by_sm_count(torch.device("cuda:0"), [8])
+    # The green stream is non-blocking, so it must be told to wait for the
+    # inputs produced above; synchronize() afterwards only waits on the consumer.
+    streams[0].wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(streams[0]):
+        squeezed = _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps=4)[0]
+    streams[0].synchronize()
+    assert torch.equal(full, squeezed)
+
+
 @pytest.mark.parametrize("case", _RING_CASES)
 def test_replayssm_ring_matches_unfused(case):
     _compare_case(case, num_warps=4, use_ring=True)

test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py

@@ -11,6 +11,7 @@ import sys
 import pytest
 import torch
 
+from sglang.srt.utils import is_gfx95_supported
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 try:
@@ -19,6 +20,7 @@ try:
         fused_recurrent_gated_delta_rule_update,
     )
     from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
+        _select_recurrent_launch_config,
         fused_sigmoid_gating_delta_rule_update,
     )
 
@@ -180,6 +182,57 @@ def test_fused_gdn_mtp_precision(N: int, T: int):
     torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
 
 
+@pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernel not available")
+@pytest.mark.parametrize("N", [1, 3, 16])
+def test_qwen35_tp4_fused_gdn_mtp_precision(N: int):
+    """Exercise the gfx950 TP4 launch shape against the reference path."""
+    T, H, HV, K, V = 4, 4, 16, 128, 128
+    A_log, dt_bias, a, b, q, k, v, state, indices, cu_seqlens = _make_tensors(
+        N, T, H, HV, K, V
+    )
+
+    out_ref = run_reference(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state.clone(),
+        indices,
+        cu_seqlens,
+        disable_state_update=True,
+    )
+    out_fused = run_fused_mtp(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        state.clone(),
+        indices,
+        cu_seqlens,
+        disable_state_update=True,
+    )
+
+    torch.testing.assert_close(out_ref, out_fused, rtol=1e-2, atol=1e-2)
+
+
+@pytest.mark.skipif(
+    not (torch.version.hip and is_gfx95_supported()), reason="requires AMD gfx95"
+)
+def test_qwen35_tp4_launch_config_is_narrow():
+    assert _select_recurrent_launch_config(1, 4, 16, 128, 128, False) == (8, 4)
+    assert _select_recurrent_launch_config(3, 4, 16, 128, 128, False) == (16, 2)
+    assert _select_recurrent_launch_config(32, 4, 16, 128, 128, False) == (16, 2)
+    assert _select_recurrent_launch_config(33, 4, 16, 128, 128, False) == (32, 1)
+    assert _select_recurrent_launch_config(3, 8, 32, 128, 128, False) == (32, 1)
+    assert _select_recurrent_launch_config(3, 4, 16, 128, 128, True) == (32, 1)
+
+
 @pytest.mark.skipif(not KERNELS_AVAILABLE, reason="Kernels not available")
 @pytest.mark.parametrize("N", [1, 16, 128])
 def test_mtp_single_step_decode(N: int):

test/registered/kernels/ops/attention/test_gdn_decode_fused_proj_conv.py

@@ -14,7 +14,7 @@ from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
 from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update
 from sglang.test.ci.ci_register import register_cuda_ci
 
-register_cuda_ci(est_time=7, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=7, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 
 def _reference(

test/registered/kernels/ops/attention/test_kda_prefill.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import patch
 
 import torch
 import torch.nn.functional as F
@@ -10,6 +11,8 @@ from sglang.kernels.ops.attention.linear.kda_nvidia_prefill import (
 from sglang.kernels.ops.attention.linear.kda_ptx_prefill import (
     chunk_kda_fwd as ptx_chunk_kda_fwd,
 )
+from sglang.srt.layers.attention.linear.kernels.kda_ptx import PtxKDAKernel
+from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -79,6 +82,45 @@ def _reference(q, k, v, gate, beta, a_log, dt_bias, state, fused_qk_norm):
 
 
 class TestKdaPrefill(CustomTestCase):
+    @torch.inference_mode()
+    def test_ptx_padded_raw_beta(self):
+        """Raw beta must match Triton, including final state after neutral padding."""
+        if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (
+            10,
+            3,
+        ):
+            self.skipTest("PTX KDA prefill requires GB300")
+        q, k, v, gate, beta, a_log, dt_bias, state = _inputs(2, seq_len=1025)
+        state.fill_(0.1)
+        actual_state = state.clone()
+        inputs = dict(
+            q=q,
+            k=k,
+            v=v,
+            g=gate,
+            beta=beta,
+            cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32),
+            query_start_loc=torch.tensor([0, 1025], device="cuda", dtype=torch.int32),
+            A_log=a_log,
+            dt_bias=dt_bias,
+            lower_bound=-5.0,
+            beta_is_raw=True,
+            extend_seq_lens_cpu=[1025],
+        )
+        kernel = PtxKDAKernel()
+        with patch.object(
+            kernel._triton,
+            "extend",
+            side_effect=AssertionError("PTX unexpectedly fell back to Triton"),
+        ):
+            actual = kernel.extend(**inputs, ssm_states=actual_state)
+        # Triton may mutate inputs, so run the reference last.
+        expected = TritonKDAKernel().extend(**inputs, ssm_states=state)
+        torch.testing.assert_close(
+            actual.float(), expected.float(), rtol=2e-2, atol=3e-2
+        )
+        torch.testing.assert_close(actual_state, state, rtol=2e-2, atol=3e-2)
+
     @torch.inference_mode()
     def test_nvidia_prefill(self):
         if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:

test/registered/kernels/ops/attention/test_topk_v2.py

@@ -22,6 +22,9 @@ cluster floor and pool size are per-arch (see topk_v2.cuh), so the (batch, seq)
 grid below brackets the fixed boundaries (8192/8193, 16384/16385) exactly and
 spans the arch-dependent ones, across k in {512,1024,2048} and identity/perm
 page tables.
+
+``test_topk_v2_packed_rows`` covers the DSA extend layout on top of that: all
+requests packed into one score buffer, sharing a table row per request.
 """
 
 from __future__ import annotations
@@ -33,9 +36,11 @@ import torch
 
 from sglang.kernels.ops.attention.dsv4.topk import (
     plan_topk_v2,
+    topk_transform_packed_v2,
     topk_transform_paged_v2,
     topk_transform_ragged_v2,
 )
+from sglang.srt.utils import is_hip
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -465,5 +470,143 @@ def test_topk_v2_ragged_no_row_starts(k: int) -> None:
         assert sorted(explicit[i]) == sorted(implicit[i]), f"row {i} differs"
 
 
+@pytest.mark.skipif(
+    not is_hip(), reason="packed layout is compiled under USE_ROCM only"
+)
+@pytest.mark.parametrize("k", [512, 2048])
+@pytest.mark.parametrize(
+    "extend_lens",
+    [
+        [7],  # one request
+        [4, 4],  # equal row counts
+        [1, 13, 2],  # ragged, including a single-row request
+    ],
+)
+@torch.inference_mode()
+def test_topk_v2_packed_rows(extend_lens: list[int], k: int) -> None:
+    """DSA extend layout: batch-global packed scores + shared page-table rows.
+
+    Rows are causal within a request; a distinct page-table permutation per request
+    catches row/request index mix-ups, and the ragged case leaves most window starts
+    off the 16-byte load boundary (the production case).
+    """
+    torch.manual_seed(4242 + k + len(extend_lens))
+    device = "cuda"
+
+    # Keep every row longer than k so no row takes the trivial path.
+    prefix = k + 1024
+    kv_lens = [prefix + e for e in extend_lens]
+    k_offsets = [0]
+    for kv in kv_lens[:-1]:
+        k_offsets.append(k_offsets[-1] + kv)
+    total_kv = sum(kv_lens)
+
+    row_starts, lengths, row_to_batch = [], [], []
+    for i, e in enumerate(extend_lens):
+        for local in range(e):
+            row_starts.append(k_offsets[i])
+            lengths.append(kv_lens[i] - e + local + 1)
+            row_to_batch.append(i)
+    rows = len(lengths)
+
+    width = (total_kv + 3) & ~3
+    scores = torch.randn(rows, width, dtype=torch.float32, device=device)[:, :total_kv]
+    lengths_t = torch.tensor(lengths, dtype=torch.int32, device=device)
+    row_starts_t = torch.tensor(row_starts, dtype=torch.int32, device=device)
+    row_to_batch_t = torch.tensor(row_to_batch, dtype=torch.int32, device=device)
+
+    num_pages = (max(kv_lens) + PAGE_SIZE - 1) // PAGE_SIZE
+    page_table, inv_cpu = _make_page_table(
+        len(extend_lens), num_pages, "perm", device, per_row=True
+    )
+
+    out = torch.full((rows, k), -1, dtype=torch.int32, device=device)
+    # The kernel masks in place, so reference values must be read before the call.
+    scores_cpu = scores.cpu()
+    topk_transform_packed_v2(
+        scores,
+        lengths_t,
+        page_table,
+        out,
+        PAGE_SIZE,
+        row_starts=row_starts_t,
+        row_to_batch=row_to_batch_t,
+    )
+    torch.cuda.synchronize()
+
+    out_cpu = out.cpu().tolist()
+    for r in range(rows):
+        L, start, req = lengths[r], row_starts[r], row_to_batch[r]
+        window = scores_cpu[r, start : start + L]
+        ref = torch.topk(window, k, sorted=False).indices.tolist()
+        our = _invert(out_cpu[r], inv_cpu[req])
+        _assert_topk_close(window.unsqueeze(0), [ref], [our], 1, [L], k)
+
+
+@pytest.mark.skipif(
+    not is_hip(), reason="packed layout is compiled under USE_ROCM only"
+)
+@pytest.mark.parametrize("residue", [1, 2, 3])
+@pytest.mark.parametrize("boundary", [8192, 16384])
+@torch.inference_mode()
+def test_topk_v2_packed_level_boundary(boundary: int, residue: int) -> None:
+    """Rows whose length sits on an implementation's max_seq_len boundary.
+
+    The masked head widens the problem by ``residue``, so a row of exactly
+    ``boundary`` tokens spills past the register implementation sized for it. The
+    packed kernel picks the implementation per row from the widened length, so
+    these must still be exact; a compile-time choice made from the un-widened
+    length would overflow.
+    """
+    torch.manual_seed(boundary + residue)
+    device = "cuda"
+    k = 512
+
+    # Request 0 exists only to push request 1's window off the 16-byte boundary.
+    kv_lens = [residue, boundary + 1]
+    lengths = [residue] + [boundary - 1, boundary, boundary + 1]
+    row_starts = [0] + [residue] * 3
+    row_to_batch = [0, 1, 1, 1]
+    rows = len(lengths)
+    total_kv = sum(kv_lens)
+
+    width = (total_kv + 3) & ~3
+    scores = torch.randn(rows, width, dtype=torch.float32, device=device)[:, :total_kv]
+    lengths_t = torch.tensor(lengths, dtype=torch.int32, device=device)
+    row_starts_t = torch.tensor(row_starts, dtype=torch.int32, device=device)
+    row_to_batch_t = torch.tensor(row_to_batch, dtype=torch.int32, device=device)
+
+    num_pages = (max(kv_lens) + PAGE_SIZE - 1) // PAGE_SIZE
+    page_table, inv_cpu = _make_page_table(
+        len(kv_lens), num_pages, "perm", device, per_row=True
+    )
+
+    out = torch.full((rows, k), -1, dtype=torch.int32, device=device)
+    scores_cpu = scores.cpu()
+    topk_transform_packed_v2(
+        scores,
+        lengths_t,
+        page_table,
+        out,
+        PAGE_SIZE,
+        row_starts=row_starts_t,
+        row_to_batch=row_to_batch_t,
+    )
+    torch.cuda.synchronize()
+
+    out_cpu = out.cpu().tolist()
+    for r in range(rows):
+        L, start, req = lengths[r], row_starts[r], row_to_batch[r]
+        window = scores_cpu[r, start : start + L]
+        our = _invert(out_cpu[r], inv_cpu[req])
+        if L <= k:
+            # Trivial path: every position, then -1 padding.
+            assert sorted(our[:L]) == list(range(L)), f"row {r} trivial output wrong"
+            assert all(v == -1 for v in out_cpu[r][L:]), f"row {r} padding wrong"
+            continue
+        ref = torch.topk(window, k, sorted=False).indices.tolist()
+        _assert_topk_close(window.unsqueeze(0), [ref], [our], 1, [L], k)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/communication/test_custom_all_reduce.py

@@ -37,6 +37,7 @@ from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kernels.utils import multigpu_pytest_main
 
@@ -129,6 +130,7 @@ def _init_cpu_group_once() -> dist.ProcessGroup:
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=coord)
     atexit.register(dist.destroy_process_group)
     cpu_group = coord.cpu_group
     assert isinstance(cpu_group, dist.ProcessGroup)

test/registered/kernels/ops/communication/test_deterministic_reduce_scatter.py

@@ -14,6 +14,7 @@ from sglang.srt.distributed import parallel_state as ps
 from sglang.srt.environ import envs
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kernels.utils import multigpu_pytest_main
+from sglang.test.test_utils import publish_build_topology
 
 register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 
@@ -37,7 +38,8 @@ def group():
         local_rank=local_rank,
         distributed_init_method="env://",
     )
-    ps.initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(tp_size=world_size, world_rank=rank)
+    ps.initialize_model_parallel()
     yield ps.get_tp_group()
     ps.destroy_model_parallel()
     ps.destroy_distributed_environment()

test/registered/kernels/ops/communication/test_symm_mem_all_gather.py

@@ -31,6 +31,7 @@ from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
     all_gather_inner,
     create_state,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kernels.utils import multigpu_pytest_main
 
@@ -70,6 +71,7 @@ def _init_cpu_group_once() -> dist.ProcessGroup:
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=ps._WORLD)
     atexit.register(dist.destroy_process_group)
     logging.disable(logging.INFO)
     torch.cuda.set_stream(torch.cuda.Stream())

test/registered/kernels/ops/communication/test_tp_qknorm.py

@@ -24,6 +24,7 @@ from sglang.kernels.ops.communication.mp import register_comm_cleanup
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kernels.utils import multigpu_pytest_main
 
@@ -92,6 +93,7 @@ def _init_cpu_group_once() -> dist.ProcessGroup:
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=coord)
     atexit.register(dist.destroy_process_group)
     cpu_group = coord.cpu_group
     assert isinstance(cpu_group, dist.ProcessGroup)

test/registered/kernels/ops/embeddings/test_qwen4_ple_offload.py

@@ -6,6 +6,7 @@ import pytest
 import torch
 from torch import nn
 
+from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig
 from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
 from sglang.srt.layers.vocab_parallel_embedding import (
     VocabParallelEmbeddingShardIndices,
@@ -15,6 +16,7 @@ from sglang.srt.models.qwen4_exp import (
     Qwen4ExpPinnedHostEmbedding,
     Qwen4ExpPLELayer,
 )
+from sglang.srt.runtime_context import get_context, get_parallel
 from sglang.srt.utils import set_weight_attrs
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -191,6 +193,52 @@ def test_qwen4_ple_prefetch_buffer_lifecycle(monkeypatch):
     assert set(layer._graph_prefetch_buffers) == {3, 5}
 
 
+@pytest.fixture
+def single_rank_runtime_context():
+    """``Qwen4ExpPLELayer.__init__`` reads the TP topology through
+    ``VocabParallelEmbedding``; pin it to one rank without a process group."""
+    override = get_context().override_server_args(tp_size=1)
+    override.install()
+    try:
+        with get_parallel().override(
+            tp_rank=0, tp_size=1, attn_tp_rank=0, attn_tp_size=1
+        ):
+            yield
+    finally:
+        override.restore()
+
+
+def test_qwen4_ple_offload_avoids_device_table(single_rank_runtime_context):
+    # sgl-project/sglang#39841: the table was built on the device before the
+    # host table existed, so the flag needed a full per-rank shard of free VRAM.
+    # Small everywhere except the n-gram table (16 heads x ~20k rows x 4 dims),
+    # which must dominate the layer's footprint for the peak check to bite.
+    config = Qwen4ExpTextConfig(
+        vocab_size=64,
+        hidden_size=16,
+        hc_count=2,
+        ple_embed_dim=64,
+        ngram_size=3,
+        heads_per_ngram=8,
+        ngram_vocab_size_base=20_000,
+        eos_token_id=1,
+        ple_offload_embedding=True,
+    )
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
+    base = torch.cuda.memory_allocated()
+    with torch.device("cuda"):  # the model loader builds every layer this way
+        layer = Qwen4ExpPLELayer(config, prefix="ple", layer_id=0, ple_layer_index=0)
+    peak = torch.cuda.max_memory_allocated() - base
+
+    emb = layer.ple_embedding.ngram_embedding
+    table_bytes = emb.weight.numel() * emb.weight.element_size()
+    assert peak < table_bytes // 2, (peak, table_bytes)
+    assert emb.weight.device.type == "cpu" and emb.weight.is_pinned()
+    assert emb.weight_scale.is_cuda
+    assert not any(t.is_meta for t in (*layer.parameters(), *layer.buffers()))
+
+
 def _file_backend_supported() -> bool:
     from sglang.srt.models.qwen4_exp_ple_table import device_uses_host_page_tables
 

test/registered/kernels/ops/kimi_k3/test_ar_fusion.py

@@ -30,6 +30,7 @@ from sglang.kernels.ops.kimi_k3 import all_reduce
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kernels.utils import multigpu_pytest_main
 
@@ -65,6 +66,7 @@ def _init_world():
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=coord)
     atexit.register(dist.destroy_process_group)
     logging.disable(logging.INFO)
     torch.cuda.set_stream(torch.cuda.Stream())

test/registered/kernels/ops/kimi_k3/test_collectives.py

@@ -20,6 +20,7 @@ from sglang.kernels.ops.kimi_k3 import (
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kernels.utils import multigpu_pytest_main
 
@@ -53,6 +54,7 @@ def _init_world():
         local_rank=local_rank,
         backend="nccl",
     )
+    get_parallel().override_permanently(world_group=coord)
     atexit.register(dist.destroy_process_group)
     cpu_group = coord.cpu_group
     assert isinstance(cpu_group, dist.ProcessGroup)

test/registered/kernels/ops/kv_canary/test_verify_hand.py

@@ -20,6 +20,7 @@ from sglang.kernels.ops.kv_canary.verify import (
 from sglang.kernels.ops.kv_canary.verify_ref import (
     _compute_real_kv_hash_scalar,
     launch_canary_verify_kernel_torch_reference,
+    materialize_real_kv_sources,
 )
 from sglang.kernels.ops.kv_canary.write_ref import (
     launch_canary_write_kernel_torch_reference,
@@ -917,14 +918,18 @@ class TestRealKvHash:
         positions = [0, 1, 2]
 
         running = splitmix64(consts.CANARY_CHAIN_ANCHOR)
+        host_sources = materialize_real_kv_sources(
+            real_kv_sources=sources_cuda,
+            real_kv_hash_mode=consts.RealKvHashMode.ALL,
+            slot_indices=slot_indices,
+            work_device=torch.device("cpu"),
+        )
         real_kv_hashes: list[int] = []
         for slot_idx in slot_indices:
             real_kv_hashes.append(
                 _compute_real_kv_hash_scalar(
-                    real_kv_sources=sources_cuda,
-                    real_kv_hash_mode=consts.RealKvHashMode.ALL,
                     slot_idx=slot_idx,
-                    work_device=torch.device("cpu"),
+                    host_sources=host_sources,
                 )
             )
 
@@ -1030,6 +1035,21 @@ class TestRealKvSource:
                 read_bytes=0,
             )
 
+    def test_real_kv_source_rejects_row_narrower_than_page(self) -> None:
+        """A row too narrow for its page must raise: neither fold reports it.
+
+        The CUDA fold reads past the row and the torch fold's dim-1 slice clamps to
+        the row end, so the tail slots of the page hash 0 bytes and the chain still
+        verifies clean.
+        """
+        with pytest.raises(ValueError, match="page_size"):
+            RealKvSource(
+                tensor=torch.zeros((1, 16), dtype=torch.uint8, device=_DEVICE),
+                page_size=2,
+                num_bytes_per_token=16,
+                read_bytes=16,
+            )
+
     def test_real_kv_source_padding_below_4(self) -> None:
         """Host wrapper pads to 4 slots when fewer sources are supplied; dummy slots are never dereferenced."""
         buf_pair = _buf_pair()
@@ -1310,12 +1330,16 @@ class TestLayoutAndScheduling:
         # byte-by-byte loop, so the stamped real_kv_hash matches what the kernel /
         # verify reference will recompute. A byte-by-byte fold was the previous bug
         # here and triggered REAL_KV_HASH violations on otherwise clean chains.
+        host_sources = materialize_real_kv_sources(
+            real_kv_sources=sources_cuda,
+            real_kv_hash_mode=consts.RealKvHashMode.ALL,
+            slot_indices=slot_indices,
+            work_device=_DEVICE,
+        )
         rkv_values = [
             _compute_real_kv_hash_scalar(
                 slot_idx=slot_idx,
-                real_kv_sources=sources_cuda,
-                real_kv_hash_mode=consts.RealKvHashMode.ALL,
-                work_device=_DEVICE,
+                host_sources=host_sources,
             )
             for slot_idx in slot_indices
         ]

test/registered/kernels/ops/kvcache/test_hicache_tma.py

@@ -0,0 +1,173 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.kvcache.hicache import _jit_hicache_tma_module
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.version.hip is not None
+    or torch.cuda.get_device_capability()[0] < 9,
+    reason="HiCache TMA kernel requires SM90+",
+)
+
+POOL_TOKENS = 8192
+NUM_LAYERS = 3
+ROW_DIM = 256  # 512-byte bf16 rows: below the register kernel's 128 B unit width x 4
+
+
+def _token_indices(num_tokens: int, page_size: int, dtype: torch.dtype, seed: int):
+    gen = torch.Generator().manual_seed(seed)
+    pages = torch.randperm(POOL_TOKENS // page_size, generator=gen)[
+        : num_tokens // page_size
+    ]
+    idx = (pages[:, None] * page_size + torch.arange(page_size)).reshape(-1)
+    return idx.to(device="cuda", dtype=dtype)
+
+
+def _fill(t: torch.Tensor, seed: int) -> None:
+    t.view(torch.int16).copy_(
+        torch.randint(
+            0,
+            30000,
+            t.shape,
+            dtype=torch.int16,
+            generator=torch.Generator().manual_seed(seed),
+        )
+    )
+
+
+def _host_view(layout: str, layer: int):
+    if layout == "layer_first":
+        return torch.empty(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True)
+    # page_first: [tokens, layers, dim]; a per-layer view has strided rows
+    return torch.empty(
+        POOL_TOKENS, NUM_LAYERS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True
+    )[:, layer]
+
+
+@pytest.mark.parametrize("host_layout", ["layer_first", "page_first"])
+@pytest.mark.parametrize("index_dtype", [torch.int64, torch.int32])
+@pytest.mark.parametrize("page_size", [128, 1])
+def test_one_layer_roundtrip(
+    host_layout: str, index_dtype: torch.dtype, page_size: int
+) -> None:
+    """H2D then D2H of one layer; page runs take the single-op paths (bulk copy,
+    tensor-map box, bulk store), scattered rows take the per-row paths, and the
+    odd token count leaves a partial tail chunk."""
+    module = _jit_hicache_tma_module(block_quota=2)
+    num_tokens = 2048 + (96 if page_size == 1 else 0)
+    k_host, v_host = _host_view(host_layout, 1), _host_view(host_layout, 2)
+    k_dev = torch.zeros(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, device="cuda")
+    v_dev = torch.zeros_like(k_dev)
+    _fill(k_host, 1)
+    _fill(v_host, 2)
+    host_idx = _token_indices(num_tokens, page_size, index_dtype, seed=3)
+    dev_idx = _token_indices(num_tokens, page_size, index_dtype, seed=4)
+
+    module.launch_one(k_dev, v_dev, dev_idx, k_host, v_host, host_idx)
+    torch.cuda.synchronize()
+    assert torch.equal(k_dev[dev_idx.long()].cpu(), k_host[host_idx.cpu().long()])
+    assert torch.equal(v_dev[dev_idx.long()].cpu(), v_host[host_idx.cpu().long()])
+    untouched = torch.ones(POOL_TOKENS, dtype=torch.bool, device="cuda")
+    untouched[dev_idx.long()] = False
+    assert not k_dev[untouched].any() and not v_dev[untouched].any()
+
+    _fill(k_dev, 5)
+    _fill(v_dev, 6)
+    k_host.zero_()
+    v_host.zero_()
+    module.launch_one(k_host, v_host, host_idx, k_dev, v_dev, dev_idx)
+    torch.cuda.synchronize()
+    assert torch.equal(k_host[host_idx.cpu().long()], k_dev[dev_idx.long()].cpu())
+    assert torch.equal(v_host[host_idx.cpu().long()], v_dev[dev_idx.long()].cpu())
+
+
+def _ptr_table(tensors) -> torch.Tensor:
+    return torch.tensor(
+        [t.data_ptr() for t in tensors], dtype=torch.uint64, device="cuda"
+    )
+
+
+def test_all_layer_tables_lf_to_pf() -> None:
+    """All-layer D2H through per-layer pointer tables into a page-first host pool
+    (strided destination rows), the write-back shape."""
+    module = _jit_hicache_tma_module(block_quota=2)
+    k_dev = [
+        torch.empty(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, device="cuda")
+        for _ in range(NUM_LAYERS)
+    ]
+    v_dev = [torch.empty_like(k_dev[0]) for _ in range(NUM_LAYERS)]
+    for i, t in enumerate(k_dev + v_dev):
+        _fill(t, 10 + i)
+    k_host = torch.zeros(
+        POOL_TOKENS, NUM_LAYERS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True
+    )
+    v_host = torch.zeros_like(k_host).pin_memory()
+    host_idx = _token_indices(2048, 128, torch.int64, seed=7)
+    dev_idx = _token_indices(2048, 128, torch.int64, seed=8)
+    row_bytes = ROW_DIM * 2
+
+    module.launch_all(
+        _ptr_table([k_host[:, l] for l in range(NUM_LAYERS)]),
+        _ptr_table([v_host[:, l] for l in range(NUM_LAYERS)]),
+        host_idx,
+        _ptr_table(k_dev),
+        _ptr_table(v_dev),
+        dev_idx,
+        row_bytes,
+        NUM_LAYERS * row_bytes,
+        row_bytes,
+    )
+    torch.cuda.synchronize()
+    for l in range(NUM_LAYERS):
+        assert torch.equal(k_host[host_idx.cpu(), l], k_dev[l][dev_idx].cpu())
+        assert torch.equal(v_host[host_idx.cpu(), l], v_dev[l][dev_idx].cpu())
+    untouched = torch.ones(POOL_TOKENS, dtype=torch.bool)
+    untouched[host_idx.cpu()] = False
+    assert not k_host[untouched].any() and not v_host[untouched].any()
+
+
+def test_mla_single_buffer() -> None:
+    """MLA rows (576 x bf16 = 1152 B, not a multiple of 128 B) through the
+    single-buffer entry points, one layer and all layers."""
+    module = _jit_hicache_tma_module(block_quota=2)
+    dim = 576
+    dev = [
+        torch.empty(POOL_TOKENS, dim, dtype=torch.bfloat16, device="cuda")
+        for _ in range(NUM_LAYERS)
+    ]
+    for i, t in enumerate(dev):
+        _fill(t, 20 + i)
+    host = torch.zeros(
+        POOL_TOKENS, NUM_LAYERS, dim, dtype=torch.bfloat16, pin_memory=True
+    )
+    host_idx = _token_indices(1024, 128, torch.int64, seed=9)
+    dev_idx = _token_indices(1024, 128, torch.int64, seed=10)
+    row_bytes = dim * 2
+
+    module.launch_all_mla(
+        _ptr_table([host[:, l] for l in range(NUM_LAYERS)]),
+        host_idx,
+        _ptr_table(dev),
+        dev_idx,
+        row_bytes,
+        NUM_LAYERS * row_bytes,
+        row_bytes,
+    )
+    torch.cuda.synchronize()
+    for l in range(NUM_LAYERS):
+        assert torch.equal(host[host_idx.cpu(), l], dev[l][dev_idx].cpu())
+
+    dev[0].zero_()
+    module.launch_one_mla(dev[0], dev_idx, host[:, 0], host_idx)
+    torch.cuda.synchronize()
+    assert torch.equal(dev[0][dev_idx].cpu(), host[host_idx.cpu(), 0])
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/kvcache/test_hisparse.py

@@ -3,26 +3,44 @@ import sys
 import pytest
 import torch
 
-from sglang.kernels.ops.kvcache.hisparse import (
-    load_cache_to_device_buffer_dsv4_mla,
-    load_cache_to_device_buffer_mla,
-    transfer_cache_dsv4_mla,
+from sglang.srt.utils import (
+    get_device,
+    get_device_module,
+    is_cuda,
+    is_hip,
+    is_xpu,
+)
+from sglang.test.ci.ci_register import (
+    register_amd_ci,
+    register_cuda_ci,
+    register_xpu_ci,
 )
-from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_amd_ci(est_time=30, stage="stage-b", runner_config="1-gpu-small-amd")
 register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
+
+if is_xpu():
+    from sgl_kernel import (
+        load_cache_to_device_buffer_dsv4_mla,
+        load_cache_to_device_buffer_mla,
+        transfer_cache_dsv4_mla,
+    )
+else:
+    from sglang.kernels.ops.kvcache.hisparse import (
+        load_blocks_to_device_buffer_mha,
+        load_cache_to_device_buffer_dsv4_mla,
+        load_cache_to_device_buffer_mla,
+        transfer_cache_dsv4_mla,
+    )
+
 
 pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or is_npu()
-    or is_xpu()
-    or not (is_cuda() or is_hip()),
-    reason="HiSparse JIT tests require CUDA/ROCm.",
+    not (is_cuda() or is_hip() or is_xpu()),
+    reason="HiSparse kernel tests require CUDA/ROCm/XPU.",
 )
 
-DEVICE = "cuda"
+DEVICE = get_device()
 DTYPE = torch.float32
 KV_DIM = 8
 HOT_BUFFER_SIZE = 4
@@ -131,7 +149,7 @@ def _run_kernel(
         block_size=256,
         num_real_reqs=torch.tensor([num_real_reqs], dtype=torch.int32, device=DEVICE),
     )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
     return out
 
 
@@ -171,7 +189,7 @@ def _make_state(
         device_buffer[device_buffer_locs[rid, HOT_BUFFER_SIZE]].copy_(
             host_cache[newest_token].to(DEVICE, non_blocking=True)
         )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     return {
         "host_cache": host_cache,
@@ -199,7 +217,7 @@ def test_transfer_cache_dsv4_mla_copies_paged_token() -> None:
         src_indices=torch.tensor([src_loc], dtype=torch.int64, device=DEVICE),
         dst_indices=torch.tensor([dst_loc], dtype=torch.int64, device=DEVICE),
     )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     assert torch.equal(
         _read_dsv4_token(dst_cache, dst_loc).to(DEVICE),
@@ -251,7 +269,7 @@ def test_dsv4_swap_in_reads_paged_host_layout() -> None:
         block_size=256,
         num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
     )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     assert out.item() == swap_loc
     assert torch.equal(
@@ -351,6 +369,84 @@ def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None:
     )
 
 
+@pytest.mark.skipif(is_xpu(), reason="MiniMax MHA block swap-in has no XPU kernel.")
+def test_load_blocks_to_device_buffer_mha_handles_partial_newest_block() -> None:
+    """A partial newest block must not consume slots for its invalid tail."""
+    sparse_block_size = 4
+    hot_buffer_size = 8
+    host_k = _host_cache()
+    host_v = _host_cache()
+    host_v.add_(1000)
+    device_k = torch.full(
+        (DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE
+    )
+    device_v = torch.full_like(device_k, -1)
+    device_buffer_locs = torch.arange(
+        hot_buffer_size + 1, dtype=torch.int32, device=DEVICE
+    ).view(1, -1)
+    device_buffer_tokens = torch.tensor(
+        [[0, 1, 2, 3, -1, -1, -1, -1, -1]],
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    for slot, token in enumerate([0, 1, 2, 3]):
+        device_k[device_buffer_locs[0, slot]].copy_(host_k[token], non_blocking=True)
+        device_v[device_buffer_locs[0, slot]].copy_(host_v[token], non_blocking=True)
+    device_k[device_buffer_locs[0, hot_buffer_size]].copy_(
+        host_k[10], non_blocking=True
+    )
+    device_v[device_buffer_locs[0, hot_buffer_size]].copy_(
+        host_v[10], non_blocking=True
+    )
+
+    top_k_blocks = torch.tensor([[0, 2]], dtype=torch.int32, device=DEVICE)
+    out = torch.full(
+        (1, top_k_blocks.size(1) * sparse_block_size),
+        -1,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    lru_slots = torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE).view(
+        1, -1
+    )
+    load_blocks_to_device_buffer_mha(
+        top_k_blocks=top_k_blocks,
+        device_buffer_tokens=device_buffer_tokens,
+        host_cache_locs=torch.arange(
+            HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE
+        ).view(1, -1),
+        device_buffer_locs=device_buffer_locs,
+        host_cache_k=host_k,
+        host_cache_v=host_v,
+        device_buffer_k=device_k,
+        device_buffer_v=device_v,
+        top_k_device_locs=out,
+        req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
+        seq_lens=torch.tensor([11], dtype=torch.int32, device=DEVICE),
+        lru_slots=lru_slots,
+        item_size_bytes=ITEM_SIZE_BYTES,
+        hot_buffer_size=hot_buffer_size,
+        sparse_block_size=sparse_block_size,
+        num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
+    )
+    get_device_module().synchronize()
+
+    assert torch.equal(
+        out.cpu(), torch.tensor([[0, 1, 2, 3, 4, 5, 8, -1]], dtype=torch.int32)
+    )
+    assert torch.equal(device_k[4].cpu(), host_k[8])
+    assert torch.equal(device_v[4].cpu(), host_v[8])
+    assert torch.equal(device_k[5].cpu(), host_k[9])
+    assert torch.equal(device_v[5].cpu(), host_v[9])
+    assert torch.equal(
+        device_buffer_tokens.cpu(),
+        torch.tensor([[0, 1, 2, 3, 8, 9, -1, -1, -1]], dtype=torch.int32),
+    )
+    assert torch.equal(
+        lru_slots.cpu(), torch.tensor([[6, 7, 4, 5, 0, 1, 2, 3]], dtype=torch.int16)
+    )
+
+
 def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None:
     state = _long_case()
 
@@ -427,7 +523,7 @@ def test_load_cache_to_device_buffer_miss_copy_is_byte_exact(
     )
     for slot in range(HOT_BUFFER_SIZE):
         device_buffer[slot].copy_(host_cache[slot].to(DEVICE))
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     top_k_tokens = torch.tensor([[miss_token]], dtype=torch.int32, device=DEVICE)
     out = torch.full_like(top_k_tokens, -1)
@@ -454,7 +550,7 @@ def test_load_cache_to_device_buffer_miss_copy_is_byte_exact(
         block_size=256,
         num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
     )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     # The miss evicts the LRU head (slot 0, physical loc 0) and lands there.
     assert torch.equal(out.cpu(), torch.tensor([[0]], dtype=torch.int32))
@@ -593,7 +689,7 @@ def test_load_cache_to_device_buffer_dsv4_mla_miss_copy_layout() -> None:
         block_size=256,
         num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
     )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32))
 
@@ -666,7 +762,7 @@ def test_load_cache_to_device_buffer_dsv4_fused_copy_multi_miss() -> None:
         block_size=256,
         num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
     )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     # Which slot each miss evicts is up to the LRU, so take the destinations
     # from the kernel; only require that they are distinct and in range.
@@ -741,7 +837,7 @@ def test_load_cache_to_device_buffer_rocm_large_lru_writeback() -> None:
         block_size=1024,
         num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
     )
-    torch.cuda.synchronize()
+    get_device_module().synchronize()
 
     expected_lru = torch.cat(
         [

test/registered/kernels/ops/kvcache/test_mxfp8_kv_reserved_slot.py

@@ -0,0 +1,138 @@
+"""MXFP8 KV cache must never write the reserved CUDA-graph padding slot.
+
+Padding lanes carry undefined activations that quantize to NaN payload and
+0xFF e8m0 scales; attention reads slot 0 back for padded page-table entries,
+so a poisoned slot 0 defeats probability masking (0 * NaN = NaN in PV).
+Asserts require slot 0 to stay exactly zero, so finite-garbage writes fail too.
+"""
+
+import pytest
+import torch
+
+from sglang.srt.utils import get_device_sm
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+requires_sm100 = pytest.mark.skipif(
+    not torch.cuda.is_available() or get_device_sm() < 100,
+    reason="MXFP8 KV cache requires SM100+",
+)
+
+DEV, HD, PS, NHKV = "cuda", 128, 128, 2
+
+
+def _make_pool(**kwargs):
+    from sglang.srt.mem_cache.memory_pool import MHATokenToKVPoolMXFP8
+
+    return MHATokenToKVPoolMXFP8(
+        size=4 * PS,
+        page_size=PS,
+        dtype=torch.float8_e4m3fn,
+        head_num=NHKV,
+        head_dim=HD,
+        layer_num=1,
+        device=DEV,
+        enable_memory_saver=False,
+        **kwargs,
+    )
+
+
+class _Layer:
+    layer_id = 0
+
+
+def _quantize(k, v):
+    from sglang.kernels.ops.quantization.mxfp8_quant import to_mxfp8
+
+    km, vm = to_mxfp8(k), to_mxfp8(v)
+    return (
+        km.data,
+        vm.data,
+        km.scale.view(torch.float8_e8m0fnu),
+        vm.scale.view(torch.float8_e8m0fnu),
+    )
+
+
+def _assert_slot0_zero(pool):
+    kc, vc = pool.get_kv_buffer(0)
+    ksf, vsf = pool.get_kv_scale_buffer(0)
+    s0k = kc.view(-1, PS, NHKV, HD)[0, 0].view(torch.uint8)
+    s0v = vc.view(-1, PS, NHKV, HD)[0, 0].view(torch.uint8)
+    assert int(s0k.sum()) == 0, "reserved slot K payload written"
+    assert int(s0v.sum()) == 0, "reserved slot V payload written"
+    zero_loc = torch.zeros(1, dtype=torch.int64, device=DEV)
+    s0_ksf = pool._read_sf_interleaved(ksf, zero_loc).view(torch.uint8)
+    s0_vsf = pool._read_sf_interleaved(vsf, zero_loc).view(torch.uint8)
+    assert int(s0_ksf.sum()) == 0, "reserved slot K scales written"
+    assert int(s0_vsf.sum()) == 0, "reserved slot V scales written"
+
+
+@requires_sm100
+def test_direct_path_skips_reserved_slot():
+    torch.manual_seed(0)
+    pool = _make_pool()
+    k = torch.randn(4, NHKV, HD, dtype=torch.bfloat16, device=DEV) * 0.5
+    v = torch.randn(4, NHKV, HD, dtype=torch.bfloat16, device=DEV) * 0.5
+    k[[0, 2]] = float("nan")
+    v[[0, 2]] = float("nan")
+    kq, vq, ks, vs = _quantize(k, v)
+    loc = torch.tensor([0, 7, 0, 9], dtype=torch.int64, device=DEV)
+
+    pool.set_kv_buffer(_Layer(), loc, kq, vq, ks, vs)
+
+    _assert_slot0_zero(pool)
+    kc, _ = pool.get_kv_buffer(0)
+    got = kc.view(-1, PS, NHKV, HD)[0, 7].view(torch.uint8)
+    assert torch.equal(got, kq[1].view(torch.uint8)), "non-reserved write corrupted"
+
+
+@requires_sm100
+def test_fused_quant_store_path_skips_reserved_slot():
+    """k_scale=None routes to the fused quant_store_kv_mxfp8 kernel."""
+    torch.manual_seed(2)
+    pool = _make_pool()
+    k = torch.randn(4, NHKV, HD, dtype=torch.bfloat16, device=DEV) * 0.5
+    v = torch.randn(4, NHKV, HD, dtype=torch.bfloat16, device=DEV) * 0.5
+    k[[1, 3]] = float("nan")
+    v[[1, 3]] = float("nan")
+    loc = torch.tensor([5, 0, 9, 0], dtype=torch.int64, device=DEV)
+
+    pool.set_kv_buffer(_Layer(), loc, k, v)
+
+    _assert_slot0_zero(pool)
+    kc, _ = pool.get_kv_buffer(0)
+    valid = kc.view(-1, PS, NHKV, HD)[0, 5].float()
+    assert not torch.isnan(valid).any() and valid.abs().sum() > 0, (
+        "fused valid write lost"
+    )
+
+
+@requires_sm100
+def test_set_kv_buffer_is_cuda_graph_capture_safe():
+    torch.manual_seed(5)
+    pool = _make_pool(enable_alt_stream=False)
+    k = torch.randn(2, NHKV, HD, dtype=torch.bfloat16, device=DEV) * 0.5
+    v = torch.randn(2, NHKV, HD, dtype=torch.bfloat16, device=DEV) * 0.5
+    kq, vq, ks, vs = _quantize(k, v)
+    loc = torch.tensor([0, 7], dtype=torch.int64, device=DEV)
+
+    for _ in range(2):  # warmup
+        pool.set_kv_buffer(_Layer(), loc, kq, vq, ks, vs)
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        pool.set_kv_buffer(_Layer(), loc, kq, vq, ks, vs)
+    g.replay()
+    torch.cuda.synchronize()
+
+    _assert_slot0_zero(pool)
+    kc, _ = pool.get_kv_buffer(0)
+    got = kc.view(-1, PS, NHKV, HD)[0, 7].view(torch.uint8)
+    assert torch.equal(got, kq[1].view(torch.uint8)), "captured valid write lost"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/kvcache/test_pd_dcp_gather.py

@@ -1,8 +1,13 @@
+import concurrent.futures
 import unittest
+from types import SimpleNamespace
 
+import numpy as np
 import torch
 
 from sglang.kernels.ops.kvcache.pd_dcp_gather import copy_mla_rows_into_pack
+from sglang.srt.disaggregation.common.staging_buffer import StagingBuffer
+from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -36,6 +41,135 @@ class TestPdDcpGather(CustomTestCase):
         torch.testing.assert_close(packed0, kv0[row_indices], rtol=0, atol=0)
         torch.testing.assert_close(packed1, kv1[row_indices], rtol=0, atol=0)
 
+    def test_packed_tp2_pp2_to_dcp4_preserves_kv(self):
+        """Packing must preserve both target rows and draft head shards across PP stages."""
+        for custom_pool in (False, True):
+            for capacity in (256 * (2 * 64 + 2 * 512), 256 * 2 * 64 // 4):
+                for rank in range(4):
+                    with self.subTest(
+                        custom_pool=custom_pool, capacity=capacity, rank=rank
+                    ):
+                        self._check_packed_transfer(rank, custom_pool, capacity)
+
+    def _check_packed_transfer(self, rank, custom_pool, capacity):
+        page, tokens, chunk = 64, 521, 256
+        src_pages = np.array([7, 1, 9, 3, 4, 11, 2, 5, 8], dtype=np.int32)
+        dst_pages = np.array([4, 1, 6], dtype=np.int32)
+        layers, widths = [3, 11, 19, 27, 28, 28], [64] * 4 + [256] * 2
+        logical = torch.arange(tokens, device="cuda")
+        src_rows = (
+            torch.as_tensor(src_pages, device="cuda")[logical // page] * page
+            + logical % page
+        )
+        values = [
+            (
+                (logical[:, None] + 256) * 13
+                + torch.arange(width, device="cuda") * 7
+                + entry * 31
+            )
+            .remainder(251)
+            .to(torch.uint8)
+            for entry, width in enumerate([64] * 4 + [1024] * 2)
+        ]
+        destinations = [
+            torch.full((2048, w), 165, dtype=torch.uint8, device="cuda") for w in widths
+        ]
+        expected = [x.clone() for x in destinations]
+        owned = logical[rank::4]
+        target_rows = (
+            torch.as_tensor(dst_pages, device="cuda")[owned // 256] * page
+            + owned % 256 // 4
+        )
+        draft_rows = (
+            torch.as_tensor(dst_pages, device="cuda")[logical // 256] * 256
+            + logical % 256
+        )
+        for entry in range(4):
+            expected[entry][target_rows] = values[entry][owned]
+        for entry in (4, 5):
+            expected[entry][draft_rows] = values[entry][
+                :, rank * 256 : (rank + 1) * 256
+            ]
+
+        pack = StagingBuffer(capacity, "cuda:0", 0)
+        with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+            for stage, entries in enumerate(([0, 1], [2, 3, 4, 5])):
+                sources = []
+                for entry in entries:
+                    data = values[entry]
+                    if entry >= 4:
+                        start = (rank // 2) * 512
+                        data = data[:, start : start + 512]
+                    source = torch.full(
+                        (1024, data.shape[1]), 165, dtype=torch.uint8, device="cuda"
+                    )
+                    source[src_rows] = data
+                    sources.append(source)
+                buffers = sources + destinations + [pack.buffer]
+
+                def transfer(session, blocks):
+                    def view(ptr, size):
+                        for tensor in buffers:
+                            offset = ptr - tensor.data_ptr()
+                            if 0 <= offset and offset + size <= tensor.numel():
+                                return tensor.flatten()[offset : offset + size]
+                        raise AssertionError(
+                            f"Transfer outside registered buffers: {ptr}, {size}"
+                        )
+
+                    for src, dst, size in blocks:
+                        view(dst, size).copy_(view(src, size))
+                    torch.cuda.synchronize()
+                    return 0
+
+                manager = SimpleNamespace(
+                    is_mla_backend=False,
+                    kv_args=SimpleNamespace(
+                        page_size=page,
+                        kv_layer_ids=[layers[e] for e in entries],
+                        kv_data_ptrs=[x.data_ptr() for x in sources],
+                        num_draft_entries=2 if stage else 0,
+                        engine_rank=stage * 2 + rank // 2,
+                    ),
+                    attn_tp_size=2,
+                    max_transfer_batch_indices=37,
+                    enable_custom_mem_pool=custom_pool,
+                    enable_deferred_decode_kv_release=False,
+                    _transfer_data=transfer,
+                )
+                manager._await_transfer_futures = lambda futures: (
+                    MooncakeKVManager._await_transfer_futures(manager, futures)
+                )
+                for start in range(0, tokens, chunk):
+                    count = min(chunk, tokens - start)
+                    result = MooncakeKVManager.send_kvcache_dcp(
+                        manager,
+                        "session",
+                        src_pages[start // page : (start + count + page - 1) // page],
+                        [x.data_ptr() for x in destinations],
+                        dst_pages,
+                        dcp_token_item_lens=[x.shape[1] for x in sources],
+                        dst_dcp_size=4,
+                        dst_dcp_rank=rank,
+                        src_page_offset=start // page,
+                        decode_prefix_len=256,
+                        num_kv_tokens=count,
+                        executor=executor,
+                        dst_layer_ids=layers,
+                        pack_buffer=pack,
+                        dst_kv_item_lens=[
+                            page * w * (4 if e >= 4 else 1)
+                            for e, w in enumerate(widths)
+                        ],
+                        dst_tp_rank=rank,
+                        dst_attn_tp_size=4,
+                    )
+                    self.assertEqual(result, 0)
+                for entry in entries:
+                    torch.testing.assert_close(
+                        destinations[entry], expected[entry], rtol=0, atol=0
+                    )
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/kernels/ops/layernorm/test_gemma4_fused_routing.py

@@ -16,7 +16,7 @@ import torch
 
 from sglang.test.ci.ci_register import register_cuda_ci
 
-register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-small")
 
 pytestmark = pytest.mark.skipif(
     not torch.cuda.is_available(),

test/registered/kernels/ops/layernorm/test_mhc_aiter_hip.py

@@ -0,0 +1,213 @@
+"""The AITER mHC route on gfx950: gate, fallback latch, and kernel numerics vs the Torch oracle."""
+
+import sys
+import types
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.kernels.ops.layernorm import mhc
+from sglang.srt.environ import envs
+from sglang.srt.utils import is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=120, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
+    "requires one gfx950 GPU",
+)
+class TestAiterMHCGLM53Flash(CustomTestCase):
+    hidden_size = 4096
+    hc_mult = 4
+    rms_eps = 1e-6
+    hc_eps = 1e-6
+
+    def setUp(self):
+        mhc._AITER_MHC_RUNTIME_DISABLED = False
+
+    def _inputs(self, tokens: int, seed: int = 0):
+        torch.manual_seed(seed)
+        device = torch.device("cuda")
+        mix_size = 2 * self.hc_mult + self.hc_mult**2
+        residual = (
+            torch.randn(
+                tokens,
+                self.hc_mult,
+                self.hidden_size,
+                device=device,
+                dtype=torch.bfloat16,
+            )
+            * 0.1
+        )
+        fn = (
+            torch.randn(
+                mix_size,
+                self.hc_mult * self.hidden_size,
+                device=device,
+                dtype=torch.float32,
+            )
+            * 0.01
+        )
+        scale = torch.tensor([0.5, 0.25, 0.25], device=device, dtype=torch.float32)
+        base = torch.zeros(mix_size, device=device, dtype=torch.float32)
+        return residual, fn, scale, base
+
+    def _rmsnorm(self, x, weight):
+        return (
+            x.float()
+            * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + self.rms_eps)
+            * weight.float()
+        ).to(x.dtype)
+
+    def test_gate_selects_aiter_on_gfx950(self):
+        """A gate that resolves False on real hardware silently serves the Torch path."""
+        with envs.SGLANG_USE_AITER.override(True):
+            self.assertTrue(mhc._use_aiter_mhc())
+
+    def test_hip_without_aiter_stays_on_torch_and_never_loads_tilelang(self):
+        """The TileLang/DeepGEMM flags default on; only the HIP gate keeps them off this device."""
+        residual, fn, scale, base = self._inputs(8)
+        x = residual.reshape(8, self.hc_mult * self.hidden_size)
+        _, _, layer_ref = mhc._mhc_pre_torch(
+            residual, fn, scale, base, self.rms_eps, self.hc_eps, self.hc_eps, 2.0, 4
+        )
+        with (
+            envs.SGLANG_USE_AITER.override(False),
+            envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.override(True),
+            envs.SGLANG_OPT_USE_TILELANG_MHC_POST.override(True),
+            envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.override(True),
+            patch.object(
+                mhc, "_load_tilelang", side_effect=AssertionError("TileLang imported")
+            ),
+        ):
+            self.assertFalse(mhc._use_aiter_mhc())
+            self.assertFalse(mhc._use_tilelang_mhc_pre())
+            self.assertFalse(mhc._use_tilelang_mhc_post())
+            self.assertFalse(mhc._use_deep_gemm_hc_prenorm())
+            layer_input, h_res, h_post, norm_fused = mhc.hc_pre(
+                x, fn, scale, base, self.hc_mult, self.rms_eps, self.hc_eps, 4
+            )
+            out = mhc.hc_post(layer_input, x, h_post, h_res, self.hc_mult)
+        self.assertFalse(norm_fused)
+        torch.testing.assert_close(layer_input, layer_ref)
+        self.assertTrue(torch.isfinite(out).all())
+
+    def test_aiter_import_and_runtime_failures_latch_to_torch(self):
+        """A missing symbol or a raising kernel must disable the route once, not fail the request."""
+        residual, fn, scale, base = self._inputs(8)
+        x = residual.reshape(8, self.hc_mult * self.hidden_size)
+        modules = {
+            "aiter": types.ModuleType("aiter"),
+            "aiter.ops": types.ModuleType("aiter.ops"),
+            "aiter.ops.mhc": types.ModuleType("aiter.ops.mhc"),
+        }
+        with patch.dict(sys.modules, modules):
+            result = mhc._try_aiter_mhc_pre(
+                residual,
+                fn,
+                scale,
+                base,
+                self.rms_eps,
+                self.hc_eps,
+                self.hc_eps,
+                2.0,
+                4,
+                None,
+                None,
+            )
+        self.assertIsNone(result)
+        self.assertTrue(mhc._AITER_MHC_RUNTIME_DISABLED)
+
+        mhc._AITER_MHC_RUNTIME_DISABLED = False
+
+        def fail_post(*_args, **_kwargs):
+            raise RuntimeError("synthetic failure")
+
+        failing = types.ModuleType("aiter.ops.mhc")
+        failing.mhc_post = fail_post
+        modules["aiter.ops.mhc"] = failing
+        with envs.SGLANG_USE_AITER.override(False):
+            layer_input, h_res, h_post, _ = mhc.hc_pre(
+                x, fn, scale, base, self.hc_mult, self.rms_eps, self.hc_eps, 4
+            )
+        with (
+            patch.dict(sys.modules, modules),
+            envs.SGLANG_USE_AITER.override(True),
+        ):
+            out = mhc.hc_post(layer_input, x, h_post, h_res, self.hc_mult)
+        self.assertTrue(mhc._AITER_MHC_RUNTIME_DISABLED)
+        self.assertTrue(torch.isfinite(out).all())
+
+    def test_aiter_pre_post_match_torch_oracle(self):
+        """A positional or kwarg mixup in the AITER call shows up only against the real kernel."""
+        norm_weight = torch.linspace(
+            0.75, 1.25, self.hidden_size, device="cuda", dtype=torch.bfloat16
+        )
+        for tokens in (1, 8, 17, 32, 64, 128):
+            for sinkhorn_iters in (2, 20):
+                for fused_norm in (False, True):
+                    with self.subTest(
+                        tokens=tokens, sinkhorn_iters=sinkhorn_iters, norm=fused_norm
+                    ):
+                        residual, fn, scale, base = self._inputs(tokens)
+                        post_ref, comb_ref, layer_ref = mhc._mhc_pre_torch(
+                            residual,
+                            fn,
+                            scale,
+                            base,
+                            self.rms_eps,
+                            self.hc_eps,
+                            self.hc_eps,
+                            2.0,
+                            sinkhorn_iters,
+                        )
+                        result = mhc._try_aiter_mhc_pre(
+                            residual,
+                            fn,
+                            scale,
+                            base,
+                            self.rms_eps,
+                            self.hc_eps,
+                            self.hc_eps,
+                            2.0,
+                            sinkhorn_iters,
+                            norm_weight if fused_norm else None,
+                            self.rms_eps if fused_norm else None,
+                        )
+                        self.assertIsNotNone(result, "AITER mHC pre fell back")
+                        post_out, comb_out, layer_out = result
+                        if fused_norm:
+                            layer_ref = self._rmsnorm(layer_ref, norm_weight)
+                        torch.cuda.synchronize()
+
+                        torch.testing.assert_close(
+                            post_out, post_ref, atol=2e-3, rtol=2e-3
+                        )
+                        torch.testing.assert_close(
+                            comb_out, comb_ref, atol=2e-3, rtol=2e-3
+                        )
+                        torch.testing.assert_close(
+                            layer_out, layer_ref, atol=2e-2, rtol=2e-2
+                        )
+
+                        x = (layer_ref.float() * 0.75).to(layer_ref.dtype)
+                        post_ref_out = mhc._mhc_post_torch(
+                            x, residual, post_ref, comb_ref
+                        )
+                        post_out_actual = mhc._try_aiter_mhc_post(
+                            x, residual, post_out, comb_out
+                        )
+                        self.assertIsNotNone(
+                            post_out_actual, "AITER mHC post fell back"
+                        )
+                        torch.testing.assert_close(
+                            post_out_actual, post_ref_out, atol=2e-2, rtol=2e-2
+                        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/layernorm/test_mhc_kernels.py

@@ -1,4 +1,5 @@
 from contextlib import nullcontext
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -12,12 +13,7 @@ register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
 
 @pytest.fixture
 def stated_tp_group():
-    """A TP group for a test that runs in a process without one.
-
-    The production call passes the group *into* `use_symmetric_memory`, so
-    stubbing that context manager does not stop the read -- the argument is
-    evaluated first. Stating it on the context answers every spelling.
-    """
+    """Provide a TP-group placeholder for kernels with mocked symmetric memory."""
     from sglang.srt.runtime_context import get_parallel
 
     with get_parallel().override(tp_group=None):
@@ -25,7 +21,7 @@ def stated_tp_group():
 
 
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
-@pytest.mark.parametrize("num_tokens", [0, 1, 8, 17, 32, 64])
+@pytest.mark.parametrize("num_tokens", [0, 1, 6, 8, 17, 32, 64])
 @pytest.mark.parametrize("use_norm", [False, True])
 def test_mhc_fused_post_pre_matches_unfused(
     monkeypatch, hidden_size, num_tokens, use_norm, stated_tp_group
@@ -34,10 +30,7 @@ def test_mhc_fused_post_pre_matches_unfused(
         pytest.skip("CUDA is required for TileLang mHC kernels")
 
     monkeypatch.setattr(mhc, "is_dsa_prefill_cp_interleave", lambda: False)
-    # This is a single-process kernel unit test with no TP group initialized.
-    # mhc_pre / mhc_fused_post_pre allocate the MoE input in the symmetric-memory
-    # pool, which asks for the TP group; bypassing the allocation is enough, and
-    # then nothing asks. Mirrors the workaround in test_mxfp4_sm90_cutlass.py.
+    # Disable symmetric-memory allocation for this single-process kernel test.
     monkeypatch.setattr(mhc, "use_symmetric_memory", lambda *a, **kw: nullcontext())
     monkeypatch.setattr(mhc, "is_allocation_symmetric", lambda: False)
     torch.manual_seed(0)
@@ -107,6 +100,18 @@ def test_mhc_fused_post_pre_matches_unfused(
         norm_eps=norm_eps,
     )
 
+    if hidden_size == 4096 and num_tokens in (0, 1, 6, 17):
+        _check_glm_boundary(
+            x,
+            residual,
+            post_prev,
+            comb_prev,
+            fn,
+            hc_scale,
+            hc_base,
+            use_norm=use_norm,
+        )
+
     torch.cuda.synchronize()
     if num_tokens == 0:
         assert residual_out.shape == residual.shape
@@ -136,6 +141,72 @@ def test_mhc_fused_post_pre_matches_unfused(
     torch.testing.assert_close(layer_out, layer_ref, atol=layer_atol, rtol=layer_rtol)
 
 
+def _check_glm_boundary(x, residual, post, comb, fn, scale, base, *, use_norm):
+    from sglang.srt.environ import envs
+    from sglang.srt.layers.communicator_mhc import MHCState
+    from sglang.srt.layers.layernorm import RMSNorm
+    from sglang.srt.models.glm5_next import Glm5NextDecoderLayer
+
+    layer = Glm5NextDecoderLayer.__new__(Glm5NextDecoderLayer)
+    torch.nn.Module.__init__(layer)
+    layer.config = SimpleNamespace(
+        mhc=True,
+        hc_mult=4,
+        rms_norm_eps=1e-6,
+        hc_eps=1e-6,
+        hc_sinkhorn_iters=20,
+    )
+    layer.hc_ffn_fn = torch.nn.Parameter(fn)
+    layer.hc_ffn_scale = torch.nn.Parameter(scale)
+    layer.hc_ffn_base = torch.nn.Parameter(base)
+    norm = RMSNorm(x.shape[-1], eps=1e-6).to(x) if use_norm else None
+    states = [
+        MHCState(
+            hc_mult=4,
+            hc_attn_pre=layer.hc_attn_pre,
+            hc_ffn_pre=layer.hc_ffn_pre,
+            hc_post=layer.hc_post,
+            hc_ffn_post_pre=callback,
+            h_res=comb.flatten(1),
+            h_post=post.flatten(1),
+        )
+        for callback in (None, layer.hc_ffn_post_pre)
+    ]
+    # Literal, not derived from the cutoff constant: deriving it makes this a
+    # mirror that stays green when the cutoff moves. None is the empty batch,
+    # which attn_to_mlp short-circuits before reaching the callback.
+    fused_expected = {1: True, 6: True, 17: False}[x.shape[0]] if x.shape[0] else None
+    with envs.SGLANG_OPT_FUSE_MHC_POST_PRE.override(True):
+        if x.shape[0] > 0:
+            declined = (
+                layer.hc_ffn_post_pre(
+                    hidden_states=x,
+                    residual=residual.flatten(1),
+                    h_res=comb.flatten(1),
+                    h_post=post.flatten(1),
+                    out_norm_weight=None,
+                    out_norm_eps=None,
+                )
+                is None
+            )
+            assert declined is not fused_expected, (
+                f"num_tokens={x.shape[0]} fused={not declined}, "
+                f"expected fused={fused_expected}"
+            )
+        outputs = [s.attn_to_mlp(x, residual.flatten(1), norm) for s in states]
+    torch.testing.assert_close(outputs[0][0], outputs[1][0], atol=2e-2, rtol=2e-2)
+    torch.testing.assert_close(outputs[0][1], outputs[1][1], atol=0, rtol=0)
+    torch.testing.assert_close(states[0].h_res, states[1].h_res, atol=1e-3, rtol=1e-3)
+    torch.testing.assert_close(states[0].h_post, states[1].h_post, atol=1e-3, rtol=1e-3)
+    # The next combine must consume the FFN mixing matrices, not attention's.
+    torch.testing.assert_close(
+        states[0].mlp_combine(x, outputs[0][1]),
+        states[1].mlp_combine(x, outputs[1][1]),
+        atol=2e-3,
+        rtol=2e-2,
+    )
+
+
 if __name__ == "__main__":
     import sys
 

test/registered/kernels/ops/mamba/test_lfm_short_conv.py

@@ -17,7 +17,7 @@ from sglang.kernels.ops.mamba.lfm_short_conv import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 PAD_SLOT_ID = -1
 requires_sm90 = unittest.skipUnless(

test/registered/kernels/ops/moe/test_fused_swiglu_epilogue.py

@@ -19,6 +19,7 @@ import pytest
 import torch
 
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import publish_build_topology
 
 register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
@@ -52,12 +53,8 @@ def _runtime_scaffolding():
     if not torch.distributed.is_initialized():
         init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="gloo")
     if not model_parallel_is_initialized():
-        initialize_model_parallel(
-            tensor_model_parallel_size=1,
-            expert_model_parallel_size=1,
-            pipeline_model_parallel_size=1,
-            backend="gloo",
-        )
+        publish_build_topology(tp_size=1, ep_size=1, pp_size=1)
+        initialize_model_parallel(backend="gloo")
 
 
 def _interleave_w13_rows(w13: torch.Tensor) -> torch.Tensor:

test/registered/kernels/ops/moe/test_kpool_topk_transform_fused.py

@@ -0,0 +1,107 @@
+"""Parity coverage for the fused DSA k-pool top-k / pool-expansion / tail JIT kernel.
+
+``fast_kpool_topk_transform_fused`` is the only implementation for the pooled
+group budgets GLM-5.3-Flash uses (``index_topk=2048`` over ``index_kpool=4``
+gives ``group_topk=512``); ``kpool_fp8_index`` has no Python fallback in that
+range, so a build or numerical break here takes the model down rather than
+making it slower.
+
+The radix selector does not specify an output order and DSA attention is
+permutation-invariant over the selected set, so the pooled columns are compared
+as a set. The tail columns are positional and are compared exactly.
+
+Registered for AMD only: the kernel had no direct coverage on any platform, and
+adding CUDA coverage for it is not this change's call to make.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.moe.kpool_topk_transform import fast_kpool_topk_transform_fused
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=60, stage="jit-kernel-unit", runner_config="amd")
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "Test requires a GPU")
+class TestKpoolTopkTransformFused(CustomTestCase):
+    POOL_SIZE = 4
+
+    def _distinct_scores(self, rows: int, groups: int) -> torch.Tensor:
+        """Strictly distinct scores per row, so top-k selection has no ties to break."""
+        return torch.stack(
+            [torch.randperm(groups, dtype=torch.float32) for _ in range(rows)]
+        ).cuda()
+
+    def _expected_tokens(
+        self, score_row: torch.Tensor, group_topk: int
+    ) -> torch.Tensor:
+        """The pooled top-k groups expanded to their ``pool_size`` token ids."""
+        groups = torch.topk(score_row.float().cpu(), group_topk).indices
+        offsets = torch.arange(self.POOL_SIZE, dtype=torch.int64)
+        return (groups.unsqueeze(1) * self.POOL_SIZE + offsets).reshape(-1)
+
+    def _run(self, rows, groups, topk, seq_lens_host=None):
+        torch.manual_seed(0)
+        score = self._distinct_scores(rows, groups)
+        lengths = torch.full((rows,), groups, dtype=torch.int32, device="cuda")
+        seq_lens = (
+            torch.tensor(seq_lens_host, dtype=torch.int32, device="cuda")
+            if seq_lens_host is not None
+            else None
+        )
+        out = fast_kpool_topk_transform_fused(
+            score=score,
+            lengths=lengths,
+            pool_size=self.POOL_SIZE,
+            topk=topk,
+            seq_lens=seq_lens,
+        )
+        return score, out.cpu()
+
+    def _assert_pooled_columns(self, score, out, topk):
+        group_topk = topk // self.POOL_SIZE
+        for row in range(score.shape[0]):
+            selected = out[row, :topk]
+            expected = self._expected_tokens(score[row], group_topk)
+            self.assertEqual(
+                sorted(selected.tolist()),
+                sorted(expected.tolist()),
+                msg=f"row {row}: selected token set differs from torch.topk",
+            )
+
+    def test_group_topk_512_matches_reference(self):
+        # GLM-5.3-Flash: index_topk=2048 over index_kpool=4.
+        score, out = self._run(rows=2, groups=1024, topk=2048)
+        self._assert_pooled_columns(score, out, topk=2048)
+
+    def test_group_topk_128_matches_reference(self):
+        score, out = self._run(rows=2, groups=512, topk=512)
+        self._assert_pooled_columns(score, out, topk=512)
+
+    def test_tail_columns_hold_the_trailing_partial_pool(self):
+        groups, topk = 1024, 2048
+        for extra in range(self.POOL_SIZE):
+            with self.subTest(tail=extra):
+                seq_len = groups * self.POOL_SIZE + extra
+                score, out = self._run(
+                    rows=2, groups=groups, topk=topk, seq_lens_host=[seq_len] * 2
+                )
+                self._assert_pooled_columns(score, out, topk=topk)
+                expected_tail = [seq_len - extra + i for i in range(extra)]
+                expected_tail += [-1] * (self.POOL_SIZE - 1 - extra)
+                for row in range(out.shape[0]):
+                    self.assertEqual(out[row, topk:].tolist(), expected_tail)
+
+    def test_output_width_carries_the_tail_columns(self):
+        # kpool_fp8_index feeds this width straight into the page-table transform,
+        # so it is 2048 + 3 = 2051 for GLM-5.3-Flash rather than a round 2048.
+        topk = 2048
+        _, out = self._run(rows=1, groups=1024, topk=topk, seq_lens_host=[1024 * 4 + 1])
+        self.assertEqual(tuple(out.shape), (1, topk + self.POOL_SIZE - 1))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/moe/test_minimax_quant_scatter.py

@@ -37,12 +37,7 @@ dev = "cuda"
 
 @pytest.fixture
 def stated_tp_group():
-    """A TP group for a test that runs in a process without one.
-
-    The production call passes the group *into* `use_symmetric_memory`, so
-    stubbing that context manager does not stop the read -- the argument is
-    evaluated first. Stating it on the context answers every spelling.
-    """
+    """Provide a TP-group placeholder for kernels with mocked symmetric memory."""
     from sglang.srt.runtime_context import get_parallel
 
     with get_parallel().override(tp_group=None):

test/registered/kernels/ops/speculative/test_draft_decode_window.py

@@ -0,0 +1,439 @@
+"""Unit test: windowed draft-decode KV index builder (StreamingLLM sink + window).
+
+Covers the ``window_size`` / ``sink_size`` additions to
+``generate_draft_decode_kv_indices`` (the shared kv-index kernel used by the
+built-in MTP/NEXTN + EAGLE draft-decode path on BOTH the Triton and FlashInfer
+draft attention backends).
+
+Properties checked:
+
+  1. IDENTITY (off-by-default is a no-op) -- ``window_size == 0`` reproduces the
+     full-KV read plan. Both the CSR offsets (``kv_indptr``) and the gathered
+     slots (``kv_indices``) are checked against an independent closed-form oracle
+     derived only from the input ``seq_lens`` (NOT from the kernel's packing
+     logic), so an offset bug is not mirrored by the oracle. This is the
+     losslessness/regression guarantee.
+  2. KV-INDPTR lengths -- with a window each draft-decode step keeps
+     ``min(seq_len, sink + window)`` base tokens + ``it + 1`` never-windowed tree
+     tokens; compared to a clamp+cumsum oracle. Parametrized past the
+     256K/512K/1M thresholds (incl. the perf-smoke W4032/S64 params).
+  3. CONTENT -- the actual gathered slots are, in order,
+     ``[first sink] + [most-recent window] + [draft tree]`` (the StreamingLLM
+     layout), read back at buffer offset 0 (num_seqs=1, topk=1).
+  3b. TREE (``topk > 1``) -- with a tree draft every ``(request, topk)`` slot gets
+     its own copy of the windowed base list plus its own branch of tree tokens.
+     The windowed per-slot write offset is ``topk_id * (kept + it + 1)``, so a
+     windowing bug there would silently overlap neighbouring branches while the
+     lengths still look right. Slices are located by the *oracle* ``kv_indptr``
+     (independent closed form) and their content checked, which ties the offset
+     and the CSR bounds together. Covered for ``page_size == 1`` and for the
+     ``page_size > 1 and topk > 1`` paged-tree branch, where the tree tokens are
+     read from the UNCAPPED ``seq_len`` but written at the capped offset.
+
+  4. RESOLUTION -- ``resolve_draft_decode_window`` maps server args to the
+     kernel's (window, sink), and returns (0, 0) for a draft model that already
+     has a sliding window of its own (that per-layer window wins over the flag,
+     since this index builder emits one KV list for all draft layers). It warns
+     only when it substitutes a different window than the one requested.
+  5. BACKEND WIRING -- both ``TritonMultiStepDraftBackend`` and
+     ``FlashInferMultiStepDraftBackend`` resolve the pair in ``__init__`` and
+     forward it to the kernel, so windowing is honored identically on both
+     backends (CPU-only guard).
+
+The kernel is imported normally from ``sglang`` (no by-file-path staging), so it
+tracks the installed tree.
+"""
+
+import inspect
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+from sglang.srt.speculative import spec_utils
+from sglang.srt.speculative.spec_utils import (
+    generate_draft_decode_kv_indices,
+    resolve_draft_decode_window,
+)
+from sglang.srt.utils import next_power_of_2
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+_HAS_CUDA = torch.cuda.is_available()
+
+
+# --- Harness: mirror {Triton,FlashInfer}MultiStepDraftBackend.common_template ---
+def _make_inputs(seq_lens, num_steps, topk, page_size=1, pool_len=None):
+    dev = "cuda"
+    num_seqs = len(seq_lens)
+    bs = num_seqs * topk
+    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64, device=dev)
+    max_seq = int(seq_lens_t.max().item())
+    if pool_len is None:
+        # Room for every branch's tree tokens: topk*num_steps slots (page_size==1)
+        # or up to topk*num_new_pages*page_size past the prefix's last page.
+        pool_len = max_seq + topk * (num_steps + page_size) + 16
+    # Unique physical slot ids per (req, position) so gather checks are unambiguous.
+    req_to_token = (
+        torch.arange(num_seqs * pool_len, dtype=torch.int64, device=dev).reshape(
+            num_seqs, pool_len
+        )
+        + 1000
+    )
+    req_pool_indices = torch.arange(num_seqs, dtype=torch.int64, device=dev)
+    positions = torch.empty(bs, dtype=torch.int64, device=dev)
+    for bid in range(num_seqs):
+        for tk in range(topk):
+            positions[bid * topk + tk] = seq_lens[bid]
+
+    width = bs * (max_seq + num_steps) + 16
+    kv_indices = torch.zeros((num_steps, width), dtype=torch.int64, device=dev)
+    kv_indptr = torch.zeros((num_steps, bs + 1), dtype=torch.int64, device=dev)
+    return dict(
+        req_pool_indices=req_pool_indices,
+        req_to_token=req_to_token,
+        seq_lens=seq_lens_t,
+        positions=positions,
+        kv_indices=kv_indices,
+        kv_indptr=kv_indptr,
+        pool_len=pool_len,
+        page_size=page_size,
+        num_seqs=num_seqs,
+        topk=topk,
+        num_steps=num_steps,
+        bs=bs,
+    )
+
+
+def _run(io, window_size=0, sink_size=0):
+    kv_indices = io["kv_indices"].clone()
+    kv_indptr = io["kv_indptr"].clone()
+    generate_draft_decode_kv_indices[(io["num_steps"], io["num_seqs"], io["topk"])](
+        io["req_pool_indices"],
+        io["req_to_token"],
+        io["seq_lens"],
+        kv_indices,
+        kv_indptr,
+        io["positions"],
+        io["pool_len"],
+        kv_indices.shape[1],
+        kv_indptr.shape[1],
+        next_power_of_2(io["num_seqs"]),
+        next_power_of_2(io["num_steps"]),
+        next_power_of_2(io["bs"]),
+        io["page_size"],
+        window_size,
+        sink_size,
+    )
+    torch.cuda.synchronize()
+    return kv_indices, kv_indptr
+
+
+def _expected_kv_indptr(io, cap=None):
+    """Closed form: kv_indptr[it][zid] = sum(clamp(pos[:zid], cap)) + zid*(it+1)."""
+    pos = io["positions"].to("cpu")
+    if cap is not None:
+        pos = torch.clamp(pos, max=cap)
+    csum = torch.cat([torch.zeros(1, dtype=torch.int64), torch.cumsum(pos, 0)])
+    out = torch.zeros((io["num_steps"], io["bs"] + 1), dtype=torch.int64)
+    for it in range(io["num_steps"]):
+        for zid in range(io["bs"] + 1):
+            out[it, zid] = csum[zid] + zid * (it + 1)
+    return out
+
+
+def _expected_base(r2t, seq_len, window, sink):
+    """StreamingLLM base gather for one request (window==0 => full context)."""
+    if window == 0:
+        return r2t[0:seq_len]
+    cap = window + sink
+    kept = min(seq_len, cap)
+    s_eff = min(sink, seq_len)
+    recent_start = seq_len - (kept - s_eff)
+    return torch.cat([r2t[0:s_eff], r2t[recent_start : recent_start + (kept - s_eff)]])
+
+
+def _expected_tree(r2t, seq_len, topk_id, topk, num_steps, n_tree, page_size):
+    """This branch's tree tokens, always read from the UNCAPPED prefix end.
+
+    Mirrors the kernel's two extend layouts (linear, and the paged per-topk page
+    stride), both of which windowing must leave untouched.
+    """
+    if page_size == 1 or topk == 1:
+        start = seq_len + topk_id * num_steps
+    else:
+        last_page_len = seq_len % page_size
+        num_new_pages_per_topk = (
+            last_page_len + num_steps + page_size - 1
+        ) // page_size
+        start = (
+            (seq_len // page_size) * page_size
+            + topk_id * num_new_pages_per_topk * page_size
+            + last_page_len
+        )
+    return r2t[start : start + n_tree]
+
+
+@unittest.skipUnless(_HAS_CUDA, "draft-decode kv-index kernel requires a CUDA GPU")
+class TestDraftDecodeWindowKernel(unittest.TestCase):
+    def test_window_zero_is_full_kv(self):
+        """window_size==0 => original full-KV read plan (indptr + content)."""
+        for seq_lens, num_steps, topk in [
+            ([128], 4, 1),
+            ([300, 130, 517], 5, 1),
+            ([1024, 777], 4, 2),
+        ]:
+            with self.subTest(seq_lens=seq_lens, num_steps=num_steps, topk=topk):
+                io = _make_inputs(seq_lens, num_steps, topk)
+                self._assert_slot_layout(io, 0, 0)
+
+    def test_windowed_kv_indptr_lengths(self):
+        for seq_lens, window, sink in [
+            ([5000, 2048, 900], 1024, 64),
+            ([5000, 2048, 900], 1024, 0),  # pure recent window
+            ([100, 200], 4096, 64),  # cap exceeds seq_len -> keep all
+            # long-context regression (perf smoke used W4032/S64 past 256K/512K/1M)
+            ([1029306], 4032, 64),
+            ([1048576, 300000, 70000], 4032, 64),
+            ([262145, 131072], 4032, 64),
+        ]:
+            with self.subTest(seq_lens=seq_lens, window=window, sink=sink):
+                io = _make_inputs(seq_lens, num_steps=4, topk=1)
+                _, indptr = _run(io, window_size=window, sink_size=sink)
+                expected = _expected_kv_indptr(io, cap=window + sink)
+                self.assertTrue(torch.equal(indptr.cpu(), expected))
+
+    def test_windowed_content_sink_plus_recent(self):
+        for seq_len, window, sink in [
+            (5000, 1024, 64),
+            (65536, 4032, 64),
+            (262145, 4032, 64),  # just past native 262144
+            (524288, 4032, 64),
+            (1029306, 4032, 64),  # exact perf-smoke seq_len (W4032/S64)
+        ]:
+            with self.subTest(seq_len=seq_len, window=window, sink=sink):
+                io = _make_inputs([seq_len], num_steps=4, topk=1)
+                kv_indices, _ = _run(io, window_size=window, sink_size=sink)
+                r2t = io["req_to_token"][0].cpu()
+                kept = min(seq_len, window + sink)
+                expected_base = _expected_base(r2t, seq_len, window, sink)
+                for step in range(io["num_steps"]):
+                    row = kv_indices[step].cpu()
+                    self.assertTrue(
+                        torch.equal(row[0:kept], expected_base),
+                        f"base gather wrong at step {step} (seq_len={seq_len})",
+                    )
+                    n_tree = step + 1
+                    self.assertTrue(
+                        torch.equal(
+                            row[kept : kept + n_tree], r2t[seq_len : seq_len + n_tree]
+                        ),
+                        f"tree tokens wrong at step {step} (seq_len={seq_len})",
+                    )
+
+    def _assert_slot_layout(self, io, window, sink):
+        """Every (request, topk) slot holds [sink + recent] + its own tree branch.
+
+        Slots are located by the closed-form kv_indptr oracle, so this also pins
+        the per-slot write offset (topk_id * (kept + iters)). ``topk == 1`` is a
+        one-branch tree and ``window == 0`` asks for the whole prefix, so this
+        covers the identity case too.
+        """
+        cap = None if window == 0 else window + sink
+        kv_indices, indptr = _run(io, window_size=window, sink_size=sink)
+        expected_indptr = _expected_kv_indptr(io, cap=cap)
+        self.assertTrue(torch.equal(indptr.cpu(), expected_indptr))
+
+        topk, num_steps = io["topk"], io["num_steps"]
+        seq_lens = io["seq_lens"].tolist()
+        for step in range(num_steps):
+            row = kv_indices[step].cpu()
+            n_tree = step + 1
+            for bid, seq_len in enumerate(seq_lens):
+                r2t = io["req_to_token"][bid].cpu()
+                kept = seq_len if cap is None else min(seq_len, cap)
+                for topk_id in range(topk):
+                    zid = bid * topk + topk_id
+                    start = int(expected_indptr[step, zid])
+                    where = f"step={step} bid={bid} topk_id={topk_id}"
+                    self.assertTrue(
+                        torch.equal(
+                            row[start : start + kept],
+                            _expected_base(r2t, seq_len, window, sink),
+                        ),
+                        f"base gather wrong at {where}",
+                    )
+                    self.assertTrue(
+                        torch.equal(
+                            row[start + kept : start + kept + n_tree],
+                            _expected_tree(
+                                r2t,
+                                seq_len,
+                                topk_id,
+                                topk,
+                                num_steps,
+                                n_tree,
+                                io["page_size"],
+                            ),
+                        ),
+                        f"tree tokens wrong at {where}",
+                    )
+
+    def test_tree_topk_layout(self):
+        """topk > 1 (tree draft): per-branch slices, windowed and unwindowed."""
+        for seq_lens, num_steps, topk, window, sink in [
+            ([600], 4, 2, 128, 16),
+            ([600, 250], 3, 4, 128, 16),
+            ([5000, 900], 4, 2, 1024, 64),
+            ([5000, 900], 4, 2, 1024, 0),  # pure recent window
+            ([300, 130], 3, 2, 4096, 64),  # cap exceeds seq_len -> keep all
+            ([65536, 1029306], 3, 2, 4032, 64),  # long-context tree
+            ([600, 250], 3, 4, 0, 0),  # off-by-default control
+        ]:
+            with self.subTest(seq_lens=seq_lens, topk=topk, window=window, sink=sink):
+                io = _make_inputs(seq_lens, num_steps, topk)
+                self._assert_slot_layout(io, window, sink)
+
+    def test_tree_topk_paged(self):
+        """page_size > 1 AND topk > 1: the paged-tree extend branch.
+
+        The tree tokens are read from the uncapped prefix end with a per-topk page
+        stride and stored at the CAPPED offset, so this is the one place where the
+        windowed and unwindowed lengths must both be respected in one statement.
+        """
+        for seq_lens, num_steps, topk, page_size, window, sink in [
+            ([600], 4, 2, 4, 128, 16),
+            ([601, 255], 3, 2, 8, 128, 16),  # prefix not page-aligned
+            ([604, 256], 3, 4, 4, 128, 0),  # page-aligned prefix, no sink
+            ([5000, 900], 4, 2, 16, 1024, 64),
+            ([601, 255], 3, 2, 8, 0, 0),  # off-by-default control
+        ]:
+            with self.subTest(
+                seq_lens=seq_lens, topk=topk, page_size=page_size, window=window
+            ):
+                io = _make_inputs(seq_lens, num_steps, topk, page_size=page_size)
+                self._assert_slot_layout(io, window, sink)
+
+
+class TestResolveDraftDecodeWindow(unittest.TestCase):
+    """CPU-only: server args -> (window, sink), incl. the opt-out for an SWA draft."""
+
+    @staticmethod
+    def _runner(window=None, sink=None, native_window=None, declared=None):
+        server_args = SimpleNamespace(
+            speculative_draft_window_size=window,
+            speculative_draft_sink_size=sink,
+        )
+        if declared is not None:
+            server_args._resolved_overrides = (
+                ("handle_speculative_decoding", declared),
+            )
+        return SimpleNamespace(
+            server_args=server_args,
+            sliding_window_size=native_window,
+        )
+
+    def test_unset_is_full_attention(self):
+        self.assertEqual(resolve_draft_decode_window(self._runner()), (0, 0))
+
+    def test_window_only(self):
+        self.assertEqual(
+            resolve_draft_decode_window(self._runner(window=4032)), (4032, 0)
+        )
+
+    def test_window_and_sink(self):
+        self.assertEqual(
+            resolve_draft_decode_window(self._runner(window=4032, sink=64)), (4032, 64)
+        )
+
+    def test_declared_values_win_over_the_raw_fields(self):
+        """handle_speculative_decoding declares these fields instead of assigning
+        them, so a raw field read would answer with the unvalidated input."""
+        runner = self._runner(
+            window=None,
+            sink=None,
+            declared={
+                "speculative_draft_window_size": 4032,
+                "speculative_draft_sink_size": 64,
+            },
+        )
+        self.assertEqual(resolve_draft_decode_window(runner), (4032, 64))
+
+    def test_native_sliding_window_draft_opts_out(self):
+        """A draft with its own window keeps it: the flag must not override it."""
+        with mock.patch.object(spec_utils.logger, "warning") as warn:
+            self.assertEqual(
+                resolve_draft_decode_window(
+                    self._runner(window=4032, sink=64, native_window=1024)
+                ),
+                (0, 0),
+            )
+        warn.assert_called_once()
+
+    def test_native_window_equal_to_flag_is_quiet(self):
+        """LlamaForCausalLMEagle3 routes this flag into its own window, so an equal
+        native window is the requested one applied per layer, not a conflict."""
+        with mock.patch.object(spec_utils.logger, "warning") as warn:
+            self.assertEqual(
+                resolve_draft_decode_window(
+                    self._runner(window=4032, sink=64, native_window=4032)
+                ),
+                (0, 0),
+            )
+        warn.assert_not_called()
+
+    def test_no_native_window_still_windows(self):
+        """Guard must not fire for full-attention drafts (0 / None / absent)."""
+        for native in (None, 0):
+            with self.subTest(native_window=native):
+                self.assertEqual(
+                    resolve_draft_decode_window(
+                        self._runner(window=4032, sink=64, native_window=native)
+                    ),
+                    (4032, 64),
+                )
+        runner = self._runner(window=4032, sink=64)
+        del runner.sliding_window_size
+        self.assertEqual(resolve_draft_decode_window(runner), (4032, 64))
+
+
+class TestDraftDecodeWindowBackendWiring(unittest.TestCase):
+    """CPU-only guard: both draft backends must forward window/sink to the kernel.
+
+    The kernel tests above prove correctness once the args reach the kernel; this
+    guards the per-backend plumbing (init resolution + kernel passthrough) so it
+    cannot silently regress on either backend without a GPU-loaded model.
+    """
+
+    def _assert_backend_wires(self, cls):
+        init_src = inspect.getsource(cls.__init__)
+        self.assertIn("resolve_draft_decode_window", init_src)
+        self.assertIn("self.draft_window_size", init_src)
+        self.assertIn("self.draft_sink_size", init_src)
+        tmpl_src = inspect.getsource(cls.common_template)
+        self.assertIn("self.draft_window_size", tmpl_src)
+        self.assertIn("self.draft_sink_size", tmpl_src)
+
+    def test_triton_backend_wires_window_sink(self):
+        try:
+            from sglang.srt.layers.attention.triton_backend import (
+                TritonMultiStepDraftBackend,
+            )
+        except ImportError as e:  # pragma: no cover
+            self.skipTest(f"triton backend import unavailable: {e}")
+        self._assert_backend_wires(TritonMultiStepDraftBackend)
+
+    def test_flashinfer_backend_wires_window_sink(self):
+        try:
+            from sglang.srt.layers.attention.flashinfer_backend import (
+                FlashInferMultiStepDraftBackend,
+            )
+        except ImportError as e:  # pragma: no cover
+            self.skipTest(f"flashinfer backend import unavailable: {e}")
+        self._assert_backend_wires(FlashInferMultiStepDraftBackend)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/speculative/test_gather_spec_extras.py

@@ -1,6 +1,6 @@
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
-register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-small")
 register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)
 
 import unittest

test/registered/kernels/ops/speculative/test_spec_topk1.py

@@ -1,6 +1,6 @@
 from sglang.test.ci.ci_register import register_cuda_ci
 
-register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-small")
 
 import unittest
 

test/registered/kv_canary/test_self_unit_capture_refusal.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import unittest
+from unittest import mock
+
+import torch
+
+from sglang.srt.kv_canary import api
+from sglang.srt.kv_canary.api import torch_reference_conflicts_with_decode_graph
+from sglang.srt.model_executor.cuda_graph_config import (
+    Backend,
+    CudaGraphConfig,
+    PhaseConfig,
+)
+from sglang.srt.runtime_context import get_context
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+
+class TestTorchReferenceConflictsWithDecodeGraph(CustomTestCase):
+    """The refusal that keeps a graph-captured torch reference from passing silently.
+
+    The reference path does host work and D2H, so its launches leave nothing in a
+    captured decode graph and every replay verifies clean. Each case below pins one
+    branch of the gate; the platform capability is patched rather than probed so the
+    CPU lane exercises all four.
+    """
+
+    def _publish_decode_backend(self, backend: str) -> None:
+        override = get_context().override_server_args(
+            cuda_graph_config=CudaGraphConfig(decode=PhaseConfig(backend=backend))
+        )
+        override.install()
+        self.addCleanup(override.restore)
+
+    def _patch_graph_support(self, supported: bool) -> None:
+        patcher = mock.patch.object(
+            api.current_platform, "support_cuda_graph", return_value=supported
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
+
+    def test_reference_device_with_captured_decode_conflicts(self) -> None:
+        self._patch_graph_support(True)
+        self._publish_decode_backend(Backend.FULL)
+        self.assertTrue(
+            torch_reference_conflicts_with_decode_graph(torch.device("xpu"))
+        )
+
+    def test_reference_device_with_decode_graph_disabled_is_allowed(self) -> None:
+        self._patch_graph_support(True)
+        self._publish_decode_backend(Backend.DISABLED)
+        self.assertFalse(
+            torch_reference_conflicts_with_decode_graph(torch.device("xpu"))
+        )
+
+    def test_platform_without_graph_capture_is_allowed(self) -> None:
+        """A device that never captures (CPU) keeps canary on the reference path."""
+        self._patch_graph_support(False)
+        self._publish_decode_backend(Backend.FULL)
+        self.assertFalse(
+            torch_reference_conflicts_with_decode_graph(torch.device("cpu"))
+        )
+
+    def test_cuda_device_is_never_refused(self) -> None:
+        """CUDA/HIP run the real kernels, so the gate must not fire on them."""
+        self._patch_graph_support(True)
+        self._publish_decode_backend(Backend.FULL)
+        self.assertFalse(
+            torch_reference_conflicts_with_decode_graph(torch.device("cuda"))
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kv_canary/test_self_unit_future_tensor.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
 import unittest
-from typing import cast
 
 import torch
 
 from sglang.srt.kv_canary.runner.future_tensor import FutureTensors
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.srt.utils import create_device_stream, get_current_device_stream_fast
+from sglang.test.ci.ci_register import (
+    register_amd_ci,
+    register_cuda_ci,
+    register_xpu_ci,
+)
+from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=10, stage="extra-a", runner_config="1-gpu-small")
 register_amd_ci(est_time=20, suite="extra-a-test-1-gpu-small-amd")
+register_xpu_ci(est_time=20, suite="stage-b-test-1-gpu-xpu")
 
 
 class _FakeEvent:
@@ -22,32 +28,30 @@ class _FakeEvent:
 
 
 class TestFutureTensors(CustomTestCase):
-    def test_cuda_stage_then_wait_returns_host_copy(self) -> None:
-        """Verify staged CUDA tensors are copied back on wait."""
-        device = torch.device("cuda")
-        alt_stream = torch.cuda.Stream(device=device)
-        default_stream = torch.cuda.current_stream(device)
+    def test_device_stage_then_wait_returns_host_copy(self) -> None:
+        """Verify staged device tensors are copied back on wait."""
+        alt_stream = create_device_stream(DEFAULT_DEVICE)
+        default_stream = get_current_device_stream_fast()
         self.assertNotEqual(alt_stream.stream_id, default_stream.stream_id)
 
-        src_first = torch.tensor([41], dtype=torch.int32, device=device)
+        src_first = torch.tensor([41], dtype=torch.int32, device=DEFAULT_DEVICE)
         future_first = FutureTensors.device_to_host(
             xs_device=src_first, d2h_stream=alt_stream
         )
         result_first = future_first.wait()
         self.assertEqual(int(result_first.item()), 41)
 
-        src_second = torch.tensor([97], dtype=torch.int32, device=device)
+        src_second = torch.tensor([97], dtype=torch.int32, device=DEFAULT_DEVICE)
         future_second = FutureTensors.device_to_host(
             xs_device=src_second, d2h_stream=alt_stream
         )
         result_second = future_second.wait()
         self.assertEqual(int(result_second.item()), 97)
 
-    def test_cuda_pinned_when_stream_is_provided(self) -> None:
-        """Verify CUDA staging uses pinned host memory with a stream."""
-        device = torch.device("cuda")
-        alt_stream = torch.cuda.Stream(device=device)
-        src = torch.tensor([5], dtype=torch.int32, device=device)
+    def test_device_pinned_when_stream_is_provided(self) -> None:
+        """Verify device staging uses pinned host memory with a stream."""
+        alt_stream = create_device_stream(DEFAULT_DEVICE)
+        src = torch.tensor([5], dtype=torch.int32, device=DEFAULT_DEVICE)
         future = FutureTensors.device_to_host(xs_device=src, d2h_stream=alt_stream)
         staged_tensors = [
             v for v in future._data.values() if isinstance(v, torch.Tensor)
@@ -56,12 +60,11 @@ class TestFutureTensors(CustomTestCase):
         self.assertTrue(all(t.is_pinned() for t in staged_tensors))
         self.assertEqual(int(future.wait().item()), 5)
 
-    def test_cuda_each_call_allocates_fresh_host(self) -> None:
-        """Verify each CUDA staging call owns a fresh host buffer."""
-        device = torch.device("cuda")
-        alt_stream = torch.cuda.Stream(device=device)
-        src_a = torch.tensor([13], dtype=torch.int32, device=device)
-        src_b = torch.tensor([29], dtype=torch.int32, device=device)
+    def test_device_each_call_allocates_fresh_host(self) -> None:
+        """Verify each device staging call owns a fresh host buffer."""
+        alt_stream = create_device_stream(DEFAULT_DEVICE)
+        src_a = torch.tensor([13], dtype=torch.int32, device=DEFAULT_DEVICE)
+        src_b = torch.tensor([29], dtype=torch.int32, device=DEFAULT_DEVICE)
         future_a = FutureTensors.device_to_host(xs_device=src_a, d2h_stream=alt_stream)
         future_b = FutureTensors.device_to_host(xs_device=src_b, d2h_stream=alt_stream)
         ptrs_a = {
@@ -77,11 +80,10 @@ class TestFutureTensors(CustomTestCase):
 
     def test_dict_of_all_tensors_roundtrip(self) -> None:
         """Verify a dict of multiple tensors round-trips entry-by-entry."""
-        device = torch.device("cuda")
-        stream = torch.cuda.Stream(device=device)
+        stream = create_device_stream(DEFAULT_DEVICE)
         src = {
-            "x": torch.tensor([11, 22], dtype=torch.int64, device=device),
-            "y": torch.tensor([99], dtype=torch.int32, device=device),
+            "x": torch.tensor([11, 22], dtype=torch.int64, device=DEFAULT_DEVICE),
+            "y": torch.tensor([99], dtype=torch.int32, device=DEFAULT_DEVICE),
         }
         future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
         out = future.wait()
@@ -93,14 +95,13 @@ class TestFutureTensors(CustomTestCase):
 
     def test_dict_mixes_tensor_and_passthrough(self) -> None:
         """Verify non-tensor dict entries ride through verbatim alongside staging."""
-        device = torch.device("cuda")
-        stream = torch.cuda.Stream(device=device)
+        stream = create_device_stream(DEFAULT_DEVICE)
         sentinel_obj = {"nested": [1, 2, 3]}
         src = {
             "step": 42,
             "label": "decode",
             "extra": sentinel_obj,
-            "counter": torch.tensor([7], dtype=torch.int32, device=device),
+            "counter": torch.tensor([7], dtype=torch.int32, device=DEFAULT_DEVICE),
         }
         future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
         out = future.wait()
@@ -113,9 +114,8 @@ class TestFutureTensors(CustomTestCase):
 
     def test_dict_passthrough_preserves_tensor_value(self) -> None:
         """Verify tensors share device memory but non-tensor types are not staged."""
-        device = torch.device("cuda")
-        stream = torch.cuda.Stream(device=device)
-        src_tensor = torch.tensor([3], dtype=torch.int32, device=device)
+        stream = create_device_stream(DEFAULT_DEVICE)
+        src_tensor = torch.tensor([3], dtype=torch.int32, device=DEFAULT_DEVICE)
         src = {"step": 100, "buf": src_tensor}
         future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
         out = future.wait()
@@ -128,8 +128,7 @@ class TestFutureTensors(CustomTestCase):
 
     def test_dict_without_tensor_raises(self) -> None:
         """Verify a tensor-less dict raises (no device to anchor the d2h sync)."""
-        device = torch.device("cuda")
-        stream = torch.cuda.Stream(device=device)
+        stream = create_device_stream(DEFAULT_DEVICE)
         with self.assertRaises(ValueError):
             FutureTensors.device_to_host(
                 xs_device={"step": 0, "label": "decode"}, d2h_stream=stream
@@ -137,9 +136,8 @@ class TestFutureTensors(CustomTestCase):
 
     def test_wait_called_twice_raises(self) -> None:
         """Verify wait() after the first drain raises (state cleared)."""
-        device = torch.device("cuda")
-        stream = torch.cuda.Stream(device=device)
-        src = torch.tensor([3], dtype=torch.int32, device=device)
+        stream = create_device_stream(DEFAULT_DEVICE)
+        src = torch.tensor([3], dtype=torch.int32, device=DEFAULT_DEVICE)
         future = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream)
         self.assertEqual(int(future.wait().item()), 3)
         with self.assertRaises(RuntimeError):
@@ -149,9 +147,7 @@ class TestFutureTensors(CustomTestCase):
         """Verify wait() syncs the event exactly once and clears internal state."""
         tensor = torch.tensor([1, 2, 3])
         event = _FakeEvent()
-        future = FutureTensors(
-            _data={"x": tensor}, _event=cast(torch.cuda.Event, event)
-        )
+        future = FutureTensors(_data={"x": tensor}, _event=event)
 
         result = future.wait()
         self.assertIs(result["x"], tensor)
@@ -166,11 +162,10 @@ class TestFutureTensors(CustomTestCase):
 
     def test_dict_anchor_picked_from_first_tensor(self) -> None:
         """Verify staging works when the first key is a non-tensor (anchor must scan)."""
-        device = torch.device("cuda")
-        stream = torch.cuda.Stream(device=device)
+        stream = create_device_stream(DEFAULT_DEVICE)
         src = {
             "step": 5,
-            "buf": torch.tensor([17], dtype=torch.int32, device=device),
+            "buf": torch.tensor([17], dtype=torch.int32, device=DEFAULT_DEVICE),
         }
         out = FutureTensors.device_to_host(xs_device=src, d2h_stream=stream).wait()
         self.assertEqual(out["step"], 5)

test/registered/kv_canary/test_self_unit_plan_input.py

@@ -6,15 +6,21 @@ from types import SimpleNamespace
 import torch
 
 from sglang.srt.kv_canary.plan_input import PlanInput
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.ci.ci_register import (
+    register_amd_ci,
+    register_cuda_ci,
+    register_xpu_ci,
+)
 from sglang.test.kv_canary.fixtures import (
     DEFAULT_DEVICE,
+    DEFAULT_DEVICE_MODULE,
     make_forward_batch,
 )
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=9, stage="extra-a", runner_config="1-gpu-small")
 register_amd_ci(est_time=30, suite="extra-a-test-1-gpu-small-amd")
+register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu")
 
 
 def _make_static_plan_input(*, bs_capacity: int, device) -> PlanInput:
@@ -122,7 +128,7 @@ class TestSelfUnitPlanInput(CustomTestCase):
         fb.req_all_ids_lens = torch.tensor([7, 9], dtype=torch.int64, pin_memory=True)
         plan = _make_static_plan_input(bs_capacity=4, device=self.device)
         plan.fill_from_forward_batch(forward_batch=fb)
-        torch.cuda.synchronize()
+        DEFAULT_DEVICE_MODULE.synchronize()
         self.assertEqual(
             plan.req_to_verify_expected_tokens_valid_lens[:2].tolist(), [7, 9]
         )

test/registered/kv_canary/test_self_unit_req_to_expected_token_ids_manager.py

@@ -10,12 +10,21 @@ from sglang.srt.kv_canary.req_to_expected_token_ids_manager import (
     compute_req_all_ids_info,
     populate_req_to_expected_token_ids,
 )
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_forward_batch
+from sglang.test.ci.ci_register import (
+    register_amd_ci,
+    register_cuda_ci,
+    register_xpu_ci,
+)
+from sglang.test.kv_canary.fixtures import (
+    DEFAULT_DEVICE,
+    DEFAULT_DEVICE_MODULE,
+    make_forward_batch,
+)
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=11, stage="extra-a", runner_config="1-gpu-small")
 register_amd_ci(est_time=15, suite="extra-a-test-1-gpu-small-amd")
+register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu")
 
 
 def _make_req(*, origin: list[int], output: list[int]) -> SimpleNamespace:
@@ -92,7 +101,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
         populate_req_to_expected_token_ids(
             forward_batch=fb, req_to_verify_expected_tokens=pool
         )
-        torch.cuda.synchronize()
+        DEFAULT_DEVICE_MODULE.synchronize()
         self.assertTrue(torch.equal(pool, original))
 
     def test_no_op_when_pool_is_none(self) -> None:
@@ -117,7 +126,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
         populate_req_to_expected_token_ids(
             forward_batch=fb, req_to_verify_expected_tokens=pool
         )
-        torch.cuda.synchronize()
+        DEFAULT_DEVICE_MODULE.synchronize()
         self.assertTrue(torch.equal(pool, original))
 
     def test_raises_when_lens_length_mismatches_batch_size(self) -> None:
@@ -154,7 +163,7 @@ class TestPopulateReqToExpectedTokenIds(CustomTestCase):
         populate_req_to_expected_token_ids(
             forward_batch=fb, req_to_verify_expected_tokens=pool
         )
-        torch.cuda.synchronize()
+        DEFAULT_DEVICE_MODULE.synchronize()
 
         pool_cpu = pool.cpu()
         self.assertEqual(pool_cpu[1, :3].tolist(), [10, 20, 30])

test/registered/kv_canary/test_self_unit_runner_swa_divergence.py

@@ -15,24 +15,28 @@ from sglang.srt.kv_canary.runner.swa_divergence import (
     SwaDivergenceReporter,
     compute_swa_full_idx_divergence,
 )
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.kv_canary.fixtures import make_buffer_group
+from sglang.srt.utils import create_device_stream
+from sglang.test.ci.ci_register import (
+    register_amd_ci,
+    register_cuda_ci,
+    register_xpu_ci,
+)
+from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_buffer_group
 from sglang.test.kv_canary.runner_test_base import CanaryManagerTestCase, make_manager
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=11, stage="extra-a", runner_config="1-gpu-small")
 register_amd_ci(est_time=45, suite="extra-a-test-1-gpu-small-amd")
-
-_DEVICE = torch.device("cuda")
+register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
 
 _EMPTY_FORWARD_BATCH = SimpleNamespace(
-    req_pool_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
-    seq_lens=torch.empty(0, dtype=torch.int64, device=_DEVICE),
+    req_pool_indices=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
+    seq_lens=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
 )
 
 
 def _make_verify_plan(value: int) -> VerifyPlan:
-    plan = VerifyPlan.allocate(verify_capacity=4, device=_DEVICE)
+    plan = VerifyPlan.allocate(verify_capacity=4, device=DEFAULT_DEVICE)
     plan.verify_num_valid.copy_(torch.tensor([value], dtype=torch.int32))
     return plan
 
@@ -46,11 +50,13 @@ def _make_req_to_token_pool_stub(req_to_token: torch.Tensor) -> SimpleNamespace:
 
 
 def _make_identity_mapping(size: int) -> torch.Tensor:
-    return torch.arange(size, dtype=torch.int64, device=_DEVICE)
+    return torch.arange(size, dtype=torch.int64, device=DEFAULT_DEVICE)
 
 
 def _make_identity_req_to_token(num_reqs: int, max_seq_len: int) -> torch.Tensor:
-    base = torch.arange(num_reqs * max_seq_len, dtype=torch.int64, device=_DEVICE)
+    base = torch.arange(
+        num_reqs * max_seq_len, dtype=torch.int64, device=DEFAULT_DEVICE
+    )
     return base.view(num_reqs, max_seq_len)
 
 
@@ -83,9 +89,9 @@ def _run_compute(
 
 class TestSwaDivergenceReporter(CustomTestCase):
     def test_swa_divergence_log_emitted(self) -> None:
-        d2h_stream = torch.cuda.Stream(device=_DEVICE)
+        d2h_stream = create_device_stream(DEFAULT_DEVICE)
         stats = SwaDivergenceReporter(
-            device=_DEVICE,
+            device=DEFAULT_DEVICE,
             d2h_stream=d2h_stream,
             interval=10,
             swa_allocator=None,
@@ -96,13 +102,13 @@ class TestSwaDivergenceReporter(CustomTestCase):
         for forward_idx in range(3):
             stats.observe_after_invoke_plan(
                 group=make_buffer_group(
-                    device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
+                    device=DEFAULT_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
                 ),
                 verify_plan=_make_verify_plan(10),
             )
             stats.observe_after_invoke_plan(
                 group=make_buffer_group(
-                    device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
+                    device=DEFAULT_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
                 ),
                 verify_plan=_make_verify_plan(3),
             )
@@ -115,13 +121,13 @@ class TestSwaDivergenceReporter(CustomTestCase):
         # the staged future hangs onto it. forward_ct is now 4.
         stats.observe_after_invoke_plan(
             group=make_buffer_group(
-                device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
+                device=DEFAULT_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
             ),
             verify_plan=_make_verify_plan(10),
         )
         stats.observe_after_invoke_plan(
             group=make_buffer_group(
-                device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
+                device=DEFAULT_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
             ),
             verify_plan=_make_verify_plan(3),
         )
@@ -150,9 +156,9 @@ class TestSwaDivergenceReporter(CustomTestCase):
         self.assertEqual(fields.swa_full_idx_divergence, 0)
 
     def test_swa_divergence_counts_monotonic_increasing(self) -> None:
-        d2h_stream = torch.cuda.Stream(device=_DEVICE)
+        d2h_stream = create_device_stream(DEFAULT_DEVICE)
         stats = SwaDivergenceReporter(
-            device=_DEVICE,
+            device=DEFAULT_DEVICE,
             d2h_stream=d2h_stream,
             interval=10,
             swa_allocator=None,
@@ -188,13 +194,19 @@ class TestSwaDivergenceReporter(CustomTestCase):
             for _ in range(5):
                 stats.observe_after_invoke_plan(
                     group=make_buffer_group(
-                        device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
+                        device=DEFAULT_DEVICE,
+                        kind=PoolKind.FULL,
+                        has_v=False,
+                        num_slots=1,
                     ),
                     verify_plan=_make_verify_plan(7),
                 )
                 stats.observe_after_invoke_plan(
                     group=make_buffer_group(
-                        device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
+                        device=DEFAULT_DEVICE,
+                        kind=PoolKind.SWA,
+                        has_v=False,
+                        num_slots=1,
                     ),
                     verify_plan=_make_verify_plan(2),
                 )
@@ -216,8 +228,8 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
         req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
 
         forward_batch = _make_forward_batch(
-            req_pool_indices=torch.empty(0, dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.empty(0, dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
+            seq_lens=torch.empty(0, dtype=torch.int64, device=DEFAULT_DEVICE),
         )
 
         self.assertEqual(
@@ -234,8 +246,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
         req_to_token = _make_identity_req_to_token(num_reqs=4, max_seq_len=16)
 
         forward_batch = _make_forward_batch(
-            req_pool_indices=torch.tensor([0, 2], dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.tensor([8, 5], dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.tensor(
+                [0, 2], dtype=torch.int64, device=DEFAULT_DEVICE
+            ),
+            seq_lens=torch.tensor([8, 5], dtype=torch.int64, device=DEFAULT_DEVICE),
         )
 
         self.assertEqual(
@@ -256,8 +270,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
         mapping[17] = 60
 
         forward_batch = _make_forward_batch(
-            req_pool_indices=torch.tensor([0, 1], dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.tensor([8, 8], dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.tensor(
+                [0, 1], dtype=torch.int64, device=DEFAULT_DEVICE
+            ),
+            seq_lens=torch.tensor([8, 8], dtype=torch.int64, device=DEFAULT_DEVICE),
         )
 
         self.assertEqual(
@@ -281,8 +297,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
         mapping[7] = 42
 
         forward_batch = _make_forward_batch(
-            req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.tensor(
+                [0], dtype=torch.int64, device=DEFAULT_DEVICE
+            ),
+            seq_lens=torch.tensor([8], dtype=torch.int64, device=DEFAULT_DEVICE),
         )
 
         self.assertEqual(
@@ -302,8 +320,10 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
         mapping[28] = 77
 
         forward_batch = _make_forward_batch(
-            req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.tensor([10], dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.tensor(
+                [0], dtype=torch.int64, device=DEFAULT_DEVICE
+            ),
+            seq_lens=torch.tensor([10], dtype=torch.int64, device=DEFAULT_DEVICE),
         )
 
         self.assertEqual(
@@ -325,12 +345,16 @@ class TestSwaFullIdxDivergenceCompute(CustomTestCase):
         mapping[33] = 100
 
         fb_req0 = _make_forward_batch(
-            req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.tensor(
+                [0], dtype=torch.int64, device=DEFAULT_DEVICE
+            ),
+            seq_lens=torch.tensor([4], dtype=torch.int64, device=DEFAULT_DEVICE),
         )
         fb_req2 = _make_forward_batch(
-            req_pool_indices=torch.tensor([2], dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.tensor([4], dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.tensor(
+                [2], dtype=torch.int64, device=DEFAULT_DEVICE
+            ),
+            seq_lens=torch.tensor([4], dtype=torch.int64, device=DEFAULT_DEVICE),
         )
 
         self.assertEqual(
@@ -363,15 +387,17 @@ class TestSwaDivergenceReporterWithCompute(CustomTestCase):
         mapping[2] = 52
 
         forward_batch = _make_forward_batch(
-            req_pool_indices=torch.tensor([0], dtype=torch.int64, device=_DEVICE),
-            seq_lens=torch.tensor([8], dtype=torch.int64, device=_DEVICE),
+            req_pool_indices=torch.tensor(
+                [0], dtype=torch.int64, device=DEFAULT_DEVICE
+            ),
+            seq_lens=torch.tensor([8], dtype=torch.int64, device=DEFAULT_DEVICE),
         )
 
         swa_allocator = _make_allocator_stub(mapping)
         req_to_token_pool = _make_req_to_token_pool_stub(req_to_token)
-        d2h_stream = torch.cuda.Stream(device=_DEVICE)
+        d2h_stream = create_device_stream(DEFAULT_DEVICE)
         stats = SwaDivergenceReporter(
-            device=_DEVICE,
+            device=DEFAULT_DEVICE,
             d2h_stream=d2h_stream,
             interval=10,
             swa_allocator=swa_allocator,
@@ -379,13 +405,13 @@ class TestSwaDivergenceReporterWithCompute(CustomTestCase):
         )
         stats.observe_after_invoke_plan(
             group=make_buffer_group(
-                device=_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
+                device=DEFAULT_DEVICE, kind=PoolKind.FULL, has_v=False, num_slots=1
             ),
             verify_plan=_make_verify_plan(11),
         )
         stats.observe_after_invoke_plan(
             group=make_buffer_group(
-                device=_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
+                device=DEFAULT_DEVICE, kind=PoolKind.SWA, has_v=False, num_slots=1
             ),
             verify_plan=_make_verify_plan(3),
         )

test/registered/layers/mamba/test_mamba2_mixer.py

@@ -17,6 +17,7 @@ from sglang.srt.distributed.parallel_state import (
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_device, get_device_count
 from sglang.test.ci.ci_register import register_cuda_ci, register_xpu_ci
+from sglang.test.test_utils import publish_build_topology
 
 register_cuda_ci(est_time=30, stage="base-b", runner_config="2-gpu-large")
 register_xpu_ci(est_time=60, suite="nightly-xpu-2-gpu", nightly=True)
@@ -105,7 +106,8 @@ def mixer2_gated_norm_tensor_parallel(
         local_rank=local_rank,
         backend=get_default_distributed_backend(device.type),
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(tp_size=world_size, world_rank=local_rank)
+    initialize_model_parallel()
 
     # create random weights an inputs
     weight = torch.rand((hidden_size,), dtype=dtype, device=device)
@@ -127,8 +129,24 @@ def mixer2_gated_norm_tensor_parallel(
         )
         mixer.weight.weight_loader(mixer.weight, weight)
 
-    # m2 reads tp via get_parallel().tp_size/rank — force it through the context.
-    with get_parallel().override(tp_size=1, tp_rank=0):
+    with get_parallel().override(
+        tp_size=1,
+        tp_rank=0,
+        tp_group=None,
+        attn_tp_size=1,
+        attn_tp_rank=0,
+        attn_tp_group=None,
+        attn_dp_size=1,
+        attn_dp_rank=0,
+        attn_cp_size=1,
+        attn_cp_rank=0,
+        attn_cp_group=None,
+        moe_ep_size=1,
+        moe_ep_rank=0,
+        moe_ep_group=None,
+        moe_dp_size=1,
+        moe_tp_size=1,
+    ):
         # create gated-norm without TP to compute reference
         mixer_single_gpu = m2.Mixer2RMSNormGated(
             full_hidden_size=hidden_size,

test/registered/layers/test_fla_layernorm_guard.py

@@ -1,396 +0,0 @@
-from __future__ import annotations
-
-import socket
-import sys
-from dataclasses import dataclass
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from sglang.kernels.ops.attention.fla.layernorm_gated import (
-    _layer_norm_fwd as layer_norm_fwd,
-)
-from sglang.kernels.ops.attention.fla.layernorm_gated import (
-    layernorm_fn,
-    rms_norm_ref,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(
-    est_time=60,
-    stage="base-b",
-    runner_config="2-gpu-large",
-    disabled="Temporarily disabled",
-)
-
-# Optional dependency in sglang repo; skip collection cleanly if absent.
-custom_all_reduce_utils = pytest.importorskip(
-    "sglang.srt.distributed.device_communicators.custom_all_reduce_utils"
-)
-parallel_state = pytest.importorskip("sglang.srt.distributed.parallel_state")
-
-update_environment_variables = custom_all_reduce_utils.update_environment_variables
-init_distributed_environment = parallel_state.init_distributed_environment
-initialize_model_parallel = parallel_state.initialize_model_parallel
-
-NUM_GPUS = 2
-
-
-def _find_free_port() -> int:
-    # Avoid hard-coded port collisions when pytest runs tests in parallel.
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("localhost", 0))
-        s.listen(1)
-        return int(s.getsockname()[1])
-
-
-def _skip_if_no_cuda_or_not_enough_gpus(required_gpus: int = NUM_GPUS) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA device not available")
-    if torch.cuda.device_count() < required_gpus:
-        pytest.skip(f"Need >= {required_gpus} GPUs, got {torch.cuda.device_count()}")
-
-
-def _skip_if_dtype_unsupported(dtype: torch.dtype) -> None:
-    if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported():
-        pytest.skip("bfloat16 not supported on this CUDA device")
-
-
-def _setup_sglang_distributed(
-    local_rank: int,
-    world_size: int,
-    master_port: int,
-    dtype: torch.dtype,
-) -> torch.device:
-    # Match sglang test style: set per-rank CUDA device + default dtype/device.
-    torch.manual_seed(0)
-    torch.cuda.manual_seed_all(0)
-
-    device = torch.device(f"cuda:{local_rank}")
-    torch.cuda.set_device(device)
-
-    if hasattr(torch, "set_default_device"):
-        torch.set_default_device(device)
-    if hasattr(torch, "set_default_dtype"):
-        torch.set_default_dtype(dtype)
-
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": str(master_port),
-        }
-    )
-
-    init_distributed_environment(
-        world_size=world_size, rank=local_rank, local_rank=local_rank
-    )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
-
-    return device
-
-
-def layer_norm_ref(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor | None,
-    z: torch.Tensor | None = None,
-    eps: float = 1e-6,
-    group_size: int | None = None,
-    norm_before_gate: bool = True,
-    is_rms_norm: bool = False,
-) -> torch.Tensor:
-    """Reference implementation for both LayerNorm and RMSNorm (supports optional gate + group norm)."""
-    if is_rms_norm:
-        return rms_norm_ref(
-            x,
-            weight,
-            bias,
-            z=z,
-            eps=eps,
-            group_size=group_size,
-            norm_before_gate=norm_before_gate,
-            upcast=True,
-        )
-
-    dtype = x.dtype
-    x_f = x.float()
-    w_f = weight.float()
-    b_f = bias.float() if bias is not None else None
-    z_f = z.float() if z is not None else None
-
-    if z_f is not None and not norm_before_gate:
-        x_f = x_f * F.silu(z_f)
-
-    if group_size is None:
-        mean = x_f.mean(dim=-1, keepdim=True)
-        var = (x_f - mean).square().mean(dim=-1, keepdim=True)
-        rstd = torch.rsqrt(var + eps)
-        out = (x_f - mean) * rstd * w_f
-        if b_f is not None:
-            out = out + b_f
-    else:
-        hidden = x_f.shape[-1]
-        assert hidden % group_size == 0
-        ng = hidden // group_size
-        xg = x_f.view(*x_f.shape[:-1], ng, group_size)
-        mean = xg.mean(dim=-1, keepdim=True)
-        var = (xg - mean).square().mean(dim=-1, keepdim=True)
-        rstd = torch.rsqrt(var + eps)
-        xg = (xg - mean) * rstd
-        out = xg.reshape(*x_f.shape[:-1], hidden) * w_f
-        if b_f is not None:
-            out = out + b_f
-
-    if z_f is not None and norm_before_gate:
-        out = out * F.silu(z_f)
-
-    return out.to(dtype)
-
-
-@dataclass(frozen=True)
-class FwdCase:
-    name: str
-    with_gate: bool
-    norm_before_gate: bool
-    group_size: int | None
-    is_rms_norm: bool
-
-
-CASES: list[FwdCase] = [
-    FwdCase(
-        "layernorm",
-        with_gate=False,
-        norm_before_gate=True,
-        group_size=None,
-        is_rms_norm=False,
-    ),
-    FwdCase(
-        "rmsnorm",
-        with_gate=False,
-        norm_before_gate=True,
-        group_size=None,
-        is_rms_norm=True,
-    ),
-    FwdCase(
-        "layernorm_gate_pre",
-        with_gate=True,
-        norm_before_gate=True,
-        group_size=None,
-        is_rms_norm=False,
-    ),
-    FwdCase(
-        "layernorm_gate_post",
-        with_gate=True,
-        norm_before_gate=False,
-        group_size=None,
-        is_rms_norm=False,
-    ),
-    FwdCase(
-        "rmsnorm_gate_pre",
-        with_gate=True,
-        norm_before_gate=True,
-        group_size=None,
-        is_rms_norm=True,
-    ),
-    FwdCase(
-        "group_layernorm",
-        with_gate=False,
-        norm_before_gate=True,
-        group_size=128,
-        is_rms_norm=False,
-    ),
-    FwdCase(
-        "group_rmsnorm",
-        with_gate=False,
-        norm_before_gate=True,
-        group_size=128,
-        is_rms_norm=True,
-    ),
-]
-
-
-@pytest.mark.parametrize("num_tokens", [128])
-@pytest.mark.parametrize("hidden_size", [256])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("case", CASES, ids=lambda c: c.name)
-def test_layernorm_guard_fwd_spawn(
-    num_tokens: int,
-    hidden_size: int,
-    dtype: torch.dtype,
-    case: FwdCase,
-    device: str = "cuda",
-):
-    _skip_if_no_cuda_or_not_enough_gpus(NUM_GPUS)
-    _skip_if_dtype_unsupported(dtype)
-
-    if case.group_size is not None and hidden_size % case.group_size != 0:
-        pytest.skip(
-            f"hidden_size {hidden_size} not divisible by group_size {case.group_size}"
-        )
-
-    master_port = _find_free_port()
-    world_size = NUM_GPUS
-
-    torch.multiprocessing.spawn(
-        _layernorm_guard_fwd_worker,
-        args=(
-            world_size,
-            master_port,
-            num_tokens,
-            hidden_size,
-            dtype,
-            case,
-            device,
-        ),
-        nprocs=world_size,
-        join=True,
-    )
-
-
-def _layernorm_guard_fwd_worker(
-    local_rank: int,
-    world_size: int,
-    master_port: int,
-    num_tokens: int,
-    hidden_size: int,
-    dtype: torch.dtype,
-    case: FwdCase,
-    device: str,
-):
-    device = _setup_sglang_distributed(local_rank, world_size, master_port, dtype)
-
-    with torch.inference_mode():
-        torch.manual_seed(42 + local_rank)
-        torch.cuda.manual_seed_all(42 + local_rank)
-
-        x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
-        z = (
-            torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
-            if case.with_gate
-            else None
-        )
-        weight = torch.randn(hidden_size, dtype=dtype, device=device)
-        bias = (
-            None
-            if case.is_rms_norm
-            else torch.randn(hidden_size, dtype=dtype, device=device)
-        )
-        eps = 1e-6
-
-        out, mean, rstd = layer_norm_fwd(
-            x,
-            weight,
-            bias,
-            eps,
-            z=z,
-            group_size=case.group_size,
-            norm_before_gate=case.norm_before_gate,
-            is_rms_norm=case.is_rms_norm,
-        )
-
-        ref_out = layer_norm_ref(
-            x,
-            weight,
-            bias,
-            z=z,
-            eps=eps,
-            group_size=case.group_size,
-            norm_before_gate=case.norm_before_gate,
-            is_rms_norm=case.is_rms_norm,
-        )
-
-        assert out.shape == x.shape
-        assert out.dtype == x.dtype
-        torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
-
-        # mean/rstd shape checks (same spirit as original vLLM tests)
-        if case.group_size is None:
-            if not case.is_rms_norm:
-                assert mean.shape == (num_tokens,)
-            assert rstd.shape == (num_tokens,)
-        else:
-            ngroups = hidden_size // case.group_size
-            if not case.is_rms_norm:
-                assert mean.shape == (ngroups * num_tokens,)
-            assert rstd.shape == (ngroups * num_tokens,)
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_layernorm_guard_misc_spawn(dtype: torch.dtype, device: str = "cuda"):
-    _skip_if_no_cuda_or_not_enough_gpus(NUM_GPUS)
-    _skip_if_dtype_unsupported(dtype)
-
-    master_port = _find_free_port()
-    world_size = NUM_GPUS
-
-    torch.multiprocessing.spawn(
-        _layernorm_guard_misc_worker,
-        args=(world_size, master_port, dtype, device),
-        nprocs=world_size,
-        join=True,
-    )
-
-
-def _layernorm_guard_misc_worker(
-    local_rank: int,
-    world_size: int,
-    master_port: int,
-    dtype: torch.dtype,
-    device: str,
-):
-    device = _setup_sglang_distributed(local_rank, world_size, master_port, dtype)
-
-    with torch.inference_mode():
-        torch.manual_seed(123 + local_rank)
-        torch.cuda.manual_seed_all(123 + local_rank)
-
-        # 1) rows_per_block-like sizes
-        hidden_size = 1024
-        weight = torch.randn(hidden_size, dtype=dtype, device=device)
-        bias = torch.randn(hidden_size, dtype=dtype, device=device)
-        eps = 1e-6
-        for num_tokens in [513]:
-            x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
-            out, _, _ = layer_norm_fwd(x, weight, bias, eps, z=None, is_rms_norm=False)
-            ref = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
-            torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
-
-        # 2) strided input (slice then contiguous)
-        num_tokens = 128
-        x_large = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
-        x = x_large[:, :hidden_size]
-        x_contig = x.contiguous()
-        out, _, _ = layer_norm_fwd(
-            x_contig, weight, bias, eps, z=None, is_rms_norm=False
-        )
-        ref = layer_norm_ref(x_contig, weight, bias, z=None, eps=eps, is_rms_norm=False)
-        torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
-
-        # 3) provided output buffer
-        num_tokens = 256
-        x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
-        out_buf = torch.empty_like(x)
-        out, _, _ = layer_norm_fwd(
-            x, weight, bias, eps, z=None, out=out_buf, is_rms_norm=False
-        )
-        assert out.data_ptr() == out_buf.data_ptr()
-        ref = layer_norm_ref(x, weight, bias, z=None, eps=eps, is_rms_norm=False)
-        torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
-
-        # 4) multidimensional input via autograd fn
-        for shape in [(4, 16, 1024)]:
-            hs = shape[-1]
-            x = torch.randn(*shape, dtype=dtype, device=device)
-            w = torch.randn(hs, dtype=dtype, device=device)
-            b = torch.randn(hs, dtype=dtype, device=device)
-            out = layernorm_fn(x, w, b, z=None, eps=eps)
-            ref = layer_norm_ref(x, w, b, z=None, eps=eps, is_rms_norm=False)
-            torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/mem_cache/test_post_capture_kv_sizing.py

@@ -47,6 +47,13 @@ class TestPostCaptureKVSizing(CustomTestCase):
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             env={**os.environ, "SGLANG_ENABLE_POST_CAPTURE_KV_SIZING": "1"},
             return_stdout_stderr=(cls.stdout, cls.stderr),
+            other_args=[
+                "--enable-hierarchical-cache",
+                "--hicache-mem-layout",
+                "page_first",
+                "--hicache-size",
+                "1",
+            ],
         )
 
     @classmethod
@@ -78,6 +85,10 @@ class TestPostCaptureKVSizing(CustomTestCase):
             "or the resize path did not run.",
         )
         self.assertGreater(float(m.group(1)), 0)
+        logs = self._server_logs()
+        staging = logs.find("HiCache staging prepared before KV sizing:")
+        self.assertGreaterEqual(staging, 0, "HiCache staging was not prepared")
+        self.assertLess(staging, m.start())
 
     def test_server_info_pool_sized(self):
         info = requests.get(f"{self.base_url}/server_info").json()

test/registered/mock_model/test_e2e_cp.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.mock_model.utils import run_mock_model_bench_serving
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=50, stage="extra-a", runner_config="2-gpu-large")
+
+
+class TestE2EContextParallel(CustomTestCase):
+    def test_cp_prefill_then_decode_no_canary_violation(self) -> None:
+        # CP prefill enters the transformer body directly, while decode calls
+        # the outer model.forward. Both need exactly one canary bracket.
+        run_mock_model_bench_serving(
+            extra_server_args=[
+                "--tp",
+                "2",
+                "--attn-cp-size",
+                "2",
+                "--enable-prefill-cp",
+                "--cp-strategy",
+                "zigzag",
+                "--attention-backend",
+                "fa3",
+                "--kv-canary-real-data",
+                "all",
+                "--mem-fraction-static",
+                "0.2",
+                "--max-total-tokens",
+                "4096",
+                "--max-running-requests",
+                "8",
+                "--context-length",
+                "256",
+                "--cuda-graph-max-bs-decode",
+                "4",
+            ],
+            num_prompts=4,
+            random_input_len=32,
+            random_output_len=8,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/npu/accuracy/deepseek_v4_flash/test_npu_deepseek_v4_flash_w8a8_8p_gpqa.py

@@ -1,3 +1,5 @@
+import os
+import subprocess
 import unittest
 
 from sglang.test.ascend.e2e.test_npu_accuracy_utils import (
@@ -8,6 +10,8 @@ from sglang.test.ascend.e2e.test_npu_performance_utils import (
     DEEPSEEK_V4_FLASH_0731_W8A8_MODEL_PATH,
 )
 from sglang.test.ci.ci_register import register_npu_ci
+from sglang.test.test_utils import DEFAULT_URL_FOR_TEST
+from sglang.utils import wait_for_server
 
 register_npu_ci(
     est_time=3600,
@@ -29,6 +33,7 @@ DEEPSEEK_V4_FLASH_W8A8_DSPARK_8P_ENVS = {
     "HCCL_SOCKET_IFNAME": "lo",
     "GLOO_SOCKET_IFNAME": "lo",
     "HCCL_OP_EXPANSION_MODE": "AIV",
+    "SGLANG_NPU_USE_MULTI_STREAM": "1",
     # skip gpu branch
     "SGLANG_OPT_FP8_WO_A_GEMM": "0",
     "SGLANG_OPT_USE_OVERLAP_STORE_CACHE": "False",
@@ -43,10 +48,13 @@ DEEPSEEK_V4_FLASH_W8A8_DSPARK_8P_ENVS = {
     # DSPARK
     "SGLANG_RAGGED_VERIFY_MODE": "static",
     "SGLANG_DSPARK_FAST_KERNEL": "0",
+    # mtp
+    "SGLANG_ENABLE_SPEC_V2": "1",
+    "SGLANG_ENABLE_OVERLAP_PLAN_STREAM": "1",
     # deepep
     "DEEP_NORMAL_MODE_USE_INT8_QUANT": "1",
-    "DEEPEP_HCCL_BUFFSIZE": "2048",
-    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "64",
+    "DEEPEP_HCCL_BUFFSIZE": "2500",
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "96",
     "DEEPEP_HYBRID_DEPLOYMENT": "1",
     # war barrier
     "SGLANG_ENABLE_WAR_BARRIER": "1",
@@ -66,15 +74,15 @@ DEEPSEEK_V4_FLASH_W8A8_DSPARK_8P_OTHER_ARGS = [
     "--watchdog-timeout",
     9000,
     "--mem-fraction-static",
-    0.62,
+    0.68,
     "--prefill-max-requests",
-    32,
+    192,
     "--max-prefill-tokens",
-    131072,
+    80000,
     "--chunked-prefill-size",
     131072,
     "--max-running-requests",
-    96,
+    192,
     "--dp-size",
     16,
     "--enable-dp-attention",
@@ -87,6 +95,8 @@ DEEPSEEK_V4_FLASH_W8A8_DSPARK_8P_OTHER_ARGS = [
     "--enable-dp-lm-head",
     "--kv-cache-dtype",
     "bfloat16",
+    "--load-balance-method",
+    "round_robin",
     "--speculative-algorithm",
     "DSPARK",
     "--speculative-draft-model-path",
@@ -96,16 +106,17 @@ DEEPSEEK_V4_FLASH_W8A8_DSPARK_8P_OTHER_ARGS = [
     "--speculative-draft-attention-backend",
     "ascend",
     "--speculative-num-draft-tokens",
-    6,
+    7,
     "--speculative-dspark-block-size",
-    5,
-    "--skip-server-warmup",
+    6,
     "--cuda-graph-bs-decode",
     1,
     2,
     4,
-    5,
     6,
+    8,
+    10,
+    "--disable-radix-cache",
 ]
 
 
@@ -136,6 +147,31 @@ class TestNPUDeepSeekV4FlashW8A88PGPQA(TestNpuAccuracyTestCaseBase):
     timeout = 6000
     seed = 1
 
+    @classmethod
+    def setUpClass(cls):
+        """Launch server via `python3 -m sglang.launch_server` instead of `sglang serve`."""
+        cls._setup_per_case_output()
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        env = os.environ.copy()
+        if cls.envs:
+            env.update(cls.envs)
+
+        _, host, port = cls.base_url.split(":")
+        command = [
+            "python3",
+            "-m",
+            "sglang.launch_server",
+            "--model-path",
+            cls.model,
+            *[str(x) for x in cls.other_args],
+            "--host",
+            host[2:],
+            "--port",
+            port,
+        ]
+        cls.process = subprocess.Popen(command, env=env)
+        wait_for_server(cls.base_url, timeout=cls.server_timeout, process=cls.process)
+
     def test_npu_deepseek_v4_flash_w8a8_8p_gpqa(self):
         """Run NPU accuracy test for DeepSeek-V4-Flash W8A8 8p DSPARK GPQA."""
         self.run_accuracy()

test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py

@@ -38,8 +38,6 @@ class TestAscendMhaHicache(CustomTestCase):
             "--attention-backend",
             "ascend",
             "--enable-hierarchical-cache",
-            "--hicache-ratio",
-            1.2,
         ]
 
     def test_a_gsm8k(self):

test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py

@@ -40,8 +40,6 @@ class TestAscendMlaHicache(CustomTestCase):
             "--tp-size",
             4,
             "--enable-hierarchical-cache",
-            "--hicache-ratio",
-            1.2,
         ]
 
     def test_a_gsm8k(self):

test/registered/npu/performance/deepseek_v4_flash/test_npu_deepseek_v4_flash_w8a8_8p_in8k_out1k_50ms.py

@@ -1,3 +1,5 @@
+import os
+import subprocess
 import unittest
 
 from sglang.test.ascend.e2e.test_npu_performance_utils import (
@@ -7,6 +9,8 @@ from sglang.test.ascend.e2e.test_npu_performance_utils import (
     TestNpuPerformanceTestCaseBase,
 )
 from sglang.test.ci.ci_register import register_npu_ci
+from sglang.test.test_utils import DEFAULT_URL_FOR_TEST
+from sglang.utils import wait_for_server
 
 register_npu_ci(est_time=1800, suite="nightly-perf-16-npu-a3", nightly=True)
 register_npu_ci(est_time=1800, suite="nightly-perf-16-npu-a3-cann910", nightly=True)
@@ -23,9 +27,11 @@ DEEPSEEK_V4_FLASH_W8A8_8P_ENVS = {
     "SGLANG_NPU_USE_MULTI_STREAM": "1",
     # deepep
     "DEEP_NORMAL_MODE_USE_INT8_QUANT": "1",
-    "DEEPEP_HCCL_BUFFSIZE": "2048",
-    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "35",
+    "DEEPEP_HCCL_BUFFSIZE": "2500",
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "96",
     "DEEPEP_HYBRID_DEPLOYMENT": "1",
+    "SGLANG_RAGGED_VERIFY_MODE": "static",
+    "SGLANG_DSPARK_FAST_KERNEL": "0",
     # war barrier
     "SGLANG_ENABLE_WAR_BARRIER": "1",
     "SGLANG_FORCE_COARSE_WAR_BARRIER": "1",
@@ -55,7 +61,7 @@ DEEPSEEK_V4_FLASH_W8A8_8P_OTHER_ARGS = [
     "--device",
     "npu",
     "--prefill-max-requests",
-    160,
+    192,
     "--max-prefill-tokens",
     80000,
     "--attention-backend",
@@ -67,7 +73,7 @@ DEEPSEEK_V4_FLASH_W8A8_8P_OTHER_ARGS = [
     "--chunked-prefill-size",
     131072,
     "--max-running-requests",
-    160,
+    192,
     "--dp-size",
     16,
     "--enable-dp-attention",
@@ -80,24 +86,28 @@ DEEPSEEK_V4_FLASH_W8A8_8P_OTHER_ARGS = [
     "--enable-dp-lm-head",
     "--kv-cache-dtype",
     "bfloat16",
-    "--skip-server-warmup",
+    "--load-balance-method",
+    "round_robin",
     "--cuda-graph-bs-decode",
     1,
     2,
     4,
+    6,
     8,
     10,
-    # MTP (EAGLE) configuration.
+    # DSPARK configuration.
     "--speculative-algorithm",
-    "EAGLE",
-    "--speculative-num-steps",
-    2,
-    "--speculative-eagle-topk",
-    1,
+    "DSPARK",
+    "--speculative-draft-model-path",
+    DEEPSEEK_V4_FLASH_0731_W8A8_MODEL_PATH,
+    "--speculative-draft-model-quantization",
+    "modelslim",
+    "--speculative-draft-attention-backend",
+    "ascend",
     "--speculative-num-draft-tokens",
-    3,
-    "--ep-size",
-    16,
+    7,
+    "--speculative-dspark-block-size",
+    6,
     "--disable-radix-cache",
 ]
 
@@ -111,18 +121,42 @@ class TestNPUDeepSeekV4FlashW8A88PIn8kOut1k50ms(TestNpuPerformanceTestCaseBase):
     other_args = DEEPSEEK_V4_FLASH_W8A8_8P_OTHER_ARGS
     envs = DEEPSEEK_V4_FLASH_W8A8_8P_ENVS
     dataset_name = "random"
-    dataset_path = "/root/.cache/modelscope/hub/datasets/gsm8k_deepseekv4/cache0_8000/formal_run1_160_8000_cache0.json"
     input_len = 8000
     output_len = 1000
     num_prompts = 160
     max_concurrency = 160
     random_range_ratio = 1
-    warmup_requests = 16
+    warmup_requests = 32
     request_rate = float("inf")
-    seed = 1
     tpot = 50
     max_attempts = 3
-    output_token_throughput = 2825
+    output_token_throughput = 3100
+    accept_rate = 0.5
+
+    @classmethod
+    def setUpClass(cls):
+        """Launch server via `python3 -m sglang.launch_server` instead of `sglang serve`."""
+        cls._setup_per_case_output()
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        env = os.environ.copy()
+        if cls.envs:
+            env.update(cls.envs)
+
+        _, host, port = cls.base_url.split(":")
+        command = [
+            "python3",
+            "-m",
+            "sglang.launch_server",
+            "--model-path",
+            cls.model,
+            *[str(x) for x in cls.other_args],
+            "--host",
+            host[2:],
+            "--port",
+            port,
+        ]
+        cls.process = subprocess.Popen(command, env=env)
+        wait_for_server(cls.base_url, timeout=cls.timeout, process=cls.process)
 
     def test_npu_deepseek_v4_flash_w8a8_8p_in8k_out1k_50ms(self):
         """Run NPU performance test for DeepSeek-V4-Flash W8A8 8p in8k out1k."""

test/registered/observability/test_encoder_server_metrics.py

@@ -2,6 +2,7 @@
 
 import unittest
 import uuid
+from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, List
 from urllib.parse import urlparse
 
@@ -61,26 +62,39 @@ class TestEncoderServerMetrics(CustomTestCase):
             self.assertEqual(health.status_code, 200)
 
             req_id = f"metrics-probe-{uuid.uuid4().hex}"
-            requests.post(
-                f"{DEFAULT_URL_FOR_TEST}/scheduler_receive_url",
-                json={
-                    "req_id": req_id,
-                    "receive_url": f"{base_host}:{recv_port}",
-                    "receive_count": 1,
-                },
-            )
-            response = requests.post(
-                f"{DEFAULT_URL_FOR_TEST}/encode",
-                json={
-                    "req_id": req_id,
-                    "modality": "IMAGE",
-                    "mm_items": [f"data:image/png;base64,{MINIMUM_PNG_PICTURE_BASE64}"],
-                    "num_parts": 1,
-                    "part_idx": 0,
-                    "embedding_port": None,
-                },
-                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            )
+            # A scheduler registers concurrently with /encode, never before
+            # it: the request state only exists once /encode dispatches.
+            with ThreadPoolExecutor(max_workers=1) as pool:
+                registration = pool.submit(
+                    requests.post,
+                    f"{DEFAULT_URL_FOR_TEST}/scheduler_receive_url",
+                    json={
+                        "req_id": req_id,
+                        "receive_url": f"{base_host}:{recv_port}",
+                        "receive_count": 1,
+                    },
+                    timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                )
+                response = requests.post(
+                    f"{DEFAULT_URL_FOR_TEST}/encode",
+                    json={
+                        "req_id": req_id,
+                        "modality": "IMAGE",
+                        "mm_items": [
+                            f"data:image/png;base64,{MINIMUM_PNG_PICTURE_BASE64}"
+                        ],
+                        "num_parts": 1,
+                        "part_idx": 0,
+                        "embedding_port": None,
+                    },
+                    timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                )
+                registration_response = registration.result(
+                    timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+                )
+            # A 200 from /encode alone does not prove the embedding reached
+            # anyone; the registration is the other half of that contract.
+            self.assertEqual(registration_response.status_code, 200)
             self.assertEqual(response.status_code, 200)
 
             metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics")

test/registered/ops/test_aiter_allreduce_fusion_amd.py

@@ -14,7 +14,7 @@ import torch
 from sglang.srt.layers import communicator as comm
 from sglang.srt.layers.communicator import LayerCommunicator, ScatterMode
 from sglang.test.ci.ci_register import register_amd_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, publish_build_topology
 
 register_amd_ci(est_time=240, suite="stage-c-test-large-8-gpu-amd")
 
@@ -64,7 +64,8 @@ def _run_residual_accuracy_check():
         distributed_init_method="env://",
         backend="nccl",
     )
-    initialize_model_parallel(tensor_model_parallel_size=world_size)
+    publish_build_topology(tp_size=world_size, world_rank=rank)
+    initialize_model_parallel()
 
     dtype = torch.bfloat16
     eps = 1e-6

test/registered/ops/test_aiter_greedy_sample_amd.py

@@ -37,10 +37,7 @@ def _mock_global_server_args(backend="pytorch"):
     class _DummyTPGroup:
         device_group = None
 
-    # `Sampler.__init__` asks the context for the group; state one for the rest
-    # of the process, since this process has no distributed init. Not the scoped
-    # `override()`: its context manager would be collected here and take the
-    # value back down with it.
+    # Provide a TP group for sampler initialization without distributed setup.
     get_parallel().override_permanently(tp_group=_DummyTPGroup())
     from sglang.srt.runtime_context import get_flags
 

test/registered/perf/test_bench_one_batch_2gpu.py

@@ -1,53 +0,0 @@
-import unittest
-
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.test_utils import (
-    DEFAULT_MODEL_NAME_FOR_TEST,
-    DEFAULT_MOE_MODEL_NAME_FOR_TEST,
-    CustomTestCase,
-    is_in_amd_ci,
-    is_in_ci,
-    run_bench_offline_throughput,
-    write_github_step_summary,
-)
-
-register_cuda_ci(est_time=162, stage="extra-a", runner_config="2-gpu-large")
-register_amd_ci(est_time=630, suite="stage-b-test-2-gpu-large-amd")
-
-
-class TestBenchOneBatch2GPU(CustomTestCase):
-    def test_moe_tp2_bs1(self):
-        output_throughput = run_bench_offline_throughput(
-            DEFAULT_MOE_MODEL_NAME_FOR_TEST,
-            ["--tp", "2", "--cuda-graph-max-bs-decode", "2"],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_moe_tp2_bs1 (Mixtral-8x7B)\n"
-                f"output_throughput: {output_throughput:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(output_throughput, 85)
-            else:
-                self.assertGreater(output_throughput, 125)
-
-    def test_torch_compile_tp2_bs1(self):
-        output_throughput = run_bench_offline_throughput(
-            DEFAULT_MODEL_NAME_FOR_TEST,
-            ["--tp", "2", "--enable-torch-compile", "--cuda-graph-max-bs-decode", "2"],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_torch_compile_tp2_bs1 (Mixtral-8x7B)\n"
-                f"output_throughput: {output_throughput:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(output_throughput, 200)
-            else:
-                self.assertGreater(output_throughput, 220)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/perf/test_bench_serving_1gpu_large.py

@@ -1,84 +0,0 @@
-"""
-Performance tests for single GPU that need H200 (80GB) - FP8 and EAGLE tests.
-"""
-
-import unittest
-
-from sglang.srt.utils import is_hip
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.test_utils import (
-    DEFAULT_DRAFT_MODEL_EAGLE,
-    DEFAULT_MODEL_NAME_FOR_TEST_FP8,
-    DEFAULT_TARGET_MODEL_EAGLE,
-    CustomTestCase,
-    is_in_amd_ci,
-    is_in_ci,
-    run_bench_serving,
-    write_github_step_summary,
-)
-
-register_cuda_ci(est_time=275, stage="extra-a", runner_config="1-gpu-large")
-register_amd_ci(est_time=300, suite="stage-b-test-1-gpu-large-amd")
-
-
-class TestBenchServing1GPULarge(CustomTestCase):
-    def test_offline_throughput_default_fp8(self):
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST_FP8,
-            num_prompts=500,
-            request_rate=float("inf"),
-            other_server_args=[],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_offline_throughput_default_fp8\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 3500)
-            else:
-                self.assertGreater(res["output_throughput"], 4300)
-
-    @unittest.skipIf(is_hip(), "Skip Eagle test for ROCm")
-    def test_online_latency_eagle(self):
-        res = run_bench_serving(
-            model=DEFAULT_TARGET_MODEL_EAGLE,
-            num_prompts=300,
-            request_rate=8,
-            sharegpt_context_len=3072,
-            disable_ignore_eos=True,
-            dataset_name="sharegpt",
-            other_server_args=[
-                "--speculative-algorithm",
-                "EAGLE",
-                "--speculative-draft-model-path",
-                DEFAULT_DRAFT_MODEL_EAGLE,
-                "--speculative-num-steps",
-                "5",
-                "--speculative-eagle-topk",
-                "4",
-                "--speculative-num-draft-tokens",
-                "16",
-                "--mem-fraction-static",
-                "0.7",
-            ],
-            need_warmup=True,
-            seed=42,
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_online_latency_eagle\n"
-                f"median_e2e_latency_ms: {res['median_e2e_latency_ms']:.2f} ms\n"
-                f"accept_length: {res['accept_length']:.2f} \n"
-            )
-            if is_in_amd_ci():
-                self.assertLess(res["median_e2e_latency_ms"], 1800)
-            else:
-                self.assertLess(res["median_e2e_latency_ms"], 900)
-            self.assertGreater(res["accept_length"], 3.0)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/perf/test_bench_serving_1gpu_part1.py

@@ -1,267 +0,0 @@
-"""
-Performance tests for single GPU - LLM throughput/latency and LoRA tests.
-Works on 5090 (32GB).
-"""
-
-import asyncio
-import itertools
-import unittest
-
-import requests
-
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.test_utils import (
-    DEFAULT_MODEL_NAME_FOR_TEST,
-    CustomTestCase,
-    is_in_amd_ci,
-    is_in_ci,
-    run_bench_serving,
-    write_github_step_summary,
-)
-
-register_cuda_ci(est_time=1264, stage="extra-a", runner_config="1-gpu-large")
-register_amd_ci(est_time=1100, suite="stage-b-test-1-gpu-large-amd")
-
-
-class TestBenchServing1GPUPart1(CustomTestCase):
-    def test_offline_throughput_default(self):
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST,
-            num_prompts=500,
-            request_rate=float("inf"),
-            other_server_args=[],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_offline_throughput_default\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 3050)
-            else:
-                self.assertGreater(res["output_throughput"], 3800)
-
-    def test_offline_throughput_non_stream_small_batch_size(self):
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST,
-            num_prompts=200,
-            request_rate=float("inf"),
-            other_server_args=["--max-running-requests", "10"],
-            dataset_name="sharegpt",
-            random_input_len=None,
-            random_output_len=None,
-            disable_stream=True,
-            need_warmup=True,
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_offline_throughput_non_stream_small_batch_size\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 1000)
-            else:
-                self.assertGreater(res["output_throughput"], 1050)
-
-    def test_offline_throughput_without_radix_cache(self):
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST,
-            num_prompts=500,
-            request_rate=float("inf"),
-            other_server_args=["--disable-radix-cache"],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_offline_throughput_without_radix_cache\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 3050)
-            else:
-                self.assertGreater(res["output_throughput"], 3800)
-
-    def test_offline_throughput_without_chunked_prefill(self):
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST,
-            num_prompts=500,
-            request_rate=float("inf"),
-            other_server_args=["--chunked-prefill-size", "-1"],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_offline_throughput_without_chunked_prefill\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            self.assertGreater(res["output_throughput"], 2600)
-
-    def test_offline_throughput_with_triton_attention_backend(self):
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST,
-            num_prompts=500,
-            request_rate=float("inf"),
-            other_server_args=[
-                "--attention-backend",
-                "triton",
-                "--context-length",
-                "8192",
-            ],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_offline_throughput_with_triton_attention_backend\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 2700)
-            else:
-                self.assertGreater(res["output_throughput"], 3700)
-
-    def test_online_latency_default(self):
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST,
-            num_prompts=100,
-            request_rate=1,
-            other_server_args=[],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_online_latency_default\n"
-                f"median_e2e_latency_ms: {res['median_e2e_latency_ms']:.2f} ms\n"
-            )
-            self.assertLess(res["median_e2e_latency_ms"], 11000)
-            if is_in_amd_ci():
-                self.assertLess(res["median_ttft_ms"], 115)
-            else:
-                self.assertLess(res["median_ttft_ms"], 86)
-            self.assertLess(res["median_itl_ms"], 10)
-
-    def test_online_lora_latency(self):
-        res = self._run_lora_latency_test(enable_background_task=False)
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_online_lora_latency\n"
-                f"median_e2e_latency_ms: {res['median_e2e_latency_ms']:.2f} ms\n"
-                f"median_ttft_ms: {res['median_ttft_ms']:.2f} ms\n"
-            )
-            if is_in_amd_ci():
-                self.assertLess(res["median_e2e_latency_ms"], 3320)
-            else:
-                self.assertLess(res["median_e2e_latency_ms"], 2400)
-            # relax for mi300x (LoRA TTFT ~2x slower than mi325)
-            if is_in_amd_ci():
-                self.assertLess(res["median_ttft_ms"], 100)
-            else:
-                self.assertLess(res["median_ttft_ms"], 58)
-
-    def test_online_lora_latency_with_concurrent_adapter_updates(self):
-        res = self._run_lora_latency_test(enable_background_task=True)
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_online_lora_latency_with_concurrent_adapter_updates\n"
-                f"median_e2e_latency_ms: {res['median_e2e_latency_ms']:.2f} ms\n"
-                f"median_ttft_ms: {res['median_ttft_ms']:.2f} ms\n"
-            )
-            if is_in_amd_ci():
-                self.assertLess(res["median_e2e_latency_ms"], 6000)
-            else:
-                self.assertLess(res["median_e2e_latency_ms"], 4000)
-            # relax for mi300x (LoRA TTFT ~2x slower than mi325)
-            if is_in_amd_ci():
-                self.assertLess(res["median_ttft_ms"], 130)
-            else:
-                self.assertLess(res["median_ttft_ms"], 80)
-
-    def _run_lora_latency_test(self, enable_background_task: bool):
-        """
-        Run a latency test for LoRA with the specified background task setting.
-        """
-
-        async def lora_loader_unloader_task(
-            base_url: str,
-            start_event: asyncio.Event,
-            stop_event: asyncio.Event,
-        ):
-            """
-            A background task that repeatedly loads and unloads a LoRA adapter.
-            """
-            await start_event.wait()
-
-            path_cycler = itertools.cycle(
-                [
-                    "pbevan11/llama-3.1-8b-ocr-correction",
-                    "faridlazuarda/valadapt-llama-3.1-8B-it-chinese",
-                    "philschmid/code-llama-3-1-8b-text-to-sql-lora",
-                ]
-            )
-            load_url = f"{base_url}/load_lora_adapter"
-            unload_url = f"{base_url}/unload_lora_adapter"
-            num_updates = 0
-
-            while not stop_event.is_set():
-                lora_path = next(path_cycler)
-                response = await asyncio.to_thread(
-                    requests.post,
-                    load_url,
-                    json={"lora_name": lora_path, "lora_path": lora_path},
-                )
-                self.assertTrue(
-                    response.ok, f"Failed to load LoRA adapter: {response.text}"
-                )
-                num_updates += 1
-
-                if stop_event.is_set():
-                    break
-
-                await asyncio.sleep(1)
-
-                response = await asyncio.to_thread(
-                    requests.post,
-                    unload_url,
-                    json={"lora_name": lora_path},
-                )
-                self.assertTrue(
-                    response.ok, f"Failed to unload LoRA adapter: {response.text}"
-                )
-                num_updates += 1
-
-                await asyncio.sleep(1)
-
-        background_task = lora_loader_unloader_task if enable_background_task else None
-        res = run_bench_serving(
-            model=DEFAULT_MODEL_NAME_FOR_TEST,
-            num_prompts=400,
-            request_rate=8,
-            other_server_args=[
-                "--enable-lora",
-                "--max-loras-per-batch",
-                "1",
-                "--disable-radix-cache",
-                "--random-seed",
-                "42",
-                "--mem-fraction-static",
-                "0.8",
-                "--lora-paths",
-                "nvidia/llama-3.1-nemoguard-8b-topic-control",
-                "--max-lora-rank",
-                "256",
-            ],
-            dataset_name="random",
-            random_input_len=256,
-            random_output_len=256,
-            lora_name=["nvidia/llama-3.1-nemoguard-8b-topic-control"],
-            background_task=background_task,
-        )
-
-        return res
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/perf/test_bench_serving_1gpu_part2.py

@@ -1,211 +0,0 @@
-"""
-Performance tests for single GPU - VLM, Score API, and Embeddings API tests.
-Works on 5090 (32GB).
-"""
-
-import unittest
-
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.test_utils import (
-    DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
-    DEFAULT_SMALL_MODEL_NAME_FOR_TEST_SCORE,
-    DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
-    CustomTestCase,
-    is_in_amd_ci,
-    is_in_ci,
-    run_bench_serving,
-    run_embeddings_benchmark,
-    run_score_benchmark,
-    write_github_step_summary,
-)
-
-register_cuda_ci(est_time=909, stage="extra-a", runner_config="1-gpu-large")
-register_amd_ci(est_time=900, suite="stage-b-test-1-gpu-large-amd")
-
-
-class TestBenchServing1GPUPart2(CustomTestCase):
-    @unittest.skip(
-        "Qwen2.5-VL server crashes with SIGBUS (exit code -7) on main; disable until fixed"
-    )
-    def test_vlm_offline_throughput(self):
-        res = run_bench_serving(
-            model=DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
-            num_prompts=200,
-            request_rate=float("inf"),
-            other_server_args=[
-                "--mem-fraction-static",
-                "0.7",
-            ],
-            dataset_name="mmmu",
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_vlm_offline_throughput\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            # relax for mi300x
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 900)
-            else:
-                self.assertGreater(res["output_throughput"], 2500)
-
-    def test_vlm_online_latency(self):
-        res = run_bench_serving(
-            model=DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
-            num_prompts=250,
-            request_rate=1,
-            other_server_args=[
-                "--mem-fraction-static",
-                "0.7",
-            ],
-            dataset_name="mmmu",
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_vlm_online_latency\n"
-                f"median_e2e_latency_ms: {res['median_e2e_latency_ms']:.2f} ms\n"
-            )
-            self.assertLess(res["median_e2e_latency_ms"], 16500)
-            if is_in_amd_ci():
-                self.assertLess(res["median_ttft_ms"], 150)
-            else:
-                self.assertLess(res["median_ttft_ms"], 100)
-            self.assertLess(res["median_itl_ms"], 8)
-
-    def test_score_api_latency_throughput(self):
-        """Test score API latency and throughput performance"""
-        res = run_score_benchmark(
-            model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST_SCORE,
-            num_requests=1000,
-            batch_size=10,
-            other_server_args=[],
-            need_warmup=True,
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_score_api_throughput\n"
-                f"Average latency: {res['avg_latency_ms']:.2f} ms\n"
-                f"P95 latency: {res['p95_latency_ms']:.2f} ms\n"
-                f"Score API throughput: {res['throughput']:.2f} req/s\n"
-                f"Successful requests: {res['successful_requests']}/{res['total_requests']}\n"
-            )
-
-        self.assertEqual(res["successful_requests"], res["total_requests"])
-        # relax for mi300x
-        if is_in_amd_ci():
-            self.assertLess(res["avg_latency_ms"], 60)
-            self.assertLess(res["p95_latency_ms"], 65)
-            self.assertGreater(res["throughput"], 16)
-        else:
-            self.assertLess(res["avg_latency_ms"], 48)
-            self.assertLess(res["p95_latency_ms"], 50)
-            self.assertGreater(res["throughput"], 20)
-
-    def test_score_api_batch_scaling(self):
-        """Test score API performance with different batch sizes"""
-        batch_sizes = [10, 25, 50]
-
-        for batch_size in batch_sizes:
-            res = run_score_benchmark(
-                model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST_SCORE,
-                num_requests=500,
-                batch_size=batch_size,
-            )
-
-            if is_in_ci():
-                write_github_step_summary(
-                    f"### test_score_api_batch_scaling_size_{batch_size}\n"
-                    f"Batch size: {batch_size}\n"
-                    f"Average latency: {res['avg_latency_ms']:.2f} ms\n"
-                    f"P95 latency: {res['p95_latency_ms']:.2f} ms\n"
-                    f"Throughput: {res['throughput']:.2f} req/s\n"
-                    f"Successful requests: {res['successful_requests']}/{res['total_requests']}\n"
-                )
-
-            self.assertEqual(res["successful_requests"], res["total_requests"])
-            # relax for mi300x
-            if is_in_amd_ci():
-                bounds = {10: (60, 65), 25: (70, 80), 50: (80, 90)}
-                default_bounds = (90, 90)
-            else:
-                bounds = {10: (45, 50), 25: (50, 60), 50: (60, 65)}
-                default_bounds = (60, 65)
-            avg_latency_bound, p95_latency_bound = bounds.get(
-                batch_size, default_bounds
-            )
-            self.assertLess(res["avg_latency_ms"], avg_latency_bound)
-            self.assertLess(res["p95_latency_ms"], p95_latency_bound)
-
-    def test_embeddings_api_latency_throughput(self):
-        """Test embeddings API latency and throughput performance"""
-        res = run_embeddings_benchmark(
-            model=DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
-            num_requests=1000,
-            batch_size=1,
-            input_tokens=500,
-            other_server_args=[],
-            need_warmup=True,
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_embeddings_api_throughput\n"
-                f"Average latency: {res['avg_latency_ms']:.2f} ms\n"
-                f"P95 latency: {res['p95_latency_ms']:.2f} ms\n"
-                f"Embeddings API throughput: {res['throughput']:.2f} req/s\n"
-                f"Successful requests: {res['successful_requests']}/{res['total_requests']}\n"
-            )
-
-        self.assertEqual(res["successful_requests"], res["total_requests"])
-        # relax for mi300x
-        if is_in_amd_ci():
-            self.assertLess(res["avg_latency_ms"], 35)
-            self.assertLess(res["p95_latency_ms"], 40)
-            self.assertGreater(res["throughput"], 30)
-        else:
-            self.assertLess(res["avg_latency_ms"], 20)
-            self.assertLess(res["p95_latency_ms"], 25)
-            self.assertGreater(res["throughput"], 60)
-
-    def test_embeddings_api_batch_scaling(self):
-        """Test embeddings API performance with different batch sizes"""
-        batch_sizes = [10, 25, 50]
-
-        for batch_size in batch_sizes:
-            res = run_embeddings_benchmark(
-                model=DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST,
-                num_requests=500,
-                batch_size=batch_size,
-                input_tokens=500,
-            )
-
-            if is_in_ci():
-                write_github_step_summary(
-                    f"### test_embeddings_api_batch_scaling_size_{batch_size}\n"
-                    f"Batch size: {batch_size}\n"
-                    f"Average latency: {res['avg_latency_ms']:.2f} ms\n"
-                    f"P95 latency: {res['p95_latency_ms']:.2f} ms\n"
-                    f"Throughput: {res['throughput']:.2f} req/s\n"
-                    f"Successful requests: {res['successful_requests']}/{res['total_requests']}\n"
-                )
-
-            self.assertEqual(res["successful_requests"], res["total_requests"])
-            # relax for mi300x
-            if is_in_amd_ci():
-                bounds = {10: (80, 90), 25: (140, 150), 50: (230, 240)}
-                default_bounds = (300, 300)
-            else:
-                bounds = {10: (60, 65), 25: (115, 120), 50: (190, 195)}
-                default_bounds = (250, 250)
-            avg_latency_bound, p95_latency_bound = bounds.get(
-                batch_size, default_bounds
-            )
-            self.assertLess(res["avg_latency_ms"], avg_latency_bound)
-            self.assertLess(res["p95_latency_ms"], p95_latency_bound)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/perf/test_bench_serving_2gpu.py

@@ -1,108 +0,0 @@
-"""
-Performance tests for 2-GPU that need large GPUs (H200 80GB) - MoE and Pipeline Parallel tests.
-"""
-
-import unittest
-
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.test_utils import (
-    DEFAULT_MOE_MODEL_NAME_FOR_TEST,
-    CustomTestCase,
-    is_in_amd_ci,
-    is_in_ci,
-    run_bench_serving,
-    write_github_step_summary,
-)
-
-register_cuda_ci(est_time=687, stage="extra-a", runner_config="2-gpu-large")
-register_amd_ci(est_time=1450, suite="stage-b-test-2-gpu-large-amd")
-
-
-class TestBenchServing2GPU(CustomTestCase):
-    def test_moe_offline_throughput_default(self):
-        res = run_bench_serving(
-            model=DEFAULT_MOE_MODEL_NAME_FOR_TEST,
-            num_prompts=300,
-            request_rate=float("inf"),
-            other_server_args=["--tp", "2"],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_moe_offline_throughput_default\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 2100)
-            else:
-                self.assertGreater(res["output_throughput"], 2200)
-
-    def test_moe_offline_throughput_without_radix_cache(self):
-        res = run_bench_serving(
-            model=DEFAULT_MOE_MODEL_NAME_FOR_TEST,
-            num_prompts=300,
-            request_rate=float("inf"),
-            other_server_args=["--tp", "2", "--disable-radix-cache"],
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_moe_offline_throughput_without_radix_cache\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["output_throughput"], 2100)
-            else:
-                self.assertGreater(res["output_throughput"], 2200)
-
-    def test_pp_offline_throughput_default_decode(self):
-        res = run_bench_serving(
-            model=DEFAULT_MOE_MODEL_NAME_FOR_TEST,
-            num_prompts=1000,
-            request_rate=float("inf"),
-            random_input_len=1,
-            random_output_len=1024,
-            other_server_args=["--pp-size", "2"],
-            need_warmup=True,
-            seed=42,
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_pp_offline_throughput_default_decode\n"
-                f"Output throughput: {res['output_throughput']:.2f} token/s\n"
-            )
-            self.assertGreater(res["output_throughput"], 6700)
-
-    def test_pp_long_context_prefill(self):
-        res = run_bench_serving(
-            model="meta-llama/Llama-3.3-70B-Instruct",
-            num_prompts=4,
-            request_rate=float("inf"),
-            random_input_len=128000,
-            random_output_len=1,
-            dataset_name="random",
-            other_server_args=[
-                "--quantization",
-                "fp8",
-                "--pp-size",
-                "2",
-            ]
-            + (["--mem-fraction-static", "0.7"] if is_in_amd_ci() else []),
-            need_warmup=False,
-            seed=42,
-        )
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_pp_long_context_latency_prefill\n"
-                f"input_throughput: {res['input_throughput']:.2f} ms\n"
-            )
-            if is_in_amd_ci():
-                self.assertGreater(res["input_throughput"], 3000)
-            else:
-                self.assertGreater(res["input_throughput"], 4000)
-
-
-if __name__ == "__main__":
-    unittest.main()