DMMHA: add unit tests; fix CPU, CUDA kernel

[mindest]

Description

Fixes:
(1) cpu kernel: applying scale before bias and mask like other MHA ops
(2) cpu kernel: correct offset during appending past to present.
(3) cuda kernel: apply mask if provided; fix output_qk offset.

Add DMMHA unit tests

[hariharans29]

When the PR is ready could you please update the PR title and description to better reflect the problem and fix ? Thanks.

[mindest]

@hariharans29, is it true that for cross attention CUDA kernel, the key layout is also reordered as [B, H, head_size/x, L, x],

onnxruntime/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu

Lines 105 to 107 in dd28f09

	// The layout of the cache is [B, H, head_size/x, L, x] with x == 4/8/16 for FP32/FP16/FP8. Since each thread
	// owns x elements, we have to decompose the linear index into chunks of x values and the posi-
	// tion of the thread in that chunk.

instead of BNSH?

onnxruntime/onnxruntime/core/graph/contrib_ops/bert_defs.cc

Lines 855 to 860 in dd28f09

	.Input(1,
	"key",
	"Key with shape (batch_size, 1, hidden_size) for self attention "
	"or past_key with shape (batch_size, num_heads, kv_sequence_length, head_size) for cross attention",
	"T",
	OpSchema::Optional)

[mindest]

/azp run Big Models,Linux Android Emulator QNN CI Pipeline,Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline

[azure-pipelines]

Azure Pipelines successfully started running 9 pipeline(s).

[mindest]

/azp run ONNX Runtime Web CI Pipeline,Windows ARM64 QNN CI Pipeline,Windows CPU CI Pipeline,Windows GPU CUDA CI Pipeline,Windows GPU DML CI Pipeline,Windows GPU Doc Gen CI Pipeline,Windows GPU TensorRT CI Pipeline,Windows x64 QNN CI Pipeline

[azure-pipelines]

Azure Pipelines successfully started running 8 pipeline(s).

[kunal-vaishnavi]

@hariharans29, is it true that for cross attention CUDA kernel, the key layout is also reordered as [B, H, head_size/x, L, x],

onnxruntime/onnxruntime/contrib_ops/cuda/bert/fastertransformer_decoder_attention/decoder_masked_multihead_attention_impl.cu

Lines 105 to 107 in dd28f09

	// The layout of the cache is [B, H, head_size/x, L, x] with x == 4/8/16 for FP32/FP16/FP8. Since each thread
	// owns x elements, we have to decompose the linear index into chunks of x values and the posi-
	// tion of the thread in that chunk.

instead of BNSH?

onnxruntime/onnxruntime/core/graph/contrib_ops/bert_defs.cc

Lines 855 to 860 in dd28f09

	.Input(1,
	"key",
	"Key with shape (batch_size, 1, hidden_size) for self attention "
	"or past_key with shape (batch_size, num_heads, kv_sequence_length, head_size) for cross attention",
	"T",
	OpSchema::Optional)

For self-attention, parameters.k_cache = present_key_data = past_key_data since past_present_share_buffer = true. For cross-attention, parameters.k_cache = key_data.

onnxruntime/onnxruntime/contrib_ops/cuda/bert/decoder_masked_multihead_attention.cc

Lines 150 to 202 in dd28f09

	if (past_key == nullptr && present_key == nullptr) {
	if (attention_bias != nullptr) {
	return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED,
	"DecoderMaskedMultiHeadAttention does not support attention bias for cross-attention");
	}
	parameters.is_cross_attention = true;
	parameters.total_sequence_length = parameters.kv_sequence_length;
	parameters.max_sequence_length = parameters.kv_sequence_length;
	// parameters.k and parameters.v are nullptr
	parameters.k_cache = const_cast<T1*>(key->Data<T1>());
	parameters.v_cache = const_cast<T1*>(value->Data<T1>());
	parameters.k_bias = nullptr;
	parameters.v_bias = nullptr;
	} else {
	// Sanity check
	ORT_ENFORCE(past_present_share_buffer_);
	ORT_ENFORCE(past_key != nullptr && past_value != nullptr);
	auto* present_key_data = present_key->MutableData<T1>();
	auto* present_value_data = present_value->MutableData<T1>();
	auto* past_key_data = past_key->Data<T1>();
	auto* past_value_data = past_value->Data<T1>();
	// No production use-case will incur this copy cost as the implementation of
	// GreedySearch/BeamSearch is written in such a way that the past and present buffers
	// will be shared.
	// This is just to circumvent the OpTester's limitation of not being able to bind a specific
	// buffer to inputs/outputs.
	if (present_key_data != past_key_data) {
	CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(present_key_data, past_key_data, past_key->SizeInBytes(),
	cudaMemcpyDeviceToDevice, cuda_stream));
	}
	if (present_value_data != past_value_data) {
	CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(present_value_data, past_value_data, past_value->SizeInBytes(),
	cudaMemcpyDeviceToDevice, cuda_stream));
	}
	parameters.is_cross_attention = false;
	bool is_packed_qkv = (key == nullptr && value == nullptr);
	parameters.is_packed_qkv = is_packed_qkv;
	parameters.k = is_packed_qkv
	? const_cast<T1*>(query->Data<T1>() + parameters.hidden_size)
	: const_cast<T1*>(key->Data<T1>());
	parameters.v = is_packed_qkv
	? const_cast<T1*>(query->Data<T1>() + 2 * static_cast<size_t>(parameters.hidden_size))
	: const_cast<T1*>(value->Data<T1>());
	parameters.k_cache = present_key_data;
	parameters.v_cache = present_value_data;
	}

I believe parameters.k_cache should be reordered before the kernel is launched so past_key is already reordered for self-attention and key is already reordered for cross-attention. This behavior would match the below comments.

onnxruntime/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_whisper.h

Lines 341 to 352 in 03ea5dc

	#ifdef USE_CUDA
	// Here we only need to reorder the past key for self-attention and cross-attention.
	for (size_t i = 0; i < 2 * static_cast<size_t>(decoder_subgraph_.num_layers); ++i) {
	ORT_RETURN_IF_ERROR(reorder_past_state_func_(cuda_device_prop_,
	*decoder_feeds[offset + 2 * i].GetMutable<Tensor>(),
	beam_state.staging_for_past_state_reorder,
	this->ort_stream_));
	}
	size_t cache_indir_input_offset = static_cast<size_t>(decoder_subgraph_.GetFirstPastInputIndex()) + 4 * static_cast<size_t>(decoder_subgraph_.num_layers) + 2;
	ORT_RETURN_IF_ERROR(init_cache_indir_func_(*decoder_feeds[cache_indir_input_offset].GetMutable<Tensor>(), this->ort_stream_));
	#endif
	}

[tianleiwu]

@kunal-vaishnavi, what's the reason that need cross attention in this op? I think cross attention shall be supported by MHA, and this op is only for decoding only. That could make logic more clearly.

[mindest]

To #22567 (comment): Makes sense to me, in the tests I also use reordered key as input for cross attention. Maybe I should also add some comments in

onnxruntime/onnxruntime/core/graph/contrib_ops/bert_defs.cc

Lines 855 to 860 in dd28f09

	.Input(1,
	"key",
	"Key with shape (batch_size, 1, hidden_size) for self attention "
	"or past_key with shape (batch_size, num_heads, kv_sequence_length, head_size) for cross attention",
	"T",
	OpSchema::Optional)

so that it is clear key in cross-attention is also reordered for CUDA EP.

[mindest]

Thanks @tianleiwu @kunal-vaishnavi for the review! Is this PR ready to merge, if I keep the following changes in another PR?

• Support float16 for out_qk
• Add comments why it is sum_tlength + 1 instead of tlength (for cross-attn total length is kv_seq_len, for self-attn it is past_len + 1)
• Update schema comments of input 1 key for cross-attn.

[kunal-vaishnavi]

@kunal-vaishnavi, what's the reason that need cross attention in this op? I think cross attention shall be supported by MHA, and this op is only for decoding only. That could make logic more clearly.

Whisper uses alternating layers of self-attention and cross-attention during decoding.

[mindest]

Thanks @tianleiwu, @kunal-vaishnavi, @hariharans29!