bugfix: preserve MTP correctness during asynchronous execution. (#1931)

Author: DongheJin

tests/core/runtime/acl_graph_executor_test.cpp

@@ -27,6 +27,7 @@ limitations under the License.
 #include "core/framework/block/block.h"
 #include "core/framework/block/block_manager_impl.h"
 #include "core/framework/config/execution_config.h"
+#include "core/framework/config/speculative_config.h"
 #include "core/framework/kv_cache/kv_cache.h"
 #include "core/framework/model/model_args.h"
 #include "core/framework/model/model_output.h"
@@ -37,6 +38,7 @@ limitations under the License.
 #include "core/layers/npu/npu_lm_head_impl.h"
 #include "core/layers/npu/npu_word_embedding_impl.h"
 #include "core/runtime/acl_graph_executor_impl.h"
+#include "core/runtime/acl_graph_persistent_param.h"
 #include "core/runtime/base_executor_impl.h"
 #include "core/runtime/options.h"
 
@@ -770,6 +772,86 @@ TEST_F(AclGraphExecutorTest, GraphDoubleBufferFlagControlsSlotCount) {
       original_enable_graph_double_buffer);
 }
 
+TEST(AclGraphPersistentParamTest, SpecVerifyMetadataUsesTokenCapacity) {
+  SpeculativeConfig& speculative_config = SpeculativeConfig::get_instance();
+  const bool original_enable_atb_spec_kernel =
+      speculative_config.enable_atb_spec_kernel();
+  speculative_config.enable_atb_spec_kernel(false);
+
+  ModelArgs args;
+  args.model_type("deepseek_v4");
+  args.dtype("float32");
+  args.hidden_size(8);
+  args.max_position_embeddings(32);
+
+  runtime::Options options;
+  options.block_size(4);
+  options.max_seqs_per_batch(10);
+  options.max_tokens_per_batch(64);
+  options.num_decoding_tokens(3);
+  options.enable_speculative_decode(true);
+  options.is_draft_engine(false);
+
+  const torch::Device device("npu:0");
+  ::xllm::npu::GraphPersistentParam persistent_param(
+      args,
+      device,
+      options,
+      /*need_update_attn_mask=*/false,
+      /*is_hybrid_linear_attention=*/false);
+  EXPECT_EQ(persistent_param.q_seq_lens().size(0), 30);
+  EXPECT_EQ(persistent_param.kv_seq_lens().size(0), 30);
+  EXPECT_EQ(persistent_param.persistent_block_tables().size(0), 30);
+
+  constexpr int64_t kValidateRows = 12;
+  const torch::TensorOptions int_options =
+      torch::dtype(torch::kInt).device(device);
+  const torch::Tensor tokens = torch::arange(kValidateRows, int_options);
+  const torch::Tensor positions = torch::arange(kValidateRows, int_options);
+  ModelInputParams params;
+  params.is_spec_verify = true;
+  params.meta.batch_forward_type = BatchForwardType::DECODE;
+  params.meta.num_sequences = kValidateRows;
+  params.attention.host.q_seq_lens.assign(kValidateRows, 1);
+  params.attention.host.kv_seq_lens.assign(kValidateRows, 8);
+  params.attention.device.q_seq_lens =
+      torch::ones({kValidateRows}, int_options);
+  params.attention.device.kv_seq_lens =
+      torch::full({kValidateRows}, 8, int_options);
+  params.attention.device.new_cache_slots =
+      torch::zeros({kValidateRows}, int_options);
+  params.attention.device.block_tables =
+      torch::zeros({kValidateRows, 2}, int_options);
+
+  std::optional<ModelInputParams> capture_params;
+  EXPECT_NO_THROW(capture_params = persistent_param.update(
+                      tokens,
+                      torch::Tensor(),
+                      torch::Tensor(),
+                      positions,
+                      params,
+                      /*padded_num_tokens=*/kValidateRows,
+                      /*return_capture_params=*/true));
+  EXPECT_TRUE(capture_params.has_value());
+  if (capture_params.has_value()) {
+    EXPECT_EQ(capture_params->attention.device.q_seq_lens.size(0),
+              kValidateRows);
+    EXPECT_EQ(capture_params->attention.device.block_tables.size(0),
+              kValidateRows);
+  }
+
+  ::xllm::npu::GraphPersistentParam hybrid_persistent_param(
+      args,
+      device,
+      options,
+      /*need_update_attn_mask=*/false,
+      /*is_hybrid_linear_attention=*/true);
+  EXPECT_EQ(hybrid_persistent_param.q_seq_lens().size(0), 10);
+  EXPECT_EQ(hybrid_persistent_param.persistent_block_tables().size(0), 10);
+
+  speculative_config.enable_atb_spec_kernel(original_enable_atb_spec_kernel);
+}
+
 TEST(AclGraphExecutorHybridTest, KvCacheSupportsLinearOnlyLayers) {
   auto conv_cache = torch::zeros({4, 32, 3}, torch::dtype(torch::kFloat32));
   auto ssm_cache = torch::zeros({4, 8, 64, 64}, torch::dtype(torch::kFloat32));

xllm/core/runtime/acl_graph_persistent_param.cpp

@@ -162,7 +162,11 @@ GraphPersistentParam::GraphPersistentParam(const ModelArgs& args,
   // only serves decode / spec-verify batches, so the relevant row upper bound
   // comes from decode graph capacity instead.
   const int64_t max_graph_tokens = get_decode_graph_token_capacity(options);
-  const int64_t max_seqs_per_batch = get_decode_graph_capacity(options);
+  // Hybrid spec verify keeps sequence-scoped metadata, while non-hybrid MTP
+  // expands every proposal token into its own decode metadata row.
+  const int64_t metadata_capacity = is_hybrid_linear_attention
+                                        ? get_decode_graph_capacity(options)
+                                        : max_graph_tokens;
 
   const int64_t max_seq_len = args_.max_position_embeddings();
 
@@ -182,18 +186,17 @@ GraphPersistentParam::GraphPersistentParam(const ModelArgs& args,
   persistent_new_cache_slots_default_ = torch::zeros(
       {max_tokens_per_batch}, torch::dtype(torch::kInt).device(device));
   persistent_linear_state_indices_ = torch::zeros(
-      {max_seqs_per_batch}, torch::dtype(torch::kInt).device(device));
+      {metadata_capacity}, torch::dtype(torch::kInt).device(device));
   persistent_num_accepted_tokens_ = torch::ones(
-      {max_seqs_per_batch}, torch::dtype(torch::kInt).device(device));
+      {metadata_capacity}, torch::dtype(torch::kInt).device(device));
 
-  // Sequence length tensors with max_seqs_per_batch
-  q_seq_lens_ = torch::zeros({max_seqs_per_batch},
+  q_seq_lens_ = torch::zeros({metadata_capacity},
                              torch::dtype(torch::kInt).device(device));
-  kv_seq_lens_ = torch::zeros({max_seqs_per_batch},
+  kv_seq_lens_ = torch::zeros({metadata_capacity},
                               torch::dtype(torch::kInt).device(device));
-  q_seq_lens_default_ = torch::ones({max_seqs_per_batch},
+  q_seq_lens_default_ = torch::ones({metadata_capacity},
                                     torch::dtype(torch::kInt).device(device));
-  kv_seq_lens_default_ = torch::ones({max_seqs_per_batch},
+  kv_seq_lens_default_ = torch::ones({metadata_capacity},
                                      torch::dtype(torch::kInt).device(device));
   expanded_kv_seq_lens_ = torch::zeros(
       {max_tokens_per_batch}, torch::dtype(torch::kInt).device(device));
@@ -203,10 +206,10 @@ GraphPersistentParam::GraphPersistentParam(const ModelArgs& args,
   const int64_t max_block_table_len =
       (max_seq_len + block_size - 1) / block_size + 1;
   persistent_block_tables_ =
-      torch::zeros({max_seqs_per_batch, max_block_table_len},
+      torch::zeros({metadata_capacity, max_block_table_len},
                    torch::dtype(torch::kInt).device(device));
   persistent_block_tables_default_ =
-      torch::zeros({max_seqs_per_batch, max_block_table_len},
+      torch::zeros({metadata_capacity, max_block_table_len},
                    torch::dtype(torch::kInt).device(device));
   persistent_expanded_block_tables_ =
       torch::zeros({max_tokens_per_batch, max_block_table_len},
@@ -243,8 +246,8 @@ GraphPersistentParam::GraphPersistentParam(const ModelArgs& args,
   }
 
   q_cu_seq_lens_default_ = torch::zeros(
-      {max_seqs_per_batch + 1}, torch::dtype(torch::kInt).device(device));
-  q_cu_seq_lens_ = torch::zeros({max_seqs_per_batch + 1},
+      {metadata_capacity + 1}, torch::dtype(torch::kInt).device(device));
+  q_cu_seq_lens_ = torch::zeros({metadata_capacity + 1},
                                 torch::dtype(torch::kInt).device(device));
 
   // Pre-allocate persistent dp/cp ep padding buffers with maximum capacity.
@@ -782,9 +785,15 @@ std::optional<ModelInputParams> GraphPersistentParam::update(
 
   if (persistent_block_tables_default_.defined() &&
       persistent_block_tables_default_.sizes() ==
-          persistent_block_tables_.sizes()) {
-    persistent_block_tables_.copy_(persistent_block_tables_default_,
-                                   /*non_blocking=*/true);
+          persistent_block_tables_.sizes() &&
+      padded_batch_size > 0) {
+    CHECK_LE(padded_batch_size, persistent_block_tables_.size(0))
+        << "padded graph metadata rows exceed persistent block table capacity";
+    persistent_block_tables_
+        .slice(/*dim=*/0, /*start=*/0, /*end=*/padded_batch_size)
+        .copy_(persistent_block_tables_default_.slice(
+                   /*dim=*/0, /*start=*/0, /*end=*/padded_batch_size),
+               /*non_blocking=*/true);
   }
   if (actual_seq_len_rows > 0 &&
       params.attention.device.block_tables.defined() &&

xllm/core/runtime/forward_params.h

@@ -908,6 +908,9 @@ struct ForwardOutput {
   // Keep no-sync input tensor handles alive until downstream consumers finish
   // using outputs on the same compute stream.
   std::shared_ptr<ForwardInput> retained_input;
+  // Composite workers retain nested no-sync inputs until their final output is
+  // synchronized or its ready event is consumed.
+  std::vector<std::shared_ptr<ForwardInput>> retained_input_dependencies;
   // Device-side readiness dependency for no-sync outputs. This local runtime
   // handle is intentionally not included in proto or shared-memory transport.
   StreamEventPtr ready_event;

xllm/core/runtime/llm_worker_impl.cpp

@@ -124,7 +124,8 @@ std::optional<ForwardOutput> LLMWorkerImpl::step_no_sync(
 
 std::optional<ForwardOutput> LLMWorkerImpl::execute_no_sync_on_stream(
     const ForwardInput& input,
-    Stream& compute_stream) {
+    Stream& compute_stream,
+    bool record_ready_event) {
   const ForwardSyncPolicy sync_policy = ForwardSyncPolicy::NO_SYNC;
   c10::StreamGuard stream_guard = compute_stream.set_stream_guard();
   if (::xllm::LoadConfig::get_instance().enable_manual_loader()) {
@@ -136,19 +137,19 @@ std::optional<ForwardOutput> LLMWorkerImpl::execute_no_sync_on_stream(
           const_cast<atb::Context*>(context_.get_atb_context());
       atb_context->SetExecuteStream(current_acl_stream);
       wait_input_ready_events(input, compute_stream);
-      return step_internal(input, sync_policy);
+      return step_internal(input, sync_policy, record_ready_event);
     } else {
       SET_ATB_EXECUTE_STREAM((&compute_stream), device_, context_);
       wait_input_ready_events(input, compute_stream);
-      return step_internal(input, sync_policy);
+      return step_internal(input, sync_policy, record_ready_event);
     }
 #else
     wait_input_ready_events(input, compute_stream);
-    return step_internal(input, sync_policy);
+    return step_internal(input, sync_policy, record_ready_event);
 #endif
   }
   wait_input_ready_events(input, compute_stream);
-  return step_internal(input, sync_policy);
+  return step_internal(input, sync_policy, record_ready_event);
 }
 
 std::optional<ForwardOutput> LLMWorkerImpl::step(const ForwardInput& input) {
@@ -213,7 +214,8 @@ LLMWorkerImpl::update_input_by_last_step_output_for_schedule_overlap(
 
 std::optional<ForwardOutput> LLMWorkerImpl::step_internal(
     const ForwardInput& input,
-    ForwardSyncPolicy sync_policy) {
+    ForwardSyncPolicy sync_policy,
+    bool record_ready_event) {
   MULTI_MODEL_STEP_LOCK(::xllm::KVCacheConfig::get_instance().enable_xtensor());
 
   Timer timer;
@@ -376,7 +378,7 @@ std::optional<ForwardOutput> LLMWorkerImpl::step_internal(
   if (sync_policy == ForwardSyncPolicy::NO_SYNC) {
     wait_kv_push();
     output.retained_input = std::make_shared<ForwardInput>(input);
-    if (enable_schedule_overlap()) {
+    if (enable_schedule_overlap() && record_ready_event) {
       output.ready_event = record_current_stream_event(device_);
     }
     return output;

xllm/core/runtime/llm_worker_impl.h

@@ -51,14 +51,16 @@ class LLMWorkerImpl : public WorkerImpl {
   std::optional<ForwardOutput> step_no_sync(const ForwardInput& input);
   std::optional<ForwardOutput> execute_no_sync_on_stream(
       const ForwardInput& input,
-      Stream& compute_stream);
+      Stream& compute_stream,
+      bool record_ready_event = true);
 
   folly::SemiFuture<std::optional<ForwardOutput>> step_async_no_sync(
       const ForwardInput& input);
 
   std::optional<ForwardOutput> step_internal(
       const ForwardInput& input,
-      ForwardSyncPolicy sync_policy = ForwardSyncPolicy::LEGACY);
+      ForwardSyncPolicy sync_policy = ForwardSyncPolicy::LEGACY,
+      bool record_ready_event = true);
 
  protected:
   std::optional<ForwardOutput> step_for_schedule_overlap(