Repository: mlc-ai/mlc-llm
File: cpp/serve/engine.cc (C++)

Code before:
      String model_path = std::get<1>(model_info);
      DLDevice device = std::get<2>(model_info);
      Model model = Model::Create(model_lib, std::move(model_path), device,
                                  kv_cache_config_->max_num_sequence);
      model->CreateKVCache(this->kv_cache_config_);
      CHECK_GE(model->GetMaxWindowSize(), this->max_single_sequence_length_)
          << "The window size of the model, " << model->GetMaxWindowSize()
          << ", is smaller than the pre-defined max single sequence length, "
          << this->max_single_sequence_length_;
      this->models_.push_back(model);
      this->model_workspaces_.push_back(ModelWorkspace{model->AllocEmbeddingTensor()});
    }
    int max_num_tokens = kv_cache_config_->max_num_sequence;
    if (engine_mode_->enable_speculative) {
      max_num_tokens *= engine_mode_->spec_draft_length;
    }
    LogitProcessor logit_processor =
        this->models_[0]->CreateLogitProcessor(max_num_tokens, trace_recorder);
    bool is_speculative = this->models_.size() > 1U;
    // TODO: need a better sampler selection logic
    std::string sampler_selection = "auto";
    if (is_speculative) {
      sampler_selection = "cpu";
    }
    Sampler sampler =
        this->models_[0]->CreateSampler(max_num_tokens, sampler_selection, trace_recorder);
    // Step 3. Initialize engine actions that represent state transitions.
    if (this->engine_mode_->enable_speculative) {
      // Speculative decoding is only possible for more than one model.
      ICHECK_GT(this->models_.size(), 1U);
      this->actions_ = {
          EngineAction::NewRequestPrefill(this->models_,            //
                                          logit_processor,          //
                                          sampler,                  //
                                          this->model_workspaces_,  //
                                          this->kv_cache_config_,   //
                                          this->engine_mode_,       //
                                          this->trace_recorder_),
          EngineAction::BatchDraft(this->models_, logit_processor, sampler, this->trace_recorder_,
                                   this->engine_mode_->spec_draft_length),
          EngineAction::BatchVerify(this->models_, logit_processor, sampler, this->kv_cache_config_,
                                    this->trace_recorder_)};
    } else {
      this->actions_ = {EngineAction::NewRequestPrefill(this->models_,            //
                                                        logit_processor,          //
                                                        sampler,                  //
                                                        this->model_workspaces_,  //
                                                        this->kv_cache_config_,   //
                                                        this->engine_mode_,       //
                                                        this->trace_recorder_),
                        EngineAction::BatchDecode(this->models_, logit_processor, sampler,

Review comment (refactor):
Let's just pass the models as `const Array<Model>&` to `CreateSampler`, and select CPU sampler inside `CreateSampler` when the number of models is more than 1. The goal is to abstract, hide and not complicate the logic.

Code after:
      String model_path = std::get<1>(model_info);
      DLDevice device = std::get<2>(model_info);
      Model model = Model::Create(model_lib, std::move(model_path), device,
                                  kv_cache_config_->max_num_sequence);
      model->CreateKVCache(this->kv_cache_config_);
      CHECK_GE(model->GetMaxWindowSize(), this->max_single_sequence_length_)
          << "The window size of the model, " << model->GetMaxWindowSize()
          << ", is smaller than the pre-defined max single sequence length, "
          << this->max_single_sequence_length_;
      this->models_.push_back(model);
      this->model_workspaces_.push_back(ModelWorkspace{model->AllocEmbeddingTensor()});
    }
    int max_num_tokens = kv_cache_config_->max_num_sequence;
    if (engine_mode_->enable_speculative) {
      max_num_tokens *= engine_mode_->spec_draft_length;
    }
    LogitProcessor logit_processor =
        this->models_[0]->CreateLogitProcessor(max_num_tokens, trace_recorder);
    Sampler sampler = this->models_[0]->CreateSampler(
        max_num_tokens, static_cast<int>(this->models_.size()), trace_recorder);
    // Step 3. Initialize engine actions that represent state transitions.
    if (this->engine_mode_->enable_speculative) {
      // Speculative decoding is only possible for more than one model.
      ICHECK_GT(this->models_.size(), 1U);
      this->actions_ = {
          EngineAction::NewRequestPrefill(this->models_,            //
                                          logit_processor,          //
                                          sampler,                  //
                                          this->model_workspaces_,  //
                                          this->kv_cache_config_,   //
                                          this->engine_mode_,       //
                                          this->trace_recorder_),
          EngineAction::BatchDraft(this->models_, logit_processor, sampler, this->trace_recorder_,
                                   this->engine_mode_->spec_draft_length),
          EngineAction::BatchVerify(this->models_, logit_processor, sampler, this->kv_cache_config_,
                                    this->trace_recorder_)};
    } else {
      this->actions_ = {EngineAction::NewRequestPrefill(this->models_,            //
                                                        logit_processor,          //
                                                        sampler,                  //
                                                        this->model_workspaces_,  //
                                                        this->kv_cache_config_,   //
                                                        this->engine_mode_,       //
                                                        this->trace_recorder_),
                        EngineAction::BatchDecode(this->models_, logit_processor, sampler,
                                                  this->trace_recorder_)};
    }
    // Step 4. Automatically set the threading backend max concurrency.
    SetThreadMaxConcurrency();
  }