win: fix cuda build

[dhiltgen]

Adjust recent CUDA changes to build on windows.

Proposed changes

HEAD on main currently fails to build on Windows with CUDA enabled. This gets the build working.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Thanks!

[dhiltgen]

While testing more deeply I noticed sporadic JIT crashes on Windows CUDA and traced it down to missing device initialization, which I've fixed in a discrete commit. Let me know if you want that split out to a different PR.

[zcbenz]

The new change looks good to me, thanks for fixing it! Still need to wait for the CI environment to get fixed before I can merge.

[dhiltgen]

I fixed another performance issue. I can move it to another PR as well if needed. It looks like cuDNN should be disabled by default on Windows due to how it interacts with WDDM. Using mlx-community/Qwen3-0.6B-4bit to test with an RTX 6000 Ada on Win 11:

  ┌───────────┬─────────────────┬───────────────────┐
  │   Mode    │ Prefill (p2048) │ Generation (g128) │
  ├───────────┼─────────────────┼───────────────────┤
  │ cuDNN ON  │ 2,022 tok/s     │ 10.4 tok/s        │
  ├───────────┼─────────────────┼───────────────────┤
  │ cuDNN OFF │ 4,822 tok/s     │ 480 tok/s         │
  └───────────┴─────────────────┴───────────────────┘

[zcbenz]

I'm good disabling cuDNN SDPA for Windows, but can you check with NVIDIA whether this is something that can be fixed? Our fallback kernel is not good for prefill and long context decoding.