Merge branch 'main' into fmhav2-bench

Author: jimmyzho

benchmarks/routines/attention.py

@@ -1305,6 +1305,7 @@ def run_backend_wrapper(
                 batch_size=batch_size,
                 cum_seq_lens_q=qo_indptr,
                 cum_seq_lens_kv=kv_indptr,
+                causal=causal,
                 kv_cache_sf=kv_cache_sf,
             )
         elif backend == "cudnn-native":

benchmarks/routines/gemm.py

@@ -148,6 +148,7 @@ def parse_gemm_args(line, parser):
             "trtllm",
             "cutlass",
             "tgv",
+            "cublaslt",
             "cute-dsl",
             "b12x",
             "auto",
@@ -1585,7 +1586,7 @@ def testMmBf16(args):
     use_pdl = getattr(args, "enable_pdl", False)
     is_cuda_graph_compatible = not args.no_cuda_graph
     run_refcheck = args.refcheck
-    autotune_supported_backends = ["cudnn", "cutlass", "tgv", "auto"]
+    autotune_supported_backends = ["cudnn", "cutlass", "tgv", "cublaslt", "auto"]
     res = []
 
     out_dtype = dtype_str_to_torch_dtype(args.out_dtype)
@@ -1650,7 +1651,7 @@ def testMmBf16(args):
         return res
 
     def run_backend(backend, a, b, bias, use_pdl, out_dtype):
-        if backend in ["cudnn", "cutlass", "tgv", "auto"]:
+        if backend in ["cudnn", "cutlass", "tgv", "cublaslt", "auto"]:
             return flashinfer.mm_bf16(
                 a=a,
                 b=b,

csrc/bmm_fp8.cu

@@ -49,7 +49,8 @@ void bmm_fp8(TensorView A, TensorView B, TensorView D, TensorView A_scale, Tenso
         auto stream = get_stream(A.device());
 
         auto status = flashinfer::bmm_fp8::bmm_fp8_internal_cublaslt(
-            workspace_buffer.data_ptr(), workspace_buffer.numel(),
+            workspace_buffer.data_ptr(),
+            workspace_buffer.numel() * get_element_size(workspace_buffer),
             static_cast<b_type*>(B.data_ptr()), static_cast<a_type*>(A.data_ptr()),
             static_cast<d_type*>(D.data_ptr()), batch_size, n, m, k,
             static_cast<float*>(B_scale.data_ptr()), static_cast<float*>(A_scale.data_ptr()),
@@ -61,3 +62,91 @@ void bmm_fp8(TensorView A, TensorView B, TensorView D, TensorView A_scale, Tenso
     });
   });
 }
+
+int64_t bmm_fp8_get_algos(TensorView A, TensorView B, TensorView D, TensorView A_scale,
+                          TensorView B_scale, TensorView workspace_buffer, int64_t cublas_handle,
+                          TensorView algo_buffer) {
+  CHECK_CUDA(A);
+  CHECK_CUDA(B);
+  CHECK_CUDA(D);
+  CHECK_DIM(3, A);
+  CHECK_DIM(3, B);
+  CHECK_DIM(3, D);
+  CHECK_CONTIGUOUS(algo_buffer);
+  TVM_FFI_ICHECK(A.size(0) == B.size(0) && A.size(0) == D.size(0)) << "Batch sizes must match";
+  TVM_FFI_ICHECK(A.size(2) == B.size(1)) << "Incompatible matrix sizes";
+  TVM_FFI_ICHECK(A.size(1) == D.size(1) && B.size(2) == D.size(2))
+      << "Result tensor has incorrect shape";
+
+  int64_t result = 0;
+  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(B.dtype(), b_type, [&] {
+    return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(A.dtype(), a_type, [&] {
+      return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(D.dtype(), d_type, [&] {
+        auto batch_size = A.size(0);
+        auto m = A.size(1);
+        auto k = A.size(2);
+        auto n = B.size(2);
+
+        auto lt_handle = reinterpret_cast<cublasLtHandle_t>(cublas_handle);
+        ffi::CUDADeviceGuard device_guard(A.device().device_id);
+
+        int max_algos = static_cast<int>(algo_buffer.numel() * get_element_size(algo_buffer) /
+                                         flashinfer::bmm_fp8::kAlgoBytes);
+        result = flashinfer::bmm_fp8::get_fp8_algorithms<b_type, a_type, d_type>(
+            batch_size, n, m, k, static_cast<float*>(B_scale.data_ptr()),
+            static_cast<float*>(A_scale.data_ptr()),
+            workspace_buffer.numel() * get_element_size(workspace_buffer), lt_handle,
+            algo_buffer.data_ptr(), max_algos);
+        return true;
+      });
+    });
+  });
+  return static_cast<int64_t>(result);
+}
+
+void bmm_fp8_run_with_algo(TensorView A, TensorView B, TensorView D, TensorView A_scale,
+                           TensorView B_scale, TensorView workspace_buffer, int64_t cublas_handle,
+                           TensorView algo_buffer, int64_t algo_idx) {
+  CHECK_CUDA(A);
+  CHECK_CUDA(B);
+  CHECK_CUDA(D);
+  CHECK_DIM(3, A);
+  CHECK_DIM(3, B);
+  CHECK_DIM(3, D);
+  CHECK_CONTIGUOUS(algo_buffer);
+  TVM_FFI_ICHECK(A.size(0) == B.size(0) && A.size(0) == D.size(0)) << "Batch sizes must match";
+  TVM_FFI_ICHECK(A.size(2) == B.size(1)) << "Incompatible matrix sizes";
+  TVM_FFI_ICHECK(A.size(1) == D.size(1) && B.size(2) == D.size(2))
+      << "Result tensor has incorrect shape";
+
+  int64_t max_algos =
+      algo_buffer.numel() * get_element_size(algo_buffer) / flashinfer::bmm_fp8::kAlgoBytes;
+  TVM_FFI_ICHECK(algo_idx >= 0 && algo_idx < max_algos)
+      << "algo_idx " << algo_idx << " out of range [0, " << max_algos << ")";
+
+  DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(B.dtype(), b_type, [&] {
+    return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP8(A.dtype(), a_type, [&] {
+      return DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(D.dtype(), d_type, [&] {
+        auto batch_size = A.size(0);
+        auto m = A.size(1);
+        auto k = A.size(2);
+        auto n = B.size(2);
+
+        auto lt_handle = reinterpret_cast<cublasLtHandle_t>(cublas_handle);
+        ffi::CUDADeviceGuard device_guard(A.device().device_id);
+        auto stream = get_stream(A.device());
+
+        auto status = flashinfer::bmm_fp8::bmm_fp8_run_with_algo<b_type, a_type, d_type>(
+            workspace_buffer.data_ptr(),
+            workspace_buffer.numel() * get_element_size(workspace_buffer),
+            static_cast<b_type*>(B.data_ptr()), static_cast<a_type*>(A.data_ptr()),
+            static_cast<d_type*>(D.data_ptr()), batch_size, n, m, k,
+            static_cast<float*>(B_scale.data_ptr()), static_cast<float*>(A_scale.data_ptr()),
+            lt_handle, stream, algo_buffer.data_ptr(), static_cast<int>(algo_idx));
+        TVM_FFI_ICHECK(status == CUBLAS_STATUS_SUCCESS)
+            << "bmm_fp8_run_with_algo failed: " << cublasGetStatusString(status);
+        return true;
+      });
+    });
+  });
+}

csrc/flashinfer_gemm_binding.cu

@@ -19,9 +19,19 @@
 void bmm_fp8(TensorView A, TensorView B, TensorView D, TensorView A_scale, TensorView B_scale,
              TensorView workspace_buffer, int64_t cublas_handle);
 
+int64_t bmm_fp8_get_algos(TensorView A, TensorView B, TensorView D, TensorView A_scale,
+                          TensorView B_scale, TensorView workspace_buffer, int64_t cublas_handle,
+                          TensorView algo_buffer);
+
+void bmm_fp8_run_with_algo(TensorView A, TensorView B, TensorView D, TensorView A_scale,
+                           TensorView B_scale, TensorView workspace_buffer, int64_t cublas_handle,
+                           TensorView algo_buffer, int64_t algo_idx);
+
 void CutlassSegmentGEMM(TensorView workspace_buffer, TensorView all_problems, TensorView x_ptr,
                         TensorView w_ptr, TensorView y_ptr, TensorView x_ld, TensorView w_ld,
                         TensorView y_ld, TensorView empty_x_data, bool weight_column_major);
 
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(cutlass_segment_gemm, CutlassSegmentGEMM);
 TVM_FFI_DLL_EXPORT_TYPED_FUNC(bmm_fp8, bmm_fp8);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(bmm_fp8_get_algos, bmm_fp8_get_algos);
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(bmm_fp8_run_with_algo, bmm_fp8_run_with_algo);

csrc/fmha_v2/fmha/gmem_tile_qkv_packed.h

@@ -889,12 +889,14 @@ struct Gmem_tile_paged_kv {
     // Do not load/store if the thread is in the padded area
     col_in_bytes_ = cta_col_offset_in_bytes + col * BYTES_PER_LDG;
 
-    int64_t kv_stride_in_bytes =
-        qkv_offset == 1 ? params.k_stride_in_bytes : params.v_stride_in_bytes;
-    // The head offset.
-    head_stride_in_bytes_ = (int64_t)(binfo.bidh / params.h_q_per_kv) * kv_stride_in_bytes;
-    // When V is padded (like MLA), we cannot use VALID_BYTES_PER_ROW
-    token_stride_in_bytes_ = kv_stride_in_bytes >> paged_kv_log2_block_size_;
+    // The head stride in bytes.
+    int64_t head_stride_in_bytes =
+        qkv_offset == 1 ? params.k_stride_in_bytes_2 : params.v_stride_in_bytes_2;
+    // The head offset in bytes.
+    head_offset_in_bytes_ = (binfo.bidh / params.h_q_per_kv) * head_stride_in_bytes;
+
+    // The token stride in bytes.
+    token_stride_in_bytes_ = qkv_offset == 1 ? params.k_stride_in_bytes : params.v_stride_in_bytes;
 
     // Take the CTA offset to modify the sequence length.
     // Actually we don't need that for flash attention.
@@ -918,7 +920,7 @@ struct Gmem_tile_paged_kv {
     void const* ptrs[LDGS];
 
     // Offset for the new paged kv pointer.
-    uint64_t const head_col_in_bytes = head_stride_in_bytes_ + col_in_bytes_;
+    uint64_t const head_col_in_bytes = head_offset_in_bytes_ + col_in_bytes_;
 
 // Update paged_kv ptr for each LDG (reuse is possible).
 #pragma unroll
@@ -984,9 +986,9 @@ struct Gmem_tile_paged_kv {
   int row_;
   int64_t col_in_bytes_;
   // Keep track of the head offset.
-  int64_t head_stride_in_bytes_;
+  int64_t head_offset_in_bytes_;
   // // for DeepSeek MLA, the stride of V tokens != VALID_BYTES_PER_ROW
-  int32_t token_stride_in_bytes_;
+  int64_t token_stride_in_bytes_;
   // The sequence length.
   int actual_seqlen_;
   // The past sequence length (kv_seqlen - q_seqlen) considering chunked context.

csrc/fmha_v2/fmha/warpspec/dma.h

@@ -795,13 +795,14 @@ struct DMA {
         uint32_t tensor_size_v[4] = {dv, tokens_per_block, h_kv, INT_MAX};
 
         uint64_t tensor_stride_k[3];
-        tensor_stride_k[0] = params.k_stride_in_bytes / tokens_per_block;  // d
-        tensor_stride_k[1] = params.k_stride_in_bytes;                     // d * 64
+        tensor_stride_k[0] = params.k_stride_in_bytes;
+        tensor_stride_k[1] = params.k_stride_in_bytes_2;
         tensor_stride_k[2] = params.paged_kv_cache.mBytesPerBlock;
         uint64_t tensor_stride_v[3];
         // we cannot use dv * Kernel_traits::ELEMENT_BYTES because V may be padded (MLA)
-        tensor_stride_v[0] = params.v_stride_in_bytes / tokens_per_block;  // dv
-        tensor_stride_v[1] = params.v_stride_in_bytes;                     // dv * 64
+        // use the values given by caller
+        tensor_stride_v[0] = params.v_stride_in_bytes;
+        tensor_stride_v[1] = params.v_stride_in_bytes_2;
         tensor_stride_v[2] = params.paged_kv_cache.mBytesPerBlock;
 
         char* kv_ptr = reinterpret_cast<char*>(params.paged_kv_cache.mPoolPtr);

csrc/fmha_v2/fused_multihead_attention.h

@@ -237,6 +237,13 @@ struct Fused_multihead_attention_params_v2 : Fused_multihead_attention_params_ba
   int64_t q_stride_in_bytes;
   int64_t k_stride_in_bytes;
   int64_t v_stride_in_bytes;
+  // Paged KV uses 4D tensor, the tensor size is:
+  //   HND = [num_pages, H, page_size, D] or NHD = [num_pages, page_size, H, D]
+  // so need another pair of stride.
+  // x_stride_in_bytes means the stride of tensor_size[1]
+  // x_stride_in_bytes_2 means the stride of tensor_size[2]
+  int64_t k_stride_in_bytes_2;
+  int64_t v_stride_in_bytes_2;
 
   // Paged KV load.
   int blocks_per_tma_load;

csrc/fmha_v2/fused_multihead_attention_demo_bert_params.h

@@ -177,4 +177,12 @@ struct Fused_multihead_attention_params_v2 {
   uint32_t* skip_softmax_total_blocks;
   uint32_t* skip_softmax_skipped_blocks;
 #endif
+
+  // Paged KV uses 4D tensor, the tensor size is:
+  //   HND = [num_pages, H, page_size, D] or NHD = [num_pages, page_size, H, D]
+  // so need another pair of stride.
+  // x_stride_in_bytes means the stride of tensor_size[1]
+  // x_stride_in_bytes_2 means the stride of tensor_size[2]
+  int64_t k_stride_in_bytes_2;
+  int64_t v_stride_in_bytes_2;
 };

csrc/fmha_v2_run.cu

@@ -50,7 +50,7 @@ static inline void set_params(
     // types
     Data_type data_type, Data_type acc_type, Data_type output_dtype,
     // attention input layout
-    Attention_input_layout input_layout,
+    Attention_input_layout input_layout, const bool is_paged_hnd,
     // sizes
     const size_t b, const size_t s_q, const size_t s_kv, const size_t h, const size_t h_kv,
     const size_t d, const size_t dv, const size_t total, const size_t num_grouped_heads,
@@ -119,8 +119,21 @@ static inline void set_params(
                          get_size_in_bytes(tokens_per_block * h_kv * std::gcd(d, dv), data_type),
                          paged_kv_pool_ptr);
       params.paged_kv_cache.mBlockOffsets = paged_block_offsets;
-      params.k_stride_in_bytes = get_size_in_bytes(tokens_per_block * d, data_type);
-      params.v_stride_in_bytes = get_size_in_bytes(tokens_per_block * dv, data_type);
+      // FMHA kernels always access the K/V tensor in 4D coordinate [num_pages, H_kv, page_size, D].
+      // The layout of HND or NHD is implemented by tensor strides to get the correct memory
+      // address. 4D tensor strides of HND: [block_size, page_size * D, D ,1] 4D tensor strides of
+      // NHD: [block_size, D, H_kv * D, 1]
+      if (is_paged_hnd) {
+        params.k_stride_in_bytes = get_size_in_bytes(d, data_type);
+        params.v_stride_in_bytes = get_size_in_bytes(dv, data_type);
+        params.k_stride_in_bytes_2 = get_size_in_bytes(tokens_per_block * d, data_type);
+        params.v_stride_in_bytes_2 = get_size_in_bytes(tokens_per_block * dv, data_type);
+      } else {
+        params.k_stride_in_bytes = get_size_in_bytes(h_kv * d, data_type);
+        params.v_stride_in_bytes = get_size_in_bytes(h_kv * dv, data_type);
+        params.k_stride_in_bytes_2 = get_size_in_bytes(d, data_type);
+        params.v_stride_in_bytes_2 = get_size_in_bytes(dv, data_type);
+      }
     } else if (input_layout == Attention_input_layout::SEPARATE_Q_K_V) {
       // Layout [B, S, H_kv, D].
       params.k_ptr = k_d;
@@ -247,10 +260,15 @@ static inline void determine_launch_params(
   launch_params.multi_processor_count = props.multiProcessorCount;
   launch_params.device_l2_cache_size = props.l2CacheSize;
 
+#if 0
   // threshold for adopting flash attention or warp_specialized kernels.
   launch_params.flash_attention =
       (data_type == DATA_TYPE_FP16 || data_type == DATA_TYPE_BF16 || data_type == DATA_TYPE_E4M3) &&
       (s >= 16 && d >= 16) && !force_non_flash_attention;
+#else
+  // Currently only flash attention kernels are generated in FlashInfer
+  launch_params.flash_attention = true;
+#endif
 
   // enable warp_speialized kernels when s >= 512 on hopper
   // note that warp_speialized kernels need flash attention + tma
@@ -304,11 +322,18 @@ static inline Attention_mask_type string_to_mask_type(const std::string& s) {
   return Attention_mask_type::CAUSAL;  // default
 }
 
-static inline Attention_input_layout string_to_input_layout(const std::string& s) {
+static inline Attention_input_layout string_to_input_layout(const std::string& s,
+                                                            bool& is_paged_hnd) {
+  is_paged_hnd = false;
   if (s == "packed_qkv") return Attention_input_layout::PACKED_QKV;
   if (s == "contiguous_q_kv") return Attention_input_layout::CONTIGUOUS_Q_KV;
-  if (s == "q_paged_kv_nhd") return Attention_input_layout::Q_PAGED_KV;
-  if (s == "q_paged_kv_hnd") return Attention_input_layout::Q_PAGED_KV;
+  if (s == "q_paged_kv_nhd") {
+    return Attention_input_layout::Q_PAGED_KV;
+  }
+  if (s == "q_paged_kv_hnd") {
+    is_paged_hnd = true;
+    return Attention_input_layout::Q_PAGED_KV;
+  }
   if (s == "separate_q_k_v") return Attention_input_layout::SEPARATE_Q_K_V;
   throw std::invalid_argument("Unsupported input_layout: " + s);
 }
@@ -330,7 +355,8 @@ void fmha_v2_run(
     float skip_softmax_threshold_scale_factor,
     Optional<ffi::TensorView> softmax_stats,  // Optional [batch, s_q, num_heads, 2] for (max, sum)
     Optional<ffi::TensorView> sinks) {
-  Attention_input_layout input_layout = string_to_input_layout(input_layout_str);
+  bool is_paged_hnd;
+  Attention_input_layout input_layout = string_to_input_layout(input_layout_str, is_paged_hnd);
   Attention_mask_type attention_mask_type = string_to_mask_type(mask_mode_str);
   Data_type output_dtype = dltype_to_data_type(o.dtype());
   // Get device properties
@@ -360,9 +386,12 @@ void fmha_v2_run(
     d = q.shape()[3];     // head_dim_qk
     dv = q.shape()[3];    // head_dim_v (same as d for standard attention)
   } else if (input_layout == Attention_input_layout::Q_PAGED_KV) {
-    // q is 3D: [total_tokens, H, D], k/v are 4D paged: [num_pages, H_kv, page_size, D]
+    // q is 3D: [total_tokens, H, D]
     h = q.shape()[1];
-    h_kv = k.shape()[1];
+    // k/v are 4D paged:
+    //   HND: [num_pages, H_kv, page_size, D]
+    //   NHD: [num_pages, page_size, H_kv, D]
+    h_kv = k.shape()[is_paged_hnd ? 1 : 2];
     d = q.shape()[2];
     dv = v.shape()[3];
   } else if (input_layout == Attention_input_layout::CONTIGUOUS_Q_KV) {

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

@@ -71,6 +71,9 @@ using namespace tensorrt_llm::kernels;
 using namespace tensorrt_llm::common;
 
 namespace tensorrt_llm::kernels::cutlass_kernels {
+
+constexpr int CVT_ELTS_PER_THREAD = 8;
+
 /**
  * Takes the input maps and prepares the expanded maps for min latency
  * @param num_active_experts_per_node: Number of active experts on current node