bugfix: Fix FLOPS calculation for bench_trtllm_gen_mla.py (#1640)

RayWang96 · web-flow · commit 17978a3d40f1 · 2025-09-05T02:48:20.000-07:00
Fix FLOPS calculation for bench_trtllm_gen_mla.py Details in #1638
diff --git a/benchmarks/bench_trtllm_gen_mla.py b/benchmarks/bench_trtllm_gen_mla.py
@@ -108,10 +108,9 @@ def bench_trtllm_mla(batch_size, q_len_per_request, seq_len, page_size, dtype):
     ms = np.median(measurements)
     flops = (
         2
-        * batch_size
         * num_q_heads
         * (2 * kv_lora_rank + qk_rope_head_dim)
-        * seq_len
+        * sum(seq_lens)
         * q_len_per_request
     )
     print(

Original file line number	Diff line number	Diff line change
`@@ -108,10 +108,9 @@ def bench_trtllm_mla(batch_size, q_len_per_request, seq_len, page_size, dtype):`
`108`	`108`	`ms = np.median(measurements)`
`109`	`109`	`flops = (`
`110`	`110`	`2`
`111`		`- * batch_size`
`112`	`111`	`* num_q_heads`
`113`	`112`	`* (2 * kv_lora_rank + qk_rope_head_dim)`
`114`		`- * seq_len`
	`113`	`+ * sum(seq_lens)`
`115`	`114`	`* q_len_per_request`
`116`	`115`	`)`
`117`	`116`	`print(`