Fix the type annotation of q_dtype and kv_dtype on ragged prefill (#798)

Author: nandor

flashinfer/prefill.py

@@ -1951,8 +1951,8 @@ def plan(
         sm_scale: Optional[float] = None,
         rope_scale: Optional[float] = None,
         rope_theta: Optional[float] = None,
-        q_data_type: str = "float16",
-        kv_data_type: Optional[str] = None,
+        q_data_type: Union[str, torch.dtype] = "float16",
+        kv_data_type: Optional[Union[str, torch.dtype]] = None,
     ) -> None:
         r"""Plan batch prefill/append attention on Ragged KV-Cache for given problem specification.