Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
A-nnonymous wants to merge 2 commits into
base: paddle
Choose a base branch
from
Draft

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4b175dd
feat(gemm): unified gemm_num_sms SM-count cap for all frontier GEMMs
A-nnonymous Jul 2, 2026
97ee27f
fix(sm-limit): import config at module top level (install-alias safe)
A-nnonymous Jul 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions sonicmoe/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions sonicmoe/quack_utils/bf16_wgrad_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand Down
29 changes: 22 additions & 7 deletions sonicmoe/quack_utils/blockscaled_fp8_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand Down
8 changes: 7 additions & 1 deletion sonicmoe/quack_utils/gemm_dgated.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 11 additions & 1 deletion sonicmoe/quack_utils/gemm_gated.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Loading