feat(moe): add--freq and --collect-routing, train predictor of experts (#386)

* feat: port --freq and --collect-routing from StreamMoE

Add expert frequency tracking and binary routing data collection for
training MLP expert predictors.

New features:
- --freq: prints per-layer expert frequency analysis at shutdown
- --collect-routing <path>: writes binary routing data (hidden states +
  expert selections) for offline predictor training

New files:
- server/src/common/moe_routing_collector.{h,cpp}: thread-safe binary
  writer for per-token routing decisions
- scripts/train_predictor.py: MLP/linear predictor training script
  (ported from StreamMoE, supports CUDA, mlp/linear/cross variants)

Binary format per sample:
  [int32 layer_idx, int32 K, float32[n_embd] hidden, int32[K] experts]

Hooks added in qwen35moe pipelined decode (cached + dynamic paths),
qwen35moe prefill, laguna hybrid decode, and laguna hybrid prefill.

Tested: 1K tokens → 38.5% top-4 hit rate (mlp-cross), 2x temporal baseline.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(predictor): ruff lint + cubic P1/P2 review fixes

CI lint (ruff):
- sort imports (I001), drop unused math / torch.nn (F401)

P1 (correctness):
- pad expert indices with -1, not 0 (0 is a valid expert id, so 0
  padding corrupts labels when K varies). Filter negatives in
  build_target_multilabel + every eval path, and count real experts
  (not padded K) in temporal/top-k/per-layer hit-rate denominators.
- guard empty train/test split (exit cleanly instead of crashing).
- MoeRoutingCollector::record now checks every fwrite; a short write
  closes the file and stops, instead of emitting misaligned/corrupt
  training data with an inflated sample count.

P2:
- int64_t sample counter (+ %lld).
- moe_hybrid_routing_stats: cap top-10/30/60 to total when
  n_expert < bucket size (was printing 0%).
- --freq/--collect-routing force routing-stats env vars when unset or
  empty (overwrite=0 left an empty value → silent no-op).
- set_routing_collector returns bool (capability check); caller warns +
  closes collector when the backend can't collect; lifetime contract
  documented.
- dedup the per-layer ffn_post D2H copy across all three decode paths
  when --collect-routing and cold compute are both active.
- correct --hidden-dim/--num-experts/--num-layers help defaults.
- validate --threshold is in (0,1); handle empty/truncated routing files.

Co-Authored-By: WOZCODE <contact@withwoz.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: mrciffa <49000955+davide221@users.noreply.github.com>
Co-authored-by: WOZCODE <contact@withwoz.com>

Author: 4 people

scripts/train_predictor.py

@@ -265,6 +265,7 @@ add_library(dflash_common STATIC
     src/common/moe_hybrid_stream.cpp
     src/common/cold_ffn_cpu.cpp
     src/common/moe_hybrid_swap_manager.cpp
+    src/common/moe_routing_collector.cpp
     src/qwen35/layer_split_forward.cpp
     src/qwen35/qwen35_target_shard_ipc.cpp
     src/qwen35/qwen35_target_shard_ipc_daemon.cpp

server/CMakeLists.txt

@@ -323,6 +323,25 @@ struct ModelBackend {
     // the server before startup.
     virtual bool supports_remote_draft() const { return false; }
 
+    // ── Routing data collection ──────────────────────────────────────
+    // Set an external routing collector that the backend will call for each
+    // token/layer during decode (hidden state + expert IDs). Used by
+    // --collect-routing for predictor training data.
+    //
+    // Lifetime: the collector pointer is borrowed, not owned. The caller must
+    // keep it alive until set_routing_collector(nullptr) is called (or the
+    // backend is destroyed), and must not pass a collector to a backend that
+    // decodes on another thread without outliving that decode.
+    //
+    // Returns true if the backend supports routing collection (MoE backends).
+    // The default returns false so the server can detect unsupported backends
+    // and warn instead of silently collecting nothing.
+    virtual bool set_routing_collector(class MoeRoutingCollector *) { return false; }
+
+    // Get the current routing stats (if tracked). Returns nullptr if the
+    // backend does not support routing stats or they are not enabled.
+    virtual const struct MoeHybridRoutingStats * get_routing_stats() const { return nullptr; }
+
     // ── Cleanup ──────────────────────────────────────────────────────
     // Release all resources (weights, cache, snapshots, drafter).
     // Called by run_daemon() before returning.

server/src/common/model_backend.h

@@ -3,6 +3,7 @@
 #include <algorithm>
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
 #include <fstream>
 #include <numeric>
 
@@ -118,6 +119,75 @@ std::vector<int> MoeHybridRoutingStats::hot_experts(int layer_idx, int hot_count
     return ranked;
 }
 
+void MoeHybridRoutingStats::print_freq_analysis() const {
+    if (n_layer <= 0 || n_expert <= 0 || counts.empty()) return;
+
+    uint64_t total_all = 0;
+    for (int il = 0; il < n_layer; ++il) total_all += layer_totals[(size_t)il];
+    if (total_all == 0) return;
+
+    // Estimate total tokens: layer_totals[0] / n_expert_used
+    uint64_t total_tokens = (n_expert_used > 0 && layer_totals[0] > 0)
+        ? layer_totals[0] / (uint64_t)n_expert_used : 0;
+    double avg_k = (n_layer > 0 && total_tokens > 0)
+        ? (double)total_all / ((double)n_layer * (double)total_tokens) : 0.0;
+
+    std::fprintf(stderr, "\n=== Expert Frequency Analysis ===\n");
+    std::fprintf(stderr, "Tokens tracked: %llu, avg-K=%.2f\n\n",
+                 (unsigned long long)total_tokens, avg_k);
+
+    int experts_for_80_total = 0;
+    std::vector<uint64_t> sorted((size_t)n_expert);
+
+    for (int il = 0; il < n_layer; ++il) {
+        uint64_t tact = layer_totals[(size_t)il];
+        if (tact == 0) {
+            std::fprintf(stderr, "Layer %2d:   0 unique\n", il);
+            continue;
+        }
+        for (int ie = 0; ie < n_expert; ++ie)
+            sorted[(size_t)ie] = counts[index_of(il, ie)];
+        std::sort(sorted.begin(), sorted.end(), std::greater<uint64_t>());
+
+        int unique = 0;
+        for (int ie = 0; ie < n_expert; ++ie) if (sorted[(size_t)ie] > 0) unique++;
+
+        uint64_t cum = 0;
+        uint64_t top10 = 0, top30 = 0, top60 = 0;
+        int n50 = 0, n80 = 0, n90 = 0;
+        for (int ie = 0; ie < n_expert; ++ie) {
+            cum += sorted[(size_t)ie];
+            if (ie == 9)  top10 = cum;
+            if (ie == 29) top30 = cum;
+            if (ie == 59) top60 = cum;
+            if (!n50 && cum * 100 >= tact * 50) n50 = ie + 1;
+            if (!n80 && cum * 100 >= tact * 80) n80 = ie + 1;
+            if (!n90 && cum * 100 >= tact * 90) n90 = ie + 1;
+        }
+        // If there are fewer experts than a bucket size, the ie==N-1 check
+        // never fires and topN stays 0, printing a bogus 0%. cum == tact here
+        // (all experts summed), so the bucket covers everything → 100%.
+        if (n_expert < 10) top10 = cum;
+        if (n_expert < 30) top30 = cum;
+        if (n_expert < 60) top60 = cum;
+        std::fprintf(stderr, "Layer %2d: %3d unique, top-10 %.0f%%, top-30 %.0f%%, top-60 %.0f%% "
+                             "(50%%@%d, 80%%@%d, 90%%@%d)\n",
+                     il, unique,
+                     100.0 * (double)top10 / (double)tact,
+                     100.0 * (double)top30 / (double)tact,
+                     100.0 * (double)top60 / (double)tact,
+                     n50, n80, n90);
+        experts_for_80_total += n80;
+    }
+
+    double avg80 = (n_layer > 0) ? (double)experts_for_80_total / n_layer : 0.0;
+    std::fprintf(stderr, "\n--- Overall Summary ---\n");
+    std::fprintf(stderr, "80%% hit rate: %d experts pinned (avg %.0f/layer)\n",
+                 experts_for_80_total, avg80);
+    std::fprintf(stderr, "Model: %d layers x %d experts, top-%d routing\n",
+                 n_layer, n_expert, n_expert_used);
+}
+
 bool MoeHybridRoutingStats::save_csv(const std::string & path, std::string * err) const {
     if (n_layer <= 0 || n_expert <= 0 || counts.size() != (size_t)n_layer * (size_t)n_expert) {
         if (err) *err = "routing stats not initialized";

server/src/common/moe_hybrid_routing_stats.cpp

@@ -38,6 +38,9 @@ struct MoeHybridRoutingStats {
     std::vector<int> ranked_experts(int layer_idx) const;
     std::vector<int> hot_experts(int layer_idx, int hot_count) const;
 
+    // Print StreamMoE-style frequency analysis to stderr.
+    void print_freq_analysis() const;
+
     bool save_csv(const std::string & path, std::string * err = nullptr) const;
     static bool load_csv(const std::string & path,
                          MoeHybridRoutingStats & out,

server/src/common/moe_hybrid_routing_stats.h

@@ -0,0 +1,59 @@
+#include "moe_routing_collector.h"
+
+#include <cerrno>
+#include <cstdio>
+#include <cstring>
+
+namespace dflash::common {
+
+bool MoeRoutingCollector::open(const std::string & path) {
+    std::lock_guard<std::mutex> lock(mu_);
+    if (fd_) {
+        std::fclose(fd_);
+        fd_ = nullptr;
+    }
+    fd_ = std::fopen(path.c_str(), "wb");
+    if (!fd_) {
+        std::fprintf(stderr, "[routing-collector] failed to open '%s': %s\n",
+                     path.c_str(), std::strerror(errno));
+        return false;
+    }
+    samples_ = 0;
+    std::fprintf(stderr, "[routing-collector] collecting routing data to '%s'\n", path.c_str());
+    return true;
+}
+
+void MoeRoutingCollector::record(int layer_idx, const float * hidden, int n_embd,
+                                 const int32_t * expert_ids, int K) {
+    std::lock_guard<std::mutex> lock(mu_);
+    if (!fd_) return;
+
+    int32_t li = layer_idx;
+    int32_t ki = K;
+    // A short write desyncs the fixed-layout binary stream: every subsequent
+    // record would be misaligned and the file silently unparseable. Treat any
+    // partial write as fatal for this file — stop writing and don't count the
+    // sample, rather than emitting corrupt training data with an inflated count.
+    if (std::fwrite(&li, sizeof(int32_t), 1, fd_) != 1 ||
+        std::fwrite(&ki, sizeof(int32_t), 1, fd_) != 1 ||
+        std::fwrite(hidden, sizeof(float), (size_t)n_embd, fd_) != (size_t)n_embd ||
+        std::fwrite(expert_ids, sizeof(int32_t), (size_t)K, fd_) != (size_t)K) {
+        std::fprintf(stderr, "[routing-collector] write error after %lld samples: %s\n",
+                     (long long)samples_, std::strerror(errno));
+        std::fclose(fd_);
+        fd_ = nullptr;
+        return;
+    }
+    samples_++;
+}
+
+void MoeRoutingCollector::close() {
+    std::lock_guard<std::mutex> lock(mu_);
+    if (!fd_) return;
+    std::fclose(fd_);
+    fd_ = nullptr;
+    std::fprintf(stderr, "[routing-collector] closed, %lld samples written\n",
+                 (long long)samples_);
+}
+
+}  // namespace dflash::common

server/src/common/moe_routing_collector.cpp

@@ -0,0 +1,50 @@
+// MoE routing data collector — writes per-token binary routing data for
+// predictor training (gate input hidden states + expert selections).
+//
+// Binary format per sample (matches StreamMoE train_predictor.py format):
+//   int32  layer_idx
+//   int32  K  (number of selected experts)
+//   float32[n_embd]  hidden_state (gate input / post-norm hidden for this layer)
+//   int32[K]  expert_indices (actual routing result)
+//
+// Enable via --collect-routing <path> on the server CLI or via env var
+// DFLASH_COLLECT_ROUTING=<path>.
+
+#pragma once
+
+#include <cstdint>
+#include <cstdio>
+#include <mutex>
+#include <string>
+
+namespace dflash::common {
+
+class MoeRoutingCollector {
+public:
+    MoeRoutingCollector() = default;
+    ~MoeRoutingCollector() { close(); }
+
+    MoeRoutingCollector(const MoeRoutingCollector &) = delete;
+    MoeRoutingCollector & operator=(const MoeRoutingCollector &) = delete;
+
+    // Open the output file. Returns false on failure.
+    bool open(const std::string & path);
+
+    // Record one routing sample: gate input + expert selection for one token at one layer.
+    // hidden must point to n_embd floats; expert_ids to K int32s.
+    void record(int layer_idx, const float * hidden, int n_embd,
+                const int32_t * expert_ids, int K);
+
+    // Close and flush the output file. Prints summary to stderr.
+    void close();
+
+    bool    is_open() const { return fd_ != nullptr; }
+    int64_t sample_count() const { return samples_; }
+
+private:
+    std::FILE * fd_ = nullptr;
+    int64_t samples_ = 0;  // wide: collection runs can exceed 2^31 samples
+    std::mutex mu_;
+};
+
+}  // namespace dflash::common

server/src/common/moe_routing_collector.h

@@ -19,6 +19,7 @@
 #include "../common/moe_hybrid_storage.h"
 #include "../common/moe_hybrid_routing_stats.h"
 #include "../common/moe_hybrid_swap_manager.h"
+#include "../common/moe_routing_collector.h"
 #include "common/step_graph.h"
 
 #include "ggml-cuda.h"
@@ -1426,6 +1427,15 @@ bool LagunaBackend::hybrid_forward_one_token(int32_t tok, int kv_pos,
                 }
             }
 
+            // Collect routing data for predictor training
+            if (routing_collector_) {
+                std::vector<float> gate_input((size_t)w_.n_embd);
+                ggml_backend_tensor_get(layer_sg.ffn_post, gate_input.data(), 0,
+                                         sizeof(float) * (size_t)w_.n_embd);
+                routing_collector_->record(il, gate_input.data(), w_.n_embd,
+                                           selected.data(), (int)selected.size());
+            }
+
             // Hybrid FFN: hot on GPU, cold on CPU, combine on GPU
             auto & storage = moe_hybrid_->layers[(size_t)il];
             { int _lc=0; for (int _k=0;_k<(int)selected.size();++_k){ int _g=selected[(size_t)_k];
@@ -1661,6 +1671,18 @@ GenerateResult LagunaBackend::generate_hybrid(const GenerateRequest & req,
                     }
                 }
 
+                // Collect routing data for predictor training (prefill path)
+                if (routing_collector_) {
+                    for (int i = 0; i < chunk_len; ++i) {
+                        routing_collector_->record(
+                            il,
+                            chunk_post.data() + (size_t)i * (size_t)w_.n_embd,
+                            w_.n_embd,
+                            chunk_selected.data() + (size_t)i * (size_t)n_expert_used,
+                            n_expert_used);
+                    }
+                }
+
                 // Batched hybrid FFN evaluation
                 auto & storage = moe_hybrid_->layers[(size_t)il];
                 MoeHybridConfig chunk_cfg = make_moe_hybrid_config(w_);

server/src/laguna/laguna_backend.cpp

@@ -17,6 +17,7 @@
 #include "../common/moe_hybrid_routing_stats.h"
 #include "../common/moe_hybrid_swap_manager.h"
 #include "../common/moe_hybrid_stream.h"
+#include "../common/moe_routing_collector.h"
 
 #include "ggml.h"
 #include "ggml-backend.h"
@@ -73,6 +74,9 @@ class LagunaBackend : public ModelBackend {
 
     void shutdown() override;
 
+    bool set_routing_collector(MoeRoutingCollector * c) override { routing_collector_ = c; return true; }
+    const MoeHybridRoutingStats * get_routing_stats() const override { return routing_stats_.get(); }
+
 private:
     LagunaBackendArgs                           args_;
     ggml_backend_t                              backend_   = nullptr;
@@ -98,6 +102,7 @@ class LagunaBackend : public ModelBackend {
     MoeHybridSwapPolicy                        swap_policy_;
     bool                                       hybrid_telemetry_ = false;
     MoeHybridStreamEngine                      stream_engine_;
+    MoeRoutingCollector *                       routing_collector_ = nullptr;
 
     bool ensure_slot(int slot);
 

server/src/laguna/laguna_backend.h

@@ -394,6 +394,7 @@ bool Qwen35MoeBackend::ensure_pipe_state(int kv_start) {
         pipe_state_.reset();
         return false;
     }
+    pipe_state_->routing_collector = routing_collector_;
     return true;
 }
 
@@ -867,6 +868,18 @@ GenerateResult Qwen35MoeBackend::generate_impl(const GenerateRequest & req,
                 }
             }
 
+            // Collect routing data for predictor training (prefill path)
+            if (routing_collector_) {
+                for (int i = 0; i < chunk_len; ++i) {
+                    routing_collector_->record(
+                        il,
+                        chunk_post.data() + (size_t)i * (size_t)hidden,
+                        hidden,
+                        chunk_selected.data() + (size_t)i * (size_t)n_expert_used,
+                        n_expert_used);
+                }
+            }
+
             // Hybrid FFN — skip batched path when cold experts exist (CUDA mul_mat_id bug on sm_75)
             MoeHybridConfig chunk_cfg = make_moe_hybrid_config(target_weights());
             MoeLayerDesc chunk_desc = make_moe_layer_desc(target_weights().layers[(size_t)il]);