[Performance] CUDA and ROCm beam search (huge) performance regression

[cloudhan]

Describe the issue

1001ec9 severely deteriorate ROCm beam search performance by >= 2.x

For an test example, before this commit the decoding performance is 0.545s for a sample with batch 2, beam 4 and min len 20, max len 128. With this commit, the time it takes is 2.274s, this is more than 2x worse!

To reproduce

git checkout f5e9625 && git cherry-pick 15f16ef36e # build and run, the perf should be good
git checkout 1001ec9 && git cherry-pick 15f16ef36e # build and run, the perf will be much worse

Urgency

No response

Platform

No response

OS Version

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1001ec9

ONNX Runtime API

Python

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

CUDA and ROCm

Model File

whisper large

Is this a quantized model?

No

[cloudhan]

Ryan Hill (@RyanUnderhill) Tianlei Wu (@tianleiwu) #16272 Introduced a severe performance degeneration, on ROCm. Either a fix or a revert is perfered.

cc. Yao Zhang (@zhangyaobit)

[cloudhan]

I also verified CUDA EP with

git checkout f5e9625 && git cherry-pick b84c63db2a # build and run, the perf should be good
git checkout 1001ec9 && git cherry-pick b84c63db2a # build and run, the perf will be much worse

The perf number is 0.488s and 1.735s respectively, on Tesla V100-DGXS-32GB

NOTE: b84c63d fixed a mistake introduced by me.

[tianleiwu]

cloudhan, what's exact command lines to reproduce the latency number? Ryan tested GPT-2 and T5 in V100, and he did not see perf regression.

[cloudhan]

I ran on the same cuda machine on latest main tip

With the same input, the E2E time scales linearly with max_length, even though the min_length is 20 and the actual outputed sequence length is less than ~30.

max_length	e2e time
32	0.373s
64	0.628s
128	1.157s
256	2.244s
512	err

max_length >= 512 (not sure about the lower bound, tho) produce the following error with whisper large.

onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Non-zero status code returned while running BeamSearch node. Name:'BeamSearch_zcode' Status Message: Non-zero status code returned while running DecoderMaskedMultiHeadAttention node. Name:'Attention_0' Status Message: Input sequence length should be 1 to use DecoderMaskedMultiHeadAttention

Ryan Hill (@RyanUnderhill) This info might be helpful.