Implementation:Mlc ai Mlc llm Logit Processor Header

Overview

The file cpp/serve/logit_processor.h defines the LogitProcessor component of the MLC-LLM serving engine. This header declares an abstract interface for modifying raw model logits in-place before token sampling occurs. Logit processing encompasses applying presence and frequency penalties, logit biases, and temperature adjustments to the model's output probability space.

The design follows the TVM object system pattern, with a pure virtual LogitProcessorObj base class and a corresponding LogitProcessor managed reference type. This separation allows for different backend implementations while maintaining a consistent interface throughout the serving pipeline.

File Location

cpp/serve/logit_processor.h

Dependencies

The header includes the following dependencies:

Header	Purpose
tvm/ffi/string.h	TVM string type for request IDs
tvm/runtime/module.h	TVM runtime module and tensor support
../base.h	MLC-LLM base definitions
config.h	GenerationConfig definitions
event_trace_recorder.h	Optional event tracing support
function_table.h	Packed function table for GPU kernel dispatch
request_state.h	RequestModelState for tracking per-request state

Namespace

All types are defined within mlc::llm::serve. The header imports tvm::Device and the entire tvm::runtime namespace for convenience.

Class: LogitProcessorObj

LogitProcessorObj is an abstract base class inheriting from tvm::runtime::Object. It defines two pure virtual methods for logit manipulation and probability computation.

Method: InplaceUpdateLogits

virtual void InplaceUpdateLogits(
    Tensor logits, const Array<GenerationConfig>& generation_cfg,
    const Array<RequestModelState>& mstates, const Array<String>& request_ids,
    const std::vector<int>* cum_num_token = nullptr,
    const Array<RequestModelState>* draft_mstates = nullptr,
    const std::vector<std::vector<int>>* draft_token_indices = nullptr) = 0;

This method modifies the logits tensor in-place. Key parameters:

• logits -- A tensor of shape (num_total_token, vocab_size). The total token count may exceed the batch size when sequences contribute more than one token (e.g., during multi-step speculation).
• generation_cfg -- An array of generation configs, one per sequence, controlling penalty parameters such as temperature, top-p, repetition penalty, and frequency penalty.
• mstates -- Per-sequence model states containing committed token histories and appeared-token counts used for penalty calculations.
• request_ids -- Unique identifiers for each request, used primarily for tracing.
• cum_num_token -- Optional cumulative token counts. When nullptr, each sequence is assumed to have exactly one token in the batch.
• draft_mstates -- Optional draft model states used during speculative decoding.
• draft_token_indices -- Optional indices mapping draft tokens to their positions. A value of -1 indicates a non-draft token. This is used to reconstruct the full sequence state including draft tokens without modifying the persisted state.

Method: ComputeProbsFromLogits

virtual Tensor ComputeProbsFromLogits(Tensor logits,
                                      const Array<GenerationConfig>& generation_cfg,
                                      const Array<String>& request_ids,
                                      const std::vector<int>* cum_num_token = nullptr) = 0;

Converts the (already-processed) logits into probability distributions, typically via a softmax operation. Returns the resulting probability tensor on GPU.

TVM Object Registration

static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LogitProcessorObj>();
}

static constexpr const bool _type_has_method_sequal_reduce = false;
static constexpr const bool _type_has_method_shash_reduce = false;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.LogitProcessor", LogitProcessorObj, Object);

The object is registered under the type key "mlc.serve.LogitProcessor". It is marked as mutable since it modifies tensors in-place. Structural equality and hash reduce are disabled.

Class: LogitProcessor

LogitProcessor is the managed reference (smart pointer wrapper) for LogitProcessorObj.

Constructor

explicit LogitProcessor(int max_num_token, int vocab_size, FunctionTable* ft, DLDevice device,
                        Optional<EventTraceRecorder> trace_recorder);

Parameter	Description
max_num_token	Maximum number of tokens the processor can handle in a single batch
vocab_size	Size of the model vocabulary
ft	Pointer to the packed function table for GPU kernel dispatch
device	The DL device (GPU) on which the model runs
trace_recorder	Optional event trace recorder for performance profiling

The reference type uses TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE, meaning it may hold a null reference.

Role in the Serving Pipeline

The LogitProcessor sits between the model forward pass (which produces raw logits) and the Sampler (which draws tokens from probability distributions). The typical flow is:

1. The model produces raw logits of shape (batch_size, vocab_size).
2. InplaceUpdateLogits applies penalties and biases to the raw logits.
3. ComputeProbsFromLogits transforms the adjusted logits into probability distributions.
4. The Sampler draws tokens from these probability distributions.

See Also

• Sampler Header -- The sampler that consumes the processed probabilities
• Request State Implementation -- RequestModelState, which tracks token histories used for penalty computation
• Metrics Header -- Metrics tracking for the serving engine