diffusion: Metal-side Stage-1 sampling + Accelerate host sampler

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)

examples/diffusion/diffusion.cpp

@@ -12,6 +12,10 @@
 #include <utility>
 #include <vector>
 
+#ifdef __APPLE__
+#include <Accelerate/Accelerate.h>
+#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<float> 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<std::thread> pool;
             const int32_t chunk = (C + (int32_t) nth - 1) / (int32_t) nth;

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)

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);

ggml/src/ggml-metal/ggml-metal-diffusion-sample.cpp

@@ -0,0 +1,99 @@
+#include "ggml-metal-device.h"
+
+#include <Accelerate/Accelerate.h>
+
+#include <algorithm>
+#include <cmath>
+#include <thread>
+#include <vector>
+
+// 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<float> 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<std::thread> 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;
+}

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;
 

src/models/diffusion-gemma.cpp

@@ -2,6 +2,7 @@
 #include "gemma4-common.h"
 
 #include <algorithm>
+#include <atomic>
 #include <cstring>
 #include <thread>
 #include <vector>
@@ -651,27 +652,60 @@ 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<const llama_model_diffusion_gemma *>(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) {
     const auto * dm = dynamic_cast<const llama_model_diffusion_gemma *>(model);
     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<dg_dev_sample_fn> 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() {