MultiHeadAttention with dynamic kv-cache attention

[turneram]

Motivation

Add support for parsing and running MHA with dynamic past-sequence-length.

Technical Details

Updates parser and makes changes necessary for running kv-cache attention.

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

• • Added: New functionality.
• • Changed: Changes to existing functionality.
• • Removed: Functionality or support that has been removed. (Compared to a previous release)
• • Optimized: Component performance that has been optimized or improved.
• • Resolved Issues: Known issues from a previous version that have been resolved.
• • Not Applicable: This PR is not to be included in the changelog.

[github-actions]

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

format.py

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_multi_head_attention.cpp

Line 1059 in 29cecde

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_multi_head_attention.cpp

Line 1098 in 29cecde

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_multi_head_attention.cpp

Line 1101 in 29cecde

qk_output = info.add_instruction(make_op("convert", {{"target_type", value->get_shape().type()}}), qk_output);

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_rotary_embedding.cpp

Line 117 in 29cecde

// check if input is dynamic, if so, check if all dims are fixed

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_rotary_embedding.cpp

Line 120 in 29cecde

if(not all_of(input->get_shape().dyn_dims(), [](const shape::dynamic_dimension& dim) { return dim.is_fixed(); }))

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_rotary_embedding.cpp

Line 124 in 29cecde

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_rotary_embedding.cpp

Line 127 in 29cecde

auto input_shape = input->get_shape();

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_rotary_embedding.cpp

Line 129 in 29cecde

auto input_dyn_dims = input_dyn_shape.dyn_dims();

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_rotary_embedding.cpp

Line 131 in 29cecde

if(std::all_of(input_dyn_dims.begin(), input_dyn_dims.end(), [](const shape::dynamic_dimension& dim) { return dim.is_fixed(); }))

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/onnx/parse_rotary_embedding.cpp

Line 133 in 29cecde

std::transform(input_dyn_dims.begin(), input_dyn_dims.end(), input_lens.begin(), [](const shape::dynamic_dimension& dim) { return dim.get_interval().max; });

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 68 in 29cecde

auto in_shape = in->get_shape();

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 70 in 29cecde

auto in_dyn_dims = in_dyn_shape.dyn_dims();

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 72 in 29cecde

if(std::all_of(in_dyn_dims.begin(), in_dyn_dims.end(), [](const shape::dynamic_dimension& dim) { return dim.is_fixed(); }))

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 74 in 29cecde

std::transform(in_dyn_dims.begin(), in_dyn_dims.end(), in_lens.begin(), [](const shape::dynamic_dimension& dim) { return dim.get_interval().max; });

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 192 in 29cecde

auto in_shape = in->get_shape();

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 194 in 29cecde

auto in_dyn_dims = in_dyn_shape.dyn_dims();

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 196 in 29cecde

if(std::all_of(in_dyn_dims.begin(), in_dyn_dims.end(), [](const shape::dynamic_dimension& dim) { return dim.is_fixed(); }))

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/src/op/builder/rotary_embedding.cpp

Line 198 in 29cecde

std::transform(in_dyn_dims.begin(), in_dyn_dims.end(), in_lens.begin(), [](const shape::dynamic_dimension& dim) { return dim.get_interval().max; });

[gh-app-migraphx-bot-pr-write]

Test	Batch	New Rate (9ffa22)	Old Rate (360c7e)	Diff	Status
torchvision-resnet50	64	3,159.60	3,138.13	0.68%	✅
torchvision-resnet50_fp16	64	6,673.44	6,626.36	0.71%	✅
torchvision-densenet121	32	1,554.01	2,694.15	-42.32%	🔴
torchvision-densenet121_fp16	32	362.43	4,530.41	-92.00%	🔴
torchvision-inceptionv3	32	1,738.65	1,796.21	-3.20%	✅
torchvision-inceptionv3_fp16	32	2,735.35	2,820.95	-3.03%	✅
cadene-inceptionv4	16	823.54	824.22	-0.08%	✅
cadene-resnext64x4	16	473.41	783.03	-39.54%	🔴
slim-mobilenet	64	8,028.78	8,387.11	-4.27%	✅
slim-nasnetalarge	64	188.83	229.03	-17.55%	🔴
slim-resnet50v2	64	588.69	3,163.68	-81.39%	🔴
bert-mrpc-onnx	8	102.86	1,172.03	-91.22%	🔴
bert-mrpc-tf	1	96.97	488.32	-80.14%	🔴
pytorch-examples-wlang-gru	1	150.29	327.58	-54.12%	🔴
pytorch-examples-wlang-lstm	1	203.77	456.34	-55.35%	🔴
torchvision-resnet50_1	1	618.25	791.11	-21.85%	🔴
cadene-dpn92_1	1	106.66	444.86	-76.02%	🔴
cadene-resnext101_1	1	338.14	365.24	-7.42%	🔴
onnx-taau-downsample	1	401.53	398.75	0.70%	✅
dlrm-criteoterabyte	1	25.13	32.34	-22.28%	🔴
dlrm-criteoterabyte_fp16	1	nan	51.84	nan	❌
agentmodel	1	742.07	9,815.28	-92.44%	🔴
unet_fp16	2	56.61	56.77	-0.28%	✅
resnet50v1_fp16	1	953.51	935.30	1.95%	✅
resnet50v1_int8	1	942.92	929.37	1.46%	✅
bert_base_cased_fp16	64	882.14	1,098.93	-19.73%	🔴
bert_large_uncased_fp16	32	288.86	346.15	-16.55%	🔴
bert_large_fp16	1	11.01	205.08	-94.63%	🔴
distilgpt2_fp16	16	954.26	2,093.37	-54.42%	🔴
yolov5s	1	127.37	596.84	-78.66%	🔴
tinyllama	1	38.70	45.91	-15.70%	🔴
vicuna-fastchat	1	42.07	44.09	-4.59%	✅
whisper-tiny-encoder	1	106.32	417.74	-74.55%	🔴
whisper-tiny-decoder	1	128.18	417.56	-69.30%	🔴
llama2_7b	1	20.06	20.38	-1.57%	✅
qwen1.5-7b	1	23.68	23.53	0.64%	✅
phi3-3.8b	1	16.73	26.74	-37.43%	🔴
llama3-8b	1	21.81	21.69	0.55%	✅
whisper-large-encoder	1	8.69	10.28	-15.41%	🔴
whisper-large-decoder	1	8.13	104.46	-92.22%	🔴
mistral-7b	1	23.86	23.68	0.74%	✅
FLUX.1-schnell	1	752.32	769.92	-2.29%	✅

Regressions detected 🔴

[gh-app-migraphx-bot-pr-write]

Test	Status	Result
bert-mrpc-onnx	✅	PASSED: MIGraphX meets tolerance
bert-mrpc-tf	✅	PASSED: MIGraphX meets tolerance
pytorch-examples-wlang-gru	✅	PASSED: MIGraphX meets tolerance
pytorch-examples-wlang-lstm	✅	PASSED: MIGraphX meets tolerance
dlrm-criteoterabyte	✅	PASSED: MIGraphX meets tolerance
agentmodel	✅	PASSED: MIGraphX meets tolerance
unet	✅	PASSED: MIGraphX meets tolerance
resnet50v1	✅	PASSED: MIGraphX meets tolerance
bert_base_cased_fp16	✅	PASSED: MIGraphX meets tolerance
bert_large_uncased_fp16	🔴	FAILED: MIGraphX is not within tolerance - check verbose output
bert_large	✅	PASSED: MIGraphX meets tolerance
yolov5s	✅	PASSED: MIGraphX meets tolerance
tinyllama	✅	PASSED: MIGraphX meets tolerance
vicuna-fastchat	✅	PASSED: MIGraphX meets tolerance
whisper-tiny-encoder	✅	PASSED: MIGraphX meets tolerance
whisper-tiny-decoder	✅	PASSED: MIGraphX meets tolerance
distilgpt2_fp16	✅	PASSED: MIGraphX meets tolerance
llama2_7b	✅	PASSED: MIGraphX meets tolerance
qwen1.5-7b	✅	PASSED: MIGraphX meets tolerance
phi3-3.8b	✅	PASSED: MIGraphX meets tolerance
llama3-8b	✅	PASSED: MIGraphX meets tolerance
whisper-large-encoder	✅	PASSED: MIGraphX meets tolerance
whisper-large-decoder	✅	PASSED: MIGraphX meets tolerance
mistral-7b	✅	PASSED: MIGraphX meets tolerance
FLUX.1-schnell	✅	PASSED: MIGraphX meets tolerance