[mergennachin]

Summary

A chat interface or coding agent can now explicitly save a serving session before closing and resume it after restarting the worker. The application keeps its conversation messages and tool results; serving persists the model state needed to continue without repeating the saved prefix's prefill. Model loading and checkpoint transfer still have a cost.

The application-facing operations are implemented on MultiplexedWorkerClient and SessionRuntime:

await client.save_session("chat-42")
# Close the application/worker. Later, start a matching worker with its store:
session_id = await client.resume_session("chat-42")

save_session captures the current committed state and atomically replaces that session's saved file. The native runtime chooses the boundary; application code does not calculate token counts or manage target/draft cache details. resume_session restores the file into an unused live session with the same ID and returns that ID. Send the complete next prompt through that session. It must extend exactly the saved token IDs, including chat formatting and committed generated tokens, for serving to prefill only the new suffix. Re-rendered histories can change those IDs and cause a cold prefill; the advanced API supports an explicit earlier match. Saving and resuming are explicit; close does not autosave or remove saved files.

One-time storage configuration

At application startup, configure storage using the artifacts and settings of the worker being launched. SessionStore.create hashes the files off the event loop and binds the store to that configuration. Include the PTE, tokenizer, external weights/adapters, and activation-affecting settings such as backend, RoPE, and draft configuration. Recreate this setup in a later process; per-session calls only need the session ID.

from executorch.examples.llm_server.python.session_store import SessionStore
from executorch.examples.llm_server.python.multiplexed_worker_client import (
    spawn_multiplexed_worker,
)

store = await SessionStore.create(
    "./saved-sessions",
    model_files=[model_path, tokenizer_path, *external_weight_paths],
    configuration=worker_configuration,
)
client = await spawn_multiplexed_worker(worker_command, session_store=store)

Save, close, and resume

After the conversation turn finishes:

from types import SimpleNamespace

await client.save_session("chat-42")
await client.close()

# In the later process, recreate the store as above and start a new worker.
client = await spawn_multiplexed_worker(worker_command, session_store=store)
session_id = await client.resume_session("chat-42")
stream = client.generate(
    full_next_prompt,
    SimpleNamespace(session_id=session_id, max_new_tokens=128, temperature=0),
)
async for text in stream:
    print(text, end="")
await stream.wait()

For applications using the serving runtime, runtime = SessionRuntime(client) exposes the same save_session(id) and resume_session(id) calls. Within its request-capacity limit, it holds its per-session lease through native and file I/O, so save waits for earlier work on that session. Direct worker-client users must finish generation before saving.

A missing or corrupt file, incompatible model/configuration, unusable checkpoint, or already-active destination session raises an error. The application chooses whether to rebuild from its independently saved conversation history. No fallback is selected implicitly. Queued saves and resumes can be cancelled without starting. After a transfer starts, worker-client cancellation waits for settlement. SessionRuntime cancellation returns immediately while the transfer settles under the session lease; later admitted same-session operations wait behind it. Neither API undoes a completed saved file or restored session. This is a Python serving API; HTTP persistence endpoints and automatic UI shutdown hooks are outside this change.

Saving requires a nonempty, idle, healthy, token-only history. Cancellation, failed generation, stop-string cutoff, reset/empty sessions, and opaque/image histories can make the state unsaveable. A failed save raises and preserves the previous file; a later successful token-only turn can be saved again. Both API levels share one client-side checkpoint queue.

SessionStore.create(..., max_bytes=...) sets a symmetric save/load payload limit, defaulting to 16 GiB. Oversized saves fail before replacing the last loadable file. Store directories on POSIX must be owned by the current user and not writable by group or others; existing owned mode-0755 directories work. File checksums detect corruption, not authenticity; load only trusted files. Nonregular files and POSIX symlinks are rejected.

Saved-state deletion is also explicit. First await outstanding worker-client save/resume tasks; client.close() alone does not wait for Python file writes. With SessionRuntime, await runtime.aclose_worker() also settles its background operations. Then:

import asyncio

removed = await asyncio.to_thread(store.delete, "chat-42")

This returns whether a saved file existed and does not change a live session. Startup does not evict files or sweep temporary files. A killed writer can leave .<session-hash>.etcp.* files; remove those only after all writers have stopped.

Advanced API and state contract

The lower-level APIs remain available when an application wants to choose a capture position, retain CPU snapshots in memory, use custom storage, or restore only part of a matching prefix. They use the same checkpoint format and native validation as save_session / resume_session.

Export a chosen committed prefix and save it. Assume the idle seed session still retains the state needed at token position 4,000. The advanced caller explicitly supplies the compatibility identity and capture position:

import asyncio
from executorch.examples.llm_server.python.checkpoint import (
    checkpoint_identity,
    load_checkpoint,
    save_checkpoint,
)

identity = await asyncio.to_thread(
    checkpoint_identity,
    [model_path, tokenizer_path, *external_weight_paths],
    worker_configuration,
)
checkpoint = await client.export_checkpoint_to_cpu(
    "seed", position=4000, identity=identity
)
await asyncio.to_thread(save_checkpoint, checkpoint, "shared-prefix.etcp")

Export returns an independent CPUCheckpoint; it does not create a durable file or close seed. The file helpers are optional if the application keeps that object in memory or supplies its own persistence.

Load and restore into a fresh session. After starting a compatible worker, load CPU data and explicitly choose how much to restore:

checkpoint = await asyncio.to_thread(
    load_checkpoint, "shared-prefix.etcp", expected_identity=identity
)
await client.restore_checkpoint_from_cpu(
    checkpoint,
    "restored",  # Must be an unused live session ID.
    position=4000,
    expected_identity=identity,
)

Loading validates the file without allocating device state. Restore copies state into the worker. Generate with session_id="restored" and the complete next prompt, including the matching 4,000-token prefix; serving prefills its new suffix.

Restore a partial match or replay an exact prompt. Use the exact token IDs the worker will consume, including BOS and chat formatting. Instead of the full-position restore above:

matched = checkpoint.common_prefix_length(next_prompt_token_ids)
# Leave a token to process because the checkpoint does not contain logits.
position = min(matched, max(0, len(next_prompt_token_ids) - 1))
# "fork" must be an unused live session ID in either branch.
if position == 0:
    await client.open_session("fork")  # Explicitly choose a fresh prefill.
else:
    checkpoint.validate_restore_position(position)
    await client.restore_checkpoint_from_cpu(
        checkpoint,
        "fork",
        position=position,
        expected_identity=identity,
    )

Then generate with session_id="fork" and the complete next prompt. A matching 4,000-token prefix followed by new input restores all 4,000 tokens. Retrying exactly those 4,000 tokens restores 3,999 and prefills one token to obtain logits. The zero-position branch explicitly chooses fresh prefill. A position whose required history was overwritten still raises; the application chooses another checkpoint or a fresh prefill.

C++ serving API. Native applications can export a named session into owned CPU memory and restore it into a compatible ServingRuntime. The following runs inside an application function, on an external thread. source owns the idle chat-42 session; destination is a configured runtime using the same model and execution settings, with that key unused. The promises adapt the existing callback API; callbacks only deliver results, and all waiting happens on the application thread.

#include <executorch/extension/llm/serving/serving_runtime.h>
#include <future>
#include <iostream>
#include <memory>
#include <utility>

namespace batching = executorch::extension::llm::batching;
namespace serving = executorch::extension::llm::serving;

// Inside an application function with source and destination runtimes:
using ExportResult = serving::SessionCheckpointResult;
auto exported = std::make_shared<std::promise<ExportResult>>();
auto export_done = exported->get_future();
source.export_checkpoint_to_cpu_async(
    "chat-42", std::nullopt,  // Current committed boundary; or an explicit position.
    [exported](ExportResult result) {
      exported->set_value(std::move(result));
    });

auto result = export_done.get();
if (auto* error = std::get_if<serving::ServingError>(&result)) {
  std::cerr << error->message << '\n';
  return;
}
auto checkpoint =
    std::get<std::shared_ptr<serving::SessionCheckpoint>>(std::move(result));

auto restored = std::make_shared<std::promise<serving::LifecycleResult>>();
auto restore_done = restored->get_future();
destination.restore_checkpoint_from_cpu_async(
    "chat-42", checkpoint, checkpoint->state.position,
    [restored](serving::LifecycleResult result) {
      restored->set_value(std::move(result));
    });
if (auto error = restore_done.get()) {
  std::cerr << error->message << '\n';
  return;
}

The exported SessionCheckpoint includes token history and independent CPU state, so it can outlive source. After restore, call destination.generate with chat-42 and the complete next prompt extending that history. An earlier retained position can be passed to restore for a partial match; the same token-match, retained-history, and one-token exact-retry requirements above apply. Do not mutate a checkpoint while a restore is reading it, or block on these operations from runtime callbacks.

C++ runner API. Applications using batching::Runner directly already have future-returning primitives. For example, this helper exports a chosen committed boundary and returns a newly restored session or its error:

#include <executorch/extension/llm/batching/runner.h>

batching::CheckpointResult<batching::Session> copy_committed_prefix(
    const batching::Session& source,
    batching::Runner& destination,
    batching::Position position) {
  auto result = source.export_checkpoint_to_cpu_async(position).get();
  if (auto* error = std::get_if<batching::CheckpointError>(&result)) {
    return *error;
  }
  auto checkpoint = std::make_shared<batching::Checkpoint>(
      std::get<batching::Checkpoint>(std::move(result)));
  return destination.restore_checkpoint_from_cpu_async(checkpoint, position).get();
}

Call this helper from an external thread and retain the returned Session through subsequent generation. The runner-level checkpoint contains model state but no prompt token history; the application owns that history and supplies the continuation suffix. The serving-level SessionCheckpoint adds the history used for full-prompt matching.

These C++ APIs perform CPU transfer and native restoration. This PR's durable file helpers and save_session / resume_session convenience API are Python APIs; it does not add a C++ file serializer or automatic disk persistence. A C++ application providing its own storage must preserve all checkpoint fields and enforce model/artifact compatibility itself. Native validation checks executor format, geometry, dtype, and retained positions; it does not fingerprint model weights.

A token match is reusable only if every layer retains the required history. Ordinary sliding-window attention needs window - 1 prior rows; DFlash needs window.

Checkpoints carry token history, target KV, retained layer ranges, draft KV, and pending committed target features needed by DFlash. Python checks bounds, integrity, and the configured model identity. Native restore checks backend format, geometry, dtype, and position before publishing a new session. MLX and CUDA implement the same transfer interface; their executor formats are not interchangeable.

Token IDs and cache buffers use binary payload, so long histories do not enlarge JSON control frames. Files use ETCP0002 and transport manifest version 2; files from earlier drafts must be regenerated. Transfers use private temporary files and host copies. Device synchronization and native transfer pause admission and prompt preparation across sessions.

examples/llm_server/CHECKPOINTS.md documents both workflows. Review the cache/executor checkpoint types and runner lifecycle first, then the MLX/CUDA implementations, worker transport, and Python session storage.

This draft was authored with AI assistance using OpenAI Codex. End-to-end CUDA save/resume validation remains outstanding.

Test plan

Validated on an Apple M2 Ultra with the real off-graph Q4_K_M Muse-Glimmer worker on MLX. The new examples/models/muse-glimmer/tests/test_checkpoint_live.py passed for 512- and 4,096-token code prompts with DFlash enabled and disabled: four cases total, rerun with the worker receiving a tokenizer directory. Each calls save_session after generation, restarts the worker, calls resume_session by ID, and re-exports it to verify byte-for-byte equality of token history, target KV, draft KV, and pending target features. Full, earliest-valid partial, and exact-retry continuations match cold generated token IDs. Restored requests prefill only the suffix; exact retries prefill one token. The tests bypass Python preflight to check native rejection one token before the valid retention boundary and verify that rejection leaves the destination key reusable.

An independent live reproduction confirmed the original DFlash boundary bug: restore reported success and the next forward failed an unrelated co-batched request. With the fix, the native worker rejects the restore while the unrelated request completes all 400 tokens. A valid boundary restore co-batched with another request also matches cold generation.

The live command used --backend=mlx --max_sessions=8 --max_session_tokens=8192 --max_decode_sequences=8 --max_inflight_requests=8 --prefix_cache_entries=0, with either --dflash=true --dflash_n_draft=7 or --dflash=false --dflash_n_draft=0. The test accepts the complete worker argv through MUSE_GLIMMER_CHECKPOINT_COMMAND and tokenizer directory through MUSE_GLIMMER_HF_DIR.

python -m pytest -q examples/models/muse-glimmer/tests/test_checkpoint_live.py

# EXECUTORCH_TEST_MULTIPLEXED_WORKER and EXECUTORCH_BATCHING_TEST_WORKER
# point to the built native test workers.
python -m pytest -q \
  examples/llm_server/python/tests/test_session_store.py \
  examples/llm_server/python/tests/test_session_runtime.py \
  examples/llm_server/python/tests/test_checkpoint.py \
  examples/llm_server/python/tests/test_checkpoint_native.py \
  examples/llm_server/python/tests/test_multiplexed_worker_client.py \
  examples/llm_server/python/tests/test_multiplexed_worker_client_native.py

The Python/native-worker tests plus the live fixture's four artifact-identity cases passed 393 tests, including save/close/restart/resume without manual positions, model/configuration mismatch, atomic file replacement, cancellation retaining the session lease, mixed advanced/high-level checkpoint queuing, size-limit preservation, explicit deletion, and nonregular file rejection. Native-worker tests cover failed saves after reset, stop-string cutoff, and cancellation, then a successful save after a clean turn. Native tests passed: cache core (62), batching (231), session table (10), executor utilities (11), serving runtime (278), multiplexed transport (46), and real MLX checkpoint layers (4). These total 642 native tests. All ten CTest targets passed, including the worker prefill/loop smoke targets. Muse-Glimmer batching preparation/materialization tests also passed. Changed native components were rebuilt in an isolated CMake harness against existing Release dependencies. The new subprocess checkpoint tests are collected by the existing native-worker CI job. Regression tests also cover a 300,000-token binary round-trip with a sub-1KB manifest, wrapped ring offsets, allocation failure during checkpoint admission, shutdown during restore, abandoned restore futures, serving error codes and empty callbacks, malformed two-cache metadata, file descriptor cleanup, deeply nested manifests, and same-session requests during restore I/O. The checkpoint Runner test file also passed a -fno-exceptions syntax check. Both C++ usage examples above passed a syntax check against the new headers. The mixed Python queue test fails in six affected cases when run against the original unlocked advanced methods.

The four MLX checkpoint tests also compiled, linked, and passed with AddressSanitizer and UndefinedBehaviorSanitizer enabled in an isolated harness using the new CMake sanitizer block. The MLX delegate sources in that harness were instrumented too.

CUDA partial restore now clears the destination ring column before uploading retained rows: masked unused rows must contain finite values because multiplying zero by a stale NaN still produces NaN. A new cache regression fills the destination with NaN bit patterns, checks zeroed unused rows and preserved retained rows, and checks that other columns are unchanged. New CUDA DFlash executor tests cover full/partial restore, re-exported state versus clone, continuation, malformed state, draft-import rollback, and poisoned-executor rejection. These new CUDA tests are uncompiled and unrun locally on this Mac and await CI; the earlier CI results below predate them.

Scoped lintrunner, including cppcheck, and git diff --check passed. ROCm CI exposed a missing cudaMemset2DAsync compatibility wrapper; the revision adds the direct hipMemset2DAsync mapping, with compilation awaiting the next ROCm CI run. The CUDA runtime CI job passed all 64 cache tests on commit 8a483b03, including both checkpoint round-trips and the DFlash retention-boundary test. The CUDA Muse-Glimmer batching job also passed on that commit: it compiled cuda_dflash_executor.cpp, passed 40 Python tests and 28 CUDA DFlash tests, and passed the native batching test target. Complete model restart/save/resume has been validated on MLX; end-to-end CUDA restore remains untested. This validation exercised the native worker API; it did not exercise Muse CLI or HTTP, and the live runs are correctness checks rather than controlled performance benchmarks.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24984

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 Unclassified Failures

As of commit 59d9416 with merge base a73396e ():

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Test CUDA Builds / test-muse-glimmer-cuda-e2e-kquant-17gb-dflash-image-batching / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]received SIGINT
• Test ROCm AOTI / test-voxtral-realtime-rocm-gfx950-rocm7.2 / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.