From 4225518cf3f942be3480c37c0dba233d593cf7fb Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Tue, 7 Jul 2026 22:28:32 +0800
Subject: [PATCH] perf: add fast-math and vectorization nvcc/cxx build flags

- centralize CXX_FLAGS/NVCC_FLAGS in csrc/build.py as single source
- add --use_fast_math, --ptxas-options=-O3,-v, --extra-device-vectorization
- add -march=native -funroll-loops host flags
- setup.py reads shared cxx_flags/nvcc_flags from registry
- sync pure-C test build commands with new flags
---
 csrc/build.py                  | 12 +++++++++++-
 csrc/tests/gqa_decode_test.cu  |  8 +++++++-
 csrc/tests/gqa_prefill_test.cu |  8 +++++++-
 setup.py                       |  5 ++++-
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/csrc/build.py b/csrc/build.py
index c33f955..a445ba2 100644
--- a/csrc/build.py
+++ b/csrc/build.py
@@ -20,13 +20,23 @@ def _arch_flags() -> list[str]:
 _kernels_dir = Path("csrc/kernels")
 REGISTRY: dict[str, dict] = {}
 
+CXX_FLAGS = ["-O3", "-march=native", "-funroll-loops"]
+NVCC_FLAGS = [
+    "-O3",
+    "--expt-relaxed-constexpr",
+    "--use_fast_math",
+    "--ptxas-options=-O3,-v",
+    "--extra-device-vectorization",
+]
+
 
 def register(name: str, sources: list[str] | None = None, **kwargs):
     if sources is None:
         sources = [str(_kernels_dir / f"{name}.cu")]
     REGISTRY[name] = {
         "sources": sources,
-        "nvcc_flags": ["-O3", "--expt-relaxed-constexpr", *_arch_flags()],
+        "cxx_flags": [*CXX_FLAGS],
+        "nvcc_flags": [*NVCC_FLAGS, *_arch_flags()],
         "extra_link_args": kwargs.pop("extra_link_args", []),
         **kwargs,
     }
diff --git a/csrc/tests/gqa_decode_test.cu b/csrc/tests/gqa_decode_test.cu
index e4756ff..a81caf0 100644
--- a/csrc/tests/gqa_decode_test.cu
+++ b/csrc/tests/gqa_decode_test.cu
@@ -1,4 +1,10 @@
-// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_decode_test.cu -o test && ./test
+/*
+Pure-C test:
+nvcc -I csrc -arch=sm_89 -O3 \
+    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
+    csrc/tests/gqa_decode_test.cu -o test && ./test
+*/
+
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>
diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu
index 7f3e8ee..19ba520 100644
--- a/csrc/tests/gqa_prefill_test.cu
+++ b/csrc/tests/gqa_prefill_test.cu
@@ -1,4 +1,10 @@
-// Pure-C test: nvcc -I csrc -arch=sm_89 csrc/tests/gqa_prefill_test.cu -o test && ./test
+/*
+Pure-C test:
+nvcc -I csrc -arch=sm_89 -O3 \
+    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
+    csrc/tests/gqa_prefill_test.cu -o test && ./test
+*/
+
 #include <cstdio>
 #include <cstdlib>
 #include <cmath>
diff --git a/setup.py b/setup.py
index 21c1ec5..2f1af11 100644
--- a/setup.py
+++ b/setup.py
@@ -41,7 +41,10 @@ if _should_build():
             CUDAExtension(
                 f"astrai.extension.{name}",
                 info["sources"],
-                extra_compile_args={"cxx": ["-O3"], "nvcc": info["nvcc_flags"]},
+                extra_compile_args={
+                    "cxx": info["cxx_flags"],
+                    "nvcc": info["nvcc_flags"],
+                },
                 extra_link_args=[f"-Wl,-rpath,{_torch_lib}"],
             )
         )