hip: RDNA3 (gfx1100) WMMA kernel for the QSA block-score select/scorer (opt-in STRATA_SELECT_WMMA, gfx12/gfx11)

CMakeLists.txt

@@ -359,6 +359,17 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
     target_link_libraries(qsa_select_bench PRIVATE strata_kernels strata_hip_runtime)
     add_test(NAME qsa_select_bench COMMAND qsa_select_bench 32768 256 3 262144)
   endif()
+  # The AMD tensor-core block scores (ROCWMMA, opt-in) against the FP64 oracle and the warp order's exact
+  # host model (PORTING.md 18).  HIP + rocWMMA only: it tests an arm that exists nowhere else.
+  if(STRATA_ENABLE_HIP AND (STRATA_BUILD_TESTS OR STRATA_PARITY_QSA_SELECT_WMMA))
+    add_executable(qsa_select_wmma_parity src/kernels/qsa_select_wmma_parity.cpp)
+    set_source_files_properties(src/kernels/qsa_select_wmma_parity.cpp PROPERTIES LANGUAGE HIP)
+    target_link_libraries(qsa_select_wmma_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target})
+    if(STRATA_ROCWMMA_INCLUDE_DIR)
+      target_include_directories(qsa_select_wmma_parity PRIVATE ${STRATA_ROCWMMA_INCLUDE_DIR})
+    endif()
+    add_test(NAME qsa_select_wmma_parity COMMAND qsa_select_wmma_parity)
+  endif()
   # plan v0.3 P6: llama.cpp's block layouts and codebook grids for the i-quant formats (MIT, third_party/ggml)
   target_include_directories(strata_kernels PRIVATE third_party/ggml)
   # Match the pinned llama.cpp CUDA arithmetic contract for this adapter only.

include/strata/kernels/qsa_select.hpp

@@ -35,6 +35,15 @@ void qsa_block_scores(const float* pooled, const float* dead, const float* q_idx
 bool qsa_block_scores_tc(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq,
                          int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, int64_t active_blocks);
 
+/// The same scores on AMD tensor cores (ROCWMMA 16x16x16, FP16 operands with FP32 accumulation: FP16-level input
+/// accuracy and another summation order; the tail block n_bid is still the warp kernel's arithmetic).  HIP/gfx1100
+/// only; false (nothing launched) on another geometry, architecture or build, so callers keep their fallbacks.
+/// Keys and queries are converted to FP16 per call into an internal scratch - that conversion is part of the arm's
+/// cost.  Scores may differ from the warp kernel's by at most the bound recorded in PORTING.md 18:
+/// B(qi,b) = 2^-9.5 * S(qi,b) + 8 * 2^-24 * |score|, with S the absolute product sum of the entry.
+bool qsa_block_scores_wmma(const float* pooled, const float* dead, const float* q_idx, const int32_t* steps, int64_t nq,
+                           int64_t max_blocks, const QsaShapes& s, float* scores, void* stream, int64_t active_blocks);
+
 /// ids [nq, cap] (cells, ascending); `cap` >= the largest selection width.
 void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap,
                     const QsaShapes& s, int32_t* ids, void* stream,

src/kernels/cuda/qsa_prompt_attn.cu

@@ -1122,19 +1122,26 @@ bool launch70(const float*, const QsaAttnPools&, const int32_t*, const int32_t*,
 #endif  // !__HIPCC__ && STRATA_EXPERIMENTAL_SM60
 
 #if defined(__HIPCC__)
-// ---- S6: the int8-KV prompt attention on RDNA4 matrix cores (opt-in: STRATA_HIP_WMMA=1, gfx12 only). The design of
+// ---- S6: the int8-KV prompt attention on RDNA4/RDNA3 matrix cores (opt-in: STRATA_HIP_WMMA=1, gfx12/gfx11). The design of
 // the v2 kernel above with gfx12's v_wmma_f32_16x16x16_f16 (wave32) in place of m16n8k16: wave w owns dims
 // [64w, 64w+64) (int8 scale group w) for q.k and p.v; q and p are split into FP16 hi + lo parts, the int8 codes enter
 // exactly as FP16, the scales are applied in FP32 and the four groups' q.k partials are added in a fixed order.
 // FP32-level accuracy, deterministic, but not bitwise equal to qsa_decode_attn_batch (another summation order).
-// Fragment layout (16x16x16, wave32, checked on gfx1201): A lane l holds A[l % 16][(l / 16) * 8 + i], B lane l holds
-// B[(l / 16) * 8 + i][l % 16], C/D lane l holds D[(l / 16) * 8 + i][l % 16], i = 0..7.
+// Fragment layout (16x16x16, wave32, checked on gfx1201; parity drives it on gfx1100 too): A lane l holds
+// A[l % 16][(l / 16) * 8 + i], B lane l holds B[(l / 16) * 8 + i][l % 16], C/D lane l holds
+// D[(l / 16) * 8 + i][l % 16], i = 0..7.  gfx11 runs the same v_wmma_f32_16x16x16_f16 instruction (RDNA3 introduced
+// it; RDNA4 kept it), through the unsuffixed builtin.
 #if defined(__gfx1200__) || defined(__gfx1201__)
 #define STRATA_PA_WMMA 1
+#define STRATA_PA_WMMA_GFX12 1
+#elif defined(__gfx1100__)
+#define STRATA_PA_WMMA 1
+#define STRATA_PA_WMMA_GFX11 1
 #else
 #define STRATA_PA_WMMA 0
 #endif
 typedef _Float16 wh8 __attribute__((ext_vector_type(8)));
+typedef _Float16 wh16 __attribute__((ext_vector_type(16)));
 typedef float wf8 __attribute__((ext_vector_type(8)));
 constexpr int WCH = 32;          // cells per chunk (two 16-cell tiles)
 constexpr int WVS = 80;          // staged V row stride in bytes (64 codes, 16-byte aligned, banks spread)
@@ -1150,15 +1157,27 @@ struct alignas(16) SmemW {
     float mrow[16];
 };
 
+// RDNA4: the fragment A/B operands are 8 halves per lane.  RDNA3's unsuffixed builtin takes 16 (8 VGPRs), with the
+// mappings from the AMD Matrix Instruction Calculator (rdna3, v_wmma_f32_16x16x16_f16, wave32, no modifiers):
+// A[l & 15][i] on lane l slot i, B[j][l & 15] on lane l slot j, D[(l >> 4) + 2*j][l & 15] on lane l slot j.
 __device__ __forceinline__ wf8 wmma_f16(wh8 a, wh8 b, wf8 c) {
-#if STRATA_PA_WMMA
+#if defined(STRATA_PA_WMMA_GFX12)
     return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a, b, c);
 #else
     __builtin_trap();
     return c;
 #endif
 }
 
+__device__ __forceinline__ wf8 wmma_f16(wh16 a, wh16 b, wf8 c) {
+#if defined(STRATA_PA_WMMA_GFX11)
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
+#else
+    __builtin_trap();
+    return c;
+#endif
+}
+
 __device__ __forceinline__ wh8 i8x8_to_h8(uint2 x) {   // exact: |code| <= 128
     wh8 h;
 #pragma unroll
@@ -1169,6 +1188,19 @@ __device__ __forceinline__ wh8 i8x8_to_h8(uint2 x) {   // exact: |code| <= 128
     return h;
 }
 
+__device__ __forceinline__ wh16 i8x16_to_h16(uint4 x) {   // exact: |code| <= 128 (RDNA3 K loads)
+    wh16 h;
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+        const unsigned w = i == 0 ? x.x : i == 1 ? x.y : i == 2 ? x.z : x.w;
+        h[4 * i + 0] = (_Float16) (int) (int8_t) (w & 0xffu);
+        h[4 * i + 1] = (_Float16) (int) (int8_t) ((w >> 8) & 0xffu);
+        h[4 * i + 2] = (_Float16) (int) (int8_t) ((w >> 16) & 0xffu);
+        h[4 * i + 3] = (_Float16) (int) (int8_t) ((w >> 24) & 0xffu);
+    }
+    return h;
+}
+
 __global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float* __restrict__ q, QsaAttnPools p,
                                                                    const int32_t* __restrict__ ids,
                                                                    const int32_t* __restrict__ steps, int n_kv_heads,
@@ -1197,6 +1229,31 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float*
     int qe = 0;
     if (qm > 0.0f) frexpf(qm, &qe);
     const float qup = ldexpf(1.0f, 14 - qe), qdown = ldexpf(scale_log2, qe - 14);
+#if defined(STRATA_PA_WMMA_GFX11)
+    // q A fragments of this wave's 64 dims: the RDNA3 layout holds A[l & 15][k] on lane l (row duplicated across
+    // the lane halves), k = dim0 + kk*16 + i, i = 0..15 (AMD Matrix Instruction Calculator, rdna3 table)
+    wh16 qh[4], ql[4];
+#pragma unroll
+    for (int kk = 0; kk < 4; ++kk) {
+        float x[16];
+        if (col < G) {
+            const float4* src = reinterpret_cast<const float4*>(q + (size_t) col * HD + dim0 + kk * 16);
+            const float4 a = src[0], b = src[1], c = src[2], d = src[3];
+            x[0] = a.x; x[1] = a.y; x[2] = a.z; x[3] = a.w; x[4] = b.x; x[5] = b.y; x[6] = b.z; x[7] = b.w;
+            x[8] = c.x; x[9] = c.y; x[10] = c.z; x[11] = c.w; x[12] = d.x; x[13] = d.y; x[14] = d.z; x[15] = d.w;
+        } else {
+#pragma unroll
+            for (int i = 0; i < 16; ++i) x[i] = 0.0f;
+        }
+#pragma unroll
+        for (int i = 0; i < 16; ++i) {
+            const float v = x[i] * qup;
+            const _Float16 hi = (_Float16) v;
+            qh[kk][i] = hi;
+            ql[kk][i] = (_Float16) (v - (float) hi);
+        }
+    }
+#else
     // q A fragments of this wave's 64 dims: row = col (heads 12..15 zero), k = dim0 + kk*16 + half*8 + i
     wh8 qh[4], ql[4];
 #pragma unroll
@@ -1218,6 +1275,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float*
             ql[kk][i] = (_Float16) (v - (float) hi);
         }
     }
+#endif
 
     wf8 acc[4];
 #pragma unroll
@@ -1252,6 +1310,23 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float*
         // q.k over this wave's 64 dims for two 16-cell tiles, times the cell's K scale for this group
 #pragma unroll
         for (int nt = 0; nt < 2; ++nt) {
+#if defined(STRATA_PA_WMMA_GFX11)
+            // RDNA3: the lane's B/K cell is nt*16 + (l & 15) (B[k][n] sits on lane n), one cell per lane as on gfx12
+            const long long rr = __shfl(myrow, nt * 16 + col);
+            wf8 s = wf8{0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
+            float ksc = 0.0f;
+            if (rr >= 0) ksc = __half2float(__ushort_as_half(__ldg(p.k_scale + rr * (HD / KV_Q8_GROUP) + warp)));
+#pragma unroll
+            for (int kk = 0; kk < 4; ++kk) {
+                uint4 kx = make_uint4(0, 0, 0, 0);
+                if (rr >= 0) kx = __ldg(reinterpret_cast<const uint4*>(p.k_q + rr * HD + dim0 + kk * 16));
+                const wh16 b = i8x16_to_h16(kx);
+                s = wmma_f16(qh[kk], b, s);
+                s = wmma_f16(ql[kk], b, s);
+            }
+#pragma unroll
+            for (int i = 0; i < 8; ++i) S.part[warp][half + 2 * i][nt * 16 + col] = s[i] * ksc;
+#else
             const long long rr = __shfl(myrow, nt * 16 + col);
             wf8 s = wf8{0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
             float ksc = 0.0f;
@@ -1266,6 +1341,7 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float*
             }
 #pragma unroll
             for (int i = 0; i < 8; ++i) S.part[warp][half * 8 + i][nt * 16 + col] = s[i] * ksc;
+#endif
         }
         __syncthreads();
         // online softmax over the four groups' sum (fixed order): row t/8, 4 cells per thread
@@ -1310,6 +1386,38 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float*
             wf8 tmp[4];
 #pragma unroll
             for (int j = 0; j < 4; ++j) tmp[j] = wf8{0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
+#if defined(STRATA_PA_WMMA_GFX11)
+            // RDNA3: A = p[head = col][cells ks*16 + i] on lane col (i = 0..15); B = V[cell = ks*16 + j][dim of this
+            // lane (col) within the tile] on lane col slot j; D slot i = (head, dim) at [(l>>4) + 2*i][l & 15]
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks) {
+                const int cb = ks * 16;
+                wh16 ah, al;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const float pv = S.p[col][cb + i] * (S.vs[warp][cb + i] * vup);
+                    const _Float16 hi = (_Float16) pv;
+                    ah[i] = hi;
+                    al[i] = (_Float16) (pv - (float) hi);
+                }
+#pragma unroll
+                for (int j = 0; j < 4; ++j) {
+                    const int d = j * 16 + col;
+                    wh16 b;
+#pragma unroll
+                    for (int i = 0; i < 16; ++i) b[i] = (_Float16) (int) (int8_t) S.v[warp][cb + i][d];
+                    tmp[j] = wmma_f16(ah, b, tmp[j]);
+                    tmp[j] = wmma_f16(al, b, tmp[j]);
+                }
+            }
+            float a[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i) a[i] = S.alpha[half + 2 * i];
+#pragma unroll
+            for (int j = 0; j < 4; ++j)
+#pragma unroll
+                for (int i = 0; i < 8; ++i) acc[j][i] = fmaf(acc[j][i], a[i], tmp[j][i] * vdown);
+#else
 #pragma unroll
             for (int ks = 0; ks < 2; ++ks) {
                 const int cb = ks * 16 + half * 8;   // this lane's 8 cells
@@ -1338,42 +1446,54 @@ __global__ void __launch_bounds__(THREADS) prompt_attn_wmma_kernel(const float*
             for (int j = 0; j < 4; ++j)
 #pragma unroll
                 for (int i = 0; i < 8; ++i) acc[j][i] = fmaf(acc[j][i], a[i], tmp[j][i] * vdown);
+#endif
         }
         // no barrier here: `v`/`vs` are this wave's own, and `p`/`alpha`/`valid` are rewritten only after the next
         // chunk's first barrier (which every wave reaches after its p.v)
     }
     float inv[8];
 #pragma unroll
     for (int i = 0; i < 8; ++i) {
+#if defined(STRATA_PA_WMMA_GFX11)
+        const float l = S.lsum[half + 2 * i];
+#else
         const float l = S.lsum[half * 8 + i];
+#endif
         inv[i] = l > 0.0f ? 1.0f / l : 0.0f;
     }
 #pragma unroll
     for (int j = 0; j < 4; ++j)
 #pragma unroll
         for (int i = 0; i < 8; ++i) {
+#if defined(STRATA_PA_WMMA_GFX11)
+            const int row = half + 2 * i;
+#else
             const int row = half * 8 + i;
+#endif
             if (row < G) attn[(size_t) row * HD + dim0 + j * 16 + col] = acc[j][i] * inv[i];
         }
 #else
     __builtin_trap();
 #endif
 }
 
-// gfx12 (RDNA4) only, and only on request: the output differs from the default kernel's in its last bits
+// RDNA4/RDNA3 matrix cores, and only on request: the output differs from the default kernel's in its last bits.
+// The arch prefixes are checked at RUNTIME on purpose: the host pass of a HIP TU sees no arch macros, and HIP only
+// loads the code object that matches the card, so a prefix match here is the same kernel the card will run.
 bool hip_wmma_usable() {
     static const bool want = [] {
         const char* e = std::getenv("STRATA_HIP_WMMA");
         return e != nullptr && e[0] == '1';
     }();
     if (!want) return false;
-    static int arch[64] = {};   // 0 unknown, 1 gfx12, 2 other
+    static int arch[64] = {};   // 0 unknown, 1 gfx12/gfx11, 2 other
     int dev = 0;
     if (hipGetDevice(&dev) != hipSuccess || dev < 0 || dev >= 64) { (void) hipGetLastError(); return false; }
     if (arch[dev] == 0) {
         hipDeviceProp_t prop{};
         if (hipGetDeviceProperties(&prop, dev) != hipSuccess) { (void) hipGetLastError(); return false; }
-        arch[dev] = std::strncmp(prop.gcnArchName, "gfx12", 5) == 0 ? 1 : 2;
+        arch[dev] = (std::strncmp(prop.gcnArchName, "gfx12", 5) == 0 ||
+                     std::strncmp(prop.gcnArchName, "gfx11", 5) == 0) ? 1 : 2;
         static bool told = false;
         if (!told) {
             told = true;
@@ -1435,8 +1555,8 @@ bool qsa_prompt_attn_batch(const float* q, const QsaAttnPools& pools, const int3
         turing = cc[dev] < 80;
     }
 #if defined(__HIPCC__)
-    // the tensor-core kernels are compiled out on AMD (its major version is not a CUDA sm); RDNA4 has its own int8-KV
-    // matrix-core kernel, opt-in (STRATA_HIP_WMMA=1); everything else keeps the old kernel
+    // the tensor-core kernels are compiled out on AMD (its major version is not a CUDA sm); RDNA4/RDNA3 have their
+    // own int8-KV matrix-core kernel, opt-in (STRATA_HIP_WMMA=1); everything else keeps the old kernel
     (void) turing; (void) volta;
     if (pools.k_q != nullptr && pools.v_q != nullptr && pools.k_scale != nullptr && pools.v_scale != nullptr &&
         pools.k_q4 == nullptr && pools.v_q4 == nullptr && s.head_dim == HD && s.n_head == (int64_t) G * s.n_head_kv &&