feat: Add Option For Fixed cta_tile_q

csrc/batch_prefill.cu

@@ -50,7 +50,7 @@ Array<int64_t> BatchPrefillWithKVCachePlan(
     TensorView kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads,
     int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk,
     int64_t head_dim_vo, bool causal, int64_t window_left, int64_t fixed_split_size,
-    bool disable_split_kv, int64_t num_colocated_ctas = 0) {
+    bool disable_split_kv, int64_t fixed_cta_tile_q = -1, int64_t num_colocated_ctas = 0) {
   size_t float_workspace_size_in_bytes =
       float_workspace_buffer.size(0) * get_element_size(float_workspace_buffer);
   size_t int_workspace_size_in_bytes =
@@ -66,8 +66,8 @@ Array<int64_t> BatchPrefillWithKVCachePlan(
       int_workspace_size_in_bytes, plan_info, static_cast<IdType*>(qo_indptr.data_ptr()),
       static_cast<IdType*>(kv_indptr.data_ptr()), total_num_rows, batch_size, num_qo_heads,
       num_kv_heads, head_dim_qk, head_dim_vo, page_size, enable_cuda_graph,
-      /*sizeof_dtype_o=*/2, window_left, fixed_split_size, disable_split_kv, num_colocated_ctas,
-      stream);
+      /*sizeof_dtype_o=*/2, window_left, fixed_split_size, disable_split_kv, fixed_cta_tile_q,
+      num_colocated_ctas, stream);
 
   TVM_FFI_ICHECK(status == cudaSuccess)
       << "Failed to plan prefill with error: " << cudaGetErrorString(status);

csrc/batch_prefill_jit_binding.cu

@@ -25,7 +25,7 @@ Array<int64_t> BatchPrefillWithKVCachePlan(
     TensorView kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads,
     int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk,
     int64_t head_dim_vo, bool causal, int64_t window_left, int64_t fixed_split_size,
-    bool disable_split_kv, int64_t num_colocated_ctas);
+    bool disable_split_kv, int64_t fixed_cta_tile_q, int64_t num_colocated_ctas);
 
 void BatchPrefillWithRaggedKVCacheRun(TensorView float_workspace_buffer,
                                       TensorView int_workspace_buffer, Array<int64_t> plan_info_vec,

flashinfer/decode.py

@@ -845,6 +845,7 @@ def plan(
         seq_lens: Optional[torch.Tensor] = None,
         fixed_split_size: Optional[int] = None,
         disable_split_kv: bool = False,
+        fixed_cta_tile_q: Optional[int] = None,
     ) -> None:
         r"""Plan batch decode for given problem specification.
 
@@ -903,6 +904,9 @@ def plan(
             and lead to a varied number of launched CTAs.
         disable_split_kv : bool,
             Whether to disable the split-kv for determinism in CUDA Graph, defaults to ``False``.
+        fixed_cta_tile_q : Optional[int]
+            Fixed CTA tile size for FA2 attention planning. Supported values are ``16``, ``64``, and
+            ``128``. Defaults to ``None`` (auto heuristic).
         Note
         ----
         The :meth:`plan` method should be called before any :meth:`run` or
@@ -981,6 +985,16 @@ def plan(
             )
         if fixed_split_size is None:
             fixed_split_size = -1
+        if fixed_cta_tile_q is None:
+            fixed_cta_tile_q = -1
+        elif fixed_cta_tile_q not in (16, 64, 128):
+            raise ValueError(
+                f"fixed_cta_tile_q should be one of {{16, 64, 128}}, got {fixed_cta_tile_q}"
+            )
+        elif fixed_cta_tile_q == 128 and head_dim >= 256:
+            raise ValueError(
+                f"fixed_cta_tile_q=128 is not supported with head_dim={head_dim} (requires head_dim < 256)"
+            )
 
         self._cached_q_data_type = q_data_type
         self._cached_kv_data_type = kv_data_type
@@ -1089,6 +1103,7 @@ def plan(
             if self._backend == "fa2":
                 args.append(fixed_split_size)
                 args.append(disable_split_kv)
+                args.append(fixed_cta_tile_q)
                 args.append(0)  # num_colocated_ctas
             self._plan_info = self._cached_module.plan(
                 *args,
@@ -2628,6 +2643,7 @@ def fast_decode_plan(
     fixed_split_size: Optional[int] = None,
     disable_split_kv: bool = False,
     global_override_indptr_cpu: Optional[torch.Tensor] = None,
+    fixed_cta_tile_q: Optional[int] = None,
 ) -> None:
     """
     A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for FlashInferMultiStepDraftBackend.
@@ -2656,6 +2672,16 @@ def fast_decode_plan(
         # Here we set fixed_split_size to -1 to avoid the assertion error in flashinfer's plan function
         if fixed_split_size is None:
             fixed_split_size = -1
+        if fixed_cta_tile_q is None:
+            fixed_cta_tile_q = -1
+        elif fixed_cta_tile_q not in (16, 64, 128):
+            raise ValueError(
+                f"fixed_cta_tile_q should be one of {{16, 64, 128}}, got {fixed_cta_tile_q}"
+            )
+        elif fixed_cta_tile_q == 128 and head_dim >= 256:
+            raise ValueError(
+                f"fixed_cta_tile_q=128 is not supported with head_dim={head_dim} (requires head_dim < 256)"
+            )
 
     if self.is_cuda_graph_enabled:
         if batch_size != self._fixed_batch_size:
@@ -2718,7 +2744,7 @@ def fast_decode_plan(
             kv_lens_arr_host = get_seq_lens(indptr_host, last_page_len_host, page_size)
 
             try:
-                # Make sure we pass exactly 19 arguments for fa2 backend and 16 arguments for fa3 backend
+                # Make sure we pass exactly 20 arguments for fa2 backend and 16 arguments for fa3 backend
                 args = [
                     self._float_workspace_buffer,
                     self._int_workspace_buffer,
@@ -2740,6 +2766,7 @@ def fast_decode_plan(
                 if self._backend == "fa2":
                     args.append(fixed_split_size)
                     args.append(disable_split_kv)
+                    args.append(fixed_cta_tile_q)
                     args.append(0)  # num_colocated_ctas
                 self._plan_info = self._cached_module.plan(
                     *args,

flashinfer/pod.py

@@ -422,6 +422,7 @@ def plan(
             window_left,
             -1,  # fixed_split_size
             False,  # disable_split_kv
+            -1,  # fixed_cta_tile_q
             0,  # num_colocated_ctas
         )
 
@@ -977,6 +978,7 @@ def plan(
             window_left,
             -1,  # fixed_split_size
             False,  # disable_split_kv
+            -1,  # fixed_cta_tile_q
             0,  # num_colocated_ctas
         )
 
@@ -1003,6 +1005,7 @@ def plan(
             window_left,
             -1,  # fixed_split_size
             False,  # disable_split_kv
+            -1,  # fixed_cta_tile_q
             num_colocated_ctas,
         )
         self._indptr_type = kv_indptr_p.dtype

flashinfer/prefill.py

@@ -1675,6 +1675,7 @@ def plan(
         max_sequence_kv: Optional[int] = None,
         fixed_split_size: Optional[int] = None,
         disable_split_kv: bool = False,
+        fixed_cta_tile_q: Optional[int] = None,
     ) -> None:
         r"""Plan batch prefill/append attention on Paged KV-Cache for given problem specification.
 
@@ -1785,6 +1786,9 @@ def plan(
             and lead to a varied number of launched CTAs.
         disable_split_kv : bool,
             Whether to disable the split-kv for determinism in CUDA Graph, defaults to ``False``.
+        fixed_cta_tile_q : Optional[int]
+            Fixed CTA tile size for FA2 prefill. Supported values are ``16``, ``64``, and ``128``.
+            Defaults to ``None`` (auto heuristic).
         Note
         ----
         The :meth:`plan` method should be called before any :meth:`run` or
@@ -1811,6 +1815,16 @@ def plan(
             head_dim_vo = head_dim_qk
         if fixed_split_size is None:
             fixed_split_size = -1
+        if fixed_cta_tile_q is None:
+            fixed_cta_tile_q = -1
+        elif fixed_cta_tile_q not in (16, 64, 128):
+            raise ValueError(
+                f"fixed_cta_tile_q should be one of {{16, 64, 128}}, got {fixed_cta_tile_q}"
+            )
+        elif fixed_cta_tile_q == 128 and head_dim_vo >= 256:
+            raise ValueError(
+                f"fixed_cta_tile_q=128 is not supported with head_dim={head_dim_vo} (requires head_dim < 256)"
+            )
 
         batch_size = len(qo_indptr) - 1
         self._batch_size = batch_size
@@ -2018,6 +2032,7 @@ def plan(
             if self._backend == "fa2":
                 args.append(fixed_split_size or -1)  # fixed_split_size
                 args.append(disable_split_kv)  # disable_split_kv
+                args.append(fixed_cta_tile_q)  # fixed_cta_tile_q
                 args.append(0)  # num_colocated_ctas
             self._plan_info = self._cached_module.plan(
                 *args,
@@ -2657,6 +2672,7 @@ def plan(
         max_sequence_kv: Optional[int] = None,
         v_indptr: Optional[torch.Tensor] = None,
         o_indptr: Optional[torch.Tensor] = None,
+        fixed_cta_tile_q: Optional[int] = None,
     ) -> None:
         r"""Plan batch prefill/append attention on Ragged KV-Cache for given problem specification.
 
@@ -2751,6 +2767,9 @@ def plan(
             and lead to a varied number of launched CTAs.
         disable_split_kv : bool,
             Whether to disable the split-kv for determinism in CUDA Graph, defaults to ``False``.
+        fixed_cta_tile_q : Optional[int]
+            Fixed CTA tile size for FA2 prefill. Supported values are ``16``, ``64``, and ``128``.
+            Defaults to ``None`` (auto heuristic).
         seq_lens: Optional[torch.Tensor]
             A uint32 1D tensor indicating the kv sequence length of each prompt. shape: ``[batch_size]``.
         seq_lens_q: Optional[torch.Tensor]
@@ -2787,6 +2806,16 @@ def plan(
             head_dim_vo = head_dim_qk
         if fixed_split_size is None:
             fixed_split_size = -1
+        if fixed_cta_tile_q is None:
+            fixed_cta_tile_q = -1
+        elif fixed_cta_tile_q not in (16, 64, 128):
+            raise ValueError(
+                f"fixed_cta_tile_q should be one of {{16, 64, 128}}, got {fixed_cta_tile_q}"
+            )
+        elif fixed_cta_tile_q == 128 and head_dim_vo >= 256:
+            raise ValueError(
+                f"fixed_cta_tile_q=128 is not supported with head_dim={head_dim_vo} (requires head_dim < 256)"
+            )
         if logits_soft_cap is None:
             logits_soft_cap = 0.0
 
@@ -2943,6 +2972,7 @@ def plan(
             if self._backend == "fa2":
                 args.append(fixed_split_size or -1)  # fixed_split_size
                 args.append(disable_split_kv)  # disable_split_kv
+                args.append(fixed_cta_tile_q)  # fixed_cta_tile_q
                 args.append(0)  # num_colocated_ctas
             self._plan_info = self._cached_module.plan(
                 *args,

flashinfer/sparse.py

@@ -443,6 +443,7 @@ def plan(
             if self._backend == "fa2":
                 args.append(-1)  # fixed_split_size
                 args.append(False)  # disable_split_kv
+                args.append(-1)  # fixed_cta_tile_q
                 args.append(0)  # num_colocated_ctas
             self._plan_info = self._cached_module.plan(
                 *args,
@@ -986,6 +987,7 @@ def _block_mask_map_to_expanded_indices(
         if self._backend == "fa2":
             args.append(-1)  # fixed_split_size
             args.append(False)  # disable_split_kv
+            args.append(-1)  # fixed_cta_tile_q
             args.append(0)  # num_colocated_ctas
         self._plan_info = self._cached_module.plan(
             *args,

include/flashinfer/attention/scheduler.cuh

@@ -497,7 +497,8 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h,
                                    uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t head_dim,
                                    uint32_t page_size, uint32_t max_batch_size_if_split,
                                    bool enable_cuda_graph, int32_t window_left,
-                                   int32_t fixed_split_size, bool disable_split_kv) {
+                                   int32_t fixed_split_size, bool disable_split_kv,
+                                   int32_t fixed_cta_tile_q) {
   std::vector<IdType> request_indices, qo_tile_indices, kv_tile_indices, merge_indptr, o_indptr;
   merge_indptr.push_back(0);
   o_indptr.push_back(0);
@@ -527,27 +528,42 @@ inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h,
   const uint32_t min_kv_chunk_size = std::max((128 / page_size), 1U);
   uint32_t cta_tile_q;
   uint32_t total_num_tiles_q;
-  if (enable_cuda_graph) {
+  if (fixed_cta_tile_q > 0) {
+    if (fixed_cta_tile_q != 16 && fixed_cta_tile_q != 64 && fixed_cta_tile_q != 128) {
+      std::ostringstream err_msg;
+      err_msg << "fixed_cta_tile_q should be one of {16, 64, 128}, but got " << fixed_cta_tile_q;
+      FLASHINFER_ERROR(err_msg.str());
+    }
+    if (fixed_cta_tile_q == 128 && head_dim >= 256) {
+      std::ostringstream err_msg;
+      err_msg << "fixed_cta_tile_q=128 is not supported with head_dim=" << head_dim
+              << " (requires head_dim < 256)";
+      FLASHINFER_ERROR(err_msg.str());
+    }
+    cta_tile_q = fixed_cta_tile_q;
+  } else if (enable_cuda_graph) {
     // When CUDA graphs are enabled, the lengths of sequences determined by
     // qo_indptr_h can vary. We assume that the dummy data based on which
     // the CUDA graph is created fixes the maximum number of tokens.
     const uint64_t max_seq_len = total_num_rows - batch_size + 1;
     uint64_t max_qo_len = uint64_t(max_seq_len) * gqa_group_size;
     cta_tile_q = FA2DetermineCtaTileQ(max_qo_len, head_dim);
-
-    // Find an upper bound for the number of tiles, derived from the total
-    // number of rows and the batch size.  The sum of qo lengths rounded
-    // up to cta_tile_q will not exceed this number derived from the total
-    // number of rows.
-    total_num_tiles_q = ceil_div(total_num_rows * gqa_group_size, cta_tile_q) + batch_size - 1;
   } else {
     int64_t sum_packed_qo_len = 0;
     for (uint32_t i = 0; i < batch_size; ++i) {
       sum_packed_qo_len += packed_qo_len_arr[i];
     }
     const int64_t avg_packed_qo_len = sum_packed_qo_len / batch_size;
     cta_tile_q = FA2DetermineCtaTileQ(avg_packed_qo_len, head_dim);
+  }
 
+  if (enable_cuda_graph) {
+    // Find an upper bound for the number of tiles, derived from the total
+    // number of rows and the batch size. The sum of qo lengths rounded
+    // up to cta_tile_q will not exceed this number derived from the total
+    // number of rows.
+    total_num_tiles_q = ceil_div(total_num_rows * gqa_group_size, cta_tile_q) + batch_size - 1;
+  } else {
     total_num_tiles_q = 0;
     for (uint32_t i = 0; i < batch_size; ++i) {
       total_num_tiles_q += ceil_div(packed_qo_len_arr[i], cta_tile_q);
@@ -699,6 +715,7 @@ inline cudaError_t PrefillPlan(void* float_buffer, size_t float_workspace_size_i
                                uint32_t head_dim_qk, uint32_t head_dim_vo, uint32_t page_size,
                                bool enable_cuda_graph, uint32_t sizeof_dtype_o, int32_t window_left,
                                int32_t fixed_split_size, bool disable_split_kv,
+                               int32_t fixed_cta_tile_q,
                                int64_t num_colocated_ctas,  // for POD attention, limit prefill
                                                             // splits by #colocated decode CTAs
                                cudaStream_t stream) {
@@ -724,7 +741,8 @@ inline cudaError_t PrefillPlan(void* float_buffer, size_t float_workspace_size_i
         qo_tile_indices_vec, kv_tile_indices_vec, merge_indptr_vec, o_indptr_vec] =
       PrefillSplitQOKVIndptr(qo_indptr_h, kv_indptr_h, total_num_rows, batch_size, num_qo_heads,
                              num_kv_heads, head_dim_vo, page_size, max_batch_size_if_split,
-                             enable_cuda_graph, window_left, fixed_split_size, disable_split_kv);
+                             enable_cuda_graph, window_left, fixed_split_size, disable_split_kv,
+                             fixed_cta_tile_q);
 
   plan_info.cta_tile_q = cta_tile_q;
   plan_info.total_num_rows = total_num_rows;

tests/attention/test_batch_invariant_fa2.py

@@ -58,6 +58,7 @@ def warmup_jit():
 @pytest.mark.parametrize("kv_len", [4096, 8192, 5000])
 @pytest.mark.parametrize("fixed_split_size", [2048])
 @pytest.mark.parametrize("disable_split_kv", [True, False])
+@pytest.mark.parametrize("fixed_cta_tile_q", [16, 64, 128])
 @pytest.mark.parametrize("page_size", [1, 8, 16])
 @pytest.mark.parametrize("num_kv_heads", [4])
 @pytest.mark.parametrize("group_size", [1, 4, 8])
@@ -70,13 +71,16 @@ def test_batch_decode_tensor_cores(
     kv_len: int,
     fixed_split_size: int,
     disable_split_kv: bool,
+    fixed_cta_tile_q: int,
     page_size: int,
     num_kv_heads: int,
     group_size: int,
     head_dim: int,
     kv_layout: str,
     pos_encoding_mode: str,
 ):
+    if head_dim >= 256 and fixed_cta_tile_q == 128:
+        pytest.skip("fixed_cta_tile_q=128 is not supported with head_dim >= 256")
     num_qo_heads = num_kv_heads * group_size
     q = torch.randn(
         batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=torch.float16
@@ -133,6 +137,7 @@ def test_batch_decode_tensor_cores(
         q_data_type=torch.float16,
         fixed_split_size=fixed_split_size if not disable_split_kv else None,
         disable_split_kv=disable_split_kv,
+        fixed_cta_tile_q=fixed_cta_tile_q,
     )
     o_tensor_cores, lse_tensor_cores = wrapper_tensor_cores.run(
         q, kv_data, return_lse=True
@@ -153,6 +158,7 @@ def test_batch_decode_tensor_cores(
         q_data_type=torch.float16,
         fixed_split_size=fixed_split_size if not disable_split_kv else None,
         disable_split_kv=disable_split_kv,
+        fixed_cta_tile_q=fixed_cta_tile_q,
     )
     o_tensor_cores_invariant, lse_tensor_cores_invariant = wrapper_tensor_cores.run(
         q[:invariant_bs], kv_data, return_lse=True
@@ -195,6 +201,7 @@ def test_batch_decode_tensor_cores(
 @pytest.mark.parametrize("qo_len", [128, 256])
 @pytest.mark.parametrize("fixed_split_size", [2048])
 @pytest.mark.parametrize("disable_split_kv", [True, False])
+@pytest.mark.parametrize("fixed_cta_tile_q", [16, 64, 128])
 @pytest.mark.parametrize("page_size", [1, 8, 16])
 @pytest.mark.parametrize("num_kv_heads", [4])
 @pytest.mark.parametrize("group_size", [1, 4, 8])
@@ -208,13 +215,16 @@ def test_batch_prefill_tensor_cores(
     qo_len: int,
     fixed_split_size: int,
     disable_split_kv: bool,
+    fixed_cta_tile_q: int,
     page_size: int,
     num_kv_heads: int,
     group_size: int,
     head_dim: int,
     kv_layout: str,
     pos_encoding_mode: str,
 ):
+    if head_dim >= 256 and fixed_cta_tile_q == 128:
+        pytest.skip("fixed_cta_tile_q=128 is not supported with head_dim >= 256")
     num_qo_heads = num_kv_heads * group_size
     q = torch.randn(
         batch_size * qo_len,
@@ -281,6 +291,7 @@ def test_batch_prefill_tensor_cores(
         kv_data_type=torch.float16,
         fixed_split_size=fixed_split_size if not disable_split_kv else None,
         disable_split_kv=disable_split_kv,
+        fixed_cta_tile_q=fixed_cta_tile_q,
     )
     o_tensor_cores, lse_tensor_cores = wrapper_tensor_cores.run(
         q, kv_data, return_lse=True
@@ -304,6 +315,7 @@ def test_batch_prefill_tensor_cores(
         kv_data_type=torch.float16,
         fixed_split_size=fixed_split_size if not disable_split_kv else None,
         disable_split_kv=disable_split_kv,
+        fixed_cta_tile_q=fixed_cta_tile_q,
     )
     o_tensor_cores_invariant, lse_tensor_cores_invariant = wrapper_tensor_cores.run(
         q[: invariant_bs * qo_len], kv_data, return_lse=True