diff --git a/.github/vllm-release-tag.commit b/.github/vllm-release-tag.commit index 0c2a959e86e..6897c006a5f 100644 --- a/.github/vllm-release-tag.commit +++ b/.github/vllm-release-tag.commit @@ -1 +1 @@ -v0.23.0 +v0.24.0 diff --git a/Dockerfile b/Dockerfile index ff4d421b2ba..a9ab6e75a7b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -39,7 +39,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 ARG VLLM_COMMIT="" RUN if [ -n "$VLLM_COMMIT" ]; then \ git init /vllm-workspace/vllm && \ diff --git a/Dockerfile.310p b/Dockerfile.310p index 1906ae61c81..3d71eeaa197 100644 --- a/Dockerfile.310p +++ b/Dockerfile.310p @@ -33,7 +33,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 ARG VLLM_COMMIT="" RUN if [ -n "$VLLM_COMMIT" ]; then \ git init /vllm-workspace/vllm && \ diff --git a/Dockerfile.310p.openEuler b/Dockerfile.310p.openEuler index 4db663cd9f9..ff4ec8dec25 100644 --- a/Dockerfile.310p.openEuler +++ b/Dockerfile.310p.openEuler @@ -32,7 +32,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 ARG VLLM_COMMIT="" RUN if [ -n "$VLLM_COMMIT" ]; then \ git init /vllm-workspace/vllm && \ diff --git a/Dockerfile.a3 b/Dockerfile.a3 index cfdd24b6a5f..b2d4cefd0b3 100644 --- a/Dockerfile.a3 +++ b/Dockerfile.a3 @@ -41,7 +41,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 ARG VLLM_COMMIT="" RUN if [ -n "$VLLM_COMMIT" ]; then \ git init /vllm-workspace/vllm && \ @@ -77,4 +77,4 @@ RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pyp RUN echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc -CMD ["/bin/bash"] \ No newline at end of file +CMD ["/bin/bash"] diff --git a/Dockerfile.a3.openEuler b/Dockerfile.a3.openEuler index dce37b06dd5..f938946f3d0 100644 --- a/Dockerfile.a3.openEuler +++ b/Dockerfile.a3.openEuler @@ -38,7 +38,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 ARG VLLM_COMMIT="" RUN if [ -n "$VLLM_COMMIT" ]; then \ git init /vllm-workspace/vllm && \ @@ -72,4 +72,4 @@ RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pyp RUN echo "export LD_PRELOAD=/usr/lib64/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc -CMD ["/bin/bash"] \ No newline at end of file +CMD ["/bin/bash"] diff --git a/Dockerfile.a5 b/Dockerfile.a5 index 99372b58007..c05781e5ddc 100644 --- a/Dockerfile.a5 +++ b/Dockerfile.a5 @@ -41,7 +41,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 RUN git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm # In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it. RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \ diff --git a/Dockerfile.a5.openEuler b/Dockerfile.a5.openEuler index 8ee3d7d9bf4..d769f51d263 100644 --- a/Dockerfile.a5.openEuler +++ b/Dockerfile.a5.openEuler @@ -38,7 +38,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 RUN git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm # In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it. RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \ diff --git a/Dockerfile.openEuler b/Dockerfile.openEuler index 2274097f6dd..f45269c2f03 100644 --- a/Dockerfile.openEuler +++ b/Dockerfile.openEuler @@ -38,7 +38,7 @@ RUN pip config set global.index-url ${PIP_INDEX_URL} && \ # Install vLLM ARG VLLM_REPO=https://github.com/vllm-project/vllm.git -ARG VLLM_TAG=v0.23.0 +ARG VLLM_TAG=v0.24.0 ARG VLLM_COMMIT="" RUN if [ -n "$VLLM_COMMIT" ]; then \ git init /vllm-workspace/vllm && \ diff --git a/README.md b/README.md index 9e3167533f2..b4c167ecee3 100644 --- a/README.md +++ b/README.md @@ -90,7 +90,7 @@ Below are the maintained branches: | Branch | Status | Note | |------------------|--------------|--------------------------------------| -| main | Maintained | CI commitment for vLLM main branch and vLLM v0.23.0 tag | +| main | Maintained | CI commitment for vLLM main branch and vLLM v0.24.0 tag | | v0.7.1-dev | Unmaintained | Outdated, no longer maintained. | | v0.7.3-dev | Unmaintained | Only bug fixes are allowed, and no new release tags anymore. | | v0.9.1-dev | Unmaintained | Only bug fixes are allowed, and no new release tags anymore. | diff --git a/README.zh.md b/README.zh.md index 813aff220bc..b78d312ebfe 100644 --- a/README.zh.md +++ b/README.zh.md @@ -84,7 +84,7 @@ vllm-ascend有主干分支和开发分支。 | 分支 | 状态 | 备注 | |------------------|--------------|----------------------| -| main | Maintained | 基于vLLM main分支和vLLM最新版本（v0.23.0）CI看护 | +| main | Maintained | 基于vLLM main分支和vLLM最新版本（v0.24.0）CI看护 | | v0.7.1-dev | Unmaintained | 不再维护 | | v0.7.3-dev | Unmaintained | 只允许Bug修复，不会再发布新版本 | | v0.9.1-dev | Unmaintained | 只允许Bug修复，不会再发布新版本 | diff --git a/docs/source/community/slash-commands.md b/docs/source/community/slash-commands.md index dc699501a5a..f399719f4fa 100644 --- a/docs/source/community/slash-commands.md +++ b/docs/source/community/slash-commands.md @@ -70,7 +70,7 @@ binary bisect for genuine failures. By default, AOP hooks are disabled. /nightly qwen3-vl-32b-instruct-w8a8 # Run on a specific release branch -/nightly qwen3-vl-32b-instruct-w8a8 --branch releases/v0.23.0 +/nightly qwen3-vl-32b-instruct-w8a8 --branch releases/v0.24.0 # Run all tests on a specific branch /nightly all --branch my-feature-branch @@ -88,7 +88,7 @@ binary bisect for genuine failures. By default, AOP hooks are disabled. /nightly all --aop_enabled # Run specific test with AOP on a release branch -/nightly test_custom_op --branch releases/v0.23.0 --aop_enabled +/nightly test_custom_op --branch releases/v0.24.0 --aop_enabled ``` This triggers `workflow_dispatch` on both `schedule_nightly_test_a2.yaml` and `schedule_nightly_test_a3.yaml`. @@ -107,7 +107,7 @@ Cherry-pick a PR's commits onto a specified target branch and create a new PR. T ```text # Cherry-pick to a release branch -/cherry-pick releases/v0.23.0 +/cherry-pick releases/v0.24.0 # Cherry-pick to main /cherry-pick main diff --git a/docs/source/developer_guide/Design_Documents/balance_schedule_refactor.md b/docs/source/developer_guide/Design_Documents/balance_schedule_refactor.md index ba074ef5820..bc451404a88 100644 --- a/docs/source/developer_guide/Design_Documents/balance_schedule_refactor.md +++ b/docs/source/developer_guide/Design_Documents/balance_schedule_refactor.md @@ -11,7 +11,7 @@ conditional activation). The `schedule()` copy is **kept for now** — upstream exposes no finer-grained hook to borrow, and deleting it depends on contributing an override seam upstream, tracked as later Phase 2B. The file therefore does not shrink to a few dozen lines: the `schedule()` body is still a verbatim -upstream copy (**aligned verbatim to release tag `v0.23.0`**, with only 3 +upstream copy (**aligned verbatim to release tag `v0.24.0`**, with only 3 balance deltas), and the file is ~830 lines. Aligning to a stable release tag (rather than a moving main-verified commit hash) makes "verbatim comparison against upstream" a reproducible drift check — a fixed tag points at the same @@ -85,24 +85,24 @@ three large upstream units verbatim: This "copy whole units" approach has three concrete harms: 1. **The `schedule()` copy is now aligned verbatim to a release tag (the - production pin, currently `v0.23.0`).** The **single source of truth** for + production pin, currently `v0.24.0`).** The **single source of truth** for the release tag is `.github/vllm-release-tag.commit` (CI reads the same file via `tr -d '[:space:]' < .github/vllm-release-tag.commit`), currently - `v0.23.0`; dev/CI actually installs the main-verified commit pointed at by - `.github/vllm-main-verified.commit` (which already carries `throttle_prefills` - and other v0.23.1+ evolution). The old patch copied a `schedule()` from an - older vLLM than v0.23.0, so it was stale as a whole. This round aligns the + `v0.24.0`; dev/CI actually installs the main-verified commit pointed at by + `.github/vllm-main-verified.commit` (which carries later scheduler + evolution). The old patch copied a `schedule()` from an + older vLLM than v0.24.0, so it was stale as a whole. This round aligns the `schedule()` copy **verbatim to the release tag's `Scheduler.schedule()`**, keeping only the 3 balance deltas (disabled-path early return, `balance_flag` gate, `if request_queue is None: break`); the `run_busy_loop()` / `run_engine_core()` copies were deleted in Phase 1. - **Note: any concrete `v0.23.0` in this document is just a snapshot of the pin + **Note: any concrete `v0.24.0` in this document is just a snapshot of the pin file's current value — it goes stale as the pin advances and must NOT be used as a version authority; any code/test that needs this tag must read the file at runtime.** - **Why align to the v0.23.0 tag rather than the installed main-verified - commit?** Two reasons: (a) production actually runs the v0.23.0 release, so + **Why align to the v0.24.0 tag rather than the installed main-verified + commit?** Two reasons: (a) production actually runs the v0.24.0 release, so aligning the copy to it keeps production behavior consistent with runtime; (b) a fixed git tag points at the **same** source on every CI run, so "verbatim comparison of the copy against upstream (allowing only the 3 @@ -110,22 +110,15 @@ This "copy whole units" approach has three concrete harms: main-verified hash makes the comparison drift forward with every commit and cannot serve as a stable guardrail. - **Cost and boundary:** the copy (v0.23.0 logic) and the main-verified + **Cost and boundary:** the copy (v0.24.0 logic) and the main-verified runtime differ slightly in behavior, but balance's real scheduling path is only reached under NPU + DP + MoE, never by CPU UT (see Test plan); and - those differences do not affect the gate's own semantics. **Key: vllm-ascend - CI runs two vllm versions at once** (release tag v0.23.0 + main-verified - commit 1f486d96), whose engines call `schedule()` differently — v0.23.0 calls - `schedule()`, 1f486d96 calls `schedule(throttle_prefills)`. So the override - signature takes the **union of both versions** - (`schedule(self, throttle_prefills=False)`, a default-carrying superset) so - both engines can call it; the disabled path delegates to `super()` and uses - **signature introspection** (`_SUPER_SCHEDULE_HAS_THROTTLE`, decided once at - import) to decide whether to forward `throttle_prefills`, rather than a - version string — this is correct on both lanes and is not affected by a dev - checkout's non-standard PEP 440 `__version__` (which would make - `vllm_version_is` raise). I.e.: **body aligned to the release tag; signature - takes the two-version union; disabled path uses signature introspection.** + those differences do not affect the gate's own semantics. Both supported + revisions — release tag v0.24.0 and main-verified commit e5588e49 — expose + `schedule(self, throttle_prefills=False)`, so the override matches that + shared signature and the disabled path forwards `throttle_prefills` + directly to `super()`. I.e.: **body aligned to the release tag; signature + matches both supported revisions; disabled path delegates directly.** 2. **It violates the `AGENTS.md` patch policy.** The policy requires patches to be "minimal and focused" with "a long-term plan to contribute upstream". A @@ -138,31 +131,11 @@ This "copy whole units" approach has three concrete harms: `if request_queue is None: break`. Such deviations make future diffs untrustworthy. -> Lesson learned this round (a pit we hit — `schedule()` signature drift across -> versions + dual-version CI). vLLM's `Scheduler.schedule()` signature changes -> across versions: release tag **v0.23.0** is `schedule(self)` (engine calls -> `schedule()`), main-verified commit **1f486d96** adds -> `throttle_prefills: bool = False` (engine calls -> `schedule(self._should_throttle_prefills())`). And **vllm-ascend CI runs both -> at once**, so the override cannot hardcode either signature — `schedule(self)` -> would `TypeError` on 1f486d96 when the engine passes the arg. The right fix: -> the override signature takes the **union of both versions** -> `def schedule(self, throttle_prefills: bool = False)` (a default-carrying -> superset callable by both engines); the disabled path delegates to `super()` -> and uses **signature introspection** (computed once at import: -> `_SUPER_SCHEDULE_HAS_THROTTLE = "throttle_prefills" in inspect.signature(Scheduler.schedule).parameters`) -> to decide whether to forward `throttle_prefills`, instead of a -> `vllm_version_is("0.23.0")` version string — the latter raises `ValueError` -> on a dev checkout (non-PEP 440 `__version__`) and fundamentally reframes the -> factual question "does super() accept this arg?" as a brittle "does the -> version string match?". Conclusion: **when CI runs multiple upstream versions -> at once, take the union for override signatures and drive version-dependent -> branches by signature introspection, not version strings.** At the unit-test -> level, "signature equals installed vLLM's" is **wrong** under dual versions -> (the two lanes' installed signatures differ by definition, so an equality -> assertion can only ever pass on one lane); it is replaced by "the override is -> bindable by both engines' call shapes and is a superset of the installed -> signature". +> Lesson learned this round: after v0.23 support is dropped, both supported +> revisions expose the same `Scheduler.schedule(self, throttle_prefills=False)` +> contract. Compatibility introspection and version-string branching therefore +> add no value and should be removed. The unit test now requires the override's +> signature to equal the installed upstream signature in each CI lane. ## Design @@ -310,17 +283,11 @@ currently no overridable seam. This is solved in two phases. Keep the `schedule()` override, but: -- **The override signature takes the two-version union:** - `def schedule(self, throttle_prefills: bool = False)`. vllm-ascend CI runs - v0.23.0 (engine calls `schedule()`) and 1f486d96 (engine calls - `schedule(throttle_prefills)`) at the same time, so the override carries - `throttle_prefills` with a default and is callable by both engines. **The - disabled path delegates to `super()` via signature introspection** — compute - `_SUPER_SCHEDULE_HAS_THROTTLE` once at import; if true call - `super().schedule(throttle_prefills)`, else `super().schedule()`. This - replaces the old `vllm_version_is("0.23.0")` version-string branch (which - `ValueError`s on a dev checkout and reframes "does super accept the arg?" as - "does the version match?"). +- **The override matches the shared supported signature:** + `def schedule(self, throttle_prefills: bool = False)`. Both v0.24.0 and + e5588e49 expose this signature, and the disabled path delegates directly via + `super().schedule(throttle_prefills)`. The old v0.23 compatibility branch and + signature introspection are no longer needed. - Collapse the balance changes into 3 clearly-commented deltas: (1) the disabled-path early return delegating to `super()`; (2) the `balance_flag` gate inside the WAITING loop; (3) `if request_queue is None: break` (upstream @@ -329,7 +296,7 @@ Keep the `schedule()` override, but: - **Verbatim comparison is now reproducible:** the `schedule()` copy is aligned to the release tag (only the 3 balance deltas differ), so the fixed tag makes "verbatim comparison against upstream" yield the same baseline on every CI - run. The "intent lock" tests (signature callable by both engines, the 3 delta + run. The "intent lock" tests (signature equality, the 3 delta lines present, upstream seams still exist) remain as CPU-reachable guardrails, and a new "verbatim comparison against the release tag (allowing only the 3 deltas)" drift test is added (see Test plan). The drift test **reads the tag @@ -426,7 +393,8 @@ deviation is a bug. `_balance_run_engine_core` restores the module-level `DPEngineCoreProc` to the upstream original and the engine core runs upstream's implementation verbatim; `BalanceScheduler` with `_balance_enabled=False` delegates - `schedule()` to `super().schedule()`, does not allocate `balance_queue`, and + `schedule(throttle_prefills)` to `super().schedule(throttle_prefills)`, does + not allocate `balance_queue`, and performs no collective communication. I.e. balance does not touch any config when off (including PD-disaggregated recompute / `AsyncRecomputeScheduler`, which is already mutually exclusive with balance via `platform.py`; this is a @@ -438,14 +406,13 @@ deviation is a bug. ## Post-refactor file shape After this round (Phase 1 + 2A + 3), the key structure is as follows. The -`schedule()` body is a verbatim copy of release tag `v0.23.0` (cannot be +`schedule()` body is a verbatim copy of release tag `v0.24.0` (cannot be deleted before Phase 2B), with three documented deltas (disabled-path early return + the `balance_flag` gate in the WAITING loop + `if request_queue is None: break`): ```python # vllm_ascend/patch/platform/patch_balance_schedule.py -import inspect import torch import torch.distributed as dist import vllm.v1.core.sched.scheduler as _sched_mod @@ -481,16 +448,11 @@ class BalanceScheduler(Scheduler): running_tensor = torch.tensor([len(self.running)], dtype=torch.int, device="cpu") dist.all_gather(self.balance_queue, running_tensor, group=self.dp_group) - def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # two-version union: both v0.23.0's schedule() and 1f486d96's schedule(throttle_prefills) bind + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # shared by v0.24.0 and e5588e49 if not self._balance_enabled: # delta 1: disabled-path early return - # Whether to forward throttle_prefills is decided by signature introspection - # (_SUPER_SCHEDULE_HAS_THROTTLE), not a version string -- correct on both CI lanes - # and unaffected by a dev checkout's bad __version__. - if _SUPER_SCHEDULE_HAS_THROTTLE: - return super().schedule(throttle_prefills) - return super().schedule() + return super().schedule(throttle_prefills) # NOTE: balance_gather is NOT called here -- see BalanceDPEngineCoreProc. - # ... upstream schedule() body (verbatim-aligned to the v0.23.0 tag) ... + # ... upstream schedule() body (verbatim-aligned to the v0.24.0 tag) ... # # inside the WAITING loop (deltas 2, 3): # if max(t.item() for t in self.balance_queue) == self.max_num_running_reqs: # delta 2: leader-at-cap => global freeze # break @@ -551,13 +513,11 @@ first inside `schedule()`, then inside `_process_engine_step`; now after ## Test plan 1. **Signature + intent lock + verbatim drift test (Phase 2A).** Assert: (a) - `BalanceScheduler.schedule`'s signature is **bindable by both engine call - shapes** (`schedule()` and `schedule(throttle_prefills=...)`) **and is a - superset of the installed `Scheduler.schedule` parameter set** — note this is - NOT "signature line equals installed", because under dual-version CI the two - lanes' installed signatures differ and an equality assertion can only pass on - one lane (this locks the "dual versions ⇒ take the union" lesson); (b) the 3 - balance delta lines must exist in the body (disabled-path `super().schedule()` + `BalanceScheduler.schedule`'s signature **equals the installed + `Scheduler.schedule` signature**; both supported revisions share this + contract, so each CI lane checks the same invariant; (b) the 3 + balance delta lines must exist in the body (disabled-path + `super().schedule(throttle_prefills)` delegation, the `balance_flag` gate in the WAITING loop, `if request_queue is None: break`); (c) the `_balance_run_engine_core` wrapper is installed and `DPEngineCoreProc` is **not** swapped at import @@ -587,8 +547,8 @@ first inside `schedule()`, then inside `_process_engine_step`; now after that each `balance_gather()` does exactly one `all_gather`, with payload `len(self.running)` and the injected dp_group (contract item 3). 4. **Disabled-path test.** With the flag off, assert `balance_queue` is not - allocated, `all_gather` is not called, and `schedule()` delegates to - `super().schedule()` (contract item 4). + allocated, `all_gather` is not called, and `schedule(throttle_prefills)` + delegates to `super().schedule(throttle_prefills)` (contract item 4). 5. **NPU performance check.** Per AGENTS.md's NPU guidance, `max(t.item() for t in self.balance_queue)` triggers one host sync per step (unavoidable, since this value drives host-side control flow). Profile to @@ -599,7 +559,7 @@ first inside `schedule()`, then inside `_process_engine_step`; now after | Phase | Scope | Risk | Depends on | Status | |-------|------------------------------------------------------------------------------------------------------------------------|------|----------------|---------------| | 1 | Hook gather onto `_has_global_unfinished_reqs` (after the cross-rank all-reduce — avoids both the schedule()-skip deadlock and the _process_engine_step wave-boundary deadlock); slim `BalanceDPEngineCoreProc` to that hook; delete the `run_engine_core`/`run_busy_loop` copies; `run_engine_core` wrapper conditionally activates `DPEngineCoreProc`; module-level `Scheduler` swap | Low | none | ✅ Done | -| 2A | Override signature takes the **two-version union** (`schedule(self, throttle_prefills=False)`, CI runs v0.23.0 + 1f486d96 at once); **body aligned verbatim to the release tag** (only the 3 balance deltas); disabled path delegates to `super()` via **signature introspection** (`_SUPER_SCHEDULE_HAS_THROTTLE`); signature-callability + intent-lock + release-tag verbatim drift tests | Low | none | ✅ Done | +| 2A | Override matches the shared supported signature (`schedule(self, throttle_prefills=False)` on v0.24.0 + e5588e49); **body aligned verbatim to the release tag** (only the 3 balance deltas); disabled path delegates directly to `super()`; signature equality + intent-lock + release-tag verbatim drift tests | Low | none | ✅ Done | | 3 | Collapse config probing to two fallbacks (AscendConfig → additional_config); remove the direct env-var read (still parsed centrally by AscendConfig) | Low | Phase 1 | ✅ Done | | 2B | Upstream `_should_stop_admitting_waiting` PR; delete the `schedule()` copy | Med | upstream review | ⏳ TODO | | Tests | Drift regression / behavior equivalence / gather cadence / disabled path / NPU performance check | Low | Phase 1 + 2A | ⏳ TODO (needs NPU) | diff --git a/docs/source/locale/zh_CN/LC_MESSAGES/developer_guide/Design_Documents/balance_schedule_refactor.po b/docs/source/locale/zh_CN/LC_MESSAGES/developer_guide/Design_Documents/balance_schedule_refactor.po index 932a7b6a0f3..08a4c19242a 100644 --- a/docs/source/locale/zh_CN/LC_MESSAGES/developer_guide/Design_Documents/balance_schedule_refactor.po +++ b/docs/source/locale/zh_CN/LC_MESSAGES/developer_guide/Design_Documents/balance_schedule_refactor.po @@ -33,7 +33,7 @@ msgid "" "exposes no finer-grained hook to borrow, and deleting it depends on contributing\n" "an override seam upstream, tracked as later Phase 2B. The file therefore does\n" "not shrink to a few dozen lines: the `schedule()` body is still a verbatim\n" -"upstream copy (**aligned verbatim to release tag `v0.23.0`**, with only 3\n" +"upstream copy (**aligned verbatim to release tag `v0.24.0`**, with only 3\n" "balance deltas), and the file is ~830 lines. Aligning to a stable release tag\n" "(rather than a moving main-verified commit hash) makes \"verbatim comparison\n" "against upstream\" a reproducible drift check — a fixed tag points at the same\n" @@ -49,7 +49,7 @@ msgstr "" "语义**的前提下，首先删除了两个已经过时的副本 `run_busy_loop()` / `run_engine_core()`（将其替换为基于 " "`_has_global_unfinished_reqs` 的引擎核心钩子，以及用于条件激活的模块级名称交换）。`schedule()` " "的副本**暂时保留**——上游没有提供更细粒度的钩子供借用，删除它依赖于向上游贡献一个覆盖接缝，计划在后续的 Phase 2B " -"中完成。因此，该文件并未缩减到几十行：`schedule()` 的主体仍然是上游的逐字副本（**与发布标签 `v0.23.0` " +"中完成。因此，该文件并未缩减到几十行：`schedule()` 的主体仍然是上游的逐字副本（**与发布标签 `v0.24.0` " "逐字对齐**，仅包含3处平衡增量），文件约830行。与稳定的发布标签（而非变动的主分支验证提交哈希）对齐，使得“与上游逐字比较”成为可复现的漂移检查——固定标签在每次CI运行中都指向相同的源代码。本轮实际移除的是" " `run_busy_loop` / `run_engine_core` 副本的陈旧漂移风险，并修复了一个启用平衡时的死锁（最初是“在 " "`schedule()` 内部收集”，后来在一次迭代中改为“在 `_process_engine_step` 内部收集”）；收集操作现在紧跟在 " @@ -159,35 +159,35 @@ msgstr "这种“复制整个单元”的方法有三个具体危害：" msgid "" "1. **The `schedule()` copy is now aligned verbatim to a release tag (the\n" -" production pin, currently `v0.23.0`).** The **single source of truth** for\n" +" production pin, currently `v0.24.0`).** The **single source of truth** for\n" " the release tag is `.github/vllm-release-tag.commit` (CI reads the same file\n" " via `tr -d '[:space:]' < .github/vllm-release-tag.commit`), currently\n" -" `v0.23.0`; dev/CI actually installs the main-verified commit pointed at by\n" -" `.github/vllm-main-verified.commit` (which already carries `throttle_prefills`\n" -" and other v0.23.1+ evolution). The old patch copied a `schedule()` from an\n" -" older vLLM than v0.23.0, so it was stale as a whole. This round aligns the\n" +" `v0.24.0`; dev/CI actually installs the main-verified commit pointed at by\n" +" `.github/vllm-main-verified.commit` (which carries later scheduler\n" +" evolution). The old patch copied a `schedule()` from an\n" +" older vLLM than v0.24.0, so it was stale as a whole. This round aligns the\n" " `schedule()` copy **verbatim to the release tag's `Scheduler.schedule()`**,\n" " keeping only the 3 balance deltas (disabled-path early return,\n" " `balance_flag` gate, `if request_queue is None: break`); the\n" " `run_busy_loop()` / `run_engine_core()` copies were deleted in Phase 1.\n" -" **Note: any concrete `v0.23.0` in this document is just a snapshot of the pin\n" +" **Note: any concrete `v0.24.0` in this document is just a snapshot of the pin\n" " file's current value — it goes stale as the pin advances and must NOT be\n" " used as a version authority; any code/test that needs this tag must read the\n" " file at runtime.**" msgstr "" -"1. **`schedule()` 的副本现已与发布标签（生产固定版本，当前为 `v0.23.0`）逐字对齐。** 发布标签的**唯一真实来源**是 " +"1. **`schedule()` 的副本现已与发布标签（生产固定版本，当前为 `v0.24.0`）逐字对齐。** 发布标签的**唯一真实来源**是 " "`.github/vllm-release-tag.commit`（CI通过 `tr -d '[:space:]' < .github/vllm-" -"release-tag.commit` 读取同一文件），当前为 `v0.23.0`；开发/CI实际安装的是 `.github/vllm-main-" -"verified.commit` 指向的主分支验证提交（该提交已包含 `throttle_prefills` 和其他 v0.23.1+ " -"的演进）。旧补丁复制了一个比 v0.23.0 更旧的 vLLM 中的 `schedule()`，因此整体已经过时。本轮将 `schedule()` " +"release-tag.commit` 读取同一文件），当前为 `v0.24.0`；开发/CI实际安装的是 `.github/vllm-main-" +"verified.commit` 指向的主分支验证提交（该提交包含后续调度器演进）。旧补丁复制了一个比 v0.24.0 更旧的 " +"vLLM 中的 `schedule()`，因此整体已经过时。本轮将 `schedule()` " "副本**逐字对齐到发布标签的 `Scheduler.schedule()`**，仅保留3处平衡增量（禁用路径的提前返回、`balance_flag` " "门控、`if request_queue is None: break`）；`run_busy_loop()` / " -"`run_engine_core()` 的副本已在 Phase 1 中删除。**注意：本文档中任何具体的 `v0.23.0` " +"`run_engine_core()` 的副本已在 Phase 1 中删除。**注意：本文档中任何具体的 `v0.24.0` " "只是固定版本文件当前值的快照——随着固定版本推进它会过时，绝不能用作版本权威依据；任何需要此标签的代码/测试必须在运行时读取该文件。**" msgid "" -"**Why align to the v0.23.0 tag rather than the installed main-verified\n" -" commit?** Two reasons: (a) production actually runs the v0.23.0 release, so\n" +"**Why align to the v0.24.0 tag rather than the installed main-verified\n" +" commit?** Two reasons: (a) production actually runs the v0.24.0 release, so\n" " aligning the copy to it keeps production behavior consistent with runtime;\n" " (b) a fixed git tag points at the **same** source on every CI run, so\n" " \"verbatim comparison of the copy against upstream (allowing only the 3\n" @@ -195,39 +195,26 @@ msgid "" " main-verified hash makes the comparison drift forward with every commit and\n" " cannot serve as a stable guardrail." msgstr "" -"**为什么对齐到 v0.23.0 标签而不是已安装的 main-verified 提交？** 两个原因：(a) 生产环境实际运行的是 v0.23.0 " +"**为什么对齐到 v0.24.0 标签而不是已安装的 main-verified 提交？** 两个原因：(a) 生产环境实际运行的是 v0.24.0 " "版本，因此将副本对齐到该版本可保持生产行为与运行时一致；(b) 固定的 git 标签在每次 CI " "运行时都指向**相同**的源码，因此“对副本与上游进行逐字比较（仅允许 3 处差异）”成为**可重现**的漂移检查——而移动的 main-" "verified 哈希会随着每次提交向前漂移，无法作为稳定的护栏。" msgid "" -"**Cost and boundary:** the copy (v0.23.0 logic) and the main-verified\n" +"**Cost and boundary:** the copy (v0.24.0 logic) and the main-verified\n" " runtime differ slightly in behavior, but balance's real scheduling path is\n" " only reached under NPU + DP + MoE, never by CPU UT (see Test plan); and\n" -" those differences do not affect the gate's own semantics. **Key: vllm-ascend\n" -" CI runs two vllm versions at once** (release tag v0.23.0 + main-verified\n" -" commit 1f486d96), whose engines call `schedule()` differently — v0.23.0 calls\n" -" `schedule()`, 1f486d96 calls `schedule(throttle_prefills)`. So the override\n" -" signature takes the **union of both versions**\n" -" (`schedule(self, throttle_prefills=False)`, a default-carrying superset) so\n" -" both engines can call it; the disabled path delegates to `super()` and uses\n" -" **signature introspection** (`_SUPER_SCHEDULE_HAS_THROTTLE`, decided once at\n" -" import) to decide whether to forward `throttle_prefills`, rather than a\n" -" version string — this is correct on both lanes and is not affected by a dev\n" -" checkout's non-standard PEP 440 `__version__` (which would make\n" -" `vllm_version_is` raise). I.e.: **body aligned to the release tag; signature\n" -" takes the two-version union; disabled path uses signature introspection.**" -msgstr "" -"**代价与边界：** 副本（v0.23.0 逻辑）与 main-verified 运行时的行为略有差异，但 balance 的实际调度路径仅在 NPU " -"+ DP + MoE 下触发，CPU 单元测试不会涉及（参见测试计划）；这些差异不影响门控本身的语义。**关键点：vllm-ascend CI " -"同时运行两个 vllm 版本**（发布标签 v0.23.0 + main-verified 提交 1f486d96），它们的引擎调用 " -"`schedule()` 的方式不同——v0.23.0 调用 `schedule()`，1f486d96 调用 " -"`schedule(throttle_prefills)`。因此，重写签名采用**两个版本的并集**（`schedule(self, " -"throttle_prefills=False)`，一个带默认值的超集），使得两个引擎都能调用它；禁用路径委托给 `super()` " -"并使用**签名自省**（`_SUPER_SCHEDULE_HAS_THROTTLE`，在导入时决定一次）来决定是否传递 " -"`throttle_prefills`，而不是使用版本字符串——这在两条路径上都是正确的，并且不受开发检出分支的非标准 PEP 440 " -"`__version__` 影响（后者会导致 `vllm_version_is` " -"抛出异常）。即：**主体对齐发布标签；签名采用两版本并集；禁用路径使用签名自省。**" +" those differences do not affect the gate's own semantics. Both supported\n" +" revisions — release tag v0.24.0 and main-verified commit e5588e49 — expose\n" +" `schedule(self, throttle_prefills=False)`, so the override matches that\n" +" shared signature and the disabled path forwards `throttle_prefills`\n" +" directly to `super()`. I.e.: **body aligned to the release tag; signature\n" +" matches both supported revisions; disabled path delegates directly.**" +msgstr "" +"**代价与边界：** 副本（v0.24.0 逻辑）与 main-verified 运行时的行为略有差异，但 balance 的实际调度路径仅在 NPU " +"+ DP + MoE 下触发，CPU 单元测试不会涉及（参见测试计划）；这些差异不影响门控本身的语义。两个受支持的版本——发布标签 " +"v0.24.0 和 main-verified 提交 e5588e49——都提供 `schedule(self, throttle_prefills=False)`；因此重写与该共享签名一致，" +"禁用路径也会直接将 `throttle_prefills` 传给 `super()`。即：**主体对齐发布标签；签名匹配两个受支持版本；禁用路径直接委托。**" msgid "" "2. **It violates the `AGENTS.md` patch policy.** The policy requires patches to\n" @@ -249,46 +236,15 @@ msgstr "" " request_queue is None: break`。此类偏差会使未来的差异变得不可信。" msgid "" -"> Lesson learned this round (a pit we hit — `schedule()` signature drift across\n" -"> versions + dual-version CI). vLLM's `Scheduler.schedule()` signature changes\n" -"> across versions: release tag **v0.23.0** is `schedule(self)` (engine calls\n" -"> `schedule()`), main-verified commit **1f486d96** adds\n" -"> `throttle_prefills: bool = False` (engine calls\n" -"> `schedule(self._should_throttle_prefills())`). And **vllm-ascend CI runs both\n" -"> at once**, so the override cannot hardcode either signature — `schedule(self)`\n" -"> would `TypeError` on 1f486d96 when the engine passes the arg. The right fix:\n" -"> the override signature takes the **union of both versions**\n" -"> `def schedule(self, throttle_prefills: bool = False)` (a default-carrying\n" -"> superset callable by both engines); the disabled path delegates to `super()`\n" -"> and uses **signature introspection** (computed once at import:\n" -"> `_SUPER_SCHEDULE_HAS_THROTTLE = \"throttle_prefills\" in inspect.signature(Scheduler.schedule).parameters`)\n" -"> to decide whether to forward `throttle_prefills`, instead of a\n" -"> `vllm_version_is(\"0.23.0\")` version string — the latter raises `ValueError`\n" -"> on a dev checkout (non-PEP 440 `__version__`) and fundamentally reframes the\n" -"> factual question \"does super() accept this arg?\" as a brittle \"does the\n" -"> version string match?\". Conclusion: **when CI runs multiple upstream versions\n" -"> at once, take the union for override signatures and drive version-dependent\n" -"> branches by signature introspection, not version strings.** At the unit-test\n" -"> level, \"signature equals installed vLLM's\" is **wrong** under dual versions\n" -"> (the two lanes' installed signatures differ by definition, so an equality\n" -"> assertion can only ever pass on one lane); it is replaced by \"the override is\n" -"> bindable by both engines' call shapes and is a superset of the installed\n" -"> signature\"." -msgstr "" -"> 本轮经验教训（我们遇到的一个坑——`schedule()` 签名在不同版本间漂移 + 双版本 CI）。vLLM 的 " -"`Scheduler.schedule()` 签名在不同版本间会变化：发布标签 **v0.23.0** 是 `schedule(self)`（引擎调用 " -"`schedule()`），main-verified 提交 **1f486d96** 增加了 `throttle_prefills: bool = " -"False`（引擎调用 `schedule(self._should_throttle_prefills())`）。而 **vllm-ascend CI" -" 同时运行两者**，因此重写不能硬编码任一签名——`schedule(self)` 在 1f486d96 上，当引擎传入参数时会引发 " -"`TypeError`。正确的修复：重写签名采用**两个版本的并集** `def schedule(self, throttle_prefills: " -"bool = False)`（一个带默认值的超集，两个引擎都能调用）；禁用路径委托给 `super()` " -"并使用**签名自省**（在导入时计算一次：`_SUPER_SCHEDULE_HAS_THROTTLE = \"throttle_prefills\" " -"in inspect.signature(Scheduler.schedule).parameters`）来决定是否传递 " -"`throttle_prefills`，而不是使用 `vllm_version_is(\"0.23.0\")` 版本字符串——后者在开发检出分支（非 " -"PEP 440 `__version__`）上会引发 `ValueError`，并且从根本上将事实性问题“super() " -"接受这个参数吗？”重构为脆弱的“版本字符串匹配吗？”。结论：**当 CI " -"同时运行多个上游版本时，重写签名采用并集，版本相关的分支由签名自省驱动，而非版本字符串。** 在单元测试层面，“签名等于已安装 vLLM " -"的签名”在双版本下是**错误**的（两条路径的已安装签名定义上就不同，因此相等性断言永远只能通过一条路径）；它被替换为“重写可被两个引擎的调用形态绑定，并且是已安装签名的超集”。" +"> Lesson learned this round: after v0.23 support is dropped, both supported\n" +"> revisions expose the same `Scheduler.schedule(self, throttle_prefills=False)`\n" +"> contract. Compatibility introspection and version-string branching therefore\n" +"> add no value and should be removed. The unit test now requires the override's\n" +"> signature to equal the installed upstream signature in each CI lane." +msgstr "" +"> 本轮经验：移除 v0.23 支持后，两个受支持版本提供相同的 " +"`Scheduler.schedule(self, throttle_prefills=False)` 契约。因此兼容性自省和版本字符串分支已无价值，应当删除。" +"单元测试现在要求重写签名在每条 CI 路径中都与已安装的上游签名一致。" msgid "## Design" msgstr "## 设计" @@ -497,17 +453,11 @@ msgid "Keep the `schedule()` override, but:" msgstr "保留`schedule()`覆盖，但：" msgid "" -"- **The override signature takes the two-version union:**\n" -" `def schedule(self, throttle_prefills: bool = False)`. vllm-ascend CI runs\n" -" v0.23.0 (engine calls `schedule()`) and 1f486d96 (engine calls\n" -" `schedule(throttle_prefills)`) at the same time, so the override carries\n" -" `throttle_prefills` with a default and is callable by both engines. **The\n" -" disabled path delegates to `super()` via signature introspection** — compute\n" -" `_SUPER_SCHEDULE_HAS_THROTTLE` once at import; if true call\n" -" `super().schedule(throttle_prefills)`, else `super().schedule()`. This\n" -" replaces the old `vllm_version_is(\"0.23.0\")` version-string branch (which\n" -" `ValueError`s on a dev checkout and reframes \"does super accept the arg?\" as\n" -" \"does the version match?\").\n" +"- **The override matches the shared supported signature:**\n" +" `def schedule(self, throttle_prefills: bool = False)`. Both v0.24.0 and\n" +" e5588e49 expose this signature, and the disabled path delegates directly via\n" +" `super().schedule(throttle_prefills)`. The old v0.23 compatibility branch and\n" +" signature introspection are no longer needed.\n" "- Collapse the balance changes into 3 clearly-commented deltas: (1) the\n" " disabled-path early return delegating to `super()`; (2) the `balance_flag`\n" " gate inside the WAITING loop; (3) `if request_queue is None: break` (upstream\n" @@ -516,7 +466,7 @@ msgid "" "- **Verbatim comparison is now reproducible:** the `schedule()` copy is aligned\n" " to the release tag (only the 3 balance deltas differ), so the fixed tag makes\n" " \"verbatim comparison against upstream\" yield the same baseline on every CI\n" -" run. The \"intent lock\" tests (signature callable by both engines, the 3 delta\n" +" run. The \"intent lock\" tests (signature equality, the 3 delta\n" " lines present, upstream seams still exist) remain as CPU-reachable guardrails,\n" " and a new \"verbatim comparison against the release tag (allowing only the 3\n" " deltas)\" drift test is added (see Test plan). The drift test **reads the tag\n" @@ -528,16 +478,11 @@ msgid "" " the release tag advances, re-apply the 3 deltas onto the new tag's\n" " `schedule()` (continues until Phase 2B deletes the copy)." msgstr "" -"- **重写签名采用双版本联合：**\n" -" `def schedule(self, throttle_prefills: bool = False)`。vllm-ascend CI 同时运行\n" -" v0.23.0（引擎调用 `schedule()`）和 1f486d96（引擎调用\n" -" `schedule(throttle_prefills)`），因此重写携带带有默认值的 `throttle_prefills`，两个引擎均可调用。**禁用路径通过签名内省委托给 `super()`** — 在导入时计算一次\n" -" `_SUPER_SCHEDULE_HAS_THROTTLE`；如果为真则调用\n" -" `super().schedule(throttle_prefills)`，否则调用 `super().schedule()`。这\n" -" 取代了旧的 `vllm_version_is(\"0.23.0\")` 版本字符串分支（该分支在开发检出上会引发\n" -" `ValueError`，并将“父类是否接受该参数？”重新表述为“版本是否匹配？”）。\n" +"- **重写匹配受支持版本的共享签名：**\n" +" `def schedule(self, throttle_prefills: bool = False)`。v0.24.0 和 e5588e49 " +"都提供此签名，禁用路径通过 `super().schedule(throttle_prefills)` 直接委托。旧的 v0.23 兼容分支和签名自省已不再需要。\n" "- 将平衡变更压缩为3个带有清晰注释的差异点：(1) 禁用路径的早期返回委托给 `super()`；(2) WAITING 循环内的 `balance_flag` 门控；(3) `if request_queue is None: break`（上游使用 `assert`）。由于上游没有更细粒度的钩子，函数体仍需复制。\n" -"- **逐字比较现在可重现：** `schedule()` 的副本与发布标签对齐（仅3个平衡差异点不同），因此固定的标签使得每次 CI 运行时“与上游的逐字比较”都能产生相同的基线。“意图锁定”测试（签名可被两个引擎调用，存在3个差异行，上游接缝仍然存在）继续作为 CPU 可及的护栏，并新增了一个“与发布标签的逐字比较（仅允许3个差异点）”的漂移测试（参见测试计划）。漂移测试**在运行时从 `.github/vllm-release-tag.commit` 读取标签**（与 CI 同源）——它不硬编码版本也不读取设计文档；当固定点前进时，测试自动与新标签比较并变红以指示“副本需要重新同步”。\n" +"- **逐字比较现在可重现：** `schedule()` 的副本与发布标签对齐（仅3个平衡差异点不同），因此固定的标签使得每次 CI 运行时“与上游的逐字比较”都能产生相同的基线。“意图锁定”测试（签名相等、存在3个差异行、上游接缝仍然存在）继续作为 CPU 可及的护栏，并新增了一个“与发布标签的逐字比较（仅允许3个差异点）”的漂移测试（参见测试计划）。漂移测试**在运行时从 `.github/vllm-release-tag.commit` 读取标签**（与 CI 同源）——它不硬编码版本也不读取设计文档；当固定点前进时，测试自动与新标签比较并变红以指示“副本需要重新同步”。\n" "- “在固定点前进时重新对齐副本”现在是常规维护：每次发布标签前进时，将3个差异点重新应用到新标签的 `schedule()` 上（持续到阶段2B删除该副本）。" msgid "**Phase 2B — target (lands with an upstream contribution):**" @@ -643,7 +588,8 @@ msgid "" " `_balance_run_engine_core` restores the module-level `DPEngineCoreProc` to\n" " the upstream original and the engine core runs upstream's implementation\n" " verbatim; `BalanceScheduler` with `_balance_enabled=False` delegates\n" -" `schedule()` to `super().schedule()`, does not allocate `balance_queue`, and\n" +" `schedule(throttle_prefills)` to `super().schedule(throttle_prefills)`, does\n" +" not allocate `balance_queue`, and\n" " performs no collective communication. I.e. balance does not touch any config\n" " when off (including PD-disaggregated recompute / `AsyncRecomputeScheduler`,\n" " which is already mutually exclusive with balance via `platform.py`; this is a\n" @@ -656,7 +602,7 @@ msgstr "" " `max(balance_queue) == max_num_running_reqs`，根据上一步收集到的各 rank 的 `len(running)` 计算。当条件为真时，所有 rank 都不再接受新的 WAITING 请求。比较操作是 `==` 与配置的 `max_num_running_reqs` 进行比较——**不是** `>=`，**也不是**“追赶领导者”。\n" "2. **相同输入 ⇒ 相同输出。** 给定相同的 `self.running`、`self.waiting`、`self.skipped_waiting`、`balance_queue` 和 token 预算，重构后的 `schedule()` 产生的 `SchedulerOutput` 与当前实现完全相同（相同的已调度/已抢占/已恢复集合，相同的 `num_scheduled_tokens`，相同的连接器元数据）。\n" "3. **收集节奏不变。** 每个活跃引擎步骤恰好执行一次 `all_gather`，在相同的 DP 组上，负载仍然是 `len(self.running)`，当所有 rank 都空闲时，所有 rank 一致地跳过。仅调用位置发生了变化。\n" -"4. **禁用路径不变。** 当 `enable_balance_scheduling` 为 false 时，`_balance_run_engine_core` 将模块级别的 `DPEngineCoreProc` 恢复为上游原始实现，并且引擎核心逐字运行上游的实现；`BalanceScheduler` 在 `_balance_enabled=False` 时将 `schedule()` 委托给 `super().schedule()`，不分配 `balance_queue`，并且不执行任何集合通信。即，当平衡功能关闭时，它不会触及任何配置（包括 PD 分离重计算 / `AsyncRecomputeScheduler`，后者已通过 `platform.py` 与平衡功能互斥；这是第二层防御）。\n" +"4. **禁用路径不变。** 当 `enable_balance_scheduling` 为 false 时，`_balance_run_engine_core` 将模块级别的 `DPEngineCoreProc` 恢复为上游原始实现，并且引擎核心逐字运行上游的实现；`BalanceScheduler` 在 `_balance_enabled=False` 时将 `schedule(throttle_prefills)` 委托给 `super().schedule(throttle_prefills)`，不分配 `balance_queue`，并且不执行任何集合通信。即，当平衡功能关闭时，它不会触及任何配置（包括 PD 分离重计算 / `AsyncRecomputeScheduler`，后者已通过 `platform.py` 与平衡功能互斥；这是第二层防御）。\n" "5. **现有约束仍然适用。** `profiling_chunk_config` 互斥锁（参见 `vllm_ascend/ascend_config.py`）和 PD 混合模式限制（参见 `vllm_ascend/platform.py`）仍在原处强制执行。" msgid "## Post-refactor file shape" @@ -664,12 +610,12 @@ msgstr "## 重构后的文件结构" msgid "" "After this round (Phase 1 + 2A + 3), the key structure is as follows. The\n" -"`schedule()` body is a verbatim copy of release tag `v0.23.0` (cannot be\n" +"`schedule()` body is a verbatim copy of release tag `v0.24.0` (cannot be\n" "deleted before Phase 2B), with three documented deltas (disabled-path early\n" "return + the `balance_flag` gate in the WAITING loop + `if request_queue is\n" "None: break`):" msgstr "" -"经过本轮（阶段 1 + 2A + 3）后，关键结构如下。`schedule()` 主体是发布标签 `v0.23.0` 的逐字副本（在阶段 2B " +"经过本轮（阶段 1 + 2A + 3）后，关键结构如下。`schedule()` 主体是发布标签 `v0.24.0` 的逐字副本（在阶段 2B " "之前不能删除），包含三个已记录的差异（禁用路径的提前返回 + WAITING 循环中的 `balance_flag` 门控 + `if " "request_queue is None: break`）：" @@ -702,13 +648,11 @@ msgstr "## 测试计划" msgid "" "1. **Signature + intent lock + verbatim drift test (Phase 2A).** Assert: (a)\n" -" `BalanceScheduler.schedule`'s signature is **bindable by both engine call\n" -" shapes** (`schedule()` and `schedule(throttle_prefills=...)`) **and is a\n" -" superset of the installed `Scheduler.schedule` parameter set** — note this is\n" -" NOT \"signature line equals installed\", because under dual-version CI the two\n" -" lanes' installed signatures differ and an equality assertion can only pass on\n" -" one lane (this locks the \"dual versions ⇒ take the union\" lesson); (b) the 3\n" -" balance delta lines must exist in the body (disabled-path `super().schedule()`\n" +" `BalanceScheduler.schedule`'s signature **equals the installed\n" +" `Scheduler.schedule` signature**; both supported revisions share this\n" +" contract, so each CI lane checks the same invariant; (b) the 3\n" +" balance delta lines must exist in the body (disabled-path\n" +" `super().schedule(throttle_prefills)`\n" " delegation, the `balance_flag` gate in the WAITING loop,\n" " `if request_queue is None: break`); (c) the `_balance_run_engine_core`\n" " wrapper is installed and `DPEngineCoreProc` is **not** swapped at import\n" @@ -738,8 +682,8 @@ msgid "" " that each `balance_gather()` does exactly one `all_gather`, with payload\n" " `len(self.running)` and the injected dp_group (contract item 3).\n" "4. **Disabled-path test.** With the flag off, assert `balance_queue` is not\n" -" allocated, `all_gather` is not called, and `schedule()` delegates to\n" -" `super().schedule()` (contract item 4).\n" +" allocated, `all_gather` is not called, and `schedule(throttle_prefills)`\n" +" delegates to `super().schedule(throttle_prefills)` (contract item 4).\n" "5. **NPU performance check.** Per AGENTS.md's NPU guidance,\n" " `max(t.item() for t in self.balance_queue)` triggers one host sync per step\n" " (unavoidable, since this value drives host-side control flow). Profile to\n" @@ -765,13 +709,11 @@ msgstr "" "| 风险 | 依赖项 | 状态 |" msgid "" -"| 2A | Override signature takes the **two-version union** " -"(`schedule(self, throttle_prefills=False)`, CI runs v0.23.0 + 1f486d96 at " -"once); **body aligned verbatim to the release tag** (only the 3 balance " -"deltas); disabled path delegates to `super()` via **signature " -"introspection** (`_SUPER_SCHEDULE_HAS_THROTTLE`); signature-callability + " -"intent-lock + release-tag verbatim drift tests | Low | none | ✅ " -"Done |" +"| 2A | Override matches the shared supported signature " +"(`schedule(self, throttle_prefills=False)` on v0.24.0 + e5588e49); **body " +"aligned verbatim to the release tag** (only the 3 balance deltas); disabled " +"path delegates directly to `super()`; signature equality + intent-lock + " +"release-tag verbatim drift tests | Low | none | ✅ Done |" msgstr "" "| 2A | 覆盖签名采用**双版本并集**（`schedule(self, throttle_prefills=False)`，CI同时运行 " "v0.23.0 + " diff --git a/examples/offline_data_parallel.py b/examples/offline_data_parallel.py index bcf8ee31006..c3fa8d4c285 100644 --- a/examples/offline_data_parallel.py +++ b/examples/offline_data_parallel.py @@ -123,13 +123,9 @@ def main( os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port) - from vllm_ascend.utils import vllm_version_is + import torch - _dp_device_ids = None - if not vllm_version_is("0.23.0"): - import torch - - _dp_device_ids = [str(i) for i in range(torch.npu.device_count())] + _dp_device_ids = [str(i) for i in range(torch.npu.device_count())] # Sample prompts. prompts = [ diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index cd2f40ae5c7..f58e584b9cf 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -881,23 +881,20 @@ def _run_vllm_runner_dp_worker(conn, llm_kwargs: dict[str, Any], dp_rank: int, d os.environ["VLLM_DP_MASTER_PORT"] = str(master_port) os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - from vllm_ascend.utils import vllm_version_is - - if not vllm_version_is("0.23.0"): - import torch + import torch - visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "") - full_device_ids: list[str] = [d for d in visible.split(",") if d] - if not full_device_ids: - full_device_ids = [str(i) for i in range(torch.npu.device_count())] + visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "") + full_device_ids: list[str] = [d for d in visible.split(",") if d] + if not full_device_ids: + full_device_ids = [str(i) for i in range(torch.npu.device_count())] - if llm_kwargs.get("distributed_executor_backend") == "ray": - devs = full_device_ids - chunk = max(len(devs) // dp_size, 1) - start = dp_rank * chunk - os.environ["ASCEND_RT_VISIBLE_DEVICES"] = ",".join(devs[start : start + chunk]) - else: - llm_kwargs["device_ids"] = full_device_ids + if llm_kwargs.get("distributed_executor_backend") == "ray": + devs = full_device_ids + chunk = max(len(devs) // dp_size, 1) + start = dp_rank * chunk + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = ",".join(devs[start : start + chunk]) + else: + llm_kwargs["device_ids"] = full_device_ids llm = LLM(**llm_kwargs) conn.send({"status": "ready", "rank": dp_rank}) diff --git a/tests/e2e/pull_request/one_card/model_runner_v2/test_basic.py b/tests/e2e/pull_request/one_card/model_runner_v2/test_basic.py index 761202807d7..bd2574ef67f 100644 --- a/tests/e2e/pull_request/one_card/model_runner_v2/test_basic.py +++ b/tests/e2e/pull_request/one_card/model_runner_v2/test_basic.py @@ -36,8 +36,8 @@ DSPARK_MODELS = ["deepseek-ai/dspark_qwen3_8b_block7"] pytestmark = pytest.mark.skipif( - vllm_version_is("0.23.0"), - reason="v2 model runner patches not supported on v0.23.0", + vllm_version_is("0.24.0"), + reason="v2 model runner patches are only supported on the verified vLLM main commit", ) diff --git a/tests/e2e/pull_request/one_card/model_runner_v2/test_uva.py b/tests/e2e/pull_request/one_card/model_runner_v2/test_uva.py index c26c87554c2..23e32051c4a 100644 --- a/tests/e2e/pull_request/one_card/model_runner_v2/test_uva.py +++ b/tests/e2e/pull_request/one_card/model_runner_v2/test_uva.py @@ -27,8 +27,8 @@ MODELS = ["Qwen/Qwen3-0.6B"] pytestmark = pytest.mark.skipif( - vllm_version_is("0.23.0"), - reason="v2 model runner patches not supported on v0.23.0", + vllm_version_is("0.24.0"), + reason="v2 model runner patches are only supported on the verified vLLM main commit", ) diff --git a/tests/e2e/pull_request/one_card/spec_decode/test_dspark.py b/tests/e2e/pull_request/one_card/spec_decode/test_dspark.py index f033bf79ba8..7a0bdb7c321 100644 --- a/tests/e2e/pull_request/one_card/spec_decode/test_dspark.py +++ b/tests/e2e/pull_request/one_card/spec_decode/test_dspark.py @@ -11,8 +11,8 @@ from vllm_ascend.utils import vllm_version_is pytestmark = pytest.mark.skipif( - vllm_version_is("0.23.0"), - reason="The community has not yet incorporated the dspark feature in 0.23.0", + vllm_version_is("0.24.0"), + reason="The community has not yet incorporated the DSpark feature in vLLM v0.24.0", ) diff --git a/tests/e2e/pull_request/one_card/test_guided_decoding.py b/tests/e2e/pull_request/one_card/test_guided_decoding.py index 75e302fcf42..4ff2630ca27 100644 --- a/tests/e2e/pull_request/one_card/test_guided_decoding.py +++ b/tests/e2e/pull_request/one_card/test_guided_decoding.py @@ -40,8 +40,8 @@ @pytest.fixture(params=[False, True], ids=["v1", "v2"]) def model_runner_env(request): use_v2_model_runner = request.param - if use_v2_model_runner and vllm_version_is("0.23.0"): - pytest.skip("No need to support v2 model runner for vLLM tag version.") + if use_v2_model_runner and vllm_version_is("0.24.0"): + pytest.skip("The v2 model runner is not supported on vLLM v0.24.0.") with patch.dict(os.environ, {"VLLM_USE_V2_MODEL_RUNNER": "1" if use_v2_model_runner else "0"}): yield @@ -90,26 +90,9 @@ def test_guided_json_completion_xgrammar(sample_json_schema, request): sampling_params = SamplingParams( temperature=1.0, max_tokens=500, structured_outputs=StructuredOutputsParams(json=sample_json_schema) ) - if not vllm_version_is("0.23.0"): - model_marker = request.node.get_closest_marker("model") - model_marker.kwargs["env_vars"] = REGEX_COMPILATION_TIMEOUT_ENV - with patch.dict(os.environ, REGEX_COMPILATION_TIMEOUT_ENV, clear=False): - vllm_runner = request.getfixturevalue("vllm_runner") - prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 - inputs = vllm_runner.get_inputs(prompts) - outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) - - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt - generated_text = output.outputs[0].text - assert generated_text is not None - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - output_json = json.loads(generated_text) - jsonschema.validate(instance=output_json, schema=sample_json_schema) - else: + model_marker = request.node.get_closest_marker("model") + model_marker.kwargs["env_vars"] = REGEX_COMPILATION_TIMEOUT_ENV + with patch.dict(os.environ, REGEX_COMPILATION_TIMEOUT_ENV, clear=False): vllm_runner = request.getfixturevalue("vllm_runner") prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 inputs = vllm_runner.get_inputs(prompts) @@ -254,26 +237,9 @@ def test_guided_json_completion_outlines(sample_json_schema, request): sampling_params = SamplingParams( temperature=1.0, max_tokens=500, structured_outputs=StructuredOutputsParams(json=sample_json_schema) ) - if not vllm_version_is("0.23.0"): - model_marker = request.node.get_closest_marker("model") - model_marker.kwargs["env_vars"] = REGEX_COMPILATION_TIMEOUT_ENV - with patch.dict(os.environ, REGEX_COMPILATION_TIMEOUT_ENV, clear=False): - vllm_runner = request.getfixturevalue("vllm_runner") - prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 - inputs = vllm_runner.get_inputs(prompts) - outputs = vllm_runner.model.generate(inputs, sampling_params=sampling_params) - - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt - generated_text = output.outputs[0].text - assert generated_text is not None - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - output_json = json.loads(generated_text) - jsonschema.validate(instance=output_json, schema=sample_json_schema) - else: + model_marker = request.node.get_closest_marker("model") + model_marker.kwargs["env_vars"] = REGEX_COMPILATION_TIMEOUT_ENV + with patch.dict(os.environ, REGEX_COMPILATION_TIMEOUT_ENV, clear=False): vllm_runner = request.getfixturevalue("vllm_runner") prompts = [f"Give an example JSON for an employee profile that fits this schema: {sample_json_schema}"] * 2 inputs = vllm_runner.get_inputs(prompts) diff --git a/tests/e2e/pull_request/two_card/test_gpt_oss_distributed.py b/tests/e2e/pull_request/two_card/test_gpt_oss_distributed.py index cfda67715c2..020fb0fe8fa 100644 --- a/tests/e2e/pull_request/two_card/test_gpt_oss_distributed.py +++ b/tests/e2e/pull_request/two_card/test_gpt_oss_distributed.py @@ -24,6 +24,7 @@ import pytest from tests.e2e.conftest import VllmRunner +from vllm_ascend.utils import vllm_version_is GPT_OSS_MODELS = [ "unsloth/gpt-oss-20b-BF16", @@ -31,6 +32,10 @@ @pytest.mark.parametrize("model", GPT_OSS_MODELS) +@pytest.mark.skipif( + vllm_version_is("0.24.0"), + reason="GPT-OSS checkpoint loading requires vLLM #45818, which is not in v0.24.0.", +) def test_gpt_oss_distributed_tp2(model): example_prompts = [ "Hello, my name is", diff --git a/tests/e2e/pull_request/two_card/test_qwen3_30b_a3b.py b/tests/e2e/pull_request/two_card/test_qwen3_30b_a3b.py index 0da49fb63cb..7ccf08ff7a9 100644 --- a/tests/e2e/pull_request/two_card/test_qwen3_30b_a3b.py +++ b/tests/e2e/pull_request/two_card/test_qwen3_30b_a3b.py @@ -23,12 +23,8 @@ from vllm.utils.network_utils import get_open_port from tests.e2e.conftest import RemoteOpenAIServer, wait_until_npu_memory_free -from vllm_ascend.utils import vllm_version_is -pytestmark = pytest.mark.skipif( - not vllm_version_is("0.23.0"), - reason="broken on main, fix me.", -) +pytestmark = pytest.mark.skip(reason="broken on vLLM v0.24.0 and the verified main commit, fix me.") @pytest.mark.e2e_model("Qwen/Qwen3-30B-A3B") diff --git a/tests/e2e/pull_request/two_card/test_qwen3_moe_eplb.py b/tests/e2e/pull_request/two_card/test_qwen3_moe_eplb.py index ec330f59a6e..2a456324ce3 100644 --- a/tests/e2e/pull_request/two_card/test_qwen3_moe_eplb.py +++ b/tests/e2e/pull_request/two_card/test_qwen3_moe_eplb.py @@ -23,12 +23,8 @@ from vllm.utils.network_utils import get_open_port from tests.e2e.conftest import RemoteOpenAIServer -from vllm_ascend.utils import vllm_version_is -pytestmark = pytest.mark.skipif( - not vllm_version_is("0.23.0"), - reason="broken on main, fix me.", -) +pytestmark = pytest.mark.skip(reason="broken on vLLM v0.24.0 and the verified main commit, fix me.") @pytest.mark.asyncio diff --git a/tests/ut/_310p/fused_moe/test_shared_fused_moe_310.py b/tests/ut/_310p/fused_moe/test_shared_fused_moe_310.py index 2a97cf622c5..bc6b6724ca6 100644 --- a/tests/ut/_310p/fused_moe/test_shared_fused_moe_310.py +++ b/tests/ut/_310p/fused_moe/test_shared_fused_moe_310.py @@ -1,109 +1,161 @@ -# -# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from unittest.mock import patch +# SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +import pytest import torch -import torch.nn.functional as F +from torch import nn +from vllm_ascend._310p.fused_moe import fused_moe as fused_moe_310_module from vllm_ascend._310p.fused_moe.fused_moe import ( - AscendFusedMoE310, + AscendMoERunner310, + AscendUnquantizedFusedMoEMethod310, ) - - -class _DummyGate(torch.nn.Module): - def forward(self, hidden_states: torch.Tensor): - # Keep gate output deterministic: sigmoid(0)=0.5. - return torch.zeros( - hidden_states.shape[0], - 1, - dtype=hidden_states.dtype, - device=hidden_states.device, - ), None - - -class _DummySharedExperts(torch.nn.Module): - def __init__(self, with_gate: bool): - super().__init__() - self.expert_gate = _DummyGate() if with_gate else None - - def forward(self, hidden_states: torch.Tensor): - out = hidden_states * 2.0 + 1.0 - if self.expert_gate is not None: - gate_out, _ = self.expert_gate(hidden_states) - out = F.sigmoid(gate_out) * out - return out - - -def _build_layer(shared_experts: torch.nn.Module | None) -> AscendFusedMoE310: - layer = AscendFusedMoE310.__new__(AscendFusedMoE310) - # The test bypasses full layer init with __new__, so we must initialize - # nn.Module internals before assigning child modules. - torch.nn.Module.__init__(layer) - layer._shared_experts = shared_experts - return layer - - -def test_forward_shared_experts_without_gate_310(): - layer = _build_layer(_DummySharedExperts(with_gate=False)) - hidden_states = torch.randn(4, 8) - output = layer._forward_shared_experts(hidden_states) - expected = hidden_states * 2.0 + 1.0 - torch.testing.assert_close(output, expected) - - -def test_forward_shared_experts_with_gate_310(): - layer = _build_layer(_DummySharedExperts(with_gate=True)) - hidden_states = torch.randn(4, 8) - output = layer._forward_shared_experts(hidden_states) - expected = 0.5 * (hidden_states * 2.0 + 1.0) +from vllm_ascend.ascend_forward_context import MoECommType +from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner +from vllm_ascend.quantization.quant_type import QuantType + + +def _build_runner() -> AscendMoERunner310: + runner = AscendMoERunner310.__new__(AscendMoERunner310) + nn.Module.__init__(runner) + return runner + + +def _build_weight_layer(): + return SimpleNamespace( + w13_weight=nn.Parameter(torch.randn(2, 3, 4)), + w2_weight=nn.Parameter(torch.randn(2, 4, 3)), + ) + + +def test_runner_310_installs_specialized_unquantized_method_and_comm(): + runner = _build_runner() + moe_config = MagicMock() + runner.moe_config = moe_config + runner._get_quant_type = MagicMock(return_value=QuantType.NONE) + routed_experts = SimpleNamespace(quant_config=None, quant_method=None) + quant_method = object() + comm_method = object() + + with ( + patch.object(AscendMoERunner, "__init__", return_value=None) as parent_init, + patch.object(fused_moe_310_module, "AscendUnquantizedFusedMoEMethod310", return_value=quant_method), + patch.object(fused_moe_310_module, "AllGatherCommImpl310", return_value=comm_method), + patch.dict(fused_moe_310_module._MoECommMethods, clear=False), + ): + AscendMoERunner310.__init__( + runner, + "model.layers.0.mlp", + moe_config, + MagicMock(), + routed_experts, + ) + + assert routed_experts.quant_method is quant_method + assert runner.quant_type == QuantType.NONE + assert runner.multistream_overlap_shared_expert is False + assert fused_moe_310_module._MoECommMethods[MoECommType.ALLGATHER] is comm_method + parent_init.assert_called_once() + + +@pytest.mark.parametrize( + "is_v024, expected_contiguous", + [(True, True), (False, False)], +) +def test_process_weights_after_loading_310_uses_version_specific_layout( + monkeypatch, + is_v024, + expected_contiguous, +): + method = AscendUnquantizedFusedMoEMethod310.__new__(AscendUnquantizedFusedMoEMethod310) + method._maybe_pad_weight = MagicMock(side_effect=lambda weight: weight) + layer = _build_weight_layer() + original_w13 = layer.w13_weight.detach().clone() + original_w2 = layer.w2_weight.detach().clone() + + monkeypatch.setattr( + fused_moe_310_module, + "vllm_version_is", + lambda version: is_v024 and version == "0.24.0", + ) + monkeypatch.setattr(fused_moe_310_module, "maybe_trans_nz", lambda weight: weight) + monkeypatch.setattr( + fused_moe_310_module.UnquantizedFusedMoEMethod, + "process_weights_after_loading", + lambda self, layer: None, + ) + + method.process_weights_after_loading(layer) + + torch.testing.assert_close(layer.w13_weight, original_w13.transpose(1, 2)) + torch.testing.assert_close(layer.w2_weight, original_w2.transpose(1, 2)) + assert layer.w13_weight.is_contiguous() is expected_contiguous + assert layer.w2_weight.is_contiguous() is expected_contiguous + + +class _Projection(nn.Module): + def forward(self, hidden_states): + return hidden_states * 2.0 + 1.0, None + + +class _Gate(nn.Module): + def forward(self, hidden_states): + return torch.zeros((*hidden_states.shape[:-1], 1), dtype=hidden_states.dtype), None + + +@pytest.mark.parametrize("with_gate", [False, True]) +def test_shared_experts_part2_310_applies_optional_gate(with_gate): + runner = _build_runner() + runner._shared_experts = SimpleNamespace( + act_fn=nn.Identity(), + down_proj=_Projection(), + expert_gate=_Gate() if with_gate else None, + ) + hidden_states = torch.randn(3, 4) + shared_gate_up = torch.randn(3, 4) + + output = runner._shared_experts_part2(hidden_states, shared_gate_up) + + expected = shared_gate_up * 2.0 + 1.0 + if with_gate: + expected = expected * 0.5 torch.testing.assert_close(output, expected) -def test_forward_impl_with_shared_experts_returns_tuple_310(): - layer = _build_layer(_DummySharedExperts(with_gate=True)) - hidden_states = torch.randn(3, 8) - router_logits = torch.randn(3, 8) - routed_out = torch.randn(3, 8) - - with patch.object(AscendFusedMoE310, "forward_impl", return_value=routed_out): - shared_out, routed = layer.shared_forward_impl(hidden_states, router_logits) - - expected_shared = 0.5 * (hidden_states * 2.0 + 1.0) - torch.testing.assert_close(shared_out, expected_shared) - torch.testing.assert_close(routed, routed_out) - - -def test_forward_impl_without_shared_experts_integration_310(): - layer = _build_layer(None) - hidden_states = torch.randn(3, 8) - assert layer._forward_shared_experts(hidden_states) is None - - -def test_forward_impl_without_shared_experts_returns_routed_only_310(): - layer = _build_layer(None) - hidden_states = torch.randn(3, 8) - router_logits = torch.randn(3, 8) - routed_out = torch.randn(3, 8) - - with patch.object(AscendFusedMoE310, "forward_impl", return_value=routed_out): - output = layer.shared_forward_impl(hidden_states, router_logits) - - torch.testing.assert_close(output, routed_out) - - -def test_is_internal_router_is_false_310(): - layer = _build_layer(_DummySharedExperts(with_gate=True)) - assert layer.is_internal_router is False +@pytest.mark.parametrize("has_shared_experts", [False, True]) +def test_shared_forward_impl_310_returns_current_runner_contract(monkeypatch, has_shared_experts): + runner = _build_runner() + runner._shared_experts = object() if has_shared_experts else None + hidden_states = torch.randn(2, 4) + router_logits = torch.randn(2, 3) + routed_out = torch.randn(2, 4) + shared_out = torch.randn(2, 4) + routed_result = SimpleNamespace( + routed_out=routed_out, + before_dispatch_evt=None, + before_gmm2_evt=None, + before_combine_evt=None, + swiglu_limit=0.0, + ) + runner.no_shared_forward_impl = MagicMock(return_value=routed_result) + runner._forward_shared_experts = MagicMock(return_value=shared_out) + current_stream = MagicMock() + + monkeypatch.setattr(AscendMoERunner310, "is_internal_router", property(lambda _: False)) + monkeypatch.setattr(fused_moe_310_module.torch.npu, "current_stream", lambda: current_stream) + + result = runner.shared_forward_impl(hidden_states, router_logits) + + runner.no_shared_forward_impl.assert_called_once_with( + hidden_states, + router_logits, + return_with_event=True, + ) + if has_shared_experts: + assert result[0] is shared_out + assert result[1] is routed_out + runner._forward_shared_experts.assert_called_once() + else: + assert result is routed_out + runner._forward_shared_experts.assert_not_called() diff --git a/tests/ut/_310p/quantization/test_modelslim_config_310.py b/tests/ut/_310p/quantization/test_modelslim_config_310.py index 2aa88abdc8b..d0c755194e4 100644 --- a/tests/ut/_310p/quantization/test_modelslim_config_310.py +++ b/tests/ut/_310p/quantization/test_modelslim_config_310.py @@ -15,6 +15,7 @@ from unittest.mock import MagicMock, patch +from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig, FusedMoEParallelConfig from vllm.model_executor.layers.linear import LinearBase @@ -22,12 +23,6 @@ from vllm_ascend._310p.fused_moe.fused_moe import AscendUnquantizedFusedMoEMethod310 from vllm_ascend._310p.quantization.modelslim_config import AscendModelSlimConfig310 from vllm_ascend.ops.linear import AscendUnquantizedLinearMethod -from vllm_ascend.utils import vllm_version_is - -if vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe import FusedMoE -else: - from vllm.model_executor.layers.fused_moe import RoutedExperts class TestAscendModelSlimConfig310(TestBase): @@ -96,11 +91,7 @@ def test_get_quant_method_maps_lm_head_prefix_310(self): ) def test_get_quant_method_for_fused_moe_310(self): - if vllm_version_is("0.23.0"): - fused_moe_cls = FusedMoE - else: - fused_moe_cls = RoutedExperts - fused_moe_layer = MagicMock(spec=fused_moe_cls) + fused_moe_layer = MagicMock(spec=RoutedExperts) fused_moe_layer.moe = MagicMock(spec=FusedMoEConfig) fused_moe_layer.moe_config = MagicMock(spec=FusedMoEConfig) fused_moe_layer.moe_config.moe_backend = "auto" diff --git a/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py b/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py index d70188a06c9..bd9c9bc6980 100644 --- a/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py +++ b/tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py @@ -72,7 +72,7 @@ def test_init_accepts_model_argument(): def test_init_passes_model_to_super(): captured: dict = {} - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): def fake_init_v0(self, config, parallel_config, model=None): captured["args"] = (config, parallel_config, model) diff --git a/tests/ut/eplb/core/a2/test_eplb_utils.py b/tests/ut/eplb/core/a2/test_eplb_utils.py index 5a9f8fde756..95b50fd614f 100644 --- a/tests/ut/eplb/core/a2/test_eplb_utils.py +++ b/tests/ut/eplb/core/a2/test_eplb_utils.py @@ -9,7 +9,6 @@ from vllm_ascend.ascend_config import init_ascend_config from vllm_ascend.eplb.core.eplb_utils import generate_log2phy_map, init_eplb_config -from vllm_ascend.utils import vllm_version_is # isort: on @@ -26,36 +25,21 @@ def setUp(self, mock_fix_incompatible_config): from vllm.model_executor.layers.fused_moe.config import RoutingMethodType moe_parallel_config = FusedMoEParallelConfig(2, 0, 1, 2, 1, 1, 1, 1, 1, True, "hccl", enable_eplb=True) - if vllm_version_is("0.23.0"): - moe_config = FusedMoEConfig( - num_experts=8, - experts_per_token=8, - hidden_dim=8192, - intermediate_size_per_partition=5, - num_local_experts=8, - num_logical_experts=8, - activation="silu", - device="npu", - routing_method=RoutingMethodType.Simulated, - moe_parallel_config=moe_parallel_config, - in_dtype=torch.float16, - ) - else: - from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.activation import MoEActivation - moe_config = FusedMoEConfig( - num_experts=8, - experts_per_token=8, - hidden_dim=8192, - intermediate_size=10, - num_local_experts=8, - num_logical_experts=8, - activation=MoEActivation.SILU, - device="npu", - routing_method=RoutingMethodType.Simulated, - moe_parallel_config=moe_parallel_config, - in_dtype=torch.float16, - ) + moe_config = FusedMoEConfig( + num_experts=8, + experts_per_token=8, + hidden_dim=8192, + intermediate_size=10, + num_local_experts=8, + num_logical_experts=8, + activation=MoEActivation.SILU, + device="npu", + routing_method=RoutingMethodType.Simulated, + moe_parallel_config=moe_parallel_config, + in_dtype=torch.float16, + ) moe_config.supports_eplb = True self.vllm_config = vllm_config self.moe_config = moe_config diff --git a/tests/ut/ops/test_fused_moe.py b/tests/ut/ops/test_fused_moe.py index a57455368ed..740603024a5 100644 --- a/tests/ut/ops/test_fused_moe.py +++ b/tests/ut/ops/test_fused_moe.py @@ -1,947 +1,228 @@ -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# This file is a part of the vllm-ascend project. -# -import ast -import inspect -import textwrap +# SPDX-License-Identifier: Apache-2.0 from types import SimpleNamespace -from typing import TypedDict -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest import torch -import torch.nn as nn -import torch.nn.functional as F -from pytest_mock import MockerFixture +from torch import nn from vllm_ascend.ascend_forward_context import MoECommType from vllm_ascend.ops.fused_moe import fused_moe as fused_moe_module -from vllm_ascend.ops.fused_moe.moe_comm_method import FusedExpertsResult -from vllm_ascend.ops.fused_moe.moe_runtime_args import ( - MoEMlpComputeInput, - MoEPrepareOutput, - MoEQuantParams, - MoEWeights, +from vllm_ascend.ops.fused_moe.fused_moe import ( + AscendMoERunner, + AscendUnquantizedFusedMoEMethod, ) from vllm_ascend.quantization.quant_type import QuantType -from vllm_ascend.utils import AscendDeviceType, adapt_patch, vllm_version_is - -if vllm_version_is("0.23.0"): - from vllm_ascend.ops.fused_moe import fused_moe_0_23_0 as fused_moe_legacy_module - from vllm_ascend.ops.fused_moe.fused_moe import ( - AscendFusedMoE, - AscendMoERunner, - AscendUnquantizedFusedMoEMethod, - ) - adapt_patch(True) -else: - pytest.skip( - "Legacy AscendFusedMoE UTs are only for vLLM 0.23.0.", - allow_module_level=True, + +def _build_weight_layer(): + return SimpleNamespace( + w13_weight=nn.Parameter(torch.randn(2, 3, 4)), + w2_weight=nn.Parameter(torch.randn(2, 4, 3)), ) -def mock_ep_and_mc2_group(mocker): - mock_group = mocker.MagicMock() - mock_group.rank_in_group = 0 - mock_group.rank = 0 - mock_group.world_size = 4 - mock_group.device_group = "mock_group_ep" - mock_group.all_to_all = MagicMock(return_value=torch.randn(8, 8)) - return mock_group - - -def mock_dp_and_tp_group(mocker): - mock_group = mocker.MagicMock() - mock_group.rank_in_group = 0 - mock_group.world_size = 2 - mock_group.device_group = "mock_group" - mock_group.all_gather = MagicMock(return_value=torch.randn(10, 32)) - return mock_group - - -def mock_npu_format_cast(weight_data, format): - return weight_data - - -def build_mlp_compute_input_fixture( - *, - hidden_states: torch.Tensor, - w1: torch.Tensor | list[torch.Tensor], - w2: torch.Tensor | list[torch.Tensor], - group_list: torch.Tensor, - with_quant: bool, - group_list_type: int = 1, - dynamic_scale: torch.Tensor | None = None, - topk_scales: torch.Tensor | None = None, - w1_scale: torch.Tensor | list[torch.Tensor] | None = None, - w2_scale: torch.Tensor | list[torch.Tensor] | None = None, - w1_scale_bias: torch.Tensor | None = None, - w2_scale_bias: torch.Tensor | None = None, - w1_offset: torch.Tensor | None = None, - w2_offset: torch.Tensor | None = None, - fusion: bool = False, - activation: str = "silu", - need_trans: bool = True, - dynamic_eplb: bool = False, -) -> MoEMlpComputeInput: - return MoEMlpComputeInput( - hidden_states=hidden_states, - group_list=group_list, - group_list_type=group_list_type, - dynamic_scale=dynamic_scale, - topk_scales=topk_scales, - weights=MoEWeights( - w1=w1, - w2=w2, - w1_scale=w1_scale, - w2_scale=w2_scale, - w1_scale_bias=w1_scale_bias, - w2_scale_bias=w2_scale_bias, - w1_offset=w1_offset, - w2_offset=w2_offset, - ), - quant=MoEQuantParams(quant_type=QuantType.W8A8 if with_quant else QuantType.NONE), - fusion=fusion, - activation=activation, - need_trans=need_trans, - dynamic_eplb=dynamic_eplb, +def _build_apply_layer(): + return SimpleNamespace( + w13_weight=nn.Parameter(torch.randn(4, 3, 8)), + w2_weight=nn.Parameter(torch.randn(4, 8, 3)), + w13_bias=None, + w2_bias=None, + zero_expert_num=0, + zero_expert_type=None, + n_shared_experts=0, + swiglu_limit=0.0, ) -@pytest.fixture(autouse=True) -def setup_vllm_config_mock(mocker: MockerFixture): - mock_hf_config = MagicMock() - mock_hf_config.model_type = "llama" +def _build_unquantized_method(*, dynamic_eplb: bool = False): + method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) + method.dynamic_eplb = dynamic_eplb + method.tid2eid = None + method.moe = SimpleNamespace(has_bias=False) + method._maybe_pad_weight = MagicMock(side_effect=lambda weight: weight) + return method - mock_model_config = MagicMock() - mock_model_config.hf_config = mock_hf_config - mock_vllm_config = MagicMock() - mock_vllm_config.model_config = mock_model_config - mock_vllm_config.parallel_config = MagicMock(tensor_parallel_size=2) - mock_vllm_config.scheduler_config = MagicMock(max_num_seqs=4) - mock_vllm_config.model_config.max_model_len = 2048 +def test_ascend_unquantized_skips_upstream_modular_kernel_init(): + method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) - mocker.patch("vllm_ascend.ops.fused_moe.fused_moe.get_current_vllm_config", return_value=mock_vllm_config) + assert method.maybe_make_prepare_finalize() is None -@pytest.fixture -def mock_dist_env(mocker: MockerFixture): - mock_moe_comm_method = MagicMock() +@pytest.mark.parametrize( + "is_v024, expected_contiguous", + [(True, True), (False, False)], +) +def test_process_weights_after_loading_uses_version_specific_layout( + monkeypatch, + is_v024, + expected_contiguous, +): + method = _build_unquantized_method() + layer = _build_weight_layer() + original_w13 = layer.w13_weight.detach().clone() + original_w2 = layer.w2_weight.detach().clone() + ascend_config = SimpleNamespace(enable_fused_mc2=False) + + monkeypatch.setattr(fused_moe_module, "vllm_version_is", lambda version: is_v024 and version == "0.24.0") + monkeypatch.setattr(fused_moe_module, "get_ascend_config", lambda: ascend_config) + monkeypatch.setattr(fused_moe_module, "maybe_trans_nz", lambda weight: weight) + upstream_method_base = AscendUnquantizedFusedMoEMethod.__mro__[2] + monkeypatch.setattr( + upstream_method_base, + "process_weights_after_loading", + lambda self, layer: None, + raising=False, + ) - def mock_prepare(hidden_states, router_logits, **kwargs): - return MoEPrepareOutput( - hidden_states=hidden_states, - router_logits=router_logits, - mc2_mask=kwargs.get("mc2_mask"), - padded_hidden_states_shape=None, - pertoken_scale=None, - ) + method.process_weights_after_loading(layer) + + torch.testing.assert_close(layer.w13_weight, original_w13.transpose(1, 2)) + torch.testing.assert_close(layer.w2_weight, original_w2.transpose(1, 2)) + assert layer.w13_weight.is_contiguous() is expected_contiguous + assert layer.w2_weight.is_contiguous() is expected_contiguous + + +@pytest.mark.parametrize("moe_comm_type", [MoECommType.ALLGATHER, MoECommType.FUSED_MC2]) +def test_unquantized_apply_builds_current_fused_experts_input(monkeypatch, moe_comm_type): + method = _build_unquantized_method() + layer = _build_apply_layer() + hidden_states = torch.randn(2, 3, dtype=torch.float16) + topk_weights = torch.tensor([[0.25, 0.75], [0.6, 0.4]], dtype=torch.float32) + topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64) + routed_out = torch.ones_like(hidden_states) + moe_comm_method = MagicMock() + moe_comm_method.fused_experts.return_value = routed_out + + monkeypatch.setattr( + fused_moe_module, + "_EXTRA_CTX", + SimpleNamespace(moe_comm_type=moe_comm_type, moe_comm_method=moe_comm_method), + ) + monkeypatch.setattr(fused_moe_module, "get_moe_num_logical_experts", lambda *args, **kwargs: 4) + monkeypatch.setattr(fused_moe_module, "get_forward_context", lambda: SimpleNamespace(input_ids=None)) + monkeypatch.setattr(fused_moe_module, "get_current_vllm_config", lambda: None) + select_experts = MagicMock(return_value=(topk_weights, topk_ids)) + monkeypatch.setattr(fused_moe_module, "select_experts", select_experts) + + result = method.apply( + layer=layer, + x=hidden_states, + use_grouped_topk=False, + top_k=2, + router_logits=torch.randn(2, 4), + renormalize=True, + num_experts=4, + apply_router_weight_on_input=True, + activation="gelu", + ) + + assert result is routed_out + select_experts.assert_called_once() + fused_input = moe_comm_method.fused_experts.call_args.kwargs["fused_experts_input"] + assert fused_input.hidden_states is hidden_states + torch.testing.assert_close(fused_input.topk_weights, topk_weights.to(hidden_states.dtype)) + assert torch.equal(fused_input.topk_ids, topk_ids) + assert fused_input.routing.apply_router_weight_on_input + assert fused_input.activation == "gelu" + assert fused_input.quant.quant_type == QuantType.NONE + if moe_comm_type == MoECommType.FUSED_MC2: + assert fused_input.weights.w1[0] is layer.w13_weight + assert fused_input.weights.w2[0] is layer.w2_weight + else: + assert fused_input.weights.w1 is layer.w13_weight + assert fused_input.weights.w2 is layer.w2_weight + + +@pytest.mark.parametrize( + "moe_comm_type, flash_comm_v1_enabled, expected", + [ + (MoECommType.ALLTOALL, False, True), + (MoECommType.MC2, False, True), + (MoECommType.FUSED_MC2, False, True), + (MoECommType.ALLGATHER, False, False), + (MoECommType.ALLGATHER, True, True), + ], +) +def test_runner_reduction_contract(monkeypatch, moe_comm_type, flash_comm_v1_enabled, expected): + runner = AscendMoERunner.__new__(AscendMoERunner) + shared_output = object() + monkeypatch.setattr( + fused_moe_module, + "_EXTRA_CTX", + SimpleNamespace(moe_comm_type=moe_comm_type, flash_comm_v1_enabled=flash_comm_v1_enabled), + ) - mock_moe_comm_method.prepare.side_effect = mock_prepare + assert runner.use_dp_chunking is False + assert runner._fused_output_is_reduced is expected + assert runner._maybe_reduce_shared_expert_output(shared_output) is shared_output - mock_fused_experts_result = torch.randn(16, 2) - mock_moe_comm_method.fused_experts.return_value = mock_fused_experts_result - def mock_finalize(hidden_states, **kwargs): - return hidden_states +class _Projection(nn.Module): + def forward(self, hidden_states): + return hidden_states * 2.0 + 1.0, None - mock_moe_comm_method.finalize.side_effect = mock_finalize - dp_metadata = MagicMock(num_tokens_across_dp_cpu=[5, 5]) - mock_weight_prefetch_method = MagicMock() - mock_forward_context_obj = MagicMock( - moe_comm_method=mock_moe_comm_method, - moe_comm_type=MoECommType.MC2, - max_tokens_across_dp=10, - dp_metadata=dp_metadata, - mc2_mask=torch.zeros(16, dtype=torch.bool), - padded_num_tokens=16, - with_quant=False, - ) - with ( - patch("torch.distributed.get_rank", return_value=0), - patch("torch.distributed.get_world_size", return_value=4), - patch("vllm_ascend.ops.fused_moe.fused_moe.get_ep_group", return_value=mock_ep_and_mc2_group(mocker)), - patch("vllm_ascend.ops.fused_moe.token_dispatcher.get_ep_group", return_value=mock_ep_and_mc2_group(mocker)), - patch("vllm_ascend.ops.fused_moe.fused_moe.get_mc2_group", return_value=mock_ep_and_mc2_group(mocker)), - patch("vllm_ascend.ops.fused_moe.fused_moe.get_tp_group", return_value=mock_dp_and_tp_group(mocker)), - patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_dp_and_tp_group(mocker)), - patch("vllm_ascend.ops.fused_moe.fused_moe.get_dp_group", return_value=mock_dp_and_tp_group(mocker)), - patch("vllm.model_executor.layers.fused_moe.layer.get_dp_group", return_value=mock_dp_and_tp_group(mocker)), - patch("vllm.model_executor.layers.fused_moe.config.get_dp_group", return_value=mock_dp_and_tp_group(mocker)), - patch( - "vllm_ascend.ops.fused_moe.fused_moe.get_ascend_config", - return_value=MagicMock(enable_multistream_moe=False, expert_map_path=None), - ), - patch( - "vllm_ascend.ops.fused_moe.fused_moe.init_eplb_config", - return_value=(torch.tensor([0, 1, 2, -1, -1, -1, -1, -1]), None, 0), - ), - patch("vllm_ascend.ops.fused_moe.fused_moe.get_forward_context", return_value=mock_forward_context_obj), - patch("vllm_ascend.ascend_forward_context.get_forward_context", return_value=mock_forward_context_obj), - patch("vllm_ascend.utils.get_ascend_device_type", return_value=AscendDeviceType.A3), - patch("vllm_ascend.ops.fused_moe.moe_comm_method.MC2CommImpl._get_token_dispatcher", return_value=None), - patch("vllm_ascend.ops.fused_moe.moe_comm_method.AlltoAllCommImpl._get_token_dispatcher", return_value=None), - patch("vllm_ascend.ops.fused_moe.moe_comm_method.AllGatherCommImpl._get_token_dispatcher", return_value=None), - patch( - "vllm_ascend.ops.fused_moe.experts_selector.get_weight_prefetch_method", - return_value=mock_weight_prefetch_method, - ), - ): - yield { - "mock_forward_context_obj": mock_forward_context_obj, - "mock_moe_comm_method": mock_moe_comm_method, - } - - -@pytest.fixture -def default_moe_config(): - return {"num_experts": 8, "top_k": 2, "hidden_size": 512, "intermediate_size": 1024} - - -@pytest.fixture -def moe_method(mock_dist_env): - moe = MagicMock() - moe.moe_parallel_config.return_value = MagicMock(ep_size=4) - moe.moe_parallel_config.use_ep = False - moe.moe_parallel_config.dp_size = 1 - return AscendUnquantizedFusedMoEMethod(moe) +class _Gate(nn.Module): + def forward(self, hidden_states): + return torch.zeros((*hidden_states.shape[:-1], 1), dtype=hidden_states.dtype), None -def test_ascend_unquantized_skips_upstream_modular_kernel_init(): - method = AscendUnquantizedFusedMoEMethod.maybe_make_prepare_finalize - - assert method(object()) is None - - -class Device(TypedDict): - device_id: int - device_expert: list[int] - - -class Layer(TypedDict): - layer_id: int - device_count: int - device_list: list[Device] - - -class MockData(TypedDict): - moe_layer_count: int - layer_list: list[Layer] - - -class MockQuantMethod(nn.Module): - def __init__(self, shared_experts, num_tokens): - super().__init__() - if shared_experts: - self.apply = MagicMock(return_value=(torch.randn(num_tokens, 32), torch.randn(num_tokens, 10))) - else: - self.apply = MagicMock(return_value=(torch.randn(num_tokens, 32))) - - -def _drop_self(signature: inspect.Signature) -> list[inspect.Parameter]: - params = list(signature.parameters.values()) - if params and params[0].name == "self": - return params[1:] - return params - - -def _format_signature_mismatch(method_name: str, issues: list[str]) -> str: - return f"{method_name} signature is not aligned with vLLM parent: " + "; ".join(issues) - - -def _assert_child_signature_accepts_parent_interface(child_method, parent_method): - child_params = _drop_self(inspect.signature(child_method)) - parent_params = _drop_self(inspect.signature(parent_method)) - child_by_name = { - param.name: param - for param in child_params - if param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) - } - child_has_var_positional = any(param.kind == inspect.Parameter.VAR_POSITIONAL for param in child_params) - child_has_var_keyword = any(param.kind == inspect.Parameter.VAR_KEYWORD for param in child_params) - issues: list[str] = [] - - for parent_param in parent_params: - if parent_param.kind == inspect.Parameter.VAR_POSITIONAL: - if not child_has_var_positional: - issues.append("child is missing *args from parent") - continue - - if parent_param.kind == inspect.Parameter.VAR_KEYWORD: - if not child_has_var_keyword: - issues.append("child is missing **kwargs from parent") - continue - - child_param = child_by_name.get(parent_param.name) - if child_param is None: - if parent_param.kind == inspect.Parameter.KEYWORD_ONLY: - if not child_has_var_keyword: - issues.append(f"missing keyword-only parameter {parent_param.name!r}") - elif not child_has_var_positional and not child_has_var_keyword: - issues.append(f"missing parameter {parent_param.name!r}") - continue - - if parent_param.kind != child_param.kind: - issues.append( - f"parameter {parent_param.name!r} has kind {child_param.kind!s}, expected {parent_param.kind!s}" - ) - - parent_param_names = {param.name for param in parent_params} - for child_param in child_params: - if child_param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): - continue - if child_param.name in parent_param_names: - continue - if child_param.default is inspect.Parameter.empty: - issues.append(f"extra parameter {child_param.name!r} must be optional") - - assert not issues, _format_signature_mismatch(parent_method.__qualname__, issues) - - -def _method_uses_super(method) -> bool: - try: - source = inspect.getsource(method) - except (OSError, TypeError): - return False - - tree = ast.parse(textwrap.dedent(source)) - return any( - isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "super" - for node in ast.walk(tree) +@pytest.mark.parametrize("with_gate", [False, True]) +def test_shared_experts_part2_applies_optional_gate(with_gate): + runner = AscendMoERunner.__new__(AscendMoERunner) + nn.Module.__init__(runner) + runner._shared_experts = SimpleNamespace( + act_fn=nn.Identity(), + down_proj=_Projection(), + expert_gate=_Gate() if with_gate else None, ) + hidden_states = torch.randn(3, 4) + shared_gate_up = torch.randn(3, 4) + + output = runner._shared_experts_part2(hidden_states, shared_gate_up) + + expected = shared_gate_up * 2.0 + 1.0 + if with_gate: + expected = expected * 0.5 + torch.testing.assert_close(output, expected) + + +@pytest.mark.parametrize("has_shared_experts", [False, True]) +def test_shared_forward_impl_returns_current_runner_contract(monkeypatch, has_shared_experts): + runner = AscendMoERunner.__new__(AscendMoERunner) + nn.Module.__init__(runner) + runner._shared_experts = object() if has_shared_experts else None + hidden_states = torch.randn(2, 4) + router_logits = torch.randn(2, 3) + routed_out = torch.randn(2, 4) + shared_out = torch.randn(2, 4) + routed_result = SimpleNamespace( + routed_out=routed_out, + before_dispatch_evt=None, + before_gmm2_evt=None, + before_combine_evt=None, + swiglu_limit=0.0, + ) + runner.no_shared_forward_impl = MagicMock(return_value=routed_result) + runner._forward_shared_experts = MagicMock(return_value=shared_out) + current_stream = MagicMock() + monkeypatch.setattr(AscendMoERunner, "is_internal_router", property(lambda _: False)) + monkeypatch.setattr(fused_moe_module.torch.npu, "current_stream", lambda: current_stream) -class TestVllmParentInterfaceCompatibility: - @pytest.mark.parametrize( - "child_cls,parent_cls,method_name", - [ - (AscendUnquantizedFusedMoEMethod, fused_moe_module.UnquantizedFusedMoEMethod, "__init__"), - ( - AscendUnquantizedFusedMoEMethod, - fused_moe_module.UnquantizedFusedMoEMethod, - "process_weights_after_loading", - ), - (AscendUnquantizedFusedMoEMethod, fused_moe_module.UnquantizedFusedMoEMethod, "apply"), - (AscendMoERunner, fused_moe_module.MoERunner, "__init__"), - (AscendMoERunner, fused_moe_module.MoERunner, "forward_impl"), - (AscendMoERunner, fused_moe_module.MoERunner, "_forward_impl"), - (AscendFusedMoE, fused_moe_module.FusedMoE, "__init__"), - (AscendFusedMoE, fused_moe_module.FusedMoE, "forward"), - (AscendFusedMoE, fused_moe_module.FusedMoE, "forward_impl"), - (AscendFusedMoE, fused_moe_module.FusedMoE, "maybe_all_reduce_tensor_model_parallel"), - ], - ) - def test_overridden_method_signature_accepts_parent_interface(self, child_cls, parent_cls, method_name): - child_method = getattr(child_cls, method_name) - if not _method_uses_super(child_method): - pytest.skip( - f"{child_cls.__name__}.{method_name} does not call " - "super(), so parent interface alignment is not " - "required" - ) - - if not hasattr(parent_cls, method_name): - pytest.fail( - f"{child_cls.__name__}.{method_name} calls super(), but {parent_cls.__name__} has no {method_name}" - ) - - _assert_child_signature_accepts_parent_interface( - child_method, - getattr(parent_cls, method_name), - ) - - -class TestAscendUnquantizedFusedMoEMethod: - def _build_layer(self, *, has_bias=True, zero_expert_num=0): - layer = MagicMock() - layer.w13_weight = nn.Parameter(torch.randn(2, 3, 4)) - layer.w2_weight = nn.Parameter(torch.randn(2, 4, 3)) - layer.w13_bias = torch.randn(2, 4) if has_bias else None - layer.w2_bias = torch.randn(2, 3) if has_bias else None - layer.zero_expert_num = zero_expert_num - layer.zero_expert_type = "identity" if zero_expert_num > 0 else None - layer.n_shared_experts = 0 - layer.moe_config = SimpleNamespace(num_logical_experts=None) - layer.layer_id = 3 - layer.vllm_config = SimpleNamespace(model_config=SimpleNamespace(enable_return_routed_experts=False)) - return layer - - @pytest.mark.parametrize("enable_fused_mc2", [True, False]) - def test_process_weights_after_loading_transposes_and_formats(self, monkeypatch, enable_fused_mc2): - method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) - method.dynamic_eplb = False - method._maybe_pad_weight = MagicMock(side_effect=lambda weight: weight) - layer = self._build_layer() - original_w13 = layer.w13_weight.detach().clone() - original_w2 = layer.w2_weight.detach().clone() - format_cast = MagicMock(side_effect=lambda weight, _: weight) - maybe_trans_nz = MagicMock(side_effect=lambda weight: weight) - - mock_ascend_config = MagicMock() - mock_ascend_config.enable_fused_mc2 = enable_fused_mc2 - monkeypatch.setattr(fused_moe_module, "get_ascend_config", lambda: mock_ascend_config) - monkeypatch.setattr(fused_moe_module.torch_npu, "npu_format_cast", format_cast) - monkeypatch.setattr(fused_moe_module, "maybe_trans_nz", maybe_trans_nz) - - method.process_weights_after_loading(layer) - - torch.testing.assert_close(layer.w13_weight, original_w13.transpose(1, 2).contiguous()) - torch.testing.assert_close(layer.w2_weight, original_w2.transpose(1, 2).contiguous()) - if enable_fused_mc2: - assert format_cast.call_count == 2 - maybe_trans_nz.assert_not_called() - else: - assert maybe_trans_nz.call_count == 2 - format_cast.assert_not_called() - - def test_process_weights_after_loading_splits_dynamic_eplb_fused_mc2_weights(self, monkeypatch): - method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) - method.dynamic_eplb = True - method._maybe_pad_weight = MagicMock(side_effect=lambda weight: weight) - layer = nn.Module() - layer.w13_weight = nn.Parameter(torch.randn(2, 3, 4)) - layer.w2_weight = nn.Parameter(torch.randn(2, 4, 3)) - expected_w13 = layer.w13_weight.detach().clone().transpose(1, 2).contiguous() - expected_w2 = layer.w2_weight.detach().clone().transpose(1, 2).contiguous() - format_cast = MagicMock(side_effect=lambda weight, _: weight) - empty_cache = MagicMock() - - mock_ascend_config = MagicMock() - mock_ascend_config.enable_fused_mc2 = True - monkeypatch.setattr(fused_moe_module, "get_ascend_config", lambda: mock_ascend_config) - monkeypatch.setattr(fused_moe_module.torch_npu, "npu_format_cast", format_cast) - monkeypatch.setattr(fused_moe_module.torch, "npu", SimpleNamespace(empty_cache=empty_cache), raising=False) - - method.process_weights_after_loading(layer) - - assert "w13_weight" not in layer._parameters - assert "w2_weight" not in layer._parameters - assert len(layer.w13_weight_list) == 2 - assert len(layer.w2_weight_list) == 2 - torch.testing.assert_close(layer.w13_weight_list[0], expected_w13[0]) - torch.testing.assert_close(layer.w2_weight_list[1], expected_w2[1]) - assert layer.w13_weight_list[0].untyped_storage().data_ptr() != expected_w13[0].untyped_storage().data_ptr() - assert format_cast.call_count == 2 - empty_cache.assert_called_once() - - @pytest.mark.parametrize("moe_comm_type", [MoECommType.MC2, MoECommType.FUSED_MC2]) - def test_apply_builds_fused_experts_input(self, monkeypatch, moe_comm_type): - method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) - method.moe = SimpleNamespace(has_bias=True) - method.dynamic_eplb = False - method.tid2eid = None - layer = self._build_layer(has_bias=True) - hidden_states = torch.randn(2, 4, dtype=torch.float16) - router_logits = torch.randn(2, 4) - topk_weights = torch.tensor([[0.25, 0.75], [0.6, 0.4]], dtype=torch.float32) - topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64) - moe_comm_method = MagicMock() - moe_comm_method.fused_experts.return_value = torch.ones_like(hidden_states) - monkeypatch.setattr( - fused_moe_module, - "_EXTRA_CTX", - SimpleNamespace(moe_comm_type=moe_comm_type, moe_comm_method=moe_comm_method), - ) - select_experts_mock = MagicMock(return_value=(topk_weights, topk_ids)) - monkeypatch.setattr(fused_moe_module, "select_experts", select_experts_mock) - monkeypatch.setattr(fused_moe_module, "get_forward_context", MagicMock(return_value=MagicMock(input_ids=None))) - - result = method.apply( - layer=layer, - x=hidden_states, - use_grouped_topk=False, - top_k=2, - router_logits=router_logits, - renormalize=True, - num_experts=4, - apply_router_weight_on_input=True, - activation="gelu", - pertoken_scale=torch.ones(2), - mc2_mask=torch.tensor([True, False]), - ) - - torch.testing.assert_close(result, torch.ones_like(hidden_states)) - select_experts_mock.assert_called_once() - fused_input = moe_comm_method.fused_experts.call_args.kwargs["fused_experts_input"] - assert fused_input.hidden_states is hidden_states - torch.testing.assert_close(fused_input.topk_weights, topk_weights.to(hidden_states.dtype)) - assert torch.equal(fused_input.topk_ids, topk_ids) - assert fused_input.weights.w1_bias is layer.w13_bias - assert fused_input.weights.w2_bias is layer.w2_bias - assert fused_input.routing.apply_router_weight_on_input - assert fused_input.activation == "gelu" - if moe_comm_type == MoECommType.FUSED_MC2: - assert fused_input.weights.w1[0] is layer.w13_weight - assert fused_input.weights.w2[0] is layer.w2_weight - assert isinstance(fused_input.weights.w1_scale, list) - assert isinstance(fused_input.weights.w2_scale, list) - assert fused_input.weights.w1_scale[0].dtype == torch.int64 - assert fused_input.weights.w2_scale[0].dtype == torch.int64 - assert fused_input.weights.w1_scale_bias[0].dtype == torch.float32 - assert fused_input.weights.w2_scale_bias[0].dtype == torch.float32 - else: - assert fused_input.weights.w1 is layer.w13_weight - assert fused_input.weights.w2 is layer.w2_weight - assert fused_input.weights.w1_scale is None - assert fused_input.weights.w2_scale is None - - @pytest.mark.parametrize("moe_comm_type", [MoECommType.MC2, MoECommType.FUSED_MC2]) - def test_apply_uses_weight_lists_when_dynamic_eplb_splits_weights(self, monkeypatch, moe_comm_type): - method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) - method.moe = SimpleNamespace(has_bias=False) - method.dynamic_eplb = True - method.tid2eid = None - layer = self._build_layer(has_bias=False) - layer.w13_weight_list = [torch.randn(4, 6), torch.randn(4, 6)] - layer.w2_weight_list = [torch.randn(3, 4), torch.randn(3, 4)] - hidden_states = torch.randn(2, 4, dtype=torch.float16) - topk_weights = torch.ones(2, 2, dtype=torch.float32) - topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64) - moe_comm_method = MagicMock() - moe_comm_method.fused_experts.return_value = torch.ones_like(hidden_states) - monkeypatch.setattr( - fused_moe_module, - "_EXTRA_CTX", - SimpleNamespace(moe_comm_type=moe_comm_type, moe_comm_method=moe_comm_method), - ) - monkeypatch.setattr(fused_moe_module, "select_experts", MagicMock(return_value=(topk_weights, topk_ids))) - monkeypatch.setattr(fused_moe_module, "get_forward_context", MagicMock(return_value=MagicMock(input_ids=None))) - - method.apply( - layer=layer, - x=hidden_states, - use_grouped_topk=False, - top_k=2, - router_logits=torch.randn(2, 4), - renormalize=True, - num_experts=4, - ) - - fused_input = moe_comm_method.fused_experts.call_args.kwargs["fused_experts_input"] - assert fused_input.weights.w1 is layer.w13_weight_list - assert fused_input.weights.w2 is layer.w2_weight_list - if moe_comm_type == MoECommType.FUSED_MC2: - assert len(fused_input.weights.w1_scale) == 1 - assert len(fused_input.weights.w2_scale) == 1 - assert fused_input.weights.w1_scale[0].dtype == torch.int64 - assert fused_input.weights.w2_scale[0].dtype == torch.int64 - assert fused_input.weights.w1_scale[0].numel() == 0 - assert fused_input.weights.w2_scale[0].numel() == 0 - assert fused_input.weights.w1_scale_bias[0].dtype == torch.float32 - assert fused_input.weights.w2_scale_bias[0].dtype == torch.float32 - assert fused_input.weights.w1_scale_bias[0].numel() == 0 - assert fused_input.weights.w2_scale_bias[0].numel() == 0 - else: - assert fused_input.weights.w1_scale is None - assert fused_input.weights.w2_scale is None - - def test_apply_warns_when_dynamic_eplb_fused_mc2_weights_are_not_split(self, monkeypatch): - method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) - method.moe = SimpleNamespace(has_bias=False) - method.dynamic_eplb = True - method.tid2eid = None - layer = self._build_layer(has_bias=False) - hidden_states = torch.randn(2, 4, dtype=torch.float16) - topk_weights = torch.ones(2, 2, dtype=torch.float32) - topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64) - moe_comm_method = MagicMock() - moe_comm_method.fused_experts.return_value = torch.ones_like(hidden_states) - warning_once = MagicMock() - monkeypatch.setattr( - fused_moe_module, - "_EXTRA_CTX", - SimpleNamespace(moe_comm_type=MoECommType.FUSED_MC2, moe_comm_method=moe_comm_method), - ) - monkeypatch.setattr(fused_moe_module, "select_experts", MagicMock(return_value=(topk_weights, topk_ids))) - monkeypatch.setattr(fused_moe_module, "get_forward_context", MagicMock(return_value=MagicMock(input_ids=None))) - monkeypatch.setattr(fused_moe_module.logger, "warning_once", warning_once) - - method.apply( - layer=layer, - x=hidden_states, - use_grouped_topk=False, - top_k=2, - router_logits=torch.randn(2, 4), - renormalize=True, - num_experts=4, - ) - - warning_once.assert_called_once() - warning_msg = warning_once.call_args.args[0] - assert "dynamic EPLB" in warning_msg - assert "not split into tensor lists" in warning_msg - fused_input = moe_comm_method.fused_experts.call_args.kwargs["fused_experts_input"] - assert fused_input.weights.w1[0] is layer.w13_weight - assert fused_input.weights.w2[0] is layer.w2_weight + result = runner.shared_forward_impl(hidden_states, router_logits) - def test_apply_adds_zero_expert_result_and_force_balances(self, monkeypatch): - method = AscendUnquantizedFusedMoEMethod.__new__(AscendUnquantizedFusedMoEMethod) - method.moe = SimpleNamespace(has_bias=False) - method.dynamic_eplb = True - method.tid2eid = None - layer = self._build_layer(has_bias=False, zero_expert_num=1) - hidden_states = torch.randn(2, 4) - topk_weights = torch.ones(2, 2) - topk_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32) - zero_hidden = torch.full_like(hidden_states, 3.0) - routed_hidden = torch.full_like(hidden_states, 5.0) - expected = routed_hidden + zero_hidden - moe_comm_method = MagicMock() - moe_comm_method.fused_experts.return_value = routed_hidden - - monkeypatch.setattr( - fused_moe_module, - "_EXTRA_CTX", - SimpleNamespace(moe_comm_type=MoECommType.MC2, moe_comm_method=moe_comm_method), - ) - monkeypatch.setattr(fused_moe_module, "select_experts", MagicMock(return_value=(topk_weights, topk_ids))) - zero_experts_mock = MagicMock(return_value=(topk_ids, topk_weights, zero_hidden)) - monkeypatch.setattr(fused_moe_module, "zero_experts_compute", zero_experts_mock) - monkeypatch.setattr(torch, "rand", MagicMock(return_value=torch.tensor([[0.2, 0.1], [0.4, 0.3]]))) - monkeypatch.setattr(fused_moe_module, "get_forward_context", MagicMock(return_value=MagicMock(input_ids=None))) - - result = method.apply( - layer=layer, - x=hidden_states, - use_grouped_topk=False, - top_k=2, - router_logits=torch.randn(2, 2), - renormalize=False, - num_experts=2, - enable_force_load_balance=True, - ) - - torch.testing.assert_close(result, expected) - zero_experts_mock.assert_called_once() - fused_input = moe_comm_method.fused_experts.call_args.kwargs["fused_experts_input"] - assert fused_input.dynamic_eplb - assert fused_input.weights.w1_bias is None - assert fused_input.weights.w2_bias is None - - -class TestAscendMoERunner: - @pytest.mark.parametrize( - "moe_comm_type, flash_comm_v1_enabled, expected", - [ - (MoECommType.ALLTOALL, False, True), - (MoECommType.MC2, False, True), - (MoECommType.FUSED_MC2, False, True), - (MoECommType.ALLGATHER, False, False), - (MoECommType.ALLGATHER, True, True), - ], + runner.no_shared_forward_impl.assert_called_once_with( + hidden_states, + router_logits, + return_with_event=True, ) - def test_runner_reduction_properties(self, monkeypatch, moe_comm_type, flash_comm_v1_enabled, expected): - runner = AscendMoERunner.__new__(AscendMoERunner) - monkeypatch.setattr(fused_moe_legacy_module, "_EXTRA_CTX", SimpleNamespace(moe_comm_type=moe_comm_type)) - monkeypatch.setattr( - fused_moe_legacy_module, - "_EXTRA_CTX", - SimpleNamespace(moe_comm_type=moe_comm_type, flash_comm_v1_enabled=flash_comm_v1_enabled), - ) - - assert runner.use_dp_chunking is False - if hasattr(type(runner), "_fused_output_is_reduced"): - assert runner._fused_output_is_reduced is expected - if hasattr(runner, "_maybe_reduce_shared_expert_output"): - assert runner._maybe_reduce_shared_expert_output("shared") == "shared" - - @pytest.mark.parametrize("has_shared_experts", [False, True]) - def test_forward_impl_delegates_to_layer(self, monkeypatch, has_shared_experts): - runner = AscendMoERunner.__new__(AscendMoERunner) - shared_experts = MagicMock() if has_shared_experts else None - shared_experts_owner = next( - (cls for cls in type(runner).__mro__ if "shared_experts" in cls.__dict__), - AscendMoERunner, - ) - monkeypatch.setattr(shared_experts_owner, "shared_experts", property(lambda _: shared_experts), raising=False) - layer = MagicMock() - hidden_states = torch.randn(2, 4) - router_logits = torch.randn(2, 3) - layer.forward_impl.return_value = "routed" - layer.shared_forward_impl.return_value = ("shared", "routed") - - result = runner.forward_impl(layer, hidden_states, router_logits, None) - - if has_shared_experts: - assert result == ("shared", "routed") - layer.shared_forward_impl.assert_called_once_with(hidden_states, router_logits) - layer.forward_impl.assert_not_called() - else: - assert result == "routed" - layer.forward_impl.assert_called_once_with(hidden_states, router_logits) - layer.shared_forward_impl.assert_not_called() - - -class TestAscendFusedMoE: - def _build_layer(self): - layer = AscendFusedMoE.__new__(AscendFusedMoE) - layer.quant_method = MagicMock() - layer.ensure_moe_quant_config_init = MagicMock() - layer.runner = MagicMock() - layer.moe_load = torch.zeros(2, dtype=torch.int64) - layer.multi_stage = False - layer.log2phy = torch.tensor([1, 0]) - return layer - - def test_simple_helpers(self, monkeypatch): - layer = self._build_layer() - layer.quant_method.quant_method = SimpleNamespace(quant_type=QuantType.W8A8) - layer.update_expert_map(torch.tensor([0, -1])) - assert torch.equal(layer._expert_map, torch.tensor([0, -1])) - assert torch.equal(layer.get_log2phy_map(), torch.tensor([1, 0])) - assert layer._get_quant_type() == QuantType.W8A8 - - layer.clear_moe_load() - assert torch.equal(layer.moe_load, torch.zeros_like(layer.moe_load)) - layer.multi_stage = True - layer.load_counter = torch.tensor(4) - layer.clear_moe_load() - assert layer.load_counter.item() == 0 - - maybe_all_reduce = MagicMock(return_value="reduced") - monkeypatch.setattr( - fused_moe_module.torch.ops, - "vllm", - SimpleNamespace(maybe_all_reduce_tensor_model_parallel=maybe_all_reduce), - raising=False, - ) - assert layer.maybe_all_reduce_tensor_model_parallel(torch.ones(1)) == "reduced" - - def test_forward_delegates_to_runner(self): - layer = self._build_layer() - hidden_states = torch.randn(2, 4) - router_logits = torch.randn(2, 3) - layer.runner.forward.return_value = "forwarded" - - assert layer.forward(hidden_states, router_logits) == "forwarded" - layer.ensure_moe_quant_config_init.assert_called_once() - layer.runner.forward.assert_called_once_with(hidden_states, router_logits) - - @pytest.mark.parametrize("return_with_event", [True, False]) - def test_forward_impl_prepare_apply_finalize(self, monkeypatch, return_with_event): - layer = self._build_layer() - layer.enable_npugraph_ex_static_kernel = True - layer.enable_shared_expert_dp = False - layer.quant_type = QuantType.NONE - layer.top_k = 2 - layer.renormalize = True - layer.use_grouped_topk = False - layer.moe_config = SimpleNamespace(num_experts=4) - layer._expert_map = None - layer.topk_group = None - layer.num_expert_group = None - layer.custom_routing_function = None - layer.scoring_func = "softmax" - layer._original_routed_scaling_factor = 1.0 - layer.routed_scaling_factor = 1.0 - layer.e_score_correction_bias = None - layer.activation = "silu" - layer.apply_router_weight_on_input = False - layer.global_redundant_expert_num = 0 - layer.dynamic_eplb = True - layer.reduce_results = True - forward_context = SimpleNamespace(moe_layer_index=5, all_moe_layers=[0, 1]) - hidden_states = torch.randn(2, 4) - router_logits = torch.randn(2, 4) - prepared_hidden = hidden_states + 1 - prepared_logits = router_logits + 1 - prepare_output = MoEPrepareOutput( - hidden_states=prepared_hidden, - router_logits=prepared_logits, - mc2_mask=torch.tensor([True, False]), - padded_hidden_states_shape=torch.Size([4, 4]), - pertoken_scale=torch.ones(2), - ) - moe_comm_method = MagicMock() - moe_comm_method.prepare.return_value = prepare_output - moe_comm_method.finalize.side_effect = lambda hidden_states, **_: hidden_states + 2 - before_dispatch_evt = MagicMock() - before_combine_evt = MagicMock() - layer.quant_method.apply.return_value = FusedExpertsResult( - routed_out=torch.ones_like(hidden_states), - before_dispatch_evt=before_dispatch_evt, - before_combine_evt=before_combine_evt, - expert_tokens=torch.tensor([2, 5]), - group_list_type=0, - ) - monkeypatch.setattr(fused_moe_legacy_module, "get_forward_context", MagicMock(return_value=forward_context)) - monkeypatch.setattr( - fused_moe_legacy_module, - "_EXTRA_CTX", - SimpleNamespace( - in_profile_run=True, - moe_comm_method=moe_comm_method, - flash_comm_v1_enabled=True, - eplb_heat_collection_status=True, - ), - ) - - result = layer.forward_impl(hidden_states, router_logits, return_with_event=return_with_event) - - assert forward_context.moe_layer_index == 1 - moe_comm_method.prepare.assert_called_once_with( - hidden_states=hidden_states, - router_logits=router_logits, - replace_allreduce=True, - enable_shared_expert_dp=False, - quant_type=QuantType.NONE, - ) - apply_kwargs = layer.quant_method.apply.call_args.kwargs - assert apply_kwargs["x"] is prepared_hidden - assert apply_kwargs["router_logits"] is prepared_logits - assert apply_kwargs["num_experts"] == 4 - assert apply_kwargs["enable_force_load_balance"] is True - assert torch.equal(apply_kwargs["mc2_mask"], prepare_output.mc2_mask) - torch.testing.assert_close(layer.moe_load, torch.tensor([2, 3])) - if return_with_event: - assert result.routed_out.shape == hidden_states.shape - assert result.before_dispatch_evt is before_dispatch_evt - assert result.before_combine_evt is before_combine_evt - else: - torch.testing.assert_close(result, torch.ones_like(hidden_states) + 2) - - def test_forward_impl_dynamic_eplb_multi_stage(self, monkeypatch): - layer = self._build_layer() - layer.enable_npugraph_ex_static_kernel = False - layer.enable_shared_expert_dp = False - layer.quant_type = QuantType.NONE - layer.top_k = 1 - layer.renormalize = False - layer.use_grouped_topk = False - layer.moe_config = SimpleNamespace(num_experts=2) - layer._expert_map = None - layer.topk_group = None - layer.num_expert_group = None - layer.custom_routing_function = None - layer.scoring_func = "softmax" - layer._original_routed_scaling_factor = 1.0 - layer.routed_scaling_factor = 1.0 - layer.e_score_correction_bias = None - layer.activation = "silu" - layer.apply_router_weight_on_input = False - layer.global_redundant_expert_num = 0 - layer.dynamic_eplb = True - layer.multi_stage = True - layer.moe_load = torch.zeros((2, 2), dtype=torch.int32) - layer.load_counter = torch.tensor([1], dtype=torch.int64) - layer.num_iter = 2 - layer.reduce_results = False - moe_comm_method = MagicMock() - moe_comm_method.prepare.return_value = MoEPrepareOutput( - hidden_states=torch.ones(2, 4), - router_logits=torch.ones(2, 2), - mc2_mask=None, - padded_hidden_states_shape=None, - ) - moe_comm_method.finalize.side_effect = lambda hidden_states, **_: hidden_states - layer.quant_method.apply.return_value = FusedExpertsResult( - routed_out=torch.ones(2, 4), - expert_tokens=torch.tensor([4, 6]), - group_list_type=1, - ) - monkeypatch.setattr(fused_moe_legacy_module, "get_forward_context", MagicMock(return_value=SimpleNamespace())) - monkeypatch.setattr( - fused_moe_legacy_module, - "_EXTRA_CTX", - SimpleNamespace( - in_profile_run=False, - moe_comm_method=moe_comm_method, - flash_comm_v1_enabled=False, - eplb_heat_collection_status=True, - ), - ) - - layer.forward_impl(torch.zeros(2, 4), torch.zeros(2, 2)) - - assert torch.equal(layer.moe_load[1], torch.tensor([4, 6], dtype=torch.int32)) - assert layer.load_counter.item() == 2 - - -class TestAscendFusedMoESharedExperts: - def test_properties_and_forward_delegate(self, monkeypatch): - layer = AscendFusedMoE.__new__(AscendFusedMoE) - if not hasattr(type(layer), "gate"): - pytest.skip("Current AscendFusedMoE does not expose gate property") - layer.multistream_overlap_shared_expert = False - layer._gate = MagicMock() - layer.use_overlapped = True - assert layer.gate is layer._gate - layer.use_overlapped = False - assert layer.gate is None - assert layer.is_internal_router is False - assert layer.use_dp_chunking is False - - monkeypatch.setattr(fused_moe_module.AscendFusedMoE, "forward", MagicMock(return_value="routed")) - layer._shared_experts = None - assert layer.forward(torch.ones(1, 2), torch.ones(1, 2)) == "routed" - - fused_moe_module.AscendFusedMoE.forward.return_value = "forwarded" - layer._shared_experts = MagicMock() - assert layer.forward(torch.ones(1, 2), torch.ones(1, 2)) == "forwarded" - - def test_shared_experts_split_with_expert_gate(self): - layer = AscendFusedMoE.__new__(AscendFusedMoE) - if not hasattr(layer, "_shared_experts_part1"): - pytest.skip("Current AscendFusedMoE does not split shared experts") - hidden_states = torch.tensor([[1.0, -1.0]]) - gate_up = torch.tensor([[2.0, -2.0]]) - down_out = torch.tensor([[3.0, 4.0]]) - gate_out = torch.tensor([[0.0, 2.0]]) - shared_experts = MagicMock() - shared_experts.gate_up_proj.return_value = (gate_up, None) - shared_experts.act_fn.side_effect = lambda tensor: tensor + 1 - shared_experts.down_proj.return_value = (down_out, None) - shared_experts.expert_gate.return_value = (gate_out, None) - layer._shared_experts = shared_experts - - part1_out = layer._shared_experts_part1(hidden_states) - part2_out = layer._shared_experts_part2(hidden_states, part1_out) - - torch.testing.assert_close(part1_out, gate_up) - torch.testing.assert_close(part2_out, F.sigmoid(gate_out) * down_out) - - @pytest.mark.parametrize("has_shared_experts", [False, True]) - def test_shared_forward_impl_routes_shared_output(self, monkeypatch, has_shared_experts): - layer = AscendFusedMoE.__new__(AscendFusedMoE) - if not hasattr(layer, "shared_forward_impl"): - pytest.skip("Current AscendFusedMoE has no shared_forward_impl") - layer.multistream_overlap_shared_expert = False - layer.use_overlapped = False - layer._shared_experts = MagicMock() if has_shared_experts else None - hidden_states = torch.randn(2, 4) - router_logits = torch.randn(2, 3) - fused_result = fused_moe_module.FusedMoEResult( - routed_out=torch.ones(2, 4), - before_dispatch_evt=MagicMock(), - before_combine_evt=MagicMock(), - ) - monkeypatch.setattr( - fused_moe_module.torch.npu, - "current_stream", - MagicMock(return_value=MagicMock(record_event=MagicMock(return_value=MagicMock()))), - ) - monkeypatch.setattr(fused_moe_module.AscendFusedMoE, "forward_impl", MagicMock(return_value=fused_result)) - layer._forward_shared_experts = MagicMock(return_value="shared_out") - - result = layer.shared_forward_impl(hidden_states, router_logits) - - if has_shared_experts: - assert result == ("shared_out", fused_result.routed_out) - layer._forward_shared_experts.assert_called_once() - else: - torch.testing.assert_close(result, fused_result.routed_out) + if has_shared_experts: + assert result[0] is shared_out + assert result[1] is routed_out + runner._forward_shared_experts.assert_called_once() + else: + assert result is routed_out + runner._forward_shared_experts.assert_not_called() diff --git a/tests/ut/ops/test_gdn_attn_builder.py b/tests/ut/ops/test_gdn_attn_builder.py index ff6f94a5acd..c9fc1fcfc85 100644 --- a/tests/ut/ops/test_gdn_attn_builder.py +++ b/tests/ut/ops/test_gdn_attn_builder.py @@ -31,7 +31,6 @@ from vllm_ascend.ops.triton.fla.utils import ( prepare_update_chunk_offsets as runtime_prepare_update_chunk_offsets, ) -from vllm_ascend.utils import vllm_version_is @pytest.fixture(autouse=True) @@ -50,12 +49,11 @@ def _no_pin_memory(): # compute_causal_conv1d_metadata uses np_to_pinned_tensor which reads # PIN_MEMORY. Without physical NPU, t.pin_memory() raises # "Please register PrivateUse1HooksInterface first". - with patch("vllm.utils.torch_utils.PIN_MEMORY", False): - if vllm_version_is("0.23.0"): - yield - else: - with patch("vllm.v1.attention.backends.utils.PIN_MEMORY", False): - yield + with ( + patch("vllm.utils.torch_utils.PIN_MEMORY", False), + patch("vllm.v1.attention.backends.utils.PIN_MEMORY", False), + ): + yield @dataclass diff --git a/tests/ut/patch/platform/test_patch_balance_schedule.py b/tests/ut/patch/platform/test_patch_balance_schedule.py index e86674add79..dc7d304319a 100644 --- a/tests/ut/patch/platform/test_patch_balance_schedule.py +++ b/tests/ut/patch/platform/test_patch_balance_schedule.py @@ -9,10 +9,8 @@ What is guarded here (everything reachable from CPU UT): -* the ``schedule`` override signature stays callable by BOTH vllm versions CI - runs (v0.23.0 calls ``schedule()``; 1f486d96 calls ``schedule(throttle_prefills)``) - and carries every parameter the installed ``schedule`` exposes -- NOT an - exact-match to the installed signature, which differs per lane; +* the ``schedule`` override signature stays aligned with the installed + scheduler signature shared by both supported vLLM refs; * the ``BalanceScheduler.__init__`` signature stays drop-in compatible with upstream's ``Scheduler.__init__`` (upstream constructs ``Scheduler(...)`` with kwargs, which after the swap constructs our subclass); @@ -108,22 +106,30 @@ def _schedule_body_ast(source: str) -> str: ) assert func is not None, "no schedule() in source" - filtered = [] - for stmt in func.body: - dump = ast.dump(stmt) - # delta 1: disabled-path early return. ``_balance_enabled`` is touched - # nowhere else, so its presence alone identifies this delta. - if "_balance_enabled" in dump: - continue - # delta 2: the balance_flag admission gate. - if "balance_queue" in dump and "max_num_running_reqs" in dump: - continue - # delta 3: request_queue None-check (ours: if-is-None-break; - # upstream: assert is-not-None). "Is" is a substring of "IsNot". - if "request_queue" in dump and "None" in dump and "Is" in dump: - continue - filtered.append(stmt) - return ast.dump(ast.Module(body=filtered, type_ignores=[])) + class _BalanceDeltaStripper(ast.NodeTransformer): + def visit_If(self, node: ast.If): # noqa: N802 + test = ast.dump(node.test) + # delta 1: disabled-path early return. + if "_balance_enabled" in test: + return None + # delta 2: the balance admission gate inside the WAITING loop. + if "balance_queue" in test and "max_num_running_reqs" in test: + return None + # delta 3 (ours): if request_queue is None: break. + if "request_queue" in test and "None" in test and "Is" in test: + return None + return self.generic_visit(node) + + def visit_Assert(self, node: ast.Assert): # noqa: N802 + test = ast.dump(node.test) + # delta 3 (upstream): assert request_queue is not None. + if "request_queue" in test and "None" in test and "IsNot" in test: + return None + return self.generic_visit(node) + + stripped = _BalanceDeltaStripper().visit(func) + assert isinstance(stripped, ast.FunctionDef) + return ast.dump(ast.Module(body=stripped.body, type_ignores=[])) def _vllm_ascend_repo_root() -> Path | None: @@ -178,43 +184,13 @@ def _pinned_release_schedule_source() -> tuple[str, str] | None: # --------------------------------------------------------------------------- -# 1. schedule() signature is callable by BOTH engine versions (dual-version CI) +# 1. schedule() signature matches the installed supported vLLM ref # --------------------------------------------------------------------------- -def test_schedule_signature_covers_both_engine_versions(): - """vllm-ascend CI runs against TWO vllm versions at once: the release tag - v0.23.0 (whose engine calls ``schedule()`` with no args) and the - main-verified commit 1f486d96 (whose engine calls - ``schedule(throttle_prefills)``). A single override signature must be - callable by BOTH engines, so it carries ``throttle_prefills`` with a - default -- a deliberate superset of v0.23.0's ``schedule(self)``. - - Asserting exact equality with the *installed* signature would be wrong: - on the v0.23.0 lane the installed signature is ``schedule(self)`` while - ours is ``schedule(self, throttle_prefills=False)``, so an equality check - can only ever pass on ONE of the two lanes. Instead we assert the two - things that actually matter on both lanes: - - * both engines' call shapes bind cleanly to our signature (callable); and - * our signature carries every parameter the installed ``schedule`` exposes - (so an upstream parameter addition is caught here regardless of lane).""" - sig = inspect.signature(BalanceScheduler.schedule) - # The engine invokes schedule() on an instance, so ``self`` is implicitly - # bound; strip it before simulating the engine's call shapes, otherwise - # sig.bind() complains about the missing ``self`` argument. - sig = sig.replace(parameters=[p for p in sig.parameters.values() if p.name != "self"]) - # v0.23.0 engine call shape, then 1f486d96 engine call shape. - sig.bind() - sig.bind(throttle_prefills=True) - - up = {k for k in inspect.signature(_UpstreamScheduler.schedule).parameters if k != "self"} - ours = {k for k in sig.parameters if k != "self"} - assert up <= ours, ( - f"BalanceScheduler.schedule is missing parameters the installed " - f"vLLM Scheduler.schedule exposes ({up - ours}); the engine's call " - f"would raise TypeError." - ) +def test_schedule_signature_matches_installed_vllm(): + """Both supported refs expose ``schedule(throttle_prefills=False)``.""" + assert inspect.signature(BalanceScheduler.schedule) == inspect.signature(_UpstreamScheduler.schedule) # --------------------------------------------------------------------------- @@ -229,13 +205,8 @@ def test_balance_deltas_present_in_schedule(): src = inspect.getsource(BalanceScheduler.schedule) # delta 1: disabled-path early return delegates to super().schedule(). - # Whether throttle_prefills is forwarded is decided by signature - # introspection (_SUPER_SCHEDULE_HAS_THROTTLE), NOT a version string, so - # the disabled path works on BOTH the v0.23.0 and 1f486d96 CI lanes. assert "if not self._balance_enabled:" in src - assert "_SUPER_SCHEDULE_HAS_THROTTLE" in src assert "super().schedule(throttle_prefills)" in src - assert "super().schedule()" in src # delta 2: the balance_flag admission gate (leader-at-cap => global freeze). assert "max(t.item() for t in self.balance_queue)" in src diff --git a/tests/ut/patch/platform/test_patch_deepseek_v4_tool_call_parser.py b/tests/ut/patch/platform/test_patch_deepseek_v4_tool_call_parser.py index 572c4817318..22d3d06b2b6 100644 --- a/tests/ut/patch/platform/test_patch_deepseek_v4_tool_call_parser.py +++ b/tests/ut/patch/platform/test_patch_deepseek_v4_tool_call_parser.py @@ -7,9 +7,9 @@ from vllm_ascend.utils import vllm_version_is -if not vllm_version_is("0.23.0"): +if not vllm_version_is("0.24.0"): pytest.skip( - "DeepSeekV4ToolParser compatibility patch only applies to vLLM 0.23.0", + "DeepSeekV4ToolParser compatibility patch only applies to vLLM 0.24.0", allow_module_level=True, ) diff --git a/tests/ut/patch/test_hunyuan_vl_processor_compat.py b/tests/ut/patch/test_hunyuan_vl_processor_compat.py index add27ec0e47..8a5346e0b11 100644 --- a/tests/ut/patch/test_hunyuan_vl_processor_compat.py +++ b/tests/ut/patch/test_hunyuan_vl_processor_compat.py @@ -9,7 +9,7 @@ import vllm_ascend.patch.hunyuan_vl_processor_compat as compat -def test_v023_imports_native_processors_without_persistent_aliases(monkeypatch): +def test_v024_imports_native_processors_without_persistent_aliases(monkeypatch): import transformers.models.hunyuan_vl.image_processing_hunyuan_vl as native_image import vllm.transformers_utils.processors as vllm_processors @@ -48,7 +48,7 @@ def import_hunyuan_vision(name: str) -> ModuleType: monkeypatch.setattr(native_image, "smart_resize", fake_smart_resize) monkeypatch.setattr(compat.importlib, "import_module", import_hunyuan_vision) - assert compat._import_v023_hunyuan_vision() is hunyuan_vision + assert compat._import_v024_hunyuan_vision() is hunyuan_vision for module_name, previous_module in previous_modules.items(): assert sys.modules.get(module_name) is previous_module @@ -56,7 +56,7 @@ def import_hunyuan_vision(name: str) -> ModuleType: assert vars(vllm_processors).get(attribute_name) is previous_attribute -def test_v023_restores_aliases_after_import_error(monkeypatch): +def test_v024_restores_aliases_after_import_error(monkeypatch): import transformers.models.hunyuan_vl.image_processing_hunyuan_vl as native_image class FakeProcessor: @@ -81,13 +81,13 @@ def fail_import(_name: str) -> ModuleType: monkeypatch.setattr(compat.importlib, "import_module", fail_import) with pytest.raises(ImportError, match="expected test failure"): - compat._import_v023_hunyuan_vision() + compat._import_v024_hunyuan_vision() for module_name, previous_module in previous_modules.items(): assert sys.modules.get(module_name) is previous_module -def test_installer_runs_v023_backports_in_order(monkeypatch): +def test_installer_runs_release_backports_in_order(monkeypatch): hunyuan_vision = object() calls: list[Any] = [] @@ -105,10 +105,10 @@ def patch_processor(module: Any) -> None: def patch_loader(module: Any) -> None: calls.append(("loader", module)) - monkeypatch.setattr(compat, "vllm_version_is", lambda version: version == "0.23.0") + monkeypatch.setattr(compat, "vllm_version_is", lambda version: version == "0.24.0") monkeypatch.setattr( compat, - "_import_v023_hunyuan_vision", + "_import_v024_hunyuan_vision", import_hunyuan_vision, ) monkeypatch.setattr( @@ -123,7 +123,7 @@ def patch_loader(module: Any) -> None: ) monkeypatch.setattr( compat, - "_patch_v023_processor_methods", + "_patch_v024_processor_methods", patch_processor, ) compat.install_hunyuan_vl_processor_compat() @@ -246,7 +246,7 @@ def native_init( ] -def test_v023_backports_native_processor_call_protocol(monkeypatch): +def test_v024_backports_native_processor_call_protocol(monkeypatch): class FakeProcessingInfo: pass @@ -258,7 +258,7 @@ class FakeMultiModalProcessor: HunYuanVLMultiModalProcessor=FakeMultiModalProcessor, ) compat._patch_hunyuan_processor_loader(hunyuan_vision) - compat._patch_v023_processor_methods(hunyuan_vision) + compat._patch_v024_processor_methods(hunyuan_vision) processor_args: list[tuple[Any, dict[str, Any]]] = [] @@ -373,15 +373,15 @@ class FakeMultiModalProcessor: hunyuan_vision = SimpleNamespace( HunYuanVLMultiModalProcessor=FakeMultiModalProcessor, ) - monkeypatch.setattr(compat, "vllm_version_is", lambda version: version == "0.23.0") + monkeypatch.setattr(compat, "vllm_version_is", lambda version: version == "0.24.0") monkeypatch.setattr( compat, - "_import_v023_hunyuan_vision", + "_import_v024_hunyuan_vision", lambda: hunyuan_vision, ) monkeypatch.setattr(compat, "_remove_stale_registry_entries", lambda: True) monkeypatch.setattr(compat, "_patch_hunyuan_processor_loader", lambda _module: None) - monkeypatch.setattr(compat, "_patch_v023_processor_methods", lambda _module: None) + monkeypatch.setattr(compat, "_patch_v024_processor_methods", lambda _module: None) compat.install_hunyuan_vl_processor_compat() diff --git a/tests/ut/quantization/test_compressed_tensors_config.py b/tests/ut/quantization/test_compressed_tensors_config.py index afbfc4d0fe6..977831d27f7 100644 --- a/tests/ut/quantization/test_compressed_tensors_config.py +++ b/tests/ut/quantization/test_compressed_tensors_config.py @@ -1,17 +1,14 @@ from unittest.mock import MagicMock, patch -import pytest from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod from tests.ut.base import TestBase from tests.ut.quantization.conftest_quantization import COMPRESSED_TENSORS_W8A8_CONFIG -from vllm_ascend.ops.fused_moe.fused_moe import AscendUnquantizedFusedMoEMethod from vllm_ascend.quantization.compressed_tensors_config import AscendCompressedTensorsConfig -from vllm_ascend.quantization.method_adapters import AscendFusedMoEMethod, AscendLinearMethod -from vllm_ascend.quantization.methods import AscendW8A8DynamicFusedMoEMethod, AscendW8A8DynamicLinearMethod -from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD, vllm_version_is +from vllm_ascend.quantization.method_adapters import AscendLinearMethod +from vllm_ascend.quantization.methods import AscendW8A8DynamicLinearMethod +from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD class TestAscendCompressedTensorsQuanType(TestBase): @@ -95,35 +92,6 @@ def test_get_linear_unquantized_method(self): self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) self.assertTrue(isinstance(result, UnquantizedLinearMethod)) - @pytest.mark.skipif( - not vllm_version_is("0.23.0"), - reason="Legacy FusedMoE quant method UT is only for vLLM 0.23.0.", - ) - @patch("vllm_ascend.quantization.methods.AscendW8A8DynamicFusedMoEMethod.__init__") - def test_get_moe_quant_method(self, mock_method): - mock_method.return_value = None - layer = MagicMock(spec=FusedMoE) - layer.moe_config = {} - result = self.config.get_quant_method(layer, "model.layers.0.mlp.experts") - self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) - self.assertTrue(isinstance(result, AscendFusedMoEMethod)) - self.assertTrue(isinstance(layer.scheme, AscendW8A8DynamicFusedMoEMethod)) - - @pytest.mark.skipif( - not vllm_version_is("0.23.0"), - reason="Legacy FusedMoE quant method UT is only for vLLM 0.23.0.", - ) - @patch("vllm_ascend.ops.fused_moe.fused_moe.AscendUnquantizedFusedMoEMethod.__init__") - @patch("vllm_ascend.quantization.compressed_tensors_config.should_ignore_layer") - def test_get_moe_unquantized_method(self, mock_ignore_layer, mock_method): - mock_method.return_value = None - mock_ignore_layer.return_value = True - layer = MagicMock(spec=FusedMoE) - layer.moe_config = {} - result = self.config.get_quant_method(layer, "model.layers.0.mlp.experts") - self.assertEqual(layer.ascend_quant_method, COMPRESSED_TENSORS_METHOD) - self.assertTrue(isinstance(result, AscendUnquantizedFusedMoEMethod)) - def test_no_quant_method(self): layer = MagicMock(spec=Attention) result = self.config.get_quant_method(layer, "attn") diff --git a/tests/ut/quantization/test_modelslim_config.py b/tests/ut/quantization/test_modelslim_config.py index 2b9ba5a8f08..50c9735b578 100644 --- a/tests/ut/quantization/test_modelslim_config.py +++ b/tests/ut/quantization/test_modelslim_config.py @@ -3,12 +3,9 @@ import tempfile from unittest.mock import MagicMock, patch -import pytest import torch from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.fused_moe import FusedMoE -from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.linear import LinearBase from tests.ut.base import TestBase @@ -19,7 +16,7 @@ get_linear_quant_type, get_packed_modules_mapping, ) -from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD, vllm_version_is +from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD class TestAscendModelSlimConfig(TestBase): @@ -159,41 +156,6 @@ def test_get_quant_method_for_c8_kv_cache_attention(self): self.assertIsInstance(args[0], AscendC8KVCacheAttentionMethod) - @pytest.mark.skipif( - not vllm_version_is("0.23.0"), - reason="Legacy FusedMoE quant method UT is only for vLLM 0.23.0.", - ) - def test_get_quant_method_for_fused_moe(self): - fused_moe_layer = MagicMock(spec=FusedMoE) - fused_moe_layer.moe = MagicMock(spec=FusedMoEConfig) - fused_moe_layer.moe_config = MagicMock(spec=FusedMoEConfig) - mock_config = MagicMock() - mock_config.model_config.hf_config.model_type = None - - # Test skipped layer - with ( - patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=True), - patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config), - patch( - "vllm_ascend.ops.fused_moe.fused_moe.AscendUnquantizedFusedMoEMethod", return_value=MagicMock() - ) as mock_ascend_moe, - ): - method = self.ascend_config.get_quant_method(fused_moe_layer, "moe_layer") - self.assertIs(method, mock_ascend_moe.return_value) - - # Test quantized layer - mock_scheme = MagicMock() - with ( - patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=False), - patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config), - patch("vllm_ascend.quantization.modelslim_config.create_scheme_for_layer", return_value=mock_scheme), - patch( - "vllm_ascend.quantization.method_adapters.AscendFusedMoEMethod", return_value=MagicMock() - ) as mock_ascend_moe, - ): - method = self.ascend_config.get_quant_method(fused_moe_layer, "moe_layer") - self.assertIs(method, mock_ascend_moe.return_value) - def test_is_layer_skipped_ascend(self): # Test non-fused layer that should be quantized self.assertFalse(self.ascend_config.is_layer_skipped_ascend("layer1")) diff --git a/tests/ut/spec_decode/a2/test_eagle_proposer.py b/tests/ut/spec_decode/a2/test_eagle_proposer.py index e47afd0de6c..1c46497dbe3 100644 --- a/tests/ut/spec_decode/a2/test_eagle_proposer.py +++ b/tests/ut/spec_decode/a2/test_eagle_proposer.py @@ -4058,6 +4058,9 @@ def test_set_inputs_first_pass_parallel_drafting(self): proposer.parallel_drafting_token_id = -2 parallel_drafting_hs = proposer.parallel_drafting_hidden_state_tensor + assert parallel_drafting_hs is not None + # Production initializes this from the loaded drafter's mask hidden state. + parallel_drafting_hs.fill_(1.0) mock_kv_cache_spec = MagicMock() mock_kv_cache_spec.block_size = block_size diff --git a/tests/ut/spec_decode/test_extract_hidden_states_proposer.py b/tests/ut/spec_decode/test_extract_hidden_states_proposer.py index 6763523984b..6580e2a486e 100644 --- a/tests/ut/spec_decode/test_extract_hidden_states_proposer.py +++ b/tests/ut/spec_decode/test_extract_hidden_states_proposer.py @@ -32,23 +32,15 @@ from vllm_ascend.spec_decode.extract_hidden_states_proposer import ( AscendExtractHiddenStatesProposer, ) -from vllm_ascend.utils import vllm_version_is @pytest.fixture(autouse=True) def _no_pin_memory(): - if vllm_version_is("0.23.0"): - with patch( - "vllm.v1.spec_decode.extract_hidden_states.is_pin_memory_available", - return_value=False, - ): - yield - else: - with patch( - "vllm.v1.spec_decode.extract_hidden_states.PIN_MEMORY", - False, - ): - yield + with patch( + "vllm.v1.spec_decode.extract_hidden_states.PIN_MEMORY", + False, + ): + yield class MockCachedRequestState: diff --git a/vllm_ascend/_310p/fused_moe/fused_moe.py b/vllm_ascend/_310p/fused_moe/fused_moe.py index 8289a8c2b26..ae89581e6fe 100644 --- a/vllm_ascend/_310p/fused_moe/fused_moe.py +++ b/vllm_ascend/_310p/fused_moe/fused_moe.py @@ -17,17 +17,13 @@ from collections.abc import Callable import torch -from vllm.distributed import get_dp_group, get_ep_group, get_tp_group from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import UnquantizedFusedMoEMethod from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType from vllm_ascend.ops.fused_moe.experts_selector import zero_experts_compute from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner -from vllm_ascend.ops.fused_moe.moe_comm_method import ( - AllGatherCommImpl, - FusedExpertsResult, - _MoECommMethods, -) +from vllm_ascend.ops.fused_moe.moe_comm_method import _MoECommMethods from vllm_ascend.ops.fused_moe.moe_runtime_args import build_fused_experts_input from vllm_ascend.quantization.quant_type import QuantType from vllm_ascend.utils import maybe_trans_nz, vllm_version_is @@ -35,20 +31,6 @@ from .experts_selector import select_experts from .moe_comm_method import AllGatherCommImpl310 -if vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe.layer import FusedMoE as _LegacyFusedMoEBase - from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod -else: - from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import UnquantizedFusedMoEMethod - - try: - from vllm.model_executor.layers.fused_moe.layer import FusedMoE as _LegacyFusedMoEBase - except ImportError: - _LegacyFusedMoEBase = torch.nn.Module - - if not isinstance(_LegacyFusedMoEBase, type): - _LegacyFusedMoEBase = torch.nn.Module - class AscendUnquantizedFusedMoEMethod310(UnquantizedFusedMoEMethod): def __init__(self, moe: FusedMoEConfig = None): @@ -66,7 +48,9 @@ def maybe_make_prepare_finalize(self, routing_tables=None): def process_weights_after_loading(self, layer): super().process_weights_after_loading(layer) - if not vllm_version_is("0.23.0"): + # vLLM PR #44589 landed after the v0.24 main-line cut point + # (798185d) and is present in the verified main commit only. + if not vllm_version_is("0.24.0"): w13_data = self._maybe_pad_weight(layer.w13_weight.data).transpose(1, 2) w13_data = maybe_trans_nz(w13_data) layer.w13_weight = torch.nn.Parameter(w13_data, requires_grad=False) @@ -148,226 +132,42 @@ def apply( return final_hidden_states -if not vllm_version_is("0.23.0"): - - class AscendMoERunner310(AscendMoERunner): - def __init__( - self, +class AscendMoERunner310(AscendMoERunner): + def __init__( + self, + layer_name, + moe_config, + router, + routed_experts, + enable_dbo=False, + gate=None, + shared_experts=None, + shared_expert_gate=None, + routed_input_transform=None, + routed_output_transform=None, + routed_scaling_factor=1, + tid2eid=None, + n_shared_experts: int = 0, + ): + super().__init__( layer_name, moe_config, router, routed_experts, - enable_dbo=False, - gate=None, - shared_experts=None, - shared_expert_gate=None, - routed_input_transform=None, - routed_output_transform=None, - routed_scaling_factor=1, - tid2eid=None, - n_shared_experts: int = 0, - ): - super().__init__( - layer_name, - moe_config, - router, - routed_experts, - enable_dbo, - gate, - shared_experts, - shared_expert_gate, - routed_input_transform, - routed_output_transform, - routed_scaling_factor, - tid2eid, - n_shared_experts, - ) - - if routed_experts.quant_config is None: - routed_experts.quant_method = AscendUnquantizedFusedMoEMethod310(self.moe_config) - self.quant_type = self._get_quant_type() - - self.multistream_overlap_shared_expert = False - _MoECommMethods[MoECommType.ALLGATHER] = AllGatherCommImpl310(self.moe_config) - - -class AscendFusedMoE310(_LegacyFusedMoEBase): - def __init__(self, *args, **kwargs): - if _LegacyFusedMoEBase is torch.nn.Module: - raise RuntimeError("AscendFusedMoE310 is only kept for the legacy FusedMoE class API.") - super().__init__(*args, **kwargs) - - self._routed_input_transform = kwargs.get("routed_input_transform") - self._shared_experts = kwargs.get("shared_experts") - self.global_num_experts = kwargs["num_experts"] - - if self.quant_config is None: - self.quant_method = AscendUnquantizedFusedMoEMethod310(self.moe_config) - else: - self.quant_method = self.quant_config.get_quant_method(self, self.layer_name) - - assert self.quant_method is not None - # Keep base_quant_method aligned with the Ascend-replaced quant_method - # so FusedMoE.maybe_init_modular_kernel doesn't dispatch into the - # upstream UnquantizedFusedMoEMethod.maybe_make_prepare_finalize. - self.base_quant_method = self.quant_method - - self.moe_config.tp_group = get_tp_group() - self.moe_config.dp_group = get_dp_group() - self.moe_config.ep_group = get_ep_group() - self.moe_config.supports_eplb = False - - # init moe - self.global_expert_map = None - self.local_expert_map = None - if self.moe_config.ep_size > 1: - raise RuntimeError("Expert Parallel is not supported on 310P. Please remove --enable-expert-parallel.") - self.local_num_experts = self.global_num_experts - - self.moe_config.num_experts = self.global_num_experts - self.moe_config.num_local_experts = self.local_num_experts - self.moe_config.global_redundant_expert_num = 0 - - moe_quant_params = { - "num_experts": self.local_num_experts, - "hidden_size": self.hidden_size, - "intermediate_size_per_partition": self.intermediate_size_per_partition, - "params_dtype": self.params_dtype, - "weight_loader": self.weight_loader, - } - - self.quant_method.create_weights(layer=self, **moe_quant_params) - self.quant_type = self.get_quant_type() - - _MoECommMethods[MoECommType.ALLGATHER] = AllGatherCommImpl310(self.moe_config) - - if vllm_version_is("0.23.0"): - self.runner = AscendMoERunner( - self.layer_name, - self.moe_config, - self.router, - self._routed_input_transform, - kwargs.pop("gate", None), - kwargs.pop("shared_experts", None), - self.quant_method, - self.vllm_config.parallel_config.enable_dbo, - ) - else: - self.runner = AscendMoERunner310( - self.layer_name, - self.moe_config, - self.router, - self._routed_input_transform, - kwargs.pop("gate", None), - kwargs.pop("shared_experts", None), - self.quant_method, - self.vllm_config.parallel_config.enable_dbo, - ) - - @property - def is_internal_router(self) -> bool: - # 310P Ascend path expects router logits from the model forward path. - return False - - def init_experts_map(self, moe_config): - """ - Initialize expert mapping for MoE (Mixture of Experts) model. - - This function creates mappings between global expert indices and local expert indices - for each rank in the expert parallel group. It divides the total experts among - different ranks and creates both global and local expert maps that are used - during MoE computation to determine which experts are handled by which rank. - - Args: - moe_config: Configuration object containing MoE parameters including - number of experts, expert parallel size, and expert parallel rank. - - Returns: - tuple: A tuple containing: - - global_expert_map: Stack of expert maps for all ranks - - local_expert_map: Expert map for the current rank (transferred to NPU) - """ - n_experts = moe_config.num_experts - ep_size = moe_config.ep_size - all_experts = torch.arange(n_experts, dtype=torch.int32) - experts_groups = all_experts.chunk(ep_size) - global_expert_map = [] - local_expert_map = None - for rankid in range(ep_size): - expert_map = torch.full((n_experts,), -1, dtype=torch.int32) - local_experts = experts_groups[rankid] - expert_map[local_experts] = torch.arange(local_experts.shape[0], dtype=torch.int32) - global_expert_map.append(expert_map) - if rankid == moe_config.ep_rank: - local_expert_map = expert_map.npu() - return torch.stack(global_expert_map), local_expert_map - - def get_quant_type(self) -> QuantType: - quant_method = self.quant_method - if not hasattr(quant_method, "quant_method") or quant_method.quant_method is None: - return QuantType.NONE - - method = quant_method.quant_method - quant_type = getattr(method, "quant_type", QuantType.NONE) - if quant_type not in [QuantType.NONE, QuantType.W8A8]: - raise RuntimeError("Only Unquant and W8A8 is supported.") - return quant_type - - def forward_impl( # type: ignore[override] - self, hidden_states: torch.Tensor, router_logits: torch.Tensor - ) -> torch.Tensor: - assert self.quant_method is not None - assert self.routed_scaling_factor == 1.0, "routed_scaling_factor != 1.0 is not supported." - - prepare_output = _EXTRA_CTX.moe_comm_method.prepare( - hidden_states=hidden_states, router_logits=router_logits, quant_type=self.quant_type - ) - hidden_states = prepare_output.hidden_states - router_logits = prepare_output.router_logits - pertoken_scale = prepare_output.pertoken_scale - padded_hidden_states_shape = prepare_output.padded_hidden_states_shape - - # Matrix multiply. - fused_experts_results: FusedExpertsResult = self.quant_method.apply( - layer=self, - x=hidden_states, - use_grouped_topk=self.use_grouped_topk, - top_k=self.top_k, - router_logits=router_logits, - renormalize=self.renormalize, - topk_group=self.topk_group, - num_expert_group=self.num_expert_group, - custom_routing_function=self.custom_routing_function, - scoring_func=self.scoring_func, - e_score_correction_bias=self.e_score_correction_bias, - num_experts=self.global_num_experts, - expert_map=self.local_expert_map, - apply_router_weight_on_input=self.apply_router_weight_on_input, - pertoken_scale=pertoken_scale, - ) - - routed_out = _EXTRA_CTX.moe_comm_method.finalize( - hidden_states=fused_experts_results.routed_out, - reduce_results=isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl), - padded_hidden_states_shape=padded_hidden_states_shape, + enable_dbo, + gate, + shared_experts, + shared_expert_gate, + routed_input_transform, + routed_output_transform, + routed_scaling_factor, + tid2eid, + n_shared_experts, ) - return routed_out - - def _forward_shared_experts(self, hidden_states: torch.Tensor): - if self._shared_experts is None: - return None - return self._shared_experts(hidden_states) + if routed_experts.quant_config is None: + routed_experts.quant_method = AscendUnquantizedFusedMoEMethod310(self.moe_config) + self.quant_type = self._get_quant_type() - def shared_forward_impl( # type: ignore[override] - self, hidden_states: torch.Tensor, router_logits: torch.Tensor - ): - routed_out = AscendFusedMoE310.forward_impl( - self, - hidden_states=hidden_states, - router_logits=router_logits, - ) - if self._shared_experts is None: - return routed_out - shared_out = self._forward_shared_experts(hidden_states) - return shared_out, routed_out + self.multistream_overlap_shared_expert = False + _MoECommMethods[MoECommType.ALLGATHER] = AllGatherCommImpl310(self.moe_config) diff --git a/vllm_ascend/_310p/quantization/modelslim_config.py b/vllm_ascend/_310p/quantization/modelslim_config.py index 33e71e7f0b1..d879072e688 100644 --- a/vllm_ascend/_310p/quantization/modelslim_config.py +++ b/vllm_ascend/_310p/quantization/modelslim_config.py @@ -22,6 +22,7 @@ import torch from vllm.config import get_current_vllm_config from vllm.logger import logger +from vllm.model_executor.layers.fused_moe import MoERunner, RoutedExperts from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization import register_quantization_config from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase @@ -38,19 +39,11 @@ get_quant_type_for_layer, packed_modules_model_mapping, ) -from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD, vllm_version_is - -if vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe import FusedMoE -else: - from vllm.model_executor.layers.fused_moe import MoERunner, RoutedExperts +from vllm_ascend.utils import ASCEND_QUANTIZATION_METHOD def _is_fused_moe_layer(layer: torch.nn.Module) -> bool: - if vllm_version_is("0.23.0"): - return isinstance(layer, FusedMoE) - else: - return isinstance(layer, (MoERunner, RoutedExperts)) + return isinstance(layer, (MoERunner, RoutedExperts)) def create_scheme_for_layer( diff --git a/vllm_ascend/_310p/worker_310p.py b/vllm_ascend/_310p/worker_310p.py index 20112ed78ea..a25d06ce696 100644 --- a/vllm_ascend/_310p/worker_310p.py +++ b/vllm_ascend/_310p/worker_310p.py @@ -25,7 +25,7 @@ from vllm.utils.torch_utils import set_random_seed # noqa: E402 from vllm_ascend._310p.model_runner_310p import NPUModelRunner310 -from vllm_ascend.utils import is_rc_device, vllm_version_is +from vllm_ascend.utils import is_rc_device from vllm_ascend.worker.worker import NPUWorker, init_workspace_manager @@ -138,10 +138,7 @@ def _init_device(self): torch.npu.empty_cache() # take current memory snapshot - if vllm_version_is("0.23.0"): - self.init_snapshot = MemorySnapshot() - else: - self.init_snapshot = MemorySnapshot(device=device) + self.init_snapshot = MemorySnapshot(device=device) self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization if is_rc_device(): self.init_snapshot.free_memory = psutil.virtual_memory().available diff --git a/vllm_ascend/core/single_type_kv_cache_manager.py b/vllm_ascend/core/single_type_kv_cache_manager.py index efd9f7ec37f..582f471a077 100644 --- a/vllm_ascend/core/single_type_kv_cache_manager.py +++ b/vllm_ascend/core/single_type_kv_cache_manager.py @@ -288,9 +288,9 @@ def get_manager_for_kv_cache_spec( # and ``full_sequence_must_fit`` admission reserves the full # ``max_model_len`` worth of blocks per request, exhausting the pool # at cc>=2 on DSv4 (see vLLM issue #40863). - token_budget = max_num_batched_tokens if vllm_version_is("0.23.0") else max_in_flight_tokens + token_budget = max_num_batched_tokens if vllm_version_is("0.24.0") else max_in_flight_tokens if token_budget is not None and max_model_len is not None: - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): kwargs["max_admission_blocks_per_request"] = kv_cache_spec.max_admission_blocks_per_request( max_num_batched_tokens=token_budget, max_model_len=max_model_len, diff --git a/vllm_ascend/distributed/weight_transfer/hccl_engine.py b/vllm_ascend/distributed/weight_transfer/hccl_engine.py index f1f24ae6755..8e896177291 100644 --- a/vllm_ascend/distributed/weight_transfer/hccl_engine.py +++ b/vllm_ascend/distributed/weight_transfer/hccl_engine.py @@ -116,7 +116,7 @@ class HCCLWeightTransferEngine(WeightTransferEngine[HCCLWeightTransferInitInfo, init_info_cls = HCCLWeightTransferInitInfo update_info_cls = HCCLWeightTransferUpdateInfo - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): def __init__( self, @@ -139,7 +139,7 @@ def __init__( # type: ignore[misc] super().__init__(config, vllm_config, device, model) self.model_update_group: PyHcclCommunicator | None = None # type: ignore[no-redef] - if not vllm_version_is("0.23.0"): + if not vllm_version_is("0.24.0"): def start_weight_update(self) -> None: from vllm.model_executor.model_loader.reload import ( diff --git a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py index 3b03e56d729..5c9e74e9b98 100644 --- a/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py +++ b/vllm_ascend/distributed/weight_transfer/npu_ipc_engine.py @@ -118,7 +118,7 @@ class NPUIPCWeightTransferEngine(WeightTransferEngine[NPUIPCWeightTransferInitIn init_info_cls = NPUIPCWeightTransferInitInfo update_info_cls = NPUIPCWeightTransferUpdateInfo - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): def __init__( self, @@ -168,7 +168,7 @@ def init_transfer_engine(self, init_info: NPUIPCWeightTransferInitInfo) -> None: """No initialization needed for NPU IPC backend.""" pass - if not vllm_version_is("0.23.0"): + if not vllm_version_is("0.24.0"): def start_weight_update(self) -> None: """No-op for NPU IPC engine (no layerwise reloading).""" diff --git a/vllm_ascend/models/deepseek_v4.py b/vllm_ascend/models/deepseek_v4.py index b99fd94db25..51e46757634 100644 --- a/vllm_ascend/models/deepseek_v4.py +++ b/vllm_ascend/models/deepseek_v4.py @@ -44,7 +44,7 @@ tensor_model_parallel_all_gather, ) from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoE, fused_moe_make_expert_params_mapping from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -83,12 +83,8 @@ extract_dsv4_layer_index, get_ascend_device_type, get_dsv4_compress_ratio, - vllm_version_is, ) -if not vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping - def _get_ascend_dsa_backend(): # Keep this lazy to avoid vLLM model-inspection circular imports. @@ -1280,26 +1276,15 @@ def compute_logits( def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) - if vllm_version_is("0.23.0"): - return FusedMoE.make_expert_params_mapping( - self.model, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - + (self.config.n_shared_experts if getattr(get_ascend_config(), "mix_placement", False) else 0), - num_redundant_experts=0, - ) - else: - return fused_moe_make_expert_params_mapping( - self.model, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - + (self.config.n_shared_experts if getattr(get_ascend_config(), "mix_placement", False) else 0), - num_redundant_experts=0, - ) + return fused_moe_make_expert_params_mapping( + self.model, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + + (self.config.n_shared_experts if getattr(get_ascend_config(), "mix_placement", False) else 0), + num_redundant_experts=0, + ) def get_mtp_target_hidden_states(self) -> torch.Tensor | None: """Pre-hc_head residual stream buffer (max_num_batched_tokens, @@ -1317,26 +1302,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) - if vllm_version_is("0.23.0"): - expert_params_mapping = FusedMoE.make_expert_params_mapping( - self.model, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - + (self.config.n_shared_experts if rocm_aiter_moe_shared_expert_enabled else 0), - num_redundant_experts=self.num_redundant_experts, - ) - else: - expert_params_mapping = fused_moe_make_expert_params_mapping( - self.model, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - + (self.config.n_shared_experts if rocm_aiter_moe_shared_expert_enabled else 0), - num_redundant_experts=self.num_redundant_experts, - ) + expert_params_mapping = fused_moe_make_expert_params_mapping( + self.model, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + + (self.config.n_shared_experts if rocm_aiter_moe_shared_expert_enabled else 0), + num_redundant_experts=self.num_redundant_experts, + ) params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() diff --git a/vllm_ascend/models/deepseek_v4_mtp.py b/vllm_ascend/models/deepseek_v4_mtp.py index b412d61158a..8ea12d25dd3 100644 --- a/vllm_ascend/models/deepseek_v4_mtp.py +++ b/vllm_ascend/models/deepseek_v4_mtp.py @@ -10,7 +10,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -23,10 +23,7 @@ from vllm.sequence import IntermediateTensors from vllm_ascend.ascend_config import get_ascend_config -from vllm_ascend.utils import enable_dsa_cp, vllm_version_is - -if not vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping +from vllm_ascend.utils import enable_dsa_cp from .deepseek_v4 import ( DeepseekV2DecoderLayer, @@ -260,26 +257,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ("gate_up_proj", "up_proj", 1), ] - if vllm_version_is("0.23.0"): - expert_params_mapping = FusedMoE.make_expert_params_mapping( - model=self.model, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - + (self.config.n_shared_experts if rocm_aiter_moe_shared_expert_enabled else 0), - num_redundant_experts=self.num_redundant_experts, - ) - else: - expert_params_mapping = fused_moe_make_expert_params_mapping( - model=self.model, - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts - + (self.config.n_shared_experts if rocm_aiter_moe_shared_expert_enabled else 0), - num_redundant_experts=self.num_redundant_experts, - ) + expert_params_mapping = fused_moe_make_expert_params_mapping( + model=self.model, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + + (self.config.n_shared_experts if rocm_aiter_moe_shared_expert_enabled else 0), + num_redundant_experts=self.num_redundant_experts, + ) tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size() diff --git a/vllm_ascend/ops/fused_moe/fused_moe.py b/vllm_ascend/ops/fused_moe/fused_moe.py index 9024280eec7..2b8046e1532 100644 --- a/vllm_ascend/ops/fused_moe/fused_moe.py +++ b/vllm_ascend/ops/fused_moe/fused_moe.py @@ -31,6 +31,7 @@ FusedMoE, # noqa: F401 MoERunner, ) +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import UnquantizedFusedMoEMethod from vllm_ascend.ascend_config import get_ascend_config from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType @@ -51,11 +52,6 @@ vllm_version_is, ) -if vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod -else: - from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import UnquantizedFusedMoEMethod - def get_compressed_expert_map(expert_map: torch.Tensor) -> str: global_indices = torch.where(expert_map != -1)[0] @@ -111,7 +107,9 @@ def maybe_make_prepare_finalize(self, routing_tables=None): def process_weights_after_loading(self, layer): super(UnquantizedFusedMoEMethod, self).process_weights_after_loading(layer) - if not vllm_version_is("0.23.0"): + # vLLM PR #44589 landed after the v0.24 main-line cut point + # (798185d) and is present in the verified main commit only. + if not vllm_version_is("0.24.0"): w13_data = self._maybe_pad_weight(layer.w13_weight.data).transpose(1, 2) layer.w13_weight = torch.nn.Parameter(w13_data, requires_grad=False) @@ -196,14 +194,11 @@ def apply( tid2eid=self.tid2eid, input_ids=input_ids, ) - if vllm_version_is("0.23.0"): - model_config = layer.vllm_config.model_config - else: - try: - _vllm_config = get_current_vllm_config() - except AssertionError: - _vllm_config = None - model_config = None if _vllm_config is None else _vllm_config.model_config + try: + _vllm_config = get_current_vllm_config() + except AssertionError: + _vllm_config = None + model_config = None if _vllm_config is None else _vllm_config.model_config if model_config is not None and model_config.enable_return_routed_experts: capturer = getattr(layer, "_ascend_routed_experts_capturer", None) if capturer is not None: @@ -293,519 +288,511 @@ def apply( return final_hidden_states -if vllm_version_is("0.23.0"): - from vllm_ascend.ops.fused_moe.fused_moe_0_23_0 import AscendFusedMoE, AscendMoERunner - - AscendFusedMoE.__module__ = __name__ - AscendMoERunner.__module__ = __name__ - -else: - - class AscendMoERunner(MoERunner): # type: ignore[no-redef] - moe_counter = -1 +class AscendMoERunner(MoERunner): # type: ignore[no-redef] + moe_counter = -1 - def __init__( - self, + def __init__( + self, + layer_name, + moe_config, + router, + routed_experts, + enable_dbo=False, + gate=None, + shared_experts=None, + shared_expert_gate=None, + routed_input_transform=None, + routed_output_transform=None, + routed_scaling_factor=1, + tid2eid=None, + n_shared_experts: int = 0, + ): + super().__init__( layer_name, moe_config, router, routed_experts, - enable_dbo=False, - gate=None, - shared_experts=None, - shared_expert_gate=None, - routed_input_transform=None, - routed_output_transform=None, - routed_scaling_factor=1, - tid2eid=None, - n_shared_experts: int = 0, - ): - super().__init__( - layer_name, - moe_config, - router, - routed_experts, - enable_dbo, - gate, - shared_experts, - shared_expert_gate, - routed_input_transform, - routed_output_transform, - routed_scaling_factor, - ) - self.top_k = moe_config.experts_per_token - self._gate = gate - self.hidden_size = moe_config.hidden_dim - - # Routing params — all stored on routed_experts by the factory function. - self.use_grouped_topk = routed_experts.use_grouped_topk - self.renormalize = routed_experts.renormalize - self.topk_group = routed_experts.topk_group - self.num_expert_group = routed_experts.num_expert_group - self.custom_routing_function = routed_experts.custom_routing_function - self.scoring_func = routed_experts.scoring_func - self._original_routed_scaling_factor = routed_experts.routed_scaling_factor - self.e_score_correction_bias = routed_experts.e_score_correction_bias - self.apply_router_weight_on_input = routed_experts.apply_router_weight_on_input - - # Ascend-specific: not stored on RoutedExperts, passed via runner_args. - self.tid2eid = tid2eid - - # Replace quant_method on routed_experts with the Ascend version. - # Must NOT set self.quant_method (instance attr) — vllm loader scans all - # nn.Module children for quant_method and would call process_weights_after_loading - # on the runner with wrong expectations. Set on routed_experts so that - # self._quant_method (the MoERunner property) returns the Ascend version. - if routed_experts.quant_config is None: - routed_experts.quant_method = AscendUnquantizedFusedMoEMethod(self.moe_config, tid2eid=self.tid2eid) - else: - routed_experts.quant_method = routed_experts.quant_config.get_quant_method( - routed_experts, self.layer_name, tid2eid=self.tid2eid - ) + enable_dbo, + gate, + shared_experts, + shared_expert_gate, + routed_input_transform, + routed_output_transform, + routed_scaling_factor, + ) + self.top_k = moe_config.experts_per_token + self._gate = gate + self.hidden_size = moe_config.hidden_dim + + # Routing params — all stored on routed_experts by the factory function. + self.use_grouped_topk = routed_experts.use_grouped_topk + self.renormalize = routed_experts.renormalize + self.topk_group = routed_experts.topk_group + self.num_expert_group = routed_experts.num_expert_group + self.custom_routing_function = routed_experts.custom_routing_function + self.scoring_func = routed_experts.scoring_func + self._original_routed_scaling_factor = routed_experts.routed_scaling_factor + self.e_score_correction_bias = routed_experts.e_score_correction_bias + self.apply_router_weight_on_input = routed_experts.apply_router_weight_on_input + + # Ascend-specific: not stored on RoutedExperts, passed via runner_args. + self.tid2eid = tid2eid - self.quant_type = self._get_quant_type() + # Replace quant_method on routed_experts with the Ascend version. + # Must NOT set self.quant_method (instance attr) — vllm loader scans all + # nn.Module children for quant_method and would call process_weights_after_loading + # on the runner with wrong expectations. Set on routed_experts so that + # self._quant_method (the MoERunner property) returns the Ascend version. + if routed_experts.quant_config is None: + routed_experts.quant_method = AscendUnquantizedFusedMoEMethod(self.moe_config, tid2eid=self.tid2eid) + else: + routed_experts.quant_method = routed_experts.quant_config.get_quant_method( + routed_experts, self.layer_name, tid2eid=self.tid2eid + ) - self.moe_config.tp_group = get_tp_group() - self.moe_config.dp_group = get_dp_group() - if self.moe_config.ep_size > 1: - self.moe_config.ep_group = get_ep_group() - self.moe_config.mc2_group = get_mc2_group() + self.quant_type = self._get_quant_type() - ascend_config = get_ascend_config() - self._shared_experts = shared_experts + self.moe_config.tp_group = get_tp_group() + self.moe_config.dp_group = get_dp_group() + if self.moe_config.ep_size > 1: + self.moe_config.ep_group = get_ep_group() + self.moe_config.mc2_group = get_mc2_group() - self.enable_npugraph_ex_static_kernel = ascend_config.ascend_compilation_config.enable_static_kernel + ascend_config = get_ascend_config() + self._shared_experts = shared_experts - vllm_config = get_current_vllm_config() + self.enable_npugraph_ex_static_kernel = ascend_config.ascend_compilation_config.enable_static_kernel - if ( - self.custom_routing_function is None - and self.e_score_correction_bias is not None - and not vllm_config.model_config.is_deepseek_mla - ): - self.e_score_correction_bias.data = self.e_score_correction_bias.data.to( - dtype=vllm_config.model_config.dtype - ) + vllm_config = get_current_vllm_config() - self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp - self.multistream_overlap_shared_expert = ( - ascend_config.multistream_overlap_shared_expert and shared_experts is not None - ) - mix_placement = getattr(ascend_config, "mix_placement", False) - - # EPLB initialization (Ascend-specific; mirrors old AscendFusedMoE logic). - AscendMoERunner.moe_counter += 1 - self.moe_instance_id = AscendMoERunner.moe_counter - - eplb_config = ascend_config.eplb_config - - if mix_placement: - moe_config.num_experts += n_shared_experts - - ( - self.global_expert_map, - self._expert_map, - self.log2phy, - self.global_redundant_expert_num, - ) = init_eplb_config( - eplb_config, - AscendMoERunner.moe_counter, - moe_config, - mix_placement, - n_shared_experts, - tp_size=vllm_config.parallel_config.tensor_parallel_size, + if ( + self.custom_routing_function is None + and self.e_score_correction_bias is not None + and not vllm_config.model_config.is_deepseek_mla + ): + self.e_score_correction_bias.data = self.e_score_correction_bias.data.to( + dtype=vllm_config.model_config.dtype ) - moe_config.global_redundant_expert_num = self.global_redundant_expert_num - local_num_experts = (moe_config.num_experts + self.global_redundant_expert_num) // moe_config.ep_size - moe_config.num_local_experts = local_num_experts - routed_experts.expert_map_manager._local_num_experts = local_num_experts - routed_experts.expert_map_manager._expert_map = self._expert_map - - self.dynamic_eplb = eplb_config.dynamic_eplb and (self.log2phy is not None) - self.multi_stage = False - self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64).npu() - if self.dynamic_eplb and eplb_config.expert_heat_collection_interval > 1: - self.multi_stage = True - self.load_counter = torch.tensor(0, dtype=torch.int32, device="npu") - self.num_iter = eplb_config.expert_heat_collection_interval - self.moe_load = torch.zeros((self.num_iter, local_num_experts), dtype=torch.int32, device="npu") - - setup_moe_comm_method(self.moe_config) - if self.multistream_overlap_shared_expert: - # Wrap the quant_method's process_weights_after_loading to validate that - # splitting shared expert computation (gate_up projection + activation, - # then down projection) yields identical results to integrated - # computation after weight loading. - original_process_weights = self._quant_method.process_weights_after_loading - - @wraps(original_process_weights) - def wrapped_process_weights(*args, **kwargs): - result = original_process_weights(*args, **kwargs) - self._validate_shared_expert_consistency() - return result - - self._quant_method.process_weights_after_loading = wrapped_process_weights # type: ignore - - # Register this MoE layer with EPLB for PP compatibility. - # PPMissingLayer (nn.Identity) never calls AscendFusedMoE.__init__, - # so only real MoE layers on this rank are registered. - VllmEplbAdaptor.register_layer(self) - - def _validate_shared_expert_consistency(self): - """Validate that split shared expert computation matches integrated computation.""" - test_input = ( - torch.rand(10, self.hidden_size, device="npu", dtype=self.moe_config.in_dtype) * 2 - 1 - ) # Random input for testing, scoped to [-1, 1] - - assert self._shared_experts is not None - integrated_out = self._shared_experts(test_input) - part1_out = self._shared_experts_part1(test_input) - split_out = self._shared_experts_part2(test_input, part1_out) - - if not torch.allclose(integrated_out, split_out): - diff = (integrated_out - split_out).abs() - logger.error( - "[fused_moe/layer] Shared expert split computation validation failed." - " The split-path computation does not match the integrated-path result." - " max_abs_diff=%s, integrated_sum=%s, integrated_norm=%s," - " split_sum=%s, split_norm=%s, hidden_size=%s, dtype=%s.", - diff.max().item(), - integrated_out.sum().item(), - integrated_out.norm().item(), - split_out.sum().item(), - split_out.norm().item(), - self.hidden_size, - self.moe_config.in_dtype, - ) - raise ValueError("FusedMoE shared experts split computation does not match the integrated computation.") - logger.info_once( - "[fused_moe/layer] Shared expert split computation validation passed." - " Integrated and split-path results are consistent." + self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp + self.multistream_overlap_shared_expert = ( + ascend_config.multistream_overlap_shared_expert and shared_experts is not None + ) + mix_placement = getattr(ascend_config, "mix_placement", False) + + # EPLB initialization (Ascend-specific; mirrors old AscendFusedMoE logic). + AscendMoERunner.moe_counter += 1 + self.moe_instance_id = AscendMoERunner.moe_counter + + eplb_config = ascend_config.eplb_config + + if mix_placement: + moe_config.num_experts += n_shared_experts + + ( + self.global_expert_map, + self._expert_map, + self.log2phy, + self.global_redundant_expert_num, + ) = init_eplb_config( + eplb_config, + AscendMoERunner.moe_counter, + moe_config, + mix_placement, + n_shared_experts, + tp_size=vllm_config.parallel_config.tensor_parallel_size, + ) + + moe_config.global_redundant_expert_num = self.global_redundant_expert_num + local_num_experts = (moe_config.num_experts + self.global_redundant_expert_num) // moe_config.ep_size + moe_config.num_local_experts = local_num_experts + routed_experts.expert_map_manager._local_num_experts = local_num_experts + routed_experts.expert_map_manager._expert_map = self._expert_map + + self.dynamic_eplb = eplb_config.dynamic_eplb and (self.log2phy is not None) + self.multi_stage = False + self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64).npu() + if self.dynamic_eplb and eplb_config.expert_heat_collection_interval > 1: + self.multi_stage = True + self.load_counter = torch.tensor(0, dtype=torch.int32, device="npu") + self.num_iter = eplb_config.expert_heat_collection_interval + self.moe_load = torch.zeros((self.num_iter, local_num_experts), dtype=torch.int32, device="npu") + + setup_moe_comm_method(self.moe_config) + if self.multistream_overlap_shared_expert: + # Wrap the quant_method's process_weights_after_loading to validate that + # splitting shared expert computation (gate_up projection + activation, + # then down projection) yields identical results to integrated + # computation after weight loading. + original_process_weights = self._quant_method.process_weights_after_loading + + @wraps(original_process_weights) + def wrapped_process_weights(*args, **kwargs): + result = original_process_weights(*args, **kwargs) + self._validate_shared_expert_consistency() + return result + + self._quant_method.process_weights_after_loading = wrapped_process_weights # type: ignore + + # Register this MoE layer with EPLB for PP compatibility. + # PPMissingLayer (nn.Identity) never calls AscendFusedMoE.__init__, + # so only real MoE layers on this rank are registered. + VllmEplbAdaptor.register_layer(self) + + def _validate_shared_expert_consistency(self): + """Validate that split shared expert computation matches integrated computation.""" + test_input = ( + torch.rand(10, self.hidden_size, device="npu", dtype=self.moe_config.in_dtype) * 2 - 1 + ) # Random input for testing, scoped to [-1, 1] + + assert self._shared_experts is not None + integrated_out = self._shared_experts(test_input) + part1_out = self._shared_experts_part1(test_input) + split_out = self._shared_experts_part2(test_input, part1_out) + + if not torch.allclose(integrated_out, split_out): + diff = (integrated_out - split_out).abs() + logger.error( + "[fused_moe/layer] Shared expert split computation validation failed." + " The split-path computation does not match the integrated-path result." + " max_abs_diff=%s, integrated_sum=%s, integrated_norm=%s," + " split_sum=%s, split_norm=%s, hidden_size=%s, dtype=%s.", + diff.max().item(), + integrated_out.sum().item(), + integrated_out.norm().item(), + split_out.sum().item(), + split_out.norm().item(), + self.hidden_size, + self.moe_config.in_dtype, ) + raise ValueError("FusedMoE shared experts split computation does not match the integrated computation.") + logger.info_once( + "[fused_moe/layer] Shared expert split computation validation passed." + " Integrated and split-path results are consistent." + ) - def _shared_experts_part1(self, hidden_states: torch.Tensor): - shared_gate_up, _ = self._shared_experts.gate_up_proj(hidden_states) # type: ignore - return shared_gate_up + def _shared_experts_part1(self, hidden_states: torch.Tensor): + shared_gate_up, _ = self._shared_experts.gate_up_proj(hidden_states) # type: ignore + return shared_gate_up - def _shared_experts_part2(self, hidden_states: torch.Tensor, shared_gate_up: torch.Tensor): - shared_act = self._shared_experts.act_fn(shared_gate_up) # type: ignore - shared_out, _ = self._shared_experts.down_proj(shared_act) # type: ignore + def _shared_experts_part2(self, hidden_states: torch.Tensor, shared_gate_up: torch.Tensor): + shared_act = self._shared_experts.act_fn(shared_gate_up) # type: ignore + shared_out, _ = self._shared_experts.down_proj(shared_act) # type: ignore - # Qwen3-Next specific gating mechanism - assert self._shared_experts is not None - if hasattr(self._shared_experts, "expert_gate") and self._shared_experts.expert_gate is not None: - gate_out, _ = self._shared_experts.expert_gate(hidden_states) # type: ignore - shared_out = F.sigmoid(gate_out) * shared_out - return shared_out + # Qwen3-Next specific gating mechanism + assert self._shared_experts is not None + if hasattr(self._shared_experts, "expert_gate") and self._shared_experts.expert_gate is not None: + gate_out, _ = self._shared_experts.expert_gate(hidden_states) # type: ignore + shared_out = F.sigmoid(gate_out) * shared_out + return shared_out - def _get_quant_type(self) -> QuantType: - quant_type = QuantType.NONE - method = getattr(self._quant_method, "quant_method", None) + def _get_quant_type(self) -> QuantType: + quant_type = QuantType.NONE + method = getattr(self._quant_method, "quant_method", None) - if method is not None: - quant_type = getattr(method, "quant_type", QuantType.NONE) + if method is not None: + quant_type = getattr(method, "quant_type", QuantType.NONE) - return quant_type + return quant_type - @property - def is_internal_router(self) -> bool: - gate = self.gate - return gate is not None and hasattr(gate, "weight_fp32") - - @property - def use_dp_chunking(self) -> bool: - """Ascend uses its own forward_impl path, not the FlashInfer Cutlass - chunked path. Always return False to stay on forward_impl.""" - return False - - @property - def _fused_output_is_reduced(self) -> bool: - # For MC2/ALLTOALL/FUSED_MC2 comm types, finalize() already includes - # TP all-reduce for the routed output, and _forward_shared_experts - # handles it for the shared output. Signal this to the upstream - # MoERunner.forward() so _maybe_reduce_final_output does not apply a - # second TP all-reduce (which would double-count the contributions). - moe_comm_type = _EXTRA_CTX.moe_comm_type - return moe_comm_type in { - MoECommType.ALLTOALL, - MoECommType.MC2, - MoECommType.FUSED_MC2, - } or (moe_comm_type == MoECommType.ALLGATHER and _EXTRA_CTX.flash_comm_v1_enabled) - - def _maybe_reduce_shared_expert_output( - self, - shared_output: torch.Tensor | None, - ) -> torch.Tensor | None: - # _forward_shared_experts already handles shared expert TP all-reduce - # for MC2/ALLTOALL/FUSED_MC2. For AllGather the reduction is done - # via _maybe_reduce_final_output on the combined (shared + routed) - # output. Skip any additional reduction here. - return shared_output - - def _maybe_reduce_final_output( - self, - states: torch.Tensor, - trunc_size: int, - ) -> torch.Tensor: - states = torch.ops.vllm.maybe_all_reduce_tensor_model_parallel(states) - return states[..., :trunc_size] - - def set_lora_context(self, lora_context): - self.routed_experts._ascend_moe_lora_context = lora_context - - def no_shared_forward_impl( # type: ignore[override] - self, hidden_states: torch.Tensor, router_logits: torch.Tensor, return_with_event: bool = False - ) -> torch.Tensor | FusedMoEResult: - forward_context = get_forward_context() - # When static kernels are enabled, the forward pass runs twice (compilation + capture), - # causing moe_layer_index to overflow. Wrap the index to prevent out-of-bounds errors. - if self.enable_npugraph_ex_static_kernel and forward_context.all_moe_layers: - moe_layer_index = forward_context.moe_layer_index % (len(forward_context.all_moe_layers)) - forward_context.moe_layer_index = moe_layer_index - - # Load balancing for token distribution among experts in dummy_run - # TODO: The community only considers load balancing when DP > 1. - # This approach may overlook some extreme scenarios. - enable_force_load_balance = _EXTRA_CTX.in_profile_run - - prepare_output = _EXTRA_CTX.moe_comm_method.prepare( - hidden_states=hidden_states, - router_logits=router_logits, - replace_allreduce=_EXTRA_CTX.flash_comm_v1_enabled, - enable_shared_expert_dp=self.enable_shared_expert_dp, - quant_type=self.quant_type, + @property + def is_internal_router(self) -> bool: + gate = self.gate + return gate is not None and hasattr(gate, "weight_fp32") + + @property + def use_dp_chunking(self) -> bool: + """Ascend uses its own forward_impl path, not the FlashInfer Cutlass + chunked path. Always return False to stay on forward_impl.""" + return False + + @property + def _fused_output_is_reduced(self) -> bool: + # For MC2/ALLTOALL/FUSED_MC2 comm types, finalize() already includes + # TP all-reduce for the routed output, and _forward_shared_experts + # handles it for the shared output. Signal this to the upstream + # MoERunner.forward() so _maybe_reduce_final_output does not apply a + # second TP all-reduce (which would double-count the contributions). + moe_comm_type = _EXTRA_CTX.moe_comm_type + return moe_comm_type in { + MoECommType.ALLTOALL, + MoECommType.MC2, + MoECommType.FUSED_MC2, + } or (moe_comm_type == MoECommType.ALLGATHER and _EXTRA_CTX.flash_comm_v1_enabled) + + def _maybe_reduce_shared_expert_output( + self, + shared_output: torch.Tensor | None, + ) -> torch.Tensor | None: + # _forward_shared_experts already handles shared expert TP all-reduce + # for MC2/ALLTOALL/FUSED_MC2. For AllGather the reduction is done + # via _maybe_reduce_final_output on the combined (shared + routed) + # output. Skip any additional reduction here. + return shared_output + + def _maybe_reduce_final_output( + self, + states: torch.Tensor, + trunc_size: int, + ) -> torch.Tensor: + states = torch.ops.vllm.maybe_all_reduce_tensor_model_parallel(states) + return states[..., :trunc_size] + + def set_lora_context(self, lora_context): + self.routed_experts._ascend_moe_lora_context = lora_context + + def no_shared_forward_impl( # type: ignore[override] + self, hidden_states: torch.Tensor, router_logits: torch.Tensor, return_with_event: bool = False + ) -> torch.Tensor | FusedMoEResult: + forward_context = get_forward_context() + # When static kernels are enabled, the forward pass runs twice (compilation + capture), + # causing moe_layer_index to overflow. Wrap the index to prevent out-of-bounds errors. + if self.enable_npugraph_ex_static_kernel and forward_context.all_moe_layers: + moe_layer_index = forward_context.moe_layer_index % (len(forward_context.all_moe_layers)) + forward_context.moe_layer_index = moe_layer_index + + # Load balancing for token distribution among experts in dummy_run + # TODO: The community only considers load balancing when DP > 1. + # This approach may overlook some extreme scenarios. + enable_force_load_balance = _EXTRA_CTX.in_profile_run + + prepare_output = _EXTRA_CTX.moe_comm_method.prepare( + hidden_states=hidden_states, + router_logits=router_logits, + replace_allreduce=_EXTRA_CTX.flash_comm_v1_enabled, + enable_shared_expert_dp=self.enable_shared_expert_dp, + quant_type=self.quant_type, + ) + hidden_states = prepare_output.hidden_states + router_logits = prepare_output.router_logits + mc2_mask = prepare_output.mc2_mask + padded_hidden_states_shape = prepare_output.padded_hidden_states_shape + pertoken_scale = prepare_output.pertoken_scale + + # Matrix multiply. + # apply() expects a RoutedExperts-like layer for weight access + # (w13_weight, w2_weight, swiglu_limit, etc.). Pass routed_experts, + # not self; the routing params come through the other kwargs. + fused_experts_results: FusedExpertsResult = self._quant_method.apply( + layer=self.routed_experts, + x=hidden_states, + router_logits=router_logits, + pertoken_scale=pertoken_scale, + top_k=self.top_k, + renormalize=self.renormalize, + use_grouped_topk=self.use_grouped_topk, + num_experts=self.moe_config.num_experts, + expert_map=self._expert_map, + topk_group=self.topk_group, + num_expert_group=self.num_expert_group, + custom_routing_function=self.custom_routing_function, + scoring_func=self.scoring_func, + routed_scaling_factor=self._original_routed_scaling_factor, + e_score_correction_bias=self.e_score_correction_bias, + activation=self.activation, + apply_router_weight_on_input=self.apply_router_weight_on_input, + enable_force_load_balance=enable_force_load_balance, + log2phy=self.log2phy, + global_redundant_expert_num=self.global_redundant_expert_num, + mc2_mask=mc2_mask, + ) + + if self.dynamic_eplb and _EXTRA_CTX.eplb_heat_collection_status: + expert_tokens = fused_experts_results.expert_tokens + group_list_type = fused_experts_results.group_list_type + assert expert_tokens is not None and group_list_type is not None, ( + "expert_tokens and group_list_type should not be None when dynamic_eplb is enabled." ) - hidden_states = prepare_output.hidden_states - router_logits = prepare_output.router_logits - mc2_mask = prepare_output.mc2_mask - padded_hidden_states_shape = prepare_output.padded_hidden_states_shape - pertoken_scale = prepare_output.pertoken_scale - - # Matrix multiply. - # apply() expects a RoutedExperts-like layer for weight access - # (w13_weight, w2_weight, swiglu_limit, etc.). Pass routed_experts, - # not self; the routing params come through the other kwargs. - fused_experts_results: FusedExpertsResult = self._quant_method.apply( - layer=self.routed_experts, - x=hidden_states, - router_logits=router_logits, - pertoken_scale=pertoken_scale, - top_k=self.top_k, - renormalize=self.renormalize, - use_grouped_topk=self.use_grouped_topk, - num_experts=self.moe_config.num_experts, - expert_map=self._expert_map, - topk_group=self.topk_group, - num_expert_group=self.num_expert_group, - custom_routing_function=self.custom_routing_function, - scoring_func=self.scoring_func, - routed_scaling_factor=self._original_routed_scaling_factor, - e_score_correction_bias=self.e_score_correction_bias, - activation=self.activation, - apply_router_weight_on_input=self.apply_router_weight_on_input, - enable_force_load_balance=enable_force_load_balance, - log2phy=self.log2phy, - global_redundant_expert_num=self.global_redundant_expert_num, - mc2_mask=mc2_mask, + local_load = ( + expert_tokens + if group_list_type == 1 + else torch.cat([expert_tokens[:1], expert_tokens[1:] - expert_tokens[:-1]]) ) - - if self.dynamic_eplb and _EXTRA_CTX.eplb_heat_collection_status: - expert_tokens = fused_experts_results.expert_tokens - group_list_type = fused_experts_results.group_list_type - assert expert_tokens is not None and group_list_type is not None, ( - "expert_tokens and group_list_type should not be None when dynamic_eplb is enabled." - ) - local_load = ( - expert_tokens - if group_list_type == 1 - else torch.cat([expert_tokens[:1], expert_tokens[1:] - expert_tokens[:-1]]) + if self.multi_stage: + cur_iter = torch.remainder(self.load_counter, self.num_iter) + self.moe_load.index_add_( + dim=0, index=cur_iter, source=local_load.to(torch.int32, non_blocking=True).view(1, -1) ) - if self.multi_stage: - cur_iter = torch.remainder(self.load_counter, self.num_iter) - self.moe_load.index_add_( - dim=0, index=cur_iter, source=local_load.to(torch.int32, non_blocking=True).view(1, -1) - ) - self.load_counter.add_(1) - else: - self.moe_load.add_(local_load) - - routed_out = _EXTRA_CTX.moe_comm_method.finalize( - hidden_states=fused_experts_results.routed_out, - reduce_results=isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl), - padded_hidden_states_shape=padded_hidden_states_shape, + self.load_counter.add_(1) + else: + self.moe_load.add_(local_load) + + routed_out = _EXTRA_CTX.moe_comm_method.finalize( + hidden_states=fused_experts_results.routed_out, + reduce_results=isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl), + padded_hidden_states_shape=padded_hidden_states_shape, + ) + + if return_with_event: + return FusedMoEResult( + routed_out=routed_out, + before_dispatch_evt=fused_experts_results.before_dispatch_evt, + before_gmm2_evt=fused_experts_results.before_gmm2_evt, + before_combine_evt=fused_experts_results.before_combine_evt, + swiglu_limit=fused_experts_results.swiglu_limit, ) + else: + # The vLLM FusedMoE forward_impl does not return events. + return routed_out + + def _forward_shared_experts(self, hidden_states: torch.Tensor, fused_moe_evts: FusedMoEEvents): + if self._shared_experts is None: + return None + + def maybe_wait_event(evt: torch.npu.Event | None): + if evt is not None: + torch.npu.current_stream().wait_event(evt) - if return_with_event: - return FusedMoEResult( - routed_out=routed_out, - before_dispatch_evt=fused_experts_results.before_dispatch_evt, - before_gmm2_evt=fused_experts_results.before_gmm2_evt, - before_combine_evt=fused_experts_results.before_combine_evt, - swiglu_limit=fused_experts_results.swiglu_limit, + with npu_stream_switch(shared_experts_calculation_stream(), enabled=self.multistream_overlap_shared_expert): + # Only used for int quantization + has_quantized_shared = hasattr(self._shared_experts.gate_up_proj, "weight_scale") and hasattr( + self._shared_experts.down_proj, "weight_scale" + ) + if has_quantized_shared and self.quant_type in (QuantType.W8A8, QuantType.W4A8): + original_dtype = hidden_states.dtype + # Execute dynamic quant concurrently with MoE gate. + torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) + quantized_x, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) + # Execute the gate projection and activation concurrently with the + # dispatch communication. + maybe_wait_event(fused_moe_evts.after_routed_experts) + hidden_states = torch_npu.npu_quant_matmul( + quantized_x, + self._shared_experts.gate_up_proj.weight, + self._shared_experts.gate_up_proj.weight_scale, + pertoken_scale=None, + bias=None, + output_dtype=torch.int32, + ) + # Execute activation concurrently with gmm2. + + maybe_wait_event(fused_moe_evts.before_gmm2) + quantized_x, swiglu_out_scale = torch.ops._C_ascend.npu_dequant_swiglu_quant( + x=hidden_states, + weight_scale=self._shared_experts.gate_up_proj.weight_scale_fp32, + activation_scale=pertoken_scale, + bias=None, + quant_scale=None, + quant_offset=None, + group_index=None, + activate_left=True, + quant_mode=1, + swiglu_mode=1, + clamp_limit=fused_moe_evts.swiglu_limit, + ) + # Execute the down projection concurrently with the combine + # communication. + maybe_wait_event(fused_moe_evts.before_combine) + shared_out = torch_npu.npu_quant_matmul( + quantized_x, + self._shared_experts.down_proj.weight, + self._shared_experts.down_proj.weight_scale, + pertoken_scale=swiglu_out_scale, + bias=None, + output_dtype=original_dtype, + ) + elif has_quantized_shared and self.quant_type == QuantType.W4A8MXFP: + original_dtype = hidden_states.dtype + # Execute dynamic quant concurrently with MoE gate. + torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) + quantized_x, pertoken_scale = torch_npu.npu_dynamic_mx_quant( + hidden_states, dst_type=torch.float8_e4m3fn + ) + # Execute the gate projection and activation concurrently with the + # dispatch communication. + maybe_wait_event(fused_moe_evts.before_dispatch) + hidden_states = self._shared_experts.gate_up_proj((quantized_x, pertoken_scale))[0] + # Execute activation concurrently with gmm2. + maybe_wait_event(fused_moe_evts.before_gmm2) + quantized_x, swiglu_out_scale, _ = torch.ops._C_ascend.npu_swiglu_group_quant( + hidden_states, + topk_weight=None, + group_index=None, + dst_type=torch.float8_e4m3fn, + quant_mode=2, + clamp_value=fused_moe_evts.swiglu_limit, ) + # Execute the down projection concurrently with the combine + # communication. + maybe_wait_event(fused_moe_evts.before_combine) + shared_out = self._shared_experts.down_proj((quantized_x, swiglu_out_scale))[0] else: - # The vLLM FusedMoE forward_impl does not return events. - return routed_out + # Ensure the shared experts wait for hidden_states to be ready. + torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) + # Execute the gate projection and activation concurrently with the + # dispatch communication. + maybe_wait_event(fused_moe_evts.before_dispatch) + part1_out = self._shared_experts_part1(hidden_states) + # Execute the down projection concurrently with the combine + # communication. + maybe_wait_event(fused_moe_evts.before_combine) + shared_out = self._shared_experts_part2(hidden_states, part1_out) + + # Make sure the default stream waits for the shared experts stream to + # finish. + if self.multistream_overlap_shared_expert: + torch.npu.current_stream().wait_stream(shared_experts_calculation_stream()) + + # NOTE: This is exactly the opposite of + # `maybe_all_reduce_tensor_model_parallel` + moe_comm_type = _EXTRA_CTX.moe_comm_type + if ( + moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_MC2} + and not shared_expert_dp_enabled() + ): + shared_out = tensor_model_parallel_all_reduce(shared_out) + return shared_out - def _forward_shared_experts(self, hidden_states: torch.Tensor, fused_moe_evts: FusedMoEEvents): - if self._shared_experts is None: - return None + def shared_forward_impl( # type: ignore[override] + self, hidden_states: torch.Tensor, router_logits: torch.Tensor + ): + if self.is_internal_router: + gate = self.gate + assert gate is not None + # NOTE(Angazenn): To make this cast explicitly, the hbm usage might + # increase with extra hidden states. We also assume that all gate + # linear is unquantized so that we the weight is pre-casted in + # process_weights_after_loading of AscendUnquantizedLinearMethod. + hidden_states_fp32 = hidden_states.float() + before_routed_experts = torch.npu.current_stream().record_event() + router_logits = F.linear(hidden_states_fp32, gate.weight_fp32) + after_routed_experts = torch.npu.current_stream().record_event() + else: + before_routed_experts = torch.npu.current_stream().record_event() + after_routed_experts = None - def maybe_wait_event(evt: torch.npu.Event | None): - if evt is not None: - torch.npu.current_stream().wait_event(evt) + fused_moe_results = self.no_shared_forward_impl( + hidden_states, + router_logits, + return_with_event=True, + ) + routed_out = fused_moe_results.routed_out + + if self._shared_experts is None: + return routed_out + + shared_out = self._forward_shared_experts( + hidden_states, + FusedMoEEvents( + after_routed_experts=after_routed_experts, + before_routed_experts=before_routed_experts, + before_dispatch=fused_moe_results.before_dispatch_evt, + before_gmm2=fused_moe_results.before_gmm2_evt, + before_combine=fused_moe_results.before_combine_evt, + swiglu_limit=fused_moe_results.swiglu_limit, + ), + ) + return shared_out, routed_out - with npu_stream_switch(shared_experts_calculation_stream(), enabled=self.multistream_overlap_shared_expert): - # Only used for int quantization - has_quantized_shared = hasattr(self._shared_experts.gate_up_proj, "weight_scale") and hasattr( - self._shared_experts.down_proj, "weight_scale" - ) - if has_quantized_shared and self.quant_type in (QuantType.W8A8, QuantType.W4A8): - original_dtype = hidden_states.dtype - # Execute dynamic quant concurrently with MoE gate. - torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) - quantized_x, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) - # Execute the gate projection and activation concurrently with the - # dispatch communication. - maybe_wait_event(fused_moe_evts.after_routed_experts) - hidden_states = torch_npu.npu_quant_matmul( - quantized_x, - self._shared_experts.gate_up_proj.weight, - self._shared_experts.gate_up_proj.weight_scale, - pertoken_scale=None, - bias=None, - output_dtype=torch.int32, - ) - # Execute activation concurrently with gmm2. - - maybe_wait_event(fused_moe_evts.before_gmm2) - quantized_x, swiglu_out_scale = torch.ops._C_ascend.npu_dequant_swiglu_quant( - x=hidden_states, - weight_scale=self._shared_experts.gate_up_proj.weight_scale_fp32, - activation_scale=pertoken_scale, - bias=None, - quant_scale=None, - quant_offset=None, - group_index=None, - activate_left=True, - quant_mode=1, - swiglu_mode=1, - clamp_limit=fused_moe_evts.swiglu_limit, - ) - # Execute the down projection concurrently with the combine - # communication. - maybe_wait_event(fused_moe_evts.before_combine) - shared_out = torch_npu.npu_quant_matmul( - quantized_x, - self._shared_experts.down_proj.weight, - self._shared_experts.down_proj.weight_scale, - pertoken_scale=swiglu_out_scale, - bias=None, - output_dtype=original_dtype, - ) - elif has_quantized_shared and self.quant_type == QuantType.W4A8MXFP: - original_dtype = hidden_states.dtype - # Execute dynamic quant concurrently with MoE gate. - torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) - quantized_x, pertoken_scale = torch_npu.npu_dynamic_mx_quant( - hidden_states, dst_type=torch.float8_e4m3fn - ) - # Execute the gate projection and activation concurrently with the - # dispatch communication. - maybe_wait_event(fused_moe_evts.before_dispatch) - hidden_states = self._shared_experts.gate_up_proj((quantized_x, pertoken_scale))[0] - # Execute activation concurrently with gmm2. - maybe_wait_event(fused_moe_evts.before_gmm2) - quantized_x, swiglu_out_scale, _ = torch.ops._C_ascend.npu_swiglu_group_quant( - hidden_states, - topk_weight=None, - group_index=None, - dst_type=torch.float8_e4m3fn, - quant_mode=2, - clamp_value=fused_moe_evts.swiglu_limit, - ) - # Execute the down projection concurrently with the combine - # communication. - maybe_wait_event(fused_moe_evts.before_combine) - shared_out = self._shared_experts.down_proj((quantized_x, swiglu_out_scale))[0] - else: - # Ensure the shared experts wait for hidden_states to be ready. - torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) - # Execute the gate projection and activation concurrently with the - # dispatch communication. - maybe_wait_event(fused_moe_evts.before_dispatch) - part1_out = self._shared_experts_part1(hidden_states) - # Execute the down projection concurrently with the combine - # communication. - maybe_wait_event(fused_moe_evts.before_combine) - shared_out = self._shared_experts_part2(hidden_states, part1_out) - - # Make sure the default stream waits for the shared experts stream to - # finish. - if self.multistream_overlap_shared_expert: - torch.npu.current_stream().wait_stream(shared_experts_calculation_stream()) - - # NOTE: This is exactly the opposite of - # `maybe_all_reduce_tensor_model_parallel` - moe_comm_type = _EXTRA_CTX.moe_comm_type - if ( - moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_MC2} - and not shared_expert_dp_enabled() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out - - def shared_forward_impl( # type: ignore[override] - self, hidden_states: torch.Tensor, router_logits: torch.Tensor - ): - if self.is_internal_router: - gate = self.gate - assert gate is not None - # NOTE(Angazenn): To make this cast explicitly, the hbm usage might - # increase with extra hidden states. We also assume that all gate - # linear is unquantized so that we the weight is pre-casted in - # process_weights_after_loading of AscendUnquantizedLinearMethod. - hidden_states_fp32 = hidden_states.float() - before_routed_experts = torch.npu.current_stream().record_event() - router_logits = F.linear(hidden_states_fp32, gate.weight_fp32) - after_routed_experts = torch.npu.current_stream().record_event() + def _forward_impl( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + shared_experts_input: torch.Tensor | None, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + with self._sequence_parallel_context(): + if self.shared_experts is None: + return self.no_shared_forward_impl(hidden_states, router_logits) else: - before_routed_experts = torch.npu.current_stream().record_event() - after_routed_experts = None - - fused_moe_results = self.no_shared_forward_impl( - hidden_states, - router_logits, - return_with_event=True, - ) - routed_out = fused_moe_results.routed_out - - if self._shared_experts is None: - return routed_out - - shared_out = self._forward_shared_experts( - hidden_states, - FusedMoEEvents( - after_routed_experts=after_routed_experts, - before_routed_experts=before_routed_experts, - before_dispatch=fused_moe_results.before_dispatch_evt, - before_gmm2=fused_moe_results.before_gmm2_evt, - before_combine=fused_moe_results.before_combine_evt, - swiglu_limit=fused_moe_results.swiglu_limit, - ), - ) - return shared_out, routed_out - - def _forward_impl( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - input_ids: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - with self._sequence_parallel_context(): - if self.shared_experts is None: - return self.no_shared_forward_impl(hidden_states, router_logits) - else: - return self.shared_forward_impl(hidden_states, router_logits) + return self.shared_forward_impl(hidden_states, router_logits) diff --git a/vllm_ascend/ops/fused_moe/fused_moe_0_23_0.py b/vllm_ascend/ops/fused_moe/fused_moe_0_23_0.py deleted file mode 100644 index ef6b8e15a87..00000000000 --- a/vllm_ascend/ops/fused_moe/fused_moe_0_23_0.py +++ /dev/null @@ -1,652 +0,0 @@ -# -# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. -# This file is a part of the vllm-ascend project. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# ruff: noqa: E501 -"""Legacy vLLM 0.23.0 FusedMoE implementation. - -Private module. Import AscendFusedMoE and AscendMoERunner through -vllm_ascend.ops.fused_moe.fused_moe only. -""" - -from __future__ import annotations - -from vllm_ascend.ops.fused_moe.fused_moe import ( - _EXTRA_CTX, - AllGatherCommImpl, - AscendUnquantizedFusedMoEMethod, - F, - FusedExpertsResult, - FusedMoE, - FusedMoEEvents, - FusedMoEResult, - MoECommType, - MoERunner, - QuantType, - VllmEplbAdaptor, - get_ascend_config, - get_compressed_expert_map, - get_current_vllm_config, - get_dp_group, - get_ep_group, - get_forward_context, - get_mc2_group, - get_tp_group, - init_eplb_config, - logger, - npu_stream_switch, - setup_moe_comm_method, - shared_expert_dp_enabled, - shared_experts_calculation_stream, - tensor_model_parallel_all_reduce, - torch, - torch_npu, - wraps, -) -from vllm_ascend.utils import enable_sp - - -class AscendMoERunner(MoERunner): - @property - def use_dp_chunking(self) -> bool: - """Ascend uses its own forward_impl path, not the FlashInfer Cutlass - chunked path. Always return False to stay on forward_impl.""" - return False - - @property - def _fused_output_is_reduced(self) -> bool: - # For MC2/ALLTOALL/FUSED_MC2 comm types, finalize() already includes - # TP all-reduce for the routed output, and _forward_shared_experts - # handles it for the shared output. Signal this to the upstream - # MoERunner.forward() so _maybe_reduce_final_output does not apply a - # second TP all-reduce (which would double-count the contributions). - moe_comm_type = _EXTRA_CTX.moe_comm_type - return moe_comm_type in { - MoECommType.ALLTOALL, - MoECommType.MC2, - MoECommType.FUSED_MC2, - } or (moe_comm_type == MoECommType.ALLGATHER and _EXTRA_CTX.flash_comm_v1_enabled) - - def _maybe_reduce_shared_expert_output( - self, - shared_output: torch.Tensor | None, - ) -> torch.Tensor | None: - # _forward_shared_experts already handles shared expert TP all-reduce - # for MC2/ALLTOALL/FUSED_MC2. For AllGather the reduction is done - # via _maybe_reduce_final_output on the combined (shared + routed) - # output. Skip any additional reduction here. - return shared_output - - def _maybe_reduce_final_output( - self, - states: torch.Tensor, - trunc_size: int, - ) -> torch.Tensor: - states = torch.ops.vllm.maybe_all_reduce_tensor_model_parallel(states) - return states[..., :trunc_size] - - # TODO: Remove this after drop v0.19.1 support - def forward_impl( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - Override the default forward_impl to use Ascend-specific implementation. - This delegates to the layer's forward_impl method which contains the - Ascend-specific MoE computation logic. - """ - if self.shared_experts is None: - result = layer.forward_impl(hidden_states, router_logits) - # If the layer has shared experts, forward_impl returns a tuple (shared_out, routed_out) - # Otherwise, it returns just routed_out - # The torch op expects the same return type based on whether it's moe_forward or moe_forward_shared - else: - result = layer.shared_forward_impl(hidden_states, router_logits) - return result - - def _forward_impl( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - input_ids: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - with self._sequence_parallel_context(): - return self.forward_impl( - layer, - hidden_states, - router_logits, - shared_experts_input, - ) - - -class AscendFusedMoE(FusedMoE): - moe_counter = -1 - - def __init__(self, *args, **kwargs): - # Save original routed_scaling_factor before super().__init__ modifies it. - # When apply_routed_scale_to_output=True, vLLM sets self.routed_scaling_factor - # to 1.0 and expects the runner to apply scaling to output. But vllm-ascend - # uses its own forward path, so we need the original value. - _ = kwargs.pop("hash") if "hash" in kwargs else None - tid2eid = kwargs.pop("tid2eid") if "tid2eid" in kwargs else None - - self._original_routed_scaling_factor = kwargs.get("routed_scaling_factor", 1.0) - super().__init__(*args, **kwargs) - self.use_overlapped = True - self._routed_input_transform = kwargs.get("routed_input_transform") - self._shared_experts = kwargs.get("shared_experts") - self.shared_expert_stream = None - has_shared_experts = self._shared_experts is not None - num_experts = kwargs["num_experts"] - intermediate_size = kwargs["intermediate_size"] - num_shared_experts = kwargs.get("n_shared_experts", 0) - - AscendFusedMoE.moe_counter += 1 - self.moe_instance_id = AscendFusedMoE.moe_counter - - self._expert_map = None - self.log2phy = None - - self.tid2eid = tid2eid - - if self.quant_config is None: - self.quant_method = AscendUnquantizedFusedMoEMethod(self.moe_config, tid2eid=self.tid2eid) - else: - self.quant_method = self.quant_config.get_quant_method(self, self.layer_name, tid2eid=self.tid2eid) - - assert self.quant_method is not None - # Keep base_quant_method in sync with the swapped-in Ascend method, - # otherwise FusedMoE.maybe_init_modular_kernel (called via the V2 - # model runner's prepare_communication_buffer_for_model) would dispatch - # to the upstream UnquantizedFusedMoEMethod.maybe_make_prepare_finalize, - # which raises by design. - self.base_quant_method = self.quant_method - - self.moe_config.tp_group = get_tp_group() - self.moe_config.dp_group = get_dp_group() - if self.moe_config.ep_size > 1: - self.moe_config.ep_group = get_ep_group() - self.moe_config.mc2_group = get_mc2_group() - self.moe_config.supports_eplb = self.quant_method.supports_eplb - ascend_config = get_ascend_config() - self.multistream_overlap_shared_expert = ascend_config.multistream_overlap_shared_expert and has_shared_experts - if self.multistream_overlap_shared_expert: - logger.info_once("[fused_moe/layer] Multistream overlap shared expert is enabled.") - if enable_sp() and has_shared_experts: - logger.info_once( - "[fused_moe/layer] Sequence parallelism is enabled, shared experts are replicated for best performance." - ) - - vllm_config = get_current_vllm_config() - if ( - self.custom_routing_function is None - and self.e_score_correction_bias is not None - and not vllm_config.model_config.is_deepseek_mla - ): - self.e_score_correction_bias.data = self.e_score_correction_bias.data.to( - dtype=vllm_config.model_config.dtype - ) - self._gate = kwargs.get("gate") - - # init moe - eplb_config = ascend_config.eplb_config - self.mix_placement = getattr(ascend_config, "mix_placement", False) - self.n_shared_experts = num_shared_experts - num_experts += num_shared_experts if self.mix_placement else 0 - self.moe_config.num_experts = num_experts - self.global_expert_map, self._expert_map, self.log2phy, self.global_redundant_expert_num = init_eplb_config( - eplb_config, - self.moe_instance_id, - self.moe_config, - self.mix_placement, - num_shared_experts, - tp_size=self.vllm_config.parallel_config.tensor_parallel_size, - ) - self.global_num_experts = num_experts + self.global_redundant_expert_num - self.dynamic_eplb = eplb_config.dynamic_eplb and (self.log2phy is not None) - self.local_num_experts = self.global_num_experts // self.ep_size - self.expert_map_manager._local_num_experts = self.local_num_experts - self.expert_map_manager._expert_map = self._expert_map - if self._expert_map is not None: - logger.info_once( - "[fused_moe/layer] Expert parallelism is enabled." - " ep_rank=%s/%s, local_num_experts=%s, global_num_experts=%s," - " expert_map=%s", - self.ep_rank, - self.ep_size, - self.local_num_experts, - self.global_num_experts, - get_compressed_expert_map(self._expert_map), - ) - if self.dynamic_eplb: - self.multi_stage = False - self.moe_load = torch.zeros(self.local_num_experts, dtype=torch.int64).npu() - if eplb_config.eplb_policy_type == 3: - self.multi_stage = True - self.load_counter = torch.tensor(0, dtype=torch.int32, device="npu") - self.num_iter = eplb_config.expert_heat_collection_interval - self.moe_load = torch.zeros((self.num_iter, self.local_num_experts), dtype=torch.int32, device="npu") - - self.moe_config.num_experts = self.global_num_experts - self.moe_config.num_local_experts = self.local_num_experts - self.moe_config.global_redundant_expert_num = self.global_redundant_expert_num - self.swiglu_limit = getattr(self.vllm_config.model_config.hf_config, "swiglu_limit", 0) - - moe_quant_params = { - "num_experts": self.local_num_experts, - "hidden_size": self.hidden_size, - "intermediate_size_per_partition": self.intermediate_size_per_partition, - "params_dtype": self.params_dtype, - "weight_loader": self.weight_loader, - } - # need full intermediate size pre-sharding for WNA16 act order - if self.quant_method.__class__.__name__ in ("GPTQMarlinMoEMethod", "CompressedTensorsWNA16MoEMethod"): - moe_quant_params["intermediate_size_full"] = intermediate_size - self.quant_method.create_weights(layer=self, **moe_quant_params) - - self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp - self.enable_npugraph_ex_static_kernel = ascend_config.ascend_compilation_config.enable_static_kernel - - setup_moe_comm_method(self.moe_config) - self.quant_type = self._get_quant_type() - - self.runner = AscendMoERunner( - self.layer_name, - self.moe_config, - self.router, - self._routed_input_transform, - kwargs.pop("gate", None), - kwargs.pop("shared_experts", None), - self.quant_method, - self.vllm_config.parallel_config.enable_dbo, - ) - - if self.multistream_overlap_shared_expert: - # Wrap the quant_method's process_weights_after_loading to validate that - # splitting shared expert computation (gate_up projection + activation, - # then down projection) yields identical results to integrated - # computation after weight loading. - original_process_weights = self.quant_method.process_weights_after_loading - - @wraps(original_process_weights) - def wrapped_process_weights(*args, **kwargs): - result = original_process_weights(*args, **kwargs) - self._validate_shared_expert_consistency() - return result - - self.quant_method.process_weights_after_loading = wrapped_process_weights # type: ignore - - # Register this MoE layer with EPLB for PP compatibility. - # PPMissingLayer (nn.Identity) never calls AscendFusedMoE.__init__, - # so only real MoE layers on this rank are registered. - VllmEplbAdaptor.register_layer(self) - - def _validate_shared_expert_consistency(self): - """Validate that split shared expert computation matches integrated - computation.""" - test_input = ( - torch.rand(10, self.hidden_size, device="npu", dtype=self.moe_config.in_dtype) * 2 - 1 - ) # Random input for testing, scoped to [-1, 1] - - assert self._shared_experts is not None - integrated_out = self._shared_experts(test_input) - part1_out = self._shared_experts_part1(test_input) - split_out = self._shared_experts_part2(test_input, part1_out) - - if not torch.allclose(integrated_out, split_out): - diff = (integrated_out - split_out).abs() - logger.error( - "[fused_moe/layer] Shared expert split computation validation failed." - " The split-path computation does not match the integrated-path result." - " max_abs_diff=%s, integrated_sum=%s, integrated_norm=%s," - " split_sum=%s, split_norm=%s, hidden_size=%s, dtype=%s.", - diff.max().item(), - integrated_out.sum().item(), - integrated_out.norm().item(), - split_out.sum().item(), - split_out.norm().item(), - self.hidden_size, - self.moe_config.in_dtype, - ) - raise ValueError("FusedMoE shared experts split computation does not match the integrated computation.") - logger.info_once( - "[fused_moe/layer] Shared expert split computation validation passed." - " Integrated and split-path results are consistent." - ) - - def _shared_experts_part1(self, hidden_states: torch.Tensor): - shared_gate_up, _ = self._shared_experts.gate_up_proj(hidden_states) # type: ignore - return shared_gate_up - - def _shared_experts_part2(self, hidden_states: torch.Tensor, shared_gate_up: torch.Tensor): - shared_act = self._shared_experts.act_fn(shared_gate_up) # type: ignore - shared_out, _ = self._shared_experts.down_proj(shared_act) # type: ignore - - # Qwen3-Next specific gating mechanism - assert self._shared_experts is not None - if hasattr(self._shared_experts, "expert_gate") and self._shared_experts.expert_gate is not None: - gate_out, _ = self._shared_experts.expert_gate(hidden_states) # type: ignore - shared_out = F.sigmoid(gate_out) * shared_out - return shared_out - - def _get_quant_type(self) -> QuantType: - quant_type = QuantType.NONE - method = getattr(self.quant_method, "quant_method", None) - - if method is not None: - quant_type = getattr(method, "quant_type", QuantType.NONE) - - return quant_type - - def set_lora_context(self, lora_context): - self._ascend_moe_lora_context = lora_context - - def update_expert_map(self, new_expert_map): - self._expert_map = new_expert_map - - def get_log2phy_map(self): - return self.log2phy - - def clear_moe_load(self): - if self.moe_load is not None: - self.moe_load.zero_() - if self.multi_stage: - self.load_counter.zero_() - - def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor): - """NOTE(Yizhou): This is to override the parent class method. In `mc2commimpl`, - and `alltoallcommimpl`, we do not need to all-reduce the final outputs since - the outputs are already aggregated across tensor parallel ranks in the - `finalize` function. In `allgathercommimpl`, we still need to all-reduce the - outputs since each rank only has partial outputs. - """ - return torch.ops.vllm.maybe_all_reduce_tensor_model_parallel(final_hidden_states) - - @property - def gate(self) -> torch.nn.Module | None: - return self._gate if self.use_overlapped else None - - @property - def is_internal_router(self) -> bool: - gate = self.gate - return gate is not None and hasattr(gate, "weight_fp32") - - @property - def use_dp_chunking(self) -> bool: - """This func routes to the chunked forward path using the FlashInfer Cutlass kernel - only when data parallelism (DP) is enabled. Thus just returning False in vllm-ascend - """ - return False - - def forward( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - self.ensure_moe_quant_config_init() - return self.runner.forward( - hidden_states, - router_logits, - ) - - def forward_impl( # type: ignore[override] - self, hidden_states: torch.Tensor, router_logits: torch.Tensor, return_with_event: bool = False - ) -> torch.Tensor | FusedMoEResult: - assert self.quant_method is not None - - forward_context = get_forward_context() - # When static kernels are enabled, the forward pass runs twice (compilation + capture), - # causing moe_layer_index to overflow. Wrap the index to prevent out-of-bounds errors. - if self.enable_npugraph_ex_static_kernel and forward_context.all_moe_layers: - moe_layer_index = forward_context.moe_layer_index % (len(forward_context.all_moe_layers)) - forward_context.moe_layer_index = moe_layer_index - - # Load balancing for token distribution among experts in dummy_run - # TODO: The community only considers load balancing when DP > 1. - # This approach may overlook some extreme scenarios. - enable_force_load_balance = _EXTRA_CTX.in_profile_run - - prepare_output = _EXTRA_CTX.moe_comm_method.prepare( - hidden_states=hidden_states, - router_logits=router_logits, - replace_allreduce=_EXTRA_CTX.flash_comm_v1_enabled, - enable_shared_expert_dp=self.enable_shared_expert_dp, - quant_type=self.quant_type, - ) - hidden_states = prepare_output.hidden_states - router_logits = prepare_output.router_logits - mc2_mask = prepare_output.mc2_mask - padded_hidden_states_shape = prepare_output.padded_hidden_states_shape - pertoken_scale = prepare_output.pertoken_scale - - # Matrix multiply. - fused_experts_results: FusedExpertsResult = self.quant_method.apply( - layer=self, - x=hidden_states, - router_logits=router_logits, - pertoken_scale=pertoken_scale, - top_k=self.top_k, - renormalize=self.renormalize, - use_grouped_topk=self.use_grouped_topk, - num_experts=self.moe_config.num_experts, - expert_map=self._expert_map, - topk_group=self.topk_group, - num_expert_group=self.num_expert_group, - custom_routing_function=self.custom_routing_function, - scoring_func=self.scoring_func, - routed_scaling_factor=self._original_routed_scaling_factor, - e_score_correction_bias=self.e_score_correction_bias, - activation=self.activation, - apply_router_weight_on_input=self.apply_router_weight_on_input, - enable_force_load_balance=enable_force_load_balance, - log2phy=self.log2phy, - global_redundant_expert_num=self.global_redundant_expert_num, - mc2_mask=mc2_mask, - ) - - if self.dynamic_eplb and _EXTRA_CTX.eplb_heat_collection_status: - expert_tokens = fused_experts_results.expert_tokens - group_list_type = fused_experts_results.group_list_type - assert expert_tokens is not None and group_list_type is not None, ( - "expert_tokens and group_list_type should not be None when dynamic_eplb is enabled." - ) - local_load = ( - expert_tokens - if group_list_type == 1 - else torch.cat([expert_tokens[:1], expert_tokens[1:] - expert_tokens[:-1]]) - ) - if self.multi_stage: - cur_iter = torch.remainder(self.load_counter, self.num_iter) - self.moe_load.index_add_( - dim=0, index=cur_iter, source=local_load.to(torch.int32, non_blocking=True).view(1, -1) - ) - self.load_counter.add_(1) - else: - self.moe_load.add_(local_load) - - routed_out = _EXTRA_CTX.moe_comm_method.finalize( - hidden_states=fused_experts_results.routed_out, - reduce_results=isinstance(_EXTRA_CTX.moe_comm_method, AllGatherCommImpl), - padded_hidden_states_shape=padded_hidden_states_shape, - ) - - if return_with_event: - return FusedMoEResult( - routed_out=routed_out, - before_dispatch_evt=fused_experts_results.before_dispatch_evt, - before_gmm2_evt=fused_experts_results.before_gmm2_evt, - before_combine_evt=fused_experts_results.before_combine_evt, - swiglu_limit=fused_experts_results.swiglu_limit, - ) - else: - # The vLLM FusedMoE forward_impl does not return events. - return routed_out - - def _forward_shared_experts(self, hidden_states: torch.Tensor, fused_moe_evts: FusedMoEEvents): - if self._shared_experts is None: - return None - - def maybe_wait_event(evt: torch.npu.Event | None): - if evt is not None: - torch.npu.current_stream().wait_event(evt) - - with npu_stream_switch(shared_experts_calculation_stream(), enabled=self.multistream_overlap_shared_expert): - # Only used for int quantization - has_quantized_shared = hasattr(self._shared_experts.gate_up_proj, "weight_scale") and hasattr( - self._shared_experts.down_proj, "weight_scale" - ) - if has_quantized_shared and self.quant_type in (QuantType.W8A8, QuantType.W4A8): - original_dtype = hidden_states.dtype - # Execute dynamic quant concurrently with MoE gate. - torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) - quantized_x, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) - # Execute the gate projection and activation concurrently with the - # dispatch communication. - maybe_wait_event(fused_moe_evts.after_routed_experts) - hidden_states = torch_npu.npu_quant_matmul( - quantized_x, - self._shared_experts.gate_up_proj.weight, - self._shared_experts.gate_up_proj.weight_scale, - pertoken_scale=None, - bias=None, - output_dtype=torch.int32, - ) - # Execute activation concurrently with gmm2. - - maybe_wait_event(fused_moe_evts.before_gmm2) - quantized_x, swiglu_out_scale = torch.ops._C_ascend.npu_dequant_swiglu_quant( - x=hidden_states, - weight_scale=self._shared_experts.gate_up_proj.weight_scale_fp32, - activation_scale=pertoken_scale, - bias=None, - quant_scale=None, - quant_offset=None, - group_index=None, - activate_left=True, - quant_mode=1, - swiglu_mode=1, - clamp_limit=fused_moe_evts.swiglu_limit, - ) - # Execute the down projection concurrently with the combine - # communication. - maybe_wait_event(fused_moe_evts.before_combine) - shared_out = torch_npu.npu_quant_matmul( - quantized_x, - self._shared_experts.down_proj.weight, - self._shared_experts.down_proj.weight_scale, - pertoken_scale=swiglu_out_scale, - bias=None, - output_dtype=original_dtype, - ) - elif has_quantized_shared and self.quant_type == QuantType.W4A8MXFP: - original_dtype = hidden_states.dtype - # Execute dynamic quant concurrently with MoE gate. - torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) - quantized_x, pertoken_scale = torch_npu.npu_dynamic_mx_quant( - hidden_states, dst_type=torch.float8_e4m3fn - ) - # Execute the gate projection and activation concurrently with the - # dispatch communication. - maybe_wait_event(fused_moe_evts.before_dispatch) - hidden_states = self._shared_experts.gate_up_proj((quantized_x, pertoken_scale))[0] - # Execute activation concurrently with gmm2. - maybe_wait_event(fused_moe_evts.before_gmm2) - quantized_x, swiglu_out_scale, _ = torch.ops._C_ascend.npu_swiglu_group_quant( - hidden_states, - topk_weight=None, - group_index=None, - dst_type=torch.float8_e4m3fn, - quant_mode=2, - clamp_value=fused_moe_evts.swiglu_limit, - ) - # Execute the down projection concurrently with the combine - # communication. - maybe_wait_event(fused_moe_evts.before_combine) - shared_out = self._shared_experts.down_proj((quantized_x, swiglu_out_scale))[0] - else: - # Ensure the shared experts wait for hidden_states to be ready. - torch.npu.current_stream().wait_event(fused_moe_evts.before_routed_experts) - # Execute the gate projection and activation concurrently with the - # dispatch communication. - maybe_wait_event(fused_moe_evts.before_dispatch) - part1_out = self._shared_experts_part1(hidden_states) - # Execute the down projection concurrently with the combine - # communication. - maybe_wait_event(fused_moe_evts.before_combine) - shared_out = self._shared_experts_part2(hidden_states, part1_out) - - # Make sure the default stream waits for the shared experts stream to - # finish. - if self.multistream_overlap_shared_expert: - torch.npu.current_stream().wait_stream(shared_experts_calculation_stream()) - - # NOTE: This is exactly the opposite of - # `maybe_all_reduce_tensor_model_parallel` - moe_comm_type = _EXTRA_CTX.moe_comm_type - if ( - moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2, MoECommType.FUSED_MC2} - and not shared_expert_dp_enabled() - ): - shared_out = tensor_model_parallel_all_reduce(shared_out) - return shared_out - - def shared_forward_impl( # type: ignore[override] - self, hidden_states: torch.Tensor, router_logits: torch.Tensor - ): - if self.is_internal_router: - gate = self.gate - assert gate is not None - # NOTE(Angazenn): To make this cast explicitly, the hbm usage might - # increase with extra hidden states. We also assume that all gate - # linear is unquantized so that we the weight is pre-casted in - # process_weights_after_loading of AscendUnquantizedLinearMethod. - hidden_states_fp32 = hidden_states.float() - before_routed_experts = torch.npu.current_stream().record_event() - router_logits = F.linear(hidden_states_fp32, gate.weight_fp32) - after_routed_experts = torch.npu.current_stream().record_event() - else: - before_routed_experts = torch.npu.current_stream().record_event() - after_routed_experts = None - - fused_moe_results = self.forward_impl( - hidden_states=hidden_states, - router_logits=router_logits, - return_with_event=True, - ) - routed_out = fused_moe_results.routed_out - - if self._shared_experts is None: - return routed_out - - shared_out = self._forward_shared_experts( - hidden_states, - FusedMoEEvents( - after_routed_experts=after_routed_experts, - before_routed_experts=before_routed_experts, - before_dispatch=fused_moe_results.before_dispatch_evt, - before_gmm2=fused_moe_results.before_gmm2_evt, - before_combine=fused_moe_results.before_combine_evt, - swiglu_limit=fused_moe_results.swiglu_limit, - ), - ) - return shared_out, routed_out - - -__all__ = ["AscendFusedMoE", "AscendMoERunner"] diff --git a/vllm_ascend/ops/fused_moe/moe_runtime_args.py b/vllm_ascend/ops/fused_moe/moe_runtime_args.py index 5705999bd5c..645abf53cd3 100644 --- a/vllm_ascend/ops/fused_moe/moe_runtime_args.py +++ b/vllm_ascend/ops/fused_moe/moe_runtime_args.py @@ -77,7 +77,6 @@ MoERoutingParams, ) from vllm_ascend.quantization.quant_type import QuantType -from vllm_ascend.utils import vllm_version_is def _build_mxfp_params( @@ -149,7 +148,7 @@ def build_fused_experts_input( swiglu_limit: float | None = 0.0, lora_context=None, ) -> MoEFusedExpertsInput: - if not vllm_version_is("0.23.0") and swiglu_limit is None: + if swiglu_limit is None: swiglu_limit = 0.0 assert swiglu_limit is not None diff --git a/vllm_ascend/ops/gdn.py b/vllm_ascend/ops/gdn.py index 72c1ec563e9..4c219d4afb7 100644 --- a/vllm_ascend/ops/gdn.py +++ b/vllm_ascend/ops/gdn.py @@ -143,7 +143,7 @@ def forward( core_attn_out = self.norm(core_attn_out, z) core_attn_out = core_attn_out.reshape(z_shape_og) core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)") - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): output[:num_tokens], _ = self.out_proj(core_attn_out) else: out, _ = self.out_proj(core_attn_out) diff --git a/vllm_ascend/patch/hunyuan_vl_processor_compat.py b/vllm_ascend/patch/hunyuan_vl_processor_compat.py index 36dd198313c..e25565e52f4 100644 --- a/vllm_ascend/patch/hunyuan_vl_processor_compat.py +++ b/vllm_ascend/patch/hunyuan_vl_processor_compat.py @@ -79,8 +79,8 @@ def __init__( ) -def _import_v023_hunyuan_vision() -> Any: - """Import v0.23's model with the native processors from vLLM PR #47872.""" +def _import_v024_hunyuan_vision() -> Any: + """Import a bundled release model with native processors from vLLM PR #47872.""" from transformers.models.hunyuan_vl.image_processing_hunyuan_vl import ( HunYuanVLImageProcessor, smart_resize, @@ -171,7 +171,7 @@ def get_hf_processor(self: Any, **kwargs: object) -> Any: hunyuan_vision.HunYuanVLProcessingInfo.get_hf_processor = get_hf_processor -def _patch_v023_processor_methods(hunyuan_vision: Any) -> None: +def _patch_v024_processor_methods(hunyuan_vision: Any) -> None: """Backport the Transformers 5.13 call protocol from vLLM PR #47872.""" def call_hf_processor( @@ -201,11 +201,11 @@ def install_hunyuan_vl_processor_compat() -> None: # Keep each target's native, image-token-only prompt replacement. The # cached processor path applies it inside an existing start/image/end # wrapper; using a full-wrapper replacement here would duplicate wrappers. - if vllm_version_is("0.23.0"): - v023_hunyuan_vision = _import_v023_hunyuan_vision() + if vllm_version_is("0.24.0"): + v024_hunyuan_vision = _import_v024_hunyuan_vision() _remove_stale_registry_entries() - _patch_hunyuan_processor_loader(v023_hunyuan_vision) - _patch_v023_processor_methods(v023_hunyuan_vision) + _patch_hunyuan_processor_loader(v024_hunyuan_vision) + _patch_v024_processor_methods(v024_hunyuan_vision) return if not _remove_stale_registry_entries(): diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py index a04b21348a4..e6c76fd2a72 100644 --- a/vllm_ascend/patch/platform/__init__.py +++ b/vllm_ascend/patch/platform/__init__.py @@ -31,11 +31,15 @@ import vllm_ascend.patch.platform.patch_minimax_m2_config # noqa import vllm_ascend.patch.platform.patch_glm_tool_call_streaming # noqa +# TODO: Remove these retained v0.23-only compatibility patches after +# their respective owners complete the cleanup. if vllm_version_is("0.23.0"): import vllm_ascend.patch.platform.patch_glm47_tool_call_parser # noqa - import vllm_ascend.patch.platform.patch_deepseek_v4_tool_call_parser # noqa import vllm_ascend.patch.platform.patch_minimax_m2_tool_call_parser # noqa import vllm_ascend.patch.platform.patch_minimax_usage_accounting # noqa + +if vllm_version_is("0.24.0"): + import vllm_ascend.patch.platform.patch_deepseek_v4_tool_call_parser # noqa import vllm_ascend.patch.platform.patch_structured_output # noqa import vllm_ascend.patch.platform.patch_weight_transfer_engine # noqa import vllm_ascend.patch.platform.patch_torch_accelerator # noqa @@ -50,6 +54,5 @@ import vllm_ascend.patch.platform.patch_kv_cache_coordinator # noqa import vllm_ascend.patch.platform.patch_speculative_config # noqa -if not vllm_version_is("0.23.0"): - import vllm_ascend.patch.platform.patch_fused_moe # noqa - import vllm_ascend.patch.platform.patch_dp_device_ids # noqa +import vllm_ascend.patch.platform.patch_fused_moe # noqa +import vllm_ascend.patch.platform.patch_dp_device_ids # noqa diff --git a/vllm_ascend/patch/platform/patch_balance_schedule.py b/vllm_ascend/patch/platform/patch_balance_schedule.py index 3b65eadec99..ab53eb30290 100644 --- a/vllm_ascend/patch/platform/patch_balance_schedule.py +++ b/vllm_ascend/patch/platform/patch_balance_schedule.py @@ -10,28 +10,15 @@ there is no leader. See ``docs/.../balance_schedule_refactor.md`` for the design. -The ``schedule()`` body is a verbatim copy of the **v0.23.0** release tag's +The ``schedule()`` body is a verbatim copy of the **v0.24.0** release tag's ``Scheduler.schedule()`` (the production pin), plus exactly three balance deltas: (1) the disabled-path early return that delegates to ``super()``, (2) the ``balance_flag`` break inside the WAITING loop (``any-rank-at-cap => global freeze``), and (3) ``if request_queue is None: break`` in place of upstream's ``assert request_queue is not None`` (so a -drained-rank schedule does not assert when balance defers admission). - -The **signature**, in contrast, must work across BOTH vllm versions that -vllm-ascend CI runs simultaneously: the release tag v0.23.0 (whose engine -calls ``schedule()`` with no args) and the main-verified commit 1f486d96 -(whose engine calls ``schedule(throttle_prefills)``). So the override carries -``throttle_prefills`` with a default -- a deliberate superset of v0.23.0's -``schedule(self)`` -- making it callable by both engines. On the disabled -fast-path it then forwards ``throttle_prefills`` only when the installed -``super().schedule`` actually accepts it, decided by introspecting the -signature once at import (``_SUPER_SCHEDULE_HAS_THROTTLE``) rather than -parsing a version string (a dev checkout's ``__version__`` is not a clean -PEP 440 release and would make a ``vllm_version_is`` check raise). The body -and the signature therefore deliberately target different things: the body -tracks the stable release tag, the signature tracks the union of both -engines' call shapes. See the design doc for the full rationale. +drained-rank schedule does not assert when balance defers admission). Both +supported vLLM refs expose ``schedule(throttle_prefills=False)``, so the +disabled fast path forwards that argument directly. The engine-core side is NOT copied: ``BalanceDPEngineCoreProc`` hooks ``_has_global_unfinished_reqs`` (called every iteration by upstream's @@ -69,7 +56,6 @@ touch configs that don't use it, e.g. PD-disaggregated recompute). """ -import inspect import time import torch @@ -79,6 +65,7 @@ from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata from vllm.logger import logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput @@ -91,19 +78,6 @@ from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext -# Whether the *installed* upstream ``Scheduler.schedule`` accepts the -# ``throttle_prefills`` argument. vllm-ascend CI runs against TWO vllm -# versions at once: the release tag v0.23.0 (``schedule(self)``, engine calls -# ``schedule()``) and the main-verified commit 1f486d96 (``schedule(self, -# throttle_prefills=False)``, engine calls ``schedule(throttle_prefills)``). -# The override signature carries ``throttle_prefills`` (with a default) so it -# is callable by BOTH engines; on the disabled path it must then forward the -# arg only when the installed super() actually accepts it. Introspecting the -# signature once at import (rather than parsing a version string) is robust to -# both lanes -- including dev checkouts whose ``__version__`` is not a clean -# PEP 440 release (which would make a ``vllm_version_is`` check raise). -_SUPER_SCHEDULE_HAS_THROTTLE = "throttle_prefills" in inspect.signature(Scheduler.schedule).parameters - def _balance_scheduling_enabled(vllm_config) -> bool: # Primary source of truth is AscendConfig. The additional_config fallback @@ -175,12 +149,7 @@ def balance_gather(self): def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if not self._balance_enabled: - # Forward throttle_prefills only when the installed super() accepts - # it (main-verified); v0.23.0's super() does not. See - # _SUPER_SCHEDULE_HAS_THROTTLE for why this is signature-based. - if _SUPER_SCHEDULE_HAS_THROTTLE: - return super().schedule(throttle_prefills) - return super().schedule() + return super().schedule(throttle_prefills) self.current_step += 1 # NOTE(woosuk) on the scheduling algorithm: # There's no "decoding phase" nor "prefill phase" in the scheduler. @@ -216,6 +185,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.kv_cache_manager.new_step_starts() + # DP prefill balancing: on a throttled (non-cadence-aligned) step, defer + # all prefill compute unless saturated. + defer_prefills = (throttle_prefills and not self.prefill_capacity_bound) and any( + not r.is_prefill_chunk for r in self.running + ) + # First, schedule the RUNNING requests. req_index = 0 while req_index < len(self.running) and token_budget > 0: @@ -243,6 +218,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req_index += 1 continue + if defer_prefills and request.is_prefill_chunk: + # DP prefill balancing: defer this in-progress prefill chunk to a + # cadence-aligned step; decodes still run to fill this step. + req_index += 1 + continue + num_new_tokens = ( request.num_tokens_with_spec + request.num_output_placeholders - request.num_computed_tokens ) @@ -252,7 +233,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Make sure the input position does not exceed the max model len. # This is necessary when using spec decoding. - num_new_tokens = min(num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens) + num_new_tokens = min( + num_new_tokens, + self.max_model_len - request.num_computed_tokens - self.num_sampled_tokens_per_step, + ) # Schedule encoder inputs. encoder_inputs_to_schedule = None @@ -286,7 +270,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # 2. The encoder budget is exhausted. # 3. The encoder cache is exhausted. # 4. Insufficient budget for a block-aligned chunk in hybrid - # models with mamba cache mode "align". + # models with mamba cache mode \"align\". # NOTE(woosuk): Here, by doing `continue` instead of `break`, # we do not strictly follow the FCFS scheduling policy and # allow the lower-priority requests to be scheduled. @@ -398,11 +382,9 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if len(self.running) == self.max_num_running_reqs: break - # >>> balance-scheduling delta (the whole point of this patch) <<< - # If any DP rank was at the running cap at the end of the - # previous step, stop admitting new WAITING requests on this - # rank too, so load stays even across ranks - # (leader-at-cap => global freeze). + # Keep admission balanced across DP ranks: if any rank was at + # the running cap after the previous step, stop admitting new + # waiting requests on every rank. if max(t.item() for t in self.balance_queue) == self.max_num_running_reqs: break @@ -444,13 +426,53 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_external_computed_tokens = 0 load_kv_async = False connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 + num_uncached_common_prefix_tokens = 0 # Get already-cached tokens. if request.num_computed_tokens == 0: # Get locally-cached tokens. - new_computed_blocks, num_new_local_computed_tokens = self.kv_cache_manager.get_computed_blocks( - request - ) + if ( + self.connector is not None + and self.has_mamba_layers + and isinstance( + self.kv_cache_manager.coordinator, + HybridKVCacheCoordinator, + ) + ): + computed, per_group_hits = self.kv_cache_manager.coordinator.find_longest_cache_hit_per_group( + request.block_hashes, + request.num_tokens - 1, + ) + new_computed_blocks = self.kv_cache_manager.create_kv_cache_blocks(computed) + # NOTE(ZhanqiuHu): For Mamba hybrid models, + # num_new_local_computed_tokens should be the FA hit + # length. This value is passed to the connector's + # get_num_new_matched_tokens which computes: + # external = total - local_computed. + # Using the FA hit skips re-transferring FA blocks + # already cached on D-side. The Mamba state (always + # the last block) is transferred unconditionally by + # _apply_prefix_caching in nixl/worker.py. + num_new_local_computed_tokens = max(per_group_hits) + if self.kv_cache_manager.log_stats: + assert self.kv_cache_manager.prefix_cache_stats is not None + self.kv_cache_manager.prefix_cache_stats.record( + num_tokens=request.num_tokens, + num_hits=num_new_local_computed_tokens, + preempted=request.num_preemptions > 0, + ) + else: + new_computed_blocks, num_new_local_computed_tokens = self.kv_cache_manager.get_computed_blocks( + request + ) + + # In case of hybrid models, obtain hint for Marconi-style APC logic + if self.has_mamba_layers: + num_uncached_common_prefix_tokens = getattr( + self.kv_cache_manager.coordinator, + "num_uncached_common_prefix_tokens", + 0, + ) # Get externally-cached tokens if using a KVConnector. if self.connector is not None: @@ -508,6 +530,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # KVTransfer: loading remote KV, do not allocate for new work. assert num_external_computed_tokens > 0 num_new_tokens = 0 + elif defer_prefills and request.num_computed_tokens == 0: + # DP prefill balancing: async KV loads (the branch above) are + # allowed to start even on throttled steps, but committing new + # prefill compute is deferred to a cadence-aligned step. + break else: # Number of tokens to be scheduled. # We use `request.num_tokens` instead of @@ -553,6 +580,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_new_tokens, num_new_local_computed_tokens, num_external_computed_tokens, + num_uncached_common_prefix_tokens, ) if num_new_tokens == 0: break @@ -589,6 +617,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=self.scheduler_reserve_full_isl, reserved_blocks=reserved_blocks, + has_scheduled_reqs=bool(self.running), ) if new_blocks is None: @@ -680,6 +709,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if step_skipped_waiting: self.skipped_waiting.prepend_requests(step_skipped_waiting) + # DP prefill balancing: on a step that admitted prefills (release), + # record whether it was capacity-bound. + if not defer_prefills: + self.prefill_capacity_bound = bool(self.waiting) + # Check if the scheduling constraints are satisfied. total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens @@ -701,8 +735,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Construct the scheduler output. if self.use_v2_model_runner: - scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs - scheduled_resumed_reqs = [] + scheduled_new_reqs.extend(scheduled_resumed_reqs) + scheduled_resumed_reqs.clear() new_reqs_data = [ NewRequestData.from_request( req, @@ -726,14 +760,20 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req_to_new_blocks, ) - # Record the request ids that were scheduled in this step. - self.prev_step_scheduled_req_ids.clear() - self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + # Record the request ids that were scheduled in this step (MRV1-only). + if not self.use_v2_model_runner: + self.prev_step_scheduled_req_ids.clear() + self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) new_block_ids_to_zero = ( (self.kv_cache_manager.take_new_block_ids() or None) if self.needs_kv_cache_zeroing else None ) + # Dynamic speculative decoding: compute optimal K + num_spec_tokens_to_schedule = self.num_spec_tokens + if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0: + num_spec_tokens_to_schedule = self.dynamic_sd_lookup[len(num_scheduled_tokens)] + scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=cached_reqs_data, @@ -750,6 +790,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=new_block_ids_to_zero, + num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, ) # NOTE(Kuntai): this function is designed for multiple purposes: @@ -765,6 +806,11 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: ec_meta: ECConnectorMetadata = self.ec_connector.build_connector_meta(scheduler_output) scheduler_output.ec_connector_metadata = ec_meta + # Advance the fence only for non-empty steps (those that actually + # write KV and have their output processed later in update_from_output). + if self.defer_block_free and total_num_scheduled_tokens > 0: + self.sched_step_seq += 1 + with record_function_or_nullcontext("schedule: update_after_schedule"): self._update_after_schedule(scheduler_output) return scheduler_output diff --git a/vllm_ascend/patch/platform/patch_deepseek_v4_tool_call_parser.py b/vllm_ascend/patch/platform/patch_deepseek_v4_tool_call_parser.py index 84401b1413c..30f2d29d770 100644 --- a/vllm_ascend/patch/platform/patch_deepseek_v4_tool_call_parser.py +++ b/vllm_ascend/patch/platform/patch_deepseek_v4_tool_call_parser.py @@ -38,9 +38,9 @@ from vllm_ascend.utils import vllm_version_is -if not vllm_version_is("0.23.0"): +if not vllm_version_is("0.24.0"): raise RuntimeError( - "patch_deepseek_v4_tool_call_parser is only for vLLM 0.23.0; " + "patch_deepseek_v4_tool_call_parser is only for vLLM 0.24.0; " "newer vLLM versions use the upstream DeepSeekV4 engine parser." ) diff --git a/vllm_ascend/patch/platform/patch_dp_device_ids.py b/vllm_ascend/patch/platform/patch_dp_device_ids.py index 0820451e65f..8c20474b6d1 100644 --- a/vllm_ascend/patch/platform/patch_dp_device_ids.py +++ b/vllm_ascend/patch/platform/patch_dp_device_ids.py @@ -29,44 +29,43 @@ # offset is out of range and the helper raises ``IndexError`` (wrapped in # the user-facing "Error computing device indices for ..." message). -from vllm_ascend.utils import vllm_version_is +import os -if not vllm_version_is("0.23.0"): - import os +from vllm.platforms import current_platform +from vllm.v1.engine import utils as _engine_utils - from vllm.platforms import current_platform - from vllm.v1.engine import utils as _engine_utils +_original_get_physical_gpu_ids = _engine_utils.get_physical_gpu_ids_for_local_dp_rank - _original_get_physical_gpu_ids = _engine_utils.get_physical_gpu_ids_for_local_dp_rank - def _patched_get_physical_gpu_ids_for_local_dp_rank( +def _patched_get_physical_gpu_ids_for_local_dp_rank( + device_control_env_var, + local_dp_rank, + world_size, + local_world_size=None, + user_assigned_gpu_ids=None, +): + if local_world_size is None: + local_world_size = world_size + + # If the caller did not pass --device-ids and the env var has + # fewer devices than the full DP range expects, the env var has + # already been pre-sharded per rank by the caller. Use it + # directly from index 0 instead of applying the DP offset again. + if user_assigned_gpu_ids is None and device_control_env_var in os.environ: + visible = [d for d in os.environ[device_control_env_var].split(",") if d] + if local_dp_rank * world_size + local_world_size > len(visible): + return [ + current_platform.device_control_id_to_physical_device_id(visible[device_id]) + for device_id in range(local_world_size) + ] + + return _original_get_physical_gpu_ids( device_control_env_var, local_dp_rank, world_size, - local_world_size=None, - user_assigned_gpu_ids=None, - ): - if local_world_size is None: - local_world_size = world_size - - # If the caller did not pass --device-ids and the env var has - # fewer devices than the full DP range expects, the env var has - # already been pre-sharded per rank by the caller. Use it - # directly from index 0 instead of applying the DP offset again. - if user_assigned_gpu_ids is None and device_control_env_var in os.environ: - visible = [d for d in os.environ[device_control_env_var].split(",") if d] - if local_dp_rank * world_size + local_world_size > len(visible): - return [ - current_platform.device_control_id_to_physical_device_id(visible[device_id]) - for device_id in range(local_world_size) - ] + local_world_size, + user_assigned_gpu_ids, + ) - return _original_get_physical_gpu_ids( - device_control_env_var, - local_dp_rank, - world_size, - local_world_size, - user_assigned_gpu_ids, - ) - _engine_utils.get_physical_gpu_ids_for_local_dp_rank = _patched_get_physical_gpu_ids_for_local_dp_rank +_engine_utils.get_physical_gpu_ids_for_local_dp_rank = _patched_get_physical_gpu_ids_for_local_dp_rank diff --git a/vllm_ascend/patch/platform/patch_fused_moe.py b/vllm_ascend/patch/platform/patch_fused_moe.py index 2fa955f2807..65aaae229bd 100644 --- a/vllm_ascend/patch/platform/patch_fused_moe.py +++ b/vllm_ascend/patch/platform/patch_fused_moe.py @@ -27,31 +27,32 @@ # 2. from vllm_ascend import ops # 3. model loading -> deepseek_v2 imported -> gets patched FusedMoE ✓ -from vllm_ascend.utils import is_310p, vllm_version_is - -if not vllm_version_is("0.23.0"): - import vllm.model_executor.layers.fused_moe as _fused_moe_pkg - import vllm.model_executor.layers.fused_moe.layer as _fused_moe_layer - - # Capture the real original before fused_moe.py's module-level code runs. - _original_FusedMoE = _fused_moe_layer.FusedMoE - - if is_310p(): - from vllm_ascend._310p.fused_moe.fused_moe import AscendMoERunner310 as _DefaultAscendMoERunner - else: - from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner as _DefaultAscendMoERunner - - def _ascend_FusedMoE(*args, runner_cls=None, runner_args=None, **kwargs): - if runner_cls is None: - runner_cls = _DefaultAscendMoERunner - # 'hash' is a DeepSeek V4 flag already consumed before FusedMoE is called; - # 'tid2eid' is Ascend-specific and must reach AscendMoERunner via runner_args. - kwargs.pop("hash", None) - tid2eid = kwargs.pop("tid2eid", None) - if tid2eid is not None: - runner_args = dict(runner_args) if runner_args is not None else {} - runner_args["tid2eid"] = tid2eid - return _original_FusedMoE(*args, runner_cls=runner_cls, runner_args=runner_args, **kwargs) - - _fused_moe_layer.FusedMoE = _ascend_FusedMoE - _fused_moe_pkg.FusedMoE = _ascend_FusedMoE +import vllm.model_executor.layers.fused_moe as _fused_moe_pkg +import vllm.model_executor.layers.fused_moe.layer as _fused_moe_layer + +from vllm_ascend.utils import is_310p + +# Capture the real original before fused_moe.py's module-level code runs. +_original_FusedMoE = _fused_moe_layer.FusedMoE + +if is_310p(): + from vllm_ascend._310p.fused_moe.fused_moe import AscendMoERunner310 as _DefaultAscendMoERunner +else: + from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner as _DefaultAscendMoERunner + + +def _ascend_FusedMoE(*args, runner_cls=None, runner_args=None, **kwargs): + if runner_cls is None: + runner_cls = _DefaultAscendMoERunner + # 'hash' is a DeepSeek V4 flag already consumed before FusedMoE is called; + # 'tid2eid' is Ascend-specific and must reach AscendMoERunner via runner_args. + kwargs.pop("hash", None) + tid2eid = kwargs.pop("tid2eid", None) + if tid2eid is not None: + runner_args = dict(runner_args) if runner_args is not None else {} + runner_args["tid2eid"] = tid2eid + return _original_FusedMoE(*args, runner_cls=runner_cls, runner_args=runner_args, **kwargs) + + +_fused_moe_layer.FusedMoE = _ascend_FusedMoE +_fused_moe_pkg.FusedMoE = _ascend_FusedMoE diff --git a/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py b/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py index ef15c46733b..c66096f553f 100644 --- a/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py +++ b/vllm_ascend/patch/platform/patch_kv_cache_coordinator.py @@ -43,7 +43,7 @@ def _select_kv_token_budget( max_in_flight_tokens: int | None, max_num_batched_tokens: int | None, ) -> int: - token_budget = max_num_batched_tokens if vllm_version_is("0.23.0") else max_in_flight_tokens + token_budget = max_num_batched_tokens if vllm_version_is("0.24.0") else max_in_flight_tokens return token_budget if token_budget is not None else max_model_len @@ -131,7 +131,7 @@ def __init__( self.eagle_group_ids = set(range(len(kv_cache_config.kv_cache_groups))) extra_mgr_kwargs: dict = {"scheduler_block_size": scheduler_block_size} - if not vllm_version_is("0.23.0"): + if not vllm_version_is("0.24.0"): extra_mgr_kwargs["needs_kv_cache_zeroing"] = kv_cache_config.needs_kv_cache_zeroing self.single_type_managers = tuple( get_manager_for_kv_cache_spec( @@ -510,7 +510,7 @@ def get_kv_cache_coordinator( hash_block_size=hash_block_size, metrics_collector=metrics_collector, ) - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): orig_kwargs["max_num_batched_tokens"] = token_budget else: orig_kwargs["max_in_flight_tokens"] = token_budget diff --git a/vllm_ascend/patch/platform/patch_kv_cache_utils.py b/vllm_ascend/patch/platform/patch_kv_cache_utils.py index c08daf54238..44efc9126ce 100644 --- a/vllm_ascend/patch/platform/patch_kv_cache_utils.py +++ b/vllm_ascend/patch/platform/patch_kv_cache_utils.py @@ -17,8 +17,6 @@ UniformTypeKVCacheSpecs, ) -from vllm_ascend.utils import vllm_version_is - _orig_resolve_kv_cache_block_sizes = vllm.v1.core.kv_cache_utils.resolve_kv_cache_block_sizes @@ -252,13 +250,10 @@ def _get_kv_cache_config_deepseek_v4( vllm.v1.core.kv_cache_utils.resolve_kv_cache_block_sizes = _ascend_resolve_kv_cache_block_sizes vllm.v1.core.kv_cache_utils.group_and_unify_kv_cache_specs = group_and_unify_kv_cache_specs vllm.v1.core.kv_cache_utils._get_kv_cache_groups_uniform_groups = _get_kv_cache_groups_uniform_groups -# vllm v0.24.0 renamed _get_kv_cache_config_deepseek_v4 to _get_kv_cache_config_packed and +# vLLM v0.24.0 renamed _get_kv_cache_config_deepseek_v4 to _get_kv_cache_config_packed and # get_kv_cache_config_from_groups now calls _get_kv_cache_config_packed directly, bypassing # the alias patch above. Patch the canonical name so Ascend's non-packed layout is used. -if vllm_version_is("0.23.0"): - vllm.v1.core.kv_cache_utils._get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_deepseek_v4 -else: - vllm.v1.core.kv_cache_utils._get_kv_cache_config_packed = _get_kv_cache_config_deepseek_v4 +vllm.v1.core.kv_cache_utils._get_kv_cache_config_packed = _get_kv_cache_config_deepseek_v4 # Also patch the reference used by engine/core.py which imports the function directly. import vllm.v1.engine.core # noqa: E402 diff --git a/vllm_ascend/patch/platform/patch_profiling_chunk.py b/vllm_ascend/patch/platform/patch_profiling_chunk.py index fa91f6a80ac..55e661243ba 100644 --- a/vllm_ascend/patch/platform/patch_profiling_chunk.py +++ b/vllm_ascend/patch/platform/patch_profiling_chunk.py @@ -32,8 +32,6 @@ from vllm.logger import logger from vllm.v1.engine.core import EngineCore, EngineCoreProc -from vllm_ascend.utils import vllm_version_is - _profiling_patches_applied = False _original_update_from_output = None _original_schedule = None @@ -170,10 +168,7 @@ def _ensure_schedule_wrapped(scheduler): _original_schedule = cls.schedule def _wrapped_schedule(self, throttle_prefills: bool = False): - if vllm_version_is("0.23.0"): - output = _original_schedule(self) - else: - output = _original_schedule(self, throttle_prefills) + output = _original_schedule(self, throttle_prefills) if getattr(self, "_profiling_timing_done", False) and output is not None: output.disable_profiling_timing = True return output diff --git a/vllm_ascend/patch/platform/patch_torch_accelerator.py b/vllm_ascend/patch/platform/patch_torch_accelerator.py index f1c36fd6d5d..e67d1cc79ad 100644 --- a/vllm_ascend/patch/platform/patch_torch_accelerator.py +++ b/vllm_ascend/patch/platform/patch_torch_accelerator.py @@ -1,7 +1,5 @@ import torch -from vllm_ascend.utils import vllm_version_is - def patch_empty_cache() -> None: torch.npu.empty_cache() @@ -16,11 +14,10 @@ def patch_empty_cache() -> None: torch.accelerator.memory_stats = torch.npu.memory_stats # type: ignore[attr-defined] torch.accelerator.memory_reserved = torch.npu.memory_reserved # type: ignore[attr-defined] torch.accelerator.reset_peak_memory_stats = torch.npu.reset_peak_memory_stats # type: ignore[attr-defined] -if not vllm_version_is("0.23.0"): - # torch.accelerator.get_memory_info() routes through c10's - # CachingDeviceAllocator and asserts the backend allocator is a - # DeviceAllocator; NPU's caching allocator is not, so it crashes with - # "Allocator for npu is not a DeviceAllocator". Redirect to the - # NPU-native API. Only needed on v0.24.0+ where MemorySnapshot - # is constructed with an explicit device arg that triggers this path. - torch.accelerator.get_memory_info = torch.npu.mem_get_info # type: ignore[attr-defined] +# torch.accelerator.get_memory_info() routes through c10's +# CachingDeviceAllocator and asserts the backend allocator is a +# DeviceAllocator; NPU's caching allocator is not, so it crashes with +# "Allocator for npu is not a DeviceAllocator". Redirect to the +# NPU-native API. This is needed on v0.24.0+ where MemorySnapshot +# is constructed with an explicit device arg that triggers this path. +torch.accelerator.get_memory_info = torch.npu.mem_get_info # type: ignore[attr-defined] diff --git a/vllm_ascend/patch/worker/__init__.py b/vllm_ascend/patch/worker/__init__.py index 3e220740fed..128d6ae7d95 100644 --- a/vllm_ascend/patch/worker/__init__.py +++ b/vllm_ascend/patch/worker/__init__.py @@ -19,16 +19,10 @@ from vllm_ascend.utils import is_310p, vllm_version_is -# The v2 model runner is intentionally NOT made compatible with the v0.23.0 -# release. vLLM v0.23.0 and the verified main commit are diverged, and the v2 -# worker patches target main-only APIs; rather than maintain a separate v0.23.0 -# compatibility path we keep v2 main-only. With v0.23.0 installed this flag is -# False, so none of the patch_v2.* / routed-experts-capture patches below are -# imported and the v2 worker stays dormant (the release uses the v1 runner). -if vllm_version_is("0.23.0"): - _V2_MODEL_RUNNER_SUPPORTED = False -else: - _V2_MODEL_RUNNER_SUPPORTED = True +# The v2 model runner tracks only the verified vLLM main commit. The v0.24.0 +# release has diverged APIs and is intentionally kept on the v1 runner instead +# of maintaining a separate v2 compatibility path. +_V2_MODEL_RUNNER_SUPPORTED = not vllm_version_is("0.24.0") if HAS_TRITON: import vllm_ascend.patch.worker.patch_triton @@ -76,8 +70,7 @@ # when the env var is explicitly set. import vllm_ascend.patch.worker.patch_v2.patch_use_v2_model_runner # noqa -if not vllm_version_is("0.23.0"): - import vllm_ascend.patch.worker.patch_fused_moe # noqa +import vllm_ascend.patch.worker.patch_fused_moe # noqa if _V2_MODEL_RUNNER_SUPPORTED: import vllm_ascend.patch.worker.patch_v2.patch_uva # noqa diff --git a/vllm_ascend/patch/worker/patch_deepseek_v2.py b/vllm_ascend/patch/worker/patch_deepseek_v2.py index 6a012383cab..a27165b6e20 100644 --- a/vllm_ascend/patch/worker/patch_deepseek_v2.py +++ b/vllm_ascend/patch/worker/patch_deepseek_v2.py @@ -1,8 +1,14 @@ +from itertools import islice + import torch from torch import nn from transformers import DeepseekV2Config, DeepseekV3Config from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -18,10 +24,13 @@ from vllm.model_executor.models.deepseek_v2 import ( DeepSeekV2FusedQkvAProjLinear, DeepseekV2MLAAttention, + DeepseekV2Model, Indexer, + _get_llama_4_scaling, yarn_get_mscale, ) from vllm.model_executor.models.utils import extract_layer_index +from vllm.sequence import IntermediateTensors from vllm_ascend.utils import vllm_version_is @@ -283,76 +292,68 @@ def _deepseek_v2_mla_attention_init( DeepseekV2MLAAttention.__init__ = _deepseek_v2_mla_attention_init -if not vllm_version_is("0.23.0"): - from itertools import islice - - from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather - from vllm.model_executor.models.deepseek_v2 import ( - DeepseekV2Model, - _get_llama_4_scaling, - ) - from vllm.sequence import IntermediateTensors - - def _patched_forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - if input_ids is None: - raise ValueError("Either input_ids or inputs_embeds must be provided to DeepseekV2Model.forward") - hidden_states = self.embed_input_ids(input_ids) - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - llama_4_scaling_config = getattr(self.config, "llama_4_scaling", None) - llama_4_scaling: torch.Tensor | None - if llama_4_scaling_config is not None: - llama_4_scaling = _get_llama_4_scaling( - original_max_position_embeddings=llama_4_scaling_config["original_max_position_embeddings"], - scaling_beta=llama_4_scaling_config["beta"], - positions=positions, - ) +def _patched_forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, +) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds else: - llama_4_scaling = None - - aux_hidden_states = [] - for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer), - start=self.start_layer, - ): - if idx in self.aux_hidden_state_layers: - aux_hidden_state = hidden_states + residual - if aux_hidden_state.shape[0] != positions.shape[0]: - aux_hidden_state = tensor_model_parallel_all_gather(aux_hidden_state, 0) - aux_hidden_state = aux_hidden_state[: positions.shape[0]] - aux_hidden_states.append(aux_hidden_state) - hidden_states, residual = layer(positions, hidden_states, residual, llama_4_scaling) - - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) - - if hidden_states.shape[0] != positions.shape[0]: - combined_states = torch.cat([hidden_states, residual], dim=-1) - combined_states = tensor_model_parallel_all_gather(combined_states, 0) - combined_states = combined_states[: positions.shape[0]] - hidden_states, residual = combined_states.split([self.hidden_size, self.hidden_size], dim=-1) - residual = residual.contiguous() - - if self.end_layer in self.aux_hidden_state_layers: - aux_hidden_states.append(hidden_states + residual) - - hidden_states, _ = self.norm(hidden_states, residual) - if len(aux_hidden_states) > 0: - return hidden_states, aux_hidden_states - return hidden_states - - DeepseekV2Model.forward = _patched_forward + if input_ids is None: + raise ValueError("Either input_ids or inputs_embeds must be provided to DeepseekV2Model.forward") + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + llama_4_scaling_config = getattr(self.config, "llama_4_scaling", None) + llama_4_scaling: torch.Tensor | None + if llama_4_scaling_config is not None: + llama_4_scaling = _get_llama_4_scaling( + original_max_position_embeddings=llama_4_scaling_config["original_max_position_embeddings"], + scaling_beta=llama_4_scaling_config["beta"], + positions=positions, + ) + else: + llama_4_scaling = None + + aux_hidden_states = [] + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): + if idx in self.aux_hidden_state_layers: + aux_hidden_state = hidden_states + residual + if aux_hidden_state.shape[0] != positions.shape[0]: + aux_hidden_state = tensor_model_parallel_all_gather(aux_hidden_state, 0) + aux_hidden_state = aux_hidden_state[: positions.shape[0]] + aux_hidden_states.append(aux_hidden_state) + hidden_states, residual = layer(positions, hidden_states, residual, llama_4_scaling) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + + if hidden_states.shape[0] != positions.shape[0]: + combined_states = torch.cat([hidden_states, residual], dim=-1) + combined_states = tensor_model_parallel_all_gather(combined_states, 0) + combined_states = combined_states[: positions.shape[0]] + hidden_size = self.config.hidden_size if vllm_version_is("0.24.0") else self.hidden_size + hidden_states, residual = combined_states.split([hidden_size, hidden_size], dim=-1) + residual = residual.contiguous() + + if self.end_layer in self.aux_hidden_state_layers: + aux_hidden_states.append(hidden_states + residual) + + hidden_states, _ = self.norm(hidden_states, residual) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + +DeepseekV2Model.forward = _patched_forward diff --git a/vllm_ascend/patch/worker/patch_qwen3_5.py b/vllm_ascend/patch/worker/patch_qwen3_5.py index 246e84e8009..dcb85c68c65 100644 --- a/vllm_ascend/patch/worker/patch_qwen3_5.py +++ b/vllm_ascend/patch/worker/patch_qwen3_5.py @@ -84,7 +84,7 @@ def forward(self, positions: torch.Tensor, hidden_states: torch.Tensor, output: gate = torch.sigmoid(gate) attn_output = attn_output * gate - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): output[:], _ = self.o_proj(attn_output) else: out, _ = self.o_proj(attn_output) @@ -105,7 +105,7 @@ def forward( else: hidden_states, residual = self.input_layernorm(hidden_states, residual) - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): if self.layer_idx == 0 and _EXTRA_CTX.flash_comm_v1_enabled: tp_size = get_tensor_model_parallel_world_size() n_out = (hidden_states.shape[0] + tp_size - 1) // tp_size diff --git a/vllm_ascend/patch/worker/patch_qwen3_dflash.py b/vllm_ascend/patch/worker/patch_qwen3_dflash.py index 64a732ebd7a..ba78d0621e4 100644 --- a/vllm_ascend/patch/worker/patch_qwen3_dflash.py +++ b/vllm_ascend/patch/worker/patch_qwen3_dflash.py @@ -70,7 +70,7 @@ def precompute_and_store_context_kv( DFlashQwen3Model.precompute_and_store_context_kv = precompute_and_store_context_kv -if not vllm_version_is("0.23.0"): +if not vllm_version_is("0.24.0"): _orig_read_mask_embedding = DFlashQwen3ForCausalLM._read_mask_embedding def _patched_read_mask_embedding(self): diff --git a/vllm_ascend/quantization/compressed_tensors_config.py b/vllm_ascend/quantization/compressed_tensors_config.py index 904ba7d56df..d5da285187a 100644 --- a/vllm_ascend/quantization/compressed_tensors_config.py +++ b/vllm_ascend/quantization/compressed_tensors_config.py @@ -22,6 +22,7 @@ import torch from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy, QuantizationType from vllm.logger import logger +from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS, register_quantization_config from vllm.model_executor.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase @@ -32,21 +33,13 @@ ) from vllm.model_executor.models.utils import WeightsMapper -from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD, vllm_version_is +from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD from .methods import AscendLinearScheme, AscendMoEScheme -if vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe import FusedMoE -else: - from vllm.model_executor.layers.fused_moe import MoERunner - def _is_fused_moe_layer(layer: torch.nn.Module) -> bool: - if vllm_version_is("0.23.0"): - return isinstance(layer, FusedMoE) - else: - return isinstance(layer, MoERunner) + return isinstance(layer, MoERunner) # Remove the original compressed_tensors method to replace with our implementation diff --git a/vllm_ascend/quantization/fp8_config.py b/vllm_ascend/quantization/fp8_config.py index c5e00186756..0a7743685bc 100644 --- a/vllm_ascend/quantization/fp8_config.py +++ b/vllm_ascend/quantization/fp8_config.py @@ -3,25 +3,18 @@ import torch from compressed_tensors.quantization import QuantizationArgs from vllm.logger import logger +from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS, register_quantization_config from vllm.model_executor.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase -from vllm_ascend.utils import FP8_METHOD, vllm_version_is - -if vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe import FusedMoE -else: - from vllm.model_executor.layers.fused_moe import MoERunner +from vllm_ascend.utils import FP8_METHOD from .methods import get_scheme_class def _is_fused_moe_layer(layer: torch.nn.Module) -> bool: - if vllm_version_is("0.23.0"): - return isinstance(layer, FusedMoE) - else: - return isinstance(layer, MoERunner) + return isinstance(layer, MoERunner) QUANTIZATION_SCHEME_MAP_TYPE = dict[str, dict[str, QuantizationArgs] | None] diff --git a/vllm_ascend/quantization/modelslim_config.py b/vllm_ascend/quantization/modelslim_config.py index 076e16a0859..14dd69d2ad6 100644 --- a/vllm_ascend/quantization/modelslim_config.py +++ b/vllm_ascend/quantization/modelslim_config.py @@ -34,6 +34,7 @@ from vllm.config import get_current_vllm_config from vllm.logger import logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.fused_moe import MoERunner, RoutedExperts from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization import register_quantization_config from vllm.model_executor.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase @@ -45,22 +46,13 @@ AscendDeviceType, calc_split_factor, get_ascend_device_type, - vllm_version_is, ) -if vllm_version_is("0.23.0"): - from vllm.model_executor.layers.fused_moe import FusedMoE -else: - from vllm.model_executor.layers.fused_moe import MoERunner, RoutedExperts - from .methods import get_scheme_class def _is_fused_moe_layer(layer: torch.nn.Module) -> bool: - if vllm_version_is("0.23.0"): - return isinstance(layer, FusedMoE) - else: - return isinstance(layer, (MoERunner, RoutedExperts)) + return isinstance(layer, (MoERunner, RoutedExperts)) # The config filename that ModelSlim generates after quantizing a model. diff --git a/vllm_ascend/spec_decode/llm_base_proposer.py b/vllm_ascend/spec_decode/llm_base_proposer.py index d3d9b1dd307..4640f6b327a 100644 --- a/vllm_ascend/spec_decode/llm_base_proposer.py +++ b/vllm_ascend/spec_decode/llm_base_proposer.py @@ -54,7 +54,7 @@ from vllm_ascend.spec_decode.utils import SlidingWindowAdapter from vllm_ascend.utils import check_gdn_layer, enable_sp, lmhead_tp_enable, shared_expert_dp_enabled, vllm_version_is -if not vllm_version_is("0.23.0"): +if not vllm_version_is("0.24.0"): from vllm.model_executor.models.qwen3_dspark import Qwen3DSparkForCausalLM else: Qwen3DSparkForCausalLM = None diff --git a/vllm_ascend/spec_decode/ngram_proposer.py b/vllm_ascend/spec_decode/ngram_proposer.py index bf876cadfd0..6b10a845cbb 100644 --- a/vllm_ascend/spec_decode/ngram_proposer.py +++ b/vllm_ascend/spec_decode/ngram_proposer.py @@ -1,8 +1,6 @@ import torch from vllm.v1.spec_decode.ngram_proposer import NgramProposer -from vllm_ascend.utils import vllm_version_is - class AscendNgramProposer(NgramProposer): def __init__(self, vllm_config, runner): @@ -28,75 +26,35 @@ def dummy_run( ): pass - if vllm_version_is("0.23.0"): - - def propose( - self, - sampled_token_ids: list[list[int]], - num_tokens_no_spec=None, - token_ids_cpu=None, - slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, - ) -> list[list[int]]: - input_batch = self.runner.input_batch - valid_ngram_requests = [] - for i, sampled_ids in enumerate(sampled_token_ids): - num_sampled_ids = len(sampled_ids) - if not num_sampled_ids: - continue - - req_id = input_batch.req_ids[i] - if req_id in input_batch.spec_decode_unsupported_reqs: - continue - - num_tokens = input_batch.num_tokens_no_spec[i] - if num_tokens >= input_batch.max_model_len: - # Skip requests that have already reached the max model length. - continue - - start_idx = input_batch.num_tokens_no_spec[i] - end_idx = start_idx + num_sampled_ids - input_batch.token_ids_cpu[i, start_idx:end_idx] = sampled_ids - - valid_ngram_requests.append(i) - - return self.batch_propose( - len(sampled_token_ids), - valid_ngram_requests, - input_batch.num_tokens_no_spec, - input_batch.token_ids_cpu, - ) - - else: - - def propose( # type: ignore[misc] - self, - num_speculative_tokens: int, - sampled_token_ids: list[list[int]], - num_tokens_no_spec=None, - token_ids_cpu=None, - slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, - ) -> list[list[int]]: - assert num_speculative_tokens <= self.k - assert num_tokens_no_spec is not None - assert token_ids_cpu is not None - - valid_ngram_requests = [] - for i, sampled_ids in enumerate(sampled_token_ids): - num_sampled_ids = len(sampled_ids) - if not num_sampled_ids: - continue - - num_tokens = num_tokens_no_spec[i] - if num_tokens >= self.max_model_len: - # Skip requests that have already reached the max model length. - continue - - valid_ngram_requests.append(i) - - return self.batch_propose( - len(sampled_token_ids), - valid_ngram_requests, - num_tokens_no_spec, - token_ids_cpu, - num_speculative_tokens, - ) + def propose( # type: ignore[misc] + self, + num_speculative_tokens: int, + sampled_token_ids: list[list[int]], + num_tokens_no_spec=None, + token_ids_cpu=None, + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, + ) -> list[list[int]]: + assert num_speculative_tokens <= self.k + assert num_tokens_no_spec is not None + assert token_ids_cpu is not None + + valid_ngram_requests = [] + for i, sampled_ids in enumerate(sampled_token_ids): + num_sampled_ids = len(sampled_ids) + if not num_sampled_ids: + continue + + num_tokens = num_tokens_no_spec[i] + if num_tokens >= self.max_model_len: + # Skip requests that have already reached the max model length. + continue + + valid_ngram_requests.append(i) + + return self.batch_propose( + len(sampled_token_ids), + valid_ngram_requests, + num_tokens_no_spec, + token_ids_cpu, + num_speculative_tokens, + ) diff --git a/vllm_ascend/spec_decode/suffix_proposer.py b/vllm_ascend/spec_decode/suffix_proposer.py index 80446c2ce32..1129203d644 100644 --- a/vllm_ascend/spec_decode/suffix_proposer.py +++ b/vllm_ascend/spec_decode/suffix_proposer.py @@ -1,8 +1,6 @@ import torch from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer -from vllm_ascend.utils import vllm_version_is - class AscendSuffixDecodingProposer(SuffixDecodingProposer): def __init__(self, vllm_config, runner): @@ -31,12 +29,9 @@ def propose( num_speculative_tokens: int = 0, slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, ): - if vllm_version_is("0.23.0"): - return super().propose(self.runner.input_batch, sampled_token_ids) - else: - return super().propose( - num_speculative_tokens, - self.runner.input_batch, - sampled_token_ids, - slot_mappings, - ) + return super().propose( + num_speculative_tokens, + self.runner.input_batch, + sampled_token_ids, + slot_mappings, + ) diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index 2a3f2fa5dc5..608ab019ce8 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -758,11 +758,6 @@ def register_ascend_customop(vllm_config: VllmConfig | None = None): "GatedDeltaNetAttention": AscendGatedDeltaNetAttention, "BailingMoELinearAttention": AscendBailingMoELinearAttention, } - if vllm_version_is("0.23.0"): - from vllm_ascend.ops.fused_moe.fused_moe import AscendFusedMoE - - REGISTERED_ASCEND_OPS["FusedMoE"] = AscendFusedMoE - if vllm_config is None: try: from vllm.config import get_current_vllm_config @@ -807,11 +802,6 @@ def register_ascend_customop(vllm_config: VllmConfig | None = None): "MRotaryEmbedding": AscendMRotaryEmbedding310, } ) - if vllm_version_is("0.23.0"): - from vllm_ascend._310p.fused_moe.fused_moe import AscendFusedMoE310 - - REGISTERED_ASCEND_OPS["FusedMoE"] = AscendFusedMoE310 - for name, op_cls in REGISTERED_ASCEND_OPS.items(): CustomOp.register_oot(_decorated_op_cls=op_cls, name=name) diff --git a/vllm_ascend/worker/block_table.py b/vllm_ascend/worker/block_table.py index a81117f25f3..83fd41861bb 100644 --- a/vllm_ascend/worker/block_table.py +++ b/vllm_ascend/worker/block_table.py @@ -167,10 +167,10 @@ def compute_slot_mapping( "PAD_ID": PAD_SLOT_ID, "BLOCK_SIZE": 1024, } - if not vllm_version_is("0.23.0"): + if not vllm_version_is("0.24.0"): # vLLM #40996 split physical KV blocks into kernel blocks in # the slot-mapping kernel. These are required constexprs on - # main; the v0.23.0 kernel does not accept them. + # main; the v0.24.0 kernel does not accept them. kernel_kwargs.update( KV_CACHE_BLOCK_SIZE=self.physical_block_size, BLOCKS_PER_KV_BLOCK=self.blocks_per_phys_block, diff --git a/vllm_ascend/worker/encoder_acl_graph.py b/vllm_ascend/worker/encoder_acl_graph.py index 186c0926f40..bfa7d67ff7a 100644 --- a/vllm_ascend/worker/encoder_acl_graph.py +++ b/vllm_ascend/worker/encoder_acl_graph.py @@ -27,7 +27,7 @@ from vllm.platforms import current_platform from vllm.v1.worker.encoder_cudagraph import BudgetGraphMetadata, EncoderCudaGraphManager -from vllm_ascend.utils import vllm_version_is, weak_ref_tensors +from vllm_ascend.utils import weak_ref_tensors # --------------------------------------------------------------------------- # Per–encoder-budget ACL graph bookkeeping (ViT FIA tasks) @@ -291,30 +291,18 @@ def _capture_budget_graph(self, token_budget: int, path: str = "default"): self.max_frames_per_batch, ) - if vllm_version_is("0.23.0"): - capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs( - token_budget, - self.max_batch_size, - self.max_frames_per_batch, - self.device, - self.dtype, - ) - else: - capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs( - token_budget, - self.max_batch_size, - self.max_frames_per_batch, - self.device, - self.dtype, - path, - ) + capture_inputs = self.model.prepare_encoder_cudagraph_capture_inputs( + token_budget, + self.max_batch_size, + self.max_frames_per_batch, + self.device, + self.dtype, + path, + ) values = capture_inputs.values with torch.inference_mode(): - if vllm_version_is("0.23.0"): - output = self.model.encoder_cudagraph_forward(dict(values)) - else: - output = self.model.encoder_cudagraph_forward(dict(values), path=path) + output = self.model.encoder_cudagraph_forward(dict(values), path=path) output_buffer = torch.empty_like(output) graph = torch.npu.NPUGraph() @@ -323,10 +311,7 @@ def _capture_budget_graph(self, token_budget: int, path: str = "default"): torch.inference_mode(), torch.npu.graph(graph, self.graph_pool), ): - if vllm_version_is("0.23.0"): - output = self.model.encoder_cudagraph_forward(dict(values)) - else: - output = self.model.encoder_cudagraph_forward(dict(values), path=path) + output = self.model.encoder_cudagraph_forward(dict(values), path=path) output_buffer.copy_(output) graph_meta = BudgetGraphMetadata( @@ -337,11 +322,8 @@ def _capture_budget_graph(self, token_budget: int, path: str = "default"): input_buffers=values, output_buffer=weak_ref_tensors(output_buffer), ) - if vllm_version_is("0.23.0"): - self.budget_graphs[token_budget] = graph_meta - else: - graph_set = self._get_graph_set(path) - graph_set[token_budget] = graph_meta + graph_set = self._get_graph_set(path) + graph_set[token_budget] = graph_meta def _run_budget_graph( self, @@ -350,33 +332,19 @@ def _run_budget_graph( path: str = "default", ) -> torch.Tensor | None: num_items = len(self._get_item_specs(mm_kwargs)) - if vllm_version_is("0.23.0"): - if token_budget not in self.budget_graphs: - self.graph_misses += num_items - return None - graph_meta = self.budget_graphs[token_budget] - else: - graph_set = self._get_graph_set(path) - if token_budget not in graph_set: - self.graph_misses += num_items - return None - graph_meta = graph_set[token_budget] - - if vllm_version_is("0.23.0"): - replay = self.model.prepare_encoder_cudagraph_replay_buffers( - mm_kwargs, - self.max_batch_size, - self.max_frames_per_batch, - ) - buffer_items = ((key, graph_meta.input_buffers[key]) for key in self.config.buffer_keys) - else: - replay = self.model.prepare_encoder_cudagraph_replay_buffers( - mm_kwargs, - self.max_batch_size, - self.max_frames_per_batch, - path, - ) - buffer_items = graph_meta.input_buffers.items() + graph_set = self._get_graph_set(path) + if token_budget not in graph_set: + self.graph_misses += num_items + return None + graph_meta = graph_set[token_budget] + + replay = self.model.prepare_encoder_cudagraph_replay_buffers( + mm_kwargs, + self.max_batch_size, + self.max_frames_per_batch, + path, + ) + buffer_items = graph_meta.input_buffers.items() for key, buf in buffer_items: src = replay.values.get(key) diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index f9028651337..cfb04f1e0c4 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -285,8 +285,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) - if not vllm_version_is("0.23.0"): - self.pin_memory = PIN_MEMORY + self.pin_memory = PIN_MEMORY # Replace the CUDA PrefetchOffloader set by parent __init__ with NPU version. offload_cfg = vllm_config.offload_config @@ -463,7 +462,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): if vllm_config.speculative_config else None ) - if not vllm_version_is("0.23.0"): + if not vllm_version_is("0.24.0"): if vllm_config.speculative_config and vllm_config.speculative_config.use_dspark(): self.use_aux_hidden_state_outputs = True # When True, run update_full_graph_params before self.model (ENPU / graph capture order). @@ -1688,23 +1687,17 @@ def propose_draft_token_ids( # Speculative decoding is not enabled. draft_token_ids = None elif isinstance(self.drafter, AscendNgramProposer): - if vllm_version_is("0.23.0"): - draft_token_ids = self.drafter.propose(valid_sampled_token_ids) - else: - draft_token_ids = self.drafter.propose( - scheduler_output.num_spec_tokens_to_schedule, - valid_sampled_token_ids, - self.input_batch.num_tokens_no_spec, - self.input_batch.token_ids_cpu, - ) + draft_token_ids = self.drafter.propose( + scheduler_output.num_spec_tokens_to_schedule, + valid_sampled_token_ids, + self.input_batch.num_tokens_no_spec, + self.input_batch.token_ids_cpu, + ) elif isinstance(self.drafter, AscendSuffixDecodingProposer): - if vllm_version_is("0.23.0"): - draft_token_ids = self.drafter.propose(valid_sampled_token_ids) - else: - draft_token_ids = self.drafter.propose( - valid_sampled_token_ids, - num_speculative_tokens=scheduler_output.num_spec_tokens_to_schedule, - ) + draft_token_ids = self.drafter.propose( + valid_sampled_token_ids, + num_speculative_tokens=scheduler_output.num_spec_tokens_to_schedule, + ) elif isinstance(self.drafter, AscendNgramProposerNPU): batch_size = min(self.input_batch.num_reqs, self.token_ids_gpu_tensor.shape[0]) @@ -1768,19 +1761,12 @@ def propose_draft_token_ids( common_attn_metadata = spec_decode_common_attn_metadata target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states] - if vllm_version_is("0.23.0"): - draft_token_ids = self.drafter.propose( - sampled_token_ids=valid_sampled_token_ids, - target_hidden_states=target_hidden_states, - common_attn_metadata=common_attn_metadata, - ) - else: - draft_token_ids = self.drafter.propose( - self.speculative_config.num_speculative_tokens, - sampled_token_ids=valid_sampled_token_ids, - target_hidden_states=target_hidden_states, - common_attn_metadata=common_attn_metadata, - ) + draft_token_ids = self.drafter.propose( + self.speculative_config.num_speculative_tokens, + sampled_token_ids=valid_sampled_token_ids, + target_hidden_states=target_hidden_states, + common_attn_metadata=common_attn_metadata, + ) next_token_ids, valid_sampled_tokens_count = ( self.drafter.prepare_next_token_ids_padded( valid_sampled_token_ids, @@ -3929,26 +3915,13 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.init_routed_experts_capturer() def _bind_routed_experts_capturer(self, capturer=None) -> None: - if vllm_version_is("0.23.0"): - # Upstream binds via ``module.router.set_capture_fn(...)`` on - # FusedMoE layers whose router is a ``BaseRouter``. Ascend's - # ``select_experts`` does not go through ``BaseRouter``, so the - # upstream hook never fires. Instead, stash the capturer as a - # plain attribute on every FusedMoE layer; ``apply()`` reads it - # back on the hot path. - from vllm.model_executor.layers.fused_moe.layer import FusedMoE - - for module in self.compilation_config.static_forward_context.values(): - if isinstance(module, FusedMoE): - module._ascend_routed_experts_capturer = capturer - else: - # test_qwen3_moe_routing_replay - from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner + # test_qwen3_moe_routing_replay + from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner - for module in self.compilation_config.static_forward_context.values(): - if isinstance(module, AscendMoERunner): - module._ascend_routed_experts_capturer = capturer - module.routed_experts._ascend_routed_experts_capturer = capturer + for module in self.compilation_config.static_forward_context.values(): + if isinstance(module, AscendMoERunner): + module._ascend_routed_experts_capturer = capturer + module.routed_experts._ascend_routed_experts_capturer = capturer def _align_memory(self, tensor: torch.Tensor, alignment: int) -> torch.Tensor: data_ptr = tensor.data_ptr() @@ -5010,7 +4983,7 @@ def _check_and_update_cudagraph_mode( min_cg_attn_backend = attn_backend.__name__ with update_pass_config(self): - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes( min_cg_support=min_cg_support, min_cg_attn_backend=min_cg_attn_backend, diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py index 46fdc7d8840..ccb4e56a8b2 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py @@ -159,8 +159,11 @@ def __init__( WEIGHT_LOADER_V2_SUPPORTED.remove("UnquantizedLinearMethod") self.use_v2_model_runner = self.vllm_config.use_v2_model_runner - if self.use_v2_model_runner and vllm_version_is("0.23.0"): - logger.warning("VLLM_USE_V2_MODEL_RUNNER is not supported on vllm 0.23.0; falling back to v1 model runner.") + if self.use_v2_model_runner and vllm_version_is("0.24.0"): + logger.warning( + "VLLM_USE_V2_MODEL_RUNNER is supported only on the verified vLLM main commit; " + "falling back to the v1 model runner." + ) self.use_v2_model_runner = False self._pp_send_work: list[Handle] = [] @@ -394,56 +397,53 @@ def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: self.cache_config.num_cpu_blocks = num_cpu_blocks def _init_device(self): - if not vllm_version_is("0.23.0"): - # vLLM v0.24.0 (PR #45026) removed automatic per-process device - # isolation for DP workers. Mirror gpu_worker.py::init_device: - # shift self.local_rank by dp_local_rank * tp_pp_world_size so - # that each DP group binds to a distinct set of NPUs. - parallel_config = self.parallel_config - if ( - parallel_config.distributed_executor_backend not in ("ray", "external_launcher") - and parallel_config.data_parallel_backend != "ray" - and parallel_config.nnodes_within_dp == 1 - # vllm-ascend: when the user pre-shards devices via - # --device-ids (which becomes assigned_physical_gpu_ids), - # each child process already binds to its own NPU(s); the - # DP local_rank shift below would push local_rank past the - # length of the per-rank device list and trip the assert - # in this same method. Skip the shift in that case. - and parallel_config.assigned_physical_gpu_ids is None - ): - dp_local_rank = parallel_config.data_parallel_rank_local - if dp_local_rank is None: - dp_local_rank = parallel_config.data_parallel_index - tp_pp_world_size = parallel_config.pipeline_parallel_size * parallel_config.tensor_parallel_size - self.local_rank += dp_local_rank * tp_pp_world_size - - # Publish the logical-to-physical mapping for topology queries. - assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids - if assigned_physical_gpu_ids is not None: - from vllm.platforms.interface import set_assigned_physical_gpu_ids - - set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) - assert self.local_rank < len(assigned_physical_gpu_ids), ( - f"local_rank {self.local_rank} is out of bounds for " - f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" - ) - if parallel_config.distributed_executor_backend not in ("ray", "external_launcher"): - assert parallel_config.local_world_size <= len(assigned_physical_gpu_ids), ( - f"local_world_size ({parallel_config.local_world_size}) " - f"exceeds assigned_physical_gpu_ids count " - f"({len(assigned_physical_gpu_ids)})" - ) - else: - visible_device_count = torch.npu.device_count() if torch.npu.is_available() else 0 - assert self.local_rank < visible_device_count, ( - f"DP adjusted local rank {self.local_rank} is out of bounds for {visible_device_count} devices." + # vLLM v0.24.0 (PR #45026) removed automatic per-process device + # isolation for DP workers. Mirror gpu_worker.py::init_device: + # shift self.local_rank by dp_local_rank * tp_pp_world_size so + # that each DP group binds to a distinct set of NPUs. + parallel_config = self.parallel_config + if ( + parallel_config.distributed_executor_backend not in ("ray", "external_launcher") + and parallel_config.data_parallel_backend != "ray" + and parallel_config.nnodes_within_dp == 1 + # vllm-ascend: when the user pre-shards devices via + # --device-ids (which becomes assigned_physical_gpu_ids), + # each child process already binds to its own NPU(s); the + # DP local_rank shift below would push local_rank past the + # length of the per-rank device list and trip the assert + # in this same method. Skip the shift in that case. + and parallel_config.assigned_physical_gpu_ids is None + ): + dp_local_rank = parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = parallel_config.data_parallel_index + tp_pp_world_size = parallel_config.pipeline_parallel_size * parallel_config.tensor_parallel_size + self.local_rank += dp_local_rank * tp_pp_world_size + + # Publish the logical-to-physical mapping for topology queries. + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" + ) + if parallel_config.distributed_executor_backend not in ("ray", "external_launcher"): + assert parallel_config.local_world_size <= len(assigned_physical_gpu_ids), ( + f"local_world_size ({parallel_config.local_world_size}) " + f"exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" ) - - visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank) - device = torch.device(f"{current_platform.device_type}:{visible_device_index}") else: - device = torch.device(f"npu:{self.local_rank}") + visible_device_count = torch.npu.device_count() if torch.npu.is_available() else 0 + assert self.local_rank < visible_device_count, ( + f"DP adjusted local rank {self.local_rank} is out of bounds for {visible_device_count} devices." + ) + + visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank) + device = torch.device(f"{current_platform.device_type}:{visible_device_index}") torch.npu.set_device(device) @@ -463,10 +463,7 @@ def _init_device(self): setup_ascend_local_comm_res(self.local_rank, self.vllm_config.kv_transfer_config) # take current memory snapshot - if vllm_version_is("0.23.0"): - self.init_snapshot = MemorySnapshot() - else: - self.init_snapshot = MemorySnapshot(device=device) + self.init_snapshot = MemorySnapshot(device=device) self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization if self.init_snapshot.free_memory < self.requested_memory: GiB = lambda b: round(b / GiB_bytes, 2) @@ -684,7 +681,7 @@ def load_model(self) -> None: WeightTransferEngineFactory, ) - if vllm_version_is("0.23.0"): + if vllm_version_is("0.24.0"): self.weight_transfer_engine = WeightTransferEngineFactory.create_engine( self.vllm_config.weight_transfer_config, self.vllm_config.parallel_config,