prefill: STRATA_HCD_EXACT=1 - the HC down projection by a WMMA kernel with hipBLASLt solution 1176/1177 k order (StaggerU included), bitwise

include/strata/prefill/gemm.hpp

@@ -27,7 +27,12 @@ public:
 
     /// Y[T, N] (fp32, row stride ldy) = X[T, K] (bf16, row-major) . W[N, K]^T (bf16, row-major).  `beta` = 1 adds.
     void bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0,
-              float beta = 0.0f);
+              float beta = 0.0f, int64_t ldx = 0);
+
+    /// S (STRATA_HCD_EXACT): the HC down projection (N 320, K 10240) by the WMMA kernel that reproduces hipBLASLt's
+    /// solution 1176 / 1177 bit for bit; false (nothing launched) unless hipBLASLt would take one of those for this
+    /// shape (so the caller falls back to bf16(), with the same ldx).
+    bool bf16_hcd_exact(const uint16_t* X, int64_t ldx, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K);
 
     /// Y = X . W^T with both in FP16 (bits).
     void f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy = 0,

src/prefill/gemm.cu

@@ -127,7 +127,7 @@ struct HipLtDescriptors {
         if (c) hipblasLtMatrixLayoutDestroy(c);
     }
 
-    bool init(hipDataType type, int t, int n, int k, int ldy) {
+    bool init(hipDataType type, int t, int n, int k, int ldy, int ldb = 0) {
         const hipblasOperation_t trans_a = HIPBLAS_OP_T;
         const hipblasOperation_t trans_b = HIPBLAS_OP_N;
         if (hipblasLtMatmulDescCreate(&op, HIPBLAS_COMPUTE_32F, HIP_R_32F) != HIPBLAS_STATUS_SUCCESS ||
@@ -136,7 +136,7 @@ struct HipLtDescriptors {
             hipblasLtMatmulDescSetAttribute(op, HIPBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b)) !=
                 HIPBLAS_STATUS_SUCCESS ||
             hipblasLtMatrixLayoutCreate(&a, type, k, n, k) != HIPBLAS_STATUS_SUCCESS ||
-            hipblasLtMatrixLayoutCreate(&b, type, k, t, k) != HIPBLAS_STATUS_SUCCESS ||
+            hipblasLtMatrixLayoutCreate(&b, type, k, t, ldb > k ? ldb : k) != HIPBLAS_STATUS_SUCCESS ||
             hipblasLtMatrixLayoutCreate(&c, HIP_R_32F, n, t, ldy) != HIPBLAS_STATUS_SUCCESS) {
             return false;
         }
@@ -243,9 +243,20 @@ HipLtCachedAlgo resolve_hipblaslt_algo(HipLtState& state, strata::prefill::hipbl
     return state.cache.emplace(key, resolved).first->second;
 }
 
+// the hipBLASLt solution index the tuning table gives this call (-1: none / unsupported)
+int hipblaslt_solution_for(void* opaque_state, strata::prefill::hipblaslt::InputType type, int64_t t, int64_t n, int64_t k,
+                           int64_t ldy, float beta) {
+    auto* state = static_cast<HipLtState*>(opaque_state);
+    if (!state || t <= 0 || n <= 0 || k <= 0 || t > INT_MAX || n > INT_MAX || k > INT_MAX || ldy > INT_MAX || ldy < n) return -1;
+    const auto resolved = resolve_hipblaslt_algo(*state, type, (int) t, (int) n, (int) k, (int) ldy, beta);
+    if (!resolved.supported || resolved.workspace_bytes > state->workspace_bytes) return -1;
+    auto algo = resolved.algo;
+    return hipblaslt_ext::getIndexFromAlgo(algo);
+}
+
 bool try_hipblaslt(void* opaque_state, strata::prefill::hipblaslt::InputType type, const uint16_t* x,
                    const uint16_t* w, float* y, int64_t t, int64_t n, int64_t k, int64_t ldy, float beta,
-                   void* stream) {
+                   void* stream, int64_t ldx = 0) {
     auto* state = static_cast<HipLtState*>(opaque_state);
     if (!state || t <= 0 || n <= 0 || k <= 0 || t > INT_MAX || n > INT_MAX || k > INT_MAX || ldy > INT_MAX ||
         ldy < n) {
@@ -269,7 +280,7 @@ bool try_hipblaslt(void* opaque_state, strata::prefill::hipblaslt::InputType typ
 
     HipLtDescriptors desc;
     const hipDataType input_type = type == strata::prefill::hipblaslt::InputType::bf16 ? HIP_R_16BF : HIP_R_16F;
-    if (!desc.init(input_type, (int) t, (int) n, (int) k, (int) ldy)) return false;
+    if (!desc.init(input_type, (int) t, (int) n, (int) k, (int) ldy, (int) ldx)) return false;
     const float alpha = 1.0f;
     const hipblasStatus_t status = hipblasLtMatmul(state->handle, desc.op, &alpha, w, desc.a, x, desc.b, &beta, y,
                                                    desc.c, y, desc.c, &resolved.algo, state->workspace,
@@ -376,16 +387,29 @@ bool Gemm::init(void* stream, int64_t scratch_elems, std::string& err) {
     return true;
 }
 
+bool Gemm::bf16_hcd_exact(const uint16_t* X, int64_t ldx, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K) {
+#if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE)
+    if (!hipblaslt_state_ || N != 320 || K != 10240 || T < 1) return false;
+    const int id = hipblaslt_solution_for(hipblaslt_state_, strata::prefill::hipblaslt::InputType::bf16, T, N, K, N, 0.0f);
+    if (id != 1176 && id != 1177) return false;
+    return strata_pf_hcdown_exact_bf16(X, ldx, W, Y, T, N, K, stream_);
+#else
+    (void) X; (void) ldx; (void) W; (void) Y; (void) T; (void) N; (void) K;
+    return false;
+#endif
+}
+
 void Gemm::bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy,
-                float beta) {
+                float beta, int64_t ldx) {
     if (T <= 0 || N <= 0) return;
     if (ldy <= 0) ldy = N;
+    if (ldx <= K) ldx = 0;
 #ifdef STRATA_USE_HIP
     // #313 (opt-in STRATA_WMMA_GEMM=1; STRATA_WMMA_BF16=0 excludes bf16): RDNA3 / RDNA3.5 WMMA dense GEMM, before
     // hipBLASLt; it returns false for shapes it does not take (and off gfx11) and this falls through
     static const bool wmma_on = [] { const char* v = std::getenv("STRATA_WMMA_GEMM"); return v && v[0] != 0 && v[0] != '0'; }();
     static const bool bf16_on = [] { const char* v = std::getenv("STRATA_WMMA_BF16"); return !v || (v[0] != 0 && v[0] != '0'); }();
-    if (wmma_on && bf16_on && T >= 16 && (beta == 0.0f || beta == 1.0f) &&
+    if (wmma_on && bf16_on && ldx == 0 && T >= 16 && (beta == 0.0f || beta == 1.0f) &&
         strata_wmma_gemm_bf16(X, W, Y, T, N, K, ldy, beta, stream_)) {
         return;
     }
@@ -393,14 +417,14 @@ void Gemm::bf16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64
     const float alpha = 1.0f;
 #if defined(__HIPCC__) && defined(STRATA_HIPBLASLT_AVAILABLE)
     if (try_hipblaslt(hipblaslt_state_, strata::prefill::hipblaslt::InputType::bf16, X, W, Y, T, N, K, ldy,
-                      beta, stream_)) {
+                      beta, stream_, ldx)) {
         STRATA_ABSORB_HIPBLAS_STICKY("hipBLASLt bf16");
         return;
     }
 #endif
     // Column-major view: Y^T[N, T] = W[N, K] (stored K x N col-major, transposed) . X^T[K, T].
     ck(cublasGemmEx((cublasHandle_t) handle_, CUBLAS_OP_T, CUBLAS_OP_N, (int) N, (int) T, (int) K, &alpha, W,
-                    CUDA_R_16BF, (int) K, X, CUDA_R_16BF, (int) K, &beta, Y, CUDA_R_32F, (int) ldy,
+                    CUDA_R_16BF, (int) K, X, CUDA_R_16BF, (int) (ldx ? ldx : K), &beta, Y, CUDA_R_32F, (int) ldy,
                     CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT),
        "cublasGemmEx");
     STRATA_ABSORB_HIPBLAS_STICKY("cublasGemmEx");

src/prefill/prefill.cpp

@@ -191,6 +191,14 @@ inline bool pf_hcdown() {
     return v;
 }
 constexpr int64_t XN_PAD = 64;
+// S (opt-in STRATA_HCD_EXACT=1, on chunks of STRATA_PF_SWITCH_MIN_T+ tokens): the HC down projection by a WMMA kernel with
+// hipBLASLt's own k order (bitwise: strata_pf_hcdown_exact_bf16), xn16 written with token stride 10240 + 64 (the inject
+// projection stays on hipBLASLt, reading that stride)
+inline bool hcd_exact() {
+    static const bool v = [] { const char* e = std::getenv("STRATA_HCD_EXACT"); return e && e[0] == '1'; }();
+    return v;
+}
+inline bool hc_pad() { return pf_hcdown() || hcd_exact(); }
 inline bool cvec_fuse() {
     static const bool v = [] { const char* e = std::getenv("STRATA_CVEC_FUSE"); return e && e[0] == '1'; }();
     return v;
@@ -755,7 +763,7 @@ bool Prefill::carve(size_t T, void* alloc) {
     m.emb = o.take<float>(T * N, ok); m.R = o.take<float>(T * D, ok);
     m.xn = gr_unfused() ? o.take<float>(T * D, ok) : nullptr;   // F-1: not needed (gr_mix_r reads R)
     m.grs = o.take<float>(T * HC, ok);
-    m.xn16 = o.take<uint16_t>(T * (D + (pf_hcdown() ? XN_PAD : 0)), ok); m.lo = o.take<float>(T * LR, ok); m.lo16 = o.take<uint16_t>(T * LR, ok);
+    m.xn16 = o.take<uint16_t>(T * (D + (hc_pad() ? XN_PAD : 0)), ok); m.lo = o.take<float>(T * LR, ok); m.lo16 = o.take<uint16_t>(T * LR, ok);
     m.gated = o.take<float>(T * D, ok); m.inj = o.take<float>(T * HC, ok);
     m.mixed = o.take<float>(T * N, ok); m.mixed_bf = o.take<uint16_t>(T * N, ok);
     m.mixed_h = o.take<uint16_t>(T * N, ok); m.bo = o.take<float>(T * N, ok);
@@ -1020,7 +1028,7 @@ uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::Session
     o.take<uint16_t>((size_t) GEMM_SCRATCH, ok);
     o.take<uint8_t>(GEMM_WS, ok);
     auto f = [&](size_t n) { o.take<float>(n, ok); };
-    f(T * N); f(T * D); f(T * D); o.take<uint16_t>(T * (D + (pf_hcdown() ? XN_PAD : 0)), ok); f(T * LR);
+    f(T * N); f(T * D); f(T * D); o.take<uint16_t>(T * (D + (hc_pad() ? XN_PAD : 0)), ok); f(T * LR);
     o.take<uint16_t>(T * LR, ok);
     f(T * D); f(T * HC); f(T * N); o.take<uint16_t>(T * N, ok); o.take<uint16_t>(T * N, ok); f(T * N);
     if (bf16x2_hc()) { o.take<uint16_t>(T * D, ok); o.take<uint16_t>(T * LR, ok); }
@@ -1063,9 +1071,9 @@ bool native_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y
     return true;
 }
 bool bf16_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y, int64_t T, const std::string& name,
-               std::string& err, int64_t ldy = 0, const uint16_t* X_lo = nullptr) {
+               std::string& err, int64_t ldy = 0, const uint16_t* X_lo = nullptr, int64_t ldx = 0) {
     if (w->kind != core::WeightKind::Bf16InF32 || !w->data) { err = "prefill: " + name + " is not a resident BF16 tensor"; return false; }
-    gm.bf16(X, (const uint16_t*) w->data, Y, T, w->ne1 > 0 ? w->ne1 : 1, w->ne0, ldy);
+    gm.bf16(X, (const uint16_t*) w->data, Y, T, w->ne1 > 0 ? w->ne1 : 1, w->ne0, ldy, 0.0f, ldx);
     if (X_lo) gm.bf16(X_lo, (const uint16_t*) w->data, Y, T, w->ne1 > 0 ? w->ne1 : 1, w->ne0, ldy, 1.0f);
     return true;
 }
@@ -1457,13 +1465,17 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                 if (!wn || !wd || !wu || !wi) return false;
                 pt.mark(kPfHc, cs);
                 // the xn16 token stride of this chunk (the writers of both halves and of the fused write-backs use it)
-                const int64_t ldx = pf_hcdown() && T >= std::max<int64_t>(pf_switch_min_t(), 64) && !gr_unfused() &&
+                const int64_t ldx = hc_pad() && T >= std::max<int64_t>(pf_switch_min_t(), 64) && !gr_unfused() &&
                                             !m.xn16_lo ? D + XN_PAD : D;
                 if (gr_unfused()) gr_norm(m.R, (const float*) wn->data, EPS, m.xn, m.xn16, T, m.cs, m.xn16_lo);
                 else if (!normed) gr_norm_rs(m.R, (const float*) wn->data, EPS, m.grs, m.xn16, T, m.cs, m.xn16_lo, ldx);
                 normed = false;
                 bool hcd = false;
-                if (ldx != D) {
+                bool hdown = false;
+                if (ldx != D && !pf_hcdown()) {   // STRATA_HCD_EXACT: the exact-order down kernel when hipBLASLt would take 1176 / 1177
+                    hdown = wd->kind == core::WeightKind::Bf16InF32 && wd->data && wd->ne0 == D && wd->ne1 == LR &&
+                            m.gemm.bf16_hcd_exact(m.xn16, ldx, (const uint16_t*) wd->data, m.lo, T, LR, D);
+                } else if (ldx != D) {
                     if (wd->kind != core::WeightKind::Bf16InF32 || !wd->data || wd->ne0 != D || wd->ne1 != LR ||
                         wi->kind != core::WeightKind::Bf16InF32 || !wi->data || wi->ne0 != D || wi->ne1 != HC ||
                         !strata_pf_hcdown_bf16(m.xn16, ldx, (const uint16_t*) wd->data, (const uint16_t*) wi->data, LR,
@@ -1473,7 +1485,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                     }
                     hcd = true;
                 }
-                if (!hcd && !bf16_proj(m.gemm, wd, m.xn16, m.lo, T, sd, err, 0, m.xn16_lo)) return false;
+                if (!hcd && !hdown && !bf16_proj(m.gemm, wd, m.xn16, m.lo, T, sd, err, 0, m.xn16_lo, ldx != D ? ldx : 0)) return false;
                 gr_silu(m.lo, m.lo16, T, m.cs, m.lo16_lo);
                 bool upmixed = false;
                 if (hc_upmix() && T >= pf_switch_min_t() && !gr_unfused() && !m.lo16_lo && !m.mixed_bf_lo &&
@@ -1503,7 +1515,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                     }
                 }
                 if (!upmixed && !bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err, 0, m.lo16_lo)) return false;
-                if (!hcd && !bf16_proj(m.gemm, wi, m.xn16, m.inj, T, si, err, 0, m.xn16_lo)) return false;
+                if (!hcd && !bf16_proj(m.gemm, wi, m.xn16, m.inj, T, si, err, 0, m.xn16_lo, ldx != D ? ldx : 0)) return false;
                 if (upmixed) {
                 } else if (gr_unfused()) {
                     gr_mix(m.xn, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, m.mixed_bf_lo);

src/prefill/wmma_gemm.cu

@@ -554,6 +554,82 @@ __global__ void __launch_bounds__(256) kernel_hcdown(const uint16_t* __restrict_
         }
 #endif
 }
+
+// S23/S (STRATA_HCD_EXACT=1): the HC down projection (N 320, K 10240, BF16 in, FP32 out) bitwise equal to hipBLASLt
+// solution 1176 / 1177 (MT32x96 / MT96x96, StaggerU 32, stride 256 B): every output is the k tiles of 32 accumulated in
+// order, 16 per WMMA, starting at k tile 4 * (token tile of 96 % 32) and wrapping (hcd_bit / hcd_perm probes: 0 of
+// 5.2 M outputs differ at T 16384).  96 x 160 blocks (two per token tile), 6 waves of 32 x 80.
+typedef uint32_t hx_u4 __attribute__((ext_vector_type(4)));
+constexpr int HX_N = 320, HX_BN = 160, HX_BM = 96, HX_BK = 64, HX_LDK = HX_BK + 8, HX_NT = 192, HX_K = 10240;
+__device__ __forceinline__ hb16 hx_frag(const uint16_t* p) {
+    const hx_u4 a = *reinterpret_cast<const hx_u4*>(p), b = *reinterpret_cast<const hx_u4*>(p + 8);
+    const uint32_t w[8] = {a.x, a.y, a.z, a.w, b.x, b.y, b.z, b.w};
+    return __builtin_bit_cast(hb16, w);
+}
+__global__ void __launch_bounds__(HX_NT) kernel_hcd_exact(const uint16_t* __restrict__ X, int ldx,
+                                                          const uint16_t* __restrict__ W, float* __restrict__ Y, int M) {
+#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1150__) || defined(__gfx1151__)
+    constexpr int K = HX_K, BN = HX_BN, BMT = HX_BM, BK = HX_BK, NT = HX_NT;
+    __shared__ __align__(16) uint16_t sA[BMT][HX_LDK];
+    __shared__ __align__(16) uint16_t sB[BN][HX_LDK];
+    const int tid = threadIdx.x, lane = tid & 31, wave = tid >> 5, l16 = lane & 15, hi = lane >> 4;
+    const int wm = wave % 3, wn = wave / 3, nbase = BN * (blockIdx.x & 1);
+    const int tt = blockIdx.x >> 1, m0 = tt * BMT;
+    constexpr int nslab = K / BK;
+    const int rot = (2 * (tt & 31)) % nslab;   // 4 k tiles of 32 = 2 slabs of 64
+    hx_u4 ra0, ra1, ra2, ra3, rb0, rb1, rb2, rb3, rb4, rb5, rb6;
+#define HX_LDA(v, i) { const int idx = tid + NT * (i), r = idx >> 3, q = idx & 7; v = *reinterpret_cast<const hx_u4*>(X + (size_t) min(m0 + r, M - 1) * ldx + k0 + 8 * q); }
+#define HX_LDB(v, i) { const int idx = min(tid + NT * (i), BN * 8 - 1), r = idx >> 3, q = idx & 7; v = *reinterpret_cast<const hx_u4*>(W + (size_t) (nbase + r) * K + k0 + 8 * q); }
+#define HX_STA(v, i) { const int idx = tid + NT * (i); *reinterpret_cast<hx_u4*>(&sA[idx >> 3][8 * (idx & 7)]) = v; }
+#define HX_STB(v, i) { const int idx = tid + NT * (i); if (idx < BN * 8) *reinterpret_cast<hx_u4*>(&sB[idx >> 3][8 * (idx & 7)]) = v; }
+#define HX_LOAD(k0v) { const int k0 = (k0v); HX_LDA(ra0, 0) HX_LDA(ra1, 1) HX_LDA(ra2, 2) HX_LDA(ra3, 3) HX_LDB(rb0, 0) HX_LDB(rb1, 1) HX_LDB(rb2, 2) HX_LDB(rb3, 3) HX_LDB(rb4, 4) HX_LDB(rb5, 5) HX_LDB(rb6, 6) }
+#define HX_STORE() { HX_STA(ra0, 0) HX_STA(ra1, 1) HX_STA(ra2, 2) HX_STA(ra3, 3) HX_STB(rb0, 0) HX_STB(rb1, 1) HX_STB(rb2, 2) HX_STB(rb3, 3) HX_STB(rb4, 4) HX_STB(rb5, 5) HX_STB(rb6, 6) }
+    f8 acc[2][5];
+#pragma unroll
+    for (int i = 0; i < 2; ++i)
+#pragma unroll
+        for (int j = 0; j < 5; ++j) acc[i][j] = f8{0, 0, 0, 0, 0, 0, 0, 0};
+    HX_LOAD(rot * BK);
+    HX_STORE();
+    __syncthreads();
+    const int ar = 32 * wm + l16, br = 80 * wn + l16;
+    for (int s = 0; s < nslab; ++s) {
+        const bool more = s + 1 < nslab;
+        if (more) HX_LOAD(((s + 1 + rot) % nslab) * BK);
+#pragma unroll
+        for (int ks = 0; ks < BK; ks += 16) {
+            const hb16 a0 = hx_frag(&sA[ar][ks]), a1 = hx_frag(&sA[ar + 16][ks]);
+            const hb16 b0 = hx_frag(&sB[br][ks]), b1 = hx_frag(&sB[br + 16][ks]), b2 = hx_frag(&sB[br + 32][ks]),
+                       b3 = hx_frag(&sB[br + 48][ks]), b4 = hx_frag(&sB[br + 64][ks]);
+#define HX_W2(j, bj) acc[0][j] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a0, bj, acc[0][j]); acc[1][j] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a1, bj, acc[1][j]);
+            HX_W2(0, b0) HX_W2(1, b1) HX_W2(2, b2) HX_W2(3, b3) HX_W2(4, b4)
+#undef HX_W2
+        }
+        if (more) {
+            __syncthreads();
+            HX_STORE();
+            __syncthreads();
+        }
+    }
+#undef HX_LDA
+#undef HX_LDB
+#undef HX_STA
+#undef HX_STB
+#undef HX_LOAD
+#undef HX_STORE
+#pragma unroll
+    for (int i = 0; i < 2; ++i)
+#pragma unroll
+        for (int j = 0; j < 5; ++j) {
+            const int n = nbase + 80 * wn + 16 * j + l16;
+#pragma unroll
+            for (int e = 0; e < 8; ++e) {
+                const int m = m0 + 32 * wm + 16 * i + 2 * e + hi;
+                if (m < M) Y[(size_t) m * HX_N + n] = acc[i][j][e];
+            }
+        }
+#endif
+}
 }  // namespace pfg
 
 bool strata_pf_hcdown_bf16(const uint16_t* X, int64_t ldx, const uint16_t* Wd, const uint16_t* Wi, int64_t nd,
@@ -574,6 +650,23 @@ bool strata_pf_hcdown_bf16(const uint16_t* X, int64_t ldx, const uint16_t* Wd, c
     return hipGetLastError() == hipSuccess;
 }
 
+bool strata_pf_hcdown_exact_bf16(const uint16_t* X, int64_t ldx, const uint16_t* W, float* Y, int64_t T, int64_t N,
+                                 int64_t K, void* stream) {
+    if (!X || !W || !Y || T < 1 || N != pfg::HX_N || K != pfg::HX_K || ldx < K || ldx % 8 != 0 || ldx > (1LL << 30) ||
+        T > (1LL << 30))
+        return false;
+    static const bool gfx11 = [] {
+        int dev = 0;
+        hipDeviceProp_t prop{};
+        if (hipGetDevice(&dev) != hipSuccess || hipGetDeviceProperties(&prop, dev) != hipSuccess) return false;
+        return std::strncmp(prop.gcnArchName, "gfx11", 5) == 0;
+    }();
+    if (!gfx11) return false;
+    const unsigned grid = (unsigned) (2 * ((T + pfg::HX_BM - 1) / pfg::HX_BM));
+    pfg::kernel_hcd_exact<<<grid, pfg::HX_NT, 0, static_cast<hipStream_t>(stream)>>>(X, (int) ldx, W, Y, (int) T);
+    return hipGetLastError() == hipSuccess;
+}
+
 bool strata_pf_gemm_f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t T, int64_t N, int64_t K, int64_t ldy,
                         float beta, void* stream) {
     return strata_pf_gemm_f16_ld(X, K, W, K, Y, T, N, K, ldy, beta, stream);
@@ -627,6 +720,9 @@ bool strata_pf_hcdown_bf16(const uint16_t*, int64_t, const uint16_t*, const uint
                            int64_t, int64_t, void*) {
     return false;
 }
+bool strata_pf_hcdown_exact_bf16(const uint16_t*, int64_t, const uint16_t*, float*, int64_t, int64_t, int64_t, void*) {
+    return false;
+}
 bool strata_wmma_gemm_f16(const uint16_t*, const uint16_t*, float*,
                           int64_t, int64_t, int64_t, int64_t, float,
                           void*) {

src/prefill/wmma_gemm.h

@@ -28,6 +28,11 @@ bool strata_pf_gemm_f16(const uint16_t* X, const uint16_t* W, float* Y, int64_t
 /// K); lo T x nd and inj T x ni, FP32.  K a multiple of 64, T >= 64.  False (nothing launched) otherwise / off gfx11.
 bool strata_pf_hcdown_bf16(const uint16_t* X, int64_t ldx, const uint16_t* Wd, const uint16_t* Wi, int64_t nd,
                            int64_t ni, float* lo, float* inj, int64_t T, int64_t K, void* stream = nullptr);
+/// S (STRATA_HCD_EXACT=1): Y[t, n] = sum_k X[t, k] W[n, k] for N 320, K 10240 (BF16, FP32 out, X token stride ldx >= K, a
+/// multiple of 8), bitwise equal to hipBLASLt solution 1176 / 1177 for that shape (its k order, StaggerU included).
+/// False (nothing launched) for any other shape / off gfx11.
+bool strata_pf_hcdown_exact_bf16(const uint16_t* X, int64_t ldx, const uint16_t* W, float* Y, int64_t T, int64_t N,
+                                 int64_t K, void* stream = nullptr);
 bool strata_pf_gemm_f16_ld(const uint16_t* X, int64_t ldx, const uint16_t* W, int64_t ldw, float* Y, int64_t T,
                            int64_t N, int64_t K, int64_t ldy, float beta, void* stream = nullptr);