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

@@ -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,
     }

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>

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>

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}"],
             )
         )