Verify window: the shared expert's gate in the combine, fewer QSA mixer launches, one PLE pass a window; bitwise

- The shared expert's branch ends with its down projection: its scalar gate's logits come from a branch of their own,
  and the gather-combine applies sigmoid and product in shared_expert.cu's instructions (inline PTX without flush to
  zero; shared_expert_gate_rows, which the one-token path now uses, is the reference).  Its last kernels no longer
  wait behind the routed experts' grids: on IQ3_XXS thinking text it ends last before the combine in 0.4-16% of the
  layers instead of 58-66%, and at 2+ tokens the prediction's bias update, dispatched behind them, no longer delays
  the next layer.
- The QSA mixer: the indexer's key and query projections in one launch (bf16_gemv_fp32_mmvf_multi takes a second
  matrix), the key tail's snapshot in the append, each group of heads' RMS norm and RoPE in one kernel
  (native_qsa_norm_rope_tokens; the queries read from their projection's rows, the split copy gone), the keys and
  values dispatched ahead of the queries' grid, and that grid after the indexer's projections, so the append goes
  ahead of it: 8-12 us less a QSA layer.
- The PLE block: both projections of a window's tokens in one pass each, then native_ple_postops_tokens, which now
  writes the history after each token for the commit.  bf16_gemv_fp32_mmvf_multi takes 4 rows a block from 2 tokens
  (each pair's inputs loaded once for them) and the one-row kernel at 1 token: the PLE section 110-174 us a window of
  1-4 tokens instead of 112-428 (IQ3_XXS).
- The verify window requires the native combine, shared expert gate, QSA norms and RoPE and PLE block (the --native
  preset's); their fallbacks are gone from it.
- verify_parity: the gated combine, the fused norm and RoPE, the two-matrix and 4-row projections, the append's tail
  snapshot, the history after each token.

Window hashes identical (thinking, code, after 16K; both models; --spec-split and --second-gpu -1).
Fixed text, both GPUs, against 962b99c, ms per round (2 rounds; UD-Q4_K_XL code 5): IQ3_XXS code -6.1%,
thinking -2.4%, after 16K -2.5%; UD-Q4_K_XL code +1.6% (its CPU pool's drifting RAM rate: the serial part
-0.3 ms, the pool +0.9), thinking -2.0%, after 16K -1.4%.

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

include/strata/core/verify.hpp

@@ -25,7 +25,8 @@
 // the CPU (`pool_takes_q8k` layers) and q8_1 rows for the GPUs, so neither quantizes it.
 //
 // Requires the default native decode configuration (native projections, fused GR, fused GDN, fast attention and
-// selection, native indexer) and a profile-filled VRAM expert tier.
+// selection, native indexer, QSA norms, RoPE, combine, shared expert gate and PLE block) and a profile-filled VRAM
+// expert tier.
 #pragma once
 
 #include "strata/core/expert_source.hpp"
@@ -241,6 +242,10 @@ private:
     cudaEvent_t g2fork_ = nullptr, g2join_[2] = {};
     cudaStream_t cps_ = nullptr;                                  // the wait for the CPU's rows, per token group
     cudaEvent_t cpfork_ = nullptr, cpjoin_[2] = {};
+    cudaStream_t kst_ = nullptr;                                  // the QSA mixer's keys
+    cudaEvent_t kjoin_ = nullptr, iproj_ = nullptr;               // and the end of its indexer's projections
+    cudaStream_t sgs_ = nullptr;                                  // the shared expert's scalar gate
+    cudaEvent_t sgjoin_ = nullptr;
     cudaEvent_t efork_ = nullptr, rdone_ = nullptr;               // the FFN read's end, the router's
     cudaEvent_t done_ = nullptr;                                  // after the last window's graph
 
@@ -258,7 +263,7 @@ private:
     float *qidx_ = nullptr, *scores_ = nullptr, *attn_ = nullptr, *attn32_ = nullptr, *attn_scratch_ = nullptr;
     float* tail_snap_ = nullptr;                              // per QSA layer
     int32_t* sel_ = nullptr;
-    float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *parts_ = nullptr, *hit_out_ = nullptr;
+    float *logits_ = nullptr, *w_ = nullptr, *shared_ = nullptr, *hit_out_ = nullptr;
     int32_t *ids_ = nullptr, *hit_slot_ = nullptr, *hit_dst_ = nullptr, *hit_count_ = nullptr;
     unsigned* ring_count_ = nullptr;                          // verify_router's block counters: the layer's, the prediction's
     float* plogits_ = nullptr;                                    // the predictions' logits, by layer parity
@@ -278,9 +283,11 @@ private:
     float* hit_xs_ = nullptr;
     void* hit_scratch_ = nullptr;
     float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr;
-    uint16_t* sh_bf16_ = nullptr;
     float *sh_gate_ = nullptr, *sh_up_ = nullptr, *sh_g_ = nullptr;
     float* hist_snap_ = nullptr;                              // T * NG_HIST * NG_HC_DIM
+    // the PLE block's projections (keys T * NG_HC_DIM, values T * n_embd) and the rest of the block's rows
+    float *ple_key_ = nullptr, *ple_val_ = nullptr, *ple_nkey_ = nullptr, *ple_norm_ = nullptr, *ple_gated_ = nullptr;
+    float* ple_gate_ = nullptr;
     int64_t cap_ = 0, max_blocks_ = 0, attn_scratch_floats_ = 0;
 
     // --window-profile: stamps [window size][layer, then one row for the window][stage]: a layer's main-stream stages,

include/strata/kernels/bf16_gemv.hpp

@@ -45,8 +45,9 @@ void bf16_gemv_split(const uint16_t* x, const uint16_t* w, float* y, int64_t n_i
 void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y,
                          int64_t n_in, int64_t n_out, void* stream);
 /// `bf16_gemv_fp32_mmvf` for n_tok (1..8) activation rows x (n_tok, n_in) -> y (n_tok, n_out), each weight read once;
-/// every row bitwise what the single-row call gives.
+/// every row bitwise what the single-row call gives.  With `w2` the same launch also gives y2 (n_tok, n_out2) from a
+/// second matrix of the same input.
 void bf16_gemv_fp32_mmvf_multi(const float* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out, int n_tok,
-                               void* stream);
+                               void* stream, const uint16_t* w2 = nullptr, float* y2 = nullptr, int64_t n_out2 = 0);
 
 }  // namespace strata::kernels

include/strata/kernels/native_moe.hpp

@@ -20,9 +20,12 @@ void native_moe_combine(const float* parts, const float* weights, const float* s
 // used to assemble: row t*k + j comes from `gpu_rows` (row-indexed) as the host's zeroed row plus the hit when an
 // entry of dst[0..*count) or dst2[0..*count2) names it, as it is when an entry of dst3[0..*count3) does, and
 // otherwise from `host_rows` (mapped memory: only these rows cross PCIe).  `dst2` and `dst3` may be null.
-// weights (n_tok, k), shared and output (n_tok, n_embd).  Graph-capturable.
+// weights (n_tok, k), shared and output (n_tok, n_embd).  With `shared_gate` (n_tok logits of the shared expert's
+// scalar gate, shared_expert_multi) the shared rows are its down projection's, and each is gated first, bitwise
+// shared_expert_gate_rows.  Graph-capturable.
 void native_moe_gather_combine(const float* gpu_rows, const float* host_rows, const int32_t* dst,
                                const int32_t* count, const int32_t* dst2, const int32_t* count2, const int32_t* dst3,
-                               const int32_t* count3, const float* weights, const float* shared, float* output,
-                               int64_t n_embd, int64_t k, int n_tok, void* stream);
+                               const int32_t* count3, const float* weights, const float* shared,
+                               const float* shared_gate, float* output, int64_t n_embd, int64_t k, int n_tok,
+                               void* stream);
 }

include/strata/kernels/native_ple_postops.hpp

@@ -27,10 +27,11 @@ void native_ple_postops(const float* projected_key, const float* hidden,
                         const PleWeights& weights, const NativePlePostopsBuffers& buffers,
                         void* stream);
 
-// n_tok consecutive tokens (the prompt path), row-major with the token slow: projected_key, hidden, key, query,
-// gated, normalized and result [n_tok, 10240], value [n_tok, 2560], gate [n_tok, 4]. Bitwise n_tok calls of
-// native_ple_postops, each followed by ple_history_advance: history is read, then left holding the last nine
-// normalized rows. result may equal hidden and query may equal normalized; nothing else may overlap.
+// n_tok consecutive tokens (the prompt path, a verify window's group), row-major with the token slow: projected_key,
+// hidden, key, query, gated, normalized and result [n_tok, 10240], value [n_tok, 2560], gate [n_tok, 4]. Bitwise
+// n_tok calls of native_ple_postops, each followed by ple_history_advance: history is read, then left holding the
+// last nine normalized rows; `snaps` (optional, [n_tok, 9 * 10240]) receives the history after each token. result
+// may equal hidden and query may equal normalized; nothing else may overlap.
 struct NativePleTokensBuffers {
     float* key;
     float* query;
@@ -41,6 +42,6 @@ struct NativePleTokensBuffers {
 };
 void native_ple_postops_tokens(const float* projected_key, const float* hidden, const float* value,
                                float* history, const PleWeights& weights, const NativePleTokensBuffers& buffers,
-                               int n_tok, void* stream);
+                               int n_tok, void* stream, float* snaps = nullptr);
 
 } // namespace strata::kernels

include/strata/kernels/native_qsa.hpp

@@ -21,6 +21,13 @@ bool native_qsa_enabled();
 void native_qsa_rms_norm_weighted(const float* input, const float* gamma, float* output,
                                   int n_cols, int n_rows, float epsilon, void* stream);
 
+// native_qsa_rms_norm_weighted, then native_rope_apply_tokens (64 rotated columns) on its output, in one launch and
+// bitwise both: n_rows rows of n_cols (128 or 256) read at input + r * in_stride, written contiguously to output;
+// row r's position is positions[(r / heads) * pos_stride + r % heads].  output == input needs in_stride == n_cols.
+void native_qsa_norm_rope_tokens(const float* input, int in_stride, const float* gamma, float* output, int n_cols,
+                                 int n_rows, float epsilon, int n_rot, float freq_base, const int* positions, int heads,
+                                 int pos_stride, void* stream);
+
 // attn/output [head_dim,n_head]; q_full [2*head_dim,n_head], each row is
 // [query channels, gate channels]. Output is attn*sigmoid(second-half gate).
 // Exact output==attn is supported; all other spans must be disjoint.

include/strata/kernels/native_qsa_indexer.hpp

@@ -32,10 +32,11 @@ void native_qsa_indexer_append(const float* raw, const int32_t* relative_pos_dev
                                const QsaIndexerBuffers& buffers, const QsaShapes& shapes,
                                int64_t max_cells, float freq_base, void* stream);
 // n_tok appends in one launch, in order: token t's raw key at raw + t * 128, its position at
-// relative_pos_device[t * pos_stride].  Bitwise n_tok calls of native_qsa_indexer_append.
+// relative_pos_device[t * pos_stride].  Bitwise n_tok calls of native_qsa_indexer_append.  With `tail_snap`
+// (3 * 128 floats) the key tail as it was before them is copied there first.
 void native_qsa_indexer_append_multi(const float* raw, const int32_t* relative_pos_device, int pos_stride,
                                      int n_tok, int32_t pos_base, const float* gamma, float epsilon,
                                      const QsaIndexerBuffers& buffers, const QsaShapes& shapes,
-                                     int64_t max_cells, float freq_base, void* stream);
+                                     int64_t max_cells, float freq_base, void* stream, float* tail_snap = nullptr);
 
 } // namespace strata::kernels

include/strata/kernels/ple.hpp

@@ -33,6 +33,7 @@ namespace strata::kernels {
 /// Opt in to pinned CUDA BF16/F32 MMVF for the PLE value projection only (default false).
 /// Configure before session capture; captured graphs retain their selected projection kernels.
 void ple_set_native_bf16(bool enabled);
+bool ple_native_bf16();
 
 /// Opt in to pinned CUDA postprojection norms, gate, convolution and residual
 /// arithmetic (default false). Set before capture; existing graphs keep their

include/strata/kernels/shared_expert.hpp

@@ -69,13 +69,19 @@ void shared_expert(const uint8_t* x_q8_0, const uint8_t* x_q8k, const uint16_t*
                    int64_t n_embd, int64_t n_ff, int tpr, void* stream, const float* x_f32 = nullptr,
                    const NativeSharedWeights* native = nullptr);
 
-/// Plan v0.3 P6: the shared expert for `n_tok` <= 8 tokens (a verify window) with all three projections native:
-/// multi-column MMVQ, so the weights are read once; every token is bitwise `shared_expert` on that token.
-/// `x` (n_tok, n_embd) f32, `x_bf16` the same rounded (only read when the native BF16 gate is off), `gate`/`up`
-/// (n_tok, n_ff) scratch, `g` n_tok floats, `out` (n_tok, n_embd).  `nw.q8_1` must hold n_tok columns of n_embd.
-void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw,
-                         const uint16_t* gate_inp_bf16, float* gate, float* up, float* g, float* out, int64_t n_embd,
-                         int64_t n_ff, void* stream);
+/// The native scalar gate applied: `g` (n_tok logits) becomes sigmoid(g), and each token's row of `out` (n_tok,
+/// n_embd) is multiplied by it.
+void shared_expert_gate_rows(float* out, float* g, int64_t n_embd, int n_tok, void* stream);
+
+/// Plan v0.3 P6: the shared expert for `n_tok` <= 8 tokens (a verify window) with all three projections and the gate
+/// native: multi-column MMVQ, so the weights are read once.  `x` (n_tok, n_embd) f32, `gate`/`up` (n_tok, n_ff)
+/// scratch; `out` (n_tok, n_embd) receives the down projection's rows and `g` the scalar gate's n_tok logits,
+/// computed on `gate_stream` (may be `stream`), so that the chain ends with the down projection:
+/// `shared_expert_gate_rows` on them, or `native_moe_gather_combine` applying them, gives every token bitwise
+/// `shared_expert` on that token.  `nw.q8_1` must hold n_tok columns of n_embd.
+void shared_expert_multi(int n_tok, const float* x, const NativeSharedWeights& nw, const uint16_t* gate_inp_bf16,
+                         float* gate, float* up, float* g, float* out, int64_t n_embd, int64_t n_ff, void* stream,
+                         void* gate_stream);
 
 /// The MoE block's final combination, `ref/moe.py::moe` L156:
 ///

src/core/layer.cpp

@@ -447,6 +447,10 @@ bool layer_verify_compatible(std::string& why) {
     else if (!g_fast_attn || native_flash_attn_short) why = "the split-K decode attention is off";
     else if (!g_fast_select) why = "the block top-k selection is off";
     else if (!strata::kernels::native_qsa_indexer_enabled()) why = "the native QSA indexer is off";
+    else if (!strata::kernels::native_qsa_enabled() || !strata::kernels::native_rope_enabled())
+        why = "the native QSA norms or RoPE are off";
+    else if (!strata::kernels::native_moe_combine_enabled()) why = "the native MoE combine is off";
+    else if (!strata::kernels::shared_expert_native_bf16()) why = "the shared expert's native gate is off";
     else return true;
     return false;
 }

src/core/verify.cpp

@@ -15,6 +15,7 @@
 #include "strata/kernels/kv_q8.hpp"
 #include "strata/kernels/native_mmvq.hpp"
 #include "strata/kernels/native_moe.hpp"
+#include "strata/kernels/native_ple_postops.hpp"
 #include "strata/kernels/native_qsa.hpp"
 #include "strata/kernels/native_qsa_indexer.hpp"
 #include "strata/kernels/native_router.hpp"
@@ -25,7 +26,6 @@
 #include "strata/kernels/qsa_decode_attn.hpp"
 #include "strata/kernels/qsa_select.hpp"
 #include "strata/kernels/quantize_act.hpp"
-#include "strata/kernels/rope.hpp"
 #include "strata/kernels/s2_expert_grouped.hpp"
 #include "strata/kernels/sampler.hpp"
 #include "strata/kernels/shared_expert.hpp"
@@ -102,8 +102,11 @@ Verifier::~Verifier() {
     if (shs_) cudaStreamDestroy(shs_);
     if (g2s_) cudaStreamDestroy(g2s_);
     if (cps_) cudaStreamDestroy(cps_);
+    if (kst_) cudaStreamDestroy(kst_);
+    if (sgs_) cudaStreamDestroy(sgs_);
     for (cudaEvent_t e : {join_, bfork_, bjoin_, pfork_, pjoin_, shfork_, shjoin_, res_ready_, done_, g2fork_,
-                          g2join_[0], g2join_[1], cpfork_, cpjoin_[0], cpjoin_[1], efork_, rdone_})
+                          g2join_[0], g2join_[1], cpfork_, cpjoin_[0], cpjoin_[1], efork_, rdone_, kjoin_, iproj_,
+                          sgjoin_})
         if (e) cudaEventDestroy(e);
     if (arena_) cudaFree(arena_);
     if (stamps_) cudaFree(stamps_);
@@ -140,6 +143,11 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         err = "verify: geometry differs from the artifact's";
         return false;
     }
+    if (ss.ple.ready() && (!strata::kernels::ple_native_postops_enabled() || !strata::kernels::ple_native_bf16() ||
+                           (ss.ple.w.key_bf16 == nullptr && ss.ple.w.key_native_data == nullptr))) {
+        err = "verify: the PLE block needs the native postops and value and a BF16 or native key";
+        return false;
+    }
     const WeightRef* wo = wt.find("output.weight");
     if (wo == nullptr) { err = "verify: output.weight is missing"; return false; }
     n_vocab_ = wo->ne1;
@@ -211,7 +219,7 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         sink_.start2 = h_plan_ + plan_ptr_off_ + 4 * cap;
         sink_.cap = cap;
         sink_.publish = &Verifier::publish_plan;
-        sink_.host_rows_only = strata::kernels::native_moe_combine_enabled();   // the gather-combine (post)
+        sink_.host_rows_only = true;                      // the gather-combine (post) reads the GPUs' from VRAM
         sink_.fetch = &Verifier::fetch_dma;
         sink_.ctx = this;
     }
@@ -249,7 +257,7 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         pids_ = b.take<int32_t>(T * K);
         pmixed_ = b.take<float>(T * N); gate_ema_ = b.take<float>((uint64_t) g.n_layers * HC * N);
         pbias_ = b.take<float>((uint64_t) (g.n_layers * g.n_expert));
-        shared_ = b.take<float>(T * N); parts_ = b.take<float>(T * K * N); hit_out_ = b.take<float>(T * K * N);
+        shared_ = b.take<float>(T * N); hit_out_ = b.take<float>(T * K * N);
         hit_slot_ = b.take<int32_t>(T * K); hit_dst_ = b.take<int32_t>(T * K); hit_count_ = b.take<int32_t>(4);
         list2_ = b.take<int32_t>(2 * (uint64_t) list2_stride_);
         ring_count_ = b.take<unsigned>(2);
@@ -262,10 +270,13 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
             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)));
         head_mixed_ = b.take<float>(T * N); head_inj_ = b.take<float>(HC);
-        sh_bf16_ = b.take<uint16_t>(T * N); sh_gate_ = b.take<float>(T * (uint64_t) g.n_ff);
+        sh_gate_ = b.take<float>(T * (uint64_t) g.n_ff);
         sh_up_ = b.take<float>(T * (uint64_t) g.n_ff); sh_g_ = b.take<float>(T + 4);
         head_logits_ = b.take<float>(T * (uint64_t) n_vocab_);
         hist_snap_ = b.take<float>(T * HS);
+        const uint64_t PD = (uint64_t) strata::kernels::NG_HC_DIM;
+        ple_key_ = b.take<float>(T * PD); ple_val_ = b.take<float>(T * N); ple_nkey_ = b.take<float>(T * PD);
+        ple_norm_ = b.take<float>(T * PD); ple_gated_ = b.take<float>(T * PD); ple_gate_ = b.take<float>(T * HC);
     };
     Bump count;
     carve(count);
@@ -306,6 +317,11 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         cudaStreamCreateWithFlags(&shs_, cudaStreamNonBlocking) != cudaSuccess ||
         cudaStreamCreateWithFlags(&g2s_, cudaStreamNonBlocking) != cudaSuccess ||
         cudaStreamCreateWithFlags(&cps_, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaStreamCreateWithFlags(&kst_, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaEventCreateWithFlags(&kjoin_, cudaEventDisableTiming) != cudaSuccess ||
+        cudaEventCreateWithFlags(&iproj_, cudaEventDisableTiming) != cudaSuccess ||
+        cudaStreamCreateWithFlags(&sgs_, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaEventCreateWithFlags(&sgjoin_, cudaEventDisableTiming) != cudaSuccess ||
         cudaEventCreateWithFlags(&cpfork_, cudaEventDisableTiming) != cudaSuccess ||
         cudaEventCreateWithFlags(&cpjoin_[0], cudaEventDisableTiming) != cudaSuccess ||
         cudaEventCreateWithFlags(&cpjoin_[1], cudaEventDisableTiming) != cudaSuccess ||
@@ -446,7 +462,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
     }
 
     // the shared expert: in a window of one token group on a branch of its own forked after the FFN read (`efork_`),
-    // beside the router and the routed experts (joined before the combine); in a split window inline
+    // beside the router and the routed experts (joined before the combine); in a split window inline.  It ends with
+    // its down projection: the combine applies its scalar gate, whose logits come from a branch of their own.
     const bool beside = G == 1;
     auto shared = [&](int64_t l, int grp) -> bool {
         const int tb = tb_[grp], te = te_[grp], n = te - tb;
@@ -462,27 +479,26 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
         nsw.up_type = wsu->native_type; nsw.up_data = wsu->native_data;
         nsw.down_type = wsd->native_type; nsw.down_data = wsd->native_data;
         nsw.q8_1 = xq_;                                   // the main stream's next use is the next layer's mixer
-        cudaStream_t ws = cs;
+        cudaStream_t ws = cs, gs = cs;
         if (beside) {
-            if (cudaStreamWaitEvent(shs_, efork_, 0) != cudaSuccess) {
+            if (cudaStreamWaitEvent(shs_, efork_, 0) != cudaSuccess ||
+                cudaStreamWaitEvent(sgs_, efork_, 0) != cudaSuccess) {
                 err = "verify: the shared expert's branch could not fork";
                 return false;
             }
             ws = shs_;
+            gs = sgs_;
         }
-        if (!shared_expert_native_bf16())                 // the scalar gate's BF16 input (the native gate reads x)
-            for (int t = tb; t < te; ++t) f32_to_bf16_bulk(mixed_ + t * N, sh_bf16_ + t * N, N, ws);
         try {
-            shared_expert_multi(n, mixed_ + tb * N, sh_bf16_ + tb * N, nsw, (const uint16_t*) wgi->data,
-                                sh_gate_ + (size_t) tb * g.n_ff, sh_up_ + (size_t) tb * g.n_ff, sh_g_ + tb, shared_ + tb * N,
-                                N, g.n_ff, ws);
+            shared_expert_multi(n, mixed_ + tb * N, nsw, (const uint16_t*) wgi->data, sh_gate_ + (size_t) tb * g.n_ff,
+                                sh_up_ + (size_t) tb * g.n_ff, sh_g_ + tb, shared_ + tb * N, N, g.n_ff, ws, gs);
         } catch (const std::exception& e) {
             err = std::string("verify shared expert: ") + e.what();
             return false;
         }
         if (beside) {
             stamp(l, kLayerStamps, shs_);
-            if (cudaEventRecord(shjoin_, shs_) != cudaSuccess) {
+            if (cudaEventRecord(shjoin_, shs_) != cudaSuccess || cudaEventRecord(sgjoin_, sgs_) != cudaSuccess) {
                 err = "verify: the shared expert's branch could not join";
                 return false;
             }
@@ -505,25 +521,46 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
             if (!wn[h] || !wd[h] || !wu[h] || !wi[h]) return false;
         }
         bool pending = l > 0;   // the previous layer's FFN write, folded into this layer's first read
+        // the projections of the group's tokens: 2+ tokens are quantized with the interleaved copy the multi-token
+        // kernels read
+        auto quant = [&](const float* x, int64_t n_in) {
+            if (n >= 2) native_quantize_q8_1_il(x, xq_, xil_, (int) n_in, n, cs);
+            else native_quantize_q8_1(x, xq_, (int) n_in, n, cs);
+        };
+        auto mm = [&](int type, const void* w, float* y, int64_t n_in, int64_t n_out, cudaStream_t ps) {
+            if (n >= 2) native_mmvq_il(type, w, xq_, xil_, y, (int) n_in, (int) n_out, n, ps);
+            else native_mmvq(type, w, xq_, y, (int) n_in, (int) n_out, n, ps);
+        };
+        auto proj = [&](const WeightRef* w, float* y, int64_t n_in, int64_t n_out, cudaStream_t ps) {
+            mm(w->native_type, w->native_data, y, n_in, n_out, ps);
+        };
         if (l == 1 && ple_on) {
             if (grp == 0) {                                // the host reads the rows while layer 0 runs
                 wait_flag_ge(m_pleflag_, 1, cs);
                 copy_from_mapped(ple_, m_ple_, (int64_t) T * N, cs);
             }
-            float* normalized = (float*) ((uint8_t*) ss.ple.scratch + ple_block_scratch_bytes());
-            for (int t = tb; t < te; ++t) {
-                gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs);
-                PleOut po;
-                po.normalized = normalized;
-                po.result = Rt(t);
-                try {
-                    ple_block(ple_ + t * N, Rt(t), ss.ple.hist, ss.ple.w, po, ss.ple.scratch, cs);
-                    ple_history_advance(ss.ple.hist, normalized, cs);
-                } catch (const std::exception& e) {
-                    err = std::string("verify PLE: ") + e.what();
-                    return false;
+            for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs);
+            // the PLE block: the group's projections with each weight read once, then the rest of the block over its
+            // tokens in order, which leaves the history after each of them for the commit
+            const PleWeights& pw = ss.ple.w;
+            const float* emb = ple_ + tb * N;
+            float* key = ple_key_ + (size_t) tb * NG_HC_DIM;
+            float* norm = ple_norm_ + (size_t) tb * NG_HC_DIM;
+            try {
+                if (pw.key_bf16 != nullptr) {
+                    bf16_gemv_fp32_mmvf_multi(emb, pw.key_bf16, key, N, NG_HC_DIM, n, cs);
+                } else {
+                    quant(emb, N);
+                    mm(pw.key_native_type, pw.key_native_data, key, N, NG_HC_DIM, cs);
                 }
-                copy_from_mapped(hist_snap_ + (size_t) t * HS, ss.ple.hist, HS, cs);
+                bf16_gemv_fp32_mmvf_multi(emb, pw.value_bf16, ple_val_ + tb * N, N, N, n, cs);
+                const NativePleTokensBuffers pb{ple_nkey_ + (size_t) tb * NG_HC_DIM, norm, ple_gate_ + tb * HC,
+                                                ple_gated_ + (size_t) tb * NG_HC_DIM, norm, Rt(tb)};
+                native_ple_postops_tokens(key, Rt(tb), ple_val_ + tb * N, ss.ple.hist, pw, pb, n, cs,
+                                          hist_snap_ + (size_t) tb * HS);
+            } catch (const std::exception& e) {
+                err = std::string("verify PLE: ") + e.what();
+                return false;
             }
             pending = false;
         }
@@ -545,16 +582,6 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
         gr_read_group(0, pending, inj2_, inj_, nullptr, nullptr, nullptr);
         stamp(l, 1, cs);
         float* xm = mixed_ + tb * N;
-        // the mixer's projections: a group of 2+ tokens is quantized with the interleaved copy the multi-token kernels
-        // read
-        auto quant = [&](const float* x, int64_t n_in) {
-            if (n >= 2) native_quantize_q8_1_il(x, xq_, xil_, (int) n_in, n, cs);
-            else native_quantize_q8_1(x, xq_, (int) n_in, n, cs);
-        };
-        auto proj = [&](const WeightRef* w, float* y, int64_t n_in, int64_t n_out, cudaStream_t ps) {
-            if (n >= 2) native_mmvq_il(w->native_type, w->native_data, xq_, xil_, y, (int) n_in, (int) n_out, n, ps);
-            else native_mmvq(w->native_type, w->native_data, xq_, y, (int) n_in, (int) n_out, n, ps);
-        };
         try {
             if (!is_qsa_layer(g, l)) {
                 // ======================= GDN =======================
@@ -606,55 +633,51 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
                 if (!native_of(wq, v.name("attn_q.weight"), err) || !native_of(wk, v.name("attn_k.weight"), err) ||
                     !native_of(wv, v.name("attn_v.weight"), err) || !native_of(wo, v.name("attn_output.weight"), err))
                     return false;
-                auto norm_rope = [&](float* data, const WeightRef* norm, int rows, int cols, const int32_t* pos,
-                                     cudaStream_t ns) {
-                    if (native_qsa_enabled()) native_qsa_rms_norm_weighted(data, (const float*) norm->data, data, cols, rows, EPS, ns);
-                    else rms_norm_weighted(data, (const float*) norm->data, rows, cols, EPS, ns);
-                    if (native_rope_enabled()) native_rope_apply(data, data, rows, cols, (int) s.n_rot, (float) qsa_freq_base(), pos, ns);
-                    else rope_neox_apply(data, data, rows, cols, (int) s.n_rot, st.cos_tab, st.sin_tab, pos, ns);
+                // the rms norm and rope of the group's heads in one launch: rows of `cols` read at in + r * stride,
+                // token t's positions at pos_ + t * NH
+                auto norm_rope = [&](const float* in, int stride, float* out, const WeightRef* norm, int heads,
+                                     int cols, cudaStream_t ns) {
+                    native_qsa_norm_rope_tokens(in, stride, (const float*) norm->data, out, cols, n * heads, EPS,
+                                                (int) s.n_rot, (float) qsa_freq_base(), pos_ + tb * NH, heads, (int) NH,
+                                                ns);
                 };
-                // the same for the n tokens' heads in one launch each: token t's positions at pos + t * NH
-                const bool batched = native_qsa_enabled() && native_rope_enabled();
-                auto norm_rope_tokens = [&](float* data, const WeightRef* norm, int heads, int cols, const int32_t* pos,
-                                            cudaStream_t ns) {
-                    native_qsa_rms_norm_weighted(data, (const float*) norm->data, data, cols, n * heads, EPS, ns);
-                    native_rope_apply_tokens(data, data, n * heads, cols, (int) s.n_rot, (float) qsa_freq_base(), pos, heads,
-                                             (int) NH, ns);
-                };
-                // two branches beside the keys and values, joined before the attention: the indexer (its keys appended,
-                // the queries' block scores and top-k), and the queries
+                // three branches beside the values, joined before the attention: the indexer (its key and query
+                // projections in one launch, the keys appended, the queries' block scores and top-k), the keys and the
+                // queries.  A grid is dispatched whole before a later kernel's blocks, so the keys' and values'
+                // projections go ahead of the queries' grid, and that grid waits for the indexer's projections: the
+                // append is dispatched ahead of it, and the indexer's last kernels end about when the queries' do.
                 if (!fork(side_, bfork_)) return false;
                 float* idx_raw = idx_raw_L_ + (size_t) qi * MT * ID;
-                bf16_gemv_fp32_mmvf_multi(xm, (const uint16_t*) wik->data, idx_raw + tb * ID, N, ID, n, side_);
-                if (grp == 0) copy_from_mapped(tail_snap_ + (size_t) qi * TS, st.idx_tail, TS, side_);
+                bf16_gemv_fp32_mmvf_multi(xm, (const uint16_t*) wik->data, idx_raw + tb * ID, N, ID, n, side_,
+                                          (const uint16_t*) wiq->data, qidx_ + tb * IQ * ID, IQ * ID);
+                if (!mark(side_, iproj_)) return false;
                 const QsaIndexerBuffers ib{st.idx_tail, st.idx_dead, st.idx_pooled, st.idx_block_pos};
                 native_qsa_indexer_append_multi(idx_raw + tb * ID, step_ + tb * kStepCount + kStepPos, (int) kStepCount, n,
                                                 0, (const float*) wikn->data, EPS, ib, s, st.max_cells,
-                                                (float) qsa_freq_base(), side_);
-                bf16_gemv_fp32_mmvf_multi(xm, (const uint16_t*) wiq->data, qidx_ + tb * IQ * ID, N, IQ * ID, n, side_);
-                if (batched) norm_rope_tokens(qidx_ + tb * IQ * ID, wiqn, (int) IQ, (int) ID, pos_ + tb * NH, side_);
-                else for (int t = tb; t < te; ++t) norm_rope(qidx_ + t * IQ * ID, wiqn, (int) IQ, (int) ID, pos_ + t * NH, side_);
+                                                (float) qsa_freq_base(), side_,
+                                                grp == 0 ? tail_snap_ + (size_t) qi * TS : nullptr);
+                norm_rope(qidx_ + tb * IQ * ID, (int) ID, qidx_ + tb * IQ * ID, wiqn, (int) IQ, (int) ID, side_);
                 qsa_block_scores(st.idx_pooled, st.idx_dead, qidx_ + tb * IQ * ID, step_ + tb * kStepCount, n, max_blocks_,
                                  s, scores_ + (size_t) tb * max_blocks_, side_, qsa_score_grid_blocks);
                 qsa_block_topk(scores_ + (size_t) tb * max_blocks_, step_ + tb * kStepCount, n, max_blocks_, cap_, s,
                                sel_ + (size_t) tb * cap_, side_);
                 if (!mark(side_, bjoin_)) return false;
                 quant(xm, N);
-                if (!fork(shs_, pfork_)) return false;
-                proj(wq, qfull_ + tb * NH * 2 * HD, N, NH * 2 * HD, shs_);
-                // the query halves of every token's heads (rows of 2 * HD: query, then gate)
-                if (cudaMemcpy2DAsync(qcur_ + tb * NH * HD, (size_t) HD * 4, qfull_ + tb * NH * 2 * HD, (size_t) HD * 2 * 4,
-                                      (size_t) HD * 4, (size_t) (n * NH), cudaMemcpyDeviceToDevice, shs_) != cudaSuccess) {
-                    err = "verify: the q/gate split failed";
+                if (cudaEventRecord(pfork_, cs) != cudaSuccess || cudaStreamWaitEvent(kst_, pfork_, 0) != cudaSuccess ||
+                    cudaStreamWaitEvent(shs_, pfork_, 0) != cudaSuccess ||
+                    cudaStreamWaitEvent(shs_, iproj_, 0) != cudaSuccess) {
+                    err = "verify: a branch could not fork";
                     return false;
                 }
-                if (batched) norm_rope_tokens(qcur_ + tb * NH * HD, wqn, (int) NH, (int) HD, pos_ + tb * NH, shs_);
-                else for (int t = tb; t < te; ++t) norm_rope(qcur_ + t * NH * HD, wqn, (int) NH, (int) HD, pos_ + t * NH, shs_);
-                if (!mark(shs_, pjoin_)) return false;
-                proj(wk, kcur_ + tb * NKV * HD, N, NKV * HD, cs);
+                proj(wk, kcur_ + tb * NKV * HD, N, NKV * HD, kst_);
+                norm_rope(kcur_ + tb * NKV * HD, (int) HD, kcur_ + tb * NKV * HD, wkn, (int) NKV, (int) HD, kst_);
+                if (!mark(kst_, kjoin_)) return false;
                 proj(wv, vcur_ + tb * NKV * HD, N, NKV * HD, cs);
-                if (batched) norm_rope_tokens(kcur_ + tb * NKV * HD, wkn, (int) NKV, (int) HD, pos_ + tb * NH, cs);
-                else for (int t = tb; t < te; ++t) norm_rope(kcur_ + t * NKV * HD, wkn, (int) NKV, (int) HD, pos_ + t * NH, cs);
+                proj(wq, qfull_ + tb * NH * 2 * HD, N, NH * 2 * HD, shs_);
+                // the query halves of every token's heads (rows of 2 * HD: query, then gate)
+                norm_rope(qfull_ + tb * NH * 2 * HD, (int) (2 * HD), qcur_ + tb * NH * HD, wqn, (int) NH, (int) HD,
+                          shs_);
+                if (!mark(shs_, pjoin_) || !wait(kjoin_)) return false;
                 if (st.kv_int8)
                     kv_append_q8_steps(st.k_q, st.v_q, st.k_scale, st.v_scale, st.page_table, step_ + tb * kStepCount,
                                        (int) kStepCount, kcur_ + tb * NKV * HD, vcur_ + tb * NKV * HD, n, s, cs);
@@ -669,12 +692,8 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
                 else { pools.k_pool = st.k_pool; pools.v_pool = st.v_pool; }
                 qsa_decode_attn_batch(qcur_ + tb * NH * HD, pools, sel_ + (size_t) tb * cap_, step_ + tb * kStepCount, cap_,
                                       s, attn_scratch_ + (size_t) tb * attn_scratch_floats_, attn_ + tb * NH * HD, n, cs);
-                if (native_qsa_enabled())
-                    native_qsa_gate_apply(attn_ + tb * NH * HD, qfull_ + tb * NH * 2 * HD, attn32_ + tb * NH * HD,
-                                          (int) (n * NH), (int) HD, cs);
-                else
-                    for (int t = tb; t < te; ++t)
-                        qsa_gate_apply_f32(attn_ + t * NH * HD, qfull_ + t * NH * 2 * HD, s, attn32_ + t * NH * HD, cs);
+                native_qsa_gate_apply(attn_ + tb * NH * HD, qfull_ + tb * NH * 2 * HD, attn32_ + tb * NH * HD,
+                                      (int) (n * NH), (int) HD, cs);
                 quant(attn32_ + tb * NH * HD, NH * HD);
                 proj(wo, bo_ + tb * N, NH * HD, N, cs);
             }
@@ -874,31 +893,21 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
         stamp(l, 9, cs);
         if (gpu2_ && !wait(g2join_[grp])) return false;        // the second GPU's, in `hit_out`
         stamp(l, 10, cs);
-        if (beside && cudaStreamWaitEvent(cs, shjoin_, 0) != cudaSuccess) {
+        if (beside && (cudaStreamWaitEvent(cs, shjoin_, 0) != cudaSuccess ||
+                       cudaStreamWaitEvent(cs, sgjoin_, 0) != cudaSuccess)) {
             err = "verify: the shared expert's branch could not join";
             return false;
         }
-        if (native_moe_combine_enabled()) {
-            // only the CPU's rows cross PCIe here; the GPUs' come from `hit_out` (the pool leaves their rows unwritten)
-            const int32_t* list2 = gpu2_ ? list2_ + grp * list2_stride_ : nullptr;
-            try {
-                native_moe_gather_combine(hit_out, m_ymiss_ + (size_t) tb * K * N, vram.dst, vram.counts + 1, pcie.dst,
-                                          pcie.counts != nullptr ? pcie.counts + 1 : nullptr,
-                                          list2 != nullptr ? list2 + 4 : nullptr, list2, w_ + tb * K, shared_ + tb * N,
-                                          bo_ + tb * N, N, K, n, cs);
-            } catch (const std::exception& e) {
-                err = "verify layer " + std::to_string(l) + " combine: " + e.what();
-                return false;
-            }
-        } else {
-            copy_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K * N, cs);
-            moe_hit_add(parts_ + (size_t) tb * K * N, hit_out, vram.dst, vram.counts + 1, cap, N, cs);
-            if (pcie_share_) moe_hit_add(parts_ + (size_t) tb * K * N, hit_out, pcie.dst, pcie.counts + 1, cap, N, cs);
-            for (int t = tb; t < te; ++t) {
-                MoEBuffers mb = ss.moe;
-                mb.weights = w_ + t * K; mb.shared = shared_ + t * N;
-                if (!moe_combine_parts(g, l, K, mb, parts_ + (size_t) t * K * N, bo_ + t * N, cs, err)) return false;
-            }
+        // only the CPU's rows cross PCIe here; the GPUs' come from `hit_out` (the pool leaves their rows unwritten)
+        const int32_t* list2 = gpu2_ ? list2_ + grp * list2_stride_ : nullptr;
+        try {
+            native_moe_gather_combine(hit_out, m_ymiss_ + (size_t) tb * K * N, vram.dst, vram.counts + 1, pcie.dst,
+                                      pcie.counts != nullptr ? pcie.counts + 1 : nullptr,
+                                      list2 != nullptr ? list2 + 4 : nullptr, list2, w_ + tb * K, shared_ + tb * N,
+                                      sh_g_ + tb, bo_ + tb * N, N, K, n, cs);
+        } catch (const std::exception& e) {
+            err = "verify layer " + std::to_string(l) + " combine: " + e.what();
+            return false;
         }
         stamp(l, 11, cs);
         if (branches(l) && cudaStreamWaitEvent(cs, join_, 0) != cudaSuccess) {

src/kernels/cuda/mmvf_multi.cuh

@@ -1,6 +1,6 @@
-// src/kernels/cuda/mmvf_multi.cuh - one block's row of the FP32-activation BF16 MMVF for up to 8 activation rows, for
-// bf16_gemv_fp32_mmvf_multi (native_bf16.cu) and the verify window's router (native_router.cu); both translation
-// units compile with --use_fast_math, the pinned contract's arithmetic.
+// src/kernels/cuda/mmvf_multi.cuh - one block's row, or rows, of the FP32-activation BF16 MMVF for up to 8 activation
+// rows, for bf16_gemv_fp32_mmvf_multi (native_bf16.cu) and the verify window's router (native_router.cu); both
+// translation units compile with --use_fast_math, the pinned contract's arithmetic.
 //
 // Adapted from llama.cpp 3cf03257f219afbe7334045ff7c6a06ac68c627d, ggml/src/ggml-cuda/{mmvf.cu,common.cuh}.
 // MIT License, Copyright (c) 2023-2026 The ggml authors (see native_bf16.cu for the full notice).
@@ -70,5 +70,72 @@ __device__ __forceinline__ void mmvf_multi_row(const float* __restrict__ x, cons
     }
 }
 
+// mmvf_multi_row for R consecutive rows of one matrix (the first `rows` of them real): each pair's inputs are loaded
+// once for the R rows, and every row's sums and their order are mmvf_multi_row's; thread 0 ends with acc[r][0 .. n_tok).
+template <int BLOCK_SIZE, int R>
+__device__ __forceinline__ void mmvf_multi_rows(const float* __restrict__ x, const uint16_t* __restrict__ row0, int n_in,
+                                                int n_tok, int rows, float (&acc)[R][kMmvfMaxRows]) {
+    const int t = threadIdx.x;
+    __shared__ float partials[R][kMmvfMaxRows][32];
+    if constexpr (BLOCK_SIZE > 32) {
+        if (t < 32)
+#pragma unroll
+            for (int r = 0; r < R; ++r)
+#pragma unroll
+                for (int k = 0; k < kMmvfMaxRows; ++k) partials[r][k][t] = 0.0f;
+        __syncthreads();
+    }
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+#pragma unroll
+        for (int k = 0; k < kMmvfMaxRows; ++k) acc[r][k] = 0.0f;
+    for (int pair = t; pair < n_in / 2; pair += BLOCK_SIZE) {
+        float w0[R], w1[R];
+#pragma unroll
+        for (int r = 0; r < R; ++r) {
+            const uint32_t weight = r < rows ? reinterpret_cast<const uint32_t*>(row0 + (size_t) r * n_in)[pair] : 0u;
+            w0[r] = f32_from_bf16((uint16_t) weight);
+            w1[r] = f32_from_bf16((uint16_t) (weight >> 16));
+        }
+#pragma unroll
+        for (int k = 0; k < kMmvfMaxRows; ++k) {
+            if (k >= n_tok) break;
+            const float2 input = reinterpret_cast<const float2*>(x + (size_t) k * n_in)[pair];
+#pragma unroll
+            for (int r = 0; r < R; ++r) {
+                acc[r][k] = __fmaf_rn(w0[r], input.x, acc[r][k]);
+                acc[r][k] = __fmaf_rn(w1[r], input.y, acc[r][k]);
+            }
+        }
+    }
+#pragma unroll
+    for (int r = 0; r < R; ++r)
+#pragma unroll
+        for (int k = 0; k < kMmvfMaxRows; ++k) {
+            if (k >= n_tok) break;
+            acc[r][k] = mmvf_warp_sum(acc[r][k]);
+        }
+    if constexpr (BLOCK_SIZE > 32) {
+        if ((t & 31) == 0)
+#pragma unroll
+            for (int r = 0; r < R; ++r)
+#pragma unroll
+                for (int k = 0; k < kMmvfMaxRows; ++k) {
+                    if (k >= n_tok) break;
+                    partials[r][k][t / 32] = acc[r][k];
+                }
+        __syncthreads();
+        if (t < 32) {
+#pragma unroll
+            for (int r = 0; r < R; ++r)
+#pragma unroll
+                for (int k = 0; k < kMmvfMaxRows; ++k) {
+                    if (k >= n_tok) break;
+                    acc[r][k] = mmvf_warp_sum(partials[r][k][t]);
+                }
+        }
+    }
+}
+
 }  // namespace
 }  // namespace strata::kernels

src/kernels/cuda/native_bf16.cu

@@ -64,15 +64,25 @@ __global__ void bf16_f32_mmvf_kernel(const float* __restrict__ x, const uint16_t
     if (t == 0) y[blockIdx.x] = acc;
 }
 
-// `bf16_f32_mmvf_kernel` for up to kMultiT activation rows (mmvf_multi_row)
-constexpr int kMultiT = kMmvfMaxRows;
+// `bf16_f32_mmvf_kernel` for up to kMultiT activation rows, kRows weight rows a block (mmvf_multi_rows); the blocks
+// past n_out's take the rows of a second matrix of the same input
+constexpr int kMultiT = kMmvfMaxRows, kRows = 4;
 template <int BLOCK_SIZE>
 __global__ void bf16_f32_mmvf_multi_kernel(const float* __restrict__ x, const uint16_t* __restrict__ w,
-                                           float* __restrict__ y, int n_in, int n_out, int n_tok) {
-    float acc[kMultiT];
-    mmvf_multi_row<BLOCK_SIZE>(x, w + (size_t) blockIdx.x * n_in, n_in, n_tok, acc);
+                                           float* __restrict__ y, int n_in, int n_out, int n_tok,
+                                           const uint16_t* __restrict__ w2, float* __restrict__ y2, int n_out2) {
+    const int blocks1 = (n_out + kRows - 1) / kRows;
+    const bool first = (int) blockIdx.x < blocks1;
+    const int row0 = (first ? (int) blockIdx.x : (int) blockIdx.x - blocks1) * kRows, n = first ? n_out : n_out2;
+    const int rows = min(kRows, n - row0);
+    float acc[kRows][kMultiT];
+    mmvf_multi_rows<BLOCK_SIZE, kRows>(x, (first ? w : w2) + (size_t) row0 * n_in, n_in, n_tok, rows, acc);
     if (threadIdx.x == 0)
-        for (int k = 0; k < n_tok; ++k) y[(size_t) k * n_out + blockIdx.x] = acc[k];
+#pragma unroll
+        for (int r = 0; r < kRows; ++r)
+#pragma unroll
+            for (int k = 0; k < kMultiT; ++k)
+                if (r < rows && k < n_tok) (first ? y : y2)[(size_t) k * n + row0 + r] = acc[r][k];
 }
 
 int mmvf_block_size(int64_t n_in) {
@@ -121,20 +131,26 @@ void bf16_gemv_fp32_mmvf(const float* x, const uint16_t* w, float* y,
 
 
 void bf16_gemv_fp32_mmvf_multi(const float* x, const uint16_t* w, float* y, int64_t n_in, int64_t n_out, int n_tok,
-                               void* stream) {
+                               void* stream, const uint16_t* w2, float* y2, int64_t n_out2) {
     if (n_tok < 1 || n_tok > kMultiT)
         throw std::invalid_argument("bf16_gemv_fp32_mmvf_multi: 1..8 activation rows");
     if (n_in <= 0 || (n_in & 1) != 0 || n_in > std::numeric_limits<int>::max() ||
-        n_out <= 0 || n_out > std::numeric_limits<int>::max())
+        n_out <= 0 || n_out2 < 0 || n_out + n_out2 > std::numeric_limits<int>::max())
         throw std::invalid_argument("bf16_gemv_fp32_mmvf_multi: require positive even n_in and positive n_out <= INT_MAX");
-    if (x == nullptr || w == nullptr || y == nullptr ||
+    if (x == nullptr || w == nullptr || y == nullptr || (n_out2 > 0 && (w2 == nullptr || y2 == nullptr)) ||
         (reinterpret_cast<uintptr_t>(x) & 7u) != 0 ||
-        (reinterpret_cast<uintptr_t>(w) & 3u) != 0 ||
-        (reinterpret_cast<uintptr_t>(y) & 3u) != 0)
+        (reinterpret_cast<uintptr_t>(w) & 3u) != 0 || (reinterpret_cast<uintptr_t>(w2) & 3u) != 0 ||
+        (reinterpret_cast<uintptr_t>(y) & 3u) != 0 || (reinterpret_cast<uintptr_t>(y2) & 3u) != 0)
         throw std::invalid_argument("bf16_gemv_fp32_mmvf_multi: null or misaligned pointer");
+    if (n_tok == 1 && n_out2 == 0) {   // the same sums: a row a block streams the weights faster
+        bf16_gemv_fp32_mmvf(x, w, y, n_in, n_out, stream);
+        return;
+    }
     const cudaStream_t st = (cudaStream_t) stream;
+    const unsigned blocks = (unsigned) ((n_out + kRows - 1) / kRows + (n_out2 + kRows - 1) / kRows);
 #define STRATA_MMVF_CASE(N) case N: \
-    bf16_f32_mmvf_multi_kernel<N><<<(unsigned) n_out, N, 0, st>>>(x, w, y, (int) n_in, (int) n_out, n_tok); break
+    bf16_f32_mmvf_multi_kernel<N><<<blocks, N, 0, st>>>(x, w, y, (int) n_in, (int) n_out, n_tok, w2, y2, (int) n_out2); \
+    break
     switch (mmvf_block_size(n_in)) {
         STRATA_MMVF_CASE(32);
         STRATA_MMVF_CASE(64);

src/kernels/cuda/native_moe.cu

@@ -44,14 +44,29 @@ __global__ void combine(const float* __restrict__ parts, const float* __restrict
     output[col] = sum;
 }
 constexpr int kMaxK = 15;
+// A shared expert's row times its scalar gate, in shared_expert_gate_rows' instructions (shared_expert.cu builds
+// without fast math: nothing flushes to zero): sigmoid(logit) as __fdividef(1, 1 + __expf(-logit)), then the product.
+// The explicit rounding keeps ptxas from fusing the product into the sum after it.
+__device__ __forceinline__ float gated_shared(float row, float logit) {
+    float e, d, g, y;
+    const float one = 1.0f;
+    asm("mul.rn.f32 %0, %1, 0fBFB8AA3B;" : "=f"(e) : "f"(logit));
+    asm("ex2.approx.f32 %0, %1;" : "=f"(d) : "f"(e));
+    asm("add.rn.f32 %0, %1, %2;" : "=f"(e) : "f"(d), "f"(one));
+    asm("div.approx.f32 %0, %1, %2;" : "=f"(g) : "f"(one), "f"(e));
+    asm("mul.rn.f32 %0, %1, %2;" : "=f"(y) : "f"(row), "f"(g));
+    return y;
+}
 // One block row per token.  The rows a GPU computed are marked from the plan's entries first; each column then
-// loads its k values (the host's over PCIe, all in flight at once) and sums them as `combine` does.
+// loads its k values (the host's over PCIe, all in flight at once) and sums them as `combine` does, then adds the
+// shared expert's row, times its gate when `shared_gate` holds the gates' logits.
 __global__ void gather_combine(const float* __restrict__ gpu_rows, const float* host_rows,
                                const int32_t* __restrict__ dst, const int32_t* __restrict__ count,
                                const int32_t* __restrict__ dst2, const int32_t* __restrict__ count2,
                                const int32_t* __restrict__ dst3, const int32_t* __restrict__ count3,
                                const float* __restrict__ weights, const float* __restrict__ shared,
-                               float* __restrict__ output, int64_t n_embd, int k) {
+                               const float* __restrict__ shared_gate, float* __restrict__ output, int64_t n_embd,
+                               int k) {
     __shared__ unsigned on_gpu, as_is;   // bit j: row t*k + j is in gpu_rows, a hit / a row taken as it is
     const int t = blockIdx.y;
     if (threadIdx.x == 0) on_gpu = as_is = 0u;
@@ -86,7 +101,10 @@ __global__ void gather_combine(const float* __restrict__ gpu_rows, const float*
         if (j >= k) break;
         sum += v[j] * w[j];
     }
-    if (shared) sum += shared[int64_t(t) * n_embd + col];
+    if (shared) {
+        const float s = shared[int64_t(t) * n_embd + col];
+        sum += shared_gate ? gated_shared(s, shared_gate[t]) : s;
+    }
     output[int64_t(t) * n_embd + col] = sum;
 }
 bool valid_span(const void* p, size_t bytes) {
@@ -119,17 +137,18 @@ void native_moe_combine(const float* parts, const float* weights, const float* s
 }
 void native_moe_gather_combine(const float* gpu_rows, const float* host_rows, const int32_t* dst,
                                const int32_t* count, const int32_t* dst2, const int32_t* count2, const int32_t* dst3,
-                               const int32_t* count3, const float* weights, const float* shared, float* output,
-                               int64_t n_embd, int64_t k, int n_tok, void* stream) {
+                               const int32_t* count3, const float* weights, const float* shared,
+                               const float* shared_gate, float* output, int64_t n_embd, int64_t k, int n_tok,
+                               void* stream) {
     if (!stream || n_embd <= 0 || n_embd > std::numeric_limits<int>::max() || k < 1 || k > kMaxK || n_tok < 1 ||
         !gpu_rows || !host_rows || !dst || !count || (dst2 != nullptr && count2 == nullptr) ||
-        (dst3 != nullptr && count3 == nullptr) || !weights || !output)
+        (dst3 != nullptr && count3 == nullptr) || !weights || !output || (shared_gate != nullptr && !shared))
         throw std::invalid_argument("native MoE gather-combine requires a stream, positive width, 1..15 experts "
                                     "and its buffers");
     const dim3 grid(unsigned((n_embd + 255) / 256), unsigned(n_tok));
     gather_combine<<<grid, 256, 0, static_cast<cudaStream_t>(stream)>>>(gpu_rows, host_rows, dst, count, dst2, count2,
-                                                                         dst3, count3, weights, shared, output, n_embd,
-                                                                         int(k));
+                                                                         dst3, count3, weights, shared, shared_gate,
+                                                                         output, n_embd, int(k));
     const auto error = cudaGetLastError();
     if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error));
 }

src/kernels/cuda/native_ple_postops.cu

@@ -110,13 +110,22 @@ __global__ void conv_residual_tokens_kernel(const float* history, const float* n
     const float activation = sum / (1.0f + expf(-sum));
     result[i] = __fadd_rn(hidden[i], __fadd_rn(gated[i], activation));
 }
-// The history after the chunk: its last nine normalized rows, older rows moved up when the chunk is shorter.
-__global__ void history_tokens_kernel(float* history, const float* normalized, int n_tok) {
+// The history after the chunk: its last nine normalized rows, older rows moved up when the chunk is shorter; with
+// `snaps`, the history after each of its tokens too.
+__global__ void history_tokens_kernel(float* history, const float* normalized, int n_tok, float* snaps) {
     const int c = blockIdx.x * blockDim.x + threadIdx.x;
     if (c >= D) return;
     float old[HISTORY];
 #pragma unroll
     for (int r = 0; r < HISTORY; ++r) old[r] = history[c * HISTORY + r];
+    if (snaps != nullptr)
+        for (int t = 1; t <= n_tok; ++t)
+#pragma unroll
+            for (int r = 0; r < HISTORY; ++r) {
+                const int j = t - HISTORY + r;
+                snaps[size_t(t - 1) * HISTORY * D + c * HISTORY + r] =
+                    j >= 0 ? normalized[size_t(j) * D + c] : old[r + t];
+            }
 #pragma unroll
     for (int r = 0; r < HISTORY; ++r) {
         const int j = n_tok - HISTORY + r;
@@ -170,7 +179,8 @@ void native_ple_postops(const float* projected_key, const float* hidden,
 }
 
 void native_ple_postops_tokens(const float* projected_key, const float* hidden, const float* value, float* history,
-                               const PleWeights& w, const NativePleTokensBuffers& b, int n_tok, void* stream) {
+                               const PleWeights& w, const NativePleTokensBuffers& b, int n_tok, void* stream,
+                               float* snaps) {
     if (!stream || n_tok < 1) throw std::invalid_argument("native PLE postops: tokens need a stream and n_tok >= 1");
     const size_t rows = size_t(n_tok) * D * 4;
     const Span inputs[] = {{projected_key,rows,4}, {hidden,rows,4}, {value,size_t(n_tok)*N*4,4},
@@ -188,6 +198,14 @@ void native_ple_postops_tokens(const float* projected_key, const float* hidden,
             if (!(i == 1 && j == 4 && b.query == b.normalized) && overlaps(outputs[i], outputs[j]))
                 throw std::invalid_argument("native PLE postops writable spans overlap");
     }
+    if (snaps != nullptr) {
+        const Span s{snaps, size_t(n_tok) * HISTORY * D * 4, 4};
+        validate(s);
+        for (const auto& span : inputs)
+            if (overlaps(s, span)) throw std::invalid_argument("native PLE postops: the snapshots overlap an input");
+        for (const auto& span : outputs)
+            if (overlaps(s, span)) throw std::invalid_argument("native PLE postops: the snapshots overlap an output");
+    }
     const int R = H * n_tok;
     const unsigned blocks = unsigned((size_t(n_tok) * D + 255) / 256);
     auto st = static_cast<cudaStream_t>(stream);
@@ -198,7 +216,7 @@ void native_ple_postops_tokens(const float* projected_key, const float* hidden,
     launch_check();
     native_gr_rms_norm_weighted_repeat(b.gated,w.norm_conv,b.normalized,N,R,H,NG_RMS_EPS,stream);
     conv_residual_tokens_kernel<<<blocks,256,0,st>>>(history,b.normalized,w.conv1d_f16,hidden,b.gated,b.result,n_tok);
-    history_tokens_kernel<<<D/256,256,0,st>>>(history,b.normalized,n_tok);
+    history_tokens_kernel<<<D/256,256,0,st>>>(history,b.normalized,n_tok,snaps);
     launch_check();
 }
 } // namespace strata::kernels

src/kernels/cuda/native_qsa.cu

@@ -23,6 +23,7 @@
 // SOFTWARE.
 
 #include "strata/kernels/native_qsa.hpp"
+#include "strata/kernels/mrope.hpp"
 #include <cuda_runtime.h>
 #include <atomic>
 #include <cmath>
@@ -66,6 +67,47 @@ __global__ void norm(const float* input, const float* __restrict__ gamma, float*
     for (std::size_t col = tid; col < std::size_t(n_cols); col += BlockSize)
         output[col] = scale * input[col] * gamma[col];
 }
+// `norm<256>`, then native_rope.cu's `apply` on its output, one block a row of at most 256 columns: the normalized
+// row waits in shared memory.  Row r is read at input + r * in_stride.
+__global__ void norm_rope(const float* input, int in_stride, const float* __restrict__ gamma, float* output,
+                          int n_cols, float epsilon, int n_rot, float theta_scale, const int* __restrict__ positions,
+                          const int32_t* __restrict__ mtab, int heads, int pos_stride) {
+    constexpr int BlockSize = 256;
+    const int tid = threadIdx.x, row = blockIdx.x;
+    input += std::size_t(row) * in_stride;
+    output += std::size_t(row) * n_cols;
+    float partial = 0.0f;
+    for (std::size_t col = tid; col < std::size_t(n_cols); col += BlockSize) {
+        const float x = input[col];
+        partial += x * x;
+    }
+    __shared__ float sums[32];
+    __shared__ float values[BlockSize];
+    partial = warp_sum(partial);
+    const int lane = tid % 32;
+    if (lane == 0) sums[tid / 32] = partial;
+    __syncthreads();
+    partial = lane < BlockSize / 32 ? sums[lane] : 0.0f;
+    partial = warp_sum(partial);
+    const float mean = partial / n_cols;
+    const float scale = rsqrtf(mean + epsilon);
+    for (std::size_t col = tid; col < std::size_t(n_cols); col += BlockSize)
+        values[col] = scale * input[col] * gamma[col];
+    __syncthreads();   // every read of input precedes this: output may be input
+    const int pair = tid;
+    if (pair >= n_cols / 2) return;
+    if (pair >= n_rot / 2) {
+        output[2 * pair] = values[2 * pair];
+        output[2 * pair + 1] = values[2 * pair + 1];
+        return;
+    }
+    const int position = positions[(row / heads) * pos_stride + row % heads];
+    const float theta = mrope_pos(mtab, position, pair) * powf(theta_scale, float(pair));
+    const float c = cosf(theta), s = sinf(theta);
+    const float a = values[pair], b = values[pair + n_rot / 2];
+    output[pair] = a * c - b * s;
+    output[pair + n_rot / 2] = a * s + b * c;
+}
 __global__ void gate(const float* attn, const float* __restrict__ q_full, float* output,
                      int n_head, int head_dim) {
     const std::size_t i = std::size_t(blockIdx.x) * blockDim.x + threadIdx.x;
@@ -117,6 +159,27 @@ void native_qsa_rms_norm_weighted(const float* input, const float* gamma, float*
         norm<1024><<<unsigned(n_rows), 1024, 0, static_cast<cudaStream_t>(stream)>>>(input, gamma, output, n_cols, epsilon);
     check_launch();
 }
+void native_qsa_norm_rope_tokens(const float* input, int in_stride, const float* gamma, float* output, int n_cols,
+                                 int n_rows, float epsilon, int n_rot, float freq_base, const int* positions, int heads,
+                                 int pos_stride, void* stream) {
+    const auto count = elements(n_cols, n_rows);
+    if ((n_cols != 128 && n_cols != 256) || n_rot != 64 || in_stride < n_cols || heads < 1 || pos_stride < heads ||
+        !std::isfinite(epsilon) || epsilon < 0.0f || !std::isfinite(freq_base) || freq_base <= 1.0f || !positions)
+        throw std::invalid_argument("native QSA norm-rope requires 128 or 256 columns, 64 rotated, valid strides, "
+                                    "epsilon and base");
+    const std::size_t in_bytes = (std::size_t(n_rows - 1) * in_stride + n_cols) * 4;
+    if (!stream || !valid(input, in_bytes) || !valid(gamma, std::size_t(n_cols) * 4) || !valid(output, count * 4) ||
+        (input != output && overlap(input, in_bytes, output, count * 4)) ||
+        (input == output && in_stride != n_cols) || overlap(output, count * 4, gamma, std::size_t(n_cols) * 4))
+        throw std::invalid_argument("native QSA norm-rope requires a stream, aligned spans, and disjoint buffers or an "
+                                    "exact input/output alias");
+    // pinned host-side float powf, as native_rope_apply_tokens
+    const float theta_scale = powf(freq_base, -2.0f / n_rot);
+    norm_rope<<<unsigned(n_rows), 256, 0, static_cast<cudaStream_t>(stream)>>>(
+        input, in_stride, gamma, output, n_cols, epsilon, n_rot, theta_scale, positions, mrope_table(), heads,
+        pos_stride);
+    check_launch();
+}
 void native_qsa_gate_apply(const float* attn, const float* q_full, float* output,
                            int n_head, int head_dim, void* stream) {
     const auto count = elements(head_dim, n_head);

src/kernels/cuda/native_qsa_indexer.cu

@@ -97,11 +97,17 @@ __device__ void append_one(const float* __restrict__ raw, const int pos,
     }
     __syncthreads();   // values and partials are reused by the next call
 }
+// With `tail_snap`, the tail as it was is copied there first.
 __global__ void append(const float* __restrict__ raw, const int32_t* __restrict__ pos_dev, int pos_stride,
                        int n_tok, int pos_base, const float* __restrict__ gamma, float epsilon,
                        float* __restrict__ tail, float* __restrict__ dead,
                        float* __restrict__ pooled, int32_t* __restrict__ block_pos,
-                       int max_cells, float theta_scale, const int32_t* __restrict__ mtab) {
+                       int max_cells, float theta_scale, const int32_t* __restrict__ mtab,
+                       float* __restrict__ tail_snap) {
+    if (tail_snap != nullptr) {
+        for (int i = threadIdx.x; i < (R - 1) * D; i += blockDim.x) tail_snap[i] = tail[i];
+        __syncthreads();
+    }
     for (int t = 0; t < n_tok; ++t) {
         append_one(raw + std::size_t(t) * D, pos_dev[std::size_t(t) * pos_stride], pos_base, gamma, epsilon, tail,
                    dead, pooled, block_pos, max_cells, theta_scale, mtab);
@@ -131,7 +137,7 @@ void native_qsa_indexer_append(const float* raw, const int32_t* relative_pos_dev
 void native_qsa_indexer_append_multi(const float* raw, const int32_t* relative_pos_device, int pos_stride,
                                      int n_tok, int32_t pos_base, const float* gamma, float epsilon,
                                      const QsaIndexerBuffers& b, const QsaShapes& s, int64_t max_cells,
-                                     float freq_base, void* stream) {
+                                     float freq_base, void* stream, float* tail_snap) {
     if (n_tok < 1 || pos_stride < 0 || (n_tok > 1 && pos_stride < 1))
         throw std::invalid_argument("native QSA indexer: n_tok >= 1 appends with a positive position stride");
     if (!stream || s.idx_dim != D || s.idx_block != R || s.n_rot != ROT ||
@@ -145,9 +151,15 @@ void native_qsa_indexer_append_multi(const float* raw, const int32_t* relative_p
     for (const auto& span : spans) validate(span);
     for (int i = 0; i < 7; ++i) for (int j = i + 1; j < 7; ++j)
         if (overlaps(spans[i], spans[j])) throw std::invalid_argument("native QSA indexer buffers overlap");
+    if (tail_snap != nullptr) {
+        const Span snap{tail_snap, (R - 1) * D * 4};
+        validate(snap);
+        for (const auto& span : spans)
+            if (overlaps(snap, span)) throw std::invalid_argument("native QSA indexer: the tail snapshot overlaps");
+    }
     const float theta_scale = powf(freq_base, -2.0f / ROT);
     append<<<1,THREADS,0,static_cast<cudaStream_t>(stream)>>>(raw,relative_pos_device,pos_stride,n_tok,pos_base,gamma,
-        epsilon,b.tail,b.dead,b.pooled,b.block_pos,int(max_cells),theta_scale,mrope_table());
+        epsilon,b.tail,b.dead,b.pooled,b.block_pos,int(max_cells),theta_scale,mrope_table(),tail_snap);
     const auto error = cudaGetLastError();
     if (error != cudaSuccess) throw std::runtime_error(cudaGetErrorString(error));
 }

src/kernels/cuda/ple.cu

@@ -197,6 +197,7 @@ void ck(cudaError_t e, const char* what) {
 }  // namespace
 
 void ple_set_native_bf16(bool enabled) { native_bf16 = enabled; }
+bool ple_native_bf16() { return native_bf16; }
 void ple_set_native_postops(bool enabled) { native_postops = enabled; }
 bool ple_native_postops_enabled() { return native_postops; }
 

src/kernels/cuda/shared_expert.cu

@@ -160,12 +160,22 @@ __global__ void scale_rows_kernel(float* __restrict__ out, const float* __restri
 }
 }  // namespace
 
-void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw,
-                         const uint16_t* gate_inp_bf16, float* gate, float* up, float* g, float* out, int64_t n_embd,
-                         int64_t n_ff, void* stream) {
-    if (n_tok < 1 || n_tok > 8 || !nw.q8_1 || !nw.gate_data || !nw.up_data || !nw.down_data || !stream)
-        throw std::invalid_argument("shared_expert_multi: needs 1..8 tokens, native weights, scratch and a stream");
+void shared_expert_gate_rows(float* out, float* g, int64_t n_embd, int n_tok, void* stream) {
     cudaStream_t cs = (cudaStream_t) stream;
+    native_scalar_sigmoid_kernel<<<1, (unsigned) n_tok, 0, cs>>>(g);
+    scale_rows_kernel<<<dim3((unsigned) ((n_embd + THREADS - 1) / THREADS), (unsigned) n_tok), THREADS, 0, cs>>>(
+        out, g, (int) n_embd);
+}
+
+void shared_expert_multi(int n_tok, const float* x, const NativeSharedWeights& nw, const uint16_t* gate_inp_bf16,
+                         float* gate, float* up, float* g, float* out, int64_t n_embd, int64_t n_ff, void* stream,
+                         void* gate_stream) {
+    if (n_tok < 1 || n_tok > 8 || !nw.q8_1 || !nw.gate_data || !nw.up_data || !nw.down_data || !stream ||
+        !gate_stream || !native_bf16)
+        throw std::invalid_argument("shared_expert_multi: needs 1..8 tokens, native weights and gate, scratch and "
+                                    "streams");
+    cudaStream_t cs = (cudaStream_t) stream;
+    bf16_gemv_fp32_mmvf_multi(x, gate_inp_bf16, g, n_embd, 1, n_tok, gate_stream);
     native_quantize_q8_1(x, nw.q8_1, (int) n_embd, n_tok, stream);
     native_mmvq(nw.gate_type, nw.gate_data, nw.q8_1, gate, (int) n_embd, (int) n_ff, n_tok, stream);
     native_mmvq(nw.up_type, nw.up_data, nw.q8_1, up, (int) n_embd, (int) n_ff, n_tok, stream);
@@ -173,15 +183,6 @@ void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, cons
     native_swiglu_kernel<<<(unsigned) ((n + THREADS - 1) / THREADS), THREADS, 0, cs>>>(gate, up, gate, n);
     native_quantize_q8_1(gate, nw.q8_1, (int) n_ff, n_tok, stream);
     native_mmvq(nw.down_type, nw.down_data, nw.q8_1, out, (int) n_ff, (int) n_embd, n_tok, stream);
-    if (native_bf16) {
-        bf16_gemv_fp32_mmvf_multi(x, gate_inp_bf16, g, n_embd, 1, n_tok, stream);
-        native_scalar_sigmoid_kernel<<<1, (unsigned) n_tok, 0, cs>>>(g);
-    } else {
-        for (int t = 0; t < n_tok; ++t)
-            scalar_gate_kernel<<<1, 256, 0, cs>>>(x_bf16 + (size_t) t * n_embd, gate_inp_bf16, g + t, (int) n_embd);
-    }
-    scale_rows_kernel<<<dim3((unsigned) ((n_embd + THREADS - 1) / THREADS), (unsigned) n_tok), THREADS, 0, cs>>>(
-        out, g, (int) n_embd);
     const cudaError_t e = cudaGetLastError();
     if (e != cudaSuccess) throw std::runtime_error(std::string("shared_expert_multi: ") + cudaGetErrorString(e));
 }
@@ -295,11 +296,11 @@ void shared_expert(const uint8_t* x_q8_0, const uint8_t* x_q8k, const uint16_t*
     // round trip.  256 threads is the reduction's width, not the problem's size.
     if (use_native) {
         bf16_gemv_fp32_mmvf(x_f32, gate_inp_bf16, g, n_embd, 1, stream);
-        native_scalar_sigmoid_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(g);
+        shared_expert_gate_rows(out, g, n_embd, 1, stream);
     } else {
         scalar_gate_kernel<<<1, 256, 0, (cudaStream_t) stream>>>(x_bf16, gate_inp_bf16, g, (int) n_embd);
+        scale_kernel<<<g_embd, THREADS, 0, (cudaStream_t) stream>>>(out, g, (int) n_embd);
     }
-    scale_kernel<<<g_embd, THREADS, 0, (cudaStream_t) stream>>>(out, g, (int) n_embd);
 
     if (stream == nullptr) {
         const cudaError_t e = cudaDeviceSynchronize();

src/kernels/verify_parity.cpp

@@ -13,11 +13,13 @@
 #include "strata/kernels/kv_q8.hpp"
 #include "strata/kernels/native_moe.hpp"
 #include "strata/kernels/native_ple_postops.hpp"
+#include "strata/kernels/native_qsa.hpp"
 #include "strata/kernels/native_qsa_indexer.hpp"
 #include "strata/kernels/native_rope.hpp"
 #include "strata/kernels/qsa_decode_attn.hpp"
 #include "strata/kernels/native_router.hpp"
 #include "strata/kernels/s2_expert_grouped.hpp"
+#include "strata/kernels/shared_expert.hpp"
 #include "strata/kernels/verify_kernels.hpp"
 #include "strata/prefill/kernels.hpp"
 
@@ -81,7 +83,8 @@ int bitwise_diff(const std::vector<float>& a, const std::vector<float>& b, const
 }
 
 // ---- the MoE combine: copy all rows + add the hits + one combine per token  vs  the second GPU's rows taken into VRAM
-// (fetch_listed_rows) and native_moe_gather_combine
+// (fetch_listed_rows) and native_moe_gather_combine; from rep 2 on the shared rows gated by shared_expert_gate_rows vs
+// the gather-combine gating them
 int test_gather_combine(std::mt19937& rng, cudaStream_t s) {
     const int64_t N = 2560, K = 10;
     int bad = 0;
@@ -94,6 +97,11 @@ int test_gather_combine(std::mt19937& rng, cudaStream_t s) {
             for (auto& v : host) v = edgy(rng);
             for (auto& v : w) v = std::uniform_real_distribution<float>(0.0f, 1.0f)(rng);
             for (auto& v : shared) v = edgy(rng);
+            std::vector<float> glog((size_t) n_tok);   // the gates' logits: sigmoid near 1, near 0 (denormal products)
+            for (auto& v : glog)
+                v = rng() % 4 == 0 ? std::uniform_real_distribution<float>(-100.0f, -80.0f)(rng)
+                                   : std::uniform_real_distribution<float>(-8.0f, 8.0f)(rng);
+            const bool gated = rep >= 2;
             // the GPU's entries, in a random order and a random share (none, some, all); from rep 4 on the second GPU
             // takes some of the rest
             std::vector<int32_t> dst, raw;
@@ -120,12 +128,16 @@ int test_gather_combine(std::mt19937& rng, cudaStream_t s) {
 
             float *d_hit = dev<float>(hit.size()), *d_parts = dev<float>(hit.size()), *d_w = dev<float>(w.size()),
                   *d_sh = dev<float>(shared.size()), *d_out_old = dev<float>(shared.size()),
-                  *d_out_new = dev<float>(shared.size());
+                  *d_out_new = dev<float>(shared.size()), *d_shg = dev<float>(shared.size()),
+                  *d_g = dev<float>(glog.size()), *d_glog = dev<float>(glog.size());
             int32_t *d_dst = dev<int32_t>(std::max<size_t>(1, dst.size()) + 64), *d_count = dev<int32_t>(3),
                     *d_list = dev<int32_t>((size_t) (4 + rows));
             up(d_hit, hit);
             up(d_w, w);
             up(d_sh, shared);
+            up(d_shg, shared);
+            up(d_g, glog);
+            up(d_glog, glog);
             if (!dst.empty()) up(d_dst, dst);
             // the new path takes them as two lists (the window's VRAM share and the pool's PCIe share), split at random
             const int32_t split = rep % 2 == 0 ? count : (int32_t) (rng() % (dst.size() + 1));
@@ -144,9 +156,10 @@ int test_gather_combine(std::mt19937& rng, cudaStream_t s) {
             std::memcpy(h_map, host_zeroed.data(), host.size() * sizeof(float));
             strata::kernels::copy_from_mapped(d_parts, m_map, rows * N, s);
             strata::kernels::moe_hit_add(d_parts, d_hit, d_dst, d_count, rows, N, s);
+            if (gated) strata::kernels::shared_expert_gate_rows(d_shg, d_g, N, n_tok, s);
             for (int t = 0; t < n_tok; ++t)
-                strata::kernels::native_moe_combine(d_parts + t * K * N, d_w + t * K, d_sh + t * N, d_out_old + t * N,
-                                                    N, K, s);
+                strata::kernels::native_moe_combine(d_parts + t * K * N, d_w + t * K, (gated ? d_shg : d_sh) + t * N,
+                                                    d_out_old + t * N, N, K, s);
             check(cudaStreamSynchronize(s), "old path");
             // new
             std::memcpy(h_map, host_poison.data(), host.size() * sizeof(float));
@@ -156,7 +169,7 @@ int test_gather_combine(std::mt19937& rng, cudaStream_t s) {
             strata::kernels::native_moe_gather_combine(d_hit, m_map, d_dst, d_count + 1,
                                                        rep % 2 == 0 ? nullptr : d_dst + split, d_count + 2,
                                                        second ? d_list + 4 : nullptr, second ? d_list : nullptr, d_w,
-                                                       d_sh, d_out_new, N, K, n_tok, s);
+                                                       d_sh, gated ? d_glog : nullptr, d_out_new, N, K, n_tok, s);
             check(cudaStreamSynchronize(s), "new path");
             int b = bitwise_diff(down(d_out_old, shared.size()), down(d_out_new, shared.size()), "gather_combine");
             if (second) {
@@ -174,13 +187,15 @@ int test_gather_combine(std::mt19937& rng, cudaStream_t s) {
             cudaFreeHost(h_map);
             cudaFreeHost(h_list);
             for (void* p : {(void*) d_hit, (void*) d_parts, (void*) d_w, (void*) d_sh, (void*) d_out_old,
-                            (void*) d_out_new, (void*) d_dst, (void*) d_count, (void*) d_list})
+                            (void*) d_out_new, (void*) d_shg, (void*) d_g, (void*) d_glog, (void*) d_dst,
+                            (void*) d_count, (void*) d_list})
                 cudaFree(p);
         }
     }
     std::printf("gather_combine: %s\n",
                 bad ? "MISMATCH"
-                    : "bitwise equal (n_tok 1-4, 8; 8 plans each, one or two lists, the second GPU's rows fetched)");
+                    : "bitwise equal (n_tok 1-4, 8; 8 plans each, one or two lists, the second GPU's rows fetched, "
+                      "the shared rows gated)");
     return bad;
 }
 
@@ -257,40 +272,52 @@ int test_hit_plan(std::mt19937& rng, cudaStream_t s) {
     return bad;
 }
 
-// ---- BF16 x FP32 MMVF: one call per row  vs  bf16_gemv_fp32_mmvf_multi (the router, indexer and gate shapes)
+// ---- BF16 x FP32 MMVF: one call per row  vs  bf16_gemv_fp32_mmvf_multi (the router, indexer and gate shapes), and
+// with a second matrix of 125 rows in the same launch (the indexer's keys and queries; partial groups of 4 rows)
 int test_mmvf_multi(std::mt19937& rng, cudaStream_t s) {
-    int bad = 0;
-    for (int64_t n_out : {512, 128, 1}) {
-        const int64_t n_in = 2560;
-        std::vector<uint16_t> w((size_t) (n_in * n_out));
+    auto bf16_matrix = [&](int64_t n) {
+        std::vector<uint16_t> w((size_t) n);
         for (auto& v : w) {
             const float f = edgy(rng);
             uint32_t b;
             std::memcpy(&b, &f, 4);
             v = (uint16_t) (b >> 16);
         }
-        uint16_t* d_w = dev<uint16_t>(w.size());
-        up(d_w, w);
+        uint16_t* d = dev<uint16_t>(w.size());
+        up(d, w);
+        return d;
+    };
+    int bad = 0;
+    const int64_t n_in = 2560, n_out2 = 125;
+    uint16_t* d_w2 = bf16_matrix(n_in * n_out2);
+    for (int64_t n_out : {512, 130, 1}) {
+        uint16_t* d_w = bf16_matrix(n_in * n_out);
         for (int n_tok = 1; n_tok <= 8; ++n_tok) {
             std::vector<float> x((size_t) (n_tok * n_in));
             for (auto& v : x) v = edgy(rng);
-            float *d_x = dev<float>(x.size()), *d_a = dev<float>((size_t) (n_tok * n_out)),
-                  *d_b = dev<float>((size_t) (n_tok * n_out));
+            const size_t ny = (size_t) (n_tok * n_out), ny2 = (size_t) (n_tok * n_out2);
+            float *d_x = dev<float>(x.size()), *d_a = dev<float>(ny), *d_b = dev<float>(ny), *d_c = dev<float>(ny),
+                  *d_a2 = dev<float>(ny2), *d_c2 = dev<float>(ny2);
             up(d_x, x);
-            for (int t = 0; t < n_tok; ++t)
+            for (int t = 0; t < n_tok; ++t) {
                 strata::kernels::bf16_gemv_fp32_mmvf(d_x + t * n_in, d_w, d_a + t * n_out, n_in, n_out, s);
+                strata::kernels::bf16_gemv_fp32_mmvf(d_x + t * n_in, d_w2, d_a2 + t * n_out2, n_in, n_out2, s);
+            }
             strata::kernels::bf16_gemv_fp32_mmvf_multi(d_x, d_w, d_b, n_in, n_out, n_tok, s);
+            strata::kernels::bf16_gemv_fp32_mmvf_multi(d_x, d_w, d_c, n_in, n_out, n_tok, s, d_w2, d_c2, n_out2);
             check(cudaStreamSynchronize(s), "mmvf");
-            const int b = bitwise_diff(down(d_a, (size_t) (n_tok * n_out)), down(d_b, (size_t) (n_tok * n_out)), "mmvf_multi");
+            const int b = bitwise_diff(down(d_a, ny), down(d_b, ny), "mmvf_multi") +
+                          bitwise_diff(down(d_a, ny), down(d_c, ny), "mmvf_multi, two matrices") +
+                          bitwise_diff(down(d_a2, ny2), down(d_c2, ny2), "mmvf_multi, the second matrix");
             if (b) std::fprintf(stderr, "mmvf_multi: n_out %lld n_tok %d: %d differ\n", (long long) n_out, n_tok, b);
             bad += b;
-            cudaFree(d_x);
-            cudaFree(d_a);
-            cudaFree(d_b);
+            for (float* p : {d_x, d_a, d_b, d_c, d_a2, d_c2}) cudaFree(p);
         }
         cudaFree(d_w);
     }
-    std::printf("mmvf_multi: %s\n", bad ? "MISMATCH" : "bitwise equal (n_out 512, 128, 1; 1-8 rows)");
+    cudaFree(d_w2);
+    std::printf("mmvf_multi: %s\n",
+                bad ? "MISMATCH" : "bitwise equal (n_out 512, 130, 1; 1-8 rows; with a second matrix)");
     return bad;
 }
 
@@ -765,6 +792,60 @@ int test_rope_tokens(std::mt19937& rng, cudaStream_t s) {
     return bad;
 }
 
+// ---- the QSA heads' norm and RoPE: native_qsa_rms_norm_weighted + native_rope_apply_tokens  vs
+// native_qsa_norm_rope_tokens, in place and from rows of twice the width (the queries' halves)
+int test_norm_rope(std::mt19937& rng, cudaStream_t s) {
+    const int NH = 24;
+    int bad = 0;
+    for (int heads : {2, 4, 24}) {
+        for (int head_dim : {128, 256}) {
+            std::vector<float> gamma((size_t) head_dim);
+            for (auto& g : gamma) g = 0.5f + std::uniform_real_distribution<float>(0.0f, 1.0f)(rng);
+            float* d_gamma = dev<float>(gamma.size());
+            up(d_gamma, gamma);
+            for (int n_tok = 1; n_tok <= 8; ++n_tok) {
+                const int rows = n_tok * heads;
+                std::vector<float> x((size_t) rows * 2 * head_dim);   // rows of 2 * head_dim: the first half is read
+                for (auto& v : x) v = std::normal_distribution<float>(0.0f, 3.0f)(rng);
+                std::vector<float> half((size_t) rows * head_dim);
+                for (int r = 0; r < rows; ++r)
+                    std::copy(x.begin() + (size_t) r * 2 * head_dim, x.begin() + (size_t) r * 2 * head_dim + head_dim,
+                              half.begin() + (size_t) r * head_dim);
+                std::vector<int32_t> pos((size_t) n_tok * NH);
+                const int p0 = (int) (rng() % 200000);
+                for (int t = 0; t < n_tok; ++t)
+                    for (int h = 0; h < NH; ++h) pos[(size_t) t * NH + h] = p0 + t;
+                float *d_a = dev<float>(half.size()), *d_b = dev<float>(half.size()), *d_x = dev<float>(x.size()),
+                      *d_c = dev<float>(half.size());
+                int32_t* d_pos = dev<int32_t>(pos.size());
+                up(d_a, half);
+                up(d_b, half);
+                up(d_x, x);
+                up(d_pos, pos);
+                strata::kernels::native_qsa_rms_norm_weighted(d_a, d_gamma, d_a, head_dim, rows, 1e-6f, s);
+                strata::kernels::native_rope_apply_tokens(d_a, d_a, rows, head_dim, 64, 5000000.0f, d_pos, heads, NH,
+                                                          s);
+                strata::kernels::native_qsa_norm_rope_tokens(d_b, head_dim, d_gamma, d_b, head_dim, rows, 1e-6f, 64,
+                                                             5000000.0f, d_pos, heads, NH, s);
+                strata::kernels::native_qsa_norm_rope_tokens(d_x, 2 * head_dim, d_gamma, d_c, head_dim, rows, 1e-6f, 64,
+                                                             5000000.0f, d_pos, heads, NH, s);
+                check(cudaStreamSynchronize(s), "norm_rope");
+                const int b = bitwise_diff(down(d_a, half.size()), down(d_b, half.size()), "norm_rope") +
+                              bitwise_diff(down(d_a, half.size()), down(d_c, half.size()), "norm_rope, strided");
+                if (b)
+                    std::fprintf(stderr, "norm_rope: %d heads of %d, n_tok %d: %d differ\n", heads, head_dim, n_tok, b);
+                bad += b;
+                for (float* p : {d_a, d_b, d_x, d_c}) cudaFree(p);
+                cudaFree(d_pos);
+            }
+            cudaFree(d_gamma);
+        }
+    }
+    std::printf("norm_rope_tokens: %s\n", bad ? "MISMATCH" : "bitwise equal (2, 4, 24 heads of 128, 256; 1-8 tokens; "
+                                                           "in place and strided)");
+    return bad;
+}
+
 // ---- the int8 KV append: kv_append_q8_step per token  vs  kv_append_q8_steps
 int test_kv_append(std::mt19937& rng, cudaStream_t s) {
     const strata::kernels::QsaShapes sh = strata::kernels::qsa_real_shapes();
@@ -867,9 +948,15 @@ int test_indexer_append(std::mt19937& rng, cudaStream_t s) {
                 for (int t = 0; t < n_tok; ++t)
                     strata::kernels::native_qsa_indexer_append(d_raw + (size_t) t * D, d_pos + t, 0, d_gamma, 1e-6f, ib0, sh,
                                                                max_cells, 5000000.0f, s);
+                float* d_snap = dev<float>(tail.size());
                 strata::kernels::native_qsa_indexer_append_multi(d_raw, d_pos, 1, n_tok, 0, d_gamma, 1e-6f, ib1, sh,
-                                                                 max_cells, 5000000.0f, s);
+                                                                 max_cells, 5000000.0f, s, d_snap);
                 check(cudaStreamSynchronize(s), "indexer");
+                if (down(d_snap, tail.size()) != tail) {
+                    std::fprintf(stderr, "indexer: n_tok %d start %d: the tail snapshot differs\n", n_tok, start);
+                    ++bad;
+                }
+                cudaFree(d_snap);
                 const size_t sizes[3] = {tail.size(), dead.size(), pooled.size()};
                 const char* names[3] = {"tail", "dead", "pooled"};
                 for (int o = 0; o < 3; ++o) {
@@ -891,7 +978,8 @@ int test_indexer_append(std::mt19937& rng, cudaStream_t s) {
         }
     }
     cudaFree(d_gamma);
-    std::printf("indexer_append_multi: %s\n", bad ? "MISMATCH" : "bitwise equal (1-8 tokens, 8 starts, rejected cells)");
+    std::printf("indexer_append_multi: %s\n",
+                bad ? "MISMATCH" : "bitwise equal (1-8 tokens, 8 starts, rejected cells), the tail snapshot the tail");
     return bad;
 }
 
@@ -923,8 +1011,9 @@ int test_ple_tokens(std::mt19937& rng, cudaStream_t s) {
         for (auto& x : hist) x = nd(rng);
         float *d_key = dev<float>(key.size()), *d_hidden = dev<float>(hidden.size()), *d_value = dev<float>(value.size());
         up(d_key, key); up(d_hidden, hidden); up(d_value, value);
-        // per token: the decode path's calls
-        float *h0 = dev<float>(hist.size()), *r0 = dev<float>(hidden.size());
+        // per token: the decode path's calls, the history after each token kept
+        float *h0 = dev<float>(hist.size()), *r0 = dev<float>(hidden.size()),
+              *s0 = dev<float>((size_t) n_tok * hist.size());
         float *k1 = dev<float>(D), *q1 = dev<float>(D), *g1 = dev<float>(H), *gd1 = dev<float>(D), *c1 = dev<float>(D);
         up(h0, hist);
         for (int t = 0; t < n_tok; ++t) {
@@ -932,22 +1021,30 @@ int test_ple_tokens(std::mt19937& rng, cudaStream_t s) {
             strata::kernels::native_ple_postops(d_key + (size_t) t * D, d_hidden + (size_t) t * D,
                                                 d_value + (size_t) t * N, h0, w, b, s);
             strata::kernels::ple_history_advance(h0, q1, s);
+            check(cudaMemcpyAsync(s0 + (size_t) t * hist.size(), h0, hist.size() * 4, cudaMemcpyDeviceToDevice, s),
+                  "ple snapshot");
         }
         // the chunk at once
         float *h2 = dev<float>(hist.size()), *r2 = dev<float>(hidden.size()), *k2 = dev<float>(key.size()),
-              *q2 = dev<float>(key.size()), *g2 = dev<float>((size_t) n_tok * H), *gd2 = dev<float>(key.size());
+              *q2 = dev<float>(key.size()), *g2 = dev<float>((size_t) n_tok * H), *gd2 = dev<float>(key.size()),
+              *s2 = dev<float>((size_t) n_tok * hist.size());
         up(h2, hist);
-        strata::kernels::native_ple_postops_tokens(d_key, d_hidden, d_value, h2, w, {k2, q2, g2, gd2, q2, r2}, n_tok, s);
+        strata::kernels::native_ple_postops_tokens(d_key, d_hidden, d_value, h2, w, {k2, q2, g2, gd2, q2, r2}, n_tok, s,
+                                                   s2);
         check(cudaStreamSynchronize(s), "ple tokens");
         const int br = bitwise_diff(down(r0, hidden.size()), down(r2, hidden.size()), "result");
         const int bh = bitwise_diff(down(h0, hist.size()), down(h2, hist.size()), "history");
-        if (br || bh) std::fprintf(stderr, "ple tokens: n_tok %d: result or history differ\n", n_tok);
-        bad += br + bh;
-        for (float* p : {d_key, d_hidden, d_value, h0, r0, k1, q1, g1, gd1, c1, h2, r2, k2, q2, g2, gd2}) cudaFree(p);
+        const int bs = bitwise_diff(down(s0, (size_t) n_tok * hist.size()), down(s2, (size_t) n_tok * hist.size()),
+                                    "histories");
+        if (br || bh || bs) std::fprintf(stderr, "ple tokens: n_tok %d: results or histories differ\n", n_tok);
+        bad += br + bh + bs;
+        for (float* p : {d_key, d_hidden, d_value, h0, r0, s0, k1, q1, g1, gd1, c1, h2, r2, k2, q2, g2, gd2, s2})
+            cudaFree(p);
     }
     for (float* p : {d_nk, d_nq, d_nc}) cudaFree(p);
     cudaFree(d_taps);
-    std::printf("ple_postops_tokens: %s\n", bad ? "MISMATCH" : "bitwise equal (1-23 tokens, results and history)");
+    std::printf("ple_postops_tokens: %s\n",
+                bad ? "MISMATCH" : "bitwise equal (1-23 tokens, results, history, the history after each token)");
     return bad;
 }
 
@@ -1130,6 +1227,7 @@ int main(int argc, char** argv) {
     bad += test_router_rows(rng, s);
     bad += test_gr_read(rng, s);
     bad += test_rope_tokens(rng, s);
+    bad += test_norm_rope(rng, s);
     bad += test_kv_append(rng, s);
     bad += test_indexer_append(rng, s);
     bad += test_ple_tokens(rng, s);