add cuda version check for jit (#1526)

<!-- .github/pull_request_template.md -->

## 📌 Description

Add cuda version check for option `-static-global-template-stub`.
`-static-global-template-stub` is only enabled for cuda >= 12.8.

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

cc: @zhyncs

Author: cyx-6

flashinfer/jit/cpp_ext.py

@@ -1,9 +1,12 @@
 # Adapted from https://github.com/pytorch/pytorch/blob/v2.7.0/torch/utils/cpp_extension.py
 
+import functools
 import os
+import re
 import subprocess
 import sys
 import sysconfig
+from packaging.version import Version
 from pathlib import Path
 from typing import List, Optional
 
@@ -19,6 +22,21 @@
 from . import env as jit_env
 
 
+@functools.cache
+def _get_cuda_version() -> Version:
+    if CUDA_HOME is None:
+        nvcc = "nvcc"
+    else:
+        nvcc = os.path.join(CUDA_HOME, "bin/nvcc")
+    txt = subprocess.check_output([nvcc, "--version"], text=True)
+    matches = re.findall(r"release (\d+\.\d+),", txt)
+    if not matches:
+        raise RuntimeError(
+            f"Could not parse CUDA version from nvcc --version output: {txt}"
+        )
+    return Version(matches[0])
+
+
 def _get_glibcxx_abi_build_flags() -> List[str]:
     glibcxx_abi_cflags = [
         "-D_GLIBCXX_USE_CXX11_ABI=" + str(int(torch._C._GLIBCXX_USE_CXX11_ABI))
@@ -78,8 +96,13 @@ def generate_ninja_build_for_op(
         "$common_cflags",
         "--compiler-options=-fPIC",
         "--expt-relaxed-constexpr",
-        "-static-global-template-stub=false",
     ]
+    cuda_version = _get_cuda_version()
+    # enable -static-global-template-stub when cuda version >= 12.8
+    if cuda_version >= Version("12.8"):
+        cuda_cflags += [
+            "-static-global-template-stub=false",
+        ]
     cuda_cflags += _get_cuda_arch_flags(extra_cuda_cflags)
     if extra_cuda_cflags is not None:
         cuda_cflags += extra_cuda_cflags