[SM120] Split block-scaled MMA emitter

LeiWang1999 · Jul 30, 2026 · Jul 3, 2026 · Jul 8, 2026 · Jul 8, 2026 · Jul 8, 2026
commit a1b21a7ee7dfae12b6ff71b8963999a53cf99196
diff --git a/maint/gemm/gemm_sm120/correctness_evaluation_nvf4_vs_cutlass.py b/maint/gemm/gemm_sm120/correctness_evaluation_nvf4_vs_cutlass.py
@@ -11,7 +11,7 @@
 import torch
 import tilelang
 import tilelang.language as T
-from tilelang.intrinsics import TensorCoreIntrinEmitter, make_mma_swizzle_layout
+from tilelang.intrinsics import TensorCoreIntrinEmitterSM120, make_mma_swizzle_layout
 from tilelang.transform import simplify_prim_func
 
 
@@ -53,7 +53,7 @@ def _make_tilelang_nvf4_kernel(m: int, n: int, k: int, micro_k: int, sf_vec_size
     warp_rows = warp_row_tiles // micro_size_x
     warp_cols = warp_col_tiles // micro_size_y
 
-    mma_emitter = TensorCoreIntrinEmitter(
+    mma_emitter = TensorCoreIntrinEmitterSM120(
         is_blockscaled=True,
         a_dtype=in_dtype,
         b_dtype=in_dtype,

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
@@ -9,8 +9,8 @@
 import tilelang.testing
 from tilelang import tvm
 from tilelang.cuda.intrinsics.layout.mma_layout import mma_load_a_32x32_to_shared_16x64_layout
-from tilelang.cuda.intrinsics.macro.mma_macro_generator import SM120BlockScaleTile
-from tilelang.intrinsics import TensorCoreIntrinEmitter, get_swizzle_layout
+from tilelang.cuda.intrinsics.macro.mma_sm120_macro_generator import SM120BlockScaleTile
+from tilelang.intrinsics import TensorCoreIntrinEmitterSM120, get_swizzle_layout
 from examples.dequantize_gemm.quantize.nvfp4 import blockscaled_chunk_kmajor_word_offset
 from tilelang.transform import simplify_prim_func
 
@@ -36,7 +36,7 @@
 
 
 def _make_blockscale_emitter(**kwargs):
-    return TensorCoreIntrinEmitter(
+    return TensorCoreIntrinEmitterSM120(
         is_blockscaled=True,
         kind="mxf4nvf4",
         scale_vec_size=4,
@@ -101,15 +101,25 @@ def test_tensor_core_intrin_emitter_mma_keeps_base_positional_signature():
     """
     import inspect
 
-    from tilelang.cuda.intrinsics.macro.mma_macro_generator import _TensorCoreIntrinEmitterBase
+    from tilelang.cuda.intrinsics.macro.mma_macro_generator import TensorCoreIntrinEmitter as MMAIntrinEmitter
 
-    base = list(inspect.signature(_TensorCoreIntrinEmitterBase.mma).parameters.values())
-    override = list(inspect.signature(TensorCoreIntrinEmitter.mma).parameters.values())
+    base = list(inspect.signature(MMAIntrinEmitter.mma).parameters.values())
+    override = list(inspect.signature(TensorCoreIntrinEmitterSM120.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_default_tensor_core_intrin_emitter_stays_arch_neutral():
+    from tilelang.cuda.intrinsics.macro.mma_macro_generator import TensorCoreIntrinEmitter as MMAIntrinEmitter
+    from tilelang.cuda.op.gemm.gemm_mma import GemmMMA
+    from tilelang.intrinsics import TensorCoreIntrinEmitter
+
+    assert TensorCoreIntrinEmitter is MMAIntrinEmitter
+    assert GemmMMA.intrin_emitter_cls is MMAIntrinEmitter
+    assert issubclass(TensorCoreIntrinEmitterSM120, MMAIntrinEmitter)
+
+
 def test_sm120_mma_blockscaled_strategy_helpers_are_not_public_api():
     # NVFP4-specific staging helpers stay out of the general T.* surface; the
     # scale-tile addressing lives in examples.dequantize_gemm.quantize.nvfp4.
@@ -319,8 +329,8 @@ def test_nvf4_mma_block_scale_fragment_layouts_match_cute():
 
 
 def test_nvf4_mma_block_scale_lane_scale_mapping_matches_cute():
-    sfa_rows = [TensorCoreIntrinEmitter._sfa_row_in_atom(tx) for tx in range(32)]
-    sfb_cols = [TensorCoreIntrinEmitter._sfb_col_in_atom(tx) for tx in range(32)]
+    sfa_rows = [TensorCoreIntrinEmitterSM120._sfa_row_in_atom(tx) for tx in range(32)]
+    sfb_cols = [TensorCoreIntrinEmitterSM120._sfb_col_in_atom(tx) for tx in range(32)]
 
     assert sfa_rows == [
         0,
@@ -402,7 +412,7 @@ def test_nvf4_mma_block_scale_lane_scale_mapping_matches_cute():
 )
 def test_nvf4_mma_block_scale_rejects_unsupported_configs(kwargs):
     with pytest.raises(ValueError, match="Unsupported SM120 block-scale MMA config"):
-        TensorCoreIntrinEmitter(
+        TensorCoreIntrinEmitterSM120(
             is_blockscaled=True,
             a_dtype=T.float4_e2m1fn,
             b_dtype=T.float4_e2m1fn,
@@ -428,7 +438,7 @@ def test_nvf4_mma_block_scale_rejects_unsupported_configs(kwargs):
 )
 def test_nvf4_mma_block_scale_rejects_incompatible_dtypes(dtype_kwargs):
     with pytest.raises(ValueError, match="mxf4nvf4 expects"):
-        TensorCoreIntrinEmitter(
+        TensorCoreIntrinEmitterSM120(
             is_blockscaled=True,
             a_transposed=False,
             b_transposed=True,
@@ -766,10 +776,6 @@ def test_nvf4_mma_block_scale_varying_scale_correctness():
     torch.testing.assert_close(C, ref, rtol=0.0, atol=0.0)
 
 
-if __name__ == "__main__":
-    tilelang.testing.main()
-
-
 # ---------------------------------------------------------------------------
 # Example tail-tile behavior (moved from test_tilelang_sm120_nvfp4_example_cli).
 # ---------------------------------------------------------------------------
@@ -823,3 +829,7 @@ def test_sm120_nvfp4_example_kernel_rejects_unsupported_tails(monkeypatch):
     # bf16 output rows must stay 16-byte aligned
     with pytest.raises(AssertionError, match="multiple of 8"):
         module.sm120_nvfp4_blockscaled_gemm(128, 130, 256)
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()
diff --git a/testing/python/quantize/test_tilelang_quantize_nvfp4_scale_layout.py b/testing/python/quantize/test_tilelang_quantize_nvfp4_scale_layout.py
@@ -793,7 +793,7 @@ def test_compact_package_issue_table_matches_cpp_macro_issue_order():
 
 
 def test_frontend_package_macro_keeps_compact_selector_issue_order():
-    macro = Path(__file__).resolve().parents[3] / "tilelang/cuda/intrinsics/macro/mma_macro_generator.py"
+    macro = Path(__file__).resolve().parents[3] / "tilelang/cuda/intrinsics/macro/mma_sm120_macro_generator.py"
     source = macro.read_text()
     issue_loop = source[source.index("        def _mma_kblock") : source.index("        def _warp_mma_blockscaled_fulltile")]
 
@@ -984,3 +984,7 @@ def test_swizzle_blockscaled_chunk_kmajor_pads_rows_to_full_tiles():
     back = unswizzle_blockscaled_chunk_kmajor_scale_words(swizzled)
     assert torch.equal(back[:rows], words)
     assert bool((back[rows:] == 0).all())
+
+
+if __name__ == "__main__":
+    tilelang_testing.main()
diff --git a/tilelang/cuda/intrinsics/__init__.py b/tilelang/cuda/intrinsics/__init__.py
@@ -8,6 +8,7 @@
     TensorCoreIntrinEmitter,
     TensorCoreIntrinEmitterWithLadderTransform,
 )
+from .macro.mma_sm120_macro_generator import TensorCoreIntrinEmitterSM120  # noqa: F401
 
 from .layout.mma_layout import get_swizzle_layout  # noqa: F401
 from .layout.mma_layout import make_mma_swizzle_layout  # noqa: F401
diff --git a/tilelang/cuda/intrinsics/macro/__init__.py b/tilelang/cuda/intrinsics/macro/__init__.py
@@ -2,5 +2,6 @@
     TensorCoreIntrinEmitter,
     TensorCoreIntrinEmitterWithLadderTransform,
 )
+from .mma_sm120_macro_generator import TensorCoreIntrinEmitterSM120  # noqa: F401
 
 from .mma_sp_macro_generator import SparseTensorCoreIntrinEmitter  # noqa: F401