Add XAttention reference operation #31864
Conversation
github-actions
bot
added
category: Core
vshampor
force-pushed
the
xattn_reference
branch
2 times, most recently
from
9a119ee to
e29e853
Compare
vshampor
force-pushed
the
xattn_reference
branch
from
894d8bc to
7e88152
Compare
vshampor
force-pushed
the
xattn_reference
branch
2 times, most recently
from
460a365 to
6be527b
Compare
vshampor
force-pushed
the
xattn_reference
branch
from
6be527b to
1113f02
Compare
There was a problem hiding this comment.
The reference tests should be stored in:
src/plugins/template/tests/functional/op_reference/
as others.
There was a problem hiding this comment.
I am aware of these tests. All of these use utilize unnecessary abstractions of OV opset ops and involve the sample plugin to do unit tests. Seeing that there will never be a counterpart to this operation in the OV opset, there is no reason to proliferate the bad design decision to have all reference ops tested through the template plugin just for the sake of consistency.
| OPENVINO_ASSERT(input_shape.size() == 3); // [num_heads, num_tokens, head_size] | ||
| OPENVINO_ASSERT(out_shape.size() == 3); | ||
| OPENVINO_ASSERT(input_shape[0] == out_shape[0]); | ||
| OPENVINO_ASSERT(input_shape[1] % m_stride == 0); | ||
| OPENVINO_ASSERT(input_shape[1] / m_stride == out_shape[1]); | ||
| OPENVINO_ASSERT(input_shape[2] * m_stride == out_shape[2]); |
There was a problem hiding this comment.
Remove these assert. the shapes and input validation is done during shape inference in operator implementation.
There was a problem hiding this comment.
Not relevant since there is no shape inference involved by design
| * @param out_shape Shape of the output tensor data. Expected shape is strictly equal to | ||
| * `reshaped_qk_product_shape`. | ||
| */ | ||
| void softmax(const T* reshaped_qk_product_data, |
There was a problem hiding this comment.
Why is required as it is just call for ov::reference::softmax(reshaped_qk_product_data, out, reshaped_qk_product_shape, {2});
There was a problem hiding this comment.
This is a reference operation. It is designed to be understandable and testable first and foremost. I provide a set of functions that each directly map to a phase or sub-phase of the intended HW-accelerated kernel. Softmax is one of these phases and therefore has a separate function. It also has an extra shape check, so having a function wrapper has some utility.
| } | ||
| } | ||
|
|
||
| /** Selects the elements of the input tensor along the last two dimensions, indepently along the first dimension, so |
There was a problem hiding this comment.
| /** Selects the elements of the input tensor along the last two dimensions, indepently along the first dimension, so | |
| /** Selects the elements of the input tensor along the last two dimensions, independently along the first dimension, so |
vshampor
force-pushed
the
xattn_reference
branch
from
4042108 to
2295ca6
Compare
vshampor
force-pushed
the
xattn_reference
branch
from
22b057f to
aa28325
Compare
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
Details:
Tickets: