Revert "fallback to triton mm_persistent kernel when deepGemm fail" (#13178)
This commit is contained in:
@@ -253,12 +253,6 @@ def _matmul_persistent_deepgemm(
|
|||||||
def matmul_persistent(
|
def matmul_persistent(
|
||||||
a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
|
a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
|
||||||
):
|
):
|
||||||
M, K = a.shape
|
|
||||||
K2, N = b.shape
|
|
||||||
|
|
||||||
# DeepGEMM requires minimum dimensions, skip DeepGEMM for small dimensions to avoid CUDA_ERROR_INVALID_VALUE
|
|
||||||
MIN_DIM_FOR_DEEPGEMM = 64
|
|
||||||
|
|
||||||
if (
|
if (
|
||||||
_ENABLE_MM_DEEPGEMM
|
_ENABLE_MM_DEEPGEMM
|
||||||
and ENABLE_JIT_DEEPGEMM
|
and ENABLE_JIT_DEEPGEMM
|
||||||
@@ -266,8 +260,6 @@ def matmul_persistent(
|
|||||||
and (b.dtype == torch.bfloat16)
|
and (b.dtype == torch.bfloat16)
|
||||||
and a.is_contiguous()
|
and a.is_contiguous()
|
||||||
and b.transpose(0, 1).is_contiguous()
|
and b.transpose(0, 1).is_contiguous()
|
||||||
and M >= MIN_DIM_FOR_DEEPGEMM
|
|
||||||
and N >= MIN_DIM_FOR_DEEPGEMM
|
|
||||||
):
|
):
|
||||||
if _ENABLE_MM_COMPARISON_TEST:
|
if _ENABLE_MM_COMPARISON_TEST:
|
||||||
out_triton = _matmul_persistent_triton(a=a, b=b, bias=bias)
|
out_triton = _matmul_persistent_triton(a=a, b=b, bias=bias)
|
||||||
@@ -284,12 +276,7 @@ def matmul_persistent(
|
|||||||
# print(f"{a=} {b=} {bias=} {out_triton=} {out_deepgemm=}")
|
# print(f"{a=} {b=} {bias=} {out_triton=} {out_deepgemm=}")
|
||||||
return out_deepgemm
|
return out_deepgemm
|
||||||
|
|
||||||
try:
|
|
||||||
return _matmul_persistent_deepgemm(a=a, b=b, bias=bias)
|
return _matmul_persistent_deepgemm(a=a, b=b, bias=bias)
|
||||||
except RuntimeError:
|
|
||||||
# DeepGEMM failed, fallback to Triton kernel silently
|
|
||||||
# (dimension checks above should prevent most errors)
|
|
||||||
pass
|
|
||||||
|
|
||||||
return _matmul_persistent_triton(a=a, b=b, bias=bias)
|
return _matmul_persistent_triton(a=a, b=b, bias=bias)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user