Static attention IO manager: fix causal mask bug for last input position

sxu · web-flow · commit ee13aa053bf3 · 2025-08-12T16:20:46.000-07:00
Differential Revision: D80098286 Pull Request resolved: #13332
diff --git a/examples/models/llama/runner/static_attention_io_manager.h b/examples/models/llama/runner/static_attention_io_manager.h
@@ -328,7 +328,7 @@ class StaticAttentionMask {
   }
 
   void set_causal_mask() {
-    for (size_t i = 0; i < input_len_ - 1; i++) {
+    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_);
       std::fill(p + cache_len_ + 1 + i, p + cache_len_ + input_len_, mask_val_);

Original file line number	Diff line number	Diff line change
`@@ -328,7 +328,7 @@ class StaticAttentionMask {`
`328`	`328`	`}`
`329`	`329`
`330`	`330`	`void set_causal_mask() {`
`331`		`- for (size_t i = 0; i < input_len_ - 1; i++) {`
	`331`	`+ for (size_t i = 0; i < input_len_; i++) {`
`332`	`332`	`auto* p = data_ + (cache_len_ + input_len_) * i;`
`333`	`333`	`std::fill(p + cache_len_, p + cache_len_ + 1 + i, zero_val_);`
`334`	`334`	`std::fill(p + cache_len_ + 1 + i, p + cache_len_ + input_len_, mask_val_);`