diff --git a/CMakeLists.txt b/CMakeLists.txt index 7c686766845..8926dc36f73 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/include/strata/prefill/moe_mmq.hpp b/include/strata/prefill/moe_mmq.hpp index fee662fcc76..6fba093e5f6 100644 --- a/include/strata/prefill/moe_mmq.hpp +++ b/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); diff --git a/include/strata/prefill/moe_ws.hpp b/include/strata/prefill/moe_ws.hpp new file mode 100644 index 00000000000..b1ebb73d902 --- /dev/null +++ b/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 +#include + +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 diff --git a/src/prefill/moe_mmq.cu b/src/prefill/moe_mmq.cu index e7ff5d45bc1..e832cffe38e 100644 --- a/src/prefill/moe_mmq.cu +++ b/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'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 +__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(row + 2 * i0), b = *reinterpret_cast(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(row + i0); + u4 = *reinterpret_cast(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(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<<>>(gu, (block_q8_1_mmq*) hq, rows); + else swiglu_quant_kernel<<>>(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<<>>(dst, n); diff --git a/src/prefill/moe_ws.cu b/src/prefill/moe_ws.cu new file mode 100644 index 00000000000..3393f70e4b4 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +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 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 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 diff --git a/src/prefill/moe_ws.cuh b/src/prefill/moe_ws.cuh new file mode 100644 index 00000000000..ef37d4340c1 --- /dev/null +++ b/src/prefill/moe_ws.cuh @@ -0,0 +1,49 @@ +// src/prefill/moe_ws.cuh - the weight-stationary expert kernels' shared pieces (see moe_ws.hpp): included after +// ggml-cuda's common.cuh and mmq.cuh by moe_ws.cu and the moe_ws_*.cu that hold the kernels (a few formats per file, +// so the instances compile in parallel). +#pragma once + +#include "strata/prefill/moe_ws.hpp" + +#include + +namespace strata::prefill::ws::detail { + +// The two ways every kernel multiplies: int8 values sign-extended to int16 pairs (one PRMT per pair: bytes 0,2 and 1,3) +// and a pair of products added to a 32-bit sum by two XMAD (mad.wide.s16) - MMQ's dp4a on sm_60, one product at a time. +// A 32-value block's integer sum is exact (|sum| < 2^22 for every format here), so the order of its MACs is free. +__device__ __forceinline__ void split2(const int v, int& lo02, int& hi13) { + asm("prmt.b32 %0, %1, 0, 0xA280;" : "=r"(lo02) : "r"(v)); + asm("prmt.b32 %0, %1, 0, 0xB391;" : "=r"(hi13) : "r"(v)); +} +__device__ __forceinline__ int mad2(const int w, const int a, int c) { + asm("{\n\t.reg .s16 al, ah, bl, bh;\n\t" + "mov.b32 {al, ah}, %1;\n\tmov.b32 {bl, bh}, %2;\n\t" + "mad.wide.s16 %0, al, bl, %0;\n\tmad.wide.s16 %0, ah, bh, %0;\n\t}" + : "+r"(c) : "r"(w), "r"(a)); + return c; +} +// float(sum) for |sum| < 2^22 without the quarter-rate I2F: a chain seeded with 1.5 * 2^23 holds the float's bits +constexpr int kMagicI = 0x4B400000; +constexpr float kMagicF = 12582912.0f; + +// CTA b's chunk and first weight row: chunk b / pieces, piece b % pieces (`pieces` = weight rows / rows per CTA) - the +// chunks' order, each chunk's pieces next to each other +template +__device__ __forceinline__ Chunk cta_work(const Chunk* __restrict__ chunks, int& r0) { + r0 = (int) (blockIdx.x % PIECES) * ROWS; + const int4 c = __ldg((const int4*) chunks + blockIdx.x / PIECES); + return Chunk{c.x, c.y, c.z, c.w}; +} + +// one launch of a layer's product, the variant picked by u (rows per routed expert); grid = nchunks x pieces +void gu_iq2_xxs(const Layer& L, double u, cudaStream_t s); +void gu_iq2_xs(const Layer& L, double u, cudaStream_t s); +void gu_iq2_s(const Layer& L, double u, cudaStream_t s); +void gu_iq3_xxs(const Layer& L, double u, cudaStream_t s); +void gu_iq3_s(const Layer& L, double u, cudaStream_t s); +void dn_q2_0(const Layer& L, double u, cudaStream_t s); +// `fast`: ws::fast() (STRATA_WS_FAST=1) - the 32-token tier's fast kernel, the same Dm rows +void dn_iq4_nl(const Layer& L, double u, cudaStream_t s, bool fast); + +} // namespace strata::prefill::ws::detail diff --git a/src/prefill/moe_ws_down.cu b/src/prefill/moe_ws_down.cu new file mode 100644 index 00000000000..346f3a820ec --- /dev/null +++ b/src/prefill/moe_ws_down.cu @@ -0,0 +1,645 @@ +// src/prefill/moe_ws_down.cu - weight-stationary down products (see moe_ws.hpp): Q2_0 and IQ4_NL, [2560 x 640] read in +// place at the blob's down offset, K stopping at 640 (nothing past a matrix is read). A CTA = 128 or 256 threads = +// as many weight rows of one expert (a row per thread, all of the chunk's H rows TPP at a time, the sums in registers); +// per K step the next step's weight bytes and q8_1 words are loaded into registers before this step's compute and +// stored to shared memory after it (one barrier pair per step). The H rows are ONE mmq::swiglu_quant over all of the +// layer's rows: [k/128][rows] block_q8_1_mmq, of which blocks 0-4 are read. (Measured on a GP100 die: ws / per-group +// MMQ 0.24-0.29 at 2.5 rows per expert, 0.38-0.49 at 10, 0.52-0.68 at 35, 0.61-0.74 at 160; bit-identical, +// compute-sanitizer clean with exact-size matrices.) +#include "common.cuh" +#include "mmq.cuh" + +#include "moe_ws.cuh" + +#include +#include +#include + +namespace strata::prefill::ws::detail { +namespace { + +constexpr int DN_ROWS = 2560, DN_K = 640; + +void ck(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "prefill ws: %s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +// ---- Q2_0: 18-B blocks {f16 d; 16 B of 2-bit codes}, 180-B rows; a step is 128 values = 2 blocks = 36 B of a row. +// MMQ's values {-1, 0, 1, 2}[code] (ggml_cuda_mmq_load_tiles_q2_0's PRMT, the selector masked to 3 bits) and the +// activations' int8 are both made exact fp16 pairs; a 32-value block is 16 HFMA2 into 2 half2 sums, each fp16 lane +// summing 8 products of magnitude <= 256 - integers <= 2048, exact - added in fp32: the float MMQ's I2F(sumi) gives. +// Epilogue as MMQ's sm_60 SASS: per 32-value block FMUL.FTZ (dx * dy), FFMA.FTZ (prod, sum, acc), blocks in K order. +namespace q20 { +constexpr int NT = 128, ROWB = 180, NSTEP = DN_K / 128, QB = 144, PIECES = DN_ROWS / NT; + +__device__ __forceinline__ uint32_t prmt(const uint32_t a, const uint32_t b, const uint32_t s) { + uint32_t d; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(d) : "r"(a), "r"(b), "r"(s)); + return d; +} +__device__ __forceinline__ uint32_t hfma2(const uint32_t a, const uint32_t b, const uint32_t c) { + uint32_t d; + asm("fma.rn.f16x2 %0, %1, %2, %3;" : "=r"(d) : "r"(a), "r"(b), "r"(c)); + return d; +} +__device__ __forceinline__ uint32_t hadd2(const uint32_t a, const uint32_t b) { + uint32_t d; + asm("add.rn.f16x2 %0, %1, %2;" : "=r"(d) : "r"(a), "r"(b)); + return d; +} +// the two fp16 lanes of h converted exactly to f32 and added +__device__ __forceinline__ float h2sum(const uint32_t h) { + float f; + asm("{\n\t.reg .f16 lo, hi;\n\t.reg .f32 a, b;\n\tmov.b32 {lo, hi}, %1;\n\t" + "cvt.f32.f16 a, lo;\n\tcvt.f32.f16 b, hi;\n\tadd.f32 %0, a, b;\n\t}" + : "=f"(f) : "r"(h)); + return f; +} +__device__ __forceinline__ float f16_to_f32(const uint32_t bits16) { + float f; + asm("{\n\t.reg .b16 h;\n\tcvt.u16.u32 h, %1;\n\tcvt.f32.f16 %0, h;\n\t}" : "=f"(f) : "r"(bits16)); + return f; +} + +// the 16 values of the int16 qa = x2 & 0xFFFF, qb = x2 >> 16 of a block's codes (value i at bits 2i) as fp16 pairs +// (v0,v4) (v1,v5) (v2,v6) (v3,v7), qa's to wa and qb's to wb: prmt(H, 0, sel) with nibbles 1 / 3 of sel the code with +// its low bit flipped and nibbles 0 / 2 zero (byte 0 of H, the zero high byte) +__device__ __forceinline__ void dec2_h16(const uint32_t x2, uint32_t (&wa)[4], uint32_t (&wb)[4]) { + constexpr uint32_t H = 0x3C40BC00u; // bytes 00 BC 40 3C = fp16 high bytes of {0, -1, +2, +1}[c ^ 1] = {-1,0,1,2}[c] + const uint32_t x = x2 ^ 0x55555555u; + const uint32_t m0 = (x << 4) & 0x30303030u, m1 = (x << 2) & 0x30303030u, m2 = x & 0x30303030u, m3 = (x >> 2) & 0x30303030u; + wa[0] = prmt(H, 0u, m0); + wa[1] = prmt(H, 0u, m1); + wa[2] = prmt(H, 0u, m2); + wa[3] = prmt(H, 0u, m3); + wb[0] = prmt(H, 0u, m0 >> 16); + wb[1] = prmt(H, 0u, m1 >> 16); + wb[2] = prmt(H, 0u, m2 >> 16); + wb[3] = prmt(H, 0u, m3 >> 16); +} + +// the 8 int8 of activation words A (values 0-3) and B (4-7) as fp16 pairs (a0,a4) (a1,a5) (a2,a6) (a3,a7): +// u = a + 128; fp16 bits 0x64uu = 1024 + u exactly; minus 1152 = a exactly +__device__ __forceinline__ void act_h16(const uint32_t A, const uint32_t B, uint32_t (&a)[4]) { + const uint32_t P1 = prmt(A, B, 0x5140u) ^ 0x80808080u; // u0 u4 u1 u5 + const uint32_t P2 = prmt(A, B, 0x7362u) ^ 0x80808080u; // u2 u6 u3 u7 + constexpr uint32_t M = 0x64646464u, NEG1152 = 0xE480E480u; + a[0] = hadd2(prmt(P1, M, 0x4140u), NEG1152); + a[1] = hadd2(prmt(P1, M, 0x4342u), NEG1152); + a[2] = hadd2(prmt(P2, M, 0x4140u), NEG1152); + a[3] = hadd2(prmt(P2, M, 0x4342u), NEG1152); +} + +// the exact integer dot of one 32-value block (weights w[16], activations a0..a3 = 16 fp16 pairs) as a float +__device__ __forceinline__ float block_sum(const uint32_t (&w)[16], const uint4 a0, const uint4 a1, const uint4 a2, + const uint4 a3) { + uint32_t p = hfma2(w[0], a0.x, 0u), q = hfma2(w[1], a0.y, 0u); + p = hfma2(w[2], a0.z, p); q = hfma2(w[3], a0.w, q); + p = hfma2(w[4], a1.x, p); q = hfma2(w[5], a1.y, q); + p = hfma2(w[6], a1.z, p); q = hfma2(w[7], a1.w, q); + p = hfma2(w[8], a2.x, p); q = hfma2(w[9], a2.y, q); + p = hfma2(w[10], a2.z, p); q = hfma2(w[11], a2.w, q); + p = hfma2(w[12], a3.x, p); q = hfma2(w[13], a3.y, q); + p = hfma2(w[14], a3.z, p); q = hfma2(w[15], a3.w, q); + return h2sum(p) + h2sum(q); +} + +// TPP tokens per pass; tokens guarded in pairs (a pair's token past the count runs on a clamped row, not written). +// Per step every thread stages exactly 9 words of the CTA's 128 x 36 B (row stride 9 words: conflict-free decode reads) +// and threads < 8 TPP one 16-value q8_1 piece each. +template +__global__ void __launch_bounds__(NT, 4) dn_kernel(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const int64_t xadd, + const char* __restrict__ xq, const int total_rows, + float* __restrict__ dst) { + constexpr int G = 2, WPR = 9, AQ = TPP * 8, AI = (AQ + NT - 1) / NT; + static_assert(TPP % G == 0, ""); + __shared__ int wsm[NT * WPR]; + __shared__ __align__(16) uint32_t act_s[TPP][4][16]; + __shared__ __align__(8) float dy_s[4][TPP]; + int r0; + const Chunk c = cta_work(chunks, r0); + const char* __restrict__ W = w + xoff[c.j] + xadd + (size_t) r0 * ROWB; + const int tid = threadIdx.x; + int woff[WPR]; // word tid + NT i of the step's NT x 9 words: row (tid + NT i) / 9, word (tid + NT i) % 9 +#pragma unroll + for (int i = 0; i < WPR; ++i) { + const int idx = tid + i * NT, r = idx / WPR; + woff[i] = r * ROWB + 4 * (idx - r * WPR); + } + int wpre[WPR]; + uint4 apre[AI]; + float dpre[AI]; + auto load_w = [&](const int st) { + const char* __restrict__ Ws = W + st * 36; +#pragma unroll + for (int i = 0; i < WPR; ++i) wpre[i] = __ldg((const int*) (Ws + woff[i])); + }; + auto store_w = [&]() { +#pragma unroll + for (int i = 0; i < WPR; ++i) wsm[tid + i * NT] = wpre[i]; + }; + auto load_a = [&](const int st, const int trow0, const int nt) { +#pragma unroll + for (int i = 0; i < AI; ++i) { + const int idx = tid + i * NT; + if (AQ % NT == 0 || idx < AQ) { + const int t = idx >> 3, m = idx & 7; + const char* blk = xq + ((size_t) st * total_rows + trow0 + min(t, nt - 1)) * QB; + apre[i] = __ldg((const uint4*) (blk + 16) + m); // values 16m .. 16m + 15 of the 128 + dpre[i] = __ldg((const float*) blk + (m >> 1)); // d of their 32-value block + } + } + }; + auto store_a = [&]() { +#pragma unroll + for (int i = 0; i < AI; ++i) { + const int idx = tid + i * NT; + if (AQ % NT == 0 || idx < AQ) { + const int t = idx >> 3, m = idx & 7, b = m >> 1, hf = m & 1; + uint32_t p0[4], p1[4]; + act_h16(apre[i].x, apre[i].y, p0); + act_h16(apre[i].z, apre[i].w, p1); + uint4* d = (uint4*) &act_s[t][b][8 * hf]; + d[0] = make_uint4(p0[0], p0[1], p0[2], p0[3]); + d[1] = make_uint4(p1[0], p1[1], p1[2], p1[3]); + dy_s[b][t] = dpre[i]; // both pieces of a block write the same value + } + } + }; + + load_w(0); + load_a(0, c.row0, min(TPP, c.rows)); + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int st = 0; st < NSTEP; ++st) { + __syncthreads(); + store_w(); + store_a(); + __syncthreads(); + if (st + 1 < NSTEP) { + load_w(st + 1); + load_a(st + 1, trow0, nt); + } else if (p0 + TPP < c.rows) { + load_w(0); + load_a(0, trow0 + TPP, min(TPP, c.rows - p0 - TPP)); + } + // the row's 9 step words: block 0 = d0 | qs0 at bytes 0..17 (the codes 2 bytes into a word: funnel shift), + // block 1 = d1 in the top half of word 4, qs1 = words 5..8 + const uint32_t* cw = (const uint32_t*) (wsm + tid * WPR); +#pragma unroll 1 + for (int b = 0; b < 4; ++b) { // 32-value block: Q2_0 block b / 2, half b % 2 (int16 4h .. 4h + 3), K order + const int pq = b >> 1, h = b & 1; + const int base = pq ? 5 + 2 * h : 2 * h, sh = pq ? 0 : 16; + const uint32_t c0 = cw[base], c1 = cw[base + 1], c2 = cw[min(base + 2, WPR - 1)]; + const float dx = f16_to_f32(pq ? cw[4] >> 16 : cw[0] & 0xFFFFu); + uint32_t wv[16], wa[4], wb[4]; + dec2_h16(__funnelshift_r(c0, c1, sh), wa, wb); // int16 0 and 1 of the half +#pragma unroll + for (int k = 0; k < 4; ++k) { wv[k] = wa[k]; wv[4 + k] = wb[k]; } + dec2_h16(__funnelshift_r(c1, c2, sh), wa, wb); // int16 2 and 3 +#pragma unroll + for (int k = 0; k < 4; ++k) { wv[8 + k] = wa[k]; wv[12 + k] = wb[k]; } +#pragma unroll + for (int t0 = 0; t0 < TPP; t0 += G) { + if (t0 < nt) { + const float2 d2 = *(const float2*) &dy_s[b][t0]; +#pragma unroll + for (int t = t0; t < t0 + G; ++t) { + const uint4* a4 = (const uint4*) act_s[t][b]; + const float s = block_sum(wv, a4[0], a4[1], a4[2], a4[3]); + acc[t] = __fmaf_rn(__fmul_rn(dx, t == t0 ? d2.x : d2.y), s, acc[t]); + } + } + } + } + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * DN_ROWS + r0 + tid] = acc[t]; + } +} +} // namespace q20 + +// ---- IQ4_NL: 18-B blocks {f16 d; 16 B of 4-bit indices into kvalues_iq4nl}, 360-B rows; a step is SB blocks (SB 4: +// 72 B of a row, 8-byte aligned; SB 2: 36 B). The int8 values are get_int_from_table_16's table lookups (the 16-byte +// table in registers) without its final byte reorder - the same values paired (v4k, v4k+1), (v4k+16, v4k+17), ... and +// the activations paired the same way (PRMT 0x9180 / 0xB3A2); one XMAD chain per (row, token, block), then MMQ's sm_60 +// FMUL.FTZ (dx * dy), I2F (sumi), FFMA.FTZ, blocks in K order. Staging: 3 rows per warp and k (lanes 0-26, a row's 9 +// 4- or 8-byte pieces), the last k guarded so no row past the CTA's is read. +namespace nl4 { +constexpr int NB = DN_K / QK4_NL, BB = (int) sizeof(block_iq4_nl), ROWB = NB * BB; // 20, 18, 360 +static_assert(NB == 20 && BB == 18 && ROWB == 360, ""); + +// bytes (0,1) and (2,3) of v, each sign-extended to 16 bits +__device__ __forceinline__ void split01(const int v, int& p01, int& p23) { + asm("prmt.b32 %0, %1, 0, 0x9180;" : "=r"(p01) : "r"(v)); + asm("prmt.b32 %0, %1, 0, 0xB3A2;" : "=r"(p23) : "r"(v)); +} +// the 16-byte kvalues_iq4nl table as 4 words, read byte by byte +__device__ __forceinline__ void load_tab(uint32_t (&tab)[4]) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + uint32_t v = 0; +#pragma unroll + for (int b = 0; b < 4; ++b) v |= (uint32_t) (uint8_t) kvalues_iq4nl[4 * i + b] << (8 * b); + tab[i] = v; + } +} +// a 32-value block from its 4 qs words (q[k]: values 4k..4k+3 in the low nibbles, 16+4k.. in the high ones): 16 int16 +// pair words against the activation pairs split01 makes of activation words m = 0..7 (at act[2m], act[2m + 1]). +// get_int_from_table_16's body (CUDA branch) up to its two lookup words tmp[0] = (v4k, v4k+16, v4k+1, v4k+17), +// tmp[1] = (v4k+2, v4k+18, v4k+3, v4k+19) +__device__ __forceinline__ void dec(const int (&q)[4], const uint32_t (&tab)[4], int (&wv)[16]) { +#pragma unroll + for (int k = 0; k < 4; ++k) { + const uint32_t q4 = (uint32_t) q[k]; + uint32_t tmp[2]; + const uint32_t low_high_selection_indices = (0x32103210 | ((q4 & 0x88888888) >> 1)); +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + const uint32_t shift = 16 * i; + const uint32_t low = __byte_perm(tab[0], tab[1], q4 >> shift); + const uint32_t high = __byte_perm(tab[2], tab[3], q4 >> shift); + tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift); + } + split2((int) tmp[0], wv[2 * k], wv[2 * (k + 4)]); // (v4k, v4k+1) | (v16+4k, v17+4k) + split2((int) tmp[1], wv[2 * k + 1], wv[2 * (k + 4) + 1]); // (v4k+2, v4k+3) | (v18+4k, v19+4k) + } +} +// block PAR of block pair p of the staged step (9 words per pair: even block d = lo16(w0), qs = bytes 2..17 by funnel +// shifts; odd block d = hi16(w4), qs = w5..w8) +template +__device__ __forceinline__ float load_blk(const int* __restrict__ cw, const int p, int (&q)[4]) { + const int* c = cw + 9 * p; + if constexpr (PAR == 0) { +#pragma unroll + for (int k = 0; k < 4; ++k) q[k] = (int) __funnelshift_r((uint32_t) c[k], (uint32_t) c[k + 1], 16); + return __half2float(__ushort_as_half((unsigned short) ((uint32_t) c[0] & 0xFFFFu))); + } else { +#pragma unroll + for (int k = 0; k < 4; ++k) q[k] = c[5 + k]; + return __half2float(__ushort_as_half((unsigned short) ((uint32_t) c[4] >> 16))); + } +} + +template +__device__ __forceinline__ void blk_compute(const int* __restrict__ cw, const int p, const uint32_t (&tab)[4], + const int (*act_s)[SB][16], const float (*dy_s)[SB], const int nt, + float (&acc)[TPP]) { + const int b = 2 * p + PAR; // block within the step + int wv[16]; + float dx; + { + int q[4]; + dx = load_blk(cw, p, q); + dec(q, tab, wv); + } + // tokens in groups of G behind one guard (a group's tokens past nt run on stale shared memory, not written) +#pragma unroll + for (int t0 = 0; t0 < TPP; t0 += G) { + if (t0 < nt) { +#pragma unroll + for (int t = t0; t < t0 + G; ++t) { + const int4* a4 = (const int4*) act_s[t][b]; + const int4 a0 = a4[0], a1 = a4[1], a2 = a4[2], a3 = a4[3]; + const float dy = dy_s[t][b]; + int s0 = 0; + s0 = mad2(wv[0], a0.x, s0); s0 = mad2(wv[1], a0.y, s0); s0 = mad2(wv[2], a0.z, s0); s0 = mad2(wv[3], a0.w, s0); + s0 = mad2(wv[4], a1.x, s0); s0 = mad2(wv[5], a1.y, s0); s0 = mad2(wv[6], a1.z, s0); s0 = mad2(wv[7], a1.w, s0); + s0 = mad2(wv[8], a2.x, s0); s0 = mad2(wv[9], a2.y, s0); s0 = mad2(wv[10], a2.z, s0); s0 = mad2(wv[11], a2.w, s0); + s0 = mad2(wv[12], a3.x, s0); s0 = mad2(wv[13], a3.y, s0); s0 = mad2(wv[14], a3.z, s0); s0 = mad2(wv[15], a3.w, s0); + acc[t] = __fmaf_rn(__fmul_rn(dx, dy), __int2float_rn(s0), acc[t]); + } + } + } +} + +template +__global__ void __launch_bounds__(NT, 1) dn_kernel(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const int64_t xadd, + const block_q8_1_mmq* __restrict__ xq, const int total_rows, + float* __restrict__ dst) { + static_assert((SB == 2 || SB == 4) && NT % 32 == 0 && TPP % G == 0, ""); + constexpr int NW = NT / 32, WPS = 9 * SB / 2, S = WPS | 1, NSTEP = NB / SB, PIECES = DN_ROWS / NT; + constexpr int KI = (NT + 3 * NW - 1) / (3 * NW); + constexpr int AW = 9 * SB; // activation words per token and step (SB d + 8 SB qs words) + constexpr int AI = (TPP * AW + NT - 1) / NT; + using PT = typename std::conditional::type; + __shared__ int wsm[NT * S]; + __shared__ __align__(16) int act_s[TPP][SB][16]; + __shared__ float dy_s[TPP][SB]; + uint32_t tab[4]; + load_tab(tab); + int r0; + const Chunk c = cta_work(chunks, r0); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int sub = lane / 9, pc = lane - 9 * sub; // sub 3: lanes 27-31, idle + const int srow0 = 3 * warp + sub; + const bool sact = sub < 3; + const char* __restrict__ Wg = w + xoff[c.j] + xadd + (size_t) (r0 + srow0) * ROWB + pc * (int) sizeof(PT); + int* __restrict__ ws_st = wsm + srow0 * S + pc * (int) (sizeof(PT) / 4); + PT pre[KI]; + int apre[AI]; + auto load_w = [&](const int st) { + if (sact) { + const char* g = Wg + st * (SB * BB); +#pragma unroll + for (int k = 0; k < KI; ++k) + if (k < KI - 1 || srow0 + 3 * NW * k < NT) pre[k] = __ldg((const PT*) (g + (size_t) k * 3 * NW * ROWB)); + } + }; + auto store_w = [&]() { + if (sact) { +#pragma unroll + for (int k = 0; k < KI; ++k) + if (k < KI - 1 || srow0 + 3 * NW * k < NT) { + if constexpr (SB == 4) { + ws_st[k * 3 * NW * S] = pre[k].x; + ws_st[k * 3 * NW * S + 1] = pre[k].y; + } else { + ws_st[k * 3 * NW * S] = pre[k]; + } + } + } + }; + // word wd of token t's step: d4[b0 + wd] for wd < SB, else qs word 8 b0 + (wd - SB) + auto load_a = [&](const int st, const int trow0, const int nt) { + const int kb = (st * SB) >> 2, b0 = (st * SB) & 3; // the 128-value q8_1 block and the first 32-block in it +#pragma unroll + for (int i = 0; i < AI; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * AW) { + const int t = idx / AW, wd = idx - t * AW; + const int* blk = (const int*) (xq + (size_t) kb * total_rows + trow0 + t); + apre[i] = __ldg(blk + (wd < SB ? b0 + wd : 4 + 8 * b0 + (wd - SB))); + } + } + }; + auto store_a = [&](const int nt) { +#pragma unroll + for (int i = 0; i < AI; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * AW) { + const int t = idx / AW, wd = idx - t * AW; + if (wd < SB) { + dy_s[t][wd] = __int_as_float(apre[i]); + } else { + const int m = wd - SB, bb = m >> 3, mm = m & 7; + int lo, hi; + split01(apre[i], lo, hi); + *(int2*) &act_s[t][bb][2 * mm] = make_int2(lo, hi); + } + } + } + }; + load_w(0); + load_a(0, c.row0, min(TPP, c.rows)); + const int* cw = wsm + threadIdx.x * S; + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int st = 0; st < NSTEP; ++st) { + __syncthreads(); + store_w(); + store_a(nt); + __syncthreads(); + if (st + 1 < NSTEP) { + load_w(st + 1); + load_a(st + 1, trow0, nt); + } else if (p0 + TPP < c.rows) { + load_w(0); + load_a(0, trow0 + TPP, min(TPP, c.rows - p0 - TPP)); + } +#pragma unroll 1 + for (int p = 0; p < SB / 2; ++p) { + blk_compute(cw, p, tab, act_s, dy_s, nt, acc); + blk_compute(cw, p, tab, act_s, dy_s, nt, acc); + } + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * DN_ROWS + r0 + threadIdx.x] = acc[t]; + } +} + +// ---- STRATA_WS_FAST=1 (ws::fast(), moe_ws.hpp): the 256 x 32 tier (2-block steps, pairs of tokens) with fewer +// instructions per pass, the same Dm rows bit for bit: every token's operations and their order are dn_kernel's. +// (Prototyped as a bench variant on one GP100 die: memcmp-equal to dn_kernel on real IQ4_NL layers at T 1024 and +// 1536, lognormal sigma 0.5 / 1.0 and a flat-32 control; 4.7-9.2% faster, 122 -> 102 registers, still two CTAs per +// SM. moe_ws_down_bench checks and times it against dn_kernel.) +// DUFF a switch on the number of token pairs falls through the 16 unrolled pairs (a jump table) instead of a guard +// per pair and block +// DYT dy as [block][token]: a pair's two dy in one LDS.64 +template +__device__ __forceinline__ void blk_compute_fast(const int* __restrict__ cw, const int p, const uint32_t (&tab)[4], + const int (*act_s)[SB][16], const float* __restrict__ dyb, const int nt, + float (&acc)[TPP]) { + static_assert(TPP / G == 16 && G == 2, "the switch below: 16 token pairs"); + const int b = 2 * p + PAR; // block within the step + int wv[16]; + float dx; + { + int q[4]; + dx = load_blk(cw, p, q); + dec(q, tab, wv); + } + // pairs past nt's run on stale shared memory and are not written, as dn_kernel's + auto grp = [&](auto K) { + constexpr int t0 = decltype(K)::value * G; + const float2 d2 = *(const float2*) &dyb[b * TPP + t0]; +#pragma unroll + for (int t = t0; t < t0 + G; ++t) { + const int4* a4 = (const int4*) act_s[t][b]; + const int4 a0 = a4[0], a1 = a4[1], a2 = a4[2], a3 = a4[3]; + int s0 = 0; + s0 = mad2(wv[0], a0.x, s0); s0 = mad2(wv[1], a0.y, s0); s0 = mad2(wv[2], a0.z, s0); s0 = mad2(wv[3], a0.w, s0); + s0 = mad2(wv[4], a1.x, s0); s0 = mad2(wv[5], a1.y, s0); s0 = mad2(wv[6], a1.z, s0); s0 = mad2(wv[7], a1.w, s0); + s0 = mad2(wv[8], a2.x, s0); s0 = mad2(wv[9], a2.y, s0); s0 = mad2(wv[10], a2.z, s0); s0 = mad2(wv[11], a2.w, s0); + s0 = mad2(wv[12], a3.x, s0); s0 = mad2(wv[13], a3.y, s0); s0 = mad2(wv[14], a3.z, s0); s0 = mad2(wv[15], a3.w, s0); + acc[t] = __fmaf_rn(__fmul_rn(dx, t == t0 ? d2.x : d2.y), __int2float_rn(s0), acc[t]); + } + }; + using std::integral_constant; + switch ((nt + G - 1) / G) { + case 16: grp(integral_constant{}); [[fallthrough]]; + case 15: grp(integral_constant{}); [[fallthrough]]; + case 14: grp(integral_constant{}); [[fallthrough]]; + case 13: grp(integral_constant{}); [[fallthrough]]; + case 12: grp(integral_constant{}); [[fallthrough]]; + case 11: grp(integral_constant{}); [[fallthrough]]; + case 10: grp(integral_constant{}); [[fallthrough]]; + case 9: grp(integral_constant{}); [[fallthrough]]; + case 8: grp(integral_constant{}); [[fallthrough]]; + case 7: grp(integral_constant{}); [[fallthrough]]; + case 6: grp(integral_constant{}); [[fallthrough]]; + case 5: grp(integral_constant{}); [[fallthrough]]; + case 4: grp(integral_constant{}); [[fallthrough]]; + case 3: grp(integral_constant{}); [[fallthrough]]; + case 2: grp(integral_constant{}); [[fallthrough]]; + case 1: grp(integral_constant{}); [[fallthrough]]; + default: break; + } +} + +// dn_kernel with blk_compute_fast and dy_s as [SB][TPP]; everything else is dn_kernel's +template +__global__ void __launch_bounds__(NT, 1) dn_kernel_fast(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const int64_t xadd, + const block_q8_1_mmq* __restrict__ xq, const int total_rows, + float* __restrict__ dst) { + static_assert((SB == 2 || SB == 4) && NT % 32 == 0 && TPP % G == 0, ""); + constexpr int NW = NT / 32, WPS = 9 * SB / 2, S = WPS | 1, NSTEP = NB / SB, PIECES = DN_ROWS / NT; + constexpr int KI = (NT + 3 * NW - 1) / (3 * NW); + constexpr int AW = 9 * SB; // activation words per token and step (SB d + 8 SB qs words) + constexpr int AI = (TPP * AW + NT - 1) / NT; + using PT = typename std::conditional::type; + __shared__ int wsm[NT * S]; + __shared__ __align__(16) int act_s[TPP][SB][16]; + __shared__ __align__(8) float dy_s[SB * TPP]; // [SB][TPP]: dy of token t, block b at b * TPP + t + uint32_t tab[4]; + load_tab(tab); + int r0; + const Chunk c = cta_work(chunks, r0); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int sub = lane / 9, pc = lane - 9 * sub; // sub 3: lanes 27-31, idle + const int srow0 = 3 * warp + sub; + const bool sact = sub < 3; + const char* __restrict__ Wg = w + xoff[c.j] + xadd + (size_t) (r0 + srow0) * ROWB + pc * (int) sizeof(PT); + int* __restrict__ ws_st = wsm + srow0 * S + pc * (int) (sizeof(PT) / 4); + PT pre[KI]; + int apre[AI]; + auto load_w = [&](const int st) { + if (sact) { + const char* g = Wg + st * (SB * BB); +#pragma unroll + for (int k = 0; k < KI; ++k) + if (k < KI - 1 || srow0 + 3 * NW * k < NT) pre[k] = __ldg((const PT*) (g + (size_t) k * 3 * NW * ROWB)); + } + }; + auto store_w = [&]() { + if (sact) { +#pragma unroll + for (int k = 0; k < KI; ++k) + if (k < KI - 1 || srow0 + 3 * NW * k < NT) { + if constexpr (SB == 4) { + ws_st[k * 3 * NW * S] = pre[k].x; + ws_st[k * 3 * NW * S + 1] = pre[k].y; + } else { + ws_st[k * 3 * NW * S] = pre[k]; + } + } + } + }; + auto load_a = [&](const int st, const int trow0, const int nt) { + const int kb = (st * SB) >> 2, b0 = (st * SB) & 3; +#pragma unroll + for (int i = 0; i < AI; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * AW) { + const int t = idx / AW, wd = idx - t * AW; + const int* blk = (const int*) (xq + (size_t) kb * total_rows + trow0 + t); + apre[i] = __ldg(blk + (wd < SB ? b0 + wd : 4 + 8 * b0 + (wd - SB))); + } + } + }; + auto store_a = [&](const int nt) { +#pragma unroll + for (int i = 0; i < AI; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * AW) { + const int t = idx / AW, wd = idx - t * AW; + if (wd < SB) { + dy_s[wd * TPP + t] = __int_as_float(apre[i]); + } else { + const int m = wd - SB, bb = m >> 3, mm = m & 7; + int lo, hi; + split01(apre[i], lo, hi); + *(int2*) &act_s[t][bb][2 * mm] = make_int2(lo, hi); + } + } + } + }; + load_w(0); + load_a(0, c.row0, min(TPP, c.rows)); + const int* cw = wsm + threadIdx.x * S; + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int st = 0; st < NSTEP; ++st) { + __syncthreads(); + store_w(); + store_a(nt); + __syncthreads(); + if (st + 1 < NSTEP) { + load_w(st + 1); + load_a(st + 1, trow0, nt); + } else if (p0 + TPP < c.rows) { + load_w(0); + load_a(0, trow0 + TPP, min(TPP, c.rows - p0 - TPP)); + } +#pragma unroll 1 + for (int p = 0; p < SB / 2; ++p) { + blk_compute_fast(cw, p, tab, act_s, dy_s, nt, acc); + blk_compute_fast(cw, p, tab, act_s, dy_s, nt, acc); + } + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * DN_ROWS + r0 + threadIdx.x] = acc[t]; + } +} +} // namespace nl4 + +} // namespace + +// Q2_0: 8 tokens per pass below 5 rows per expert, then 16 (equal at 5) +void dn_q2_0(const Layer& L, double u, cudaStream_t s) { + if (L.nchunks <= 0) return; + const unsigned grid = (unsigned) (L.nchunks * q20::PIECES); + if (u < 5.0) q20::dn_kernel<8><<>>(L.chunks, (const char*) L.w, L.xoff, L.down_off, (const char*) L.hq, + (int) L.rows, L.dm); + else q20::dn_kernel<16><<>>(L.chunks, (const char*) L.w, L.xoff, L.down_off, (const char*) L.hq, + (int) L.rows, L.dm); + ck(cudaGetLastError(), "down Q2_0"); +} + +// IQ4_NL: 128 rows x 8 tokens (4-block steps) below 5 rows per expert, 256 x 16 below 20, then 256 x 32 (2-block steps); +// `fast` (STRATA_WS_FAST=1): the 256 x 32 tier runs dn_kernel_fast instead - the same Dm rows +void dn_iq4_nl(const Layer& L, double u, cudaStream_t s, bool fast) { + if (L.nchunks <= 0) return; + const char* w = (const char*) L.w; + const block_q8_1_mmq* x = (const block_q8_1_mmq*) L.hq; + if (u < 5.0) + nl4::dn_kernel<128, 8, 4, 2><<<(unsigned) (L.nchunks * (DN_ROWS / 128)), 128, 0, s>>>(L.chunks, w, L.xoff, L.down_off, x, + (int) L.rows, L.dm); + else if (u < 20.0) + nl4::dn_kernel<256, 16, 4, 2><<<(unsigned) (L.nchunks * (DN_ROWS / 256)), 256, 0, s>>>(L.chunks, w, L.xoff, L.down_off, x, + (int) L.rows, L.dm); + else if (fast) + nl4::dn_kernel_fast<256, 32, 2, 2><<<(unsigned) (L.nchunks * (DN_ROWS / 256)), 256, 0, s>>>(L.chunks, w, L.xoff, + L.down_off, x, + (int) L.rows, L.dm); + else + nl4::dn_kernel<256, 32, 2, 2><<<(unsigned) (L.nchunks * (DN_ROWS / 256)), 256, 0, s>>>(L.chunks, w, L.xoff, L.down_off, x, + (int) L.rows, L.dm); + ck(cudaGetLastError(), "down IQ4_NL"); +} + +} // namespace strata::prefill::ws::detail diff --git a/src/prefill/moe_ws_down_bench.cu b/src/prefill/moe_ws_down_bench.cu new file mode 100644 index 00000000000..086bf2528cd --- /dev/null +++ b/src/prefill/moe_ws_down_bench.cu @@ -0,0 +1,270 @@ +// src/prefill/moe_ws_down_bench.cu - moe_ws_down_bench: the weight-stationary down (ws::down) on one GP100 die with +// and without STRATA_WS_FAST (ws::set_fast), byte for byte and timed. Not built by default (EXCLUDE_FROM_ALL; the +// STRATA_EXPERIMENTAL_SM60 build; needs STRATA_MMQ_WS=1 and a cc 6.0 device). +// +// moe_ws_down_bench [--layers L,..] [--t T,..] [--models ln0.5,ln1.0,flat32] [--reps N] +// +// Per layer whose down experts are IQ4_NL or Q2_0: all 512 experts' down matrices read from the GGUF into one VRAM +// blob, per T and routing model the counts of the T * 10 routed rows over the experts ("lnS": lognormal popularity, +// sigma S, each <= T; "flatN": experts of exactly N rows), ws::chunks' table, the GU rows random and H's q8_1 rows +// made by mmq::swiglu_quant as the prompt path makes them. ws::down runs once with set_fast(0) (the reference) and +// once with set_fast(1); the two Dm buffers (prefilled with different bytes) must be equal. Then both are timed, +// interleaved, a warm-up round first; medians. Exit 1 on any difference. +#include "strata/artifact/gguf_reader.hpp" +#include "strata/prefill/moe_mmq.hpp" +#include "strata/prefill/moe_ws.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CK(x) \ + do { \ + cudaError_t e_ = (x); \ + if (e_ != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d\n", cudaGetErrorString(e_), __FILE__, __LINE__); \ + std::exit(3); \ + } \ + } while (0) + +namespace ws = strata::prefill::ws; +namespace mmq = strata::prefill::mmq; + +namespace { + +constexpr int64_t kEmbd = 2560, kFf = 640, kK = 10, kNE = 512; +constexpr int kIQ4_NL = 20, kQ2_0 = 42; + +__global__ void fill_gu(float* x, long long n, unsigned long long seed) { + const long long i = (long long) blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + unsigned long long h = seed ^ (unsigned long long) i * 0x9E3779B97F4A7C15ull; + h ^= h >> 33; h *= 0xFF51AFD7ED558CCDull; h ^= h >> 33; h *= 0xC4CEB9FE1A85EC53ull; h ^= h >> 33; + const float u = (float) (h & 0xFFFFFF) / 16777216.0f - 0.5f; + x[i] = ((h >> 40) % 97 == 0) ? 24.0f * u : 4.0f * u; +} + +std::vector split_list(const char* s) { + std::vector v; + std::string t(s); + size_t p = 0; + while (p < t.size()) { + size_t q = t.find(',', p); + if (q == std::string::npos) q = t.size(); + v.push_back(t.substr(p, q - p)); + p = q + 1; + } + return v; +} + +// the counts of T * K rows over the 512 experts +std::vector make_counts(const std::string& model, int T, int layer) { + const int rows = T * (int) kK; + std::vector cnt(kNE, 0); + if (model.rfind("flat", 0) == 0) { + const int n = std::atoi(model.c_str() + 4); + int left = rows; + for (int e = 0; e < kNE && left > 0; ++e) { + cnt[e] = std::min(n, left); + left -= cnt[e]; + } + return cnt; + } + const double sigma = std::atof(model.c_str() + 2); + std::mt19937_64 rng(1000003ull * (uint64_t) T + (uint64_t) (sigma * 1000) + 7 + (uint64_t) layer); + std::normal_distribution nd(0.0, 1.0); + std::vector wgt(kNE); + for (auto& x : wgt) x = std::exp(sigma * nd(rng)); + const double S = std::accumulate(wgt.begin(), wgt.end(), 0.0); + std::vector> fr; + int sum = 0; + for (int e = 0; e < kNE; ++e) { + const double cc = rows * wgt[e] / S; + cnt[e] = std::min((int) std::floor(cc), T); + sum += cnt[e]; + fr.push_back({ cc - std::floor(cc), e }); + } + std::sort(fr.begin(), fr.end(), [](auto& a, auto& b) { return a.first > b.first; }); + for (size_t i = 0; sum < rows && i < 4 * (size_t) kNE; ++i) { + const int e = fr[i % kNE].second; + if (cnt[e] < T) { ++cnt[e]; ++sum; } + } + return cnt; +} + +double median(std::vector v) { + std::sort(v.begin(), v.end()); + return v.empty() ? 0.0 : v[v.size() / 2]; +} + +} // namespace + +int main(int argc, char** argv) { + setvbuf(stdout, nullptr, _IONBF, 0); + std::vector shards, models = { "ln0.5", "ln1.0", "flat32" }; + std::vector layers = { 0, 2 }, ts = { 768, 1024, 1536, 2048, 4096 }; + int reps = 15; + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto next = [&]() -> const char* { + if (i + 1 >= argc) { + std::fprintf(stderr, "%s needs a value\n", a.c_str()); + std::exit(2); + } + return argv[++i]; + }; + if (a == "--layers") { layers.clear(); for (auto& s : split_list(next())) layers.push_back(std::atoi(s.c_str())); } + else if (a == "--t") { ts.clear(); for (auto& s : split_list(next())) ts.push_back(std::atoi(s.c_str())); } + else if (a == "--models") models = split_list(next()); + else if (a == "--reps") reps = std::atoi(next()); + else if (a.size() > 5 && a.substr(a.size() - 5) == ".gguf") shards.push_back(a); + else { + std::fprintf(stderr, "usage: moe_ws_down_bench [--layers L,..] [--t T,..] " + "[--models ln0.5,ln1.0,flat32] [--reps N]\n"); + return 2; + } + } + if (shards.empty()) { + std::fprintf(stderr, "moe_ws_down_bench: give a model shard (.gguf)\n"); + return 2; + } + cudaDeviceProp p{}; + CK(cudaGetDeviceProperties(&p, 0)); + std::printf("device 0: %s, cc %d.%d, %d SMs, clock %d kHz; ws::mode() %d\n", p.name, p.major, p.minor, + p.multiProcessorCount, p.clockRate, ws::mode()); + if (ws::mode() != 1) { + std::fprintf(stderr, "needs the WS kernels (STRATA_EXPERIMENTAL_SM60 build, STRATA_MMQ_WS=1, a cc 6.0 device)\n"); + return 2; + } + std::unique_ptr model; + try { + model = std::make_unique(strata::GgufModel::open(shards[0])); + } catch (const std::exception& ex) { + std::fprintf(stderr, "%s\n", ex.what()); + return 2; + } + cudaEvent_t e0, e1; + CK(cudaEventCreate(&e0)); + CK(cudaEventCreate(&e1)); + auto time_once = [&](auto&& f) { + CK(cudaEventRecord(e0)); + f(); + CK(cudaEventRecord(e1)); + CK(cudaEventSynchronize(e1)); + CK(cudaGetLastError()); + float ms = 0; + CK(cudaEventElapsedTime(&ms, e0, e1)); + return (double) ms; + }; + long long fails = 0, checks = 0; + for (int l : layers) { + size_t sh = 0; + const strata::TensorInfo* t = model->find("blk." + std::to_string(l) + ".ffn_down_exps.weight", &sh); + if (t == nullptr) { + std::printf("layer %d: no ffn_down_exps\n", l); + ++fails; + continue; + } + const int dt = (int) t->type; + if (dt != kIQ4_NL && dt != kQ2_0) { + std::printf("layer %d: down type %d is not IQ4_NL / Q2_0, skipped\n", l, dt); + continue; + } + const size_t per = mmq::matrix_bytes(dt, kEmbd, kFf); + uint8_t* W = nullptr; + CK(cudaMalloc(&W, (size_t) kNE * per)); + CK(cudaMemcpy(W, model->shard(sh).tensor_data(*t), (size_t) kNE * per, cudaMemcpyHostToDevice)); + std::printf("\n##### layer %d: down %s (%d), %zu B per expert\n", l, t->type_name(), dt, per); + int64_t* dxo = nullptr; + CK(cudaMalloc(&dxo, kNE * 8)); + for (int T : ts) + for (const std::string& mdl : models) { + const std::vector cnt = make_counts(mdl, T, l); + std::vector order; + int R = 0; + for (int e = 0; e < kNE; ++e) + if (cnt[e] > 0) { order.push_back(e); R += cnt[e]; } + std::stable_sort(order.begin(), order.end(), + [&](int a, int b) { return cnt[a] != cnt[b] ? cnt[a] < cnt[b] : a < b; }); + const int n = (int) order.size(); + std::vector bounds(n + 1, 0); + for (int j = 0; j < n; ++j) bounds[j + 1] = bounds[j] + cnt[order[j]]; + std::vector ch(ws::max_chunks(R, n)); + const size_t nch = ws::chunks(order.data(), (size_t) n, cnt.data(), bounds.data(), ch.data()); + const double u = (double) R / n; + std::vector hxo(n); + for (int j = 0; j < n; ++j) hxo[j] = (int64_t) order[j] * (int64_t) per; + CK(cudaMemcpy(dxo, hxo.data(), (size_t) n * 8, cudaMemcpyHostToDevice)); + ws::Chunk* dch = nullptr; + CK(cudaMalloc(&dch, nch * sizeof(ws::Chunk))); + CK(cudaMemcpy(dch, ch.data(), nch * sizeof(ws::Chunk), cudaMemcpyHostToDevice)); + const size_t gub = (size_t) R * 2 * kFf * 4, dmb = (size_t) R * kEmbd * 4, hqb = mmq::q8_bytes(R, kFf); + float *GU = nullptr, *DMa = nullptr, *DMb = nullptr; + void* HQ = nullptr; + CK(cudaMalloc(&GU, gub)); + CK(cudaMalloc(&HQ, hqb)); + CK(cudaMalloc(&DMa, dmb)); + CK(cudaMalloc(&DMb, dmb)); + fill_gu<<<(unsigned) (((long long) R * 2 * kFf + 255) / 256), 256>>>(GU, (long long) R * 2 * kFf, + 0xD0A11ull + (uint64_t) R + (uint64_t) l); + CK(cudaGetLastError()); + CK(cudaMemset(HQ, 0, hqb)); + mmq::swiglu_quant(GU, HQ, R, false, nullptr); + CK(cudaDeviceSynchronize()); + CK(cudaGetLastError()); + ws::Layer L; + L.gu_type = -1; L.d_type = dt; L.w = W; L.xoff = dxo; L.down_off = 0; + L.chunks = dch; L.nchunks = (int64_t) nch; L.rows = R; L.n = n; L.hq = HQ; + ws::Layer La = L, Lb = L; + La.dm = DMa; + Lb.dm = DMb; + CK(cudaMemset(DMa, 0xFF, dmb)); + CK(cudaMemset(DMb, 0x00, dmb)); + ws::set_fast(0); + ws::down(La, nullptr); + ws::set_fast(1); + ws::down(Lb, nullptr); + ws::set_fast(0); + CK(cudaDeviceSynchronize()); + CK(cudaGetLastError()); + std::vector a(dmb), b(dmb); + CK(cudaMemcpy(a.data(), DMa, dmb, cudaMemcpyDeviceToHost)); + CK(cudaMemcpy(b.data(), DMb, dmb, cudaMemcpyDeviceToHost)); + long long rows_diff = 0; + for (int r = 0; r < R; ++r) + rows_diff += std::memcmp(a.data() + (size_t) r * kEmbd * 4, b.data() + (size_t) r * kEmbd * 4, kEmbd * 4) != 0; + const bool tier32 = dt == kIQ4_NL && u >= 20.0; + ++checks; + if (rows_diff) ++fails; + std::vector ta, tb; + for (int rep = -1; rep < reps; ++rep) { + const double x = time_once([&] { ws::set_fast(0); ws::down(La, nullptr); }); + const double y = time_once([&] { ws::set_fast(1); ws::down(Lb, nullptr); ws::set_fast(0); }); + if (rep >= 0) { ta.push_back(x); tb.push_back(y); } + } + const double ma = median(ta), mb = median(tb); + std::printf("L%d %-6s T %5d rows %6d over %3d experts (u %6.2f, %zu chunks; %s): default %8.3f ms, FAST %8.3f ms, " + "x%.3f; Dm %s\n", + l, mdl.c_str(), T, R, n, u, nch, tier32 ? "32-token tier, fast kernel" : "default kernel either way", + ma, mb, mb / ma, rows_diff ? ("DIFFER in " + std::to_string(rows_diff) + " rows").c_str() : "memcmp equal"); + CK(cudaFree(dch)); + CK(cudaFree(GU)); + CK(cudaFree(HQ)); + CK(cudaFree(DMa)); + CK(cudaFree(DMb)); + } + CK(cudaFree(dxo)); + CK(cudaFree(W)); + } + std::printf("\n%s: %lld checks, %lld failures\n", fails ? "FAIL" : "ALL CHECKS EQUAL", checks, fails); + return fails ? 1 : 0; +} diff --git a/src/prefill/moe_ws_gu_iq2.cu b/src/prefill/moe_ws_gu_iq2.cu new file mode 100644 index 00000000000..4df68b6d698 --- /dev/null +++ b/src/prefill/moe_ws_gu_iq2.cu @@ -0,0 +1,736 @@ +// src/prefill/moe_ws_gu_iq2.cu - weight-stationary gate/up products (see moe_ws.hpp): IQ2_XXS, IQ2_XS and IQ2_S, +// [1280 x 2560] read in place. A CTA = 128 threads = 128 weight rows of one expert (a row per thread, all of the +// chunk's activation rows TPP at a time, the sums in registers); K advances 128 values a step (half a 256-value block +// = 4 sub-blocks): the step's bytes of the 128 rows and the step's q8_1 block of each token go to shared memory (the +// activations split to int16 pairs once), then each thread decodes its row's 4 sub-blocks once per pass and runs them +// against the tokens. Bits: per 32-value sub-block MMQ's int8 values (ggml_cuda_mmq_load_tiles_iq2_*) and scale +// expressions, then MMQ's epilogue (IQ2_XXS: acc = fma(dx * dy, float(sum), acc); IQ2_XS / IQ2_S, a scale per 16 +// values: t = fma(float(sum 0-15), d0, 0), t = fma(float(sum 16-31), d1, t), acc = fma(t, dy, acc)), sub-blocks in K +// order. (Measured on a GP100 die: ws / per-group MMQ 0.26-0.33 at 2.5 +// rows per expert, 0.46-0.52 at 10, 0.67-0.76 at 35, 0.75-0.82 at 160; bit-identical, compute-sanitizer clean with +// exact-size matrices.) +// +// Decode: every iq2xxs / iq2xs / iq2s_grid byte is 0x08, 0x19 or 0x2B, so a grid entry is 8 nibble codes (0, 1, 2) +// and a value's sign adds 4 to its nibble: __byte_perm(0x002B1908, 0x00D5E7F8, codes | signs) is MMQ's +// __vsub4(g ^ s, s) word, one PRMT per 4 values (checked against MMQ's expressions for every entry and sign byte). +// Staging without per-lane branches: rows are 660 / 740 / 820 B, so block kc starts 2 * (kc & 1) bytes into a 4-byte +// word for the whole CTA - an even block's words are two aligned loads and a funnel shift (the second load ends at most +// 2 bytes into block kc + 1 of the same row), an odd block's are aligned; d sits the other way round (an odd block's d +// is the high half of the word 2 bytes before it, inside block kc - 1). Every load of a lane is issued before its +// stores, and nothing outside a row is read (the last row's last block is odd). Token groups of TG: one `t < nt` +// branch per group, so the XMAD chains of the group's tokens interleave (a group's tokens past the count are not +// written). +#include "common.cuh" +#include "mmq.cuh" + +#include "moe_ws.cuh" + +#include +#include + +namespace strata::prefill::ws::detail { +namespace { + +constexpr int GU_ROWS = 1280, GU_K = 2560, NB = GU_K / QK_K, NSTEP = 2 * NB, NT = 128, NW = NT / 32, PIECES = GU_ROWS / NT; +constexpr uint32_t PX = 0x002B1908u, PY = 0x00D5E7F8u; // codes 0-2: +8 +25 +43 | codes 4-6: -8 -25 -43 + +void ck(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "prefill ws: %s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +// the step's activations of IQ2_XS / IQ2_S: nt consecutive 144-B block_q8_1_mmq, d4 -> dy_s, int8 -> int16 pairs +template +__device__ __forceinline__ void stage_act(const block_q8_1_mmq* __restrict__ xq, const int hc, const int total_rows, + const int trow0, const int nt, int (*act_s)[4][16], float (*dy_s)[4]) { + constexpr int NA = (TPP * 36 + NT - 1) / NT; + int av[NA]; + const int* xr = (const int*) (xq + (size_t) hc * total_rows + trow0); +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * 36) av[i] = __ldg(xr + idx); + } +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * 36) { + const int t = idx / 36, wd = idx - t * 36; + if (wd < 4) { + dy_s[t][wd] = __int_as_float(av[i]); + } else { + const int m = wd - 4; + int lo, hi; + split2(av[i], lo, hi); + *(int2*) &act_s[t][m / 8][2 * (m % 8)] = make_int2(lo, hi); + } + } + } +} + +// IQ2_XS / IQ2_S: sub-block s's 16 int16-pair words wv and scales d0 / d1 against the tokens (vec_dot_q8_0_16_q8_1's +// epilogue). CVT 0: four XMAD chains and I2F; 1: two chains seeded with kMagicI (exact for the 16-value sums, +// |sum| <= 16 * 43 * 127 < 2^22) +template +__device__ __forceinline__ void mac16(const int (&wv)[16], const float dx0, const float dx1, const int s, + const int (*act_s)[4][16], const float (*dy_s)[4], const int nt, float (&acc)[TPP]) { +#pragma unroll + for (int t0 = 0; t0 < TPP; t0 += TG) { + if (t0 < nt) { +#pragma unroll + for (int t = t0; t < t0 + TG; ++t) { + const int4* a4 = (const int4*) act_s[t][s]; + const int4 a0 = a4[0], a1 = a4[1], a2 = a4[2], a3 = a4[3]; + const float dy = dy_s[t][s]; + float fa, fb; // float(sum of values 0-15), float(sum of values 16-31): exact + if constexpr (CVT == 0) { + int s0 = 0, s1 = 0, s2 = 0, s3 = 0; + s0 = mad2(wv[0], a0.x, s0); s1 = mad2(wv[1], a0.y, s1); s0 = mad2(wv[2], a0.z, s0); s1 = mad2(wv[3], a0.w, s1); + s0 = mad2(wv[4], a1.x, s0); s1 = mad2(wv[5], a1.y, s1); s0 = mad2(wv[6], a1.z, s0); s1 = mad2(wv[7], a1.w, s1); + s2 = mad2(wv[8], a2.x, s2); s3 = mad2(wv[9], a2.y, s3); s2 = mad2(wv[10], a2.z, s2); s3 = mad2(wv[11], a2.w, s3); + s2 = mad2(wv[12], a3.x, s2); s3 = mad2(wv[13], a3.y, s3); s2 = mad2(wv[14], a3.z, s2); s3 = mad2(wv[15], a3.w, s3); + fa = __int2float_rn(s0 + s1); + fb = __int2float_rn(s2 + s3); + } else { + int sa = kMagicI, sb = kMagicI; + sa = mad2(wv[0], a0.x, sa); sb = mad2(wv[8], a2.x, sb); + sa = mad2(wv[1], a0.y, sa); sb = mad2(wv[9], a2.y, sb); + sa = mad2(wv[2], a0.z, sa); sb = mad2(wv[10], a2.z, sb); + sa = mad2(wv[3], a0.w, sa); sb = mad2(wv[11], a2.w, sb); + sa = mad2(wv[4], a1.x, sa); sb = mad2(wv[12], a3.x, sb); + sa = mad2(wv[5], a1.y, sa); sb = mad2(wv[13], a3.y, sb); + sa = mad2(wv[6], a1.z, sa); sb = mad2(wv[14], a3.z, sb); + sa = mad2(wv[7], a1.w, sa); sb = mad2(wv[15], a3.w, sb); + fa = __fsub_rn(__int_as_float(sa), kMagicF); + fb = __fsub_rn(__int_as_float(sb), kMagicF); + } + float tf = __fmaf_rn(fa, dx0, 0.0f); + tf = __fmaf_rn(fb, dx1, tf); + acc[t] = __fmaf_rn(tf, dy, acc[t]); + } + } + } +} + +// d0 / d1 = ((ls & 15 | ls >> 4) * d + d / 2) / 4 as MMQ's sm_60 SASS computes them +__device__ __forceinline__ void scales16(const int ls, const int dbits, float& d0, float& d1) { + const float d = __half2float(__ushort_as_half((unsigned short) (dbits & 0xFFFF))); + const float dh = __fmul_rn(d, 0.5f); + d0 = __fmul_rn(__fmaf_rn(d, __int2float_rn(ls & 0x0F), dh), 0.25f); + d1 = __fmul_rn(__fmaf_rn(d, __int2float_rn(ls >> 4), dh), 0.25f); +} + +// ---- IQ2_XXS: 66-B blocks; per row and step 8 qs words (4 grid bytes, then aux32: 4 x 7 sign bits + the scale) and +// d (stride 9). Every lane slot is one word type: qs words rows tid/8 + 16 i (word tid%8), d the thread's own row. +// A partial pass's token groups run past the count on copies of row nt - 1 (not written); a full pass has no branch. +namespace xxs2 { +constexpr int BB = (int) sizeof(block_iq2_xxs), ROWB = NB * BB, CWS = 9; +static_assert(BB == 66 && ROWB % 4 == 0, "block_iq2_xxs"); + +// nibble b = the code of byte b of iq2xxs_grid[i] +__device__ __forceinline__ uint32_t grid_codes(const int i) { + const uint2 g = ((const uint2*) iq2xxs_grid)[i]; + uint32_t w = 0; +#pragma unroll + for (int b = 0; b < 8; ++b) { + const uint32_t v = ((b < 4 ? g.x : g.y) >> (8 * (b & 3))) & 0xFF; + w |= ((uint32_t) (v > 0x08u) + (uint32_t) (v > 0x19u)) << (4 * b); + } + return w; +} +// a 7-bit sign field: MMQ's unpack_ksigns byte (bit 7 = parity) spread to bit 2 of nibble b +__device__ __forceinline__ uint32_t sign_spread(const int v) { + const uint32_t s = (uint32_t) v ^ ((uint32_t) (__popc(v) & 1) << 7); + uint32_t w = 0; +#pragma unroll + for (int b = 0; b < 8; ++b) w |= ((s >> b) & 1u) << (4 * b + 2); + return w; +} + +template +__device__ __forceinline__ void sb_compute(const int* __restrict__ cw, const int s, const uint32_t* __restrict__ tabA, + const uint32_t* __restrict__ tabB, const int (*act_s)[4][16], + const float (*dy_s)[4], const int nt, float (&acc)[TPP]) { + int wv[16]; + float dx; + { + const uint32_t q2 = (uint32_t) cw[2 * s + 0]; + const uint32_t aux32 = (uint32_t) cw[2 * s + 1]; +#pragma unroll + for (int l = 0; l < QR2_XXS; ++l) { + const uint32_t sel = tabA[(q2 >> (8 * l)) & 0xFF] | tabB[(aux32 >> (7 * l)) & 0x7F]; + split2((int) __byte_perm(PX, PY, sel), wv[4 * l + 0], wv[4 * l + 1]); + split2((int) __byte_perm(PX, PY, sel >> 16), wv[4 * l + 2], wv[4 * l + 3]); + } + const int ls = aux32 >> 27 | 1; + const float d = __ushort_as_half((unsigned short) (cw[8] & 0xFFFF)); + dx = d * ls / 8; // MMQ: (d * scale + d / 2) / 4 + } +#pragma unroll + for (int t0 = 0; t0 < TPP; t0 += TG) { + if (FULL || t0 < nt) { +#pragma unroll + for (int t = t0; t < t0 + TG; ++t) { + const int4* a4 = (const int4*) act_s[t][s]; + const int4 a0 = a4[0], a1 = a4[1], a2 = a4[2], a3 = a4[3]; + const int a[16] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w, a2.x, a2.y, a2.z, a2.w, a3.x, a3.y, a3.z, a3.w}; + const float dy = dy_s[t][s]; + int sc[CH]; + sc[0] = kMagicI; +#pragma unroll + for (int c = 1; c < CH; ++c) sc[c] = 0; +#pragma unroll + for (int m = 0; m < 16; ++m) sc[m % CH] = mad2(wv[m], a[m], sc[m % CH]); + int sum = sc[0]; +#pragma unroll + for (int c = 1; c < CH; ++c) sum += sc[c]; + acc[t] = __fmaf_rn(__fmul_rn(dx, dy), __fsub_rn(__int_as_float(sum), kMagicF), acc[t]); + } + } + } +} + +template +__global__ void __launch_bounds__(NT, 1) gu_kernel(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const block_q8_1_mmq* __restrict__ xq, + const int total_rows, float* __restrict__ dst) { + static_assert(TPP % TG == 0 && (TPP * 32) % NT == 0 && TPP * 4 <= NT, ""); + constexpr int NDS = 8; // qs-word slots per lane: rows tid/8 + i*NT/8, word tid%8 + constexpr int RS = NT / 8; // rows per slot + constexpr int NA = TPP * 32 / NT; // activation qs-word slots per lane + __shared__ __align__(16) uint32_t tabA[256]; + __shared__ __align__(16) uint32_t tabB[128]; + __shared__ __align__(16) int act_s[TPP][4][16]; // [token][32-block of the step][8 words split to 16] + __shared__ float dy_s[TPP][4]; + __shared__ int wsm[NT * CWS]; + const int tid = threadIdx.x; + for (int i = tid; i < 256; i += NT) tabA[i] = grid_codes(i); + for (int i = tid; i < 128; i += NT) tabB[i] = sign_spread(i); + int r0; + const Chunk c = cta_work(chunks, r0); + const char* __restrict__ W = w + xoff[c.j] + (size_t) r0 * ROWB; + const char* __restrict__ Wq = W + (size_t) (tid >> 3) * ROWB + 4 * (tid & 7); // + i*RS*ROWB + kc*BB + 32h + const char* __restrict__ Wd = W + (size_t) tid * ROWB; // + kc*BB + int* __restrict__ sq = wsm + (tid >> 3) * CWS + (tid & 7); // + i*RS*CWS + int* __restrict__ sd = wsm + tid * CWS + 8; + const int* cw = wsm + tid * CWS; + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int hc = 0; hc < NSTEP; ++hc) { + const int kc = hc >> 1, h = hc & 1; + const int* xr = (const int*) (xq + (size_t) hc * total_rows + trow0); // nt consecutive 36-word blocks + __syncthreads(); + int av[NA], dv = 0; +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = tid + i * NT, t = idx >> 5; + av[i] = __ldg(xr + min(t, nt - 1) * 36 + 4 + (idx & 31)); // clamped: no branch, never past row nt-1 + } + if (TPP * 4 == NT || tid < TPP * 4) dv = __ldg(xr + min(tid >> 2, nt - 1) * 36 + (tid & 3)); + const char* bq = Wq + kc * BB + 32 * h; + if (kc & 1) { // odd block: qs words aligned; d = high half of the word 2 bytes before the block + int v[NDS]; +#pragma unroll + for (int i = 0; i < NDS; ++i) v[i] = __ldg((const int*) (bq + 2 + (size_t) i * RS * ROWB)); + const int vd = __ldg((const int*) (Wd + kc * BB - 2)); +#pragma unroll + for (int i = 0; i < NDS; ++i) sq[i * RS * CWS] = v[i]; + *sd = (int) ((uint32_t) vd >> 16); + } else { // even block: qs words 2 bytes into a word: two aligned words + funnel shift; d = low half + int lo[NDS], hi[NDS]; +#pragma unroll + for (int i = 0; i < NDS; ++i) { + lo[i] = __ldg((const int*) (bq + (size_t) i * RS * ROWB)); + hi[i] = __ldg((const int*) (bq + 4 + (size_t) i * RS * ROWB)); + } + const int vd = __ldg((const int*) (Wd + kc * BB)); +#pragma unroll + for (int i = 0; i < NDS; ++i) sq[i * RS * CWS] = (int) __funnelshift_r((uint32_t) lo[i], (uint32_t) hi[i], 16); + *sd = vd; + } +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = tid + i * NT, t = idx >> 5, wd = idx & 31; + int lo, hi; + split2(av[i], lo, hi); + *(int2*) &act_s[t][wd >> 3][2 * (wd & 7)] = make_int2(lo, hi); + } + if (TPP * 4 == NT || tid < TPP * 4) dy_s[tid >> 2][tid & 3] = __int_as_float(dv); + __syncthreads(); + if (nt == TPP) { +#pragma unroll 1 + for (int s = 0; s < 4; ++s) sb_compute(cw, s, tabA, tabB, act_s, dy_s, nt, acc); + } else { +#pragma unroll 1 + for (int s = 0; s < 4; ++s) sb_compute(cw, s, tabA, tabB, act_s, dy_s, nt, acc); + } + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * GU_ROWS + r0 + tid] = acc[t]; + } +} +} // namespace xxs2 + +// ---- IQ2_XS: 74-B blocks; per row and step the qs u16 [16h, 16h + 16) (8 words), the scales [4h, 4h + 4) and d: +// 10 words, row stride 11, staged by lanes 0-9 / 16-25 of each warp (two rows per warp and k). The grid codes and +// ksigns' spread signs (q2 & 0x1FF, q2 >> 9) are tables generated from ggml's iq2xs_grid. +namespace xs2 { +constexpr int BB = 74, ROWB = NB * BB, CW = 10, CWS = 11, KR = NT / (2 * NW); + +__device__ const uint32_t kIq2xsCode[512] = { + 0x00000000u, 0x00000002u, 0x00000011u, 0x00000020u, 0x00000022u, 0x00000101u, 0x00000110u, 0x00000112u, + 0x00000121u, 0x00000200u, 0x00000202u, 0x00000211u, 0x00000220u, 0x00001001u, 0x00001010u, 0x00001012u, + 0x00001021u, 0x00001100u, 0x00001102u, 0x00001111u, 0x00001120u, 0x00001201u, 0x00001210u, 0x00002000u, + 0x00002002u, 0x00002011u, 0x00002020u, 0x00002101u, 0x00002110u, 0x00002121u, 0x00002200u, 0x00010001u, + 0x00010010u, 0x00010012u, 0x00010021u, 0x00010100u, 0x00010102u, 0x00010111u, 0x00010120u, 0x00010122u, + 0x00010201u, 0x00010210u, 0x00011000u, 0x00011002u, 0x00011011u, 0x00011020u, 0x00011101u, 0x00011110u, + 0x00011200u, 0x00011220u, 0x00012001u, 0x00012010u, 0x00012100u, 0x00020000u, 0x00020002u, 0x00020011u, + 0x00020020u, 0x00020101u, 0x00020110u, 0x00020200u, 0x00021001u, 0x00021010u, 0x00021100u, 0x00021111u, + 0x00022000u, 0x00022022u, 0x00100001u, 0x00100010u, 0x00100012u, 0x00100021u, 0x00100100u, 0x00100102u, + 0x00100111u, 0x00100120u, 0x00100201u, 0x00100210u, 0x00101000u, 0x00101002u, 0x00101011u, 0x00101020u, + 0x00101101u, 0x00101110u, 0x00101112u, 0x00101200u, 0x00102001u, 0x00102010u, 0x00102100u, 0x00110000u, + 0x00110002u, 0x00110011u, 0x00110020u, 0x00110101u, 0x00110110u, 0x00110200u, 0x00111001u, 0x00111010u, + 0x00111100u, 0x00111201u, 0x00112000u, 0x00120001u, 0x00120010u, 0x00120100u, 0x00120212u, 0x00121000u, + 0x00121002u, 0x00122010u, 0x00200000u, 0x00200002u, 0x00200011u, 0x00200020u, 0x00200022u, 0x00200101u, + 0x00200110u, 0x00200200u, 0x00200211u, 0x00201001u, 0x00201010u, 0x00201100u, 0x00201120u, 0x00202000u, + 0x00202200u, 0x00202222u, 0x00210001u, 0x00210010u, 0x00210100u, 0x00211000u, 0x00212001u, 0x00212021u, + 0x00220000u, 0x00220200u, 0x00220220u, 0x00222112u, 0x00222200u, 0x01000001u, 0x01000010u, 0x01000012u, + 0x01000021u, 0x01000100u, 0x01000102u, 0x01000111u, 0x01000120u, 0x01000201u, 0x01000210u, 0x01001000u, + 0x01001002u, 0x01001011u, 0x01001020u, 0x01001101u, 0x01001110u, 0x01001200u, 0x01001222u, 0x01002001u, + 0x01002010u, 0x01002100u, 0x01010000u, 0x01010002u, 0x01010011u, 0x01010020u, 0x01010101u, 0x01010110u, + 0x01010200u, 0x01011001u, 0x01011010u, 0x01011100u, 0x01012000u, 0x01012110u, 0x01012112u, 0x01020001u, + 0x01020010u, 0x01020012u, 0x01020100u, 0x01021000u, 0x01021200u, 0x01100000u, 0x01100002u, 0x01100011u, + 0x01100020u, 0x01100101u, 0x01100110u, 0x01100200u, 0x01101001u, 0x01101010u, 0x01101021u, 0x01101100u, + 0x01101210u, 0x01102000u, 0x01110001u, 0x01110010u, 0x01110100u, 0x01111000u, 0x01120000u, 0x01120110u, + 0x01121021u, 0x01200001u, 0x01200010u, 0x01200100u, 0x01200102u, 0x01201000u, 0x01201110u, 0x01202012u, + 0x01210000u, 0x01210011u, 0x01211212u, 0x01221101u, 0x01222221u, 0x02000000u, 0x02000002u, 0x02000011u, + 0x02000020u, 0x02000022u, 0x02000101u, 0x02000110u, 0x02000200u, 0x02001001u, 0x02001010u, 0x02001100u, + 0x02002000u, 0x02002200u, 0x02010001u, 0x02010010u, 0x02010100u, 0x02011000u, 0x02011020u, 0x02011211u, + 0x02020000u, 0x02020202u, 0x02022000u, 0x02022220u, 0x02100001u, 0x02100010u, 0x02100100u, 0x02100221u, + 0x02101000u, 0x02110000u, 0x02111001u, 0x02111102u, 0x02112121u, 0x02120001u, 0x02120122u, 0x02122212u, + 0x02200000u, 0x02200020u, 0x02200022u, 0x02200200u, 0x02201111u, 0x02202020u, 0x02202202u, 0x02211220u, + 0x02212100u, 0x02220020u, 0x02220200u, 0x02222002u, 0x02222020u, 0x02222022u, 0x10000001u, 0x10000010u, + 0x10000012u, 0x10000021u, 0x10000100u, 0x10000102u, 0x10000111u, 0x10000120u, 0x10000201u, 0x10000210u, + 0x10001000u, 0x10001002u, 0x10001011u, 0x10001020u, 0x10001022u, 0x10001101u, 0x10001110u, 0x10001200u, + 0x10001211u, 0x10002001u, 0x10002010u, 0x10002100u, 0x10010000u, 0x10010002u, 0x10010011u, 0x10010020u, + 0x10010101u, 0x10010110u, 0x10010200u, 0x10011001u, 0x10011010u, 0x10011100u, 0x10012000u, 0x10012011u, + 0x10012202u, 0x10020001u, 0x10020010u, 0x10020100u, 0x10020102u, 0x10020221u, 0x10021000u, 0x10100000u, + 0x10100002u, 0x10100011u, 0x10100020u, 0x10100101u, 0x10100110u, 0x10100121u, 0x10100200u, 0x10101001u, + 0x10101010u, 0x10101100u, 0x10102000u, 0x10102110u, 0x10110001u, 0x10110010u, 0x10110100u, 0x10110210u, + 0x10111000u, 0x10112122u, 0x10120000u, 0x10120022u, 0x10121010u, 0x10121100u, 0x10200001u, 0x10200010u, + 0x10200100u, 0x10201000u, 0x10201011u, 0x10201110u, 0x10201202u, 0x10210000u, 0x10210101u, 0x10211010u, + 0x10211100u, 0x10211221u, 0x10220010u, 0x11000000u, 0x11000002u, 0x11000011u, 0x11000020u, 0x11000101u, + 0x11000110u, 0x11000200u, 0x11000220u, 0x11001001u, 0x11001010u, 0x11001100u, 0x11002000u, 0x11010001u, + 0x11010010u, 0x11010100u, 0x11010111u, 0x11011000u, 0x11011002u, 0x11020000u, 0x11021010u, 0x11022222u, + 0x11100001u, 0x11100010u, 0x11100100u, 0x11100201u, 0x11101000u, 0x11101200u, 0x11102001u, 0x11102201u, + 0x11110000u, 0x11110020u, 0x11112000u, 0x11112020u, 0x11120201u, 0x11121220u, 0x11122201u, 0x11200000u, + 0x11200110u, 0x11201001u, 0x11201100u, 0x11202121u, 0x11210122u, 0x11211000u, 0x11211002u, 0x11222011u, + 0x12000001u, 0x12000010u, 0x12000100u, 0x12001000u, 0x12001110u, 0x12001202u, 0x12002012u, 0x12002221u, + 0x12010000u, 0x12020210u, 0x12021022u, 0x12022102u, 0x12100000u, 0x12100112u, 0x12110100u, 0x12111000u, + 0x12111011u, 0x12112210u, 0x12200001u, 0x12201222u, 0x12210211u, 0x12220012u, 0x12221110u, 0x12221202u, + 0x20000000u, 0x20000002u, 0x20000011u, 0x20000020u, 0x20000101u, 0x20000110u, 0x20000200u, 0x20000222u, + 0x20001001u, 0x20001010u, 0x20001100u, 0x20002000u, 0x20002002u, 0x20002220u, 0x20002222u, 0x20010001u, + 0x20010010u, 0x20010012u, 0x20010100u, 0x20011000u, 0x20011101u, 0x20011121u, 0x20020000u, 0x20020200u, + 0x20022000u, 0x20022002u, 0x20022200u, 0x20022220u, 0x20100001u, 0x20100010u, 0x20100100u, 0x20100102u, + 0x20100111u, 0x20101000u, 0x20101200u, 0x20102021u, 0x20110000u, 0x20111010u, 0x20112211u, 0x20120120u, + 0x20121222u, 0x20200000u, 0x20200020u, 0x20200211u, 0x20201122u, 0x20202000u, 0x20202002u, 0x20202220u, + 0x20210012u, 0x20220202u, 0x20222000u, 0x20222020u, 0x20222112u, 0x20222220u, 0x21000001u, 0x21000010u, + 0x21000100u, 0x21001000u, 0x21001112u, 0x21002010u, 0x21010000u, 0x21010202u, 0x21011210u, 0x21021112u, + 0x21022021u, 0x21100000u, 0x21100011u, 0x21101010u, 0x21101100u, 0x21101120u, 0x21110221u, 0x21112100u, + 0x21112102u, 0x21121001u, 0x21201101u, 0x21202212u, 0x21211021u, 0x21220111u, 0x21221200u, 0x22000000u, + 0x22000002u, 0x22000020u, 0x22000022u, 0x22000200u, 0x22000222u, 0x22002200u, 0x22011101u, 0x22011121u, + 0x22012212u, 0x22020000u, 0x22020002u, 0x22020020u, 0x22020222u, 0x22022000u, 0x22022200u, 0x22101000u, + 0x22102111u, 0x22121211u, 0x22122120u, 0x22200022u, 0x22200200u, 0x22200202u, 0x22200220u, 0x22202200u, + 0x22202220u, 0x22210010u, 0x22212010u, 0x22212012u, 0x22220220u, 0x22220222u, 0x22222101u, 0x22222222u, +}; +__device__ const uint32_t kIq2xsSign[128] = { + 0x00000000u, 0x40000004u, 0x40000040u, 0x00000044u, 0x40000400u, 0x00000404u, 0x00000440u, 0x40000444u, + 0x40004000u, 0x00004004u, 0x00004040u, 0x40004044u, 0x00004400u, 0x40004404u, 0x40004440u, 0x00004444u, + 0x40040000u, 0x00040004u, 0x00040040u, 0x40040044u, 0x00040400u, 0x40040404u, 0x40040440u, 0x00040444u, + 0x00044000u, 0x40044004u, 0x40044040u, 0x00044044u, 0x40044400u, 0x00044404u, 0x00044440u, 0x40044444u, + 0x40400000u, 0x00400004u, 0x00400040u, 0x40400044u, 0x00400400u, 0x40400404u, 0x40400440u, 0x00400444u, + 0x00404000u, 0x40404004u, 0x40404040u, 0x00404044u, 0x40404400u, 0x00404404u, 0x00404440u, 0x40404444u, + 0x00440000u, 0x40440004u, 0x40440040u, 0x00440044u, 0x40440400u, 0x00440404u, 0x00440440u, 0x40440444u, + 0x40444000u, 0x00444004u, 0x00444040u, 0x40444044u, 0x00444400u, 0x40444404u, 0x40444440u, 0x00444444u, + 0x44000000u, 0x04000004u, 0x04000040u, 0x44000044u, 0x04000400u, 0x44000404u, 0x44000440u, 0x04000444u, + 0x04004000u, 0x44004004u, 0x44004040u, 0x04004044u, 0x44004400u, 0x04004404u, 0x04004440u, 0x44004444u, + 0x04040000u, 0x44040004u, 0x44040040u, 0x04040044u, 0x44040400u, 0x04040404u, 0x04040440u, 0x44040444u, + 0x44044000u, 0x04044004u, 0x04044040u, 0x44044044u, 0x04044400u, 0x44044404u, 0x44044440u, 0x04044444u, + 0x04400000u, 0x44400004u, 0x44400040u, 0x04400044u, 0x44400400u, 0x04400404u, 0x04400440u, 0x44400444u, + 0x44404000u, 0x04404004u, 0x04404040u, 0x44404044u, 0x04404400u, 0x44404404u, 0x44404440u, 0x04404444u, + 0x44440000u, 0x04440004u, 0x04440040u, 0x44440044u, 0x04440400u, 0x44440404u, 0x44440440u, 0x04440444u, + 0x04444000u, 0x44444004u, 0x44444040u, 0x04444044u, 0x44444400u, 0x04444404u, 0x04444440u, 0x44444444u, +}; + +template +__global__ void __launch_bounds__(NT, 1) gu_kernel(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const block_q8_1_mmq* __restrict__ xq, + const int total_rows, float* __restrict__ dst) { + static_assert(TPP % TG == 0, ""); + __shared__ uint32_t tabA_s[512]; + __shared__ uint32_t tabB_s[128]; + __shared__ __align__(16) int act_s[TPP][4][16]; // [token][sub-block][8 int8x4 words split to 16 int16 pairs] + __shared__ float dy_s[TPP][4]; + __shared__ int wsm[NT * CWS]; + for (int i = threadIdx.x; i < 512; i += NT) tabA_s[i] = kIq2xsCode[i]; + for (int i = threadIdx.x; i < 128; i += NT) tabB_s[i] = kIq2xsSign[i]; + int r0; + const Chunk c = cta_work(chunks, r0); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int ci = lane & 15, rsub = lane >> 4; // lanes 0-9 stage one row, 16-25 the next; ci = the compact word + const bool cact = ci < CW; + const char* __restrict__ Wr = w + xoff[c.j] + (size_t) (r0 + 2 * warp + rsub) * ROWB; // + k * 2 NW rows + int* __restrict__ wsl = wsm + (2 * warp + rsub) * CWS + ci; // + k * 2 NW * CWS + const int* cw = wsm + threadIdx.x * CWS; + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int hc = 0; hc < NSTEP; ++hc) { + const int kc = hc >> 1, h = hc & 1; + // compact word ci's byte offset inside the block: qs words (ci 0-7), scales (ci 8); d (ci 9) apart + const int boff = ci < 8 ? 2 + 32 * h + 4 * ci : 66 + 4 * h; + __syncthreads(); + if (cact) { + const char* blk0 = Wr + kc * BB; + if (kc & 1) { // odd block: words aligned, one load each; d = high half of the word at blk - 2 + const int o = ci < 9 ? boff : -2, sh = ci < 9 ? 0 : 16; + int v[KR]; +#pragma unroll + for (int k = 0; k < KR; ++k) v[k] = __ldg((const int*) (blk0 + (size_t) k * 2 * NW * ROWB + o)); +#pragma unroll + for (int k = 0; k < KR; ++k) wsl[k * 2 * NW * CWS] = __funnelshift_r(v[k], v[k], sh); + } else { // even block: words 2 bytes into a word, two loads + funnel shift; d = low half at blk + const int o0 = ci < 9 ? boff - 2 : 0, o1 = ci < 9 ? boff + 2 : 0, sh = ci < 9 ? 16 : 0; + int v0[KR], v1[KR]; +#pragma unroll + for (int k = 0; k < KR; ++k) { + const char* b = blk0 + (size_t) k * 2 * NW * ROWB; + v0[k] = __ldg((const int*) (b + o0)); + v1[k] = __ldg((const int*) (b + o1)); + } +#pragma unroll + for (int k = 0; k < KR; ++k) wsl[k * 2 * NW * CWS] = __funnelshift_r(v0[k], v1[k], sh); + } + } + stage_act(xq, hc, total_rows, trow0, nt, act_s, dy_s); + __syncthreads(); +#pragma unroll 1 + for (int s = 0; s < 4; ++s) { // cw: [0..7] qs u16 pairs (sub-block s = words 2s, 2s+1), [8] scales, [9] d + int wv[16]; + float dx0, dx1; +#pragma unroll + for (int l = 0; l < 4; ++l) { + const uint32_t q2 = ((uint32_t) cw[2 * s + (l >> 1)] >> (16 * (l & 1))) & 0xFFFFu; + const uint32_t sel = tabA_s[q2 & 0x1FFu] | tabB_s[q2 >> 9]; + split2((int) __byte_perm(PX, PY, sel), wv[4 * l + 0], wv[4 * l + 1]); + split2((int) __byte_perm(PX, PY, sel >> 16), wv[4 * l + 2], wv[4 * l + 3]); + } + scales16((cw[8] >> (8 * s)) & 0xFF, cw[9], dx0, dx1); + mac16(wv, dx0, dx1, s, act_s, dy_s, nt, acc); + } + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * GU_ROWS + r0 + threadIdx.x] = acc[t]; + } +} +} // namespace xs2 + +// ---- IQ2_S: 82-B blocks; per row and step the grid bytes qs[16h..] (4 words), the sign bytes qs[32 + 16h..] (4), +// qh[4h..4h + 3], scales[4h..4h + 3] and d: 11 words, row stride 11, staged by lanes 0-10 / 16-26 of each warp. The grid +// index is 10 bits (qs | qh bits), the sign byte is explicit: tables generated from ggml's iq2s_grid. +namespace s2 { +constexpr int BB = (int) sizeof(block_iq2_s), ROWB = NB * BB, CW = 11, CWS = 11, KR = NT / (2 * NW); +static_assert(BB == 82, "block_iq2_s"); + +__device__ const uint32_t kIq2sCode[1024] = { + 0x00000000u, 0x00000002u, 0x00000011u, 0x00000020u, 0x00000022u, 0x00000101u, 0x00000110u, 0x00000112u, + 0x00000121u, 0x00000200u, 0x00000202u, 0x00000211u, 0x00000220u, 0x00001001u, 0x00001010u, 0x00001012u, + 0x00001021u, 0x00001100u, 0x00001102u, 0x00001111u, 0x00001120u, 0x00001201u, 0x00001210u, 0x00001212u, + 0x00001221u, 0x00002000u, 0x00002002u, 0x00002011u, 0x00002020u, 0x00002101u, 0x00002110u, 0x00002200u, + 0x00002211u, 0x00002222u, 0x00010001u, 0x00010010u, 0x00010012u, 0x00010021u, 0x00010100u, 0x00010102u, + 0x00010111u, 0x00010120u, 0x00010201u, 0x00010210u, 0x00011000u, 0x00011002u, 0x00011011u, 0x00011020u, + 0x00011101u, 0x00011110u, 0x00011112u, 0x00011121u, 0x00011200u, 0x00011211u, 0x00011220u, 0x00012001u, + 0x00012010u, 0x00012100u, 0x00012102u, 0x00012111u, 0x00012201u, 0x00012210u, 0x00020000u, 0x00020002u, + 0x00020011u, 0x00020020u, 0x00020101u, 0x00020110u, 0x00020200u, 0x00020222u, 0x00021001u, 0x00021010u, + 0x00021012u, 0x00021021u, 0x00021100u, 0x00021111u, 0x00022000u, 0x00022011u, 0x00022022u, 0x00022110u, + 0x00022202u, 0x00100001u, 0x00100010u, 0x00100012u, 0x00100021u, 0x00100100u, 0x00100102u, 0x00100111u, + 0x00100120u, 0x00100201u, 0x00100210u, 0x00100212u, 0x00100221u, 0x00101000u, 0x00101002u, 0x00101011u, + 0x00101020u, 0x00101022u, 0x00101101u, 0x00101110u, 0x00101112u, 0x00101121u, 0x00101200u, 0x00101202u, + 0x00101211u, 0x00102001u, 0x00102010u, 0x00102012u, 0x00102021u, 0x00102100u, 0x00102111u, 0x00102120u, + 0x00102201u, 0x00102210u, 0x00110000u, 0x00110002u, 0x00110011u, 0x00110020u, 0x00110022u, 0x00110101u, + 0x00110110u, 0x00110112u, 0x00110121u, 0x00110200u, 0x00110211u, 0x00110220u, 0x00111001u, 0x00111010u, + 0x00111012u, 0x00111021u, 0x00111100u, 0x00111102u, 0x00111111u, 0x00111120u, 0x00111201u, 0x00111210u, + 0x00112000u, 0x00112002u, 0x00112011u, 0x00112020u, 0x00112101u, 0x00112110u, 0x00112200u, 0x00120001u, + 0x00120010u, 0x00120012u, 0x00120021u, 0x00120100u, 0x00120111u, 0x00121000u, 0x00121011u, 0x00121020u, + 0x00121101u, 0x00121110u, 0x00121200u, 0x00122001u, 0x00122010u, 0x00122100u, 0x00200000u, 0x00200002u, + 0x00200011u, 0x00200020u, 0x00200101u, 0x00200110u, 0x00200112u, 0x00200121u, 0x00200200u, 0x00200211u, + 0x00200222u, 0x00201001u, 0x00201010u, 0x00201012u, 0x00201021u, 0x00201100u, 0x00201102u, 0x00201111u, + 0x00201120u, 0x00201201u, 0x00201210u, 0x00202000u, 0x00202011u, 0x00202110u, 0x00202222u, 0x00210001u, + 0x00210010u, 0x00210100u, 0x00210102u, 0x00210111u, 0x00210120u, 0x00210201u, 0x00211000u, 0x00211011u, + 0x00211020u, 0x00211101u, 0x00211110u, 0x00211200u, 0x00212001u, 0x00212100u, 0x00220000u, 0x00220101u, + 0x00220110u, 0x00220202u, 0x00220220u, 0x00220222u, 0x00221100u, 0x00222121u, 0x01000001u, 0x01000010u, + 0x01000012u, 0x01000021u, 0x01000100u, 0x01000102u, 0x01000111u, 0x01000120u, 0x01000201u, 0x01000210u, + 0x01000212u, 0x01001000u, 0x01001002u, 0x01001011u, 0x01001020u, 0x01001101u, 0x01001110u, 0x01001112u, + 0x01001121u, 0x01001200u, 0x01001202u, 0x01001211u, 0x01001220u, 0x01002001u, 0x01002010u, 0x01002012u, + 0x01002100u, 0x01002111u, 0x01002120u, 0x01002201u, 0x01002210u, 0x01010000u, 0x01010002u, 0x01010011u, + 0x01010020u, 0x01010022u, 0x01010101u, 0x01010110u, 0x01010112u, 0x01010121u, 0x01010200u, 0x01010202u, + 0x01010211u, 0x01010220u, 0x01011001u, 0x01011010u, 0x01011012u, 0x01011021u, 0x01011100u, 0x01011102u, + 0x01011111u, 0x01011120u, 0x01011201u, 0x01011210u, 0x01012000u, 0x01012002u, 0x01012011u, 0x01012020u, + 0x01012101u, 0x01012110u, 0x01020001u, 0x01020010u, 0x01020021u, 0x01020100u, 0x01020111u, 0x01020201u, + 0x01020210u, 0x01021000u, 0x01021011u, 0x01021101u, 0x01021110u, 0x01022001u, 0x01022010u, 0x01022100u, + 0x01100000u, 0x01100002u, 0x01100011u, 0x01100020u, 0x01100101u, 0x01100110u, 0x01100112u, 0x01100121u, + 0x01100200u, 0x01100211u, 0x01100220u, 0x01101001u, 0x01101010u, 0x01101012u, 0x01101021u, 0x01101100u, + 0x01101102u, 0x01101111u, 0x01101120u, 0x01101201u, 0x01101210u, 0x01102000u, 0x01102002u, 0x01102011u, + 0x01102020u, 0x01102101u, 0x01102110u, 0x01102200u, 0x01110001u, 0x01110010u, 0x01110012u, 0x01110021u, + 0x01110100u, 0x01110102u, 0x01110111u, 0x01110120u, 0x01110201u, 0x01110210u, 0x01111000u, 0x01111002u, + 0x01111011u, 0x01111020u, 0x01111101u, 0x01111110u, 0x01111200u, 0x01112001u, 0x01112010u, 0x01112100u, + 0x01120000u, 0x01120011u, 0x01120020u, 0x01120101u, 0x01120110u, 0x01120200u, 0x01121001u, 0x01121010u, + 0x01121100u, 0x01122000u, 0x01122222u, 0x01200001u, 0x01200010u, 0x01200012u, 0x01200021u, 0x01200100u, + 0x01200111u, 0x01200120u, 0x01200201u, 0x01201000u, 0x01201002u, 0x01201011u, 0x01201020u, 0x01201101u, + 0x01201110u, 0x01201200u, 0x01202001u, 0x01202010u, 0x01210000u, 0x01210002u, 0x01210011u, 0x01210020u, + 0x01210101u, 0x01210110u, 0x01210200u, 0x01211001u, 0x01211010u, 0x01211100u, 0x01211221u, 0x01212202u, + 0x01220010u, 0x01220100u, 0x01221000u, 0x01221112u, 0x02000000u, 0x02000002u, 0x02000011u, 0x02000020u, + 0x02000101u, 0x02000110u, 0x02000112u, 0x02000121u, 0x02000200u, 0x02000211u, 0x02000222u, 0x02001001u, + 0x02001010u, 0x02001100u, 0x02001102u, 0x02001111u, 0x02001210u, 0x02002000u, 0x02002022u, 0x02002110u, + 0x02002222u, 0x02010001u, 0x02010010u, 0x02010100u, 0x02010102u, 0x02010111u, 0x02010201u, 0x02011000u, + 0x02011002u, 0x02011011u, 0x02011101u, 0x02011110u, 0x02011200u, 0x02012001u, 0x02012010u, 0x02012100u, + 0x02020000u, 0x02020022u, 0x02020202u, 0x02020220u, 0x02020222u, 0x02021010u, 0x02021100u, 0x02022020u, + 0x02022022u, 0x02022220u, 0x02100001u, 0x02100010u, 0x02100012u, 0x02100021u, 0x02100100u, 0x02100111u, + 0x02100120u, 0x02100201u, 0x02100210u, 0x02101000u, 0x02101002u, 0x02101011u, 0x02101020u, 0x02101101u, + 0x02101110u, 0x02101200u, 0x02102001u, 0x02102010u, 0x02102100u, 0x02110000u, 0x02110011u, 0x02110020u, + 0x02110101u, 0x02110110u, 0x02110200u, 0x02111001u, 0x02111010u, 0x02111100u, 0x02111212u, 0x02112000u, + 0x02120001u, 0x02120010u, 0x02120100u, 0x02121000u, 0x02121121u, 0x02200000u, 0x02200011u, 0x02200101u, + 0x02200110u, 0x02201001u, 0x02201010u, 0x02201100u, 0x02202022u, 0x02202222u, 0x02210001u, 0x02210010u, + 0x02210100u, 0x02212111u, 0x02220022u, 0x02220202u, 0x02221210u, 0x02222020u, 0x02222022u, 0x10000001u, + 0x10000010u, 0x10000012u, 0x10000021u, 0x10000100u, 0x10000102u, 0x10000111u, 0x10000120u, 0x10000122u, + 0x10000201u, 0x10000210u, 0x10000212u, 0x10001000u, 0x10001002u, 0x10001011u, 0x10001020u, 0x10001022u, + 0x10001101u, 0x10001110u, 0x10001112u, 0x10001121u, 0x10001200u, 0x10001202u, 0x10001211u, 0x10002001u, + 0x10002010u, 0x10002100u, 0x10002111u, 0x10002120u, 0x10002201u, 0x10002210u, 0x10010000u, 0x10010002u, + 0x10010011u, 0x10010020u, 0x10010101u, 0x10010110u, 0x10010112u, 0x10010121u, 0x10010200u, 0x10010202u, + 0x10010211u, 0x10011001u, 0x10011010u, 0x10011012u, 0x10011021u, 0x10011100u, 0x10011102u, 0x10011111u, + 0x10011120u, 0x10011201u, 0x10011210u, 0x10012000u, 0x10012002u, 0x10012011u, 0x10012020u, 0x10012101u, + 0x10012110u, 0x10012200u, 0x10020001u, 0x10020010u, 0x10020100u, 0x10020102u, 0x10020111u, 0x10020120u, + 0x10020210u, 0x10021000u, 0x10021011u, 0x10021020u, 0x10021101u, 0x10021110u, 0x10021200u, 0x10022001u, + 0x10022010u, 0x10100000u, 0x10100002u, 0x10100011u, 0x10100020u, 0x10100022u, 0x10100101u, 0x10100110u, + 0x10100112u, 0x10100121u, 0x10100200u, 0x10100202u, 0x10100211u, 0x10100220u, 0x10101001u, 0x10101010u, + 0x10101012u, 0x10101021u, 0x10101100u, 0x10101102u, 0x10101111u, 0x10101120u, 0x10101201u, 0x10101210u, + 0x10102000u, 0x10102002u, 0x10102011u, 0x10102020u, 0x10102101u, 0x10102110u, 0x10102200u, 0x10110001u, + 0x10110010u, 0x10110012u, 0x10110021u, 0x10110100u, 0x10110102u, 0x10110111u, 0x10110120u, 0x10110201u, + 0x10110210u, 0x10111000u, 0x10111002u, 0x10111011u, 0x10111020u, 0x10111101u, 0x10111110u, 0x10111200u, + 0x10111222u, 0x10112001u, 0x10112010u, 0x10112100u, 0x10120000u, 0x10120002u, 0x10120011u, 0x10120020u, + 0x10120101u, 0x10120110u, 0x10120200u, 0x10121001u, 0x10121010u, 0x10121100u, 0x10122000u, 0x10122211u, + 0x10200001u, 0x10200010u, 0x10200021u, 0x10200100u, 0x10200102u, 0x10200111u, 0x10200120u, 0x10200201u, + 0x10200210u, 0x10201000u, 0x10201002u, 0x10201011u, 0x10201020u, 0x10201101u, 0x10201110u, 0x10201200u, + 0x10202010u, 0x10202100u, 0x10210000u, 0x10210002u, 0x10210011u, 0x10210020u, 0x10210101u, 0x10210110u, + 0x10210200u, 0x10211001u, 0x10211010u, 0x10211100u, 0x10212000u, 0x10212112u, 0x10220001u, 0x10220010u, + 0x10220100u, 0x10221000u, 0x11000000u, 0x11000002u, 0x11000011u, 0x11000020u, 0x11000101u, 0x11000110u, + 0x11000112u, 0x11000121u, 0x11000200u, 0x11000202u, 0x11000211u, 0x11000220u, 0x11001001u, 0x11001010u, + 0x11001012u, 0x11001021u, 0x11001100u, 0x11001102u, 0x11001111u, 0x11001120u, 0x11001201u, 0x11001210u, + 0x11002000u, 0x11002002u, 0x11002011u, 0x11002020u, 0x11002101u, 0x11002110u, 0x11010001u, 0x11010010u, + 0x11010012u, 0x11010021u, 0x11010100u, 0x11010102u, 0x11010111u, 0x11010120u, 0x11010201u, 0x11010210u, + 0x11011000u, 0x11011002u, 0x11011011u, 0x11011020u, 0x11011101u, 0x11011110u, 0x11011200u, 0x11012001u, + 0x11012010u, 0x11012100u, 0x11020000u, 0x11020011u, 0x11020020u, 0x11020101u, 0x11020110u, 0x11020200u, + 0x11021001u, 0x11021010u, 0x11021100u, 0x11021221u, 0x11022000u, 0x11100001u, 0x11100010u, 0x11100012u, + 0x11100021u, 0x11100100u, 0x11100102u, 0x11100111u, 0x11100120u, 0x11100201u, 0x11100210u, 0x11101000u, + 0x11101002u, 0x11101011u, 0x11101020u, 0x11101101u, 0x11101110u, 0x11101200u, 0x11102001u, 0x11102010u, + 0x11102100u, 0x11110000u, 0x11110002u, 0x11110011u, 0x11110020u, 0x11110101u, 0x11110110u, 0x11110200u, + 0x11111001u, 0x11111010u, 0x11111100u, 0x11112000u, 0x11120001u, 0x11120010u, 0x11120100u, 0x11120212u, + 0x11121000u, 0x11200000u, 0x11200002u, 0x11200011u, 0x11200020u, 0x11200101u, 0x11200110u, 0x11200200u, + 0x11201001u, 0x11201010u, 0x11201100u, 0x11201122u, 0x11202000u, 0x11210001u, 0x11210010u, 0x11210100u, + 0x11211000u, 0x11220000u, 0x11220121u, 0x11222011u, 0x11222220u, 0x12000001u, 0x12000010u, 0x12000012u, + 0x12000100u, 0x12000102u, 0x12000111u, 0x12000120u, 0x12000201u, 0x12000210u, 0x12001000u, 0x12001011u, + 0x12001020u, 0x12001101u, 0x12001110u, 0x12001200u, 0x12002010u, 0x12002100u, 0x12010000u, 0x12010002u, + 0x12010011u, 0x12010020u, 0x12010101u, 0x12010110u, 0x12010200u, 0x12011001u, 0x12011010u, 0x12011100u, + 0x12012000u, 0x12012121u, 0x12020010u, 0x12020100u, 0x12021000u, 0x12021112u, 0x12022201u, 0x12100000u, + 0x12100011u, 0x12100020u, 0x12100101u, 0x12100110u, 0x12100200u, 0x12101001u, 0x12101010u, 0x12101100u, + 0x12102000u, 0x12110001u, 0x12110010u, 0x12110100u, 0x12111000u, 0x12111022u, 0x12111220u, 0x12112102u, + 0x12120000u, 0x12122110u, 0x12200001u, 0x12200010u, 0x12200100u, 0x12201211u, 0x12202120u, 0x12210000u, + 0x12210222u, 0x12221002u, 0x12222201u, 0x20000000u, 0x20000002u, 0x20000011u, 0x20000020u, 0x20000101u, + 0x20000110u, 0x20000121u, 0x20000200u, 0x20000211u, 0x20001001u, 0x20001010u, 0x20001100u, 0x20001102u, + 0x20001111u, 0x20001120u, 0x20001201u, 0x20002000u, 0x20002011u, 0x20002101u, 0x20002110u, 0x20010001u, + 0x20010010u, 0x20010021u, 0x20010100u, 0x20010102u, 0x20010111u, 0x20010120u, 0x20010201u, 0x20010210u, + 0x20011000u, 0x20011002u, 0x20011011u, 0x20011020u, 0x20011101u, 0x20011110u, 0x20012001u, 0x20012010u, + 0x20012100u, 0x20012221u, 0x20020000u, 0x20020011u, 0x20020022u, 0x20020101u, 0x20020110u, 0x20021001u, + 0x20021010u, 0x20021100u, 0x20100001u, 0x20100010u, 0x20100012u, 0x20100021u, 0x20100100u, 0x20100102u, + 0x20100111u, 0x20100120u, 0x20100201u, 0x20101000u, 0x20101002u, 0x20101011u, 0x20101020u, 0x20101101u, + 0x20101110u, 0x20101200u, 0x20102001u, 0x20102010u, 0x20102100u, 0x20110000u, 0x20110002u, 0x20110011u, + 0x20110020u, 0x20110101u, 0x20110110u, 0x20110200u, 0x20111001u, 0x20111010u, 0x20111100u, 0x20112000u, + 0x20112022u, 0x20120001u, 0x20120010u, 0x20120100u, 0x20120221u, 0x20121000u, 0x20200000u, 0x20200011u, + 0x20200101u, 0x20200110u, 0x20201001u, 0x20201010u, 0x20201100u, 0x20202202u, 0x20210001u, 0x20210010u, + 0x20211000u, 0x20211211u, 0x20220202u, 0x20221120u, 0x20221122u, 0x20222002u, 0x20222202u, 0x21000001u, + 0x21000010u, 0x21000021u, 0x21000100u, 0x21000102u, 0x21000111u, 0x21000120u, 0x21000210u, 0x21001000u, + 0x21001002u, 0x21001011u, 0x21001020u, 0x21001101u, 0x21001110u, 0x21001200u, 0x21002001u, 0x21002010u, + 0x21002100u, 0x21010000u, 0x21010011u, 0x21010101u, 0x21010110u, 0x21011001u, 0x21011010u, 0x21011100u, + 0x21011122u, 0x21020001u, 0x21020010u, 0x21020100u, 0x21021000u, 0x21022212u, 0x21100000u, 0x21100002u, + 0x21100011u, 0x21100020u, 0x21100101u, 0x21100110u, 0x21100200u, 0x21101001u, 0x21101010u, 0x21101100u, + 0x21102000u, 0x21102112u, 0x21110001u, 0x21110010u, 0x21110100u, 0x21111000u, 0x21112120u, 0x21112201u, + 0x21120000u, 0x21121012u, 0x21121210u, 0x21200001u, 0x21200010u, 0x21200100u, 0x21200212u, 0x21201000u, + 0x21202221u, 0x21210000u, 0x21211021u, 0x21211102u, 0x21222100u, 0x22000000u, 0x22000011u, 0x22000022u, + 0x22000110u, 0x22000202u, 0x22000222u, 0x22001001u, 0x22001010u, 0x22001100u, 0x22002202u, 0x22002222u, + 0x22011000u, 0x22011211u, 0x22020002u, 0x22020022u, 0x22020202u, 0x22020220u, 0x22020222u, 0x22022002u, + 0x22022020u, 0x22022022u, 0x22022220u, 0x22100001u, 0x22100010u, 0x22100100u, 0x22101000u, 0x22102021u, + 0x22102210u, 0x22110000u, 0x22110121u, 0x22121101u, 0x22200022u, 0x22200220u, 0x22202202u, 0x22211110u, + 0x22212012u, 0x22220020u, 0x22220022u, 0x22220200u, 0x22220202u, 0x22220220u, 0x22222020u, 0x22222222u, +}; +__device__ const uint32_t kIq2sSign[256] = { + 0x00000000u, 0x00000004u, 0x00000040u, 0x00000044u, 0x00000400u, 0x00000404u, 0x00000440u, 0x00000444u, + 0x00004000u, 0x00004004u, 0x00004040u, 0x00004044u, 0x00004400u, 0x00004404u, 0x00004440u, 0x00004444u, + 0x00040000u, 0x00040004u, 0x00040040u, 0x00040044u, 0x00040400u, 0x00040404u, 0x00040440u, 0x00040444u, + 0x00044000u, 0x00044004u, 0x00044040u, 0x00044044u, 0x00044400u, 0x00044404u, 0x00044440u, 0x00044444u, + 0x00400000u, 0x00400004u, 0x00400040u, 0x00400044u, 0x00400400u, 0x00400404u, 0x00400440u, 0x00400444u, + 0x00404000u, 0x00404004u, 0x00404040u, 0x00404044u, 0x00404400u, 0x00404404u, 0x00404440u, 0x00404444u, + 0x00440000u, 0x00440004u, 0x00440040u, 0x00440044u, 0x00440400u, 0x00440404u, 0x00440440u, 0x00440444u, + 0x00444000u, 0x00444004u, 0x00444040u, 0x00444044u, 0x00444400u, 0x00444404u, 0x00444440u, 0x00444444u, + 0x04000000u, 0x04000004u, 0x04000040u, 0x04000044u, 0x04000400u, 0x04000404u, 0x04000440u, 0x04000444u, + 0x04004000u, 0x04004004u, 0x04004040u, 0x04004044u, 0x04004400u, 0x04004404u, 0x04004440u, 0x04004444u, + 0x04040000u, 0x04040004u, 0x04040040u, 0x04040044u, 0x04040400u, 0x04040404u, 0x04040440u, 0x04040444u, + 0x04044000u, 0x04044004u, 0x04044040u, 0x04044044u, 0x04044400u, 0x04044404u, 0x04044440u, 0x04044444u, + 0x04400000u, 0x04400004u, 0x04400040u, 0x04400044u, 0x04400400u, 0x04400404u, 0x04400440u, 0x04400444u, + 0x04404000u, 0x04404004u, 0x04404040u, 0x04404044u, 0x04404400u, 0x04404404u, 0x04404440u, 0x04404444u, + 0x04440000u, 0x04440004u, 0x04440040u, 0x04440044u, 0x04440400u, 0x04440404u, 0x04440440u, 0x04440444u, + 0x04444000u, 0x04444004u, 0x04444040u, 0x04444044u, 0x04444400u, 0x04444404u, 0x04444440u, 0x04444444u, + 0x40000000u, 0x40000004u, 0x40000040u, 0x40000044u, 0x40000400u, 0x40000404u, 0x40000440u, 0x40000444u, + 0x40004000u, 0x40004004u, 0x40004040u, 0x40004044u, 0x40004400u, 0x40004404u, 0x40004440u, 0x40004444u, + 0x40040000u, 0x40040004u, 0x40040040u, 0x40040044u, 0x40040400u, 0x40040404u, 0x40040440u, 0x40040444u, + 0x40044000u, 0x40044004u, 0x40044040u, 0x40044044u, 0x40044400u, 0x40044404u, 0x40044440u, 0x40044444u, + 0x40400000u, 0x40400004u, 0x40400040u, 0x40400044u, 0x40400400u, 0x40400404u, 0x40400440u, 0x40400444u, + 0x40404000u, 0x40404004u, 0x40404040u, 0x40404044u, 0x40404400u, 0x40404404u, 0x40404440u, 0x40404444u, + 0x40440000u, 0x40440004u, 0x40440040u, 0x40440044u, 0x40440400u, 0x40440404u, 0x40440440u, 0x40440444u, + 0x40444000u, 0x40444004u, 0x40444040u, 0x40444044u, 0x40444400u, 0x40444404u, 0x40444440u, 0x40444444u, + 0x44000000u, 0x44000004u, 0x44000040u, 0x44000044u, 0x44000400u, 0x44000404u, 0x44000440u, 0x44000444u, + 0x44004000u, 0x44004004u, 0x44004040u, 0x44004044u, 0x44004400u, 0x44004404u, 0x44004440u, 0x44004444u, + 0x44040000u, 0x44040004u, 0x44040040u, 0x44040044u, 0x44040400u, 0x44040404u, 0x44040440u, 0x44040444u, + 0x44044000u, 0x44044004u, 0x44044040u, 0x44044044u, 0x44044400u, 0x44044404u, 0x44044440u, 0x44044444u, + 0x44400000u, 0x44400004u, 0x44400040u, 0x44400044u, 0x44400400u, 0x44400404u, 0x44400440u, 0x44400444u, + 0x44404000u, 0x44404004u, 0x44404040u, 0x44404044u, 0x44404400u, 0x44404404u, 0x44404440u, 0x44404444u, + 0x44440000u, 0x44440004u, 0x44440040u, 0x44440044u, 0x44440400u, 0x44440404u, 0x44440440u, 0x44440444u, + 0x44444000u, 0x44444004u, 0x44444040u, 0x44444044u, 0x44444400u, 0x44444404u, 0x44444440u, 0x44444444u, +}; + +template +__global__ void __launch_bounds__(NT, 1) gu_kernel(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const block_q8_1_mmq* __restrict__ xq, + const int total_rows, float* __restrict__ dst) { + static_assert(TPP % TG == 0, ""); + __shared__ __align__(16) uint32_t tabA_s[1024]; + __shared__ __align__(16) uint32_t tabB_s[256]; + __shared__ __align__(16) int act_s[TPP][4][16]; + __shared__ float dy_s[TPP][4]; + __shared__ int wsm[NT * CWS]; + for (int i = threadIdx.x; i < 1024; i += NT) tabA_s[i] = kIq2sCode[i]; + for (int i = threadIdx.x; i < 256; i += NT) tabB_s[i] = kIq2sSign[i]; + int r0; + const Chunk c = cta_work(chunks, r0); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int ci = lane & 15, rsub = lane >> 4; // lanes 0-10 stage one row, 16-26 the next; ci = the compact word + const bool cact = ci < CW; + const char* __restrict__ Wr = w + xoff[c.j] + (size_t) (r0 + 2 * warp + rsub) * ROWB; // + k * 2 NW rows + int* __restrict__ wsl = wsm + (2 * warp + rsub) * CWS + ci; // + k * 2 NW * CWS + const int* cw = wsm + threadIdx.x * CWS; + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int hc = 0; hc < NSTEP; ++hc) { + const int kc = hc >> 1, h = hc & 1; + // compact word ci's byte offset inside the block (ci 0-9); d (ci 10) sits at +0 + const int boff = ci < 4 ? 2 + 16 * h + 4 * ci : ci < 8 ? 18 + 16 * h + 4 * ci : ci == 8 ? 66 + 4 * h : 74 + 4 * h; + __syncthreads(); + if (cact) { + const char* blk0 = Wr + kc * BB; + if (kc & 1) { // odd block: words aligned, one load each; d = high half of the word at blk - 2 + const int o = ci < 10 ? boff : -2, sh = ci < 10 ? 0 : 16; + int v[KR]; +#pragma unroll + for (int k = 0; k < KR; ++k) v[k] = __ldg((const int*) (blk0 + (size_t) k * 2 * NW * ROWB + o)); +#pragma unroll + for (int k = 0; k < KR; ++k) wsl[k * 2 * NW * CWS] = __funnelshift_r(v[k], v[k], sh); + } else { // even block: words 2 bytes into a word, two loads + funnel shift; d = low half at blk + const int o0 = ci < 10 ? boff - 2 : 0, o1 = ci < 10 ? boff + 2 : 0, sh = ci < 10 ? 16 : 0; + int v0[KR], v1[KR]; +#pragma unroll + for (int k = 0; k < KR; ++k) { + const char* b = blk0 + (size_t) k * 2 * NW * ROWB; + v0[k] = __ldg((const int*) (b + o0)); + v1[k] = __ldg((const int*) (b + o1)); + } +#pragma unroll + for (int k = 0; k < KR; ++k) wsl[k * 2 * NW * CWS] = __funnelshift_r(v0[k], v1[k], sh); + } + } + stage_act(xq, hc, total_rows, trow0, nt, act_s, dy_s); + __syncthreads(); +#pragma unroll 1 + for (int s = 0; s < 4; ++s) { // cw: [0..3] grid bytes, [4..7] sign bytes, [8] qh, [9] scales, [10] d + int wv[16]; + float dx0, dx1; + const int qs_packed = cw[s], signs_packed_32 = cw[4 + s], qh = (cw[8] >> (8 * s)) & 0xFF; +#pragma unroll + for (int l = 0; l < QR2_S; ++l) { + const int idx = ((qs_packed >> (8 * l)) & 0xFF) | ((qh << (8 - 2 * l)) & 0x300); + const uint32_t sel = tabA_s[idx] | tabB_s[(signs_packed_32 >> (8 * l)) & 0xFF]; + split2((int) __byte_perm(PX, PY, sel), wv[4 * l + 0], wv[4 * l + 1]); + split2((int) __byte_perm(PX, PY, sel >> 16), wv[4 * l + 2], wv[4 * l + 3]); + } + scales16((cw[9] >> (8 * s)) & 0xFF, cw[10], dx0, dx1); + mac16(wv, dx0, dx1, s, act_s, dy_s, nt, acc); + } + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * GU_ROWS + r0 + threadIdx.x] = acc[t]; + } +} +} // namespace s2 + +template +void go(F kernel, const Layer& L, cudaStream_t s) { + kernel<<<(unsigned) (L.nchunks * PIECES), NT, 0, s>>>(L.chunks, (const char*) L.w, L.xoff, (const block_q8_1_mmq*) L.xq, + (int) L.rows, L.gu); +} + +} // namespace + +// tokens per pass by u (the measured crossovers): IQ2_XXS 8 below 5 rows per expert, 16 below 25, then 32 (4-token +// groups); IQ2_XS 8 below 6, 16 below 20, then 16 in 4-token groups; IQ2_S 8 below 5, 16 below 25, then 32 +void gu_iq2_xxs(const Layer& L, double u, cudaStream_t s) { + if (L.nchunks <= 0) return; + if (u < 5.0) go(xxs2::gu_kernel<8, 3, 2>, L, s); + else if (u < 25.0) go(xxs2::gu_kernel<16, 3, 2>, L, s); + else go(xxs2::gu_kernel<32, 2, 4>, L, s); + ck(cudaGetLastError(), "gate/up IQ2_XXS"); +} + +void gu_iq2_xs(const Layer& L, double u, cudaStream_t s) { + if (L.nchunks <= 0) return; + if (u < 6.0) go(xs2::gu_kernel<8, 1, 2>, L, s); + else if (u < 20.0) go(xs2::gu_kernel<16, 0, 2>, L, s); + else go(xs2::gu_kernel<16, 1, 4>, L, s); + ck(cudaGetLastError(), "gate/up IQ2_XS"); +} + +void gu_iq2_s(const Layer& L, double u, cudaStream_t s) { + if (L.nchunks <= 0) return; + if (u < 5.0) go(s2::gu_kernel<8, 0, 2>, L, s); + else if (u < 25.0) go(s2::gu_kernel<16, 1, 2>, L, s); + else go(s2::gu_kernel<32, 1, 4>, L, s); + ck(cudaGetLastError(), "gate/up IQ2_S"); +} + +} // namespace strata::prefill::ws::detail diff --git a/src/prefill/moe_ws_gu_iq3.cu b/src/prefill/moe_ws_gu_iq3.cu new file mode 100644 index 00000000000..efd551dcdc6 --- /dev/null +++ b/src/prefill/moe_ws_gu_iq3.cu @@ -0,0 +1,370 @@ +// src/prefill/moe_ws_gu_iq3.cu - weight-stationary gate/up products (see moe_ws.hpp): IQ3_XXS and IQ3_S, [1280 x 2560] +// read in place. A CTA = 128 threads = 128 weight rows of one expert (a row per thread, all of the chunk's activation +// rows TPP at a time, the sums in registers); K advances 128 values a step (half a 256-value block = 4 sub-blocks): +// the step's bytes of the 128 rows and the step's q8_1 block of each token go to shared memory (the activations split +// to int16 pairs once), then each thread decodes its row's 4 sub-blocks once per pass and runs them against the tokens. +// Bits: per 32-value sub-block MMQ's int8 values (ggml_cuda_mmq_load_tiles_iq3_*: grid entries, its __vcmpne4 sign +// masks applied as (g ^ s) + (s & 0x01010101) == __vsub4(g ^ s, s), every grid byte >= 1) and scale expression, then +// MMQ's epilogue acc = fma(dx * dy, float(sum), acc), sub-blocks in K order. (Measured on a GP100 die: +// ws / per-group MMQ 0.31-0.35 at 2.5 rows per expert, 0.49-0.52 at 10, 0.67-0.76 at 35, +// 0.78-0.79 at 160; bit-identical, compute-sanitizer clean with exact-size matrices.) +// +// Staging without per-lane branches: rows are 980 / 1100 B, so block kc starts 2 * (kc & 1) bytes into a 4-byte word +// for the whole CTA - an even block's words are two aligned loads and a funnel shift (the second load ends at most 2 +// bytes into block kc + 1 of the same row), an odd block's are aligned; d sits the other way round (an odd block's d is +// the high half of the word 2 bytes before it, inside block kc - 1). Every load of a lane is issued before its stores, +// and nothing outside a row is read (the last row's last block is odd). +#include "common.cuh" +#include "mmq.cuh" + +#include "moe_ws.cuh" + +#include +#include + +namespace strata::prefill::ws::detail { +namespace { + +constexpr int GU_ROWS = 1280, GU_K = 2560, NB = GU_K / QK_K, NSTEP = 2 * NB, NT = 128, PIECES = GU_ROWS / NT; + +void ck(cudaError_t e, const char* what) { + if (e != cudaSuccess) { + std::fprintf(stderr, "prefill ws: %s: %s\n", what, cudaGetErrorString(e)); + std::exit(1); + } +} + +// ---- IQ3_XXS: 98-B blocks; per row and step 8 grid-index words, 4 aux words (signs + scale) and d (stride 13). +// Every lane slot is one word type: grid words rows tid/8 + 16 i (word tid%8), aux words rows tid/4 + 32 i (word +// tid%4), d the thread's own row. Token groups of TG: one branch per group on a partial pass (a group's tokens past +// the count run on copies of row nt - 1 and are not written), none on a full pass. +namespace xxs3 { +constexpr int BB = (int) sizeof(block_iq3_xxs), ROWB = NB * BB, CWS = 13; +static_assert(BB == 98 && ROWB % 4 == 0, "block_iq3_xxs"); + +template +__device__ __forceinline__ void sb_compute(const int* __restrict__ cw, const int s, const uint32_t* __restrict__ grid, + const int2* __restrict__ lut, const int (*act_s)[4][16], + const float (*dy_s)[4], const int nt, float (&acc)[TPP]) { + int wv[16]; + float dx; + { + const uint32_t qa = (uint32_t) cw[2 * s + 0], qb = (uint32_t) cw[2 * s + 1]; // 8 grid indices + const uint32_t aux32 = (uint32_t) cw[8 + s]; +#pragma unroll + for (int l = 0; l < QR3_XXS; ++l) { + const uint32_t qw = l < 2 ? qa : qb; + const uint32_t g0 = grid[(qw >> (16 * (l & 1))) & 0xFF], g1 = grid[(qw >> (16 * (l & 1) + 8)) & 0xFF]; + const int2 sg = lut[(aux32 >> (7 * l)) & 0x7F]; + split2((int) ((g0 ^ (uint32_t) sg.x) + ((uint32_t) sg.x & 0x01010101u)), wv[4 * l + 0], wv[4 * l + 1]); + split2((int) ((g1 ^ (uint32_t) sg.y) + ((uint32_t) sg.y & 0x01010101u)), wv[4 * l + 2], wv[4 * l + 3]); + } + const int ls = aux32 >> 28; + const float d = __ushort_as_half((unsigned short) (cw[12] & 0xFFFF)); + dx = (ls*d + d/2)/2; // MMQ's expression, verbatim + } +#pragma unroll + for (int t0 = 0; t0 < TPP; t0 += TG) { + if (FULL || t0 < nt) { +#pragma unroll + for (int t = t0; t < t0 + TG; ++t) { + const int4* a4 = (const int4*) act_s[t][s]; + const int4 a0 = a4[0], a1 = a4[1], a2 = a4[2], a3 = a4[3]; + const int a[16] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w, a2.x, a2.y, a2.z, a2.w, a3.x, a3.y, a3.z, a3.w}; + const float dy = dy_s[t][s]; + int sc[CH]; + sc[0] = kMagicI; +#pragma unroll + for (int c = 1; c < CH; ++c) sc[c] = 0; +#pragma unroll + for (int m = 0; m < 16; ++m) sc[m % CH] = mad2(wv[m], a[m], sc[m % CH]); + int sum = sc[0]; +#pragma unroll + for (int c = 1; c < CH; ++c) sum += sc[c]; + acc[t] = __fmaf_rn(__fmul_rn(dx, dy), __fsub_rn(__int_as_float(sum), kMagicF), acc[t]); + } + } + } +} + +template +__global__ void __launch_bounds__(NT, 1) gu_kernel(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const block_q8_1_mmq* __restrict__ xq, + const int total_rows, float* __restrict__ dst) { + static_assert(TPP % TG == 0 && (TPP * 32) % NT == 0 && TPP * 4 <= NT, ""); + constexpr int NG = 8, RG = NT / 8; // grid-index-word slots: rows tid/8 + i*RG, word tid%8 + constexpr int NX = 4, RX = NT / 4; // aux-word slots: rows tid/4 + i*RX, word tid%4 + constexpr int NA = TPP * 32 / NT; + __shared__ uint32_t grid_s[256]; + __shared__ int2 lut_s[128]; + __shared__ __align__(16) int act_s[TPP][4][16]; + __shared__ float dy_s[TPP][4]; + __shared__ int wsm[NT * CWS]; + const int tid = threadIdx.x; + for (int i = tid; i < 256; i += NT) grid_s[i] = iq3xxs_grid[i]; + for (int i = tid; i < 128; i += NT) { + const uint32_t signs = unpack_ksigns((uint8_t) i); + lut_s[i] = make_int2(__vcmpne4(signs & 0x08040201, 0), __vcmpne4(signs & 0x80402010, 0)); + } + int r0; + const Chunk c = cta_work(chunks, r0); + const char* __restrict__ W = w + xoff[c.j] + (size_t) r0 * ROWB; + const char* __restrict__ Wg = W + (size_t) (tid >> 3) * ROWB + 4 * (tid & 7); // + i*RG*ROWB + kc*BB + 32h + const char* __restrict__ Wx = W + (size_t) (tid >> 2) * ROWB + 64 + 4 * (tid & 3); // + i*RX*ROWB + kc*BB + 16h + const char* __restrict__ Wd = W + (size_t) tid * ROWB; + int* __restrict__ sg_ = wsm + (tid >> 3) * CWS + (tid & 7); + int* __restrict__ sx_ = wsm + (tid >> 2) * CWS + 8 + (tid & 3); + int* __restrict__ sd_ = wsm + tid * CWS + 12; + const int* cw = wsm + tid * CWS; + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int hc = 0; hc < NSTEP; ++hc) { + const int kc = hc >> 1, h = hc & 1; + const int* xr = (const int*) (xq + (size_t) hc * total_rows + trow0); + __syncthreads(); + int av[NA], dv = 0; +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = tid + i * NT, t = idx >> 5; + av[i] = __ldg(xr + min(t, nt - 1) * 36 + 4 + (idx & 31)); // clamped: no branch, never past row nt-1 + } + if (TPP * 4 == NT || tid < TPP * 4) dv = __ldg(xr + min(tid >> 2, nt - 1) * 36 + (tid & 3)); + const char* bg = Wg + kc * BB + 32 * h; // even kc: the aligned word 2 bytes before the data word + const char* bx = Wx + kc * BB + 16 * h; + if (kc & 1) { // odd block: data words aligned (at +2); d = high half of the word 2 bytes before the block + int vg[NG], vx[NX]; +#pragma unroll + for (int i = 0; i < NG; ++i) vg[i] = __ldg((const int*) (bg + 2 + (size_t) i * RG * ROWB)); +#pragma unroll + for (int i = 0; i < NX; ++i) vx[i] = __ldg((const int*) (bx + 2 + (size_t) i * RX * ROWB)); + const int vd = __ldg((const int*) (Wd + kc * BB - 2)); +#pragma unroll + for (int i = 0; i < NG; ++i) sg_[i * RG * CWS] = vg[i]; +#pragma unroll + for (int i = 0; i < NX; ++i) sx_[i * RX * CWS] = vx[i]; + *sd_ = (int) ((uint32_t) vd >> 16); + } else { // even block: data words 2 bytes into a word: two aligned words + funnel shift; d = low half + int lg[NG], hg[NG], lx[NX], hx[NX]; +#pragma unroll + for (int i = 0; i < NG; ++i) { + lg[i] = __ldg((const int*) (bg + (size_t) i * RG * ROWB)); + hg[i] = __ldg((const int*) (bg + 4 + (size_t) i * RG * ROWB)); + } +#pragma unroll + for (int i = 0; i < NX; ++i) { + lx[i] = __ldg((const int*) (bx + (size_t) i * RX * ROWB)); + hx[i] = __ldg((const int*) (bx + 4 + (size_t) i * RX * ROWB)); + } + const int vd = __ldg((const int*) (Wd + kc * BB)); +#pragma unroll + for (int i = 0; i < NG; ++i) sg_[i * RG * CWS] = (int) __funnelshift_r((uint32_t) lg[i], (uint32_t) hg[i], 16); +#pragma unroll + for (int i = 0; i < NX; ++i) sx_[i * RX * CWS] = (int) __funnelshift_r((uint32_t) lx[i], (uint32_t) hx[i], 16); + *sd_ = vd; + } +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = tid + i * NT, t = idx >> 5, wd = idx & 31; + int lo, hi; + split2(av[i], lo, hi); + *(int2*) &act_s[t][wd >> 3][2 * (wd & 7)] = make_int2(lo, hi); + } + if (TPP * 4 == NT || tid < TPP * 4) dy_s[tid >> 2][tid & 3] = __int_as_float(dv); + __syncthreads(); + if (nt == TPP) { +#pragma unroll 1 + for (int s = 0; s < 4; ++s) sb_compute(cw, s, grid_s, lut_s, act_s, dy_s, nt, acc); + } else { +#pragma unroll 1 + for (int s = 0; s < 4; ++s) sb_compute(cw, s, grid_s, lut_s, act_s, dy_s, nt, acc); + } + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * GU_ROWS + r0 + tid] = acc[t]; + } +} +} // namespace xxs3 + +// ---- IQ3_S: 110-B blocks; per row and step a compact copy of the bytes the 4 sub-blocks read - qs 32 B, signs 16 B, +// qh 4 B, d and the 2 scale bytes: 14 words, row stride 15 - staged by lanes 0-13 / 16-29 of each warp (two rows per +// warp and k), each word built from two aligned loads by one PRMT. The sign masks come from a 256-entry table of MMQ's +// own __vcmpne4 expressions. Token groups of TG behind one `t < nt` branch. +namespace s3 { +constexpr int BB = 110, ROWB = NB * BB, CW = 14, CWS = 15, NW = NT / 32, KN = NT / (2 * NW); + +__device__ __forceinline__ int2 mmq_signs(const int sb) { // ggml_cuda_mmq_load_tiles_iq3_s's, verbatim + const int s0 = __vcmpne4(((sb & 0x03) << 7) | ((sb & 0x0C) << 21), 0x00000000); + const int s1 = __vcmpne4(((sb & 0x30) << 3) | ((sb & 0xC0) << 17), 0x00000000); + return make_int2(s0, s1); +} + +// sub-block s of the staged half: kqsx = 4 h + s's 8 int8x4 words and scale. cw: [0..7] qs, [8..11] sign bytes, +// [12] qh, [13] d | the half's 2 scale bytes << 16 +__device__ __forceinline__ float dec(const int* __restrict__ cw, const int s, const uint32_t* __restrict__ grid, + const int2* __restrict__ lut, int (&q)[8]) { + const int2 qs_packed = make_int2(cw[2 * s + 0], cw[2 * s + 1]); + const uint8_t* qs = (const uint8_t*) &qs_packed; + const int qh = (cw[12] >> (8 * s)) & 0xFF; + const int signs_packed_32 = cw[8 + s]; + const int w13 = cw[13]; +#pragma unroll + for (int l = 0; l < QR3_S; ++l) { + const int2 grid_pos = make_int2(grid[qs[2 * l + 0] | ((qh << (8 - 2 * l)) & 0x100)], + grid[qs[2 * l + 1] | ((qh << (7 - 2 * l)) & 0x100)]); + const int2 sg = lut[(signs_packed_32 >> (8 * l)) & 0xFF]; + q[2 * l + 0] = (grid_pos.x ^ sg.x) + (sg.x & 0x01010101); + q[2 * l + 1] = (grid_pos.y ^ sg.y) + (sg.y & 0x01010101); + } + const int ls = 1 + 2 * (((w13 >> (16 + 8 * (s >> 1))) >> (4 * (s & 1))) & 0x0F); + const float d = __ushort_as_half((unsigned short) (w13 & 0xFFFF)); + return ls * d; +} + +// compact word ci of step half h in a block of parity odd: its two aligned words (offsets from the block) and selector +__device__ __forceinline__ void stage_map(const int ci, const int h, const int odd, int& o0, int& o1, int& sel) { + const int off = ci < 8 ? 2 + 32 * h + 4 * ci : ci < 12 ? 74 + 16 * h + 4 * (ci - 8) : 66 + 4 * h; + const bool dw = ci == 13; + o0 = dw ? (odd ? -2 : 0) : (odd ? off : off - 2); + o1 = dw ? (odd ? 106 : 104 + 4 * h) : (odd ? off : off + 2); + sel = dw ? (odd ? (h ? 0x7632 : 0x5432) : (h ? 0x5410 : 0x7610)) : (odd ? 0x3210 : 0x5432); +} + +template +__device__ __forceinline__ void sb_compute(const int* __restrict__ cw, const int s, const uint32_t* __restrict__ grid, + const int2* __restrict__ lut, const int (*act_s)[4][16], const float (*dy_s)[4], + const int nt, float (&acc)[TPP]) { + int q[8]; + const float dx = dec(cw, s, grid, lut, q); + int wv[16]; +#pragma unroll + for (int m = 0; m < 8; ++m) split2(q[m], wv[2 * m], wv[2 * m + 1]); +#pragma unroll + for (int t0 = 0; t0 < TPP; t0 += TG) { + if (t0 < nt) { +#pragma unroll + for (int t = t0; t < t0 + TG; ++t) { + const int4* a4 = (const int4*) act_s[t][s]; + const int4 a0 = a4[0], a1 = a4[1], a2 = a4[2], a3 = a4[3]; + const float dy = dy_s[t][s]; + int s0 = 0, s1 = 0, s2 = 0, s3 = 0; + s0 = mad2(wv[0], a0.x, s0); s1 = mad2(wv[1], a0.y, s1); s2 = mad2(wv[2], a0.z, s2); s3 = mad2(wv[3], a0.w, s3); + s0 = mad2(wv[4], a1.x, s0); s1 = mad2(wv[5], a1.y, s1); s2 = mad2(wv[6], a1.z, s2); s3 = mad2(wv[7], a1.w, s3); + s0 = mad2(wv[8], a2.x, s0); s1 = mad2(wv[9], a2.y, s1); s2 = mad2(wv[10], a2.z, s2); s3 = mad2(wv[11], a2.w, s3); + s0 = mad2(wv[12], a3.x, s0); s1 = mad2(wv[13], a3.y, s1); s2 = mad2(wv[14], a3.z, s2); s3 = mad2(wv[15], a3.w, s3); + acc[t] = __fmaf_rn(__fmul_rn(dx, dy), __int2float_rn((s0 + s1) + (s2 + s3)), acc[t]); + } + } + } +} + +template +__global__ void __launch_bounds__(NT) gu_kernel(const Chunk* __restrict__ chunks, const char* __restrict__ w, + const int64_t* __restrict__ xoff, const block_q8_1_mmq* __restrict__ xq, + const int total_rows, float* __restrict__ dst) { + static_assert(TPP % TG == 0, ""); + constexpr size_t KSTR = (size_t) 2 * NW * ROWB; // a staging lane's rows are 2 * NW apart + __shared__ uint32_t grid_s[512]; + __shared__ int2 lut_s[256]; + __shared__ __align__(16) int act_s[TPP][4][16]; + __shared__ float dy_s[TPP][4]; + __shared__ int wsm[NT * CWS]; + for (int i = threadIdx.x; i < 512; i += NT) grid_s[i] = iq3s_grid[i]; + for (int i = threadIdx.x; i < 256; i += NT) lut_s[i] = mmq_signs(i); + int r0; + const Chunk c = cta_work(chunks, r0); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int ci = lane & 15, rsub = lane >> 4; // lanes 0-13 stage one row, 16-29 the next; ci = the compact word + const bool cact = ci < CW; + const char* __restrict__ Wr = w + xoff[c.j] + (size_t) (r0 + 2 * warp + rsub) * ROWB; // + k * KSTR + int* __restrict__ wsl = wsm + (2 * warp + rsub) * CWS + ci; // + k * 2 NW * CWS + for (int p0 = 0; p0 < c.rows; p0 += TPP) { + const int nt = min(TPP, c.rows - p0); + const int trow0 = c.row0 + p0; + float acc[TPP]; +#pragma unroll + for (int t = 0; t < TPP; ++t) acc[t] = 0.0f; +#pragma unroll 1 + for (int hc = 0; hc < NSTEP; ++hc) { + __syncthreads(); + if (cact) { // branch-free inside: every lane issues its 2 * KN loads, then its KN stores + int o0, o1, sel; + stage_map(ci, hc & 1, (hc >> 1) & 1, o0, o1, sel); + const char* b = Wr + (hc >> 1) * BB; + int v0[KN], v1[KN]; +#pragma unroll + for (int k = 0; k < KN; ++k) v0[k] = __ldg((const int*) (b + k * KSTR + o0)); +#pragma unroll + for (int k = 0; k < KN; ++k) v1[k] = __ldg((const int*) (b + k * KSTR + o1)); +#pragma unroll + for (int k = 0; k < KN; ++k) wsl[k * 2 * NW * CWS] = (int) __byte_perm(v0[k], v1[k], sel); + } + { // the step's activations: nt consecutive 144-B blocks, loads first, then split + store + constexpr int NA = (TPP * 36 + NT - 1) / NT; + int av[NA]; + const int* xr = (const int*) (xq + (size_t) hc * total_rows + trow0); +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * 36) av[i] = __ldg(xr + idx); + } +#pragma unroll + for (int i = 0; i < NA; ++i) { + const int idx = threadIdx.x + i * NT; + if (idx < nt * 36) { + const int t = idx / 36, wd = idx - t * 36; + if (wd < 4) { + dy_s[t][wd] = __int_as_float(av[i]); + } else { + const int m = wd - 4; + int lo, hi; + split2(av[i], lo, hi); + *(int2*) &act_s[t][m / 8][2 * (m % 8)] = make_int2(lo, hi); + } + } + } + } + __syncthreads(); + const int* cw = wsm + threadIdx.x * CWS; +#pragma unroll 1 + for (int s = 0; s < 4; ++s) sb_compute(cw, s, grid_s, lut_s, act_s, dy_s, nt, acc); + } +#pragma unroll + for (int t = 0; t < TPP; ++t) + if (t < nt) dst[(size_t) (trow0 + t) * GU_ROWS + r0 + threadIdx.x] = acc[t]; + } +} +} // namespace s3 + +template +void go(F kernel, const Layer& L, cudaStream_t s) { + kernel<<<(unsigned) (L.nchunks * PIECES), NT, 0, s>>>(L.chunks, (const char*) L.w, L.xoff, (const block_q8_1_mmq*) L.xq, + (int) L.rows, L.gu); +} + +} // namespace + +// tokens per pass by u (the measured crossovers): IQ3_XXS 8 below 5 rows per expert, 16 below 25, then 32 (4-token +// groups); IQ3_S 8 below 4, 16 below 15, then 32 +void gu_iq3_xxs(const Layer& L, double u, cudaStream_t s) { + if (L.nchunks <= 0) return; + if (u < 5.0) go(xxs3::gu_kernel<8, 4, 2>, L, s); + else if (u < 25.0) go(xxs3::gu_kernel<16, 4, 2>, L, s); + else go(xxs3::gu_kernel<32, 2, 4>, L, s); + ck(cudaGetLastError(), "gate/up IQ3_XXS"); +} + +void gu_iq3_s(const Layer& L, double u, cudaStream_t s) { + if (L.nchunks <= 0) return; + if (u < 4.0) go(s3::gu_kernel<8, 2>, L, s); + else if (u < 15.0) go(s3::gu_kernel<16, 2>, L, s); + else go(s3::gu_kernel<32, 4>, L, s); + ck(cudaGetLastError(), "gate/up IQ3_S"); +} + +} // namespace strata::prefill::ws::detail diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index c9b4535407a..4bb08898672 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -39,6 +39,7 @@ #include "strata/kernels/qsa_select.hpp" #include "strata/prefill/gemm.hpp" #include "strata/prefill/moe_fused.hpp" +#include "strata/prefill/moe_ws.hpp" #include "strata/prefill/moe_fused_iq.hpp" #include "strata/prefill/moe_mmq.hpp" #include "strata/core/peer_experts.hpp" @@ -82,6 +83,8 @@ bool gather_native_group(const GatherGroup&, size_t, size_t, size_t, size_t, voi void*, size_t) { return false; } void gather_strata_q2(const uint8_t*, void*, void*, void*) {} void swiglu(const float*, float*, int64_t, int64_t, bool, void*) {} +bool swiglu_quant_ok(int) { return false; } +void swiglu_quant(const float*, void*, int64_t, bool, void*) {} void swiglu_scaled(const float*, float*, int64_t, int64_t, const int32_t*, int, const float*, int64_t, void*) {} void down_row_scales(float*, const int32_t*, int, const float*, int64_t, void*) {} void iota(int32_t*, int64_t, void*) {} @@ -112,6 +115,17 @@ void experts_native(const Batch&, const NativeGeom&, int, int64_t, const void*, float*, void*) {} } // namespace strata::prefill::fused #endif +#ifndef STRATA_PREFILL_WS +// the weight-stationary experts (moe_ws.hpp) are in the CUDA STRATA_EXPERIMENTAL_SM60 build only +namespace strata::prefill::ws { +int mode() { return 0; } +bool supported(int, int) { return false; } +bool faster(int, double) { return false; } +size_t chunks(const int32_t*, size_t, const int32_t*, const int32_t*, Chunk*) { return 0; } +bool gate_up(const Layer&, void*) { return false; } +void down(const Layer&, void*) {} +} // namespace strata::prefill::ws +#endif #if defined(STRATA_GPU_METAL) // The Metal backend reads a prompt through its decode window (metal/src/engine/window.cpp): init keeps the chunk, run @@ -918,6 +932,13 @@ struct Prefill::Impl { float *Xscale = nullptr, *Hscale = nullptr, *grp_tail = nullptr; // NVFP4: activation scales, expert tails float* row_sd = nullptr; // NVFP4 in MMQ: each Dm row's s_down, applied by the combine as it reads the row std::vector bounds_host; + // STRATA_MMQ_WS=1 (moe_ws.hpp): a layer whose experts are all resident - their byte offsets from the expert cache's + // base in MMQ order and the CTAs' chunks, built on the host per layer and uploaded with the bounds (the device + // buffers exist only where ws::mode() != 0; the host copies are the unmapped path's) + int64_t* xoff_dev = nullptr; + std::vector xoff_host; + ws::Chunk* ws_chunks = nullptr; + std::vector ws_chunks_host; std::unique_ptr mmq_ctx; std::vector ids_host, slot_host, src_host, cnt, off; // set_cpu_pool: a small chunk's CPU experts. The layer's MoE input (T x N, from `mixed`), the rows the pool writes @@ -1474,7 +1495,10 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co m.steps_up.init(T * strata::kernels::kStepCount); m.ids_host.resize(T * K); m.slot_host.resize(T * K); m.src_host.resize(T * K); m.cnt.resize(m.g->n_expert); m.off.resize(m.g->n_expert + 1); { - const size_t need = 3 * T * K + (size_t) (2 * (m.g->n_expert + m.g->n_expert / MMQ_GROUP + 2)); + // the routing tables, the MMQ group bounds; STRATA_MMQ_WS: the in-place offsets (n_expert int64) and the + // weight-stationary layers' chunks (4 int32 each, 16-byte aligned) + const size_t need = 3 * T * K + (size_t) (2 * (m.g->n_expert + m.g->n_expert / MMQ_GROUP + 2)) + + (ws::mode() != 0 ? 2 * (size_t) m.g->n_expert + 4 * ws::max_chunks((int64_t) T * K, m.g->n_expert) + 3 : 0); const char* gc = std::getenv("STRATA_GROUP_COPY"); if (m.grp_n < need && !(gc && gc[0] == '1')) { if (m.grp_host) cudaFreeHost(m.grp_host); @@ -1629,6 +1653,10 @@ bool Prefill::carve(size_t T, void* alloc) { m.bounds_dev = o.take((size_t) (2 * (m.g->n_expert + m.g->n_expert / MMQ_GROUP + 2)), ok); m.grp_gu = o.take(MMQ_GROUP * mp.gu_max + MMQ_TAIL, ok); m.grp_d = o.take(MMQ_GROUP * mp.d_max + MMQ_TAIL, ok); + if (ws::mode() != 0) { // STRATA_MMQ_WS: the offsets and chunks (bytes_needed_impl takes the same) + m.xoff_dev = o.take((size_t) m.g->n_expert, ok); + m.ws_chunks = (ws::Chunk*) o.take(4 * ws::max_chunks((int64_t) T * K, m.g->n_expert), ok); + } // (written at every run's start, not here: when serving, these are live expert-cache slots until a request // lends them - a write now would corrupt a resident expert) if (!m.mmq_ctx) m.mmq_ctx = std::make_unique(); @@ -2362,6 +2390,10 @@ uint64_t Prefill::bytes_needed_impl(const core::ModelGeometry& g, const core::Se o.take((size_t) (2 * (g.n_expert + g.n_expert / MMQ_GROUP + 2)), ok); o.take(MMQ_GROUP * mp.gu_max + MMQ_TAIL, ok); o.take(MMQ_GROUP * mp.d_max + MMQ_TAIL, ok); + if (ws::mode() != 0) { // carve's STRATA_MMQ_WS offsets and chunks + o.take((size_t) g.n_expert, ok); + o.take(4 * ws::max_chunks((int64_t) T * K, g.n_expert), ok); + } } if (owned_pages) { if (stage_slots(T) > 0) o.take((size_t) stage_slots(T) * (size_t) MAXBLOB(), ok); // one allocation @@ -4023,6 +4055,8 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri } const size_t mmq_gub = use_mmq ? mmq::matrix_bytes(mmq_gt, 1280, N) : 0; const size_t mmq_db = use_mmq ? mmq::matrix_bytes(mmq_dt, N, 640) : 0; + bool ws_layer = false; // STRATA_MMQ_WS: this layer's products weight-stationary (below) + ws::Layer wl; pt.mark(kPfGather, cs); // the routed experts' input: `mixed`, or (APR) its rotated copy, which the peer / the layer // split's helper also take (one rotation for both cards; `mixed` itself feeds the shared expert) @@ -4065,6 +4099,67 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri return false; } } + // STRATA_MMQ_WS=1 (moe_ws.hpp, GP100): a layer whose routed experts are ALL resident, on this + // GPU alone (no peer, CPU or streamed share), runs gate/up and down as one weight-stationary + // launch each, with ONE swiglu_quant over all its rows in between, instead of a gate/up, + // swiglu and down MMQ launch per 16-expert group - the same GU, H and Dm bits. The experts + // are read where they lie in the expert cache (offsets from its base). Only where both + // products beat the per-group MMQ at this layer's rows per expert, H can be q8_1 in one pass + // (no Hadamard, NVFP4 or FP4 rows; not STRATA_DBG_NAN, which reads H in floats), and not the + // fused layout (its Hq holds stream_all_min() - 1 tokens' rows). + static const bool ws_dbg = std::getenv("STRATA_DBG_NAN") != nullptr; + const double ws_u = n > 0 ? (double) rows_local / (double) n : 0.0; + if (ws::mode() != 0 && m.xoff_dev != nullptr && m.ws_chunks != nullptr && n > 0 && !ws_dbg && + lay.native && m.cache != nullptr && m.host_res != nullptr && !wmma_l && + !mmq_plan().direct[(size_t) l] && lay.fmt[(size_t) l].up_off == mmq_gub / 2 && + order_peer.empty() && rows_cpu == 0 && rows_local == T * K && !m.fused_bufs && + (!stream_all || seq_start[(size_t) l] == seq_start[(size_t) l + 1]) && !had_gu && !had_d && + !lay.fmt[(size_t) l].tail_off && !mmq::fp4_activations(mmq_gt) && + !mmq::fp4_activations(mmq_dt) && mmq::swiglu_quant_ok(mmq_dt) && + ws::supported(mmq_gt, mmq_dt) && ws::faster(mmq_gt, ws_u) && ws::faster(mmq_dt, ws_u)) { + const uint8_t* ip_base = m.cache->device_slot(0); + m.xoff_host.assign(n, -1); + bool all = true; // every expert resident + uint64_t at = 0; // their offsets' low bits + for (size_t j = 0; j < n && all; ++j) { + const int32_t s = m.host_res[(size_t) l * m.g->n_expert + order[j]]; + if (s < 0) { all = false; break; } + m.xoff_host[j] = (int64_t) (m.cache->device_slot(s) - ip_base); + at |= (uint64_t) m.xoff_host[j]; + } + ws_layer = all && ((uint64_t) (uintptr_t) ip_base | at | lay.fmt[(size_t) l].down_off) % ws::kAlign == 0; + if (ws_layer) { + if (static bool said_ws = false; !said_ws) { + said_ws = true; + std::fprintf(stderr, "strata prefill: a resident layer's experts weight-stationary, one " + "launch per product (STRATA_MMQ_WS=1)\n"); + } + // the offsets after the bounds in the mapped table, the chunks after them (16-byte + // aligned); or the host copies through pinned slots + const size_t xo = 3 * m.grp_tk + 2 * (m.g->n_expert + m.g->n_expert / MMQ_GROUP + 2); + ws::Chunk* ch = nullptr; + if (grp_mapped) { + int32_t* xh = m.grp_host + xo; + std::memcpy(xh, m.xoff_host.data(), n * sizeof(int64_t)); + copy_i32((int32_t*) m.xoff_dev, m.grp_dev + xo, (int64_t) (2 * n), m.cs); + ch = (ws::Chunk*) (m.grp_host + (xo + 2 * m.g->n_expert + 3) / 4 * 4); + } else { + m.ws_chunks_host.resize(ws::max_chunks(T * K, m.g->n_expert)); + ch = m.ws_chunks_host.data(); + } + wl.nchunks = (int64_t) ws::chunks(order.data(), n, m.cnt.data(), m.bounds_host.data(), ch); + if (grp_mapped) { + copy_i32((int32_t*) m.ws_chunks, m.grp_dev + ((int32_t*) ch - m.grp_host), 4 * wl.nchunks, m.cs); + } else if (m.up_ring.h2d(m.xoff_dev, m.xoff_host.data(), n * sizeof(int64_t), m.cs) != cudaSuccess || + m.up_ring.h2d(m.ws_chunks, ch, (size_t) wl.nchunks * sizeof(ws::Chunk), m.cs) != cudaSuccess) { + err = "prefill: the weight-stationary tables' upload failed"; + return false; + } + wl.gu_type = mmq_gt; wl.d_type = mmq_dt; wl.w = ip_base; wl.xoff = m.xoff_dev; + wl.down_off = (int64_t) lay.fmt[(size_t) l].down_off; wl.chunks = m.ws_chunks; + wl.rows = T * K; wl.n = (int) n; wl.xq = m.Xq; wl.gu = m.GU; wl.hq = m.Hq; wl.dm = m.Dm; + } + } } else { gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * K, N, m.cs); // APR: Hadamard-folded gate/up take H (s * x); Xs is this layer's private gathered copy @@ -4639,7 +4734,21 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri m.gemm.f16(m.Hh + o0 * 640, m.dq_d[q], m.Dm + o0 * N, ne, N, 640); return true; }; - if (!stream_all) { + if (ws_layer) { + // STRATA_MMQ_WS: the layer's gate/up in one launch, H of all its rows in one pass, its down + // in one launch (no expert of the layer is staged or streamed: nothing to wait for or give + // back beyond what the walk would) + stats_.experts_resident += (int64_t) order.size(); + pt.mark(kPfGemmGU, cs); + if (!ws::gate_up(wl, m.cs)) { // the layer was taken because ws::faster said so + err = "prefill: the weight-stationary gate/up declined a layer it was picked for"; + return false; + } + mmq::swiglu_quant(m.GU, m.Hq, T * K, !lay.native, m.cs); + pt.mark(kPfGemmD, cs); + ws::down(wl, m.cs); + if (stream_all && group_gather) give_back(consumed); + } else if (!stream_all) { size_t staged = 0; size_t pending = 0; const bool stream_ahead = stream_ahead_enabled(); diff --git a/tests/cuda/prefill_swiglu_quant_test.cpp b/tests/cuda/prefill_swiglu_quant_test.cpp new file mode 100644 index 00000000000..0d7b9d686f6 --- /dev/null +++ b/tests/cuda/prefill_swiglu_quant_test.cpp @@ -0,0 +1,94 @@ +// prefill_swiglu_quant_test - the prompt path's fused MMQ expert step against the kernels it replaces, byte for byte: +// mmq::swiglu_quant against swiglu + quantize (the down product's q8_1 rows of H), for down types the packs have +// (Q2_0, IQ4_NL, Q8_0), with the gate/up rows split (GGUF) and interleaved (the Strata pack), a zero row and a row +// with one huge value. Exit 77 without a CUDA device. +#include "strata/prefill/moe_mmq.hpp" + +#include "ggml.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { +namespace mmq = strata::prefill::mmq; + +void ck(cudaError_t e, const char* what) { + if (e != cudaSuccess) throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(e)); +} + +struct Dev { + void* p = nullptr; + explicit Dev(size_t n) { ck(cudaMalloc(&p, n), "cudaMalloc"); } + ~Dev() { cudaFree(p); } + Dev(const Dev&) = delete; + Dev& operator=(const Dev&) = delete; + template T* as() const { return (T*) p; } +}; + +int swiglu_quant_case(ggml_type t, int64_t rows, bool interleaved, uint32_t seed) { + constexpr int64_t NFF = 640; + std::mt19937 rng(seed); + std::normal_distribution nd(0.f, 1.f); + std::uniform_real_distribution mag(0.05f, 40.f); + std::vector gu((size_t) rows * 2 * NFF); + for (int64_t r = 0; r < rows; ++r) { + const float m = mag(rng); + for (int64_t k = 0; k < 2 * NFF; ++k) gu[(size_t) (r * 2 * NFF + k)] = nd(rng) * m; + } + if (rows > 3) { // a zero row, a row with one huge value + std::memset(gu.data() + 2 * NFF, 0, 2 * NFF * sizeof(float)); + gu[(size_t) (3 * 2 * NFF + 5)] = 3e4f; + } + const size_t qb = mmq::q8_bytes(rows, NFF); + Dev dgu(gu.size() * 4), dh((size_t) rows * NFF * 4 + 16), dq0(qb), dq1(qb); + ck(cudaMemcpy(dgu.p, gu.data(), gu.size() * 4, cudaMemcpyHostToDevice), "upload"); + ck(cudaMemset(dq0.p, 0x5a, qb), "memset"); + ck(cudaMemset(dq1.p, 0x5a, qb), "memset"); + // the two kernels it replaces, as prefill.cpp launches them + mmq::swiglu(dgu.as(), dh.as(), rows, NFF, interleaved, nullptr); + mmq::quantize(dh.as(), nullptr, dq0.p, (int) t, NFF, NFF, rows, nullptr); + mmq::swiglu_quant(dgu.as(), dq1.p, rows, interleaved, nullptr); + ck(cudaDeviceSynchronize(), "run"); + const size_t used = (size_t) rows * (1024 / 128) * 144; // the rows' 8 blocks of 128 values (144 bytes each) + std::vector a(used), b(used); + ck(cudaMemcpy(a.data(), dq0.p, used, cudaMemcpyDeviceToHost), "download"); + ck(cudaMemcpy(b.data(), dq1.p, used, cudaMemcpyDeviceToHost), "download"); + size_t bad = 0; + for (size_t i = 0; i < used; ++i) bad += a[i] != b[i]; + std::printf(" swiglu_quant %-6s %5lld rows%s: %zu of %zu bytes differ\n", ggml_type_name(t), (long long) rows, + interleaved ? " interleaved" : "", bad, used); + return bad == 0 ? 0 : 1; +} + +} // namespace + +int main() { + int ndev = 0; + if (cudaGetDeviceCount(&ndev) != cudaSuccess || ndev == 0) { + std::printf("prefill_swiglu_quant_test: no CUDA device (skipped)\n"); + return 77; + } + int fails = 0; + try { + for (ggml_type t : {GGML_TYPE_Q2_0, GGML_TYPE_IQ4_NL, GGML_TYPE_Q8_0}) { + if (!mmq::swiglu_quant_ok((int) t)) { std::printf(" %s: not on the fused path here\n", ggml_type_name(t)); continue; } + for (bool il : {false, true}) { + fails += swiglu_quant_case(t, 1, il, 1); + fails += swiglu_quant_case(t, 37, il, 2); + fails += swiglu_quant_case(t, 4113, il, 3); + } + } + } catch (const std::exception& e) { + std::fprintf(stderr, "prefill_swiglu_quant_test: %s\n", e.what()); + return 1; + } + std::printf("prefill_swiglu_quant_test: %s\n", fails ? "FAILED" : "ok"); + return fails ? 1 : 0; +}