Implementation:InternLM Lmdeploy OutputProcessor
| Knowledge Sources | |
|---|---|
| Domains | Inference Engine, Output Processing |
| Last Updated | 2026-02-07 15:00 GMT |
Overview
Manages the extraction and output of hidden states and logits from the language model, routing full-sequence and per-token outputs to the appropriate request output buffers.
Description
The OutputProcessor class handles the complex task of extracting hidden states and logits from the model's intermediate representations and writing them to per-request output tensors. It uses the pimpl idiom with a private Impl struct.
The processor supports two output types: type 1 (selected token outputs, e.g., the last token per request) and type 2 (full sequence outputs across all input tokens).
Add (kAdd): Initializes output intervals for each new request. Based on GenerationConfig::output_logits and output_last_hidden_state, it sets up Interval ranges tracking which positions need output. For "all" mode, output starts from the initial step; for "generation" mode, output starts from the prompt boundary.
Setup (kSetup): Computes which token positions from the current batch need hidden states and logits output. Uses a Matching helper that performs interval intersection to map batch positions to request output positions. Builds lists of (request_index, type, source_interval, dest_interval) tuples for both hidden states and logits. The full hidden states range is extended to cover logits needs (since logits are computed from hidden states).
Prepare (kPrepare): If full hidden states output is needed, signals the model to produce them by adding an "output_hidden_states" entry to the environment.
OutputHiddenStatesAndLogits(): The main output method, called after the model forward pass. For type 2 (full), it retrieves the full hidden states tensor, copies matching intervals to request output buffers, and computes logits by running the LM head in chunks (bounded by max_logits_len_) to manage GPU memory. For type 1 (selected), it outputs from the per-batch selected token hidden states and logits. Logits are copied via 2D CUDA memcpy to handle stride differences between source and destination tensors.
Usage
Instantiated during engine initialization. Called at kAdd, kSetup, and kPrepare stages, plus directly for hidden states and logits output after the model forward pass.
Code Reference
Source Location
- Repository: InternLM_Lmdeploy
- File: src/turbomind/models/output_processor.h
- File: src/turbomind/models/output_processor.cc
- Lines: output_processor.h 1-27, output_processor.cc 1-313
Signature
class OutputProcessor {
public:
~OutputProcessor();
OutputProcessor(const ModelParam& model,
int max_logits_len,
int tp_rank,
int phases,
std::function<Tensor(const Tensor&)> lm_head);
void Run(BatchOp op, int phase, TensorMap& env);
void OutputHiddenStatesAndLogits(int phase, TensorMap& env, int type);
private:
struct Impl;
std::unique_ptr<Impl> impl_;
};
Import
#include "src/turbomind/models/output_processor.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| model | ModelParam | Yes | Model parameters (vocab_size) |
| max_logits_len | int | Yes | Maximum chunk size for logits computation (memory limit) |
| tp_rank | int | Yes | Tensor-parallel rank (only rank 0 writes outputs) |
| phases | int | Yes | Number of pipeline phases |
| lm_head | std::function<Tensor(const Tensor&)> | Yes | LM head function mapping hidden states to logits |
| env["requests"] | Buffer_<RequestCache*> | Yes (Add) | New requests for output interval initialization |
| env["batch"] | BatchData* | Yes (Setup/Output) | Current batch data with request caches |
| env["token_num"] | int* | Yes (Setup) | Total input token count for the batch |
| env["hidden_states"] | Tensor | Yes (Output type 1) | Selected token hidden states |
| env["logits"] | Tensor | Yes (Output type 1) | Selected token logits |
| env["full_hidden_states"] | Tensor | Yes (Output type 2) | Full sequence hidden states |
Outputs
| Name | Type | Description |
|---|---|---|
| Request.outputs["last_hidden_state"] | Tensor | Per-request hidden state output (when requested) |
| Request.outputs["logits"] | Tensor | Per-request logits output (when requested) |
| env["output_hidden_states"] | Tensor (signal) | Empty tensor signaling the model to produce full hidden states |
Usage Examples
// Construction
auto lm_head = [&](const Tensor& h) { return model.lm_head(h); };
OutputProcessor output_proc(model_param, max_logits_len, tp_rank, phases, lm_head);
// During batch processing:
output_proc.Run(BatchOp::kAdd, phase, env); // Initialize output intervals
output_proc.Run(BatchOp::kSetup, phase, env); // Compute output ranges
output_proc.Run(BatchOp::kPrepare, phase, env); // Signal model for full outputs
// After model forward pass:
output_proc.OutputHiddenStatesAndLogits(phase, env, 2); // Full sequence outputs
output_proc.OutputHiddenStatesAndLogits(phase, env, 1); // Selected token outputs