feat(overlap): global gemm_num_sms SM cap for frontier GEMMs

sonicmoe/config.py

@@ -79,6 +79,12 @@ 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: Cap the number of SMs the frontier CuTe-DSL GEMMs may
+            occupy, by shrinking each persistent kernel's ``max_active_clusters``
+            (grid.z).  Modeled after DeepGEMM's ``set_num_sms`` single knob.
+            Used to leave SMs free for communication kernels when overlapping
+            compute and comm on separate CUDA streams.  Env:
+            ``SONIC_MOE_GEMM_NUM_SMS``. Default: None (use all SMs).
     """
 
     use_fp8: Optional[bool] = None
@@ -96,6 +102,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 +201,23 @@ 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]:
+        """Return the SM cap for frontier GEMMs, or None to use all SMs.
+
+        Priority: explicit field > env ``SONIC_MOE_GEMM_NUM_SMS`` > None.
+        A value <= 0 is treated as unset (None).
+        """
+        if self.gemm_num_sms is not None:
+            return self.gemm_num_sms if self.gemm_num_sms > 0 else None
+        raw = os.getenv("SONIC_MOE_GEMM_NUM_SMS", "").strip()
+        if not raw:
+            return None
+        try:
+            n = int(raw)
+        except ValueError:
+            return None
+        return n if n > 0 else None
+
     # --- Context manager for temporary activation ----------------------------
 
     @contextmanager

sonicmoe/quack_utils/bf16_wgrad_gemm.py

@@ -9,6 +9,8 @@
 from cutlass import Float32
 from cutlass.cute.runtime import from_dlpack
 from quack.cute_dsl_utils import get_device_capacity, get_max_active_clusters
+
+from .sm_limit import capped_max_active_clusters
 from quack.gemm_config import GemmConfig
 from quack.gemm_default_epi import GemmDefaultSm100
 from quack.gemm_wrapper_utils import GemmTensorInfo, GemmWrapperBase
@@ -168,7 +170,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,

sonicmoe/quack_utils/blockscaled_fp8_gemm.py

@@ -24,6 +24,8 @@
 from cutlass.cute.runtime import from_dlpack
 from cutlass.utils.blockscaled_layout import BlockScaledBasicChunk
 from quack.cute_dsl_utils import get_device_capacity, get_max_active_clusters
+
+from .sm_limit import capped_max_active_clusters
 from quack.gemm_default_epi import GemmDefaultSm100
 from quack.gemm_interface import default_config
 from quack.gemm_wrapper_utils import GemmTensorInfo, GemmWrapperBase
@@ -2265,7 +2267,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,
@@ -2433,7 +2435,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,
@@ -2596,7 +2598,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,
@@ -2875,7 +2877,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,
@@ -4591,7 +4593,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,
@@ -4842,7 +4844,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,
@@ -5038,7 +5040,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,

sonicmoe/quack_utils/gemm_dgated.py

@@ -49,6 +49,8 @@
     GemmDGatedFP8CLoadSm100ZeroMat,
 )
 
+from .sm_limit import capped_max_active_clusters
+
 _E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", torch.uint8)
 
 
@@ -194,7 +196,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(

sonicmoe/quack_utils/gemm_gated.py

@@ -17,6 +17,8 @@
 from cutlass._mlir.dialects import llvm
 from cutlass.cute.runtime import from_dlpack
 from quack.cute_dsl_utils import get_device_capacity, get_max_active_clusters, mlir_namedtuple
+
+from .sm_limit import capped_max_active_clusters
 from quack.epi_ops import TileStore, EpiOp, assume_stride_divisibility
 from quack.gemm_act import GemmActMixin
 from quack.gemm_default_epi import GemmDefaultEpiMixin
@@ -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

sonicmoe/quack_utils/gemm_sm100_fp8_zeromat.py

@@ -454,6 +454,8 @@ class GemmDGatedFP8CLoadSm100ZeroMat(GemmDGatedFP8CLoadMixin, _GemmSm100ZeroMatM
 from cutlass.cute.runtime import from_dlpack
 from quack.cute_dsl_utils import get_device_capacity, get_max_active_clusters
 
+from .sm_limit import capped_max_active_clusters
+
 _TORCH_TO_CUTLASS = {
     torch.float8_e4m3fn: cutlass.Float8E4M3FN,
     _E8M0_DTYPE: cutlass.Float8E8M0FNU,
@@ -532,7 +534,7 @@ def gemm_gated_zeromat(
 
     tile_M, tile_N = 128, 128
     cluster_M, cluster_N = 1, 1
-    max_active_clusters = get_max_active_clusters(cluster_M * cluster_N)
+    max_active_clusters = capped_max_active_clusters(cluster_M * cluster_N)
 
     for name, info in tensor_infos.items():
         if info.tensor is not None and name in major_configs:
@@ -657,7 +659,7 @@ def blockscaled_fp8_gemm_zeromat_quant(
 
     tile_M, tile_N = 128, 128
     cluster_M, cluster_N = 1, 1
-    max_active_clusters = get_max_active_clusters(cluster_M * cluster_N)
+    max_active_clusters = capped_max_active_clusters(cluster_M * cluster_N)
 
     for name, info in tensor_infos.items():
         if info.tensor is not None and name in major_configs:
@@ -781,7 +783,7 @@ def blockscaled_fp8_gemm_zeromat_bf16(
     GemmCls = GemmSm100ZeroMatBf16
 
     tile_M, tile_N, cluster_M, cluster_N = 128, 128, 1, 1
-    max_active_clusters = get_max_active_clusters(cluster_M * cluster_N)
+    max_active_clusters = capped_max_active_clusters(cluster_M * cluster_N)
 
     for name, info in tensor_infos.items():
         if info.tensor is not None and name in major_configs:

sonicmoe/quack_utils/sm_limit.py

@@ -0,0 +1,82 @@
+# ********************************************************************************
+# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao
+# ********************************************************************************
+
+"""Unified SM-occupancy cap for the frontier CuTe-DSL persistent GEMMs.
+
+Every persistent GEMM in the FP8/BF16 frontier launches with a
+``max_active_clusters`` value that upper-bounds how many persistent clusters
+(hence SMs, since one CTA == one SM) the grid may occupy.  By default that is
+the hardware full-occupancy value from ``get_max_active_clusters(cluster_size)``.
+
+To overlap compute and communication on separate CUDA streams we need to leave
+some SMs free for the comm kernels.  ``capped_max_active_clusters`` shrinks the
+per-launch ``max_active_clusters`` to honor a global SM budget
+(``SonicMoEConfig.gemm_num_sms`` / env ``SONIC_MOE_GEMM_NUM_SMS``), mirroring
+DeepGEMM's single ``set_num_sms`` knob.
+
+The cap only lowers grid.z; it never changes the reduction order, so results
+stay bit-identical to the uncapped run.  ``max_active_clusters`` is passed to
+the kernel as a runtime dynamic ``Int32`` (via ``TileSchedulerOptions``), so
+changing the cap does NOT force a recompile.
+
+This module deliberately depends only on ``..config`` (single-directional, no
+cycle) and ``quack.cute_dsl_utils``.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from quack.cute_dsl_utils import get_max_active_clusters
+
+from ..config import get_active_config
+
+
+def _resolve_gemm_num_sms() -> Optional[int]:
+    """Return the global SM budget for frontier GEMMs, or None for all SMs.
+
+    Priority: active ``SonicMoEConfig`` > env ``SONIC_MOE_GEMM_NUM_SMS`` > None.
+    Delegates entirely to ``SonicMoEConfig.resolve_gemm_num_sms`` (which also
+    reads the env fallback) so there is exactly one place that knows the
+    resolution order.
+    """
+    cfg = get_active_config()
+    if cfg is not None:
+        return cfg.resolve_gemm_num_sms()
+    # No active config: still honor the env var for scripts / standalone use.
+    import os
+
+    raw = os.getenv("SONIC_MOE_GEMM_NUM_SMS", "").strip()
+    if not raw:
+        return None
+    try:
+        n = int(raw)
+    except ValueError:
+        return None
+    return n if n > 0 else None
+
+
+def capped_max_active_clusters(cluster_size: int, *, persistent: bool = True) -> int:
+    """Hardware full-occupancy clusters, capped to the global SM budget.
+
+    Args:
+        cluster_size: number of CTAs per cluster (``cluster_M * cluster_N``).
+        persistent: when False the caller uses a non-persistent scheduler and
+            expects 0 (matching the historical ``... if persistent else 0``).
+
+    Returns:
+        The ``max_active_clusters`` to pass to the tile scheduler.  When a SM
+        budget is set, this is ``min(hw, num_sms // cluster_size)`` clamped to
+        at least 1 cluster; otherwise the hardware value.
+    """
+    if not persistent:
+        return 0
+    hw = get_max_active_clusters(cluster_size)
+    n = _resolve_gemm_num_sms()
+    if n is None:
+        return hw
+    cs = max(1, int(cluster_size))
+    # SM budget -> cluster budget: each cluster occupies `cluster_size` SMs.
+    budget_clusters = n // cs
+    return max(1, min(hw, budget_clusters))

tests/ops/bench_gemm_sm_cap.py

@@ -0,0 +1,71 @@
+"""Micro-benchmark: gemm_gated up-proj latency at several SM caps.
+
+Demonstrates the compute/comm-overlap lever: capping SMs frees SMs for a
+concurrent communication kernel on another stream, at the cost of a modest
+per-GEMM slowdown.  Reports per-call latency for the up-proj gated GEMM at
+gemm_num_sms in {None(all), 128, 112, 96, 64}.
+
+Run: USE_QUACK_GEMM=1 SONIC_MOE_FP8_MODE=perf with quack on PYTHONPATH.
+"""
+import os
+import torch
+
+os.environ.setdefault("USE_QUACK_GEMM", "1")
+os.environ.setdefault("SONIC_MOE_FP8_MODE", "perf")
+os.environ.setdefault("SONIC_MOE_FP8_ASSUME_ALIGNED", "1")
+
+from sonicmoe.config import SonicMoEConfig, set_active_config
+from sonicmoe.quack_utils.gemm_interface import gemm_gated
+from sonicmoe.quack_utils.blockscaled_fp8_gemm import (
+    quantize_and_pack_activation,
+    precompute_weight_fp8_for_fused_gated,
+)
+
+
+def _setup(T=8192, H=3072, I=1536, E=8, K=8):
+    TK = T * K // E
+    total_M = TK * E
+    cu = torch.arange(0, (E + 1) * TK, TK, dtype=torch.int32, device="cuda")
+    x = torch.randn(total_M, H, dtype=torch.bfloat16, device="cuda") * 0.02
+    w1 = torch.randn(2 * I, H, E, dtype=torch.bfloat16, device="cuda") * 0.02
+    x_fp8, a_scales = quantize_and_pack_activation(x)
+    w_fp8, b_scales = precompute_weight_fp8_for_fused_gated(w1)
+    return x_fp8, w_fp8, a_scales, b_scales, cu
+
+
+def _bench(cap, x_fp8, w_fp8, a_scales, b_scales, cu, iters=50):
+    set_active_config(SonicMoEConfig(gemm_num_sms=cap) if cap else None)
+    try:
+        def run():
+            gemm_gated(x_fp8, w_fp8, activation="swiglu", cu_seqlens_m=cu,
+                       a_scales=a_scales, b_scales=b_scales)
+        for _ in range(10):
+            run()
+        torch.cuda.synchronize()
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(iters):
+            run()
+        end.record()
+        torch.cuda.synchronize()
+        return start.elapsed_time(end) / iters * 1000.0  # us/call
+    finally:
+        set_active_config(None)
+
+
+def main():
+    x_fp8, w_fp8, a_scales, b_scales, cu = _setup()
+    print(f"up-proj gated GEMM (ERNIE shape T=8192,H=3072,I=1536,E=8,K=8), fp8")
+    print(f"{'cap':>8} {'us/call':>10}")
+    base = None
+    for cap in (None, 128, 112, 96, 64):
+        us = _bench(cap, x_fp8, w_fp8, a_scales, b_scales, cu)
+        if base is None:
+            base = us
+        label = "all" if cap is None else str(cap)
+        print(f"{label:>8} {us:>10.2f}  ({us/base:.2f}x baseline)")
+
+
+if __name__ == "__main__":
+    main()