[Refactor] Update example_mla_decode.py and add tests for block_sparse_attn_tilelang

* Refactor example_mla_decode.py to define a main function for better structure and clarity.
* Introduce test_example_mla_decode.py to validate the functionality of example_mla_decode.
* Refactor block_sparse_attn_tilelang.py to define a main function and add test_block_sparse_attn_tilelang.py for testing.
* Ensure all new test files are integrated with tilelang testing framework.

Author: chengyupku

examples/deepseek_mla/example_mla_decode.py

@@ -272,7 +272,7 @@ def ref_program(q, q_pe, kv, k_pe, glse, Output_partial):
     return out
 
 
-if __name__ == "__main__":
+def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--batch', type=int, default=128, help='batch size')
     parser.add_argument('--heads', type=int, default=128, help='q heads number')
@@ -296,3 +296,7 @@ def ref_program(q, q_pe, kv, k_pe, glse, Output_partial):
     latency = profiler.do_bench(warmup=500)
     print(f"Latency: {latency} ms")
     print(f"TFlops: {total_flops / latency * 1e-9} TFlops")
+
+
+if __name__ == "__main__":
+    main()

examples/deepseek_mla/test_example_mla_decode.py

@@ -0,0 +1,14 @@
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+import tilelang.testing
+
+import example_mla_decode
+
+
+@tilelang.testing.requires_cuda
+def test_example_mla_decode():
+    example_mla_decode.main()
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()

examples/seer_attention/block_sparse_attn_tilelang.py

@@ -261,6 +261,10 @@ def test_topk_sparse_attention_qlen_lt_klen():
     print("Pass topk sparse attention test with qlen < klen")
 
 
-if __name__ == "__main__":
+def main():
     test_topk_sparse_attention()
     test_topk_sparse_attention_qlen_lt_klen()
+
+
+if __name__ == "__main__":
+    main()

examples/seer_attention/test_block_sparse_attn_tilelang.py

@@ -0,0 +1,14 @@
+# Copyright (c) Tile-AI Corporation.
+# Licensed under the MIT License.
+import tilelang.testing
+
+import block_sparse_attn_tilelang
+
+
+@tilelang.testing.requires_cuda
+def test_block_sparse_attn_tilelang():
+    block_sparse_attn_tilelang.main()
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()