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

@@ -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()

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
+  }
+}

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