fix(models): compute shared experts before in-place routed experts - #463
Conversation
06db1f2 to
1cc51e3
Compare
|
Confirmed: |
How Can I. contact you |
|
email?
|
yeap |
|
you can find it in my GitHub profile |
…lashML-org#463) (cherry picked from commit e0886cc)
… stack Upstream FlashML-org#454 serves images on the Qwen families: requests carry MMItems whose placeholder rows hold content pad ids (so the radix cache keys them by content), the tokenizer precomputes 3-axis rope positions, and the engine encodes on the GPU from pinned host banks right before the LM forward. Kai had its own image path (a CPU vision tower in the tokenizer worker, per-request cos/sin tables, image prompts kept out of the prefix cache and scheduled alone). The two cannot coexist -- they define Req, Batch, UserMsg and the rope differently -- so this merge takes upstream's and removes Kai's: vision_cpu.py, mrope.py, tokenizer/mm_host.py, tokenizer/qwen_vl_lite.py, the scheduler's _encode_multimodal and rope tables, the data:-only image_url renderer, and their tests. What Kai keeps is adapted to the new shape: - --spec-mtp: the verify-window graph, the draft-head chain graph and the eager draft and check-step batches carry [3, n] rope positions (logical + the request's mrope_delta) on an mrope model; the draft head embeds the placeholder token where an image row's successor is a content pad id past the vocab. - --prefill-mixer-pieces: a piece takes its columns of the 3-axis positions; an image chunk splits like a text one (its soft tokens are in the stream before the pieces run). - --prefill-chunk-budget: the transient probe feeds 3-axis positions, like upstream's warmup. - PLE on disk: the hash windows go through upstream's placeholder restore, the verify window's drafts included. - The pipeline window, --dense-quant and --host-embedding config paths sit on upstream's active_encoders / hf_config stripping; the loader passes include_mtp and include_vision. - A chunk checkpoint and --prefix-disk-cache no longer exclude image requests: their ids now name the image. FlashML-org#462 fixes the W4A16_NVFP4 input_scale wait Kai had fixed in e4ffedf, more completely (it also infers W4A16 from config_groups and skips a stored input_scale the scheme does not declare); upstream's modelopt.py is taken as is and Kai's duplicate test is dropped. FlashML-org#463 does not touch a model Kai runs (qwen3_5_moe and qwen4_exp already compute the shared expert first). The merge base is upstream 9535656: the history rewrite of 2026-09-13 dropped GitHub's signature from the previous sync's upstream commits, so git would otherwise take fb7f732. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
BF16 prefill and resident decode overwrite
hidden_stateswith the routed expert output. GLM-4 MoE, GLM MoE DSA, GLM5 Next and MiniMax-M3 currently evaluate their shared experts afterwards, producingR(x) + S(R(x))instead ofR(x) + S(x).Compute each shared branch before
routed_forward, preserving the kernel's existing memory reuse, and document the input mutation contract on both routed entry points.Local validation (the regression script is retained outside this PR):
main(08d728d) and 24 passed with this change: 16 CPU cases with mutating/non-mutating expert implementations, and 8 real BF16 GPU cases covering resident prefill and decode against a PyTorch routed-expert reference.git diff --checkpasses.Exact local test command, run from the checkout's
python/directory: