Add CuTe DSL Rubin (sm107) grouped GEMM GLU/dGLU/quant kernels

python/cudnn/api_base.py

@@ -33,6 +33,16 @@ def is_power_of_2(n: int) -> bool:
     return n > 0 and (n & (n - 1)) == 0
 
 
+def is_sm107_device() -> bool:
+    """Return True when the current CUDA device is Rubin (SM107)."""
+    return torch.cuda.is_available() and torch.cuda.get_device_capability(torch.cuda.current_device()) == (10, 7)
+
+
+def get_device_type() -> str:
+    """Return the architecture family used by SM100 grouped GEMM wrappers."""
+    return "rubin" if is_sm107_device() else "blackwell"
+
+
 _experimental_api_warnings_emitted = set()
 _experimental_api_warnings_lock = threading.Lock()
 
@@ -379,12 +389,16 @@ def __init__(self):
         - self._is_supported: Flag indicating if configuration is validated
         - self._kernel: Kernel instance
         - self._compiled_kernel: Cache for compiled kernel
+        - self._device_type: Architecture family used for dispatch/cache keys
+        - self._is_rubin_kernel: True when running on Rubin (SM107)
         - self._logger: Logger instance for this class
         """
         self._is_supported = False
         self._kernel = None
         self._compiled_kernel = None
         self._interpret_uint8_as_fp4x2 = False
+        self._device_type = get_device_type()
+        self._is_rubin_kernel = self._device_type == "rubin"
         self._logger = logging.getLogger(self.__class__.__name__)
 
     def _warn_experimental_api(self) -> None:

python/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.py

@@ -45,6 +45,7 @@
 
 from .moe_blockscaled_grouped_gemm_dglu_dbias import BlockScaledMoEGroupedGemmDgluDbiasKernel
 from ..moe_utils import MoEWeightMode
+from ..grouped_gemm_utils import rubin_single_group_offsets_kwarg
 from cuda.bindings import driver as cuda
 import os
 import torch
@@ -59,6 +60,29 @@
 from cudnn.api_base import APIBase, ceil_div, is_power_of_2
 
 
+def _get_rubin_kernel():
+    from .moe_blockscaled_grouped_gemm_dglu_rubin import (
+        BlockScaledMoEGroupedGemmDgluKernel as RubinBlockScaledMoEGroupedGemmDgluKernel,
+    )
+
+    return RubinBlockScaledMoEGroupedGemmDgluKernel
+
+
+_GEGGLU_ALPHA_DEFAULT = 1.702
+_GLU_CLAMP_MAX_DEFAULT = 7.0
+_GLU_CLAMP_MIN_DEFAULT = -7.0
+
+
+def _reject_unsupported_rubin_glu_tune_params(
+    is_rubin_kernel: bool,
+    geglu_alpha: float,
+    glu_clamp_max: float,
+    glu_clamp_min: float,
+) -> None:
+    if is_rubin_kernel and (geglu_alpha != _GEGGLU_ALPHA_DEFAULT or glu_clamp_max != _GLU_CLAMP_MAX_DEFAULT or glu_clamp_min != _GLU_CLAMP_MIN_DEFAULT):
+        raise NotImplementedError("Rubin grouped GEMM dGLU does not support geglu_alpha, glu_clamp_max, or glu_clamp_min tuning")
+
+
 class GroupedGemmDgluBlockScaledAPI(APIBase):
     """Unified API for grouped GEMM dGLU backward operation on SM100+ GPUs.
 
@@ -279,10 +303,16 @@ def __init__(
         self.geglu_alpha = geglu_alpha
         self.glu_clamp_max = glu_clamp_max
         self.glu_clamp_min = glu_clamp_min
+        _reject_unsupported_rubin_glu_tune_params(
+            self._is_rubin_kernel,
+            self.geglu_alpha,
+            self.glu_clamp_max,
+            self.glu_clamp_min,
+        )
 
         self._interpret_uint8_as_fp4x2 = True
         self._has_dbias = self.dbias_desc is not None
-        self._kernel = BlockScaledMoEGroupedGemmDgluDbiasKernel
+        self._kernel = _get_rubin_kernel() if self._is_rubin_kernel else BlockScaledMoEGroupedGemmDgluDbiasKernel
 
         self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0"))
         self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}")
@@ -605,8 +635,8 @@ def check_support(self) -> bool:
             f"m_aligned must be divisible by mma_tiler_mn[0], got {self.m_aligned} % {self.mma_tiler_mn[0]} != 0",
         )
         self._value_error_if(
-            self.m_aligned != BlockScaledMoEGroupedGemmDgluDbiasKernel.FIX_PAD_SIZE,
-            f"m_aligned must be {BlockScaledMoEGroupedGemmDgluDbiasKernel.FIX_PAD_SIZE} (FIX_PAD_SIZE), got {self.m_aligned}",
+            self.m_aligned != self._kernel.FIX_PAD_SIZE,
+            f"m_aligned must be {self._kernel.FIX_PAD_SIZE} (FIX_PAD_SIZE), got {self.m_aligned}",
         )
 
         # ---- Tensor alignment ----
@@ -688,7 +718,7 @@ def compile(self) -> None:
             weight_mode=self.weight_mode,
             act_func=self.act_func,
             use_dynamic_sched=self.use_dynamic_sched,
-            use_single_group_runtime_offsets=self.use_single_group_runtime_offsets,
+            **rubin_single_group_offsets_kwarg(self._is_rubin_kernel, self.use_single_group_runtime_offsets),
         )
 
         hardware_info = cutlass.utils.HardwareInfo()
@@ -903,8 +933,7 @@ def _compile_dense(self, gemm_dglu, max_active_clusters, fake_stream) -> None:
         # Compile with keyword args (dense mode uses the unified __call__ positional order).
         dbias_fake = self._make_fake_cute_tensor_from_desc(self.dbias_desc, assumed_align=16)
 
-        _compiled_kernel = cute.compile(
-            gemm_dglu,
+        compile_kwargs = dict(
             a=a_cute_fake,
             b=b_cute_fake,
             sfb=sfb_cute_fake,
@@ -930,12 +959,18 @@ def _compile_dense(self, gemm_dglu, max_active_clusters, fake_stream) -> None:
             max_active_clusters=max_active_clusters,
             stream=fake_stream,
             epilogue_op=self.epilogue_op,
-            linear_offset=self.linear_offset,
-            geglu_alpha=self.geglu_alpha,
-            glu_clamp_max=self.glu_clamp_max,
-            glu_clamp_min=self.glu_clamp_min,
+            linear_offset=cutlass.Float32(self.linear_offset) if self._is_rubin_kernel else self.linear_offset,
             options="--enable-tvm-ffi",
         )
+        if not self._is_rubin_kernel:
+            compile_kwargs.update(
+                {
+                    "geglu_alpha": self.geglu_alpha,
+                    "glu_clamp_max": self.glu_clamp_max,
+                    "glu_clamp_min": self.glu_clamp_min,
+                }
+            )
+        _compiled_kernel = cute.compile(gemm_dglu, **compile_kwargs)
 
         # Cache workspace pointer for the tensor_api closure
         cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator
@@ -961,7 +996,7 @@ def tensor_api(
             stream: cuda.CUstream,
         ) -> None:
             norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const")
-            _compiled_kernel(
+            kernel_args = (
                 a_tensor,
                 b_tensor,
                 sfb_tensor,
@@ -982,9 +1017,11 @@ def tensor_api(
                 beta_tensor,
                 prob_tensor,
                 dprob_tensor,
-                dbias_tensor,
-                stream,
             )
+            if self._is_rubin_kernel:
+                _compiled_kernel(*kernel_args, cutlass.Float32(self.linear_offset), dbias_tensor, stream)
+            else:
+                _compiled_kernel(*kernel_args, dbias_tensor, stream)
 
         self._compiled_kernel = tensor_api
 
@@ -1093,8 +1130,7 @@ def _compile_discrete(self, gemm_dglu, max_active_clusters, fake_stream) -> None
         workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator
 
         self._logger.debug("Compiling discrete grouped GEMM dGLU kernel")
-        _compiled_kernel = cute.compile(
-            gemm_dglu,
+        compile_kwargs = dict(
             a=a_tensor,
             b=b_ptrs_cute,
             sfb=sfb_ptrs_cute,
@@ -1120,12 +1156,18 @@ def _compile_discrete(self, gemm_dglu, max_active_clusters, fake_stream) -> None
             max_active_clusters=max_active_clusters,
             stream=fake_stream,
             epilogue_op=self.epilogue_op,
-            linear_offset=self.linear_offset,
-            geglu_alpha=self.geglu_alpha,
-            glu_clamp_max=self.glu_clamp_max,
-            glu_clamp_min=self.glu_clamp_min,
+            linear_offset=cutlass.Float32(self.linear_offset) if self._is_rubin_kernel else self.linear_offset,
             options="--enable-tvm-ffi",
         )
+        if not self._is_rubin_kernel:
+            compile_kwargs.update(
+                {
+                    "geglu_alpha": self.geglu_alpha,
+                    "glu_clamp_max": self.glu_clamp_max,
+                    "glu_clamp_min": self.glu_clamp_min,
+                }
+            )
+        _compiled_kernel = cute.compile(gemm_dglu, **compile_kwargs)
 
         self._n = n
         self._k = k
@@ -1161,7 +1203,7 @@ def tensor_api(
             b_ptrs_addr = int(b_ptrs_device.data_ptr())
             sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr())
 
-            _compiled_kernel(
+            kernel_args = (
                 a_tensor,
                 b_ptrs_addr,
                 sfb_ptrs_addr,
@@ -1182,9 +1224,11 @@ def tensor_api(
                 beta_tensor,
                 prob_tensor,
                 dprob_tensor,
-                dbias_tensor,
-                stream,
             )
+            if self._is_rubin_kernel:
+                _compiled_kernel(*kernel_args, cutlass.Float32(self.linear_offset), dbias_tensor, stream)
+            else:
+                _compiled_kernel(*kernel_args, dbias_tensor, stream)
 
         self._compiled_kernel = tensor_api
 

python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py

@@ -58,7 +58,7 @@
 import torch
 from typing import Any, Tuple, Optional, overload
 
-from cudnn.api_base import APIBase, TupleDict, ceil_div
+from cudnn.api_base import APIBase, TupleDict, ceil_div, get_device_type
 
 _BLOCK_SCALED_DTYPE_PAIRS = {
     (dtype, dtype)
@@ -72,7 +72,10 @@
 
 
 from ._bf16_api import GroupedGemmDgluBf16API
-from ._blockscaled_api import GroupedGemmDgluBlockScaledAPI
+from ._blockscaled_api import (
+    GroupedGemmDgluBlockScaledAPI,
+    _reject_unsupported_rubin_glu_tune_params,
+)
 
 
 @dataclass(frozen=True)
@@ -521,6 +524,13 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict:
     # default" (1.0 for dgeglu, 0.0 for dswiglu).
     if linear_offset is None:
         linear_offset = 1.0 if act_func == "dgeglu" else 0.0
+    device_type = get_device_type()
+    _reject_unsupported_rubin_glu_tune_params(
+        device_type == "rubin",
+        geglu_alpha,
+        glu_clamp_max,
+        glu_clamp_min,
+    )
     dgeglu_cache_signature = None
     if act_func == "dgeglu":
         dgeglu_cache_signature = (
@@ -632,6 +642,7 @@ def dynamic_m_tensor_signature(
 
     if is_dense:
         cache_key = (
+            device_type,
             weight_mode,
             act_func,
             dgeglu_cache_signature,
@@ -677,6 +688,7 @@ def dynamic_m_tensor_signature(
         )
     else:
         cache_key = (
+            device_type,
             weight_mode,
             act_func,
             dgeglu_cache_signature,

python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py

@@ -588,8 +588,10 @@ def _setup_attributes(self):
         # Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case
         self.overlapping_accum = self.num_acc_stage == 1 and self.mma_tiler[1] == 256
 
-        # To prefetch more accumulator when overlapping_accum is enabled in epilogue
-        self.epilogue_prefetch_more = self.d_dtype.width == 8 and self.a_dtype.width == 8
+        # The ping-pong prefetch path selects between two rmem tensor objects in
+        # the epilogue loop. Recent CuTe DSL lowers that to an arith.select over
+        # memrefs and leaves an unrealized_conversion_cast during LLVM lowering.
+        self.epilogue_prefetch_more = False
 
         # Use ptx fp8 fp32 convert
         self.use_fp8_ptx_cvt = True