Fix BFloat16 logits returned as garbage in Logits::Get()

[justinchuby]

Summary

Logits::Get() only converted the model's raw logits to float32 when the output type was Float16. For BFloat16 the conversion was skipped, and the subsequent WrapTensor<float> reinterpreted the raw 2-byte bf16 values as 4-byte float32 — corrupting every logit (wrong argmax, incoherent generation). The identical model in Float16 worked correctly.

Fixes #2202.

Changes

• src/models/logits.cpp: treat BFloat16 the same as Float16 in both the fp32 staging-buffer allocation and the Cast to float32.
• src/cuda/model_kernels.cu + src/cuda/kernels.h: add LaunchBf16ToFp32 (on-device bf16→f32 conversion).
• src/cuda/interface.cpp: route BFloat16 → float through the new kernel in the CUDA Cast. (The CPU Cast already supported bf16→f32, so the generic Cast helper otherwise falls back to a host round-trip.)

Verification

On a bf16 decoder (vocab 262144), the first-token argmax now matches the Float16 / HuggingFace reference, and greedy generation is byte-identical to the fp16 export. Before the fix, genai returned argmax 4539 (a value that did not correspond to any position in the model's actual output); after, it returns 496 (correct).

[Copilot]

Pull request overview

Fixes incorrect logits returned by Logits::Get() when the model output tensor is BFloat16 by ensuring bf16 logits are cast to float32 (matching the existing fp16 behavior) and by adding a CUDA bf16→fp32 conversion path to keep the cast on-device.

Changes:

• Treat Ort::BFloat16_t the same as Ort::Float16_t in Logits::Get() by allocating/using an fp32 staging tensor and invoking Cast() before wrapping logits as float.
• Add a CUDA kernel entry point LaunchBf16ToFp32 for device-side bf16→fp32 conversion.
• Route CUDA Cast() for BFloat16 → float through the new kernel.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
src/models/logits.cpp	Enables bf16 logits to be converted to fp32 before exposing them via WrapTensor<float>().
src/cuda/model_kernels.cu	Adds device kernel implementation for bf16→fp32 conversion.
src/cuda/kernels.h	Declares the new CUDA bf16→fp32 launcher.
src/cuda/interface.cpp	Adds CUDA Cast() dispatch for bf16→fp32 using the new kernel.

[kunal-vaishnavi]

Typically, we have upcasted the logits to float32 inside the model itself.

onnxruntime-genai/src/python/py/models/builders/base.py

Lines 385 to 387 in 13addde

	# Always use float32 logits to improve accuracy in the case of bf16 models.
	if self.io_dtype == ir.DataType.BFLOAT16:
	self.output_types["logits"] = ir.DataType.FLOAT

But I think it is still good to have the upcasting option inside ORT GenAI as a fallback.

[justinchuby]

@kunal-vaishnavi @baijumeswani could you approve? Thanks.

[baijumeswani]

@justinchuby can you merge main to trigger the integration test pipeline?

[justinchuby]

/azp run Integration Tests

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[kunal-vaishnavi]

We need to sync with the main branch of ORT GenAI to get the changes needed to run the integration tests.

[justinchuby]

Doesn't seem to be triggered 🤔

[kunal-vaishnavi]

/azp run Integration Tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).