Fix TopK v2 fallback when 16-block cluster capacity is zero (#40163)

Co-authored-by: Hank Han <hanhan7630@outlook.com>

python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh

@@ -328,7 +328,8 @@ TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKPagedParams params
 #endif
 
 #ifndef SGL_TOPK_V2_MAX_C16_OCC1
-#define SGL_TOPK_V2_MAX_C16_OCC1 7
+// Non-portable clusters require a positive device probe.
+#define SGL_TOPK_V2_MAX_C16_OCC1 0
 #endif
 
 constexpr uint32_t kNumPersistentClusters = SGL_TOPK_V2_MAX_C8_OCC2;

python/sglang/kernels/jit/utils/occupancy.py

@@ -10,7 +10,11 @@ from sglang.kernels.jit.utils.compile import load_jit
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
 
-__all__ = ["get_max_active_clusters"]
+__all__ = ["NoSchedulableClustersError", "get_max_active_clusters"]
+
+
+class NoSchedulableClustersError(ValueError):
+    """The occupancy query succeeded, but no cluster fits the requested shape."""
 
 
 @cache_once
@@ -36,11 +40,13 @@ def get_max_active_clusters(cluster_size: int, occupancy: int) -> int:
     dividing a GPC evenly. The probe kernel is pinned to ``occupancy`` blocks per
     SM, so pass the occupancy the real kernel reaches (its second
     ``__launch_bounds__`` argument). Raises ``RuntimeError`` before sm90, which
-    has no clusters, and ``ValueError`` when nothing is schedulable.
+    has no clusters, and ``NoSchedulableClustersError`` (a ``ValueError``)
+    when the query succeeds but nothing is schedulable. Other probe errors
+    propagate to the caller.
     """
     result = _get_max_active_clusters(cluster_size, occupancy)
     if result == 0:
-        raise ValueError(
+        raise NoSchedulableClustersError(
             f"no cluster of {cluster_size} fits at occupancy {occupancy}; "
             "the cluster width is likely beyond what this device supports"
         )

python/sglang/kernels/ops/attention/dsv4/topk.py

@@ -29,27 +29,25 @@ def _jit_topk_v1_module():
 
 @cache_once
 def _jit_topk_v2_module():
-    from sglang.kernels.jit.utils.occupancy import get_max_active_clusters
+    from sglang.kernels.jit.utils.occupancy import (
+        NoSchedulableClustersError,
+        get_max_active_clusters,
+    )
 
     args = make_cpp_args(is_arch_support_pdl())
-    # Leave these undefined if the probe fails: topk_v2.cuh carries per-arch
+    # Enable each cluster path only when its occupancy probe reports capacity.
-    # defaults, and a 0 would size the persistent pool to an empty grid.
     extra_cuda_cflags = []
     if is_arch_support_pdl():  # set the persistent cluster size after hopper
-        try:
+        for cluster_size, occupancy in ((8, 2), (16, 1)):
-            occ_8_2 = get_max_active_clusters(8, occupancy=2)
+            try:
-        except Exception:
+                max_active_clusters = get_max_active_clusters(
-            pass
+                    cluster_size, occupancy=occupancy
-        else:
+                )
-            if occ_8_2 > 0:
+            except NoSchedulableClustersError:
-                extra_cuda_cflags.append(f"-DSGL_TOPK_V2_MAX_C8_OCC2={occ_8_2}")
+                max_active_clusters = 0
-        try:
+            extra_cuda_cflags.append(
-            occ_16_1 = get_max_active_clusters(16, occupancy=1)
+                f"-DSGL_TOPK_V2_MAX_C{cluster_size}_OCC{occupancy}={max_active_clusters}"
-        except Exception:
+            )
-            pass
-        else:
-            if occ_16_1 > 0:
-                extra_cuda_cflags.append(f"-DSGL_TOPK_V2_MAX_C16_OCC1={occ_16_1}")
     kernel = f"TopKKernel<{args}>"
     return load_jit(
         make_name("topk_v2"),