Skip to content

Carry ggml-org#24423 (DiffusionGemma) merged onto b10630 - #107

Closed
danielhanchen wants to merge 25 commits into
diffusiongemma-24423-upstream-basefrom
diffusiongemma-24423-b10630
Closed

Carry ggml-org#24423 (DiffusionGemma) merged onto b10630#107
danielhanchen wants to merge 25 commits into
diffusiongemma-24423-upstream-basefrom
diffusiongemma-24423-b10630

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

ggml-org/llama.cpp#24423 stopped 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 daca8075d8 with b10630 merged into it and the conflict resolved by hand, so the pin set merges again. The base branch is b10630 itself, 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/--conversation argument. Upstream 9fee29e943 ("arg: remove -no-cnv from cli", ggml-org#27542) removed LLAMA_EXAMPLE_CLI from it; the PR added LLAMA_EXAMPLE_DIFFUSION to the same list.

-    ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_DIFFUSION}));
+    ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_DIFFUSION}));

Resolved by honouring the removal and keeping the addition. DIFFUSION is load-bearing: examples/diffusion/diffusion-cli.cpp parses under LLAMA_EXAMPLE_DIFFUSION and reads params.conversation_mode. CLI is not: conversation_mode appears nowhere in the PR's other new examples, and llama-diffusion-cli never parses under LLAMA_EXAMPLE_CLI.

Verified behaviourally against the built binaries:

binary advertises -cnv
base llama-cli no upstream's removal
carry llama-cli no removal honoured
carry llama-diffusion-cli yes PR intent kept

One extra commit

examples/diffusion-gemma-server/diffusion-gemma-visual-server.cpp did not compile against the base. Upstream d9f918d2d0 ("common: add json.h abstraction", ggml-org#27511) put common_json in front of nlohmann and common_chat_msgs_parse_oaicompat now takes it, so the raw nlohmann::ordered_json the server built no longer converts:

error: call of overloaded 'common_json_value(const nlohmann::...::basic_json<...>&)' is ambiguous

Ported the four call sites to common_json. This does not conflict textually, so nothing would have caught it before the build. It matters because unsloth-prebuilt-*.yml builds llama-diffusion-gemma-visual-server and llama-diffusion-cli best-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), plus LLAMA_BUILD_TESTS=ON which the nightly does not use. Everything was run against the bare b10630 tag as well, and compared.

  • Build clean; all four diffusion binaries produced (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-archs plus the three test-recurrent-state-rollback variants: pass, same as base.
  • Real weights: unsloth/diffusiongemma-26B-A4B-it-GGUF Q4_K_M on a B200, 64 diffusion steps, coherent output, 167.9 tok/s.

Once ggml-org#24423 lands upstream or a base tag carries it, the pin should move back to the upstream commit and this branch can be deleted.

danielhanchen and others added 25 commits June 10, 2026 15:55
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.
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.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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".

Comment on lines +95 to +97
if (N <= 0 || N > MAXTOK) {
printf("ERR badN %d\n", N); fflush(stdout); continue;
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment on lines +546 to +552
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;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment on lines +208 to +210
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);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment on lines +265 to +269
// 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);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment on lines +561 to +563
if (llama_decode(ctx, batch) != 0) {
LOG_ERR("%s: failed to decode at step %d\n", __func__, step_idx);
break;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Comment on lines +566 to +573
// 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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Comment on lines +36 to +39
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];

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

@danielhanchen

Copy link
Copy Markdown
Member Author

Closing. The pin in scripts/unsloth/pr-set.json still points at this branch's head, and the nightly merges pinned commits regardless of PR state, so the release is unaffected. The commit is also mirrored to refs/pins/<sha>, so it survives even if the branch is deleted.

Note that repin.py skips closed PRs, so this pin will no longer be repinned automatically when the base tag moves; it will need a manual merge once it stops applying.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant