chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make various updates and fixes:

Browse Source 
- Add support for legacy CUDA versions; now compatible with CUDA 12.3 and newer
- Add support for NVRTC compilation
- Other fixes and code refactoring
This commit is contained in:
Ray Wang
2025-08-02 19:52:22 -07:00
parent aff9da0aba
commit d9c363f86f
36 changed files with 592 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
tests/generators.py
View File

@@ -14,6 +14,7 @@ class KernelType(enum.Enum):
# For SM100 GEMMs
Kernel1D1D = 0
Kernel1D2D = 1
KernelNoSF = 2
def is_1d1d(self):
return self.value == 0
@@ -21,6 +22,9 @@ class KernelType(enum.Enum):
def is_1d2d(self):
return self.value == 1
def is_nosf(self):
return self.value == 2
class MajorTypeAB(enum.Enum):
KMajor = 0
@@ -44,7 +48,9 @@ def get_ue8m0_usage(kernel_type: KernelType) -> bool:
return kernel_type.is_1d1d()
def get_kernel_types() -> tuple:
def get_kernel_types(use_bf16: bool = False) -> tuple:
if use_bf16:
return (KernelType.KernelNoSF, )
return (KernelType.Kernel1D2D, ) if get_arch_major() == 9 else (KernelType.Kernel1D1D, KernelType.Kernel1D2D)
@@ -61,13 +67,13 @@ def get_major_ab(freeze_a: bool) -> tuple:
(MajorTypeAB.MNMajor, MajorTypeAB.KMajor), (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor)
def enumerate_normal() -> Generator:
for kernel_type in get_kernel_types():
def enumerate_normal(use_bf16: bool = False) -> Generator:
for kernel_type in get_kernel_types(use_bf16):
for m in (128, 4096):
for n, k in [(2112, 7168), (24576, 1536), (32768, 512), (7168, 16384), (4096, 7168), (7168, 2048)]:
for n, k in [(2112, 7168), (24576, 1536), (32768, 512), (7168, 16384), (4096, 7168), (7168, 2048), (129280, 7168)]:
for major_a, major_b in get_major_ab(False):
for out_dtype in get_out_dtype():
for accumulate in (False, ) if out_dtype == torch.bfloat16 or kernel_type.is_1d2d() else (False, True):
for accumulate in (False, ) if out_dtype == torch.bfloat16 or not kernel_type.is_1d1d() else (False, True):
yield kernel_type, m, n, k, major_a, major_b, accumulate, out_dtype
@@ -123,7 +129,7 @@ def enumerate_k_grouped_sf_layout():
def generate_normal(m: int, n: int, k: int,
major_a: MajorTypeAB, major_b: MajorTypeAB,
accumulate: bool, out_dtype: torch.dtype,
use_ue8m0: bool):
use_ue8m0: bool = False, use_bf16: bool = False):
a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16)
d = torch.randn((m, n), device='cuda', dtype=out_dtype) * 32 if accumulate else \
@@ -131,6 +137,11 @@ def generate_normal(m: int, n: int, k: int,
c = d if accumulate else None
ref_d = (a.float() @ b.float().t() + (c if accumulate else 0)).to(out_dtype)
if use_bf16:
a = a if major_a.is_k_major() else a.T.contiguous().T
b = b if major_b.is_k_major() else b.T.contiguous().T
return a, b, c, d, ref_d
a_fp8, b_fp8 = per_token_cast_to_fp8(a, use_ue8m0=use_ue8m0), per_block_cast_to_fp8(b, use_ue8m0=use_ue8m0)
a_fp8 = a_fp8 if major_a.is_k_major() else (a_fp8[0].T.contiguous().T, a_fp8[1])
b_fp8 = b_fp8 if major_b.is_k_major() else (b_fp8[0].T.contiguous().T, b_fp8[1])

6
tests/test_core.py → tests/test_fp8.py
View File

@@ -51,10 +51,10 @@ def test_gemm() -> None:
deep_gemm.fp8_gemm_nt(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast)
t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
print(f' > Perf (m={m:5}, n={n:5}, k={k:5}, {kernel_opt}, layout={major_opt}, {out_opt}, {acc_opt}):'
f' launch {(launch_end_t - launch_start_t) / 1e3:4.0f} us | {t * 1e6:4.0f} us | '
print(f' > Perf (m={m:5}, n={n:5}, k={k:5}, {kernel_opt}, layout={major_opt}, {out_opt}, {acc_opt}): '
f'launch {(launch_end_t - launch_start_t) / 1e3:4.0f} us | {t * 1e6:4.0f} us | '
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes(a, b, c, d) / 1e9 / t:4.0f} GB/s')
f'{(count_bytes(a, b, d) + count_bytes(c) * int(accumulate)) / 1e9 / t:4.0f} GB/s')
print()