72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
"""Diagnostic: Print exact SMEM budget for various HEAD_DIM values."""
|
|
import math
|
|
from dsv4.kernels.attention.fmha import FmhaKernel
|
|
|
|
|
|
def smem_budget(hd, s_k=128, use_smem_p=None):
|
|
"""Reproduce FmhaKernel._setup SMEM allocation sizes."""
|
|
kv_stage = 1 if hd > 128 else 2
|
|
q_stage = 1
|
|
k_tile = min(hd, 256)
|
|
pv_n_tile = min(hd, 256)
|
|
pv_mma_tiler = (128, pv_n_tile, k_tile)
|
|
qk_mma_tiler = (128, 128, k_tile)
|
|
|
|
# Rough SMEM sizes (actual CuTe layouts have padding/alignment)
|
|
# sQ: (M, K_tile) BF16 * q_stage
|
|
sQ = 128 * k_tile * 2 * q_stage
|
|
# sK: (K_tile, N) BF16 * kv_stage
|
|
sK = k_tile * 128 * 2 * kv_stage
|
|
# sV: (K_tile_V, N) BF16 * kv_stage, where K_tile_V = pv_n_tile
|
|
sV = pv_n_tile * 128 * 2 * kv_stage
|
|
# sC: (M, N) BF16 * num_c_stage
|
|
num_c_stage = 1 if hd > 256 else 2
|
|
sC = 128 * pv_n_tile * 2 * num_c_stage
|
|
# sP: only if SMEM-P
|
|
sP = 0
|
|
if use_smem_p is None:
|
|
use_smem_p = hd > 64
|
|
if use_smem_p:
|
|
sP = 128 * pv_n_tile * 2 * 1 # 1 stage
|
|
|
|
total = sQ + sK + sV + sC + sP
|
|
limit = 232 * 1024
|
|
|
|
print(f"hd={hd}: k_tile={k_tile}, kv_stage={kv_stage}, pv_n_tile={pv_n_tile}, "
|
|
f"num_c_stage={num_c_stage}, use_smem_p={use_smem_p}")
|
|
print(f" sQ={sQ//1024}KB, sK={sK//1024}KB, sV={sV//1024}KB, sC={sC//1024}KB, sP={sP//1024}KB")
|
|
print(f" Total={total//1024}KB vs {limit//1024}KB limit → {'✅ FITS' if total <= limit else '❌ OVER'}")
|
|
print()
|
|
|
|
# Now compute with overlap: sQ/sV share same region
|
|
if not use_smem_p:
|
|
overlap_total = max(sQ, sV) + sK + sC
|
|
print(f" With sQ/sV overlap: {overlap_total//1024}KB → {'✅ FITS' if overlap_total <= limit else '❌ OVER'}")
|
|
else:
|
|
overlap_total = max(sQ, sV) + sK + sC + sP
|
|
print(f" With sQ/sV overlap (+sP): {overlap_total//1024}KB → {'✅ FITS' if overlap_total <= limit else '❌ OVER'}")
|
|
|
|
# Another option: also overlap sK with sC (K consumed before C written)
|
|
# After QK GEMM, sK is done. Then PV starts, writes O to TMEM, then epilogue writes C.
|
|
# But sK is needed for ALL kv_tiles... and sC is written after all PV is done.
|
|
# Actually sK is needed per kv_tile, so it stays alive through the whole QK loop.
|
|
# Overlap sK/sC doesn't work.
|
|
|
|
# Option: reduce pv_n_tile to 128
|
|
pv_n_tile_small = 128
|
|
sV_small = pv_n_tile_small * 128 * 2 * kv_stage
|
|
sC_small = 128 * pv_n_tile_small * 2 * num_c_stage
|
|
total_small = sQ + sK + sV_small + sC_small + sP
|
|
print(f" With pv_n_tile=128: sV={sV_small//1024}KB, sC={sC_small//1024}KB, "
|
|
f"Total={total_small//1024}KB → {'✅ FITS' if total_small <= limit else '❌ OVER'}")
|
|
overlap_small = max(sQ, sV_small) + sK + sC_small
|
|
print(f" With pv_n_tile=128 + sQ/sV overlap: {overlap_small//1024}KB → {'✅ FITS' if overlap_small <= limit else '❌ OVER'}")
|
|
print()
|
|
|
|
|
|
print("=== SMEM Budget Analysis ===\n")
|
|
smem_budget(64)
|
|
smem_budget(128)
|
|
smem_budget(256)
|
|
smem_budget(512)
|