Download src/prefill/prefill.cpp from WineryLabs/Winery-Strata: direct link, hf CLI and curl.
- Browser
- Download file 173 kB
-
https://huggingface.co/WineryLabs/Winery-Strata/resolve/main/src/prefill/prefill.cpp
- Command line
-
hf download hf://WineryLabs/Winery-Strata/src/prefill/prefill.cpp
-
curl -L -o prefill.cpp https://huggingface.co/WineryLabs/Winery-Strata/resolve/main/src/prefill/prefill.cpp
173 kB
| // src/prefill/prefill.cpp - see include/strata/prefill/prefill.hpp. | |
| // A build without the llama.cpp sources (no STRATA_NATIVE_EXPERTS): no MMQ, the FP16 expert path everywhere. | |
| namespace strata::prefill::mmq { | |
| bool built() { return false; } | |
| bool supported(int) { return false; } | |
| bool fits(int, int64_t) { return false; } | |
| size_t matrix_bytes(int, int64_t, int64_t) { return 0; } | |
| size_t q8_bytes(int64_t, int64_t) { return 0; } | |
| void quantize(const float*, const int32_t*, void*, int, int64_t, int64_t, int64_t, void*) {} | |
| 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*) {} | |
| } // namespace strata::prefill::mmq | |
| // #136: the fused int8 experts are CUDA-only (HIP and builds without MMQ keep the MMQ / FP16 paths) | |
| namespace strata::prefill::fused { | |
| bool built() { return false; } | |
| bool available() { return false; } | |
| bool enabled() { return false; } | |
| bool requested() { return false; } | |
| size_t act_bytes(int64_t, int64_t) { return 0; } | |
| size_t group_bytes(int64_t, int) { return 0; } | |
| void quantize_act(const float*, int64_t, int64_t, void*, void*) {} | |
| void group(const int32_t*, int64_t, int, int, void*, int32_t*, int32_t*, void*) {} | |
| void experts(const Batch&, int, int64_t, const void*, const void*, const int32_t*, void*, float*, void*) {} | |
| bool native_supported(int, int) { return false; } | |
| void quantize_act_native(const float*, int64_t, int64_t, void*, void*) {} | |
| void experts_native(const Batch&, const NativeGeom&, int, int64_t, const void*, const void*, const int32_t*, void*, | |
| float*, void*) {} | |
| } // namespace strata::prefill::fused | |
| namespace strata::prefill { | |
| namespace { | |
| using Clock = std::chrono::steady_clock; | |
| constexpr float EPS = 1e-6f; | |
| constexpr int64_t N = 2560, HC = 4, D = N * HC, LR = 320, K = 10, NE = 512; | |
| constexpr int64_t C = 10240, ZV = 6144, HV = 48; | |
| // plan v0.3 P6: staging holds the largest blob of the pack (a native pack's blobs differ per layer) | |
| inline int64_t MAXBLOB() { return (int64_t) strata::kernels::cpu::expert_layout().max_blob; } | |
| constexpr int STAGE = 8; // host->device expert staging ring (chunks below stream_all_min()) | |
| // Step 3: from this chunk size on, every non-resident expert of every layer streams in a fixed order through a | |
| // ring_slots()-slot ring (nearly all 512 are routed at such a chunk), so the copy engine keeps working through the | |
| // attention halves instead of waiting for each layer's routing. | |
| constexpr int RING_MAX = 1024; // the arrays; the ring itself is ring_slots(), at most ring_cap() | |
| // The chunk size from which every expert streams: 1024 since 0.1.30 (was 2048). Measured on the 5070, Q2_0 / IQ2_XS, | |
| // fixed cache: 1,500-token prompts 621 -> 785 / 612 -> 735 tok/s, 2,000 727 -> 934 / 712 -> 892, 4,000 (its last | |
| // chunk) 779 -> 912 / 766 -> 844, the same output. Below ~1,000 tokens the output changed on Q2_0 (a smaller chunk | |
| // takes other kernels), so 1024 is the floor. STRATA_PREFILL_STREAM_MIN overrides (A/B). | |
| inline int64_t stream_all_min() { | |
| static const int64_t v = [] { const char* e = std::getenv("STRATA_PREFILL_STREAM_MIN"); return e ? (int64_t) std::atoll(e) : (int64_t) 1024; }(); | |
| return v; | |
| } | |
| double g_pinned_share = 1.0; | |
| // The streamed ring: 384 slots when (nearly) every streamed expert is DMA'd from pinned RAM - measured on Q2_0, | |
| // 8192-token chunks: 96 slots 1153 tok/s, 384 1294 (the next layer's experts arrive during its attention half) - | |
| // and 96 when a large share goes through host copies (IQ3_S on 64 GB, a third unpinned: 96 slots 1216, 256 1070 - | |
| // the host copies are the limit and the bigger ring only takes cache slots). STRATA_PREFILL_RING overrides. | |
| int g_ring_override = 0; // #340: set by a layer split (Prefill::set_ring_override); 0 = the rule below | |
| // #136: the fused experts (STRATA_PF_FUSED=1) launch on a batch of a layer's streamed experts at once, so the ring | |
| // should hold a whole layer's (~460 of 512 on Q2_0): with 384 slots a layer's last batch waits for slots its own | |
| // first batch frees. Measured on the 5070, Q2_0, the 4K / 32K code-agent prompts (one run each): fused at 384 slots | |
| // +5% / +3% over MMQ, at 512 +19-22% / +11-12% (MMQ itself at 512: -2% / -1%). P3: with the fused path's smaller | |
| // buffers (moe_bufs) 1024 slots - two layers' experts - fit too; 2 pairs each, prompt tok/s at 512 / 1024 slots: 4K | |
| // 1,512 / 1,605, 32K 2,485 / 2,660 (both chunk 8192), 128K with KV streaming 2,346 / 2,392 (the chunk falls from | |
| // 8192 to 6144, but more experts stay resident). A native pack likewise when the native kernels (moe_fused_iq.hpp) | |
| // take any of its layers: IQ2_XS, 4K / 32K, their first version at 384 slots -3% / -6% against MMQ, at 512 +8% / 0%. | |
| inline bool fused_ring() { | |
| if (!fused::enabled()) return false; | |
| if (core::peer_portable()) return false; // multi-GPU: --peer-device keeps the MMQ path and its buffer sizes | |
| const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout(); | |
| if (!lay.native) return true; | |
| // EVERY layer: fused_layout() shrinks the MoE buffers to the fused path's needs, so a layer the native kernels do | |
| // not cover (Unsloth UD-IQ4_XS's Q8_0 down projections) would run MMQ in them at the full chunk and overflow them | |
| // (garbage, an illegal memory access or a hung prompt on gfx1151). A pack with such a layer keeps MMQ's buffers; | |
| // its covered layers still take the fused kernels. | |
| static const bool all = [&lay] { | |
| for (const auto& f : lay.fmt) | |
| if (!fused::native_supported(f.gu_type, f.d_type)) return false; | |
| return !lay.fmt.empty(); | |
| }(); | |
| return all; | |
| } | |
| // the largest ring: 512 slots; 1024 with the Q2_0 pack's fused experts (P3's smaller buffers, measured there) - the | |
| // native packs' fused layers were measured at 512 | |
| inline int ring_cap() { return fused_ring() && !strata::kernels::cpu::expert_layout().native ? RING_MAX : 512; } | |
| inline int ring_slots(size_t T) { | |
| const char* v = std::getenv("STRATA_PREFILL_RING"); | |
| // S6: with the opt-in RDNA4 matrix-core attention (STRATA_HIP_WMMA=1) a 96-slot ring: measured with it, 9070 XT | |
| // 4K prompts 718 -> 1,211 tok/s (16K 1,949 -> 2,032), R9700 4K 2,426 -> 2,483 (16K the same) | |
| static const bool wmma = [] { | |
| const char* e = std::getenv("STRATA_HIP_WMMA"); | |
| return e != nullptr && e[0] == '1'; | |
| }(); | |
| if (!v && g_ring_override <= 0 && wmma) return (int64_t) T >= stream_all_min() ? 96 : STAGE; | |
| const int pinned_ring = fused_ring() ? 1024 : 384; | |
| const int r = v ? std::atoi(v) : g_ring_override > 0 ? g_ring_override : (g_pinned_share >= 0.9 ? pinned_ring : 96); | |
| if (v && r == STAGE) return STAGE; // Explicit opt-in to routed-only staging, including large chunks. | |
| const int big = r < 16 ? 16 : r > ring_cap() ? ring_cap() : r; | |
| return (int64_t) T >= stream_all_min() ? big : STAGE; | |
| } | |
| constexpr int DQ = 2; // dequantized-expert ring (FP16 gate/up + down) | |
| // The BF16-weight projections (hyper-connection, SSM alpha/beta, indexer, router, shared gate, PLE key/value) take | |
| // BF16 activations here and FP32 ones in decode. STRATA_PREFILL_BF16X2=1 adds each activation's BF16 remainder as a | |
| // second GEMM (Y = W.hi + W.lo, ~16 mantissa bits): a router that picks its top 10 from the same x decode would. | |
| // 2 = all but the hyper-connection's; 1 = the hyper-connection's too (its activations are 10240 wide and its up | |
| // projection writes as much: slower); 0 (the default: opt-in, it changes the prompt path's numbers) = off. | |
| inline int bf16x2_mode() { | |
| static const int v = [] { | |
| const char* e = std::getenv("STRATA_PREFILL_BF16X2"); | |
| return e != nullptr ? std::atoi(e) : 0; | |
| }(); | |
| return v; | |
| } | |
| inline bool bf16x2() { return bf16x2_mode() != 0; } | |
| inline bool bf16x2_hc() { return bf16x2_mode() == 1; } | |
| // F-1: STRATA_GR_UNFUSED=1 keeps the FP32 copy of the normalized rows (gr_norm + gr_mix), the A/B arm | |
| inline bool gr_unfused() { | |
| static const bool v = [] { const char* e = std::getenv("STRATA_GR_UNFUSED"); return e && e[0] == '1'; }(); | |
| return v; | |
| } | |
| double ms_since(Clock::time_point t) { return std::chrono::duration<double, std::milli>(Clock::now() - t).count(); } | |
| // Either cudaMalloc (owned, freed with the object) or a bump allocation from a borrowed region; with no base and | |
| // no region it only counts, which is how `bytes_needed` sizes the region. | |
| struct Alloc { | |
| uint8_t* base = nullptr; | |
| uint64_t cap = 0, used = 0; | |
| bool count_only = false; | |
| std::vector<void*>* owned = nullptr; | |
| template <typename T> T* take(size_t n, bool& ok) { | |
| const uint64_t bytes = ((uint64_t) n * sizeof(T) + 256 + 255) & ~255ull; | |
| if (count_only) { used += bytes; return nullptr; } | |
| if (base != nullptr) { | |
| if (used + bytes > cap) { ok = false; return nullptr; } | |
| T* p = (T*) (base + used); | |
| used += bytes; | |
| return p; | |
| } | |
| void* p = nullptr; | |
| if (cudaMalloc(&p, bytes) != cudaSuccess) { ok = false; return nullptr; } | |
| owned->push_back(p); | |
| used += bytes; | |
| return (T*) p; | |
| } | |
| }; | |
| } // namespace | |
| // Step 4 of the prompt-speed plan: the experts the arena could not pin (a third of the streamed ones on IQ3_S) are | |
| // copied into pinned buffers by these threads, ahead of the launches. Copied in line by the launching thread they | |
| // left the GPU without queued work while each ~2 MB memcpy ran (~15 s of a 32K prompt on IQ3_S). Job j - a layer's | |
| // j-th unpinned expert, in launch order - lands in host buffer j % kRing, which is free again once the DMA of job | |
| // j - kRing (recorded by the launching thread, `issued`) is done. | |
| struct Stager { | |
| // D-5: the pinned ring's depth (STRATA_STAGER_RING, default 16) - how far the host copies can run ahead of the | |
| // DMAs of the unpinned experts' blobs | |
| int kRing = 16; | |
| // `from` set: the blob is copied by the source itself (CS-T: a GGUF read in place assembles it from its three | |
| // role slices; a pointer to it would not live as long as the queue) | |
| struct Job { const uint8_t* src; size_t bytes; core::ExpertSource* from = nullptr; int32_t l = 0, e = 0; }; | |
| std::vector<uint8_t*> buf; | |
| std::vector<char> pinned; | |
| std::vector<std::vector<uint8_t>> pageable; // the fallback when no more RAM can be pinned | |
| std::vector<cudaEvent_t> dma_done; | |
| std::vector<Job> jobs; | |
| std::unique_ptr<std::atomic<int>[]> ready; | |
| size_t ready_cap = 0; | |
| // gen << 32 | n << 16 | next index: a claim is a CAS on the generation it woke for (a thread late from the | |
| // previous layer can never take a job of this one - the expert pool's issue #29 lesson) | |
| std::atomic<uint64_t> head{0}; | |
| std::atomic<int> issued{0}, active{0}; | |
| uint32_t gen = 0; | |
| bool quit = false; | |
| std::mutex mu; | |
| std::condition_variable cv; | |
| std::vector<std::thread> threads; | |
| int device = 0; | |
| bool init(size_t blob_bytes, int nthreads) { | |
| if (const char* v = std::getenv("STRATA_STAGER_RING")) kRing = std::clamp(std::atoi(v), 2, 256); | |
| buf.assign((size_t) kRing, nullptr); | |
| pinned.assign((size_t) kRing, 0); | |
| dma_done.assign((size_t) kRing, nullptr); | |
| pageable.resize(kRing); | |
| for (int i = 0; i < kRing; ++i) { | |
| pinned[i] = cudaHostAlloc((void**) &buf[i], blob_bytes, cudaHostAllocDefault) == cudaSuccess; | |
| if (!pinned[i]) { | |
| cudaGetLastError(); | |
| pageable[(size_t) i].resize(blob_bytes); | |
| buf[i] = pageable[(size_t) i].data(); | |
| } | |
| if (cudaEventCreateWithFlags(&dma_done[i], cudaEventDisableTiming) != cudaSuccess) return false; | |
| } | |
| cudaGetDevice(&device); | |
| for (int t = 0; t < nthreads; ++t) threads.emplace_back([this] { work(); }); | |
| return true; | |
| } | |
| ~Stager() { | |
| finish(); | |
| { std::lock_guard<std::mutex> lk(mu); quit = true; } | |
| cv.notify_all(); | |
| for (auto& t : threads) t.join(); | |
| for (int i = 0; i < kRing; ++i) { | |
| if (dma_done[i]) cudaEventDestroy(dma_done[i]); | |
| if (buf[i] && pinned[i]) cudaFreeHost(buf[i]); | |
| } | |
| } | |
| void work() { | |
| cudaSetDevice(device); | |
| uint32_t seen = 0; | |
| for (;;) { | |
| { | |
| std::unique_lock<std::mutex> lk(mu); | |
| cv.wait(lk, [&] { return quit || gen != seen; }); | |
| if (quit) return; | |
| seen = gen; | |
| } | |
| for (;;) { | |
| active.fetch_add(1, std::memory_order_acq_rel); | |
| const int j = claim(seen); | |
| if (j < 0) { active.fetch_sub(1, std::memory_order_acq_rel); break; } | |
| const int b = j % kRing; | |
| if (j >= kRing) // job j - kRing's DMA from this buffer is queued | |
| while (issued.load(std::memory_order_acquire) <= j - kRing) std::this_thread::yield(); | |
| // and done - for a generation's first kRing jobs that is the previous generation's last DMA from | |
| // the buffer, which nothing else waits for when a chunk ends without a sync (no MTP) or the DMA | |
| // was a ring entry the routing skipped (an event never recorded returns at once) | |
| cudaEventSynchronize(dma_done[b]); | |
| const Job& jb = jobs[(size_t) j]; | |
| if (jb.from == nullptr) std::memcpy(buf[b], jb.src, jb.bytes); | |
| else if (!jb.from->copy_blob(jb.l, jb.e, buf[b])) { | |
| std::fprintf(stderr, "prefill: the expert source could not copy expert %d of layer %d\n", jb.e, jb.l); | |
| std::abort(); | |
| } | |
| ready[(size_t) j].store(1, std::memory_order_release); | |
| active.fetch_sub(1, std::memory_order_acq_rel); | |
| } | |
| } | |
| } | |
| int claim(uint32_t g) { | |
| uint64_t cur = head.load(std::memory_order_acquire); | |
| for (;;) { | |
| if ((uint32_t) (cur >> 32) != g) return -1; | |
| const int n = (int) ((cur >> 16) & 0xffff), j = (int) (cur & 0xffff); | |
| if (j >= n) return -1; | |
| if (head.compare_exchange_weak(cur, cur + 1, std::memory_order_acq_rel, std::memory_order_acquire)) return j; | |
| } | |
| } | |
| /// A layer's jobs; the previous layer's are finished (finish()). | |
| void start(std::vector<Job>&& js) { | |
| if (js.empty()) return; | |
| std::lock_guard<std::mutex> lk(mu); | |
| jobs = std::move(js); | |
| if (ready_cap < jobs.size()) { | |
| ready_cap = jobs.size() * 2; | |
| ready.reset(new std::atomic<int>[ready_cap]); | |
| } | |
| for (size_t i = 0; i < jobs.size(); ++i) ready[i].store(0, std::memory_order_relaxed); | |
| issued.store(0); | |
| ++gen; | |
| head.store((uint64_t) gen << 32 | (uint64_t) jobs.size() << 16, std::memory_order_release); | |
| cv.notify_all(); | |
| } | |
| /// Job j's bytes, in a pinned buffer (waits for the copy). | |
| const uint8_t* wait(int j) { | |
| while (!ready[(size_t) j].load(std::memory_order_acquire)) std::this_thread::yield(); | |
| return buf[j % kRing]; | |
| } | |
| /// The launching thread queued job j's DMA on `copy`: its buffer is free once that is done. | |
| void issued_one(int j, cudaStream_t copy) { | |
| cudaEventRecord(dma_done[j % kRing], copy); | |
| issued.store(j + 1, std::memory_order_release); | |
| } | |
| /// No job is running after this (the end of a layer, or an early return in the middle of one). | |
| void finish() { | |
| head.store((uint64_t) gen << 32, std::memory_order_release); // n = 0: nothing more to claim | |
| issued.store(1 << 30, std::memory_order_release); | |
| while (active.load(std::memory_order_acquire) != 0) std::this_thread::yield(); | |
| } | |
| }; | |
| // multi-GPU: the peer GPU's share of a prompt chunk's experts. Per MoE layer the primary copies its normed | |
| // activations over P2P, the peer quantizes the rows routed to the experts it holds, runs the same MMQ products the | |
| // primary would (gathered from its own slots), and copies the result rows back into the primary's Dm rows - the rows | |
| // are laid out local-first, so the peer's are one contiguous block at the end. | |
| struct PeerPrefill { | |
| core::PeerExperts* peer = nullptr; | |
| int dev = -1; | |
| int64_t cap_rows = 0, T_max = 0; | |
| cudaStream_t s = nullptr; | |
| cudaEvent_t ev_in = nullptr; // on the primary: the activations and the row tables are ready | |
| cudaEvent_t ev_done = nullptr; // on the peer: its rows have landed in the primary's Dm | |
| // multi-GPU: the result rows go back group by group on a second stream while the next group computes | |
| cudaStream_t s_out = nullptr; | |
| static constexpr int kGrpEv = NE / 16 + 2; | |
| cudaEvent_t ev_grp[kGrpEv] = {}; | |
| bool out_pending = false; // s_out may still read Dm (the next layer's products wait for it) | |
| bool out_pipe = true; | |
| // multi-GPU COMPACT: the peer's per-row buffers hold one GROUP's rows instead of the layer's (a group | |
| // = up to 16 experts and at most G rows), the result rows double-buffered: ~0.9 GB -> ~0.26 GB on the helper | |
| // card at 8192-token chunks, and no row cap (every peer-held expert's rows go to the peer). | |
| // STRATA_PF_PEER_COMPACT=0: the layer-sized buffers (the A/B). | |
| bool compact = true; | |
| int64_t G = 0; | |
| void *Xq_g = nullptr, *Hq_g = nullptr; | |
| float *GU_g = nullptr, *H_g = nullptr, *Dm_b[2] = {}; | |
| cudaEvent_t ev_dm[2] = {}; | |
| bool dm_live[2] = {}; | |
| std::vector<std::pair<size_t, size_t>> groups; | |
| // multi-GPU PEER STREAMING: a share of the experts the primary would stream over ITS PCIe link in a | |
| // big chunk is streamed by the peer over its own link instead (into its own ring) and computed there - it halves | |
| // the primary's copy load and moves expert work to the card that idles through most of the MoE half. | |
| // STRATA_PF_PEER_STREAM = the share (0 = off), STRATA_PF_PEER_RING = its ring slots. | |
| double ps_frac = 0.0; | |
| int RP = 0; | |
| std::vector<uint8_t*> pstage; | |
| std::vector<cudaEvent_t> pcopied, pused; | |
| std::vector<char> plive; | |
| cudaStream_t s_cp = nullptr; | |
| struct PsEntry { int32_t l, e; const uint8_t* blob; }; | |
| std::vector<PsEntry> pseq; // this chunk's peer-streamed experts, layer by layer in id order | |
| std::vector<size_t> pseq_start; | |
| std::vector<char> ps_flag; // [layer * NE + e]: streamed by the peer in this chunk | |
| size_t p_issued = 0, pk = 0; | |
| int64_t ps_experts = 0; | |
| float *mixed = nullptr, *GU = nullptr, *H = nullptr, *Dm = nullptr; | |
| void *Xq = nullptr, *Hq = nullptr; | |
| int32_t *src = nullptr, *bounds = nullptr, *ident = nullptr; | |
| uint8_t *grp_gu = nullptr, *grp_d = nullptr; | |
| std::unique_ptr<mmq::Context> ctx; | |
| std::vector<int32_t> bounds_host; | |
| std::vector<void*> owned; | |
| int64_t layers = 0, experts = 0, rows = 0, over_cap = 0; // stats | |
| ~PeerPrefill() { | |
| if (dev < 0) return; | |
| int prev = 0; | |
| cudaGetDevice(&prev); | |
| cudaSetDevice(dev); | |
| if (s) cudaStreamSynchronize(s); | |
| ctx.reset(); | |
| for (void* p : owned) cudaFree(p); | |
| if (ev_done) cudaEventDestroy(ev_done); | |
| if (s_out) cudaStreamSynchronize(s_out); | |
| for (cudaEvent_t e : ev_grp) if (e) cudaEventDestroy(e); | |
| for (cudaEvent_t e : ev_dm) if (e) cudaEventDestroy(e); | |
| if (s_cp) cudaStreamSynchronize(s_cp); | |
| for (cudaEvent_t e : pcopied) if (e) cudaEventDestroy(e); | |
| for (cudaEvent_t e : pused) if (e) cudaEventDestroy(e); | |
| if (s_cp) cudaStreamDestroy(s_cp); | |
| if (s_out) cudaStreamDestroy(s_out); | |
| if (s) cudaStreamDestroy(s); | |
| cudaSetDevice(prev); | |
| if (ev_in) cudaEventDestroy(ev_in); | |
| } | |
| }; | |
| struct Prefill::Impl { | |
| const core::WeightTable* wt = nullptr; | |
| const core::ModelGeometry* g = nullptr; | |
| core::SessionState* ss = nullptr; | |
| core::ExpertSource* src = nullptr; | |
| const core::ExpertCache* cache = nullptr; | |
| const int32_t* host_res = nullptr; | |
| int64_t T = 0, T_max = 0; | |
| bool borrowed = false; | |
| cudaStream_t cs = nullptr, copy = nullptr; | |
| Gemm gemm; | |
| std::vector<void*> owned; | |
| // chunk buffers | |
| float *emb = nullptr, *R = nullptr, *xn = nullptr, *lo = nullptr, *gated = nullptr, *inj = nullptr; | |
| float* grs = nullptr; // F-1: the hyper-connection read's row scales (T x 4) | |
| uint16_t *xn16 = nullptr, *lo16 = nullptr; | |
| float* mixed = nullptr; | |
| uint16_t *mixed_bf = nullptr, *mixed_h = nullptr; | |
| uint16_t *xn16_lo = nullptr, *lo16_lo = nullptr, *mixed_bf_lo = nullptr; // bf16x2(): the BF16 GEMMs' low parts | |
| float* bo = nullptr; | |
| // GDN | |
| float *qkv = nullptr, *z = nullptr, *ab = nullptr, *gate = nullptr, *beta = nullptr, *hbuf = nullptr, *y = nullptr; | |
| uint16_t* y_h = nullptr; | |
| // QSA | |
| float *Kc = nullptr, *Vc = nullptr, *Qf = nullptr, *q = nullptr, *idx_raw = nullptr, *q_idx = nullptr, *attn = nullptr; | |
| uint16_t* attn_h = nullptr; | |
| int32_t* steps_dev = nullptr; | |
| std::vector<int32_t> steps_host; | |
| int32_t* sel_ids = nullptr; | |
| float* sel_scores = nullptr; // [sel_batch, max_blocks] | |
| int64_t sel_batch = 256, max_blocks = 0; | |
| float* attn_scratch = nullptr; | |
| int64_t attn_batch = 32, cap = 0; | |
| // MoE | |
| float *logits = nullptr, *w = nullptr, *GU = nullptr, *Dm = nullptr, *sgate = nullptr, *sup = nullptr, | |
| *shared = nullptr, *sg = nullptr; | |
| int32_t *ids = nullptr, *slot_dev = nullptr, *src_dev = nullptr; | |
| uint16_t *Xs = nullptr, *Hh = nullptr, *sh_h = nullptr; | |
| // step 2b (MMQ): the activations quantized per layer, H in FP32 and its group's quantized rows, the identity | |
| // row map, the group bounds, the group buffers of gathered experts | |
| void *Xq = nullptr, *Hq = nullptr; | |
| float* H = nullptr; | |
| int32_t *ids_identity = nullptr, *bounds_dev = nullptr; | |
| uint8_t *grp_gu = nullptr, *grp_d = nullptr; | |
| std::vector<int32_t> bounds_host; | |
| std::unique_ptr<mmq::Context> mmq_ctx; | |
| std::vector<int32_t> ids_host, slot_host, src_host, cnt, off; | |
| // The grouping tables in mapped pinned memory, [ids | slot | src] of T_max * K each, then the MMQ bounds: kernels | |
| // read and write them in place. A cudaMemcpyAsync of them queues behind the expert blobs the copy stream already | |
| // holds (up to `ring` of them, ~70 us each), and the GPU idles meanwhile - measured 4.2 s of a 128K prompt's | |
| // 70 s at the default ring, 0.7 s with a 16-slot one. STRATA_GROUP_COPY=1: the copies (the A/B arm). | |
| int32_t* grp_host = nullptr; | |
| int32_t* grp_dev = nullptr; // its device alias | |
| size_t grp_n = 0, grp_tk = 0; // int32s allocated; T_max * K (the offset of slot, and of src past it) | |
| uint16_t* dq_gu[DQ] = {}; | |
| uint16_t* dq_d[DQ] = {}; | |
| uint8_t* stage_dev[RING_MAX] = {}; | |
| int ring = STAGE; // the slots of this layout's ring (ring_slots) | |
| 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 | |
| float* ple_emb_host[2] = {}; // pinned, double-buffered: the next chunk's rows are read while this | |
| cudaEvent_t ple_copied[2] = {}; // one runs; the event marks that buffer's upload done | |
| std::vector<uint32_t> ple_rows[2]; | |
| float* ple_norm = nullptr; | |
| uint8_t* region = nullptr; // the attention/MoE scratch region (idle while the PLE block runs) | |
| uint64_t region_bytes = 0; | |
| PrefillStats* stats = nullptr; | |
| // KV streaming: one layer's whole K/V, staged from the host copy per layer and chunk (identity layout) | |
| strata::kernels::KvHostPools stage; | |
| int32_t* ident_table = nullptr; | |
| // layer split: the device, and the hand-off to the next stage (two pinned chunk buffers, used in turn) | |
| int device = -1; | |
| float* hand[2] = {}; | |
| // C-4: the chunk's token ids on the device, for one batched embedding gather | |
| int32_t* tok_dev = nullptr; | |
| std::vector<int32_t> tok_host; | |
| std::unique_ptr<PeerPrefill> pp; // multi-GPU: the peer GPU's expert share (set_peer) | |
| }; | |
| namespace { | |
| // the staging pool of a streamed session: every page of one layer (same sequence in init and bytes_needed) | |
| // STRATA_KV_STAGE_OWN (A/B only): the staging pool gets its own allocation instead of borrowed expert slots, so a | |
| // streamed run lends the prompt path exactly the slots a resident one does (a lent expert runs on the CPU, which | |
| // rounds differently: without this an A/B compares two expert placements as well as two KV placements) | |
| bool stage_own() { static const bool v = std::getenv("STRATA_KV_STAGE_OWN") != nullptr; return v; } | |
| void take_stage(Alloc& o_borrowed, const core::SessionState& ss, const strata::kernels::QsaShapes& s, | |
| strata::kernels::KvHostPools& st, bool& ok) { | |
| const core::QsaState& q0 = ss.qsa_states[ss.qsa_primary()]; | |
| if (q0.kv_mode != 1) return; | |
| if (stage_own() && o_borrowed.count_only) return; | |
| Alloc own; | |
| own.owned = o_borrowed.owned; | |
| Alloc& o = stage_own() ? own : o_borrowed; | |
| const size_t rows = (size_t) q0.n_pages * s.n_head_kv * s.page_size; | |
| if (q0.kv_q4) { | |
| st.k_q4 = o.take<uint8_t>(rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim), ok); | |
| st.v_q4 = o.take<uint8_t>(rows * strata::kernels::kv_q4_bytes_per_head((int) s.head_dim), ok); | |
| } else if (q0.kv_int8) { | |
| st.k_q = o.take<int8_t>(rows * s.head_dim, ok); | |
| st.v_q = o.take<int8_t>(rows * s.head_dim, ok); | |
| st.k_scale = o.take<uint16_t>(rows * (s.head_dim / 64), ok); | |
| st.v_scale = o.take<uint16_t>(rows * (s.head_dim / 64), ok); | |
| } else { | |
| st.k_pool = o.take<uint16_t>(rows * s.head_dim, ok); | |
| st.v_pool = o.take<uint16_t>(rows * s.head_dim, ok); | |
| } | |
| } | |
| strata::kernels::QsaAttnPools pools_of(const strata::kernels::KvHostPools& h, const int32_t* table) { | |
| strata::kernels::QsaAttnPools p; | |
| p.k_pool = h.k_pool; p.v_pool = h.v_pool; p.k_q = h.k_q; p.v_q = h.v_q; p.k_scale = h.k_scale; p.v_scale = h.v_scale; | |
| p.k_q4 = h.k_q4; p.v_q4 = h.v_q4; | |
| p.page_table = table; | |
| return p; | |
| } | |
| } // namespace | |
| Prefill::Prefill() : impl_(new Impl) {} | |
| Prefill::~Prefill() { release(); } | |
| void Prefill::reset() { | |
| release(); | |
| impl_.reset(new Impl); | |
| stats_ = PrefillStats{}; | |
| } | |
| void Prefill::release() { | |
| if (!impl_) return; | |
| if (impl_->cs) cudaStreamSynchronize(impl_->cs); | |
| if (impl_->copy) cudaStreamSynchronize(impl_->copy); | |
| for (int i = 0; i < RING_MAX; ++i) { | |
| if (impl_->copied[i]) cudaEventDestroy(impl_->copied[i]); | |
| if (impl_->used[i]) cudaEventDestroy(impl_->used[i]); | |
| } | |
| for (int b = 0; b < 2; ++b) { | |
| if (impl_->hand[b]) cudaFreeHost(impl_->hand[b]); | |
| if (impl_->ple_copied[b]) cudaEventDestroy(impl_->ple_copied[b]); | |
| if (impl_->ple_emb_host[b] && impl_->ple_pageable[b].empty()) cudaFreeHost(impl_->ple_emb_host[b]); | |
| } | |
| if (impl_->copy) cudaStreamDestroy(impl_->copy); | |
| if (impl_->grp_host) cudaFreeHost(impl_->grp_host); | |
| for (void* p : impl_->owned) cudaFree(p); | |
| } | |
| namespace { | |
| constexpr int64_t GEMM_SCRATCH = 32ll << 20; // FP16 elements for the largest dequantized dense weight | |
| constexpr size_t GEMM_WS = 32u << 20; // cuBLAS workspace | |
| // THE ATTENTION HALF AND THE MoE HALF SHARE THEIR BUFFERS. A layer runs its attention (GDN or QSA), writes it back | |
| // into the residual, and only then its MoE, so the three sets of scratch are never live at once: one region the size | |
| // of the largest holds them all. That is ~260 KB of the ~680 KB a prompt token cost - which is what lets a chunk | |
| // grow (every expert is streamed once per chunk, so a bigger chunk streams fewer bytes per token). The sizes are | |
| // counted with the same `take` sequence `init` uses; a mismatch makes `init` fail with "do not fit", never overlap. | |
| uint64_t gdn_set_bytes(size_t T) { | |
| Alloc a; a.count_only = true; bool ok = true; | |
| a.take<float>(T * C, ok); a.take<float>(T * ZV, ok); a.take<float>(T * 2 * HV, ok); a.take<float>(T * HV, ok); | |
| a.take<float>(T * HV, ok); a.take<float>(T * C, ok); a.take<float>(T * ZV, ok); a.take<uint16_t>(T * ZV, ok); | |
| return a.used; | |
| } | |
| uint64_t qsa_set_bytes(size_t T, int64_t cap, int64_t max_blocks, int64_t sel_batch, int64_t attn_batch, | |
| const strata::kernels::QsaShapes& s) { | |
| Alloc a; a.count_only = true; bool ok = true; | |
| a.take<float>(T * 512, ok); a.take<float>(T * 512, ok); a.take<float>(T * 12288, ok); a.take<float>(T * ZV, ok); | |
| a.take<float>(T * 128, ok); a.take<float>(T * 512, ok); a.take<float>(T * ZV, ok); a.take<uint16_t>(T * ZV, ok); | |
| a.take<int32_t>(T * (size_t) cap, ok); | |
| a.take<float>((size_t) sel_batch * (size_t) max_blocks, ok); | |
| a.take<float>((size_t) attn_batch * strata::kernels::qsa_decode_attn_scratch_floats(cap, s), ok); | |
| return a.used; | |
| } | |
| // Step 2b: which layers' experts go through MMQ (both weight types covered; the Strata Q2_0 pack always - its blob | |
| // is converted to GGUF Q2_0 blocks on the gather), whether any layer keeps the FP16 path (IQ1_M), and the largest | |
| // gate/up and down matrices a group buffer slot holds. STRATA_PREFILL_MMQ=0: the FP16 path everywhere (the A/B). | |
| constexpr int MMQ_GROUP = 16; // experts per MMQ launch (the gather is per expert, as blobs arrive) | |
| // MMQ reads up to one 256-value tile past a matrix's last row when the row length is not a multiple of it (the down | |
| // product: 640 values). Those bytes meet zero activations, which is harmless only if they decode to finite numbers - | |
| // llama.cpp zero-pads after every tensor, and so does a group buffer: this many zeroed bytes follow its last expert. | |
| constexpr size_t MMQ_TAIL = 4096; | |
| struct MmqPlan { | |
| bool any = false, fallback = true; | |
| std::vector<char> layer; // per layer: MMQ | |
| size_t gu_max = 0, d_max = 0; | |
| }; | |
| const MmqPlan& mmq_plan() { | |
| static const MmqPlan plan = [] { | |
| MmqPlan p; | |
| const auto& lay = strata::kernels::cpu::expert_layout(); | |
| const char* env = std::getenv("STRATA_PREFILL_MMQ"); | |
| const bool on = mmq::built() && (env == nullptr || std::atoi(env) != 0); | |
| const int64_t layers = lay.native ? (int64_t) lay.fmt.size() : lay.n_layers; | |
| p.layer.assign((size_t) std::max<int64_t>(layers, 0), 0); | |
| p.fallback = !on || layers <= 0; | |
| for (int64_t l = 0; on && l < layers; ++l) { | |
| const int gt = lay.native ? lay.fmt[(size_t) l].gu_type : 42, dt = lay.native ? lay.fmt[(size_t) l].d_type : 42; | |
| // #420: a tile on every GPU for these shapes (gate+up: 1280 rows, down: N rows), else the FP16 path | |
| if (!mmq::fits(gt, 1280) || !mmq::fits(dt, N)) { p.fallback = true; continue; } | |
| p.layer[(size_t) l] = 1; | |
| p.any = true; | |
| p.gu_max = std::max(p.gu_max, mmq::matrix_bytes(gt, 1280, N)); | |
| p.d_max = std::max(p.d_max, mmq::matrix_bytes(dt, N, 640)); | |
| } | |
| return p; | |
| }(); | |
| return plan; | |
| } | |
| // #136 P3: a layout whose chunks run the fused experts (STRATA_PF_FUSED=1, the Q2_0 pack, a streamed chunk of | |
| // stream_all_min() tokens or more). Its GU, H and Xq hold only the fused path's grouping tables, int8 H and per-token | |
| // int8 activations, and Hq nothing: ~100 KB a token less than MMQ's FP32 GU / H and per-slot q8_1 rows, which is what | |
| // lets a bigger chunk or ring fit in the slots the prompt path borrows. A last chunk below stream_all_min() still runs | |
| // MMQ in the same buffers, so each keeps MMQ's size for stream_all_min() - 1 tokens. `src`: the layout streams experts | |
| // (Prefill::init got an ExpertSource; without one no chunk takes the streamed walk, so no chunk is fused). | |
| bool fused_layout(size_t T, bool src) { | |
| return src && fused_ring() && mmq_plan().any && ring_slots(T) > STAGE && (int64_t) T >= stream_all_min(); | |
| } | |
| // The MoE buffers MMQ and the fused path share: GU and H in floats, Xq and Hq in bytes. Without `fused` (the | |
| // default): MMQ's, for T tokens. | |
| struct MoeBufs { size_t gu, h, xq, hq; }; | |
| MoeBufs moe_bufs(size_t T, int64_t n_expert, bool fused) { | |
| if (!fused) return {T * K * 1280, T * K * 640, mmq::q8_bytes((int64_t) (T * K), N), mmq::q8_bytes((int64_t) (T * K), 640)}; | |
| const size_t ts = (size_t) std::min<int64_t>((int64_t) T, stream_all_min() - 1); // MMQ's last small chunk | |
| return {std::max(ts * K * 1280, (fused::group_bytes((int64_t) (T * K), (int) n_expert) + 3) / 4), | |
| std::max(ts * K * 640, (fused::act_bytes((int64_t) (T * K), 640) + 3) / 4), | |
| std::max(mmq::q8_bytes((int64_t) (ts * K), N), fused::act_bytes((int64_t) T, N)), | |
| mmq::q8_bytes((int64_t) (ts * K), 640)}; | |
| } | |
| uint64_t moe_set_bytes(size_t T, int64_t n_expert, bool fused) { | |
| const MmqPlan& mp = mmq_plan(); | |
| const MoeBufs mb = moe_bufs(T, n_expert, fused); | |
| Alloc a; a.count_only = true; bool ok = true; | |
| a.take<float>(T * n_expert, ok); a.take<float>(T * K, ok); a.take<int32_t>(T * K, ok); a.take<int32_t>(T * K, ok); | |
| a.take<int32_t>(T * K, ok); | |
| if (mp.fallback) a.take<uint16_t>(T * K * N, ok); | |
| a.take<float>(mb.gu, ok); | |
| if (mp.fallback) a.take<uint16_t>(T * K * 640, ok); | |
| a.take<float>(T * K * N, ok); a.take<float>(T * 640, ok); | |
| a.take<float>(T * 640, ok); a.take<uint16_t>(T * 640, ok); a.take<float>(T * N, ok); a.take<float>(T, ok); | |
| if (mp.any) { | |
| a.take<uint8_t>(mb.xq, ok); | |
| a.take<float>(mb.h, ok); | |
| a.take<uint8_t>(mb.hq, ok); | |
| } | |
| return a.used; | |
| } | |
| } | |
| bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, core::SessionState& ss, | |
| core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t chunk, | |
| void* stream, std::string& err, void* borrow, uint64_t borrow_bytes) { | |
| Impl& m = *impl_; | |
| m.wt = &wt; m.g = &g; m.ss = &ss; m.src = src; m.cache = cache; m.host_res = host_res; | |
| m.T = chunk; m.cs = (cudaStream_t) stream; m.stats = &stats_; | |
| if (g.n_embd != N || g.hc != HC || g.hc_lr != LR || g.n_expert < 1 || ss.k != K) { | |
| err = "prefill: geometry differs from the artifact's"; return false; | |
| } | |
| cudaGetDevice(&m.device); | |
| if (stage_le_ < 0) stage_le_ = g.n_layers; | |
| if (stage_lb_ < 0 || stage_lb_ >= stage_le_ || stage_le_ > g.n_layers || (stage_le_ < g.n_layers) != (next_ != nullptr)) { | |
| err = "prefill: the stage's layer range is wrong"; | |
| return false; | |
| } | |
| for (int b = 0; next_ != nullptr && b < 2; ++b) | |
| if (!m.hand[b] && cudaHostAlloc((void**) &m.hand[b], (size_t) chunk * D * 4, cudaHostAllocPortable) != cudaSuccess) { | |
| err = "prefill: the layer split's hand-off buffers"; | |
| return false; | |
| } | |
| if (m.tok_dev == nullptr) { | |
| if (const cudaError_t e = cudaMalloc((void**) &m.tok_dev, (size_t) chunk * sizeof(int32_t)); e != cudaSuccess) { | |
| err = std::string("prefill: the token id buffer (") + cudaGetErrorString(e) + ")"; | |
| return false; | |
| } | |
| m.owned.push_back(m.tok_dev); | |
| m.tok_host.resize((size_t) chunk); | |
| } | |
| if (cudaStreamCreateWithFlags(&m.copy, cudaStreamNonBlocking) != cudaSuccess) { err = "prefill: copy stream"; return false; } | |
| const size_t T = (size_t) chunk; | |
| m.T_max = chunk; | |
| m.borrowed = borrow != nullptr; | |
| bool ok = true; | |
| // one-time: events, the stager, the host buffers (for the largest chunk), the identity page table | |
| for (int i = 0; i < ring_cap(); ++i) { | |
| if (cudaEventCreateWithFlags(&m.copied[i], cudaEventDisableTiming) != cudaSuccess) ok = false; | |
| if (cudaEventCreateWithFlags(&m.used[i], cudaEventDisableTiming) != cudaSuccess) ok = false; | |
| } | |
| if (!m.stager) { | |
| m.stager = std::make_unique<Stager>(); | |
| const int hw = (int) std::thread::hardware_concurrency(); | |
| const char* stv = std::getenv("STRATA_STAGER_THREADS"); // D-5: the host copy threads of unpinned blobs | |
| // A GGUF read in place (UD-Q4_K_XL beyond its RAM budget): most of a chunk's blobs are page faults on the | |
| // SSD, so the copies need many reads in flight - 32 threads and a 128-deep ring read a 4K chunk in 29 s | |
| // instead of 71 s on an RTX 5070 / NVMe PC (4 threads, 16 deep: the defaults, kept for every other source) | |
| bool files = false; | |
| for (int64_t l = 0; src != nullptr && !files && l < g.n_layers; ++l) | |
| for (int64_t e = 0; !files && e < g.n_expert; ++e) files = src->transient(l, e); | |
| const int threads = stv ? std::clamp(std::atoi(stv), 1, 32) : files ? 32 : std::max(2, std::min(4, hw / 4)); | |
| if (files && std::getenv("STRATA_STAGER_RING") == nullptr) m.stager->kRing = 4 * threads; | |
| if (!m.stager->init((size_t) MAXBLOB(), threads)) ok = false; | |
| } | |
| m.steps_host.resize(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)); | |
| 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); | |
| m.grp_host = m.grp_dev = nullptr; | |
| m.grp_n = m.grp_tk = 0; | |
| void *h = nullptr, *d = nullptr; | |
| if (cudaHostAlloc(&h, need * 4, cudaHostAllocMapped | (core::peer_portable() ? cudaHostAllocPortable : 0)) == cudaSuccess && | |
| cudaHostGetDevicePointer(&d, h, 0) == cudaSuccess) { | |
| m.grp_host = (int32_t*) h; | |
| m.grp_dev = (int32_t*) d; | |
| m.grp_n = need; | |
| m.grp_tk = T * K; | |
| } else { // the copies, as before | |
| if (h) cudaFreeHost(h); | |
| cudaGetLastError(); | |
| } | |
| } | |
| } | |
| for (int b = 0; b < 2; ++b) { | |
| if (!m.ple_emb_host[b] && | |
| cudaHostAlloc((void**) &m.ple_emb_host[b], (size_t) T * N * 4, cudaHostAllocDefault) != cudaSuccess) { | |
| cudaGetLastError(); | |
| m.ple_pageable[b].resize(T * N); // pageable: the upload is staged before it returns | |
| m.ple_emb_host[b] = m.ple_pageable[b].data(); | |
| } | |
| if (!m.ple_copied[b] && cudaEventCreateWithFlags(&m.ple_copied[b], cudaEventDisableTiming) != cudaSuccess) | |
| ok = false; | |
| m.ple_rows[b].resize(T * strata::kernels::PLE_N_HEADS); | |
| } | |
| if (ss.qsa_states[ss.qsa_primary()].kv_mode == 1) { // KV streaming: the staging pool's identity page table | |
| const int64_t pages = ss.qsa_states[ss.qsa_primary()].n_pages; | |
| std::vector<int32_t> ident((size_t) pages); | |
| for (int64_t i = 0; i < pages; ++i) ident[(size_t) i] = (int32_t) i; | |
| if (cudaMalloc((void**) &m.ident_table, ident.size() * 4) != cudaSuccess || | |
| cudaMemcpy(m.ident_table, ident.data(), ident.size() * 4, cudaMemcpyHostToDevice) != cudaSuccess) | |
| ok = false; | |
| else | |
| m.owned.push_back(m.ident_table); | |
| } | |
| if (!ok) { err = "prefill: host buffers or events for a chunk of " + std::to_string(chunk) + " tokens"; return false; } | |
| Alloc o; | |
| o.base = (uint8_t*) borrow; | |
| o.cap = borrow_bytes; | |
| o.owned = &m.owned; | |
| { | |
| uint16_t* gs = o.take<uint16_t>((size_t) GEMM_SCRATCH, ok); | |
| void* ws = o.take<uint8_t>(GEMM_WS, ok); | |
| if (!ok) { err = "prefill: GEMM scratch does not fit"; return false; } | |
| if (!m.gemm.init_external(stream, gs, GEMM_SCRATCH, ws, GEMM_WS, err)) return false; | |
| } | |
| if (!carve(T, &o)) { | |
| err = "prefill: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; | |
| return false; | |
| } | |
| return true; | |
| } | |
| // Every device buffer of a chunk of T tokens, from the Alloc `alloc` (after the GEMM scratch and workspace): `init` | |
| // once, and `relayout` for a request's own chunk. The order is `bytes_needed`'s. | |
| bool Prefill::carve(size_t T, void* alloc) { | |
| Impl& m = *impl_; | |
| Alloc& o = *static_cast<Alloc*>(alloc); | |
| const core::ModelGeometry& g = *m.g; | |
| core::SessionState& ss = *m.ss; | |
| bool ok = true; | |
| m.emb = o.take<float>(T * N, ok); m.R = o.take<float>(T * D, ok); | |
| m.xn = gr_unfused() ? o.take<float>(T * D, ok) : nullptr; // F-1: not needed (gr_mix_r reads R) | |
| m.grs = o.take<float>(T * HC, ok); | |
| m.xn16 = o.take<uint16_t>(T * D, ok); m.lo = o.take<float>(T * LR, ok); m.lo16 = o.take<uint16_t>(T * LR, ok); | |
| m.gated = o.take<float>(T * D, ok); m.inj = o.take<float>(T * HC, ok); | |
| m.mixed = o.take<float>(T * N, ok); m.mixed_bf = o.take<uint16_t>(T * N, ok); | |
| m.mixed_h = o.take<uint16_t>(T * N, ok); m.bo = o.take<float>(T * N, ok); | |
| if (bf16x2_hc()) { m.xn16_lo = o.take<uint16_t>(T * D, ok); m.lo16_lo = o.take<uint16_t>(T * LR, ok); } | |
| if (bf16x2()) m.mixed_bf_lo = o.take<uint16_t>(T * N, ok); | |
| m.steps_dev = o.take<int32_t>(T * strata::kernels::kStepCount, ok); | |
| strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); | |
| s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; | |
| s.idx_dim = g.idx_key_dim; | |
| m.cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); | |
| m.max_blocks = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2; | |
| { | |
| // one region for the attention half's and the MoE half's scratch (see gdn_set_bytes) | |
| const bool fz = fused_layout(T, m.src != nullptr); | |
| const MoeBufs mb = moe_bufs(T, m.g->n_expert, fz); | |
| const uint64_t region = std::max({gdn_set_bytes(T), qsa_set_bytes(T, m.cap, m.max_blocks, m.sel_batch, | |
| m.attn_batch, s), moe_set_bytes(T, m.g->n_expert, fz)}); | |
| uint8_t* base = o.take<uint8_t>((size_t) region, ok); | |
| m.region = base; | |
| m.region_bytes = region; | |
| Alloc a; | |
| a.base = base; a.cap = region; a.owned = &m.owned; | |
| m.qkv = a.take<float>(T * C, ok); m.z = a.take<float>(T * ZV, ok); m.ab = a.take<float>(T * 2 * HV, ok); | |
| m.gate = a.take<float>(T * HV, ok); m.beta = a.take<float>(T * HV, ok); m.hbuf = a.take<float>(T * C, ok); | |
| m.y = a.take<float>(T * ZV, ok); m.y_h = a.take<uint16_t>(T * ZV, ok); | |
| Alloc b; | |
| b.base = base; b.cap = region; b.owned = &m.owned; | |
| m.Kc = b.take<float>(T * 512, ok); m.Vc = b.take<float>(T * 512, ok); m.Qf = b.take<float>(T * 12288, ok); | |
| m.q = b.take<float>(T * ZV, ok); m.idx_raw = b.take<float>(T * 128, ok); m.q_idx = b.take<float>(T * 512, ok); | |
| m.attn = b.take<float>(T * ZV, ok); m.attn_h = b.take<uint16_t>(T * ZV, ok); | |
| m.sel_ids = b.take<int32_t>(T * (size_t) m.cap, ok); | |
| m.sel_scores = b.take<float>((size_t) m.sel_batch * (size_t) m.max_blocks, ok); | |
| m.attn_scratch = b.take<float>((size_t) m.attn_batch * strata::kernels::qsa_decode_attn_scratch_floats(m.cap, s), ok); | |
| Alloc c; | |
| c.base = base; c.cap = region; c.owned = &m.owned; | |
| m.logits = c.take<float>(T * m.g->n_expert, ok); m.w = c.take<float>(T * K, ok); m.ids = c.take<int32_t>(T * K, ok); | |
| m.slot_dev = c.take<int32_t>(T * K, ok); m.src_dev = c.take<int32_t>(T * K, ok); | |
| const MmqPlan& mp = mmq_plan(); | |
| m.Xs = mp.fallback ? c.take<uint16_t>(T * K * N, ok) : nullptr; | |
| m.GU = c.take<float>(mb.gu, ok); | |
| m.Hh = mp.fallback ? c.take<uint16_t>(T * K * 640, ok) : nullptr; | |
| m.Dm = c.take<float>(T * K * N, ok); | |
| m.sgate = c.take<float>(T * 640, ok); m.sup = c.take<float>(T * 640, ok); m.sh_h = c.take<uint16_t>(T * 640, ok); | |
| m.shared = c.take<float>(T * N, ok); m.sg = c.take<float>(T, ok); | |
| if (mp.any) { | |
| m.Xq = c.take<uint8_t>(mb.xq, ok); | |
| m.H = c.take<float>(mb.h, ok); | |
| m.Hq = c.take<uint8_t>(mb.hq, ok); | |
| } | |
| if (base == nullptr) ok = false; | |
| } | |
| for (int i = 0; i < DQ; ++i) { m.dq_gu[i] = o.take<uint16_t>(1280 * 2560, ok); m.dq_d[i] = o.take<uint16_t>(2560 * 640, ok); } | |
| if (mmq_plan().any) { | |
| const MmqPlan& mp = mmq_plan(); | |
| m.ids_identity = o.take<int32_t>(T * K, ok); | |
| m.bounds_dev = o.take<int32_t>((size_t) (2 * (m.g->n_expert + m.g->n_expert / MMQ_GROUP + 2)), ok); | |
| m.grp_gu = o.take<uint8_t>(MMQ_GROUP * mp.gu_max + MMQ_TAIL, ok); | |
| m.grp_d = o.take<uint8_t>(MMQ_GROUP * mp.d_max + MMQ_TAIL, 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<mmq::Context>(); | |
| } | |
| m.ring = ring_slots(T); | |
| 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); | |
| take_stage(o, ss, s, m.stage, ok); | |
| m.T = (int64_t) T; | |
| return ok; | |
| } | |
| bool Prefill::relayout(int64_t chunk, void* borrow, uint64_t borrow_bytes, std::string& err) { | |
| Impl& m = *impl_; | |
| if (!m.borrowed || borrow == nullptr || chunk <= 0 || chunk > m.T_max) { | |
| err = "prefill: relayout needs borrowed buffers and a chunk of at most " + std::to_string(m.T_max); | |
| return false; | |
| } | |
| if (cudaStreamSynchronize(m.cs) != cudaSuccess || cudaStreamSynchronize(m.copy) != cudaSuccess) { | |
| err = "prefill: relayout: the stream failed"; | |
| return false; | |
| } | |
| bool ok = true; | |
| Alloc o; | |
| o.base = (uint8_t*) borrow; | |
| o.cap = borrow_bytes; | |
| o.owned = &m.owned; | |
| uint16_t* gs = o.take<uint16_t>((size_t) GEMM_SCRATCH, ok); | |
| void* ws = o.take<uint8_t>(GEMM_WS, ok); | |
| if (ok) m.gemm.rebind(gs, GEMM_SCRATCH, ws, GEMM_WS); | |
| if (!ok || !carve((size_t) chunk, &o)) { | |
| err = "prefill: device buffers for a chunk of " + std::to_string(chunk) + " tokens do not fit"; | |
| return false; | |
| } | |
| return true; | |
| } | |
| int64_t Prefill::chunk() const { return impl_->T; } | |
| bool Prefill::draft_kv(core::MtpDrafter& mtp, const float* R_rows, const int32_t* next_tokens, int64_t n, int64_t cell0, | |
| std::string& err) { | |
| Impl& m = *impl_; | |
| static const bool off = [] { const char* v = std::getenv("STRATA_MTP_BATCH"); return v != nullptr && v[0] == '0'; }(); | |
| // A ring (KV streaming: the drafter's window, page p in slot p % n_slots over a host copy) takes the same appends | |
| // with its own page table and host copy, as a streamed main layer does; the cells written are those the window can | |
| // still reach (r0 below), which the ring holds, so no two of them share a slot. STRATA_MTP_BATCH_RING=0: the | |
| // drafter's own pass for a ring (the A/B). | |
| static const bool ring_ok = [] { const char* v = std::getenv("STRATA_MTP_BATCH_RING"); return v == nullptr || v[0] != '0'; }(); | |
| core::QsaState& st = mtp.kv_state_rw(); | |
| if (off || n <= 0 || m.g == nullptr || m.region == nullptr || (st.kv_mode != 0 && !(st.kv_mode == 2 && ring_ok)) || | |
| st.kv_hybrid || mtp.device() != m.device) | |
| return false; | |
| const auto t0 = Clock::now(); | |
| const core::ModelGeometry& g = *m.g; | |
| const int64_t Nn = g.n_embd, HCN = g.hc * g.n_embd, KV = g.n_head_kv * g.head_dim; | |
| constexpr int kQ8_0 = 8; // GGML_TYPE_Q8_0 | |
| const float* w_ne = mtp.tensor_f32("pre_fc_norm_embedding.weight"); | |
| const void* w_fe = mtp.tensor_q8("fc_embedding.weight"); | |
| const float* w_nh = mtp.tensor_f32("pre_fc_norm_hidden.weight"); | |
| const void* w_fh = mtp.tensor_q8("fc_hidden.weight"); | |
| const float* w_hn = mtp.tensor_f32("attn_hyper_connection.hc_norm.weight"); | |
| const uint16_t* w_dn = mtp.tensor_bf16("attn_hyper_connection.input_mix_weight_down.weight"); | |
| const uint16_t* w_up = mtp.tensor_bf16("attn_hyper_connection.input_mix_weight_up.weight"); | |
| const void* w_k = mtp.tensor_q8("self_attn.k_proj.weight"); | |
| const void* w_v = mtp.tensor_q8("self_attn.v_proj.weight"); | |
| const float* w_kn = mtp.tensor_f32("self_attn.k_norm.weight"); | |
| if (!w_ne || !w_fe || !w_nh || !w_fh || !w_hn || !w_dn || !w_up || !w_k || !w_v || !w_kn) return false; | |
| const core::NativeEmbed* nemb = core::native_embed(); | |
| const core::WeightRef* wemb = nemb ? nullptr : m.wt->find("token_embd.weight"); | |
| if (!nemb && (wemb == nullptr || wemb->codebook_iq4nl || wemb->ne0 != g.n_embd || wemb->group_elems <= 0 || | |
| (wemb->code_bits != 2 && wemb->code_bits != 4 && wemb->code_bits != 8))) | |
| return false; | |
| // the cells the drafter's window can still reach | |
| const int64_t r0 = std::max<int64_t>(0, mtp.first_needed() - cell0); | |
| if (r0 >= n) return true; | |
| // per row: emb/e2 (N), en16 (N half), hn/h2/Rm/gated (HCN), hn16/xn16 (HCN half), lo (LR) + lo16, grs, mixed (N) + | |
| // mixed_h, K and V (KV each), the token id | |
| // E-9: the drafter's Q8_0 matrices through Q8_1 x Q8_0 MMQ - its own pass's integer dot products (mmvq), so | |
| // its K/V stay close to what the drafter computes itself; STRATA_MTP_BATCH_F16=1: FP16 GEMMs (the A/B) | |
| static const bool f16_only = [] { const char* v = std::getenv("STRATA_MTP_BATCH_F16"); return v && v[0] == '1'; }(); | |
| const bool q8 = !f16_only && mmq::built() && mmq::fits(kQ8_0, Nn) && mmq::fits(kQ8_0, KV); // #420 | |
| const uint64_t per_row = 4 * (2 * Nn + 4 * HCN + LR + HC + Nn + 2 * KV + 1) + 2 * (Nn + 2 * HCN + LR + Nn) + 64 + | |
| (q8 ? (uint64_t) mmq::q8_bytes(g.hc, Nn) + 4 * g.hc : 0); | |
| int64_t B = std::min<int64_t>(n - r0, (int64_t) (m.region_bytes / per_row) & ~(int64_t) 63); | |
| if (st.kv_mode == 2) // a ring: one batch's cells must not share a slot (a batch can straddle one page more) | |
| B = std::min<int64_t>(B, ((st.n_slots - 1) * strata::kernels::qsa_real_shapes().page_size) & ~(int64_t) 63); | |
| if (B < 64) return false; | |
| uint8_t* q = m.region; | |
| auto carve = [&](size_t bytes) { void* p = q; q += (bytes + 255) & ~(size_t) 255; return p; }; | |
| float* emb = (float*) carve((size_t) B * Nn * 4); | |
| float* e2 = (float*) carve((size_t) B * Nn * 4); | |
| uint16_t* en16 = (uint16_t*) carve((size_t) B * Nn * 2); | |
| float* hn = (float*) carve((size_t) B * HCN * 4); | |
| float* h2 = (float*) carve((size_t) B * HCN * 4); | |
| float* Rm = (float*) carve((size_t) B * HCN * 4); | |
| float* gated = (float*) carve((size_t) B * HCN * 4); | |
| uint16_t* hn16 = (uint16_t*) carve((size_t) B * HCN * 2); | |
| uint16_t* xn16 = (uint16_t*) carve((size_t) B * HCN * 2); | |
| float* lo = (float*) carve((size_t) B * LR * 4); | |
| uint16_t* lo16 = (uint16_t*) carve((size_t) B * LR * 2); | |
| float* grs = (float*) carve((size_t) B * HC * 4); | |
| float* mixed = (float*) carve((size_t) B * Nn * 4); | |
| uint16_t* mixed_h = (uint16_t*) carve((size_t) B * Nn * 2); | |
| float* Kc = (float*) carve((size_t) B * KV * 4); | |
| float* Vc = (float*) carve((size_t) B * KV * 4); | |
| int32_t* tok = (int32_t*) carve((size_t) B * 4); | |
| void* xq = q8 ? carve(mmq::q8_bytes(B * g.hc, Nn)) : nullptr; | |
| int32_t* ident = q8 ? (int32_t*) carve((size_t) B * g.hc * 4) : nullptr; | |
| int32_t* bnd = q8 ? (int32_t*) carve(16) : nullptr; | |
| if ((uint64_t) (q - m.region) > m.region_bytes) return false; | |
| static const bool timing = std::getenv("STRATA_DRAFT_TIMING") != nullptr; // debug: where this pass's time goes | |
| if (timing) cudaStreamSynchronize(m.cs); | |
| const auto ti0 = Clock::now(); | |
| if (!mtp.idle(err)) return false; // the drafter's own stream (its graph uploads) before this writes its K/V | |
| const double ms_idle = ms_since(ti0); | |
| const auto tl0 = Clock::now(); | |
| strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); | |
| s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; | |
| s.idx_dim = g.idx_key_dim; | |
| std::vector<int32_t> tk((size_t) B); | |
| if (q8) { | |
| if (!m.mmq_ctx) m.mmq_ctx = std::make_unique<mmq::Context>(); | |
| mmq::iota(ident, B * g.hc, m.cs); | |
| } | |
| // y[rows, n_out] = x[rows, k] . w^T for a Q8_0 matrix: MMQ from the FP32 rows (bounds slot 0: rows, 1: rows*hc) | |
| auto proj = [&](const float* x, const uint16_t* x16, const void* w, float* y, int64_t rows, int64_t n_out, | |
| int64_t k, int slot) { | |
| if (!q8) { m.gemm.native(x16, kQ8_0, w, y, rows, n_out, k); return; } | |
| mmq::quantize(x, nullptr, xq, kQ8_0, k, k, rows, m.cs); | |
| mmq::Product p; | |
| p.w = w; p.type = kQ8_0; p.w_rows = n_out; p.w_cols = k; p.expert_bytes = mmq::matrix_bytes(kQ8_0, n_out, k); | |
| p.n = 1; p.xq = xq; p.bounds = bnd + 2 * slot; p.ids = ident; p.total_rows = rows; p.max_rows = rows; | |
| p.dst = y; p.ld_dst = n_out; | |
| m.mmq_ctx->run(p, m.cs); | |
| }; | |
| for (int64_t b0 = r0; b0 < n; b0 += B) { | |
| const int64_t nb = std::min(B, n - b0), c0 = cell0 + b0; | |
| for (int64_t i = 0; i < nb; ++i) tk[(size_t) i] = next_tokens[b0 + i]; | |
| const int32_t bh[4] = {0, (int32_t) nb, 0, (int32_t) (nb * g.hc)}; | |
| if (q8 && cudaMemcpyAsync(bnd, bh, sizeof bh, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { | |
| err = "prefill: the draft bounds' upload failed"; | |
| return false; | |
| } | |
| if (cudaMemcpyAsync(tok, tk.data(), (size_t) nb * 4, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { | |
| err = "prefill: the draft tokens' upload failed"; | |
| return false; | |
| } | |
| // the input branches: the next token's embedding, and this cell's final residual rows | |
| if (nemb) { | |
| nemb->gather_dev(tok, nb, emb, m.cs); | |
| } else { | |
| const auto* codes = (const uint8_t*) wemb->data; | |
| const auto* scales = (const float*) (codes + wemb->codes_bytes); | |
| const auto* offsets = wemb->has_offset ? (const float*) (codes + wemb->codes_bytes + wemb->scales_bytes) | |
| : nullptr; | |
| strata::kernels::embedding_gather_dev(codes, scales, offsets, tok, (int) nb, wemb->ne0, wemb->code_bits, | |
| wemb->code_bias, wemb->group_elems, | |
| (uint64_t) (wemb->ne0 / (8 / wemb->code_bits)), | |
| (uint64_t) (wemb->ne0 / wemb->group_elems), emb, m.cs); | |
| } | |
| rms_rows(emb, w_ne, nb, Nn, Nn, EPS, m.cs); | |
| if (!q8) to_f16(emb, en16, nb * Nn, m.cs); | |
| proj(emb, en16, w_fe, e2, nb, Nn, Nn, 0); | |
| cudaMemcpyAsync(hn, R_rows + (size_t) b0 * HCN, (size_t) nb * HCN * 4, cudaMemcpyDeviceToDevice, m.cs); | |
| rms_rows(hn, w_nh, nb, HCN, HCN, EPS, m.cs); | |
| if (!q8) to_f16(hn, hn16, nb * HCN, m.cs); | |
| proj(hn, hn16, w_fh, h2, nb * g.hc, Nn, Nn, 1); // every stream through fc_hidden | |
| strata::kernels::add_streams_broadcast(h2, e2, Rm, Nn, (int) g.hc, (int) nb, m.cs); | |
| // the attention hyper-connection's read (its mixed input only: this pass writes nothing back) | |
| gr_norm_rs(Rm, w_hn, EPS, grs, xn16, nb, m.cs); | |
| m.gemm.bf16(xn16, w_dn, lo, nb, LR, HCN); | |
| gr_silu(lo, lo16, nb, m.cs); | |
| m.gemm.bf16(lo16, w_up, gated, nb, HCN, LR); | |
| gr_mix_r(Rm, grs, w_hn, gated, mixed, nullptr, nb, m.cs, mixed_h); | |
| // K and V into the drafter's cache, as the prompt path's QSA layers append theirs | |
| proj(mixed, mixed_h, w_k, Kc, nb, KV, Nn, 0); | |
| proj(mixed, mixed_h, w_v, Vc, nb, KV, Nn, 0); | |
| rms_rows(Kc, w_kn, nb * g.n_head_kv, g.head_dim, g.head_dim, EPS, m.cs); | |
| rope(Kc, nb, g.n_head_kv, g.head_dim, KV, c0, strata::kernels::rope_scaling(), m.cs); | |
| if (st.kv_rot) { // rotated as the drafter's own decode stores them (mtp.cpp) | |
| strata::kernels::fwht256_inplace_cuda(Kc, nb * g.n_head_kv, m.cs); | |
| strata::kernels::fwht256_inplace_cuda(Vc, nb * g.n_head_kv, m.cs); | |
| } | |
| if (st.kv_q4) { | |
| strata::kernels::kv_append_q4(st.k_q4, st.v_q4, st.page_table, c0, nb, Kc, Vc, s, m.cs, &st.host); | |
| } else { | |
| kv_append(Kc, Vc, nb, c0, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, | |
| st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs, &st.host); | |
| } | |
| } | |
| if (cudaStreamSynchronize(m.cs) != cudaSuccess) { | |
| err = std::string("prefill: the draft layer's K/V: ") + cudaGetErrorString(cudaGetLastError()); | |
| return false; | |
| } | |
| mtp.ms_prefill += ms_since(t0); | |
| if (timing) | |
| std::fprintf(stderr, "strata draft kv: %lld cells from %lld (first needed %lld), batch %lld: drafter idle %.1f ms, " | |
| "the batches %.1f ms, all %.1f ms\n", (long long) n, (long long) cell0, (long long) (cell0 + r0), | |
| (long long) B, ms_idle, ms_since(tl0), ms_since(t0)); | |
| return true; | |
| } | |
| bool Prefill::set_peer(core::PeerExperts* peer, int64_t cap_rows, std::string& err) { | |
| Impl& m = *impl_; | |
| if (peer == nullptr || !peer->valid()) { m.pp.reset(); return true; } | |
| if (!peer->p2p()) { err = "prefill peer: the two GPUs cannot access each other (no P2P)"; return false; } | |
| if (!mmq_plan().any) { err = "prefill peer: needs the MMQ prompt path"; return false; } | |
| auto pp = std::make_unique<PeerPrefill>(); | |
| pp->peer = peer; | |
| pp->T_max = m.T_max; | |
| pp->cap_rows = std::max<int64_t>(1, std::min<int64_t>(cap_rows, m.T_max * K)); | |
| const int64_t R = pp->cap_rows; | |
| const MmqPlan& mp = mmq_plan(); | |
| if (cudaEventCreateWithFlags(&pp->ev_in, cudaEventDisableTiming) != cudaSuccess) { err = "prefill peer: event"; return false; } | |
| int prev = 0; | |
| cudaGetDevice(&prev); | |
| pp->dev = peer->device(); | |
| cudaSetDevice(pp->dev); | |
| bool ok = cudaStreamCreateWithFlags(&pp->s, cudaStreamNonBlocking) == cudaSuccess && | |
| cudaEventCreateWithFlags(&pp->ev_done, cudaEventDisableTiming) == cudaSuccess; | |
| { | |
| const char* v = std::getenv("STRATA_PF_PEER_OUT_PIPE"); // =0: one copy after the last group (the A/B) | |
| pp->out_pipe = (v == nullptr || std::atoi(v) != 0) || pp->compact; // compact always pipes | |
| } | |
| if (ok && pp->out_pipe) { | |
| ok = cudaStreamCreateWithFlags(&pp->s_out, cudaStreamNonBlocking) == cudaSuccess; | |
| for (int i = 0; ok && i < PeerPrefill::kGrpEv; ++i) | |
| ok = cudaEventCreateWithFlags(&pp->ev_grp[i], cudaEventDisableTiming) == cudaSuccess; | |
| } | |
| auto take = [&](size_t bytes) -> void* { | |
| void* p = nullptr; | |
| if (!ok || cudaMalloc(&p, bytes) != cudaSuccess) { ok = false; return nullptr; } | |
| pp->owned.push_back(p); | |
| return p; | |
| }; | |
| pp->mixed = (float*) take((size_t) m.T_max * N * 4); | |
| { | |
| const char* v = std::getenv("STRATA_PF_PEER_COMPACT"); | |
| pp->compact = v == nullptr || std::atoi(v) != 0; | |
| } | |
| if (pp->compact) { | |
| pp->cap_rows = m.T_max * K; // no row cap: only the row tables grow with it | |
| pp->G = m.T_max; // one expert never has more rows than the chunk has tokens | |
| const int64_t Gr = pp->G; | |
| pp->Xq_g = take(mmq::q8_bytes(Gr, N)); | |
| pp->Hq_g = take(mmq::q8_bytes(Gr, 640)); | |
| pp->GU_g = (float*) take((size_t) Gr * 1280 * 4); | |
| pp->H_g = (float*) take((size_t) Gr * 640 * 4); | |
| pp->Dm_b[0] = (float*) take((size_t) Gr * N * 4); | |
| pp->Dm_b[1] = (float*) take((size_t) Gr * N * 4); | |
| for (int b = 0; b < 2 && ok; ++b) ok = cudaEventCreateWithFlags(&pp->ev_dm[b], cudaEventDisableTiming) == cudaSuccess; | |
| const char* fs = std::getenv("STRATA_PF_PEER_STREAM"); | |
| pp->ps_frac = fs ? std::atof(fs) : 0.35; // measured: 0.25-0.5 all ~1950-1970 at 32K, 0.35 best | |
| if (pp->ps_frac > 0.0 && !mp.fallback && m.T_max >= stream_all_min()) { | |
| pp->RP = 48; | |
| if (const char* pr = std::getenv("STRATA_PF_PEER_RING"); pr != nullptr) pp->RP = std::atoi(pr); | |
| pp->pstage.assign((size_t) pp->RP, nullptr); | |
| pp->pcopied.assign((size_t) pp->RP, nullptr); | |
| pp->pused.assign((size_t) pp->RP, nullptr); | |
| pp->plive.assign((size_t) pp->RP, 0); | |
| ok = ok && cudaStreamCreateWithFlags(&pp->s_cp, cudaStreamNonBlocking) == cudaSuccess; | |
| for (int i = 0; ok && i < pp->RP; ++i) { | |
| pp->pstage[(size_t) i] = (uint8_t*) take((size_t) MAXBLOB()); | |
| ok = ok && cudaEventCreateWithFlags(&pp->pcopied[(size_t) i], cudaEventDisableTiming) == cudaSuccess && | |
| cudaEventCreateWithFlags(&pp->pused[(size_t) i], cudaEventDisableTiming) == cudaSuccess; | |
| } | |
| } else { | |
| pp->ps_frac = 0.0; | |
| } | |
| } else { | |
| pp->Xq = take(mmq::q8_bytes(R, N)); | |
| pp->Hq = take(mmq::q8_bytes(R, 640)); | |
| pp->GU = (float*) take((size_t) R * 1280 * 4); | |
| pp->H = (float*) take((size_t) R * 640 * 4); | |
| pp->Dm = (float*) take((size_t) R * N * 4); | |
| } | |
| const int64_t Rt = pp->cap_rows; | |
| pp->src = (int32_t*) take((size_t) Rt * 4); | |
| pp->ident = (int32_t*) take((size_t) Rt * 4); | |
| pp->bounds = (int32_t*) take((size_t) (2 * (NE + NE / MMQ_GROUP + 2)) * 4); | |
| pp->grp_gu = (uint8_t*) take(MMQ_GROUP * mp.gu_max + MMQ_TAIL); | |
| pp->grp_d = (uint8_t*) take(MMQ_GROUP * mp.d_max + MMQ_TAIL); | |
| if (ok) { | |
| pp->ctx = std::make_unique<mmq::Context>(); | |
| mmq::iota(pp->ident, pp->cap_rows, pp->s); | |
| ok = cudaStreamSynchronize(pp->s) == cudaSuccess; | |
| } | |
| size_t fb = 0, tb = 0; | |
| cudaMemGetInfo(&fb, &tb); | |
| cudaSetDevice(prev); | |
| if (!ok) { err = "prefill peer: the peer's buffers do not fit (raise --peer-reserve-mib or lower --peer-prefill-rows)"; return false; } | |
| std::fprintf(stderr, "strata prefill: peer GPU %d computes its experts' rows of each prompt chunk (up to %lld rows per " | |
| "layer%s); %zu MiB left free on it\n", pp->dev, (long long) pp->cap_rows, | |
| pp->compact ? (pp->ps_frac > 0.0 ? (", compact group buffers, streams " + std::to_string((int) (pp->ps_frac * 100 + 0.5)) + | |
| "% of the primary's streamed experts through a " + std::to_string(pp->RP) + "-slot ring").c_str() | |
| : ", compact group buffers") : "", fb >> 20); | |
| m.pp = std::move(pp); | |
| return true; | |
| } | |
| void Prefill::set_pinned_share(double share) { g_pinned_share = share; } | |
| void Prefill::set_ring_override(int slots) { g_ring_override = slots > 0 ? slots : 0; } | |
| double Prefill::pinned_share() { return g_pinned_share; } | |
| uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk) { | |
| // the same allocation sequence as `init`, counted | |
| const size_t T = (size_t) chunk; | |
| bool ok = true; | |
| Alloc o; | |
| o.count_only = true; | |
| o.take<uint16_t>((size_t) GEMM_SCRATCH, ok); | |
| o.take<uint8_t>(GEMM_WS, ok); | |
| auto f = [&](size_t n) { o.take<float>(n, ok); }; | |
| f(T * N); f(T * D); f(T * D); o.take<uint16_t>(T * D, ok); f(T * LR); o.take<uint16_t>(T * LR, ok); | |
| f(T * D); f(T * HC); f(T * N); o.take<uint16_t>(T * N, ok); o.take<uint16_t>(T * N, ok); f(T * N); | |
| if (bf16x2_hc()) { o.take<uint16_t>(T * D, ok); o.take<uint16_t>(T * LR, ok); } | |
| if (bf16x2()) o.take<uint16_t>(T * N, ok); | |
| o.take<int32_t>(T * strata::kernels::kStepCount, ok); | |
| strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); | |
| s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; | |
| s.idx_dim = g.idx_key_dim; | |
| const int64_t cap = strata::kernels::qsa_selection_width(strata::kernels::kTopkMaxCells, s); | |
| const int64_t max_blocks = ss.qsa_states[ss.qsa_primary()].max_cells / s.idx_block + 2; | |
| o.take<uint8_t>((size_t) std::max({gdn_set_bytes(T), qsa_set_bytes(T, cap, max_blocks, 256, 32, s), | |
| moe_set_bytes(T, g.n_expert, fused_layout(T, true))}), ok); | |
| for (int i = 0; i < DQ; ++i) { o.take<uint16_t>(1280 * 2560, ok); o.take<uint16_t>(2560 * 640, ok); } | |
| if (mmq_plan().any) { | |
| const MmqPlan& mp = mmq_plan(); | |
| o.take<int32_t>(T * K, ok); | |
| o.take<int32_t>((size_t) (2 * (g.n_expert + g.n_expert / MMQ_GROUP + 2)), ok); | |
| o.take<uint8_t>(MMQ_GROUP * mp.gu_max + MMQ_TAIL, ok); | |
| o.take<uint8_t>(MMQ_GROUP * mp.d_max + MMQ_TAIL, ok); | |
| } | |
| for (int i = 0; i < ring_slots(T); ++i) o.take<uint8_t>((size_t) MAXBLOB(), ok); | |
| f(T * N); | |
| f((size_t) strata::kernels::NG_HC_DIM); | |
| strata::kernels::KvHostPools stage; | |
| take_stage(o, ss, s, stage, ok); | |
| return o.used + (8u << 20); // alignment slack | |
| } | |
| namespace { | |
| const core::WeightRef* need(const core::LayerView& v, const char* suffix, std::string& err) { | |
| const core::WeightRef* r = v.get(suffix); | |
| if (!r) err = v.name(suffix) + " is missing"; | |
| return r; | |
| } | |
| bool native_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y, int64_t T, const std::string& name, | |
| std::string& err, int64_t ldy = 0) { | |
| if (!w->native_data) { err = "prefill: " + name + " has no native GGUF blocks (run with --native)"; return false; } | |
| gm.native(X, w->native_type, w->native_data, Y, T, w->ne1, w->ne0, ldy); | |
| return true; | |
| } | |
| bool bf16_proj(Gemm& gm, const core::WeightRef* w, const uint16_t* X, float* Y, int64_t T, const std::string& name, | |
| std::string& err, int64_t ldy = 0, const uint16_t* X_lo = nullptr) { | |
| if (w->kind != core::WeightKind::Bf16InF32 || !w->data) { err = "prefill: " + name + " is not a resident BF16 tensor"; return false; } | |
| gm.bf16(X, (const uint16_t*) w->data, Y, T, w->ne1 > 0 ? w->ne1 : 1, w->ne0, ldy); | |
| if (X_lo) gm.bf16(X_lo, (const uint16_t*) w->data, Y, T, w->ne1 > 0 ? w->ne1 : 1, w->ne0, ldy, 1.0f); | |
| return true; | |
| } | |
| } // namespace | |
| namespace { | |
| // STRATA_PREFILL_TIMING=1: the prompt path's GPU time by phase. Events are recorded on the compute stream in order; | |
| // the time between two consecutive marks is charged to the phase of the first, so a gap where the GPU waits (for the | |
| // host's expert grouping, or for an expert's copy) lands on the phase that was waiting. Events are reused: the marks | |
| // are folded at every MoE layer's host sync, after which all of them have completed. | |
| enum PfPhase { kPfStart, kPfHc, kPfGdn, kPfQsa, kPfQsaIdx, kPfQsaSel, kPfQsaAttn, kPfRouter, kPfHostGroup, kPfGather, | |
| kPfWaitCopy, kPfDequant, kPfGemmGU, kPfGemmD, kPfCombine, kPfPle, kPfKvStage, kPfGdnConv, kPfGdnRec, kPfGdnOut, | |
| kPfCount }; | |
| const char* const kPfNames[kPfCount] = {"embed+steps", "hc read", "gdn", "qsa proj", "qsa indexer", "qsa select", | |
| "qsa attn", "router+shared", "host grouping", "gather", "wait copy", "dequant", | |
| "gemm gate/up", "gemm down", "combine", "ple", "kv stage", "gdn conv+gates", | |
| "gdn recurrence", "gdn out proj"}; | |
| struct PfTimer { | |
| bool on = std::getenv("STRATA_PREFILL_TIMING") != nullptr; | |
| std::vector<cudaEvent_t> ev; | |
| std::vector<int> ph; | |
| size_t used = 0; | |
| double ms[kPfCount] = {}; | |
| void mark(int phase, cudaStream_t s) { | |
| if (!on) return; | |
| if (used == ev.size()) { | |
| cudaEvent_t e = nullptr; | |
| cudaEventCreate(&e); | |
| ev.push_back(e); | |
| ph.push_back(0); | |
| } | |
| ph[used] = phase; | |
| cudaEventRecord(ev[used], s); | |
| ++used; | |
| } | |
| // every recorded mark has completed (the stream was synchronized): charge the gaps, keep the last mark | |
| void fold() { | |
| if (!on || used < 2) return; | |
| for (size_t i = 0; i + 1 < used; ++i) { | |
| float t = 0.0f; | |
| if (cudaEventElapsedTime(&t, ev[i], ev[i + 1]) == cudaSuccess) ms[ph[i]] += t; | |
| } | |
| std::swap(ev[0], ev[used - 1]); | |
| std::swap(ph[0], ph[used - 1]); | |
| used = 1; | |
| } | |
| ~PfTimer() { | |
| for (cudaEvent_t e : ev) cudaEventDestroy(e); | |
| } | |
| }; | |
| // multi-GPU: the peer's own timeline (STRATA_PREFILL_TIMING): marks on the peer stream, folded with the primary's | |
| enum PePhase { kPeIdle, kPeMoeIn, kPeMoeGemm, kPeMoeOut, kPeCount }; | |
| const char* const kPeNames[kPeCount] = {"idle", "moe in", "moe gemm", "moe out"}; | |
| struct PeTimer { | |
| bool on = std::getenv("STRATA_PREFILL_TIMING") != nullptr; | |
| int dev = -1; | |
| std::vector<cudaEvent_t> ev; | |
| std::vector<int> ph; | |
| size_t used = 0; | |
| double ms[kPeCount] = {}; | |
| void mark(int phase, cudaStream_t s) { // the peer device is current | |
| if (!on) return; | |
| if (used == ev.size()) { | |
| cudaEvent_t e = nullptr; | |
| cudaEventCreate(&e); | |
| ev.push_back(e); | |
| ph.push_back(0); | |
| } | |
| ph[used] = phase; | |
| cudaEventRecord(ev[used], s); | |
| ++used; | |
| } | |
| void fold() { // every mark has completed | |
| if (!on || used < 2) return; | |
| for (size_t i = 0; i + 1 < used; ++i) { | |
| float t = 0.0f; | |
| if (cudaEventElapsedTime(&t, ev[i], ev[i + 1]) == cudaSuccess) ms[ph[i]] += t; | |
| } | |
| std::swap(ev[0], ev[used - 1]); | |
| std::swap(ph[0], ph[used - 1]); | |
| used = 1; | |
| } | |
| ~PeTimer() { | |
| if (dev < 0) return; | |
| int prev = 0; | |
| cudaGetDevice(&prev); | |
| cudaSetDevice(dev); | |
| for (cudaEvent_t e : ev) cudaEventDestroy(e); | |
| cudaSetDevice(prev); | |
| } | |
| }; | |
| } // namespace | |
| bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) { | |
| err.clear(); | |
| Impl& m = *impl_; | |
| const core::OnDevice on_device(m.device); | |
| const core::ModelGeometry& g = *m.g; | |
| core::SessionState& ss = *m.ss; | |
| const auto t_start = Clock::now(); | |
| const int64_t LB = stage_lb_, LE = stage_le_; | |
| // the next stage reads chunk c on a thread while this one reads chunk c + 1 (declared first: an early return | |
| // waits for it before anything it reads goes away) | |
| std::string next_err; | |
| std::future<bool> next_run; | |
| int hand_buf = 0; | |
| double host_sync_ms = 0, host_chunk_ms = 0, host_setup_ms = 0; // STRATA_PREFILL_TIMING: the host's share | |
| strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes(); | |
| s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads; | |
| s.idx_dim = g.idx_key_dim; | |
| const uint64_t gdn_floats = (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size + | |
| (uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1); | |
| int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]}; | |
| PfTimer pt; | |
| PeTimer pe; | |
| if (m.pp) pe.dev = m.pp->dev; else pe.on = false; | |
| const cudaStream_t cs = (cudaStream_t) m.cs; | |
| // the MMQ row table lives in the borrowed cache slots, which the refill after a prompt overwrites with experts: | |
| // write it again for every prompt (a layout is reused as long as the chunk and the slots are the same) | |
| if (m.ids_identity != nullptr) mmq::iota(m.ids_identity, m.T * K, m.cs); | |
| // The PLE rows of a chunk are read from the model file on the host (an SSD read per missed row): the chunk | |
| // after this one is read on a thread while the GPU runs this one, into the other of two buffers. The rows | |
| // depend only on the tokens (the two before a position name its n-grams), so this is the same data. | |
| const bool ple_on = ss.ple.ready() && LB <= 1 && 1 < LE; | |
| // the PLE block batched over the chunk: the pinned postops and a BF16 or GGUF-native key (else token by token); | |
| // STRATA_PLE_BATCH=0 keeps the per-token block (the A/B) | |
| static const bool ple_batch_env = [] { | |
| const char* v = std::getenv("STRATA_PLE_BATCH"); | |
| return v == nullptr || std::atoi(v) != 0; | |
| }(); | |
| const bool ple_batch = ple_on && ple_batch_env && strata::kernels::ple_native_postops_enabled() && | |
| (ss.ple.w.key_bf16 != nullptr || ss.ple.w.key_native_data != nullptr) && | |
| m.region_bytes / ((uint64_t) (3 * strata::kernels::NG_HC_DIM + N + 4) * 4 + (uint64_t) N * 2 + 4096) >= 64; | |
| const int32_t prev0[2] = {prev[0], prev[1]}; | |
| auto ple_gather = [&m, &ss, tokens, n, prev0](int64_t c0, int buf, std::string& e) -> bool { | |
| const int64_t T = std::min(m.T, n - c0); | |
| auto at = [&](int64_t i) { return i < 2 ? prev0[i] : (int32_t) tokens[i - 2]; }; // prev0, then the tokens | |
| int32_t pv[2] = {at(c0), at(c0 + 1)}; | |
| for (int64_t t = 0; t < T; ++t) { | |
| const int32_t tok = (int32_t) tokens[c0 + t]; | |
| strata::kernels::ngram_rows(&tok, pv, 1, ss.ple.consts, | |
| m.ple_rows[buf].data() + t * strata::kernels::PLE_N_HEADS); | |
| pv[0] = pv[1]; | |
| pv[1] = tok; | |
| } | |
| return ss.ple.table->gather_batch(m.ple_rows[buf].data(), (size_t) T, m.ple_emb_host[buf], e); | |
| }; | |
| std::string ple_next_err; | |
| std::future<bool> ple_next; // declared after everything it reads: an early return waits for it | |
| int ple_buf = 0; | |
| for (int64_t c0 = 0; c0 < n; c0 += m.T) { | |
| if (should_stop && should_stop()) { err = "cancelled"; return false; } | |
| if (std::getenv("STRATA_TRACE")) { std::fprintf(stderr, "strata trace: prompt chunk %lld of %lld\n", (long long) c0, (long long) n); std::fflush(stderr); } | |
| const int64_t T = std::min(m.T, n - c0), p0 = pos0 + c0; | |
| core::progress_at("reading the prompt (batched): preparing the chunk from token", p0); // #251 | |
| ++stats_.chunks; | |
| pt.mark(kPfStart, cs); | |
| const auto tsetup = Clock::now(); | |
| // ---- embeddings, broadcast to the four streams - or, in a later stage of a layer split, the rows the | |
| // previous stage handed on | |
| if (hand_in_ != nullptr) { | |
| if (cudaMemcpyAsync(m.R, hand_in_ + (size_t) c0 * D, (size_t) T * D * 4, cudaMemcpyHostToDevice, m.cs) != | |
| cudaSuccess) { | |
| err = "prefill: the layer split's hand-off upload failed"; | |
| return false; | |
| } | |
| } | |
| // C-4: the whole chunk's rows in one gather (the same per-element arithmetic as the per-token path, so the | |
| // same bits); a chunk with picture rows, or a token outside the table, takes the per-token path | |
| bool batched = hand_in_ == nullptr && m.tok_dev != nullptr; | |
| const core::NativeEmbed* nemb = core::native_embed(); | |
| const core::WeightRef* wemb = nemb ? nullptr : m.wt->find("token_embd.weight"); | |
| if (batched && nemb == nullptr && | |
| (wemb == nullptr || wemb->codebook_iq4nl || wemb->ne0 != g.n_embd || wemb->group_elems <= 0 || | |
| (wemb->code_bits != 2 && wemb->code_bits != 4 && wemb->code_bits != 8))) | |
| batched = false; | |
| for (int64_t t = 0; batched && t < T; ++t) { | |
| const int64_t tok = tokens[c0 + t]; | |
| if ((embd_rows && embd_rows[p0 + t]) || tok < 0 || (wemb && tok >= wemb->ne1)) batched = false; | |
| else m.tok_host[(size_t) t] = (int32_t) tok; | |
| } | |
| if (batched) { | |
| if (cudaMemcpyAsync(m.tok_dev, m.tok_host.data(), (size_t) T * sizeof(int32_t), cudaMemcpyHostToDevice, | |
| m.cs) != cudaSuccess) { | |
| err = "prefill: the token id upload failed"; | |
| return false; | |
| } | |
| if (nemb) { | |
| nemb->gather_dev(m.tok_dev, T, m.emb, m.cs); | |
| } else { | |
| const auto* codes = (const uint8_t*) wemb->data; | |
| const auto* scales = (const float*) (codes + wemb->codes_bytes); | |
| const auto* offsets = wemb->has_offset ? (const float*) (codes + wemb->codes_bytes + wemb->scales_bytes) | |
| : nullptr; | |
| strata::kernels::embedding_gather_dev(codes, scales, offsets, m.tok_dev, (int) T, wemb->ne0, | |
| wemb->code_bits, wemb->code_bias, wemb->group_elems, | |
| (uint64_t) (wemb->ne0 / (8 / wemb->code_bits)), | |
| (uint64_t) (wemb->ne0 / wemb->group_elems), m.emb, m.cs); | |
| } | |
| } | |
| for (int64_t t = 0; hand_in_ == nullptr && !batched && t < T; ++t) { | |
| const float* row = embd_rows ? embd_rows[p0 + t] : nullptr; | |
| if (row) { | |
| if (cudaMemcpyAsync(m.emb + t * N, row, (size_t) N * 4, cudaMemcpyHostToDevice, m.cs) != cudaSuccess) { | |
| err = "prefill: the image embedding upload failed"; | |
| return false; | |
| } | |
| } else if (!core::embed_row(*m.wt, g, tokens[c0 + t], m.emb + t * N, m.cs, err)) { | |
| return false; | |
| } | |
| } | |
| if (hand_in_ == nullptr) gr_broadcast(m.emb, m.R, T, m.cs); | |
| // ---- the PLE rows of the whole chunk, one batched SSD request on a thread (see ple_gather). A later chunk's | |
| // were read ahead during the previous chunk; the first chunk's are read beside layer 0 - they are needed from | |
| // layer 1 on, and gathering them here first left the GPU idle for the whole read (~0.4 s of a 32K prompt) | |
| if (ple_on && !ple_next.valid()) | |
| ple_next = std::async(std::launch::async, [&ple_gather, &ple_next_err, c0, b = ple_buf] { | |
| return ple_gather(c0, b, ple_next_err); | |
| }); | |
| bool ple_pending = ple_on; | |
| // the chunk's rows onto the device just before layer 1 reads them, and the next chunk's gather started | |
| auto ple_land = [&]() -> bool { | |
| if (!ple_pending) return true; | |
| ple_pending = false; | |
| const auto tp = Clock::now(); | |
| if (!ple_next.get()) { | |
| err = ple_next_err; | |
| return false; | |
| } | |
| if (cudaMemcpyAsync(m.ple_emb, m.ple_emb_host[ple_buf], (size_t) T * N * 4, cudaMemcpyHostToDevice, m.cs) != | |
| cudaSuccess || | |
| cudaEventRecord(m.ple_copied[ple_buf], m.cs) != cudaSuccess) { | |
| err = std::string("prefill: the PLE rows' upload failed: ") + cudaGetErrorString(cudaGetLastError()); | |
| return false; | |
| } | |
| if (c0 + m.T < n) { | |
| // the other buffer's upload (a chunk ago) is done before the SSD thread refills it | |
| if (cudaEventSynchronize(m.ple_copied[ple_buf ^ 1]) != cudaSuccess) { | |
| err = std::string("prefill: the PLE rows' upload failed: ") + cudaGetErrorString(cudaGetLastError()); | |
| return false; | |
| } | |
| ple_next = std::async(std::launch::async, [&ple_gather, &ple_next_err, c1 = c0 + m.T, b = ple_buf ^ 1] { | |
| return ple_gather(c1, b, ple_next_err); | |
| }); | |
| } | |
| ple_buf ^= 1; | |
| stats_.ms_ple += ms_since(tp); | |
| return true; | |
| }; | |
| for (int64_t t = 0; t < T; ++t) { prev[0] = prev[1]; prev[1] = (int32_t) tokens[c0 + t]; } | |
| // ---- the QSA step records of every position in the chunk | |
| for (int64_t t = 0; t < T; ++t) strata::kernels::qsa_step_fill(m.steps_host.data() + t * strata::kernels::kStepCount, p0 + t, s); | |
| cudaMemcpyAsync(m.steps_dev, m.steps_host.data(), (size_t) T * strata::kernels::kStepCount * 4, | |
| cudaMemcpyHostToDevice, m.cs); | |
| int64_t qsa_index = 0, gdn_index = 0; | |
| for (int64_t l = 0; l < LB; ++l) (core::is_qsa_layer(g, l) ? qsa_index : gdn_index) += 1; | |
| // step 3: this chunk's stream - every non-resident expert of every layer, layer by layer in id order (entry | |
| // k lands in ring slot k % ring); a copy is issued once the entry `ring` before it is consumed (its slot's | |
| // `used` event recorded), so the copy stream never waits on an event that is not queued yet | |
| const strata::kernels::cpu::ExpertLayout& lay0 = strata::kernels::cpu::expert_layout(); | |
| const bool stream_all = m.ring > STAGE && T >= stream_all_min() && m.src != nullptr; | |
| const bool ps_on = stream_all && m.pp && m.pp->ps_frac > 0.0; | |
| if (m.pp && !ps_on) m.pp->ps_flag.clear(); | |
| // multi-GPU: the peer's ring - issue its copies up to `limit` / give entries back (the peer device is current) | |
| auto p_issue_until = [&](size_t limit) { | |
| PeerPrefill& P = *m.pp; | |
| limit = std::min(limit, P.pseq.size()); | |
| while (P.p_issued < limit) { | |
| const PeerPrefill::PsEntry& en = P.pseq[P.p_issued]; | |
| const size_t sl = P.p_issued % (size_t) P.RP; | |
| if (P.plive[sl]) cudaStreamWaitEvent(P.s_cp, P.pused[sl], 0); | |
| cudaMemcpyAsync(P.pstage[sl], en.blob, (size_t) strata::kernels::cpu::expert_layout().blob_bytes(en.l), | |
| cudaMemcpyHostToDevice, P.s_cp); | |
| cudaEventRecord(P.pcopied[sl], P.s_cp); | |
| P.plive[sl] = 1; | |
| ++P.p_issued; | |
| } | |
| }; | |
| auto p_release_to = [&](int64_t l, int32_t e_stop, cudaStream_t ps) { | |
| PeerPrefill& P = *m.pp; | |
| while (P.pk < P.pseq_start[(size_t) l + 1] && P.pseq[P.pk].e < e_stop) { | |
| cudaEventRecord(P.pused[P.pk % (size_t) P.RP], ps); | |
| ++P.pk; | |
| p_issue_until(P.pk + (size_t) P.RP); | |
| } | |
| }; | |
| struct StreamEntry { int32_t l, e; const uint8_t* blob; int job; }; | |
| std::vector<StreamEntry> seq; | |
| std::vector<size_t> seq_start; | |
| size_t issued = 0, consumed = 0; | |
| if (stream_all) { | |
| seq_start.assign((size_t) g.n_layers + 1, 0); | |
| std::vector<Stager::Job> js; | |
| if (ps_on) { | |
| m.pp->pseq.clear(); | |
| m.pp->pseq_start.assign((size_t) g.n_layers + 1, 0); | |
| m.pp->ps_flag.assign((size_t) g.n_layers * m.g->n_expert, 0); | |
| m.pp->p_issued = m.pp->pk = 0; | |
| } | |
| double ps_acc = 0.0; | |
| for (int64_t l = LB; l < LE; ++l) { | |
| seq_start[(size_t) l] = seq.size(); | |
| if (ps_on) m.pp->pseq_start[(size_t) l] = m.pp->pseq.size(); | |
| for (int32_t e = 0; e < m.g->n_expert; ++e) { | |
| if (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) continue; | |
| if (m.pp && m.pp->peer->has(l, e)) continue; // multi-GPU: computed on (or read from) the peer | |
| int job = -1; | |
| const uint8_t* b = nullptr; | |
| if (m.src->transient(l, e)) { // CS-T: copied by the source into the stager's buffer | |
| job = (int) js.size(); | |
| js.push_back({nullptr, (size_t) lay0.blob_bytes(l), m.src, (int32_t) l, e}); | |
| } else { | |
| b = m.src->blob(l, e); | |
| if (!b) { err = "prefill: expert source has no blob"; return false; } | |
| if (ps_on && m.src->pinned(l, e)) { // multi-GPU: every ps_frac-th one goes to the peer's ring | |
| ps_acc += m.pp->ps_frac; | |
| if (ps_acc >= 1.0) { | |
| ps_acc -= 1.0; | |
| m.pp->pseq.push_back({(int32_t) l, e, b}); | |
| m.pp->ps_flag[(size_t) l * m.g->n_expert + e] = 1; | |
| continue; | |
| } | |
| } | |
| if (!m.src->pinned(l, e)) { | |
| job = (int) js.size(); | |
| js.push_back({b, (size_t) lay0.blob_bytes(l)}); | |
| } | |
| } | |
| seq.push_back({(int32_t) l, e, b, job}); | |
| } | |
| } | |
| for (int64_t l = LE; l <= g.n_layers; ++l) { | |
| seq_start[(size_t) l] = seq.size(); | |
| if (ps_on) m.pp->pseq_start[(size_t) l] = m.pp->pseq.size(); | |
| } | |
| m.stager->start(std::move(js)); | |
| } | |
| struct StagerDone { | |
| Stager* st; | |
| ~StagerDone() { if (st) st->finish(); } | |
| } chunk_stager_done{stream_all ? m.stager.get() : nullptr}; | |
| auto issue_until = [&](size_t limit) { | |
| limit = std::min(limit, seq.size()); | |
| while (issued < limit) { | |
| const StreamEntry& en = seq[issued]; | |
| 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[m.used_of[sl]], 0); | |
| if (en.job < 0) { | |
| cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); | |
| ++stats_.experts_dma; | |
| } else { | |
| const uint8_t* hb = m.stager->wait(en.job); | |
| cudaMemcpyAsync(m.stage_dev[sl], hb, bytes, cudaMemcpyHostToDevice, m.copy); | |
| m.stager->issued_one(en.job, m.copy); | |
| } | |
| cudaEventRecord(m.copied[sl], m.copy); | |
| m.stage_live[sl] = true; | |
| stats_.ms_experts_host += ms_since(th); | |
| ++stats_.experts_streamed; | |
| ++issued; | |
| } | |
| }; | |
| // D-5: the stream is issued by its own host thread, so the thread launching the layers' kernels never waits | |
| // behind a host copy of an unpinned blob (that wait left the GPU idle: the 'wait copy' / 'dequant' time of the | |
| // i-quant prompts). The same copies in the same order into the same slots, and a slot is refilled only once | |
| // the compute stream has recorded that it is done with it: the same results. STRATA_PREFILL_ISSUER=0: inline. | |
| static const bool issuer_on = [] { | |
| const char* v = std::getenv("STRATA_PREFILL_ISSUER"); | |
| return v == nullptr || std::atoi(v) != 0; | |
| }(); | |
| std::atomic<size_t> a_issued{0}, a_consumed{0}; | |
| std::atomic<bool> a_stop{false}; | |
| double iss_ms = 0; | |
| int64_t iss_streamed = 0, iss_dma = 0; | |
| std::thread issuer; | |
| struct IssuerJoin { | |
| std::atomic<bool>* stop; | |
| std::thread* t; | |
| ~IssuerJoin() { if (t->joinable()) { stop->store(true); t->join(); } } | |
| } issuer_join{&a_stop, &issuer}; | |
| const bool threaded_issue = stream_all && issuer_on; | |
| if (threaded_issue) { | |
| issuer = std::thread([&] { | |
| const core::OnDevice od(m.device); | |
| for (size_t idx = 0; idx < seq.size(); ++idx) { | |
| while (idx >= a_consumed.load(std::memory_order_acquire) + (size_t) m.ring) { | |
| if (a_stop.load(std::memory_order_acquire)) return; | |
| std::this_thread::yield(); | |
| } | |
| const StreamEntry& en = seq[idx]; | |
| 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[m.used_of[sl]], 0); | |
| if (en.job < 0) { | |
| cudaMemcpyAsync(m.stage_dev[sl], en.blob, bytes, cudaMemcpyHostToDevice, m.copy); | |
| ++iss_dma; | |
| } else { | |
| const uint8_t* hb = m.stager->wait(en.job); | |
| cudaMemcpyAsync(m.stage_dev[sl], hb, bytes, cudaMemcpyHostToDevice, m.copy); | |
| m.stager->issued_one(en.job, m.copy); | |
| } | |
| cudaEventRecord(m.copied[sl], m.copy); | |
| m.stage_live[sl] = true; | |
| iss_ms += ms_since(th); | |
| ++iss_streamed; | |
| a_issued.store(idx + 1, std::memory_order_release); | |
| } | |
| }); | |
| } else if (stream_all) { | |
| issue_until((size_t) m.ring); // layer 0's first experts, behind the embedding and the PLE | |
| } | |
| // the consumer's side: entry k's copy is on the copy stream (the thread issued it), then k is given back | |
| auto wait_issued = [&](size_t k) { | |
| if (!threaded_issue) return; | |
| while (a_issued.load(std::memory_order_acquire) <= k) std::this_thread::yield(); | |
| }; | |
| auto give_back = [&](size_t upto) { | |
| if (threaded_issue) a_consumed.store(upto, std::memory_order_release); | |
| else issue_until(upto + (size_t) m.ring); | |
| }; | |
| if (ps_on) { | |
| int pd = 0; | |
| cudaGetDevice(&pd); | |
| cudaSetDevice(m.pp->dev); | |
| p_issue_until((size_t) m.pp->RP); | |
| cudaSetDevice(pd); | |
| } | |
| host_setup_ms += ms_since(tsetup); | |
| bool normed = false; // F-2: the previous half's write already normed R for this half (grs, xn16) | |
| for (int64_t l = LB; l < LE; ++l) { | |
| core::progress_beat(); // the serve watchdog: a prompt chunk of 8192 tokens is still moving | |
| core::progress_at("reading the prompt (batched): layer", l, p0); // #251: a stall names layer and chunk | |
| const core::LayerView v(*m.wt, l); | |
| if (l == std::max<int64_t>(LB, 1) && !ple_land()) return false; // the PLE rows, read from layer 1 on | |
| // ---- the PLE block at layer 1, token by token (its conv reads the previous tokens' rows) | |
| if (l == 1 && ple_on && ple_batch) { | |
| // the whole chunk at once, in sub-batches carved from the idle scratch region: the key and value | |
| // projections as GEMMs (a token at a time they re-read ~52 MB of BF16 key per token on the IQ | |
| // files), the rest with the per-token kernels' arithmetic (native_ple_postops_batch) | |
| pt.mark(kPfPle, cs); | |
| const auto tp = Clock::now(); | |
| const strata::kernels::PleWeights& pw = ss.ple.w; | |
| constexpr int64_t HD = strata::kernels::NG_HC_DIM; | |
| const uint64_t per_token = (uint64_t) (3 * HD + N + 4) * 4 + (uint64_t) N * (bf16x2() ? 4 : 2) + 4096; | |
| const int64_t SB = std::min<int64_t>(T, (int64_t) (m.region_bytes / per_token)); | |
| for (int64_t s0 = 0; s0 < T; s0 += SB) { | |
| const int64_t nb = std::min(SB, T - s0); | |
| uint8_t* q = m.region; | |
| auto carve_f = [&](size_t n) { float* p = (float*) q; q += (n * 4 + 255) & ~(size_t) 255; return p; }; | |
| float* key = carve_f((size_t) nb * HD); | |
| float* qn = carve_f((size_t) nb * HD); | |
| float* gated = carve_f((size_t) nb * HD); | |
| float* val = carve_f((size_t) nb * N); | |
| float* gate = carve_f((size_t) nb * 4); | |
| uint16_t* e16 = (uint16_t*) carve_f((size_t) nb * N / 2); | |
| uint16_t* e16_lo = bf16x2() ? (uint16_t*) carve_f((size_t) nb * N / 2) : nullptr; | |
| const float* emb = m.ple_emb + s0 * N; | |
| if (pw.key_bf16 != nullptr) { | |
| to_bf16(emb, e16, nb * N, m.cs, e16_lo); | |
| m.gemm.bf16(e16, pw.key_bf16, key, nb, HD, N); | |
| if (e16_lo) m.gemm.bf16(e16_lo, pw.key_bf16, key, nb, HD, N, 0, 1.0f); | |
| } else { | |
| to_f16(emb, e16, nb * N, m.cs); | |
| m.gemm.native(e16, pw.key_native_type, pw.key_native_data, key, nb, HD, N); | |
| to_bf16(emb, e16, nb * N, m.cs, e16_lo); | |
| } | |
| m.gemm.bf16(e16, pw.value_bf16, val, nb, N, N); | |
| if (e16_lo) m.gemm.bf16(e16_lo, pw.value_bf16, val, nb, N, N, 0, 1.0f); | |
| try { | |
| strata::kernels::native_ple_postops_batch(key, m.R + s0 * D, val, ss.ple.hist, pw, qn, gated, | |
| gate, (int) nb, m.cs); | |
| } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } | |
| } | |
| stats_.ms_ple += ms_since(tp); | |
| } else if (l == 1 && ple_on) { | |
| pt.mark(kPfPle, cs); | |
| const auto tp = Clock::now(); | |
| for (int64_t t = 0; t < T; ++t) { | |
| strata::kernels::PleOut po; | |
| po.normalized = m.ple_norm; | |
| po.result = m.R + t * D; | |
| try { | |
| strata::kernels::ple_block(m.ple_emb + t * N, m.R + t * D, ss.ple.hist, ss.ple.w, po, | |
| ss.ple.scratch, m.cs); | |
| } catch (const std::exception& e) { err = std::string("prefill PLE: ") + e.what(); return false; } | |
| strata::kernels::ple_history_advance(ss.ple.hist, m.ple_norm, m.cs); | |
| } | |
| stats_.ms_ple += ms_since(tp); | |
| } | |
| for (int half = 0; half < 2; ++half) { | |
| // ---- the hyper-connection read of this half | |
| const char* pre = half == 0 ? "hc_attn_" : "hc_ffn_"; | |
| const std::string sn = std::string(pre) + "norm.weight", sd = std::string(pre) + "down.weight", | |
| su = std::string(pre) + "up.weight", si = std::string(pre) + "inject.weight"; | |
| const core::WeightRef *wn = need(v, sn.c_str(), err), *wd = need(v, sd.c_str(), err), | |
| *wu = need(v, su.c_str(), err), *wi = need(v, si.c_str(), err); | |
| if (!wn || !wd || !wu || !wi) return false; | |
| pt.mark(kPfHc, cs); | |
| if (gr_unfused()) gr_norm(m.R, (const float*) wn->data, EPS, m.xn, m.xn16, T, m.cs, m.xn16_lo); | |
| else if (!normed) gr_norm_rs(m.R, (const float*) wn->data, EPS, m.grs, m.xn16, T, m.cs, m.xn16_lo); | |
| normed = false; | |
| if (!bf16_proj(m.gemm, wd, m.xn16, m.lo, T, sd, err, 0, m.xn16_lo)) return false; | |
| gr_silu(m.lo, m.lo16, T, m.cs, m.lo16_lo); | |
| if (!bf16_proj(m.gemm, wu, m.lo16, m.gated, T, su, err, 0, m.lo16_lo)) return false; | |
| if (!bf16_proj(m.gemm, wi, m.xn16, m.inj, T, si, err, 0, m.xn16_lo)) return false; | |
| if (gr_unfused()) gr_mix(m.xn, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, m.mixed_bf_lo); | |
| else gr_mix_r(m.R, m.grs, (const float*) wn->data, m.gated, m.mixed, m.mixed_bf, T, m.cs, m.mixed_h, | |
| m.mixed_bf_lo); | |
| if (half == 0 && !core::is_qsa_layer(g, l)) { | |
| // ======================= GDN ======================= | |
| const core::WeightRef *wqkv = need(v, "attn_qkv.weight", err), *wg = need(v, "attn_gate.weight", err), | |
| *wo = need(v, "ssm_out.weight", err), *wa = need(v, "ssm_alpha.weight", err), | |
| *wb = need(v, "ssm_beta.weight", err), *wc = need(v, "ssm_conv1d.weight", err), | |
| *wnm = need(v, "ssm_norm.weight", err), *wdt = need(v, "ssm_dt.bias", err), | |
| *wsa = need(v, "ssm_a", err); | |
| if (!wqkv || !wg || !wo || !wa || !wb || !wc || !wnm || !wdt || !wsa) return false; | |
| pt.mark(kPfGdn, cs); | |
| float* state = ss.gdn_state + (size_t) (gdn_index - ss.gdn_ord0) * gdn_floats; | |
| float* conv = state + (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size; | |
| if (!native_proj(m.gemm, wqkv, m.mixed_h, m.qkv, T, v.name("attn_qkv.weight"), err)) return false; | |
| if (!native_proj(m.gemm, wg, m.mixed_h, m.z, T, v.name("attn_gate.weight"), err)) return false; | |
| if (!bf16_proj(m.gemm, wa, m.mixed_bf, m.ab, T, v.name("ssm_alpha.weight"), err, 2 * HV, m.mixed_bf_lo)) return false; | |
| if (!bf16_proj(m.gemm, wb, m.mixed_bf, m.ab + HV, T, v.name("ssm_beta.weight"), err, 2 * HV, m.mixed_bf_lo)) return false; | |
| pt.mark(kPfGdnConv, cs); // "gdn" is the projections in; the rest on their own lines | |
| gdn_gates(m.ab, (const float*) wdt->data, (const float*) wsa->data, m.gate, m.beta, T, m.cs); | |
| gdn_conv(conv, m.qkv, (const float*) wc->data, m.hbuf, T, EPS, m.cs); | |
| pt.mark(kPfGdnRec, cs); | |
| gdn_recurrence(state, m.hbuf, m.gate, m.beta, m.z, (const float*) wnm->data, EPS, m.y, m.y_h, T, m.cs); | |
| pt.mark(kPfGdnOut, cs); | |
| if (!native_proj(m.gemm, wo, m.y_h, m.bo, T, v.name("ssm_out.weight"), err)) return false; | |
| ++gdn_index; | |
| } else if (half == 0) { | |
| // ======================= QSA ======================= | |
| const core::QsaState& st = ss.qsa_states[qsa_index]; | |
| const core::WeightRef *wq = need(v, "attn_q.weight", err), *wk = need(v, "attn_k.weight", err), | |
| *wv = need(v, "attn_v.weight", err), *wo = need(v, "attn_output.weight", err), | |
| *wik = need(v, "indexer.k_proj.weight", err), | |
| *wiq = need(v, "indexer.q_proj.weight", err), | |
| *wqn = need(v, "attn_q_norm.weight", err), *wkn = need(v, "attn_k_norm.weight", err), | |
| *wiqn = need(v, "indexer.q_norm.weight", err), | |
| *wikn = need(v, "indexer.k_norm.weight", err); | |
| if (!wq || !wk || !wv || !wo || !wik || !wiq || !wqn || !wkn || !wiqn || !wikn) return false; | |
| pt.mark(kPfQsa, cs); | |
| if (!native_proj(m.gemm, wk, m.mixed_h, m.Kc, T, v.name("attn_k.weight"), err)) return false; | |
| if (!native_proj(m.gemm, wv, m.mixed_h, m.Vc, T, v.name("attn_v.weight"), err)) return false; | |
| if (!native_proj(m.gemm, wq, m.mixed_h, m.Qf, T, v.name("attn_q.weight"), err)) return false; | |
| if (!bf16_proj(m.gemm, wik, m.mixed_bf, m.idx_raw, T, v.name("indexer.k_proj.weight"), err, 0, m.mixed_bf_lo)) return false; | |
| if (!bf16_proj(m.gemm, wiq, m.mixed_bf, m.q_idx, T, v.name("indexer.q_proj.weight"), err, 0, m.mixed_bf_lo)) return false; | |
| rms_rows(m.Kc, (const float*) wkn->data, T * 2, 256, 256, EPS, m.cs); | |
| rope(m.Kc, T, 2, 256, 512, p0, strata::kernels::rope_scaling(), m.cs); | |
| // KV streaming: this layer's cells [0, p0) come in from the host copy to the staging pool, and the | |
| // chunk's cells go to the host copy, the staging pool, and the VRAM slots of resident blocks | |
| const bool staged = st.kv_mode == 1; | |
| if (staged) { | |
| pt.mark(kPfKvStage, cs); | |
| strata::kernels::kv_stage_from_host(pools_of(m.stage, m.ident_table), st.host, | |
| core::qsa_kv_format(st), | |
| (p0 + s.page_size - 1) / s.page_size, s, m.cs); | |
| pt.mark(kPfQsa, cs); | |
| } | |
| if (st.kv_hybrid) { // K8V4: K INT8 unrotated, V rotated Q4_0 (only V and the output rotate) | |
| strata::kernels::fwht256_inplace_cuda(m.Vc, T * 2, m.cs); | |
| kv_append(m.Kc, m.Kc, T, p0, st.page_table, s.page_size, nullptr, nullptr, | |
| st.k_q, st.k_q, st.k_scale, st.k_scale, m.cs, nullptr, // mode 0: no host mirror | |
| staged ? &m.stage : nullptr); | |
| strata::kernels::kv_append_q4(st.v_q4, st.v_q4, st.page_table, p0, T, m.Vc, m.Vc, s, m.cs, | |
| nullptr, staged ? &m.stage : nullptr); | |
| } else { | |
| if (st.kv_rot) { // rotated K and V (kv_q4.hpp), the queries below too, the output back | |
| strata::kernels::fwht256_inplace_cuda(m.Kc, T * 2, m.cs); | |
| strata::kernels::fwht256_inplace_cuda(m.Vc, T * 2, m.cs); | |
| } | |
| if (st.kv_q4) | |
| strata::kernels::kv_append_q4(st.k_q4, st.v_q4, st.page_table, p0, T, m.Kc, m.Vc, s, m.cs, | |
| &st.host, staged ? &m.stage : nullptr); | |
| else | |
| kv_append(m.Kc, m.Vc, T, p0, st.page_table, s.page_size, st.kv_int8 ? nullptr : st.k_pool, | |
| st.kv_int8 ? nullptr : st.v_pool, st.k_q, st.v_q, st.k_scale, st.v_scale, m.cs, | |
| &st.host, staged ? &m.stage : nullptr); | |
| } | |
| split_q(m.Qf, m.q, T, m.cs); | |
| rms_rows(m.q, (const float*) wqn->data, T * 24, 256, 256, EPS, m.cs); | |
| rope(m.q, T, 24, 256, 6144, p0, strata::kernels::rope_scaling(), m.cs); | |
| if (st.kv_rot) strata::kernels::fwht256_inplace_cuda(m.q, T * 24, m.cs); | |
| rms_rows(m.q_idx, (const float*) wiqn->data, T * 4, 128, 128, EPS, m.cs); | |
| rope(m.q_idx, T, 4, 128, 512, p0, strata::kernels::rope_scaling(), m.cs); | |
| // the indexer appends, token by token; then scores + selection for many queries at once: | |
| // a query reads completed blocks (final once completed) and `dead` for its own tail block | |
| const strata::kernels::QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos}; | |
| pt.mark(kPfQsaIdx, cs); | |
| // C-2: the chunk's appends in three launches instead of one per token (the same end state: | |
| // the queries below read it only after the whole chunk is appended). STRATA_INDEXER_PER_TOKEN=1: the old | |
| try { | |
| static const bool per_token = std::getenv("STRATA_INDEXER_PER_TOKEN") != nullptr; | |
| if (!per_token) { | |
| strata::kernels::native_qsa_indexer_append_batch(m.idx_raw, T, p0, 0, (const float*) wikn->data, | |
| EPS, ib, s, st.max_cells, | |
| strata::kernels::rope_scaling(), m.cs); | |
| } | |
| for (int64_t t = 0; per_token && t < T; ++t) { | |
| const int32_t* step_t = m.steps_dev + t * strata::kernels::kStepCount; | |
| strata::kernels::native_qsa_indexer_append(m.idx_raw + t * 128, step_t + strata::kernels::kStepPos, 0, | |
| (const float*) wikn->data, EPS, ib, s, st.max_cells, | |
| strata::kernels::rope_scaling(), m.cs); | |
| } | |
| } catch (const std::exception& e) { err = std::string("prefill indexer: ") + e.what(); return false; } | |
| pt.mark(kPfQsaSel, cs); | |
| for (int64_t t0 = 0; t0 < T; t0 += m.sel_batch) { | |
| const int64_t nb = std::min(m.sel_batch, T - t0); | |
| const int32_t* steps0 = m.steps_dev + t0 * strata::kernels::kStepCount; | |
| // C-1: the grid reaches the batch's last query's n_bid (they rise with the position) | |
| const int64_t active = (int64_t) m.steps_host[(size_t) ((t0 + nb - 1) * strata::kernels::kStepCount + | |
| strata::kernels::kStepNBid)] + 1; | |
| // the scores on tensor cores (3xTF32: FP32-level, not bitwise); STRATA_SELECT_OLD=1: the warp kernel | |
| static const bool old_sel = std::getenv("STRATA_SELECT_OLD") != nullptr; | |
| if (old_sel || !strata::kernels::qsa_block_scores_tc(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512, | |
| steps0, nb, m.max_blocks, s, m.sel_scores, | |
| m.cs, active)) | |
| strata::kernels::qsa_block_scores(st.idx_pooled, st.idx_dead, m.q_idx + t0 * 512, steps0, nb, | |
| m.max_blocks, s, m.sel_scores, m.cs, active); | |
| strata::kernels::qsa_block_topk(m.sel_scores, steps0, nb, m.max_blocks, m.cap, s, | |
| m.sel_ids + t0 * m.cap, m.cs, active); | |
| } | |
| // STRATA_SEL_OVERLAP (debug, D-1's question): how much do neighbouring queries' selections share? | |
| // Per tile of 16 queries: the union of their selected cells against the sum of their widths. | |
| if (static const bool ovl = std::getenv("STRATA_SEL_OVERLAP") != nullptr; ovl && qsa_index == 0) { | |
| std::vector<int32_t> ids((size_t) (T * m.cap)); | |
| cudaMemcpyAsync(ids.data(), m.sel_ids, ids.size() * 4, cudaMemcpyDeviceToHost, m.cs); | |
| cudaStreamSynchronize(m.cs); | |
| double sum_w = 0, sum_u = 0; | |
| for (int64_t t0 = 0; t0 + 16 <= T; t0 += 16) { | |
| std::vector<int32_t> u; | |
| for (int64_t t = t0; t < t0 + 16; ++t) { | |
| const int64_t w = m.steps_host[(size_t) (t * strata::kernels::kStepCount + strata::kernels::kStepWidth)]; | |
| sum_w += (double) w; | |
| u.insert(u.end(), ids.begin() + t * m.cap, ids.begin() + t * m.cap + w); | |
| } | |
| std::sort(u.begin(), u.end()); | |
| sum_u += (double) (std::unique(u.begin(), u.end()) - u.begin()); | |
| } | |
| std::fprintf(stderr, "strata prefill: selection overlap at %lld: 16-query tiles read %.1f%% of the " | |
| "cells one query at a time does\n", (long long) p0, sum_w > 0 ? 100.0 * sum_u / sum_w : 0.0); | |
| } | |
| // STRATA_IDX_FP16_CHECK: would FP16 pooled indexer keys select the same cells? (the KV-streaming | |
| // design's last question). Every query is selected again from the pooled keys and `dead` rounded | |
| // to fp16 (exactly what an fp16 store reads back); the agreement with the fp32 selection is | |
| // printed cumulatively after each chunk's last QSA layer. Debug: syncs per layer. | |
| if (static const bool f16chk = std::getenv("STRATA_IDX_FP16_CHECK") != nullptr; f16chk) { | |
| static float *pooled16 = nullptr, *dead16 = nullptr; | |
| static int32_t* ids16 = nullptr; | |
| static double shared = 0, cells = 0; | |
| static long long queries = 0, same = 0, sel_queries = 0; | |
| const int64_t rows = st.idx_pooled_rows; | |
| if (pooled16 == nullptr && | |
| (cudaMalloc((void**) &pooled16, (size_t) rows * s.idx_dim * 4) != cudaSuccess || | |
| cudaMalloc((void**) &dead16, (size_t) s.idx_dim * 4) != cudaSuccess || | |
| cudaMalloc((void**) &ids16, (size_t) (m.T * m.cap) * 4) != cudaSuccess)) { | |
| err = "STRATA_IDX_FP16_CHECK: no room for its buffers"; | |
| return false; | |
| } | |
| round_f16(st.idx_pooled, pooled16, rows * s.idx_dim, m.cs); | |
| round_f16(st.idx_dead, dead16, s.idx_dim, m.cs); | |
| for (int64_t t0 = 0; t0 < T; t0 += m.sel_batch) { | |
| const int64_t nb = std::min(m.sel_batch, T - t0); | |
| const int32_t* steps0 = m.steps_dev + t0 * strata::kernels::kStepCount; | |
| strata::kernels::qsa_block_scores(pooled16, dead16, m.q_idx + t0 * 512, steps0, nb, | |
| m.max_blocks, s, m.sel_scores, m.cs); | |
| strata::kernels::qsa_block_topk(m.sel_scores, steps0, nb, m.max_blocks, m.cap, s, | |
| ids16 + t0 * m.cap, m.cs); | |
| } | |
| std::vector<int32_t> a((size_t) (T * m.cap)), b((size_t) (T * m.cap)); | |
| cudaMemcpyAsync(a.data(), m.sel_ids, a.size() * 4, cudaMemcpyDeviceToHost, m.cs); | |
| cudaMemcpyAsync(b.data(), ids16, b.size() * 4, cudaMemcpyDeviceToHost, m.cs); | |
| cudaStreamSynchronize(m.cs); | |
| for (int64_t t = 0; t < T; ++t) { | |
| const int64_t w = m.steps_host[(size_t) (t * strata::kernels::kStepCount + strata::kernels::kStepWidth)]; | |
| const int32_t *x = a.data() + t * m.cap, *y = b.data() + t * m.cap; | |
| int64_t i = 0, j = 0, c = 0; | |
| while (i < w && j < w) { | |
| if (x[i] == y[j]) { ++c; ++i; ++j; } else if (x[i] < y[j]) ++i; else ++j; | |
| } | |
| ++queries; | |
| same += c == w; | |
| if (p0 + t + 1 > m.cap) { ++sel_queries; shared += (double) c; cells += (double) w; } | |
| } | |
| if (qsa_index + 1 == g.n_qsa_layers()) | |
| std::fprintf(stderr, "strata prefill: FP16 indexer keys: %lld of %lld selections identical; " | |
| "where the selection is sparse, %.4f%% of cells shared (%lld queries)\n", | |
| same, queries, cells > 0 ? 100.0 * shared / cells : 100.0, sel_queries); | |
| } | |
| // STRATA_QSA_DUMP=<file>: append every QSA layer's selected cells for the prompt's last | |
| // STRATA_QSA_DUMP_LAST (4096) positions - records of int32 {qsa layer, pos0, T, cap} + T*cap cells, | |
| // for tools/qsa_locality.py (how local the sparse attention's reads are: the KV-streaming question) | |
| if (static const char* dump = std::getenv("STRATA_QSA_DUMP"); dump != nullptr) { | |
| static const long long last = std::getenv("STRATA_QSA_DUMP_LAST") | |
| ? std::atoll(std::getenv("STRATA_QSA_DUMP_LAST")) : 4096; | |
| if (p0 + T > pos0 + n - last) { | |
| std::vector<int32_t> h((size_t) (T * m.cap)); | |
| cudaMemcpyAsync(h.data(), m.sel_ids, h.size() * 4, cudaMemcpyDeviceToHost, m.cs); | |
| cudaStreamSynchronize(m.cs); | |
| if (std::FILE* f = std::fopen(dump, "ab")) { | |
| const int32_t hdr[4] = {(int32_t) qsa_index, (int32_t) p0, (int32_t) T, (int32_t) m.cap}; | |
| std::fwrite(hdr, 4, 4, f); | |
| std::fwrite(h.data(), 4, h.size(), f); | |
| std::fclose(f); | |
| } | |
| } | |
| } | |
| const strata::kernels::QsaAttnPools pools = staged ? pools_of(m.stage, m.ident_table) | |
| : core::qsa_attn_pools(st); | |
| pt.mark(kPfQsaAttn, cs); | |
| // perf-review D-1: the whole chunk on tensor cores, one block per (query, KV head), FP32-level | |
| // accuracy but not bitwise (qsa_prompt_attn.hpp). Q4_0 KV, or STRATA_PROMPT_ATTN_OLD=1: the | |
| // decode kernel, 32 queries at a time (K8V4 runs the tensor kernel's mode 3: INT8 K, | |
| // V dequantized from its q4_0 blocks to fp16 at gather) | |
| static const bool old_attn = std::getenv("STRATA_PROMPT_ATTN_OLD") != nullptr; | |
| if (old_attn || !strata::kernels::qsa_prompt_attn_batch(m.q, pools, m.sel_ids, m.steps_dev, m.cap, s, | |
| m.attn, T, m.cs)) | |
| for (int64_t t0 = 0; t0 < T; t0 += m.attn_batch) { | |
| const int64_t nb = std::min(m.attn_batch, T - t0); | |
| strata::kernels::qsa_decode_attn_batch(m.q + t0 * ZV, pools, m.sel_ids + t0 * m.cap, | |
| m.steps_dev + t0 * strata::kernels::kStepCount, m.cap, | |
| s, m.attn_scratch, m.attn + t0 * ZV, nb, m.cs); | |
| } | |
| if (st.kv_rot || st.kv_hybrid) strata::kernels::fwht256_inplace_cuda(m.attn, T * 24, m.cs); | |
| pt.mark(kPfQsa, cs); | |
| gate_attn(m.attn, m.Qf, m.attn_h, T, m.cs); | |
| if (!native_proj(m.gemm, wo, m.attn_h, m.bo, T, v.name("attn_output.weight"), err)) return false; | |
| ++qsa_index; | |
| } else { | |
| // ======================= MoE ======================= | |
| const core::WeightRef *wr = need(v, "ffn_gate_inp.weight", err), | |
| *wgi = need(v, "ffn_gate_inp_shexp.weight", err), | |
| *wsg = need(v, "ffn_gate_shexp.weight", err), | |
| *wsu = need(v, "ffn_up_shexp.weight", err), | |
| *wsd = need(v, "ffn_down_shexp.weight", err); | |
| if (!wr || !wgi || !wsg || !wsu || !wsd) return false; | |
| pt.mark(kPfRouter, cs); | |
| if (!bf16_proj(m.gemm, wr, m.mixed_bf, m.logits, T, v.name("ffn_gate_inp.weight"), err, 0, m.mixed_bf_lo)) return false; | |
| route(m.logits, m.ids, m.w, T, m.g->n_expert, m.cs); | |
| // the shared expert and its scalar gate | |
| if (!native_proj(m.gemm, wsg, m.mixed_h, m.sgate, T, v.name("ffn_gate_shexp.weight"), err)) return false; | |
| if (!native_proj(m.gemm, wsu, m.mixed_h, m.sup, T, v.name("ffn_up_shexp.weight"), err)) return false; | |
| swiglu_pair(m.sgate, m.sup, m.sh_h, T, m.cs); | |
| if (!native_proj(m.gemm, wsd, m.sh_h, m.shared, T, v.name("ffn_down_shexp.weight"), err)) return false; | |
| if (wgi->kind != core::WeightKind::Bf16InF32) { err = "prefill: shared gate is not BF16"; return false; } | |
| m.gemm.bf16(m.mixed_bf, (const uint16_t*) wgi->data, m.sg, T, 1, N); | |
| if (m.mixed_bf_lo) m.gemm.bf16(m.mixed_bf_lo, (const uint16_t*) wgi->data, m.sg, T, 1, N, 0, 1.0f); | |
| // #136: STRATA_PF_FUSED=1 - the Q2_0 pack's experts on the fused int8 kernels (moe_fused.hpp), | |
| // grouped on the GPU: no host sync. Only where every expert's place is known before the routing - | |
| // the streamed walk, in which every non-resident expert of the layer comes through the ring in id | |
| // order - and where the MMQ buffers exist: they hold the fused path's own (the per-token int8 | |
| // activations in Xq, the int8 H in H, the grouping tables in GU). Chunks below stream_all_min() | |
| // keep MMQ; without the variable nothing here runs. A native pack's layer takes the native kernels | |
| // (moe_fused_iq.hpp) where they cover its two formats, else MMQ (or the FP16 path: IQ1_M). | |
| const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout(); | |
| const bool use_mmq = mmq_plan().any && mmq_plan().layer[(size_t) l]; | |
| const int mmq_gt = lay.native ? lay.fmt[(size_t) l].gu_type : 42; | |
| const int mmq_dt = lay.native ? lay.fmt[(size_t) l].d_type : 42; | |
| // --peer-device: MMQ only, whether or not the peer took the prompt path (set_peer can decline), as | |
| // fused_ring() sized the ring and the buffers for it | |
| const bool no_peer = !core::peer_portable(); | |
| const bool fused_nat = use_mmq && stream_all && no_peer && lay.native && fused::native_supported(mmq_gt, mmq_dt); | |
| const bool fused_l = (use_mmq && stream_all && no_peer && !lay.native && fused::enabled()) || fused_nat; | |
| size_t n_order = 0; // the routed experts (the debug report; unknown when fused) | |
| bool peer_now = false; // multi-GPU: the peer computed rows of this layer (MMQ path only) | |
| if (fused_l) { | |
| if (static bool said = false; !said) { | |
| said = true; | |
| std::fprintf(stderr, "strata: prompt experts on the fused int8 kernels (STRATA_PF_FUSED=1, " | |
| "#136)\n"); | |
| } | |
| pt.mark(kPfGather, cs); | |
| // the layer's input to int8 once per token; the rows of each expert from the router's ids | |
| if (fused_nat) fused::quantize_act_native(m.mixed, T, N, m.Xq, m.cs); | |
| else fused::quantize_act(m.mixed, T, N, m.Xq, m.cs); | |
| fused::group(m.ids, T * K, (int) K, (int) m.g->n_expert, m.GU, m.slot_dev, m.src_dev, m.cs); | |
| // launches over the experts in id order, each at most kMaxBatch experts of which at most a | |
| // third of the ring streamed: the next batch's blobs arrive while one computes. An expert | |
| // the routing did not pick has no tiles; its ring slot is given back with its batch. | |
| const size_t per = (size_t) std::max(1, m.ring / 3); | |
| size_t k = seq_start[(size_t) l]; | |
| const size_t kend = seq_start[(size_t) l + 1]; | |
| for (int32_t e = 0; e < m.g->n_expert;) { | |
| fused::Batch b; | |
| b.e0 = e; | |
| const size_t k0 = k; | |
| for (; e < m.g->n_expert && e - b.e0 < fused::kMaxBatch; ++e) { | |
| if (k < kend && seq[k].e == e) { | |
| if (k - k0 == per) break; | |
| b.blob[e - b.e0] = m.stage_dev[k % (size_t) m.ring]; | |
| ++k; | |
| } else { // not streamed: resident (the walk streams all others) | |
| b.blob[e - b.e0] = m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]); | |
| // counted whether routed or not (the routing stays on the GPU): at a streamed | |
| // chunk's size (>= 1024 tokens x 10 of 512) nearly every expert is routed | |
| ++stats_.experts_resident; | |
| } | |
| } | |
| b.e1 = e; | |
| if (k > k0) { | |
| // one copy stream, in order: the batch's last blob covers the others | |
| pt.mark(kPfWaitCopy, cs); | |
| wait_issued(k - 1); | |
| cudaStreamWaitEvent(m.cs, m.copied[(k - 1) % (size_t) m.ring], 0); | |
| } | |
| pt.mark(kPfGemmGU, cs); | |
| if (fused_nat) { | |
| const auto& f = lay.fmt[(size_t) l]; | |
| const fused::NativeGeom ng{f.gu_type, f.d_type, f.gu_row, f.d_row, f.up_off, f.down_off}; | |
| fused::experts_native(b, ng, (int) m.g->n_expert, T * K, m.GU, m.Xq, m.src_dev, m.H, | |
| m.Dm, m.cs); | |
| } 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); | |
| m.used_of[kk % (size_t) m.ring] = (int) (kk % (size_t) m.ring); | |
| } | |
| if (k > k0) { | |
| consumed = k; | |
| give_back(consumed); | |
| } | |
| } | |
| } else { | |
| // group the (token, k) pairs by expert on the host | |
| pt.mark(kPfHostGroup, cs); | |
| // (the sync below also orders this layer's writes of slot/src/bounds after the previous | |
| // layer's kernels that read them) | |
| const bool grp_mapped = m.grp_host != nullptr; | |
| int32_t* ids_h = grp_mapped ? m.grp_host : m.ids_host.data(); | |
| int32_t* slot_h = grp_mapped ? m.grp_host + m.grp_tk : m.slot_host.data(); | |
| int32_t* src_h = grp_mapped ? m.grp_host + 2 * m.grp_tk : m.src_host.data(); | |
| if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * K, m.cs); | |
| else cudaMemcpyAsync(m.ids_host.data(), m.ids, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs); | |
| cudaStreamSynchronize(m.cs); | |
| pt.fold(); | |
| if (pe.on) { // the peer's marks so far are done: the primary waited for its last rows | |
| int pd = 0; cudaGetDevice(&pd); cudaSetDevice(pe.dev); cudaStreamSynchronize(m.pp->s); pe.fold(); cudaSetDevice(pd); | |
| } | |
| std::fill(m.cnt.begin(), m.cnt.end(), 0); | |
| for (int64_t i = 0; i < T * K; ++i) { | |
| const int32_t e = ids_h[(size_t) i]; | |
| if (e < 0 || e >= m.g->n_expert) { err = "prefill: routed id out of range"; return false; } | |
| ++m.cnt[(size_t) e]; | |
| } | |
| // multi-GPU: the rows of the experts the peer computes go last, as one block [rows_local, T*K) | |
| const bool pre_mmq = mmq_plan().any && mmq_plan().layer[(size_t) l]; | |
| std::vector<char> on_peer; | |
| int64_t rows_local = T * K, rows_peer = 0; | |
| if (m.pp && pre_mmq) { | |
| on_peer.assign((size_t) m.g->n_expert, 0); | |
| for (int32_t e = 0; e < m.g->n_expert; ++e) { | |
| const int32_t c = m.cnt[(size_t) e]; | |
| if (c > 0 && !m.pp->ps_flag.empty() && m.pp->ps_flag[(size_t) l * m.g->n_expert + e]) { // peer-streamed | |
| on_peer[(size_t) e] = 2; | |
| rows_peer += c; | |
| continue; | |
| } | |
| if (c == 0 || (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) || | |
| !m.pp->peer->has(l, e)) | |
| continue; | |
| if (rows_peer + c > m.pp->cap_rows) { ++m.pp->over_cap; continue; } | |
| on_peer[(size_t) e] = 1; | |
| rows_peer += c; | |
| } | |
| rows_local = T * K - rows_peer; | |
| } | |
| { | |
| int32_t r = 0; | |
| for (int32_t e = 0; e < m.g->n_expert; ++e) | |
| if (on_peer.empty() || !on_peer[(size_t) e]) { m.off[(size_t) e] = r; r += m.cnt[(size_t) e]; } | |
| // the peer's rows in order_peer's order (peer-held, then peer-streamed): its groups need them contiguous | |
| for (int kind = 1; kind <= 2 && !on_peer.empty(); ++kind) | |
| for (int32_t e = 0; e < m.g->n_expert; ++e) | |
| if (on_peer[(size_t) e] == kind) { m.off[(size_t) e] = r; r += m.cnt[(size_t) e]; } | |
| m.off[(size_t) m.g->n_expert] = r; | |
| } | |
| std::vector<int32_t> fill(m.off.begin(), m.off.end() - 1); | |
| for (int64_t i = 0; i < T * K; ++i) { | |
| const int32_t e = ids_h[(size_t) i]; | |
| const int32_t p = fill[(size_t) e]++; | |
| slot_h[(size_t) i] = p; | |
| src_h[(size_t) p] = (int32_t) (i / K); | |
| } | |
| if (grp_mapped) { | |
| copy_i32(m.slot_dev, m.grp_dev + m.grp_tk, T * K, m.cs); | |
| copy_i32(m.src_dev, m.grp_dev + 2 * m.grp_tk, T * K, m.cs); | |
| } else { | |
| cudaMemcpyAsync(m.slot_dev, m.slot_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs); | |
| cudaMemcpyAsync(m.src_dev, m.src_host.data(), (size_t) T * K * 4, cudaMemcpyHostToDevice, m.cs); | |
| } | |
| // the experts, in id order: resident ones from VRAM, the others through the staging ring | |
| std::vector<int32_t> order, order_peer; | |
| for (int32_t e = 0; e < m.g->n_expert; ++e) | |
| if (m.cnt[(size_t) e] > 0 && (on_peer.empty() || on_peer[(size_t) e] != 2)) | |
| (!on_peer.empty() && on_peer[(size_t) e] ? order_peer : order).push_back(e); | |
| for (int32_t e = 0; e < m.g->n_expert && !on_peer.empty(); ++e) // the peer-streamed ones last: | |
| if (on_peer[(size_t) e] == 2) order_peer.push_back(e); // their copies get the most time | |
| n_order = order.size(); | |
| 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; | |
| pt.mark(kPfGather, cs); | |
| if (use_mmq) { | |
| // step 2b: the layer's activations as q8_1 rows in expert order, straight from `mixed` | |
| mmq::quantize(m.mixed, m.src_dev, m.Xq, mmq_gt, N, N, T * K, m.cs); | |
| // each group's rows: absolute bounds (gate/up reads the layer's rows), relative ones (down | |
| // reads the group's own quantized H) | |
| const size_t n = order.size(), ng = (n + MMQ_GROUP - 1) / MMQ_GROUP; | |
| m.bounds_host.resize(n + 1 + ng * (MMQ_GROUP + 1)); | |
| for (size_t j = 0; j < n; ++j) m.bounds_host[j] = m.off[(size_t) order[j]]; | |
| m.bounds_host[n] = (int32_t) rows_local; | |
| for (size_t g = 0; g < ng; ++g) | |
| for (size_t i = 0; i <= MMQ_GROUP; ++i) | |
| m.bounds_host[n + 1 + g * (MMQ_GROUP + 1) + i] = | |
| m.bounds_host[std::min(n, g * MMQ_GROUP + i)] - m.bounds_host[g * MMQ_GROUP]; | |
| if (grp_mapped) { | |
| int32_t* bh = m.grp_host + 3 * m.grp_tk; | |
| std::memcpy(bh, m.bounds_host.data(), m.bounds_host.size() * 4); | |
| copy_i32(m.bounds_dev, m.grp_dev + 3 * m.grp_tk, (int64_t) m.bounds_host.size(), m.cs); | |
| } else { | |
| cudaMemcpyAsync(m.bounds_dev, m.bounds_host.data(), m.bounds_host.size() * 4, | |
| cudaMemcpyHostToDevice, m.cs); | |
| } | |
| } else { | |
| gather_rows16(m.mixed_h, m.src_dev, m.Xs, T * K, N, m.cs); | |
| } | |
| // multi-GPU: the peer's share, enqueued before the primary's own experts so both cards work at once | |
| peer_now = use_mmq && !order_peer.empty(); | |
| if (peer_now) { | |
| PeerPrefill& P = *m.pp; | |
| cudaEventRecord(P.ev_in, m.cs); | |
| int prevd = 0; | |
| cudaGetDevice(&prevd); | |
| cudaSetDevice(P.dev); | |
| const cudaStream_t ps = P.s; | |
| cudaStreamWaitEvent(ps, P.ev_in, 0); | |
| if (P.out_pending) { cudaStreamWaitEvent(ps, P.ev_done, 0); P.out_pending = false; } | |
| pe.mark(kPeMoeIn, ps); | |
| cudaMemcpyPeerAsync(P.mixed, P.dev, m.mixed, prevd, (size_t) T * N * 4, ps); | |
| cudaMemcpyAsync(P.src, src_h + rows_local, // multi-GPU: the mapped table when on | |
| (size_t) rows_peer * 4, cudaMemcpyHostToDevice, ps); | |
| const size_t n = order_peer.size(); | |
| if (P.compact) { | |
| // groups of up to MMQ_GROUP experts and at most G rows, each computed from its own rows | |
| P.groups.clear(); | |
| { | |
| size_t g0 = 0; | |
| int64_t gr = 0; | |
| for (size_t j = 0; j < n; ++j) { | |
| const int64_t c = m.cnt[(size_t) order_peer[j]]; | |
| if (j > g0 && (j - g0 == (size_t) MMQ_GROUP || gr + c > P.G)) { P.groups.push_back({g0, j}); g0 = j; gr = 0; } | |
| gr += c; | |
| } | |
| if (n > g0) P.groups.push_back({g0, n}); | |
| } | |
| const size_t ng = P.groups.size(); | |
| P.bounds_host.resize(n + 1 + ng * (MMQ_GROUP + 1)); | |
| for (size_t j = 0; j < n; ++j) P.bounds_host[j] = m.off[(size_t) order_peer[j]] - (int32_t) rows_local; | |
| P.bounds_host[n] = (int32_t) rows_peer; | |
| for (size_t g2 = 0; g2 < ng; ++g2) | |
| for (size_t i = 0; i <= MMQ_GROUP; ++i) | |
| P.bounds_host[n + 1 + g2 * (MMQ_GROUP + 1) + i] = | |
| P.bounds_host[std::min(P.groups[g2].second, P.groups[g2].first + i)] - | |
| P.bounds_host[P.groups[g2].first]; | |
| cudaMemcpyAsync(P.bounds, P.bounds_host.data(), P.bounds_host.size() * 4, cudaMemcpyHostToDevice, ps); | |
| pe.mark(kPeMoeGemm, ps); | |
| const auto& f = lay.fmt[(size_t) l]; | |
| for (size_t g2 = 0; g2 < ng; ++g2) { | |
| const size_t j0 = P.groups[g2].first, j1 = P.groups[g2].second; | |
| const int ngx = (int) (j1 - j0); | |
| int64_t maxr = 0; | |
| for (size_t j = j0; j < j1; ++j) { | |
| const int32_t e = order_peer[j]; | |
| maxr = std::max<int64_t>(maxr, m.cnt[(size_t) e]); | |
| const uint8_t* bd = nullptr; | |
| int psl = -1; | |
| if (on_peer[(size_t) e] == 2) { // from the peer's ring | |
| p_release_to(l, e, ps); | |
| if (P.pk >= P.pseq_start[(size_t) l + 1] || P.pseq[P.pk].e != e) { | |
| err = "prefill: a peer-streamed expert is not next in the peer's ring"; | |
| cudaSetDevice(prevd); | |
| return false; | |
| } | |
| psl = (int) (P.pk % (size_t) P.RP); | |
| cudaStreamWaitEvent(ps, P.pcopied[(size_t) psl], 0); | |
| bd = P.pstage[(size_t) psl]; | |
| ++P.ps_experts; | |
| } else { | |
| bd = P.peer->slot_ptr(l, e); | |
| } | |
| const size_t q = j - j0; | |
| if (lay.native) | |
| mmq::gather_native(bd, bd + f.up_off, mmq_gub / 2, bd + f.down_off, mmq_db, | |
| P.grp_gu + q * mmq_gub, P.grp_d + q * mmq_db, ps); | |
| else | |
| mmq::gather_strata_q2(bd, P.grp_gu + q * mmq_gub, P.grp_d + q * mmq_db, ps); | |
| if (psl >= 0) { // the slot is free once gathered | |
| cudaEventRecord(P.pused[(size_t) psl], ps); | |
| ++P.pk; | |
| p_issue_until(P.pk + (size_t) P.RP); | |
| } | |
| } | |
| const int64_t r0 = P.bounds_host[j0], nr = P.bounds_host[j1] - r0; | |
| if (nr <= 0) continue; | |
| const int32_t* rel = P.bounds + n + 1 + g2 * (MMQ_GROUP + 1); | |
| cudaMemsetAsync(P.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, ps); | |
| cudaMemsetAsync(P.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, ps); | |
| mmq::quantize(P.mixed, P.src + r0, P.Xq_g, mmq_gt, N, N, nr, ps); | |
| mmq::Product gu; | |
| gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub; | |
| gu.n = ngx; gu.xq = P.Xq_g; gu.bounds = rel; gu.ids = P.ident; | |
| gu.total_rows = nr; gu.max_rows = maxr; gu.dst = P.GU_g; gu.ld_dst = 1280; | |
| P.ctx->run(gu, ps); | |
| mmq::swiglu(P.GU_g, P.H_g, nr, 640, !lay.native, ps); | |
| mmq::quantize(P.H_g, nullptr, P.Hq_g, mmq_dt, 640, 640, nr, ps); | |
| const int b = (int) (g2 & 1); | |
| if (P.dm_live[b]) cudaStreamWaitEvent(ps, P.ev_dm[b], 0); // its last rows have left | |
| mmq::Product dn; | |
| dn.w = P.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db; | |
| dn.n = ngx; dn.xq = P.Hq_g; dn.bounds = rel; | |
| dn.ids = P.ident; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = P.Dm_b[b]; | |
| dn.ld_dst = N; | |
| P.ctx->run(dn, ps); | |
| cudaEventRecord(P.ev_grp[g2 % PeerPrefill::kGrpEv], ps); | |
| cudaStreamWaitEvent(P.s_out, P.ev_grp[g2 % PeerPrefill::kGrpEv], 0); | |
| cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * N, prevd, P.Dm_b[b], P.dev, | |
| (size_t) nr * N * 4, P.s_out); | |
| cudaEventRecord(P.ev_dm[b], P.s_out); | |
| P.dm_live[b] = true; | |
| } | |
| } else { | |
| const size_t ng = (n + MMQ_GROUP - 1) / MMQ_GROUP; | |
| P.bounds_host.resize(n + 1 + ng * (MMQ_GROUP + 1)); | |
| for (size_t j = 0; j < n; ++j) P.bounds_host[j] = m.off[(size_t) order_peer[j]] - (int32_t) rows_local; | |
| P.bounds_host[n] = (int32_t) rows_peer; | |
| for (size_t g2 = 0; g2 < ng; ++g2) | |
| for (size_t i = 0; i <= MMQ_GROUP; ++i) | |
| P.bounds_host[n + 1 + g2 * (MMQ_GROUP + 1) + i] = | |
| P.bounds_host[std::min(n, g2 * MMQ_GROUP + i)] - P.bounds_host[g2 * MMQ_GROUP]; | |
| cudaMemcpyAsync(P.bounds, P.bounds_host.data(), P.bounds_host.size() * 4, cudaMemcpyHostToDevice, ps); | |
| mmq::quantize(P.mixed, P.src, P.Xq, mmq_gt, N, N, rows_peer, ps); | |
| pe.mark(kPeMoeGemm, ps); | |
| const auto& f = lay.fmt[(size_t) l]; | |
| for (size_t j = 0; j < n; ++j) { | |
| const int32_t e = order_peer[j]; | |
| const uint8_t* bd = P.peer->slot_ptr(l, e); | |
| const size_t q = j % MMQ_GROUP; | |
| if (lay.native) | |
| mmq::gather_native(bd, bd + f.up_off, mmq_gub / 2, bd + f.down_off, mmq_db, | |
| P.grp_gu + q * mmq_gub, P.grp_d + q * mmq_db, ps); | |
| else | |
| mmq::gather_strata_q2(bd, P.grp_gu + q * mmq_gub, P.grp_d + q * mmq_db, ps); | |
| if (q + 1 < MMQ_GROUP && j + 1 < n) continue; | |
| const size_t j0 = j - q, g2 = j0 / MMQ_GROUP; | |
| const int ngx = (int) (q + 1); | |
| const int64_t r0 = P.bounds_host[j0], nr = P.bounds_host[j + 1] - r0; | |
| int64_t maxr = 0; | |
| for (size_t i = j0; i <= j; ++i) maxr = std::max<int64_t>(maxr, m.cnt[(size_t) order_peer[i]]); | |
| cudaMemsetAsync(P.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, ps); | |
| cudaMemsetAsync(P.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, ps); | |
| mmq::Product gu; | |
| gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub; | |
| gu.n = ngx; gu.xq = P.Xq; gu.bounds = P.bounds + j0; gu.ids = P.ident; | |
| gu.total_rows = rows_peer; gu.max_rows = maxr; gu.dst = P.GU; gu.ld_dst = 1280; | |
| P.ctx->run(gu, ps); | |
| mmq::swiglu(P.GU + r0 * 1280, P.H + r0 * 640, nr, 640, !lay.native, ps); | |
| mmq::quantize(P.H + r0 * 640, nullptr, P.Hq, mmq_dt, 640, 640, nr, ps); | |
| mmq::Product dn; | |
| dn.w = P.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db; | |
| dn.n = ngx; dn.xq = P.Hq; dn.bounds = P.bounds + n + 1 + g2 * (MMQ_GROUP + 1); | |
| dn.ids = P.ident; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = P.Dm + r0 * N; | |
| dn.ld_dst = N; | |
| P.ctx->run(dn, ps); | |
| if (P.out_pipe && nr > 0) { // this group's rows go back while the next group computes | |
| cudaEventRecord(P.ev_grp[g2], ps); | |
| cudaStreamWaitEvent(P.s_out, P.ev_grp[g2], 0); | |
| cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * N, prevd, P.Dm + (size_t) r0 * N, | |
| P.dev, (size_t) nr * N * 4, P.s_out); | |
| } | |
| } | |
| } | |
| pe.mark(kPeMoeOut, ps); | |
| if (P.out_pipe) { | |
| cudaEventRecord(P.ev_done, P.s_out); | |
| P.out_pending = true; | |
| } else { | |
| cudaMemcpyPeerAsync(m.Dm + (size_t) rows_local * N, prevd, P.Dm, P.dev, (size_t) rows_peer * N * 4, ps); | |
| cudaEventRecord(P.ev_done, ps); | |
| } | |
| pe.mark(kPeIdle, ps); | |
| cudaSetDevice(prevd); | |
| ++P.layers; | |
| P.experts += (int64_t) n; | |
| P.rows += rows_peer; | |
| } | |
| if (ps_on) { // multi-GPU: this layer's peer-ring entries the routing did not pick give their slots back | |
| int pd = 0; | |
| cudaGetDevice(&pd); | |
| cudaSetDevice(m.pp->dev); | |
| p_release_to(l, (int32_t) m.g->n_expert, m.pp->s); | |
| cudaSetDevice(pd); | |
| } | |
| // Stage ahead: the copy stream moves blobs host -> device while the compute stream works. | |
| int stage_next = 0; | |
| std::vector<int> stage_of(order.size(), -1); | |
| // the unpinned ones are copied to pinned buffers by the stager's threads, in this order | |
| std::vector<int> job_of(order.size(), -1); | |
| if (!stream_all) { | |
| std::vector<Stager::Job> js; | |
| for (size_t j = 0; j < order.size(); ++j) { | |
| const int32_t e = order[j]; | |
| if (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) continue; | |
| if (m.src->pinned(l, e)) continue; | |
| job_of[j] = (int) js.size(); | |
| if (m.src->transient(l, e)) { // CS-T: copied by the source | |
| js.push_back({nullptr, (size_t) lay.blob_bytes(l), m.src, (int32_t) l, e}); | |
| continue; | |
| } | |
| const uint8_t* b = m.src->blob(l, e); | |
| if (!b) { err = "prefill: expert source has no blob"; return false; } | |
| js.push_back({b, (size_t) lay.blob_bytes(l)}); | |
| } | |
| m.stager->start(std::move(js)); | |
| } | |
| StagerDone stager_done{stream_all ? nullptr : m.stager.get()}; | |
| auto stage_one = [&](size_t j) -> bool { | |
| const int32_t e = order[j]; | |
| const bool resident = m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0; | |
| if (resident) return true; | |
| const int sl = stage_next; | |
| stage_next = (stage_next + 1) % STAGE; | |
| const auto th = Clock::now(); | |
| const bool pinned = m.src->pinned(l, e); // pinned: never transient | |
| const uint8_t* b = pinned ? m.src->blob(l, e) : nullptr; | |
| 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[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[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); | |
| } | |
| cudaEventRecord(m.copied[sl], m.copy); | |
| m.stage_live[sl] = true; | |
| stage_of[j] = sl; | |
| stats_.ms_experts_host += ms_since(th); | |
| ++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 { | |
| const int32_t e = order[j]; | |
| pt.mark(kPfDequant, cs); | |
| if (use_mmq) { | |
| // gather the expert into its group slot (GGUF blocks, unchanged or converted) | |
| const size_t q = j % MMQ_GROUP; | |
| 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 && !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(); | |
| const int ngx = (int) (q + 1); | |
| const int64_t r0 = m.bounds_host[j0], nr = m.bounds_host[j + 1] - r0; | |
| int64_t maxr = 0; | |
| for (size_t i = j0; i <= j; ++i) maxr = std::max<int64_t>(maxr, m.cnt[(size_t) order[i]]); | |
| pt.mark(kPfGemmGU, cs); | |
| // the zeroed tail after the group's last expert (see MMQ_TAIL) | |
| cudaMemsetAsync(m.grp_gu + (size_t) ngx * mmq_gub, 0, MMQ_TAIL, m.cs); | |
| cudaMemsetAsync(m.grp_d + (size_t) ngx * mmq_db, 0, MMQ_TAIL, m.cs); | |
| mmq::Product gu; | |
| gu.w = m.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub; | |
| gu.n = ngx; gu.xq = m.Xq; gu.bounds = m.bounds_dev + j0; gu.ids = m.ids_identity; | |
| gu.total_rows = T * K; gu.max_rows = maxr; gu.dst = m.GU; gu.ld_dst = 1280; | |
| m.mmq_ctx->run(gu, m.cs); | |
| mmq::swiglu(m.GU + r0 * 1280, m.H + r0 * 640, nr, 640, !lay.native, m.cs); | |
| pt.mark(kPfGemmD, cs); | |
| mmq::quantize(m.H + r0 * 640, nullptr, m.Hq, mmq_dt, 640, 640, nr, m.cs); | |
| mmq::Product dn; | |
| dn.w = m.grp_d; dn.type = mmq_dt; dn.w_rows = N; dn.w_cols = 640; dn.expert_bytes = mmq_db; | |
| dn.n = ngx; dn.xq = m.Hq; dn.bounds = m.bounds_dev + n + 1 + g * (MMQ_GROUP + 1); | |
| dn.ids = m.ids_identity; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = m.Dm + r0 * N; | |
| dn.ld_dst = N; | |
| m.mmq_ctx->run(dn, m.cs); | |
| return true; | |
| } | |
| const int q = (int) (j % DQ); | |
| if (lay.native) { | |
| // plan v0.3 P6: a native pack's layer, dequantized by llama.cpp's own formulas | |
| const auto& f = lay.fmt[(size_t) l]; | |
| strata::kernels::iq_dequant_gu_f16(f.gu_type, blob_dev, blob_dev + f.up_off, f.n_ff, f.n_embd, | |
| m.dq_gu[q], m.cs); | |
| strata::kernels::iq_dequant_f16(f.d_type, blob_dev + f.down_off, f.n_embd * f.n_ff, m.dq_d[q], m.cs); | |
| } 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); 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); | |
| swiglu_interleaved(m.GU + o0 * 1280, m.Hh + o0 * 640, ne, m.cs); | |
| pt.mark(kPfGemmD, cs); | |
| m.gemm.f16(m.Hh + o0 * 640, m.dq_d[q], m.Dm + o0 * N, ne, N, 640); | |
| return true; | |
| }; | |
| if (!stream_all) { | |
| size_t staged = 0; | |
| const size_t lookahead = STAGE - 1; | |
| for (size_t j = 0; j < order.size(); ++j) { | |
| while (staged < order.size() && staged <= j + lookahead) { | |
| if (!stage_one(staged)) return false; | |
| ++staged; | |
| } | |
| const int32_t e = order[j]; | |
| if (stage_of[j] < 0) { | |
| ++stats_.experts_resident; | |
| if (!compute(j, m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]), -1)) return false; | |
| } else { | |
| pt.mark(kPfWaitCopy, cs); | |
| cudaStreamWaitEvent(m.cs, m.copied[stage_of[j]], 0); | |
| if (!compute(j, m.stage_dev[stage_of[j]], stage_of[j])) return false; | |
| } | |
| } | |
| } else { | |
| // the streamed walk: this layer's entries [k, kend) in id order; an entry the routing did not | |
| // pick only gives its slot back | |
| size_t k = seq_start[(size_t) l]; | |
| 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) { | |
| 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); | |
| } | |
| }; | |
| for (size_t j = 0; j < order.size(); ++j) { | |
| const int32_t e = order[j]; | |
| release_to(e); | |
| if (k < kend && seq[k].e == e) { | |
| const int sl = (int) (k % (size_t) m.ring); | |
| wait_issued(k); | |
| 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; | |
| if (!group_gather || gg_nslots == 0) give_back(consumed); // its group was gathered | |
| } else { | |
| const bool r0 = m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0; | |
| const uint8_t* bp = r0 ? m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e]) | |
| : (m.pp ? m.pp->peer->slot_ptr(l, e) : nullptr); // over the cap: P2P | |
| if (bp == nullptr) { err = "prefill: an expert is neither resident, streamed nor on the peer"; return false; } | |
| ++stats_.experts_resident; | |
| if (!compute(j, bp, -1)) return false; | |
| if (group_gather && gg_nslots == 0) give_back(consumed); | |
| } | |
| } | |
| release_to(m.g->n_expert); | |
| } | |
| } | |
| pt.mark(kPfCombine, cs); | |
| if (peer_now) cudaStreamWaitEvent(m.cs, m.pp->ev_done, 0); // multi-GPU: the peer's rows are in Dm | |
| moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs); | |
| // debug: STRATA_DBG_NAN=1 reports the first layer of a chunk whose MoE produced non-finite values | |
| if (static const bool dbg = std::getenv("STRATA_DBG_NAN") != nullptr; dbg) { | |
| cudaStreamSynchronize(m.cs); | |
| auto bad = [&](const float* d, int64_t n) { | |
| std::vector<float> h((size_t) n); | |
| cudaMemcpy(h.data(), d, (size_t) n * 4, cudaMemcpyDeviceToHost); | |
| int64_t c = 0; | |
| for (float v : h) c += !std::isfinite(v); | |
| return c; | |
| }; | |
| // (fused: GU and H hold the grouping tables and int8 rows, not floats) | |
| const int64_t bgu = fused_l ? 0 : bad(m.GU, T * K * 1280), bdm = bad(m.Dm, T * K * N), | |
| bbo = bad(m.bo, T * N); | |
| const int64_t bh = m.H && !fused_l ? bad(m.H, T * K * 640) : -1; | |
| static int64_t reported = -1; | |
| if ((bgu || bdm || bbo || bh > 0) && reported != stats_.chunks) { | |
| reported = stats_.chunks; | |
| std::fprintf(stderr, "strata dbg: layer %lld (mmq %d, types %d/%d, %zu experts): non-finite GU %lld " | |
| "H %lld Dm %lld bo %lld of T %lld\n", (long long) l, (int) use_mmq, mmq_gt, mmq_dt, | |
| n_order, (long long) bgu, (long long) bh, (long long) bdm, (long long) bbo, | |
| (long long) T); | |
| } | |
| } | |
| } | |
| // ---- the hyper-connection write of this half; F-2: fused with the next half's norm when nothing else | |
| // touches R in between (not the stage's last half, not before the PLE block of layer 1, not under a | |
| // control vector) | |
| const int64_t nl = half == 0 ? l : l + 1; | |
| const bool fuse = !gr_unfused() && nl < LE && !(half == 1 && nl == 1 && ple_on) && | |
| !(half == 1 && strata::kernels::cvec().covers(l)); | |
| const core::WeightRef* wnn = nullptr; | |
| if (fuse) { | |
| const core::LayerView vn(*m.wt, nl); | |
| wnn = need(vn, half == 0 ? "hc_ffn_norm.weight" : "hc_attn_norm.weight", err); | |
| if (!wnn) return false; | |
| } | |
| if (wnn) { | |
| gr_write_norm_rs(m.R, m.bo, m.inj, HC, (const float*) wnn->data, EPS, m.grs, m.xn16, T, m.cs, | |
| m.xn16_lo); | |
| normed = true; | |
| } else { | |
| gr_write(m.R, m.bo, m.inj, HC, T, m.cs); | |
| } | |
| if (half == 1 && strata::kernels::cvec().covers(l)) // --control-vector-scaled | |
| strata::kernels::cvec_apply(m.R, l, T, D, nullptr, 0, nullptr, 0, false, m.cs); | |
| } | |
| } | |
| if (!ple_land()) return false; // a stage that ends before layer 1: the rows land anyway, the next gather starts | |
| if (issuer.joinable()) { | |
| issuer.join(); | |
| stats_.ms_experts_host += iss_ms; | |
| stats_.experts_streamed += iss_streamed; | |
| stats_.experts_dma += iss_dma; | |
| } | |
| stats_.tokens += T; | |
| core::progress_at("reading the prompt (batched): finishing the chunk from token", p0); | |
| pt.mark(kPfStart, cs); | |
| if (next_ != nullptr) { | |
| // the rows to the host buffer the next stage read two chunks ago (it has finished: waited below) | |
| float* h = m.hand[hand_buf]; | |
| if (cudaMemcpyAsync(h, m.R, (size_t) T * D * 4, cudaMemcpyDeviceToHost, m.cs) != cudaSuccess || | |
| cudaStreamSynchronize(m.cs) != cudaSuccess) { | |
| err = std::string("prefill: the layer split's hand-off: ") + cudaGetErrorString(cudaGetLastError()); | |
| return false; | |
| } | |
| // this stage's state is at the chunk's end now (synced) and moves on with the next chunk below | |
| if (on_stage_chunk && !on_stage_chunk(p0 + T, err)) return false; | |
| if (next_run.valid() && !next_run.get()) { err = next_err; return false; } | |
| next_->hand_in_ = h; | |
| next_run = std::async(std::launch::async, [this, tokens, c0, T, p0, &next_err] { | |
| return next_->run(tokens + c0, T, p0, next_err); | |
| }); | |
| hand_buf ^= 1; | |
| continue; // the last stage reports the chunk (on_chunk) | |
| } | |
| if (const char* dump = std::getenv("STRATA_PREFILL_DUMP_R")) { // debug: the final residuals, every 64th | |
| cudaStreamSynchronize(m.cs); // position (A/B quality of this path) | |
| if (std::FILE* f = std::fopen(dump, c0 == 0 ? "wb" : "ab")) { | |
| std::vector<float> row((size_t) D); | |
| for (int64_t t = (64 - p0 % 64) % 64; t < T; t += 64) { | |
| cudaMemcpy(row.data(), m.R + t * D, (size_t) D * 4, cudaMemcpyDeviceToHost); | |
| const int64_t pos = p0 + t; | |
| std::fwrite(&pos, sizeof pos, 1, f); | |
| std::fwrite(row.data(), 4, row.size(), f); | |
| } | |
| std::fclose(f); | |
| } | |
| } | |
| if (on_chunk || on_stage_chunk) { | |
| const auto toc = Clock::now(); | |
| if (cudaStreamSynchronize(m.cs) != cudaSuccess) { | |
| err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); | |
| return false; | |
| } | |
| const auto toc2 = Clock::now(); | |
| if (on_stage_chunk && !on_stage_chunk(p0 + T, err)) return false; | |
| if (on_chunk && !on_chunk(m.R, T, p0, err)) return false; | |
| host_sync_ms += std::chrono::duration<double, std::milli>(toc2 - toc).count(); | |
| host_chunk_ms += ms_since(toc2); | |
| } | |
| } | |
| if (next_run.valid() && !next_run.get()) { err = next_err; return false; } | |
| ss.ple_prev[0] = prev[0]; | |
| ss.ple_prev[1] = prev[1]; | |
| if (std::getenv("STRATA_DBG_NAN") != nullptr) { // debug: the state the prompt leaves for the token path | |
| cudaStreamSynchronize(m.cs); | |
| auto bad = [&](const float* d, int64_t n) { | |
| std::vector<float> h((size_t) n); | |
| cudaMemcpy(h.data(), d, (size_t) n * 4, cudaMemcpyDeviceToHost); | |
| int64_t c = 0; | |
| double mx = 0; | |
| for (float v : h) { c += !std::isfinite(v); if (std::isfinite(v)) mx = std::max(mx, (double) std::fabs(v)); } | |
| std::fprintf(stderr, " %lld non-finite (max |x| %.3g)", (long long) c, mx); | |
| }; | |
| const int64_t last = (n - 1) % m.T; | |
| std::fprintf(stderr, "strata dbg: prompt end: last residual row"); | |
| bad(m.R + last * D, D); | |
| if (ss.ple.ready()) { std::fprintf(stderr, "; PLE history"); bad(ss.ple.hist, (int64_t) strata::kernels::NG_HIST * strata::kernels::NG_HC_DIM); } | |
| std::fprintf(stderr, "; GDN state 0"); | |
| bad(ss.gdn_state, 64 * 1024); | |
| std::fprintf(stderr, "\n"); | |
| } | |
| if (cudaStreamSynchronize(m.cs) != cudaSuccess) { | |
| err = std::string("prefill: ") + cudaGetErrorString(cudaGetLastError()); | |
| return false; | |
| } | |
| // (PR #121) an expert copy that failed on the copy stream surfaces here, not in the next request | |
| if (const cudaError_t cst = cudaStreamSynchronize(m.copy); cst != cudaSuccess) { | |
| err = std::string("prefill: expert copy stream: ") + cudaGetErrorString(cst); | |
| return false; | |
| } | |
| stats_.ms_total += ms_since(t_start); | |
| if (pt.on) { | |
| pt.fold(); | |
| double total = 0.0; | |
| for (double v : pt.ms) total += v; | |
| std::string line; | |
| char b[96]; | |
| for (int i = 0; i < kPfCount; ++i) { | |
| if (pt.ms[i] <= 0.0) continue; | |
| std::snprintf(b, sizeof b, " %s %.0f (%.1f%%)", kPfNames[i], pt.ms[i], total > 0 ? 100.0 * pt.ms[i] / total : 0.0); | |
| line += b; | |
| } | |
| std::fprintf(stderr, "strata prefill timing: %lld tokens, GPU timeline %.0f ms, wall %.0f ms, host staging %.0f ms:%s\n", | |
| (long long) n, total, ms_since(t_start), stats_.ms_experts_host, line.c_str()); | |
| std::fprintf(stderr, "strata prefill timing: host: chunk setup (PLE rows, the expert stream plan) %.0f ms, " | |
| "waiting for each chunk %.0f ms, after each chunk (the draft layer, progress) %.0f ms, " | |
| "PLE %.0f ms\n", host_setup_ms, host_sync_ms, host_chunk_ms, stats_.ms_ple); | |
| if (pe.on) { | |
| int pd = 0; cudaGetDevice(&pd); cudaSetDevice(pe.dev); cudaStreamSynchronize(m.pp->s); pe.fold(); cudaSetDevice(pd); | |
| std::string pl; | |
| for (int i = 0; i < kPeCount; ++i) { | |
| std::snprintf(b, sizeof b, " %s %.0f", kPeNames[i], pe.ms[i]); | |
| pl += b; | |
| } | |
| std::fprintf(stderr, "strata prefill timing (peer GPU, ms):%s; MoE layers %lld, rows/layer %.0f of %lld cap, experts/layer %.0f, " | |
| "over the cap %lld\n", pl.c_str(), (long long) m.pp->layers, | |
| m.pp->layers ? (double) m.pp->rows / m.pp->layers : 0.0, (long long) m.pp->cap_rows, | |
| m.pp->layers ? (double) m.pp->experts / m.pp->layers : 0.0, (long long) m.pp->over_cap); | |
| if (m.pp->ps_frac > 0.0) | |
| std::fprintf(stderr, "strata prefill timing: peer-streamed experts %lld (%.0f per MoE layer)\n", | |
| (long long) m.pp->ps_experts, m.pp->layers ? (double) m.pp->ps_experts / m.pp->layers : 0.0); | |
| } | |
| } | |
| if (std::getenv("STRATA_STATE_HASH_GDN") != nullptr) { // debug: the GDN states as the prompt path leaves them | |
| cudaStreamSynchronize(m.cs); | |
| std::vector<uint8_t> b((size_t) gdn_floats * 4); | |
| std::string line; | |
| char h[8]; | |
| for (int64_t i = 0; i < ss.gdn_alloc; ++i) { | |
| cudaMemcpy(b.data(), ss.gdn_state + (size_t) i * gdn_floats, b.size(), cudaMemcpyDeviceToHost); | |
| uint64_t x = 1469598103934665603ull; | |
| for (uint8_t c : b) x = (x ^ c) * 1099511628211ull; | |
| std::snprintf(h, sizeof(h), "%04llx ", (unsigned long long) (x & 0xffff)); | |
| line += h; | |
| } | |
| std::fprintf(stderr, "strata prefill: GDN_HASH %s\n", line.c_str()); | |
| } | |
| return true; | |
| } | |
| } // namespace strata::prefill | |