qsa_select: FP32 tiled block scores below sm_80 (opt-in: STRATA_SELECT_SIMT=1)

Based on konijiwa110's PR #742. Below sm_80 there is no TF32 tensor-core scorer, so the prompt path scores the QSA
blocks on the warp kernel: 5 shuffles per 4 products, ~1 TFLOPS on an RTX 2080 Ti, 24% of a 128K prompt. With
STRATA_SELECT_SIMT=1, qsa_block_scores_tc runs an FP32 tiled GEMM there (rows = query x indexer head, K = 128 through
shared memory, 2 queries x 4 heads x 4 blocks per thread, relu per head summed in registers), ~6.9 TFLOPS by the
author's measurement (128K prompt 1,016 to 1,242 tok/s, 248K 715 to 952 on one 2080 Ti).

The tiled kernel sums in another order than the warp kernel, so its scores are not bitwise equal. The PR made it the
default below sm_80; here it is opt-in, so default output is unchanged until an interleaved A/B of at least 10 pairs
supports a default. The tail block keeps the warp kernel's arithmetic. STRATA_SELECT_OLD=1 still forces the warp kernel.

qsa_select_bench is now also a CUDA test: at 32K, and at 131K with STRATA_QSA_WARP=1 and STRATA_SELECT_SIMT=1 so any
card runs the FP32 tiled scorer (FP64 gate, selections compared with the warp kernel's).

Co-authored-by: konijiwa110 <[email protected]>

CMakeLists.txt

@@ -378,6 +378,16 @@ 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 same on CUDA: the tensor-core scorer (sm_80+) or the FP32 tiled one (below sm_80, and forced on any card by
+  # STRATA_QSA_WARP=1 in the second test) against the warp kernel and FP64, at 32K and 131K
+  if(STRATA_BUILD_TESTS AND STRATA_ENABLE_CUDA)
+    add_executable(qsa_select_bench src/kernels/qsa_select_bench.cpp)
+    target_include_directories(qsa_select_bench PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
+    target_link_libraries(qsa_select_bench PRIVATE strata_kernels CUDA::cudart)
+    add_test(NAME qsa_select_bench COMMAND qsa_select_bench 32768 256 3 262144)
+    add_test(NAME qsa_select_bench_fp32_tiled COMMAND qsa_select_bench 131072 256 3 262144)
+    set_tests_properties(qsa_select_bench_fp32_tiled PROPERTIES ENVIRONMENT "STRATA_QSA_WARP=1;STRATA_SELECT_SIMT=1")
+  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

@@ -32,6 +32,7 @@ void qsa_block_scores(const float* pooled, const float* dead, const float* q_idx
 /// The same scores on tensor cores (3xTF32, FP32-level accuracy but another summation order: not bitwise; the tail
 /// block n_bid is the warp kernel's arithmetic). For the prompt path; false (nothing launched) on another geometry.
 /// On AMD it is the gfx12 (RDNA4) WMMA kernel (a three-way bf16 split, six products); false on any other AMD target.
+/// On CUDA below sm_80 (no TF32 tensor cores) it is an FP32 tiled kernel, also not bitwise, and only with STRATA_SELECT_SIMT=1 (otherwise false).
 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);
 

src/kernels/cuda/qsa_select.cu

@@ -484,6 +484,74 @@ __global__ void __launch_bounds__(32) block_scores_tail_kernel(const float* __re
     }
 }
 
+// ---- the block scores as an FP32 tiled GEMM for CUDA cards without TF32 tensor cores (below sm_80): rows (query,
+// indexer head) x columns (blocks), K = 128 through shared memory, 2 queries x 4 heads x 4 blocks per thread, relu per
+// head summed in registers.  The warp kernel spends 5 shuffles per 4 products and reached ~1 TFLOPS on an RTX 2080
+// Ti; this one ~6.9 (256 queries: 4.5x at 4K, 6.8x at 120K-250K).  FP32 like the warp kernel in another summation
+// order, so not bitwise (as the tensor-core scorer).  Blocks < n_bid only; the tail block is block_scores_tail_kernel's.
+constexpr int SM_QT = 32, SM_NB = 64, SM_KC = 32, SM_QS = SM_QT * IDX_HEADS + 4, SM_KS = SM_NB + 4;
+__global__ void __launch_bounds__(256) block_scores_simt_kernel(const float* __restrict__ pooled,
+                                                                const float* __restrict__ q_idx,
+                                                                const int32_t* __restrict__ steps, int64_t nq,
+                                                                int64_t max_blocks, int64_t reach,
+                                                                float* __restrict__ out) {
+    __shared__ __align__(16) float Qs[SM_KC][SM_QS];
+    __shared__ __align__(16) float Ks[SM_KC][SM_KS];
+    const int64_t q0 = (int64_t) blockIdx.y * SM_QT, b0 = (int64_t) blockIdx.x * SM_NB;
+    const int64_t qlast = (q0 + SM_QT < nq ? q0 + SM_QT : nq) - 1;
+    if (b0 >= steps[qlast * kStepCount + kStepNBid]) return;   // n_bid rises with the position
+    const int tid = threadIdx.x, tx = tid & 15, ty = tid >> 4;
+    float acc[8][4] = {};
+    for (int kc = 0; kc < IDX_DIM; kc += SM_KC) {
+#pragma unroll
+        for (int i = 0; i < 4; ++i) {   // the query rows: 128 x 8 float4
+            const int idx = tid + i * 256, r = idx >> 3, k4 = idx & 7;
+            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
+            if (q0 + (r >> 2) < nq)
+                v = *reinterpret_cast<const float4*>(q_idx + (q0 * IDX_HEADS + r) * IDX_DIM + kc + k4 * 4);
+            Qs[k4 * 4 + 0][r] = v.x; Qs[k4 * 4 + 1][r] = v.y; Qs[k4 * 4 + 2][r] = v.z; Qs[k4 * 4 + 3][r] = v.w;
+        }
+#pragma unroll
+        for (int i = 0; i < 2; ++i) {   // the pooled keys: 64 x 8 float4
+            const int idx = tid + i * 256, c = idx >> 3, k4 = idx & 7;
+            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
+            if (b0 + c < reach) v = *reinterpret_cast<const float4*>(pooled + (b0 + c) * IDX_DIM + kc + k4 * 4);
+            Ks[k4 * 4 + 0][c] = v.x; Ks[k4 * 4 + 1][c] = v.y; Ks[k4 * 4 + 2][c] = v.z; Ks[k4 * 4 + 3][c] = v.w;
+        }
+        __syncthreads();
+#pragma unroll 8
+        for (int k = 0; k < SM_KC; ++k) {
+            const float4 qa = *reinterpret_cast<const float4*>(&Qs[k][ty * 8]);
+            const float4 qb = *reinterpret_cast<const float4*>(&Qs[k][ty * 8 + 4]);
+            const float4 kv = *reinterpret_cast<const float4*>(&Ks[k][tx * 4]);
+            const float qr[8] = {qa.x, qa.y, qa.z, qa.w, qb.x, qb.y, qb.z, qb.w}, kr[4] = {kv.x, kv.y, kv.z, kv.w};
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+#pragma unroll
+                for (int j = 0; j < 4; ++j) acc[i][j] = fmaf(qr[i], kr[j], acc[i][j]);
+        }
+        __syncthreads();
+    }
+#pragma unroll
+    for (int i = 0; i < 2; ++i) {
+        const int64_t q = q0 + ty * 2 + i;
+        if (q >= nq) continue;
+        const int64_t n_bid = steps[q * kStepCount + kStepNBid];
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+            const int64_t b = b0 + tx * 4 + j;
+            if (b >= n_bid || b >= max_blocks) continue;
+            float sc = 0.0f;
+#pragma unroll
+            for (int h = 0; h < IDX_HEADS; ++h) {
+                const float d = acc[i * 4 + h][j];
+                sc += d > 0.0f ? d : 0.0f;
+            }
+            out[q * max_blocks + b] = sc;
+        }
+    }
+}
+
 // ---- 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.
@@ -1084,7 +1152,25 @@ bool qsa_block_scores_tc(const float* pooled, const float* dead, const float* q_
             const char* w = std::getenv("STRATA_QSA_WARP");
             cc_major[dev] = w && (!std::strcmp(w, "1") || !std::strcmp(w, "select")) ? 7 : strata::cc_major_of(major);
         }
-        if (cc_major[dev] < 8) return false;
+        if (cc_major[dev] < 8) {
+            // no TF32 tensor cores: the FP32 tiled kernel, opt-in (STRATA_SELECT_SIMT=1): its sum order differs from the
+            // warp kernel's, so the default stays the caller's warp kernel and the output bytes stay the same
+            static const bool simt = [] {
+                const char* v = std::getenv("STRATA_SELECT_SIMT");
+                return v != nullptr && v[0] == '1';
+            }();
+            if (!simt) return false;
+            const int64_t reach = active_blocks > 0 && active_blocks < max_blocks ? active_blocks : max_blocks;
+            const dim3 grid((unsigned) ((reach + SM_NB - 1) / SM_NB), (unsigned) ((nq + SM_QT - 1) / SM_QT));
+            if (grid.y > 65535) return false;
+            block_scores_simt_kernel<<<grid, 256, 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 (simt): %s\n", cudaGetErrorString(e)); std::exit(1); }
+            return true;
+        }
     }
     static bool attr[64] = {};   // the shared-memory opt-in is per device (a layer split runs it on several)
     int adev = 0;

src/kernels/qsa_select_bench.cpp

@@ -1,5 +1,5 @@
 // src/kernels/qsa_select_bench.cpp - the prompt path's QSA selection (qsa_select.hpp) timed per stage, block scores
-// (the warp kernel and the tensor-core one) and top-k, for a batch of consecutive queries at a given context, and
+// (the warp kernel and the fast one: tensor cores, or the FP32 tiled kernel on CUDA below sm_80) and top-k, for a batch of consecutive queries at a given context, and
 // the two scorers compared: score difference and how many selections differ (GPU, synthetic, no model).
 // Usage: qsa_select_bench [context=131072] [queries=256] [reps=10] [capacity_cells]
 // capacity_cells (the engine's --max-context): the score buffers and the top-k dispatch follow the CAPACITY
@@ -163,9 +163,9 @@ int main(int argc, char** argv) {
     const float t_tk2 = timed([&] { k::qsa_block_topk(sc_old, d_steps, nq, max_blocks, cap, s, ids_reg, nullptr, active); });
     std::printf("top-k %.3f -> %.3f ms (%.1fx), register top-k identical to the reference %lld/%lld\n", t_tk, t_tk2,
                 t_tk / t_tk2, (long long) reg_same, (long long) nq);
-    std::printf("%s accuracy vs FP64 (score scale %.3g): warp kernel max err %.3g, tensor-core max err %.3g (%.2g of scale)\n",
+    std::printf("%s accuracy vs FP64 (score scale %.3g): warp kernel max err %.3g, fast scorer max err %.3g (%.2g of scale)\n",
                 !have_tc ? "SKIP" : acc_ok ? "PASS" : "FAIL", scale, err_old, err_new, scale > 0 ? err_new / scale : 0.0);
-    if (!have_tc) std::printf("tensor-core scorer not available on this device: warp scorer %.3f ms, top-k %.3f ms (%.0f%% of the two)\n", t_old, t_tk2, 100.0 * t_tk2 / (t_old + t_tk2));
+    if (!have_tc) std::printf("fast scorer not available on this device: warp scorer %.3f ms, top-k %.3f ms (%.0f%% of the two)\n", t_old, t_tk2, 100.0 * t_tk2 / (t_old + t_tk2));
     std::printf("ctx %lld, %lld queries x %lld blocks: scores %.3f -> %.3f ms (%.1fx), top-k %.3f ms; score rel diff "
                 "mean %.2g max %.2g; selections identical %lld/%lld, cells differing %.4f%%\n", (long long) ctx,
                 (long long) nq, (long long) active, t_old, t_new, t_old / t_new, t_tk, sum_rel / std::max(1.0, n_rel),

src/prefill/prefill.cpp

@@ -2363,7 +2363,8 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
                         // C-1: the grid reaches the batch's last query's n_bid (they rise with the position)
                         const int64_t active = (int64_t) m.steps_host[(size_t) ((t0 + nb - 1) * strata::kernels::kStepCount +
                                                                                 strata::kernels::kStepNBid)] + 1;
-                        // the scores on tensor cores (3xTF32: FP32-level, not bitwise); STRATA_SELECT_OLD=1: the warp kernel
+                        // the scores on tensor cores (3xTF32: FP32-level, not bitwise), below sm_80 the FP32 tiled kernel with STRATA_SELECT_SIMT=1;
+                        // STRATA_SELECT_OLD=1: the warp kernel
                         static const bool old_sel = std::getenv("STRATA_SELECT_OLD") != nullptr;
                         if (old_sel || !strata::kernels::qsa_block_scores_tc(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512,
                                                                              steps0, nb, m.max_blocks, s, m.sel_scores,