[UPDATE] Update

[ghstack-poisoned]

Author: mergennachin

extension/llm/runner/llm_runner_helper.cpp

@@ -325,7 +325,12 @@ static std::unique_ptr<TextLLMRunner> assemble_text_llm_runner(
       temperature);
 }
 
-std::unique_ptr<TextLLMRunner> create_text_llm_runner_from_program(
+// Builds a TextLLMRunner over an already-loaded Program: the runner's Module
+// reuses `program` while owning its own method state and KV cache. File-local —
+// the per-session construction path for TextLLMEngine (which keeps the backing
+// DataLoader alive for the runners' lifetime). External callers go through
+// LLMEngine -> LLMSession, not a raw shared-Program runner.
+static std::unique_ptr<TextLLMRunner> create_text_llm_runner_from_program(
     std::shared_ptr<Program> program,
     std::unique_ptr<::tokenizers::Tokenizer> tokenizer,
     float temperature,
@@ -413,6 +418,11 @@ Error TextLLMSession::reset() {
 void TextLLMSession::stop() {
   runner_->stop();
 }
+
+std::unique_ptr<LLMSession> make_text_llm_session(
+    std::unique_ptr<TextLLMRunner> runner) {
+  return std::make_unique<TextLLMSession>(std::move(runner));
+}
 } // namespace detail
 
 TextLLMEngine::TextLLMEngine(
@@ -504,8 +514,7 @@ TextLLMEngine::create_session() {
     ET_LOG(Error, "TextLLMEngine: failed to build session runner");
     return Error::InvalidState;
   }
-  return std::unique_ptr<LLMSession>(
-      std::make_unique<detail::TextLLMSession>(std::move(runner)));
+  return detail::make_text_llm_session(std::move(runner));
 }
 
 std::unique_ptr<MultimodalRunner> create_multimodal_runner(

extension/llm/runner/llm_runner_helper.h

@@ -142,33 +142,6 @@ ET_EXPERIMENTAL std::unique_ptr<TextLLMRunner> create_text_llm_runner(
     const std::string& method_name = "forward",
     Module::LoadMode load_mode = Module::LoadMode::MmapUseMlockIgnoreErrors);
 
-/**
- * @brief Creates a TextLLMRunner over an already-loaded Program.
- *
- * Unlike create_text_llm_runner(model_path, ...), this does not load the model
- * file again: the resulting runner's Module reuses `program` while owning its
- * own method state and KV cache. This is the per-session construction path for
- * TextLLMEngine — N sessions reuse one loaded Program but isolate their mutable
- * KV state. Whether they also avoid re-materializing packed weights per session
- * is backend-dependent (serving_capacity() is authoritative).
- *
- * The caller must keep the DataLoader backing `program` alive for the lifetime
- * of every runner created from it (TextLLMEngine holds the loader Module).
- *
- * @param program Shared, already-loaded program.
- * @param tokenizer Initialized tokenizer instance (owned by the new runner).
- * @param temperature Optional temperature (deprecated; prefer
- * GenerationConfig).
- * @param method_name Name of the method to execute in the model.
- * @return std::unique_ptr<TextLLMRunner> on success, or nullptr on failure.
- */
-ET_EXPERIMENTAL std::unique_ptr<TextLLMRunner>
-create_text_llm_runner_from_program(
-    std::shared_ptr<Program> program,
-    std::unique_ptr<::tokenizers::Tokenizer> tokenizer,
-    float temperature = -1.0f,
-    const std::string& method_name = "forward");
-
 /**
  * @brief Engine for multi-session text generation over one loaded Program.
  *
@@ -225,6 +198,20 @@ class ET_EXPERIMENTAL TextLLMEngine : public LLMEngine {
   std::unordered_map<std::string, int64_t> metadata_;
 };
 
+namespace detail {
+// Implementation detail (not a public API): wraps a TextLLMRunner in an
+// LLMSession (the runner -> session seam). The supported entry point is
+// LLMEngine::create_session(); this exists only so TextLLMEngine can build its
+// sessions and so unit tests can drive the runner's token-step primitives
+// through the public LLMSession surface (the concrete adapter type is private).
+// Do not depend on wrapping arbitrary runners.
+//
+// @param runner A loaded TextLLMRunner; ownership transfers to the session.
+// @return std::unique_ptr<LLMSession> wrapping `runner`.
+std::unique_ptr<LLMSession> make_text_llm_session(
+    std::unique_ptr<TextLLMRunner> runner);
+} // namespace detail
+
 /**
  * @brief Creates a MultimodalRunner instance with dependency injection
  *

extension/llm/runner/llm_session.h

@@ -36,12 +36,19 @@ struct SamplingConfig {
 };
 
 /// One decoded step: the exact sampled token id (for prefix-cache id tracking
-/// and batching), its decoded text piece (raw bytes; may be a partial UTF-8
-/// sequence the caller assembles), and whether it is an EOS token.
+/// and batching) and its decoded text piece (raw bytes; may be a partial UTF-8
+/// sequence the caller assembles).
+///
+/// `is_eos` is literal: the sampled token is an end-of-sequence token (use it
+/// for the "stop" finish reason, metrics, cache/accounting). `is_terminal` is
+/// the loop signal: generation ended at this step — either EOS or a cooperative
+/// stop() took effect. A decode loop should end when is_terminal is set; every
+/// EOS step is also terminal, but a stop step is terminal without being EOS.
 struct DecodeResult {
   uint64_t token_id;
   std::string text_piece;
   bool is_eos;
+  bool is_terminal;
 };
 
 /// How many physical sessions an engine can host, so the server admits logical
@@ -96,8 +103,10 @@ class ET_EXPERIMENTAL LLMSession {
 
   /// Request that a decode_one() loop stop. This is a TOKEN-BOUNDARY,
   /// cooperative stop: it is safe to call from another thread, but it does not
-  /// abort a decode_one() that is already running — it takes effect before the
-  /// next decode_one() (the loop driver checks between tokens).
+  /// abort a decode_one() that is already running. It takes effect at the next
+  /// decode_one(), which then returns a terminal step (is_terminal set, is_eos
+  /// false) without forwarding a new token. The stop is cleared by the next
+  /// prefill_tokens() or reset().
   virtual void stop() = 0;
 };
 

extension/llm/runner/targets.bzl

@@ -113,11 +113,16 @@ def define_common_targets():
             name = "runner_lib" + aten_suffix,
             exported_headers = [
                 "text_llm_runner.h",
-                "text_llm_session.h",
                 "llm_runner_helper.h",
                 "llm_session.h",
                 "constants.h",
             ],
+            # Internal: the detail::TextLLMSession adapter (sole friended caller
+            # of TextLLMRunner's token-step hooks). Private so dependents reach
+            # it only through LLMEngine/LLMSession + make_text_llm_session().
+            headers = [
+                "text_llm_session.h",
+            ],
             srcs = [
                 "text_llm_runner.cpp",
                 "llm_runner_helper.cpp",