mstar.model.cosyvoice3.components.sampling

mstar.model.cosyvoice3.components.sampling#

Repetition-aware sampling (RAS) for the CosyVoice3 speech-token LM.

The reference (cosyvoice.utils.common.ras_sampling) draws a token from the nucleus: ids in descending probability while the mass before them is under top_p and fewer than top_k are taken. If that token already appears >= window * tau times among the last window generated tokens, it is masked and a new token is drawn from the whole distribution. Without the redraw, a nucleus that collapses onto a silence token never reaches EOS.

Here both draws are Gumbel-max over noise from a counter-based hash of (seed, step, id): one fixed-shape tensor program per batch that a CUDA graph can capture, the same tokens whatever the batch around a request, and a request’s seed reproduces its tokens. The draws follow the reference’s distributions; they cannot follow its random stream, which is the global torch.multinomial state.

Functions

hashed_uniform(seed, step, salt, n)

[B, n] uniforms in (0, 1), a function of (seed[b], step[b], salt, j) only.

ras_sample(logits, *, seed, step, history, ...)

One token per row.
mstar.model.cosyvoice3.components.sampling.hashed_uniform(seed, step, salt, n)[source]#

[B, n] uniforms in (0, 1), a function of (seed[b], step[b], salt, j) only.

Parameters:
Return type:

Tensor

mstar.model.cosyvoice3.components.sampling.ras_sample(logits, *, seed, step, history, temperature, top_p, top_k, window, repeat_threshold, blocked)[source]#

One token per row.

logits [B, V] float32. Per row: seed/step (int64) key the noise; history [B, W] holds the last W generated tokens (-1 where fewer exist); temperature (0 = greedy), top_p, top_k (<= MAX_TOP_K), window (the row’s RAS window, <= W: repeats count in its last window history slots only) and repeat_threshold (window * tau; <= 0 disables the redraw) are [B]; blocked [B, V] bool marks ids that may not be drawn this step (the reference’s minimum-length rule).

Parameters:
Return type:

Tensor