[BugFix][Spec Decode] Resolve draft ACL graph params per layer for Gemma4 MTP Problem - #16737
swimming2007-doge wants to merge 2 commits into
Conversation
…mma4 MTP Problem ------- Gemma4 MTP acceptance on vLLM 0.28 + main collapses. The probe (measure_acceptance.py, k=3, greedy, 5 coding prompts x 3 repeats) drops from the 94.8% baseline (pos0/pos1/pos2 = 98.2/96.3/89.9) to 5%~57%, and the value differs between runs of the same build: 11.7/3.0/1.3 = 5.3% and 77.8/56.0/38.4 = 57.4% Background ---------- The Gemma4 draft spans two KV cache groups: draft layers 0-2 share the KV of target layer 58 (sliding attention), draft layer 3 shares target layer 59 (full attention). A draft step therefore carries a different block table per group. The feature was developed and validated on the v0.26.0rc1 base, where the host-side attention metadata of a captured draft graph was refreshed by AscendAttentionBackendImpl.update_graph_params(), which resolved metadata per layer *and* per draft step. It was merged one day after vllm-project#15908 ("[MRV1][MRV2] Refactor Host-Side Parameter Updates for ACL Graph Replay", later re-landed as vllm-project#16425), which replaced that path with UpdatableGraph. FIA/PA backends now require UpdatableGraph -- update_graph_params() raises NotImplementedError -- so there is no per-layer path left to fall back to. Root cause ---------- AscendSpecDecodeBaseProposer._maybe_update_metadata() flattens every draft step to a single representative layer before handing params to the replay: metadata = next(iter(per_layer_metadata.values())) UpdatableGraph.resolve_tasks() then binds that one block table to every provider, so all four draft layers read the sliding group's blocks. Which group comes "first" depends on the set iteration order in Gemma4Proposer.initialize_attn_backend(), which is why the observed acceptance varies from process to process. Evidence -------- * cudagraph_mode=NONE (the UpdatableGraph path is not entered) measures 98.2/96.3/89.9 = 94.8%, identical to the v0.26.0rc1 baseline. * The flattening is already present in ffc2112^ (main immediately before this feature merged), i.e. it predates the merge; v0.26.0rc1 has no _maybe_update_metadata at all. Fix --- Resolve the replay params per layer instead of per step: SharedSource delegates to a ParamSource when one is supplied, and AscendGemma4Proposer installs a source that resolves each provider against its own layer for every draft step. Verification ------------ Same card, TP=1, FULL_DECODE_ONLY, cold compile cache: pos0 98.2 / pos1 96.3 / pos2 89.9 = 94.8%, avg accepted 2.84/3 -- byte-identical to the v0.26.0rc1 baseline. Reproduced on a second, freshly started process (a different set-iteration order, which the buggy build was sensitive to), and the ACL graph is still captured and replayed ("Replaying aclgraph"), so this is not a workaround that falls back to eager. Sara workload (CONC=16, NUM=300, long prompts; same server): spec_accept 63.4% (pos0 81.6 / pos1 60.2 / pos2 48.5), 300/300 ok, 1.99 req/s -- matches the documented 63~64% baseline for this workload. Co-Authored-By: Claude Code <noreply@anthropic.com>
…4 MTP Guards the fix in the previous commit: - _LayerResolvedSource resolves each provider against its own layer, and keeps the representative-layer fallback for providers without a layer name; - AscendGemma4Proposer._maybe_update_metadata installs that source instead of the flattened per-step params; - SharedSource delegates to a ParamSource rather than returning it as params. Without these, a multi-KV-group draft silently serves one group's block table to every layer again, and that is not visible from the acceptance probe alone (it only shows up as a lower acceptance rate). Co-Authored-By: Claude Code <noreply@anthropic.com>
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! The pull request fix/gemma4, developed by swimming2007-doge, addresses critical speculative decoding collapses in Gemma4 MTP models when running on the vLLM and vllm-ascend backend with compilation enabled. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Attention][Feature] Support layer-resolved metadata updates for Gemma4 speculative decodingSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces `_LayerResolvedSource` to support layer-resolved metadata updates for Gemma4 speculative decoding in `vllm-ascend`. This allows different attention layers (such as sliding vs. full attention) to resolve their own block tables and sequence lengths during speculative decoding steps. Additionally, `SharedSource` is updated to support delegation to a `ParamSource`.
Feedback: In `vllm_ascend/compilation/updatable_graph.py`, using `isinstance` with a `@runtime_checkable` protocol in the hot path can introduce significant runtime overhead. It is recommended to use `hasattr(self.params, 'get')` instead to avoid this overhead.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New unit tests were added in `tests/ut/compilation/test_updatable_graph.py` and `tests/ut/spec_decode/test_gemma4_proposer.py` to verify `SharedSource` delegation, `_LayerResolvedSource` parameter resolution, and metadata updates.| if isinstance(self.params, ParamSource): | ||
| return self.params.get(provider) |
There was a problem hiding this comment.
Using isinstance with a @runtime_checkable protocol in the hot path (called for every provider on every speculative decoding step) introduces significant runtime overhead due to structural checks. Since self.params is either a Sequence (which does not have a get method) or a ParamSource (which does), we can use a much faster attribute check like hasattr(self.params, 'get') to avoid this overhead.
| if isinstance(self.params, ParamSource): | |
| return self.params.get(provider) | |
| if hasattr(self.params, "get"): | |
| return self.params.get(provider) |
swimming2007-doge
changed the title
swimming2007-doge
changed the title
|
upstream has fixed this problem. |
2 participants
[BugFix][Spec Decode] Resolve draft ACL graph params per layer for Gemma4 MTP
Problem
Gemma4 MTP acceptance on vLLM 0.28 + main collapses. The probe
(measure_acceptance.py, k=3, greedy, 5 coding prompts x 3 repeats) drops from
the 94.8% baseline (pos0/pos1/pos2 = 98.2/96.3/89.9) to 5%~57%, and the value
differs between runs of the same build:
11.7/3.0/1.3 = 5.3% and 77.8/56.0/38.4 = 57.4%
Background
The Gemma4 draft spans two KV cache groups: draft layers 0-2 share the KV of
target layer 58 (sliding attention), draft layer 3 shares target layer 59
(full attention). A draft step therefore carries a different block table per
group.
The feature was developed and validated on the v0.26.0rc1 base, where the
host-side attention metadata of a captured draft graph was refreshed by
AscendAttentionBackendImpl.update_graph_params(), which resolved metadata per
layer and per draft step. It was merged one day after #15908 ("[MRV1][MRV2]
Refactor Host-Side Parameter Updates for ACL Graph Replay", later re-landed as
#16425), which replaced that path with UpdatableGraph. FIA/PA backends now
require UpdatableGraph -- update_graph_params() raises NotImplementedError --
so there is no per-layer path left to fall back to.
Root cause
AscendSpecDecodeBaseProposer._maybe_update_metadata() flattens every draft step
to a single representative layer before handing params to the replay:
UpdatableGraph.resolve_tasks() then binds that one block table to every
provider, so all four draft layers read the sliding group's blocks. Which group
comes "first" depends on the set iteration order in
Gemma4Proposer.initialize_attn_backend(), which is why the observed acceptance
varies from process to process.
Evidence
98.2/96.3/89.9 = 94.8%, identical to the v0.26.0rc1 baseline.
feature merged), i.e. it predates the merge; v0.26.0rc1 has no
_maybe_update_metadata at all.
Fix
Resolve the replay params per layer instead of per step: SharedSource delegates
to a ParamSource when one is supplied, and AscendGemma4Proposer installs a
source that resolves each provider against its own layer for every draft step.
Verification
Same card, TP=1, FULL_DECODE_ONLY, cold compile cache:
pos0 98.2 / pos1 96.3 / pos2 89.9 = 94.8%, avg accepted 2.84/3
-- byte-identical to the v0.26.0rc1 baseline. Reproduced on a second, freshly
started process (a different set-iteration order, which the buggy build was
sensitive to), and the ACL graph is still captured and replayed
("Replaying aclgraph"), so this is not a workaround that falls back to eager.
Sara workload (CONC=16, NUM=300, long prompts; same server):
spec_accept 63.4% (pos0 81.6 / pos1 60.2 / pos2 48.5), 300/300 ok, 1.99 req/s
-- matches the documented 63~64% baseline for this workload.