fix(dflash): unblock Qwen3.5-MoE + DFlash + ContextPilot serving hang (multi-GPU device placement) - #147
Open
drunkcoding wants to merge 5 commits into
Open
fix(dflash): unblock Qwen3.5-MoE + DFlash + ContextPilot serving hang (multi-GPU device placement)#147drunkcoding wants to merge 5 commits into
drunkcoding wants to merge 5 commits into
Conversation
added 3 commits
August 11, 2026 14:31
Serving Qwen3.5-MoE with a DFlash drafter on more than one GPU hung
during generation. The offloaded text backbone (embed_tokens/attention/
lm_head) is resident on the LAST visible GPU (mirrors
ContinuousBatchingEngine._resolve_device), but the DFlash native path
assumed cuda:0:
- MoE._native_model_forward{,_rich} placed inputs on cuda:0, so the very
first target forward's embed_tokens(input_ids) raised a device
mismatch (index on cuda:0 vs weight on cuda:N).
- DFlashSpeculator._infer_cuda_device returned the first cuda parameter
(or cuda:0), so the drafter and its block landed on a different GPU
than the bound shared embed_tokens/lm_head.
Both RuntimeErrors propagated out of engine.step() and were swallowed by
the async engine loop, so the client saw an indefinite hang (HTTP 000)
instead of an error.
Resolve the native input device from the resident input embedding
(fallback cuda:{last_gpu}); infer the speculator device from the bound
shared embedding. Verified end-to-end on 2 GPUs: Qwen3.5-35B-A3B target +
z-lab/Qwen3.5-35B-A3B-DFlash drafter + ContextPilot generate over
api_server_v2. dflash (264) / qwen3_5_moe (16) / contextpilot (75)
suites stay green.
Refs: #146
Qwen3.5-MoE ships a vision-language config that nests the text backbone dimensions under text_config, so _build_engine_config raised "RuntimeError: unable to resolve model num_layers". Resolve num_layers/heads/kv_heads/head_dim/eos via config.get_text_config(), which returns self for text-only configs so existing checkpoints are unaffected. Verified: num_layers=40, num_kv_heads=2, head_dim=256, eos=248044. Refs: #146
Mirror test_tokenizer_compat's importorskip: the CP import test and the live-middleware tests now skip (instead of failing) when the optional contextpilot package is not installed, so the suite is green with or without it.
added 2 commits
August 11, 2026 18:20
…to_experts transformers 5.15 removed GlmMoeDsaMoE.route_tokens_to_experts, which glm_moe_dsa.py and test_glm_routing.py rely on. CI installs the latest 5.x (requirements pin transformers>=5.3.0,<6), so unit-tests (3.10)/(3.12) fail on dev (and every open PR) with "AttributeError: type object 'GlmMoeDsaMoE' has no attribute route_tokens_to_experts". Skip the GLM routing module when the method is absent (mirrors the existing importorskip guard), and resolve it via getattr so the block raises a clear, actionable error at routing time instead of a cryptic AttributeError at construction. Refs: #146
_initialize_model passes speculative_draft to ContinuousBatchingEngine (since the DFlash integration), but the watchdog test's _FakeRuntimeEngine mock never accepted it. On transformers 5.15, once the earlier GLM failure stops masking it (pytest runs fail-fast), test_watchdog_integration fails with "unexpected keyword argument 'speculative_draft'". Match the real ContinuousBatchingEngine signature (same fix already in #148). Refs: #146
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
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.
Summary
Fixes the primary open bug in #146: serving the full trifecta —
Qwen3.5-MoE target + DFlash speculative drafter + ContextPilot — via
api_server_v2reachedHEALTHYbut token generation hung (client gotHTTP 000, process wedged in the CUDA driver).Root cause: two independent multi-GPU device-placement bugs in the DFlash
native path. MoE-Infinity keeps the offloaded model's resident backbone
(
embed_tokens/ attention /lm_head) on the last visible GPU (this isexactly what
ContinuousBatchingEngine._resolve_devicedoes), but the DFlashspeculative path assumed
cuda:0:MoE._native_model_forward/_native_model_forward_richplaced theinput on
cuda:0, so the first target forward'sembed_tokens(input_ids)raised
RuntimeError: Expected all tensors to be on the same device (index cuda:0 vs weight cuda:1).DFlashSpeculator._infer_cuda_devicereturned the first CUDA parameter(or
cuda:0), so the drafter and itsblocklanded on a different GPU thanthe bound shared
embed_tokens/lm_head, failing in_run_drafter'sself.embed_tokens(block).Both
RuntimeErrors propagated out ofengine.step()and were silentlyswallowed by the async engine loop, so the request hung indefinitely instead
of returning an error. (This silent-hang-on-step-exception is the mechanism
behind #146 Bug 2 and is why the failure looked like a CUDA wedge rather than a
500.)
Changes
fix(dflash)— resolve the native-forward input device from the residentinput embedding (fallback
cuda:{last_gpu}), and infer the speculator devicefrom the bound shared embedding. Single-GPU / CPU behavior is unchanged
(
last_gpu == 0; CPU path preserved).fix(serving)— resolve Qwen3.5-MoE architecture dims viaconfig.get_text_config()(its VL config nests text dims undertext_config),fixing
RuntimeError: unable to resolve model num_layers. Text-only configsare unaffected (
get_text_config()returnsself).test(contextpilot)—importorskip/skipifguards so the CP suite isgreen with or without the optional
contextpilotpackage (mirrorstest_tokenizer_compat).Verification
Live 2×GPU run (
CUDA_VISIBLE_DEVICES=0,1) viaapi_server_v2with--speculative-draft z-lab/Qwen3.5-35B-A3B-DFlash --enable-contextpilot:HTTP 000after 60s (hang). After: generates end-to-end.clean code.
Regression suites (unchanged from baseline):
tests/python/dflash— 264 passed, 3 skipped (GPU-gated harnesses)tests/python/unit/test_qwen3_5_moe.py— 16 passedtests/python/contextpilot— 75 passed, 1 skippedNotes / out of scope
(
/v1/completionsContextPilot single-promptlist index out of range, caughtby fallback) are not addressed here; this PR fixes the generation hang
(Bug 1). Surfacing engine-step exceptions as request errors is a sensible
follow-up.
Refs: #146