Implementation:Mlc ai Mlc llm Engine Action

Knowledge Sources	Mlc_ai_Mlc_llm
Domains	LLM Serving, Engine Architecture, Speculative Decoding
Last Updated	2026-02-09 19:00 GMT

Overview

The Engine Action header defines the abstract action interface and declares the full catalog of concrete actions that the MLC LLM serving engine can execute at each time step. Actions encapsulate distinct engine behaviors such as prefilling new requests, decoding running requests, generating speculative draft tokens, verifying speculative proposals, and handling disaggregated serving.

Description

This header file (cpp/serve/engine_actions/action.h) defines two main types:

EngineActionObj (abstract base class):

Inherits from Object (TVM runtime object system).
Declares a single pure virtual method Step(EngineState estate) that takes the current engine state, performs the action (invoking model functions, running samplers), updates the engine state, and returns the set of processed requests.
Marked as mutable (_type_mutable = true) since actions may hold and modify internal state.

EngineAction (managed reference class with static factory methods):

The factory methods create the following action types:

Action	Description
`NewRequestPrefill`	Prefills requests from the engine's waiting queue using the main model. Runs batched prefill, applies logit processing, and samples initial tokens.
`EagleNewRequestPrefill`	Prefills requests for Eagle-style speculative decoding. Additionally manages draft token workspace allocation.
`BatchDecode`	Runs one-step decode for all requests in the running queue. Preempts low-priority requests when resources are insufficient. Not used for speculative decoding with multiple models.
`BatchDraft`	Generates speculative draft proposals for running requests using the draft model.
`EagleBatchDraft`	Eagle-specific batch draft proposal generation with draft token workspace management.
`BatchVerify`	Verifies speculative draft proposals using the main model. Accepts or rejects proposed tokens.
`EagleBatchVerify`	Eagle-specific verification of draft proposals.
`BatchJumpForward`	Predicts next tokens according to grammar constraints via jump-forward decoding. Handles retokenization when predicted strings cross tokenization boundaries.
`AutoSpecDecode`	A meta-action that dynamically decides between speculative decoding and normal batch decode based on runtime conditions.
`DisaggPrepareReceive`	Prepares for disaggregated prefill by computing KV cache metadata and prefix cache match information.
`DisaggRemoteSend`	Runs prefill and sends the resulting KV data to a remote decode instance in disaggregated serving.

Usage

Engine actions are created during engine initialization and composed into an action pipeline. The engine's Step() method delegates to the appropriate action sequence:

Normal mode: NewRequestPrefill followed by BatchDecode (with optional BatchJumpForward).
Speculative mode: NewRequestPrefill, then BatchDraft + BatchVerify (or AutoSpecDecode for adaptive switching).
Eagle mode: Uses EagleNewRequestPrefill, EagleBatchDraft, and EagleBatchVerify.
Disaggregated mode: Uses DisaggPrepareReceive and DisaggRemoteSend.

Code Reference

Source Location

Property	Value
File	`cpp/serve/engine_actions/action.h`
Namespace	`mlc::llm::serve`
Lines	258
Include Guard	`MLC_LLM_SERVE_ENGINE_ACTIONS_ACTION_H_`

Signature

namespace mlc {
namespace llm {
namespace serve {

class EngineActionObj : public Object {
 public:
  virtual Array<Request> Step(EngineState estate) = 0;
};

class EngineAction : public ObjectRef {
 public:
  static EngineAction NewRequestPrefill(
      Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
      std::vector<ModelWorkspace> model_workspaces, EngineConfig engine_config,
      std::vector<picojson::object> model_configs,
      Optional<EventTraceRecorder> trace_recorder);

  static EngineAction EagleNewRequestPrefill(
      Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
      std::vector<ModelWorkspace> model_workspaces,
      DraftTokenWorkspaceManager draft_token_workspace_manager,
      EngineConfig engine_config, std::vector<picojson::object> model_configs,
      Optional<EventTraceRecorder> trace_recorder);

  static EngineAction BatchDecode(
      Array<Model> models, Tokenizer tokenizer, LogitProcessor logit_processor,
      Sampler sampler, EngineConfig engine_config,
      Optional<EventTraceRecorder> trace_recorder);

  static EngineAction BatchDraft(
      Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
      std::vector<ModelWorkspace> model_workspaces,
      DraftTokenWorkspaceManager draft_token_workspace_manager,
      EngineConfig engine_config, Optional<EventTraceRecorder> trace_recorder);

  static EngineAction EagleBatchDraft(
      Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
      std::vector<ModelWorkspace> model_workspaces,
      DraftTokenWorkspaceManager draft_token_workspace_manager,
      EngineConfig engine_config, Optional<EventTraceRecorder> trace_recorder);

  static EngineAction BatchVerify(
      Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
      std::vector<ModelWorkspace> model_workspaces,
      DraftTokenWorkspaceManager draft_token_workspace_manager,
      EngineConfig engine_config, Optional<EventTraceRecorder> trace_recorder);

  static EngineAction EagleBatchVerify(
      Array<Model> models, LogitProcessor logit_processor, Sampler sampler,
      std::vector<ModelWorkspace> model_workspaces,
      DraftTokenWorkspaceManager draft_token_workspace_manager,
      EngineConfig engine_config, Optional<EventTraceRecorder> trace_recorder);

  static EngineAction BatchJumpForward(
      Array<Model> models, Tokenizer tokenizer,
      Optional<EventTraceRecorder> trace_recorder);

  static EngineAction AutoSpecDecode(
      std::vector<EngineAction> spec_decode_actions,
      std::vector<EngineAction> batch_decode_actions,
      EngineConfig engine_config);

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

  static 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);
};

}  // namespace serve
}  // namespace llm
}  // namespace mlc

Import

#include "serve/engine_actions/action.h"

Dependencies:

../config.h for EngineConfig and related configuration types
../draft_token_workspace_manager.h for DraftTokenWorkspaceManager
../engine.h for Engine and FRequestStreamCallback
../engine_state.h for EngineState
../event_trace_recorder.h for EventTraceRecorder
../model.h for Model and ModelWorkspace
../sampler/sampler.h for Sampler and LogitProcessor

I/O Contract

EngineActionObj::Step

Direction	Name	Type	Description
Input	estate	`EngineState`	The engine state containing waiting_queue, running_queue, and all associated state
Output	(return)	`Array<Request>`	The requests that were processed in this step
Side Effect	estate (mutated)	`EngineState`	The engine state is updated (requests moved between queues, tokens appended, etc.)

Common Factory Method Parameters

Parameter	Type	Used By	Description
`models`	`Array<Model>`	All actions	The model(s) to run inference on
`logit_processor`	`LogitProcessor`	Prefill, Decode, Draft, Verify	Processes raw model logits before sampling
`sampler`	`Sampler`	Prefill, Decode, Draft, Verify	Samples tokens from processed logit distributions
`tokenizer`	`Tokenizer`	BatchDecode, JumpForward	Tokenizer for detokenization and retokenization
`model_workspaces`	`std::vector<ModelWorkspace>`	Prefill, Draft, Verify, DisaggSend	Per-model workspace buffers
`draft_token_workspace_manager`	`DraftTokenWorkspaceManager`	Eagle actions, Draft, Verify	Manages draft token workspace slots
`engine_config`	`EngineConfig`	All actions	Engine-wide configuration
`model_configs`	`std::vector<picojson::object>`	Prefill, Disagg actions	Per-model JSON configs
`trace_recorder`	`Optional<EventTraceRecorder>`	All actions	Optional event tracing
`request_stream_callback`	`FRequestStreamCallback`	Disagg actions	Callback for streaming results

Usage Examples

Creating actions for normal serving:

#include "serve/engine_actions/action.h"

// Create prefill action
EngineAction prefill = EngineAction::NewRequestPrefill(
    models, logit_processor, sampler, model_workspaces,
    engine_config, model_configs, trace_recorder);

// Create decode action
EngineAction decode = EngineAction::BatchDecode(
    models, tokenizer, logit_processor, sampler,
    engine_config, trace_recorder);

// Execute a step
Array<Request> processed = prefill->Step(engine_state);
Array<Request> decoded = decode->Step(engine_state);

Creating actions for Eagle speculative decoding:

// Eagle prefill
EngineAction eagle_prefill = EngineAction::EagleNewRequestPrefill(
    models, logit_processor, sampler, model_workspaces,
    draft_token_workspace_manager, engine_config, model_configs, trace_recorder);

// Eagle draft
EngineAction eagle_draft = EngineAction::EagleBatchDraft(
    models, logit_processor, sampler, model_workspaces,
    draft_token_workspace_manager, engine_config, trace_recorder);

// Eagle verify
EngineAction eagle_verify = EngineAction::EagleBatchVerify(
    models, logit_processor, sampler, model_workspaces,
    draft_token_workspace_manager, engine_config, trace_recorder);

Creating an adaptive speculative decoding action:

std::vector<EngineAction> spec_actions = {eagle_draft, eagle_verify};
std::vector<EngineAction> normal_actions = {decode};

EngineAction auto_action = EngineAction::AutoSpecDecode(
    spec_actions, normal_actions, engine_config);

// The auto action dynamically selects between speculative and normal decoding
Array<Request> processed = auto_action->Step(engine_state);

Related Pages

Mlc_ai_Mlc_llm_Engine_Interface - The engine that composes and executes these actions
Mlc_ai_Mlc_llm_Draft_Token_Workspace - Workspace manager used by speculative decoding actions
Mlc_ai_Mlc_llm_Serve_Data_Header - Data types processed during action execution
Mlc_ai_Mlc_llm_Serve_Data - Data implementations used in prefill and decode
Mlc_ai_Mlc_llm_Model_Metadata_Header - Model metadata that influences action behavior

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment