qwen 3.8 rebase (#35758)

Co-authored-by: cherichy <cherichy@outlook.com>
Co-authored-by: guangyunh-nv <guangyunh@nvidia.com>
Co-authored-by: jiahanc <jiahanc@nvidia.com>
Co-authored-by: jinyangyuan-nvidia <joyuan@nvidia.com>
Co-authored-by: Cheng Hang <chang@nvidia.com>
Co-authored-by: Yicheng Qiang <yqiang@nvidia.com>
Co-authored-by: Sam Li <lsam@nvidia.com>
Co-authored-by: Tom-Zheng <tizheng@nvidia.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: xiaoweiw-nv <xiaoweiw@nvidia.com>
Co-authored-by: Zheng Li <lizheng.cs@zju.edu.cn>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Zijie Xia <zijie.xia@radixark.ai>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

test/registered/attention/unittests/mamba/test_replay_state_indices_validator.py

@@ -0,0 +1,64 @@
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.mamba.replay_state_indices_validator import (
+    validate_replay_state_indices_cpu,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestReplayStateIndicesValidator(unittest.TestCase):
+    def test_valid_unique_live_slots_and_padding(self):
+        validate_replay_state_indices_cpu(
+            torch.tensor([0, 2, 9, -1, -1], dtype=torch.int32),
+            valid_bs=3,
+            total_bs=5,
+            num_state_slots=10,
+        )
+
+    def test_rejects_duplicate_live_slot(self):
+        with self.assertRaisesRegex(AssertionError, r"duplicate_slots=\[7\]"):
+            validate_replay_state_indices_cpu(
+                torch.tensor([7, 2, 7, -1], dtype=torch.int32),
+                valid_bs=3,
+                total_bs=4,
+                num_state_slots=10,
+            )
+
+    def test_rejects_out_of_range_live_slots(self):
+        for bad_slot in (-1, -2, 10):
+            with self.subTest(bad_slot=bad_slot):
+                with self.assertRaisesRegex(AssertionError, "live rows"):
+                    validate_replay_state_indices_cpu(
+                        torch.tensor([3, bad_slot, -1], dtype=torch.int64),
+                        valid_bs=2,
+                        total_bs=3,
+                        num_state_slots=10,
+                    )
+
+    def test_rejects_non_sentinel_padding(self):
+        with self.assertRaisesRegex(AssertionError, "padded rows"):
+            validate_replay_state_indices_cpu(
+                torch.tensor([3, 5, 5], dtype=torch.int32),
+                valid_bs=2,
+                total_bs=3,
+                num_state_slots=10,
+            )
+
+    def test_requires_cpu_tensor(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA is unavailable")
+        with self.assertRaisesRegex(ValueError, "copied to CPU"):
+            validate_replay_state_indices_cpu(
+                torch.tensor([1], dtype=torch.int32, device="cuda"),
+                valid_bs=1,
+                total_bs=1,
+                num_state_slots=2,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,392 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
+    can_use_fused_qkvzba_causal_conv1d_update_contiguous,
+    fused_qkvzba_causal_conv1d_update_contiguous,
+    fused_qkvzba_split_reshape_cat_contiguous,
+)
+
+# This is also the update implementation imported directly by GDNBackend on
+# CUDA; the presence of the optional sgl_kernel AOT extension does not reroute
+# GDN decode through srt.layers.attention.mamba.causal_conv1d.
+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=8, stage="base-b", runner_config="1-gpu-large")
+
+
+def _reference(
+    qkvz,
+    ba,
+    state,
+    weight,
+    bias,
+    indices,
+    *,
+    qkv_dim,
+    v_dim,
+    num_v_heads,
+    head_v_dim,
+    activation,
+):
+    qkv = qkvz[:, :qkv_dim]
+    out = torch.empty_like(qkv)
+    state_out = state.clone()
+    width = weight.shape[1]
+    for row, slot_tensor in enumerate(indices.cpu()):
+        slot = int(slot_tensor)
+        if slot < 0 or slot >= state.shape[0]:
+            out[row].copy_(qkv[row])
+            continue
+        # The deployed direct-Triton decode wrapper uses an effective
+        # state_len=width-1 even if the physical cache envelope is wider.
+        history = state[slot, :, : width - 1].float()
+        values = torch.cat((history, qkv[row, :, None].float()), dim=-1)
+        acc = (values * weight.float()).sum(dim=-1)
+        if bias is not None:
+            acc = acc + bias.float()
+        if activation in ("silu", "swish"):
+            acc = torch.nn.functional.silu(acc)
+        out[row].copy_(acc.to(qkv.dtype))
+        if width > 2:
+            state_out[slot, :, : width - 2].copy_(state[slot, :, 1 : width - 1])
+        state_out[slot, :, width - 2].copy_(qkv[row])
+
+    z = qkvz[:, qkv_dim:].reshape(-1, num_v_heads, head_v_dim).contiguous()
+    b, a = ba.split([num_v_heads, num_v_heads], dim=-1)
+    return out, z, b.contiguous(), a.contiguous(), state_out
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
+class TestGDNDecodeFusedProjectionConv1D(unittest.TestCase):
+    def test_contiguous_unpack_ratio8_microbenchmark_baseline(self):
+        batch = 9
+        num_qk_heads = 1
+        num_v_heads = 8
+        head_dim = 128
+        qkv_dim = (2 * num_qk_heads + num_v_heads) * head_dim
+        v_dim = num_v_heads * head_dim
+        qkvz = torch.randn(
+            batch,
+            qkv_dim + v_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        ba = torch.randn(
+            batch,
+            2 * num_v_heads,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous(
+            qkvz,
+            ba,
+            num_qk_heads,
+            num_v_heads,
+            head_dim,
+            head_dim,
+        )
+        torch.testing.assert_close(mixed_qkv, qkvz[:, :qkv_dim])
+        torch.testing.assert_close(
+            z, qkvz[:, qkv_dim:].reshape(batch, num_v_heads, head_dim)
+        )
+        torch.testing.assert_close(b, ba[:, :num_v_heads])
+        torch.testing.assert_close(a, ba[:, num_v_heads:])
+
+    def _run_case(
+        self,
+        *,
+        batch,
+        q_dim,
+        k_dim,
+        v_dim,
+        num_v_heads,
+        head_v_dim,
+        width,
+        state_len,
+        dtype,
+        with_bias,
+        activation,
+        strided_state=False,
+        with_padding=False,
+    ):
+        torch.manual_seed(17)
+        device = "cuda"
+        qkv_dim = q_dim + k_dim + v_dim
+        qkvz = torch.randn(batch, qkv_dim + v_dim, device=device, dtype=dtype)
+        ba = torch.randn(batch, 2 * num_v_heads, device=device, dtype=dtype)
+        weight = torch.randn(qkv_dim, width, device=device, dtype=dtype) * 0.1
+        bias = (
+            torch.randn(qkv_dim, device=device, dtype=dtype) * 0.1
+            if with_bias
+            else None
+        )
+        slots = batch + 3
+        if strided_state:
+            backing = torch.randn(
+                slots,
+                state_len,
+                qkv_dim * 2,
+                device=device,
+                dtype=dtype,
+            )
+            state = backing[:, :, ::2].transpose(1, 2)
+            self.assertFalse(state.is_contiguous())
+        else:
+            state = torch.randn(slots, qkv_dim, state_len, device=device, dtype=dtype)
+        indices = torch.randperm(slots, device=device, dtype=torch.int64)[:batch]
+        if with_padding:
+            indices[-1] = -1
+        indices = indices.to(torch.int32)
+
+        ref = _reference(
+            qkvz,
+            ba,
+            state,
+            weight,
+            bias,
+            indices,
+            qkv_dim=qkv_dim,
+            v_dim=v_dim,
+            num_v_heads=num_v_heads,
+            head_v_dim=head_v_dim,
+            activation=activation,
+        )
+        state_test = state.clone(memory_format=torch.preserve_format)
+        out, z, b, a = fused_qkvzba_causal_conv1d_update_contiguous(
+            qkvz,
+            ba,
+            state_test,
+            weight,
+            bias,
+            indices,
+            qkv_dim=qkv_dim,
+            v_dim=v_dim,
+            num_v_heads=num_v_heads,
+            head_v_dim=head_v_dim,
+            activation=activation,
+        )
+        torch.cuda.synchronize()
+
+        atol = 2e-2 if dtype == torch.bfloat16 else 3e-3
+        output_max_diff = (out.float() - ref[0].float()).abs().max().item()
+        state_max_diff = (state_test.float() - ref[4].float()).abs().max().item()
+        print(
+            "case "
+            f"B={batch} QKV={qkv_dim} V={v_dim} W={width} "
+            f"dtype={dtype} output_max_diff={output_max_diff:.8g} "
+            f"state_max_diff={state_max_diff:.8g}",
+            flush=True,
+        )
+        torch.testing.assert_close(out, ref[0], rtol=0, atol=atol)
+        torch.testing.assert_close(z, ref[1], rtol=0, atol=0)
+        torch.testing.assert_close(b, ref[2], rtol=0, atol=0)
+        torch.testing.assert_close(a, ref[3], rtol=0, atol=0)
+        torch.testing.assert_close(state_test, ref[4], rtol=0, atol=0)
+        self.assertEqual(b.data_ptr() % 32, 0)
+        self.assertEqual(a.data_ptr() % 32, 0)
+
+    def test_random_shapes_widths_dtypes_and_state_updates(self):
+        cases = (
+            # Small boundary shapes.
+            dict(
+                batch=1,
+                q_dim=16,
+                k_dim=16,
+                v_dim=32,
+                num_v_heads=2,
+                head_v_dim=16,
+                width=2,
+                state_len=1,
+                dtype=torch.float16,
+                with_bias=False,
+                activation=None,
+            ),
+            dict(
+                batch=17,
+                q_dim=32,
+                k_dim=32,
+                v_dim=64,
+                num_v_heads=4,
+                head_v_dim=16,
+                width=3,
+                state_len=5,
+                dtype=torch.bfloat16,
+                with_bias=True,
+                activation="silu",
+                strided_state=True,
+                with_padding=True,
+            ),
+            # Qwen3.5-35B TP16 local GDN dimensions.
+            dict(
+                batch=32,
+                q_dim=128,
+                k_dim=128,
+                v_dim=256,
+                num_v_heads=2,
+                head_v_dim=128,
+                width=4,
+                state_len=3,
+                dtype=torch.bfloat16,
+                with_bias=False,
+                activation="silu",
+            ),
+            # Large TP-local GDN dimensions with an 8:1 value/key head ratio.
+            dict(
+                batch=8,
+                q_dim=128,
+                k_dim=128,
+                v_dim=1024,
+                num_v_heads=8,
+                head_v_dim=128,
+                width=4,
+                state_len=3,
+                dtype=torch.bfloat16,
+                with_bias=False,
+                activation="silu",
+            ),
+        )
+        for case in cases:
+            with self.subTest(case=case):
+                self._run_case(**case)
+
+    def test_cuda_graph_replay(self):
+        batch = 4
+        qkv_dim, v_dim, num_v_heads, head_v_dim = 128, 64, 2, 32
+        qkvz = torch.randn(batch, qkv_dim + v_dim, device="cuda", dtype=torch.bfloat16)
+        ba = torch.randn(batch, 2 * num_v_heads, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(qkv_dim, 4, device="cuda", dtype=torch.bfloat16)
+        state = torch.randn(batch + 1, qkv_dim, 3, device="cuda", dtype=torch.bfloat16)
+        initial_state = state.clone()
+        indices = torch.arange(batch, device="cuda", dtype=torch.int32)
+        # Compile before capture; Triton compilation itself is not graph-safe.
+        fused_qkvzba_causal_conv1d_update_contiguous(
+            qkvz,
+            ba,
+            state,
+            weight,
+            None,
+            indices,
+            qkv_dim=qkv_dim,
+            v_dim=v_dim,
+            num_v_heads=num_v_heads,
+            head_v_dim=head_v_dim,
+            activation="silu",
+        )
+        state.copy_(initial_state)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = fused_qkvzba_causal_conv1d_update_contiguous(
+                qkvz,
+                ba,
+                state,
+                weight,
+                None,
+                indices,
+                qkv_dim=qkv_dim,
+                v_dim=v_dim,
+                num_v_heads=num_v_heads,
+                head_v_dim=head_v_dim,
+                activation="silu",
+            )
+        state.copy_(initial_state)
+        graph.replay()
+        ref_state = initial_state.clone()
+        ref_qkv, ref_z, ref_b, ref_a = fused_qkvzba_split_reshape_cat_contiguous(
+            qkvz,
+            ba,
+            1,
+            num_v_heads,
+            32,
+            head_v_dim,
+        )
+        ref_qkv = causal_conv1d_update(
+            ref_qkv,
+            ref_state,
+            weight,
+            None,
+            "silu",
+            conv_state_indices=indices,
+        )
+        torch.testing.assert_close(captured[0], ref_qkv, rtol=0, atol=0)
+        torch.testing.assert_close(captured[1], ref_z, rtol=0, atol=0)
+        torch.testing.assert_close(captured[2], ref_b, rtol=0, atol=0)
+        torch.testing.assert_close(captured[3], ref_a, rtol=0, atol=0)
+        torch.testing.assert_close(state, ref_state, rtol=0, atol=0)
+
+    def test_out_of_range_state_slots_are_safely_masked(self):
+        torch.manual_seed(29)
+        batch = 4
+        qkv_dim, v_dim, num_v_heads, head_v_dim = 128, 64, 2, 32
+        qkvz = torch.randn(batch, qkv_dim + v_dim, device="cuda", dtype=torch.bfloat16)
+        ba = torch.randn(batch, 2 * num_v_heads, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(qkv_dim, 4, device="cuda", dtype=torch.bfloat16)
+        state = torch.randn(3, qkv_dim, 3, device="cuda", dtype=torch.bfloat16)
+        # -1 is the expected padding sentinel; -2 and len(state) exercise the
+        # hard lower/upper bounds.  Slot 1 remains a normal live update.
+        indices = torch.tensor([-2, -1, state.shape[0], 1], device="cuda")
+        indices = indices.to(torch.int32)
+
+        ref = _reference(
+            qkvz,
+            ba,
+            state,
+            weight,
+            None,
+            indices,
+            qkv_dim=qkv_dim,
+            v_dim=v_dim,
+            num_v_heads=num_v_heads,
+            head_v_dim=head_v_dim,
+            activation="silu",
+        )
+        state_test = state.clone()
+        out, z, b, a = fused_qkvzba_causal_conv1d_update_contiguous(
+            qkvz,
+            ba,
+            state_test,
+            weight,
+            None,
+            indices,
+            qkv_dim=qkv_dim,
+            v_dim=v_dim,
+            num_v_heads=num_v_heads,
+            head_v_dim=head_v_dim,
+            activation="silu",
+        )
+        torch.cuda.synchronize()
+
+        torch.testing.assert_close(out, ref[0], rtol=0, atol=2e-2)
+        torch.testing.assert_close(z, ref[1], rtol=0, atol=0)
+        torch.testing.assert_close(b, ref[2], rtol=0, atol=0)
+        torch.testing.assert_close(a, ref[3], rtol=0, atol=0)
+        torch.testing.assert_close(state_test, ref[4], rtol=0, atol=0)
+
+    def test_fp8_activation_is_an_explicit_fallback(self):
+        if not hasattr(torch, "float8_e4m3fn"):
+            self.skipTest("PyTorch has no FP8 dtype")
+        qkvz = torch.empty(1, 96, device="cuda", dtype=torch.float8_e4m3fn)
+        ba = torch.empty(1, 4, device="cuda", dtype=torch.bfloat16)
+        state = torch.empty(2, 64, 3, device="cuda", dtype=torch.bfloat16)
+        weight = torch.empty(64, 4, device="cuda", dtype=torch.bfloat16)
+        indices = torch.zeros(1, device="cuda", dtype=torch.int32)
+        eligible, reason = can_use_fused_qkvzba_causal_conv1d_update_contiguous(
+            qkvz,
+            ba,
+            state,
+            weight,
+            None,
+            indices,
+            qkv_dim=64,
+            v_dim=32,
+            num_v_heads=2,
+            activation="silu",
+        )
+        self.assertFalse(eligible)
+        self.assertIn("FP16, BF16, or FP32", reason)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -1,6 +1,7 @@
 import random
 import sys
 from contextlib import nullcontext
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -159,6 +160,42 @@ def test_compact_all_tokens_uses_tight_routing_independent_bound(
     )
 
 
+def test_compact_eager_keeps_masked_layout_for_cuda_graph(monkeypatch):
+    config = MoeRunnerConfig(
+        num_experts=128,
+        num_local_experts=16,
+        hidden_size=2048,
+        intermediate_size_per_partition=4096,
+        top_k=4,
+        activation="silu",
+        is_gated=True,
+        inplace=False,
+    )
+    monkeypatch.setattr(
+        deep_gemm_runner.envs.SGLANG_OPT_DG_COMPACT_EAGER, "get", lambda: True
+    )
+    capture = SimpleNamespace(disable_dispose_tensor=False)
+    monkeypatch.setattr(
+        deep_gemm_runner, "get_flags", lambda: SimpleNamespace(capture=capture)
+    )
+    hidden_states = torch.empty((128, 2048), device="meta")
+    quant_info = DeepGemmMoeQuantInfo(
+        w13_weight=torch.empty((1, 4096, 1), dtype=torch.float8_e4m3fn),
+        w2_weight=torch.empty((1, 2048, 1), dtype=torch.float8_e4m3fn),
+        use_fp8=True,
+        block_shape=[128, 128],
+    )
+    with envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.override("masked"):
+        assert not deep_gemm_runner._should_use_masked_standard_layout(
+            config, quant_info, hidden_states
+        )
+
+        capture.disable_dispose_tensor = True
+        assert deep_gemm_runner._should_use_masked_standard_layout(
+            config, quant_info, hidden_states
+        )
+
+
 def test_standard_layout_auto_memory_policy(monkeypatch):
     config = MoeRunnerConfig(
         num_experts=512,

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -434,6 +434,41 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
         self.assertEqual(future_map.dsa_topk_indices_buf.shape, (4, 3))
         self.assertEqual(future_map.dsa_topk_indices_buf.dtype, torch.int32)
 
+    @patch(
+        "sglang.srt.speculative.spec_utils.spec_need_hidden_states",
+        return_value=False,
+    )
+    def test_future_map_initializes_topk_after_prefill_payload(self, _):
+        future_map = object.__new__(FutureMap)
+        future_map.spec_algo = SimpleNamespace(
+            is_some=Mock(return_value=True),
+            need_topk=Mock(return_value=True),
+        )
+        future_map.req_pool_size = 4
+        future_map.device = "cpu"
+        future_map.need_topk = False
+        future_map.need_hidden_states = False
+        future_map.topk_p_buf = None
+        future_map.topk_index_buf = None
+        future_map.hidden_states_buf = None
+        future_map.draft_probs_buf = None
+
+        future_map._maybe_init_forward_bufs(
+            RelayPayload(bonus_tokens=torch.zeros((2,), dtype=torch.int64))
+        )
+        self.assertFalse(future_map.need_topk)
+
+        future_map._maybe_init_forward_bufs(
+            RelayPayload(
+                bonus_tokens=torch.zeros((2,), dtype=torch.int64),
+                topk_p=torch.zeros((2, 3), dtype=torch.float32),
+                topk_index=torch.zeros((2, 3), dtype=torch.int64),
+            )
+        )
+        self.assertTrue(future_map.need_topk)
+        self.assertEqual(future_map.topk_p_buf.shape, (4, 3))
+        self.assertEqual(future_map.topk_index_buf.shape, (4, 3))
+
 
 class TestDSV4C128StateIndices(unittest.TestCase):
     def test_online_aligned_boundary_has_no_partial_state(self):

test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py

@@ -0,0 +1,284 @@
+"""Unit tests for full-attention KV transfer with prefill pp_size > 1 on
+hybrid-linear models (HybridLinearKVPool)."""
+
+import unittest
+from types import SimpleNamespace
+
+import numpy as np
+
+from sglang.srt.disaggregation.common.conn import CommonKVManager
+from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
+from sglang.srt.disaggregation.prefill import _transfer_start_layer
+from sglang.srt.disaggregation.utils import (
+    build_kv_layer_ids,
+    build_transfer_entry_pairs,
+)
+from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
+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")
+
+
+def _full_attention_ids(*, num_layers: int, interval: int) -> list:
+    return [i for i in range(num_layers) if i % interval == interval - 1]
+
+
+def _hybrid_pool(*, start_layer: int) -> HybridLinearKVPool:
+    pool = HybridLinearKVPool.__new__(HybridLinearKVPool)
+    pool.start_layer = start_layer
+    return pool
+
+
+class TestTransferStartLayer(CustomTestCase):
+    """Bug regression: with prefill pp_size=2 on a 60-layer hybrid-linear model
+    (full_attention_interval=4), stage 1's pool.start_layer is 30 — a global
+    layer index counting linear layers. The decode peer's KV pointer list is
+    dense over the 15 full-attention layers only, so slicing dst[30:38] yielded
+    [] and an IndexError in mooncake send_kvcache_slice. The transfer offset
+    must be the count of full-attention layers before the stage boundary."""
+
+    def test_hybrid_stage1_translates_to_full_attention_offset(self):
+        cfg = SimpleNamespace(
+            full_attention_layer_ids=_full_attention_ids(num_layers=60, interval=4)
+        )
+        self.assertEqual(
+            _transfer_start_layer(
+                pool=_hybrid_pool(start_layer=30), hf_text_config=cfg
+            ),
+            7,
+        )
+
+    def test_hybrid_stage0_is_zero(self):
+        cfg = SimpleNamespace(
+            full_attention_layer_ids=_full_attention_ids(num_layers=60, interval=4)
+        )
+        self.assertEqual(
+            _transfer_start_layer(pool=_hybrid_pool(start_layer=0), hf_text_config=cfg),
+            0,
+        )
+
+    def test_non_hybrid_pool_keeps_global_start_layer(self):
+        cfg = SimpleNamespace(full_attention_layer_ids=[])
+        self.assertEqual(
+            _transfer_start_layer(
+                pool=SimpleNamespace(start_layer=30), hf_text_config=cfg
+            ),
+            30,
+        )
+
+
+class _RecordingKVManager:
+    get_mha_kv_ptrs_with_pp = CommonKVManager.get_mha_kv_ptrs_with_pp
+
+    def __init__(self, *, prefill_start_layer: int, pp_size: int):
+        self.is_mla_backend = False
+        self.is_hybrid_mla_backend = False
+        self.enable_custom_mem_pool = False
+        self.pp_size = pp_size
+        self.kv_args = SimpleNamespace(prefill_start_layer=prefill_start_layer)
+        self.blocks = []
+
+    def _transfer_data(self, mooncake_session_id, transfer_blocks):
+        self.blocks.extend(transfer_blocks)
+        return 0
+
+
+class TestHybridSendUsesLayerIdPairing(CustomTestCase):
+    """Bug regression: a hybrid-linear (non-MLA-flagged) backend fell into the
+    positional MHA slicing path of _send_kvcache_generic even when both peers
+    published layer ids. For a stage with F full-attention layers against a
+    decode peer with N (F < N, F not dividing N), the draft-KV modulo heuristic
+    silently placed the V block at F * (N // F) instead of N — wrong layers
+    transferred, no error. With layer ids published on both sides the pairing
+    must be exact."""
+
+    def _run_case(
+        self, *, model_full_ids: list, stage_full_ids: list, start_offset: int
+    ):
+        num_stage = len(stage_full_ids)
+        num_model = len(model_full_ids)
+        src_ptrs = [1000 + i for i in range(2 * num_stage)]
+        dst_ptrs = [2000 + i for i in range(2 * num_model)]
+        item_lens = [10 + i for i in range(2 * num_stage)]
+        manager = _RecordingKVManager(prefill_start_layer=start_offset, pp_size=2)
+        rc = MooncakeKVManager._send_kvcache_generic(
+            manager,
+            mooncake_session_id="session",
+            src_data_ptrs=src_ptrs,
+            dst_data_ptrs=dst_ptrs,
+            item_lens=item_lens,
+            prefill_data_indices=np.array([0], dtype=np.int32),
+            dst_data_indices=np.array([0], dtype=np.int32),
+            executor=None,
+            src_layer_ids=stage_full_ids * 2,
+            dst_layer_ids=model_full_ids * 2,
+        )
+        self.assertEqual(rc, 0)
+        expected = [
+            (src_ptrs[i], dst_ptrs[start_offset + i], item_lens[i])
+            for i in range(num_stage)
+        ] + [
+            (
+                src_ptrs[num_stage + i],
+                dst_ptrs[num_model + start_offset + i],
+                item_lens[num_stage + i],
+            )
+            for i in range(num_stage)
+        ]
+        self.assertEqual(manager.blocks, expected)
+
+    def test_stage1_f8_of_n15(self):
+        ids = _full_attention_ids(num_layers=60, interval=4)
+        self._run_case(model_full_ids=ids, stage_full_ids=ids[7:], start_offset=7)
+
+    def test_stage0_f7_of_n15(self):
+        ids = _full_attention_ids(num_layers=60, interval=4)
+        self._run_case(model_full_ids=ids, stage_full_ids=ids[:7], start_offset=0)
+
+    def test_f5_of_n12(self):
+        ids = _full_attention_ids(num_layers=48, interval=4)
+        self._run_case(model_full_ids=ids, stage_full_ids=ids[:5], start_offset=0)
+
+
+class TestGetMhaKvPtrsWithPp(CustomTestCase):
+    """Derived property: the modulo heuristic in get_mha_kv_ptrs_with_pp exists
+    for the decode-has-draft-KV layout [K_main, V_main, draft_K, draft_V]. Pin
+    that geometry (15 main + 1 draft layer) so a future rewrite of the
+    heuristic (e.g. to fix the plain-MHA pp>1 F-not-dividing-N case) keeps the
+    draft case intact."""
+
+    def test_draft_kv_geometry_selects_main_v_block(self):
+        manager = SimpleNamespace(kv_args=SimpleNamespace(prefill_start_layer=0))
+        src_kv_ptrs = list(range(30))
+        dst_kv_ptrs = list(range(100, 132))
+        src_k, src_v, dst_k, dst_v, num_layers = (
+            CommonKVManager.get_mha_kv_ptrs_with_pp(manager, src_kv_ptrs, dst_kv_ptrs)
+        )
+        self.assertEqual(src_k, src_kv_ptrs[:15])
+        self.assertEqual(src_v, src_kv_ptrs[15:])
+        self.assertEqual(dst_k, dst_kv_ptrs[:15])
+        self.assertEqual(dst_v, dst_kv_ptrs[15:30])
+        self.assertEqual(num_layers, 15)
+
+
+class TestBuildTransferEntryPairsDuplicateIds(CustomTestCase):
+    """Derived property: layer ids repeat across the K and V tensor groups, so
+    pairing must consume dst occurrences in order (K with K, V with V) rather
+    than by plain id lookup."""
+
+    def test_k_then_v_occurrence_ordering(self):
+        pairs = build_transfer_entry_pairs(
+            src_layer_ids=[3, 7, 3, 7],
+            dst_layer_ids=[3, 7, 11, 3, 7, 11],
+            n_src=4,
+            n_dst=6,
+            allow_positional_fallback=False,
+        )
+        self.assertEqual(pairs, [(0, 0), (1, 1), (2, 3), (3, 4)])
+
+
+def _hybrid_pool_with_ids(*, layer_ids: list) -> HybridLinearKVPool:
+    pool = HybridLinearKVPool.__new__(HybridLinearKVPool)
+    pool.full_attention_layer_id_mapping = layer_ids
+    pool.use_mla = False
+    return pool
+
+
+class TestBuildKvLayerIds(CustomTestCase):
+    """Bug regression: enabling EAGLE appended draft KV buffers to kv_data_ptrs
+    while kv_layer_ids described only the target's entries, so the ids were
+    suppressed entirely and the transfer fell back to positional slicing. Under
+    prefill pp_size > 1 that slices the wrong layers -- prefill pp=2 + EAGLE
+    produced garbled decode output while pp=1 + EAGLE did not."""
+
+    def _stage1_ids(self) -> list:
+        full = _full_attention_ids(num_layers=60, interval=4)
+        return [lid for lid in full if lid >= 30]
+
+    def test_draft_entries_get_a_reserved_band_above_the_target_range(self):
+        """A draft pool that only reports a layer count, not ids."""
+        ids = build_kv_layer_ids(
+            token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=self._stage1_ids()),
+            draft_token_to_kv_pool=SimpleNamespace(layer_num=1),
+            num_draft_entries=2,
+            num_hidden_layers=60,
+        )
+        stage1 = self._stage1_ids()
+        # k0..k(L-1) then v0..v(L-1) per pool, and the pools are concatenated --
+        # so the band repeats per group after the target's ids, not interleaved.
+        self.assertEqual(ids, stage1 + stage1 + [60, 60])
+
+    def test_hybrid_draft_pool_is_remapped_out_of_the_target_range(self):
+        """The EAGLE draft pool for a hybrid-linear model is itself a
+        HybridLinearKVPool that numbers its single MTP layer from zero, so its
+        raw ids collide with target layer 0 and must be remapped into the band."""
+        ids = build_kv_layer_ids(
+            token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=self._stage1_ids()),
+            draft_token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=[0]),
+            num_draft_entries=2,
+            num_hidden_layers=60,
+        )
+        stage1 = self._stage1_ids()
+        self.assertEqual(ids, stage1 + stage1 + [60, 60])
+
+    def test_non_hybrid_pool_publishes_nothing(self):
+        self.assertEqual(
+            build_kv_layer_ids(
+                token_to_kv_pool=SimpleNamespace(),
+                draft_token_to_kv_pool=None,
+                num_draft_entries=0,
+                num_hidden_layers=60,
+            ),
+            [],
+        )
+
+    def test_ragged_draft_registration_is_rejected(self):
+        with self.assertRaises(RuntimeError):
+            build_kv_layer_ids(
+                token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=self._stage1_ids()),
+                draft_token_to_kv_pool=SimpleNamespace(layer_num=2),
+                num_draft_entries=3,
+                num_hidden_layers=60,
+            )
+
+
+class TestDraftBandPairsAcrossPipelineStages(CustomTestCase):
+    """Derived property: a pp=2 prefill stage and a pp=1 decode peer, both with
+    an EAGLE draft pool, must pair on layer id -- the stage's 8 full-attention
+    layers land on the decode peer's matching K and V entries, and the draft
+    band lands on the decode peer's draft entries rather than on layer 0."""
+
+    def test_stage1_pairs_onto_the_decode_layout(self):
+        full = _full_attention_ids(num_layers=60, interval=4)
+        stage1 = [lid for lid in full if lid >= 30]
+        src = build_kv_layer_ids(
+            token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=stage1),
+            draft_token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=[0]),
+            num_draft_entries=2,
+            num_hidden_layers=60,
+        )
+        dst = build_kv_layer_ids(
+            token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=full),
+            draft_token_to_kv_pool=_hybrid_pool_with_ids(layer_ids=[0]),
+            num_draft_entries=2,
+            num_hidden_layers=60,
+        )
+        pairs = build_transfer_entry_pairs(
+            src, dst, len(src), len(dst), allow_positional_fallback=False
+        )
+        k_offset = len(full) - len(stage1)
+        self.assertEqual(
+            pairs,
+            # K block, then V block, then the two draft entries at the tail.
+            [(i, k_offset + i) for i in range(len(stage1))]
+            + [(len(stage1) + i, len(full) + k_offset + i) for i in range(len(stage1))]
+            + [
+                (2 * len(stage1), 2 * len(full)),
+                (2 * len(stage1) + 1, 2 * len(full) + 1),
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_staging_draft_kv_slots.py

@@ -0,0 +1,133 @@
+"""Staging slot ids stay aligned once a draft KV pool is registered.
+
+The staging gather writes every k_buffer and then every v_buffer, while
+kv_data_ptrs (and therefore kv_layer_ids) is ordered
+[K target, V target, K draft, V draft]. Labelling slots with kv_layer_ids
+silently pairs a layer's KV with another layer's staging slot as soon as a
+draft pool exists.
+"""
+
+import unittest
+
+from sglang.srt.disaggregation.utils import (
+    build_staging_slot_metadata,
+    build_transfer_entry_pairs,
+)
+from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool
+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 _Pool(MHATokenToKVPool):
+    def __init__(self, tag, layer_ids):
+        self.k_buffer = [f"{tag}K{i}" for i in layer_ids]
+        self.v_buffer = [f"{tag}V{i}" for i in layer_ids]
+
+
+class _Wrapper(HybridLinearKVPool):
+    def __init__(self, inner):
+        self.full_kv_pool = inner
+
+
+def _kv_layer_ids(target_ids, draft_ids):
+    """kv_data_ptrs order: K target, V target, K draft, V draft."""
+    return list(target_ids) + list(target_ids) + list(draft_ids) + list(draft_ids)
+
+
+class TestStagingDraftKvSlots(CustomTestCase):
+    def test_draft_slots_follow_gather_order(self):
+        target, draft = [87, 91], [92]
+        k_buffers, v_buffers, slot_ids = build_staging_slot_metadata(
+            kv_layer_ids=_kv_layer_ids(target, draft),
+            num_draft_entries=2,
+            kv_pool=_Pool("t", target),
+            draft_kv_pool=_Pool("d", draft),
+        )
+        self.assertEqual(k_buffers, ["tK87", "tK91", "dK92"])
+        self.assertEqual(v_buffers, ["tV87", "tV91", "dV92"])
+        self.assertEqual(slot_ids, [87, 91, 92, 87, 91, 92])
+        self.assertNotEqual(slot_ids, _kv_layer_ids(target, draft))
+
+    def test_without_draft_matches_kv_layer_ids(self):
+        # The two orders coincide with no draft pool, so every deployment that
+        # predates draft KV must keep its exact slot labelling.
+        target = [3, 7]
+        _, _, slot_ids = build_staging_slot_metadata(
+            kv_layer_ids=_kv_layer_ids(target, []),
+            num_draft_entries=0,
+            kv_pool=_Pool("t", target),
+            draft_kv_pool=None,
+        )
+        self.assertEqual(slot_ids, _kv_layer_ids(target, []))
+
+    def test_pp_stage_pairs_against_full_decode(self):
+        # A prefill stage holds a slice of the layers while decode holds them
+        # all, so the ids -- not the positions -- have to drive the pairing.
+        src = build_staging_slot_metadata(
+            kv_layer_ids=_kv_layer_ids([87, 91], [92]),
+            num_draft_entries=2,
+            kv_pool=_Pool("t", [87, 91]),
+            draft_kv_pool=_Pool("d", [92]),
+        )[2]
+        decode_target = [3, 7, 11, 87, 91]
+        dst = build_staging_slot_metadata(
+            kv_layer_ids=_kv_layer_ids(decode_target, [92]),
+            num_draft_entries=2,
+            kv_pool=_Pool("t", decode_target),
+            draft_kv_pool=_Pool("d", [92]),
+        )[2]
+        pairs = build_transfer_entry_pairs(src, dst, len(src), len(dst))
+        self.assertEqual(len(pairs), len(src))
+        for i, j in pairs:
+            self.assertEqual(src[i], dst[j])
+        self.assertEqual(len({j for _, j in pairs}), len(pairs))
+
+    def test_hybrid_wrapper_pools_are_unwrapped(self):
+        # A hybrid draft pool that is left wrapped looks exactly like a draft
+        # pool with no buffers, which drops draft KV out of staging.
+        target, draft = [87, 91], [92]
+        k_buffers, _, slot_ids = build_staging_slot_metadata(
+            kv_layer_ids=_kv_layer_ids(target, draft),
+            num_draft_entries=2,
+            kv_pool=_Wrapper(_Pool("t", target)),
+            draft_kv_pool=_Wrapper(_Pool("d", draft)),
+        )
+        self.assertEqual(k_buffers, ["tK87", "tK91", "dK92"])
+        self.assertEqual(slot_ids, [87, 91, 92, 87, 91, 92])
+
+    def test_undescribable_draft_still_yields_target_buffers(self):
+        # Returning nothing here left the caller skipping set_kv_buffer_tensors
+        # entirely, and staging then came up with no buffers at all.
+        class _NoBuffers:
+            pass
+
+        k_buffers, v_buffers, slot_ids = build_staging_slot_metadata(
+            kv_layer_ids=_kv_layer_ids([87], [92]),
+            num_draft_entries=2,
+            kv_pool=_Pool("t", [87]),
+            draft_kv_pool=_NoBuffers(),
+        )
+        self.assertEqual(k_buffers, ["tK87"])
+        self.assertEqual(v_buffers, ["tV87"])
+        self.assertEqual(slot_ids, [])
+
+    def test_pool_without_contiguous_tensors_is_declined(self):
+        # MLA pools have no k_buffer/v_buffer to stage; the caller relies on None
+        # to skip the registration rather than register empty lists.
+        class _NoBuffers:
+            pass
+
+        self.assertIsNone(
+            build_staging_slot_metadata(
+                kv_layer_ids=[],
+                num_draft_entries=0,
+                kv_pool=_NoBuffers(),
+                draft_kv_pool=None,
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/attention/test_gdn_flashinfer_alignment.py

@@ -0,0 +1,412 @@
+import unittest
+from unittest import mock
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
+    FlashInferGDNKernel,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def _view_with_pointer_mod(
+    shape: tuple[int, ...], dtype: torch.dtype, pointer_mod: int
+) -> torch.Tensor:
+    numel = 1
+    for dim in shape:
+        numel *= dim
+    element_size = dtype.itemsize
+    base = torch.empty(numel + 32 // element_size, dtype=dtype)
+    for offset in range(32 // element_size):
+        view = base[offset : offset + numel]
+        if view.data_ptr() % 32 == pointer_mod:
+            return view.view(shape)
+    raise AssertionError(f"Could not construct a pointer with mod32={pointer_mod}")
+
+
+def _make_kernel_without_flashinfer() -> FlashInferGDNKernel:
+    kernel = object.__new__(FlashInferGDNKernel)
+    # Match the SM100 path used by the CPU-only fake prefill tests. Real
+    # instances initialize this from the detected SM architecture in __init__.
+    kernel._prefill_needs_fp32_state = False
+    kernel._aligned_input_buffers = {}
+    kernel._aligned_parameter_cache = {}
+    kernel._verify_intermediate_buffers = {}
+    kernel._alignment_fallback_warned = False
+    return kernel
+
+
+class TestFlashInferGDNAlignment(unittest.TestCase):
+    def test_extend_writes_directly_to_preallocated_output(self):
+        kernel = _make_kernel_without_flashinfer()
+        kernel.use_state_pool = True
+        captured = {}
+
+        def fake_prefill(**kwargs):
+            captured.update(kwargs)
+            kwargs["output"].fill_(7.0)
+            kwargs["output_state"].copy_(kwargs["initial_state"])
+            return kwargs["output"], kwargs["output_state"]
+
+        kernel._prefill_fn = fake_prefill
+        q = torch.ones((1, 3, 1, 4), dtype=torch.bfloat16)
+        k = torch.ones_like(q)
+        v = torch.ones((1, 3, 2, 4), dtype=torch.bfloat16)
+        g = torch.zeros((1, 3, 2), dtype=torch.bfloat16)
+        beta = torch.ones_like(g)
+        ssm_states = torch.zeros((3, 2, 4, 4), dtype=torch.bfloat16)
+        physical_output = torch.empty((1, 5, 2, 4), dtype=v.dtype)
+        preallocated_output = physical_output[:, :3]
+
+        with mock.patch(
+            "sglang.kernels.ops.attention.fla.l2norm.l2norm_fwd",
+            side_effect=lambda tensor: tensor,
+        ):
+            result, _, checkpoints = kernel.extend(
+                q,
+                k,
+                v,
+                g,
+                beta,
+                ssm_states=ssm_states,
+                cache_indices=torch.tensor([1], dtype=torch.int32),
+                query_start_loc=torch.tensor([0, 3], dtype=torch.int32),
+                output=preallocated_output,
+            )
+
+        self.assertEqual(captured["output"].data_ptr(), preallocated_output.data_ptr())
+        self.assertEqual(result.data_ptr(), preallocated_output.data_ptr())
+        torch.testing.assert_close(result, torch.full_like(result, 7.0))
+        self.assertIsNone(checkpoints)
+
+    def test_ratio8_bs1_split_view_reproduces_under_alignment(self):
+        # In a BF16 [b_local(8)|a_local(8)] projection, a begins 16 bytes in;
+        # contiguous() is a no-op at BS=1 and cannot meet FlashInfer's 32-byte ABI.
+        projected_ba = torch.empty((2, 16), dtype=torch.bfloat16)
+        _, a = projected_ba.split((8, 8), dim=-1)
+
+        a_bs1 = a[:1]
+        self.assertEqual(a_bs1.stride(), (16, 1))
+        self.assertTrue(a_bs1.is_contiguous())
+        self.assertEqual(a_bs1.data_ptr() % 32, 16)
+        self.assertEqual(a_bs1.contiguous().data_ptr(), a_bs1.data_ptr())
+
+        # BS>1 exposes the row gap, so contiguous() does allocate a rebased,
+        # allocator-aligned tensor. This explains why only BS=1 failed.
+        a_bs2 = a[:2]
+        self.assertFalse(a_bs2.is_contiguous())
+        repaired = a_bs2.contiguous()
+        self.assertNotEqual(repaired.data_ptr(), a_bs2.data_ptr())
+        self.assertEqual(repaired.data_ptr() % 32, 0)
+
+    def test_dynamic_repair_buffer_is_reused_without_allocator_churn(self):
+        kernel = _make_kernel_without_flashinfer()
+        source = _view_with_pointer_mod((1, 1, 8), torch.bfloat16, 16)
+        source.fill_(1)
+
+        first = kernel._prepare_dynamic_input("decode_a", source)
+        first_ptr = first.data_ptr()
+        self.assertEqual(first_ptr % 32, 0)
+        torch.testing.assert_close(first, source)
+        self.assertEqual(len(kernel._aligned_input_buffers), 1)
+
+        source.fill_(2)
+        second = kernel._prepare_dynamic_input("decode_a", source)
+        self.assertIs(second, first)
+        self.assertEqual(second.data_ptr(), first_ptr)
+        torch.testing.assert_close(second, source)
+        self.assertEqual(len(kernel._aligned_input_buffers), 1)
+
+        # Distinct kernel arguments cannot alias because both are live at the
+        # FlashInfer call boundary.
+        other = kernel._prepare_dynamic_input("decode_b", source)
+        self.assertNotEqual(other.data_ptr(), first_ptr)
+        self.assertEqual(len(kernel._aligned_input_buffers), 2)
+
+    def test_decode_repairs_read_only_arguments_before_flashinfer(self):
+        kernel = _make_kernel_without_flashinfer()
+        kernel.use_state_pool = True
+        captured = {}
+
+        def fake_decode(**kwargs):
+            captured.update(kwargs)
+            v = kwargs["v"]
+            return (
+                torch.zeros(
+                    v.shape[0],
+                    1,
+                    v.shape[2],
+                    v.shape[3],
+                    dtype=v.dtype,
+                ),
+                None,
+            )
+
+        kernel._decode_fn = fake_decode
+
+        q = torch.empty(1, 1, 1, 128, dtype=torch.bfloat16)
+        k = torch.empty_like(q)
+        v = torch.empty(1, 1, 8, 128, dtype=torch.bfloat16)
+        a = _view_with_pointer_mod((1, 1, 8), torch.bfloat16, 16)
+        b = _view_with_pointer_mod((1, 1, 8), torch.bfloat16, 16)
+        A_log = _view_with_pointer_mod((8,), torch.float32, 4)
+        dt_bias = _view_with_pointer_mod((8,), torch.bfloat16, 2)
+        state = torch.zeros(2, 8, 128, 128, dtype=torch.bfloat16)
+        cache_indices = _view_with_pointer_mod((1,), torch.int32, 4)
+
+        result = kernel.decode(
+            q,
+            k,
+            v,
+            a,
+            b,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            ssm_states=state,
+            cache_indices=cache_indices,
+            query_start_loc=torch.tensor([0, 1], dtype=torch.int32),
+        )
+
+        self.assertEqual(result.shape, (1, 1, 8, 128))
+        for name in (
+            "q",
+            "k",
+            "v",
+            "A_log",
+            "a",
+            "dt_bias",
+            "b",
+            "initial_state",
+            "initial_state_indices",
+        ):
+            with self.subTest(name=name):
+                self.assertEqual(captured[name].data_ptr() % 32, 0)
+        torch.testing.assert_close(captured["a"], a)
+        torch.testing.assert_close(captured["b"], b)
+
+    def test_gate_parameter_cache_preserves_backend_dtype_contract(self):
+        kernel = _make_kernel_without_flashinfer()
+        A_log = torch.empty(8, dtype=torch.bfloat16)
+        dt_bias = torch.empty(8, dtype=torch.bfloat16)
+
+        A_log_sm90, _ = kernel._prepare_gate_parameters(A_log, dt_bias)
+        A_log_sm100, _ = kernel._prepare_gate_parameters(
+            A_log, dt_bias, A_log_dtype=torch.float32
+        )
+
+        self.assertEqual(A_log_sm90.dtype, torch.bfloat16)
+        self.assertEqual(A_log_sm100.dtype, torch.float32)
+        self.assertEqual(A_log_sm90.data_ptr() % 32, 0)
+        self.assertEqual(A_log_sm100.data_ptr() % 32, 0)
+        self.assertIs(
+            kernel._prepare_gate_parameters(A_log, dt_bias)[0],
+            A_log_sm90,
+        )
+
+    def test_mutable_state_falls_back_without_losing_writeback(self):
+        kernel = _make_kernel_without_flashinfer()
+        captured = {}
+        expected = torch.empty(1)
+
+        class FakeFallback:
+            def decode(self, *args, **kwargs):
+                captured["args"] = args
+                captured["kwargs"] = kwargs
+                return expected
+
+        kernel._alignment_fallback_kernel = FakeFallback()
+        state = _view_with_pointer_mod((2, 8, 4, 4), torch.bfloat16, 16)
+        q = torch.empty(1, 1, 1, 4, dtype=torch.bfloat16)
+        k = torch.empty_like(q)
+        v = torch.empty(1, 1, 8, 4, dtype=torch.bfloat16)
+        a = torch.empty(1, 1, 8, dtype=torch.bfloat16)
+        b = torch.empty_like(a)
+        cache_indices = torch.zeros(1, dtype=torch.int32)
+        query_start_loc = torch.tensor([0, 1], dtype=torch.int32)
+
+        result = kernel.decode(
+            q,
+            k,
+            v,
+            a,
+            b,
+            A_log=torch.zeros(8),
+            dt_bias=torch.zeros(8, dtype=torch.bfloat16),
+            ssm_states=state,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+        )
+
+        self.assertIs(result, expected)
+        self.assertIs(captured["kwargs"]["ssm_states"], state)
+        self.assertEqual(len(kernel._aligned_input_buffers), 0)
+
+    def test_mutable_mtp_workspace_falls_back_without_copying(self):
+        kernel = _make_kernel_without_flashinfer()
+        kernel.use_state_pool = True
+        captured = {}
+        expected = torch.empty(1)
+
+        class FakeFallback:
+            def target_verify(self, **kwargs):
+                captured.update(kwargs)
+                return expected
+
+        kernel._alignment_fallback_kernel = FakeFallback()
+        q = torch.empty(1, 2, 1, 4, dtype=torch.bfloat16)
+        k = torch.empty_like(q)
+        v = torch.empty(1, 2, 8, 4, dtype=torch.bfloat16)
+        a = torch.empty(1, 2, 8, dtype=torch.bfloat16)
+        b = torch.empty_like(a)
+        state = torch.empty(2, 8, 4, 4, dtype=torch.bfloat16)
+        workspace = _view_with_pointer_mod((2, 2, 8, 4, 4), torch.bfloat16, 16)
+
+        result = kernel.target_verify(
+            torch.zeros(8),
+            torch.zeros(8, dtype=torch.bfloat16),
+            q,
+            k,
+            v,
+            a,
+            b,
+            ssm_states=state,
+            cache_indices=torch.zeros(1, dtype=torch.int32),
+            query_start_loc=torch.tensor([0, 2], dtype=torch.int32),
+            intermediate_states_buffer=workspace,
+            intermediate_state_indices=torch.zeros(1, 2, dtype=torch.int32),
+            cache_steps=2,
+            retrieve_parent_token=None,
+        )
+
+        self.assertIs(result, expected)
+        self.assertIs(captured["intermediate_states_buffer"], workspace)
+        self.assertEqual(len(kernel._aligned_input_buffers), 0)
+
+    def test_mtp_padded_capture_uses_stable_exact_batch_workspace_and_copies_back(self):
+        kernel = _make_kernel_without_flashinfer()
+        kernel.use_state_pool = True
+        captured_ptrs = []
+
+        def fake_mtp(**kwargs):
+            workspace = kwargs["intermediate_states_buffer"]
+            captured_ptrs.append(workspace.data_ptr())
+            self.assertEqual(workspace.shape[0], 8)
+            for row in range(workspace.shape[0]):
+                workspace[row].fill_(row + 1)
+            return torch.zeros_like(kwargs["v"]), None
+
+        kernel._mtp_fn = fake_mtp
+        workspace = torch.zeros((7, 2, 8, 4, 4), dtype=torch.bfloat16)
+
+        def run_once():
+            return kernel.target_verify(
+                torch.zeros(8),
+                torch.zeros(8, dtype=torch.bfloat16),
+                torch.empty(1, 16, 1, 4, dtype=torch.bfloat16),
+                torch.empty(1, 16, 1, 4, dtype=torch.bfloat16),
+                torch.empty(1, 16, 8, 4, dtype=torch.bfloat16),
+                torch.empty(1, 16, 8, dtype=torch.bfloat16),
+                torch.empty(1, 16, 8, dtype=torch.bfloat16),
+                ssm_states=torch.zeros(8, 8, 4, 4, dtype=torch.bfloat16),
+                cache_indices=torch.zeros(8, dtype=torch.int32),
+                query_start_loc=torch.arange(0, 18, 2, dtype=torch.int32),
+                intermediate_states_buffer=workspace,
+                intermediate_state_indices=torch.arange(8, dtype=torch.int32),
+                cache_steps=2,
+                retrieve_parent_token=None,
+            )
+
+        self.assertEqual(run_once().shape, (1, 16, 8, 4))
+        for row in range(workspace.shape[0]):
+            torch.testing.assert_close(
+                workspace[row], torch.full_like(workspace[row], row + 1)
+            )
+        self.assertEqual(len(kernel._verify_intermediate_buffers), 1)
+
+        workspace.zero_()
+        run_once()
+        self.assertEqual(captured_ptrs[0], captured_ptrs[1])
+        self.assertEqual(len(kernel._verify_intermediate_buffers), 1)
+
+    def test_mtp_pool_sized_batch_keeps_zero_copy_fast_path(self):
+        kernel = _make_kernel_without_flashinfer()
+        kernel.use_state_pool = True
+        workspace = torch.zeros((7, 2, 8, 4, 4), dtype=torch.bfloat16)
+        captured = {}
+
+        def fake_mtp(**kwargs):
+            captured.update(kwargs)
+            return torch.zeros_like(kwargs["v"]), None
+
+        kernel._mtp_fn = fake_mtp
+        result = kernel.target_verify(
+            torch.zeros(8),
+            torch.zeros(8, dtype=torch.bfloat16),
+            torch.empty(1, 4, 1, 4, dtype=torch.bfloat16),
+            torch.empty(1, 4, 1, 4, dtype=torch.bfloat16),
+            torch.empty(1, 4, 8, 4, dtype=torch.bfloat16),
+            torch.empty(1, 4, 8, dtype=torch.bfloat16),
+            torch.empty(1, 4, 8, dtype=torch.bfloat16),
+            ssm_states=torch.zeros(7, 8, 4, 4, dtype=torch.bfloat16),
+            cache_indices=torch.zeros(2, dtype=torch.int32),
+            query_start_loc=torch.arange(0, 6, 2, dtype=torch.int32),
+            intermediate_states_buffer=workspace,
+            intermediate_state_indices=torch.arange(7, dtype=torch.int32),
+            cache_steps=2,
+            retrieve_parent_token=None,
+        )
+
+        self.assertEqual(result.shape, (1, 4, 8, 4))
+        self.assertEqual(
+            captured["intermediate_states_buffer"].data_ptr(), workspace.data_ptr()
+        )
+        self.assertEqual(len(kernel._verify_intermediate_buffers), 0)
+
+    def test_mtp_padded_workspace_is_reused_across_sequential_layer_pools(self):
+        kernel = _make_kernel_without_flashinfer()
+        kernel.use_state_pool = True
+        captured_ptrs = []
+        call_value = 0
+
+        def fake_mtp(**kwargs):
+            nonlocal call_value
+            call_value += 1
+            scratch = kwargs["intermediate_states_buffer"]
+            captured_ptrs.append(scratch.data_ptr())
+            scratch.fill_(call_value)
+            return torch.zeros_like(kwargs["v"]), None
+
+        kernel._mtp_fn = fake_mtp
+
+        def run(pool):
+            kernel.target_verify(
+                torch.zeros(8),
+                torch.zeros(8, dtype=torch.bfloat16),
+                torch.empty(1, 16, 1, 4, dtype=torch.bfloat16),
+                torch.empty(1, 16, 1, 4, dtype=torch.bfloat16),
+                torch.empty(1, 16, 8, 4, dtype=torch.bfloat16),
+                torch.empty(1, 16, 8, dtype=torch.bfloat16),
+                torch.empty(1, 16, 8, dtype=torch.bfloat16),
+                ssm_states=torch.zeros(8, 8, 4, 4, dtype=torch.bfloat16),
+                cache_indices=torch.zeros(8, dtype=torch.int32),
+                query_start_loc=torch.arange(0, 18, 2, dtype=torch.int32),
+                intermediate_states_buffer=pool,
+                intermediate_state_indices=torch.arange(8, dtype=torch.int32),
+                cache_steps=2,
+                retrieve_parent_token=None,
+            )
+
+        first_pool = torch.zeros((7, 2, 8, 4, 4), dtype=torch.bfloat16)
+        second_pool = torch.zeros_like(first_pool)
+        run(first_pool)
+        run(second_pool)
+
+        self.assertEqual(captured_ptrs[0], captured_ptrs[1])
+        torch.testing.assert_close(first_pool, torch.ones_like(first_pool))
+        torch.testing.assert_close(second_pool, torch.full_like(second_pool, 2))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/moe/test_flashinfer_a2a_wide_ep.py

@@ -0,0 +1,87 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.moe.ep_moe_kernels import fused_moe_dispatch_index
+from sglang.srt.layers.moe.moe_runner.base import (
+    FusedOpPool,
+    PermuteMethodPool,
+)
+from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
+    _max_tokens_per_scattered_source,
+    _scattered_source_token_counts,
+    _workspace_size_for_namespace,
+)
+from sglang.srt.layers.quantization import fp8  # noqa: F401
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestFlashinferA2AWideEPPlumbing(CustomTestCase):
+    def test_runner_paths_are_registered(self):
+        self.assertIn(("flashinfer", "flashinfer_trtllm"), FusedOpPool._fused_funcs)
+        self.assertIn(
+            ("flashinfer", "flashinfer_trtllm_routed"), FusedOpPool._fused_funcs
+        )
+        self.assertIn(
+            ("flashinfer", "deep_gemm"), PermuteMethodPool._pre_permute_methods
+        )
+        self.assertIn(
+            ("deep_gemm", "flashinfer"), PermuteMethodPool._post_permute_methods
+        )
+
+    def test_dp4_tp4_uses_physical_source_rank_geometry(self):
+        self.assertEqual(_max_tokens_per_scattered_source([2048] * 4, 4), 512)
+        self.assertEqual(_max_tokens_per_scattered_source([1, 0, 0, 0], 4), 1)
+        self.assertEqual(_max_tokens_per_scattered_source([7, 3, 2, 1], 4), 2)
+        self.assertEqual(_max_tokens_per_scattered_source([512] * 16, 1), 512)
+
+    def test_target_and_draft_decode_use_distinct_workspaces(self):
+        sizes = {
+            _workspace_size_for_namespace(4096, speculative=speculative)
+            for speculative in (False, True)
+        }
+        self.assertEqual(sizes, {4096, 4224})
+
+    def test_prefill_ag_expands_dp_counts_to_physical_source_ranks(self):
+        self.assertEqual(
+            _scattered_source_token_counts([7, 3], 4),
+            [2, 2, 2, 1, 1, 1, 1, 0],
+        )
+        self.assertEqual(_scattered_source_token_counts([4] * 16, 1), [4] * 16)
+
+    def test_deepgemm_dispatch_marks_empty_expert_lanes_invalid(self):
+        topk_ids = torch.tensor([-1, 0, -1, 1], dtype=torch.int32, device="cuda")
+        masked_m, src2dst = fused_moe_dispatch_index(
+            topk_ids, num_local_experts=2, m_max=4
+        )
+
+        torch.testing.assert_close(
+            masked_m, torch.tensor([1, 1], dtype=torch.int32, device="cuda")
+        )
+        torch.testing.assert_close(
+            src2dst,
+            torch.tensor([-1, 0, -1, 4], dtype=torch.int32, device="cuda"),
+        )
+
+    def test_global_expert_mapping_is_fused_into_dispatch_index(self):
+        global_ids = torch.tensor(
+            [-1, 15, 16, 17, 31, 32], dtype=torch.int32, device="cuda"
+        )
+        masked_m, src2dst = fused_moe_dispatch_index(
+            global_ids, num_local_experts=2, m_max=4, expert_start=16
+        )
+
+        torch.testing.assert_close(
+            masked_m, torch.tensor([1, 1], dtype=torch.int32, device="cuda")
+        )
+        torch.testing.assert_close(
+            src2dst,
+            torch.tensor([-1, -1, 0, 4, -1, -1], dtype=torch.int32, device="cuda"),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/moe/test_flashinfer_cutedsl_dispatch.py

@@ -0,0 +1,63 @@
+import sys
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+import sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl as cutedsl_runner
+from sglang.srt.layers.moe.token_dispatcher.standard import (
+    StandardCombineInput,
+    StandardDispatchOutput,
+)
+from sglang.srt.layers.moe.topk import StandardTopKOutput
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def test_flashinfer_prefill_returns_standard_combine_input():
+    dispatch_output = StandardDispatchOutput(
+        hidden_states=torch.empty(2, 16, dtype=torch.bfloat16),
+        hidden_states_scale=None,
+        topk_output=StandardTopKOutput(
+            topk_weights=torch.empty(2, 1),
+            topk_ids=torch.zeros(2, 1, dtype=torch.int32),
+            router_logits=None,
+        ),
+    )
+    expected_output = torch.empty(2, 16, dtype=torch.bfloat16)
+    wrapper = Mock()
+    wrapper.run.return_value = expected_output
+    quant_info = SimpleNamespace(
+        wrapper=wrapper,
+        use_per_token_activation=False,
+        a1_scale=torch.tensor(1.0),
+        a2_scale=torch.tensor(1.0),
+        w13_weight=object(),
+        w13_weight_sf=object(),
+        w1_alpha=object(),
+        w2_weight=object(),
+        w2_weight_sf=object(),
+        w2_alpha=object(),
+    )
+    runner_config = SimpleNamespace(activation="silu")
+
+    with patch(
+        "sglang.srt.layers.quantization.fp4_utils.fp4_quantize",
+        return_value=(
+            torch.empty(2, 8, dtype=torch.uint8),
+            torch.empty(2, 1, dtype=torch.float8_e4m3fn),
+        ),
+    ):
+        result = cutedsl_runner.fused_experts_flashinfer_to_flashinfer_cutedsl_fp4(
+            dispatch_output, quant_info, runner_config
+        )
+
+    assert isinstance(result, StandardCombineInput)
+    assert result.hidden_states is expected_output
+    wrapper.run.assert_called_once()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/layers/moe/test_qwen35_flashinfer_fusion.py

@@ -0,0 +1,205 @@
+from dataclasses import dataclass
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.layers.flashinfer_mnnvl_cutedsl import (
+    FlashInferMNNVLCuteDSLARFusion,
+    _with_early_finalize_shared_load,
+)
+from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
+    Qwen35MoeFinalizeHandoff,
+    is_supported_forward_mode,
+    resolve_max_m,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.models.qwen3_5_text import Qwen3_5ForCausalLM
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=20, suite="base-c-test-cpu")
+
+
+@dataclass(frozen=True)
+class _TestPreset:
+    load_shared_expert_before_pdl: bool = False
+
+
+@dataclass(frozen=True)
+class _TestTarget:
+    preset: object
+
+
+@dataclass(frozen=True)
+class _TestRoutes:
+    targets: tuple[_TestTarget, ...]
+
+
+@dataclass(frozen=True)
+class _TestProfile:
+    finalize_routes: _TestRoutes
+
+
+@dataclass(frozen=True)
+class _TestConfig:
+    profiles: tuple[_TestProfile, ...]
+
+
+@pytest.mark.parametrize(
+    ("forward_mode", "expected"),
+    [
+        (ForwardMode.DECODE, True),
+        (ForwardMode.EXTEND, True),
+        (ForwardMode.IDLE, False),
+        (ForwardMode.TARGET_VERIFY, True),
+        (ForwardMode.DRAFT_EXTEND_V2, False),
+    ],
+)
+def test_supported_forward_modes(forward_mode, expected):
+    assert is_supported_forward_mode(forward_mode) is expected
+
+
+def test_framework_capacity_is_maximum_of_all_sources():
+    graph = SimpleNamespace(
+        decode=SimpleNamespace(max_bs=512, bs=[1, 64, 256]),
+        prefill=SimpleNamespace(max_bs=4096, bs=[1024, 2048, 4096]),
+    )
+    server_args = SimpleNamespace(
+        cuda_graph_config=graph,
+        cutedsl_moe_max_num_tokens=lambda: 8192,
+    )
+    runner = SimpleNamespace(server_args=server_args, max_running_requests=2048)
+
+    assert resolve_max_m(runner) == 8192
+
+
+def test_deferred_handoff_reuses_producer_storage():
+    m, top_k, hidden_size = 3, 10, 16
+    gemm2_out = torch.empty(m * top_k + 4, hidden_size, dtype=torch.bfloat16)
+    expert_weights = torch.empty(m, top_k, dtype=torch.bfloat16)
+    permuted_indices = torch.empty(m, top_k, dtype=torch.int32)
+    gated_shared_output = torch.empty(m, hidden_size, dtype=torch.bfloat16)
+    deferred = SimpleNamespace(
+        gemm2_out=gemm2_out,
+        expert_weights=expert_weights,
+        expanded_idx_to_permuted_idx=permuted_indices,
+        top_k=top_k,
+    )
+
+    handoff = Qwen35MoeFinalizeHandoff.from_flashinfer(
+        deferred,
+        gated_shared_output=gated_shared_output,
+        m=m,
+    )
+
+    assert handoff.routed_output.data_ptr() == gemm2_out.data_ptr()
+    assert handoff.expert_weights.data_ptr() == expert_weights.data_ptr()
+    assert handoff.permuted_indices.data_ptr() == permuted_indices.data_ptr()
+    assert handoff.gated_shared_output is gated_shared_output
+
+
+def test_qwen_workspace_config_enables_only_supported_finalize_presets():
+    untouched_preset = object()
+    default_config = _TestConfig(
+        profiles=(
+            _TestProfile(
+                finalize_routes=_TestRoutes(
+                    targets=(
+                        _TestTarget(_TestPreset()),
+                        _TestTarget(untouched_preset),
+                    )
+                )
+            ),
+        )
+    )
+
+    qwen_config = _with_early_finalize_shared_load(default_config)
+
+    assert qwen_config is not default_config
+    assert (
+        default_config.profiles[0]
+        .finalize_routes.targets[0]
+        .preset.load_shared_expert_before_pdl
+        is False
+    )
+    assert (
+        qwen_config.profiles[0]
+        .finalize_routes.targets[0]
+        .preset.load_shared_expert_before_pdl
+        is True
+    )
+    assert qwen_config.profiles[0].finalize_routes.targets[1].preset is untouched_preset
+
+
+def test_wrapper_calls_only_the_stable_unified_api():
+    calls = []
+    wrapper = object.__new__(FlashInferMNNVLCuteDSLARFusion)
+    wrapper.hidden_size = 8
+    wrapper.top_k = 2
+    wrapper.max_m = 4
+    wrapper.rms_epsilon = 1e-5
+    wrapper.weight_bias = 0.0
+    wrapper.device = torch.device("cpu")
+    wrapper.workspace = object()
+    wrapper.supports = lambda m: True
+    wrapper._patterns = SimpleNamespace(
+        kARResidualRMSNorm=1,
+        kMoEFinalizeARResidualRMSNorm=7,
+    )
+    wrapper._allreduce_fusion = lambda **kwargs: calls.append(kwargs)
+
+    routed_output = torch.empty(8, 8, dtype=torch.bfloat16)
+    expert_weights = torch.empty(4, 2, dtype=torch.bfloat16)
+    permuted_indices = torch.empty(4, 2, dtype=torch.int32)
+    gated_shared_output = torch.empty(4, 8, dtype=torch.bfloat16)
+    residual = torch.empty(4, 8, dtype=torch.bfloat16)
+    gamma = torch.empty(8, dtype=torch.bfloat16)
+    norm_output = torch.empty_like(residual)
+    residual_output = torch.empty_like(residual)
+
+    wrapper.moe_finalize_all_reduce_rms_norm(
+        routed_output=routed_output,
+        expert_weights=expert_weights,
+        permuted_indices=permuted_indices,
+        gated_shared_output=gated_shared_output,
+        residual=residual,
+        gamma=gamma,
+        norm_output=norm_output,
+        residual_output=residual_output,
+    )
+
+    assert calls[0]["launch_with_pdl"] is True
+    assert "routed_scaling_factor" not in calls[0]
+
+    wrapper.all_reduce_residual_rms_norm(
+        local_contribution=residual,
+        residual=residual,
+        gamma=gamma,
+        norm_output=norm_output,
+        residual_output=residual_output,
+    )
+
+    assert calls[1]["pattern"] == 1
+    assert calls[1]["launch_with_pdl"] is True
+    assert "routed_scaling_factor" not in calls[1]
+    assert "expanded_idx_to_permuted_idx" not in calls[1]
+
+
+def test_text_entry_wrapper_delegates_pre_capture_prepare():
+    calls = []
+    runner = object()
+    wrapper = SimpleNamespace(
+        model=SimpleNamespace(
+            prepare_before_cuda_graph_capture=lambda value: calls.append(value)
+        )
+    )
+
+    Qwen3_5ForCausalLM.prepare_before_cuda_graph_capture(wrapper, runner)
+
+    assert calls == [runner]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/layers/quantization/test_flashinfer_pr4266_bf16_gemm.py

@@ -0,0 +1,62 @@
+import pytest
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.quantization.unquant import (
+    _FLASHINFER_PR4266_TUNED_TACTICS,
+    Bf16GemmBackend,
+    should_enable_bf16_splitk_gemm,
+    use_flashinfer_pr4266_bf16_gemm,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+@pytest.mark.parametrize("m,n,k", _FLASHINFER_PR4266_TUNED_TACTICS)
+def test_flashinfer_pr4266_selects_tuned_oakhaven_shape(m: int, n: int, k: int):
+    assert use_flashinfer_pr4266_bf16_gemm(m, n, k)
+
+
+@pytest.mark.parametrize("m", [0, 33, 64])
+@pytest.mark.parametrize("n,k", [(256, 8192), (512, 8192), (2304, 8192), (2560, 8192)])
+def test_flashinfer_pr4266_keeps_large_m_on_existing_path(m: int, n: int, k: int):
+    assert not use_flashinfer_pr4266_bf16_gemm(m, n, k)
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (1, 1024, 2048),
+        (3, 256, 8192),
+        (16, 8192, 4096),
+        (32, 4096, 8192),
+    ],
+)
+def test_flashinfer_pr4266_rejects_unmeasured_shapes(shape: tuple[int, int, int]):
+    assert not use_flashinfer_pr4266_bf16_gemm(*shape)
+
+
+def test_flashinfer_pr4266_backend_is_explicit():
+    assert Bf16GemmBackend.FLASHINFER_PR4266.value == "flashinfer_pr4266"
+
+
+def test_bf16_splitk_is_enabled_by_default():
+    assert envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.default is True
+    with envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.override(True):
+        assert should_enable_bf16_splitk_gemm(Bf16GemmBackend.CUTEDSL)
+
+
+def test_bf16_splitk_env_kill_switch():
+    with envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.override(False):
+        assert not should_enable_bf16_splitk_gemm(Bf16GemmBackend.CUTEDSL)
+
+
+def test_bf16_splitk_does_not_override_torch_backend():
+    with envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.override(True):
+        assert not should_enable_bf16_splitk_gemm(Bf16GemmBackend.TORCH)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/layers/test_radix_linear_attention.py

@@ -0,0 +1,232 @@
+"""CPU regression coverage for padded linear-attention inputs and outputs."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+import sglang.srt.layers.radix_linear_attention as radix_linear_attention
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class _FakeAttentionBackend:
+    def forward(
+        self,
+        *,
+        layer,
+        forward_batch,
+        mixed_qkv,
+        a,
+        b,
+        linear_attn_output=None,
+    ):
+        del layer
+        torch.testing.assert_close(forward_batch.out_cache_loc, torch.arange(3))
+        assert mixed_qkv.shape[0] == 3
+        assert a.shape[0] == 3
+        assert b.shape[0] == 3
+        if linear_attn_output is None:
+            return torch.full((1, 3, 2, 4), 5.0)
+        linear_attn_output.fill_(5.0)
+        return linear_attn_output
+
+
+class _FailingAttentionBackend:
+    def forward(self, **kwargs):
+        del kwargs
+        raise RuntimeError("backend failure")
+
+
+class _ExtendMode:
+    def is_extend(self):
+        return True
+
+    def is_target_verify(self):
+        return False
+
+
+class _TargetVerifyMode:
+    def is_extend(self):
+        return True
+
+    def is_target_verify(self):
+        return True
+
+
+class _PhysicalAttentionBackend:
+    def forward(self, *, layer, forward_batch, mixed_qkv, a, b):
+        del layer, forward_batch
+        assert mixed_qkv.shape[0] == 5
+        assert a.shape[0] == 5
+        assert b.shape[0] == 5
+        return torch.full((1, 5, 2, 4), 9.0)
+
+
+class TestRadixLinearAttentionPadding(CustomTestCase):
+    def test_eager_padded_input_is_sliced_and_output_shape_is_restored(self):
+        layer = radix_linear_attention.RadixLinearAttention(
+            layer_id=0,
+            num_q_heads=1,
+            num_k_heads=1,
+            num_v_heads=2,
+            head_q_dim=4,
+            head_k_dim=4,
+            head_v_dim=4,
+        )
+        original_out_cache_loc = torch.arange(5)
+        forward_batch = SimpleNamespace(
+            forward_mode=_ExtendMode(),
+            num_token_non_padded_cpu=3,
+            out_cache_loc=original_out_cache_loc,
+        )
+
+        with (
+            patch.object(
+                radix_linear_attention,
+                "get_tc_piecewise_forward_context",
+                return_value=None,
+            ),
+            patch.object(
+                radix_linear_attention,
+                "get_attn_backend",
+                return_value=_FakeAttentionBackend(),
+            ),
+        ):
+            output = layer.forward(
+                forward_batch=forward_batch,
+                mixed_qkv=torch.zeros((5, 8)),
+                a=torch.zeros((5, 2)),
+                b=torch.zeros((5, 2)),
+            )
+
+        torch.testing.assert_close(output[:, :3], torch.full((1, 3, 2, 4), 5.0))
+        torch.testing.assert_close(output[:, 3:], torch.zeros((1, 2, 2, 4)))
+        self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
+
+    def test_target_verify_keeps_physical_rows_matching_its_metadata(self):
+        layer = radix_linear_attention.RadixLinearAttention(
+            layer_id=0,
+            num_q_heads=1,
+            num_k_heads=1,
+            num_v_heads=2,
+            head_q_dim=4,
+            head_k_dim=4,
+            head_v_dim=4,
+        )
+        original_out_cache_loc = torch.arange(5)
+        forward_batch = SimpleNamespace(
+            forward_mode=_TargetVerifyMode(),
+            num_token_non_padded_cpu=3,
+            out_cache_loc=original_out_cache_loc,
+        )
+
+        with (
+            patch.object(
+                radix_linear_attention,
+                "get_tc_piecewise_forward_context",
+                return_value=None,
+            ),
+            patch.object(
+                radix_linear_attention,
+                "get_attn_backend",
+                return_value=_PhysicalAttentionBackend(),
+            ),
+        ):
+            output = layer.forward(
+                forward_batch=forward_batch,
+                mixed_qkv=torch.zeros((5, 8)),
+                a=torch.zeros((5, 2)),
+                b=torch.zeros((5, 2)),
+            )
+
+        torch.testing.assert_close(output, torch.full((1, 5, 2, 4), 9.0))
+        self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
+
+    def test_eager_backend_failure_restores_out_cache_loc(self):
+        layer = radix_linear_attention.RadixLinearAttention(
+            layer_id=0,
+            num_q_heads=1,
+            num_k_heads=1,
+            num_v_heads=2,
+            head_q_dim=4,
+            head_k_dim=4,
+            head_v_dim=4,
+        )
+        original_out_cache_loc = torch.arange(5)
+        forward_batch = SimpleNamespace(
+            forward_mode=_ExtendMode(),
+            num_token_non_padded_cpu=3,
+            out_cache_loc=original_out_cache_loc,
+        )
+
+        with (
+            patch.object(
+                radix_linear_attention,
+                "get_tc_piecewise_forward_context",
+                return_value=None,
+            ),
+            patch.object(
+                radix_linear_attention,
+                "get_attn_backend",
+                return_value=_FailingAttentionBackend(),
+            ),
+            self.assertRaisesRegex(RuntimeError, "backend failure"),
+        ):
+            layer.forward(
+                forward_batch=forward_batch,
+                mixed_qkv=torch.zeros((5, 8)),
+                a=torch.zeros((5, 2)),
+                b=torch.zeros((5, 2)),
+            )
+
+        self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
+
+    def test_padded_output_tail_is_initialized(self):
+        for padded_num_tokens in (3, 5):
+            with self.subTest(padded_num_tokens=padded_num_tokens):
+                original_out_cache_loc = torch.arange(padded_num_tokens)
+                forward_batch = SimpleNamespace(
+                    num_token_non_padded_cpu=3,
+                    out_cache_loc=original_out_cache_loc,
+                )
+                context = SimpleNamespace(
+                    forward_batch=forward_batch,
+                    attention_layers=[object()],
+                )
+                output = torch.full((1, padded_num_tokens, 2, 4), float("nan"))
+
+                with (
+                    patch.object(
+                        radix_linear_attention,
+                        "get_tc_piecewise_forward_context",
+                        return_value=context,
+                    ),
+                    patch.object(
+                        radix_linear_attention,
+                        "get_attn_backend",
+                        return_value=_FakeAttentionBackend(),
+                    ),
+                ):
+                    radix_linear_attention._unified_linear_attention_with_output_impl(
+                        mixed_qkv=torch.zeros((padded_num_tokens, 8)),
+                        a=torch.zeros((padded_num_tokens, 2)),
+                        b=torch.zeros((padded_num_tokens, 2)),
+                        output=output,
+                        layer_id=0,
+                    )
+
+                torch.testing.assert_close(output[:, :3], torch.full((1, 3, 2, 4), 5.0))
+                torch.testing.assert_close(
+                    output[:, 3:],
+                    torch.zeros((1, padded_num_tokens - 3, 2, 4)),
+                )
+                self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
+
+
+if __name__ == "__main__":
+    import unittest
+
+    unittest.main()

test/registered/unit/managers/test_scheduler_decision_batch_params.py

@@ -1,5 +1,7 @@
 import inspect
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import maybe_stub_sgl_kernel
@@ -48,5 +50,39 @@ class TestDecisionMethodsHaveNoHiddenBatchChannel(unittest.TestCase):
                 )
 
 
+class TestMtpPhaseBoundaryOverlap(unittest.TestCase):
+    @staticmethod
+    def _batch(*, is_extend: bool, is_speculative: bool = True):
+        return SimpleNamespace(
+            is_extend_in_batch=is_extend,
+            forward_mode=SimpleNamespace(
+                is_extend=lambda: is_extend,
+                is_decode=lambda: not is_extend,
+            ),
+            spec_algorithm=SimpleNamespace(is_none=lambda: not is_speculative),
+            grammar_needs_sync=lambda: False,
+        )
+
+    def _scheduler(self, *, require_mlp_sync: bool):
+        scheduler = object.__new__(Scheduler)
+        scheduler.require_mlp_sync = require_mlp_sync
+        scheduler.result_queue = [object()]
+        return scheduler
+
+    @patch(
+        "sglang.srt.managers.scheduler.envs."
+        "SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get",
+        return_value=False,
+    )
+    def test_mtp_phase_crossing_keeps_overlap(self, _disable_consecutive_prefill):
+        extend = self._batch(is_extend=True)
+        decode = self._batch(is_extend=False)
+
+        for require_mlp_sync in (False, True):
+            scheduler = self._scheduler(require_mlp_sync=require_mlp_sync)
+            self.assertFalse(scheduler.is_disable_overlap_for_batch(decode, extend))
+            self.assertFalse(scheduler.is_disable_overlap_for_batch(extend, decode))
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_replayssm_ring_accounting.py

@@ -19,6 +19,7 @@ from sglang.srt.configs.mamba_utils import (
     Mamba2StateDType,
     Mamba2StateShape,
 )
+from sglang.srt.mem_cache.kv_cache_configurator import _pp_local_per_request_bytes
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -81,6 +82,19 @@ class TestReplaySSMRingAccounting(CustomTestCase):
         self.assertEqual(_gdn_params().replayssm_ring_bytes_per_req(record_len=0), 0)
         self.assertEqual(_kda_params().replayssm_ring_bytes_per_req(record_len=0), 0)
 
+    def test_pp_local_state_budget(self):
+        # Four equal-cost linear layers globally, two owned by this PP stage.
+        self.assertEqual(
+            _pp_local_per_request_bytes(4096, [0, 1, 3, 4], 1, 4),
+            2048,
+        )
+
+    def test_pp_local_state_budget_empty_stage(self):
+        self.assertEqual(
+            _pp_local_per_request_bytes(4096, [0, 1, 3, 4], 5, 8),
+            0,
+        )
+
 
 if __name__ == "__main__":
     import sys

test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py

@@ -7,12 +7,47 @@ from sglang.srt.model_executor.model_runner_components import cuda_graph_setup
 from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import (
     _align_pipeline_layers,
     capture_decode_graph,
+    has_standard_gqa_for_all_local_layers,
+    index_attention_layers_by_global_id,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
+def test_standard_gqa_gate_uses_pipeline_local_layer_range():
+    # PP rank owns layers [23, 46), while the full model has 92 layers.
+    assert has_standard_gqa_for_all_local_layers(
+        attention_layer_count=23, start_layer=23, end_layer=46
+    )
+    assert not has_standard_gqa_for_all_local_layers(
+        attention_layer_count=22, start_layer=23, end_layer=46
+    )
+
+
+def test_standard_gqa_gate_is_unchanged_without_pipeline_parallelism():
+    assert has_standard_gqa_for_all_local_layers(
+        attention_layer_count=92, start_layer=0, end_layer=92
+    )
+
+
+def test_pipeline_attention_metadata_is_indexed_by_global_layer_id():
+    layer23 = SimpleNamespace(layer_id=23)
+    layer24 = SimpleNamespace(layer_id=24)
+    companion24 = object()
+
+    attention, companions = index_attention_layers_by_global_id(
+        [layer23, layer24], [None, companion24]
+    )
+
+    assert len(attention) == 25
+    assert all(layer is None for layer in attention[:23])
+    assert attention[23] is layer23
+    assert attention[24] is layer24
+    assert companions[23] is None
+    assert companions[24] is companion24
+
+
 def test_model_runner_can_override_decode_graph_runner(monkeypatch):
     from sglang.srt.runtime_context import get_context
 

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -1195,6 +1195,76 @@ class TestBuildPrefillRegistry(unittest.TestCase):
         reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=3, padded_num_tokens=8)
         self.assertTrue(torch.equal(idx, torch.tensor([3, 4], dtype=torch.int64)))
 
+    def test_pp_proxy_token_slots_copy_head_and_zero_bucket_tail(self):
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_prefill_registry,
+        )
+        from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+
+        hidden = torch.full((16, 4), 7.0)
+        residual = torch.full((16, 4), 7.0)
+        src = self._src(
+            pp_proxy_tensors={
+                "hidden_states": hidden,
+                "residual": residual,
+            }
+        )
+        reg = build_prefill_registry(
+            device=torch.device("cpu"),
+            max_bs=1,
+            max_num_token=16,
+            cache_loc_dtype=torch.int64,
+            source=src,
+        )
+        self.assertTrue(reg.has_slot("pp_proxy_tensors.hidden_states"))
+        fb = _MiniForwardBatch(
+            input_ids=torch.zeros(3, dtype=torch.int64),
+            positions=torch.zeros(3, dtype=torch.int64),
+            out_cache_loc=torch.zeros(3, dtype=torch.int64),
+        )
+        pp_proxy = PPProxyTensors(
+            {
+                "hidden_states": torch.ones((3, 4)),
+                "residual": torch.full((3, 4), 2.0),
+            }
+        )
+        reg.fill_from(
+            fb,
+            raw_bs=1,
+            padded_bs=1,
+            raw_num_tokens=3,
+            padded_num_tokens=8,
+            pp_proxy_tensors=pp_proxy,
+        )
+        self.assertTrue(torch.all(hidden[:3] == 1.0))
+        self.assertTrue(torch.all(residual[:3] == 2.0))
+        self.assertTrue(torch.all(hidden[3:8] == 0.0))
+        self.assertTrue(torch.all(residual[3:8] == 0.0))
+        self.assertTrue(torch.all(hidden[8:] == 7.0))
+
+    def test_prefill_input_buffers_allocate_pp_proxy_by_token(self):
+        from sglang.srt.model_executor.runner_utils.buffers import (
+            PrefillInputBuffers,
+        )
+
+        buffers = PrefillInputBuffers.create(
+            device=torch.device("cpu"),
+            max_bs=4,
+            max_num_tokens=16,
+            cache_loc_dtype=torch.int64,
+            is_multimodal=False,
+            hidden_size=8,
+            dtype=torch.bfloat16,
+            enable_mamba_track=False,
+            pp_size=2,
+            pp_proxy_topk_size=3,
+        )
+        self.assertEqual(
+            tuple(buffers.pp_proxy_tensors["hidden_states"].shape), (16, 8)
+        )
+        self.assertEqual(tuple(buffers.pp_proxy_tensors["residual"].shape), (16, 8))
+        self.assertEqual(tuple(buffers.pp_proxy_tensors["topk_indices"].shape), (16, 3))
+
     def test_source_none_owns_allocated_buffers(self):
         # source=None -> the registry allocates (owns) every slot.
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -9,7 +9,12 @@ import torch
 import sglang.srt.model_executor.model_runner_components.cuda_graph_setup as graph_setup
 import sglang.srt.model_executor.runner.prefill_cuda_graph_runner as runner_module
 from sglang.srt.model_executor.cuda_graph_config import Backend
-from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+    PPProxyTensors,
+)
 from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import (
     capture_prefill_graph,
 )
@@ -61,6 +66,32 @@ class _FakeKVIndexKernel:
         return run
 
 
+class _FakeGraphSlot:
+    def __init__(self, buffer):
+        self.buffer = buffer
+
+    def slice_for(self, _batch_size, num_tokens):
+        return self.buffer[:num_tokens]
+
+
+class _FakeBatchRegistry:
+    def __init__(self):
+        self.slots = {
+            "input_ids": _FakeGraphSlot(torch.arange(4, dtype=torch.int64)),
+            "positions": _FakeGraphSlot(torch.arange(4, dtype=torch.int64)),
+            "out_cache_loc": _FakeGraphSlot(torch.arange(4, dtype=torch.int64)),
+        }
+
+    def fill_from(self, *_args, **_kwargs):
+        return None
+
+    def has_slot(self, name):
+        return name in self.slots
+
+    def get_slot(self, name):
+        return self.slots[name]
+
+
 class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
     def test_low_free_memory_still_captures_prefill_graph(self):
         eager_runner = object()
@@ -86,6 +117,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             server_args=SimpleNamespace(),
             model=SimpleNamespace(),
             model_config=SimpleNamespace(context_len=8192, num_hidden_layers=1),
+            layer_info=SimpleNamespace(start_layer=0, end_layer=1),
             req_to_token_pool=SimpleNamespace(size=1),
         )
         language_model = SimpleNamespace(layers=[object()])
@@ -98,7 +130,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             patch.object(
                 graph_setup,
                 "compute_attention_and_moe_layers",
-                return_value=([object()], [], [], [], []),
+                return_value=([object()], [], [], [], [None]),
             ),
             patch.object(
                 graph_setup,
@@ -145,6 +177,63 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
 
         self.assertIs(capture.runner, eager_runner)
 
+    def test_pp_proxy_output_is_trimmed_to_raw_prefill_tokens(self):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.raw_num_tokens = 3
+        output = PPProxyTensors(
+            {
+                "hidden_states": torch.arange(32).view(8, 4),
+                "residual": torch.arange(32, 64).view(8, 4),
+            }
+        )
+
+        trimmed = runner._finalize_execute_output(output)
+
+        self.assertIsInstance(trimmed, PPProxyTensors)
+        self.assertEqual(tuple(trimmed["hidden_states"].shape), (3, 4))
+        self.assertEqual(tuple(trimmed["residual"].shape), (3, 4))
+
+    def test_static_batch_preserves_consumed_multimodal_embeddings(self):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner.capture_num_tokens = [4]
+        runner.buffer_registry = _FakeBatchRegistry()
+        runner.enable_cp_v2_bcg_capture = False
+        runner._is_full_backend = False
+        runner.backend = SimpleNamespace()
+        runner.has_mha_companion_layers = False
+        runner._prefill_static_buffers = None
+        runner.static_draft_hidden_states = None
+        runner.capture_return_pooled_hidden_states = False
+        runner._next_token_logits_buffer = lambda _rows: None
+        runner._prefill_logits_buffer_rows = lambda _batch: 1
+        runner._prepare_forward_metadata_for_replay = lambda *_args: None
+
+        mm_input_embeds = torch.randn(3, 8)
+        forward_batch = ForwardBatch(
+            forward_mode=ForwardMode.EXTEND,
+            batch_size=1,
+            input_ids=torch.arange(3, dtype=torch.int64),
+            req_pool_indices=torch.zeros(1, dtype=torch.int64),
+            seq_lens=torch.tensor([3], dtype=torch.int32),
+            out_cache_loc=torch.arange(3, dtype=torch.int64),
+            seq_lens_sum=3,
+            positions=torch.arange(3, dtype=torch.int64),
+            seq_lens_cpu=torch.tensor([3], dtype=torch.int32),
+            extend_seq_lens=torch.tensor([3], dtype=torch.int32),
+            extend_prefix_lens=torch.zeros(1, dtype=torch.int32),
+            extend_start_loc=torch.zeros(1, dtype=torch.int32),
+            extend_seq_lens_cpu=[3],
+            extend_prefix_lens_cpu=[0],
+            mm_inputs=None,
+            mm_input_embeds=mm_input_embeds,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+            global_forward_mode=ForwardMode.EXTEND,
+        )
+
+        static_batch = runner.load_batch(forward_batch)
+
+        self.assertIs(static_batch.mm_input_embeds, mm_input_embeds)
+
     def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self):
         graph_config = SimpleNamespace(
             prefill=SimpleNamespace(full_prefill_prefix_chunk_tokens=None, max_bs=8)

test/registered/unit/test_runtime_context.py

@@ -581,9 +581,11 @@ class TestMoeFlagsGroup(_IsolatedServerArgs):
             self.assertTrue(get_moe_a2a_backend().is_none())
             # MTP layers are unquantized: fp4 allgather is forced off
             self.assertTrue(get_flags().moe.disable_fp4_allgather)
+            self.assertTrue(get_flags().moe.speculative_context)
         self.assertEqual(get_moe_runner_backend().name, "TRITON")
         self.assertTrue(get_moe_a2a_backend().is_deepep())
         self.assertFalse(get_flags().moe.disable_fp4_allgather)
+        self.assertFalse(get_flags().moe.speculative_context)
 
     def test_swap_restores_on_exception(self):
         from sglang.srt.layers.moe.utils import (