Prefill: STRATA_DEQUANT_SIDE=1 dequantizes the dense weights on a side stream into a double-buffered scratch

Gemm::native ran the dequant serially in front of every dense GEMM on the compute stream. With the env the dequant
of the next projection runs on a side stream into the other scratch (64 MiB more of the prompt path's region), with
ready/free events; the sliced path stays serial. Same kernels, same bits. Off on HIP.

- the side stream's event waits / records and ~Gemm's destroys are (void) (HIP's [[nodiscard]] hipError_t)
- metal: gemm_dequant_side() answers false and Gemm::set_dequant_side(nullptr) the serial form (stubs_0145.cpp)

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>

include/strata/prefill/gemm.hpp

@@ -66,6 +66,13 @@ public:
     /// Caller-owned buffers only: the scratch and workspace moved (the prompt path laid its buffers out again).
     void rebind(uint16_t* scratch, int64_t scratch_elems, void* workspace, size_t ws_bytes);
 
+    /// STRATA_DEQUANT_SIDE=1 (gemm_dequant_side): a second dequant scratch of scratch_elems() elements, caller-owned
+    /// like the first.  native() then dequantizes on a side stream into the two in turn, so a weight's dequant runs
+    /// while the kernels queued ahead of its GEMM do; events order every write and read of each scratch, so the
+    /// products are the serial form's, bit for bit.  nullptr: the serial form.  Called after init_external and
+    /// after every rebind (the "scratch free" events are re-armed on the compute stream).  CUDA only.
+    bool set_dequant_side(uint16_t* scratch2, std::string& err);
+
     uint16_t* scratch() const { return scratch_; }
     int64_t scratch_elems() const { return scratch_elems_; }
     void* stream() const { return stream_; }
@@ -77,6 +84,14 @@ private:
     int64_t scratch_elems_ = 0;
     void* workspace_ = nullptr;
     bool external_ = false;
+    // STRATA_DEQUANT_SIDE (set_dequant_side): the second scratch, the side stream and, per scratch, "the dequant
+    // landed" (recorded on the side stream) and "the GEMM read it" (recorded on the compute stream); side_cur_ is
+    // the scratch the next native() dequantizes into (0: scratch_, 1: side_scratch_)
+    uint16_t* side_scratch_ = nullptr;
+    void* side_stream_ = nullptr;
+    void* side_ready_[2] = {};
+    void* side_free_[2] = {};
+    int side_cur_ = 0;
     void* hipblaslt_state_ = nullptr;
     void* rocblas_state_ = nullptr;   // HIP: the rocBLAS solution table of f16_inplace (STRATA_ROCBLAS_TUNING)
     bool f16_io_ = false;
@@ -122,6 +137,11 @@ void gemm_set_dense_f32(int mode);
 /// set_act_f16).  Cached per device (a layer split can mix cards).  STRATA_HIP_PROMPT_F16=0/1 overrides.
 bool prompt_f16();
 
+/// STRATA_DEQUANT_SIDE=1: the prompt path's dense-weight dequant on a side stream with a second GEMM scratch
+/// (Gemm::set_dequant_side; the prompt path takes the scratch from its borrowed region when this is on).  Off by
+/// default; always off on HIP (its padded / MMQ paths use the scratch on the compute stream without the events).
+bool gemm_dequant_side();
+
 
 
 }  // namespace strata::prefill

metal/src/engine/stubs_0145.cpp

@@ -151,3 +151,6 @@ namespace strata::kernels { void native_q8_0_sm60_release(void const* p0) { } }
 namespace strata::prefill { void swiglu_split_f16(float const* p0, unsigned short* p1, long long p2, int p3, bool p4, void* p5) { } }
 // STRATA_PLE_COMPACT (opt-in, CUDA / HIP prompt path): prefill.cpp references it; the Metal prompt runs through its window and never reaches it.
 namespace strata::kernels { void ple_iq4nl_dequant(const uint8_t* p0, int64_t p1, float* p2, void* p3) { strata::metal_engine::missing("strata::kernels::ple_iq4nl_dequant"); } }
+// STRATA_DEQUANT_SIDE (opt-in, CUDA only): off on Metal; the prompt path takes no second scratch and set_dequant_side(nullptr) keeps the serial form.
+namespace strata::prefill { bool gemm_dequant_side() { return false; } }
+auto strata::prefill::Gemm::set_dequant_side(uint16_t* p0, std::string& p1) -> bool { if (p0 == nullptr) return true; strata::metal_engine::missing("strata::prefill::Gemm::set_dequant_side"); return false; }

src/prefill/gemm.cu

@@ -582,7 +582,21 @@ const uint16_t* Gemm::unpad_x(const uint16_t* X, int64_t T, int64_t K, int64_t l
     return unpadded_x_;
 }
 
+bool gemm_dequant_side() {
+#if defined(__HIPCC__)
+    return false;
+#else
+    static const bool on = [] { const char* v = std::getenv("STRATA_DEQUANT_SIDE"); return v && v[0] == '1'; }();
+    return on;
+#endif
+}
+
 Gemm::~Gemm() {
+    for (void* e : side_ready_)
+        if (e) (void) cudaEventDestroy((cudaEvent_t) e);
+    for (void* e : side_free_)
+        if (e) (void) cudaEventDestroy((cudaEvent_t) e);
+    if (side_stream_) (void) cudaStreamDestroy((cudaStream_t) side_stream_);
 #if defined(STRATA_PREFILL_MMQ) && defined(__HIPCC__)
     delete static_cast<strata::prefill::mmq::Context*>(mmq_ctx_);
     if (mmq_buf_) cudaFree(mmq_buf_);
@@ -642,6 +656,38 @@ void Gemm::rebind(uint16_t* scratch, int64_t scratch_elems, void* workspace, siz
 #endif
 }
 
+bool Gemm::set_dequant_side(uint16_t* scratch2, std::string& err) {
+    side_scratch_ = scratch2;
+    side_cur_ = 0;
+    if (scratch2 == nullptr) return true;
+    if (side_stream_ == nullptr) {   // once: the stream and the four events outlive every rebind
+        cudaStream_t s = nullptr;
+        if (const cudaError_t e = cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking); e != cudaSuccess) {
+            err = std::string("prefill gemm: dequant side stream: ") + cudaGetErrorString(e);
+            side_scratch_ = nullptr;
+            return false;
+        }
+        side_stream_ = s;
+        for (int i = 0; i < 2; ++i) {
+            cudaEvent_t r = nullptr, f = nullptr;
+            const bool made = cudaEventCreateWithFlags(&r, cudaEventDisableTiming) == cudaSuccess &&
+                              cudaEventCreateWithFlags(&f, cudaEventDisableTiming) == cudaSuccess;
+            side_ready_[i] = r;   // kept even on a failure: the destructor frees what was made
+            side_free_[i] = f;
+            if (!made) {
+                err = "prefill gemm: dequant side events";
+                side_scratch_ = nullptr;
+                return false;
+            }
+        }
+    }
+    // both scratches are free once what the compute stream holds now is done (a rebind follows a stream sync, so
+    // this is belt and braces; a wait on an event never recorded would pass anyway)
+    for (int i = 0; i < 2; ++i)
+        (void) cudaEventRecord((cudaEvent_t) side_free_[i], (cudaStream_t) stream_);
+    return true;
+}
+
 bool Gemm::init(void* stream, int64_t scratch_elems, std::string& err) {
     // #240: every failure names the call and the real status, so "no VRAM" can be told from a broken install
     cublasHandle_t h = nullptr;
@@ -1311,6 +1357,24 @@ void Gemm::native(const uint16_t* X, int ggml_type, const void* W_blocks, float*
             strata::kernels::dequant_f16(ggml_type, W_blocks, r0, n, K, scratch_, stream_);
             f16(X, scratch_, Y + r0, T, n, K, ldy, beta);
         }
+        // the side stream's next dequant into scratch_ waits for these slices' reads and writes of it
+        if (side_scratch_) (void) cudaEventRecord((cudaEvent_t) side_free_[0], (cudaStream_t) stream_);
+        return;
+    }
+    if (side_scratch_) {
+        // STRATA_DEQUANT_SIDE: the dequant on the side stream, into the scratch the compute stream read last but one.
+        // free[cur] -> dequant -> ready[cur] -> GEMM -> free[cur] orders the dequant after the GEMM that last read
+        // this scratch and the GEMM after the dequant that filled it, so every read and write of a scratch is where
+        // the serial form had it - the same kernels, the same bits - and the dequant runs while the kernels queued
+        // ahead of this GEMM do.
+        uint16_t* const s = side_cur_ == 0 ? scratch_ : side_scratch_;
+        (void) cudaStreamWaitEvent((cudaStream_t) side_stream_, (cudaEvent_t) side_free_[side_cur_], 0);
+        strata::kernels::dequant_f16(ggml_type, W_blocks, 0, N, K, s, side_stream_);
+        (void) cudaEventRecord((cudaEvent_t) side_ready_[side_cur_], (cudaStream_t) side_stream_);
+        (void) cudaStreamWaitEvent((cudaStream_t) stream_, (cudaEvent_t) side_ready_[side_cur_], 0);
+        f16(X, s, Y, T, N, K, ldy, beta);
+        (void) cudaEventRecord((cudaEvent_t) side_free_[side_cur_], (cudaStream_t) stream_);
+        side_cur_ ^= 1;
         return;
     }
     strata::kernels::dequant_f16(ggml_type, W_blocks, 0, N, K, scratch_, stream_);

src/prefill/prefill.cpp

@@ -1561,8 +1561,11 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co
     {
         uint16_t* gs = o.take<uint16_t>((size_t) GEMM_SCRATCH, ok);
         void* ws = o.take<uint8_t>(GEMM_WS, ok);
+        // STRATA_DEQUANT_SIDE: the second dequant scratch (counted by bytes_needed_impl, retaken by relayout)
+        uint16_t* gs2 = gemm_dequant_side() ? o.take<uint16_t>((size_t) GEMM_SCRATCH, ok) : nullptr;
         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 (!m.gemm.set_dequant_side(gs2, err)) return false;
         // this stage's device decides, once, here (init runs on the stage's own device): its gr_* kernels write the
         // image its GEMMs read, and set_act_f16 writes that device's symbol
         const core::OnDevice on_stage(m.device);
@@ -1718,7 +1721,9 @@ bool Prefill::relayout(int64_t chunk, void* borrow, uint64_t borrow_bytes, std::
     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);
+    uint16_t* gs2 = gemm_dequant_side() ? o.take<uint16_t>((size_t) GEMM_SCRATCH, ok) : nullptr;   // STRATA_DEQUANT_SIDE
     if (ok) m.gemm.rebind(gs, GEMM_SCRATCH, ws, GEMM_WS);
+    if (ok && !m.gemm.set_dequant_side(gs2, err)) return false;
     if (!ok || !carve((size_t) chunk, &o)) {
         size_t fb = 0, tb = 0;
         cudaMemGetInfo(&fb, &tb);
@@ -2367,6 +2372,7 @@ uint64_t Prefill::bytes_needed_impl(const core::ModelGeometry& g, const core::Se
     if (owned_pages) o.granule = 2ull << 20;
     o.take<uint16_t>((size_t) GEMM_SCRATCH, ok);
     o.take<uint8_t>(GEMM_WS, ok);
+    if (gemm_dequant_side()) o.take<uint16_t>((size_t) GEMM_SCRATCH, ok);   // STRATA_DEQUANT_SIDE: the second scratch
     auto f = [&](size_t n) { o.take<float>(n, ok); };
     // `carve`'s order, buffer for buffer: emb, R, xn, grs, xn16, lo, lo16, gated, inj, mixed, mixed_bf, mixed_h,
     // bo aliases emb. This counted `xn` unconditionally (carve takes it only under STRATA_GR_UNFUSED) and never counted