StaticAttetnion runtime support for generate_full_logits=False

Summary: Prefill will produce the logits at position 0 when not generating full logits. Lookahead decoding requires full logits.

Differential Revision: D88790445

Author: sxu

examples/models/llama/runner/static_attention_io_manager.h

@@ -438,6 +438,7 @@ class StaticAttentionIOManager {
     RopeT* rope_freqs_cos;
     RopeT* rope_freqs_sin;
     StaticAttentionUpdateStyle style = StaticAttentionUpdateStyle::SMART_MASK;
+    bool generate_full_logits = true;
   };
 
   StaticAttentionIOManager(StaticAttentionIOConfig config)
@@ -607,7 +608,7 @@ class StaticAttentionIOManager {
       batch_len = std::min(input_len, tokens.size() - i);
       if (input_pos_ + batch_len > config_.max_context_len) {
         ET_LOG(Error, "Maximum context size reached, stopping prefill.");
-        return input_len - 1;
+        return config_.generate_full_logits ? input_len - 1 : 0;
       }
       std::copy(&tokens[i], &tokens[i + batch_len], input_buffer.begin());
       prepare(method);
@@ -622,10 +623,12 @@ class StaticAttentionIOManager {
         auto* logits = logits_tensor.const_data_ptr<LogitT>();
         logits_callback(executorch::runtime::Span(
             logits,
-            logits + batch_len * logits_tensor.size(logits_tensor.dim() - 1)));
+            logits +
+                (config_.generate_full_logits ? batch_len : 1) *
+                    logits_tensor.size(logits_tensor.dim() - 1)));
       }
     }
-    return batch_len - 1;
+    return config_.generate_full_logits ? batch_len - 1 : 0;
   }
 
   /**
@@ -685,6 +688,7 @@ class StaticAttentionIOManager {
       size_t window_size,
       size_t n_verifications,
       std::unordered_map<TokenT, SuffixCache<TokenT>> suffix_caches) {
+    ET_CHECK(config_.generate_full_logits);
     ET_LOG(
         Info,
         "Decoding with lookahead and verification at position %zu",