mstar.model.omnivoice.components.backbone

mstar.model.omnivoice.components.backbone#

The OmniVoice backbone, run as one packed bidirectional forward per step.

A request in a step becomes two documents — the conditional canvas and its CFG counterpart — packed end to end into a single row and separated by flashinfer’s ragged attention rather than by a mask, or one document when it turned CFG off. Nothing is padded, so a batch mixing a 2-second and a 25-second utterance wastes nothing, and the head GEMM runs only at the positions whose logits are read.

The per-step kernel path is the reference’s own fastest mode, reused rather than re-derived: apply_flashinfer fuses QKV, swaps in flashinfer RoPE, RMSNorm and silu_and_mul, and keeps the whole layer in (S, H, D) so no transpose copies happen. Re-implementing that here would mean a second copy of tuned kernels whose only job is to end up numerically identical.

What M* adds on top is what the reference cannot do from inside one generate() call: the documents in a step come from different concurrent requests, at different diffusion steps, with different per-request knobs. Its own server has no batching at all (k2-fsa/OmniVoice#251).

The coupling to omnivoice.models.omnivoice_flashinfer is deliberate but narrow — four module-private names, checked at load by assert_flashinfer_api so an upstream rename fails at boot rather than at the first request.

Functions

assert_flashinfer_api()

Fail at boot if the reference's private surface moved.

build_packed_canvas(items, audio_mask_id, device)

Pack a step's requests into one ragged row.

Classes

CanvasItem(request_id, prefix_ids, ...[, ...])

One request's contribution to a step's packed batch.

OmniVoiceBackbone(model, plan_dtype)

The patched reference model, driven one packed step at a time.

PackedCanvas(packed_ids, audio_mask, ...)

One step's packed inputs, plus the index that gathers its logits.
class mstar.model.omnivoice.components.backbone.CanvasItem(request_id, prefix_ids, prefix_audio_mask, tokens, guidance_scale, flat_start=-1, flat_u_start=-1)[source]#

Bases: object

One request’s contribution to a step’s packed batch.

prefix_ids is [C, N] — style tokens, then text tokens, then the reference audio tokens when cloning. tokens is [1, C, T], the live target canvas: all MASK at step 0, fully revealed when the loop ends.

An item at guidance_scale == 0 contributes one document instead of two. Scoring throws the unconditional logits away in that case, and the unconditional document is the whole target region plus its share of the head GEMM, so building it would be close to twice the work for nothing.

Parameters:
property cond_len: int#
property does_cfg: bool#
flat_start: int = -1#
flat_u_start: int = -1#
guidance_scale: float#
prefix_audio_mask: Tensor#
prefix_ids: Tensor#
request_id: str#
property target_len: int#
tokens: Tensor#
class mstar.model.omnivoice.components.backbone.OmniVoiceBackbone(model, plan_dtype)[source]#

Bases: Module

The patched reference model, driven one packed step at a time.

Holding the reference instance rather than borrowing its submodules is what keeps the embedding merge, the audio head reshape and the fused attention path byte-for-byte the upstream ones; every line of per-step math that could drift is theirs.

Parameters:
forward(canvas)[source]#

One packed forward; returns [1, C, 2 * flat_target_total, V].

Only the logits-consuming positions go through the 1024 -> 8200 head GEMM and the float32 upcast. Running them over the full packed length would be several times the work, most of it on prefix positions whose logits are never read.

Parameters:

canvas (PackedCanvas)

Return type:

Tensor

class mstar.model.omnivoice.components.backbone.PackedCanvas(packed_ids, audio_mask, position_ids, doc_lens, tgt_index, flat_target_total, flat_uncond_total, items)[source]#

Bases: object

One step’s packed inputs, plus the index that gathers its logits.

Parameters:
audio_mask: Tensor#
doc_lens: list[int]#
flat_target_total: int#
flat_uncond_total: int#
items: list[CanvasItem]#
packed_ids: Tensor#
position_ids: Tensor#
slice_logits(logits, item)[source]#

Conditional and unconditional target logits for one item.

Laid out as every item’s conditional target block followed by the unconditional blocks of the items that asked for CFG, which is the gather order tgt_index was built in.

The unconditional half is None for an item at guidance_scale == 0: it has no unconditional document, because scoring would discard those logits anyway.

Parameters:
Return type:

tuple[Tensor, Tensor | None]

tgt_index: Tensor#
mstar.model.omnivoice.components.backbone.assert_flashinfer_api()[source]#

Fail at boot if the reference’s private surface moved.

These names are private to omnivoice. Pinning them here means an upgrade that renames one produces a clear error on startup instead of a silent fallback or a first-request crash.

mstar.model.omnivoice.components.backbone.build_packed_canvas(items, audio_mask_id, device)[source]#

Pack a step’s requests into one ragged row.

Documents are laid out [cond_0, uncond_0, cond_1, uncond_1, ...] and kept apart by the attention plan’s doc_lens, not by a mask — so there are no pad positions at all, and therefore none of the fully-masked-row hazard a dense padded batch has to guard against.

An item at guidance_scale == 0 contributes only its conditional document, so a batch of those is close to half the packed length and half the head GEMM. Mixing the two kinds in one step is fine: the layout is driven per item.

position_ids restart at 0 in every document. Without that, RoPE would read a later document’s canvas as a continuation of the previous request’s.

Parameters:
Return type:

PackedCanvas