Skip to content

[Bugfix][Diffusion] Detach the idle MoE LoRA context on Ascend - #9512

Closed
tzhouam wants to merge 1 commit into
release/v0.32.0from
fix-moe-lora-idle-context-v0.32.0
Closed

tzhouam wants to merge 1 commit into
release/v0.32.0from
fix-moe-lora-idle-context-v0.32.0

Conversation

@tzhouam

@tzhouam tzhouam commented Oct 28, 2026

Copy link
Copy Markdown
Collaborator

Purpose

Fixes the HunyuanImage-3 NPU performance regression introduced by #9443 (reached release/v0.32.0 through batch-7 backport #9470) without reverting it, so #9232 stays fixed.

#9443 wraps MoE runners eagerly and publishes a no-adapter (slot-0) mapping before the first forward, so LoRA-typed MoE kernels have a context on HunyuanImage-3's baseline generate(). On GPU this is free: punica sets no_lora_flag for slot 0 and the TRTLLM / Triton LoRA experts return early.

Ascend's unquantized MoE path instead branches on whether a context is published at all, not on whether any token selects an adapter:

  • vllm_ascend/ops/fused_moe/routed_experts.py runs the argsort-based routing recovery plus moe_lora_apply_w13 / moe_lora_apply_w2 under if lora_context is not None:
  • vllm_ascend/ops/fused_moe/token_dispatcher.py adds an all_to_all of the LoRA indices per layer, plus permute/unpermute
  • vllm_ascend/lora/punica_npu.py: add_lora_fused_moe has no no_lora short-circuit, so bgmv_shrink / bgmv_expand_slice still launch over every row for a zero delta

(The Ascend quantized path already gates on has_lora(lora_context) in ops/fused_moe/moe_mlp.py; only the unquantized path keys on presence.)

Change

Keep the eager wrapping and the GPU publish; detach the context while no adapter is active on runners that expose set_lora_context(). The hook is attached by capability, so GPU wrappers — which have no such entry point and whose experts assert the context exists — are unchanged. Activation, re-arm and rollback all republish through _publish_moe_lora_context(), so an active adapter still gets a real context; deactivation now also releases the cost instead of leaving it published.

.buildkite/common/ci_source_file_dependencies.yml: add vllm_omni/diffusion/lora/ to the hunyuan_image3 anchor. HunyuanImage-3 is the MoE + PEFT model whose baseline path runs this code, and without it the HunyuanImage-3 perf/accuracy steps are filtered out of a LoRA-only diff — which is how this regression reached the release branch unnoticed.

Test Plan

vLLM version: 0.32.0    Base: release/v0.32.0 @ 2dbc92bc0d

Test Result

metric (tp2_cfgp2 / tp2_sp2 / tp4) pre-regression median (13 nightlies) with #9443 (control #9505) revert of #9443 (#9502) this PR
e2e mean (s) 5.43 / 5.70 / 5.81 7.90 / 8.25 / 8.44 5.88 / 5.74 / 5.81 pending
DiT stage_0_gen (s) 4.81 / 5.05 / 5.19 7.30 / 7.62 / 7.82 5.28 / 5.09 / 5.19 pending
model.forward (ms) 3.79 / 3.99 / 4.13 6.31 / 6.53 / 6.67 3.83 / 3.98 / 4.13 pending

H200 is unaffected either way: the 10-28 main nightly (18 commits ahead of #9443) is flat at 2189 / 2381 / 2265 ms.

Follow-up filed upstream: vllm-ascend should gate its unquantized MoE path on has_lora(lora_context) the way the quantized one already does, and/or short-circuit add_lora_fused_moe on punica_wrapper.no_lora. That is the general fix and also helps LLM serving batches that contain no LoRA tokens, but it needs a vllm-ascend release plus an image bump, so it cannot land for v0.32.0. The two changes compose.

#9443 wraps MoE runners eagerly and publishes a no-adapter (slot-0)
mapping before the first forward so LoRA-typed MoE kernels have a context
on HunyuanImage-3's baseline generate() path (#9232). On GPU that is free:
punica sets no_lora_flag for slot 0 and the TRTLLM / Triton LoRA experts
return early.

Ascend's unquantized MoE path instead branches on whether a context is
published at all, not on whether any token selects an adapter:
vllm_ascend/ops/fused_moe/routed_experts.py runs the argsort-based routing
recovery plus moe_lora_apply_w13/w2 under `if lora_context is not None`,
token_dispatcher.py adds an all_to_all of the LoRA indices per layer, and
vllm_ascend/lora/punica_npu.py:add_lora_fused_moe has no no_lora
short-circuit, so bgmv shrink/expand still launch over every row for a
zero delta.

Measured on A3 (910) HunyuanImage-3 nightly perf, same nodes and day:
with the idle context published, e2e is 7.90 / 8.25 / 8.44 s and
HunyuanImage3Pipeline.model.forward 6.31 / 6.53 / 6.67 ms, against a
pre-#9443 baseline of 5.43 / 5.70 / 5.81 s and 3.79 / 3.99 / 4.13 ms
(tp2_cfgp2 / tp2_sp2 / tp4). H200 is unaffected.

Keep the eager wrapping and the GPU publish, and detach the context while
no adapter is active on runners that expose set_lora_context(). The hook
is attached by capability, so GPU wrappers, which have no such entry point
and whose experts assert the context exists, are unchanged. Activation,
re-arm and rollback all republish through _publish_moe_lora_context, so an
active adapter still gets a real context; deactivation now also releases
the cost instead of keeping it published.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

nightly-test label to trigger buildkite nightly test CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant