From 9b549c646c87f8547e02014828bb7482a266c9ce Mon Sep 17 00:00:00 2001 From: PeixuanZuo <94887879+PeixuanZuo@users.noreply.github.com> Date: Tue, 18 Jul 2023 16:47:39 +0800 Subject: [PATCH] [ROCm] fix kernel explorer GemmSoftmaxGemm test (#16735) GemmSoftmaxGemmTunble occasionally broken with large numerical error. The root cause of this error is CK's Strided Batched Gemm has larger error under a specific initialization distribution `(multinormal_distribution)`. Generic(Gemm1 + Softmax + Gemm2) implementation is one instance of GemmSoftmaxGemmTunble. Gemm1 and Gemm2 in Generic implementation are TunableOps when tuning enabled. In some case GemmSoftmaxGemmTunble select Generic implentation, while Gemm1 or Gemm2 select ck implementation, the result of GemmSoftmaxGemmTunble affect by CK. - Make tolerance more loosen. - Add `GemmSoftmaxGemmPermuteGenericNestedTunable` to test Generic implementation with tuning enabled. --- .../kernels/gemm_softmax_gemm_permute_test.py | 22 +++++++++++-- .../kernels/rocm/gemm_softmax_gemm_permute.cu | 32 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py index ed19398f72..64c7c76a1a 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py @@ -179,7 +179,7 @@ def _test_gemm_softmax_gemm_permute( # KERNEL_EXPLORER_STRICT_TEST=1 pytest ... -s -v np.testing.assert_allclose(out, ref) else: - is_zero_tol, atol, rtol = 1e-3, 1e-2, 1e-2 + is_zero_tol, atol, rtol = 1e-3, 2e-2, 1e-2 not_close_to_zeros = np.abs(ref) > is_zero_tol np.testing.assert_allclose(out[not_close_to_zeros], ref[not_close_to_zeros], atol=atol, rtol=rtol) except Exception as err: @@ -200,7 +200,7 @@ def _test_gemm_softmax_gemm_permute( @pytest.mark.parametrize("total_seqlen", total_seqlens) @pytest.mark.parametrize("seqlen", seqlens) @pytest.mark.parametrize("batch", [16]) -@pytest.mark.parametrize("dtype", dtypes) +@pytest.mark.parametrize("dtype", ["float16", "float32"]) def test_gemm_softmax_gemm_permute_generic(dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim): f = getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype)) scale = 1.0 / np.sqrt(head_size) @@ -209,6 +209,24 @@ def test_gemm_softmax_gemm_permute_generic(dtype, batch, seqlen, total_seqlen, n ) +@pytest.mark.parametrize("mask_dim", [2], ids=get_mask_dim_id) +@pytest.mark.parametrize("biased", [False], ids=get_biased_id) +@pytest.mark.parametrize("head_size", [64]) +@pytest.mark.parametrize("nhead", [8]) +@pytest.mark.parametrize("total_seqlen", [128]) +@pytest.mark.parametrize("seqlen", [64]) +@pytest.mark.parametrize("batch", [16]) +@pytest.mark.parametrize("dtype", ["float16", "float32"]) +def test_gemm_softmax_gemm_permute_generic_nested_tunable( + dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim +): + f = getattr(ke, "GemmSoftmaxGemmPermuteGenericNestedTunable_" + dtype_to_suffix(dtype)) + scale = 1.0 / np.sqrt(head_size) + _test_gemm_softmax_gemm_permute( + f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale, ke.qkv_format.Q_K_V_BNSH + ) + + @pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled") @pytest.mark.parametrize("mask_dim", mask_dims, ids=get_mask_dim_id) @pytest.mark.parametrize("biased", biaseds, ids=get_biased_id) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu index caf55f5a5d..5e60bad776 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu @@ -175,6 +175,32 @@ class GemmSoftmaxGemmPermuteGeneric : public IGemmSoftmaxGemmPermuteKernelExplor } }; +template +class GemmSoftmaxGemmPermuteGenericNestedTunable : public GemmSoftmaxGemmPermuteGeneric { + public: + GemmSoftmaxGemmPermuteGenericNestedTunable( + int64_t batch, + int64_t seqlen, + int64_t total_seqlen, + std::optional max_seqlen, + int64_t num_heads, + int64_t head_size, + int64_t mask_dim, + double scale, + contrib::AttentionQkvFormat qkv_format, + DeviceArray& Q, + std::optional& K, + std::optional& V, + std::optional& attn_bias, + std::optional& attn_mask, + DeviceArray& out) + : GemmSoftmaxGemmPermuteGeneric(batch, seqlen, total_seqlen, max_seqlen, + num_heads, head_size, mask_dim, scale, qkv_format, + Q, K, V, attn_bias, attn_mask, out) { + this->params_.TuningContext()->EnableTunableOpAndTuning(); + } +}; + #ifdef USE_COMPOSABLE_KERNEL template class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer { @@ -301,6 +327,9 @@ class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplor #define REGISTER_GENERIC(dtype) \ REGISTER_COMMON("GemmSoftmaxGemmPermuteGeneric_" #dtype, GemmSoftmaxGemmPermuteGeneric, dtype) +#define REGISTER_GENERIC_NESTEDTUNABLE(dtype) \ + REGISTER_COMMON("GemmSoftmaxGemmPermuteGenericNestedTunable_" #dtype, GemmSoftmaxGemmPermuteGenericNestedTunable, dtype) + #define REGISTER_CK(dtype, biased, masked, mask_bias_suffix) \ REGISTER_COMMON( \ "GemmSoftmaxGemmPermuteCK" mask_bias_suffix "_" #dtype, GemmSoftmaxGemmPermuteCK, dtype, biased, masked) @@ -318,6 +347,9 @@ KE_REGISTER(m) { .export_values(); REGISTER_GENERIC(half); + REGISTER_GENERIC(float); + REGISTER_GENERIC_NESTEDTUNABLE(half); + REGISTER_GENERIC_NESTEDTUNABLE(float); #ifdef USE_COMPOSABLE_KERNEL REGISTER_CK(half, false, false, "");