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>
2026-09-18 02:55:30 -07:00
parent 1b200ffaaa
commit a6cf05817f
103 changed files with 8807 additions and 723 deletions
@@ -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
@@ -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()
@@ -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()