Method backbone_cached_step

src/models/heartmula/mula.cpp:1922–1978

}

HeartMuLaBackboneResult HeartMuLaWeightsRuntime::backbone_cached_step(
    const std::vector<float> & embedding,
    int64_t batch_size,
    HeartMuLaBackboneCachedState & state,
    int64_t cache_capacity) const {
    const auto & config = assets_->mula_config.backbone;
    if (batch_size <= 0) {
        throw std::runtime_error("HeartMuLa backbone cached step requires positive batch size");
    }
    if (cache_capacity <= 0) {
        throw std::runtime_error("HeartMuLa backbone cached step requires positive cache capacity");
    }
    if (cache_capacity > config.max_seq_len) {
        throw std::runtime_error("HeartMuLa backbone cached step exceeds model context length");
    }
    if (static_cast<int64_t>(embedding.size()) != batch_size * config.embed_dim) {
        throw std::runtime_error("HeartMuLa backbone cached step embedding payload size mismatch");
    }
    const int64_t current_end = state.graph_has_state_ && backbone_cached_step_graph_ != nullptr
        ? backbone_cached_step_graph_->current_end()
        : state.pending_state_.current_end;
    const int64_t required_capacity = current_end + 1;
    if (required_capacity > config.max_seq_len) {
        throw std::runtime_error("HeartMuLa backbone cached step capacity exceeds model context length");
    }
    const int64_t prompt_steps = state.prefill_steps_;
    const int64_t old_generated_capacity = backbone_cached_step_graph_ == nullptr
        ? 0
        : std::max<int64_t>(1, backbone_cached_step_graph_->cache_steps() - prompt_steps);
    const int64_t required_generated_capacity = std::max<int64_t>(1, required_capacity - prompt_steps);
    const int64_t max_generated_capacity = std::max<int64_t>(1, cache_capacity - prompt_steps);
    const int64_t next_generated_capacity = std::min<int64_t>(
        max_generated_capacity,
        std::max<int64_t>(
            required_generated_capacity,
            old_generated_capacity > 0 ? old_generated_capacity * 2 : kInitialBackboneGeneratedCacheFrames));
    const int64_t graph_capacity = prompt_steps + next_generated_capacity;
    if (backbone_cached_step_graph_ != nullptr && state.graph_has_state_ &&
        !backbone_cached_step_graph_->can_run(*this, batch_size, required_capacity)) {
        state.pending_state_ = backbone_cached_step_graph_->export_state();
        state.graph_has_state_ = false;
    }
    if (backbone_cached_step_graph_ == nullptr ||
        !backbone_cached_step_graph_->can_run(*this, batch_size, required_capacity)) {
        backbone_cached_step_graph_.reset();
        backbone_cached_step_graph_ = std::make_unique<HeartMuLaBackboneCachedStepGraph>(
            *this,
            batch_size,
            graph_capacity);
    }
    if (!state.graph_has_state_) {
        backbone_cached_step_graph_->import_state(state.pending_state_);
        state.pending_state_ = {};
        state.graph_has_state_ = true;
    }
    return backbone_cached_step_graph_->run_step(embedding);
}