QSA select: RDNA3 (gfx1100) WMMA block-scores arm on #337's arch dispatch (opt-in STRATA_SELECT_WMMA=1)

CMakeLists.txt

@@ -283,6 +283,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_select.cu

@@ -1,11 +1,46 @@
 // src/kernels/cuda/qsa_select.cu - see include/strata/kernels/qsa_select.hpp.
+
 #include "strata/core/emulate.hpp"
 #include <cstdlib>
 #include <cstring>
+//
+// AMD tensor cores: ROCWMMA (16x16x16, FP16 operands, FP32 accumulation) where the CUDA TF32 path cannot run
+// (gfx11 has no mma.sync).  The CUDA-compat shim (hip_compat/intrinsics.hpp, force-included before this file)
+// renames the CUDA warp intrinsics with macros - and hip's own bf16/fp8 headers, which rocWMMA pulls in, use
+// those same tokens as function NAMES.  The macros are therefore lifted while rocWMMA parses and put back
+// after; they must match intrinsics.hpp exactly.  Without the rocWMMA headers the arm compiles out and the
+// host entry refuses - the warp kernel stays the fallback everywhere.
+#if defined(__HIPCC__) && defined(__has_include)
+#if __has_include(<rocwmma/rocwmma.hpp>)
+#define STRATA_QSA_HAVE_ROCWMMA 1
+#undef __dp4a
+#undef __vsub4
+#undef __vsubss4
+#undef __vcmpne4
+#undef __shfl_xor_sync
+#undef __shfl_down_sync
+#undef __shfl_up_sync
+#undef __shfl_sync
+#undef __ballot_sync
+#include <rocwmma/rocwmma.hpp>
+#define __dp4a(a, b, c) (::strata::hip_compat::dp4a((a), (b), (c)))
+#define __vsub4(a, b) (::strata::hip_compat::vsub4((a), (b)))
+#define __vsubss4(a, b) (::strata::hip_compat::vsubss4((a), (b)))
+#define __vcmpne4(a, b) (::strata::hip_compat::vcmpne4((a), (b)))
+#define __shfl_xor_sync(...) (::strata::hip_compat::shfl_xor_sync(__VA_ARGS__))
+#define __shfl_down_sync(...) (::strata::hip_compat::shfl_down_sync(__VA_ARGS__))
+#define __shfl_up_sync(...) (::strata::hip_compat::shfl_up_sync(__VA_ARGS__))
+#define __shfl_sync(...) (::strata::hip_compat::shfl_sync(__VA_ARGS__))
+#define __ballot_sync(mask, predicate) (::strata::hip_compat::ballot_sync((mask), (predicate)))
+#endif
+#endif
+
+
 #include "strata/kernels/qsa_select.hpp"
 
 #include <cuda_runtime.h>
 
+#include <algorithm>
 #include <cfloat>
 #include <cstdio>
 #include <cstdlib>
@@ -461,6 +496,62 @@ __global__ void __launch_bounds__(32) block_scores_tail_kernel(const float* __re
     }
 }
 
+// ---- AMD tensor cores (rocWMMA): the same scores as block_scores_kernel with FP16 operands and FP32 accumulation
+// in another (unspecified) summation order - the PORTING.md 18 bound covers exactly that difference.  The tail
+// block is NOT computed here: block_scores_tail_kernel runs unchanged and keeps it bitwise the warp kernel's.
+// Layout: one wave (32 lanes) per CTA computes a 16-row x 16-block tile; a row is one (query, indexer head) pair
+// (4 consecutive queries x IDX_HEADS heads, q_idx's own row order), the K axis is IDX_DIM in 16-wide slices.
+#if defined(STRATA_QSA_HAVE_ROCWMMA)
+constexpr int WM_BT = 16;                                // key blocks per tile (mma n16)
+constexpr int WM_ROWS = 4 * IDX_HEADS;                   // (query, head) rows per tile: 4 queries
+
+__global__ void __launch_bounds__(256) f32_to_f16_poison_kernel(const float* __restrict__ in, size_t n,
+                                                                __half* __restrict__ out, size_t total) {
+    for (size_t i = blockIdx.x * (size_t) blockDim.x + threadIdx.x; i < total; i += (size_t) gridDim.x * blockDim.x)
+        out[i] = __float2half(i < n ? in[i] : 60000.0f);   // poisoned slack must never reach a score
+}
+
+__global__ void __launch_bounds__(32) block_scores_wmma_kernel(const __half* __restrict__ pooled16,
+                                                               const __half* __restrict__ q16,
+                                                               const int32_t* __restrict__ steps, int nq,
+                                                               int64_t max_blocks, float* __restrict__ out) {
+    const int q0 = (int) blockIdx.y * 4;
+    const int64_t b0 = (int64_t) blockIdx.x * WM_BT;
+    if (q0 >= nq || b0 >= max_blocks) return;
+    int64_t n_bid[4];
+#pragma unroll
+    for (int q = 0; q < 4; ++q)
+        n_bid[q] = (q0 + q) < nq ? steps[(size_t) (q0 + q) * kStepCount + kStepNBid] : (int64_t) -1;
+    rocwmma::fragment<rocwmma::accumulator, 16, 16, 16, float> acc;
+    rocwmma::fill_fragment(acc, 0.0f);
+#pragma unroll
+    for (int ks = 0; ks < IDX_DIM / 16; ++ks) {
+        rocwmma::fragment<rocwmma::matrix_a, 16, 16, 16, __half, rocwmma::row_major> a;
+        rocwmma::fragment<rocwmma::matrix_b, 16, 16, 16, __half, rocwmma::col_major> b;
+        rocwmma::load_matrix_sync(a, q16 + (size_t) q0 * IDX_HEADS * IDX_DIM + ks * 16, IDX_DIM);
+        rocwmma::load_matrix_sync(b, pooled16 + b0 * IDX_DIM + ks * 16, IDX_DIM);
+        rocwmma::mma_sync(acc, a, b, acc);
+    }
+    __shared__ float tile[WM_ROWS * WM_BT];
+    rocwmma::store_matrix_sync(tile, acc, WM_BT, rocwmma::mem_row_major);
+    __syncthreads();
+    // tile[q * IDX_HEADS + h][c] = dot(query q0 + q, head h, pooled block b0 + c)
+    for (int i = threadIdx.x; i < 4 * WM_BT; i += blockDim.x) {
+        const int q = i / WM_BT, c = i - q * WM_BT;
+        const int64_t b = b0 + c;
+        const int qi = q0 + q;
+        if (qi >= nq || b >= max_blocks || b >= n_bid[q]) continue;   // b == n_bid is the tail kernel's
+        float s = 0.0f;
+#pragma unroll
+        for (int h = 0; h < IDX_HEADS; ++h) {
+            const float d = tile[(q * IDX_HEADS + h) * WM_BT + c];
+            s += d > 0.0f ? d : 0.0f;
+        }
+        out[(size_t) qi * max_blocks + b] = s;
+    }
+}
+#endif   // STRATA_QSA_HAVE_ROCWMMA
+
 // ---- the same top-k with each query's keys read once: 1,024 threads hold up to TK_PER consecutive blocks' keys in
 // registers (contexts up to 4 * 1024 * TK_PER cells), per-warp histograms, block-wide scans. The selection rule is
 // block_topk_kernel's (radix threshold, ties to the lowest index, cells ascending): identical ids.
@@ -904,17 +995,25 @@ bool qsa_block_scores_tc(const float* pooled, const float* dead, const float* q_
         const char* v = std::getenv("STRATA_SELECT_WMMA");
         return v != nullptr && v[0] != '\0' && v[0] != '0';
     }();
-    if (!wmma_on || !sel_gfx12_device()) return false;
-    {
+    if (wmma_on && sel_gfx12_device()) {
         const int64_t reach = active_blocks > 0 && active_blocks < max_blocks ? active_blocks : max_blocks;
         const int64_t per = (int64_t) 4 * WITER * 16;
         const dim3 grid((unsigned) ((reach + per - 1) / per), (unsigned) ((nq + WQT - 1) / WQT));
         block_scores_wmma_kernel<<<grid, 128, 0, (cudaStream_t) stream>>>(pooled, q_idx, steps, nq, max_blocks, reach, scores);
         block_scores_tail_kernel<<<(unsigned) nq, 32, 0, (cudaStream_t) stream>>>(dead, q_idx, steps, max_blocks, scores);
         const cudaError_t e = cudaGetLastError();
         if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_scores_tc: %s\n", cudaGetErrorString(e)); std::exit(1); }
+        static int said12 = 0;
+        if (!said12) { said12 = 1; std::fprintf(stderr, "strata select: prompt scorer arm = wmma (gfx12, opt-in)\n"); }
+        return true;
+    }
+    if (wmma_on &&
+        qsa_block_scores_wmma(pooled, dead, q_idx, steps, nq, max_blocks, s, scores, stream, active_blocks)) {
+        static int said11 = 0;
+        if (!said11) { said11 = 1; std::fprintf(stderr, "strata select: prompt scorer arm = wmma (gfx1100, opt-in)\n"); }
         return true;
     }
+    return false;
 #else
     {   // sm_80 or newer (TF32 MMA); an older card keeps the warp kernel
         static int cc_major[64] = {};
@@ -956,6 +1055,67 @@ bool qsa_block_scores_tc(const float* pooled, const float* dead, const float* q_
 #endif
 }
 
+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) {
+    if (nq <= 0) return true;
+    if (s.idx_dim != IDX_DIM || s.idx_n_head != IDX_HEADS || s.idx_block != R || nq > 65535 * 4) return false;
+#if defined(STRATA_QSA_HAVE_ROCWMMA)
+    static int arch[64] = {};   // 1 = gfx1100 (verified), -1 = anything else: refuse, the caller keeps its fallback
+    int dev = 0;
+    if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) { cudaGetLastError(); return false; }
+    if (arch[dev] == 0) {
+        hipDeviceProp_t p{};    // gcnArchName is a HIP field; this whole branch is HIP-only
+        if (hipGetDeviceProperties(&p, dev) != hipSuccess) { cudaGetLastError(); return false; }
+        arch[dev] = p.gcnArchName[0] == 'g' && p.gcnArchName[1] == 'f' && p.gcnArchName[2] == 'x' &&
+                            p.gcnArchName[3] == '1' && p.gcnArchName[4] == '1' && p.gcnArchName[5] == '0' &&
+                            p.gcnArchName[6] == '0'
+                        ? 1
+                        : -1;
+    }
+    if (arch[dev] < 0) return false;
+    // the conversion scratch: keys up to whole tiles, query rows up to whole 4-query tiles, poison in the slack
+    const int64_t reach = active_blocks > 0 && active_blocks < max_blocks ? active_blocks : max_blocks;
+    const size_t kb = (size_t) ((reach + WM_BT - 1) / WM_BT) * WM_BT * IDX_DIM;
+    const size_t qb = (size_t) ((nq + 3) / 4) * 4 * IDX_HEADS * IDX_DIM;
+    static __half* scratch[64] = {};
+    static size_t scratch_n[64] = {};
+    if (scratch_n[dev] < kb + qb) {
+        cudaFree(scratch[dev]);
+        cudaGetLastError();
+        scratch[dev] = nullptr;
+        scratch_n[dev] = 0;
+        if (cudaMalloc((void**) &scratch[dev], (kb + qb) * sizeof(__half)) != cudaSuccess) {
+            cudaGetLastError();
+            return false;
+        }
+        scratch_n[dev] = kb + qb;
+    }
+    const auto conv = [&](const float* in, size_t n, __half* out, size_t total) {
+        const unsigned blocks = (unsigned) std::min<size_t>((total + 255) / 256, 65535);
+        f32_to_f16_poison_kernel<<<blocks, 256, 0, (cudaStream_t) stream>>>(in, n, out, total);
+    };
+    conv(pooled, (size_t) reach * IDX_DIM, scratch[dev], kb);
+    conv(q_idx, (size_t) nq * IDX_HEADS * IDX_DIM, scratch[dev] + kb, qb);
+    block_scores_wmma_kernel<<<dim3((unsigned) ((reach + WM_BT - 1) / WM_BT), (unsigned) ((nq + 3) / 4)), 32, 0,
+                               (cudaStream_t) stream>>>(scratch[dev], scratch[dev] + kb, steps, (int) nq, max_blocks,
+                                                        scores);
+    block_scores_tail_kernel<<<(unsigned) nq, 32, 0, (cudaStream_t) stream>>>(dead, q_idx, steps, max_blocks, scores);
+    const cudaError_t e = cudaGetLastError();
+    if (e != cudaSuccess) { std::fprintf(stderr, "qsa_block_scores_wmma: %s\n", cudaGetErrorString(e)); std::exit(1); }
+    return true;
+#else
+    (void) pooled;
+    (void) dead;
+    (void) q_idx;
+    (void) steps;
+    (void) max_blocks;
+    (void) scores;
+    (void) stream;
+    (void) active_blocks;
+    return false;   // an AMD-only arm: a CUDA build keeps the TF32 path and the warp kernel
+#endif
+}
+
 void qsa_block_topk_ref(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) {
     if (nq <= 0) return;