Skip to content

Commit 4bd18ec

Browse files
authored
[Minor] Fix type annotation in fused moe (#3045)
1 parent 2410e32 commit 4bd18ec

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

vllm/model_executor/layers/fused_moe/fused_moe.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import functools
33
import json
44
import os
5-
from typing import Any, Dict, Optional
5+
from typing import Any, Dict, Optional, Tuple
66

77
import torch
88
import triton
@@ -137,7 +137,7 @@ def fused_moe_kernel(
137137

138138
def moe_align_block_size(
139139
topk_ids: torch.Tensor, block_size: int,
140-
num_experts: int) -> (torch.Tensor, torch.Tensor, torch.Tensor):
140+
num_experts: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
141141
"""
142142
Aligns the token distribution across experts to be compatible with block size for matrix multiplication.
143143
@@ -185,7 +185,8 @@ def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
185185
sorted_token_ids: torch.Tensor,
186186
expert_ids: torch.Tensor,
187187
num_tokens_post_padded: torch.Tensor,
188-
mul_routed_weight: bool, top_k: int, config: dict):
188+
mul_routed_weight: bool, top_k: int,
189+
config: Dict[str, Any]) -> None:
189190
assert topk_weights.stride(1) == 1
190191
assert sorted_token_ids.stride(0) == 1
191192

0 commit comments

Comments
 (0)