Add num_colocated_ctas to BlockSparseAttention

Aditya K Kamath · Aditya K Kamath · commit 6e4fedb6cd26 · 2025-11-13T00:10:34.000Z
diff --git a/flashinfer/sparse.py b/flashinfer/sparse.py
@@ -473,6 +473,7 @@ def plan(
             if self._backend == "fa2":
                 args.append(-1)  # fixed_split_size
                 args.append(False)  # disable_split_kv
+                args.append(0)  # num_colocated_ctas
             self._plan_info = self._cached_module.plan(
                 *args,
             )

Original file line number	Diff line number	Diff line change
`@@ -473,6 +473,7 @@ def plan(`
`473`	`473`	`if self._backend == "fa2":`
`474`	`474`	`args.append(-1) # fixed_split_size`
`475`	`475`	`args.append(False) # disable_split_kv`
	`476`	`+ args.append(0) # num_colocated_ctas`
`476`	`477`	`self._plan_info = self._cached_module.plan(`
`477`	`478`	`*args,`
`478`	`479`	`)`