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

Code before:

  LogitProcessor CreateLogitProcessor(int max_num_token,
                                      Optional<EventTraceRecorder> trace_recorder) {
    return LogitProcessor(max_num_token, vocab_size_, &this->ft_, device_,
                          std::move(trace_recorder));
  }

  Sampler CreateSampler(int max_num_sample, std::string required_devices,
                        Optional<EventTraceRecorder> trace_recorder) {
    if (required_devices == "auto") {
      if (Sampler::SupportGPUSampler(device_)) {
        return Sampler::CreateGPUSampler(max_num_sample, vocab_size_, &this->ft_, device_,
                                         std::move(trace_recorder));
      } else {
        return Sampler::CreateCPUSampler(std::move(trace_recorder));
      }
    } else {
      CHECK(required_devices == "cpu" || required_devices == "gpu");
      if (required_devices == "cpu") {
        return Sampler::CreateCPUSampler(std::move(trace_recorder));
      } else if (required_devices == "gpu") {
        return Sampler::CreateGPUSampler(max_num_sample, vocab_size_, &this->ft_, device_,
                                         std::move(trace_recorder));
      }
    }
    CHECK(false) << "Invalid Sampler arguments.";
  }

  void CreateKVCache(KVCacheConfig kv_cache_config) final {
    IntTuple max_num_sequence{kv_cache_config->max_num_sequence};
    IntTuple max_total_sequence_length{kv_cache_config->max_total_sequence_length};
    IntTuple prefill_chunk_size{kv_cache_config->prefill_chunk_size};
    IntTuple page_size{kv_cache_config->page_size};
    IntTuple support_sliding_window{sliding_window_size_ != -1};
    kv_cache_ = ft_.create_kv_cache_func_(max_num_sequence, max_total_sequence_length,
                                          prefill_chunk_size, page_size, support_sliding_window);
    local_kv_cache_ = ft_.use_disco ? Downcast<DRef>(kv_cache_)->DebugGetFromRemote(0) : kv_cache_;
  }

  void AddNewSequence(int64_t seq_id) final { ft_.kv_cache_add_sequence_func_(kv_cache_, seq_id); }

  void ForkSequence(int64_t parent_seq_id, int64_t child_seq_id) final {
    ft_.kv_cache_fork_sequence_func_(kv_cache_, parent_seq_id, child_seq_id);
  }

  void RemoveSequence(int64_t seq_id) final {
    ft_.kv_cache_remove_sequence_func_(kv_cache_, seq_id);
  }

  void PopNFromKVCache(int64_t seq_id, int num_tokens) final {
    ft_.kv_cache_popn_func_(kv_cache_, seq_id, num_tokens);

Review comment (suggestion):
See the other comment. Take in `const Array<Model>& models`, update the function body to
```c++
if (models.size() == 1 && Sampler::SupportGPUSampler(device_)) {
  // create GPU sampler and return
} else {
  // create CPU sampler and return
}
```

Code after:

  LogitProcessor CreateLogitProcessor(int max_num_token,
                                      Optional<EventTraceRecorder> trace_recorder) {
    return LogitProcessor(max_num_token, vocab_size_, &this->ft_, device_,
                          std::move(trace_recorder));
  }

  Sampler CreateSampler(int max_num_sample, int num_models,
                        Optional<EventTraceRecorder> trace_recorder) {
    if (num_models > 1) {  // speculative decoding uses cpu sampler
      return Sampler::CreateCPUSampler(std::move(trace_recorder));
    } else if (Sampler::SupportGPUSampler(device_)) {
      return Sampler::CreateGPUSampler(max_num_sample, vocab_size_, &this->ft_, device_,
                                       std::move(trace_recorder));
    } else {
      return Sampler::CreateCPUSampler(std::move(trace_recorder));
    }
  }

  void CreateKVCache(KVCacheConfig kv_cache_config) final {
    IntTuple max_num_sequence{kv_cache_config->max_num_sequence};
    IntTuple max_total_sequence_length{kv_cache_config->max_total_sequence_length};
    IntTuple prefill_chunk_size{kv_cache_config->prefill_chunk_size};
    IntTuple page_size{kv_cache_config->page_size};
    IntTuple support_sliding_window{sliding_window_size_ != -1};
    kv_cache_ = ft_.create_kv_cache_func_(max_num_sequence, max_total_sequence_length,
                                          prefill_chunk_size, page_size, support_sliding_window);
    local_kv_cache_ = ft_.use_disco ? Downcast<DRef>(kv_cache_)->DebugGetFromRemote(0) : kv_cache_;
  }

  void AddNewSequence(int64_t seq_id) final { ft_.kv_cache_add_sequence_func_(kv_cache_, seq_id); }

  void ForkSequence(int64_t parent_seq_id, int64_t child_seq_id) final {
    ft_.kv_cache_fork_sequence_func_(kv_cache_, parent_seq_id, child_seq_id);
  }

  void RemoveSequence(int64_t seq_id) final {
    ft_.kv_cache_remove_sequence_func_(kv_cache_, seq_id);
  }

  void PopNFromKVCache(int64_t seq_id, int num_tokens) final {
    ft_.kv_cache_popn_func_(kv_cache_, seq_id, num_tokens);
  }

  void EnableSlidingWindowForSeq(int64_t seq_id) final {
    if (sliding_window_size_ != -1) {
      ft_.kv_cache_enable_sliding_window_for_seq_(kv_cache_, seq_id, sliding_window_size_,
                                                  attention_sink_size_);
    }
  }