From 4a750ac30b52794e4d574c089c7047fa82713e32 Mon Sep 17 00:00:00 2001 From: Li Baoming <1508269885@qq.com> Date: Wed, 2 Sep 2026 07:15:56 +0000 Subject: [PATCH] feat(cambricon): add flash attention KV-cache provider --- src/common/op_utils/paged_kv_cache.h | 44 ++++ src/linked/torch/cambricon/flash_attn.yaml | 2 + .../ops/flash_attn_with_kvcache/flash_attn.cc | 186 ++++++++++++++ .../ops/flash_attn_with_kvcache/flash_attn.h | 49 ++++ .../flash_attn_with_kvcache/flash_attn.yaml | 4 + tests/test_flash_attn_with_kvcache.py | 229 ++++++++++++++---- 6 files changed, 471 insertions(+), 43 deletions(-) create mode 100644 src/common/op_utils/paged_kv_cache.h create mode 100644 src/linked/torch/cambricon/flash_attn.yaml create mode 100644 src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.cc create mode 100644 src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.h create mode 100644 src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.yaml diff --git a/src/common/op_utils/paged_kv_cache.h b/src/common/op_utils/paged_kv_cache.h new file mode 100644 index 000000000..47fc54793 --- /dev/null +++ b/src/common/op_utils/paged_kv_cache.h @@ -0,0 +1,44 @@ +#ifndef INFINI_OPS_COMMON_OP_UTILS_PAGED_KV_CACHE_H_ +#define INFINI_OPS_COMMON_OP_UTILS_PAGED_KV_CACHE_H_ + +#include + +#include +#include +#include + +namespace infini::ops::paged_kv_cache { + +// These helpers intentionally operate on ATen tensors because linked Torch +// FlashAttention providers expose paged KV caches through ATen. Keep their use +// within linked Torch providers until a backend-independent abstraction is +// required. + +inline std::vector ToHostInt32Vector(const at::Tensor& tensor) { + const auto cpu = tensor.to(at::kCPU).contiguous(); + const auto* data = cpu.data_ptr(); + return {data, data + cpu.numel()}; +} + +inline at::Tensor GatherSequence(const at::Tensor& cache, + const std::vector& block_table, + int64_t table_width, int64_t batch, + int64_t length) { + if (length == 0) { + return cache.new_empty({0, cache.size(2), cache.size(3)}); + } + + const int64_t page_size = cache.size(1); + const int64_t block_count = (length + page_size - 1) / page_size; + std::vector pages; + pages.reserve(block_count); + for (int64_t i = 0; i < block_count; ++i) { + assert(i < table_width && "KV cache block table is too small"); + pages.push_back(cache.select(0, block_table[batch * table_width + i])); + } + return at::cat(pages, 0).slice(0, 0, length); +} + +} // namespace infini::ops::paged_kv_cache + +#endif // INFINI_OPS_COMMON_OP_UTILS_PAGED_KV_CACHE_H_ diff --git a/src/linked/torch/cambricon/flash_attn.yaml b/src/linked/torch/cambricon/flash_attn.yaml new file mode 100644 index 000000000..65ab869fc --- /dev/null +++ b/src/linked/torch/cambricon/flash_attn.yaml @@ -0,0 +1,2 @@ +python_distribution_package: flash-attn +library_glob: flash_attn_2_bang*.so diff --git a/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.cc b/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.cc new file mode 100644 index 000000000..d4eddfe17 --- /dev/null +++ b/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.cc @@ -0,0 +1,186 @@ +#include "linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.h" + +#include + +#include +#include +#include + +#include "common/op_utils/paged_kv_cache.h" + +std::vector mha_varlen_fwd( + const at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& alibi_slopes, int max_seqlen_q, int max_seqlen_k, + float dropout_p, float softmax_scale, bool zero_tensors, bool causal, + int window_size_left, int window_size_right, bool return_softmax, + std::optional generator); + +namespace infini::ops::linked::torch::cambricon { +namespace { + +namespace paged_kv_cache = ::infini::ops::paged_kv_cache; +using paged_kv_cache::ToHostInt32Vector; + +at::Tensor ToDeviceIndices(const std::vector& values, + const at::Device& device) { + auto cpu = at::from_blob(const_cast(values.data()), + {static_cast(values.size())}, + at::TensorOptions().dtype(at::kInt)) + .clone(); + return cpu.to(device); +} + +void UpdatePagedCache(const at::Tensor& cache, const at::Tensor& values, + const std::vector& block_table, + int64_t table_width, int64_t batch, int64_t offset) { + const int64_t page_size = cache.size(1); + int64_t source_offset = 0; + while (source_offset < values.size(1)) { + const int64_t logical_offset = offset + source_offset; + const int64_t table_column = logical_offset / page_size; + const int64_t page_offset = logical_offset % page_size; + assert(table_column < table_width && "KV cache block table is too small"); + const int64_t block = block_table[batch * table_width + table_column]; + const int64_t count = + std::min(page_size - page_offset, values.size(1) - source_offset); + cache.select(0, block) + .slice(0, page_offset, page_offset + count) + .copy_(values.select(0, batch).slice(0, source_offset, + source_offset + count)); + source_offset += count; + } +} + +} // namespace + +std::vector FlashAttnKvcache::Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits) { + assert(!rotary_cos.has_value() && !rotary_sin.has_value() && + "Cambricon KV-cache attention does not yet support rotary tables"); + assert(!cache_leftpad.has_value() && + "Cambricon KV-cache attention does not yet support left padding"); + assert(softcap == 0.0f && + "Cambricon KV-cache attention does not support softcap"); + + const int64_t batch_size = q.size(0); + const int64_t query_length = q.size(1); + const int64_t num_heads = q.size(2); + const int64_t head_size = q.size(3); + const bool paged = block_table.has_value(); + const int64_t append_length = k.has_value() ? k->size(1) : 0; + assert(!k.has_value() || cache_seqlens.has_value()); + assert(!paged || !cache_batch_idx.has_value()); + + std::vector lengths( + batch_size, static_cast(paged ? 0 : k_cache.size(1))); + if (cache_seqlens.has_value()) { + lengths = ToHostInt32Vector(*cache_seqlens); + } + std::vector cache_rows(batch_size); + for (int64_t i = 0; i < batch_size; ++i) { + cache_rows[i] = static_cast(i); + } + if (cache_batch_idx.has_value()) { + cache_rows = ToHostInt32Vector(*cache_batch_idx); + } + + std::vector host_block_table; + int64_t table_width = 0; + if (paged) { + host_block_table = ToHostInt32Vector(*block_table); + table_width = block_table->size(1); + } + + if (k.has_value()) { + for (int64_t batch = 0; batch < batch_size; ++batch) { + if (paged) { + UpdatePagedCache(k_cache, *k, host_block_table, table_width, batch, + lengths[batch]); + UpdatePagedCache(v_cache, *v, host_block_table, table_width, batch, + lengths[batch]); + } else { + const int64_t row = cache_rows[batch]; + k_cache.select(0, row) + .slice(0, lengths[batch], lengths[batch] + append_length) + .copy_(k->select(0, batch)); + v_cache.select(0, row) + .slice(0, lengths[batch], lengths[batch] + append_length) + .copy_(v->select(0, batch)); + } + } + } + + std::vector packed_k; + std::vector packed_v; + std::vector cu_seqlens_q{0}; + std::vector cu_seqlens_k{0}; + int64_t max_key_length = 0; + packed_k.reserve(batch_size); + packed_v.reserve(batch_size); + for (int64_t batch = 0; batch < batch_size; ++batch) { + const int64_t length = lengths[batch] + append_length; + assert(length > 0 && "KV-cache attention requires a non-empty cache"); + max_key_length = std::max(max_key_length, length); + cu_seqlens_q.push_back(cu_seqlens_q.back() + query_length); + cu_seqlens_k.push_back(cu_seqlens_k.back() + length); + if (paged) { + packed_k.push_back(paged_kv_cache::GatherSequence( + k_cache, host_block_table, table_width, batch, length)); + packed_v.push_back(paged_kv_cache::GatherSequence( + v_cache, host_block_table, table_width, batch, length)); + } else { + const int64_t row = cache_rows[batch]; + packed_k.push_back(k_cache.select(0, row).slice(0, 0, length)); + packed_v.push_back(v_cache.select(0, row).slice(0, 0, length)); + } + } + + auto packed_q = + q.contiguous().view({batch_size * query_length, num_heads, head_size}); + auto at_cu_seqlens_q = ToDeviceIndices(cu_seqlens_q, q.device()); + auto at_cu_seqlens_k = ToDeviceIndices(cu_seqlens_k, q.device()); + auto at_packed_k = at::cat(packed_k, 0).contiguous(); + auto at_packed_v = at::cat(packed_v, 0).contiguous(); + std::optional packed_out; + std::optional seqused_k; + std::optional generator; + const auto result = ::mha_varlen_fwd( + packed_q, at_packed_k, at_packed_v, packed_out, at_cu_seqlens_q, + at_cu_seqlens_k, seqused_k, alibi_slopes, static_cast(query_length), + static_cast(max_key_length), 0.0f, softmax_scale, false, causal, + window_size_left, window_size_right, false, generator); + assert(result.size() > 5 && + "Cambricon FlashAttention returned an incomplete result"); + + auto result_out = result[0].view(q.sizes()); + auto result_lse = result[5] + .view({num_heads, batch_size, query_length}) + .permute({1, 0, 2}) + .contiguous(); + (void)out; + (void)rotary_interleaved; + (void)num_splits; + return {result_out, result_lse}; +} + +} // namespace infini::ops::linked::torch::cambricon + +namespace infini::ops::linked::torch { + +template class TorchFlashAttnWithKvcache< + ::infini::ops::linked::torch::cambricon::FlashAttnKvcache>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.h b/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.h new file mode 100644 index 000000000..3c235371d --- /dev/null +++ b/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.h @@ -0,0 +1,49 @@ +#ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ +#define INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ + +#include "linked/torch/ops/flash_attn_with_kvcache.h" +#include "torch/cambricon/c10.h" + +namespace infini::ops::linked::torch::cambricon { + +struct FlashAttnKvcache : C10 { + static std::vector Call( + at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, + std::optional& k, std::optional& v, + std::optional& cache_seqlens, + std::optional& rotary_cos, + std::optional& rotary_sin, + std::optional& cache_batch_idx, + std::optional& cache_leftpad, + std::optional& block_table, + std::optional& alibi_slopes, std::optional& out, + float softmax_scale, bool causal, int window_size_left, + int window_size_right, float softcap, bool rotary_interleaved, + int num_splits); +}; + +} // namespace infini::ops::linked::torch::cambricon + +namespace infini::ops::linked::torch { + +extern template class TorchFlashAttnWithKvcache< + ::infini::ops::linked::torch::cambricon::FlashAttnKvcache>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchFlashAttnWithKvcache< + linked::torch::cambricon::FlashAttnKvcache> { + public: + using linked::torch::TorchFlashAttnWithKvcache< + linked::torch::cambricon::FlashAttnKvcache>::TorchFlashAttnWithKvcache; + using linked::torch::TorchFlashAttnWithKvcache< + linked::torch::cambricon::FlashAttnKvcache>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_ diff --git a/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.yaml b/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.yaml new file mode 100644 index 000000000..f45ca5758 --- /dev/null +++ b/src/linked/torch/cambricon/ops/flash_attn_with_kvcache/flash_attn.yaml @@ -0,0 +1,4 @@ +library: flash_attn +required_symbols: + - >- + mha_varlen_fwd(at::Tensor const&, at::Tensor const&, at::Tensor const&, std::optional&, at::Tensor const&, at::Tensor const&, std::optional&, std::optional&, int, int, float, float, bool, bool, int, int, bool, std::optional) diff --git a/tests/test_flash_attn_with_kvcache.py b/tests/test_flash_attn_with_kvcache.py index 8c885548d..366f1c795 100644 --- a/tests/test_flash_attn_with_kvcache.py +++ b/tests/test_flash_attn_with_kvcache.py @@ -1,3 +1,5 @@ +import math + import infini.ops import pytest import torch @@ -33,8 +35,8 @@ def test_flash_attn_with_kvcache_dense( rtol, atol, ): - if device != "cuda": - pytest.skip("FlashAttention FA2 requires the NVIDIA backend") + if device not in ("cuda", "mlu"): + pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") batch_size, cache_size = 2, 16 num_heads, num_kv_heads, head_size = 4, 2, 64 @@ -64,19 +66,32 @@ def test_flash_attn_with_kvcache_dense( expected_v_cache = v_cache.clone() actual_k_cache = k_cache.clone() actual_v_cache = v_cache.clone() - expected, expected_softmax_lse = flash_attn.flash_attn_with_kvcache( - q, - expected_k_cache, - expected_v_cache, - k, - v, - cache_seqlens=cache_seqlens, - softmax_scale=0.125, - causal=True, - window_size=(4, 0), - num_splits=1, - return_softmax_lse=True, - ) + if device == "mlu": + expected, expected_softmax_lse = _reference_flash_attn_with_kvcache( + q, + expected_k_cache, + expected_v_cache, + k, + v, + cache_seqlens=cache_seqlens, + softmax_scale=0.125, + causal=True, + window_size=(4, 0), + ) + else: + expected, expected_softmax_lse = flash_attn.flash_attn_with_kvcache( + q, + expected_k_cache, + expected_v_cache, + k, + v, + cache_seqlens=cache_seqlens, + softmax_scale=0.125, + causal=True, + window_size=(4, 0), + num_splits=1, + return_softmax_lse=True, + ) actual = torch.empty_like(q) actual_softmax_lse = torch.empty( (q.size(0), q.size(2), q.size(1)), @@ -122,8 +137,8 @@ def test_flash_attn_with_kvcache_dense( def test_flash_attn_with_kvcache_paged(device, implementation_index): - if device != "cuda": - pytest.skip("FlashAttention FA2 requires the NVIDIA backend") + if device not in ("cuda", "mlu"): + pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") batch_size, page_size = 2, 256 num_heads, num_kv_heads, head_size = 4, 2, 64 @@ -140,14 +155,24 @@ def test_flash_attn_with_kvcache_paged(device, implementation_index): v_cache = torch.randn_like(k_cache) cache_seqlens = torch.tensor((130, 300), dtype=torch.int32, device=device) block_table = torch.tensor(((0, 1), (2, 3)), dtype=torch.int32, device=device) - expected = flash_attn.flash_attn_with_kvcache( - q, - k_cache, - v_cache, - cache_seqlens=cache_seqlens, - block_table=block_table, - causal=True, - ) + if device == "mlu": + expected, _ = _reference_flash_attn_with_kvcache( + q, + k_cache, + v_cache, + cache_seqlens=cache_seqlens, + block_table=block_table, + causal=True, + ) + else: + expected = flash_attn.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + cache_seqlens=cache_seqlens, + block_table=block_table, + causal=True, + ) actual = torch.empty_like(q) infini.ops.flash_attn_with_kvcache( @@ -182,8 +207,8 @@ def test_flash_attn_with_kvcache_paged(device, implementation_index): def test_flash_attn_with_kvcache_scalar_seqlens_with_cache_batch_idx( device, implementation_index ): - if device != "cuda": - pytest.skip("FlashAttention FA2 requires the NVIDIA backend") + if device not in ("cuda", "mlu"): + pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") q = torch.randn((2, 1, 4, 60), dtype=torch.float16, device=device) k_cache = torch.randn((3, 8, 2, 60), dtype=torch.float16, device=device) @@ -191,13 +216,22 @@ def test_flash_attn_with_kvcache_scalar_seqlens_with_cache_batch_idx( cache_batch_idx = torch.tensor((2, 0), dtype=torch.int32, device=device) # Dao's scalar wrapper sizes by cache batch, so use the query-batch tensor # expected by its C++ kernel when `cache_batch_idx` remaps a smaller batch. - expected = flash_attn.flash_attn_with_kvcache( - q, - k_cache, - v_cache, - cache_seqlens=torch.full((q.size(0),), 5, dtype=torch.int32, device=device), - cache_batch_idx=cache_batch_idx, - ) + if device == "mlu": + expected, _ = _reference_flash_attn_with_kvcache( + q, + k_cache, + v_cache, + cache_seqlens=5, + cache_batch_idx=cache_batch_idx, + ) + else: + expected = flash_attn.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + cache_seqlens=torch.full((q.size(0),), 5, dtype=torch.int32, device=device), + cache_batch_idx=cache_batch_idx, + ) actual = torch.empty_like(q) infini.ops.flash_attn_with_kvcache( @@ -230,13 +264,16 @@ def test_flash_attn_with_kvcache_scalar_seqlens_with_cache_batch_idx( def test_flash_attn_with_kvcache_defaults(device, implementation_index): - if device != "cuda": - pytest.skip("FlashAttention FA2 requires the NVIDIA backend") + if device not in ("cuda", "mlu"): + pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") q = torch.randn((2, 1, 4, 64), dtype=torch.float16, device=device) k_cache = torch.randn((2, 8, 2, 64), dtype=torch.float16, device=device) v_cache = torch.randn_like(k_cache) - expected = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache) + if device == "mlu": + expected, _ = _reference_flash_attn_with_kvcache(q, k_cache, v_cache) + else: + expected = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache) actual = torch.empty_like(q) infini.ops.flash_attn_with_kvcache( @@ -252,25 +289,131 @@ def test_flash_attn_with_kvcache_defaults(device, implementation_index): def test_flash_attn_with_kvcache_non_default_stream(device, implementation_index): - if device != "cuda": - pytest.skip("non-default CUDA streams require the NVIDIA backend") + if device == "cuda": + accelerator = torch.cuda + stream_attribute = "cuda_stream" + elif device == "mlu": + accelerator = torch.mlu + stream_attribute = "mlu_stream" + else: + pytest.skip("stream coverage requires an accelerator backend") q = torch.randn((2, 1, 4, 64), dtype=torch.float16, device=device) k_cache = torch.randn((2, 8, 2, 64), dtype=torch.float16, device=device) v_cache = torch.randn_like(k_cache) - expected = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache) + if device == "mlu": + expected, _ = _reference_flash_attn_with_kvcache(q, k_cache, v_cache) + else: + expected = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache) actual = torch.empty_like(q) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) + stream = accelerator.Stream() + stream.wait_stream(accelerator.current_stream()) infini.ops.flash_attn_with_kvcache( q, k_cache, v_cache, actual, - stream=stream.cuda_stream, + stream=getattr(stream, stream_attribute), implementation_index=implementation_index, ) stream.synchronize() torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) + + +def _reference_flash_attn_with_kvcache( + q, + k_cache, + v_cache, + k=None, + v=None, + cache_seqlens=None, + cache_batch_idx=None, + block_table=None, + softmax_scale=None, + causal=False, + window_size=(-1, -1), +): + batch_size, query_length, num_heads, _ = q.shape + if cache_seqlens is None: + lengths = [k_cache.size(1)] * batch_size + elif isinstance(cache_seqlens, int): + lengths = [cache_seqlens] * batch_size + else: + lengths = cache_seqlens.cpu().tolist() + cache_rows = ( + list(range(batch_size)) + if cache_batch_idx is None + else cache_batch_idx.cpu().tolist() + ) + append_length = 0 if k is None else k.size(1) + + outputs = [] + softmax_lses = [] + for batch in range(batch_size): + length = lengths[batch] + if k is not None: + row = cache_rows[batch] + k_cache[row, length : length + append_length].copy_(k[batch]) + v_cache[row, length : length + append_length].copy_(v[batch]) + length += append_length + + if block_table is None: + row = cache_rows[batch] + k_seq = k_cache[row, :length] + v_seq = v_cache[row, :length] + else: + page_size = k_cache.size(1) + block_count = (length + page_size - 1) // page_size + blocks = block_table[batch, :block_count].cpu().tolist() + k_seq = torch.cat(tuple(k_cache[index] for index in blocks))[:length] + v_seq = torch.cat(tuple(v_cache[index] for index in blocks))[:length] + + q_seq = q[batch].transpose(0, 1) + k_seq = k_seq.transpose(0, 1) + v_seq = v_seq.transpose(0, 1) + groups = num_heads // k_seq.size(0) + k_seq = k_seq.repeat_interleave(groups, dim=0) + v_seq = v_seq.repeat_interleave(groups, dim=0) + scale = softmax_scale if softmax_scale is not None else q.size(-1) ** -0.5 + scores = torch.matmul(q_seq.float(), k_seq.float().transpose(-2, -1)) + scores *= scale + mask = _attention_mask( + query_length, + length, + causal, + window_size, + q.device, + ) + if mask is not None: + scores.masked_fill_(~mask.unsqueeze(0), -math.inf) + softmax_lse = torch.logsumexp(scores, dim=-1) + softmax_lses.append( + torch.where(torch.isneginf(softmax_lse), math.inf, softmax_lse) + ) + probabilities = torch.nan_to_num(torch.softmax(scores, dim=-1), nan=0.0) + outputs.append(torch.matmul(probabilities, v_seq.float()).to(q.dtype)) + + return ( + torch.stack(outputs).transpose(1, 2), + torch.stack(softmax_lses), + ) + + +def _attention_mask(q_len, k_len, causal, window_size, device): + left, right = window_size + if not causal and left < 0 and right < 0: + return None + + query_positions = torch.arange(q_len, device=device).unsqueeze(1) + key_positions = torch.arange(k_len, device=device).unsqueeze(0) + aligned_query_positions = query_positions + k_len - q_len + mask = torch.ones((q_len, k_len), dtype=torch.bool, device=device) + if left >= 0: + mask &= key_positions >= aligned_query_positions - left + if causal: + right = 0 + if right >= 0: + mask &= key_positions <= aligned_query_positions + right + return mask