mha fp8 support descale

aiter/ops/mha.py

@@ -28,6 +28,9 @@ def cmdGenFunc_mha_fwd(
     out: Optional[Tensor] = None,
     bias: Optional[Tensor] = None,
     alibi_slopes: Optional[Tensor] = None,
+    q_descale: Optional[Tensor] = None,
+    k_descale: Optional[Tensor] = None,
+    v_descale: Optional[Tensor] = None,
     gen: Optional[Generator] = None,
 ):
     (_, seqlen_q, _, _) = q.shape
@@ -45,9 +48,11 @@ def cmdGenFunc_mha_fwd(
         md_name += "_bf16"
         filter += "bf16*"
     elif q.dtype == dtypes.fp8:
-        # only support bf16 out for fp8 input
-        md_name += "_fp8bf16"
-        filter += "fp8bf16*"
+        if out is None or out.dtype == dtypes.bf16:
+            md_name += "_fp8bf16"
+            filter += "fp8bf16*"
+        else:
+            raise NotImplementedError("Unsupported output dtype for FP8 MHA")
     if bias is not None:
         md_name += "_bias"
         filter += "_bias*"
@@ -75,6 +80,13 @@ def cmdGenFunc_mha_fwd(
     else:
         md_name += "_dropout"
         filter += "_dropout*"
+    if q_descale is None or k_descale is None or v_descale is None:
+        md_name += "_nqscale"
+        filter += "_nqscale*"
+    else:
+        # only support per-tensor quantization for now
+        md_name += "_pertensor"
+        filter += "_pertensor*"
 
     blob_gen_cmd = [
         f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d fwd "
@@ -103,7 +115,8 @@ def common_mha_fwd_fake_tensors(
     seqlen_k = k.size(1)
 
     if out is not None:
-        assert out.dtype == q.dtype, "Output must have the same dtype as inputs"
+        if q.dtype != dtypes.fp8:
+            assert out.dtype == q.dtype, "Output must have the same dtype as inputs"
         assert out.device == q.device, "Output must be on the same device as inputs"
         assert out.stride(-1) == 1, "Output tensor must have contiguous last dimension"
         assert out.shape == (
@@ -158,6 +171,9 @@ def gen_mha_fwd_fake_tensors(
     out: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
     alibi_slopes: Optional[torch.Tensor] = None,
+    q_descale: Optional[Tensor] = None,
+    k_descale: Optional[Tensor] = None,
+    v_descale: Optional[Tensor] = None,
     gen: Optional[torch.Generator] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     return common_mha_fwd_fake_tensors(
@@ -187,6 +203,9 @@ def mha_fwd(
     out: Optional[Tensor] = None,
     bias: Optional[Tensor] = None,
     alibi_slopes: Optional[Tensor] = None,
+    q_descale: Optional[Tensor] = None,
+    k_descale: Optional[Tensor] = None,
+    v_descale: Optional[Tensor] = None,
     gen: Optional[Generator] = None,
 ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: ...
 
@@ -257,6 +276,9 @@ def cmdGenFunc_mha_varlen_fwd(
     block_table: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
     alibi_slopes: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
     gen: Optional[torch.Generator] = None,
     cu_seqlens_q_padded: Optional[torch.Tensor] = None,
     cu_seqlens_k_padded: Optional[torch.Tensor] = None,
@@ -275,9 +297,11 @@ def cmdGenFunc_mha_varlen_fwd(
             md_name += "_bf16"
             filter_fwd += "bf16*"
         elif q.dtype == dtypes.fp8:
-            # only support bf16 out for fp8 input
-            md_name += "_fp8bf16"
-            filter_fwd += "fp8bf16*"
+            if out is None or out.dtype == dtypes.bf16:
+                md_name += "_fp8bf16"
+                filter_fwd += "fp8bf16*"
+            else:
+                raise NotImplementedError("Unsupported output dtype for FP8 MHA")
         if 0.0 < logits_soft_cap:
             md_name += "_logits"
             filter_fwd += "_logits*"
@@ -317,6 +341,13 @@ def cmdGenFunc_mha_varlen_fwd(
         else:
             md_name += "_skip"
             filter_fwd += "_skip*"
+        if q_descale is None or k_descale is None or v_descale is None:
+            md_name += "_nqscale"
+            filter_fwd += "_nqscale*"
+        else:
+            # only support per-tensor quantization for now
+            md_name += "_pertensor"
+            filter_fwd += "_pertensor*"
         blob_gen_cmd = [
             f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d fwd "
             "--receipt 200 --filter {} --output_dir {{}}".format(filter_fwd)
@@ -410,6 +441,9 @@ def gen_mha_varlen_fwd_fake_tensor(
     block_table: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
     alibi_slopes: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
     gen: Optional[torch.Generator] = None,
     cu_seqlens_q_padded: Optional[torch.Tensor] = None,
     cu_seqlens_k_padded: Optional[torch.Tensor] = None,
@@ -474,6 +508,9 @@ def mha_varlen_fwd(
     block_table: Optional[torch.Tensor] = None,
     bias: Optional[torch.Tensor] = None,
     alibi_slopes: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
     gen: Optional[torch.Generator] = None,
     cu_seqlens_q_padded: Optional[torch.Tensor] = None,
     cu_seqlens_k_padded: Optional[torch.Tensor] = None,
@@ -1176,6 +1213,9 @@ def _flash_attn_forward(
     window_size_right: int,
     bias: Optional[torch.Tensor],
     alibi_slopes: Optional[torch.Tensor],
+    q_descale: Optional[torch.Tensor],
+    k_descale: Optional[torch.Tensor],
+    v_descale: Optional[torch.Tensor],
     return_lse: bool,
     return_softmax: bool,
     how_v3_bf16_cvt: Optional[int] = 1,
@@ -1255,6 +1295,9 @@ def _validate_cu(name: str, x: Optional[torch.Tensor]):
             None,
             bias,
             alibi_slopes,
+            q_descale,
+            k_descale,
+            v_descale,
             None,
             # custom_build_args={"md_name": md_name, "blob_gen_cmd": blob_gen_cmd},
         )
@@ -1665,6 +1708,9 @@ def forward(
             window_size_right=int(window_size[1]),
             bias=bias,
             alibi_slopes=alibi_slopes,
+            q_descale=None,
+            k_descale=None,
+            v_descale=None,
             return_lse=return_lse,
             return_softmax=return_softmax and dropout_p > 0,
             how_v3_bf16_cvt=how_v3_bf16_cvt,
@@ -1879,6 +1925,9 @@ def _flash_attn_varlen_forward(
     window_size_right: int = -1,
     bias: Optional[torch.Tensor] = None,
     alibi_slopes: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
     return_lse: bool = False,
     return_softmax: bool = False,
     how_v3_bf16_cvt: Optional[int] = 1,
@@ -1982,6 +2031,9 @@ def _validate(name: str, t: torch.Tensor):
             block_table=block_table,
             bias=bias,
             alibi_slopes=alibi_slopes,
+            q_descale=q_descale,
+            k_descale=k_descale,
+            v_descale=v_descale,
             gen=None,
             cu_seqlens_q_padded=cu_seqlens_q_padded,
             cu_seqlens_k_padded=cu_seqlens_k_padded,
@@ -2242,6 +2294,9 @@ def forward(
             window_size_right=window_size[1],
             bias=bias,
             alibi_slopes=alibi_slopes,
+            q_descale=None,
+            k_descale=None,
+            v_descale=None,
             return_lse=return_lse,
             return_softmax=return_softmax and dropout_p > 0,
             how_v3_bf16_cvt=how_v3_bf16_cvt,
@@ -2696,6 +2751,9 @@ def flash_attn_fp8_pertensor_func(
     q,
     k,
     v,
+    q_descale,
+    k_descale,
+    v_descale,
     causal=False,
     window_size=(-1, -1),  # -1 means infinite context window
     softmax_scale=None,
@@ -2720,6 +2778,9 @@ def flash_attn_fp8_pertensor_func(
         window_size_right=int(window_size[1]),
         bias=None,
         alibi_slopes=None,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
         return_lse=False,
         return_softmax=False,
     )
@@ -2731,6 +2792,9 @@ def flash_attn_varlen_fp8_pertensor_func(
     q,
     k,
     v,
+    q_descale,
+    k_descale,
+    v_descale,
     cu_seqlens_q,
     cu_seqlens_k,
     max_seqlen_q,
@@ -2769,6 +2833,9 @@ def flash_attn_varlen_fp8_pertensor_func(
         window_size_right=int(window_size[1]),
         bias=None,
         alibi_slopes=None,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
         return_lse=False,
         return_softmax=False,
     )

csrc/cpp_itfs/mha_fwd_generate.py

@@ -37,6 +37,7 @@
                                   bias_enum bias_type,
                                   bool has_lse,
                                   bool has_dropout,
+                                  quant_scale_enum qscale_type,
                                   bool use_ext_asm,
                                   int how_v3_bf16_cvt = 1,
                                   bool skip_min_seqlen_q = false)
@@ -50,6 +51,7 @@
                           bias_type,
                           has_lse,
                           has_dropout,
+                          qscale_type,
                           use_ext_asm,
                           how_v3_bf16_cvt,
                           skip_min_seqlen_q);
@@ -87,6 +89,7 @@
               mask_enum mask_type,
               bias_enum bias_type,
               bool has_lse,
+              quant_scale_enum qscale_type,
               bool use_ext_asm,
               int how_v3_bf16_cvt,
               const void* seqstart_q_padding_ptr,
@@ -105,6 +108,7 @@
                                      bias_type,
                                      has_lse,
                                      has_dropout,
+                                     qscale_type,
                                      use_ext_asm,
                                      how_v3_bf16_cvt,
                                      args.min_seqlen_q != 0);

csrc/include/mha_fwd.h

@@ -20,6 +20,7 @@ struct mha_fwd_traits : public fmha_fwd_traits
                    bias_enum bias_type,
                    bool has_lse,
                    bool has_dropout,
+                   quant_scale_enum qscale_type,
                    bool use_ext_asm,
                    int how_v3_bf16_cvt,
                    bool skip_min_seqlen_q)
@@ -33,7 +34,7 @@ struct mha_fwd_traits : public fmha_fwd_traits
                           bias_type,
                           has_lse,
                           has_dropout,
-                          false, // do_fp8_static_quant
+                          qscale_type,
                           skip_min_seqlen_q},
           use_ext_asm(use_ext_asm),
           how_v3_bf16_cvt(how_v3_bf16_cvt)
@@ -78,6 +79,7 @@ __attribute__((visibility("default"))) float mha_fwd(mha_fwd_args args,
                                                      mask_enum mask_type,
                                                      bias_enum bias_type,
                                                      bool has_lse,
+                                                     quant_scale_enum qscale_type,
                                                      bool use_ext_asm,
                                                      int how_v3_bf16_cvt                = 1,
                                                      const void* seqstart_q_padding_ptr = nullptr,

csrc/include/rocm_ops.hpp

@@ -665,6 +665,9 @@ namespace py = pybind11;
           py::arg("out")           = std::nullopt, \
           py::arg("bias")          = std::nullopt, \
           py::arg("alibi_slopes")  = std::nullopt, \
+          py::arg("q_descale")     = std::nullopt, \
+          py::arg("k_descale")     = std::nullopt, \
+          py::arg("v_descale")     = std::nullopt, \
           py::arg("gen")           = std::nullopt);
 
 #define MHA_VARLEN_FWD_ASM_PYBIND                 \
@@ -821,6 +824,9 @@ namespace py = pybind11;
           py::arg("block_table")         = std::nullopt, \
           py::arg("bias")                = std::nullopt, \
           py::arg("alibi_slopes")        = std::nullopt, \
+          py::arg("q_descale")           = std::nullopt, \
+          py::arg("k_descale")           = std::nullopt, \
+          py::arg("v_descale")           = std::nullopt, \
           py::arg("gen")                 = std::nullopt, \
           py::arg("cu_seqlens_q_padded") = std::nullopt, \
           py::arg("cu_seqlens_k_padded") = std::nullopt);
@@ -1376,27 +1382,27 @@ namespace py = pybind11;
           py::arg("stride0"),      \
           py::arg("stride1"));
 
-#define MLA_METADATA_PYBIND                     \
-    m.def("get_mla_metadata_v1",                \
-          &get_mla_metadata_v1,                 \
-          "get_mla_metadata_v1",                \
-          py::arg("seqlens_qo_indptr"),         \
-          py::arg("seqlens_kv_indptr"),         \
-          py::arg("num_heads_per_head_k"),      \
-          py::arg("num_heads_k"),               \
-          py::arg("is_causal"),                 \
-          py::arg("work_metadata_ptrs"),        \
-          py::arg("work_info_set"),             \
-          py::arg("work_indptr"),               \
-          py::arg("reduce_indptr"),             \
-          py::arg("reduce_final_map"),          \
-          py::arg("reduce_partial_map"),        \
-          py::arg("kv_granularity") = 16,       \
-          py::arg("max_seqlen_qo")  = -1,       \
-          py::arg("uni_seqlen_qo")  = -1,       \
-          py::arg("fast_mode")      = true,     \
-          py::arg("topk")           = -1,       \
-          py::arg("max_split_per_batch") = -1); \
+#define MLA_METADATA_PYBIND                      \
+    m.def("get_mla_metadata_v1",                 \
+          &get_mla_metadata_v1,                  \
+          "get_mla_metadata_v1",                 \
+          py::arg("seqlens_qo_indptr"),          \
+          py::arg("seqlens_kv_indptr"),          \
+          py::arg("num_heads_per_head_k"),       \
+          py::arg("num_heads_k"),                \
+          py::arg("is_causal"),                  \
+          py::arg("work_metadata_ptrs"),         \
+          py::arg("work_info_set"),              \
+          py::arg("work_indptr"),                \
+          py::arg("reduce_indptr"),              \
+          py::arg("reduce_final_map"),           \
+          py::arg("reduce_partial_map"),         \
+          py::arg("kv_granularity")      = 16,   \
+          py::arg("max_seqlen_qo")       = -1,   \
+          py::arg("uni_seqlen_qo")       = -1,   \
+          py::arg("fast_mode")           = true, \
+          py::arg("topk")                = -1,   \
+          py::arg("max_split_per_batch") = -1);  \
     m.def("get_mla_metadata_v1_no_redundant", &get_mla_metadata_v1_no_redundant);
 
 #define MLA_REDUCE_PYBIND                \

csrc/include/torch/mha_fwd.h

@@ -20,6 +20,9 @@ std::vector<at::Tensor> mha_fwd(at::Tensor& q,       // [b, sq, hq, d]
                                 std::optional<at::Tensor> out,                // [b, sq, hq, d]
                                 std::optional<const at::Tensor> bias,         // [sq, sk]
                                 std::optional<const at::Tensor> alibi_slopes, // [hq] or [b, hq]
+                                std::optional<const at::Tensor> q_descale,    // [1]
+                                std::optional<const at::Tensor> k_descale,    // [1]
+                                std::optional<const at::Tensor> v_descale,    // [1]
                                 std::optional<at::Generator> gen);
 } // namespace torch_itfs
 } // namespace aiter

csrc/include/torch/mha_varlen_fwd.h

@@ -27,6 +27,9 @@ mha_varlen_fwd(at::Tensor& q,                                 // [total_q, hq, d
                std::optional<const at::Tensor> block_table,  // [hq] or [b, hq]
                std::optional<const at::Tensor> bias,         // [total_q, max_seqlen_k]
                std::optional<const at::Tensor> alibi_slopes, // [hq] or [b, hq]
+               std::optional<const at::Tensor> q_descale,    // [1]
+               std::optional<const at::Tensor> k_descale,    // [1]
+               std::optional<const at::Tensor> v_descale,    // [1]
                std::optional<at::Generator> gen,
                std::optional<const at::Tensor> cu_seqlens_q_padded = std::nullopt,
                std::optional<const at::Tensor> cu_seqlens_k_padded = std::nullopt);