wehub-resource-sync
237 lines
9.4 KiB
Python
237 lines
9.4 KiB
Python
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
class FakeLinearOp(torch.autograd.Function):
|
|
@staticmethod
|
|
def symbolic(g, input, in_features, out_features, has_bias, name):
|
|
# These become the operator attributes.
|
|
kwargs = {
|
|
"in_features_i": in_features,
|
|
"out_features_i": out_features,
|
|
"has_bias_i": has_bias,
|
|
"name_s": name
|
|
}
|
|
from torch.onnx.symbolic_helper import _get_tensor_sizes
|
|
out_sizes = _get_tensor_sizes(input)[:-1] + [out_features]
|
|
output_type = input.type().with_sizes(out_sizes)
|
|
return g.op("LlmExporter::FakeLinear", input, **kwargs).setType(output_type)
|
|
|
|
@staticmethod
|
|
def forward(ctx, input, in_features, out_features, has_bias, name):
|
|
out_shape = list(input.shape)[:-1] + [out_features]
|
|
return input.new_zeros(out_shape)
|
|
|
|
class FakeLinear(torch.nn.Module):
|
|
def __init__(self, in_features, out_features, has_bias, name):
|
|
super(FakeLinear, self).__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.has_bias = has_bias
|
|
self.name = name
|
|
|
|
def forward(self, x):
|
|
return FakeLinearOp.apply(x, self.in_features, self.out_features, self.has_bias, self.name)
|
|
|
|
class FusedAttentionOp(torch.autograd.Function):
|
|
@staticmethod
|
|
def symbolic(g, query, key, value, attention_mask, output_dim, kv_cache, name, layer_index, kv_shared_layer_index):
|
|
# These become the operator attributes.
|
|
kwargs = {
|
|
"output_dim_i": output_dim,
|
|
"kv_cache_i": kv_cache,
|
|
"name_s": name,
|
|
"layer_index_i": layer_index,
|
|
"kv_shared_layer_index_i": kv_shared_layer_index,
|
|
}
|
|
from torch.onnx.symbolic_helper import _get_tensor_sizes
|
|
out_sizes = _get_tensor_sizes(query)
|
|
out_sizes[-1] = output_dim
|
|
output_type = query.type().with_sizes(out_sizes)
|
|
return g.op("LlmExporter::FusedAttention", query, key, value, attention_mask, **kwargs).setType(output_type)
|
|
|
|
@staticmethod
|
|
def forward(ctx, query, key, value, attention_mask, output_dim, kv_cache, name, layer_index, kv_shared_layer_index):
|
|
out_shape = list(query.shape)[:2] + [output_dim]
|
|
return query.new_zeros(out_shape)
|
|
|
|
class FusedAttention(torch.nn.Module):
|
|
def __init__(self, hidden_size, kv_cache, name, layer_index=-1, kv_shared_layer_index=-1):
|
|
super(FusedAttention, self).__init__()
|
|
self.hidden_size = hidden_size
|
|
self.kv_cache = int(kv_cache)
|
|
self.name = name
|
|
self.layer_index = layer_index
|
|
self.kv_shared_layer_index = kv_shared_layer_index
|
|
|
|
def forward(self, query, key, value, attention_mask):
|
|
return FusedAttentionOp.apply(query, key, value, attention_mask, self.hidden_size, self.kv_cache, self.name, self.layer_index, self.kv_shared_layer_index)
|
|
|
|
class FusedRoPEOp(torch.autograd.Function):
|
|
@staticmethod
|
|
def symbolic(g, query, key, cos, sin, q_norm_weight, k_norm_weight, rope_cut_head_dim, q_norm_eps, k_norm_eps,
|
|
q_norm, k_norm, name):
|
|
kwargs = {
|
|
"rope_cut_head_dim_i": rope_cut_head_dim,
|
|
"q_norm_eps_f": q_norm_eps,
|
|
"k_norm_eps_f": k_norm_eps,
|
|
"q_norm_i": q_norm,
|
|
"k_norm_i": k_norm,
|
|
"name_s": name,
|
|
}
|
|
query_output, key_output = g.op(
|
|
"LlmExporter::FusedRoPE",
|
|
query,
|
|
key,
|
|
cos,
|
|
sin,
|
|
q_norm_weight,
|
|
k_norm_weight,
|
|
**kwargs,
|
|
outputs=2,
|
|
)
|
|
query_output.setType(query.type())
|
|
key_output.setType(key.type())
|
|
return query_output, key_output
|
|
|
|
@staticmethod
|
|
def forward(ctx, query, key, cos, sin, q_norm_weight, k_norm_weight, rope_cut_head_dim, q_norm_eps, k_norm_eps,
|
|
q_norm, k_norm, name):
|
|
return query, key
|
|
|
|
class FusedRoPE(torch.nn.Module):
|
|
def __init__(self, rope_cut_head_dim, name):
|
|
super(FusedRoPE, self).__init__()
|
|
self.rope_cut_head_dim = int(rope_cut_head_dim)
|
|
self.name = name
|
|
|
|
@staticmethod
|
|
def norm_eps(norm):
|
|
if hasattr(norm, 'variance_epsilon'):
|
|
return float(norm.variance_epsilon)
|
|
return float(norm.eps)
|
|
|
|
def forward(self, query, key, cos, sin, q_norm=None, k_norm=None):
|
|
q_norm_weight = query.new_empty((0,)) if q_norm is None else q_norm.weight
|
|
k_norm_weight = key.new_empty((0,)) if k_norm is None else k_norm.weight
|
|
q_norm_eps = 0.0 if q_norm is None else self.norm_eps(q_norm)
|
|
k_norm_eps = 0.0 if k_norm is None else self.norm_eps(k_norm)
|
|
return FusedRoPEOp.apply(
|
|
query,
|
|
key,
|
|
cos,
|
|
sin,
|
|
q_norm_weight,
|
|
k_norm_weight,
|
|
self.rope_cut_head_dim,
|
|
q_norm_eps,
|
|
k_norm_eps,
|
|
int(q_norm is not None),
|
|
int(k_norm is not None),
|
|
self.name,
|
|
)
|
|
|
|
class MoEOp(torch.autograd.Function):
|
|
@staticmethod
|
|
def symbolic(g, hidden_states, routing_weights, selected_experts, num_experts, top_k, layer_id):
|
|
kwargs = {
|
|
"num_experts_i": num_experts,
|
|
"top_k_i": top_k,
|
|
"layer_id_i": layer_id
|
|
}
|
|
from torch.onnx.symbolic_helper import _get_tensor_sizes
|
|
out_sizes = _get_tensor_sizes(hidden_states)
|
|
output_type = hidden_states.type().with_sizes(out_sizes)
|
|
return g.op("LlmExporter::MoE", hidden_states, routing_weights, selected_experts, **kwargs).setType(output_type)
|
|
|
|
@staticmethod
|
|
def forward(ctx, hidden_states, routing_weights, selected_experts, num_experts, top_k, layer_id):
|
|
return hidden_states
|
|
|
|
class MoE(torch.nn.Module):
|
|
def __init__(self, num_experts, top_k, layer_id):
|
|
super(MoE, self).__init__()
|
|
self.num_experts = num_experts
|
|
self.top_k = top_k
|
|
self.layer_id = layer_id
|
|
|
|
def forward(self, hidden_states, routing_weights, selected_experts):
|
|
return MoEOp.apply(hidden_states, routing_weights, selected_experts, self.num_experts, self.top_k, self.layer_id)
|
|
|
|
class FusedLinearAttentionOp(torch.autograd.Function):
|
|
"""
|
|
Unified Custom Op for Linear Attention variants.
|
|
|
|
Inputs (Tensors):
|
|
0: qkv [B, D, L] - QKV projection output (before conv)
|
|
1: gate [B, L, H] - Pre-computed decay factor
|
|
2: beta [B, L, H] - Pre-computed learning rate (optional)
|
|
3: conv_weight [D, 1, K] - Causal conv weight (optional)
|
|
4: conv_bias [D] - Causal conv bias (optional)
|
|
|
|
Attributes:
|
|
type: "gated_delta_rule" | "mamba" | "rwkv" | "gla" | "retnet"
|
|
key_dim: K total dim = num_k_heads * head_k_dim
|
|
value_dim: V total dim = num_v_heads * head_v_dim
|
|
num_k_heads: number of K/Q heads
|
|
num_v_heads: number of V heads (may differ from K for GQA)
|
|
head_k_dim: per-head K dimension
|
|
head_v_dim: per-head V dimension
|
|
use_qk_l2norm: whether to L2-normalize Q and K
|
|
|
|
Output:
|
|
0: attn_out [B, L, num_v_heads, head_v_dim]
|
|
|
|
Internal State (managed by MNN Execution, not in graph):
|
|
conv_state: [B, D, kernel_size - 1]
|
|
rnn_state: [B, num_v_heads, head_k_dim, head_v_dim]
|
|
"""
|
|
@staticmethod
|
|
def symbolic(g, qkv, gate, beta, conv_weight, name, attn_type,
|
|
num_k_heads, num_v_heads, head_k_dim, head_v_dim, use_qk_l2norm):
|
|
kwargs = {
|
|
"name_s": name,
|
|
"attn_type_s": attn_type,
|
|
"num_k_heads_i": num_k_heads,
|
|
"num_v_heads_i": num_v_heads,
|
|
"head_k_dim_i": head_k_dim,
|
|
"head_v_dim_i": head_v_dim,
|
|
"use_qk_l2norm_i": int(use_qk_l2norm)
|
|
}
|
|
inputs = [qkv, gate, beta, conv_weight]
|
|
from torch.onnx.symbolic_helper import _get_tensor_sizes
|
|
qkv_sizes = _get_tensor_sizes(qkv)
|
|
# qkv shape is [Batch, Dim, SeqLen]
|
|
batch_size = qkv_sizes[0]
|
|
seq_len = qkv_sizes[2]
|
|
|
|
out_sizes = [batch_size, seq_len, num_v_heads, head_v_dim]
|
|
output_type = qkv.type().with_sizes(out_sizes)
|
|
|
|
return g.op("LlmExporter::FusedLinearAttention", *inputs, **kwargs).setType(output_type)
|
|
|
|
@staticmethod
|
|
def forward(ctx, qkv, gate, beta, conv_weight, name, attn_type,
|
|
num_k_heads, num_v_heads, head_k_dim, head_v_dim, use_qk_l2norm):
|
|
# Dummy forward: return correct output shape
|
|
# qkv: [B, D, L] -> output: [B, L, num_v_heads, head_v_dim]
|
|
batch_size = qkv.shape[0]
|
|
seq_len = qkv.shape[2]
|
|
return qkv.new_zeros([batch_size, seq_len, num_v_heads, head_v_dim])
|
|
|
|
class FusedLinearAttention(torch.nn.Module):
|
|
def __init__(self, name, attn_type, num_k_heads, num_v_heads, head_k_dim, head_v_dim, use_qk_l2norm):
|
|
super(FusedLinearAttention, self).__init__()
|
|
self.name = name
|
|
self.attn_type = attn_type
|
|
self.num_k_heads = num_k_heads
|
|
self.num_v_heads = num_v_heads
|
|
self.head_k_dim = head_k_dim
|
|
self.head_v_dim = head_v_dim
|
|
self.use_qk_l2norm = use_qk_l2norm
|
|
|
|
def forward(self, qkv, gate, beta, conv_weight):
|
|
return FusedLinearAttentionOp.apply(qkv, gate, beta, conv_weight, self.name,
|
|
self.attn_type, self.num_k_heads, self.num_v_heads,
|
|
self.head_k_dim, self.head_v_dim, self.use_qk_l2norm)
|