Skip to content

Prompt path: an MMQ group's experts gathered in one launch, after one wait, released by one event (-9% on a 32K prompt, same output) #372

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
sergqwer wants to merge 1 commit into from
Closed

Prompt path: an MMQ group's experts gathered in one launch, after one wait, released by one event (-9% on a 32K prompt, same output) #372

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions include/strata/prefill/moe_mmq.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
35 changes: 35 additions & 0 deletions src/prefill/moe_mmq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<<<dim3(blocks(2 * na + nc), (unsigned) (g.n - g.first)), 256, 0, (cudaStream_t) stream>>>(
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<<<blocks(1280LL * 40 + 2560LL * 10), 256, 0, (cudaStream_t) stream>>>(blob, (uint16_t*) gu_dst,
(uint16_t*) d_dst);
Expand Down
85 changes: 73 additions & 12 deletions src/prefill/prefill.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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*) {}
Expand Down Expand Up @@ -389,6 +390,9 @@ struct Prefill::Impl {
std::unique_ptr<Stager> 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<float> ple_pageable[2]; // the fallback when no more RAM can be pinned
Expand Down Expand Up @@ -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<uint8_t>((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<float>(T * N, ok);
m.ple_norm = o.take<float>((size_t) strata::kernels::NG_HC_DIM, ok);
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
Expand All @@ -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 {
Expand All @@ -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();
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
Expand Down