Implementation:Mlc ai Mlc llm Request State
| Knowledge Sources | |
|---|---|
| Domains | LLM Serving, Request Management, Speculative Decoding |
| Last Updated | 2026-02-09 19:00 GMT |
Overview
A header file defining the data structures that maintain the generation states of user requests in the MLC LLM serving engine, including per-model state and per-generation-entry state.
Description
The request_state.h header defines several interrelated classes that track the complete lifecycle and generation state of each request being processed by the MLC LLM serving engine.
RequestModelStateNode holds the state of a single request with respect to a single model. In MLC LLM, the engine may use multiple models (e.g., a small draft model and a large verification model for speculative decoding). This class tracks:
- committed_tokens: Tokens that have been finalized and will not change.
- inputs: Input data yet to be prefilled.
- draft_output_tokens: Draft tokens generated by a speculative model, awaiting verification.
- appeared_token_ids: A frequency map of all tokens that have appeared, used for repetition penalty.
- grammar_matcher: An optional xgrammar grammar matcher for grammar-guided generation (e.g., JSON mode).
- Methods for committing tokens, rolling back tokens, managing draft tokens, and producing token bitmasks for constrained generation.
RequestStateEntryNode represents the state of a single generation within a request. A request with parallel generations (n > 1 in the OpenAI API sense) will have multiple state entries organized in a tree structure: a root entry for the shared prompt prefix and child entries for each parallel generation. Each entry contains:
- An array of RequestModelState objects (one per model).
- A RandomGenerator for sampling.
- A StopStrHandler for detecting stop strings.
- Callback position tracking for streaming output.
RequestStateNode groups all state entries for a single request, along with metrics tracking and post-processing workspace data structures. The post-processing workspace is maintained as state to avoid repetitive memory allocation during action post-processing.
DeltaRequestReturn is a helper struct for streaming output, containing delta token IDs, log-probability JSON strings, an optional finish reason, and an extra prefix string.
RequestStateStatus is an enum with values kPending, kAlive, and kFinished.
Usage
These data structures are used internally by the MLC LLM serving engine to track the state of every active request. They are created when a request is added to the engine and updated during each decode step. The multi-model design supports speculative decoding workflows, while the grammar matcher integration enables constrained generation.
Code Reference
Source Location
- Repository: Mlc_ai_Mlc_llm
- File: cpp/serve/request_state.h
Signature
class RequestModelStateNode : public Object {
public:
Request request;
int64_t internal_id = -1;
int model_id = -1;
std::vector<SampleResult> committed_tokens;
Array<Data> inputs;
std::optional<xgrammar::GrammarMatcher> grammar_matcher;
int GetInputLength() const;
bool RequireNextTokenBitmask();
void GetNextTokenBitmask(DLTensor* bitmask);
void CommitToken(SampleResult sampled_token);
void RollbackTokens(int count);
void AddDraftToken(SampleResult sampled_token, int draft_token_slot, int64_t parent_idx);
void RemoveAllDraftTokens(std::vector<int>* removed_draft_token_slots = nullptr);
};
class RequestStateEntryNode : public Object {
public:
RequestStateStatus status;
Request request;
int parent_idx = -1;
std::vector<int> child_indices;
Array<RequestModelState> mstates;
RandomGenerator rng;
StopStrHandler stop_str_handler;
void GetDeltaRequestReturn(const Tokenizer& tokenizer, int64_t max_single_sequence_length,
RequestStreamOutput* delta_stream_output, int idx);
};
class RequestStateNode : public Object {
public:
std::vector<RequestStateEntry> entries;
RequestMetrics metrics;
RequestActionPostProcWorkspace postproc_states;
};
Import
#include "request_state.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| request | Request | Yes | The original user request object this state corresponds to. |
| model_id | int | Yes | The index of the model this state tracks (for multi-model setups). |
| internal_id | int64_t | Yes | The physical index of the request in the running queue (-1 if not running). |
| inputs | Array | Yes | Input data to be prefilled by the model. |
| compiled_grammar | optional<xgrammar::CompiledGrammar> | No | Optional compiled grammar for constrained generation. |
| num_models | int | Yes (for RequestStateEntry) | Number of models being used. |
| rng_seed | int | Yes (for RequestStateEntry) | Seed for the random number generator. |
| token_table | vector<string> | Yes (for RequestStateEntry) | The tokenizer vocabulary table. |
Outputs
| Name | Type | Description |
|---|---|---|
| committed_tokens | vector<SampleResult> | The finalized generated tokens and their probabilities. |
| draft_output_tokens | vector<SampleResult> | Draft tokens pending verification (speculative decoding). |
| GetInputLength() | int | Total length of all input data items. |
| RequireNextTokenBitmask() | bool | Whether grammar-guided generation bitmask is needed. |
| GetDeltaRequestReturn | void (DPS) | Writes delta token IDs, logprob strings, and finish reason into the output. |
Usage Examples
// Create a request model state for model 0
RequestModelState mstate(request, /*model_id=*/0, /*internal_id=*/42, inputs, compiled_grammar);
// Commit a generated token
mstate->CommitToken(sample_result);
// Check if grammar-guided bitmask is needed
if (mstate->RequireNextTokenBitmask()) {
mstate->GetNextTokenBitmask(bitmask_tensor);
}
// For speculative decoding: add a draft token
mstate->AddDraftToken(draft_result, slot, parent_idx);
// Create a request state entry
RequestStateEntry entry(request, num_models, internal_id, rng_seed, token_table, grammar);
// Create the full request state
std::vector<RequestStateEntry> entries = {root_entry, child_entry_1, child_entry_2};
RequestState rstate(entries, /*num_response=*/2, add_time_point);