From ed6ae1e36a03bed4a29287163e051a7772b1d8b1 Mon Sep 17 00:00:00 2001 From: Xiao Li Date: Thu, 20 Nov 2025 17:54:35 -0800 Subject: [PATCH] [AITER] [ROCm] Fix crash when loading llama4 model with old aiter version installed, fallback to forward_native implementation (#29124) Signed-off-by: Xiao Li --- vllm/v1/sample/ops/topk_topp_sampler.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index c6c7e924175f7..5b2d130b0ea42 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -60,13 +60,20 @@ class TopKTopPSampler(nn.Module): logprobs_mode not in ("processed_logits", "processed_logprobs") and rocm_aiter_ops.is_enabled() ): - import aiter.ops.sampling # noqa: F401 + try: + import aiter.ops.sampling # noqa: F401 - self.aiter_ops = torch.ops.aiter - logger.info_once( - "Using aiter sampler on ROCm (lazy import, sampling-only)." - ) - self.forward = self.forward_hip + self.aiter_ops = torch.ops.aiter + logger.info_once( + "Using aiter sampler on ROCm (lazy import, sampling-only)." + ) + 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: self.forward = self.forward_native