⚡ llada.rs

A diffusion-native language model runtime, written in Rust.

Think llama.cpp — rebuilt around denoising blocks instead of next-token decoding.

No PyTorch runtime. No training stack. No autoregressive-first assumptions.

Quick start · Architecture · Entry points · Roadmap

---

llada.rs is a standalone inference runtime for masked-diffusion and block-diffusion language models. Its first execution contract targets LLaDA2-style generation: keep a clean prefix cache, denoise a mutable token block with bidirectional attention inside that block, perform one unchanged final forward, commit its KV, and move to the next block.

This is not another general-purpose autoregressive transformer engine. A small causal path exists only where it is needed to prefill the clean prefix and prove parity with imported model backbones.

Important

The project is pre-alpha and correctness-first. The reference CPU graph, runtime state machine, strict GGUF contract and checkpoint converter exist; production tokenization, a generate CLI, SIMD and GPU backends are still on the roadmap.

Why a diffusion-native runtime?

Autoregressive engines are built around one dominant operation: append a token to a causal KV cache. A block-diffusion model has a different state machine.

flowchart LR
    P[Clean prefix cache] --> F[Masked block forward]
    B[Mutable masked block] --> F
    F --> D{Denoise policy}
    D -->|masks remain| B
    D -->|fully denoised| C[Unchanged final forward]
    C --> K[Atomically commit block KV]
    K --> P

Loading

The final unchanged forward is not an implementation detail: it is the cache validity boundary. llada.rs makes that boundary explicit in the model, runtime and trace APIs.

What is already here

Area	Current capability
Runtime	Allocation-stable BlockSession, clean-prefix/mutable-block lifecycle, forward limits, EOS stopping and exact final-KV commit
Denoising	Progress-safe low-confidence decoding, dInfer-compatible adaptive threshold and uniform fixed-count schedule
LLaDA2 graph	Embedding, untied head, ordinary RMSNorm, partial-RoPE GQA, dense SwiGLU, grouped sigmoid routing, merged MoE experts and shared expert
GGUF	Safe GGUF v3 parsing, bounds/type/shape validation, read-only memory mapping and strict namespaced diffusion metadata
Tensor formats	F32, F16, BF16, Q8_0, Q4_0 and Q4_K reference tensor views/kernels
Conversion	Streaming sharded safetensors → GGUF v3 conversion for official per-expert and dFactory/VeOmni merged-expert LLaDA2 layouts
Parity	Versioned dInfer/SGLang reference traces with exact token/top-1 comparison and configurable confidence tolerance

Architecture

flowchart TB
    HF[Hugging Face checkpoint<br/>config + tokenizer + safetensors]
    CONVERT[Offline streaming converter]
    GGUF[Portable GGUF v3<br/>weights + tokenizer + llada.* contract]
    MAP[Safe read-only mapping]
    MODEL[LLaDA2 reference graph]
    SESSION[BlockSession]
    POLICY[DenoisePolicy]
    TRACE[ReferenceTrace]

    HF --> CONVERT --> GGUF --> MAP --> MODEL --> SESSION
    POLICY <--> SESSION
    SESSION --> TRACE
    DINFER[dInfer / SGLang] --> TRACE

Loading

The module boundaries intentionally keep training frameworks out of the runtime:

src/
├── gguf.rs              safe GGUF parser and tensor directory
├── gguf_writer.rs       streaming GGUF v3 writer
├── convert.rs           Hugging Face/safetensors importer
├── tensor.rs            mapped tensor views and quantized layouts
├── cpu.rs               scalar correctness kernels
├── model/               architecture graphs, workspaces and caches
├── runtime.rs           block generation, policies and request state
├── trace.rs             cross-runtime parity trace contract
└── bin/llada.rs         thin CLI over the library

Quick start

You need a current stable Rust toolchain.

git clone https://github.com/getStRiCtd/llada.rs.git
cd llada.rs
cargo test
cargo build --release

Convert a LLaDA2 checkpoint

The converter reads sharded safetensors directly. It does not import PyTorch and does not materialize the full checkpoint in memory.

cargo run --release --bin llada-cli -- \
  convert /path/to/hf-checkpoint /path/to/model.gguf \
  --block-length 32 \
  --revision inclusionAI/LLaDA2.0-mini

Optional flags:

Flag	Meaning
--block-length N	Diffusion block length written into the runtime contract; default: 32
--revision REV	Source checkpoint revision recorded in GGUF metadata
--mask-id N	Override a mask token ID that cannot be resolved from tokenizer files
--eos-id N	Override an EOS token ID that cannot be resolved from tokenizer files
--allow-extra-tensors	Permit source tensors outside the strict LLaDA2 manifest after manual inspection

The correctness-first converter preserves F32/F16/BF16 storage. It does not silently cast FP8 checkpoints and does not quantize weights. It refuses to overwrite an existing output, writes through a sibling .partial file, computes a canonical SHA-256 while streaming and reopens the result through the Rust manifest validator before reporting success.

Inspect and validate GGUF

cargo run --release --bin llada-cli -- inspect model.gguf
cargo run --release --bin llada-cli -- validate model.gguf

inspect checks the generic container and prints its directory summary. validate dispatches by general.architecture and checks the complete model contract before model allocation.

Entry points

CLI

Command	Entry point	Purpose
llada-cli convert …	src/bin/llada.rs	Convert an exported Hugging Face LLaDA2 checkpoint to GGUF v3
llada-cli inspect MODEL	Gguf::open	Parse and inspect a generic GGUF container
llada-cli validate MODEL	Llada2GgufMetadata::validate_gguf	Validate diffusion semantics and the full tensor manifest
llada-cli validate-trace TRACE	ReferenceTrace::validate_complete	Validate a dInfer/SGLang or Rust parity trace

Rust library

Layer	Primary API	Use it for
Checkpoint import	convert_hf_to_gguf	Programmatic safetensors → GGUF conversion
Container	Gguf / MappedGguf	Safe metadata parsing and zero-copy tensor access
Model	Llada2Model::from_mapped	Construct the complete validated reference graph
Runtime	BlockModel / BlockSession	Integrate another backend or run one block-diffusion request
Policies	LowConfidence, DinferThreshold, FixedCount	Select the reveal schedule independently from the model
Parity	ReferenceTrace	Record, validate and compare forward-by-forward token decisions
GGUF output	write_gguf_v3	Build a streaming offline exporter without checkpoint-sized allocations

The core runtime can be embedded without the CLI:

use llada::{
    model::{
        llada2_gguf::Llada2GgufMetadata,
        llada2_model::Llada2Model,
    },
    runtime::{BlockGenerationConfig, BlockSession, DinferThreshold},
    tensor::MappedGguf,
};

# fn main() -> llada::Result<()> {
let mapped = MappedGguf::open("model.gguf")?;
let metadata = Llada2GgufMetadata::validate_gguf(mapped.metadata())?;
let model = Llada2Model::from_mapped(&mapped)?;
let policy = DinferThreshold::new(0.9)?;

// Tokenizer execution is not public yet: supply token IDs from the embedded
// tokenizer contract or from the checkpoint tokenizer.
let prompt_ids = vec![1_u32, 2, 3];
let config = BlockGenerationConfig {
    mask_token_id: metadata.mask_token_id,
    eos_token_id: metadata.eos_token_id,
    max_new_tokens: 64,
    max_steps_per_block: 64,
};

let mut session = BlockSession::new(model, policy, &prompt_ids, config)?;
let generated_ids = session.run_to_completion()?;
println!("{generated_ids:?}");
# Ok(())
# }

dInfer / SGLang trace export

tools/dinfer_trace_export.py is a small hook, not another model runner. Create one DinferTraceRecorder per request, record every decoder update, then record the unchanged cache-populating final forward with committed=True.

from tools.dinfer_trace_export import DinferTraceRecorder

trace = DinferTraceRecorder(
    model_revision="checkpoint-revision",
    architecture="llada2_moe",
    mask_token_id=mask_id,
    vocab_size=vocab_size,
    block_size=block_size,
    prompt_tokens=prompt_ids,
)

# Inside the dInfer/SGLang block loop:
trace.record(
    block_offset=block_offset,
    forward_index=forward_index,
    input_tokens=before_update,
    output_tokens=after_update,
    logits=logits,
)

# After the block is fully denoised, run it unchanged and commit that exact KV:
trace.record(
    block_offset=block_offset,
    forward_index=final_forward_index,
    input_tokens=final_tokens,
    output_tokens=final_tokens,
    logits=final_logits,
    committed=True,
)
trace.finish("reference-trace.json")

Then validate it locally:

cargo run --release --bin llada-cli -- validate-trace reference-trace.json

Design laws

• Diffusion-first. Block denoising semantics are the product; AR exists only as a narrowly scoped parity/prefill tool.
• Portable weights. GGUF v3 is the base container; diffusion extensions live under namespaced llada.* metadata.
• Explicit state. Shared weights, request-local state, mutable block KV and committed prefix KV have distinct lifetimes.
• No surprise allocations in the hot path. Workspaces and trace storage are allocated before repeated forwards.
• Reference before acceleration. Every SIMD, parallel CPU or GPU kernel must match a simple tested implementation.
• Framework independence. FSDP2, VeOmni, dFactory and training code stop at the offline export boundary.

Project status

The next high-value gates are:

1. Run a full official dInfer/SGLang trace through the complete LLaDA2 graph.
2. Execute the tokenizer embedded in GGUF and expose tokenize / generate.
3. Split shared model weights from request-local state for FDFO-style batching.
4. Profile, then add SIMD/parallel CPU kernels and a WGPU backend.

See the full correctness gates and milestone order in ROADMAP.md.

Contributing

Keep changes small, measurable and parity-driven. For every new kernel: land a clear reference implementation and test first, then optimize it and prove equivalence.

Before opening a change:

cargo fmt --check
cargo check
cargo test

Rust is ready for diffusion-native inference. Let's build the runtime it deserves.