A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 13:30 UTC
[ROCm] Enable the aiter top-k/top-p sampler by default (#43331)
Signed-off-by: John Qin <yanyuan.qin@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
JohnQinAMD committed
May 28, 2026 at 11:19 UTC
ed7fe831da79d8db2dfb3ac53108305c5b852e5b
1 file changed
-4
vllm/v1/sample/ops/topk_topp_sampler.py
-4
@@ -193,8 +193,6 @@ class TopKTopPSampler(nn.Module):
193
k: torch.Tensor | None,
194
p: torch.Tensor | None,
195
) -> tuple[torch.Tensor, torch.Tensor | None]:
196
- # FIXME: Fix aiter_sampler's accuracy issue and remove this flag
197
- DISABLE_AITER_SAMPLER = True
196
"""Optimized ROCm/aiter path (same structure as forward_cuda)."""
197
if (k is None and p is None) or generators:
198
if generators:
@@ -207,8 +205,6 @@ class TopKTopPSampler(nn.Module):
205
"processed_logits",
206
"processed_logprobs",
207
), "aiter sampler does not support returning logits/logprobs."
210
- if DISABLE_AITER_SAMPLER:
211
- return self.forward_native(logits, generators, k, p)
208
return self.aiter_sample(logits, k, p, generators), None
209
210
def aiter_sample(