Local inference for diffusion language models, written in Mojo.

Introduction

mojo-dllm runs LLaDA-8B and Dream-7B, two diffusion language models, from quantized GGUF files on your own CPU or on an NVIDIA GPU. The whole inference path is Mojo: the GGUF reader, the tokenizer, the quantized matrix multiply, the GPU kernels, the transformer and the denoising loop. You need no Python at runtime, and the GPU is optional.

A diffusion language model writes differently from the models you know. An autoregressive model appends one token at a time. LLaDA starts from a canvas of [MASK] tokens and fills it in over a fixed number of steps, committing the tokens it is most sure about first. Diffusion decoding in five minutes walks through it with a picture.

What works today

• mojo-dllm run generates text from LLaDA-8B-Instruct and Dream-v0-Instruct-7B (Q4_K_M) on x86 CPUs with AVX2, using AVX-VNNI when present.
• --device gpu runs both models on NVIDIA GPUs with compute capability 8.0 or newer. It has been measured on an RTX 3050 Laptop GPU only.
• Both tokenizers match llama.cpp id for id on a 36-case multilingual corpus.
• Logits match llama.cpp's implementation of both models; the measured agreement is in Correctness.
• Benchmarks against llama.cpp and diffuse-cpp on the same machine, with every number generated from a result file in the repository.

What does not work yet

• The GPU path is slower than llama.cpp's CUDA build; Benchmarks has the numbers. It supports NVIDIA only.
• Two architectures: LLaDA and Dream. Others stop with unsupported model architecture: <name> (supported: llada, dream).
• One quantization mix: the F32, Q4_K and Q6_K tensors a Q4_K_M file contains. Any other tensor type stops the load with an error that names it.

Where to go next

• Run it: Quick start.
• Read the numbers: Benchmarks.
• Read the code: Architecture, then the source under src/mojo_dllm/.

Quick start

You need Linux on x86-64 with AVX2, glibc 2.34 or newer, about 6 GB of free memory, and uv. uv installs the pinned Mojo toolchain (Mojo 1.1.0, max-core 26.6.0) into the project's .venv; nothing goes system-wide.

1. Set up

curl -LsSf https://astral.sh/uv/install.sh | sh    # skip if you have uv
git clone https://github.com/vladimirrott/mojo-dllm
cd mojo-dllm
./mojo-dllm setup

The uv installer puts uv in ~/.local/bin. If your shell does not find it afterwards, open a new terminal or run source ~/.local/bin/env.

setup runs four steps and prints each one:

1. Toolchain. It checks the platform, glibc, AVX2 and free memory, stops on anything that cannot work, and runs uv sync --frozen.
2. Build. It compiles build/mojo-dllm, or reports that the binary is newer than the sources.
3. Model. It asks before downloading LLaDA-8B-Instruct Q4_K_M (4.9 GB) into the Hugging Face cache, and checks the file's sha256 against the one the hub publishes. --yes answers the question for you; --no-model skips the step; --model picks another model.
4. Check. It generates 16 tokens to prove the binary, the model and the device work together.

Without a terminal to ask on (a script, CI), setup does not download anything unless you pass --yes. It then exits 3 and prints the command that finishes the job.

./mojo-dllm doctor runs the same checks without changing anything, and adds the GPU, the binary's age and the model cache. Every line that is not ok names the command that fixes it.

2. Generate

./mojo-dllm run "Explain speculative decoding in two sentences."
./mojo-dllm run "Write a haiku about rain." --max-tokens 64 --steps 32 --verbose

The prompt is the first word that is not a flag; quote it. run rebuilds the binary if a source file changed, finds the model, picks the device and hands the rest of the flags to build/mojo-dllm run. Command line lists them all.

--max-tokens sets the canvas size, --steps the number of denoising passes. Fewer steps run faster and commit more tokens per step, which costs quality. Each step runs the full model over prompt and canvas, so expect seconds per step on a laptop CPU. --verbose prints load time, per-step time and peak memory.

The prompt goes through the model's chat template. Pass --no-chat to feed raw text instead.

The device

--device auto is the default. It picks the first NVIDIA GPU with compute capability 8.0 or newer that nvidia-smi lists, and the CPU when there is none. Status lines go to stderr, so stdout holds only the model's text:

device: gpu (NVIDIA GeForce RTX 3050 6GB Laptop GPU, compute 8.6, 585 MiB free)

auto does not compare the card's free memory with the model's size. In that line another program held most of the card, so the run that followed would have said how much memory it needs and exited. Pass --device cpu to run anyway.

3. Other models

-m (or --model) takes one of three forms:

Without :QUANT, a repo resolves to its Q4_K_M file, then Q8_0, then the first .gguf it lists, the order llama.cpp uses. A quant that matches nothing lists the files the repo has. Split (sharded) GGUF files are not supported.

mojo-dllm runs LLaDA and Dream GGUF files made by llama.cpp's converter. Another architecture downloads and then stops with an error naming the architecture.

./mojo-dllm pull MODEL downloads without generating. Models live in the Hugging Face cache (~/.cache/huggingface/hub, or $HF_HUB_CACHE), shared with every other tool that uses it. A gated or private repo needs HF_TOKEN.

The sha256 check runs on download. A file already in the cache is checked against the hub's sha256 the first time it is used, and the launcher remembers the result in ~/.cache/mojo-dllm/verified.json (under $XDG_CACHE_HOME when that is set). A file that fails the check is deleted, and the next run downloads it again. With --offline the check cannot run, so an unchecked file is used after a warning.

4. Offline

./mojo-dllm pull                      # on a connected machine
./mojo-dllm run "Hello" --offline     # later, with no network

--offline (or HF_HUB_OFFLINE=1) uses the cache only and fails with the pull command to run when the model is missing.

5. Look inside a model

./mojo-dllm pull                             # prints the model file's path
./mojo-dllm inspect PATH.gguf | head -40

You get the metadata, a count of tensor types, and every tensor with its shape. The binary's own commands (inspect, tokenize, logits) go to build/mojo-dllm unchanged, after the same rebuild check. Any other word stops with the closest command name.

Without the launcher

The launcher is a convenience. The binary works on its own:

uv sync --frozen
scripts/build.sh
build/mojo-dllm run --model PATH.gguf --prompt "Hello" --device cpu

build/mojo-dllm defaults to --device cpu and takes only file paths.

Command line

mojo-dllm run      --model PATH --prompt TEXT [options]
mojo-dllm tokenize --model PATH --prompt TEXT [--no-chat]
mojo-dllm inspect  PATH
mojo-dllm logits   --model PATH --tokens IDS --rows POSITIONS --out FILE [--threads N] [--device D]

mojo-dllm COMMAND --help (or -h) prints that command's usage and exits 0. Help wins over every other problem in the arguments, wherever it appears, except as the value of a flag (--prompt --help sets the prompt to --help). mojo-dllm, mojo-dllm help and mojo-dllm --help print the usage of every command. Each command declares its flags: a flag the command does not declare stops the run (see Errors).

run

The printed answer ends at the first end-of-text or end-of-turn token in the canvas. Generation stops as soon as that answer is final: once the first position after the prompt that is still masked or holds such a token holds such a token, no later step can change the printed text, because steps only fill masked positions. The result is the same text with fewer steps, and --json reports "stopped_early": true. --no-early-stop runs every step.

tokenize

Prints the token ids of the chat-formatted prompt, the same ids run feeds the model. Useful to compare against llama-tokenize.

inspect

Prints GGUF metadata, a count per tensor type, and every tensor's type and shape. It needs no supported architecture, so it works on any GGUF file.

logits

Runs one forward pass over the given token ids and writes f32 logits for the requested positions to a file, [positions, vocab], row major. The parity tests use it; see Correctness.

Errors

An error names the thing that is wrong. Usage errors exit with status 2 and change nothing; every other error exits with status 1.

The remaining errors exit with status 1 and come from the model file itself:

ERROR: unsupported GGUF tensor type: IQ2_XXS (tensor blk.0.attn_q.weight)
ERROR: tensor blk.0.attn_q.weight is Q8_0: matrix kernels exist for Q4_K and Q6_K
ERROR: unsupported model architecture: qwen2 (supported: llada, dream)
ERROR: required metadata key missing: tokenizer.ggml.model
ERROR: tensor shape mismatch for blk.0.attn_q.weight: [4096, 2048], expected [4096, 4096]
ERROR: no GPU found: the GPU backend needs an NVIDIA GPU with compute capability 8.0 or newer (run with --device cpu)

With --device gpu, setting MOJO_DLLM_GPU_PROFILE=1 makes --verbose split a step into GEMM, attention and LM head time. It synchronizes the GPU after each kernel group to measure them, so leave it off when you benchmark.

Benchmarks

LLaDA-8B-Instruct, Q4_K_M. 13th Gen Intel(R) Core(TM) i5-13420H, 12 threads, 33.3 GB RAM, performance power profile. Each prompt: 128 generated tokens, 32 denoising steps; median of 3 runs x 3 prompts.

Versions: mojo-dllm bdd268b, Mojo 1.1.0 (8189361e), llama.cpp 836d571, diffuse-cpp 1d2bd6a. Evidence: bench/results/2026-10-03-cpu-diffuse.json, bench/results/2026-10-03-cpu-llada.json.

† Not the same work: the inter-step cache reuses stale K/V and entropy_exit can stop early. Shown because it is diffuse-cpp's recommended mode.

LLaDA-8B-Instruct on the GPU, Q4_K_M. NVIDIA GeForce RTX 3050 6GB Laptop GPU (6144 MiB, driver 595.84), host 13th Gen Intel(R) Core(TM) i5-13420H. Each prompt: 128 generated tokens, 32 denoising steps; median of 3 runs x 3 prompts. Peak RSS is host memory.

Versions: mojo-dllm 78014b4, Mojo 1.1.0 (8189361e), llama.cpp 836d571. Evidence: bench/results/2026-10-12-gpu-llada.json.

Dream-v0-Instruct-7B, Q4_K_M. 13th Gen Intel(R) Core(TM) i5-13420H, 12 threads, 33.3 GB RAM, performance power profile. Each prompt: 128 generated tokens, 32 denoising steps; median of 3 runs x 3 prompts.

Versions: mojo-dllm bdd268b, Mojo 1.1.0 (8189361e), llama.cpp 836d571. Evidence: bench/results/2026-10-03-cpu-dream.json.

Dream-v0-Instruct-7B on the GPU, Q4_K_M. NVIDIA GeForce RTX 3050 6GB Laptop GPU (6144 MiB, driver 595.84), host 13th Gen Intel(R) Core(TM) i5-13420H. Each prompt: 128 generated tokens, 32 denoising steps; median of 3 runs x 3 prompts. Peak RSS is host memory.

Versions: mojo-dllm 036ef79, Mojo 1.1.0 (8189361e), llama.cpp 836d571. Evidence: bench/results/2026-10-12-gpu-dream.json.

LLaDA-8B-Instruct, llama.cpp thread sweep and forward pass alone: 13th Gen Intel(R) Core(TM) i5-13420H; median of 2 runs, one prompt, canvas 148 tokens.

The ms / step column runs llama-diffusion-cli. The forward pass alone is mojo-dllm's measured forward pass against llama.cpp's llama-bench at the same token count. llama-bench applies a causal mask and computes logits for one position; mojo-dllm's figure includes logits for up to 32 positions and copying the canvas in. CPU busy before a run: at most 6.5%. Evidence: bench/results/2026-10-09-fair-cpu-llada.json.

LLaDA-8B-Instruct on the GPU, forward pass alone: NVIDIA GeForce RTX 3050 6GB Laptop GPU; median of 3 runs, one prompt, canvas 148 tokens.

The forward pass alone is mojo-dllm's measured forward pass against llama.cpp's llama-bench at the same token count. llama-bench applies a causal mask and computes logits for one position; mojo-dllm's figure includes logits for up to 32 positions and copying the canvas in. CPU busy before a run: at most 8.9%. Evidence: bench/results/2026-10-12-fair-gpu-llada.json.

What is compared

For LLaDA, all three runtimes load LLaDA-8B-Instruct Q4_K_M and denoise a canvas of L = prompt + 128 tokens in 32 steps. For Dream, mojo-dllm and llama.cpp (diffuse-cpp's Dream support needs its own converted file, not benchmarked here) do the same with Dream-v0-Instruct-7B Q4_K_M; both run Dream's timestep schedule over the whole canvas. A step is one full forward pass over the canvas in each of them, so the step counts and shapes match.

diffuse-cpp cannot read llama.cpp's GGUF files, so it loads its own conversion of the same model (diffuse-cpp/LLaDA-8B-Instruct-GGUF, llada-8b-q4km.gguf, 5.5 GB): F16 embeddings and a slightly different Q4_K / Q6_K mix. It has no tokenizer and receives the token ids mojo-dllm produces. It prints no timings either, so the harness timestamps the two lines it writes around its generate call.

llama.cpp runs its timestep schedule because its block mode needs the prompt length to be a multiple of the block size. The schedules commit tokens in a different order, which changes the text but not the cost of a step.

† diffuse-cpp's recommended mode reuses K/V between steps and stops early when the canvas has settled. It does less work, so it is not a like-for-like row.

On the GPU

bench/config-gpu.json and bench/config-dream-gpu.json compare mojo-dllm --device gpu with the same llama-diffusion-cli built with CUDA and run with -ngl 99; the harness refuses a CUDA run unless llama.cpp reports every layer on the GPU. Canvas, steps and prompts match the CPU runs. These configs leave the pytest watchdog off ("kill_pytest": false), so each result records how busy the CPU was before every run. Peak RSS is host memory; neither runtime reports GPU memory. The GPU tables above come from these configs.

The test GPU is a laptop part held to 30 W, whose clocks fall under sustained load (bench/checks/2026-10-12-gpu-power-under-load.txt) and vary from day to day: the same llama.cpp build ran LLaDA steps much slower on 2026-10-12 than on 2026-10-09 (compare bench/results/archive/2026-10-09-gpu-llada.json with the current file). Both runtimes in one result file run interleaved, prompt by prompt, so compare them within a file, not across dates.

bench/config-fair-cpu.json and bench/config-fair-gpu.json address two questions about fairness: whether llama.cpp does better below 12 threads on a hybrid CPU (a sweep over -t 4, 6, 8, 12), and how the forward passes compare without either runtime's sampler (mojo-dllm's measured forward pass against llama.cpp's llama-bench at the same token count).

Method

bench/run_bench.py does the measuring:

• three prompts, three repetitions, runtimes interleaved within each repetition;
• every run starts only when the CPU has been under 10% busy for two seconds;
• the power profile is performance for the duration and restored after;
• a watchdog kills any pytest process (other work on the machine started test suites mid-run), and the kills are logged in the result file;
• generation time excludes model loading, which is reported as startup;
• peak RSS comes from /usr/bin/time.

The table above is generated from bench/results/*.json by scripts/bench_table.py, and CI fails if the two disagree.

Steps and quality

A diffusion model commits tokens / steps tokens per step. The table above uses 4 tokens per step, which is fast and shows in the text as repeated or dropped words. More steps cost proportionally more time:

LLaDA-8B-Instruct.Q4_K_M.gguf, 128 generated tokens.

Evidence: bench/quality/2026-10-03-LLaDA-8B-Instruct.Q4_K_M.json (mojo-dllm bdd268b).

"Repeated words" counts immediate repetitions such as "the the", a rough proxy for parallel-decoding damage. bench/quality.py produces this table.

Accuracy of --threshold and --dual-cache

Both options change what a step computes, so their answers can differ from the plain sampler's. A paired GSM8K run checks how often they still get the answer right (bench/eval_gsm8k.py, see Development):

LLaDA-8B-Instruct.Q4_K_M.gguf, GSM8K, 200 problems, flags --max-tokens 256 --block-length 32 --steps 256 --device gpu --json, mojo-dllm 3b2dac4.

Discordant counts problems only the baseline solved / only the arm solved. Time per problem is the mean over the first 50 problems of each arm, a round that ran with nothing else on the machine.

Evidence: bench/eval/2026-10-29-baseline-vs-dualcache.json, bench/eval/2026-10-29-baseline-vs-threshold09.json, bench/eval/2026-10-29-baseline-vs-dualcache-first50.json, bench/eval/2026-10-29-baseline-vs-threshold09-first50.json.

Reading the numbers

• ms / step is the cost of one denoising pass over the canvas.
• tokens / s is 128 divided by generation time. A diffusion model's tokens/s depends on the step count you choose, so compare ms / step first.
• The ablation row shows what computing logits only for candidate positions saves.
• On the GPU, read ms / step and peak RSS together. llama.cpp is faster per step; mojo-dllm holds far less in host memory, because it uploads the weights once and releases the file's pages behind it.
• The thread sweep and the forward-pass-alone tables test the CPU comparison. If llama.cpp ran faster below 12 threads, or mojo-dllm's lead came only from skipping logits, those tables would show it. The sweep stops at 12 because the test CPU runs 12 threads.
• A note under a table counts the runs that started on a busy CPU and gives the medians over the quiet runs alone.

Diffusion decoding in five minutes

An autoregressive model such as Llama writes left to right. Each forward pass produces one token, and an append-only KV cache makes the next pass cheap.

A masked diffusion model such as LLaDA writes the way you fill in a crossword. It starts from a canvas of [MASK] tokens after the prompt and, at every step, looks at the whole canvas at once, proposes a token for every blank, and commits the few it is most confident about. The rest stay masked for the next step.

prompt: "Explain speculative decoding"           canvas (8 positions shown)

step 0   [M] [M]   [M]     [M]  [M]    [M] [M]     [M]
step 1   [M] [M]   [M]     [M]  [M]    [M] [M]     .
step 2   A   [M]   [M]     [M]  [M]    [M] [M]     .
step 3   A   draft [M]     [M]  [M]    [M] tokens  .
  ...
step 8   A   draft model   guesses     several     tokens .

Each step runs the full transformer over prompt plus canvas. Attention is bidirectional: position 3 can read position 7, which a causal model forbids.

The loop LLaDA uses

mojo-dllm implements the reference sampler from the LLaDA repository (generate.py), with its defaults:

1. Split the generated region into blocks of --block-length positions and fill them left to right (semi-autoregressive decoding).
2. Give each block steps / blocks steps and spread its masked positions over them evenly: with 32 masks and 8 steps, 4 tokens per step.
3. At each step, for every masked position in the block, take the argmax token and its softmax probability.
4. Commit the k positions with the highest probability ("low-confidence remasking": the uncertain ones stay masked).

At temperature 0 the loop is deterministic. Above 0, LLaDA adds Gumbel noise to the logits before the argmax; mojo-dllm does the same.

Why the runtime looks different

• No KV cache. Any position can change between steps, so every step recomputes keys and values for the whole canvas. A step costs a full prompt pass, about 2 × 7.5 billion operations per token on the canvas.
• GEMMs, not GEMVs. A step multiplies weight matrices by a hundred or more activation rows at once. The CPU kernel is built for that shape; see The quantized GEMM.
• Few rows need logits. Only the masked positions of the current block are candidates, so mojo-dllm computes the 126 464-way output projection for those rows alone. The benchmark ablation measures what that saves.

Architecture

                 mojo-dllm run
                      |
        +-------------+--------------+
        |             |              |
   GGUF reader    tokenizer     diffusion loop
   (mmap)         (BPE)         (diffusion/sampler.mojo)
        |                            |
        |                       LLaDA forward
        |                       (models/llada.mojo)
        |                            |
        +----> repacked weights ---> quantized GEMM, RMSNorm, RoPE,
               (kernels/packed.mojo)  attention, SwiGLU (kernels/)

Source map

A forward pass

For a canvas of L tokens and a list of positions whose logits the caller needs:

1. Embedding rows are decoded from the memory-mapped Q4_K table, one row per token.
2. For each of 32 layers: RMSNorm, then the activations are quantized to int8 once and reused by the Q, K and V projections. RoPE rotates Q and K; attention runs over all L positions with no mask; the output projection and residual follow. The FFN quantizes once for gate and up, applies SiLU, quantizes again for down.
3. In the last layer, only the requested positions continue past attention.
4. Final RMSNorm and the output projection run on the requested rows only.

The forward pass allocates nothing: every activation buffer is sized once for the longest canvas the run will use.

Memory

The GGUF file is memory-mapped. At load, every projection matrix is repacked into a heap block (see The quantized GEMM) and its mapped pages are released, so the model is resident once. Q4_K stays 4-bit after repacking. Q6_K expands to one byte per weight, which costs about 0.3 GB on LLaDA-8B and buys a simpler inner loop. The embedding table and norm vectors stay mapped and fault in as needed.

Threads

sys/parallel.mojo hands out work items from a shared atomic counter. On a hybrid CPU (performance and efficiency cores) a static split makes every parallel region wait for the slowest core; claiming items one at a time lets the fast cores take more. On the development laptop the static split left every GEMM waiting on the four efficiency cores.

The quantized GEMM

Most of a denoising step is matrix multiplication: seven projections per layer, 32 layers, every step, over the whole canvas. run --verbose prints the split for your machine. This page explains the kernel that does it.

The shape of the problem

A step multiplies each weight matrix W (N rows × K columns, 4 to 6 bits per weight) by an activation matrix X (L rows × K). L is the canvas length, a few hundred at most. That is a skinny GEMM, and it is compute-bound: each weight is read once per step and used L times.

Activations: int8 with a scale per 256

Before a projection, each activation row is quantized to int8 with one f32 scale per 256 values, and the sum of each group of 16 values is kept as a float. This is llama.cpp's Q8_K scheme, which keeps the comparison with it fair. Integer dot products are 4× denser than f32 FMAs on AVX2.

Weights: 8 rows per register

GGUF stores each row contiguously. A row-at-a-time dot product needs a horizontal sum per output, and the block scales arrive one scalar at a time. So the loader repacks every matrix once:

for each group of 8 rows, for each 256-weight super-block:
    Q4_K:  d[8], dmin[8] as f32 | scales[8][8], mins[8][8] as u8 | 1024 bytes of nibbles
    Q6_K:  d[8] as f32          | scales[16][8] as i8           | 2048 bytes, one value per byte

Inside the weight bytes, a 32-byte load holds 4 consecutive weights of each of the 8 rows. One vpdpbusd (AVX-VNNI) multiplies those 32 unsigned weights by 4 activation bytes broadcast to every lane and adds into 8 int32 accumulators: 8 rows × 4 multiply-adds in one instruction. The scale of a sub-block is an 8-lane vector, one entry per row, applied once per 16 or 32 weights.

The zero points do not touch the inner loop. Q4_K computes w = d·sc·q − dmin·m, so the second term becomes dmin·m times the activation sum of the sub-block, which the quantizer already stored. Q6_K stores q + 32, and the 32 comes off the same way.

Tiling and threads

A work item is 16 output rows. It walks the tokens 4 at a time, keeping 4 int32 and 4 f32 accumulators in registers, so each loaded weight vector feeds 4 tokens. Items go to threads through an atomic counter (see Architecture).

Checks

tests/test_quant.mojo compares the decoders with the gguf Python package's dequantization, and both GEMMs with an f64 reference over the same dequantized weights. Swapping the nibble halves or dropping the Q6_K zero point fails these tests; that mutation check runs before changes to this file land.

bench/bench_qgemm.mojo measures the kernel alone on LLaDA's layer shapes.

The GPU path

src/mojo_dllm/gpu/kernels.mojo does the same arithmetic on an NVIDIA GPU, with a layout chosen for tensor cores instead of AVX registers.

Weights stay in GGUF blocks. Upload copies each Q4_K matrix as it is and pads each 210-byte Q6_K block to 224 bytes, so every 16-byte chunk of a block is aligned and the kernels can load 16 bytes at a time.

Activations get one scale per 32. The GPU quantizes activations to int8 with one f32 scale per 32 values (llama.cpp's Q8_1 shape) and keeps scale x sum(q) next to it. One warp quantizes one 32-value block, a value per lane, and reduces its max and integer sum with shuffles, so loads and stores are coalesced. A 32-weight Q4_K sub-block then meets exactly one activation scale:

sum w x = d*sc*dx * sum(q * qx)  -  dmin*m * (dx * sum(qx))

One mma per scale group. A Q6_K thread block computes 64 weight rows by a tile of 32, 64 or 80 tokens. For each 256-wide super-block it copies the weights into shared memory (see the Q6_K paragraph below), copies the matching activations next to them, and each of its 4 warps runs one mma.m16n8k16.s8 per 16-weight scale group for its 16 rows and every 8 tokens. (A Q4_K block scales every 32 weights, so it runs mma.m16n8k32.s8 per sub-block instead.) The int32 result goes to f32 with the two scales and, for Q4_K, the min term.

Q4_K: packed in shared memory, unpacked in registers. Removing the staging work cut the GEMM roughly in half (bench/checks/2026-10-13-gemm-cost-split.json: with staging done once, the GEMM took about 100 ms instead of about 220), while removing the mma saved only a few milliseconds. So a Q4_K block copies each row's 144-byte super-block as it is, in 16-byte loads: the 16-byte header (d, dmin and the packed scales) to one array, the 128 bytes of nibbles to another. A warp unpacks its A fragments in registers: word i of nibble pair p holds sub-block 2p in its low nibbles and 2p + 1 in its high ones at the same position, so one shared load feeds two mmas, and the 6-bit scales are decoded from the header per sub-block. With the weights packed, a block holds 128 rows and runs 8 warps in the same 48 KiB, which halves the activation copies per output. The arithmetic is unchanged, so the logits are too (bench/checks/2026-10-13-gemm3-logits.json).

Prefetching the next super-block. On the kernel before the packed layout, removing staging saved about 120 ms of the GEMM and removing the mma 4 to 7 ms, both measured against the integer-accumulation variant (bench/checks/2026-10-13-gemm-cost-split.json). The packed layout (128-row blocks) took that harness from 258.7 to 215.1 ms without a breakdown of the staging share. The loop still copied a super-block, waited at a barrier, and only then computed. Now each thread reads super-block b + 1 from global memory into registers (weights, activation bytes and activation scales) before the mma loop of b, and stores it to shared memory at the top of the next iteration, after the barrier that ends b. The Q4_K GEMM went from 148.8 to 87.2 ms per forward (median launch time times launches, middle value of three interleaved traces, bench/profiles/2026-10-13-gemm4-*), and the logits match main on all 84 rows (bench/checks/2026-10-13-gemm4-logits.json). That the gain comes from overlapping the loads with compute is the reading the numbers fit; no hardware counters were collected to confirm it.

Q6_K: the same layout and prefetch. The Q6_K kernel took the route the Q4_K one did. A block of 64 rows copies each row's 224 bytes (low nibbles, high bits, 16 int8 group scales, d) into shared memory as they are, in 16-byte loads; a warp rebuilds each 16-weight group's A fragment in registers (_q6k_frag: nibble, two high bits, minus 32 by an xor and a sign extension), and the next super-block waits in registers while the mma loop runs. Groups accumulate in the same order as before, so the result is unchanged. The padded block is 12 bytes larger than the 212-byte padding it replaces: with 6,217,728 Q6_K blocks in LLaDA-8B Q4_K_M (output, 16 attn_v and 16 ffn_down tensors) that is 74.6 MB (71.2 MiB) more device memory.

Tile width per call. Unpacking a weight tile costs the same however many tokens then use it, so the widest tile that saves a tile wins: a 148-token canvas runs two 80-wide tiles rather than five 32-wide ones, and each weight tile is unpacked twice instead of five times. The LM head's 32 rows keep the 32-wide tile. 80 is the widest tile whose shared memory fits the 48 KiB a kernel gets without asking for more. Token tiles of one weight tile run back to back, so they read it through L2.

Shared-memory layout. Weight rows are 272 bytes apart, so the 8 rows a fragment load reads fall in different banks. Activation scales are stored by sub-block, then token, then (scale, sum), so the pair of tokens a thread folds after an mma is one 128-bit load, and the scale rows are padded to 9 and 17 floats for the same reason. tests/gpu/test_gpu.mojo checks the fragment layout against a CPU matmul on random int8 matrices, and the whole kernel against an f64 GEMM.

Row statistics on the device. A greedy step needs, per masked row, the argmax and a confidence. Copying the logits to the host (32 rows of 126,464 floats for a 32-token block) and looping over them there is the largest piece of a step outside the forward pass. One block per row now reads its logits twice: the first pass finds the first index of the maximum (a tie goes to the smaller index, as the host loop picks it), the second accumulates, in f64, the sum of exp(x - max), the sum of exp(x - max)(x - max) and the largest logit at any other index. The sampler finishes from those four numbers: p = 1 / z for LLaDA and maskgit_plus, p minus the runner-up's probability for topk_margin, and sum e (x - max) / z - log z for Dream's entropy. The exponentials are f32, so a confidence differs from the host's f64 loop in the sixth digit; tests bound it at 1e-5 relative.

Fused elementwise kernels. The quantized GEMMs want int8 activations, so each norm and the FFN activation end in a quantize. Run separately, a layer wrote and re-read the normalized rows and the product of the gate and up projections through global memory. Three kernels remove that traffic: add_rmsnorm_quantize adds the projection output into the residual stream, normalizes the sum and quantizes it; rmsnorm_quantize does the same for layer 0, whose input has no add; silu_mul_quantize quantizes silu(g) * u without storing it. The norm kernels give thread t the values t, t + 256, ... so a warp holds 32 consecutive values at each step, which is one quantization block, and the quantize step is the same _quantize_lane code the plain kernel runs. The float expressions and their order match the unfused kernels', so the fused output is identical to the chain's, which the tests assert exactly. The attention output still goes through the plain quantize kernel.

Attention on f16 tensor cores. Each warp owns 16 queries and walks the keys 32 at a time. Q is scaled by 1/sqrt(d) and held in registers as f16 fragments; K is staged row-major and V transposed, both as f16, so every operand load is one 32-bit word. mma.m16n8k16 computes S = Q K^T in f32, the four threads that share a row reduce its max with two shuffles, and the online softmax rescales the running sums. S's accumulator for 16 keys has the same register layout as the A operand of the next mma, so P goes from S to P V without leaving the registers. On LLaDA-8B the logits stay as close to the f32 CPU path as with f32 attention (bench/checks/2026-10-13-attention-f16-logits.json). The kernel never stores a score row, so the canvas length has no fixed limit.

Correctness

mojo-dllm is checked at four levels, each against an implementation it does not share code with.

1. Formats and kernels (unit tests, every commit)

• GGUF parsing, Q4_K / Q5_K / Q6_K / Q8_0 / F16 / BF16 decoding: compared with the gguf Python package's dequantization on fixtures it generated (max abs difference 1e-6).
• Both GEMMs: compared with an f64 reference over the same dequantized weights; relative error under 1% with int8 activations.
• RMSNorm, RoPE, softmax, bidirectional attention, SwiGLU: closed-form cases.
• The forward pass of a 2-layer model with random K-quant weights: compared with a numpy forward pass; relative logit error under 2% on the LLaDA model and under 5% on the Dream model (its large Q/K biases sharpen the softmax, which amplifies int8 error), and the same argmax on every position.

The GPU has its own suite, tests/gpu/test_gpu.mojo, which runs on a machine with an NVIDIA GPU (scripts/run-gpu-tests.sh; CI only builds it). It checks each GPU kernel against the CPU function it replaces, the int8 mma fragment layout against a CPU matmul, and the GPU forward pass of both tiny models against the numpy reference with a 2% bound on both models.

2. Tokenizer (unit tests, every commit)

36 strings covering English, contractions, numbers, accents, combining marks, CJK, Cyrillic, Arabic, Devanagari, emoji, CRLF, non-breaking and em spaces, code, and LLaDA's chat-template special tokens. Every id matches llama-tokenize on the real LLaDA GGUF, and decoding restores the input byte for byte.

3. Logits on the real model

tools/ref/llada_ref.cpp links libllama and dumps llama.cpp's logits for a fixed token sequence; mojo-dllm logits dumps ours; scripts/compare_logits.py compares them. The sections marked GPU run mojo-dllm logits --device gpu.

The GPU sections also compare against llama.cpp built with CUDA. That build quantizes activations with one scale per 32 values, as the GPU path does, while llama.cpp on the CPU uses one per 256, as the CPU path does. On Dream's prompt rows the two llama.cpp builds disagree with each other, so a single reference would mix that spread into mojo-dllm's numbers. The table under each GPU section shows all three pairs.

LLaDA-8B-Instruct.Q4_K_M.gguf

Evidence: bench/parity/2026-10-03-llada.json (mojo-dllm 0b3e17a).

LLaDA-8B-Instruct.Q4_K_M.gguf, GPU

The same canvases against a second reference, llama.cpp CUDA, and the two references against each other:

Evidence: bench/parity/2026-10-09-llada-gpu.json (mojo-dllm cedc0b5).

Dream-v0-Instruct-7B-Q4_K_M.gguf

Evidence: bench/parity/2026-10-03-dream.json (mojo-dllm 0b3e17a).

Dream-v0-Instruct-7B-Q4_K_M.gguf, GPU

The same canvases against a second reference, llama.cpp CUDA, and the two references against each other:

Evidence: bench/parity/2026-10-09-dream-gpu.json (mojo-dllm cedc0b5).

Exact equality is not the goal: both runtimes quantize activations to int8, but they block and accumulate in different orders. The llama.cpp maintainers recommend the same kind of statistical comparison for diffusion models.

The table scores every position of each canvas. On Dream the median agrees as closely as on LLaDA, but a few prompt positions do not: they sit on special tokens and newlines, where Qwen-family models carry very large outlier activations that magnify int8 rounding. Masked positions, the ones the sampler reads, agree closely. Which of the two runtimes is nearer an f32 reference at those outlier rows is not measured yet.

4. Generated tokens on the real model

llama.cpp's own diffusion example cannot serve as a token-level reference. On the pinned commit it never applies --diffusion-cfg-scale or --diffusion-alg-temp, adds Gumbel noise to one row only, ranks entropy in the wrong direction, and samples from the distribution even at temperature 0. So llada_ref generate re-implements LLaDA's generate.py (temperature 0, low-confidence remasking) on top of llama.cpp's logits, and the result is compared token by token with mojo-dllm run. For Dream, llada_ref generate-dream does the same with Dream's diffusion_generate.

A diffusion loop amplifies small differences: when two candidates are nearly tied, the runtimes can commit different tokens at one step, and every later step then sees a different canvas. A run that matches only up to some position therefore says where the first near-tie fell, not that the algorithm differs. A run that matches every token, as the prime-checking prompt does on LLaDA, shows that the loops agree.

LLaDA-8B-Instruct.Q4_K_M.gguf

Evidence: bench/parity/2026-10-03-llada.json (mojo-dllm 0b3e17a).

LLaDA-8B-Instruct.Q4_K_M.gguf, GPU

Evidence: bench/parity/2026-10-09-llada-gpu.json (mojo-dllm cedc0b5).

Dream-v0-Instruct-7B-Q4_K_M.gguf

Evidence: bench/parity/2026-10-03-dream.json (mojo-dllm 0b3e17a).

Dream-v0-Instruct-7B-Q4_K_M.gguf, GPU

Evidence: bench/parity/2026-10-09-dream-gpu.json (mojo-dllm cedc0b5).

Reproduce

c++ -O2 -std=c++17 tools/ref/llada_ref.cpp -I$LLAMA/include -I$LLAMA/ggml/include \
    -L$LLAMA/build/bin -lllama -lggml -lggml-base -Wl,-rpath,$LLAMA/build/bin -o build/llada_ref
uv run --only-group fixtures python bench/parity.py --arch llada --model models/LLaDA-8B-Instruct.Q4_K_M.gguf
uv run --only-group fixtures python bench/parity.py --arch dream --model models/Dream-v0-Instruct-7B-Q4_K_M.gguf
uv run --no-project python scripts/bench_table.py --write

Development

Setup

uv sync --frozen                      # the pinned Mojo toolchain, in .venv
scripts/ci-local.sh --install-hooks

scripts/ci-local.sh also needs shellcheck, markdownlint-cli2 (npm install --global markdownlint-cli2), mdBook and an authenticated gh. A missing tool fails its step; it never skips it.

The hooks run scripts/ci-local.sh --fast before each commit (hygiene scan, mojo format, -Werror build) and the full set before each push.

The gates

scripts/ci-local.sh runs the same steps as .github/workflows/ci.yml:

GitHub's runners have no GPU, so CI builds the GPU suite without running it. ci-local.sh runs it when nvidia-smi lists a GPU and otherwise prints a SKIP line, never a PASS; run-gpu-tests.sh on its own exits 1 without a GPU.

The test count is pinned because Mojo's TestSuite exits 0 after running zero tests. If you add tests, run UPDATE_TEST_COUNT=1 scripts/run-tests.sh and commit the new count with them.

Tests

scripts/run-tests.sh compiles each tests/test_*.mojo and runs it. Fixtures live in tests/fixtures/ and come from scripts, never from hand:

tests/gpu/test_gpu.mojo holds the GPU tests: each kernel against the CPU function it replaces, and the GPU forward pass on both tiny models against the f32 reference.

The tiny model is 2 layers of hidden size 256 with random K-quant blocks. Its expected logits come from a float numpy forward pass over the weights that the gguf package dequantizes, so the Mojo forward pass is checked against the reference formats, not against itself.

Benchmarks

uv run --no-project python bench/run_bench.py --config bench/config.json --reps 3
uv run --no-project python scripts/bench_table.py --write

bench/config.json names the model files and the llama.cpp and diffuse-cpp checkouts. The harness waits for an idle CPU before each run, interleaves the runtimes, switches the power profile to performance for the duration, and writes a result file with every raw run. It also kills any process named pytest while it runs, because test suites from other work on the development machine kept starting mid-benchmark. Set "kill_pytest": false in a config to turn that off; the GPU configs (bench/config-gpu.json, bench/config-dream-gpu.json) do, since a GPU run leaves the CPU nearly idle. They compare --device gpu with llama.cpp's diffusion example built with CUDA in build-cuda/:

cmake -B build-cuda -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=86
cmake --build build-cuda --target llama-diffusion-cli

GPU kernels

scripts/gpu_profile.py traces one run --device gpu with Nsight Systems (nsys --trace=cuda, which ran here without root although the driver keeps hardware counters for admins) and writes the device time per forward pass of each kernel group to bench/profiles/. It does not synchronize the stream, unlike MOJO_DLLM_GPU_PROFILE=1. --root points it at another checkout's build for before/after comparisons, and it reports a median-based total next to the plain sum, because a laptop GPU's clocks move under sustained load. --runtime llama.cpp --seq-len N traces llama.cpp's CUDA llama-diffusion-cli on the same prompt instead (it replays its forward pass as a CUDA graph, so the trace expands graph nodes), with the same kernel groups. bench/bench_gpu_gemm.mojo times each LLaDA weight shape alone, adding up the per-forward total from the tensor types in the GGUF.

On the development laptop the GPU is capped at 30 W and its clocks drop under sustained load (bench/checks/2026-10-12-gpu-power-under-load.txt), so compare kernels with interleaved runs (A, B, A, B) in one session rather than against a number from another day.

Quality

bench/eval_gsm8k.py answers a fixed 200-problem GSM8K subset (bench/eval/gsm8k-200.jsonl, see bench/eval/SOURCE.md) and compares two arms problem by problem: accuracy difference with a 95% interval and McNemar's exact test. Any change that is not exact (a sampler that commits more per step, a K/V cache) needs a comparison against the one-token-per-step baseline before it becomes a default.

uv run --no-project python bench/eval_gsm8k.py run --model M.gguf --arm baseline
uv run --no-project python bench/eval_gsm8k.py run --model M.gguf --arm candidate -- --some-flag
uv run --no-project python bench/eval_gsm8k.py compare baseline candidate
uv run --no-project python bench/eval_gsm8k.py compare baseline candidate --limit 50

compare writes bench/eval/<date>-<a>-vs-<b>.json, and scripts/bench_table.py turns those files into the accuracy table in Benchmarks. Time per problem depends on what else the machine runs, so --limit N compares each arm's first N problems in run order and writes a -firstN file; the table takes its times from that file. Run that round with nothing else on the machine.

Toolchain

Mojo is pinned to 1.1.0 and max-core to 26.6.0 in pyproject.toml, locked in uv.lock; both come from PyPI. Bump them together by editing the two pins, running uv lock, and committing that on its own with a green run: Mojo 1.1 removed fn, alias and @parameter closures, and each release has moved APIs this code uses. max-core supplies max.algorithm.parallelize and the GPU API (max.gpu, max.gpu.host). The Mojo wheels need glibc 2.34 or newer.

The Python scripts run through uv too. uv run --no-project python script.py runs a script that uses only the standard library. The fixture generators and bench/parity.py need numpy and gguf, pinned in the fixtures group: uv run --only-group fixtures python scripts/gen_fixtures.py.

Porting a model

mojo-dllm supports LLaDA and Dream. This checklist is the path Dream took, which is the path the next model will take.

1. Read the metadata

Run mojo-dllm inspect model.gguf and note general.architecture, the <arch>.* hyperparameters, tokenizer.ggml.model and tokenizer.ggml.pre, the mask token id, and every tensor type. Dream differs from LLaDA in four places that matter:

Metadata contract

<arch> is llada or dream. Required keys fail with a message naming the missing key. This table describes ModelConfig; the tokenizer also needs its own metadata.

The head dimension is embedding_length / head_count, with exact divisibility required. Other scaling metadata is not interpreted. These compatibility checks do not add general validation of metadata ranges. Whether to constrain canvases by context_length or reject attention.causal = true remains a separate policy decision.

To convert a Hugging Face source with llama.cpp and quantize the result:

python convert_hf_to_gguf.py <hf-dir> --outtype f16 --outfile model-f16.gguf
llama-quantize model-f16.gguf model-Q4_K_M.gguf Q4_K_M

The converter accepts f32, f16, bf16, q8_0, tq1_0, tq2_0, and auto as output types. Conversion success alone does not establish mojo-dllm support; inspect the metadata and tensor types before running the model.

2. Load the weights

Add a config struct and a weight loader beside models/llada.mojo. Reuse PackedMatrix for every projection. A tensor type it does not support stops the load with unsupported GGUF tensor type: <type> (tensor <name>) when no decoder exists. A decodable type without a matrix kernel reports tensor <name> is <type>: matrix kernels exist for Q4_K and Q6_K. Reference decoders accept F32, F16, BF16, Q8_0, Q4_K, Q5_K and Q6_K for embedding rows and vectors; matrix kernels still require Q4_K or Q6_K. Add the decoder and packed kernel with fixtures from scripts/gen_fixtures.py before advertising a new matrix type as supported end to end.

Before loading a large conversion, compare its tensor shapes with the source:

uv run --only-group fixtures python scripts/check_weights.py \
  model-Q4_K_M.gguf <hf-dir>/model.safetensors.index.json
# For an unsharded source, pass model.safetensors instead of the index.

The checker reads safetensors headers only and reverses GGUF dimension order before comparing. It reports mismatches, missing tensors on either side, and unmapped source names; dtypes are shown for reference, not compared. Use --ignore 'rotary.*inv_freq' for source buffers intentionally omitted during conversion; ignored N counts these source entries, and their mapped GGUF counterparts are also omitted. Index shard names must be bare .safetensors filenames resolving inside the index directory. The checker accepts 1–1024 blocks to bound name-map construction. Malformed inputs return 2, including JSON nesting errors. An output tensor absent from a tied source is still reported as missing, with an explanation. Exit 0 means at least one tensor was compared and all comparisons passed, 1 means findings, and 2 means no comparison or invalid input. Shape agreement does not establish matching weights or logits.

3. Prove the forward pass

1. Extend scripts/gen_fixtures.py with a tiny model of the new layout and its numpy reference logits. Write the test first and watch it fail.
2. On the real model, compare logits with llama.cpp through tools/ref/llada_ref.cpp (logits mode) and scripts/compare_logits.py.

4. Tokenizer

If tokenizer.ggml.pre is new, add its pre-tokenizer to tokenizer/bpe.mojo and golden cases to scripts/gen_tokenizer_fixture.py. llama.cpp's regex for each pre-tokenizer is in src/llama-vocab.cpp.

5. Sampler

Add the model's reference sampler next to LLaDA's in diffusion/sampler.mojo, with unit tests for its schedule.

Development plan

Written 2026-10-03 from the draft spec (docs/spec.md) and three research passes: the Mojo 1.1 toolchain, the diffusion-LLM landscape, and the house style of the author's other repositories. Where research contradicted the spec, the correction is recorded here and the code follows the correction.

Corrections to the spec

Model facts used (LLaDA-8B-Instruct, GGUF llada)

32 layers, d = 4096, 32 heads × 128 (MHA), FFN 12288 SwiGLU, RMSNorm ε = 1e-5, RoPE θ = 500 000, vocab 126 464, mask 126 336, BOS 126 080, EOS 126 081, EOT 126 348, non-causal, untied output. Tokenizer: byte-level BPE (gpt2), pre-tokenizer bailingmoe.

Compute design

A denoising step runs the full transformer over the whole canvas (prompt + generation, L ≈ 100 to 160 tokens), so linear layers are GEMMs, not GEMVs. Cost per step ≈ 2 × 7.5·10⁹ × L FLOPs: compute-bound, not bandwidth-bound.

• Activations are quantized per token to int8 with one f32 scale per 256 values (the Q8_K scheme), plus int sums per 16 values.
• Each weight row is decoded once per step into unsigned bytes, a per-16 int16 scale and a per-16 f32 offset. Q4_K and Q6_K share this decoded form, so one micro-kernel serves both.
• The inner product uses maddubs (u8×s8→s16 pairs) and madd with the scale (s16×s16→s32): two AVX2 instructions per 32 MACs, no horizontal sum until the row ends.
• The LM head runs only on the masked positions of the current block (≤ 32 rows instead of L). This is a diffusion-specific saving that a generic runtime does not take.
• Threads split output rows into many small chunks so the 4 P-cores and 4 E-cores of a hybrid CPU balance.

Correctness ladder

1. Unit: GGUF parse (synthetic files written by the tests), Q4_K/Q6_K decode against hand-built blocks, quantized dot vs f32 dot, RMSNorm, RoPE, softmax, attention, scheduler transfer counts against generate.py.
2. Tokenizer: golden token IDs produced by llama-tokenize on a fixed corpus.
3. Logits: a small C++ tool over libllama dumps logits for fixed token sequences; mojo-dllm must match on argmax and cosine similarity.
4. Generation: the same tool runs LLaDA's generate.py greedy algorithm over llama.cpp logits; mojo-dllm must produce the same tokens.

Benchmark protocol

Same machine, same prompts, same thread count, release builds, model load excluded from generation time and reported on its own. Configuration: gen 128, block 32, steps 32 (4 tokens per step, 32 forward passes), plus per-step latency. Three runs each, median reported, raw numbers in bench/results/*.json. The README table is rendered from that JSON and a gate fails when they disagree. No estimated numbers are published.

Repository gates

Sized to what this repo can break, copied from sysknife's reasoning:

Local first: scripts/ci-local.sh runs everything; pre-commit runs the fast set, pre-push the full set. GitHub Actions runs the same script on push and PR (free-plan private repos get 2 000 minutes a month). Pages deploy is gated on the repo being public.

Milestones

Going public needs: M6 done, parity documented, license review of shipped binaries (Modular runtime libraries are not Apache), Pages switched on.

mojo-dllm — Technical Specification

Status: Draft v0.1
Project type: Open-source ML inference systems project
Primary language: Mojo
Initial model: LLaDA-8B, quantized GGUF
Initial target: x86 CPU
Secondary target: NVIDIA consumer GPU
Long-term targets: Apple Silicon, AMD GPU
License target: Apache-2.0 or MIT-compatible, subject to dependency review

1. Summary

mojo-dllm is a lightweight, Mojo-native runtime for local inference of diffusion language models (dLLMs).

The project will begin with LLaDA-8B in quantized GGUF form and implement the complete inference path in Mojo, including:

• GGUF model loading;
• tokenization;
• quantized matrix multiplication;
• transformer execution;
• bidirectional attention;
• diffusion masking and remasking;
• confidence/entropy-based token selection;
• iterative denoising;
• deterministic sampling;
• CLI inference.

The first milestone is deliberately CPU-only and correctness-focused.

The primary long-term technical thesis is:

A single Mojo inference implementation can execute diffusion language models efficiently across heterogeneous consumer hardware, while sharing substantially more kernel/runtime code than traditional C++ + CUDA + Metal + ROCm stacks.

The project is not intended to introduce a new diffusion-model algorithm initially. Its novelty is primarily systems-oriented: portable, quantized, local dLLM inference in Mojo.

2. Motivation

Most current high-performance language-model runtimes are optimized around autoregressive decoding.

Diffusion language models use a fundamentally different inference loop:

1. initialize part or all of the generation region with mask tokens;
2. execute a bidirectional transformer pass;
3. estimate token distributions for unresolved positions;
4. select some positions to commit;
5. optionally remask low-confidence positions;
6. repeat until the generation region is resolved.

This creates runtime requirements that differ from ordinary autoregressive inference:

• repeated whole-sequence or block-level transformer evaluation;
• bidirectional attention rather than a purely causal mask;
• diffusion-specific scheduling state;
• mask/remask operations;
• per-position confidence or entropy calculations;
• opportunities for inter-step state reuse;
• potentially different memory/computation tradeoffs.

Existing implementations primarily use:

• PyTorch;
• C/C++;
• CUDA;
• Triton;
• specialized serving frameworks.

Mojo is a strong candidate for this workload because it aims to support portable low-level CPU and accelerator programming from a common language and type system.

3. Project Goals

3.1 v0.1 goals

The first release MUST:

• run LLaDA-8B Q4_K_M GGUF locally;
• require no Python at inference time;
• execute entirely on x86 CPU;
• parse GGUF directly in Mojo;
• tokenize prompts locally;
• perform iterative diffusion generation;
• expose a simple CLI;
• support deterministic generation with a fixed seed;
• produce numerically reasonable parity with a reference runtime;
• provide reproducible benchmark tooling.

Target invocation:

mojo-dllm run \
  --model ./LLaDA-8B-Q4_K_M.gguf \
  --prompt "Explain speculative decoding." \
  --max-tokens 128 \
  --steps 128 \
  --seed 42

3.2 v0.2 goals

• optimized CPU SIMD;
• thread-parallel transformer execution;
• faster Q4_K matrix multiplication;
• fused normalization;
• improved sampling kernels;
• benchmark parity against diffuse-cpp and llama.cpp.

3.3 v0.3 goals

• NVIDIA GPU backend;
• RTX 30-series support as an initial consumer target;
• device-side quantized GEMM;
• device-side attention;
• device-side diffusion sampling operations;
• optional CPU/GPU split execution.

3.4 v0.4+ goals

• Dream model support;
• Apple Silicon GPU backend;
• AMD GPU backend;
• inter-step cache experiments;
• block diffusion;
• runtime integration layer for projects such as molla;
• API/library mode in addition to CLI mode.

4. Non-Goals

The initial project MUST NOT attempt to implement:

• an OpenAI-compatible HTTP server;
• multi-user continuous batching;
• distributed inference;
• tensor parallelism;
• speculative decoding;
• training;
• fine-tuning;
• arbitrary Hugging Face model loading;
• every GGUF quantization format;
• MAX serving integration;
• custom model download infrastructure;
• a general-purpose LLM runtime;
• novel dLLM research algorithms before baseline correctness.

This project should remain a small inference core, not become another general serving framework.

5. Initial Model

5.1 Model

Initial target:

LLaDA-8B

Preferred representation:

GGUF Q4_K_M

Reasons:

• laptop-friendly model size;
• existing public quantized artifacts;
• simple enough transformer structure;
• multiple independent reference implementations;
• existing C/C++ implementations provide useful correctness oracles;
• bidirectional transformer execution makes it a representative dLLM baseline.

5.2 Second model

The second supported model SHOULD be:

Dream-7B

Adding Dream forces the runtime to generalize beyond a single LLaDA-specific model layout and helps validate the abstraction around:

• GQA;
• model metadata;
• tokenizer configuration;
• diffusion scheduler differences;
• transformer architecture differences.

6. High-Level Architecture

                    +----------------------+
                    |      mojo-dllm       |
                    +----------+-----------+
                               |
             +-----------------+-----------------+
             |                 |                 |
             v                 v                 v
        GGUF Loader        Tokenizer      Diffusion Runtime
             |                                   |
             |                          +--------+--------+
             |                          |                 |
             v                          v                 v
       Tensor Store              Diffusion State      Scheduler
             |                          |                 |
             +-------------+------------+-----------------+
                           |
                           v
                    Model Execution
                           |
             +-------------+-------------+
             |                           |
             v                           v
       Quantized Linear           Bidirectional Attention
             |                           |
             +-------------+-------------+
                           |
                           v
                     Device Backend
                  CPU / NVIDIA / ...

7. Repository Layout

Suggested initial structure:

mojo-dllm/
├── README.md
├── LICENSE
├── pyproject.toml              # only if needed for tooling; not runtime
├── mojoproject.toml
├── src/
│   ├── main.mojo
│   ├── cli/
│   │   └── run.mojo
│   ├── formats/
│   │   ├── gguf.mojo
│   │   └── tensor_store.mojo
│   ├── tokenizer/
│   │   ├── tokenizer.mojo
│   │   └── bpe.mojo
│   ├── quant/
│   │   ├── q4_k.mojo
│   │   ├── q8_0.mojo
│   │   └── dequant.mojo
│   ├── kernels/
│   │   ├── matmul.mojo
│   │   ├── rmsnorm.mojo
│   │   ├── rope.mojo
│   │   ├── attention.mojo
│   │   ├── softmax.mojo
│   │   └── sampling.mojo
│   ├── models/
│   │   ├── model.mojo
│   │   └── llada.mojo
│   ├── diffusion/
│   │   ├── state.mojo
│   │   ├── scheduler.mojo
│   │   ├── masking.mojo
│   │   ├── confidence.mojo
│   │   └── sampler.mojo
│   └── runtime/
│       ├── device.mojo
│       ├── cpu.mojo
│       └── executor.mojo
├── tests/
│   ├── test_gguf.mojo
│   ├── test_quant.mojo
│   ├── test_attention.mojo
│   ├── test_scheduler.mojo
│   └── test_llada.mojo
├── benchmarks/
│   ├── bench_matmul.mojo
│   ├── bench_attention.mojo
│   ├── bench_generation.mojo
│   └── scripts/
└── docs/
    ├── architecture.md
    ├── model-porting.md
    └── benchmarks.md

8. Core Data Structures

8.1 Model configuration

struct ModelConfig:
    var vocab_size: Int
    var hidden_size: Int
    var intermediate_size: Int
    var num_layers: Int
    var num_attention_heads: Int
    var num_kv_heads: Int
    var head_dim: Int
    var max_position_embeddings: Int
    var rope_theta: Float64
    var mask_token_id: Int

8.2 Diffusion state

struct DiffusionState:
    var tokens: Tensor[Int32]
    var masked: BitSet
    var committed: BitSet
    var confidence: Tensor[Float32]
    var step: Int
    var total_steps: Int

Future versions MAY add:

var dirty_positions: BitSet
var block_state: BlockState
var cache_state: DiffusionCache

but cache-related structures SHOULD NOT be required for v0.1.

8.3 Tensor descriptor

struct TensorInfo:
    var name: String
    var shape: List[Int]
    var ggml_type: Int
    var offset: UInt64
    var byte_length: UInt64

9. GGUF Support

9.1 Required v0.1 support

The GGUF reader MUST support:

• GGUF header validation;
• metadata parsing;
• tensor metadata;
• tensor offsets;
• memory-mapped or direct file-backed reads;
• little-endian decoding;
• required tokenizer metadata;
• LLaDA tensor-name mapping.

Initial quantization support:

• Q4_K_M;
• optionally Q8_0;
• optionally F16 for correctness tests.

9.2 Unsupported formats

Unknown tensor types MUST fail explicitly:

unsupported GGUF tensor type: IQ3_XS (tensor blk.0.attn_q.weight)

No silent fallback.

A type with a reference decoder but without a matrix kernel currently fails at matrix load, for example:

tensor blk.0.attn_q.weight is Q8_0: matrix kernels exist for Q4_K and Q6_K

9.3 Loading strategy

Prefer memory mapping where practical:

GGUF file
   |
   v
memory map
   |
   +--> metadata
   |
   +--> tensor descriptors
   |
   +--> quantized tensor views

Avoid eagerly dequantizing the entire model into FP32.

10. Quantized Compute

10.1 Initial requirement

Implement Q4_K_M matrix multiplication directly against quantized weights.

Conceptually:

Q4_K_M weights
      |
      v
block decode
      |
      v
vectorized dot product
      |
      v
FP32/BF16 accumulation

The baseline MAY initially use a simple reference implementation.

The optimized implementation SHOULD:

• exploit SIMD;
• operate blockwise;
• avoid full-tensor dequantization;
• reuse scales efficiently;
• minimize temporary allocations.

10.2 Kernel API

Example:

fn linear_q4k[
    T: DType
](
    x: TensorView[T],
    weight: QuantizedTensorView,
    out: TensorView[T],
):
    ...

10.3 Benchmark requirements

Q4_K_M GEMM MUST be benchmarked independently.

Measurements:

• matrix dimensions;
• latency;
• effective memory bandwidth;
• effective FLOP/s where meaningful;
• allocations;
• thread count.

11. Transformer Execution

A LLaDA transformer layer will initially include:

input
  |
RMSNorm
  |
QKV projection
  |
RoPE
  |
bidirectional attention
  |
output projection
  |
residual
  |
RMSNorm
  |
MLP
  |
residual

11.1 Critical difference from autoregressive execution

The attention mask is not purely causal.

The runtime must support full/bidirectional attention over the relevant diffusion region.

Therefore, the initial implementation SHOULD NOT assume an append-only KV cache.

12. Diffusion Inference Loop

The core execution loop should be explicit and readable.

Pseudo-code:

tokens = tokenize(prompt)
canvas = append_masks(tokens, max_new_tokens)

for step in 0 .. num_steps:
    logits = model.forward(canvas)

    unresolved = get_masked_positions(canvas)

    scores = confidence(logits, unresolved)

    selected = scheduler.select(
        unresolved,
        scores,
        step,
        num_steps
    )

    sampled = sample(logits[selected])

    canvas[selected] = sampled

    if scheduler.requires_remasking:
        remask_low_confidence(canvas, scores)

return decode(canvas)

The scheduler MUST be modular.

13. Scheduler Interface

trait DiffusionScheduler:
    fn initialize(
        tokens: Tensor[Int32],
        generated_length: Int
    ) -> DiffusionState

    fn select_positions(
        state: DiffusionState,
        logits: TensorView[Float32]
    ) -> Selection

    fn update(
        state: DiffusionState,
        selection: Selection,
        sampled_tokens: TensorView[Int32]
    )

Initial scheduler:

• reference implementation matching LLaDA inference behavior.

Future schedulers:

• confidence-based;
• entropy-based;
• fixed block count;
• adaptive steps;
• block diffusion.

14. Sampling

Required primitives:

• softmax;
• log-softmax where needed;
• entropy;
• argmax;
• top-k;
• categorical sampling;
• temperature;
• deterministic PRNG;
• masked-position gather/scatter.

Possible future fused kernel:

logits
  |
softmax
  |
entropy/confidence
  |
top-k
  |
sample
  |
scatter

15. Tokenizer

v0.1 MUST perform tokenization without invoking Python.

Possible implementation paths:

1. reuse or adapt an existing Mojo BPE tokenizer;
2. implement the minimum tokenizer behavior required by the target model;
3. parse tokenizer metadata from GGUF where available.

The tokenizer MUST correctly support:

• prompt encoding;
• mask token;
• BOS/EOS if required;
• decoding generated IDs;
• byte fallback if required by the model tokenizer.

16. CPU Runtime

The initial backend is x86 CPU.

Priority order:

Stage 1 — correctness

• scalar/reference kernels;
• single thread;
• deterministic output.

Stage 2 — vectorization

• SIMD vector operations;
• vectorized dequantization;
• vectorized dot products;
• optimized softmax.

Stage 3 — threading

Potential parallel regions:

• output channels in linear layers;
• attention heads;
• MLP projections;
• batch-independent work.

Threading MUST be benchmark-driven.

17. NVIDIA Backend

The NVIDIA backend begins only after the CPU runtime passes correctness tests.

Initial GPU target:

• consumer NVIDIA RTX 30-series;
• FP16/BF16 where supported;
• quantized model weights;
• device-resident hot tensors.

Required kernels:

• Q4_K_M GEMM;
• RMSNorm;
• RoPE;
• attention;
• softmax;
• diffusion sampling operations.

17.1 Device abstraction

The runtime SHOULD expose:

trait Device:
    fn allocate(...)
    fn copy_to_device(...)
    fn copy_from_device(...)
    fn synchronize(...)

Kernel implementation should avoid unnecessary backend-specific code where Mojo abstractions are sufficient.

18. CPU/GPU Hybrid Execution

Future laptop-focused mode:

system RAM
   |
model weights
   |
+--+--------------------+
|                       |
CPU                  GPU VRAM
|                       |
fallback layers      hot layers
|                       |
+-----------+-----------+
            |
       diffusion state

Potential strategies:

• layer offload;
• operator offload;
• GPU-only attention;
• GPU-only matrix multiplication;
• dynamic VRAM budget.

This is NOT required for v0.1.

19. Correctness Strategy

The project MUST use at least two external reference implementations when available.

Recommended references:

• official PyTorch LLaDA implementation;
• llama.cpp diffusion implementation;
• diffuse-cpp.

Validation levels:

Level 1 — tensor-level

Compare:

• embeddings;
• norm outputs;
• Q/K/V projections;
• RoPE outputs;
• attention outputs;
• MLP outputs;
• final logits.

Use small contexts and fixed inputs.

Level 2 — diffusion-step parity

Compare:

• mask positions;
• confidence values;
• selected positions;
• sampled token IDs;
• updated canvas.

Level 3 — generation parity

With fixed:

• model;
• prompt;
• seed;
• temperature;
• diffusion steps;
• generation length;

compare final output and intermediate states.

Exact parity may not always be possible across floating-point implementations, so tolerances MUST be documented.

20. Testing

Minimum unit tests:

File format

• valid GGUF header;
• invalid magic;
• metadata parsing;
• tensor offset calculations;
• Q4_K block parsing.

Quantization

• known Q4_K block decode;
• scalar vs SIMD result;
• matrix-vector correctness.

Transformer

• RMSNorm;
• RoPE;
• softmax;
• attention;
• MLP;
• full layer output.

Diffusion

• mask initialization;
• confidence computation;
• selection logic;
• remasking;
• deterministic RNG.

Integration

• tiny synthetic model;
• single transformer layer;
• short prompt;
• one diffusion step;
• full LLaDA run.

21. Benchmark Plan

21.1 Comparison targets

Primary:

• llama.cpp;
• diffuse-cpp.

Secondary:

• official PyTorch/Hugging Face implementation.

21.2 Metrics

Measure:

• time to first completed denoising step;
• time per denoising step;
• total generation latency;
• generated tokens per second;
• denoising steps per second;
• peak RSS;
• peak VRAM;
• model load time;
• CPU utilization;
• GPU utilization;
• energy consumption where practical.

21.3 Benchmark matrix

Example:

22. Performance Philosophy

Correctness comes before optimization.

Optimization order:

correct model execution
        |
        v
Q4_K compute
        |
        v
CPU SIMD
        |
        v
threading
        |
        v
fused kernels
        |
        v
GPU
        |
        v
inter-step caching
        |
        v
block diffusion

Do not introduce caching before a fully correct uncached baseline exists.

23. Future Caching Research

After the baseline runtime is stable, experiment with:

23.1 Inter-step cache

Track which intermediate state can be reused between adjacent denoising iterations.

23.2 Dirty-position tracking

struct DirtyState:
    var changed_positions: BitSet
    var changed_blocks: BitSet

23.3 Layer-level reuse

Explore whether selected intermediate layer outputs can be reused under configurable error thresholds.

23.4 Block diffusion

Support models that denoise blocks rather than the entire generation canvas.

All cache optimizations MUST quantify:

• speedup;
• additional memory use;
• output divergence;
• benchmark quality impact.

24. Integration with Other Mojo Runtimes

The core should eventually expose a small library API.

Example:

var model = DiffusionModel.load("model.gguf")

var config = GenerationConfig(
    max_new_tokens=128,
    steps=128,
    temperature=0.2,
    seed=42,
)

var output = model.generate(
    "Explain KV-cache quantization.",
    config
)

A runtime such as molla could then provide:

• HTTP serving;
• OpenAI-compatible APIs;
• model lifecycle;
• chat templates;
• request scheduling.

mojo-dllm should remain focused on dLLM execution.

25. CLI

Initial interface:

mojo-dllm run [OPTIONS]

Required flags:

--model PATH
--prompt TEXT
--max-tokens N
--steps N
--seed N

Optional:

--temperature FLOAT
--threads N
--device cpu
--verbose
--dump-step-stats

Example:

mojo-dllm run \
  --model ./llada-8b-q4_k_m.gguf \
  --prompt "What is paged attention?" \
  --max-tokens 128 \
  --steps 128 \
  --seed 1234 \
  --threads 8

26. Logging and Diagnostics

Verbose mode SHOULD print:

model: LLaDA-8B
quantization: Q4_K_M
device: CPU
threads: 8
prompt tokens: 17
generation tokens: 128
diffusion steps: 128

load time: 1.21 s
step 1: 83.2 ms
step 2: 81.7 ms
...
generation: 10.4 s
peak RSS: 6.1 GB

Debug mode MAY expose:

• selected token positions;
• confidence distributions;
• per-layer timings;
• kernel timings.

27. Error Handling

The runtime should fail with actionable errors.

Examples:

ERROR: unsupported model architecture: diffusion_gemma (supported: llada, dream)
ERROR: unsupported GGUF tensor type: IQ2_XXS (tensor blk.0.attn_q.weight)
ERROR: required metadata key missing: tokenizer.ggml.model
ERROR: tensor shape mismatch for blk.0.attn_q.weight
ERROR: model requires 8.1 GB but available memory is 5.9 GB

Avoid assertion-only failure paths in user-facing code.

28. Memory Strategy

For the initial CPU backend:

• memory-map GGUF;
• keep quantized weights in-place;
• allocate reusable activation buffers;
• avoid layer-by-layer heap churn;
• preallocate diffusion state;
• reuse attention workspaces.

Ideal high-level layout:

mapped GGUF
    |
    +--> quantized weights

runtime arena
    |
    +--> hidden state
    +--> Q/K/V
    +--> attention workspace
    +--> MLP workspace
    +--> logits
    +--> diffusion scheduler state

29. Milestones

M0 — repository bootstrap

Deliverables:

• build system;
• CLI skeleton;
• CI;
• basic tests;
• coding conventions.

M1 — GGUF inspection

Deliverables:

mojo-dllm inspect model.gguf

Outputs:

• architecture;
• tensor count;
• tensor names;
• shapes;
• quantization formats;
• tokenizer metadata.

M2 — Q4_K_M primitives

Deliverables:

• parser;
• dequantization;
• reference matvec;
• numerical tests;
• microbenchmarks.

M3 — transformer primitives

Deliverables:

• RMSNorm;
• RoPE;
• softmax;
• attention;
• MLP;
• transformer layer tests.

M4 — LLaDA forward pass

Deliverables:

• model configuration;
• tensor mapping;
• embedding;
• transformer stack;
• logits;
• comparison against reference implementation.

M5 — diffusion generation

Deliverables:

• mask canvas;
• scheduler;
• confidence;
• sampling;
• iterative generation;
• complete text output.

At this milestone:

mojo-dllm run ...

must work.

M6 — CPU optimization

Deliverables:

• SIMD;
• multithreading;
• fused kernels;
• benchmark report.

M7 — NVIDIA backend

Deliverables:

• device allocation;
• GPU kernels;
• end-to-end GPU inference;
• consumer RTX benchmark.

M8 — second architecture

Deliverables:

• Dream support;
• model abstraction cleanup;
• generalized dLLM scheduler interface.

30. Initial Definition of Done

v0.1 is complete when all of the following are true:

• LLaDA Q4_K_M GGUF loads directly in Mojo.
• No Python process is required during inference.
• Prompt tokenization works.
• Full transformer forward pass executes.
• Bidirectional attention is correct.
• Diffusion scheduler is implemented.
• The model generates readable text.
• Fixed-seed runs are deterministic.
• Tests cover quantization, attention and scheduler primitives.
• Peak memory fits within a reasonable consumer laptop RAM budget.
• Benchmark comparison with diffuse-cpp is published.
• Benchmark comparison with llama.cpp is published.
• README documents known differences and limitations.

31. Risks

31.1 Mojo ecosystem churn

Mojo APIs may still evolve quickly.

Mitigation:

• keep backend abstraction small;
• pin compiler/toolchain versions;
• isolate unstable APIs;
• use CI against a documented supported version.

31.2 Q4_K performance

Naive Q4_K dequantization may make CPU inference unusably slow.

Mitigation:

• optimize Q4_K early;
• benchmark it independently;
• borrow mathematical layout ideas from permissively licensed implementations while respecting licenses.

31.3 GGUF complexity

GGUF supports many metadata conventions and quantization formats.

Mitigation:

• support only the exact subset required by LLaDA initially;
• reject unsupported formats explicitly;
• expand only after v0.1.

31.4 Tokenizer mismatch

Small tokenizer differences can invalidate inference.

Mitigation:

• validate token IDs against a known reference tokenizer;
• add golden prompt/token test cases.

31.5 Floating-point divergence

Different kernel implementations may diverge slightly.

Mitigation:

• use tensor-level tolerances;
• compare logits statistically;
• separate numerical mismatch from algorithmic mismatch.

31.6 GPU portability may underperform

A single portable kernel may not match highly specialized vendor kernels.

Mitigation:

• measure rather than assume portability;
• allow compile-time specialization;
• permit small backend-specific sections if required;
• quantify shared vs backend-specific LOC.

32. Success Criteria

Minimum success

A working pure-Mojo CPU implementation of LLaDA-8B Q4_K_M.

Strong success

CPU performance within a reasonable factor of diffuse-cpp or llama.cpp, with a considerably smaller and more readable runtime.

High-impact success

One shared Mojo implementation runs efficiently on:

• x86 CPU;
• NVIDIA GPU;
• Apple GPU;
• AMD GPU;

with limited backend-specific code.

Research-quality success

The project demonstrates a useful dLLM-specific optimization—such as efficient inter-step reuse or portable block diffusion—that produces a measurable speedup while preserving output quality.

33. Project Positioning

Recommended one-line description:

Portable local inference for diffusion language models, written in Mojo.

Alternative:

Run LLaDA and Dream locally from GGUF using a Mojo-native diffusion inference runtime.

Avoid claims such as:

• "first diffusion inference engine";
• "first fast dLLM runtime";
• "novel diffusion caching";
• "faster than autoregressive LLMs";

unless validated by reproducible benchmarks and a fresh prior-art review.

A defensible early claim is:

A standalone Mojo-native runtime for quantized local diffusion-language-model inference.

34. README Benchmark Target

The eventual repository front page should contain a table like:

No synthetic or estimated performance numbers should be published.

35. Immediate Implementation Order

The recommended sequence is:

1. GGUF metadata reader
2. LLaDA tensor mapping
3. Q4_K_M block decoder
4. reference Q4_K matvec
5. RMSNorm
6. RoPE
7. softmax
8. bidirectional attention
9. MLP
10. single transformer layer
11. full transformer
12. tokenizer
13. diffusion canvas
14. confidence calculation
15. scheduler
16. sampling
17. full generation
18. correctness comparison
19. SIMD optimization
20. NVIDIA backend

This sequence minimizes the number of moving pieces during debugging.

36. First Engineering Spike

Before committing to the full runtime, implement a one-week feasibility spike:

Deliverable A

mojo-dllm inspect llada.gguf

Correctly prints model metadata and tensor layout.

Deliverable B

Load one Q4_K_M matrix and reproduce a reference matrix-vector multiplication within numerical tolerance.

Deliverable C

Implement one complete LLaDA transformer layer and compare output tensors against a PyTorch or C++ reference.

If these three succeed cleanly, proceed with the full project.

If Q4_K_M or GGUF support proves disproportionately difficult because of current Mojo limitations, reconsider whether the first version should use Q8_0 before returning to Q4_K_M.

37. Core Principle

The project should optimize for one question:

Can Mojo provide a clean, performant, portable execution layer for diffusion language models on ordinary local hardware?

Everything that does not help answer that question should remain outside the initial scope.