feat: generate and vendor flashinfer kernels

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+.bak
+__pycache__
+result

README.md

@@ -0,0 +1,10 @@
+---
+license: apache-2.0
+tags:
+  - kernel
+---
+
+
+This kernel is a work in progress and requires more work to correctly add all of the FlashInfer kernels.
+
+Please see the [generate-source.md](generate-source.md) for instructions on how to generate the source files that are contained in this kernel.

build.toml

@@ -0,0 +1,224 @@
+[general]
+name = "flashinfer"
+universal = false
+
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h"
+]
+
+[kernel.flashinfer]
+backend = "cuda"
+cuda-capabilities = [
+    # "7.0",
+    # "7.2",
+    "7.5",
+    "8.0",
+    "8.6",
+    "8.7",
+    "8.9",
+    "9.0",
+    # "10.0",
+    # "10.1",
+    # "11.8",
+    # "12.0"
+]
+cuda-flags = [
+    "-O3",
+    "-std=c++17",
+    "--use_fast_math",
+    "--expt-relaxed-constexpr",
+    "--expt-extended-lambda",
+    "-DFLASHINFER_ENABLE_F16",
+    "-DFLASHINFER_ENABLE_BF16",
+    "-DFLASHINFER_ENABLE_FP8_E4M3",
+    "-DFLASHINFER_ENABLE_FP8_E5M2",
+    "-DNDEBUG",
+]
+cxx-flags = [
+    "-DFLASHINFER_ENABLE_F16",
+    "-DFLASHINFER_ENABLE_BF16",
+    "-DFLASHINFER_ENABLE_FP8_E4M3",
+    "-DFLASHINFER_ENABLE_FP8_E5M2",
+]
+include = [ "include" ]
+depends = ["torch", "cutlass_3_8"]
+src = [
+    "include/flashinfer/page.cuh",
+    "include/flashinfer/arch_condition.h",
+    "include/flashinfer/cubin_loader.h",
+    "include/flashinfer/permuted_smem.cuh",
+    "include/flashinfer/exception.h",
+    "include/flashinfer/trtllm/common.h",
+    "include/flashinfer/trtllm/fmha/fmhaRunnerParams.h",
+    "include/flashinfer/trtllm/fmha/decoder_impl_common.h",
+    "include/flashinfer/trtllm/fmha/decoder_params.h",
+    "include/flashinfer/trtllm/fmha/kernelParams.h",
+    "include/flashinfer/trtllm/fmha/gen_kernel_launcher.cuh",
+    "include/flashinfer/trtllm/fmha/fmhaKernels.cuh",
+    "include/flashinfer/trtllm/fmha/cubin/kernelMetaInfo.h",
+    "include/flashinfer/trtllm/fmha/fmhaRunner.cuh",
+    "include/flashinfer/trtllm/common/cudaTypeUtils.cuh",
+    "include/flashinfer/trtllm/common/cudaUtils.h",
+    "include/flashinfer/trtllm/common/cudaBf16Wrapper.h",
+    "include/flashinfer/trtllm/common/cudaFp8Utils.h",
+    "include/flashinfer/trtllm/common/cudaBf16Fallbacks.cuh",
+    "include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh",
+    "include/flashinfer/trtllm/fused_moe/DevKernel.h",
+    "include/flashinfer/trtllm/fused_moe/RoutingKernel.h",
+    "include/flashinfer/trtllm/fused_moe/runner.h",
+    "include/flashinfer/trtllm/fused_moe/IntFastDiv.h",
+    "include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh",
+    "include/flashinfer/trtllm/batched_gemm/KernelRunner.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelMetaInfo.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h",
+    "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/MmaDecl.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/SfLayoutDecl.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/DtypeDecl.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CommonUtils.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CudaKernelLauncher.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/TmaDescriptor.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelTraits.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmInterface.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelParams.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/Enums.h",
+    "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmOptions.h",
+    "include/flashinfer/page.cuh",
+    "include/flashinfer/vec_dtypes.cuh",
+    "include/flashinfer/sampling.cuh",
+    "include/flashinfer/logging.h",
+    "include/flashinfer/fp16.h",
+    "include/flashinfer/attention_impl.cuh",
+    "include/flashinfer/allocator.h",
+    "include/flashinfer/utils.cuh",
+    "include/flashinfer/fastdiv.cuh",
+    "include/flashinfer/norm.cuh",
+    "include/flashinfer/layout.cuh",
+    "include/flashinfer/math.cuh",
+    "include/flashinfer/gemm/group_gemm_sm90.cuh",
+    "include/flashinfer/gemm/group_gemv.cuh",
+    "include/flashinfer/gemm/fp8_gemm_cutlass_template.h",
+    "include/flashinfer/gemm/group_gemm_mxfp4_groupwise_sm100.cuh",
+    "include/flashinfer/gemm/fp4_gemm_cutlass.h",
+    "include/flashinfer/gemm/fp4_gemm_cutlass_template.h",
+    "include/flashinfer/gemm/cutlass_gemm_configs.h",
+    "include/flashinfer/gemm/fp4_gemm_template_sm100.h",
+    "include/flashinfer/gemm/group_gemm_fp8_groupwise_sm100.cuh",
+    "include/flashinfer/gemm/gemm_groupwise_sm100.cuh",
+    "include/flashinfer/gemm/group_gemm_lora.cuh",
+    "include/flashinfer/gemm/group_gemm.cuh",
+    "include/flashinfer/gemm/fp8_gemm_cutlass.h",
+    "include/flashinfer/gemm/fp8_gemm_template_sm100.h",
+    "include/flashinfer/gemm/bmm_fp8.cuh",
+    "include/flashinfer/activation.cuh",
+    "include/flashinfer/semaphore_utils.cuh",
+    "include/flashinfer/quantization.cuh",
+    "include/flashinfer/mma.cuh",
+    "include/flashinfer/attention/variant_helper.cuh",
+    "include/flashinfer/attention/hopper.cuh",
+    "include/flashinfer/attention/prefill.cuh",
+    "include/flashinfer/attention/persistent_template.cuh",
+    "include/flashinfer/attention/mla_hopper.cuh",
+    "include/flashinfer/attention/pod.cuh",
+    "include/flashinfer/attention/variants.cuh",
+    "include/flashinfer/attention/decode.cuh",
+    "include/flashinfer/attention/heap.h",
+    "include/flashinfer/attention/state.cuh",
+    "include/flashinfer/attention/mask.cuh",
+    "include/flashinfer/attention/default_prefill_params.cuh",
+    "include/flashinfer/attention/cutlass_mla.cuh",
+    "include/flashinfer/attention/scheduler.cuh",
+    "include/flashinfer/attention/blackwell/kernel/fmha_options.hpp",
+    "include/flashinfer/attention/blackwell/kernel/gather_tensor.hpp",
+    "include/flashinfer/attention/blackwell/kernel/fmha_tile_scheduler.hpp",
+    "include/flashinfer/attention/blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_tma_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/kernel/sm100_mla_tile_scheduler.hpp",
+    "include/flashinfer/attention/blackwell/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_reduction.hpp",
+    "include/flashinfer/attention/blackwell/plan.cuh",
+    "include/flashinfer/attention/blackwell/common/pow_2.hpp",
+    "include/flashinfer/attention/blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_mainloop_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/collective/sm100_fmha_load_cpasync_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_epilogue_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp",
+    "include/flashinfer/attention/blackwell/collective/fmha_fusion.hpp",
+    "include/flashinfer/attention/blackwell/collective/fmha_common.hpp",
+    "include/flashinfer/attention/blackwell/device/sm100_mla.hpp",
+    "include/flashinfer/attention/blackwell/device/fmha.hpp",
+    "include/flashinfer/attention/blackwell/fmha_cutlass_sm100.cuh",
+    "include/flashinfer/attention/mla.cuh",
+    "include/flashinfer/attention/decode_mla_cute_sm80.cuh",
+    "include/flashinfer/attention/cascade.cuh",
+    "include/flashinfer/attention/hopper/variant_helper.cuh",
+    "include/flashinfer/attention/hopper/default_params.cuh",
+    "include/flashinfer/attention/hopper/attention_updater.cuh",
+    "include/flashinfer/attention/hopper/epilogue.cuh",
+    "include/flashinfer/attention/hopper/variants.cuh",
+    "include/flashinfer/attention/hopper/mainloop.cuh",
+    "include/flashinfer/attention/hopper/block_sparse_gather.cuh",
+    "include/flashinfer/attention/hopper/utils.cuh",
+    "include/flashinfer/attention/hopper/prefill_sm90.cuh",
+    "include/flashinfer/attention/hopper/named_barrier.cuh",
+    "include/flashinfer/attention/hopper/tile_scheduler.cuh",
+    "include/flashinfer/attention/hopper/sparse_mainloop.cuh",
+    "include/flashinfer/attention/hopper/kernel_traits.cuh",
+    "include/flashinfer/attention/hopper/mainloop_mma.cuh",
+    "include/flashinfer/attention/hopper/quantization/epilogue.cuh",
+    "include/flashinfer/attention/hopper/quantization/prefill_sm90.cuh",
+    "include/flashinfer/attention/hopper/quantization/mainloop_load.cuh",
+    "include/flashinfer/attention/hopper/quantization/mainloop_sparse_load.cuh",
+    "include/flashinfer/attention/hopper/quantization/kernel_traits.cuh",
+    "include/flashinfer/attention/hopper/quantization/mainloop_mma.cuh",
+    "include/flashinfer/attention/default_decode_params.cuh",
+    "include/flashinfer/attention/mla_params.cuh",
+    "include/flashinfer/attention/persistent.cuh",
+    "include/flashinfer/pos_enc.cuh",
+    "include/flashinfer/cutlass_utils.cuh",
+    "include/flashinfer/comm/trtllm_moe_allreduce_fusion.cuh",
+    "include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh",
+    "include/flashinfer/comm/vllm_custom_all_reduce.cuh",
+    "include/flashinfer/comm/trtllm_allreduce.cuh",
+    "include/flashinfer/comm/trtllm_allreduce_fusion.cuh",
+    "include/flashinfer/comm/trtllm_alltoall.cuh",
+    "include/flashinfer/frag_layout_swizzle.cuh",
+    "include/flashinfer/profiler.cuh",
+    "include/flashinfer/cp_async.cuh",
+
+    "include/pytorch_conversion_utils.h",
+    "include/pytorch_extension_utils.h",
+
+    # # batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc",
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu",
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu",
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu",
+
+    # # batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc",
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu",
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu",
+    # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu",
+    
+    
+    "csrc/generated/gelu_and_mul.cu",
+    "csrc/generated/gelu_tanh_and_mul.cu",
+    "csrc/generated/silu_and_mul.cu",
+]

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, false, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = nv_bfloat16;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 128;
+constexpr int HEAD_DIM_VO = 128;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = false;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<128, PosEncodingMode::kNone, DefaultAttention<false, false, false, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, true, true, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = nv_bfloat16;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 256;
+constexpr int HEAD_DIM_VO = 256;
+constexpr auto USE_LOGITS_SOFT_CAP = true;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = true;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<256, PosEncodingMode::kNone, DefaultAttention<false, true, true, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, false, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = nv_bfloat16;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 64;
+constexpr int HEAD_DIM_VO = 64;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = false;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention<false, false, false, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, true, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = nv_bfloat16;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 64;
+constexpr int HEAD_DIM_VO = 64;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = true;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention<false, true, false, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, false, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = __nv_fp8_e4m3;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 128;
+constexpr int HEAD_DIM_VO = 128;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = false;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<128, PosEncodingMode::kNone, DefaultAttention<false, false, false, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, true, true, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = __nv_fp8_e4m3;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 256;
+constexpr int HEAD_DIM_VO = 256;
+constexpr auto USE_LOGITS_SOFT_CAP = true;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = true;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<256, PosEncodingMode::kNone, DefaultAttention<false, true, true, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, false, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = __nv_fp8_e4m3;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 64;
+constexpr int HEAD_DIM_VO = 64;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = false;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention<false, false, false, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, true, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = nv_bfloat16;
+using DTypeKV = __nv_fp8_e4m3;
+using DTypeO = nv_bfloat16;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 64;
+constexpr int HEAD_DIM_VO = 64;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = true;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention<false, true, false, false>, Params>(
+    Params params, nv_bfloat16* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, false, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = half;
+using DTypeKV = __nv_fp8_e4m3;
+using DTypeO = half;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 128;
+constexpr int HEAD_DIM_VO = 128;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = false;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<128, PosEncodingMode::kNone, DefaultAttention<false, false, false, false>, Params>(
+    Params params, half* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, true, true, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = half;
+using DTypeKV = __nv_fp8_e4m3;
+using DTypeO = half;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 256;
+constexpr int HEAD_DIM_VO = 256;
+constexpr auto USE_LOGITS_SOFT_CAP = true;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = true;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<256, PosEncodingMode::kNone, DefaultAttention<false, true, true, false>, Params>(
+    Params params, half* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu

@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <flashinfer/attention/scheduler.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/utils.cuh>
+#include <optional>
+
+#include "batch_decode_config.inc"
+#include "pytorch_conversion_utils.h"
+#include "pytorch_extension_utils.h"
+
+namespace flashinfer {
+
+template <uint32_t HEAD_DIM, PosEncodingMode POS_ENCODING_MODE, typename AttentionVariant,
+          typename Params>
+cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v,
+                                                  float* tmp_s, bool enable_pdl,
+                                                  cudaStream_t stream);
+
+}  // namespace flashinfer
+
+using namespace flashinfer;
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data) {
+  size_t float_workspace_size_in_bytes =
+      float_workspace_buffer.size(0) * float_workspace_buffer.element_size();
+  size_t int_workspace_size_in_bytes =
+      int_workspace_buffer.size(0) * int_workspace_buffer.element_size();
+
+  DecodePlanInfo plan_info;
+
+  auto q_scalar_type = empty_q_data.scalar_type();
+  auto kv_scalar_type = empty_kv_data.scalar_type();
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device());
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, {
+          auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched<
+              GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>;
+          cudaError_t status = DecodePlan<HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>(
+              static_cast<void*>(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes,
+              static_cast<void*>(int_workspace_buffer.data_ptr()),
+              static_cast<void*>(page_locked_int_workspace_buffer.data_ptr()),
+              int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(indptr.data_ptr()),
+              batch_size, num_qo_heads, page_size, enable_cuda_graph,
+              /*stream=*/stream, work_estimation_func);
+
+          TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                      cudaGetErrorString(status));
+          return true;
+        });
+      });
+
+  return vec_to_tensor(plan_info.ToVector());
+}
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS) {
+  DecodePlanInfo plan_info;
+  plan_info.FromVector(tensor_to_vec(plan_info_vec));
+  QKVLayout kv_layout = static_cast<QKVLayout>(kv_layout_code);
+  auto device = q.device();
+  int64_t batch_size = q.size(0);
+  int64_t num_qo_heads = q.size(1);
+  int64_t num_kv_heads, page_size;
+
+  if (kv_layout == QKVLayout::kHND) {
+    num_kv_heads = paged_k_cache.size(1);
+    page_size = paged_k_cache.size(2);
+  } else {
+    page_size = paged_k_cache.size(1);
+    num_kv_heads = paged_k_cache.size(2);
+  }
+  uint32_t head_dim_qk = q.size(2);
+  uint32_t head_dim_vo = paged_v_cache.size(3);
+
+  TORCH_CHECK(head_dim_qk == head_dim_vo,
+              "CUDA cores template only supports equal head dim for QK and VO, please use tensor "
+              "cores template for different head dim");
+
+  if (maybe_lse) {
+    const auto& lse = *maybe_lse;
+    TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0));
+    TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1));
+  }
+
+  void* float_buffer = static_cast<void*>(float_workspace_buffer.data_ptr());
+  void* int_buffer = static_cast<void*>(int_workspace_buffer.data_ptr());
+
+  // get q_scalar_type and kv_scalar_type
+  auto q_scalar_type = q.scalar_type();
+  auto kv_scalar_type = paged_k_cache.scalar_type();
+
+  // get q_stride_n and q_stride_h
+  const auto q_stride_n = q.stride(0);
+  const auto q_stride_h = q.stride(1);
+
+  // get kv_cache_strides
+  const int64_t* kv_cache_strides = nullptr;
+  auto k_strides = paged_k_cache.strides();
+  auto v_strides = paged_v_cache.strides();
+  TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical");
+  kv_cache_strides = k_strides.data();
+
+  const c10::cuda::OptionalCUDAGuard device_guard(device);
+  const cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
+
+  DISPATCH_context(
+      DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE,
+      USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] {
+        paged_kv_t<DTypeKV, IdType> paged_kv(
+            num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout,
+            static_cast<DTypeKV*>(paged_k_cache.data_ptr()),
+            static_cast<DTypeKV*>(paged_v_cache.data_ptr()), kv_cache_strides,
+            static_cast<IdType*>(paged_kv_indices.data_ptr()),
+            static_cast<IdType*>(paged_kv_indptr.data_ptr()),
+            static_cast<IdType*>(paged_kv_last_page_len.data_ptr()));
+
+        Params params;
+        params.q = static_cast<DTypeQ*>(q.data_ptr());
+        params.paged_kv = paged_kv;
+        params.o = static_cast<DTypeO*>(o.data_ptr());
+        params.lse = maybe_lse ? static_cast<float*>(maybe_lse->data_ptr()) : nullptr;
+        params.padded_batch_size = 0;
+        params.num_qo_heads = num_qo_heads;
+        params.q_stride_n = q_stride_n;
+        params.q_stride_h = q_stride_h;
+        params.window_left = window_left;
+        params.request_indices = nullptr;
+        params.kv_tile_indices = nullptr;
+        params.o_indptr = nullptr;
+        params.kv_chunk_size_ptr = nullptr;
+        params.block_valid_mask = nullptr;
+        params.partition_kv = false;
+
+        ADDITIONAL_PARAMS_SETTER
+
+        DTypeO* tmp_v = nullptr;
+        float* tmp_s = nullptr;
+        params.request_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.request_indices_offset);
+        params.kv_tile_indices =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_tile_indices_offset);
+        params.o_indptr = GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.o_indptr_offset);
+        params.kv_chunk_size_ptr =
+            GetPtrFromBaseOffset<IdType>(int_buffer, plan_info.kv_chunk_size_ptr_offset);
+        if (plan_info.split_kv) {
+          tmp_v = GetPtrFromBaseOffset<DTypeO>(float_buffer, plan_info.v_offset);
+          tmp_s = GetPtrFromBaseOffset<float>(float_buffer, plan_info.s_offset);
+          if (plan_info.enable_cuda_graph) {
+            params.block_valid_mask =
+                GetPtrFromBaseOffset<bool>(int_buffer, plan_info.block_valid_mask_offset);
+          }
+        }
+        params.padded_batch_size = plan_info.padded_batch_size;
+
+        cudaError_t status =
+            flashinfer::BatchDecodeWithPagedKVCacheDispatched<HEAD_DIM_QK, POS_ENCODING_MODE,
+                                                              AttentionVariant>(params, tmp_v,
+                                                                                tmp_s, enable_pdl,
+                                                                                /*stream=*/stream);
+        TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ",
+                    cudaGetErrorString(status));
+        return true;
+      });
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc

@@ -0,0 +1,71 @@
+#pragma once
+#include <flashinfer/page.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/pos_enc.cuh>
+#include <flashinfer/attention/variant_helper.cuh>
+
+#define ADDITIONAL_FUNC_PARAMS , std::optional<at::Tensor> maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta
+#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast<float*>(maybe_alibi_slopes->data_ptr()): nullptr; \
+params.logits_soft_cap = logits_soft_cap; \
+params.sm_scale = sm_scale; \
+params.rope_rcp_scale = rope_rcp_scale; \
+params.rope_rcp_theta = rope_rcp_theta;
+
+#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \
+  using AttentionVariant = DefaultAttention<false, false, false, false>; \
+  __VA_ARGS__(); \
+}
+
+using namespace flashinfer;
+
+using DTypeQ = half;
+using DTypeKV = __nv_fp8_e4m3;
+using DTypeO = half;
+using IdType = int32_t;
+constexpr int HEAD_DIM_QK = 64;
+constexpr int HEAD_DIM_VO = 64;
+constexpr auto USE_LOGITS_SOFT_CAP = false;
+constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone;
+constexpr auto USE_SLIDING_WINDOW = false;
+
+struct Params {
+  using DTypeQ = DTypeQ;
+  using DTypeKV = DTypeKV;
+  using DTypeO = DTypeO;
+  using IdType = IdType;
+
+  DTypeQ* q;
+  paged_kv_t<DTypeKV, IdType> paged_kv;
+  DTypeO* o;
+  float* lse;
+
+  float* maybe_alibi_slopes;
+double logits_soft_cap;
+double sm_scale;
+double rope_rcp_scale;
+double rope_rcp_theta;
+
+
+  uint32_t padded_batch_size;
+  uint32_t num_qo_heads;
+  IdType q_stride_n;
+  IdType q_stride_h;
+  int32_t window_left;
+  bool enable_pdl;
+
+  IdType* request_indices;
+  IdType* kv_tile_indices;
+  IdType* o_indptr;
+  IdType* kv_chunk_size_ptr;
+  bool* block_valid_mask;
+  bool partition_kv;
+
+  __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; }
+
+  __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const {
+    return paged_kv.get_length(batch_idx);
+  }
+};
+
+#include<flashinfer/attention/variants.cuh>

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2023-2025 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "batch_decode_config.inc"
+#include "pytorch_extension_utils.h"
+
+at::Tensor BatchDecodeWithPagedKVCachePlan(
+    at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer,
+    at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size,
+    int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph,
+    int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo,
+    at::Tensor empty_q_data, at::Tensor empty_kv_data);
+
+void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer,
+                                    at::Tensor int_workspace_buffer, at::Tensor plan_info_vec,
+                                    at::Tensor q, at::Tensor paged_k_cache,
+                                    at::Tensor paged_v_cache, at::Tensor paged_kv_indptr,
+                                    at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len,
+                                    at::Tensor o, std::optional<at::Tensor> maybe_lse,
+                                    int64_t kv_layout_code, int64_t window_left,
+                                    bool enable_pdl ADDITIONAL_FUNC_PARAMS);
+
+TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) {
+  // Batched decode with paged KV-Cache plan
+  m.def("plan", BatchDecodeWithPagedKVCachePlan);
+  // Batched decode with paged KV-Cache run
+  m.def("run", BatchDecodeWithPagedKVCacheRun);
+}

csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu

@@ -0,0 +1,13 @@
+#include <flashinfer/attention/decode.cuh>
+#include "batch_decode_config.inc"
+
+using namespace flashinfer;
+
+namespace flashinfer {
+
+template cudaError_t
+BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention<false, false, false, false>, Params>(
+    Params params, half* tmp_v,
+    float* tmp_s, bool enable_pdl, cudaStream_t stream);
+
+};