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 (fix review issues #4-#10)

Browse Source
This commit is contained in:
admin
2026-09-03 11:34:27 +00:00
parent 1385d01816
commit f756729e43
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -145,11 +145,12 @@ class StreamKBranch(Branch):
t_mmad = t_mmad_tile / grid_k t_mmad = t_mmad_tile / grid_k
t_mte2 = t_mte2_tile / grid_k t_mte2 = t_mte2_tile / grid_k
# 归约: 部分和 4B 驻留 L2, AIV 归约 # 归约: 部分和 4B 驻留 L2, AIV 归约 (含最终按 C dtype 写回)
t_reduce = eval_streamk_reduce(tile_elems, grid_k, out_b, s) t_reduce = eval_streamk_reduce(tile_elems, grid_k, out_b, s)
# 写出: 最终归约结果按 C dtype; 中间部分和按 4B # Fixpipe: 部分和写出按 4B (L0C dtype, 防精度丢失), 驻留 L2.
fix_bytes = tile_elems * out_b # 最终归约结果的 C dtype 写回已在 t_reduce 内计, 此处不重复 (issue#9 口径对齐).
fix_bytes = grid_k * tile_elems * 4
t_fix = fix_bytes / s.bw_l2_pc t_fix = fix_bytes / s.bw_l2_pc
flops_pc = 2.0 * tile_elems * k / grid_k flops_pc = 2.0 * tile_elems * k / grid_k
@@ -158,7 +159,8 @@ class StreamKBranch(Branch):
return assemble_timing( return assemble_timing(
t_mte2_gm=t_mte2, t_mte2_l2=0.0, t_dma_cmd=0.0, t_mte2_gm=t_mte2, t_mte2_l2=0.0, t_dma_cmd=0.0,
t_mmad=t_mmad, t_fixpipe=t_fix, t_reduce=t_reduce, t_mmad=t_mmad, t_fixpipe=t_fix, t_reduce=t_reduce,
t_drain=t_reduce, # 归约串行追加 t_drain=t_reduce, # 归约串行追加 (reduce_serial 默认 True, 不进稳态 max)
gm_read_bytes=gm_bytes, l2_read_bytes=0.0, dma_cmd_count=0.0, gm_read_bytes=gm_bytes, l2_read_bytes=0.0, dma_cmd_count=0.0,
cube_flops=flops_pc, fixpipe_bytes=fix_bytes, cube_flops=flops_pc, fixpipe_bytes=fix_bytes,
reduce_serial=True,
) )