update unitest

This commit is contained in:
wangzhe_ant
2025-06-24 18:53:19 +08:00
parent 7db1b0ef63
commit e29e996a42

View File

@@ -374,14 +374,15 @@ def test_k_grouped_wgrad_gemm():
def test_m_grouped_gemm_offset() -> None:
print('Testing grouped offset GEMM:')
for num_groups, expected_m_per_group, k, n in ((8, 32, 7168, 4096),(9, 64, 7168, 4096),(32, 32, 7168, 4096)):
for num_groups, expected_m_per_group in ((2, 16), (4, 16), (2, 32), (9, 32), (2, 32), (4, 32), (32, 64)):
for k, n in ((7168, 4096),):
# NOTES: we should mask the unfilled part before calculating difference
ms, m_offset, x_fp8_offset, y_fp8_offset, offset, out_offset, ref_out_offset = construct_offset_grouped(num_groups, expected_m_per_group, k, n)
pad_x_fp8 = change_to_offset_layout(ms, x_fp8_offset[0], x_fp8_offset[1])
deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_offset(pad_x_fp8, y_fp8_offset, offset, out_offset, expected_m_per_group)
diff = calc_diff(out_offset, ref_out_offset)
assert diff < 0.001.1, f'{m_offset=}, {k=}, {n=}, {diff:.5f}'
assert diff < 0.001, f'{m_offset=}, {k=}, {n=}, {diff:.5f}'
# noinspection PyShadowingNames
def test_func():