diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 7485938f..c643237f 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -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. diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index 2bad5df5..c012c404 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -12,7 +12,7 @@ namespace SD_CUDA_NAMESPACE { v = __fadd_rn(v, cuda_shuffle_xor(v, mask)); return v; } - template + template 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 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(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(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><<((rows + 3) / 4), dim3(32, 4), 0, stream>>>( + static_cast(x->data), nullptr, static_cast(pe->data), static_cast(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; diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index d4e1c324..881fd019 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -335,6 +335,18 @@ namespace MiniMaxH3VAE { return {q, k, v}; } + std::array project_sol(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe) { + auto to_qkv = std::dynamic_pointer_cast(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(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]); };