API.md

Native API contract

Entry points

• load_model(local_snapshot, device="cuda:0", dtype=torch.bfloat16, attention="sdpa") returns a Session with public .model and .processor. Missing/unexpected tensors fail. MOSS uses vendored pinned source and instance-only compatibility patches; no remote trust or external base snapshot is required after packaging.
• Request(prompt, audios) holds one user turn and 0--5 ordered local files or Waveform(samples: float[N], sample_rate: int). Waveforms must be mono. Files are mean-downmixed. URLs, fullsong directories, native control tokens and multi-turn are intentionally outside this initial contract.
• prepare(requests) -> PreparedBatch; generate(requests, profile=..., max_new_tokens=...) -> list[str], one completion per request.
• extract_audio_embeddings(batch, grad=False) -> list[AudioEmbedding] in sample/source order. extract_contextual_embeddings(batch, layers=(-1,), grad=False) -> per-source dictionaries of requested layer tensors.
• forward(batch, output_hidden_states=False, grad=False) returns native ModelOutput, logits [B,1,V], optional language hidden states [B,L,D]. .model(**batch.inputs, ...) permits advanced custom calls/full logits.

All core functions document input/output shapes in source docstrings. No method quietly averages tokens, detaches gradients requested by grad=True, sorts sources or inserts silence between different slots.

Missing markers follow the paper-reference wrapper: MF single audio uses <audio>\n{prompt} without an added label; MOSS and multiple audios use numbered labels. Explicit markers and arbitrary labels remain unchanged.

MF retains even a subsecond trailing window. A window below the conv/pool token resolution can produce zero tokens. An instance-only RoTE patch avoids upstream out-of-range indexing for that unused row; waveform/mel/masks are not truncated. A whole source producing zero tokens is rejected explicitly.

Shapes and source mappings

B=requests; A=total sources; W=total independent MF windows; L=left-padded expanded prompt width; F=mel frames; U=all valid audio tokens; U_i=one source's valid tokens.

Item	MF2601	MOSS-Music
Input IDs / attention mask	[B,L]	[B,L]
Mel / feature mask	[W,128,3000] / [W,3000]	[A,128,F_max] / lengths [A]
Source grouping	slot lengths + windows/source	group sizes [B], token lengths [A]
Raw encoder/source	[U_i,1280] before RoTE	[U_i,1280] after final encoder norm
Projected/source	[U_i,3584] after RoTE/projector	[U_i,4096] after audio adapter
DeepStack/source	none	3 raw [U_i,1280], 3 projected [U_i,4096]
Language hidden states	[B,L,3584]	[B,L,4096]
Slot token_positions	[U_i] padded-row indexes	[U_i] padded-row indexes; exclude time digits

AudioSlot.sample_index/audio_index and AudioEmbedding identify source boundaries. No output embedding padding is returned. PreparedBatch.inputs.attention_mask marks language padding; slot positions specify actual audio tokens, which may be noncontiguous because timestamps/control tokens appear inside native spans.

MOSS encodes complete-source mel, splits at 400 frames without crossing a source boundary and uses reset positions plus padding attention masks per chunk. The last short chunk is preserved. conv_chunksize=64 batches independent chunks; it is not an audio duration. Valid flat final and DeepStack streams [1,U,1280] are regrouped by explicit source metadata before injection into language layers.

MF uses separate 30-second windows with a 1,200-second cap per source, and resets RoTE per source. Sources with 2/3 windows have timeline indexes [0,1,0,1,2]. Reference SoundFile/SciPy 16-kHz preprocessing preserves subsecond tails.

File decoding uses SoundFile float32 and an arithmetic channel mean. Resampling uses scipy.signal.resample_poly with rates reduced by their greatest common divisor and SciPy's pinned default filter; it does not use soxr/librosa defaults. Resampling precedes the MF duration cap. Do not replace this with legacy MF evaluation helpers that use soxr or discard a final subsecond 30-second window.

Raw/projected extraction does not run the language model or produce a pooled retrieval vector. Contextual extraction performs a complete causal language prefill, then selects only the audio-token positions at the requested layers (layer 0 is the decoder input; negative indexes follow Python indexing). An audio token can see preceding text/audio, not a question placed after it. For explicitly question-conditioned audio states, place the question before the audio markers. No pooling or L2 normalization is applied; callers choose those research policies.

Generation, caching and hooks

paper_reference is greedy with repetition=1 and output budget 16,384 by default; generic_sampled is MF temperature .7/top_p .9 or MOSS 1/top_p .8, both top_k 50. GenerationConfig is constructed explicitly, not inherited from training's saved sampled/cache-disabled defaults. KV cache is explicitly enabled for generation, disabled for embedding prefill. Prefix caching is not implemented here.

Native contexts are 32,768 (MF) / 40,960 (MOSS). Batched output budget conservatively uses the largest padded prompt width; near the limit, call with a single request for the paper's per-request remaining-context budget. Generic sampled batches share one seed-3407 RNG stream; caller RNG is restored afterwards. This differs from vLLM's independently seeded requests. Greedy cross-kernel text identity and paper-score reproduction have not been certified.

Use session.model.named_modules() and register_forward_hook to inspect/edit forward execution. Projector hooks avoid ambiguity about DeepStack injection. MOSS's temporary language-layer hooks add DeepStack after each early layer and use prepend=True, before Transformers' persistent output recorder and normally registered user hooks. Thus ordinary hooks and returned hidden states observe post-injection boundaries on both the first and subsequent forwards. To observe pre-injection outputs, instrument the editable layer forward itself; validate the intended observation point for your research use. Remove handles in finally. The same MOSS instance must not run concurrent calls.

Ordinary autograd is supported by grad=True and direct model calls. This wrapper is not a fine-tuning implementation: checkpointed training is rejected because temporary DeepStack hooks would otherwise be removed before backward recomputation. Long-prompt full hidden-state collection can be expensive; select a small request for exploratory embedding work. Multi-turn and efficient streaming extraction of only selected hidden layers are follow-ups, not silently claimed capabilities.

Optional vLLM reference

vLLM is not a pip dependency of this native package. The paper used locally patched vLLM 0.23.1.dev0 at source commit 0fc695fc6d1d82e9a5ac6835ac8e4e1c83703665, not stock vLLM. MF requires multi-audio processor registration, independent RoTE timestamps and HF weight-name mapping; MOSS requires the compatible custom audio model, Music weights/processor identities and regenerated position buffer. The historical MOSS receipts did not seal the custom model source file itself; an installed-version string does not prove its historical byte identity. This release does not advertise stock-vLLM compatibility or ship a working vLLM server. Stage 3 compares native features with the available audited local route.