From ea3b9ab6918cb3c6f9b44173c30c0640eb83b7c8 Mon Sep 17 00:00:00 2001 From: Rayfxl Date: Mon, 8 Dec 2025 21:21:09 +0800 Subject: [PATCH] [Tests] Fix assert bug for fp8e5 dot on sm120. --- python/test/unit/language/test_core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/test/unit/language/test_core.py b/python/test/unit/language/test_core.py index cada30bbc0f2..86bb268b7999 100644 --- a/python/test/unit/language/test_core.py +++ b/python/test/unit/language/test_core.py @@ -3480,7 +3480,8 @@ def kernel(X, stride_xm, stride_xk, Y, stride_yk, stride_yn, W, stride_wn, strid if capability[0] == 9 and M >= 64 and N >= 8: assert 'wgmma.mma_async.sync.aligned.m64n128k32.f32.e5m2.e5m2' in ptx elif capability[0] >= 8 and M < 64: - assert 'mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32' in ptx + assert ('mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32' + in ptx) or ('mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32' in ptx) elif in_dtype == "float8e4nv" and out_dtype == tl.float32: if capability[0] == 9 and M >= 64 and N >= 8: assert 'wgmma.mma_async.sync.aligned.m64n128k32.f32.e4m3.e4m3' in ptx