Qwen causal GQA rounds scaled queries to BF16 before the QK contraction
- Lenguaje dominante
- C
- Estrellas
- 2.7k
- Forks
- 210
- Métricas de merge de PR
- Sin PR fusionados en 30 d
Descripción
## Problem
`h3_gqa_causal_bf16` scales each BF16 query value in F32, but then rounds the product back to BF16 before the F32 QK contraction:
```metal
shared_query[d] = h3_bf16_to_f32(h3_f32_to_bf16(
h3_bf16_to_f32(query[q_base + d]) * args.scale));
```
`shared_query` is a `threadgroup float` array and the dot product uses F32 `fma`, so this intermediate BF16 round loses precision without reducing storage or bandwidth.
```text
Current: BF16 Q -> F32 scale -> BF16 round -> F32 QK contraction
Expected: BF16 Q -> F32 scale --------------> F32 QK contraction
```
The lost bits perturb the attention scores before softmax. This path is used by the Qwen3-VL prompt encoder's causal GQA kernel; the DiT attention path uses MPSGraph SDPA and is not affected by this specific expression.
## Reproduction and evidence
I exercised the real Metal kernel on an Apple M5 with deterministic synthetic BF16 inputs at the production head width:
- sequence: 32
- query heads: 8
- KV heads: 2
- head dimension: 128
- scale: `1 / sqrt(128)`
- 64 trials, 2,097,152 output elements
Each GPU result was compared at the kernel's BF16 output boundary with an explicit CPU attention reference that keeps the scaled query in F32.
| Metric versus F32 reference | Current BF16-rounded Q | F32-scaled Q | Change |
|---|---:|---:|---:|
| Maximum absolute error | 0.0924071 | 0.00390625 | 95.77% lower |
| RMSE | 0.00194536 | 0.000868307 | 55.37% lower |
| BF16 output mismatch rate | 21.6468% | 0.007486% | 99.97% lower |
The two variants differed in 453,957 of 2,097,152 outputs (21.65%). Both were bit-for-bit deterministic across repeated runs.
A smaller eight-trial regression configuration also distinguishes the behavior reliably:
- F32-scaled implementation: max error 0.00390613, RMSE 0.000864323, mismatch rate 0.00839% — passes
- Current BF16 re-rounding: max error 0.0924071, RMSE 0.00197766, mismatch rate 21.5992% — fails
Repeated GPU timestamp measurements found no performance regression: median dispatch time was 21.274 us with the BF16 re-round and 20.491 us while retaining F32. The distributions overlap, so this should be read as precision improvement at neutral measured cost rather than a speedup claim.
## Expected behavior
Keep `bf16_to_f32(query) * scale` in F32 through the QK contraction, and cover the production-width kernel with a fixture-free numerical regression test.
## Scope
This confirms a material kernel-level numerical divergence from explicit F32 attention mathematics. Downstream prompt-encoder effects are outside the scope of this report.
Guía de contribución
No hay ninguna guía de contribución indexada para este repositorio
Evaluación
Este issue todavía no se ha evaluado.