Jais-3-0.9B-Uno

Jais-3-0.9B-Uno makes inception42/Jais-3-0.9B generate text 1.51× faster on average, with exactly the same output distribution as the base model.

It is a small diffusion adapter (LoRA, rank 128) trained with the Uno recipe on top of the frozen Jais-3-0.9B weights. This repository contains the adapter only; the base weights are in inception42/Jais-3-0.9B.

Paper: Unlocking Lossless Speedups in LLMs via Discrete Diffusion · Code / inference engine: github.com/ifm-ai/uno · Project page: s-sahoo.github.io/uno

How it works

Uno keeps two pathways inside one model:

• AR pathway: the original Jais-3-0.9B weights, unchanged. This is the model whose answers you get.
• Diffusion pathway: the same weights plus this LoRA adapter. It drafts several future tokens in a single forward pass.

At each step the diffusion pathway proposes a block of tokens, and the AR pathway checks the whole block in one forward pass with the Ψ-Spec sampler. Accepted tokens are kept, and the first rejected token is replaced by a token sampled from the AR model. The generated text therefore follows exactly the distribution of inception42/Jais-3-0.9B, in greedy decoding as well as in sampling. No separate draft model is needed, and quality is never traded for speed.

Quality: because Uno is lossless, the benchmark results of the base model apply unchanged. Please refer to the inception42/Jais-3-0.9B model card.

Speed

All speed numbers are measured with greedy decoding. Results for Jais-3-0.9B-Uno:

Task	AR (tok/s)	Uno (tok/s)	Speedup
Math (MATH-500)	368	595	1.62×
Agentic, multi-turn tool use (ToolHop)	359	567	1.58×
Code (LiveCodeBench v6)	366	553	1.51×
Arabic chat (Arabic AlpacaEval)	371	458	1.24×
Average (256 prompts)	366	552	1.51×

Where the numbers come from:

• Hardware and engine: one H100 80GB GPU running nano_vllm_uno, with FlashAttention-3 and CUDA graphs.
• Baseline (AR): the base model decoded token by token in the same engine, with the same kernels and no adapter.
• Uno: tree sampler with tree_verify_size=32 and tree_candidate_top_k=32. We report the faster of diffusion block sizes 8 and 16 (block 8 for 0.9B, 3.7B and block 16 for 7B, 32B); the two are within 2.5% of each other.
• Workload: batch size 1, greedy decoding, generation stopped at the natural end-of-sequence token or at 4,096 new tokens. We used 256 prompts, 64 per task.
• Speedup: total generated tokens divided by total wall-clock time, relative to AR.
• Sampling: with sampling (T=1, top-p 0.95, top-k 50) the speedups are 10–15% lower.

Usage

Install

git clone https://github.com/ifm-ai/uno && cd uno
pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cu128
pip install -e .
# The tree sampler (recommended on Hopper GPUs) needs FlashAttention-3:
#   see "Installation" in the repository README.

Command line

python inference.py \
  --model inception42/Jais-3-0.9B \
  --gated-lora-path inception42/Jais-3-0.9B-Uno \
  --prompt "What is the capital of the United Arab Emirates? Answer in one sentence." \
  --attention-backend fa3 \
  --diffusion-block-size 8 \
  --tree-verify-size 32 --tree-candidate-top-k 32 \
  --temperature 0.6 --top-p 0.95 --top-k 50 --max-tokens 1024

On GPUs without FlashAttention-3, drop --attention-backend fa3 and the two --tree-* options to use the linear sampler.

Python

import json
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from nano_vllm_uno import LLM, SamplingParams

base = snapshot_download("inception42/Jais-3-0.9B")
adapter = snapshot_download("inception42/Jais-3-0.9B-Uno")
tok = AutoTokenizer.from_pretrained(base)
cfg = json.load(open(f"{base}/config.json"))
gen = json.load(open(f"{base}/generation_config.json"))
eos = gen["eos_token_id"] if isinstance(gen["eos_token_id"], list) else [gen["eos_token_id"]]

llm = LLM(
    model=base, tokenizer_path=base,
    gated_lora_path=adapter,            # the Uno adapter; None = plain AR decoding
    attention_backend="fa3",
    max_diffusion_block_size=32,        # must equal tree_verify_size for the tree sampler
    tree_verify_size=32, tree_candidate_top_k=32,
    cuda_graph_block_sizes=[1, 8],
)
params = SamplingParams(
    temperature=0.6, top_p=0.95, top_k=50, max_tokens=1024,
    diffusion_block_size=8,
    mask_token_id=cfg["vocab_size"],    # Uno mask token = vocab size
    stop_token_ids=eos,
)
messages = [{"role": "user", "content": "Explain speculative decoding in two sentences."}]
prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
prompt_ids = tok.encode(prompt, add_special_tokens=False)
out = llm.generate([prompt_ids], params)
print(tok.decode(out[0]["token_ids"], skip_special_tokens=True))

The adapter tensors are stored without the base_model.model. prefix, using the same layout as K2-Horizon-7B-Uno. They are meant for nano_vllm_uno; plain PeftModel.from_pretrained will not load them.

Limitations

• Batch size. Uno accelerates latency-bound decoding. The gain shrinks as concurrency grows, and with many parallel requests (around 64) plain AR decoding is faster.
• Engine support. Uno needs nano_vllm_uno or another Ψ-Spec implementation. Standard vLLM, SGLang and Transformers generate cannot use this adapter.
• Content dependence. Speedups are largest on structured, predictable text such as math and tool calls, and smallest on open-ended chat.
• The model's capabilities, biases and limitations are those of inception42/Jais-3-0.9B.

Citation

@misc{jais3_2026,
  title  = {Jais 3: Introducing Arabic-centric Agentic AI},
  author = {{IFM Team, Inception42 Team, and Cerebras Team}},
  year   = {2026},
  url    = {https://ifm.ai/blog/jais3/},
}