[Native] Select the logits rows an off-graph export returns

backends/native/examples/hf_llm/export_hf_llm.py

@@ -44,6 +44,7 @@
 from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter
 from executorch.extension.llm.export.model_metadata import (
     write_cache_geometry,
+    write_logits_to_keep_mode,
     write_max_context_len,
 )
 from torch.export import Dim, export
@@ -113,14 +114,16 @@ def forward(
         self,
         input_ids: torch.Tensor,
         position_ids: torch.Tensor | None = None,
-        logits_to_keep: int = 0,
+        logits_to_keep: int | torch.Tensor = 0,
         **_,
     ) -> CausalLMOutput:
         text_config = self.config.text_config
         hidden = self.model["language_model"](
             input_ids=input_ids, position_ids=position_ids, use_cache=False
         ).last_hidden_state
-        if logits_to_keep:
+        if isinstance(logits_to_keep, torch.Tensor):
+            hidden = hidden[:, logits_to_keep]
+        elif logits_to_keep:
             hidden = hidden[:, -logits_to_keep:]
         logits = self.lm_head(hidden) * text_config.output_multiplier
         cap = text_config.final_logit_softcapping
@@ -289,6 +292,11 @@ def forward(
             logits_to_keep=1,
         ).logits
 
+    def call_inputs(
+        self, input_ids: torch.Tensor, positions: torch.Tensor
+    ) -> list[torch.Tensor]:
+        return [input_ids, positions]
+
     def reset_cache(self) -> None:
         self.cache.reset()
 
@@ -318,7 +326,9 @@ def _cache_windows(config) -> list[int]:
 
 
 class OffGraphLastTokenLM(nn.Module):
-    """Last-position logits with the KV cache outside the graph.
+    """Last-position logits with the KV cache outside the graph, or with
+    ``selected`` the logits of the token indices a third input names, which a
+    batched forward over several sequences needs.
 
     Each attention layer emits ``kvcache::update_and_attend`` over this step's
     K/V; the cache, not the graph, applies a sliding layer's window, so the
@@ -327,9 +337,10 @@ class OffGraphLastTokenLM(nn.Module):
     ``capacity`` positions the active one.
     """
 
-    def __init__(self, model: nn.Module, capacity: int) -> None:
+    def __init__(self, model: nn.Module, capacity: int, selected: bool = False) -> None:
         super().__init__()
         self.model = model
+        self.selected = selected
         config = model.config.get_text_config()
         self._windows = _cache_windows(config)
         self._head_dim = (
@@ -362,13 +373,27 @@ def metadata(self) -> dict[str, int | list[int]]:
         geometry = write_cache_geometry(
             [self._kv_heads] * n, [self._head_dim] * n, self._windows
         )
-        return {
-            name: value.tolist() if isinstance(value, torch.Tensor) else value
-            for name, value in geometry.items()
-        } | write_max_context_len(self._max_context_len)
+        return (
+            {
+                name: value.tolist() if isinstance(value, torch.Tensor) else value
+                for name, value in geometry.items()
+            }
+            | write_max_context_len(self._max_context_len)
+            | write_logits_to_keep_mode("selected" if self.selected else "last")
+        )
+
+    def call_inputs(
+        self, input_ids: torch.Tensor, positions: torch.Tensor
+    ) -> list[torch.Tensor]:
+        """A call's inputs, selecting the last token's logits."""
+        last = torch.tensor([input_ids.shape[1] - 1], dtype=torch.long)
+        return [input_ids, positions, last] if self.selected else [input_ids, positions]
 
     def forward(
-        self, input_ids: torch.Tensor, cache_position: torch.Tensor
+        self,
+        input_ids: torch.Tensor,
+        cache_position: torch.Tensor,
+        logits_to_keep: torch.Tensor | None = None,
     ) -> torch.Tensor:
         return self.model(
             input_ids=input_ids,
@@ -377,7 +402,7 @@ def forward(
             position_ids=cache_position.unsqueeze(0),
             use_cache=False,
             past_key_values=None,
-            logits_to_keep=1,
+            logits_to_keep=logits_to_keep if self.selected else 1,
         ).logits
 
     def reset_cache(self) -> None:
@@ -460,8 +485,9 @@ def golden_samples(
             positions = torch.arange(len(ids), dtype=torch.long)
             calls = []
             for _ in range(decode_steps + 1):
-                logits = wrapper(input_ids, positions)
-                calls.append(Call("forward", [input_ids, positions], [logits]))
+                inputs = wrapper.call_inputs(input_ids, positions)
+                logits = wrapper(*inputs)
+                calls.append(Call("forward", inputs, [logits]))
                 input_ids = logits[:, -1].argmax(dim=-1, keepdim=True)
                 positions = positions[-1:] + 1
             samples.append(Sample(calls))
@@ -505,6 +531,14 @@ def _parse_args() -> argparse.Namespace:
         help="Export attention as kvcache::update_and_attend with no KV cache in "
         "the graph; the runner owns the cache and sets its capacity and dtype.",
     )
+    parser.add_argument(
+        "--logits_to_keep",
+        choices=["last", "selected"],
+        default="last",
+        help="With --offgraph_cache: 'selected' adds a third int64[K] input of "
+        "token indices whose logits to return, for a batching runtime that "
+        "puts several sequences in one forward.",
+    )
     parser.add_argument(
         "--samples",
         type=int,
@@ -520,6 +554,8 @@ def _parse_args() -> argparse.Namespace:
         parser.error(
             "--gguf loads every weight; it excludes --random_weights and --quantize"
         )
+    if args.logits_to_keep == "selected" and not args.offgraph_cache:
+        parser.error("--logits_to_keep selected needs --offgraph_cache")
     if args.offgraph_cache and (
         args.max_context_length or args.kv_cache_dtype != "fp32"
     ):
@@ -560,7 +596,9 @@ def main() -> None:
     if args.quantize == "8da4w":
         _quantize_8da4w(model, args.group_size, args.embedding_group_size)
     if args.offgraph_cache:
-        wrapper = OffGraphLastTokenLM(model, cache_len).eval()
+        wrapper = OffGraphLastTokenLM(
+            model, cache_len, selected=args.logits_to_keep == "selected"
+        ).eval()
     else:
         kv_cache_dtype = (
             torch.float16 if args.kv_cache_dtype == "fp16" else torch.float32
@@ -570,11 +608,19 @@ def main() -> None:
     input_ids = torch.zeros((1, 6), dtype=torch.long)
     cache_position = torch.arange(input_ids.shape[1], dtype=torch.long)
     seq = Dim("seq", min=1, max=args.max_seq_length - 1)
+    inputs = wrapper.call_inputs(input_ids, cache_position)
+    dynamic_shapes = {"input_ids": {1: seq}, "cache_position": {0: seq}}
+    if len(inputs) == 3:
+        # Trace two rows: export specializes a size of 1.
+        inputs[2] = torch.tensor([0, input_ids.shape[1] - 1])
+        dynamic_shapes["logits_to_keep"] = {
+            0: Dim("rows", min=1, max=args.max_seq_length - 1)
+        }
     with torch.no_grad():
         ep = export(
             wrapper,
-            (input_ids, cache_position),
-            dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}},
+            tuple(inputs),
+            dynamic_shapes=dynamic_shapes,
             strict=not args.non_strict,
         )
     report_lowering(ep, "forward")

backends/native/examples/hf_llm/runtime/BUCK

@@ -38,6 +38,7 @@ runtime.cxx_library(
         "//executorch/backends/native/runtime/graph:scalar_type",
         "//executorch/examples/models/llama/tokenizer:tiktoken",
         "//executorch/extension/llm/cache:kv_cache",
+        "//executorch/extension/llm/runner:constants",
     ],
     exported_deps = ["//executorch/backends/native/runtime/engine:engine"],
     external_deps = ["gflags"],

backends/native/examples/hf_llm/runtime/LlmRunner.cpp

@@ -1,3 +1,3 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 // All rights reserved.
 //
@@ -35,6 +35,7 @@
 #include <executorch/backends/native/runtime/graph/ScalarType.h>
 #include <executorch/examples/models/llama/tokenizer/llama_tiktoken.h>
 #include <executorch/extension/llm/cache/cache_registry.h>
+#include <executorch/extension/llm/runner/constants.h>
 #include <executorch/runtime/backend/options.h>
 #include <gflags/gflags.h>
 #include <pytorch/tokenizers/error.h>
@@ -93,6 +94,9 @@
 
 namespace cache = executorch::extension::llm::cache;
 
+// get_logits_to_keep_mode's value for selected logits.
+constexpr int64_t kSelectedLogits = 2;
+
 using Clock = std::chrono::steady_clock;
 
 double elapsed_ms(Clock::time_point start, Clock::time_point end) {
@@ -304,6 +308,13 @@
   auto run_tokens = [&](const std::vector<uint64_t>& tokens, int64_t position) {
     set_tokens(tokens);
     set_positions(position, tokens.size());
+    // A selected-logits method takes the token indices to return; the runner
+    // wants only the last.
+    if (executable.num_inputs() == 3) {
+      const int64_t last = static_cast<int64_t>(tokens.size()) - 1;
+      executable.resize_input(2, {1});
+      executable.set_input(2, &last, 1, ptn::kLong);
+    }
     executable.execute();
     const std::vector<int64_t> output_shape =
         executable.output_sizes(logits_output);
@@ -536,8 +547,18 @@
         elapsed_ms(load_start, load_end));
     std::printf("constants: %zu package bytes\n", package->constant_bytes());
 
-    if (executable->num_inputs() != 2) {
-      throw std::runtime_error("runner requires token and position inputs");
+    // A selected-logits method's third input names the token indices whose
+    // logits to return.
+    const bool selected_logits =
+        package->metadata_int(executorch::extension::llm::kLogitsToKeepMode) ==
+        kSelectedLogits;
+    const size_t num_inputs = executable->num_inputs();
+    if (num_inputs != (selected_logits ? 3 : 2)) {
+      throw std::runtime_error(
+          selected_logits
+              ? "a selected-logits method must take tokens, positions and the "
+                "token indices whose logits to return"
+              : "runner requires token and position inputs");
     }
     if (logits_output >= executable->num_outputs()) {
       throw std::runtime_error("logits output is outside engine outputs");
@@ -556,11 +577,11 @@
       throw std::runtime_error(
           "position input must have one element or one per token");
     }
-    for (size_t i = 0; i < 2; ++i) {
+    for (size_t i = 0; i < num_inputs; ++i) {
       const ptn::ScalarType dtype = executable->input_dtype(i);
       if (dtype != ptn::kInt && dtype != ptn::kLong) {
         throw std::runtime_error(
-            "token and position inputs must be Int or Long, got " +
+            "token, position and index inputs must be Int or Long, got " +
             std::string(ptn::scalar_type_name(dtype)));
       }
     }
@@ -601,7 +622,7 @@
                 " runs method " + calls[j].method +
                 ", but the runner replays only " + method_name);
           }
-          if (calls[j].inputs.size() != 2 || !calls[j].outputs ||
+          if (calls[j].inputs.size() != num_inputs || !calls[j].outputs ||
               calls[j].outputs->size() <= logits_output ||
               (*calls[j].outputs)[logits_output].dtype != ptn::kFloat ||
               (*calls[j].outputs)[logits_output].sizes.empty() ||
@@ -627,8 +648,12 @@
               static_cast<int64_t>(j == 0 ? 0 : prompt.size() + j - 1);
           std::vector<int64_t> positions(per_token_positions ? count : 1);
           std::iota(positions.begin(), positions.end(), start);
+          const std::vector<int64_t> last_index = {
+              static_cast<int64_t>(ids.size()) - 1};
           if (ids.empty() || (j > 0 && ids.size() != 1) ||
-              integer_values(calls[j].inputs[1]) != positions) {
+              integer_values(calls[j].inputs[1]) != positions ||
+              (num_inputs == 3 &&
+               integer_values(calls[j].inputs[2]) != last_index)) {
             throw std::runtime_error(
                 "sample " + std::to_string(i) + " call " + std::to_string(j) +
                 " is not a prefill or the next decode step");

backends/native/examples/hf_llm/test_export_hf_llm.py

@@ -407,6 +407,64 @@ def test_muse_glimmer_exports_text_only_with_an_offgraph_cache(self) -> None:
         ops = _ops_except_casts(ep)
         self.assertEqual(ops["kvcache.update_and_attend.default"], 4)
         self.assertNotIn("aten.eq.Scalar", ops)
+        selected = OffGraphLastTokenLM(model, 32, selected=True).eval()
+        ids = torch.tensor([[3, 4, 5]])
+        with torch.no_grad():
+            with selected.cache_scope():
+                got = selected(ids, torch.arange(3), torch.tensor([2]))
+            with offgraph.cache_scope():
+                offgraph.reset_cache()
+                want = offgraph(ids, torch.arange(3))
+        torch.testing.assert_close(got, want)
+
+    def test_offgraph_selected_logits_match_the_last_token(self) -> None:
+        torch.manual_seed(7)
+        register_offgraph_attention("offgraph")
+        model = _tiny()
+        model.config._attn_implementation = "offgraph"
+        last = OffGraphLastTokenLM(model, 32).eval()
+        selected = OffGraphLastTokenLM(model, 32, selected=True).eval()
+        prompts = [[3, 4, 5, 6, 7, 8], [9, 10]]
+
+        want = golden_samples(last, prompts, 2)
+        got = golden_samples(selected, prompts, 2)
+
+        for sample, expected in zip(got.samples, want.samples, strict=True):
+            for call, ref in zip(sample.calls, expected.calls, strict=True):
+                ids = call.inputs[0]
+                self.assertTrue(
+                    torch.equal(call.inputs[2], torch.tensor([ids.shape[1] - 1]))
+                )
+                self.assertTrue(torch.equal(call.outputs[0], ref.outputs[0]))
+        self.assertEqual(selected.metadata()["get_logits_to_keep_mode"], 2)
+        self.assertEqual(last.metadata()["get_logits_to_keep_mode"], 1)
+        seq = torch.export.Dim("seq", min=1, max=31)
+        rows = torch.export.Dim("rows", min=1, max=31)
+        with torch.no_grad():
+            ep = torch.export.export(
+                selected,
+                (torch.tensor([[3, 4, 5]]), torch.arange(3), torch.tensor([0, 2])),
+                dynamic_shapes={
+                    "input_ids": {1: seq},
+                    "cache_position": {0: seq},
+                    "logits_to_keep": {0: rows},
+                },
+                strict=True,
+            )
+            with selected.cache_scope():
+                logits = ep.module()(
+                    torch.tensor([[3, 4, 5]]), torch.arange(3), torch.tensor([0, 2])
+                )
+        self.assertEqual(tuple(logits.shape), (1, 2, 128))
+        self.assertEqual(len(ep.graph_signature.user_inputs), 3)
+
+    def test_selected_logits_need_an_offgraph_cache(self) -> None:
+        argv = ["export_hf_llm", "--model", "unused", "--logits_to_keep", "selected"]
+        with mock.patch.object(sys, "argv", argv):
+            with mock.patch.object(export_hf_llm, "_load") as load:
+                with self.assertRaises(SystemExit):
+                    export_hf_llm.main()
+        load.assert_not_called()
 
     def test_offgraph_cache_excludes_in_graph_cache_flags(self) -> None:
         for flags in (
@@ -536,6 +594,7 @@ def test_offgraph_metadata_gives_each_layer_its_window(self) -> None:
                         "get_head_dims": [16, 16],
                         "get_windows": windows,
                         "get_max_context_len": 64,
+                        "get_logits_to_keep_mode": 1,
                     },
                 )
                 self.assertEqual(