Repository: microsoft/onnxruntime
File: onnxruntime/core/graph/contrib_ops/bert_defs.cc (C++)

Code before:
void GroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int past_key_index, int qk_output_index) {
  // TODO(aciddelgado): propagate output shapes depending if kv-share buffer is on or not
  constexpr int use_max_past_present_buffer = -1;
  BaseGroupQueryAttentionTypeAndShapeInference(ctx, past_key_index, use_max_past_present_buffer, qk_output_index);
}

void SparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int past_key_index) {
  constexpr int use_max_past_present_buffer = 1;
  constexpr int qk_output_index = -1;
  BaseGroupQueryAttentionTypeAndShapeInference(ctx, past_key_index, use_max_past_present_buffer, qk_output_index);
}

constexpr const char* Attention_ver1_doc = R"DOC(
Multi-Head Attention that can be either unidirectional (like GPT-2) or bidirectional (like BERT).

The weights for input projection of Q, K and V are merged. The data is stacked on the second dimension. Its shape
is (input_hidden_size, hidden_size + hidden_size + v_hidden_size). Here hidden_size is the hidden dimension of Q and K,
and v_hidden_size is that of V.

The mask_index is optional. Besides raw attention mask with shape (batch_size, total_sequence_length)
or (batch_size, sequence_length, total_sequence_length) with value 0 for masked and 1 otherwise,
we support other two formats: When input has right-side padding, mask_index is one dimension with shape (batch_size),
where value is actual sequence length excluding padding. When input has left-side padding, mask_index has
shape (2 * batch_size), where the values are the exclusive end positions followed by the inclusive start positions.

When unidirectional is 1, each token only attends to previous tokens.

Both past and present state are optional. They shall be used together, and not allowed to use only one of them.
The qkv_hidden_sizes is required only when K and V have different hidden sizes.

When there is past state, hidden dimension for Q, K and V shall be the same.

The total_sequence_length is past_sequence_length + kv_sequence_length. Here kv_sequence_length is the length of K or V.
For self attention, kv_sequence_length equals to sequence_length (sequence length of Q).
For cross attention, query and key might have different lengths.
)DOC";

// Currently, the `convert_generation.py` script renames the `Attention` nodes to `DecoderMaskedSelfAttention`
// if the user requests it. Hence, the schemas of `DecoderMaskedSelfAttention` and `Attention` schemas
// are tightly coupled. A change in Attention also needs corresponding schema updates in `DecoderMaskedSelfAttention`
// and its kernel.
// TODO(hasesh): Decouple the schema of `DecoderMaskedSelfAttention` from the schema of the `Attention` operator
// by making appropriate tool changes.

ONNX_MS_OPERATOR_SET_SCHEMA(
    Attention, 1,
    OpSchema()
        .SetDoc(Attention_ver1_doc)
        .Attr("num_heads", "Number of attention heads", AttributeProto::INT)
        .Attr("unidirectional",
              "Whether every token can only attend to previous tokens. Default value is 0.",

Review comment (performance):
The spec says "When past_key uses same tensor as present_key, it is of length max_sequence_length... otherwise of length past_sequence_length +kv_sequence_length.". That means the shape depends on whether past and present buffer sharing is used. If so, it is max_seq_length, otherwise total_seq_length.

Usually, cuda/cpu EP will allocate a buffer with max_seq_length for kv cache. So the shape is like [batch_size, kv_num_heads, max_seq_len, head_size] and use I/O binding to bind to past_key and present_key tensors.

Is this use case for webgpu that only allocate total_seq_len buffer for prefill?

Code after:
      auto& query_shape = getInputShape(ctx, 0);
      auto& query_dims = query_shape.dim();

      int64_t num_heads = getAttribute(ctx, "num_heads", 0);
      int64_t kv_num_heads = getAttribute(ctx, "kv_num_heads", 0);

      if (num_heads > 0 && kv_num_heads > 0 && query_dims.size() == 3 && query_dims[2].has_dim_value()) {
        int64_t hidden_size = query_dims[2].dim_value();
        int64_t head_size = 0;

        if (hasInputShape(ctx, 2)) {
          // query shape is (batch_size, sequence_length, num_heads * head_size)
          head_size = hidden_size / num_heads;
        } else {
          // Packed QKV: query shape is (batch_size, sequence_length, (num_heads + 2 * kv_num_heads) * head_size)
          head_size = hidden_size / (num_heads + 2 * kv_num_heads);
        }

        if (head_size > 0) {
          // Determine present_sequence_length from total_sequence_length or kv_sequence_length
          int64_t present_sequence_length = 0;
          if (total_sequence_length_value > 0) {
            present_sequence_length = total_sequence_length_value;
          } else if (kv_sequence_length > 0) {
            present_sequence_length = kv_sequence_length;
          }

          // present key/value shape is (batch_size, kv_num_heads, present_sequence_length, head_size)
          ONNX_NAMESPACE::TensorShapeProto present_shape;
          *present_shape.add_dim() = query_dims[0];  // batch_size
          present_shape.add_dim()->set_dim_value(kv_num_heads);
          if (present_sequence_length > 0) {
            present_shape.add_dim()->set_dim_value(present_sequence_length);
          } else {
            // Fallback: use query sequence_length (dim 1) as present_sequence_length for prefill
            *present_shape.add_dim() = query_dims[1];
          }
          present_shape.add_dim()->set_dim_value(head_size);

          updateOutputShape(ctx, 1, present_shape);
          updateOutputShape(ctx, 2, present_shape);
        }
      }
    }
  }
}

void GroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int past_key_index, int qk_output_index) {
  // TODO(aciddelgado): propagate output shapes depending if kv-share buffer is on or not
  constexpr int use_max_past_present_buffer = -1;
  BaseGroupQueryAttentionTypeAndShapeInference(ctx, past_key_index, use_max_past_present_buffer, qk_output_index);