Skip to content

perf(cuda): fuse MOSS modulation and direct RoPE F16 output #2269

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
leejet merged 1 commit into from
Nov 3, 2026
+164 −59
Merged

perf(cuda): fuse MOSS modulation and direct RoPE F16 output #2269

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
206 changes: 152 additions & 54 deletions src/core/cuda/rms_affine.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#include <cstdlib>
#include <cstring>
#include "bf16.h"
#include "core/ggml_extensions.h"
#include "rms_affine.h"

namespace SD_CUDA_NAMESPACE {
Expand All @@ -18,6 +19,28 @@ namespace SD_CUDA_NAMESPACE {
return t && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t);
}

static int modulate_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
auto* g = const_cast<ggml_cgraph*>(graph);
if (disabled() || first + 2 >= ggml_graph_n_nodes(g))
return 0;
const auto* norm = ggml_graph_node(g, first);
const auto* affine = ggml_graph_node(g, first + 1);
const auto* modulate = ggml_graph_node(g, first + 2);
if (norm->op != GGML_OP_RMS_NORM || affine->op != GGML_OP_MUL || affine->src[0] != norm ||
!sd_ggml_modulate_describe(modulate) || modulate->src[0] != affine)
return 0;
const auto* x = norm->src[0];
if (!f32(x) || !f32(norm) || !f32(affine) || x->ne[0] < 32 || x->ne[0] > 8192 ||
ggml_nrows(x) > INT_MAX || !ggml_are_same_shape(x, norm) || !ggml_are_same_shape(x, affine))
return 0;
for (const auto* parameter : {affine->src[1], modulate->src[1], modulate->src[2]}) {
if (parameter && (!f32(parameter) || parameter == norm || parameter == affine ||
parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0]))
return 0;
}
return affine->src[1] ? 3 : 0;
}

static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
auto* g = const_cast<ggml_cgraph*>(graph);
if (disabled() || first + 3 >= ggml_graph_n_nodes(g))
Expand All @@ -44,13 +67,17 @@ namespace SD_CUDA_NAMESPACE {
static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
if (disabled())
return;
for (int i = 2; i + 1 < ggml_graph_n_nodes(graph); ++i) {
for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
const auto* mul = ggml_graph_node(graph, i);
const auto* add = ggml_graph_node(graph, i + 1);
if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
continue;
if (!sd_ggml_modulate_describe(mul)) {
if (i + 1 >= ggml_graph_n_nodes(graph))
continue;
const auto* add = ggml_graph_node(graph, i + 1);
if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
continue;
}
const auto* affine = mul->src[0];
if (affine->op != GGML_OP_MUL || affine->src[0]->op != GGML_OP_RMS_NORM)
if (!affine || affine->op != GGML_OP_MUL || !affine->src[0] || affine->src[0]->op != GGML_OP_RMS_NORM)
continue;
for (int j = i - 2; j >= 0 && j >= i - 32; --j) {
if (ggml_graph_node(graph, j) == affine->src[0] && ggml_graph_node(graph, j + 1) == affine) {
Expand All @@ -68,64 +95,136 @@ namespace SD_CUDA_NAMESPACE {
return value;
}

template <int Block, bool Modulate, typename T = float>
static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, T* dst, int width, float eps) {
enum class Kind { Residual,
SelfResidual,
Modulate };

struct Operation {
Kind kind = Kind::Residual;
const ggml_tensor* input;
const ggml_tensor* weight;
const ggml_tensor* gate;
const ggml_tensor* residual = nullptr;
const ggml_tensor* shift = nullptr;
ggml_tensor* output;
float eps;
};

static Operation describe(ggml_tensor* const* nodes, int count) {
Operation op;
op.input = nodes[0]->src[0];
op.weight = nodes[1]->src[1];
op.gate = nodes[2]->src[1];
op.output = nodes[count - 1];
std::memcpy(&op.eps, nodes[0]->op_params, sizeof(op.eps));
if (count == 3) {
op.kind = Kind::Modulate;
op.shift = op.output->src[2];
} else {
op.residual = op.output->src[0] == nodes[2] ? op.output->src[1] : op.output->src[0];
if (op.residual == nodes[1]) {
op.kind = Kind::SelfResidual;
op.residual = nullptr;
}
}
return op;
}

template <int Block, int Groups, Kind Op, typename T>
static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, const float* shift, T* dst, int width, float eps) {
const int col0 = threadIdx.x;
const int64_t base = int64_t(blockIdx.x) * width;
float values[8], sum = 0.f;
for (int j = 0, col = col0; col < width; ++j, col += Block) {
values[j] = x[base + col];
sum = __fmaf_rn(values[j], values[j], sum);
}
sum = warp_sum(sum);
float values[Groups][8];
__shared__ float sums[32];
const int lane = col0 % 32;
if (lane == 0)
sums[col0 / 32] = sum;
#pragma unroll
for (int group = 0; group < Groups; ++group) {
float sum = 0.f;
for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
values[group][j] = x[base + col];
sum = __fmaf_rn(values[group][j], values[group][j], sum);
}
sum = warp_sum(sum);
if (lane == 0)
sums[(col0 + group * Block) / 32] = sum;
}
// Groups preserve the native reduction tree when using fewer physical threads.
__syncthreads();
sum = warp_sum(lane < Block / 32 ? sums[lane] : 0.f);
const float sum = warp_sum(lane < Block * Groups / 32 ? sums[lane] : 0.f);
const float scale = rsqrtf(sum / width + eps);
for (int j = 0, col = col0; col < width; ++j, col += Block) {
const float value = __fmul_rn(__fmul_rn(scale, values[j]), weight[col]);
const float r = Modulate ? value : residual[base + col];
dst[base + col] = T(__fadd_rn(__fmul_rn(value, gate[col]), r));
#pragma unroll
for (int group = 0; group < Groups; ++group) {
for (int j = 0, col = col0 + group * Block; col < width; ++j, col += Block * Groups) {
float value = __fmul_rn(__fmul_rn(scale, values[group][j]), weight[col]);
if constexpr (Op == Kind::Modulate) {
value = __fadd_rn(value, __fmul_rn(value, gate[col]));
if (shift)
value = __fadd_rn(value, shift[col]);
} else {
const float r = Op == Kind::SelfResidual ? value : residual[base + col];
value = __fadd_rn(__fmul_rn(value, gate[col]), r);
}
dst[base + col] = T(value);
}
}
}

static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
if (!b)
return false;
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
}

static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
*workspace = ggml_nbytes(nodes[3]);
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
*workspace = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}

static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
template <typename T, Kind Op>
static CUDAError launch(CUDAStream stream, const Operation& op, T* output) {
auto compute = kernel<256, 1, Op, T>;
int threads = 256;
if (op.input->ne[0] >= 1024) {
if constexpr (Op == Kind::Modulate) {
compute = kernel<512, 2, Op, T>;
threads = 512;
} else {
compute = kernel<1024, 1, Op, T>;
threads = 1024;
}
}
compute<<<unsigned(ggml_nrows(op.input)), threads, 0, stream>>>(
static_cast<const float*>(op.input->data), static_cast<const float*>(op.weight->data),
static_cast<const float*>(op.gate->data), op.residual ? static_cast<const float*>(op.residual->data) : nullptr,
op.shift ? static_cast<const float*>(op.shift->data) : nullptr, output, int(op.input->ne[0]), op.eps);
return cuda_last_error();
}

static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
const auto* x = nodes[0]->src[0];
const auto* w = nodes[1]->src[1];
const auto* gate = nodes[2]->src[1];
auto* dst = nodes[3];
const auto* residual = dst->src[0] == nodes[2] ? dst->src[1] : dst->src[0];
const bool modulation = residual == nodes[1];
const bool alias = (overlaps(dst, x) && dst->data != x->data) || overlaps(dst, w) || overlaps(dst, gate) ||
(!modulation && overlaps(dst, residual) && dst->data != residual->data);
auto* output = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
float eps;
std::memcpy(&eps, nodes[0]->op_params, sizeof(eps));
auto launch = x->ne[0] < 1024 ? (modulation ? kernel<256, true> : kernel<256, false>)
: (modulation ? kernel<1024, true> : kernel<1024, false>);
const auto op = describe(nodes, count);
auto* dst = op.output;
const bool alias = (overlaps(dst, op.input) && dst->data != op.input->data) ||
(overlaps(dst, op.residual) && dst->data != op.residual->data) ||
overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift);
auto* output = static_cast<float*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
const auto stream = static_cast<CUDAStream>(context.stream);
const int threads = x->ne[0] < 1024 ? 256 : 1024;
launch<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
static_cast<const float*>(x->data), static_cast<const float*>(w->data), static_cast<const float*>(gate->data),
modulation ? nullptr : static_cast<const float*>(residual->data), output, int(x->ne[0]), eps);
auto error = cuda_last_error();
CUDAError error = cuda_invalid_value;
switch (op.kind) {
case Kind::Residual:
error = launch<float, Kind::Residual>(stream, op, output);
break;
case Kind::SelfResidual:
error = launch<float, Kind::SelfResidual>(stream, op, output);
break;
case Kind::Modulate:
error = launch<float, Kind::Modulate>(stream, op, output);
break;
}
if (error == cuda_success && alias)
error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
Expand All @@ -137,19 +236,12 @@ namespace SD_CUDA_NAMESPACE {
}

ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output) {
const auto* x = nodes[0]->src[0];
const auto* residual = nodes[3]->src[0] == nodes[2] ? nodes[3]->src[1] : nodes[3]->src[0];
const bool modulation = residual == nodes[1];
auto launch = x->ne[0] < 1024 ? (modulation ? sd_rms_affine::kernel<256, true, CUDABFloat16> : sd_rms_affine::kernel<256, false, CUDABFloat16>)
: (modulation ? sd_rms_affine::kernel<1024, true, CUDABFloat16> : sd_rms_affine::kernel<1024, false, CUDABFloat16>);
float eps;
std::memcpy(&eps, nodes[0]->op_params, sizeof(eps));
const int threads = x->ne[0] < 1024 ? 256 : 1024;
launch<<<unsigned(ggml_nrows(x)), threads, 0, stream>>>(
static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),
static_cast<const float*>(nodes[2]->src[1]->data), modulation ? nullptr : static_cast<const float*>(residual->data),
static_cast<CUDABFloat16*>(output), int(x->ne[0]), eps);
return cuda_last_error() == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
const auto op = sd_rms_affine::describe(nodes, 4);
auto* dst = static_cast<CUDABFloat16*>(output);
const auto error = op.kind == sd_rms_affine::Kind::SelfResidual
? sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::SelfResidual>(stream, op, dst)
: sd_rms_affine::launch<CUDABFloat16, sd_rms_affine::Kind::Residual>(stream, op, dst);
return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
}

void register_rms_affine(ggml_backend_dev_t device) {
Expand All @@ -164,5 +256,11 @@ namespace SD_CUDA_NAMESPACE {
desc.destroy = sd_rms_affine::destroy;
desc.execute = sd_rms_affine::execute;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
desc.name = "sd.rms_affine_modulate." SD_CUDA_NAME ".v1";
// Prefer the existing fusion that also produces ConvRot input when available.
desc.priority = 200;
desc.match = sd_rms_affine::modulate_match;
desc.optimize = nullptr;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
11 changes: 9 additions & 2 deletions src/core/cuda/rope.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ namespace SD_CUDA_NAMESPACE {
return false;
const auto* x = norm->src[0];
const auto* weight = mul->src[1];
return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
return x && weight && weight != norm && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
Expand Down Expand Up @@ -232,6 +232,13 @@ namespace SD_CUDA_NAMESPACE {
ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
return 5;
}
if (first + 3 < ggml_graph_n_nodes(g)) {
const auto* rope = ggml_graph_node(g, first + 2);
const auto* copy = ggml_graph_node(g, first + 3);
if (copy->op == GGML_OP_CPY && copy->src[0] == rope && copy->type == GGML_TYPE_F16 &&
ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
return 4;
}
return 3;
}

Expand Down Expand Up @@ -293,7 +300,7 @@ namespace SD_CUDA_NAMESPACE {
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
CUDAError error;
if (count > 1) {
error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f);
} else {
RopeApplyStrides strides;
const auto* pe = rope->src[1];
Expand Down
6 changes: 3 additions & 3 deletions src/model/diffusion/moss_image.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ namespace MOSSImage {
auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<RMSNorm>(blocks[name])->forward(ctx, h); };
auto part = [&](int i) { return mod[i]; };
auto y = norm("norm1", x);
y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(1))), part(0));
y = ggml_ext_modulate(ctx, y, part(0), part(1));
auto qkv = linear("proj.qkv", y);
auto split = [&](int i) {
return ggml_view_4d(c, qkv, config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2],
Expand All @@ -130,7 +130,7 @@ namespace MOSSImage {
y = ggml_mul(c, y, part(2));
x = ggml_add(c, x, y);
y = norm("norm2", x);
y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(4))), part(3));
y = ggml_ext_modulate(ctx, y, part(3), part(4));
auto up_gate = linear("proj.ffn_up_gate", y);
y = ggml_ext_situ(ctx, up_gate);
y = linear("proj.ffn_down", y);
Expand Down Expand Up @@ -191,7 +191,7 @@ namespace MOSSImage {
x = std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm_video"])->forward(ctx, x);
auto shift = ggml_ext_slice(c, head, 0, 0, config.dim);
auto scale = ggml_ext_slice(c, head, 0, config.dim, 2 * config.dim);
x = ggml_add(c, ggml_add(c, x, ggml_mul(c, x, scale)), shift);
x = ggml_ext_modulate(ctx, x, shift, scale);
x = linear("final_linear_video", x);
return DiT::unpatchify(c, x, image->ne[1], image->ne[0], 1, 1);
}
Expand Down
Loading