sglang/test/registered/cpu/test_topk.py

import itertools
import unittest

import sgl_kernel  # noqa: F401
import torch

from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers.moe.topk import (
    biased_grouped_topk_impl as native_biased_grouped_topk,
)
from sglang.srt.layers.moe.topk import biased_topk_impl as native_biased_topk
from sglang.srt.layers.moe.topk import fused_topk_torch_native as native_fused_topk
from sglang.srt.layers.moe.topk import grouped_topk_gpu as native_grouped_topk
from sglang.srt.models.llama4 import Llama4MoE
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=30, suite="base-b-test-cpu")
register_cpu_ci(est_time=10, suite="base-b-test-cpu-arm64")


# This is used by the Deepseek-V2 model
class TestGroupedTopK(CustomTestCase):
    def _run_single_test(self, M, E, G, topk, topk_group, renormalize, dtype):
        torch.manual_seed(12)

        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
        hidden_states = torch.randn(M, 100, dtype=dtype)
        gating_output = torch.randn(M, E, dtype=dtype) * 2 * M

        ref_topk_weights, ref_topk_ids = native_grouped_topk(
            hidden_states.float(),
            gating_output.float(),
            topk,
            renormalize,
            G,
            topk_group,
        )

        # fused version
        topk_weights, topk_ids = torch.ops.sgl_kernel.grouped_topk_cpu(
            hidden_states,
            gating_output,
            topk,
            renormalize,
            G,
            topk_group,
            0,
            None,
            None,
        )

        res = torch.zeros(M, E, dtype=torch.float)
        ref = torch.zeros(M, E, dtype=torch.float)
        res.scatter_(1, topk_ids.long(), topk_weights)
        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
        torch.testing.assert_close(res, ref)

    def test_grouped_topk(self):
        for renormalize in [True, False]:
            self._run_single_test(123, 8, 2, 2, 1, renormalize, torch.bfloat16)
            self._run_single_test(123, 16, 4, 3, 2, renormalize, torch.bfloat16)
            self._run_single_test(123, 32, 4, 3, 2, renormalize, torch.bfloat16)
            self._run_single_test(1123, 32, 4, 3, 2, renormalize, torch.bfloat16)
            self._run_single_test(123, 64, 1, 6, 1, renormalize, torch.bfloat16)
            self._run_single_test(123, 256, 8, 4, 8, renormalize, torch.bfloat16)
            self._run_single_test(123, 160, 8, 6, 2, renormalize, torch.bfloat16)


# DeepSeek V2/V3/R1 uses biased_grouped_top
class TestBiasedGroupedTopK(CustomTestCase):
    def _run_single_test(
        self,
        M,
        E,
        G,
        topk,
        topk_group,
        renormalize,
        gating_dtype,
        bias_dtype,
        routed_scaling_factor,
    ):
        torch.manual_seed(1024)

        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
        hidden_states = torch.randn(M, 100, dtype=torch.bfloat16)
        gating_output = torch.randn(M, E, dtype=gating_dtype) * 2 * M
        correction_bias = torch.randn(E, dtype=bias_dtype)

        ref_topk_weights, ref_topk_ids = native_biased_grouped_topk(
            hidden_states.float(),
            gating_output.float(),
            correction_bias.float(),
            topk,
            renormalize,
            G,
            topk_group,
        )
        ref_topk_weights = (
            ref_topk_weights * routed_scaling_factor
            if routed_scaling_factor is not None
            else ref_topk_weights
        )
        # fused version
        topk_weights, topk_ids = torch.ops.sgl_kernel.biased_grouped_topk_cpu(
            hidden_states,
            gating_output,
            correction_bias,
            topk,
            renormalize,
            G,
            topk_group,
            0,
            routed_scaling_factor,
            None,
        )

        res = torch.zeros(M, E, dtype=torch.float)
        ref = torch.zeros(M, E, dtype=torch.float)
        res.scatter_(1, topk_ids.long(), topk_weights)
        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
        torch.testing.assert_close(res, ref)

    def test_biased_grouped_topk(self):
        for renormalize in [False]:
            for bias_dtype in [torch.float32, torch.bfloat16]:
                for gating_dtype in [torch.float32, torch.bfloat16]:
                    for routed_scaling_factor in [None, 1.125]:
                        for E_num in [128, 192, 256, 384]:
                            self._run_single_test(
                                34,
                                E_num,
                                8,
                                8,
                                2,
                                renormalize,
                                gating_dtype,
                                bias_dtype,
                                routed_scaling_factor,
                            )


class TestBiasedTopK(CustomTestCase):
    def test_biased_topk_returns_logical_ids_with_eplb_info(self):
        hidden_states = torch.ones(1, 4)
        gating_output = torch.tensor([[10.0, 9.0, 1.0, 0.0]])
        correction_bias = torch.zeros(4)
        dispatch_info = ExpertLocationDispatchInfo(
            ep_dispatch_algorithm="static",
            partial_logical_to_rank_dispatch_physical_map=torch.tensor(
                [2, 3, 0, 1], dtype=torch.int64
            ),
            partial_logical_to_all_physical_map=torch.tensor(
                [[2], [3], [0], [1]], dtype=torch.int64
            ),
            partial_logical_to_all_physical_map_num_valid=torch.ones(
                4, dtype=torch.int64
            ),
            num_physical_experts=4,
        )

        _, topk_ids = native_biased_topk(
            hidden_states=hidden_states,
            gating_output=gating_output,
            correction_bias=correction_bias,
            topk=2,
            renormalize=False,
            scoring_func="sqrtsoftplus",
            expert_location_dispatch_info=dispatch_info,
        )

        torch.testing.assert_close(topk_ids, torch.tensor([[0, 1]], dtype=torch.int32))


class TestTopK(CustomTestCase):
    def _run_single_test(self, M, E, topk, renormalize, dtype):
        torch.manual_seed(1998)

        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
        hidden_states = torch.randn(M, 100, dtype=dtype)
        gating_output = torch.randn(M, E, dtype=dtype) * 2 * M

        ref_topk_weights, ref_topk_ids = native_fused_topk(
            hidden_states.float(),
            gating_output.float(),
            topk,
            renormalize,
        )

        # fused version
        topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
            hidden_states, gating_output, topk, renormalize
        )

        res = torch.zeros(M, E, dtype=torch.float)
        ref = torch.zeros(M, E, dtype=torch.float)
        res.scatter_(1, topk_ids.long(), topk_weights)
        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
        torch.testing.assert_close(res, ref)

    def test_topk(self):
        for renormalize in [True, False]:
            self._run_single_test(123, 8, 2, renormalize, torch.bfloat16)
            self._run_single_test(123, 16, 3, renormalize, torch.bfloat16)
            self._run_single_test(123, 32, 3, renormalize, torch.bfloat16)
            self._run_single_test(123, 32, 3, renormalize, torch.bfloat16)
            self._run_single_test(123, 64, 6, renormalize, torch.bfloat16)
            self._run_single_test(123, 256, 4, renormalize, torch.bfloat16)
            self._run_single_test(123, 160, 6, renormalize, torch.bfloat16)

    def test_topk_softmax_mixed_input_dtypes(self):
        torch.manual_seed(0)
        hidden_states = torch.randn((17, 16), dtype=torch.bfloat16)
        gating_output = torch.randn((17, 128), dtype=torch.float32)
        correction_bias = torch.randn(128, dtype=torch.float32)

        topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
            hidden_states=hidden_states,
            gating_output=gating_output,
            topk=8,
            renormalize=True,
            correction_bias=correction_bias,
        )

        scores = torch.softmax(gating_output, dim=-1)
        expected_ids = torch.topk(
            scores + correction_bias.unsqueeze(0), k=8, dim=-1
        ).indices
        expected_weights = scores.gather(1, topk_ids.to(torch.int64))
        expected_weights /= expected_weights.sum(dim=-1, keepdim=True)

        self.assertEqual(
            torch.sort(topk_ids.to(torch.int64), dim=-1).values.tolist(),
            torch.sort(expected_ids, dim=-1).values.tolist(),
        )
        torch.testing.assert_close(topk_weights, expected_weights)

    def test_topk_softmax_with_correction_bias(self):
        """Bias must affect expert selection without becoming a routing weight."""
        for num_tokens, num_experts, topk, with_bias, renormalize in itertools.product(
            [1, 17, 128],
            [16, 128, 384, 512],
            [1, 2, 4, 8],
            [False, True],
            [False, True],
        ):
            torch.manual_seed(0)
            hidden_states = torch.randn((num_tokens, 16), dtype=torch.bfloat16)
            gating_output = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16)
            correction_bias = torch.randn(num_experts) if with_bias else None

            topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
                hidden_states=hidden_states,
                gating_output=gating_output,
                topk=topk,
                renormalize=renormalize,
                correction_bias=correction_bias,
            )

            scores = torch.softmax(gating_output.float(), dim=-1)
            scores_for_choice = scores
            if correction_bias is not None:
                scores_for_choice = scores_for_choice + correction_bias.unsqueeze(0)
            expected_choice_scores = torch.topk(
                scores_for_choice, k=topk, dim=-1, sorted=True
            ).values
            selected_choice_scores = torch.sort(
                scores_for_choice.gather(1, topk_ids.to(torch.int64)),
                dim=-1,
                descending=True,
            ).values

            expected_weights = scores.gather(1, topk_ids.to(torch.int64))
            if renormalize:
                expected_weights = expected_weights / expected_weights.sum(
                    dim=-1, keepdim=True
                )

            self.assertEqual(topk_ids.dtype, torch.int32)
            self.assertEqual(topk_weights.dtype, torch.float32)
            self.assertTrue(torch.all((topk_ids >= 0) & (topk_ids < num_experts)))
            sorted_ids = torch.sort(topk_ids, dim=-1).values
            self.assertTrue(torch.all(sorted_ids[:, 1:] != sorted_ids[:, :-1]))
            torch.testing.assert_close(
                selected_choice_scores,
                expected_choice_scores,
                atol=1e-4,
                rtol=1e-4,
            )
            torch.testing.assert_close(
                topk_weights, expected_weights, atol=1e-4, rtol=1e-4
            )


class TestCustomTopK(CustomTestCase):
    def _run_single_test(
        self, M, E, topk, renormalize, dtype, native_custom_f, fused_custom_f
    ):
        torch.manual_seed(16)

        # expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
        hidden_states = torch.randn(M, 100, dtype=dtype)
        gating_output = torch.randn(M, E, dtype=dtype) * 2 * M

        ref_topk_weights, ref_topk_ids = native_custom_f(
            hidden_states.float(),
            gating_output.float(),
            topk,
            renormalize,
        )

        # fused version
        topk_weights, topk_ids = fused_custom_f(
            hidden_states, gating_output, topk, renormalize
        )

        res = torch.zeros(M, E, dtype=torch.float)
        ref = torch.zeros(M, E, dtype=torch.float)
        res.scatter_(1, topk_ids.long(), topk_weights)
        ref.scatter_(1, ref_topk_ids.long(), ref_topk_weights)
        torch.testing.assert_close(res, ref)

    def test_custom_topk(self):
        test_custom_functions = [
            (Llama4MoE.custom_routing_function, torch.ops.sgl_kernel.topk_sigmoid_cpu)
        ]
        for native_custom_f, fused_custom_f in test_custom_functions:
            self._run_single_test(
                123, 8, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
            )
            self._run_single_test(
                123, 16, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
            )
            self._run_single_test(
                123, 32, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
            )

    def test_topk_sigmoid_with_correction_bias(self):
        """Biased scores must select experts while returned weights stay unbiased."""
        for num_tokens, num_experts, topk, with_bias, renormalize in itertools.product(
            [1, 17, 128],
            [16, 128, 256, 384, 512],
            [1, 2, 4, 8],
            [False, True],
            [False, True],
        ):
            torch.manual_seed(0)
            hidden_states = torch.randn((num_tokens, 16), dtype=torch.bfloat16)
            gating_output = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16)
            correction_bias = torch.randn(num_experts) if with_bias else None

            topk_weights, topk_ids = torch.ops.sgl_kernel.topk_sigmoid_cpu(
                hidden_states=hidden_states,
                gating_output=gating_output,
                topk=topk,
                renormalize=renormalize,
                correction_bias=correction_bias,
            )

            scores = torch.sigmoid(gating_output.float())
            scores_for_choice = scores
            if correction_bias is not None:
                scores_for_choice = scores_for_choice + correction_bias.unsqueeze(0)

            expected_choice_scores = torch.topk(
                scores_for_choice, k=topk, dim=-1
            ).values
            selected_choice_scores = torch.sort(
                scores_for_choice.gather(1, topk_ids.to(torch.int64)),
                dim=-1,
                descending=True,
            ).values

            expected_weights = scores.gather(1, topk_ids.to(torch.int64))
            if renormalize:
                expected_weights /= expected_weights.sum(dim=-1, keepdim=True)

            self.assertEqual(topk_ids.dtype, torch.int32)
            self.assertEqual(topk_weights.dtype, torch.float32)
            self.assertTrue(torch.equal(selected_choice_scores, expected_choice_scores))
            torch.testing.assert_close(
                topk_weights, expected_weights, atol=1e-4, rtol=1e-4
            )

    def test_topk_sigmoid_mixed_input_dtypes(self):
        torch.manual_seed(0)
        hidden_states = torch.randn((17, 16), dtype=torch.bfloat16)
        gating_output = torch.randn((17, 256), dtype=torch.float32)

        topk_weights, topk_ids = torch.ops.sgl_kernel.topk_sigmoid_cpu(
            hidden_states=hidden_states,
            gating_output=gating_output,
            topk=8,
            renormalize=True,
            correction_bias=None,
        )

        scores = torch.sigmoid(gating_output)
        expected_ids = torch.topk(scores, k=8, dim=-1).indices
        expected_weights = scores.gather(1, topk_ids.to(torch.int64))
        expected_weights /= expected_weights.sum(dim=-1, keepdim=True)

        self.assertTrue(
            torch.equal(
                torch.sort(topk_ids.to(torch.int64), dim=-1).values,
                torch.sort(expected_ids, dim=-1).values,
            )
        )
        torch.testing.assert_close(topk_weights, expected_weights, atol=1e-5, rtol=1e-5)


if __name__ == "__main__":
    unittest.main()