Skip to content

Fix CUDA BeamSearch fp16 score mapping and add regression test - #31640

Open
apsonawane wants to merge 3 commits into
mainfrom
msrc/cuda-beamsearch-scorecopy-fix
Open

Fix CUDA BeamSearch fp16 score mapping and add regression test#31640
apsonawane wants to merge 3 commits into
mainfrom
msrc/cuda-beamsearch-scorecopy-fix

Conversation

@apsonawane

Copy link
Copy Markdown
Contributor

This pull request introduces improvements to the Beam Search implementation and its test coverage, specifically for FP16 (half-precision float) support. The main changes include a bug fix in the CUDA implementation and the addition of a new test to verify the output type and shape for FP16 scores.

CUDA implementation fix:

  • Changed the type mapping in LaunchBeamSearchScoreCopy from always using float to using the template parameter T, ensuring correct type handling for FP16 and other types. (onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu)

Test improvements:

  • Added a new test, GptBeamSearchFp16_ScoresOutputTypeAndShape, which runs the GPT-2 Beam Search model with FP16 outputs and verifies that the output tensor has the correct type (FLOAT16) and expected shape. (onnxruntime/test/contrib_ops/beam_search_test.cc)

Use ToCudaType<T> in LaunchBeamSearchScoreCopy so MLFloat16 outputs map to half writes instead of float writes.
Cover optional scores output path for tiny_gpt2_beamsearch_fp16 and assert output[2] is float16 with expected rank/leading dims.
@apsonawane
apsonawane requested review from tianleiwu and a lite review from Copilot August 4, 2026 21:42

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR fixes FP16 score output handling in the CUDA BeamSearch path by correcting the device-side output type mapping during score copying, and adds a regression test to validate FP16 scores output type/shape using a tiny GPT-2 beam search model.

Changes:

  • Fix CUDA LaunchBeamSearchScoreCopy to map the output pointer type using the template parameter T (instead of always float) so FP16 scores are written with the correct element type.
  • Add a new BeamSearch test that runs the FP16 GPT-2 beam search model and validates the scores output tensor element type is FLOAT16 and the shape matches expectations.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File Description
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu Corrects the CUDA output type mapping used when converting/copying final beam scores into the output tensor.
onnxruntime/test/contrib_ops/beam_search_test.cc Adds a regression test for FP16 scores output type/shape (CUDA-only execution).
Suppressed comments (1)

onnxruntime/test/contrib_ops/beam_search_test.cc:290

  • Inside the new CUDA-guarded block, the indentation is inconsistent (e.g., Ort::SessionOptions is not indented under the if, and #ifdef USE_CUDA has leading spaces). This reduces readability and deviates from the pattern used in the preceding GptBeamSearchFp16 test.
  if (enable_cuda) {
    Ort::SessionOptions session_options;
#ifdef USE_CUDA
    OrtCUDAProviderOptionsV2 cuda_options;
    cuda_options.use_tf32 = false;
    session_options.AppendExecutionProvider_CUDA_V2(cuda_options);
#endif

Comment thread onnxruntime/test/contrib_ops/beam_search_test.cc Outdated

const auto scores_shape = scores_ts.GetShape();
ASSERT_EQ(scores_shape.size(), static_cast<size_t>(4));
ASSERT_EQ(scores_shape[0], max_length[0] - static_cast<int32_t>(input_ids_shape[1]));
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants