Add softmax_scale to blackwell cutlass fmha

Summary: Adding softmax_scale plumbing to unblock D82490887 for tritonbench

Differential Revision: D82788784

Author: henrylhtsang

fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py

@@ -100,6 +100,7 @@ def _cutlass_blackwell_fmha_backward(
     cu_seqlens_k: torch.Tensor | None = None,
     max_seq_len_q: int | None = None,
     max_seq_len_k: int | None = None,
+    softmax_scale: float | None = None,
     causal: bool = False,
     window_left: int = -1,
     window_right: int = -1,
@@ -123,6 +124,7 @@ def _cutlass_blackwell_fmha_backward(
         cu_seqlens_k=cu_seqlens_k,
         max_seq_len_q=max_seq_len_q,
         max_seq_len_k=max_seq_len_k,
+        softmax_scale=softmax_scale,
         causal=causal,
         window_size_left=window_left,
         window_size_right=window_right,
@@ -274,6 +276,7 @@ def backward(ctx, dout: torch.Tensor, *args: Any) -> Tuple[  # type: ignore
             ctx.cu_seqlens_k,
             ctx.max_seq_len_q,
             ctx.max_seq_len_k,
+            ctx.softmax_scale,
             ctx.causal,
             window_left,
             window_right,

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_bwd.cu

@@ -19,6 +19,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
     const std::optional<const at::Tensor>& cu_seqlens_k,
     std::optional<int> max_seq_len_q,
     std::optional<int> max_seq_len_k,
+    const std::optional<double> softmax_scale,
     const int window_size_left,
     const int window_size_right
 ) {
@@ -207,8 +208,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
     get<2, 1>(stride_dV) = 0;
   }
 
-  // TODO: pass in softmax_scale?
-  ElementAccumulator softmax_scale = 1.0f / sqrtf(D);
+  ElementAccumulator softmax_scale_value = softmax_scale.has_value() ? softmax_scale.value() : (1.0f / sqrtf(D));
 
   at::Tensor dQ = torch::empty_like(q);
   at::Tensor dK = torch::empty_like(k);
@@ -253,7 +253,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fmha_bwd(
     stride_dK,
     static_cast<Element*>(dV.data_ptr()),
     stride_dV,
-    softmax_scale,
+    softmax_scale_value,
     dq_semaphore_ptr,
     window_size_left,
     window_size_right,
@@ -276,6 +276,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
     const std::optional<at::Tensor>& cu_seqlens_k,
     std::optional<int64_t> max_seq_len_q,
     std::optional<int64_t> max_seq_len_k,
+    std::optional<double> softmax_scale,
     bool causal,
     int64_t window_size_left,
     int64_t window_size_right,
@@ -334,6 +335,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> dispatch_fmha_bwd(
         cu_seqlens_k,
         max_seq_len_q,
         max_seq_len_k,
+        softmax_scale,
         window_size_left,
         window_size_right);
     };
@@ -419,6 +421,7 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
         "    Tensor? cu_seqlens_k=None, "
         "    int? max_seq_len_q=None, "
         "    int? max_seq_len_k=None, "
+        "    float? softmax_scale=None, "
         "    bool causal=False, "
         "    int window_size_left=-1, "
         "    int window_size_right=-1, "

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd.cu

@@ -225,7 +225,7 @@ std::tuple<at::Tensor, at::Tensor> fmha_fwd(
               static_cast<Element*>(v.data_ptr()), stride_V,
               window_size_left, window_size_right
           },
-          0.0f /* softmax_scale */,
+          static_cast<float>(softmax_scale.value_or(0.0f)) /* softmax_scale */,
           1.0f /* scale_q */,
           1.0f /* scale_k */,
           1.0f /* scale_v */,

fbgemm_gpu/experimental/gen_ai/test/attention/blackwell_fmha_test.py

@@ -6,7 +6,7 @@
 
 import random
 import unittest
-from typing import Tuple
+from typing import Optional, Tuple
 
 import hypothesis.strategies as st
 import torch
@@ -117,10 +117,15 @@ def _execute_cutlass_blackwell_attn_dense(
         window_size: tuple[int, int],
         fwd_only: bool,
         deterministic: bool,
+        sm_scale: Optional[float],
     ) -> None:
         device = torch.accelerator.current_accelerator()
         assert device is not None
         assert seqlen_q <= seqlen_k
+
+        # Initialize deterministic variables
+        out_d = None
+
         q, k, v = self._generate_qkv(
             batch_size,
             seqlen_q,
@@ -144,7 +149,13 @@ def _execute_cutlass_blackwell_attn_dense(
 
         # Run reference attention
         out_baseline, _ = attention_ref(
-            q, k, v, causal=causal, window_size=window_size, upcast=True
+            q,
+            k,
+            v,
+            causal=causal,
+            window_size=window_size,
+            upcast=True,
+            softmax_scale=sm_scale,
         )
         if dtype == torch.float8_e4m3fn:
             # reference implementation only supports decode case (seqlen_q == 1)
@@ -161,6 +172,7 @@ def _execute_cutlass_blackwell_attn_dense(
                 window_size=window_size,
                 reorder_ops=True,
                 upcast=False,
+                softmax_scale=sm_scale,
             )
 
         # Run tested kernel
@@ -172,6 +184,7 @@ def _execute_cutlass_blackwell_attn_dense(
             window_size=window_size,
             seqlen_kv=seqlen_kv,
             deterministic=deterministic,
+            softmax_scale=sm_scale,
         )
         if DEBUG:
             print("cutlass_blackwell_fmha_func completed successfully!")
@@ -190,6 +203,7 @@ def _execute_cutlass_blackwell_attn_dense(
                 window_size=window_size,
                 seqlen_kv=seqlen_kv,
                 deterministic=deterministic,
+                softmax_scale=sm_scale,
             )
             assert torch.equal(out, out_d)
 
@@ -244,9 +258,13 @@ def _execute_cutlass_blackwell_attn_varlen(
         window_size: tuple[int, int],
         fwd_only: bool,
         deterministic: bool,
+        sm_scale: Optional[float],
     ) -> None:
         device = torch.accelerator.current_accelerator()
         assert device is not None
+
+        # Initialize deterministic variables
+        out_unpad_d = None
         q_ref, k_ref, v_ref = self._generate_qkv(
             batch_size,
             seqlen_q,
@@ -306,6 +324,7 @@ def _execute_cutlass_blackwell_attn_varlen(
             key_padding_mask,
             causal=causal,
             window_size=window_size,
+            softmax_scale=sm_scale,
         )
 
         out_pt, _ = attention_ref(
@@ -318,6 +337,7 @@ def _execute_cutlass_blackwell_attn_varlen(
             window_size=window_size,
             upcast=False,
             reorder_ops=True,
+            softmax_scale=sm_scale,
         )
 
         out_unpad = cutlass_blackwell_fmha_func(
@@ -331,6 +351,7 @@ def _execute_cutlass_blackwell_attn_varlen(
             max_seq_len_k=max_seqlen_k,
             window_size=window_size,
             deterministic=deterministic,
+            softmax_scale=sm_scale,
         )
         out = output_pad_fn(out_unpad)
 
@@ -351,6 +372,7 @@ def _execute_cutlass_blackwell_attn_varlen(
                 max_seq_len_k=max_seqlen_k,
                 window_size=window_size,
                 deterministic=deterministic,
+                softmax_scale=sm_scale,
             )
             out_d = output_pad_fn(out_unpad_d)
             assert torch.equal(out, out_d)
@@ -396,11 +418,13 @@ def _execute_cutlass_blackwell_attn_varlen(
                 batch_size,
                 is_mqa,
                 window_size,
+                sm_scale,
             )
             for seqlen_k in [64, 128, 256, 1024]
             for batch_size in [1, 2]
             for is_mqa in [True]
             for window_size in [(-1, -1), (0, 0), (0, 128), (128, 0), (1024, 0)]
+            for sm_scale in [None, 1.0 / 128]
         ]
     )
     def test_decode(
@@ -409,6 +433,7 @@ def test_decode(
         batch_size: int,
         is_mqa: bool,
         window_size: tuple[int, int],
+        sm_scale: Optional[float],
         q_heads: int = 8,
         dtype: torch.dtype = torch.float8_e4m3fn,
     ) -> None:
@@ -429,6 +454,7 @@ def test_decode(
             window_size=window_size,
             fwd_only=True,
             deterministic=False,
+            sm_scale=sm_scale,
         )
 
     @skip_cuda_lt_sm100
@@ -441,12 +467,14 @@ def test_decode(
                 q_heads,
                 causal,
                 window_size,
+                sm_scale,
             )
             for kv_padding in [128, 256, 512, 1024]
             for batch_size in [2, 8]
             for q_heads in [8, 16]
             for causal in [True, False]
             for window_size in [(-1, -1), (0, 0), (0, 128), (128, 0), (1024, 0)]
+            for sm_scale in [None, 1.0 / 128]
         ]
     )
     def test_jagged_vs_padded_kv(
@@ -455,7 +483,8 @@ def test_jagged_vs_padded_kv(
         batch_size: int,
         q_heads: int,
         causal: bool,
-        window_size: tuple[int, int] = (-1, -1),
+        window_size: tuple[int, int],
+        sm_scale: Optional[float],
     ) -> None:
         """
         Test comparing two scenarios:
@@ -565,6 +594,7 @@ def test_jagged_vs_padded_kv(
             max_seq_len_k=max_seqlen_k,
             causal=causal,
             window_size=window_size,
+            softmax_scale=sm_scale,
         )
 
         # # Scenario B: Padded KV with seqlen_kv
@@ -583,6 +613,7 @@ def test_jagged_vs_padded_kv(
             causal=causal,
             window_size=window_size,
             seqlen_kv=seqused_k,
+            softmax_scale=sm_scale,
         )
         if DEBUG:
             print(f"out_jagged: {out_jagged}")
@@ -611,6 +642,7 @@ def test_jagged_vs_padded_kv(
                 is_varlen,
                 kv_heads,
                 window_size,
+                sm_scale,
             )
             for seqlen_q, offset_q in [
                 (101, 0),
@@ -629,6 +661,7 @@ def test_jagged_vs_padded_kv(
             for is_varlen in [False, True]
             for kv_heads in [1, 2, 3, 4]
             for window_size in [(-1, -1), (0, 0), (0, 128), (128, 0), (1024, 0)]
+            for sm_scale in [None, 1.0 / 128]
         ]
     )
     def test_forward(
@@ -641,6 +674,7 @@ def test_forward(
         is_varlen: bool,
         kv_heads: int,
         window_size: tuple[int, int],
+        sm_scale: Optional[float],
         dtype: torch.dtype = torch.bfloat16,
     ) -> None:
         seqlen_k = offset_q + seqlen_q
@@ -664,6 +698,7 @@ def test_forward(
             window_size=window_size,
             fwd_only=True,
             deterministic=False,
+            sm_scale=sm_scale,
         )
 
     @skip_cuda_lt_sm100
@@ -680,6 +715,7 @@ def test_forward(
             [(-1, -1), (128, 0), (256, 0), (128, 128), (512, 0)]
         ),
         deterministic=st.booleans(),
+        sm_scale=st.sampled_from([None, 1.0 / 128]),
     )
     @settings(**common_settings)
     def test_backward(
@@ -693,6 +729,7 @@ def test_backward(
         is_gqa: bool,
         window_size: tuple[int, int],
         deterministic: bool,
+        sm_scale: Optional[float],
     ) -> None:
         test_func = (
             self._execute_cutlass_blackwell_attn_varlen
@@ -712,4 +749,5 @@ def test_backward(
             window_size=window_size,
             fwd_only=False,
             deterministic=deterministic,
+            sm_scale=sm_scale,
         )

fbgemm_gpu/experimental/gen_ai/test/attention/test_utils.py

@@ -243,6 +243,7 @@ def attention_ref(  # noqa
     upcast=True,
     reorder_ops=False,
     key_leftpad=None,
+    softmax_scale=None,
 ):
     """
     Arguments:
@@ -261,6 +262,7 @@ def attention_ref(  # noqa
         reorder_ops: whether to change the order of operations (scaling k instead of scaling q, etc.)
             without changing the math. This is to estimate the numerical error from operation
             reordering.
+        softmax_scale: float, scale for softmax. If None, use 1/sqrt(head_dim)
     Output:
         output: (batch_size, seqlen_q, nheads, head_dim)
         attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout
@@ -274,10 +276,12 @@ def attention_ref(  # noqa
     k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2])
     v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2])
     d = q.shape[-1]
+    # Use provided softmax_scale or default to 1/sqrt(d)
+    scale = softmax_scale if softmax_scale is not None else 1 / math.sqrt(d)
     if not reorder_ops:
-        scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k)
+        scores = torch.einsum("bthd,bshd->bhts", q * scale, k)
     else:
-        scores = torch.einsum("bthd,bshd->bhts", q, k / math.sqrt(d))
+        scores = torch.einsum("bthd,bshd->bhts", q, k * scale)
     if softcap > 0:
         scores /= softcap
         scores = scores.tanh()