add ops
This commit is contained in:
Chranos
627
torch_mlu_ops-v1.3.2/torch_mlu_ops/abstract.py
Normal file
627
torch_mlu_ops-v1.3.2/torch_mlu_ops/abstract.py
Normal file
@@ -0,0 +1,627 @@
|
||||
from typing import Tuple, List, Dict, Optional
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import torch._custom_ops
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::attention_project")
|
||||
def attention_project_abstract(
|
||||
input: Tensor,
|
||||
q_weight: Tensor,
|
||||
q_bias: Optional[Tensor],
|
||||
k_weight: Optional[Tensor],
|
||||
k_bias: Optional[Tensor],
|
||||
v_weight: Optional[Tensor],
|
||||
v_bias: Optional[Tensor],
|
||||
norm_weight: Optional[Tensor],
|
||||
norm_bias: Optional[Tensor],
|
||||
residual: Optional[Tensor],
|
||||
out_layout: str,
|
||||
head_size: int,
|
||||
eps: float,
|
||||
alpha: float,
|
||||
beta: float,
|
||||
norm_out: bool,
|
||||
) -> List[Tensor]:
|
||||
input_view = input
|
||||
if input.dim() == 2:
|
||||
input_view = input.unsqueeze(0)
|
||||
n = input_view.size(0)
|
||||
t = input_view.size(1)
|
||||
hidden_size_q = q_weight.size(0)
|
||||
hidden_size_k = k_weight.size(0) if k_weight is not None else 0
|
||||
hidden_size_v = v_weight.size(0) if v_weight is not None else 0
|
||||
head_num_q = hidden_size_q // head_size
|
||||
head_num_k = hidden_size_k // head_size
|
||||
head_num_v = hidden_size_v // head_size
|
||||
|
||||
out_q = torch.empty(n, t, hidden_size_q, dtype=input_view.dtype, device=input_view.device)
|
||||
out_k = torch.empty(n, t, hidden_size_k, dtype=input_view.dtype, device=input_view.device) if hidden_size_k > 0 else None
|
||||
out_v = torch.empty(n, t, hidden_size_v, dtype=input_view.dtype, device=input_view.device) if hidden_size_v > 0 else None
|
||||
|
||||
if out_layout == "nhtc":
|
||||
out_q = torch.empty(n, head_num_q, t, head_size, dtype=input_view.dtype, device=input_view.device)
|
||||
out_k = torch.empty(n, head_num_k, t, head_size, dtype=input_view.dtype, device=input_view.device) if hidden_size_k > 0 else None
|
||||
out_v = torch.empty(n, head_num_v, t, head_size, dtype=input_view.dtype, device=input_view.device) if hidden_size_v > 0 else None
|
||||
|
||||
out_ln = torch.empty_like(input_view) if norm_out else None
|
||||
res = [out_q]
|
||||
if k_weight is not None:
|
||||
res.append(out_k)
|
||||
if v_weight is not None:
|
||||
res.append(out_v)
|
||||
if norm_out:
|
||||
res.append(out_ln)
|
||||
|
||||
return res
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::ffn")
|
||||
def ffn_abstract(
|
||||
input: Tensor,
|
||||
up_fc_weight: Tensor,
|
||||
up_fc_bias: Optional[Tensor],
|
||||
down_proj_weight: Tensor,
|
||||
down_proj_bias: Optional[Tensor],
|
||||
gate_up_proj_weight: Optional[Tensor],
|
||||
gate_up_proj_bias: Optional[Tensor],
|
||||
layernorm_weight: Optional[Tensor],
|
||||
layernorm_bias: Optional[Tensor],
|
||||
act_mode: str,
|
||||
residual_is: str,
|
||||
eps: float,
|
||||
alpha: float,
|
||||
beta: float,
|
||||
) -> Tensor:
|
||||
return torch.empty_like(input)
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::flash_attention")
|
||||
def flash_attention_abstract(
|
||||
q: Tensor,
|
||||
k: Tensor,
|
||||
v: Tensor,
|
||||
out: Tensor,
|
||||
output_lse: Optional[Tensor],
|
||||
cu_seq_lens_q: Optional[Tensor],
|
||||
cu_seq_lens_kv: Optional[Tensor],
|
||||
alibi_slope: Optional[Tensor],
|
||||
attn_bias: Optional[Tensor],
|
||||
k_quant_scale: Optional[Tensor],
|
||||
v_quant_scale: Optional[Tensor],
|
||||
block_tables: Optional[Tensor],
|
||||
max_seq_len_q: int,
|
||||
max_seq_len_kv: int,
|
||||
softmax_scale: float,
|
||||
is_causal: bool,
|
||||
window_size_left: int,
|
||||
window_size_right: int,
|
||||
compute_dtype: str,
|
||||
return_lse: bool,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::single_query_cached_kv_attn")
|
||||
def single_query_cached_kv_attn_abstract(
|
||||
q_ori: Tensor,
|
||||
k_cache: Tensor,
|
||||
v_cache: Tensor,
|
||||
output: Tensor,
|
||||
block_tables: Tensor,
|
||||
context_lens: Tensor,
|
||||
output_lse: Optional[Tensor],
|
||||
k_cache_quant_scale: Optional[Tensor],
|
||||
v_cache_quant_scale: Optional[Tensor],
|
||||
alibi_slopes: Optional[Tensor],
|
||||
max_contxt_len: int,
|
||||
windows_size_left: int,
|
||||
windows_size_right: int,
|
||||
softmax_scale: float,
|
||||
return_lse: bool,
|
||||
kv_cache_quant_bit_size: int
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::apply_rotary")
|
||||
def apply_rotary_abstract(
|
||||
input: Tensor,
|
||||
sin_cache: Tensor,
|
||||
cos_cache: Tensor,
|
||||
position_ids: Optional[Tensor],
|
||||
cu_seqlens: Optional[Tensor],
|
||||
interleaved: bool,
|
||||
discrete: bool,
|
||||
dynamic_ntk: bool,
|
||||
max_seqlen: int,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::reshape_linear_cache")
|
||||
def reshape_linear_cache_abstract(
|
||||
key: Tensor,
|
||||
value: Optional[Tensor],
|
||||
key_cache: Tensor,
|
||||
value_cache: Optional[Tensor],
|
||||
context_lengths: Tensor,
|
||||
max_context_len: int,
|
||||
packed: bool,
|
||||
context_seq_offset: Optional[Tensor],
|
||||
cache_bs_id: Optional[Tensor],
|
||||
cache_seqlen_offset: Optional[Tensor],
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::reshape_paged_cache")
|
||||
def reshape_paged_cache_abstract(
|
||||
k: Tensor,
|
||||
v: Optional[Tensor],
|
||||
k_cache: Tensor,
|
||||
v_cache: Optional[Tensor],
|
||||
slot_mapping: Tensor
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::quant_to_paged_cache")
|
||||
def quant_to_paged_cache_abstract(
|
||||
k: Tensor,
|
||||
v: Optional[Tensor],
|
||||
k_cache: Tensor,
|
||||
v_cache: Optional[Tensor],
|
||||
k_cache_scale: Tensor,
|
||||
v_cache_scale: Optional[Tensor],
|
||||
slot_mapping: Tensor,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::offline_quant_to_paged_cache")
|
||||
def offline_quant_to_paged_cache_abstract(
|
||||
k: Tensor,
|
||||
v: Optional[Tensor],
|
||||
k_cache_scale: Tensor,
|
||||
v_cache_scale: Optional[Tensor],
|
||||
slot_mapping: Tensor,
|
||||
k_cache: Tensor,
|
||||
v_cache: Optional[Tensor],
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::quant_to_linear_cache")
|
||||
def quant_to_linear_cache_abstract(
|
||||
key: Tensor,
|
||||
value: Optional[Tensor],
|
||||
key_cache: Tensor,
|
||||
value_cache: Optional[Tensor],
|
||||
key_cache_scale: Tensor,
|
||||
value_cache_scale: Optional[Tensor],
|
||||
context_lengths: Tensor,
|
||||
max_context_len: int,
|
||||
packed: bool,
|
||||
context_seq_offset: Optional[Tensor],
|
||||
cache_bs_id: Optional[Tensor],
|
||||
cache_seqlen_offset: Optional[Tensor],
|
||||
quant_bit: int = 8,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::offline_quant_to_linear_cache")
|
||||
def offline_quant_to_linear_cache_abstract(
|
||||
key: Tensor,
|
||||
value: Optional[Tensor],
|
||||
key_cache: Tensor,
|
||||
value_cache: Optional[Tensor],
|
||||
key_cache_scale: Tensor,
|
||||
value_cache_scale: Optional[Tensor],
|
||||
context_lengths: Tensor,
|
||||
max_context_len: int,
|
||||
quant_mode: int,
|
||||
packed: bool,
|
||||
context_seq_offset: Optional[Tensor],
|
||||
cache_bs_id: Optional[Tensor],
|
||||
cache_seqlen_offset: Optional[Tensor],
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::swap_blocks")
|
||||
def swap_blocks_abstract(
|
||||
dst: Tensor, src: Tensor, block_mapping: Dict[int, int]
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::copy_blocks")
|
||||
def copy_blocks_abstract(
|
||||
k_caches: List[Tensor], v_caches: List[Tensor], block_mapping: Dict[int, List[int]]
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::copy_blocks_out_of_place")
|
||||
def copy_blocks_out_of_place_abstract(k_caches: List[Tensor], v_caches: List[Tensor], block_mapping: Dict[int, List[int]]) -> (List[Tensor], List[Tensor]):
|
||||
return ([torch.empty_like(k) for k in k_caches], [torch.empty_like(v) for v in v_caches])
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::quant_matmul")
|
||||
def quant_matmul_abstract(
|
||||
a_tensor: Tensor,
|
||||
a_scale: Optional[Tensor],
|
||||
a_zero: Optional[Tensor],
|
||||
b_tensor: Tensor,
|
||||
b_scale: Optional[Tensor],
|
||||
b_zero: Optional[Tensor],
|
||||
bias: Optional[Tensor],
|
||||
c_tensor: Optional[Tensor],
|
||||
c_scale: Optional[Tensor],
|
||||
c_zero: Optional[Tensor],
|
||||
gemm_output_scale: Optional[Tensor],
|
||||
gemm_output_zero: Optional[Tensor],
|
||||
data_type: Optional[str],
|
||||
d: Optional[Tensor],
|
||||
quant_algo: str,
|
||||
a_quant_layout: str,
|
||||
b_quant_layout: str,
|
||||
quant_bit_size: int = 8,
|
||||
act_mode: str = "none",
|
||||
use_hp_active: bool = False,
|
||||
act_coef: float = 1.0,
|
||||
alpha: float = 1.0,
|
||||
beta: float = 1.0,
|
||||
trans_a: bool = False,
|
||||
trans_b: bool = True,
|
||||
) -> Tensor:
|
||||
|
||||
if data_type is None:
|
||||
output_type = a_tensor.dtype
|
||||
elif data_type == "float":
|
||||
output_type = torch.float32
|
||||
elif data_type == "bfloat16":
|
||||
output_type = torch.bfloat16
|
||||
else:
|
||||
output_type = torch.float16
|
||||
return torch.empty(a_tensor.size(0), b_tensor.size(0), dtype=output_type, device=a_tensor.device)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::quant_matmul_allreduce")
|
||||
def quant_matmul_allreduce_abstract(
|
||||
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],
|
||||
d: Optional[torch.Tensor],
|
||||
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:
|
||||
output_type = torch.float16
|
||||
if data_type == "float":
|
||||
output_type = torch.float32
|
||||
elif data_type == "bfloat16":
|
||||
output_type = torch.bfloat16
|
||||
return torch.empty(a_tensor.size(0), b_tensor.size(0), dtype=output_type, device=a_tensor.device)
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::active")
|
||||
def active_abstract(input: Tensor,
|
||||
output: Tensor,
|
||||
bias: Optional[Tensor],
|
||||
cusum_token_count: Optional[Tensor],
|
||||
act_mode: str,
|
||||
is_gated: bool,
|
||||
start_expert_id: int = 0,
|
||||
expert_size: int = 0,
|
||||
active_coef: float = 1.0) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::smooth_quant")
|
||||
def smooth_quant_abstract(
|
||||
input: Tensor,
|
||||
input_scale: Tensor,
|
||||
output: Tensor,
|
||||
output_scale: Tensor,
|
||||
input_zero: Optional[Tensor],
|
||||
token_count: Optional[Tensor],
|
||||
gather_index: Optional[Tensor],
|
||||
gather_index_start_position: Optional[Tensor],
|
||||
quant_mode: str,
|
||||
dynamic_quant: bool
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::fused_layernorm")
|
||||
def fused_layernorm_abstract(
|
||||
input: Tensor,
|
||||
output: Tensor,
|
||||
residual: Optional[Tensor],
|
||||
beta: Optional[Tensor],
|
||||
gamma: Optional[Tensor],
|
||||
bias: Optional[Tensor],
|
||||
quant_scale: Optional[Tensor],
|
||||
residual_out: Optional[Tensor],
|
||||
smooth_quant_scale: Optional[Tensor],
|
||||
norm_mode: str,
|
||||
eps: float,
|
||||
store_output_before_norm: bool,
|
||||
dynamic_quant: bool,
|
||||
)-> None:
|
||||
return None
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::fused_moe")
|
||||
def fused_moe_abstract(
|
||||
hidden_states: Tensor,
|
||||
gating_output: Tensor,
|
||||
w1: Tensor,
|
||||
w2: Tensor,
|
||||
bias1: Optional[Tensor],
|
||||
bias2: Optional[Tensor],
|
||||
residual: Optional[Tensor],
|
||||
input_smooth: Optional[Tensor],
|
||||
act_smooth: Optional[Tensor],
|
||||
w1_scale: Optional[Tensor],
|
||||
w2_scale: Optional[Tensor],
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
gated: bool,
|
||||
act_mode: str,
|
||||
start_expert_id: int,
|
||||
block_n: int,
|
||||
cncl_comm: int,
|
||||
w1_quant_flag: Optional[List],
|
||||
w2_quant_flag: Optional[List]
|
||||
) -> Tensor:
|
||||
return torch.empty_like(hidden_states)
|
||||
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::matmul")
|
||||
def matmul_abstract(
|
||||
a: Tensor,
|
||||
b: Tensor,
|
||||
d: Optional[Tensor],
|
||||
bias: Optional[Tensor],
|
||||
c: Optional[Tensor],
|
||||
data_type: Optional[str],
|
||||
act_mode: str,
|
||||
alpha: float,
|
||||
beta: float,
|
||||
fast_act: bool,
|
||||
approximate: bool,
|
||||
a_scale: float,
|
||||
b_scale: float,
|
||||
trans_a: bool,
|
||||
trans_b: bool
|
||||
) -> Tensor:
|
||||
m = a.size(1) if trans_a else a.size(0)
|
||||
n = b.size(0) if trans_b else b.size(1)
|
||||
if data_type is None:
|
||||
output_type = a.dtype
|
||||
elif data_type == "float":
|
||||
output_type = torch.float32
|
||||
elif data_type == "bfloat16":
|
||||
output_type = torch.bfloat16
|
||||
else:
|
||||
output_type = torch.half
|
||||
return torch.empty(m, n, dtype=output_type, device=a.device)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::batch_matmul")
|
||||
def batch_matmul_abstract(
|
||||
a: Tensor,
|
||||
b: Tensor,
|
||||
c: Tensor,
|
||||
alpha: float,
|
||||
beta: float,
|
||||
a_scale: float,
|
||||
b_scale: float,
|
||||
trans_a: bool,
|
||||
trans_b: bool
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::matmul_allreduce")
|
||||
def matmul_allreduce_abstract(
|
||||
cncl_comm,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
c: Optional[torch.Tensor] = None,
|
||||
d: Optional[torch.Tensor] = None,
|
||||
alpha: float = 1.0,
|
||||
beta: float = .0,
|
||||
block_m: int = 0
|
||||
) -> Tensor:
|
||||
return torch.empty(a.size(0), b.size(0), dtype=a.dtype, device=a.device)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::group_gemm")
|
||||
def group_gemm_abstract(
|
||||
a: Tensor,
|
||||
b: Tensor,
|
||||
m_list: Tensor,
|
||||
expand_idx: Optional[Tensor],
|
||||
c: Optional[Tensor],
|
||||
alpha: Optional[Tensor],
|
||||
beta: Optional[Tensor],
|
||||
a_scale: Optional[Tensor],
|
||||
b_scale: Optional[Tensor],
|
||||
bias: Optional[Tensor],
|
||||
data_type: Optional[str],
|
||||
quant_flag: Optional[List],
|
||||
b_offset: Optional[Tensor],
|
||||
max_m: int
|
||||
) -> Tensor:
|
||||
if data_type is None:
|
||||
output_type = a.dtype
|
||||
elif data_type == "float":
|
||||
output_type = torch.float32
|
||||
elif data_type == "bfloat16":
|
||||
output_type = torch.bfloat16
|
||||
else:
|
||||
output_type = torch.half
|
||||
total_m = a.size(0) if expand_idx is None else expand_idx.size(0)
|
||||
|
||||
return torch.empty(total_m, b.size(1), dtype=output_type, device=a.device)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::preload")
|
||||
def preload_abstract(
|
||||
weight: Tensor,
|
||||
size: int
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::flash_attn_sq_mm_allreduce")
|
||||
def flash_attn_sq_mm_allreduce_abstract(cncl_comm,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seq_lens_q,
|
||||
cu_seq_lens_k,
|
||||
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) -> torch.Tensor:
|
||||
res_q = q.unsqueeze(0) if cu_seq_lens_q is not None else q
|
||||
res_q = res_q.flatten(-2, -1).flatten(0, 1)
|
||||
return torch.empty(res_q.size(0), weight.size(0), dtype=q.dtype, device=q.device)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::moe_softmax_topk")
|
||||
def moe_softmax_topk_abstract(input,
|
||||
topk,
|
||||
num_expert_group,
|
||||
topk_group,
|
||||
normalize,
|
||||
mask: Optional[torch.Tensor] = None,
|
||||
normed_by: str = "topk_logit") -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
out_shape = list(input.size())[:-1] + [topk]
|
||||
reduce_weight = torch.empty(out_shape, dtype=torch.float32, device=input.device)
|
||||
expert_id = torch.empty(out_shape, dtype=torch.int, device=input.device)
|
||||
return (reduce_weight, expert_id)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::moe_expand_input")
|
||||
def moe_expand_input_abstract(input: Tensor,
|
||||
gather_idx: Tensor,
|
||||
cusum_token_count: Optional[Tensor] = None,
|
||||
start_expert_id: int = 0,
|
||||
expert_size: int = 0) -> Tensor:
|
||||
return torch.empty(gather_idx.size(0), input.size(-1), dtype=input.dtype, device=input.device)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::moe_gen_idx")
|
||||
def moe_gen_idx_abstract(expert_id: Tensor,
|
||||
expert_num: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
token_num, topk = expert_id.size(0), expert_id.size(1)
|
||||
expand_idx = torch.empty((token_num * topk), dtype=torch.int32, device=expert_id.device)
|
||||
combine_idx = torch.empty((token_num * topk), dtype=torch.int32, device=expert_id.device)
|
||||
token_count = torch.empty((expert_num,), dtype=torch.int32, device=expert_id.device)
|
||||
cusum_token_count = torch.empty((expert_num + 1,), dtype=torch.int32, device=expert_id.device)
|
||||
return (expand_idx, combine_idx, token_count, cusum_token_count)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::moe_combine_result")
|
||||
def moe_combine_result_abstract(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:
|
||||
num_tokens, hidden_size, topk = input.size(0), input.size(1), reduce_weight.size(1)
|
||||
num_token = num_tokens // topk
|
||||
return torch.empty(num_token, hidden_size, dtype=input.dtype, device=input.device)
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::fused_rope")
|
||||
def fused_rope_abstract(qkv: torch.Tensor,
|
||||
key_cache_hp: torch.Tensor,
|
||||
value_cache_hp: torch.Tensor,
|
||||
key_cache_lp: Optional[torch.Tensor],
|
||||
value_cache_lp: Optional[torch.Tensor],
|
||||
sin_table: torch.Tensor,
|
||||
cos_table: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
gamma: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
key_scale_hp: Optional[torch.Tensor],
|
||||
value_scale_hp: Optional[torch.Tensor],
|
||||
key_scale_lp: Optional[torch.Tensor],
|
||||
value_scale_lp: Optional[torch.Tensor],
|
||||
cache_bs_id_hp: Optional[torch.Tensor],
|
||||
cache_seq_offsets_hp: Optional[torch.Tensor],
|
||||
cache_bs_id_lp: Optional[torch.Tensor],
|
||||
cache_seq_offsets_lp: Optional[torch.Tensor],
|
||||
slot_mapping_hp: Optional[torch.Tensor],
|
||||
slot_mapping_lp: Optional[torch.Tensor],
|
||||
eps: float) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::moe_cast_gating")
|
||||
def moe_cast_gating_abstract(input: torch.Tensor,
|
||||
weight: torch.Tensor) ->Tensor:
|
||||
output_shape = input.shape[:-1] + (weight.shape[0],)
|
||||
output = torch.empty(output_shape, dtype=torch.float, device="mlu")
|
||||
return output
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::update_out_and_lse")
|
||||
def update_out_and_lse_abstract(out: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
block_out: torch.Tensor,
|
||||
block_lse: torch.Tensor,
|
||||
seq_offsets: Optional[torch.Tensor] = None,
|
||||
cu_seqs: Optional[torch.Tensor] = None,
|
||||
block_cu_seqs: Optional[torch.Tensor] = None) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::dequant_from_linear_cache")
|
||||
def dequant_from_linear_cache_abstract(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,
|
||||
context_seq_offset: Optional[torch.Tensor],
|
||||
cache_bs_id: Optional[torch.Tensor],
|
||||
cache_seq_offset: Optional[torch.Tensor],
|
||||
quant_mode: int = 0,
|
||||
quant_bit: int = 8) -> None:
|
||||
return None
|
||||
|
||||
@torch._custom_ops.impl_abstract("torch_mlu_ops::dequant_from_paged_cache")
|
||||
def dequant_from_paged_cache_abstract(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,
|
||||
context_seq_offset: Optional[torch.Tensor],
|
||||
block_tables: torch.Tensor,
|
||||
quant_mode: int = 0,
|
||||
quant_bit: int = 8) -> None:
|
||||
return None
|
||||
Reference in New Issue
Block a user