Runtime Memory Checkpointing (Prototype)#
TensorRT-LLM can temporarily release selected GPU allocations while keeping an LLM runtime process alive. This enables a different GPU workload to run while the inference runtime is parked.
Release can preserve allocation contents in host memory and free their GPU physical memory. Resume allocates GPU physical memory again, maps it at the same virtual addresses, and restores the saved contents. The virtual addresses remain reserved while the runtime is parked. This is an in-process memory checkpoint: it is not written to disk and does not survive process termination or support restart in a different process.
Warning
Sleep and wakeup are prototype APIs and are subject to change.
Python API#
Enable runtime memory checkpointing with SleepConfig. Both synchronous
LLM and asynchronous AsyncLLM expose release(), resume(), and
get_memory_status().
from tensorrt_llm import LLM
from tensorrt_llm.llmapi import ExecutorMemoryType, SleepConfig
llm = LLM(
model="Qwen/Qwen3-8B",
sleep_config=SleepConfig(),
)
llm.release()
try:
# The selected GPU physical memory is available to another workload.
...
finally:
llm.resume()
With no tags, release() selects every usable ExecutorMemoryType and
resume() reactivates every tag still parked using its configured restore
mode. Reactivation does not always restore the previous contents: NONE
discards them, while MEMSET maps and zeroes new memory. You can release or
reactivate a specific subset:
tags = [
ExecutorMemoryType.MODEL_ENGINE_MAIN,
ExecutorMemoryType.MODEL_WEIGHTS_MAIN,
]
llm.release(tags)
llm.resume([ExecutorMemoryType.MODEL_WEIGHTS_MAIN])
status = llm.get_memory_status()
assert status.state == "parked"
assert status.parked_tags == [ExecutorMemoryType.MODEL_ENGINE_MAIN]
llm.resume() # Reactivate the remaining parked tag and reopen request admission.
Tags can be enum members or their string values. TensorRT-LLM normalizes them,
removes duplicates while retaining their first-seen order, and rejects an
explicitly empty list. The two _no_capture_init_* enum values are
initialization-only markers and cannot be selected.
The asynchronous API has the same arguments and results:
from tensorrt_llm import AsyncLLM
from tensorrt_llm.llmapi import SleepConfig
async with AsyncLLM(model="Qwen/Qwen3-8B",
sleep_config=SleepConfig()) as llm:
await llm.release()
status = await llm.get_memory_status()
await llm.resume()
release() first closes request admission and drains work already in flight.
Generation remains rejected until all parked tags have been restored. The
synchronous MPI/local and asynchronous Ray paths use the same admission states:
State |
Meaning |
|---|---|
|
Requests are admitted and no tags are parked |
|
Admission is closed and release is in progress |
|
At least one tag remains released |
|
A restore is in progress |
|
A memory operation might have partially mutated runtime state |
Distributed status replies must agree. A disagreement is reported as a runtime failure instead of choosing one worker’s view.
Retries are idempotent. Releasing tags that are already parked or resuming tags that are already active succeeds without repeating a memory operation. Releasing additional tags while already parked is rejected; resume the runtime first. New operations are also rejected during transitional and failed states.
Restore modes#
SleepConfig.restore_modes controls how each tagged allocation is restored:
Mode |
Sleep behavior |
Wakeup behavior |
|---|---|---|
|
Copy to pageable host RAM, then release GPU physical memory |
Copy the host backup to new GPU memory |
|
Copy to pinned host RAM, then release GPU physical memory |
Copy the pinned backup to new GPU memory |
|
Release GPU physical memory without preserving contents |
Map new memory without restoring old contents |
|
Release GPU physical memory without a host copy |
Map and zero new GPU memory |
For example, preserve model weights in pinned host memory and discard the KV cache:
sleep_config = SleepConfig(
restore_modes={
ExecutorMemoryType.MODEL_WEIGHTS_MAIN: "PINNED",
ExecutorMemoryType.KV_CACHE: "NONE",
}
)
The KV-cache restore mode defaults to NONE. Other unlisted tags default to
PINNED when the runtime prefers pinned backup memory, and otherwise to CPU.
Model weights are therefore preserved in host RAM by default. Host memory must
have enough capacity for all allocations using CPU or PINNED.
Enabling sleep_config changes model loading because TensorRT-LLM must create
tagged CUDA virtual-memory pools and route allocations through them. Enable it
only for instances that will use sleep and wakeup.
GPU Memory Service weights#
Weights loaded with load_format="gms" are managed by GPU Memory Service
(GMS), not by the TensorRT-LLM sleep virtual-memory allocator. They have no
model_weights or draft_model_weights sleep tag.
When a sleep request contains a weight tag:
Ordinary virtual-memory-managed weights follow their restore mode, and their GPU physical memory is released.
GMS-managed weights are skipped. They remain in the shared GMS GPU memory pool and are not copied to host RAM by TensorRT-LLM sleep.
Other requested, non-GMS allocations remain eligible for release.
GPU-memory usage might therefore not fall by the full model-weight size with GMS loading. GMS mappings and sessions have a separate lifecycle.
Failure and lifecycle behavior#
Invalid or explicitly empty tag lists raise
ValueError.Release and resume require
sleep_config.Calls synchronize CUDA work before changing mappings and before returning.
Distributed operations coordinate all ranks. An error after ranks begin changing memory can put the worker in a terminal failed state because a partially changed distributed state cannot be safely reconciled.
A host-memory backup is usable only by the same live runtime process.
Put
resume()in afinallyblock if the LLM must remain usable when the caller’s intervening work fails.
Control-plane HTTP endpoints#
trtllm-serve can expose authenticated runtime-control routes for the
PyTorch HTTP text-generation server. The feature is opt-in and requires the
secret to come from the environment:
export TRTLLM_RUNTIME_CONTROL_API_KEY='replace-with-a-long-random-secret'
trtllm-serve Qwen/Qwen3-8B --enable_sleep_mode
If the YAML configuration has no sleep_config, the command creates a default
SleepConfig. An explicitly configured sleep_config is preserved. Sleep
mode is rejected for gRPC.
The enabled routes are:
POST /release_memory
POST /resume_memory
GET /memory_status
The POST body is optional. An omitted body or {} selects the defaults
described in the Python API. To select specific resources, send
{"tags":["model","model_weights"]}.
Every request must include these headers:
x-trtllm-runtime-control-auth: sha256=<hex digest>x-trtllm-runtime-control-timestamp: <Unix seconds>x-trtllm-runtime-control-nonce: <32 lowercase hexadecimal characters>
The digest is HMAC-SHA256 over this byte sequence:
b"trtllm-runtime-control-v1\n"
+ uppercase_method_ascii + b"\n"
+ registered_route_path_utf8 + b"\n"
+ decimal_timestamp_ascii + b"\n"
+ lowercase_hex_nonce_ascii + b"\n"
+ exact_transmitted_body
The route path has one leading slash, no trailing slash, and excludes any deployment root path and query string. Sign an empty byte string for a GET or a body-less POST.
The server accepts timestamps up to five minutes old or five minutes in the future and accepts each nonce only once. It retains accepted nonces through their complete validity period. If the bounded replay cache fills with live entries, authentication fails closed until an entry expires. A retried control request therefore needs a new timestamp, nonce, and signature.
This Python example signs and sends matching bytes:
import hashlib
import hmac
import json
import os
import secrets
import time
import requests
key = os.environ["TRTLLM_RUNTIME_CONTROL_API_KEY"].encode()
def auth(method: str, path: str, body: bytes) -> dict[str, str]:
timestamp = str(int(time.time()))
nonce = secrets.token_hex(16)
payload = (
b"trtllm-runtime-control-v1\n"
+ method.upper().encode("ascii")
+ b"\n"
+ path.encode()
+ b"\n"
+ timestamp.encode("ascii")
+ b"\n"
+ nonce.encode("ascii")
+ b"\n"
+ body
)
digest = hmac.new(key, payload, hashlib.sha256).hexdigest()
return {
"x-trtllm-runtime-control-auth": f"sha256={digest}",
"x-trtllm-runtime-control-timestamp": timestamp,
"x-trtllm-runtime-control-nonce": nonce,
}
body = json.dumps(
{"tags": ["model", "model_weights"]},
separators=(",", ":"),
).encode()
response = requests.post(
"http://localhost:8000/release_memory",
data=body,
headers={
"content-type": "application/json",
**auth("POST", "/release_memory", body),
},
)
response.raise_for_status()
status = requests.get(
"http://localhost:8000/memory_status",
headers=auth("GET", "/memory_status", b""),
)
print(status.json())
resume_body = json.dumps({"tags": ["model_weights"]},
separators=(",", ":")).encode()
requests.post(
"http://localhost:8000/resume_memory",
data=resume_body,
headers={
"content-type": "application/json",
**auth("POST", "/resume_memory", resume_body),
},
).raise_for_status()
A status response has this shape:
{
"state": "parked",
"parked_tags": ["model"]
}
The parked_tags list is sorted lexicographically by tag value.
Successful POST requests return {"status":"success"}. Invalid tags return
HTTP 400, conflicting transitions return 409, authentication failures return
401, and unexpected worker failures return 500.
Treat the control key as an administrative secret and expose these routes only
to a trusted control plane. Coordinate traffic draining before release when
clients can race with the operation, and wait for running before resuming
normal traffic. Host memory must remain available until restore completes.
The older enable_rl_control_endpoints, x-trtllm-rl-control-auth, and
/update_weights surface remains as a deprecated compatibility path for
custom AsyncLLM integrations. It cannot be enabled together with generic
runtime control, and --enable_sleep_mode never exposes /update_weights.
Implementation overview#
At model creation, eligible allocations enter tagged CUDA virtual-memory pools.
Sleep selects allocations by tag, optionally copies contents to CPU or pinned
storage, unmaps them, and calls cuMemRelease on their CUDA physical allocation
handles. Wakeup calls cuMemCreate, remaps new handles at the reserved virtual
addresses, and restores or initializes contents according to SleepConfig.
Principal implementation locations are:
Synchronous API:
tensorrt_llm/llmapi/llm.pyAsynchronous API:
tensorrt_llm/_torch/async_llm.pyServing and authentication:
tensorrt_llm/serve/openai_server.pyandtensorrt_llm/serve/runtime_control_auth.pyWorkers:
tensorrt_llm/executor/base_worker.pyandtensorrt_llm/executor/ray/gpu_worker.pyAllocation tagging:
tensorrt_llm/_torch/pyexecutor/py_executor_creator.pyandtensorrt_llm/_torch/pyexecutor/model_loader.pyPython VMM wrapper:
tensorrt_llm/_torch/virtual_memory.pyCUDA VMM:
cpp/tensorrt_llm/runtime/virtualMemory.cppandcpp/include/tensorrt_llm/runtime/virtualMemory.hGMS adapter:
tensorrt_llm/_torch/memory/gpu_memory_backend.py