DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 03:24:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(cuda): fuse RMSNorm and RoPE in Sol-H3-VAE (#2323)

Browse Source
This commit is contained in:
losewayy authored and GitHub committed 2026-11-11 01:06:53 +08:00
1 parent 465adf389f
commit f8417d02e1
3 files changed
+58 -10
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/ggml_extensions.md
+5 -3
View File
@@ -215,11 +215,13 @@ 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
and `RMSNorm -> Mul -> RoPE [-> Scale -> FP16 Copy]` fusion. CuTe attention shares
`src/core/cuda/rope.cu` supplies standalone CUDA RoPE,
`RMSNorm -> RoPE` fusion for non-affine normalization, 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
requires contiguous position encoding and weights. Copy fusion uses workspace
requires contiguous position encoding and, when present, weights. The non-affine
fusion uses 64-channel F32 inputs and split-half RoPE. Copy fusion uses workspace
when the destination overlaps an input. These registrations do not require
CUTLASS. The old CUDA/HIP/MUSA RoPE dispatch is removed.
src/core/cuda/rope.cu
+40 -6
View File
@@ -12,7 +12,7 @@ namespace SD_CUDA_NAMESPACE {
v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
return v;
}
template <int D, bool interleaved, typename T>
template <int D, bool interleaved, typename T, bool affine = true>
static __global__ void rms_norm_mul_rope_apply_fused(
const float* x,
const float* weight,
@@ -51,7 +51,7 @@ namespace SD_CUDA_NAMESPACE {
#pragma unroll
for (int j = 0; j < D / 32; ++j) {
const int col = j * 32 + lane;
const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
const float value = affine ? __fmul_rn(__fmul_rn(scale, v[j]), weight[col]) : __fmul_rn(scale, v[j]);
float rotated = value;
if constexpr (interleaved) {
const float other = cuda_shuffle_xor(value, 1, 32);
@@ -65,7 +65,7 @@ namespace SD_CUDA_NAMESPACE {
} else {
if (col < 2 * rot_half) {
const int partner = col < rot_half ? col + rot_half : col - rot_half;
const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
const float other = affine ? __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]) : __fmul_rn(scale, x_row[partner]);
const float* m = pe + 4 * (token * rot_half + col % rot_half);
const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
rotated = col < rot_half
@@ -178,6 +178,20 @@ namespace SD_CUDA_NAMESPACE {
(ggml_nrows(x) + 3) / 4 <= INT_MAX;
}
static bool rope_non_affine_fusion_supports(const ggml_tensor* norm, const ggml_tensor* rope) {
if (!norm || norm->op != GGML_OP_RMS_NORM || !rope_is_node(rope) || rope->src[0] != norm ||
!rope_supports(nullptr, rope, nullptr) || rope_interleaved(rope))
return false;
const auto* x = norm->src[0];
const auto* pe = rope->src[1];
// The kernel reads contiguous channels and position matrices, with row strides in floats.
return x && x->type == GGML_TYPE_F32 && x->ne[0] == 64 &&
x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 &&
x->nb[2] % sizeof(float) == 0 && x->nb[3] % sizeof(float) == 0 &&
ggml_are_same_shape(x, norm) && ggml_is_contiguous(pe) &&
ggml_nrows(x) > 0 && (ggml_nrows(x) + 3) / 4 <= INT_MAX;
}
template <typename T, bool Interleaved>
static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
const auto* x = nodes[0]->src[0];
@@ -220,6 +234,9 @@ namespace SD_CUDA_NAMESPACE {
if (disabled)
return 0;
auto* g = const_cast<ggml_cgraph*>(graph);
if (first + 1 < ggml_graph_n_nodes(g) &&
rope_non_affine_fusion_supports(ggml_graph_node(g, first), ggml_graph_node(g, first + 1)))
return 2;
if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
return 0;
@@ -253,6 +270,16 @@ namespace SD_CUDA_NAMESPACE {
for (int i = 2; i + 1 < total; ++i) {
const auto* mul = ggml_graph_node(graph, i);
const auto* rope = ggml_graph_node(graph, i + 1);
if (rope_non_affine_fusion_supports(rope->src[0], rope)) {
for (int j = i; j >= 0 && j >= i - 64; --j) {
if (ggml_graph_node(graph, j) == rope->src[0]) {
if (j < i)
ggml_ext_graph_delay(graph, j, 1, i + 1);
break;
}
}
continue;
}
if (mul->op != GGML_OP_MUL || !rope_is_node(rope) || rope->src[0] != mul ||
!mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
continue;
@@ -296,13 +323,20 @@ namespace SD_CUDA_NAMESPACE {
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
auto stream = static_cast<CUDAStream>(context.stream);
const auto* rope = nodes[count == 1 ? 0 : 2];
const auto* rope = nodes[count <= 2 ? count - 1 : 2];
const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0];
auto* dst = nodes[count - 1];
const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 2 && overlaps(dst, nodes[1]->src[1]));
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
CUDAError error;
if (count > 1) {
if (count == 2) {
const auto* pe = rope->src[1];
const int64_t rows = ggml_nrows(x);
rms_norm_mul_rope_apply_fused<64, false, float, false><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>(
static_cast<const float*>(x->data), nullptr, static_cast<const float*>(pe->data), static_cast<float*>(output),
x->ne[1], x->ne[2], rows, x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), 1.f);
error = cuda_last_error();
} else if (count > 2) {
error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f);
} else {
RopeApplyStrides strides;
src/model/vae/minimax_h3_vae.hpp
+13 -1
View File
@@ -335,6 +335,18 @@ namespace MiniMaxH3VAE {
return {q, k, v};
}
std::array<ggml_tensor*, 3> project_sol(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe) {
auto to_qkv = std::dynamic_pointer_cast<Linear>(blocks["to_qkv"]);
auto qkv_projection = to_qkv->forward(ctx, x);
qkv_projection = ggml_reshape_4d(ctx->ggml_ctx, qkv_projection, 3 * head_dim, num_head, x->ne[1], x->ne[2] * x->ne[3]);
auto qkv = ggml_ext_chunk(ctx, qkv_projection, 3, 0, false);
auto q = ggml_rms_norm(ctx->ggml_ctx, qkv[0], 1e-5f);
auto k = ggml_rms_norm(ctx->ggml_ctx, qkv[1], 1e-5f);
return {Rope::apply_rope(ctx, q, pe, false),
Rope::apply_rope(ctx, k, pe, false),
ggml_ext_cont(ctx, qkv[2])};
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
@@ -343,7 +355,7 @@ namespace MiniMaxH3VAE {
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
if (sol != nullptr) {
auto project_rows = [&](ggml_tensor* rows, ggml_tensor* local_pe) {
auto [q, k, v] = project(ctx, norm.forward(ctx, rows), local_pe);
auto [q, k, v] = project_sol(ctx, norm.forward(ctx, rows), local_pe);
auto as_rows = [&](ggml_tensor* tensor) {
return ggml_reshape_2d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, tensor), dim, rows->ne[1]);
};