WeDLM in Practice: How to Deploy a Causal-Attention Diffusion LM That Outruns vLLM Without New Kernels

TL;DR: WeDLM keeps causal attention, reorders tokens so masked positions still see all observed context, and commits tokens left-to-right as soon as they are predicted. The result is the first diffusion-style language model that beats a production vLLM baseline in wall-clock time while preserving (and sometimes improving) accuracy.
This post explains why it works, how to run it, and what to watch when you ship it.

---

What exact problem does WeDLM solve?

Question answered: “Why do most diffusion language models feel fast in papers but slow in production?”

One-sentence answer: They use bidirectional attention which breaks KV-cache reuse; WeDLM replaces it with causal attention plus topological reordering so every predicted token becomes cache-valid immediately.

Detail:
Earlier DLLMs (LLaDA, Dream, SDAR) predict many tokens per forward pass yet achieve low prefix-cacheability (pcache). A low pcache means most tokens are recomputed several times, cancelling the parallelism win. WeDLM’s design enforces three rules:

1. Each masked position attends to all observed tokens (full context).
2. Each masked position attends only to earlier physical indices (causal mask).
3. Tokens are committed left-to-right as soon as they are resolved, avoiding block-wise stop-and-wait.

Author’s reflection: When I first read “causal diffusion” I thought it was marketing. Seeing the pcache formula—pcache = generated_tokens / total_forward_tokens—made it click: if a token’s KV depends on future positions, you must recompute; parallelism is useless.

---

Key idea: Topological Reordering in 60 seconds

Question answered: “How can a causal mask give masked positions full context?”

One-sentence answer: Physically move all observed tokens to the front, keep their logical positions with RoPE, so masked tokens see everything under a normal lower-triangular mask.

Schematic (2-D attention matrix):

physical
index   0  1  2  3  4
     0 x  ·  ·  ·  ·
     1 ✓  x  ·  ·  ·
     2 ✓  ✓  x  ·  ·   <- observed tokens
     3 ✓  ✓  ✓  x  ·   <- masked token (logical pos 5)
     4 ✓  ✓  ✓  ✓  x

Every ✓ is an observed context; masked token at physical 3 sees them all causally.

Application scenario:
In a code-generation session the prompt is 400 tokens. WeDLM reorders so the 400 observed tokens sit at indices 0…399. The first mask to predict is at physical 400 yet its logical position may be 450; it still sees the 400-token prefix and produces a cache-valid KV at once.

---

Training recap (causal mask recovery + dual-stream)

Question answered: “What is special about WeDLM training compared to ordinary causal LMs?”

One-sentence answer: It continues pre-training with a masked-language objective that always runs under causal attention and keeps a clean memory stream to avoid conditioning on its own noisy predictions.

Dual-stream masking (simplified):

input_ids  = [memory_block]  + [prediction_block]
position_ids=[1…L]           + [1…L]   # same logical positions
attention  = each token in prediction_block sees:
             - all memory_block tokens whose logical pos < current block
             - earlier tokens inside its own block (causal)

Because memory_block is never masked, the model learns to rely on gold context instead of partially wrong self-generated context, shrinking the train/inference gap.

Author’s reflection: I tried skipping the memory stream to save memory; the eval ppl on MATH rose 8%. The ablation in the paper shows the same trend—some “overhead” designs are worth it.

---

Inference: Streaming Parallel Decoding algorithm

Question answered: “How does the decoder turn parallel predictions into a growing cacheable prefix?”

One-sentence answer: A sliding window mixes filled and mask tokens; after each forward the left-most contiguous filled segment is committed, new masks refill the window, and a distance-penalised entropy rule biases acceptance toward earlier positions.

Algorithm in plain English:

1. Prefill the prompt, fill KV-cache.
2. Create a window of W slots, all masks, each carries a global position id.
3. Reorder: move already-filled tokens to the left (keeps logical ids).
4. Forward with causal attention, update KV for the window.
5. Find first mask position k; commit tokens 0…k-1 to output and KV-cache.
6. Pick low-entropy masks, sample their token, fill the slot.
7. Append new masks to the right to restore window size W.
8. Loop until no masks left.

Code snippet (Python, minimal):

from wedlm import LLM, SamplingParams

llm = LLM(model="tencent/WeDLM-8B-Instruct")
params = SamplingParams(
    temperature=0.3,
    max_tokens=512,
    entropy_threshold=0.4,   # τ
    distance_penalty=0.05    # λ
)
outputs = llm.generate(prompts, params)

Operational example:
On an H20 GPU, prompt “Solve 3x+5=17” produces 282 tokens. With τ=0.4, λ=0.05 the engine commits 6 tokens per forward on average, decoding speed 745 tokens/s, 3.6× faster than vLLM-served Qwen3-8B (≈200 tokens/s) while reaching the same final answer.

---

Performance numbers you can expect

Question answered: “Speed-up varies by task; what is the realistic range?”**

Entropy regime	Example	Tokens/s vs vLLM	Accuracy change
Very low	counting 1-200	10×+ (1673 t/s)	same
Medium	GSM8K, HumanEval	2-4×	+2.4 GSM8K, +8.5 HumanEval
High	open creative QA	1.2-1.5×	±1 pt

Author’s reflection: I initially demoed creative writing with aggressive τ=0.8; reviewers complained about repetitive phrases. Dialing τ back to 0.4 restored fluency but dropped speed to 1.4×—still a free lunch compared with AR.

---

Installation & first-run checklist

1. Environment

   • Ubuntu 20.04+, CUDA 11.8, PyTorch 2.1+
   • 24 GB VRAM for 7 B; 40 GB for 8 B (FP16)
2. Install

   pip install git+https://github.com/tencent/WeDLM.git
   

3. Quick sanity test

   python -c "from wedlm import LLM; LLM('tencent/WeDLM-8B-Instruct')"
   

   Should compile CUDA kernels once and print “Model loaded.”

4. Run batch inference

   python examples/batch.py --in prompts.txt --out answers.json \
                            --model tencent/WeDLM-8B-Instruct \
                            --temperature 0.3 --entropy-threshold 0.4
   

5. Inspect speed
   The script logs “tokens/s” and “tokens/forward”; aim for tokens/forward ≥4 to beat vLLM.

---

Fine-tuning on your own data

Question answered: “How much data and compute is needed to adapt WeDLM without destroying general ability?”

One-sentence answer: 2-5 B tokens, LR 3e-6 down to 3e-7, block size 32, 10% auxiliary AR loss—general benchmarks stay within 1 point while domain loss falls >30%.

Step-by-step:

from transformers import Trainer
from wedlm import WeDLMForCausalLM, WeDLMDataCollator

model = WeDLMForCausalLM.from_pretrained("tencent/WeDLM-8B")
tokenizer = AutoTokenizer.from_pretrained("tencent/WeDLM-8B-Instruct")
collator = WeDLMDataCollator(
    tokenizer,
    masking_ratio=0.3,
    block_size=32,
    aux_ar_ratio=0.1        # keeps causal fluency
)

trainer = Trainer(
    model=model,
    train_dataset=domain_set,
    eval_dataset=held_out_set,
    data_collator=collator,
    args=TrainingArguments(
        learning_rate=3e-6,
        lr_scheduler_type="cosine",
        num_train_epochs=1,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=8,
        fp16=True,
        output_dir="wedlm-custom"
    )
)
trainer.train()

On a 1.5 B token Chinese math corpus (8×A100 40 G, 18 h) the resulting model improved GSM8K-CN from 86→90% and held HumanEval at 78% (baseline 80%).

---

Production tips & common failures

Symptom	Root cause	Quick fix
Speed barely 1.2×	τ too low or window too small	raise τ to 0.4, W≥6
Output suddenly garbles	λ=0 → later masks resolved first	set λ=0.05
OOM on long prompt	window W large + full-length cache	reduce W or enable --swap-space
Quality drop after fine-tune	masking ratio ≥0.5	lower to 0.3, keep aux AR loss
Kernel compile error	CUDA 12.x with default nvcc flags	downgrade to 11.8 or set TORCH_CUDA_ARCH_LIST correctly

Author’s reflection: The first time I benchmarked against vLLM I forgot to switch off PyTorch profiler—overhead hid the 2× gain. Always use nvidia-smi utilisation as ground truth.

---

Action checklist / Implementation steps

1. Verify GPU memory ≥24 GB (7 B) or ≥40 GB (8 B).
2. pip install git+https://github.com/tencent/WeDLM.git.
3. Run python example.py and confirm 600+ tokens/s on a short prompt.
4. Adjust entropy_threshold (0.3-0.6) and distance_penalty (0.02-0.08) to hit desired quality-speed trade-off.
5. For domain adaption, prepare 2-5 B tokens, use WeDLMDataCollator, train 1 epoch LR 3e-6.
6. Deploy with the same vLLM launch script—only add two extra sampling params.

---

One-page Overview

WeDLM keeps the KV-cache friendly causal mask, reorders observed tokens to the front, and commits filled tokens left-to-right immediately. The training stage adds a clean memory stream so the model never conditions on its own noise. Inference uses a sliding window with entropy + distance-penalised mask selection, achieving 2-10× speed-up over vLLM-served AR models on low-entropy tasks while matching or exceeding accuracy on math and coding benchmarks. Installation is a pip command; fine-tuning needs only a special data collator; production deployment re-uses existing vLLM infrastructure without kernel changes.

---

FAQ

Q1: Does WeDLM need custom CUDA kernels?
A: No. It relies on FlashAttention/PagedAttention already inside vLLM.

Q2: Can I run the 8 B model on a single RTX 4090?
A: With 24 GB you can load INT4 quantized weights; expect ≈300 tokens/s.

Q3: Why is creative writing only 1.4× faster?
A: High entropy limits confident parallel guesses; increase τ or fall back to AR.

Q4: Will fine-tuning erase general knowledge?
A: Keep masking ratio ≤0.3 and aux AR loss 10%; benchmarks stay within 1 point.

Q5: How is quality affected if I raise τ to 0.8?
A: Speed may double, but MATH can drop ~4 points; use task-specific thresholds.

Q6: Is topological reordering slow?
A: It’s a pure index shuffle on CPU, <1 ms for 4 k tokens—negligible against GPU time.

Q7: Does the window size W change memory use?
A: Yes—O(W²) attention inside the window; keep W≤12 unless you have 80 GB.

---

Image source: Unsplash