mirror of
https://github.com/deepseek-ai/DeepGEMM
synced 2025-04-06 01:34:54 +00:00
100 lines
4.3 KiB
Python
100 lines
4.3 KiB
Python
import hashlib
|
|
from typing import Tuple
|
|
from .utils import get_num_sms, ceil_div, get_col_major_tma_aligned_tensor, get_m_alignment_for_contiguous_layout
|
|
|
|
def is_tma_multicast_legal(n: int, block_n: int, num_tma_multicast: int, num_sms: int) -> bool:
|
|
if num_tma_multicast == 1:
|
|
return True
|
|
return (n % (block_n * num_tma_multicast) == 0) and num_sms % num_tma_multicast == 0
|
|
|
|
|
|
def get_smem_size(num_stages: int, k: int, block_m: int, block_n: int, block_k: int = 128) -> int:
|
|
smem_d = block_m * block_n * 2
|
|
smem_a_per_stage = block_m * block_k
|
|
smem_scales_a_per_stage = block_m * 4
|
|
smem_b_per_stage = block_n * block_k
|
|
smem_scales_b = ceil_div(k, block_k) * 4
|
|
smem_barrier = num_stages * 8 * 2
|
|
|
|
smem_size = 0
|
|
smem_size += smem_d
|
|
smem_size += num_stages * smem_a_per_stage
|
|
smem_size += num_stages * smem_scales_a_per_stage
|
|
smem_size += num_stages * smem_b_per_stage
|
|
smem_size += ceil_div(smem_scales_b * (1 if block_k % block_n == 0 else 2), 8) * 8
|
|
smem_size += smem_barrier
|
|
return smem_size
|
|
|
|
|
|
def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
|
|
is_grouped_contiguous: bool = False) -> Tuple[int, int, int, int, int]:
|
|
if not is_grouped_contiguous:
|
|
# TODO: for some cases, smaller M block is better, add them into tuning space
|
|
block_ms = (64 if m <= 64 else 128, )
|
|
else:
|
|
block_ms = (get_m_alignment_for_contiguous_layout(), )
|
|
block_ns = tuple(range(16, 129, 8))
|
|
|
|
fix_wave_saturate = lambda x: num_sms if x == 0 else x
|
|
get_num_waves = lambda bm, bn: (ceil_div(ceil_div(m, bm) * ceil_div(n, bn) * num_groups, num_sms) if bm else None)
|
|
get_last_wave_util = lambda bm, bn: fix_wave_saturate((ceil_div(m, bm) * ceil_div(n, bn) * num_groups) % num_sms)
|
|
|
|
# Decide block sizes by waves
|
|
best_block_m, best_block_n = None, None
|
|
for block_m in block_ms:
|
|
for block_n in block_ns:
|
|
success = False
|
|
num_waves, best_num_waves = get_num_waves(block_m, block_n), get_num_waves(best_block_m, best_block_n)
|
|
if best_block_m is None or best_block_n is None:
|
|
success = True
|
|
elif num_waves < best_num_waves:
|
|
success = True
|
|
elif num_waves == best_num_waves:
|
|
# Check last wave utilization
|
|
util = get_last_wave_util(block_m, block_n)
|
|
best_util = get_last_wave_util(best_block_m, best_block_n)
|
|
success = util > best_util or (util == best_util and (block_m > best_block_m or (block_m == best_block_m and block_n < best_block_n)))
|
|
best_block_m, best_block_n = (block_m, block_n) if success else (best_block_m, best_block_n)
|
|
assert best_block_m is not None and best_block_n is not None
|
|
|
|
# Always pick the longest one
|
|
# NOTES: for double B scales, the best number of stages may be reduced
|
|
best_num_stages, best_smem_size, sm90_capacity = None, None, 232448
|
|
for num_stages in (6, 5, 4) if 128 % best_block_n != 0 else (8, 7, 6, 5, 4):
|
|
best_smem_size = get_smem_size(num_stages, k, best_block_m, best_block_n)
|
|
if best_smem_size <= sm90_capacity:
|
|
best_num_stages = num_stages
|
|
break
|
|
assert best_num_stages is not None
|
|
|
|
# Decide the number of TMA multicast
|
|
best_num_tma_multicast = 1
|
|
if m >= 1024 and is_tma_multicast_legal(n, best_block_n, 2, num_sms) and num_groups == 1:
|
|
best_num_tma_multicast = 2
|
|
|
|
return best_block_m, best_block_n, best_num_stages, best_num_tma_multicast, best_smem_size
|
|
|
|
|
|
def hash_to_hex(s: str) -> str:
|
|
md5 = hashlib.md5()
|
|
md5.update(s.encode('utf-8'))
|
|
return md5.hexdigest()[0:12]
|
|
|
|
class ConfigCache:
|
|
|
|
def __init__(self) -> None:
|
|
self.cached = {}
|
|
|
|
def compute_and_cache(self, m: int, n: int, k: int, num_groups: int, num_sms: int,
|
|
is_grouped_contiguous: bool = False) -> Tuple[int, int, int, int]:
|
|
signature = str((m, n, k, num_groups, num_sms, is_grouped_contiguous))
|
|
signature = hash_to_hex(signature)
|
|
if signature in self.cached:
|
|
return self.cached[signature]
|
|
best_config = get_best_configs(m, n, k, num_groups, num_sms, is_grouped_contiguous)
|
|
self.cached[signature] = best_config
|
|
return best_config
|
|
|
|
config_cache = ConfigCache()
|
|
|