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]