rename

Signed-off-by: Enwei Zhu <[email protected]>

Author: syuoni

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -22,12 +22,12 @@
     import cutlass
     import cutlass.cute as cute
 
+    from ..cute_dsl_kernels.blackwell.blockscaled_contiguous_grouped_gemm import \
+        Sm100BlockScaledContiguousGroupedGemmKernel
+    from ..cute_dsl_kernels.blackwell.blockscaled_contiguous_grouped_gemm_finalize_fusion import \
+        Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel
     from ..cute_dsl_kernels.blackwell.dense_blockscaled_gemm_persistent import \
         Sm100BlockScaledPersistentDenseGemmKernel
-    from ..cute_dsl_kernels.blackwell.grouped_blockscaled_gemm_finalize_fusion import \
-        Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel
-    from ..cute_dsl_kernels.blackwell.grouped_blockscaled_gemm_persistent import \
-        Sm100BlockScaledPersistentGroupedGemmKernel
     from ..cute_dsl_kernels.blackwell.utils import make_ptr
 
     class CuteDSLNVFP4BlackwellRunner(TunableRunner):
@@ -499,8 +499,8 @@ def inputs_pre_hook_finalize_fusion(
                 device=num_non_exiting_tiles.device)
             return a, b, a_sf, b_sf, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, num_non_exiting_tiles, token_final_scales
 
-    class Sm100BlockScaledPersistentGroupedGemmRunner(TunableRunner):
-        kernel_class = Sm100BlockScaledPersistentGroupedGemmKernel
+    class Sm100BlockScaledContiguousGroupedGemmRunner(TunableRunner):
+        kernel_class = Sm100BlockScaledContiguousGroupedGemmKernel
         kernel_cache = dict()
         tuning_config_cache = dict()
 
@@ -730,7 +730,7 @@ def cute_dsl_nvfp4_grouped_gemm_blackwell(
     ) -> torch.Tensor:
         tuner = AutoTuner.get()
 
-        runner = Sm100BlockScaledPersistentGroupedGemmRunner(
+        runner = Sm100BlockScaledContiguousGroupedGemmRunner(
             num_experts, top_k, num_local_experts, local_expert_offset,
             tile_size, output_dtype, scaling_vector_size)
         inputs = [
@@ -769,9 +769,9 @@ def _(
         n = weight.size(1)
         return torch.empty(m, n, dtype=output_dtype, device=input.device)
 
-    class Sm100BlockScaledPersistentGroupedGemmFinalizeFusionRunner(
+    class Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner(
             TunableRunner):
-        kernel_class = Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel
+        kernel_class = Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel
         kernel_cache = dict()
         tuning_config_cache = dict()
 
@@ -1038,7 +1038,7 @@ def cute_dsl_nvfp4_grouped_gemm_finalize_blackwell(
     ) -> torch.Tensor:
         tuner = AutoTuner.get()
 
-        runner = Sm100BlockScaledPersistentGroupedGemmFinalizeFusionRunner(
+        runner = Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner(
             num_experts, top_k, num_local_experts, local_expert_offset,
             tile_size, output_dtype, scaling_vector_size)
         inputs = [

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/grouped_blockscaled_gemm_persistent.py renamed to tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm.py

@@ -53,7 +53,7 @@
 from cutlass.cute.nvgpu import cpasync, tcgen05
 
 
-class Sm100BlockScaledPersistentGroupedGemmKernel:
+class Sm100BlockScaledContiguousGroupedGemmKernel:
     """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types
     and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
 
@@ -88,7 +88,7 @@ class Sm100BlockScaledPersistentGroupedGemmKernel:
         - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
 
     Example:
-        >>> gemm = Sm100BlockScaledPersistentGroupedGemmKernel(
+        >>> gemm = Sm100BlockScaledContiguousGroupedGemmKernel(
         ...     sf_vec_size=16, mma_tiler_mn=(256, 128), cluster_shape_mn=(2, 1)
         ... )
         >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, max_active_clusters, stream)
@@ -2138,8 +2138,9 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
             is_valid = False
         return is_valid
 
-    @staticmethod
+    @classmethod
     def can_implement(
+        cls,
         ab_dtype: Type[cutlass.Numeric],
         sf_dtype: Type[cutlass.Numeric],
         sf_vec_size: int,
@@ -2198,24 +2199,22 @@ def can_implement(
         """
         can_implement = True
         # Skip unsupported types
-        if not Sm100BlockScaledPersistentGroupedGemmKernel.is_valid_dtypes_and_scale_factor_vec_size(
+        if not cls.is_valid_dtypes_and_scale_factor_vec_size(
             ab_dtype, sf_dtype, sf_vec_size, acc_dtype, c_dtype
         ):
             can_implement = False
 
         # Skip unsupported layouts
-        if not Sm100BlockScaledPersistentGroupedGemmKernel.is_valid_layouts(
-            ab_dtype, c_dtype, a_major, b_major, c_major
-        ):
+        if not cls.is_valid_layouts(ab_dtype, c_dtype, a_major, b_major, c_major):
             can_implement = False
 
         # Skip invalid mma tile shape and cluster shape
-        if not Sm100BlockScaledPersistentGroupedGemmKernel.is_valid_mma_tiler_and_cluster_shape(
+        if not cls.is_valid_mma_tiler_and_cluster_shape(
             use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, m_aligned
         ):
             can_implement = False
         # Skip illegal problem shape for load/store alignment
-        if not Sm100BlockScaledPersistentGroupedGemmKernel.is_valid_tensor_alignment(
+        if not cls.is_valid_tensor_alignment(
             m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major
         ):
             can_implement = False

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/grouped_blockscaled_gemm_finalize_fusion.py renamed to tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py

@@ -210,7 +210,7 @@ def atomic_add_func(rOut_epi_packed, scatter_out_offset, loc=None, ip=None):
         )
 
 
-class Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel:
+class Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel:
     """This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types
     and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
 
@@ -245,7 +245,7 @@ class Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel:
         - Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
 
     Example:
-        >>> gemm = Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel(
+        >>> gemm = Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel(
         ...     sf_vec_size=16, mma_tiler_mn=(256, 128), cluster_shape_mn=(2, 1)
         ... )
         >>> gemm(
@@ -2147,8 +2147,9 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
             is_valid = False
         return is_valid
 
-    @staticmethod
+    @classmethod
     def can_implement(
+        cls,
         ab_dtype: Type[cutlass.Numeric],
         sf_dtype: Type[cutlass.Numeric],
         sf_vec_size: int,
@@ -2207,24 +2208,22 @@ def can_implement(
         """
         can_implement = True
         # Skip unsupported types
-        if not Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel.is_valid_dtypes_and_scale_factor_vec_size(
+        if not cls.is_valid_dtypes_and_scale_factor_vec_size(
             ab_dtype, sf_dtype, sf_vec_size, acc_dtype, out_dtype
         ):
             can_implement = False
 
         # Skip unsupported layouts
-        if not Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel.is_valid_layouts(
-            ab_dtype, out_dtype, a_major, b_major, c_major
-        ):
+        if not cls.is_valid_layouts(ab_dtype, out_dtype, a_major, b_major, c_major):
             can_implement = False
 
         # Skip invalid mma tile shape and cluster shape
-        if not Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel.is_valid_mma_tiler_and_cluster_shape(
+        if not cls.is_valid_mma_tiler_and_cluster_shape(
             use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, m_aligned
         ):
             can_implement = False
         # Skip illegal problem shape for load/store alignment
-        if not Sm100BlockScaledPersistentGroupedGemmFinalizeFusionKernel.is_valid_tensor_alignment(
+        if not cls.is_valid_tensor_alignment(
             m, n, k, l, ab_dtype, out_dtype, a_major, b_major, c_major
         ):
             can_implement = False

tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py

@@ -1817,8 +1817,9 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
             is_valid = False
         return is_valid
 
-    @staticmethod
+    @classmethod
     def can_implement(
+        cls,
         ab_dtype: Type[cutlass.Numeric],
         sf_dtype: Type[cutlass.Numeric],
         sf_vec_size: int,
@@ -1859,20 +1860,20 @@ def can_implement(
         """
         can_implement = True
         # Skip unsupported types
-        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size(
+        if not cls.is_valid_dtypes_and_scale_factor_vec_size(
                 ab_dtype, sf_dtype, sf_vec_size, c_dtype):
             can_implement = False
         # Skip unsupported layouts
-        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_layouts(
-                ab_dtype, c_dtype, a_major, b_major, c_major):
+        if not cls.is_valid_layouts(ab_dtype, c_dtype, a_major, b_major,
+                                    c_major):
             can_implement = False
         # Skip invalid mma tile shape and cluster shape
-        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape(
-                mma_tiler_mn, cluster_shape_mn):
+        if not cls.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn,
+                                                        cluster_shape_mn):
             can_implement = False
         # Skip illegal problem shape for load/store alignment
-        if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment(
-                m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major):
+        if not cls.is_valid_tensor_alignment(m, n, k, l, ab_dtype, c_dtype,
+                                             a_major, b_major, c_major):
             can_implement = False
         return can_implement