From 04acc81bcc17a636366d425a3b60ac3d04a9c52a Mon Sep 17 00:00:00 2001 From: andreinknv Date: Thu, 11 Jun 2026 19:39:31 -0400 Subject: [PATCH] metal : add Stage-1 diffusion dense sampler for DiffusionGemma Port of the CUDA diffusion_dense_sample kernel to Metal so Apple Silicon gets on-device sampling instead of the host fallback. One 256-thread threadgroup per canvas row: parallel max/argmax reduction, parallel Z and T (T = sum d*e) reductions with entropy = logZ - T/Z, then the slice-scanned multinomial first-crossing walk in vocab order - the same structure as the CUDA kernel, so argmax matches the host bit-for-bit and Z/entropy differ only by reduction order. Wiring mirrors the CUDA backend boundary: the kernel is dispatched by ggml_metal_device_diffusion_sample (grow-only shared scratch for u/argmax/entropy/sampled), exposed through the Metal backend reg via get_proc_address("ggml_backend_metal_diffusion_sample"), and llama_diffusion_device_sample now resolves CUDA first, then MTL. Verified with DG_DEVSAMPLE_CHECK=1 on diffusiongemma-26B-A4B-it Q4_K_M (M4 Max): amax_mismatch 0/256 on every step, max|dH| <= 2e-4, tok_diff 0-4/256 near ties - the same tolerance profile as CUDA-vs-host. Per-step time 462ms -> 390ms, effective throughput 36.9 -> 43.8 tok/s with -fa on. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-metal/ggml-metal-device.h | 5 ++ ggml/src/ggml-metal/ggml-metal-device.m | 104 ++++++++++++++++++++++++ ggml/src/ggml-metal/ggml-metal.cpp | 15 ++++ ggml/src/ggml-metal/ggml-metal.metal | 104 ++++++++++++++++++++++++ src/models/diffusion-gemma.cpp | 29 ++++--- 5 files changed, 247 insertions(+), 10 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 4a3ebb5569d9..6228d7db0e4a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -292,6 +292,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te const struct ggml_metal_device_props * ggml_metal_device_get_props(ggml_metal_device_t dev); +// Stage-1 diffusion sampler over a device-resident logits tensor (DiffusionGemma entropy-bound decoder). +// Synchronous: encodes one kernel launch on the device queue and waits; outputs land in the host arrays. +bool ggml_metal_device_diffusion_sample(ggml_metal_device_t dev, struct ggml_tensor * logits, + const float * u, int * argmax, float * entropy, int * sampled, int n_tokens, float inv_temp); + // // device buffers // diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 05d7f43051ba..6af1d2b50bc4 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1877,3 +1877,107 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co return res; } + +// Stage-1 diffusion sampler over a device-resident logits tensor (DiffusionGemma entropy-bound decoder). +// Synchronous one-shot dispatch on the device queue; outputs land in the host arrays. Grow-only shared +// scratch (u in; argmax/entropy/sampled out) mirrors the CUDA implementation's cached per-device scratch. +bool ggml_metal_device_diffusion_sample(ggml_metal_device_t dev, struct ggml_tensor * logits, + const float * u, int * argmax, float * entropy, int * sampled, int n_tokens, float inv_temp) { + if (!dev || !logits || !u || !argmax || !entropy || !sampled || n_tokens <= 0) { + return false; + } + if (logits->type != GGML_TYPE_F32 || !ggml_is_contiguous(logits) || logits->buffer == NULL) { + return false; + } + const int32_t n_vocab = (int32_t) logits->ne[0]; + if (n_vocab <= 0 || (int) ggml_nrows(logits) < n_tokens) { + return false; + } + + ggml_metal_library_t lib = ggml_metal_device_get_library(dev); + if (!lib) { + return false; + } + struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline( + lib, "kernel_diffusion_dense_sample_f32", "kernel_diffusion_dense_sample_f32", NULL); + if (!ppl.pipeline) { + return false; + } + + struct ggml_metal_buffer_id bid = + ggml_metal_buffer_get_id((ggml_metal_buffer_t) logits->buffer->context, logits); + if (!bid.metal) { + return false; + } + + static id g_dg_scratch = nil; + static int g_dg_scratch_cap = 0; // tokens + static size_t g_dg_scratch_sec = 0; // bytes per section (u/argmax/entropy/sampled) + static NSLock * g_dg_scratch_lock = nil; + static dispatch_once_t g_dg_scratch_once; + dispatch_once(&g_dg_scratch_once, ^{ g_dg_scratch_lock = [[NSLock alloc] init]; }); + + [g_dg_scratch_lock lock]; + + bool ok = false; + + @autoreleasepool { + if (g_dg_scratch_cap < n_tokens) { + if (g_dg_scratch) { + [g_dg_scratch release]; + g_dg_scratch = nil; + } + id mtl_dev = (id) ggml_metal_device_get_obj(dev); + const size_t sec = ((size_t) n_tokens * sizeof(float) + 255) & ~(size_t) 255; + g_dg_scratch = [mtl_dev newBufferWithLength:4*sec options:MTLResourceStorageModeShared]; + if (!g_dg_scratch) { + [g_dg_scratch_lock unlock]; + return false; + } + g_dg_scratch_cap = n_tokens; + g_dg_scratch_sec = sec; + } + + const size_t sec = g_dg_scratch_sec; + char * base = (char *) [g_dg_scratch contents]; + + memcpy(base, u, (size_t) n_tokens * sizeof(float)); + + id queue = (id) ggml_metal_device_get_queue(dev); + id cb = [queue commandBuffer]; + + ggml_metal_encoder_t enc = ggml_metal_encoder_init((ggml_metal_cmd_buf_t) cb, false); + + struct ggml_metal_buffer_id bid_u = { (void *) g_dg_scratch, 0*sec }; + struct ggml_metal_buffer_id bid_argmax = { (void *) g_dg_scratch, 1*sec }; + struct ggml_metal_buffer_id bid_entropy = { (void *) g_dg_scratch, 2*sec }; + struct ggml_metal_buffer_id bid_sampled = { (void *) g_dg_scratch, 3*sec }; + + ggml_metal_encoder_set_pipeline(enc, ppl); + ggml_metal_encoder_set_buffer (enc, bid, 0); + ggml_metal_encoder_set_buffer (enc, bid_u, 1); + ggml_metal_encoder_set_buffer (enc, bid_argmax, 2); + ggml_metal_encoder_set_buffer (enc, bid_entropy, 3); + ggml_metal_encoder_set_buffer (enc, bid_sampled, 4); + ggml_metal_encoder_set_bytes (enc, (void *) &n_vocab, sizeof(n_vocab), 5); + ggml_metal_encoder_set_bytes (enc, (void *) &inv_temp, sizeof(inv_temp), 6); + + ggml_metal_encoder_dispatch_threadgroups(enc, n_tokens, 1, 1, 256, 1, 1); + ggml_metal_encoder_end_encoding(enc); + ggml_metal_encoder_free(enc); + + [cb commit]; + [cb waitUntilCompleted]; + + if ([cb status] == MTLCommandBufferStatusCompleted) { + memcpy(argmax, base + 1*sec, (size_t) n_tokens * sizeof(int)); + memcpy(entropy, base + 2*sec, (size_t) n_tokens * sizeof(float)); + memcpy(sampled, base + 3*sec, (size_t) n_tokens * sizeof(int)); + ok = true; + } + } + + [g_dg_scratch_lock unlock]; + + return ok; +} diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index a1003b3acff8..fd1788aca984 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -868,11 +868,26 @@ static ggml_backend_feature * ggml_backend_metal_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } +// Stage-1 DiffusionGemma sampler, exposed via get_proc_address so llama can reach it across the backend +// boundary (mirrors ggml_backend_cuda_diffusion_sample). Single Metal GPU -> device 0. +static 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) { + ggml_metal_device_t dev = ggml_metal_device_get(0); + if (!dev) { + return false; + } + return ggml_metal_device_diffusion_sample(dev, logits, u, argmax, entropy, sampled, n_tokens, 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; GGML_UNUSED(reg); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 2bd310d94506..5dae80473eb0 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -10733,3 +10733,107 @@ kernel void kernel_count_equal( typedef decltype(kernel_count_equal) kernel_count_equal_t; template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal; + +// Stage-1 diffusion sampler (DiffusionGemma entropy-bound decoder): one threadgroup of 256 threads per +// canvas row. Mirrors ggml-cuda/diffusion-sampling.cu: parallel max->argmax, parallel Z and T (T = sum d*e), +// entropy = logZ - T/Z, then a slice-scanned multinomial first-crossing walk in vocab order. Only the +// reduction order differs from the host worker, so argmax is exact and Z/entropy match to FP tolerance. +kernel void kernel_diffusion_dense_sample_f32( + device const float * logits [[buffer(0)]], + device const float * u [[buffer(1)]], + device int32_t * argmax_o [[buffer(2)]], + device float * entropy_o [[buffer(3)]], + device int32_t * sampled_o [[buffer(4)]], + constant int32_t & n_vocab [[buffer(5)]], + constant float & inv_temp [[buffer(6)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], + uint ntg [[threads_per_threadgroup]]) { + threadgroup float s_val[256]; + threadgroup float s_sum[256]; + threadgroup int32_t s_idx[256]; + threadgroup int32_t s_tok; + + device const float * row_logits = logits + (ulong) row * (ulong) n_vocab; + + float local_max = -INFINITY; + int32_t local_idx = 0; + for (int v = (int) tid; v < n_vocab; v += (int) ntg) { + const float x = row_logits[v] * inv_temp; + if (x > local_max) { local_max = x; local_idx = v; } + } + s_val[tid] = local_max; + s_idx[tid] = local_idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = ntg >> 1; stride > 0; stride >>= 1) { + if (tid < stride && s_val[tid + stride] > s_val[tid]) { + s_val[tid] = s_val[tid + stride]; + s_idx[tid] = s_idx[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float max_l = s_val[0]; + const int32_t amax = s_idx[0]; + + float local_sum = 0.0f; + float local_t = 0.0f; + for (int v = (int) tid; v < n_vocab; v += (int) ntg) { + const float d = row_logits[v] * inv_temp - max_l; + const float e = exp(d); + local_sum += e; + local_t += d * e; + } + s_sum[tid] = local_sum; + s_val[tid] = local_t; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = ntg >> 1; stride > 0; stride >>= 1) { + if (tid < stride) { + s_sum[tid] += s_sum[tid + stride]; + s_val[tid] += s_val[tid + stride]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + const float z = s_sum[0]; + const float t = s_val[0]; + if (tid == 0) { + argmax_o[row] = amax; + entropy_o[row] = log(z) - t / z; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // multinomial draw (first v with cumulative exp(d) >= r, in vocab order): per-thread contiguous slice + // sums, exclusive scan on thread 0 to find the crossing slice, only that thread walks its slice. + const float r = u[row] * z; + const int chunk = (n_vocab + (int) ntg - 1) / (int) ntg; + const int beg = (int) tid * chunk; + const int end = min(beg + chunk, n_vocab); + + float slice_sum = 0.0f; + for (int v = beg; v < end; ++v) { + slice_sum += exp(row_logits[v] * inv_temp - max_l); + } + s_sum[tid] = slice_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tid == 0) { + s_tok = n_vocab - 1; // host default if cum never reaches r (FP guard) + s_idx[0] = -1; // no crossing slice -> no thread walks, default stands + float pref = 0.0f; + for (uint i = 0; i < ntg; ++i) { // exclusive scan + locate the crossing slice + const float next = pref + s_sum[i]; + if (next >= r) { s_idx[0] = (int) i; s_val[0] = pref; break; } + pref = next; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if ((int) tid == s_idx[0]) { // only the crossing thread walks its slice from its prefix + float cum = s_val[0]; + for (int v = beg; v < end; ++v) { + cum += exp(row_logits[v] * inv_temp - max_l); + if (cum >= r) { s_tok = v; break; } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { sampled_o[row] = s_tok; } +} diff --git a/src/models/diffusion-gemma.cpp b/src/models/diffusion-gemma.cpp index 46c610945d3c..3e9596ad05bb 100644 --- a/src/models/diffusion-gemma.cpp +++ b/src/models/diffusion-gemma.cpp @@ -566,10 +566,24 @@ size_t llama_diffusion_debug_get_sc_dev(const struct llama_model * model, float return n; } -// 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); +// Stage-1 device sampling entry. Fetches the backend's dense sampler via the backend-reg proc address +// (keeps the llama<->ggml backend link at the existing boundary) and runs it on sc_dev. Returns false +// for non-DiffusionGemma / no sc_dev / unsupported backends so the caller falls back to the host path. +typedef bool (*dg_dense_sample_fn)(struct ggml_tensor *, const float *, int *, float *, int *, int, float); + +static dg_dense_sample_fn dg_resolve_dense_sample_fn() { + if (ggml_backend_reg_t reg = ggml_backend_reg_by_name("CUDA")) { + if (void * fn = ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_diffusion_sample")) { + return (dg_dense_sample_fn) fn; + } + } + if (ggml_backend_reg_t reg = ggml_backend_reg_by_name("MTL")) { + if (void * fn = ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_diffusion_sample")) { + return (dg_dense_sample_fn) fn; + } + } + 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) { @@ -577,12 +591,7 @@ 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; - } - static dg_cuda_sample_fn fn = - (dg_cuda_sample_fn) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_diffusion_sample"); + static dg_dense_sample_fn fn = dg_resolve_dense_sample_fn(); if (!fn) { return false; }