[BugFix][0.18.0][310p] fix post-sampling not working in graph mode on 310p - #8077
Conversation
… 310p Signed-off-by: Tflowers-0129 <2906339855@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 resolves a significant stability issue on 310P devices where post-sampling with temperature enabled would lead to system hangs. The solution involves implementing a dedicated sampling pipeline for 310P that intelligently handles device-specific limitations by moving certain operations to the CPU, thereby preventing the unsupported kernel calls and ensuring robust model execution. Highlights
🧠 New Feature in Public Preview: You can now enable Memory 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 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 counter productive. 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
|
![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
This pull request introduces a specialized sampler for 310P devices to resolve a hang issue associated with NPU-based exponential noise generation by moving the operation to the CPU. The changes include the implementation of _random_sample_310p and the AscendSampler310 class. Feedback suggests adhering to the repository's PR title and summary formatting standards, optimizing the CPU-to-NPU data transfer using pinned memory and non-blocking operations, and overriding the do_async_exponential method in AscendSampler310 to prevent the base class from triggering the buggy NPU operator.
| MambaSpec, | ||
| UniformTypeKVCacheSpecs, | ||
| ) | ||
| from vllm.v1.sample.rejection_sampler import RejectionSampler |
There was a problem hiding this comment.
The Pull Request title and summary should follow the repository's style guide format.
Suggested PR Title:
[0.18.0][310p][BugFix] fix post-sampling not working in graph mode on 310PSuggested PR Summary:
### What this PR does / why we need it?
This PR fixes a hang issue on 310P devices when temperature is enabled in post-processing. The root cause was identified as the `q.exponential_()` operator on NPU, which is buggy on 310P. The fix involves moving the exponential noise generation to the CPU.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Tested locally with accuracy-dataset and stress tests.References
- Pull Request titles must follow the [Branch][Module][Action] format, and the summary must follow the specified markdown structure. (link)
| with npu_stream_switch(global_stream()): | ||
| q = torch.empty_like(probs) | ||
| q = q.cpu() | ||
| if len(generators) != q.shape[0]: | ||
| q.exponential_() | ||
| if generators: | ||
| for i, generator in generators.items(): | ||
| q[i].exponential_(generator=generator) | ||
| q = q.npu() |
There was a problem hiding this comment.
The current implementation is inefficient because it allocates q on the NPU and then copies it to the CPU while it's still uninitialized. It's better to allocate q directly on the CPU. Additionally, using pinned memory (pin_memory=True) and an asynchronous transfer (non_blocking=True) will improve performance when moving the tensor back to the NPU.
| with npu_stream_switch(global_stream()): | |
| q = torch.empty_like(probs) | |
| q = q.cpu() | |
| if len(generators) != q.shape[0]: | |
| q.exponential_() | |
| if generators: | |
| for i, generator in generators.items(): | |
| q[i].exponential_(generator=generator) | |
| q = q.npu() | |
| q = torch.empty(probs.shape, | |
| dtype=probs.dtype, | |
| device='cpu', | |
| pin_memory=True) | |
| if len(generators) != q.shape[0]: | |
| q.exponential_() | |
| if generators: | |
| for i, generator in generators.items(): | |
| q[i].exponential_(generator=generator) | |
| with npu_stream_switch(global_stream()): | |
| q = q.to(device=probs.device, non_blocking=True) |
| def __init__(self, logprobs_mode=DEFAULT_LOGPROBS_MODE): | ||
| super().__init__(logprobs_mode=logprobs_mode) | ||
| self.topk_topp_sampler = AscendTopKTopPSampler310(logprobs_mode=logprobs_mode) |
There was a problem hiding this comment.
The base class AscendSampler implements do_async_exponential, which performs q.exponential_() on the NPU. If enable_async_exponential is enabled in the configuration, this method will still be called during model execution, which could trigger the 310P hang issue even if the result is not used by the sampler. This method should be overridden to avoid executing the buggy NPU operator.
| def __init__(self, logprobs_mode=DEFAULT_LOGPROBS_MODE): | |
| super().__init__(logprobs_mode=logprobs_mode) | |
| self.topk_topp_sampler = AscendTopKTopPSampler310(logprobs_mode=logprobs_mode) | |
| def __init__(self, logprobs_mode=DEFAULT_LOGPROBS_MODE): | |
| super().__init__(logprobs_mode=logprobs_mode) | |
| self.topk_topp_sampler = AscendTopKTopPSampler310(logprobs_mode=logprobs_mode) | |
| def do_async_exponential(self, b_s, head_dim, generators): | |
| # Disable async NPU exponential on 310P to avoid hangs. | |
| pass |
wangxiyuan
changed the title
82e17f6
into
vllm-project:releases/v0.18.0
… 310p (vllm-project#8077) ### What this PR does / why we need it? Enabling temperature in post-processing on 310P devices can cause the service to stall and eventually hang. We first traced the issue to a timeout where the temperature-related `div` operator was waiting for results from a sub-stream. After investigating the preceding operators, we finally identified the root cause as the `q.exponential_()` operator, which is not well supported on 310P and triggers an internal issue in the `add` kernel. ### Does this PR introduce _any_ user-facing change? NA ### How was this patch tested? This patch was thoroughly tested locally（accuracy-dataset test and stress test）. It is not easy to design a proper unit test for this case, and I appreciate your understanding. Signed-off-by: Tflowers-0129 <2906339855@qq.com>
… 310p (vllm-project#8077) ### What this PR does / why we need it? Enabling temperature in post-processing on 310P devices can cause the service to stall and eventually hang. We first traced the issue to a timeout where the temperature-related `div` operator was waiting for results from a sub-stream. After investigating the preceding operators, we finally identified the root cause as the `q.exponential_()` operator, which is not well supported on 310P and triggers an internal issue in the `add` kernel. ### Does this PR introduce _any_ user-facing change? NA ### How was this patch tested? This patch was thoroughly tested locally（accuracy-dataset test and stress test）. It is not easy to design a proper unit test for this case, and I appreciate your understanding. Signed-off-by: Tflowers-0129 <2906339855@qq.com>
…V2 on the 310P (#16503) ### What this PR does / why we need it? Enable **temperature / top-k / top-p** sampling on Ascend **310P Model Runner V2**. - First-version `Ascend310PSampler` only supported greedy (`argmax`) and rejected non-zero temperature. - Mainline MRV2 uses Triton Gumbel sampling; 310P has no Triton, and NPU `exponential_` / large RNG can hang under ACLGraph (same issue fixed in MRV1). - This PR reuses the **MRV1 inverse-CDF path** (`_random_sample_310p`: CPU uniform per request → NPU `softmax` + `cumsum` + `searchsorted`), while keeping the MRV2 `sampling_states` surface required by MTP draft `propose()`. - Still rejects min_p / penalties / logprobs / logit processors (follow-up). Refs: MRV1 post-sampling fixes (#8017 / #8077), CDF optimization (#12966). RFC: #15577 ### Does this PR introduce _any_ user-facing change? No. On 310P with `VLLM_USE_V2_MODEL_RUNNER=1`, requests may now use `temperature`, `top_k`, and `top_p` (previously raised `NotImplementedError`). Default greedy behavior is unchanged. ### How was this patch tested? - UT: `tests/ut/_310p/test_model_runner_v2_310p.py` (temperature accept / scale / greedy; still rejects penalties) - vLLM main: vllm-project/vllm@84030bb --------- Signed-off-by: Thiagor2002 <13476117628@163.com>
2 participants
What this PR does / why we need it?
Enabling temperature in post-processing on 310P devices can cause the service to stall and eventually hang. We first traced the issue to a timeout where the temperature-related
divoperator was waiting for results from a sub-stream. After investigating the preceding operators, we finally identified the root cause as theq.exponential_()operator, which is not well supported on 310P and triggers an internal issue in theaddkernel.Does this PR introduce any user-facing change?
NA
How was this patch tested?
This patch was thoroughly tested locally（accuracy-dataset test and stress test）. It is not easy to design a proper unit test for this case, and I appreciate your understanding.