Fix: TP+EP produces silently corrupted outputs (2 bugs in qwen3_5.py load_weights)

#8
by juanginer - opened

Hi Escha team,

Running Qwen3.6-35B-A3B-Escha-W2 on 2×RTX 3090 with --tp-size 2 --ep-size 2 produced silently corrupted outputs (arithmetic drift, degenerate loops with non-English fragments, <think> blocks that never close). --dp-size 2 was fine, and DP=1 baseline was fine.

My motivation for chasing this: I care more about total KV cache pool than raw throughput. DP=N replicates the model per rank and carves KV pool out of the remaining VRAM per replica — great for concurrency, but the max context I can serve per request is bound by one GPU's leftover. TP+EP distributes the model across ranks so ~half the VRAM per GPU is freed for KV, and the pool from all ranks is effectively shared for long-context requests. On 2×3090 this went from ~340 K tokens per DP replica to 1.36 M tokens aggregated under TP=2+EP=2 — decisive for agentic / long-context workloads.

Bisecting the runtime narrowed the silent corruption to two bugs in sglang/srt/models/qwen3_5.py (the sglang fork inside escha-1.0.2+qwen3moe.whl):

  1. eschamoe checkpoint loader ignores the EP rank offset. The manual eschamoe branch in load_weights falls into a shape-mismatch else that was written for mixed-bit K3-vs-K2 last-dim cases and does _qp[_pn].data = _w.to(...) without slicing on dim 0. Under ep_size>1 both ranks end up with experts [0:num_local); experts [num_local:num_experts) are never loaded on any rank. My mask (topk_w=0 for topk_ids<0) turned the missing contributions into silent zeros, so 4/8 top-k slots per token contributed nothing — visible as consistent ~5-15 % numerical drift.

  2. SGLANG_INT8_EMBED install ignores TP sharding. The embed_tokens int8 branch registers the full-vocab checkpoint tensor as a buffer directly into a VocabParallelEmbedding whose .weight was preallocated as [vocab/tp, hidden]. Every rank then indexes a full-vocab table with local row indices → garbage embeddings from step 0 → degenerate output.

Both are drop-in fixable in Python without touching the compiled kernels. I published the patch, a bisect log, and full-runtime reproduction steps at:

https://github.com/ginerJuanUdesa/escha-tp-fix

Verified on 2×3090: post-fix TP=2+EP=2 matches DP=1 greedy tokens; ~5× larger total KV pool at ~10 % single-request latency cost (all-reduce over PCIe, no NVLink here). Full stack works: attn/GDN/shared-expert/embed/lm_head TP-sharded, routed MoE EP-sharded.

Happy to open a PR against a public source tree if one exists — I could only find the precompiled wheel.

Thanks for the runtime — this is by far the fastest way to serve the 35B on consumer 3090s.

Escha Labs org

Thanks for this! Confirmed both bugs in the shipped source.

One warning for anyone applying the patch: it isn't sufficient on a stock install. Under ep_size>1 sglang rewrites topk_ids to -1 for experts on the other rank, and our MoE kernel corrupts its routing table on those. Your topk_w=0 for topk_ids<0 mask is what keeps your build out of that, and it isn't in the diff. Worth adding to the repo.

We also found pure TP (ep_size=1) is silently wrong the same way, so we'll add a guard rather than let it run.

On the PR offer: the sglang runtime isn't open sourced, so there's no tree to send it to, sorry. We're folding your fix into the next wheel along with the -1 remap and the guard.

We now only have single GPU setups, so if you'd run a release candidate on your 2x3090 against your DP=1 baseline, that's the gate for shipping. We will credit your findings in our release note either way. Thanks again!

Amazing!! Thanks for the quick and friendly response! I am currently using this model as my local daily driver so I am happy to help to make it better in any way I can! Let me know if u will make versions of 3.6-27 and 3.8-27, would love to test and add support for both if u need someone with 2 GPUs to do it.

Escha Labs org

We are working hard on 3.8-27B release. Also we have decided to opensource the whole runtime stack soon (because to make it the best, we really need the whole community). So hopefully by next week your stuff will be just merged in as PR. Stay tuned!

Amazing sounds like a plan! Good luck!

Sign up or log in to comment