[SM120] Simplify NVFP4 lowering internals

LeiWang1999 · Jul 30, 2026 · Jul 3, 2026 · Jul 8, 2026 · Jul 8, 2026 · Jul 8, 2026
commit 2d692571b38afd59afea4cab58d4d101633f25fe
diff --git a/maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py b/maint/gemm/gemm_sm120/benchmark_sm120_nvfp4_blockscaled_gemm.py
@@ -535,7 +535,7 @@ def run_tilelang(args: argparse.Namespace) -> tuple[float, float]:
         kernel.export_sass(str(sass_path))
         kernel.compile_flags = saved_compile_flags
         print(f"TileLang SASS: {sass_path}")
-    if "sm120_mma_blockscaled_kblock_fulltile_package_pingpong" not in source:
+    if "sm120_mma_blockscaled_fulltile" not in source:
         raise RuntimeError("TileLang source did not lower to the SM120 full-tile MMA helper")
 
     if args.input_mode == "ones":

diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc
@@ -3270,9 +3270,7 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
     replacer.register_rule("(scale_b_byte_id)", scale_b_byte_id);
     replacer.register_rule("(scale_b_thread_id)", scale_b_thread_id);
     this->stream << replacer.rewrite(mma_call);
-  } else if (
-      op->op.same_as(
-          tl::sm120_mma_blockscaled_kblock_fulltile_package_pingpong())) {
+  } else if (op->op.same_as(tl::sm120_mma_blockscaled_fulltile())) {
     need_mma_block_scale_instruction_h_ = true;
     // arg 0: C fragment pointer
     // arg 1: C fragment offset
@@ -3293,11 +3291,10 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
     };
 
     this->PrintIndent();
-    this->stream
-        << "tl::sm120_mma_blockscaled_kblock_fulltile_package_pingpong("
-        << "reinterpret_cast<float*>((" << c_ref << ") + (" << c_offset
-        << ")), " << void_ptr_arg(2) << ", " << void_ptr_arg(3) << ", "
-        << uint32_ptr_arg(4) << ", " << uint32_ptr_arg(5) << ");\n";
+    this->stream << "tl::sm120_mma_blockscaled_fulltile("
+                 << "reinterpret_cast<float*>((" << c_ref << ") + (" << c_offset
+                 << ")), " << void_ptr_arg(2) << ", " << void_ptr_arg(3) << ", "
+                 << uint32_ptr_arg(4) << ", " << uint32_ptr_arg(5) << ");\n";
   } else if (op->op.same_as(builtin::ptx_mma_sp())) {
     // arg 0: shape: mXnXkX
     // arg 1: A layout: row/col

diff --git a/src/cuda/op/copy.cc b/src/cuda/op/copy.cc
@@ -1647,21 +1647,12 @@ Stmt Copy::LowerBulk(const CopyNode &op, const LowerArgs &lower_args,
                      "to normal copy";
     return fallback_to_normal("undecodable shared swizzle layout");
   }
-  // Recast the swizzle atom into byte-address space because CuTe/TMA
-  // swizzles are based on byte addresses. For scalar NVFP4, only the swizzle
-  // atom can be recast cleanly from 4-bit logical elements; the plain layout
-  // may contain half-byte contiguous strides and remains in element space for
-  // offset/coordinate analysis below.
+  // Recast element-space layout into byte-address space.
+  // Because CuTe swizzle are based on byte addresses.
   int elem_bits = shared_tensor->dtype.bits();
-  cute::Swizzle swizzle_bytes;
-  if (shared_tensor->dtype.is_float4_e2m1fn() &&
-      shared_tensor->dtype.is_scalar()) {
-    swizzle_bytes = composed.value()->swizzle.Recast(elem_bits, /*new_bits=*/8);
-  } else {
-    // E.g., composed_bytes = Sw<3,4,3> o 0 o (64,64,8):(128,2,8192)
-    swizzle_bytes = composed.value().Recast(elem_bits, /*new_bits=*/8)->swizzle;
-  }
-  const auto *sw = swizzle_bytes.get();
+  // E.g., composed_bytes = Sw<3,4,3> o 0 o (64,64,8):(128,2,8192)
+  auto composed_bytes = composed.value().Recast(elem_bits, /*new_bits=*/8);
+  const auto *sw = composed_bytes->swizzle.get();
   int b_bits = sw->b_bits, m_base = sw->m_base, s_shift = sw->s_shift;
   if (!sw->IsSwizzled()) {
     desc.swizzle = static_cast<int>(CU_TENSOR_MAP_SWIZZLE_NONE);

diff --git a/src/op/builtin.cc b/src/op/builtin.cc
@@ -63,12 +63,6 @@ DataType CuTensorMapType() { return DataType::UInt(8, 128); }
   TVM_REGISTER_OP("tl." #OpName)                                               \
       .set_attr<TScriptPrinterName>("TScriptPrinterName", #OpName)
 
-#define TIR_DEFINE_TL_OPAQUE_BUILTIN(OpName, NumInputs)                        \
-  TIR_DEFINE_TL_BUILTIN(OpName)                                                \
-      .set_num_inputs(NumInputs)                                               \
-      .set_attr<TCallEffectKind>("TCallEffectKind",                            \
-                                 Integer(CallEffectKind::kOpaque))
-
 // Pointer access metadata op (frontend-only, lowered later).
 TIR_DEFINE_TL_BUILTIN(access_ptr)
     .set_num_inputs(3)
@@ -258,10 +252,15 @@ TIR_DEFINE_TL_BUILTIN(ptx_wgmma_sp_rs)
     .set_attr<TCallEffectKind>("TCallEffectKind",
                                Integer(CallEffectKind::kOpaque));
 
-TIR_DEFINE_TL_OPAQUE_BUILTIN(ptx_mma_block_scale, 21);
+TIR_DEFINE_TL_BUILTIN(ptx_mma_block_scale)
+    .set_num_inputs(21)
+    .set_attr<TCallEffectKind>("TCallEffectKind",
+                               Integer(CallEffectKind::kOpaque));
 
-TIR_DEFINE_TL_OPAQUE_BUILTIN(
-    sm120_mma_blockscaled_kblock_fulltile_package_pingpong, 6);
+TIR_DEFINE_TL_BUILTIN(sm120_mma_blockscaled_fulltile)
+    .set_num_inputs(6)
+    .set_attr<TCallEffectKind>("TCallEffectKind",
+                               Integer(CallEffectKind::kOpaque));
 
 TIR_DEFINE_TL_BUILTIN(ptx_tcgen05_mma_ss)
     .set_num_inputs(14)

diff --git a/src/op/builtin.h b/src/op/builtin.h
@@ -447,9 +447,9 @@ TVM_DLL const Op &ptx_wgmma_sp_rs();
 TVM_DLL const Op &ptx_mma_block_scale();
 
 /*!
- * \brief Private SM120a NVF4 full-tile package-lifecycle helper.
+ * \brief Private SM120a NVF4 full-tile MMA helper.
  */
-TVM_DLL const Op &sm120_mma_blockscaled_kblock_fulltile_package_pingpong();
+TVM_DLL const Op &sm120_mma_blockscaled_fulltile();
 
 /*!
  * \brief tvm intrinsic for tcgen05 mma shared-shared instructions.

diff --git a/src/tl_templates/cuda/instruction/mma_block_scale.h b/src/tl_templates/cuda/instruction/mma_block_scale.h
@@ -309,9 +309,10 @@ TL_DEVICE void sm120_gemm_fulltile_ab_owner_wide_package(
 #undef TL_SM120_PKG_C_OFFSET
 }
 
-TL_DEVICE void sm120_mma_blockscaled_kblock_fulltile_package_pingpong(
-    float *c, const void *a_smem_base, const void *b_smem_base,
-    const uint32_t *sfa_smem_base, const uint32_t *sfb_smem_base) {
+TL_DEVICE void sm120_mma_blockscaled_fulltile(float *c, const void *a_smem_base,
+                                              const void *b_smem_base,
+                                              const uint32_t *sfa_smem_base,
+                                              const uint32_t *sfb_smem_base) {
   const char *a_base = static_cast<const char *>(a_smem_base);
   const char *b_base = static_cast<const char *>(b_smem_base);
   SM120FulltileABOwnerWidePackage pkg0;

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
@@ -115,7 +115,7 @@ def test_sm120_mma_blockscaled_strategy_helpers_are_not_public_api():
         T,
         "sm120_mma_blockscaled_kblock_fulltile_afull_bpanel_owner_wide",
     )
-    assert not hasattr(T, "sm120_mma_blockscaled_kblock_fulltile_package_pingpong")
+    assert not hasattr(T, "sm120_mma_blockscaled_fulltile")
     assert not hasattr(T, "sm120_mma_blockscaled_cute_consumer_bridge")
 
 
@@ -617,7 +617,7 @@ def test_nvf4_mma_block_scale_package_pingpong_contract_lowers_fulltile():
     )
 
     src = kernel.get_kernel_source()
-    assert "sm120_mma_blockscaled_kblock_fulltile_package_pingpong" in src
+    assert "sm120_mma_blockscaled_fulltile" in src
 
 
 @tilelang.testing.requires_cuda

diff --git a/tilelang/cuda/intrinsics/macro/mma_macro_generator.py b/tilelang/cuda/intrinsics/macro/mma_macro_generator.py
@@ -2017,7 +2017,7 @@ def _warp_backend_package_pingpong(
             )
             T.call_intrin(
                 "handle",
-                tirx.op.Op.get("tl.sm120_mma_blockscaled_kblock_fulltile_package_pingpong"),
+                tirx.op.Op.get("tl.sm120_mma_blockscaled_fulltile"),
                 C_local_buf.data,
                 0,
                 a_base,