metal : add Stage-1 diffusion dense sampler for DiffusionGemma (Apple Silicon)#1
Open
andreinknv wants to merge 1 commit into
Conversation
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 <noreply@anthropic.com>
danielhanchen
pushed a commit
that referenced
this pull request
Jul 3, 2026
* Sycl tp stage1 (#1) * SYCL: tensor parallelism (--split-mode tensor) for dual-GPU Adds the comm_init/comm_free/comm_allreduce_tensor trio that the meta-backend queries via get_proc_address to enable backend-specific all-reduce, mirroring the pattern used by ggml-cuda.cu. For N=2 (the common dual-GPU case) implements a degenerate ring all-reduce with two size-branched paths: * Small (nelem < 32768): FP32 direct memcpy + per-device ADD kernel chained via depends_on(memcpy_event). 4 SYCL submissions/call. * Large (nelem >= 32768): BF16-compressed. Each device compresses FP32 -> BF16 in a local outbox, cross-device memcpys to the peer's inbox (HALF the PCIe bytes), then decompresses + adds into the local FP32 partial. 6 SYCL submissions/call but PCIe bytes halved -- wins for any tensor where PCIe dominates kernel time. Threshold and BF16 path pattern mirror the CUDA NCCL allreduce. Storage: ONE persistent uint8_t buffer per device, 4 * nelem bytes (matches both path layouts: FP32 nelem floats; BF16 outbox+inbox = 2 * nelem uint16_t each). Single alloc+free per device keeps the SYCL pool's strict-LIFO invariant trivial. Initial impl handles N=2 FP32 contiguous tensors. Other cases return false, causing the meta-backend to use its generic butterfly fallback. Per-call sync is intentionally omitted. SYCL in-order queue semantics ensure that the meta-backend's next compute on the same per-device queue waits for our final ADD, and the next allreduce's first op on the same persistent buffer waits via the same queue. Only comm_free does an explicit final wait. OneCCL is NOT used: OneCCL 2021.17 hardcodes single-device-per-process in communicator_impl.hpp:47 (condition devices.size() == 1), which is incompatible with llama.cpp's single-process multi-GPU model. Measured on dual Intel Arc Pro B70 (NEO 26.05.x, oneAPI 2025.3 + DPC++ nightly): Llama-3.3-70B Q4_K_M, -sm tensor -fa 1 -ctk f16 -ctv f16: pp512 = 377.08 t/s (vs 313.65 layer mode = +20.2%) tg128 = 17.40 t/s (vs 9.74 layer mode = +78.6%) Qwen3-Coder-Next-80B-A3B Q3_K_M (MoE): pp512 = 216.56 t/s (vs 156.58 meta-backend butterfly = +38.3%) tg128 = 17.60 t/s (vs 14.31 meta-backend butterfly = +23.0%) Qwen3-4B Q4_K_M: pp64 = 984.51 t/s, tg16 = 49.29 t/s Llama-3.3-70B in SYCL TP now comfortably beats production layer mode on both prefill and decode. Coder-Next-80B-A3B (MoE) also wins on both — the BF16 path is what unlocks the many-medium-allreduces prefill pattern. Build/CMake: no changes. No new dependencies. ~210 lines added across ggml-sycl.h and ggml-sycl.cpp. * Fix comments * documentation update to address PR feedback * Bring over my device-to-device memcpy chagnes * move the dev2dev_memcpy calls to the upstream 7-parameter variety * Fix a typo and remove a trailing whitespace
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This targets the
diffusion-visual-updatesbranch behind ggml-org#24423, adding a Metal port of the CUDA Stage-1 sampler so Apple Silicon runs on-device sampling instead of the per-step host fallback ("on-device sampling unsupported on this backend").What
kernel_diffusion_dense_sample_f32(ggml-metal.metal) — line-for-line structural mirror ofdiffusion_dense_sample_kernelinggml-cuda/diffusion-sampling.cu: one 256-thread threadgroup per canvas row; parallel max→argmax tree reduction; parallel Z and T (T = Σ d·e) reductions with entropy = logZ − T/Z; slice-scanned multinomial first-crossing walk in vocab order. Same reduction structure ⇒ argmax matches the host bit-for-bit, Z/entropy differ only by FP reduction order.ggml_metal_device_diffusion_sample(ggml-metal-device.m/.h) — synchronous one-shot dispatch on the device queue; grow-only shared scratch for u/argmax/entropy/sampled mirroring the CUDA per-device scratch cache; pipeline compiled lazily viaggml_metal_library_compile_pipeline.get_proc_address("ggml_backend_metal_diffusion_sample"), mirroring the CUDA boundary.llama_diffusion_device_sample(diffusion-gemma.cpp) — resolves CUDA first, then MTL, so the CLI gpu-sample-reduce path lights up on Metal with no flag changes.Verification
DG_DEVSAMPLE_CHECK=1ondiffusiongemma-26B-A4B-it-Q4_K_M(M4 Max,-ngl 99), every step of a 15-step run:— the same tolerance profile documented for CUDA-vs-host.
Performance (M4 Max, Q4_K_M, 256-token canvas)
-fa on)The remaining step cost is the MoE forward pass, so the win is bounded (~+19%), but it also frees the CPU from the per-step 256 x n_vocab softmax/entropy reduction.
Happy to adjust style/structure to fit the PR — thanks for the DiffusionGemma work!
🤖 Generated with Claude Code