Skip to content

fix(dflash): unblock Qwen3.5-MoE + DFlash + ContextPilot serving hang (multi-GPU device placement) - #147

Open
drunkcoding wants to merge 5 commits into
devfrom
fix/qwen35-dflash-multigpu-device-hang
Open

fix(dflash): unblock Qwen3.5-MoE + DFlash + ContextPilot serving hang (multi-GPU device placement)#147
drunkcoding wants to merge 5 commits into
devfrom
fix/qwen35-dflash-multigpu-device-hang

Conversation

@drunkcoding

Copy link
Copy Markdown
Contributor

Summary

Fixes the primary open bug in #146: serving the full trifecta —
Qwen3.5-MoE target + DFlash speculative drafter + ContextPilot — via
api_server_v2 reached HEALTHY but token generation hung (client got
HTTP 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 is
exactly what ContinuousBatchingEngine._resolve_device does), but the DFlash
speculative path assumed cuda:0:

  1. MoE._native_model_forward / _native_model_forward_rich placed the
    input on cuda:0, so the first target forward's embed_tokens(input_ids)
    raised RuntimeError: Expected all tensors to be on the same device (index cuda:0 vs weight cuda:1).
  2. 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, failing in _run_drafter's
    self.embed_tokens(block).

Both RuntimeErrors propagated out of engine.step() and were silently
swallowed 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 resident
    input embedding (fallback cuda:{last_gpu}), and infer the speculator device
    from the bound shared embedding. Single-GPU / CPU behavior is unchanged
    (last_gpu == 0; CPU path preserved).
  • fix(serving) — resolve Qwen3.5-MoE architecture dims via
    config.get_text_config() (its VL config nests text dims under text_config),
    fixing RuntimeError: unable to resolve model num_layers. Text-only configs
    are unaffected (get_text_config() returns self).
  • test(contextpilot)importorskip/skipif guards so the CP suite is
    green with or without the optional contextpilot package (mirrors
    test_tokenizer_compat).

Verification

Live 2×GPU run (CUDA_VISIBLE_DEVICES=0,1) via api_server_v2 with
--speculative-draft z-lab/Qwen3.5-35B-A3B-DFlash --enable-contextpilot:

POST /v1/completions {"prompt":"The capital of France is","max_tokens":16,"temperature":0}
-> HTTP 200 in ~7s: " Paris.\nThe capital of France is Paris. ..."  (16 tokens, 14 verify steps)
  • Before: HTTP 000 after 60s (hang). After: generates end-to-end.
  • Bisected spec-off (works) vs spec-on (hung) to localize; confirmed the fix on
    clean code.

Regression suites (unchanged from baseline):

  • tests/python/dflash264 passed, 3 skipped (GPU-gated harnesses)
  • tests/python/unit/test_qwen3_5_moe.py16 passed
  • tests/python/contextpilot75 passed, 1 skipped

Notes / out of scope

Refs: #146

drunkcoding 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.
drunkcoding 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant