[Engine] GP100 prompt path: a fully resident layer's experts weight-stationary, one launch per product (STRATA_MMQ_WS=1, opt-in)

CMakeLists.txt

@@ -1718,6 +1718,12 @@ if(STRATA_NATIVE_EXPERTS)
     add_library(strata_mmq STATIC src/prefill/moe_mmq.cu src/prefill/moe_fused.cu src/prefill/moe_fused_iq.cu
                 src/prefill/mmq_nvfp4_w4a8.cu
                 src/prefill/ggml_cuda_host.cu ${_strata_quantize_cu})
+    if(STRATA_EXPERIMENTAL_SM60 AND NOT STRATA_HIP_GFX906)
+      # GP100 (sm_60): a resident layer's experts weight-stationary, one launch per product (moe_ws.hpp; opt-in at
+      # run time, STRATA_MMQ_WS=1, and only where every visible GPU is cc 6.0)
+      target_sources(strata_mmq PRIVATE src/prefill/moe_ws.cu src/prefill/moe_ws_gu_iq2.cu src/prefill/moe_ws_gu_iq3.cu
+                     src/prefill/moe_ws_down.cu)
+    endif()
     # q8_0: the drafter's matrices (E-9); STRATA_ORCA_Q4KS_MMQ adds q5_0 (#296, in _strata_mmq_types), and
     # STRATA_MMQ_KQUANTS q4_k q5_k q5_1 q6_k, the experts of Unsloth's UD-Q4/Q5/Q6_K_XL (docs/QUANTS.md; the dispatch
     # follows eddoursul/Strata 7706c1b).  MODEL-UPDATE Phase 1: on CUDA they are a library of their own, strata_kq
@@ -1850,6 +1856,13 @@ if(STRATA_NATIVE_EXPERTS)
     endif()
     target_link_libraries(strata_prefill PUBLIC strata_mmq)
     target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_MMQ=1 STRATA_PREFILL_FUSED=1)
+    if(STRATA_EXPERIMENTAL_SM60 AND NOT STRATA_HIP_GFX906)
+      # the weight-stationary experts (moe_ws.hpp) are in strata_mmq: the prompt path may take them (STRATA_MMQ_WS=1)
+      target_compile_definitions(strata_prefill PUBLIC STRATA_PREFILL_WS=1)
+      # moe_ws_down_bench: the WS down with and without STRATA_WS_FAST on one GP100 die, byte for byte and timed
+      add_executable(moe_ws_down_bench EXCLUDE_FROM_ALL src/prefill/moe_ws_down_bench.cu)
+      target_link_libraries(moe_ws_down_bench PRIVATE strata_mmq strata_artifact ggml-base ${_strata_gpu_runtime_target})
+    endif()
     if(STRATA_BUILD_TESTS)
       # #136: the fused int8 experts against a double-precision reference and the MMQ path (exit 77 below sm_80)
       add_executable(prefill_fused_moe_test tests/cuda/prefill_fused_moe_test.cpp)
@@ -1867,6 +1880,13 @@ if(STRATA_NATIVE_EXPERTS)
       target_link_libraries(prefill_mmq_scatter_test PRIVATE strata_mmq ggml-base ${_strata_gpu_runtime_target})
       add_test(NAME prefill_mmq_scatter_test COMMAND prefill_mmq_scatter_test)
       set_tests_properties(prefill_mmq_scatter_test PROPERTIES TIMEOUT 120)
+      if(STRATA_EXPERIMENTAL_SM60)
+        # mmq::swiglu_quant: swiglu + the down product's q8_1 rows in one kernel against the two, byte for byte
+        add_executable(prefill_swiglu_quant_test tests/cuda/prefill_swiglu_quant_test.cpp)
+        target_link_libraries(prefill_swiglu_quant_test PRIVATE strata_mmq ggml-base ${_strata_gpu_runtime_target})
+        add_test(NAME prefill_swiglu_quant_test COMMAND prefill_swiglu_quant_test)
+        set_tests_properties(prefill_swiglu_quant_test PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 120)
+      endif()
     endif()
     if(STRATA_BUILD_TESTS AND _strata_mmq_tab_on)
       # STRATA_MMQ_DIRECT: the pointer-table MMQ bit for bit against the gathered MMQ (exit 77 below sm_80)

include/strata/prefill/moe_mmq.hpp

@@ -122,6 +122,11 @@ void gather_strata_q2(const uint8_t* blob, void* gu_dst, void* d_dst, void* stre
 /// h[r, k] = silu(gate) * up of GU rows [2 n_ff wide]: interleaved (gate 2k, up 2k+1: the Strata pack) or split
 /// (gate k, up n_ff + k: GGUF).  FP32 out (the down product's quantizer reads floats).
 void swiglu(const float* gu, float* h, int64_t rows, int64_t n_ff, bool interleaved, void* stream);
+/// swiglu (n_ff 640), then quantize() of its rows for a down product of `ggml_type`, in one pass: H never stored,
+/// the same bytes in hq.  Only where swiglu_quant_ok (a q8_1 D4 down type; gu 16-byte aligned): CUDA, the
+/// STRATA_EXPERIMENTAL_SM60 build.
+bool swiglu_quant_ok(int ggml_type);
+void swiglu_quant(const float* gu, void* hq, int64_t rows, bool interleaved, void* stream);
 /// swiglu of an NVFP4 group's gate/up rows (split halves) with scale_gu_rows' scales applied as they are read.
 void swiglu_scaled(const float* gu, float* h, int64_t rows, int64_t n_ff, const int32_t* bounds, int n,
                    const float* tails, int64_t row0, void* stream);

include/strata/prefill/moe_ws.hpp

@@ -0,0 +1,70 @@
+// include/strata/prefill/moe_ws.hpp - GP100 (sm_60, the STRATA_EXPERIMENTAL_SM60 build): the prompt path's expert
+// products of a layer whose experts are all resident in VRAM, read in place there and weight-stationary: ONE launch
+// for the layer's gate/up and ONE for its down instead of a gate/up, swiglu and down MMQ launch per 16-expert group.
+// A CTA holds 128 (or 256) weight rows of one expert and runs them against up to 64 of the expert's activation rows, a
+// thread per weight row: each weight block is decoded once per pass instead of once per MMQ tile, and an expert's
+// rows are its own count, not the group's largest.  Every product keeps MMQ's int8 values, scales and per-32-value
+// epilogue in its K order (the down K loop stops at 640), so GU, H and Dm are the per-group MMQ's bits.  Swift
+// Flash's shapes and formats only: gate/up [1280 x 2560] IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS or IQ3_S, down [2560 x 640]
+// Q2_0 or IQ4_NL.
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+namespace strata::prefill::ws {
+
+/// The products of a layer: 1 with STRATA_MMQ_WS=1 (opt-in) gate/up and down, else 0 (the per-group MMQ).  0 also where
+/// the build has no kernels (HIP, no ggml sources, not the STRATA_EXPERIMENTAL_SM60 build) or a visible GPU is not
+/// cc 6.0 (the kernels were tuned and checked on GP100 only).
+int mode();
+/// mode() != 0 and both formats have kernels.
+bool supported(int gu_type, int d_type);
+/// STRATA_WS_FAST=1 (opt-in, read once): the IQ4_NL down's 32-token tier (u >= 20 rows per routed expert) runs a
+/// kernel that issues fewer instructions per pass - a switch into the unrolled token pairs, dy transposed - with every
+/// token's operations and their order unchanged, so Dm is the same bits.  0 (unset): the kernels as before.
+int fast();
+/// Benches and tests: 0 / 1 overrides fast(), -1 goes back to the environment.
+void set_fast(int v);
+/// The weight-stationary product of `ggml_type` beats the per-group MMQ at u rows per routed expert (the measured
+/// ratios, moe_ws.cu; u is clamped to the measured 2.5 .. 160).
+bool faster(int ggml_type, double u);
+/// Every gate/up and down matrix starts this aligned (the IQ4_NL down reads 8-byte pieces; the expert cache's slots are
+/// 256-byte aligned and the Swift gate/up sizes multiples of 8).
+constexpr uint64_t kAlign = 8;
+
+/// The CTAs' work: rows [row0, row0 + rows) of Xq / GU / Hq / Dm, all of the expert at position j of the layer's order
+/// (MMQ's: its xoff entry), at most kChunkRows.  A CTA takes one chunk and one 128- or 256-row piece of the weights.
+struct alignas(16) Chunk {
+    int32_t j, row0, rows, pad;
+};
+constexpr int kChunkRows = 64;
+/// The most chunks a layer of `rows` routed rows over at most `n_expert` experts makes.
+inline size_t max_chunks(int64_t rows, int64_t n_expert) { return (size_t) (rows / kChunkRows + n_expert); }
+/// The chunks of the n experts of `order` (expert e has cnt[e] rows from bounds[j], j its position), the largest
+/// expert first (the long CTAs start first); returns their number (<= max_chunks).
+size_t chunks(const int32_t* order, size_t n, const int32_t* cnt, const int32_t* bounds, Chunk* out);
+
+/// One layer: its n routed experts' matrices at w + xoff[j] (gate/up, rows 0-639 gate then up) and w + xoff[j] +
+/// down_off (down), xoff on the device in the layer's order (Product::xoff's table, all n entries); `rows` = T*K.
+struct Layer {
+    int gu_type = -1, d_type = -1;
+    const void* w = nullptr;
+    const int64_t* xoff = nullptr;
+    int64_t down_off = 0;
+    const Chunk* chunks = nullptr;   ///< on the device
+    int64_t nchunks = 0;
+    int64_t rows = 0;
+    int n = 0;
+    const void* xq = nullptr;        ///< the gate/up's q8_1 rows (mmq::quantize_scatter's, rows of them)
+    float* gu = nullptr;             ///< [rows][1280], the rows MMQ writes through identity ids
+    const void* hq = nullptr;        ///< H's q8_1 rows: ONE mmq::swiglu_quant over all `rows` rows of gu
+    float* dm = nullptr;             ///< [rows][2560]
+};
+/// The layer's gate/up into gu, one launch; false (nothing launched) where the per-group MMQ is to run instead (mode 0,
+/// or not faster at this u) - the same GU rows either way.
+bool gate_up(const Layer& L, void* stream);
+/// The layer's down from hq into dm, one launch.
+void down(const Layer& L, void* stream);
+
+}  // namespace strata::prefill::ws

src/prefill/moe_mmq.cu

@@ -144,6 +144,54 @@ __global__ void swiglu_kernel(const float* __restrict__ gu, float* __restrict__
     h[i] = g / (1.0f + __expf(-g)) * u;
 }
 
+#if !defined(__HIPCC__) && defined(STRATA_EXPERIMENTAL_SM60)
+// swiglu_kernel (n_ff 640), then quantize_mmq_q8_1<D4>'s rows of its output, in one pass: a block per row (128
+// threads, 4 values each, blockIdx.y the row's 512-value half), H never stored.  Each value is swiglu_kernel's
+// expression and each 32-value block quantize_mmq_q8_1's (the same unit and flags: the same bits).
+template <bool IL>
+__global__ void __launch_bounds__(128) swiglu_quant_kernel(const float* __restrict__ gu, block_q8_1_mmq* __restrict__ y,
+                                                           int64_t rows) {
+    constexpr int64_t n_ff = 640, ne0 = 1024;   // the padded row: pad512(640)
+    const int64_t i0 = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * 4;
+    if (i0 >= ne0) return;
+    const int64_t r = blockIdx.x;
+    float4 xi = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    if (i0 < n_ff) {
+        const float* row = gu + r * 2 * n_ff;
+        float4 g4, u4;
+        if constexpr (IL) {   // gate 2k, up 2k + 1
+            const float4 a = *reinterpret_cast<const float4*>(row + 2 * i0), b = *reinterpret_cast<const float4*>(row + 2 * i0 + 4);
+            g4 = make_float4(a.x, a.z, b.x, b.z);
+            u4 = make_float4(a.y, a.w, b.y, b.w);
+        } else {
+            g4 = *reinterpret_cast<const float4*>(row + i0);
+            u4 = *reinterpret_cast<const float4*>(row + n_ff + i0);
+        }
+        xi.x = g4.x / (1.0f + __expf(-g4.x)) * u4.x;
+        xi.y = g4.y / (1.0f + __expf(-g4.y)) * u4.y;
+        xi.z = g4.z / (1.0f + __expf(-g4.z)) * u4.z;
+        xi.w = g4.w / (1.0f + __expf(-g4.w)) * u4.w;
+    }
+    float amax = fabsf(xi.x);
+    amax = fmaxf(amax, fabsf(xi.y));
+    amax = fmaxf(amax, fabsf(xi.z));
+    amax = fmaxf(amax, fabsf(xi.w));
+#pragma unroll
+    for (int offset = 32 / 8; offset > 0; offset >>= 1) amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
+    const float d_inv = 127.0f / amax;
+    char4 c;
+    c.x = roundf(xi.x * d_inv);
+    c.y = roundf(xi.y * d_inv);
+    c.z = roundf(xi.z * d_inv);
+    c.w = roundf(xi.w * d_inv);
+    const float d = 1.0f / d_inv;
+    const int64_t k_block = i0 / QK8_1_MMQ, iqs = i0 % QK8_1_MMQ;
+    block_q8_1_mmq* yb = y + k_block * rows + r;
+    reinterpret_cast<char4*>(yb->qs)[iqs / 4] = c;
+    if (iqs % 32 == 0) yb->d4[iqs / 32] = d;
+}
+#endif
+
 // swiglu_kernel on the MMQ gate/up output with each expert's gate and up scales applied as it is read: the products
 // scale_gu_rows stored, without the pass that stored them
 __global__ void swiglu_scaled_kernel(const float* __restrict__ gu, float* __restrict__ h, int64_t rows, int64_t n_ff,
@@ -665,6 +713,28 @@ void swiglu(const float* gu, float* h, int64_t rows, int64_t n_ff, bool interlea
     ck(cudaGetLastError(), "swiglu");
 }
 
+bool swiglu_quant_ok(int t) {
+#if !defined(__HIPCC__) && defined(STRATA_EXPERIMENTAL_SM60)
+    // the q8_1 D4 rows quantize() writes for a down product's H (K 640)
+    return supported(t) && !fp4_activations(t) && mmq_get_q8_1_ds_layout((ggml_type) t) == MMQ_Q8_1_DS_LAYOUT_D4;
+#else
+    (void) t;
+    return false;   // HIP: not checked against the HIP build's quantizer; outside the SM60 build: not built
+#endif
+}
+
+void swiglu_quant(const float* gu, void* hq, int64_t rows, bool interleaved, void* stream) {
+#if !defined(__HIPCC__) && defined(STRATA_EXPERIMENTAL_SM60)
+    if (rows <= 0) return;
+    const dim3 grid((unsigned) rows, 2);
+    if (interleaved) swiglu_quant_kernel<true><<<grid, 128, 0, (cudaStream_t) stream>>>(gu, (block_q8_1_mmq*) hq, rows);
+    else swiglu_quant_kernel<false><<<grid, 128, 0, (cudaStream_t) stream>>>(gu, (block_q8_1_mmq*) hq, rows);
+    ck(cudaGetLastError(), "swiglu_quant");
+#else
+    (void) gu; (void) hq; (void) rows; (void) interleaved; (void) stream;
+#endif
+}
+
 void iota(int32_t* dst, int64_t n, void* stream) {
     if (n <= 0) return;
     iota_kernel<<<blocks(n), 256, 0, (cudaStream_t) stream>>>(dst, n);

src/prefill/moe_ws.cu

@@ -0,0 +1,157 @@
+// src/prefill/moe_ws.cu - see include/strata/prefill/moe_ws.hpp: which layers and products take the weight-stationary
+// kernels (moe_ws_gu_iq2.cu, moe_ws_gu_iq3.cu, moe_ws_down.cu), their work tables and launches.
+#include "common.cuh"
+#include "mmq.cuh"
+
+#include "moe_ws.cuh"
+
+#include <algorithm>
+#include <atomic>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+namespace strata::prefill::ws {
+namespace {
+
+// ws / per-group MMQ (one 512-expert top-10 layer, lognormal counts sigma 0.5 and 1.0, the worse of the two; isolated
+// sustained medians measured on a GP100 die, the variant moe_ws_*.cu picks) at 2.5, 10, 35 and 160 rows per routed
+// expert
+struct Measured { int type; float r[4]; };
+constexpr double kU[4] = {2.5, 10.0, 35.0, 160.0};
+constexpr Measured kMeasured[] = {
+    {GGML_TYPE_IQ2_XXS, {0.276f, 0.494f, 0.728f, 0.754f}},
+    {GGML_TYPE_IQ2_XS, {0.299f, 0.513f, 0.750f, 0.824f}},
+    {GGML_TYPE_IQ2_S, {0.317f, 0.520f, 0.754f, 0.816f}},
+    {GGML_TYPE_IQ3_XXS, {0.329f, 0.514f, 0.755f, 0.786f}},
+    {GGML_TYPE_IQ3_S, {0.347f, 0.521f, 0.715f, 0.793f}},
+    {GGML_TYPE_Q2_0, {0.246f, 0.389f, 0.555f, 0.614f}},
+    {GGML_TYPE_IQ4_NL, {0.292f, 0.492f, 0.682f, 0.742f}},
+};
+
+const Measured* measured(int t) {
+    for (const Measured& m : kMeasured)
+        if (m.type == t) return &m;
+    return nullptr;
+}
+bool gu_type(int t) {
+    return t == GGML_TYPE_IQ2_XXS || t == GGML_TYPE_IQ2_XS || t == GGML_TYPE_IQ2_S || t == GGML_TYPE_IQ3_XXS ||
+           t == GGML_TYPE_IQ3_S;
+}
+
+}  // namespace
+
+int mode() {
+#if defined(__HIPCC__)
+    return 0;
+#else
+    static const int m = [] {
+        const char* e = std::getenv("STRATA_MMQ_WS");   // opt-in
+        if (e == nullptr || std::atoi(e) != 1) return 0;
+        // GP100 (cc 6.0) only: tuned and checked there (any visible GPU that is not: off for all)
+        int n = 0;
+        if (cudaGetDeviceCount(&n) != cudaSuccess || n <= 0) { cudaGetLastError(); return 0; }
+        for (int d = 0; d < n; ++d) {
+            cudaDeviceProp p{};
+            if (cudaGetDeviceProperties(&p, d) != cudaSuccess || p.major != 6 || p.minor != 0) { cudaGetLastError(); return 0; }
+        }
+        return 1;
+    }();
+    return m;
+#endif
+}
+
+bool supported(int gt, int dt) { return mode() != 0 && gu_type(gt) && (dt == GGML_TYPE_Q2_0 || dt == GGML_TYPE_IQ4_NL); }
+
+namespace {
+std::atomic<int> g_fast_override{-1};   // set_fast (benches, tests); -1: the environment
+}  // namespace
+
+int fast() {
+    const int o = g_fast_override.load(std::memory_order_relaxed);
+    if (o >= 0) return o;
+    static const int f = [] {
+        const char* e = std::getenv("STRATA_WS_FAST");   // opt-in
+        return e != nullptr && std::atoi(e) == 1 ? 1 : 0;
+    }();
+    return f;
+}
+
+void set_fast(int v) { g_fast_override.store(v < 0 ? -1 : (v != 0 ? 1 : 0), std::memory_order_relaxed); }
+
+bool faster(int t, double u) {
+    const Measured* m = measured(t);
+    if (m == nullptr) return false;
+    // linear in log u between the measured points, the ends held
+    if (u <= kU[0]) return m->r[0] < 1.0f;
+    for (int i = 1; i < 4; ++i)
+        if (u <= kU[i]) {
+            const double f = std::log(u / kU[i - 1]) / std::log(kU[i] / kU[i - 1]);
+            return m->r[i - 1] + f * (m->r[i] - m->r[i - 1]) < 1.0;
+        }
+    return m->r[3] < 1.0f;
+}
+
+size_t chunks(const int32_t* order, size_t n, const int32_t* cnt, const int32_t* bounds, Chunk* out) {
+    // largest expert first: MMQ's order is ascending by count where the resident sort ran (walked backwards), else
+    // by id (sorted here)
+    bool asc = true;
+    for (size_t j = 1; j < n && asc; ++j) asc = cnt[order[j - 1]] <= cnt[order[j]];
+    static thread_local std::vector<int32_t> by;
+    by.resize(n);
+    for (size_t j = 0; j < n; ++j) by[j] = (int32_t) (n - 1 - j);
+    if (!asc)
+        std::stable_sort(by.begin(), by.end(), [&](int32_t a, int32_t b) { return cnt[order[a]] > cnt[order[b]]; });
+    size_t k = 0;
+    for (const int32_t j : by) {
+        const int32_t c = cnt[order[j]], b = bounds[j];
+        for (int32_t tk = 0; tk < c; tk += kChunkRows) out[k++] = Chunk{j, b + tk, std::min(kChunkRows, c - tk), 0};
+    }
+    return k;
+}
+
+bool gate_up(const Layer& L, void* stream) {
+    if (L.n <= 0 || L.nchunks <= 0) return true;
+    const double u = (double) L.rows / L.n;
+    if (mode() != 1 || !faster(L.gu_type, u)) return false;
+    const cudaStream_t s = (cudaStream_t) stream;
+    switch (L.gu_type) {
+        case GGML_TYPE_IQ2_XXS: detail::gu_iq2_xxs(L, u, s); break;
+        case GGML_TYPE_IQ2_XS: detail::gu_iq2_xs(L, u, s); break;
+        case GGML_TYPE_IQ2_S: detail::gu_iq2_s(L, u, s); break;
+        case GGML_TYPE_IQ3_XXS: detail::gu_iq3_xxs(L, u, s); break;
+        case GGML_TYPE_IQ3_S: detail::gu_iq3_s(L, u, s); break;
+        default:
+            std::fprintf(stderr, "prefill ws: gate/up type %d has no kernel\n", L.gu_type);
+            std::exit(1);
+    }
+    return true;
+}
+
+void down(const Layer& L, void* stream) {
+    if (L.n <= 0 || L.nchunks <= 0) return;
+    const double u = (double) L.rows / L.n;
+    const cudaStream_t s = (cudaStream_t) stream;
+    switch (L.d_type) {
+        case GGML_TYPE_Q2_0: detail::dn_q2_0(L, u, s); break;
+        case GGML_TYPE_IQ4_NL: {
+            const bool f = fast() == 1;
+            if (f) {   // once per process: the switch is on and the WS path runs
+                static const bool said = [] {
+                    std::fprintf(stderr, "strata prefill: STRATA_WS_FAST=1: the weight-stationary IQ4_NL down's 32-token "
+                                         "tier with switched token pairs and transposed dy (the same Dm bits)\n");
+                    return true;
+                }();
+                (void) said;
+            }
+            detail::dn_iq4_nl(L, u, s, f);
+            break;
+        }
+        default:
+            std::fprintf(stderr, "prefill ws: down type %d has no kernel\n", L.d_type);
+            std::exit(1);
+    }
+}
+
+}  // namespace strata::prefill::ws