diff --git a/.github/workflows/test-code-build.yml b/.github/workflows/test-code-build.yml
index f87845d..6e0765d 100644
--- a/.github/workflows/test-code-build.yml
+++ b/.github/workflows/test-code-build.yml
@@ -7,23 +7,47 @@ jobs:
     name: Build wheels and Upload
     # runs-on: windows-latest
     runs-on: codebuild-flash-attention-pre-build-wheel-windows-${{ github.run_id }}-${{ github.run_attempt }}
+    env:
+      MAX_JOBS: 4
+      NVCC_THREADS: 4
     steps:
       - uses: actions/checkout@v4
-
       - name: Enable Git long paths
+        shell: pwsh
         run: git config --system core.longpaths true
-        shell: pwsh
-
-      - name: Install VS2022 BuildTools 17.9.7
-        run: choco install -y visualstudio2022buildtools --version=117.9.7.0 --params "--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64 --installChannelUri https://aka.ms/vs/17/release/180911598_-255012421/channel"
-        shell: pwsh
 
       - uses: actions/setup-python@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: "3.11"
 
-      - uses: Jimver/cuda-toolkit@master
+      - uses: Jimver/cuda-toolkit@v0.2.24
         with:
-          cuda: ${{ matrix.cuda-version }}
+          cuda: "12.4.1"
           method: "network"
+          use-github-cache: false
+          use-local-cache: false
 
+      - name: Install VS2022 BuildTools
+        shell: pwsh
+        run: |
+          choco install -y visualstudio2022buildtools `
+            --version=117.14.1 `
+            --params "--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64"
+
+      - name: Install build dependencies
+        shell: pwsh
+        run: |
+          pip install -U pip setuptools==75.8.0 wheel setuptools packaging psutil ninja
+
+      - name: Build wheels
+        shell: pwsh
+        run: |
+          .\build_windows.ps1 -FlashAttnVersion "2.5.9" -PythonVersion "3.11" -TorchVersion "2.6.0" -CudaVersion "12.4.1"
+          $wheelName = Get-ChildItem -Path "flash-attention\dist\*.whl" | Select-Object -First 1 | ForEach-Object { $_.Name }
+          echo "wheel_name=$wheelName" >> $env:GITHUB_ENV
+
+      - name: Install Test
+        shell: pwsh
+        run: |
+          pip install flash-attention/dist/$env:wheel_name
+          python -c "import flash_attn; print(flash_attn.__version__)"