[BugFix][0.18.0][310p] fix post-sampling not working in graph mode on 310p (#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>
This commit is contained in:
Shaoxu Cheng
2026-04-09 16:31:38 +08:00
committed by GitHub
parent 0d1424d81a
commit 82e17f693a
3 changed files with 74 additions and 0 deletions

View File

@@ -35,7 +35,9 @@ from vllm.v1.kv_cache_interface import (
MambaSpec,
UniformTypeKVCacheSpecs,
)
from vllm.v1.sample.rejection_sampler import RejectionSampler
from vllm_ascend._310p.sample.sampler import AscendSampler310
from vllm_ascend.attention.attention_v1 import AscendAttentionState
from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ
from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
@@ -52,6 +54,9 @@ class NPUModelRunner310(NPUModelRunner):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._acl_format = ACL_FORMAT_FRACTAL_NZ
self.sampler = AscendSampler310()
if getattr(self, "rejection_sampler", None) is not None:
self.rejection_sampler = RejectionSampler(self.sampler)
if self.speculative_config is not None and self.speculative_config.method == "ngram":
# 310P ngram requires decode-only graph shapes to be built with q_len=1.
# Keep dispatcher's internal query_len in sync to avoid key-init assert.