DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-06-26 23:15:49 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add TMA optimization hint for large FFMA segments

Browse Source 
This check is added to identify segments with a high number of FFMA (Fused Multiply-Add) instructions, which may benefit from Tensor Memory Accelerator (TMA) optimizations on NVIDIA Hopper architecture. TMA enables asynchronous data transfers from global memory to shared memory, reducing register pressure and improving data access efficiency. Segments with a large number of FFMA instructions often involve significant data movement, making them prime candidates for TMA to accelerate memory operations and enhance overall GEMM performance.
This commit is contained in:
A-transformer
2025-03-07 19:37:55 +04:00
committed by GitHub
parent 9d3222a93e
commit 82d4b996f7
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
deep_gemm/jit/interleave_ffma.py
View File

@@ -76,6 +76,11 @@ def modify_segment(m, name, ffma_lines):
num_lines = (len(ffma_lines) * 9 // 16) // 2 * 2 num_lines = (len(ffma_lines) * 9 // 16) // 2 * 2
assert num_lines % 2 == 0 assert num_lines % 2 == 0
tma_threshold = 32
if num_lines // 2 > tma_threshold:
if os.getenv('DG_PRINT_REG_REUSE', None):
print(f' > segment `{name}` may benefit from TMA optimization due to large FFMA count ({num_lines // 2})')
le_bytes, new_le_bytes = [], [] le_bytes, new_le_bytes = [], []
reused_list = [] reused_list = []
dst_reg_set = set() dst_reg_set = set()