Method build

src/framework/modules/optimizations/fast_kv_modules.cpp:37–69

const core::ModuleSchema & FastKVSetRowsModule::schema() const noexcept {
    return static_schema();
}

core::TensorValue FastKVSetRowsModule::build(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & cache,
    const core::TensorValue & row,
    const core::TensorValue & row_index) const {
    if (ctx.ggml == nullptr) {
        throw std::runtime_error("ModuleBuildContext.ggml is null");
    }
    core::validate_rank_between(cache, 4, 4, "cache");
    core::validate_shape(
        row,
        core::TensorShape::from_dims({cache.shape.dims[0], 1, cache.shape.dims[2], cache.shape.dims[3]}),
        "row");
    const int64_t batch = cache.shape.dims[0];
    if (row_index.shape.rank != 1 || (row_index.shape.dims[0] != 1 && row_index.shape.dims[0] != batch)) {
        throw std::runtime_error("FastKVSetRowsModule row_index must have shape {1} or {batch}");
    }
    const bool optimized = config_.mode == FastKVSetRowsMode::BackendViewOptimized;
    if (((!optimized && cache.type != GGML_TYPE_F32) ||
         (optimized && cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16 && cache.type != GGML_TYPE_BF16)) ||
        row.type != GGML_TYPE_F32) {
        throw std::runtime_error(
            optimized
                ? "FastKVSetRowsModule requires an f32/f16/bf16 cache and an f32 row tensor"
                : "FastKVSetRowsModule requires f32 cache and row tensors");
    }
    if (row_index.type != GGML_TYPE_I32 && row_index.type != GGML_TYPE_I64) {
        throw std::runtime_error("FastKVSetRowsModule requires i32 or i64 row_index tensor");
    }
    if (!core::has_backend_addressable_layout(cache.tensor)) {
        throw std::runtime_error("FastKVSetRowsModule requires a contiguous cache tensor");
    }

    const int64_t steps = cache.shape.dims[1];