Enable gfx950 FP8 HD128/256 fused-attn ASM forward from aiter

transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp

@@ -16,9 +16,10 @@ namespace ck_fused_attn{
 
 // input qkv dtypes
 enum class DType {
-  kFloat16    = 0,  /*!< 16-bit float (E5M10) */
-  kBFloat16   = 1,  /*!< 16-bit bfloat (E8M7) */
-  kNumTypes         /*!< Number of supported types */
+  kFloat16  = 0,  /*!< 16-bit float (E5M10) */
+  kBFloat16 = 1,  /*!< 16-bit bfloat (E8M7) */
+  kFloat8   = 2,  /*!< 8-bit float (E4M3) input */
+  kNumTypes       /*!< Number of supported types */
 };
 
 // keep sync with mask_enum in mask.hpp
@@ -99,6 +100,13 @@ struct CKAttnFwdArgs : CKAttnCommonArgs {
   // Output (writable)
   void* o_ptr = nullptr;
   void* lse_ptr = nullptr;
+  // For FP8 E4M3 input, the AITER fp8bf16 kernels require BF16.  
+  DType o_dtype = DType::kNumTypes;
+
+  // Per-tensor inverse scales for FP8 Q/K/V: one FP32 value per tensor.
+  const void* q_descale_ptr = nullptr;
+  const void* k_descale_ptr = nullptr;
+  const void* v_descale_ptr = nullptr;
 
   // V3 ASM kernel selection
   bool uses_fwd_v3 = false;

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp

@@ -292,11 +292,14 @@ aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){
   fmha_args.is_group_mode   = args.is_group_mode();
   fmha_args.bias_type       = static_cast<int>(bias_type);
   fmha_args.has_lse         = args.lse_ptr!=nullptr;
-  fmha_args.qscale_type     = static_cast<int>(quant_scale_enum::no_scale);
+  fmha_args.qscale_type     = static_cast<int>(
+      args.dtype == DType::kFloat8
+          ? quant_scale_enum::pertensor
+          : quant_scale_enum::no_scale);
   fmha_args.has_sink        = args.has_sink;
-  fmha_args.q_descale_ptr    = nullptr;
-  fmha_args.k_descale_ptr    = nullptr;
-  fmha_args.v_descale_ptr    = nullptr;
+  fmha_args.q_descale_ptr    = args.q_descale_ptr;
+  fmha_args.k_descale_ptr    = args.k_descale_ptr;
+  fmha_args.v_descale_ptr    = args.v_descale_ptr;
   // sink_size is CK's StreamingLLM sink *prefix width* in key columns, which is a
   // different feature from the learnable softmax offset NVTE asks for here. The
   // offset only needs has_sink + sink_ptr (CK folds it into the softmax
@@ -313,9 +316,10 @@ aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){
   return fmha_args;
 }
 
-//Split-KV and Sink dispatchers do not honor v3_api_check
-//so we cannot rely on the AITER backend to tell us whether v3 is available.
-//Currently this method is not used so just opt it out.
+// Split-KV and sink dispatchers do not honor v3_api_check, so a naive
+// mha_fwd(..., v3_api_check=true) can still launch a real kernel when those
+// paths are compiled in. Unused today; the FP8 probe in ck_attn_fwd clears
+// sink/split-KV fields before using the same dry-run.
 #if !(FA_WITH_SINK || FA_WITH_NATIVE_SPLITKV)
 // Probe whether AITER's v3 (asm) forward path will run for this config, without
 // launching any kernel. Builds the same args as ck_attn_fwd and relies on AITER's
@@ -340,6 +344,29 @@ hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream){
 
   aiter::mha_fwd_args fmha_args = build_fwd_fmha_args(args);
 
+  // This PR wires AITER gfx950 FP8 FMHA through the v3 ASM kernels only.
+  // Upstream CK has an FP8 tile path, but it is not enabled here, so a failed
+  // ASM probe must not fall through to an unvalidated CK launch. Enabling CK
+  // tile FP8 is a separate change.
+  //
+  // QoLA builds with FA_WITH_SINK and FA_WITH_NATIVE_SPLITKV, and those
+  // dispatchers run before fmha_fwd_v3 without honoring v3_api_check. Clear
+  // sink/split-KV fields so the probe reaches the v3 dry-run and cannot launch
+  // a real kernel (same reason ck_attn_fwd_uses_v3 is compiled out).
+  if(args.dtype == DType::kFloat8){
+    aiter::mha_fwd_args probe_args = fmha_args;
+    probe_args.v3_api_check = true;
+    probe_args.num_splits = 0;
+    probe_args.splitkv_workspace_ptr = nullptr;
+    probe_args.has_sink = false;
+    probe_args.sink_ptr = nullptr;
+    ck_tile::stream_config probe_config{nullptr, false, false};
+    if(QOLA_NS(mha_fwd)(probe_args, probe_config) != 1){
+      throw std::runtime_error(
+        "ck_fused_attn fwd: no AITER fp8bf16 ASM kernel for this configuration.");
+    }
+  }
+
   if(const char* env_p = std::getenv("NVTE_CK_RUNTIME_MAX_SEQLEN")){
     if(args.is_group_mode() && std::string(env_p) == "1"){
       if(log_file){

transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp

@@ -106,14 +106,18 @@ std::ostream* get_ck_log_stream() {
   return log_stream;
 }
 
+// Return AITER's data_type kernel-selection tag. The tag may encode a
+// combination of input and output types rather than the input dtype alone.
 std::string get_data_type_str(DType dtype){
   std::string data_type_str;
   if(dtype==DType::kFloat16){
     data_type_str = "fp16";
   }else if(dtype==DType::kBFloat16){
     data_type_str = "bf16";
+  }else if(dtype==DType::kFloat8){
+    data_type_str = "fp8bf16";
+    // "fp8bf16" selects FP8 input with BF16 output; "fp8" selects FP8 output.
   }else{
-    //TODO: better error out system
     throw std::runtime_error("Invalid dtype in ck_fused_attn.");
   }
   return data_type_str;

transformer_engine/common/fused_attn_rocm/fused_attn.cpp

@@ -299,6 +299,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
 
   // first check whether ck can be used, then check aotriton
   if(nvte_fused_attn_ck && fused_attn_rocm::is_ck_backend_supported(
+        is_training,
         q_dtype,
         kv_dtype,
         qkv_layout,

transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp

@@ -21,6 +21,7 @@ namespace fused_attn_rocm {
 // check the fused attn config to see whether it's ck backend supported
 // single filtering followed by joint filtering
 bool is_ck_backend_supported(
+  bool is_training,
   NVTEDType q_dtype,
   NVTEDType kv_dtype,
   NVTE_QKV_Layout qkv_layout,
@@ -49,11 +50,13 @@ bool is_ck_backend_supported(
     return false;
   }
 
-  // filter based on data type
-  // Q and KV must have the same data type, in fp16 or bf16
-  if((q_dtype!=kv_dtype) || !((q_dtype==NVTEDType::kNVTEFloat16) || (q_dtype == NVTEDType::kNVTEBFloat16))){
+  const bool is_fp8 = q_dtype == NVTEDType::kNVTEFloat8E4M3;
+  const bool is_16bit =
+      q_dtype == NVTEDType::kNVTEFloat16 || q_dtype == NVTEDType::kNVTEBFloat16;
+  //  q and kv have to have same dtype: FP16/BF16  or FP8 E4M3 
+  if(q_dtype != kv_dtype || (!is_16bit && !is_fp8)){
     if(nvte_log_ck_config){
-      std::cout<<"q, k, v data type has to be fp16 or bf16"<<std::endl;
+      std::cout<<"q, k, v data type has to be fp16, bf16, or fp8 e4m3"<<std::endl;
     }
     return false;
   }
@@ -136,6 +139,33 @@ bool is_ck_backend_supported(
     }
     return false;
   }
+  if(is_fp8){
+    // FP8 is wired to AITER v3 ASM only. CK tile FP8 is not enabled, so turning
+    // ASM off must report unsupported and let the caller select NVTE_No_Backend.
+    // A later PR that enables CK FP8 can fall back to that path here instead.
+    if(getenv<int>("NVTE_CK_USES_FWD_V3", 1) == 0){
+      if(nvte_log_ck_config){
+        std::cout<<"AITER FP8 ASM fwd is disabled by NVTE_CK_USES_FWD_V3=0"<<std::endl;
+      }
+      return false;
+    }
+    const bool fp8_hdim_supported =
+        (head_dim_qk == 128 && head_dim_v == 128) ||
+        (head_dim_qk == 256 && head_dim_v == 256);
+    const bool fp8_window_supported =
+        is_causal ? (window_size_left == -1 && window_size_right == 0)
+                  : (window_size_left == -1 && window_size_right == -1);
+    if(cuda::sm_arch() != 95 || is_training || !fp8_hdim_supported ||
+       !fp8_window_supported || dropout != 0.f ||
+       bias_type != NVTE_Bias_Type::NVTE_NO_BIAS ||
+       (is_padding && !is_ragged)){
+      if(nvte_log_ck_config){
+        std::cout<<"AITER FP8 ASM fwd requires gfx950 inference, hd128/hd256, "
+                   "no bias/dropout/sliding-window, and THD for padding masks"<<std::endl;
+      }
+      return false;
+    }
+  }
   return true;
 #else
   return false;
@@ -176,6 +206,7 @@ ck_fused_attn::DType nvte_to_ck_dtype(DType t_dtype){
 #define CAST_TYPE(aname, dtname) if (t_dtype == DType::aname) return ck_fused_attn::DType::dtname
   CAST_TYPE(kFloat16, kFloat16);
   CAST_TYPE(kBFloat16, kBFloat16);
+  CAST_TYPE(kFloat8E4M3, kFloat8);
   return ck_fused_attn::DType::kNumTypes;
 #undef CAST_TYPE
 }
@@ -490,7 +521,8 @@ void fused_attn_ck_fwd_impl(
   void* devPtrDropoutSeed, void* devPtrDropoutOffset,
   void* devPtrCuSeqlensQ, void* devPtrCuSeqlensKV,
   void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV,
-  DType dtype,
+  void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV,
+  DType dtype, DType output_dtype,
   void *workspace, 
   size_t *workspace_size,
   cudaStream_t stream){
@@ -520,11 +552,18 @@ void fused_attn_ck_fwd_impl(
   size_t q_storage_bytes = max_tokens_q*h*d_qk*nvte_dtype_size(dtype);
   size_t k_storage_bytes = max_tokens_kv*hg*d_qk*nvte_dtype_size(dtype);
   size_t v_storage_bytes = max_tokens_kv*hg*d_v*nvte_dtype_size(dtype);
-  size_t o_storage_bytes = max_tokens_q*h*d_v*nvte_dtype_size(dtype);
+  size_t o_storage_bytes = max_tokens_q*h*d_v*nvte_dtype_size(output_dtype);
 
   // Reserve workspace chunks. Same allocation sequence runs in sizing mode
   // (planner returns nullptr, accumulates total) and execution mode.
   WorkspacePlanner planner(workspace);
+  // A real FP8 launch must carry the inverse scale with each input. Workspace
+  // sizing passes null tensors, so that pass does not require scales.
+  if(dtype == DType::kFloat8E4M3 && !planner.is_sizing()){
+    NVTE_CHECK(devPtrDescaleQ != nullptr && devPtrDescaleK != nullptr &&
+                   devPtrDescaleV != nullptr,
+               "AITER FP8 ASM fused attention requires one FP32 scale_inv for each of Q, K, and V.");
+  }
 
   if(ck_small_seq_env_enabled) {
     if(cuda::sm_arch() == 94 || cuda::sm_arch() == 95) {
@@ -589,6 +628,10 @@ void fused_attn_ck_fwd_impl(
   // Common fields filled here; mode-specific fields are overwritten below.
   ck_fused_attn::CKAttnFwdArgs ck_args;
   ck_args.dtype = nvte_to_ck_dtype(dtype);
+  ck_args.o_dtype = nvte_to_ck_dtype(output_dtype);
+  ck_args.q_descale_ptr = devPtrDescaleQ;
+  ck_args.k_descale_ptr = devPtrDescaleK;
+  ck_args.v_descale_ptr = devPtrDescaleV;
   ck_args.b = b; ck_args.h = h; ck_args.hg = hg;
   ck_args.s_q = s_q; ck_args.s_kv = s_kv; ck_args.d_qk = d_qk; ck_args.d_v = d_v;
   ck_args.is_training = is_training;
@@ -611,7 +654,10 @@ void fused_attn_ck_fwd_impl(
     ck_args.cu_seqlen_q_ptr = devPtrCuSeqlensQ;
   }
 
-  ck_args.num_splits = (ck_args.uses_fwd_v3 && !has_sink) ? ck_attn_fwd_num_splits(ck_args) : 0;
+  ck_args.num_splits =
+      (ck_args.uses_fwd_v3 && !has_sink && dtype != DType::kFloat8E4M3)
+          ? ck_attn_fwd_num_splits(ck_args)
+          : 0;
   if (ck_args.num_splits > 0)
   {
     size_t splitkv_workspace_bytes = ck_attn_fwd_workspace_size(ck_args);
@@ -1352,11 +1398,41 @@ void fused_attn_ck_fwd(
 
 #ifdef USE_FUSED_ATTN_CK
   const DType QKV_type = input_Q->data.dtype;
+  const DType O_type = output_O->data.dtype;
+  const bool is_fp8 = QKV_type == DType::kFloat8E4M3;
+
+  NVTE_CHECK(input_K->data.dtype == QKV_type && input_V->data.dtype == QKV_type,
+             "ROCm fused attention requires Q, K, and V to have the same dtype.");
+  if(is_fp8){
+    NVTE_CHECK(!is_training, "AITER FP8 ASM fused attention is forward-inference only.");
+    // FP8 E4M3 input is accepted only with BF16 output for AITER fp8bf16 kernels.
+    NVTE_CHECK(O_type == DType::kBFloat16,
+               "AITER fp8bf16 fused attention requires BF16 output.");
+    NVTE_CHECK(input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING &&
+                   input_K->scaling_mode == NVTE_DELAYED_TENSOR_SCALING &&
+                   input_V->scaling_mode == NVTE_DELAYED_TENSOR_SCALING,
+               "AITER FP8 ASM fused attention supports delayed tensor scaling only.");
+    auto valid_descale = [](const Tensor* tensor) {
+      return tensor->scale_inv.dptr != nullptr &&
+             tensor->scale_inv.dtype == DType::kFloat32 &&
+             tensor->scale_inv.numel() == 1;
+    };
+    const bool all_descales_present =
+        valid_descale(input_Q) && valid_descale(input_K) && valid_descale(input_V);
+    // Workspace sizing uses null tensors. A real launch must carry the scales.
+    if(input_Q->data.dptr != nullptr){
+      NVTE_CHECK(all_descales_present,
+                 "AITER FP8 ASM fused attention requires one FP32 scale_inv for each of Q, K, and V.");
+    }
+  }
 
   void *devPtrQ = input_Q->data.dptr;
   void *devPtrK = input_K->data.dptr;
   void *devPtrV = input_V->data.dptr;
   void *devPtrO = output_O->data.dptr;
+  void *devPtrDescaleQ = is_fp8 ? input_Q->scale_inv.dptr : nullptr;
+  void *devPtrDescaleK = is_fp8 ? input_K->scale_inv.dptr : nullptr;
+  void *devPtrDescaleV = is_fp8 ? input_V->scale_inv.dptr : nullptr;
   void *devPtrS = nullptr;
   void *devPtrBias = nullptr;
   size_t bias_b = 0;
@@ -1444,7 +1520,8 @@ void fused_attn_ck_fwd(
     reinterpret_cast<void *>(reinterpret_cast<uint64_t *>(rng_state->data.dptr) + 1),
     devPtrCuSeqlensQ, devPtrCuSeqlensKV,
     devPtrSeqOffsetsQ, devPtrSeqOffsetsKV,
-    QKV_type,
+    devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV,
+    QKV_type, O_type,
     workspace->data.dptr,
     &workspace_size,
     stream);

transformer_engine/common/fused_attn_rocm/fused_attn_ck.h

@@ -17,6 +17,7 @@ namespace transformer_engine {
 namespace fused_attn_rocm {
 // check the fused attn config to see whether it's ck backend supported
 bool is_ck_backend_supported(
+  bool is_training,
   NVTEDType q_dtype,
   NVTEDType kv_dtype,
   NVTE_QKV_Layout qkv_layout,