[UPDATE] Update

[ghstack-poisoned]

Author: mergennachin

examples/models/gemma4_31b/CMakeLists.txt

@@ -79,9 +79,7 @@ if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug")
 endif()
 
 if(EXECUTORCH_BUILD_CUDA)
-  add_executable(
-    gemma4_31b_worker gemma4_31b_worker.cpp gemma4_31b_engine.cpp
-  )
+  add_executable(gemma4_31b_worker gemma4_31b_worker.cpp gemma4_31b_engine.cpp)
   target_include_directories(
     gemma4_31b_worker PUBLIC ${_common_include_directories} ${_json_include}
   )

examples/models/gemma4_31b/README.md

@@ -181,9 +181,9 @@ python -m executorch.examples.models.gemma4_31b.serve \
     --max-sessions 1
 ```
 
-The launcher defaults to the Hermes `<tool_call>{...}</tool_call>` parser. Use
-`--tool-parser qwen` or `--tool-parser none` if the model/template you are
-testing emits a different tool-call format.
+The launcher defaults to Gemma's `<|tool_call>call:...<tool_call|>` parser. Use
+`--tool-parser hermes`, `--tool-parser qwen`, or `--tool-parser none` if the
+model/template you are testing emits a different tool-call format.
 
 Named sessions and warm resume require worker capacity above one. CUDA exports
 with `get_mutable_buffer_metadata` can use per-session mutable rebinding and

examples/models/gemma4_31b/gemma4_31b_engine.cpp

@@ -457,7 +457,6 @@ Result<std::unique_ptr<Gemma4_31BEngine>> Gemma4_31BEngine::create(
 
   auto eos_ids = get_eos_ids(tokenizer.get(), meta_module.get());
   eos_ids.insert(static_cast<uint64_t>(config.eos_id));
-  add_token_piece(tokenizer.get(), eos_ids, "<end_of_turn>");
   add_token_piece(tokenizer.get(), eos_ids, "<turn|>");
 
   const auto& metadata = metadata_result.get();

examples/models/gemma4_31b/gemma4_31b_worker.cpp

@@ -24,7 +24,7 @@ DEFINE_bool(
     warm_resume,
     true,
     "Warm append-only resume for named sessions when the engine supports them.");
-DEFINE_int32(bos_id, 2, "BOS token id to prepend to every Gemma prompt.");
+DEFINE_int32(bos_id, 2, "BOS token id to prepend to server-rendered prompts.");
 DEFINE_int32(eos_id, 1, "EOS token id (Gemma convention: 1).");
 
 namespace {

examples/models/gemma4_31b/serve.py

@@ -9,19 +9,32 @@
 import argparse
 import logging
 import os
+import re
 from pathlib import Path
 
 from executorch.extension.llm.server.python.chat_template import ChatTemplate
 from executorch.extension.llm.server.python.serving_chat import ServingChat
 from executorch.extension.llm.server.python.session_runtime import SessionRuntime
 from executorch.extension.llm.server.python.tool_parsers import (
+    GemmaToolCallDetector,
     HermesDetector,
     QwenFunctionCallDetector,
 )
 from executorch.extension.llm.server.python.worker_client import spawn_worker
 
 logger = logging.getLogger(__name__)
 
+_GEMMA_CHANNEL_SPECIALS = {"<|channel>", "<channel|>", "<|think|>"}
+_GEMMA_CHANNEL_BLOCK = re.compile(r"<\|channel>.*?<channel\|>", re.DOTALL)
+
+
+def _strip_gemma_channels(text: str) -> str:
+    text = _GEMMA_CHANNEL_BLOCK.sub("", text)
+    open_idx = text.find("<|channel>")
+    if open_idx != -1:
+        text = text[:open_idx]
+    return text.replace("<channel|>", "").replace("<|think|>", "").strip()
+
 
 def _default_worker_bin() -> str:
     repo_root = Path(__file__).resolve().parents[3]
@@ -62,6 +75,8 @@ def _spawn(args):
 
 
 def _tool_detector(name: str):
+    if name == "gemma":
+        return GemmaToolCallDetector
     if name == "hermes":
         return HermesDetector
     if name == "qwen":
@@ -72,7 +87,13 @@ def _tool_detector(name: str):
 
 
 def build_app_from_args(args):
-    template = ChatTemplate(args.hf_tokenizer)
+    template = ChatTemplate(
+        args.hf_tokenizer,
+        assistant_header="<|turn>model\n",
+        # Gemma's HF template starts with bos_token text. Strip that text before
+        # C++ tokenization; the worker prepends the numeric BOS id.
+        strip_rendered_bos=True,
+    )
     worker = _spawn(args)
     runtime = SessionRuntime(worker)
     serving = ServingChat(
@@ -82,6 +103,8 @@ def build_app_from_args(args):
         max_context=args.max_context,
         tool_detector_cls=_tool_detector(args.tool_parser),
         prompt_token_offset=1,
+        content_filter=_strip_gemma_channels,
+        content_filter_specials=_GEMMA_CHANNEL_SPECIALS,
     )
 
     from executorch.extension.llm.server.python.server import build_app
@@ -132,11 +155,17 @@ def main() -> None:
     )
     p.add_argument(
         "--tool-parser",
-        choices=("hermes", "qwen", "none"),
-        default="hermes",
+        choices=("gemma", "hermes", "qwen", "none"),
+        default="gemma",
         help="Tool-call format parser to apply to model output.",
     )
-    p.add_argument("--bos-id", type=int, default=2)
+    p.add_argument(
+        "--bos-id",
+        type=int,
+        default=2,
+        help="BOS token id to prepend in the worker. The launcher strips the "
+        "HF template's literal <bos> before C++ tokenization.",
+    )
     p.add_argument("--eos-id", type=int, default=1)
     p.add_argument(
         "--worker-bin",

examples/models/gemma4_31b/test_ondevice_serving.py

@@ -0,0 +1,64 @@
+# 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.
+
+import json
+import os
+import urllib.request
+
+import pytest
+
+from executorch.extension.llm.server.python.chat_template import ChatTemplate
+from executorch.extension.llm.server.python.protocol import ChatMessage
+
+_SERVER = os.environ.get("GEMMA_SERVER_URL")
+_HF_DIR = os.environ.get(
+    "GEMMA_HF_DIR", "/home/mnachin/local/scripts/models/gemma-4-31B-it-HQQ-INT4"
+)
+
+
+pytestmark = pytest.mark.skipif(
+    not _SERVER or not os.path.isdir(_HF_DIR),
+    reason="set GEMMA_SERVER_URL and GEMMA_HF_DIR to run Gemma on-device tests",
+)
+
+
+def _post(path: str, payload: dict) -> dict:
+    data = json.dumps(payload).encode("utf-8")
+    req = urllib.request.Request(
+        _SERVER.rstrip("/") + path,
+        data=data,
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+    with urllib.request.urlopen(req, timeout=120) as resp:
+        return json.loads(resp.read().decode("utf-8"))
+
+
+def test_prompt_tokens_match_real_template_with_numeric_bos_prefix():
+    _ = pytest.importorskip("transformers")
+    from transformers import AutoTokenizer
+
+    template = ChatTemplate(
+        _HF_DIR,
+        assistant_header="<|turn>model\n",
+        strip_rendered_bos=True,
+    )
+    tok = AutoTokenizer.from_pretrained(_HF_DIR)
+    messages = [ChatMessage(role="user", content="Say ok.")]
+    rendered = template.render(messages)
+    expected_ids = [tok.bos_token_id] + tok.encode(rendered, add_special_tokens=False)
+
+    body = _post(
+        "/v1/chat/completions",
+        {
+            "model": "gemma4_31b",
+            "messages": [{"role": "user", "content": "Say ok."}],
+            "max_tokens": 1,
+            "temperature": 0,
+            "session_id": "gemma-bos-regression",
+        },
+    )
+    assert body["usage"]["prompt_tokens"] == len(expected_ids)

examples/models/gemma4_31b/test_serve.py

@@ -89,6 +89,19 @@ def test_spawn_defaults_worker_bin_and_omits_empty_data_path(monkeypatch):
     assert "--warm_resume=false" in cmd
 
 
+def test_strip_gemma_channels_returns_visible_answer():
+    text = "<|channel>thought\nscratch work\n<channel|>The answer."
+    assert serve._strip_gemma_channels(text) == "The answer."
+
+
+def test_strip_gemma_channels_cuts_unclosed_channel():
+    assert serve._strip_gemma_channels("Lead <|channel>thought") == "Lead"
+
+
+def test_strip_gemma_channels_removes_stray_close():
+    assert serve._strip_gemma_channels("Visible<channel|>") == "Visible"
+
+
 def test_rejects_multiple_runners(monkeypatch):
     import sys
 

extension/llm/server/cpp/test_worker_loop.cpp

@@ -58,6 +58,7 @@ class FakeSession : public LLMSession {
 
   int prefill_calls = 0;
   std::vector<size_t> prefill_sizes; // size of each prefill_tokens() call
+  std::vector<std::vector<uint64_t>> prefill_batches;
   int fail_prefill_on = -1; // 0-based call index to fail (-1 = never)
   int decode_calls = 0;
   int fail_decode_on = -1;
@@ -68,6 +69,7 @@ class FakeSession : public LLMSession {
       std::vector<uint64_t> tokens,
       const SamplingConfig* /*initial_sampling*/ = nullptr) override {
     prefill_sizes.push_back(tokens.size());
+    prefill_batches.push_back(tokens);
     if (prefill_calls++ == fail_prefill_on) {
       return ETError::Internal; // failed AFTER (notionally) mutating state
     }
@@ -173,13 +175,14 @@ Emitted run(
     WorkerSessionState& st,
     bool warm,
     const nlohmann::json& req,
-    const std::unordered_map<std::string, int64_t>& md = {}) {
+    const std::unordered_map<std::string, int64_t>& md = {},
+    const std::vector<uint64_t>& prefix = {}) {
   static FakeTokenizer tok;
   std::ostringstream cap;
   std::streambuf* old = std::cout.rdbuf(cap.rdbuf());
   Emitted em;
   try {
-    worker_handle_request(st, warm, tok, md, req);
+    worker_handle_request(st, warm, tok, md, req, prefix);
   } catch (const std::exception&) {
     em.threw = true;
   }
@@ -325,6 +328,23 @@ void test_generated_token_ids_excludes_terminal() {
       st.resident_token_ids.size() == (size_t)st.session->position());
 }
 
+void test_prompt_prefix_ids_prepend_text_prompt_once() {
+  auto st = makeState();
+  fake(st).steps = {{0, "", true, true}};
+  auto em =
+      run(st,
+          /*warm=*/true,
+          {{"max_new_tokens", 1}, {"prompt", "ab"}},
+          {},
+          {2});
+  check("prefix: prompt_tokens includes prefix", em.done["prompt_tokens"] == 3);
+  check(
+      "prefix: prefilled ids == [2,'a','b']",
+      fake(st).prefill_batches ==
+          std::vector<std::vector<uint64_t>>{
+              {2, static_cast<uint64_t>('a'), static_cast<uint64_t>('b')}});
+}
+
 void test_stop_string_marks_dirty_and_omits_ids() {
   auto st = makeState();
   fake(st).steps = {
@@ -430,6 +450,7 @@ int main() {
   test_equal_prompt_no_empty_prefill();
   test_anonymous_never_warm();
   test_generated_token_ids_excludes_terminal();
+  test_prompt_prefix_ids_prepend_text_prompt_once();
   test_stop_string_marks_dirty_and_omits_ids();
   test_prefill_failure_marks_dirty();
   test_decode_failure_marks_dirty();