Commit fbc0abba authored by zhuwenwen's avatar zhuwenwen
Browse files

fix moe error

parent 029da5e8
...@@ -4,7 +4,7 @@ ...@@ -4,7 +4,7 @@
import functools import functools
import json import json
import os import os
from typing import Any, Callable, Optional, List, Optional, Tuple from typing import Any, Callable, Optional, List, Optional, Tuple, Dict
import torch import torch
...@@ -502,7 +502,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor, ...@@ -502,7 +502,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
use_int8_w8a16: bool, use_int8_w8a16: bool,
use_int4_w4a16: bool, use_int4_w4a16: bool,
per_channel_quant: bool, per_channel_quant: bool,
block_shape: Optional[List[int],] = None, block_shape: Optional[List[int]] = None,
use_nn_moe: Optional[bool]=False) -> None: use_nn_moe: Optional[bool]=False) -> None:
assert topk_weights is not None or not mul_routed_weight assert topk_weights is not None or not mul_routed_weight
assert topk_weights is None or topk_weights.stride(1) == 1 assert topk_weights is None or topk_weights.stride(1) == 1
...@@ -679,7 +679,7 @@ def get_moe_configs( ...@@ -679,7 +679,7 @@ def get_moe_configs(
block_n: Optional[int] = None, block_n: Optional[int] = None,
block_k: Optional[int] = None, block_k: Optional[int] = None,
use_nn_moe: Optional[bool] = False, use_nn_moe: Optional[bool] = False,
) -> Optional[dict[int, Any]]: ) -> Optional[Dict[int, Any]]:
""" """
Return optimized configurations for the fused MoE kernel. Return optimized configurations for the fused MoE kernel.
......
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