Repository navigation
[Bugfix][Diffusion] Detach the idle MoE LoRA context on Ascend - #9515
Conversation
#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>
|
(English version of the review.) Keeping wrap at init and only changing publish is the right layer — please do not add a worker The idle switch is wrong. Gating on Please branch per wrapper, not all-or-nothing:
Activation should still |
| # skips rows inside the kernel). Expose an idle hook so the manager can | ||
| # detach the context while no adapter is active. GPU wrappers have no | ||
| # such entry point and keep the published no-adapter mapping, which | ||
| # TRTLLM / Triton LoRA experts assert on (issue #9232). |
There was a problem hiding this comment.
English: this sentence treats TRTLLM and Triton as the same kernel. Only TRTLLM asserts lora_context is not None. Triton gets its context from FusedMoEWithLoRA.set_mapping() → fused_experts.set_lora_context(...); the GPU runner has no set_lora_context, so the idle hook is never installed and H100 keeps publishing slot 0. Triton treats presence of _lora_context as the LoRA switch — same class of idle cost as Ascend, just smaller.
# GPU set_mapping (vllm lora/layers/fused_moe.py)
def set_mapping(self, punica_wrapper):
super().set_mapping(punica_wrapper)
fused_experts.set_lora_context(lora_context) # both Triton and TRTLLM| # their runner supports detaching; the rest keep the slot-0 | ||
| # mapping below. | ||
| idle_hooks = [getattr(ll, "_set_diffusion_moe_lora_idle", None) for ll in moe_wrappers] | ||
| if all(callable(hook) for hook in idle_hooks): |
There was a problem hiding this comment.
English: all(callable(hook)) is all-or-nothing. If every wrapper is Ascend, idle detach works. If any GPU wrapper lacks the hook, all of them take set_mapping(slot 0) below, including layers that should detach.
If the idle hook is later attached on GPU Triton, TRTLLM in the same process would be detached too and hit assert lora_context is not None.
Please split per wrapper:
def _moe_kernel_needs_lora_context(layer) -> bool:
experts = getattr(getattr(layer, "_moe_kernel", None), "fused_experts", None)
return callable(getattr(experts, "_batch_has_no_lora", None)) # TRTLLM only
if self._active_adapter_id is None:
for layer in moe_wrappers:
if _moe_kernel_needs_lora_context(layer):
continue
idle = getattr(layer, "_set_diffusion_moe_lora_idle", None)
if callable(idle):
idle() # Ascend: runner.set_lora_context(None)
else:
experts = getattr(getattr(layer, "_moe_kernel", None), "fused_experts", None)
set_ctx = getattr(experts, "set_lora_context", None)
if callable(set_ctx):
set_ctx(None) # Triton
moe_wrappers = [ll for ll in moe_wrappers if _moe_kernel_needs_lora_context(ll)]
if not moe_wrappers:
return
# set_mapping(slot 0) for TRTLLM onlyOn activate, skip this branch and set_mapping(slot 1) on every wrapper.
Review feedback on the first revision: Triton keys on context presence too, so detaching only on Ascend left H100 paying for the idle context. ``TrtLlmBf16LoRAExperts`` is the only kernel that needs the no-adapter mapping: ``apply()`` asserts the context exists and then returns through ``_batch_has_no_lora()``. ``TritonExperts`` instead disables ``skip_invalid`` in ``_prepare_expert_assignment`` and still runs ``apply_w13_lora`` / ``apply_w2_lora``, whose config lookup and ``moe_lora_align_block_size`` precede the ``no_lora_flag`` return inside ``PunicaWrapperGPU.add_lora_fused_moe``; reported at ~6% of HunyuanImage-3 DiT on H100. Branch per wrapper while idle: keep ``set_mapping(slot 0)`` for kernels exposing ``_batch_has_no_lora``, detach the rest through the runner hook (Ascend) or ``fused_experts`` / ``prepare_finalize`` ``set_lora_context(None)`` (Triton), and keep the mapping for anything that cannot be detached. Activation is unchanged: every wrapper gets the slot-1 mapping. One process may mix kernels, so the decision cannot be all-or-nothing. Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Co-authored-by: Cursor <cursoragent@cursor.com>
|
Thanks — you're right on all three points, and I verified each against vLLM 0.32.0 before changing anything:
Pushed if self._active_adapter_id is None:
needs_context = []
for lora_layer in moe_wrappers:
if moe_lora_kernel_requires_context(lora_layer) or not detach_moe_lora_context(lora_layer):
needs_context.append(lora_layer)
if not needs_context:
return
moe_wrappers = needs_context
# set_mapping(slot 0) for the remaining wrappers only
Tests: four new cases for Ascend / Triton / TRT-LLM / mixed-in-one-process, plus active-adapter and non-detachable fallbacks (68 passing in PR description updated: the "GPU / Triton early-returns" claim is gone, replaced by a per-kernel table including your H100 number (1.293s → 1.376s, ~6%), and the H200 note now says flat e2e does not rule out a Hopper DiT hit. One thing I can't cover from here: I have no H100 to A/B the Triton side. If you can re-run your #9443 A/B/A on this branch I'll fold the number into the PR; otherwise I'll report the NPU numbers only and flag the H100 side as verified-by-construction. |
NPU result — build #10116, A3 (910), all jobs exit 0, 10/10 requestsMeasured on
Every metric is back inside the pre-regression distribution, with one honest exception: Peak memory is unchanged throughout (46.1 / 46.1 / 44.4 GB). Still outstanding: the Triton/H100 side is verified by construction only ( |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Confirmation on the current head — build #10126,
|
| test | pre-regression (13 nightlies) median [min–max] | #10116 (9494d496a) |
#10126 (a63a7eaa0) |
|---|---|---|---|
tp2_cfgp2 e2e mean |
5.43 [5.24–6.06] s | 5.30 s | 5.95 s |
tp2_cfgp2 stage_0_gen |
4.81 [4.64–5.41] s | 4.70 s | 5.31 s |
tp2_cfgp2 model.forward |
3.79 [3.78–3.93] ms | 3.86 ms | 3.79 ms |
tp2_sp2 e2e mean |
5.70 [5.60–6.28] s | 5.69 s | 5.79 s |
tp2_sp2 stage_0_gen |
5.05 [4.98–5.62] s | 5.02 s | 5.14 s |
tp2_sp2 model.forward |
3.99 [3.97–4.04] ms | 3.99 ms | 3.99 ms |
tp4 e2e mean |
5.81 [5.71–6.18] s | 5.91 s | 5.77 s |
tp4 stage_0_gen |
5.19 [5.10–5.56] s | 5.27 s | 5.13 s |
tp4 model.forward |
4.13 [4.10–4.13] ms | 4.23 ms | 4.11 ms |
Every metric on the current head is inside the pre-regression range, and the tp4 model.forward 4.23 ms flagged on the previous run was run-to-run noise — it is 4.11 ms here, against the historical 4.10–4.13 ms. Peak memory unchanged (46.1 / 46.1 / 44.4 GB).
For reference, the same three configs with #9443's idle context published: 7.90 / 8.25 / 8.44 s e2e and 6.31 / 6.53 / 6.67 ms model.forward.
… idle MoE LoRA context on Ascend (#9690) Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: zzh <943967662@qq.com> Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk> Co-authored-by: dengyunyang <584797741@qq.com>
Purpose
Fixes #9514 — the HunyuanImage-3 NPU performance regression introduced by #9443 (reached
release/v0.32.0through 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(). Only one kernel actually wants that mapping while idle:TrtLlmBf16LoRAExperts(B200)apply()assertslora_context is not None_batch_has_no_lora()TritonExperts(H100)skip_invaliddisabled in_prepare_expert_assignment, plusapply_w13_lora/apply_w2_loraconfig lookup andmoe_lora_align_block_sizeahead of theno_lora_flagreturn inPunicaWrapperGPU.add_lora_fused_moe— ~6% of HunyuanImage-3 DiT, reported in reviewAscend'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.pyruns theargsort-based routing recovery plusmoe_lora_apply_w13/moe_lora_apply_w2underif lora_context is not None:vllm_ascend/ops/fused_moe/token_dispatcher.pyadds anall_to_allof the LoRA indices per layer, plus permute/unpermutevllm_ascend/lora/punica_npu.py: add_lora_fused_moehas nono_lorashort-circuit, sobgmv_shrink/bgmv_expand_slicestill launch over every row for a zero delta(The Ascend quantized path already gates on
has_lora(lora_context)inops/fused_moe/moe_mlp.py; only the unquantized path keys on presence.)Change
Keep the eager wrapping, and decide per wrapper while no adapter is active, since one process can mix kernels:
_batch_has_no_lora()(TRT-LLM) keepset_mapping(slot 0)set_lora_context(None)), elsefused_experts/prepare_finalizeset_lora_context(None)on GPUActivation is unchanged: every wrapper gets the slot-1 mapping. 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: addvllm_omni/diffusion/lora/to thehunyuan_image3anchor. 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.Alternative considered: an
if/elseinDiffusionWorkerA smaller-looking fix is to branch at the
init_lora_manager()call site and skip the eager MoE context on NPU, e.g.if current_platform.device_type == "npu": self.init_lora_manager(prepare_moe_lora_context=False). That also restores the baseline numbers, but it is weaker in three ways:_ensure_moe_lora_context()runs inDiffusionLoRAManager.__init__, so the worker can only decline to build the context at startup. Onceadd_adapter()wraps the runners and publishes a context, deactivating or unloading the adapter leaves it published with the slot-0 mapping, and the cost returns for the rest of the process lifetime. Gating at the publish site on_active_adapter_idmakes every idle state cheap, and also releases the cost on deactivation — which [Bugfix][Diffusion] Prepare the MoE LoRA context #9443 does not do today.device_type == "npu"encodes the wrong invariant. The reason is not the device, it is that the kernel treats context presence as the LoRA switch — which Triton does on CUDA too, so a platform branch would fix NPU and leave H100 broken. Probing forset_lora_context()keeps this change self-limiting: once [Performance]: Unquantized MoE path takes the LoRA branch on context presence instead of has_lora(), so idle adapters cost ~40% e2e vllm-ascend#20078 lands and the unquantized path gates onhas_lora(), nothing here needs revisiting, whereas a platform branch would keep NPU on a separate init path indefinitely.add_adapter()on NPU only, outsideset_current_vllm_config, which is safe today only becauseAscendFusedMoEWithLoRA.__init__deliberately skips the GPU modular-kernel construction. Removing exactly that divergence was part of [Bugfix][Diffusion] Prepare the MoE LoRA context #9443. This PR keeps the wrap where [Bugfix][Diffusion] Prepare the MoE LoRA context #9443 put it and changes only whether the context is published.The worker variant also skips the eager wrap along with the publish, so a green run would not say which of the two costs. This PR skips only the publish, so its numbers additionally confirm the mechanism reported upstream.
Test Plan
vLLM version: 0.32.0 Base:
release/v0.32.0@2dbc92bc0dpytest -q tests/diffusion/lora/test_lora_manager.py(two new cases: idle detaches and publishes nothing; active still publishes the slot-1 mapping)test_hunyuan_image_tp2_cfgp2,test_hunyuan_image_tp2_sp2,test_hunyuan_image_tp4Test Result
stage_0_gen(s)model.forward(ms)H200 e2e is flat across #9443 (the 10-28
mainnightly, 18 commits ahead of it, sits at 2189 / 2381 / 2265 ms), but that does not rule out the Triton DiT cost reported on H100 — e2e there is dominated by other stages.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-circuitadd_lora_fused_moeonpunica_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.Supersedes #9512, whose Buildkite jobs could not be scheduled: every job failed in the bootstrap with
fatal: unable to read tree (7a6c826f3d…)/ exit 128 while checking out the PR merge ref, and GitHub stopped syncing that PR's head. This PR carries the identical diff on a fresh branch.