antirez / antirez/h3.c

Qwen causal GQA rounds scaled queries to BF16 before the QK contraction

Aperta
#3 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
C
Stelle
2.7k
Fork
210
Metriche di merge delle PR
Nessuna PR unita negli ultimi 30g

Descrizione

## 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.

Guida per i contributori

Nessuna guida per i contributori indicizzata per questo repository

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.