bugix: cp+mtp get hidden from lmhead

Author: shifengmin

xllm/core/framework/model/causal_lm.h

@@ -74,6 +74,17 @@ class CausalLM : public torch::nn::Module {
   virtual torch::Tensor logits(const torch::Tensor& hidden_states,
                                const torch::Tensor& seleted_idxes) = 0;
 
+  // hidden_states: [num_tokens, hidden_size]
+  // seleted_idxes: [num_tokens]
+  // out_hidden: [num_selected_tokens, hidden_size]
+  // returns: [num_selected_tokens, vocab_size]
+  virtual torch::Tensor logits(const torch::Tensor& hidden_states,
+                               const torch::Tensor& seleted_idxes,
+                               torch::Tensor& out_hidden) {
+    NOT_IMPLEMENTED();
+    return torch::Tensor();
+  }
+
   virtual void load_model(std::unique_ptr<ModelLoader> loader) = 0;
 
   virtual torch::Device device() const = 0;
@@ -158,6 +169,16 @@ class CausalLMImpl : public CausalLM {
     return model_->logits(hidden_states, seleted_idxes);
   }
 
+  torch::Tensor logits(const torch::Tensor& hidden_states,
+                       const torch::Tensor& seleted_idxes,
+                       torch::Tensor& out_hidden) override {
+    if constexpr (detail::has_logits_with_hidden<Model>::value) {
+      return model_->logits(hidden_states, seleted_idxes, out_hidden);
+    } else {
+      return CausalLM::logits(hidden_states, seleted_idxes, out_hidden);
+    }
+  }
+
   void load_model(std::unique_ptr<ModelLoader> loader) override {
     model_->load_model(std::move(loader));
   }

xllm/core/framework/model/model_traits.h

@@ -67,6 +67,17 @@ struct has_set_word_embedding<
     std::void_t<decltype(std::declval<T>()->set_word_embedding(
         std::declval<layer::WordEmbedding&>()))>> : std::true_type {};
 
+template <typename T, typename = void>
+struct has_logits_with_hidden : std::false_type {};
+
+template <typename T>
+struct has_logits_with_hidden<
+    T,
+    std::void_t<decltype(std::declval<T>()->logits(
+        std::declval<const torch::Tensor&>(),
+        std::declval<const torch::Tensor&>(),
+        std::declval<torch::Tensor&>()))>> : std::true_type {};
+
 template <typename T, typename = void>
 struct has_lazy_load_model : std::false_type {};
 

xllm/core/framework/parallel_state/collective_communicator.cpp

@@ -131,25 +131,6 @@ void CollectiveCommunicator::create_process_groups(
   net::parse_host_port_from_addr(master_addr, host, port);
 
 #if defined(USE_NPU)
-  // ATB backend usually relies on ATB/HCCL communication. However CP+MTP
-  // needs torch process group collectives in worker runtime.
-  if (cp_size > 1) {
-    CHECK_EQ(world_size % cp_size, 0)
-        << "world_size must be divisible by cp_size.";
-    const int cp_group_count = world_size / cp_size;
-    int port_offset = global_rank % cp_group_count + 1;
-    cp_group_ = create_process_group(global_rank,
-                                     world_size,
-                                     cp_size,
-                                     port + port_offset,
-                                     true,
-                                     host,
-                                     "cp_group",
-                                     device);
-    parallel_args_->cp_group_ = cp_group_.get();
-    port += cp_group_count;
-  }
-
   if (FLAGS_npu_kernel_backend == "ATB") {
     return;
   }

xllm/core/framework/parallel_state/collective_communicator.h

@@ -48,7 +48,6 @@ class CollectiveCommunicator {
   // aliases ParallelArgs::sp_group_ to tp_group_ instead of constructing a
   // separate communicator, so this stays empty for now.
   std::unique_ptr<ProcessGroup> sp_group_;
-  std::unique_ptr<ProcessGroup> cp_group_;
   std::unique_ptr<ProcessGroup> moe_tp_group_;
   std::unique_ptr<ProcessGroup> moe_ep_group_;
 };

xllm/core/framework/parallel_state/parallel_args.h

@@ -128,8 +128,6 @@ struct ParallelArgs {
   // the same rank set during prefill, but it remains a separate handle so the
   // SP communication policy can evolve independently from TP.
   ProcessGroup* sp_group_ = nullptr;
-  // Context-parallel communication group for CP-specific runtime collectives.
-  ProcessGroup* cp_group_ = nullptr;
   ProcessGroup* moe_ep_group_ = nullptr;
   ProcessGroup* moe_tp_group_ = nullptr;
 };

xllm/core/layers/npu/npu_lm_head_impl.cpp

@@ -28,6 +28,7 @@ void NpuLmHeadImpl::param_from_args(atb_speed::common::LmHeadParam& param,
                                     const ParallelArgs& parallel_args,
                                     bool isPrefill) {
   const bool use_column_parallel = cp_size_ > 1;
+  param.outputHidden = cp_size_ > 1;
   param.unpadInputs = true;
   param.gatherAhead = isPrefill;
   param.hiddenSizePerAttentionHead = args.hidden_size() / args.n_heads();
@@ -101,7 +102,7 @@ NpuLmHeadImpl::NpuLmHeadImpl(const ModelContext& context) : BaseLayer(context) {
                   false);
 
   atb_weight_tensors_.resize(1);
-  atOutTensors_.resize(1);
+  atOutTensors_.resize(2);
 
   auto options = context.get_tensor_options();
   dtype_ = c10::typeMetaToScalarType(options.dtype());
@@ -146,26 +147,38 @@ int64_t NpuLmHeadImpl::init_node(atb_speed::Model::Node& node,
     return -1;
   }
   node.inTensors.resize(node.operation->GetInputNum());
-  node.outTensors.resize(1);
+  node.outTensors.resize(node.operation->GetOutputNum());
 
   node.inTensors.at(1) = &atb_weight_tensors_[0];
 
   node.variantPack.inTensors.reserve(node.inTensors.size());
   node.variantPack.inTensors.resize(node.inTensors.size());
-  node.variantPack.outTensors.reserve(1);
-  node.variantPack.outTensors.resize(1);
+  node.variantPack.outTensors.reserve(node.outTensors.size());
+  node.variantPack.outTensors.resize(node.outTensors.size());
 
   return atb::NO_ERROR;
 }
 
 torch::Tensor NpuLmHeadImpl::forward(const torch::Tensor& hidden_states,
                                      const torch::Tensor& seleted_idxes,
                                      int nodeId) {
+  torch::Tensor out_hidden;
+  return forward_with_hidden(hidden_states, seleted_idxes, out_hidden, nodeId);
+}
+
+torch::Tensor NpuLmHeadImpl::forward_with_hidden(
+    const torch::Tensor& hidden_states,
+    const torch::Tensor& seleted_idxes,
+    torch::Tensor& out_hidden,
+    int nodeId) {
   atb::Status st;
   build_node_variant_pack(lm_head_node_prefill_, hidden_states, seleted_idxes);
   st = execute_node(lm_head_node_prefill_, nodeId);
   LOG_IF(FATAL, st != 0) << model_name_
                          << "execute lmhead node fail, error code: " << st;
+  if (atOutTensors_.size() > 1) {
+    out_hidden = atOutTensors_[1];
+  }
   return atOutTensors_[0];
 }
 
@@ -212,12 +225,16 @@ void NpuLmHeadImpl::build_node_variant_pack(
   inTensorDescs.at(8) = placeholder_.desc;
 
   atb::Status st = node.operation->InferShape(inTensorDescs, outTensorDescs);
-  at::Tensor newTensor =
-      atb_speed::Utils::CreateAtTensorFromTensorDesc(outTensorDescs.at(0));
-
-  atOutTensors_.at(0) = newTensor;
-  node.variantPack.outTensors.at(0) =
-      atb_speed::Utils::AtTensor2Tensor(atOutTensors_.at(0));
+  LOG_IF(FATAL, st != atb::NO_ERROR)
+      << model_name_ << " infer lmhead shape fail, error code: " << st;
+
+  atOutTensors_.resize(node.variantPack.outTensors.size());
+  for (size_t i = 0; i < node.variantPack.outTensors.size(); ++i) {
+    atOutTensors_.at(i) =
+        atb_speed::Utils::CreateAtTensorFromTensorDesc(outTensorDescs.at(i));
+    node.variantPack.outTensors.at(i) =
+        atb_speed::Utils::AtTensor2Tensor(atOutTensors_.at(i));
+  }
 }
 
 }  // namespace layer

xllm/core/layers/npu/npu_lm_head_impl.h

@@ -53,6 +53,11 @@ class NpuLmHeadImpl : public BaseLayer {
                         const torch::Tensor& seleted_idxes,
                         int nodeId);
 
+  torch::Tensor forward_with_hidden(const torch::Tensor& hidden_states,
+                                    const torch::Tensor& seleted_idxes,
+                                    torch::Tensor& out_hidden,
+                                    int nodeId);
+
  private:
   void param_from_args(atb_speed::common::LmHeadParam& param,
                        const ModelArgs& args,

xllm/core/runtime/llm_worker_impl.cpp

@@ -31,7 +31,6 @@ limitations under the License.
 #include "core/common/global_flags.h"
 #include "framework/kv_cache/kv_cache.h"
 #include "framework/model/model_input_params.h"
-#include "framework/parallel_state/parallel_state.h"
 #include "framework/state_dict/state_dict.h"
 #if defined(USE_CUDA) || defined(USE_ILU) || defined(USE_MUSA)
 #include "layers/cuda/flashinfer_workspace.h"
@@ -125,9 +124,16 @@ std::optional<ForwardOutput> LLMWorkerImpl::step_internal(
   }
 
   torch::Tensor logits;
+  torch::Tensor selected_hidden_from_lm_head;
   if (sampling_params.selected_token_idxes.defined()) {
-    logits = model_->logits(model_output.hidden_states,
-                            sampling_params.selected_token_idxes);
+    if (options_.cp_size() > 1) {
+      logits = model_->logits(model_output.hidden_states,
+                              sampling_params.selected_token_idxes,
+                              selected_hidden_from_lm_head);
+    } else {
+      logits = model_->logits(model_output.hidden_states,
+                              sampling_params.selected_token_idxes);
+    }
   }
 
   ForwardOutput output;
@@ -199,16 +205,15 @@ std::optional<ForwardOutput> LLMWorkerImpl::step_internal(
     if (!input.input_params.batch_forward_type.is_decode() && !is_spec_draft_) {
       output.sample_output.embeddings = embeddings;
     } else if (sampling_params.selected_token_idxes.defined()) {
-      if (context_.get_parallel_args().cp_size() > 1) {
-        auto* cp_group = context_.get_parallel_args().cp_group_;
-        CHECK(cp_group != nullptr)
-            << "cp_group_ must be initialized for CP+MTP embeddings gather.";
-        CHECK_EQ(cp_group->world_size(), context_.get_parallel_args().cp_size())
-            << "cp_group_ world size mismatch with cp_size.";
-        embeddings = parallel_state::gather(embeddings, cp_group, /*dim=*/0);
-      }
-      output.sample_output.embeddings = embeddings.index_select(
+      if (options_.cp_size() > 1) {
+        CHECK(selected_hidden_from_lm_head.defined())
+            << "selected_hidden_from_lm_head must be defined when "
+              "selected_token_idxes is defined.";
+        output.sample_output.embeddings = selected_hidden_from_lm_head;
+      } else {
+        output.sample_output.embeddings = embeddings.index_select(
           /*dim=*/0, sampling_params.selected_token_idxes);
+      }
     }
   }
 

xllm/core/runtime/mtp_worker_impl.cpp

@@ -323,7 +323,6 @@ std::optional<ForwardOutput> MTPWorkerImpl::step_prefill(
     auto mask = (token_ids == -1);
     token_ids.masked_scatter_(mask, next_tokens);
   }
-
   // generate kv cache for draft model
   timer.reset();
   auto draft_future = draft_impl_->step_async(prefill_input);
@@ -360,7 +359,8 @@ void MTPWorkerImpl::prepare_prefill_inputs(const ForwardInput& input,
   }
 
   auto& extra_token_ids = input_params.extra_token_ids;
-  CHECK_EQ(extra_token_ids.size(), static_cast<size_t>(input_params.num_sequences))
+  CHECK_EQ(extra_token_ids.size(),
+           static_cast<size_t>(input_params.num_sequences))
       << "extra_token_ids size should match num_sequences";
 
   torch::Tensor token_ids = safe_to(input.token_ids, torch::kCPU);

xllm/models/llm/npu/llm_model_base.h

@@ -420,6 +420,17 @@ class LlmForCausalLMImplBase : public torch::nn::Module {
     return npu_lm_head_(hidden_states, seleted_idxes, 0);
   }
 
+  // hidden_states: [num_tokens, hidden_size]
+  // seleted_idxes: [num_tokens]
+  // out_hidden: [num_seqs, hidden_size]
+  // returns: [num_tokens, vocab_size]
+  virtual torch::Tensor logits(const torch::Tensor& hidden_states,
+                               const torch::Tensor& seleted_idxes,
+                               torch::Tensor& out_hidden) {
+    return npu_lm_head_->forward_with_hidden(
+        hidden_states, seleted_idxes, out_hidden, 0);
+  }
+
   // hidden_states: [num_tokens, hidden_size]
   // seleted_idxes: [num_tokens]
   // returns: [num_seqs, hidden_size]