diff --git a/examples/diffusion/CMakeLists.txt b/examples/diffusion/CMakeLists.txt index 42a84b2dfe5f..7f789f0f7b5d 100644 --- a/examples/diffusion/CMakeLists.txt +++ b/examples/diffusion/CMakeLists.txt @@ -2,6 +2,9 @@ set(TARGET llama-diffusion) add_library(${TARGET} STATIC diffusion.cpp diffusion.h) target_link_libraries(${TARGET} PUBLIC llama llama-common ${CMAKE_THREAD_LIBS_INIT}) target_compile_features(${TARGET} PUBLIC cxx_std_17) +if (APPLE) + target_link_libraries(${TARGET} PUBLIC "-framework Accelerate") +endif() set(TARGET llama-diffusion-cli) add_executable(${TARGET} diffusion-cli.cpp) diff --git a/examples/diffusion/diffusion.cpp b/examples/diffusion/diffusion.cpp index 73649473f263..b909154e78cf 100644 --- a/examples/diffusion/diffusion.cpp +++ b/examples/diffusion/diffusion.cpp @@ -12,6 +12,10 @@ #include #include +#ifdef __APPLE__ +#include +#endif + static float calculate_confidence(const llama_token_data_array & cur_p, diffusion_algorithm algorithm, std::mt19937 & rng) { @@ -580,6 +584,72 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } // per position: argmax, entropy of softmax(raw/t), and a multinomial sample; stash raw row for SC +#ifdef __APPLE__ + // Accelerate-vectorized host path: argmax matches the scalar loop bit-for-bit; Z, H, and the + // sampled cumulative differ only by reduction order (same tolerance as the device sampler). + // H = -sum(p log p) with p = e_v/Z and log p = z_v - log Z ==> H = log Z - dot(e, z)/Z. + auto worker = [&](int32_t p0, int32_t p1) { + const int n = (int) n_vocab; + const int blk = 8192; + const int nblk = (n + blk - 1) / blk; + std::vector zbuf((size_t) n), ebuf((size_t) n), bsum((size_t) nblk); + for (int32_t pos = p0; pos < p1; pos++) { + const float * row = logits + (size_t) (logit_off + pos) * n_vocab; + + vDSP_vsmul(row, 1, &temp_inv, zbuf.data(), 1, (vDSP_Length) n); // z = row * temp_inv + float m; vDSP_Length amax; + vDSP_maxvi(zbuf.data(), 1, &m, &amax, (vDSP_Length) n); // m, argmax + const float neg_m = -m; + vDSP_vsadd(zbuf.data(), 1, &neg_m, zbuf.data(), 1, (vDSP_Length) n); // z -= m + // floor z so -inf logits (masked tokens) give e = exp(-80) ~ 1.8e-35 instead of + // e = 0 with z = -inf, whose 0 * -inf product would turn the dot into NaN entropy. + const float z_floor = -80.0f; + vDSP_vthr(zbuf.data(), 1, &z_floor, zbuf.data(), 1, (vDSP_Length) n); + vvexpf(ebuf.data(), zbuf.data(), &n); // e = exp(z) + + // per-block partial sums serve both Z and the multinomial crossing search + float Z = 0.0f; + for (int ib = 0; ib < nblk; ib++) { + const int b0 = ib * blk; + const int bn = std::min(blk, n - b0); + vDSP_sve(ebuf.data() + b0, 1, &bsum[(size_t) ib], (vDSP_Length) bn); + Z += bsum[(size_t) ib]; + } + float S; + vDSP_dotpr(ebuf.data(), 1, zbuf.data(), 1, &S, (vDSP_Length) n); // S = dot(e, z) + const float H = logf(Z) - S / Z; + + // multinomial: first v (vocab order) with cumsum(e) >= u*Z. Block sums skip ahead; + // if FP reduction order makes a claimed crossing vanish inside a block, the + // sequential cum carries into the next block instead of falling back, so the + // n_vocab-1 default only remains when the cumulative sum truly never reaches target. + const float target = u[pos] * Z; + int32_t sampled = (int32_t) n_vocab - 1; + bool picked = false; + float cum = 0.0f; + for (int ib = 0; ib < nblk && !picked; ib++) { + const int b0 = ib * blk; + const int bn = std::min(blk, n - b0); + if (cum + bsum[(size_t) ib] < target) { + cum += bsum[(size_t) ib]; + continue; + } + for (int v = b0; v < b0 + bn; v++) { + cum += ebuf[(size_t) v]; + if (cum >= target) { sampled = v; picked = true; break; } + } + } + + entropy[pos] = H; + argmax_canvas[pos] = (int32_t) amax; + denoiser[pos] = sampled; + // device SC keeps prev-step logits on-device (cpy in-graph), so no host stash needed + if (!dev_sc) { + std::memcpy(sc_buffer.data() + (size_t) pos * n_vocab, row, n_vocab * sizeof(float)); + } + } + }; +#else auto worker = [&](int32_t p0, int32_t p1) { for (int32_t pos = p0; pos < p1; pos++) { const float * row = logits + (size_t) (logit_off + pos) * n_vocab; @@ -611,6 +681,7 @@ void diffusion_generate_entropy_bound(llama_context * ctx, } } }; +#endif auto run_host_worker = [&]() { std::vector pool; const int32_t chunk = (C + (int32_t) nth - 1) / (int32_t) nth; diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt index 42054d841aa3..e6e11358d0a8 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt @@ -1,6 +1,7 @@ -find_library(FOUNDATION_LIBRARY Foundation REQUIRED) -find_library(METAL_FRAMEWORK Metal REQUIRED) -find_library(METALKIT_FRAMEWORK MetalKit REQUIRED) +find_library(FOUNDATION_LIBRARY Foundation REQUIRED) +find_library(METAL_FRAMEWORK Metal REQUIRED) +find_library(METALKIT_FRAMEWORK MetalKit REQUIRED) +find_library(ACCELERATE_FRAMEWORK Accelerate REQUIRED) message(STATUS "Metal framework found") @@ -11,12 +12,14 @@ ggml_add_backend_library(ggml-metal ggml-metal-common.cpp ggml-metal-context.m ggml-metal-ops.cpp + ggml-metal-diffusion-sample.cpp ) target_link_libraries(ggml-metal PRIVATE ${FOUNDATION_LIBRARY} ${METAL_FRAMEWORK} ${METALKIT_FRAMEWORK} + ${ACCELERATE_FRAMEWORK} ) if (GGML_METAL_NDEBUG) diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 4a3ebb5569d9..f7f2c3173d2a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -305,6 +305,13 @@ void ggml_metal_buffer_free (ggml_metal_buffer_t buf); void * ggml_metal_buffer_get_base (ggml_metal_buffer_t buf); bool ggml_metal_buffer_is_shared(ggml_metal_buffer_t buf); +// DiffusionGemma Stage-1 dense sampler reduction (ggml-metal-diffusion-sample.cpp): rows of a +// host-visible logits buffer are reduced in place on the CPU. Declared here so the defining and +// calling TUs share one prototype. +bool ggml_backend_metal_diffusion_sample_impl( + const float * base, const float * u, int * argmax, float * entropy, + int * sampled, int n_tokens, int n_vocab, float inv_temp); + void ggml_metal_buffer_memset_tensor(ggml_metal_buffer_t buf, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); void ggml_metal_buffer_set_tensor (ggml_metal_buffer_t buf, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); void ggml_metal_buffer_get_tensor (ggml_metal_buffer_t buf, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size); diff --git a/ggml/src/ggml-metal/ggml-metal-diffusion-sample.cpp b/ggml/src/ggml-metal/ggml-metal-diffusion-sample.cpp new file mode 100644 index 000000000000..8716113e983a --- /dev/null +++ b/ggml/src/ggml-metal/ggml-metal-diffusion-sample.cpp @@ -0,0 +1,99 @@ +#include "ggml-metal-device.h" + +#include + +#include +#include +#include +#include + +// Mirrors the CUDA diffusion_dense_sample_kernel contract: per canvas row, argmax (bit-exact with +// the host worker), entropy of softmax(logits * inv_temp) via H = log Z - dot(e, z)/Z, and a +// multinomial draw picking the first vocab-order crossing of u[row] * Z. Differences from the host +// path are limited to FP reduction order (same tolerance contract as the CUDA kernel). Unlike CUDA +// there is no device kernel or scratch upload: the shared MTLBuffer rows are reduced in place on +// the CPU with vDSP/vvexpf, which also skips the per-step full-logits fetch in the caller. +// Validation (tensor type, contiguity, metal buffer, shared storage) happens in ggml-metal.cpp. +extern "C" bool ggml_backend_metal_diffusion_sample_impl( + const float * base, + const float * u_host, + int * argmax_host, + float * entropy_host, + int * sampled_host, + int n_tokens, + int n_vocab, + float inv_temp) { + if (!base || !u_host || !argmax_host || !entropy_host || !sampled_host || n_tokens <= 0 || n_vocab <= 0) { + return false; + } + + const unsigned nth = std::max(1u, std::min(std::thread::hardware_concurrency(), (unsigned) n_tokens)); + + auto worker = [&](int p0, int p1) { + const int n = n_vocab; + const int blk = 8192; + const int nblk = (n + blk - 1) / blk; + std::vector zbuf((size_t) n), ebuf((size_t) n), bsum((size_t) nblk); + for (int pos = p0; pos < p1; pos++) { + const float * row = base + (size_t) pos * n_vocab; + + vDSP_vsmul(row, 1, &inv_temp, zbuf.data(), 1, (vDSP_Length) n); // z = row * inv_temp + float m; vDSP_Length amax; + vDSP_maxvi(zbuf.data(), 1, &m, &amax, (vDSP_Length) n); // m, argmax (first max) + const float neg_m = -m; + vDSP_vsadd(zbuf.data(), 1, &neg_m, zbuf.data(), 1, (vDSP_Length) n); // z -= m + // floor z so -inf logits (masked tokens) give e = exp(-80) ~ 1.8e-35 instead of e = 0 + // with z = -inf, whose 0 * -inf product would turn the dot below into NaN entropy. + const float z_floor = -80.0f; + vDSP_vthr(zbuf.data(), 1, &z_floor, zbuf.data(), 1, (vDSP_Length) n); + vvexpf(ebuf.data(), zbuf.data(), &n); // e = exp(z) + + // per-block partial sums serve both Z and the multinomial crossing search + float Z = 0.0f; + for (int ib = 0; ib < nblk; ib++) { + const int b0 = ib * blk; + const int bn = std::min(blk, n - b0); + vDSP_sve(ebuf.data() + b0, 1, &bsum[(size_t) ib], (vDSP_Length) bn); + Z += bsum[(size_t) ib]; + } + float S; + vDSP_dotpr(ebuf.data(), 1, zbuf.data(), 1, &S, (vDSP_Length) n); // S = dot(e, z) + + argmax_host[pos] = (int) amax; + entropy_host[pos] = logf(Z) - S / Z; + + // multinomial: first v (vocab order) with cumsum(e) >= u*Z. Block sums skip ahead; if FP + // reduction order makes a claimed crossing vanish inside a block, the sequential cum + // carries into the next block instead of falling back, so the n_vocab-1 default only + // remains when the cumulative sum truly never reaches the target. + const float target = u_host[pos] * Z; + int sampled = n_vocab - 1; + bool picked = false; + float cum = 0.0f; + for (int ib = 0; ib < nblk && !picked; ib++) { + const int b0 = ib * blk; + const int bn = std::min(blk, n - b0); + if (cum + bsum[(size_t) ib] < target) { + cum += bsum[(size_t) ib]; + continue; + } + for (int v = b0; v < b0 + bn; v++) { + cum += ebuf[(size_t) v]; + if (cum >= target) { sampled = v; picked = true; break; } + } + } + sampled_host[pos] = sampled; + } + }; + + std::vector pool; + const int chunk = (n_tokens + (int) nth - 1) / (int) nth; + for (unsigned ti = 0; ti < nth; ti++) { + const int p0 = (int) ti * chunk; + const int p1 = std::min(p0 + chunk, n_tokens); + if (p0 < p1) { pool.emplace_back(worker, p0, p1); } + } + for (auto & th : pool) { th.join(); } + + return true; +} diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index a1003b3acff8..d963c4ad612c 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -868,10 +868,41 @@ static ggml_backend_feature * ggml_backend_metal_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } +// DiffusionGemma Stage-1 dense sampler (reduction lives in ggml-metal-diffusion-sample.cpp). +// Shared and mapped Metal buffers are unified memory, so the sc_dev rows are reduced in place on +// the host with no device kernel and no logits fetch; private-storage tensors fall back to the +// host path. +extern "C" bool ggml_backend_metal_diffusion_sample( + struct ggml_tensor * logits, const float * u, int * argmax, float * entropy, + int * sampled, int n_tokens, float inv_temp) { + if (!logits || n_tokens <= 0) { + return false; + } + if (logits->type != GGML_TYPE_F32 || !ggml_is_contiguous(logits) || logits->data == nullptr) { + return false; + } + const int n_vocab = (int) logits->ne[0]; + if (n_vocab <= 0 || (int) ggml_nrows(logits) < n_tokens) { + return false; + } + if (!logits->buffer || !ggml_backend_buffer_is_metal(logits->buffer)) { + return false; + } + ggml_metal_buffer_t buf_ctx = (ggml_metal_buffer_t) logits->buffer->context; + if (!ggml_metal_buffer_is_shared(buf_ctx)) { + return false; // private storage is not host-visible -> caller falls back to the host path + } + return ggml_backend_metal_diffusion_sample_impl( + (const float *) logits->data, u, argmax, entropy, sampled, n_tokens, n_vocab, inv_temp); +} + static void * ggml_backend_metal_get_proc_address(ggml_backend_reg_t reg, const char * name) { if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_metal_get_features; } + if (strcmp(name, "ggml_backend_metal_diffusion_sample") == 0) { + return (void *)ggml_backend_metal_diffusion_sample; + } return NULL; diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index 096f5878bc11..d53c6b1974d7 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -2,6 +2,7 @@ #include "gemma4-common.h" #include +#include #include #include #include @@ -651,10 +652,38 @@ void llama_diffusion_set_device_sc(struct llama_model * model, bool enabled) { dm->sc_device_resident = enabled; } -// Stage-1 device sampling entry. Fetches the CUDA backend's dense sampler via the backend-reg proc address -// (keeps the llama<->ggml-cuda link at the existing backend boundary) and runs it on sc_dev. Returns false -// for non-DiffusionGemma / no sc_dev / non-CUDA builds so the caller falls back to the host path. -typedef bool (*dg_cuda_sample_fn)(struct ggml_tensor *, const float *, int *, float *, int *, int, float); +// Debug only: copy the device SC buffer (sc_dev) to host; returns floats copied (0 if none). Verifies +// sc_dev == the canvas logits the host path would upload. Off the hot path. +size_t llama_diffusion_debug_get_sc_dev(const struct llama_model * model, float * dst, size_t max_elems) { + const auto * dm = dynamic_cast(model); + if (!dm || dm->sc_dev == nullptr || dst == nullptr) { + return 0; + } + const size_t n = std::min(max_elems, (size_t) ggml_nelements(dm->sc_dev)); + ggml_backend_tensor_get(dm->sc_dev, dst, 0, n * sizeof(float)); + return n; +} + +// Stage-1 device sampling entry. Fetches a backend dense sampler via the backend-reg proc address +// (keeps the llama<->ggml-backend link at the existing backend boundary) and runs it on sc_dev. CUDA +// runs a device kernel; Metal reduces the shared-memory tensor in place on the host (unified memory), +// which skips the per-step full-logits fetch. Returns false for non-DiffusionGemma / no sc_dev / +// unsupported backends so the caller falls back to the host path. +typedef bool (*dg_dev_sample_fn)(struct ggml_tensor *, const float *, int *, float *, int *, int, float); + +static dg_dev_sample_fn dg_resolve_dev_sample_fn() { + if (ggml_backend_reg_t reg = ggml_backend_reg_by_name("CUDA")) { + if (void * p = ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_diffusion_sample")) { + return (dg_dev_sample_fn) p; + } + } + if (ggml_backend_reg_t reg = ggml_backend_reg_by_name("MTL")) { + if (void * p = ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_diffusion_sample")) { + return (dg_dev_sample_fn) p; + } + } + return nullptr; +} bool llama_diffusion_device_sample(const struct llama_model * model, const float * u, int * argmax, float * entropy, int * sampled, int n_tokens, float inv_temp) { @@ -662,16 +691,21 @@ bool llama_diffusion_device_sample(const struct llama_model * model, const float if (!dm || dm->sc_dev == nullptr || !u || !argmax || !entropy || !sampled || n_tokens <= 0) { return false; } - ggml_backend_reg_t reg = ggml_backend_reg_by_name("CUDA"); - if (!reg) { - return false; + // re-resolve while null so a backend registered after the first call (dynamic backend + // loading) is still picked up; only a successful resolution is cached. Atomic because this + // is a public llama.h entry point that external callers may hit from multiple threads. + static std::atomic fn{nullptr}; + dg_dev_sample_fn f = fn.load(std::memory_order_relaxed); + if (f == nullptr) { + f = dg_resolve_dev_sample_fn(); + if (f != nullptr) { + fn.store(f, std::memory_order_relaxed); + } } - static dg_cuda_sample_fn fn = - (dg_cuda_sample_fn) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_diffusion_sample"); - if (!fn) { + if (f == nullptr) { return false; } - return fn(dm->sc_dev, u, argmax, entropy, sampled, n_tokens, inv_temp); + return f(dm->sc_dev, u, argmax, entropy, sampled, n_tokens, inv_temp); } llama_model_diffusion_gemma::~llama_model_diffusion_gemma() {