Add XAttention reference operation

[vshampor]

Details:

• Implements the XAttention operation using STL and OV reference operations.
• Meant to be used for testing and debugging purposes only - does not bring performance benefits if used on its own. Actual HW-accelerated implementation would be using this code as reference.
• [DO NOT MERGE] Example of the XAttention integration in the CPU plugin #31955 contains an example of integrating with the CPU plugin.

Tickets:

• CVS-170365

[praasz]

The reference tests should be stored in:
src/plugins/template/tests/functional/op_reference/

as others.

[vshampor]

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.

[praasz]

Remove these assert. the shapes and input validation is done during shape inference in operator implementation.

[vshampor]

Not relevant since there is no shape inference involved by design

[praasz]

Why is required as it is just call for ov::reference::softmax(reshaped_qk_product_data, out, reshaped_qk_product_shape, {2});

[vshampor]

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.

[l-bat]

Suggested change

	/** 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]

Done

[WeldonWangwang]

Hi @vshampor It looks like this implementation selects blocks by globally sorting all blocks per head and accumulating their scores until the threshold is reached. In the X-Attention paper (Section 2.2, https://arxiv.org/abs/2503.16428
), the algorithm selects blocks per query block based on antidiagonal scoring. Could this global accumulation approach lead to differences from the paper’s intended block coverage for each query(https://github.com/mit-han-lab/x-attention/blob/e37988770b9d1bebd489eba011d615f35587ba08/xattn/src/utils.py#L44)?

[vshampor]

Please consider the relative position of this function in the frontend select_blocks function and observe that the block selection is ultimately done among the antidiagonal subset of the full attention matrix. I didn't understand the concerns about globality - in which sense do you consider the sorting to be global and why does it not match to the original algo?

[WeldonWangwang]

Simplify the problem, I use https://github.com/mit-han-lab/x-attention/tree/main?tab=readme -Use your test case data in ov file # quick use with your test date in

openvino/src/core/tests/reference/xattention.cpp

Line 27 in 231ae69

{2, 4, 4},

The selected blocks are inconsistent:

THRESHOLD = 0.8
BLOCK_SIZE = 2
STRIDE = 2

# batch=1, head=2, seq_len=4, dim=4
bsz = 1
heads = 2
seq_len = 4
dim = 4

q_data = [
     3.144,  8.512,  8.518, -8.386,
     7.889, -5.721,  5.507,  4.295,
    -6.624, -8.463,  7.474,  9.879,
     4.534, -5.908, -9.388,  2.356,

     7.497,  8.186, -8.658, -4.796,
    -8.248, -9.797, -7.907, -4.513,
     3.469,  7.633,  7.244, -6.844,
    -7.173,  4.450,  6.705, -7.035
]

q = torch.tensor(q_data, dtype=torch.bfloat16).reshape(bsz, heads, seq_len, dim).to("cuda")
k = q.clone()  # K = Q
v = torch.randn((bsz, heads, seq_len, dim), dtype=torch.bfloat16).to("cuda")

attention_output = Xattention_prefill(
    query_states=q,
    key_states=k,
    value_states=v,
    stride=STRIDE,
    block_size=BLOCK_SIZE,
    use_triton=True,
    chunk_size=4
)

your code:
Head 0 selected blocks: (0,1) (1,0)
Head 1 selected blocks: (0,1) (1,0)
mit-han-lab/x-attention:
Head 0 selected blocks: (0,0) (1,0) (1,1)
Head 1 selected blocks: (0,0) (1,0) (1,1)
So may I know what is the source of the reference results for your test?