[Feature]: Add benchmark scripts for examples

.github/workflows/pr-perfbench-bot.yml

@@ -6,7 +6,7 @@ on:
       - created
 
 permissions:
-  contents: read
+  contents: write
 
 concurrency:
   group: "${{ github.workflow }}-${{ github.ref }}"
@@ -53,8 +53,12 @@ jobs:
         run: |
           python -m venv tll
           source tll/bin/activate
-          pip install -r requirements-test.txt
-          pip install .
+          export PIP_CONFIG_FILE=/dev/null
+          export PYTHONUSERBASE=""
+          pip config unset global.user
+          pip config unset user.user
+          pip install --no-user -r requirements-test.txt
+          pip install --no-user .
 
       - name: Install original version
         run: |
@@ -64,25 +68,70 @@ jobs:
           git checkout main
           python -m venv tl
           source tl/bin/activate
-          pip install -r requirements-test.txt
-          pip install .
+          export PIP_CONFIG_FILE=/dev/null
+          export PYTHONUSERBASE=""           
+          pip config unset global.user || true
+          pip config unset user.user || true
+          pip install --no-user -r requirements-test.txt
+          pip install --no-user .
 
       - name: Run performance test
         id: perfbench
         run: |
           source tl/bin/activate
           python maint/scripts/ci_performance.py
+      - name: Read markdown table
+        id: read_md
+        run: |
+          echo "content<<EOF" >> $GITHUB_OUTPUT
+          cat bench.md >> $GITHUB_OUTPUT
+          echo "EOF" >> $GITHUB_OUTPUT
+      - name: Upload PNG to GitHub and get URL
+        id: upload_png
+        uses: actions/github-script@v8
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const fs = require('fs');
+            const content = fs.readFileSync('bench.png').toString('base64');
+            // Create blob in the repo
+            const blob = await github.rest.git.createBlob({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              content: content,
+              encoding: "base64",
+            });
+            // Attach blob as a tree item
+            const tree = await github.rest.git.createTree({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              tree: [{
+                path: `bench_${context.runId}.png`,
+                mode: '100644',
+                type: 'blob',
+                sha: blob.data.sha
+              }]
+            });
+            // Raw file URL (works for embedding image)
+            const url = `https://raw.githubusercontent.com/${context.repo.owner}/${context.repo.repo}/${tree.data.sha}/bench_${context.runId}.png`
+            core.setOutput("url", url);
 
       - name: Post test results as PR comment
         uses: actions/github-script@v8
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           script: |
+            const md = `${{ steps.read_md.outputs.content }}`;
+            const img = `${{ steps.upload_png.outputs.url }}`;
             github.rest.issues.createComment({
               owner: context.repo.owner,
               repo: context.repo.repo,
               issue_number: context.issue.number,
-              body: '📊 ​**Performance Test Results** (triggered by @' + context.payload.comment.user.login + '):\n\n' +
-                'Run listed here: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}\n\n' +
-                "${{ steps.perfbench.outputs.stdout }}"
+              body:
+                '📊 **Performance Test Results** (triggered by @' +
+                context.payload.comment.user.login + ')\n\n' +
+                'Run: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}\n\n' +
+                md +
+                '\n\n📈 **Speedup Plot:**\n\n' +
+                `![Speedup Plot](${img})`
             })

examples/analyze/bench_example_analyze.py

@@ -0,0 +1,15 @@
+import tilelang.tools.bench
+import example_conv_analyze
+import example_gemm_analyze
+
+
+def bench_example_gemm_analyze():
+    tilelang.tools.bench.process_func(example_gemm_analyze.main)
+
+
+def bench_example_conv_analyze():
+    tilelang.tools.bench.process_func(example_conv_analyze.main)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/attention_sink/bench_example_attention_sink.py

@@ -0,0 +1,52 @@
+import tilelang.tools.bench
+import example_gqa_sink_bwd_bhsd
+import example_gqa_sink_fwd_bhsd_wgmma_pipelined
+import example_mha_sink_bwd_bhsd
+import example_mha_sink_fwd_bhsd
+import example_mha_sink_fwd_bhsd_wgmma_pipelined
+
+
+def bench_example_mha_sink_fwd_bhsd():
+    tilelang.tools.bench.process_func(example_mha_sink_fwd_bhsd.main)
+
+
+def bench_example_mha_sink_fwd_bhsd_sliding_window():
+    tilelang.tools.bench.process_func(example_mha_sink_fwd_bhsd.main, window_size=128)
+
+
+def bench_example_mha_sink_fwd_bhsd_wgmma_pipelined():
+    tilelang.tools.bench.process_func(example_mha_sink_fwd_bhsd_wgmma_pipelined.main)
+
+
+def bench_example_mha_sink_fwd_bhsd_wgmma_pipelined_sliding_window():
+    tilelang.tools.bench.process_func(
+        example_mha_sink_fwd_bhsd_wgmma_pipelined.main, window_size=128)
+
+
+def bench_example_gqa_sink_fwd_bhsd_wgmma_pipelined():
+    tilelang.tools.bench.process_func(example_gqa_sink_fwd_bhsd_wgmma_pipelined.main)
+
+
+def bench_example_gqa_sink_fwd_bhsd_wgmma_pipelined_sliding_window():
+    tilelang.tools.bench.process_func(
+        example_gqa_sink_fwd_bhsd_wgmma_pipelined.main, window_size=128)
+
+
+def bench_example_mha_sink_bwd_bhsd():
+    tilelang.tools.bench.process_func(example_mha_sink_bwd_bhsd.main)
+
+
+def bench_example_mha_sink_bwd_bhsd_sliding_window():
+    tilelang.tools.bench.process_func(example_mha_sink_bwd_bhsd.main, window_size=128)
+
+
+def bench_example_gqa_sink_bwd_bhsd():
+    tilelang.tools.bench.process_func(example_gqa_sink_bwd_bhsd.main)
+
+
+def bench_example_gqa_sink_bwd_bhsd_sliding_window():
+    tilelang.tools.bench.process_func(example_gqa_sink_bwd_bhsd.main, window_size=128)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/blocksparse_attention/bench_example_blocksparse_attention.py

@@ -0,0 +1,55 @@
+import tilelang.tools.bench
+import block_sparse_attn_triton
+import example_tilelang_block_sparse_attn
+import example_tilelang_sparse_gqa_decode_varlen_indice
+import example_tilelang_sparse_gqa_decode_varlen_mask
+import example_triton_sparse_gqa_decode_varlen_indice
+import example_triton_sparse_gqa_decode_varlen_mask
+
+
+def bench_block_sparse_attn_triton():
+    tilelang.tools.bench.process_func(block_sparse_attn_triton.main)
+
+
+def bench_example_tilelang_block_sparse_attn():
+    tilelang.tools.bench.process_func(example_tilelang_block_sparse_attn.main)
+
+
+def bench_example_tilelang_sparse_gqa_decode_varlen_indice():
+    tilelang.tools.bench.process_func(
+        example_tilelang_sparse_gqa_decode_varlen_indice.main, batch=1, max_cache_seqlen=2048)
+
+
+def bench_example_tilelang_sparse_gqa_decode_varlen_mask():
+    tilelang.tools.bench.process_func(
+        example_tilelang_sparse_gqa_decode_varlen_mask.main, batch=1, max_cache_seqlen=2048)
+
+
+def bench_example_triton_sparse_gqa_decode_varlen_indice():
+    tilelang.tools.bench.process_func(
+        example_triton_sparse_gqa_decode_varlen_indice.main,
+        batch=8,
+        heads=8,
+        heads_kv=4,
+        max_cache_seqlen=2048,
+        dim=128,
+        dim_v=128,
+        sparse_ratio=0.8,
+        block_size=32)
+
+
+def bench_example_triton_sparse_gqa_decode_varlen_mask():
+    tilelang.tools.bench.process_func(
+        example_triton_sparse_gqa_decode_varlen_mask.main,
+        batch=8,
+        heads=8,
+        heads_kv=4,
+        max_cache_seqlen=2048,
+        dim=128,
+        dim_v=128,
+        sparse_ratio=0.8,
+        block_size=32)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/blocksparse_gemm/bench_example_blocksparse_gemm.py

@@ -0,0 +1,10 @@
+import tilelang.tools.bench
+import example_blocksparse_gemm
+
+
+def bench_example_blocksparse_gemm():
+    tilelang.tools.bench.process_func(example_blocksparse_gemm.main)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/cast/bench_example_cast.py

@@ -0,0 +1,21 @@
+import tilelang.tools.bench
+import example_group_per_split_token_cast_to_fp8
+import example_per_token_cast_to_fp8
+
+
+def bench_example_group_per_split_token_cast_to_fp8():
+    tilelang.tools.bench.process_func(
+        example_group_per_split_token_cast_to_fp8.main,
+        M=1024,
+        N=1024,
+        BG=2,
+        blk_m=4,
+        batch_sizes=[128, 896])
+
+
+def bench_example_per_token_cast_to_fp8():
+    tilelang.tools.bench.process_func(example_per_token_cast_to_fp8.main, M=2048, N=512, blk_m=8)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/convolution/bench_example_convolution.py

@@ -0,0 +1,15 @@
+import tilelang.tools.bench
+import example_convolution
+import example_convolution_autotune
+
+
+def bench_example_convolution():
+    tilelang.tools.bench.process_func(example_convolution.main)
+
+
+def bench_example_convolution_autotune():
+    tilelang.tools.bench.process_func(example_convolution_autotune.main)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/deepseek_deepgemm/bench_example_deepgemm_fp8_2xAcc.py

@@ -0,0 +1,10 @@
+import tilelang.tools.bench
+import example_deepgemm_fp8_2xAcc
+
+
+def bench_example_deepgemm_fp8_2xAcc():
+    tilelang.tools.bench.process_func(example_deepgemm_fp8_2xAcc.main)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/deepseek_mla/bench_example_mla_decode.py

@@ -0,0 +1,10 @@
+import tilelang.tools.bench
+import example_mla_decode
+
+
+def bench_example_mla_decode():
+    tilelang.tools.bench.process_func(example_mla_decode.main)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/deepseek_nsa/bench_example_tilelang_nsa.py

@@ -0,0 +1,15 @@
+import tilelang.tools.bench
+import example_tilelang_nsa_fwd
+import example_tilelang_nsa_decode
+
+
+def bench_example_tilelang_nsa_fwd():
+    tilelang.tools.bench.process_func(example_tilelang_nsa_fwd.main)
+
+
+def bench_example_tilelang_nsa_fwd_decode():
+    tilelang.tools.bench.process_func(example_tilelang_nsa_decode.main)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()

examples/deepseek_v32/bench_tilelang_example_deepseek_v32.py

@@ -0,0 +1,64 @@
+import tilelang.tools.bench
+import fp8_lighting_indexer
+import sparse_mla_bwd
+import sparse_mla_fwd
+import sparse_mla_fwd_pipelined
+import topk_selector
+
+
+def bench_topk_selector():
+    tilelang.tools.bench.process_func(topk_selector.test_topk_selector)
+
+
+def bench_fp8_lighting_indexer():
+    tilelang.tools.bench.process_func(
+        fp8_lighting_indexer.test_fp8_lighting_indexer,
+        S=512,
+        SKV=1024,
+        H=32,
+        HKV=1,
+        D=64,
+        kv_stride=1)
+
+
+def bench_sparse_mla_fwd():
+    tilelang.tools.bench.process_func(
+        sparse_mla_fwd.test_sparse_mla_fwd,
+        S=256,
+        SKV=1024,
+        H=64,
+        HKV=1,
+        DQK=576,
+        DV=512,
+        topk=256,
+        check_correctness=False)
+
+
+def bench_sparse_mla_fwd_pipelined():
+    tilelang.tools.bench.process_func(
+        sparse_mla_fwd_pipelined.test_sparse_mla_fwd_pipelined,
+        S=256,
+        SKV=512,
+        H=64,
+        HKV=1,
+        DQK=576,
+        DV=512,
+        topk=256,
+        check_correctness=False)
+
+
+def bench_sparse_mla_bwd():
+    tilelang.tools.bench.process_func(
+        sparse_mla_bwd.test_sparse_mla_bwd,
+        S=256,
+        SKV=512,
+        H=64,
+        HKV=1,
+        DQKV=576,
+        DV=512,
+        topk=256,
+        check_correctness=False)
+
+
+if globals().get("__name__") == "__main__":
+    tilelang.tools.bench.main()