Skip to content

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

Merged
tzhouam merged 2 commits into
release/v0.32.0from
fix/9514-moe-lora-idle-context
Oct 28, 2026
Merged

tzhouam merged 2 commits into
release/v0.32.0from
fix/9514-moe-lora-idle-context

Conversation

@tzhouam

@tzhouam tzhouam commented Oct 28, 2026 •

Copy link
Copy Markdown
Collaborator

Purpose

Fixes #9514 — 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(). Only one kernel actually wants that mapping while idle:

kernel idle context cost of publishing it
TrtLlmBf16LoRAExperts (B200) required — apply() asserts lora_context is not None none: returns through _batch_has_no_lora()
TritonExperts (H100) not required skip_invalid disabled in _prepare_expert_assignment, plus apply_w13_lora / apply_w2_lora config lookup and moe_lora_align_block_size ahead of the no_lora_flag return in PunicaWrapperGPU.add_lora_fused_moe — ~6% of HunyuanImage-3 DiT, reported in review
Ascend unquantized MoE not required the full LoRA path, below — 40–50% e2e

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 decide per wrapper while no adapter is active, since one process can mix kernels:

  • experts exposing _batch_has_no_lora() (TRT-LLM) keep set_mapping(slot 0)
  • otherwise detach: the runner idle hook on Ascend (set_lora_context(None)), else fused_experts / prepare_finalize set_lora_context(None) on GPU
  • anything that cannot be detached keeps the mapping rather than being left without one

Activation 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: 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.

Alternative considered: an if/else in DiffusionWorker

A 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:

  • It only covers the idle state before the first adapter. _ensure_moe_lora_context() runs in DiffusionLoRAManager.__init__, so the worker can only decline to build the context at startup. Once add_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_id makes 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 for set_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 on has_lora(), nothing here needs revisiting, whereas a platform branch would keep NPU on a separate init path indefinitely.
  • It re-splits NPU and GPU initialization. The MoE wrap would move back to the first add_adapter() on NPU only, outside set_current_vllm_config, which is safe today only because AscendFusedMoEWithLoRA.__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 @ 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 e2e is flat across #9443 (the 10-28 main nightly, 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-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.


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.

#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>
@tzhouam tzhouam added the nightly-test label to trigger buildkite nightly test CI label Oct 28, 2026
BLANKETusers

This comment was marked as outdated.

@BLANKETusers

Copy link
Copy Markdown
Contributor

(English version of the review.)

Keeping wrap at init and only changing publish is the right layer — please do not add a worker if npu branch. Detaching the idle context on Ascend is the correct fix for #9514.

The idle switch is wrong. Gating on runner.set_lora_context plus all(idle_hooks) only hits Ascend; GPU Triton still gets set_mapping(slot 0). Once _lora_context is set, Triton does not take the original base path (skip_invalid / apply_w13_lora stay on even with no_lora_flag). We measured ~6% HunyuanImage-3 DiT on H100 for #9443 (1.293s → 1.376s, A/B/A). Only TRTLLM asserts a context and early-returns on _batch_has_no_lora; slot-0 is correct there. H200 e2e being flat does not rule out a Hopper DiT hit.

Please branch per wrapper, not all-or-nothing:

  • experts have _batch_has_no_lora → still set_mapping(slot 0) (TRTLLM)
  • otherwise do not set_mapping: call the idle hook if present (Ascend), else fused_experts.set_lora_context(None) (Triton)

Activation should still set_mapping(slot 1) on every wrapper. Wrap stays at init. That fixes NPU and H100 together. Please drop the claim that GPU / Triton early-returns on the no-adapter mapping.

Comment thread vllm_omni/diffusion/lora/utils.py Outdated
# 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).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread vllm_omni/diffusion/lora/manager.py Outdated
# 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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 only

On 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>
@tzhouam

tzhouam commented Oct 28, 2026

Copy link
Copy Markdown
Collaborator Author

Thanks — you're right on all three points, and I verified each against vLLM 0.32.0 before changing anything:

  • vllm/model_executor/layers/fused_moe/experts/triton_moe.py:353-356 — skip_invalid = expert_map is not None and self._lora_context is None and not a2_scale_is_global, so a published idle context also disables ignore_invalid_experts in _prepare_expert_assignment, i.e. it costs something in the base path too.
  • apply_w13_lora → PunicaWrapperGPU.add_lora_w13 does the tile-config lookup and moe_lora_align_block_size before add_lora_fused_moe, and the no_lora_flag return is only at punica_gpu.py:485-486 — so "Triton early-returns" was wrong as written.
  • _batch_has_no_lora exists only in trtllm_lora_moe.py, which is exactly the right capability probe for "needs the slot-0 mapping".

Pushed a63a7eaa0 with the per-wrapper split you asked for:

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
  • moe_lora_kernel_requires_context() = callable(getattr(layer._moe_kernel.fused_experts, "_batch_has_no_lora", None)) → TRT-LLM keeps set_mapping(slot 0).
  • detach_moe_lora_context() → Ascend idle hook if present, otherwise set_lora_context(None) on fused_experts and prepare_finalize (mirroring what FusedMoEWithLoRA.set_mapping publishes on the way in).
  • A wrapper that can be neither identified nor detached keeps the no-adapter mapping instead of being left with none, so a kernel that does require one never ends up contextless.
  • Activation is untouched: every wrapper gets set_mapping(slot 1). Wrap stays at init. No worker if npu.

Tests: four new cases for Ascend / Triton / TRT-LLM / mixed-in-one-process, plus active-adapter and non-detachable fallbacks (68 passing in tests/diffusion/lora/test_lora_manager.py).

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.

@tzhouam tzhouam added nightly-test label to trigger buildkite nightly test CI and removed nightly-test label to trigger buildkite nightly test CI labels Oct 28, 2026
@tzhouam

tzhouam commented Oct 28, 2026

Copy link
Copy Markdown
Collaborator Author

NPU result — build #10116, A3 (910), all jobs exit 0, 10/10 requests

Measured on 9494d496a (the Ascend-only revision of this branch). a63a7eaa0 adds the Triton branch, which does not change the NPU path — Ascend wrappers carry the runner idle hook and take the same detach — but I'm re-running NPU CI on the current head to confirm rather than assert it.

test metric pre-regression (13 nightlies, 10-15..10-26) median [min–max] 10-27 nightlies control, with #9443 (#10096) revert of #9443 (#10092) this PR
tp2_cfgp2 e2e mean 5.43 [5.24–6.06] s 7.95 s 7.90 s 5.88 s 5.30 s
stage_0_gen 4.81 [4.64–5.41] s 7.32 s 7.30 s 5.28 s 4.70 s
model.forward 3.79 [3.78–3.93] ms 6.31 ms 6.31 ms 3.83 ms 3.86 ms
throughput 0.184 [0.165–0.191] qps 0.126 qps 0.127 qps 0.170 qps 0.189 qps
tp2_sp2 e2e mean 5.70 [5.60–6.28] s 8.27 s 8.25 s 5.74 s 5.69 s
stage_0_gen 5.05 [4.98–5.62] s 7.62 s 7.62 s 5.09 s 5.02 s
model.forward 3.99 [3.97–4.04] ms 6.50 ms 6.53 ms 3.98 ms 3.99 ms
throughput 0.175 [0.159–0.179] qps 0.121 qps 0.121 qps 0.174 qps 0.176 qps
tp4 e2e mean 5.81 [5.71–6.18] s 8.33 s 8.44 s 5.81 s 5.91 s
stage_0_gen 5.19 [5.10–5.56] s 7.70 s 7.82 s 5.19 s 5.27 s
model.forward 4.13 [4.10–4.13] ms 6.65 ms 6.67 ms 4.13 ms 4.23 ms
throughput 0.172 [0.162–0.175] qps 0.120 qps 0.120 qps 0.172 qps 0.169 qps

Every metric is back inside the pre-regression distribution, with one honest exception: tp4 model.forward is 4.23 ms against a historical max of 4.13 ms (+2.4%, single run; stage_0_gen and e2e for the same config are both inside range, and the revert measured 4.13 ms). I'll see whether it persists on the re-run of the current head.

Peak memory is unchanged throughout (46.1 / 46.1 / 44.4 GB).

Still outstanding: the Triton/H100 side is verified by construction only (_lora_context is cleared, so skip_invalid returns to its pre-#9443 value and apply_w13_lora / apply_w2_lora are not called). @BLANKETusers if you can re-run your #9443 A/B/A on this branch, I'll add the number here.

@tzhouam tzhouam added nightly-test label to trigger buildkite nightly test CI and removed nightly-test label to trigger buildkite nightly test CI labels Oct 28, 2026
@tzhouam
tzhouam marked this pull request as ready for review October 28, 2026 08:05
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ⚠️ Failed 2026-10-28T08:05:25.199578Z
We couldn’t complete this review. Comment @codex review to try again.
a63a7ea Draft marked ready
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

@tzhouam

tzhouam commented Oct 28, 2026

Copy link
Copy Markdown
Collaborator Author

Confirmation on the current head — build #10126, a63a7eaa0 (per-wrapper revision), A3 (910)

All jobs exit 0, 10/10 requests, accuracy test included.

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.

@hsliuustc0106 hsliuustc0106 added the high priority high priority issue, needs to be done asap label Oct 28, 2026 — with ChatGPT Codex Connector
@tzhouam
tzhouam merged commit 5917730 into release/v0.32.0 Oct 28, 2026
10 of 11 checks passed
Bounty-hunter added a commit that referenced this pull request Nov 3, 2026
… 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

high priority high priority issue, needs to be done asap nightly-test label to trigger buildkite nightly test CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants