From 014d647ba30703ccaaf29eb48425f659301ea070 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 26 May 2026 15:03:49 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20sink=20bias=20domain=20correction=20?= =?UTF-8?q?=E2=80=94=20add=20attn=5Fsink/scale=20to=20raw=20logits?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The softmax scales by scale_log2 = scale * log2(e). Adding sink_val to raw logits causes it to be scaled too. Fix: add sink_val/scale instead, so after scaling: (sink_val/scale) * scale_log2 = sink_val * log2(e). This correctly multiplies attention weights by exp(sink_val). --- dsv4/kernels/attention/fmha.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 0992413e..8a0676db 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -434,9 +434,12 @@ class FmhaKernel: kt_offset = Int32(kt * 128) # KV position offset for this tile # D5c: Read sink bias once (same for all positions in this head). # Define unconditionally for CuTeDSL scoping (used when n_comp > 0). + # The bias must be added in the SCALED-LOG2 domain: attn_sink * log2(e). + # But we add to the RAW logits before the scale_log2 multiply. + # Raw correction: attn_sink / scale → after * scale_log2 → attn_sink * log2(e) sink_val = Float32(0.0) if const_expr(self.n_comp > 0): - sink_val = mSinkBias[Int32(0)] + sink_val = mSinkBias[Int32(0)] / Float32(self.scale_softmax) for j0 in range(32): for j1 in range(4): coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0]