diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh index 6e0033e20..99ed41ce7 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh +++ b/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; diff --git a/python/sglang/kernels/jit/utils/occupancy.py b/python/sglang/kernels/jit/utils/occupancy.py index bdae2169b..0e01f688a 100644 --- a/python/sglang/kernels/jit/utils/occupancy.py +++ b/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" ) diff --git a/python/sglang/kernels/ops/attention/dsv4/topk.py b/python/sglang/kernels/ops/attention/dsv4/topk.py index 1ed069eff..d40391655 100644 --- a/python/sglang/kernels/ops/attention/dsv4/topk.py +++ b/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 - # defaults, and a 0 would size the persistent pool to an empty grid. + # Enable each cluster path only when its occupancy probe reports capacity. extra_cuda_cflags = [] if is_arch_support_pdl(): # set the persistent cluster size after hopper - try: - occ_8_2 = get_max_active_clusters(8, occupancy=2) - except Exception: - pass - else: - if occ_8_2 > 0: - extra_cuda_cflags.append(f"-DSGL_TOPK_V2_MAX_C8_OCC2={occ_8_2}") - try: - occ_16_1 = get_max_active_clusters(16, occupancy=1) - except Exception: - pass - else: - if occ_16_1 > 0: - extra_cuda_cflags.append(f"-DSGL_TOPK_V2_MAX_C16_OCC1={occ_16_1}") + for cluster_size, occupancy in ((8, 2), (16, 1)): + try: + max_active_clusters = get_max_active_clusters( + cluster_size, occupancy=occupancy + ) + except NoSchedulableClustersError: + max_active_clusters = 0 + extra_cuda_cflags.append( + f"-DSGL_TOPK_V2_MAX_C{cluster_size}_OCC{occupancy}={max_active_clusters}" + ) kernel = f"TopKKernel<{args}>" return load_jit( make_name("topk_v2"),