[Spec Decode][MRV2] Support DFlash2 in MRV2 - #14594
YzTongNiar wants to merge 5 commits into
Conversation
Route DFlash2DraftModel architectures to AscendDFlash2Speculator on the V2 model runner. The speculator combines the upstream DFlash2Speculator (candidate top-k, low-rank bilinear selector scoring, sequential path walk, and draft-logits scatter for lossless rejection sampling) with the Ascend NPU wiring inherited from AscendDFlashSpeculator (attention metadata building, aclgraph manager, set_attn bookkeeping, propose wrapper) via cooperative MRO. Includes: - triton-ascend adaptation of the selector walk kernel: tldevice.log1p has no NPU lowering and is replaced by the algebraically equivalent tl.log(1 - u); philox (tl.rand) requires int32 seed/offset operands, so the int64 seeds, sample positions and candidate ids are cast (an int64 counter lowers to 64-bit multiplies whose runtime helper __multi3 is unavailable); fp64 gumbel noise is unavailable on NPU. The kernel is swapped onto the upstream dflash2 speculator module from patch_v2/patch_dflash_speculator. - guard the seq_lens update in the shared autoregressive speculator for hybrid linear-attention layers (GDN metadata has no seq_lens), which otherwise crashes the first propose on Qwen3.5/Qwen3.8-family models. Verified on 910B: both kernels match torch references (greedy walk elementwise, T=1 sampling path executes, full-vocab scatter including duplicate-token overwrite and -inf reset); e2e reaches the draft loop. Signed-off-by: YzTongNiar <1667927948@qq.com>
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! This pull request introduces support for DFlash2 speculative decoding within the MRV2 framework for Ascend NPUs. It includes the necessary infrastructure to integrate DFlash2, provides a specialized Triton kernel to handle NPU-specific constraints, and addresses a sequence length tracking issue that affected hybrid attention models. 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
|
|
👋 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. |
![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:
[Ops][Feature] Add AscendDFlash2Speculator and support for DFlash2 draft modelsSuggested PR Summary:
### What this PR does / why we need it?
This PR introduces support for the `DFlash2` speculative drafter on Ascend NPUs by implementing `AscendDFlash2Speculator`. It includes:
- Swapping in an Ascend-compatible Triton kernel (`_selector_walk_kernel_ascend`) that replaces `tldevice.log1p` with an algebraically equivalent `log(1 - u)` transform and uses `int32` operands for Philox compatibility.
- Adding initialization logic for `AscendDFlash2Speculator` when `DFlash2DraftModel` is present in the architectures.
- Adding a safety check in `_ascend_update_seq_lens` to skip sequence length updates for hybrid models that do not track `seq_lens`.
Feedback:
A review comment suggests improving the safety of the `hasattr(attn_meta, "seq_lens")` check by using `getattr(attn_meta, "seq_lens", None) is None` to handle cases where `attn_meta` or `seq_lens` itself is `None`.
### Does this PR introduce _any_ user-facing change?
No, this is an internal enablement of DFlash2 speculative decoding on Ascend NPUs.
### How was this patch tested?
CI tests and verification on Ascend hardware.| if not hasattr(attn_meta, "seq_lens"): | ||
| continue |
There was a problem hiding this comment.
Using hasattr(attn_meta, "seq_lens") is not fully safe if attn_meta is None or if the seq_lens attribute itself is None. To prevent potential TypeError or AttributeError crashes in hybrid models or edge cases, it is safer to use getattr(attn_meta, "seq_lens", None) is None to handle both cases defensively.
| if not hasattr(attn_meta, "seq_lens"): | |
| continue | |
| if getattr(attn_meta, "seq_lens", None) is None: | |
| continue |
Upstream model states (MambaHybridAttnMetadata) pass is_prefilling per request through get_extra_common_attn_kwargs, while build_attn_metadata also passed it explicitly, so the first real request on any hybrid attention model crashes with: AscendCommonAttentionMetadata() got multiple values for keyword argument 'is_prefilling'. Merge via setdefault: the model-specific tensor wins when present, the explicit parameter remains the fallback. Verified: Qwen3.8-27B (48 linear-attention + 16 full-attention layers) serves correctly on MRV2 after the fix. Signed-off-by: YzTongNiar <1667927948@qq.com>
With the target on FULL_DECODE_ONLY, the draft previously went through dynamo twice and crashed: torch_npu fake tensors mis-derive the selector's hidden[:, :, None] broadcast (mul((s, 7, 16, 256), (s, 256, 1)) fails dimension checks), and the spec-level enforce_eager is not consulted by the V2 speculator graph decision. - AscendDFlash2Speculator.init_cudagraph_manager: honor speculative_config.enforce_eager by forcing CUDAGraphMode.NONE for the draft speculator while the target keeps its own cudagraph mode. - patch_qwen3_dflash: mark the DFlash/DFlash2 draft model classes and CandidateSelector with _ignore_compile_vllm so the draft skips torch.compile entirely (submodule decorators compile independently of the outer model, so each decorated class must be flagged). Verified on 910B: target FULL_DECODE_ONLY + eager DFlash2 draft serves greedy output identical to full-eager, acceptance length ~5.6 @ num_speculative_tokens=7. Signed-off-by: YzTongNiar <1667927948@qq.com>
Signed-off-by: lilinsiman <lilinsiman@gmail.com>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
| # Hybrid models mix full-attention layers with linear-attention | ||
| # ones (e.g. GDN), whose metadata tracks recurrent state via | ||
| # indices and has no seq_lens to bump. | ||
| if not hasattr(attn_meta, "seq_lens"): |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
| ) | ||
|
|
||
| DFlashQwen3Model._ignore_compile_vllm = True | ||
| DFlash2Qwen3Model._ignore_compile_vllm = True |
There was a problem hiding this comment.
open a new patch file named patch_qwen3_dflash2
| ) | ||
| # Upstream model states (e.g. MambaHybrid) pass is_prefilling per | ||
| # request through the extra kwargs; merge instead of passing it twice. | ||
| common_attn_metadata_extra_kwargs.setdefault("is_prefilling", is_prefilling) |
Signed-off-by: lilinsiman <lilinsiman@gmail.com>
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
3 participants
What this PR does / why we need it?
Main Idea is supporting DFlash2 (https://inco.ai/blog/dflash2/) in MRV2
Currently, only eager mode is supported
Refering:
vllm-project/vllm#52816
vllm-project/vllm#52883
Includes:
Does this PR introduce any user-facing change?
No change, accoridng to vllm upsteram PR, use the same argument:
--speculative-config '{"method": "dflash", "model": "draft_model_path"How was this patch tested?
Currently, only Qwen3.8-27B-DFlash2 is available in MRV2, thus the PR is tested under this model
Use this script for testing in A3:
vllm v0.27.1
Need Upstream Cherry-Pick:
vllm-project/vllm#52816
vllm-project/vllm#52883
vllm-ascend: main @ 1ba7ee4