refactor: use custom RoPE Apply nodes exclusively

This commit is contained in:
leejet committed 2026-11-01 01:11:17 +08:00
1 parent d07d48d466
commit 051b49751e
8 files changed
+43 -55

No files matched your search

+7 -11
View File
@@ -23,7 +23,7 @@ namespace sd_cpu {
const auto* x = execution.norm ? execution.norm->src[0] : rope->src[0];
const auto* pe = rope->src[1];
int32_t interleaved;
std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved));
std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
float eps = 0.0f;
if (execution.norm) {
std::memcpy(&eps, execution.norm->op_params, sizeof(eps));
@@ -90,7 +90,7 @@ namespace sd_cpu {
if (norm->op != GGML_OP_RMS_NORM || norm->type != GGML_TYPE_F32 || !norm->src[0] ||
norm->src[0]->type != GGML_TYPE_F32 || norm->src[0]->nb[0] != sizeof(float) ||
mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || rope->src[0] != mul ||
(rope->op != GGML_OP_ROPE_APPLY && ggml_ext_op_get(rope) != sd_ggml_rope_op()) ||
ggml_ext_op_get(rope) != sd_ggml_rope_op() ||
!sd_ggml_rope_describe(rope) || !ggml_are_same_shape(norm, mul)) {
return 0;
}
@@ -126,16 +126,12 @@ namespace sd_cpu {
desc.execute = rope_execute;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
desc.name = "sd.rope_apply.cpu.override.v1";
desc.op = GGML_OP_ROPE_APPLY;
desc.name = "sd.rms_norm_mul_rope.cpu.v1";
desc.op = GGML_OP_RMS_NORM;
desc.custom_op = nullptr;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
desc.name = "sd.rms_norm_mul_rope.cpu.v1";
desc.op = GGML_OP_RMS_NORM;
desc.priority = 200;
desc.supports = nullptr;
desc.match = rope_fusion_match;
desc.priority = 200;
desc.supports = nullptr;
desc.match = rope_fusion_match;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
}
+1 -1
View File
@@ -211,7 +211,7 @@ bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat,
!sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
return false;
int32_t interleaved;
std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved));
std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
return interleaved != 0;
}
+8 -11
View File
@@ -132,7 +132,7 @@ namespace SD_CUDA_NAMESPACE {
static bool rope_interleaved(const ggml_tensor* rope) {
int32_t value;
std::memcpy(&value, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(value));
std::memcpy(&value, ggml_ext_op_attributes(rope), sizeof(value));
return value != 0;
}
@@ -143,7 +143,7 @@ namespace SD_CUDA_NAMESPACE {
}
static bool rope_is_node(const ggml_tensor* node) {
return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op());
return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
}
static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
@@ -324,16 +324,13 @@ namespace SD_CUDA_NAMESPACE {
desc.destroy = rope_destroy;
desc.execute = rope_execute;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.rope_apply." SD_CUDA_NAME ".override.v1";
desc.op = GGML_OP_ROPE_APPLY;
desc.name = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_RMS_NORM;
desc.custom_op = nullptr;
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." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_RMS_NORM;
desc.priority = 250;
desc.supports = nullptr;
desc.match = rope_match;
desc.optimize = rope_optimize;
desc.priority = 250;
desc.supports = nullptr;
desc.match = rope_match;
desc.optimize = rope_optimize;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
+3 -13
View File
@@ -265,19 +265,9 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved) {
sd_ggml_register_extensions(backend);
if (backend) {
auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")
#ifdef SD_BACKEND_EXTENSIONS
|| ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1") || ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")
#endif
) {
ggml_tensor* inputs[] = {x, pe};
int32_t attributes = interleaved ? 1 : 0;
return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
}
}
return ggml_rope_apply(ctx, x, pe, interleaved);
ggml_tensor* inputs[] = {x, pe};
int32_t attributes = interleaved ? 1 : 0;
return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
}
#else
+6
View File
@@ -10,6 +10,12 @@ void sd_ggml_register_extensions(ggml_backend_t backend);
ggml_ext_op_t sd_ggml_rope_op();
bool sd_ggml_rope_describe(const ggml_tensor* node);
// Apply precomputed rotary matrices, with shapes in ne[] order:
// a: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens].
// pe stores [[cos, -sin], [sin, cos]] for each pair and token, shared across heads/batches.
// Returns contiguous F32 [head_dim, tokens, heads*batch]. Input strides are preserved.
// Interleaved pairs are (2*i, 2*i+1); otherwise pairs are (i, i+rot_dim/2).
// Channels from rot_dim to head_dim are copied unchanged.
ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
struct SDI8Operation {
+8 -11
View File
@@ -32,7 +32,7 @@ namespace sd_vulkan_rope {
return value;
}
static bool is_rope(const ggml_tensor* node) {
return node && (node->op == GGML_OP_ROPE_APPLY || ggml_ext_op_get(node) == sd_ggml_rope_op());
return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
}
static bool aligned(const ggml_tensor* tensor, size_t alignment) {
size_t offset = 0;
@@ -166,7 +166,7 @@ namespace sd_vulkan_rope {
pc.x_offset = buffers[0].shader_offset / 4;
pc.pe_offset = buffers[1].shader_offset / 4;
pc.dst_offset = buffers[2].shader_offset / uint32_t(ggml_type_size(output->type));
std::memcpy(&pc.interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, 4);
std::memcpy(&pc.interleaved, ggml_ext_op_attributes(rope), 4);
uint32_t groups[3] = {1, 1, 1};
ggml_ext_vulkan_pipeline pipeline = {};
pipeline.name = "sd.rope_apply.vulkan";
@@ -213,15 +213,12 @@ void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) {
desc.execute = sd_vulkan_rope::execute;
desc.destroy = sd_vulkan_rope::destroy;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
desc.name = "sd.rope_apply.vulkan.override.v1";
desc.op = GGML_OP_ROPE_APPLY;
desc.name = "sd.rms_rope_apply.vulkan.v1";
desc.op = GGML_OP_RMS_NORM;
desc.custom_op = nullptr;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
desc.name = "sd.rms_rope_apply.vulkan.v1";
desc.op = GGML_OP_RMS_NORM;
desc.priority = 200;
desc.supports = nullptr;
desc.match = sd_vulkan_rope::match;
desc.optimize = sd_vulkan_rope::optimize;
desc.priority = 200;
desc.supports = nullptr;
desc.match = sd_vulkan_rope::match;
desc.optimize = sd_vulkan_rope::optimize;
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
}