feat: add onerec worker impl.

Author: DragonFive

xllm/core/common/CMakeLists.txt

@@ -14,6 +14,7 @@ cc_library(
     $<$<BOOL:${USE_NPU}>:mspti_helper.h>
     options.h
     rate_limiter.h
+    rec_model_utils.h
     types.h
     device_monitor.h
     version_singleton.h
@@ -66,4 +67,3 @@ cc_test(
 target_link_libraries(common PRIVATE OpenSSL::SSL OpenSSL::Crypto protobuf::libprotobuf)
 add_dependencies(common brpc-static)
 
-

xllm/core/common/rec_model_utils.h

@@ -0,0 +1,48 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <cstdint>
+#include <string_view>
+
+namespace xllm {
+
+enum class RecModelKind : int8_t {
+  kNone = 0,
+  kOneRec = 1,
+  kLlmRec = 2,
+};
+
+inline constexpr bool is_onerec_model_type(std::string_view model_type) {
+  return model_type == "onerec";
+}
+
+inline constexpr bool is_llmrec_model_type(std::string_view model_type) {
+  return model_type == "qwen2" || model_type == "qwen3";
+}
+
+inline constexpr RecModelKind get_rec_model_kind(std::string_view model_type) {
+  if (is_onerec_model_type(model_type)) {
+    return RecModelKind::kOneRec;
+  }
+  if (is_llmrec_model_type(model_type)) {
+    return RecModelKind::kLlmRec;
+  }
+  return RecModelKind::kNone;
+}
+
+}  // namespace xllm
+

xllm/core/distributed_runtime/master.cpp

@@ -22,6 +22,7 @@ limitations under the License.
 #include <boost/algorithm/string.hpp>
 #include <csignal>
 #include <memory>
+#include <optional>
 #include <thread>
 #include <utility>
 
@@ -30,6 +31,7 @@ limitations under the License.
 #include "common/types.h"
 #include "dit_master.h"
 #include "framework/model/model_args.h"
+#include "framework/model_loader.h"
 #include "framework/request/request.h"
 #include "llm_engine.h"
 #include "llm_master.h"
@@ -237,8 +239,13 @@ Master::Master(const Options& options, EngineType type) : options_(options) {
     options_.enable_schedule_overlap(false);
     LOG(WARNING) << "Force to disable schedule overlap for REC model, not "
                     "supported yet.";
+
+    auto model_loader = ModelLoader::create(options_.model_path());
+    const std::string rec_model_type = model_loader->model_args().model_type();
+
     runtime::Options eng_options;
     eng_options.model_path(options_.model_path())
+        .rec_model_type(std::make_optional(rec_model_type))
         .devices(devices)
         .backend(options_.backend())
         .block_size(options_.block_size())

xllm/core/distributed_runtime/rec_master.cpp

@@ -25,6 +25,7 @@ limitations under the License.
 
 #include "common/macros.h"
 #include "common/metrics.h"
+#include "common/rec_model_utils.h"
 #include "common/types.h"
 #include "framework/request/mm_data.h"
 #include "models/model_registry.h"
@@ -42,12 +43,14 @@ namespace {
 constexpr int32_t kDefaultPlaceholderToken = 20152019;
 
 RecType get_rec_type(const ModelArgs& model_args) {
-  const auto& model_type = model_args.model_type();
-  if (model_type == "onerec") {
-    return RecType::kOneRec;
-  }
-  if (model_type == "qwen2" || model_type == "qwen3") {
-    return RecType::kLlmRec;
+  const auto kind = get_rec_model_kind(model_args.model_type());
+  switch (kind) {
+    case RecModelKind::kOneRec:
+      return RecType::kOneRec;
+    case RecModelKind::kLlmRec:
+      return RecType::kLlmRec;
+    case RecModelKind::kNone:
+      return RecType::kNone;
   }
   return RecType::kNone;
 }

xllm/core/runtime/CMakeLists.txt

@@ -22,6 +22,9 @@ cc_library(
     dit_worker.h
     embed_worker_impl.h
     embed_vlm_worker_impl.h
+    rec_worker_impl.h
+    llmrec_worker_impl.h
+    onerec_worker_impl.h
     worker_client.h
     xservice_client.h
     speculative_worker_impl.h
@@ -38,6 +41,9 @@ cc_library(
     dit_worker.cpp
     embed_worker_impl.cpp
     embed_vlm_worker_impl.cpp
+    rec_worker_impl.cpp
+    llmrec_worker_impl.cpp
+    onerec_worker_impl.cpp
     worker_client.cpp
     xservice_client.cpp
     params_utils.cpp

xllm/core/runtime/llmrec_worker_impl.cpp

@@ -0,0 +1,131 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "llmrec_worker_impl.h"
+
+#include <glog/logging.h>
+#include <torch/torch.h>
+
+#include <algorithm>
+#include <vector>
+
+#include "common/types.h"
+#include "core/layers/word_embedding.h"
+
+namespace xllm {
+
+LlmRecWorkerImpl::LlmRecWorkerImpl(const ParallelArgs& parallel_args,
+                                   const torch::Device& device,
+                                   const runtime::Options& options)
+    : RecWorkerImpl(parallel_args, device, options) {}
+
+void LlmRecWorkerImpl::prepare_work_before_execute(
+    const ForwardInput& inputs,
+    ForwardInput& processed_inputs) {
+  WorkerImpl::prepare_work_before_execute(inputs, processed_inputs);
+
+  if (!inputs.input_params.mm_data.valid()) {
+    return;
+  }
+
+  torch::Tensor input_embedding;
+  torch::Tensor input_tokens_tensor;
+  torch::Tensor input_indices_tensor;
+
+  const auto& mm_data = inputs.input_params.mm_data;
+  const auto& processed_mm_data = processed_inputs.input_params.mm_data;
+
+  if (auto res =
+          processed_mm_data.get<torch::Tensor>(LLM_REC_INPUT_TOKENS)) {
+    input_tokens_tensor = res.value();
+  }
+
+  // input indices 需要在 Host 侧生成位置索引
+  if (auto res = mm_data.get<torch::Tensor>(LLM_REC_INPUT_INDICES)) {
+    input_indices_tensor = res.value();
+  }
+
+  if (auto res =
+          processed_mm_data.get<torch::Tensor>(LLM_REC_INPUT_EMBEDDING)) {
+    input_embedding = res.value();
+  }
+
+  if (input_embedding.defined()) {
+    input_embedding = input_embedding.to(dtype());
+  }
+
+  if (input_indices_tensor.defined()) {
+    layer::WordEmbedding word_embedding = get_word_embedding();
+    torch::Tensor input_tokens_embedding =
+        word_embedding(input_tokens_tensor, 0);
+
+    if (input_embedding.defined()) {
+      std::vector<int> input_indices(
+          input_indices_tensor.data_ptr<int>(),
+          input_indices_tensor.data_ptr<int>() + input_indices_tensor.numel());
+
+      processed_inputs.input_params.input_embedding =
+          merge_embeddings_by_indices(
+              input_tokens_embedding, input_embedding, input_indices);
+    } else {
+      processed_inputs.input_params.input_embedding = input_tokens_embedding;
+    }
+  } else if (input_embedding.defined()) {
+    processed_inputs.input_params.input_embedding = input_embedding;
+  }
+}
+
+torch::Tensor LlmRecWorkerImpl::merge_embeddings_by_indices(
+    const torch::Tensor& input_tokens_embedding,
+    const torch::Tensor& input_embedding,
+    const std::vector<int>& input_indices) {
+  CHECK_EQ(input_embedding.dim(), 2);
+  CHECK_EQ(input_tokens_embedding.dim(), 2);
+  CHECK_EQ(input_tokens_embedding.size(1), input_embedding.size(1));
+  CHECK_EQ(input_tokens_embedding.dtype(), input_embedding.dtype());
+  CHECK_EQ(input_tokens_embedding.device(), input_embedding.device());
+
+  const int64_t total_rows =
+      input_tokens_embedding.size(0) + input_embedding.size(0);
+  const int64_t cols = input_embedding.size(1);
+
+  torch::Device device = input_embedding.device();
+  torch::Tensor merged = torch::empty(
+      {total_rows, cols}, torch::dtype(input_embedding.dtype()).device(device));
+
+  std::vector<int> input_embedding_indices;
+  for (int i = 0; i < static_cast<int>(total_rows); ++i) {
+    if (std::find(input_indices.begin(), input_indices.end(), i) ==
+        input_indices.end()) {
+      input_embedding_indices.push_back(i);
+    }
+  }
+
+  CHECK_EQ(input_embedding_indices.size(), input_embedding.size(0));
+
+  torch::Tensor input_embedding_indices_tensor =
+      torch::tensor(input_embedding_indices, torch::kInt64).to(device);
+  merged.index_put_({input_embedding_indices_tensor, torch::indexing::Ellipsis},
+                    input_embedding);
+
+  torch::Tensor input_indices_tensor =
+      torch::tensor(input_indices, torch::kInt64).to(device);
+  merged.index_put_({input_indices_tensor, torch::indexing::Ellipsis},
+                    input_tokens_embedding);
+
+  return merged;
+}
+
+}  // namespace xllm

xllm/core/runtime/llmrec_worker_impl.h

@@ -0,0 +1,44 @@
+/* Copyright 2025 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include <vector>
+
+#include "runtime/rec_worker_impl.h"
+
+namespace xllm {
+
+class LlmRecWorkerImpl final : public RecWorkerImpl {
+ public:
+  LlmRecWorkerImpl(const ParallelArgs& parallel_args,
+                   const torch::Device& device,
+                   const runtime::Options& options);
+
+  ~LlmRecWorkerImpl() override = default;
+
+  void prepare_work_before_execute(const ForwardInput& inputs,
+                                   ForwardInput& processed_inputs) override;
+
+ private:
+  torch::Tensor merge_embeddings_by_indices(
+      const torch::Tensor& input_tokens_embedding,
+      const torch::Tensor& input_embedding,
+      const std::vector<int>& input_indices);
+};
+
+}  // namespace xllm