Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,467 @@
|
||||
"""ONNX operators as native torch.fx operators.
|
||||
|
||||
This module provides a set of functions to create ONNX operators in the FX graph
|
||||
which are exportable to ONNX.
|
||||
"""
|
||||
|
||||
# flake8: noqa: B950
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
__all__ = [
|
||||
"aten_decompositions",
|
||||
"symbolic",
|
||||
"symbolic_multi_out",
|
||||
"rotary_embedding",
|
||||
"attention",
|
||||
]
|
||||
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.onnx.ops import _impl, _symbolic_impl
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
|
||||
# https://github.com/onnx/onnx/blob/f542e1f06699ea7e1db5f62af53355b64338c723/onnx/onnx.proto#L597
|
||||
_TORCH_DTYPE_TO_ONNX_DTYPE = {
|
||||
torch.float32: 1, # FLOAT
|
||||
torch.uint8: 2, # UINT8
|
||||
torch.int8: 3, # INT8
|
||||
torch.uint16: 4, # UINT16
|
||||
torch.int16: 5, # INT16
|
||||
torch.int32: 6, # INT32
|
||||
torch.int64: 7, # INT64
|
||||
str: 8, # STRING
|
||||
torch.bool: 9, # BOOL
|
||||
torch.float16: 10, # FLOAT16
|
||||
torch.double: 11, # DOUBLE
|
||||
torch.uint32: 12, # UINT32
|
||||
torch.uint64: 13, # UINT64
|
||||
torch.complex64: 14, # COMPLEX64
|
||||
torch.complex128: 15, # COMPLEX128
|
||||
torch.bfloat16: 16, # BFLOAT16
|
||||
torch.float8_e4m3fn: 17, # FLOAT8E4M3FN
|
||||
torch.float8_e4m3fnuz: 18, # FLOAT8E4M3FNUZ
|
||||
torch.float8_e5m2: 19, # FLOAT8E5M2
|
||||
torch.float8_e5m2fnuz: 20, # FLOAT8E5M2FNUZ
|
||||
# 21 = UINT4
|
||||
# 22 = INT4
|
||||
torch.float4_e2m1fn_x2: 23, # FLOAT4E2M1
|
||||
}
|
||||
|
||||
|
||||
def aten_decompositions() -> dict[torch._ops.OpOverload, Callable]:
|
||||
"""Return the ONNX to ATen decomp table."""
|
||||
return _impl.ONNX_ATEN_DECOMP_TABLE
|
||||
|
||||
|
||||
def _parse_domain_op_type(domain_op: str) -> tuple[str, str]:
|
||||
split = domain_op.split("::", 1)
|
||||
if len(split) == 1:
|
||||
domain = ""
|
||||
op_type = split[0]
|
||||
else:
|
||||
domain = split[0]
|
||||
op_type = split[1]
|
||||
return domain, op_type
|
||||
|
||||
|
||||
def symbolic(
|
||||
domain_op: str,
|
||||
/,
|
||||
inputs: Sequence[torch.Tensor | None],
|
||||
attrs: dict[
|
||||
str,
|
||||
int
|
||||
| float
|
||||
| str
|
||||
| bool
|
||||
| Sequence[int]
|
||||
| Sequence[float]
|
||||
| Sequence[str]
|
||||
| Sequence[bool],
|
||||
]
|
||||
| None = None,
|
||||
*,
|
||||
dtype: torch.dtype | int,
|
||||
shape: Sequence[int | torch.SymInt],
|
||||
version: int | None = None,
|
||||
metadata_props: dict[str, str] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Create a symbolic FX operator to represent an arbitrary ONNX operator.
|
||||
|
||||
This function is used to create a symbolic operator with a single output.
|
||||
To create an operator with multiple outputs, use :func:`symbolic_multi_out`.
|
||||
|
||||
You may use ``if torch.onnx.is_in_onnx_export()`` to conditionally enable the
|
||||
symbolic logic only during ``torch.onnx.export()``.
|
||||
|
||||
Example::
|
||||
|
||||
class CustomOp(torch.nn.Module):
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# Normal torch operators can interleave with the symbolic ops during ONNX export
|
||||
x = x + 1
|
||||
|
||||
# Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
|
||||
# The output tensor will have the specified dtype and shape
|
||||
val = torch.onnx.ops.symbolic(
|
||||
"custom_domain::CustomOp",
|
||||
(x,),
|
||||
dict(attr_key="attr_value"),
|
||||
dtype=x.dtype,
|
||||
shape=x.shape,
|
||||
version=1,
|
||||
)
|
||||
|
||||
# The result of the symbolic op can be used in normal torch operations during ONNX export
|
||||
return torch.nn.functional.relu(val)
|
||||
|
||||
|
||||
# You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).
|
||||
|
||||
Args:
|
||||
domain_op: The domain and operator name, separated by "::". For example,
|
||||
"custom_domain::CustomOp".
|
||||
inputs: The input tensors to the operator.
|
||||
attrs: The attributes of the operator. The keys are attribute names and
|
||||
the values are attribute values. Valid attribute types are int, float,
|
||||
str, bool, and lists of int, float, str, and bool. Tensor attributes
|
||||
are unsupported.
|
||||
dtype: The data type of the output tensor.This can be either a torch.dtype
|
||||
or an integer representing the ONNX data type.
|
||||
shape: The shape of the output tensor. This can be a list of integers or
|
||||
SymInt values.
|
||||
version: The version of the opset used for the operator.
|
||||
metadata_props: Metadata properties for the ONNX node.
|
||||
This is a dictionary of str-str pairs.
|
||||
|
||||
Returns:
|
||||
The output tensor of the operator.
|
||||
"""
|
||||
if not isinstance(dtype, int):
|
||||
torch._check(
|
||||
dtype in _TORCH_DTYPE_TO_ONNX_DTYPE, lambda: f"Unsupported dtype: {dtype}"
|
||||
)
|
||||
dtype = _TORCH_DTYPE_TO_ONNX_DTYPE[dtype]
|
||||
domain, op_type = _parse_domain_op_type(domain_op)
|
||||
if attrs is None:
|
||||
attrs = {}
|
||||
encoded_attrs = _symbolic_impl.EncodedAttrs.from_dict(attrs)
|
||||
# TODO: Parse domain
|
||||
return _symbolic_impl._symbolic(
|
||||
inputs,
|
||||
op_type,
|
||||
dtype,
|
||||
shape=shape,
|
||||
attr_keys=encoded_attrs.attr_keys,
|
||||
attr_types=encoded_attrs.attr_types,
|
||||
attr_pos=encoded_attrs.attr_pos,
|
||||
attr_ints=encoded_attrs.attr_ints,
|
||||
attr_floats=encoded_attrs.attr_floats,
|
||||
attr_strs=encoded_attrs.attr_strs,
|
||||
metadata_props_keys=metadata_props.keys() if metadata_props else [],
|
||||
metadata_props_values=metadata_props.values() if metadata_props else [],
|
||||
domain=domain,
|
||||
version=version,
|
||||
)
|
||||
|
||||
|
||||
def symbolic_multi_out(
|
||||
domain_op: str,
|
||||
/,
|
||||
inputs: Sequence[torch.Tensor | None],
|
||||
attrs: dict[
|
||||
str,
|
||||
int
|
||||
| float
|
||||
| str
|
||||
| bool
|
||||
| Sequence[int]
|
||||
| Sequence[float]
|
||||
| Sequence[str]
|
||||
| Sequence[bool],
|
||||
]
|
||||
| None = None,
|
||||
*,
|
||||
dtypes: Sequence[torch.dtype | int],
|
||||
shapes: Sequence[Sequence[int | torch.SymInt]],
|
||||
version: int | None = None,
|
||||
metadata_props: dict[str, str] | None = None,
|
||||
) -> Sequence[torch.Tensor]:
|
||||
"""Create a symbolic FX operator to represent an arbitrary ONNX operator with multiple outputs.
|
||||
|
||||
You may use ``if torch.onnx.is_in_onnx_export()`` to conditionally enable the
|
||||
symbolic logic only during ``torch.onnx.export()``.
|
||||
|
||||
Example::
|
||||
|
||||
class CustomOp(torch.nn.Module):
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# Normal torch operators can interleave with the symbolic ops during ONNX export
|
||||
x = x + 1
|
||||
|
||||
# Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
|
||||
# The output tensors will have the specified dtypes and shapes
|
||||
(out1, out2) = torch.onnx.ops.symbolic_multi_out(
|
||||
"custom_domain::CustomOp",
|
||||
(x,),
|
||||
dict(attr_key="attr_value"),
|
||||
dtypes=(x.dtype, torch.float32),
|
||||
shapes=(x.shape, [1, 2, 3]),
|
||||
version=1,
|
||||
)
|
||||
|
||||
# The result of the symbolic op can be used in normal torch operations during ONNX export
|
||||
return torch.nn.functional.relu(out1 + out2)
|
||||
|
||||
|
||||
# You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).
|
||||
|
||||
Args:
|
||||
domain_op: The domain and operator name, separated by "::". For example,
|
||||
"custom_domain::CustomOp".
|
||||
inputs: The input tensors to the operator.
|
||||
attrs: The attributes of the operator. The keys are attribute names and
|
||||
the values are attribute values. Valid attribute types are int, float,
|
||||
str, bool, and lists of int, float, str, and bool. Tensor attributes
|
||||
are unsupported.
|
||||
dtypes: The data types of the output tensors. This can be a list of
|
||||
torch.dtype or integers representing the ONNX data types. The length
|
||||
of this list must be the number of outputs.
|
||||
shapes: The shapes of the output tensors. This can be a list of lists of
|
||||
integers or SymInt values. The length of this list must be the number of outputs.
|
||||
version: The version of the opset used for the operator.
|
||||
metadata_props: Metadata properties for the ONNX node.
|
||||
This is a dictionary of str-str pairs.
|
||||
|
||||
Returns:
|
||||
A list of output tensors of the operator.
|
||||
"""
|
||||
torch._check(
|
||||
len(shapes) == len(dtypes),
|
||||
lambda: f"Number of shapes ({len(shapes)}) must match number of dtypes ({len(dtypes)})",
|
||||
)
|
||||
onnx_dtypes = []
|
||||
for dtype in dtypes:
|
||||
if not isinstance(dtype, int):
|
||||
torch._check(
|
||||
dtype in _TORCH_DTYPE_TO_ONNX_DTYPE,
|
||||
lambda: f"Unsupported dtype: {dtype}",
|
||||
)
|
||||
onnx_dtypes.append(_TORCH_DTYPE_TO_ONNX_DTYPE[dtype])
|
||||
else:
|
||||
onnx_dtypes.append(dtype)
|
||||
domain, op_type = _parse_domain_op_type(domain_op)
|
||||
if attrs is None:
|
||||
attrs = {}
|
||||
encoded_attrs = _symbolic_impl.EncodedAttrs.from_dict(attrs)
|
||||
# Use the size of dtypes to determine the number of outputs
|
||||
return _symbolic_impl._symbolic_multi_out(
|
||||
inputs,
|
||||
op_type,
|
||||
onnx_dtypes,
|
||||
shapes=shapes,
|
||||
attr_keys=encoded_attrs.attr_keys,
|
||||
attr_types=encoded_attrs.attr_types,
|
||||
attr_pos=encoded_attrs.attr_pos,
|
||||
attr_ints=encoded_attrs.attr_ints,
|
||||
attr_floats=encoded_attrs.attr_floats,
|
||||
attr_strs=encoded_attrs.attr_strs,
|
||||
metadata_props_keys=metadata_props.keys() if metadata_props else [],
|
||||
metadata_props_values=metadata_props.values() if metadata_props else [],
|
||||
domain=domain,
|
||||
version=version,
|
||||
)
|
||||
|
||||
|
||||
def rotary_embedding(
|
||||
X: torch.Tensor,
|
||||
cos_cache: torch.Tensor,
|
||||
sin_cache: torch.Tensor,
|
||||
position_ids: torch.Tensor | None = None,
|
||||
*,
|
||||
interleaved: bool = False,
|
||||
num_heads: int = 0,
|
||||
rotary_embedding_dim: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""RotaryEmbedding op in ONNX.
|
||||
|
||||
https://onnx.ai/onnx/operators/onnx__RotaryEmbedding.html
|
||||
|
||||
RotaryEmbedding is the implementation of rotary positional embeddings (RoPE) based on the paper https://arxiv.org/pdf/2104.09864.
|
||||
The key advantage of RoPE is that it allows the model to understand both the absolute position of a token and the relative distances
|
||||
between tokens. This is achieved through a rotational mechanism where the extent of rotation is computed based on the token's absolute position (position_ids).
|
||||
|
||||
The rotational mechanism is defined by sine and cosine functions that are used to represent the rotation angles.
|
||||
For each token in the sequence, its positional embedding is computed by rotating its embedding vector. This is done by splitting the
|
||||
embedding vector either into two halves or interleaving every alternate token and applying the rotation matrix to each half of the embedding vector.
|
||||
The rotation matrix is parameterized by the token's position in the sequence. The rotated halves of the embedding vector are concatenated
|
||||
to form the final positional embedding for each token. The rotated positional embeddings are used in the self-attention mechanism.
|
||||
The rotation ensures that the model captures both absolute and relative positional information.
|
||||
|
||||
Args:
|
||||
X: The input tensor representing the token embeddings. 4D tensor with
|
||||
shape `(batch_size, num_heads, sequence_length, head_size)` or 3D tensor
|
||||
with shape `(batch_size, sequence_length, hidden_size)`. For cases with
|
||||
a 4D input tensor, `head_size` has to be even. For cases with a 3D input
|
||||
tensor, `num_heads` attribute must be provided and `hidden_size` must
|
||||
be an even multiple of `num_heads` where `hidden_size = num_heads * head_size`
|
||||
cos_cache: The cosine values for the rotation. 2D tensor with shape `(max_position_id_plus_1, head_size / 2)`
|
||||
for full rotation or `(max_position_id_plus_1, rotary_embedding_dim / 2)`
|
||||
for partial rotation when `position_ids` are provided. 3D tensor with shape
|
||||
`(batch_size, sequence_length, head_size / 2)` for full rotation or
|
||||
`(batch_size, sequence_length, rotary_embedding_dim / 2)` for partial
|
||||
rotation when `position_ids` are not provided. `max_position_id_plus_1`
|
||||
is a parameter to the model.
|
||||
sin_cache: The sine values for the rotation. 2D tensor with shape `(max_position_id_plus_1, head_size / 2)`
|
||||
for full rotation or `(max_position_id_plus_1, rotary_embedding_dim / 2)`
|
||||
for partial rotation when `position_ids` are provided. 3D tensor with shape
|
||||
`(batch_size, sequence_length, head_size / 2)` for full rotation or
|
||||
`(batch_size, sequence_length, rotary_embedding_dim / 2)` for partial rotation
|
||||
when `position_ids` are not provided. `max_position_id_plus_1` is a parameter
|
||||
to the model.
|
||||
position_ids: The position indices for the tokens. 2D tensor with shape
|
||||
`(batch_size, sequence_length)`.
|
||||
interleaved: Rotate using interleaved pattern. Default value is 0 (False).
|
||||
num_heads: Number of attention heads. Must be provided when input is a 3D tensor.
|
||||
rotary_embedding_dim: Rotary embedding dimension used to apply partial rotary embeddings.
|
||||
|
||||
Returns:
|
||||
Tensor with same shape as input.
|
||||
"""
|
||||
return _impl.rotary_embedding_23(
|
||||
X,
|
||||
cos_cache,
|
||||
sin_cache,
|
||||
position_ids=position_ids,
|
||||
interleaved=interleaved,
|
||||
num_heads=num_heads,
|
||||
rotary_embedding_dim=rotary_embedding_dim,
|
||||
)
|
||||
|
||||
|
||||
def attention(
|
||||
Q: torch.Tensor,
|
||||
K: torch.Tensor,
|
||||
V: torch.Tensor,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
past_key: torch.Tensor | None = None,
|
||||
past_value: torch.Tensor | None = None,
|
||||
*,
|
||||
is_causal: bool = False,
|
||||
kv_num_heads: int = 0,
|
||||
q_num_heads: int = 0,
|
||||
qk_matmul_output_mode: int = 0,
|
||||
scale: float | None = None,
|
||||
softcap: float = 0.0,
|
||||
softmax_precision: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Attention op in ONNX.
|
||||
|
||||
https://onnx.ai/onnx/operators/onnx__Attention.html
|
||||
|
||||
Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed.
|
||||
|
||||
This operator covers self and cross variants of the attention operation based on sequence lengths of K, Q and V.
|
||||
|
||||
For self attention, ``kv_sequence_length`` equals to ``q_sequence_length``.
|
||||
|
||||
For cross attention, query and key might have different lengths.
|
||||
|
||||
This operator also covers the 3 following variants based on the number of heads:
|
||||
|
||||
1. Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762, `q_num_heads = kv_num_heads`.
|
||||
2. Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245, `q_num_heads > kv_num_heads`, `q_num_heads % kv_num_heads == 0`.
|
||||
3. Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150, `q_num_heads > kv_num_heads`, `kv_num_heads=1`.
|
||||
|
||||
Attention bias to be added is calculated based on ``attn_mask`` input and ``is_causal` `attribute``, only one of which can be provided.
|
||||
|
||||
1. If ``is_causal`` is set to `1`, the attention masking is a lower triangular matrix when the mask is a square matrix. The attention masking has the form of the upper left causal bias due to the alignment.
|
||||
2. `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score.
|
||||
|
||||
Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them.
|
||||
The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided::
|
||||
|
||||
The following pattern is applied by this operator:
|
||||
Q K V
|
||||
| | |
|
||||
Q*sqrt(scale) K*sqrt(scale) |
|
||||
| | |
|
||||
| Transpose |
|
||||
| | |
|
||||
---MatMul--- |
|
||||
| |
|
||||
at_mask---Add |
|
||||
| |
|
||||
softcap (if provided) |
|
||||
| |
|
||||
Softmax |
|
||||
| |
|
||||
-----MatMul------
|
||||
|
|
||||
Y
|
||||
|
||||
Args:
|
||||
Q: Query tensor. 4D tensor with shape `(batch_size, q_num_heads, q_sequence_length, head_size)` or 3D tensor
|
||||
with shape `(batch_size, q_sequence_length, q_hidden_size)`. For cases with a 3D input tensor,
|
||||
`q_hidden_size = q_num_heads * head_size`
|
||||
K: Key tensor. 4D tensor with shape `(batch_size, kv_num_heads, kv_sequence_length, head_size)` or 3D tensor
|
||||
with shape `(batch_size, kv_sequence_length, k_hidden_size)`. For cases with a 3D input tensor,
|
||||
`k_hidden_size = kv_num_heads * head_size`
|
||||
V: Value tensor. 4D tensor with shape `(batch_size, kv_num_heads, kv_sequence_length, v_head_size)` or 3D tensor
|
||||
with shape `(batch_size, kv_sequence_length, v_hidden_size)`. For cases with a 3D input tensor,
|
||||
`v_hidden_size = kv_num_heads * v_head_size`
|
||||
attn_mask: Attention mask. Shape must be broadcastable to 4D tensor with shape
|
||||
`(batch_size, q_num_heads, q_sequence_length, total_sequence_length)` where
|
||||
`total_sequence_length = past_sequence_length + kv_sequence_length`. Two types of masks are supported.
|
||||
A boolean mask where a value of True indicates that the element should take part in attention.
|
||||
Also supports a float mask of the same type as query, key, value that is added to the attention score.
|
||||
past_key: Past state cache for key with shape `(batch_size, kv_num_heads, past_sequence_length, head_size)`
|
||||
past_value: Past state cache for value with shape `(batch_size, kv_num_heads, past_sequence_length, v_head_size)`
|
||||
is_causal: If set to True, the attention masking is a lower triangular matrix when the mask is a square matrix.
|
||||
The attention masking has the form of the upper left causal bias due to the alignment.
|
||||
kv_num_heads: Number of heads of key and value. Must be used with 3D inputs of Q, K and V.
|
||||
q_num_heads: Number of heads of query. Must be used with 3D inputs of Q, K and V.
|
||||
qk_matmul_output_mode: If set to 0, qk_matmul_output is the output of qk matmul. If set to 1,
|
||||
qk_matmul_output includes the addition of the attention mask to the output of qk matmul.
|
||||
If set to 2, qk_matmul_output is the output after the softcap operation. If set to 3,
|
||||
qk_matmul_output is the output after the softmax operation. Default value is 0.
|
||||
scale: Scaling factor applied to Q*K^T. Default value is 1/sqrt(head_size). To prevent numerical overflow,
|
||||
scale Q, K by sqrt(scale) before matmul.
|
||||
softcap: Softcap value for attention weights. Default value is 0.
|
||||
softmax_precision: The floating-point precision used in softmax computation. If softmax precision is not provided,
|
||||
the same precision as the input of softmax (Q and K) is used.
|
||||
|
||||
Returns:
|
||||
A tuple containing:
|
||||
- The output tensor. 4D tensor with shape `(batch_size, q_num_heads, q_sequence_length, v_head_size)` or 3D tensor
|
||||
with shape `(batch_size, q_sequence_length, hidden_size)`. For cases with a 3D input tensor,
|
||||
`hidden_size = q_num_heads * v_head_size`
|
||||
- Updated key cache with shape `(batch_size, kv_num_heads, total_sequence_length, head_size)` where
|
||||
`total_sequence_length = past_sequence_length + kv_sequence_length`.
|
||||
- Updated value cache with shape `(batch_size, kv_num_heads, total_sequence_length, v_head_size)` where
|
||||
`total_sequence_length = past_sequence_length + kv_sequence_length`.
|
||||
- The output of QK matmul. 4D tensor with shape `(batch_size, q_num_heads, q_sequence_length, total_sequence_length)`
|
||||
where `total_sequence_length = past_sequence_length + kv_sequence_length`.
|
||||
"""
|
||||
return _impl.attention_23(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
attn_mask=attn_mask,
|
||||
past_key=past_key,
|
||||
past_value=past_value,
|
||||
is_causal=is_causal,
|
||||
kv_num_heads=kv_num_heads,
|
||||
q_num_heads=q_num_heads,
|
||||
qk_matmul_output_mode=qk_matmul_output_mode,
|
||||
scale=scale,
|
||||
softcap=softcap,
|
||||
softmax_precision=softmax_precision,
|
||||
)
|
||||
@@ -0,0 +1,27 @@
|
||||
import torch
|
||||
|
||||
|
||||
ONNX_DTYPE_TO_TORCH_DTYPE: dict[int, torch.dtype] = {
|
||||
1: torch.float32, # FLOAT
|
||||
2: torch.uint8, # UINT8
|
||||
3: torch.int8, # INT8
|
||||
4: torch.uint16, # UINT16
|
||||
5: torch.int16, # INT16
|
||||
6: torch.int32, # INT32
|
||||
7: torch.int64, # INT64
|
||||
9: torch.bool, # BOOL
|
||||
10: torch.float16, # FLOAT16
|
||||
11: torch.double, # DOUBLE
|
||||
12: torch.uint32, # UINT32
|
||||
13: torch.uint64, # UINT64
|
||||
14: torch.complex64, # COMPLEX64
|
||||
15: torch.complex128, # COMPLEX128
|
||||
16: torch.bfloat16, # BFLOAT16
|
||||
17: torch.float8_e4m3fn, # FLOAT8E4M3FN
|
||||
18: torch.float8_e4m3fnuz, # FLOAT8E4M3FNUZ
|
||||
19: torch.float8_e5m2, # FLOAT8E5M2
|
||||
20: torch.float8_e5m2fnuz, # FLOAT8E5M2FNUZ
|
||||
21: torch.uint8, # UINT4
|
||||
22: torch.uint8, # INT4
|
||||
23: torch.float4_e2m1fn_x2, # FLOAT4E2M1
|
||||
}
|
||||
@@ -0,0 +1,547 @@
|
||||
"""Implementations of ONNX operators as native Torch ops.
|
||||
|
||||
NOTE: Fake implementations:
|
||||
Refer to https://docs.pytorch.org/docs/stable/library.html#torch.library.register_fake
|
||||
for more details on how to create fake kernels.
|
||||
"""
|
||||
|
||||
# flake8: noqa: B950
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from typing import TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
import torch
|
||||
from torch.onnx.ops import _dtype_mappings
|
||||
|
||||
|
||||
# Use ParamSpec for better type preservation instead of bound Callable TypeVar
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
# ONNX to ATen decomp table
|
||||
ONNX_ATEN_DECOMP_TABLE: dict[torch._ops.OpOverload, Callable] = {}
|
||||
_ATTENTION_23_ALLOWED_INTERMEDIATE_PRECISIONS = frozenset(
|
||||
{
|
||||
1, # FLOAT
|
||||
10, # FLOAT16
|
||||
11, # DOUBLE
|
||||
16, # BFLOAT16
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _onnx_op(
|
||||
op_type: str, opset_version: int, fake_impl: Callable[_P, _R]
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
|
||||
"""Decorator to register an ONNX operator with a custom implementation."""
|
||||
|
||||
def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
overload = f"opset{opset_version}"
|
||||
torch_op = torch.library.custom_op(
|
||||
f"onnx::{op_type}.{overload}", mutates_args=()
|
||||
)(func)
|
||||
ONNX_ATEN_DECOMP_TABLE[getattr(getattr(torch.ops.onnx, op_type), overload)] = (
|
||||
func # type: ignore[assignment]
|
||||
)
|
||||
torch_op.register_fake(fake_impl)
|
||||
return torch_op # type: ignore[return-value]
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def _rotary_embedding_23_fake_impl(
|
||||
x: torch.Tensor,
|
||||
cos_cache: torch.Tensor,
|
||||
sin_cache: torch.Tensor,
|
||||
position_ids: torch.Tensor | None = None,
|
||||
*,
|
||||
interleaved: bool = False,
|
||||
num_heads: int = 0,
|
||||
rotary_embedding_dim: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""Fake implementation for RotaryEmbedding-23 for torch.compile purposes."""
|
||||
return x.clone()
|
||||
|
||||
|
||||
@_onnx_op("RotaryEmbedding", 23, _rotary_embedding_23_fake_impl)
|
||||
def rotary_embedding_23(
|
||||
x: torch.Tensor,
|
||||
cos_cache: torch.Tensor,
|
||||
sin_cache: torch.Tensor,
|
||||
position_ids: torch.Tensor | None = None,
|
||||
*,
|
||||
interleaved: bool = False,
|
||||
num_heads: int = 0,
|
||||
rotary_embedding_dim: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""RotaryEmbedding-23 https://onnx.ai/onnx/operators/onnx__RotaryEmbedding.html#rotaryembedding-23"""
|
||||
# x has shape (batch_size, num_heads, sequence_length, head_size)
|
||||
# or (batch_size, sequence_length, hidden_size)
|
||||
input_shape = x.shape
|
||||
input_rank = len(input_shape)
|
||||
batch_size = input_shape[0]
|
||||
sequence_length = input_shape[-2]
|
||||
|
||||
# Validate position_ids and caches match x
|
||||
if position_ids is not None:
|
||||
torch._check(
|
||||
position_ids.dim() == 2,
|
||||
lambda: f"position_ids must be 2D when provided. Received shape {position_ids.shape}",
|
||||
)
|
||||
torch._check(
|
||||
position_ids.shape[0] == batch_size,
|
||||
lambda: f"position_ids first dim (batch) must match x.shape[0] ({batch_size}). Received {position_ids.shape[0]}",
|
||||
)
|
||||
torch._check(
|
||||
position_ids.shape[1] == sequence_length,
|
||||
lambda: f"position_ids second dim (sequence) must match x.shape[-2] ({sequence_length}). Received {position_ids.shape[1]}",
|
||||
)
|
||||
torch._check(
|
||||
cos_cache.dim() == 2 and sin_cache.dim() == 2,
|
||||
lambda: "cos_cache/sin_cache must be 2D when position_ids is provided. "
|
||||
f"Received cos_cache shape {cos_cache.shape}, sin_cache shape {sin_cache.shape}",
|
||||
)
|
||||
else:
|
||||
torch._check(
|
||||
cos_cache.dim() == 3 and sin_cache.dim() == 3,
|
||||
lambda: "cos_cache/sin_cache must be 3D when position_ids is not provided. "
|
||||
f"Received cos_cache shape {cos_cache.shape}, sin_cache shape {sin_cache.shape}",
|
||||
)
|
||||
|
||||
# First ensure x has shape [batch_size, num_heads, seq_len, head_size]
|
||||
# So that the rotation logic can be shared with reshaped 3D inputs
|
||||
if input_rank == 4:
|
||||
# Reshape from (batch_size, num_heads, seq_len, head_size)
|
||||
# to [batch_size, seq_len, num_heads, head_size]
|
||||
x = torch.permute(x, (0, 2, 1, 3))
|
||||
elif input_rank == 3:
|
||||
torch._check(
|
||||
num_heads != 0,
|
||||
lambda: f"num_heads must be provided for 3D inputs. Received input tensor with shape {input_shape}",
|
||||
)
|
||||
hidden_size = input_shape[2]
|
||||
head_size = hidden_size // num_heads
|
||||
new_shape = [batch_size, sequence_length, num_heads, head_size]
|
||||
x = torch.reshape(x, new_shape)
|
||||
|
||||
torch._check(len(x.shape) == 4, lambda: "x should be a 4D tensor by now")
|
||||
head_size = x.shape[3]
|
||||
|
||||
# Fully or partially perform rotation on x based on rotary_embedding_dim attribute
|
||||
if rotary_embedding_dim == 0:
|
||||
# If rotary_embedding_dim not provided, perform full rotation by using head_size
|
||||
rotary_embedding_dim = head_size
|
||||
x_rotate = x[:, :, :, :rotary_embedding_dim]
|
||||
x_not_rotate = x[:, :, :, rotary_embedding_dim:]
|
||||
rotary_embedding_dim_half = rotary_embedding_dim // 2
|
||||
|
||||
# Retrieve sin and cos caches using position ids
|
||||
if position_ids is not None:
|
||||
cos = cos_cache[
|
||||
position_ids
|
||||
] # Shape: [batch_size, sequence_length, head_size/2]
|
||||
sin = sin_cache[
|
||||
position_ids
|
||||
] # Shape: [batch_size, sequence_length, head_size/2]
|
||||
else:
|
||||
cos = cos_cache # Shape: [batch_size, sequence_length, rotary_embedding_dim/2]
|
||||
sin = sin_cache # Shape: [batch_size, sequence_length, rotary_embedding_dim/2]
|
||||
|
||||
torch._check(
|
||||
cos.shape[0] == batch_size and cos.shape[1] == sequence_length,
|
||||
lambda: f"cos has shape {cos.shape} but expected (batch={batch_size}, seq={sequence_length}, ...)",
|
||||
)
|
||||
torch._check(
|
||||
sin.shape[0] == batch_size and sin.shape[1] == sequence_length,
|
||||
lambda: f"sin has shape {sin.shape} but expected (batch={batch_size}, seq={sequence_length}, ...)",
|
||||
)
|
||||
torch._check(
|
||||
cos.shape[-1] == rotary_embedding_dim_half,
|
||||
lambda: f"Last dimension of cos cache ({cos.shape[-1]}) should match rotary_embedding_dim/2 ({rotary_embedding_dim_half}).",
|
||||
)
|
||||
torch._check(
|
||||
sin.shape[-1] == rotary_embedding_dim_half,
|
||||
lambda: f"Last dimension of sin cache ({sin.shape[-1]}) should match rotary_embedding_dim/2 ({rotary_embedding_dim_half}).",
|
||||
)
|
||||
cos = torch.unsqueeze(
|
||||
cos, 2
|
||||
) # Shape: [batch_size, sequence_length, 1, rotary_embedding_dim/2]
|
||||
sin = torch.unsqueeze(
|
||||
sin, 2
|
||||
) # Shape: [batch_size, sequence_length, 1, rotary_embedding_dim/2]
|
||||
|
||||
# Either divide the x in halves or interleave (based on interleaved attribute)
|
||||
if interleaved:
|
||||
x1 = x_rotate[:, :, :, 0::2]
|
||||
x2 = x_rotate[:, :, :, 1::2]
|
||||
else:
|
||||
x1, x2 = torch.chunk(x_rotate, 2, dim=-1)
|
||||
|
||||
# Calculate real and imaginary values
|
||||
real = cos * x1 - sin * x2
|
||||
imag = sin * x1 + cos * x2
|
||||
|
||||
# Inserted rotated embeddings back to the original x
|
||||
if interleaved:
|
||||
# x_rotate[:, :, :, 0::2] = real
|
||||
# x_rotate[:, :, :, 1::2] = imag
|
||||
real = torch.unsqueeze(real, -1)
|
||||
imag = torch.unsqueeze(imag, -1)
|
||||
x_rotate_concat = torch.cat((real, imag), dim=-1)
|
||||
x_rotate = torch.reshape(x_rotate_concat, x_rotate.shape)
|
||||
else:
|
||||
x_rotate = torch.cat((real, imag), dim=-1)
|
||||
output = torch.cat((x_rotate, x_not_rotate), dim=-1)
|
||||
if input_rank == 3:
|
||||
return torch.reshape(output, input_shape)
|
||||
|
||||
# Return the dimensions to the original order
|
||||
return torch.permute(output, (0, 2, 1, 3))
|
||||
|
||||
|
||||
def _get_scale_factor(scale: float | None, head_size: int) -> float:
|
||||
"""Get the scale factor for attention computation."""
|
||||
return scale if scale is not None else (1.0 / math.sqrt(head_size))
|
||||
|
||||
|
||||
def _reshape_3d_to_4d(
|
||||
tensor: torch.Tensor, batch_size: int, num_heads: int
|
||||
) -> torch.Tensor:
|
||||
"""Reshape 3D tensor to 4D for multi-head attention."""
|
||||
sequence_length, hidden_size = tensor.shape[1], tensor.shape[2]
|
||||
head_size = hidden_size // num_heads
|
||||
return (
|
||||
tensor.view(batch_size, sequence_length, num_heads, head_size)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
|
||||
def _get_qk_output_for_aten_spda(
|
||||
Q: torch.Tensor,
|
||||
K: torch.Tensor,
|
||||
current_q_num_heads: int,
|
||||
current_kv_num_heads: int,
|
||||
scale: float | None,
|
||||
qk_matmul_output_mode: int,
|
||||
) -> torch.Tensor:
|
||||
"""Get QK output tensor based on the specified mode."""
|
||||
if qk_matmul_output_mode == 0:
|
||||
return _compute_qk_output_for_mode_0(
|
||||
Q, K, current_q_num_heads, current_kv_num_heads, scale
|
||||
)
|
||||
else:
|
||||
# For other modes, return a zero tensor with correct shape
|
||||
return torch.zeros_like(torch.matmul(Q, K.transpose(-2, -1)))
|
||||
|
||||
|
||||
def _validate_gqa_configuration(
|
||||
current_q_num_heads: int, current_kv_num_heads: int
|
||||
) -> None:
|
||||
"""Validate Group Query Attention configuration."""
|
||||
torch._check(
|
||||
current_q_num_heads % current_kv_num_heads == 0,
|
||||
lambda: f"q_num_heads ({current_q_num_heads}) must be divisible by kv_num_heads ({current_kv_num_heads}) for GQA",
|
||||
)
|
||||
|
||||
|
||||
def _compute_qk_output_for_mode_0(
|
||||
Q: torch.Tensor,
|
||||
K: torch.Tensor,
|
||||
current_q_num_heads: int,
|
||||
current_kv_num_heads: int,
|
||||
scale: float | None,
|
||||
) -> torch.Tensor:
|
||||
"""Helper function to compute QK output for qk_matmul_output_mode == 0."""
|
||||
# Handle GQA manually for QK output
|
||||
K_for_qk = K
|
||||
if current_q_num_heads != current_kv_num_heads:
|
||||
repeat_factor = current_q_num_heads // current_kv_num_heads
|
||||
K_for_qk = K.repeat_interleave(repeat_factor, dim=1)
|
||||
|
||||
scale_factor = _get_scale_factor(scale, Q.shape[3])
|
||||
# Scale both Q and K by sqrt(scale_factor) for numerical stability
|
||||
sqrt_scale = math.sqrt(scale_factor)
|
||||
Q_scaled = Q * sqrt_scale
|
||||
K_scaled = K_for_qk * sqrt_scale
|
||||
return torch.matmul(Q_scaled, K_scaled.transpose(-2, -1))
|
||||
|
||||
|
||||
def _attention_23_fake_impl(
|
||||
Q: torch.Tensor,
|
||||
K: torch.Tensor,
|
||||
V: torch.Tensor,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
past_key: torch.Tensor | None = None,
|
||||
past_value: torch.Tensor | None = None,
|
||||
*,
|
||||
is_causal: bool = False,
|
||||
kv_num_heads: int = 0,
|
||||
q_num_heads: int = 0,
|
||||
qk_matmul_output_mode: int = 0,
|
||||
scale: float | None = None,
|
||||
softcap: float = 0.0,
|
||||
softmax_precision: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Fake implementation for Attention-23 for torch.compile purposes."""
|
||||
batch_size = Q.shape[0]
|
||||
|
||||
# Handle 3D vs 4D input shapes
|
||||
if len(Q.shape) == 3:
|
||||
# 3D input: (batch_size, sequence_length, hidden_size)
|
||||
q_sequence_length = Q.shape[1]
|
||||
output_shape = Q.shape # Same shape as Q for 3D output
|
||||
|
||||
# For present_key and present_value, we need 4D shapes
|
||||
if past_key is not None:
|
||||
present_key_shape = (
|
||||
batch_size,
|
||||
kv_num_heads,
|
||||
past_key.shape[2] + K.shape[1], # Combined sequence length
|
||||
K.shape[2] // kv_num_heads, # head_size
|
||||
)
|
||||
else:
|
||||
present_key_shape = (
|
||||
batch_size,
|
||||
kv_num_heads,
|
||||
K.shape[1], # sequence_length
|
||||
K.shape[2] // kv_num_heads, # head_size
|
||||
)
|
||||
present_value_shape = present_key_shape # Same shape as present_key
|
||||
|
||||
# QK output shape for 3D input (reshaped to 4D internally)
|
||||
qk_output_shape = (
|
||||
batch_size,
|
||||
q_num_heads,
|
||||
q_sequence_length,
|
||||
present_key_shape[2], # kv_sequence_length
|
||||
)
|
||||
else:
|
||||
# 4D input: (batch_size, num_heads, sequence_length, head_size)
|
||||
q_sequence_length = Q.shape[2]
|
||||
# Same shape as Q for 4D output
|
||||
output_shape = Q.shape # type: ignore[assignment]
|
||||
|
||||
# Handle past key/value concatenation
|
||||
if past_key is not None:
|
||||
present_key_shape = (
|
||||
K.shape[0], # batch_size
|
||||
K.shape[1], # num_heads
|
||||
past_key.shape[2] + K.shape[2], # Combined sequence length
|
||||
K.shape[3], # head_size
|
||||
)
|
||||
else:
|
||||
present_key_shape = K.shape # type: ignore[assignment]
|
||||
present_value_shape = present_key_shape # Same shape as present_key
|
||||
|
||||
# QK output shape
|
||||
qk_output_shape = (
|
||||
Q.shape[0], # batch_size
|
||||
Q.shape[1], # q_num_heads
|
||||
Q.shape[2], # q_sequence_length
|
||||
present_key_shape[2], # kv_sequence_length
|
||||
)
|
||||
|
||||
# Create fake tensors with correct shapes and dtypes
|
||||
output = torch.empty(output_shape, dtype=Q.dtype, device=Q.device)
|
||||
present_key = torch.empty(present_key_shape, dtype=K.dtype, device=K.device)
|
||||
present_value = torch.empty(present_value_shape, dtype=V.dtype, device=V.device)
|
||||
qk_output = torch.empty(qk_output_shape, dtype=Q.dtype, device=Q.device)
|
||||
|
||||
return output, present_key, present_value, qk_output
|
||||
|
||||
|
||||
@_onnx_op("Attention", 23, _attention_23_fake_impl)
|
||||
def attention_23(
|
||||
Q: torch.Tensor,
|
||||
K: torch.Tensor,
|
||||
V: torch.Tensor,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
past_key: torch.Tensor | None = None,
|
||||
past_value: torch.Tensor | None = None,
|
||||
*,
|
||||
is_causal: bool = False,
|
||||
kv_num_heads: int = 0,
|
||||
q_num_heads: int = 0,
|
||||
qk_matmul_output_mode: int = 0,
|
||||
scale: float | None = None,
|
||||
softcap: float = 0.0,
|
||||
softmax_precision: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Attention-23 https://onnx.ai/onnx/operators/onnx__Attention.html#attention-23"""
|
||||
|
||||
num_head_dim, sequence_dim, head_dim = 1, 2, 3
|
||||
|
||||
# Store original input shape to determine output shape
|
||||
input_shape_len = len(Q.shape)
|
||||
batch_size = Q.shape[0]
|
||||
|
||||
# Reshape 3D inputs to 4D format
|
||||
if len(Q.shape) == 3:
|
||||
torch._check(
|
||||
q_num_heads != 0 and kv_num_heads != 0,
|
||||
lambda: "q_num_heads and kv_num_heads must be provided for 3D inputs",
|
||||
)
|
||||
q_sequence_length = Q.shape[1]
|
||||
Q = _reshape_3d_to_4d(Q, batch_size, q_num_heads)
|
||||
K = _reshape_3d_to_4d(K, batch_size, kv_num_heads)
|
||||
V = _reshape_3d_to_4d(V, batch_size, kv_num_heads)
|
||||
|
||||
torch._check(
|
||||
len(Q.shape) == 4 and len(K.shape) == 4 and len(V.shape) == 4,
|
||||
lambda: "Q, K, and V should be 4D tensors by now",
|
||||
)
|
||||
|
||||
# Calculate scale factor if not provided
|
||||
q_head_size = Q.shape[head_dim]
|
||||
scale = _get_scale_factor(scale, q_head_size)
|
||||
|
||||
# Handle past key/value caches
|
||||
present_key = (
|
||||
torch.cat([past_key, K], dim=sequence_dim)
|
||||
if past_key is not None
|
||||
else K.clone()
|
||||
)
|
||||
present_value = (
|
||||
torch.cat([past_value, V], dim=sequence_dim)
|
||||
if past_value is not None
|
||||
else V.clone()
|
||||
)
|
||||
|
||||
# Update K and V to include past states
|
||||
K, V = present_key, present_value
|
||||
|
||||
# Get current dimensions
|
||||
current_q_num_heads = Q.shape[num_head_dim]
|
||||
current_kv_num_heads = K.shape[num_head_dim]
|
||||
q_sequence_length = Q.shape[sequence_dim]
|
||||
kv_sequence_length = K.shape[sequence_dim]
|
||||
|
||||
# Check if we can use the optimized scaled_dot_product_attention (most optimized)
|
||||
can_use_sdpa = (
|
||||
softcap == 0.0 # No softcap
|
||||
and qk_matmul_output_mode == 0 # Default QK output mode
|
||||
and softmax_precision is None # No custom softmax precision
|
||||
and (attn_mask is None or attn_mask.dtype == torch.bool)
|
||||
)
|
||||
|
||||
_validate_gqa_configuration(current_q_num_heads, current_kv_num_heads)
|
||||
|
||||
if can_use_sdpa:
|
||||
# Use PyTorch's optimized scaled_dot_product_attention
|
||||
output = torch.nn.functional.scaled_dot_product_attention(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
attn_mask=attn_mask,
|
||||
dropout_p=0.0,
|
||||
is_causal=is_causal,
|
||||
scale=scale,
|
||||
enable_gqa=bool(
|
||||
current_q_num_heads != current_kv_num_heads
|
||||
), # Ensure enable_gqa is not SymBool
|
||||
)
|
||||
|
||||
qk_output = _get_qk_output_for_aten_spda(
|
||||
Q,
|
||||
K,
|
||||
current_q_num_heads,
|
||||
current_kv_num_heads,
|
||||
scale,
|
||||
qk_matmul_output_mode,
|
||||
)
|
||||
else:
|
||||
# Fallback to manual implementation for complex cases
|
||||
|
||||
# Handle Group Query Attention (GQA) and Multi-Query Attention (MQA)
|
||||
if current_q_num_heads != current_kv_num_heads:
|
||||
repeat_factor = current_q_num_heads // current_kv_num_heads
|
||||
K = K.repeat_interleave(repeat_factor, dim=num_head_dim)
|
||||
V = V.repeat_interleave(repeat_factor, dim=num_head_dim)
|
||||
|
||||
# Create attention bias
|
||||
attn_bias = torch.zeros(
|
||||
q_sequence_length, kv_sequence_length, dtype=Q.dtype, device=Q.device
|
||||
)
|
||||
|
||||
# Apply causal masking
|
||||
if is_causal:
|
||||
torch._check(
|
||||
attn_mask is None, lambda: "Cannot use both is_causal and attn_mask"
|
||||
)
|
||||
causal_mask = torch.tril(
|
||||
torch.ones(
|
||||
q_sequence_length,
|
||||
kv_sequence_length,
|
||||
dtype=torch.bool,
|
||||
device=Q.device,
|
||||
)
|
||||
)
|
||||
attn_bias = attn_bias.masked_fill(~causal_mask, float("-inf"))
|
||||
|
||||
# Apply attention mask
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
# Boolean mask: True means participate in attention
|
||||
attn_bias = attn_bias.masked_fill(~attn_mask, float("-inf"))
|
||||
else:
|
||||
# Float mask: added to attention scores
|
||||
attn_bias = attn_bias + attn_mask
|
||||
|
||||
# Apply scaling factor
|
||||
scale_factor = _get_scale_factor(scale, Q.shape[3])
|
||||
|
||||
# Scale both Q and K by sqrt(scale_factor) for numerical stability
|
||||
sqrt_scale = math.sqrt(scale_factor)
|
||||
Q_scaled = Q * sqrt_scale
|
||||
K_scaled = K * sqrt_scale
|
||||
|
||||
# Compute Q @ K^T
|
||||
qk_matmul_output = torch.matmul(Q_scaled, K_scaled.transpose(-2, -1))
|
||||
|
||||
# Initialize QK output based on mode
|
||||
qk_output = qk_matmul_output # Default case for mode 0
|
||||
|
||||
# Add attention bias
|
||||
qk_with_bias = qk_matmul_output + attn_bias
|
||||
|
||||
if qk_matmul_output_mode == 1:
|
||||
qk_output = qk_with_bias
|
||||
|
||||
# Apply softcap if provided
|
||||
if softcap > 0.0:
|
||||
qk_with_bias = softcap * torch.tanh(qk_with_bias / softcap)
|
||||
|
||||
if qk_matmul_output_mode == 2:
|
||||
qk_output = qk_with_bias
|
||||
|
||||
# Apply softmax with optional precision casting
|
||||
if softmax_precision is not None:
|
||||
# Map ONNX data type to torch dtype
|
||||
if softmax_precision in _ATTENTION_23_ALLOWED_INTERMEDIATE_PRECISIONS:
|
||||
original_dtype = qk_with_bias.dtype
|
||||
qk_with_bias = qk_with_bias.to(
|
||||
_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[softmax_precision]
|
||||
)
|
||||
qk_softmax = torch.softmax(qk_with_bias, dim=-1)
|
||||
qk_softmax = qk_softmax.to(original_dtype)
|
||||
else:
|
||||
qk_softmax = torch.softmax(qk_with_bias, dim=-1)
|
||||
else:
|
||||
qk_softmax = torch.softmax(qk_with_bias, dim=-1)
|
||||
|
||||
if qk_matmul_output_mode == 3:
|
||||
qk_output = qk_softmax
|
||||
|
||||
# Compute attention output
|
||||
output = torch.matmul(qk_softmax, V)
|
||||
|
||||
# Reshape output back to 3D if input was 3D
|
||||
if input_shape_len == 3:
|
||||
# output: (batch_size, q_num_heads, q_sequence_length, v_head_size) -> (batch_size, q_sequence_length, hidden_size)
|
||||
output = (
|
||||
output.transpose(1, 2).contiguous().view(batch_size, q_sequence_length, -1)
|
||||
)
|
||||
|
||||
return output, present_key, present_value, qk_output
|
||||
@@ -0,0 +1,304 @@
|
||||
"""Implementation of symbolic FX ops to represent arbitrary ONNX ops.
|
||||
|
||||
This module provides a way to create symbolic FX operators that can represent
|
||||
arbitrary ONNX operators.
|
||||
|
||||
The operators are called "symbolic" because they don't do any actual computation
|
||||
but instead serve as placeholders in the computation graph.
|
||||
|
||||
Each implementation contains two parts: A "real" implementation that produce all
|
||||
zeros based on the input shape and dtype, and a "fake" implementation that does more
|
||||
or less the same thing but is required by the `torch.library.custom_op` interface.
|
||||
"""
|
||||
|
||||
# flake8: noqa: B950
|
||||
import dataclasses
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch
|
||||
from torch.onnx.ops import _dtype_mappings
|
||||
|
||||
|
||||
_INT_TYPE = "i"
|
||||
_FLOAT_TYPE = "f"
|
||||
_STRING_TYPE = "s"
|
||||
_INT_SEQ_TYPE = "is"
|
||||
_FLOAT_SEQ_TYPE = "fs"
|
||||
_STRING_SEQ_TYPE = "ss"
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class EncodedAttrs:
|
||||
"""Class to encode attributes from dictionary into lists of FX compatible attributes.
|
||||
|
||||
Since FX does not support dictionaries, we need to encode the attributes into
|
||||
lists. This class provides a way to encode and decode the attributes.
|
||||
|
||||
Attributes:
|
||||
attr_keys: List of attribute keys.
|
||||
attr_types: List of attribute types. Values can be "i" (int), "f" (float),
|
||||
"s" (string), "is" (int sequence), "fs" (float sequence), or "ss" (string sequence).
|
||||
attr_pos: List of tuples representing the start and end positions of each
|
||||
attribute in the corresponding list.
|
||||
attr_ints: List of integer attributes.
|
||||
attr_floats: List of float attributes.
|
||||
attr_strs: List of string attributes.
|
||||
"""
|
||||
|
||||
attr_keys: list[str]
|
||||
attr_types: list[str]
|
||||
attr_pos: list[tuple[int, int]]
|
||||
attr_ints: list[int]
|
||||
attr_floats: list[float]
|
||||
attr_strs: list[str]
|
||||
|
||||
@classmethod
|
||||
def from_dict(
|
||||
cls,
|
||||
attrs: dict[
|
||||
str,
|
||||
int
|
||||
| float
|
||||
| str
|
||||
| bool
|
||||
| Sequence[int]
|
||||
| Sequence[float]
|
||||
| Sequence[str]
|
||||
| Sequence[bool],
|
||||
],
|
||||
) -> "EncodedAttrs":
|
||||
encoded = cls(
|
||||
attr_keys=[],
|
||||
attr_types=[],
|
||||
attr_pos=[],
|
||||
attr_ints=[],
|
||||
attr_floats=[],
|
||||
attr_strs=[],
|
||||
)
|
||||
for k, v in attrs.items():
|
||||
encoded.attr_keys.append(k)
|
||||
if isinstance(v, int):
|
||||
start_pos = len(encoded.attr_ints)
|
||||
encoded.attr_ints.append(v)
|
||||
encoded.attr_pos.append((start_pos, start_pos + 1))
|
||||
encoded.attr_types.append(_INT_TYPE)
|
||||
elif isinstance(v, float):
|
||||
start_pos = len(encoded.attr_floats)
|
||||
encoded.attr_floats.append(v)
|
||||
encoded.attr_pos.append((start_pos, start_pos + 1))
|
||||
encoded.attr_types.append(_FLOAT_TYPE)
|
||||
elif isinstance(v, str):
|
||||
start_pos = len(encoded.attr_strs)
|
||||
encoded.attr_strs.append(v)
|
||||
encoded.attr_pos.append((start_pos, start_pos + 1))
|
||||
encoded.attr_types.append(_STRING_TYPE)
|
||||
elif isinstance(v, Sequence):
|
||||
if len(v) == 0:
|
||||
raise ValueError(f"Empty sequence for attribute {k}")
|
||||
if any(isinstance(elem, float) for elem in v):
|
||||
start_pos = len(encoded.attr_floats)
|
||||
encoded.attr_floats.extend([float(elem) for elem in v])
|
||||
encoded.attr_pos.append((start_pos, start_pos + len(v)))
|
||||
encoded.attr_types.append(_FLOAT_SEQ_TYPE)
|
||||
elif isinstance(v[0], int):
|
||||
start_pos = len(encoded.attr_ints)
|
||||
encoded.attr_ints.extend([int(elem) for elem in v])
|
||||
encoded.attr_pos.append((start_pos, start_pos + len(v)))
|
||||
encoded.attr_types.append(_INT_SEQ_TYPE)
|
||||
elif isinstance(v[0], str):
|
||||
start_pos = len(encoded.attr_strs)
|
||||
encoded.attr_strs.extend([str(elem) for elem in v])
|
||||
encoded.attr_pos.append((start_pos, start_pos + len(v)))
|
||||
encoded.attr_types.append(_STRING_SEQ_TYPE)
|
||||
else:
|
||||
raise ValueError(f"Unsupported sequence type for attribute {k}")
|
||||
else:
|
||||
raise ValueError(f"Unsupported attribute type for {k}: {type(v)}")
|
||||
if len(encoded.attr_keys) != len(encoded.attr_types):
|
||||
raise AssertionError(
|
||||
f"Mismatch between number of attribute keys and types: {len(encoded.attr_keys)} != {len(encoded.attr_types)}"
|
||||
)
|
||||
if len(encoded.attr_keys) != len(encoded.attr_pos):
|
||||
raise AssertionError(
|
||||
f"Mismatch between number of attribute keys and positions: {len(encoded.attr_keys)} != {len(encoded.attr_pos)}"
|
||||
)
|
||||
return encoded
|
||||
|
||||
def to_dict(
|
||||
self,
|
||||
) -> dict[
|
||||
str,
|
||||
int | float | str | list[int] | list[float] | list[str],
|
||||
]:
|
||||
"""Convert the encoded attributes back to a dictionary for creating an ONNX node."""
|
||||
attrs: dict[
|
||||
str,
|
||||
int | float | str | list[int] | list[float] | list[str],
|
||||
] = {}
|
||||
for i, key in enumerate(self.attr_keys):
|
||||
attr_type = self.attr_types[i]
|
||||
if attr_type == _INT_TYPE:
|
||||
attrs[key] = self.attr_ints[self.attr_pos[i][0]]
|
||||
elif attr_type == _FLOAT_TYPE:
|
||||
attrs[key] = self.attr_floats[self.attr_pos[i][0]]
|
||||
elif attr_type == _STRING_TYPE:
|
||||
attrs[key] = self.attr_strs[self.attr_pos[i][0]]
|
||||
elif attr_type == _FLOAT_SEQ_TYPE:
|
||||
attrs[key] = self.attr_floats[self.attr_pos[i][0] : self.attr_pos[i][1]]
|
||||
elif attr_type == _INT_SEQ_TYPE:
|
||||
attrs[key] = self.attr_ints[self.attr_pos[i][0] : self.attr_pos[i][1]]
|
||||
elif attr_type == _STRING_SEQ_TYPE:
|
||||
attrs[key] = self.attr_strs[self.attr_pos[i][0] : self.attr_pos[i][1]]
|
||||
else:
|
||||
raise ValueError(f"Unsupported attribute type: {attr_type}")
|
||||
return attrs
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"onnx_symbolic::_symbolic",
|
||||
mutates_args=(),
|
||||
schema=(
|
||||
"(Tensor?[] inputs, str op_type, int onnx_dtype, *,"
|
||||
" SymInt[] shape, str[] attr_keys, str[] attr_types, int[][] attr_pos,"
|
||||
" int[] attr_ints, float[] attr_floats, str[] attr_strs, str[] metadata_props_keys,"
|
||||
" str[] metadata_props_values, str domain='', int? version=None"
|
||||
") -> Tensor"
|
||||
),
|
||||
)
|
||||
def _symbolic(
|
||||
inputs: Sequence[torch.Tensor | None],
|
||||
op_type: str,
|
||||
onnx_dtype: int,
|
||||
*,
|
||||
shape: Sequence[int | torch.SymInt],
|
||||
attr_keys: Sequence[str],
|
||||
attr_types: Sequence[str],
|
||||
attr_pos: Sequence[tuple[int, int]],
|
||||
attr_ints: Sequence[int],
|
||||
attr_floats: Sequence[float],
|
||||
attr_strs: Sequence[str],
|
||||
metadata_props_keys: Sequence[str] = (),
|
||||
metadata_props_values: Sequence[str] = (),
|
||||
domain: str = "",
|
||||
version: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
torch._check(
|
||||
onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE,
|
||||
lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}",
|
||||
)
|
||||
return torch.zeros(
|
||||
shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype]
|
||||
)
|
||||
|
||||
|
||||
@_symbolic.register_fake
|
||||
def _(
|
||||
inputs: Sequence[torch.Tensor],
|
||||
op_type: str,
|
||||
onnx_dtype: int,
|
||||
*,
|
||||
shape: Sequence[int | torch.SymInt],
|
||||
attr_keys: Sequence[str],
|
||||
attr_types: Sequence[str],
|
||||
attr_pos: Sequence[tuple[int, int]],
|
||||
attr_ints: Sequence[int],
|
||||
attr_floats: Sequence[float],
|
||||
attr_strs: Sequence[str],
|
||||
metadata_props_keys: Sequence[str] = (),
|
||||
metadata_props_values: Sequence[str] = (),
|
||||
domain: str = "",
|
||||
version: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
torch._check(
|
||||
onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE,
|
||||
lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}",
|
||||
)
|
||||
# NOTE(justinchuby): Use zeros instead of torch.empty because I haven't figured
|
||||
# out how it can handle empty shapes
|
||||
return torch.zeros(
|
||||
shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype]
|
||||
)
|
||||
|
||||
|
||||
@torch.library.custom_op(
|
||||
"onnx_symbolic::_symbolic_multi_out",
|
||||
mutates_args=(),
|
||||
schema=(
|
||||
"(Tensor?[] inputs, str op_type, int[] onnx_dtypes, *,"
|
||||
" SymInt[][] shapes, str[] attr_keys, str[] attr_types, int[][] attr_pos,"
|
||||
" int[] attr_ints, float[] attr_floats, str[] attr_strs, str[] metadata_props_keys,"
|
||||
" str[] metadata_props_values, str domain='', int? version=None"
|
||||
") -> Tensor[]"
|
||||
),
|
||||
)
|
||||
def _symbolic_multi_out(
|
||||
inputs: Sequence[torch.Tensor | None],
|
||||
op_type: str,
|
||||
onnx_dtypes: Sequence[int],
|
||||
*,
|
||||
shapes: Sequence[Sequence[int | torch.SymInt]],
|
||||
attr_keys: Sequence[str],
|
||||
attr_types: Sequence[str],
|
||||
attr_pos: Sequence[tuple[int, int]],
|
||||
attr_ints: Sequence[int],
|
||||
attr_floats: Sequence[float],
|
||||
attr_strs: Sequence[str],
|
||||
metadata_props_keys: Sequence[str] = (),
|
||||
metadata_props_values: Sequence[str] = (),
|
||||
domain: str = "",
|
||||
version: int | None = None,
|
||||
) -> list[torch.Tensor]:
|
||||
outputs = []
|
||||
torch._check(
|
||||
len(shapes) == len(onnx_dtypes),
|
||||
lambda: f"Number of shapes ({len(shapes)}) must match number of ONNX dtypes ({len(onnx_dtypes)})",
|
||||
)
|
||||
for shape, onnx_dtype in zip(shapes, onnx_dtypes):
|
||||
torch._check(
|
||||
onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE,
|
||||
lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}",
|
||||
)
|
||||
outputs.append(
|
||||
torch.zeros(
|
||||
shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype]
|
||||
)
|
||||
)
|
||||
return outputs
|
||||
|
||||
|
||||
@_symbolic_multi_out.register_fake
|
||||
def _(
|
||||
inputs: Sequence[torch.Tensor],
|
||||
op_type: str,
|
||||
onnx_dtypes: Sequence[int],
|
||||
*,
|
||||
shapes: Sequence[Sequence[int | torch.SymInt]],
|
||||
attr_keys: Sequence[str],
|
||||
attr_types: Sequence[str],
|
||||
attr_pos: Sequence[tuple[int, int]],
|
||||
attr_ints: Sequence[int],
|
||||
attr_floats: Sequence[float],
|
||||
attr_strs: Sequence[str],
|
||||
metadata_props_keys: Sequence[str] = (),
|
||||
metadata_props_values: Sequence[str] = (),
|
||||
domain: str = "",
|
||||
version: int | None = None,
|
||||
) -> list[torch.Tensor]:
|
||||
outputs = []
|
||||
torch._check(
|
||||
len(shapes) == len(onnx_dtypes),
|
||||
lambda: f"Number of shapes ({len(shapes)}) must match number of ONNX dtypes ({len(onnx_dtypes)})",
|
||||
)
|
||||
for shape, onnx_dtype in zip(shapes, onnx_dtypes):
|
||||
torch._check(
|
||||
onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE,
|
||||
lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}",
|
||||
)
|
||||
# NOTE(justinchuby): Use zeros instead of torch.empty because I haven't figured
|
||||
# out how it can handle empty shapes
|
||||
outputs.append(
|
||||
torch.zeros(
|
||||
shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype]
|
||||
)
|
||||
)
|
||||
return outputs
|
||||
Reference in New Issue
Block a user