Model Layer

OASR's model layer is architecture-agnostic in the vLLM / SGLang sense: the engine touches a model only through a small base-class surface, and a new architecture plugs in by subclassing and registering. No engine edits.

This document covers the contracts, the registry, and the seven built-in architectures. The registry seams themselves are in architecture.md; checkpoint conversion is in checkpoints.md; parity and WER results are in .artifacts/model_validation.md.

The engine-facing surface

oasr/models/base.py defines what the engine may assume.

Member	Meaning
encode_offline / encode_chunk_paged	Encoder only → raw hidden states
forward_offline / forward_chunk_paged	Fused encoder + head → log-probs (the CTC fast path CUDA-graph capture preserves)
forward_offline_packed	Gapless varlen offline forward, used when sequence packing is enabled
cache_spec	CacheSpec describing the streaming cache; None when streaming_kind == "none"
decode_type / default_decode_type / capabilities	Which decode families this checkpoint can serve (DecodeType in base.py)
encoder.streaming_kind	"paged" / "slot" / "stateful" / "none"
encoder.subsampling_rate, encoder.right_context	Streaming geometry inputs
load_weights(state_dict) → LoadReport	Weight loading, reporting mapped / dropped / missing

BaseEncoder.n_kv_head and head_dim are paged-streaming only — they are non-abstract with a self-explaining raise, so an offline-only encoder implements neither.

Three optional encoder declarations

All three default to "nothing declared", so an encoder that needs none of them is unaffected by their existence. Together they are what makes a new streaming cache data rather than a new manager.

Declaration	What it says
streaming_state_specs	Extra fixed-extent per-stream tensors → CacheSpec.stream_states. See cache_manager.md §10.
fixed_attention_window	A trained attention span. The engine pre-fills the paged K/V window so cache_seqlens is uniform across the cohort — the precondition for one shared relative-position table, since a Transformer-XL table's distances are cache + i - j and a per-row table would run the positional projection per row. It also makes cache_t1 constant, so the graph cache captures one graph per batch size instead of one per (B, cache_t1 bucket).
streaming_geometry(chunk_size)	A front-end the generic (chunk_size - 1) * sub + right_context + 1 window formula does not describe — a cached causal subsampling consumes exactly chunk_size * sub frames with no lookahead. Also the place to refuse a chunk size the encoder cannot serve; the backend calls it once at construction.

A fourth optional declaration is model-level: quantization_layout() returns a QuantLayout — the SmoothGroups (which norm feeds which GEMMs, and how a SmoothQuant scale is absorbed) and the layers kept at full precision. Pure data; the default declares nothing smoothable, and online round-to-nearest quantization works without it. See quantization.md. A model that infers its dtype from a parameter uses compute_dtype, not next(parameters()).dtype, which reads int8 storage on a quantized model.

streaming_kind must describe what this config's weights can actually do, not what the class implements. Returning "none" makes streaming service mode fail at construction with a clear message and makes cache_spec None, so no paged pool is allocated. An encoder that over-claims builds an engine that raises on its first request instead.

Capabilities

oasr/models/interfaces.py::CAPABILITIES is a declarative table: for each decode family, the dotted model-attribute paths it requires plus a one-line why.

build_decode_strategy calls require_capability once, so a checkpoint advertising a capability it cannot serve fails at engine construction naming the missing members. tests/models/test_contract.py validates the table against every registered architecture, built tiny on CPU.

This is the answer to "what must a model implement to support family X".

Registry and loading

oasr/models/registry.py provides register_model, build_model_from_checkpoint, load_checkpoint_bundle, and the CheckpointConverter protocol.

• Built-ins come from one _BUILTIN_PACKAGES list.
• Out-of-tree architectures arrive through oasr.models entry points, so adding one needs no registry edit. A broken plugin warns and is skipped rather than taking the built-ins down.
• oasr/models/loaders.py::from_pretrained (re-exported as oasr.from_pretrained) resolves a local directory or a Hugging Face Hub id. load_pretrained additionally returns the tokenizer / feature / decoding specs plus a LoadReport, and is what ASREngine uses.

Config deserialization

Model configs inherit one type-driven from_dict (oasr.models.base.coerce_config): it restores Tuple fields, recurses into nested dataclasses, and unwraps Optional. A config overrides _from_dict_overrides only for a flat encoder dict or a polymorphic field — e.g. TransducerModelConfig.encoder, whose class comes from the sibling encoder_type.

Weight-load accounting

load_weights returns a LoadReport{mapped, dropped, missing}, built via LoadReport.build (which folds unexpected into dropped and uses set membership). The registry then warns on dropped weights according to the converter's declarations:

Converter declaration	Effect
expected_unused_prefixes	Silent — e.g. icefall simple_*_proj
capability_drop_hints	Named warning — e.g. the U2++ decoder.* rescoring branch
(none): a model tensor left missing	Load error, carrying the drop hints that fired

Autoregressive decoder contracts

oasr/models/decoders/base.py holds BaseDecoder, PredictionNetwork, Joiner and TransducerPredictor — the transducer / AED / LLM extension points.

Two batched decoder shapes exist, and both are driven from oasr/engine/decode/:

• Frame-synchronous (transducer): init_state + step.
• Label-synchronous AR (Whisper, speech-LLM): the incremental surface prefill / step / select / merge. select(state, idx) is an index_select, so repeated indices are legal — which is what lets beam search both expand a prefilled batch into a B × k grid and reorder it onto each slot's parent, with no new model method. merge(a, b) joins two prefilled states into one forward, guarded by can_merge; it works because the KV is indexed per row rather than off a shared generation offset (see decoding.md § Groups merge).

Two class-level declarations go with it, both answered by the decoder rather than probed. supports_paged_kv says prefill accepts a kv_manager and will page its decoder KV out of that pool — for an AED, the cross-attention over the encoder window as well (build_cross_kv); supports_step_graphs says a step can be recorded into a CUDA graph — which requires that everything it reads either lives at a process-stable address or is small enough to copy into a static buffer per step. Whisper meets it only paged: dense, its cross-attention K/V is allocated per prefill and a graph would bake in a pointer to the batch it was captured on, so the graph cache refuses a dense state whatever the class declares.

supports_step_graphs also decides storage under the default kv_storage="auto": paging is what makes a step capturable, and a cost otherwise, so a decoder that declares it gets a paged pool and one that does not keeps its dense capacity buffer. Declare it from what this config's weights can do, not from what the class implements — the failure mode is the same as over-claiming streaming_kind, an engine that builds and then misbehaves instead of one that refuses at construction.

An AR decoder may optionally add cross_attention(enc_out, token_ids, heads, max_frames), returning the attention probabilities of a declared (layer, head) set plus the pass's logits. It is what backs word timestamps for the aed family (DTW over the attention — see decoding.md § Word timings), and it is optional in the real sense: a decoder without it makes AedDecodeStrategy.word_timing_modes report (), and the engine then refuses a word_timestamps request instead of answering it with the field missing. It is deliberately a second teacher-forced forward rather than a hook on generation, so the cost lands on the request that asked rather than on every decode step of every request.

The transducer predictor state is opaque

TransducerPredictor declares init_state / predict / advance(state, tokens, emit) / stack_states / unstack_states. One greedy loop therefore serves both:

• a stateless label-window predictor (icefall: recomputable from the last k labels), and
• a recurrent one (NeMo's 2-layer LSTM: not).

Two consequences:

1. A recurrent predictor's start state is not zeros. NeMo runs the LSTM once on the blank embedding (the zero row, via padding_idx) from a zero hidden state, and that response is what the first frame's joint sees.
2. beam_size > 1 is refused at construction for such a predictor (label_window_state = False), because modified beam search keeps the beam's states in one (B, k, ctx) buffer and gather-reorders them onto their parents — which only expresses a label window.

oasr/models/decoders/transformer_decoder.py holds the WeNet/ESPnet-compatible TransformerDecoder / BiTransformerDecoder (state-dict keys mirror WeNet's embed.0.weight / decoders.N.self_attn.linear_q.* layout, so U2++ decoder.* weights load 1:1), the add_sos_eos / reverse_pad_list helpers, and an incremental forward_one_step. It is built on oasr.layers: WeNet's linear_q / linear_k / linear_v / linear_out names over the shared Attention core, and feed_forward as a FeedForward(names=("w_1", "w_2")) whose ReLU folds into the GEMM epilogue. Cross-attention takes memory_lens rather than a materialized memory mask — the form the fused kernel enforces directly.

Heads

oasr/models/heads/ctc.py::CTCHead wraps oasr.layers.ctc.CtcProjection, the fused log_softmax(x @ Wᵀ + b) via oasr.gemm_log_softmax.

Output projections are allocated at a GEMM-aligned width. CUTLASS's alignment-8 iterators mean an unaligned head cannot reach the kernel at all, so oasr.models.base.align_out_features rounds the width up and pad_output_projection widens the checkpoint on load, biasing the padding rows to PAD_LOGIT so they can never win an argmax. test_no_architecture_needs_an_unaligned_gemm keeps this closed.

The seven built-in architectures

Registry key	Package	Decode families	Streaming	Source format
conformer	models/conformer/	ctc, ctc_aed_rescoring	paged	WeNet
zipformer	models/zipformer/	ctc	slot (config-dependent)	icefall
transducer	models/transducer/	transducer	follows the encoder	icefall (explicit only)
nemotron	models/nemotron/	transducer	paged	Hugging Face
whisper	models/whisper/	aed	none	Hugging Face
paraformer	models/paraformer/	paraformer	none	FunASR
speech_llm	models/speech_llm/	llm	none	Hugging Face

list_models() prints this table's key column as the registry sees it at runtime, including any out-of-tree architecture that arrived through the entry point group.

Every architecture is built from oasr.layers, not from bare nn.Linear / nn.LayerNorm / nn.Embedding. tests/models/test_layer_waist.py enforces it, and its tiny-config table is keyed off list_models() so a newly registered architecture with no entry fails rather than going uncovered.

Conformer (WeNet U2/U2++)

model.py, config.py, convert.py::WenetConverter. streaming_kind="paged"; supports sequence packing.

A U2/U2++ directory whose train.yaml declares a (bi)transformer decoder gets the AED branch loaded as self.decoder (decoder.left_decoder.* keys map 1:1; plain-transformer decoder.decoders.* keys are remapped into left_decoder), giving capabilities={"ctc", "ctc_aed_rescoring"} with default_decode_type="ctc". Rescoring is opt-in via EngineConfig.decode_method. The decoder config — including sos/eos (raw vocab − 1) and the trained reverse_weight — lives on ConformerModelConfig.decoder.

The converter builds only a WeNet asr_model or k2_model with a ConformerEncoder, because OASR decodes through the CTC head those two train. It refuses, at load: - another encoder or model class, and model_conf.ctc_weight: 0; - an encoder_conf value this Conformer does not implement, such as a pos_enc_layer_type other than rel_pos, a selfattention_layer_type other than rel_selfattn, RMSNorm, concat_after, or a static chunk without a dynamic one; - an activation other than swish or relu. WeNet's gelu is exact erf, while the convolution module's fused one is the tanh approximation; - an encoder_conf key it does not know.

A decoder_conf it cannot build drops the AED branch with a warning instead. The tables of keys are in convert.py.

Zipformer (icefall CTC)

model.py, encoder.py, subsampling.py, scaling.py, config.py, convert.py::IcefallConverter (which infers the config from checkpoint shapes).

streaming_kind is "slot" only when the config is streaming-capable (causal=True and chunk_size > 0), otherwise "none" — so a non-causal release such as zipformer-large-cr-ctc is refused in streaming service mode at engine construction rather than raising on its first request. The converter tells a causal release from its conv modules (icefall's ChunkCausalDepthwiseConv1d has causal_conv + chunkwise_conv where a non-causal one has depthwise_conv.weight) and gives it the decode geometry the releases are published at, chunk 16 / left context 128.

Streaming follows icefall's recipe exactly, pinned against the release's own TorchScript streaming export (TestCausalRelease):

• windows of 2 * chunk + 13 input frames at a stride of 2 * chunk (streaming_window_frames / streaming_chunk_frames) — the embed turns T frames into (T - 7) // 2 - 3, so that is what yields exactly chunk embed frames, the unit the downsampled stacks need whole; a short final window is padded with icefall's LOG_EPS (streaming_pad_value);
• a processed-length mask: the state's last tensor counts the embed frames a stream has been through, and the part of the left-context cache that is still initial zeros is masked.

The whole per-stream cache — the embed's cached frames, six tensors per layer (key, nonlin-attention, two value caches, two conv tails), the processed-length counter — is declared as slot_state_specs, one StreamStateSpec per tensor with its icefall batch axis as the slot axis. The engine's slot runtime owns it (see cache_manager.md § 10) and graph-captures the whole step; the list API (get_streaming_init_states / streaming_forward, plus stack_streaming_states / unstack_streaming_states) stays for the "stateful" runtime, which is the parity oracle. The NonlinAttention gate uses the standalone Tanh waist module; its gate is one chunk of a three-way projection output, so a whole-GEMM epilogue cannot represent the operation.

The attention is decomposed on purpose and cannot be fused. RelPositionMultiheadAttentionWeights materializes the probabilities once and both SelfAttention and NonlinAttention consume the same tensor, so FMHA — which never materializes them — has no way to express the sharing. The five products per layer therefore go through oasr.bmm's general lane: 4-D permuted views, an operand broadcast over the request batch, head dims of 32 / 4 / 12 and an always-odd relative-position extent. See kernels.md § "The BMM general lane"; the remaining fusion (score add, relative shift, masks, softmax) is the open item.

Transducer (icefall RNNT)

model.py::TransducerModel, decoder.py (stateless predictor), joiner.py, config.py, convert.py::IcefallTransducerConverter. encoder_type ∈ {"conformer", "zipformer"} selects the acoustic front-end, and streaming follows the encoder (paged, or slot for a causal Zipformer).

Not auto-detected. icefall directories sniff as zipformer (CTC) and hybrid checkpoints carry both branches, so load with from_pretrained(dir, architecture="transducer"). The config is shape-inferred from decoder.* / joiner.*; simple_*_proj is declared-dropped and ctc_output.* is a named capability hint.

The predictor starts from icefall's window. Before the first label the state is [-1] * (context_size - 1) + [blank] (StatelessDecoder.init_state), and a negative label embeds to zero in StatelessDecoder.forward and in both fused decode kernels. That is the input icefall trains on: [blank] + y, with the embeddings left-padded by zero vectors. It is also where icefall's own decoders and sherpa-onnx start. A window of blanks is a different input wherever blank's embedding row is not zero. The Zipformer recipes declare no padding_idx, so their blank row is learned, and a blank-filled start window changes the first emitted tokens.

The predictor's post-embedding/conv ReLU and the joiner's tanh-over-projection-sum use standalone activation modules. Neither can use a GEMM epilogue: the former also covers the no-conv context-size-one case, and the latter precedes the vocabulary projection.

Nemotron ASR (FastConformer + RNN-T)

encoder.py, subsampling.py, predictor.py, convert.py::HFNemotronConverter (auto-detects config.json: model_type=nemotron3_5_asr). Offline and cache-aware streaming. Reference release: nvidia/nemotron-3.5-asr-streaming-0.6b.

Every subsampling-stage pointwise Conv2D fuses ReLU into its existing epilogue. The RNN-T joint still uses standalone Relu because activation applies to the sum of encoder and decoder projections before the vocabulary head.

• 24-layer macaron Conformer over causal depthwise-separable 8× Conv2d subsampling, kept in NHWC so the projection's input columns are permuted once at load.
• Transformer-XL relative position (rel_shift). WeNet's convention needs no shift, which is why oasr.layers.RelPositionMultiHeadedAttention has none and this is a separate module.
• A chunked_limited attention mask that applies offline too — it is trained, not a streaming-only device: a query sees its own chunk of num_lookahead_tokens + 1 frames plus (sliding_window - 1) // chunk earlier ones.
• Offline attention runs in query blocks. The same window bounds what the offline pass computes: each block of OFFLINE_BLOCK_FRAMES queries (64, rounded up to whole chunks) is scored against its own frames plus the trained history, i.e. one streaming step with cache_t1 = left, run for every block at once. Scores and the relative table then grow with T, not T². Upstream builds the whole T x (2T - 1) table and masks almost all of it, which ran long audio out of memory. An utterance that blocks would not cut to half the dense scores (T <= 2 * keys, 240 frames on the release) stays one dense block, because copying K and V into overlapping blocks costs more than it saves there.
• predictor.py — 2-layer LSTM predictor on oasr.layers.LSTM's fused recurrent kernels + additive joint head(relu(enc_proj + dec_proj)) + the language-prompt projector, with no residual — its output replaces the hidden state. Upstream concatenates a 128-wide one-hot onto every encoder frame; here the one-hot's columns of the first projection are folded with its bias into a per-prompt table at load (W_h h + (W_p[:, id] + b)), so each row adds its own prompt's row of it. That drops the concat and shrinks the GEMM's K from 1152 to 1024, and is within one rounding of the concat rather than bit-identical (the GEMM output is rounded before the bias is added).
• The prompt is per request. DecodingOptions.language (a bare ISO-639 code) resolves against the checkpoint's own processor_config.json table: the code if it is listed (en → 0), else the table's xx-XX spelling (zh → zh-ZH → 4), else the one slot its locales share (vi → vi-VN → 33). A code the table does not know, or whose locales disagree, is refused at admission. No language means default_prompt_id (auto, 101: the model's own language ID). The model declares it through BaseAsrModel.encoder_options, so the transducer family admits language here and refuses it on a checkpoint that has no prompt. Offline, each batch row carries its id into encode_offline (encoder_conditioning); a batch where nobody names a language takes the plain call. Streaming, the id is a per-stream persistent slot state (StreamStateSpec.persistent), written at admission, read inside the chunk forward so captured graphs gather each replay's rows, and kept across a VAD turn reset, which rewinds what the encoder heard but not what the request asked for. num_lookahead_tokens is still fixed at construction: it changes the streaming geometry, not a per-row input.
• TokenizerSpec(kind="huggingface", special_ids=[blank]) — blank 13087 sits past the tokenizer's 13087 pieces, so it must be filtered before the backend sees it.

Streaming carries four kinds of state, each on a declared axis rather than as a special case: the three subsampling-stage tails (streaming_state_specs), the per-layer post-GLU conv tail (the engine's existing "conv" slot cache — conv_kernel_size - 1 is exactly what this encoder needs), the prefilled paged K/V window (fixed_attention_window), and the frontend frame grid (on the extractor, not here). The model adds a fifth slot state, the stream's prompt id, which describes the request rather than the audio.

Two alignment preconditions, both silent if violated, both refused at engine construction by streaming_geometry:

1. chunk_size must be a whole number of trained attention chunks (num_lookahead_tokens + 1) — the mask groups absolute frame positions, so a partial trained chunk would need keys from the future.
2. The resulting feature window must be a multiple of the subsampling factor, so every stage's input length is a multiple of its stride.

Whisper (Hugging Face)

model.py (encoder + decoder in HF key layout), config.py (including generation-control ids), convert.py::HFWhisperConverter (auto-detects config.json: model_type=whisper, reads model.safetensors + generation_config.json).

Offline-only (streaming_kind="none"), decoded by the incremental aed strategy with suppress / begin-suppress lists and the SOT prompt from the config. The decoder exposes a batched incremental prefill / step / select surface.

fc1 / fc2 stay flat because HF's layout puts them on the layer; fc1 is an exact-erf LinearActivation("gelu"), which preserves the flat checkpoint keys while applying GELU in the GEMM epilogue. The two frontend convs use standalone Gelu over their BTC output.

Paraformer (FunASR, non-autoregressive)

• encoder.py — SANM encoder: FSMN-memory self-attention over the shared oasr.layers.Attention, 560 → 512 first layer, sinusoidal PE, LayerNorm eps 1e-12 (ESPnet convention; parity breaks at PyTorch's 1e-5).
• predictor.py — CifPredictor (CifPredictorV2) with the vectorized cif_v1 prefix-sum integrate-and-fire, always fp32. Sigmoid runs over the contiguous aligned alpha head before selecting its real column; the post-sigmoid ReLU was redundant and is omitted.
• decoder.py — SANM NAR decoder: FFN-first layer order, FSMN-only self-attention, one parallel pass over the CIF acoustic embeddings. Cross-attention takes encoder lengths, not a mask.
• convert.py::FunASRParaformerConverter — auto-detects config.yaml: model: Paraformer, parses am.mvn CMVN into synthetic encoder.cmvn_shift / cmvn_scale state-dict buffers so native round-trip is automatic.

CIF fire positions become per-token RequestOutput.timestamps.

speech_llm (Qwen2-Audio-style LLM-ASR)

• audio_tower.py — Whisper-geometry encoder (×32, d=1280) with key-padding-only attention mask + AvgPool1d(2) + post-pool LayerNorm. Valid lengths follow HF's two-stage formula: (mel − 1) // 2 + 1, then (feat − 2) // 2 + 1.
• llm.py::Qwen2Lm — a faithful Qwen2 causal LM built on oasr.layers (RMSNorm, NeoxRotaryEmbedding, GatedMLP, Attention). RMSNorm accumulates in fp32 to the store where HF rounds before the weight multiply — the more accurate order, one rounding step apart under bf16. prefill(inputs_embeds, valid[, capacity]) / step / select handle left-padded variable-length prompts (per-row positions = cumsum(mask) − 1, HF's masked-generate convention). With capacity, the per-layer K/V buffers are preallocated and each step writes its slot in place.
• model.py — module names mirror HF 1:1 (audio_tower / multi_modal_projector / language_model); load_weights normalizes both the 4.x (language_model.model.*) and 5.x-resave (language_model.model.model.*) key layouts.
• convert.py::HFQwen2AudioConverter — auto-detects config.json: model_type=qwen2_audio, fills omitted text_config fields from Qwen2Config defaults (the published 7B relies on them), and reads sharded safetensors.

Checkpoint bundles load host-side (map_location="cpu") with the dtype cast before the device move — an 8.4B checkpoint double-books the GPU otherwise.

Pre-norm residual streams use the waist's fused add+norm operations: Whisper, the Qwen2-Audio tower, Paraformer, and Nemotron reuse add+LayerNorm, while the Qwen2 LM uses add+RMSNorm. The residual-returning forms carry the unnormalized sum into the next sublayer without a separate elementwise launch. Zipformer is intentionally different: its learned bypass is a lerp and its endpoint is BiasNorm, so it cannot use the LayerNorm/RMSNorm fusion without changing model semantics.

Adding an architecture

See the cookbook in architecture.md. In brief:

1. class FooEncoder(BaseEncoder) — forward, the introspection properties, and (for streaming) either forward_chunk_paged or get_streaming_init_states / streaming_forward.
2. class FooModel(BaseAsrModel) with from_config + load_weights.
3. A CheckpointConverter + register_model("foo", ...) in the package __init__ — or an oasr.models entry point, out of tree.
4. Build it from oasr.layers and add a tiny config to tests/models/test_layer_waist.py.