Fix native static attention sliding-window masks

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(

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 <algorithm>
+#include <cstddef>
+#include <tuple>
+#include <vector>
+
+#include <executorch/examples/models/llama/runner/static_attention_io_manager.h>
+
+#include <gtest/gtest.h>
+
+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<float> local_mask(
+      window,
+      input_len,
+      1,
+      kZeroVal,
+      kMaskVal,
+      StaticAttentionUpdateStyle::SMART_MASK,
+      true);
+  StaticAttentionMask<float> 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<std::tuple<size_t, size_t, size_t>> 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<float> 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<std::vector<float>> 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

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"],
+    )

examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp

@@ -62,6 +62,58 @@ void fill_mask(
     }
   }
 }
+
+template <typename T>
+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<uint16_t*>(buf),
+          size,
+          count,
+          static_cast<uint16_t>(65535),
+          static_cast<uint16_t>(0));
+      break;
+    case executorch::aten::ScalarType::Byte:
+      mask_oldest_visible_impl(
+          reinterpret_cast<uint8_t*>(buf),
+          size,
+          count,
+          static_cast<uint8_t>(255),
+          static_cast<uint8_t>(0));
+      break;
+    case executorch::aten::ScalarType::Float:
+      mask_oldest_visible_impl(
+          reinterpret_cast<float*>(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<MethodMeta> 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<int32_t>& position_offset) {
+    const std::vector<int32_t>& /*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<size_t>(sliding_window);
+  std::vector<size_t> 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,27 +289,24 @@ 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(
         attention_mask_dtype_,
         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<int32_t>& position_offset) {
   std::byte* cur_ptr =
       attention_mask + n_past * getDtypeSize(attention_mask_dtype_);
+  const size_t window_size = static_cast<size_t>(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<size_t>(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_);
   }