Commit fbc0abba authored by zhuwenwen's avatar zhuwenwen
Browse files

fix moe error

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