Merge branch 'f-139b-spd' into f-139-rel

# Conflicts:
#	src/core/expert_source.cpp
#	src/kernels/cuda/iq_kernels.cu

CMakeLists.txt

@@ -415,6 +415,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   # library is the honest shape of the dependency graph: core knows about devices and weights, kernels know
   # about arithmetic, and this one knows the ORDER the arithmetic happens in.
   add_library(strata_engine STATIC src/core/layer.cpp src/core/session.cpp src/core/expert_source.cpp src/core/remote_experts.cpp
+    src/core/remote_expert_opt.cu
     src/core/expert_cache.cpp src/core/native_head.cpp src/core/native_dense.cpp src/core/verify.cpp src/core/mtp.cpp
     src/core/conversation_snapshot.cpp src/core/conversation_state.cpp src/core/conversation_memory.cpp src/core/peer_experts.cpp)
   target_include_directories(strata_engine PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
@@ -762,6 +763,13 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   target_link_libraries(qsa_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target})
   add_test(NAME qsa_parity COMMAND qsa_parity --selftest)
 
+  # the QSA top-k the engine dispatches against the original kernel: identical ids on both sides of the register
+  # kernel's reach, with tied and with equal scores (synthetic, no model)
+  add_executable(qsa_topk_parity src/kernels/qsa_topk_parity.cpp)
+  target_include_directories(qsa_topk_parity PRIVATE ${_strata_gpu_include_directories})
+  target_link_libraries(qsa_topk_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target})
+  add_test(NAME qsa_topk_parity COMMAND qsa_topk_parity --selftest)
+
   # ---- P2.S4: the PLE n-gram hash, the IQ4_NL table read, and the PLE block ----
   # `ngram.cpp` is host-only but lives in `strata_kernels` because `ple.cu` shares its geometry constants, and
   # splitting them across two libraries would only add a layer.
@@ -1181,6 +1189,15 @@ if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS AND EXISTS "${CMAKE_CURRENT_SOURCE_
 endif()
 
 if(STRATA_BUILD_TESTS)
+  if(WIN32)
+    add_executable(pool_affinity_test src/kernels/cpu/pool_affinity_test.cpp)
+    target_link_libraries(pool_affinity_test PRIVATE strata_kernels_cpu)
+    add_test(NAME pool_affinity_test COMMAND pool_affinity_test)
+  elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    add_executable(pool_affinity_test src/kernels/cpu/pool_affinity_linux_test.cpp)
+    target_link_libraries(pool_affinity_test PRIVATE strata_kernels_cpu)
+    add_test(NAME pool_affinity_test COMMAND pool_affinity_test)
+  endif()
   # Plan v0.3 P6: one expert applied to several tokens, bitwise vs the single-token kernel (synthetic, CPU only).
   # plan v0.3 P5: the device dequantizers vs the artifact's CPU dequantizers on the real shards (not a ctest: needs
   # the 38 GB model; run it by hand or from tools/gpu_session.py)

docs/DETAILS.md

@@ -1031,7 +1031,19 @@ the document, +0.4% on the chat. Details: `bench/results/2026-09-27-esp/`.
   0.1.7) drafts up to 5 tokens from the earlier copy, but only where its measured acceptance and cost say it pays:
   code edits 6-11% faster, other text unchanged. The drafts are checked like the MTP's, so the output is the same.
 - **Prompts** are processed in chunks of up to 8,192 tokens (`--prefill auto`; 32,768 opt-in) with the experts
-  streamed to the GPU over PCIe.
+  streamed to the GPU over PCIe. Since engine 0.1.39b (#583) the streamed ring is sized in bytes for the pack (a
+  native pack with bigger blobs gets fewer ring slots and keeps more cache slots) and `--prefill auto` picks the
+  largest chunk that keeps that ring full. It only gives up ring slots where 0.1.39's rule held the chunk at 4,096
+  tokens or less; from 6,144 on it keeps 0.1.39's ring and only looks for a bigger chunk beside it (shrinking the
+  ring there measured slower: the Coder -12%), and a prompt that fits 0.1.39's chunk keeps 0.1.39's ring (one
+  chunk: a smaller ring only slowed it). RTX 5070, 32K prompts against `STRATA_RING_BYTES=0`, 2-3 interleaved
+  rounds: IQ3_XXS +18.5% at a 1,500-slot cache and +9% in another memory state, the Coder +3%, IQ3_S unchanged; 4K
+  prompts unchanged (-0.7% to +1.1%, bit-identical). Q2_0 (a fixed `--prefill 2048`) is not affected. It changes
+  the bits of a long prompt against 0.1.39 (the experts go through a different
+  mix of cached and streamed groups). Measured quality, teacher-forced over the next 2,001 tokens of a long document
+  against the FP16 prompt path (IQ3_XXS, `--prefill auto`): 8K prompt KL 0.042 (0.1.39: 0.054), top-1 agreement
+  93.5% (92.2%); 32K prompt KL 0.020 (0.019), top-1 95.3% (95.0%) - the same band as before. `STRATA_RING_BYTES=0`
+  restores 0.1.39's ring, loan and chunk choice.
 
 The full story, with measurements, bottlenecks and what comes next: **[docs/paper/Strata-Paper.pdf](paper/Strata-Paper.pdf)**.
 

docs/MULTI_GPU.md

@@ -76,6 +76,20 @@ VRAM keep its own prompt buffers - the same output as 0.1.31, measured on an R97
 differs from the default's (stable and coherent); `STRATA_SPLIT_OWN=auto` does that only where the buffers are at
 most 12% of each card's VRAM.
 
+**The idle card can help one-chunk prompts (opt-in, `STRATA_PREFILL_HELP=1`).** A prompt that fits one chunk runs the stages one after the other, so while
+one card reads its layers the other idles. With it on, each stage hands a share of its streamed experts to the idle card: it
+streams them over its own PCIe link into its own (lent) prompt buffers, computes their rows on the MMQ path and sends
+them back - `--peer-device`'s peer streaming, without P2P (the activations and the rows go through mapped host memory,
+read by copy kernels, so they do not queue behind the expert blobs on either card's copy engine). The share falls with
+the prompt (0.41 of the streamed experts at 1.5K tokens, 0.32 at 3K) and is off from ~3.3K tokens, where the stages
+overlap anyway and no share paid. Measured on 2x RTX 3090 (UD-Q4_K_XL, no P2P), prompt tok/s without / with: 1.1K 496 /
+723, 1.5K 674 / 833, 2K 878 / 1,128, 2.5K 1,017 / 1,296, 3K 1,260 / 1,440; 4K and 8K unchanged, decode unchanged. It
+costs no VRAM (the idle card's own prompt buffers) and ~110 KB of mapped host memory per token of the largest chunk it
+helped (~360 MB at 3.3K tokens). The rows it computes round like a different MMQ grouping, so the output is not
+bit-identical to the default's (it is repeatable: same prompt, same output), which is why it is **opt-in**:
+`STRATA_PREFILL_HELP=1` turns it on. Native packs on the MMQ prompt path only (not with the fused prompt kernels,
+`STRATA_PF_FUSED=1`, nor with `--peer-device`). With it on, `STRATA_PREFILL_HELP_FRAC=f` fixes the share.
+
 The engine flags behind it: `--layer-split K1[,K2..]|auto` and `--split-device D1[,D2..]` (the later stages'
 devices; default the next visible ones). `--layer-split K --split-device 0` runs both stages on one card sharing
 everything - the bit-exact check of the hand-off, not a speed mode.

docs/SECOND_GPU.md

@@ -84,6 +84,44 @@ The existing warning about the CUDA0 expert-cache GPU hit path still applies:
 its outputs diverge from cache-off runs. Treat performance as experimental
 until the generated tokens have been validated.
 
+## Optional helper decode optimization
+
+`--remote-expert-opt` (`--serve` only) optimizes the CUDA1-3 helper caches
+above. The engine's default is off; since 0.1.39b setup adds it to a config on
+two or more GPUs (`--gpus`, or "use both" at start). It acts only when a helper
+cache is configured; a layer split runs exactly as before. To leave it out:
+setup's `--no-remote-expert-opt`, or `"remote_expert_opt": false` in the
+model's `strata-*.json` (kept when setup runs again). Measured by the PR's
+author: dual RTX 4090 +63% mixed / +132% code decode over the plain helper
+path; RTX 5090 + 4090 +28% / +63%. The primary cache avoids admitting experts already held
+by a helper, and helpers replace cold experts with frequently routed CPU
+misses using their existing same-layer slots. Each helper reduces its expert
+outputs to a weighted partial sum on its GPU before returning one vector per
+token. Tokens with no CPU expert work skip CPU activation quantization.
+
+For an existing server configuration with CUDA0 as the primary and CUDA1 as a
+helper, use these engine arguments alongside the model and profile arguments:
+
+```text
+--expert-cache auto --expert-cache-device1 auto --remote-expert-opt
+```
+
+`--expert-cache-device1`, `--expert-cache-device2` and
+`--expert-cache-device3` now also accept `auto`: fill each helper from its
+assigned ranking using actual aligned expert bytes and free VRAM, retaining
+the existing 512 MiB allowance. Explicit numeric budgets still work. This is
+startup capacity sizing, not throughput balancing; with several helpers a
+card's truncated candidate tail is not redistributed to another card.
+
+The optimization uses the existing host scheduling and pinned-host transport,
+not P2P or tensor parallelism. Each layer still waits for its participating
+helpers. It changes floating-point summation order, so enabled output is not
+claimed to be bitwise identical. Only two-card CUDA operation has been measured;
+three/four cards and HIP have not been validated. Without the switch, the
+existing decode path remains in use. This does not optimize the separate
+`--peer-device` path below or change its existing incompatibility with helper
+caches.
+
 ## Peer tier (`--peer-device`)
 
 `--peer-device N` puts a second adaptive expert cache on CUDA device N. It

include/strata/core/expert_cache.hpp

@@ -135,6 +135,13 @@ public:
     /// sweep) and a placeholder policy would set the hit rate that everything downstream is then sized against.
     int32_t admit(int64_t layer, int64_t expert);
 
+    /// Publish a same-layer replacement after its slot copy has completed.
+    void replace(int64_t layer, int32_t old_expert, int32_t new_expert) {
+        auto& old = residency_[(size_t) layer * n_expert_ + old_expert];
+        residency_[(size_t) layer * n_expert_ + new_expert] = old;
+        old = kNotResident;
+    }
+
     /// **R4.2g: GIVE EACH LAYER ITS OWN SLOTS.  ROUND 328 MEASURED WHY THE GLOBAL FORM CANNOT WORK.**
     ///
     /// `admit`'s original policy hands out slots in ARRIVAL ORDER from one counter shared by every layer

include/strata/core/mtp.hpp

@@ -82,6 +82,7 @@ public:
     /// Coupled mode with penalties: the history the next draft() chain starts from - what the session holds after the
     /// window's commit (`tail`) and the window's pick at its last accepted row (`next`, the next window's row 0).
     void set_draft_history(const int32_t* tail, int64_t n_tail, int32_t next);
+    void set_ple_session(SessionState* s) { ple_ss_ = s; }
     bool coupled() const { return coupled_active_; }
 
     double ms_draft = 0, ms_prefill = 0;
@@ -124,6 +125,7 @@ private:
 
     const ModelGeometry* g_ = nullptr;
     SessionState* ss_ = nullptr;
+    SessionState* ple_ss_ = nullptr;
     const WeightTable* wt_ = nullptr;
     const NativeHead* head_ = nullptr;
     const float* window_R_ = nullptr;

include/strata/core/remote_expert_opt.hpp

@@ -0,0 +1,53 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace strata::core {
+class RemoteExperts;
+class ExpertSource;
+
+// Remote Expert Decode Optimization v1 on existing helper caches. One instance per serial engine.
+class RemoteExpertOpt {
+public:
+    RemoteExpertOpt() = default;
+    ~RemoteExpertOpt();
+    RemoteExpertOpt(const RemoteExpertOpt&) = delete;
+    RemoteExpertOpt& operator=(const RemoteExpertOpt&) = delete;
+
+    void attach(RemoteExperts& remote); // before RemoteExperts::open
+    bool init(std::string& err);        // after all helper caches are open
+    bool owns(int64_t layer, int32_t expert) const;
+    bool adapt(const std::vector<float>& usage, const std::vector<int32_t>& primary,
+               const std::vector<std::pair<int32_t, int32_t>>& pending, int max_swaps,
+               ExpertSource& source);
+
+    // Bracket the existing pool callback; its signature and dispatch remain unchanged.
+    void begin(const float* weights, int token_begin, int tokens);
+    void end() { weights_ = nullptr; }
+    void copy_rows(float* dst, const float* src, int token_begin, int tokens,
+                   const int32_t* hit_rows, const int32_t* count, void* stream) const;
+    void combine(float* dst, float* scratch, int token_begin, int tokens,
+                 const uint32_t* skip, uint32_t ring, void* stream) const;
+
+private:
+    friend class RemoteExperts;
+    static size_t metadata_bytes();
+    bool active() const { return weights_ != nullptr; }
+    void prepare(const RemoteExperts& remote, void* metadata) const;
+    bool reduce(RemoteExperts& remote, const void* metadata, std::string& err);
+    void accumulate(const RemoteExperts& remote);
+
+    struct Peer { RemoteExperts* remote; float* sum = nullptr; };
+    std::vector<Peer> peers_;
+    float* h_sum_ = nullptr;
+    float* m_sum_ = nullptr;
+    int32_t* h_mask_ = nullptr;
+    int32_t* m_mask_ = nullptr;
+    const float* weights_ = nullptr;
+    int token_begin_ = 0, tokens_ = 0;
+};
+} // namespace strata::core

include/strata/core/remote_experts.hpp

@@ -11,6 +11,7 @@
 #include <vector>
 
 namespace strata::core {
+class RemoteExpertOpt;
 
 /// A static, profile-filled expert tier on another CUDA device. CUDA0 keeps all
 /// dense weights and state; results return through the existing pinned CPU rows.
@@ -27,7 +28,7 @@ public:
     bool open(int device, int slots, int64_t layers, int64_t experts,
               const std::vector<std::pair<int32_t, int32_t>>& ranked,
               const ExpertCache& primary, ExpertSource& source,
-              std::vector<uint8_t>& claimed, std::string& err);
+              std::vector<uint8_t>& claimed, std::string& err, bool auto_size = false);
     void close();
 
     /// `kind` is the primary verifier's classification (-1 = CPU candidate),
@@ -36,6 +37,7 @@ public:
                int64_t k, const int32_t* kind, const int32_t* primary_res,
                std::string& err);
     bool owns(int64_t index) const { return owned_[(size_t) index] != 0; }
+    bool optimized_decode() const { return remote_opt_ != nullptr; }
     bool finish(float* out, std::string& err);
     int64_t resident() const { return cache_.resident(); }
     int64_t computed() const { return computed_; }
@@ -48,6 +50,8 @@ public:
     double ms_wait() const { return ms_wait_; }
 
 private:
+    friend class RemoteExpertOpt;
+    RemoteExpertOpt* remote_opt_ = nullptr;
     int device_ = -1;
     int64_t n_expert_ = 0;
     int32_t groups_ = 0;

include/strata/core/session.hpp

@@ -16,6 +16,7 @@
 
 #include "strata/core/hit_hook.hpp"
 #include "strata/core/layer.hpp"
+#include "strata/kernels/cpu/pool.hpp"
 
 #include <cuda_runtime.h>
 
@@ -369,7 +370,7 @@ struct SessionLoopScratch {
     float* y_miss = nullptr;        ///< pinned host staging for the pool's answer, `parts_bytes` long
     size_t parts_bytes = 0;
     cudaEvent_t probe = nullptr;
-    long long pinned_core = -1;     ///< the affinity to restore, or -1 if the host was never pinned
+    strata::kernels::cpu::ThreadAffinity pinned_core{};  ///< affinity to restore, or invalid if the host was not pinned
     bool pinned = false;
 
     /// Allocates the buffers and pins the host thread.  Call ONCE, at session setup.

include/strata/core/verify.hpp

@@ -38,6 +38,7 @@
 namespace strata::core {
 
 class NativeHead;
+class RemoteExpertOpt;
 
 /// The CPU pool for a window: x_f (n_tok, n_embd), ids (n_tok, k) -> out (n_tok * k, n_embd), hit rows zeroed.
 using PoolMultiFn = void (*)(void* user, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k, float* out,
@@ -45,6 +46,7 @@ using PoolMultiFn = void (*)(void* user, const float* x_f, const int32_t* ids, i
 
 struct VerifyHits {
     const int32_t* d_res = nullptr;      ///< device [n_layers * n_expert] slot or -1
+    const int32_t* h_res = nullptr;      ///< host [n_layers * n_expert] slot or -1 (when all >= 0 in stage: zero-doorbell)
     const uint8_t* cache_base = nullptr; ///< slot 0 of the VRAM expert arena
     const uint64_t* slot_off = nullptr;   ///< E-6: host per-slot offsets when slots differ in size (null: slot * blob)
     int64_t n_slots = 0;                  ///< E-6: how many (for the device copy)
@@ -57,6 +59,7 @@ public:
     ~Verifier();
     Verifier(const Verifier&) = delete;
     Verifier& operator=(const Verifier&) = delete;
+    void set_remote_expert_opt(RemoteExpertOpt* opt) { remote_opt_ = opt; } // before init/capture
 
     /// The watchdog's view of the window in flight (issue #31): the layer, the GPU's sequence, the flags.
     void diag(std::FILE* f) const;
@@ -164,6 +167,7 @@ public:
     std::string profile_report();
 
 private:
+    RemoteExpertOpt* remote_opt_ = nullptr;
     bool capture(int T, std::string& err);
     strata::kernels::SamplerParams sampling_ = [] {
         strata::kernels::SamplerParams s;
@@ -176,6 +180,7 @@ private:
     bool head_sampling_ = true;          ///< set_head_sampling
     int device_ = -1;                    ///< the device `init` ran on: run/commit switch to it (layer split)
     std::atomic<bool> released_{false};  ///< #267: release_gpu_waits ran (maybe on the watchdog thread): no more windows
+    bool all_resident_ = false;           ///< 100% of experts in [lb_, le_) resident in VRAM: zero-doorbell graph
     bool device_plan_ = false;            ///< E-6: resident-only layers planned on the device (STRATA_VERIFY_DEVICE_PLAN)
     uint32_t* skip_ = nullptr;            ///< E-6: per group, the ring whose plan the device built (0: the host's)
     unsigned long long* slot_off_d_ = nullptr;   ///< E-6: the slot offsets on the device
@@ -185,6 +190,9 @@ private:
     Verifier* next_ = nullptr;
     void* next_user_ = nullptr;
     bool ple_stage() const { return lb_ <= 1 && 1 < le_; }   ///< holds layer 1, where the PLE block runs
+    void stage_inputs(int T, const int32_t* tokens, int64_t pos0);
+    bool staged_ = false;
+    bool copy_used_ = false;
     bool capture_commit(std::string& err);
     bool record_window(int T, cudaStream_t cs, std::string& err);
     // #649: STRATA_VERIFY_TRACE=1 - a host event ring (trace_ev) and GPU breadcrumbs: the profiler's stamp points,
@@ -212,6 +220,8 @@ private:
     int32_t last_tokens_[8] = {};
     int64_t n_vocab_ = 0;
     cudaStream_t cs_ = nullptr;
+    cudaStream_t sh_cs_ = nullptr;
+    cudaEvent_t ev_fork_ = nullptr, ev_join_ = nullptr;
     cudaGraphExec_t exec_[9] = {};
     cudaGraphExec_t commit_exec_ = nullptr;
 
@@ -252,6 +262,7 @@ private:
     float *ple_ = nullptr, *emb_ = nullptr, *R_ = nullptr, *mixed_ = nullptr, *bo_ = nullptr;
     float *inj_ = nullptr, *inj2_ = nullptr, *lo_ = nullptr, *rs_ = nullptr, *xn_ = nullptr;
     uint8_t* xq_ = nullptr;                                   // T columns of q8_1
+    uint8_t* sh_xq_ = nullptr;                                // T columns of q8_1 for shared expert branch
     float *qkv_L_ = nullptr, *h_L_ = nullptr, *gate_L_ = nullptr, *beta_L_ = nullptr;   // per GDN layer
     float *z_ = nullptr, *y_ = nullptr, *y_dummy_ = nullptr;
     float *qfull_ = nullptr, *qcur_ = nullptr, *kcur_ = nullptr, *vcur_ = nullptr, *idx_raw_L_ = nullptr;

include/strata/hip_compat/cuda_runtime.h

@@ -46,6 +46,7 @@
 #define cudaFreeHost hipFreeHost
 #endif
 #define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize
+#define cudaFuncAttributePreferredSharedMemoryCarveout hipFuncAttributePreferredSharedMemoryCarveout
 #define cudaGetDevice hipGetDevice
 #define cudaGetDeviceCount hipGetDeviceCount
 #define cudaGetDeviceProperties hipGetDeviceProperties

include/strata/kernels/cpu/expert_layout.hpp

@@ -44,6 +44,8 @@ struct ExpertLayout {
 /// Plan v0.3 P6: whether this CPU (and its OS) runs the AVX-512 kernels (F, BW, VL, VNNI, VBMI).  Probed in a
 /// file compiled without AVX-512, so asking is safe everywhere; STRATA_FORCE_AVX2=1 answers no (for tests).
 bool cpu_avx512_ok();
+/// Whether this CPU runs AVX-512F + AVX-512BW + AVX-512VL + AVX-512DQ (e.g. Skylake-X i9-9920X) for iq_avx512.cpp.
+bool cpu_avx512bw_ok();
 /// Whether this CPU (and its OS) runs the AVX2 kernels (AVX, AVX2, FMA, F16C): the floor of every expert kernel
 /// (q2_avx2.cpp, iq_avx2.cpp, and ggml-cpu in the portable build).  STRATA_FORCE_AVX2 does not change it.
 bool cpu_avx2_ok();

include/strata/kernels/cpu/pool.hpp

@@ -76,7 +76,7 @@ struct CpuTopology {
     int p_cores = 0;                ///< Physical performance cores
     int p_threads = 0;              ///< Total logical threads on performance cores
     int e_cores = 0;                ///< Efficient cores
-    std::vector<int> worker_cores;  ///< Ordered logical core IDs for workers (excluding host if skip_first)
+    std::vector<int> worker_cores;  ///< Ordered CPU IDs; on Windows, group * 64 + processor within the group
     int host_core = -1;             ///< Logical core reserved for host thread
 };
 
@@ -100,9 +100,23 @@ std::vector<int> physical_cores(bool skip_first, PoolAffinity affinity = PoolAff
 /// - exactly 5/6 of L9's 44.14 on 6 - and at **26.9 GB/s inside the host loop**, where the unpinned spinning
 /// host is free to land on a worker's core or its SMT sibling.  That 1.35x is not the kernel.
 ///
-/// Returns the PREVIOUS affinity mask, or -1 if the platform refused; pass it to `restore_thread_affinity`.
-long long pin_current_thread(int core);
-void restore_thread_affinity(long long previous);
+/// The previous host placement. `valid` is false when querying or setting placement failed.
+/// Windows uses a reversible CPU Set selection, leaving hard affinity (including Windows 11's implicit
+/// all-group eligibility) untouched. An empty selection restores inheritance from the process defaults.
+/// Linux stores the complete, dynamically sized native CPU mask, including CPU IDs above 63.
+struct ThreadAffinity {
+#if defined(_WIN32)
+    std::vector<unsigned long> cpu_sets;
+#else
+    std::vector<unsigned long> mask;
+#endif
+    bool valid = false;
+};
+
+/// Selects one encoded processor for the caller; Windows CPU Sets respect existing hard-affinity limits.
+/// Restore on the same calling thread. Pool-owned workers use hard group affinity instead.
+ThreadAffinity pin_current_thread(int core);
+void restore_thread_affinity(const ThreadAffinity& previous);
 
 class ExpertPool {
 public:

include/strata/kernels/elementwise.hpp

@@ -1,4 +1,4 @@
-// include/strata/kernels/elementwise.hpp - the small ops a LAYER needs between its GEMVs, P2.S5.
+// include/strata/kernels/elementwise.hpp - the small ops a LAYER needs between its GEMVs, P2.S5.
 //
 // These are not interesting kernels and that is why they are in one file with one parity test: they are the
 // glue in the per-layer graph, and glue that is written inline at each call site is glue with several
@@ -118,6 +118,8 @@ void doorbell_publish(const float* x, const int32_t* ids, const float* weights,
 void doorbell_publish_value(const float* x, const int32_t* ids, const float* weights, int64_t n, int64_t k,
                             float* x_out, int32_t* ids_out, float* weights_out, uint32_t* d_seq, uint32_t value,
                             void* stream);
+void doorbell_publish_res(const float* x, const int32_t* ids, const int32_t* d_res, int n_expert, int64_t n, int64_t k,
+                          float* x_out, int32_t* ids_out, uint32_t* d_seq, void* stream);
 
 /// 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.

include/strata/kernels/native_mmvq.hpp

@@ -38,6 +38,9 @@ bool native_mmvq_multi_exact();
 void native_quantize_q8_1(const float* x, void* x_q8_1, int n_in, int ncols,
                           void* stream);
 
+void native_swiglu_quantize_q8_1(const float* gate, const float* up, void* x_q8_1,
+                                 int n_in, int ncols, void* stream);
+
 void native_q5_k_mmvq(const void* weights, const void* x_q8_1, float* y,
                       int n_in, int n_out, int ncols, void* stream);
 

include/strata/kernels/ngram.hpp

@@ -143,6 +143,9 @@ public:
     /// needs the rows. `gather` is `issue` followed by `collect`. In Mmap mode `issue` only prefetches.
     bool issue(const uint32_t* rows16);
     bool collect(float* out2560, std::string& err);
+    /// Non-blocking prefetch of 16 rows for one token into the reader's row cache while the GPU is busy.
+    void prefetch_rows(const uint32_t* rows16);
+    void wait_prefetches();
     /// Plan v0.3 P5: the rows of `n_tokens` tokens (16 each, `rows` token-major) into `out` (2560 floats per token),
     /// as ONE reader request - page dedupe and sort across the whole batch, the reader's full queue depth.  Not
     /// while a single-token `issue` is pending.  The mapped mode gathers row by row.

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_enabled();
 
 /// Opt in to pinned CUDA postprojection norms, gate, convolution and residual
 /// arithmetic (default false). Set before capture; existing graphs keep their
@@ -107,6 +108,9 @@ uint64_t ple_block_scratch_bytes();
 void ple_block(const float* emb, const float* hidden, const float* hist_rows, const PleWeights& w,
                PleOut& out, void* scratch, void* stream);
 
+void ple_block_projected(const float* projected_key, const float* projected_value, const float* hidden,
+                         const float* hist_rows, const PleWeights& w, PleOut& out, void* scratch, void* stream);
+
 /// Advance the row-fastest normalized history by one token: hist[r,c] = old_hist[r+1,c], then append
 /// normalized[c] at row NG_HIST-1. One thread owns each channel, making the in-place shift well-defined.
 /// hist and normalized must be nonnull, disjoint device regions of NG_HIST*NG_HC_DIM and NG_HC_DIM floats.

include/strata/kernels/qsa_select.hpp

@@ -49,7 +49,8 @@ void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64
 /// build below sm_90, or a capacity whose keys do not fit one cluster's shared memory. Capturable.
 bool qsa_block_topk_cluster(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap,
                             const QsaShapes& s, int32_t* ids, void* stream);
-/// The original kernel (keys read from memory on every radix pass), for tests: the same ids.
+/// The original kernel (256 threads, one shared histogram, keys read from memory on every radix pass), for tests and
+/// STRATA_TOPK_OLD=1: the same ids.
 void qsa_block_topk_ref(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap,
                         const QsaShapes& s, int32_t* ids, void* stream);
 

include/strata/kernels/shared_expert.hpp

@@ -24,6 +24,7 @@ namespace strata::kernels {
 /// Configure before session capture; captured graphs retain their selected kernels.
 /// In this mode shared_expert requires its optional unrounded x_f32 input.
 void shared_expert_set_native_bf16(bool enabled);
+bool shared_expert_native_bf16_enabled();
 
 /// Optional native GGUF projections. Each supported type with nonnull data
 /// replaces only that canonical projection; absent or unsupported entries fall
@@ -37,6 +38,7 @@ struct NativeSharedWeights {
     const void* up_data = nullptr;
     const void* down_data = nullptr;
     void* q8_1 = nullptr;
+    const void* x_q8_1 = nullptr;   // optional pre-quantized Q8_1 input for gate/up in shared_expert_multi
 };
 
 /// Bytes of caller-owned scratch `shared_expert` needs.  **THE KERNEL USED TO `cudaMalloc` FOUR BUFFERS ON

include/strata/prefill/kernels.hpp

@@ -60,6 +60,9 @@ void swiglu_pair(const float* g, const float* u, uint16_t* h16, int64_t n, void*
 /// dst[i] = src[i] for n int32s, as a kernel: either side may be mapped host memory, and the copy never waits
 /// behind the copy engine's queue (the prompt path's grouping tables, while the expert stream fills it).
 void copy_i32(int32_t* dst, const int32_t* src, int64_t n, void* stream);
+/// dst[0, n) = src[0, n) in 16-byte loads (both 16-byte aligned; src may be mapped host memory: read over PCIe by
+/// the kernel, not queued on a copy engine).
+void copy_f32_wide(float* dst, const float* src, int64_t n, void* stream);
 /// Gather rows: dst16[i, :] = x16[src[i], :] (n rows of `width` BF16).
 void gather_rows16(const uint16_t* x16, const int32_t* src, uint16_t* dst16, int64_t n, int64_t width, void* stream);
 /// bo[t, :] = shared[t, :] * sigmoid(sg[t]) + sum_k w[t, k] * D[slot[t, k], :]

include/strata/prefill/prefill.hpp

@@ -18,6 +18,7 @@
 
 #include <cstdint>
 #include <functional>
+#include <future>
 #include <memory>
 #include <string>
 
@@ -78,10 +79,37 @@ public:
     /// #340: the streamed ring's slot count for chunks that stream every expert, instead of the pinned-share rule
     /// (0 = that rule). Set before any `bytes_needed`/`init` (both count the ring); STRATA_PREFILL_RING still wins.
     static void set_ring_override(int slots);
+    /// #583: the auto chunk scan's byte-budget ring, for chunks above `small_max` (0.1.39's auto chunk: a prompt that
+    /// fits it keeps 0.1.39's ring).  0 slots = none.  A layer split's set_ring_override and STRATA_PREFILL_RING win.
+    static void set_ring_budget(int slots, int64_t small_max);
 
     /// Device bytes `init` needs for a chunk of `chunk` tokens (what a borrowed region must hold).
     static uint64_t bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk);
 
+    /// The same without the streamed ring: what the chunk's own buffers cost.  The auto chunk scan sizes the chunk
+    /// first and hands the ring what the chunk leaves over, so it needs the chunk priced on its own.
+    static uint64_t bytes_needed_no_ring(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk);
+
+    /// The streamed ring's byte budget as a slot count for this pack (the measured slot count x Q2_0's blob, over
+    /// max_blob, never past ring_cap()):
+    /// what the auto chunk scan treats as a full ring.  A slot is one whole blob, so a pack with bigger blobs than
+    /// Q2_0's gets fewer of them for the same bytes - 384 on Q2_0, 199 on a 2.54 MiB-blob IQ3_S pack.
+    static int64_t ring_max_slots();
+    /// 0.1.39's ring for this PC (1024 fused / 384 pinned / 96), within ring_cap().
+    static int64_t ring_default_slots();
+    /// #583: the ring the auto scan keeps full, given the chunk 0.1.39's rule picked: its byte budget where that rule's
+    /// chunk was small (< 6144), else 0.1.39's ring (measured: shrinking it for a bigger chunk lost there).
+    static int64_t ring_cap_for(int64_t old_chunk);
+
+    /// What the ring actually resolves to for a chunk of `chunk` tokens, after the override, STRATA_PREFILL_RING
+    /// and the pinned-share rule - the slot count `init` lays out.  The engine reports it on its INFO line so the
+    /// Monitor tab shows the pair the run really got, not what it asked for.
+    static int64_t ring_slots_for(int64_t chunk);
+
+    /// 0.1.39b (#583, the default): the ring as a byte budget, the loan's corrected count and the auto chunk scan that
+    /// keeps the ring full.  STRATA_RING_BYTES=0: 0.1.39's ring, loan and chunk list.
+    static bool ring_bytes_enabled();
+
     /// Positions [pos0, pos0 + n) holding `tokens`; `ss.ple_prev` must be the two tokens before pos0 (oldest
     /// first, -1 for none) and is advanced to the last two of these.
     bool run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err);
@@ -120,10 +148,32 @@ public:
         stage_lb_ = layer_begin; stage_le_ = layer_end; next_ = next;
     }
 
+    /// LAYER SPLIT: `helper` is another stage's prompt path (another GPU).  A prompt of one chunk runs the stages one
+    /// after the other, so while this stage reads it the helper's GPU idles: it then streams a share of this stage's
+    /// non-resident experts over its own PCIe link into its own (lent) prompt buffers, computes their rows, and sends
+    /// them back - the --peer-device peer's streaming, without P2P (activations and rows through mapped host memory,
+    /// read by copy kernels).  Only on the MMQ prompt path, only for chunks of stream_all_min() tokens and more, up to
+    /// the size where the measured share stops paying.  Opt-in: STRATA_PREFILL_HELP=1 (not bit-identical to the default).  Both stages must have
+    /// run `init`.
+    bool set_stage_helper(Prefill* helper, std::string& err);
+
 private:
+    // Stage-1 pipeline: intermediate stages return after handing their chunk to
+    // the direct successor. The public run() drains the chain once at prompt end.
+    bool run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err);
+    bool drain_pipeline(std::string& err);
+
     int64_t stage_lb_ = 0, stage_le_ = -1;
     Prefill* next_ = nullptr;
+    Prefill* helper_ = nullptr;         ///< set_stage_helper
+    bool single_chunk_ = false;         ///< a later stage: the prompt is one chunk (set by the stage before)
+    bool bind_stage_helper(int64_t T);  // binds the helper's buffers for a one-chunk prompt of T tokens
     const float* hand_in_ = nullptr;    ///< the previous stage's rows of the chunk being read (host, pinned)
+
+    std::string next_err_;
+    std::future<bool> next_run_;
+    int hand_buf_ = 0;
+
     bool carve(std::size_t T, void* alloc);   // the device buffers of a chunk (prefill.cpp's Alloc)
     void release();                          // the destructor's cleanup (also `reset`'s)
     struct Impl;

setup.py

@@ -855,6 +855,29 @@ def split_budget(cfg: dict) -> bool:
     return True
 
 
+REMOTE_EXPERT_OPT = "--remote-expert-opt"
+
+
+def recommend_remote_expert_opt(cfg: dict, off: bool = False) -> None:
+    """0.1.39b (#578): a config on two or more GPUs gets --remote-expert-opt - the helper expert caches
+    (--expert-cache-device1..3) then stay complementary to the main GPU's, return their rows already weighted and skip
+    the CPU's activation quantization where no expert is left to it (dual RTX 4090: +63% mixed, +132% code over the
+    plain helper path).  The engine uses it only with a helper cache; a layer split runs as before.  A recommendation:
+    `off` (setup's --no-remote-expert-opt) or "remote_expert_opt": false in the config keeps it out, and a single-GPU
+    config is not touched."""
+    if not isinstance(cfg.get("gpu"), list) or len(cfg["gpu"]) < 2:
+        return
+    args = cfg.setdefault("args", [])
+    if off or cfg.get("remote_expert_opt") is False:
+        if REMOTE_EXPERT_OPT in args:
+            args.remove(REMOTE_EXPERT_OPT)
+        return
+    if REMOTE_EXPERT_OPT not in args:
+        args.append(REMOTE_EXPERT_OPT)
+        ok("multi-GPU: --remote-expert-opt (helper expert caches complementary to the main GPU's, #578; "
+           "--no-remote-expert-opt leaves it out)")
+
+
 def offer_together(cfg_path: Path, cfg: dict, yes: bool) -> dict:
     """Starting a model set up for one card on a PC with two or more that can share it: asked once (the answer is
     saved in its config)."""
@@ -897,6 +920,7 @@ def offer_together(cfg_path: Path, cfg: dict, yes: bool) -> dict:
         cfg["layer_split"] = cfg.get("layer_split") or "auto"
         split_mmap(cfg)
         split_budget(cfg)
+        recommend_remote_expert_opt(cfg)
         ok("from now on this model runs on " + " + ".join(gpu_name(g) for g in pair))
     else:
         ok("staying on one GPU (START-HERE.bat --gpus " + ",".join(str(g["index"]) for g in pair) + " switches)")
@@ -3106,6 +3130,7 @@ def start(cfg_path: Path, port: int | None, gpu: int | list | None = None, open_
         cfg["gpu"], cfg["gpus_asked"] = gpu, True
         cfg["layer_split"] = layer_split or cfg.get("layer_split") or "auto"
         split_budget(cfg)                              # #498: before it is saved (it stops when the RAM is short)
+        recommend_remote_expert_opt(cfg)
         write_config(cfg_path, cfg)
         gpu = None
     elif gpu is not None:                              # --gpu N: this start only, on that card
@@ -3480,6 +3505,8 @@ def main() -> int:
     ap.add_argument("--layer-split", help="with --gpus: where each later GPU's layers start (\"18\", \"16,32\"), one "
                                           "rising number per GPU after the first - not layers per card; default "
                                           "auto, placed from each GPU's free VRAM")
+    ap.add_argument("--no-remote-expert-opt", action="store_true",
+                    help="with two or more GPUs: leave out --remote-expert-opt, which setup adds there (#578)")
     ap.add_argument("--host", help="where the server listens: 127.0.0.1 = this PC only (default), 0.0.0.0 = also other "
                                    "devices on your network (issue #26; set --api-key too)")
     ap.add_argument("--api-key", help="require this key from clients (recommended with --host 0.0.0.0)")
@@ -4248,6 +4275,7 @@ def main() -> int:
         cfg["gpu"] = multi
         cfg["layer_split"] = a.layer_split or "auto"
         ok(f"layer split across GPUs {multi} ({cfg['layer_split']})")
+        recommend_remote_expert_opt(cfg, off=a.no_remote_expert_opt)
     if a.host:
         cfg["host"] = a.host
     if a.api_key:

src/core/expert_source.cpp

@@ -2127,16 +2127,30 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
             for (int64_t i = 0; i < n; ++i) if (d.remote[r]->owns(i)) kind[i] = 2;
         }
     }
+    static const bool dec_batch = [] { const char* v = std::getenv("STRATA_DEC_BATCH"); return v == nullptr || std::atoi(v) != 0; }();
+    // the GPU's copy_rows_from_mapped zeroes the rows it computed itself (dec_batch); STRATA_VERIFY_DEVICE_PLAN's
+    // copy_or_zero_from_mapped copies every row, so there the host still zeroes them
+    static const bool device_plan = [] { const char* v = std::getenv("STRATA_VERIFY_DEVICE_PLAN"); return v != nullptr && std::atoi(v) != 0; }();
+    const bool gpu_zeroes_hits = (d.plan != nullptr && n <= kMaxWindowEntries && n <= d.plan->cap && dec_batch && !device_plan);
+    bool any_cpu = false;
+    for (int64_t i = 0; i < n; ++i)
+        if (kind[i] < 0) { any_cpu = true; break; }
     const auto c1 = std::chrono::steady_clock::now();
-    if (native && strata::kernels::cpu::q2_native_kernels(lay.fmt[(size_t) d.layers].gu_type))   // a native Q2_0 pack: the Q2_0 kernels' activations
-        for (int64_t t = 0; t < n_tok; ++t) act_quant_any(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]);
-    else if (native)
-        for (int64_t t = 0; t < n_tok; ++t)
-            native_quant_act(lay.fmt[(size_t) d.layers], x_f + (size_t) t * H, d.nact_multi.data() + (size_t) t * kNativeActBytes);
-    else
-        for (int64_t t = 0; t < n_tok; ++t) act_quant_q8_1(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]);
+    if (any_cpu) {
+        // #578 --remote-expert-opt: a token whose experts all run on a GPU (CUDA0 or a helper) needs no CPU activation
+        const bool ep = d.remote_count > 0 && d.remote[0]->optimized_decode();
+        for (int64_t t = 0; t < n_tok; ++t) {
+            if (ep && std::all_of(kind + t * k, kind + (t + 1) * k, [](int32_t v) { return v >= 0; })) continue;
+            if (native && strata::kernels::cpu::q2_native_kernels(lay.fmt[(size_t) d.layers].gu_type))   // a native Q2_0 pack: the Q2_0 kernels' activations
+                act_quant_any(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]);
+            else if (native)
+                native_quant_act(lay.fmt[(size_t) d.layers], x_f + (size_t) t * H, d.nact_multi.data() + (size_t) t * kNativeActBytes);
+            else
+                act_quant_q8_1(x_f + (size_t) t * H, H, d.act_multi[(size_t) t]);
+        }
+    }
     const auto c2 = std::chrono::steady_clock::now();
-    {   // CS-T: the experts the CPU computes, fetched together (the GGUF in place reads them on several threads)
+    if (any_cpu) {   // CS-T: the experts the CPU computes, fetched together (the GGUF in place reads them on several threads)
         static thread_local std::vector<int64_t> miss;
         miss.clear();
         for (int64_t i = 0; i < n_tok * k; ++i)
@@ -2162,8 +2176,10 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
                 if (kind[i] == 0) ++d.cache_hits;
                 else ++d.offload_entries;                       // #588: PCIe or another GPU
                 if (kind[i] == 2 && d.peer != nullptr) ++d.peer_entries;
-                // multi-GPU: a direct peer launch is writing this row right now - zeroing it would race it
-                if (!(kind[i] == 2 && d.peer != nullptr && d.peer->launched_direct()))
+                // multi-GPU: a direct peer launch is writing this row right now - zeroing it would race it;
+                // dec_batch: copy_rows_from_mapped_kernel already zeroes kind 0/1 rows on the GPU
+                if (!((kind[i] == 2 && d.peer != nullptr && d.peer->launched_direct()) ||
+                      (gpu_zeroes_hits && (kind[i] == 0 || kind[i] == 1))))
                     std::memset(row, 0, (size_t) H * sizeof(float));
                 continue;
             }
@@ -2193,8 +2209,10 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
         }
     const auto c3 = std::chrono::steady_clock::now();
     pt("run", njobs);
-    if (native) d.pool->run_split_multi_native(lay.fmt[(size_t) d.layers], d.jobs_multi.data(), njobs);
-    else d.pool->run_split_multi(d.jobs_multi.data(), njobs);
+    if (njobs > 0) {
+        if (native) d.pool->run_split_multi_native(lay.fmt[(size_t) d.layers], d.jobs_multi.data(), njobs);
+        else d.pool->run_split_multi(d.jobs_multi.data(), njobs);
+    }
     if (d.remote_count > 0) {
         static thread_local std::string remote_error;
         for (int r = 0; r < d.remote_count; ++r)

src/core/mtp.cpp

@@ -26,6 +26,7 @@
 #include "strata/kernels/sampler.hpp"
 #include "strata/kernels/shared_expert.hpp"
 #include "strata/kernels/verify_kernels.hpp"
+#include "strata/kernels/ngram.hpp"
 
 #include <algorithm>
 #include <atomic>
@@ -37,6 +38,11 @@
 #include <fstream>
 #include <sstream>
 #include <vector>
+#if defined(_WIN32)
+#include <intrin.h>
+#elif defined(__x86_64__)
+#include <immintrin.h>
+#endif
 
 namespace strata::core {
 namespace {
@@ -146,6 +152,7 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session
     cudaGetDevice(&device_);   // a layer split's last stage on another GPU: the drafter lives there
     g_ = &g;
     ss_ = &ss;
+    if (ple_ss_ == nullptr) ple_ss_ = &ss;
     max_t_ = max_t;
     rt_dir_ = rt_dir;
     if (max_t < 1 || max_t > strata::kernels::kVerifyMaxT) { err = "mtp: max_t out of range"; return false; }
@@ -580,8 +587,13 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
         nsw.down_type = GGML_Q8_0; nsw.down_data = q8("mlp.shared_expert.down_proj.weight");
         nsw.q8_1 = xq_;
         const SForm none{};
+        const bool need_bf16_x = !shared_expert_native_bf16_enabled();
+        static const bool fuse_head_gr = [] {
+            const char* v = std::getenv("STRATA_FUSE_HEAD_GR");
+            return v != nullptr && std::atoi(v) != 0;
+        }();
         for (int t = 0; t < T; ++t) {
-            f32_to_bf16_bulk(mixed_ + t * N, x_bf16_, N, cs);
+            if (need_bf16_x) 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);
@@ -589,14 +601,30 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
                 native_moe_combine(parts_ + (size_t) t * K * N, w_ + t * K, shared_ + t * N, y_ + t * N, N, K, cs);
             else
                 moe_combine(parts_ + (size_t) t * K * N, w_ + t * K, shared_ + t * N, y_ + t * N, N, K, cs);
-            gr_write(R_ + (size_t) t * HC * N, y_ + t * N, inj2_ + t * HC, gs, R_ + (size_t) t * HC * N, cs);
+            if (!fuse_head_gr)
+                gr_write(R_ + (size_t) t * HC * N, y_ + t * N, inj2_ + t * HC, gs, R_ + (size_t) t * HC * N, cs);
         }
         // ---- the final mixer and the main model's head
-        for (int t = 0; t < T; ++t)
-            gr_read(R_ + (size_t) t * HC * N, f32("hyper_connection_mixer.hc_norm.weight"),
-                    bf16("hyper_connection_mixer.input_mix_weight_down.weight"),
-                    bf16("hyper_connection_mixer.input_mix_weight_up.weight"), nullptr, EPS, gs, ss.block.gr,
-                    sample_ + t * N, dummy_inj_, cs);
+        if (fuse_head_gr) {
+            FusedGrArgs fa[kFusedGrMaxT];
+            for (int t = 0; t < T; ++t) {
+                fa[t].R = R_ + (size_t) t * HC * N; fa[t].R_out = R_ + (size_t) t * HC * N; fa[t].apply = true;
+                fa[t].bo_prev = y_ + t * N; fa[t].inj_prev = inj2_ + t * HC;
+                fa[t].w_norm = f32("hyper_connection_mixer.hc_norm.weight");
+                fa[t].w_down = bf16("hyper_connection_mixer.input_mix_weight_down.weight");
+                fa[t].w_up = bf16("hyper_connection_mixer.input_mix_weight_up.weight");
+                fa[t].w_inject = nullptr;
+                fa[t].eps = EPS; fa[t].lo = lo_ + t * g.hc_lr; fa[t].rs = rs_ + t * HC;
+                fa[t].inject_out = dummy_inj_; fa[t].mixed = sample_ + t * N;
+            }
+            fused_gr_read_multi(fa, T, xn_, cs);
+        } else {
+            for (int t = 0; t < T; ++t)
+                gr_read(R_ + (size_t) t * HC * N, f32("hyper_connection_mixer.hc_norm.weight"),
+                        bf16("hyper_connection_mixer.input_mix_weight_down.weight"),
+                        bf16("hyper_connection_mixer.input_mix_weight_up.weight"), nullptr, EPS, gs, ss.block.gr,
+                        sample_ + t * N, dummy_inj_, cs);
+        }
         native_quantize_q8_1(sample_, xq_, (int) N, T, cs);
         const bool sub = dhead_ != nullptr;
         const int64_t nv = sub ? n_dvocab_ : n_vocab_;
@@ -829,6 +857,9 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t*
     if (T < 1 || T > max_t_ || a < 0 || a >= T) { err = "mtp: draft arguments out of range"; return false; }
     const bool cp = coupled_active_;   // coupled draft sampling for this request: its own graphs
     if (!capture_round(T, cp, err)) return false;
+    const int max_steps = std::min(max_t_ - 1, max_drafts_);
+    for (int j = 1; j < max_steps; ++j)
+        if (!capture_step(j, cp, err)) return false;
     const Clock::time_point t0 = Clock::now();
     const int64_t NH = g_->n_head;
     auto put = [&](int row, int64_t cell) {
@@ -843,32 +874,99 @@ bool MtpDrafter::draft(int T, const int32_t* tokens, int64_t p, int a, int32_t*
         put(t, p + t);
     }
     put(2 * max_t_ - 1, coupled_draft_cell(p, a, 0));   // p + a: draft 0's cell
+    for (int j = 1; j < max_steps; ++j)
+        put(max_t_ + j - 1, coupled_draft_cell(p, a, j));
+    for (int j = 0; j < max_steps; ++j) ((volatile int32_t*) h_out_)[j] = -1;
     h_row_[0] = a;
     h_row_[1] = 0;
     std::atomic_thread_fence(std::memory_order_seq_cst);
-    if (cudaGraphLaunch(cp ? round_exec_c_[T] : round_exec_[T], cs_) != cudaSuccess ||
-        cudaStreamSynchronize(cs_) != cudaSuccess) {
-        err = std::string("mtp draft: ") + cudaGetErrorString(cudaGetLastError());
-        return false;
-    }
-    drafts[0] = ((volatile int32_t*) h_out_)[0];
-    float pj = ((volatile float*) h_prob_)[0];
-    if (probs) probs[0] = pj;
-    int n = 1;
-    // the chain continues while the last draft is likely enough to be verified
-    for (int j = 1; j < std::min(max_t_ - 1, max_drafts_) && pj >= min_p; ++j) {
-        if (!capture_step(j, cp, err)) return false;
-        put(max_t_ + j - 1, coupled_draft_cell(p, a, j));   // p + a + j; coupled: drawn at counter cell + 1
-        std::atomic_thread_fence(std::memory_order_seq_cst);
-        if (cudaGraphLaunch(cp ? step_exec_c_[j] : step_exec_[j], cs_) != cudaSuccess ||
-            cudaStreamSynchronize(cs_) != cudaSuccess) {
-            err = std::string("mtp draft step: ") + cudaGetErrorString(cudaGetLastError());
+
+    SessionState* const pss = ple_ss_ ? ple_ss_ : ss_;
+    const bool do_ple = pss != nullptr && pss->ple.ready() && pss->ple.table != nullptr;
+    int32_t ple_prev[2] = {do_ple ? pss->ple_prev[0] : 0, do_ple ? pss->ple_prev[1] : 0};
+    auto prefetch_ple = [&](int32_t tok) {
+        if (!do_ple || tok < 0) return;
+        uint32_t rows16[strata::kernels::PLE_N_HEADS];
+        strata::kernels::ngram_rows(&tok, ple_prev, 1, pss->ple.consts, rows16);
+        ple_prev[0] = ple_prev[1];
+        ple_prev[1] = tok;
+        pss->ple.table->prefetch_rows(rows16);
+    };
+
+    int n = 0;
+    if (min_p <= 0.0f && max_steps > 0) {
+        if (cudaGraphLaunch(cp ? round_exec_c_[T] : round_exec_[T], cs_) != cudaSuccess) {
+            err = std::string("mtp draft: ") + cudaGetErrorString(cudaGetLastError());
+            return false;
+        }
+        for (int j = 1; j < max_steps; ++j) {
+            if (cudaGraphLaunch(cp ? step_exec_c_[j] : step_exec_[j], cs_) != cudaSuccess) {
+                err = std::string("mtp draft step: ") + cudaGetErrorString(cudaGetLastError());
+                return false;
+            }
+        }
+        (void) cudaStreamQuery(cs_);
+        prefetch_ple(tokens[a]);
+        for (int j = 0; j + 1 < max_steps; ++j) {
+            uint32_t spins = 0;
+            while (((volatile int32_t*) h_out_)[j] < 0) {
+#if defined(_WIN32) || defined(__x86_64__)
+                _mm_pause();
+#endif
+                if ((++spins & 1023u) == 0 && cudaStreamQuery(cs_) != cudaErrorNotReady) break;
+            }
+            drafts[j] = ((volatile int32_t*) h_out_)[j];
+            if (probs) probs[j] = ((volatile float*) h_prob_)[j];
+            prefetch_ple(drafts[j]);
+        }
+        if (cudaStreamSynchronize(cs_) != cudaSuccess) {
+            err = std::string("mtp draft: ") + cudaGetErrorString(cudaGetLastError());
+            return false;
+        }
+        const int last = max_steps - 1;
+        drafts[last] = ((volatile int32_t*) h_out_)[last];
+        if (probs) probs[last] = ((volatile float*) h_prob_)[last];
+        prefetch_ple(drafts[last]);
+        n = max_steps;
+    } else if (max_steps > 0) {
+        if (cudaGraphLaunch(cp ? round_exec_c_[T] : round_exec_[T], cs_) != cudaSuccess) {
+            err = std::string("mtp draft: ") + cudaGetErrorString(cudaGetLastError());
+            return false;
+        }
+        (void) cudaStreamQuery(cs_);
+        prefetch_ple(tokens[a]);
+        auto wait_step = [&](int j) {
+            uint32_t spins = 0;
+            while (((volatile int32_t*) h_out_)[j] < 0) {
+#if defined(_WIN32) || defined(__x86_64__)
+                _mm_pause();
+#endif
+                if ((++spins & 1023u) == 0 && cudaStreamQuery(cs_) != cudaErrorNotReady) break;
+            }
+        };
+        wait_step(0);
+        drafts[0] = ((volatile int32_t*) h_out_)[0];
+        float pj = ((volatile float*) h_prob_)[0];
+        if (probs) probs[0] = pj;
+        n = 1;
+        for (int j = 1; j < max_steps && pj >= min_p; ++j) {
+            if (cudaGraphLaunch(cp ? step_exec_c_[j] : step_exec_[j], cs_) != cudaSuccess) {
+                err = std::string("mtp draft step: ") + cudaGetErrorString(cudaGetLastError());
+                return false;
+            }
+            (void) cudaStreamQuery(cs_);
+            prefetch_ple(drafts[j - 1]);
+            wait_step(j);
+            drafts[j] = ((volatile int32_t*) h_out_)[j];
+            pj = ((volatile float*) h_prob_)[j];
+            if (probs) probs[j] = pj;
+            ++n;
+        }
+        prefetch_ple(drafts[n - 1]);
+        if (cudaStreamSynchronize(cs_) != cudaSuccess) {
+            err = std::string("mtp draft: ") + cudaGetErrorString(cudaGetLastError());
             return false;
         }
-        drafts[j] = ((volatile int32_t*) h_out_)[j];
-        pj = ((volatile float*) h_prob_)[j];
-        if (probs) probs[j] = pj;
-        ++n;
     }
     for (int j = n; j < max_t_ - 1; ++j) { drafts[j] = 0; if (probs) probs[j] = 0.0f; }
     if (n_drafts) *n_drafts = n;

src/core/pinned.cu

@@ -224,9 +224,27 @@ void* reserve(uint64_t bytes, PageBacking& got, std::string& note, const std::st
         note = "MAP_HUGETLB unavailable (needed " + std::to_string(need) + " 2 MiB pages, vm.nr_hugepages=" +
                (have_pool ? std::to_string(pool) : std::string("?")) + "); using 4 KB pages";
     }
-    void* p = mmap(nullptr, bytes, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    // No hugetlb mapping: a 2 MiB-aligned anonymous mapping with MADV_HUGEPAGE, so transparent huge pages back
+    // it where THP is "madvise" (the common distro default). The CPU expert pool streams whole experts out of
+    // this arena; with 4 KB pages every 3 MB expert costs ~750 TLB misses.
+    constexpr uint64_t kAlign = 2ull << 20;
+    const uint64_t padded = bytes + kAlign;
+    void* raw = mmap(nullptr, padded, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
     got = PageBacking::NormalPages;
-    return p == MAP_FAILED ? nullptr : p;
+    if (raw == MAP_FAILED) return nullptr;
+    const uintptr_t start = (uintptr_t) raw;
+    const uintptr_t aligned = (start + kAlign - 1) & ~(uintptr_t) (kAlign - 1);
+    if (aligned > start) munmap(raw, aligned - start);
+    const uintptr_t end = aligned + bytes, raw_end = start + padded;
+    if (raw_end > end) munmap((void*) end, raw_end - end);
+    void* p = (void*) aligned;
+    if (std::getenv("STRATA_NO_LARGEPAGES") == nullptr && madvise(p, bytes, MADV_HUGEPAGE) == 0) {
+        const std::string four_k = "; using 4 KB pages";
+        if (note.size() >= four_k.size() && note.compare(note.size() - four_k.size(), four_k.size(), four_k) == 0)
+            note.resize(note.size() - four_k.size());
+        note += "; transparent huge pages requested (MADV_HUGEPAGE)";
+    }
+    return p;
 #endif
 }
 

src/core/remote_expert_opt.cu

@@ -0,0 +1,194 @@
+#include "strata/core/remote_expert_opt.hpp"
+#include "strata/core/remote_experts.hpp"
+#include "strata/core/on_device.hpp"
+#include "strata/kernels/cpu/expert_layout.hpp"
+#include "strata/kernels/elementwise.hpp"
+#include "strata/kernels/verify_kernels.hpp"
+
+#include <algorithm>
+#include <cstring>
+
+namespace strata::core {
+namespace {
+constexpr int64_t H = strata::kernels::cpu::H;
+constexpr int K = 10, MAXT = strata::kernels::cpu::MAXT, CAP = MAXT * K;
+struct ReduceMeta { int32_t row[CAP]; float weight[CAP]; };
+
+bool check(cudaError_t status, std::string& err) {
+    if (status == cudaSuccess) return true;
+    err = std::string("remote expert decode optimization: ") + cudaGetErrorString(status);
+    return false;
+}
+
+__global__ void reduce_experts(const float* parts, const ReduceMeta* meta, float* sum) {
+    const int t = blockIdx.y, c = blockIdx.x * blockDim.x + threadIdx.x;
+    if (c >= H) return;
+    float v = 0.0f;
+    for (int j = 0; j < K; ++j) {
+        const int i = t * K + j, r = meta->row[i];
+        if (r >= 0) v += parts[(int64_t) r * H + c] * meta->weight[i];
+    }
+    sum[(int64_t) t * H + c] = v;
+}
+
+// Like the upstream mapped-row copy, also skip rows returned in the remote sum.
+__global__ void copy_cpu_rows(float4* dst, const volatile float4* src, const int32_t* mask,
+                              const int32_t* hit_rows, const int32_t* count) {
+    const int row = blockIdx.x;
+    __shared__ int skip;
+    if (threadIdx.x == 0) {
+        int s = mask[row];
+        for (int i = 0; i < *count; ++i) s |= hit_rows[i] == row;
+        skip = s;
+    }
+    __syncthreads();
+    const int64_t start = (int64_t) row * (H / 4);
+    for (int i = threadIdx.x; i < H / 4; i += blockDim.x)
+        dst[start + i] = skip ? make_float4(0, 0, 0, 0) : const_cast<const float4*>(src)[start + i];
+}
+} // namespace
+
+void RemoteExpertOpt::attach(RemoteExperts& remote) {
+    remote.remote_opt_ = this;
+    peers_.push_back({&remote});
+}
+
+bool RemoteExpertOpt::init(std::string& err) {
+    void* mapped = nullptr;
+    if (!check(cudaHostAlloc((void**) &h_sum_, MAXT * (H + K) * sizeof(float),
+                             cudaHostAllocPortable | cudaHostAllocMapped), err) ||
+        !check(cudaHostGetDevicePointer(&mapped, h_sum_, 0), err)) return false;
+    m_sum_ = (float*) mapped;
+    h_mask_ = reinterpret_cast<int32_t*>(h_sum_ + MAXT * H);
+    m_mask_ = reinterpret_cast<int32_t*>(m_sum_ + MAXT * H);
+    for (auto& p : peers_) {
+        const OnDevice on(p.remote->device_);
+        if (!check(cudaMalloc((void**) &p.sum, MAXT * H * sizeof(float)), err)) return false;
+    }
+    return true;
+}
+
+RemoteExpertOpt::~RemoteExpertOpt() {
+    for (auto& p : peers_) {
+        const OnDevice on(p.remote->device_);
+        if (p.remote->stream_) cudaStreamSynchronize(p.remote->stream_);
+        if (p.sum) cudaFree(p.sum);
+        p.remote->remote_opt_ = nullptr;
+    }
+    if (h_sum_) cudaFreeHost(h_sum_);
+}
+
+size_t RemoteExpertOpt::metadata_bytes() { return sizeof(ReduceMeta); }
+
+void RemoteExpertOpt::begin(const float* weights, int token_begin, int tokens) {
+    weights_ = weights;
+    token_begin_ = token_begin;
+    tokens_ = tokens;
+    std::memset(h_sum_ + token_begin * H, 0, tokens * H * sizeof(float));
+    std::memset(h_mask_ + token_begin * K, 0, tokens * K * sizeof(int32_t));
+}
+
+void RemoteExpertOpt::prepare(const RemoteExperts& remote, void* metadata) const {
+    auto* m = (ReduceMeta*) metadata;
+    std::fill(m->row, m->row + tokens_ * K, -1);
+    for (size_t i = 0; i < remote.original_row_.size(); ++i)
+        m->row[remote.original_row_[i]] = (int32_t) i;
+    std::memcpy(m->weight, weights_, tokens_ * K * sizeof(float));
+}
+
+bool RemoteExpertOpt::reduce(RemoteExperts& remote, const void* metadata, std::string& err) {
+    const auto& p = *std::find_if(peers_.begin(), peers_.end(), [&](const Peer& p) { return p.remote == &remote; });
+    reduce_experts<<<dim3((H + 255) / 256, tokens_), 256, 0, remote.stream_>>>(
+        remote.d_out_, (const ReduceMeta*) metadata, p.sum);
+    return check(cudaMemcpyAsync(remote.h_out_, p.sum, tokens_ * H * sizeof(float),
+                                  cudaMemcpyDeviceToHost, remote.stream_), err);
+}
+
+void RemoteExpertOpt::accumulate(const RemoteExperts& remote) {
+    for (int64_t i = 0; i < tokens_ * H; ++i) h_sum_[token_begin_ * H + i] += remote.h_out_[i];
+    for (int i = 0; i < tokens_ * K; ++i)
+        h_mask_[token_begin_ * K + i] |= remote.owned_[(size_t) i] != 0;
+}
+
+void RemoteExpertOpt::copy_rows(float* dst, const float* src, int token_begin, int tokens,
+                         const int32_t* hit_rows, const int32_t* count, void* stream) const {
+    copy_cpu_rows<<<tokens * K, 128, 0, (cudaStream_t) stream>>>(
+        (float4*) dst, (const volatile float4*) src, m_mask_ + token_begin * K, hit_rows, count);
+}
+
+void RemoteExpertOpt::combine(float* dst, float* scratch, int token_begin, int tokens,
+                       const uint32_t* skip, uint32_t ring, void* stream) const {
+    const float* sum = m_sum_ + token_begin * H;
+    if (skip) {
+        // An all-primary device-planned group bypasses the host: ignore its stale sum.
+        strata::kernels::copy_or_zero_from_mapped(scratch, sum, tokens * H, skip, ring, stream);
+        sum = scratch;
+    }
+    strata::kernels::add_inplace(dst, sum, tokens * H, stream);
+}
+
+bool RemoteExpertOpt::owns(int64_t layer, int32_t expert) const {
+    for (const auto& p : peers_)
+        if (p.remote->cache_.slot_of(layer, expert) >= 0) return true;
+    return false;
+}
+
+bool RemoteExpertOpt::adapt(const std::vector<float>& usage, const std::vector<int32_t>& primary,
+                     const std::vector<std::pair<int32_t, int32_t>>& pending, int max_swaps,
+                     ExpertSource& source) {
+    if (max_swaps <= 0) return true;
+    const auto& lay = strata::kernels::cpu::expert_layout();
+    std::vector<uint8_t> resident(primary.size(), 0);
+    for (size_t i = 0; i < primary.size(); ++i) resident[i] = primary[i] >= 0;
+    for (const auto& [i, slot] : pending) resident[(size_t) i] = 1;
+    for (const auto& p : peers_) {
+        const auto& r = *p.remote;
+        for (int32_t l = 0; l < (int32_t) r.layers_present_.size(); ++l)
+            for (int32_t e = 0; e < r.n_expert_; ++e)
+                if (r.cache_.slot_of(l, e) >= 0) resident[(size_t) l * r.n_expert_ + e] = 1;
+    }
+    struct Swap { float gain; int32_t layer, in, out; };
+    std::string err;
+    for (auto& p : peers_) {
+        auto& r = *p.remote;
+        std::vector<Swap> swaps;
+        std::vector<std::pair<float, int32_t>> candidates, victims;
+        for (int32_t l = 0; l < (int32_t) r.layers_present_.size(); ++l) {
+            candidates.clear(); victims.clear();
+            const size_t base = (size_t) l * r.n_expert_;
+            float primary_cold = 1.0e30f;
+            for (int32_t e = 0; e < r.n_expert_; ++e) {
+                const float u = usage[base + e];
+                if (primary[base + e] >= 0) primary_cold = std::min(primary_cold, u);
+                if (r.cache_.slot_of(l, e) >= 0) victims.emplace_back(u, e);
+                else if (!resident[base + e] && u >= 2.0f) candidates.emplace_back(u, e);
+            }
+            const size_t n = std::min(candidates.size(), victims.size());
+            if (!n) continue;
+            std::sort(candidates.begin(), candidates.end(), [](auto& a, auto& b) { return a.first > b.first; });
+            std::partial_sort(victims.begin(), victims.begin() + n, victims.end(),
+                              [](auto& a, auto& b) { return a.first < b.first; });
+            for (size_t i = 0; i < n; ++i) {
+                if (victims[i].first > primary_cold) break;
+                if (candidates[i].first < victims[i].first + 1.5f) break;
+                swaps.push_back({candidates[i].first - victims[i].first, l, candidates[i].second, victims[i].second});
+            }
+        }
+        std::sort(swaps.begin(), swaps.end(), [](auto& a, auto& b) { return a.gain > b.gain; });
+        if ((int) swaps.size() > max_swaps) swaps.resize((size_t) max_swaps);
+        if (swaps.empty()) continue;
+        const OnDevice on(r.device_);
+        for (const auto& s : swaps)
+            if (!r.cache_.fill_slot(r.cache_.slot_of(s.layer, s.out), source.blob(s.layer, s.in),
+                                    r.stream_, err, (int64_t) lay.blob_bytes(s.layer))) return false;
+        if (!check(cudaStreamSynchronize(r.stream_), err)) return false;
+        for (const auto& s : swaps) {
+            r.cache_.replace(s.layer, s.out, s.in);
+            const size_t base = (size_t) s.layer * r.n_expert_;
+            resident[base + s.out] = 0;
+            resident[base + s.in] = 1;
+        }
+    }
+    return true;
+}
+} // namespace strata::core

src/core/remote_experts.cpp

@@ -1,4 +1,5 @@
 #include "strata/core/remote_experts.hpp"
+#include "strata/core/remote_expert_opt.hpp"
 
 #include "strata/kernels/cpu/expert_layout.hpp"
 #include "strata/kernels/iq_kernels.hpp"
@@ -120,7 +121,7 @@ void RemoteExperts::close() {
 bool RemoteExperts::open(int device, int slots, int64_t layers, int64_t experts,
                          const std::vector<std::pair<int32_t, int32_t>>& ranked,
                          const ExpertCache& primary, ExpertSource& source,
-                         std::vector<uint8_t>& claimed, std::string& err) {
+                         std::vector<uint8_t>& claimed, std::string& err, bool auto_size) {
     close();
     int count = 0;
     if (!check(cudaGetDeviceCount(&count), "cudaGetDeviceCount", err, device)) return false;
@@ -134,29 +135,34 @@ bool RemoteExperts::open(int device, int slots, int64_t layers, int64_t experts,
     device_ = device;
     n_expert_ = experts;
     const auto& lay = strata::kernels::cpu::expert_layout();
+    size_t free_bytes = 0, total_bytes = 0;
+    if (!check(cudaMemGetInfo(&free_bytes, &total_bytes), "free memory", err, device)) { close(); return false; }
+    uint64_t needed = 0;
     std::vector<std::pair<int32_t, int32_t>> selected;
-    selected.reserve((size_t) slots);
+    selected.reserve((size_t) std::min<int64_t>(slots, layers * experts));
     std::vector<uint8_t> picked(claimed.size(), 0);
     for (const auto& pair : ranked) {
         if (pair.first < 0 || pair.first >= layers || pair.second < 0 || pair.second >= experts) continue;
         const size_t index = (size_t) pair.first * (size_t) experts + (size_t) pair.second;
         if (primary.slot_of(pair.first, pair.second) < 0 && !claimed[index] && !picked[index]) {
+            const uint64_t bytes = lay.native ? (lay.blob_bytes(pair.first) + 255) / 256 * 256 : lay.max_blob;
+            if (auto_size && needed + bytes + (512ull << 20) > free_bytes) break;
             selected.push_back(pair);
+            needed += bytes;
             picked[index] = 1;
             if ((int) selected.size() >= slots) break;
         }
     }
-    if (selected.empty()) { err = "CUDA" + std::to_string(device) + " experts: no unclaimed experts remain"; close(); return false; }
+    if (selected.empty()) {
+        err = "CUDA" + std::to_string(device) + (auto_size ? " experts: no unclaimed expert fits with 512 MiB free"
+                                                          : " experts: no unclaimed experts remain");
+        close(); return false;
+    }
     std::vector<int64_t> sizes;
     if (lay.native) {
         sizes.reserve(selected.size());
         for (const auto& pair : selected) sizes.push_back((int64_t) lay.blob_bytes(pair.first));
     }
-    size_t free_bytes = 0, total_bytes = 0;
-    if (!check(cudaMemGetInfo(&free_bytes, &total_bytes), "free memory", err, device)) { close(); return false; }
-    uint64_t needed = 0;
-    for (const auto& pair : selected)
-        needed += lay.native ? (lay.blob_bytes(pair.first) + 255) / 256 * 256 : lay.max_blob;
     // Leave room for the CUDA context, staging and later driver allocations, especially under WDDM.
     if (needed + (512ull << 20) > free_bytes) {
         err = "CUDA" + std::to_string(device) + " experts: slots leave less than 512 MiB free; reduce --expert-cache-device" + std::to_string(device);
@@ -185,17 +191,18 @@ bool RemoteExperts::open(int device, int slots, int64_t layers, int64_t experts,
     const size_t scratch = std::max<size_t>(
         (size_t) strata::kernels::moe_hit_grouped_scratch_bytes(CAP, H, FF),
         strata::kernels::native_expert_scratch_bytes(CAP, FF));
+    const size_t meta_bytes = sizeof(RemoteMeta) + (remote_opt_ ? remote_opt_->metadata_bytes() : 0);
     const bool allocated =
         check(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking), "stream", err, device) &&
         check(cudaHostAlloc((void**) &h_x_, (size_t) CAP * H * sizeof(float), cudaHostAllocPortable | cudaHostAllocMapped), "input staging", err, device) &&
         check(cudaHostAlloc((void**) &h_out_, (size_t) CAP * H * sizeof(float), cudaHostAllocPortable | cudaHostAllocMapped), "result staging", err, device) &&
-        check(cudaHostAlloc(&h_meta_, sizeof(RemoteMeta), cudaHostAllocPortable), "metadata staging", err, device) &&
+        check(cudaHostAlloc(&h_meta_, meta_bytes, cudaHostAllocPortable), "metadata staging", err, device) &&
         check(cudaMalloc((void**) &d_x_, (size_t) CAP * H * sizeof(float)), "input", err, device) &&
         check(cudaMalloc((void**) &d_out_, (size_t) CAP * H * sizeof(float)), "result", err, device) &&
         check(cudaMalloc((void**) &d_q8_, (size_t) CAP * (H / 32) * 36), "activation", err, device) &&
         check(cudaMalloc((void**) &d_scales_, (size_t) CAP * (H / 32) * sizeof(float)), "activation scales", err, device) &&
         check(cudaMalloc(&d_scratch_, scratch), "scratch", err, device) &&
-        check(cudaMalloc(&d_meta_, sizeof(RemoteMeta)), "group metadata", err, device);
+        check(cudaMalloc(&d_meta_, meta_bytes), "group metadata", err, device);
     if (!allocated) { close(); return false; }
     // Zero-copy: the helper reads its input from, and writes its compact rows into, the pinned host buffers
     // directly - two copies fewer per layer, each of which is a PCIe round trip.  STRATA_REMOTE_ZEROCOPY=0 copies.
@@ -284,13 +291,16 @@ bool RemoteExperts::begin(int64_t layer, const float* x, const int32_t* ids, int
     std::memcpy(meta->dst, dst_.data(), dst_.size() * sizeof(dst_[0]));
     std::memcpy(meta->tok, tok_.data(), tok_.size() * sizeof(tok_[0]));
     meta->count = (int32_t) group_id_.size();
+    const bool reduce = remote_opt_ && remote_opt_->active();
+    if (reduce) remote_opt_->prepare(*this, meta + 1);
     DeviceScope scope(device_);
     if (!scope.ok) { err = scope.error(device_); return false; }
     groups_ = (int32_t) group_id_.size();
     const cudaStream_t s = stream_;
     const bool staged =
         (zero_copy_ || check(cudaMemcpyAsync(d_x_, h_x_, (size_t) n_tok * H * sizeof(float), cudaMemcpyHostToDevice, s), "copy input", err, device_)) &&
-        check(cudaMemcpyAsync(d_meta_, h_meta_, sizeof(RemoteMeta), cudaMemcpyHostToDevice, s), "copy group metadata", err, device_);
+        check(cudaMemcpyAsync(d_meta_, h_meta_, sizeof(RemoteMeta) + (reduce ? remote_opt_->metadata_bytes() : 0),
+                              cudaMemcpyHostToDevice, s), "copy group metadata", err, device_);
     if (!staged) return false;
     const auto& lay = strata::kernels::cpu::expert_layout();
     if (lay.native) {
@@ -298,14 +308,16 @@ bool RemoteExperts::begin(int64_t layer, const float* x, const int32_t* ids, int
         const auto& fmt = lay.fmt[(size_t) layer];
         auto L = strata::kernels::native_expert_layout(fmt.gu_type, fmt.d_type, fmt.n_embd, fmt.n_ff);
         strata::kernels::native_expert_grouped(L, d_ptr_, d_start_, d_count_, d_dst_, d_tok_,
-                                               groups_, (int64_t) dst_.size(), d_q8_, d_scratch_, zero_copy_ ? z_out_ : d_out_, s);
+                                               groups_, (int64_t) dst_.size(), d_q8_, d_scratch_, zero_copy_ && !reduce ? z_out_ : d_out_, s);
     } else {
         strata::kernels::quantize_q8_0_scaled(zero_copy_ ? z_x_ : d_x_, d_q8_, d_scales_, n_tok * H, s);
         strata::kernels::moe_grouped_s2(d_ptr_, d_start_, d_count_, d_dst_, d_tok_,
-                                        groups_, (int64_t) dst_.size(), d_q8_, d_scales_, d_scratch_, zero_copy_ ? z_out_ : d_out_, s);
+                                        groups_, (int64_t) dst_.size(), d_q8_, d_scales_, d_scratch_, zero_copy_ && !reduce ? z_out_ : d_out_, s);
     }
-    const uint64_t compact_bytes = (uint64_t) dst_.size() * H * sizeof(float);
-    if (!zero_copy_ && !check(cudaMemcpyAsync(h_out_, d_out_, (size_t) compact_bytes, cudaMemcpyDeviceToHost, s),
+    const uint64_t compact_bytes = (uint64_t) (reduce ? n_tok : dst_.size()) * H * sizeof(float);
+    if (reduce) {
+        if (!remote_opt_->reduce(*this, (RemoteMeta*) d_meta_ + 1, err)) return false;
+    } else if (!zero_copy_ && !check(cudaMemcpyAsync(h_out_, d_out_, (size_t) compact_bytes, cudaMemcpyDeviceToHost, s),
                "copy results", err, device_)) return false;
     ++launched_layers_;
     returned_bytes_ += compact_bytes;
@@ -320,6 +332,7 @@ bool RemoteExperts::finish(float* out, std::string& err) {
     const auto w0 = std::chrono::steady_clock::now();
     if (!check(cudaStreamSynchronize(stream_), "finish", err, device_)) return false;
     ms_wait_ += std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - w0).count();
+    if (remote_opt_ && remote_opt_->active()) { remote_opt_->accumulate(*this); return true; }
     for (size_t i = 0; i < original_row_.size(); ++i)
         std::memcpy(out + (size_t) original_row_[i] * H, h_out_ + i * H, (size_t) H * sizeof(float));
     return true;

src/core/session.cpp

@@ -547,7 +547,7 @@ bool SessionLoopScratch::init(size_t parts_bytes_in, std::string& err) {
     const std::vector<int> cores = strata::kernels::cpu::physical_cores(false);
     if (!cores.empty()) {
         pinned_core = strata::kernels::cpu::pin_current_thread(cores[0]);
-        pinned = true;
+        pinned = pinned_core.valid;
     }
     return true;
 }
@@ -558,7 +558,7 @@ void SessionLoopScratch::free() {
     if (pinned) {
         strata::kernels::cpu::restore_thread_affinity(pinned_core);
         pinned = false;
-        pinned_core = -1;
+        pinned_core = {};
     }
     if (probe != nullptr) { cudaEventDestroy(probe); probe = nullptr; }
     if (y_miss != nullptr) { cudaFreeHost(y_miss); y_miss = nullptr; }

src/core/verify.cpp

@@ -1,5 +1,6 @@
 // src/core/verify.cpp - see include/strata/core/verify.hpp.
 #include "strata/core/verify.hpp"
+#include "strata/core/remote_expert_opt.hpp"
 #if defined(_WIN32)
 #include <intrin.h>
 #endif
@@ -274,12 +275,16 @@ Verifier::~Verifier() {
         slot.compare_exchange_strong(me, nullptr);
     }
     if (cs_) cudaStreamSynchronize(cs_);
+    if (sh_cs_) cudaStreamSynchronize(sh_cs_);
     for (auto& e : exec_)
         if (e) cudaGraphExecDestroy(e);
     if (commit_exec_) cudaGraphExecDestroy(commit_exec_);
     if (cs_) cudaStreamDestroy(cs_);
+    if (sh_cs_) cudaStreamDestroy(sh_cs_);
     if (copy_) { cudaStreamSynchronize(copy_); cudaStreamDestroy(copy_); }
     if (commit_done_) cudaEventDestroy(commit_done_);
+    if (ev_fork_) cudaEventDestroy(ev_fork_);
+    if (ev_join_) cudaEventDestroy(ev_join_);
     if (arena_) cudaFree(arena_);
     void* hosts[] = {h_tok_, h_step_, h_pos_, h_commit_, h_ple_, h_out_, h_x_, h_ids_, h_w_, h_seq_, h_flag_, h_ymiss_,
                      h_flagA_, h_plan_, h_flagB_};
@@ -399,6 +404,7 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         inj_ = b.take<float>(T * HC); inj2_ = b.take<float>(T * HC);
         lo_ = b.take<float>(T * (uint64_t) g.hc_lr); rs_ = b.take<float>(T * HC); xn_ = b.take<float>(T * HC * N);
         xq_ = b.take<uint8_t>(strata::kernels::native_q8_1_bytes(max_in, (int) T));
+        sh_xq_ = b.take<uint8_t>(strata::kernels::native_q8_1_bytes(max_in, (int) T));
         qkv_L_ = b.take<float>(nG * T * C); h_L_ = b.take<float>(nG * T * C);
         gate_L_ = b.take<float>(nG * T * HV); beta_L_ = b.take<float>(nG * T * HV);
         z_ = b.take<float>(T * ZV); y_ = b.take<float>(T * ZV); y_dummy_ = b.take<float>(T * ZV);
@@ -465,31 +471,56 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         err = "verify: copy stream create failed";
         return false;
     }
-    if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess) {
+    if (cudaStreamCreateWithFlags(&cs_, cudaStreamNonBlocking) != cudaSuccess ||
+        cudaStreamCreateWithFlags(&sh_cs_, cudaStreamNonBlocking) != cudaSuccess) {
         err = "verify: stream create failed";
         return false;
     }
-    if (cudaEventCreateWithFlags(&commit_done_, cudaEventDisableTiming) != cudaSuccess) {
+    if (cudaEventCreateWithFlags(&commit_done_, cudaEventDisableTiming) != cudaSuccess ||
+        cudaEventCreateWithFlags(&ev_fork_, cudaEventDisableTiming) != cudaSuccess ||
+        cudaEventCreateWithFlags(&ev_join_, cudaEventDisableTiming) != cudaSuccess) {
         err = "verify: event create failed";
         return false;
     }
+    // Check whether 100% of experts across [lb_, le_) are resident in this stage's VRAM cache.
+    // When true, every layer plans on device and writes directly into parts_ without any CPU doorbells,
+    // wait_flag_ge spins, PCIe empty launches, or moe_hit_add copies.
+    all_resident_ = false;
+    if (hits.h_res != nullptr && hits.d_res != nullptr && hits.cache_base != nullptr) {
+        const char* v_ar = std::getenv("STRATA_VERIFY_ALL_RESIDENT");
+        if (v_ar == nullptr || std::atoi(v_ar) != 0) {
+            bool all_ok = true;
+            for (int64_t l = lb_; l < le_ && all_ok; ++l) {
+                for (int64_t e = 0; e < g.n_expert; ++e) {
+                    if (hits.h_res[l * g.n_expert + e] < 0) {
+                        all_ok = false;
+                        break;
+                    }
+                }
+            }
+            all_resident_ = all_ok;
+        }
+    }
     // E-6: a layer whose routed experts are all resident is planned on the device (STRATA_VERIFY_DEVICE_PLAN=1: on;
     // exact, but neutral on RIBPC 1-2 GPUs: off by default)
     {
         const char* v = std::getenv("STRATA_VERIFY_DEVICE_PLAN");
-        device_plan_ = v != nullptr && std::atoi(v) != 0;
+        device_plan_ = !all_resident_ && (v != nullptr && std::atoi(v) != 0);
     }
-    if (device_plan_) {
-        bool ok2 = cudaMalloc((void**) &skip_, 64) == cudaSuccess && cudaMemset(skip_, 0, 64) == cudaSuccess;
+    if (all_resident_ || device_plan_) {
+        bool ok2 = true;
+        if (device_plan_)
+            ok2 = cudaMalloc((void**) &skip_, 64) == cudaSuccess && cudaMemset(skip_, 0, 64) == cudaSuccess;
         if (ok2 && hits.slot_off != nullptr && hits.n_slots > 0) {
             ok2 = cudaMalloc((void**) &slot_off_d_, (size_t) hits.n_slots * sizeof(unsigned long long)) == cudaSuccess &&
                   cudaMemcpy(slot_off_d_, hits.slot_off, (size_t) hits.n_slots * sizeof(unsigned long long),
                              cudaMemcpyHostToDevice) == cudaSuccess;
         }
-        if (!ok2) { cudaGetLastError(); device_plan_ = false; }
+        if (!ok2) { cudaGetLastError(); all_resident_ = false; device_plan_ = false; }
     }
-    std::fprintf(stderr, "strata verify: window up to %d tokens, %.1f MiB of device buffers\n", max_t,
-                 (double) count.used / 1048576.0);
+    std::fprintf(stderr, "strata verify: window up to %d tokens, %.1f MiB of device buffers%s\n", max_t,
+                 (double) count.used / 1048576.0,
+                 all_resident_ ? " (100% VRAM resident: zero-doorbell graph)" : "");
     return true;
 }
 
@@ -537,17 +568,14 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
     // per-ROW positions of the K rows [t][NKV] and the indexer query rows [t][IQ] (for batched RoPE)
     const int32_t* pos_k = pos_ + MT * NH;
     const int32_t* pos_i = pos_ + MT * (NH + NKV);
-    if (ple_on) copy_from_mapped(ple_, m_ple_, (int64_t) T * N, cs);
 
     // ---- the embeddings, broadcast to the hc streams - or, in a later stage of a layer split, the previous stage's
     // residual, pending write and inject (see set_stage)
     const int64_t HB = Verifier::handoff_floats(g);
     if (lb_ > 0) {
-        for (int t = 0; t < T; ++t) {
-            copy_from_mapped(Rt(t), hand_in_ + (size_t) t * HB, HC * N, cs);
-            copy_from_mapped(bo_ + (size_t) t * N, hand_in_ + (size_t) t * HB + HC * N, N, cs);
-            copy_from_mapped(inj2_ + (size_t) t * HC, hand_in_ + (size_t) t * HB + HC * N + N, HC, cs);
-        }
+        copy_from_mapped(R_, hand_in_, (int64_t) T * HC * N, cs);
+        copy_from_mapped(bo_, hand_in_ + (size_t) T * HC * N, (int64_t) T * N, cs);
+        copy_from_mapped(inj2_, hand_in_ + (size_t) T * (HC + 1) * N, (int64_t) T * HC, cs);
     } else if (const NativeEmbed* ne = native_embed()) {       // plan v0.3 P6: the GGUF-form table
         ne->gather_dev(tok_, T, emb_, cs);
         broadcast_streams(emb_, R_, N, (int) HC, T, cs);
@@ -595,14 +623,41 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
         // already applied it)
         bool pending = l > 0 && !cvec().covers(l - 1);
         if (l == 1 && ple_on) {
+            if (grp == 0) {
+                if (all_resident_) wait_flag_ge(m_flag_, 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());
+            static const bool ple_batch_env = [] {
+                const char* e = std::getenv("STRATA_PLE_BATCH");
+                return !e || e[0] != '0';
+            }();
+            const bool ple_batch_kv = ple_batch_env && dec_batch && n > 1 && ss.ple.w.key_bf16 != nullptr &&
+                                      ss.ple.w.value_bf16 != nullptr && ple_native_bf16_enabled() &&
+                                      ple_native_postops_enabled();
+            if (ple_batch_kv) {
+                try {
+                    bf16_gemv_fp32_mmvf_multi(ple_ + (size_t) tb * N, N, ss.ple.w.key_bf16,
+                                              xn_ + (size_t) tb * HC * N, HC * N, N, HC * N, n, cs);
+                    bf16_gemv_fp32_mmvf_multi(ple_ + (size_t) tb * N, N, ss.ple.w.value_bf16,
+                                              z_ + (size_t) tb * N, N, N, N, n, cs);
+                } catch (const std::exception& e) {
+                    err = std::string("verify PLE multi: ") + 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);
                 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);
+                    if (ple_batch_kv) {
+                        ple_block_projected(xn_ + (size_t) t * HC * N, z_ + (size_t) t * N, Rt(t), ss.ple.hist,
+                                            ss.ple.w, po, ss.ple.scratch, cs);
+                    } else {
+                        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();
@@ -792,6 +847,16 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
         }
         stamp(l, 16, grp);
         gr_read_group(1, true, inj_, inj2_);
+        static const bool sh_stream_env = [] {
+            const char* e = std::getenv("STRATA_SH_STREAM");
+            return !e || e[0] != '0';
+        }();
+        const bool sh_fork = sh_stream_env && !prof_on_ && sh_cs_ != nullptr && ev_fork_ != nullptr && ev_join_ != nullptr;
+        cudaStream_t sh_stream = sh_fork ? sh_cs_ : cs;
+        if (sh_fork) {
+            cudaEventRecord(ev_fork_, cs);
+            cudaStreamWaitEvent(sh_cs_, ev_fork_, 0);
+        }
         // the window's rows routed in 2 launches (one router GEMV reading the weight once, one
         // top-10) instead of 2 per token; every row's arithmetic is the single-token call's (STRATA_DEC_BATCH=0: old)
         const WeightRef* w_router = v.get("ffn_gate_inp.weight");
@@ -816,19 +881,31 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
             mb.logits = logits_ + t * NE; mb.ids = ids_ + t * K; mb.weights = w_ + t * K;
             if (!moe_route(wt, g, l, K, mb, mixed_ + t * N, cs, err, nullptr)) return false;
         }
-        if (device_plan_)   // E-6: every routed expert resident: this group's plan without the host
+        if (all_resident_) {
             resident_plan(ids_ + tb * K, n * (int) K, (int) K, hits_.d_res + l * g.n_expert, (int) g.n_expert,
                           hits_.cache_base, slot_off_d_, (long long) hits_.blob,
-                          plan_ + (size_t) grp * (size_t) (plan_i32_ + 16), (long long) max_t_ * K, skip_ + grp,
-                          (uint32_t) ((l - lb_) * G + grp + 1), cs);
+                          plan_ + (size_t) grp * (size_t) (plan_i32_ + 16), (long long) max_t_ * K, nullptr, 0, cs);
+        } else {
+            if (device_plan_)   // E-6: every routed expert resident: this group's plan without the host
+                resident_plan(ids_ + tb * K, n * (int) K, (int) K, hits_.d_res + l * g.n_expert, (int) g.n_expert,
+                              hits_.cache_base, slot_off_d_, (long long) hits_.blob,
+                              plan_ + (size_t) grp * (size_t) (plan_i32_ + 16), (long long) max_t_ * K, skip_ + grp,
+                              (uint32_t) ((l - lb_) * G + grp + 1), cs);
 #if defined(STRATA_USE_HIP)
-        if (g_doorbell_store)   // #649 A/B: the step's ring stored, not incremented over PCIe
-            doorbell_publish_value(xm, ids_ + tb * K, w_ + tb * K, (int64_t) n * N, (int64_t) n * K, m_x_ + tb * N,
-                                   m_ids_ + tb * K, m_w_ + tb * K, m_seq_, (uint32_t) ((l - lb_) * G + grp + 1), cs);
-        else
+            if (g_doorbell_store)   // #649 A/B: the step's ring stored, not incremented over PCIe
+                doorbell_publish_value(xm, ids_ + tb * K, w_ + tb * K, (int64_t) n * N, (int64_t) n * K, m_x_ + tb * N,
+                                       m_ids_ + tb * K, m_w_ + tb * K, m_seq_, (uint32_t) ((l - lb_) * G + grp + 1), cs);
+            else
 #endif
-        doorbell_publish(xm, ids_ + tb * K, w_ + tb * K, (int64_t) n * N, (int64_t) n * K, m_x_ + tb * N,
-                         m_ids_ + tb * K, m_w_ + tb * K, m_seq_, cs);
+            if (remote_opt_)   // #578: the helper GPUs reduce with the routing weights - publish them
+                doorbell_publish(xm, ids_ + tb * K, w_ + tb * K, (int64_t) n * N, (int64_t) n * K, m_x_ + tb * N,
+                                 m_ids_ + tb * K, m_w_ + tb * K, m_seq_, cs);
+            else {
+                const int32_t* layer_res = hits_.d_res != nullptr ? (hits_.d_res + l * g.n_expert) : nullptr;
+                doorbell_publish_res(xm, ids_ + tb * K, layer_res, (int) g.n_expert, (int64_t) n * N, (int64_t) n * K,
+                                     m_x_ + tb * N, m_ids_ + tb * K, m_seq_, cs);
+            }
+        }
         stamp(l, 17, grp);
         {
             const WeightRef *wgi = need(v, "ffn_gate_inp_shexp.weight", err), *wsg = need(v, "ffn_gate_shexp.weight", err),
@@ -841,16 +918,19 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
             nsw.gate_type = wsg->native_type; nsw.gate_data = wsg->native_data;
             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_;
-            if (dec_batch) f32_to_bf16_bulk(mixed_ + tb * N, sh_bf16_ + tb * N, (int64_t) n * N, cs);   // contiguous rows
-            else for (int t = tb; t < te; ++t) f32_to_bf16_bulk(mixed_ + t * N, sh_bf16_ + t * N, N, cs);
+            nsw.q8_1 = sh_fork ? sh_xq_ : xq_;
+            if (!shared_expert_native_bf16_enabled()) {
+                if (dec_batch) f32_to_bf16_bulk(mixed_ + tb * N, sh_bf16_ + tb * N, (int64_t) n * N, sh_stream);   // contiguous rows
+                else for (int t = tb; t < te; ++t) f32_to_bf16_bulk(mixed_ + t * N, sh_bf16_ + t * N, N, sh_stream);
+            }
             try {
                 shared_expert_multi(n, xm, 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, cs);
+                                    sh_up_ + (size_t) tb * g.n_ff, sh_g_ + tb, shared_ + tb * N, N, g.n_ff, sh_stream);
             } catch (const std::exception& e) {
                 err = std::string("verify shared expert: ") + e.what();
                 return false;
             }
+            if (sh_fork) cudaEventRecord(ev_join_, sh_cs_);
         }
         if (strata::kernels::cpu::expert_layout().native)
             quantize_q8_1_rows(xm, n, N, nat_xq_ + (size_t) tb * (N / 32) * 36, cs);
@@ -861,19 +941,17 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
     };
 
     // ---------------------------------------------------------------- post(l, group): experts, combine
+    static const bool fuse_head_gr_env = [] {
+        const char* e = std::getenv("STRATA_FUSE_HEAD_GR");
+        return e && e[0] == '1';
+    }();
+    const bool fuse_head_gr = fuse_head_gr_env && (le_ == g.n_layers) && (head_ != nullptr && head_->loaded()) &&
+                              !cvec().covers(g.n_layers - 1);
     auto post = [&](int64_t l, int grp) -> bool {
         const int tb = tb_[grp], te = te_[grp], n = te - tb;
         const uint32_t ring = (uint32_t) ((l - lb_) * G + grp + 1);
         const int64_t cap = (int64_t) n * K, capx = (int64_t) max_t_ * K;
         int32_t* pl = plan_ + (size_t) grp * (size_t) (plan_i32_ + 16);
-        if (device_plan_) {   // E-6: skipped when the device planned this group (all its experts resident)
-            wait_flag_ge_or(m_flagA_, ring, skip_ + grp, cs);
-            copy_i32_from_mapped_unless(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, skip_ + grp, ring, cs);
-        } else {
-            wait_flag_ge(m_flagA_, ring, cs);                  // the pool published this group's GPU plan
-            copy_i32_from_mapped(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, cs);
-        }
-        stamp(l, 19, grp);
         const int32_t* p_counts = pl;
         const int32_t* p_start = pl + 4;
         const int32_t* p_dst = p_start + capx + 1;
@@ -883,50 +961,70 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
         const unsigned long long* p_ptr2 = p_ptr + capx;
         const int32_t* p_start2 = pl + ptr_off + 4 * capx;
         float* hit_out = hit_out_ + (size_t) tb * K * N;
+        float* parts_out = parts_ + (size_t) tb * K * N;
         const auto& lay = strata::kernels::cpu::expert_layout();
         // plan v0.3 P6: the VRAM groups now; the PCIe groups once the copy engine has landed them in staging.
         // `gy`: the native launch's groups side by side (0: cap, one block row per possible group).
-        auto grouped = [&](const unsigned long long* gp, const int32_t* gs, const int32_t* gn, int64_t gy) {
+        auto grouped = [&](const unsigned long long* gp, const int32_t* gs, const int32_t* gn, int64_t gy, float* dst_buf) {
             if (lay.native) {
                 // the layer's GGUF formats (i-quant gate/up, Q2_0 / IQ4_NL down)
                 const auto& f = lay.fmt[(size_t) l];
                 const NativeExpertLayout L = native_expert_layout(f.gu_type, f.d_type, f.n_embd, f.n_ff);
                 native_expert_grouped(L, gp, gs, gn, p_dst, p_tok, cap, cap,
-                                      nat_xq_ + (size_t) tb * (N / 32) * 36, hit_scratch_, hit_out, cs, gy);
+                                      nat_xq_ + (size_t) tb * (N / 32) * 36, hit_scratch_, dst_buf, cs, gy);
             } else {
                 moe_grouped_s2(gp, gs, gn, p_dst, p_tok, cap, cap, hit_xq_ + (size_t) tb * (N / 32) * 34,
-                               hit_xs_ + (size_t) tb * (N / 32), hit_scratch_, hit_out, cs);
+                               hit_xs_ + (size_t) tb * (N / 32), hit_scratch_, dst_buf, cs);
             }
         };
-        grouped(p_ptr, p_start, p_counts, 0);
-        stamp(l, 20, grp);
-        if (device_plan_) wait_flag_ge_or(m_flagB_, ring, skip_ + grp, cs);
-        else wait_flag_ge(m_flagB_, ring, cs);                 // the PCIe share is in staging (DMA) or mapped
-        if (sink_.pcie_mode == 2) {                            // stage it with a copy kernel, then point at staging
-            const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs;
-            uint8_t* stage = staging_ + (size_t) (grp * per) * lay.max_blob;
-            fetch_blobs(p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs);
-            rebase_ptrs((unsigned long long*) p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), cs);
-        }
-        stamp(l, 21, grp);
-        // the PCIe share is pcie_frac of the misses: a few groups when the cache is cold, usually none (always none at
-        // pcie_frac 0), so its launch is kPcieGroupRows block rows striding over the groups, not cap of them
-        grouped(p_ptr2, p_start2, p_counts + 2, kPcieGroupRows);
-        stamp(l, 22, grp);
-        if (device_plan_) {   // no CPU share when the device planned the group: its rows are zeros
-            wait_flag_ge_or(m_flag_, ring, skip_ + grp, cs);
-            copy_or_zero_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (long long) n * K * N,
-                                     skip_ + grp, ring, cs);
+        if (all_resident_) {
+            stamp(l, 19, grp);
+            grouped(p_ptr, p_start, p_counts, 0, parts_out);
+            stamp(l, 20, grp);
         } else {
-            wait_flag_ge(m_flag_, ring, cs);               // the CPU's share is in the mapped rows
-            stamp(l, 23, grp);
-            if (dec_batch)   // only the CPU rows cross PCIe (p_dst[0, counts[1]) = the GPU's own rows)
-                copy_rows_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K, N,
-                                      p_dst, p_counts + 1, cs);
-            else
-                copy_from_mapped(parts_ + (size_t) tb * K * N, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K * N, cs);
+            if (device_plan_) {   // E-6: skipped when the device planned this group (all its experts resident)
+                wait_flag_ge_or(m_flagA_, ring, skip_ + grp, cs);
+                copy_i32_from_mapped_unless(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, skip_ + grp, ring, cs);
+            } else {
+                wait_flag_ge(m_flagA_, ring, cs);                  // the pool published this group's GPU plan
+                copy_i32_from_mapped(pl, m_plan_ + (size_t) grp * (size_t) plan_i32_, plan_i32_, cs);
+            }
+            stamp(l, 19, grp);
+            grouped(p_ptr, p_start, p_counts, 0, hit_out);
+            stamp(l, 20, grp);
+            if (device_plan_) wait_flag_ge_or(m_flagB_, ring, skip_ + grp, cs);
+            else wait_flag_ge(m_flagB_, ring, cs);                 // the PCIe share is in staging (DMA) or mapped
+            if (sink_.pcie_mode == 2) {                            // stage it with a copy kernel, then point at staging
+                const int64_t per = G == 2 ? kStagingBlobs / 2 : kStagingBlobs;
+                uint8_t* stage = staging_ + (size_t) (grp * per) * lay.max_blob;
+                fetch_blobs(p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), (int) per, cs);
+                rebase_ptrs((unsigned long long*) p_ptr2, p_counts + 2, stage, (int64_t) lay.blob_bytes(l), cs);
+            }
+            stamp(l, 21, grp);
+            // the PCIe share is pcie_frac of the misses: a few groups when the cache is cold, usually none (always none at
+            // pcie_frac 0), so its launch is kPcieGroupRows block rows striding over the groups, not cap of them
+            grouped(p_ptr2, p_start2, p_counts + 2, kPcieGroupRows, hit_out);
+            stamp(l, 22, grp);
+            if (device_plan_) {   // no CPU share when the device planned the group: its rows are zeros
+                wait_flag_ge_or(m_flag_, ring, skip_ + grp, cs);
+                copy_or_zero_from_mapped(parts_out, m_ymiss_ + (size_t) tb * K * N, (long long) n * K * N,
+                                         skip_ + grp, ring, cs);
+            } else {
+                wait_flag_ge(m_flag_, ring, cs);               // the CPU's share is in the mapped rows
+                stamp(l, 23, grp);
+                if (remote_opt_)   // #578: the helper's rows come back reduced; skip them as well
+                    remote_opt_->copy_rows(parts_out, m_ymiss_ + (size_t) tb * K * N, tb, n, p_dst, p_counts + 1, cs);
+                else if (dec_batch)   // only the CPU rows cross PCIe (p_dst[0, counts[1]) = the GPU's own rows)
+                    copy_rows_from_mapped(parts_out, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K, N,
+                                          p_dst, p_counts + 1, cs);
+                else
+                    copy_from_mapped(parts_out, m_ymiss_ + (size_t) tb * K * N, (int64_t) n * K * N, cs);
+            }
+            moe_hit_add(parts_out, hit_out, p_dst, p_counts + 1, cap, N, cs);
+        }
+        if (!prof_on_ && sh_cs_ != nullptr && ev_fork_ != nullptr && ev_join_ != nullptr) {
+            cudaStreamWaitEvent(cs, ev_join_, 0);
         }
-        moe_hit_add(parts_ + (size_t) tb * K * N, hit_out, p_dst, p_counts + 1, cap, N, cs);
         if (dec_batch && n > 1 && native_moe_combine_enabled()) {   // one launch for the window's rows
             try {
                 native_moe_combine_multi(parts_ + (size_t) tb * K * N, w_ + tb * K, shared_ + tb * N, bo_ + tb * N, N, K, n, cs);
@@ -937,10 +1035,14 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
             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;
         }
+        if (remote_opt_) remote_opt_->combine(bo_ + tb * N, y_dummy_ + tb * N, tb, n,
+                              device_plan_ ? skip_ + grp : nullptr, ring, cs);
         stamp(l, 24, grp);
         if (l == g.n_layers - 1) {
-            for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs);
-            if (cvec().covers(l)) cvec_apply(Rt(tb), l, n, HC * N, nullptr, 0, nullptr, 0, false, cs);
+            if (!fuse_head_gr) {
+                for (int t = tb; t < te; ++t) gr_write(Rt(t), bo_ + t * N, inj2_ + t * HC, gs, Rt(t), cs);
+                if (cvec().covers(l)) cvec_apply(Rt(tb), l, n, HC * N, nullptr, 0, nullptr, 0, false, cs);
+            }
         } else if (cvec().covers(l)) {
             cvec_apply(Rt(tb), l, n, HC * N, bo_ + tb * N, N, inj2_ + tb * HC, HC, true, cs);
         }
@@ -955,11 +1057,9 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
             if (l + 1 < le_ && !pre(l + 1, grp)) return false;
         }
     if (le_ < g.n_layers) {   // a layer split's earlier stage: hand the residual on, no head
-        for (int t = 0; t < T; ++t) {
-            copy_from_mapped(hand_out_ + (size_t) t * HB, Rt(t), HC * N, cs);
-            copy_from_mapped(hand_out_ + (size_t) t * HB + HC * N, bo_ + (size_t) t * N, N, cs);
-            copy_from_mapped(hand_out_ + (size_t) t * HB + HC * N + N, inj2_ + (size_t) t * HC, HC, cs);
-        }
+        copy_from_mapped(hand_out_, R_, (int64_t) T * HC * N, cs);
+        copy_from_mapped(hand_out_ + (size_t) T * HC * N, bo_, (int64_t) T * N, cs);
+        copy_from_mapped(hand_out_ + (size_t) T * (HC + 1) * N, inj2_, (int64_t) T * HC, cs);
         return true;
     }
 
@@ -969,14 +1069,33 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
         const WeightRef *hn = wt.find("output_hc_norm.weight"), *hd = wt.find("output_hc_down.weight"),
                         *hu = wt.find("output_hc_up.weight");
         if (!hn || !hd || !hu) { err = "verify: an output_hc_* weight is missing"; return false; }
-        for (int t = 0; t < T; ++t) {
-            BlockBuffers bb = ss.block;
-            bb.R = Rt(t);
-            bb.mixed = head_mixed_ + t * N;
-            if (head_ != nullptr && head_->loaded()) {
-                if (!lm_head_mix(wt, g, bb, cs, err)) return false;
-            } else if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) {
-                return false;
+        if (fuse_head_gr) {
+            FusedGrArgs fa[kFusedGrMaxT];
+            for (int t = 0; t < T; ++t) {
+                FusedGrArgs& a = fa[t];
+                a.R = Rt(t); a.R_out = Rt(t); a.apply = true;
+                a.bo_prev = bo_ + t * N; a.inj_prev = inj2_ + t * HC;
+                a.w_norm = (const float*) hn->data;
+                a.w_down = (const uint16_t*) hd->data;
+                a.w_up = (const uint16_t*) hu->data;
+                a.w_inject = nullptr;
+                a.eps = EPS;
+                a.lo = lo_ + t * g.hc_lr;
+                a.rs = rs_ + t * HC;
+                a.inject_out = head_inj_;
+                a.mixed = head_mixed_ + t * N;
+            }
+            fused_gr_read_multi(fa, T, xn_, cs);
+        } else {
+            for (int t = 0; t < T; ++t) {
+                BlockBuffers bb = ss.block;
+                bb.R = Rt(t);
+                bb.mixed = head_mixed_ + t * N;
+                if (head_ != nullptr && head_->loaded()) {
+                    if (!lm_head_mix(wt, g, bb, cs, err)) return false;
+                } else if (!lm_head(wt, g, bb, head_logits_ + (size_t) t * n_vocab_, cs, err)) {
+                    return false;
+                }
             }
         }
         if (head_ != nullptr && head_->loaded()) {
@@ -1120,7 +1239,7 @@ bool Verifier::capture_commit(std::string& err) {
                 gdn_conv_commit(conv, qkv, (int) C, commit_, cs_);
                 gdn_step_norm_multi(state, h_L_ + (size_t) gdn_index * MT * C, (int) C, gate_L_ + (size_t) gdn_index * MT * HV,
                                     beta_L_ + (size_t) gdn_index * MT * HV, z_, (const float*) wnm->data, EPS, y_dummy_,
-                                    (int) g.ssm_k_heads, (int) HV, (int) MT, commit_, cs_);
+                                    (int) g.ssm_k_heads, (int) HV, (int) MT, commit_, cs_, (int) MT);
                 ++gdn_index;
             } else {
                 const QsaState& st = ss.qsa_states[qsa_index];
@@ -1155,6 +1274,31 @@ bool Verifier::capture_commit(std::string& err) {
     return true;
 }
 
+void Verifier::stage_inputs(int T, const int32_t* tokens, int64_t pos0) {
+    using namespace strata::kernels;
+    const ModelGeometry& g = *g_;
+    const QsaShapes s = shapes_of(g);
+    int32_t* const pk = h_pos_ + (size_t) max_t_ * g.n_head;
+    int32_t* const pi = pk + (size_t) max_t_ * g.n_head_kv;
+    for (int t = 0; t < T; ++t) {
+        h_tok_[t] = tokens[t];
+        qsa_step_fill(h_step_ + t * kStepCount, pos0 + t, s);
+        const int32_t pos_t = (int32_t) (pos0 + t);
+        for (int64_t h = 0; h < g.n_head; ++h) h_pos_[t * g.n_head + h] = pos_t;
+        for (int64_t h = 0; h < g.n_head_kv; ++h) pk[t * g.n_head_kv + h] = pos_t;
+        for (int64_t h = 0; h < g.idx_q_heads; ++h) pi[t * g.idx_q_heads + h] = pos_t;
+    }
+    *(volatile uint32_t*) h_seq_ = 0;
+    *(volatile uint32_t*) h_flag_ = 0;
+    *(volatile uint32_t*) h_flagA_ = 0;
+    *(volatile uint32_t*) h_flagB_ = 0;
+    std::atomic_thread_fence(std::memory_order_seq_cst);
+    last_t_ = T;
+    last_pos0_ = pos0;
+    for (int t = 0; t < T; ++t) last_tokens_[t] = tokens[t];
+    staged_ = true;
+}
+
 bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out,
                    std::string& err) {
     using namespace strata::kernels;
@@ -1167,36 +1311,22 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
     if (!capture(T, err) || !capture_commit(err)) return false;
     VDBG("captured; staging\n");
     const Clock::time_point t0 = Clock::now();
-    const QsaShapes s = shapes_of(g);
-    for (int t = 0; t < T; ++t) {
-        h_tok_[t] = tokens[t];
-        qsa_step_fill(h_step_ + t * kStepCount, pos0 + t, s);
-        for (int64_t h = 0; h < g.n_head; ++h) h_pos_[t * g.n_head + h] = (int32_t) (pos0 + t);
-        int32_t* pk = h_pos_ + (size_t) max_t_ * g.n_head;
-        int32_t* pi = pk + (size_t) max_t_ * g.n_head_kv;
-        for (int64_t h = 0; h < g.n_head_kv; ++h) pk[t * g.n_head_kv + h] = (int32_t) (pos0 + t);
-        for (int64_t h = 0; h < g.idx_q_heads; ++h) pi[t * g.idx_q_heads + h] = (int32_t) (pos0 + t);
-    }
-    if (ss.ple.ready() && ple_stage()) {
-        uint32_t rows[kVerifyMaxT * PLE_N_HEADS];
+    if (!staged_) stage_inputs(T, tokens, pos0);
+    staged_ = false;
+    const bool do_ple = ss.ple.ready() && ple_stage();
+    uint32_t ple_rows[kVerifyMaxT * PLE_N_HEADS];
+    if (do_ple) {
         int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]};
         for (int t = 0; t < T; ++t) {
-            ngram_rows(&tokens[t], prev, 1, ss.ple.consts, rows + t * PLE_N_HEADS);
+            ngram_rows(&tokens[t], prev, 1, ss.ple.consts, ple_rows + t * PLE_N_HEADS);
             prev[0] = prev[1];
             prev[1] = tokens[t];
+            ss.ple.table->prefetch_rows(ple_rows + t * PLE_N_HEADS);
         }
-        if (!ss.ple.table->gather_batch(rows, (size_t) T, h_ple_, err)) return false;
     }
-    *(volatile uint32_t*) h_seq_ = 0;
-    *(volatile uint32_t*) h_flag_ = 0;
-    *(volatile uint32_t*) h_flagA_ = 0;
-    *(volatile uint32_t*) h_flagB_ = 0;
     if (trace_h_ != nullptr) std::memset(trace_h_, 0, trace_n_ * 8);   // #649: this window's breadcrumbs only
     std::atomic_thread_fence(std::memory_order_seq_cst);
     trace_ev("WINDOW", -1, -1, pos0 * 16 + T);
-    last_t_ = T;
-    last_pos0_ = pos0;
-    for (int t = 0; t < T; ++t) last_tokens_[t] = tokens[t];
     ms_host += ms_since(t0);
     VDBG("staged; launching\n");
     const cudaError_t le = cudaGraphLaunch(exec_[T], cs_);
@@ -1210,6 +1340,18 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
     const int gtb[2] = {0, (T + 1) / 2}, gte[2] = {G == 2 ? (T + 1) / 2 : T, T};
     const int64_t steps = (le_ - lb_) * G;
     const bool test_stall = g_test_stall > 0 && windows + 1 == g_test_stall;   // #267 test hook (off: false)
+    if (all_resident_ && !test_stall) {
+        if (do_ple) {
+            const Clock::time_point tp = Clock::now();
+            if (!ss.ple.table->gather_batch(ple_rows, (size_t) T, h_ple_, err)) return false;
+            std::atomic_thread_fence(std::memory_order_seq_cst);
+            _mm_sfence();
+            *flag = 1;
+            ms_host += ms_since(tp);
+        } else {
+            *flag = 1;
+        }
+    } else
     for (int64_t k = 0; k < steps; ++k) {
         const int64_t l = lb_ + k / G;
         const int grp = (int) (k % G);
@@ -1254,9 +1396,11 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
         set_plan_slot(grp);
         const int tb = gtb[grp], n = gte[grp] - gtb[grp];
         progress_at("verify window: the CPU experts of layer", l);
+        if (remote_opt_) remote_opt_->begin(h_w_ + (size_t) tb * ss.k, tb, n);
         if (pool != nullptr)
             pool(user, h_x_ + (size_t) tb * g.n_embd, h_ids_ + (size_t) tb * ss.k, n, ss.k,
                  h_ymiss_ + (size_t) tb * ss.k * g.n_embd, l);
+        if (remote_opt_) remote_opt_->end();
         VDBG("layer %lld served\n", (long long) l);
         if (g_trace) trace_ev(*(volatile uint32_t*) h_flagA_ == want ? "SERVED" : "SERVED-NO-PLAN-YET", k, l,
                               (int64_t) ms_since(b));   // aux: ms the CPU experts took
@@ -1273,10 +1417,22 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
             *(volatile uint32_t*) h_flagA_ = want;
             raise_flag(h_flagB_, want);
         }
+        // Layer 1's pre(1, 0) copies h_ple_ -> ple_ after Layer 0's wait_flag_ge(m_flag_, 1).
+        // By collecting PLE here at k == 0 (after publishing flagA/flagB for Layer 0 so the GPU can run
+        // Layer 0's VRAM experts, and before raising *flag = 1), the NVMe PLE reads overlap with both
+        // MtpDrafter::draft and Layer 0's attention + router + expert execution!
+        if (k == 0 && do_ple) {
+            const Clock::time_point tp = Clock::now();
+            if (!ss.ple.table->gather_batch(ple_rows, (size_t) T, h_ple_, err)) return false;
+            std::atomic_thread_fence(std::memory_order_seq_cst);
+            _mm_sfence();
+            ms_host += ms_since(tp);
+        }
         if (!(test_stall && k + 1 == steps)) *flag = want;
         ms_wait += std::chrono::duration<double, std::milli>(b - a).count();
         ms_pool += ms_since(b);
     }
+    // (#646 staged the next stage's inputs here; 0.1.39b keeps the layer split's order: each stage stages its own)
     progress_at("verify window: waiting for the GPU to finish the window (flags A/B/M raised)", (int64_t) T);
     // #267: a window the GPU never finishes (a spin kernel that never sees its flag) holds the host here; the stall
     // watchdog then releases every verifier's GPU waits (release_live_verifiers) before it ends the engine, so no
@@ -1287,8 +1443,12 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
     const cudaError_t se = cudaStreamSynchronize(cs_);
     trace_ev("SYNCED", -1, -1, (int64_t) se);
     if (se != cudaSuccess) { err = std::string("verify: ") + cudaGetErrorString(se); return false; }
+    commit_pending_ = false;
     progress_at("verify window: waiting for the expert copies", (int64_t) T);
-    cudaStreamSynchronize(copy_);   // no host function of this window may raise flag B in the next one
+    if (copy_used_) {
+        cudaStreamSynchronize(copy_);   // no host function of this window may raise flag B in the next one
+        copy_used_ = false;
+    }
     if (prof_on_ && G == 1) {       // the window's GPU stage stamps
         cudaMemcpy(prof_h_.data(), prof_, prof_h_.size() * 8, cudaMemcpyDeviceToHost);
         const int64_t L = g.n_layers;
@@ -1404,6 +1564,7 @@ void Verifier::fetch_dma(void* ctx, const uint8_t* const* src, int n, size_t byt
     Verifier* v = (Verifier*) ctx;
     const uint32_t want = v->cur_layer_ + 1;
     if (n <= 0) { raise_flag(v->h_flagB_, want); return; }
+    v->copy_used_ = true;
     uint8_t* stage = (uint8_t*) v->sink_.staging;                  // this group's half in a split window
     for (int i = 0; i < n; ++i) cudaMemcpyAsync(stage + (size_t) i * bytes, src[i], bytes, cudaMemcpyHostToDevice, v->copy_);
     FlagSet& fs = v->flag_sets_[v->cur_layer_ % (sizeof v->flag_sets_ / sizeof v->flag_sets_[0])];
@@ -1496,6 +1657,7 @@ bool Verifier::commit(int n_keep, std::string& err) {
     } else {
         const cudaError_t re = cudaEventRecord(commit_done_, cs_);
         if (re != cudaSuccess) { err = std::string("verify: commit event: ") + cudaGetErrorString(re); return false; }
+        (void) cudaStreamQuery(cs_);
         commit_pending_ = true;
     }
     if (ple_stage())   // stages that share one session must advance it once

src/kernels/cpu/expert_layout.cpp

@@ -70,6 +70,38 @@ bool cpu_avx512_ok() {
     return ok;
 }
 
+bool cpu_avx512bw_ok() {
+    static const bool ok = [] {
+        if (const char* f = std::getenv("STRATA_FORCE_AVX2"); f != nullptr && f[0] == '1') return false;
+        unsigned r[4] = {0, 0, 0, 0};
+        auto cpuid = [&](unsigned leaf, unsigned sub) {
+#if defined(_MSC_VER)
+            int x[4];
+            __cpuidex(x, (int) leaf, (int) sub);
+            for (int i = 0; i < 4; ++i) r[i] = (unsigned) x[i];
+#else
+            __cpuid_count(leaf, sub, r[0], r[1], r[2], r[3]);
+#endif
+        };
+        cpuid(0, 0);
+        if (r[0] < 7) return false;
+        cpuid(1, 0);
+        if (!((r[2] >> 27) & 1u)) return false;             // OSXSAVE
+#if defined(_MSC_VER)
+        const unsigned long long xcr0 = _xgetbv(0);
+#else
+        unsigned lo = 0, hi = 0;
+        __asm__ volatile("xgetbv" : "=a"(lo), "=d"(hi) : "c"(0));
+        const unsigned long long xcr0 = ((unsigned long long) hi << 32) | lo;
+#endif
+        if ((xcr0 & 0xE6) != 0xE6) return false;          // the OS saves the AVX-512 state
+        cpuid(7, 0);
+        const unsigned ebx = r[1];
+        return ((ebx >> 16) & 1u) && ((ebx >> 17) & 1u) && ((ebx >> 30) & 1u) && ((ebx >> 31) & 1u);
+    }();
+    return ok;
+}
+
 bool cpu_avx2_ok() {
     static const bool ok = [] {
         if (cpu_isa_cap() < 2) return false;

src/kernels/cpu/iq_avx2.cpp

@@ -106,6 +106,11 @@ inline void rows_ahead(const uint8_t* p) {
     _mm_prefetch((const char*) p + 64, _MM_HINT_T0);
 }
 
+// E-2 on the AVX-2 path (the AVX-512 kernels' STRATA_IQ_GATHER): the IQ3 grids by one AVX2 gather instead of eight
+// scalar loads assembled with set_epi32.  AVX2 gather is a different instruction with worse throughput on some cores
+// (opt-in, as on AVX-512); on the i7-12850HX it measures ~1.3x on the two 32-bit-grid formats, bit-exact.
+static const bool gather = [] { const char* v = std::getenv("STRATA_IQ256_GATHER"); return v != nullptr && std::atoi(v) != 0; }();
+
 // ---- per format: one 32-value half (values 64*j + 32*half .. +31) -> grid magnitudes, sign vector, scales
 template <int TY> struct Fmt32;
 
@@ -163,8 +168,13 @@ template <> struct Fmt32<18> {   // IQ3_XXS: d, qs[64] grid bytes, 8 x u32 (4 x
     static constexpr float K = 0.25f;
     static inline void decode(const uint8_t* b, int j, int half, __m256i& g, __m256i& sgn, __m256i& sc) {
         const uint8_t* q = b + 2 + 16 * j + 8 * half;
-        g = _mm256_set_epi32((int) iq3xxs_grid[q[7]], (int) iq3xxs_grid[q[6]], (int) iq3xxs_grid[q[5]], (int) iq3xxs_grid[q[4]],
-                             (int) iq3xxs_grid[q[3]], (int) iq3xxs_grid[q[2]], (int) iq3xxs_grid[q[1]], (int) iq3xxs_grid[q[0]]);
+        if (gather) {
+            const __m256i idx = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i*) q));
+            g = _mm256_i32gather_epi32((const int*) iq3xxs_grid, idx, 4);
+        } else {
+            g = _mm256_set_epi32((int) iq3xxs_grid[q[7]], (int) iq3xxs_grid[q[6]], (int) iq3xxs_grid[q[5]], (int) iq3xxs_grid[q[4]],
+                                 (int) iq3xxs_grid[q[3]], (int) iq3xxs_grid[q[2]], (int) iq3xxs_grid[q[1]], (int) iq3xxs_grid[q[0]]);
+        }
         const uint32_t w = u32(b + 2 + 64 + 8 * j + 4 * half);
         sgn = _mm256_set_epi64x((long long) even_signs.v[(w >> 21) & 127], (long long) even_signs.v[(w >> 14) & 127],
                                 (long long) even_signs.v[(w >> 7) & 127], (long long) even_signs.v[w & 127]);
@@ -178,9 +188,17 @@ template <> struct Fmt32<21> {   // IQ3_S: d, qs[64], qh[8], signs[32], scales[4
     static inline void decode(const uint8_t* b, int j, int half, __m256i& g, __m256i& sgn, __m256i& sc) {
         const uint8_t* q = b + 2 + 16 * j + 8 * half;
         const uint32_t h = b[66 + 2 * j + half];
+        if (gather) {
+            const __m256i bits = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7);
+            __m256i idx = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i*) q));
+            idx = _mm256_add_epi32(idx, _mm256_slli_epi32(
+                _mm256_and_si256(_mm256_srlv_epi32(_mm256_set1_epi32((int) h), bits), _mm256_set1_epi32(1)), 8));
+            g = _mm256_i32gather_epi32((const int*) iq3s_grid, idx, 4);
+        } else {
 #define G3(k) (int) iq3s_grid[q[k] | (((h >> k) & 1u) << 8)]
-        g = _mm256_set_epi32(G3(7), G3(6), G3(5), G3(4), G3(3), G3(2), G3(1), G3(0));
+            g = _mm256_set_epi32(G3(7), G3(6), G3(5), G3(4), G3(3), G3(2), G3(1), G3(0));
 #undef G3
+        }
         const uint64_t m = u64(b + 74 + 8 * j);
         sgn = sgn_vec(half ? (uint32_t) (m >> 32) : (uint32_t) m);
         const uint8_t s = b[106 + j];

src/kernels/cpu/pool.cpp

@@ -14,9 +14,11 @@
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
 #include <windows.h>
+#include "pool_affinity_win.hpp"
 #else
 #include <pthread.h>
 #include <sched.h>
+#include "pool_affinity_linux.hpp"
 #endif
 
 namespace strata::kernels::cpu {
@@ -36,7 +38,13 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
     DWORD len = 0;
     GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &len);
     if (len == 0) {
-        for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) topo.worker_cores.push_back((int) i);
+        // Preserve group identity even when detailed core topology is unavailable.
+        const WORD groups = GetActiveProcessorGroupCount();
+        for (WORD group = 0; group < groups; ++group) {
+            const DWORD count = GetActiveProcessorCount(group);
+            if (count == 0 || count == (DWORD) -1 || count > 64) continue;
+            for (DWORD i = 0; i < count; ++i) topo.worker_cores.push_back((int) group * 64 + (int) i);
+        }
         if (skip_first && !topo.worker_cores.empty()) {
             topo.host_core = topo.worker_cores.front();
             topo.worker_cores.erase(topo.worker_cores.begin());
@@ -60,10 +68,12 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
                 CoreDesc cd;
                 cd.efficiency = e->Processor.EfficiencyClass;
                 cd.has_smt = (e->Processor.Flags & LTP_PC_SMT) != 0;
-                const GROUP_AFFINITY& g = e->Processor.GroupMask[0];
-                for (int bit = 0; bit < 64; ++bit) {
-                    if (g.Mask & (1ull << bit)) {
-                        cd.lps.push_back((int) (g.Group * 64 + bit));
+                for (WORD group = 0; group < e->Processor.GroupCount; ++group) {
+                    const GROUP_AFFINITY& g = e->Processor.GroupMask[group];
+                    for (int bit = 0; bit < 64; ++bit) {
+                        if (g.Mask & (KAFFINITY(1) << bit)) {
+                            cd.lps.push_back((int) (g.Group * 64 + bit));
+                        }
                     }
                 }
                 if (!cd.lps.empty()) {
@@ -168,12 +178,8 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
     };
 
     std::vector<int> allowed;
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    if (sched_getaffinity(0, sizeof set, &set) == 0) {
-        for (int i = 0; i < CPU_SETSIZE; ++i)
-            if (CPU_ISSET(i, &set)) allowed.push_back(i);
-    } else {
+    std::vector<unsigned long> allowed_mask;
+    if (detail::get_thread_affinity(allowed_mask, &allowed) != 0) {
         for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) allowed.push_back((int) i);
     }
 
@@ -273,49 +279,60 @@ std::vector<int> physical_cores(bool skip_first, PoolAffinity affinity) {
 
 namespace {
 
-void pin_this_thread(int core) {
-    if (core < 0) return;
+bool pin_this_thread(int core, [[maybe_unused]] int worker = -1) {
+    if (core < 0) return false;
 #if defined(_WIN32)
-    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) 1 << (core & 63));
+    return detail::set_thread_group_affinity(core, worker);
 #else
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    CPU_SET(core, &set);
-    pthread_setaffinity_np(pthread_self(), sizeof set, &set);
+    const int error = detail::pin_thread_to_cpu(core);
+    if (error != 0)
+        std::fprintf(stderr, "strata cpu pool: affinity for worker %d (CPU %d) failed: %d; previous affinity kept\n",
+                     worker, core, error);
+    return error == 0;
 #endif
 }
 
 }  // namespace
 
-long long pin_current_thread(int core) {
-    if (core < 0) return -1;
+ThreadAffinity pin_current_thread(int core) {
+    if (core < 0) return {};
 #if defined(_WIN32)
-    // `SetThreadAffinityMask` RETURNS the previous mask, or 0 on failure - so 0 doubles as the error, which is
-    // why the caller must not treat it as a restorable value.
-    const DWORD_PTR prev = SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) 1 << (core & 63));
-    return prev == 0 ? -1 : (long long) prev;
+    ThreadAffinity previous;
+    ULONG target = 0;
+    if (!detail::get_thread_cpu_sets(previous.cpu_sets) || !detail::cpu_set_for_core(core, target) ||
+        !SetThreadSelectedCpuSets(GetCurrentThread(), &target, 1)) {
+        std::fprintf(stderr, "strata cpu pool: host CPU Set selection for processor %d failed: %lu; previous placement kept\n",
+                     core, (unsigned long) GetLastError());
+        return {};
+    }
+    previous.valid = true;
+    return previous;
 #else
-    cpu_set_t prev;
-    CPU_ZERO(&prev);
-    if (pthread_getaffinity_np(pthread_self(), sizeof prev, &prev) != 0) return -1;
-    unsigned long mask = 0;
-    for (int i = 0; i < CPU_SETSIZE && i < 64; ++i)
-        if (CPU_ISSET(i, &prev)) mask |= 1ul << i;
-    pin_this_thread(core);
-    return (long long) mask;
+    ThreadAffinity previous;
+    int error = detail::get_thread_affinity(previous.mask);
+    if (error == 0) error = detail::pin_thread_to_cpu(core);
+    if (error != 0) {
+        std::fprintf(stderr, "strata cpu pool: host affinity for CPU %d failed: %d; previous affinity kept\n", core, error);
+        return {};
+    }
+    previous.valid = true;
+    return previous;
 #endif
 }
 
-void restore_thread_affinity(long long previous) {
-    if (previous <= 0) return;
+void restore_thread_affinity(const ThreadAffinity& previous) {
+    if (!previous.valid) return;
 #if defined(_WIN32)
-    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) previous);
+    // Clearing an originally empty selection restores process-default/all-group eligibility without
+    // turning the caller's implicit Windows 11 affinity into an explicit single-group hard mask.
+    if (!SetThreadSelectedCpuSets(GetCurrentThread(), previous.cpu_sets.empty() ? nullptr : previous.cpu_sets.data(),
+                                 (ULONG) previous.cpu_sets.size()))
+        std::fprintf(stderr, "strata cpu pool: host CPU Set restoration failed: %lu\n",
+                     (unsigned long) GetLastError());
 #else
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    for (int i = 0; i < 64; ++i)
-        if ((previous >> i) & 1) CPU_SET(i, &set);
-    pthread_setaffinity_np(pthread_self(), sizeof set, &set);
+    const int error = detail::set_thread_affinity(previous.mask);
+    if (error != 0)
+        std::fprintf(stderr, "strata cpu pool: host affinity restoration failed: %d\n", error);
 #endif
 }
 
@@ -375,7 +392,7 @@ ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works, PoolAffinity af
     for (int i = 0; i < n_; ++i) {
         const int core = pin ? (i < (int) topo_.worker_cores.size() ? topo_.worker_cores[(size_t) i] : -1) : -1;
         threads_.emplace_back([this, i, core] {
-            pin_this_thread(core);
+            pin_this_thread(core, i);
             worker(i);
         });
     }

src/kernels/cpu/pool_affinity_linux.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#if defined(__linux__)
+#include <pthread.h>
+#include <sched.h>
+
+#include <cerrno>
+#include <climits>
+#include <cstring>
+#include <vector>
+
+namespace strata::kernels::cpu::detail {
+
+// CPU_ALLOC and the sized macros also handle kernels whose CPU mask exceeds CPU_SETSIZE.
+// Keep the native mask opaque: copy its bytes for storage and use CPU_*_S to inspect it.
+class LinuxCpuSet {
+public:
+    explicit LinuxCpuSet(int cpus) : bytes(CPU_ALLOC_SIZE(cpus)), set(CPU_ALLOC(cpus)) {
+        if (set) CPU_ZERO_S(bytes, set);
+    }
+    ~LinuxCpuSet() { CPU_FREE(set); }
+    LinuxCpuSet(const LinuxCpuSet&) = delete;
+    LinuxCpuSet& operator=(const LinuxCpuSet&) = delete;
+
+    const size_t bytes;
+    cpu_set_t* const set;
+};
+
+// pthread affinity functions return an error number directly; they do not set errno.
+// A too-small query buffer returns EINVAL, even when the number of online CPUs is small.
+inline int get_thread_affinity(std::vector<unsigned long>& mask, std::vector<int>* cpus = nullptr) {
+    for (int count = CPU_SETSIZE;;) {
+        LinuxCpuSet current(count);
+        if (!current.set) return ENOMEM;
+        const int error = pthread_getaffinity_np(pthread_self(), current.bytes, current.set);
+        if (error == 0) {
+            mask.resize(current.bytes / sizeof(unsigned long));
+            std::memcpy(mask.data(), current.set, current.bytes);
+            if (cpus) {
+                cpus->clear();
+                for (int cpu = 0; cpu < count; ++cpu)
+                    if (CPU_ISSET_S(cpu, current.bytes, current.set)) cpus->push_back(cpu);
+            }
+            return 0;
+        }
+        if (error != EINVAL || count > INT_MAX / 2) return error;
+        count *= 2;
+    }
+}
+
+inline int set_thread_affinity(const std::vector<unsigned long>& mask) {
+    constexpr size_t word_bits = sizeof(unsigned long) * CHAR_BIT;
+    if (mask.empty() || mask.size() > INT_MAX / word_bits) return EINVAL;
+    LinuxCpuSet target(static_cast<int>(mask.size() * word_bits));
+    if (!target.set) return ENOMEM;
+    std::memcpy(target.set, mask.data(), target.bytes);
+    return pthread_setaffinity_np(pthread_self(), target.bytes, target.set);
+}
+
+inline int pin_thread_to_cpu(int core) {
+    if (core < 0 || core == INT_MAX) return EINVAL;
+    LinuxCpuSet target(core < CPU_SETSIZE ? CPU_SETSIZE : core + 1);
+    if (!target.set) return ENOMEM;
+    CPU_SET_S(core, target.bytes, target.set);
+    return pthread_setaffinity_np(pthread_self(), target.bytes, target.set);
+}
+
+}  // namespace strata::kernels::cpu::detail
+#endif

src/kernels/cpu/pool_affinity_linux_test.cpp

@@ -0,0 +1,290 @@
+#ifndef _GNU_SOURCE
+#define _GNU_SOURCE
+#endif
+
+#include "strata/kernels/cpu/pool.hpp"
+
+#include <algorithm>
+#include <cerrno>
+#include <climits>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <thread>
+#include <vector>
+
+#if defined(__linux__)
+#include "pool_affinity_linux.hpp"
+
+#include <pthread.h>
+#include <sched.h>
+#include <unistd.h>
+
+namespace cpu = strata::kernels::cpu;
+
+class DynamicCpuSet {
+public:
+    explicit DynamicCpuSet(size_t count)
+        : count_(count), bytes_(CPU_ALLOC_SIZE(count)), set_(CPU_ALLOC(count)) {
+        if (set_) CPU_ZERO_S(bytes_, set_);
+    }
+
+    ~DynamicCpuSet() {
+        if (set_) CPU_FREE(set_);
+    }
+
+    DynamicCpuSet(const DynamicCpuSet&) = delete;
+    DynamicCpuSet& operator=(const DynamicCpuSet&) = delete;
+
+    cpu_set_t* get() { return set_; }
+    const cpu_set_t* get() const { return set_; }
+    size_t bytes() const { return bytes_; }
+    size_t count() const { return count_; }
+    explicit operator bool() const { return set_ != nullptr; }
+
+private:
+    size_t count_;
+    size_t bytes_;
+    cpu_set_t* set_;
+};
+
+static size_t initial_cpu_count() {
+    const long configured = sysconf(_SC_NPROCESSORS_CONF);
+    return (std::max)(size_t(128), configured > 0 ? (size_t) configured : size_t(0));
+}
+
+// Query with CPU_ALLOC rather than cpu_set_t: this must observe CPUs beyond CPU_SETSIZE too.
+static bool get_affinity(std::vector<int>& ids) {
+    size_t count = initial_cpu_count();
+    for (int attempt = 0; attempt < 8; ++attempt) {
+        DynamicCpuSet set(count);
+        if (!set) return false;
+        const int error = pthread_getaffinity_np(pthread_self(), set.bytes(), set.get());
+        if (error == 0) {
+            ids.clear();
+            for (size_t i = 0; i < set.count(); ++i)
+                if (CPU_ISSET_S(i, set.bytes(), set.get())) ids.push_back((int) i);
+            return !ids.empty();
+        }
+        if (error != EINVAL || count > (size_t(1) << 20)) {
+            std::fprintf(stderr, "pthread_getaffinity_np failed: %s\n", std::strerror(error));
+            return false;
+        }
+        count *= 2;
+    }
+    std::fprintf(stderr, "pthread_getaffinity_np still needs a larger CPU set\n");
+    return false;
+}
+
+static bool set_affinity(const std::vector<int>& ids) {
+    if (ids.empty()) return false;
+    const int highest = *std::max_element(ids.begin(), ids.end());
+    if (highest < 0) return false;
+    const size_t count = (std::max)(initial_cpu_count(), (size_t) highest + 1);
+    DynamicCpuSet set(count);
+    if (!set) return false;
+    for (int id : ids) {
+        if (id < 0 || (size_t) id >= count) return false;
+        CPU_SET_S((size_t) id, set.bytes(), set.get());
+    }
+    const int error = pthread_setaffinity_np(pthread_self(), set.bytes(), set.get());
+    if (error != 0) {
+        std::fprintf(stderr, "pthread_setaffinity_np failed: %s\n", std::strerror(error));
+        return false;
+    }
+    return true;
+}
+
+class OriginalAffinity {
+public:
+    bool capture() { return get_affinity(ids_); }
+    const std::vector<int>& ids() const { return ids_; }
+    bool restore() { return ids_.empty() || set_affinity(ids_); }
+    ~OriginalAffinity() { if (!ids_.empty()) set_affinity(ids_); }
+
+private:
+    std::vector<int> ids_;
+};
+
+static bool expect_affinity(const std::vector<int>& expected, const char* label) {
+    std::vector<int> actual;
+    if (!get_affinity(actual)) {
+        std::fprintf(stderr, "Linux affinity test: %s: cannot query current affinity\n", label);
+        return false;
+    }
+    if (actual != expected) {
+        std::fprintf(stderr, "Linux affinity test: %s: expected %zu CPUs, observed %zu\n",
+                     label, expected.size(), actual.size());
+        return false;
+    }
+    return true;
+}
+
+static bool round_trip_full_mask(const std::vector<int>& allowed) {
+    const int target = allowed.front();
+    const auto previous = cpu::pin_current_thread(target);
+    const bool pinned = previous.valid && expect_affinity({target}, "pin from full allowed mask");
+    cpu::restore_thread_affinity(previous);
+    return pinned && expect_affinity(allowed, "restore full allowed mask");
+}
+
+static bool topology_respects_current_mask() {
+    std::vector<int> allowed;
+    if (!get_affinity(allowed)) return false;
+    const auto topology = cpu::detect_cpu_topology(false);
+    if (topology.worker_cores.empty()) return false;
+    for (int core : topology.worker_cores) {
+        if (!std::binary_search(allowed.begin(), allowed.end(), core)) {
+            std::fprintf(stderr, "Linux affinity test: topology selected disallowed CPU %d\n", core);
+            return false;
+        }
+    }
+    return true;
+}
+
+static bool round_trip_sparse_nested(const std::vector<int>& allowed) {
+    if (allowed.size() < 2) {
+        std::puts("SKIP sparse/nested affinity: fewer than two allowed CPUs");
+        return true;
+    }
+    const std::vector<int> sparse{allowed.front(), allowed.back()};
+    if (!set_affinity(sparse) || !expect_affinity(sparse, "establish sparse mask") ||
+        !topology_respects_current_mask()) return false;
+
+    const auto outer = cpu::pin_current_thread(sparse.front());
+    const bool outer_pinned = outer.valid && expect_affinity({sparse.front()}, "outer pin");
+    if (!outer_pinned) {
+        cpu::restore_thread_affinity(outer);
+        return false;
+    }
+
+    const auto inner = cpu::pin_current_thread(sparse.back());
+    const bool inner_pinned = inner.valid && expect_affinity({sparse.back()}, "inner pin");
+    cpu::restore_thread_affinity(inner);
+    const bool outer_restored = expect_affinity({sparse.front()}, "restore inner pin to outer pin");
+    cpu::restore_thread_affinity(outer);
+    const bool sparse_restored = expect_affinity(sparse, "restore outer pin to sparse mask");
+    const bool full_restored = set_affinity(allowed) && expect_affinity(allowed, "restore full mask after sparse test");
+    return inner_pinned && outer_restored && sparse_restored && full_restored;
+}
+
+static bool round_trip_high_cpu(const std::vector<int>& allowed, bool require_high) {
+    const auto it = std::find_if(allowed.begin(), allowed.end(), [](int id) { return id >= 64; });
+    if (it == allowed.end()) {
+        if (require_high) {
+            std::fprintf(stderr, "Linux affinity test: --require-high-cpu requested, but no allowed CPU ID is >= 64\n");
+            return false;
+        }
+        std::puts("SKIP high-CPU affinity path: no allowed CPU ID is >= 64 (use --require-high-cpu to require it)");
+        return true;
+    }
+    const int target = allowed.front() == *it ? (allowed.size() > 1 ? allowed.back() : *it) : allowed.front();
+    if (target == *it) {
+        if (require_high) {
+            std::fprintf(stderr, "Linux affinity test: --require-high-cpu needs at least two allowed CPUs to test restoration\n");
+            return false;
+        }
+        std::puts("SKIP high-CPU affinity round trip: only one CPU is allowed, so no meaningful pin transition is possible");
+        return true;
+    }
+
+    // Restrict the caller to one high-numbered CPU, then pin to a different originally allowed CPU.
+    // Restoring the high-only mask catches truncation even when the original process mask is sparse.
+    const std::vector<int> high_only{*it};
+    if (!set_affinity(high_only) || !expect_affinity(high_only, "establish high-only mask") ||
+        !topology_respects_current_mask()) return false;
+    const auto previous = cpu::pin_current_thread(target);
+    const bool pinned = previous.valid && expect_affinity({target}, "pin away from high-only mask");
+    cpu::restore_thread_affinity(previous);
+    const bool restored = pinned && expect_affinity(high_only, "restore high-only mask");
+    if (restored) std::printf("High-CPU affinity round trip: saved CPU %d, pin target %d\n", *it, target);
+    return restored;
+}
+
+static bool negative_cpu_is_noop() {
+    std::vector<int> before;
+    if (!get_affinity(before)) return false;
+    const auto invalid = cpu::pin_current_thread(-1);
+    std::vector<int> after;
+    const bool unchanged = get_affinity(after) && before == after;
+    cpu::restore_thread_affinity(invalid);
+    return !invalid.valid && unchanged && expect_affinity(before, "negative CPU no-op");
+}
+
+static bool worker_thread_can_pin(int target) {
+    bool ok = false;
+    std::thread worker([&] {
+        std::vector<int> before;
+        if (!get_affinity(before)) return;
+        std::vector<unsigned long> saved;
+        if (cpu::detail::get_thread_affinity(saved) != 0) return;
+        const int pin_error = cpu::detail::pin_thread_to_cpu(target);
+        const bool pinned = pin_error == 0 && expect_affinity({target}, "worker thread helper pin") &&
+                            sched_getcpu() == target;
+        const int restore_error = cpu::detail::set_thread_affinity(saved);
+        ok = pinned && restore_error == 0 && expect_affinity(before, "worker thread helper restore");
+    });
+    worker.join();
+    return ok;
+}
+
+static bool failed_pin_preserves_affinity() {
+    std::vector<int> before;
+    if (!get_affinity(before)) return false;
+    // An ID outside the caller's current mask may still be online and legal, so it is not a
+    // reliable failure case. INT_MAX is rejected before allocating or changing any CPU set.
+    const int error = cpu::detail::pin_thread_to_cpu(INT_MAX);
+    std::vector<int> after;
+    const bool helper_unchanged = get_affinity(after) && before == after;
+    if (error == 0) {
+        std::fprintf(stderr, "Linux affinity test: pin_thread_to_cpu(INT_MAX) unexpectedly succeeded\n");
+        return false;
+    }
+    const auto invalid = cpu::pin_current_thread(INT_MAX);
+    cpu::restore_thread_affinity(invalid);
+    std::vector<int> after_public;
+    const bool public_unchanged = get_affinity(after_public) && before == after_public;
+    if (!helper_unchanged || invalid.valid || !public_unchanged)
+        std::fprintf(stderr, "Linux affinity test: failed pin changed the observed affinity or returned a valid snapshot\n");
+    return helper_unchanged && !invalid.valid && public_unchanged;
+}
+
+int main(int argc, char** argv) {
+    bool require_high = false;
+    for (int i = 1; i < argc; ++i) {
+        if (std::strcmp(argv[i], "--require-high-cpu") == 0) require_high = true;
+        else {
+            std::fprintf(stderr, "Unknown argument: %s\n", argv[i]);
+            return 2;
+        }
+    }
+
+    OriginalAffinity original;
+    if (!original.capture()) {
+        std::fprintf(stderr, "Linux affinity test: cannot capture original thread affinity\n");
+        return 1;
+    }
+    const auto allowed = original.ids();
+
+    bool ok = round_trip_full_mask(allowed) && worker_thread_can_pin(allowed.back());
+    if (ok) ok = round_trip_sparse_nested(allowed);
+    if (ok) ok = round_trip_high_cpu(allowed, require_high);
+    if (ok) ok = negative_cpu_is_noop();
+    if (ok) ok = failed_pin_preserves_affinity();
+    const bool restored = original.restore() && expect_affinity(allowed, "final original affinity cleanup");
+    if (!ok || !restored) return 1;
+
+    std::printf("Linux affinity test: full, sparse/nested, worker helper, invalid-target, failed-pin, and cleanup checks passed (%zu allowed CPUs)\n",
+                allowed.size());
+    return 0;
+}
+
+#else
+
+int main() {
+    std::puts("SKIP Linux affinity test: Linux is required");
+    return 0;
+}
+
+#endif

src/kernels/cpu/pool_affinity_test.cpp

@@ -0,0 +1,148 @@
+#include "strata/kernels/cpu/pool.hpp"
+#include "pool_affinity_win.hpp"
+
+#include <algorithm>
+#include <cstdio>
+#include <thread>
+#include <vector>
+
+namespace cpu = strata::kernels::cpu;
+
+static bool check(bool ok, const char* message) {
+    if (!ok) std::fprintf(stderr, "Windows affinity test: %s (error %lu)\n", message, (unsigned long) GetLastError());
+    return ok;
+}
+
+static bool select(const std::vector<ULONG>& ids) {
+    return SetThreadSelectedCpuSets(GetCurrentThread(), ids.empty() ? nullptr : ids.data(), (ULONG) ids.size()) != 0;
+}
+
+static bool selected_is(std::vector<ULONG> expected) {
+    std::vector<ULONG> actual;
+    if (!cpu::detail::get_thread_cpu_sets(actual)) return false;
+    std::sort(actual.begin(), actual.end());
+    std::sort(expected.begin(), expected.end());
+    return actual == expected;
+}
+
+// Establish reachability before the round trip: CPU Sets are soft affinity and cannot override an
+// existing job/hard-affinity restriction. Afterwards, require the same observable reachability.
+static bool observe(int core) {
+    for (int i = 0; i < 200; ++i) {
+        Sleep(1);
+        PROCESSOR_NUMBER current{};
+        GetCurrentProcessorNumberEx(&current);
+        if ((int) current.Group * 64 + current.Number == core) return true;
+    }
+    return false;
+}
+
+static bool host_round_trips(const std::vector<int>& cores, const std::vector<ULONG>& ids) {
+    if (!check(select({}), "clear initial thread CPU Sets")) return false;
+    std::vector<bool> reachable;
+    for (size_t i = 0; i < cores.size(); ++i) {
+        if (!check(select({ids[i]}), "select baseline CPU Set")) return false;
+        reachable.push_back(observe(cores[i]));
+    }
+    if (!check(select({}), "clear baseline CPU Sets")) return false;
+
+    for (size_t i = 0; i < cores.size(); ++i) {
+        const auto previous = cpu::pin_current_thread(cores[i]);
+        const bool pinned = previous.valid && previous.cpu_sets.empty() && selected_is({ids[i]});
+        cpu::restore_thread_affinity(previous);
+        if (!check(pinned && selected_is({}), "restore empty thread assignment")) return false;
+    }
+
+    // GetThreadGroupAffinity alone misses a lost implicit all-group state. Selecting each CPU Set
+    // again must still permit execution there, without ever setting a hard group affinity.
+    for (size_t i = 0; i < cores.size(); ++i) {
+        if (!check(select({ids[i]}), "select CPU Set after restore")) return false;
+        if (reachable[i]) {
+            if (!check(observe(cores[i]), "lost CPU/group eligibility after host restore")) return false;
+            std::printf("Host restore: still executes on group %d processor %d\n", cores[i] / 64, cores[i] & 63);
+        } else {
+            std::printf("SKIP scheduling observation for group %d: not reachable before test\n", cores[i] / 64);
+        }
+    }
+
+    // A nonempty selection may span groups. Nested host selections must restore the entire set.
+    if (!check(select(ids), "set existing multi-CPU selection")) return false;
+    const auto outer = cpu::pin_current_thread(cores.front());
+    const auto inner = cpu::pin_current_thread(cores.back());
+    const bool inner_ok = inner.valid && selected_is({ids.back()});
+    cpu::restore_thread_affinity(inner);
+    const bool outer_ok = outer.valid && selected_is({ids.front()});
+    cpu::restore_thread_affinity(outer);
+    if (!check(inner_ok && outer_ok && selected_is(ids), "restore nested / existing CPU Set selections")) return false;
+    const auto invalid = cpu::pin_current_thread(-1);
+    cpu::restore_thread_affinity(invalid);
+    if (!check(!invalid.valid && selected_is(ids), "invalid target must not mutate selection")) return false;
+    return check(select({}), "clear final host selection");
+}
+
+static bool inherited_default(int core, ULONG id, int temporary_core, ULONG temporary_id) {
+    ULONG count = 0;
+    if (!GetProcessDefaultCpuSets(GetCurrentProcess(), nullptr, 0, &count) &&
+        GetLastError() != ERROR_INSUFFICIENT_BUFFER) return false;
+    std::vector<ULONG> saved(count);
+    if (count && !GetProcessDefaultCpuSets(GetCurrentProcess(), saved.data(), count, &count)) return false;
+    saved.resize(count);
+    if (!check(SetProcessDefaultCpuSets(GetCurrentProcess(), &id, 1) != 0, "set process default")) return false;
+    bool ok = false;
+    std::thread host([&] {
+        if (!select({})) return;
+        const bool reachable = observe(core);
+        const auto previous = cpu::pin_current_thread(temporary_core);
+        const bool overridden = previous.valid && selected_is({temporary_id});
+        cpu::restore_thread_affinity(previous);
+        ULONG inherited = 0, required = 0;
+        ok = overridden && previous.cpu_sets.empty() && selected_is({}) &&
+             GetProcessDefaultCpuSets(GetCurrentProcess(), &inherited, 1, &required) &&
+             required == 1 && inherited == id && (!reachable || observe(core));
+    });
+    host.join();
+    const bool restored = SetProcessDefaultCpuSets(GetCurrentProcess(), saved.empty() ? nullptr : saved.data(),
+                                                  (ULONG) saved.size()) != 0;
+    return check(ok && restored, "restore inherited process default");
+}
+
+int main() {
+    const auto all = cpu::physical_cores(false);
+    std::vector<int> cores;
+    std::vector<ULONG> ids;
+    for (int core : all) {
+        if (std::any_of(cores.begin(), cores.end(), [core](int other) { return other / 64 == core / 64; })) continue;
+        ULONG id = 0;
+        if (!cpu::detail::cpu_set_for_core(core, id)) continue;
+        cores.push_back(core);
+        ids.push_back(id);
+    }
+    if (!check(!cores.empty(), "no usable CPU Sets")) return 1;
+    if (!check(cores.size() == GetActiveProcessorGroupCount(), "CPU Sets do not cover all active groups")) return 1;
+
+    // Hard pins belong on pool-owned threads; never contaminate the caller under test with one.
+    for (size_t i = 0; i < cores.size(); ++i) {
+        bool ok = false;
+        std::thread worker([&] {
+            if (!cpu::detail::set_thread_group_affinity(cores[i], (int) i)) return;
+            GROUP_AFFINITY pinned{};
+            if (!GetThreadGroupAffinity(GetCurrentThread(), &pinned)) return;
+            ok = pinned.Group == cores[i] / 64 && pinned.Mask == (KAFFINITY(1) << (cores[i] & 63));
+            // Host placement must also preserve an existing restrictive hard affinity.
+            const auto previous = cpu::pin_current_thread(cores.back());
+            cpu::restore_thread_affinity(previous);
+            GROUP_AFFINITY after{};
+            ok = ok && previous.valid && GetThreadGroupAffinity(GetCurrentThread(), &after) &&
+                 after.Group == pinned.Group && after.Mask == pinned.Mask;
+        });
+        worker.join();
+        if (!check(ok, "worker hard pin / caller hard-affinity preservation")) return 1;
+    }
+
+    bool host_ok = false;
+    std::thread host([&] { host_ok = host_round_trips(cores, ids); });
+    host.join();
+    if (!host_ok || !inherited_default(cores.back(), ids.back(), cores.front(), ids.front())) return 1;
+    std::printf("Windows affinity test: worker pins, host CPU Sets and process-default restoration passed (%zu groups)\n", cores.size());
+    return 0;
+}

src/kernels/cpu/pool_affinity_win.hpp

@@ -0,0 +1,64 @@
+#pragma once
+
+#if defined(_WIN32)
+#ifndef WIN32_LEAN_AND_MEAN
+#define WIN32_LEAN_AND_MEAN
+#endif
+#include <windows.h>
+
+#include <cstdio>
+#include <cstddef>
+#include <vector>
+
+namespace strata::kernels::cpu::detail {
+
+// Hard affinity is only for pool-owned threads, which terminate with the pool. A caller's implicit
+// all-group affinity cannot be restored from PreviousGroupAffinity, so the host uses CPU Sets below.
+inline bool set_thread_group_affinity(int core, int worker) {
+    if (core < 0) return false;
+    GROUP_AFFINITY target{};
+    target.Group = (WORD) (core / 64);
+    target.Mask = KAFFINITY(1) << (core & 63);
+    if (SetThreadGroupAffinity(GetCurrentThread(), &target, nullptr)) return true;
+    std::fprintf(stderr, "strata cpu pool: SetThreadGroupAffinity for worker %d (group %u, mask 0x%llx) failed: %lu; previous affinity kept\n",
+                 worker, (unsigned) target.Group, (unsigned long long) target.Mask, (unsigned long) GetLastError());
+    return false;
+}
+
+inline bool get_thread_cpu_sets(std::vector<ULONG>& ids) {
+    ULONG count = 0;
+    if (!GetThreadSelectedCpuSets(GetCurrentThread(), nullptr, 0, &count) &&
+        GetLastError() != ERROR_INSUFFICIENT_BUFFER) return false;
+    ids.resize(count);
+    if (count == 0) return true;
+    if (!GetThreadSelectedCpuSets(GetCurrentThread(), ids.data(), count, &count)) return false;
+    ids.resize(count);
+    return true;
+}
+
+// CPU Set IDs are opaque: find the entry by both group and group-relative processor number.
+inline bool cpu_set_for_core(int core, ULONG& id) {
+    ULONG bytes = 0;
+    if (!GetSystemCpuSetInformation(nullptr, 0, &bytes, GetCurrentProcess(), 0) &&
+        GetLastError() != ERROR_INSUFFICIENT_BUFFER) return false;
+    std::vector<unsigned char> buffer(bytes);
+    if (bytes && !GetSystemCpuSetInformation(reinterpret_cast<PSYSTEM_CPU_SET_INFORMATION>(buffer.data()),
+                                            bytes, &bytes, GetCurrentProcess(), 0)) return false;
+    constexpr size_t header_size = offsetof(SYSTEM_CPU_SET_INFORMATION, CpuSet);
+    for (size_t offset = 0; offset + header_size <= bytes;) {
+        const auto* info = reinterpret_cast<const SYSTEM_CPU_SET_INFORMATION*>(buffer.data() + offset);
+        if (info->Size < header_size || info->Size > bytes - offset) break;
+        if (info->Type == CpuSetInformation && info->Size >= sizeof(SYSTEM_CPU_SET_INFORMATION) && info->CpuSet.Group == core / 64 &&
+            info->CpuSet.LogicalProcessorIndex == (core & 63) &&
+            (!info->CpuSet.Allocated || info->CpuSet.AllocatedToTargetProcess)) {
+            id = info->CpuSet.Id;
+            return true;
+        }
+        offset += info->Size;
+    }
+    SetLastError(ERROR_NOT_FOUND);
+    return false;
+}
+
+}  // namespace strata::kernels::cpu::detail
+#endif

src/kernels/cuda/elementwise.cu

@@ -294,11 +294,34 @@ __global__ void doorbell_publish_kernel(const float* __restrict__ x, const int32
                                         const float* __restrict__ w, int n, int k, float* x_out, int32_t* ids_out,
                                         float* w_out, uint32_t* seq) {
     for (int i = threadIdx.x; i < n; i += blockDim.x) x_out[i] = x[i];
-    if ((int) threadIdx.x < k) { ids_out[threadIdx.x] = ids[threadIdx.x]; w_out[threadIdx.x] = w[threadIdx.x]; }
+    if ((int) threadIdx.x < k) {
+        ids_out[threadIdx.x] = ids[threadIdx.x];
+        if (w_out != nullptr) w_out[threadIdx.x] = w[threadIdx.x];
+    }
     __threadfence_system();
     __syncthreads();
     if (threadIdx.x == 0) {
+        *(volatile uint32_t*) seq = *(volatile uint32_t*) seq + 1u;
+    }
+}
+
+__global__ void doorbell_publish_res_kernel(const float* __restrict__ x, const int32_t* __restrict__ ids,
+                                            const int32_t* __restrict__ d_res, int n_expert, int n, int k,
+                                            float* x_out, int32_t* ids_out, uint32_t* seq) {
+    int any_miss = 0;
+    if ((int) threadIdx.x < k) {
+        const int32_t id = ids[threadIdx.x];
+        ids_out[threadIdx.x] = id;
+        if (d_res == nullptr || id < 0 || id >= n_expert || d_res[id] < 0) any_miss = 1;
+    }
+    if (__syncthreads_or(any_miss)) {
+        for (int i = threadIdx.x; i < n; i += blockDim.x) x_out[i] = x[i];
         __threadfence_system();
+    } else if ((int) threadIdx.x < k) {
+        __threadfence_system();
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
         *(volatile uint32_t*) seq = *(volatile uint32_t*) seq + 1u;
     }
 }
@@ -335,6 +358,14 @@ void doorbell_publish(const float* x, const int32_t* ids, const float* weights,
     check_launch("doorbell_publish");
 }
 
+void doorbell_publish_res(const float* x, const int32_t* ids, const int32_t* d_res, int n_expert, int64_t n, int64_t k,
+                          float* x_out, int32_t* ids_out, uint32_t* d_seq, void* stream) {
+    if (k > 1024) { std::fprintf(stderr, "doorbell_publish_res: k too large\n"); std::exit(1); }
+    doorbell_publish_res_kernel<<<1, 1024, 0, (cudaStream_t) stream>>>(x, ids, d_res, n_expert, (int) n, (int) k,
+                                                                        x_out, ids_out, d_seq);
+    check_launch("doorbell_publish_res");
+}
+
 __global__ void copy_i32_from_mapped_kernel(int32_t* __restrict__ dst, const volatile int32_t* src, int n) {
     for (int i = threadIdx.x; i < n; i += blockDim.x) dst[i] = src[i];
 }

src/kernels/cuda/fused_gr.cu

@@ -320,13 +320,13 @@ __global__ void __launch_bounds__(THREADS) gr_up_multi_kernel(GrMulti m) {
 constexpr int PR = LR + HC;                        // partial rows per (token, stream): 320 down + 4 inject
 constexpr int TQ3 = N / 8 / 32;                    // uint4 weight chunks per lane in one stream's slice (10)
 
-template <int S>
+template <int S, int MAX_T = kFusedGrMaxT, bool EXACT_T = false>
 __global__ void __launch_bounds__(THREADS) gr_down_v3_kernel(GrMulti m, float* __restrict__ part, float* __restrict__ ssg) {
     extern __shared__ __align__(16) float xs[];    // [T][N / S]
     constexpr int R2 = 1;                          // down rows per warp
-    __shared__ float red[WARPS][kFusedGrMaxT];
+    __shared__ float red[WARPS][MAX_T];
     const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
-    const int T = m.T;
+    const int T = EXACT_T ? MAX_T : m.T;
     // S = 2: each stream's 2560 columns in two halves (blockIdx.y = stream * S + half): twice the blocks
     constexpr int SL = N / S, TQS = SL / 8 / 32;
     const int rg = blockIdx.x, c = blockIdx.y / S, h = blockIdx.y - (blockIdx.y / S) * S;
@@ -337,19 +337,11 @@ __global__ void __launch_bounds__(THREADS) gr_down_v3_kernel(GrMulti m, float* _
     const int nrows = inject_block ? ((m.a[0].w_inject != nullptr && warp < HC) ? 1 : 0) : R2;
     const bool active = nrows > 0;
     const uint16_t* wbase = inject_block ? m.a[0].w_inject : m.a[0].w_down;
-    uint4 wv[R2][TQS];
+    float ssp[MAX_T];
 #pragma unroll
-    for (int r = 0; r < R2; ++r) {
-        if (r >= nrows) break;
-        const uint4* w4 = reinterpret_cast<const uint4*>(wbase + (size_t) (row0 + r) * D + (size_t) c * N + (size_t) h * SL);
-#pragma unroll
-        for (int q = 0; q < TQS; ++q) wv[r][q] = __ldg(w4 + lane + 32 * q);
-    }
-    float ssp[kFusedGrMaxT];
-#pragma unroll
-    for (int k = 0; k < kFusedGrMaxT; ++k) {
+    for (int k = 0; k < MAX_T; ++k) {
         ssp[k] = 0.0f;
-        if (k >= T) continue;
+        if (!EXACT_T && k >= T) continue;
         const FusedGrArgs& a = m.a[k];
         const float gw = a.apply ? 2.0f * sigmoidf_(a.inj_prev[c] / (float) HC) : 0.0f;
         const float4* R4 = reinterpret_cast<const float4*>(a.R + (size_t) c * N + (size_t) h * SL);
@@ -369,8 +361,8 @@ __global__ void __launch_bounds__(THREADS) gr_down_v3_kernel(GrMulti m, float* _
         }
     }
 #pragma unroll
-    for (int k = 0; k < kFusedGrMaxT; ++k) {
-        if (k >= T) break;
+    for (int k = 0; k < MAX_T; ++k) {
+        if (!EXACT_T && k >= T) break;
         const float v = warp_sum(ssp[k]);
         if (lane == 0) red[warp][k] = v;
     }
@@ -384,34 +376,36 @@ __global__ void __launch_bounds__(THREADS) gr_down_v3_kernel(GrMulti m, float* _
 #pragma unroll
     for (int r = 0; r < R2; ++r) {
         if (r >= nrows) break;
-        float acc[kFusedGrMaxT];
+        const uint4* w4 = reinterpret_cast<const uint4*>(wbase + (size_t) (row0 + r) * D + (size_t) c * N + (size_t) h * SL);
+        float acc[MAX_T];
 #pragma unroll
-        for (int k = 0; k < kFusedGrMaxT; ++k) acc[k] = 0.0f;
+        for (int k = 0; k < MAX_T; ++k) acc[k] = 0.0f;
 #pragma unroll
         for (int q = 0; q < TQS; ++q) {
             const int j = lane + 32 * q;
+            const uint4 wvq = __ldg(w4 + j);
 #pragma unroll
-            for (int k = 0; k < kFusedGrMaxT; ++k)
-                if (k < T) acc[k] += dot8(wv[r][q], xs + (size_t) k * SL + j * 8);
+            for (int k = 0; k < MAX_T; ++k)
+                if (EXACT_T || k < T) acc[k] += dot8(wvq, xs + (size_t) k * SL + j * 8);
         }
         const int prow = inject_block ? LR + warp : row0 + r;
 #pragma unroll
-        for (int k = 0; k < kFusedGrMaxT; ++k) {
-            if (k >= T) break;
+        for (int k = 0; k < MAX_T; ++k) {
+            if (!EXACT_T && k >= T) break;
             const float v = warp_sum(acc[k]);
             if (lane == 0) part[(((size_t) k * HC + c) * S + h) * PR + prow] = v;
         }
     }
 }
 
-template <int S>
+template <int S, int MAX_T = kFusedGrMaxT, bool EXACT_T = false>
 __global__ void __launch_bounds__(THREADS) gr_up_v3_kernel(GrMulti m, const float* __restrict__ part,
                                                            const float* __restrict__ ssg) {
-    __shared__ __align__(16) float lo[kFusedGrMaxT][LR];
-    __shared__ float rsS[kFusedGrMaxT][HC];
-    __shared__ float g[kFusedGrMaxT][HC][UPM_COLS];
+    __shared__ __align__(16) float lo[MAX_T][LR];
+    __shared__ float rsS[MAX_T][HC];
+    __shared__ float g[MAX_T][HC][UPM_COLS];
     const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
-    const int T = m.T;
+    const int T = EXACT_T ? MAX_T : m.T;
     const int d0 = blockIdx.x * UPM_COLS;
     constexpr int RPW = HC * UPM_COLS / WARPS;     // 8 rows per warp
     if (t < T * HC) {
@@ -470,8 +464,8 @@ __global__ void __launch_bounds__(THREADS) gr_up_v3_kernel(GrMulti m, const floa
         }
         float mine = 0.0f;
 #pragma unroll
-        for (int k = 0; k < kFusedGrMaxT; ++k) {
-            if (k >= T) break;
+        for (int k = 0; k < MAX_T; ++k) {
+            if (!EXACT_T && k >= T) break;
             float acc = dot8(wa, lo[k] + lane * 8);
             if (lane < LR / 8 - 32) acc += dot8(wb, lo[k] + (32 + lane) * 8);
             acc = warp_sum(acc);
@@ -1089,12 +1083,16 @@ void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, voi
             int optin = 0;
             cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev3);
             const int limit = optin > 0 ? optin : 48 * 1024;
-            const int need1 = (int) (kFusedGrMaxT * N * sizeof(float)), need2 = need1 / 2;
+            const int need1 = (int) (kFusedGrMaxT * N * sizeof(float)), need6 = (int) (6 * N * sizeof(float)), need2 = need1 / 2;
             int split = -1;
             if (need1 <= limit &&
-                cudaFuncSetAttribute(gr_down_v3_kernel<1>, cudaFuncAttributeMaxDynamicSharedMemorySize, need1) == cudaSuccess)
+                cudaFuncSetAttribute(gr_down_v3_kernel<1, kFusedGrMaxT, false>, cudaFuncAttributeMaxDynamicSharedMemorySize, need1) == cudaSuccess &&
+                cudaFuncSetAttribute(gr_down_v3_kernel<1, 6, true>, cudaFuncAttributeMaxDynamicSharedMemorySize, need6) == cudaSuccess) {
+                cudaFuncSetAttribute(gr_down_v3_kernel<1, 4, true>, cudaFuncAttributePreferredSharedMemoryCarveout, 100);
+                cudaFuncSetAttribute(gr_down_v3_kernel<1, 4, false>, cudaFuncAttributePreferredSharedMemoryCarveout, 100);
+                cudaFuncSetAttribute(gr_down_v3_kernel<1, 2, true>, cudaFuncAttributePreferredSharedMemoryCarveout, 100);
                 split = 1;
-            else if (need2 <= limit)
+            } else if (need2 <= limit)
                 // #375 (kenh0u): the S = 2 split (a 64 KB opt-in card: Turing) disagrees with itself in gr_parity
                 // (graph replay vs direct call) - such a card keeps the default read until that split is fixed
                 std::fprintf(stderr, "strata: STRATA_GR_V3=1 needs the two-half split on this card, which fails its "
@@ -1110,10 +1108,18 @@ void fused_gr_read_multi(const FusedGrArgs* a, int n_tok, float* xn_scratch, voi
         const size_t sm = (size_t) n_tok * (N / split) * sizeof(float);
         if (stamp_buf) gpu_stamp(stamp_buf, stamp_i0, stream);
         if (split == 2) gr_down_v3_kernel<2><<<dim3(LR / WARPS + 1, HC * 2), THREADS, sm, st>>>(m, part, ssg);
-        else gr_down_v3_kernel<1><<<dim3(LR / WARPS + 1, HC), THREADS, sm, st>>>(m, part, ssg);
+        else if (n_tok == 2) gr_down_v3_kernel<1, 2, true><<<dim3(LR / WARPS + 1, HC), THREADS, sm, st>>>(m, part, ssg);
+        else if (n_tok == 4) gr_down_v3_kernel<1, 4, true><<<dim3(LR / WARPS + 1, HC), THREADS, sm, st>>>(m, part, ssg);
+        else if (n_tok == 6) gr_down_v3_kernel<1, 6, true><<<dim3(LR / WARPS + 1, HC), THREADS, sm, st>>>(m, part, ssg);
+        else if (n_tok < 4) gr_down_v3_kernel<1, 4, false><<<dim3(LR / WARPS + 1, HC), THREADS, sm, st>>>(m, part, ssg);
+        else gr_down_v3_kernel<1, kFusedGrMaxT, false><<<dim3(LR / WARPS + 1, HC), THREADS, sm, st>>>(m, part, ssg);
         if (stamp_buf) gpu_stamp(stamp_buf, stamp_i0 + 1, stream);
         if (split == 2) gr_up_v3_kernel<2><<<UPM_BLOCKS, THREADS, 0, st>>>(m, part, ssg);
-        else gr_up_v3_kernel<1><<<UPM_BLOCKS, THREADS, 0, st>>>(m, part, ssg);
+        else if (n_tok == 2) gr_up_v3_kernel<1, 2, true><<<UPM_BLOCKS, THREADS, 0, st>>>(m, part, ssg);
+        else if (n_tok == 4) gr_up_v3_kernel<1, 4, true><<<UPM_BLOCKS, THREADS, 0, st>>>(m, part, ssg);
+        else if (n_tok == 6) gr_up_v3_kernel<1, 6, true><<<UPM_BLOCKS, THREADS, 0, st>>>(m, part, ssg);
+        else if (n_tok < 4) gr_up_v3_kernel<1, 4, false><<<UPM_BLOCKS, THREADS, 0, st>>>(m, part, ssg);
+        else gr_up_v3_kernel<1, kFusedGrMaxT, false><<<UPM_BLOCKS, THREADS, 0, st>>>(m, part, ssg);
         const cudaError_t e3 = cudaGetLastError();
         if (e3 != cudaSuccess) {
             std::fprintf(stderr, "fused_gr_read_multi v3: %s\n", cudaGetErrorString(e3));

src/kernels/cuda/iq_kernels.cu

@@ -572,18 +572,24 @@ template<int TY> struct Split;
 // per-entry ones, which call Fmt<TY>::dot per column exactly as before #242 (the launchers test kSplit at compile
 // time, so the multi kernels are never instantiated for a type without a Split).
 template<int TY> inline constexpr bool kSplit = false;
+template<int TY> inline constexpr bool kStageIqGrid = (TY == 16 || TY == 17 || TY == 22 || TY == 29);
+
 template<> inline constexpr bool kSplit<16> = true;
 template<> struct Split<16> {   // IQ2_XXS
     struct W { int g[8]; int ls; float dw; };
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ s_grid = nullptr) {
         const block_iq2_xxs* bq2 = (const block_iq2_xxs*) vbq + kbx;
         const int q2 = get_int_b2(bq2->qs, iqs);
         const uint8_t* aux8 = (const uint8_t*) &q2;
         const uint32_t aux32 = get_int_b2(bq2->qs, iqs + 1);
+        const uint2* grid_lut;
+        if constexpr (STAGE_GRID) grid_lut = (const uint2*) s_grid;
+        else grid_lut = (const uint2*) iq2xxs_grid;
         W r;
 #pragma unroll
         for (int k0 = 0; k0 < 8; k0 += 2) {
-            const uint2 grid_pos = ((const uint2*) iq2xxs_grid)[aux8[k0 / 2]];
+            const uint2 grid_pos = grid_lut[aux8[k0 / 2]];
             const uint32_t signs = unpack_ksigns(aux32 >> (7 * k0 / 2));
             const int signs0 = __vcmpne4(signs & 0x08040201, 0);
             r.g[k0 + 0] = __vsub4(grid_pos.x ^ signs0, signs0);
@@ -619,16 +625,20 @@ struct SplitLs2 {
 };
 template<> inline constexpr bool kSplit<17> = true;
 template<> struct Split<17> : SplitLs2 {   // IQ2_XS
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ s_grid = nullptr) {
         const block_iq2_xs* bq2 = (const block_iq2_xs*) vbq + kbx;
         const int2 q2_packed = make_int2(get_int_b2(bq2->qs, iqs + 0), get_int_b2(bq2->qs, iqs + 1));
         const uint16_t* q2 = (const uint16_t*) &q2_packed;
+        const uint2* grid_lut;
+        if constexpr (STAGE_GRID) grid_lut = (const uint2*) s_grid;
+        else grid_lut = (const uint2*) iq2xs_grid;
         W r;
         r.ls0 = bq2->scales[iqs / 2] & 0x0F;
         r.ls1 = bq2->scales[iqs / 2] >> 4;
 #pragma unroll
         for (int l0 = 0; l0 < 8; l0 += 2) {
-            const uint2 grid_pos = ((const uint2*) iq2xs_grid)[q2[l0 / 2] & 0x1FF];
+            const uint2 grid_pos = grid_lut[q2[l0 / 2] & 0x1FF];
             const uint32_t signs = unpack_ksigns(q2[l0 / 2] >> 9);
             const int signs0 = __vcmpne4(signs & 0x08040201, 0);
             r.g[l0 + 0] = __vsub4(grid_pos.x ^ signs0, signs0);
@@ -641,19 +651,23 @@ template<> struct Split<17> : SplitLs2 {   // IQ2_XS
 };
 template<> inline constexpr bool kSplit<22> = true;
 template<> struct Split<22> : SplitLs2 {   // IQ2_S
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ s_grid = nullptr) {
         const block_iq2_s* bq2 = (const block_iq2_s*) vbq + kbx;
         const int qs_packed = get_int_b2(bq2->qs, iqs / 2);
         const uint8_t* qs = (const uint8_t*) &qs_packed;
         const int qh = bq2->qh[iqs / 2];
         const int signs_packed_32 = get_int_b2(bq2->qs, QK_K / 32 + iqs / 2);
         const uint8_t* signs_packed_8 = (const uint8_t*) &signs_packed_32;
+        const uint64_t* grid_lut;
+        if constexpr (STAGE_GRID) grid_lut = (const uint64_t*) s_grid;
+        else grid_lut = iq2s_grid;
         W r;
         r.ls0 = bq2->scales[iqs / 2] & 0x0F;
         r.ls1 = bq2->scales[iqs / 2] >> 4;
 #pragma unroll
         for (int l0 = 0; l0 < 8; l0 += 2) {
-            const int* grid_pos = (const int*) (iq2s_grid + (qs[l0 / 2] | ((qh << (8 - l0)) & 0x300)));
+            const int* grid_pos = (const int*) (grid_lut + (qs[l0 / 2] | ((qh << (8 - l0)) & 0x300)));
             const int signs0 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x03) << 7) | ((signs_packed_8[l0 / 2] & 0x0C) << 21), 0x00000000);
             const int signs1 = __vcmpne4(((signs_packed_8[l0 / 2] & 0x30) << 3) | ((signs_packed_8[l0 / 2] & 0xC0) << 17), 0x00000000);
             r.g[l0 + 0] = __vsub4(grid_pos[0] ^ signs0, signs0);
@@ -666,7 +680,8 @@ template<> struct Split<22> : SplitLs2 {   // IQ2_S
 template<> inline constexpr bool kSplit<18> = true;
 template<> struct Split<18> {   // IQ3_XXS
     struct W { int g[8]; int ls; float dw; };
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) {
         const block_iq3_xxs* bq3 = (const block_iq3_xxs*) vbq + kbx;
         const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs + 1));
         const uint8_t* q3 = (const uint8_t*) &q3_packed;
@@ -697,7 +712,8 @@ template<> struct Split<18> {   // IQ3_XXS
 template<> inline constexpr bool kSplit<21> = true;
 template<> struct Split<21> {   // IQ3_S
     struct W { int g[8]; int ls; float dw; };
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) {
         const block_iq3_s* bq3 = (const block_iq3_s*) vbq + kbx;
         const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1));
         const uint8_t* qs = (const uint8_t*) &qs_packed;
@@ -730,15 +746,19 @@ template<> struct Split<21> {   // IQ3_S
 template<> inline constexpr bool kSplit<29> = true;
 template<> struct Split<29> {   // IQ1_M
     struct W { int g[8]; float delta[4]; int sc0, sc1; float dw; };
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ s_grid = nullptr) {
         const block_iq1_m* bq1 = (const block_iq1_m*) vbq + kbx;
         const int qs_packed = get_int_b4(bq1->qs, iqs);
         const uint8_t* qs = (const uint8_t*) &qs_packed;
+        const uint32_t* grid_lut;
+        if constexpr (STAGE_GRID) grid_lut = s_grid;
+        else grid_lut = iq1s_grid_gpu;
         W r;
 #pragma unroll
         for (int l0 = 0; l0 < 8; l0 += 2) {
             const int qhl = bq1->qh[2 * iqs + l0 / 4] >> (4 * ((l0 / 2) % 2));
-            const int grid = iq1s_grid_gpu[qs[l0 / 2] | ((qhl & 0x07) << 8)];
+            const int grid = grid_lut[qs[l0 / 2] | ((qhl & 0x07) << 8)];
             r.g[l0 + 0] = (grid >> 0) & 0x0F0F0F0F;
             r.g[l0 + 1] = (grid >> 4) & 0x0F0F0F0F;
             r.delta[l0 / 2] = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f * IQ1M_DELTA / 0x08);
@@ -773,7 +793,8 @@ template<> struct Split<29> {   // IQ1_M
 template<> inline constexpr bool kSplit<20> = true;
 template<> struct Split<20> {   // IQ4_NL
     struct W { int2 v[2]; float dw; };
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) {
         const block_iq4_nl* bq4 = (const block_iq4_nl*) vbq + kbx;
         W r;
 #pragma unroll
@@ -796,7 +817,8 @@ template<> struct Split<20> {   // IQ4_NL
 template<> inline constexpr bool kSplit<23> = true;
 template<> struct Split<23> {   // IQ4_XS
     struct W { int2 v[4]; int ls; float dw; };
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) {
         const block_iq4_xs* bq4 = (const block_iq4_xs*) vbq + kbx;
         W r;
 #pragma unroll
@@ -822,7 +844,8 @@ template<> struct Split<23> {   // IQ4_XS
 template<> inline constexpr bool kSplit<42> = true;
 template<> struct Split<42> {   // Q2_0
     struct W { int qx[4], qy[4]; float d2; };
-    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs) {
+    template<bool STAGE_GRID = false>
+    __device__ static W load(const void* __restrict__ vbq, int kbx, int iqs, const uint32_t* __restrict__ = nullptr) {
         const block_q2_0* bq2_0 = (const block_q2_0*) vbq + kbx;
         W r;
         r.d2 = bq2_0->d;
@@ -855,31 +878,166 @@ template<> struct Split<42> {   // Q2_0
 // One row against the n <= NC activations x + off[0..n) (n >= 1, warp-uniform; offsets in q8_1 blocks, 32-bit to
 // spare registers), the whole warp.  Per activation this is row_dot: the same calls k, lane-strided the same way,
 // summed in the same order, then the same warp_sum.  Only the weight side moves out of the per-activation loop.
-template<int TY, int NC>
+template<int TY, int NC, bool EXACT_N = false, bool STAGE_GRID = false>
 __device__ __forceinline__ void row_dot_multi(const uint8_t* row, const block_q8_1* x, const int (&off)[NC], int n,
-                                              int nb, int lane, float (&s)[NC]) {
+                                              int nb, int lane, float (&s)[NC], const uint32_t* __restrict__ s_grid = nullptr) {
     using F = Fmt<TY>;
     using S = Split<TY>;
 #pragma unroll
     for (int c = 0; c < NC; ++c) s[c] = 0.0f;
     for (int k = lane; k < nb * F::ipb; k += 32) {
         const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb);
-        const typename S::W w = S::load(row, kbx, iqs);
+        const typename S::W w = S::template load<STAGE_GRID>(row, kbx, iqs, s_grid);
 #pragma unroll
         for (int c = 0; c < NC; ++c)
-            if (c < n) s[c] += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
+            if (EXACT_N || c < n) s[c] += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
     }
 #pragma unroll
     for (int c = 0; c < NC; ++c)
-        if (c < n) s[c] = warp_sum(s[c]);
+        if (EXACT_N || c < n) s[c] = warp_sum(s[c]);
 }
 
+// 8-lane sub-warp row dot for K = nb * ipb == 40 (TD == 20, IQ4_NL with n_ff == 640):
+// 4 rows per warp (32 rows per 256-thread block), 100% active lanes, bitwise identical addition tree to row_dot_multi.
+template<int TY, int NC, bool EXACT_N = false>
+__device__ __forceinline__ void row_dot_40_sub8(const uint8_t* row, const block_q8_1* x, const int (&off)[NC], int n,
+                                                int t, float (&s)[NC]) {
+    using F = Fmt<TY>;
+    using S = Split<TY>;
+#pragma unroll
+    for (int c = 0; c < NC; ++c) s[c] = 0.0f;
+    for (int k = t; k < 40; k += 32) {
+        const int kbx = k / F::ipb, iqs = F::step * (k % F::ipb);
+        const typename S::W w = S::load(row, kbx, iqs);
+#pragma unroll
+        for (int c = 0; c < NC; ++c)
+            if (EXACT_N || c < n) s[c] += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
+    }
+    {
+        const int k = t + 16, kbx = k / F::ipb, iqs = F::step * (k % F::ipb);
+        const typename S::W w = S::load(row, kbx, iqs);
+#pragma unroll
+        for (int c = 0; c < NC; ++c) {
+            if (EXACT_N || c < n) {
+                float s16 = 0.0f;
+                s16 += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
+                s[c] = __fadd_rn(s[c], s16);
+            }
+        }
+    }
+    float s8[NC];
+#pragma unroll
+    for (int c = 0; c < NC; ++c) s8[c] = 0.0f;
+    {
+        const int k = t + 8, kbx = k / F::ipb, iqs = F::step * (k % F::ipb);
+        const typename S::W w = S::load(row, kbx, iqs);
+#pragma unroll
+        for (int c = 0; c < NC; ++c)
+            if (EXACT_N || c < n) s8[c] += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
+    }
+    {
+        const int k = t + 24, kbx = k / F::ipb, iqs = F::step * (k % F::ipb);
+        const typename S::W w = S::load(row, kbx, iqs);
+#pragma unroll
+        for (int c = 0; c < NC; ++c) {
+            if (EXACT_N || c < n) {
+                float s24 = 0.0f;
+                s24 += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
+                s[c] = __fadd_rn(s[c], __fadd_rn(s8[c], s24));
+            }
+        }
+    }
+#pragma unroll
+    for (int c = 0; c < NC; ++c) {
+        if (EXACT_N || c < n) {
+#pragma unroll
+            for (int o = 4; o > 0; o >>= 1) s[c] += __shfl_xor_sync(0xffffffffu, s[c], o);
+        }
+    }
+}
+
+// 16-lane sub-warp row dot for K = nb * ipb == 20 (TD == 42, Q2_0 with n_ff == 640):
+// 2 rows per warp (16 rows per 256-thread block), bitwise identical addition tree to row_dot_multi.
+template<int TY, int NC, bool EXACT_N = false>
+__device__ __forceinline__ void row_dot_20_sub16(const uint8_t* row, const block_q8_1* x, const int (&off)[NC], int n,
+                                                 int t, float (&s)[NC]) {
+    using F = Fmt<TY>;
+    using S = Split<TY>;
+#pragma unroll
+    for (int c = 0; c < NC; ++c) s[c] = 0.0f;
+    {
+        const int k = t, kbx = k / F::ipb, iqs = F::step * (k % F::ipb);
+        const typename S::W w = S::load(row, kbx, iqs);
+#pragma unroll
+        for (int c = 0; c < NC; ++c)
+            if (EXACT_N || c < n) s[c] += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
+    }
+    float s16[NC];
+#pragma unroll
+    for (int c = 0; c < NC; ++c) s16[c] = 0.0f;
+    if (t < 4) {
+        const int k = t + 16, kbx = k / F::ipb, iqs = F::step * (k % F::ipb);
+        const typename S::W w = S::load(row, kbx, iqs);
+#pragma unroll
+        for (int c = 0; c < NC; ++c)
+            if (EXACT_N || c < n) s16[c] += S::apply(w, x + off[c] + kbx * (F::qk / 32), iqs);
+    }
+#pragma unroll
+    for (int c = 0; c < NC; ++c) {
+        if (EXACT_N || c < n) {
+            s[c] = __fadd_rn(s[c], s16[c]);
+#pragma unroll
+            for (int o = 8; o > 0; o >>= 1) s[c] += __shfl_xor_sync(0xffffffffu, s[c], o);
+        }
+    }
+}
+
+template<int TY, bool STAGE_GRID = kStageIqGrid<TY>>
+__device__ __forceinline__ const uint32_t* stage_iq_grid(uint32_t* s_buf, int tid, int nthreads) {
+    if constexpr (!STAGE_GRID) {
+        return nullptr;
+    } else if constexpr (TY == 16) {
+        uint2* dst = (uint2*) s_buf;
+        const uint2* src = (const uint2*) iq2xxs_grid;
+        for (int i = tid; i < 256; i += nthreads) dst[i] = src[i];
+        __syncthreads();
+        return s_buf;
+    } else if constexpr (TY == 17) {
+        uint2* dst = (uint2*) s_buf;
+        const uint2* src = (const uint2*) iq2xs_grid;
+        for (int i = tid; i < 512; i += nthreads) dst[i] = src[i];
+        __syncthreads();
+        return s_buf;
+    } else if constexpr (TY == 22) {
+        uint2* dst = (uint2*) s_buf;
+        const uint2* src = (const uint2*) iq2s_grid;
+        for (int i = tid; i < 1024; i += nthreads) dst[i] = src[i];
+        __syncthreads();
+        return s_buf;
+    } else if constexpr (TY == 29) {
+        uint2* dst = (uint2*) s_buf;
+        const uint2* src = (const uint2*) iq1s_grid_gpu;
+        for (int i = tid; i < 1024; i += nthreads) dst[i] = src[i];
+        __syncthreads();
+        return s_buf;
+    } else {
+        return nullptr;
+    }
+}
+
+template<int TY, bool STAGE_GRID = kStageIqGrid<TY>>
+struct IqGridWords {
+    static constexpr int value = !STAGE_GRID ? 4 : (TY == 22 || TY == 29) ? 2048 : (TY == 17) ? 1024 : (TY == 16) ? 512 : 4;
+};
+
 // mmvq_kernel with the columns taken NC at a time.  After warp_sum every lane holds the same sum, so lane c stores
 // column c.
-template<int TY, int NC>
+template<int TY, int NC, bool STAGE_GRID = kStageIqGrid<TY>>
 __global__ void __launch_bounds__(128) mmvq_multi_kernel(const uint8_t* __restrict__ w, size_t row_bytes,
                                                          const block_q8_1* __restrict__ x, float* __restrict__ y,
                                                          int n_in, int n_out, int ncols) {
+    __shared__ alignas(16) uint32_t s_grid_buf[IqGridWords<TY, STAGE_GRID>::value];
+    const uint32_t* s_grid = stage_iq_grid<TY, STAGE_GRID>(s_grid_buf, threadIdx.y * 32 + threadIdx.x, 128);
     const int row = blockIdx.x * 4 + threadIdx.y;
     if (row >= n_out) return;
     const int lane = threadIdx.x;
@@ -891,7 +1049,7 @@ __global__ void __launch_bounds__(128) mmvq_multi_kernel(const uint8_t* __restri
 #pragma unroll
         for (int c = 0; c < NC; ++c) off[c] = (c0 + min(c, n - 1)) * xb;
         float s[NC];
-        row_dot_multi<TY, NC>(wr, x, off, n, nb, lane, s);
+        row_dot_multi<TY, NC, false, STAGE_GRID>(wr, x, off, n, nb, lane, s, s_grid);
 #pragma unroll
         for (int c = 0; c < NC; ++c)
             if (c < n && lane == c) y[(size_t) (c0 + c) * n_out + row] = s[c];
@@ -937,13 +1095,17 @@ __global__ void __launch_bounds__(256) native_gu_kernel(const unsigned long long
 // has halves of <= 4), so 4 takes such windows in one pass; 8 would take ~64-80 registers against ~48 (ptxas -v).
 constexpr int GRP_NC = 4;
 
-template<int TG>
+template<int TG, bool STAGE_GRID = kStageIqGrid<TG>>
 __global__ void __launch_bounds__(256) native_gu_multi_kernel(const unsigned long long* __restrict__ grp_ptr,
                                                               const int32_t* __restrict__ grp_start,
                                                               const int32_t* __restrict__ n_groups,
                                                               const int32_t* __restrict__ ent_tok,
                                                               const block_q8_1* __restrict__ xq, NativeExpertLayout L,
                                                               float* __restrict__ gate, float* __restrict__ up) {
+    const int ng = *n_groups;
+    if (blockIdx.y >= ng) return;
+    __shared__ alignas(16) uint32_t s_grid_buf[IqGridWords<TG, STAGE_GRID>::value];
+    const uint32_t* s_grid = stage_iq_grid<TG, STAGE_GRID>(s_grid_buf, threadIdx.x, 256);
     const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
     const int row = blockIdx.x * GU_ROWS + warp;             // 0 .. 2*n_ff
     if (row >= 2 * L.n_ff) return;
@@ -952,20 +1114,31 @@ __global__ void __launch_bounds__(256) native_gu_multi_kernel(const unsigned lon
     const size_t w_off = (is_up ? L.up_off : 0) + (size_t) r * L.gu_row;
     const int nb = (int) (L.n_embd / Fmt<TG>::qk), xb = (int) (L.n_embd / 32);
     float* dst = is_up ? up : gate;
-    const int ng = *n_groups;
     for (int g = blockIdx.y; g < ng; g += gridDim.y) {   // the group stride, as native_gu_kernel
         const uint8_t* wr = (const uint8_t*) grp_ptr[g] + w_off;
         const int e0 = grp_start[g], e1 = grp_start[g + 1];
         for (int e = e0; e < e1; e += GRP_NC) {
             const int n = min(GRP_NC, e1 - e);
-            int off[GRP_NC];
+            if (n == 1) {
+                const int off1[1] = { ent_tok[e] * xb };
+                float s1[1];
+                row_dot_multi<TG, 1, true, STAGE_GRID>(wr, xq, off1, 1, nb, lane, s1, s_grid);
+                if (lane == 0) dst[(size_t) e * L.n_ff + r] = s1[0];
+            } else if (n == 2) {
+                const int off2[2] = { ent_tok[e] * xb, ent_tok[e + 1] * xb };
+                float s2[2];
+                row_dot_multi<TG, 2, true, STAGE_GRID>(wr, xq, off2, 2, nb, lane, s2, s_grid);
+                if (lane < 2) dst[(size_t) (e + lane) * L.n_ff + r] = s2[lane];
+            } else {
+                int off[GRP_NC];
 #pragma unroll
-            for (int c = 0; c < GRP_NC; ++c) off[c] = ent_tok[e + min(c, n - 1)] * xb;
-            float s[GRP_NC];
-            row_dot_multi<TG, GRP_NC>(wr, xq, off, n, nb, lane, s);
+                for (int c = 0; c < GRP_NC; ++c) off[c] = ent_tok[e + min(c, n - 1)] * xb;
+                float s[GRP_NC];
+                row_dot_multi<TG, GRP_NC, false, STAGE_GRID>(wr, xq, off, n, nb, lane, s, s_grid);
 #pragma unroll
-            for (int c = 0; c < GRP_NC; ++c)
-                if (c < n && lane == c) dst[(size_t) (e + c) * L.n_ff + r] = s[c];
+                for (int c = 0; c < GRP_NC; ++c)
+                    if (c < n && lane == c) dst[(size_t) (e + c) * L.n_ff + r] = s[c];
+            }
         }
     }
 }
@@ -1009,25 +1182,148 @@ __global__ void __launch_bounds__(256) native_down_multi_kernel(const unsigned l
                                                                 const int32_t* __restrict__ ent_dst,
                                                                 const block_q8_1* __restrict__ hq, NativeExpertLayout L,
                                                                 float* __restrict__ out) {
+    const int ng = *n_groups;
+    if (blockIdx.y >= ng) return;
+    __shared__ uint32_t s_hq_buf[GRP_NC * 20 * 9];
+    const int nb = (int) (L.n_ff / Fmt<TD>::qk), hb = (int) (L.n_ff / 32);
+    if constexpr (TD == 20) {
+        if (hb == 20) {
+            const int r = blockIdx.x * 32 + (threadIdx.x >> 3);
+            const int t = threadIdx.x & 7;
+            const bool valid_r = (r < L.n_embd);
+            const size_t w_off = L.down_off + (size_t) (valid_r ? r : 0) * L.d_row;
+            const block_q8_1* s_hq = reinterpret_cast<const block_q8_1*>(s_hq_buf);
+            for (int g = blockIdx.y; g < ng; g += gridDim.y) {
+                const uint8_t* wr = (const uint8_t*) grp_ptr[g] + w_off;
+                const int e0 = grp_start[g], e1 = grp_start[g + 1];
+                for (int e = e0; e < e1; e += GRP_NC) {
+                    const int n = min(GRP_NC, e1 - e);
+                    const uint32_t* src = reinterpret_cast<const uint32_t*>(hq + (size_t) e * 20);
+                    const int words = n * (20 * 9);
+                    __syncthreads();
+                    for (int i = threadIdx.x; i < words; i += 256) s_hq_buf[i] = src[i];
+                    __syncthreads();
+                    if (valid_r) {
+                        if (n == 1) {
+                            const int off1[1] = { 0 };
+                            float s1[1];
+                            row_dot_40_sub8<TD, 1, true>(wr, s_hq, off1, 1, t, s1);
+                            if (t == 0) out[(size_t) ent_dst[e] * L.n_embd + r] = s1[0];
+                        } else if (n == 2) {
+                            const int off2[2] = { 0, 20 };
+                            float s2[2];
+                            row_dot_40_sub8<TD, 2, true>(wr, s_hq, off2, 2, t, s2);
+                            if (t < 2) out[(size_t) ent_dst[e + t] * L.n_embd + r] = s2[t];
+                        } else {
+                            int off[GRP_NC];
+#pragma unroll
+                            for (int c = 0; c < GRP_NC; ++c) off[c] = min(c, n - 1) * 20;
+                            float s[GRP_NC];
+                            row_dot_40_sub8<TD, GRP_NC>(wr, s_hq, off, n, t, s);
+#pragma unroll
+                            for (int c = 0; c < GRP_NC; ++c)
+                                if (c < n && t == c) out[(size_t) ent_dst[e + c] * L.n_embd + r] = s[c];
+                        }
+                    }
+                }
+            }
+            return;
+        }
+    } else if constexpr (TD == 42) {
+        if (hb == 20) {
+            const int r = blockIdx.x * 16 + (threadIdx.x >> 4);
+            const int t = threadIdx.x & 15;
+            const bool valid_r = (r < L.n_embd);
+            const size_t w_off = L.down_off + (size_t) (valid_r ? r : 0) * L.d_row;
+            const block_q8_1* s_hq = reinterpret_cast<const block_q8_1*>(s_hq_buf);
+            for (int g = blockIdx.y; g < ng; g += gridDim.y) {
+                const uint8_t* wr = (const uint8_t*) grp_ptr[g] + w_off;
+                const int e0 = grp_start[g], e1 = grp_start[g + 1];
+                for (int e = e0; e < e1; e += GRP_NC) {
+                    const int n = min(GRP_NC, e1 - e);
+                    const uint32_t* src = reinterpret_cast<const uint32_t*>(hq + (size_t) e * 20);
+                    const int words = n * (20 * 9);
+                    __syncthreads();
+                    for (int i = threadIdx.x; i < words; i += 256) s_hq_buf[i] = src[i];
+                    __syncthreads();
+                    if (valid_r) {
+                        if (n == 1) {
+                            const int off1[1] = { 0 };
+                            float s1[1];
+                            row_dot_20_sub16<TD, 1, true>(wr, s_hq, off1, 1, t, s1);
+                            if (t == 0) out[(size_t) ent_dst[e] * L.n_embd + r] = s1[0];
+                        } else if (n == 2) {
+                            const int off2[2] = { 0, 20 };
+                            float s2[2];
+                            row_dot_20_sub16<TD, 2, true>(wr, s_hq, off2, 2, t, s2);
+                            if (t < 2) out[(size_t) ent_dst[e + t] * L.n_embd + r] = s2[t];
+                        } else {
+                            int off[GRP_NC];
+#pragma unroll
+                            for (int c = 0; c < GRP_NC; ++c) off[c] = min(c, n - 1) * 20;
+                            float s[GRP_NC];
+                            row_dot_20_sub16<TD, GRP_NC>(wr, s_hq, off, n, t, s);
+#pragma unroll
+                            for (int c = 0; c < GRP_NC; ++c)
+                                if (c < n && t == c) out[(size_t) ent_dst[e + c] * L.n_embd + r] = s[c];
+                        }
+                    }
+                }
+            }
+            return;
+        }
+    }
     const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
     const int r = blockIdx.x * 8 + warp;
-    if (r >= L.n_embd) return;
-    const size_t w_off = L.down_off + (size_t) r * L.d_row;
-    const int nb = (int) (L.n_ff / Fmt<TD>::qk), hb = (int) (L.n_ff / 32);
-    const int ng = *n_groups;
+    const bool valid_r = (r < L.n_embd);
+    const size_t w_off = L.down_off + (size_t) (valid_r ? r : 0) * L.d_row;
     for (int g = blockIdx.y; g < ng; g += gridDim.y) {   // the group stride, as native_gu_kernel
         const uint8_t* wr = (const uint8_t*) grp_ptr[g] + w_off;
         const int e0 = grp_start[g], e1 = grp_start[g + 1];
-        for (int e = e0; e < e1; e += GRP_NC) {
-            const int n = min(GRP_NC, e1 - e);
-            int off[GRP_NC];
+        if (hb == 20) {
+            const block_q8_1* s_hq = reinterpret_cast<const block_q8_1*>(s_hq_buf);
+            for (int e = e0; e < e1; e += GRP_NC) {
+                const int n = min(GRP_NC, e1 - e);
+                const uint32_t* src = reinterpret_cast<const uint32_t*>(hq + (size_t) e * 20);
+                const int words = n * (20 * 9);
+                __syncthreads();
+                for (int i = threadIdx.x; i < words; i += 256) s_hq_buf[i] = src[i];
+                __syncthreads();
+                if (valid_r) {
+                    if (n == 1) {
+                        const int off1[1] = { 0 };
+                        float s1[1];
+                        row_dot_multi<TD, 1, true>(wr, s_hq, off1, 1, nb, lane, s1);
+                        if (lane == 0) out[(size_t) ent_dst[e] * L.n_embd + r] = s1[0];
+                    } else if (n == 2) {
+                        const int off2[2] = { 0, 20 };
+                        float s2[2];
+                        row_dot_multi<TD, 2, true>(wr, s_hq, off2, 2, nb, lane, s2);
+                        if (lane < 2) out[(size_t) ent_dst[e + lane] * L.n_embd + r] = s2[lane];
+                    } else {
+                        int off[GRP_NC];
 #pragma unroll
-            for (int c = 0; c < GRP_NC; ++c) off[c] = (e + min(c, n - 1)) * hb;
-            float s[GRP_NC];
-            row_dot_multi<TD, GRP_NC>(wr, hq, off, n, nb, lane, s);
+                        for (int c = 0; c < GRP_NC; ++c) off[c] = min(c, n - 1) * 20;
+                        float s[GRP_NC];
+                        row_dot_multi<TD, GRP_NC>(wr, s_hq, off, n, nb, lane, s);
 #pragma unroll
-            for (int c = 0; c < GRP_NC; ++c)
-                if (c < n && lane == c) out[(size_t) ent_dst[e + c] * L.n_embd + r] = s[c];
+                        for (int c = 0; c < GRP_NC; ++c)
+                            if (c < n && lane == c) out[(size_t) ent_dst[e + c] * L.n_embd + r] = s[c];
+                    }
+                }
+            }
+        } else if (valid_r) {
+            for (int e = e0; e < e1; e += GRP_NC) {
+                const int n = min(GRP_NC, e1 - e);
+                int off[GRP_NC];
+#pragma unroll
+                for (int c = 0; c < GRP_NC; ++c) off[c] = (e + min(c, n - 1)) * hb;
+                float s[GRP_NC];
+                row_dot_multi<TD, GRP_NC>(wr, hq, off, n, nb, lane, s);
+#pragma unroll
+                for (int c = 0; c < GRP_NC; ++c)
+                    if (c < n && lane == c) out[(size_t) ent_dst[e + c] * L.n_embd + r] = s[c];
+            }
         }
     }
 }
@@ -1410,6 +1706,18 @@ bool env_on(const char* name) {
 }
 // STRATA_OLD_IQ_MMVQ=1 keeps the per-column kernels (bitwise equal to the new ones; kept for A/B timing)
 bool g_old_kernels = env_on("STRATA_OLD_IQ_MMVQ");
+// STRATA_IQ_STAGE_GRID=0 disables staging 64-bit i-quant codebook tables into shared memory
+bool g_stage_grid = [] {
+    const char* v = std::getenv("STRATA_IQ_STAGE_GRID");
+    return v == nullptr || v[0] == '\0' || v[0] != '0';
+}();
+// The single-matrix mmvq (128-thread blocks, one per 4 rows) pays the 2-8 KB staging per block: on the RTX 5070 its
+// IQ2_XS / IQ2_S / IQ1_M calls were 10-20% slower staged at 3-8 columns (iq_multi_parity --bench), while the grouped
+// expert kernels gain.  So mmvq stages only with STRATA_IQ_STAGE_GRID_MMVQ=1 (bitwise the same either way).
+bool g_stage_grid_mmvq = g_stage_grid && [] {
+    const char* v = std::getenv("STRATA_IQ_STAGE_GRID_MMVQ");
+    return v != nullptr && v[0] == '1';
+}();
 
 template<int TY>
 void launch_mmvq(const uint8_t* W, size_t rb, const block_q8_1* X, float* y, int n_in, int n_out, int ncols,
@@ -1417,10 +1725,24 @@ void launch_mmvq(const uint8_t* W, size_t rb, const block_q8_1* X, float* y, int
     const dim3 grid((unsigned) ((n_out + 3) / 4)), block(32, 4);
     if constexpr (!kSplit<TY>) mmvq_kernel<TY><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
     else if (g_old_kernels) mmvq_kernel<TY><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
-    else if (ncols <= 1) mmvq_multi_kernel<TY, 1><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
-    else if (ncols == 2) mmvq_multi_kernel<TY, 2><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
-    else if (ncols <= 4) mmvq_multi_kernel<TY, 4><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
-    else mmvq_multi_kernel<TY, 8><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);   // 8 at a time past 8
+    else if constexpr (kStageIqGrid<TY>) {
+        if (!g_stage_grid_mmvq) {
+            if (ncols <= 1) mmvq_multi_kernel<TY, 1, false><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+            else if (ncols == 2) mmvq_multi_kernel<TY, 2, false><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+            else if (ncols <= 4) mmvq_multi_kernel<TY, 4, false><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+            else mmvq_multi_kernel<TY, 8, false><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+            return;
+        }
+        if (ncols <= 1) mmvq_multi_kernel<TY, 1><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+        else if (ncols == 2) mmvq_multi_kernel<TY, 2><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+        else if (ncols <= 4) mmvq_multi_kernel<TY, 4><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+        else mmvq_multi_kernel<TY, 8><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+    } else {
+        if (ncols <= 1) mmvq_multi_kernel<TY, 1><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+        else if (ncols == 2) mmvq_multi_kernel<TY, 2><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+        else if (ncols <= 4) mmvq_multi_kernel<TY, 4><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);
+        else mmvq_multi_kernel<TY, 8><<<grid, block, 0, s>>>(W, rb, X, y, n_in, n_out, ncols);   // 8 at a time past 8
+    }
 }
 
 template<int TG>
@@ -1429,7 +1751,13 @@ void launch_gu(dim3 grid, cudaStream_t s, const unsigned long long* grp_ptr, con
                float* gate, float* up) {
     if constexpr (!kSplit<TG>) native_gu_kernel<TG><<<grid, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up);
     else if (g_old_kernels) native_gu_kernel<TG><<<grid, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up);
-    else native_gu_multi_kernel<TG><<<grid, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up);
+    else if constexpr (kStageIqGrid<TG>) {
+        if (!g_stage_grid) {
+            native_gu_multi_kernel<TG, false><<<grid, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up);
+            return;
+        }
+        native_gu_multi_kernel<TG><<<grid, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up);
+    } else native_gu_multi_kernel<TG><<<grid, 256, 0, s>>>(grp_ptr, grp_start, n_groups, ent_tok, X, L, gate, up);
 }
 
 template<int TD>
@@ -2212,7 +2540,8 @@ void native_expert_grouped(const NativeExpertLayout& L, const unsigned long long
     check("native_expert_grouped/swiglu");
     }
     if (g_exp_phase == 1) return;
-    const dim3 gd((unsigned) ((L.n_embd + 7) / 8), (unsigned) gy);
+    const int d_rows = (!g_old_kernels && L.n_ff == 640) ? (L.d_type == 20 ? 32 : (L.d_type == 42 ? 16 : 8)) : 8;
+    const dim3 gd((unsigned) ((L.n_embd + d_rows - 1) / d_rows), (unsigned) gy);
 #if defined(STRATA_HIP_GFX906)
     if ((em0 == 7 || em0 == 8) && (L.d_type == 20 || L.d_type == 42)) {
         const dim3 gl((unsigned) ((L.n_embd + LDS_RB - 1) / LDS_RB), (unsigned) cap_groups);

src/kernels/cuda/native_mmvq.cu

@@ -153,6 +153,22 @@ __global__ void native_quantize_q8_1_kernel(const float* __restrict__ x,
     if (i % Q8K == 0) y[i / Q8K].ds = q8_1_ds(d, sum);
 }
 
+__launch_bounds__(QUANT_THREADS, 1)
+__global__ void native_swiglu_quantize_q8_1_kernel(const float* __restrict__ gate,
+                                                   const float* __restrict__ up,
+                                                   Q81Block* __restrict__ y, int n_in) {
+    const int i = int(blockIdx.x) * QUANT_THREADS + int(threadIdx.x);
+    if (i >= n_in) return; // n_in is a multiple of 32: only whole warps return.
+    const float gi = gate[i];
+    const float xi = __fmul_rn(__fdividef(gi, __fadd_rn(1.0f, __expf(-gi))), up[i]);
+    const float amax = warp_max(fabsf(xi));
+    const float sum = warp_sum(xi);
+    const float d = amax / 127.0f;
+    const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
+    y[i / Q8K].qs[i % Q8K] = q;
+    if (i % Q8K == 0) y[i / Q8K].ds = make_half2(d, sum);
+}
+
 // Exact pinned vec_dot_q5_K_q8_1_impl_vmmq expression and integer dot order.
 __device__ __forceinline__ float q5_q8_dot_impl(
     const int* __restrict__ vl, const int* __restrict__ vh, const int* __restrict__ u,
@@ -1011,7 +1027,7 @@ struct SmallTraits {
 bool g_multi_exact = true;   // until the upstream layout is timed on an idle GPU (plan rule: default only what is measured)
 
 template<typename F, int NCOLS, int NW, int ROWS>
-__launch_bounds__(NW * WARP, 1)
+__launch_bounds__(NW * WARP, (ROWS <= 2 ? 4 : 1))
 __global__ void native_mmvq_multi_kernel(const typename F::Block* __restrict__ w,
                                          const Q81Block* __restrict__ x,
                                          float* __restrict__ y, int n_in, int n_out) {
@@ -1068,8 +1084,9 @@ void launch_multi_n(const void* weights, const void* x_q8_1, float* y, int n_in,
     }
     const dim3 threads(WARP, WARPS);
     if (n_in / F::DIV < F::BPI) {
-        const unsigned blocks = unsigned((std::size_t(n_out) + WARPS - 1) / WARPS);
-        native_mmvq_multi_kernel<F, NCOLS, WARPS, WARPS><<<blocks, threads, 0, s>>>(w, x, y, n_in, n_out);
+        constexpr int ROWS = 2;
+        const unsigned blocks = unsigned((std::size_t(n_out) + ROWS - 1) / ROWS);
+        native_mmvq_multi_kernel<F, NCOLS, WARPS, ROWS><<<blocks, threads, 0, s>>>(w, x, y, n_in, n_out);
     } else {
         native_mmvq_multi_kernel<F, NCOLS, WARPS, 1><<<unsigned(n_out), threads, 0, s>>>(w, x, y, n_in, n_out);
     }
@@ -1249,6 +1266,21 @@ void native_quantize_q8_1(const float* x, void* x_q8_1, int n_in, int ncols, voi
     launch_check();
 }
 
+void native_swiglu_quantize_q8_1(const float* gate, const float* up, void* x_q8_1,
+                                 int n_in, int ncols, void* stream) {
+    validate_shape(n_in, ncols);
+    validate_pointer(gate);
+    validate_pointer(up);
+    validate_pointer(x_q8_1);
+    validate_stream(stream);
+    const int n_total = n_in * ncols;
+    const unsigned blocks = unsigned((std::size_t(n_total) + QUANT_THREADS - 1) / QUANT_THREADS);
+    native_swiglu_quantize_q8_1_kernel<<<blocks, QUANT_THREADS, 0,
+                                         static_cast<cudaStream_t>(stream)>>>(
+        gate, up, static_cast<Q81Block*>(x_q8_1), n_total);
+    launch_check();
+}
+
 void native_q5_k_mmvq(const void* weights, const void* x_q8_1, float* y,
                       int n_in, int n_out, int ncols, void* stream) {
     validate_shape(n_in, ncols, QK);

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_enabled() { return native_bf16; }
 void ple_set_native_postops(bool enabled) { native_postops = enabled; }
 bool ple_native_postops_enabled() { return native_postops; }
 
@@ -365,4 +366,21 @@ void ple_block(const float* emb, const float* hidden, const float* hist_rows, co
     // itself, and the engine's caller does not want that at all.
 }
 
+void ple_block_projected(const float* projected_key, const float* projected_value, const float* hidden,
+                         const float* hist_rows, const PleWeights& w, PleOut& out, void* scratch, void* stream) {
+    if (!projected_key || !projected_value || !hidden || !hist_rows || !out.result || !scratch) return;
+    const int n_embd = NG_N_EMBD, hc = NG_HC, hc_dim = NG_HC_DIM;
+    float* d_scratch = (float*) scratch;
+    float* d_key = d_scratch;
+    float* d_query = d_key + hc_dim;
+    float* d_norm = d_query + hc_dim;
+    float* d_gated = d_norm + hc_dim;
+    float* d_conv = d_gated + hc_dim;
+    float* d_value = d_conv + hc_dim;
+    float* d_gate = d_value + n_embd;
+    float* norm_dst = out.normalized ? out.normalized : d_norm;
+    NativePlePostopsBuffers buffers{d_query, norm_dst, d_gate, d_gated, norm_dst, d_conv, out.result};
+    native_ple_postops(projected_key, hidden, projected_value, hist_rows, w, buffers, stream);
+}
+
 }  // namespace strata::kernels

src/kernels/cuda/qsa_select.cu

@@ -596,6 +596,125 @@ __global__ void __launch_bounds__(TK_T) block_topk_reg_kernel(const float* __res
     }
 }
 
+#if !defined(__HIPCC__)
+// ---- a query with more blocks than the register kernel holds (contexts past 4 * 1024 * TK_PER cells): the register
+// kernel's threads, per-warp histograms and warp scans, with each key read from memory again on every pass.  A
+// histogram has no order, so on the four radix passes thread t reads blocks t, t + 1024, ...: a warp reads 32
+// neighbours at a time.  The cells are emitted ascending in the order (warp, row, lane): warp w holds the `per` rows
+// of 32 consecutive blocks from block w * 32 * per.  block_topk_kernel's selection rule (radix threshold, ties to the
+// lowest index): identical ids.
+__global__ void __launch_bounds__(TK_T) block_topk_wide_kernel(const float* __restrict__ scores,
+                                                               const int32_t* __restrict__ steps, int64_t max_blocks,
+                                                               int64_t cap, int32_t* __restrict__ ids) {
+    __shared__ int hist[TK_T / 32][256];
+    __shared__ int s_gt[TK_T / 32], s_eq[TK_T / 32];
+    __shared__ int s_digit, s_above;
+    const int64_t qi = blockIdx.x;
+    const int32_t* st = steps + qi * kStepCount;
+    const int64_t n_kv = st[kStepNKv], n_bid = st[kStepNBid], width = st[kStepWidth];
+    int32_t* out = ids + qi * cap;
+    const int t = threadIdx.x, lane = t & 31, warp = t >> 5;
+    if (n_kv <= width) {
+        for (int64_t j = t; j < n_kv; j += TK_T) out[j] = (int32_t) j;
+        return;
+    }
+    const float* sc = scores + qi * max_blocks;
+    const int64_t nb = n_bid + 1;
+    const int64_t per = (nb + TK_T - 1) / TK_T;       // rows of 32 blocks per warp
+    auto weight = [&](int64_t b) -> int { return b < n_bid ? R : (int) (n_kv - n_bid * R); };
+    uint32_t prefix = 0;
+    int above = 0;
+    for (int shift = 24; shift >= 0; shift -= 8) {
+        for (int i = lane; i < 256; i += 32) hist[warp][i] = 0;
+        __syncwarp();
+        const uint32_t hi_mask = shift == 24 ? 0u : (0xffffffffu << (shift + 8));
+        for (int64_t b = t; b < nb; b += TK_T) {
+            const int w = weight(b);
+            if (w == 0) continue;
+            const uint32_t k = order_key(sc[b]);
+            if ((k & hi_mask) == (prefix & hi_mask)) atomicAdd(&hist[warp][(k >> shift) & 255], w);
+        }
+        __syncthreads();
+        if (t < 256) {                                // fold the warps' histograms into warp 0's
+            int sum = 0;
+            for (int w2 = 0; w2 < TK_T / 32; ++w2) sum += hist[w2][t];
+            hist[0][t] = sum;
+        }
+        __syncthreads();
+        if (t == 0) {
+            int cum = above, d = 255;
+            for (; d > 0; --d) {
+                if (cum + hist[0][d] >= width) break;
+                cum += hist[0][d];
+            }
+            s_digit = d;
+            s_above = cum;
+        }
+        __syncthreads();
+        prefix |= (uint32_t) s_digit << shift;
+        above = s_above;
+        __syncthreads();
+    }
+    const uint32_t thr = prefix;
+    const int eq_budget = (int) (width - above);      // cells equal to thr that fit, lowest index first
+    const int64_t w0 = (int64_t) warp * 32 * per;
+    int gt = 0, eq = 0;
+    for (int64_t b = w0 + lane; b < nb && b < w0 + 32 * per; b += 32) {
+        const int w = weight(b);
+        if (w == 0) continue;
+        const uint32_t k = order_key(sc[b]);
+        if (k > thr) gt += w;
+        else if (k == thr) eq += w;
+    }
+#pragma unroll
+    for (int o = 16; o > 0; o >>= 1) {
+        gt += __shfl_xor_sync(0xffffffffu, gt, o);
+        eq += __shfl_xor_sync(0xffffffffu, eq, o);
+    }
+    if (lane == 0) { s_gt[warp] = gt; s_eq[warp] = eq; }
+    __syncthreads();
+    if (t == 0) {                                     // per warp: the selected cells and the tied cells before it
+        int eb = 0, sb = 0;
+        for (int w2 = 0; w2 < TK_T / 32; ++w2) {
+            const int g2 = s_gt[w2], e2 = s_eq[w2];
+            const int take = eq_budget - eb < 0 ? 0 : (eq_budget - eb > e2 ? e2 : eq_budget - eb);
+            s_gt[w2] = sb;
+            s_eq[w2] = eb;
+            sb += g2 + take;
+            eb += e2;
+        }
+    }
+    __syncthreads();
+    int run_sel = s_gt[warp], run_eq = s_eq[warp];
+    for (int64_t r0 = w0; r0 < nb && r0 < w0 + 32 * per; r0 += 32) {
+        const int64_t b = r0 + lane;
+        const int w = b < nb ? weight(b) : 0;
+        const uint32_t k = w ? order_key(sc[b]) : 0u;
+        const int my_gt = (w && k > thr) ? w : 0, my_eq = (w && k == thr) ? w : 0;
+        if (__ballot_sync(0xffffffffu, (my_gt | my_eq) != 0) == 0u) continue;   // a row without a selected block
+        int pe = my_eq;                               // tied cells up to and with this lane
+#pragma unroll
+        for (int o = 1; o < 32; o <<= 1) {
+            const int y = __shfl_up_sync(0xffffffffu, pe, o);
+            if (lane >= o) pe += y;
+        }
+        const int left = eq_budget - (run_eq + pe - my_eq);
+        const int take = left < 0 ? 0 : (left > my_eq ? my_eq : left);
+        const int my_sel = my_gt + take;
+        int ps = my_sel;
+#pragma unroll
+        for (int o = 1; o < 32; o <<= 1) {
+            const int y = __shfl_up_sync(0xffffffffu, ps, o);
+            if (lane >= o) ps += y;
+        }
+        int32_t* dst = out + run_sel + ps - my_sel;
+        for (int c = 0; c < my_sel; ++c) dst[c] = (int32_t) (b * R + c);
+        run_eq += __shfl_sync(0xffffffffu, pe, 31);
+        run_sel += __shfl_sync(0xffffffffu, ps, 31);
+    }
+}
+#endif
+
 
 // Block scores with every key block read ONCE for all of a call's queries (block_scores_kernel's grid is
 // (max_blocks / 8) x nq: ~24,600 mostly-idle blocks per layer at a decode window, each key re-read per query).  A fixed
@@ -1054,8 +1173,8 @@ bool qsa_block_topk_cluster(const float* scores, const int32_t* steps, int64_t n
 
 void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64_t max_blocks, int64_t cap,
                     const QsaShapes& s, int32_t* ids, void* stream, int64_t active_blocks) {
-    // keys in registers when every query's blocks fit (contexts up to ~135K cells); the same ids. STRATA_TOPK_OLD=1:
-    // the kernel that reads them from memory on every pass
+    // keys in registers when every query's blocks fit (contexts up to ~135K cells), else (CUDA) the same threads
+    // reading them from memory; the same ids. STRATA_TOPK_OLD=1: the original kernel
     static const bool old = std::getenv("STRATA_TOPK_OLD") != nullptr;
     if (nq <= 0) return;
 #if !defined(__HIPCC__)
@@ -1089,10 +1208,12 @@ void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64
 #if defined(__HIPCC__)
     constexpr int64_t kRegMinBlocks = 7168;   // gfx1201: below ~28K cells the 1,024-thread kernel's fixed cost loses to the ref
     const bool too_small = counted && reach < kRegMinBlocks;
-#else
-    const bool too_small = false;
-#endif
     if (old || too_small || reach > fit) {
+#else
+    constexpr int64_t kWideMinBlocks = 4608;   // RTX 3060: a prompt batch below ~18K cells is faster on the ref
+    const bool too_small = reach > fit && active_blocks > 0 && active_blocks < kWideMinBlocks;
+    if (old || too_small) {
+#endif
         qsa_block_topk_ref(scores, steps, nq, max_blocks, cap, s, ids, stream);
         return;
     }
@@ -1100,6 +1221,11 @@ void qsa_block_topk(const float* scores, const int32_t* steps, int64_t nq, int64
         std::fprintf(stderr, "qsa_block_topk: unsupported geometry or cap\n");
         std::exit(1);
     }
+#if !defined(__HIPCC__)
+    if (reach > fit)   // past the register kernel's reach
+        block_topk_wide_kernel<<<(unsigned) nq, TK_T, 0, (cudaStream_t) stream>>>(scores, steps, max_blocks, cap, ids);
+    else
+#endif
     if (reach <= (int64_t) TK_T * TK_PER)
         block_topk_reg_kernel<TK_PER><<<(unsigned) nq, TK_T, 0, (cudaStream_t) stream>>>(scores, steps, max_blocks, cap, ids);
     else

src/kernels/cuda/shared_expert.cu

@@ -27,6 +27,7 @@
 #include "strata/kernels/s_gemv.hpp"
 #include "strata/kernels/native_mmvq.hpp"
 
+#include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
 #include <cmath>
@@ -157,6 +158,7 @@ __global__ void moe_combine_kernel(const float* __restrict__ parts, const float*
 }  // namespace
 
 void shared_expert_set_native_bf16(bool enabled) { native_bf16 = enabled; }
+bool shared_expert_native_bf16_enabled() { return native_bf16; }
 
 namespace {
 __global__ void scale_rows_kernel(float* __restrict__ out, const float* __restrict__ g, int n) {
@@ -164,6 +166,23 @@ __global__ void scale_rows_kernel(float* __restrict__ out, const float* __restri
     const int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i < n) out[(size_t) t * n + i] *= g[t];
 }
+
+__global__ void sigmoid_scale_rows_kernel(float* __restrict__ out, const float* __restrict__ g, int n) {
+    const int t = blockIdx.y;
+    const int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        const float gt = __fdividef(1.0f, 1.0f + __expf(-g[t]));
+        out[(size_t) t * n + i] *= gt;
+    }
+}
+
+bool fused_swiglu_q81_enabled() {
+    static const bool on = [] {
+        const char* v = std::getenv("STRATA_FUSED_SWIGLU_Q81");
+        return v == nullptr || std::atoi(v) != 0;
+    }();
+    return on;
+}
 }  // namespace
 
 void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, const NativeSharedWeights& nw,
@@ -172,28 +191,42 @@ void shared_expert_multi(int n_tok, const float* x, const uint16_t* x_bf16, cons
     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");
     cudaStream_t cs = (cudaStream_t) 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);
+    const void* x_q8_1 = nw.x_q8_1;
+    if (!x_q8_1) {
+        native_quantize_q8_1(x, nw.q8_1, (int) n_embd, n_tok, stream);
+        x_q8_1 = nw.q8_1;
+    }
+    native_mmvq(nw.gate_type, nw.gate_data, x_q8_1, gate, (int) n_embd, (int) n_ff, n_tok, stream);
+    native_mmvq(nw.up_type, nw.up_data, x_q8_1, up, (int) n_embd, (int) n_ff, n_tok, stream);
     const int n = (int) (n_ff * n_tok);
-    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);
+    if (fused_swiglu_q81_enabled()) {
+        native_swiglu_quantize_q8_1(gate, up, nw.q8_1, (int) n_ff, n_tok, stream);
+    } else {
+        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);
     static const bool batch = [] { const char* v = std::getenv("STRATA_DEC_BATCH"); return v == nullptr || std::atoi(v) != 0; }();
-    if (native_bf16 && batch && n_tok > 1) {   // one gemv for all rows (outputs identical), one sigmoid launch
+    if (native_bf16 && batch && n_tok > 1) {   // one gemv for all rows (outputs identical), fused sigmoid+scale
         bf16_gemv_fp32_mmvf_multi(x, n_embd, gate_inp_bf16, g, 1, n_embd, 1, n_tok, stream);
-        native_scalar_sigmoid_multi_kernel<<<1, n_tok, 0, cs>>>(g);
-    } else
-    for (int t = 0; t < n_tok; ++t) {
-        if (native_bf16) {
-            bf16_gemv_fp32_mmvf(x + (size_t) t * n_embd, gate_inp_bf16, g + t, n_embd, 1, stream);
-            native_scalar_sigmoid_kernel<<<1, 1, 0, cs>>>(g + t);
-        } else {
-            scalar_gate_kernel<<<1, 256, 0, cs>>>(x_bf16 + (size_t) t * n_embd, gate_inp_bf16, g + t, (int) n_embd);
+        sigmoid_scale_rows_kernel<<<dim3((unsigned) ((n_embd + THREADS - 1) / THREADS), (unsigned) n_tok), THREADS, 0, cs>>>(
+            out, g, (int) n_embd);
+    } else if (native_bf16 && n_tok == 1) {
+        bf16_gemv_fp32_mmvf(x, gate_inp_bf16, g, n_embd, 1, stream);
+        sigmoid_scale_rows_kernel<<<dim3((unsigned) ((n_embd + THREADS - 1) / THREADS), 1u), THREADS, 0, cs>>>(
+            out, g, (int) n_embd);
+    } else {
+        for (int t = 0; t < n_tok; ++t) {
+            if (native_bf16) {
+                bf16_gemv_fp32_mmvf(x + (size_t) t * n_embd, gate_inp_bf16, g + t, n_embd, 1, stream);
+                native_scalar_sigmoid_kernel<<<1, 1, 0, cs>>>(g + t);
+            } else {
+                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);
     }
-    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));
 }
@@ -273,31 +306,36 @@ void shared_expert(const uint8_t* x_q8_0, const uint8_t* x_q8k, const uint16_t*
         native_mmvq(native->up_type, native->up_data, native->q8_1, up, (int) n_embd, (int) n_ff, 1, stream);
     else
         gemv(up_form, up_codes, up_scales, up_off, x_q8_0, x_q8k, up, n_embd, n_ff);
-    if (native_projection)
-        native_swiglu_kernel<<<g_ff, THREADS, 0, (cudaStream_t) stream>>>(gate, up, gate, (int) n_ff);
-    else
-        swiglu_kernel<<<g_ff, THREADS, 0, (cudaStream_t) stream>>>(gate, up, gate, (int) n_ff);
-
-    // down: (n_ff) -> (n_embd), and THE INTERMEDIATE IS QUANTIZED TO THE DOWN WEIGHT'S OWN CONTRACT - which is
-    // what `ggml_mul_mat` does for every matmul in the model.  It used to be rounded to fp16 with no
-    // justification beyond "the kernel takes fp16".
-    if (native_down) {
-        native_quantize_q8_1(gate, native->q8_1, (int) n_ff, 1, stream);
+    if (native_projection && native_down && fused_swiglu_q81_enabled()) {
+        native_swiglu_quantize_q8_1(gate, up, native->q8_1, (int) n_ff, 1, stream);
         native_mmvq(native->down_type, native->down_data, native->q8_1, out, (int) n_ff, (int) n_embd, 1, stream);
-    } else if (down_form.act_kind == 1) {
-        if (n_ff % 256 != 0) {
-            std::fprintf(stderr, "shared_expert: the down weight wants Q8_K but n_ff %lld is not a multiple "
-                                 "of 256; Q8_K is structurally impossible here\n", (long long) n_ff);
-            std::exit(1);
-        }
-        quantize_q8_K(gate, h_q8k, n_ff, stream);
-        gemv(down_form, down_codes, down_scales, down_off, h_q8_0, h_q8k, out, n_ff, n_embd);
-    } else if (down_form.code_bits == 2) {
-        quantize_q8_0(gate, h_q8_0, n_ff, stream);
-        gemv(down_form, down_codes, down_scales, down_off, h_q8_0, h_q8k, out, n_ff, n_embd);
     } else {
-        quantize_q8_0(gate, h_q8_0, n_ff, stream);
-        gemv(down_form, down_codes, down_scales, down_off, h_q8_0, h_q8k, out, n_ff, n_embd);
+        if (native_projection)
+            native_swiglu_kernel<<<g_ff, THREADS, 0, (cudaStream_t) stream>>>(gate, up, gate, (int) n_ff);
+        else
+            swiglu_kernel<<<g_ff, THREADS, 0, (cudaStream_t) stream>>>(gate, up, gate, (int) n_ff);
+
+        // down: (n_ff) -> (n_embd), and THE INTERMEDIATE IS QUANTIZED TO THE DOWN WEIGHT'S OWN CONTRACT - which is
+        // what `ggml_mul_mat` does for every matmul in the model.  It used to be rounded to fp16 with no
+        // justification beyond "the kernel takes fp16".
+        if (native_down) {
+            native_quantize_q8_1(gate, native->q8_1, (int) n_ff, 1, stream);
+            native_mmvq(native->down_type, native->down_data, native->q8_1, out, (int) n_ff, (int) n_embd, 1, stream);
+        } else if (down_form.act_kind == 1) {
+            if (n_ff % 256 != 0) {
+                std::fprintf(stderr, "shared_expert: the down weight wants Q8_K but n_ff %lld is not a multiple "
+                                     "of 256; Q8_K is structurally impossible here\n", (long long) n_ff);
+                std::exit(1);
+            }
+            quantize_q8_K(gate, h_q8k, n_ff, stream);
+            gemv(down_form, down_codes, down_scales, down_off, h_q8_0, h_q8k, out, n_ff, n_embd);
+        } else if (down_form.code_bits == 2) {
+            quantize_q8_0(gate, h_q8_0, n_ff, stream);
+            gemv(down_form, down_codes, down_scales, down_off, h_q8_0, h_q8k, out, n_ff, n_embd);
+        } else {
+            quantize_q8_0(gate, h_q8_0, n_ff, stream);
+            gemv(down_form, down_codes, down_scales, down_off, h_q8_0, h_q8k, out, n_ff, n_embd);
+        }
     }
 
     // the per-token scalar gate, then the multiply.  Note the gate is computed from `x`, the ORIGINAL hidden
@@ -307,11 +345,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);
+        sigmoid_scale_rows_kernel<<<dim3(g_embd, 1u), THREADS, 0, (cudaStream_t) stream>>>(out, g, (int) n_embd);
     } 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/cuda/verify_kernels.cu

@@ -67,23 +67,24 @@ __global__ void gdn_conv_commit_kernel(float* __restrict__ hist, const float* __
     hist[c * 3 + 2] = seq[2];
 }
 
-__global__ void __launch_bounds__(256) gdn_ab_multi_kernel(const float* __restrict__ x, const uint16_t* __restrict__ wa,
-                                                           const uint16_t* __restrict__ wb,
-                                                           const float* __restrict__ dt,
-                                                           const float* __restrict__ ssm_a, float* __restrict__ gate,
-                                                           float* __restrict__ beta, int n, int h_v, int T) {
-    const int row = blockIdx.x * 8 + (threadIdx.x >> 5), lane = threadIdx.x & 31;
+template <int MAX_T = kVerifyMaxT>
+__global__ void __launch_bounds__(64) gdn_ab_multi_kernel(const float* __restrict__ x, const uint16_t* __restrict__ wa,
+                                                          const uint16_t* __restrict__ wb,
+                                                          const float* __restrict__ dt,
+                                                          const float* __restrict__ ssm_a, float* __restrict__ gate,
+                                                          float* __restrict__ beta, int n, int h_v, int T) {
+    const int row = blockIdx.x * 2 + (threadIdx.x >> 5), lane = threadIdx.x & 31;
     if (row >= 2 * h_v) return;
     const bool is_beta = row >= h_v;
     const int r = is_beta ? row - h_v : row;
     const uint4* w4 = reinterpret_cast<const uint4*>((is_beta ? wb : wa) + (size_t) r * n);
-    float acc[kVerifyMaxT];
+    float acc[MAX_T];
 #pragma unroll
-    for (int t = 0; t < kVerifyMaxT; ++t) acc[t] = 0.0f;
+    for (int t = 0; t < MAX_T; ++t) acc[t] = 0.0f;
     for (int j = lane; j < n / 8; j += 32) {
         const uint4 wv = __ldg(w4 + j);
 #pragma unroll
-        for (int t = 0; t < kVerifyMaxT; ++t) {
+        for (int t = 0; t < MAX_T; ++t) {
             if (t >= T) break;
             const float* xt = x + (size_t) t * n;
             const float4 xa = *reinterpret_cast<const float4*>(xt + j * 8);
@@ -97,7 +98,7 @@ __global__ void __launch_bounds__(256) gdn_ab_multi_kernel(const float* __restri
         }
     }
 #pragma unroll
-    for (int t = 0; t < kVerifyMaxT; ++t) {
+    for (int t = 0; t < MAX_T; ++t) {
         if (t >= T) break;
         float a = acc[t];
         for (int o = 16; o > 0; o >>= 1) a += __shfl_xor_sync(0xffffffffu, a, o);
@@ -112,6 +113,50 @@ __global__ void __launch_bounds__(256) gdn_ab_multi_kernel(const float* __restri
     }
 }
 
+// State-only commit kernel: each head's 128 independent state columns are split across 4 blocks of 32 columns
+// (48 * 4 = 192 blocks of 128 threads across all SMs, vs 48 blocks of 512 threads), with no sq/o/norm/z/y work.
+__global__ void __launch_bounds__(32 * RG) gdn_step_commit_kernel(float* __restrict__ state,
+                                                                  const float* __restrict__ hbuf, int C,
+                                                                  const float* __restrict__ gate,
+                                                                  const float* __restrict__ beta,
+                                                                  int h_k, int h_v,
+                                                                  const int32_t* __restrict__ n_keep) {
+    const int n = *n_keep;
+    if (n <= 0) return;
+    __shared__ float sk[S];
+    __shared__ float red[RG][32];
+    const int head = blockIdx.x;
+    const int col_local = threadIdx.x;
+    const int col = blockIdx.y * 32 + col_local;
+    const int rg = threadIdx.y;
+    const int tid = rg * 32 + col_local;    // 0 .. 127
+    const int qh = head % h_k;
+    const int qk = S * h_k;
+    float s[RPG];
+    float* base = state + ((size_t) (rg * RPG) * h_v + head) * S + col;
+    const size_t row_stride = (size_t) h_v * S;
+#pragma unroll
+    for (int r = 0; r < RPG; ++r) s[r] = base[r * row_stride];
+    for (int t = 0; t < n; ++t) {
+        const float* ht = hbuf + (size_t) t * C;
+        __syncthreads();
+        sk[tid] = ht[qk + qh * S + tid];
+        __syncthreads();
+        const float g = __expf(gate[(size_t) t * h_v + head]);
+        float kv = 0.0f;
+#pragma unroll
+        for (int r = 0; r < RPG; ++r) kv = fmaf(s[r], sk[rg * RPG + r], kv);
+        red[rg][col_local] = kv;
+        __syncthreads();
+        const float kv_col = red[0][col_local] + red[1][col_local] + red[2][col_local] + red[3][col_local];
+        const float delta = (ht[2 * qk + head * S + col] - g * kv_col) * beta[(size_t) t * h_v + head];
+#pragma unroll
+        for (int r = 0; r < RPG; ++r) s[r] = fmaf(g, s[r], sk[rg * RPG + r] * delta);
+    }
+#pragma unroll
+    for (int r = 0; r < RPG; ++r) base[r * row_stride] = s[r];
+}
+
 __global__ void __launch_bounds__(S * RG) gdn_step_norm_multi_kernel(float* __restrict__ state,
                                                                      const float* __restrict__ hbuf, int C,
                                                                      const float* __restrict__ gate,
@@ -138,8 +183,12 @@ __global__ void __launch_bounds__(S * RG) gdn_step_norm_multi_kernel(float* __re
     for (int r = 0; r < RPG; ++r) s[r] = base[r * row_stride];
     for (int t = 0; t < n; ++t) {
         const float* ht = hbuf + (size_t) t * C;
+        const bool need_out = (t >= t_out_begin);
         __syncthreads();                // the previous token is done with sk/sq/red/wsum
-        if (tid < S) { sk[tid] = ht[qk + qh * S + tid]; sq[tid] = ht[qh * S + tid]; }
+        if (tid < S) {
+            sk[tid] = ht[qk + qh * S + tid];
+            if (need_out) sq[tid] = ht[qh * S + tid];
+        }
         __syncthreads();
         const float g = __expf(gate[(size_t) t * h_v + head]);
         float kv = 0.0f;
@@ -149,6 +198,11 @@ __global__ void __launch_bounds__(S * RG) gdn_step_norm_multi_kernel(float* __re
         __syncthreads();
         const float kv_col = red[0][col] + red[1][col] + red[2][col] + red[3][col];
         const float delta = (ht[2 * qk + head * S + col] - g * kv_col) * beta[(size_t) t * h_v + head];
+        if (!need_out) {
+#pragma unroll
+            for (int r = 0; r < RPG; ++r) s[r] = fmaf(g, s[r], sk[rg * RPG + r] * delta);
+            continue;
+        }
         float o = 0.0f;
 #pragma unroll
         for (int r = 0; r < RPG; ++r) {
@@ -163,7 +217,6 @@ __global__ void __launch_bounds__(S * RG) gdn_step_norm_multi_kernel(float* __re
             oc = (red[0][col] + red[1][col] + red[2][col] + red[3][col]) * rsqrtf((float) S);
             sq_part = oc * oc;
         }
-        if (t < t_out_begin) continue;   // a replayed token: its state update is needed, its output is not
         for (int o2 = 16; o2 > 0; o2 >>= 1) sq_part += __shfl_xor_sync(0xffffffffu, sq_part, o2);
         if ((tid & 31) == 0) wsum[tid >> 5] = sq_part;
         __syncthreads();
@@ -285,11 +338,13 @@ __global__ void mtp_select_kernel(const float* __restrict__ R_src, int64_t strid
     const int row = *row_dev;
     for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < stride; i += (int64_t) gridDim.x * blockDim.x)
         R_dst[i] = R_src[(size_t) row * stride + i];
+    __syncthreads();
     if (blockIdx.x == 0 && threadIdx.x == 0) {
         const int32_t tok = ids[row];
         *tok_dst = tok;
-        if (out != nullptr) ((volatile int32_t*) out)[j] = tok;
         if (probs != nullptr && out_p != nullptr) ((volatile float*) out_p)[j] = probs[row];
+        __threadfence_system();
+        if (out != nullptr) ((volatile int32_t*) out)[j] = tok;
     }
 }
 
@@ -331,8 +386,8 @@ void ident_hits(const int32_t* ids, int n, int32_t* slot, int32_t* dst, int32_t*
 
 void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const int32_t* row_dev, float* R_dst,
                 int32_t* tok_dst, int32_t* out, int j, void* stream, const float* probs, float* out_p) {
-    mtp_select_kernel<<<16, 256, 0, (cudaStream_t) stream>>>(R_src, R_stride, ids, row_dev, R_dst, tok_dst, out, j,
-                                                             probs, out_p);
+    mtp_select_kernel<<<1, 256, 0, (cudaStream_t) stream>>>(R_src, R_stride, ids, row_dev, R_dst, tok_dst, out, j,
+                                                            probs, out_p);
     check("mtp_select");
 }
 
@@ -404,8 +459,14 @@ void gdn_ab_multi(const float* x, const uint16_t* w_alpha, const uint16_t* w_bet
         std::fprintf(stderr, "gdn_ab_multi: invalid arguments\n");
         std::exit(1);
     }
-    gdn_ab_multi_kernel<<<(unsigned) ((2 * h_v + 7) / 8), 256, 0, (cudaStream_t) stream>>>(
-        x, w_alpha, w_beta, dt, ssm_a, gate, beta, n_embd, h_v, n_tok);
+    const unsigned blocks = (unsigned) ((2 * h_v + 1) / 2);
+    if (n_tok <= 4) {
+        gdn_ab_multi_kernel<4><<<blocks, 64, 0, (cudaStream_t) stream>>>(
+            x, w_alpha, w_beta, dt, ssm_a, gate, beta, n_embd, h_v, n_tok);
+    } else {
+        gdn_ab_multi_kernel<kVerifyMaxT><<<blocks, 64, 0, (cudaStream_t) stream>>>(
+            x, w_alpha, w_beta, dt, ssm_a, gate, beta, n_embd, h_v, n_tok);
+    }
     check("gdn_ab_multi");
 }
 
@@ -417,8 +478,17 @@ void gdn_step_norm_multi(float* state, const float* h, int conv_channels, const
         std::fprintf(stderr, "gdn_step_norm_multi: invalid arguments\n");
         std::exit(1);
     }
-    gdn_step_norm_multi_kernel<<<(unsigned) h_v, dim3(S, RG), 0, (cudaStream_t) stream>>>(
-        state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, n_keep, t_out_begin);
+    static const bool commit_split = [] {
+        const char* e = std::getenv("STRATA_GDN_COMMIT_SPLIT");
+        return !e || e[0] != '0';
+    }();
+    if (commit_split && n_keep != nullptr && t_out_begin >= n_tok) {
+        gdn_step_commit_kernel<<<dim3((unsigned) h_v, 4u), dim3(32, RG), 0, (cudaStream_t) stream>>>(
+            state, h, conv_channels, gate, beta, h_k, h_v, n_keep);
+    } else {
+        gdn_step_norm_multi_kernel<<<(unsigned) h_v, dim3(S, RG), 0, (cudaStream_t) stream>>>(
+            state, h, conv_channels, gate, beta, z, gamma, eps, y, h_k, h_v, n_tok, n_keep, t_out_begin);
+    }
     check("gdn_step_norm_multi");
 }
 
@@ -430,15 +500,58 @@ __global__ void wait_flag_ge_kernel(const volatile uint32_t* flag, uint32_t valu
 }  // namespace
 
 namespace {
-__global__ void resident_plan_kernel(const int32_t* __restrict__ ids, int n, int k, const int32_t* __restrict__ res,
+// one thread per window entry: up to kVerifyMaxT tokens x 10 routed experts (80), so 128 (#646 had 64: a window of
+// 7+ tokens lost its last entries)
+constexpr int kResidentPlanMax = 128;
+static_assert(kVerifyMaxT * 10 <= kResidentPlanMax, "resident_plan: one thread per entry");
+__global__ void __launch_bounds__(kResidentPlanMax) resident_plan_kernel(const int32_t* __restrict__ ids, int n, int k, const int32_t* __restrict__ res,
                                      int n_expert, const uint8_t* cache_base, const unsigned long long* slot_off,
                                      long long blob, int32_t* __restrict__ pl, long long capx, uint32_t* skip,
                                      uint32_t ring) {
-    // one thread: at most kVerifyMaxT * 10 entries, the host's exact loop
-    for (int i = 0; i < n; ++i) {
-        const int32_t e = ids[i];
-        if (e < 0 || e >= n_expert || res[e] < 0) { *skip = 0; return; }
+    __shared__ int32_t s_ids[kResidentPlanMax];
+    __shared__ unsigned long long s_ptr[kResidentPlanMax];
+    __shared__ int32_t s_first[kResidentPlanMax];
+    __shared__ int32_t s_cnt[kResidentPlanMax];
+    __shared__ int32_t s_gstart[kResidentPlanMax];
+    __shared__ int s_bad;
+    const int tid = threadIdx.x;
+    if (tid == 0) s_bad = 0;
+    __syncthreads();
+
+    int32_t eid = -1;
+    if (tid < n) {
+        eid = ids[tid];
+        s_ids[tid] = eid;
+        const int32_t slot = (eid >= 0 && eid < n_expert) ? res[eid] : -1;
+        if (slot < 0) {
+            atomicOr(&s_bad, 1);
+        } else {
+            s_ptr[tid] = (unsigned long long) (cache_base + (slot_off ? (size_t) slot_off[slot] : (size_t) slot * (size_t) blob));
+        }
     }
+    __syncthreads();
+    if (s_bad) {
+        if (tid == 0 && skip != nullptr) *skip = 0;
+        return;
+    }
+
+    int first_j = tid;
+    int rank_in_group = 0;
+    int count_same = 0;
+    if (tid < n) {
+        for (int j = 0; j < n; ++j) {
+            if (s_ids[j] == eid) {
+                if (j < first_j) first_j = j;
+                if (j < tid) ++rank_in_group;
+                ++count_same;
+            }
+        }
+        const bool is_first = (first_j == tid);
+        s_first[tid] = is_first ? 1 : 0;
+        s_cnt[tid] = is_first ? count_same : 0;
+    }
+    __syncthreads();
+
     int32_t* counts = pl;
     int32_t* start = pl + 4;
     int32_t* dst = start + capx + 1;
@@ -446,30 +559,38 @@ __global__ void resident_plan_kernel(const int32_t* __restrict__ ids, int n, int
     const long long ptr_off = ((4 + (capx + 1) + 2 * capx) + 1) & ~1ll;
     unsigned long long* ptr = (unsigned long long*) (pl + ptr_off);
     int32_t* start2 = pl + ptr_off + 4 * capx;
-    int groups = 0, entries = 0;
-    for (int i0 = 0; i0 < n; ++i0) {
-        bool first = true;
-        for (int j = 0; j < i0; ++j) if (ids[j] == ids[i0]) { first = false; break; }
-        if (!first) continue;
-        const int32_t slot = res[ids[i0]];
-        ptr[groups] = (unsigned long long) (cache_base + (slot_off ? (size_t) slot_off[slot] : (size_t) slot * (size_t) blob));
-        start[groups] = entries;
-        for (int i = i0; i < n; ++i)
-            if (ids[i] == ids[i0]) {
-                // an entry belongs to i0's group when its first occurrence is i0: the same expert id
-                dst[entries] = i;
-                tok[entries] = i / k;
-                ++entries;
-            }
-        ++groups;
+
+    if (tid < n && first_j == tid) {
+        int grp_idx = 0;
+        int ent_start = 0;
+        for (int j = 0; j < tid; ++j) {
+            grp_idx += s_first[j];
+            ent_start += s_cnt[j];
+        }
+        ptr[grp_idx] = s_ptr[tid];
+        start[grp_idx] = ent_start;
+        s_gstart[tid] = ent_start;
+    }
+    __syncthreads();
+
+    if (tid < n) {
+        const int out_idx = s_gstart[first_j] + rank_in_group;
+        dst[out_idx] = tid;
+        tok[out_idx] = tid / k;
+    }
+    if (tid == 0) {
+        int groups = 0;
+        for (int j = 0; j < n; ++j) groups += s_first[j];
+        start[groups] = n;
+        start2[0] = n;
+        counts[0] = groups;
+        counts[1] = n;
+        counts[2] = 0;
+        if (skip != nullptr) {
+            __threadfence();
+            *skip = ring;
+        }
     }
-    start[groups] = entries;
-    start2[0] = entries;
-    counts[0] = groups;
-    counts[1] = entries;
-    counts[2] = 0;
-    __threadfence();
-    *skip = ring;
 }
 __global__ void wait_flag_ge_or_kernel(const volatile uint32_t* flag, uint32_t value, const volatile uint32_t* skip) {
     if (*skip == value) return;
@@ -492,8 +613,8 @@ __global__ void copy_or_zero_kernel(float4* __restrict__ dst, const volatile flo
 void resident_plan(const int32_t* ids, int n_entries, int k, const int32_t* res_layer, int n_expert,
                    const uint8_t* cache_base, const unsigned long long* slot_off, long long blob, int32_t* plan,
                    long long capx, uint32_t* skip, uint32_t ring, void* stream) {
-    resident_plan_kernel<<<1, 1, 0, (cudaStream_t) stream>>>(ids, n_entries, k, res_layer, n_expert, cache_base, slot_off,
-                                                             blob, plan, capx, skip, ring);
+    resident_plan_kernel<<<1, kResidentPlanMax, 0, (cudaStream_t) stream>>>(ids, n_entries, k, res_layer, n_expert, cache_base, slot_off,
+                                                              blob, plan, capx, skip, ring);
     check("resident_plan");
 }
 void wait_flag_ge_or(const uint32_t* flag, uint32_t value, const uint32_t* skip, void* stream) {

src/kernels/ngram.cpp

@@ -150,6 +150,11 @@ struct PleTable::Impl {
     bool locked = false;
     uint32_t rows[PLE_N_HEADS] = {};
     uint8_t raw[PLE_N_HEADS * PLE_ROW_BYTES_MAX] = {};
+    static constexpr size_t kMaxPrefetch = 16;
+    size_t n_prefetch = 0;
+    strata::ngram::PleReader::Ticket prefetch_tickets[kMaxPrefetch] = {};
+    uint32_t prefetch_keys[kMaxPrefetch][PLE_N_HEADS] = {};
+    uint8_t prefetch_raw[kMaxPrefetch][PLE_N_HEADS * PLE_ROW_BYTES_MAX] = {};
     bool fp8 = false;                 // F8_E4M3 rows (tools/ple_fp8_pack.py); else IQ4_NL
     float scale = 1.0f;               // the FP8 table's one scale
     uint32_t rb = PLE_ROW_BYTES;      // bytes per row
@@ -275,7 +280,40 @@ bool PleTable::open(const std::string& gguf_path, std::string& err, const PleIoO
     return true;
 }
 
+void PleTable::wait_prefetches() {
+    if (impl_->n_prefetch == 0) return;
+    std::string dummy;
+    for (size_t i = 0; i < impl_->n_prefetch; ++i) (void) impl_->reader.collect(impl_->prefetch_tickets[i], dummy);
+    impl_->n_prefetch = 0;
+}
+
+void PleTable::prefetch_rows(const uint32_t* rows16) {
+    if (impl_->mode == PleIo::Direct && impl_->reader.is_open()) {
+        for (size_t i = 0; i < impl_->n_prefetch; ++i)
+            if (std::memcmp(impl_->prefetch_keys[i], rows16, sizeof(uint32_t) * PLE_N_HEADS) == 0) return;
+        if (impl_->n_prefetch >= Impl::kMaxPrefetch) wait_prefetches();
+        const size_t slot = impl_->n_prefetch++;
+        std::memcpy(impl_->prefetch_keys[slot], rows16, sizeof(uint32_t) * PLE_N_HEADS);
+        impl_->prefetch_tickets[slot] = impl_->reader.issue(impl_->prefetch_keys[slot], PLE_N_HEADS, impl_->prefetch_raw[slot]);
+        return;
+    }
+#if defined(_WIN32)
+    if (impl_->data != nullptr && g_ple_prefetch) {
+        WIN32_MEMORY_RANGE_ENTRY ranges[PLE_N_HEADS];
+        ULONG_PTR n = 0;
+        for (int h = 0; h < PLE_N_HEADS; ++h) {
+            if (rows16[h] >= impl_->n_rows) continue;
+            ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->rb);
+            ranges[n].NumberOfBytes = impl_->rb;
+            ++n;
+        }
+        if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0);
+    }
+#endif
+}
+
 void PleTable::close() {
+    wait_prefetches();
     impl_->reader.close();
     impl_->pending = false;
     impl_->locked = false;   // the unmap below releases the lock
@@ -318,6 +356,7 @@ void PleTable::read_row(uint32_t row, float* out160) const {
 }
 
 bool PleTable::issue(const uint32_t* rows16) {
+    wait_prefetches();
     std::memcpy(impl_->rows, rows16, sizeof impl_->rows);
     if (impl_->mode == PleIo::Direct) {
         if (impl_->pending) return false;              // one token in flight per table
@@ -360,6 +399,35 @@ bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, s
     if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; }
     const size_t n = n_tokens * (size_t) PLE_N_HEADS;
     if (impl_->mode == PleIo::Direct) {
+        // Fast path: if all requested tokens match our in-flight/completed prefetch slots, collect them directly
+        // without issuing a second PleReader ticket or copying through RowCache.
+        if (impl_->n_prefetch >= n_tokens && n_tokens <= Impl::kMaxPrefetch) {
+            bool exact_match = true;
+            for (size_t t = 0; t < n_tokens; ++t) {
+                if (std::memcmp(impl_->prefetch_keys[t], rows + t * PLE_N_HEADS, sizeof(uint32_t) * PLE_N_HEADS) != 0) {
+                    exact_match = false;
+                    break;
+                }
+            }
+            if (exact_match) {
+                for (size_t i = 0; i < impl_->n_prefetch; ++i) {
+                    if (!impl_->reader.collect(impl_->prefetch_tickets[i], err)) {
+                        impl_->n_prefetch = 0;
+                        return false;
+                    }
+                }
+                impl_->n_prefetch = 0;
+                for (size_t t = 0; t < n_tokens; ++t) {
+                    for (int h = 0; h < PLE_N_HEADS; ++h) {
+                        impl_->decode(impl_->prefetch_raw[t] + (size_t) h * impl_->rb,
+                                      out + (t * PLE_N_HEADS + (size_t) h) * PLE_HEAD_DIM);
+                    }
+                }
+                impl_->bytes_read += (uint64_t) n * impl_->rb;
+                return true;
+            }
+        }
+        wait_prefetches();
         std::vector<uint8_t> raw(n * impl_->rb);
         const auto ticket = impl_->reader.issue(rows, n, raw.data());
         if (!impl_->reader.collect(ticket, err)) return false;
@@ -367,6 +435,7 @@ bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, s
         impl_->bytes_read += (uint64_t) n * impl_->rb;
         return true;
     }
+    wait_prefetches();
     for (size_t i = 0; i < n; ++i) read_row(rows[i], out + i * PLE_HEAD_DIM);
     return true;
 }

src/kernels/qsa_topk_parity.cpp

@@ -0,0 +1,139 @@
+// src/kernels/qsa_topk_parity.cpp - the QSA top-k the engine dispatches (qsa_block_topk) against the original kernel
+// (qsa_block_topk_ref): the selected cells must be identical (GPU, synthetic, no model).
+//
+// A case is a context under a capacity (the engine's --max-context; max_blocks = capacity / 4 + 2, and without an
+// active-block count the dispatch follows it), with consecutive queries at the context's end so that the tail block
+// holds 0 to 3 cells.  Three score sets per case: continuous, 64 levels (many ties at the threshold) and all equal
+// (the rule's ties go to the lowest index).  The dispatch is called as a decode window calls it (no block count) and as
+// the prompt path does (with the call's largest block count).  The reference and the dispatched kernel are also timed.
+//
+// Usage: qsa_topk_parity --selftest
+//        qsa_topk_parity CONTEXT [QUERIES=256] [REPS=10] [CAPACITY=CONTEXT] [COUNT=1]   (COUNT=0: no block count)
+#include "strata/kernels/qsa.hpp"
+#include "strata/kernels/qsa_select.hpp"
+
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+namespace k = strata::kernels;
+
+namespace {
+void ck(cudaError_t e, const char* w) {
+    if (e != cudaSuccess) { std::fprintf(stderr, "%s: %s\n", w, cudaGetErrorString(e)); std::exit(2); }
+}
+
+bool run_case(int64_t ctx, int64_t nq, int reps, int64_t capacity, bool counted) {
+    const k::QsaShapes s = k::qsa_real_shapes();
+    const int64_t max_blocks = capacity / s.idx_block + 2, cap = k::qsa_selection_width(k::kTopkMaxCells, s);
+    std::vector<int32_t> steps((size_t) (nq * k::kStepCount));
+    for (int64_t i = 0; i < nq; ++i) {
+        int32_t* st = steps.data() + i * k::kStepCount;
+        const int64_t pos = ctx - nq + i;
+        st[k::kStepPos] = (int32_t) pos;
+        st[k::kStepNKv] = (int32_t) (pos + 1);
+        st[k::kStepNBid] = (int32_t) ((pos + 1) / s.idx_block);
+        st[k::kStepWidth] = (int32_t) k::qsa_selection_width(pos + 1, s);
+    }
+    std::mt19937 rng((uint32_t) (ctx * 31 + capacity));
+    std::normal_distribution<float> nd(0.f, 1.f);
+    std::vector<float> scores((size_t) (nq * max_blocks)), tied(scores.size()), equal(scores.size(), 0.25f);
+    for (int64_t i = 0; i < nq; ++i) {
+        const int32_t* st = steps.data() + i * k::kStepCount;
+        float* row = scores.data() + i * max_blocks;
+        for (int64_t b = 0; b < max_blocks; ++b) row[b] = nd(rng);
+        // the tail block as the scorers leave it: far above the others when it holds cells
+        if (st[k::kStepNKv] % s.idx_block != 0) row[st[k::kStepNBid]] += 1e9f;
+        for (int64_t b = 0; b < max_blocks; ++b)
+            tied[(size_t) (i * max_blocks + b)] = row[b] > 1e8f ? row[b] : std::floor(row[b] * 8.0f) / 8.0f;   // ~64 levels
+    }
+    int32_t *d_steps = nullptr, *ids_ref = nullptr, *ids_new = nullptr;
+    float* d_scores = nullptr;
+    ck(cudaMalloc(&d_steps, steps.size() * 4), "malloc");
+    ck(cudaMalloc(&d_scores, scores.size() * 4), "malloc");
+    ck(cudaMalloc(&ids_ref, (size_t) (nq * cap) * 4), "malloc");
+    ck(cudaMalloc(&ids_new, (size_t) (nq * cap) * 4), "malloc");
+    ck(cudaMemcpy(d_steps, steps.data(), steps.size() * 4, cudaMemcpyHostToDevice), "upload");
+    const int64_t active = counted ? steps[(size_t) ((nq - 1) * k::kStepCount + k::kStepNBid)] + 1 : -1;
+    std::vector<int32_t> a((size_t) (nq * cap)), b(a.size());
+    auto same = [&](const std::vector<float>& sc) -> int64_t {
+        ck(cudaMemcpy(d_scores, sc.data(), sc.size() * 4, cudaMemcpyHostToDevice), "upload");
+        ck(cudaMemset(ids_ref, 0xff, a.size() * 4), "memset");
+        ck(cudaMemset(ids_new, 0xff, a.size() * 4), "memset");
+        k::qsa_block_topk_ref(d_scores, d_steps, nq, max_blocks, cap, s, ids_ref, nullptr);
+        k::qsa_block_topk(d_scores, d_steps, nq, max_blocks, cap, s, ids_new, nullptr, active);
+        ck(cudaDeviceSynchronize(), "sync");
+        ck(cudaMemcpy(a.data(), ids_ref, a.size() * 4, cudaMemcpyDeviceToHost), "down");
+        ck(cudaMemcpy(b.data(), ids_new, b.size() * 4, cudaMemcpyDeviceToHost), "down");
+        int64_t n = 0;
+        for (int64_t i = 0; i < nq; ++i) {
+            const int64_t w = steps[(size_t) (i * k::kStepCount + k::kStepWidth)];
+            n += std::equal(a.begin() + i * cap, a.begin() + i * cap + w, b.begin() + i * cap);
+        }
+        return n;
+    };
+    const int64_t same_tied = same(tied), same_equal = same(equal), same_cont = same(scores);   // d_scores: continuous
+    cudaEvent_t e0, e1;
+    ck(cudaEventCreate(&e0), "event");
+    ck(cudaEventCreate(&e1), "event");
+    auto timed = [&](auto f) {
+        ck(cudaEventRecord(e0), "record");
+        for (int r = 0; r < reps; ++r) f();
+        ck(cudaEventRecord(e1), "record");
+        ck(cudaEventSynchronize(e1), "time");
+        float ms = 0;
+        ck(cudaEventElapsedTime(&ms, e0, e1), "elapsed");
+        return ms / (float) reps;
+    };
+    const float t_ref = timed([&] { k::qsa_block_topk_ref(d_scores, d_steps, nq, max_blocks, cap, s, ids_ref, nullptr); });
+    const float t_new = timed([&] { k::qsa_block_topk(d_scores, d_steps, nq, max_blocks, cap, s, ids_new, nullptr, active); });
+    cudaEventDestroy(e0); cudaEventDestroy(e1);
+    cudaFree(d_steps); cudaFree(d_scores); cudaFree(ids_ref); cudaFree(ids_new);
+    const bool ok = same_cont == nq && same_tied == nq && same_equal == nq;
+    std::printf("%s context %lld under a capacity of %lld (%lld blocks), %lld queries %s: reference %.3f ms, "
+                "dispatched %.3f ms (%.1fx); identical ids: continuous %lld/%lld, 64 levels %lld/%lld, equal scores "
+                "%lld/%lld\n", ok ? "PASS" : "FAIL", (long long) ctx, (long long) capacity, (long long) max_blocks,
+                (long long) nq, counted ? "with their block count" : "without a block count", t_ref, t_new, t_ref / t_new,
+                (long long) same_cont, (long long) nq, (long long) same_tied,
+                (long long) nq, (long long) same_equal, (long long) nq);
+    return ok;
+}
+}  // namespace
+
+int main(int argc, char** argv) {
+    if (argc > 1 && std::strcmp(argv[1], "--selftest") == 0) {
+        // {context, capacity}: both sides of the register kernel's reach (33,792 blocks = 135,168 cells), short
+        // contexts under a long capacity (fewer blocks than the kernel's threads; a prompt batch below the dispatch's
+        // 4,608 blocks), and the long contexts themselves
+        const int64_t cases[][2] = {{9000, 9000},    {135160, 135160}, {135176, 135176}, {2100, 262144},
+                                    {9000, 262144},  {32768, 262144},  {262144, 262144}, {300001, 524288},
+                                    {524288, 524288}};
+        bool ok = true;
+        for (const auto& c : cases) {
+            ok = run_case(c[0], 8, 3, c[1], false) && ok;     // a decode window's call
+            ok = run_case(c[0], 64, 3, c[1], true) && ok;     // the prompt path's
+        }
+        std::printf("%s\n", ok ? "qsa_topk_parity: PASS" : "qsa_topk_parity: FAIL");
+        return ok ? 0 : 1;
+    }
+    if (argc < 2) {
+        std::fprintf(stderr, "usage: qsa_topk_parity --selftest | CONTEXT [QUERIES=256] [REPS=10] [CAPACITY=CONTEXT] "
+                             "[COUNT=1]\n");
+        return 2;
+    }
+    const int64_t ctx = std::atoll(argv[1]);
+    const int64_t nq = argc > 2 ? std::atoll(argv[2]) : 256;
+    const int reps = argc > 3 ? std::atoi(argv[3]) : 10;
+    const int64_t capacity = argc > 4 ? std::atoll(argv[4]) : ctx;
+    if (ctx < nq || nq <= 0 || reps <= 0 || capacity < ctx) {
+        std::fprintf(stderr, "qsa_topk_parity: CONTEXT >= QUERIES > 0, REPS > 0, CAPACITY >= CONTEXT\n");
+        return 2;
+    }
+    return run_case(ctx, nq, reps, capacity, argc <= 5 || std::atoi(argv[5]) != 0) ? 0 : 1;
+}

src/prefill/kernels.cu

@@ -5,6 +5,7 @@
 
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
+#include <algorithm>
 
 #include <cfloat>
 #include <cmath>
@@ -956,6 +957,19 @@ __global__ void copy_i32_kernel(int32_t* __restrict__ dst, const int32_t* __rest
         dst[i] = src[i];
 }
 }  // namespace
+namespace {
+__global__ void copy_f4_kernel(float4* __restrict__ dst, const float4* __restrict__ src, int64_t n4) {
+    for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n4; i += (int64_t) gridDim.x * blockDim.x)
+        dst[i] = src[i];
+}
+}  // namespace
+void copy_f32_wide(float* dst, const float* src, int64_t n, void* stream) {
+    if (n <= 0) return;
+    const int64_t n4 = n / 4, b = std::min<int64_t>((n4 + 255) / 256, 4096);
+    if (n4 > 0) copy_f4_kernel<<<(unsigned) b, 256, 0, (cudaStream_t) stream>>>((float4*) dst, (const float4*) src, n4);
+    if (n % 4) copy_i32((int32_t*) dst + n4 * 4, (const int32_t*) src + n4 * 4, n % 4, stream);
+    check("copy_f32_wide");
+}
 void copy_i32(int32_t* dst, const int32_t* src, int64_t n, void* stream) {
     if (n <= 0) return;
     const int64_t b = (n + 255) / 256;

src/prefill/prefill.cpp

@@ -92,6 +92,7 @@ constexpr int64_t C = 10240, ZV = 6144, HV = 48;
 // plan v0.3 P6: staging holds the largest blob of the pack (a native pack's blobs differ per layer)
 inline int64_t MAXBLOB() { return (int64_t) strata::kernels::cpu::expert_layout().max_blob; }
 constexpr int STAGE = 8;           // host->device expert staging ring (chunks below stream_all_min())
+constexpr int kSplitHelpRing = 48;  // layer split help: the helper stage's ring slots it streams through (at most)
 // Step 3: from this chunk size on, every non-resident expert of every layer streams in a fixed order through a
 // ring_slots()-slot ring (nearly all 512 are routed at such a chunk), so the copy engine keeps working through the
 // attention halves instead of waiting for each layer's routing.
@@ -109,7 +110,20 @@ double g_pinned_share = 1.0;
 // 8192-token chunks: 96 slots 1153 tok/s, 384 1294 (the next layer's experts arrive during its attention half) -
 // and 96 when a large share goes through host copies (IQ3_S on 64 GB, a third unpinned: 96 slots 1216, 256 1070 -
 // the host copies are the limit and the bigger ring only takes cache slots).  STRATA_PREFILL_RING overrides.
+//
+// Those are slots on the pack they were measured on, and a slot is one WHOLE expert blob - a Q2_0 blob is
+// 1,382,400 B, so 384 slots is 506 MiB.  On a pack with bigger blobs the same slot count is a different amount of
+// memory: 1,912 MiB on Q8_0 (5,222,400 B), which is more than half of an 8 GB card's expert cache and what held
+// that card at a 1,024-token chunk when its own buffers would have fitted 6,144.  Measured on the 4-way rig, one
+// env var and nothing else: chunk 1,024 -> 6,144, prefill 87 -> 402 tok/s.  So the budget is kept in BYTES and the
+// slot count is derived from the pack (`ring_budget_slots`); Q2_0 still resolves to exactly 1024 (fused) and 384,
+// so the pack all of this was tuned on does not move.
 int g_ring_override = 0;   // #340: set by a layer split (Prefill::set_ring_override); 0 = the rule below
+// #583 (0.1.39b): the auto chunk scan's byte-budget ring (Prefill::set_ring_budget), for requests whose chunk is above
+// g_ring_small_max - the chunk 0.1.39's rule would have picked.  A prompt that fits that chunk keeps 0.1.39's ring:
+// on one chunk the smaller ring only slowed it (RTX 5070, 4K prompts: Coder -15%, IQ3_S -5%, IQ3_XXS -2%).
+int g_ring_budget = 0;
+int64_t g_ring_small_max = 0;
 // #136: the fused experts (STRATA_PF_FUSED=1) launch on a batch of a layer's streamed experts at once, so the ring
 // should hold a whole layer's (~460 of 512 on Q2_0): with 384 slots a layer's last batch waits for slots its own
 // first batch frees.  Measured on the 5070, Q2_0, the 4K / 32K code-agent prompts (one run each): fused at 384 slots
@@ -137,6 +151,30 @@ inline bool fused_ring() {
 // the largest ring: 512 slots; 1024 with the Q2_0 pack's fused experts (P3's smaller buffers, measured there) - the
 // native packs' fused layers were measured at 512
 inline int ring_cap() { return fused_ring() && !strata::kernels::cpu::expert_layout().native ? RING_MAX : 512; }
+// The ring's budget in bytes: the measured slot counts above, at the blob size they were measured with.  It is
+// bytes and not slots because a slot is one whole blob and the blob is the pack's - see the note above.
+inline constexpr uint64_t Q2_0_BLOB = 1382400ull;   // the blob of the pack the ring was tuned on
+inline uint64_t ring_bytes() {
+    const uint64_t slots = fused_ring() ? 1024ull : 384ull;   // #136: a fused ring holds two layers' experts
+    return (g_pinned_share >= 0.9 ? slots : slots / 4) * Q2_0_BLOB;
+}
+// ...and what that buys on THIS pack, never past ring_cap(): the slot count `init` lays out, and what the auto
+// chunk scan treats as a full ring.  A pack whose blobs are larger than Q2_0's gets fewer slots for the same
+// bytes, which is the point - the ring competes with the expert cache for the same VRAM.
+// 0.1.39b: on by default (#583); STRATA_RING_BYTES=0 restores 0.1.39's ring, loan and auto chunk list.  It moves the
+// prompt path's loan on a native pack (fewer ring slots, the rest kept as cache slots, a larger auto chunk), so a long
+// prompt's experts are read through a different mix of resident and streamed groups and its bits differ from 0.1.39's
+// (RTX 5070, IQ3_XXS, 32K prompt: +14% to +26%; teacher-forced against the FP16 prompt path in the same band).
+inline bool ring_bytes_on() {
+    static const bool on = [] { const char* v = std::getenv("STRATA_RING_BYTES"); return v == nullptr || v[0] != '0'; }();
+    return on;
+}
+inline int ring_budget_slots() {
+    const int64_t per = MAXBLOB();
+    const int64_t n = per > 0 ? (int64_t) (ring_bytes() / (uint64_t) per) : 0;
+    const int cap = ring_cap();
+    return (int) (n <= 0 ? 0 : (n > cap ? cap : n));
+}
 inline int ring_slots(size_t T) {
     const char* v = std::getenv("STRATA_PREFILL_RING");
 #if defined(STRATA_USE_HIP)
@@ -148,8 +186,14 @@ inline int ring_slots(size_t T) {
     }();
     if (!v && g_ring_override <= 0 && wmma) return (int64_t) T >= stream_all_min() ? 96 : STAGE;
 #endif
-    const int pinned_ring = fused_ring() ? 1024 : 384;
-    const int r = v ? std::atoi(v) : g_ring_override > 0 ? g_ring_override : (g_pinned_share >= 0.9 ? pinned_ring : 96);
+    // the byte budget as this pack's slots: 1024 fused / 384 not on Q2_0, fewer on a pack with bigger blobs.  The
+    // unpinned arm stays a slot count (96): it was measured where the host copies are the limit, and there the ring
+    // is not what is competing for VRAM.
+    // 0.1.39's ring (1024 fused / 384 pinned, 96 when a large share goes through host copies), and the #583 byte
+    // budget the auto scan chose for chunks past the size 0.1.39's rule would have picked (set_ring_budget)
+    const int pinned_ring = g_pinned_share >= 0.9 ? (fused_ring() ? 1024 : 384) : 96;
+    const bool budget = ring_bytes_on() && g_ring_budget > 0 && (int64_t) T > g_ring_small_max;
+    const int r = v ? std::atoi(v) : g_ring_override > 0 ? g_ring_override : budget ? g_ring_budget : pinned_ring;
     if (v && r == STAGE) return STAGE; // Explicit opt-in to routed-only staging, including large chunks.
     const int big = r < 16 ? 16 : r > ring_cap() ? ring_cap() : r;
     return (int64_t) T >= stream_all_min() ? big : STAGE;
@@ -386,6 +430,21 @@ struct PeerPrefill {
     int32_t *src = nullptr, *bounds = nullptr, *ident = nullptr;
     uint8_t *grp_gu = nullptr, *grp_d = nullptr;
     std::unique_ptr<mmq::Context> ctx;
+    mmq::Context* run_ctx = nullptr;      // the context its products run in: ctx, or a layer split helper stage's
+    // Without P2P (a layer split's helper) the activations and the result rows go through MAPPED pinned host buffers: the device that needs them reads them with a copy kernel.  A copy-engine
+    // upload would queue behind the expert blobs that device's streams already hold (up to a ring of them, ~70 us
+    // each - the grouping tables' reason for mapped memory too), and cudaMemcpyPeerAsync would be staged by the
+    // driver.  `back_at` / `back_rows`: the rows this layer sends back.
+    bool p2p = true;
+    // The row table and the group bounds go the same way: the primary's are pinned for its own device only (or
+    // pageable), and a copy from them on the peer's stream would hold up the host until that stream is idle.
+    float* host_x = nullptr;              // T x N activations (mapped, portable)
+    float* host_rows = nullptr;           // the peer's result rows (mapped, portable)
+    int32_t* host_src = nullptr;          // the peer's rows' tokens (mapped, portable)
+    int32_t* host_bounds = nullptr;       // its group bounds (mapped, portable)
+    size_t host_x_cap = 0, host_rows_cap = 0, host_src_cap = 0;   // elements
+    static constexpr size_t kHostBounds = 2 * (NE + NE / 16 + 2) + 64;
+    int64_t back_at = 0, back_rows = 0;
     std::vector<int32_t> bounds_host;
     std::vector<void*> owned;
     int64_t layers = 0, experts = 0, rows = 0, over_cap = 0;   // stats
@@ -409,6 +468,10 @@ struct PeerPrefill {
         if (s) cudaStreamDestroy(s);
         cudaSetDevice(prev);
         if (ev_in) cudaEventDestroy(ev_in);
+        if (host_x) cudaFreeHost(host_x);
+        if (host_rows) cudaFreeHost(host_rows);
+        if (host_src) cudaFreeHost(host_src);
+        if (host_bounds) cudaFreeHost(host_bounds);
     }
 };
 
@@ -496,6 +559,9 @@ struct Prefill::Impl {
     int32_t* tok_dev = nullptr;
     std::vector<int32_t> tok_host;
     std::unique_ptr<PeerPrefill> pp;         // multi-GPU: the peer GPU's expert share (set_peer)
+    // layer split: the next stage's GPU as a stream-only peer (set_stage_helper); its buffers are that stage's own
+    // prompt buffers, bound per prompt, and it takes the place of `pp` for a run it helps
+    std::unique_ptr<PeerPrefill> help_pp;
 };
 
 namespace {
@@ -1051,6 +1117,137 @@ bool Prefill::draft_kv(core::MtpDrafter& mtp, const float* R_rows, const int32_t
                      (long long) B, ms_idle, ms_since(tl0), ms_since(t0));
     return true;
 }
+namespace {
+// Layer split: STRATA_PREFILL_HELP=1 lends the idle stage's GPU (opt-in: the helped rows are computed in another MMQ
+// grouping and round differently from the default, like STRATA_SPLIT_OWN). Unset or 0: every stage's prompt experts on
+// its own GPU, as before.
+bool split_help_env() {
+    static const bool v = [] { const char* e = std::getenv("STRATA_PREFILL_HELP"); return e != nullptr && std::atoi(e) != 0; }();
+    return v;
+}
+// The share of a stage's streamed experts the idle stage's GPU takes for a one-chunk prompt of T tokens (0: none).
+// The streamed bytes are the same at any T; the hand-over (the activations there, the rows back) grows with T and is
+// paid per layer whatever the share.  Measured on 2x RTX 3090 (UD-Q4_K_XL, PCIe 4.0 x16 each, no P2P), fixed shares
+// 0.2 / 0.3 / 0.4 against none: 1.5K tokens 749 / 820 / 908 (none 675), 3K 1335 / 1437 / 1328 (1262), 4K 1622 / 1620 /
+// 1440 (1609), 6K 1811 / 1660 / 1471 (2038), 8K 1895 / 1804 / 1663 (2230) - the best share falls with T and from ~4K
+// none pays, which this rule follows (0.41 at 1.5K, 0.32 at 3K, off from ~3.3K).
+// STRATA_PREFILL_HELP_FRAC sets the share for every T.
+double split_help_frac(int64_t T) {
+    static const double fixed = [] {
+        const char* e = std::getenv("STRATA_PREFILL_HELP_FRAC");
+        return e ? std::clamp(std::atof(e), 0.0, 1.0) : -1.0;
+    }();
+    if (fixed >= 0.0) return fixed;
+    const double f = std::min(0.45, 0.5 - (double) T / 16384.0);
+    return f >= 0.3 ? f : 0.0;
+}
+}  // namespace
+
+bool Prefill::set_stage_helper(Prefill* helper, std::string& err) {
+    Impl& m = *impl_;
+    m.help_pp.reset();
+    if (helper == nullptr || helper == this || !split_help_env()) return true;
+    const Impl& h = *helper->impl_;
+    if (m.device < 0 || h.device < 0 || h.device == m.device || !mmq_plan().any) return true;
+    auto pp = std::make_unique<PeerPrefill>();
+    pp->peer = nullptr;   // stream-only: it holds no experts, it streams a share of this stage's
+    pp->compact = true;
+    pp->p2p = false;      // the mapped host route, whether or not the cards could reach each other (measured there)
+    pp->out_pipe = true;
+    bool ok;
+    {
+        const core::OnDevice on(m.device);
+        ok = cudaEventCreateWithFlags(&pp->ev_in, cudaEventDisableTiming) == cudaSuccess;
+    }
+    int prev = 0;
+    cudaGetDevice(&prev);
+    pp->dev = h.device;
+    cudaSetDevice(pp->dev);
+    ok = ok && cudaStreamCreateWithFlags(&pp->s, cudaStreamNonBlocking) == cudaSuccess &&
+         cudaStreamCreateWithFlags(&pp->s_out, cudaStreamNonBlocking) == cudaSuccess &&
+         cudaStreamCreateWithFlags(&pp->s_cp, cudaStreamNonBlocking) == cudaSuccess &&
+         cudaEventCreateWithFlags(&pp->ev_done, cudaEventDisableTiming) == cudaSuccess;
+    for (int i = 0; ok && i < PeerPrefill::kGrpEv; ++i)
+        ok = cudaEventCreateWithFlags(&pp->ev_grp[i], cudaEventDisableTiming) == cudaSuccess;
+    for (int b = 0; ok && b < 2; ++b) ok = cudaEventCreateWithFlags(&pp->ev_dm[b], cudaEventDisableTiming) == cudaSuccess;
+    // the ring's events; its slots are the helper stage's own ring, bound per prompt (bind_stage_helper)
+    pp->pcopied.assign((size_t) kSplitHelpRing, nullptr);
+    pp->pused.assign((size_t) kSplitHelpRing, nullptr);
+    for (int i = 0; ok && i < kSplitHelpRing; ++i)
+        ok = cudaEventCreateWithFlags(&pp->pcopied[(size_t) i], cudaEventDisableTiming) == cudaSuccess &&
+             cudaEventCreateWithFlags(&pp->pused[(size_t) i], cudaEventDisableTiming) == cudaSuccess;
+    cudaSetDevice(prev);
+    if (!ok) { err = "prefill: the layer split's prompt help (streams and events)"; return false; }
+    helper_ = helper;
+    m.help_pp = std::move(pp);
+    return true;
+}
+
+// A one-chunk prompt of T tokens on a layer split: binds the stream-only peer to the helper stage's prompt buffers
+// (that stage idles until this one hands its rows over).  False when it cannot help this prompt.
+bool Prefill::bind_stage_helper(int64_t T) {
+    Impl& m = *impl_;
+    if (!m.help_pp || helper_ == nullptr || m.pp) return false;
+    const double frac = split_help_frac(T);
+    if (frac <= 0.0 || T < stream_all_min() || m.src == nullptr || m.ring <= STAGE || fused_ring()) return false;
+    const strata::kernels::cpu::ExpertLayout& lay = strata::kernels::cpu::expert_layout();
+    for (int64_t l = stage_lb_; l < stage_le_; ++l) {
+        if (!(mmq_plan().any && mmq_plan().layer[(size_t) l])) return false;   // the peer computes on MMQ only
+        if (lay.native && fused::native_supported(lay.fmt[(size_t) l].gu_type, lay.fmt[(size_t) l].d_type)) return false;
+    }
+    Impl& h = *helper_->impl_;
+    // the helper's buffers must be laid out for this chunk (the serve path lends every stage's at once)
+    if (h.T < T || h.mixed == nullptr || h.Xq == nullptr || h.Hq == nullptr || h.GU == nullptr || h.H == nullptr ||
+        h.Dm == nullptr || h.src_dev == nullptr || h.ids_identity == nullptr || h.bounds_dev == nullptr ||
+        h.grp_gu == nullptr || h.grp_d == nullptr || !h.mmq_ctx || h.ring <= STAGE)
+        return false;
+    PeerPrefill& P = *m.help_pp;
+    P.T_max = T;
+    P.cap_rows = T * K;
+    P.G = T;                                  // one expert never has more rows than the chunk has tokens
+    P.mixed = h.mixed;
+    P.Xq_g = h.Xq;                            // the helper's per-layer buffers hold T * K rows: a group's T fit
+    P.Hq_g = h.Hq;
+    P.GU_g = h.GU;
+    P.H_g = h.H;
+    P.Dm_b[0] = h.Dm;
+    P.Dm_b[1] = h.Dm + (size_t) T * N;
+    P.src = h.src_dev;
+    P.ident = h.ids_identity;
+    P.bounds = h.bounds_dev;
+    P.grp_gu = h.grp_gu;
+    P.grp_d = h.grp_d;
+    P.run_ctx = h.mmq_ctx.get();
+    P.RP = std::min(h.ring, kSplitHelpRing);
+    P.pstage.assign(h.stage_dev, h.stage_dev + P.RP);
+    P.plive.assign((size_t) P.RP, 0);
+    P.dm_live[0] = P.dm_live[1] = false;
+    P.out_pending = false;
+    P.ps_frac = frac;
+    if (!P.p2p) {   // grown once to the largest chunk it helped
+        auto grow = [](auto*& p, size_t& cap, size_t n) {
+            if (cap >= n) return true;
+            if (p) cudaFreeHost(p);
+            p = nullptr;
+            cap = 0;
+            if (cudaHostAlloc((void**) &p, n * sizeof(*p), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) {
+                cudaGetLastError();
+                p = nullptr;
+                return false;
+            }
+            cap = n;
+            return true;
+        };
+        size_t bounds_cap = P.host_bounds ? PeerPrefill::kHostBounds : 0;
+        if (!grow(P.host_x, P.host_x_cap, (size_t) T * N) || !grow(P.host_rows, P.host_rows_cap, (size_t) T * K * N) ||
+            !grow(P.host_src, P.host_src_cap, (size_t) T * K) || !grow(P.host_bounds, bounds_cap, PeerPrefill::kHostBounds))
+            return false;
+    }
+    const core::OnDevice on(P.dev);
+    mmq::iota(P.ident, T * K, P.s);           // the helper's row table: a refill may have overwritten its slots
+    return true;
+}
+
 bool Prefill::set_peer(core::PeerExperts* peer, int64_t cap_rows, std::string& err) {
     Impl& m = *impl_;
     if (peer == nullptr || !peer->valid()) { m.pp.reset(); return true; }
@@ -1133,6 +1330,7 @@ bool Prefill::set_peer(core::PeerExperts* peer, int64_t cap_rows, std::string& e
     pp->grp_d = (uint8_t*) take(MMQ_GROUP * mp.d_max + MMQ_TAIL);
     if (ok) {
         pp->ctx = std::make_unique<mmq::Context>();
+        pp->run_ctx = pp->ctx.get();
         mmq::iota(pp->ident, pp->cap_rows, pp->s);
         ok = cudaStreamSynchronize(pp->s) == cudaSuccess;
     }
@@ -1151,6 +1349,10 @@ bool Prefill::set_peer(core::PeerExperts* peer, int64_t cap_rows, std::string& e
 
 void Prefill::set_pinned_share(double share) { g_pinned_share = share; }
 void Prefill::set_ring_override(int slots) { g_ring_override = slots > 0 ? slots : 0; }
+void Prefill::set_ring_budget(int slots, int64_t small_max) {
+    g_ring_budget = slots > 0 ? slots : 0;
+    g_ring_small_max = small_max > 0 ? small_max : 0;
+}
 double Prefill::pinned_share() { return g_pinned_share; }
 
 uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk) {
@@ -1162,7 +1364,19 @@ uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::Session
     o.take<uint16_t>((size_t) GEMM_SCRATCH, ok);
     o.take<uint8_t>(GEMM_WS, ok);
     auto f = [&](size_t n) { o.take<float>(n, ok); };
-    f(T * N); f(T * D); f(T * D); o.take<uint16_t>(T * D, ok); f(T * LR); o.take<uint16_t>(T * LR, ok);
+    // `carve`'s order, buffer for buffer: emb, R, xn, grs, xn16, lo, lo16, gated, inj, mixed, mixed_bf, mixed_h,
+    // bo.  This counted `xn` unconditionally (carve takes it only under STRATA_GR_UNFUSED) and never counted
+    // `grs`.  Net over-count T*(D-HC)*4 bytes: 42 MB at a 1024-token chunk, 252 MB (48 Q8_0 slots) at 6144 - the
+    // prompt path was told it had less room than it did.  Safe - the direction is over-estimating, and `take`
+    // still bounds-checks - but it under-sizes every loan, so every chunk the scan picks is one step smaller.
+    if (ring_bytes_on()) {
+        f(T * N); f(T * D);
+        if (gr_unfused()) f(T * D);
+        f(T * HC);
+    } else {   // STRATA_RING_BYTES=0: 0.1.39's count
+        f(T * N); f(T * D); f(T * D);
+    }
+    o.take<uint16_t>(T * D, ok); f(T * LR); o.take<uint16_t>(T * LR, ok);
     f(T * D); f(T * HC); f(T * N); o.take<uint16_t>(T * N, ok); o.take<uint16_t>(T * N, ok); f(T * N);
     if (bf16x2_hc()) { o.take<uint16_t>(T * D, ok); o.take<uint16_t>(T * LR, ok); }
     if (bf16x2()) o.take<uint16_t>(T * N, ok);
@@ -1190,6 +1404,33 @@ uint64_t Prefill::bytes_needed(const core::ModelGeometry& g, const core::Session
     return o.used + (8u << 20);   // alignment slack
 }
 
+uint64_t Prefill::bytes_needed_no_ring(const core::ModelGeometry& g, const core::SessionState& ss, int64_t chunk) {
+    return bytes_needed(g, ss, chunk) - (uint64_t) ring_slots((size_t) chunk) * (uint64_t) MAXBLOB();
+}
+
+int64_t Prefill::ring_default_slots() {
+    const int r = g_pinned_share >= 0.9 ? (fused_ring() ? 1024 : 384) : 96;
+    return (int64_t) std::min(r, ring_cap());
+}
+
+int64_t Prefill::ring_cap_for(int64_t old_chunk) {
+    // 0.1.39b (#583, measured on the RTX 5070): giving up ring slots for a bigger chunk paid where 0.1.39's ring held
+    // the chunk at 4096 or less (IQ3_XXS 32K prompts +14% to +26%: 4096/384 -> 6656/56) and lost where 0.1.39 already
+    // read 6144-token chunks (the Coder: 6144/384 -> 7936/199, 32K -12%).  From kKeepRingChunk on the ring keeps its
+    // 0.1.39 size and the scan only looks for a bigger chunk next to it (IQ3_XXS unpinned 6144/96 -> 6912/96: +9%).
+    constexpr int64_t kKeepRingChunk = 6144;
+    return old_chunk >= kKeepRingChunk ? ring_default_slots() : ring_max_slots();
+}
+
+// the unpinned arm keeps its measured 96 (the PR's rule; the byte budget would have been ~49 slots on IQ3_S)
+int64_t Prefill::ring_max_slots() {
+    return g_pinned_share >= 0.9 ? (int64_t) ring_budget_slots() : (int64_t) std::min(96, ring_cap());
+}
+
+int64_t Prefill::ring_slots_for(int64_t chunk) { return ring_slots((size_t) chunk); }
+
+bool Prefill::ring_bytes_enabled() { return ring_bytes_on(); }
+
 namespace {
 
 const core::WeightRef* need(const core::LayerView& v, const char* suffix, std::string& err) {
@@ -1301,7 +1542,7 @@ struct PeTimer {
 };
 }  // namespace
 
-bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) {
+bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) {
     err.clear();
     Impl& m = *impl_;
     const core::OnDevice on_device(m.device);
@@ -1309,11 +1550,8 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
     core::SessionState& ss = *m.ss;
     const auto t_start = Clock::now();
     const int64_t LB = stage_lb_, LE = stage_le_;
-    // the next stage reads chunk c on a thread while this one reads chunk c + 1 (declared first: an early return
-    // waits for it before anything it reads goes away)
-    std::string next_err;
-    std::future<bool> next_run;
-    int hand_buf = 0;
+    // The direct successor's future lives on the Prefill object. Intermediate
+    // stages therefore do not drain the complete remaining GPU chain here.
     double host_sync_ms = 0, host_chunk_ms = 0, host_setup_ms = 0;   // STRATA_PREFILL_TIMING: the host's share
     strata::kernels::QsaShapes s = strata::kernels::qsa_real_shapes();
     s.n_head = g.n_head; s.n_head_kv = g.n_head_kv; s.head_dim = g.head_dim; s.idx_n_head = g.idx_q_heads;
@@ -1321,6 +1559,24 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
     const uint64_t gdn_floats = (uint64_t) g.ssm_state_size * g.ssm_v_heads * g.ssm_state_size +
                                 (uint64_t) g.ssm_conv_channels * (g.ssm_d_conv - 1);
     int32_t prev[2] = {ss.ple_prev[0], ss.ple_prev[1]};
+    // layer split: a prompt of one chunk runs the stages one after the other, so the next stage's GPU idles while
+    // this one reads; it streams and computes a share of this stage's experts then (set_stage_helper), in the
+    // place of a --peer-device peer, and gives the stage back its own buffers when the run ends
+    const bool single_chunk = hand_in_ == nullptr ? n <= m.T : single_chunk_;
+    const bool helped = single_chunk && bind_stage_helper(n);
+    if (helped) std::swap(m.pp, m.help_pp);
+    struct HelpScope {
+        Impl& m;
+        bool on;
+        ~HelpScope() {
+            if (!on) return;
+            const core::OnDevice d(m.pp->dev);   // its work is done before the stage hands over; settle its streams
+            cudaStreamSynchronize(m.pp->s_cp);
+            cudaStreamSynchronize(m.pp->s_out);
+            cudaStreamSynchronize(m.pp->s);
+            std::swap(m.pp, m.help_pp);
+        }
+    } help_scope{m, helped};
     PfTimer pt;
     PeTimer pe;
     if (m.pp) pe.dev = m.pp->dev; else pe.on = false;
@@ -1515,7 +1771,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                 if (ps_on) m.pp->pseq_start[(size_t) l] = m.pp->pseq.size();
                 for (int32_t e = 0; e < m.g->n_expert; ++e) {
                     if (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) continue;
-                    if (m.pp && m.pp->peer->has(l, e)) continue;   // multi-GPU: computed on (or read from) the peer
+                    if (m.pp && m.pp->peer && m.pp->peer->has(l, e)) continue;   // multi-GPU: computed on (or read from) the peer
                     int job = -1;
                     const uint8_t* b = nullptr;
                     if (m.src->transient(l, e)) {   // CS-T: copied by the source into the stager's buffer
@@ -2112,7 +2368,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                     continue;
                                 }
                                 if (c == 0 || (m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0) ||
-                                    !m.pp->peer->has(l, e))
+                                    !m.pp->peer || !m.pp->peer->has(l, e))
                                     continue;
                                 if (rows_peer + c > m.pp->cap_rows) { ++m.pp->over_cap; continue; }
                                 on_peer[(size_t) e] = 1;
@@ -2183,6 +2439,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                         peer_now = use_mmq && !order_peer.empty();
                         if (peer_now) {
                             PeerPrefill& P = *m.pp;
+                            if (!P.p2p) cudaMemcpyAsync(P.host_x, m.mixed, (size_t) T * N * 4, cudaMemcpyDeviceToHost, m.cs);
                             cudaEventRecord(P.ev_in, m.cs);
                             int prevd = 0;
                             cudaGetDevice(&prevd);
@@ -2191,9 +2448,17 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                             cudaStreamWaitEvent(ps, P.ev_in, 0);
                             if (P.out_pending) { cudaStreamWaitEvent(ps, P.ev_done, 0); P.out_pending = false; }
                             pe.mark(kPeMoeIn, ps);
-                            cudaMemcpyPeerAsync(P.mixed, P.dev, m.mixed, prevd, (size_t) T * N * 4, ps);
-                            cudaMemcpyAsync(P.src, src_h + rows_local,   // multi-GPU: the mapped table when on
-                                             (size_t) rows_peer * 4, cudaMemcpyHostToDevice, ps);
+                            if (P.p2p) cudaMemcpyPeerAsync(P.mixed, P.dev, m.mixed, prevd, (size_t) T * N * 4, ps);
+                            else copy_f32_wide(P.mixed, P.host_x, T * N, ps);
+                            P.back_at = rows_local;
+                            P.back_rows = rows_peer;
+                            if (P.p2p) {
+                                cudaMemcpyAsync(P.src, src_h + rows_local,   // multi-GPU: the mapped table when on
+                                                (size_t) rows_peer * 4, cudaMemcpyHostToDevice, ps);
+                            } else {
+                                std::memcpy(P.host_src, src_h + rows_local, (size_t) rows_peer * 4);
+                                copy_i32(P.src, P.host_src, rows_peer, ps);
+                            }
                             const size_t n = order_peer.size();
                             if (P.compact) {
                                 // groups of up to MMQ_GROUP experts and at most G rows, each computed from its own rows
@@ -2217,7 +2482,12 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                         P.bounds_host[n + 1 + g2 * (MMQ_GROUP + 1) + i] =
                                             P.bounds_host[std::min(P.groups[g2].second, P.groups[g2].first + i)] -
                                             P.bounds_host[P.groups[g2].first];
-                                cudaMemcpyAsync(P.bounds, P.bounds_host.data(), P.bounds_host.size() * 4, cudaMemcpyHostToDevice, ps);
+                                if (P.p2p || P.bounds_host.size() > PeerPrefill::kHostBounds) {
+                                    cudaMemcpyAsync(P.bounds, P.bounds_host.data(), P.bounds_host.size() * 4, cudaMemcpyHostToDevice, ps);
+                                } else {
+                                    std::memcpy(P.host_bounds, P.bounds_host.data(), P.bounds_host.size() * 4);
+                                    copy_i32(P.bounds, P.host_bounds, (int64_t) P.bounds_host.size(), ps);
+                                }
                                 pe.mark(kPeMoeGemm, ps);
                                 const auto& f = lay.fmt[(size_t) l];
                                 for (size_t g2 = 0; g2 < ng; ++g2) {
@@ -2265,7 +2535,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                     gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub;
                                     gu.n = ngx; gu.xq = P.Xq_g; gu.bounds = rel; gu.ids = P.ident;
                                     gu.total_rows = nr; gu.max_rows = maxr; gu.dst = P.GU_g; gu.ld_dst = 1280;
-                                    P.ctx->run(gu, ps);
+                                    P.run_ctx->run(gu, ps);
                                     mmq::swiglu(P.GU_g, P.H_g, nr, 640, !lay.native, ps);
                                     mmq::quantize(P.H_g, nullptr, P.Hq_g, mmq_dt, 640, 640, nr, ps);
                                     const int b = (int) (g2 & 1);
@@ -2275,11 +2545,15 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                     dn.n = ngx; dn.xq = P.Hq_g; dn.bounds = rel;
                                     dn.ids = P.ident; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = P.Dm_b[b];
                                     dn.ld_dst = N;
-                                    P.ctx->run(dn, ps);
+                                    P.run_ctx->run(dn, ps);
                                     cudaEventRecord(P.ev_grp[g2 % PeerPrefill::kGrpEv], ps);
                                     cudaStreamWaitEvent(P.s_out, P.ev_grp[g2 % PeerPrefill::kGrpEv], 0);
-                                    cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * N, prevd, P.Dm_b[b], P.dev,
-                                                        (size_t) nr * N * 4, P.s_out);
+                                    if (P.p2p)
+                                        cudaMemcpyPeerAsync(m.Dm + (size_t) (rows_local + r0) * N, prevd, P.Dm_b[b], P.dev,
+                                                            (size_t) nr * N * 4, P.s_out);
+                                    else   // the primary reads them when it combines (below)
+                                        cudaMemcpyAsync(P.host_rows + (size_t) r0 * N, P.Dm_b[b], (size_t) nr * N * 4,
+                                                        cudaMemcpyDeviceToHost, P.s_out);
                                     cudaEventRecord(P.ev_dm[b], P.s_out);
                                     P.dm_live[b] = true;
                                 }
@@ -2292,7 +2566,12 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                     for (size_t i = 0; i <= MMQ_GROUP; ++i)
                                         P.bounds_host[n + 1 + g2 * (MMQ_GROUP + 1) + i] =
                                             P.bounds_host[std::min(n, g2 * MMQ_GROUP + i)] - P.bounds_host[g2 * MMQ_GROUP];
-                                cudaMemcpyAsync(P.bounds, P.bounds_host.data(), P.bounds_host.size() * 4, cudaMemcpyHostToDevice, ps);
+                                if (P.p2p || P.bounds_host.size() > PeerPrefill::kHostBounds) {
+                                    cudaMemcpyAsync(P.bounds, P.bounds_host.data(), P.bounds_host.size() * 4, cudaMemcpyHostToDevice, ps);
+                                } else {
+                                    std::memcpy(P.host_bounds, P.bounds_host.data(), P.bounds_host.size() * 4);
+                                    copy_i32(P.bounds, P.host_bounds, (int64_t) P.bounds_host.size(), ps);
+                                }
                                 mmq::quantize(P.mixed, P.src, P.Xq, mmq_gt, N, N, rows_peer, ps);
                                 pe.mark(kPeMoeGemm, ps);
                                 const auto& f = lay.fmt[(size_t) l];
@@ -2317,7 +2596,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                     gu.w = P.grp_gu; gu.type = mmq_gt; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = mmq_gub;
                                     gu.n = ngx; gu.xq = P.Xq; gu.bounds = P.bounds + j0; gu.ids = P.ident;
                                     gu.total_rows = rows_peer; gu.max_rows = maxr; gu.dst = P.GU; gu.ld_dst = 1280;
-                                    P.ctx->run(gu, ps);
+                                    P.run_ctx->run(gu, ps);
                                     mmq::swiglu(P.GU + r0 * 1280, P.H + r0 * 640, nr, 640, !lay.native, ps);
                                     mmq::quantize(P.H + r0 * 640, nullptr, P.Hq, mmq_dt, 640, 640, nr, ps);
                                     mmq::Product dn;
@@ -2325,7 +2604,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                     dn.n = ngx; dn.xq = P.Hq; dn.bounds = P.bounds + n + 1 + g2 * (MMQ_GROUP + 1);
                                     dn.ids = P.ident; dn.total_rows = nr; dn.max_rows = maxr; dn.dst = P.Dm + r0 * N;
                                     dn.ld_dst = N;
-                                    P.ctx->run(dn, ps);
+                                    P.run_ctx->run(dn, ps);
                                     if (P.out_pipe && nr > 0) {   // this group's rows go back while the next group computes
                                         cudaEventRecord(P.ev_grp[g2], ps);
                                         cudaStreamWaitEvent(P.s_out, P.ev_grp[g2], 0);
@@ -2565,7 +2844,7 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                                 } else {
                                     const bool r0 = m.host_res && m.cache && m.host_res[(size_t) l * m.g->n_expert + e] >= 0;
                                     const uint8_t* bp = r0 ? m.cache->device_slot(m.host_res[(size_t) l * m.g->n_expert + e])
-                                                           : (m.pp ? m.pp->peer->slot_ptr(l, e) : nullptr);   // over the cap: P2P
+                                                           : (m.pp && m.pp->peer ? m.pp->peer->slot_ptr(l, e) : nullptr);   // over the cap: P2P
                                     if (bp == nullptr) { err = "prefill: an expert is neither resident, streamed nor on the peer"; return false; }
                                     ++stats_.experts_resident;
                                     if (!compute(j, bp, -1)) return false;
@@ -2576,7 +2855,11 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
                         }
                     }
                     pt.mark(kPfCombine, cs);
-                    if (peer_now) cudaStreamWaitEvent(m.cs, m.pp->ev_done, 0);   // multi-GPU: the peer's rows are in Dm
+                    if (peer_now) {   // multi-GPU: the peer's rows are in Dm (or, without P2P, in host memory)
+                        cudaStreamWaitEvent(m.cs, m.pp->ev_done, 0);
+                        if (!m.pp->p2p)
+                            copy_f32_wide(m.Dm + (size_t) m.pp->back_at * N, m.pp->host_rows, m.pp->back_rows * N, m.cs);
+                    }
                     moe_combine(m.Dm, m.slot_dev, m.w, m.shared, m.sg, m.bo, T, m.cs);
                     // debug: STRATA_DBG_NAN=1 reports the first layer of a chunk whose MoE produced non-finite values
                     if (static const bool dbg = std::getenv("STRATA_DBG_NAN") != nullptr; dbg) {
@@ -2636,21 +2919,24 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
         core::progress_at("reading the prompt (batched): finishing the chunk from token", p0);
         pt.mark(kPfStart, cs);
         if (next_ != nullptr) {
-            // the rows to the host buffer the next stage read two chunks ago (it has finished: waited below)
-            float* h = m.hand[hand_buf];
+            // The current hand-off slot was used two chunks ago.
+            float* h = m.hand[hand_buf_];
             if (cudaMemcpyAsync(h, m.R, (size_t) T * D * 4, cudaMemcpyDeviceToHost, m.cs) != cudaSuccess ||
                 cudaStreamSynchronize(m.cs) != cudaSuccess) {
                 err = std::string("prefill: the layer split's hand-off: ") + cudaGetErrorString(cudaGetLastError());
                 return false;
             }
-            // this stage's state is at the chunk's end now (synced) and moves on with the next chunk below
+            // Wait only for the DIRECT successor's previous chunk. That successor
+            // may already have forwarded its older chunk to later GPUs.
             if (on_stage_chunk && !on_stage_chunk(p0 + T, err)) return false;
-            if (next_run.valid() && !next_run.get()) { err = next_err; return false; }
+            if (next_run_.valid() && !next_run_.get()) { err = next_err_; return false; }
+            next_err_.clear();
             next_->hand_in_ = h;
-            next_run = std::async(std::launch::async, [this, tokens, c0, T, p0, &next_err] {
-                return next_->run(tokens + c0, T, p0, next_err);
+            next_->single_chunk_ = single_chunk;
+            next_run_ = std::async(std::launch::async, [this, tokens, c0, T, p0] {
+                return next_->run_impl(tokens + c0, T, p0, next_err_);
             });
-            hand_buf ^= 1;
+            hand_buf_ ^= 1;
             continue;   // the last stage reports the chunk (on_chunk)
         }
         if (const char* dump = std::getenv("STRATA_PREFILL_DUMP_R")) {   // debug: the final residuals, every 64th
@@ -2679,7 +2965,8 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
             host_chunk_ms += ms_since(toc2);
         }
     }
-    if (next_run.valid() && !next_run.get()) { err = next_err; return false; }
+    // Do not drain the successor here. This is the overlap: an intermediate
+    // stage can return while later GPUs are still processing the previous chunk.
     ss.ple_prev[0] = prev[0];
     ss.ple_prev[1] = prev[1];
     if (std::getenv("STRATA_DBG_NAN") != nullptr) {   // debug: the state the prompt leaves for the token path
@@ -2759,4 +3046,39 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
     return true;
 }
 
+bool Prefill::drain_pipeline(std::string& err) {
+    bool ok = true;
+
+    if (next_run_.valid()) {
+        if (!next_run_.get()) {
+            err = next_err_;
+            ok = false;
+        }
+    }
+
+    if (next_ != nullptr) {
+        std::string tail_err;
+        if (!next_->drain_pipeline(tail_err)) {
+            if (ok) err = tail_err;
+            ok = false;
+        }
+    }
+
+    return ok;
+}
+
+bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& err) {
+    const bool body_ok = run_impl(tokens, n, pos0, err);
+
+    std::string drain_err;
+    const bool drain_ok = drain_pipeline(drain_err);
+
+    if (!body_ok) return false;
+    if (!drain_ok) {
+        err = drain_err;
+        return false;
+    }
+    return true;
+}
+
 }  // namespace strata::prefill

src/program/generate.cpp

@@ -15,6 +15,7 @@
 // (`phase-2-correct-engine.md:5-9`).  What it is FOR is the honest tok/s figure and the logit dump.
 
 #include "strata/core/device.hpp"
+#include "strata/core/remote_expert_opt.hpp"
 #include "strata/core/expert_cache.hpp"
 #include "strata/core/conversation_snapshot.hpp"
 #include "strata/core/conversation_memory.hpp"
@@ -369,6 +370,8 @@ struct Options {
     /// measurement found nothing - see `expert_cache_per_layer`.
     int expert_cache = 0;
     std::array<int, 3> expert_cache_remote{}; ///< CUDA1..3 slots; CUDA0 keeps dense/state/MTP
+    std::array<bool, 3> expert_cache_remote_auto{};
+    bool remote_expert_opt = false; ///< optional optimization of the existing remote-expert decode path
     std::string expert_cache_remote_placement = "stripe"; ///< stripe experts or assign complete layers to CUDA1..3
     bool expert_cache_cpu_order = false;
     /// **R4.2g.  ROUND 328 MEASURED THAT THE GLOBAL ADMISSION POLICY CANNOT WORK, AND THIS IS THE FIX.**
@@ -656,9 +659,10 @@ void usage() {
                  "                       GPU via `moe_hit_grouped_s2`.  DEFAULT 0.  Measured at 4096 slots\n"
                  "                       with --expert-cache-per-layer: 54.4%% hits, CPU pool drain 19.1 -> 10.3\n"
                  "                       ms/token, -2.7 ms/token end to end.\n"
-                 "  --expert-cache-device1 N  pre-fill N experts on CUDA1 (experimental)\n"
-                 "  --expert-cache-device2 N  pre-fill N more experts on CUDA2\n"
-                 "  --expert-cache-device3 N  pre-fill N more experts on CUDA3\n"
+                 "  --expert-cache-device1 N|auto  fill CUDA1 with N experts or as many as fit\n"
+                 "  --expert-cache-device2 N|auto  fill CUDA2 with more experts\n"
+                 "  --expert-cache-device3 N|auto  fill CUDA3 with more experts\n"
+                 "  --remote-expert-opt   enable remote-expert decode optimizations (serve mode)\n"
                  "  --expert-cache-remote-placement stripe|layer  distribute expert ranks or whole\n"
                  "                       layers across CUDA1..3 (default: stripe)\n"
                  "  --peer-device N      a second GPU as an adaptive expert-cache tier (rows over P2P; it also\n"
@@ -1348,9 +1352,13 @@ int main(int argc, char** argv) {
             const std::string v = next("--expert-cache");
             o.expert_cache = (v == "auto") ? -1 : std::atoi(v.c_str());
         }
-        else if (a == "--expert-cache-device1") o.expert_cache_remote[0] = std::atoi(next("--expert-cache-device1"));
-        else if (a == "--expert-cache-device2") o.expert_cache_remote[1] = std::atoi(next("--expert-cache-device2"));
-        else if (a == "--expert-cache-device3") o.expert_cache_remote[2] = std::atoi(next("--expert-cache-device3"));
+        else if (a == "--expert-cache-device1" || a == "--expert-cache-device2" || a == "--expert-cache-device3") {
+            const size_t r = (size_t) (a.back() - '1');
+            const std::string v = next(a.c_str());
+            o.expert_cache_remote_auto[r] = v == "auto";
+            o.expert_cache_remote[r] = v == "auto" ? std::numeric_limits<int>::max() : std::atoi(v.c_str());
+        }
+        else if (a == "--remote-expert-opt") o.remote_expert_opt = true;
         else if (a == "--expert-cache-remote-placement")
             o.expert_cache_remote_placement = next("--expert-cache-remote-placement");
         else if (a == "--vram-reserve-mib") { o.vram_reserve_mib = std::atoi(next("--vram-reserve-mib")); o.vram_reserve_given = true; }
@@ -1873,6 +1881,10 @@ int main(int argc, char** argv) {
         std::fprintf(stderr, "strata generate: --expert-cache-remote-placement must be stripe or layer\n");
         return 2;
     }
+    if (o.remote_expert_opt && !o.serve) {
+        std::fprintf(stderr, "strata generate: --remote-expert-opt requires --serve\n");
+        return 2;
+    }
     // the peer tier is the second card's only user: a layer split or a remote expert cache would put a second engine
     // part (and a second copy of the same experts) on it
     if (o.peer_device >= 1 && (!o.layer_split.empty() || o.expert_cache_remote[0] > 0)) {
@@ -2791,7 +2803,7 @@ int main(int argc, char** argv) {
             for (int i = 0; i + 1 < ns; ++i) {
                 acc += 1.0 / layer_ms[(size_t) i];
                 at[(size_t) i] = std::clamp<int64_t>((int64_t) std::llround(acc / total * (double) L),
-                                                    i == 0 ? 2 : at[(size_t) i - 1] + 1, L - (ns - 1 - i));
+                                                     i == 0 ? 2 : at[(size_t) i - 1] + 1, L - (ns - 1 - i));
             }
             consider();
         }
@@ -2905,6 +2917,7 @@ int main(int argc, char** argv) {
         // with a layer split across GPUs the drafter reads the last stage's residual: it lives on that device
         const strata::core::OnDevice on_mtp(last_st ? last_st->dev : -1);
         if (!o.mtp.empty() && !mtp.load(o.mtp, draft_geometry, last_st ? last_st->ss : ss, o.spec, err, o.mtp_window)) { std::fprintf(stderr, "strata generate: %s%s\n", err.c_str(), vram_free_note().c_str()); return 1; }
+        mtp.set_ple_session(&ss);
     }
     // THE HEAD BEFORE THE CACHE, AND BEFORE THE ARENA.  The expert cache takes what is free minus the reserve, so
     // everything allocated after it comes out of the reserve.  The native head (~0.5 GB with IQ3_S) was loaded after
@@ -3442,6 +3455,8 @@ int main(int argc, char** argv) {
         std::fprintf(stderr, "strata generate: layer split: CUDA0 runs layers 0-%lld\n", (long long) (split_at[0] - 1));
 
     std::array<strata::core::RemoteExperts, 3> remote_experts;
+    std::unique_ptr<strata::core::RemoteExpertOpt> remote_opt;
+    if (o.remote_expert_opt && o.expert_cache_remote[0] > 0) remote_opt = std::make_unique<strata::core::RemoteExpertOpt>();
     const bool multi_remote = o.expert_cache_remote[1] > 0 || o.expert_cache_remote[2] > 0;
     if (o.expert_cache_remote[0] > 0) {
         if (o.expert_cache <= 0 || profile.empty() || o.no_pool) {
@@ -3455,7 +3470,8 @@ int main(int argc, char** argv) {
                 for (const auto& pr : st->profile)
                     if (st->cache.slot_of(pr.first, pr.second) < 0) ranked.push_back(pr);
         }
-        if (multi_remote) {
+        if (multi_remote || std::any_of(o.expert_cache_remote_auto.begin(), o.expert_cache_remote_auto.end(),
+                                        [](bool automatic) { return automatic; })) {
             // The shipped frequency profile names only 8000 of 24576 experts. Once exhausted,
             // fill remaining VRAM from unranked pairs in expert-then-layer order: this spreads
             // the tail across all layers instead of concentrating it on layer zero.
@@ -3513,8 +3529,10 @@ int main(int argc, char** argv) {
                 if (st->cache.slot_of(pr.first, pr.second) >= 0)
                     claimed[(size_t) pr.first * (size_t) g.n_expert + (size_t) pr.second] = 1;
         for (int r = 0; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0) {
+            if (remote_opt) remote_opt->attach(remote_experts[(size_t) r]);
             if (!remote_experts[(size_t) r].open(remote_dev[r], o.expert_cache_remote[(size_t) r],
-                     g.n_layers, g.n_expert, by_device[(size_t) r], xcache, *srcp, claimed, err)) {
+                     g.n_layers, g.n_expert, by_device[(size_t) r], xcache, *srcp, claimed, err,
+                     o.expert_cache_remote_auto[(size_t) r])) {
                 std::fprintf(stderr, "strata generate: %s\n", err.c_str());
                 return 1;
             }
@@ -3543,6 +3561,7 @@ int main(int argc, char** argv) {
                      (long long) (peer.resident() + xcache.slots()), (long long) (g.n_layers * g.n_expert));
     }
 
+    if (remote_opt && !remote_opt->init(err)) { std::fprintf(stderr, "strata generate: %s\n", err.c_str()); return 1; }
     Drive drive;
     for (int r = 0; r < 3; ++r) if (o.expert_cache_remote[(size_t) r] > 0)
         drive.d.remote[drive.d.remote_count++] = &remote_experts[(size_t) r];
@@ -4271,8 +4290,9 @@ int main(int argc, char** argv) {
             }
         strata::prefill::Prefill::set_pinned_share(total ? (double) pinned / (double) total : 1.0);
     }
-    auto lend_slots = [&](int64_t c) -> int64_t {
-        const uint64_t need = strata::prefill::Prefill::bytes_needed(g, ss, c);
+    // bytes -> slots for CUDA0's cache: exact when it knows its per-slot offsets (a native pack's blobs differ
+    // per layer), otherwise max_blob each.
+    auto slots_from_bytes = [&](uint64_t need) -> int64_t {
         const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob;
         int64_t k = (int64_t) ((need + (uint64_t) blob - 1) / (uint64_t) blob);
         if (xcache.slot_offsets() != nullptr) {   // sized slots: take slots from the end until they hold `need`
@@ -4282,6 +4302,18 @@ int main(int argc, char** argv) {
         }
         return k;
     };
+    // ...and its inverse, for the ring's room: the bytes the cache's last `k` slots hold.  k can arrive <= 0: the
+    // caller's lend budget is `min(slots - 128, pct * slots / 100)` and goes negative on a cache smaller than 128
+    // slots, where `slots - k` would index slot_offsets() past its end.  Such a cache lends nothing, so say so.
+    auto bytes_from_slots = [&](int64_t k) -> uint64_t {
+        if (k <= 0) return 0;
+        if (xcache.slot_offsets() != nullptr)
+            return (uint64_t) (xcache.bytes() - (int64_t) xcache.slot_offsets()[xcache.slots() - k]);
+        return (uint64_t) k * (uint64_t) strata::kernels::cpu::expert_layout().max_blob;
+    };
+    auto lend_slots = [&](int64_t c) -> int64_t {
+        return slots_from_bytes(strata::prefill::Prefill::bytes_needed(g, ss, c));
+    };
     // `lend_bytes` went with the single-cache serve loan: a participant's loan is priced by `part_bytes` from its
     // OWN cache, and the only other user of the old helper was the serve path's own relayout.
     auto request_chunk = [](int64_t tokens, int64_t max_chunk) -> int64_t {
@@ -4292,7 +4324,7 @@ int main(int argc, char** argv) {
         return std::min(max_chunk, rounded);
     };
     // The prompt path's chunk and the slots it borrows for its buffers: the requested chunk halved until it fits,
-    // or with --prefill auto the largest of kAutoChunks whose buffers take at most kAutoLendPct % of the slots (a
+    // or with --prefill auto the largest chunk whose buffers take at most kAutoLendPct % of the slots (a
     // lent slot's expert is streamed during the prompt and refilled after it; measured on a 12 GB card, 32K Q2_0
     // prompt: 4096 791 tok/s, 6144 878, 8192 973 with 69% of the slots lent).  A request lends only what its own
     // prompt needs (Prefill::relayout), so a big chunk costs short prompts nothing.  0 = none fits.
@@ -4305,20 +4337,95 @@ int main(int argc, char** argv) {
         return v ? (int64_t) std::atoi(v)
                  : (int64_t) (strata::prefill::Prefill::pinned_share() >= 0.9 ? 90 : 85);
     }();
+    // The largest chunk --prefill auto may pick: the operator's ceiling (--prefill auto:N) and never past the
+    // context, but a bare `auto` always reaches 8192.  32768 and 16384 stay opt-in (#282): a 32K prompt with
+    // IQ2_XS (RTX 5090, 64K context) read at 5,624 tok/s in 8192-token chunks and 6,465 in one 32768 chunk (40K
+    // -> 18K experts streamed; an NVFP4 pack at 262K: 3,535 -> 5,201).
+    const int64_t auto_ceiling = std::max<int64_t>(8192, std::min<int64_t>(o.prefill_auto_max, o.max_context));
+    // The auto scan used to walk a fixed list of sizes - 32768, 16384, 8192, 6144, ... - and take the first that
+    // fit.  That list is coarse exactly where a rig needs it: this one affords ~8,700 tokens and was handed 8192,
+    // and 8704 is not on it.  `bytes_needed` is a sum of (T x positive constant) terms plus a max of such sums,
+    // so it rises monotonically with T, and so does every test the scan applies - which makes the largest size
+    // that fits a bisection on the 256-token grid the prompt path already works on (`request_chunk` rounds up to
+    // it).  log2(32768/256) = 7 probes, against the list's 10, at one `bytes_needed` per probe.
+    auto biggest_chunk = [](int64_t ceiling, auto&& ok) -> int64_t {
+        if (ceiling < 256) return 0;
+        int64_t lo = 1, hi = ceiling / 256, best = 0;   // n = T / 256, and ok() is monotone in n
+        while (lo <= hi) {
+            const int64_t mid = lo + (hi - lo) / 2;
+            if (ok(mid * 256)) { best = mid * 256; lo = mid + 1; } else hi = mid - 1;
+        }
+        return best;
+    };
     auto plan_lend = [&](int64_t& chunk) -> int64_t {
-        // 32768 and 16384 (#282, opt-in: --prefill auto:32768): a 32K prompt with IQ2_XS (RTX 5090, 64K context)
-        // read at 5,624 tok/s in 8192-token chunks and 6,465 in one 32768 chunk (40K -> 18K experts streamed; an
-        // NVFP4 pack at 262K: 3,535 -> 5,201)
-        static constexpr int64_t kAutoChunks[] = {32768, 16384, 8192, 6144, 4096, 3072, 2048, 1024, 512, 256};
         auto slots_for = lend_slots;
-        if (o.prefill_auto) {
+        // 0.1.39's list with its ring (also the size up to which a prompt keeps that ring under #583)
+        auto old_rule = [&]() -> int64_t {
+            strata::prefill::Prefill::set_ring_budget(0, 0);
+            static constexpr int64_t kAutoChunks[] = {32768, 16384, 8192, 6144, 4096, 3072, 2048, 1024, 512, 256};
             for (const int64_t c : kAutoChunks) {
                 // above 8192: only when asked for, and only when a prompt of the context can use it
                 if (c > 8192 && (c > o.prefill_auto_max || c > o.max_context)) continue;
                 const int64_t k = slots_for(c);
-                if (k + 128 <= xcache.slots() && k * 100 <= kAutoLendPct * xcache.slots()) { chunk = c; return k; }
+                if (k + 128 <= xcache.slots() && k * 100 <= kAutoLendPct * xcache.slots()) return c;
             }
             return 0;
+        };
+        if (o.prefill_auto && !strata::prefill::Prefill::ring_bytes_enabled()) {   // STRATA_RING_BYTES=0: 0.1.39's list
+            const int64_t c = old_rule();
+            if (c > 0) { chunk = c; return slots_for(c); }
+            return 0;
+        }
+        if (o.prefill_auto) {   // the default since 0.1.39b (#583)
+            // The chunk and the ring are one budget, and the ring is the better buy.  Measured on this 4-way
+            // IQ3_S rig (2x RTX 3060 + 2x RTX 5060, CUDA3 lending at its 90% cap) on a 120K prompt: 8960 tokens
+            // with the 17-slot ring that leaves reads at 963 tok/s, 8192/130 at 1,008, 7168 with the ring full at
+            // its 199-slot byte budget at 1,000, and 5632/199 at 915.  A ring slot is worth ~0.53 tok/s there and
+            // a chunk token ~0.05, so the 69 slots between a full ring and 8192's 130 are worth more than the
+            // 512 chunk tokens they cost - and once the ring IS full, shrinking the chunk further buys nothing.
+            // So: the largest chunk that still leaves the ring full.  Only a rig where no chunk at all can afford
+            // one falls back to the old rule, which takes the largest chunk whose ring clears kRingMin (the
+            // value at or below which ring_slots() returns STAGE and streaming is off).
+            constexpr int64_t kRingMin = 16;
+            const int64_t small = old_rule();   // 0.1.39's chunk (and its ring for the prompts that fit it)
+            const int64_t ring_max = strata::prefill::Prefill::ring_cap_for(small);
+            const int64_t budget = std::min(xcache.slots() - 128, kAutoLendPct * xcache.slots() / 100);
+            // The room is a BYTE budget.  A ring slot is max_blob bytes (`carve` lays out one whole blob each),
+            // while these cache slots hold their own layer's blob, which is smaller than max_blob unless the cache
+            // happens to hold the pack's biggest layer - on the 4-way IQ3_S rig that is 2.15 MiB a slot against a
+            // 2.54 MiB max_blob.  Counting the room in slots priced the ring at 0.85x what it costs, `fits` then
+            // rejected the whole chunk instead of shrinking the ring, and the scan stopped a step short: measured
+            // against chunk 6144/ring 199, 8192 was refused where the same bytes afford a ~87-slot ring.
+            const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob;
+            const uint64_t avail = bytes_from_slots(budget);
+            auto room_of = [&](int64_t t) -> int64_t {
+                const uint64_t nonring = strata::prefill::Prefill::bytes_needed_no_ring(g, ss, t);
+                return std::min((int64_t) ((avail - std::min(avail, nonring)) / (uint64_t) blob), ring_max);
+            };
+            // `room_of` is capped at ring_max, so "the ring is full" is exactly room == ring_max, and both that
+            // test and the ones below it only get harder as t grows - the bisection stays valid.
+            auto scan = [&](int64_t floor) -> int64_t {
+                return biggest_chunk(auto_ceiling, [&](int64_t t) -> bool {
+                    const int64_t room = room_of(t);
+                    if (room < floor) return false;
+                    strata::prefill::Prefill::set_ring_budget((int) room, 0);
+                    // ring_slots() now returns `room` - unless STRATA_PREFILL_RING overrides it, in which case the
+                    // chunk has to fit THAT ring, which is the pre-fix rule and the A/B arm
+                    const int64_t k = slots_for(t);
+                    return k + 128 <= xcache.slots() && k * 100 <= kAutoLendPct * xcache.slots();
+                });
+            };
+            int64_t c = scan(ring_max);
+            if (c == 0 && ring_max < strata::prefill::Prefill::ring_default_slots()) c = scan(kRingMin);
+            // a prompt that fits 0.1.39's chunk keeps 0.1.39's ring (one chunk: the smaller ring only slowed it)
+            if (small >= c) {   // the scan bought nothing: 0.1.39's choice
+                strata::prefill::Prefill::set_ring_budget(0, 0);
+                if (small > 0) { chunk = small; return slots_for(small); }
+                return 0;
+            }
+            strata::prefill::Prefill::set_ring_budget((int) room_of(c), small);
+            chunk = c;
+            return slots_for(c);
         }
         for (int64_t c = chunk; c >= 256; c /= 2) {
             const int64_t k = slots_for(c);
@@ -4432,9 +4539,9 @@ int main(int argc, char** argv) {
             int64_t lent_chunk = 0;
             std::vector<std::pair<int32_t, int32_t>> lent;
         };
-        auto part_slots = [&](const PfPart& p, int64_t c) -> int64_t {
-            const uint64_t need = strata::prefill::Prefill::bytes_needed(g, *p.ses, c);
-            strata::core::ExpertCache& xc = *p.cache;
+        // bytes -> slots for one cache: exact when it knows its per-slot offsets (a native pack's blobs differ
+        // per layer), otherwise max_blob each.
+        auto cache_slots_for = [&](const strata::core::ExpertCache& xc, uint64_t need) -> int64_t {
             if (xc.slot_offsets() != nullptr) {   // sized slots: from the end until they hold `need`
                 int64_t k = 0;
                 while (k < xc.slots() &&
@@ -4444,6 +4551,9 @@ int main(int argc, char** argv) {
             const int64_t blob = (int64_t) strata::kernels::cpu::expert_layout().max_blob;
             return (int64_t) ((need + (uint64_t) blob - 1) / (uint64_t) blob);
         };
+        auto part_slots = [&](const PfPart& p, int64_t c) -> int64_t {
+            return cache_slots_for(*p.cache, strata::prefill::Prefill::bytes_needed(g, *p.ses, c));
+        };
         auto part_bytes = [&](const PfPart& p, int32_t first) -> uint64_t {
             strata::core::ExpertCache& xc = *p.cache;
             return xc.slot_offsets() ? (uint64_t) (xc.bytes() - (int64_t) xc.slot_offsets()[first])
@@ -4490,54 +4600,121 @@ int main(int argc, char** argv) {
                     if (!fits_one(p, c, cap)) return false;
                 return true;
             };
-            static constexpr int64_t kAutoChunks[] = {32768, 16384, 8192, 6144, 4096, 3072, 2048, 1024, 512, 256};
-            auto pick = [&](const PfPart* only) -> int64_t {
-                if (o.prefill_auto) {
-                    for (const int64_t c : kAutoChunks) {
-                        if (c > 8192 && (c > o.prefill_auto_max || c > o.max_context)) continue;   // #282, as plan_lend
-                        if (fits(c, true, only)) return c;
-                    }
-                } else {
-                    for (int64_t c = o.prefill_chunk; c >= 256; c /= 2)
-                        if (fits(c, false, only)) return c;
+            // The chunk and the ring are one budget.  The ring used to be a fixed 384 slots taken off the top
+            // before the chunk was considered: 506 MiB on a Q2_0 pack, where it was tuned, but 1,912 MiB on
+            // Q8_0, more than half an 8 GB card's cache - which is what kept that card at a 1024-token chunk.
+            // Size the chunk first and hand the ring what the smallest participant leaves over, up to the ring's
+            // byte budget - but hand it ALL of it, as in plan_lend: the ring is the better buy per byte (the
+            // measurements are there), so the chunk to take is the largest one that still leaves the ring full,
+            // not the largest one that leaves it anything at all.  kRingMin is the fallback floor for a rig where
+            // no chunk can afford a full ring; ring_slots() returns STAGE at or below it, which turns streaming
+            // off.
+            constexpr int64_t kRingMin = 16;
+            // The room is a BYTE budget, as in plan_lend: a ring slot is max_blob bytes, but a cache slot holds
+            // its own layer's blob, which on this rig averages 0.85 of it.  In slots the ring looked 18% cheaper
+            // than it is, and a chunk whose ring only fitted after that discount was refused outright - the scan
+            // stopped one step short while the chunk it did take had room to spare for the cap-sized ring.
+            const int64_t kBlob = (int64_t) strata::kernels::cpu::expert_layout().max_blob;
+            auto ring_room = [&](int64_t c, const PfPart* only, int64_t cap) -> int64_t {
+                int64_t room = cap;
+                for (const PfPart& p : pf_parts) {
+                    if (only != nullptr && &p != only) continue;
+                    // clamped at 0: a cache under 128 slots has a negative lend budget, and `slots - budget` is
+                    // then `slots`, which is off_[slots] - the end of the prefix sum, so that cache lends 0 bytes
+                    // rather than reading one entry past its offsets
+                    const int64_t budget = std::max<int64_t>(0, std::min(p.cache->slots() - 128,
+                                                                        kAutoLendPct * p.cache->slots() / 100));
+                    const uint64_t avail = part_bytes(p, (int32_t) (p.cache->slots() - budget));
+                    const uint64_t nonring = strata::prefill::Prefill::bytes_needed_no_ring(g, *p.ses, c);
+                    room = std::min(room, (int64_t) ((avail - std::min(avail, nonring)) / (uint64_t) kBlob));
+                }
+                return room;
+            };
+            // The largest chunk on the 256-token grid that `only` - or, with null, every stage - can lend with the
+            // ring still at its full byte budget (ring_room is capped at ring_max, so that is room == ring_max),
+            // under the operator's ceiling (auto_ceiling keeps a bare `auto` at 8192, #282).  Failing that, the
+            // largest chunk whose ring clears kRingMin.
+            // 0.1.39's list with its ring: the chunk it would pick, and the size up to which a prompt keeps that ring
+            auto old_pick = [&](const PfPart* only) -> int64_t {
+                strata::prefill::Prefill::set_ring_budget(0, 0);
+                static constexpr int64_t kAutoChunks[] = {32768, 16384, 8192, 6144, 4096, 3072, 2048, 1024, 512, 256};
+                for (const int64_t c : kAutoChunks) {
+                    if (c > 8192 && (c > o.prefill_auto_max || c > o.max_context)) continue;   // #282, as plan_lend
+                    if (fits(c, true, only)) return c;
                 }
                 return 0;
             };
-            const int64_t chunk = pick(nullptr);
+            auto scan = [&](const PfPart* only) -> int64_t {
+                if (!o.prefill_auto) {   // an explicit --prefill is the operator's number, and a loan of it only has to fit
+                    for (int64_t c = o.prefill_chunk; c >= 256; c /= 2)
+                        if (fits(c, false, only)) return c;
+                    return 0;
+                }
+                if (!strata::prefill::Prefill::ring_bytes_enabled()) return old_pick(only);   // STRATA_RING_BYTES=0
+                const int64_t small = old_pick(only);
+                const int64_t cap = strata::prefill::Prefill::ring_cap_for(small);
+                auto probe = [&](int64_t floor) -> int64_t {
+                    return biggest_chunk(auto_ceiling, [&](int64_t t) -> bool {
+                        const int64_t room = ring_room(t, only, cap);
+                        if (room < floor) return false;
+                        strata::prefill::Prefill::set_ring_budget((int) room, 0);
+                        // ring_slots() now returns `room` - unless STRATA_PREFILL_RING overrides it, in which
+                        // case the chunk has to fit THAT ring, which is the pre-fix rule and the A/B arm
+                        return fits(t, true, only);
+                    });
+                };
+                int64_t c = probe(cap);
+                if (c == 0 && cap < strata::prefill::Prefill::ring_default_slots()) c = probe(kRingMin);
+                return std::max(c, small);   // never a smaller chunk than 0.1.39's
+            };
+            const int64_t chunk = scan(nullptr);
             // #448: a small card in a layer split caps every stage's chunk (an RTX 3080's 512-slot cache held a
             // 32 GB card's split to 512 tokens: prompts 6.2x slower, decode the same).  Named when it bites, so the
             // regression is one log line: each stage that cannot fund the chunk CUDA0 alone would read in.
-            if (pf_parts.size() > 1) {
-                const int64_t alone = pick(&pf_parts[0]);
-                if (alone > chunk) {
-                    for (size_t i = 1; i < pf_parts.size(); ++i) {
-                        const PfPart& p = pf_parts[i];
-                        if (fits_one(p, alone, o.prefill_auto)) continue;
-                        const int dev = p.dev < 0 ? 0 : p.dev;
-                        cudaDeviceProp prop{};
-                        if (cudaGetDeviceProperties(&prop, dev) != cudaSuccess) {
-                            (void) cudaGetLastError();
-                            prop.name[0] = 0;
-                        }
-                        const std::string pct =
-                            o.prefill_auto ? ", and lend at most " + std::to_string(kAutoLendPct) + "%" : "";
-                        std::fprintf(stderr, "strata serve: WARNING: prompt chunk %lld tokens, not %lld: CUDA%d (%s) "
-                                             "has %lld expert-cache slots, and a %lld-token chunk borrows %lld of them "
-                                             "(it must keep 128%s) - prompts read slower than on CUDA0 alone (#448)\n",
-                                     (long long) chunk, (long long) alone, dev, prop.name,
-                                     (long long) p.cache->slots(), (long long) alone,
-                                     (long long) part_slots(p, alone), pct.c_str());
-                        // the helper tiers start at CUDA1 without a split (and are enabled in order)
-                        std::fprintf(stderr, "strata serve:   a card this small can serve as a helper expert cache "
-                                             "instead of a split stage: without --layer-split, with %s "
-                                             "(docs/SECOND_GPU.md)\n",
-                                     dev == 1 ? "--expert-cache-device1 N" : "--expert-cache-device1..3 N, in order");
+            const int64_t alone = pf_parts.size() > 1 ? scan(&pf_parts[0]) : chunk;
+            // both scans left the ring override on their own last trial: put it back on the chunk that won, which
+            // is what `init` lays out and what the INFO line reports
+            if (chunk > 0 && o.prefill_auto && strata::prefill::Prefill::ring_bytes_enabled()) {
+                // a prompt that fits 0.1.39's chunk keeps 0.1.39's ring (one chunk: the smaller ring only slowed
+                // it); the byte-budget ring is for the chunks past it
+                const int64_t small = old_pick(nullptr);
+                strata::prefill::Prefill::set_ring_budget(
+                    small >= chunk ? 0 : (int) ring_room(chunk, nullptr, strata::prefill::Prefill::ring_cap_for(small)),
+                    small);
+            }
+            if (chunk > 0 && alone > chunk) {
+                for (size_t i = 1; i < pf_parts.size(); ++i) {
+                    const PfPart& p = pf_parts[i];
+                    if (fits_one(p, alone, o.prefill_auto)) continue;
+                    const int dev = p.dev < 0 ? 0 : p.dev;
+                    cudaDeviceProp prop{};
+                    if (cudaGetDeviceProperties(&prop, dev) != cudaSuccess) {
+                        (void) cudaGetLastError();
+                        prop.name[0] = 0;
                     }
+                    const std::string pct =
+                        o.prefill_auto ? ", and lend at most " + std::to_string(kAutoLendPct) + "%" : "";
+                    std::fprintf(stderr, "strata serve: WARNING: prompt chunk %lld tokens, not %lld: CUDA%d (%s) "
+                                         "has %lld expert-cache slots, and a %lld-token chunk borrows %lld of them "
+                                         "(it must keep 128%s) - prompts read slower than on CUDA0 alone (#448)\n",
+                                 (long long) chunk, (long long) alone, dev, prop.name,
+                                 (long long) p.cache->slots(), (long long) alone,
+                                 (long long) part_slots(p, alone), pct.c_str());
+                    // the helper tiers start at CUDA1 without a split (and are enabled in order)
+                    std::fprintf(stderr, "strata serve:   a card this small can serve as a helper expert cache "
+                                         "instead of a split stage: without --layer-split, with %s "
+                                         "(docs/SECOND_GPU.md)\n",
+                                 dev == 1 ? "--expert-cache-device1 N" : "--expert-cache-device1..3 N, in order");
                 }
             }
             if (chunk > 0) {
                 if (o.prefill_auto)
-                    std::fprintf(stderr, "strata serve: prompt chunk auto: %lld tokens\n", (long long) chunk);
+                    // the RESOLVED ring, not the room offered: ring_slots() clamps to [16, RING_MAX] and a chunk
+                    // under stream_all_min() gets STAGE, and STRATA_PREFILL_RING overrides both.  This is what
+                    // `init` lays out, and what the INFO line reports to the Monitor tab.
+                    std::fprintf(stderr, "strata serve: prompt chunk auto: %lld tokens, a %lld-slot ring\n",
+                                 (long long) chunk,
+                                 (long long) strata::prefill::Prefill::ring_slots_for(chunk));
                 else if (chunk != o.prefill_chunk)
                     std::fprintf(stderr, "strata serve: prompt chunk %lld -> %lld tokens so its buffers fit in "
                                          "every expert cache\n", (long long) o.prefill_chunk, (long long) chunk);
@@ -4719,6 +4896,19 @@ int main(int argc, char** argv) {
                 err.clear();
             }
         }
+        // a layer split: a one-chunk prompt runs the stages one after the other, so each stage's prompt path gets the
+        // next stage's GPU (the last one the first's) to stream and compute a share of its experts meanwhile
+        if (multi_gpu && !stages.empty()) {
+            std::vector<std::pair<strata::prefill::Prefill*, int>> paths{{&sp, 0}};
+            for (const auto& st : stages) paths.emplace_back(&st->sp, st->dev);
+            for (size_t i = 0; i < paths.size(); ++i) {
+                const strata::core::OnDevice on(paths[i].second);
+                if (!paths[i].first->set_stage_helper(paths[(i + 1) % paths.size()].first, err)) {
+                    std::fprintf(stderr, "strata serve: %s - stage %zu reads its prompts alone\n", err.c_str(), i);
+                    err.clear();
+                }
+            }
+        }
         mem_mark("the head and the prompt path");
         // #340: STRATA_SPLIT_SMALL_OWN=S (tokens): on a layer split, every stage that borrows keeps the slots for an
         // S-token chunk's buffers for the whole session (0.1.29's own buffers, carved from the tail of its cache):
@@ -4773,6 +4963,7 @@ int main(int argc, char** argv) {
         strata::core::Verifier ver;
         strata::core::VerifyHits vh;
         vh.d_res = thits.d_res;
+        vh.h_res = host_res.empty() ? nullptr : host_res.data();
         vh.cache_base = thits.cache_base;
         vh.blob = thits.blob;
         vh.slot_off = xcache.slot_offsets();   // E-6: the device plan's pointers
@@ -4819,10 +5010,12 @@ int main(int argc, char** argv) {
                     const strata::core::OnDevice on(gs.dev);
                     strata::core::VerifyHits vs;
                     vs.d_res = gs.d_res;
+                    vs.h_res = host_res.empty() ? nullptr : host_res.data();
                     vs.cache_base = gs.cache.device_slot(0);
                     vs.blob = thits.blob;
                     vs.slot_off = gs.cache.slot_offsets();
                     vs.n_slots = gs.cache.slots();
+                    gs.ver.set_remote_expert_opt(remote_opt.get());
                     ok_s = gs.ver.init(gs.wt, g, gs.ss, vs, gs.head.loaded() ? &gs.head : nullptr, o.spec, err);
                     split_drive.cache_base[st] = gs.cache.device_slot(0);
                     split_drive.cache_slot_off[st] = gs.cache.slot_offsets();
@@ -4840,6 +5033,7 @@ int main(int argc, char** argv) {
                           " (CUDA" + std::to_string(split_same ? 0 : stages[(size_t) st - 1]->dev) + ")";
             std::fprintf(stderr, "strata serve: layer split: layers %s, one hand-off per window\n", plan_s.c_str());
         }
+        ver.set_remote_expert_opt(remote_opt.get());
         if (!ver.init(wt, g, ss, vh, native_head.loaded() ? &native_head : nullptr, o.spec, err) ||
             !mtp.bind(last_st ? last_st->wt : wt, last_st ? &last_st->head : &native_head, ver.final_R_all(), err)) {
             std::fprintf(stderr, "strata serve: %s\n", err.c_str());
@@ -5073,11 +5267,14 @@ int main(int argc, char** argv) {
         }
         drive.d.plan = ver.plan_sink();
         drive.d.pcie_num = std::max(0, std::min(256, (int) (o.pcie_frac * 256.0 + 0.5)));
-        if (o.adapt_every > 0 && o.adapt_swaps > 0) drive.d.usage.assign((size_t) (g.n_layers * g.n_expert), 0.0f);
+        const bool all_experts_resident = !host_res.empty() &&
+            std::all_of(host_res.begin(), host_res.end(), [](int32_t r) { return r >= 0; });
+        if (o.adapt_every > 0 && o.adapt_swaps > 0 && !all_experts_resident)
+            drive.d.usage.assign((size_t) (g.n_layers * g.n_expert), 0.0f);
         // #477 --expert-profile-save: what the adaptive tier learned, kept across restarts (opt-in; off: `heat` stays
         // empty and nothing below runs).  It needs the adaptive tier's counts and the residency table.
         std::vector<double> heat;
-        if (!o.expert_profile_save.empty()) {
+        if (!o.expert_profile_save.empty() && !all_experts_resident) {
             if (drive.d.usage.empty() || host_res.empty())
                 std::fprintf(stderr, "strata serve: --expert-profile-save needs the adaptive tier (--adapt-every and "
                                      "--adapt-swaps above 0) and --expert-profile: nothing will be saved\n");
@@ -5127,6 +5324,7 @@ int main(int argc, char** argv) {
         // the VRAM tier follows the conversation (the same rule as the speculative loop below)
         auto adapt = [&]() -> bool {
             if (!pending.empty()) return true;   // the previous swaps are still in flight
+            if (remote_opt && !remote_opt->adapt(drive.d.usage, host_res, pending, o.adapt_swaps, *srcp)) return false;
             struct Swap { float gain; int32_t layer, in, out; };
             std::vector<Swap> swaps;
             std::vector<std::pair<float, int32_t>> cand, vict;
@@ -5136,7 +5334,7 @@ int main(int argc, char** argv) {
                 const float* u = drive.d.usage.data() + l * g.n_expert;
                 const int32_t* r = host_res.data() + l * g.n_expert;
                 for (int32_t e = 0; e < (int32_t) g.n_expert; ++e) {
-                    if (r[e] < 0) { if (u[e] >= 2.0f && !(peer.valid() && peer.has(l, e))) cand.emplace_back(u[e], e); }
+                    if (r[e] < 0) { if (u[e] >= 2.0f && !(peer.valid() && peer.has(l, e)) && !(remote_opt && remote_opt->owns(l, e))) cand.emplace_back(u[e], e); }
                     else vict.emplace_back(u[e], e);
                 }
                 if (cand.empty() || vict.empty()) continue;
@@ -6884,6 +7082,7 @@ int main(int argc, char** argv) {
         strata::core::Verifier ver;
         strata::core::VerifyHits vh;
         vh.d_res = thits.d_res;
+        vh.h_res = host_res.empty() ? nullptr : host_res.data();
         vh.cache_base = thits.cache_base;
         vh.blob = thits.blob;
         vh.slot_off = xcache.slot_offsets();   // E-6: the device plan's pointers

tools/test_setup_golden.json

@@ -482,7 +482,8 @@
     "--max-context",
     "65536",
     "--kv",
-    "int8"
+    "int8",
+    "--remote-expert-opt"
    ],
    "cwd": "<T>",
    "exe": "<T>/engine/<EXE>",
@@ -526,7 +527,8 @@
     "--max-context",
     "65536",
     "--kv",
-    "int8"
+    "int8",
+    "--remote-expert-opt"
    ],
    "cwd": "<T>",
    "exe": "<T>/engine/<EXE>",

tools/test_setup_remote_opt.py

@@ -0,0 +1,74 @@
+"""0.1.39b (#578): setup recommends --remote-expert-opt for a config on two or more GPUs, never for one, and keeps it
+out when asked (setup's --no-remote-expert-opt, or "remote_expert_opt": false in the config).
+
+    python -m unittest tools.test_setup_remote_opt
+"""
+from __future__ import annotations
+
+import sys
+import unittest
+from pathlib import Path
+from unittest import mock
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+sys.path.insert(0, str(ROOT / "tools"))
+import setup  # noqa: E402
+from test_setup_golden import PROFILES, install  # noqa: E402
+
+FLAG = "--remote-expert-opt"
+
+
+class Rule(unittest.TestCase):
+    def setUp(self):
+        p = mock.patch.object(setup, "ok", lambda *a: None)
+        p.start()
+        self.addCleanup(p.stop)
+
+    def test_two_gpus_get_it_once(self):
+        cfg = {"args": ["--spec", "4"], "gpu": [0, 1]}
+        setup.recommend_remote_expert_opt(cfg)
+        setup.recommend_remote_expert_opt(cfg)
+        self.assertEqual(cfg["args"], ["--spec", "4", FLAG])
+
+    def test_one_gpu_is_untouched(self):
+        for gpu in (None, 0, [0]):
+            cfg = {"args": ["--spec", "4"]} if gpu is None else {"args": ["--spec", "4"], "gpu": gpu}
+            setup.recommend_remote_expert_opt(cfg)
+            self.assertEqual(cfg["args"], ["--spec", "4"], gpu)
+
+    def test_overrides(self):
+        cfg = {"args": ["--spec", "4", FLAG], "gpu": [0, 1]}
+        setup.recommend_remote_expert_opt(cfg, off=True)
+        self.assertEqual(cfg["args"], ["--spec", "4"])
+        cfg = {"args": ["--spec", "4"], "gpu": [0, 1, 2], "remote_expert_opt": False}
+        setup.recommend_remote_expert_opt(cfg)
+        self.assertEqual(cfg["args"], ["--spec", "4"])
+
+    def test_the_users_key_is_kept_on_a_rerun(self):
+        self.assertNotIn("remote_expert_opt", setup.SETUP_KEYS)   # #629: carried over, so the opt-out survives
+
+
+class Install(unittest.TestCase):
+    def test_install_on_two_gpus(self):
+        ram, found = PROFILES["47GB-2x16GB"]
+        args = ["--family", "qwen", "--no-start", "--model", "Q2_0"]
+        code, text, cfg, _ = install(ram, found, args)
+        self.assertEqual(code, 0, text)
+        self.assertIsInstance(cfg.get("gpu"), list)
+        self.assertEqual(cfg["args"].count(FLAG), 1)
+        self.assertIn("--no-remote-expert-opt leaves it out", text)
+        code, text, off, _ = install(ram, found, args + ["--no-remote-expert-opt"])
+        self.assertEqual(code, 0, text)
+        self.assertNotIn(FLAG, off["args"])
+        self.assertEqual([a for a in cfg["args"] if a != FLAG], off["args"])
+
+    def test_install_on_one_gpu(self):
+        ram, found = PROFILES["96GB-1x16GB"]
+        code, text, cfg, _ = install(ram, found, ["--family", "qwen", "--no-start", "--model", "Q2_0"])
+        self.assertEqual(code, 0, text)
+        self.assertNotIn(FLAG, cfg["args"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tools/test_setup_risk.py

@@ -276,7 +276,7 @@ class StartOnSeveralGpus(unittest.TestCase):
         code, out, asked, cfg = self.offer(["--resident-experts"], "y")
         self.assertIn("[n]", asked[0])
         self.assertEqual(cfg["gpu"], [0, 1])
-        self.assertEqual(cfg["args"], ["--mmap-experts"])
+        self.assertEqual(cfg["args"], ["--mmap-experts", "--remote-expert-opt"])   # 0.1.39b: #578 on 2+ GPUs
         code, out, asked, cfg = self.offer(["--mmap-experts"], None)           # other configs: as before
         self.assertEqual(cfg["gpu"], [0, 1])
 
@@ -296,7 +296,7 @@ class StartOnSeveralGpus(unittest.TestCase):
             cfg = json.loads(p.read_text())
         self.assertIsNone(code, out)
         self.assertEqual(cfg["gpu"], [0, 1])
-        self.assertEqual(cfg["args"], ["--mmap-experts"])
+        self.assertEqual(cfg["args"], ["--mmap-experts", "--remote-expert-opt"])   # 0.1.39b: #578 on 2+ GPUs
         self.assertIn("no layer split yet", out)
         self.assertTrue(call.called)
 

tools/test_setup_unsloth.py

@@ -480,7 +480,7 @@ class LayerSplit(unittest.TestCase):
         self.assertIsNone(code, out)
         self.assertTrue(started)
         self.assertEqual(cfg["gpu"], [0, 1])
-        self.assertEqual(cfg["args"], ["--pack", "p", "--kv", "int8"])
+        self.assertEqual(cfg["args"], ["--pack", "p", "--kv", "int8", "--remote-expert-opt"])   # 0.1.39b: #578
         self.assertIn("no RAM budget", out)
 
     def test_start_with_gpus_refused_without_the_ram(self):