Sampler support for discrete diffusion

[NSagan271]

Currently, the Sampler resource only supports autoregressive-type generation: one token per request, and only returns the sampled tokens. For discrete diffusion, we need to be able to sample multiple (an arbitrary number) of tokens per request, and also return the log probabilities. This would probably make sense as a new Resource, but is also could be incorporated into the existing Sampler.

[Hert4]

#249 is a use case for this and currently sidesteps the sampler entirely: get_node_resources() returns [], and the token choice sits in mstar/model/omnivoice/components/unmask.py.

One note on the API shape: "multiple tokens per request" does not quite fit masked diffusion. OmniVoice does not ask for N tokens. Each step it scores the whole masked canvas, then ranks by confidence and commits only the top k:

pred_tokens = log_probs.argmax(dim=-1)    # [1, C, T]  every cell
confidence  = log_probs.max(dim=-1)[0]    # [1, C, T]  drives the ranking
_, topk_idx = torch.topk(scores.flatten(), reveal_count)

The winners are only known after the scores exist, so the contract wants to be "sample every position handed in, return (tokens, logprobs)" and let the caller pick. Logits are [C, T, V] with T ragged across requests, not [1, V].

Two things that are sampling but are not in Sampler today: CFG combine (two logit tensors in, one out) and top-ratio vocab filtering (ceil(0.1 * V)) with Gumbel. Both are per request config the way temperature is, and both are elementwise over [C, T, V], which is the expensive part of the step. The confidence ranking and the scatter should stay in the model: they use a layer penalty and a position-level temperature that are decoder specific.

Separate resource or a mode on Sampler? Repetition penalty and seen-token masks are dead weight here, and SamplerResource bakes APPLY_PENALTY as a Triton constexpr at capture, so I lean separate. Happy to prototype against #249 either way.

[NSagan271]

I also think that a separate resource is the way to go. @Hert4 I'll assign you to this issue, happy to help out or review.

[Hert4]

The resource is up as #290. While building it I measured the three things this issue was supposed to unlock, on a H200 with the real checkpoint, and all three came back negative. Posting because it changes what is worth doing after this.

CUDA graphs buy nothing at batch 8.

eager   8.54 ms
replay  9.35 ms

Do not read that as a 9% regression: in the same run, batch 8 measured 9.59 ms in one table and 8.54 ms in another, so the noise floor here is about 12% and the honest reading is "no change".

Upstream's README already says this, which I had missed: "CUDA graphs are recommended for single-stream (batch=1) usage, where kernel-launch overhead dominates; at batch >= 4 the plain FlashInfer path is already the fastest configuration." Their own table puts graphs at 1.17x and only at batch 1 (RTF 0.0430 to 0.0367). So the ~40x realtime they advertise comes from batch, not from graphs: batch 8 is RTF 0.0115, which is 87x. I claimed the opposite in #249 and was wrong.

Batch scaling is already near its knee. target_len 100, GPU time per step:

batch	gpu_ms	ms per request	vs batch 1
1	7.01	7.007	1.00x
4	6.76	1.689	0.96x
8	9.59	1.198	1.37x
16	16.31	1.020	2.33x
32	30.68	0.959	4.38x
64	60.51	0.946	8.64x

Batch 4 is free, batch 8 nearly so, and past 16 it is linear. Raising max_batch_size from 16 to 64 is worth 7%.

Porting the backbone onto mstar's components has nothing to cut. Kernel profile at batch 8:

kernel	calls	us/call	share
nvjet GEMM (two shapes)	560	33-37	48.8%
flashinfer norm	560	6.6-22.8	20.7%
elementwise	280	16.4	11.5%
flashinfer act_and_mul	140	24.9	8.7%
flashinfer prefill attention	140	14.2	5.0%
flashinfer rotary	140	10.4	3.6%

Attention is 5%. Everything that is not GEMM is already on flashinfer's fused kernels, which is what a port would have replaced them with. There is a gap to roofline, roughly 4x at this batch, but it sits inside the GEMM shapes, not in the engine layer, so it is a kernel problem rather than something the port would fix.

For completeness on where the time in a served step goes, from --log-stats, 154 requests: the Python half of a step is 1.8 ms out of 33, about 5%, and the scheduler does form real batches (measured at the seam: conc 8 gives mean batch 7.88, conc 16 gives 12.19). Neither is the bottleneck.

So my read is: the resource in #290 is worth having for the contract, and the follow-up port and capture work is not worth doing on this model. Happy to be told otherwise, and glad to re-run any of this if a configuration looks wrong.