lint

LeiWang1999 · Mar 24, 2026 · Feb 26, 2026 · Feb 26, 2026 · Mar 2, 2026 · Mar 3, 2026
commit dcc1bae4f9f48ab56ad3be50c1215b673cea5058
diff --git a/examples/gemm_sm100/gemm_tcgen5mma_2sm.py b/examples/gemm_sm100/gemm_tcgen5mma_2sm.py
@@ -2,6 +2,7 @@
 import tilelang
 import tilelang.language as T
 from tilelang.engine import register_cuda_postproc
+
 tilelang.disable_cache()
 
 
@@ -130,7 +131,7 @@ def main(
 
 
 M, N, K = 2048, 2048, 2048  # FIXME: buggy when size is lager
-print(f'M: {M}, N: {N}, K: {K}')
+print(f"M: {M}, N: {N}, K: {K}")
 block_M, block_N, block_K = 128, 256, 64
 in_dtype, out_dtype, accum_dtype = T.bfloat16, T.bfloat16, T.float
 num_stages = 0 if block_N >= 256 or block_M >= 256 or block_K >= 256 else 2
@@ -156,7 +157,7 @@ def main(
     c = jit_kernel(a, b)
     ref_c = (a @ b).to(torch.bfloat16)
     torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
-print('ALL CHECK PASSED. ✅')
+print("ALL CHECK PASSED. ✅")
 profiler = jit_kernel.get_profiler()
 latency = profiler.do_bench()
 print(f"Latency: {latency} ms")

diff --git a/examples/gemm_sm100/gemm_tcgen5mma_ws_2sm.py b/examples/gemm_sm100/gemm_tcgen5mma_ws_2sm.py
@@ -5,6 +5,7 @@
 import tilelang.language as T
 from tilelang.profiler import do_bench
 from tilelang.engine import register_cuda_postproc
+
 tilelang.disable_cache()
 
 
@@ -151,7 +152,10 @@ def gemm(A, B, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, num_
             for k in T.serial(k_iters):
                 T.mbarrier_wait_parity(consumed[k % num_stages], ((k // num_stages) & 1) ^ 1)
                 T.copy(A[bx * block_M : (bx + 1) * block_M, k * block_K : (k + 1) * block_K], A_shared[k % num_stages, :, :])
-                T.copy(B[k * block_K : (k + 1) * block_K, (by * 2 + cta_id) * (block_N // 2) : (by * 2 + cta_id + 1) * (block_N // 2)], B_shared[k % num_stages, :, :])
+                T.copy(
+                    B[k * block_K : (k + 1) * block_K, (by * 2 + cta_id) * (block_N // 2) : (by * 2 + cta_id + 1) * (block_N // 2)],
+                    B_shared[k % num_stages, :, :],
+                )
                 T.mbarrier_arrive(loaded[k % num_stages], 0)  # arrive on leader cta's barrier
         elif cta_id == 0 and tx < 64:  # Only warp 1 on leader cta issues tcgen5
             for k in T.serial(k_iters):
@@ -190,13 +194,11 @@ def main():
     b = torch.randn(K, N, device="cuda", dtype=torch.bfloat16)
     print(gemm.get_kernel_source(a, b, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, num_stages))
     c = gemm(a, b, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, num_stages)
-
 
     ref_c = (a.to(torch.float) @ b.to(torch.float)).to(torch.bfloat16)
-    # torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
     print("All checks passed. ✅")
 
-
     tl_latency = do_bench(lambda: gemm(a, b, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, num_stages), backend="cupti")
     torch_latency = do_bench(lambda: a @ b, backend="cupti")
     print(f"Tilelang latency: {tl_latency} ms")

diff --git a/src/op/builtin.h b/src/op/builtin.h
@@ -62,7 +62,8 @@ static constexpr const char *kEnableAsyncCopy = "tl.enable_async_copy";
 static constexpr const char *kEnableVectorizePlannerVerbose =
     "tl.enable_vectorize_planner_verbose";
 static constexpr const char *kDisableWGMMA = "tl.disable_wgmma";
-static constexpr const char *kDisable2CTATcgen5MMA = "tl.disable_2cta_tcgen5mma";
+static constexpr const char *kDisable2CTATcgen5MMA =
+    "tl.disable_2cta_tcgen5mma";
 static constexpr const char *kDisableShuffleElect = "tl.disable_shuffle_elect";
 static constexpr const char *kDisableLoopUnswitching =
     "tl.disable_loop_unswitching";

diff --git a/src/op/gemm_py.cc b/src/op/gemm_py.cc
@@ -318,9 +318,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def(
-      "tl.get_tcgen5_mma_meta",
-      [](int M, int N, int K, DataType ab_dtype, DataType c_dtype, bool disable_2cta) {
-        auto [success, meta] = GetTCGEN5MMAMeta(M, N, K, ab_dtype, c_dtype, disable_2cta);
+      "tl.get_tcgen5_mma_meta", [](int M, int N, int K, DataType ab_dtype,
+                                   DataType c_dtype, bool disable_2cta) {
+        auto [success, meta] =
+            GetTCGEN5MMAMeta(M, N, K, ab_dtype, c_dtype, disable_2cta);
         Array<Integer> result;
         if (success) {
           result.push_back(Integer(meta.atom_m));

diff --git a/src/op/tcgen5_meta.h b/src/op/tcgen5_meta.h
@@ -19,9 +19,10 @@ struct TCGEN5MMAMeta {
 };
 
 inline std::pair<bool, TCGEN5MMAMeta>
-GetTCGEN5MMAMeta(int M, int N, int K, DataType ab_dtype, DataType c_dtype, bool disable_2cta = false) {
+GetTCGEN5MMAMeta(int M, int N, int K, DataType ab_dtype, DataType c_dtype,
+                 bool disable_2cta = false) {
 // TODO (lei) Currently not all shapes / dtypes are supported for TCGEN5MMA.
-// NOTE: In case that other tcgen5mma in the same kernel must use 1-cta, 
+// NOTE: In case that other tcgen5mma in the same kernel must use 1-cta,
 // we should disable 2cta for the current tcgen5mma.
 #define FAIL                                                                   \
   return {                                                                     \

diff --git a/src/target/codegen_cuda.cc b/src/target/codegen_cuda.cc
@@ -436,7 +436,8 @@ class ClusterInfoExtractor : public tir::StmtVisitor {
       cluster_grid_x_ext = cluster_dims[0].as<IntImmNode>()->value;
       cluster_grid_y_ext = cluster_dims[1].as<IntImmNode>()->value;
       cluster_grid_z_ext = cluster_dims[2].as<IntImmNode>()->value;
-      ICHECK(cluster_grid_x_ext > 0 && cluster_grid_y_ext > 0 && cluster_grid_z_ext > 0);
+      ICHECK(cluster_grid_x_ext > 0 && cluster_grid_y_ext > 0 &&
+             cluster_grid_z_ext > 0);
     }
     StmtVisitor::VisitStmt(f->body);
   }
@@ -447,10 +448,12 @@ class ClusterInfoExtractor : public tir::StmtVisitor {
   int64_t cluster_grid_z_ext = 1;
 
 public:
-  std::optional<std::tuple<int64_t, int64_t, int64_t>> extract(const PrimFunc &f) {
+  std::optional<std::tuple<int64_t, int64_t, int64_t>>
+  extract(const PrimFunc &f) {
     this->VisitStmt(f);
     if (launch_with_cluster) {
-      return std::make_tuple(cluster_grid_x_ext, cluster_grid_y_ext, cluster_grid_z_ext);
+      return std::make_tuple(cluster_grid_x_ext, cluster_grid_y_ext,
+                             cluster_grid_z_ext);
     }
     return std::nullopt;
   }
@@ -1895,16 +1898,16 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
   } else if (op->op.same_as(tl::ptx_init_tensor_memory())) {
     std::ostringstream ss;
     ss << "tl::tmem_allocate";
-    if (op->annotations.find("use_2cta") != op->annotations.end()
-        && Downcast<Bool>(op->annotations["use_2cta"])->value) {
+    if (op->annotations.find("use_2cta") != op->annotations.end() &&
+        Downcast<Bool>(op->annotations["use_2cta"])->value) {
       ss << "<true>";
     }
     print_extern_call_stmt(ss.str());
   } else if (op->op.same_as(tl::ptx_deallocate_tensor_memory())) {
     std::ostringstream ss;
     ss << "tl::tmem_deallocate";
-    if (op->annotations.find("use_2cta") != op->annotations.end()
-        && Downcast<Bool>(op->annotations["use_2cta"])->value) {
+    if (op->annotations.find("use_2cta") != op->annotations.end() &&
+        Downcast<Bool>(op->annotations["use_2cta"])->value) {
       ss << "<true>";
     }
     print_extern_call_stmt(ss.str());
@@ -1917,9 +1920,11 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
         this->eviction_policy_names_
             [op->args[op->args.size() - 1].as<IntImmNode>()->value];
     // Simplify the code by using the default eviction policy
-    if (op->annotations.find("use_2cta") != op->annotations.end() && Downcast<Bool>(op->annotations["use_2cta"])->value) {
+    if (op->annotations.find("use_2cta") != op->annotations.end() &&
+        Downcast<Bool>(op->annotations["use_2cta"])->value) {
       if (eviction_policy != "EVICT_NORMAL") {
-        ss << "tl::tma_load_2sm<tl::CacheHintSm100::" << eviction_policy << ">(";
+        ss << "tl::tma_load_2sm<tl::CacheHintSm100::" << eviction_policy
+           << ">(";
       } else {
         ss << "tl::tma_load_2sm(";
       }
@@ -2473,15 +2478,13 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
     bool enable_2cta = Downcast<Bool>(op->args[14])->value;
 
     auto dtype_enum = tl::codegen::ptx::DTypeFromString(kind_dtype);
-    std::string ab_type_str =
-        tl::codegen::ptx::DTypeEnumToString(dtype_enum);
+    std::string ab_type_str = tl::codegen::ptx::DTypeEnumToString(dtype_enum);
 
-    // Currently tcgen05mma_ss<Float16|BFloat16> has use_2cta template param; 
+    // Currently tcgen05mma_ss<Float16|BFloat16> has use_2cta template param;
     // others don't. tcgen05mma_ws_ss has no use_2cta.
     std::string use_2cta_suffix;
-    if (!enable_ws &&
-        (dtype_enum == tl::codegen::ptx::DataType::kFloat16 ||
-         dtype_enum == tl::codegen::ptx::DataType::kBFloat16)) {
+    if (!enable_ws && (dtype_enum == tl::codegen::ptx::DataType::kFloat16 ||
+                       dtype_enum == tl::codegen::ptx::DataType::kBFloat16)) {
       use_2cta_suffix = std::string(", ") + (enable_2cta ? "true" : "false");
     }
 
@@ -2564,8 +2567,8 @@ void CodeGenTileLangCUDA::VisitExpr_(const CallNode *op, std::ostream &os) {
     need_tcgen05_common_h_ = true;
     std::ostringstream ss;
     ss << "tl::tcgen05_mma_arrive";
-    if (op->annotations.find("use_2cta") != op->annotations.end()
-        && Downcast<Bool>(op->annotations["use_2cta"])->value) {
+    if (op->annotations.find("use_2cta") != op->annotations.end() &&
+        Downcast<Bool>(op->annotations["use_2cta"])->value) {
       ss << "<true>";
     }
     print_extern_call_stmt(ss.str());
@@ -3378,25 +3381,31 @@ void CodeGenTileLangCUDA::VisitStmt_(const AttrStmtNode *op) {
     std::string func_name;
     int panel_size = 0;
     if (const auto *call = op->value.as<CallNode>()) {
-      if (call->op.same_as(tir::builtin::tvm_tuple()) && call->args.size() >= 2) {
+      if (call->op.same_as(tir::builtin::tvm_tuple()) &&
+          call->args.size() >= 2) {
         const auto *name_node = call->args[0].as<StringImmNode>();
         const auto *size_node = call->args[1].as<IntImmNode>();
-        ICHECK(name_node && size_node)
-            << "threadblock_swizzle_pattern expects tvm_tuple(device_func, panel_size)";
+        ICHECK(name_node && size_node) << "threadblock_swizzle_pattern expects "
+                                          "tvm_tuple(device_func, panel_size)";
         func_name = name_node->value;
         panel_size = static_cast<int>(size_node->value);
       }
     }
     ICHECK(!func_name.empty() && panel_size > 0);
     if (this->cluster_dims.has_value()) {
-      auto [cluster_grid_x_ext, cluster_grid_y_ext, cluster_grid_z_ext] = this->cluster_dims.value();
+      auto [cluster_grid_x_ext, cluster_grid_y_ext, cluster_grid_z_ext] =
+          this->cluster_dims.value();
       ICHECK(cluster_grid_y_ext == 1 && cluster_grid_z_ext == 1)
-        << "Only support annotate threadblock swizzle for cluster on X dimension for now!";
-      ICHECK(panel_size % cluster_grid_x_ext == 0) << "panel_size must be divisible by clusterDim.x";
-      this->stream << "const dim3 blockIdx = tl::" << func_name << "WithCluster<" 
-        << panel_size / cluster_grid_x_ext << ", " << cluster_grid_x_ext << ">();\n";
+          << "Only support annotate threadblock swizzle for cluster on X "
+             "dimension for now!";
+      ICHECK(panel_size % cluster_grid_x_ext == 0)
+          << "panel_size must be divisible by clusterDim.x";
+      this->stream << "const dim3 blockIdx = tl::" << func_name
+                   << "WithCluster<" << panel_size / cluster_grid_x_ext << ", "
+                   << cluster_grid_x_ext << ">();\n";
     } else {
-      this->stream << "const dim3 blockIdx = tl::" << func_name << "<" << panel_size << ">();\n";
+      this->stream << "const dim3 blockIdx = tl::" << func_name << "<"
+                   << panel_size << ">();\n";
     }
     this->VisitStmt(op->body);
     return;

diff --git a/src/tl_templates/cuda/cluster.h b/src/tl_templates/cuda/cluster.h
@@ -3,8 +3,9 @@
 #include "common.h"
 
 // Config
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && \
-  ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8))))
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) &&                       \
+     ((__CUDACC_VER_MAJOR__ >= 12) ||                                          \
+      ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8))))
 #define TILELANG_CLUSTER_ENABLED
 #endif
 

diff --git a/src/tl_templates/cuda/common.h b/src/tl_templates/cuda/common.h
@@ -1,9 +1,9 @@
 #pragma once
 
 #ifndef __CUDACC_RTC__
-#include <cuda_runtime.h>
 #include <cstdio>
 #include <cstdlib>
+#include <cuda_runtime.h>
 #endif
 
 #include "atomic.h"
@@ -60,14 +60,22 @@ using int4_t = int4;
   } while (0)
 
 #if defined(__CUDA_ARCH__)
-#define TILELANG_UNREACHABLE(msg) \
-  do { printf("%s, %s:%d\n", msg, __FILE__, __LINE__); __trap(); } while(0)
+#define TILELANG_UNREACHABLE(msg)                                              \
+  do {                                                                         \
+    printf("%s, %s:%d\n", msg, __FILE__, __LINE__);                            \
+    __trap();                                                                  \
+  } while (0)
 #elif defined(__CUDACC_RTC__)
-#define TILELANG_UNREACHABLE(msg) \
-  do { __builtin_trap(); } while(0)
+#define TILELANG_UNREACHABLE(msg)                                              \
+  do {                                                                         \
+    __builtin_trap();                                                          \
+  } while (0)
 #else
-#define TILELANG_UNREACHABLE(msg) \
-  do { fprintf(stderr, "%s, %s:%d\n", msg, __FILE__, __LINE__); abort(); } while(0)
+#define TILELANG_UNREACHABLE(msg)                                              \
+  do {                                                                         \
+    fprintf(stderr, "%s, %s:%d\n", msg, __FILE__, __LINE__);                   \
+    abort();                                                                   \
+  } while (0)
 #endif
 
 // using cutlass abs function for half_t

diff --git a/src/tl_templates/cuda/copy_sm100.h b/src/tl_templates/cuda/copy_sm100.h
@@ -4,12 +4,12 @@
 #include <cuda.h>
 #endif
 
+#include "barrier.h"
+#include "common.h"
 #include "cuda_fp8.h"
 #include "tcgen_05.h"
 #include "tcgen_05_ld.h"
 #include "tcgen_05_st.h"
-#include "barrier.h"
-#include "common.h"
 
 namespace tl {
 
@@ -302,8 +302,8 @@ constexpr uint32_t Sm100MmaPeerBitMask = 0xFEFFFFFF;
 template <CacheHintSm100 cache_hint = CacheHintSm100::EVICT_NORMAL,
           typename BarrierType = uint64_t>
 TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
-                            BarrierType &smem_mbar,
-                            void const *const smem_ptr, int32_t const &crd0) {
+                            BarrierType &smem_mbar, void const *const smem_ptr,
+                            int32_t const &crd0) {
   uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
   // Executed by both CTAs. Set peer bit to 0 so that the
   // transaction bytes will update CTA0's barrier.
@@ -327,9 +327,8 @@ TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
 template <CacheHintSm100 cache_hint = CacheHintSm100::EVICT_NORMAL,
           typename BarrierType = uint64_t>
 TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
-                            BarrierType &smem_mbar,
-                            void const *const smem_ptr, int32_t const &crd0,
-                            int32_t const &crd1) {
+                            BarrierType &smem_mbar, void const *const smem_ptr,
+                            int32_t const &crd0, int32_t const &crd1) {
   uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
   // Executed by both CTAs. Set peer bit to 0 so that the
   // transaction bytes will update CTA0's barrier.
@@ -353,9 +352,9 @@ TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
 template <CacheHintSm100 cache_hint = CacheHintSm100::EVICT_NORMAL,
           typename BarrierType = uint64_t>
 TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
-                            BarrierType &smem_mbar,
-                            void const *const smem_ptr, int32_t const &crd0,
-                            int32_t const &crd1, int32_t const &crd2) {
+                            BarrierType &smem_mbar, void const *const smem_ptr,
+                            int32_t const &crd0, int32_t const &crd1,
+                            int32_t const &crd2) {
   uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
   // Executed by both CTAs. Set peer bit to 0 so that the
   // transaction bytes will update CTA0's barrier.
@@ -379,10 +378,9 @@ TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
 template <CacheHintSm100 cache_hint = CacheHintSm100::EVICT_NORMAL,
           typename BarrierType = uint64_t>
 TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
-                            BarrierType &smem_mbar,
-                            void const *const smem_ptr, int32_t const &crd0,
-                            int32_t const &crd1, int32_t const &crd2,
-                            int32_t const &crd3) {
+                            BarrierType &smem_mbar, void const *const smem_ptr,
+                            int32_t const &crd0, int32_t const &crd1,
+                            int32_t const &crd2, int32_t const &crd3) {
   uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
   // Executed by both CTAs. Set peer bit to 0 so that the
   // transaction bytes will update CTA0's barrier.
@@ -406,10 +404,10 @@ TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
 template <CacheHintSm100 cache_hint = CacheHintSm100::EVICT_NORMAL,
           typename BarrierType = uint64_t>
 TL_DEVICE void tma_load_2sm(const CUtensorMap &descriptor,
-                            BarrierType &smem_mbar,
-                            void const *const smem_ptr, int32_t const &crd0,
-                            int32_t const &crd1, int32_t const &crd2,
-                            int32_t const &crd3, int32_t const &crd4) {
+                            BarrierType &smem_mbar, void const *const smem_ptr,
+                            int32_t const &crd0, int32_t const &crd1,
+                            int32_t const &crd2, int32_t const &crd3,
+                            int32_t const &crd4) {
   uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor);
   // Executed by both CTAs. Set peer bit to 0 so that the
   // transaction bytes will update CTA0's barrier.