flashinfer-python 0.2.7.post1__tar.gz → 0.2.8rc1__tar.gz

{flashinfer_python-0.2.7.post1/flashinfer_python.egg-info → flashinfer_python-0.2.8rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.2.7.post1
+Version: 0.2.8rc1
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 Project-URL: Homepage, https://github.com/flashinfer-ai/flashinfer
@@ -93,6 +93,9 @@ Alternatively, build FlashInfer from source:
 git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
 cd flashinfer
 python -m pip install -v .
+
+# for development & contribution, install in editable mode
+python -m pip install --no-build-isolation -e . -v
 ```
 
 To pre-compile essential kernels ahead-of-time (AOT), run the following command:

{flashinfer_python-0.2.7.post1 → flashinfer_python-0.2.8rc1}/README.md

@@ -73,6 +73,9 @@ Alternatively, build FlashInfer from source:
 git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
 cd flashinfer
 python -m pip install -v .
+
+# for development & contribution, install in editable mode
+python -m pip install --no-build-isolation -e . -v
 ```
 
 To pre-compile essential kernels ahead-of-time (AOT), run the following command:

{flashinfer_python-0.2.7.post1 → flashinfer_python-0.2.8rc1}/csrc/cudnn_sdpa_kernel_launcher.cu

@@ -68,16 +68,28 @@ constexpr int32_t BYTES_PER_ELEMENT = 2;
 
 enum KernelType { PREFILL, PREFILL_DEEPSEEK, DECODE };
 
-void init_cudnn_cubin(std::map<KernelType, std::string>& cubin_map) {
-  cubin_map[PREFILL] = getCubin("fmha/cudnn/cudnn_sm100_fprop_sdpa_prefill_d128_bf16",
-                                "ff14e8dcfc04d9b3a912dd44056be37d9aa8a85976e0070494ca0cce0524f2a1");
-
-  cubin_map[DECODE] = getCubin("fmha/cudnn/cudnn_sm100_fprop_sdpa_decode_d128_bf16",
-                               "e7ce0408b4c3a36c42616498228534ee64cab785ef570af5741deaf9dd1b475c");
+enum PrefillType {
+  KERNEL_PREFILL,
+  KERNEL_PREFILL_DEEPSEEK,
+  KERNEL_PREFILL_CAUSAL,
+  KERNEL_PREFILL_DEEPSEEK_CAUSAL,
+  KERNEL_NUM_PREFILL_TYPES
+};
 
-  cubin_map[PREFILL_DEEPSEEK] =
-      getCubin("fmha/cudnn/cudnn_sm100_fprop_sdpa_prefill_d192_bf16",
-               "2190967b8733e193cdcecc054eeb7c2907080a158a33fe7ba2004523a4aff6f9");
+void init_cudnn_cubin(std::map<KernelType, std::string>& cubin_map) {
+  cubin_map[PREFILL] = getCubin(
+      "4c623163877c8fef5751c9c7a59940cd2baae02e/fmha/cudnn/"
+      "cudnn_sm100_fprop_sdpa_prefill_d128_bf16",
+      "ff14e8dcfc04d9b3a912dd44056be37d9aa8a85976e0070494ca0cce0524f2a1");
+
+  cubin_map[DECODE] = getCubin(
+      "4c623163877c8fef5751c9c7a59940cd2baae02e/fmha/cudnn/cudnn_sm100_fprop_sdpa_decode_d128_bf16",
+      "e7ce0408b4c3a36c42616498228534ee64cab785ef570af5741deaf9dd1b475c");
+
+  cubin_map[PREFILL_DEEPSEEK] = getCubin(
+      "4c623163877c8fef5751c9c7a59940cd2baae02e/fmha/cudnn/"
+      "cudnn_sm100_fprop_sdpa_prefill_d192_bf16",
+      "2190967b8733e193cdcecc054eeb7c2907080a158a33fe7ba2004523a4aff6f9");
 }
 
 auto get_cudnn_cubin(KernelType kernel_type) -> std::string {
@@ -385,19 +397,28 @@ static void create_packed_tma_desc_qo_prefill(int b, int32_t* actual_seq_lens_q_
   }
 }
 
-void setup_prefill(CUfunction* hfunc, CUfunction* hfunc_deepseek) {
+void setup_prefill(CUfunction* prefill_func) {
   // Use cu++filt to get the kernel name
-  std::string kernel_name_deepseek =
-
+  std::string kernel_name_deepseek_causal =
       "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_"
       "128x128x192ILb1ELb0EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_"
       "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE";
 
-  std::string kernel_name =
+  std::string kernel_name_causal =
       "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_"
       "128x128x128ILb1ELb1EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_"
       "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE";
 
+  std::string kernel_name_deepseek =
+      "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_"
+      "128x128x192ILb0ELb0EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_"
+      "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE";
+
+  std::string kernel_name =
+      "_Z47cudnn_sm100_fprop_sdpa_prefill_bf16_"
+      "128x128x128ILb0ELb1EEvN4fmha19AttentionDescriptorEPKN3tma11cudaTmaDescES5_fPfNS0_"
+      "7stridesES5_S5_PKjS9_S9_jjNS0_11FastDivisorE";
+
   std::string cubin = get_cudnn_cubin(PREFILL);
   std::string cubin_deepseek = get_cudnn_cubin(PREFILL_DEEPSEEK);
 
@@ -418,12 +439,23 @@ void setup_prefill(CUfunction* hfunc, CUfunction* hfunc_deepseek) {
     throw std::runtime_error("Failed to cuModuleLoadData for prefill");
   }
 
-  if (cuModuleGetFunction(hfunc, hmod, kernel_name.c_str()) != CUDA_SUCCESS) {
+  if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL], hmod, kernel_name.c_str()) !=
+      CUDA_SUCCESS) {
     throw std::runtime_error("Failed to cuModuleGetFunction for prefill");
   }
 
-  if (cuModuleGetFunction(hfunc_deepseek, hmod_deepseek, kernel_name_deepseek.c_str()) !=
+  if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL_DEEPSEEK], hmod_deepseek,
+                          kernel_name_deepseek.c_str()) != CUDA_SUCCESS) {
+    throw std::runtime_error("Failed to cuModuleGetFunction for prefill_deepseek");
+  }
+
+  if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL_CAUSAL], hmod, kernel_name_causal.c_str()) !=
       CUDA_SUCCESS) {
+    throw std::runtime_error("Failed to cuModuleGetFunction for prefill");
+  }
+
+  if (cuModuleGetFunction(&prefill_func[KERNEL_PREFILL_DEEPSEEK_CAUSAL], hmod_deepseek,
+                          kernel_name_deepseek_causal.c_str()) != CUDA_SUCCESS) {
     throw std::runtime_error("Failed to cuModuleGetFunction for prefill_deepseek");
   }
 };
@@ -506,30 +538,24 @@ void prefill(int64_t b, int64_t s_qo, int64_t max_s_kv, at::Tensor q, at::Tensor
 
   // Step 1: Setup the kernel pointer
 
-  static CUfunction hfunc{nullptr};
-  static CUfunction hfunc_deepseek{nullptr};
+  static CUfunction prefill_func[KERNEL_NUM_PREFILL_TYPES] = {nullptr, nullptr, nullptr, nullptr};
 
   int64_t d_qk = q.size(2);
 
   int64_t d_vo = v_cache.dim() == 3 ? v_cache.size(2) : v_cache.size(3);
 
-  if (hfunc == nullptr || hfunc_deepseek == nullptr) {
-    setup_prefill(&hfunc, &hfunc_deepseek);
+  if (prefill_func[0] == nullptr) {
+    setup_prefill(prefill_func);
 
-    if (hfunc != nullptr) {
-      cuErrCheck(
-          cuFuncSetAttribute(hfunc, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, SMEM_SIZE));
-      cuErrCheck(
-          cuFuncSetAttribute(hfunc, CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, 100));
-      cuErrCheck(cuFuncSetAttribute(hfunc, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
-    }
-    if (hfunc_deepseek != nullptr) {
-      cuErrCheck(cuFuncSetAttribute(hfunc_deepseek, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
-                                    SMEM_SIZE));
-      cuErrCheck(cuFuncSetAttribute(hfunc_deepseek,
-                                    CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, 100));
-      cuErrCheck(cuFuncSetAttribute(hfunc_deepseek,
-                                    CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
+    for (int i = 0; i < KERNEL_NUM_PREFILL_TYPES; i++) {
+      if (prefill_func[i] != nullptr) {
+        cuErrCheck(cuFuncSetAttribute(prefill_func[i],
+                                      CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, SMEM_SIZE));
+        cuErrCheck(cuFuncSetAttribute(prefill_func[i],
+                                      CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT, 100));
+        cuErrCheck(cuFuncSetAttribute(prefill_func[i],
+                                      CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
+      }
     }
   }
 
@@ -582,7 +608,7 @@ void prefill(int64_t b, int64_t s_qo, int64_t max_s_kv, at::Tensor q, at::Tensor
     config.gridDimX = actual_num_tiles_per_head;
 
   } else {
-    config.gridDimX = static_cast<int>(std::ceil(q.size(0) / (TILE_M_1 * 2.0f))) * b;
+    config.gridDimX = static_cast<int>(std::ceil(s_qo / (TILE_M_1 * 2.0f))) * b;
   }
 
   config.gridDimY = h_qo;
@@ -666,7 +692,15 @@ void prefill(int64_t b, int64_t s_qo, int64_t max_s_kv, at::Tensor q, at::Tensor
     dim3 grid(1, 1, 1);
     dim3 block(128, 1, 1);
 
-    qkv_tma_setup_prefill<<<grid, block, 0, stream>>>(
+    at::cuda::CUDAStream cuda_stream = at::cuda::getCurrentCUDAStream(device.index());
+    cudaStream_t raw_stream = cuda_stream.stream();
+
+    cudaError_t err = cudaStreamQuery(raw_stream);
+    if (!(err == cudaSuccess || err == cudaErrorNotReady)) {
+      throw std::runtime_error("CUDA cudnn stream error" + std::string(cudaGetErrorString(err)));
+    }
+
+    qkv_tma_setup_prefill<<<grid, block, 0, raw_stream>>>(
         b, h_qo, h_kv, d_qk, d_vo, is_kv_ragged, page_size, total_num_pages,
         k_cache.strides().data()[2], k_cache.strides().data()[1], k_cache.strides().data()[0],
         v_cache.strides().data()[2], v_cache.strides().data()[1], v_cache.strides().data()[0],
@@ -737,12 +771,16 @@ void prefill(int64_t b, int64_t s_qo, int64_t max_s_kv, at::Tensor q, at::Tensor
   args[13] = &page_size_div;
 
   auto err_launch = CUDA_SUCCESS;
-  if (d_qk == 192) {
-    err_launch = cuLaunchKernelEx(&config, hfunc_deepseek, (void**)args, nullptr);
+
+  auto choice = KERNEL_PREFILL;
+  if (causal) {
+    choice = d_qk == 192 ? KERNEL_PREFILL_DEEPSEEK_CAUSAL : KERNEL_PREFILL_CAUSAL;
   } else {
-    err_launch = cuLaunchKernelEx(&config, hfunc, (void**)args, nullptr);
+    choice = d_qk == 192 ? KERNEL_PREFILL_DEEPSEEK : KERNEL_PREFILL;
   }
 
+  err_launch = cuLaunchKernelEx(&config, prefill_func[choice], (void**)args, nullptr);
+
   if (err_launch != CUDA_SUCCESS) {
     const char* errstr = NULL;
     cuGetErrorString(err_launch, &errstr);
@@ -882,19 +920,20 @@ void setup_tma_desc_decode(int64_t b, int64_t s_kv, int64_t h_qo, int64_t h_kv,
                                 tma::cudaTmaDescSwizzle::SWIZZLE_128B);
 }
 
-void decode(at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, double scale,
+void decode(int64_t max_s_kv, at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, double scale,
             at::Tensor workspace_buffer, at::Tensor actual_seq_lens_kv,
             at::Tensor actual_seq_lens_kv_gpu, at::Tensor block_tables, at::Tensor out,
-            std::optional<at::Tensor> batch_offset_array, bool is_cuda_graph_compatible) {
+            std::optional<at::Tensor> batch_offset_q_array,
+            std::optional<at::Tensor> batch_offset_o_array, bool is_cuda_graph_compatible) {
   constexpr size_t SMEM_SIZE = 227 * 1024;  // All smem
   constexpr size_t REDUCTION_MEM_SIZE = 128 * 1024;
   constexpr int64_t TILE_N_1 = 128;
 
   constexpr int32_t NUM_THREADS = 384;
 
-  int64_t* batch_offset_array_data = nullptr;
-  if (batch_offset_array.has_value()) {
-    batch_offset_array_data = batch_offset_array.value().data_ptr<int64_t>();
+  int64_t* batch_offset_q_array_data = nullptr;
+  if (batch_offset_q_array.has_value()) {
+    batch_offset_q_array_data = batch_offset_q_array.value().data_ptr<int64_t>();
   }
 
   auto device = q.device();
@@ -940,20 +979,21 @@ void decode(at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, double scale,
 
   int64_t b = q.size(0);
   int64_t h_qo = q.size(1);
-  int64_t h_kv = k_cache.size(1);
-  int64_t page_size = k_cache.size(2);
   int64_t d = q.size(2);
-  int64_t total_num_pages = k_cache.size(0);
 
-  int64_t max_skv = (k_cache.size(0) / b) * page_size;
+  int64_t h_kv = k_cache.size(1);
+
+  int64_t page_size = k_cache.dim() == 4 ? k_cache.size(2) : 1;
 
-  int64_t s_qo = 1;
+  int64_t total_num_pages = k_cache.dim() == 4 ? k_cache.size(0) : 1;
+
+  int64_t s_kv = max_s_kv;
 
-  int32_t split_factor = compute_split_factor(b, h_kv, h_qo, max_skv, sm_count);
+  int64_t s_qo = 1;
 
-  split_factor = 1;
+  int32_t split_factor = compute_split_factor(b, h_kv, h_qo, s_kv, sm_count);
 
-  // TORCH_CHECK(split_factor == 1, "Split factor > 1 not supported yet");
+  split_factor = 1;  // Fix split factor. Setting it to 1 for now
 
   // Set up TMA descriptors for Q, K, V, O
   auto qo_strides = q.strides();
@@ -982,14 +1022,13 @@ void decode(at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, double scale,
 
   config.hStream = stream;
   config.numAttrs = 1;
-  config.attrs = attrs;
 
   int8_t* workspace_start = workspace_buffer.data_ptr<int8_t>();
   int8_t* partial_o_dev = workspace_start;
   int8_t* tma_descriptor_start =
       partial_o_dev + (b * s_qo * h_qo * d * sizeof(float) * split_factor);
 
-  int8_t* batch_strides_dev = tma_descriptor_start + ((3 * b + 2) * sizeof(tma::cudaTmaDesc));
+  int8_t* batch_strides_dev = tma_descriptor_start + ((5 * b) * sizeof(tma::cudaTmaDesc));
 
   tma::cudaTmaDesc* packed_tma_desc_q_dev =
       reinterpret_cast<tma::cudaTmaDesc*>(tma_descriptor_start);
@@ -999,8 +1038,8 @@ void decode(at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, double scale,
       reinterpret_cast<tma::cudaTmaDesc*>(tma_descriptor_start + b * sizeof(tma::cudaTmaDesc) * 2);
   tma::cudaTmaDesc* tma_desc_k_dev =
       reinterpret_cast<tma::cudaTmaDesc*>(tma_descriptor_start + b * sizeof(tma::cudaTmaDesc) * 3);
-  tma::cudaTmaDesc* tma_desc_v_dev = reinterpret_cast<tma::cudaTmaDesc*>(
-      tma_descriptor_start + (sizeof(tma::cudaTmaDesc) * ((b * 3) + 1)));
+  tma::cudaTmaDesc* tma_desc_v_dev =
+      reinterpret_cast<tma::cudaTmaDesc*>(tma_descriptor_start + b * sizeof(tma::cudaTmaDesc) * 4);
 
   int8_t* lse_dev = batch_strides_dev + (b * sizeof(int64_t));
 
@@ -1034,15 +1073,15 @@ void decode(at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, double scale,
         tma_desc_v_dev, packed_tma_desc_o_dev, packed_tma_desc_partial_o_dev,
         reinterpret_cast<int64_t*>(batch_strides_dev));
   } else {
-    std::unique_ptr<tma::cudaTmaDesc[]> tma_desc_host(new tma::cudaTmaDesc[(3 * b) + 2]);
+    std::unique_ptr<tma::cudaTmaDesc[]> tma_desc_host(new tma::cudaTmaDesc[5 * b]);
 
     tma::cudaTmaDesc* tma_desc_q = tma_desc_host.get();
     tma::cudaTmaDesc* tma_desc_o = tma_desc_host.get() + b;
     tma::cudaTmaDesc* tma_desc_partial_o = tma_desc_host.get() + b * 2;
     tma::cudaTmaDesc* tma_desc_k = tma_desc_host.get() + b * 3;
-    tma::cudaTmaDesc* tma_desc_v = tma_desc_host.get() + (b * 3) + 1;
+    tma::cudaTmaDesc* tma_desc_v = tma_desc_host.get() + b * 4;
 
-    setup_tma_desc_decode(b, max_skv, h_qo, h_kv, d, total_num_pages, q, out, k_cache, v_cache,
+    setup_tma_desc_decode(b, max_s_kv, h_qo, h_kv, d, total_num_pages, q, out, k_cache, v_cache,
                           split_factor, page_size, partial_o_dev, tma_desc_q, tma_desc_o,
                           tma_desc_partial_o, tma_desc_k, tma_desc_v);
 
@@ -1053,17 +1092,18 @@ void decode(at::Tensor q, at::Tensor k_cache, at::Tensor v_cache, double scale,
     cudaMemcpyAsync(batch_strides_dev, batch_strides.get(), sizeof(int64_t) * b,
                     cudaMemcpyHostToDevice, stream);
 
-    cudaMemcpyAsync(tma_descriptor_start, tma_desc_host.get(),
-                    sizeof(tma::cudaTmaDesc) * (3 * b + 2), cudaMemcpyHostToDevice, stream);
+    cudaMemcpyAsync(tma_descriptor_start, tma_desc_host.get(), sizeof(tma::cudaTmaDesc) * (5 * b),
+                    cudaMemcpyHostToDevice, stream);
   }
 
-  cudnn_sdpa::AttentionDescriptor_t attnDesc{b, h_qo, h_kv, h_kv, s_qo, max_skv, d, h_qo / h_kv, 0};
+  cudnn_sdpa::AttentionDescriptor_t attnDesc{b,        h_qo, h_kv,        h_kv, s_qo,
+                                             max_s_kv, d,    h_qo / h_kv, 0};
 
   cudnn_sdpa::FastDivisor_t page_size_div;
   setFastDivisor(page_size_div, page_size);
 
   uint32_t page_size32 = static_cast<uint32_t>(page_size);
-  uint32_t num_pages_per_seq32 = static_cast<uint32_t>(max_skv / page_size);
+  uint32_t num_pages_per_seq32 = static_cast<uint32_t>(max_s_kv / page_size);
 
   void* args[15];
 

{flashinfer_python-0.2.7.post1 → flashinfer_python-0.2.8rc1}/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_sm100_ops.cu

@@ -182,7 +182,8 @@ class FusedMoeRunner : public torch::CustomClassHolder {
   FusedMoeRunner(FusedMoeRunner const&) = delete;
   void operator=(FusedMoeRunner const&) = delete;
 
-  at::Tensor runMoe(at::Tensor const& input, at::Tensor const& token_selected_experts,
+  at::Tensor runMoe(at::Tensor& output, at::Tensor const& input,
+                    at::Tensor const& token_selected_experts,
                     std::optional<at::Tensor> token_final_scales,
                     at::Tensor const& fc1_expert_weights, at::Tensor const& fc2_expert_weights,
                     std::optional<c10::ArrayRef<at::Tensor>> quant_scales,
@@ -241,9 +242,6 @@ class FusedMoeRunner : public torch::CustomClassHolder {
 
     auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
 
-    std::vector<int64_t> output_shape = {num_rows, hidden_size};
-    auto output = torch::empty(output_shape, input.options().dtype(mOutputDtype));
-
     WorkspaceInfo workspace_info = getWorkspaceInfo(
         num_rows, hidden_size, inter_size, num_experts_total, static_cast<int>(experts_per_token),
         activation_type, parallelism_config, min_latency_mode);
@@ -271,7 +269,8 @@ class FusedMoeRunner : public torch::CustomClassHolder {
   }
 
   std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> runMoeMinLantency(
-      torch::Tensor const& input, torch::Tensor const& token_selected_experts,
+      torch::Tensor& output, torch::Tensor const& input,
+      torch::Tensor const& token_selected_experts,
       torch::optional<torch::Tensor> token_final_scales, torch::Tensor const& fc1_expert_weights,
       torch::Tensor const& fc2_expert_weights,
       torch::optional<c10::ArrayRef<torch::Tensor>> quant_scales,
@@ -328,9 +327,6 @@ class FusedMoeRunner : public torch::CustomClassHolder {
 
     auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
 
-    std::vector<int64_t> output_shape = {num_rows * num_experts_on_rank, hidden_size};
-    auto output = torch::empty(output_shape, input.options().dtype(mOutputDtype));
-
     auto num_active_experts_per_node =
         torch::empty({1}, input.options().dtype(at::ScalarType::Int));
     auto experts_to_token_score =

{flashinfer_python-0.2.7.post1 → flashinfer_python-0.2.8rc1}/csrc/trtllm_fmha_kernel_launcher.cu

@@ -16,7 +16,6 @@
 #include <c10/cuda/CUDAStream.h>
 #include <flashinfer/exception.h>
 #include <nvrtc.h>
-#include <torch/all.h>
 
 #include <algorithm>
 #include <cmath>
@@ -25,23 +24,19 @@
 #include <flashinfer/trtllm/fmha/gen_kernel_launcher.cuh>
 #include <iostream>
 
-__constant__ float out_scale = 1.0f;
-__constant__ float log2e_h64 = 0.18033688011112042;
-__constant__ float log2e_h128 = 0.12751743082459868;
-__constant__ float log2e_h192 = 0.10411754627697266;
-__constant__ float log2e_h256 = 0.09016844005556021;
-
 namespace flashinfer {
 template <typename T, Data_type CACHE_T>
 void trtllm_paged_attention_launcher(at::Tensor& out, at::Tensor& query,
                                      at::Tensor& key_value_cache, at::Tensor& workspace_buffer,
-                                     int64_t num_heads_dummy, int64_t num_kv_heads, double scale,
+                                     int64_t num_q_heads, int64_t num_kv_heads, double scale,
                                      at::Tensor& block_tables, at::Tensor& seq_lens,
                                      int64_t block_size, int64_t max_seq_len,
                                      const std::string kv_cache_dtype, double k_scale,
                                      double v_scale) {
   int num_seqs = query.size(0);
   int num_heads = query.size(1);
+  TORCH_CHECK(num_heads == static_cast<int>(num_q_heads),
+              "num_q_heads params and query shape does not match!");
   int head_size = query.size(2);
   int max_num_blocks_per_seq = block_tables.size(-1);
 
@@ -55,6 +50,12 @@ void trtllm_paged_attention_launcher(at::Tensor& out, at::Tensor& query,
 
   uint32_t num_k_heads = num_kv_heads;
   uint32_t num_v_heads = num_k_heads;
+  if (num_heads % num_k_heads != 0) {
+    std::ostringstream err_msg;
+    err_msg << "num_heads must be a multiple of num_k_heads, got num_k_heads: " << num_k_heads
+            << "and num_heads: " << num_heads;
+    FLASHINFER_ERROR(err_msg.str());
+  }
   auto batch_size = num_seqs;
 
   int const beam_width = num_seqs / batch_size;  // always 1
@@ -91,35 +92,19 @@ void trtllm_paged_attention_launcher(at::Tensor& out, at::Tensor& query,
   // num_kv_heads should be enough, but num_heads for safty at long seq len.
   size_t num_semaphores = batch_size * num_heads;
 
-  runner_params.multiCtasKvScratchPtr =
-      reinterpret_cast<void*>(workspace_buffer.data_ptr() + num_semaphores * sizeof(uint32_t));
+  runner_params.multiCtasKvScratchPtr = reinterpret_cast<void*>(
+      static_cast<char*>(workspace_buffer.data_ptr()) + num_semaphores * sizeof(uint32_t));
   runner_params.multiCtasKvCounterPtr = reinterpret_cast<int32_t*>(workspace_buffer.data_ptr());
 
   zero_gmem_semaphore_launcher(runner_params.multiCtasKvCounterPtr, num_semaphores,
                                /*enable_pdl=*/true, stream);
 
-  float* log_scale_ptr;
-  float* output_scale_ptr;
-
-  if (head_size == 64) {
-    cudaGetSymbolAddress((void**)&log_scale_ptr, log2e_h64);
-  } else if (head_size == 128) {
-    cudaGetSymbolAddress((void**)&log_scale_ptr, log2e_h128);
-  } else if (head_size == 192) {
-    cudaGetSymbolAddress((void**)&log_scale_ptr, log2e_h192);
-  } else if (head_size == 256) {
-    cudaGetSymbolAddress((void**)&log_scale_ptr, log2e_h256);
-  } else {
+  if (head_size != 64 && head_size != 128 && head_size != 192 && head_size != 256) {
     std::ostringstream err_msg;
     err_msg << "head_size " << head_size << " is not supported!";
     FLASHINFER_ERROR(err_msg.str());
   }
 
-  cudaGetSymbolAddress((void**)&output_scale_ptr, out_scale);
-
-  runner_params.scaleSoftmaxLog2Ptr = log_scale_ptr;
-  runner_params.outputScalePtr = output_scale_ptr;
-
   runner_params.seqLensKvPtr = reinterpret_cast<int const*>(seq_lens.data_ptr<int>());
 
   runner_params.oPtr = output_ptr;
@@ -127,12 +112,22 @@ void trtllm_paged_attention_launcher(at::Tensor& out, at::Tensor& query,
   runner_params.mHeadDimV = head_size;
   runner_params.mNumHeadsQ = num_heads;
   runner_params.mNumHeadsKv = num_k_heads;
-  runner_params.mNumHeadsQPerKv = 8;
+  runner_params.mNumHeadsQPerKv = num_heads / num_k_heads;
   runner_params.mBatchSize = batch_size;
   runner_params.mMaxSeqLenQ = 1;
   runner_params.mMaxSeqLenKv = max_seq_len;
   runner_params.mSumOfSeqLensQ = int(batch_size * runner_params.mMaxSeqLenQ);
   runner_params.mScaleQ = 1.0;
+  // Set the chunked attention size and sliding window size to INT_MAX to disable them when checking
+  // if the kernel is supported.
+  runner_params.mChunkedAttentionSize = INT_MAX;
+  runner_params.mAttentionWindowSize = INT_MAX;
+  auto [foundKernels, kinfo] = fmha_runner.isSupportedWithInfo(runner_params);
+  if (!foundKernels) {
+    std::ostringstream err_msg;
+    err_msg << "Missing TRTLLM-GEN decode kernel:" << kinfo;
+    FLASHINFER_ERROR(err_msg.str());
+  }
 
   runner_params.mMultiProcessorCount = getMultiProcessorCount();
   auto const [free_memory, total_memory] = getDeviceMemoryInfo(false);
@@ -146,10 +141,10 @@ void trtllm_paged_attention_launcher(at::Tensor& out, at::Tensor& query,
   fmha_runner.run(runner_params);
 }
 
-#define CALL_GEN_LAUNCHER(T, CACHE_T_ENUM)                                                         \
-  trtllm_paged_attention_launcher<T, CACHE_T_ENUM>(                                                \
-      out, query, key_value_cache, workspace_buffer, num_heads, num_kv_heads, scale, block_tables, \
-      seq_lens, block_size, max_seq_len, kv_cache_dtype, k_scale, v_scale);
+#define CALL_GEN_LAUNCHER(T, CACHE_T_ENUM)                                             \
+  trtllm_paged_attention_launcher<T, CACHE_T_ENUM>(                                    \
+      out, query, key_value_cache, workspace_buffer, num_q_heads, num_kv_heads, scale, \
+      block_tables, seq_lens, block_size, max_seq_len, kv_cache_dtype, k_scale, v_scale);
 
 // The following macro is used to dispatch the conversion function based on
 // the data type of the key and value cache. The FN is a macro that calls a
@@ -159,7 +154,7 @@ void trtllm_paged_attention_launcher(at::Tensor& out, at::Tensor& query,
     if (SRC_DTYPE == at::ScalarType::Half) {                                   \
       FN(half, Data_type::DATA_TYPE_FP16);                                     \
     } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                        \
-      FN(__nv_bfloat16, Data_type::DATA_TYPE_FP16);                            \
+      FN(__nv_bfloat16, Data_type::DATA_TYPE_BF16);                            \
     } else {                                                                   \
       TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE);   \
     }                                                                          \
@@ -178,7 +173,7 @@ void trtllm_paged_attention_launcher(at::Tensor& out, at::Tensor& query,
   }
 
 void trtllm_paged_attention(at::Tensor& out, at::Tensor& query, at::Tensor& key_value_cache,
-                            at::Tensor& workspace_buffer, int64_t num_heads, int64_t num_kv_heads,
+                            at::Tensor& workspace_buffer, int64_t num_q_heads, int64_t num_kv_heads,
                             double scale, at::Tensor& block_tables, at::Tensor& seq_lens,
                             int64_t block_size, int64_t max_seq_len,
                             const std::string kv_cache_dtype, double k_scale, double v_scale) {

{flashinfer_python-0.2.7.post1 → flashinfer_python-0.2.8rc1}/custom_backend.py

@@ -2,7 +2,6 @@ import shutil
 from pathlib import Path
 
 from setuptools import build_meta as orig
-from setuptools.build_meta import *  # noqa: F403
 
 _root = Path(__file__).parent.resolve()
 _data_dir = _root / "flashinfer" / "data"
@@ -63,12 +62,7 @@ def _prepare_for_editable():
     _aot_ops_package_dir.symlink_to(_aot_ops_dir)
 
 
-def get_requires_for_build_wheel(config_settings=None):
-    _prepare_for_wheel()
-    return _requires_for_aot
-
-
-def get_requires_for_build_sdist(config_settings=None):
+def _prepare_for_sdist():
     # Remove data directory
     if _data_dir.exists():
         shutil.rmtree(_data_dir)
@@ -78,6 +72,14 @@ def get_requires_for_build_sdist(config_settings=None):
     _aot_ops_package_dir.parent.mkdir(parents=True, exist_ok=True)
     _aot_ops_package_dir.mkdir(parents=True)
 
+
+def get_requires_for_build_wheel(config_settings=None):
+    _prepare_for_wheel()
+    return _requires_for_aot
+
+
+def get_requires_for_build_sdist(config_settings=None):
+    _prepare_for_sdist()
     return []
 
 
@@ -94,3 +96,18 @@ def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None):
 def prepare_metadata_for_build_editable(metadata_directory, config_settings=None):
     _prepare_for_editable()
     return orig.prepare_metadata_for_build_editable(metadata_directory, config_settings)
+
+
+def build_editable(wheel_directory, config_settings=None, metadata_directory=None):
+    _prepare_for_editable()
+    return orig.build_editable(wheel_directory, config_settings, metadata_directory)
+
+
+def build_sdist(sdist_directory, config_settings=None):
+    _prepare_for_sdist()
+    return orig.build_sdist(sdist_directory, config_settings)
+
+
+def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
+    _prepare_for_wheel()
+    return orig.build_wheel(wheel_directory, config_settings, metadata_directory)

flashinfer_python-0.2.8rc1/flashinfer/_build_meta.py

@@ -0,0 +1 @@
+__version__ = '0.2.8rc1'

{flashinfer_python-0.2.7.post1 → flashinfer_python-0.2.8rc1}/flashinfer/attention.py

@@ -39,6 +39,7 @@ class BatchAttention:
     def __init__(
         self,
         kv_layout: str = "NHD",
+        device: str = "cuda",
     ):
         _check_kv_layout(kv_layout)
         self._kv_layout = kv_layout
@@ -46,12 +47,12 @@ class BatchAttention:
         self.float_workspace_buffer = torch.empty(
             256 * 1024 * 1024,
             dtype=torch.uint8,
-            device=torch.device("cuda"),
+            device=torch.device(device),
         )
         self.int_workspace_buffer = torch.empty(
             8 * 1024 * 1024,
             dtype=torch.uint8,
-            device=torch.device("cuda"),
+            device=torch.device(device),
         )
         self.page_locked_int_workspace_buffer = torch.empty(
             8 * 1024 * 1024,