From 02cd44c59a69f4176978698c5e907f808af288a0 Mon Sep 17 00:00:00 2001
From: yvbbrjdr <xyyqzy@hotmail.com>
Date: Wed, 5 Aug 2026 15:06:59 -0700
Subject: [PATCH] feat(dgx-spark): add inkling-small MoE support for sm_121
 (#33108)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
---
 python/sglang/kernels/ops/moe/inkling_moe.py  |  7 +++-
 .../silu_and_mul_interleaved_sm_121.json      | 39 +++++++++++++++++++
 python/sglang/srt/utils/common.py             |  7 ++++
 3 files changed, 52 insertions(+), 1 deletion(-)
 create mode 100644 python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_121.json

diff --git a/python/sglang/kernels/ops/moe/inkling_moe.py b/python/sglang/kernels/ops/moe/inkling_moe.py
index 2ec710989..768e2b666 100644
--- a/python/sglang/kernels/ops/moe/inkling_moe.py
+++ b/python/sglang/kernels/ops/moe/inkling_moe.py
@@ -11,6 +11,7 @@ from sglang.srt.layers.moe.moe_runner.triton_utils.helion_utils import (
     get_model_depths,
     helion_aot_autotune,
 )
+from sglang.srt.utils.common import is_sm121
 
 DEFAULT_BLOCK_SIZE = 4096
 BLOCK_SIZE_M = 128
@@ -884,7 +885,10 @@ def grouped_gemm_triton(
             "BLOCK_SIZE_K": 128,
             "GROUP_SIZE_M": 8,
             "num_warps": 4,
-            "num_stages": 4,
+            # sm_121 (GB10 / DGX Spark) caps shared memory at 99 KB per block and
+            # num_stages=4 needs 108 KB. The BLOCK_SIZE_M=128 branch below fits at
+            # its default 3 (96 KB) and needs no gate.
+            "num_stages": 3 if is_sm121() else 4,
         }
     else:
         assert block_size_m == BLOCK_SIZE_M, f"{block_size_m=}"
@@ -896,6 +900,7 @@ def grouped_gemm_triton(
             "num_warps": 8,
             "num_stages": 3,
         }
+
     # Set grid_m to the max number of M blocks and skip padding-only blocks
     # in the kernel based on expert_block_offs[-1]
     grid_m = expert_block_schedule.numel()
diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_121.json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_121.json
new file mode 100644
index 000000000..badc7a9ee
--- /dev/null
+++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/silu_and_mul_interleaved_sm_121.json
@@ -0,0 +1,39 @@
+{
+  "useful_configs": {
+    "0": "helion.Config(block_sizes=[2, 512], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
+    "1": "helion.Config(block_sizes=[1, 1024], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
+    "4": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])"
+  },
+  "hash_configs": {
+    "(512, (2,), (False,))": 0,
+    "(512, (2,), (True,))": 1,
+    "(1024, (2,), (False,))": 0,
+    "(1024, (2,), (True,))": 0,
+    "(1536, (2,), (False,))": 1,
+    "(1536, (2,), (True,))": 1,
+    "(2048, (2,), (False,))": 4,
+    "(2048, (2,), (True,))": 1,
+    "(3072, (2,), (False,))": 0,
+    "(3072, (2,), (True,))": 0,
+    "(4096, (2,), (False,))": 4,
+    "(4096, (2,), (True,))": 1,
+    "(4608, (2,), (False,))": 0,
+    "(4608, (2,), (True,))": 0,
+    "(6144, (2,), (False,))": 0,
+    "(6144, (2,), (True,))": 1,
+    "(7680, (2,), (False,))": 1,
+    "(7680, (2,), (True,))": 1,
+    "(8192, (2,), (False,))": 4,
+    "(8192, (2,), (True,))": 1,
+    "(9216, (2,), (False,))": 0,
+    "(9216, (2,), (True,))": 0,
+    "(10240, (2,), (False,))": 0,
+    "(10240, (2,), (True,))": 1,
+    "(12288, (2,), (False,))": 0,
+    "(12288, (2,), (True,))": 1,
+    "(14336, (2,), (False,))": 0,
+    "(14336, (2,), (True,))": 1,
+    "(16384, (2,), (False,))": 4,
+    "(16384, (2,), (True,))": 1
+  }
+}
diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py
index 6636d06ab..116e93b07 100644
--- a/python/sglang/srt/utils/common.py
+++ b/python/sglang/srt/utils/common.py
@@ -309,6 +309,13 @@ is_sm90_supported = lru_cache(maxsize=1)(
 )
 
 
+# GB10 (DGX Spark and OEM equivalents). Not expressible via
+# _check_cuda_device_version, which only matches on the major.
+@lru_cache(maxsize=1)
+def is_sm121() -> bool:
+    return is_cuda() and torch.cuda.get_device_capability() == (12, 1)
+
+
 try:
     import sgl_kernel  # noqa: F401