Enable prompt-end checkpoints by default for Muse-Glimmer serving

Author: mergennachin

examples/llm_server/python/README.md

@@ -68,6 +68,7 @@ limit; choose a value supported by the export. Capacity defaults are
 bounded by `min(max_sessions, max_inflight_requests)`. A positive value overrides
 this allowance. Alternatively, `--prefix-checkpoint-prompt-end` retains initial,
 continuation, and replacement prompt boundaries separately for each named session.
+This mode is opt-in.
 `--prefix-cache-entries` defaults to eight snapshots per session; duplicates
 are skipped and LRU eviction stays within the owning session. Close/reset clears
 that session's snapshots. Anonymous requests neither use nor retain them.
@@ -99,8 +100,7 @@ event loop. Post-spawn startup failures also await worker cleanup. CLI validatio
 and chat-template preflight still happen before any worker is started.
 
 The batching guide covers the native build, supported export contract, and
-integration checks. The existing generic launcher below and Muse-specific
-launcher are unchanged.
+integration checks. The legacy launcher below is unchanged.
 
 ## Run a legacy worker
 
@@ -257,8 +257,8 @@ worker JSONL protocol is covered separately by `tests/test_worker_client.py`,
 and `tests/test_stream_disconnect.py` uses real loopback Uvicorn/TCP plus a
 model-free subprocess to verify disconnect cancellation end to end.
 The BPE splice tests use an in-memory tokenizer with no model downloads. Optional
-integration tests use local tokenizer directories set with `QWEN_HF_DIR`,
-`GEMMA_HF_DIR`, or `MUSE_GLIMMER_HF_DIR` and require `transformers`.
+integration tests require `transformers` and local tokenizer directories; see the
+environment variables in `tests/test_warm_resume_scaffold.py`.
 
 ## Architecture
 
@@ -313,8 +313,8 @@ the generic control plane never imports an example. Backend specifics
 
 The dedicated batching launcher supports concurrent text requests on one shared
 runner, ordered persistent-session continuation, targeted in-band cancellation,
-and optional prefix reuse for new sessions and clean prompt replacements. It
-does not yet migrate Muse image or DFlash execution. The serialized-worker
+and optional prefix reuse for new sessions and clean prompt replacements. Image
+and speculative execution require a model-specific worker. The serialized-worker
 limitations below describe the legacy path, not the new batching worker.
 
 Deliberately narrow (reliability-first): Hermes/Qwen tool calling only;

examples/llm_server/spec/README.md

@@ -139,7 +139,7 @@ keepalive comments every 15 seconds. Model launchers may bound initial readiness
 waiting to return headers during prefill; errors detected after headers arrive
 as SSE error events.
 
-The generic and Muse batching launchers accept `--api-dialect openai|openrouter`.
+The batching launcher accepts `--api-dialect openai|openrouter`.
 The `openrouter` dialect attaches requested usage to the finish chunk, once, for
 clients that stop consuming on `finish_reason`, and supplies overflow metadata
 for context recovery. The default `openai` dialect retains the separate
@@ -385,8 +385,8 @@ loads only the program before ModuleExecutor construction, reads activation dtyp
 and context length, resolves model/tokenizer EOS, binds the backend's default
 batched cache via `cache::kind::kBatched`, derives scheduler width from
 `preferred_batch_tokens()`, and constructs
-one `ServingRuntime`. It does not preload `forward`, run token-step loops, or use
-Glimmer's legacy session implementation. Prefix caching is disabled by default.
+one `ServingRuntime`. It does not preload `forward` or run token-step loops.
+Prefix caching is disabled by default.
 `--prefix_cache_entries=N` enables retained-prefix reuse for both greedy and
 sampled requests. With `--prefix_capture_slots=C`, it adds `N + C` physical rows
 for retained snapshots and transient captures; logical `--max_sessions` remains
@@ -426,7 +426,7 @@ contains multiple session slices. Overlapping HTTP requests or sequential
 single-session forwards are not proof. Real-model smoke testing is explicitly
 deferred: no compatible local PTE/tokenizer has been confirmed, and this stack
 must not download/export weights or claim that a real combined forward was
-verified. A Glimmer `embed_text`/`forward_from_embeddings` last-logits artifact is
+verified. An `embed_text`/`forward_from_embeddings` last-logits artifact is
 not assumed compatible with this packed construction path.
 
 Idle-session recovery retries at most once and only for native

examples/llm_server/spec/batching.md

@@ -36,7 +36,8 @@ of N and at the initial prompt end. It requires `--prefix-cache-entries > 0`.
 Zero (the default) disables periodic capture. With prefix caching enabled and no
 other capture options, only initial prompts are retained.
 `--prefix-checkpoint-prompt-end` instead retains prompt ends within each named
-session, including continuations and replacements. `--prefix-cache-entries`
+session, including continuations and replacements. This mode is opt-in.
+`--prefix-cache-entries`
 defaults to eight per session. Checkpoints cannot be reused or evicted by another
 session, and close/reset clears them. Anonymous requests do not use this cache.
 This mode cannot be combined with a nonzero interval or custom checkpoint policy.

examples/models/muse-glimmer/BATCHING_SERVING.md

@@ -219,34 +219,41 @@ submitted history; server-side image references are not supported.
 
 The base64/framing check does not guarantee that an arbitrarily long text history
 fits. Whole-frame, decoded-pixel, vision-patch, and model context limits still
-apply. Text/image prefix snapshots are opt-in via `--prefix-cache-entries`
-(default 0).
+apply. Text/image prompt-end checkpoints are enabled by default, with
+`--prefix-cache-entries` defaulting to eight snapshots per named session.
 
 ## Session Behavior
 
-Without prompt-end mode, nonzero `--prefix-cache-entries` enables shared prefix
-caching; without other capture options it retains initial prompts. `--prefix-checkpoint-prompt-end` instead
-retains prompt boundaries separately for each named session, including tool-result
+The Muse-Glimmer batching launcher and native workers enable
+`--prefix-checkpoint-prompt-end` by default. General LLM serving keeps it opt-in.
+This mode retains prompt boundaries separately for each named session, including tool-result
 continuations and replacements. `--prefix-cache-entries` defaults to eight
 checkpoints per session. Repeated or branched requests using the same session ID can
 reuse its snapshots; other sessions and anonymous requests cannot. New tool-result
 text still needs prefill. Exact prompt ends already retained in that session are
 not cloned again. Close/reset, including idle-session eviction, clears its snapshots.
 
+Use `--prefix-checkpoint-prompt-end=false` to disable this mode in the Python launcher,
+or `--prefix_checkpoint_prompt_end=false` in a native worker. With it disabled,
+`--prefix-cache-entries` defaults to zero; a positive value enables shared prefix
+caching, which retains initial prompts unless periodic capture is configured.
+
 Prompt-end-only capture preserves per-request capture concurrency and the default
 transient-row budget. LRU eviction stays within each session, so one session cannot
 evict another's checkpoints. Capture slots are released after clone collection,
 without waiting for decoding to finish. Capture is best effort: concurrent captures
 or session replacements can occupy the available rows. Measure capture overhead
-and reuse for your workload before enabling it.
+and reuse to choose the entry limit for your workload.
 
 `--prefix-checkpoint-prompt-end` cannot be combined with a nonzero
 `--prefix-checkpoint-interval`; both the launcher and native worker reject that combination.
+To use periodic capture, disable prompt-end capture and set a positive
+`--prefix-cache-entries` limit.
 With prompt-end mode, `--prefix-cache-entries` is the per-session budget (default eight);
 without it, the same option is the shared-cache budget (default zero). This mode reserves
 `max_sessions * prefix_cache_entries` snapshot rows plus working sessions
 and capture slots: eight sessions, eight checkpoints each, and eight automatic
-capture slots require 80 physical rows. The feature remains disabled by default.
+capture slots require 80 physical rows.
 Periodic checkpoints use one serialized lane and may skip a boundary while an
 earlier clone is pending. Automatic transient-row sizing is capped at two in
 periodic mode; prompt-end capture retains the ordinary concurrency budget.

examples/models/muse-glimmer/runtime/runners/batching_bootstrap.cpp

@@ -35,7 +35,7 @@ DEFINE_int32(max_inflight_requests, 4, "Maximum admitted requests");
 DEFINE_int32(
     prefix_cache_entries,
     0,
-    "Retained snapshots: shared by default (0), per session with prompt-end checkpoints (8)");
+    "Retained snapshots per session (default: 8); shared with prompt-end checkpoints disabled (default: 0)");
 DEFINE_int32(
     prefix_checkpoint_interval,
     0,
@@ -46,8 +46,8 @@ DEFINE_int32(
     "Transient capture rows; 0: size automatically from session/request limits");
 DEFINE_bool(
     prefix_checkpoint_prompt_end,
-    false,
-    "Retain prompt-end checkpoints separately for each named session");
+    true,
+    "Retain prompt-end checkpoints separately for each named session (enabled by default)");
 DEFINE_uint64(max_image_bytes, 20 * 1024 * 1024, "Maximum encoded image bytes");
 DEFINE_uint64(
     max_input_frame_bytes,

examples/models/muse-glimmer/serving/serve_batching.py

@@ -76,7 +76,7 @@ def _parse_args(argv=None):
     parser.add_argument(
         "--prefix-cache-entries",
         type=int,
-        help="Retained snapshots: shared by default (0), per named session with prompt-end checkpoints (8)",
+        help="Retained snapshots per named session (default: 8); shared with prompt-end checkpoints disabled (default: 0)",
     )
     parser.add_argument(
         "--prefix-capture-slots",
@@ -87,8 +87,11 @@ def _parse_args(argv=None):
     parser.add_argument("--max-vision-patches", type=int, default=4096)
     parser.add_argument(
         "--prefix-checkpoint-prompt-end",
-        action="store_true",
-        help="Retain prompt-end checkpoints separately for each named session",
+        nargs="?",
+        choices=("true", "false"),
+        const="true",
+        default="true",
+        help="Retain prompt-end checkpoints separately for each named session (default: true; false disables)",
     )
     parser.add_argument(
         "--prefix-checkpoint-interval",
@@ -124,6 +127,7 @@ def _parse_args(argv=None):
     parser.add_argument("--host", default="127.0.0.1")
     parser.add_argument("--port", type=int, default=8000)
     args = parser.parse_args(argv)
+    args.prefix_checkpoint_prompt_end = args.prefix_checkpoint_prompt_end == "true"
     if args.prefix_cache_entries is None:
         args.prefix_cache_entries = 8 if args.prefix_checkpoint_prompt_end else 0
     _validate_limits(
@@ -207,7 +211,7 @@ def _worker_command(args):
         ("prefix_checkpoint_interval", args.prefix_checkpoint_interval or None),
         (
             "prefix_checkpoint_prompt_end",
-            "true" if args.prefix_checkpoint_prompt_end else None,
+            str(args.prefix_checkpoint_prompt_end).lower(),
         ),
         ("max_vision_patches", args.max_vision_patches),
         ("max_image_bytes", args.max_image_bytes),

examples/models/muse-glimmer/tests/test_serve_batching.py

@@ -53,47 +53,51 @@ def test_prefix_capture_override_reaches_worker(arguments, captures):
     assert f"--prefix_capture_slots={captures}" in serve_batching._worker_command(args)
 
 
-def test_prompt_end_capture_is_opt_in_and_reaches_worker(arguments, capsys):
-    default = serve_batching._parse_args(arguments + ["--prefix-cache-entries", "4"])
-    assert not default.prefix_checkpoint_prompt_end
-    assert not any(
-        "prefix_checkpoint_prompt_end" in flag
-        for flag in serve_batching._worker_command(default)
-    )
-    enabled = serve_batching._parse_args(arguments + ["--prefix-checkpoint-prompt-end"])
+@pytest.mark.parametrize(
+    "options",
+    [[], ["--prefix-checkpoint-prompt-end"], ["--prefix-checkpoint-prompt-end=true"]],
+)
+def test_prompt_end_capture_defaults_on_and_reaches_worker(arguments, options):
+    enabled = serve_batching._parse_args(arguments + options)
+    assert enabled.prefix_checkpoint_prompt_end
     assert enabled.prefix_cache_entries == 8
     assert "--prefix_checkpoint_prompt_end=true" in serve_batching._worker_command(
         enabled
     )
     assert "--prefix_cache_entries=8" in serve_batching._worker_command(enabled)
+
+
+@pytest.mark.parametrize("capacity", [None, 0, 3])
+def test_prompt_end_capture_opt_out_reaches_worker(arguments, capacity):
+    options = ["--prefix-checkpoint-prompt-end=false"]
+    if capacity is not None:
+        options += ["--prefix-cache-entries", str(capacity)]
+    disabled = serve_batching._parse_args(arguments + options)
+    assert not disabled.prefix_checkpoint_prompt_end
+    assert disabled.prefix_cache_entries == (capacity or 0)
+    command = serve_batching._worker_command(disabled)
+    assert "--prefix_checkpoint_prompt_end=false" in command
+    assert f"--prefix_cache_entries={capacity or 0}" in command
+    assert ("--prefix_capture_slots=0" in command) == bool(capacity)
+
+
+def test_prompt_end_capture_rejects_invalid_combinations(arguments, capsys):
     with pytest.raises(SystemExit):
-        serve_batching._parse_args(
-            arguments
-            + ["--prefix-checkpoint-prompt-end", "--prefix-cache-entries", "0"]
-        )
+        serve_batching._parse_args(arguments + ["--prefix-cache-entries", "0"])
+    assert "requires --prefix-cache-entries > 0" in capsys.readouterr().err
     with pytest.raises(SystemExit):
-        serve_batching._parse_args(
-            arguments
-            + [
-                "--prefix-checkpoint-prompt-end",
-                "--prefix-checkpoint-interval",
-                "64",
-            ]
-        )
+        serve_batching._parse_args(arguments + ["--prefix-checkpoint-interval", "64"])
     assert "cannot be combined" in capsys.readouterr().err
 
 
 def test_session_checkpoint_capacity_override_and_row_overflow(arguments):
-    args = serve_batching._parse_args(
-        arguments + ["--prefix-checkpoint-prompt-end", "--prefix-cache-entries", "2"]
-    )
+    args = serve_batching._parse_args(arguments + ["--prefix-cache-entries", "2"])
     assert "--prefix_cache_entries=2" in serve_batching._worker_command(args)
     for capacity in (0, -1, (1 << 31) - 1):
         with pytest.raises(SystemExit):
             serve_batching._parse_args(
                 arguments
                 + [
-                    "--prefix-checkpoint-prompt-end",
                     "--prefix-cache-entries",
                     str(capacity),
                 ]
@@ -110,6 +114,7 @@ def test_worker_forwards_mg_options(arguments):
             "2048",
             "--prefix-cache-entries",
             "3",
+            "--prefix-checkpoint-prompt-end=false",
             "--prefix-checkpoint-interval",
             "512",
         ]
@@ -127,6 +132,7 @@ def test_worker_forwards_mg_options(arguments):
     assert flags["--bos_id"] == "200000"
     assert flags["--eos_id"] == "200001"
     assert flags["--prefix_checkpoint_interval"] == "512"
+    assert flags["--prefix_checkpoint_prompt_end"] == "false"
 
 
 @pytest.mark.parametrize("options,enabled", [([], "false"), (["--dflash"], "true")])
@@ -160,6 +166,7 @@ def test_draft_count_forwarded(arguments, n_draft):
         ("--max-vision-patches", str(1 << 31)),
         ("--prefix-checkpoint-interval", "8"),
         ("--prefix-checkpoint-interval", "-1"),
+        ("--prefix-checkpoint-prompt-end", "invalid"),
         ("--prefix-capture-slots", "-1"),
         ("--prefix-capture-slots", str(1 << 31)),
     ],