dsv4.1: remaining model and runtime integration (#38798)

Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Xiaoyu Zhang <xiaoyu.zhang@radixark.ai>
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Zhichen Zeng <zczeng@uw.edu>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -126,6 +126,26 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase):
             extend_lens=(16,),
             compress_ratio=128,
         ),
+        DSV4AttentionCase(
+            name="dsv4_c2_extend",
+            backend="dsv4",
+            forward_mode=ForwardMode.EXTEND,
+            num_heads=64,
+            page_size=DSV4_PAGE_SIZE,
+            # Odd lengths: the ratio-2 causal count (pos + 1) // 2 rounds down.
+            prefix_lens=(33,),
+            extend_lens=(7,),
+            compress_ratio=2,
+        ),
+        DSV4AttentionCase(
+            name="dsv4_c2_decode",
+            backend="dsv4",
+            forward_mode=ForwardMode.DECODE,
+            num_heads=64,
+            page_size=DSV4_PAGE_SIZE,
+            prefix_lens=(65,),
+            compress_ratio=2,
+        ),
         DSV4AttentionCase(
             name="dsv4_c128_decode",
             backend="dsv4",
@@ -533,6 +553,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
                 backend.model_runner = SimpleNamespace(
                     spec_algorithm=SpeculativeAlgorithm.DFLASH
                 )
+                backend.token_to_kv_pool = SimpleNamespace(request_window=None)
                 backend.forward_metadata = DSV4Metadata(
                     self._make_core_metadata(0), indexer_metadata=None
                 )
@@ -577,6 +598,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
         backend.model_runner = SimpleNamespace(
             spec_algorithm=SpeculativeAlgorithm.DFLASH
         )
+        backend.token_to_kv_pool = SimpleNamespace(request_window=None)
         backend.forward_metadata = DSV4Metadata(
             self._make_core_metadata(0), indexer_metadata=None
         )
@@ -794,7 +816,8 @@ class TestDSV4SwaOutCacheLocResolution(CustomTestCase):
         backend = object.__new__(DeepseekV4AttnBackend)
         backend.forward_metadata = None
         backend.token_to_kv_pool = SimpleNamespace(
-            translate_loc_from_full_to_swa=lambda loc: mapping[loc]
+            translate_loc_from_full_to_swa=lambda loc: mapping[loc],
+            request_window=None,
         )
         return backend
 

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

@@ -175,7 +175,9 @@ def _make_backend(
     dsv4_prefill_backend: str = "auto",
 ) -> DeepseekV4AttnBackend:
     backend = DeepseekV4AttnBackend.__new__(DeepseekV4AttnBackend)
-    backend.forward_metadata = SimpleNamespace(sparse_prefill_cache=None)
+    backend.forward_metadata = SimpleNamespace(
+        sparse_prefill_cache=None, late_layer_tail=None
+    )
     backend.req_to_token = req_to_token
     backend.sparse_prefill_workspace = SparsePrefillWorkspace(device)
     backend.softmax_scale = 512**-0.5

test/registered/spec/dspark/test_dspark_draft_path_default.py

@@ -1,6 +1,8 @@
 import unittest
 from types import SimpleNamespace
 
+import torch
+
 from sglang.srt.arg_groups.overrides import resolution_result
 from sglang.srt.arg_groups.speculative_hook import (
     _handle_dspark,
@@ -120,5 +122,37 @@ class TestDsparkDpAttentionMoeA2aGate(CustomTestCase):
                 _handle_dspark(server_args)
 
 
+class TestDsparkFoldedSamplingDefault(CustomTestCase):
+    def test_sharded_greedy_default_and_sampling_override(self):
+        from sglang.srt.environ import DsparkFoldedSampling, envs
+        from sglang.srt.speculative.dspark_components.dspark_draft_sampler import (
+            _resolve_folded_sampling,
+        )
+
+        model = SimpleNamespace(
+            lm_head=SimpleNamespace(org_vocab_size=128, weight=torch.empty(1)),
+            markov_head=SimpleNamespace(supports_sharded_greedy=True),
+        )
+        args = dict(
+            model=model,
+            gamma=5,
+            max_bs=64,
+            device="cpu",
+            tp_rank=0,
+            available_memory_gb=16,
+        )
+        with envs.SGLANG_DSPARK_FOLDED_SAMPLING.override(
+            DsparkFoldedSampling.AUTO.value
+        ):
+            self.assertFalse(_resolve_folded_sampling(**args))
+            model.markov_head.supports_sharded_greedy = False
+            self.assertTrue(_resolve_folded_sampling(**args))
+        model.markov_head.supports_sharded_greedy = True
+        with envs.SGLANG_DSPARK_FOLDED_SAMPLING.override(
+            DsparkFoldedSampling.FORCE.value
+        ):
+            self.assertTrue(_resolve_folded_sampling(**args))
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py

@@ -86,6 +86,7 @@ def _prefill_queue(sched):
 def _decode_queue(sched):
     q = SimpleNamespace(
         scheduler=sched,
+        token_to_kv_pool_allocator=MagicMock(),
         retracted_queue=[],
         pending_reqs=[],
         _check_if_req_exceed_kv_capacity=MagicMock(return_value=False),

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -906,7 +906,7 @@ def _buf_infos(*ptrs):
 
 def _make_dsv4_target(*, unified, mapping=None):
     pool = object.__new__(DeepSeekV4TokenToKVPool)
-    pool.compression_ratios = [0, 4, 128]
+    pool.compression_ratios = [0, 2, 1, 4, 128]
     pool._unified_kv = unified
     pool.page_size = 256
     pool.sliding_window = 128

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

@@ -7,6 +7,7 @@ SM90 / SM100 / SM120.
 """
 
 import unittest
+from types import SimpleNamespace
 from unittest import mock
 
 import torch
@@ -267,24 +268,26 @@ class TestMxfp8LinearBackends(_LinearBackendCheck):
             is_backend_supported.assert_called_once_with("cute-dsl", 107)
 
 
+def _build_block32_layer(n: int, k: int, keep_plain_weight_layout: bool = False):
+    quant_config = Fp8Config(
+        is_checkpoint_fp8_serialized=True,
+        activation_scheme="dynamic",
+        weight_block_size=[32, 32],
+        scale_fmt="ue8m0",
+    )
+    layer = _make_linear(quant_config, n, k)
+    if keep_plain_weight_layout:
+        layer.keep_plain_weight_layout = True
+    w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10
+    w_fp8, scale_e8m0, w_dequant = _quantize_fp8_block32_ue8m0(w)
+    load_linear_weights(layer, weight=w_fp8, weight_scale_inv=scale_e8m0)
+    return layer, w_dequant
+
+
 class TestBlockFp8AsMxfp8Linear(_LinearBackendCheck):
     """A 32-wide-K ue8m0 block-fp8 weight served through the MXFP8 GEMMs."""
 
-    @staticmethod
-    def _build_layer(n: int, k: int, keep_plain_weight_layout: bool = False):
-        quant_config = Fp8Config(
-            is_checkpoint_fp8_serialized=True,
-            activation_scheme="dynamic",
-            weight_block_size=[32, 32],
-            scale_fmt="ue8m0",
-        )
-        layer = _make_linear(quant_config, n, k)
-        if keep_plain_weight_layout:
-            layer.keep_plain_weight_layout = True
-        w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10
-        w_fp8, scale_e8m0, w_dequant = _quantize_fp8_block32_ue8m0(w)
-        load_linear_weights(layer, weight=w_fp8, weight_scale_inv=scale_e8m0)
-        return layer, w_dequant
+    _build_layer = staticmethod(_build_block32_layer)
 
     def _run(self, backend: str):
         self._check_backend(
@@ -339,6 +342,120 @@ class TestBlockFp8AsMxfp8Linear(_LinearBackendCheck):
                 plain_layer.quant_method.apply(plain_layer, Mxfp8SwizzledInput(q, s))
 
 
+@unittest.skipUnless(
+    "flashinfer_cutedsl" in _block32_backends(),
+    "block-fp8-as-MXFP8 prefill tuning needs the FlashInfer CuTe-DSL kernel",
+)
+class TestBlockFp8AsMxfp8PrefillAutotune(_LinearBackendCheck):
+    """The startup hook that tunes those layers for the prefill M buckets."""
+
+    def setUp(self):
+        super().setUp()
+        patcher = mock.patch.object(
+            fp8_utils,
+            "FP8_GEMM_RUNNER_BACKEND",
+            Fp8GemmRunnerBackend.FLASHINFER_CUTEDSL,
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        torch.manual_seed(7)
+
+    @staticmethod
+    def _ready_layer(n: int, k: int, keep_plain_weight_layout: bool = False):
+        layer, _ = _build_block32_layer(n, k, keep_plain_weight_layout)
+        layer.quant_method.process_weights_after_loading(layer)
+        return layer
+
+    def test_model_hook_deduplicates_ready_block_fp8_weights(self):
+        from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
+
+        layers = torch.nn.ModuleList()
+        methods = []
+        for _ in range(2):
+            layer = self._ready_layer(128, 128)
+            methods.append(layer.quant_method)
+            layer.quant_method.apply = mock.Mock()
+            layers.append(layer)
+        # An unprepared layer intentionally has no swizzled scale buffer.
+        fallback = self._ready_layer(128, 128, keep_plain_weight_layout=True)
+        layers.append(fallback)
+        model = SimpleNamespace(
+            config=SimpleNamespace(model_type="deepseek_v41"), model=layers
+        )
+        count = DeepseekV4ForCausalLM.autotune_prefill_kernels(
+            model, 4096, dtype=torch.bfloat16
+        )
+        self.assertEqual(count, 1)
+        methods[0].apply.assert_called_once()
+        self.assertEqual(methods[0].apply.call_args.args[1].shape, (4096, 128))
+        methods[1].apply.assert_not_called()
+        for method in methods:
+            self.assertEqual(method.mxfp8_prefill_autotune_min_tokens, 4096)
+        self.assertIsNone(fallback.quant_method.mxfp8_prefill_autotune_min_tokens)
+
+    def test_block_fp8_dispatch_keeps_decode_and_determinism_pinned(self):
+        layer = self._ready_layer(128, 128)
+        method = layer.quant_method
+        method.mxfp8_prefill_autotune_min_tokens = 4096
+        call = mock.Mock(return_value=torch.empty(0))
+        method.w8a8_mxfp8_linear = call
+        for rows, invariant, deterministic, expected in (
+            (6, False, False, None),
+            (4096, False, False, False),
+            (4096, True, False, True),
+            (4096, False, True, True),
+        ):
+            with self.subTest(
+                rows=rows, invariant=invariant, deterministic=deterministic
+            ):
+                with (
+                    mock.patch(
+                        "sglang.srt.batch_invariant_ops.is_batch_invariant_mode_enabled",
+                        return_value=invariant,
+                    ),
+                    mock.patch(
+                        "sglang.srt.runtime_context.get_exec",
+                        return_value=SimpleNamespace(
+                            deterministic=SimpleNamespace(
+                                enable_deterministic_inference=deterministic
+                            )
+                        ),
+                    ),
+                ):
+                    method.apply(layer, torch.empty(rows, 128, device="cuda"))
+                self.assertEqual(call.call_args.kwargs.get("pin_tactic"), expected)
+
+    def test_prefill_tuning_leaves_decode_bit_identical(self):
+        """Tuning the prefill buckets must not move the decode tactic: below the
+        stamped min_tokens the output has to stay bit-for-bit what it was."""
+        from flashinfer.autotuner import autotune
+
+        from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
+
+        runtime_patch = mock.patch(
+            "sglang.srt.runtime_context.get_exec",
+            return_value=SimpleNamespace(
+                deterministic=SimpleNamespace(enable_deterministic_inference=False)
+            ),
+        )
+        runtime_patch.start()
+        self.addCleanup(runtime_patch.stop)
+        layer = self._ready_layer(1792, 5120)
+        method = layer.quant_method
+        x = torch.randn(6, 5120, device="cuda", dtype=torch.bfloat16)
+        original = method.apply(layer, x)
+        model = SimpleNamespace(
+            config=SimpleNamespace(model_type="deepseek_v41"),
+            model=torch.nn.ModuleList([layer]),
+        )
+        with autotune(True):
+            DeepseekV4ForCausalLM.autotune_prefill_kernels(
+                model, 4096, dtype=torch.bfloat16
+            )
+        self.assertEqual(method.mxfp8_prefill_autotune_min_tokens, 4096)
+        torch.testing.assert_close(method.apply(layer, x), original, rtol=0, atol=0)
+
+
 @unittest.skipIf(get_device_sm() < 90, "FP8 GEMM backends require SM90+")
 class TestModeloptFp8PerTensorLinear(_LinearBackendCheck):
     """Per-tensor FP8 (ModelOptFp8LinearMethod, static scales) on the auto

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

@@ -0,0 +1,114 @@
+"""A TP-sharded MXFP4 trtllm-gen MoE whose per-rank intermediate size needs
+padding must sum to the unsharded experts' output."""
+
+import unittest
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import torch
+
+from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+from sglang.srt.layers.moe.topk import StandardTopKOutput
+from sglang.srt.layers.quantization import mxfp4_flashinfer_trtllm_moe as mxfp4
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+def make_layer(weights):
+    layer = torch.nn.Module()
+    names = (
+        "w13_weight",
+        "w2_weight",
+        "w13_weight_scale_inv",
+        "w2_weight_scale_inv",
+    )
+    for name, tensor in zip(names, weights):
+        layer.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False))
+    layer.num_experts = weights[0].shape[0]
+    layer.num_local_experts = layer.num_experts
+    layer.moe_ep_rank = 0
+    return layer
+
+
+def make_weights(intermediate, hidden=256, device="cpu"):
+    experts = 8
+
+    def fp4_packed(*shape):
+        return torch.randint(-128, 128, shape, dtype=torch.int8, device=device)
+
+    def e8m0_scales(*shape):
+        return torch.randint(-6, -3, shape, device=device).float().exp2()
+
+    return (
+        fp4_packed(experts, 2 * intermediate, hidden // 2),
+        fp4_packed(experts, hidden, intermediate // 2),
+        e8m0_scales(experts, 2 * intermediate, hidden // 32),
+        e8m0_scales(experts, hidden, intermediate // 32),
+    )
+
+
+class TestMxfp4TrtllmPadding(CustomTestCase):
+    @unittest.skipUnless(
+        torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10,
+        "Requires Blackwell",
+    )
+    def test_tp4_matches_unsharded_experts(self):
+        torch.manual_seed(42)
+        weights = make_weights(2304, hidden=5120, device="cuda")
+
+        def prepare(tensors):
+            layer = make_layer(tensors)
+            method = object.__new__(mxfp4.Mxfp4FlashinferTrtllmMoEMethod)
+            method._fp8 = Mock()
+            method.prefix = "test.experts"
+            method.flashinfer_mxfp4_moe_precision = "default"
+            method.process_weights_after_loading(layer)
+            method.create_moe_runner(layer, SimpleNamespace(swiglu_limit=10.0))
+            return method, layer
+
+        full = prepare([tensor.clone() for tensor in weights])
+        shards = []
+        for rank in range(4):
+            start = rank * 576
+            end = start + 576
+            w13, w2, s13, s2 = weights
+            shard = (
+                torch.cat(
+                    (w13[:, start:end], w13[:, 2304 + start : 2304 + end]), dim=1
+                ),
+                w2[..., start // 2 : end // 2].contiguous(),
+                torch.cat(
+                    (s13[:, start:end], s13[:, 2304 + start : 2304 + end]), dim=1
+                ),
+                s2[..., start // 32 : end // 32].contiguous(),
+            )
+            shards.append(prepare(shard))
+
+        with (
+            patch.object(mxfp4, "get_tp_group", return_value=None),
+            patch.object(mxfp4, "is_allocation_symmetric", return_value=False),
+            patch.object(mxfp4, "use_symmetric_memory", return_value=nullcontext()),
+        ):
+            for tokens in (1, 64):
+                with self.subTest(tokens=tokens):
+                    x = torch.randn(tokens, 5120, dtype=torch.bfloat16, device="cuda")
+                    logits = torch.randn(tokens, 8, device="cuda")
+                    scores, ids = logits.softmax(-1).topk(6, dim=-1)
+                    topk = StandardTopKOutput(scores, ids.to(torch.int32), logits)
+                    dispatch = StandardDispatchOutput(x, None, topk)
+                    reference = full[0].apply(full[1], dispatch).hidden_states.float()
+                    actual = sum(
+                        method.apply(layer, dispatch).hidden_states.float()
+                        for method, layer in shards
+                    )
+                    rmse = torch.linalg.norm(actual - reference) / torch.linalg.norm(
+                        reference
+                    )
+                    self.assertLess(rmse.item(), 0.01)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_dsv4_nonpaged_indexer.py

@@ -531,5 +531,26 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
         self.assertEqual(call.kwargs, {"clean_logits": False, "max_seqlen_k": 128})
 
 
+class TestCandidateIndexerGating(CustomTestCase):
+    def test_candidate_indexer_gating(self):
+        from sglang.srt.layers.attention.dsv4 import candidate_indexer
+
+        def platform(sm):
+            return patch.object(
+                candidate_indexer, "get_platform", lambda: SimpleNamespace(device_sm=sm)
+            )
+
+        flag = "sglang.srt.layers.deep_gemm_wrapper.configurer.DEEPGEMM_PAGED_SPARSE_MQA_LOGITS"
+        # V4 models have no candidate source; Hopper selects through masks inline.
+        with platform(100), patch(flag, True):
+            self.assertIsNone(candidate_indexer.make_candidate_indexer(0, 8))
+        with platform(90), patch(flag, False):
+            self.assertIsNone(candidate_indexer.make_candidate_indexer(2048, 8))
+        # Blackwell without DeepGEMM's sparse logits fails instead of falling back.
+        with platform(100), patch(flag, False):
+            with self.assertRaises(RuntimeError):
+                candidate_indexer.make_candidate_indexer(2048, 8)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/managers/test_priority_scheduling_disaggregation.py

@@ -529,6 +529,8 @@ class TestDecodePrebuilt(unittest.TestCase):
         scheduler.policy = MagicMock()
         scheduler.schedule_stream = MagicMock()
         scheduler.forward_stream = MagicMock()
+        scheduler.ngram_embedding_manager = MagicMock()
+        scheduler.chunked_req = None
         return scheduler
 
     def test_waiting_queue_is_sorted_before_prebuilt_selection(self):

test/registered/unit/mem_cache/test_dsv4_compressed_pools.py

@@ -5,11 +5,17 @@ from unittest.mock import MagicMock, patch
 
 import torch
 
+from sglang.kernels.ops.attention.dsv4.kv_layout import (
+    KVLayout,
+    is_valid_kv_layout_pair,
+)
+from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
     DeepSeekV4SingleKVPool,
     DeepSeekV4TokenToKVPool,
     _CompressedPoolConfig,
 )
+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
 
@@ -25,6 +31,8 @@ class TestDSV4CompressedPools(CustomTestCase):
                 pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
                 pool._unified_kv = unified
                 pool.uniform_fp8 = False
+                pool.kv_layout = KVLayout.V4
+                pool.compressed_kv_layout_option = None
                 pool.compressed_pool_configs = {
                     4: _CompressedPoolConfig(
                         256, 64, torch.bfloat16, indexer_size=1024
@@ -176,5 +184,151 @@ class TestDSV4CompressedPools(CustomTestCase):
             pool.get_index_k_page_size(128)
 
 
+HEAD_DIM = 512
+ROPE_DIM = 64
+PAGE_SIZE = 256
+FULL_SIZE = 4 * PAGE_SIZE
+
+
+class TestV41KVPoolLayouts(CustomTestCase):
+    """A V4.1-layout pool hands the attention kernel page-aligned buffers and
+    picks the compressed layout each ratio asks for."""
+
+    def setUp(self):
+        super().setUp()
+        override = get_context().override_server_args(page_size=PAGE_SIZE)
+        override.install()
+        self.addCleanup(override.restore)
+
+    def make_pool(self, ratios, kv_source_layers, kv_layout, compressed=None, **sizes):
+        return DeepSeekV4TokenToKVPool(
+            max_num_reqs=16,
+            swa_size=FULL_SIZE,
+            c4_size=sizes.get("c4_size", 0),
+            c128_size=sizes.get("c128_size", 0),
+            c4_state_pool_size=sizes.get("c4_state_pool_size", 0),
+            c128_state_pool_size=sizes.get("c128_state_pool_size", 0),
+            page_size=PAGE_SIZE,
+            swa_page_size=PAGE_SIZE,
+            dtype=torch.float8_e4m3fn,
+            c4_state_dtype=torch.float32,
+            c128_state_dtype=torch.float32,
+            qk_nope_head_dim=HEAD_DIM - ROPE_DIM,
+            qk_rope_head_dim=ROPE_DIM,
+            indexer_head_dim=128,
+            layer_num=len(ratios),
+            device="cpu",
+            enable_memory_saver=False,
+            compression_ratios=ratios,
+            kv_source_layers=kv_source_layers,
+            full_size=FULL_SIZE,
+            kv_layout=kv_layout,
+            compressed_kv_layout=compressed,
+        )
+
+    def assert_kernel_requirements(self, pool, layout):
+        """Pages start on the kernel's alignment, and its
+        (num_pages, page_size, 1, bytes_per_token) view walks one token per row."""
+        for buf in pool.kv_buffer:
+            self.assertEqual(buf.stride(0) % layout.page_align, 0)
+            bpt = layout.bytes_per_token
+            view = buf[:, : pool.page_size * bpt].view(
+                buf.shape[0], pool.page_size, 1, bpt
+            )
+            self.assertEqual(view.stride(1), bpt)
+            self.assertEqual(view.stride(0), pool.bytes_per_page_padded)
+
+    def test_v41_pool_buffers(self):
+        for option, expect in ((None, KVLayout.V41_FP4), ("fp8", KVLayout.V41)):
+            with self.subTest(compressed=option):
+                pool = self.make_pool([0, 0, 2, 1, 1], [2, 3], KVLayout.V41, option)
+                self.assert_kernel_requirements(pool.swa_kv_pool, KVLayout.V41)
+                self.assertEqual(pool.get_swa_key_bytes_per_token(), 528)
+                for ratio in (1, 2):
+                    layer_id = pool.sources_by_ratio[ratio][0]
+                    self.assertIs(pool.get_extra_key_layout(layer_id), expect)
+                    self.assertEqual(
+                        pool.get_extra_key_bytes_per_token(layer_id),
+                        expect.bytes_per_token,
+                    )
+                    self.assertTrue(is_valid_kv_layout_pair(pool.kv_layout, expect))
+                    self.assert_kernel_requirements(pool.kv_pools[ratio], expect)
+        # A pool of the fp4 layout cannot be the main cache.
+        with self.assertRaises(AssertionError):
+            self.make_pool([0], [], KVLayout.V41_FP4)
+
+    def test_v41_pool_with_c4_c128(self):
+        pool = self.make_pool(
+            [0, 4, 128],
+            [],
+            KVLayout.V41,
+            c4_size=PAGE_SIZE,
+            c128_size=PAGE_SIZE,
+            c4_state_pool_size=16,
+            c128_state_pool_size=16,
+        )
+        for ratio in (4, 128):
+            self.assertEqual(pool.kv_pools[ratio].page_size, PAGE_SIZE // ratio)
+        # The 2-token c128 page is the only production page that pads.
+        self.assertEqual(pool.kv_pools[128].bytes_per_page_padded, 1536)
+
+
+class TestPagedDSparkWithEncoderReplay(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        override = get_context().override_server_args(
+            enable_encoder_swa_bounded_replay=True,
+            speculative_algorithm="DSPARK",
+            speculative_num_draft_tokens=6,
+            speculative_dspark_block_size=5,
+            page_size=256,
+            max_running_requests=2,
+            chunked_prefill_size=256,
+        )
+        override.install()
+        self.addCleanup(override.restore)
+
+    def make_pool(self, *, draft):
+        return DeepSeekV4TokenToKVPool(
+            max_num_reqs=2,
+            num_req_slots=3,
+            swa_size=1024,
+            c4_size=0,
+            c128_size=0,
+            c4_state_pool_size=0,
+            c128_state_pool_size=0,
+            page_size=256,
+            swa_page_size=256,
+            dtype=torch.float8_e4m3fn,
+            c4_state_dtype=torch.float32,
+            c128_state_dtype=torch.bfloat16,
+            qk_nope_head_dim=448,
+            qk_rope_head_dim=64,
+            indexer_head_dim=128,
+            layer_num=3,
+            device="cpu",
+            enable_memory_saver=False,
+            compression_ratios=[0, 0, 0],
+            online_mtp_max_draft_tokens=6,
+            full_size=2048,
+            is_draft_worker=draft,
+        )
+
+    def test_target_window_and_draft_paged_storage_share_allocator_mapping(self):
+        target = self.make_pool(draft=False)
+        draft = self.make_pool(draft=True)
+        allocator = SWATokenToKVPoolAllocator(
+            2048, 1024, 256, torch.float8_e4m3fn, "cpu", target, False
+        )
+        draft.register_mapping(allocator.full_to_swa_index_mapping)
+        allocator.full_to_swa_index_mapping[256:512] = torch.arange(768, 1024)
+        self.assertEqual(
+            draft.translate_loc_from_full_to_swa(
+                torch.tensor([256, 300, 511])
+            ).tolist(),
+            [768, 812, 1023],
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py

@@ -48,6 +48,7 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase):
         )
 
         kvcache = _mock_kvcache(DeepSeekV4TokenToKVPool)
+        kvcache.swa_kv_pool = MagicMock()
         strategy = _select_strategy(kvcache, {FULL, SWA})
         self.assertIsInstance(strategy, _DeepSeekV4Strategy)
 
@@ -141,6 +142,7 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase):
 
         for cls in (SWAKVPool, DeepSeekV4TokenToKVPool):
             kvcache = _mock_kvcache(cls)
+            kvcache.swa_kv_pool = MagicMock()
             with self.assertRaises(AssertionError) as cm:
                 _select_strategy(kvcache, {FULL})
             self.assertIn("No matching HiCache strategy", str(cm.exception))

test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py

@@ -6,9 +6,10 @@ reduction holds only if ranks also enter tuning with the same cache, so these
 cover that gate and the digest it decides on.
 """
 
-from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
 
-register_cpu_ci(est_time=52, suite="base-a-test-cpu")
+register_cpu_ci(est_time=57, suite="base-a-test-cpu")
+register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 
 import json
 import multiprocessing
@@ -16,11 +17,15 @@ import os
 import tempfile
 import traceback
 import unittest
+from contextlib import nullcontext
 from pathlib import Path
 from types import SimpleNamespace
+from unittest.mock import Mock, patch
 
+import torch
 import torch.distributed as dist
 
+from sglang.srt.model_executor.runner import flashinfer_autotune as autotune
 from sglang.srt.model_executor.runner.flashinfer_autotune import (
     _autotune_cache_digest,
     _autotune_tactic_sync_group,
@@ -160,5 +165,113 @@ class TestDropDivergedAutotuneCache(CustomTestCase):
         )
 
 
+class TestModelPrefillAutotune(CustomTestCase):
+    """Model kernel warmup must cover prefill without a speculative dummy batch."""
+
+    def setUp(self):
+        self.hook = Mock(return_value=1)
+        self.mr = SimpleNamespace(
+            model=SimpleNamespace(autotune_prefill_kernels=self.hook),
+            is_generation=True,
+            is_draft_worker=False,
+            dtype=torch.bfloat16,
+        )
+        self.runner = SimpleNamespace(model_runner=self.mr)
+        # No dummy-buffer or attention APIs: this path must not build a
+        # TARGET_VERIFY batch or mutate request/KV state.
+        for target, kwargs in (
+            ("max_prefill_buffer_tokens", {"return_value": 65536}),
+            (
+                "flashinfer_autotune_context",
+                {"side_effect": lambda *a, **k: nullcontext()},
+            ),
+        ):
+            p = patch.object(autotune, target, **kwargs)
+            setattr(self, target, p.start())
+            self.addCleanup(p.stop)
+        p = patch.object(
+            autotune.envs.SGLANG_FLASHINFER_AUTOTUNE_EXTEND, "get", return_value=False
+        )
+        p.start()
+        self.addCleanup(p.stop)
+
+    def test_declining_model_never_enters_the_autotune_context(self):
+        self.mr.model.wants_prefill_autotune = lambda: False
+        autotune.maybe_flashinfer_autotune_extend(self.runner, decode_num_tokens=384)
+        self.hook.assert_not_called()
+        self.flashinfer_autotune_context.assert_not_called()
+
+    def test_extend_pass_is_opt_in(self):
+        # A draft worker keeps its own warmup; a model without the hook opts out.
+        for draft, has_hook in ((True, True), (False, False)):
+            with self.subTest(draft=draft, has_hook=has_hook):
+                self.mr.is_draft_worker = draft
+                if not has_hook:
+                    del self.mr.model.autotune_prefill_kernels
+                autotune.maybe_flashinfer_autotune_extend(
+                    self.runner, decode_num_tokens=384
+                )
+                self.hook.assert_not_called()
+                self.flashinfer_autotune_context.assert_not_called()
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "FlashInfer requires CUDA")
+class TestAutotuneCachePhases(CustomTestCase):
+    """Loaded target tactics survive draft warmup, unless cache reuse is off."""
+
+    def test_target_and_draft_cache_reuse(self):
+        from flashinfer.autotuner import AutoTuner, _collect_metadata
+
+        tuner = AutoTuner.get()
+        tuner.clear_cache()
+        self.addCleanup(tuner.clear_cache)
+        runner = SimpleNamespace(
+            device="cuda",
+            forward_stream=torch.cuda.Stream(),
+            tp_group=SimpleNamespace(world_size=1),
+        )
+        with tempfile.TemporaryDirectory() as directory:
+            target, draft = (
+                Path(directory) / name for name in ("target.json", "draft.json")
+            )
+            for path, key, tactic in (
+                (target, "target_prefill", 7),
+                (draft, "draft_decode", 3),
+            ):
+                path.write_text(
+                    json.dumps(
+                        {"_metadata": _collect_metadata(), key: ["TestRunner", tactic]}
+                    )
+                )
+            with (
+                patch.object(
+                    autotune,
+                    "flashinfer_autotune_cache_path",
+                    side_effect=[target, draft, draft],
+                ),
+                patch.object(
+                    autotune, "get_flashinfer_autotune_skip_ops", return_value=set()
+                ),
+                autotune.envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.override(True),
+            ):
+                with autotune.flashinfer_autotune_context(runner, run_lm_head=False):
+                    self.assertEqual(
+                        tuner._file_configs["target_prefill"], ("TestRunner", 7)
+                    )
+                # No profiling: this models a restart that loads tactics from disk.
+                self.assertFalse(tuner.profiling_cache)
+                with autotune.flashinfer_autotune_context(runner, run_lm_head=False):
+                    pass
+                saved = json.loads(draft.read_text())
+                self.assertEqual(saved["target_prefill"], ["TestRunner", 7])
+                self.assertEqual(saved["draft_decode"], ["TestRunner", 3])
+                with (
+                    autotune.envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.override(False),
+                    autotune.flashinfer_autotune_context(runner, run_lm_head=False),
+                ):
+                    self.assertNotIn("target_prefill", tuner._file_configs)
+                    self.assertNotIn("draft_decode", tuner._file_configs)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_compensated_mhc_update_guard.py

@@ -0,0 +1,61 @@
+"""Weight-update entry points refuse a model carrying compensated-mHC derived
+weight caches before any weight is written."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import torch
+
+from sglang.srt.model_executor.model_runner_components.weight_updater import (
+    WeightUpdater,
+    _unsupported_derived_weight_cache_error,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestCompensatedMhcUpdateGuard(CustomTestCase):
+    def test_all_update_entries_reject_before_writes(self):
+        for field in ("_hc_attn_tf32_parts", "_hc_ffn_tf32_parts"):
+            for method, args in (
+                ("update_weights_from_tensor", ([], "direct")),
+                ("update_weights_from_distributed", ([], [], [], "unused")),
+                ("update_weights_from_disk", ("unused", "auto")),
+                ("update_weights_from_ipc", (SimpleNamespace(),)),
+            ):
+                with self.subTest(field=field, method=method, args=args):
+                    model = torch.nn.Sequential(torch.nn.Linear(1, 1))
+                    original = model[0].weight.detach().clone()
+                    setattr(model[0], field, (torch.ones(1), torch.zeros(1)))
+                    model.load_weights = Mock()
+                    updater = SimpleNamespace(
+                        get_model=lambda: model, _assert_weight_cache_inactive=Mock()
+                    )
+                    with patch(
+                        "sglang.srt.model_executor.model_runner_components.weight_updater.default_weight_loader"
+                    ) as loader:
+                        ok, message = getattr(WeightUpdater, method)(updater, *args)
+                    self.assertFalse(ok)
+                    self.assertIn("compensated mHC", message)
+                    loader.assert_not_called()
+                    model.load_weights.assert_not_called()
+                    torch.testing.assert_close(
+                        model[0].weight, original, rtol=0, atol=0
+                    )
+
+    def test_models_without_derived_splits_keep_update_support(self):
+        model = torch.nn.Sequential(torch.nn.Linear(1, 1))
+        model[0]._hc_attn_tf32_parts = model[0]._hc_ffn_tf32_parts = None
+        with patch(
+            "sglang.kernels.ops.attention.dsv4.gemm.hpc_bf16xfp32_gemm_enabled",
+            return_value=False,
+        ):
+            self.assertIsNone(_unsupported_derived_weight_cache_error(model))
+            self.assertIsNone(_unsupported_derived_weight_cache_error())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/model_executor/test_pool_configurator.py

@@ -1152,6 +1152,9 @@ class TestSWAPoolFloor(CustomTestCase):
         cfg.c4_ring_size = 8
         cfg.c4_shrink_factor = 1
         cfg._unified = unified
+        cfg.operator_swa_ratio = None
+        cfg.swa_cap_tokens = None
+        cfg.swa_prefix_tails = 0
         return cfg._compute_dsv4_sizes(max_tokens, page_size)
 
     def test_dsv4_rejects_single_page_pool(self):
@@ -1214,6 +1217,11 @@ class TestSWAPoolFloor(CustomTestCase):
         cfg.disaggregation_mode = None
         cfg.disaggregation_decode_extra_slots = 0
         cfg._unified = True
+        cfg.operator_swa_ratio = None
+        cfg.swa_cap_tokens = None
+        cfg.swa_prefix_tails = 0
+        cfg.request_window_bytes = 0
+        cfg.bytes_per_swa_token = 0.0
         cfg._unified_fp8 = False
         # object.__new__ skips __init__; bf16 unified row is 2B * latent
         cfg._unified_row_bytes = cfg.attn_head_dim * 2
@@ -1231,6 +1239,50 @@ class TestSWAPoolFloor(CustomTestCase):
             + cfg._get_c128_state_fixed_bytes(max_running_requests)
         )
 
+    def test_dsv4_paged_dspark_budget_reserves_window_and_draft_layers(self):
+        from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator
+
+        _publish_config(
+            self,
+            enable_encoder_swa_bounded_replay=True,
+            speculative_algorithm="DSPARK",
+            speculative_num_draft_tokens=6,
+            speculative_dspark_block_size=5,
+            page_size=256,
+            max_running_requests=2,
+            chunked_prefill_size=256,
+        )
+        cfg = SimpleNamespace(
+            qk_nope_head_dim=448,
+            qk_rope_head_dim=64,
+            index_head_dim=128,
+            context_len=131072,
+            compress_ratios=[0, 0] + [2] * 18 + [1] * 20,
+            window_size=128,
+            hf_config=SimpleNamespace(kv_source_layer_ids=[2, 8, 14, 20]),
+        )
+        spec = SimpleNamespace(is_dspark=lambda: True, is_none=lambda: False)
+        kvc = SimpleNamespace(
+            kv_cache_dtype_str="fp8_e4m3",
+            model_config=cfg,
+            layer_info=SimpleNamespace(start_layer=0, end_layer=40),
+            ps=SimpleNamespace(pp_size=1, attn_dp_size=1),
+            sliding_window_size=128,
+            page_size=256,
+            spec_algorithm=spec,
+            spec_aux_config=SimpleNamespace(dflash_draft_num_layers=3),
+        )
+        planner = DSV4PoolConfigurator(kvc)
+        self.assertEqual(planner.bytes_per_swa_token, 3 * 584)
+        budget = 256 * 1024 * 1024
+        sizes = planner.calculate_pool_sizes(budget, 256)
+        self.assertEqual(sizes.swa_max_total_num_tokens, planner.swa_cap_tokens)
+        self.assertLessEqual(
+            sizes.full_max_total_num_tokens * planner.bytes_per_full_token
+            + planner._get_swa_fixed_bytes(),
+            budget,
+        )
+
     def test_dsv4_unified_c4_state_not_token_scaled(self):
         # Unified-KV sizes the c4 state ring from max_running_requests in
         # finalize_with_max_running_requests, so it must not scale here.

test/registered/unit/models/test_deepseek_v4_rope_policy.py

@@ -55,6 +55,7 @@ class TestDeepseekV4RoPEPolicy(CustomTestCase):
             o_lora_rank=8,
             rms_norm_eps=1e-6,
             compress_ratios=[compress_ratio],
+            q_head_norm=True,
             rope_theta=10_000,
             compress_rope_theta=160_000,
             max_position_embeddings=128,

test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py

@@ -57,6 +57,10 @@ class TestDeepseekV4SharedExpertFusionPolicy(CustomTestCase):
             quantization_config={},
             rope_scaling={},
             compress_ratios=[],
+            kv_source_layer_ids=[],
+            index_source_layer_ids=[],
+            engram_layer_ids=[],
+            engram_num_embeddings=[],
             n_shared_experts=1,
             dspark_markov_rank=1,
             num_nextn_predict_layers=1,

test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py

@@ -88,7 +88,7 @@ class _Harness(deepseek_v4.MQALayer):
                 dtype=torch.bfloat16,
             )
         )
-        self.wo_b = lambda value: (value, None)
+        self.wo_b = lambda value, skip_all_reduce=False: (value, None)
         self.prepare_kwargs = None
 
     def _forward_prepare(

test/registered/unit/test_model_overrides.py

@@ -1718,6 +1718,14 @@ class TestGoldenModelOverrides(_IsolatedPublish):
             "swa_full_tokens_ratio",
             _deepseek_v4_overrides(_args(swa_full_tokens_ratio=0.5), hf),
         )
+        # V4.1 leaves the ratio unset (cap-mode SWA sizing).
+        hf41 = SimpleNamespace(
+            architectures=["DeepseekV4ForCausalLM"], model_type="deepseek_v41"
+        )
+        self.assertNotIn(
+            "swa_full_tokens_ratio",
+            _deepseek_v4_overrides(_args(fp8_gemm_runner_backend="triton"), hf41),
+        )
         # An explicit user choice takes precedence over the model default.
         self.assertNotIn(
             "moe_runner_backend",