[Benchmark CI] Allow specifying custom env vars via UI (#972)

Author: yf225

.github/workflows/benchmark.yml

@@ -24,6 +24,10 @@ on:
       kernels:
         required: true
         type: string
+      env-vars:
+        required: false
+        type: string
+        default: ""
 
 jobs:
   benchmark:
@@ -142,7 +146,7 @@ jobs:
             echo "Available implementations for $kernel: $IMPLS"
 
             # Do autotuning but do not record the results
-            python benchmarks/run.py \
+            ${{ inputs.env-vars }} python benchmarks/run.py \
                 --op $kernel \
                 --metrics speedup,accuracy \
                 --latency-measure-mode triton_do_bench \
@@ -157,7 +161,7 @@ jobs:
             sleep 2m
 
             # Run again with cache and record results
-            python benchmarks/run.py \
+            ${{ inputs.env-vars }} python benchmarks/run.py \
                 --op $kernel \
                 --metrics speedup,accuracy \
                 --latency-measure-mode triton_do_bench \

.github/workflows/benchmark_dispatch.yml

@@ -23,6 +23,11 @@ on:
         required: false
         type: string
         default: "softmax,jsd,welford,kl_div,int4_gemm,layer_norm,layer_norm-bwd,rms_norm,rms_norm-bwd,cross_entropy,flash_attention,gemm,grouped_gemm"
+      env_vars:
+        description: 'Environment variables for benchmark runner'
+        required: false
+        type: string
+        default: ""
 
 jobs:
   gen-matrix-h100:
@@ -49,6 +54,7 @@ jobs:
       container-options: --gpus all
       alias: h100
       kernels: ${{ matrix.kernels }}
+      env-vars: ${{ github.event.inputs.env_vars }}
 
   gen-matrix-b200:
     uses: ./.github/workflows/compute-benchmark-matrix.yml
@@ -74,6 +80,7 @@ jobs:
       container-options: --gpus all
       alias: b200
       kernels: ${{ matrix.kernels }}
+      env-vars: ${{ github.event.inputs.env_vars }}
 
   gen-matrix-mi325x:
     uses: ./.github/workflows/compute-benchmark-matrix.yml
@@ -99,3 +106,4 @@ jobs:
       container-options: --device=/dev/kfd --device=/dev/dri
       alias: mi325x
       kernels: ${{ matrix.kernels }}
+      env-vars: ${{ github.event.inputs.env_vars }}