Integration of SageAttn & QkBf16PvFp8 Attn

Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com>

Author: xrq-phys

csrc/trtllm_fmha_kernel_launcher.cu

@@ -44,20 +44,25 @@ enum class TllmPagedAttentionMode {
 
 class TllmGenFmhaRunnerCache {
  public:
-  using Key = std::tuple<Data_type, Data_type, Data_type>;
+  using Key = std::tuple<Data_type, Data_type, Data_type, Data_type, int, int, int, int>;
 
-  static std::shared_ptr<TllmGenFmhaRunner> get(Data_type q_data_type, Data_type kv_data_type,
-                                                Data_type o_data_type) {
+  static std::shared_ptr<TllmGenFmhaRunner> get(Data_type q_data_type, Data_type k_data_type,
+                                                Data_type v_data_type, Data_type o_data_type,
+                                                int num_elts_sage_q = 0, int num_elts_sage_k = 0,
+                                                int num_elts_sage_p = 0, int num_elts_sage_v = 0) {
     static std::unordered_map<Key, std::shared_ptr<TllmGenFmhaRunner>, KeyHash> cache;
     static std::mutex cache_mutex;
-    Key key = std::make_tuple(q_data_type, kv_data_type, o_data_type);
+    Key key = std::make_tuple(q_data_type, k_data_type, v_data_type, o_data_type, num_elts_sage_q,
+                              num_elts_sage_k, num_elts_sage_p, num_elts_sage_v);
 
     std::lock_guard<std::mutex> lock(cache_mutex);
     auto it = cache.find(key);
     if (it != cache.end()) {
       return it->second;
     } else {
-      auto runner = std::make_shared<TllmGenFmhaRunner>(q_data_type, kv_data_type, o_data_type);
+      auto runner = std::make_shared<TllmGenFmhaRunner>(
+          q_data_type, k_data_type, v_data_type, o_data_type, num_elts_sage_q, num_elts_sage_k,
+          num_elts_sage_p, num_elts_sage_v);
       cache.emplace(key, runner);
       return runner;
     }
@@ -68,7 +73,10 @@ class TllmGenFmhaRunnerCache {
     std::size_t operator()(const Key& k) const {
       return std::hash<int>()(static_cast<int>(std::get<0>(k))) ^
              (std::hash<int>()(static_cast<int>(std::get<1>(k))) << 1) ^
-             (std::hash<int>()(static_cast<int>(std::get<2>(k))) << 2);
+             (std::hash<int>()(static_cast<int>(std::get<2>(k))) << 2) ^
+             (std::hash<int>()(static_cast<int>(std::get<3>(k))) << 3) ^
+             (std::hash<int>()(std::get<4>(k)) << 4) ^ (std::hash<int>()(std::get<5>(k)) << 5) ^
+             (std::hash<int>()(std::get<6>(k)) << 6) ^ (std::hash<int>()(std::get<7>(k)) << 7);
     }
   };
 };
@@ -96,7 +104,9 @@ void trtllm_paged_attention_launcher(
     FLASHINFER_ERROR(err_msg.str());
   }
 
-  auto fmha_runner = TllmGenFmhaRunnerCache::get(q_data_type, kv_data_type, o_data_type);
+  // For paged attention, K and V have the same dtype (kv_data_type).
+  auto fmha_runner =
+      TllmGenFmhaRunnerCache::get(q_data_type, kv_data_type, kv_data_type, o_data_type);
   TllmGenFmhaRunnerParams runner_params;
 
   // Common params
@@ -225,6 +235,8 @@ inline Data_type dl_dtype_to_tllm_data_type(const DLDataType dtype) {
     return Data_type::DATA_TYPE_E4M3;
   } else if (dtype == dl_float8_e5m2) {
     return Data_type::DATA_TYPE_E5M2;
+  } else if (dtype == dl_int8) {
+    return Data_type::DATA_TYPE_INT8;
   } else if (dtype == dl_uint8) {
     // fp4 tensor is not supported in torch and use uint8_t as container.
     return Data_type::DATA_TYPE_E2M1;
@@ -493,22 +505,27 @@ void trtllm_paged_attention_context(
 void trtllm_ragged_attention_launcher(
     void* out, void* query, void* key, void* value, void* workspace_buffer, int* seq_lens,
     int* cum_seq_lens_q, int* cum_seq_lens_kv, float* attention_sinks, float* lse,
-    Data_type q_data_type, Data_type kv_data_type, Data_type o_data_type, int64_t max_q_len,
-    int64_t max_kv_len, int64_t num_qo_heads, int64_t num_kv_heads, int64_t head_dim_qk,
-    int64_t head_dim_v, int64_t sum_seq_q, int64_t sum_seq_kv, double bmm1_scale, double bmm2_scale,
-    const float* bmm1_scale_log2_ptr, const float* bmm2_scale_ptr, double o_sf_scale,
-    int64_t batch_size, int64_t window_left, int64_t sm_count, bool enable_pdl, bool is_causal,
-    int64_t k_stride_keys_values, int64_t k_stride_heads, int64_t k_stride_batch,
-    int64_t v_stride_keys_values, int64_t v_stride_heads, int64_t v_stride_batch,
-    float skip_softmax_threshold_scale_factor, bool skips_softmax, int64_t workspace_size,
+    Data_type q_data_type, Data_type k_data_type, Data_type v_data_type, Data_type o_data_type,
+    int64_t max_q_len, int64_t max_kv_len, int64_t num_qo_heads, int64_t num_kv_heads,
+    int64_t head_dim_qk, int64_t head_dim_v, int64_t sum_seq_q, int64_t sum_seq_kv,
+    double bmm1_scale, double bmm2_scale, const float* bmm1_scale_log2_ptr,
+    const float* bmm2_scale_ptr, double o_sf_scale, int64_t batch_size, int64_t window_left,
+    int64_t sm_count, bool enable_pdl, bool is_causal, int64_t k_stride_keys_values,
+    int64_t k_stride_heads, int64_t k_stride_batch, int64_t v_stride_keys_values,
+    int64_t v_stride_heads, int64_t v_stride_batch, float skip_softmax_threshold_scale_factor,
+    bool skips_softmax, int64_t workspace_size, const float* sage_attn_sfs_q,
+    const float* sage_attn_sfs_k, const float* sage_attn_sfs_p, const float* sage_attn_sfs_v,
+    int num_elts_sage_q, int num_elts_sage_k, int num_elts_sage_p, int num_elts_sage_v,
     cudaStream_t stream) {
   if (num_qo_heads % num_kv_heads != 0) {
     std::ostringstream err_msg;
     err_msg << "num_qo_heads must be a multiple of num_kv_heads, got num_kv_heads: " << num_kv_heads
             << " and num_qo_heads: " << num_qo_heads;
     FLASHINFER_ERROR(err_msg.str());
   }
-  auto fmha_runner = TllmGenFmhaRunnerCache::get(q_data_type, kv_data_type, o_data_type);
+  auto fmha_runner = TllmGenFmhaRunnerCache::get(q_data_type, k_data_type, v_data_type, o_data_type,
+                                                 num_elts_sage_q, num_elts_sage_k, num_elts_sage_p,
+                                                 num_elts_sage_v);
   TllmGenFmhaRunnerParams runner_params;
 
   runner_params.qPtr = query;
@@ -576,6 +593,12 @@ void trtllm_ragged_attention_launcher(
   runner_params.mSkipsSoftmaxWhenPossible = skips_softmax;
   runner_params.mSkipSoftmaxThresholdScaleFactor = skip_softmax_threshold_scale_factor;
 
+  // SageAttention scaling factors.
+  runner_params.ptrSageAttnSfsQ = sage_attn_sfs_q;
+  runner_params.ptrSageAttnSfsK = sage_attn_sfs_k;
+  runner_params.ptrSageAttnSfsP = sage_attn_sfs_p;
+  runner_params.ptrSageAttnSfsV = sage_attn_sfs_v;
+
   auto [foundKernels, kinfo] = fmha_runner->isSupportedWithInfo(runner_params);
   if (!foundKernels) {
     std::ostringstream err_msg;
@@ -586,16 +609,18 @@ void trtllm_ragged_attention_launcher(
   fmha_runner->run(runner_params);
 }
 
-void trtllm_ragged_attention(TensorView out, TensorView query, TensorView key, TensorView value,
-                             TensorView workspace_buffer, TensorView seq_lens, int64_t max_q_len,
-                             int64_t max_kv_len, Variant<double, ffi::Tensor> bmm1_scale,
-                             Variant<double, ffi::Tensor> bmm2_scale, double o_sf_scale,
-                             int64_t batch_size, int64_t window_left, TensorView cum_seq_lens_q,
-                             TensorView cum_seq_lens_kv, int64_t sm_count, bool enable_pdl,
-                             bool is_causal, int64_t workspace_size,
-                             Optional<TensorView> attention_sinks,
-                             Optional<float> skip_softmax_threshold_scale_factor,
-                             Optional<TensorView> lse) {
+void trtllm_ragged_attention(
+    TensorView out, TensorView query, TensorView key, TensorView value, TensorView workspace_buffer,
+    TensorView seq_lens, int64_t max_q_len, int64_t max_kv_len,
+    Variant<double, ffi::Tensor> bmm1_scale, Variant<double, ffi::Tensor> bmm2_scale,
+    double o_sf_scale, int64_t batch_size, int64_t window_left, TensorView cum_seq_lens_q,
+    TensorView cum_seq_lens_kv, int64_t sm_count, bool enable_pdl, bool is_causal,
+    int64_t workspace_size, Optional<TensorView> attention_sinks,
+    Optional<float> skip_softmax_threshold_scale_factor, Optional<TensorView> lse,
+    Optional<TensorView> sage_attn_sfs_q, Optional<TensorView> sage_attn_sfs_k,
+    Optional<TensorView> sage_attn_sfs_p, Optional<TensorView> sage_attn_sfs_v,
+    int64_t num_elts_per_sage_attn_blk_q, int64_t num_elts_per_sage_attn_blk_k,
+    int64_t num_elts_per_sage_attn_blk_p, int64_t num_elts_per_sage_attn_blk_v) {
   float* attention_sinks_ptr = nullptr;
   if (attention_sinks.has_value()) {
     TVM_FFI_ICHECK_EQ(attention_sinks.value().dtype(), dl_float32)
@@ -613,7 +638,8 @@ void trtllm_ragged_attention(TensorView out, TensorView query, TensorView key, T
   TVM_FFI_ICHECK_EQ(value.ndim(), 3) << "value must be a 3D tensor";
 
   auto q_data_type = dl_dtype_to_tllm_data_type(query.dtype());
-  auto kv_data_type = dl_dtype_to_tllm_data_type(key.dtype());
+  auto k_data_type = dl_dtype_to_tllm_data_type(key.dtype());
+  auto v_data_type = dl_dtype_to_tllm_data_type(value.dtype());
   auto o_data_type = dl_dtype_to_tllm_data_type(out.dtype());
   const auto stream = get_stream(query.device());
   int num_qo_heads = query.size(1);
@@ -629,6 +655,20 @@ void trtllm_ragged_attention(TensorView out, TensorView query, TensorView key, T
   int v_stride_heads = value.stride(1);
   int v_stride_batch = value.numel();
 
+  // SageAttention scaling factor pointers.
+  const float* sage_attn_sfs_q_ptr =
+      sage_attn_sfs_q.has_value() ? static_cast<const float*>(sage_attn_sfs_q.value().data_ptr())
+                                  : nullptr;
+  const float* sage_attn_sfs_k_ptr =
+      sage_attn_sfs_k.has_value() ? static_cast<const float*>(sage_attn_sfs_k.value().data_ptr())
+                                  : nullptr;
+  const float* sage_attn_sfs_p_ptr =
+      sage_attn_sfs_p.has_value() ? static_cast<const float*>(sage_attn_sfs_p.value().data_ptr())
+                                  : nullptr;
+  const float* sage_attn_sfs_v_ptr =
+      sage_attn_sfs_v.has_value() ? static_cast<const float*>(sage_attn_sfs_v.value().data_ptr())
+                                  : nullptr;
+
   auto maybe_bmm1_scale_value = bmm1_scale.as<double>();
   auto maybe_bmm2_scale_value = bmm2_scale.as<double>();
   auto maybe_bmm1_scale_log2_tensor = bmm1_scale.as<ffi::Tensor>();
@@ -658,12 +698,17 @@ void trtllm_ragged_attention(TensorView out, TensorView query, TensorView key, T
       out.data_ptr(), query.data_ptr(), key.data_ptr(), value.data_ptr(),
       workspace_buffer.data_ptr(), static_cast<int*>(seq_lens.data_ptr()),
       static_cast<int*>(cum_seq_lens_q.data_ptr()), static_cast<int*>(cum_seq_lens_kv.data_ptr()),
-      attention_sinks_ptr, lse_ptr, q_data_type, kv_data_type, o_data_type, max_q_len, max_kv_len,
-      num_qo_heads, num_kv_heads, head_dim_qk, head_dim_v, sum_seq_q, sum_seq_kv, bmm1_scale_value,
-      bmm2_scale_value, bmm1_scale_log2_ptr, bmm2_scale_ptr, o_sf_scale, batch_size, window_left,
-      sm_count, enable_pdl, is_causal, k_stride_keys_values, k_stride_heads, k_stride_batch,
-      v_stride_keys_values, v_stride_heads, v_stride_batch,
-      skip_softmax_threshold_scale_factor_value, skips_softmax, workspace_size, stream);
+      attention_sinks_ptr, lse_ptr, q_data_type, k_data_type, v_data_type, o_data_type, max_q_len,
+      max_kv_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_v, sum_seq_q, sum_seq_kv,
+      bmm1_scale_value, bmm2_scale_value, bmm1_scale_log2_ptr, bmm2_scale_ptr, o_sf_scale,
+      batch_size, window_left, sm_count, enable_pdl, is_causal, k_stride_keys_values,
+      k_stride_heads, k_stride_batch, v_stride_keys_values, v_stride_heads, v_stride_batch,
+      skip_softmax_threshold_scale_factor_value, skips_softmax, workspace_size, sage_attn_sfs_q_ptr,
+      sage_attn_sfs_k_ptr, sage_attn_sfs_p_ptr, sage_attn_sfs_v_ptr,
+      static_cast<int>(num_elts_per_sage_attn_blk_q),
+      static_cast<int>(num_elts_per_sage_attn_blk_k),
+      static_cast<int>(num_elts_per_sage_attn_blk_p),
+      static_cast<int>(num_elts_per_sage_attn_blk_v), stream);
 }
 
 namespace trtllm_cubin_loader {

flashinfer/artifacts.py

@@ -135,7 +135,7 @@ class ArtifactPath:
     When compiling new cubins for backend directories, update the corresponding path.
     """
 
-    TRTLLM_GEN_FMHA: str = "82f4c77d9cf83e3fcf105feda4ce3445100ab491/fmha/trtllm-gen/"
+    TRTLLM_GEN_FMHA: str = "ce9168f3a3f60ffaccbaf6b2ee23642d8207b3b7/fmha/trtllm-gen/"
     TRTLLM_GEN_BMM: str = (
         "39a9d28268f43475a757d5700af135e1e58c9849/batched_gemm-5ee61af-2b9855b/"
     )
@@ -155,7 +155,7 @@ class CheckSumHash:
     """
 
     TRTLLM_GEN_FMHA: str = (
-        "56c95fbe5d1b5d0d9ded7706e1c0b7ebf0582d9cfd2f9382acd878b6b9d58c89"
+        "6a05a464e0101612a7598273bda641633b6db15abd76bd7f9a94a07646c7127c"
     )
     TRTLLM_GEN_BMM: str = (
         "db06db7f36a2a9395a2041ff6ac016fe664874074413a2ed90797f91ef17e0f6"

flashinfer/prefill.py

@@ -3692,6 +3692,13 @@ def trtllm_ragged_attention_deepseek(
     skip_softmax_threshold_scale_factor: Optional[float] = None,
     out: Optional[torch.Tensor] = None,
     lse: Optional[torch.Tensor] = None,
+    sage_attn_sfs: Tuple[
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+    ] = (None, None, None, None),
+    num_elts_per_sage_attn_blk: Tuple[int, int, int, int] = (0, 0, 0, 0),
 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
     """
     Parameters
@@ -3766,7 +3773,7 @@ def trtllm_ragged_attention_deepseek(
     if out is None:
         # FP8 inputs produce bfloat16 output by default (TRT-LLM kernels
         # do not support FP8 output for ragged attention)
-        if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+        if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2, torch.int8):
             out_dtype = torch.bfloat16
         else:
             out_dtype = query.dtype
@@ -3806,6 +3813,8 @@ def trtllm_ragged_attention_deepseek(
         assert bmm2_scale.dtype == torch.float32
 
     workspace_size = workspace_buffer.numel() * workspace_buffer.element_size()
+    sage_attn_sfs_q, sage_attn_sfs_k, sage_attn_sfs_p, sage_attn_sfs_v = sage_attn_sfs
+    num_elts_sage_q, num_elts_sage_k, num_elts_sage_p, num_elts_sage_v = num_elts_per_sage_attn_blk
     run_func(
         out,
         query,
@@ -3829,6 +3838,14 @@ def trtllm_ragged_attention_deepseek(
         attention_sinks,
         skip_softmax_threshold_scale_factor,
         lse,
+        sage_attn_sfs_q,
+        sage_attn_sfs_k,
+        sage_attn_sfs_p,
+        sage_attn_sfs_v,
+        num_elts_sage_q,
+        num_elts_sage_k,
+        num_elts_sage_p,
+        num_elts_sage_v,
     )
     if return_lse:
         assert lse is not None, (