DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-05-05 09:24:22 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

add lru-cache to avoid repeated calculation

Browse Source
This commit is contained in:
Yi Zhang 2025-04-04 12:44:26 +08:00 committed by GitHub
parent c187c23ba8
commit 776bd0cccc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
deep_gemm/jit_kernels/gemm.py
View File

@ -1,5 +1,6 @@
import math
import torch
from functools import lru_cache
from typing import Tuple
from .tuner import jit_tuner
@ -66,7 +67,7 @@ def get_smem_size(num_stages: int, k: int, block_m: int, block_n: int, block_k:
smem_size += smem_barrier
return smem_size
@lru_cache(maxsize=None)
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, Tuple[int, bool], int]:
if not is_grouped_contiguous: