Method prefill_npu

cactus/engine/engine_model.cpp:925–997

}

void Model::prefill_npu(const std::vector<uint32_t>& tokens) {
    if (!npu_prefill_ || !npu_prefill_->is_available()) {
        throw std::runtime_error("NPU prefill not available");
    }

    const int chunk_size = npu_prefill_->get_chunk_size();
    const int hidden_dim = npu_prefill_->get_hidden_dim();
    const int num_layers = npu_prefill_->get_num_layers();
    const int fallback_num_kv_heads = npu_prefill_->get_num_kv_heads();
    const int fallback_head_dim = npu_prefill_->get_head_dim();

    const std::vector<size_t> layer_dims = get_kv_layer_dims();
    const std::vector<size_t> layer_heads = get_kv_layer_heads();
    const int layers_to_update = std::min<int>(num_layers, static_cast<int>(config_.num_layers));

    std::vector<__fp16> all_embeddings = get_token_embeddings(tokens);
    if (all_embeddings.empty()) {
        throw std::runtime_error("Failed to get token embeddings for NPU prefill");
    }

    if (Config::is_gemma_family(config_.model_type)) {
        float scale = std::sqrt(static_cast<float>(hidden_dim));
        for (size_t i = 0; i < all_embeddings.size(); i++) {
            all_embeddings[i] = __fp16(static_cast<float>(all_embeddings[i]) * scale);
        }
    }

    size_t num_tokens = tokens.size();
    size_t num_chunks = (num_tokens + chunk_size - 1) / chunk_size;

    for (size_t c = 0; c < num_chunks; c++) {
        size_t start = c * chunk_size;
        size_t actual_tokens = std::min(static_cast<size_t>(chunk_size), num_tokens - start);

        std::vector<__fp16> chunk_embeddings(chunk_size * hidden_dim, __fp16(0));
        std::copy(all_embeddings.begin() + start * hidden_dim,
                  all_embeddings.begin() + (start + actual_tokens) * hidden_dim,
                  chunk_embeddings.begin());

        int position_offset = static_cast<int>(start);

        npu::NPUPrefillDirectResult direct_result = npu_prefill_->prefill_chunk_direct(chunk_embeddings, position_offset);

        if (direct_result.valid) {
            for (int layer_idx = 0; layer_idx < layers_to_update; layer_idx++) {
                const auto& k_ref = direct_result.k_caches[layer_idx];
                const auto& v_ref = direct_result.v_caches[layer_idx];

                if (k_ref.data && v_ref.data) {
                    size_t layer_kv_heads = layer_idx < static_cast<int>(layer_heads.size())
                        ? layer_heads[layer_idx]
                        : static_cast<size_t>(fallback_num_kv_heads);
                    size_t layer_head_dim = layer_idx < static_cast<int>(layer_dims.size())
                        ? layer_dims[layer_idx]
                        : static_cast<size_t>(fallback_head_dim);

                    size_t expected = static_cast<size_t>(chunk_size) * layer_kv_heads * layer_head_dim;
                    if (expected > 0 && (k_ref.count < expected || v_ref.count < expected)) {