bugfix: remove the append "a" logic if user specifies cuda arch explicitly (#1798)

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

## 📌 Description

In our design of compilation context we will append "a" to cuda archs if
the major arch is >= 9, however, we shouldn't change it if user specify
them explicitly.

This PR changes the logic so that we only append "a" to cuda arch if
`FLASHINFER_CUDA_ARCH_LIST` is not set explicitly, otherwise follow the
user specified cuda arch.

## 🔍 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.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

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

Author: yzh119

flashinfer/compilation_context.py

@@ -36,9 +36,6 @@ def __init__(self):
             for arch in os.environ["FLASHINFER_CUDA_ARCH_LIST"].split(" "):
                 major, minor = arch.split(".")
                 major = int(major)
-                if major >= 9:
-                    if minor.isdigit():
-                        minor = str(minor) + "a"
                 self.TARGET_CUDA_ARCHS.add((int(major), str(minor)))
         else:
             try:

scripts/task_test_aot_build_import.sh

@@ -20,13 +20,13 @@ export MAX_JOBS
 export FLASHINFER_CUDA_ARCH_LIST=$(python3 -c '
 import torch
 cuda_ver = torch.version.cuda
-arches = ["7.5", "8.0", "8.9", "9.0"]
+arches = ["7.5", "8.0", "8.9", "9.0a"]
 if cuda_ver is not None:
     try:
         major, minor = map(int, cuda_ver.split(".")[:2])
         if (major, minor) >= (12, 8):
-            arches.append("10.0")
-            arches.append("12.0")
+            arches.append("10.0a")
+            arches.append("12.0a")
     except Exception:
         pass
 print(" ".join(arches))