chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,643 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=invalid-name
|
||||
"""Common patterns used in BYOC"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
|
||||
from tvm.relax.dpl.pattern import (
|
||||
DFPattern,
|
||||
GlobalVarPattern,
|
||||
TuplePattern,
|
||||
is_const,
|
||||
is_op,
|
||||
is_tuple_get_item,
|
||||
wildcard,
|
||||
)
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def _with_bias_activation_pattern(
|
||||
out: DFPattern,
|
||||
annotations: dict[str, DFPattern],
|
||||
with_bias: bool = False,
|
||||
activation: str | None = None,
|
||||
allow_reshape: bool = False,
|
||||
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
||||
if with_bias:
|
||||
annotations["bias"] = bias = wildcard()
|
||||
if allow_reshape:
|
||||
reshaped_bias = is_op("relax.reshape")(bias, wildcard(), varg_default_wildcard=True)
|
||||
out = is_op("relax.add")(out, reshaped_bias, varg_default_wildcard=True)
|
||||
else:
|
||||
out = is_op("relax.add")(out, bias)
|
||||
|
||||
if activation:
|
||||
out = is_op(activation)(out)
|
||||
|
||||
return out, annotations
|
||||
|
||||
|
||||
def make_fused_bias_activation_pattern(
|
||||
op_name: str,
|
||||
with_bias: bool = False,
|
||||
activation: str | None = None,
|
||||
allow_reshape: bool = False,
|
||||
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
||||
"""
|
||||
A simple utility to create patterns for an operation fused with bias addition and activation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
op_name: str
|
||||
The name of a Relax op, such as "relax.nn.conv2d"
|
||||
|
||||
with_bias: bool
|
||||
Whether or not to include bias addition
|
||||
|
||||
activation: str
|
||||
The name of an activation Relax op, such as "relax.nn.relu"
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a fused operation
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract
|
||||
important expressions from match result, to power the partition
|
||||
check function and codegen.
|
||||
"""
|
||||
lhs = wildcard()
|
||||
rhs = wildcard()
|
||||
out = is_op(op_name)(lhs, rhs)
|
||||
annotations = {"lhs": lhs, "rhs": rhs, "root": out}
|
||||
|
||||
return _with_bias_activation_pattern(out, annotations, with_bias, activation, allow_reshape)
|
||||
|
||||
|
||||
def make_residual_block_pattern(
|
||||
node_output: DFPattern | tuple[DFPattern, Mapping[str, DFPattern]],
|
||||
binary_op="relax.add",
|
||||
activation=None,
|
||||
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
||||
"""
|
||||
Create pattern for residual block.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node_output: Union[DFPattern, Tuple[DFPattern, Mapping[str, DFPattern]]]
|
||||
The output of previous node.
|
||||
|
||||
binary_op: str
|
||||
The op used to combine previous node output and residual input.
|
||||
|
||||
activation: str
|
||||
The activation function of this residual block. It should be a name of
|
||||
activation Relax op, such as "relax.nn.relu".
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a matrix multiplication.
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract
|
||||
important expressions from match result, to power the partition
|
||||
check function and codegen.
|
||||
"""
|
||||
|
||||
if isinstance(node_output, tuple):
|
||||
node_output, arg_patterns = node_output
|
||||
else:
|
||||
arg_patterns = {}
|
||||
|
||||
residual_input = wildcard()
|
||||
op = is_op(binary_op)
|
||||
output = op(node_output, residual_input) | op(residual_input, node_output)
|
||||
|
||||
if activation is not None:
|
||||
output = is_op(activation)(output)
|
||||
|
||||
return output, {**arg_patterns, "residual": residual_input}
|
||||
|
||||
|
||||
def make_conv2d_pattern(
|
||||
with_bias: bool = False,
|
||||
activation: str | None = None,
|
||||
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
||||
"""
|
||||
Create pattern for 2D convolution.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
with_bias: bool
|
||||
Whether or not to include bias addition
|
||||
|
||||
activation: str
|
||||
The name of an activation Relax op, such as "relax.nn.relu"
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a 2D convolution.
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract
|
||||
important expressions from match result, to power the partition
|
||||
check function and codegen.
|
||||
"""
|
||||
|
||||
input_tensor = wildcard()
|
||||
kernel = wildcard()
|
||||
annotations = {"input": input_tensor, "weight": kernel}
|
||||
|
||||
conv2d = is_op("relax.nn.conv2d")(input_tensor, kernel)
|
||||
annotations["root"] = conv2d
|
||||
|
||||
return _with_bias_activation_pattern(conv2d, annotations, with_bias, activation)
|
||||
|
||||
|
||||
def make_matmul_pattern(
|
||||
with_bias: bool = False,
|
||||
activation: str | None = None,
|
||||
transposed_rhs: bool = False,
|
||||
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
||||
"""
|
||||
Create pattern for matrix multiplication.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
with_bias: bool
|
||||
Whether or not to include bias addition
|
||||
|
||||
activation: str
|
||||
The name of an activation Relax op, such as "relax.nn.relu"
|
||||
|
||||
transposed_rhs: bool
|
||||
Whether the right hand side of multiplication is transposed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a matrix multiplication.
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract
|
||||
important expressions from match result, to power the partition
|
||||
check function and codegen.
|
||||
"""
|
||||
|
||||
lhs = wildcard()
|
||||
rhs = wildcard()
|
||||
annotations = {"lhs": lhs, "rhs": rhs}
|
||||
|
||||
if transposed_rhs:
|
||||
rhs = is_op("relax.permute_dims")(rhs)
|
||||
|
||||
out = is_op("relax.matmul")(lhs, rhs)
|
||||
annotations["root"] = out
|
||||
|
||||
return _with_bias_activation_pattern(out, annotations, with_bias, activation)
|
||||
|
||||
|
||||
def make_attention_pattern(with_bias: bool = False, var_len: bool = False):
|
||||
"""
|
||||
Create pattern for fused multi head attention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
with_bias: bool
|
||||
Whether or not to include bias addition.
|
||||
|
||||
var_len: bool
|
||||
Whether or not to make a pattern for batched attention with variable sequence lengths.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a fused multi head attention.
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract
|
||||
important expressions from match result, to power the partition
|
||||
check function and codegen.
|
||||
"""
|
||||
query = wildcard()
|
||||
key = wildcard()
|
||||
value = wildcard()
|
||||
annotations = {"query": query, "key": key, "value": value}
|
||||
if with_bias:
|
||||
bias = wildcard()
|
||||
annotations["bias"] = bias
|
||||
out = is_op("relax.nn.attention_bias")(query, key, value, bias)
|
||||
elif var_len:
|
||||
seqstart_q = wildcard()
|
||||
seqstart_k = wildcard()
|
||||
max_seqlen_q = wildcard()
|
||||
max_seqlen_k = wildcard()
|
||||
annotations.update(
|
||||
{
|
||||
"seqstart_q": seqstart_q,
|
||||
"seqstart_k": seqstart_k,
|
||||
"max_seqlen_q": max_seqlen_q,
|
||||
"max_seqlen_k": max_seqlen_k,
|
||||
}
|
||||
)
|
||||
out = is_op("relax.nn.attention_var_len")(
|
||||
query, key, value, seqstart_q, seqstart_k, max_seqlen_q, max_seqlen_k
|
||||
)
|
||||
else:
|
||||
out = is_op("relax.nn.attention")(query, key, value)
|
||||
|
||||
return out, annotations
|
||||
|
||||
|
||||
def make_stacked_attention_pattern(start_op: str, with_bias: bool = False, layout="BS3NH"):
|
||||
"""
|
||||
Create pattern for fused multi head attention with stacked input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
start_op: str
|
||||
The starting op for pattern, i.e. `R.split` or `R.strided_slice`.
|
||||
|
||||
with_bias: bool
|
||||
Whether or not to include bias addition
|
||||
|
||||
layout: str
|
||||
The layout of the stacked input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a fused multi head attention.
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract
|
||||
important expressions from match result, to power the partition
|
||||
check function and codegen.
|
||||
"""
|
||||
stacked_qkv = wildcard()
|
||||
ops = {}
|
||||
if start_op == "split":
|
||||
ops["split"] = qkv_tuple = is_op("relax.split")(stacked_qkv)
|
||||
query_raw = is_tuple_get_item(qkv_tuple, 0)
|
||||
key_raw = is_tuple_get_item(qkv_tuple, 1)
|
||||
value_raw = is_tuple_get_item(qkv_tuple, 2)
|
||||
elif start_op == "strided_slice":
|
||||
ops["strided_slice_query"] = query_raw = is_op("relax.strided_slice")(
|
||||
stacked_qkv, varg_default_wildcard=True
|
||||
)
|
||||
ops["strided_slice_key"] = key_raw = is_op("relax.strided_slice")(
|
||||
stacked_qkv, varg_default_wildcard=True
|
||||
)
|
||||
ops["strided_slice_value"] = value_raw = is_op("relax.strided_slice")(
|
||||
stacked_qkv, varg_default_wildcard=True
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
query_reshape_list = wildcard()
|
||||
key_reshape_list = wildcard()
|
||||
value_reshape_list = wildcard()
|
||||
if layout == "BS3NH":
|
||||
query = is_op("relax.reshape")(query_raw, query_reshape_list)
|
||||
key = is_op("relax.reshape")(key_raw, key_reshape_list)
|
||||
value = is_op("relax.reshape")(value_raw, value_reshape_list)
|
||||
elif layout == "SBN3H":
|
||||
ops["q_transpose"] = query = is_op("relax.permute_dims")(query_raw)
|
||||
ops["k_transpose"] = key = is_op("relax.permute_dims")(key_raw)
|
||||
ops["v_transpose"] = value = is_op("relax.permute_dims")(value_raw)
|
||||
annotations = {
|
||||
"stacked_qkv": stacked_qkv,
|
||||
"query_reshape_list": query_reshape_list,
|
||||
"key_reshape_list": key_reshape_list,
|
||||
"value_reshape_list": value_reshape_list,
|
||||
**ops,
|
||||
}
|
||||
if with_bias:
|
||||
bias = wildcard()
|
||||
annotations["bias"] = bias
|
||||
out = is_op("relax.nn.attention_bias")(query, key, value, bias)
|
||||
else:
|
||||
out = is_op("relax.nn.attention")(query, key, value)
|
||||
|
||||
if layout == "SBN3H":
|
||||
out = is_op("relax.permute_dims")(out)
|
||||
|
||||
return out, annotations
|
||||
|
||||
|
||||
def make_layer_norm_pattern():
|
||||
"""Create a layer norm pattern."""
|
||||
inp = wildcard()
|
||||
gamma = wildcard()
|
||||
beta = wildcard()
|
||||
|
||||
return is_op("relax.nn.layer_norm")(inp, gamma, beta), {}
|
||||
|
||||
|
||||
def make_rms_norm_pattern():
|
||||
"""Create a layer norm pattern."""
|
||||
inp = wildcard()
|
||||
weight = wildcard()
|
||||
gv = GlobalVarPattern()
|
||||
out = is_op("relax.call_tir")(gv, TuplePattern([inp, weight]))
|
||||
annotations = {"gv": gv, "inp": inp, "rms_norm": out}
|
||||
return out, annotations
|
||||
|
||||
|
||||
def make_matmul_dequantize_pattern(
|
||||
transposed_rhs: bool = False,
|
||||
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
||||
"""
|
||||
Create pattern for matrix multiplication and dequantize operation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
transposed_rhs: bool
|
||||
Whether the right hand side of multiplication is transposed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a matrix multiplication.
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract important expressions from
|
||||
match result, to power the partition check function and codegen.
|
||||
"""
|
||||
|
||||
lhs = wildcard()
|
||||
rhs = wildcard()
|
||||
annotations = {"lhs": lhs, "rhs": rhs}
|
||||
|
||||
if transposed_rhs:
|
||||
rhs = is_op("relax.permute_dims")(rhs)
|
||||
|
||||
out = is_op("relax.matmul")(lhs, rhs)
|
||||
annotations["root"] = out
|
||||
|
||||
scale = is_const()
|
||||
zp = is_const()
|
||||
annotations.update({"scale": scale, "zp": zp})
|
||||
|
||||
out = is_op("relax.dequantize")(out, scale, zp)
|
||||
|
||||
return out, annotations
|
||||
|
||||
|
||||
def make_matmul_multiply_pattern(
|
||||
transposed_rhs: bool = False,
|
||||
) -> tuple[DFPattern, Mapping[str, DFPattern]]:
|
||||
"""
|
||||
Create pattern for matrix multiplication and multiply operation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
transposed_rhs: bool
|
||||
Whether the right hand side of multiplication is transposed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing a matrix multiplication.
|
||||
|
||||
annotations: Mapping[str, DFPattern]
|
||||
A mapping from name to sub pattern. It can be used to extract important expressions from
|
||||
match result, to power the partition check function and codegen.
|
||||
"""
|
||||
|
||||
lhs = wildcard()
|
||||
rhs = wildcard()
|
||||
scaleA = wildcard()
|
||||
scaleB = wildcard()
|
||||
annotations = {"lhs": lhs, "rhs": rhs, "scaleA": scaleA, "scaleB": scaleB}
|
||||
|
||||
if transposed_rhs:
|
||||
rhs = is_op("relax.permute_dims")(rhs)
|
||||
out = is_op("relax.matmul")(lhs, rhs)
|
||||
annotations["root"] = out
|
||||
scale = is_op("relax.multiply")(scaleA.has_shape((1,)), scaleB.has_shape((1,)))
|
||||
out = is_op("relax.multiply")(out, scale)
|
||||
out = is_op("relax.astype")(out)
|
||||
|
||||
return out, annotations
|
||||
|
||||
|
||||
def make_attention_rewrite_pattern(
|
||||
qkv_layout: str, out_layout: str, with_bias: bool, with_cast: bool, with_kv_repeat: bool = False
|
||||
):
|
||||
"""
|
||||
Create pattern for implicit fused multi head attention rewriting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
qkv_layout: str
|
||||
The layout of the query, key and value tensor, i.e. BSNH or BSH.
|
||||
|
||||
out_layout: str
|
||||
The layout of the output tensor, i.e. BSNH or BSH.
|
||||
|
||||
with_bias: bool
|
||||
Whether or not to include bias addition.
|
||||
|
||||
with_cast: bool
|
||||
Whether or not rewriting is intended to be applied to a module after the FP16 conversion
|
||||
pass.
|
||||
|
||||
with_kv_repeat: bool
|
||||
Whether or not to include the Relax repeat op in the pattern, which is typically used
|
||||
in a Relax module to support multi-query attention.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pattern: DFPattern
|
||||
The resulting pattern describing an implicit fused multi head attention.
|
||||
|
||||
rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
|
||||
The rewriter for the pattern. It will check the matched patterns, and rewrite.
|
||||
If the matched pattern is not able to be rewritten to `R.nn.attention`, the rewriter
|
||||
returns the original IR.
|
||||
"""
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
def handle_input(tensor, layout, transpose, repeat=False):
|
||||
if repeat:
|
||||
tensor = is_op("relax.repeat")(tensor)
|
||||
|
||||
if layout == "BSNH":
|
||||
permuted = is_op("relax.permute_dims")(tensor)
|
||||
shape = wildcard()
|
||||
reshaped = is_op("relax.reshape")(permuted, shape)
|
||||
if transpose:
|
||||
transposed = is_op("relax.permute_dims")(reshaped)
|
||||
|
||||
def rewriter(matchings, x):
|
||||
if matchings[tensor].ty.ndim != 4:
|
||||
return None
|
||||
if list(matchings[permuted].attrs.axes) != [0, 2, 1, 3]:
|
||||
return None
|
||||
before_reshape = matchings[permuted].ty.shape.values
|
||||
after_reshape = matchings[shape].ty.values
|
||||
if not (
|
||||
len(before_reshape) == 4
|
||||
and len(after_reshape) == 3
|
||||
and before_reshape[-2:] == after_reshape[-2:]
|
||||
):
|
||||
return None
|
||||
if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
|
||||
return None
|
||||
return x, x.ty.shape
|
||||
|
||||
if transpose:
|
||||
return transposed, rewriter
|
||||
else:
|
||||
return reshaped, rewriter
|
||||
elif layout == "BSH":
|
||||
if transpose:
|
||||
transposed = is_op("relax.permute_dims")(tensor)
|
||||
|
||||
def rewriter(matchings, x):
|
||||
if matchings[tensor].ty.ndim != 3:
|
||||
return None
|
||||
if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
|
||||
return None
|
||||
before_reshape = x.ty.shape.values
|
||||
after_reshape = [before_reshape[0], before_reshape[1], 1, before_reshape[2]]
|
||||
return R.reshape(x, after_reshape), after_reshape
|
||||
|
||||
if transpose:
|
||||
return transposed, rewriter
|
||||
else:
|
||||
return tensor, rewriter
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
def handle_output(tensor, layout):
|
||||
if layout == "BSNH":
|
||||
shape = wildcard()
|
||||
reshaped = is_op("relax.reshape")(tensor, shape)
|
||||
permuted = is_op("relax.permute_dims")(reshaped)
|
||||
|
||||
def rewriter(matchings, x):
|
||||
if matchings[tensor].ty.ndim != 3:
|
||||
return None
|
||||
before_reshape = matchings[tensor].ty.shape.values
|
||||
after_reshape = matchings[shape].ty.values
|
||||
if not (
|
||||
len(before_reshape) == 3
|
||||
and len(after_reshape) == 4
|
||||
and before_reshape[-2:] == after_reshape[-2:]
|
||||
):
|
||||
return None
|
||||
if list(matchings[permuted].attrs.axes) != [0, 2, 1, 3]:
|
||||
return None
|
||||
return x
|
||||
|
||||
return permuted, rewriter
|
||||
elif layout == "BSH":
|
||||
|
||||
def rewriter(matchings, x):
|
||||
if matchings[tensor].ty.ndim != 3:
|
||||
return None
|
||||
return R.reshape(x, matchings[tensor].ty.shape.values)
|
||||
|
||||
return tensor, rewriter
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
q_raw, k_raw, v_raw = wildcard(), wildcard(), wildcard()
|
||||
q, q_rewriter = handle_input(q_raw, qkv_layout, False)
|
||||
k, k_rewriter = handle_input(k_raw, qkv_layout, True, repeat=with_kv_repeat)
|
||||
v, v_rewriter = handle_input(v_raw, qkv_layout, False, repeat=with_kv_repeat)
|
||||
matmul_1 = is_op("relax.matmul")(q, k)
|
||||
scale = is_const()
|
||||
|
||||
if with_cast:
|
||||
multiply = is_op("relax.multiply")(matmul_1, is_op("relax.astype")(scale))
|
||||
else:
|
||||
multiply = is_op("relax.multiply")(matmul_1, scale)
|
||||
|
||||
if with_bias:
|
||||
bias_raw = wildcard()
|
||||
add = is_op("relax.add")(multiply, bias_raw)
|
||||
softmax_input = add
|
||||
else:
|
||||
softmax_input = multiply
|
||||
|
||||
if with_cast:
|
||||
softmax_input = is_op("relax.astype")(softmax_input)
|
||||
|
||||
softmax = is_op("relax.nn.softmax")(softmax_input)
|
||||
|
||||
if with_cast:
|
||||
softmax_output = is_op("relax.astype")(softmax)
|
||||
else:
|
||||
softmax_output = softmax
|
||||
|
||||
matmul_2 = is_op("relax.matmul")(softmax_output, v)
|
||||
|
||||
out, out_rewriter = handle_output(matmul_2, out_layout)
|
||||
|
||||
def rewriter(original, matchings):
|
||||
query, query_shape = q_rewriter(matchings, matchings[q_raw])
|
||||
key, key_shape = k_rewriter(matchings, matchings[k_raw])
|
||||
value, _ = v_rewriter(matchings, matchings[v_raw])
|
||||
if query is None or key is None or value is None:
|
||||
return original
|
||||
softmax_axis = matchings[softmax].attrs.axis
|
||||
softmax_input_rank = len(matchings[softmax].ty.shape)
|
||||
if softmax_axis == -1:
|
||||
softmax_axis += softmax_input_rank
|
||||
if softmax_axis != softmax_input_rank - 1:
|
||||
return original
|
||||
b, s, n, _ = query_shape
|
||||
_, s_kv, _, _ = key_shape
|
||||
if with_bias:
|
||||
bias = matchings[bias_raw]
|
||||
bias_shape = list(bias.ty.shape)
|
||||
if bias_shape == [b * n, s, s_kv]:
|
||||
bias = R.reshape(bias, [b, n, s, s_kv])
|
||||
elif bias_shape == [b * n, 1, s_kv]:
|
||||
bias = R.reshape(bias, [b, n, 1, s_kv])
|
||||
elif bias_shape == [b, s, s_kv]:
|
||||
bias = R.reshape(bias, [b, 1, s, s_kv])
|
||||
elif bias_shape == [b, 1, s_kv]:
|
||||
bias = R.reshape(bias, [b, 1, 1, s_kv])
|
||||
elif bias_shape in [[1, s, s_kv], [s, s_kv]]:
|
||||
bias = R.reshape(bias, [1, 1, s, s_kv])
|
||||
elif bias_shape in [[1, 1, s_kv], [1, s_kv], [s_kv]]:
|
||||
bias = R.reshape(bias, [1, 1, 1, s_kv])
|
||||
else:
|
||||
return original
|
||||
else:
|
||||
bias = None
|
||||
out = out_rewriter(
|
||||
matchings,
|
||||
R.nn.attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
bias,
|
||||
T.FloatImm(matchings[scale].data.dtype, float(matchings[scale].data.numpy())),
|
||||
),
|
||||
)
|
||||
return out
|
||||
|
||||
return out, rewriter
|
||||
Reference in New Issue
Block a user