BFloat16 logits returned as garbage — Logits::Get() only casts Float16 to float32
### Describe the bug
For models with **BFloat16** logits output, `onnxruntime-genai` returns garbage logits, producing a wrong/incoherent first token and degenerate generation. The identical model exported in **Float16** works correctly.
### Root cause
`src/models/logits.cpp` (`Logits::Get()`) only converts the model's raw logits to `float32` when the output type is **Float16**. For **BFloat16** it skips the conversion, then `WrapTensor<float>(...)` reinterprets the raw 2-byte bf16 bytes as 4-byte float32, corrupting every logit.
Two places are gated on `Float16` only:
```cpp
// ~line 47 — fp32 staging buffer is allocated only for Float16
if (type_ == Ort::TypeToTensorType<Ort::Float16_t>)
logits_of_last_token_fp32_ = OrtValue::CreateTensor<float>(...);
// ~line 73 — Cast to float32 runs only for Float16
if (type_ == Ort::TypeToTensorType<Ort::Float16_t>) {
Cast(*logits_of_last_token, logits_of_last_token_fp32_, *model_.p_device_inputs_, Ort::TypeToTensorType<float>);
logits_of_last_token = logits_of_last_token_fp32_.get();
}
```
For a bf16 model the `Cast` never runs, so `logits_` ends up wrapping raw bf16 bytes as `float`.
### How I verified
Isolated a bf16 vs fp16 model (Gemma-style decoder, vocab 262144):
- HF bf16 reference, raw ORT `sess.run` of the bf16 decoder, and genai's own dumped model-output logits **all agree** (argmax 496, top-5 `[496,886,506,9079,992]`).
- genai's `get_logits()` for the **same** bf16 model returns argmax 4539 — a value that does **not** correspond to any position in the model's actual output. The corruption is purely in the logits read-out, not the model compute or the ORT kernels.
### Proposed fix
Treat `BFloat16` the same as `Float16` in both spots, and add a bf16→f32 conversion path. The existing `Cast()` helper already falls back to a CPU bf16→f32 cast; I also added a CUDA `LaunchBf16ToFp32` kernel so the cast stays on-device.
```diff
--- a/src/models/logits.cpp
+++ b/src/models/logits.cpp
- if (type_ == Ort::TypeToTensorType<Ort::Float16_t>)
+ if (type_ == Ort::TypeToTensorType<Ort::Float16_t> || type_ == Ort::TypeToTensorType<Ort::BFloat16_t>)
logits_of_last_token_fp32_ = OrtValue::CreateTensor<float>(...);
...
- // Convert from float16 to float32 if necessary
- if (type_ == Ort::TypeToTensorType<Ort::Float16_t>) {
+ // Convert from float16/bfloat16 to float32 if necessary
+ if (type_ == Ort::TypeToTensorType<Ort::Float16_t> || type_ == Ort::TypeToTensorType<Ort::BFloat16_t>) {
Cast(*logits_of_last_token, logits_of_last_token_fp32_, *model_.p_device_inputs_, Ort::TypeToTensorType<float>);
logits_of_last_token = logits_of_last_token_fp32_.get();
}
```
Plus a CUDA bf16→f32 case (`src/cuda/interface.cpp` `Cast`, `src/cuda/model_kernels.cu` / `kernels.h` `LaunchBf16ToFp32`).
After the fix, the bf16 model's first-token argmax is 496 (matches HF/fp16) and generation is byte-identical to fp16.
### Urgency
Medium — blocks deploying any bf16 model via genai. fp16 is an available workaround but loses bf16's wider dynamic range.
### Platform / Version
- onnxruntime-genai built from source, CUDA EP, Linux.
- Reproduces independent of EP/kernel dispatch (model compute is correct; only the logits read-out is affected).
I have a tested fix and can open a PR if useful.
关闭于 2026-06-18 0 条评论