From ddc5774d522133e27392b515fc77650f635a7b11 Mon Sep 17 00:00:00 2001 From: Ruibin Cheung Date: Wed, 4 Sep 2024 00:25:28 +0800 Subject: [PATCH] [PyTorch] Add contiguous check for `te_grouped_gemm` (#1146) [PyTorch] Add contiguous check for grouped gemm Signed-off-by: beinggod Co-authored-by: beinggod Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/csrc/extensions/gemm.cu | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cu b/transformer_engine/pytorch/csrc/extensions/gemm.cu index 01fb94cab4..7405914a0e 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cu +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cu @@ -115,6 +115,11 @@ void te_grouped_gemm(std::vector A, at::Tensor A_scale_inverse, int if (pre_gelu_out[i].data_ptr() != nullptr) pre_gelu_out[i].zero_(); continue; } + + NVTE_CHECK(A[i].is_contiguous(), "A[", i, "] must be contiguous."); + NVTE_CHECK(B[i].is_contiguous(), "B[", i, "] must be contiguous."); + NVTE_CHECK(D[i].is_contiguous(), "D[", i, "] must be contiguous."); + te_A.emplace_back(make_tensor( A[i].data_ptr(), {static_cast(A[i].size(0)), static_cast(A[i].size(1))}, A_type, nullptr, nullptr, getDataPtr(A_scale_inverse, A_offset + i)));