From 95d13f300431a4c9a01cf25c6e253cfd7bb28e60 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 12 Nov 2026 11:22:19 -0800 Subject: [PATCH] Select the logits rows an off-graph export returns The off-graph export returned the last token's logits, which a forward over one sequence needs. A batched forward packs several sequences' tokens, and each wants the logits of its own last token. --logits_to_keep selected adds a third int64[K] input of token indices, passed to Hugging Face's logits_to_keep, so the method returns those rows ([1, K, vocab]) without running the LM head over every token. The package records the mode as get_logits_to_keep_mode, as MLX packages do, and Muse Glimmer's text-only module takes the index tensor too. The native runner feeds a selected method the last token's index, so these packages still run and verify their golden samples standalone. --- .../native/examples/hf_llm/export_hf_llm.py | 76 +++++++++++++++---- backends/native/examples/hf_llm/runtime/BUCK | 1 + .../examples/hf_llm/runtime/LlmRunner.cpp | 37 +++++++-- .../examples/hf_llm/test_export_hf_llm.py | 59 ++++++++++++++ 4 files changed, 152 insertions(+), 21 deletions(-) diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index 92d0dc3dfd2..a9e8239aaba 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/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") diff --git a/backends/native/examples/hf_llm/runtime/BUCK b/backends/native/examples/hf_llm/runtime/BUCK index 61772593fde..5dfca41daec 100644 --- a/backends/native/examples/hf_llm/runtime/BUCK +++ b/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"], diff --git a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp index 4ee7a854c4d..23af3bcd0ad 100644 --- a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp +++ b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -93,6 +94,9 @@ namespace { 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 @@ Generation generate( auto run_tokens = [&](const std::vector& 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(tokens.size()) - 1; + executable.resize_input(2, {1}); + executable.set_input(2, &last, 1, ptn::kLong); + } executable.execute(); const std::vector output_shape = executable.output_sizes(logits_output); @@ -536,8 +547,18 @@ int run_llm( 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 @@ int run_llm( 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 @@ int run_llm( " 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 @@ int run_llm( static_cast(j == 0 ? 0 : prompt.size() + j - 1); std::vector positions(per_token_positions ? count : 1); std::iota(positions.begin(), positions.end(), start); + const std::vector last_index = { + static_cast(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"); diff --git a/backends/native/examples/hf_llm/test_export_hf_llm.py b/backends/native/examples/hf_llm/test_export_hf_llm.py index 489571c5e34..dd468fcca65 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/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(