[LLM] Extract text-serving contracts and helpers

Prepare a batching-backed serving path without changing current worker behavior.
Share the existing prompt-history planner and define experimental serving types
for subsequent runtime and transport changes.

Test Plan: All 8 existing worker prompt-planner cases passed. The serving
contract and planner headers passed standalone C++17 compilation. Behavioral
serving tests are introduced alongside their implementations in later commits.

Authored with AI assistance from Devmate.

Author: metascroy

CMakeLists.txt

@@ -1015,6 +1015,8 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM)
   list(APPEND _executorch_extensions extension_llm_batching
        extension_llm_batching_module
   )
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/serving)
+  list(APPEND _executorch_extensions extension_llm_serving)
 endif()
 
 if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)

examples/llm_server/cpp/BUCK

@@ -13,6 +13,7 @@ fbcode_target(_kind = runtime.cxx_library,
     header_namespace = "executorch/examples/llm_server/cpp",
     exported_deps = [
         "//executorch/extension/llm/runner:runner_lib",
+        "//executorch/extension/llm/serving:prompt_history",
         "//pytorch/tokenizers:headers",
     ],
     external_deps = ["nlohmann_json"],

examples/llm_server/cpp/worker_prefill_plan.h

@@ -8,65 +8,4 @@
 
 #pragma once
 
-// Pure warm-resume prefill decision for the model worker (no model/session/ET
-// dependency, so it is unit-testable in isolation). Given a named session's
-// resident token ids (exactly the tokens currently in its KV/recurrent state),
-// the new request's prompt token ids, and whether the session is dirty, decide
-// whether to reset + full-prefill or to keep the state and prefill only the
-// suffix. The decision is exact-token (never string / retokenized text), so a
-// kSuffix plan is always a correct continuation; anything uncertain falls back
-// to kFull. See worker_loop.h for how the plan is executed.
-
-#include <cstddef>
-#include <cstdint>
-#include <vector>
-
-namespace executorch {
-namespace extension {
-namespace llm {
-
-struct PrefillPlan {
-  enum Action {
-    kFull, // reset + prefill the whole prompt
-    kSuffix // keep state, prefill prompt_ids[suffix_start:] at pos>0
-  } action;
-  size_t suffix_start; // index in prompt_ids where prefill begins (0 for kFull)
-  // Reported as session_reset_reason: "new" (no resident), "exact_prefix"
-  // (suffix reuse), "dirty", "mismatch", "equal" (prompt == resident).
-  const char* reason;
-};
-
-inline PrefillPlan plan_prefill(
-    const std::vector<uint64_t>& resident,
-    const std::vector<uint64_t>& prompt,
-    bool dirty) {
-  if (dirty) {
-    return {PrefillPlan::kFull, 0, "dirty"};
-  }
-  if (resident.empty()) {
-    return {PrefillPlan::kFull, 0, "new"};
-  }
-  if (prompt.size() < resident.size()) {
-    return {PrefillPlan::kFull, 0, "mismatch"};
-  }
-  for (size_t i = 0; i < resident.size(); ++i) {
-    if (prompt[i] != resident[i]) {
-      return {PrefillPlan::kFull, 0, "mismatch"};
-    }
-  }
-  if (prompt.size() == resident.size()) {
-    // Prompt is exactly the resident state (no new tokens). The ideal would be
-    // to skip prefill and decode straight from the session's pending token, but
-    // the LLMSession API exposes no "is there a valid pending token?" query, so
-    // we conservatively reset + full prefill rather than risk
-    // prefill_tokens([]) or decoding from a stale/absent pending. Rare in
-    // practice (a new turn adds tokens); a session pending-state query is a
-    // possible later optimization.
-    return {PrefillPlan::kFull, 0, "equal"};
-  }
-  return {PrefillPlan::kSuffix, resident.size(), "exact_prefix"};
-}
-
-} // namespace llm
-} // namespace extension
-} // namespace executorch
+#include <executorch/extension/llm/serving/prompt_history.h>

extension/llm/runner/targets.bzl

@@ -37,11 +37,13 @@ def define_common_targets():
             exported_headers = [
                 "model_metadata.h",
                 "stats.h",
+                "text_stream.h",
                 "util.h",
             ],
             visibility = ["PUBLIC"],
             exported_deps = [
                 ":constants",
+                "//pytorch/tokenizers:headers",
                  "//executorch/extension/module:module" + aten_suffix,
                  "//executorch/extension/tensor:tensor" + aten_suffix,
             ],

extension/llm/serving/BUCK

@@ -0,0 +1,13 @@
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+# The legacy worker only needs the pure planner, not serving or batching types.
+fbcode_target(_kind = runtime.cxx_library,
+    name = "prompt_history",
+    exported_headers = ["prompt_history.h"],
+    exported_deps = ["//executorch/runtime/platform:platform"],
+    header_namespace = "executorch/extension/llm/serving",
+    visibility = ["PUBLIC"],
+)

extension/llm/serving/CMakeLists.txt

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+add_library(extension_llm_serving INTERFACE)
+target_compile_features(extension_llm_serving INTERFACE cxx_std_17)
+target_include_directories(
+  extension_llm_serving INTERFACE ${_common_include_directories}
+)
+
+install(TARGETS extension_llm_serving EXPORT ExecuTorchTargets)
+install(FILES types.h prompt_history.h
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/serving
+)

extension/llm/serving/prompt_history.h

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+// Pure exact-token continuation decision. `resident` is the history against
+// which the full prompt is compared. It can be committed KV/recurrent tokens,
+// or logical history including a pending token if the caller's continuation
+// path feeds that token before the suffix. Dirty or uncertain histories replay
+// the full prompt; this helper neither inspects nor changes execution state.
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include <executorch/runtime/platform/compiler.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+
+struct ET_EXPERIMENTAL PrefillPlan {
+  enum Action {
+    kFull, // reset + prefill the whole prompt
+    kSuffix // keep state, prefill prompt_ids[suffix_start:] at pos>0
+  } action;
+  size_t suffix_start; // index in prompt_ids where prefill begins (0 for kFull)
+  // Reported as session_reset_reason: "new" (no resident), "exact_prefix"
+  // (suffix reuse), "dirty", "mismatch", "equal" (prompt == resident).
+  const char* reason;
+};
+
+ET_EXPERIMENTAL inline PrefillPlan plan_prefill(
+    const std::vector<uint64_t>& resident,
+    const std::vector<uint64_t>& prompt,
+    bool dirty) {
+  if (dirty) {
+    return {PrefillPlan::kFull, 0, "dirty"};
+  }
+  if (resident.empty()) {
+    return {PrefillPlan::kFull, 0, "new"};
+  }
+  if (prompt.size() < resident.size()) {
+    return {PrefillPlan::kFull, 0, "mismatch"};
+  }
+  for (size_t i = 0; i < resident.size(); ++i) {
+    if (prompt[i] != resident[i]) {
+      return {PrefillPlan::kFull, 0, "mismatch"};
+    }
+  }
+  if (prompt.size() == resident.size()) {
+    // Equal histories conservatively replay: this helper cannot establish that
+    // decoding without a non-empty suffix is supported or safe for the caller.
+    return {PrefillPlan::kFull, 0, "equal"};
+  }
+  return {PrefillPlan::kSuffix, resident.size(), "exact_prefix"};
+}
+
+} // namespace llm
+} // namespace extension
+} // namespace executorch

extension/llm/serving/types.h

@@ -0,0 +1,119 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <variant>
+#include <vector>
+
+#include <executorch/extension/llm/batching/types.h>
+#include <executorch/runtime/platform/compiler.h>
+
+namespace executorch {
+namespace extension {
+namespace llm {
+namespace serving {
+
+struct ET_EXPERIMENTAL TextSegment {
+  std::string text;
+};
+
+struct ET_EXPERIMENTAL TokenIdsSegment {
+  std::vector<batching::Token> ids;
+};
+
+// Extend this variant with typed segments when another modality is supported.
+using PromptSegment ET_EXPERIMENTAL =
+    std::variant<TextSegment, TokenIdsSegment>;
+
+// The full prompt, not a session delta. Preparation encodes each text segment
+// separately and appends ID segments verbatim, in order, without adding
+// BOS/EOS. Segment boundaries therefore matter even between adjacent text
+// segments.
+struct ET_EXPERIMENTAL PromptInput {
+  std::vector<PromptSegment> segments;
+};
+
+struct ET_EXPERIMENTAL GenerationOptions {
+  // Unset uses the available context, or the service default if unknown.
+  std::optional<std::int32_t> max_new_tokens;
+  batching::SamplingParams sampling;
+  std::vector<batching::Token> stop_tokens;
+  std::vector<std::string> stop_strings;
+  std::optional<std::uint64_t> seed;
+};
+
+enum class ET_EXPERIMENTAL FinishReason {
+  Stop, // EOS, a stop token, or a stop string
+  Length,
+  Cancelled,
+  Failed,
+};
+
+enum class ET_EXPERIMENTAL ErrorCode {
+  InvalidArgument,
+  NotReady,
+  SessionNotFound,
+  SessionBusy,
+  CapacityExceeded,
+  Internal,
+};
+
+struct ET_EXPERIMENTAL ServingError {
+  ErrorCode code = ErrorCode::Internal;
+  std::string message;
+};
+
+struct ET_EXPERIMENTAL GenerationStats {
+  std::size_t prompt_tokens = 0;
+  // Non-terminal generated tokens, including tokens hidden by a string stop.
+  // Neither this count nor visible text describes the session's logical
+  // history.
+  std::size_t completion_tokens = 0;
+  std::size_t reused_prompt_tokens = 0;
+  // Actually consumed prompt tokens; may be partial on cancellation or failure.
+  std::size_t prefilled_prompt_tokens = 0;
+  double prefill_ms = 0.0;
+  double decode_ms = 0.0;
+  double total_ms = 0.0;
+  std::string session_reset_reason;
+  // Absent when output trimming or cancellation makes exact replay unsafe.
+  std::optional<std::vector<batching::Token>> generated_token_ids;
+};
+
+struct ET_EXPERIMENTAL TextEvent {
+  std::string text;
+};
+
+struct ET_EXPERIMENTAL TerminalEvent {
+  FinishReason finish_reason = FinishReason::Stop;
+  GenerationStats stats;
+  // Present exactly when finish_reason is Failed.
+  std::optional<ServingError> error;
+};
+
+// Per-request order: zero or more text events, then exactly one terminal event.
+using GenerationEvent ET_EXPERIMENTAL = std::variant<TextEvent, TerminalEvent>;
+
+struct ET_EXPERIMENTAL ServingInfo {
+  bool ready = false;
+  std::size_t max_context_length = 0;
+  std::size_t max_sessions = 0;
+  // Logical serving slots, including slots reserved during reset/reopen.
+  // This is not a count of physically live executor sessions.
+  std::size_t active_sessions = 0;
+};
+
+} // namespace serving
+} // namespace llm
+} // namespace extension
+} // namespace executorch