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

Multiple updates and refactorings (#231)

Browse Source
This commit is contained in:
Ray Wang
2025-11-21 17:49:47 +08:00
committed by GitHub
parent bb4424aad4
commit 38f8ef73a4
80 changed files with 3767 additions and 2103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
tests/generators.py
View File

@@ -3,6 +3,7 @@ import random
import torch
from typing import Generator, List
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,
@@ -36,11 +37,6 @@ class MajorTypeAB(enum.Enum):
return self.value == 1
def get_arch_major() -> int:
major, minor = torch.cuda.get_device_capability()
return major
def get_ue8m0_usage(kernel_type: KernelType) -> bool:
if get_arch_major() == 9:
return False
@@ -51,9 +47,6 @@ def get_kernel_types(dtype: torch.dtype) -> tuple:
if dtype == torch.bfloat16:
return (KernelType.KernelNoSF, )
# TODO: SM100 1D2D kernels are going to be deprecated
# But if you want to test it, please use:
# `(KernelType.Kernel1D2D, ) if get_arch_major() == 9 else (KernelType.Kernel1D1D, KernelType.Kernel1D2D)`
return (KernelType.Kernel1D2D, ) if get_arch_major() == 9 else (KernelType.Kernel1D1D, )
@@ -72,8 +65,8 @@ def enumerate_normal(dtype: torch.dtype) -> Generator:
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 = [128, 4096], [4096, ]
nk_list = bf16_output_nk
m_fwd_list, m_bwd_list = [1, 128, 4096], [4096, ]
nk_list = list(bf16_output_nk)
# Only BF16 GEMM needs FP32 outputs
if dtype == torch.bfloat16:
@@ -82,14 +75,11 @@ def enumerate_normal(dtype: torch.dtype) -> Generator:
for kernel_type in get_kernel_types(dtype):
# Forward
for m in m_fwd_list:
for n, k in nk_list:
out_dtype = torch.float if (n, k) in fp32_output_nk else torch.bfloat16
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
# TODO: support BF16 SM90 MN-major kernels
if dtype == torch.bfloat16 and get_arch_major() == 9:
continue
# Backward
for m in m_bwd_list:
for n, k in nk_list:
@@ -106,7 +96,7 @@ def enumerate_normal(dtype: torch.dtype) -> Generator:
def enumerate_m_grouped_contiguous(dtype: torch.dtype) -> Generator:
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):
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
@@ -118,9 +108,9 @@ def enumerate_m_grouped_masked(dtype: torch.dtype) -> Generator:
yield kernel_type, num_groups, max_m, m, n, k
def enumerate_k_grouped_contiguous():
# Only K-major is supported for SM90
major_a, major_b = (MajorTypeAB.KMajor, MajorTypeAB.KMajor) if get_arch_major() == 9 \
def enumerate_k_grouped_contiguous(dtype: torch.dtype):
# Only K-major is supported for SM90 FP8
major_a, major_b = (MajorTypeAB.KMajor, MajorTypeAB.KMajor) if get_arch_major() == 9 and dtype == torch.float8_e4m3fn \
else (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor)
# Must with FP32 accumulation and 1D1D kernels
for num_groups, m, n, expected_k_per_group in (( 4, 4096, 7168, 8192), ( 4, 7168, 2048, 8192), # EP64
@@ -241,7 +231,8 @@ def generate_m_grouped_masked(num_groups: int, max_m: int, expected_m_per_group:
return a_fp8, b_fp8, masked_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], use_ue8m0: bool):
def generate_k_grouped_contiguous(num_groups: int, m: int, n: int, major_a: MajorTypeAB, major_b: MajorTypeAB, ks: List[int],
use_ue8m0: bool = False, use_bf16: bool = False):
assert get_mk_alignment_for_contiguous_layout() % 128 == 0
k = sum(ks)
@@ -257,6 +248,10 @@ def generate_k_grouped_contiguous(num_groups: int, m: int, n: int, major_a: Majo
ref_d[i] = c[i] + (a[start:end].T @ b[start:end])
start = end
if use_bf16:
assert (major_a, major_b) == (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor)
return k, a, b, c, d, ref_d
a_fp8 = per_channel_cast_to_fp8(a, use_ue8m0=use_ue8m0)
b_fp8 = per_channel_cast_to_fp8(b, use_ue8m0=use_ue8m0)

224
tests/test_attention.py
View File

@@ -3,10 +3,14 @@ import torch
from typing import Tuple
import deep_gemm
from deep_gemm.testing import bench_kineto, calc_diff, count_bytes
from deep_gemm.testing import (
bench_kineto,
calc_diff, count_bytes,
ignore_env, get_arch_major
)
from deep_gemm.utils import ceil_div, per_custom_dims_cast_to_fp8
from generators import get_arch_major, generate_normal, get_ue8m0_usage, get_kernel_types, MajorTypeAB
from generators import generate_normal, get_ue8m0_usage, get_kernel_types, MajorTypeAB
def apply_skip_head_mid(d: torch.Tensor, head_splits: Tuple[int, int, int]):
@@ -107,137 +111,169 @@ def ref_fp8_mqa_logits(q: torch.Tensor, kv: torch.Tensor, weights: torch.Tensor,
return logits, cost
@ignore_env('DG_JIT_PTXAS_CHECK', lambda: get_arch_major() == 10)
def test_mqa_logits():
print('Testing FP8 MQA Logits:')
num_heads, head_dim = 64, 128
for seq_len in (2048, 4096):
for seq_len_kv in (4096, 8192, 16384, 32768, 65536, 131072):
for disable_cp in (False, True):
q = torch.randn(seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
kv = torch.randn(seq_len_kv, head_dim, device='cuda', dtype=torch.bfloat16)
weights = torch.randn(seq_len, num_heads, device='cuda', dtype=torch.float32)
for compressed_logits in (False, True):
for seq_len_kv in (4096, 8192):
for disable_cp in (False, True):
q = torch.randn(seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
kv = torch.randn(seq_len_kv, head_dim, device='cuda', dtype=torch.bfloat16)
weights = torch.randn(seq_len, num_heads, device='cuda', dtype=torch.float32)
if disable_cp:
ks = torch.zeros(seq_len, dtype=torch.int, device='cuda')
ke = torch.arange(seq_len, dtype=torch.int, device='cuda') + (seq_len_kv - seq_len)
else:
ks, ke = generate_cp_test_data(seq_len, seq_len_kv)
if disable_cp:
ks = torch.zeros(seq_len, dtype=torch.int, device='cuda')
ke = torch.arange(seq_len, dtype=torch.int, device='cuda') + (seq_len_kv - seq_len)
else:
ks, ke = generate_cp_test_data(seq_len, seq_len_kv)
q_fp8 = q.to(torch.float8_e4m3fn)
kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0, ), False)
logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke)
q_fp8 = q.to(torch.float8_e4m3fn)
kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0, ), False)
do_check = (seq_len_kv < 32768)
if do_check:
ref_logits, ref_cost = ref_fp8_mqa_logits(q=q, kv=kv, weights=weights, cu_seqlen_ks=ks, cu_seqlen_ke=ke)
if compressed_logits:
max_seqlen_k = (ke - ks).max().item()
logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke, max_seqlen_k=max_seqlen_k, clean_logits=False)
assert logits.size() == (seq_len, max_seqlen_k)
tmp = torch.full((seq_len, seq_len_kv), float('-inf'), device='cuda')
for i in range(seq_len):
tmp[i, ks[i] : ke[i]] = logits[i, : ke[i] - ks[i]]
logits = tmp
else:
logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke)
ref_neginf_mask = (ref_logits == float('-inf'))
neginf_mask = (logits == float('-inf'))
assert torch.equal(neginf_mask, ref_neginf_mask)
do_check = (seq_len_kv < 32768)
if do_check:
ref_logits, ref_cost = ref_fp8_mqa_logits(q=q, kv=kv, weights=weights, cu_seqlen_ks=ks, cu_seqlen_ke=ke)
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=}"
else:
ref_cost = ref_fp8_mqa_logits(q=q, kv=kv, weights=weights, cu_seqlen_ks=ks, cu_seqlen_ke=ke, cost_only=True)
ref_neginf_mask = (ref_logits == float('-inf'))
neginf_mask = (logits == float('-inf'))
assert torch.equal(neginf_mask, ref_neginf_mask)
tflops = 2 * ref_cost * num_heads * head_dim / 1e12
t, clean_t = bench_kineto(lambda: deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke),
('fp8_mqa_logits', 'clean_logits'))
clean_bytes = (seq_len * seq_len_kv - ref_cost) * 4 + count_bytes(ks, ke)
print(f' > S={seq_len:4}, SKV={seq_len_kv:6}, H={num_heads:3}, D={head_dim:3}, CP={0 if disable_cp else 1}: '
f'{tflops / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, '
f'{(count_bytes(q_fp8, kv_fp8, weights, ks, ke) + ref_cost * 4) / t / 1e9:4.0f} GB/s | '
f'clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s')
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=}"
else:
ref_cost = ref_fp8_mqa_logits(q=q, kv=kv, weights=weights, cu_seqlen_ks=ks, cu_seqlen_ke=ke, cost_only=True)
tflops = 2 * ref_cost * num_heads * head_dim / 1e12
if compressed_logits:
t = bench_kineto(lambda: deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke, max_seqlen_k=max_seqlen_k, clean_logits=False), 'fp8_mqa_logits')
else:
t, clean_t = bench_kineto(lambda: deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke), ('fp8_mqa_logits', 'clean_logits'))
clean_bytes = (seq_len * seq_len_kv - ref_cost) * 4 + count_bytes(ks, ke)
print(f' > S={seq_len:4}, SKV={seq_len_kv:6}, H={num_heads:3}, D={head_dim:3}, CP={0 if disable_cp else 1}: '
f'{tflops / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, '
f'{(count_bytes(q_fp8, kv_fp8, weights, ks, ke) + ref_cost * 4) / t / 1e9:4.0f} GB/s', end='')
# noinspection PyUnboundLocalVariable
print(f' | clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s' if not compressed_logits else '')
print()
def ref_fp8_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor,
weights: torch.Tensor, context_lens: torch.Tensor, block_tables: torch.Tensor,
max_model_len: int):
max_model_len: int, is_context_lens_2d: bool):
batch_size, next_n, heads, dim = q.size()
num_block, block_size, _, dim = kv_cache.size()
logits = torch.full([batch_size * next_n, max_model_len], float('-inf'), device=q.device, dtype=torch.float32)
context_lens = context_lens.tolist()
for i in range(batch_size):
context_len = context_lens[i]
q_offsets = torch.arange(context_len - next_n, context_len, device='cuda')
q_offsets = torch.full((next_n, ), context_len, device='cuda', dtype=torch.int32) if is_context_lens_2d \
else torch.arange(context_len - next_n, context_len, device='cuda')
weight_slice = weights[i * next_n:(i + 1) * next_n, :].transpose(0, 1).contiguous()
for block_rk in range(ceil_div(context_len, block_size)):
block_idx = block_tables[i][block_rk]
qx, kx = q[i], kv_cache[block_idx]
k_offsets = torch.arange(block_rk * block_size, (block_rk + 1) * block_size, device='cuda')
mask = (k_offsets[None, :] < context_len) & (k_offsets[None, :] <= q_offsets[:, None])
s = torch.where(mask[None, :, :], (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(logits.dtype), float('-inf'))
s = torch.relu(s) * weight_slice[..., None]
s = s.sum(dim=0)
logits[i * next_n:(i + 1) * next_n, block_rk * block_size: (block_rk + 1) * block_size] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float('-inf'))
num_blocks = (context_len + block_size - 1) // block_size
block_idxs = block_tables[i][:num_blocks]
kv_slice = kv_cache[block_idxs] # [num_blocks, block_size, kv_heads, dim]
kx = kv_slice.permute(2, 3, 0, 1).reshape(kv_slice.size(2), dim, -1) # [kv_heads, dim, total_tokens]
qx = q[i].transpose(0, 1) # q[i]: [next_n, heads, dim] -> [heads, next_n, dim]
s = torch.matmul(qx, kx).to(logits.dtype) # [heads, next_n, dim] @ [1, dim, total_tokens] -> [heads, next_n, total_tokens]
total_len = num_blocks * block_size
k_offsets = torch.arange(0, total_len, device=q.device)
mask = (k_offsets[None, :] < context_len) & (k_offsets[None, :] <= q_offsets[:, None])
s = torch.where(mask[None, :, :], s, float('-inf')) # mask shape: [1, next_n, total_tokens]
s = torch.relu(s) * weight_slice[..., None] # weight_slice: [heads, next_n] -> [heads, next_n, 1]
s = s.sum(dim=0) # [next_n, total_tokens]
logits[i * next_n:(i + 1) * next_n, :total_len] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float('-inf'))
return logits
def test_paged_mqa_logits():
# TODO: fully refactor with PyTest
print('Testing FP8 Paged MQA Logits:')
max_model_len = 111 * 1000
for batch_size, next_n in [(64, 1), (64, 2), (128, 1)]:
for heads, index_dim in [(64, 128)]:
for avg_kv in (8192, 32768):
num_blocks, blocksize = max_model_len * 3, 64
for is_context_lens_2d in (False, True):
for batch_size, next_n in [(64, 1), (64, 2), (128, 1)]:
for heads, index_dim in [(64, 128)]:
for avg_kv in (8192, 32768):
num_blocks, blocksize = max_model_len * 3, 64
q = torch.randn((batch_size, next_n, heads, index_dim), device='cuda', dtype=torch.bfloat16)
kv_cache = torch.randn((num_blocks, blocksize, 1, index_dim), device='cuda', dtype=torch.bfloat16)
weights = torch.randn((batch_size * next_n, heads), device='cuda', dtype=torch.float32)
q = torch.randn((batch_size, next_n, heads, index_dim), device='cuda', dtype=torch.bfloat16)
kv_cache = torch.randn((num_blocks, blocksize, 1, index_dim), device='cuda', dtype=torch.bfloat16)
weights = torch.randn((batch_size * next_n, heads), device='cuda', dtype=torch.float32)
q_fp8 = q.to(torch.float8_e4m3fn)
kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
context_lens = torch.randint(int(0.7 * avg_kv), int(1.3 * avg_kv), (batch_size, )).cuda().to(torch.int32)
max_block_len = (context_lens.max().item() + blocksize - 1) // blocksize * blocksize
block_tables = torch.zeros((batch_size, max_block_len), device='cuda', dtype=torch.int32)
context_lens = torch.randint(int(0.7 * avg_kv), int(1.3 * avg_kv), (batch_size, )).cuda().to(torch.int32)
context_lens_list = context_lens.tolist()
max_block_len = (max(context_lens_list) + blocksize - 1) // blocksize * blocksize
block_tables = torch.zeros((batch_size, max_block_len), device='cuda', dtype=torch.int32)
counter = 0
block_idx_pool = list(range(num_blocks))
random.shuffle(block_idx_pool)
for i in range(batch_size):
ctx_len = context_lens[i].item()
for j in range(ceil_div(ctx_len, blocksize)):
block_tables[i][j] = block_idx_pool[counter]
counter += 1
counter, block_idx_pool = 0, torch.randperm(num_blocks, device='cuda', dtype=torch.int32)
for i in range(batch_size):
num_blocks = ceil_div(context_lens_list[i], blocksize)
block_tables[i][:num_blocks] = block_idx_pool[counter: counter+num_blocks]
counter += num_blocks
q_fp8 = q.to(torch.float8_e4m3fn)
kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
ref_logits = ref_fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_tables, max_model_len, is_context_lens_2d)
positions = torch.arange(max_model_len, device='cuda').unsqueeze(0).expand(batch_size * next_n, -1)
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(context_lens, blocksize, deep_gemm.get_num_sms())
logits = deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens, block_tables, schedule_metadata, max_model_len, clean_logits=True)
if is_context_lens_2d:
context_lens_2d = ((context_lens.unsqueeze(1) + 1) * torch.rand(batch_size, next_n, device='cuda')).int()
context_lens_2d[:, next_n-1] = context_lens
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(context_lens_2d, blocksize, deep_gemm.get_num_sms())
logits = deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens_2d, block_tables, schedule_metadata, max_model_len, clean_logits=False)
ref_neginf_mask = ~(positions < context_lens_2d.view(-1).unsqueeze(1))
else:
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(context_lens, blocksize, deep_gemm.get_num_sms())
logits = deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens, block_tables, schedule_metadata, max_model_len, clean_logits=True)
row_indices = torch.arange(batch_size * next_n, device='cuda') // next_n
next_n_offset = torch.arange(batch_size * next_n, device='cuda') % next_n
ref_neginf_mask = ~(positions <= (context_lens[row_indices] - next_n + next_n_offset).unsqueeze(1))
neginf_mask = (logits == float('-inf'))
assert torch.equal(neginf_mask, ref_neginf_mask)
ref_logits = ref_fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_tables, max_model_len)
positions = torch.arange(max_model_len, device='cuda').unsqueeze(0).expand(batch_size * next_n, -1)
row_indices = torch.arange(batch_size * next_n, device='cuda') // next_n
next_n_offset = torch.arange(batch_size * next_n, device='cuda') % next_n
ref_neginf_mask = ~(positions <= (context_lens[row_indices] - next_n + next_n_offset).unsqueeze(1))
logits = logits.masked_fill(ref_neginf_mask, 0)
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
diff = calc_diff(logits, ref_logits)
assert diff < 1e-3, f"{diff=}"
neginf_mask = (logits == float('-inf'))
assert torch.equal(neginf_mask, ref_neginf_mask)
logits = logits.masked_fill(neginf_mask, 0)
ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
diff = calc_diff(logits, ref_logits)
assert diff < 1e-3, f"{diff=}"
sum_lens = sum(context_lens.to(torch.int64))
tflops = 2 * sum_lens * next_n * heads * index_dim / 1e12
input_bytes = count_bytes(q_fp8, weights, context_lens) + sum_lens * (index_dim + 4) + (sum_lens / blocksize) * 4
output_bytes = sum_lens * next_n * 4
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)
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 | '
f'clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s')
sum_lens = sum(context_lens.to(torch.int64))
tflops = 2 * sum_lens * next_n * heads * index_dim / 1e12
input_bytes = count_bytes(q_fp8, weights, context_lens) + sum_lens * (index_dim + 4) + (sum_lens / blocksize) * 4
output_bytes = sum_lens * next_n * 4
if is_context_lens_2d:
t = bench_kineto(lambda: deep_gemm.fp8_paged_mqa_logits(q_fp8, kv_cache_fp8, weights, context_lens_2d, block_tables, schedule_metadata, max_model_len, clean_logits=False),
'fp8_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)
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='')
# noinspection PyUnboundLocalVariable
print(f' | clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s' if not is_context_lens_2d else '')
print()
if __name__ == '__main__':
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.manual_seed(0)
random.seed(0)

62
tests/test_bf16.py
View File

@@ -1,5 +1,7 @@
import torch
import copy
import numpy as np
import random
import torch
import deep_gemm
from deep_gemm.testing import (
@@ -8,13 +10,14 @@ from deep_gemm.testing import (
)
from generators import (
get_arch_major,
enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, generate_normal,
generate_m_grouped_contiguous, generate_m_grouped_masked
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
)
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:
@@ -34,18 +37,20 @@ def test_gemm() -> None:
assert a.is_contiguous() and b.is_contiguous()
getattr(deep_gemm, func_name)(a, b, d, c=c)
diff = calc_diff(d, ref_d)
assert diff < 0.0001, (f'{m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=}, '
assert diff < 1e-5, (f'{m=}, {n=}, {k=}, {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_bf16=True)
t = bench_kineto(lambda: deep_gemm.bf16_gemm_nt(a, b, d, c=c), 'bf16_gemm', suppress_kineto_output=True)
cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), ('nvjet', 'reduce'), suppress_kineto_output=True)
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'{t * 1e6:7.1f} us | '
f'{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')
print()
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:
@@ -65,7 +70,7 @@ def test_m_grouped_gemm_contiguous() -> None:
getattr(deep_gemm, func_name)(a, b, d, m_indices)
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}'
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)
# noinspection PyShadowingNames
@@ -91,7 +96,7 @@ def test_m_grouped_gemm_masked() -> None:
deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group)
for j in range(num_groups):
diff = calc_diff(d[j, :masked_m[j].item()], ref_d[j, :masked_m[j].item()])
assert diff < 0.001, f'{m=}, {n=}, {k=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}'
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)
@@ -110,6 +115,37 @@ def test_m_grouped_gemm_masked() -> None:
print()
def test_k_grouped_gemm_contiguous() -> None:
print('Testing k-grouped contiguous GEMM:')
for num_groups, m, n, major_a, major_b, ks, expected_k_per_group in enumerate_k_grouped_contiguous(torch.bfloat16):
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_bf16=True)
new_ks_tensor = torch.tensor(new_ks, dtype=torch.int, device='cuda')
deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, new_ks, new_ks_tensor, c)
diff = calc_diff(d, ref_d)
assert diff < 1e-5, f'{m=}, {n=}, {k=}, {ks=}, {diff:.7f}'
# Test performance
k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_bf16=True)
ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda')
# noinspection PyShadowingNames
def test_func():
deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, ks_tensor, c)
t = bench_kineto(test_func, 'bf16_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()
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):
@@ -121,7 +157,7 @@ def test_cublaslt_gemm() -> None:
a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True)
deep_gemm.cublaslt_gemm_nt(a, b, d, c=c)
diff = calc_diff(d, ref_d)
assert diff < 5e-7, f'{diff=}, ({m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=})'
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,)
print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): '
@@ -132,18 +168,16 @@ def test_cublaslt_gemm() -> None:
if __name__ == '__main__':
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.manual_seed(0)
random.seed(0)
print('Library path:')
print(f' > {deep_gemm.__path__}\n')
test_gemm()
# TODO: support SM100
if get_arch_major() == 9:
if get_arch_major() >= 9:
test_gemm()
test_m_grouped_gemm_contiguous()
test_m_grouped_gemm_masked()
test_k_grouped_gemm_contiguous()
test_cublaslt_gemm()

125
tests/test_einsum.py
View File

@@ -3,8 +3,13 @@ import torch
import deep_gemm
from deep_gemm.testing import (
bench, bench_kineto,
calc_diff, count_bytes
bench_kineto,
calc_diff, count_bytes,
get_arch_major, test_filter
)
from deep_gemm.utils.math import (
ceil_div,
per_block_cast_to_fp8, per_channel_cast_to_fp8, per_token_cast_to_fp8
)
@@ -33,8 +38,8 @@ def test_bmk_bnk_mn() -> None:
def test_bhr_hdr_bhd():
print('Testing "bhr, hdr -> bhd":')
for b in (128, 4096, 8192):
for h, r, d in [(128, 512, 128)]:
for h, r, d in [(128, 512, 128), (8, 4096, 1024)]:
for b in (4, 32, 128, 4096, 8192):
x = torch.randn((b, h, r), device='cuda', dtype=torch.bfloat16)
fy = torch.randn((h, d, r + 128), device='cuda', dtype=torch.bfloat16)
y = fy[:, :, :r]
@@ -43,18 +48,20 @@ def test_bhr_hdr_bhd():
deep_gemm.einsum('bhr,hdr->bhd', x, y, z)
assert calc_diff(z, ref_z) < 1e-10
t = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z), 'nvjet', suppress_kineto_output=True)
t = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z), 'gemm', suppress_kineto_output=True)
t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True)
print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ',
f'{t * 1e6:4.0f} us | '
f'{2 * b * h * r * d / t / 1e12:.0f} TFLOPS | '
f'{count_bytes((x, y, z)) / t / 1e9:.0f} GB/s')
f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes((x, y, z)) / t / 1e9:4.0f} GB/s | '
f'{t_cublaslt / t:4.2f} x')
print()
def test_bhd_hdr_bhr():
print('Testing "bhd, hdr -> bhr":')
for b in (128, 4096, 8192):
for h, r, d in [(128, 512, 128)]:
for h, r, d in [(128, 512, 128), (8, 4096, 1024)]:
for b in (4, 32, 128, 4096, 8192):
x = torch.randn((b, h, d), device='cuda', dtype=torch.bfloat16)
fy = torch.randn((h, d, r + 128), device='cuda', dtype=torch.bfloat16)
y = fy[:, :, :r]
@@ -63,17 +70,103 @@ def test_bhd_hdr_bhr():
deep_gemm.einsum('bhd,hdr->bhr', x, y, z)
assert calc_diff(z, ref_z) < 1e-10
t = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z), 'nvjet', suppress_kineto_output=True)
t = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z), 'gemm', suppress_kineto_output=True)
t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True)
print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ',
f'{t * 1e6:4.0f} us | '
f'{2 * b * h * r * d / t / 1e12:.0f} TFLOPS | '
f'{count_bytes((x, y, z)) / t / 1e9:.0f} GB/s')
f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes((x, y, z)) / t / 1e9:4.0f} GB/s | '
f'{t_cublaslt / t:4.2f} x')
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)]:
for b in (4, 32, 128, 4096, 8192):
x = torch.randn((b, h, r), device='cuda', dtype=torch.bfloat16)
y = torch.randn((h, d, r), device='cuda', dtype=torch.bfloat16)
ref_z = torch.einsum('bhr,hdr->bhd', x, y)
x_fp8 = per_token_cast_to_fp8(x.view(-1, r), use_ue8m0=use_ue8m0)
x_fp8 = x_fp8[0].view(b, h, r), x_fp8[1].view(b, h, ceil_div(r, 128))
y_fp8 = (torch.empty_like(y, dtype=torch.float8_e4m3fn),
torch.empty((h, ceil_div(d, 128), ceil_div(r, 128)), device='cuda', dtype=torch.float))
for i in range(h):
y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i], use_ue8m0=use_ue8m0)
z = torch.empty((b, h, d), device='cuda', dtype=torch.bfloat16)
deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, z)
assert calc_diff(z, ref_z) < 1e-3
t = bench_kineto(lambda: deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, z), 'fp8_gemm', suppress_kineto_output=True)
t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True)
print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ',
f'{t * 1e6:4.0f} us | '
f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes((x_fp8, y_fp8, z)) / t / 1e9:4.0f} GB/s | '
f'{t_cublaslt / t:4.2f} x')
print()
@test_filter(lambda: get_arch_major() >= 10)
def test_fp8_bhd_hdr_bhr(use_ue8m0: bool = True):
print('Testing FP8 "bhd, hdr -> bhr":')
for h, r, d in [(8, 4096, 1024)]:
for b in (4, 32, 128, 4096, 8192):
x = torch.randn((b, h, d), device='cuda', dtype=torch.bfloat16)
y = torch.randn((h, d, r), device='cuda', dtype=torch.bfloat16)
ref_z = torch.einsum('bhd,hdr->bhr', x, y)
x_fp8 = per_token_cast_to_fp8(x.view(-1, d), use_ue8m0=use_ue8m0)
x_fp8 = x_fp8[0].view(b, h, d), x_fp8[1].view(b, h, ceil_div(d, 128))
y_fp8 = (torch.empty_like(y, dtype=torch.float8_e4m3fn),
torch.empty((h, ceil_div(d, 128), ceil_div(r, 128)), device='cuda', dtype=torch.float))
for i in range(h):
y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i], use_ue8m0=use_ue8m0)
z = torch.empty((b, h, r), device='cuda', dtype=torch.bfloat16)
deep_gemm.fp8_einsum('bhd,hdr->bhr', x_fp8, y_fp8, z)
assert calc_diff(z, ref_z) < 1e-3
t = bench_kineto(lambda: deep_gemm.fp8_einsum('bhd,hdr->bhr', x_fp8, y_fp8, z), 'fp8_gemm', suppress_kineto_output=True)
t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True)
print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ',
f'{t * 1e6:4.0f} us | '
f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes((x_fp8, y_fp8, z)) / t / 1e9:4.0f} GB/s | '
f'{t_cublaslt / t:4.2f} x')
print()
@test_filter(lambda: get_arch_major() >= 10)
def test_fp8_bhd_bhr_hdr(use_ue8m0: bool = True):
print('Testing FP8 "bhd, bhr -> hdr":')
for h, r, d in [(8, 4096, 1024)]:
for b in (4096, 8192):
x = torch.randn((b, h, d), device='cuda', dtype=torch.bfloat16)
y = torch.randn((b, h, r), device='cuda', dtype=torch.bfloat16)
z_0 = torch.randn((h, d, r), device='cuda', dtype=torch.float32) * 10
ref_z = z_0 + torch.einsum('bhd,bhr->hdr', x, y)
x_fp8 = per_channel_cast_to_fp8(x.view(b, -1), use_ue8m0=use_ue8m0)
y_fp8 = per_channel_cast_to_fp8(y.view(b, -1), use_ue8m0=use_ue8m0)
x_fp8 = (x_fp8[0].view(b, h, d), x_fp8[1].view(ceil_div(b, 128), h, d))
y_fp8 = (y_fp8[0].view(b, h, r), y_fp8[1].view(ceil_div(b, 128), h, r))
z = z_0.clone()
deep_gemm.fp8_einsum('bhd,bhr->hdr', x_fp8, y_fp8, z, z, recipe=(1, 1, 128))
assert calc_diff(z, ref_z) < 1e-3
t = bench_kineto(lambda: deep_gemm.fp8_einsum('bhd,bhr->hdr', x_fp8, y_fp8, z, z, recipe=(1, 1, 128)), 'fp8_gemm', suppress_kineto_output=True)
print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ',
f'{t * 1e6:4.0f} us | '
f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | '
f'{count_bytes((x_fp8, y_fp8, z, z)) / t / 1e9:4.0f} GB/s')
print()
if __name__ == '__main__':
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.manual_seed(0)
random.seed(0)
@@ -83,3 +176,7 @@ if __name__ == '__main__':
test_bmk_bnk_mn()
test_bhr_hdr_bhd()
test_bhd_hdr_bhr()
test_fp8_bhr_hdr_bhd()
test_fp8_bhd_hdr_bhr()
test_fp8_bhd_bhr_hdr()

29
tests/test_fp8.py
View File

@@ -1,23 +1,26 @@
import copy
import numpy as np
import random
import time
import torch
import deep_gemm
from deep_gemm.testing import (
bench, bench_kineto,
calc_diff, count_bytes
bench_kineto,
calc_diff, count_bytes,
ignore_env, get_arch_major
)
from generators import (
KernelType, get_arch_major, get_ue8m0_usage,
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'
@@ -45,10 +48,12 @@ def test_gemm() -> None:
'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:4.0f} us | {2 * m * n * k / t / 1e12:4.0f} TFLOPS | '
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')
print()
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:
@@ -100,6 +105,8 @@ def test_m_grouped_gemm_masked() -> None:
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}'
@@ -125,7 +132,7 @@ def test_k_grouped_gemm_contiguous() -> None:
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():
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):
@@ -136,10 +143,8 @@ def test_k_grouped_gemm_contiguous() -> None:
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)
do_check = True
if do_check:
diff = calc_diff(d, ref_d)
assert diff < 0.001, f'{m=}, {n=}, {k=}, {ks=}, {diff:.5f}'
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)
@@ -158,8 +163,6 @@ def test_k_grouped_gemm_contiguous() -> None:
if __name__ == '__main__':
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.manual_seed(0)
random.seed(0)

6
tests/test_layout.py
View File

@@ -1,7 +1,6 @@
import time
import torch
import random
from deep_gemm.testing import bench_kineto, count_bytes, calc_diff
from deep_gemm.testing import bench_kineto, count_bytes
from deep_gemm.utils import (
align, ceil_div,
per_token_cast_to_fp8, per_channel_cast_to_fp8,
@@ -12,7 +11,6 @@ from deep_gemm.utils import (
)
from generators import (
enumerate_transpose,
enumerate_sf_layout,
enumerate_k_grouped_sf_layout
)
@@ -107,8 +105,6 @@ def test_k_grouped_sf_layout_kernels() -> None:
if __name__ == '__main__':
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.manual_seed(1)
random.seed(1)

90
tests/test_legacy.py Normal file
View File

@@ -0,0 +1,90 @@
import torch
import random
import deep_gemm
from deep_gemm.testing import (
bench_kineto,
calc_diff, count_bytes
)
from generators import (
enumerate_m_grouped_contiguous, enumerate_k_grouped_contiguous,
generate_m_grouped_contiguous, generate_k_grouped_contiguous,
)
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):
major_opt = 'N' if major_a.is_k_major() else 'T'
major_opt += 'T' if major_b.is_k_major() else 'N'
for expand in (False, True):
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)
func_name = f"{'a_fused_' if expand else ''}m_grouped_bf16_gemm_{major_opt.lower() if test_alias else 'nt'}_contiguous_tl"
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()
if expand:
m_row_indices = torch.arange(0, m, dtype=torch.int32, device='cuda')
getattr(deep_gemm.legacy, func_name)(a, b, d, (m_indices, m_row_indices))
else:
getattr(deep_gemm.legacy, func_name)(a, b, d, m_indices)
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}, {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)
# noinspection PyShadowingNames
def test_func():
deep_gemm.legacy.m_grouped_bf16_gemm_nt_contiguous_tl(a, b, d, m_indices)
t = bench_kineto(test_func, 'm_grouped_bf16_gemm_contiguous_tl_impl', suppress_kineto_output=True)
print(f' > Perf ({num_groups=}, m={m:5}, n={n:5}, k={k:5}, 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_k_grouped_gemm_contiguous_tl() -> None:
print('Testing k-grouped contiguous Triton GEMM:')
for num_groups, m, n, major_a, major_b, ks, expected_k_per_group in enumerate_k_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 fused_operand in ('a', 'b'):
k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_ue8m0=False, use_bf16=True)
func_name = f"{fused_operand}_fused_k_grouped_bf16_gemm_{major_opt.lower()}_contiguous_tl"
k_indices = torch.arange(0, k, dtype=torch.int32, device='cuda')
k_start = torch.empty(len(ks), dtype=torch.int32, device='cuda')
k_end = torch.empty(len(ks), dtype=torch.int32, device='cuda')
for i, group_k in enumerate(ks):
k_start[i] = k_end[i-1] if i > 0 else 0
k_end[i] = k_start[i] + group_k
getattr(deep_gemm.legacy, func_name)(a, b, c, (k_indices, k_start, k_end), True)
diff = calc_diff(c, ref_d)
assert diff < 0.001, f'{m=}, {n=}, {k=}, {major_opt}, {diff:.5f}'
k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_ue8m0=False, use_bf16=True)
# noinspection PyShadowingNames
def test_func():
deep_gemm.legacy.b_fused_k_grouped_bf16_gemm_tn_contiguous_tl(a, b, c, (k_indices, k_start, k_end), True)
t = bench_kineto(test_func, 'b_fused_k_grouped_bf16_gemm_contiguous_tl_impl', suppress_kineto_output=True)
print(f' > Perf ({num_groups=}, m={m:5}, n={n:5}, k={k:5}, 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()
if __name__ == '__main__':
torch.manual_seed(0)
random.seed(0)
print('Library path:')
print(f' > {deep_gemm.__path__}\n')
test_m_grouped_gemm_contiguous_tl()
test_k_grouped_gemm_contiguous_tl()

78
tests/test_sanitizer.py Normal file
View File

@@ -0,0 +1,78 @@
import argparse
import importlib
import inspect
import os
import subprocess
import sys
import deep_gemm
# Single test template
script_dir = os.path.dirname(os.path.abspath(__file__))
test_template = """
import random
import sys
import torch
# Necessary for `generators.py`
sys.path.append('{script_dir}')
torch.manual_seed(0)
random.seed(0)
from tests.{module_name} import {func_name}
{func_name}()
"""
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--funcs', type=str, default='all')
parser.add_argument('--tools', type=str, default='memcheck,synccheck')
args = parser.parse_args()
if args.funcs != 'all':
funcs = []
for name in [x.strip() for x in args.funcs.split(',')]:
module_name, func_name = name.split('.')
funcs.append((module_name, func_name))
else:
# Get all test functions except those related to cuBLAS
files = [f for f in os.listdir(script_dir) if f.endswith('.py')]
exclude_files = ['test_sanitizer.py', 'generators.py']
funcs = [
(module_name, name)
for module_name in [os.path.splitext(f)[0] for f in files if f not in exclude_files]
for name, obj in inspect.getmembers(importlib.import_module(module_name))
if inspect.isfunction(obj) and name.startswith('test') and 'test_filter' not in name
]
tools = [x.strip() for x in args.tools.split(',')]
env = os.environ.copy()
env['CUDA_LAUNCH_BLOCKING'] = '1'
env['DG_JIT_PTXAS_CHECK'] = '1'
env['DG_USE_NVIDIA_TOOLS'] = '1'
env['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1'
env['TORCH_SHOW_CPP_STACKTRACES'] = '1'
print(f'Library path: {deep_gemm.__path__}')
for module_name, func_name in funcs:
for tool in tools:
cmd = [
'/usr/local/cuda/bin/compute-sanitizer',
f'--tool={tool}',
'--target-processes=application-only',
'--destroy-on-device-error=context',
'--force-blocking-launches',
'--check-api-memory-access=no',
'--kernel-name-exclude', 'kns=nvjet',
'python',
'-c',
test_template.format(module_name=module_name, func_name=func_name, script_dir=script_dir)
]
print(f'\n{"=" * 60}')
print(f'Running {module_name}.{func_name} with compute-sanitizer {tool}')
result = subprocess.run(cmd, env=env)
if result.returncode != 0:
sys.exit(result.returncode)