[Bugfix][Kernel] Fix fp8 dense attention: softmax normalisation and an int32 dim overflow (#1020)

Bug 1: rows lose softmax mass
Fix: scale P up before the cast so it uses e4m3's whole range. l_row scales with it, so nothing changes but the tail that no longer underflows. Free — the scale rides an FMA that already runs, as an optional bias on _scale_sub_score_pair; bf16 never passes one.

Bug 2: int32 overflow on the shape
Fix: one launch per batch entry when the total would overflow. Entries are independent and a leading slice of a contiguous tensor is still contiguous, so no copy and no kernel change.

Co-authored-by: yanyuan.qin@amd.com <yanyuan.qin@amd.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: JohnQinAMD

kernels/attention/flash_attn_fp8_gfx950.py

@@ -43,6 +43,7 @@ def build_flash_attn_dualwave_swp_fp8_module(
     waves_per_eu=2,
     daz=True,
     dualwave_swp_lazy_rescale=True,
+    rescale_threshold=6.0,
     dualwave_swp_setprio=True,
     dualwave_swp_debug_lazy_counts=False,
     dualwave_swp_enable_stagger=True,
@@ -90,6 +91,7 @@ def build_flash_attn_dualwave_swp_fp8_module(
         waves_per_eu=waves_per_eu,
         daz=daz,
         dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale,
+        rescale_threshold=rescale_threshold,
         dualwave_swp_setprio=dualwave_swp_setprio,
         dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts,
         dualwave_swp_enable_stagger=dualwave_swp_enable_stagger,
@@ -219,6 +221,13 @@ def _mask_pair(v_s_a, v_s_b, j):
                 return softmax_helper.causal_mask_pair_if_needed(v_s_a, v_s_b, j)
             return v_s_a, v_s_b
 
+        def _correct_o(v_o, m_row, l_row, m_tile):
+            if const_expr(traits.DUALWAVE_SWP_LAZY_RESCALE):
+                return softmax_helper.lazy_correct_o(v_o, m_row, l_row, m_tile)
+            m_new, corr = softmax_helper.rescale_from_tile_max(m_row, m_tile)
+            softmax_helper.scale_o(v_o, corr)
+            return v_o, m_new, softmax_helper.apply_l_rescale(l_row, corr)
+
         def _merge_tile_max(v_s_a, v_s_b):
             m_tile = softmax_helper.max2(softmax_helper.reduce_max(v_s_a), softmax_helper.reduce_max(v_s_b))
             if const_expr(traits.CAUSAL):
@@ -298,7 +307,7 @@ def _ring_wrap(x):
                 v_s_b = _mask_sub(v_s_b, j + fx.Index(1))
                 v_s_a, v_s_b = _mask_pair(v_s_a, v_s_b, j)
                 m_tile = _merge_tile_max(v_s_a, v_s_b)
-                v_o, m_new, l_row = softmax_helper.lazy_correct_o(v_o, m_row, l_row, m_tile)
+                v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile)
                 v_o = softmax_helper.anchor_v_o(v_o)
                 v_p_a, l_row = _softmax_part(v_s_a, l_row, m_new)
                 _phase_bar()
@@ -326,7 +335,7 @@ def _ring_wrap(x):
                     v_o = _pv_part(v_p_b, v_v_b, v_o)
             else:
                 m_tile = _merge_tile_max(v_s_a, v_s_b)
-                v_o, m_new, l_row = softmax_helper.lazy_correct_o(v_o, m_row, l_row, m_tile)
+                v_o, m_new, l_row = _correct_o(v_o, m_row, l_row, m_tile)
                 v_o = softmax_helper.anchor_v_o(v_o)
 
                 v_o, l_row = _subtile_tail(v_s_a, v_v_a, v_o, l_row, m_new)
@@ -624,6 +633,7 @@ def _compile(
             waves_per_eu=waves_per_eu,
             daz=daz,
             dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale,
+            rescale_threshold=rescale_threshold,
             dualwave_swp_setprio=dualwave_swp_setprio,
             dualwave_swp_debug_lazy_counts=dualwave_swp_debug_lazy_counts,
             dualwave_swp_enable_stagger=dualwave_swp_enable_stagger,

kernels/attention/flash_attn_interface.py

@@ -43,13 +43,24 @@
 
 # Short varlen/paged cases use the lightweight generic path.
 _VARLEN_LIGHT_MAX_SEQ = 256
+# Largest flat element count the fp8 C-ABI can address; see the split below.
+_FP8_MAX_FLAT_ELEMS = 2**31
+# fp8 lifts P by log2(448) - RESCALE_THRESHOLD. Past this KV length enough tiles
+# sit far below the running max that the extra two log2 units matter more than
+# the ~0.3% the lower threshold costs there; below it the two are equally
+# accurate and 6 is cheaper.
+_FP8_LONG_SEQ = 4096
 _DENSE_LIGHT_CU_FALLBACK = 256
 _DENSE_DUALWAVE_MIN_SEQ = 256
 _DENSE_DUALWAVE_LARGE_BATCH = 8
 _DENSE_DUALWAVE_MIN_SEQ_LARGE_BATCH = 192
 _DENSE_M256_MIN_TOKENS = 4096
 
 
+def _fp8_rescale_threshold(seqlen_kv: int) -> float:
+    return 6.0 if seqlen_kv <= _FP8_LONG_SEQ else 4.0
+
+
 def _dtype_str(t: torch.Tensor) -> str:
     s = _DTYPE_MAP.get(t.dtype)
     if s is None:
@@ -173,6 +184,7 @@ def _build_dense_fp8(
     num_heads: int,
     num_kv_heads: int,
     causal: bool,
+    rescale_threshold: float,
     waves_per_eu: int,
     daz: bool,
     lazy_rescale: bool,
@@ -190,6 +202,7 @@ def _build_dense_fp8(
         num_kv_heads=num_kv_heads,
         waves_per_eu=waves_per_eu,
         daz=daz,
+        rescale_threshold=rescale_threshold,
         dualwave_swp_lazy_rescale=lazy_rescale,
         dualwave_swp_setprio=setprio,
         dualwave_swp_enable_stagger=enable_stagger,
@@ -821,6 +834,58 @@ def flydsl_flash_attn_func(
             raise NotImplementedError(f"flydsl_flash_attn_func: {_name} is not supported for fp8")
         if not _t.is_cuda or _t.device != q.device:
             raise ValueError(f"flydsl_flash_attn_func: {_name} must be a CUDA tensor on {q.device}, got {_t.device}")
+
+    # The fp8 path flattens Q/K/V/O to 1-D and the C-ABI packs a dynamic dim as
+    # int32, so a launch aborts once any of them reaches 2**31 (S >= 131072 at
+    # D=128, H=64). K/V are checked too: cross-attention can hold a short Q and
+    # an over-long KV. Batch entries are independent and a leading slice of a
+    # contiguous tensor is still contiguous, so one launch per entry divides the
+    # flat dim by B at no copy. bf16 passes the natural 4-D shape and is exempt.
+    if (
+        dtype_str == "fp8"
+        and not paged_kv
+        and cu_seqlens_q is None
+        and cu_seqlens_kv is None
+        and q.dim() == 4
+        and max(q.numel(), k.numel(), v.numel()) >= _FP8_MAX_FLAT_ELEMS
+    ):
+        if q.shape[0] == 1:
+            # Out of batch to divide by. Launching would abort inside the C ABI
+            # with a struct.error naming neither the tensor nor the limit.
+            raise NotImplementedError(
+                "flydsl_flash_attn_func: fp8 flattens Q/K/V/O and packs the dynamic dim as int32, so a "
+                f"single batch entry cannot exceed {_FP8_MAX_FLAT_ELEMS} elements; got q={q.numel()}, "
+                f"k={k.numel()}, v={v.numel()}. Shorten the sequence or use bf16."
+            )
+        kw = dict(
+            causal=causal,
+            num_kv_heads=num_kv_heads,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_kv=max_seqlen_kv,
+            cross_seqlen=cross_seqlen,
+            kv_cache_layout=kv_cache_layout,
+            num_kv_splits=num_kv_splits,
+            q_descale=q_descale,
+            k_descale=k_descale,
+            v_descale=v_descale,
+            waves_per_eu=waves_per_eu,
+            daz=daz,
+            dualwave_swp_lazy_rescale=dualwave_swp_lazy_rescale,
+            dualwave_swp_setprio=dualwave_swp_setprio,
+            dualwave_swp_enable_stagger=dualwave_swp_enable_stagger,
+            debug_counts=debug_counts,
+            stream=stream,
+        )
+        if out is None:
+            # Allocate once and hand each launch its own slice. Concatenating
+            # afterwards would consume the parts on the ambient stream while the
+            # kernels are still running on `stream`, and would hold two full
+            # outputs at a size where one is already several GB.
+            out = torch.empty(q.shape, dtype=torch.bfloat16 if dtype_str == "fp8" else q.dtype, device=q.device)
+        for i in range(q.shape[0]):
+            sl = slice(i, i + 1)
+            flydsl_flash_attn_func(q[sl].contiguous(), k[sl].contiguous(), v[sl].contiguous(), out=out[sl], **kw)
+        return out
     if has_bias:
         if bias.dtype != q.dtype:
             raise ValueError(f"flydsl_flash_attn_func: bias dtype must match q dtype {q.dtype}, got {bias.dtype}")
@@ -1046,6 +1111,7 @@ def flydsl_flash_attn_func(
                     num_heads=H,
                     num_kv_heads=num_kv_heads,
                     causal=causal,
+                    rescale_threshold=_fp8_rescale_threshold(int(Skv)),
                     waves_per_eu=waves_per_eu,
                     daz=daz,
                     lazy_rescale=dualwave_swp_lazy_rescale,

kernels/attention/flash_attn_utils.py

@@ -411,17 +411,24 @@ def _sub_score_pair(v_s, row_max, fm_fast):
     return Vec.from_elements(lo_sub, fx.Float32).ir_value(), Vec.from_elements(hi_sub, fx.Float32).ir_value()
 
 
-def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast):
+def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast, bias=None):
     """Fused softmax-scale + row-max subtraction (optimization 1-A).
 
-    Returns ``scale * (v_s - row_max_raw)`` per element via a single FMA
-    (``fma(s, scale, -scale*row_max_raw)``), so the fp8 QK MMA can emit raw
+    Returns ``scale * (v_s - row_max_raw) + bias`` per element via a single FMA
+    (``fma(s, scale, bias - scale*row_max_raw)``), so the fp8 QK MMA can emit raw
     (un-scaled) logits and reduce_max can run in the raw domain (scale > 0 is
     order-preserving). Replaces the separate post-QK scale multiply + subtract.
     ``-inf`` masked lanes stay ``-inf`` (scale > 0), matching the un-fused path.
+
+    ``bias`` lands in the FMA's addend, so a caller needing ``exp2`` to produce
+    ``2**bias * P`` pays nothing -- see ``DualwaveFp8SoftmaxHelper.sub_m``.
     """
     s_lo, s_hi = v_s
     neg_scaled_max = zero_f - scale * row_max_raw
+    if bias is not None:
+        # exp2 lands on 2**bias * P instead of P, at no extra instruction: the
+        # FMA's addend absorbs it.
+        neg_scaled_max = neg_scaled_max + bias
     scale_v = Vec.from_elements([scale], fx.Float32).broadcast_to(16)
     nsm_v = Vec.from_elements([neg_scaled_max], fx.Float32).broadcast_to(16)
     lo = fx.fma(Vec(s_lo), scale_v, nsm_v, fastmath=fm_fast)
@@ -1857,6 +1864,7 @@ def cache_tag(self):
             self.WAVES_PER_EU,
             self.DAZ,
             self.DUALWAVE_SWP_LAZY_RESCALE,
+            self.DUALWAVE_SWP_RESCALE_THRESHOLD,
             self.DUALWAVE_SWP_SETPRIO,
             self.DUALWAVE_SWP_DEBUG_LAZY_COUNTS,
             self.DUALWAVE_SWP_ENABLE_STAGGER,
@@ -1885,6 +1893,7 @@ def _make_dualwave_swp_fp8_traits(
     num_heads,
     num_kv_heads,
     head_dim,
+    rescale_threshold,
     causal=True,
     waves_per_eu=2,
     daz=True,
@@ -2035,7 +2044,7 @@ def _make_dualwave_swp_fp8_traits(
         URV_DC_AXIS0_BF=snrpt_bf * vls_bf,
         URV_DC_AXIS1_BF=32,
         URV_I5_BF=d128_bf,
-        DUALWAVE_SWP_RESCALE_THRESHOLD=8.0,
+        DUALWAVE_SWP_RESCALE_THRESHOLD=rescale_threshold,
         SCHED_MFMA_MASK=0x008,
         SCHED_VALU_MASK=0x002,
         SCHED_EXP_MASK=0x400,
@@ -4406,7 +4415,7 @@ def init_types_and_constants(self):
         self.c_neg_inf = fx.Float32(float("-inf"))
         self.c_neg_floor = fx.Float32(-3.0e38)
         self.c_zero_f = fx.Float32(0.0)
-        self.c_eight_f = fx.Float32(traits.DUALWAVE_SWP_RESCALE_THRESHOLD)
+        self.c_rescale_thr_f = fx.Float32(traits.DUALWAVE_SWP_RESCALE_THRESHOLD)
         self.c_zero_v16f32 = Vec.filled(16, 0.0, fx.Float32)
 
     def init_runtime_indices(self):
@@ -5101,8 +5110,24 @@ def max2(self, a, b):
     def floor_masked_max(self, row_max):
         return fx.maxnumf(row_max, self.c_neg_floor)
 
+    # log2 of e4m3's largest finite value, 448.
+    _P_HEADROOM_LOG2 = 8.807354922057604
+
     def sub_m(self, v_s, row_max):
-        return _scale_sub_score_pair(v_s, row_max, self.c_logit_scale, self.c_zero_f, self.fm_fast)
+        # P is cast to e4m3, whose smallest subnormal is 2**-9, so a softmax
+        # over thousands of keys loses its tail to flush-to-zero -- while l_row,
+        # summed before the cast, still counts it. Scaling P up first uses the
+        # format's whole range; l_row scales with it, so the output is unchanged
+        # apart from the tail that survives. Free: it rides the FMA's addend.
+        #
+        # Available headroom is bounded by how large exp2 gets: the lazy path
+        # holds the running max until a tile exceeds it by RESCALE_THRESHOLD, so
+        # exp2 <= 2**THRESHOLD there; the eager path rebases every tile.
+        headroom = self._P_HEADROOM_LOG2
+        if const_expr(self.traits.DUALWAVE_SWP_LAZY_RESCALE):
+            headroom -= self.traits.DUALWAVE_SWP_RESCALE_THRESHOLD
+        bias = fx.Float32(headroom) if headroom > 0.0 else None
+        return _scale_sub_score_pair(v_s, row_max, self.c_logit_scale, self.c_zero_f, self.fm_fast, bias)
 
     def exp2(self, v_s, start, length):
         return _exp2_score_slice(v_s, start, length)
@@ -5180,7 +5205,7 @@ def lazy_rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p):
         def _run(v_o, m_row, l_row, m_tile_max, v_p):
             m_diff = m_tile_max - m_row
             m_diff_scaled = m_diff * self.c_logit_scale
-            below = fx.Float32(m_diff_scaled) <= self.c_eight_f
+            below = fx.Float32(m_diff_scaled) <= self.c_rescale_thr_f
             ballot = rocdl.ballot(T.i64, as_mlir_value(below))
             all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64())
             all_below = llvm.intr_expect(all_below, arith.constant(1, type=ir.IntegerType.get_signless(1)))
@@ -5210,7 +5235,7 @@ def lazy_correct_o(self, v_o, m_row, l_row, m_tile_max):
         def _run(v_o, m_row, l_row, m_tile_max):
             m_diff = m_tile_max - m_row
             m_diff_scaled = m_diff * self.c_logit_scale
-            below = fx.Float32(m_diff_scaled) <= self.c_eight_f
+            below = fx.Float32(m_diff_scaled) <= self.c_rescale_thr_f
             ballot = rocdl.ballot(T.i64, as_mlir_value(below))
             all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64())
             all_below = llvm.intr_expect(all_below, arith.constant(1, type=ir.IntegerType.get_signless(1)))