feat: add fused next draft preparation.

xllm_ops/build_aclnn.sh

@@ -150,6 +150,7 @@ elif [[ "$SOC_VERSION" =~ ^(ascend)?910b ]]; then
         "pp_matmul_opt"
         "recurrent_gated_delta_rule"
         "replace_token"
+        "mtp_prepare_next_draft"
         "select_unshared_kv"
         "x_attention_tl"
         "x_flash_attention_infer"
@@ -265,6 +266,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then
         "pp_matmul_opt"
         "recurrent_gated_delta_rule"
         "replace_token"
+        "mtp_prepare_next_draft"
         "select_unshared_kv"
         "x_attention_tl"
         "x_flash_attention_infer"

xllm_ops/mtp_prepare_next_draft/CMakeLists.txt

@@ -0,0 +1,9 @@
+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+if(NOT ENABLE_TEST AND NOT BENCHMARK)
+    list(REMOVE_ITEM CURRENT_DIRS tests)
+endif()
+foreach(SUB_DIR ${CURRENT_DIRS})
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
+        add_subdirectory(${SUB_DIR})
+    endif()
+endforeach()

xllm_ops/mtp_prepare_next_draft/op_host/CMakeLists.txt

@@ -0,0 +1,18 @@
+add_op_to_compiled_list()
+
+if(BUILD_OPEN_PROJECT)
+    target_sources(op_host_aclnn PRIVATE
+        mtp_prepare_next_draft_def.cpp
+    )
+endif()
+
+add_ops_compile_options(
+    OP_NAME MtpPrepareNextDraft
+    OPTIONS --cce-auto-sync=on
+            -Wno-deprecated-declarations
+            -Werror
+)
+
+if(NOT BUILD_OPS_RTY_KERNEL)
+    add_modules_sources(OPTYPE mtp_prepare_next_draft ACLNNTYPE aclnn)
+endif()

xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_def.cpp

@@ -0,0 +1,79 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved. */
+
+#include "register/op_def_registry.h"
+
+namespace ops {
+
+class MtpPrepareNextDraft : public OpDef {
+ public:
+  explicit MtpPrepareNextDraft(const char* name) : OpDef(name) {
+    this->Input("accepted_tokens")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT64, ge::DT_INT64})
+        .FormatList({ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("accepted_embeddings")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16})
+        .FormatList({ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("embedding_placeholder")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16})
+        .FormatList({ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("base_positions")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32})
+        .FormatList({ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("base_kv_seq_lens")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32})
+        .FormatList({ge::FORMAT_ND})
+        .AutoContiguous();
+    this->Input("block_tables")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32})
+        .FormatList({ge::FORMAT_ND})
+        .AutoContiguous();
+
+    this->Output("draft_token_ids")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32})
+        .FormatList({ge::FORMAT_ND});
+    this->Output("draft_embeddings")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_FLOAT16, ge::DT_BF16})
+        .FormatList({ge::FORMAT_ND});
+    this->Output("draft_positions")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32})
+        .FormatList({ge::FORMAT_ND});
+    this->Output("draft_kv_seq_lens")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32})
+        .FormatList({ge::FORMAT_ND});
+    this->Output("draft_cache_slots")
+        .ParamType(REQUIRED)
+        .DataType({ge::DT_INT32, ge::DT_INT32})
+        .FormatList({ge::FORMAT_ND});
+
+    this->Attr("block_size").AttrType(REQUIRED).Int();
+
+    OpAICoreConfig config;
+    config.DynamicCompileStaticFlag(true)
+        .DynamicFormatFlag(false)
+        .DynamicRankSupportFlag(false)
+        .DynamicShapeSupportFlag(true)
+        .NeedCheckSupportFlag(false)
+        .PrecisionReduceFlag(false)
+        .ExtendCfgInfo("coreType.value", "AiCore");
+    this->AICore().AddConfig("ascend910b", config);
+    this->AICore().AddConfig("ascend910_93", config);
+  }
+};
+
+OP_ADD(MtpPrepareNextDraft);
+
+}  // namespace ops

xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_proto.cpp

@@ -0,0 +1,51 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved. */
+
+#include "register/op_impl_registry.h"
+
+namespace ge {
+namespace {
+
+void Set1D(gert::Shape* shape, int64_t dim0) {
+  shape->SetDimNum(1);
+  shape->SetDim(0, dim0);
+}
+
+}  // namespace
+
+static ge::graphStatus InferShape(gert::InferShapeContext* context) {
+  const gert::Shape* token_shape = context->GetInputShape(0);
+  const gert::Shape* embedding_shape = context->GetInputShape(1);
+  if (token_shape == nullptr || embedding_shape == nullptr ||
+      token_shape->GetDimNum() != 2 || embedding_shape->GetDimNum() != 3) {
+    return GRAPH_FAILED;
+  }
+
+  const int64_t batch = token_shape->GetDim(0);
+  const int64_t doubled_batch = batch < 0 ? -1 : batch * 2;
+  const int64_t hidden = embedding_shape->GetDim(2);
+
+  Set1D(context->GetOutputShape(0), doubled_batch);
+  gert::Shape* draft_embedding_shape = context->GetOutputShape(1);
+  draft_embedding_shape->SetDimNum(2);
+  draft_embedding_shape->SetDim(0, doubled_batch);
+  draft_embedding_shape->SetDim(1, hidden);
+  Set1D(context->GetOutputShape(2), doubled_batch);
+  Set1D(context->GetOutputShape(3), batch);
+  Set1D(context->GetOutputShape(4), doubled_batch);
+  return GRAPH_SUCCESS;
+}
+
+static ge::graphStatus InferDataType(gert::InferDataTypeContext* context) {
+  context->SetOutputDataType(0, ge::DT_INT32);
+  context->SetOutputDataType(1, context->GetInputDataType(1));
+  context->SetOutputDataType(2, ge::DT_INT32);
+  context->SetOutputDataType(3, ge::DT_INT32);
+  context->SetOutputDataType(4, ge::DT_INT32);
+  return GRAPH_SUCCESS;
+}
+
+IMPL_OP_INFERSHAPE(MtpPrepareNextDraft)
+    .InferShape(InferShape)
+    .InferDataType(InferDataType);
+
+}  // namespace ge

xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.cpp

@@ -0,0 +1,54 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved. */
+
+#include "mtp_prepare_next_draft_tiling.h"
+
+#include "register/op_def_registry.h"
+
+namespace optiling {
+
+static ge::graphStatus TilingFunc(gert::TilingContext* context) {
+  const gert::StorageShape* token_shape = context->GetInputShape(0);
+  const gert::StorageShape* embedding_shape = context->GetInputShape(1);
+  const gert::StorageShape* block_table_shape = context->GetInputShape(5);
+  if (token_shape == nullptr || embedding_shape == nullptr ||
+      block_table_shape == nullptr) {
+    return ge::GRAPH_FAILED;
+  }
+
+  const gert::Shape& tokens = token_shape->GetStorageShape();
+  const gert::Shape& embeddings = embedding_shape->GetStorageShape();
+  const gert::Shape& block_tables = block_table_shape->GetStorageShape();
+  if (tokens.GetDimNum() != 2 || embeddings.GetDimNum() != 3 ||
+      block_tables.GetDimNum() != 2) {
+    return ge::GRAPH_FAILED;
+  }
+
+  const int64_t batch = tokens.GetDim(0);
+  const int64_t speculative_width = tokens.GetDim(1);
+  const int64_t hidden = embeddings.GetDim(2);
+  const int64_t num_blocks = block_tables.GetDim(1);
+  const auto* attrs = context->GetAttrs();
+  const int64_t* block_size_attr =
+      attrs == nullptr ? nullptr : attrs->GetAttrPointer<int64_t>(0);
+  if (batch <= 0 || speculative_width <= 0 || hidden <= 0 ||
+      num_blocks <= 0 || block_size_attr == nullptr ||
+      *block_size_attr <= 0) {
+    return ge::GRAPH_FAILED;
+  }
+
+  MtpPrepareNextDraftTilingData tiling;
+  tiling.set_batchSize(static_cast<uint32_t>(batch));
+  tiling.set_speculativeWidth(static_cast<uint32_t>(speculative_width));
+  tiling.set_hiddenSize(static_cast<uint32_t>(hidden));
+  tiling.set_numBlocksPerSequence(static_cast<uint32_t>(num_blocks));
+  tiling.set_blockSize(static_cast<int32_t>(*block_size_attr));
+  context->SetBlockDim(1);
+  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
+                      context->GetRawTilingData()->GetCapacity());
+  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
+  return ge::GRAPH_SUCCESS;
+}
+
+IMPL_OP_OPTILING(MtpPrepareNextDraft).Tiling(TilingFunc);
+
+}  // namespace optiling

xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.h

@@ -0,0 +1,20 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved. */
+
+#pragma once
+
+#include "register/tilingdata_base.h"
+
+namespace optiling {
+
+BEGIN_TILING_DATA_DEF(MtpPrepareNextDraftTilingData)
+TILING_DATA_FIELD_DEF(uint32_t, batchSize);
+TILING_DATA_FIELD_DEF(uint32_t, speculativeWidth);
+TILING_DATA_FIELD_DEF(uint32_t, hiddenSize);
+TILING_DATA_FIELD_DEF(uint32_t, numBlocksPerSequence);
+TILING_DATA_FIELD_DEF(int32_t, blockSize);
+END_TILING_DATA_DEF;
+
+REGISTER_TILING_DATA_CLASS(MtpPrepareNextDraft,
+                           MtpPrepareNextDraftTilingData)
+
+}  // namespace optiling

xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.cpp

@@ -0,0 +1,35 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved. */
+
+#include "kernel_operator.h"
+#include "mtp_prepare_next_draft.h"
+
+extern "C" __global__ __aicore__ void mtp_prepare_next_draft(
+    GM_ADDR acceptedTokens,
+    GM_ADDR acceptedEmbeddings,
+    GM_ADDR embeddingPlaceholder,
+    GM_ADDR basePositions,
+    GM_ADDR baseKvSeqLens,
+    GM_ADDR blockTables,
+    GM_ADDR draftTokenIds,
+    GM_ADDR draftEmbeddings,
+    GM_ADDR draftPositions,
+    GM_ADDR draftKvSeqLens,
+    GM_ADDR draftCacheSlots,
+    GM_ADDR workspace,
+    GM_ADDR tiling) {
+  GET_TILING_DATA(tiling_data, tiling);
+  MtpPrepareNextDraftKernel<DTYPE_ACCEPTED_EMBEDDINGS> op;
+  op.Init(acceptedTokens,
+          acceptedEmbeddings,
+          embeddingPlaceholder,
+          basePositions,
+          baseKvSeqLens,
+          blockTables,
+          draftTokenIds,
+          draftEmbeddings,
+          draftPositions,
+          draftKvSeqLens,
+          draftCacheSlots,
+          &tiling_data);
+  op.Process();
+}