Execution-State Capsules: 用于低延迟设备端AI服务的图绑执行状态检查点与恢复

Execution-State Capsules: Graph-Bound Execution-State Checkpoint and Restore for Low-Latency, Small-Batch, On-Device Physical-AI Serving

精选理由

FlashRT用执行状态胶囊实现了亚毫秒级恢复,比传统KV缓存多保存循环状态,对交互式AI和机器人很关键。

AI 摘要

论文提出execution-state capsules,一种图绑定的检查点与恢复机制,能保存KV缓存、循环状态、卷积状态等完整执行状态。基于此实现的FlashRT运行时在RTX 5090上,恢复操作亚毫秒级,TTFT相比冷预填充在2k tokens时加速3.9倍,16k tokens时加速27倍。在Jetson AGX Thor和DGX Spark上保持相同正确性和结构特性。该方法不取代高吞吐KV缓存服务,而是为显式执行状态复用提供互补的低延迟方案。

原文 · arXiv cs.LG

Execution-State Capsules: Graph-Bound Execution-State Checkpoint and Restore for Low-Latency, Small-Batch, On-Device Physical-AI Serving

Mainstream LLM serving systems reuse prefix work mainly through paged or radix key-value (KV) caches. This is highly effective for high-throughput, high-concurrency serving, but it manages only one positional fragment of execution state: the KV cache. We study the opposite regime: low-latency, small-batch, on-device physical-AI serving, where interactive LLM agents, speech systems, and robot policies repeatedly branch, reset, interrupt, and re-enter under tight responsiveness budgets. We introduce execution-state capsules, a graph-bound checkpoint and restore mechanism for the complete restorable state at a committed boundary. FlashRT is a white-box, backend-facing kernel runtime whose evaluated NVIDIA CUDA backend runs captured graph plans over contiguous static buffers with no block-table indirection. Because the live state is a closed set of named buffers, a capsule can snapshot, restore, fork, or roll back the whole execution boundary, including KV, recurrent state, convolution state, MTP state, and metadata. This moves reuse from token-addressed KV fragments to graph-bound execution-state boundaries. On an RTX 5090, capsule restore is byte-exact at the stored-state level and token-identical under greedy decode. A KV-only ablation diverges, showing that recurrent state is load-bearing. GPU-resident snapshot and restore are sub-millisecond, and TTFT speedup over cold prefill grows from 3.9x at 2k tokens to 27x at 16k tokens. On Jetson AGX Thor and DGX Spark, the same correctness and structural properties hold. Capsules are not a replacement for high-throughput KV-cache serving; they define a complementary latency-first serving point for explicit execution-state reuse.