[FlyDSL] Support runtime softmax scale in gfx950 FP8 FMHA

aiter/ops/flydsl/fmha_kernels.py

@@ -28,7 +28,6 @@
 
 from __future__ import annotations
 
-import math
 from functools import lru_cache
 
 import torch
@@ -253,14 +252,15 @@ def _fp8_gfx950_supported(
 ) -> bool:
     """Gate for the gfx950 fp8 kernel.
 
-    It needs e4m3fn Q/K/V with per-tensor descales, hard-wires ``1/sqrt(D)``,
-    and writes bf16. Reject anything else so it falls through rather than
+    It needs e4m3fn Q/K/V with per-tensor descales and a positive, finite
+    softmax scale, and writes bf16. Reject anything else so it falls through rather than
     silently dropping the feature.
     """
     if q.dtype is not torch.float8_e4m3fn or q_descale is None:
         return False
 
     from ...jit.utils.chip_info import get_gfx
+    from .kernels.flash_attn_func_fp8_gfx950 import _is_valid_softmax_scale
 
     if get_gfx() != "gfx950":
         return False
@@ -280,9 +280,7 @@ def _fp8_gfx950_supported(
     ):
         return False
     qk_hdim = q.shape[-1]
-    if softmax_scale is not None and not math.isclose(
-        softmax_scale, 1.0 / math.sqrt(qk_hdim), rel_tol=1e-6
-    ):
+    if not _is_valid_softmax_scale(softmax_scale):
         return False
     if not _fp8_gfx950_buildable(qk_hdim, v.shape[-1]):
         return False
@@ -361,6 +359,7 @@ def flydsl_flash_attn_varlen_func(
             q,
             k,
             v,
+            softmax_scale=softmax_scale,
             causal=causal,
             cu_seqlens_q=cu_seqlens_q,
             cu_seqlens_kv=cu_seqlens_k,
@@ -489,6 +488,7 @@ def flydsl_flash_attn_batch_func(
             q,
             k,
             v,
+            softmax_scale=softmax_scale,
             causal=causal,
             q_descale=q_descale,
             k_descale=k_descale,

aiter/ops/flydsl/kernels/flash_attn_func_fp8_gfx950.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import functools
+import math
 
 import torch
 
@@ -48,6 +49,11 @@
 _FP8_BATCH_INTERLEAVE_GROUP = 2
 
 
+def _is_valid_softmax_scale(softmax_scale: float | None) -> bool:
+    """Accept the default scale or a positive, finite custom scale."""
+    return softmax_scale is None or (math.isfinite(softmax_scale) and softmax_scale > 0)
+
+
 def _fp8_rescale_threshold(seqlen_kv: int) -> float:
     return 6.0 if seqlen_kv <= _FP8_LONG_SEQ else 4.0
 
@@ -203,6 +209,7 @@ def flydsl_flash_attn_fp8_func(
     k: torch.Tensor,
     v: torch.Tensor,
     *,
+    softmax_scale: float | None = None,
     causal: bool = True,
     num_kv_heads: int | None = None,
     cu_seqlens_q: torch.Tensor | None = None,
@@ -231,6 +238,8 @@ def flydsl_flash_attn_fp8_func(
            Dense: ``[B, Sq, H, D]`` (BSHD). Varlen: ``[total_q, H, D]`` (packed).
         k: Key tensor. Dense: ``[B, Skv, Hkv, D]``. Varlen: ``[total_kv, Hkv, D]``.
         v: Value tensor, same shape as k except the last dim may be ``Dv != D``.
+        softmax_scale: Positive, finite scale applied to QK logits, independent
+            of the Q/K descales. Defaults to ``1 / sqrt(q.shape[-1])``.
         causal: Bottom-right aligned causal mask when True.
         num_kv_heads: KV head count for GQA/MQA; defaults to k's head count.
         cu_seqlens_q / cu_seqlens_kv: Int32 ``[B+1]`` cumulative token counts (varlen).
@@ -291,6 +300,7 @@ def flydsl_flash_attn_fp8_func(
                 "or use bf16."
             )
         kw = {
+            "softmax_scale": softmax_scale,
             "causal": causal,
             "num_kv_heads": num_kv_heads,
             "max_seqlen_q": max_seqlen_q,
@@ -409,6 +419,13 @@ def flydsl_flash_attn_fp8_func(
             f"flydsl_flash_attn_fp8_func: head_dim ({D}) must be >= 64 and a multiple of 32"
         )
 
+    if softmax_scale is None:
+        softmax_scale = D**-0.5
+    if not _is_valid_softmax_scale(softmax_scale):
+        raise ValueError(
+            "flydsl_flash_attn_fp8_func: softmax_scale must be positive and finite"
+        )
+
     Dv = int(v.shape[-1])
     if k.shape[-1] != D:
         raise ValueError(
@@ -525,6 +542,7 @@ def flydsl_flash_attn_fp8_func(
 
         kwargs = {
             "stream": launch_stream,
+            "softmax_scale": softmax_scale,
             "q_descale": q_descale,
             "k_descale": k_descale,
             "v_descale": v_descale,

aiter/ops/flydsl/kernels/fmha_gfx950/flash_attn_fp8_gfx950.py

@@ -129,7 +129,7 @@ def flash_attn_dualwave_swp_fp8_bn128_kernel(
         seq_len_kv: fx.Int32,
         stride_q_n: fx.Int32,
         stride_kv_n: fx.Int32,
-        head_dim_runtime: fx.Int32,
+        softmax_scale: fx.Float32,
         lse_stride_h: fx.Int32,
     ):
         ctx = DualwaveFp8KernelContext(
@@ -149,7 +149,7 @@ def flash_attn_dualwave_swp_fp8_bn128_kernel(
             seq_len_kv,
             stride_q_n,
             stride_kv_n,
-            head_dim_runtime,
+            softmax_scale,
             lse_stride_h,
         )
         ctx.init_types_and_constants()
@@ -411,7 +411,7 @@ def launch_flash_attn_dualwave_swp(
         seq_len_kv: fx.Int32,
         stride_q_n: fx.Int32,
         stride_kv_n: fx.Int32,
-        head_dim_runtime: fx.Int32,
+        softmax_scale: fx.Float32,
         lse_stride_h: fx.Int32,
         stream: fx.Stream = fx.Stream(None),  # noqa: B008  framework idiom
     ):
@@ -452,7 +452,7 @@ def launch_flash_attn_dualwave_swp(
             seq_len_kv,
             stride_q_n,
             stride_kv_n,
-            head_dim_runtime,
+            softmax_scale,
             lse_stride_h,
             value_attrs={
                 "rocdl.waves_per_eu": 1,
@@ -503,8 +503,8 @@ def _prepare(
         seq_len,
         stride_kv_n=None,
         stride_q_n=None,
-        head_dim_runtime=None,
         *,
+        softmax_scale=None,
         seq_len_kv=None,
         workspace=None,
         cu_seqlens_q=None,
@@ -521,8 +521,8 @@ def _prepare(
             stride_kv_n = DEFAULT_STRIDE_KV_N
         if stride_q_n is None:
             stride_q_n = DEFAULT_STRIDE_Q_N
-        if head_dim_runtime is None:
-            head_dim_runtime = HEAD_DIM
+        if softmax_scale is None:
+            softmax_scale = HEAD_DIM**-0.5
         # seq_len_kv defaults to seq_len (self-attention / equal Q,KV lengths).
         if seq_len_kv is None:
             seq_len_kv = seq_len
@@ -562,7 +562,7 @@ def _prepare(
             seq_len_kv,
             stride_q_n,
             stride_kv_n,
-            head_dim_runtime,
+            softmax_scale,
             lse_stride_h,
             fx.Stream(stream),
         )

aiter/ops/flydsl/kernels/fmha_gfx950/pipeline.py

@@ -694,7 +694,7 @@ def __init__(
         seq_len_kv=None,
         stride_q_n=None,
         stride_kv_n=None,
-        head_dim_runtime=None,
+        softmax_scale=None,
         lse_stride_h=None,
     ):
         if isinstance(traits_or_ctx, DualwaveFp8KernelContext):
@@ -718,7 +718,7 @@ def __init__(
         self.seq_len_kv = seq_len_kv
         self.stride_q_n = stride_q_n
         self.stride_kv_n = stride_kv_n
-        self.head_dim_runtime = head_dim_runtime
+        self.softmax_scale = softmax_scale
         self.lse_stride_h = lse_stride_h
 
     def init_types_and_constants(self):
@@ -882,9 +882,8 @@ def _load_scale_scalar(tensor):
             )
             return fx.Float32(Vec(_v, (1,), fx.Float32)[0])
 
-        head_dim_f32 = fx.Float32(self.head_dim_runtime)
         c_log2e_f = fx.Float32(_LOG2E)
-        c_sm_scale_log2e = fx.rsqrt(head_dim_f32, fastmath=self.fm_fast) * c_log2e_f
+        c_sm_scale_log2e = self.softmax_scale * c_log2e_f
         _qd = _load_scale_scalar(self.QDescale)
         _kd = _load_scale_scalar(self.KDescale)
         self.vd_fp8 = _load_scale_scalar(self.VDescale)