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

Fixed the bug in get_swizzle_mode function related to elem_size setting.

Browse Source
This commit is contained in:
Ther-LF 2025-06-17 08:49:36 +00:00
parent 0c88cd0139
commit 88a2b6c490
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
deep_gemm/jit_kernels/gemm.py
View File

@ -17,8 +17,8 @@ def is_tma_multicast_legal(shape_dim: int, block_dim: int, num_tma_multicast: in
return divisible and num_sms % num_tma_multicast == 0
def get_swizzle_mode(block_n: int) -> int:
elem_size = 2
def get_swizzle_mode(block_n: int, is_fp32_out: bool) -> int:
elem_size = 4 if is_fp32_out else 2
for mode_bytes in (128, 64, 32):
if (block_n * elem_size) % mode_bytes == 0:
return mode_bytes
@ -38,7 +38,7 @@ def get_smem_config(num_stages: int, k: int, block_m: int, block_n: int, block_k
assert block_k == 128
# Try swizzle first, as it does not waste shared memory
swizzle_mode = get_swizzle_mode(block_n)
swizzle_mode = get_swizzle_mode(block_n, is_fp32_out)
block_n_padding = get_block_n_padding_for_smem_d(
block_n) if swizzle_mode == 0 else 0