doc: remove misleading docstring about non_blocking (#966)

As noted in #965 , we have some misleading docstring about the use of
`non_blocking` option in plan functions of attention wrappers (they are
only necessary in our old designs which we deprecate later), this PR
fixes the issue.

Author: yzh119

flashinfer/decode.py

@@ -746,7 +746,7 @@ def plan(
         sm_scale: Optional[float] = None,
         rope_scale: Optional[float] = None,
         rope_theta: Optional[float] = None,
-        non_blocking: bool = False,
+        non_blocking: bool = True,
     ) -> None:
         r"""Plan batch decode for given problem specification.
 
@@ -789,8 +789,7 @@ def plan(
             The data type of both the query and key/value tensors. Defaults to torch.float16.
             data_type is deprecated, please use q_data_type and kv_data_type instead.
         non_blocking : bool
-            Whether to copy the input tensors to the device asynchronously, defaults to ``False``.
-            If ``True``, user should synchronize before calling :meth:`run` or cuda graph replay.
+            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.
 
 
         Note
@@ -823,12 +822,12 @@ def plan(
                     "The size of indices should be less than or equal to the allocated buffer"
                 )
             self._paged_kv_indptr_buf.copy_(indptr, non_blocking=non_blocking)
-            self._paged_kv_indices_buf[: len(indices)].copy_(
-                indices, non_blocking=non_blocking
-            )
             self._paged_kv_last_page_len_buf.copy_(
                 last_page_len, non_blocking=non_blocking
             )
+            self._paged_kv_indices_buf[: len(indices)].copy_(
+                indices, non_blocking=(indices.device == self.device) and non_blocking
+            )
         else:
             self._paged_kv_indptr_buf = indptr.to(
                 self.device, non_blocking=non_blocking

flashinfer/pod.py

@@ -298,7 +298,7 @@ def plan(
         sm_scale: Optional[float] = None,
         rope_scale: Optional[float] = None,
         rope_theta: Optional[float] = None,
-        non_blocking: bool = False,
+        non_blocking: bool = True,
     ) -> None:
         r"""Plan POD's batch decode for given problem specification.
 
@@ -335,8 +335,7 @@ def plan(
             The data type of both the query and key/value tensors. Defaults to torch.float16.
             data_type is deprecated, please use q_data_type and kv_data_type instead.
         non_blocking : bool
-            Whether to copy the input tensors to the device asynchronously, defaults to ``False``.
-            If ``True``, user should synchronize before calling :meth:`run` or cuda graph replay.
+            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.
 
 
         Note
@@ -371,12 +370,12 @@ def plan(
                     "The size of indices should be less than or equal to the allocated buffer"
                 )
             self._paged_kv_indptr_buf.copy_(indptr, non_blocking=non_blocking)
-            self._paged_kv_indices_buf[: len(indices)].copy_(
-                indices, non_blocking=non_blocking
-            )
             self._paged_kv_last_page_len_buf.copy_(
                 last_page_len, non_blocking=non_blocking
             )
+            self._paged_kv_indices_buf[: len(indices)].copy_(
+                indices, non_blocking=(indices.device == self.device) and non_blocking
+            )
         else:
             self._paged_kv_indptr_buf = indptr.to(
                 self.device, non_blocking=non_blocking

flashinfer/prefill.py

@@ -1195,7 +1195,7 @@ def plan(
         rope_theta: Optional[float] = None,
         q_data_type: Union[str, torch.dtype] = "float16",
         kv_data_type: Optional[Union[str, torch.dtype]] = None,
-        non_blocking: bool = False,
+        non_blocking: bool = True,
     ) -> None:
         r"""Plan batch prefill/append attention on Paged KV-Cache for given problem specification.
 
@@ -1269,8 +1269,7 @@ def plan(
         kv_data_type : Optional[Union[str, torch.dtype]]
             The data type of the key/value tensor. If None, will be set to :attr:`q_data_type`.
         non_blocking : bool
-            Whether to copy the input tensors to the device asynchronously, defaults to ``False``.
-            If ``True``, user should synchronize before calling :meth:`run` or cuda graph replay.
+            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.
 
         Note
         ----
@@ -1349,12 +1348,13 @@ def plan(
 
             self._qo_indptr_buf.copy_(qo_indptr, non_blocking=non_blocking)
             self._paged_kv_indptr_buf.copy_(paged_kv_indptr, non_blocking=non_blocking)
-            self._paged_kv_indices_buf[: len(paged_kv_indices)].copy_(
-                paged_kv_indices, non_blocking=non_blocking
-            )
             self._paged_kv_last_page_len_buf.copy_(
                 paged_kv_last_page_len, non_blocking=non_blocking
             )
+            self._paged_kv_indices_buf[: len(paged_kv_indices)].copy_(
+                paged_kv_indices,
+                non_blocking=(paged_kv_indices.device == self.device) and non_blocking,
+            )
 
             if packed_custom_mask is not None:
                 if not torch.is_tensor(self._custom_mask_buf):
@@ -1366,7 +1366,9 @@ def plan(
                         "mask_indptr_buf must be initialized with a torch.Tensor in cuda graph mode if we use custom mask in attention computation."
                     )
                 self._custom_mask_buf[: len(packed_custom_mask)].copy_(
-                    packed_custom_mask, non_blocking=non_blocking
+                    packed_custom_mask,
+                    non_blocking=(packed_custom_mask.device == self.device)
+                    and non_blocking,
                 )
                 # NOTE(Zihao): mask_indptr has the same length as qo_indptr
                 self._mask_indptr_buf.copy_(mask_indptr, non_blocking=non_blocking)

flashinfer/sparse.py

@@ -207,7 +207,7 @@ def plan(
         rope_theta: Optional[float] = None,
         q_data_type: Union[str, torch.dtype] = "float16",
         kv_data_type: Optional[Union[str, torch.dtype]] = None,
-        non_blocking: bool = False,
+        non_blocking: bool = True,
     ) -> None:
         r"""Create auxiliary data structures for block sparse attention.
 
@@ -270,8 +270,7 @@ def plan(
         kv_data_type : Optional[Union[str, torch.dtype]]
             The data type of the key/value tensor. If None, will be set to :attr:`q_data_type`.
         non_blocking : bool
-            Whether to copy the input tensors to the device asynchronously, defaults to ``False``.
-            If ``True``, user should synchronize before calling :meth:`run` or cuda graph replay.
+            Whether to copy the input tensors to the device asynchronously, defaults to ``True``.
 
 
         The :meth:`plan` method should be called before any :meth:`run` or
@@ -414,7 +413,7 @@ def plan(
 
             kv_lens_arr_host = (kv_indptr_host[1:] - kv_indptr_host[:-1]) * self.C
             self._kv_lens_buffer[: len(kv_lens_arr_host)].copy_(
-                kv_lens_arr_host, non_blocking=non_blocking
+                kv_lens_arr_host,
             )
 
             if self._backend == "fa3":