From ed60d7ad8a493a9297f047c2340c74748beb6d89 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 22 Jun 2026 14:34:29 -0700 Subject: [PATCH 1/2] Fix GroupQueryAttention right-padded rotary prefill CUDA test BatchedRightPaddedRotaryPrefill_CUDA fed fp32 inputs via AddInput. The CUDA/WebGPU GroupQueryAttention kernels only register for MLFloat16/BFloat16, so the fp32 node silently fell back to the CPU EP and the _CUDA test never exercised the CUDA kernel it is named for. Make RunGQAPackedQKVRotaryPrefill feed fp16 tensors when targeting a GPU EP (matching the existing RunGQASharedKVFp16 convention and the test's own fp16 tolerance), so the test runs on the actual CUDA kernel. The CPU path is unchanged. Verified the CUDA fp16 path passes the right-padded prefill. --- .../group_query_attention_op_test.cc | 77 ++++++++++++++----- 1 file changed, 59 insertions(+), 18 deletions(-) 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..3376dd9d9b973 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,31 @@ static std::vector RunGQAPackedQKVRotaryPrefill( const int half_rotary = head_size / 2; const int max_seq_len = sequence_length + 8; + // The CUDA/WebGPU GQA kernels only register for MLFloat16/BFloat16, so float + // inputs silently fall back to the CPU EP. To genuinely exercise the GPU + // kernels (which is the point of the *_CUDA / *_WebGPU tests), feed fp16 + // tensors when targeting a GPU EP. + const bool use_fp16 = target_ep != GqaTargetEp::kCpu; + 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 +2503,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 +2545,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 From 0b7708629aabadcbd0eecdd531fd48bb08a8a6ed Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 22 Jun 2026 14:50:59 -0700 Subject: [PATCH 2/2] update --- .../test/contrib_ops/group_query_attention_op_test.cc | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) 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 3376dd9d9b973..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,11 +2463,12 @@ static std::vector RunGQAPackedQKVRotaryPrefill( const int half_rotary = head_size / 2; const int max_seq_len = sequence_length + 8; - // The CUDA/WebGPU GQA kernels only register for MLFloat16/BFloat16, so float - // inputs silently fall back to the CPU EP. To genuinely exercise the GPU - // kernels (which is the point of the *_CUDA / *_WebGPU tests), feed fp16 - // tensors when targeting a GPU EP. - const bool use_fp16 = target_ep != GqaTargetEp::kCpu; + // 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));