feat(gemm): unified gemm_num_sms SM-count cap for all frontier GEMMs

Add a single knob `SonicMoEConfig.gemm_num_sms` (env
`SONIC_MOE_GEMM_NUM_SMS`) that caps the number of SMs every frontier
CUTLASS-DSL persistent GEMM may occupy, mirroring DeepGEMM's
`set_num_sms`. Intended for compute/communication multi-stream overlap:
capping the GEMM leaves SMs free for a comm kernel on another stream.

Mechanism (new helper `quack_utils/sm_limit.py`), applied identically at
every frontier call site:
  1. `capped_max_active_clusters` narrows the runtime-dynamic
     max_active_clusters (SM budget -> cluster budget, floor div, min 1).
  2. `clc_persistence_default` forces use_clc_persistence=False (STATIC
     scheduler) whenever the cap is set. On Blackwell (SM100+) the default
     CLC scheduler ignores max_active_clusters entirely (grid is always
     the full problem grid), so the cap only bites under STATIC. This is a
     compile-time constant, so `sm_cap_enabled()` is folded into every
     compile key (only the STATIC-vs-CLC bool, not the SM count -- the
     count only drives the runtime-dynamic grid, no recompile).

When gemm_num_sms is None, behavior is unchanged bit-for-bit (CLC path,
full grid). Covers: gemm_gated, gemm_dgated, the 3 gather_A+blockscaled
zeromat variants, the 7 blockscaled_fp8_gemm call sites, and
bf16_wgrad_gemm. SM90/Hopper is out of scope (no CLC; STATIC already
honors max_active_clusters there).

Verified on B300 (SM103, 148 SM, CUDA 13.2):
  - All GEMM unit tests pass with cap off AND cap forced to 64 (grid-only
    change, no reduction-order change -> numerics bit-identical).
  - ncu (real knob): grid.z tracks the cap (none=CLC full / 112->56 /
    64->32 / 32->16 clusters), SM_active% ~= residentCTA/148, per-SM
    efficiency stays ~1.0 (near-linear -- this is a compute-bound GEMM, so
    capping trades throughput ~linearly; the knob is for freeing SMs, not
    a free lunch).
  - Compile cache grows by exactly one STATIC entry when toggling cap on;
    different non-None caps reuse the same STATIC compilation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: adam-xiaoyao

sonicmoe/config.py

@@ -79,6 +79,10 @@ class SonicMoEConfig:
             Env: ``SONIC_MOE_STAGEWISE_MEMORY``. Default: False.
         swiglu_clamp_value: Optional user-controlled SwiGLU clamp value.
             Default: 0.0 (disabled).
+        gemm_num_sms: Absolute SM-count cap for all frontier CUTLASS GEMMs
+            (single knob, à la DeepGEMM ``set_num_sms``). None = use all SMs.
+            Non-None auto-switches the scheduler off Blackwell CLC onto STATIC
+            so the cap takes effect. Env: ``SONIC_MOE_GEMM_NUM_SMS``.
     """
 
     use_fp8: Optional[bool] = None
@@ -96,6 +100,7 @@ class SonicMoEConfig:
     stagewise_memory: Optional[bool] = None
     iso32_weight: Optional[bool] = None
     swiglu_clamp_value: Optional[float] = None
+    gemm_num_sms: Optional[int] = None
 
     def __post_init__(self) -> None:
         # Auto-enable quack_gemm when fp8 is explicitly enabled.
@@ -194,6 +199,30 @@ def resolve_swiglu_clamp_value(self) -> float:
             return max(float(self.swiglu_clamp_value), 0.0)
         return 0.0
 
+    def resolve_gemm_num_sms(self) -> Optional[int]:
+        """Global SM-count cap for all frontier CUTLASS GEMMs.
+
+        Returns an absolute upper bound on the number of SMs a persistent
+        GEMM may occupy, or ``None`` for "use all SMs" (unchanged behavior).
+        Intended for compute/communication multi-stream overlap: cap the
+        GEMM so a DeepEP/HybridEP comm kernel on another stream can claim
+        the freed SMs.  Mirrors DeepGEMM's ``set_num_sms`` single-knob API.
+
+        When non-None, the GEMM wrappers additionally switch the persistent
+        tile scheduler off the Blackwell CLC path (which ignores the SM cap)
+        onto STATIC scheduling, where the cap actually shrinks the grid.
+        Env fallback: ``SONIC_MOE_GEMM_NUM_SMS`` (int).
+        """
+        if self.gemm_num_sms is not None:
+            return int(self.gemm_num_sms)
+        raw = os.getenv("SONIC_MOE_GEMM_NUM_SMS", "").strip()
+        if raw:
+            try:
+                return int(raw)
+            except ValueError:
+                return None
+        return None
+
     # --- Context manager for temporary activation ----------------------------
 
     @contextmanager

sonicmoe/quack_utils/bf16_wgrad_gemm.py

@@ -14,6 +14,7 @@
 from quack.gemm_wrapper_utils import GemmTensorInfo, GemmWrapperBase
 
 from ..cache_manager import InstrumentedCompileCache as _ICC
+from .sm_limit import capped_max_active_clusters, clc_persistence_default, sm_cap_enabled
 
 
 _MAX_FAST_PATH_ENTRIES = 64
@@ -168,7 +169,7 @@ def _run_bf16_wgrad_varlen_k(
     ):
         raise TypeError("Unsupported BF16 wgrad type/major combination for varlen_k")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
         tile_count_semaphore=None,
@@ -198,6 +199,7 @@ def _run_bf16_wgrad_varlen_k(
         True,
         config.is_dynamic_persistent,
         config.device_capacity,
+        sm_cap_enabled(),
     )
     compiled = compile_cache.get(compile_key)
     if compiled is None:
@@ -207,7 +209,7 @@ def _run_bf16_wgrad_varlen_k(
             tile_shape_mn,
             cluster_shape_mnk,
             gather_A=True,
-            use_clc_persistence=config.is_dynamic_persistent,
+            use_clc_persistence=clc_persistence_default(config.is_dynamic_persistent),
         )
         compiled = cute.compile(
             gemm_obj,

sonicmoe/quack_utils/blockscaled_fp8_gemm.py

@@ -112,6 +112,7 @@ def _tile_atom_to_shape_SF_rank_aware(
     tuple[torch.Tensor, torch.Tensor],
 ] = {}
 from ..cache_manager import InstrumentedCompileCache as _ICC
+from .sm_limit import capped_max_active_clusters, clc_persistence_default, sm_cap_enabled
 _COMPILE_CACHE = _ICC("blockscaled_grouped")
 _PAD_PLAN_CACHE: dict = {}       # content-key -> plan
 # Fast-path cache: skip validation/tensor-info/compile-key on steady-state calls.
@@ -2265,7 +2266,7 @@ def _run_cutlass_blockscaled_gemm_varlen_k(
     ):
         raise TypeError("Unsupported FP8 blockscaled type/major combination for varlen_k")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
         tile_count_semaphore=None, batch_idx_permute=None,
@@ -2292,13 +2293,15 @@ def _run_cutlass_blockscaled_gemm_varlen_k(
         tensor_infos["A"].major, tensor_infos["B"].major,
         tensor_infos["D"].major,
         config.pingpong, _SF_VEC_SIZE,
+        sm_cap_enabled(),  # STATIC vs CLC scheduler (compile-time)
     )
     compiled = _COMPILE_CACHE_VK.get(compile_key)
     if compiled is None:
         gemm_obj = GemmDefaultSm100(
             Float32, tensor_infos["A"].dtype,
             tile_shape_mn, cluster_shape_mnk,
             sf_vec_size=_SF_VEC_SIZE, gather_A=False,
+            use_clc_persistence=clc_persistence_default(True),
         )
         compiled = cute.compile(
             gemm_obj,
@@ -2433,7 +2436,7 @@ def _run_cutlass_blockscaled_gemm_varlen_k_accumulate(
     ):
         raise TypeError("Unsupported FP8 blockscaled type/major combination for varlen_k accumulate")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
         tile_count_semaphore=None, batch_idx_permute=None,
@@ -2458,13 +2461,15 @@ def _run_cutlass_blockscaled_gemm_varlen_k_accumulate(
         tensor_infos["A"].major, tensor_infos["B"].major,
         tensor_infos["D"].major,
         config.pingpong, _SF_VEC_SIZE,
+        sm_cap_enabled(),  # STATIC vs CLC scheduler (compile-time)
     )
     compiled = _COMPILE_CACHE_VK_ACCUM.get(compile_key)
     if compiled is None:
         gemm_obj = GemmDefaultSm100(
             Float32, tensor_infos["A"].dtype,
             tile_shape_mn, cluster_shape_mnk,
             sf_vec_size=_SF_VEC_SIZE, gather_A=False,
+            use_clc_persistence=clc_persistence_default(True),
         )
         compiled = cute.compile(
             gemm_obj,
@@ -2596,7 +2601,7 @@ def _run_cutlass_blockscaled_gemm_varlen_k_tma_add(
     ):
         raise TypeError("Unsupported FP8 blockscaled type/major combination for varlen_k tma_add")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
         tile_count_semaphore=None, batch_idx_permute=None,
@@ -2621,13 +2626,15 @@ def _run_cutlass_blockscaled_gemm_varlen_k_tma_add(
         tensor_infos["A"].major, tensor_infos["B"].major,
         tensor_infos["D"].major,
         config.pingpong, _SF_VEC_SIZE,
+        sm_cap_enabled(),  # STATIC vs CLC scheduler (compile-time)
     )
     compiled = _COMPILE_CACHE_VK_TMA_ADD.get(compile_key)
     if compiled is None:
         gemm_obj = GemmDefaultSm100(
             Float32, tensor_infos["A"].dtype,
             tile_shape_mn, cluster_shape_mnk,
             sf_vec_size=_SF_VEC_SIZE, gather_A=False,
+            use_clc_persistence=clc_persistence_default(True),
         )
         compiled = cute.compile(
             gemm_obj,
@@ -2875,7 +2882,7 @@ def blockscaled_fp8_gemm_grouped(
     ):
         raise TypeError("Skipping due to unsupported FP8 blockscaled type/major combination")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
         tile_count_semaphore=None,
@@ -2911,6 +2918,7 @@ def blockscaled_fp8_gemm_grouped(
         config.pingpong,
         True,
         _SF_VEC_SIZE,
+        sm_cap_enabled(),  # STATIC vs CLC scheduler (compile-time)
     )
     compiled = _COMPILE_CACHE.get(compile_key)
     if compiled is None:
@@ -2921,6 +2929,7 @@ def blockscaled_fp8_gemm_grouped(
             cluster_shape_mnk,
             sf_vec_size=_SF_VEC_SIZE,
             gather_A=False,
+            use_clc_persistence=clc_persistence_default(True),
         )
         compiled = cute.compile(
             gemm_obj,
@@ -4591,7 +4600,7 @@ def _run_cutlass_blockscaled_gemm(
     ):
         raise TypeError("Unsupported FP8 blockscaled type/major combination")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
 
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
@@ -4631,6 +4640,7 @@ def _run_cutlass_blockscaled_gemm(
         config.pingpong,
         True,
         _SF_VEC_SIZE,
+        sm_cap_enabled(),  # STATIC vs CLC scheduler (compile-time)
     )
     compiled = _COMPILE_CACHE.get(compile_key)
     if compiled is None:
@@ -4641,6 +4651,7 @@ def _run_cutlass_blockscaled_gemm(
             cluster_shape_mnk,
             sf_vec_size=_SF_VEC_SIZE,
             gather_A=False,
+            use_clc_persistence=clc_persistence_default(True),
         )
         compiled = cute.compile(
             gemm_obj,
@@ -4842,7 +4853,7 @@ def blockscaled_fp8_weight_grad_gemm(
     ):
         raise TypeError("Unsupported FP8 blockscaled type/major combination for weight-grad GEMM")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
         tile_count_semaphore=None,
@@ -4873,6 +4884,7 @@ def blockscaled_fp8_weight_grad_gemm(
         config.pingpong,
         True,
         _SF_VEC_SIZE,
+        sm_cap_enabled(),  # STATIC vs CLC scheduler (compile-time)
     )
     compiled = _COMPILE_CACHE.get(compile_key)
     if compiled is None:
@@ -4883,6 +4895,7 @@ def blockscaled_fp8_weight_grad_gemm(
             cluster_shape_mnk,
             sf_vec_size=_SF_VEC_SIZE,
             gather_A=False,
+            use_clc_persistence=clc_persistence_default(True),
         )
         compiled = cute.compile(
             gemm_obj,
@@ -5038,7 +5051,7 @@ def blockscaled_fp8_weight_grad_gemm_fast(
     ):
         raise TypeError("Unsupported FP8 blockscaled type/major combination for weight-grad GEMM")
 
-    max_active_clusters = get_max_active_clusters(config.cluster_m * config.cluster_n)
+    max_active_clusters = capped_max_active_clusters(config.cluster_m * config.cluster_n)
     scheduler_args = GemmWrapperBase.create_scheduler_args(
         max_active_clusters,
         tile_count_semaphore=None,
@@ -5067,6 +5080,7 @@ def blockscaled_fp8_weight_grad_gemm_fast(
         config.pingpong,
         True,
         _SF_VEC_SIZE,
+        sm_cap_enabled(),  # STATIC vs CLC scheduler (compile-time)
     )
     compiled = _COMPILE_CACHE.get(compile_key)
     if compiled is None:
@@ -5077,6 +5091,7 @@ def blockscaled_fp8_weight_grad_gemm_fast(
             cluster_shape_mnk,
             sf_vec_size=_SF_VEC_SIZE,
             gather_A=False,
+            use_clc_persistence=clc_persistence_default(True),
         )
         compiled = cute.compile(
             gemm_obj,

sonicmoe/quack_utils/gemm_dgated.py

@@ -48,6 +48,7 @@
     GemmDGatedSm100ZeroMat,
     GemmDGatedFP8CLoadSm100ZeroMat,
 )
+from .sm_limit import capped_max_active_clusters, clc_persistence_default, sm_cap_enabled
 
 _E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", torch.uint8)
 
@@ -194,7 +195,7 @@ def gemm_dgated(
     ):
         raise TypeError("Skipping due to unsupported combination of types and majors")
 
-    max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
+    max_active_clusters = capped_max_active_clusters(cluster_M * cluster_N, persistent=persistent)
     for name, info in tensor_infos.items():
         if info.tensor is not None and name in major_configs:
             info.cute_tensor = _make_cute_tensor_dynamic(
@@ -265,19 +266,24 @@ def gemm_dgated(
         blockscaled,
         fp8_preact_mode,
         float(swiglu_clamp_value),
+        sm_cap_enabled(),
         key_tensor_names=("A", "B", "D", "PostAct", "C"),
     )
     cache = gemm_dgated.compile_cache
     if compile_key not in cache:
         if device_capacity[0] == 9:
             GemmCls = partial(GemmCls, pingpong=pingpong, is_persistent=persistent)
+            extra_kwargs = {}
+        else:
+            extra_kwargs = {"use_clc_persistence": clc_persistence_default(True)}
         gemm_obj = GemmCls(
             acc_dtype,
             tensor_infos["A"].dtype,
             tile_shape_mn,
             cluster_shape_mnk,
             gather_A=gather_A,
             sf_vec_size=sf_vec_size,
+            **extra_kwargs,
         )
         cache[compile_key] = cute.compile(
             gemm_obj,

sonicmoe/quack_utils/gemm_gated.py

@@ -48,6 +48,8 @@
     GemmSm100ZeroMatBlockscaledQuant,
 )
 
+from .sm_limit import capped_max_active_clusters, clc_persistence_default, sm_cap_enabled
+
 _E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", torch.uint8)
 
 
@@ -198,7 +200,7 @@ def gemm_gated(
     ):
         raise TypeError("Skipping due to unsupported combination of types and majors")
 
-    max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
+    max_active_clusters = capped_max_active_clusters(cluster_M * cluster_N, persistent=persistent)
     for name, info in tensor_infos.items():
         if info.tensor is not None and name in major_configs:
             leading_dim = 1 if info.major == major_configs[name][1] else 0
@@ -274,19 +276,27 @@ def gemm_gated(
         postact_quant,
         float(swiglu_clamp_value),
         bool(postact_bf16_trunc),
+        sm_cap_enabled(),
         key_tensor_names=("A", "B", "D", "PostAct", "C"),
     )
     cache = gemm_gated.compile_cache
     if compile_key not in cache:
         if device_capacity[0] == 9:
             GemmCls = partial(GemmCls, pingpong=pingpong, is_persistent=persistent)
+            extra_kwargs = {}
+        else:
+            # SM100+: force STATIC scheduling (use_clc_persistence=False) when
+            # the SM cap is active, so the narrowed max_active_clusters shrinks
+            # the persistent grid (the CLC path ignores it).
+            extra_kwargs = {"use_clc_persistence": clc_persistence_default(True)}
         gemm_obj = GemmCls(
             acc_dtype,
             tensor_infos["A"].dtype,
             tile_shape_mn,
             cluster_shape_mnk,
             gather_A=gather_A,
             sf_vec_size=sf_vec_size,
+            **extra_kwargs,
         )
         cache[compile_key] = cute.compile(
             gemm_obj,