diff --git a/3rdparty/composable_kernel b/3rdparty/composable_kernel index e31a7a4f29..a38aeceb21 160000 --- a/3rdparty/composable_kernel +++ b/3rdparty/composable_kernel @@ -1 +1 @@ -Subproject commit e31a7a4f29b371c32ea9daf9211b6ae1fed2fa40 +Subproject commit a38aeceb2164f9d1807bda1a19d59636bafd4f31 diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py index a05a940c9c..5453fb6055 100644 --- a/aiter/ops/mha.py +++ b/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, ) diff --git a/csrc/cpp_itfs/mha_fwd_generate.py b/csrc/cpp_itfs/mha_fwd_generate.py index 5fe5190064..24dc5b9f70 100644 --- a/csrc/cpp_itfs/mha_fwd_generate.py +++ b/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); @@ -157,6 +161,7 @@ bias_type, has_lse, has_dropout, + quant_scale_enum::no_scale, use_ext_asm); return fmha_batch_prefill(traits, args, stream_config); }""" diff --git a/csrc/include/mha_fwd.h b/csrc/include/mha_fwd.h index 6384cb14e5..84e84e386a 100644 --- a/csrc/include/mha_fwd.h +++ b/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, diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index b0bf11ef56..8646633588 100644 --- a/csrc/include/rocm_ops.hpp +++ b/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); diff --git a/csrc/include/torch/mha_fwd.h b/csrc/include/torch/mha_fwd.h index b998498479..48fdfc2b8f 100644 --- a/csrc/include/torch/mha_fwd.h +++ b/csrc/include/torch/mha_fwd.h @@ -20,6 +20,9 @@ std::vector mha_fwd(at::Tensor& q, // [b, sq, hq, d] std::optional out, // [b, sq, hq, d] std::optional bias, // [sq, sk] std::optional alibi_slopes, // [hq] or [b, hq] + std::optional q_descale, // [1] + std::optional k_descale, // [1] + std::optional v_descale, // [1] std::optional gen); } // namespace torch_itfs } // namespace aiter diff --git a/csrc/include/torch/mha_varlen_fwd.h b/csrc/include/torch/mha_varlen_fwd.h index b9c0483102..3b7ab8dea8 100644 --- a/csrc/include/torch/mha_varlen_fwd.h +++ b/csrc/include/torch/mha_varlen_fwd.h @@ -27,6 +27,9 @@ mha_varlen_fwd(at::Tensor& q, // [total_q, hq, d std::optional block_table, // [hq] or [b, hq] std::optional bias, // [total_q, max_seqlen_k] std::optional alibi_slopes, // [hq] or [b, hq] + std::optional q_descale, // [1] + std::optional k_descale, // [1] + std::optional v_descale, // [1] std::optional gen, std::optional cu_seqlens_q_padded = std::nullopt, std::optional cu_seqlens_k_padded = std::nullopt); diff --git a/csrc/py_itfs_ck/mha_fwd_kernels.cu b/csrc/py_itfs_ck/mha_fwd_kernels.cu index d53678360d..1ab440790d 100644 --- a/csrc/py_itfs_ck/mha_fwd_kernels.cu +++ b/csrc/py_itfs_ck/mha_fwd_kernels.cu @@ -27,6 +27,9 @@ mha_fwd_args get_ck_fmha_fwd_args(bool has_lse, const at::Tensor v, std::optional &bias_, std::optional &alibi_slopes_, + std::optional &q_descale_, + std::optional &k_descale_, + std::optional &v_descale_, at::Tensor out, at::Tensor softmax_lse, at::Tensor dropout_randval, @@ -87,61 +90,66 @@ mha_fwd_args get_ck_fmha_fwd_args(bool has_lse, stride_bias = alibi_slopes.dim() == 2 ? alibi_slopes.stride(0) : 0; } + const void *q_descale_ptr = q_descale_.has_value() ? q_descale_.value().data_ptr() : nullptr; + const void *k_descale_ptr = k_descale_.has_value() ? k_descale_.value().data_ptr() : nullptr; + const void *v_descale_ptr = v_descale_.has_value() ? v_descale_.value().data_ptr() : nullptr; + const ck_tile::index_t *cu_seqlen_q_ptr = cu_seqlens_q_.has_value() ? reinterpret_cast(cu_seqlens_q_.value().data_ptr()) : nullptr; const ck_tile::index_t *cu_seqlen_kv_ptr = cu_seqlens_kv_.has_value() ? reinterpret_cast(cu_seqlens_kv_.value().data_ptr()) : nullptr; return mha_fwd_args{q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - bias_ptr, - has_dropout_randval ? dropout_randval.data_ptr() : nullptr, - has_lse ? softmax_lse.data_ptr() : nullptr, - out.data_ptr(), - nullptr, // seqstart_q_ptr - nullptr, // seqstart_k_ptr - nullptr, // seqlen_q_ptr - nullptr, // seqlen_k_ptr - cu_seqlen_q_ptr, // cu_seqlen_q_ptr - cu_seqlen_kv_ptr, // cu_seqlen_k_ptr - seqlen_q, - seqlen_k, - b, - seqlen_q, // max_seqlen_q - d, // hdim_q - d_v, // hdim_v - h, // nhead_q - h_k, // nhead_k - softmax_scale, // scale_s - 1, // scale_p - 1, // scale_o - 0.0, // logits_soft_cap - stride_q, - stride_k, - stride_v, - stride_bias, - stride_randval, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - 0, // nhead_stride_bias - nhead_stride_randval, - nhead_stride_lse, - nhead_stride_o, - batch_stride_q, - batch_stride_k, - batch_stride_v, - 0, // batch_stride_bias - batch_stride_randval, - batch_stride_lse, - batch_stride_o, - mask.left, - mask.right, - static_cast(mask.type), - 0, // min_seqlen_q - p_dropout, - has_dropout_randval, - drop_seed_offset}; + k.data_ptr(), + v.data_ptr(), + bias_ptr, + q_descale_ptr, + k_descale_ptr, + v_descale_ptr, + has_dropout_randval ? dropout_randval.data_ptr() : nullptr, + has_lse ? softmax_lse.data_ptr() : nullptr, + out.data_ptr(), + nullptr, // seqstart_q_ptr + nullptr, // seqstart_k_ptr + nullptr, // seqlen_q_ptr + nullptr, // seqlen_k_ptr + cu_seqlen_q_ptr, // cu_seqlen_q_ptr + cu_seqlen_kv_ptr, // cu_seqlen_k_ptr + seqlen_q, + seqlen_k, + b, + seqlen_q, // max_seqlen_q + d, // hdim_q + d_v, // hdim_v + h, // nhead_q + h_k, // nhead_k + softmax_scale, // scale_s + 0.0, // logits_soft_cap + stride_q, + stride_k, + stride_v, + stride_bias, + stride_randval, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + 0, // nhead_stride_bias + nhead_stride_randval, + nhead_stride_lse, + nhead_stride_o, + batch_stride_q, + batch_stride_k, + batch_stride_v, + 0, // batch_stride_bias + batch_stride_randval, + batch_stride_lse, + batch_stride_o, + mask.left, + mask.right, + static_cast(mask.type), + 0, // min_seqlen_q + p_dropout, + has_dropout_randval, + drop_seed_offset}; } std::vector @@ -160,27 +168,38 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d] std::optional out_, // [b, sq, hq, d_v] std::optional bias_, // [sq, sk] std::optional alibi_slopes_, // [hq] or [b, hq] + std::optional q_descale_, // [1] + std::optional k_descale_, // [1] + std::optional v_descale_, // [1] std::optional gen_) { auto q_dtype = q.scalar_type(); - bool isQKVFp8 = q_dtype == at::ScalarType::Float8_e4m3fn || q_dtype == at::ScalarType::Float8_e4m3fnuz; + bool is_qkv_fp8 = q_dtype == at::ScalarType::Float8_e4m3fn || q_dtype == at::ScalarType::Float8_e4m3fnuz; - TORCH_CHECK(q_dtype == at::ScalarType::Half || q_dtype == at::ScalarType::BFloat16 || isQKVFp8, + TORCH_CHECK(q_dtype == at::ScalarType::Half || q_dtype == at::ScalarType::BFloat16 || is_qkv_fp8, "FlashAttention only support fp16, bf16 and fp8_e4m3 data type"); TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype"); TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); - std::string q_dtype_str; + std::string dtype_str; if (q_dtype == at::ScalarType::Half) - q_dtype_str = "fp16"; + dtype_str = "fp16"; else if (q_dtype == at::ScalarType::BFloat16) - q_dtype_str = "bf16"; - else if (isQKVFp8) - q_dtype_str = "fp8bf16"; // only support bf16 out for fp8 + dtype_str = "bf16"; + else if (is_qkv_fp8) { + if (!out_.has_value() || out_.value().dtype() == at::ScalarType::BFloat16) + dtype_str = "fp8bf16"; // only support bf16 out for fp8 + else + TORCH_CHECK(false, "For FP8 input, output must have dtype BF16 for now"); + } + + TORCH_CHECK(q_descale_.has_value() == k_descale_.has_value() && + k_descale_.has_value() == v_descale_.has_value(), + "q_descale, k_descale, v_descale must be all provided or all not provided"); - // TODO - support descale - // TODO - set do_fp8_static_quant to true in fmha_fwd_traits + quant_scale_enum qscale_type = + q_descale_.has_value() ? quant_scale_enum::pertensor : quant_scale_enum::no_scale; CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v); @@ -246,8 +265,9 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d] CHECK_SHAPE(k, batch_size, seqlen_k, num_heads_k, head_size_q); CHECK_SHAPE(v, batch_size, seqlen_k, num_heads_k, head_size_v); + // TODO - Support fp8fp16 auto opts = q.options(); - auto out_type = isQKVFp8 ? at::ScalarType::BFloat16 : q_dtype; + auto out_type = dtype_str == "fp8bf16" ? at::ScalarType::BFloat16 : q_dtype; at::Tensor out; if (out_.has_value()) { out = out_.value(); @@ -322,6 +342,9 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d] v, bias_, alibi_slopes_, + q_descale_, + k_descale_, + v_descale_, out, softmax_lse, p, @@ -333,11 +356,12 @@ mha_fwd(at::Tensor &q, // [b, sq, hq, d] float t = aiter::mha_fwd(args, stream_config, - q_dtype_str, + dtype_str, false, // is_group_mode mask.type, bias_type, has_lse, + qscale_type, false); TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd"); } diff --git a/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu b/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu index 54ca060dc9..7f18a24bf9 100644 --- a/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu +++ b/csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu @@ -30,6 +30,9 @@ mha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, std::optional &seqlens_k, std::optional &bias_, std::optional &alibi_slopes_, + std::optional &q_descale_, + std::optional &k_descale_, + std::optional &v_descale_, at::Tensor out, at::Tensor softmax_lse, at::Tensor dropout_randval, @@ -96,7 +99,11 @@ mha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, bias_ptr = alibi_slopes.data_ptr(); stride_bias = alibi_slopes.dim() == 2 ? alibi_slopes.stride(0) : 0; } - + + const void *q_descale_ptr = q_descale_.has_value() ? q_descale_.value().data_ptr() : nullptr; + const void *k_descale_ptr = k_descale_.has_value() ? k_descale_.value().data_ptr() : nullptr; + const void *v_descale_ptr = v_descale_.has_value() ? v_descale_.value().data_ptr() : nullptr; + const void* seqstart_k_ptr = nullptr; const void* seqstart_q_ptr = nullptr; const void* cu_seqlen_k_ptr = nullptr; @@ -117,57 +124,58 @@ mha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, } return mha_fwd_args{q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - bias_ptr, - has_dropout_randval ? dropout_randval.data_ptr() : nullptr, - has_lse ? softmax_lse.data_ptr() : nullptr, - out.data_ptr(), - seqstart_q_ptr, // seqstart_q_ptr (cumulative physical with padding) - seqstart_k_ptr, // seqstart_k_ptr (cumulative physical with padding) - nullptr, // seqlen_q_ptr (per-sequence logical, alternative to cu_seqlen_q_ptr) - seqlens_k.has_value() ? seqlens_k.value().data_ptr() : nullptr, // seqlen_k_ptr (per-sequence logical K lengths) - cu_seqlen_q_ptr, // cu_seqlen_q_ptr - cu_seqlen_k_ptr, // cu_seqlen_k_ptr - total_q, - total_k, - b, - max_seqlen_q, - d, // hdim_q - d_v, // hdim_v - h, // nhead_q - h_k, // nhead_k - softmax_scale, // scale_s - 1, // scale_p - 1, // scale_o - logits_soft_cap, - stride_q, - stride_k, - stride_v, - stride_bias, - stride_randval, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - nhead_stride_bias, - nhead_stride_randval, - nhead_stride_lse, - nhead_stride_o, - batch_stride_q, - batch_stride_k, - batch_stride_v, - batch_stride_bias, - batch_stride_randval, - batch_stride_lse, - batch_stride_o, - mask.left, - mask.right, - static_cast(mask.type), - min_seqlen_q, - p_dropout, - has_dropout_randval, - drop_seed_offset}; + k.data_ptr(), + v.data_ptr(), + bias_ptr, + q_descale_ptr, + k_descale_ptr, + v_descale_ptr, + has_dropout_randval ? dropout_randval.data_ptr() : nullptr, + has_lse ? softmax_lse.data_ptr() : nullptr, + out.data_ptr(), + seqstart_q_ptr, // seqstart_q_ptr (cumulative physical with padding) + seqstart_k_ptr, // seqstart_k_ptr (cumulative physical with padding) + nullptr, // seqlen_q_ptr (per-sequence logical, alternative to cu_seqlen_q_ptr) + seqlens_k.has_value() ? seqlens_k.value().data_ptr() : nullptr, // seqlen_k_ptr (per-sequence logical K lengths) + cu_seqlen_q_ptr, // cu_seqlen_q_ptr + cu_seqlen_k_ptr, // cu_seqlen_k_ptr + total_q, + total_k, + b, + max_seqlen_q, + d, // hdim_q + d_v, // hdim_v + h, // nhead_q + h_k, // nhead_k + softmax_scale, // scale_s + logits_soft_cap, + stride_q, + stride_k, + stride_v, + stride_bias, + stride_randval, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + nhead_stride_bias, + nhead_stride_randval, + nhead_stride_lse, + nhead_stride_o, + batch_stride_q, + batch_stride_k, + batch_stride_v, + batch_stride_bias, + batch_stride_randval, + batch_stride_lse, + batch_stride_o, + mask.left, + mask.right, + static_cast(mask.type), + min_seqlen_q, + p_dropout, + has_dropout_randval, + drop_seed_offset}; } fmha_fwd_splitkv_args get_ck_fmha_varlen_fwd_splitkv_args(bool has_lse, @@ -319,7 +327,7 @@ fmha_fwd_splitkv_args get_ck_fmha_varlen_fwd_splitkv_args(bool has_lse, return args; } -std::tuple +std::tuple mha_varlen_fwd( at::Tensor& q, // [total_q, hq, d] const at::Tensor& k, // [total_k, hk, d] @@ -342,28 +350,40 @@ mha_varlen_fwd( std::optional block_table_, // [hq] or [b, hq] std::optional bias_, // [total_q, max_seqlen_k] std::optional alibi_slopes_, // [hq] or [b, hq] + std::optional q_descale_, // [1] + std::optional k_descale_, // [1] + std::optional v_descale_, // [1] std::optional gen_, std::optional cu_seqlens_q_padded_, // [b+1] physical starts with PAD std::optional cu_seqlens_k_padded_) // [b+1] { auto q_dtype = q.scalar_type(); - bool isQKVFp8 = q_dtype == at::ScalarType::Float8_e4m3fn || q_dtype == at::ScalarType::Float8_e4m3fnuz; + bool is_qkv_fp8 = q_dtype == at::ScalarType::Float8_e4m3fn || q_dtype == at::ScalarType::Float8_e4m3fnuz; - TORCH_CHECK(q_dtype == at::ScalarType::Half || q_dtype == at::ScalarType::BFloat16 || isQKVFp8, + TORCH_CHECK(q_dtype == at::ScalarType::Half || q_dtype == at::ScalarType::BFloat16 || is_qkv_fp8, "FlashAttention only support fp16, bf16 and fp8_e4m3 data type"); TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype"); TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype"); - std::string q_dtype_str; + std::string dtype_str; if (q_dtype == at::ScalarType::Half) - q_dtype_str = "fp16"; + dtype_str = "fp16"; else if (q_dtype == at::ScalarType::BFloat16) - q_dtype_str = "bf16"; - else if (isQKVFp8) - q_dtype_str = "fp8bf16"; // only support bf16 out for fp8 + dtype_str = "bf16"; + else if (is_qkv_fp8) { + if (!out_.has_value() || out_.value().dtype() == at::ScalarType::BFloat16) + dtype_str = "fp8bf16"; // only support bf16 out for fp8 + else + TORCH_CHECK(false, "For FP8 input, output must have dtype BF16 for now"); + } + + TORCH_CHECK(q_descale_.has_value() == k_descale_.has_value() && + k_descale_.has_value() == v_descale_.has_value(), + "q_descale, k_descale, v_descale must be all provided or all not provided"); - // TODO - support descale + quant_scale_enum qscale_type = + q_descale_.has_value() ? quant_scale_enum::pertensor : quant_scale_enum::no_scale; TORCH_CHECK(cu_seqlens_q.dtype() == torch::kInt32, "cu_seqlens_q must have dtype int32"); if (cu_seqlens_k.has_value()) { @@ -463,8 +483,10 @@ mha_varlen_fwd( if (cu_seqlens_k.has_value()) { CHECK_SHAPE(cu_seqlens_k.value(), batch_size + 1); } + + // TODO - Support fp8fp16 auto opts = q.options(); - auto out_type = isQKVFp8 ? at::ScalarType::BFloat16 : q_dtype; + auto out_type = dtype_str == "fp8bf16" ? at::ScalarType::BFloat16 : q_dtype; at::Tensor out; if (out_.has_value()) { out = out_.value(); @@ -568,7 +590,7 @@ mha_varlen_fwd( float t = aiter::mha_fwd_splitkv(args, stream_config, - q_dtype_str, + dtype_str, true, //is_group_mode mask.type, bias_type, @@ -599,6 +621,9 @@ mha_varlen_fwd( seqlens_k, bias_, alibi_slopes_, + q_descale_, + k_descale_, + v_descale_, out, softmax_lse, p, @@ -611,11 +636,12 @@ mha_varlen_fwd( float t = aiter::mha_fwd(args, stream_config, - q_dtype_str, + dtype_str, true, //is_group_mode mask.type, bias_type, has_lse, + qscale_type, false, // use_ext_asm 1); // how_v3_bf16_cvt TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd"); diff --git a/csrc/py_itfs_cu/asm_mha_fwd.cu b/csrc/py_itfs_cu/asm_mha_fwd.cu index 33cd53ca6b..e68848136e 100644 --- a/csrc/py_itfs_cu/asm_mha_fwd.cu +++ b/csrc/py_itfs_cu/asm_mha_fwd.cu @@ -86,57 +86,58 @@ mha_fwd_args get_asm_fmha_fwd_args(bool has_lse, } return mha_fwd_args{q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - bias_ptr, - has_dropout_randval ? dropout_randval.data_ptr() : nullptr, - has_lse ? softmax_lse.data_ptr() : nullptr, - out.data_ptr(), - nullptr, // seqstart_q_ptr - nullptr, // seqstart_k_ptr - nullptr, // seqlen_q_ptr - nullptr, // seqlen_k_ptr - nullptr, // cu_seqlen_q_ptr - nullptr, // cu_seqlen_k_ptr - seqlen_q, - seqlen_k, - b, - seqlen_q, // max_seqlen_q - d, // hdim_q - d_v, // hdim_v - h, // nhead_q - h_k, // nhead_k - softmax_scale, // scale_s - 1, // scale_p - 1, // scale_o - 0.0, // logits_soft_cap - stride_q, - stride_k, - stride_v, - stride_bias, - stride_randval, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - 0, // nhead_stride_bias - nhead_stride_randval, - nhead_stride_lse, - nhead_stride_o, - batch_stride_q, - batch_stride_k, - batch_stride_v, - 0, // batch_stride_bias - batch_stride_randval, - batch_stride_lse, - batch_stride_o, - mask.left, - mask.right, - static_cast(mask.type), - 0, // min_seqlen_q - p_dropout, - has_dropout_randval, - drop_seed_offset}; + k.data_ptr(), + v.data_ptr(), + bias_ptr, + nullptr, // q_descale_ptr + nullptr, // k_descale_ptr + nullptr, // v_descale_ptr + has_dropout_randval ? dropout_randval.data_ptr() : nullptr, + has_lse ? softmax_lse.data_ptr() : nullptr, + out.data_ptr(), + nullptr, // seqstart_q_ptr + nullptr, // seqstart_k_ptr + nullptr, // seqlen_q_ptr + nullptr, // seqlen_k_ptr + nullptr, // cu_seqlen_q_ptr + nullptr, // cu_seqlen_k_ptr + seqlen_q, + seqlen_k, + b, + seqlen_q, // max_seqlen_q + d, // hdim_q + d_v, // hdim_v + h, // nhead_q + h_k, // nhead_k + softmax_scale, // scale_s + 0.0, // logits_soft_cap + stride_q, + stride_k, + stride_v, + stride_bias, + stride_randval, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + 0, // nhead_stride_bias + nhead_stride_randval, + nhead_stride_lse, + nhead_stride_o, + batch_stride_q, + batch_stride_k, + batch_stride_v, + 0, // batch_stride_bias + batch_stride_randval, + batch_stride_lse, + batch_stride_o, + mask.left, + mask.right, + static_cast(mask.type), + 0, // min_seqlen_q + p_dropout, + has_dropout_randval, + drop_seed_offset}; } std::vector fmha_v3_fwd(at::Tensor &q, // [b, sq, hq, d] @@ -317,7 +318,7 @@ std::vector fmha_v3_fwd(at::Tensor &q, // [b, sq, hq, d] mask.type, bias_type, has_lse, - true, + quant_scale_enum::no_scale, how_v3_bf16_cvt); TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd"); } diff --git a/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu b/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu index faa1a662cb..8447c40e01 100644 --- a/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu +++ b/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu @@ -93,59 +93,60 @@ mha_fwd_args get_asm_mha_varlen_fwd_args(bool has_lse, bias_ptr = alibi_slopes.data_ptr(); stride_bias = alibi_slopes.dim() == 2 ? alibi_slopes.stride(0) : 0; } - + return mha_fwd_args{q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - bias_ptr, - has_dropout_randval ? dropout_randval.data_ptr() : nullptr, - has_lse ? softmax_lse.data_ptr() : nullptr, - out.data_ptr(), - cu_seqlens_q.data_ptr(), // seqstart_q_ptr (cumulative physical) - cu_seqlens_k.has_value() ? cu_seqlens_k.value().data_ptr() : nullptr, // seqstart_k_ptr - nullptr, // seqlen_q_ptr (per-sequence logical, not used here) - seqlens_k.has_value() ? seqlens_k.value().data_ptr() : nullptr, // seqlen_k_ptr - nullptr, // cu_seqlen_q_ptr (not used in this mode) - nullptr, // cu_seqlen_k_ptr (not used in this mode) - total_q, - total_k, - b, - max_seqlen_q, - d, // hdim_q - d_v, // hdim_v - h, // nhead_q - h_k, // nhead_k - softmax_scale, // scale_s - 1, // scale_p - 1, // scale_o - logits_soft_cap, - stride_q, - stride_k, - stride_v, - stride_bias, - stride_randval, - stride_o, - nhead_stride_q, - nhead_stride_k, - nhead_stride_v, - nhead_stride_bias, - nhead_stride_randval, - nhead_stride_lse, - nhead_stride_o, - batch_stride_q, - batch_stride_k, - batch_stride_v, - batch_stride_bias, - batch_stride_randval, - batch_stride_lse, - batch_stride_o, - mask.left, - mask.right, - static_cast(mask.type), - min_seqlen_q, - p_dropout, - has_dropout_randval, - drop_seed_offset}; + k.data_ptr(), + v.data_ptr(), + bias_ptr, + nullptr, // q_descale_ptr + nullptr, // k_descale_ptr + nullptr, // v_descale_ptr + has_dropout_randval ? dropout_randval.data_ptr() : nullptr, + has_lse ? softmax_lse.data_ptr() : nullptr, + out.data_ptr(), + cu_seqlens_q.data_ptr(), // seqstart_q_ptr (cumulative physical) + cu_seqlens_k.has_value() ? cu_seqlens_k.value().data_ptr() : nullptr, // seqstart_k_ptr + nullptr, // seqlen_q_ptr (per-sequence logical, not used here) + seqlens_k.has_value() ? seqlens_k.value().data_ptr() : nullptr, // seqlen_k_ptr + nullptr, // cu_seqlen_q_ptr (not used in this mode) + nullptr, // cu_seqlen_k_ptr (not used in this mode) + total_q, + total_k, + b, + max_seqlen_q, + d, // hdim_q + d_v, // hdim_v + h, // nhead_q + h_k, // nhead_k + softmax_scale, // scale_s + logits_soft_cap, + stride_q, + stride_k, + stride_v, + stride_bias, + stride_randval, + stride_o, + nhead_stride_q, + nhead_stride_k, + nhead_stride_v, + nhead_stride_bias, + nhead_stride_randval, + nhead_stride_lse, + nhead_stride_o, + batch_stride_q, + batch_stride_k, + batch_stride_v, + batch_stride_bias, + batch_stride_randval, + batch_stride_lse, + batch_stride_o, + mask.left, + mask.right, + static_cast(mask.type), + min_seqlen_q, + p_dropout, + has_dropout_randval, + drop_seed_offset}; } @@ -341,7 +342,7 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] aiter::ParsePhiloxCudaState, dim3(1), dim3(64), 0, 0, philox_args, rng_state_ptr); } std::optional seqlens_k = std::nullopt; - + if (max_seqlen_k > 0) { auto stream = at::hip::getCurrentHIPStream(); ck_tile::stream_config stream_config{stream}; @@ -383,6 +384,7 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] mask.type, bias_type, has_lse, + quant_scale_enum::no_scale, true, how_v3_bf16_cvt); TORCH_CHECK(t >= 0, "invalid argument for fmha_v3_varlen_fwd 3"); @@ -392,7 +394,7 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] out.zero_(); softmax_lse.fill_(std::numeric_limits::infinity()); } - + return {out, softmax_lse, p, rng_state}; } diff --git a/op_tests/test_mha_fp8.py b/op_tests/test_mha_fp8.py index 62319e623e..78377b6a17 100644 --- a/op_tests/test_mha_fp8.py +++ b/op_tests/test_mha_fp8.py @@ -15,10 +15,20 @@ def run_ck( v, causal=False, window_size=(-1, -1), # -1 means infinite context window + q_descale=None, + k_descale=None, + v_descale=None, ): if q.dtype == dtypes.fp8 and k.dtype == dtypes.fp8 and v.dtype == dtypes.fp8: return aiter.flash_attn_fp8_pertensor_func( - q, k, v, causal=causal, window_size=window_size + q, + k, + v, + q_descale, + k_descale, + v_descale, + causal=causal, + window_size=window_size, ) else: return aiter.flash_attn_func( @@ -98,12 +108,20 @@ def test_flash_attn_output( dtype=dtype, ) - # TODO - Support descale - q_quant, _ = per_tensor_quant(q, scale=torch.tensor(1), quant_dtype=quant_dtype) - k_quant, _ = per_tensor_quant(k, scale=torch.tensor(1), quant_dtype=quant_dtype) - v_quant, _ = per_tensor_quant(v, scale=torch.tensor(1), quant_dtype=quant_dtype) + q_quant, q_descale = per_tensor_quant(q, quant_dtype=quant_dtype) + k_quant, k_descale = per_tensor_quant(k, quant_dtype=quant_dtype) + v_quant, v_descale = per_tensor_quant(v, quant_dtype=quant_dtype) - out = run_ck(q_quant, k_quant, v_quant, causal, window_size) + out = run_ck( + q_quant, + k_quant, + v_quant, + causal, + window_size, + q_descale, + k_descale, + v_descale, + ) out_ref = run_ck(q, k, v, causal, window_size) max_diff = (out - out_ref).abs().max().item() @@ -137,7 +155,7 @@ def test_flash_attn_output( type=int, default=5, help="""Number of heads. Default is 5. - e.g.: -n 1""", + e.g.: -nk 1""", ) parser.add_argument( "-q", @@ -156,20 +174,12 @@ def test_flash_attn_output( e.g.: -k 1024""", ) parser.add_argument( - "-qk", - "--d_qk", + "-d", + "--d_qkv", type=int, default=128, help="""Dimension of query and key. Default is 128. - e.g.: -qk 256""", -) -parser.add_argument( - "-v", - "--d_v", - type=int, - default=128, - help="""Dimension of value. Default is 128. - e.g.: -v 256""", + e.g.: -d 128""", ) parser.add_argument( "-c", @@ -194,8 +204,8 @@ def test_flash_attn_output( args.nheads_k, args.seqlen_q, args.seqlen_k, - args.d_qk, - args.d_v, + args.d_qkv, + args.d_qkv, args.causal, args.local, ) diff --git a/op_tests/test_mha_varlen_fp8.py b/op_tests/test_mha_varlen_fp8.py index cf29db1038..faa5fefef0 100644 --- a/op_tests/test_mha_varlen_fp8.py +++ b/op_tests/test_mha_varlen_fp8.py @@ -29,12 +29,18 @@ def run_ck( min_seqlen_q, causal=False, window_size=(-1, -1), # -1 means infinite context window + q_descale=None, + k_descale=None, + v_descale=None, ): if q.dtype == dtypes.fp8 and k.dtype == dtypes.fp8 and v.dtype == dtypes.fp8: return aiter.flash_attn_varlen_fp8_pertensor_func( q, k, v, + q_descale, + k_descale, + v_descale, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, @@ -157,10 +163,9 @@ def test_flash_attn_varlen_output( k.requires_grad_(False) v.requires_grad_(False) - # TODO - Support descale - q_quant, _ = per_tensor_quant(q, scale=torch.tensor(1), quant_dtype=quant_dtype) - k_quant, _ = per_tensor_quant(k, scale=torch.tensor(1), quant_dtype=quant_dtype) - v_quant, _ = per_tensor_quant(v, scale=torch.tensor(1), quant_dtype=quant_dtype) + q_quant, q_descale = per_tensor_quant(q, quant_dtype=quant_dtype) + k_quant, k_descale = per_tensor_quant(k, quant_dtype=quant_dtype) + v_quant, v_descale = per_tensor_quant(v, quant_dtype=quant_dtype) out = run_ck( q_quant, @@ -173,6 +178,9 @@ def test_flash_attn_varlen_output( min_seqlen_q, causal, window_size, + q_descale, + k_descale, + v_descale, ) out_ref = run_ck( @@ -219,7 +227,7 @@ def test_flash_attn_varlen_output( type=int, default=5, help="""Number of heads. Default is 5. - e.g.: -n 1""", + e.g.: -nk 1""", ) parser.add_argument( "-q", @@ -238,20 +246,12 @@ def test_flash_attn_varlen_output( e.g.: -k 1024""", ) parser.add_argument( - "-qk", - "--d_qk", + "-d", + "--d_qkv", type=int, default=128, help="""Dimension of query and key. Default is 128. - e.g.: -qk 128""", -) -parser.add_argument( - "-v", - "--d_v", - type=int, - default=128, - help="""Dimension of value. Default is 128. - e.g.: -v 128""", + e.g.: -d 128""", ) parser.add_argument( "-msq", @@ -285,8 +285,8 @@ def test_flash_attn_varlen_output( args.nheads_k, args.seqlen_q, args.seqlen_k, - args.d_qk, - args.d_v, + args.d_qkv, + args.d_qkv, args.min_seqlen_q, args.causal, args.local,