refactor: enable CUDA graph replay for extensions

This commit is contained in:
leejet committed 2026-10-31 21:52:28 +08:00
1 parent ba2188624e
commit d8964d0302
8 files changed
+54 -18

No files matched your search

+1 -1
View File
@@ -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));
}
}
+3 -3
View File
@@ -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));
}
+2 -2
View File
@@ -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));
}
+1 -1
View File
@@ -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));
}
+2 -2
View File
@@ -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));
}
+3 -3
View File
@@ -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));
}