From 34a7502f4d5cd8e018201d482708915db2d40884 Mon Sep 17 00:00:00 2001 From: Tushar Sharma Date: Tue, 18 Aug 2026 08:12:10 -0700 Subject: [PATCH 1/3] fix: fall back when fused top-k lacks device support Signed-off-by: Tushar Sharma --- python/sglang/srt/layers/moe/topk.py | 14 +++++++++ .../ops/moe/test_fused_topk_deepseek.py | 31 ++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 951dc2238061..29a0447948d1 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -125,6 +125,7 @@ def routing( _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_musa = is_musa() +_fused_topk_deepseek = None # Epsilon added to the top-k weight sum before renormalization, matching the # DeepSeek reference gate (modeling_deepseek.py: `topk_weight.sum(...) + 1e-20`) @@ -145,6 +146,18 @@ def routing( _skip_hip_pad_mask = get_bool_env_var("SGLANG_MORI_NO_PAD_MASK", "False") +def _flashinfer_fused_topk_supports_device(device: torch.device) -> bool: + """Return whether FlashInfer's DeepSeek top-k kernel supports ``device``.""" + support_check = getattr( + _fused_topk_deepseek, "is_compute_capability_supported", None + ) + if support_check is None: + return True + + major, minor = torch.cuda.get_device_capability(device) + return support_check(major * 10 + minor) + + if _is_cuda: try: from flashinfer.fused_moe import fused_topk_deepseek as _fused_topk_deepseek @@ -1487,6 +1500,7 @@ def biased_grouped_topk_gpu( if ( _is_cuda and fused_topk_deepseek is not None + and _flashinfer_fused_topk_supports_device(gating_output.device) and is_power_of_two(num_experts) # flashinfer constraints (applied to routed experts only) and topk_routed <= 8 diff --git a/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py b/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py index 1afd8744b3d5..5f1afb31378f 100644 --- a/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py +++ b/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py @@ -1,8 +1,8 @@ import sys import pytest +import sglang.srt.layers.moe.topk as topk_module import torch - from sglang.srt.layers.moe.topk import biased_grouped_topk_gpu, biased_grouped_topk_impl from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci @@ -10,6 +10,35 @@ register_cuda_ci(est_time=40, stage="nightly", runner_config="1-gpu-large") +@pytest.mark.parametrize( + "compute_capability,expected", + [((8, 6), False), ((8, 9), True), ((9, 0), True)], +) +def test_flashinfer_fused_topk_checks_device_support( + monkeypatch, compute_capability, expected +): + class Kernel: + @staticmethod + def is_compute_capability_supported(capability): + return capability in {89, 90} + + monkeypatch.setattr(topk_module, "_fused_topk_deepseek", Kernel()) + monkeypatch.setattr( + torch.cuda, "get_device_capability", lambda device: compute_capability + ) + + assert ( + topk_module._flashinfer_fused_topk_supports_device(torch.device("cuda")) + is expected + ) + + +def test_flashinfer_fused_topk_without_support_metadata_is_accepted(monkeypatch): + monkeypatch.setattr(topk_module, "_fused_topk_deepseek", object()) + + assert topk_module._flashinfer_fused_topk_supports_device(torch.device("cuda")) + + @pytest.mark.parametrize( "seq_length", list(range(1, 10)) From 5330defe00669c5c689de764c88aa7690fc77f22 Mon Sep 17 00:00:00 2001 From: Tushar Sharma Date: Tue, 18 Aug 2026 08:29:11 -0700 Subject: [PATCH 2/3] style: fix fused top-k test imports Signed-off-by: Tushar Sharma --- test/registered/kernels/ops/moe/test_fused_topk_deepseek.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py b/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py index 5f1afb31378f..3a1d84bb792f 100644 --- a/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py +++ b/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py @@ -1,8 +1,9 @@ import sys import pytest -import sglang.srt.layers.moe.topk as topk_module import torch + +import sglang.srt.layers.moe.topk as topk_module from sglang.srt.layers.moe.topk import biased_grouped_topk_gpu, biased_grouped_topk_impl from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci From c9de8ab920b940104b29e47645d92d7d6da827d3 Mon Sep 17 00:00:00 2001 From: Tushar Sharma Date: Wed, 19 Aug 2026 14:00:07 -0700 Subject: [PATCH 3/3] test: remove fused top-k capability unit tests Signed-off-by: Tushar Sharma --- .../ops/moe/test_fused_topk_deepseek.py | 30 ------------------- 1 file changed, 30 deletions(-) diff --git a/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py b/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py index 3a1d84bb792f..1afd8744b3d5 100644 --- a/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py +++ b/test/registered/kernels/ops/moe/test_fused_topk_deepseek.py @@ -3,7 +3,6 @@ import pytest import torch -import sglang.srt.layers.moe.topk as topk_module from sglang.srt.layers.moe.topk import biased_grouped_topk_gpu, biased_grouped_topk_impl from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci @@ -11,35 +10,6 @@ register_cuda_ci(est_time=40, stage="nightly", runner_config="1-gpu-large") -@pytest.mark.parametrize( - "compute_capability,expected", - [((8, 6), False), ((8, 9), True), ((9, 0), True)], -) -def test_flashinfer_fused_topk_checks_device_support( - monkeypatch, compute_capability, expected -): - class Kernel: - @staticmethod - def is_compute_capability_supported(capability): - return capability in {89, 90} - - monkeypatch.setattr(topk_module, "_fused_topk_deepseek", Kernel()) - monkeypatch.setattr( - torch.cuda, "get_device_capability", lambda device: compute_capability - ) - - assert ( - topk_module._flashinfer_fused_topk_supports_device(torch.device("cuda")) - is expected - ) - - -def test_flashinfer_fused_topk_without_support_metadata_is_accepted(monkeypatch): - monkeypatch.setattr(topk_module, "_fused_topk_deepseek", object()) - - assert topk_module._flashinfer_fused_topk_supports_device(torch.device("cuda")) - - @pytest.mark.parametrize( "seq_length", list(range(1, 10))