perf: accelerate INT8 diffusion inference with CUDA fusions (#2235)

This commit is contained in:
leejet authored and GitHub committed 2026-10-28 22:02:46 +08:00
1 parent 8cbe905fdd
commit 77b72a2efe
12 files changed
+113 -128

No files matched your search

+16 -21
View File
@@ -52,8 +52,13 @@ void GGMLRunner::free_params_ctx() {
void GGMLRunner::alloc_compute_ctx() {
ggml_init_params params;
params.mem_size = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
params.mem_buffer = nullptr;
params.mem_size = static_cast<size_t>(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead());
if (compute_ctx_buffer_ == nullptr) {
const size_t count = (params.mem_size + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t);
compute_ctx_buffer_.reset(new std::max_align_t[count]);
}
// Keep tensor addresses stable between steps so CUDA graphs can finish warming up.
params.mem_buffer = compute_ctx_buffer_.get();
params.no_alloc = true;
compute_ctx = ggml_init(params);
@@ -482,6 +487,8 @@ void GGMLRunner::runner_end() {
return;
}
workspace_.release();
free_compute_ctx();
compute_ctx_buffer_.reset();
cache_.clear();
logged_compute_bytes_.clear();
logged_segment_count_ = 0;
@@ -705,7 +712,7 @@ static size_t add_bytes(size_t a, size_t b) {
return b > SIZE_MAX - a ? SIZE_MAX : a + b;
}
ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) {
ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph) {
auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
if (!params_tensor_set_.count(tensor)) {
return nullptr;
@@ -716,7 +723,7 @@ ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t dir
auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
pin_multi_device_nodes(scheduler, copy, graph);
};
return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes);
return workspace_.measure(graph, external_backend, assign_nodes);
}
std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(
@@ -767,18 +774,9 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads);
}
auto scheduler = workspace_.scheduler();
ggml_status status;
if (scheduler != nullptr) {
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str());
multi_device_eval_callback_warned = true;
}
status = ggml_backend_sched_graph_compute(scheduler, graph);
} else {
status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
}
GGML_ASSERT(scheduler != nullptr);
ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
const auto status = ggml_backend_sched_graph_compute(scheduler, graph);
workspace_.synchronize();
if (status != GGML_STATUS_SUCCESS) {
last_compute_status_ = status;
@@ -832,7 +830,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
}
const auto params = collect_used_param_tensors(graph);
const auto& cached_plan = resolve_graph_cut_plan(graph);
const auto full_measurement = measure(graph, cached_plan.compute_buffer_size);
const auto full_measurement = measure(graph);
if (full_measurement.buffers.empty()) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return std::nullopt;
@@ -924,7 +922,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
}
} segment_cleanup{*this, weights, bindings, segment_context};
auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement;
auto measurement = segmented ? measure(segment_graph) : full_measurement;
if (!workspace_.prepare(measurement)) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return fail_segment("workspace preparation");
@@ -951,9 +949,6 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
return fail_segment("weight preparation");
}
// Preparing weights can execute LoRA graphs and reclaim an idle workspace.
if (!workspace_.measurement_matches(segment_graph, measurement)) {
measurement = measure(segment_graph, segment.compute_buffer_size);
}
if (!workspace_.prepare(measurement)) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return fail_segment("workspace preparation");