# Thread Model for Handling Asynchronous Inference Operations in LiteRT-LM > Explore the LiteRT-LM thread model for asynchronous inference. Discover how its task-centric pipeline with dedicated ThreadPools ensures non-blocking LLM operations and efficient callback dispatch. - Repository: [google-ai-edge/LiteRT-LM](https://github.com/google-ai-edge/LiteRT-LM) - Tags: internals - Published: 2026-04-06 --- **LiteRT-LM employs a task-centric pipeline using two single-threaded ThreadPools—one dedicated to executing LLM prefill and decode operations, and another for dispatching user callbacks—to ensure non-blocking asynchronous inference.** The `google-ai-edge/LiteRT-LM` runtime separates inference computation from user-facing callbacks through a deterministic thread architecture. This article examines the thread model for handling asynchronous inference operations in LiteRT-LM, detailing how the **ExecutionManager** coordinates work between single-threaded pools to maintain responsiveness and thread safety. ## Core Architecture: Dual Single-Threaded ThreadPools LiteRT-LM implements a **task-centric thread model** built on two specialized `ThreadPool` instances, both configured with exactly one worker thread. This design creates a clear separation between computational work and notification delivery. ### Execution Thread Pool The **execution thread pool** handles the heavy lifting of LLM operations, including prefill, decode, text-scoring, and session cloning. Defined in [`runtime/framework/resource_management/execution_manager.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/framework/resource_management/execution_manager.h), this pool runs a single **WorkerThread** that pulls tasks from the `ExecutionManager` queue. By limiting the pool size to one thread, LiteRT-LM ensures deterministic ordering of inference operations without contention. ### Callback Thread Pool The **callback thread pool** invokes user-provided `absl::AnyInvocable` callbacks once tasks complete. Also single-threaded and configured in the `ExecutionManager` constructor, this pool prevents the execution worker from blocking on user code. The separation ensures that long-running client callbacks do not stall the inference pipeline. ### WorkerThread Implementation Underlying each `ThreadPool` is a concrete `std::thread` running the `RunWorker()` loop. Defined in [`runtime/framework/worker_thread.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/framework/worker_thread.h), each **WorkerThread** manages its own execution lifecycle, incrementing `num_active_tasks_` when active and decrementing upon completion. The generic `ThreadPool` class in [`runtime/framework/threadpool.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/framework/threadpool.h) stores a queue of `absl::AnyInvocable` callbacks protected by an `absl::Mutex`. ## Lifecycle of an Asynchronous Inference Request The flow from API call to callback execution follows a strict **single-producer → single-consumer → single-consumer** pipeline: ### 1. Session API Entry Points Public asynchronous methods reside in `SessionAdvanced`. Functions like `RunPrefillAsync`, `RunDecodeAsync`, and `RunTextScoringAsync` allocate a unique **TaskId**, construct a **TaskInfo** struct containing the user callback and work lambda, and forward the request to the `ExecutionManager`. *Source:* [`runtime/core/session_advanced.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/core/session_advanced.h) (lines 22-33 for task controller, lines 123-138 for async APIs). ### 2. Task Registration in ExecutionManager The `ExecutionManager` registers the `TaskInfo` in a thread-safe `task_lookup_` map protected by `session_and_task_lookup_mutex_`. Methods such as `AddPrefillTask` and `AddDecodeTask` create the task context before queuing it for execution. *Source:* [`execution_manager.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/execution_manager.h) (lines 96-122). ### 3. Scheduling to the Execution Thread Pool The `QueueTask` method pushes the task's executable lambda onto `execution_thread_pool_->tasks_` and wakes the worker thread. The `ThreadPool::Schedule` method handles synchronization using the pool's mutex. *Source:* [`threadpool.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/threadpool.h) (lines 64-69). ### 4. WorkerThread Execution Loop Each `WorkerThread` runs inside `RunWorker()`, pulling callbacks from the queue under the pool mutex. The thread increments `num_active_tasks_`, executes the callback (which runs the actual LLM executor like `LlmExecutor::Prefill` or `Decode`), then decrements the counter upon completion. *Source:* [`worker_thread.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/worker_thread.h) (lines 31-44). ### 5. Task Completion and Callback Dispatch Upon finishing the LLM operation, the system calls `FinishTask`. This method moves the task's callback from `TaskInfo` to the **callback thread pool** via `callback_thread_pool_->Schedule`. This handoff is critical: it guarantees that user code runs on a separate thread from the execution worker, preventing deadlocks and keeping the execution pipeline responsive. *Source:* [`execution_manager.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/execution_manager.h) (lines 223-237). ### 6. Client Callback Invocation The user-provided lambda receives the `absl::StatusOr` result on the callback thread. Because this pool is also single-threaded, callbacks execute in completion order, though not necessarily in the original scheduling order. ## Thread Safety and Synchronization Mechanisms ### Mutex Protection All mutable shared structures—including `task_lookup_`, `session_lookup_`, `tasks_`, and `num_active_tasks_`—are guarded by `absl::Mutex` with `ABSL_GUARDED_BY` annotations. `SessionAdvanced` protects its internal `session_state_` and `last_task_ids_` with a separate mutex, preventing cross-component deadlocks. ### Non-Blocking Callback Design The **execution thread** never blocks on user callbacks. By offloading completion notifications to the dedicated callback pool, the system maintains throughput even when client code performs I/O or heavy processing. ### Atomic Cancellation Cancellation is handled via a `std::shared_ptr>` stored in `TaskInfo`. The **AdvancedTaskController** can set this flag atomically, and execution workers check the `cancelled` status before proceeding with heavy computation. This allows safe termination of long-running inference tasks without corrupting shared state. ## Practical Code Examples ### Submitting an Async Prefill Request The following example demonstrates how to initiate an asynchronous prefill operation and handle the result: ```cpp // Assume we already have a SessionAdvanced* `session` and a vector of InputData. std::vector inputs = {/* ... */}; session->RunPrefillAsync( inputs, [](absl::StatusOr result) { if (!result.ok()) { LOG(ERROR) << "Prefill failed: " << result.status(); return; } // Handle the successful prefill response. LOG(INFO) << "Prefill completed, token count: " << result->token_ids_size(); }); ``` This call creates a `TaskInfo`, queues it on the execution thread pool, and schedules the callback to run on the callback thread pool. *Source:* `RunPrefillAsync` declaration (lines 123-126) in [`session_advanced.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/session_advanced.h). ### Running Async Decode with Custom Configuration To run generation with specific parameters: ```cpp litert::lm::DecodeConfig decode_cfg; decode_cfg.max_output_tokens = 64; session->RunDecodeAsync( decode_cfg, [](absl::StatusOr result) { if (!result.ok()) { LOG(ERROR) << "Decode error: " << result.status(); return; } for (const auto& token : result->token_ids()) { std::cout << token << ' '; } std::cout << std::endl; }); ``` Internally, this invokes `ExecutionManager::AddDecodeTask`, which schedules the work on the execution pool and forwards the callback to the callback pool. *Source:* `RunDecodeAsync` overloads (lines 133-140) in [`session_advanced.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/session_advanced.h). ### Cancelling an Ongoing Task Store the task controller returned by the async call to cancel execution: ```cpp // Store the controller returned by the async call. std::unique_ptr ctrl; session->RunPrefillAsync( inputs, [&](absl::StatusOr r) { // This will not be called if we cancel before completion. }, &ctrl); // Later... if (ctrl) { ctrl->Cancel(); // Sets the atomic cancellation flag. } ``` The `AdvancedTaskController` holds a shared `std::atomic` that the execution worker checks before proceeding. *Source:* `AdvancedTaskController` definition (lines 46-71) in [`session_advanced.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/session_advanced.h). ## Summary - **Dual-pool architecture**: LiteRT-LM uses two single-threaded `ThreadPool` instances—one for execution and one for callbacks—to isolate inference work from user code. - **Deterministic ordering**: The single-consumer design ensures tasks execute sequentially on the execution thread, preventing race conditions in the LLM engine. - **Safe callback dispatch**: By moving callbacks to a separate thread via `FinishTask`, the system prevents user code from blocking the inference pipeline. - **Comprehensive synchronization**: All shared state uses `absl::Mutex` with thread-safety annotations, and cancellation leverages atomic flags for safe task termination. - **Clear source locations**: Key logic resides in [`execution_manager.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/execution_manager.h) (orchestration), [`threadpool.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/threadpool.h) (queue management), [`worker_thread.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/worker_thread.h) (thread lifecycle), and [`session_advanced.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/session_advanced.h) (public API). ## Frequently Asked Questions ### How does LiteRT-LM prevent deadlocks between inference execution and user callbacks? LiteRT-LM prevents deadlocks by strictly separating execution and notification concerns into two distinct **single-threaded ThreadPools**. When a task completes in the execution pool, the `FinishTask` method moves the user callback to the callback thread pool rather than invoking it directly. This ensures that slow or blocking user code never stalls the execution thread, maintaining pipeline throughput and preventing circular wait conditions. ### Why does LiteRT-LM use single-threaded pools instead of multi-threaded thread pools? The **single-threaded design** (one worker per pool) provides deterministic ordering and eliminates contention for LLM inference tasks. According to the source in [`execution_manager.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/execution_manager.h), both pools are configured with `max_num_threads = 1`. This design suits the sequential nature of autoregressive language models where decode steps depend on previous outputs, and it simplifies thread-safety reasoning by removing the need for complex load balancing or work-stealing algorithms. ### What happens if I call Cancel() on a TaskController after the task has already started? The cancellation mechanism uses a `std::shared_ptr>` stored in the `TaskInfo` struct. When `AdvancedTaskController::Cancel()` is called, it sets this atomic flag. The execution worker checks this flag in the `RunWorker` loop before performing heavy computation. If the task is already mid-execution, the worker detects the cancellation at the next checkpoint and terminates gracefully, ensuring that `FinishTask` is not called for cancelled work, or that the callback receives a cancellation status. ### Where is the ThreadPool implementation located in the LiteRT-LM repository? The **ThreadPool** class is defined in [`runtime/framework/threadpool.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/framework/threadpool.h), while the concrete **WorkerThread** logic resides in [`runtime/framework/worker_thread.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/framework/worker_thread.h). The **ExecutionManager**, which instantiates and coordinates both pools, is located at [`runtime/framework/resource_management/execution_manager.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/framework/resource_management/execution_manager.h). The public async APIs that interact with these components are declared in [`runtime/core/session_advanced.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/core/session_advanced.h).