Skip to content

Carry ggml-org#24423 (DiffusionGemma) onto b10775, and make the arch testable - #177

Closed
danielhanchen wants to merge 30 commits into
base/upstream-de8656bd9from
diffusiongemma-24423-b10775
Closed

Carry ggml-org#24423 (DiffusionGemma) onto b10775, and make the arch testable#177
danielhanchen wants to merge 30 commits into
base/upstream-de8656bd9from
diffusiongemma-24423-b10775

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

Re-carries #107 (DiffusionGemma, ggml-org#24423) onto b10775, and makes the architecture testable for the first time.

The existing pin is 74acc40c, forked from a b10630-era base on 08-26. It still merges, so this is not urgent, but it is five weeks stale and the PR is closed, which is why the nightly log reads #107 is closed without being merged (upstream declined it). Upstream has not declined anything: ggml-org#24423 is still open, and #107 is our own carry branch, closed because that is what carry branches do.

The test gap

DiffusionGemma is one of the features the nightly ships and it is the only one with no coverage of any kind. test-llama-archs excluded it:

if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT || arch == LLM_ARCH_DIFFUSION_GEMMA) {
    return false; // FIXME @ngxson; diffusion-gemma additionally needs canvas/ISWA fixture params
}

so no DiffusionGemma model has ever been constructed by anything.

Why it could not be tested

canvas_length was read as a raw string, bypassing the arch key table:

ml.get_key(std::string("diffusion.canvas_length"), canvas_length, true);

llama_model_saver::add_kv only takes an llm_kv, so no fixture can emit that key. Giving it LLM_KV_DIFFUSION_CANVAS_LENGTH mapped to the same "diffusion.canvas_length" leaves the name on disk unchanged, leaves gguf-py alone (it already has Keys.Diffusion.CANVAS_LENGTH), and brings the last hparam that sidestepped the table back into it.

The fixture

  • the ISWA head lengths, which the loader reads unconditionally and which are separate keys from the non-SWA pair
  • canvas_length of 16, kept well under the 128 tokens the harness decodes, so both sides of the P = n_tokens - canvas_length split are exercised
  • MoE-mandatory: load_arch_tensors always creates the router and expert tensors with no dense fallback, so a config without experts sizes them zero and the loader aborts on ne >= 1
  • encoder list: causal_attn = false makes the whole batch one encoder pass, so n_ubatch must not be capped below n_tokens, exactly as for the other diffusion decoders

With that, save_models emits diffusion-gemma-moe.gguf (9.4 MB) and the model loads and reserves its graph.

What is still not covered, and why

The backend comparison is still skipped, but now for a specific reason instead of a missing fixture. It aborts in the arch's own graph input:

GGML_ASSERT(self_kq_mask && ggml_backend_buffer_is_host(self_kq_mask->buffer)) failed

self_kq_mask->buffer is null, so the mask input was not allocated for the ubatch being decoded. That is in llm_graph_input_attn_diffusion::set_input, not in the fixture, and I have not worked around it: a mask the graph did not allocate is a question for whoever owns the graph. The exclusion comment now records the exact assertion so the next person starts where this stopped.

scripts/unsloth/feature-checks.json lists this pin under unchecked with the same reason, so the nightly says out loud that DiffusionGemma has no runtime proof rather than implying it has one.

Diff

git diff b10775..HEAD is 28 files, +3248/-123: the DiffusionGemma carry, plus the llm_kv id and the test fixture. Base branch is base/upstream-67a17c17c, which is b10775 verbatim.

danielhanchen and others added 27 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.
The arch was excluded from test-llama-archs with "needs canvas/ISWA fixture
params", so nothing has ever constructed a DiffusionGemma model. It is one of
the features the nightly ships and the only one with no coverage of any kind.

canvas_length was read as a raw string:

    ml.get_key(std::string("diffusion.canvas_length"), canvas_length, true);

which is why it could not be tested: llama_model_saver::add_kv only takes an
llm_kv, so no fixture can emit the key. Giving it LLM_KV_DIFFUSION_CANVAS_LENGTH
mapped to the same "diffusion.canvas_length" leaves the name on disk unchanged
and gguf-py alone, and brings the last hparam that bypassed the arch table back
into it.

The fixture supplies the ISWA head lengths, which the loader reads
unconditionally and which are separate keys from the non-SWA pair, plus a
canvas_length of 16, kept well under the 128 tokens the harness decodes so both
sides of the P = n_tokens - canvas_length split are exercised.

The arch is also MoE-mandatory: load_arch_tensors always creates the router and
expert tensors with no dense fallback, so a config without experts sizes them
zero and the loader aborts on ne >= 1. And it joins the encoder list, for the
same reason as the other diffusion decoders: causal_attn = false makes the
whole batch one encoder pass, so n_ubatch must not be capped below n_tokens.

With that, save_models emits diffusion-gemma-moe.gguf and the model loads and
reserves its graph.

The backend comparison is still skipped, now for a specific reason rather than
a missing fixture. It aborts in the arch's own graph input:

  GGML_ASSERT(self_kq_mask && ggml_backend_buffer_is_host(self_kq_mask->buffer))

with self_kq_mask->buffer null, so the mask input was not allocated for the
ubatch being decoded. That is in llm_graph_input_attn_diffusion, not in the
fixture, and it is left to the DiffusionGemma authors rather than worked
around. The exclusion now records exactly that.
@danielhanchen
danielhanchen requested a review from CISC as a code owner September 3, 2026 22:19
@chatgpt-codex-connector

Copy link
Copy Markdown

You have reached your Codex usage limits for security reviews. Please try again later.

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Sep 3, 2026

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review Completed 2026-09-04T00:46:47.979649Z 4137a7d New commits
ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@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: 6851db7ab1

ℹ️ 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 +231 to +232
// canvas_length splits the forward (P = n_tokens - canvas_length); must be positive
ml.get_key(LLM_KV_DIFFUSION_CANVAS_LENGTH, canvas_length);

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 Preserve the required canvas length when saving

This makes diffusion.canvas_length mandatory when loading, but llama_model_saver::add_kv_from_model() never writes LLM_KV_DIFFUSION_CANVAS_LENGTH. Consequently, llama_model_save_to_file() produces a DiffusionGemma GGUF that cannot be loaded again; this is reproducible by saving the new test-llama-archs fixture and reopening it. Add this architecture-specific value to the saver before advertising save/round-trip support.

Useful? React with 👍 / 👎.

Comment on lines +273 to +275
// Stage 1 + Stage 2 are single-device features (sc_dev / prompt-KV store are single-GPU). Auto-enable
// them for one CUDA device, exactly like the CLI's --diffusion-* auto resolution.
base.kv_cache = one_gpu;

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 Keep the visual server on a viable multi-GPU path

When two or more GPU/IGPU backends are present, this disables the KV-cached path, but make_cparams() caps n_ubatch at 2048 and n_outputs_max at the canvas length while the uncached generator submits all P+C tokens and requests logits for every row. For a nonempty prompt, encode() therefore either asserts when P+C > n_ubatch or fails output reservation because P+C > C, so the visual server cannot serve requests on multi-GPU hosts.

Useful? React with 👍 / 👎.

Comment on lines +639 to +642
dm->sc_logits_ptr = sc_logits;
dm->sc_use = use_sc;
dm->sc_temp_inv = temp_inv;
dm->sc_enabled = enabled;

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 Store request-specific diffusion state per context

When multiple contexts share one model, as supported and exercised by tests/test-thread-safety.cpp, concurrent requests overwrite these model-wide pointer and gate fields before their respective graph builds. One context can consequently upload another request's logits, or dereference a local sc_buffer after its owning generation call returns; the phase and persistent KV/SC fields have the same cross-context problem. Keep this mutable request state on llama_context or another per-context object instead of the shared model.

Useful? React with 👍 / 👎.

Comment on lines +28 to +29
def _create_vocab_sentencepiece(self):
tokens, scores, toktypes = super()._create_vocab_sentencepiece()

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 Apply the special-token correction through set_vocab

This override is never called because DiffusionGemmaModel inherits Gemma4Model.set_vocab(), which constructs LlamaHfVocab directly rather than calling _create_vocab_sentencepiece(). Checkpoints needing the stated control-token correction therefore retain the original token types in the converted GGUF, affecting special-token parsing and rendering. Override the vocabulary path that Gemma4 actually uses.

Useful? React with 👍 / 👎.

CI caught what the local run did not: test-save-load-state reads every
gguf that test-llama-archs -o writes, and the diffusion-gemma one does
not load back.

    error loading model hyperparameters:
      key not found in model: diffusion-gemma.attention.sliding_window_pattern

The fixture does set that key. The saver drops it: llama-model-saver.cpp
has

    // add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,  ???);

and the ??? is not laziness. The key means a period when it is a scalar
and per-layer flags when it is an array, different loaders read it each
way, and hparams keeps only the expanded per-layer form, so there is no
single correct thing to write back. diffusion-gemma reads the key as
required, so the round trip fails where an arch with an optional read
would only lose its SWA layout quietly.

Same reason gemma4 is excluded from save_models, so diffusion-gemma is
excluded alongside it, with the reason written down.

Verified locally: 110 fixture models written, none for diffusion-gemma,
and test-save-load-state passes.
@danielhanchen

Copy link
Copy Markdown
Member Author

CI found something the local run could not: test-save-load-state loads every gguf that test-llama-archs -o writes, and the diffusion-gemma one does not load back.

error loading model hyperparameters:
  key not found in model: diffusion-gemma.attention.sliding_window_pattern

The fixture does set that key. The saver drops it, and its own note says why it has to:

// add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN,  ???);

The key means a period as a scalar and per-layer flags as an array, different loaders read it each way, and hparams keeps only the expanded per-layer form, so there is no single correct value to write back. diffusion-gemma reads it as required, so the round trip fails outright rather than quietly losing the SWA layout the way an arch with an optional read would.

So diffusion-gemma is excluded from save_models for the same reason gemma4 already is, with the reason recorded at the exclusion. Verified locally: 110 fixture models written, none for diffusion-gemma, test-save-load-state passes.

Worth saying plainly: with this and the existing test_backends exclusion, the fixture is complete but nothing runs it yet. It is blocked on the arch's own graph input, not on the fixture, and scripts/unsloth/feature-checks.json in #178 lists DiffusionGemma as knowingly unchecked for exactly that reason rather than leaving the gap unrecorded.

@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: 47073dfa09

ℹ️ 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".

} else if ((int) req.size() != 2 + N) {
printf("ERR badsize %d %d\n", N, (int) req.size()); fflush(stdout); continue;
}
if (N <= 0 || N > MAXTOK) {

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 request canvas against the model metadata

Reject requests unless P >= 0 and C equals the model's fixed diffusion.canvas_length. The graph derives its canvas width from GGUF rather than this request header, while sc_cache is sized from the supplied C; for example, a request with a sufficiently long prompt and 0 < C < canvas_length makes llm_graph_input_sc upload canvas_length * n_vocab floats from a C * n_vocab allocation, causing an out-of-bounds host read even when use_sc is zero because the SC graph remains enabled.

Useful? React with 👍 / 👎.

llama_diffusion_set_device_sc(model, false); // restore host SC path for later turns
}
llama_batch_free(batch);
n_generated = params.max_length;

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 Report failed denoising decodes as failures

Set n_generated only after the denoising loop completes successfully. If any step's llama_decode() returns an error, the loop breaks but this unconditional assignment tells the CLI and visual server that a complete canvas was produced; on a first-step failure they consequently detokenize zero-initialized or stale output-buffer contents and may emit DONE instead of reporting generation failure.

Useful? React with 👍 / 👎.


for (int32_t cur_step = S; cur_step >= 1 && !finished; --cur_step) {
const int32_t step_idx = S - cur_step; // 0-based
const float t = params.t_min + (params.t_max - params.t_min) * ((float) cur_step / (float) S);

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 Reach the configured minimum temperature on the last step

Interpolate over the executed step endpoints rather than dividing cur_step by S. The loop ends at cur_step == 1, so this expression uses t_min + (t_max - t_min) / S for the last denoising step and never uses the documented t_min; with the default 48-step schedule the terminal temperature is about 0.4083 instead of 0.4, changing entropy, sampling, acceptance, and reproducibility relative to the configured schedule.

Useful? React with 👍 / 👎.

Comment on lines +442 to +444
} else {
// pure-canvas (no prompt) path
inpL = ggml_rms_norm(ctx0, inpL, hparams.f_norm_rms_eps);

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 Apply self-conditioning to promptless canvas batches

Route the pure-canvas branch through dg_canvas_embed() as the decode and mixed branches do. The raw server explicitly accepts P == 0 and still calls llama_diffusion_set_sc(), but a batch whose size equals the model canvas takes this branch and only applies RMS normalization, so every use_sc=1 request silently ignores the previous step's logits and produces a different denoising trajectory.

Useful? React with 👍 / 👎.

Conflict-free merge, and the tree does not compile: upstream turned
n_ff_exp into a per-layer array behind an accessor, and this file still
reads and uses the old scalar field. Read into n_ff_exp_arr and take the
value through the accessor, keeping the read optional as it was.
@danielhanchen
danielhanchen changed the base branch from base/upstream-67a17c17c to base/upstream-de8656bd9 September 3, 2026 23:46
danielhanchen pushed a commit that referenced this pull request Sep 3, 2026
The base tag moved from b10775 to b10786 while this was open, and five
pins needed work to survive it. Four of the five merge without a single
conflict and produce a tree that does not compile, which is the failure
mode the compile gate in #175 exists for; three of those four are the
same upstream change.

- #172 inkling, #173 glm5next, #177 diffusion-gemma: n_ff_exp became a
  per-layer array behind an accessor. Reading the old scalar field is a
  compile error, and inside a create_tensor dimension list the member
  function quietly decays to a pointer-to-member instead. #172 and #173
  also override preprocess(), which the mtmd base classes made const.
- #152 per-run buffers: b10786 added a load-ordering pass that reads a
  llama_buf_map entry as one buffer, and this pin made an entry a run of
  buffers.
- #144 qwen4exp MTP: the only one that conflicts, in both places it
  touches, over the same n_ff_exp change.

Verified on b10786: all 13 pins merge (11 clean, 2 additive), the CPU
llama target builds, and the CUDA build plus the feature matrix are in
the PR comment.
@danielhanchen

Copy link
Copy Markdown
Member Author

Refreshed onto b10786, which is the base tonight's nightly will pick. The PR base branch moved to base/upstream-de8656bd9, which is that tag verbatim, so the diff here is still this feature and nothing else.

What the move required is in the merge commit message. Verified with the whole 13-pin set replayed on b10786: every pin merges, pin_contract.py reports all 13 intact, the CUDA sm_100 build is clean, and the feature matrix passes all six features on a B200. The full output is on #174.

@chatgpt-codex-connector

Copy link
Copy Markdown

💡 Codex Review

add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp());

P1 Badge Persist the per-layer expert metadata

Saving the newly supported Nemotron Puzzle model collapses n_ff_exp_arr to n_ff_exp(), which selects only layer 0; the same collapse occurs for n_expert_used_arr below. Puzzle relies on heterogeneous values across layers, so a saved model advertises the first layer's dimensions and routing count for every layer, causing later expert tensors to fail shape validation or use the wrong top-k when reloaded. Write these keys as per-layer arrays, reducing them to scalars only when uniform.


if (canvas_length > 0 && params.n_predict > 0) {

P1 Badge Size the default canvas invocation

When -n is omitted, n_predict retains its documented default of -1, so this sizing block is skipped and n_ubatch remains the common default of 512. For any DiffusionGemma whose fixed canvas plus prompt exceeds 512 tokens, the later max_length > max_ub check rejects an otherwise normal CLI invocation before the first decode. The context must be sized for at least one canvas block even when prediction length is unlimited.


const int32_t blocks = (params.n_predict + cl - 1) / cl;
const int32_t needed = blocks * cl + 2048; // + headroom for the prompt / chat history
params.diffusion.blocks = blocks;

P2 Badge Honor the requested prediction-token limit

For a positive -n that is not an exact multiple of canvas_length, rounding up to whole blocks also raises the actual output limit: absent an early EOG or repetition, -n 1 returns an entire canvas and -n C+1 can return 2*C tokens. Since the shared option is documented as the number of tokens to predict, cap the committed final block to the remaining requested token count rather than exposing all rounded-up canvas tokens.

ℹ️ 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".

The exclusion said the backend comparison aborts inside the arch's own
graph input:

    GGML_ASSERT(self_kq_mask && ggml_backend_buffer_is_host(self_kq_mask->buffer))

with a null buffer. That is not an arch bug. build_attn picks the SWA
mask or the plain one per layer, and the fixture set the sliding-window
pattern to a scalar, which fills every layer as sliding. The plain mask
was then built and used by nothing, so ggml-alloc never gave it a
buffer, and set_input filled it anyway. A real DiffusionGemma has both
kinds of layer.

Give it a per-layer pattern, like cohere2moe and the other pattern archs
in the same branch, and the arch decodes:

    | diffusion-gemma|      NVIDIA B200|MoE|  OK (9.45e-08)|SKIP|
    | diffusion-gemma|Xeon Platinum 8559C|MoE|  OK (0.00e+00)|SKIP|
    | diffusion-gemma|             Meta|MoE|  OK (9.45e-08)|SKIP|

Roundtrip stays SKIP, and that is the second half of this. The saver has
no way to write attention.sliding_window_pattern back: the key is a
period as a scalar and per-layer flags as an array, and hparams keeps
only the expanded form. Every other arch reading it per layer is already
on llama_model_saver_supports_arch's list for that reason, including
dots3note with the same note; diffusion-gemma belongs there too. That
also drops it from save_models, which is what test-save-load-state reads,
so the separate exclusion added there earlier is now redundant and gone.

Verified: 110 fixture models written, none for diffusion-gemma, and
test-save-load-state passes.
danielhanchen pushed a commit that referenced this pull request Sep 4, 2026
The exclusion blamed the arch's graph input for a null self_kq_mask
buffer. It was the fixture: a scalar sliding-window pattern makes every
layer sliding, so the plain mask is built and used by nothing and never
gets allocated. With a per-layer pattern the arch decodes and matches
CPU on all three devices.
@danielhanchen

Copy link
Copy Markdown
Member Author

DiffusionGemma runs now. The exclusion was wrong about where the fault was, and I would rather say that plainly than leave it standing.

The assertion is real:

GGML_ASSERT(self_kq_mask && ggml_backend_buffer_is_host(self_kq_mask->buffer))

but the cause is the fixture, not the arch. build_attn picks the SWA mask or the plain one per layer. The fixture set the sliding-window pattern as a scalar, which fills every layer as sliding, so the plain mask was built, used by no op, never allocated by ggml-alloc, and then filled by set_input through a null buffer. A real DiffusionGemma has both kinds of layer, and so should the fixture. Giving it a per-layer pattern, exactly like cohere2moe and the other pattern archs in the same branch:

|     Model arch.|                         Device|Config|   NMSE vs. CPU|Roundtrip|
| diffusion-gemma|                    NVIDIA B200|   MoE|  OK (9.45e-08)|     SKIP|
| diffusion-gemma|Intel(R) Xeon(R) Platinum 8559C|   MoE|  OK (0.00e+00)|     SKIP|
| diffusion-gemma|                           Meta|   MoE|  OK (9.45e-08)|     SKIP|

Roundtrip stays SKIP, and that is the same story from the other end. The saver cannot write attention.sliding_window_pattern back: the key is a period as a scalar and per-layer flags as an array, and hparams keeps only the expanded form. Every arch that reads it per layer is already on llama_model_saver_supports_arch's list for that reason, dots3note with the same note in the same words. diffusion-gemma belongs there, so it goes there. That also drops it from save_models, which is what test-save-load-state reads, so the separate exclusion I added there earlier is redundant and is gone again.

Verified: 110 fixture models written, none for diffusion-gemma, test-save-load-state passes, and the arch decodes on all three devices.

DiffusionGemma moves out of unchecked and into features in #178 as a result. That map exists so a gap has to be argued for, and this one no longer can be.

@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: 4137a7d366

ℹ️ 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 +109 to +110
if (!use_kv) {
// UNIFIED forward over [prompt | canvas] (default, recomputes the prompt every step)

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 Reset the KV phase before pure-canvas requests

When DG_KVCACHE=1, a prompt-bearing request leaves the model in PKV_DECODE; if a later request has P == 0, this branch is documented as the unified pure-canvas path but never calls llama_diffusion_set_phase(...PKV_UNIFIED...). The graph therefore reuses the previous request's cached prompt length and K/V tensors, producing logits conditioned on stale prompt data with incorrect positions. Reset the phase before executing this branch.

Useful? React with 👍 / 👎.

Comment thread src/llama-context.cpp
Comment on lines 1508 to +1511
GGML_ASSERT(backend_res != nullptr);
GGML_ASSERT(logits.data != nullptr);

ggml_backend_tensor_get_async(backend_res, t_logits, logits.data, 0, n_tokens*n_vocab*sizeof(float));
ggml_backend_tensor_get_async(backend_res, t_logits, logits.data, 0, n_outputs_enc*n_vocab*sizeof(float));

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-side sampling

With gpu_sample_reduce enabled, the entropy-bound path calls llama_decode(), synchronizes, and then samples from sc_dev, but this encode path has already queued an unconditional device-to-host copy of all C * n_vocab logits. Since output_reserve() always supplies logits.data, the synchronization waits for the same roughly 268 MB transfer that device-side sampling is intended to eliminate on every denoising step, defeating the feature's main performance benefit.

Useful? React with 👍 / 👎.

danielhanchen pushed a commit that referenced this pull request Sep 4, 2026
The preflight on this branch reports that #172 does not merge onto
b10786, and it is right: #172 and #177 both add a fixture block to
tests/test-llama-archs.cpp and share exactly one line, `}`, which
master's additive_merge.py reads as the same change made twice.

#170 fixes that and lands first. Merging it here so the preflight on
this branch tests the combination that will actually exist on master,
rather than a state nobody will ever run.
@danielhanchen

Copy link
Copy Markdown
Member Author

Retired. Every piece of this carry is now in ggml-org#24423 itself (c6f8d604b):

  • the DiffusionGemma fixture hparams (SWA head lengths and diffusion.canvas_length)
  • the per-layer SWA pattern branch, since with a scalar pattern every layer slides, the non-SWA mask is built and used by nothing, ggml-alloc never gives it a buffer, and set_input dereferences a null one
  • the moe_mandatory entry and removal of the three GEMMA4 || GEMMA4_ASSISTANT || DIFFUSION_GEMMA exclusions
  • the llama_model_saver deny-list arm, so the save/reload half stops failing on key not found in model: diffusion-gemma.attention.sliding_window_pattern
  • LLM_KV_DIFFUSION_CANVAS_LENGTH as a registered key rather than a raw string lookup
  • the encode list entry: the arch sets causal_attn = false, so the whole batch is one encoder pass and n_ubatch must not be capped to 64, otherwise the harness aborts on encoder requires n_ubatch >= n_tokens

Verified on the merged pin set against b10795: test-llama-archs -a diffusion-gemma reports OK (0.00e+00) where it previously printed an empty table and exited 0. Full arch sweep on the upstream head is rc=0 and gemma4 stays excluded.

The pin now points at ggml-org#24423 directly (#174).

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