diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 1aa3f2a7..50ea815d 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -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. diff --git a/ggml b/ggml index a2940033..394d41dd 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit a29400335191f983544d14df03773a50420224d7 +Subproject commit 394d41dd400cd4a9b1529bb11bc34671cc5c48d2 diff --git a/src/core/cpu/rope.cpp b/src/core/cpu/rope.cpp index 70626fdb..ad1ac79b 100644 --- a/src/core/cpu/rope.cpp +++ b/src/core/cpu/rope.cpp @@ -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)); } diff --git a/src/core/cuda/i8_projection.cu b/src/core/cuda/i8_projection.cu index 4a782bf6..5784db43 100644 --- a/src/core/cuda/i8_projection.cu +++ b/src/core/cuda/i8_projection.cu @@ -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; } diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index ffb7fa31..46002790 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -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)); } diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index 877cef93..4dfe3763 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -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 diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h index 456b7e0e..615a780b 100644 --- a/src/core/ggml_extensions.h +++ b/src/core/ggml_extensions.h @@ -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 { diff --git a/src/core/vulkan/rope.cpp b/src/core/vulkan/rope.cpp index 4ed0811d..462b6a6f 100644 --- a/src/core/vulkan/rope.cpp +++ b/src/core/vulkan/rope.cpp @@ -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)); }