From d8964d030279ddce432fb25442a5e032aa955d24 Mon Sep 17 00:00:00 2001 From: leejet Date: Sat, 31 Oct 2026 21:52:28 +0800 Subject: [PATCH] refactor: enable CUDA graph replay for extensions --- docs/ggml_extensions.md | 46 ++++++++++++++++++++++++++---- ggml | 2 +- src/core/cuda/convrot.cu | 2 +- src/core/cuda/cute_attention.cu | 6 ++-- src/core/cuda/i8_gated_residual.cu | 4 +-- src/core/cuda/i8_projection.cu | 2 +- src/core/cuda/int8.cu | 4 +-- src/core/cuda/rope.cu | 6 ++-- 8 files changed, 54 insertions(+), 18 deletions(-) diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 56ba9dc1..4077aebf 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -236,11 +236,47 @@ ConvRot extension leaves that chain together to preserve its existing fusion. Standalone builtin ConvRot and its native GEMM integration are still pending migration. -CUDA extension kernels currently execute between native graph fragments. Native -fragments retain CUDA Graph caching, but the extension execution plan is not -captured as a whole. Moving small fusions out of native fragments can therefore -increase host submission overhead even when kernel counts and GPU work stay the -same. Whole-plan CUDA Graph capture remains pending. +CUDA kernels can opt into device replay with `ggml_ext_register_kernel_ex` and +`GGML_EXT_KERNEL_CAPTURE_SAFE`. The original registration function retains its +per-execution callback behavior, and the descriptor ABI is unchanged. Opting in +means the callback enqueues replayable device work: it has no required +per-execution host side effects or decisions based on tensor contents, and any +private streams rejoin the borrowed stream. Lazy setup must complete during the +first ordinary execution. Tensor metadata, attributes, and host plan parameters +remain frozen until scheduler reallocation; tensor contents may change normally. + +When every selected extension opts in, the CUDA backend can capture the whole +scheduled plan, including its native fragments, into one CUDA Graph. The first +execution warms up kernels and library state while retaining native fragment +caches. The next stable execution captures +and launches the graph; later calls replay it. Changing a tensor address or the +execution stream invalidates the capture and repeats warmup, even when the graph +UID is unchanged. Capture failures propagate without retrying partially executed +work. An enclosing stream capture records the plan directly instead of starting +a nested capture. + +`GGML_EXT_KERNEL_REPLAY_REUSABLE` additionally permits reuse across scheduler +plan rebuilds and requires `CAPTURE_SAFE`. Recorded work must depend only on +complete tensor metadata, addresses, the device and immutable registration data. +Captured resources must belong to the backend or the supplied workspace, never +per-plan host state. Workspace contents must be initialized by recorded device +work on each execution. Prepare/destroy still run for every rebuilt host plan. +The backend cache retains only the executor and workspace; it compares graph +structure, tensor metadata/addresses, selected kernels and workspace requirements +exactly, without dereferencing old graph objects. A new metadata arena can reuse +an equivalent capture. Each backend retains at most 16 retired plans and 256 MiB +of their workspaces, evicting the oldest entries first. Larger workspaces retain +plan-local capture only. Backend destruction synchronizes and releases the cache. +This allows sd.cpp's per-step scheduler rebuilds to retain device replay. + +Mixed plans containing an extension without the flag retain ordinary dispatch. +Evaluation callbacks, multiple pipeline copies, and direct backend graph calls +also retain their previous behavior. `GGML_CUDA_DISABLE_GRAPHS` disables whole-plan +capture as well as native CUDA Graph caching. Native operations that cannot be +captured retain the existing fragment execution path. Whole-plan replay requires +CUDA VMM so native scratch addresses remain valid for the backend lifetime; +legacy scratch pools retain the existing fragment path. HIP and MUSA do not yet +provide this plan executor. sd.cpp's CUDA registrations opt in to reusable replay. The builtin CUDA/Vulkan INT8 GEMM, standalone ConvRot quantization, and remaining fusion rules still need to be externalized. Vulkan execution services and kernels also remain pending. diff --git a/ggml b/ggml index b34bc427..92408021 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit b34bc427955068ecd7fe080df214dc94a89e3d80 +Subproject commit 924080214ed35e9a1251ee61bc3185f105c52b41 diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu index 705c1e00..82d6554a 100644 --- a/src/core/cuda/convrot.cu +++ b/src/core/cuda/convrot.cu @@ -315,6 +315,6 @@ void sd_ggml_register_cuda_convrot(ggml_backend_dev_t device) { desc.prepare = sd_convrot::prepare; desc.execute = sd_convrot::execute; desc.destroy = sd_convrot::destroy; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } } diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu index 15f210a8..db7924db 100644 --- a/src/core/cuda/cute_attention.cu +++ b/src/core/cuda/cute_attention.cu @@ -773,16 +773,16 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) { desc.prepare = attention_prepare; desc.execute = attention_execute; desc.destroy = attention_destroy; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); desc.name = "sd.rms_rope_cute_attention.cuda.v1"; desc.op = GGML_OP_RMS_NORM; desc.priority = 300; desc.supports = nullptr; desc.preferred = nullptr; desc.optimize = nullptr; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); desc.name = "sd.i8_rope_cute_attention.cuda.v1"; desc.op = GGML_OP_MUL_MAT; desc.priority = 400; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu index 37d67792..cfcfe52f 100644 --- a/src/core/cuda/i8_gated_residual.cu +++ b/src/core/cuda/i8_gated_residual.cu @@ -456,10 +456,10 @@ void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) { desc.prepare = sd_i8_gated::prepare; desc.execute = sd_i8_gated::execute; desc.destroy = sd_i8_gated::destroy; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); desc.name = "sd.custom_i8_gated_residual.cuda.v1"; desc.op = GGML_OP_CUSTOM; desc.custom_op = sd_ggml_i8_op(); desc.optimize = nullptr; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } diff --git a/src/core/cuda/i8_projection.cu b/src/core/cuda/i8_projection.cu index 819ea751..cc2d9d3d 100644 --- a/src/core/cuda/i8_projection.cu +++ b/src/core/cuda/i8_projection.cu @@ -385,5 +385,5 @@ void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device) { desc.prepare = projection_prepare; desc.execute = projection_execute; desc.destroy = projection_destroy; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index 361dbf78..e1ca6cba 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -559,9 +559,9 @@ void sd_ggml_register_cuda_i8(ggml_backend_dev_t device) { desc.prepare = sd_i8_cuda::prepare; desc.execute = sd_i8_cuda::execute; desc.destroy = sd_i8_cuda::destroy; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); desc.name = "sd.quantize_i8_convrot.cuda.v1"; desc.custom_op = sd_ggml_convrot_op(); desc.supports = sd_i8_cuda::convrot_supports; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index 2217f19f..a0aa7f08 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -325,16 +325,16 @@ void sd_ggml_register_cuda_rope(ggml_backend_dev_t device) { desc.prepare = rope_prepare; desc.destroy = rope_destroy; desc.execute = rope_execute; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); desc.name = "sd.rope_apply.cuda.override.v1"; desc.op = GGML_OP_ROPE_APPLY; desc.custom_op = nullptr; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); desc.name = "sd.rms_norm_mul_rope.cuda.v1"; desc.op = GGML_OP_RMS_NORM; desc.priority = 250; desc.supports = nullptr; desc.match = rope_match; desc.optimize = rope_optimize; - GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); }