A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 02:30 UTC
[CI][AMD][BugFix] Use torch.testing.assert_close instead of assert torch.allclose in test_rocm_skinny_gemms.py (#34181)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Feb 18, 2026 at 17:10 UTC
2b84ac669cfd8a4b6433b4ae4505028d9082c3a7
1 file changed
+7
-7
tests/kernels/quantization/test_rocm_skinny_gemms.py
+7
-7
@@ -155,9 +155,9 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
155
out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)
156
157
if xnorm:
158
- assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
158
+ torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
159
else:
160
- assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-2)
160
+ torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
161
162
163
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
@@ -177,7 +177,7 @@ def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
177
ref_out = torch.matmul(A, B.t())
178
out = ops.LLMM1(B, A, rows_per_block)
179
180
- assert torch.allclose(out, ref_out, rtol=0.01)
180
+ torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
181
182
183
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -194,7 +194,7 @@ def test_rocm_wvsplitk_kernel(n, k, m, dtype, seed):
194
ref_out = torch.nn.functional.linear(A, B)
195
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count)
196
197
- assert torch.allclose(out, ref_out, rtol=0.01)
197
+ torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
198
199
200
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -213,7 +213,7 @@ def test_rocm_wvsplitk_bias1D_kernel(n, k, m, dtype, seed):
213
ref_out = torch.nn.functional.linear(A, B, BIAS)
214
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
215
216
- assert torch.allclose(out, ref_out, rtol=0.01)
216
+ torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
217
218
219
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
@@ -232,7 +232,7 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
232
ref_out = torch.nn.functional.linear(A, B, BIAS)
233
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
234
235
- assert torch.allclose(out, ref_out, rtol=0.01)
235
+ torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
236
237
238
@pytest.mark.parametrize("xnorm", [False, True])
@@ -275,4 +275,4 @@ def test_rocm_wvsplitk_fp8_kernel(
275
# wider pytrch thresh for large-K & no xnorm
276
torch.testing.assert_close(out, ref_out, atol=0.07, rtol=5e-2)
277
else:
278
- torch.testing.assert_close(out, ref_out, atol=0.01, rtol=0.01)
278
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)