Skip to content

metal : add Stage-1 diffusion dense sampler for DiffusionGemma (Apple Silicon) #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andreinknv wants to merge 1 commit into
base: diffusion-visual-updates
Choose a base branch
from
Open

metal : add Stage-1 diffusion dense sampler for DiffusionGemma (Apple Silicon) #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-device.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
//
Expand Down
104 changes: 104 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-device.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<MTLBuffer> 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<MTLDevice> mtl_dev = (id<MTLDevice>) 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<MTLCommandQueue> queue = (id<MTLCommandQueue>) ggml_metal_device_get_queue(dev);
id<MTLCommandBuffer> 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;
}
15 changes: 15 additions & 0 deletions ggml/src/ggml-metal/ggml-metal.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
104 changes: 104 additions & 0 deletions ggml/src/ggml-metal/ggml-metal.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10733,3 +10733,107 @@ kernel void kernel_count_equal(
typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;

template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;

// 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; }
}
29 changes: 19 additions & 10 deletions src/models/diffusion-gemma.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -566,23 +566,32 @@ 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) {
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;
}
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;
}
Expand Down