from typing import List, Optional, Tuple import torch import math import triton import triton.language as tl from vllm.logger import init_logger logger = init_logger(__name__) try: import torch_mlu_ops as tmo except ImportError as e: logger.warning("Failed to import from TMO OPS with %r", e) def rotary_embedding( input: torch.Tensor, sin_cache: torch.Tensor, cos_cache: torch.Tensor, position_ids: Optional[torch.Tensor], cu_seqlens: Optional[torch.Tensor], interleaved: bool, discrete: bool, dynamic_ntk: bool, max_seqlen: int, ) -> torch.Tensor: return tmo.apply_rotary( input, sin_cache, cos_cache, position_ids, cu_seqlens, interleaved, discrete, dynamic_ntk, max_seqlen) def fused_rms_norm( x: torch.Tensor, residual: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, bias: torch.Tensor, eps: float, store_output_before_norm: bool, quant_scale: torch.Tensor = None, dynamic_quant: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: return tmo.fused_rms_norm( x, residual, gamma, beta, bias, eps, store_output_before_norm, quant_scale, None, dynamic_quant) def fused_layer_norm( x: torch.Tensor, residual: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, bias: torch.Tensor, eps: float, store_output_before_norm: bool, quant_scale: torch.Tensor = None, dynamic_quant: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: return tmo.fused_layer_norm( x, residual, gamma, beta, bias, eps, store_output_before_norm, quant_scale, None, dynamic_quant) def flash_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor, cu_seq_lens_q: torch.Tensor, cu_seq_lens_kv: torch.Tensor, alibi_slope: torch.Tensor, attn_bias: torch.Tensor, max_seq_len_q: int, max_seq_len_kv: int, softmax_scale: float, is_causal: bool, window_size_left: int = -1, window_size_right: int = -1, compute_dtype: torch.dtype = torch.float, return_lse: bool = False, block_tables: torch.Tensor = None, k_cache_quant_scale: torch.Tensor = None, v_cache_quant_scale: torch.Tensor = None ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: return tmo.flash_attention( q, k, v, out, cu_seq_lens_q, cu_seq_lens_kv, alibi_slope, attn_bias, max_seq_len_q, max_seq_len_kv, softmax_scale, is_causal, window_size_left, window_size_right, compute_dtype, return_lse, block_tables, k_cache_quant_scale, v_cache_quant_scale) def split_head_nums(q_head_num, kv_head_num, max_q_head_num): """ 对 q_head_num 进行切分，使得： 1. 切分后的 q_head_num 最大值不超过 max_q_head_num 2. kv_head_num 按 q_head_num 相同份数拆分； 3. 每个切分后的 q_head_num 可以被对应的 kv_head_num 整除； 4. 若 kv_head_num < 1，则调整为 1。 参数： - q_head_num: int, 需要切分的 q_head_num。 - kv_head_num: int, 需要切分的 kv_head_num。 - max_q_head_num: int, 支持切分后最大的 q_head_num 返回： - q_splits: list, 切分后的 q_head_num。 - kv_splits: list, 切分后的 kv_head_num。 """ if q_head_num <= 0 or kv_head_num <= 0: return "q_head_num 和 kv_head_num 必须是正整数！" q_splits = [] kv_splits = [] # 剩余值 remaining_q = q_head_num remaining_kv = kv_head_num while remaining_q > 0: # 尝试切分 q_head_num，最大值不超过 max_q_head_num for q_part in range(min(max_q_head_num, remaining_q), 0, -1): # 确保 q_part 能被分配并且对应的 kv_part >= 1 if remaining_q % q_part == 0: kv_part = max(remaining_kv // (remaining_q // q_part), 1) # 确保 kv_part >= 1 if q_part % kv_part == 0: # 确保 q_part 可以被 kv_part 整除 # 记录切分值 q_splits.append(q_part) kv_splits.append(kv_part) remaining_q -= q_part remaining_kv -= kv_part break else: err_msg = f"Unable to find split method for q_head_num:{q_head_num}, kv_head_num:{kv_head_num}" raise RuntimeError(err_msg) return q_splits, kv_splits def repeat_elements(input_list, n): """ 将列表的每个成员连续重复 n 次。 参数： - input_list: list，输入的列表。 - n: int，每个元素需要重复的次数。 返回： - list，包含重复元素的新列表。 """ if not isinstance(input_list, list) or not isinstance(n, int) or n < 0: raise ValueError("输入必须是一个列表，并且重复次数 n 必须是大于或等于 0 的整数。") # 使用列表推导式重复每个元素 n 次 return [item for item in input_list for _ in range(n)] def single_query_cached_kv_attn( q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, out: torch.Tensor, block_tables: torch.Tensor, context_lens: torch.Tensor, k_cache_quant_scale: Optional[torch.Tensor], v_cache_quant_scale: Optional[torch.Tensor], alibi_slopes: Optional[torch.Tensor], max_contxt_len: int, windows_size_left: int, windows_size_right: int, softmax_scale: float, q_head_dim: Optional[int] = 2, kv_head_dim: Optional[int] = 1, seq_q_dim: Optional[int] = 1, max_seq_q_mul_q_divide_kv: Optional[int] = 48, ) -> None: # FIXME(chenxiaobing): TMO only support windows_size_right = -1 yet. windows_size_right = -1 # singleQwithkvCache limits seq_q * q_divide_kv <= max_seq_q_mul_q_divide_kv now. # When the limitation is fixed, we should delete the split process. seq_q = q.shape[seq_q_dim] q_head_num = q.shape[q_head_dim] kv_head_num = k_cache.shape[kv_head_dim] q_divide_kv = q_head_num // kv_head_num if seq_q * q_divide_kv <= max_seq_q_mul_q_divide_kv: tmo.single_query_cached_kv_attn( q, k_cache, v_cache, out, block_tables, context_lens, k_cache_quant_scale, v_cache_quant_scale, alibi_slopes, max_contxt_len, windows_size_left, windows_size_right, softmax_scale) else: max_q_head_num = max_seq_q_mul_q_divide_kv * kv_head_num // seq_q q_head_num_sizes, kv_head_num_sizes = split_head_nums(q_head_num, kv_head_num, max_q_head_num) parts_num = len(q_head_num_sizes) q_parts = torch.split(q, q_head_num_sizes, dim=q_head_dim) out_parts = torch.split(out, q_head_num_sizes, dim=q_head_dim) alibi_slopes_parts = [None] * parts_num if alibi_slopes: alibi_slopes_parts = torch.split(alibi_slopes, q_head_num_sizes, dim=0) kv_parts_num = parts_num if parts_num > kv_head_num: assert parts_num % kv_head_num == 0, f"parts_num:{parts_num} need by divided by kv_head_num:{kv_head_num} when parts_num > kv_head_num" kv_parts_num = kv_head_num kv_head_num_sizes = kv_head_num_sizes[:kv_parts_num] if len(kv_head_num_sizes) > 1: k_cache_parts = torch.split(k_cache, kv_head_num_sizes, dim=kv_head_dim) v_cache_parts = torch.split(v_cache, kv_head_num_sizes, dim=kv_head_dim) k_cache_quant_scale_parts = [None] * kv_parts_num v_cache_quant_scale_parts = [None] * kv_parts_num if k_cache_quant_scale: k_cache_quant_scale_dim = 1 if k_cache_quant_scale.dim() == 2 else kv_head_dim k_cache_quant_scale_parts = torch.split(k_cache_quant_scale, kv_head_num_sizes, dim=k_cache_quant_scale_dim) if v_cache_quant_scale: v_cache_quant_scale_dim = 1 if v_cache_quant_scale.dim() == 2 else kv_head_dim v_cache_quant_scale_parts = torch.split(v_cache_quant_scale, kv_head_num_sizes, dim=v_cache_quant_scale_dim) else: k_cache_parts = [k_cache] v_cache_parts = [v_cache] k_cache_quant_scale_parts = [k_cache_quant_scale] v_cache_quant_scale_parts = [v_cache_quant_scale] if parts_num > kv_parts_num: repeate_num = parts_num // kv_parts_num k_cache_parts = repeat_elements(k_cache_parts, repeate_num) v_cache_parts = repeat_elements(v_cache_parts, repeate_num) k_cache_quant_scale_parts = repeat_elements(k_cache_quant_scale_parts, repeate_num) v_cache_quant_scale_parts = repeat_elements(v_cache_quant_scale_parts, repeate_num) for q_value, k_cache_value, v_cache_value, out_value, k_cache_quant_scale_value, v_cache_quant_scale_value, alibi_slopes_value in zip( q_parts, k_cache_parts, v_cache_parts, out_parts, k_cache_quant_scale_parts, v_cache_quant_scale_parts, alibi_slopes_parts): tmo.single_query_cached_kv_attn( q_value, k_cache_value.contiguous(), v_cache_value.contiguous(), out_value, block_tables, context_lens, k_cache_quant_scale_value, v_cache_quant_scale_value, alibi_slopes_value, max_contxt_len, windows_size_left, windows_size_right, softmax_scale) def reshape_linear_cache( key: torch.Tensor, value: Optional[torch.Tensor], key_cache: torch.Tensor, value_cache: Optional[torch.Tensor], context_lengths: torch.Tensor, max_context_len: int, packed: bool, context_seq_offset: Optional[torch.Tensor], cache_bs_id: Optional[torch.Tensor], cache_seqlen_offset: Optional[torch.Tensor], ) -> None: tmo.reshape_linear_cache( key, value, key_cache, value_cache, context_lengths, max_context_len, packed, context_seq_offset, cache_bs_id, cache_seqlen_offset) def reshape_paged_cache( k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor ) -> None: tmo.reshape_paged_cache(k, v, k_cache, v_cache, slot_mapping) def swap_blocks( dst: torch.Tensor, src: torch.Tensor, block_mapping: torch.Tensor ) -> None: # FIXME: Remove this conversion after # tmo.swap_blocks support block_mapping tensor. block_mapping = block_mapping.tolist() block_mapping = {src: dst for src, dst in block_mapping} return tmo.swap_blocks(dst, src, block_mapping) def copy_blocks( k_caches: List[torch.Tensor], v_caches: List[torch.Tensor], block_mapping: torch.Tensor ) -> None: # FIXME: Remove this conversion after # tmo.swap_blocks support block_mapping tensor. block_mapping = block_mapping.tolist() result_dict = {} for row in block_mapping: key = row[0] values = row[1:] if key in result_dict: result_dict[key].extend(values) else: result_dict[key] = values return tmo.copy_blocks(k_caches, v_caches, result_dict) def ffn( input: torch.Tensor, up_fc_weight: torch.Tensor, up_fc_bias: Optional[torch.Tensor], down_proj_weight: torch.Tensor, down_proj_bias: Optional[torch.Tensor], gate_up_proj_weight: Optional[torch.Tensor] = None, gate_up_proj_bias: Optional[torch.Tensor] = None, act_mode: str = "none" ) -> torch.Tensor: return tmo.ffn(input, up_fc_weight, up_fc_bias, down_proj_weight, down_proj_bias, gate_up_proj_weight, gate_up_proj_bias, act_mode) def active( input: torch.Tensor, act_mode: str, is_gated: bool ) -> torch.Tensor: return tmo.active(input, act_mode, is_gated) def fused_moe( hidden_states: torch.Tensor, gating_output: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, bias1: Optional[torch.Tensor], bias2: Optional[torch.Tensor], residual: Optional[torch.Tensor], input_smooth: Optional[torch.Tensor], act_smooth: Optional[torch.Tensor], w1_scale: Optional[torch.Tensor], w2_scale: Optional[torch.Tensor], topk: int, renormalize: bool, gated: bool, act_mode: str, start_expert_id: int = 0, block_n: int = 0, cncl_comm: int = 0 ) -> torch.Tensor: return tmo.fused_moe( hidden_states, gating_output, w1, w2, bias1, bias2, residual, input_smooth, act_smooth, w1_scale, w2_scale, topk, renormalize, gated, act_mode, start_expert_id, block_n, cncl_comm) def matmul( a: torch.Tensor, b: torch.Tensor, bias: Optional[torch.Tensor] = None, c: Optional[torch.Tensor] = None, act_mode: str = 'none', alpha: float = 1.0, beta: float = .0 ) -> torch.Tensor: return tmo.matmul(a, b, bias, c, act_mode, alpha, beta) def weight_only_quant_matmul( a: torch.Tensor, b: torch.Tensor, scale: torch.Tensor, zero: torch.Tensor = None, bias: torch.Tensor = None, c: torch.Tensor = None, act_mode: str = "none", quant_bit_size: int = 8, alpha: float = 1.0, beta: float = 1.0 ) -> torch.Tensor: return tmo.weight_only_quant_matmul( a, b, scale, zero, bias, c, act_mode, quant_bit_size, alpha, beta) def smooth_quant_matmul( a: torch.Tensor, a_scale: torch.Tensor, b: torch.Tensor, b_scale: torch.Tensor, dtype: torch.dtype, bias: torch.Tensor = None, c: torch.Tensor = None, act_mode: str = "none", alpha: float = 1.0, beta: float = 1.0 ) -> torch.Tensor: return tmo.smooth_quant_matmul( a, a_scale, b, b_scale, dtype, bias, c, act_mode, alpha, beta) def per_token_smooth_quantize( x: torch.Tensor, smooth: torch.Tensor, zero: torch.Tensor = None, token_count: torch.Tensor = None ) -> Tuple[torch.Tensor, torch.Tensor]: return tmo.per_token_smooth_quantize(x, smooth, zero, token_count) def quantize( x: torch.Tensor, scale: torch.Tensor, zero: torch.Tensor = None ) -> torch.Tensor: return tmo.quantize(x, scale, zero) def quant_to_paged_cache( k: torch.Tensor, v: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, k_cache_quant_scale: torch.Tensor, v_cache_quant_scale: torch.Tensor, slot_mapping: torch.Tensor, ) -> None: return tmo.quant_to_paged_cache( k, v, k_cache, v_cache, k_cache_quant_scale, v_cache_quant_scale, slot_mapping ) def quant_to_linear_cache( key: torch.Tensor, value: Optional[torch.Tensor], key_cache: torch.Tensor, value_cache: Optional[torch.Tensor], key_cache_quant_scale: torch.Tensor, value_cache_quant_scale: Optional[torch.Tensor], context_lengths: torch.Tensor, max_context_len: int, packed: bool, context_seq_offset: Optional[torch.Tensor], cache_bs_id: Optional[torch.Tensor], cache_seqlen_offset: Optional[torch.Tensor], ) -> None: return tmo.quant_to_linear_cache( key, value, key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale, context_lengths, max_context_len, packed, context_seq_offset, cache_bs_id, cache_seqlen_offset, ) def advance_step(num_seqs: int, num_queries: int, block_size: int, input_tokens: torch.Tensor, sampled_token_ids: torch.Tensor, input_positions: torch.Tensor, seq_lens: torch.Tensor, slot_mapping: torch.Tensor, block_tables: torch.Tensor, TILE_SIZE: int = 64) -> None: """ Advance a step on MLU for existing inputs for a multi-step runner, which will update input_tokens/seq_lens/input_positions/slot_mapping inplace. """ def verify_tensor( name: str, tensor: torch.Tensor, size_0: int, size_1: int, dtype: torch.dtype, ): """ Auxiliary function to check whether input is valid. """ size_0_cond = (size_0 == -1 or tensor.size(0) == size_0) size_1_cond = (size_1 == -1 or tensor.size(1) == size_1) if not (size_0_cond and size_1_cond and tensor.is_contiguous and tensor.dtype == dtype): raise ValueError( f"The input to advance_step is invalid with tensor name = {name}, " f"shape = {tensor.shape}, " f"is_cont = {tensor.is_contiguous()}, " f"type = {tensor.dtype}, " f"is not as expected: shape[{size_0}, {size_1}], type = {dtype}" ) @triton.jit def _triton_advance_step(input_tokens_ptr, sampled_token_ids_ptr, input_positions_ptr, seq_lens_ptr, slot_mapping_ptr, block_tables_ptr, block_tables_stride, num_seqs, num_queries, block_size, TILE_SIZE: tl.constexpr, ): """ The triton implementation of advance step. Reference: https://github.com/vllm-project/vllm/blob/v0.6.1/csrc/prepare_inputs/advance_step.cu#L14-L55 """ # Set meta info. pid = tl.program_id(axis=0) offsets = pid * TILE_SIZE + tl.arange(0, TILE_SIZE) mask = offsets < num_queries # Update input_tokens. sampled_token_ids = tl.load(sampled_token_ids_ptr + offsets, mask=mask) tl.store(input_tokens_ptr + offsets, sampled_token_ids, mask=mask) seq_lens = tl.load(seq_lens_ptr + offsets, mask=mask) next_seq_lens = seq_lens + 1 next_input_pos = next_seq_lens - 1 # Update seq_lens. tl.store(seq_lens_ptr + offsets, next_seq_lens, mask=mask) # Update input_positions. tl.store(input_positions_ptr + offsets, next_input_pos, mask=mask) # Calculate slot num. block_index = next_input_pos // block_size block_offset = next_input_pos % block_size block_tables = tl.load(block_tables_ptr + block_tables_stride * offsets + block_index, mask=mask) slot_num = block_tables * block_size + block_offset # Update slot_mapping. tl.store(slot_mapping_ptr + offsets, slot_num, mask=mask) verify_tensor("input_tokens", input_tokens, num_seqs, -1, torch.int64) verify_tensor("sampled_token_ids", sampled_token_ids, num_queries, 1, torch.int64) verify_tensor("input_positions", input_positions, num_seqs, -1, torch.int32) verify_tensor("seq_lens", seq_lens, num_seqs, -1, torch.int32) verify_tensor("slot_mapping", slot_mapping, num_seqs, -1, torch.int32) verify_tensor("block_tables", block_tables, num_seqs, -1, torch.int32) grid = (math.ceil(num_queries / TILE_SIZE), ) _triton_advance_step[grid](input_tokens, sampled_token_ids, input_positions, seq_lens, slot_mapping, block_tables, block_tables.stride(0), num_seqs, num_queries, block_size, TILE_SIZE) def preload( weight: torch.Tensor, size: int ) -> None: """ Preload weights of layer. Args: weight (torch.Tensor): Weight to preload。 size (int): Preload size (byte)。 Returns: None """ return tmo.preload(weight, size) def matmul_allreduce( cncl_comm, a: torch.Tensor, b: torch.Tensor, bias: Optional[torch.Tensor] = None, c: Optional[torch.Tensor] = None, alpha: float = 1.0, beta: float = .0, block_m: int = 0 ) -> torch.Tensor: return tmo.matmul_allreduce(cncl_comm=cncl_comm, a=a, b=b, bias=bias, c=c, alpha=alpha, beta=beta, block_m=block_m) def smooth_quant_matmul_allreduce( cncl_comm, a: torch.Tensor, a_scale: torch.Tensor, b: torch.Tensor, b_scale: torch.Tensor, dtype: torch.dtype, bias: torch.Tensor = None, c: torch.Tensor = None, alpha: float = 1.0, beta: float = 1.0, block_m: int = 0): return tmo.smooth_quant_matmul_allreduce( cncl_comm=cncl_comm, a=a, a_scale=a_scale, b=b, b_scale=b_scale, dtype=dtype, bias=bias, c=c, alpha=alpha, beta=beta, block_m=block_m) def quant_matmul_allreduce( cncl_comm, a_tensor: torch.Tensor, a_scale: Optional[torch.Tensor], a_zero: Optional[torch.Tensor], b_tensor: torch.Tensor, b_scale: Optional[torch.Tensor], b_zero: Optional[torch.Tensor], bias: Optional[torch.Tensor], c_tensor: Optional[torch.Tensor], c_scale: Optional[torch.Tensor], c_zero: Optional[torch.Tensor], gemm_output_scale: Optional[torch.Tensor], gemm_output_zero: Optional[torch.Tensor], data_type: Optional[str], quant_algo: str, a_quant_layout: str, b_quant_layout: str, quant_bit_size: int = 8, alpha: float = 1.0, beta: float = 1.0, trans_a: bool = False, trans_b: bool = True, block_m: int = 0 ) -> torch.Tensor: return tmo.quant_matmul_allreduce( cncl_comm=cncl_comm, a_tensor=a_tensor, a_scale=a_scale, a_zero=a_zero, b_tensor=b_tensor, b_scale=b_scale, b_zero=b_zero, bias=bias, c_tensor=c_tensor, c_scale=c_scale, c_zero=c_zero, gemm_output_scale=gemm_output_scale, gemm_output_zero=gemm_output_zero, data_type=data_type, quant_algo=quant_algo, a_quant_layout=a_quant_layout, b_quant_layout=b_quant_layout, quant_bit_size=quant_bit_size, alpha=alpha, beta=beta, trans_a=trans_a, trans_b=trans_b, block_m=block_m) def flash_attn_sq_mm_allreduce( cncl_comm: int, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seq_lens_q: Optional[torch.Tensor], cu_seq_lens_kv: Optional[torch.Tensor], alibi_slope: Optional[torch.Tensor], attn_bias: Optional[torch.Tensor], smooth: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, bias: Optional[torch.Tensor], max_seq_len_q: int, max_seq_len_kv: int, softmax_scale: float, is_causal: bool, window_size_left: int = -1, window_size_right: int = -1, compute_dtype: torch.dtype = torch.float, block_seq: int = 0) -> torch.Tensor: return tmo.flash_attn_sq_mm_allreduce(cncl_comm, q, k, v, cu_seq_lens_q, cu_seq_lens_kv, alibi_slope, attn_bias, smooth, weight, weight_scale, bias, max_seq_len_q, max_seq_len_kv, softmax_scale, is_causal, window_size_left, window_size_right, compute_dtype, block_seq) #Moe inner kernels def moe_softmax_topk(input: torch.Tensor, topk: int, normalize: bool = False, num_expert_group: int = -1, topk_group: int = 0) -> Tuple[torch.Tensor]: return tmo.moe_softmax_topk(input, topk, normalize, num_expert_group, topk_group) def moe_gen_idx(expert_id: torch.Tensor, expert_num: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: return tmo.moe_gen_idx(expert_id, expert_num) def moe_expand_input(input: torch.Tensor, gather_idx: torch.Tensor, cusum_token_count: Optional[torch.Tensor] = None, start_expert_id: int = 0, expert_size: int = 0) -> torch.Tensor: return tmo.moe_expand_input(input, gather_idx, cusum_token_count, start_expert_id, expert_size) def moe_active(input: torch.Tensor, act_mode: str, is_gated: bool, output: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, cusum_token_count: Optional[torch.Tensor] = None, start_expert_id: int = 0, expert_size: int = 0) -> torch.Tensor: return tmo.moe_active(input, act_mode, is_gated, output, bias, cusum_token_count, start_expert_id, expert_size) def group_gemm(a: torch.Tensor, b: torch.Tensor, m_list: torch.Tensor, expand_idx: Optional[torch.Tensor], c: Optional[torch.Tensor], alpha: Optional[torch.Tensor], beta: Optional[torch.Tensor], max_m: int = 0 ) -> torch.Tensor: return tmo.group_gemm(a, b, m_list, expand_idx, c, alpha, beta, max_m) def smooth_quant_group_gemm(a: torch.Tensor, b: torch.Tensor, m_list: torch.Tensor, expand_idx: Optional[torch.Tensor], c: Optional[torch.Tensor], alpha: Optional[torch.Tensor], beta: Optional[torch.Tensor], a_scale: torch.Tensor, b_scale: torch.Tensor, dtype, max_m: int = 0 ) -> torch.Tensor: return tmo.smooth_quant_group_gemm(a, b, m_list, expand_idx, c, alpha, beta, a_scale, b_scale, dtype, max_m) def moe_combine_result(input: torch.Tensor, reduce_weight: torch.Tensor, gather_ids: torch.Tensor, residual: Optional[torch.Tensor], cusum_token_count: Optional[torch.Tensor], start_expert_id: int, expert_size: int, bias: Optional[torch.Tensor] = None) -> torch.Tensor: return tmo.moe_combine_result(input, reduce_weight, gather_ids, residual, cusum_token_count, start_expert_id, expert_size, bias) def moe_quantize(x: torch.Tensor, smooth: torch.Tensor, zero: Optional[torch.Tensor] = None, token_count: Optional[torch.Tensor] = None, gather_index: Optional[torch.Tensor] = None, gather_index_start_position: Optional[torch.Tensor] = None, output: Optional[torch.Tensor] = None, output_scale: Optional[torch.Tensor] = None, dynamic_quant: bool = True ) -> Tuple[torch.Tensor, torch.Tensor]: return tmo.moe_quantize(x, smooth, zero, token_count, gather_index, gather_index_start_position, output, output_scale, dynamic_quant)