Fix TensorCoreIntrinEmitter.mma base signature regression

Merging the SM120 block-scale mode into TensorCoreIntrinEmitter changed mma()'s 4th/5th positional parameters to SFA_buf/SFB_buf, so the plain gemm lowering's positional mma(A_local, B_local, C_buf, ki) call bound ki to SFA_buf and raised 'Scale buffers require TensorCoreIntrinEmitter block-scale mode' for every non-blockscaled MMA GEMM (CI examples: attention_sink, deepseek_nsa/mhc/v32). Keep the base (A, B, C, k_inner) positional signature, make all scale-related parameters keyword-only, forward k_inner to the base implementation, and drop the unused blockscaled mma_atom override that shadowed the base method the same way. Add a signature-contract regression test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
LeiWang1999 · Jul 30, 2026 · Jul 3, 2026 · Jul 8, 2026 · Jul 8, 2026 · Jul 8, 2026
commit 1a1062da7bcb523e0bcff25f7ef3ff40e84fc465
diff --git a/testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py b/testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py
@@ -82,6 +82,23 @@ def transform_func(i, j):
     return T.Layout(shape, transform_func)
 
 
+def test_tensor_core_intrin_emitter_mma_keeps_base_positional_signature():
+    """The block-scale extension must not shift the base positional signature.
+
+    The non-blockscaled gemm lowering calls ``emitter.mma(A, B, C, ki)``
+    positionally, so every scale-related parameter has to stay keyword-only.
+    """
+    import inspect
+
+    from tilelang.cuda.intrinsics.macro.mma_macro_generator import _TensorCoreIntrinEmitterBase
+
+    base = list(inspect.signature(_TensorCoreIntrinEmitterBase.mma).parameters.values())
+    override = list(inspect.signature(TensorCoreIntrinEmitter.mma).parameters.values())
+    assert [p.name for p in override[: len(base)]] == [p.name for p in base]
+    for extra in override[len(base) :]:
+        assert extra.kind == inspect.Parameter.KEYWORD_ONLY, extra.name
+
+
 def test_sm120_mma_blockscaled_strategy_helpers_are_not_public_api():
     assert not hasattr(T, "sm120_mma_blockscaled")
     assert not hasattr(T, "sm120_mma_blockscaled_kblock_fulltile")

diff --git a/tilelang/cuda/intrinsics/macro/mma_macro_generator.py b/tilelang/cuda/intrinsics/macro/mma_macro_generator.py
@@ -1762,18 +1762,21 @@ def mma(
         A_local_buf,
         B_local_buf,
         C_local_buf,
+        k_inner: PrimExpr | None = 0,
+        *,
         SFA_buf=None,
         SFB_buf=None,
-        ki: PrimExpr = 0,
         k_start: PrimExpr = 0,
         sf_a_granularity_k: int | None = None,
         sf_b_granularity_k: int | None = None,
         sf_layout: str = "rowmajor",
     ):
+        # Keep the base-class positional signature (A, B, C, k_inner): the
+        # non-blockscaled gemm lowering calls mma(A_local, B_local, C_buf, ki).
         if not self.is_blockscaled:
             if SFA_buf is not None or SFB_buf is not None:
                 raise ValueError("Scale buffers require TensorCoreIntrinEmitter block-scale mode")
-            return super().mma(A_local_buf, B_local_buf, C_local_buf)
+            return super().mma(A_local_buf, B_local_buf, C_local_buf, k_inner)
         if SFA_buf is None or SFB_buf is None:
             raise ValueError("Block-scaled MMA requires SFA and SFB buffers")
         warp_rows = self.warp_rows
@@ -1795,8 +1798,8 @@ def mma(
         sf_vec_size = self.sf_vec_size
         sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k
         sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k
-        scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k)
-        scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k)
+        scale_a_word_k = self._scale_word_k(k_start, k_inner, sf_a_granularity_k)
+        scale_b_word_k = self._scale_word_k(k_start, k_inner, sf_b_granularity_k)
         thread_binding = self.get_thread_binding()
         SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf)
         SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf)
@@ -3309,103 +3312,3 @@ def _warp_mma_block_scale_full_b_atom_prefetched(
             SFB_local_buf,
             SFB_rep_local_buf,
         )
-
-    def mma_atom(
-        self,
-        A_local_buf,
-        B_local_buf,
-        C_local_buf,
-        SFA_buf,
-        SFB_buf,
-        inst_m_idx: PrimExpr | int,
-        inst_n_idx: PrimExpr | int,
-        ki: PrimExpr = 0,
-        k_start: PrimExpr = 0,
-        sf_a_granularity_k: int | None = None,
-        sf_b_granularity_k: int | None = None,
-    ):
-        local_size_a = self.local_size_a
-        local_size_b = self.local_size_b
-        local_size_out = self.local_size_out
-        kind = self.kind
-        scale_vec_size = self.scale_vec_size
-        stype = self.stype
-        accum_dtype = self.accum_dtype
-        a_dtype_abbrv = self.a_dtype_abbrv
-        b_dtype_abbrv = self.b_dtype_abbrv
-        mma_prefix = self.mma_prefix
-        warp_cols = self.warp_cols
-        warp_row_tiles = self.warp_row_tiles
-        warp_col_tiles = self.warp_col_tiles
-        micro_size_x = self.micro_size_x
-        micro_size_y = self.micro_size_y
-        sf_vec_size = self.sf_vec_size
-        sf_a_granularity_k = sf_vec_size if sf_a_granularity_k is None else sf_a_granularity_k
-        sf_b_granularity_k = sf_vec_size if sf_b_granularity_k is None else sf_b_granularity_k
-        scale_a_word_k = self._scale_word_k(k_start, ki, sf_a_granularity_k)
-        scale_b_word_k = self._scale_word_k(k_start, ki, sf_b_granularity_k)
-        thread_binding = self.get_thread_binding()
-        SFA_data, SFA_other, SFA_base_m, SFA_base_k = self._scale_region_parts(SFA_buf)
-        SFB_data, SFB_other, SFB_base_n, SFB_base_k = self._scale_region_parts(SFB_buf)
-        replicate_b = self.n_dim == 16
-
-        @T.macro
-        def _warp_mma_block_scale_atom(A_local_buf, B_local_buf, C_local_buf, SFA_data, SFB_data, thread_binding):
-            tx, warp_n, warp_m = self.extract_thread_binding(thread_binding)
-            sfa_row = self._sfa_row_in_atom(tx)
-            sfb_col = self._sfb_col_in_atom(tx)
-            scale_m = warp_m * warp_row_tiles + inst_m_idx * micro_size_x + sfa_row
-            scale_n = warp_n * warp_col_tiles + inst_n_idx * micro_size_y + sfb_col
-            scale_a_ptr = T.access_ptr(
-                SFA_data[tuple(SFA_other) + (SFA_base_m + scale_m, SFA_base_k + scale_a_word_k)],
-                "r",
-            )
-            scale_b_ptr = T.access_ptr(
-                SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n, SFB_base_k + scale_b_word_k)],
-                "r",
-            )
-            T.ptx_mma_block_scale(
-                accum_dtype,
-                mma_prefix,
-                "row",
-                "col",
-                kind,
-                scale_vec_size,
-                a_dtype_abbrv,
-                b_dtype_abbrv,
-                stype,
-                A_local_buf.data,
-                inst_m_idx * local_size_a,
-                B_local_buf.data,
-                0,
-                C_local_buf.data,
-                inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out,
-                scale_a_ptr,
-                scale_b_ptr,
-            )
-            if replicate_b:
-                scale_b_rep_ptr = T.access_ptr(
-                    SFB_data[tuple(SFB_other) + (SFB_base_n + scale_n + 8, SFB_base_k + scale_b_word_k)],
-                    "r",
-                )
-                T.ptx_mma_block_scale(
-                    accum_dtype,
-                    mma_prefix,
-                    "row",
-                    "col",
-                    kind,
-                    scale_vec_size,
-                    a_dtype_abbrv,
-                    b_dtype_abbrv,
-                    stype,
-                    A_local_buf.data,
-                    inst_m_idx * local_size_a,
-                    B_local_buf.data,
-                    lift(local_size_b) // 2,
-                    C_local_buf.data,
-                    inst_m_idx * warp_cols * local_size_out + inst_n_idx * local_size_out + lift(local_size_out) // 2,
-                    scale_a_ptr,
-                    scale_b_rep_ptr,
-                )
-
-        return _warp_mma_block_scale_atom(A_local_buf, B_local_buf, C_local_buf, SFA_data, SFB_data, thread_binding)
diff --git a/tilelang/cuda/op/gemm/gemm_mma.py b/tilelang/cuda/op/gemm/gemm_mma.py
@@ -299,9 +299,9 @@ def _gemm_ss_blockscaled() -> None:
                         A_local,
                         B_local,
                         C_buf,
-                        self.SFARegion,
-                        self.SFBRegion,
-                        ki=ki,
+                        ki,
+                        SFA_buf=self.SFARegion,
+                        SFB_buf=self.SFBRegion,
                         k_start=self.sf_k_start,
                         sf_a_granularity_k=int(sf_a_granularity_k),
                         sf_b_granularity_k=int(sf_b_granularity_k),