admin/matmul-analysis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update BMM_Theory: bmm_theory/branches/stream_k.py

Browse Source
This commit is contained in:
admin
2026-09-03 12:43:55 +00:00
parent 49397260f6
commit c05cbeafad
1 changed files with 6 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
BMM/BMM_Theory/bmm_theory/branches/stream_k.py
View File

@@ -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