Carry ggml-org#24423 (DiffusionGemma) merged onto b10630 - #107
Carry ggml-org#24423 (DiffusionGemma) merged onto b10630#107danielhanchen wants to merge 25 commits into
Conversation
Some diffusion cli and visual updates
…s, drop debug hooks - guard sys/ioctl.h behind _WIN32 and add a GetConsoleScreenBufferInfo fallback for the visual viewport size, so diffusion-cli builds on Windows - skip diffusion-gemma in test-llama-archs like gemma4 (shared ISWA backbone, no synthetic fixture params yet) - remove the DG_DUMP_KV_LAYER / DG_NSWA debug scaffolding and its llama.h API - fix flake8 E306 in conversion/diffusion_gemma.py
The runner sizes n_ctx/n_ubatch/n_batch from -n and the canvas and loads the model directly instead of going through common_init_from_params, so --fit was silently ignored. Print a one-line notice pointing at -ngl / --n-cpu-moe for controlling device memory.
…to model params) The CLI hand-builds llama_model_params and never copied tensor_buft_overrides, so -ot and --n-cpu-moe were parsed but silently dropped - the MoE experts stayed on the GPU and OOMed small-VRAM cards. Mirror common_model_params_to_llama.
… /clear
- --diffusion-gpu-sampling {auto,on,off} (default auto = on for single-GPU):
keep the prev step's canvas logits in a device buffer (sc_dev) and read
self-conditioning from it instead of a 268 MB host upload each step. SC
inputs are bit-identical to the host path; auto-disables on multi-GPU like
--diffusion-kv-cache. ~1.3x per step.
- cli: add effective + in-step-parallel throughput to the timing summary.
- cli: add /help and /clear in conversation mode.
Sample argmax/entropy/multinomial per canvas position directly from the
device sc_dev buffer instead of copying the [C, n_vocab] canvas logits to
host (268 MB/step) and reducing on the CPU. Removes the last per-step bus
copy on the entropy-bound path.
- new ggml-cuda kernel (dense, top_k==0), reached from llama via the
backend-reg proc-address boundary (no new llama<->cuda link); falls back
to the host path on non-CUDA / multi-GPU / no sc_dev.
- --diffusion-gpu-sample-reduce {auto,on,off}, auto=on for single-GPU,
requires --diffusion-gpu-sampling. byte-identical when off.
- argmax bit-identical to host every step; Z/entropy differ only by the
parallel-reduction order (~1e-4), same FP-equivalence class as
--diffusion-kv-cache. greedy decode identical; stochastic output
identical on every prompt tested. ~1.42x per step on B200 Q8_0.
cudaPointerGetAttributes / cudaPointerAttributes / cudaMemoryTypeDevice are not mapped by the hip/musa vendor layer. Drop the pointer-attribute device probe (the sampler is gated to a single CUDA device, so the tensor is already on the current device) and route the runtime calls through CUDA_CHECK.
Persistent forward server that runs diffusion_generate_entropy_bound and streams the per-step argmax canvas (plus each committed block) over stdin/stdout, so a host can render the denoise without reloading the model. Reuses the entropy-bound decoder; links llama-diffusion.
Take chat messages as JSON and apply the GGUF chat template + tokenizer in
the server (common_chat_templates + common_tokenize), and stream the per-step
canvas and committed blocks back as detokenized text. Drops the need for any
client-side tokenizer; the request is now {seed, n_blocks, messages}.
When a backend cannot run the on-device sampler (e.g. Metal), latch the fallback after the first failure: warn once and use the host reduction for the rest of the run instead of retrying and logging an error every step. Output is unchanged (host sampling was already the fallback); only the per-step error spam is removed.
…oolong budget - time the template/tokenize and denoise phases and emit a STATS summary (prompt_n, predicted_n, ms, blocks, steps) before DONE - when MAXTOK is unset/0, probe the largest non-causal context that fits VRAM (capped at the training context); report it on the READY line - ERR toolong now carries both the needed token count and the budget
…verhead in STATS The STATS line reported prompt_per_second from the host tokenize wall (~2ms), yielding a meaningless ~14000 tok/s, and the decode wall folded in the per-step frame emission (detok + json + flush). Time the visualization separately and emit prompt_prepare_ms, wall_ms and decode_ms so the shim can derive honest throughput. STATS stays additive; READY/F/C/DONE unchanged.
The visual/server/eval mains call llama_backend_init() but not ggml_backend_load_all(), so on GGML_BACKEND_DL builds no GPU backend registers and NGL is ignored, running on CPU.
Remove the env-gated device-vs-host diff harness from the denoise loop and the debug-only llama_diffusion_debug_get_sc_dev export it used. These compared the on-device sampler/SC buffer against the host path during bring-up and are not needed at runtime.
…ize fallback - detok the committed answer with special=true so the <|channel>thought ... <channel|> markers survive for the client to split reasoning from the answer - if no context meets the VRAM headroom margin, reuse the floor context when it allocates (with a warning) and report free/total VRAM instead of a bare failure
…lls off the GPU The auto-sizer gated every candidate context on free VRAM, so a small GPU running the model from system RAM (NGL exceeds what fits) collapsed to the 2048 floor even though much larger contexts allocate fine in RAM. Probe the VRAM budget first (unchanged on ample VRAM); when it finds nothing, re-probe against a RAM budget (free RAM minus resident weights) and keep the largest context that actually allocates. Explicit MAXTOK now degrades through the same probe instead of hard-failing the runner.
…anvas The visual server forced n_ubatch == n_ctx, so the whole prompt went through one non-causal encode. The O(prompt^2) attention overflows the 32-bit CUDA softcap index past ~12k tokens (n_head * N^2 > 2^31, a crash), and the encode also built an [n_tokens, n_vocab] fp32 logits buffer it then discarded. - llama-context: encode() honors cparams.n_outputs_max and reserves/copies only the flagged rows (no-op when n_outputs_max >= n_tokens). - diffusion-gemma: prefill the prompt in n_ubatch-sized causal chunks into a grow-only K/V store at an offset; off=0 is the single-shot prefill. - visual server: cap n_outputs_max to the canvas and size the prefill chunk so n_head * chunk * n_ctx stays under 2^31 (2048 up to ~32k, smaller past that). The per-turn compute buffer is now flat ~566 MiB regardless of context, output is byte-identical when the prompt fits one ubatch, and prompts to 60k+ tokens work where the single-shot encode crashed. Chunked-prefill approach from potto007.
…dates # Conflicts: # include/llama.h
The Stage-1 device sampler looked up the ggml-cuda backend by the literal name "CUDA", so on HIP and MUSA builds (registered as "ROCm"/"MUSA") the lookup failed and every step fell back to the host logits path. The backend exports the same ggml_backend_cuda_diffusion_sample proc address from shared source regardless of build, so probe ROCm and MUSA as well. Reported by aaronsb.
# Conflicts: # common/arg.cpp
Upstream d9f918d (ggml-org#27511) put a common_json wrapper in front of nlohmann and common_chat_msgs_parse_oaicompat now takes it, so the raw ordered_json the server built no longer converts.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 74acc40c37
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (N <= 0 || N > MAXTOK) { | ||
| printf("ERR badN %d\n", N); fflush(stdout); continue; | ||
| } |
There was a problem hiding this comment.
Validate the requested canvas dimensions
Reject requests unless P >= 0 and C equals the model's diffusion.canvas_length, not merely when their sum is invalid. The graph always splits the sequence using the metadata canvas length, while this server sizes sc_cache and chooses output rows using the request's C; a smaller positive C therefore makes the self-conditioning upload read past sc_cache, and a negative C can turn the allocation on line 100 into a huge size_t allocation.
Useful? React with 👍 / 👎.
| batch.n_tokens = params.max_length; | ||
| for (int32_t i = 0; i < params.max_length; i++) { | ||
| batch.token[i] = (i < n_input) ? input_tokens[i] : current_canvas[i - n_input]; | ||
| batch.pos[i] = i; | ||
| batch.n_seq_id[i] = 1; | ||
| batch.seq_id[i][0] = 0; | ||
| batch.logits[i] = 1; |
There was a problem hiding this comment.
Request logits only for canvas rows
In the visual server's multi-GPU path, base.kv_cache is false while its context sets n_outputs_max to canvas_length. This branch nevertheless marks all P + C rows as outputs, so encode() calls output_reserve(P + C) and hits GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max) whenever the prompt is nonempty. Mark only i >= n_input as output in the unified path.
Useful? React with 👍 / 👎.
| params.n_ubatch = std::max(params.n_ubatch, needed); | ||
| params.n_batch = std::max(params.n_batch, params.n_ubatch); | ||
| params.n_ctx = std::max(params.n_ctx, needed); |
There was a problem hiding this comment.
Cap the CLI output reserve for canvas models
Growing n_ubatch and n_batch here without setting ctx_params.n_outputs_max leaves the new capped-output path disabled. During cached prompt prefill, encode() consequently overrides the single requested logit row and reserves/downloads logits for the entire chunk; with DiffusionGemma's roughly 262k-token vocabulary, a 2048-token prompt consumes about 2 GiB just for host logits and can make context creation or generation fail. Set the cap to the canvas length and retain only the actually needed output flags.
Useful? React with 👍 / 👎.
| // self-conditioning gated MLP (optional; unused in the zero-SC exactness forward) | ||
| sc_pre_norm = create_tensor(tn(LLM_TENSOR_SC_PRE_NORM, "weight"), {n_embd}, TENSOR_NOT_REQUIRED); | ||
| sc_gate = create_tensor(tn(LLM_TENSOR_SC_GATE, "weight"), {n_embd, n_ff}, TENSOR_NOT_REQUIRED); | ||
| sc_up = create_tensor(tn(LLM_TENSOR_SC_UP, "weight"), {n_embd, n_ff}, TENSOR_NOT_REQUIRED); | ||
| sc_down = create_tensor(tn(LLM_TENSOR_SC_DOWN, "weight"), {n_ff, n_embd}, TENSOR_NOT_REQUIRED); |
There was a problem hiding this comment.
Require the self-conditioning tensors used by the CLI
These tensors are accepted as optional, but every canvas model in diffusion-cli.cpp enables self-conditioning before context creation, and graph construction unconditionally passes sc_gate, sc_up, and sc_down to ggml_mul_mat. A GGUF missing any one of these weights therefore loads successfully and then aborts while the context is being reserved; either make the four tensors mandatory for this architecture or gate the self-conditioning graph on all four being present.
Useful? React with 👍 / 👎.
| if (llama_decode(ctx, batch) != 0) { | ||
| LOG_ERR("%s: failed to decode at step %d\n", __func__, step_idx); | ||
| break; |
There was a problem hiding this comment.
Report a failed denoising decode as failure
When llama_decode() fails, the loop breaks but the function still sets n_generated = params.max_length at exit. Both the CLI and visual server interpret that as success and consume an uninitialized or previous-block canvas, producing bogus output after backend errors or OOM. Preserve n_generated == 0 on this path and only report the full length after at least one successful final canvas.
Useful? React with 👍 / 👎.
| // Stage-1: when on, skip the 268 MB logits D2H + host reductions and sample on the GPU from sc_dev. | ||
| const bool gpu_reduce = dev_sc && device_sample_ok; | ||
| const bool want_logits = !gpu_reduce; | ||
| const float * logits = nullptr; // canvas rows packed: [C or max_length, n_vocab] | ||
| if (want_logits) { | ||
| logits = llama_get_logits(ctx); | ||
| } else { | ||
| llama_synchronize(ctx); // sc_dev write must complete before we read it |
There was a problem hiding this comment.
Avoid downloading logits before device sampling
Skipping llama_get_logits() here does not skip the advertised full-canvas device-to-host transfer: llama_context::encode() unconditionally schedules ggml_backend_tensor_get_async() for every requested logit row before llama_decode() returns. Thus the device-reduction path still transfers the approximately 268 MB logits buffer and then adds its own kernel and small transfers, defeating its primary performance purpose. The decode/output path needs an explicit way to retain these logits only on-device.
Useful? React with 👍 / 👎.
| for (int stride = blockDim.x >> 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]; |
There was a problem hiding this comment.
Preserve lowest-token tie breaking in CUDA argmax
The host argmax keeps the first vocabulary index on equal logits, but this reduction compares only values. For example, equal maxima at indices 1 and 256 leave index 256 in thread 0 and index 1 in thread 1, after which the equality retains thread 0's index. Exact ties can therefore change the output canvas and stability decision when GPU reduction is enabled; compare indices as a secondary key so ties select the lower token id.
Useful? React with 👍 / 👎.
|
Closing. The pin in Note that |
ggml-org/llama.cpp#24423stopped merging onto the aged base tag, which killed the whole nightly mix (the merge loop stops at the first conflict, so the other four pins never got applied either).This branch is
daca8075d8withb10630merged into it and the conflict resolved by hand, so the pin set merges again. The base branch isb10630itself, so the diff here is only the DiffusionGemma work, matching the arrangement used by #70, #91 and #99.The conflict
common/arg.cpp, the-cnv/--conversationargument. Upstream9fee29e943("arg: remove -no-cnv from cli", ggml-org#27542) removedLLAMA_EXAMPLE_CLIfrom it; the PR addedLLAMA_EXAMPLE_DIFFUSIONto the same list.Resolved by honouring the removal and keeping the addition.
DIFFUSIONis load-bearing:examples/diffusion/diffusion-cli.cppparses underLLAMA_EXAMPLE_DIFFUSIONand readsparams.conversation_mode.CLIis not:conversation_modeappears nowhere in the PR's other new examples, andllama-diffusion-clinever parses underLLAMA_EXAMPLE_CLI.Verified behaviourally against the built binaries:
-cnvllama-clillama-clillama-diffusion-cliOne extra commit
examples/diffusion-gemma-server/diffusion-gemma-visual-server.cppdid not compile against the base. Upstreamd9f918d2d0("common: add json.h abstraction", ggml-org#27511) putcommon_jsonin front of nlohmann andcommon_chat_msgs_parse_oaicompatnow takes it, so the rawnlohmann::ordered_jsonthe server built no longer converts:Ported the four call sites to
common_json. This does not conflict textually, so nothing would have caught it before the build. It matters becauseunsloth-prebuilt-*.ymlbuildsllama-diffusion-gemma-visual-serverandllama-diffusion-clibest-effort on every platform and never fails the job, so this would have silently shipped bundles missing both binaries.Verification
Built in the nightly's configuration (
GGML_BACKEND_DL=ON,GGML_CPU_ALL_VARIANTS=ON,GGML_RPC=ON,GGML_CUDA=ON, tools + server), plusLLAMA_BUILD_TESTS=ONwhich the nightly does not use. Everything was run against the bareb10630tag as well, and compared.llama-diffusion-cli,llama-diffusion-gemma-eval,llama-diffusion-gemma-server,llama-diffusion-gemma-visual-server).test-backend-ops test -b CUDA0 -o MUL_MAT,MUL_MAT_ID,FLASH_ATTN_EXT,FLASH_ATTN_EXT_BANDED: 4998/4998, identical to base.test-llama-archsplus the threetest-recurrent-state-rollbackvariants: pass, same as base.unsloth/diffusiongemma-26B-A4B-it-GGUFQ4_K_M on a B200, 64 diffusion steps, coherent output, 167.9 tok/s.Once
ggml-org#24423lands upstream or a base tag carries it, the pin should move back to the upstream commit and this branch can be deleted.