From e0613702ade9ace874feabb7b6f080cdfd181f4b Mon Sep 17 00:00:00 2001 From: "Chuan (Richard) Li" Date: Thu, 9 Apr 2026 17:56:17 -0700 Subject: [PATCH] [ROCm] Fix AITER ops fake impl and minor bugs (#36092) Signed-off-by: Li --- vllm/_aiter_ops.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index d59b74782be..8c2659b9c7e 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -336,9 +336,13 @@ def _rocm_aiter_fused_topk_fake( router_logits: torch.Tensor, top_k: int, gate_up: bool, -) -> None: - # tuple[torch.Tensor, torch.Tensor]: - pass +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = x.shape[0] + topk_weights = torch.empty( + (num_tokens, top_k), dtype=torch.float32, device=x.device + ) + topk_indices = torch.empty((num_tokens, top_k), dtype=torch.int32, device=x.device) + return topk_weights, topk_indices # Cache whether aiter supports FP8 MLA parameters @@ -1918,7 +1922,7 @@ class rocm_aiter_ops: @staticmethod def shuffle_weight( - self, tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) + tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) ) -> torch.Tensor: from aiter.ops.shuffle import shuffle_weight