diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 84d3b18de73fe..405b6133fa474 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2463,18 +2463,32 @@ static std::vector RunGQAPackedQKVRotaryPrefill( const int half_rotary = head_size / 2; const int max_seq_len = sequence_length + 8; + // The CUDA GQA kernel only registers for MLFloat16/BFloat16, so float inputs + // silently fall back to the CPU EP. Feed fp16 tensors when targeting CUDA so + // the *_CUDA test genuinely exercises the CUDA kernel. The CPU and WebGPU + // kernels both support float (WebGpuSupportedFloatTypes = {float, MLFloat16}), + // so those paths keep fp32 for tighter numeric comparison. + const bool use_fp16 = target_ep == GqaTargetEp::kCuda; + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); tester.AddAttribute("do_rotary", static_cast(1)); // Packed QKV: pass through `query` input, leave key/value as optional edges. - tester.AddInput("query", {batch_size, sequence_length, qkv_hidden}, packed_qkv_data); - tester.AddOptionalInputEdge(); // key (signals packed) - tester.AddOptionalInputEdge(); // value (signals packed) - - tester.AddOptionalInputEdge(); // past_key - tester.AddOptionalInputEdge(); // past_value + if (use_fp16) { + tester.AddInput("query", {batch_size, sequence_length, qkv_hidden}, ToFloat16(packed_qkv_data)); + tester.AddOptionalInputEdge(); // key (signals packed) + tester.AddOptionalInputEdge(); // value (signals packed) + tester.AddOptionalInputEdge(); // past_key + tester.AddOptionalInputEdge(); // past_value + } else { + tester.AddInput("query", {batch_size, sequence_length, qkv_hidden}, packed_qkv_data); + tester.AddOptionalInputEdge(); // key (signals packed) + tester.AddOptionalInputEdge(); // value (signals packed) + tester.AddOptionalInputEdge(); // past_key + tester.AddOptionalInputEdge(); // past_value + } tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); tester.AddInput("total_sequence_length", {1}, {total_sequence_length}, @@ -2490,21 +2504,40 @@ static std::vector RunGQAPackedQKVRotaryPrefill( sin_cache[pos * half_rotary + d] = std::sin(static_cast(pos) * freq); } } - tester.AddInput("cos_cache", {max_seq_len, half_rotary}, cos_cache); - tester.AddInput("sin_cache", {max_seq_len, half_rotary}, sin_cache); + if (use_fp16) { + tester.AddInput("cos_cache", {max_seq_len, half_rotary}, ToFloat16(cos_cache)); + tester.AddInput("sin_cache", {max_seq_len, half_rotary}, ToFloat16(sin_cache)); + } else { + tester.AddInput("cos_cache", {max_seq_len, half_rotary}, cos_cache); + tester.AddInput("sin_cache", {max_seq_len, half_rotary}, sin_cache); + } tester.AddOptionalInputEdge(); // position_ids - tester.AddOptionalInputEdge(); // attention_bias - tester.AddOptionalInputEdge(); // head_sink + if (use_fp16) { + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + } else { + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + } const int output_size = batch_size * sequence_length * hidden_size; - tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, - std::vector(output_size, 0.0f)); const int present_size = batch_size * kv_num_heads * total_sequence_length * head_size; - tester.AddOutput("present_key", {batch_size, kv_num_heads, total_sequence_length, head_size}, - std::vector(present_size, 0.0f)); - tester.AddOutput("present_value", {batch_size, kv_num_heads, total_sequence_length, head_size}, - std::vector(present_size, 0.0f)); + if (use_fp16) { + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(output_size, MLFloat16(0.0f))); + tester.AddOutput("present_key", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_size, MLFloat16(0.0f))); + tester.AddOutput("present_value", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_size, MLFloat16(0.0f))); + } else { + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(output_size, 0.0f)); + tester.AddOutput("present_key", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_size, 0.0f)); + } tester.SetOutputTolerance(1e6f); // We fetch and compare outputs ourselves. @@ -2513,8 +2546,17 @@ static std::vector RunGQAPackedQKVRotaryPrefill( tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); auto fetches = tester.GetFetches(); - const float* out_data = fetches[0].Get().Data(); - return std::vector(out_data, out_data + output_size); + std::vector result(output_size); + if (use_fp16) { + const MLFloat16* out_data = fetches[0].Get().Data(); + for (int i = 0; i < output_size; ++i) { + result[i] = out_data[i].ToFloat(); + } + } else { + const float* out_data = fetches[0].Get().Data(); + std::copy_n(out_data, output_size, result.begin()); + } + return result; } // Inner helper: builds packed-QKV inputs, computes per-prompt references, runs