feat(moe): add--freq and --collect-routing, train predictor of experts#386
Merged
Conversation
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>
howard0su
changed the title
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
14 issues found across 15 files
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
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>
Contributor
There was a problem hiding this comment.
1 issue found across 9 files (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="server/src/common/moe_routing_collector.h">
<violation number="1" location="server/src/common/moe_routing_collector.h:42">
P2: `sample_count()` reads `samples_` without synchronization while other threads mutate it under `mu_`, causing a data race.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| void close(); | ||
|
|
||
| bool is_open() const { return fd_ != nullptr; } | ||
| int64_t sample_count() const { return samples_; } |
Contributor
There was a problem hiding this comment.
P2: sample_count() reads samples_ without synchronization while other threads mutate it under mu_, causing a data race.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At server/src/common/moe_routing_collector.h, line 42:
<comment>`sample_count()` reads `samples_` without synchronization while other threads mutate it under `mu_`, causing a data race.</comment>
<file context>
@@ -38,12 +38,12 @@ class MoeRoutingCollector {
- bool is_open() const { return fd_ != nullptr; }
- int sample_count() const { return samples_; }
+ bool is_open() const { return fd_ != nullptr; }
+ int64_t sample_count() const { return samples_; }
private:
</file context>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add expert frequency tracking and binary routing data collection for training MLP expert predictors.
New features:
New files:
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 , Qwen3.6-35B-A3B
the mlp-cross predictor achieves 2x the baseline with only 1K token of training data.