biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[AITER] [ROCm] Fix crash when loading llama4 model with old aiter version installed, fallback to forward_native implementation (#29124)

Browse Source 
Signed-off-by: Xiao Li <ilx@meta.com>
This commit is contained in:
Xiao Li
2025-11-20 17:54:35 -08:00
committed by GitHub
parent 9875be6431
commit ed6ae1e36a
1 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/v1/sample/ops/topk_topp_sampler.py
View File

@@ -60,6 +60,7 @@ class TopKTopPSampler(nn.Module):
logprobs_mode not in ("processed_logits", "processed_logprobs") logprobs_mode not in ("processed_logits", "processed_logprobs")
and rocm_aiter_ops.is_enabled() and rocm_aiter_ops.is_enabled()
): ):
try:
import aiter.ops.sampling # noqa: F401 import aiter.ops.sampling # noqa: F401
self.aiter_ops = torch.ops.aiter self.aiter_ops = torch.ops.aiter
@@ -67,6 +68,12 @@ class TopKTopPSampler(nn.Module):
"Using aiter sampler on ROCm (lazy import, sampling-only)." "Using aiter sampler on ROCm (lazy import, sampling-only)."
) )
self.forward = self.forward_hip self.forward = self.forward_hip
except ImportError:
logger.warning_once(
"aiter.ops.sampling is not available on ROCm. "
"Falling back to forward_native implementation."
)
self.forward = self.forward_native
else: else:
self.forward = self.forward_native self.forward = self.forward_native