[Model] Removing unnecessary reshapes in get_logits (#2314)

Author: vinx13

cpp/serve/engine_actions/eagle_batch_draft.cc

@@ -116,20 +116,16 @@ class EagleBatchDraftActionObj : public EngineActionObj {
                                                                    request_internal_ids);
         NDArray logits;
         if (models_[model_id]->CanGetLogits()) {
-          logits = models_[model_id]->GetLogits(hidden_states, /*batch_size*/ num_rsentries,
-                                                /*seq_len*/ 1);
+          logits = models_[model_id]->GetLogits(hidden_states);
         } else {
           // - Use base model's head.
-          logits =
-              models_[0]->GetLogits(hidden_states, /*batch_size*/ num_rsentries, /*seq_len*/ 1);
+          logits = models_[0]->GetLogits(hidden_states);
         }
         RECORD_EVENT(trace_recorder_, request_ids, "finish proposal decode");
-        ICHECK_EQ(logits->ndim, 3);
+        ICHECK_EQ(logits->ndim, 2);
         ICHECK_EQ(logits->shape[0], num_rsentries);
-        ICHECK_EQ(logits->shape[1], 1);
 
         // - Update logits.
-        logits = logits.CreateView({num_rsentries, logits->shape[2]}, logits->dtype);
         logit_processor_->InplaceUpdateLogits(logits, generation_cfg, mstates, request_ids);
 
         // - Compute probability distributions.

cpp/serve/engine_actions/eagle_batch_verify.cc

@@ -114,16 +114,12 @@ class EagleBatchVerifyActionObj : public EngineActionObj {
     RECORD_EVENT(trace_recorder_, request_ids, "start verify");
     ObjectRef hidden_states = models_[verify_model_id_]->BatchVerifyToLastHidden(
         embeddings, request_internal_ids, verify_lengths);
-    NDArray logits =
-        models_[verify_model_id_]->GetLogits(hidden_states, 1, cum_verify_lengths[num_rsentries]);
+    NDArray logits = models_[verify_model_id_]->GetLogits(hidden_states);
     RECORD_EVENT(trace_recorder_, request_ids, "finish verify");
-    ICHECK_EQ(logits->ndim, 3);
-    ICHECK_EQ(logits->shape[0], 1);
-    ICHECK_EQ(logits->shape[1], cum_verify_lengths[num_rsentries]);
+    ICHECK_EQ(logits->ndim, 2);
+    ICHECK_EQ(logits->shape[0], cum_verify_lengths.back());
 
     // - Update logits.
-    logits =
-        logits.CreateView({cum_verify_lengths[num_rsentries], logits->shape[2]}, logits->dtype);
     logit_processor_->InplaceUpdateLogits(logits, generation_cfg, verify_request_mstates,
                                           request_ids, &cum_verify_lengths, &draft_output_tokens);
 
@@ -273,19 +269,16 @@ class EagleBatchVerifyActionObj : public EngineActionObj {
           fused_embedding_hidden_states, request_internal_ids);
 
       if (models_[draft_model_id_]->CanGetLogits()) {
-        logits = models_[draft_model_id_]->GetLogits(hidden_states, /*batch_size*/ num_rsentries,
-                                                     /*seq_len*/ 1);
+        logits = models_[draft_model_id_]->GetLogits(hidden_states);
       } else {
         // - Use base model's head.
-        logits = models_[0]->GetLogits(hidden_states, /*batch_size*/ num_rsentries, /*seq_len*/ 1);
+        logits = models_[0]->GetLogits(hidden_states);
       }
       RECORD_EVENT(trace_recorder_, request_ids, "finish proposal decode");
-      ICHECK_EQ(logits->ndim, 3);
+      ICHECK_EQ(logits->ndim, 2);
       ICHECK_EQ(logits->shape[0], num_rsentries);
-      ICHECK_EQ(logits->shape[1], 1);
 
       // - Update logits.
-      logits = logits.CreateView({num_rsentries, logits->shape[2]}, logits->dtype);
       logit_processor_->InplaceUpdateLogits(logits, generation_cfg, mstates, request_ids);
 
       // - Compute probability distributions.

cpp/serve/engine_actions/eagle_new_request_prefill.cc

@@ -183,8 +183,7 @@ class EagleNewRequestPrefillActionObj : public EngineActionObj {
       hidden_states_for_sample = models_[sample_model_id]->GatherHiddenStates(
           hidden_states, logit_positions, &model_workspaces_[model_id].hidden_states);
       // logits_for_sample: (b * s, v)
-      logits_for_sample =
-          models_[sample_model_id]->GetLogits(hidden_states_for_sample, 1, num_rsentries);
+      logits_for_sample = models_[sample_model_id]->GetLogits(hidden_states_for_sample);
       // - Update logits.
       ICHECK(logits_for_sample.defined());
       Array<GenerationConfig> generation_cfg;
@@ -195,8 +194,6 @@ class EagleNewRequestPrefillActionObj : public EngineActionObj {
         generation_cfg.push_back(prefill_inputs[i].rsentry->request->generation_cfg);
         mstates_for_logitproc.push_back(prefill_inputs[i].rsentry->mstates[sample_model_id]);
       }
-      logits_for_sample = logits_for_sample.CreateView({num_rsentries, logits_for_sample->shape[2]},
-                                                       logits_for_sample->dtype);
       logit_processor_->InplaceUpdateLogits(logits_for_sample, generation_cfg,
                                             mstates_for_logitproc, request_ids);
 

cpp/serve/model.cc

@@ -125,7 +125,7 @@ class ModelImpl : public ModelObj {
     return ft_.get_logits_func_.defined() && ft_.batch_get_logits_func_.defined();
   }
 
-  NDArray GetLogits(const ObjectRef& hidden_states, int batch_size, int seq_len) final {
+  NDArray GetLogits(const ObjectRef& hidden_states) final {
     NVTXScopedRange nvtx_scope("GetLogits");
     CHECK(ft_.get_logits_func_.defined()) << "`get_logits` function is not found in the model.";
 
@@ -139,18 +139,14 @@ class ModelImpl : public ModelObj {
     if (trace_enabled_) {
       TVMSynchronize(device_.device_type, device_.device_id, nullptr);
     }
-
     NDArray logits{nullptr};
     if (ft_.use_disco) {
       logits = Downcast<DRef>(ret)->DebugGetFromRemote(0);
     } else {
       logits = Downcast<NDArray>(ret);
     }
-    CHECK(logits.defined());
     // logits: (b * s, v)
-    ICHECK_EQ(logits->ndim, 2);
-    ICHECK_EQ(logits->shape[0], batch_size * seq_len);
-    return logits.CreateView({batch_size, seq_len, logits->shape[1]}, logits->dtype);
+    return logits;
   }
 
   ObjectRef FuseEmbedHidden(const ObjectRef& embeddings, const ObjectRef& previous_hidden_states,

cpp/serve/model.h

@@ -135,11 +135,9 @@ class ModelObj : public Object {
   /*!
    * \brief Compute logits for last hidden_states.
    * \param last_hidden_states The last hidden_states to compute logits for.
-   * \param batch_size The batch size of last_hidden_states
-   * \param seq_len The length of tokens in last_hidden_states
    * \return The computed logits.
    */
-  virtual NDArray GetLogits(const ObjectRef& last_hidden_states, int batch_size, int seq_len) = 0;
+  virtual NDArray GetLogits(const ObjectRef& last_hidden_states) = 0;
 
   /*!
    * \brief Batch prefill function. Embedding in, logits out.

python/mlc_llm/model/llama/llama_model.py

@@ -352,7 +352,7 @@ def get_default_spec(self):
                 },
             },
             "get_logits": {
-                "hidden_states": nn.spec.Tensor(["batch_size", self.hidden_size], self.dtype),
+                "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype),
                 "$": {
                     "param_mode": "packed",
                     "effect_mode": "none",