upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-29 01:18:05 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add LLaDA-Image support (#1968)

Browse Source 
Co-authored-by: leejet <leejet714@gmail.com>
This commit is contained in:
fszontagh
2026-09-21 01:14:20 +08:00
committed by GitHub
co-authored by leejet
parent b8248a869c
commit 15f335daa5
24 changed files with 2224 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/model/common/ggml_block.hpp
+12 -3
View File
@@ -839,21 +839,30 @@ class RMSNorm : public UnaryBlock {
protected:
int64_t hidden_size;
float eps;
bool elementwise_affine;
std::string prefix;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
this->prefix = prefix;
this->prefix = prefix;
if (!elementwise_affine) {
return;
}
enum ggml_type wtype = GGML_TYPE_F32;
params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
}
public:
RMSNorm(int64_t hidden_size,
float eps = 1e-06f)
float eps = 1e-06f,
bool elementwise_affine = true)
: hidden_size(hidden_size),
eps(eps) {}
eps(eps),
elementwise_affine(elementwise_affine) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
if (!elementwise_affine) {
return ggml_rms_norm(ctx->ggml_ctx, x, eps);
}
ggml_tensor* w = params["weight"];
if (ctx->weight_adapter) {
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
src/model/common/rope.hpp
+139
View File
@@ -929,6 +929,145 @@ namespace Rope {
return ids;
}
// LLaDA-Image shares Lumina2/z_image's axes layout, but assigns position (0,0,0) to the
// padding slots of the caption stream instead of continuing the caption ramp through them.
__STATIC_INLINE__ std::vector<std::vector<float>> gen_llada_image_ids(int h,
int w,
int patch_size,
int bs,
int context_len,
int seq_multi_of) {
int context_pad_len = bound_mod(context_len, seq_multi_of);
int padded_context_len = context_len + context_pad_len;
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
for (int i = 0; i < bs * padded_context_len; i++) {
int pos = i % padded_context_len;
if (pos < context_len) {
txt_ids[i][0] = pos + 1.f;
}
}
int axes_dim_num = 3;
int index = padded_context_len + 1;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
img_ids = concat_ids(img_ids, img_pad_ids, bs);
}
return concat_ids(txt_ids, img_ids, bs);
}
// LLaDA-Image editing packs two caption copies (clean and noisy), the source and target
// latents anchored at their own caption's end position, and the SigVQ stream after both.
// Padding slots keep position (0,0,0), as in the text-only layout.
__STATIC_INLINE__ std::vector<std::vector<float>> gen_llada_image_edit_ids(int h,
int w,
int patch_size,
int context_len,
int sigvq_len,
int seq_multi_of) {
const int context_pad = bound_mod(context_len, seq_multi_of);
const int padded_context = context_len + context_pad;
const int h_len = (h + (patch_size / 2)) / patch_size;
const int w_len = (w + (patch_size / 2)) / patch_size;
const int image_len = h_len * w_len;
const int image_pad = bound_mod(image_len, seq_multi_of);
const int padded_image = image_len + image_pad;
const int sigvq_pad = bound_mod(sigvq_len, seq_multi_of);
std::vector<std::vector<float>> cap_ids;
std::vector<int> cap_end_positions;
int cursor = 1;
for (int copy = 0; copy < 2; ++copy) {
for (int i = 0; i < padded_context; ++i) {
std::vector<float> id(3, 0.f);
if (i < context_len) {
id[0] = static_cast<float>(cursor + i);
}
cap_ids.push_back(id);
}
cursor += context_len;
cap_end_positions.push_back(cursor);
cursor += 2;
}
std::vector<std::vector<float>> img_ids;
for (int copy = 0; copy < 2; ++copy) {
auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy]);
img_ids.insert(img_ids.end(), ids.begin(), ids.end());
img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f));
}
const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
std::vector<std::vector<float>> sigvq_ids;
for (int i = 0; i < sigvq_len + sigvq_pad; ++i) {
std::vector<float> id(3, 0.f);
if (i < sigvq_len) {
id[0] = static_cast<float>(sigvq_start + i);
}
sigvq_ids.push_back(id);
}
std::vector<std::vector<float>> ids;
ids.reserve(cap_ids.size() + img_ids.size() + sigvq_ids.size());
ids.insert(ids.end(), cap_ids.begin(), cap_ids.end());
ids.insert(ids.end(), img_ids.begin(), img_ids.end());
ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end());
SD_UNUSED(padded_image);
return ids;
}
__STATIC_INLINE__ std::vector<float> gen_llada_image_edit_pe(int h,
int w,
int patch_size,
int context_len,
int sigvq_len,
int seq_multi_of,
int theta,
const std::vector<int>& axes_dim) {
auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of);
return embed_nd(ids, 1, static_cast<float>(theta), axes_dim, {});
}
__STATIC_INLINE__ std::vector<float> gen_llada_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int seq_multi_of,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len + bound_mod(context_len, seq_multi_of);
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
}
// Generate z_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
int w,
src/model/diffusion/llada_image.hpp
+527
View File
@@ -0,0 +1,527 @@
#ifndef __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
#define __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
#include <algorithm>
#include <cinttypes>
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
#include "core/util.h"
#include "model/common/ggml_block.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/z_image.hpp"
#include "model_loader.h"
// Ref: https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py
//
// The denoiser is Lumina2/z_image's NextDiT with identical hyperparameters, so the blocks are
// reused from ZImage. Two things differ: every norm here is non-parametric (the checkpoint
// carries no norm weights at all), and latents arrive already patchified from the Flux2 VAE,
// so patch_size is 1 over 128 channels.
namespace LLaDAImage {
constexpr int LLADA_IMAGE_GRAPH_SIZE = 20480;
struct LLaDAImageConfig {
int patch_size = 1;
int64_t hidden_size = 3840;
int64_t in_channels = 128;
int64_t out_channels = 128;
int64_t num_layers = 30;
int64_t num_refiner_layers = 2;
int64_t head_dim = 128;
int64_t num_heads = 30;
int64_t num_kv_heads = 30;
int64_t multiple_of = 256;
float ffn_dim_multiplier = 8.0f / 3.0f;
float norm_eps = 1e-5f;
bool qk_norm = true;
int64_t cap_feat_dim = 2560;
int64_t semantic_feat_dim = 4096;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int64_t axes_dim_sum = 128;
static int64_t count_blocks(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
const std::string& block_prefix) {
int64_t count = 0;
for (const auto& [name, _] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
size_t pos = name.find(block_prefix);
if (pos == std::string::npos) {
continue;
}
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
count = std::max<int64_t>(count, atoi(items[1].c_str()) + 1);
}
}
return count;
}
static LLaDAImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
LLaDAImageConfig config;
int64_t detected_q_dim = 0;
int64_t detected_kv_dim = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.in_channels = tensor_storage.ne[0] / patch_area;
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.cap_feat_dim = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.semantic_feat_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) {
detected_q_dim = tensor_storage.ne[1];
} else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) {
detected_kv_dim = tensor_storage.ne[1];
} else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.out_channels = tensor_storage.ne[1] / patch_area;
}
}
int64_t detected_layers = count_blocks(tensor_storage_map, prefix, "layers.");
int64_t detected_refiner = std::max(count_blocks(tensor_storage_map, prefix, "noise_refiner."),
count_blocks(tensor_storage_map, prefix, "context_refiner."));
if (detected_layers > 0) {
config.num_layers = detected_layers;
}
if (detected_refiner > 0) {
config.num_refiner_layers = detected_refiner;
}
if (detected_q_dim > 0) {
config.num_heads = detected_q_dim / config.head_dim;
}
if (detected_kv_dim > 0) {
config.num_kv_heads = detected_kv_dim / config.head_dim;
} else if (detected_q_dim > 0) {
config.num_kv_heads = config.num_heads;
}
LOG_VERBOSE("llada_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64 ", cap_feat_dim = %" PRId64 ", semantic_feat_dim = %" PRId64,
config.num_layers,
config.num_refiner_layers,
config.hidden_size,
config.num_heads,
config.num_kv_heads,
config.in_channels,
config.out_channels,
config.cap_feat_dim,
config.semantic_feat_dim);
return config;
}
};
class LLaDAImageModel : public GGMLBlock {
protected:
LLaDAImageConfig config;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
params["sigvq_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
}
std::shared_ptr<ZImage::JointTransformerBlock> make_block(bool modulation) {
return std::make_shared<ZImage::JointTransformerBlock>(0,
config.hidden_size,
config.head_dim,
config.num_heads,
config.num_kv_heads,
config.multiple_of,
config.ffn_dim_multiplier,
config.norm_eps,
config.qk_norm,
modulation,
false,
true);
}
public:
LLaDAImageModel() = default;
LLaDAImageModel(LLaDAImageConfig config)
: config(config) {
blocks["x_embedder"] = std::make_shared<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, ZImage::ADALN_EMBED_DIM);
blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps, false);
blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
blocks["semantic_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
blocks["semantic_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
blocks["sigvq_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
blocks["sigvq_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
for (int i = 0; i < config.num_refiner_layers; i++) {
blocks["noise_refiner." + std::to_string(i)] = make_block(true);
blocks["context_refiner." + std::to_string(i)] = make_block(false);
blocks["sigvq_refiner." + std::to_string(i)] = make_block(false);
}
for (int i = 0; i < config.num_layers; i++) {
blocks["layers." + std::to_string(i)] = make_block(true);
}
blocks["final_layer"] = std::make_shared<ZImage::FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
}
ggml_tensor* forward_core(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
auto final_layer = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
auto txt_pad_token = params["cap_pad_token"];
auto img_pad_token = params["x_pad_token"];
int64_t N = x->ne[2];
int64_t n_img_token = x->ne[1];
int64_t n_txt_token = context->ne[1];
// sdcpp's flow denoiser already hands over sigma * 1000, which is the range the
// reference reaches via its own t_scale, so no further scaling here.
auto t_emb = t_embedder->forward(ctx, timestep);
auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size]
auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.prelude", "txt");
sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.prelude", "img");
sd::ggml_graph_cut::mark_graph_cut(t_emb, "llada_image.prelude", "t_emb");
int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), ZImage::SEQ_MULTI_OF);
if (n_txt_pad_token > 0) {
auto txt_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, txt_pad_token, txt_pad_token->ne[0], n_txt_pad_token, N, 1);
txt = ggml_concat(ctx->ggml_ctx, txt, txt_pad_tokens, 1);
}
int64_t n_img_pad_token = Rope::bound_mod(static_cast<int>(n_img_token), ZImage::SEQ_MULTI_OF);
if (n_img_pad_token > 0) {
auto img_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_img_pad_token, N, 1);
img = ggml_concat(ctx->ggml_ctx, img, img_pad_tokens, 1);
}
GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.context_refiner." + std::to_string(i), "txt");
}
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
img = block->forward(ctx, img, img_pe, nullptr, t_emb);
sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.noise_refiner." + std::to_string(i), "img");
}
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.prelude", "txt_img");
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.layers." + std::to_string(i), "txt_img");
}
txt_img = final_layer->forward(ctx, txt_img, t_emb);
return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
}
ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) {
int64_t n_pad = Rope::bound_mod(static_cast<int>(x->ne[1]), ZImage::SEQ_MULTI_OF);
if (n_pad == 0) {
return x;
}
auto pads = ggml_repeat_4d(ctx->ggml_ctx, pad_token, pad_token->ne[0], n_pad, x->ne[2], 1);
return ggml_concat(ctx->ggml_ctx, x, pads, 1);
}
// Editing runs one joint sequence carrying two timesteps: the caption and source latent
// are clean (t = 0) while the second caption copy and the target latent are noisy. adaLN
// is a linear map of the timestep embedding, so feeding a per-token embedding selects the
// right modulation exactly, without duplicating the modulation projections.
ggml_tensor* forward_editing(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* semantic,
ggml_tensor* source_latent,
ggml_tensor* pe) {
ggml_context* gctx = ctx->ggml_ctx;
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
auto sigvq_embed_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["sigvq_embedder.0"]);
auto sigvq_embed_1 = std::dynamic_pointer_cast<Linear>(blocks["sigvq_embedder.1"]);
auto final_layer = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
auto t_noisy = t_embedder->forward(ctx, timestep);
auto t_clean = t_embedder->forward(ctx, ggml_scale(gctx, timestep, 0.f));
auto per_token = [&](ggml_tensor* emb, int64_t n) {
return ggml_repeat_4d(gctx, emb, emb->ne[0], n, 1, 1);
};
auto cap = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context));
cap = pad_stream(ctx, cap, params["cap_pad_token"]);
int64_t cap_len = cap->ne[1];
cap = ggml_concat(gctx, cap, cap, 1);
auto src = pad_stream(ctx, x_embedder->forward(ctx, source_latent), params["x_pad_token"]);
auto tgt_embed = x_embedder->forward(ctx, x);
int64_t n_img_token = tgt_embed->ne[1];
auto tgt = pad_stream(ctx, tgt_embed, params["x_pad_token"]);
int64_t img_len = tgt->ne[1];
auto img = ggml_concat(gctx, src, tgt, 1);
ggml_tensor* sig = nullptr;
int64_t sig_len = 0;
if (semantic != nullptr) {
sig = sigvq_embed_1->forward(ctx, sigvq_embed_0->forward(ctx, semantic));
sig = pad_stream(ctx, sig, params["sigvq_pad_token"]);
sig_len = sig->ne[1];
}
GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]);
auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2);
auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
cap = block->forward(ctx, cap, cap_pe, nullptr, nullptr);
}
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
img = block->forward(ctx, img, img_pe, nullptr, img_adaln);
}
if (sig != nullptr) {
auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["sigvq_refiner." + std::to_string(i)]);
sig = block->forward(ctx, sig, sig_pe, nullptr, nullptr);
}
}
auto seq = ggml_concat(gctx, cap, img, 1);
auto cap_adaln = ggml_concat(gctx, per_token(t_clean, cap_len), per_token(t_noisy, cap_len), 1);
auto seq_adaln = ggml_concat(gctx, cap_adaln, img_adaln, 1);
if (sig != nullptr) {
seq = ggml_concat(gctx, seq, sig, 1);
seq_adaln = ggml_concat(gctx, seq_adaln, per_token(t_clean, sig_len), 1);
}
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
seq = block->forward(ctx, seq, pe, nullptr, seq_adaln);
sd::ggml_graph_cut::mark_graph_cut(seq, "llada_image.layers." + std::to_string(i), "seq");
}
seq = final_layer->forward(ctx, seq, seq_adaln);
// Only the target latent is denoised; the source half of the image stream is context.
// The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of
// the latent grid that unpatchify reconstructs.
int64_t target_start = cap_len * 2 + img_len;
return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
// x: [N, C, H, W]
// timestep: [N,]
// context: [N, L, cap_feat_dim]
// pe: [L, d_head/2, 2, 2]
// return: [N, C, H, W]
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int patch_size = config.patch_size;
auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
auto out = forward_core(ctx, img, timestep, context, pe);
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);
// The reference pipeline negates the model output before the scheduler step.
return ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
}
};
struct LLaDAImageRunner : public DiffusionModelRunner {
public:
LLaDAImageConfig config;
LLaDAImageModel llada_image;
std::vector<float> pe_vec;
LLaDAImageRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(LLaDAImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
llada_image = LLaDAImageModel(config);
llada_image.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
llada_image.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor) {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x->ne[3] == 1);
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
ZImage::SEQ_MULTI_OF,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, cap_feat_dim]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
}
ggml_cgraph* build_edit_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const sd::Tensor<float>& semantic_tensor,
const sd::Tensor<float>& source_tensor) {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_input(context_tensor);
ggml_tensor* semantic = make_optional_input(semantic_tensor);
ggml_tensor* source = make_input(source_tensor);
GGML_ASSERT(x->ne[3] == 1);
pe_vec = Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(context->ne[1]),
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
ZImage::SEQ_MULTI_OF,
config.theta,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
int64_t W = x->ne[0];
int64_t H = x->ne[1];
auto target = DiT::pad_and_patchify(&runner_ctx, x, config.patch_size, config.patch_size, false);
auto src = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false);
auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe);
out = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false);
out = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
const auto* extra = std::get_if<LLaDAImageDiffusionExtra>(&diffusion_params.extra);
bool has_semantic = extra != nullptr && extra->semantic != nullptr && !extra->semantic->empty();
bool has_ref_latent = diffusion_params.ref_latents != nullptr && !diffusion_params.ref_latents->empty();
if (has_semantic && !has_ref_latent) {
LOG_WARN("llada_image: SigVQ features without a reference latent are not supported; falling back to text to image");
}
if (has_ref_latent) {
const auto& source = diffusion_params.ref_latents->front();
if (source.shape() != diffusion_params.x->shape()) {
LOG_ERROR("llada_image: reference latent must match the target shape; use resize_vae_to_target=1");
return {};
}
auto get_graph = [&]() -> ggml_cgraph* {
return build_edit_graph(*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
tensor_or_empty(extra != nullptr ? extra->semantic : nullptr),
source);
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
diffusion_params.x->dim());
}
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context));
}
};
} // namespace LLaDAImage
#endif // __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
src/model/diffusion/model.hpp
+10 -1
View File
@@ -39,6 +39,9 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
{"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
// pass_to_vlm routes the reference image to the conditioner, which is where LLaDA-Image's
// SigVQ encoder lives; it does its own half-resolution resize.
{"llada_image", {true, true, Rope::RefIndexMode::FIXED, true, true, -1, RefImageResizeMode::NONE, -1, -1, true}},
{"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
};
@@ -131,6 +134,11 @@ struct HunyuanVideoDiffusionExtra {
const sd::Tensor<float>* timestep_r = nullptr;
};
struct LLaDAImageDiffusionExtra {
// SigVQ semantic features of the reference image; present only in editing mode.
const sd::Tensor<float>* semantic = nullptr;
};
using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra,
SkipLayerDiffusionExtra,
@@ -143,7 +151,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
MiniMaxH3DiffusionExtra,
MiniT2IDiffusionExtra,
SenseNovaU1DiffusionExtra,
HunyuanVideoDiffusionExtra>;
HunyuanVideoDiffusionExtra,
LLaDAImageDiffusionExtra>;
struct DiffusionParams {
const sd::Tensor<float>* x = nullptr;
src/model/diffusion/z_image.hpp
+64 -18
View File
@@ -131,16 +131,30 @@ namespace ZImage {
int64_t num_heads;
int64_t num_kv_heads;
bool qk_norm;
bool split_qkv;
public:
JointAttention(int64_t hidden_size, int64_t head_dim, int64_t num_heads, int64_t num_kv_heads, bool qk_norm)
: head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm) {
blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
float scale = 1.f;
blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
JointAttention(int64_t hidden_size,
int64_t head_dim,
int64_t num_heads,
int64_t num_kv_heads,
bool qk_norm,
bool norm_elementwise_affine = true,
bool split_qkv = false)
: head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm), split_qkv(split_qkv) {
float scale = 1.f;
if (split_qkv) {
blocks["to_q"] = std::make_shared<Linear>(hidden_size, num_heads * head_dim, false);
blocks["to_k"] = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
blocks["to_v"] = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
blocks["to_out.0"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
} else {
blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
}
if (qk_norm) {
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-06f, norm_elementwise_affine);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-06f, norm_elementwise_affine);
}
}
@@ -151,8 +165,35 @@ namespace ZImage {
// x: [N, n_token, hidden_size]
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks[split_qkv ? "to_out.0" : "out"]);
if (split_qkv) {
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
if (sd_backend_is(ctx->backend, "ROCm")) {
out_proj->set_scale(1.f / 16.f);
out_proj->set_force_prec_f32(true);
q_proj->set_force_prec_f32(true);
k_proj->set_force_prec_f32(true);
v_proj->set_force_prec_f32(true);
}
auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N);
auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
auto v = ggml_reshape_4d(ctx->ggml_ctx, v_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
if (qk_norm) {
q = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"])->forward(ctx, q);
k = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"])->forward(ctx, k);
}
auto out = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f);
return out_proj->forward(ctx, out);
}
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
if (sd_backend_is(ctx->backend, "ROCm")) {
out_proj->set_scale(1.f / 16.f);
@@ -252,9 +293,12 @@ namespace ZImage {
ggml_tensor* x,
ggml_tensor* scale) {
// x: [N, L, C]
// scale: [N, C]
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
// scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing
// feeds a per-token timestep embedding so each segment carries its own modulation).
if (scale->ne[1] != x->ne[1]) {
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
}
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
return x;
}
@@ -272,14 +316,16 @@ namespace ZImage {
float ffn_dim_multiplier,
float norm_eps,
bool qk_norm,
bool modulation = true)
bool modulation = true,
bool norm_elementwise_affine = true,
bool split_qkv = false)
: modulation(modulation) {
blocks["attention"] = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm);
blocks["attention"] = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv);
blocks["feed_forward"] = std::make_shared<FeedForward>(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier);
blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
if (modulation) {
blocks["adaLN_modulation.0"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size);
}
src/model/te/llada_image_te.hpp
+604
View File
@@ -0,0 +1,604 @@
#ifndef __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
#define __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
#include <algorithm>
#include <array>
#include <cmath>
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
#include "model/common/ggml_block.hpp"
#include "model_loader.h"
// The conditioning components LLaDA-Image puts around its LLaDA2-MoE backbone.
// Ref: LLaDAImageQueryFormerModel / LLaDAImageTextProjectionModel in
// https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py
//
// QueryFormer turns the LLaDA token embeddings into 256 learned queries that the pipeline
// appends to the backbone input; TextProjection maps the backbone hidden states to the
// denoiser's caption dimension. Neither uses RoPE, and every norm is parameter-free.
// Both MLPs use the tanh GELU approximation, so ggml_gelu (not ggml_gelu_erf).
//
// SigVQ is the editing-only image encoder: a 40-layer ViT whose output is quantized against a
// 16384-entry codebook, with the resulting ids embedded and projected into the semantic features
// the denoiser consumes. Its MLP uses the exact erf GELU, unlike the two above.
namespace LLaDAImageTE {
constexpr int LLADA_IMAGE_TE_GRAPH_SIZE = 16384;
struct QueryFormerConfig {
int64_t num_queries = 256;
int64_t hidden_size = 2048;
int64_t num_layers = 1;
int64_t num_heads = 16;
int64_t intermediate_size = 8192;
float norm_eps = 1e-6f;
};
struct TextProjectionConfig {
int64_t hidden_size = 2048;
int64_t intermediate_size = 8960;
int64_t num_layers = 6;
int64_t num_heads = 32;
int64_t projection_dim = 2560;
float norm_eps = 1e-6f;
};
// Cross-attention with a single fused in_proj over q (from the queries) and k/v (from the
// token embeddings). The checkpoint stores in_proj as one [3*hidden, hidden] parameter.
struct QueryAttention : public GGMLBlock {
protected:
int64_t hidden_size;
int64_t num_heads;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
enum ggml_type wtype = get_type(prefix + "in_proj_weight", tensor_storage_map, GGML_TYPE_F32);
params["in_proj_weight"] = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size * 3);
params["in_proj_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size * 3);
}
public:
QueryAttention(int64_t hidden_size, int64_t num_heads)
: hidden_size(hidden_size), num_heads(num_heads) {
blocks["out_proj"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* query,
ggml_tensor* context,
ggml_tensor* mask = nullptr) {
// query: [N, num_queries, hidden_size], context: [N, n_token, hidden_size]
ggml_context* gctx = ctx->ggml_ctx;
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
auto w = params["in_proj_weight"];
auto b = params["in_proj_bias"];
auto slice_w = [&](int64_t index) {
return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size);
};
auto slice_b = [&](int64_t index) {
return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size);
};
auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0));
auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1));
auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2));
auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, num_queries, hidden_size]
return out_proj->forward(ctx, x);
}
};
struct QueryFormerBlock : public GGMLBlock {
protected:
QueryFormerConfig config;
public:
QueryFormerBlock(const QueryFormerConfig& config)
: config(config) {
blocks["norm_q"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["norm_k"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["cross_attn"] = std::make_shared<QueryAttention>(config.hidden_size, config.num_heads);
blocks["norm1"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* query,
ggml_tensor* context,
ggml_tensor* mask = nullptr) {
auto norm_q = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_k"]);
auto cross_attn = std::dynamic_pointer_cast<QueryAttention>(blocks["cross_attn"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
// The reference overwrites query_embeds with its normalized value before the
// residual add, so both residuals here are on normalized activations.
query = norm_q->forward(ctx, query);
auto ctx_n = norm_k->forward(ctx, context);
query = ggml_add(ctx->ggml_ctx, query, cross_attn->forward(ctx, query, ctx_n, mask));
query = norm1->forward(ctx, query);
auto h = fc1->forward(ctx, query);
h = ggml_gelu(ctx->ggml_ctx, h);
h = fc2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, query, h);
}
};
struct QueryFormerModel : public GGMLBlock {
protected:
QueryFormerConfig config;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["meta_queries"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, config.num_queries);
}
public:
QueryFormerModel() = default;
QueryFormerModel(const QueryFormerConfig& config)
: config(config) {
for (int i = 0; i < config.num_layers; i++) {
blocks["query_blocks." + std::to_string(i)] = std::make_shared<QueryFormerBlock>(config);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* inputs_embeds,
ggml_tensor* mask = nullptr) {
// inputs_embeds: [N, n_token, hidden_size] -> [N, num_queries, hidden_size]
auto query = params["meta_queries"];
query = ggml_reshape_3d(ctx->ggml_ctx, query, config.hidden_size, config.num_queries, 1);
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<QueryFormerBlock>(blocks["query_blocks." + std::to_string(i)]);
query = block->forward(ctx, query, inputs_embeds, mask);
}
return query;
}
};
struct TextProjectionAttention : public GGMLBlock {
protected:
int64_t num_heads;
int64_t head_dim;
public:
TextProjectionAttention(const TextProjectionConfig& config)
: num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) {
blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["out_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, config.norm_eps, false);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, config.norm_eps, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size]
ggml_context* gctx = ctx->ggml_ctx;
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto q = q_proj->forward(ctx, x);
auto k = k_proj->forward(ctx, x);
auto v = v_proj->forward(ctx, x);
q = ggml_reshape_4d(gctx, q, head_dim, num_heads, n_token, N);
k = ggml_reshape_4d(gctx, k, head_dim, num_heads, n_token, N);
q = q_norm->forward(ctx, q);
k = k_norm->forward(ctx, k);
q = ggml_reshape_3d(gctx, q, head_dim * num_heads, n_token, N);
k = ggml_reshape_3d(gctx, k, head_dim * num_heads, n_token, N);
auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
return out_proj->forward(ctx, out);
}
};
struct TextProjectionBlock : public GGMLBlock {
public:
TextProjectionBlock(const TextProjectionConfig& config) {
blocks["self_attn"] = std::make_shared<TextProjectionAttention>(config);
blocks["layer_norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
blocks["layer_norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto self_attn = std::dynamic_pointer_cast<TextProjectionAttention>(blocks["self_attn"]);
auto layer_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm1"]);
auto layer_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm2"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
x = ggml_add(ctx->ggml_ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x)));
auto h = fc1->forward(ctx, layer_norm2->forward(ctx, x));
h = ggml_gelu(ctx->ggml_ctx, h);
h = fc2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct TextProjectionModel : public GGMLBlock {
protected:
TextProjectionConfig config;
public:
TextProjectionModel() = default;
TextProjectionModel(const TextProjectionConfig& config)
: config(config) {
for (int i = 0; i < config.num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::make_shared<TextProjectionBlock>(config);
}
blocks["projector"] = std::make_shared<Linear>(config.hidden_size, config.projection_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size] -> [N, n_token, projection_dim]
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<TextProjectionBlock>(blocks["layers." + std::to_string(i)]);
x = block->forward(ctx, x);
}
auto projector = std::dynamic_pointer_cast<Linear>(blocks["projector"]);
return projector->forward(ctx, x);
}
};
struct SigVQConfig {
int64_t image_size = 2048;
int64_t patch_size = 16;
int64_t in_channels = 3;
int64_t hidden_size = 1536;
int64_t intermediate_size = 6144;
int64_t num_layers = 40;
int64_t num_heads = 16;
int64_t codebook_size = 16384;
int64_t codebook_embed_dim = 2048;
int64_t semantic_embed_dim = 4096;
float norm_eps = 1e-6f;
};
struct SigVQAttention : public GGMLBlock {
protected:
int64_t num_heads;
int64_t head_dim;
public:
SigVQAttention(const SigVQConfig& config)
: num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) {
blocks["qkv"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size * 3, true);
blocks["proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size]
ggml_context* gctx = ctx->ggml_ctx;
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
int64_t hidden_size = num_heads * head_dim;
auto qkv = qkv_proj->forward(ctx, x);
auto q = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size);
auto k = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2);
auto v = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3);
auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
return out_proj->forward(ctx, out);
}
};
struct SigVQBlock : public GGMLBlock {
public:
SigVQBlock(const SigVQConfig& config) {
blocks["norm1"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps);
blocks["norm2"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps);
blocks["attn"] = std::make_shared<SigVQAttention>(config);
blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto attn = std::dynamic_pointer_cast<SigVQAttention>(blocks["attn"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, norm1->forward(ctx, x)));
auto h = fc1->forward(ctx, norm2->forward(ctx, x));
h = ggml_gelu_erf(ctx->ggml_ctx, h);
h = fc2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct SigVQModel : public GGMLBlock {
protected:
SigVQConfig config;
public:
SigVQModel() = default;
SigVQModel(const SigVQConfig& config)
: config(config) {
blocks["visual.patch_embed.proj"] = std::make_shared<Conv2d>(config.in_channels,
config.hidden_size,
std::make_pair(config.patch_size, config.patch_size),
std::make_pair(config.patch_size, config.patch_size));
for (int i = 0; i < config.num_layers; i++) {
blocks["visual.blocks." + std::to_string(i)] = std::make_shared<SigVQBlock>(config);
}
blocks["vqmodel.quant_conv"] = std::make_shared<Conv2d>(config.hidden_size,
config.codebook_embed_dim,
std::make_pair(1, 1));
blocks["prior_projector.net.0.proj"] = std::make_shared<Linear>(config.semantic_embed_dim, config.semantic_embed_dim, true);
blocks["prior_projector.net.2"] = std::make_shared<Linear>(config.semantic_embed_dim, config.semantic_embed_dim, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["visual.embeddings.position_embedding.weight"] =
ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, (config.image_size / config.patch_size) * (config.image_size / config.patch_size));
params["vqmodel.quantize.embedding.weight"] =
ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.codebook_embed_dim, config.codebook_size);
params["prior_token_embedding.weight"] =
ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.semantic_embed_dim, config.codebook_size);
}
// Bilinear-resamples the square position-embedding grid onto the image's patch grid.
// The reference uses grid_sample(align_corners=False, padding_mode="border"); the source
// coordinate for output index j is therefore (j + 0.5) * side / out - 0.5, clamped.
ggml_tensor* resample_pos_embed(GGMLRunnerContext* ctx,
ggml_tensor* pos_idx,
ggml_tensor* pos_weight) {
auto pos_embed = params["visual.embeddings.position_embedding.weight"];
auto gathered = ggml_get_rows(ctx->ggml_ctx, pos_embed, pos_idx);
return ggml_mul(ctx->ggml_ctx, gathered, pos_weight);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* pixel_values,
const std::vector<ggml_tensor*>& pos_idx,
const std::vector<ggml_tensor*>& pos_weight) {
// pixel_values: [N, in_channels, H, W] -> [N, grid_h * grid_w, semantic_embed_dim]
ggml_context* gctx = ctx->ggml_ctx;
auto patch_embed = std::dynamic_pointer_cast<Conv2d>(blocks["visual.patch_embed.proj"]);
auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["vqmodel.quant_conv"]);
auto proj_0 = std::dynamic_pointer_cast<Linear>(blocks["prior_projector.net.0.proj"]);
auto proj_2 = std::dynamic_pointer_cast<Linear>(blocks["prior_projector.net.2"]);
auto x = patch_embed->forward(ctx, pixel_values); // [N, hidden_size, grid_h, grid_w]
int64_t grid_w = x->ne[0];
int64_t grid_h = x->ne[1];
int64_t n_token = grid_h * grid_w;
int64_t N = x->ne[3];
x = ggml_reshape_3d(gctx, x, n_token, config.hidden_size, N);
x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, hidden_size]
ggml_tensor* pos = nullptr;
for (size_t i = 0; i < pos_idx.size(); i++) {
auto corner = resample_pos_embed(ctx, pos_idx[i], pos_weight[i]);
pos = pos == nullptr ? corner : ggml_add(gctx, pos, corner);
}
x = ggml_add(gctx, x, ggml_reshape_3d(gctx, pos, config.hidden_size, n_token, N));
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<SigVQBlock>(blocks["visual.blocks." + std::to_string(i)]);
x = block->forward(ctx, x);
}
// quant_conv is 1x1, so run it as a per-token projection rather than reshaping to 2-D.
x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, hidden_size, n_token]
x = ggml_reshape_4d(gctx, x, n_token, 1, config.hidden_size, N);
x = quant_conv->forward(ctx, x); // [N, codebook_embed_dim, 1, n_token]
x = ggml_reshape_3d(gctx, x, n_token, config.codebook_embed_dim, N);
x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, codebook_embed_dim]
// Both sides are L2-normalized, so the nearest codebook entry by euclidean distance
// is the one with the largest dot product.
auto codebook = ggml_l2_norm(gctx, params["vqmodel.quantize.embedding.weight"], 1e-12f);
auto normed = ggml_l2_norm(gctx, x, 1e-12f);
auto logits = ggml_mul_mat(gctx, codebook, normed); // [N, n_token, codebook_size]
auto token_ids = ggml_argmax(gctx, ggml_reshape_2d(gctx, logits, config.codebook_size, n_token * N));
auto semantic = ggml_get_rows(gctx, params["prior_token_embedding.weight"], token_ids);
semantic = ggml_reshape_3d(gctx, semantic, config.semantic_embed_dim, n_token, N);
auto h = proj_0->forward(ctx, semantic);
h = ggml_silu(gctx, h);
return proj_2->forward(ctx, h);
}
};
struct QueryFormerRunner : public GGMLRunner {
public:
QueryFormerConfig config;
QueryFormerModel query_former;
QueryFormerRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
query_former = QueryFormerModel(config);
query_former.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image_queryformer";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
query_former.get_param_tensors(tensors, prefix);
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& inputs_embeds) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
ggml_tensor* x = make_input(inputs_embeds);
auto runner_ctx = get_context();
ggml_tensor* out = query_former.forward(&runner_ctx, x);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
inputs_embeds.dim());
}
};
struct TextProjectionRunner : public GGMLRunner {
public:
TextProjectionConfig config;
TextProjectionModel text_projection;
TextProjectionRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
text_projection = TextProjectionModel(config);
text_projection.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image_text_projection";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
text_projection.get_param_tensors(tensors, prefix);
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& hidden_states) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
ggml_tensor* x = make_input(hidden_states);
auto runner_ctx = get_context();
ggml_tensor* out = text_projection.forward(&runner_ctx, x);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
hidden_states.dim());
}
};
struct SigVQRunner : public GGMLRunner {
public:
SigVQConfig config;
SigVQModel sigvq;
std::array<std::vector<int32_t>, 4> pos_idx_data;
std::array<std::vector<float>, 4> pos_weight_data;
SigVQRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
sigvq = SigVQModel(config);
sigvq.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image_sigvq";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
sigvq.get_param_tensors(tensors, prefix);
}
// Precomputes the four bilinear taps that resample the square position-embedding grid
// onto a grid_h x grid_w patch grid, matching grid_sample(align_corners=False,
// padding_mode="border").
void build_pos_embed_taps(int64_t grid_h, int64_t grid_w) {
const int64_t side = config.image_size / config.patch_size;
for (auto& v : pos_idx_data) {
v.clear();
}
for (auto& v : pos_weight_data) {
v.clear();
}
auto clamp_index = [side](int64_t v) {
return static_cast<int32_t>(std::min<int64_t>(std::max<int64_t>(v, 0), side - 1));
};
for (int64_t i = 0; i < grid_h; ++i) {
double src_h = (static_cast<double>(i) + 0.5) * side / static_cast<double>(grid_h) - 0.5;
int64_t h_floor = static_cast<int64_t>(std::floor(src_h));
double dh = src_h - static_cast<double>(h_floor);
for (int64_t j = 0; j < grid_w; ++j) {
double src_w = (static_cast<double>(j) + 0.5) * side / static_cast<double>(grid_w) - 0.5;
int64_t w_floor = static_cast<int64_t>(std::floor(src_w));
double dw = src_w - static_cast<double>(w_floor);
int32_t h0 = clamp_index(h_floor);
int32_t h1 = clamp_index(h_floor + 1);
int32_t w0 = clamp_index(w_floor);
int32_t w1 = clamp_index(w_floor + 1);
pos_idx_data[0].push_back(h0 * static_cast<int32_t>(side) + w0);
pos_idx_data[1].push_back(h0 * static_cast<int32_t>(side) + w1);
pos_idx_data[2].push_back(h1 * static_cast<int32_t>(side) + w0);
pos_idx_data[3].push_back(h1 * static_cast<int32_t>(side) + w1);
pos_weight_data[0].push_back(static_cast<float>((1.0 - dh) * (1.0 - dw)));
pos_weight_data[1].push_back(static_cast<float>((1.0 - dh) * dw));
pos_weight_data[2].push_back(static_cast<float>(dh * (1.0 - dw)));
pos_weight_data[3].push_back(static_cast<float>(dh * dw));
}
}
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& pixel_values) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
ggml_tensor* x = make_input(pixel_values);
int64_t grid_h = x->ne[1] / config.patch_size;
int64_t grid_w = x->ne[0] / config.patch_size;
build_pos_embed_taps(grid_h, grid_w);
std::vector<ggml_tensor*> pos_idx;
std::vector<ggml_tensor*> pos_weight;
for (int i = 0; i < 4; i++) {
auto idx = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast<int64_t>(pos_idx_data[i].size()));
set_backend_tensor_data(idx, pos_idx_data[i].data());
auto w = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, 1, static_cast<int64_t>(pos_weight_data[i].size()));
set_backend_tensor_data(w, pos_weight_data[i].data());
pos_idx.push_back(idx);
pos_weight.push_back(w);
}
auto runner_ctx = get_context();
ggml_tensor* out = sigvq.forward(&runner_ctx, x, pos_idx, pos_weight);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), 3);
}
};
} // namespace LLaDAImageTE
#endif // __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
src/model/te/llm.hpp
+346 -23
View File
@@ -49,6 +49,7 @@ namespace LLM {
GEMMA2_2B,
GEMMA4_12B,
GPT_OSS_20B,
LLADA2_MOE,
ARCH_COUNT,
};
@@ -62,6 +63,7 @@ namespace LLM {
"gemma2_2b",
"gemma4_12b",
"gpt_oss_20b",
"llada2_moe",
};
enum class MLPActivation {
@@ -125,6 +127,17 @@ namespace LLM {
std::vector<int> sliding_attention;
int64_t num_experts = 0;
int64_t num_experts_per_tok = 0;
bool qkv_fused = false;
bool bidirectional = false;
float partial_rotary = 1.f;
// DeepSeek-V3-style grouped-sigmoid MoE routing (LLaDA2)
int64_t moe_intermediate_size = 0;
int64_t num_shared_experts = 0;
int64_t first_k_dense_replace = 0;
int64_t n_group = 0;
int64_t topk_group = 0;
float routed_scaling_factor = 1.f;
LLMVisionConfig vision;
bool have_vision_weight = false;
bool llama_cpp_style = false;
@@ -212,6 +225,31 @@ namespace LLM {
config.intermediate_size = 9216;
config.num_layers = 26;
config.vocab_size = 256000;
} else if (arch == LLMArch::LLADA2_MOE) {
config.head_dim = 128;
config.num_heads = 16;
config.num_kv_heads = 4;
config.qkv_bias = false;
config.attention_out_bias = false;
config.qk_norm = true;
config.rms_norm_eps = 1e-6f;
config.hidden_size = 2048;
config.intermediate_size = 5120;
config.num_layers = 20;
config.vocab_size = 173568;
config.max_position_embeddings = 16384;
config.rope_thetas = {600000.f};
config.qkv_fused = true;
config.bidirectional = true;
config.partial_rotary = 0.5f;
config.num_experts = 256;
config.num_experts_per_tok = 8;
config.moe_intermediate_size = 512;
config.num_shared_experts = 1;
config.first_k_dense_replace = 1;
config.n_group = 8;
config.topk_group = 4;
config.routed_scaling_factor = 2.5f;
} else if (arch == LLMArch::GPT_OSS_20B) {
config.head_dim = 64;
config.num_heads = 64;
@@ -419,6 +457,195 @@ namespace LLM {
}
};
// LLaDA2's MoE differs from GPT-OSS's in three ways that all change the result:
// routing scores are sigmoid (not softmax over the selected logits), expert selection is
// group-limited and uses a bias term that the returned weights do NOT include, and the
// experts carry no biases. Ref: LLaDA2MoeGate / LLaDA2MoeSparseMoeBlock in
// modeling_llada2uni_moe.py.
struct LLaDA2MoEMLP : public GGMLBlock {
protected:
int64_t hidden_size;
int64_t moe_intermediate_size;
int64_t num_experts;
int64_t num_experts_per_tok;
int64_t n_group;
int64_t topk_group;
float routed_scaling_factor;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
auto supported_type = [](ggml_type wtype, int64_t in_features) {
if (in_features % ggml_blck_size(wtype) != 0) {
return GGML_TYPE_F32;
}
return wtype;
};
// The reference runs the router in fp32; keep the weight in fp32 so the sigmoid
// scores and the group sums match.
params["gate.weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts);
params["gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
ggml_type gate_type = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
ggml_type up_type = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
ggml_type down_type = supported_type(get_type(prefix + "experts.down_proj.weight", tensor_storage_map, GGML_TYPE_F32), moe_intermediate_size);
// HF ships the stacked experts as 3-D nn.Parameters, while the ComfyUI GGUF repack
// flattens the expert axis into ne[1]. Declare whichever the file holds - the two are
// bit-identical, and forward() reshapes to 3-D for ggml_mul_mat_id either way.
auto declare_experts = [&](const std::string& name, ggml_type type, int64_t in_dim, int64_t out_dim) {
auto storage = tensor_storage_map.find(prefix + name);
if (storage != tensor_storage_map.end() && storage->second.n_dims == 2) {
GGML_ASSERT(storage->second.nelements() == in_dim * out_dim * num_experts);
params[name] = ggml_new_tensor_2d(ctx, type, in_dim, out_dim * num_experts);
} else {
params[name] = ggml_new_tensor_3d(ctx, type, in_dim, out_dim, num_experts);
}
};
declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size);
declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size);
declare_experts("experts.down_proj.weight", down_type, moe_intermediate_size, hidden_size);
}
public:
LLaDA2MoEMLP(const LLMConfig& config)
: hidden_size(config.hidden_size),
moe_intermediate_size(config.moe_intermediate_size),
num_experts(config.num_experts),
num_experts_per_tok(config.num_experts_per_tok),
n_group(config.n_group),
topk_group(config.topk_group),
routed_scaling_factor(config.routed_scaling_factor) {
if (config.num_shared_experts > 0) {
blocks["shared_experts"] = std::make_shared<MLP>(config.hidden_size,
config.moe_intermediate_size * config.num_shared_experts,
false,
config.mlp_activation);
}
}
// Reproduces group_limited_topk(): keep the topk_group groups with the highest
// "sum of the two best scores in the group", then take the global top-k among them.
ggml_tensor* group_limited_mask(GGMLRunnerContext* ctx,
ggml_tensor* routing_scores,
int64_t n_token_total) {
ggml_context* gctx = ctx->ggml_ctx;
const int64_t per_group = num_experts / n_group;
// [experts_per_group, n_group * tokens] so top-2 runs per (group, token) row.
auto grouped = ggml_reshape_2d(gctx, routing_scores, per_group, n_group * n_token_total);
auto best2_idx = ggml_argsort_top_k(gctx, grouped, 2); // [2, n_group * tokens]
auto grouped_val = ggml_reshape_3d(gctx, grouped, 1, per_group, n_group * n_token_total);
auto best2 = ggml_get_rows(gctx, grouped_val, best2_idx); // [1, 2, n_group * tokens]
best2 = ggml_reshape_2d(gctx, best2, 2, n_group * n_token_total);
auto group_score = ggml_reshape_2d(gctx, ggml_sum_rows(gctx, best2), n_group, n_token_total); // [n_group, tokens]
// Threshold = the topk_group-th largest group score, taken from the sorted top-k.
auto top_groups = ggml_argsort_top_k(gctx, group_score, (int)topk_group); // [topk_group, tokens]
auto group_val = ggml_reshape_3d(gctx, group_score, 1, n_group, n_token_total);
auto top_scores = ggml_get_rows(gctx, group_val, top_groups); // [1, topk_group, tokens]
top_scores = ggml_reshape_2d(gctx, top_scores, topk_group, n_token_total);
auto threshold = ggml_view_2d(gctx,
top_scores,
1,
n_token_total,
top_scores->nb[1],
(topk_group - 1) * top_scores->nb[0]); // [1, tokens]
threshold = ggml_cont(gctx, threshold);
// keep = 1 - step(threshold - score). step(0) == 0, so the group sitting exactly on
// the threshold is kept without needing an epsilon.
auto diff = ggml_sub(gctx, ggml_repeat(gctx, threshold, group_score), group_score);
auto keep = ggml_scale_bias(gctx, ggml_step(gctx, diff), -1.f, 1.f); // [n_group, tokens]
// 0 for kept groups, a large negative for dropped ones, broadcast over the group.
auto additive = ggml_scale_bias(gctx, keep, 1e30f, -1e30f);
additive = ggml_reshape_3d(gctx, additive, 1, n_group, n_token_total);
auto expanded = ggml_repeat_4d(gctx, additive, per_group, n_group, n_token_total, 1);
return ggml_reshape_2d(gctx, expanded, num_experts, n_token_total);
}
ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
const std::string& weight_name,
ggml_tensor* x,
ggml_tensor* selected_experts) {
ggml_tensor* w = params[weight_name];
if (w->ne[2] != num_experts) {
// Flattened layout: split the expert axis back out. ne[0] is untouched, so this
// stays valid for quantized types.
w = ggml_reshape_3d(ctx->ggml_ctx, w, w->ne[0], w->ne[1] / num_experts, num_experts);
}
return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size]
GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0);
GGML_ASSERT(n_group > 0 && topk_group > 0 && num_experts % n_group == 0);
ggml_context* gctx = ctx->ggml_ctx;
const int64_t n_token = x->ne[1];
const int64_t N = x->ne[2];
const int64_t n_token_total = n_token * N;
auto identity = x;
auto logits = ggml_mul_mat(gctx, params["gate.weight"], x);
logits = ggml_reshape_2d(gctx, logits, num_experts, n_token_total);
auto scores = ggml_sigmoid(gctx, logits); // [num_experts, tokens]
// The bias steers selection only; the combine weights come from the unbiased scores.
auto routing = ggml_add(gctx, scores, params["gate.expert_bias"]);
routing = ggml_add(gctx, routing, group_limited_mask(ctx, routing, n_token_total));
auto selected_experts = ggml_argsort_top_k(gctx, routing, (int)num_experts_per_tok); // [top_k, tokens]
auto score_rows = ggml_reshape_3d(gctx, scores, 1, num_experts, n_token_total);
auto weights = ggml_get_rows(gctx, score_rows, selected_experts); // [1, top_k, tokens]
weights = ggml_reshape_2d(gctx, weights, num_experts_per_tok, n_token_total);
if (num_experts_per_tok > 1) {
auto denom = ggml_scale_bias(gctx, ggml_sum_rows(gctx, weights), 1.f, 1e-20f); // [1, tokens]
weights = ggml_div(gctx, weights, ggml_repeat(gctx, denom, weights));
}
weights = ggml_scale(gctx, weights, routed_scaling_factor);
weights = ggml_reshape_3d(gctx, weights, 1, num_experts_per_tok, n_token_total);
auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total);
auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts);
auto up = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts);
auto activated = ggml_swiglu_split(gctx, gate, up);
auto experts = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts);
experts = ggml_mul(gctx, experts, weights);
ggml_tensor* out = nullptr;
for (int64_t i = 0; i < num_experts_per_tok; ++i) {
auto expert_out = ggml_view_2d(gctx,
experts,
hidden_size,
n_token_total,
experts->nb[2],
i * experts->nb[1]);
out = out == nullptr ? expert_out : ggml_add(gctx, out, expert_out);
}
if (num_experts_per_tok == 1) {
out = ggml_cont(gctx, out);
}
out = ggml_reshape_3d(gctx, out, hidden_size, n_token, N);
auto shared_it = blocks.find("shared_experts");
if (shared_it != blocks.end()) {
auto shared_experts = std::dynamic_pointer_cast<MLP>(shared_it->second);
out = ggml_add(gctx, out, shared_experts->forward(ctx, identity));
}
return out;
}
};
struct GPTOSSMLP : public GGMLBlock {
protected:
int64_t hidden_size;
@@ -605,21 +832,31 @@ namespace LLM {
}
txt_token_end = image_embeds[i].first;
auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
if (input_embed == nullptr) {
input_embed = txt_embed;
} else {
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
// An embed can sit flush against the previous one or at the very start/end of the
// sequence, leaving no text tokens to splice around it.
if (txt_token_end > txt_token_start) {
auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
if (input_embed == nullptr) {
input_embed = txt_embed;
} else {
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
}
}
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
if (input_embed == nullptr) {
input_embed = image_embeds[i].second;
} else {
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
}
}
txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1];
txt_token_end = raw_x->ne[1];
auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
if (txt_token_end > txt_token_start) {
auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
}
GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]);
return input_embed;
}
@@ -1122,6 +1359,7 @@ namespace LLM {
bool k_eq_v;
bool v_norm;
bool unscaled_attention;
bool qkv_fused;
float rms_norm_eps;
int rope_pairs;
@@ -1147,12 +1385,20 @@ namespace LLM {
k_eq_v(global_layer && config.global_k_eq_v),
v_norm(config.v_norm),
unscaled_attention(config.unscaled_attention),
qkv_fused(config.qkv_fused),
rms_norm_eps(config.rms_norm_eps),
rope_pairs(0) {
blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
if (!k_eq_v) {
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
if (qkv_fused) {
// The checkpoint ships q, k and v as one tensor and the loader cannot split a
// source tensor, so keep it fused and slice it in forward().
GGML_ASSERT(!k_eq_v);
blocks["query_key_value"] = std::make_shared<Linear>(config.hidden_size, (num_heads + num_kv_heads * 2) * head_dim, config.qkv_bias);
} else {
blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
if (!k_eq_v) {
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
}
}
blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, config.hidden_size, config.attention_out_bias);
if (config.qk_norm) {
@@ -1161,7 +1407,7 @@ namespace LLM {
}
// Proportional RoPE rotates only the leading `rope_pairs` dimension pairs of the head;
// the rest are left unrotated through freq_factors (see rope_freq_factors()).
float partial = global_layer ? config.global_partial_rotary : 1.f;
float partial = global_layer && config.global_partial_rotary != 1.f ? config.global_partial_rotary : config.partial_rotary;
rope_pairs = static_cast<int>(partial * head_dim / 2.f);
}
@@ -1186,14 +1432,28 @@ namespace LLM {
// x: [N, n_token, hidden_size]
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o_proj"]);
auto q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim]
auto k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
auto v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
ggml_tensor* q = nullptr;
ggml_tensor* k = nullptr;
ggml_tensor* v = nullptr;
if (qkv_fused) {
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["query_key_value"]);
auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
int64_t q_len = num_heads * head_dim;
int64_t k_len = num_kv_heads * head_dim;
q = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len);
k = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len);
v = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2);
} else {
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim]
k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
}
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, num_heads, head_dim]
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N); // [N, n_token, num_kv_heads, head_dim]
@@ -1336,6 +1596,38 @@ namespace LLM {
1.f,
32.f,
1.f);
} else if (arch == LLMArch::LLADA2_MOE) {
// LLaDA2 slices the head (query[..., :rotary_dim]) instead of zero-padding
// inv_freq like gemma does, so rotate_half pairs i with i + rotary_dim/2 and the
// frequencies use rotary_dim as the exponent denominator. Passing n_dims =
// rotary_dim reproduces both; freq_factors would give the wrong pairing.
int rotary_dim = rope_pairs * 2;
q = ggml_rope_ext(ctx->ggml_ctx,
q,
input_pos,
nullptr,
rotary_dim,
GGML_ROPE_TYPE_NEOX,
static_cast<int>(max_position_embeddings),
rope_thetas[0],
1.f,
0.f,
1.f,
32.f,
1.f);
k = ggml_rope_ext(ctx->ggml_ctx,
k,
input_pos,
nullptr,
rotary_dim,
GGML_ROPE_TYPE_NEOX,
static_cast<int>(max_position_embeddings),
rope_thetas[0],
1.f,
0.f,
1.f,
32.f,
1.f);
} else if (arch == LLMArch::QWEN3_VL) {
int sections[4] = {24, 20, 20, 0};
q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@@ -1432,6 +1724,8 @@ namespace LLM {
blocks["self_attn"] = std::make_shared<Attention>(config, sliding_attention == 0);
if (config.arch == LLMArch::GPT_OSS_20B) {
blocks["mlp"] = std::make_shared<GPTOSSMLP>(config);
} else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) {
blocks["mlp"] = std::make_shared<LLaDA2MoEMLP>(config);
} else {
blocks["mlp"] = std::make_shared<MLP>(config.hidden_size,
config.intermediate_size,
@@ -1485,6 +1779,10 @@ namespace LLM {
if (arch == LLMArch::GPT_OSS_20B) {
auto mlp = std::dynamic_pointer_cast<GPTOSSMLP>(blocks["mlp"]);
x = mlp->forward(ctx, x);
} else if (auto moe_mlp = std::dynamic_pointer_cast<LLaDA2MoEMLP>(blocks["mlp"])) {
// LLaDA2 is dense for the first first_k_dense_replace layers and MoE afterwards,
// so the block type varies per layer rather than per arch.
x = moe_mlp->forward(ctx, x);
} else {
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
x = mlp->forward(ctx, x);
@@ -1650,6 +1948,11 @@ namespace LLM {
return x;
}
ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) {
auto model = std::dynamic_pointer_cast<TextModel>(blocks["model"]);
return model->embed(ctx, input_ids);
}
std::shared_ptr<VisionModel> vision_model() {
GGML_ASSERT(enable_vision);
return std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
@@ -1990,7 +2293,8 @@ namespace LLM {
config.arch == LLMArch::GEMMA3_12B ||
config.arch == LLMArch::GEMMA4_12B ||
config.arch == LLMArch::GEMMA2_2B ||
config.arch == LLMArch::GPT_OSS_20B) {
config.arch == LLMArch::GPT_OSS_20B ||
config.arch == LLMArch::LLADA2_MOE) {
input_pos_vec.resize(n_tokens);
for (int i = 0; i < n_tokens; ++i) {
input_pos_vec[i] = i;
@@ -2042,8 +2346,9 @@ namespace LLM {
attention_mask_vec.resize(n_tokens * n_tokens);
for (int i0 = 0; i0 < n_tokens; i0++) {
for (int i1 = 0; i1 < n_tokens; i1++) {
// Diffusion LLMs attend in both directions; only causal LMs get the triangle.
float value = 0.f;
if (i0 > i1) {
if (!config.bidirectional && i0 > i1) {
value = -INFINITY;
}
attention_mask_vec[i1 * n_tokens + i0] = value;
@@ -2115,6 +2420,22 @@ namespace LLM {
input_ids.dim() + 1);
}
// LLaDA-Image's QueryFormer consumes the raw token embeddings before the backbone runs,
// so it needs the embedding lookup on its own.
sd::Tensor<float> compute_input_embeds(const int n_threads,
const sd::Tensor<int32_t>& input_ids) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* ids = make_input(input_ids);
auto runner_ctx = get_context();
ggml_tensor* out = model.embed(&runner_ctx, ids);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
input_ids.dim() + 1);
}
int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {
int64_t grid_t = 1;
int64_t grid_h = h / config.vision.patch_size;
@@ -2370,11 +2691,13 @@ namespace LLM {
pad_id = 199999;
} else if (arch == LLMArch::GEMMA2_2B) {
pad_id = 0;
} else if (arch == LLMArch::LLADA2_MOE) {
pad_id = 156892;
}
tokenizer = tokenizers.create(TokenizerConfig::MAIN, model.config.vocab_size, pad_id);
if (!tokenizer) {
if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B) {
throw std::runtime_error("GPT-OSS and Gemma 2 require an external tokenizer.json in the main tokenizer slot");
if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B || arch == LLMArch::LLADA2_MOE) {
throw std::runtime_error("GPT-OSS, Gemma 2 and LLaDA2 require an external tokenizer.json in the main tokenizer slot");
}
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
tokenizer = std::make_shared<MistralTokenizer>();