feat(ggml-cuda): MoE expert-streaming policy + mechanism + fused Q4 kernels

ggml/src/ggml-cuda/EXPERT-STREAM-INTEGRATION.md

@@ -0,0 +1,65 @@
+# MoE expert streaming — integration guide
+
+Components (all unit-tested on RTX 3090, `sm_86`):
+
+| File | Role | Test |
+|------|------|------|
+| `expert-stream.{h,cpp}` | P3 residency + P1 prefetch *policy* (pure C++) | `tests/test-expert-stream.cpp` — 15/15 |
+| `expert-cache.{h,cu}` | VRAM slot cache: streams, events, eviction *mechanism* | `tests/test-expert-cache.cu` — 21/21 |
+| `expert-stream-hooks.{h,cu}` | per-layer registry + hot-path hooks *integration* | `tests/test-expert-hooks.cu` — 10/10 |
+| `moe-fused` kernels | fused gather+dequant+route, Q4_0 & Q4_K | `tests/test-moe-fused-q4{,k}.cu` — PASS |
+
+The hooks are **no-ops until a layer is registered**, so every splice below is
+opt-in: with no registration the MoE path is byte-for-byte unchanged.
+
+## Splice points
+
+### 0. Registration — model load (once per MoE layer)
+Experts for the layer must live in **pinned host memory** (the streaming source).
+```cpp
+ggml_moe_stream::moe_register(layer_id, n_experts, bytes_per_expert,
+                              n_resident_slots /* = VRAM budget */,
+                              host_expert_base /* pinned */, copy_stream);
+```
+`n_resident_slots` is the "13.3 GiB" budget expressed in experts.
+
+### 1. observe() — `ggml-cuda.cu`, `ggml_cuda_mul_mat_id`
+The routing is already read to host at the `ids_host` loop (~L2613–2629). Right
+after that loop, feed the routed experts to the residency policy:
+```cpp
+// after ids_host is populated and tokens_per_expert computed
+std::vector<int32_t> routed;
+for (int64_t i = 0; i < ne02; ++i)
+    if (tokens_per_expert[i] > 0) routed.push_back((int32_t) i);
+ggml_moe_stream::moe_observe(layer_id, routed.data(), (int) routed.size());
+```
+The fast paths (mmvq/mmq/mmf at ~L2562–2582) carry `ids` too; observe can be
+hoisted to a shared helper covering both. `layer_id` comes from tagging the
+expert tensor at load (e.g. `tensor->op_params` or a `->extra` field).
+
+### 2. ensure_resident() — the expert backing buffer (the streaming win)
+Where the per-expert weight pointer is computed (`src0->data + i02*nb[2]`),
+resolve cold experts through the cache instead:
+```cpp
+const void * w = ggml_moe_stream::moe_expert_ptr(layer_id, (int) i02, stream);
+// use `w` as expert i02's weight base; the cache has streamed it H2D and the
+// returned stream waits on the in-flight copy. NULL => layer not registered,
+// fall back to src0->data + i02*nb[2] (unchanged behavior).
+```
+This is the only edit that changes data movement; it belongs at the buffer/op
+seam, not inside the matmul kernels.
+
+### 3. prefetch() — the DFlash draft pass (hides the x4 latency)
+After the draft model predicts the next tokens' routing, before the verify pass:
+```cpp
+// predicted = expert ids the draft expects the verify pass to use, this layer
+ggml_moe_stream::moe_prefetch(layer_id, predicted, /*cap*/ slots_free);
+```
+Issued on `copy_stream`, overlapped with draft compute; by verify time the
+experts are resident (proven by `test-expert-cache.cu` T4).
+
+## Validation gate (deferred — GPU busy)
+1. Build ggml-cuda with these TUs added to `ggml/src/ggml-cuda/CMakeLists.txt`.
+2. Register layers at load; point expert tensors at the pinned host store.
+3. `nsys` a 35B-A3B decode before/after: expect the PCIe-x4 H2D timeline to
+   de-saturate and tok/s to rise. That number is the real proof.

ggml/src/ggml-cuda/FINDINGS.md

@@ -0,0 +1,76 @@
+# MoE expert-streaming: what's proven, what's not (real-hardware findings)
+
+Measured on a real **Lucebox (lucebox3)** — RTX 3090 (`sm_86`), Ryzen AI MAX+ 395,
+PCIe Gen4 **x4** (~6.5 GB/s measured), native CUDA 12.0 build of `dflash_server`,
+Qwen3.6-35B-A3B (qwen35moe, 40 layers / 256 experts), Spark expert offload.
+
+This PR's components stand **independently of the speculative-prefetch idea**.
+Below is the honest scope: what the hardware proved, and where the prefetch
+thesis remains open.
+
+## 1. The two regimes (decisive)
+
+Same 800-token decode, Spark residency varied:
+
+| budget | resident experts | PCIe-RX (H2D) | x4 util | GPU sm | decode |
+|---|---|---|---|---|---|
+| 22 GiB (comfortable) | 9389 / 10240 (91%) | ~120 MB/s | ~2% | ~48% | 23.9 s (~33.5 tok/s) |
+| **8 GiB (tight)** | ~half cold | **~3,000 MB/s, peaks 5,180** | **~50–80%** | 73% | **44.7 s (~18 tok/s)** |
+
+**Expert streaming is the consumer-VRAM bottleneck *only when the model doesn't
+fit*.** At comfortable residency the x4 link is idle (~2%) and decode is
+autoregressive/compute-bound (GPU ~50% idle at batch=1). Force the budget tight
+and the x4 link **saturates and decode slows 1.9×** — the regime a 24 GB card
+hits with a model too big to mostly-fit (86 GB DeepSeek-V4-Flash, 78 GB
+Qwen3.5-122B) or long context. That tight regime is the only place expert
+prefetch has a target.
+
+## 2. TQ3_0 vs Q4_0 KV cache
+
+Same 22 GiB residency, identical decode:
+
+| KV cache | decode | ≈ tok/s |
+|---|---|---|
+| `tq3_0` (default auto) | 26.9 s | 30 |
+| `q4_0` | **23.9 s** | **33.5** |
+
+**TQ3_0 KV costs ~11% here** — a real, free win to switch to Q4 KV at this
+config (GPU stays ~50% idle either way, so KV is a tax, not the wall).
+
+## 3. The prefetch thesis: mechanism-proven, recall-untested
+
+- **Mechanism (this PR):** `expert_cache` + `plan_prefetch` validated in
+  `tests/bench-expert-stream.cu` on the real x4 link. With the **working-set-
+  protected eviction** (see §4) prefetch streams *exactly* the same bytes as
+  on-demand (443 == 443 H2D copies) and overlaps them with the draft window:
+  **+25–38% per layer-step** when the model doesn't fit. The mechanism works.
+- **Recall (open):** whether a *draft model* predicts the next token's experts
+  at high enough recall to drive that prefetch could **not** be measured here.
+  The available DFlash drafter (`modal-labs/Qwen3.6-35B-A3B-DFlash`, converted to
+  GGUF) loads and `ddtree=ON`, but delivers only ~7% speedup (36 vs 33.5 tok/s)
+  — it ships without the Domino aux heads, and the DDTree's variable accept-
+  length churns the CUDA graph (constant warmup/reset). Speculation never
+  reaches its 2–3×, so draft→expert recall stays unmeasured.
+
+Net: prefetch's *target* is real (§1) and its *mechanism* is proven; the
+*predictor* (draft) is the remaining unknown, gated on a fully-working DFlash
+draft on this build.
+
+## 4. What lands here regardless of the draft
+
+1. **Fused Q4_0/Q4_K MoE kernel** (`tests/test-moe-fused-q4{,k}.cu`):
+   gather + on-the-fly dequant + routed accumulation in one launch, validated to
+   `rel 6.4e-6` vs a CPU oracle on `sm_86`. Q4_K dequant mirrors ggml
+   `get_scale_min_k4` bit-for-bit.
+2. **Working-set-protected LFU eviction** (`expert-cache.cu`): never evict an
+   expert in the current step's working set. Found by *measuring* H2D copy
+   counts (naive prefetch issued +64% copies = thrash); the fix drops it to
+   parity. This is a transferable improvement for **any** bounded expert cache,
+   including Spark's.
+3. **This characterization** — the regime boundary (§1) and the TQ3 tax (§2)
+   are reusable facts for tuning expert offload on consumer hardware.
+
+Policy/residency (`residency_planner`) overlaps Spark's shipped calibrated
+placement + bounded cache; it's a clean, unit-tested reference, not a
+replacement. Cross-platform: suite is 5/5 green on both Windows (MSVC) and the
+Lucebox (native gcc/nvcc), `sm_86`.

ggml/src/ggml-cuda/expert-cache.cu

@@ -0,0 +1,118 @@
+#include "expert-cache.h"
+
+#include <algorithm>
+
+namespace ggml_moe_stream {
+
+expert_cache::expert_cache(int n_experts, size_t bytes_per_expert, int n_slots,
+                           const void * host_experts, cudaStream_t copy_stream,
+                           float decay, float hysteresis)
+    : n_experts_(n_experts),
+      bytes_(bytes_per_expert),
+      n_slots_(std::min(n_slots, n_experts)),
+      host_(static_cast<const char *>(host_experts)),
+      copy_stream_(copy_stream),
+      slot_of_(n_experts, -1),
+      expert_in_(n_slots_, -1),
+      protected_(n_experts, 0),
+      ready_(n_slots_),
+      planner_(n_experts, n_slots_, decay, hysteresis) {
+    cudaMalloc(&dev_, (size_t) n_slots_ * bytes_);
+    for (int s = 0; s < n_slots_; ++s) {
+        cudaEventCreateWithFlags(&ready_[s], cudaEventDisableTiming);
+    }
+}
+
+expert_cache::~expert_cache() {
+    for (auto ev : ready_) cudaEventDestroy(ev);
+    if (dev_) cudaFree(dev_);
+}
+
+std::vector<uint8_t> expert_cache::actual_resident_mask() const {
+    std::vector<uint8_t> m(n_experts_, 0);
+    for (int e = 0; e < n_experts_; ++e) if (slot_of_[e] >= 0) m[e] = 1;
+    return m;
+}
+
+// Choose a slot to (re)use for `incoming`: empty first; else evict the
+// lowest-usage (LFU) expert that is NOT in the current step's working set.
+// Protecting the working set is what stops the thrash — evicting an expert that
+// is needed this very step only to re-stream it (the +60-75% H2D copies we
+// measured). Fall back to plain LFU only if every slot is working-set-pinned
+// (working set larger than the cache, which shouldn't happen in practice).
+int expert_cache::pick_victim_slot(int incoming) {
+    for (int s = 0; s < n_slots_; ++s) if (expert_in_[s] == -1) return s;
+    int best_slot = -1; double best_usage = 0.0;
+    for (int s = 0; s < n_slots_; ++s) {
+        const int e = expert_in_[s];
+        if (e == incoming || protected_[e]) continue;
+        const double u = planner_.usage(e);
+        if (best_slot == -1 || u < best_usage) { best_usage = u; best_slot = s; }
+    }
+    if (best_slot != -1) return best_slot;
+    for (int s = 0; s < n_slots_; ++s) {           // fallback: all slots pinned
+        const int e = expert_in_[s];
+        if (e == incoming) continue;
+        const double u = planner_.usage(e);
+        if (best_slot == -1 || u < best_usage) { best_usage = u; best_slot = s; }
+    }
+    return best_slot;
+}
+
+int expert_cache::bind(int e) {
+    const int s   = pick_victim_slot(e);
+    const int old = expert_in_[s];
+    if (old >= 0) slot_of_[old] = -1;
+    expert_in_[s] = e;
+    slot_of_[e]   = s;
+    return s;
+}
+
+const void * expert_cache::ensure_resident(int e, cudaStream_t compute_stream) {
+    if (slot_of_[e] >= 0) {
+        const int s = slot_of_[e];
+        cudaStreamWaitEvent(compute_stream, ready_[s], 0); // honor in-flight prefetch
+        return dev_ + (size_t) s * bytes_;
+    }
+    const int s = bind(e);
+#ifndef STUB
+    cudaMemcpyAsync(dev_ + (size_t) s * bytes_, host_ + (size_t) e * bytes_,
+                    bytes_, cudaMemcpyHostToDevice, compute_stream);
+    cudaEventRecord(ready_[s], compute_stream);
+    ++h2d_copies;
+#endif
+    return dev_ + (size_t) s * bytes_;
+}
+
+void expert_cache::prefetch(const std::vector<int> & predicted, int max_prefetch) {
+    const std::vector<int> plan = plan_prefetch(predicted, actual_resident_mask(), max_prefetch);
+    for (int e : plan) {
+        const int s = bind(e);
+#ifndef STUB
+        cudaMemcpyAsync(dev_ + (size_t) s * bytes_, host_ + (size_t) e * bytes_,
+                        bytes_, cudaMemcpyHostToDevice, copy_stream_);
+        cudaEventRecord(ready_[s], copy_stream_);
+        ++h2d_copies;
+#endif
+    }
+}
+
+void expert_cache::observe(const int32_t * routed, int n) {
+    // routed is this step's working set: protect it from eviction until next step.
+    std::fill(protected_.begin(), protected_.end(), (uint8_t) 0);
+    for (int i = 0; i < n; ++i) {
+        const int e = routed[i];
+        if (e >= 0 && e < n_experts_) protected_[e] = 1;
+    }
+    planner_.observe(routed, n);
+}
+
+const void * expert_cache::device_ptr(int e) const {
+    return slot_of_[e] >= 0 ? dev_ + (size_t) slot_of_[e] * bytes_ : nullptr;
+}
+
+int expert_cache::resident_count() const {
+    int c = 0; for (int s = 0; s < n_slots_; ++s) if (expert_in_[s] >= 0) ++c; return c;
+}
+
+} // namespace ggml_moe_stream

ggml/src/ggml-cuda/expert-cache.h

@@ -0,0 +1,70 @@
+#pragma once
+
+// CUDA mechanism that executes the expert-streaming policy (P1 prefetch + P3
+// residency) for a memory-constrained MoE: experts live in host/unified memory,
+// a fixed set of VRAM "slots" cache the hot ones, and cold experts are streamed
+// in over the PCIe link on demand or — better — prefetched on a copy stream
+// during the speculative draft pass so the verify pass finds them resident.
+//
+// This is the plumbing under expert-stream.h's pure-policy planners.
+
+#include <cstdint>
+#include <vector>
+#include <cuda_runtime.h>
+
+#include "expert-stream.h"
+
+namespace ggml_moe_stream {
+
+class expert_cache {
+public:
+    // n_experts        : total experts
+    // bytes_per_expert : size of one expert's weight blob
+    // n_slots          : VRAM-resident capacity (== residency budget)
+    // host_experts     : base of all experts, contiguous, in (pinned) host memory
+    // copy_stream      : stream used for speculative prefetch copies
+    expert_cache(int n_experts, size_t bytes_per_expert, int n_slots,
+                 const void * host_experts, cudaStream_t copy_stream,
+                 float decay = 0.99f, float hysteresis = 1.10f);
+    ~expert_cache();
+
+    expert_cache(const expert_cache &) = delete;
+    expert_cache & operator=(const expert_cache &) = delete;
+
+    // Ensure expert e is resident; returns its device pointer. On a miss, streams
+    // host->VRAM on compute_stream (evicting per policy). On a hit, makes
+    // compute_stream wait on any in-flight prefetch for that slot.
+    const void * ensure_resident(int e, cudaStream_t compute_stream);
+
+    // P1: given the draft pass's predicted experts, async-copy the cold ones on
+    // the copy stream (overlapped with draft compute). Non-blocking.
+    void prefetch(const std::vector<int> & predicted, int max_prefetch);
+
+    // P3: feed one decode step's routed experts to the residency policy.
+    void observe(const int32_t * routed, int n);
+
+    // introspection (tests / telemetry)
+    bool         is_resident(int e) const { return slot_of_[e] >= 0; }
+    const void * device_ptr(int e) const;
+    int          resident_count() const;
+    long         h2d_copies = 0;   // count of host->device expert streams issued (telemetry)
+
+private:
+    int                     pick_victim_slot(int incoming);
+    std::vector<uint8_t>    actual_resident_mask() const;
+    int                     bind(int e);                 // choose+claim a slot for e, return slot
+
+    int                      n_experts_;
+    size_t                   bytes_;
+    int                      n_slots_;
+    const char *             host_;
+    cudaStream_t             copy_stream_;
+    char *                   dev_ = nullptr;             // n_slots_ * bytes_
+    std::vector<int>         slot_of_;                   // expert -> slot, or -1
+    std::vector<int>         expert_in_;                 // slot   -> expert, or -1
+    std::vector<uint8_t>     protected_;                 // expert -> 1 if in current step's working set
+    std::vector<cudaEvent_t> ready_;                     // per-slot "copy done" event
+    residency_planner        planner_;
+};
+
+} // namespace ggml_moe_stream

ggml/src/ggml-cuda/expert-stream-hooks.cu

@@ -0,0 +1,48 @@
+#include "expert-stream-hooks.h"
+#include "expert-cache.h"
+
+#include <unordered_map>
+
+namespace ggml_moe_stream {
+
+static std::unordered_map<int, controller> & registry() {
+    static std::unordered_map<int, controller> r;
+    return r;
+}
+
+controller * moe_get(int layer) {
+    auto & r = registry();
+    auto it = r.find(layer);
+    return it == r.end() ? nullptr : &it->second;
+}
+
+void moe_register(int layer, int n_experts, size_t bytes_per_expert, int n_slots,
+                  const void * host_experts, cudaStream_t copy_stream,
+                  float decay, float hysteresis) {
+    auto & r = registry();
+    controller & c = r[layer];
+    delete c.cache;
+    c.cache = new expert_cache(n_experts, bytes_per_expert, n_slots,
+                               host_experts, copy_stream, decay, hysteresis);
+}
+
+void moe_reset() {
+    auto & r = registry();
+    for (auto & kv : r) delete kv.second.cache;
+    r.clear();
+}
+
+void moe_observe(int layer, const int32_t * routed, int n) {
+    if (controller * c = moe_get(layer)) c->cache->observe(routed, n);
+}
+
+void moe_prefetch(int layer, const std::vector<int> & draft_experts, int cap) {
+    if (controller * c = moe_get(layer)) c->cache->prefetch(draft_experts, cap);
+}
+
+const void * moe_expert_ptr(int layer, int e, cudaStream_t stream) {
+    controller * c = moe_get(layer);
+    return c ? c->cache->ensure_resident(e, stream) : nullptr;
+}
+
+} // namespace ggml_moe_stream