Pass ForwardOptions from top level module and also return any relevant state as output (#8186)

Summary:

Pass a `ForwardOptions` argument (introduced by #8128) from the top level transformer, consolidate some existing inputs into it, and return any optional updates from the attention implementation.

Reviewed By: iseeyuan

Differential Revision: D69080123

Author: sxu

examples/models/llama/attention.py

@@ -236,7 +236,7 @@ def forward(
             assert input_pos is not None
             k, v = self.kv_cache.update(input_pos, k, v)
             output = self.SDPA(input_pos, q, k, v, bsz, seqlen, self.mask)
-            return self.wo(output)
+            return self.wo(output), None
 
         # grouped multiquery attention: expand out keys and values
         k = k.repeat_interleave(self.n_rep, dim=1)
@@ -252,4 +252,4 @@ def forward(
 
         output = self.wo(output)
 
-        return output
+        return output, None

examples/models/llama/eval_llama_lib.py

@@ -65,7 +65,9 @@ def _model_call(self, inps):
                 result_logits = []
                 for pos in range(inps.shape[-1]):
                     pos_tensor = torch.tensor([pos], dtype=torch.int64)
-                    logits = self._model(inps[:, pos : pos + 1], pos_tensor)
+                    logits = self._model(
+                        inps[:, pos : pos + 1], None, {"input_pos": pos_tensor}
+                    )
                     result_logits.append(logits)
                 if self._generate_full_logits:
                     return torch.cat(result_logits, dim=1)
@@ -74,7 +76,9 @@ def _model_call(self, inps):
             else:
                 pos_tensor = torch.tensor([0], dtype=torch.int64, device=self.device)
                 # Batch process the whole sequence.
-                logits = self._model(inps[:, : self._max_seq_length], pos_tensor)
+                logits = self._model(
+                    inps[:, : self._max_seq_length], None, {"input_pos": pos_tensor}
+                )
                 return logits
 
         else:

examples/models/llama/evaluate/eager_eval.py

@@ -77,7 +77,9 @@ def _model_call(self, inps):
         if self._use_kv_cache:
             pos_tensor = torch.tensor([0], dtype=torch.int64, device=self.device)
             # Batch process the whole sequence.
-            logits = self._model(inps[:, : self._max_seq_length], pos_tensor)
+            logits = self._model(
+                inps[:, : self._max_seq_length], None, {"input_pos": pos_tensor}
+            )
             return logits
         else:
             return self._model(inps)

examples/models/llama/llama_transformer.py

@@ -7,12 +7,15 @@
 
 # Please refer to README.md in the same folder for more information.
 
-from typing import Optional
+from typing import Any, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
 
-from executorch.examples.models.llama.attention import ATTENTION_REGISTRY
+from executorch.examples.models.llama.attention import (
+    ATTENTION_REGISTRY,
+    ForwardOptions,
+)
 
 from executorch.examples.models.llama.model_args import ModelArgs
 
@@ -148,17 +151,17 @@ def __init__(self, layer_id: int, args: ModelArgs, rope: Rope):
         self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
         self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
 
-    def forward(self, x, freqs_cos, freqs_sin, input_pos=None):  # x: 1xN
-        h = self.attention.forward(
-            self.attention_norm(x), freqs_cos, freqs_sin, input_pos=input_pos
+    def forward(self, x, freqs_cos, freqs_sin, attn_options: ForwardOptions):  # x: 1xN
+        h, attn_options_update = self.attention.forward(
+            self.attention_norm(x), freqs_cos, freqs_sin, **attn_options
         )
 
         h = x + h
         if hasattr(self, "block_sparse_moe"):
             out = h + self.block_sparse_moe(self.ffn_norm(h))
         else:
             out = h + self.feed_forward(self.ffn_norm(h))
-        return out
+        return out, attn_options_update
 
 
 class Transformer(nn.Module):
@@ -185,27 +188,28 @@ def __init__(self, params: ModelArgs):
     def forward(
         self,
         tokens: Optional[torch.LongTensor] = None,  # tokens
-        input_pos: Optional[
-            torch.LongTensor
-        ] = None,  # Scalar tensor indicating size of window of the caches
         h: Optional[torch.FloatTensor] = None,  # embeddings
-    ) -> torch.Tensor:
+        attn_options: Optional[ForwardOptions] = None,
+    ) -> Tuple[torch.Tensor, Optional[Any]]:
         if (tokens is None) ^ (h is not None):
             raise ValueError(
                 "You cannot specify both tokens and h at the same time, and must specify either one"
             )
         if tokens is not None and h is None:
             h = self.tok_embeddings(tokens)
+
+        if attn_options is None:
+            attn_options = {}
         seqlen = h.shape[1]
-        freqs_cos, freqs_sin = self.rope.get_freqs(input_pos, seqlen)
+        freqs_cos, freqs_sin = self.rope.get_freqs(
+            attn_options.get("input_pos"), seqlen
+        )
 
+        attn_options_update = None
         for layer in self.layers:
-            h = layer(
-                h,
-                freqs_cos,
-                freqs_sin,
-                input_pos,
-            )
+            h, attn_options_update = layer(h, freqs_cos, freqs_sin, attn_options)
+            if attn_options_update is not None:
+                attn_options.update(**attn_options_update)
 
         if not self.generate_full_logits:
             # Only the last logit is used for the new generated token
@@ -237,4 +241,4 @@ def forward(
                 expanded_logits[:, list(self.output_prune_map.values())] = logits
             logits = expanded_logits
 
-        return logits
+        return logits, attn_options_update

examples/models/llama/model.py

@@ -289,16 +289,20 @@ def get_example_inputs_kvcache_sdpa(self):
         if self.enable_dynamic_shape:
             return (
                 torch.tensor([[2, 3, 4]], dtype=torch.long),
-                torch.tensor([0], dtype=torch.long),
+                None,
+                {"input_pos": torch.tensor([0], dtype=torch.long)},
             )
         else:
             return (
                 torch.tensor(
                     [[1]], dtype=torch.long
                 ),  # tokens, with kv cache our input token length is always just 1 token.
-                torch.tensor(
-                    [0], dtype=torch.long
-                ),  # start_pos, what token of output are we on.
+                None,  # hidden state
+                {
+                    "input_pos": torch.tensor(
+                        [0], dtype=torch.long
+                    )  # start_pos, what token of output are we on.
+                },
             )
 
     def _transform_for_pre_quantization(self, checkpoint, model_args):

extension/llm/export/builder.py

@@ -170,7 +170,7 @@ def _get_dynamic_shape(self) -> Any:
             self.dynamic_shapes = ({1: dim},)
         elif self.enable_dynamic_shape:
             # Two input arguments: tokens and input_pos but input_pos is static shape
-            self.dynamic_shapes = ({1: dim}, {0: 1})
+            self.dynamic_shapes = ({1: dim}, None, {"input_pos": {0: 1}})
         else:
             # Two input arguments: tokens and input_pos but both are of static shape
             self.dynamic_shapes = None