[Benchmark CI] Use fewer num_inputs for flash_attention to avoid timeout (#857)

yf225 · web-flow · commit 05fb47de5f10 · 2025-10-09T10:51:52.000-07:00
diff --git a/benchmarks/run.py b/benchmarks/run.py
@@ -193,7 +193,7 @@ class RunResult:
         "attention",
         {
             "d_head": 128,  # Set default head dimension to 128 for TLX attention compatibility
-            "num_inputs": 10,  # flash_attention takes long time on Benchmark CI, so use fewer inputs instead.
+            "num_inputs": 6,  # flash_attention takes long time on Benchmark CI, so use fewer inputs instead.
         },
     ),
     "cross_entropy": (

Original file line number	Diff line number	Diff line change
`@@ -193,7 +193,7 @@ class RunResult:`
`193`	`193`	`"attention",`
`194`	`194`	`{`
`195`	`195`	`"d_head": 128, # Set default head dimension to 128 for TLX attention compatibility`
`196`		`- "num_inputs": 10, # flash_attention takes long time on Benchmark CI, so use fewer inputs instead.`
	`196`	`+ "num_inputs": 6, # flash_attention takes long time on Benchmark CI, so use fewer inputs instead.`
`197`	`197`	`},`
`198`	`198`	`),`
`199`	`199`	`"cross_entropy": (`