return_attention on CUDA float16: scores are rounded to fp16 before the softmax, so close candidates tie #2132

Follow-up to #2122 issue, as agreed there. #2131 fixes the NaN error: the score GEMM now accumulates in fp32. This issue is about what is left after that fix.

Summary

With compute_type="float16" or "int8_float16" on CUDA, some rows of the matrix returned by return_attention=True have an exact tie at the top: the probability is split evenly over 2–8 source tokens (0.50 / 0.33 / 0.25 / 0.20 / 0.17 / 0.12). In CUDA float32 the same rows have a clear winner, for example 0.998 against 0.002. Rows that are not tied almost always match float32.

Same model and call as in #2122 (step_30000_model, the 22-row example sentence, reference target fed back with target_prefix), cp312 wheel from the #2131 CI run:

On the RTX 4060 every row that disagrees with float32 is a tied row, at beams 1, 2 and 4 and for both compute types. On the A4500 the counts agree with that: there are never more disagreeing rows than tied ones. Any argmax over such a row breaks the tie by index, so which source token "wins" is arbitrary. float16 keeps the weights in fp16 with no int8 quantization and loses as many rows as int8_float16, so this is not a quantization effect.

At scale

The same model on the English side of the public NTREX-128 en→my test set (newstest2019-src.eng, 1997 sentences, 46,756 target rows), beam 2, #2131 wheel, RTX 4060. CUDA float32 translates, and its output is fed back as target_prefix to the other compute types:

• int8_float32 is the control: int8 weights, float32 scores. It has no ties and agrees on 95.8% of rows, so int8 quantization alone costs about 4 points.
• With float16 scores, almost a third of the rows end in a tie, and nearly every row that disagrees with float32 is one of them.
• Copy-anchor accuracy is the share of target pieces copied verbatim from the source (numbers, Latin words) whose argmax lands on that fragment in the source. It drops less (98.0% → 96.2–96.6%) than the overall agreement: copied numbers and names are the easiest links, so fewer of them end in a tie, but they are not immune (▁SIX below).
• The implied score magnitude over the 14,223 tied float16 rows: median above 2582, 90th percentile above 6461.

Cause

In dot_product_attention the scores are written to output, which has the dtype of the queries:

const ops::MatMul keys_matmul(/*trans_a=*/false, /*trans_b=*/true, queries_scale);
keys_matmul(queries, keys, output);   // output: float16
...
ops::SoftMax()(attn, values_lengths, attn);

After #2131 the products are accumulated in fp32 and queries_scale is applied in fp32, but the result is still rounded to float16 when the GEMM writes it. The softmax kernel accumulates in float internally (AccumT = float in softmax_gpu.cu), but its input is already rounded.

For a sharp, supervised alignment head the scaled scores are in the thousands. The fp16 spacing there is 1 at 1024–2048, 2 at 2048–4096, 4 at 4096–8192 and 8 at 8192–16384. Two candidates whose scores differ by less than that become equal, and the softmax splits the mass between them.

The returned matrix itself gives a lower bound on the magnitude. Two values that round to the same fp16 number differ by less than one ulp, and ulp(x) <= |x| / 1024, so a tie means |score| > 1024 * gap. The gap is the score difference of the tied pair, which float32 gives exactly as ln(p_i / p_j). For the example sentence (float16, beam 1, RTX 4060, #2131 wheel):

The bounds of 870–7700 are consistent with each other and with #2122. Unscaled, these scores are 8× larger (dim_per_head 64), up to about 62,000, right below the fp16 limit of 65,504 that the accumulator crossed before #2131.

The same matrix feeds attn · V, so the rounding also reaches the decoder. In this example the translations still match float32, so the visible damage is to the alignment.

return_attention is the only way to get word alignments out of CTranslate2, and they are used to carry tags, placeholders, numbers and named entities into the translation. A tie silently replaces a confident link (▁SIX 0.998) with a coin flip between two or more source tokens. Nothing in the output marks the row as degraded. Sharp heads are exactly what guided alignment training produces, so the better the alignment head, the more rows are affected.

Environment

• ctranslate2 4.8.4 + Fix FP16 overflow in CUDA attention score accumulation #2131 (cp312 Linux wheel from CI run 40276546050)
• RTX A4500, driver 550.127.05, CUDA 12.9
• RTX 4060 Laptop, drivers 580.178.04 and 580.196.07 (same results), cuBLAS 12.9.2
• Model: step_30000_model from mcialhn/ct2-alignment-test-model. Seq2seq transformer, relative positions, guided alignment on decoder layer -2 head 0, int8_float16 export.

Reproduction

import ctranslate2
src = "▁THOMAS ▁PRESENTED ▁BLUE ▁PETER ▁FOR ▁SIX ▁YEARS ▁AND ▁LEFT ▁FOR ▁SKY ▁SPORTS ▁IN ▁ 2 0 0 5 .".split()
tgt = "▁သောမတ် စ်သည် ▁BL UE ▁P ETER ▁ကို ▁ခြောက် နှစ်ကြာ ▁ကမ်းလှမ်း ခဲ့ပြီး ▁၂၀၀၅ ▁ခုနှစ်တွင် ▁SK Y ▁SP OR TS ▁သို့ ▁ထွက်ခွာ သွားခဲ့သည် ။".split()
for compute_type in ("float32", "float16"):
    tr = ctranslate2.Translator("step_30000_model", device="cuda", compute_type=compute_type)
    att = tr.translate_batch([src], beam_size=1, target_prefix=[tgt], return_attention=True)[0].attention[0]
    print(compute_type, [round(p, 4) for p in att[7]])   # row of "▁ခြောက်" (six)
# float32: ▁SIX 0.9982, ▁YEARS 0.0017   float16: ▁SIX 0.4993, ▁YEARS 0.4993

The attached repro_ct2_attention_fp16_pr2131.py runs float32 / float16 / int8_float16 at beams 1, 2 and 4 in separate processes and prints the summary table. --rows prints the per-row table above. The attached repro_ct2_attention_ties.py repro_ct2_attention_ties.py <model> <source.txt> --beam 2 produces the NTREX table.

Suggested fix

Keep the score GEMM output in fp32 and run the softmax on fp32 input, then cast the probabilities to the activation dtype for attn · V. The probabilities themselves are in [0, 1] and survive the cast to fp16; it is the scores that need the range.

With half inputs, cublasGemmEx can write a CUDA_R_32F C matrix with CUBLAS_COMPUTE_32F, so this is a change of output type, not of the GEMM itself. The extra cost is one fp32 batch × heads × query_len × key_len buffer per attention call. During incremental decoding query_len is 1.

If that cost matters, the fp32 path could be limited to the alignment layer when return_attention=True. The rounding then still reaches the decoder through attn · V, but the returned alignments are exact.