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
+26 -38

No files matched your search

+9 -7
View File
@@ -136,13 +136,15 @@ in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their
respective backend directories.
`src/core/cpu/rope.cpp` registers the CPU implementation of
`sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a
`sd.rope_apply.v1` and a
`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
vector. Both interleaved and split-half RoPE, partial rotation, and strided input
rows are supported. Model construction uses the custom operator on CPU, CUDA,
HIP, MUSA, and Vulkan, and retains existing backend implementations elsewhere.
HIP, MUSA, and Vulkan. Other backends can use the scheduler's CPU fallback.
The former `ggml_rope_apply` API, builtin opcode, and CPU implementation have
been removed from ggml.
`src/core/cuda/rope.cu` supplies standalone CUDA RoPE, the same builtin override,
`src/core/cuda/rope.cu` supplies standalone CUDA RoPE
and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
position-encoding strides are supported by the standalone kernel; RMSNorm fusion
@@ -180,7 +182,7 @@ Independent affine operands, including broadcasted constants, are computed befor
their normalization/projection
chain when graph dependencies allow it. The attention optimizer can then move
the complete five-node projection when safe.
These rules accept both builtin and custom RoPE nodes. They preserve the original
These rules accept custom RoPE nodes. They preserve the original
large-matrix selection thresholds, full interleaved head dimension 128, packed
activation layout, and separate weight scales. External consumers prevent fusion;
aliasing terminal copies use scheduler workspace. The CUDA projection kernels
@@ -379,7 +381,7 @@ Callback errors propagate without retry. Evaluation callbacks, direct backend
calls, synchronous devices, serialized diagnostics, and timestamp profiling
retain fragment execution.
`src/core/vulkan/rope.cpp` registers custom and builtin RoPE implementations,
`src/core/vulkan/rope.cpp` registers the custom RoPE implementation,
plus `RMSNorm -> Mul -> RoPE` and `RMSNorm -> Mul -> RoPE -> Scale -> FP16 Copy`
fusion. These kernels retain the original strides, partial-rotation behavior,
128-channel subgroup reduction, and separate arithmetic rounding. Their SPIR-V
@@ -390,8 +392,8 @@ External intermediate consumers retain ordinary operators, and overlapping
outputs use plan workspace. `GGML_VK_DISABLE_FUSION` is honored.
The corresponding Vulkan shader sources and native fusion rules have been
removed from ggml. Existing builtin RoPE graphs use the registered override;
unsupported cases retain scheduler fallback to another supported backend.
removed from ggml. Unsupported custom nodes retain scheduler fallback to
another supported backend.
Remaining model-specific CUDA paths and Vulkan INT8/ConvRot implementations still
need auditing or migration.
+1 -1
Submodule ggml updated: a294003351...394d41dd40.
+3 -7
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,13 +126,9 @@ 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.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.custom_op = nullptr;
desc.priority = 200;
desc.supports = nullptr;
desc.match = rope_fusion_match;
+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;
}
+3 -6
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,12 +324,9 @@ 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.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.custom_op = nullptr;
desc.priority = 250;
desc.supports = nullptr;
desc.match = rope_match;
-10
View File
@@ -265,20 +265,10 @@ 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);
}
#else
void sd_ggml_register_extensions(ggml_backend_t) {}
+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 {
+3 -6
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,12 +213,9 @@ 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.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.custom_op = nullptr;
desc.priority = 200;
desc.supports = nullptr;
desc.match = sd_vulkan_rope::match;