feat: implement causal prompt prefill

Author: dirmeier

README.md

@@ -1,23 +1,65 @@
 # d3pm
 
-[![ci](https://github.com/dirmeier/block-diffusion-mlx/actions/workflows/ci.yaml/badge.svg)](https://github.com/dirmeier/block-diffusion-mlx/actions/workflows/ci.yaml)
+[![ci](https://github.com/dirmeier/d3pm-mlx/actions/workflows/ci.yaml/badge.svg)](https://github.com/dirmeier/d3pm-mlx/actions/workflows/ci.yaml)
 
 > Discrete denoising diffusion probabilistic models in MLX
 
+`d3pm` is a small library for discrete denoising diffusion probabilistic models (D3PM) in
+[MLX](https://github.com/ml-explore/mlx).
+
+> [!IMPORTANT]
+> ### 🚀 DiffusionGemma
+> 🔥 ** `d3pm` now features a 💎 DiffusionGemma 💎 example.**
+> DiffusionGemma trains on `tiny_shakespeare` and generates text in a block-autoregressive fashion with a Gemma-4 MoE backbone, uniform (mask-free) diffusion, self-conditioning, and D-CFG guidance. It supports prompt-conditioned generation via causal prompt prefill with bidirectional denoising.
+
+
 ## Examples
 
-Self-contained examples can be found in [examples](examples/)
+Self-contained examples can be found in [examples](examples/).
 
 ## Installation
 
 To install the latest GitHub <RELEASE>, just call the following on the command line:
 
 ```bash
-pip install git+https://github.com/dirmeier/block-diffusion-mlx@<RELEASE>
+pip install git+https://github.com/dirmeier/d3pm@<RELEASE>
 ```
 
 Please do not directly install from the `main` branch, since these commits are not tested, but prefer using PyPI or a stable release.
 
+## Contributing
+
+Contributions in the form of pull requests are more than welcome. A good way to
+start is to check out issues labelled
+[good first issue](https://github.com/dirmeier/d3pm-mlx/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22).
+
+In order to contribute:
+
+1) Clone `d3pm-mlx` and install `uv` from [here](https://docs.astral.sh/uv/getting-started/installation/).
+2) Install all dependencies using `uv sync --all-groups`.
+3) Install the Git hooks:
+
+   ```bash
+   uv run pre-commit install -t pre-commit -t commit-msg
+   ```
+4) Create a new branch locally, e.g. `git checkout -b feature/my-new-feature`.
+5) Implement your contribution and ideally a test case.
+6) Check your work (see below).
+7) Submit a PR 🙂.
+
+### Development commands
+
+The project uses `uv` for everything:
+
+```bash
+uv sync --all-groups
+uv run ruff format src examples
+uv run ruff check --fix src examples
+uv run mypy src examples
+uv run pre-commit run --all-files
+```
+
+
 ## References
 
 - Arriola et al., [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://arxiv.org/abs/2503.09573) (BD3LM), 2025.

examples/diffusiongemma/README.md

@@ -1,80 +1,41 @@
-This repository implements the [DiffusionGemma](https://deepmind.google/models/gemma/diffusiongemma/)
-denoiser in [MLX](https://github.com/ml-explore/mlx): a *uniform-state* discrete
-diffusion language model ([UDLM](https://arxiv.org/abs/2412.10193)) on a Gemma-4
-MoE backbone, intended to sit inside the *block-diffusion*
-([BD3LM](https://arxiv.org/abs/2503.09573)) generation framework.
+# DiffusionGemma
+
+This example trains [DiffusionGemma](https://deepmind.google/models/gemma/diffusiongemma/)
+on `tiny_shakespeare`.
 
-You can find a minimal experiment in [experiments/diffusiongemma](experiments/diffusiongemma) where we train DiffusionGemma on `tiny_shakespeare`.
-For reference on DiffusionGemma itself, see the
-[DeepMind implementation](https://github.com/google-deepmind/gemma/tree/main/gemma/diffusion),
-the [vLLM writeup](https://github.com/vllm-project/vllm-project.github.io/blob/main/_posts/2026-06-10-diffusion-gemma.md),
-and the [NeMo fine-tuning guide](https://github.com/NVIDIA-NeMo/Automodel/blob/main/docs/guides/dllm/diffusiongemma.md).
 To run the experiment, first download the latest release and install all dependencies via:
 
 ```bash
-wget -qO- https://github.com/dirmeier/block-diffusion-mlx/archive/refs/tags/<TAG>.tar.gz | tar zxvf -
+wget -qO- https://github.com/dirmeier/d3pm-mlx/archive/refs/tags/<TAG>.tar.gz | tar zxvf -
 uv sync --all-groups
 ```
 
 To train a model and generate some text, call:
 
 ```bash
-uv run python experiments/diffusiongemma/main.py
+uv run python examples/diffusiongemma/main.py
 ```
 
-Pass `--prompt "..."` to condition generation on a prefix, and `--entropy-bound`
-to set the sampler's cumulative acceptance budget. The budget does not scale
-with canvas length, so a small or under-trained model on a wide `--block-size`
-needs a looser bound (raise it if samples stay noisy).
-
-
 ## Method
 
-DiffusionGemma combines two lines of work from the Kuleshov group: the
-block-autoregressive *structure* of BD3LM with the uniform-state *noise* of
-UDLM (plus self-conditioning and an entropy-bounded sampler). The axes below
-show where each model — and this repository — sits.
-
-| Axis | MDLM | BD3LM | UDLM | Duo | DiffusionGemma | This repo |
-|------|------|-------|------|-----|----------------|-----------|
-| Noise state | masked | masked | uniform | uniform | uniform (no mask) | **uniform (no mask)** |
-| Noise schedule | loglinear | loglinear | loglinear | loglinear | — | **linear or loglinear (+ importance sampling)** |
-| Generation | single block | block-autoregressive | single block | single block | block-autoregressive | **block-autoregressive** |
-| Attention | bidirectional | block-causal | bidirectional | bidirectional | causal prefill + bidir. denoise | **block-causal + bidir. denoise** |
-| Self-conditioning | no | no | no | yes | yes | **yes (`encode_logits` signal, GeGLU FFN + post-norm)** |
-| Guidance | — | — | D-CFG | — | D-CFG | **D-CFG (toy label)** |
-| Variable length | no | yes | no | no | yes | **yes (growing KV cache)** |
-| Backbone | encoder | encoder | encoder | encoder | Gemma-4 MoE | **Gemma-4 MoE** |
+The axes below show the model details of the original DiffusionGemma implementation (as well as
+several models), and how they compare to `d3pm`'s implementation.
 
-The block-autoregressive loop is implemented: the sampler commits each denoised
-block to a growing KV cache and starts the next block conditioned on that
-history (`BlockSampler`), and training scores multiple blocks under a
-block-causal mask (`block_diffusion_loss`). Two invariants from DeepMind's
-reference hold — attention is **bidirectional within a block** (every query
-attends to the whole current block, no triangular structure) and
-**block-causal across blocks** (block `i` attends to blocks `j <= i`, never
-future blocks). That cross-block causality is realised either by an explicit
-block-causal mask (parallel processing / training) or implicitly by the
-streaming sampler, where future blocks are simply not in the cache yet; the
-variable length comes from appending keys/values to the growing cache. An
-optional prompt is committed to the cache first (block-causal prefill), so
-generation continues from it (`BlockSampler.generate(prompt=...)`).
-
-### Gemma-4 parity
-
-The backbone tracks the reference Gemma-4 MoE block: interleaved
-local-sliding / global attention with dual RoPE bases, QK- and value-norm, a
-per-block `skip_scale`, a normed-and-scaled MoE router with per-expert scaling
-and a shared dense branch, and projection-only per-layer embeddings. The
-training recipe uses a UDLM-faithful stratified/antithetic noise sampler with an
-optional loglinear schedule and importance-sampling reparametrisation
-(transcribed from MDLM `LogLinearNoise`), and the sampler uses an entropy-budget
-token-acceptance rule with annealed temperature. Self-conditioning feeds the
-previous step's predicted distribution back through the embedding table
-(`encode_logits`), matching the reference.
-
-**Deferred (documented):** EMA weights, Duo consistency distillation
-([2506.10892](https://arxiv.org/abs/2506.10892)) for few-step sampling, a real
-(SentencePiece) tokenizer, and batched multi-sequence sampling.
+| Axis | MDLM | BD3LM | UDLM | Duo | DiffusionGemma | `d3pm` package |
+|------|------|-------|------|-----|----------------|------|
+| Noise state | masked | masked | uniform | uniform | uniform (no mask) | uniform (no mask) |
+| Noise schedule | loglinear | loglinear | loglinear | loglinear | ? | linear or loglinear |
+| Generation | single block | block-autoregressive | single block | single block | block-autoregressive | block-autoregressive |
+| Attention | bidirectional | block-causal | bidirectional | bidirectional | causal prefill + bidir. denoise | causal prefill + bidir. denoise |
+| Self-conditioning | no | no | no | yes | yes | yes |
+| Guidance | — | — | D-CFG | — | D-CFG | D-CFG |
+| Variable length | no | yes | no | no | yes | yes |
+| Backbone | encoder | encoder | encoder | encoder | Gemma-4 MoE | Gemma-4 MoE |
 
 ## Additional references
+
+- [DeepMind reference implementation](https://github.com/google-deepmind/gemma/tree/main/gemma/diffusion)
+- [DiffusionGemma model card](https://ai.google.dev/gemma/docs/diffusiongemma/model_card)
+- [vLLM integration write-up](https://github.com/vllm-project/vllm-project.github.io/blob/main/_posts/2026-06-10-diffusion-gemma.md)
+- [NeMo fine-tuning guide](https://github.com/NVIDIA-NeMo/Automodel/blob/main/docs/guides/dllm/diffusiongemma.md)
+- [How to Build a Diffusion Language Model](https://kuleshov-group.github.io/blog/blog/2026/how-to-build-a-diffusion-language-model/), 2026.

examples/diffusiongemma/labels.py

@@ -11,7 +11,7 @@
 import mlx.optimizers as optim
 from mlx import nn
 
-import d3pm as dg
+import d3pm
 
 if TYPE_CHECKING:
   from collections.abc import Callable, Generator
@@ -22,6 +22,17 @@
 )
 
 
+def uppercase_label(
+  x0: mx.array, block_size: int, threshold: float = 0.3
+) -> mx.array:
+  """Derive a per-block uppercase label for CFG conditioning."""
+  b, t = x0.shape
+  nb = t // block_size
+  is_upper = (x0 >= ord("A")) & (x0 <= ord("Z"))
+  frac = is_upper.reshape(b, nb, block_size).mean(axis=-1)
+  return (frac > threshold).astype(mx.int32)
+
+
 def load_corpus(path: Path) -> str:
   """Download tiny_shakespeare to ``path`` if missing and read it."""
   if not path.exists():
@@ -66,51 +77,50 @@ def main() -> None:
   parser.add_argument("--lr", type=float, default=1e-3)
   parser.add_argument("--warmup", type=int, default=100)
   parser.add_argument("--prompt", type=str, default="Romeo, my Romeo!")
-  # Total (not per-position) entropy budget for token acceptance; the
-  # reference uses 0.1 for a confident 26B model, but a small under-trained
-  # model on a wide canvas needs a looser budget to accept anything.
   parser.add_argument("--entropy-bound", type=float, default=1.0)
+  parser.add_argument("--prompt-len", type=int, default=0)
   args = parser.parse_args()
 
-  tokenizer = dg.ByteTokenizer()
+  tokenizer = d3pm.ByteTokenizer()
   text = load_corpus(Path("input.txt"))
   ids = tokenizer.encode(text)
 
-  cfg = dg.ModelConfig(
+  cfg = d3pm.ModelConfig(
     vocab_size=tokenizer.vocab_size,
     d_model=256,
     num_layers=6,
     num_heads=8,
     num_kv_heads=2,
     head_dim=32,
-    moe=dg.MoEConfig(num_experts=4, top_k=2, ffn_hidden=512, dense_hidden=256),
+    moe=d3pm.MoEConfig(
+      num_experts=4, top_k=2, ffn_hidden=512, dense_hidden=256
+    ),
     max_seq_len=args.seq_len,
     block_size=args.block_size,
     attention_pattern=("local", "global"),
     sliding_window_size=args.block_size,
     use_qk_norm=True,
   )
-  model = dg.DiffusionGemma(cfg)
-  process = dg.DiffusionProcess(dg.DiffusionConfig(cfg.vocab_size))
+  model = d3pm.DiffusionGemma(cfg)
+  process = d3pm.DiffusionProcess(d3pm.DiffusionConfig(cfg.vocab_size))
   opt = optim.AdamW(learning_rate=_lr_schedule(args))
 
   batches = make_batches(ids, args.seq_len, args.batch_size, mx.random.key(0))
-  # Fixed held-out batch + key: a deterministic, comparable eval loss.
   val_x0 = next(
     make_batches(ids, args.seq_len, args.batch_size, mx.random.key(7))
   )
   eval_key = mx.random.key(1234)
 
-  sampler = dg.BlockSampler(
+  sampler = d3pm.BlockSampler(
     model,
     process,
-    dg.SamplerConfig(
+    d3pm.SamplerConfig(
       num_steps=64,
       seq_len=args.block_size,
       entropy_bound=args.entropy_bound,
     ),
-    early_stop=dg.ChainedEarlyStop(
-      [dg.TokenStabilityEarlyStop(), dg.EntropyEarlyStop()]
+    early_stop=d3pm.ChainedEarlyStop(
+      [d3pm.TokenStabilityEarlyStop(), d3pm.EntropyEarlyStop()]
     ),
   )
   prompt_ids = tokenizer.encode(args.prompt) if args.prompt else None
@@ -128,24 +138,30 @@ def sample_text(sample_key: mx.array, max_blocks: int) -> str:
     key, k = mx.random.split(key)
 
     def loss_fn(m: mx.array) -> mx.array:
-      return dg.block_diffusion_loss(
+      return d3pm.block_diffusion_loss(
         m,
         process,
         x0,  # noqa: B023
         k,  # noqa: B023
         block_size=args.block_size,
+        prompt_len=args.prompt_len,
       )
 
     loss, grads = nn.value_and_grad(model, loss_fn)(model)
     opt.update(model, grads)
     mx.eval(model.parameters(), opt.state)
     if step % 100 == 0:
-      val = dg.block_diffusion_loss(
-        model, process, val_x0, eval_key, block_size=args.block_size
+      val = d3pm.block_diffusion_loss(
+        model,
+        process,
+        val_x0,
+        eval_key,
+        block_size=args.block_size,
+        prompt_len=args.prompt_len,
       )
       mx.eval(val)
-      print(f"step {step:5d}  train {float(loss):.4f}  eval {float(val):.4f}")
       key, k = mx.random.split(key)
+      print(f"step {step:5d}  train {float(loss):.4f}  eval {float(val):.4f}")
       print(f"  sample: {sample_text(k, 1)!r}")
 
 

examples/diffusiongemma/main.py

@@ -92,7 +92,7 @@ extend-select = [
     "TCH",  # Type-checking blocks
     "UP",   # Pyupgrade (modern syntax)
 ]
-ignore = ["S101", "ANN101"]
+ignore = ["S101"]
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"