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

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 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
SessionBasic	Orchestrates pre-fill and decode streaming; manages asynchronous execution.	runtime/core/session_basic.h
Pipeline	Forwards the user callback to low-level tasks during DecodeStreaming.	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

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:

#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:

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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}