The draft pass's kernels in the main layer's forms

- quantize_q8_0_scaled runs a warp a 32-value block (it ran a thread a block: 80 threads for the draft layer's input):
  the block's |max| over the same values, then each value's bytes as before.  6.2 -> 2.4 us and 9.2 -> 1.6 us for the
  Q2_0 draft experts' two quantizations (CUPTI).
- The draft graphs' inputs from mapped memory in one copy kernel (copy_from_mapped_multi): 7 launches in the round
  graph, 2 in a step, became one each.
- The router in one kernel (verify_router: the logits, the top 10, and the native draft experts' q8_1 rows), bitwise
  the gemv, top 10 and quantizer it replaces; as fast as the gemv and top 10 were.
- The shared expert on a branch of the draft graphs beside the router and the routed experts, joined before the combine.

Window hashes on fixed text equal the previous build's (IQ3_S, UD-Q4_K_XL; thinking and code text; IQ3_S with the
IQ4_NL draft experts too).  CUPTI, IQ3_S thinking text (--adapt-every 0): the round graph 450 / 661 / 687 / 693 / 824 ->
411 / 592 / 620 / 627 / 760 us after windows of 1 / 2 / 3 / 4 / 8 tokens, each step 435 -> 405-410 us.

Claude-Session: https://claude.ai/code/session_019XxRxQQNyk7Ycdr4ComSs1

include/strata/core/mtp.hpp

@@ -115,6 +115,10 @@ private:
     int64_t n_vocab_ = 0;
     uint64_t vram_ = 0;
     cudaStream_t cs_ = nullptr;
+    cudaStream_t side_ = nullptr;                    // the shared expert's branch of the draft graphs
+    cudaEvent_t sh_fork_ = nullptr, sh_join_ = nullptr;
+    uint32_t* rseq_ = nullptr;                       // verify_router's sequence word and counter (device)
+    unsigned* rcount_ = nullptr;
     cudaGraphExec_t round_exec_[9] = {};
 
     struct Tensor { std::string name, kind; int64_t rows = 0, cols = 0; uint64_t off = 0, bytes = 0; };

include/strata/kernels/elementwise.hpp

@@ -116,5 +116,13 @@ void doorbell_publish(const float* x, const int32_t* ids, const float* weights,
 /// Plan v0.3 P3: copy `n` int32 from mapped pinned host memory into device memory with a kernel (the QSA
 /// per-token step and positions), instead of a host-to-device memcpy node in the middle of a layer.
 void copy_i32_from_mapped(int32_t* dst, const int32_t* src, int64_t n, void* stream);
+/// Up to 8 copies from mapped host memory in one launch: `words` 4-byte words each (16 bytes a load where both
+/// ends are 16-byte aligned and the count a multiple of 4).  Graph-capturable.
+struct MappedCopy {
+    void* dst = nullptr;
+    const void* src = nullptr;
+    int64_t words = 0;
+};
+void copy_from_mapped_multi(const MappedCopy* copies, int n, void* stream);
 
 }  // namespace strata::kernels

src/core/mtp.cpp

@@ -27,6 +27,7 @@
 #include <atomic>
 #include <chrono>
 #include <cstdio>
+#include <cstdlib>
 #include <cstring>
 #include <exception>
 #include <fstream>
@@ -83,6 +84,9 @@ MtpDrafter::~MtpDrafter() {
     for (auto& e : round_exec_) if (e) cudaGraphExecDestroy(e);
     for (auto& e : step_exec_) if (e) cudaGraphExecDestroy(e);
     if (cs_) cudaStreamDestroy(cs_);
+    if (side_) cudaStreamDestroy(side_);
+    if (sh_fork_) cudaEventDestroy(sh_fork_);
+    if (sh_join_) cudaEventDestroy(sh_join_);
     if (dense_) cudaFree(dense_);
     if (experts_) cudaFree(experts_);
     if (state_arena_) cudaFree(state_arena_);
@@ -213,6 +217,8 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session
         grp_counts_ = b.take<int32_t>(4);
         hit_xq_ = b.take<uint8_t>(T * (N / 32) * 34); hit_xs_ = b.take<float>(T * (N / 32));
         nat_xq_ = b.take<uint8_t>(T * (N / 32) * 36);
+        rseq_ = b.take<uint32_t>(4);
+        rcount_ = b.take<unsigned>(4);
         hit_scratch_ = b.take<uint8_t>(std::max<uint64_t>(
             strata::kernels::moe_hit_grouped_scratch_bytes((int64_t) (T * K), g.n_embd, g.n_ff),
             strata::kernels::native_expert_scratch_bytes((int64_t) (T * K), g.n_ff)));
@@ -238,7 +244,13 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session
             for (int64_t i = 0; i < cap_; ++i) id[(size_t) (t * (uint64_t) cap_ + (uint64_t) i)] = (int32_t) i;
         cudaMemcpy(ident_, id.data(), id.size() * 4, cudaMemcpyHostToDevice);
     }
-    if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess) { err = "mtp: stream"; return false; }
+    if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaStreamCreateWithFlags(&side_, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaEventCreateWithFlags(&sh_fork_, cudaEventDisableTiming) != cudaSuccess ||
+        cudaEventCreateWithFlags(&sh_join_, cudaEventDisableTiming) != cudaSuccess) {
+        err = "mtp: streams";
+        return false;
+    }
     char fmts[48] = "Q2_0 blobs";
     if (nat_.gu_type >= 0) std::snprintf(fmts, sizeof fmts, "ggml types %d/%d", nat_.gu_type, nat_.d_type);
     std::fprintf(stderr, "strata mtp: draft layer loaded, %.0f MiB of VRAM (experts %.0f as %s, dense %.0f)\n",
@@ -398,16 +410,45 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) {
             }
             fused_gr_read_multi(fa, T, grs_, cs);
         }
-        // ---- MoE: router, the 512 resident experts, the shared expert, the combine, the write
+        // ---- MoE: the shared expert on a branch beside the router and the 512 resident experts, then the combine
+        // and the write
+        NativeSharedWeights nsw;
+        nsw.gate_type = GGML_Q8_0; nsw.gate_data = q8("mlp.shared_expert.gate_proj.weight");
+        nsw.up_type = GGML_Q8_0; nsw.up_data = q8("mlp.shared_expert.up_proj.weight");
+        nsw.down_type = GGML_Q8_0; nsw.down_data = q8("mlp.shared_expert.down_proj.weight");
+        nsw.q8_1 = xq_;   // the routed experts read nat_xq_ / hit_xq_
+        const SForm none{};
+        if (cudaEventRecord(sh_fork_, cs) != cudaSuccess || cudaStreamWaitEvent(side_, sh_fork_, 0) != cudaSuccess) {
+            err = "mtp: the shared expert's branch";
+            return false;
+        }
         for (int t = 0; t < T; ++t) {
-            bf16_gemv_fp32_mmvf(mixed_ + t * N, bf16("mlp.gate.weight"), logits_ + t * g.n_expert, (int) N, (int) g.n_expert, cs);
-            if (native_router_enabled()) native_router_top10(logits_ + t * g.n_expert, ids_ + t * K, w_ + t * K, cs);
-            else router_top10(logits_ + t * g.n_expert, 1, (int) g.n_expert, (int) K, ids_ + t * K, w_ + t * K, cs);
+            f32_to_bf16_bulk(mixed_ + t * N, x_bf16_, N, side_);
+            shared_expert(nullptr, nullptr, x_bf16_, none, nullptr, nullptr, nullptr, none, nullptr, nullptr, nullptr, none,
+                          nullptr, nullptr, nullptr, bf16("mlp.shared_expert_gate.weight"), sh_scratch_, shared_ + t * N,
+                          N, g.n_ff, 32, side_, mixed_ + t * N, &nsw);
+        }
+        if (cudaEventRecord(sh_join_, side_) != cudaSuccess) { err = "mtp: the shared expert's branch"; return false; }
+        bool xq1 = false;   // the router kernel wrote the native experts' q8_1 rows
+        if (native_router_enabled()) {
+            VerifyRouterArgs ra;
+            ra.x = mixed_; ra.w = bf16("mlp.gate.weight"); ra.logits = logits_; ra.ids = ids_; ra.weights = w_;
+            ra.xq1 = nat_.gu_type >= 0 ? nat_xq_ : nullptr;
+            ra.seq = rseq_; ra.counter = rcount_;
+            ra.n_tok = T; ra.n_embd = (int) N; ra.n_expert = (int) g.n_expert;
+            verify_router(ra, cs);
+            xq1 = ra.xq1 != nullptr;
+        } else {
+            for (int t = 0; t < T; ++t) {
+                bf16_gemv_fp32_mmvf(mixed_ + t * N, bf16("mlp.gate.weight"), logits_ + t * g.n_expert, (int) N,
+                                    (int) g.n_expert, cs);
+                router_top10(logits_ + t * g.n_expert, 1, (int) g.n_expert, (int) K, ids_ + t * K, w_ + t * K, cs);
+            }
         }
         moe_group_resident(ids_, (int) (T * K), (int) K, experts_, (int64_t) blob_, grp_ptr_, grp_start_, grp_counts_,
                            hit_dst_, hit_slot_, cs);
         if (nat_.gu_type >= 0) {
-            quantize_q8_1_rows(mixed_, T, N, nat_xq_, cs);
+            if (!xq1) quantize_q8_1_rows(mixed_, T, N, nat_xq_, cs);
             native_expert_grouped(nat_, grp_ptr_, grp_start_, grp_counts_, hit_dst_, hit_slot_, (int64_t) T * K,
                                   (int64_t) T * K, nat_xq_, hit_scratch_, parts_, cs);
         } else {
@@ -415,17 +456,8 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) {
             moe_grouped_s2(grp_ptr_, grp_start_, grp_counts_, hit_dst_, hit_slot_, (int64_t) T * K, (int64_t) T * K,
                            hit_xq_, hit_xs_, hit_scratch_, parts_, cs);
         }
-        NativeSharedWeights nsw;
-        nsw.gate_type = GGML_Q8_0; nsw.gate_data = q8("mlp.shared_expert.gate_proj.weight");
-        nsw.up_type = GGML_Q8_0; nsw.up_data = q8("mlp.shared_expert.up_proj.weight");
-        nsw.down_type = GGML_Q8_0; nsw.down_data = q8("mlp.shared_expert.down_proj.weight");
-        nsw.q8_1 = xq_;
-        const SForm none{};
+        if (cudaStreamWaitEvent(cs, sh_join_, 0) != cudaSuccess) { err = "mtp: the shared expert's join"; return false; }
         for (int t = 0; t < T; ++t) {
-            f32_to_bf16_bulk(mixed_ + t * N, x_bf16_, N, cs);
-            shared_expert(nullptr, nullptr, x_bf16_, none, nullptr, nullptr, nullptr, none, nullptr, nullptr, nullptr, none,
-                          nullptr, nullptr, nullptr, bf16("mlp.shared_expert_gate.weight"), sh_scratch_, shared_ + t * N,
-                          N, g.n_ff, 32, cs, mixed_ + t * N, &nsw);
             if (native_moe_combine_enabled())
                 native_moe_combine(parts_ + (size_t) t * K * N, w_ + t * K, shared_ + t * N, y_ + t * N, N, K, cs);
             else
@@ -479,24 +511,21 @@ bool MtpDrafter::capture_round(int T, std::string& err) {
     const int64_t HCN = g_->hc * g_->n_embd;
     if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "mtp: begin capture"; return false; }
     bool ok = true;
-    copy_i32_from_mapped(tok_, m_tok_, T, cs_);
-    copy_i32_from_mapped(step_, m_step_, (int64_t) 2 * T * 4, cs_);
-    copy_i32_from_mapped(pos_, m_pos_, (int64_t) 2 * T * g_->n_head, cs_);
-    copy_i32_from_mapped(row_, m_row_, 2, cs_);
-    copy_from_mapped(Rin_, window_R_, (int64_t) T * HCN, cs_);
     // the catch-up: the layer's front for the window's T cells (their K/V), then its rest for row a only, on row
     // a's intermediates copied to row 0, at the cell the host staged in step row 2*max_t - 1; the draft chain is
     // one graph per step (`capture_step`) so the host can stop it when a draft is unlikely
     const int ra = 2 * max_t_ - 1;
-    const int64_t N = g_->n_embd;
+    const int64_t N = g_->n_embd, NH = g_->n_head;
+    const MappedCopy in[7] = {{tok_, m_tok_, T}, {step_, m_step_, (int64_t) 2 * T * 4}, {pos_, m_pos_, 2 * T * NH},
+                              {row_, m_row_, 2}, {Rin_, window_R_, T * HCN}, {step_ + ra * 4, m_step_ + ra * 4, 4},
+                              {pos_ + ra * NH, m_pos_ + ra * NH, NH}};
+    copy_from_mapped_multi(in, 7, cs_);
     ok = record_front(T, 0, cs_, err);
     if (ok) {
         if (T > 1) {
             copy_row_to_first(row_, R_, HCN, inj_, g_->hc, mixed_, N, cs_);
             native_quantize_q8_1(mixed_, xq_, (int) N, 1, cs_);
         }
-        copy_i32_from_mapped(step_ + ra * 4, m_step_ + ra * 4, 4, cs_);
-        copy_i32_from_mapped(pos_ + ra * g_->n_head, m_pos_ + ra * g_->n_head, g_->n_head, cs_);
         ok = record_rest(ra, cs_, err);
     }
     if (ok) mtp_select(R_, HCN, out_ids_, row_ + 1, Rin_, tok_, m_out_, 0, cs_, probs_, m_prob_);
@@ -511,8 +540,9 @@ bool MtpDrafter::capture_step(int j, std::string& err) {
     const int64_t HCN = g_->hc * g_->n_embd;
     const int row = max_t_ + j - 1;
     if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "mtp: begin capture"; return false; }
-    copy_i32_from_mapped(step_ + row * 4, m_step_ + row * 4, 4, cs_);
-    copy_i32_from_mapped(pos_ + row * g_->n_head, m_pos_ + row * g_->n_head, g_->n_head, cs_);
+    const MappedCopy in[2] = {{step_ + row * 4, m_step_ + row * 4, 4},
+                              {pos_ + row * g_->n_head, m_pos_ + row * g_->n_head, g_->n_head}};
+    copy_from_mapped_multi(in, 2, cs_);
     bool ok = record_forward(1, row, cs_, err);
     if (ok) mtp_select(R_, HCN, out_ids_, row_ + 1, Rin_, tok_, m_out_, j, cs_, probs_, m_prob_);
     return finish_capture(cs_, ok, step_exec_[j], "step", err);

src/kernels/cuda/elementwise.cu

@@ -293,6 +293,37 @@ __global__ void copy_i32_from_mapped_kernel(int32_t* __restrict__ dst, const vol
     for (int i = threadIdx.x; i < n; i += blockDim.x) dst[i] = src[i];
 }
 
+struct MappedCopies {
+    MappedCopy c[8];
+};
+
+__global__ void copy_from_mapped_multi_kernel(MappedCopies a) {
+    const MappedCopy c = a.c[blockIdx.y];
+    const int64_t i0 = (int64_t) blockIdx.x * blockDim.x + threadIdx.x, st = (int64_t) gridDim.x * blockDim.x;
+    if ((c.words & 3) == 0 && ((uintptr_t) c.dst & 15) == 0 && ((uintptr_t) c.src & 15) == 0) {
+        const volatile float4* s = (const volatile float4*) c.src;
+        for (int64_t i = i0; i < c.words / 4; i += st) ((float4*) c.dst)[i] = const_cast<const float4*>(s)[i];
+    } else {
+        const volatile int32_t* s = (const volatile int32_t*) c.src;
+        for (int64_t i = i0; i < c.words; i += st) ((int32_t*) c.dst)[i] = s[i];
+    }
+}
+
+void copy_from_mapped_multi(const MappedCopy* copies, int n, void* stream) {
+    if (n <= 0) return;
+    if (n > 8) { std::fprintf(stderr, "copy_from_mapped_multi: at most 8 copies\n"); std::exit(1); }
+    MappedCopies a;
+    int64_t most = 0;
+    for (int i = 0; i < n; ++i) {
+        a.c[i] = copies[i];
+        most = copies[i].words > most ? copies[i].words : most;
+    }
+    const int64_t units = (most + 3) / 4;
+    const unsigned bx = (unsigned) ((units + 255) / 256 < 32 ? (units + 255) / 256 : 32);
+    copy_from_mapped_multi_kernel<<<dim3(bx < 1 ? 1 : bx, (unsigned) n), 256, 0, (cudaStream_t) stream>>>(a);
+    check_launch("copy_from_mapped_multi");
+}
+
 void copy_i32_from_mapped(int32_t* dst, const int32_t* src, int64_t n, void* stream) {
     if (n <= 0) return;
     copy_i32_from_mapped_kernel<<<1, 128, 0, (cudaStream_t) stream>>>(dst, (const volatile int32_t*) src, (int) n);

src/kernels/cuda/quantize_act.cu

@@ -97,33 +97,35 @@ __global__ void quantize_q8_0_kernel(const float* __restrict__ x, uint8_t* __res
 /// `moe_hit_parity`.  This function's job is different: it must match **this engine's own CPU reference**,
 /// because a hit and a miss for the same expert on the same layer have to produce the same number.  The CPU
 /// path is the reference - C1 passes on it - so the hit path is brought to it, not the reverse.
+// A warp a block, lane i on value i: the block's |max| is a max over the same values (0 for none above it, as the
+// serial scan from 0 through fmaxf), so every byte is the per-thread loop's.
 __global__ void quantize_q8_0_scaled_kernel(const float* __restrict__ x, uint8_t* __restrict__ blocks,
                                             float* __restrict__ scales, long long n_blocks) {
-    const long long b = (long long) blockIdx.x * blockDim.x + threadIdx.x;
+    const long long b = ((long long) blockIdx.x * blockDim.x + threadIdx.x) >> 5;
+    const int i = (int) (threadIdx.x & 31);
     if (b >= n_blocks) return;
-    const float* xb = x + b * QK8_0;
+    const float xi = x[b * QK8_0 + i];
     uint8_t* out = blocks + b * 34;
 
-    float amax = 0.0f;
-    for (int i = 0; i < QK8_0; ++i) amax = fmaxf(amax, fabsf(xb[i]));
+    float amax = fmaxf(0.0f, fabsf(xi));
+    for (int o = 16; o > 0; o >>= 1) amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, o));
     // VERBATIM from `cpu/expert.cpp:144-145`, including the `amax > 0` guard, so the fp32 value written here
     // is bit-identical to the `s` the CPU path used.
     const float s = amax > 0.f ? amax / 127.f : 0.f;
     const float inv = s > 0.f ? 1.f / s : 0.f;
-    scales[b] = s;
-
-    const uint16_t d16bits = f16_from_f32(s);
-    out[0] = (uint8_t) (d16bits & 0xFF);
-    out[1] = (uint8_t) (d16bits >> 8);
-    for (int i = 0; i < QK8_0; ++i) {
-        // VERBATIM from `cpu/expert.cpp:159-162`: reciprocal multiply, then `t + copysign(0.5, t)` truncated
-        // toward zero, which is `lround`'s rule - round half away from zero.
-        const float t = xb[i] * inv;
-        const float r = t + (t >= 0.f ? 0.5f : -0.5f);
-        int v = (int) r;
-        v = v < -127 ? -127 : (v > 127 ? 127 : v);
-        out[2 + i] = (uint8_t) (int8_t) v;
+    if (i == 0) {
+        scales[b] = s;
+        const uint16_t d16bits = f16_from_f32(s);
+        out[0] = (uint8_t) (d16bits & 0xFF);
+        out[1] = (uint8_t) (d16bits >> 8);
     }
+    // VERBATIM from `cpu/expert.cpp:159-162`: reciprocal multiply, then `t + copysign(0.5, t)` truncated
+    // toward zero, which is `lround`'s rule - round half away from zero.
+    const float t = xi * inv;
+    const float r = t + (t >= 0.f ? 0.5f : -0.5f);
+    int v = (int) r;
+    v = v < -127 ? -127 : (v > 127 ? 127 : v);
+    out[2 + i] = (uint8_t) (int8_t) v;
 }
 
 __global__ void dequant_q8_0_kernel(const uint8_t* __restrict__ blocks, float* __restrict__ x,
@@ -254,7 +256,7 @@ void quantize_q8_0_scaled(const float* x, uint8_t* blocks, float* scales, int64_
     }
     const long long nb = n / QK8_0;
     const int threads = 128;
-    const unsigned grid = (unsigned) ((nb + threads - 1) / threads);
+    const unsigned grid = (unsigned) ((nb + threads / 32 - 1) / (threads / 32));   // a warp a block
     quantize_q8_0_scaled_kernel<<<grid, threads, 0, (cudaStream_t) stream>>>(x, blocks, scales, nb);
     const cudaError_t e = cudaGetLastError();
     if (e != cudaSuccess) {