Method load

transformers/llm/engine/src/speculative_decoding/mtp.cpp:18–38

}
    
void MtpGeneration::load(Module::Config module_config) {
    mMtpMeta.reset(new KVMeta);
    mLlm->mRuntimeManager->setHintPtr(Interpreter::KVCACHE_INFO, mMtpMeta.get());
    
    mMtpModules.resize(1);
    auto mtp_path = mLlm->mConfig->mtp_model();
    
    std::vector<std::string> inputNames{"input_embed", "hidden_states", "attention_mask", \
        "position_ids", "logits_index"};
    std::vector<std::string> outputNames {"logits"};
    mMtpModules[0].reset(Module::load(inputNames, outputNames, mtp_path.c_str(), mLlm->mRuntimeManager, &module_config));
    
    int verify_length = mLlm->mDraftLength + 1;
    // speculative decode module
    for(int i = 1; i <= verify_length; i++) {
        mMtpModulePool[std::make_pair(i, true)].reset(Module::clone(mMtpModules[0].get()));
    }
    // prefill module
    mMtpModulePool[std::make_pair(mLlm->mPrefillKey, false)] = mMtpModules[0];
    mHiddenStateIndex = mLlm->getOutputIndex("hidden_states");
}

std::vector<VARP> MtpGeneration::mtpForward(const std::vector<int>& input_ids, VARP hidden_states) {
    CHECK_LLM_RUNNING_RET(mContext, std::vector<VARP>());