I ran Nemotron-Labs-Diffusion on DGX Spark and actually measured tri-mode generation

Introduction

Hello, I'm Morishige from Classmethod's Manufacturing Business Technology Department.

In May 2026, NVIDIA released a new model family called "Nemotron-Labs-Diffusion." Since the name includes "diffusion," some of you might have thought this was about image generation, but this is actually a "diffusion language model" that generates text. Keeping that in mind from the start will make the rest of this article easier to follow.

https://research.nvidia.com/publication/2026-05_nemotron-labs-diffusion-tri-mode-language-model-unifying-autoregressive

The LLMs we typically use write text one token at a time from left to right. Since each token must be determined before moving to the next, generation is inherently sequential. Diffusion language models change this. They generate multiple tokens in parallel in block units, and fill in uncertain parts in later steps.

What makes Nemotron-Labs-Diffusion interesting is that a single model can switch between three decoding modes (AR / Diffusion / Linear Self-Speculation). NVIDIA calls this tri-mode. Both autoregressive (AR) and diffusion modes coexist within the same checkpoint rather than as separate models.

Reading through the published information, the three things that personally caught my attention were:

• Can all three modes run smoothly in a local DGX Spark environment?
• What does the diffusion mode's "generate in parallel and fill in later" behavior actually look like when running it?
• What happens to the official figures like "2.7× over AR" and "6× tokens per forward" when running raw without quantization?

So I loaded the Nemotron-Labs-Diffusion 8B model onto a DGX Spark and measured all three modes. To give the conclusion upfront: all three modes ran on DGX Spark. The speed of Linear Self-Speculation was 1.75 to 1.98× faster than AR. However, the impressive numbers NVIDIA advertises come with conditions, and when generating longform text in BF16 without quantization, a different picture emerges. I'll cover those differences as well.

As a related note, I previously wrote an article about measuring speculative decoding with an inference engine on DGX Spark using Gemma 4 MTP. Gemma 4 MTP uses a separate draft model, while Linear Self-Speculation is self-contained within the same model — but the skeleton of "quickly producing a draft that the main model verifies" is the same. Reading them together should help you understand where diffusion language models fit in.

https://dev.classmethod.jp/articles/dgx-spark-gemma4-mtp-multi-token-prediction-bench/

What is Nemotron-Labs-Diffusion's tri-mode?

Model Family and Background

Nemotron-Labs-Diffusion offers text generation models in three sizes — 3B / 8B / 14B — each with a Base version (pretrained only) and an Instruct version (with instruction tuning). A VLM-8B that also handles images is also available. For this article, I focused on the instruction-tuned 8B (nvidia/Nemotron-Labs-Diffusion-8B).

The architecture is a straightforward Transformer, not a Mamba hybrid. According to the technical report, training started from the existing Ministral3-8B as a base, followed by continued pretraining with autoregressive-only objectives on 1T tokens, then continued pretraining with a combined autoregressive and diffusion objective on 300B tokens, and finally instruction tuning on 45B tokens. The license is the NVIDIA Nemotron Open Model License.

One subtly important point here is that training autoregressive and diffusion together doesn't degrade AR accuracy. The technical report states that compared to a model trained under the same conditions without the diffusion loss, AR accuracy actually improved by 0.14–0.43% after joint training. The nice thing about this design is that autoregressive capability isn't sacrificed in order to add diffusion capability.

Three Decoding Modes

Let me organize the contents of tri-mode. By switching how attention is applied at inference time, the same weights support three different generation methods.

AR mode is conventional autoregressive decoding. It confirms one token per forward pass, from left to right. You call ar_generate().

Diffusion mode divides the generation span into blocks of block_length tokens, masks them all at once, and gradually fills in positions with the highest confidence. You call generate() and pass diffusion-specific parameters block_length and threshold. threshold determines how many tokens to confirm at each step — higher means more cautious, lower means filling in more at once.

Linear Self-Speculation mode is a self-speculative decoding approach combining diffusion and autoregressive. Diffusion generates candidate tokens in parallel as a draft, which autoregressive then verifies and confirms up to the correct point. It's the same idea as speculative decoding, but without needing a separate small draft model — the diffusion mode of the same model doubles as the drafter. You call linear_spec_generate().

In terms of not requiring separate weight files for the draft role, this is similar in spirit to DeepSeek's MTP, which embeds a dedicated drafting module in the main model. However, whereas DeepSeek adds a dedicated module, Nemotron-Labs-Diffusion doesn't add even that — it repurposes its built-in diffusion mode directly as the draft. Gemma 4 MTP distributes a separate draft model, so comparing all three — "separate model," "built-in module," and "mode repurposing" — reveals subtle differences in where the draft comes from.

All three methods return nfe (number of forward evaluations — how many times the forward pass was run) in addition to the generated result. If the same length of text can be produced with fewer forward passes, that means it's faster. Throughout this article, I'll frequently use nfe and the derived metric tokens per forward (how many tokens are confirmed per forward pass) as speed indicators.

Test Environment

I used a single DGX Spark for testing. The hardware and software configuration is as follows:

The three tri-mode modes are switched using the custom methods ar_generate() / generate() / linear_spec_generate(). These are implemented in custom modeling (modeling_*.py) bundled with the model repository, and cannot be easily called from vLLM's standard serving functionality. For this article, I loaded the model from transformers with trust_remote_code=True and called the methods directly.

Setup and Model Loading

I created a virtual environment for testing using uv and installed the necessary libraries. PyTorch is specified with the cu130 build to match the CUDA 13 environment on DGX Spark.

cd ~/works/nemotron-labs-diffusion
uv venv --python 3.13 .venv
uv pip install --python .venv/bin/python \
    torch 'transformers>=5.0' peft accelerate datasets \
    --torch-backend=cu130

Model loading follows the model card sample, with trust_remote_code=True being required. This is needed to load the custom modeling.

from transformers import AutoModel, AutoTokenizer
import torch

repo = "nvidia/Nemotron-Labs-Diffusion-8B"
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModel.from_pretrained(repo, trust_remote_code=True)
model = model.cuda().to(torch.bfloat16).eval()

The three modes are called as follows. Applying the chat template to create prompt_ids is common across all modes.

history = [{"role": "user", "content": "Please explain diffusion language models in one sentence."}]
text = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
prompt_ids = tokenizer(text, return_tensors="pt").input_ids.cuda()
eos = tokenizer.eos_token_id

# AR mode
out_ids, nfe = model.ar_generate(prompt_ids, max_new_tokens=512)

# Diffusion mode
out_ids, nfe = model.generate(
    prompt_ids, max_new_tokens=512, block_length=32, threshold=0.9, eos_token_id=eos,
)

# Linear Self-Speculation mode
out_ids, nfe = model.linear_spec_generate(
    prompt_ids, max_new_tokens=512, block_length=32, eos_token_id=eos,
)

A LoRA adapter (linear_spec_lora) bundled in the model repository is available for Linear Self-Speculation to further extend acceptance. When using this, you attach it with PeftModel and then call the method from the unwrapped base model.

from peft import PeftModel

model = PeftModel.from_pretrained(model, repo, subfolder="linear_spec_lora").eval()
base = model.model  # Call linear_spec_generate from the unwrapped base
out_ids, nfe = base.linear_spec_generate(
    prompt_ids, max_new_tokens=512, block_length=32, eos_token_id=eos,
)

Actually, in this straightforward procedure, the diffusion mode and Linear Self-Speculation mode initially stop with an exception in transformers 5.9. The cause and workaround are explained together in the "Common Pitfalls on DGX Spark" section. For now, I'll just note that "inserting one shim makes all three modes work."

Measuring the Speed of tri-mode's Three Modes

Here's the main part. I ran the same 12 longform prompts (asking for explanations of technical topics, mixing Japanese and English) through four configurations, and took the average excluding warmup. Generation used max_new_tokens=512 with temperature 0 greedy decoding.

Warm tok/s for 4 configurations. Using AR as the baseline: Diffusion 1.20×, Linear Self-Speculation 1.75×, with LoRA 1.98×. All use the same 8B model with the same weights, only switching the decoding method.

The numerical results are as follows:

A clean staircase pattern emerged. What I want to highlight is tokens/forward and nfe. Since AR produces 1 token per forward, generating 512 tokens requires 512 forward passes. Linear Self-Speculation + LoRA confirms an average of 2.08 tokens per forward, generating roughly the same length of text in just 252 forward passes. The fact that the number of forward passes is cut in half directly translates into the speed difference.

Let me put the official figures in context here. The technical report includes measurements on DGX Spark, showing the 8B diffusion mode at 77.5 tok/s with FP8 quantization (3.14× over AR) and 112.5 tok/s with INT4 quantization (2.69× over AR). My measurements here look considerably lower, but that's because I'm running BF16 without quantization, through transformers without inference engine optimizations. The absolute values change substantially based on the runtime and quantization. What I care about in this article isn't that — it's the relative speed difference between AR and diffusion-based approaches under the same conditions. In that sense, Linear Self-Speculation achieving 1.75–1.98× feels like a straightforward representation of the structural benefits of diffusion language models.

I also looked at measurement variance. AR stayed between 12.6–12.7 tok/s across all 12 prompts with virtually no measurement noise. The tokens/forward staircase (AR 1.00 → Diffusion 1.23 → Linear Spec 1.81 → +LoRA 2.08) also reproduced stably.

Left shows tokens per forward, right shows the number of forward passes required for the same generation. Diffusion-based modes confirm more tokens per forward pass and correspondingly reduce the total number of passes.

All three modes ran on DGX Spark, and the quality of the output text was comparable to AR. Here's a summary table of the results:

tri-mode operation availability for 8B / BF16 / transformers 5.9. All three modes work, with Linear Self-Speculation being the fastest.

Visualizing Parallel Generation in Diffusion Mode

Now I know it's fast. But in what order does diffusion mode actually fill in the tokens? This was personally what I was most curious about.

The generate() function in diffusion mode internally repeats a process for each block of "confirming positions with high confidence among those still masked." So I instrumented this confirmation process to record "which positions are still masked" and "which positions were confirmed in this step" every time it's called, and tracked how one block gets filled in.

A block of block_length=32 filling in over denoising steps. The horizontal axis shows token positions within the block, the vertical axis shows steps. Dark purple indicates positions confirmed in that step, light purple indicates already confirmed, and gray indicates masked. With threshold=0.9, mainly 1 token is confirmed per step, but in the first step multiple positions are confirmed at once.

Looking at the heatmap, you can see multiple positions confirmed simultaneously in the first step of the block, followed by filling in one at a time from the highest-confidence positions. Since threshold=0.9 is a cautious setting, most steps confirm only one token at a time. Even so, the key difference from autoregressive — where position 0 must always precede 1 must always precede 2 — is that here, positions with higher confidence are filled in first.

The next figure shows in what order positions within the block were confirmed.

At which denoising step (vertical axis) each token position within the block (horizontal axis) was confirmed. You can see that positions are filled in non-sequentially from positions with high confidence, rather than strictly from left to right.

There's one thing I want to be precise about here. NVIDIA's description highlights as a feature of diffusion language models that "tokens are not permanently committed and can be revised as they go." In the scope of tracing through the 8B generate(), once a position was confirmed it did not become masked again in the remaining steps of that block — in that sense, confirmation itself was irreversible. The essential difference from autoregressive is not so much "revision" as the fact that the order of confirmation is not fixed (left-to-right) but is based on confidence, and multiple positions can be confirmed simultaneously in one step. That's the honest reading of the actual behavior. This parallelism is what leads to the reduction in forward pass count we saw earlier.

Sweeping threshold and block_length

Diffusion mode has two tuning knobs: threshold and block_length. I ran a grid to see how speed and quality change as these are varied. I tested 4 values of threshold (0.7 / 0.8 / 0.9 / 0.95) and 3 values of block_length (8 / 16 / 32), for a total of 12 combinations, running on 8 longform prompts.

Choosing a quality metric was tricky. I initially tried measuring accuracy on a multiple-choice QA benchmark (JCommonsenseQA), but with single-character answers, results barely moved regardless of how threshold was changed, making it useless as an indicator. I ultimately used the rate of repetition in the output (duplicate rate of the same 4-gram) as a proxy for quality. The technical report also mentions that diffusion mode limits generation length to avoid repetition and hallucination, and repetition is the typical failure mode when diffusion mode breaks down.

Left shows threshold vs. parallelism (tokens per forward), right shows threshold vs. repetition rate. Lowering threshold increases parallelism and makes generation faster. Repetition rate is actually lower at lower threshold values.

The results were straightforward. Lowering threshold consistently increases speed. With block_length=32, lowering threshold from 0.95 to 0.70 increases tokens per forward from 1.14 to 1.32, reducing the total number of forward passes accordingly.

What surprised me was the quality side. I expected that confirming tokens more aggressively at lower thresholds would lead to sloppier output and more repetition, but the opposite was true in practice. The repetition rate was lower at lower threshold values, and there was even a slight tendency for it to increase when threshold was raised to 0.9 or 0.95. The values themselves were low at 0.05–0.10, not at a level where text would break down.

The technical report describes threshold as "a knob that determines the tradeoff between speed and token error rate." However, at least within the scope of running longform explanation tasks with 8B BF16 in this experiment, I found no compelling reason to use a high threshold. There was also a tendency for the impact of threshold to diminish as block_length increased, so in practice it seems sensible to start with a somewhat large block_length and a low threshold. Since this will vary by task and model size, it's worth sweeping once for your own use case.

Common Pitfalls on DGX Spark

The biggest headache during testing was the incompatibility between transformers 5.9 and the bundled custom modeling. The Nemotron-Labs-Diffusion model repository includes a Python file called modeling_ministral.py, which is executed when loading with trust_remote_code=True. However, this file calls create_causal_mask() for creating attention masks using slightly older argument names (input_embeds= and cache_position=), and since transformers 5.9.0 changed these argument names, it stops with a TypeError. Diffusion mode and Linear Self-Speculation mode go through this code path (AR mode uses a different path and was unaffected).

Directly editing files in the Hugging Face cache is something I wanted to avoid, as it means tampering with what the library manages. Instead, I prepared a shim on my script side that thinly wraps the mask generation function in transformers.masking_utils to absorb the old argument names. By applying this wrap before from_pretrained() loads the custom modeling, it works without touching either the dependent code or the cache.

def apply_transformers_compat() -> None:
    import transformers.masking_utils as mu

    def _wrap(fn):
        def wrapped(*args, **kwargs):
            if "input_embeds" in kwargs and "inputs_embeds" not in kwargs:
                kwargs["inputs_embeds"] = kwargs.pop("input_embeds")
            kwargs.pop("cache_position", None)
            return fn(*args, **kwargs)
        return wrapped

    for name in ("create_causal_mask", "create_sliding_window_causal_mask"):
        setattr(mu, name, _wrap(getattr(mu, name)))

Running a freshly released model with the latest transformers tends to hit these kinds of version mismatches fairly often. Since this type of issue goes away once the custom modeling side is updated, it makes sense to treat the shim as a temporary workaround and keep it self-contained within your own script without touching dependent code.

One more note: while the model card doesn't mention a context length, looking at config.json shows max_position_embeddings of 262144 (256K). This is useful to know when planning to feed in long documents.

Summary

I loaded Nemotron-Labs-Diffusion 8B onto a DGX Spark and measured all three tri-mode modes. Here's a summary of what I found:

• All three modes — AR / Diffusion / Linear Self-Speculation — ran in the BF16 environment on DGX Spark (transformers 5.9 requires inserting one shim)
• Speed: Linear Self-Speculation was 1.75× over AR, and 1.98× with LoRA. Self-speculation, where diffusion generates the draft in parallel and autoregressive verifies it, had the biggest impact in the base configuration
• Diffusion mode fills in positions from highest-confidence first. Confirmation itself is irreversible; the essential difference from autoregressive lies in "parallelism and confirmation order" rather than "revision"
• Lowering threshold increases speed and actually decreases repetition rate. In this experiment's tasks, there was no visible reason to use a higher threshold
• NVIDIA's advertised "6× tokens per forward" and "2.7× over AR with INT4" figures include quantization and runtime optimizations; with BF16 longform generation without quantization, tokens per forward was 1.2–2.1 and speedup was just under 2×

Diffusion language models showed a clear direction: self-speculation can provide straightforward speedups without needing a separate draft model. They also seem well-suited to environments like DGX Spark where you run locally on a single node. There's still a lot I haven't tested — quantized versions, the VLM variant, behavior with longer contexts — so I hope to cover those in future articles.

Reference Links

Technical report and public resources for Nemotron-Labs-Diffusion:

https://research.nvidia.com/publication/2026-05_nemotron-labs-diffusion-tri-mode-language-model-unifying-autoregressive

https://huggingface.co/collections/nvidia/nemotron-labs-diffusion

https://huggingface.co/nvidia/Nemotron-Labs-Diffusion-8B

Related article measuring speculative decoding on DGX Spark:

https://dev.classmethod.jp/articles/dgx-spark-gemma4-mtp-multi-token-prediction-bench/