# How to Use Streaming Inference with LiteRT-LM's Callback-Based API

> Learn how LiteRT-LM's callback-based API enables non-blocking streaming inference for real-time LLM output processing.

- Repository: [google-ai-edge/LiteRT-LM](https://github.com/google-ai-edge/LiteRT-LM)
- Tags: how-to-guide
- Published: 2026-04-06

---

**LiteRT-LM implements non-blocking streaming inference by accepting a user-supplied callback that receives incremental `Responses` objects via `absl::StatusOr<Responses>`, allowing real-time processing of partial LLM outputs without blocking the caller thread.**

The `google-ai-edge/LiteRT-LM` repository provides a high-performance inference stack for large language models that supports streaming inference through a callback-driven architecture. Unlike blocking APIs that return complete results, the streaming interface in [`runtime/engine/engine.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/engine/engine.h) enables applications to process generated tokens as they arrive, reducing latency for interactive use cases.

## Architecture and Component Flow

Streaming inference in LiteRT-LM follows a layered architecture that separates orchestration from execution. The system schedules work on a background **ThreadPool** while delivering results through the user callback.

| Component | Role in Streaming | Key Source File |
|-----------|-------------------|-----------------|
| **Engine** | Factory for sessions; declares the streaming API. | [`runtime/engine/engine.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/engine/engine.h) |
| **SessionBasic** | Orchestrates pre-fill and decode streaming; manages asynchronous execution. | [`runtime/core/session_basic.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/core/session_basic.h) |
| **Pipeline** | Forwards the user callback to low-level tasks during `DecodeStreaming`. | [`runtime/core/pipeline.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/core/pipeline.h) |
| **Tasks** | Executes decode steps and invokes the callback with intermediate `Responses`. | `runtime/core/tasks.cc` |
| **ThreadPool** | Runs prefill and decode operations on background threads. | [`runtime/framework/threadpool.h`](https://github.com/google-ai-edge/LiteRT-LM/blob/main/runtime/framework/threadpool.h) |

The call graph begins at `Engine::Session::GenerateContentStream`, which delegates to `SessionBasic::GenerateContentStream` in `runtime/core/session_basic.cc`. This method first schedules `RunPrefillAsync` on the thread pool, then chains into `DecodeInternalStreaming` once prefill completes.

## Execution Flow

The streaming pipeline executes in distinct phases to maintain non-blocking behavior:

1. **User invokes the streaming API** through `session->GenerateContentStream(contents, user_callback)`.

2. **SessionBasic initiates prefill** by calling `RunPrefillAsync`, which schedules a lambda on the internal `ThreadPool` to execute `PrefillInternal`.

3. **Prefill completion triggers decoding** via the prefill callback, which captures the original user callback and invokes `DecodeInternalStreaming(user_callback, decode_cfg)`.

4. **Pipeline forwards to Tasks** where `pipeline::DecodeStreaming` validates the callback and calls `Tasks::Decode` in `runtime/core/tasks.cc`.

5. **Incremental delivery** occurs as `Tasks::Decode` runs the model's decode signature step-by-step, invoking `user_cb(Responses(state=kProcessing,…))` after each token generation.

6. **Terminal state** signals completion through a final callback containing `Responses(state=kDone,…)` or an error status if cancellation or model failure occurs.

All callbacks conform to the signature `absl::AnyInvocable<void(absl::StatusOr<Responses>)>`, defined in the engine headers.

## C++ Implementation Example

Below is a complete example demonstrating callback registration and response handling:

```cpp
#include "runtime/engine/engine.h"
#include "runtime/core/session_basic.h"
#include "absl/log/log.h"

using namespace litert::lm;

int main() {
  // Initialize engine with model assets
  auto engine = Engine::CreateEngine(EngineSettings::CreateDefault(
      ModelAssets::Create("my_model_dir").value(),
      Backend::CPU)).value();

  // Create session
  std::unique_ptr<Engine::Session> session =
      engine->CreateSession(SessionConfig::CreateDefault()).value();

  // Prepare input prompt
  std::vector<InputData> prompt = { InputText("Explain quantum entanglement.") };

  // Start streaming inference with callback
  session->GenerateContentStream(
      prompt,
      [](absl::StatusOr<Responses> resp_status) {
        if (!resp_status.ok()) {
          LOG(ERROR) << "Streaming error: " << resp_status.status();
          return;
        }

        const Responses& resp = *resp_status;
        LOG(INFO) << "Partial output: " << resp.text();

        // Check for completion
        if (resp.state() == TaskState::kDone) {
          LOG(INFO) << "Generation finished.";
        }
      });

  // Optional: Block until final callback delivered
  session->WaitUntilDone();
}

```

The callback receives both incremental and final results. Filter by `resp.state()` to distinguish between `kProcessing`, `kDone`, `kCancelled`, or `kError` states.

## Callback Signatures and Threading

The **callback-based API** requires an `absl::AnyInvocable<void(absl::StatusOr<Responses>)>` callable that LiteRT-LM invokes on the internal thread pool. According to the source in `runtime/core/tasks.cc`, the callback executes on the worker thread performing the decode step unless a separate callback thread pool is configured.

Key response accessors include:
- `resp.text()` for the generated string content
- `resp.state()` for the task progress enum
- `resp.token_ids()` for raw token identifiers

Error handling occurs through the `absl::StatusOr` wrapper. Check `!resp_status.ok()` to detect model failures, cancellation, or runtime errors before dereferencing the response object.

## Cancellation and Lifecycle Management

The session maintains an internal `std::atomic<bool> cancelled_` flag checked during each decode step. To abort streaming:

```cpp
session->CancelProcess();  // Sets cancelled_ = true

```

Subsequent decode steps return early, and the callback receives a terminal response with a cancelled error status. The `Session::WaitUntilDone()` method blocks the caller thread until the final callback executes, ensuring clean shutdown.

## Summary

- **Non-blocking architecture**: `GenerateContentStream` delegates to `ThreadPool::Schedule` for prefill and decode operations, keeping the caller thread responsive.
- **Callback contract**: User code provides `absl::AnyInvocable<void(absl::StatusOr<Responses>)>` which `Tasks::Decode` invokes after each token generation step.
- **State management**: Inspect `Responses::state()` to distinguish between streaming chunks (`kProcessing`) and terminal conditions (`kDone`, `kError`).
- **Cancellation support**: Call `Session::CancelProcess()` to set the atomic cancellation flag and trigger early termination.
- **Source locations**: Core logic resides in `runtime/core/session_basic.cc`, `runtime/core/pipeline.cc`, and `runtime/core/tasks.cc`.

## Frequently Asked Questions

### How do I handle errors in the streaming callback?

Check the `absl::StatusOr<Responses>` parameter using `resp_status.ok()` before accessing the response. If the status is not OK, log the error and return early to avoid dereferencing invalid data. The callback may receive errors from model execution failures, invalid input shapes, or cancellation requests.

### Can I cancel an ongoing streaming inference request?

Yes. Call `session->CancelProcess()` from any thread to set the internal `cancelled_` atomic flag. The next decode step in `runtime/core/tasks.cc` checks this flag and returns early, triggering a final callback with a cancelled status. This is useful for implementing stop-generation buttons in user interfaces.

### What thread executes the user callback?

The callback executes on the background worker thread performing the decode operation, as implemented in `runtime/framework/threadpool.cc`. If your callback performs heavy work, offload it to your own thread pool to avoid blocking the inference pipeline. LiteRT-LM does not guarantee which specific thread invokes the callback, only that it occurs after each decode step.

### How does prefill relate to the streaming callback?

Prefill runs asynchronously via `RunPrefillAsync` before decode streaming begins. The user callback is not invoked during prefill; it only receives tokens once the decode phase starts. The prefill callback is an internal mechanism that chains into `DecodeInternalStreaming`, ensuring the transition from prompt processing to token generation is seamless.