fix(rocke): carry the raw running max in unified attention

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py

@@ -6447,21 +6447,18 @@ def _scalar_band_ok(
     context_len,
     query_pos,
     kpos,
-    *,
-    window: bool = True,
 ) -> Optional[Value]:
     """Visibility of key ``kpos`` for the query: the right bound ``k <= d + right``
     (``d = context_len + query_pos``) AND the left bound ``d - k < sliding_window``.
     ``None`` when neither bounds anything (the loop range already limits to kv_len).
-    ``window=False`` skips the left bound.
     """
     ok = None
     if p.right_bound >= 0:
         upper = b.add(context_len, query_pos)
         if p.right_bound > 0:
             upper = b.add(upper, b.const_i32(p.right_bound))
         ok = b.cmp_le(kpos, upper)
-    if window and p.sliding_window > 0:
+    if p.sliding_window > 0:
         dist = b.sub(b.add(context_len, query_pos), kpos)
         sw_ok = b.cmp_lt(dist, b.const_i32(p.sliding_window))
         ok = sw_ok if ok is None else b.land(ok, sw_ok)
@@ -6637,9 +6634,11 @@ def build_unified_attention_2d(
         if band_ok is not None:
             score = b.select(band_ok, score, neg_inf)
         new_m_raw = b.fmax(m_val, score)
-        # If both running max and current score are -inf, the row is fully
-        # masked; force m to 0 so the resulting alpha/prob are 0 instead of NaN
-        # (matches Triton's `m_j = tl.where(m_j > -inf, m_j, 0.0)`).
+        # If both running max and current score are -inf, the row has seen no key
+        # yet; use 0 as the exponent reference so alpha/prob are 0 instead of NaN
+        # (matches Triton's `m_j = tl.where(m_j > -inf, m_j, 0.0)`). Carry the raw
+        # max: carrying the 0 pins m at 0 past masked leading keys (sliding
+        # window), and exp2(s - 0) then underflows P for very negative scores.
         is_finite = b.fcmp("ogt", new_m_raw, neg_inf)
         new_m = b.select(is_finite, new_m_raw, zero_f)
         alpha = b.exp2(b.fsub(m_val, new_m))
@@ -6657,14 +6656,14 @@ def build_unified_attention_2d(
         if p.use_fp8:
             vv = b.fmul(vv, v_scale)
         new_acc = b.fadd(b.fmul(acc_val, alpha), b.fmul(prob, vv))
-        b.scf_yield(new_m, new_l, new_acc)
+        b.scf_yield(new_m_raw, new_l, new_acc)
 
     out_val = b.fmul(loop.results[2], b.rcp(loop.results[1]))
-    if not p.default_mask:
-        # Rows can see no key at all (e.g. top-left rows past ``kv_len`` with a window
-        # or a right bound): emit 0 instead of acc * rcp(0) == NaN, like the tiled
-        # kernels. The default mask never produces such a row, so its IR is unchanged.
-        out_val = b.select(b.fcmp("ogt", loop.results[1], zero_f), out_val, zero_f)
+    # Rows can see no key at all under any mask (bottom-right causal with Sq > Skv,
+    # or a window / right bound): store 0 instead of acc * rcp(0) == NaN. This is
+    # the tiled 2D epilogue's rule: only an exact-zero sum counts as keyless, so a
+    # NaN ``l`` still reaches O.
+    out_val = b.select(b.fcmp("oeq", loop.results[1], zero_f), zero_f, out_val)
     out_cast = b.cast_f32_to(out_val, dtype)
     out_off, _ = q_desc.offset(b, token=q_tok, head=q_head, dim=dim)
     valid = b.land(active, b.cmp_lt(dim, b.const_i32(p.head_size)))
@@ -6693,9 +6692,7 @@ def kernel_name(self) -> str:
             flags={
                 "tl": p.causal_top_left,
                 **_right_bound_flag(p),
-                # The default-mask 3D scalar kernel never applied the window (its name
-                # and IR are pinned); every other mask does.
-                "sw": p.sliding_window > 0 and not p.default_mask,
+                "sw": p.sliding_window > 0,
             },
         )
 
@@ -6801,28 +6798,28 @@ def build_unified_attention_3d(
             scale,
             rcp_ln2,
         )
-        # The scalar 3D segment kernel has never applied the left bound; keep that
-        # for the default mask (its IR is pinned) and honour it for every other mask.
-        band_ok = _scalar_band_ok(
-            b, p, context_len, query_pos, kpos, window=not p.default_mask
-        )
+        # Every mask, the default one included: the right (causal) bound, and
+        # the left (sliding-window) bound whenever sliding_window > 0.
+        band_ok = _scalar_band_ok(b, p, context_len, query_pos, kpos)
         if band_ok is not None:
             score = b.select(band_ok, score, neg_inf)
-        new_m = b.fmax(m_val, score)
-        if not p.default_mask:
-            # Non-default masks can leave whole segments fully masked (top-left
-            # decode sees only key 0; a window or right bound skips others). Force m
-            # to 0 there so alpha/prob are 0 rather than NaN, which the reduce would
-            # propagate as NaN * 0. The default mask never hits this: IR unchanged.
-            new_m = b.select(b.fcmp("ogt", new_m, neg_inf), new_m, zero_f)
+        new_m_raw = b.fmax(m_val, score)
+        # A segment, or a segment's leading keys, can be fully masked for a row:
+        # a default-mask prefill row whose causal limit ends before a trailing
+        # segment, top-left decode (only key 0 visible), a window or right bound.
+        # Use 0 as the exponent reference there so alpha/prob are 0 rather than
+        # exp2(-inf - -inf) = NaN, but carry and store the raw max: it stays -inf
+        # until a key is seen, so the reduce ignores an empty segment and a
+        # masked prefix cannot pin m at 0.
+        new_m = b.select(b.fcmp("ogt", new_m_raw, neg_inf), new_m_raw, zero_f)
         alpha = b.exp2(b.fsub(m_val, new_m))
         prob = b.exp2(b.fsub(score, new_m))
         new_l = b.fadd(b.fmul(l_val, alpha), prob)
         vv = _emit_v_load(
             b, p, dtype, value, block_tables, bt_stride, seq_idx, kv_head, kpos, dim
         )
         new_acc = b.fadd(b.fmul(acc_val, alpha), b.fmul(prob, vv))
-        b.scf_yield(new_m, new_l, new_acc)
+        b.scf_yield(new_m_raw, new_l, new_acc)
 
     ml_desc, out_desc = _segm_descriptors(p, spec.num_segments)
     base, _ = out_desc.offset(b, token=q_tok, head=q_head, seg=segm_idx, dim=dim)
@@ -6907,11 +6904,20 @@ def build_unified_attention_reduce(
         [("mx", neg_inf)],
         iv_name="seg",
     )
+    # A segment that saw no key stores max = -inf with sum = 0: the segment kernel
+    # carries the raw max, so the max pass ignores that segment as is, and the
+    # combine pass weights it by exp2(-inf - ref) = 0 whenever ref is finite. ref
+    # is the overall max, or 0 when every segment is empty (overall = -inf), the
+    # same exponent-reference clamp the segment kernel uses: a row with no key at
+    # all then has den = 0 and gets O = 0 instead of NaN from exp2(-inf - -inf) or
+    # acc * rcp(0). Only an exact zero den counts as empty, so a NaN sum or
+    # denominator still yields NaN.
     with max_loop as (seg, (mx,)):
         idx, _ = ml_desc.offset(b, token=q_tok, head=q_head, seg=seg)
         mv = b.global_load_f32(segm_max, idx)
         b.scf_yield(b.fmax(mx, mv))
     overall = max_loop.results[0]
+    ref = b.select(b.fcmp("ogt", overall, neg_inf), overall, zero)
     red = b.scf_for_iter(
         b.const_i32(0),
         b.const_i32(spec.num_segments),
@@ -6923,7 +6929,7 @@ def build_unified_attention_reduce(
         idx, _ = ml_desc.offset(b, token=q_tok, head=q_head, seg=seg)
         mv = b.global_load_f32(segm_max, idx)
         lv = b.global_load_f32(segm_expsum, idx)
-        factor = b.exp2(b.fsub(mv, overall))
+        factor = b.exp2(b.fsub(mv, ref))
         den2 = b.fadd(den, b.fmul(lv, factor))
         out_idx, _ = seg_out_desc.offset(
             b,
@@ -6936,6 +6942,7 @@ def build_unified_attention_reduce(
         acc2 = b.fadd(acc, b.fmul(ov, factor))
         b.scf_yield(den2, acc2)
     result = b.fmul(red.results[1], b.rcp(red.results[0]))
+    result = b.select(b.fcmp("oeq", red.results[0], zero), zero, result)
     cast = b.cast_f32_to(result, dtype)
     out_idx, _ = q_desc.offset(b, token=q_tok, head=q_head, dim=dim)
     with b.scf_if(active):

dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py

@@ -4564,6 +4564,7 @@ def _kslot(group_idx: int) -> int:
                 # full ``l_new = alpha * l_old + l_sum`` updates are applied
                 # by that shared downstream block, not here.
                 m_new = [st_m_new for _ in range(SOFTMAX_STATE_SLOTS)]
+                m_carry = [st_m_raw for _ in range(SOFTMAX_STATE_SLOTS)]
                 l_local = [st_l_sum for _ in range(SOFTMAX_STATE_SLOTS)]
                 # In the transposed orientation we already have masked /
                 # scaled / softcapped scores in ``st_scores`` and softmax
@@ -4732,6 +4733,7 @@ def _kslot(group_idx: int) -> int:
                     masked[(n, reg)] = score
 
             m_new = []
+            m_carry = []
             s_local = {}
             for reg in range(REGS_PER_LANE):
                 local_max = neg_inf
@@ -4743,6 +4745,7 @@ def _kslot(group_idx: int) -> int:
                 full_max_raw = b.fmax(m_vals[reg], tile_max)
                 ok = b.fcmp("ogt", full_max_raw, neg_inf)
                 m_new.append(b.select(ok, full_max_raw, zero_f))
+                m_carry.append(full_max_raw)
 
             l_local = []
             for reg in range(REGS_PER_LANE):
@@ -4821,6 +4824,7 @@ def _kslot(group_idx: int) -> int:
             # repeated for every N-tile. Local lane-max across N-tiles, then 4-stage
             # XOR butterfly across the 16 lanes in the row-group.
             m_new = []
+            m_carry = []
             s_local = {}  # (reg, n) -> the lane's masked score (still owned per-lane)
             for reg in range(REGS_PER_LANE):
                 local_max = neg_inf
@@ -4833,6 +4837,7 @@ def _kslot(group_idx: int) -> int:
                 full_max_raw = b.fmax(m_vals[reg], tile_max)
                 ok = b.fcmp("ogt", full_max_raw, neg_inf)
                 m_new.append(b.select(ok, full_max_raw, zero_f))
+                m_carry.append(full_max_raw)
 
             # ---- compute P = exp2(S - m_new) and l_local = sum(P) per row ----
             # P_lds[row, col] = exp2(S[row, col] - m_new[row]). Each warp publishes
@@ -5431,9 +5436,14 @@ def _strided_v_b_operand(k_iter: int) -> Value:
             for atom in range(M_ATOMS_PER_WARP):
                 new_acc[n * M_ATOMS_PER_WARP + atom] = acc_per_atom[atom]
 
+        # Carry the unclamped running max: it stays -inf until the row has seen a
+        # key. ``m_new`` (-inf replaced by 0) is only the exponent reference for
+        # P and alpha. Carrying the clamped 0 instead pins m at 0 when a row's
+        # first tile is fully masked (sliding window), and exp2(s - 0) then
+        # underflows P for very negative scores so the row's O collapses.
         yields = []
         for r in range(SOFTMAX_STATE_SLOTS):
-            yields.append(m_new[r])
+            yields.append(m_carry[r])
             yields.append(new_l_vals[r])
         for n in range(ACC_N_TILES):
             for atom in range(ACC_M_ATOMS):
@@ -5505,9 +5515,10 @@ def _strided_v_b_operand(k_iter: int) -> Value:
     def _acc_final_get(n: int, atom: int) -> Value:
         return acc_final[n * ACC_M_ATOMS + atom]
 
-    # Per-row reciprocal of L (computed once, reused across stripes).
+    # Per-row reciprocal of L (computed once, reused across stripes). Only an
+    # exact-zero L marks a keyless row (stored as 0); a NaN L stays NaN in O.
     rcp_l = [b.rcp(l_final[r]) for r in range(SOFTMAX_STATE_SLOTS)]
-    l_nonzero = [b.fcmp("ogt", l_final[r], zero_f) for r in range(SOFTMAX_STATE_SLOTS)]
+    l_zero = [b.fcmp("oeq", l_final[r], zero_f) for r in range(SOFTMAX_STATE_SLOTS)]
 
     if USE_MFMA_32X32:
         if TRANSPOSED_QK_32X32:
@@ -5542,7 +5553,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                 dim=b.const_i32(0),
             )
             inv_l_t = b.rcp(l_final[0])
-            l_nonzero_t = b.fcmp("ogt", l_final[0], zero_f)
+            l_zero_t = b.fcmp("oeq", l_final[0], zero_f)
             for n in range(ACC_N_TILES):
                 acc32 = _acc_final_get(n, 0)
                 for reg in range(REGS_PER_LANE):
@@ -5552,7 +5563,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                     v = b.vec_extract(acc32, reg)
                     normalized = b.fmul(v, inv_l_t)
                     final_h = b.cast_f32_to(
-                        b.select(l_nonzero_t, normalized, zero_f), dtype
+                        b.select(l_zero_t, zero_f, normalized), dtype
                     )
                     with b.scf_if(op_mask_t):
                         b.global_store(
@@ -5602,7 +5613,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                 v = b.vec_extract(acc32, reg)
                 normalized = b.fmul(v, rcp_l[reg])
                 final_h = b.cast_f32_to(
-                    b.select(l_nonzero[reg], normalized, zero_f), dtype
+                    b.select(l_zero[reg], zero_f, normalized), dtype
                 )
                 b.smem_store_vN(Acc_lds, [row, col_in_stripe], final_h, 1)
             b.sync()
@@ -5712,7 +5723,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                 v = b.vec_extract(_acc_final_get(n, atom), in_atom)
                 normalized = b.fmul(v, rcp_l[reg])
                 final_h = b.cast_f32_to(
-                    b.select(l_nonzero[reg], normalized, zero_f), dtype
+                    b.select(l_zero[reg], zero_f, normalized), dtype
                 )
                 b.smem_store_vN(Acc_lds, [row, col_in_stripe], final_h, 1)
         b.sync()

dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py

@@ -3656,6 +3656,7 @@ def _mod3_add(x, c):
                 # full ``l_new = alpha * l_old + l_sum`` updates are applied
                 # by that shared downstream block, not here.
                 m_new = [st_m_new for _ in range(SOFTMAX_STATE_SLOTS)]
+                m_carry = [st_m_raw for _ in range(SOFTMAX_STATE_SLOTS)]
                 l_local = [st_l_sum for _ in range(SOFTMAX_STATE_SLOTS)]
                 # In the transposed orientation we already have masked /
                 # scaled / softcapped scores in ``st_scores`` and softmax
@@ -3849,6 +3850,7 @@ def _mod3_add(x, c):
                     masked[(n, reg)] = score
 
             m_new = []
+            m_carry = []
             s_local = {}
             for reg in range(REGS_PER_LANE):
                 local_max = neg_inf
@@ -3860,6 +3862,7 @@ def _mod3_add(x, c):
                 full_max_raw = b.fmax(m_vals[reg], tile_max)
                 ok = b.fcmp("ogt", full_max_raw, neg_inf)
                 m_new.append(b.select(ok, full_max_raw, zero_f))
+                m_carry.append(full_max_raw)
 
             l_local = []
             for reg in range(REGS_PER_LANE):
@@ -3952,6 +3955,7 @@ def _mod3_add(x, c):
             # repeated for every N-tile. Local lane-max across N-tiles, then 4-stage
             # XOR butterfly across the 16 lanes in the row-group.
             m_new = []
+            m_carry = []
             s_local = {}  # (reg, n) -> the lane's masked score (still owned per-lane)
             for reg in range(REGS_PER_LANE):
                 local_max = neg_inf
@@ -3964,6 +3968,7 @@ def _mod3_add(x, c):
                 full_max_raw = b.fmax(m_vals[reg], tile_max)
                 ok = b.fcmp("ogt", full_max_raw, neg_inf)
                 m_new.append(b.select(ok, full_max_raw, zero_f))
+                m_carry.append(full_max_raw)
 
             # ---- compute P = exp2(S - m_new) and l_local = sum(P) per row ----
             # P_lds[row, col] = exp2(S[row, col] - m_new[row]). Each warp publishes
@@ -4423,9 +4428,14 @@ def _apply_transposed_pv_regs(acc32: Value, n: int, p_regs) -> Value:
             for atom in range(M_ATOMS_PER_WARP):
                 new_acc[n * M_ATOMS_PER_WARP + atom] = acc_per_atom[atom]
 
+        # Carry the unclamped running max: it stays -inf until the row has seen a
+        # key. ``m_new`` (-inf replaced by 0) is only the exponent reference for
+        # P and alpha. Carrying the clamped 0 instead pins m at 0 when a row's
+        # first tile is fully masked (sliding window), and exp2(s - 0) then
+        # underflows P for very negative scores so the row's O collapses.
         yields = []
         for r in range(SOFTMAX_STATE_SLOTS):
-            yields.append(m_new[r])
+            yields.append(m_carry[r])
             yields.append(new_l_vals[r])
         for n in range(ACC_N_TILES):
             for atom in range(ACC_M_ATOMS):
@@ -4522,9 +4532,10 @@ def _apply_transposed_pv_regs(acc32: Value, n: int, p_regs) -> Value:
     def _acc_final_get(n: int, atom: int) -> Value:
         return acc_final[n * ACC_M_ATOMS + atom]
 
-    # Per-row reciprocal of L (computed once, reused across stripes).
+    # Per-row reciprocal of L (computed once, reused across stripes). Only an
+    # exact-zero L marks a keyless row (stored as 0); a NaN L stays NaN in O.
     rcp_l = [b.rcp(l_final[r]) for r in range(SOFTMAX_STATE_SLOTS)]
-    l_nonzero = [b.fcmp("ogt", l_final[r], zero_f) for r in range(SOFTMAX_STATE_SLOTS)]
+    l_zero = [b.fcmp("oeq", l_final[r], zero_f) for r in range(SOFTMAX_STATE_SLOTS)]
 
     if USE_MFMA_32X32:
         if TRANSPOSED_QK_32X32:
@@ -4559,7 +4570,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                 dim=b.const_i32(0),
             )
             inv_l_t = b.rcp(l_final[0])
-            l_nonzero_t = b.fcmp("ogt", l_final[0], zero_f)
+            l_zero_t = b.fcmp("oeq", l_final[0], zero_f)
             for n in range(ACC_N_TILES):
                 acc32 = _acc_final_get(n, 0)
                 for reg in range(REGS_PER_LANE):
@@ -4569,7 +4580,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                     v = b.vec_extract(acc32, reg)
                     normalized = b.fmul(v, inv_l_t)
                     final_h = b.cast_f32_to(
-                        b.select(l_nonzero_t, normalized, zero_f), dtype
+                        b.select(l_zero_t, zero_f, normalized), dtype
                     )
                     with b.scf_if(op_mask_t):
                         b.global_store(
@@ -4619,7 +4630,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                 v = b.vec_extract(acc32, reg)
                 normalized = b.fmul(v, rcp_l[reg])
                 final_h = b.cast_f32_to(
-                    b.select(l_nonzero[reg], normalized, zero_f), dtype
+                    b.select(l_zero[reg], zero_f, normalized), dtype
                 )
                 b.smem_store_vN(Acc_lds, [row, col_in_stripe], final_h, 1)
             b.sync()
@@ -4729,7 +4740,7 @@ def _acc_final_get(n: int, atom: int) -> Value:
                 v = b.vec_extract(_acc_final_get(n, atom), in_atom)
                 normalized = b.fmul(v, rcp_l[reg])
                 final_h = b.cast_f32_to(
-                    b.select(l_nonzero[reg], normalized, zero_f), dtype
+                    b.select(l_zero[reg], zero_f, normalized), dtype
                 )
                 b.smem_store_vN(Acc_lds, [row, col_in_stripe], final_h, 1)
         b.sync()