From c05cbeafad2b40351580f1245bec5ac87bd3fa9a Mon Sep 17 00:00:00 2001 From: admin Date: Thu, 3 Sep 2026 12:43:55 +0000 Subject: [PATCH] Update BMM_Theory: bmm_theory/branches/stream_k.py --- BMM/BMM_Theory/bmm_theory/branches/stream_k.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/BMM/BMM_Theory/bmm_theory/branches/stream_k.py b/BMM/BMM_Theory/bmm_theory/branches/stream_k.py index 8b1455b..321a85c 100644 --- a/BMM/BMM_Theory/bmm_theory/branches/stream_k.py +++ b/BMM/BMM_Theory/bmm_theory/branches/stream_k.py @@ -145,14 +145,13 @@ class StreamKBranch(Branch): t_mmad = t_mmad_tile / grid_k t_mte2 = t_mte2_tile / grid_k - # 归约: 部分和 4B 驻留 L2, AIV 归约 (含部分和写/读回/求和/最终按 C dtype 写回) - # 口径 (issue#11): 归约整体为串行追加 (t_drain=t_reduce, reduce_serial=True), - # 部分和写出已计入 eval_streamk_reduce 的 t_write_partial —— 稳态 Fixpipe 不再 - # 重复计账. 此前按 grid_k*tile*4B/单核带宽份额另计一次, 既重复计账又把整组 - # 部分和串行压到单核写口, 高估 grid_k 倍 (streamk_demo 曾虚高到 55us/FIXPIPE). + # 归约: 部分和 4B 驻留 L2, AIV 归约 (含最终按 C dtype 写回) t_reduce = eval_streamk_reduce(tile_elems, grid_k, out_b, s) - fix_bytes = 0.0 - t_fix = 0.0 + + # Fixpipe: 部分和写出按 4B (L0C dtype, 防精度丢失), 驻留 L2. + # 最终归约结果的 C dtype 写回已在 t_reduce 内计, 此处不重复 (issue#9 口径对齐). + fix_bytes = grid_k * tile_elems * 4 + t_fix = fix_bytes / s.bw_l2_pc flops_pc = 2.0 * tile_elems * k / grid_k gm_bytes = k * (2 * math.sqrt(tile_elems)) * dt / grid_k