Implementation:Mlc ai Mlc llm Disagg Remote Send

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

Overview

DisaggRemoteSend implements the engine action that performs prefill computation on a prefill instance and marks the resulting KV cache data for transmission to a remote decode instance in a disaggregated LLM serving system.

Description

The disagg_remote_send.cc file implements the DisaggRemoteSendActionObj class, an engine action that runs on the prefill instance of a disaggregated serving architecture. While DisaggPrepareReceive prepares the decode side to receive KV data, this action performs the actual prefill computation and marks the KV data for sending.

The class extends BatchPrefillBaseActionObj and operates in the mlc::llm::serve namespace. Its constructor accepts model instances, model workspaces (for embeddings), engine config, model configs, a trace recorder, a stream callback, and the device. If the device is CUDA or ROCm, it captures the compute stream for synchronization.

The Step method implements the full prefill-and-send pipeline:

1. Calls GetRequestStateEntriesToPrefill to find eligible requests from the waiting queue (filtered to only kRemoteSend kind).
2. Matches prefix cache entries for each request via MatchPrefixCache.
3. Updates request states from pending to alive.
4. For each model, computes embeddings by iterating through the input data. Token data is batched together for efficient embedding computation, while non-token data (e.g., image embeddings) is processed separately. The embedding accumulation is done at tracked offsets using GetEmbedding with offset parameters.
5. Calls BatchPrefill to run the actual prefill forward pass and generate logits. Note that logits are not used for sampling in this action since this is only the prefill side.
6. Commits prefix cache changes.
7. Performs an explicit StreamSync to ensure the prefill computation completes (since there is no sampling step that would implicitly synchronize).
8. Removes processed requests from the waiting queue and updates running request state entries.

The private GetRequestStateEntriesToPrefill method selects requests from the waiting queue where disagg_config.kind == DisaggRequestKind::kRemoteSend. It performs a two-attempt strategy for each request: first attempting the full input length, and if that fails due to resource constraints, attempting a chunked input that fits within the remaining prefill chunk size. It handles sliding window considerations, page requirements, and reduces across all models to find the common maximum prefill inputs.

The MatchPrefixCache method is copied from NewRequestPrefillActionObj with the addition of DisaggMarkKVSend calls. When a sequence is added, forked, or reused, it also marks the KV data for sending by calling models_[model_id]->DisaggMarkKVSend with the appropriate kv_window_begin, kv_append_metadata, and dst_group_offset from the disaggregation configuration.

Usage

Use DisaggRemoteSend on the prefill instance of a disaggregated serving deployment. It is registered as an EngineAction via the static factory method EngineAction::DisaggRemoteSend and runs as part of the engine's action loop.

Code Reference

Source Location

• Repository: Mlc_ai_Mlc_llm
• File: cpp/serve/engine_actions/disagg_remote_send.cc
• Lines: 1-498

Signature

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

  Array<Request> Step(EngineState estate) final;

private:
  std::vector<PrefillInput> GetRequestStateEntriesToPrefill(EngineState estate);
  int MatchPrefixCache(EngineState estate, PrefillInput* input) final;

  std::vector<ModelWorkspace> model_workspaces_;
  FRequestStreamCallback request_stream_callback_;
  Device device_;
  TVMStreamHandle compute_stream_ = nullptr;
};

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

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 prefill computation and KV cache operations
model_workspaces	std::vector<ModelWorkspace>	Yes	Per-model workspaces containing embedding tensors
engine_config	EngineConfig	Yes	Engine configuration with page size, prefill chunk size, 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 results back
device	Device	Yes	Target device for computation and stream synchronization
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 (prefilled and marked for KV send)

Usage Examples

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

// The action is used in the engine loop on the prefill instance:
Array<Request> processed = action->Step(engine_state);
// After Step completes, KV data has been computed and marked for
// transmission to the decode instance.