build(SM70): precompile GDN CUDA kernel to .so during docker build

precompile_gdn.py: calls torch.utils.cpp_extension.load with build_directory
to produce .so at build time. If build env has no GPU/compiler, fails
gracefully — kernel JIT compiles at runtime instead.

fused_fwd.py: _load_ext() now checks build/ dir for precompiled .so first,
skips 2-minute JIT compilation if found.

qwen3_6_scripts/flash_qla_sm70/fused_fwd.py

@@ -20,6 +20,24 @@ def _load_ext():
         raise RuntimeError("SM70 FlashQLA backend requires CUDA.")
 
     os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "7.0;7.5")
+
+    # Try precompiled .so first (built during docker build)
+    build_dir = Path(__file__).with_name("build")
+    if build_dir.is_dir():
+        so_files = list(build_dir.glob("*.so"))
+        if so_files:
+            try:
+                _EXT = load(
+                    name="flash_qla_sm70_gdn_strided",
+                    sources=[],  # empty — just load from build_directory
+                    build_directory=str(build_dir),
+                    verbose=False,
+                )
+                return _EXT
+            except Exception:
+                pass  # fall through to JIT
+
+    # JIT compile (slow, ~2min first time)
     src = Path(__file__).with_name("csrc") / "gdn_forward.cu"
     _EXT = load(
         name="flash_qla_sm70_gdn_strided",

qwen3_6_scripts/patch_ops.sh

@@ -188,9 +188,12 @@ if [ -d "./flash_qla_sm70" ]; then
     rm -rf "$FLASH_QLA_DST" 2>/dev/null
     cp -r ./flash_qla_sm70 "$FLASH_QLA_DST" 2>/dev/null && \
         echo "[patch_ops] flash_qla_sm70 deployed to $FLASH_QLA_DST" || true
+    # Pre-compile CUDA kernel → .so (skipped if no GPU/compiler at build time)
+    python3 ./precompile_gdn.py "$FLASH_QLA_DST" 2>&1 || \
+        echo "[patch_ops] WARNING: precompile failed — kernel will JIT at runtime"
     # Also deploy to VLLM2 if present
     if [ -n "$VLLM2" ]; then
         rm -rf "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null
-        cp -r ./flash_qla_sm70 "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null || true
+        cp -r "$FLASH_QLA_DST" "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null || true
     fi
 fi

qwen3_6_scripts/precompile_gdn.py

@@ -0,0 +1,53 @@
+"""
+Pre-compile SM70 GDN CUDA kernel → .so at Docker build time.
+Avoids 2-minute JIT delay at runtime.
+
+Usage: python3 precompile_gdn.py /path/to/flash_qla_sm70/
+"""
+import os
+import sys
+
+def main():
+    if len(sys.argv) < 2:
+        print("[precompile] Usage: python3 precompile_gdn.py <flash_qla_sm70_dir>")
+        sys.exit(1)
+
+    flash_dir = sys.argv[1]
+    cu_src = os.path.join(flash_dir, "csrc", "gdn_forward.cu")
+    if not os.path.exists(cu_src):
+        print(f"[precompile] ERROR: {cu_src} not found")
+        sys.exit(1)
+
+    # Set arch for BI-V100 (SM70 compatible)
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "7.0;7.5"
+
+    build_dir = os.path.join(flash_dir, "build")
+    os.makedirs(build_dir, exist_ok=True)
+
+    print(f"[precompile] Compiling {cu_src} → .so in {build_dir}")
+    print(f"[precompile] TORCH_CUDA_ARCH_LIST = {os.environ['TORCH_CUDA_ARCH_LIST']}")
+
+    try:
+        from torch.utils.cpp_extension import load
+        ext = load(
+            name="flash_qla_sm70_gdn_strided",
+            sources=[cu_src],
+            extra_cuda_cflags=["-O3"],
+            extra_cflags=["-O3"],
+            build_directory=build_dir,
+            verbose=True,
+        )
+        print(f"[precompile] SUCCESS — compiled .so in {build_dir}")
+        # List the built files
+        for f in os.listdir(build_dir):
+            if f.endswith(".so"):
+                full = os.path.join(build_dir, f)
+                print(f"[precompile]   {f} ({os.path.getsize(full)} bytes)")
+    except Exception as e:
+        print(f"[precompile] FAILED: {e}")
+        print("[precompile] Kernel will JIT compile at runtime instead (~2min)")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()