saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] fix kernel explorer GemmSoftmaxGemm test (#16735)

Browse source 
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.
This commit is contained in:
PeixuanZuo 2023-07-18 16:47:39 +08:00 committed by GitHub
parent 9ba5cdbaa4
commit 9b549c646c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 52 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py
View file

@ -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)

32
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu
View file

@ -175,6 +175,32 @@ class GemmSoftmaxGemmPermuteGeneric : public IGemmSoftmaxGemmPermuteKernelExplor
}
};
template <typename T>
class GemmSoftmaxGemmPermuteGenericNestedTunable : public GemmSoftmaxGemmPermuteGeneric<T> {
public:
GemmSoftmaxGemmPermuteGenericNestedTunable(
int64_t batch,
int64_t seqlen,
int64_t total_seqlen,
std::optional<int64_t> max_seqlen,
int64_t num_heads,
int64_t head_size,
int64_t mask_dim,
double scale,
contrib::AttentionQkvFormat qkv_format,
DeviceArray& Q,
std::optional<DeviceArray>& K,
std::optional<DeviceArray>& V,
std::optional<DeviceArray>& attn_bias,
std::optional<DeviceArray>& attn_mask,
DeviceArray& out)
: GemmSoftmaxGemmPermuteGeneric<T>(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 <typename T, bool USE_BIAS, bool USE_MASK>
class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer<T> {
@ -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, "");