Add a sampler resource for discrete diffusion

[Hert4]

What does this PR do?

Closes #267. Adds DiffusionSamplerResource, and moves OmniVoice onto it.

The autoregressive Sampler cannot express what a diffusion step needs: it draws one token per request and returns tokens only, while a diffusion step scores every unrevealed position at once and the caller needs the log-probabilities back to decide which of them to keep. It also bakes APPLY_PENALTY as a Triton constexpr at capture, around seen-token buffers and a repetition penalty that have no meaning over a canvas rewritten every step. Hence a separate resource rather than a mode, which is what we agreed on the issue.

Contract. sample(request_ids, c_logits, u_logits, seq_lens, iterations) takes [rows, positions, vocab] packed across requests and returns (tokens, logprobs) for every position. Guidance scale and temperature are per-request, broadcast to per-position vectors, so a batch mixing them is one path rather than a Python loop. Which positions are still masked, how many to reveal, the confidence ranking and the write all stay in the model: they depend on a schedule and a layer penalty that belong to the decoder.

Three things worth flagging in the review.

Scoring moved from postprocess to forward_batched, because that is where the batch still exists.

predict_tokens_with_scoring and filter_top_k are gone from unmask.py. Keeping them would have left the parity tiers testing a copy of the maths that production no longer runs; all three now go through the resource.

An item at guidance_scale == 0 no longer builds an unconditional document at all. does_cfg existed but nothing called it, so every request paid for logits scoring discards.

How was it tested?

ruff check . passes. On a H200 with the checkpoint on a PVC, the whole file including the three GPU parity tiers:

20 passed in 19.75s

Tier A is still token-exact against the reference through the new path, which is the part that says the port is unchanged.

The live server too, since unit tests never exercise the engine building the resource from the spec, injecting it into engine_inputs.resources, or the conductor calling get_request_resource_configs:

health: 200
  greedy (default)     OK  5.845s  audio 2.44s  peak 16383   (first call, includes warmup)
  sampled temp 0.9     OK  0.644s  audio 2.49s  peak 16383
  cfg off              OK  0.557s  audio 2.43s  peak 16383
  16 steps             OK  0.378s  audio 2.41s  peak 16383

This is not a performance change, and I measured it rather than assuming. Per-exec medians of the backbone node from --log-stats, 154 requests each side, same pod, same weights, back to back:

	all	fwd	pre	post
main	33.00 ms	31.15	0.78	0.95
this branch	31.65 ms	29.50	0.80	1.03

The Python side of a step is 1.8 ms out of 33, about 5%, and batching the scoring does not move it. I had expected otherwise. A closed-loop throughput sweep agreed that the two are indistinguishable, but that harness had a counting bug so I am not quoting its numbers.

Not covered: the mixed-guidance path in _unconditional_half, where one step carries both a CFG and a non-CFG request. Unit tests cover the layout and the slicing, but no live request has gone through it; the server runs above are each uniform. Flagging rather than letting the test names imply more.

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

---

One correction to something I wrote in #249: "the gap to upstream's ~40x is almost all CUDA graphs". That was wrong. Their README says CUDA graphs are for batch 1 and that "at batch >= 4 the plain FlashInfer path is already the fastest configuration", and their own table puts graphs at 1.17x (RTF 0.0430 to 0.0367). The 40x comes from batch, not from graphs. I measured this on the hardware and will post the numbers on #267, since they change what is worth doing next.

🤖 Generated with Claude Code