Commit 3bad892 for stable-diffusion.cpp
commit 3bad89267b780b320a40f34d3a0b2d70e3005bba
Author: leejet <leejet714@gmail.com>
Date: Sat Nov 14 02:16:15 2026 +0800
perf(cuda): optimize MOSS-Image inference (#2346)
diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp
index cf9d04c..128e0c3 100644
--- a/src/core/compute_workspace.cpp
+++ b/src/core/compute_workspace.cpp
@@ -53,10 +53,12 @@ namespace sd {
if (measurement.buffers.empty()) {
return false;
}
- const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
- [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
- if (grows) {
- release();
+ GGML_ASSERT(peak_scheduler_bytes_.empty() || peak_scheduler_bytes_.size() == measurement.buffers.size());
+ // The reported size includes per-graph scratch, which can shrink while the scheduler is retained.
+ for (size_t i = 0; i < peak_scheduler_bytes_.size(); ++i) {
+ if (measurement.buffers[i].bytes > peak_scheduler_bytes_[i]) {
+ return release_scheduler();
+ }
}
return true;
}
@@ -64,14 +66,14 @@ namespace sd {
bool ComputeWorkspace::release_excess(const Measurement& measurement) {
return std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
[&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) &&
- release();
+ release_scheduler();
}
- bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
+ bool ComputeWorkspace::allocate(ggml_cgraph* graph, const Measurement& measurement, const AssignNodes& assign_nodes) {
GGML_ASSERT(!active_);
const size_t capacity = static_cast<size_t>(ggml_graph_n_nodes(graph) + ggml_graph_n_leafs(graph)) + 8;
if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
- release();
+ release_scheduler();
scheduler_ = make_scheduler(capacity);
scheduler_capacity_ = capacity;
}
@@ -85,6 +87,10 @@ namespace sd {
release();
return false;
}
+ peak_scheduler_bytes_.resize(measurement.buffers.size());
+ for (size_t i = 0; i < measurement.buffers.size(); ++i) {
+ peak_scheduler_bytes_[i] = std::max(peak_scheduler_bytes_[i], bytes(measurement.buffers[i].backend));
+ }
active_ = true;
return true;
}
@@ -154,8 +160,13 @@ namespace sd {
append(uses);
}
// Match the full allocation metadata, including topology and external placement.
- if (key == measurement_key_) {
- return {measured_buffers_};
+ for (auto& cached : measurement_cache_) {
+ if (key == cached.key) {
+ if (&cached != &measurement_cache_.front()) {
+ std::swap(cached, measurement_cache_.front());
+ }
+ return {measurement_cache_.front().buffers};
+ }
}
}
const size_t graph_size = tensors.size() + 8;
@@ -229,8 +240,8 @@ namespace sd {
}
ggml_free(context);
if (cache_measurement && !result.buffers.empty()) {
- measurement_key_ = std::move(key);
- measured_buffers_ = result.buffers;
+ measurement_cache_.back() = std::move(measurement_cache_.front());
+ measurement_cache_.front() = {std::move(key), result.buffers};
}
return result;
}
@@ -263,14 +274,21 @@ namespace sd {
}
bool ComputeWorkspace::release() {
+ if (!release_scheduler()) {
+ return false;
+ }
+ measurement_cache_ = {};
+ return true;
+ }
+
+ bool ComputeWorkspace::release_scheduler() {
if (active_) {
return false;
}
- measurement_key_.clear();
- measured_buffers_.clear();
ggml_backend_sched_free(scheduler_);
scheduler_ = nullptr;
scheduler_capacity_ = 0;
+ peak_scheduler_bytes_.clear();
return true;
}
diff --git a/src/core/compute_workspace.h b/src/core/compute_workspace.h
index 8dee5da..6cfbd5c 100644
--- a/src/core/compute_workspace.h
+++ b/src/core/compute_workspace.h
@@ -1,6 +1,7 @@
#ifndef __SD_CORE_COMPUTE_WORKSPACE_H__
#define __SD_CORE_COMPUTE_WORKSPACE_H__
+#include <array>
#include <functional>
#include <vector>
@@ -19,10 +20,15 @@ namespace sd {
ggml_backend_sched_t scheduler_ = nullptr;
size_t scheduler_capacity_ = 0;
bool active_ = false;
- std::vector<uint8_t> measurement_key_;
- std::vector<BackendBufferSize> measured_buffers_;
+ std::vector<size_t> peak_scheduler_bytes_;
+ struct CachedMeasurement {
+ std::vector<uint8_t> key;
+ std::vector<BackendBufferSize> buffers;
+ };
+ std::array<CachedMeasurement, 2> measurement_cache_;
ggml_backend_sched_t make_scheduler(size_t graph_size);
+ bool release_scheduler();
public:
struct Measurement {
@@ -39,7 +45,7 @@ namespace sd {
void set_extra_backends(const std::vector<ggml_backend_t>& backends);
bool prepare(const Measurement& measurement);
bool release_excess(const Measurement& measurement);
- bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes);
+ bool allocate(ggml_cgraph* graph, const Measurement& measurement, const AssignNodes& assign_nodes);
Measurement measure(
ggml_cgraph* graph,
const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
diff --git a/src/core/cuda/affine.cu b/src/core/cuda/affine.cu
index 2cd660b..0bd7989 100644
--- a/src/core/cuda/affine.cu
+++ b/src/core/cuda/affine.cu
@@ -29,6 +29,12 @@ namespace SD_CUDA_NAMESPACE {
static bool bounded(const ggml_tensor* tensor) {
return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024;
}
+ static bool cast_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+ const auto* src = node->src[0];
+ return node->src[1] == node && !node->view_src && node->type == GGML_TYPE_F16 && ggml_is_contiguous(node) &&
+ src && src->type == GGML_TYPE_F32 && ggml_are_same_shape(src, node) && !ggml_is_contiguous(src) && src->ne[0] % 4 == 0 && src->nb[0] == sizeof(float) &&
+ (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 && aligned(src, 16) && aligned(node, 4) && bounded(node);
+ }
static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
auto* g = const_cast<ggml_cgraph*>(graph);
if (disabled() || first + 1 >= ggml_graph_n_nodes(g))
@@ -106,15 +112,19 @@ namespace SD_CUDA_NAMESPACE {
dst[2 * i] = __floats2half2_rn(value.x, value.y);
dst[2 * i + 1] = __floats2half2_rn(value.z, value.w);
}
- static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) {
+ template <bool Scaled>
+ static __global__ void cast_strided(const float* src, half2* dst, int64_t dim, int64_t ne1, int64_t ne2, int64_t count, int64_t stride1, int64_t stride2, int64_t stride3, float scale) {
const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4;
if (index >= count)
return;
const int64_t col = index % dim, row = index / dim;
- const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
- const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
- dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
- dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
+ const int64_t i1 = row % ne1, i2 = row / ne1 % ne2, i3 = row / (ne1 * ne2);
+ float4 value = *reinterpret_cast<const float4*>(src + col + i1 * stride1 + i2 * stride2 + i3 * stride3);
+ if constexpr (Scaled) {
+ value = make_float4(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale), __fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
+ }
+ dst[index / 2] = __floats2half2_rn(value.x, value.y);
+ dst[index / 2 + 1] = __floats2half2_rn(value.z, value.w);
}
static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
@@ -129,15 +139,17 @@ namespace SD_CUDA_NAMESPACE {
ggml_ext_gpu_context context;
if (!cuda_context(execution->backend, &context))
return GGML_STATUS_FAILED;
- auto stream = static_cast<CUDAStream>(context.stream);
- const auto* node = nodes[0];
- const auto* a = node->src[0];
- const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1];
- auto* dst = nodes[count - 1];
- const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
- const bool gate = residual != nullptr;
- const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
- (residual && overlaps(dst, residual) && dst->data != residual->data);
+ auto stream = static_cast<CUDAStream>(context.stream);
+ const auto* node = nodes[0];
+ const auto* a = node->src[0];
+ const auto* b = node->op == GGML_OP_CONT || node->op == GGML_OP_CPY ? nullptr : node->src[1];
+ auto* dst = nodes[count - 1];
+ const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
+ const bool gate = residual != nullptr;
+ const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
+ (residual && overlaps(dst, residual) && dst->data != residual->data);
+ if (alias && !execution->workspace)
+ return GGML_STATUS_FAILED;
void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
const int64_t count4 = ggml_nelements(dst) / 4;
const unsigned blocks = unsigned((count4 + 255) / 256);
@@ -150,9 +162,12 @@ namespace SD_CUDA_NAMESPACE {
kernel<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
static_cast<half2*>(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1],
node->ne[permuted ? 1 : 2], count4);
+ } else if (node->op == GGML_OP_CONT) {
+ cast_strided<true><<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
+ a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[2] / 4, a->nb[1] / 4, a->nb[3] / 4, param(nodes[1], 0));
} else {
- scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
- a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0));
+ cast_strided<false><<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
+ a->ne[0], a->ne[1], a->ne[2], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, 1.f);
}
auto error = cuda_last_error();
if (error == cuda_success && alias)
@@ -178,6 +193,12 @@ namespace SD_CUDA_NAMESPACE {
desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
desc.op = GGML_OP_CONT;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.cast_strided_f32_f16." SD_CUDA_NAME ".v1";
+ desc.op = GGML_OP_CPY;
+ desc.match = nullptr;
+ desc.supports = sd_affine_cuda::cast_supports;
+ desc.preferred = sd_affine_cuda::cast_supports;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu
index 26920ae..c069d46 100644
--- a/src/core/cuda/convrot.cu
+++ b/src/core/cuda/convrot.cu
@@ -30,6 +30,7 @@ namespace SD_CUDA_NAMESPACE {
MOD,
RMS_MOD,
SWIGLU,
+ SITU,
SILU,
SIGMOID };
struct Operation {
@@ -81,8 +82,11 @@ namespace SD_CUDA_NAMESPACE {
op.eps = param_f32(first, 0);
if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || (op.shift && !vector(op.shift, op.k)))
return false;
- } else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
- op.kind = SWIGLU;
+ } else if (count == 2 && first->op == GGML_OP_GLU && value == first) {
+ const auto glu = ggml_get_glu_op(first);
+ if (glu != GGML_GLU_OP_SWIGLU && glu != GGML_GLU_OP_SITU)
+ return false;
+ op.kind = glu == GGML_GLU_OP_SITU ? SITU : SWIGLU;
op.up = first->src[1] ? first->src[1] : op.input;
if (!f32(op.up))
return false;
@@ -217,26 +221,26 @@ namespace SD_CUDA_NAMESPACE {
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);
}
- template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false, bool Modulate = false>
+ template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false, bool Modulate = false, bool Situ = false>
static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) {
constexpr bool warp = Scaled || RMS || Modulate;
constexpr int extra = Normalize ? 384 : RMS ? 256
: 128;
if (!p.configured) {
- auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate>,
+ auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate, Situ>,
int(p.device.shared_memory - extra));
if (status != cuda_success)
return status;
p.configured = true;
}
- const int threads = Normalize ? 1024 : RMS ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
- : op.k <= 8192 ? 512
- : 1024)
- : Swiglu && !Sigmoid && op.k >= 16384 ? 512
- : 256;
+ const int threads = Normalize ? 1024 : RMS ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
+ : op.k <= 8192 ? 512
+ : 1024)
+ : (Situ ? op.k >= 8192 : Swiglu && !Sigmoid && op.k >= 16384) ? 512
+ : 256;
const int64_t padded = (op.rows + 3) / 4 * 4;
auto data = [](const ggml_tensor* t) { return t ? static_cast<const float*>(t->data) : nullptr; };
- quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
+ quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS, Modulate, Situ><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
return cuda_last_error();
@@ -274,6 +278,9 @@ namespace SD_CUDA_NAMESPACE {
case SILU:
error = launch<true>(plan, op, output, stream);
break;
+ case SITU:
+ error = launch<true, false, false, false, false, false, true>(plan, op, output, stream);
+ break;
case SIGMOID:
error = launch<true, false, false, true>(plan, op, output, stream);
break;
diff --git a/src/core/cuda/convrot_kernel.h b/src/core/cuda/convrot_kernel.h
index 04abf97..fc0d147 100644
--- a/src/core/cuda/convrot_kernel.h
+++ b/src/core/cuda/convrot_kernel.h
@@ -40,7 +40,13 @@ namespace SD_CUDA_NAMESPACE {
return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
}
- template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false, bool modulate = false>
+ static __device__ __forceinline__ float situ(float gate, float up) {
+ const float a = __fmul_rn(tanhf(__fmul_rn(0.25f, gate)), 4.0f);
+ const float b = __fmul_rn(tanhf(__fmul_rn(1.0f / 25.0f, up)), 25.0f);
+ return __fmul_rn(__fmul_rn(a, 1.0f / (1.0f + expf(-gate))), b);
+ }
+
+ template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false, bool modulate = false, bool situ_gate = false>
__global__ void quantize(
const T* src,
int8_t* dst,
@@ -125,8 +131,12 @@ namespace SD_CUDA_NAMESPACE {
gate_value = src_row[i];
up_value = up[row * up_stride + i];
}
- const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value)));
- v[j] = (activated * up_value) * (1.0f / 16.0f);
+ if constexpr (situ_gate) {
+ v[j] = situ(gate_value, up_value) * (1.0f / 16.0f);
+ } else {
+ const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value)));
+ v[j] = (activated * up_value) * (1.0f / 16.0f);
+ }
}
}
// One logical warp owns H256, keeping every butterfly in registers.
diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu
index 3f753d7..a422432 100644
--- a/src/core/cuda/cute_attention.cu
+++ b/src/core/cuda/cute_attention.cu
@@ -243,7 +243,7 @@ namespace sd_fattn_cute {
bool scaled_output = false;
};
- template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4>
+ template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4, bool QRegs = false>
__global__ __launch_bounds__(Warps * 32) void kernel(Params p) {
using T = Traits<BM, BN, D, Warps>;
using Element = typename T::Element;
@@ -295,6 +295,11 @@ namespace sd_fattn_cute {
copy_rows<false, true, FullDim>(global_copy, last_k, tKsK, tKVcKV, p.nk - block * BN, p.dim);
cute::cp_async_fence();
clear(acc_o);
+ if constexpr (QRegs) {
+ cute::cp_async_wait<0>();
+ __syncthreads();
+ cute::copy(copyQ, tSsQ, threadQ.retile_D(rQ));
+ }
Softmax<2 * size<1>(acc_o)> softmax;
auto step = [&](auto first_flag, int n) {
constexpr bool First = decltype(first_flag)::value;
@@ -305,7 +310,7 @@ namespace sd_fattn_cute {
auto v_tile = tVgV(_, _, _, n);
copy_rows<!First, First, FullDim>(global_copy, v_tile, tVsV, tKVcKV, p.nk - n * BN, p.dim);
cute::cp_async_fence();
- gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
+ gemm<QRegs>(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
if constexpr (First) {
#pragma unroll
for (int i = 0; i < size(acc_s); ++i) {
@@ -681,9 +686,9 @@ static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes,
return GGML_STATUS_SUCCESS;
}
-template <int D, bool FullDim>
+template <int D, bool FullDim, bool QRegs = false>
static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
- auto kernel = sd_fattn_cute::kernel<64, 64, D, FullDim>;
+ auto kernel = sd_fattn_cute::kernel<64, 64, D, FullDim, 4, QRegs>;
constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
static std::mutex mutex;
static std::unordered_map<int, bool> initialized;
@@ -706,11 +711,12 @@ static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const
}
static ggml_status attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
- const auto& plan = *static_cast<AttentionPlan*>(state);
- const auto* attn = nodes[plan.index];
- auto device = ggml_backend_get_device(execution->backend);
- auto reg = ggml_backend_dev_backend_reg(device);
- auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+ const auto& plan = *static_cast<AttentionPlan*>(state);
+ const auto* attn = nodes[plan.index];
+ auto device = ggml_backend_get_device(execution->backend);
+ const auto& properties = cuda_device(device)->properties;
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
ggml_ext_cuda_context context;
if (!query || !query(execution->backend, &context))
return GGML_STATUS_FAILED;
@@ -775,7 +781,10 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
error = attention_launch<112, true>(context, p, q->ne[3]);
break;
case 128:
- error = attention_launch<128>(context, p, q->ne[3]);
+ error = p.dim == 128 && p.nq >= 4096 && p.nq < 8192 && p.nk >= 4096 && p.nk < 8192 &&
+ properties.major == 8 && properties.minor == 9
+ ? attention_launch<128, true, true>(context, p, q->ne[3])
+ : attention_launch<128>(context, p, q->ne[3]);
break;
case 192:
error = attention_launch<192>(context, p, q->ne[3]);
@@ -795,7 +804,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
if (rows != padded_rows && cudaMemsetAsync(packed + width * rows, 0, width * (padded_rows - rows), stream) != cudaSuccess)
return GGML_STATUS_FAILED;
if (cudaFuncSetAttribute(quantize_attention, cudaFuncAttributeMaxDynamicSharedMemorySize,
- int(cuda_device(device)->properties.sharedMemPerBlockOptin - 128)) != cudaSuccess)
+ int(properties.sharedMemPerBlockOptin - 128)) != cudaSuccess)
return GGML_STATUS_FAILED;
quantize_attention<<<rows, 256, width * sizeof(float), stream>>>(work[3], packed,
reinterpret_cast<float*>(packed + width * padded_rows), width, rows, scale, bias);
diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu
index 71477d7..d5764c6 100644
--- a/src/core/cuda/i8_gated_residual.cu
+++ b/src/core/cuda/i8_gated_residual.cu
@@ -35,10 +35,12 @@ namespace sd_i8_gated {
}
};
- template <int M, int N, int WM>
+ template <int M, int N, int WM, int Stages = 4, int Band = 1>
struct TileConfig {
- using Tile = cutlass::gemm::GemmShape<M, N, 64>;
- using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+ using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+ using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+ static constexpr int kStages = Stages;
+ static constexpr int kBand = Band;
};
enum Tile {
@@ -100,7 +102,7 @@ namespace sd_i8_gated {
struct KernelTypes {
using Tile = typename Config::Tile;
using Warp = typename Config::Warp;
- using Swizzle = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+ using Swizzle = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<Config::kBand>>;
using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
@@ -139,7 +141,7 @@ namespace sd_i8_gated {
cutlass::gemm::GemmShape<16, 8, 32>,
Epilogue,
Swizzle,
- 4,
+ Config::kStages,
cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
};
@@ -147,6 +149,7 @@ namespace sd_i8_gated {
struct DeviceInfo {
int ordinal = -1;
int major = 0;
+ int minor = 0;
int sms = 0;
int threads_per_sm = 0;
size_t shared = 0;
@@ -164,7 +167,7 @@ namespace sd_i8_gated {
int ordinal = query ? query(device) : -1;
cudaDeviceProp props = {};
if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
- info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+ info = {ordinal, props.major, props.minor, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
cache.emplace(device, info);
return info;
}
@@ -286,6 +289,10 @@ namespace sd_i8_gated {
return op.rows >= 8192 ? configure<Config, false, true>(op, plan) : configure<Config, false, false>(op, plan);
}
static bool configure(const Operation& op, Plan& plan) {
+ if (plan.device.major == 8 && plan.device.minor == 9 && op.scale == 1.f &&
+ op.rows >= 4096 && op.rows < 8192 && op.n == 3072 && (op.k == 3072 || op.k == 8192) &&
+ configure<TileConfig<128, 128, 64, 3, 4>, false, false>(op, plan))
+ return true;
auto try_tile = [&](Tile tile) {
switch (tile) {
case TILE_128_256:
diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu
index 7c2702d..ee68199 100644
--- a/src/core/cuda/i8_gemm.cu
+++ b/src/core/cuda/i8_gemm.cu
@@ -7,6 +7,7 @@
#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
#include <mutex>
#include <unordered_map>
+#include "convrot_kernel.h"
#include "i8_gemm.h"
namespace sd_i8_cutlass {
@@ -33,10 +34,12 @@ namespace sd_i8_cutlass {
}
};
- template <int M, int N, int WM>
+ template <int M, int N, int WM, int Stages = 4, int Band = 1>
struct TileConfig {
- using Tile = cutlass::gemm::GemmShape<M, N, 64>;
- using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+ using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+ using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+ static constexpr int kStages = Stages;
+ static constexpr int kBand = Band;
};
enum Tile {
@@ -73,6 +76,7 @@ namespace sd_i8_cutlass {
struct DeviceInfo {
int ordinal = -1;
int major = 0;
+ int minor = 0;
int sms = 0;
int threads_per_sm = 0;
size_t shared = 0;
@@ -90,7 +94,7 @@ namespace sd_i8_cutlass {
int ordinal = query ? query(device) : -1;
cudaDeviceProp props = {};
if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
- info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+ info = {ordinal, props.major, props.minor, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
cache.emplace(device, info);
return info;
}
@@ -139,13 +143,30 @@ namespace sd_i8_cutlass {
};
template <class Map>
- struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
+ struct GatedStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
- using Base::Base;
+ struct Arguments : Base::Arguments {
+ bool situ = false;
+ };
+ using Params = Arguments;
+
+ template <class Shape>
+ static constexpr Params to_underlying_arguments(Shape const&, Arguments const& args, void*) {
+ return args;
+ }
+
+ GatedStore() = default;
+ CUTLASS_HOST_DEVICE GatedStore(Params const& params, typename Base::SharedStorage const& shared)
+ : Base(params, shared), situ(params.situ) {}
+
+ bool situ = false;
+
template <class CB>
struct Callbacks : CB {
- CUTLASS_DEVICE Callbacks(CB const& base)
- : CB(base) {}
+ bool situ;
+
+ CUTLASS_DEVICE Callbacks(CB const& base, bool situ)
+ : CB(base), situ(situ) {}
CUTLASS_DEVICE void end_step(int step) {
auto src = cute::filter(this->tC_rAux);
auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
@@ -153,9 +174,10 @@ namespace sd_i8_cutlass {
for (int i = 0; i < cute::size(src); ++i) {
auto coord = coords(i);
if (cute::elem_less(coord, this->problem_shape)) {
- // Each vector contains two gate/up pairs; only their SwiGLU results are stored.
- const float4 a = *reinterpret_cast<const float4*>(&src(i));
- float2 value = make_float2(__fmul_rn(a.x / (1.f + expf(-a.x)), a.y), __fmul_rn(a.z / (1.f + expf(-a.z)), a.w));
+ // SwiGLU packs [gate | up] and SiTU packs [up | gate].
+ const float4 a = *reinterpret_cast<const float4*>(&src(i));
+ const float2 value = situ ? make_float2(sd_cuda::sd_convrot::situ(a.y, a.x), sd_cuda::sd_convrot::situ(a.w, a.z))
+ : make_float2(__fmul_rn(a.x / (1.f + expf(-a.x)), a.y), __fmul_rn(a.z / (1.f + expf(-a.z)), a.w));
int64_t n = cute::get<1>(this->problem_shape), m = cute::get<0>(this->problem_shape);
int64_t offset = cute::get<0>(coord) * (n / 2) + cute::get<1>(coord) / 2 + cute::get<2>(coord) * m * (n / 2);
*reinterpret_cast<float2*>(this->params_ptr->ptr_aux + offset) = value;
@@ -166,7 +188,7 @@ namespace sd_i8_cutlass {
template <class Shape>
CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
auto cb = Base::get_callbacks(tile, tid, shape);
- return Callbacks<decltype(cb)>(cb);
+ return Callbacks<decltype(cb)>(cb, situ);
}
};
@@ -191,6 +213,7 @@ namespace sd_i8_cutlass {
template <class Config, bool StreamK, bool PackedW4 = false>
struct DenseTypes {
static constexpr bool kSplit = false;
+ static constexpr bool kGated = false;
static constexpr bool kBias = !PackedW4;
using Weight = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
using Tile = typename Config::Tile;
@@ -232,22 +255,23 @@ namespace sd_i8_cutlass {
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
};
template <class Config, bool StreamK, bool Split = false>
- struct SwigluTypes {
+ struct GatedTypes {
static constexpr bool kSplit = Split;
static constexpr bool kBias = false;
+ static constexpr bool kGated = true;
using Weight = int8_t;
using Tile = typename Config::Tile;
using Warp = typename Config::Warp;
- using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+ using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<Config::kBand>>;
using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
- using Store = SwigluStore<Map>;
+ using Store = GatedStore<Map>;
using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
- using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
+ using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, Config::kStages, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
using BaseMma = typename GemmBase::Mma;
using SplitMma = cutlass::gemm::threadblock::MmaMultistage<
@@ -271,7 +295,7 @@ namespace sd_i8_cutlass {
struct Plan {
DeviceInfo device;
int64_t rows, n, k;
- bool swiglu;
+ SDI8GemmActivation activation;
bool packed_w4;
bool split;
size_t kernel_bytes = 0, scratch_offset = 0;
@@ -285,6 +309,8 @@ namespace sd_i8_cutlass {
if constexpr (Types::kBias) {
// A negative zero default preserves signed-zero outputs.
epilogue = {{scales, {bias, -0.f, {}}, {}}, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+ } else if constexpr (Types::kGated) {
+ epilogue = {scales, {{out, {p.n, cute::_1{}, p.rows * p.n}}, p.activation == SDI8GemmActivation::Situ}};
} else {
epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
}
@@ -298,7 +324,7 @@ namespace sd_i8_cutlass {
16)
return GGML_STATUS_FAILED;
auto* paired = static_cast<float*>(scratch);
- auto args = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out, bias);
+ auto args = arguments<Types>(p, x, w, xs, p.activation != SDI8GemmActivation::None ? paired : ws, out, bias);
if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
return GGML_STATUS_FAILED;
if (!p.occupancy) {
@@ -318,7 +344,7 @@ namespace sd_i8_cutlass {
}
if (params.get_workspace_size() > p.kernel_bytes)
return GGML_STATUS_FAILED;
- if (p.swiglu) {
+ if (p.activation != SDI8GemmActivation::None) {
pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, up_scales, paired, int(p.n));
if (cudaGetLastError() != cudaSuccess)
return GGML_STATUS_FAILED;
@@ -346,7 +372,7 @@ namespace sd_i8_cutlass {
bytes = std::max(bytes, params.get_workspace_size());
}
p.kernel_bytes = bytes;
- p.scratch_offset = p.swiglu ? (size_t(p.n) * sizeof(float) + 255) / 256 * 256 : 0;
+ p.scratch_offset = p.activation != SDI8GemmActivation::None ? (size_t(p.n) * sizeof(float) + 255) / 256 * 256 : 0;
p.launch = launch<Types>;
return true;
}
@@ -355,12 +381,17 @@ namespace sd_i8_cutlass {
if (p.packed_w4)
return configure_kernel<DenseTypes<Config, false, true>>(p);
if (p.split)
- return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true, true>>(p) : configure_kernel<SwigluTypes<Config, false, true>>(p);
- if (p.swiglu)
- return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true>>(p) : configure_kernel<SwigluTypes<Config, false>>(p);
+ return p.rows >= 8192 ? configure_kernel<GatedTypes<Config, true, true>>(p) : configure_kernel<GatedTypes<Config, false, true>>(p);
+ if (p.activation != SDI8GemmActivation::None)
+ return p.rows >= 8192 ? configure_kernel<GatedTypes<Config, true>>(p) : configure_kernel<GatedTypes<Config, false>>(p);
return p.rows >= 8192 ? configure_kernel<DenseTypes<Config, true>>(p) : configure_kernel<DenseTypes<Config, false>>(p);
}
static bool configure(Plan& p) {
+ // This SiTU shape fits two resident CTAs per SM with the narrower, three-stage tile on Ada.
+ if (p.device.major == 8 && p.device.minor == 9 && p.activation == SDI8GemmActivation::Situ && !p.split &&
+ p.rows >= 4096 && p.rows < 8192 && p.n == 16384 && p.k == 3072 &&
+ configure_kernel<GatedTypes<TileConfig<128, 128, 64, 3, 4>, false>>(p))
+ return true;
auto try_tile = [&](Tile tile) {
switch (tile) {
case TILE_128_256:
@@ -384,22 +415,23 @@ namespace sd_i8_cutlass {
}
}
struct SDI8CudaGemmPlan : sd_i8_cutlass::Plan {};
-static SDI8CudaGemmPlan* prepare_gemm(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4, bool split) {
+static SDI8CudaGemmPlan* prepare_gemm(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, SDI8GemmActivation activation, size_t* workspace, bool packed_w4, bool split) {
#ifdef SD_CUDA_FORCE_CUBLAS
return nullptr;
#else
- const auto info = sd_i8_cutlass::device_info(device);
- if ((packed_w4 && (swiglu || k > INT_MAX / 1024)) || info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
- rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128))
+ const auto info = sd_i8_cutlass::device_info(device);
+ const bool gated = activation != SDI8GemmActivation::None;
+ if ((packed_w4 && (gated || k > INT_MAX / 1024)) || info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
+ rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (gated ? 256 : 128))
return nullptr;
- auto* p = new SDI8CudaGemmPlan;
- p->device = info;
- p->rows = rows;
- p->n = n;
- p->k = k;
- p->swiglu = swiglu;
- p->packed_w4 = packed_w4;
- p->split = split;
+ auto* p = new SDI8CudaGemmPlan;
+ p->device = info;
+ p->rows = rows;
+ p->n = n;
+ p->k = k;
+ p->activation = activation;
+ p->packed_w4 = packed_w4;
+ p->split = split;
if (!sd_i8_cutlass::configure(*p)) {
delete p;
return nullptr;
@@ -408,16 +440,16 @@ static SDI8CudaGemmPlan* prepare_gemm(ggml_backend_dev_t device, int64_t rows, i
return p;
#endif
}
-SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4) {
- return prepare_gemm(device, rows, n, k, swiglu, workspace, packed_w4, false);
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, SDI8GemmActivation activation, size_t* workspace, bool packed_w4) {
+ return prepare_gemm(device, rows, n, k, activation, workspace, packed_w4, false);
}
SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace) {
if (n > INT_MAX / 2 || k > INT_MAX / 16384)
return nullptr;
- return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true);
+ return prepare_gemm(device, rows, 2 * n, k, SDI8GemmActivation::Swiglu, workspace, false, true);
}
ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const float* bias) {
- if (bias && (plan->swiglu || plan->packed_w4))
+ if (bias && (plan->activation != SDI8GemmActivation::None || plan->packed_w4))
return GGML_STATUS_FAILED;
return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr, bias);
}
diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h
index 756238e..db9bfe7 100644
--- a/src/core/cuda/i8_gemm.h
+++ b/src/core/cuda/i8_gemm.h
@@ -3,7 +3,10 @@
#include <cuda_runtime.h>
#include "ggml-ext.h"
struct SDI8CudaGemmPlan;
-SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false);
+enum class SDI8GemmActivation { None,
+ Swiglu,
+ Situ };
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, SDI8GemmActivation activation, size_t* workspace, bool packed_w4 = false);
ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const float* bias = nullptr);
// n is the width of each branch; gate/up weights and scales remain separate.
SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace);
diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu
index 8e92468..13203b5 100644
--- a/src/core/cuda/int8.cu
+++ b/src/core/cuda/int8.cu
@@ -438,7 +438,7 @@ namespace SD_CUDA_NAMESPACE {
}
enum Kind { MAT,
QUANTIZE,
- SWIGLU,
+ GATED,
SWIGLU_QUANTIZE,
CONCAT };
static bool fusion_disabled() {
@@ -465,10 +465,25 @@ namespace SD_CUDA_NAMESPACE {
next->type == GGML_TYPE_F32 && mat->ne[2] == 1 && mat->ne[3] == 1 &&
next->ne[2] == 1 && next->ne[3] == 1 && ggml_is_contiguous(next) && ggml_is_contiguous(next->src[0]))
return 2;
- if (next->op != GGML_OP_GLU || ggml_get_glu_op(next) != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0 ||
+ if (next->op != GGML_OP_GLU ||
next->src[0] != mat || next->src[1] || next->type != GGML_TYPE_F32 || !ggml_is_contiguous(next) ||
ggml_nelements(next) > int64_t(INT_MAX) * 256)
return 0;
+ const auto activation = ggml_get_glu_op(next);
+ if (activation == GGML_GLU_OP_SITU && next->op_params[1] != 0) {
+#ifdef SD_CUDA_I8_CUTLASS
+ if (!next->data && !next->buffer && !next->view_src && cutlass_compatible(op, next)) {
+ size_t workspace = 0;
+ auto* gemm = sd_cuda_i8_gemm_prepare(device, ggml_nrows(mat), mat->ne[0], op.weight->ne[0], SDI8GemmActivation::Situ, &workspace);
+ const bool supported = gemm != nullptr;
+ sd_cuda_i8_gemm_free(gemm);
+ return supported ? 2 : 0;
+ }
+#endif
+ return 0;
+ }
+ if (activation != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0)
+ return 0;
if (first + 2 < ggml_graph_n_nodes(g)) {
const auto* quant = ggml_graph_node(g, first + 2);
int group;
@@ -477,7 +492,7 @@ namespace SD_CUDA_NAMESPACE {
#ifdef SD_CUDA_I8_CUTLASS
if (!next->data && !next->buffer && !next->view_src && cutlass_compatible(op, next)) {
size_t workspace = 0;
- auto* gemm = sd_cuda_i8_gemm_prepare(device, ggml_nrows(mat), mat->ne[0], op.weight->ne[0], true, &workspace);
+ auto* gemm = sd_cuda_i8_gemm_prepare(device, ggml_nrows(mat), mat->ne[0], op.weight->ne[0], SDI8GemmActivation::Swiglu, &workspace);
const bool fused_epilogue = gemm != nullptr;
sd_cuda_i8_gemm_free(gemm);
// ConvRot is a separate launch either way. Keep its input in the
@@ -498,7 +513,8 @@ namespace SD_CUDA_NAMESPACE {
int64_t k, n, rows, padded;
size_t quant, scales, accum, partial, output, gemm_offset;
SDI8CudaGemmPlan* gemm = nullptr;
- bool fast_swiglu = false;
+ bool fused_activation = false;
+ bool situ = false;
};
static size_t reserve(size_t& total, size_t bytes) {
const size_t offset = total;
@@ -516,12 +532,13 @@ namespace SD_CUDA_NAMESPACE {
p->kind = !mat ? QUANTIZE : count == 3 ? SWIGLU_QUANTIZE
: count == 1 ? MAT
: dst->op == GGML_OP_CONCAT ? CONCAT
- : SWIGLU;
+ : GATED;
p->k = mat ? op.weight->ne[0] : dst->ne[0];
p->n = mat ? nodes[0]->ne[0] : 0;
p->rows = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]);
p->padded = (p->rows + 3) / 4 * 4;
p->group = mat ? op.group_size : 0;
+ p->situ = p->kind == GATED && ggml_get_glu_op(dst) == GGML_GLU_OP_SITU;
if (!mat)
sd_ggml_convrot_describe(dst, &p->group);
size_t total = 0;
@@ -533,17 +550,18 @@ namespace SD_CUDA_NAMESPACE {
// Retain cuBLAS when I32 accumulation can overflow.
if (cutlass_compatible(op, dst) && (!op.bias || p->kind == MAT) &&
(p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) {
- if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) {
- p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes);
- p->fast_swiglu = p->gemm != nullptr;
+ if (p->kind == GATED || p->kind == SWIGLU_QUANTIZE) {
+ const auto activation = p->situ ? SDI8GemmActivation::Situ : SDI8GemmActivation::Swiglu;
+ p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, activation, &gemm_bytes);
+ p->fused_activation = p->gemm != nullptr;
}
if (!p->gemm)
- p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, false, &gemm_bytes);
+ p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, SDI8GemmActivation::None, &gemm_bytes);
}
p->gemm_offset = reserve(total, gemm_bytes);
#endif
- const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == SWIGLU && !p->fast_swiglu);
- p->accum = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
+ const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == GATED && !p->fused_activation);
+ p->accum = reserve(total, intermediate ? size_t(p->fused_activation ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
}
const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
p->partial = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float)
@@ -705,7 +723,7 @@ namespace SD_CUDA_NAMESPACE {
return cuda_last_error();
}
template <bool Dequantize>
- static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) {
+ static __global__ void gated_activation(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows, bool situ) {
const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
const int64_t half = n / 2;
if (i >= rows * half)
@@ -721,7 +739,7 @@ namespace SD_CUDA_NAMESPACE {
gate = src[row * n + col];
up = src[row * n + half + col];
}
- dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up);
+ dst[i] = situ ? sd_convrot::situ(up, gate) : __fmul_rn(gate / (1.f + expf(-gate)), up);
}
template <bool Swiglu, bool Dequantize, bool Rotated>
static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) {
@@ -815,7 +833,7 @@ namespace SD_CUDA_NAMESPACE {
const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0;
auto* gemm_output = reinterpret_cast<float*>(static_cast<char*>(output) + prefix);
auto* accum = reinterpret_cast<int32_t*>(scratch_at(p.accum));
- const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu);
+ const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == GATED && !p.fused_activation);
#ifdef SD_CUDA_I8_CUTLASS
if (p.gemm) {
auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
@@ -856,18 +874,18 @@ namespace SD_CUDA_NAMESPACE {
}
const auto* values = reinterpret_cast<const float*>(accum);
if (p.kind == SWIGLU_QUANTIZE) {
- if (p.fast_swiglu)
+ if (p.fused_activation)
error = quantize_epilogue<false, false, true>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
else if (p.gemm)
error = quantize_epilogue<true, false, false>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
else
error = quantize_epilogue<true, true, false>(p, values, scales, ws, static_cast<int8_t*>(output), stream);
- } else if (p.kind == SWIGLU && !p.fast_swiglu) {
+ } else if (p.kind == GATED && !p.fused_activation) {
const unsigned blocks = unsigned((p.rows * (p.n / 2) + 255) / 256);
if (p.gemm)
- swiglu<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+ gated_activation<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows, p.situ);
else
- swiglu<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+ gated_activation<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows, p.situ);
error = cuda_last_error();
} else if (!p.gemm) {
const unsigned gy = unsigned(std::min<int64_t>(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535);
diff --git a/src/core/cuda/w4.cu b/src/core/cuda/w4.cu
index 98f810e..5ccb3ee 100644
--- a/src/core/cuda/w4.cu
+++ b/src/core/cuda/w4.cu
@@ -856,7 +856,7 @@ namespace SD_CUDA_NAMESPACE {
#ifdef SD_CUDA_W4_CUTLASS
if (plan->mma && op.kind == SDW4Kind::W4A4) {
if (!op.bias && aligned(dst, 16) && aligned(op.weight_scale, 16))
- plan->gemm = sd_cuda_i8_gemm_prepare(ggml_backend_get_device(backend), ggml_nrows(dst), dst->ne[0], op.input->ne[0], false, &gemm_bytes, true);
+ plan->gemm = sd_cuda_i8_gemm_prepare(ggml_backend_get_device(backend), ggml_nrows(dst), dst->ne[0], op.input->ne[0], SDI8GemmActivation::None, &gemm_bytes, true);
plan->mma = plan->gemm != nullptr;
if (plan->gemm && gemm_bytes == 0)
gemm_bytes = 256;
diff --git a/src/core/execution_engine.cpp b/src/core/execution_engine.cpp
index 5ad241a..867fc73 100644
--- a/src/core/execution_engine.cpp
+++ b/src/core/execution_engine.cpp
@@ -84,9 +84,12 @@ namespace sd {
if (sd_backend_is(backend, "Vulkan")) {
return true;
}
- auto buft = ggml_backend_get_default_buffer_type(backend);
+ auto buft = ggml_backend_get_default_buffer_type(backend);
+ size_t max_size = ggml_backend_buft_get_max_size(buft);
+ if (max_size == SIZE_MAX) {
+ return true;
+ }
size_t alignment = ggml_backend_buft_get_alignment(buft);
- size_t max_size = ggml_backend_buft_get_max_size(buft);
max_size -= max_size % alignment;
if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) {
return false;
@@ -159,19 +162,21 @@ namespace sd {
if (sched == nullptr || gf == nullptr) {
return;
}
- ggml_backend_t current = session_.backend_;
- const int n_nodes = ggml_graph_n_nodes(gf);
+ ggml_backend_t current = session_.backend_;
+ const int n_nodes = ggml_graph_n_nodes(gf);
+ const bool multi_device = is_multi_device();
for (int i = 0; i < n_nodes; i++) {
- ggml_tensor* node = ggml_graph_node(gf, i);
- auto node_assignment = session_.node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
- if (node_assignment != session_.node_assignments_.end()) {
- current = node_assignment->second;
- }
- for (int s = 0; s < GGML_MAX_SRC; s++) {
- ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
- if (weight_backend != nullptr) {
- if (node_assignment == session_.node_assignments_.end()) {
- current = weight_backend;
+ ggml_tensor* node = ggml_graph_node(gf, i);
+ if (multi_device) {
+ auto node_assignment = session_.node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
+ if (node_assignment != session_.node_assignments_.end()) {
+ current = node_assignment->second;
+ } else {
+ for (auto* source : node->src) {
+ auto weight_backend = backend_for_weight(source);
+ if (weight_backend != nullptr) {
+ current = weight_backend;
+ }
}
}
}
@@ -665,7 +670,7 @@ namespace sd {
return fail_segment("output cache allocation");
}
new_cache_bytes = add_bytes(last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment), session_.cache_.pending_bytes(segment_graph));
- if (!session_.workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
+ if (!session_.workspace_.allocate(segment_graph, measurement, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
pin_multi_device_nodes(scheduler, current);
})) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
diff --git a/src/core/pe_cache.cpp b/src/core/pe_cache.cpp
index df4aa62..3f53a0a 100644
--- a/src/core/pe_cache.cpp
+++ b/src/core/pe_cache.cpp
@@ -20,11 +20,17 @@ namespace sd {
}
const Tensor<float>& PECache::get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build) {
- auto entry = entries_.find(key.name());
- if (entry != entries_.end() && entry->second.key == key) {
- return entry->second.value;
+ auto& entries = entries_[key.name()];
+ for (auto entry = entries.begin(); entry != entries.end(); ++entry) {
+ if (entry->key == key) {
+ entries.splice(entries.begin(), entries, entry);
+ return entries.front().value;
+ }
+ }
+ entries.push_front({key, build()});
+ if (entries.size() > entries_per_name_) {
+ entries.pop_back();
}
- auto value = build();
- return entries_.insert_or_assign(key.name(), Entry{key, std::move(value)}).first->second.value;
+ return entries.front().value;
}
}
diff --git a/src/core/pe_cache.h b/src/core/pe_cache.h
index 3539707..90b5d3a 100644
--- a/src/core/pe_cache.h
+++ b/src/core/pe_cache.h
@@ -2,6 +2,7 @@
#define __SD_CORE_PE_CACHE_H__
#include <functional>
+#include <list>
#include <map>
#include <set>
#include <type_traits>
@@ -71,9 +72,14 @@ namespace sd {
Tensor<float> value;
};
- std::map<std::string, Entry> entries_;
+ std::map<std::string, std::list<Entry>> entries_;
+ size_t entries_per_name_;
public:
+ explicit PECache(size_t entries_per_name = 1)
+ : entries_per_name_(entries_per_name) {
+ GGML_ASSERT(entries_per_name_ > 0);
+ }
const Tensor<float>& get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build);
};
diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp
index bcca79b..4bb9077 100644
--- a/src/model/diffusion/moss_image.hpp
+++ b/src/model/diffusion/moss_image.hpp
@@ -9,7 +9,6 @@
#include "model/common/ggml_block.hpp"
#include "model/common/rope.hpp"
-#include "model/diffusion/dit.hpp"
#include "model/diffusion/model.hpp"
#include "model_loader.h"
@@ -164,9 +163,10 @@ namespace MOSSImage {
auto c = ctx->ggml_ctx;
auto linear = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, h); };
int64_t image_len = image->ne[0] * image->ne[1];
- auto x = linear("video_embedder", DiT::patchify(ctx, image, 1, 1));
+ auto pixels = ggml_reshape_3d(c, image, image_len, image->ne[2], image->ne[3]);
+ auto x = linear("video_embedder", ggml_cont(c, ggml_transpose(c, pixels)));
text = linear("text_embedder.proj", text);
- auto text_pe = ggml_ext_slice(ctx, pe, 3, image_len, image_len + text->ne[1]);
+ auto text_pe = ggml_ext_slice(ctx, pe, 3, image_len, image_len + text->ne[1], false);
for (int i = 0; i < 2; ++i) {
text = std::dynamic_pointer_cast<TextBlock>(blocks["text_embedder.blocks." + std::to_string(i)])->forward(ctx, text, text_pe);
}
@@ -175,7 +175,7 @@ namespace MOSSImage {
auto head = linear("video_head_modulation", time);
std::array<ggml_tensor*, 6> parts;
for (int i = 0; i < 6; ++i) {
- parts[i] = ggml_ext_slice(ctx, mod, 0, i * config.dim, (i + 1) * config.dim);
+ parts[i] = ggml_ext_slice(ctx, mod, 0, i * config.dim, (i + 1) * config.dim, false);
if (ctx->graph != nullptr) {
ggml_build_forward_expand(ctx->graph, parts[i]);
}
@@ -187,18 +187,18 @@ namespace MOSSImage {
sd::ggml_graph_cut::mark_graph_cut(x, "moss_image.layers." + std::to_string(i), "x");
}
}
- x = ggml_ext_slice(ctx, x, 1, 0, image_len);
+ x = ggml_ext_slice(ctx, x, 1, 0, image_len, false);
x = std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm_video"])->forward(ctx, x);
- auto shift = ggml_ext_slice(ctx, head, 0, 0, config.dim);
- auto scale = ggml_ext_slice(ctx, head, 0, config.dim, 2 * config.dim);
+ auto shift = ggml_ext_slice(ctx, head, 0, 0, config.dim, false);
+ auto scale = ggml_ext_slice(ctx, head, 0, config.dim, 2 * config.dim, false);
x = ggml_ext_modulate(ctx, x, shift, scale);
x = linear("final_linear_video", x);
- return DiT::unpatchify(ctx, x, image->ne[1], image->ne[0], 1, 1);
+ return ggml_cont_4d(c, ggml_transpose(c, x), image->ne[0], image->ne[1], image->ne[2], image->ne[3]);
}
};
struct MOSSImageRunner : public DiffusionModelRunner {
- sd::PECache pe_cache_;
+ sd::PECache pe_cache_{2};
MOSSImageConfig config;
Model model;
std::vector<float> time_data;