Implementation:Mlc ai Mlc llm Disagg Prepare Recv

Knowledge Sources	Mlc_ai_Mlc_llm
Domains	Disaggregated Inference, KV Cache Management, LLM Serving
Last Updated	2026-02-09 19:00 GMT

Overview

DisaggPrepareReceive implements the engine action that prepares KV cache data locations on the receiving side of a disaggregated LLM serving system, reserving pages and matching prefix cache entries before actual KV data transfer.

Description

The disagg_prepare_recv.cc file implements the DisaggPrepareReceiveActionObj class, an engine action within MLC LLM's disaggregated serving architecture. In a disaggregated system, prefill computation and decode computation can happen on separate instances. This action runs on the decode instance to prepare for receiving KV cache data from a remote prefill instance.

The class extends BatchPrefillBaseActionObj and operates in the mlc::llm::serve namespace. Its constructor validates that the KV state kind is kKVCache (PagedKVCache), since only paged KV caches support KV migration between instances.

The Step method is the main entry point, called by the engine loop. It iterates over the waiting queue to find requests with DisaggRequestKind::kPrepareReceive. For each qualifying request, it:

1. Calls GetRequestStateEntriesToPrefill to determine if the request can be prefilled given current memory constraints.
2. Invokes MatchPrefixCache to check for prefix cache matches, potentially reusing or forking existing sequences.
3. Updates request state from pending to alive via UpdateRequestToAlive, then immediately removes the request from the running queue since it is only preparing to receive.
4. Reserves KV cache pages by calling DisaggPrepareKVRecv on each model, which returns compressed KV append metadata describing the page layout.
5. Commits prefix cache changes and removes the request from the waiting queue.
6. Constructs a response containing prompt_length, prefix_matched_length, and base64-encoded kv_append_metadata, and sends it via the stream callback.

The private GetRequestStateEntriesToPrefill method filters the waiting queue for kPrepareReceive requests and checks resource availability. It respects KV window parameters (kv_window_begin and kv_window_end) from the disaggregation config, which define the range of KV data to prepare. It verifies that the full input is not being prefilled on the remote machine (kv_window_end < input_length). It considers sliding window constraints, page availability, and maximum total sequence length.

The CanPrefill method applies heuristic conditions including a 400-page buffer to ensure at least one decode can be performed after prefill, along with checks on maximum number of sequences and speculative decoding factors.

The MatchPrefixCache method handles three scenarios: adding a new sequence (no prefix match), forking from an existing active sequence (partial prefix match), or reusing a recycling sequence (recycled prefix match). It pops already-prefilled input data and updates the max prefill length accordingly.

Usage

Use DisaggPrepareReceive as part of a disaggregated serving setup where the decode instance needs to pre-allocate KV cache pages before receiving actual KV data from a prefill instance. It is registered as an EngineAction via the static factory method EngineAction::DisaggPrepareReceive.

Code Reference

Source Location

• Repository: Mlc_ai_Mlc_llm
• File: cpp/serve/engine_actions/disagg_prepare_recv.cc
• Lines: 1-441

Signature

class DisaggPrepareReceiveActionObj : public BatchPrefillBaseActionObj {
public:
  explicit DisaggPrepareReceiveActionObj(
      Array<Model> models, EngineConfig engine_config,
      std::vector<picojson::object> model_configs,
      Optional<EventTraceRecorder> trace_recorder,
      FRequestStreamCallback request_stream_callback);

  Array<Request> Step(EngineState estate) final;

private:
  std::optional<PrefillInput> GetRequestStateEntriesToPrefill(EngineState estate);
  bool CanPrefill(EngineState estate, int num_prefill_rsentries,
                  int total_input_length, int num_required_pages,
                  int num_available_pages, int current_total_seq_len,
                  int num_running_rsentries, KVStateKind kv_state_kind,
                  bool sliding_window_enabled);
  int MatchPrefixCache(EngineState estate, PrefillInput* input) final;
};

// Factory method
EngineAction EngineAction::DisaggPrepareReceive(
    Array<Model> models, EngineConfig engine_config,
    std::vector<picojson::object> model_configs,
    Optional<EventTraceRecorder> trace_recorder,
    FRequestStreamCallback request_stream_callback);

Import

#include "batch_prefill_base.h"
#include "../sampler/sampler.h"

I/O Contract

Inputs

Name	Type	Required	Description
models	Array<Model>	Yes	Array of model instances for KV cache operations
engine_config	EngineConfig	Yes	Engine configuration with page size, max sequences, etc.
model_configs	std::vector<picojson::object>	Yes	Per-model configuration objects
trace_recorder	Optional<EventTraceRecorder>	No	Optional event trace recorder for profiling
request_stream_callback	FRequestStreamCallback	Yes	Callback to stream KV metadata back to the caller
estate	EngineState	Yes (Step)	The current engine state with waiting/running queues and prefix cache

Outputs

Name	Type	Description
Array<Request>	TVM Array	The requests that were processed in this step
Stream callback	JSON via FRequestStreamCallback	Contains prompt_length, prefix_matched_length, and base64-encoded kv_append_metadata

Usage Examples

// Creating the DisaggPrepareReceive engine action
EngineAction action = EngineAction::DisaggPrepareReceive(
    models,
    engine_config,
    model_configs,
    trace_recorder,
    request_stream_callback
);

// The action is then used in the engine loop:
Array<Request> processed = action->Step(engine_state);