Commit

Author: bkryu

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py

@@ -89,8 +89,8 @@
     st_global_u64,
     scatter_add_bf16x2,
 )
-from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120 import (
-    Sm120BlockScaledDenseGemmKernel as DenseGemmKernel,
+from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import (
+    Sm120B12xBlockScaledDenseGemmKernel as DenseGemmKernel,
 )
 
 

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_micro_kernel.py

@@ -122,8 +122,8 @@
     st_global_u64,
     scatter_add_bf16x2,
 )
-from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120 import (
-    Sm120BlockScaledDenseGemmKernel as DenseGemmKernel,
+from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import (
+    Sm120B12xBlockScaledDenseGemmKernel as DenseGemmKernel,
 )
 
 

flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_static_kernel.py

@@ -120,8 +120,8 @@
     st_global_u64,
     scatter_add_bf16x2,
 )
-from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120 import (
-    Sm120BlockScaledDenseGemmKernel as DenseGemmKernel,
+from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import (
+    Sm120B12xBlockScaledDenseGemmKernel as DenseGemmKernel,
 )
 
 

flashinfer/gemm/__init__.py

@@ -61,11 +61,11 @@
     from flashinfer.cute_dsl.utils import is_cute_dsl_available
 
     if is_cute_dsl_available():
-        from .kernels.dense_blockscaled_gemm_sm120 import (
-            Sm120BlockScaledDenseGemmKernel as Sm120BlockScaledDenseGemmKernel,
+        from .kernels.dense_blockscaled_gemm_sm120_b12x import (
+            Sm120B12xBlockScaledDenseGemmKernel as Sm120B12xBlockScaledDenseGemmKernel,
         )
 
-        _cute_dsl_kernels.append("Sm120BlockScaledDenseGemmKernel")
+        _cute_dsl_kernels.append("Sm120B12xBlockScaledDenseGemmKernel")
 except ImportError:
     pass
 

flashinfer/gemm/gemm_base.py

@@ -4858,8 +4858,8 @@ def _b12x_gemm_fp4_runner(
     """
     import cutlass
 
-    from .kernels.dense_blockscaled_gemm_sm120 import (
-        Sm120BlockScaledDenseGemmKernel,
+    from .kernels.dense_blockscaled_gemm_sm120_b12x import (
+        Sm120B12xBlockScaledDenseGemmKernel,
     )
 
     cutlass_dtype_attr = _TORCH_TO_CUTLASS_DTYPE_ATTR.get(out_dtype)
@@ -4905,7 +4905,7 @@ def get_valid_tactics(
             ]
             swap_ab = False
             for mma_tiler_mn in sm120_mma_tiler_candidates:
-                if not Sm120BlockScaledDenseGemmKernel.can_implement(
+                if not Sm120B12xBlockScaledDenseGemmKernel.can_implement(
                     ab_dtype,
                     sf_dtype,
                     sf_vec_size,
@@ -4945,11 +4945,10 @@ def forward(
             batch_size = 1
 
             if tactic is None or tactic == -1:
-                _sm_count = torch.cuda.get_device_properties(
-                    a.device
-                ).multi_processor_count
                 tactic = (
-                    _select_default_sm120_mma_tiler(m, n, _sm_count),
+                    _select_default_sm120_mma_tiler(
+                        m, n, get_device_sm_count(a.device)
+                    ),
                     (1, 1),
                     False,
                     False,
@@ -4987,7 +4986,7 @@ def forward(
                 out_dtype,
             )
 
-            make_kernel = lambda: Sm120BlockScaledDenseGemmKernel(
+            make_kernel = lambda: Sm120B12xBlockScaledDenseGemmKernel(
                 sf_vec_size,
                 mma_tiler_mn,
                 cluster_shape_mn,

…/kernels/dense_blockscaled_gemm_sm120.py …els/dense_blockscaled_gemm_sm120_b12x.pyflashinfer/gemm/kernels/dense_blockscaled_gemm_sm120.py renamed to flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120.py renamed to flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py

@@ -1550,7 +1550,7 @@ def wrapper(
 
 
 # Alias for FlashInfer integration
-Sm120BlockScaledDenseGemmKernel = DenseGemmKernel
+Sm120B12xBlockScaledDenseGemmKernel = DenseGemmKernel
 
 
 class _DenseGemmLaunch: