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 5, 08:30 UTC
[ROCm] Fix AITER ops fake impl and minor bugs (#36092)
Signed-off-by: Li <chuali@amd.com>
Chuan (Richard) Li committed
Apr 9, 2026 at 17:56 UTC
e0613702ade9ace874feabb7b6f080cdfd181f4b
1 file changed
+8
-4
vllm/_aiter_ops.py
+8
-4
@@ -336,9 +336,13 @@ def _rocm_aiter_fused_topk_fake(
336
router_logits: torch.Tensor,
337
top_k: int,
338
gate_up: bool,
339
-) -> None:
340
- # tuple[torch.Tensor, torch.Tensor]:
341
- pass
339
+) -> tuple[torch.Tensor, torch.Tensor]:
340
+ num_tokens = x.shape[0]
341
+ topk_weights = torch.empty(
342
+ (num_tokens, top_k), dtype=torch.float32, device=x.device
343
+ )
344
+ topk_indices = torch.empty((num_tokens, top_k), dtype=torch.int32, device=x.device)
345
+ return topk_weights, topk_indices
346
347
348
# Cache whether aiter supports FP8 MLA parameters
@@ -1918,7 +1922,7 @@ class rocm_aiter_ops:
1922
1923
@staticmethod
1924
def shuffle_weight(
1921
- self, tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
1925
+ tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
1926
) -> torch.Tensor:
1927
from aiter.ops.shuffle import shuffle_weight
1928