[LLM serving 1/x] Define text-serving contracts and share prompt planning (pytorch#23270)

Prepare a batching-backed serving layer without changing existing worker
behavior.

Introduce experimental C++ types for text and exact-token prompts,
generation options, streamed text and terminal events, errors, usage
statistics, and serving metadata. Reuse batching’s token and sampling
types.

Move the existing exact-token prompt-history planner into the shared
serving layer, retaining the legacy worker’s include path and unchanged
continuation/replay decisions. Add CMake and Buck build wiring.

This PR introduces contracts and shared planning only. Runtime
implementation and behavioral tests follow in subsequent PRs; no
scheduling or generation behavior changes here.

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/test_worker_prefill_plan.cpp

@@ -13,7 +13,6 @@
 #include <executorch/examples/llm_server/cpp/worker_prefill_plan.h>
 
 #include <cstdio>
-#include <cstring>
 #include <string>
 #include <vector>
 
@@ -30,15 +29,15 @@ void expect(
     size_t suffix_start,
     const char* reason) {
   bool ok = p.action == action && p.suffix_start == suffix_start &&
-      std::strcmp(p.reason, reason) == 0;
+      p.reason == reason;
   if (!ok) {
     ++g_failures;
     printf(
         "  [FAIL] %s: got action=%d suffix_start=%zu reason=%s\n",
         name,
         (int)p.action,
         p.suffix_start,
-        p.reason);
+        p.reason.c_str());
   } else {
     printf("  [PASS] %s\n", name);
   }

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,
             ],
@@ -101,12 +103,23 @@ def define_common_targets():
         )
 
         runtime.cxx_library(
-            name = "multimodal_runner_lib" + aten_suffix,
+            name = "multimodal_input" + aten_suffix,
             exported_headers = [
                 "audio.h",
                 "image.h",
-                "wav_loader.h",
                 "multimodal_input.h",
+            ],
+            visibility = ["PUBLIC"],
+            exported_deps = [
+                "//executorch/extension/tensor:tensor" + aten_suffix,
+                "//executorch/runtime/platform:platform",
+            ],
+        )
+
+        runtime.cxx_library(
+            name = "multimodal_runner_lib" + aten_suffix,
+            exported_headers = [
+                "wav_loader.h",
                 "multimodal_runner.h",
                 "multimodal_prefiller.h",
                 "multimodal_decoder_runner.h",
@@ -118,6 +131,7 @@ def define_common_targets():
                 ":text_decoder_runner" + aten_suffix,
                 ":text_prefiller" + aten_suffix,
                 ":image_prefiller" + aten_suffix,
+                ":multimodal_input" + aten_suffix,
                 ":text_token_generator" + 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,28 @@
+# 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}
+)
+
+target_link_libraries(extension_llm_serving INTERFACE executorch_core)
+
+install(TARGETS extension_llm_serving EXPORT ExecuTorchTargets)
+install(FILES types.h prompt_history.h
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/serving
+)
+
+# Public serving inputs do not require the legacy runner.
+if(NOT EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER)
+  install(
+    FILES ${CMAKE_CURRENT_SOURCE_DIR}/../runner/multimodal_input.h
+          ${CMAKE_CURRENT_SOURCE_DIR}/../runner/image.h
+          ${CMAKE_CURRENT_SOURCE_DIR}/../runner/audio.h
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/runner
+  )
+endif()

extension/llm/serving/prompt_history.h

@@ -0,0 +1,67 @@
+/*
+ * 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 <string>
+#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).
+  std::string 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