bumping sgl-deep-gemm to 0.2.0 (#39371)

python/pyproject.toml

@@ -76,7 +76,7 @@ dependencies = [
   "sentencepiece",
   "setproctitle",
   "sgl-deep-ep==0.1.2",
-  "sgl-deep-gemm==0.1.7",
+  "sgl-deep-gemm==0.2.0",
   "sglang-kernel==0.4.7",
   "smg-grpc-servicer>=0.9.0",
   "soundfile==0.13.1",

python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py

@@ -255,6 +255,12 @@ def _matmul_persistent_deepgemm(
     dtype = a.dtype
     out = torch.empty((M, N), device=a.device, dtype=dtype)
 
+    # DeepGEMM 0.2 defaults BF16 GEMMs to cuBLASLt, whose reduction can
+    # depend on the batch size. Older wheels always use the invariant kernel.
+    get_deterministic = getattr(deep_gemm, "get_deterministic_algorithms", None)
+    deterministic = get_deterministic() if get_deterministic else True
+    if not deterministic:
+        deep_gemm.use_deterministic_algorithms(True)
     try:
         deep_gemm.bf16_gemm_nn(a, b, out)
     except RuntimeError as e:
@@ -264,6 +270,9 @@ def _matmul_persistent_deepgemm(
             f"Consider increasing MIN_DEEPGEMM_DIM in matmul_persistent() or disabling DeepGEMM "
             f"for small matrices. Original error: {e}"
         ) from e
+    finally:
+        if not deterministic:
+            deep_gemm.use_deterministic_algorithms(False)
 
     # TODO can this be put in DeepGEMM's `c`?
     if bias is not None:

test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py

@@ -101,6 +101,34 @@ class TestBatchInvariantOps(CustomTestCase):
             f"{test_name}: diff_range must be 0 in batch-invariant mode, got {diff_range} for {dtype}",
         )
 
+    def test_deepgemm_random_bf16_batch_invariance(self):
+        if not batch_invariant_ops.ENABLE_JIT_DEEPGEMM:
+            self.skipTest("DeepGEMM is unavailable on this device")
+
+        deep_gemm = batch_invariant_ops.deep_gemm
+        get_deterministic = getattr(deep_gemm, "get_deterministic_algorithms", None)
+        original_mode = get_deterministic() if get_deterministic else None
+        modes = (False, True) if get_deterministic else (None,)
+        generator = torch.Generator(device=device_type).manual_seed(42)
+        a = torch.randn(257, 4096, dtype=torch.bfloat16, generator=generator)
+        b = torch.randn(4096, 4096, dtype=torch.bfloat16, generator=generator).T
+        try:
+            for mode in modes:
+                if mode is not None:
+                    deep_gemm.use_deterministic_algorithms(mode)
+                with self.subTest(deterministic=mode):
+                    ref = batch_invariant_ops._matmul_persistent_deepgemm(a[:1], b)
+                    for batch_size in (16, 64, 257):
+                        out = batch_invariant_ops._matmul_persistent_deepgemm(
+                            a[:batch_size], b
+                        )
+                        torch.testing.assert_close(out[:1], ref, rtol=0, atol=0)
+                    if get_deterministic:
+                        self.assertEqual(get_deterministic(), mode)
+        finally:
+            if original_mode is not None:
+                deep_gemm.use_deterministic_algorithms(original_mode)
+
     def test_small_matrices(self):
         """Test batch invariance with small matrix sizes"""
         test_cases = [