Method load

transformers/llm/engine/src/speculative_decoding/dflash.cpp:26–85

}

void DFlashGeneration::load(Module::Config module_config) {
    // Check if separate lm_head model exists
    std::string lmheadPath = mLlm->mConfig->dflash_lmhead();
    // Guard against empty config value: base_dir + "" = directory path, which MNNFileExist would match
    bool hasSeparateLmHead =
        !lmheadPath.empty() && MNNFileExist(lmheadPath.c_str()) && !MNNDirExist(lmheadPath.c_str());

    // Load dflash main module
    std::vector<std::string> dflashInputNames{"noise_embedding", "context_hidden", "attention_mask", "q_position_ids", "k_position_ids"};
    std::vector<std::string> dflashOutputNames{hasSeparateLmHead ? "hidden_states" : "logits"};
    mDFlashModule.reset(Module::load(
        dflashInputNames, dflashOutputNames,
        mLlm->mConfig->dflash_model().c_str(),
        mLlm->mRuntimeManager, &module_config));

    // Load separate lm_head module if available (allows fp16 transformer + int4 lm_head)
    if (hasSeparateLmHead) {
        std::vector<std::string> lmInputNames{"hidden_states"};
        std::vector<std::string> lmOutputNames{"logits"};
        mLmHeadModule.reset(Module::load(
            lmInputNames, lmOutputNames,
            lmheadPath.c_str(),
            mLlm->mRuntimeManager, &module_config));
        MNN_PRINT("DFlash: loaded separate lm_head from %s\n", lmheadPath.c_str());
    }

    // Load fc module with dedicated CPU runtime to ensure fp32 precision
    // The fc layer has very high input dimension (num_layers * hidden_size),
    // which can cause NaN in fp16 dot products during prefill
    {
        // Create a dedicated CPU runtime for FC to guarantee fp32 execution
        ScheduleConfig fc_schedule;
        fc_schedule.type = MNN_FORWARD_CPU;
        fc_schedule.numThread = 4;
        BackendConfig fc_backend_config;
        fc_backend_config.precision = BackendConfig::Precision_High;
        fc_schedule.backendConfig = &fc_backend_config;
        mFcRuntimeManager.reset(Executor::RuntimeManager::createRuntimeManager(fc_schedule));
        mFcRuntimeManager->setHint(Interpreter::MEM_ALLOCATOR_TYPE, 0);

        Module::Config fc_config;
        fc_config.shapeMutable = true;
        fc_config.rearrange = true;
        std::vector<std::string> fcInputNames{"target_hidden"};
        std::vector<std::string> fcOutputNames{"context_hidden"};
        mFcModule.reset(Module::load(
            fcInputNames, fcOutputNames,
            mLlm->mConfig->dflash_fc().c_str(),
            mFcRuntimeManager, &fc_config));
        MNN_PRINT("DFlash: FC module loaded with dedicated CPU runtime (fp32)\n");
    }

    mHiddenStateIndex = mLlm->getOutputIndex("hidden_states");

    // Disable thinking mode for better draft acceptance rate.
    // Qwen3's chat template enables thinking by default, generating unpredictable
    // <think>...</think> tokens that the draft model cannot predict well.
    // Setting enable_thinking=false via jinja context skips the <think> prefix.