[Refactor] Update main function structure in example scripts and add tests (#475)

* [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.

* [Test] Enhance test_example_mla_decode with argument mocking

* Update test_example_mla_decode.py to mock sys.argv for better test isolation.
* Ensure the main function of example_mla_decode is called with the correct arguments during testing.

examples/deepseek_mla/example_mla_decode.py

@@ -270,7 +270,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')
@@ -294,3 +294,7 @@ if __name__ == "__main__":
     latency = profiler.do_bench(warmup=500)
     print(f"Latency: {latency} ms")
     print(f"TFlops: {total_flops / latency * 1e-9} TFlops")
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file

examples/deepseek_mla/test_example_mla_decode.py

+import tilelang.testing
+
+import example_mla_decode
+from unittest import mock
+import sys
+
+
+@tilelang.testing.requires_cuda
+def test_example_mla_decode():
+    with mock.patch.object(sys, 'argv', ["example_mla_decode.py"]):
+        example_mla_decode.main()
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()

examples/seer_attention/block_sparse_attn_tilelang.py

@@ -259,6 +259,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()
\ No newline at end of file

examples/seer_attention/test_block_sparse_attn_tilelang.py

+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()
\ No newline at end of file