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