antirez / antirez/h3.c

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

Ouverte
#3 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
Langage dominant
C
Étoiles
2.7k
Forks
210
Métriques de merge des PR
Aucune PR mergée en 30 j

Description

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

Guide de contribution

Aucun guide de contribution indexé pour ce dépôt

Évaluation

Cette issue n'a pas encore été évaluée.

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.