chore: upgrade flashinfer 0.5.0 (#12523)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

python/pyproject.toml

@@ -26,7 +26,9 @@ dependencies = [
   "datasets",
   "einops",
   "fastapi",
-  "flashinfer_python==0.4.1",
+  "flashinfer_python==0.5.0",
+  "flashinfer_cubin==0.5.0",
+  "flashinfer_jit_cache==0.5.0",
   "gguf",
   "hf_transfer",
   "huggingface_hub",

python/sglang/check_env.py

@@ -22,6 +22,8 @@ PACKAGE_LIST = [
     "sglang",
     "sgl_kernel",
     "flashinfer_python",
+    "flashinfer_cubin",
+    "flashinfer_jit_cache",
     "triton",
     "transformers",
     "torchao",

python/sglang/srt/entrypoints/engine.py

@@ -712,7 +712,7 @@ def _set_envs_and_config(server_args: ServerArgs):
     if server_args.attention_backend == "flashinfer":
         assert_pkg_version(
             "flashinfer_python",
-            "0.4.1",
+            "0.5.0",
             "Please uninstall the old version and "
             "reinstall the latest version by following the instructions "
             "at https://docs.flashinfer.ai/installation.html.",

python/sglang/srt/utils/common.py

@@ -2386,7 +2386,9 @@ def set_cuda_arch():
     if is_flashinfer_available():
         capability = torch.cuda.get_device_capability()
         arch = f"{capability[0]}.{capability[1]}"
-        os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
+        os.environ["FLASHINFER_CUDA_ARCH_LIST"] = (
+            f"{arch}{'a' if capability[0] >= 9 else ''}"
+        )
 
 
 def next_power_of_2(n: int):