Skip to content

Make DMMHA kernel inside MHA optional for Whisper #25166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kunal-vaishnavi merged 1 commit into from
Jun 25, 2025
Merged

Make DMMHA kernel inside MHA optional for Whisper #25166

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions onnxruntime/python/tools/transformers/models/whisper/convert_to_onnx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,14 @@ def parse_arguments(argv=None):
)
conversion_args.set_defaults(no_beam_search_op=False)

conversion_args.add_argument(
"--use_decoder_masked_mha",
required=False,
action="store_true",
help="Use DecoderMaskedMultiHeadAttention kernel for improved performance. This is currently an experimental feature.",
)
conversion_args.set_defaults(use_decoder_masked_mha=False)

#############################################################
# Optional inputs for Whisper
# (listed below in the order that WhisperBeamSearch expects)
Expand Down Expand Up @@ -305,8 +313,13 @@ def parse_arguments(argv=None):
quant_args.set_defaults(quantize_reduce_range=False)

args = parser.parse_args(argv)

# Collect cross QKs if either flag is enabled
args.collect_cross_qk = args.collect_cross_qk or args.output_cross_qk

# FP32 CPU can be supported here once the DMMHA CPU kernel bugs are fixed
args.use_decoder_masked_mha = args.use_decoder_masked_mha and args.provider == "cuda"

return args


Expand All @@ -323,6 +336,7 @@ def export_onnx_models(
use_forced_decoder_ids: bool = False,
merge_encoder_and_decoder_init: bool = True,
no_beam_search_op: bool = False,
use_decoder_masked_mha: bool = False,
output_qk: bool = False,
overwrite: bool = False,
use_int32_inputs: bool = True,
Expand Down Expand Up @@ -402,6 +416,7 @@ def export_onnx_models(
provider=provider,
is_decoder=(name == "decoder"),
no_beam_search_op=no_beam_search_op,
use_decoder_masked_mha=use_decoder_masked_mha,
output_qk=output_qk,
)
# Remove old ONNX model and old data file
Expand Down Expand Up @@ -474,6 +489,7 @@ def main(argv=None):
args.use_forced_decoder_ids,
not args.separate_encoder_and_decoder_init,
args.no_beam_search_op,
args.use_decoder_masked_mha,
args.output_cross_qk,
args.overwrite,
not args.use_int64_inputs,
Expand Down Expand Up @@ -541,6 +557,7 @@ def main(argv=None):
args.model_name_or_path,
args.provider,
args.separate_encoder_and_decoder_init,
args.use_decoder_masked_mha,
args.output_cross_qk,
next(iter(filter(lambda path: "encoder" in path, output_paths))),
next(iter(filter(lambda path: "decoder" in path, output_paths))),
Expand Down
17 changes: 8 additions & 9 deletions onnxruntime/python/tools/transformers/models/whisper/whisper_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ def save_processing(
model_name_or_path: str,
provider: str,
separate_encoder_and_decoder_init: bool,
use_decoder_masked_mha: bool,
output_qk: bool,
encoder_path: str,
decoder_path: str,
Expand Down Expand Up @@ -596,24 +597,21 @@ def save_processing(
"no_repeat_ngram_size": 0,
"num_beams": 1,
"num_return_sequences": 1,
"past_present_share_buffer": provider == "cuda",
"past_present_share_buffer": use_decoder_masked_mha,
"repetition_penalty": 1.0,
"temperature": 1.0,
"top_k": 1,
"top_p": 1.0,
},
}

# Requirements for the DMMHA kernel (which is currently
# enabled for CUDA only):
# Requirements for the DMMHA kernel:
# - Buffer sharing = true
# - New input: past_sequence_length
# - New input: cache_indirection
# Otherwise, buffer sharing should be false and the new inputs
# should not be added for beam search to work in ORT GenAI.

if provider == "cuda":
# Add inputs for DMMHA kernel
# Otherwise, buffer sharing should be false and the new inputs should not be added
# for beam search to work in ORT GenAI.
if use_decoder_masked_mha:
genai_config["model"]["decoder"]["inputs"].update(
{
"past_sequence_length": "past_sequence_length",
Expand Down Expand Up @@ -771,6 +769,7 @@ def optimize_onnx(
provider: str = "cpu",
is_decoder: bool = False,
no_beam_search_op: bool = False,
use_decoder_masked_mha: bool = False,
output_qk: bool = False,
):
"""Optimize ONNX model with an option to convert it to use mixed precision."""
Expand All @@ -794,7 +793,7 @@ def optimize_onnx(

# Add `past_sequence_length`, `cache_indirection`, and `output_qk` to `MultiHeadAttention` ops
if is_decoder and no_beam_search_op:
if provider == "cuda": # FP32 CPU can be supported here once the DMMHA CPU kernel bugs are fixed
if use_decoder_masked_mha:
# FP16 CUDA, FP32 CUDA, and FP32 CPU use the `DecoderMaskedMultiHeadAttention` kernel
# via `MultiHeadAttention`, which requires the `past_sequence_length` and
# `cache_indirection` inputs
Expand Down