fix

yyttt6 · yyttt6 · commit 4dcdeefec185 · 2025-11-12T11:08:03.000Z
diff --git a/examples/analyze/bench_example_analyze.py b/examples/analyze/bench_example_analyze.py
@@ -11,9 +11,5 @@ def bench_example_conv_analyze():
     tilelang.tools.bench.process_func(example_conv_analyze.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/attention_sink/bench_example_attention_sink.py b/examples/attention_sink/bench_example_attention_sink.py
@@ -48,9 +48,5 @@ def bench_example_gqa_sink_bwd_bhsd_sliding_window():
     tilelang.tools.bench.process_func(example_gqa_sink_bwd_bhsd.main, window_size=128)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/blocksparse_attention/bench_example_blocksparse_attention.py b/examples/blocksparse_attention/bench_example_blocksparse_attention.py
@@ -31,9 +31,5 @@ def bench_example_triton_sparse_gqa_decode_varlen_mask():
     tilelang.tools.bench.process_func(example_triton_sparse_gqa_decode_varlen_mask.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/blocksparse_gemm/bench_example_blocksparse_gemm.py b/examples/blocksparse_gemm/bench_example_blocksparse_gemm.py
@@ -6,9 +6,5 @@ def bench_example_blocksparse_gemm():
     tilelang.tools.bench.process_func(example_blocksparse_gemm.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/cast/bench_example_cast.py b/examples/cast/bench_example_cast.py
@@ -11,9 +11,5 @@ def bench_example_per_token_cast_to_fp8():
     tilelang.tools.bench.process_func(example_per_token_cast_to_fp8.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/convolution/bench_example_convolution.py b/examples/convolution/bench_example_convolution.py
@@ -11,9 +11,5 @@ def bench_example_convolution_autotune():
     tilelang.tools.bench.process_func(example_convolution_autotune.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/deepseek_deepgemm/bench_example_deepgemm_fp8_2xAcc.py b/examples/deepseek_deepgemm/bench_example_deepgemm_fp8_2xAcc.py
@@ -6,9 +6,5 @@ def bench_example_deepgemm_fp8_2xAcc():
     tilelang.tools.bench.process_func(example_deepgemm_fp8_2xAcc.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/deepseek_mla/bench_example_mla_decode.py b/examples/deepseek_mla/bench_example_mla_decode.py
@@ -6,9 +6,5 @@ def bench_example_mla_decode():
     tilelang.tools.bench.process_func(example_mla_decode.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/deepseek_nsa/bench_example_tilelang_nsa.py b/examples/deepseek_nsa/bench_example_tilelang_nsa.py
@@ -11,9 +11,5 @@ def bench_example_tilelang_nsa_fwd_decode():
     tilelang.tools.bench.process_func(example_tilelang_nsa_decode.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/deepseek_v32/bench_tilelang_example_deepseek_v32.py b/examples/deepseek_v32/bench_tilelang_example_deepseek_v32.py
@@ -26,9 +26,5 @@ def bench_sparse_mla_bwd():
     tilelang.tools.bench.process_func(sparse_mla_bwd.test_sparse_mla_bwd)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/dequantize_gemm/bench_example_dequantize_gemm.py b/examples/dequantize_gemm/bench_example_dequantize_gemm.py
@@ -31,9 +31,5 @@ def bench_example_dequant_gemm_w4a8():
     tilelang.tools.bench.process_func(example_dequant_gemm_w4a8.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/dynamic_shape/bench_example_dynamic.py b/examples/dynamic_shape/bench_example_dynamic.py
@@ -6,9 +6,5 @@ def bench_example_dynamic():
     tilelang.tools.bench.process_func(example_dynamic.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/elementwise/bench_example_elementwise.py b/examples/elementwise/bench_example_elementwise.py
@@ -6,9 +6,5 @@ def bench_example_elementwise_add():
     tilelang.tools.bench.process_func(example_elementwise_add.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/flash_attention/bench_example_flash_attention.py b/examples/flash_attention/bench_example_flash_attention.py
@@ -66,9 +66,5 @@ def bench_example_mha_fwd_varlen():
     tilelang.tools.bench.process_func(example_mha_fwd_varlen.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/flash_decoding/bench_example_flash_decoding.py b/examples/flash_decoding/bench_example_flash_decoding.py
@@ -11,9 +11,5 @@ def bench_example_mha_inference():
     tilelang.tools.bench.process_func(example_mha_inference.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/fusedmoe/bench_example_fusedmoe.py b/examples/fusedmoe/bench_example_fusedmoe.py
@@ -6,9 +6,5 @@ def bench_example_fusedmoe_tilelang():
     tilelang.tools.bench.process_func(example_fusedmoe_tilelang.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/gemm/bench_example_gemm.py b/examples/gemm/bench_example_gemm.py
@@ -21,9 +21,5 @@ def bench_example_gemm():
     tilelang.tools.bench.process_func(example_gemm.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/gemm/bench_gemm.py b/examples/gemm/bench_gemm.py
@@ -21,9 +21,5 @@ def bench_example_gemm():
     tilelang.tools.bench.process_func(example_gemm.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/gemm_fp8/bench_example_gemm_fp8.py b/examples/gemm_fp8/bench_example_gemm_fp8.py
@@ -16,9 +16,5 @@ def bench_example_tilelang_gemm_fp8():
     tilelang.tools.bench.process_func(example_tilelang_gemm_fp8.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/gemm_splitk/bench_example_gemm_splitk.py b/examples/gemm_splitk/bench_example_gemm_splitk.py
@@ -11,9 +11,5 @@ def bench_example_tilelang_gemm_splitk_vectorize_atomicadd():
     tilelang.tools.bench.process_func(example_tilelang_gemm_splitk_vectorize_atomicadd.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/gemm_streamk/bench_example_tilelang_gemm_splitk.py b/examples/gemm_streamk/bench_example_tilelang_gemm_splitk.py
@@ -6,9 +6,5 @@ def bench_example_tilelang_gemm_streamk():
     tilelang.tools.bench.process_func(example_tilelang_gemm_streamk.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/gemv/bench_example_gemv.py b/examples/gemv/bench_example_gemv.py
@@ -6,9 +6,5 @@ def bench_example_gemv():
     tilelang.tools.bench.process_func(example_gemv.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/linear_attention/bench_linear_attn.py b/examples/linear_attention/bench_linear_attn.py
@@ -11,9 +11,5 @@ def bench_example_linear_attn_bwd():
     tilelang.tools.bench.process_func(example_linear_attn_bwd.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/minference/bench_vs_sparse_attn.py b/examples/minference/bench_vs_sparse_attn.py
@@ -6,9 +6,5 @@ def bench_example_vertical_slash_sparse_attn():
     tilelang.tools.bench.process_func(example_vertical_slash_sparse_attn.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/seer_attention/bench_block_sparse_attn_tilelang.py b/examples/seer_attention/bench_block_sparse_attn_tilelang.py
@@ -6,9 +6,5 @@ def bench_block_sparse_attn_tilelang():
     tilelang.tools.bench.process_func(block_sparse_attn_tilelang.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/sparse_tensorcore/bench_example_sparse_tensorcore.py b/examples/sparse_tensorcore/bench_example_sparse_tensorcore.py
@@ -7,9 +7,5 @@ def bench_example_sparse_tensorcore():
     tilelang.tools.bench.process_func(tilelang_example_sparse_tensorcore.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/topk/bench_topk_tilelang.py b/examples/topk/bench_topk_tilelang.py
@@ -6,9 +6,5 @@ def bench_example_topk():
     tilelang.tools.bench.process_func(example_topk.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/warp_specialize/bench_example_warp_specialize.py b/examples/warp_specialize/bench_example_warp_specialize.py
@@ -21,9 +21,5 @@ def bench_example_warp_specialize_gemm_softpipe_stage2():
     tilelang.tools.bench.process_func(example_warp_specialize_gemm_softpipe_stage2.main)
 
 
-def main():
+if globals().get("__name__") == "__main__":
     tilelang.tools.bench.main()
-
-
-if __name__ == "__main__":
-    main()
