[ROCDL] Fix wave size for gfx1250 (wave32, not wave64)

CLAUDE.md

@@ -171,17 +171,7 @@ helper code that is not part of the traced closure.
 | `gfx950` / `gfx95*` | MI350 / MI355X | 64 | MFMA | CDNA4 path; FP4, MFMA scale, wider LDS copy paths, 160KB LDS |
 | `gfx11*` | RDNA3 / RDNA3.5 (Strix Halo, e.g. gfx1151) | 32 | WMMA | No MFMA; f16/bf16 (and i8/i4) WMMA GEMM; legacy v16-operand WMMA ABI; **no native FP8** (kernels fail-fast); `kernels/rdna3_f16_gemm.py`. `is_rdna_arch()` returns True. |
 | `gfx120*` | RDNA4 (gfx1201 = Radeon AI PRO R9700) | 32 | WMMA | RDNA path, wave32; new v8-operand WMMA ABI with the gfx11 16x16x16 shapes (`gfx120x.wmma` atom, `lib/Dialect/FlyROCDL/GFX120X/`); native FP8. `is_rdna_arch()` returns True. |
-| `gfx1250` | — | 32 | WMMA / TDM | FP8/FP4 GEMM, MoE, async/TDM copy helpers, 320KB LDS. NOTE: `is_rdna_arch('gfx1250')` returns **False** and `get_warp_size` returns 64 — the gfx1250 kernels hardcode `WAVE_SIZE = 32` themselves. |
-
-Use `from flydsl.runtime.device import get_rocm_arch, is_rdna_arch` rather than
-hard-coding behavior when possible. `is_rdna_arch()`
-(`python/flydsl/runtime/device.py`) is the single source of truth for CDNA vs
-RDNA and is wave32-true only for `gfx10*`/`gfx11*`/`gfx120*` prefixes; it does
-**not** match `gfx1250`. Shared wave-size logic is `get_warp_size(arch)` in
-`kernels/kernels_common.py` (`32 if is_rdna_arch(arch) else 64`), so it returns
-64 for gfx1250 — gfx1250 kernels set wave32 explicitly.
-`tests/kernels/test_rdna_gemm.py` shows the gfx11* (v16 ABI) vs gfx120* (v8 ABI)
-kernel-selection pattern.
+| `gfx1250` | — | 32 | WMMA / TDM | FP8/FP4 GEMM, MoE, async/TDM copy helpers, 320KB LDS. |
 
 ## Kernel Entry Points
 

docs/prebuilt_kernels_guide.md

@@ -258,7 +258,7 @@ Shared kernel utilities used across GEMM/MoE/norm kernels.
 
 | Function | Description |
 |---|---|
-| `get_warp_size(arch=None)` | Wave size for the arch: `32` on RDNA, else `64` |
+| `get_warp_size(arch=None)` | Wave size for the arch: `32` on gfx10/11/12, else `64` |
 | `dtype_to_elem_type(dtype_str)` | Map a dtype string to the Fly element type |
 | `validate_moe_dtypes(a_dtype, b_dtype)` | Validate an allowed MoE A/B dtype pairing |
 | `get_llvm_ptr(ptr, offset, dtype_bytes, ...)` | Compute a byte-offset LLVM pointer |

kernels/common/kernels_common.py

@@ -14,7 +14,8 @@
 from flydsl._mlir.dialects import builtin
 from flydsl._mlir.dialects import gpu as _gpu
 from flydsl._mlir.dialects import scf as _scf
-from flydsl.runtime.device import get_rocm_arch, is_rdna_arch
+from flydsl.runtime.device import get_rocm_arch
+from flydsl.runtime.device import get_warp_size as get_warp_size
 
 # Memory/atomic primitives now live in mem_ops; re-exported here for back-compat.
 from kernels.common.mem_ops import _create_llvm_ptr
@@ -94,16 +95,6 @@ def cvt_sr_f32_to_bf16(x, rand):
     return fx.Uint16(bits >> fx.Uint32(16)).bitcast(fx.BFloat16)
 
 
-def get_warp_size(arch=None):
-    """Return the wavefront/warp size for the given GPU architecture.
-
-    CDNA (gfx9xx) uses wave64, RDNA (gfx10xx/gfx11xx/gfx12xx) uses wave32.
-    """
-    if arch is None:
-        arch = get_rocm_arch()
-    return 32 if is_rdna_arch(arch) else 64
-
-
 def default_f8_type() -> ir.Type:
     """Select E4M3 f8 type compatible with the current GPU arch.
 

python/flydsl/compiler/backends/rocm.py

@@ -3,7 +3,7 @@
 
 from typing import List, Tuple
 
-from ...runtime.device import get_rocm_arch, is_rdna_arch
+from ...runtime.device import get_rocm_arch, get_warp_size
 from ...utils import env
 from .base import BaseBackend, GPUTarget
 
@@ -18,13 +18,11 @@ def supports_target(target: GPUTarget) -> bool:
     @staticmethod
     def detect_target() -> GPUTarget:
         arch = env.compile.arch or get_rocm_arch()
-        warp_size = 32 if is_rdna_arch(arch) else 64
-        return GPUTarget(backend="rocm", arch=arch, warp_size=warp_size)
+        return GPUTarget(backend="rocm", arch=arch, warp_size=get_warp_size(arch))
 
     @classmethod
     def make_target(cls, arch: str) -> GPUTarget:
-        warp_size = 32 if is_rdna_arch(arch) else 64
-        return GPUTarget(backend="rocm", arch=arch, warp_size=warp_size)
+        return GPUTarget(backend="rocm", arch=arch, warp_size=get_warp_size(arch))
 
     @classmethod
     def llvm_address_space(cls, address_space) -> int:
@@ -74,7 +72,7 @@ def _pipeline_parts(self, *, compile_hints: dict) -> Tuple[List[str], str]:
             "module": "",
             "triple": "amdgcn-amd-amdhsa",
             "unsafe-math": "true" if compile_hints.get("unsafe_fp_math") else "false",
-            "wave64": "false" if is_rdna_arch(chip) else "true",
+            "wave64": "true" if get_warp_size(chip) == 64 else "false",
         }
 
         pre_binary_fragments = [

python/flydsl/runtime/device.py

@@ -97,3 +97,18 @@ def is_rdna_arch(arch: Optional[str] = None) -> bool:
     if arch.startswith("gfx120"):
         return True
     return False
+
+
+def get_warp_size(arch: Optional[str] = None) -> int:
+    """Lanes per warp for an architecture.
+
+    If arch is None, the current GPU arch is auto-detected.
+    """
+    if arch is None:
+        arch = get_rocm_arch()
+    if not arch:
+        return 64
+    arch = arch.lower()
+    if arch.startswith("gfx10") or arch.startswith("gfx11") or arch.startswith("gfx12"):
+        return 32
+    return 64