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

Multiple updates and refactorings (#280)

Browse Source
This commit is contained in:
Zhean Xu
2026-01-16 17:06:52 +08:00
committed by GitHub
parent 3ccf40c53a
commit 0f5f266202
55 changed files with 2706 additions and 891 deletions
Download Patch File Download Diff File Expand all files Collapse all files

228
tests/generators.py
View File

@@ -1,12 +1,13 @@
import enum
import random
import torch
from typing import Generator, List
from typing import Generator, List, Optional, Tuple
from deep_gemm.testing import get_arch_major
from deep_gemm.utils import (
align, ceil_div,
per_token_cast_to_fp8, per_channel_cast_to_fp8, per_block_cast_to_fp8,
per_token_cast_to_fp4, transpose_packed_fp4,
get_mk_alignment_for_contiguous_layout
)
@@ -35,6 +36,51 @@ class MajorTypeAB(enum.Enum):
def is_mn_major(self):
return self.value == 1
class QuantConfig:
_legacy_quant_config = (128, 128, False, False)
def __init__(self, value: Tuple[int, int, bool, bool] = _legacy_quant_config):
self.gran_k_a, self.gran_k_b, self.is_fp4_a, self.is_fp4_b = value
def print(self):
print(f' > Testing with gran_k_a={self.gran_k_a}, gran_k_b={self.gran_k_b}, '
f'is_fp4_a={self.is_fp4_a}, is_fp4_b={self.is_fp4_b}')
def is_legacy(self) -> bool:
return (self.gran_k_a, self.gran_k_b, self.is_fp4_a, self.is_fp4_b) == self._legacy_quant_config
def get_recipes(self, is_wgrad: bool = False) -> Tuple[Tuple, Tuple, Tuple]:
recipe, recipe_a, recipe_b = None, None, None
if self.is_legacy():
recipe = (1, 1, 128) if is_wgrad else None
else:
recipe_a = (1, self.gran_k_a)
recipe_b = (1, self.gran_k_b) if self.is_fp4_b or is_wgrad else (self.gran_k_b, self.gran_k_b)
return recipe, recipe_a, recipe_b
def max_diff(self) -> float:
if self.is_fp4_a and self.is_fp4_b:
return 0.02
if self.is_fp4_a or self.is_fp4_b:
return 0.01
return 0.001
@staticmethod
def get_list_from_dtype(dtype: torch.dtype) -> List:
if dtype == torch.bfloat16:
return [None]
quant_config_list = [QuantConfig()]
if get_arch_major() == 10:
quant_config_list.append(QuantConfig((128, 32, False, True)))
return quant_config_list
def reset_seed(seed: int = 0):
random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
def get_ue8m0_usage(kernel_type: KernelType) -> bool:
@@ -60,9 +106,14 @@ def get_major_ab(allow_a_mn_major: bool, allow_b_mn_major: bool) -> Generator:
yield major_a, major_b
def get_psum_layout_usage() -> tuple:
return (False, True) if get_arch_major() == 10 else (False, )
def enumerate_normal(dtype: torch.dtype) -> Generator:
assert dtype in (torch.float8_e4m3fn, torch.bfloat16)
quant_config_list = QuantConfig.get_list_from_dtype(dtype)
fp32_output_nk = [(256, 7168), (129280, 7168)]
bf16_output_nk = [(2112, 7168), (576, 7168), (24576, 1536), (32768, 512), (7168, 16384), (4096, 7168), (7168, 2048)]
m_fwd_list, m_bwd_list = [1, 128, 4096], [4096, ]
@@ -73,39 +124,61 @@ def enumerate_normal(dtype: torch.dtype) -> Generator:
nk_list += fp32_output_nk
for kernel_type in get_kernel_types(dtype):
# Forward
for m in m_fwd_list:
for i in range(len(nk_list)):
n, k = nk_list[i]
out_dtype = torch.bfloat16 if i < len(bf16_output_nk) else torch.float
yield kernel_type, m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, False, out_dtype
for quant_config in quant_config_list:
if len(quant_config_list) > 1:
quant_config.print()
reset_seed()
# Backward
for m in m_bwd_list:
for n, k in nk_list:
override_major = MajorTypeAB.MNMajor
override_kernel_type = kernel_type
if get_arch_major() == 9 and dtype == torch.float8_e4m3fn:
override_major = MajorTypeAB.KMajor
override_kernel_type = KernelType.Kernel1D1D
yield kernel_type, m, k, n, MajorTypeAB.KMajor, override_major, False, torch.bfloat16 # Dgrad
yield override_kernel_type, n, m, k, override_major, override_major, True, torch.float # Wgrad
yield override_kernel_type, n, m, k, override_major, override_major, False, torch.bfloat16 # Wgrad
# Forward
for m in m_fwd_list:
for i in range(len(nk_list)):
n, k = nk_list[i]
out_dtype = torch.bfloat16 if i < len(bf16_output_nk) else torch.float
yield kernel_type, quant_config, m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, False, out_dtype
# Backward
for m in m_bwd_list:
for n, k in nk_list:
override_major = MajorTypeAB.MNMajor
override_kernel_type = kernel_type
if get_arch_major() == 9 and dtype == torch.float8_e4m3fn:
override_major = MajorTypeAB.KMajor
override_kernel_type = KernelType.Kernel1D1D
yield kernel_type, quant_config, m, k, n, MajorTypeAB.KMajor, override_major, False, torch.bfloat16 # Dgrad
yield override_kernel_type, quant_config, n, m, k, override_major, override_major, True, torch.float # Wgrad
yield override_kernel_type, quant_config, n, m, k, override_major, override_major, False, torch.bfloat16 # Wgrad
def enumerate_m_grouped_contiguous(dtype: torch.dtype) -> Generator:
quant_config_list = QuantConfig.get_list_from_dtype(dtype)
m_group_list = [(4, 8192), (8, 4096)]
n_k_list = [(6144, 7168), (7168, 3072), (4096, 4096), (4096, 2048)]
for kernel_type in get_kernel_types(dtype):
for num_groups, expected_m_per_group, n, k in ((4, 8192, 4096, 7168), (4, 8192, 7168, 2048), (8, 4096, 4096, 7168), (8, 4096, 7168, 2048)):
for major_a, major_b in get_major_ab(False, get_arch_major() != 9 or dtype != torch.float8_e4m3fn):
yield kernel_type, num_groups, expected_m_per_group, n, k, major_a, major_b
for quant_config in quant_config_list:
if len(quant_config_list) > 1:
quant_config.print()
for use_psum_layout in get_psum_layout_usage():
reset_seed()
for num_groups, expected_m_per_group in m_group_list:
for n, k in n_k_list:
for major_a, major_b in get_major_ab(False, get_arch_major() != 9 or dtype != torch.float8_e4m3fn):
yield kernel_type, quant_config, num_groups, expected_m_per_group, n, k, major_a, major_b, use_psum_layout
def enumerate_m_grouped_masked(dtype: torch.dtype) -> Generator:
quant_config_list = QuantConfig.get_list_from_dtype(dtype)
max_m = 4096
m_group_list = [(6, 1024), (32, 192), (32, 50)]
n_k_list = [(6144, 7168), (7168, 3072), (4096, 4096), (4096, 2048)]
for kernel_type in get_kernel_types(dtype):
for num_groups, m in ((1, 1024), (2, 512), (4, 256)):
for n, k in ((4096, 7168), (7168, 2048), ):
yield kernel_type, num_groups, max_m, m, n, k
for quant_config in quant_config_list:
if len(quant_config_list) > 1:
quant_config.print()
for use_psum_layout in get_psum_layout_usage():
reset_seed()
for num_groups, m in m_group_list:
for n, k in n_k_list:
yield kernel_type, quant_config, num_groups, max_m, m, n, k, use_psum_layout
def enumerate_k_grouped_contiguous(dtype: torch.dtype):
@@ -145,11 +218,46 @@ def enumerate_transpose():
yield mn + delta, k
def cast_fp8_fp4_with_major(x: torch.Tensor, major: MajorTypeAB, gran_k: int, is_fp4: bool,
use_ue8m0: bool, use_block_cast_for_fp8: bool = False):
if is_fp4:
x_fp4 = per_token_cast_to_fp4(x, use_ue8m0=use_ue8m0, gran_k=gran_k)
x = x_fp4 if major.is_k_major() else (transpose_packed_fp4(x_fp4[0]).T, x_fp4[1])
else:
x_fp8 = per_block_cast_to_fp8(x, use_ue8m0=use_ue8m0, gran_k=gran_k) if use_block_cast_for_fp8 \
else per_token_cast_to_fp8(x, use_ue8m0=use_ue8m0, gran_k=gran_k)
x = x_fp8 if major.is_k_major() else (x_fp8[0].T.contiguous().T, x_fp8[1])
return x
def grouped_cast_fp8_fp4_with_major(x: torch.Tensor, major: MajorTypeAB, gran_k: int, is_fp4: bool,
use_ue8m0: bool, use_block_cast_for_fp8: bool = False):
num_groups, mn, k = x.size()
if is_fp4:
x_fp4 = (torch.empty((num_groups, mn, k // 2), device='cuda', dtype=torch.uint8) if major.is_k_major() else \
torch.empty((num_groups, k, mn // 2), device='cuda', dtype=torch.uint8),
torch.empty((num_groups, mn, ceil_div(k, gran_k)), device='cuda', dtype=torch.float))
for i in range(num_groups):
x_i_fp4 = per_token_cast_to_fp4(x[i], use_ue8m0=use_ue8m0, gran_k=gran_k)
x_fp4[0][i], x_fp4[1][i] = x_i_fp4 if major.is_k_major() else (transpose_packed_fp4(x_i_fp4[0]), x_i_fp4[1])
x = x_fp4 if major.is_k_major() else (x_fp4[0].mT, x_fp4[1])
else:
x_fp8 = (torch.empty_like(x, dtype=torch.float8_e4m3fn),
torch.empty((num_groups, ceil_div(mn, gran_k), ceil_div(k, gran_k)), device='cuda', dtype=torch.float) if use_block_cast_for_fp8 \
else torch.empty((num_groups, mn, ceil_div(k, gran_k)), device='cuda', dtype=torch.float))
for i in range(num_groups):
x_fp8[0][i], x_fp8[1][i] = per_block_cast_to_fp8(x[i], use_ue8m0=use_ue8m0, gran_k=gran_k) if use_block_cast_for_fp8 \
else per_token_cast_to_fp8(x[i], use_ue8m0=use_ue8m0, gran_k=gran_k)
x = x_fp8 if major.is_k_major() else (x_fp8[0].mT.contiguous().mT, x_fp8[1])
return x
def generate_normal(m: int, n: int, k: int,
major_a: MajorTypeAB, major_b: MajorTypeAB,
accumulate: bool, out_dtype: torch.dtype,
kernel_type: KernelType,
use_ue8m0: bool = False, use_bf16: bool = False):
use_ue8m0: bool = False, use_bf16: bool = False,
quant_config: Optional[QuantConfig] = None):
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 \
@@ -161,25 +269,28 @@ def generate_normal(m: int, n: int, k: int,
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
quant_config = QuantConfig() if quant_config is None else quant_config
a = cast_fp8_fp4_with_major(a, major_a, quant_config.gran_k_a, quant_config.is_fp4_a, use_ue8m0)
b = cast_fp8_fp4_with_major(b, major_b, quant_config.gran_k_b, quant_config.is_fp4_b, use_ue8m0,
use_block_cast_for_fp8=not (kernel_type.is_1d1d() and accumulate))
a_fp8 = per_token_cast_to_fp8(a, use_ue8m0=use_ue8m0)
b_fp8 = per_token_cast_to_fp8(b, use_ue8m0=use_ue8m0) if kernel_type.is_1d1d() and accumulate \
else 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])
return a_fp8, b_fp8, c, d, ref_d
return a, b, c, d, ref_d
def generate_m_grouped_contiguous(num_groups: int, expected_m_per_group: int, n: int, k: int,
major_a: MajorTypeAB, major_b: MajorTypeAB,
use_ue8m0: bool = False, use_bf16: bool = False):
use_ue8m0: bool = False, use_bf16: bool = False,
use_psum_layout: bool = False,
quant_config: Optional[QuantConfig] = None):
actual_ms = [int(expected_m_per_group * random.uniform(0.7, 1.3)) for _ in range(num_groups)]
aligned_ms = [align(actual_m, get_mk_alignment_for_contiguous_layout()) for actual_m in actual_ms]
m = sum(aligned_ms)
a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
b = torch.randn((num_groups, n, k), device='cuda', dtype=torch.bfloat16)
m_indices = torch.empty(m, device='cuda', dtype=torch.int32)
grouped_layout = torch.empty(num_groups, device='cuda', dtype=torch.int32) if use_psum_layout \
else torch.empty(m, device='cuda', dtype=torch.int32)
d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)
ref_d = torch.randn((m, n), device='cuda', dtype=torch.bfloat16)
@@ -187,48 +298,61 @@ def generate_m_grouped_contiguous(num_groups: int, expected_m_per_group: int, n:
for i, (actual_m, aligned_m) in enumerate(zip(actual_ms, aligned_ms)):
actual_end = start + actual_m
aligned_end = start + aligned_m
m_indices[start:actual_end] = i
m_indices[actual_end:aligned_end] = -1
ref_d[start:aligned_end] = a[start:aligned_end] @ b[i].t()
if use_psum_layout:
grouped_layout[i] = actual_end
else:
grouped_layout[start: actual_end] = i
grouped_layout[actual_end: aligned_end] = -1
a[actual_end: aligned_end] = 0
ref_d[start: aligned_end] = a[start: aligned_end] @ b[i].t()
start = aligned_end
ref_d = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(ref_d), ref_d)
if use_bf16:
b = b if major_b.is_k_major() else b.mT.contiguous().mT
return m, a, b, m_indices, d, ref_d
return m, a, b, grouped_layout, d, ref_d
assert major_a.is_k_major()
a_fp8 = per_token_cast_to_fp8(a, use_ue8m0=use_ue8m0)
b_fp8 = (torch.empty_like(b, dtype=torch.float8_e4m3fn),
torch.empty((num_groups, ceil_div(n, 128), ceil_div(k, 128)), device='cuda', dtype=torch.float))
quant_config = QuantConfig() if quant_config is None else quant_config
a = cast_fp8_fp4_with_major(a, major_a, quant_config.gran_k_a, quant_config.is_fp4_a, use_ue8m0)
b = grouped_cast_fp8_fp4_with_major(b, major_b, quant_config.gran_k_b, quant_config.is_fp4_b, use_ue8m0, use_block_cast_for_fp8=True)
return m, a, b, grouped_layout, d, ref_d
def layout_masked_to_psum(x: torch.Tensor, psum_m: torch.Tensor):
num_groups, max_m, _ = x.size()
x_psum = torch.empty_like(x).view(num_groups * max_m, -1)
last_psum_m = 0
for i in range(num_groups):
b_fp8[0][i], b_fp8[1][i] = per_block_cast_to_fp8(b[i], use_ue8m0=use_ue8m0)
b_fp8 = b_fp8 if major_b.is_k_major() else (b_fp8[0].mT.contiguous().mT, b_fp8[1])
return m, a_fp8, b_fp8, m_indices, d, ref_d
x_psum[last_psum_m: psum_m[i]] = x[i, :psum_m[i] - last_psum_m]
last_psum_m = align(psum_m[i], 128)
return x_psum
def generate_m_grouped_masked(num_groups: int, max_m: int, expected_m_per_group: int, n: int, k: int,
use_ue8m0: bool = False, use_bf16: bool = False):
use_ue8m0: bool = False, use_bf16: bool = False,
use_psum_layout: bool = False,
quant_config: Optional[QuantConfig] = None):
a = torch.randn((num_groups, max_m, k), device='cuda', dtype=torch.bfloat16)
b = torch.randn((num_groups, n, k), device='cuda', dtype=torch.bfloat16)
d = torch.empty((num_groups, max_m, n), device='cuda', dtype=torch.bfloat16)
ref_d = torch.einsum('gmk,gnk->gmn', a, b)
masked_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int)
psum_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int)
for j in range(num_groups):
masked_m[j] = int(expected_m_per_group * random.uniform(0.7, 1.3))
psum_m[j] = (0 if j == 0 else align(psum_m[j - 1], 128)) + masked_m[j]
assert masked_m.amax().item() <= max_m
if use_bf16:
return a, b, masked_m, d, ref_d
return a, b, masked_m, psum_m, d, ref_d
a_fp8 = (torch.empty_like(a, dtype=torch.float8_e4m3fn), torch.empty((num_groups, max_m, ceil_div(k, 128)), device='cuda', dtype=torch.float))
b_fp8 = (torch.empty_like(b, dtype=torch.float8_e4m3fn), torch.empty((num_groups, ceil_div(n, 128), ceil_div(k, 128)), device='cuda', dtype=torch.float))
for i in range(num_groups):
a_fp8[0][i], a_fp8[1][i] = per_token_cast_to_fp8(a[i], use_ue8m0=use_ue8m0)
b_fp8[0][i], b_fp8[1][i] = per_block_cast_to_fp8(b[i], use_ue8m0=use_ue8m0)
quant_config = QuantConfig() if quant_config is None else quant_config
a = grouped_cast_fp8_fp4_with_major(a, MajorTypeAB.KMajor, quant_config.gran_k_a, quant_config.is_fp4_a, use_ue8m0)
b = grouped_cast_fp8_fp4_with_major(b, MajorTypeAB.KMajor, quant_config.gran_k_b, quant_config.is_fp4_b, use_ue8m0, use_block_cast_for_fp8=True)
return a_fp8, b_fp8, masked_m, d, ref_d
return a, b, masked_m, psum_m, d, ref_d
def generate_k_grouped_contiguous(num_groups: int, m: int, n: int, major_a: MajorTypeAB, major_b: MajorTypeAB, ks: List[int],

14
tests/test_attention.py
View File

@@ -1,12 +1,14 @@
import dataclasses
import random
import torch
from typing import Tuple
from typing import Tuple, List
import deep_gemm
from deep_gemm.testing import (
bench_kineto,
calc_diff, count_bytes,
ignore_env, get_arch_major
ignore_env, get_arch_major,
test_filter
)
from deep_gemm.utils import ceil_div, per_custom_dims_cast_to_fp8
@@ -154,7 +156,7 @@ def test_mqa_logits():
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
logits = logits.masked_fill(neginf_mask, 0)
diff = calc_diff(logits, ref_logits)
assert diff < 1e-3, f"{diff=}"
assert diff < 1e-3, f'{diff=}'
else:
ref_cost = ref_fp8_mqa_logits(q=q, kv=kv, weights=weights, cu_seqlen_ks=ks, cu_seqlen_ke=ke, cost_only=True)
@@ -204,8 +206,6 @@ def ref_fp8_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor,
def test_paged_mqa_logits():
# TODO: fully refactor with PyTest
print('Testing FP8 Paged MQA Logits:')
max_model_len = 111 * 1000
for is_context_lens_2d in (False, True):
@@ -264,7 +264,7 @@ def test_paged_mqa_logits():
else:
t, clean_t = bench_kineto(lambda: deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens, block_tables, schedule_metadata, max_model_len, clean_logits=True),
('fp8_paged_mqa_logits', 'clean_logits'))
clean_bytes = (batch_size * next_n * max_model_len - neginf_mask.sum().item()) * 4 + count_bytes(context_lens)
clean_bytes = (batch_size * next_n * max_model_len - neginf_mask.sum().item()) * 4 + count_bytes(context_lens)
print(f' > BSZ={batch_size:3}, NextN={next_n:1}, H={heads:2}, D={index_dim:2}, L={avg_kv:6}: '
f'{tflops / t:4.0f} TFLOPS, {t * 1e6:3.0f} us, '
f'{(input_bytes + output_bytes) / t / 1e9:4.0f} GB/s', end='')
@@ -273,6 +273,8 @@ def test_paged_mqa_logits():
print()
if __name__ == '__main__':
torch.manual_seed(0)
random.seed(0)

87
tests/test_bf16.py
View File

@@ -9,7 +9,7 @@ from deep_gemm.testing import (
calc_diff, count_bytes
)
from generators import (
get_arch_major,
get_arch_major, layout_masked_to_psum, align,
enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous,
generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous
)
@@ -18,11 +18,7 @@ from generators import (
def test_gemm() -> None:
print('Testing GEMM:')
scores = []
for kernel_type, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(torch.bfloat16):
# TODO: support accumulation for SM90 BF16 GEMM
if get_arch_major() == 9 and accumulate:
continue
for kernel_type, _, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(torch.bfloat16):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
out_opt = 'FP32' if out_dtype == torch.float else 'BF16'
@@ -56,29 +52,30 @@ def test_gemm() -> None:
def test_m_grouped_gemm_contiguous() -> None:
print('Testing m-grouped contiguous GEMM:')
for _, num_groups, expected_m_per_group, n, k, major_a, major_b in enumerate_m_grouped_contiguous(torch.bfloat16):
for _, _, num_groups, expected_m_per_group, n, k, major_a, major_b, use_psum_layout in enumerate_m_grouped_contiguous(torch.bfloat16):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
for test_alias in (False, True):
m, a, b, m_indices, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_bf16=True)
m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b,
use_bf16=True, use_psum_layout=use_psum_layout)
func_name = f"m_grouped_bf16_gemm_{(major_opt.lower() if test_alias else 'nt')}_contiguous"
if test_alias:
assert major_a.is_k_major()
b = b if major_b.is_k_major() else b.mT
assert a[0].is_contiguous() and b[0].is_contiguous()
getattr(deep_gemm, func_name)(a, b, d, m_indices)
d = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d)
getattr(deep_gemm, func_name)(a, b, d, grouped_layout, use_psum_layout=use_psum_layout)
diff = calc_diff(d, ref_d)
assert diff < 1e-5, f'{m=}, {n=}, {k=}, {major_opt}, {diff:.5f}, alias={test_alias}'
m, a, b, m_indices, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_bf16=True)
m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b,
use_bf16=True, use_psum_layout=use_psum_layout)
# noinspection PyShadowingNames
def test_func():
deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, m_indices)
deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, use_psum_layout=use_psum_layout)
t = bench_kineto(test_func, 'bf16_gemm', suppress_kineto_output=True)
print(f' > Perf ({num_groups=}, m={m:5}, n={n:5}, k={k:5}, layout={major_opt}): '
print(f' > Perf ({num_groups=}, m={m:5}, n={n:5}, k={k:5}, layout={major_opt}, psum={use_psum_layout}): '
f'{t * 1e6:4.0f} us | '
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s')
@@ -89,29 +86,52 @@ def test_m_grouped_gemm_masked() -> None:
print('Testing m-grouped masked GEMM:')
# TODO: when the actual `m` is greater than `expected_m_per_group`, efficiency may significantly decrease.
for _, num_groups, max_m, expected_m_per_group, n, k in enumerate_m_grouped_masked(torch.bfloat16):
# Test correctness
for i in range(10):
a, b, masked_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k, use_bf16=True)
deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group)
for _, _, num_groups, max_m, expected_m_per_group, n, k, use_psum_layout in enumerate_m_grouped_masked(torch.bfloat16):
num_tests = 8
sum_t, max_t = 0, 0
sum_ops, sum_bytes = 0, 0
for i in range(num_tests):
a, b, masked_m, psum_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k,
use_bf16=True, use_psum_layout=use_psum_layout)
if use_psum_layout:
a_psum = layout_masked_to_psum(a, psum_m)
d_psum = layout_masked_to_psum(d, psum_m)
# noinspection PyShadowingNames
def test_func():
if use_psum_layout:
deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a_psum, b, d_psum, psum_m,
use_psum_layout=True, expected_m_for_psum_layout=expected_m_per_group)
else:
deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group)
test_func()
for j in range(num_groups):
diff = calc_diff(d[j, :masked_m[j].item()], ref_d[j, :masked_m[j].item()])
if masked_m[j].item() == 0:
continue
if use_psum_layout:
d_slice = d_psum[: psum_m[j]] if j == 0 else d_psum[align(psum_m[j - 1], 128): psum_m[j]]
else:
d_slice = d[j, :masked_m[j].item()]
diff = calc_diff(d_slice, ref_d[j, :masked_m[j].item()])
assert diff < 1e-5, f'{max_m=}, {n=}, {k=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}'
# Construct full cases
a, b, masked_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k, use_bf16=True)
# noinspection PyShadowingNames
def test_func():
deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group)
# Test performance with fixed shapes
valid_m = masked_m.sum().item()
t = bench_kineto(test_func, 'bf16_gemm', suppress_kineto_output=True)
# Test performance with fixed shapes
valid_m = masked_m.sum().item()
t = bench_kineto(test_func, 'bf16_gemm', suppress_kineto_output=True)
print(f' > Perf ({num_groups=}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}): '
f'{t * 1e6:4.0f} us | '
f'{2 * valid_m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{(count_bytes(a, d) * valid_m / (max_m * num_groups) + count_bytes(b)) / 1e9 / t:4.0f} GB/s')
sum_t += t
max_t = max(max_t, t)
sum_ops += 2 * valid_m * n * k
sum_bytes += count_bytes(a, d) * valid_m / (max_m * num_groups) + count_bytes(b)
print(f' > Perf (num_groups={num_groups:2}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}, '
f'psum={1 if use_psum_layout else 0}): '
f'{sum_t / num_tests * 1e6:4.0f} us (max: {max_t * 1e6:3.0f} us) | '
f'{sum_ops / sum_t / 1e12:4.0f} TFLOPS | '
f'{sum_bytes / sum_t / 1e9:4.0f} GB/s')
print()
@@ -148,7 +168,7 @@ def test_k_grouped_gemm_contiguous() -> None:
def test_cublaslt_gemm() -> None:
print('Testing cuBLASLt GEMM:')
for kernel_type, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(dtype=torch.bfloat16):
for kernel_type, _, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(dtype=torch.bfloat16):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
out_opt = 'FP32' if out_dtype == torch.float else 'BF16'
@@ -159,7 +179,8 @@ def test_cublaslt_gemm() -> None:
diff = calc_diff(d, ref_d)
assert diff < 6e-7, f'{diff=}, ({m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=})'
t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), 'nvjet', suppress_kineto_output=True,)
t_nvjet, t_gemv, t_gemm = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True)
t = t_nvjet + t_gemv + t_gemm
print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): '
f'{t * 1e6:5.0f} us | '
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '

1
tests/test_einsum.py
View File

@@ -80,7 +80,6 @@ def test_bhd_hdr_bhr():
print()
@test_filter(lambda: get_arch_major() >= 10)
def test_fp8_bhr_hdr_bhd(use_ue8m0: bool = True):
print('Testing FP8 "bhr, hdr -> bhd":')
for h, r, d in [(8, 4096, 1024)]:

175
tests/test_fp8.py
View File

@@ -1,175 +0,0 @@
import copy
import numpy as np
import random
import torch
import deep_gemm
from deep_gemm.testing import (
bench_kineto,
calc_diff, count_bytes,
ignore_env, get_arch_major
)
from generators import (
KernelType, get_ue8m0_usage,
enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous,
generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous
)
@ignore_env('DG_JIT_PTXAS_CHECK', lambda: get_arch_major() == 9)
def test_gemm() -> None:
print('Testing GEMM:')
scores = []
for kernel_type, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(torch.float8_e4m3fn):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
out_opt = 'FP32' if out_dtype == torch.float else 'BF16'
acc_opt = f'acc={int(accumulate)}'
kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D'
use_ue8m0 = get_ue8m0_usage(kernel_type)
disable_ue8m0_cast = not use_ue8m0
recipe = (1, 1, 128) if kernel_type.is_1d1d() and accumulate else None
for test_alias in (False, True):
a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0)
func_name = f'fp8_gemm_{major_opt.lower() if test_alias else "nt"}'
if test_alias:
a = a if major_a.is_k_major() else (a[0].T, a[1].T)
b = b if major_b.is_k_major() else (b[0].T, b[1].T)
assert a[0].is_contiguous() and b[0].is_contiguous()
getattr(deep_gemm, func_name)(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe)
diff = calc_diff(d, ref_d)
assert diff < 0.001, (f'{m=}, {n=}, {k=}, {kernel_opt}, {major_opt=}, {accumulate=}, {out_dtype=}, '
f'{diff:.5f}, alias={test_alias}')
a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0)
t = bench_kineto(lambda: deep_gemm.fp8_gemm_nt(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe),
'fp8_gemm', suppress_kineto_output=True)
cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a[0], b[0], d, c=c), ('nvjet', 'reduce'), suppress_kineto_output=True)
print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, {kernel_opt}, layout={major_opt}, {out_opt}, {acc_opt}): '
f'{t * 1e6:6.1f} us | {2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{(count_bytes(a, b, d) + count_bytes(c) * int(accumulate)) / 1e9 / t:4.0f} GB/s | '
f'{(cublas_t + split_k_t) / t:.2f}x cuBLAS')
if cublas_t > 0:
scores.append((cublas_t + split_k_t) / t)
print(f"Average speedup over cuBLASLt: {float(np.prod(scores)) ** (1.0 / len(scores)):.3f}x\n")
def test_m_grouped_gemm_contiguous() -> None:
print('Testing m-grouped contiguous GEMM:')
for kernel_type, num_groups, expected_m_per_group, n, k, major_a, major_b in enumerate_m_grouped_contiguous(dtype=torch.float8_e4m3fn):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D'
use_ue8m0 = get_ue8m0_usage(kernel_type)
disable_ue8m0_cast = not use_ue8m0
for test_alias in (False, True):
m, a, b, m_indices, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_ue8m0=use_ue8m0)
func_name = f"m_grouped_fp8_gemm_{(major_opt.lower() if test_alias else 'nt')}_contiguous"
if test_alias:
assert major_a.is_k_major()
b = b if major_b.is_k_major() else (b[0].mT, b[1].mT)
assert a[0].is_contiguous() and b[0].is_contiguous()
getattr(deep_gemm, func_name)(a, b, d, m_indices, disable_ue8m0_cast=disable_ue8m0_cast)
d = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d)
diff = calc_diff(d, ref_d)
assert diff < 0.001, f'{m=}, {n=}, {k=}, {major_opt}, {kernel_opt}, {diff:.5f}, alias={test_alias}'
m, a, b, m_indices, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_ue8m0=use_ue8m0)
# noinspection PyShadowingNames
def test_func():
deep_gemm.m_grouped_fp8_gemm_nt_contiguous(a, b, d, m_indices, disable_ue8m0_cast=disable_ue8m0_cast)
t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
print(f' > Perf ({num_groups=}, m={m:5}, n={n:6}, k={k:5}, {kernel_opt}, layout={major_opt}): '
f'{t * 1e6:4.0f} us | '
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s')
print()
def test_m_grouped_gemm_masked() -> None:
print('Testing m-grouped masked GEMM:')
# TODO: when the actual `m` is greater than `expected_m_per_group`, efficiency may significantly decrease.
for kernel_type, num_groups, max_m, expected_m_per_group, n, k in enumerate_m_grouped_masked(torch.float8_e4m3fn):
kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D'
use_ue8m0 = get_ue8m0_usage(kernel_type)
disable_ue8m0_cast = not use_ue8m0
# Test correctness
for i in range(10):
a, b, masked_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k, use_ue8m0=use_ue8m0)
deep_gemm.m_grouped_fp8_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group, disable_ue8m0_cast=disable_ue8m0_cast)
for j in range(num_groups):
if masked_m[j].item() == 0:
continue
diff = calc_diff(d[j, :masked_m[j].item()], ref_d[j, :masked_m[j].item()])
assert diff < 0.001, f'{max_m=}, {n=}, {k=}, {j=}, masked_m={masked_m[j]}, {kernel_opt}, {num_groups=}, {diff:.5f}'
# Construct full cases
a, b, masked_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k, use_ue8m0=use_ue8m0)
# noinspection PyShadowingNames
def test_func():
deep_gemm.m_grouped_fp8_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group, disable_ue8m0_cast=disable_ue8m0_cast)
# Test performance with fixed shapes
valid_m = masked_m.sum().item()
t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
print(f' > Perf ({num_groups=}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}, {kernel_opt}): '
f'{t * 1e6:4.0f} us | '
f'{2 * valid_m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{(count_bytes(a, d) * valid_m / (max_m * num_groups) + count_bytes(b)) / 1e9 / t:4.0f} GB/s')
print()
def test_k_grouped_gemm_contiguous() -> None:
print('Testing k-grouped contiguous GEMM:')
k_grouped_fp8_gemm_contiguous = deep_gemm.k_grouped_fp8_gemm_nt_contiguous if get_arch_major() == 9 \
else deep_gemm.k_grouped_fp8_gemm_tn_contiguous
for num_groups, m, n, major_a, major_b, ks, expected_k_per_group in enumerate_k_grouped_contiguous(torch.float8_e4m3fn):
use_ue8m0 = get_ue8m0_usage(KernelType.Kernel1D1D)
for test_empty_groups in (False, True):
new_ks = copy.deepcopy(ks)
if test_empty_groups and len(ks) > 1:
new_ks[random.randint(0, num_groups - 1)] = 0
k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, new_ks, use_ue8m0=use_ue8m0)
new_ks_tensor = torch.tensor(new_ks, dtype=torch.int, device='cuda')
k_grouped_fp8_gemm_contiguous(a, b, d, new_ks, new_ks_tensor, c)
diff = calc_diff(d, ref_d)
assert diff < 0.001, f'{m=}, {n=}, {k=}, {ks=}, {diff:.5f}'
# Test performance
k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_ue8m0=use_ue8m0)
ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda')
# noinspection PyShadowingNames
def test_func():
k_grouped_fp8_gemm_contiguous(a, b, d, ks, ks_tensor, c)
t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
print(f' > Perf ({num_groups=:2}, m={m:5}, n={n:5}, k={k:5}): '
f'{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')
print()
if __name__ == '__main__':
torch.manual_seed(0)
random.seed(0)
print('Library path:')
print(f' > {deep_gemm.__path__}\n')
test_gemm()
test_m_grouped_gemm_contiguous()
test_m_grouped_gemm_masked()
test_k_grouped_gemm_contiguous()

207
tests/test_fp8_fp4.py Normal file
View File

@@ -0,0 +1,207 @@
import copy
import numpy as np
import random
import torch
import deep_gemm
from deep_gemm.testing import (
bench_kineto,
calc_diff, count_bytes,
ignore_env, get_arch_major
)
from generators import (
KernelType, get_ue8m0_usage, layout_masked_to_psum, align,
enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous,
generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous
)
@ignore_env('DG_JIT_PTXAS_CHECK', lambda: get_arch_major() == 9)
def test_gemm() -> None:
print('Testing GEMM:')
scores = []
for kernel_type, quant_config, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(torch.float8_e4m3fn):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
out_opt = 'FP32' if out_dtype == torch.float else 'BF16'
acc_opt = f'acc={int(accumulate)}'
kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D'
use_ue8m0 = get_ue8m0_usage(kernel_type)
disable_ue8m0_cast = not use_ue8m0
recipe, recipe_a, recipe_b = quant_config.get_recipes(is_wgrad=(kernel_type.is_1d1d() and accumulate))
for test_alias in (False, True):
a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0, quant_config=quant_config)
func_name = f'fp8_fp4_gemm_{major_opt.lower() if test_alias else "nt"}'
if test_alias:
a = a if major_a.is_k_major() else (a[0].T, a[1].T)
b = b if major_b.is_k_major() else (b[0].T, b[1].T)
assert a[0].is_contiguous() and b[0].is_contiguous()
getattr(deep_gemm, func_name)(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b)
diff = calc_diff(d, ref_d)
assert diff < quant_config.max_diff(), (f'{m=}, {n=}, {k=}, {kernel_opt}, {major_opt=}, {accumulate=}, {out_dtype=}, '
f'{diff:.5f}, alias={test_alias}')
a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0, quant_config=quant_config)
t = bench_kineto(lambda: deep_gemm.fp8_fp4_gemm_nt(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b),
'fp8_gemm', suppress_kineto_output=True)
cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a[0], b[0], d, c=c), ('nvjet', 'reduce'), suppress_kineto_output=True) \
if not quant_config.is_fp4_a and not quant_config.is_fp4_b else (0, 0)
print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, {kernel_opt}, layout={major_opt}, {out_opt}, {acc_opt}): '
f'{t * 1e6:6.1f} us | {2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{(count_bytes(a, b, d) + count_bytes(c) * int(accumulate)) / 1e9 / t:4.0f} GB/s | '
f'{(cublas_t + split_k_t) / t:.2f}x cuBLAS')
if cublas_t > 0:
scores.append((cublas_t + split_k_t) / t)
print(f"Average FP8xFP8 GEMM speedup over cuBLASLt: {float(np.prod(scores)) ** (1.0 / len(scores)):.3f}x\n")
def test_m_grouped_gemm_contiguous() -> None:
print('Testing m-grouped contiguous GEMM:')
for kernel_type, quant_config, num_groups, expected_m_per_group, n, k, major_a, major_b, use_psum_layout in enumerate_m_grouped_contiguous(dtype=torch.float8_e4m3fn):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D'
use_ue8m0 = get_ue8m0_usage(kernel_type)
disable_ue8m0_cast = not use_ue8m0
recipe, recipe_a, recipe_b = quant_config.get_recipes()
for test_alias in (False, True):
m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b,
use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout,
quant_config=quant_config)
func_name = f"m_grouped_fp8_fp4_gemm_{(major_opt.lower() if test_alias else 'nt')}_contiguous"
if test_alias:
assert major_a.is_k_major()
b = b if major_b.is_k_major() else (b[0].mT, b[1].mT)
assert a[0].is_contiguous() and b[0].is_contiguous()
getattr(deep_gemm, func_name)(a, b, d, grouped_layout, disable_ue8m0_cast=disable_ue8m0_cast, use_psum_layout=use_psum_layout,
recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b)
diff = calc_diff(d, ref_d)
assert diff < quant_config.max_diff(), f'{m=}, {n=}, {k=}, {major_opt}, {kernel_opt}, {diff:.5f}, alias={test_alias}'
m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b,
use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout,
quant_config=quant_config)
# noinspection PyShadowingNames
def test_func():
deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, grouped_layout, disable_ue8m0_cast=disable_ue8m0_cast, use_psum_layout=use_psum_layout,
recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b)
t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
print(f' > Perf ({num_groups=}, m={m:5}, n={n:6}, k={k:5}, {kernel_opt}, layout={major_opt}, psum={use_psum_layout}): '
f'{t * 1e6:4.0f} us | '
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s')
print()
def test_m_grouped_gemm_masked() -> None:
print('Testing m-grouped masked GEMM:')
# TODO: when the actual `m` is greater than `expected_m_per_group`, efficiency may significantly decrease.
for kernel_type, quant_config, num_groups, max_m, expected_m_per_group, n, k, use_psum_layout in enumerate_m_grouped_masked(torch.float8_e4m3fn):
kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D'
use_ue8m0 = get_ue8m0_usage(kernel_type)
disable_ue8m0_cast = not use_ue8m0
recipe, recipe_a, recipe_b = quant_config.get_recipes()
num_tests = 8
sum_t, max_t = 0, 0
sum_ops, sum_bytes = 0, 0
for i in range(num_tests):
a, b, masked_m, psum_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k,
use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout,
quant_config=quant_config)
if use_psum_layout:
a_psum = (layout_masked_to_psum(a[0], psum_m), layout_masked_to_psum(a[1], psum_m))
d_psum = layout_masked_to_psum(d, psum_m)
# noinspection PyShadowingNames
def test_func():
if use_psum_layout:
deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(a_psum, b, d_psum, psum_m, disable_ue8m0_cast=disable_ue8m0_cast,
use_psum_layout=True, expected_m_for_psum_layout=expected_m_per_group,
recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b)
else:
deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group, disable_ue8m0_cast=disable_ue8m0_cast,
recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b)
test_func()
for j in range(num_groups):
if masked_m[j].item() == 0:
continue
if use_psum_layout:
d_slice = d_psum[: psum_m[j]] if j == 0 else d_psum[align(psum_m[j - 1], 128): psum_m[j]]
else:
d_slice = d[j, :masked_m[j].item()]
diff = calc_diff(d_slice, ref_d[j, :masked_m[j].item()])
assert diff < quant_config.max_diff(), f'{max_m=}, {n=}, {k=}, {j=}, masked_m={masked_m[j]}, {kernel_opt}, {num_groups=}, {diff:.5f}'
# Test performance with fixed shapes
valid_m = masked_m.sum().item()
t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
sum_t += t
max_t = max(max_t, t)
sum_ops += 2 * valid_m * n * k
sum_bytes += count_bytes(a, d) * valid_m / (max_m * num_groups) + count_bytes(b)
print(f' > Perf (num_groups={num_groups:2}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}, '
f'{kernel_opt}, psum={1 if use_psum_layout else 0}): '
f'{sum_t / num_tests * 1e6:4.0f} us (max: {max_t * 1e6:3.0f} us) | '
f'{sum_ops / sum_t / 1e12:4.0f} TFLOPS | '
f'{sum_bytes / sum_t / 1e9:4.0f} GB/s')
print()
@ignore_env('DG_JIT_PTXAS_CHECK', lambda: get_arch_major() == 9)
def test_k_grouped_gemm_contiguous() -> None:
print('Testing k-grouped contiguous GEMM:')
k_grouped_fp8_gemm_contiguous = deep_gemm.k_grouped_fp8_gemm_nt_contiguous if get_arch_major() == 9 \
else deep_gemm.k_grouped_fp8_gemm_tn_contiguous
for num_groups, m, n, major_a, major_b, ks, expected_k_per_group in enumerate_k_grouped_contiguous(torch.float8_e4m3fn):
use_ue8m0 = get_ue8m0_usage(KernelType.Kernel1D1D)
for test_empty_groups in (False, True):
new_ks = copy.deepcopy(ks)
if test_empty_groups and len(ks) > 1:
new_ks[random.randint(0, num_groups - 1)] = 0
k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, new_ks, use_ue8m0=use_ue8m0)
new_ks_tensor = torch.tensor(new_ks, dtype=torch.int, device='cuda')
k_grouped_fp8_gemm_contiguous(a, b, d, new_ks, new_ks_tensor, c)
diff = calc_diff(d, ref_d)
assert diff < 0.001, f'{m=}, {n=}, {k=}, {ks=}, {diff:.5f}'
# Test performance
k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_ue8m0=use_ue8m0)
ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda')
# noinspection PyShadowingNames
def test_func():
k_grouped_fp8_gemm_contiguous(a, b, d, ks, ks_tensor, c)
t = bench_kineto(test_func, 'fp8_gemm', suppress_kineto_output=True)
print(f' > Perf ({num_groups=:2}, m={m:5}, n={n:5}, k={k:5}): '
f'{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')
print()
if __name__ == '__main__':
torch.manual_seed(0)
random.seed(0)
print('Library path:')
print(f' > {deep_gemm.__path__}\n')
test_gemm()
test_m_grouped_gemm_contiguous()
test_m_grouped_gemm_masked()
test_k_grouped_gemm_contiguous()

57
tests/test_hyperconnection.py Normal file
View File

@@ -0,0 +1,57 @@
import torch
import random
import deep_gemm
from deep_gemm.testing import (
test_filter,
bench_kineto,
calc_diff, count_bytes
)
from deep_gemm.utils import align
from generators import get_arch_major
@test_filter(lambda: get_arch_major() >= 9)
def test_hc_prenorm_gemm() -> None:
# Needs TF32 precision for PyTorch GEMMs
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
print('Testing hyperconnection prenorm GEMM:')
for m in (13, 137, 4096, 8192):
for n, k in [(24, 28672), (24, 7680), (24, 7168)]:
for num_splits in [None, 16]:
a = torch.randn((m, k), dtype=torch.bfloat16, device='cuda')
b = torch.randn((n, k), dtype=torch.float, device='cuda')
d = torch.empty((m, n), dtype=torch.float, device='cuda') if num_splits is None else \
torch.empty((num_splits, m, n), dtype=torch.float, device='cuda')
s = torch.empty((m, ), dtype=torch.float, device='cuda') if num_splits is None else \
torch.empty((num_splits, m), dtype=torch.float, device='cuda')
deep_gemm.tf32_hc_prenorm_gemm(a, b, d, s, num_splits=num_splits)
final_d = d if num_splits is None else d.sum(0)
final_s = s if num_splits is None else s.sum(0)
ref_d = a.float() @ b.T
ref_s = a.float().square().sum(-1)
diff = max(calc_diff(final_d, ref_d), calc_diff(final_s, ref_s))
assert diff < 1e-8, f'{m=}, {n=}, {k=}, {diff:.10f}'
t = bench_kineto(lambda: deep_gemm.tf32_hc_prenorm_gemm(a, b, d, s, num_splits=num_splits), 'tf32_hc_prenorm_gemm', suppress_kineto_output=True)
print(f' > Perf (m={m:5}, n={n:5}, k={k:5}, num_splits={(num_splits or 0):2}): '
f'{t * 1e6:4.0f} us | '
f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes(a, b, d, s) / 1e9 / t:4.0f} GB/s')
print()
if __name__ == '__main__':
torch.manual_seed(0)
random.seed(0)
print('Library path:')
print(f' > {deep_gemm.__path__}\n')
test_hc_prenorm_gemm()

2
tests/test_legacy.py
View File

@@ -13,7 +13,7 @@ from generators import (
def test_m_grouped_gemm_contiguous_tl() -> None:
print('Testing m-grouped contiguous Triton GEMM:')
for _, num_groups, expected_m_per_group, n, k, major_a, major_b in enumerate_m_grouped_contiguous(torch.bfloat16):
for _, _, num_groups, expected_m_per_group, n, k, major_a, major_b, _ in enumerate_m_grouped_contiguous(torch.bfloat16):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'