diff --git a/include/strata/prefill/moe_mmq.hpp b/include/strata/prefill/moe_mmq.hpp index d9f0eba6ff1..0859dca35a8 100644 --- a/include/strata/prefill/moe_mmq.hpp +++ b/include/strata/prefill/moe_mmq.hpp @@ -64,6 +64,16 @@ class Context { /// slot: gate rows then up rows at `gu_dst`, down at `d_dst`. void gather_native(const void* gate, const void* up, size_t gu_half_bytes, const void* down, size_t d_bytes, void* gu_dst, void* d_dst, void* stream); +/// gather_native for an MMQ group's experts [first, n) in ONE launch: expert q's blob (`blob[q]`; gate at +0, up at +/// +up_off, down at +down_off) to gu_dst + q * gu_stride and d_dst + q * d_stride - the same bytes as one gather_native +/// each. Every pointer, offset and size 16-byte aligned (false otherwise: nothing launched, gather one at a time). +constexpr int kGatherGroupMax = 16; +struct GatherGroup { + const uint8_t* blob[kGatherGroupMax] = {}; + int first = 0, n = 0; +}; +bool gather_native_group(const GatherGroup& g, size_t up_off, size_t gu_half_bytes, size_t down_off, size_t d_bytes, + void* gu_dst, size_t gu_stride, void* d_dst, size_t d_stride, void* stream); /// A Strata-pack Q2_0 expert blob (codes and fp16 scales in separate planes, gate/up rows interleaved) into GGUF /// Q2_0 blocks: gate/up [1280, 2560] at `gu_dst` (rows stay interleaved), down [2560, 640] at `d_dst`. Same values. void gather_strata_q2(const uint8_t* blob, void* gu_dst, void* d_dst, void* stream); diff --git a/src/prefill/moe_mmq.cu b/src/prefill/moe_mmq.cu index 1f0221c7195..c61e707a591 100644 --- a/src/prefill/moe_mmq.cu +++ b/src/prefill/moe_mmq.cu @@ -29,6 +29,22 @@ __global__ void copy16_kernel(const uint4* __restrict__ a, int64_t na, const uin else if (i < na + nb) ab_dst[i] = b[i - na]; else if (i < na + nb + nc) c_dst[i - na - nb] = c[i - na - nb]; } +// copy16_kernel for an MMQ group: blockIdx.y is the expert (first + y) +struct GroupArgs { + const uint8_t* blob[kGatherGroupMax]; + int64_t up_off, down_off, gu_stride, d_stride; // in uint4 +}; +__global__ void copy16_group_kernel(GroupArgs ga, int first, int64_t na, int64_t nc, uint4* __restrict__ gu_dst, + uint4* __restrict__ d_dst) { + const int q = first + (int) blockIdx.y; + const uint4* src = (const uint4*) ga.blob[q]; + uint4* ab = gu_dst + (int64_t) q * ga.gu_stride; + uint4* cd = d_dst + (int64_t) q * ga.d_stride; + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i < na) ab[i] = src[i]; + else if (i < 2 * na) ab[i] = src[ga.up_off + (i - na)]; + else if (i < 2 * na + nc) cd[i - 2 * na] = src[ga.down_off + (i - 2 * na)]; +} __global__ void copy1_kernel(const uint8_t* __restrict__ a, int64_t na, const uint8_t* __restrict__ b, int64_t nb, uint8_t* __restrict__ ab_dst, const uint8_t* __restrict__ c, int64_t nc, uint8_t* __restrict__ c_dst) { const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; @@ -210,6 +226,25 @@ void gather_native(const void* gate, const void* up, size_t gu_half_bytes, const ck(cudaGetLastError(), "gather_native"); } +bool gather_native_group(const GatherGroup& g, size_t up_off, size_t gu_half_bytes, size_t down_off, size_t d_bytes, + void* gu_dst, size_t gu_stride, void* d_dst, size_t d_stride, void* stream) { + if (g.first < 0 || g.n <= g.first || g.n > kGatherGroupMax) return false; + uintptr_t a = (uintptr_t) gu_dst | (uintptr_t) d_dst | up_off | gu_half_bytes | down_off | d_bytes | gu_stride | d_stride; + for (int q = g.first; q < g.n; ++q) a |= (uintptr_t) g.blob[q]; + if (a % 16 != 0) return false; + GroupArgs ga{}; + for (int q = g.first; q < g.n; ++q) ga.blob[q] = g.blob[q]; + ga.up_off = (int64_t) up_off / 16; + ga.down_off = (int64_t) down_off / 16; + ga.gu_stride = (int64_t) gu_stride / 16; + ga.d_stride = (int64_t) d_stride / 16; + const int64_t na = (int64_t) gu_half_bytes / 16, nc = (int64_t) d_bytes / 16; + copy16_group_kernel<<>>( + ga, g.first, na, nc, (uint4*) gu_dst, (uint4*) d_dst); + ck(cudaGetLastError(), "gather_native_group"); + return true; +} + void gather_strata_q2(const uint8_t* blob, void* gu_dst, void* d_dst, void* stream) { strata_q2_kernel<<>>(blob, (uint16_t*) gu_dst, (uint16_t*) d_dst); diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index dc35c2be319..afbacc4a382 100644 --- a/src/prefill/prefill.cpp +++ b/src/prefill/prefill.cpp @@ -56,6 +56,7 @@ Context::Context() {} Context::~Context() {} void Context::run(const Product&, void*) {} void gather_native(const void*, const void*, size_t, const void*, size_t, void*, void*, void*) {} +bool gather_native_group(const GatherGroup&, size_t, size_t, size_t, size_t, void*, size_t, void*, size_t, void*) { return false; } void gather_strata_q2(const uint8_t*, void*, void*, void*) {} void swiglu(const float*, float*, int64_t, int64_t, bool, void*) {} void iota(int32_t*, int64_t, void*) {} @@ -389,6 +390,9 @@ struct Prefill::Impl { std::unique_ptr stager; // the unpinned experts' host copies (step 4) cudaEvent_t copied[RING_MAX] = {}, used[RING_MAX] = {}; bool stage_live[RING_MAX] = {}; + // the event that releases each ring slot: its own `used`, or - when an MMQ group is gathered in one launch - the + // `used` of the last slot gathered with it, recorded once for all of them (a later record only waits longer) + int used_of[RING_MAX] = {}; // PLE float* ple_emb = nullptr; std::vector ple_pageable[2]; // the fallback when no more RAM can be pinned @@ -765,6 +769,7 @@ bool Prefill::carve(size_t T, void* alloc) { for (int i = 0; i < m.ring; ++i) { m.stage_dev[i] = o.take((size_t) MAXBLOB(), ok); m.stage_live[i] = false; // a new buffer: nothing of an earlier layout to wait for + m.used_of[i] = i; } m.ple_emb = o.take(T * N, ok); m.ple_norm = o.take((size_t) strata::kernels::NG_HC_DIM, ok); @@ -1249,7 +1254,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e const int sl = (int) (issued % (size_t) m.ring); const auto th = Clock::now(); const size_t bytes = (size_t) lay0.blob_bytes(en.l); - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); + if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); if (en.job < 0) { cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); ++stats_.experts_dma; @@ -1296,7 +1301,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e const int sl = (int) (idx % (size_t) m.ring); const auto th = Clock::now(); const size_t bytes = (size_t) lay0.blob_bytes(en.l); - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); + if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); if (en.job < 0) { cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); ++iss_dma; @@ -1708,7 +1713,10 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e } else { fused::experts(b, (int) m.g->n_expert, T * K, m.GU, m.Xq, m.src_dev, m.H, m.Dm, m.cs); } - for (size_t kk = k0; kk < k; ++kk) cudaEventRecord(m.used[kk % (size_t) m.ring], m.cs); + for (size_t kk = k0; kk < k; ++kk) { + cudaEventRecord(m.used[kk % (size_t) m.ring], m.cs); + m.used_of[kk % (size_t) m.ring] = (int) (kk % (size_t) m.ring); + } if (k > k0) { consumed = k; give_back(consumed); @@ -1815,13 +1823,13 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e if (pinned && !b) { err = "prefill: expert source has no blob"; return false; } if (pinned) { // DMA straight from the page-locked arena: the copy stream only waits for the slot - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); + if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); cudaMemcpyAsync(m.stage_dev[sl], b, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy); ++stats_.experts_dma; } else { // copied to a pinned buffer by the stager (waits only if it is behind), then DMA const uint8_t* hb = m.stager->wait(job_of[j]); - if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[sl], 0); + if (m.stage_live[sl]) cudaStreamWaitEvent(m.copy, m.used[m.used_of[sl]], 0); cudaMemcpyAsync(m.stage_dev[sl], hb, (size_t) lay.blob_bytes(l), cudaMemcpyHostToDevice, m.copy); m.stager->issued_one(job_of[j], m.copy); } @@ -1832,6 +1840,46 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e ++stats_.experts_streamed; return true; }; + // In the streamed walk an MMQ group is gathered in ONE launch, after ONE wait on its last + // streamed copy (the copy stream is in order), and its ring slots are released by ONE event. A + // wait and a record per expert put ~10 us of GPU bubble each on the compute stream under WDDM + // (226 ms of a 32K prompt on an NVFP4 pack). The same bytes into the same group slots. A ring + // entry the routing skipped inside an open group first gathers what the group holds so far + // (`flush`), so no more than a group's entries are ever held back from the issuer. + // STRATA_PREFILL_GROUP_GATHER=0: one gather, one wait and one record per expert. + static const bool group_env = [] { + const char* v = std::getenv("STRATA_PREFILL_GROUP_GATHER"); + return v == nullptr || std::atoi(v) != 0; + }(); + const bool group_gather = group_env && stream_all && use_mmq && lay.native && + MMQ_GROUP <= mmq::kGatherGroupMax; + mmq::GatherGroup gg; + int gg_slots[MMQ_GROUP]; + int gg_nslots = 0; // ring slots gathered by the next flush + auto flush = [&]() { + if (gg.n <= gg.first) return; + const auto& f = lay.fmt[(size_t) l]; + if (gg_nslots > 0) { // the copies land in order: the last one covers the others + pt.mark(kPfWaitCopy, cs); + cudaStreamWaitEvent(m.cs, m.copied[gg_slots[gg_nslots - 1]], 0); + pt.mark(kPfDequant, cs); + } + if (!mmq::gather_native_group(gg, f.up_off, mmq_gub / 2, f.down_off, mmq_db, m.grp_gu, mmq_gub, + m.grp_d, mmq_db, m.cs)) { + for (int i = gg.first; i < gg.n; ++i) { // not 16-byte aligned: one at a time + const uint8_t* b = gg.blob[i]; + mmq::gather_native(b, b + f.up_off, mmq_gub / 2, b + f.down_off, mmq_db, + m.grp_gu + i * mmq_gub, m.grp_d + i * mmq_db, m.cs); + } + } + if (gg_nslots > 0) { + const int rel = gg_slots[gg_nslots - 1]; + cudaEventRecord(m.used[rel], m.cs); + for (int i = 0; i < gg_nslots; ++i) m.used_of[gg_slots[i]] = rel; + } + gg_nslots = 0; + gg.first = gg.n; + }; // one expert's products from its blob on the device; `slot` (a ring slot, or -1 for a resident // expert) is released once the blob is read auto compute = [&](size_t j, const uint8_t* blob_dev, int slot) -> bool { @@ -1840,14 +1888,21 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e if (use_mmq) { // gather the expert into its group slot (GGUF blocks, unchanged or converted) const size_t q = j % MMQ_GROUP; - if (lay.native) { + if (group_gather) { + gg.blob[q] = blob_dev; + gg.n = (int) q + 1; + if (slot >= 0) gg_slots[gg_nslots++] = slot; + if (q + 1 < MMQ_GROUP && j + 1 < order.size()) return true; + flush(); + gg = mmq::GatherGroup{}; + } else if (lay.native) { const auto& f = lay.fmt[(size_t) l]; mmq::gather_native(blob_dev, blob_dev + f.up_off, mmq_gub / 2, blob_dev + f.down_off, mmq_db, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs); } else { mmq::gather_strata_q2(blob_dev, m.grp_gu + q * mmq_gub, m.grp_d + q * mmq_db, m.cs); } - if (slot >= 0) cudaEventRecord(m.used[slot], m.cs); + if (slot >= 0 && !group_gather) { cudaEventRecord(m.used[slot], m.cs); m.used_of[slot] = slot; } if (q + 1 < MMQ_GROUP && j + 1 < order.size()) return true; // the group's products: gate/up, swiglu, the group's H to q8_1, down const size_t j0 = j - q, g = j0 / MMQ_GROUP, n = order.size(); @@ -1885,7 +1940,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e } else { blob_dequant_f16(blob_dev, m.dq_gu[q], m.dq_d[q], m.cs); } - if (slot >= 0) cudaEventRecord(m.used[slot], m.cs); + if (slot >= 0) { cudaEventRecord(m.used[slot], m.cs); m.used_of[slot] = slot; } const int64_t o0 = m.off[(size_t) e], ne = m.cnt[(size_t) e]; pt.mark(kPfGemmGU, cs); m.gemm.f16(m.Xs + o0 * N, m.dq_gu[q], m.GU + o0 * 1280, ne, 1280, N); @@ -1919,7 +1974,10 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e const size_t kend = seq_start[(size_t) l + 1]; auto release_to = [&](int32_t e_stop) { while (k < kend && seq[k].e < e_stop) { - cudaEventRecord(m.used[k % (size_t) m.ring], m.cs); + if (gg_nslots > 0) flush(); // the open group's slots get their event first + const int sl = (int) (k % (size_t) m.ring); + cudaEventRecord(m.used[sl], m.cs); + m.used_of[sl] = sl; consumed = ++k; give_back(consumed); } @@ -1929,15 +1987,18 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e release_to(e); if (k < kend && seq[k].e == e) { const int sl = (int) (k % (size_t) m.ring); - pt.mark(kPfWaitCopy, cs); wait_issued(k); - cudaStreamWaitEvent(m.cs, m.copied[sl], 0); + if (!group_gather) { + pt.mark(kPfWaitCopy, cs); + cudaStreamWaitEvent(m.cs, m.copied[sl], 0); + } if (!compute(j, m.stage_dev[sl], sl)) return false; consumed = ++k; - give_back(consumed); + if (!group_gather || gg_nslots == 0) give_back(consumed); // its group was gathered } else { ++stats_.experts_resident; if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false; + if (group_gather && gg_nslots == 0) give_back(consumed); } } release_to(m.g->n_expert);