From 7d7e29fe5cf25c5277e24a18aa134d4ae8c57592 Mon Sep 17 00:00:00 2001 From: Super User Date: Fri, 17 Jul 2026 15:18:26 +0800 Subject: [PATCH] perf(mtp): add fused next draft preparation --- xllm_ops/build_aclnn.sh | 2 + .../mtp_prepare_next_draft/CMakeLists.txt | 9 + .../op_host/CMakeLists.txt | 18 ++ .../op_host/mtp_prepare_next_draft_def.cpp | 79 +++++++ .../op_host/mtp_prepare_next_draft_proto.cpp | 51 +++++ .../op_host/mtp_prepare_next_draft_tiling.cpp | 54 +++++ .../op_host/mtp_prepare_next_draft_tiling.h | 20 ++ .../op_kernel/mtp_prepare_next_draft.cpp | 35 ++++ .../op_kernel/mtp_prepare_next_draft.h | 194 ++++++++++++++++++ 9 files changed, 462 insertions(+) create mode 100644 xllm_ops/mtp_prepare_next_draft/CMakeLists.txt create mode 100644 xllm_ops/mtp_prepare_next_draft/op_host/CMakeLists.txt create mode 100644 xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_def.cpp create mode 100644 xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_proto.cpp create mode 100644 xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.cpp create mode 100644 xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.h create mode 100644 xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.cpp create mode 100644 xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.h diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index 30cddda..2f3fb9f 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -146,6 +146,7 @@ elif [[ "$SOC_VERSION" =~ ^(ascend)?910b ]]; then "multi_latent_attention" "pp_matmul_opt" "replace_token" + "mtp_prepare_next_draft" "select_unshared_kv" "x_attention_tl" "x_flash_attention_infer" @@ -256,6 +257,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then "multi_latent_attention" "pp_matmul_opt" "replace_token" + "mtp_prepare_next_draft" "select_unshared_kv" "x_attention_tl" "x_flash_attention_infer" diff --git a/xllm_ops/mtp_prepare_next_draft/CMakeLists.txt b/xllm_ops/mtp_prepare_next_draft/CMakeLists.txt new file mode 100644 index 0000000..93f0082 --- /dev/null +++ b/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() diff --git a/xllm_ops/mtp_prepare_next_draft/op_host/CMakeLists.txt b/xllm_ops/mtp_prepare_next_draft/op_host/CMakeLists.txt new file mode 100644 index 0000000..29797ed --- /dev/null +++ b/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() diff --git a/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_def.cpp b/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_def.cpp new file mode 100644 index 0000000..42089b5 --- /dev/null +++ b/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 diff --git a/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_proto.cpp b/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_proto.cpp new file mode 100644 index 0000000..39a618f --- /dev/null +++ b/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 diff --git a/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.cpp b/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.cpp new file mode 100644 index 0000000..98d6f62 --- /dev/null +++ b/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(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(batch)); + tiling.set_speculativeWidth(static_cast(speculative_width)); + tiling.set_hiddenSize(static_cast(hidden)); + tiling.set_numBlocksPerSequence(static_cast(num_blocks)); + tiling.set_blockSize(static_cast(*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 diff --git a/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.h b/xllm_ops/mtp_prepare_next_draft/op_host/mtp_prepare_next_draft_tiling.h new file mode 100644 index 0000000..4b9c3ae --- /dev/null +++ b/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 diff --git a/xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.cpp b/xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.cpp new file mode 100644 index 0000000..fe63d1a --- /dev/null +++ b/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 op; + op.Init(acceptedTokens, + acceptedEmbeddings, + embeddingPlaceholder, + basePositions, + baseKvSeqLens, + blockTables, + draftTokenIds, + draftEmbeddings, + draftPositions, + draftKvSeqLens, + draftCacheSlots, + &tiling_data); + op.Process(); +} diff --git a/xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.h b/xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.h new file mode 100644 index 0000000..caa85d2 --- /dev/null +++ b/xllm_ops/mtp_prepare_next_draft/op_kernel/mtp_prepare_next_draft.h @@ -0,0 +1,194 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. */ + +#pragma once + +#include "kernel_operator.h" + +using namespace AscendC; + +template +class MtpPrepareNextDraftKernel { + public: + __aicore__ inline MtpPrepareNextDraftKernel() = default; + + __aicore__ inline void Init( + GM_ADDR accepted_tokens, + GM_ADDR accepted_embeddings, + GM_ADDR embedding_placeholder, + GM_ADDR base_positions, + GM_ADDR base_kv_seq_lens, + GM_ADDR block_tables, + GM_ADDR draft_token_ids, + GM_ADDR draft_embeddings, + GM_ADDR draft_positions, + GM_ADDR draft_kv_seq_lens, + GM_ADDR draft_cache_slots, + const MtpPrepareNextDraftTilingData* tiling) { + tiling_ = *tiling; + accepted_tokens_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int64_t*>(accepted_tokens)); + accepted_embeddings_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ T*>(accepted_embeddings)); + embedding_placeholder_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ T*>(embedding_placeholder)); + base_positions_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(base_positions)); + base_kv_seq_lens_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(base_kv_seq_lens)); + block_tables_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(block_tables)); + draft_token_ids_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(draft_token_ids)); + draft_embeddings_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ T*>(draft_embeddings)); + draft_positions_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(draft_positions)); + draft_kv_seq_lens_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(draft_kv_seq_lens)); + draft_cache_slots_gm_.SetGlobalBuffer( + reinterpret_cast<__gm__ int32_t*>(draft_cache_slots)); + + const uint32_t embedding_bytes = + ((tiling_.hiddenSize * sizeof(T) + 31U) / 32U) * 32U; + pipe_.InitBuffer(embedding_buffer_, embedding_bytes); + } + + __aicore__ inline void Process() { + if (GetBlockIdx() != 0) { + return; + } + + mte2_to_mte3_event_ = + GetTPipePtr()->FetchEventID(HardEvent::MTE2_MTE3); + mte3_to_mte2_event_ = + GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2); + + // Scalar SetValue stores from different AIV cores can target the same + // 32-byte GM cache line and overwrite adjacent rows. This preparation is + // launch-bound and the batch is small, so one core processes every row to + // keep the five compact outputs deterministic while retaining one fused + // launch. + for (row_ = 0; row_ < tiling_.batchSize; ++row_) { + ProcessRow(); + } + + GetTPipePtr()->ReleaseEventID( + mte2_to_mte3_event_); + GetTPipePtr()->ReleaseEventID( + mte3_to_mte2_event_); + } + + private: + __aicore__ inline void ProcessRow() { + + const uint32_t token_row_offset = row_ * tiling_.speculativeWidth; + int64_t accepted_length = 0; + for (uint32_t i = 0; i < tiling_.speculativeWidth; ++i) { + if (accepted_tokens_gm_.GetValue(token_row_offset + i) >= 0) { + ++accepted_length; + } + } + + const uint32_t last_index = + accepted_length > 0 ? static_cast(accepted_length - 1) : 0U; + const uint32_t previous_index = + accepted_length > 1 ? static_cast(accepted_length - 2) : 0U; + const int64_t last_token = + accepted_tokens_gm_.GetValue(token_row_offset + last_index); + const int64_t previous_token = + accepted_length > 1 + ? accepted_tokens_gm_.GetValue(token_row_offset + previous_index) + : last_token; + + const uint32_t output_row = row_ * 2U; + draft_token_ids_gm_.SetValue(output_row, + static_cast(previous_token)); + draft_token_ids_gm_.SetValue(output_row + 1U, + static_cast(last_token)); + + const int32_t base_position = + base_positions_gm_.GetValue(row_) + + static_cast(accepted_length); + const int32_t previous_position = base_position - 1; + draft_positions_gm_.SetValue(output_row, previous_position); + draft_positions_gm_.SetValue(output_row + 1U, base_position); + draft_kv_seq_lens_gm_.SetValue( + row_, + base_kv_seq_lens_gm_.GetValue(row_) + + static_cast(accepted_length)); + const int32_t previous_cache_position = + accepted_length == tiling_.speculativeWidth + ? previous_position + : base_position + 1; + draft_cache_slots_gm_.SetValue( + output_row, PositionToCacheSlot(previous_cache_position)); + draft_cache_slots_gm_.SetValue(output_row + 1U, + PositionToCacheSlot(base_position)); + + const uint32_t last_embedding_offset = + (token_row_offset + last_index) * tiling_.hiddenSize; + const uint32_t previous_embedding_offset = + (token_row_offset + previous_index) * tiling_.hiddenSize; + if (accepted_length > 1) { + CopyEmbedding(accepted_embeddings_gm_[previous_embedding_offset], + draft_embeddings_gm_[output_row * tiling_.hiddenSize]); + } else { + CopyEmbedding(embedding_placeholder_gm_, + draft_embeddings_gm_[output_row * tiling_.hiddenSize]); + } + CopyEmbedding(accepted_embeddings_gm_[last_embedding_offset], + draft_embeddings_gm_[(output_row + 1U) * + tiling_.hiddenSize]); + } + + __aicore__ inline int32_t PositionToCacheSlot(int32_t position) const { + // Graph warmup and batch-transition templates can contain placeholder + // positions before a real KV block has been assigned. Do not pass an + // invalid slot to reshape_and_cache; steady-state decode positions and + // block ids are non-negative, so this guard does not alter real requests. + if (position < 0) { + return 0; + } + const int32_t block_index = position / tiling_.blockSize; + if (block_index < 0 || + block_index >= static_cast(tiling_.numBlocksPerSequence)) { + return 0; + } + const int32_t block_id = block_tables_gm_.GetValue( + row_ * tiling_.numBlocksPerSequence + + static_cast(block_index)); + if (block_id < 0) { + return 0; + } + return block_id * tiling_.blockSize + position % tiling_.blockSize; + } + + __aicore__ inline void CopyEmbedding(GlobalTensor source, + GlobalTensor destination) { + LocalTensor local = embedding_buffer_.Get(); + DataCopy(local, source, tiling_.hiddenSize); + SetFlag(mte2_to_mte3_event_); + WaitFlag(mte2_to_mte3_event_); + DataCopy(destination, local, tiling_.hiddenSize); + SetFlag(mte3_to_mte2_event_); + WaitFlag(mte3_to_mte2_event_); + } + + MtpPrepareNextDraftTilingData tiling_; + uint32_t row_ = 0; + TPipe pipe_; + TBuf embedding_buffer_; + TEventID mte2_to_mte3_event_ = 0; + TEventID mte3_to_mte2_event_ = 0; + GlobalTensor accepted_tokens_gm_; + GlobalTensor accepted_embeddings_gm_; + GlobalTensor embedding_placeholder_gm_; + GlobalTensor base_positions_gm_; + GlobalTensor base_kv_seq_lens_gm_; + GlobalTensor block_tables_gm_; + GlobalTensor draft_token_ids_gm_; + GlobalTensor draft_embeddings_gm_; + GlobalTensor draft_positions_gm_; + GlobalTensor draft_kv_seq_lens_gm_; + GlobalTensor draft_cache_slots_gm_; +};