diff --git a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py index 4fb74dee5..c01f8bb33 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py @@ -241,7 +241,15 @@ def _matmul_persistent_deepgemm( dtype = a.dtype out = torch.empty((M, N), device=a.device, dtype=dtype) - deep_gemm.bf16_gemm_nn(a, b, out) + try: + deep_gemm.bf16_gemm_nn(a, b, out) + except RuntimeError as e: + raise RuntimeError( + f"DeepGEMM failed for matrix shapes M={M}, N={N}, K={K}. " + f"This typically occurs when dimensions are too small for DeepGEMM's TMA descriptors. " + f"Consider increasing MIN_DEEPGEMM_DIM in matmul_persistent() or disabling DeepGEMM " + f"for small matrices. Original error: {e}" + ) from e # TODO can this be put in DeepGEMM's `c`? if bias is not None: @@ -253,6 +261,11 @@ def _matmul_persistent_deepgemm( def matmul_persistent( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None ): + K, N = b.shape + + # DeepGEMM has minimum dimension requirements for TMA descriptors + MIN_DEEPGEMM_DIM = 16 + if ( _ENABLE_MM_DEEPGEMM and ENABLE_JIT_DEEPGEMM @@ -260,6 +273,7 @@ def matmul_persistent( and (b.dtype == torch.bfloat16) and a.is_contiguous() and b.transpose(0, 1).is_contiguous() + and N >= MIN_DEEPGEMM_DIM ): if _ENABLE_MM_COMPARISON_TEST: out_triton = _matmul_persistent_triton(a=a, b=b, bias=bias)