Fix bug in Squeeze for getting the value of total_seq_len

[Honry]

The Squeeze op is used for removing single-dimensional entries from the shape of a tensor. In this node the axes input is set to [0] which would only eliminate the first axis and lead to the output shape to be [1] if the batch_size is 1.This would cause ShapeInference error at https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/graph/graph.cc#L111-L132 if it is not a strict mode.

This PR fixes the issue by:

• Changing the ReduceSum output shape to [batch_size] by adding keepdims=0 attribute
• Using ReduceMax instead of Squeeze to get the value of total_seq_len and make it as a scalar, this would cover scenarios when batch_size > 1

[Honry]

@qjia7, PTAL, thanks!

[tianleiwu]

Since the reduceSum output shape is [batch_size, 1]. Squeeze might not be the right way to get total_seq_len unless we assume that batch_size==1.

A better way is to let ReduceSum output shape [batch_size] by adding keepdims=0 attribute. Then use ReduceMax to get the maximum one (assumption is the longest sequence does not have padding), or use Gather to get first item if we assume that batch_size==1 (if we only enable graph capture for one batch scenario)

If we do not use keepdims=0, the seqlen_k will have shape [batch_size, 1], while the expected shape is [batch_size]. That's not expected as well.

[qjia7]

A better way is to let ReduceSum output shape [batch_size] by adding keepdims=0 attribute. Then use ReduceMax to get the maximum one (assumption is the longest sequence does not have padding), or use Gather to get first item if we assume that batch_size==1 (if we only enable graph capture for one batch scenario)

Great idea! I did assume that batch_size was one when adding those code. Your suggestion looks great. @tianleiwu one more question: why expose a scalar total_seq_len since the batch can be larger than one?

        #          attention_mask
        #               |
        #         Cast to int32
        #               |
        #    ReduceSum (keepdims=0)
        #              /    \
        #             /      \
        #           Sub    ReduceMax
        #            |        |
        #       seqlens_k  total_seq_len
        #         (1D)       (scalar)

[Honry]

Thanks @tianleiwu @qjia7, I also had the concern about how to handle batch_size > 1.

@tianleiwu's suggestion is really a great idea! I've addressed it in the new commit, PTAL again, thanks!

[qjia7]

@Honry You may also need to update the ReduceSum in make_attention_mask_standard_reformatting_for_gqa in the similar way to get the correct 1D seqlens_k.

[Honry]

@Honry You may also need to update the ReduceSum in make_attention_mask_standard_reformatting_for_gqa in the similar way to get the correct 1D seqlens_k.

@qjia7 thanks! Done, pls. take another look.

[kunal-vaishnavi]

When the original logic to obtain seqlens_k and total_seq_len from the attention mask for the GQA op was added, it was assumed that batch_size = 1 since most inference workloads with ORT GenAI are for batch size = 1.