Commit 10184690 authored by zhuwenwen's avatar zhuwenwen
Browse files

fix moe run error

parent 675d3e87
...@@ -5,7 +5,7 @@ import functools ...@@ -5,7 +5,7 @@ import functools
import json import json
import os import os
import math import math
from typing import Any, Callable, Optional, List, Optional, Tuple from typing import Any, Callable, Dict, Optional, List, Optional, Tuple
import torch import torch
......
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Optional from typing import Optional, Tuple
import torch import torch
...@@ -153,9 +153,9 @@ def moe_align_block_size( ...@@ -153,9 +153,9 @@ def moe_align_block_size(
block_size: int, block_size: int,
num_experts: int, num_experts: int,
expert_map: Optional[torch.Tensor] = None, expert_map: Optional[torch.Tensor] = None,
pad_sorted_ids: bool = False pad_sorted_ids: bool = False,
num_token: Optional[int] = None num_token: Optional[int] = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
""" """
Aligns the token distribution across experts to be compatible with block Aligns the token distribution across experts to be compatible with block
size for matrix multiplication. size for matrix multiplication.
...@@ -253,4 +253,4 @@ def moe_align_block_size( ...@@ -253,4 +253,4 @@ def moe_align_block_size(
if expert_map is not None: if expert_map is not None:
expert_ids = expert_map[expert_ids] expert_ids = expert_map[expert_ids]
return sorted_ids, expert_ids, num_tokens_post_pad return sorted_ids, expert_ids, num_tokens_post_pad
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment