From 0181d8a38d2f652263456859794f0432e6fcb02b Mon Sep 17 00:00:00 2001 From: Yiwen Xie Date: Fri, 28 Aug 2026 15:52:56 -0700 Subject: [PATCH] Fix native static attention sliding-window masks Summary: The native static-attention mask kept every valid cache slot visible for every input row. For multi-token autoregressive chunks, row k therefore attended to too many cached tokens, and chunks longer than the local window also retained too many in-chunk keys. The QNN runner made the same assumption for long and sparse chunks. Derive SMART_MASK cache visibility from ring age and each row's remaining window budget, bound the causal in-chunk region, and apply equivalent oldest-visible trimming in the QNN runner. Add regression coverage for cache wrapping, mixed local/global masks, and chunk lengths both above and below the sliding window. Update the SNR evaluator to use the public decode path so it exercises the same per-row masking behavior. Differential Revision: D117411412 --- .../runner/static_attention_io_manager.h | 53 ++++++++- .../static_attention_io_manager_test.cpp | 107 ++++++++++++++++++ examples/models/llama/runner/targets.bzl | 6 + .../oss_scripts/llama/runner/kv_manager.cpp | 93 +++++++++++---- 4 files changed, 236 insertions(+), 23 deletions(-) create mode 100644 examples/models/llama/runner/static_attention_io_manager_test.cpp diff --git a/examples/models/llama/runner/static_attention_io_manager.h b/examples/models/llama/runner/static_attention_io_manager.h index 543b38277df..ed9c7a61e19 100644 --- a/examples/models/llama/runner/static_attention_io_manager.h +++ b/examples/models/llama/runner/static_attention_io_manager.h @@ -297,14 +297,16 @@ class StaticAttentionMask { size_t head_dim, T zero_val, T mask_val, - StaticAttentionUpdateStyle style = StaticAttentionUpdateStyle::SMART_MASK) + StaticAttentionUpdateStyle style = StaticAttentionUpdateStyle::SMART_MASK, + bool is_sliding_window = false) : cache_len_(cache_len), input_len_(input_len), head_dim_(head_dim), cache_valid_len_(0), zero_val_(zero_val), mask_val_(mask_val), - style_(style) { + style_(style), + is_sliding_window_(is_sliding_window) { data_size_ = input_len_ * (cache_len_ + input_len_); data_ = allocator_.allocate(data_size_); ET_CHECK(data_ != nullptr); @@ -351,11 +353,42 @@ class StaticAttentionMask { void set_causal_mask() { for (size_t i = 0; i < input_len_; i++) { auto* p = data_ + (cache_len_ + input_len_) * i; - std::fill(p + cache_len_, p + cache_len_ + 1 + i, zero_val_); + size_t first_visible = 0; + if (is_sliding_window_ && i + 1 > cache_len_) { + first_visible = i + 1 - cache_len_; + } + std::fill(p + cache_len_, p + cache_len_ + first_visible, mask_val_); + std::fill( + p + cache_len_ + first_visible, p + cache_len_ + 1 + i, zero_val_); std::fill(p + cache_len_ + 1 + i, p + cache_len_ + input_len_, mask_val_); } } + void set_sliding_window_mask(size_t input_pos) { + if (!is_sliding_window_ || cache_len_ == 0) { + return; + } + + const size_t valid_cache_len = std::min(input_pos, cache_len_); + const size_t cache_pos = input_pos % cache_len_; + for (size_t row = 0; row < input_len_; row++) { + auto* p = data_ + (cache_len_ + input_len_) * row; + std::fill(p, p + cache_len_, mask_val_); + // Row k already sees k + 1 in-chunk keys, leaving W - k - 1 + // cache keys. cache_pos is the next ring slot, so the previous slot has + // age zero. + const size_t max_age = row + 1 < cache_len_ + ? std::min(valid_cache_len, cache_len_ - row - 1) + : 0; + for (size_t col = 0; col < cache_len_; col++) { + const size_t age = (cache_pos + cache_len_ - 1 - col) % cache_len_; + if (age < max_age) { + p[col] = zero_val_; + } + } + } + } + T* get() { return data_; } @@ -376,6 +409,7 @@ class StaticAttentionMask { T zero_val_; T mask_val_; StaticAttentionUpdateStyle style_; + bool is_sliding_window_; AllocatorT allocator_; size_t data_size_ = 0; T* data_; @@ -496,6 +530,10 @@ class StaticAttentionIOManager { */ PerCacheLenMasks& add_mask(size_t input_len, MaskT zero_val, MaskT mask_val) { PerCacheLenMasks masks; + size_t global_cache_len = 0; + for (const auto& pair : config_.cache_len_to_mask_idx) { + global_cache_len = std::max(global_cache_len, pair.first); + } for (auto& pair : config_.cache_len_to_mask_idx) { masks.emplace_back( pair.first, @@ -505,7 +543,8 @@ class StaticAttentionIOManager { config_.head_dim, zero_val, mask_val, - config_.style)); + config_.style, + pair.first > 0 && pair.first < global_cache_len)); } auto it = attentionMasks_.emplace(input_len, std::move(masks)); return it.first->second; @@ -688,6 +727,9 @@ class StaticAttentionIOManager { *config_.last_valid_token_pos_index, &last_valid_token_pos_); } + for (auto& pair : masks) { + pair.second->set_sliding_window_mask(input_pos_); + } prepare(method); ET_CHECK(method.execute() == executorch::runtime::Error::Ok); update( @@ -740,6 +782,9 @@ class StaticAttentionIOManager { ET_LOG(Error, "Maximum context size reached, stopping decode."); break; } + for (auto& pair : masks) { + pair.second->set_sliding_window_mask(input_pos_); + } prepare(method); ET_CHECK(method.execute() == executorch::runtime::Error::Ok); update( diff --git a/examples/models/llama/runner/static_attention_io_manager_test.cpp b/examples/models/llama/runner/static_attention_io_manager_test.cpp new file mode 100644 index 00000000000..595cf1612ff --- /dev/null +++ b/examples/models/llama/runner/static_attention_io_manager_test.cpp @@ -0,0 +1,107 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include + +namespace example { +namespace { + +size_t count_visible(const float* mask, size_t row, size_t row_size) { + return std::count(mask + row * row_size, mask + (row + 1) * row_size, 0.0f); +} + +void check_mask_invariant(size_t input_len, size_t window, size_t prompt_len) { + constexpr float kZeroVal = 0.0f; + constexpr float kMaskVal = -1.0f; + const size_t global_cache_len = prompt_len + input_len; + StaticAttentionMask local_mask( + window, + input_len, + 1, + kZeroVal, + kMaskVal, + StaticAttentionUpdateStyle::SMART_MASK, + true); + StaticAttentionMask global_mask( + global_cache_len, input_len, 1, kZeroVal, kMaskVal); + local_mask.set_causal_mask(); + global_mask.set_causal_mask(); + + for (size_t pos = 0; pos < prompt_len; pos += input_len) { + const size_t update_len = std::min(input_len, prompt_len - pos); + local_mask.set_sliding_window_mask(pos); + global_mask.set_sliding_window_mask(pos); + for (size_t row = 0; row < update_len; row++) { + EXPECT_EQ( + count_visible(local_mask.get(), row, window + input_len), + std::min(pos + row + 1, window)) + << "input_len=" << input_len << ", window=" << window + << ", prompt_len=" << prompt_len << ", pos=" << pos + << ", row=" << row; + EXPECT_EQ( + count_visible(global_mask.get(), row, global_cache_len + input_len), + pos + row + 1) + << "input_len=" << input_len << ", prompt_len=" << prompt_len + << ", pos=" << pos << ", row=" << row; + } + local_mask.unmask(update_len); + global_mask.unmask(update_len); + } +} + +TEST(StaticAttentionMaskTest, SlidingWindowInvariant) { + const std::vector> cases = { + {8, 4, 24}, + {6, 12, 24}, + {64, 256, 1242}, + {1024, 256, 1242}, + {1, 256, 600}, + }; + for (const auto& [input_len, window, prompt_len] : cases) { + check_mask_invariant(input_len, window, prompt_len); + } +} + +TEST(StaticAttentionMaskTest, SlidingWindowMaskRotatesWithCacheRing) { + constexpr size_t kWindow = 4; + constexpr size_t kInputLen = 4; + StaticAttentionMask mask( + kWindow, + kInputLen, + 1, + 0.0f, + -1.0f, + StaticAttentionUpdateStyle::SMART_MASK, + true); + mask.set_causal_mask(); + mask.set_sliding_window_mask(9); + + const std::vector> expected_cache = { + {0.0f, -1.0f, 0.0f, 0.0f}, + {0.0f, -1.0f, -1.0f, 0.0f}, + {0.0f, -1.0f, -1.0f, -1.0f}, + {-1.0f, -1.0f, -1.0f, -1.0f}, + }; + for (size_t row = 0; row < kInputLen; row++) { + for (size_t col = 0; col < kWindow; col++) { + EXPECT_EQ( + mask.get()[row * (kWindow + kInputLen) + col], + expected_cache[row][col]); + } + } +} + +} // namespace +} // namespace example diff --git a/examples/models/llama/runner/targets.bzl b/examples/models/llama/runner/targets.bzl index 81a2df117f4..ece081fc749 100644 --- a/examples/models/llama/runner/targets.bzl +++ b/examples/models/llama/runner/targets.bzl @@ -83,3 +83,9 @@ def define_common_targets(): "//executorch/runtime/executor:program", ] ) + + runtime.cxx_test( + name = "static_attention_io_manager_test", + srcs = ["static_attention_io_manager_test.cpp"], + deps = [":static_attention_io_manager"], + ) diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp index 7406ca8887b..03daac6e8cd 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp @@ -62,6 +62,58 @@ void fill_mask( } } } + +template +void mask_oldest_visible_impl( + T* buf, + size_t size, + size_t count, + T visible_value, + T mask_value) { + // Lookahead rows can be sparse, and an AR chunk can be longer than the + // window, so the entries to drop are not necessarily a contiguous prefix. + for (size_t i = 0; i < size && count > 0; i++) { + if (buf[i] == visible_value) { + buf[i] = mask_value; + --count; + } + } +} + +void mask_oldest_visible( + executorch::aten::ScalarType scalar_type, + std::byte* buf, + size_t size, + size_t count) { + switch (scalar_type) { + case executorch::aten::ScalarType::UInt16: + mask_oldest_visible_impl( + reinterpret_cast(buf), + size, + count, + static_cast(65535), + static_cast(0)); + break; + case executorch::aten::ScalarType::Byte: + mask_oldest_visible_impl( + reinterpret_cast(buf), + size, + count, + static_cast(255), + static_cast(0)); + break; + case executorch::aten::ScalarType::Float: + mask_oldest_visible_impl( + reinterpret_cast(buf), size, count, 0.0f, -65535.0f); + break; + default: + ET_CHECK_MSG( + false, + "Unsupported scalar type %s", + executorch::runtime::toString(scalar_type)); + break; + } +} } // namespace KVManager::KVManager(Metadata metadata, std::unique_ptr method_meta) @@ -196,7 +248,7 @@ void KVManager::init_attention_mask( int32_t ar_len, int32_t n_past, int32_t sliding_window, - const std::vector& position_offset) { + const std::vector& /*position_offset*/) { ET_CHECK_MSG( attention_map.size() <= ar_len, "The size of attention_map (%zu) doesn't match with ar_len (%d)", @@ -213,8 +265,11 @@ void KVManager::init_attention_mask( std::byte* past_ptr = attention_mask; std::byte* new_ptr = attention_mask + (metadata_.context_len - ar_len) * getDtypeSize(attention_mask_dtype_); + const size_t window_size = static_cast(sliding_window); + std::vector visible_counts(ar_len); // All inputs will necessarily attend to n_past and itself for (int i = 0; i < ar_len; i++) { + size_t visible_count; // Iterate across ar_len if (attention_map[i] < 0) { // If negative, attend to only past tokens @@ -223,6 +278,7 @@ void KVManager::init_attention_mask( past_ptr, n_past, /*use_pos_value=*/true); + visible_count = n_past; } else { // If positive, copy attention map from (relative to 0th input) parent // Parent token index @@ -233,6 +289,7 @@ void KVManager::init_attention_mask( past_ptr, parent_ptr, metadata_.context_len * getDtypeSize(attention_mask_dtype_)); + visible_count = visible_counts[pidx]; } // Attend to itself fill_mask( @@ -240,20 +297,16 @@ void KVManager::init_attention_mask( new_ptr + i * getDtypeSize(attention_mask_dtype_), 1, /*use_pos_value=*/true); - - // mask by limitation of sliding_window - int32_t available_context_len = position_offset.empty() - ? sliding_window - (i + 1) - n_past - : sliding_window - (position_offset[i] + 1) - n_past; - // if available_context_len is less than 0, it means we need to mask some - // tokens in the past to avoid exceeding the sliding window - if (available_context_len < 0) { - fill_mask( + ++visible_count; + if (visible_count > window_size) { + mask_oldest_visible( attention_mask_dtype_, past_ptr, - -available_context_len, - /*use_pos_value=*/false); + metadata_.context_len, + visible_count - window_size); + visible_count = window_size; } + visible_counts[i] = visible_count; past_ptr += metadata_.context_len * getDtypeSize(attention_mask_dtype_); new_ptr += metadata_.context_len * getDtypeSize(attention_mask_dtype_); @@ -283,18 +336,20 @@ void KVManager::update_attention_mask( const std::vector& position_offset) { std::byte* cur_ptr = attention_mask + n_past * getDtypeSize(attention_mask_dtype_); + const size_t window_size = static_cast(sliding_window); for (int i = 0; i < ar_len; i++) { fill_mask(attention_mask_dtype_, cur_ptr, n_update, /*use_pos_value=*/true); - int32_t available_cache_len = position_offset.empty() - ? sliding_window - (i + 1) - : sliding_window - (position_offset[i] + 1); - if (n_past + n_update > available_cache_len) { - fill_mask( + const int32_t position = position_offset.empty() ? i : position_offset[i]; + const size_t visible_count = + std::min(static_cast(n_past + position + 1), window_size) + + n_update; + if (visible_count > window_size) { + mask_oldest_visible( attention_mask_dtype_, cur_ptr - n_past * getDtypeSize(attention_mask_dtype_), - n_past + n_update - available_cache_len, - /*use_pos_value=*/false); + metadata_.context_len, + visible_count - window_size); } cur_ptr += metadata_.context_len * getDtypeSize(attention_mask_dtype_); }