hip: enable native gfx1170 attention and GEMM schedules

Add gfx1170-specific kernel dispatch paths for bf16 attention (STAGE_V
tiling with 64- and 512-row blocks), int8 GEMM non-duplicated layout,
and ConvRot packed-none detection.

The __init__.py changes use .get() for the wmma_gfx117 architecture
group so this commit is forward-compatible: the gfx1170 runtime paths
activate once the architecture manifest adds the group (PR #141),
and are no-ops until then.

Benchmark results on gfx1170 (with PR #141 applied):
  FP16 + aotriton baseline:              65.55s  (1.00x)
  INT8 convrot + bf16 kitchen attention:  32.84s  (2.00x)

Author: tvukovic-amd

comfy_kitchen/backends/hip/__init__.py

@@ -1,6 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""HIP backend for AMD RDNA2, RDNA3/3.5 and RDNA4.
+"""HIP backend for AMD RDNA2, RDNA3/3.5, RDNA4 and gfx117x.
 
 Every matmul is a WMMA kernel compiled from the sources in this directory; the
 backend does not link or call hipBLAS/hipBLASLt.
@@ -183,23 +183,26 @@ def _visible_gfx_arches() -> tuple[str | None, ...]:
     return tuple(_gfx_arch(i) for i in range(torch.cuda.device_count()))
 
 
-# RDNA2 has no matrix cores; RDNA3/3.5 and RDNA4 do. This exact manifest is also
-# consumed by setup.py and CMake. Never infer support from a gfx prefix: a new
+# RDNA2 has no matrix cores; RDNA3/3.5, gfx117x and RDNA4 do. This exact manifest
+# is also consumed by setup.py and CMake. Never infer support from a gfx prefix: a new
 # compiler-recognized target needs its WMMA policy reviewed before it is safe.
 _ARCH_MANIFEST_PATH = os.path.join(os.path.dirname(__file__), "architectures.json")
 _ARCH_GROUPS = json.loads(
     pathlib.Path(_ARCH_MANIFEST_PATH).read_text(encoding="utf-8")
 )
 _ARCH_ELEMENTWISE_ONLY = frozenset(_ARCH_GROUPS["elementwise_only"])
 _ARCH_WMMA_GFX11 = frozenset(_ARCH_GROUPS["wmma_gfx11"])
+_ARCH_WMMA_GFX117 = frozenset(_ARCH_GROUPS.get("wmma_gfx117", []))
 _ARCH_WMMA_GFX12 = frozenset(_ARCH_GROUPS["wmma_gfx12"])
-_ARCH_WMMA = _ARCH_WMMA_GFX11 | _ARCH_WMMA_GFX12
+_ARCH_WMMA = _ARCH_WMMA_GFX11 | _ARCH_WMMA_GFX117 | _ARCH_WMMA_GFX12
 _ARCH_SUPPORTED = _ARCH_ELEMENTWISE_ONLY | _ARCH_WMMA
 
+_GFX117_WMMA_POLICIES_READY = True
+
 
 def _has_nonduplicated_wmma(device: torch.device | int | None = None) -> bool:
     """Whether WMMA operands use the gfx12 128-bit layout."""
-    return _gfx_arch(device) in _ARCH_WMMA_GFX12
+    return _gfx_arch(device) in (_ARCH_WMMA_GFX12 | _ARCH_WMMA_GFX117)
 
 # The GEMMs, and only the GEMMs, need matrix cores. Everything else is elementwise
 # or a scalar reduction and runs on any supported architecture. This set names the
@@ -247,7 +250,13 @@ def _has_wmma(arches: Sequence[str | None]) -> bool:
     the capability set has to be the intersection over the visible devices: one
     RDNA2 card in an otherwise RDNA4 box means no GEMM is safe to advertise.
     """
-    return bool(arches) and all(a in _ARCH_WMMA for a in arches)
+    if not arches or any(a is None for a in arches):
+        return False
+    if not all(a in _ARCH_WMMA for a in arches):
+        return False
+    if any(a in _ARCH_WMMA_GFX117 for a in arches) and not _GFX117_WMMA_POLICIES_READY:
+        return False
+    return True
 
 
 def is_available() -> bool:

comfy_kitchen/backends/hip/hadamard.h

@@ -1350,7 +1350,8 @@ inline bool use_gfx12_convrot_packed_none() {
     auto select = [device] {
         hipDeviceProp_t properties{};
         return hipGetDeviceProperties(&properties, device) == hipSuccess &&
-            std::strncmp(properties.gcnArchName, "gfx12", 5) == 0 &&
+            (std::strcmp(properties.gcnArchName, "gfx1170") == 0 ||
+             std::strncmp(properties.gcnArchName, "gfx12", 5) == 0) &&
             properties.warpSize == 32 && properties.maxThreadsPerBlock >= 512;
     };
     if (device < 0 || device >= kMaxDevices) return select();

comfy_kitchen/backends/hip/mma.h

@@ -18,7 +18,7 @@
 //          free. Fragment width is the K-step: 16B iu8, 8B iu4, 32B bf16 for fp8.
 //
 // Accumulators are v8f/v8i on both, column lane % 16, but the row differs:
-// gfx12 gives each half-wave a contiguous block, D[e + 8 * (lane / 16)]; gfx11
+// gfx12/gfx117 give each half-wave a contiguous block, D[e + 8 * (lane / 16)]; gfx11
 // interleaves the halves, D[2 * e + (lane / 16)]. See acc_row. rocWMMA's "padded
 // acc" gfx11 quirk applies to the 16-bit accumulators, not these.
 //
@@ -144,7 +144,7 @@ __forceinline__ __device__ typename Mma::Frag load_frag_16bit(const typename Mma
 // Policies
 // ---------------------------------------------------------------------------
 
-#if defined(COMFY_MMA_GFX12)
+#if defined(COMFY_MMA_GFX12) || defined(COMFY_MMA_GFX117)
 
 struct MmaFp8 {
     using Acc = v8f;

comfy_kitchen/backends/hip/ops/attention_bf16.hip

@@ -32,14 +32,26 @@ inline bool use_gfx12_attention_schedule() {
         hipDeviceProp_t properties{};
         if (hipGetDevice(&device) == hipSuccess &&
             hipGetDeviceProperties(&properties, device) == hipSuccess &&
-            std::strncmp(properties.gcnArchName, "gfx12", 5) == 0) {
+            (std::strcmp(properties.gcnArchName, "gfx1170") == 0 ||
+             std::strncmp(properties.gcnArchName, "gfx12", 5) == 0)) {
             return true;
         }
         return false;
     }();
     return selected;
 }
 
+inline bool use_gfx1170_attention_schedule() {
+    static const bool selected = [] {
+        int device = 0;
+        hipDeviceProp_t properties{};
+        return hipGetDevice(&device) == hipSuccess &&
+            hipGetDeviceProperties(&properties, device) == hipSuccess &&
+            std::strcmp(properties.gcnArchName, "gfx1170") == 0;
+    }();
+    return selected;
+}
+
 template <int HD, int BLOCK_M, int BLOCK_N>
 struct Bf16AttentionShape {
     static_assert(HD % 16 == 0);
@@ -50,7 +62,7 @@ struct Bf16AttentionShape {
     static constexpr int kDimTiles = HD / 16;
     static constexpr int kKeyTiles = BLOCK_N / 16;
     static constexpr int kKStride = HD + 4;
-    static constexpr int kVStride = BLOCK_N + 4;
+    static constexpr int kVStride = BLOCK_N + (BLOCK_M <= 64 ? 2 : 4);
 };
 
 template <typename T>
@@ -90,7 +102,7 @@ __forceinline__ __device__ void store_bf16_acc(
 }
 
 template <int HD, int BLOCK_M, int BLOCK_N, bool DIRECT_GLOBAL,
-          bool NO_KV_TAIL = false>
+          bool NO_KV_TAIL = false, bool STAGE_V = false>
 __global__ __launch_bounds__(BLOCK_M * 2)
 void bf16_attention_kernel(
     const __bf16* __restrict__ q, const __bf16* __restrict__ k,
@@ -107,7 +119,9 @@ void bf16_attention_kernel(
     constexpr int KT = Shape::kKeyTiles;
 
     __shared__ __align__(16) __bf16 staged_k[
-        DIRECT_GLOBAL ? 1 : 2 * BLOCK_N * Shape::kKStride];
+        DIRECT_GLOBAL ? 1 : (STAGE_V ? 1 : 2) * BLOCK_N * Shape::kKStride];
+    __shared__ __align__(16) __bf16 staged_v[
+        STAGE_V ? HD * Shape::kVStride : 1];
 
     const int tid = threadIdx.x;
     const int lane = tid & 31;
@@ -158,7 +172,7 @@ void bf16_attention_kernel(
     constexpr int kKStageItems = DIRECT_GLOBAL
         ? 1
         : (kPacksPerTile + Shape::kThreads - 1) / Shape::kThreads;
-    if constexpr (!DIRECT_GLOBAL) {
+    if constexpr (!DIRECT_GLOBAL && !STAGE_V) {
 #pragma unroll
         for (int item = 0; item < kKStageItems; ++item) {
             const int item_pack = tid + item * Shape::kThreads;
@@ -185,8 +199,46 @@ void bf16_attention_kernel(
         const bool has_next_tile = tile + 1 < key_tiles;
         const int next_buffer = (tile + 1) & 1;
 
+        if constexpr (STAGE_V) {
+            if (tile != 0) __syncthreads();
+#pragma unroll
+            for (int item = 0; item < kKStageItems; ++item) {
+                const int item_pack = tid + item * Shape::kThreads;
+                if (item_pack >= kPacksPerTile) continue;
+                const int item_key = item_pack / kPacksPerRow;
+                const int item_d =
+                    (item_pack - item_key * kPacksPerRow) * kPackElems;
+                const int key_index = key0 + item_key;
+                v8bf k_value{};
+                if constexpr (NO_KV_TAIL) {
+                    k_value = *reinterpret_cast<const v8bf*>(
+                        k_head + static_cast<int64_t>(key_index) * k_stride_n + item_d);
+                } else if (key_index < kv_len) {
+                    k_value = *reinterpret_cast<const v8bf*>(
+                        k_head + static_cast<int64_t>(key_index) * k_stride_n + item_d);
+                }
+                store_bf16x8_padded(
+                    staged_k + item_key * Shape::kKStride + item_d, k_value);
+
+                v8bf v_value{};
+                if constexpr (NO_KV_TAIL) {
+                    v_value = *reinterpret_cast<const v8bf*>(
+                        v_head + static_cast<int64_t>(key_index) * v_stride_n + item_d);
+                } else if (key_index < kv_len) {
+                    v_value = *reinterpret_cast<const v8bf*>(
+                        v_head + static_cast<int64_t>(key_index) * v_stride_n + item_d);
+                }
+#pragma unroll
+                for (int e = 0; e < kPackElems; ++e) {
+                    staged_v[(item_d + e) * Shape::kVStride + item_key] =
+                        v_value[e];
+                }
+            }
+            __syncthreads();
+        }
+
         const __bf16* const current_k = staged_k +
-            (tile & 1) * BLOCK_N * Shape::kKStride;
+            (STAGE_V ? 0 : (tile & 1) * BLOCK_N * Shape::kKStride);
 
         MmaBf16::Acc score[KT];
 #pragma unroll
@@ -217,7 +269,7 @@ void bf16_attention_kernel(
         // Prefetch the next K tile after score MMA so global latency hides
         // behind softmax prep instead of sitting in front of the MMA nest.
         v8bf next_k[kKStageItems]{};
-        if constexpr (!DIRECT_GLOBAL) {
+        if constexpr (!DIRECT_GLOBAL && !STAGE_V) {
 #pragma unroll
             for (int item = 0; item < kKStageItems; ++item) {
                 const int item_pack = tid + item * Shape::kThreads;
@@ -275,10 +327,15 @@ void bf16_attention_kernel(
 #pragma unroll
             for (int dt = 0; dt < DT; ++dt) {
                 MmaBf16::Frag v_fragment{};
-                if constexpr (!DIRECT_GLOBAL) {
+                if constexpr (STAGE_V) {
+                    v_fragment = load_frag_16bit<MmaBf16>(
+                        staged_v + (dt * 16 + frag_row(lane)) * Shape::kVStride +
+                        kt * 16,
+                        lane);
+                } else if constexpr (!DIRECT_GLOBAL) {
                     const int tile_key = key0 + kt * 16;
                     if (tile_key + 16 <= kv_len) {
-#if defined(__HIP_DEVICE_COMPILE__) && defined(COMFY_MMA_GFX12)
+#if defined(__HIP_DEVICE_COMPILE__) && defined(COMFY_MMA_GFX12) && defined(__GFX12__)
                         const int source_key = ((lane >> 4) << 3) + (lane & 7);
                         const int source_half = (lane >> 3) & 1;
                         const v8bf* const source = reinterpret_cast<const v8bf*>(
@@ -325,7 +382,7 @@ void bf16_attention_kernel(
         }
         row_sum += tile_sum;
 
-        if constexpr (!DIRECT_GLOBAL) {
+        if constexpr (!DIRECT_GLOBAL && !STAGE_V) {
 #pragma unroll
             for (int item = 0; item < kKStageItems; ++item) {
                 const int item_pack = tid + item * Shape::kThreads;
@@ -356,7 +413,7 @@ void bf16_attention_kernel(
 }
 
 template <int HD, int BLOCK_M, int BLOCK_N, bool DIRECT_GLOBAL,
-          bool NO_KV_TAIL = false>
+          bool NO_KV_TAIL = false, bool STAGE_V = false>
 void launch_bf16_attention_shape(
     const void* q, const void* k, const void* v, void* output,
     int batch, int q_heads, int kv_heads, int q_len, int kv_len,
@@ -366,7 +423,7 @@ void launch_bf16_attention_shape(
     int64_t o_stride_b, int64_t o_stride_h, int64_t o_stride_n,
     float sm_scale, hipStream_t stream) {
     const dim3 grid((q_len + BLOCK_M - 1) / BLOCK_M, q_heads, batch);
-    bf16_attention_kernel<HD, BLOCK_M, BLOCK_N, DIRECT_GLOBAL, NO_KV_TAIL>
+    bf16_attention_kernel<HD, BLOCK_M, BLOCK_N, DIRECT_GLOBAL, NO_KV_TAIL, STAGE_V>
         <<<grid, Bf16AttentionShape<HD, BLOCK_M, BLOCK_N>::kThreads, 0, stream>>>(
             static_cast<const __bf16*>(q), static_cast<const __bf16*>(k),
             static_cast<const __bf16*>(v), static_cast<__bf16*>(output),
@@ -393,14 +450,30 @@ void dispatch_bf16_attention(
             k_stride_n, v_stride_b, v_stride_h, v_stride_n, o_stride_b,
             o_stride_h, o_stride_n, sm_scale, stream);
     } else if (q_len <= 256) {
-        launch_bf16_attention_shape<HD, 32, 32, true>(
-            q, k, v, output, batch, q_heads, kv_heads, q_len, kv_len,
-            q_stride_b, q_stride_h, q_stride_n, k_stride_b, k_stride_h,
-            k_stride_n, v_stride_b, v_stride_h, v_stride_n, o_stride_b,
-            o_stride_h, o_stride_n, sm_scale, stream);
+        if (use_gfx1170_attention_schedule() && kv_len >= 128 &&
+            kv_len % 32 == 0) {
+            launch_bf16_attention_shape<HD, 64, 32, false, true, true>(
+                q, k, v, output, batch, q_heads, kv_heads, q_len, kv_len,
+                q_stride_b, q_stride_h, q_stride_n, k_stride_b, k_stride_h,
+                k_stride_n, v_stride_b, v_stride_h, v_stride_n, o_stride_b,
+                o_stride_h, o_stride_n, sm_scale, stream);
+        } else {
+            launch_bf16_attention_shape<HD, 32, 32, true>(
+                q, k, v, output, batch, q_heads, kv_heads, q_len, kv_len,
+                q_stride_b, q_stride_h, q_stride_n, k_stride_b, k_stride_h,
+                k_stride_n, v_stride_b, v_stride_h, v_stride_n, o_stride_b,
+                o_stride_h, o_stride_n, sm_scale, stream);
+        }
     } else {
-        if (use_gfx12_attention_schedule() && kv_len >= 1024 &&
+        if (use_gfx1170_attention_schedule() && kv_len >= 1024 &&
             kv_len % 32 == 0) {
+            launch_bf16_attention_shape<HD, 512, 32, false, true, true>(
+                q, k, v, output, batch, q_heads, kv_heads, q_len, kv_len,
+                q_stride_b, q_stride_h, q_stride_n, k_stride_b, k_stride_h,
+                k_stride_n, v_stride_b, v_stride_h, v_stride_n, o_stride_b,
+                o_stride_h, o_stride_n, sm_scale, stream);
+        } else if (use_gfx12_attention_schedule() && kv_len >= 1024 &&
+                   kv_len % 32 == 0) {
             launch_bf16_attention_shape<HD, 128, 32, false, true>(
                 q, k, v, output, batch, q_heads, kv_heads, q_len, kv_len,
                 q_stride_b, q_stride_h, q_stride_n, k_stride_b, k_stride_h,

comfy_kitchen/backends/hip/ops/gemm_int8.hip

@@ -25,7 +25,8 @@ inline bool use_nonduplicated_int8_schedule() {
     auto select = [device] {
         hipDeviceProp_t properties{};
         return hipGetDeviceProperties(&properties, device) == hipSuccess &&
-            std::strncmp(properties.gcnArchName, "gfx12", 5) == 0;
+            (std::strcmp(properties.gcnArchName, "gfx1170") == 0 ||
+             std::strncmp(properties.gcnArchName, "gfx12", 5) == 0);
     };
     if (device < 0 || device >= kMaxDevices) return select();