minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Test] Fix DeepGEMM batch invariance test output dtype (#39636)

Browse Source
This commit is contained in:
YAMY
2026-09-15 12:40:28 -07:00
committed by GitHub
parent 073fdf97cc
commit 24874f90a3
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py
+4 -2
View File
@@ -117,10 +117,12 @@ class TestBatchInvariantOps(CustomTestCase):
if mode is not None: if mode is not None:
deep_gemm.use_deterministic_algorithms(mode) deep_gemm.use_deterministic_algorithms(mode)
with self.subTest(deterministic=mode): with self.subTest(deterministic=mode):
ref = batch_invariant_ops._matmul_persistent_deepgemm(a[:1], b) ref = batch_invariant_ops._matmul_persistent_deepgemm(
a[:1], b, out_dtype=torch.bfloat16
)
for batch_size in (16, 64, 257): for batch_size in (16, 64, 257):
out = batch_invariant_ops._matmul_persistent_deepgemm( out = batch_invariant_ops._matmul_persistent_deepgemm(
a[:batch_size], b a[:batch_size], b, out_dtype=torch.bfloat16
) )
torch.testing.assert_close(out[:1], ref, rtol=0, atol=0) torch.testing.assert_close(out[:1], ref, rtol=0, atol=0)
if get_deterministic: if get_deterministic: