Add native Windows RDNA HIP and CK support

[Yasei-no-otoko]

Summary

• add an opt-in native HIP/Composable Kernel build path on Windows while preserving the default Triton-only install
• port JIT extension building to Windows .pyd, MSVC, and hipcc linking conventions
• detect Windows ROCm GPU architecture and CU count through PyTorch
• pass RDNA targets into CK FMHA generation for fixed, varlen, split-KV, and batch-prefill paths
• add offline build defaults for gfx1100, gfx1101, gfx1102, gfx1103, gfx1151, and gfx1201
• exclude Linux/CDNA-only HSA code objects and unavailable FlyDSL dependencies from Windows wheels
• document the experimental Windows native flow and add target-selection regression tests

Root cause

Windows was unconditionally forced into Triton-only mode, and the extension builder assumed POSIX shared libraries, toolchain flags, filesystem commands, and GPU discovery. After enabling the native toolchain, CK FMHA still generated only its default gfx9/gfx950 instances unless the active RDNA target was passed explicitly, which caused RDNA dispatch to fail with invalid argument for fmha_fwd.

Impact

Windows ROCm users can opt into native HIP and CK operators with AITER_ENABLE_HIP=1. The existing Windows Triton-only behavior remains the default, but it is a compatibility baseline rather than the Windows performance path: on gfx1151 it was slower than PyTorch SDPA at every tested sequence length. CDNA/Linux behavior and the CK generator defaults for gfx942/gfx950 are preserved.

Why native CK matters on Windows

The current minimal Windows support is Triton-only. On gfx1151, Aiter's default Triton FMHA was slower than PyTorch SDPA at every tested sequence length: 4.51x slower at S=128, 2.64x slower at S=512, and 1.06x slower at S=1024. Triton-only should therefore be treated as a compatibility baseline, not the Windows performance path. Native CK is needed to avoid this regression and achieved 1.14x-1.28x speedups over PyTorch SDPA in the same test.

Lower is better.

Sequence length	Aiter Triton-only	PyTorch SDPA	Triton-only relative to SDPA
128	189.39 us	41.97 us	4.51x slower
512	192.56 us	72.85 us	2.64x slower
1024	263.80 us	249.15 us	1.06x slower

All Triton-only outputs matched PyTorch SDPA (allclose=True; maximum absolute differences were 0.000977, 0.000488, and 0.000977 respectively).

Validation

• built and installed a CPython 3.13 win_amd64 wheel with MAX_JOBS=32
• ran the native wheel on an AMD Radeon 8060S (gfx1151, ROCm 7.14):
  • BF16 activation matched the PyTorch reference
  • CK BF16 FMHA matched PyTorch SDPA (max_abs=0.00542)
• attention benchmarks used B=1/H=16/D=64, BF16, non-causal attention, 50 warmup iterations, 500 measured iterations across 7 alternating-order rounds, GPU synchronization, and the median wall-clock latency
• CK FMHA latency versus PyTorch SDPA on gfx1151:
  • S=128: 43.78 us vs 50.34 us (1.15x)
  • S=512: 65.57 us vs 74.94 us (1.14x)
  • S=1024: 195.69 us vs 250.03 us (1.28x)
• python -B op_tests/test_windows_rdna_ck_targets.py: 3 tests passed
• CK codegen generated non-empty target-specific sets for gfx1100/1101/1102/1103/1151/1201
• Windows hipcc cross-compiled a representative CK BF16 FMHA translation unit for gfx1100, gfx1101, gfx1102, gfx1103, and gfx1201
• git diff --check passed

Physical-GPU runtime validation was performed on gfx1151. The gfx110x and gfx1201 coverage is codegen plus Windows hipcc cross-compilation because those GPUs were not available on the test host.

Representative offline CU defaults follow the ROCm GPU specification table; binned devices can override them with CU_NUM.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4340 --add-label <label>

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b0b15a39b9

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[Copilot]

Pull request overview

This PR adds an opt-in native HIP + Composable Kernel (CK) build/runtime path for Windows RDNA GPUs while keeping the default Windows experience as Triton-only, and updates JIT/build utilities to follow Windows toolchain and binary conventions.

Changes:

• Introduce AITER_ENABLE_HIP=1 as the Windows opt-in gate for building/using native HIP/CK, and adjust packaging to avoid Linux/CDNA-only artifacts on Windows.
• Port/adjust JIT extension build + GPU detection to work on Windows (MSVC env setup, .pyd naming, temp paths, rocminfo replacements via PyTorch where needed).
• Ensure CK FMHA codegen receives explicit RDNA targets and add regression tests for target/default selection.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
setup.py	Adds Windows opt-in gating for HIP/CK, skips FlyDSL + HSA bundling on Windows, and includes .pyd artifacts in package data.
README.md	Documents the experimental native Windows ROCm flow and the AITER_ENABLE_HIP opt-in.
pyproject.toml	Makes flydsl a non-Windows-only build requirement via environment marker.
op_tests/test_windows_rdna_ck_targets.py	Adds regression tests validating RDNA offline defaults and CK --targets propagation behavior.
csrc/kernels/activation_kernels.cu	Adds Windows-specific math header handling for hipcc/UCRT.
aiter/utility/dtypes.py	Uses PyTorch-based arch detection on Windows for FP8 dtype selection and defers some imports.
aiter/ops/mha.py	Passes runtime-derived CK --targets into FMHA codegen commands.
aiter/jit/utils/torch_guard.py	Adjusts torch.library registration naming to match fragment-library conventions.
aiter/jit/utils/mha_recipes.py	Adds helper to emit CK --targets for non-gfx9 arch at runtime.
aiter/jit/utils/file_baton.py	Handles Windows PID liveness checks where WinError behavior differs from POSIX.
aiter/jit/utils/cpp_extension.py	Adds Windows-aware extension/link conventions and ninja rule differences for MSVC/hipcc builds.
aiter/jit/utils/chip_info.py	Adds Windows GPU arch/CU detection via PyTorch and supports direct-import patterns.
aiter/jit/utils/build_targets.py	Adds RDNA arch → default CU mappings for offline/cross builds.
aiter/jit/core.py	Uses cross-platform temp dirs, updates module filename handling to use platform extension, and removes POSIX-only shell/file ops.
aiter/init.py	Makes Windows native enablement opt-in (or auto-on if bundled .pyd exists) and adjusts imports for Windows distributed availability.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.