mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 11:56:36 +00:00
refactor: use custom RoPE Apply nodes exclusively
This commit is contained in:
1 parent
d07d48d466
commit
051b49751e
8 files changed
+26
-38
No files matched your search
@@ -136,13 +136,15 @@ in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their
|
|||||||
respective backend directories.
|
respective backend directories.
|
||||||
|
|
||||||
`src/core/cpu/rope.cpp` registers the CPU implementation of
|
`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
|
`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
|
||||||
vector. Both interleaved and split-half RoPE, partial rotation, and strided input
|
vector. Both interleaved and split-half RoPE, partial rotation, and strided input
|
||||||
rows are supported. Model construction uses the custom operator on CPU, CUDA,
|
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
|
and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
|
||||||
its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
|
its RMSNorm/RoPE kernel. Partial rotation and arbitrary float-aligned input and
|
||||||
position-encoding strides are supported by the standalone kernel; RMSNorm fusion
|
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
|
their normalization/projection
|
||||||
chain when graph dependencies allow it. The attention optimizer can then move
|
chain when graph dependencies allow it. The attention optimizer can then move
|
||||||
the complete five-node projection when safe.
|
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
|
large-matrix selection thresholds, full interleaved head dimension 128, packed
|
||||||
activation layout, and separate weight scales. External consumers prevent fusion;
|
activation layout, and separate weight scales. External consumers prevent fusion;
|
||||||
aliasing terminal copies use scheduler workspace. The CUDA projection kernels
|
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
|
calls, synchronous devices, serialized diagnostics, and timestamp profiling
|
||||||
retain fragment execution.
|
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`
|
plus `RMSNorm -> Mul -> RoPE` and `RMSNorm -> Mul -> RoPE -> Scale -> FP16 Copy`
|
||||||
fusion. These kernels retain the original strides, partial-rotation behavior,
|
fusion. These kernels retain the original strides, partial-rotation behavior,
|
||||||
128-channel subgroup reduction, and separate arithmetic rounding. Their SPIR-V
|
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.
|
outputs use plan workspace. `GGML_VK_DISABLE_FUSION` is honored.
|
||||||
|
|
||||||
The corresponding Vulkan shader sources and native fusion rules have been
|
The corresponding Vulkan shader sources and native fusion rules have been
|
||||||
removed from ggml. Existing builtin RoPE graphs use the registered override;
|
removed from ggml. Unsupported custom nodes retain scheduler fallback to
|
||||||
unsupported cases retain scheduler fallback to another supported backend.
|
another supported backend.
|
||||||
|
|
||||||
Remaining model-specific CUDA paths and Vulkan INT8/ConvRot implementations still
|
Remaining model-specific CUDA paths and Vulkan INT8/ConvRot implementations still
|
||||||
need auditing or migration.
|
need auditing or migration.
|
||||||
|
|||||||
+1
-1
Submodule ggml updated: a294003351...394d41dd40.
@@ -23,7 +23,7 @@ namespace sd_cpu {
|
|||||||
const auto* x = execution.norm ? execution.norm->src[0] : rope->src[0];
|
const auto* x = execution.norm ? execution.norm->src[0] : rope->src[0];
|
||||||
const auto* pe = rope->src[1];
|
const auto* pe = rope->src[1];
|
||||||
int32_t interleaved;
|
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;
|
float eps = 0.0f;
|
||||||
if (execution.norm) {
|
if (execution.norm) {
|
||||||
std::memcpy(&eps, execution.norm->op_params, sizeof(eps));
|
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] ||
|
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) ||
|
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 ||
|
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)) {
|
!sd_ggml_rope_describe(rope) || !ggml_are_same_shape(norm, mul)) {
|
||||||
return 0;
|
return 0;
|
||||||
}
|
}
|
||||||
@@ -126,13 +126,9 @@ namespace sd_cpu {
|
|||||||
desc.execute = rope_execute;
|
desc.execute = rope_execute;
|
||||||
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
|
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.name = "sd.rms_norm_mul_rope.cpu.v1";
|
||||||
desc.op = GGML_OP_RMS_NORM;
|
desc.op = GGML_OP_RMS_NORM;
|
||||||
|
desc.custom_op = nullptr;
|
||||||
desc.priority = 200;
|
desc.priority = 200;
|
||||||
desc.supports = nullptr;
|
desc.supports = nullptr;
|
||||||
desc.match = rope_fusion_match;
|
desc.match = rope_fusion_match;
|
||||||
|
|||||||
@@ -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]))
|
!sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
|
||||||
return false;
|
return false;
|
||||||
int32_t interleaved;
|
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;
|
return interleaved != 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -132,7 +132,7 @@ namespace SD_CUDA_NAMESPACE {
|
|||||||
|
|
||||||
static bool rope_interleaved(const ggml_tensor* rope) {
|
static bool rope_interleaved(const ggml_tensor* rope) {
|
||||||
int32_t value;
|
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;
|
return value != 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -143,7 +143,7 @@ namespace SD_CUDA_NAMESPACE {
|
|||||||
}
|
}
|
||||||
|
|
||||||
static bool rope_is_node(const ggml_tensor* node) {
|
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*) {
|
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.destroy = rope_destroy;
|
||||||
desc.execute = rope_execute;
|
desc.execute = rope_execute;
|
||||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
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.name = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
|
||||||
desc.op = GGML_OP_RMS_NORM;
|
desc.op = GGML_OP_RMS_NORM;
|
||||||
|
desc.custom_op = nullptr;
|
||||||
desc.priority = 250;
|
desc.priority = 250;
|
||||||
desc.supports = nullptr;
|
desc.supports = nullptr;
|
||||||
desc.match = rope_match;
|
desc.match = rope_match;
|
||||||
|
|||||||
@@ -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) {
|
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);
|
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};
|
ggml_tensor* inputs[] = {x, pe};
|
||||||
int32_t attributes = interleaved ? 1 : 0;
|
int32_t attributes = interleaved ? 1 : 0;
|
||||||
return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
|
return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
|
||||||
}
|
}
|
||||||
}
|
|
||||||
return ggml_rope_apply(ctx, x, pe, interleaved);
|
|
||||||
}
|
|
||||||
|
|
||||||
#else
|
#else
|
||||||
void sd_ggml_register_extensions(ggml_backend_t) {}
|
void sd_ggml_register_extensions(ggml_backend_t) {}
|
||||||
|
|||||||
@@ -10,6 +10,12 @@ void sd_ggml_register_extensions(ggml_backend_t backend);
|
|||||||
|
|
||||||
ggml_ext_op_t sd_ggml_rope_op();
|
ggml_ext_op_t sd_ggml_rope_op();
|
||||||
bool sd_ggml_rope_describe(const ggml_tensor* node);
|
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);
|
ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
|
||||||
|
|
||||||
struct SDI8Operation {
|
struct SDI8Operation {
|
||||||
|
|||||||
@@ -32,7 +32,7 @@ namespace sd_vulkan_rope {
|
|||||||
return value;
|
return value;
|
||||||
}
|
}
|
||||||
static bool is_rope(const ggml_tensor* node) {
|
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) {
|
static bool aligned(const ggml_tensor* tensor, size_t alignment) {
|
||||||
size_t offset = 0;
|
size_t offset = 0;
|
||||||
@@ -166,7 +166,7 @@ namespace sd_vulkan_rope {
|
|||||||
pc.x_offset = buffers[0].shader_offset / 4;
|
pc.x_offset = buffers[0].shader_offset / 4;
|
||||||
pc.pe_offset = buffers[1].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));
|
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};
|
uint32_t groups[3] = {1, 1, 1};
|
||||||
ggml_ext_vulkan_pipeline pipeline = {};
|
ggml_ext_vulkan_pipeline pipeline = {};
|
||||||
pipeline.name = "sd.rope_apply.vulkan";
|
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.execute = sd_vulkan_rope::execute;
|
||||||
desc.destroy = sd_vulkan_rope::destroy;
|
desc.destroy = sd_vulkan_rope::destroy;
|
||||||
GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
|
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.name = "sd.rms_rope_apply.vulkan.v1";
|
||||||
desc.op = GGML_OP_RMS_NORM;
|
desc.op = GGML_OP_RMS_NORM;
|
||||||
|
desc.custom_op = nullptr;
|
||||||
desc.priority = 200;
|
desc.priority = 200;
|
||||||
desc.supports = nullptr;
|
desc.supports = nullptr;
|
||||||
desc.match = sd_vulkan_rope::match;
|
desc.match = sd_vulkan_rope::match;
|
||||||
|
|||||||
Reference in new issue
Block a user