DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-04-10 18:55:28 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #2 from deepseek-ai/main

Browse Source 
sync
This commit is contained in:
A-transformer 2025-02-27 21:47:31 +04:00 committed by GitHub
commit a2e0d68eed
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 15 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
deep_gemm/include/deep_gemm/fp8_gemm.cuh
View File

@ -109,7 +109,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout,
}
// Initialize barriers
DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "To many TMA multicast");
DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast");
if (threadIdx.x == kNumMathThreads) {
#pragma unroll
for (int i = 0; i < kNumStages; ++ i) {
@ -406,7 +406,8 @@ public:
template <typename T>
static CUtensorMap make_2d_tma_d_desc(T* global_address, uint32_t shape_m) {
return make_2d_tma_desc(global_address, Layout::RowMajor,
shape_m * (kGemmType == GemmType::GroupedMasked ? kNumGroups : 1), SHAPE_N, BLOCK_M, BLOCK_N,
shape_m * (kGemmType == GemmType::GroupedMasked ? kNumGroups : 1), SHAPE_N,
min(BLOCK_M, shape_m), BLOCK_N,
CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);
}

9
deep_gemm/jit/interleave_ffma.py
View File

@ -38,7 +38,7 @@ def extract_ffma(sass):
current = []
if os.getenv('DG_PRINT_REG_REUSE', None):
print(f"Found {len(collected)} FFMA segments")
print(f'Found {len(collected)} FFMA segments')
return collected
@ -58,7 +58,6 @@ def validate(m, offset, le_bytes, num_lines):
def parse_registers(line):
import re
line = re.sub(r'/\*.*?\*/', '', line)
line = line.replace(';', '')
tokens = line.strip().split(',')
@ -92,7 +91,7 @@ def modify_segment(m, name, ffma_lines):
is_first_occurred = dst_reg not in dst_reg_set
if is_first_occurred or (last_reused and dst_reg == last_dst_reg):
# Modify the `reuse` and `yield` bits
assert high_hex & 0x0800200000000000, f"{hex(high_hex)}"
assert high_hex & 0x0800200000000000, f'{hex(high_hex)}'
high_hex ^= 0x0800200000000000
reused = False
num_changed += 1
@ -102,7 +101,7 @@ def modify_segment(m, name, ffma_lines):
new_le_bytes.append(low_hex.to_bytes(8, 'little') + high_hex.to_bytes(8, 'little'))
last_reused, last_dst_reg = reused, dst_reg
if os.getenv('DG_PRINT_REG_REUSE', None):
print(f" > segment `{name}` new reused list ({num_changed} changed): {reused_list}")
print(f' > segment `{name}` new reused list ({num_changed} changed): {reused_list}')
# Find the offset
offsets = []
@ -130,7 +129,7 @@ def process(path):
mm.close()
if __name__ == "__main__":
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Interleave FFMA reg reuse')
parser.add_argument('--so', help='Path to the SO file')
args = parser.parse_args()

4
deep_gemm/jit_kernels/gemm.py
View File

@ -79,10 +79,12 @@ def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
elif num_waves < best_num_waves:
success = True
elif num_waves == best_num_waves:
div_n = bool(128 % block_n)
best_div_n = bool(128 % best_block_n)
# 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_n >= best_block_n and block_m <= best_block_m))
success = util > best_util or (util == best_util and (block_m > best_block_m or block_m == best_block_m and (div_n < best_div_n or div_n == best_div_n 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

5
deep_gemm/jit_kernels/m_grouped_gemm.py
View File

@ -160,6 +160,11 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
global includes, template
num_sms = get_num_sms()
block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(expected_m, n, k, num_groups, num_sms)
# Extra checks for TMA store
if num_groups > 1 and m > block_m:
assert m % block_m == 0, f'For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m})'
args = (lhs, lhs_scales, rhs, rhs_scales, out,
masked_m, m,
torch.cuda.current_stream(), num_sms, smem_size)