[UPDATE] Update

[ghstack-poisoned]

Author: mergennachin

extension/llm/runner/llm_session.h

@@ -8,8 +8,8 @@
 
 // Model-agnostic Engine/Session interfaces. Model-specific execution lives in
 // adapters that implement these (TextLLMSession over TextLLMRunner today;
-// Gemma4Session etc. later); the server and pybind layer depend only on these
-// interfaces, never on a concrete runner.
+// Gemma4Session etc. later); the serving code (HTTP control plane + C++ worker
+// binaries) depends only on these interfaces, never on a concrete runner.
 
 #pragma once
 

extension/llm/runner/test/test_util.cpp

@@ -18,6 +18,7 @@ namespace {
 using ::executorch::aten::ScalarType;
 using ::executorch::extension::make_tensor_ptr;
 using ::executorch::extension::llm::convert_to_bfloat16;
+using ::executorch::extension::llm::stop_safe_prefix_len;
 using ::executorch::extension::llm::utf8_complete_prefix_len;
 
 class ConvertToBFloat16Test : public ::testing::Test {
@@ -86,4 +87,39 @@ TEST(Utf8CompletePrefixLenTest, HandlesAsciiAndMultiByteBoundaries) {
   EXPECT_EQ(utf8_complete_prefix_len("\x80"), 1u);
 }
 
+TEST(StopSafePrefixLenTest, NoStopsEmitsEverything) {
+  bool hit = true;
+  EXPECT_EQ(stop_safe_prefix_len("hello world", {}, hit), 11u);
+  EXPECT_FALSE(hit);
+}
+
+TEST(StopSafePrefixLenTest, StopFoundReturnsEarliestOffsetAndExcludesIt) {
+  bool hit = false;
+  // "STOP" begins at offset 6; emit "Hello " (6 bytes), drop the stop and rest.
+  EXPECT_EQ(stop_safe_prefix_len("Hello STOP there", {"STOP"}, hit), 6u);
+  EXPECT_TRUE(hit);
+  // Earliest of several wins.
+  hit = false;
+  EXPECT_EQ(stop_safe_prefix_len("aXbY", {"Y", "X"}, hit), 1u);
+  EXPECT_TRUE(hit);
+}
+
+TEST(StopSafePrefixLenTest, HoldsBackPossiblePartialStopTail) {
+  bool hit = false;
+  // No full stop yet, but the trailing "ST" could become "STOP": hold back
+  // len("STOP")-1 == 3 bytes, so of "hi ST" (5 bytes) only "hi" (2) is safe.
+  EXPECT_EQ(stop_safe_prefix_len("hi ST", {"STOP"}, hit), 2u);
+  EXPECT_FALSE(hit);
+}
+
+TEST(StopSafePrefixLenTest, HoldBackSnapsToUtf8Boundary) {
+  bool hit = false;
+  // "ab" + "€"(3 bytes). Stop "XX" => hold back 1 byte, which would land inside
+  // the euro sign; snap down so the multi-byte char isn't split.
+  const std::string text = "ab\xe2\x82\xac";
+  const size_t safe = stop_safe_prefix_len(text, {"XX"}, hit);
+  EXPECT_FALSE(hit);
+  EXPECT_EQ(safe, 2u); // only "ab"; the € is held whole
+}
+
 } // namespace

extension/llm/runner/util.h

@@ -13,6 +13,7 @@
 #include <executorch/runtime/platform/compiler.h>
 #include <stdio.h>
 #include <time.h>
+#include <algorithm>
 #include <cctype>
 #include <string>
 #include <vector>
@@ -98,6 +99,55 @@ ET_EXPERIMENTAL size_t inline utf8_complete_prefix_len(const std::string& s) {
   return i;
 }
 
+// How many leading bytes of `text` a streaming consumer may safely emit given a
+// set of `stops` strings, and whether a stop was hit (`stop_hit`).
+//   * If any stop occurs, returns the byte offset of the EARLIEST occurrence and
+//     sets stop_hit=true — text before it is safe; the stop and everything after
+//     are dropped (the stop is excluded from output).
+//   * Otherwise returns the length minus the longest possible partial-stop tail
+//     (max(len(stop))-1 bytes), snapped DOWN to a UTF-8 boundary so a multi-byte
+//     character is never split; stop_hit=false. Holding back that tail lets a
+//     stop that straddles the next piece still be caught.
+// `text` is expected to be complete-UTF-8 (e.g. the assembled output of
+// utf8_complete_prefix_len). Empty `stops` => emit everything, no hold-back.
+ET_EXPERIMENTAL size_t inline stop_safe_prefix_len(
+    const std::string& text,
+    const std::vector<std::string>& stops,
+    bool& stop_hit) {
+  stop_hit = false;
+  if (stops.empty()) {
+    return text.size();
+  }
+  size_t earliest = std::string::npos;
+  size_t max_len = 0;
+  for (const auto& s : stops) {
+    if (s.empty()) {
+      continue;
+    }
+    max_len = std::max(max_len, s.size());
+    const size_t p = text.find(s);
+    if (p != std::string::npos &&
+        (earliest == std::string::npos || p < earliest)) {
+      earliest = p;
+    }
+  }
+  if (earliest != std::string::npos) {
+    stop_hit = true;
+    return earliest;
+  }
+  const size_t hold = max_len > 0 ? max_len - 1 : 0;
+  if (text.size() <= hold) {
+    return 0;
+  }
+  size_t end = text.size() - hold;
+  // Don't cut in the middle of a UTF-8 character: back up over continuation
+  // bytes (10xxxxxx).
+  while (end > 0 && (static_cast<unsigned char>(text[end]) & 0xC0) == 0x80) {
+    --end;
+  }
+  return end;
+}
+
 // ----------------------------------------------------------------------------
 // utilities: time
 

extension/llm/server/cpp/CMakeLists.txt

@@ -45,6 +45,19 @@ set(link_libraries executorch gflags)
 list(APPEND link_libraries optimized_native_cpu_ops_lib cpublas eigen_blas)
 executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib)
 
+# Custom + quantized kernels that export_llm models need, whole-archived so the
+# static op registrations survive the linker: llama::custom_sdpa (from
+# use_sdpa_with_kv_cache) and quantized_decomposed ops (from quantized exports).
+# Without these the model loads but execution fails with "Missing operator".
+if(TARGET custom_ops)
+  executorch_target_link_options_shared_lib(custom_ops)
+  list(APPEND link_libraries custom_ops)
+endif()
+if(TARGET quantized_ops_lib)
+  list(APPEND link_libraries quantized_kernels quantized_ops_lib)
+  executorch_target_link_options_shared_lib(quantized_ops_lib)
+endif()
+
 # Extensions (Engine/Session lives in extension_llm_runner)
 list(
   APPEND

extension/llm/server/cpp/text_llm_worker.cpp

@@ -16,10 +16,11 @@
 // Protocol (one JSON object per line; identical to worker_client.py):
 //   worker -> stdout, once:         {"ready": true}
 //   client -> stdin,  per request:  {"prompt": str, "max_new_tokens": int,
-//                                    "temperature": float}
+//                                    "temperature": float, "stop": [str, ...]}
 //   worker -> stdout, per request:  {"token": str} *   (streamed)
 //                                   {"done": true, "prompt_tokens": int,
-//                                    "completion_tokens": int}
+//                                    "completion_tokens": int,
+//                                    "finish_reason": "stop" | "length"}
 //                               or  {"error": str}
 //
 // stdout carries ONLY protocol JSON; all logs go to stderr (ET_LOG). One
@@ -74,6 +75,11 @@ void handle_request(
   const std::string prompt = req.at("prompt").get<std::string>();
   int64_t max_new = req.value("max_new_tokens", static_cast<int64_t>(-1));
   const float temperature = req.value("temperature", 0.0f);
+  // Stop strings (the request's `stop` sequences): terminate at the token
+  // boundary where one appears so we don't generate to EOS/max_new past it. The
+  // control plane also enforces these as a backstop.
+  const std::vector<std::string> stops =
+      req.value("stop", std::vector<std::string>{});
 
   if (session.reset() != Error::Ok) {
     throw std::runtime_error("session reset failed");
@@ -90,12 +96,21 @@ void handle_request(
   }
   const int64_t num_prompt = static_cast<int64_t>(ids.size());
 
-  if (max_new <= 0) {
-    // Fill the remaining context window when the client doesn't bound it.
-    const auto it = metadata.find(llm::kMaxContextLen);
-    max_new = (it != metadata.end())
-        ? std::max<int64_t>(1, it->second - num_prompt)
-        : 2048;
+  // Bound generation to the context window: default to filling the remaining
+  // room, and clamp an explicit max_new_tokens too, so decode never steps past
+  // the window (which would error mid-generation after partial output).
+  const auto ctx_it = metadata.find(llm::kMaxContextLen);
+  if (ctx_it != metadata.end()) {
+    const int64_t room = ctx_it->second - num_prompt;
+    if (room <= 0) {
+      throw std::runtime_error(
+          "prompt fills the context window; no room to generate");
+    }
+    if (max_new <= 0 || max_new > room) {
+      max_new = room;
+    }
+  } else if (max_new <= 0) {
+    max_new = 2048;
   }
 
   llm::SamplingConfig sampling;
@@ -104,32 +119,55 @@ void handle_request(
     throw std::runtime_error("prefill failed");
   }
 
-  std::string buf; // holds bytes not yet forming a complete UTF-8 prefix
+  std::string buf; // bytes not yet forming a complete UTF-8 prefix
+  std::string pending; // complete-UTF-8 text held back for stop-string matching
   int64_t num_generated = 0;
+  std::string finish = "length"; // EOS or stop string -> "stop"
+  bool stop_string = false; // a request stop string was matched
   for (int64_t step = 0; step < max_new; ++step) {
     auto step_result = session.decode_one(sampling);
     if (step_result.error() != Error::Ok) {
       throw std::runtime_error("decode failed");
     }
     const auto& d = step_result.get();
     if (d.is_terminal) {
-      break; // terminal step: not generated output
+      finish = "stop";
+      break; // terminal step (EOS / cooperative stop): not emitted or counted
     }
     ++num_generated;
     buf += d.text_piece;
     const size_t cut = llm::utf8_complete_prefix_len(buf);
     if (cut > 0) {
-      emit({{"token", buf.substr(0, cut)}});
+      pending += buf.substr(0, cut);
       buf.erase(0, cut);
     }
+    bool stop_hit = false;
+    const size_t safe = llm::stop_safe_prefix_len(pending, stops, stop_hit);
+    if (safe > 0) {
+      emit({{"token", pending.substr(0, safe)}});
+      pending.erase(0, safe);
+    }
+    if (stop_hit) {
+      finish = "stop"; // reached a stop string: drop it and everything after
+      stop_string = true;
+      break;
+    }
   }
-  if (!buf.empty()) {
-    emit({{"token", buf}}); // flush any trailing bytes (replaced if incomplete)
+  if (!stop_string) {
+    // EOS or length: flush held-back text + any trailing incomplete bytes
+    // (replaced if invalid). A stop-string hit drops the remainder instead.
+    pending += buf;
+    if (!pending.empty()) {
+      emit({{"token", pending}});
+    }
   }
+  // finish_reason: "stop" if the model emitted EOS or hit a stop string, else
+  // "length" — it ran to max_new (possibly clamped to the context window).
   emit(
       {{"done", true},
        {"prompt_tokens", num_prompt},
-       {"completion_tokens", num_generated}});
+       {"completion_tokens", num_generated},
+       {"finish_reason", finish}});
 }
 
 } // namespace

extension/llm/server/python/errors.py

@@ -30,13 +30,26 @@ def body(self) -> dict:
 
 
 class ContextLengthExceeded(APIError):
-    def __init__(self, num_tokens: int, max_context: int):
-        super().__init__(
-            status=400,
-            message=(
+    def __init__(self, num_tokens: int, max_context: int, completion_tokens: int = 0):
+        # completion_tokens > 0: the prompt fits but prompt + requested
+        # max_tokens would run past the window — reject up front rather than
+        # fail (or truncate) mid-generation.
+        if completion_tokens > 0:
+            message = (
+                f"This model's maximum context length is {max_context} tokens. "
+                f"However, you requested {num_tokens + completion_tokens} tokens "
+                f"({num_tokens} in the messages, {completion_tokens} in the "
+                f"completion). Please reduce the length of the messages or "
+                f"completion."
+            )
+        else:
+            message = (
                 f"This model's maximum context length is {max_context} tokens, "
                 f"but the request has {num_tokens} prompt tokens."
-            ),
+            )
+        super().__init__(
+            status=400,
+            message=message,
             err_type="invalid_request_error",
             code="context_length_exceeded",
         )

extension/llm/server/python/runner_pool.py

@@ -38,6 +38,8 @@
 class GenStats:
     prompt_tokens: int = 0
     completion_tokens: int = 0
+    # Worker-reported stop reason ("stop" | "length"), or None if not reported.
+    finish_reason: Optional[str] = None
 
 
 class RunnerPool:
@@ -98,6 +100,7 @@ def token_cb(token: str) -> None:
         def stats_cb(s) -> None:
             out_stats.prompt_tokens = s.num_prompt_tokens
             out_stats.completion_tokens = s.num_generated_tokens
+            out_stats.finish_reason = getattr(s, "finish_reason", None)
 
         def run() -> None:
             try: